Add cross_entropy example and unit test

yf225 · yf225 · commit 28a597822f1a · 2025-07-15T00:11:54.000-07:00
stack-info: PR: #320, branch: yf225/stack/28
diff --git a/examples/cross_entropy.py b/examples/cross_entropy.py
@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import os
+
+import torch
+
+import helion
+from helion._testing import run_example
+import helion.language as hl
+
+# TritonBench configuration - adjust based on HELION_DEV_LOW_VRAM environment variable
+if os.environ.get("HELION_DEV_LOW_VRAM", "0") == "1":
+    # Low memory configuration
+    TRITONBENCH_ARGS = {"B": 4, "T": 512, "v_range": "10,15"}
+
+
+@helion.kernel(ignore_warnings=[helion.exc.TensorOperationInWrapper])
+def cross_entropy(
+    logits: torch.Tensor,  # [N, V] input logits
+    labels: torch.Tensor,  # [N] target labels
+) -> torch.Tensor:
+    n, v = logits.shape
+    losses = torch.zeros([n], dtype=logits.dtype, device=logits.device)
+
+    # Flatten logits once at the beginning
+    logits_flat = logits.view(-1)
+
+    for tile_n in hl.tile(n):
+        # Get data for this tile
+        labels_tile = labels[tile_n]  # [tile_size]
+        base_indices_tile = tile_n.index * v  # [tile_size]
+
+        # Compute the actual flat indices by adding the label offset
+        # flat_index[i] = base_indices[i] + labels[i] = i*V + labels[i]
+        flat_indices = base_indices_tile + labels_tile
+
+        # Load the logits at the target indices
+        logits_at_target = hl.load(logits_flat, [flat_indices])
+
+        # Compute log_softmax for numerical stability
+        # Load the full rows for this tile
+        logits_rows = logits[tile_n, :]  # [tile_size, V]
+
+        # Compute log-sum-exp
+        max_logits = torch.amax(logits_rows, dim=-1, keepdim=True)
+        shifted = logits_rows - max_logits
+        exp_shifted = torch.exp(shifted)
+        sum_exp = torch.sum(exp_shifted, dim=-1, keepdim=True)
+        log_sum_exp = max_logits.squeeze(-1) + torch.log(sum_exp.squeeze(-1))
+
+        # Cross entropy loss: log_sum_exp - logit_at_target
+        losses[tile_n] = log_sum_exp - logits_at_target
+
+    return losses.mean()
+
+
+def main() -> None:
+    """Run cross entropy benchmark with different input sizes."""
+    # Test with moderate size
+    n, v = 128, 1000
+    logits = torch.randn(n, v, device="cuda", dtype=torch.float32)
+    labels = torch.randint(0, v, (n,), device="cuda", dtype=torch.long)
+
+    run_example(
+        cross_entropy,
+        torch.nn.functional.cross_entropy,
+        (logits, labels),
+        kernel_name="helion",
+        baseline_name="torch",
+        rtol=1e-4,
+        atol=1e-4,
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/test/test_examples.expected b/test/test_examples.expected
@@ -469,6 +469,50 @@ def concat2d_dim1(x: torch.Tensor, y: torch.Tensor, *, _launcher=_default_launch
     _launcher(_concat2d_dim1_kernel, (triton.cdiv(x.size(0), _BLOCK_SIZE_0) * triton.cdiv(out.size(1), _BLOCK_SIZE_1),), x, out, y, out.size(0), out.size(1), x.size(0), x.size(1), out.stride(0), out.stride(1), x.stride(0), x.stride(1), y.stride(0), y.stride(1), _BLOCK_SIZE_0, _BLOCK_SIZE_1, num_warps=4, num_stages=3)
     return out
 
+--- assertExpectedJournal(TestExamples.test_cross_entropy)
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from torch._inductor.runtime.triton_helpers import math as tl_math
+from helion.runtime import default_launcher as _default_launcher
+
+@triton.jit
+def _cross_entropy_kernel(labels, logits_flat, logits, losses, labels_stride_0, logits_stride_0, logits_stride_1, logits_flat_stride_0, losses_stride_0, v, _RDIM_SIZE_1: tl.constexpr):
+    pid_0 = tl.program_id(0)
+    offset_0 = pid_0
+    indices_0 = offset_0 + tl.zeros([1], tl.int32)
+    indices_1 = tl.arange(0, _RDIM_SIZE_1).to(tl.int32)
+    mask_1 = indices_1 < v
+    labels_tile = tl.load(labels + indices_0 * labels_stride_0, None)
+    v_0 = v.to(tl.int32)
+    v_1 = indices_0 * v_0
+    v_2 = v_1.to(tl.int64)
+    v_3 = v_2 + labels_tile
+    logits_at_target = tl.load(logits_flat + v_3 * logits_flat_stride_0, None)
+    logits_rows = tl.load(logits + (indices_0[:, None] * logits_stride_0 + indices_1[None, :] * logits_stride_1), mask_1[None, :], other=0)
+    _mask_to = tl.where(tl.broadcast_to(mask_1[None, :], [1, _RDIM_SIZE_1]), logits_rows, float('-inf'))
+    max_logits = tl.reshape(tl.max(_mask_to, 1), [1, 1])
+    v_4 = logits_rows - max_logits
+    v_5 = tl_math.exp(v_4)
+    _mask_to_1 = tl.where(tl.broadcast_to(mask_1[None, :], [1, _RDIM_SIZE_1]), v_5, 0)
+    sum_exp = tl.reshape(tl.sum(_mask_to_1, 1), [1, 1])
+    squeeze = tl.reshape(max_logits, [1])
+    squeeze_1 = tl.reshape(sum_exp, [1])
+    v_6 = tl_math.log(squeeze_1)
+    v_7 = squeeze + v_6
+    v_8 = v_7 - logits_at_target
+    tl.store(losses + indices_0 * losses_stride_0, v_8, None)
+
+def cross_entropy(logits: torch.Tensor, labels: torch.Tensor, *, _launcher=_default_launcher):
+    n, v = logits.shape
+    losses = torch.zeros([n], dtype=logits.dtype, device=logits.device)
+    logits_flat = logits.view(-1)
+    _RDIM_SIZE_1 = triton.next_power_of_2(v)
+    _launcher(_cross_entropy_kernel, (n,), labels, logits_flat, logits, losses, labels.stride(0), logits.stride(0), logits.stride(1), logits_flat.stride(0), losses.stride(0), v, _RDIM_SIZE_1, num_warps=4, num_stages=3)
+    return losses.mean()
+
 --- assertExpectedJournal(TestExamples.test_embedding_block_ptr)
 from __future__ import annotations
 
@@ -530,6 +574,94 @@ def embedding(x: torch.Tensor, weight: torch.Tensor, *, _launcher=_default_launc
     _launcher(_embedding_kernel, (x_flat.size(0) * triton.cdiv(embedding_dim, _BLOCK_SIZE_1),), x_flat, weight, out, x_flat.size(0), out.stride(0), out.stride(1), weight.stride(0), weight.stride(1), x_flat.stride(0), embedding_dim, _BLOCK_SIZE_1, num_warps=4, num_stages=3)
     return out.view(*x.size(), embedding_dim)
 
+--- assertExpectedJournal(TestExamples.test_fp8_attention)
+from __future__ import annotations
+
+import math
+import torch
+import triton
+import triton.language as tl
+from torch._inductor.runtime import triton_helpers
+from torch._inductor.runtime.triton_compat import libdevice
+
+@triton.jit
+def _fp8_attention_kernel_kernel(q, k, v, out, _RDIM_SIZE_2: tl.constexpr, _BLOCK_SIZE_1: tl.constexpr, _BLOCK_SIZE_3: tl.constexpr):
+    pid_0 = tl.program_id(0)
+    offset_0 = pid_0
+    indices_5 = tl.arange(0, _RDIM_SIZE_2).to(tl.int32)
+    for offset_4 in tl.range(0, 256, _BLOCK_SIZE_1):
+        indices_4 = offset_4 + tl.arange(0, _BLOCK_SIZE_1).to(tl.int32)
+        m_i = tl.full([_BLOCK_SIZE_1], float('-inf'), tl.float32)
+        l_i = tl.full([_BLOCK_SIZE_1], 0.0, tl.float32)
+        acc = tl.full([_BLOCK_SIZE_1, 64], 0.0, tl.float32)
+        q_tile = tl.load(q + (offset_0 * 16384 + indices_4[:, None] * 64 + indices_5[None, :] * 1), None)
+        for offset_2 in tl.range(0, 256, _BLOCK_SIZE_3):
+            indices_2 = offset_2 + tl.arange(0, _BLOCK_SIZE_3).to(tl.int32)
+            q_tile_copy = q_tile
+            m_i_copy = m_i
+            l_i_copy = l_i
+            acc_copy = acc
+            q_tile_copy_0 = q_tile_copy
+            m_i_copy_0 = m_i_copy
+            l_i_copy_0 = l_i_copy
+            acc_copy_0 = acc_copy
+            k_tile = tl.load(k + (offset_0 * 16384 + indices_2[:, None] * 64 + indices_5[None, :] * 1), None)
+            k_tile_t = tl.permute(k_tile, [1, 0])
+            mm = tl.dot(q_tile_copy_0, k_tile_t)
+            v_0 = mm.to(tl.float32)
+            v_1 = 0.18033688
+            v_2 = v_0 * v_1
+            qk_max = tl.max(v_2, 1)
+            v_3 = triton_helpers.maximum(m_i_copy_0, qk_max)
+            subscript = v_3[:, None]
+            v_4 = v_2 - subscript
+            v_5 = libdevice.exp2(v_4)
+            l_ij = tl.sum(v_5, 1)
+            v_6 = m_i_copy_0 - v_3
+            v_7 = libdevice.exp2(v_6)
+            v_8 = l_i_copy_0 * v_7
+            l_i = v_8 + l_ij
+            subscript_1 = v_7[:, None]
+            v_10 = acc_copy_0 * subscript_1
+            v_tile = tl.load(v + (offset_0 * 16384 + indices_5[:, None] * 1 + indices_2[None, :] * 64), None)
+            v_11 = v_5.to(tl.float8e5)
+            v_t = tl.permute(v_tile, [1, 0])
+            mm_1 = tl.dot(v_11, v_t)
+            v_12 = mm_1.to(tl.float32)
+            acc = v_10 + v_12
+            m_i = v_3
+        subscript_2 = l_i[:, None]
+        v_14 = acc / subscript_2
+        tl.store(out + (offset_0 * 16384 + indices_4[:, None] * 64 + indices_5[None, :] * 1), v_14, None)
+
+def fp8_attention_kernel(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
+    """FP8 attention kernel processing batch*heads in parallel."""
+    batch_heads = q.size(0)
+    seq_len = q.size(1)
+    head_dim = q.size(2)
+    out = torch.empty([batch_heads, seq_len, head_dim], dtype=torch.float32, device=q.device)
+    sm_scale = 1.0 / math.sqrt(float(head_dim))
+    sm_scale = sm_scale * 1.44269504
+    _RDIM_SIZE_2 = 64
+    _BLOCK_SIZE_1 = 64
+    _BLOCK_SIZE_3 = 64
+    _fp8_attention_kernel_kernel[8,](q, k, v, out, _RDIM_SIZE_2, _BLOCK_SIZE_1, _BLOCK_SIZE_3, num_warps=4, num_stages=3)
+    return out
+
+def _fp8_attention_kernel_make_precompiler(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
+    """FP8 attention kernel processing batch*heads in parallel."""
+    batch_heads = q.size(0)
+    seq_len = q.size(1)
+    head_dim = q.size(2)
+    out = torch.empty([batch_heads, seq_len, head_dim], dtype=torch.float32, device=q.device)
+    sm_scale = 1.0 / math.sqrt(float(head_dim))
+    sm_scale = sm_scale * 1.44269504
+    _RDIM_SIZE_2 = 64
+    _BLOCK_SIZE_1 = 64
+    _BLOCK_SIZE_3 = 64
+    from helion.runtime.precompile_shim import make_precompiler
+    return make_precompiler(_fp8_attention_kernel_kernel)(q, k, v, out, _RDIM_SIZE_2, _BLOCK_SIZE_1, _BLOCK_SIZE_3, num_warps=4, num_stages=3)
+
 --- assertExpectedJournal(TestExamples.test_fp8_gemm)
 from __future__ import annotations
 
@@ -762,6 +894,139 @@ def jagged_mean_kernel(x_data: torch.Tensor, x_offsets: torch.Tensor, x_feature_
     _launcher(_jagged_mean_kernel_kernel, (triton.cdiv(num_rows, _BLOCK_SIZE_0),), x_offsets, x_feature_counts, x_flat, out, out.stride(0), out.stride(1), x_feature_counts.stride(0), x_flat.stride(0), x_offsets.stride(0), num_rows, max_M, _BLOCK_SIZE_0, _BLOCK_SIZE_1, _BLOCK_SIZE_2, num_warps=4, num_stages=3)
     return out
 
+--- assertExpectedJournal(TestExamples.test_jagged_mean_2d)
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def _jagged_mean_kernel_2d_kernel(x_offsets, x_feature_counts, x_flat, out, out_stride_0, out_stride_1, x_feature_counts_stride_0, x_flat_stride_0, x_offsets_stride_0, num_rows, max_M, _BLOCK_SIZE_0: tl.constexpr, _BLOCK_SIZE_1: tl.constexpr, _BLOCK_SIZE_2: tl.constexpr):
+    pid_0 = tl.program_id(0)
+    offset_0 = pid_0 * _BLOCK_SIZE_0
+    indices_0 = (offset_0 + tl.arange(0, _BLOCK_SIZE_0)).to(tl.int32)
+    mask_0 = indices_0 < num_rows
+    starts = tl.load(x_offsets + indices_0 * x_offsets_stride_0, mask_0, other=0)
+    v_0 = tl.full([], 1, tl.int32)
+    v_1 = indices_0 + v_0
+    ends = tl.load(x_offsets + v_1 * x_offsets_stride_0, mask_0, other=0)
+    v_2 = ends - starts
+    _mask_to = tl.where(mask_0, v_2, -9223372036854775808)
+    max_nnz = tl.max(_mask_to, 0)
+    feature_counts = tl.load(x_feature_counts + indices_0 * x_feature_counts_stride_0, mask_0, other=0)
+    for offset_1 in tl.range(0, max_M.to(tl.int32), step=_BLOCK_SIZE_1):
+        indices_1 = offset_1 + tl.arange(0, _BLOCK_SIZE_1).to(tl.int32)
+        mask_1 = indices_1 < max_M
+        feature_counts_copy = feature_counts
+        max_nnz_copy = max_nnz
+        starts_copy = starts
+        v_2_copy = v_2
+        feature_counts_copy_0 = feature_counts_copy
+        max_nnz_copy_0 = max_nnz_copy
+        starts_copy_0 = starts_copy
+        v_2_copy_0 = v_2_copy
+        subscript = feature_counts_copy_0[:, None]
+        v_3 = indices_1[None, :]
+        v_4 = v_3 < subscript
+        row_sums = tl.full([_BLOCK_SIZE_0, _BLOCK_SIZE_1], 0.0, tl.float32)
+        for offset_2 in tl.range(0, max_nnz_copy_0.to(tl.int32), step=_BLOCK_SIZE_2):
+            indices_2 = offset_2 + tl.arange(0, _BLOCK_SIZE_2).to(tl.int32)
+            mask_2 = indices_2 < max_nnz_copy_0
+            starts_copy_0_copy = starts_copy_0
+            v_2_copy_0_copy = v_2_copy_0
+            v_4_copy = v_4
+            row_sums_copy = row_sums
+            starts_copy_0_copy_0 = starts_copy_0_copy
+            v_2_copy_0_copy_0 = v_2_copy_0_copy
+            v_4_copy_0 = v_4_copy
+            row_sums_copy_0 = row_sums_copy
+            subscript_1 = starts_copy_0_copy_0[:, None]
+            subscript_2 = indices_2[None, :]
+            v_5 = subscript_2.to(tl.int64)
+            v_6 = subscript_1 + v_5
+            subscript_3 = v_6[:, :, None]
+            v_7 = subscript_3 * max_M
+            subscript_4 = indices_1[None, None, :]
+            v_8 = subscript_4.to(tl.int64)
+            v_9 = v_7 + v_8
+            subscript_5 = indices_2[None, :]
+            subscript_6 = v_2_copy_0_copy_0[:, None]
+            v_10 = subscript_5.to(tl.int64)
+            v_11 = v_10 < subscript_6
+            subscript_7 = v_11[:, :, None]
+            subscript_8 = v_4_copy_0[:, None, :]
+            v_12 = subscript_7 & subscript_8
+            x_slice = tl.load(x_flat + v_9 * x_flat_stride_0, mask_0[:, None, None] & mask_2[None, :, None] & mask_1[None, None, :] & v_12, other=0)
+            sum_1 = tl.sum(x_slice, 1)
+            row_sums = row_sums_copy_0 + sum_1
+        v_14 = v_2_copy_0.to(tl.float32)
+        nnz_expanded = v_14[:, None]
+        v_15 = 0.0
+        v_16 = nnz_expanded > v_15
+        v_17 = row_sums / nnz_expanded
+        v_18 = 0.0
+        v_19 = v_18[None, None]
+        v_20 = tl.where(v_16, v_17, v_19)
+        v_21 = 0.0
+        v_22 = v_21[None, None]
+        v_23 = tl.where(v_4, v_20, v_22)
+        tl.store(out + (indices_0[:, None] * out_stride_0 + indices_1[None, :] * out_stride_1), v_23, mask_0[:, None] & mask_1[None, :])
+
+def jagged_mean_kernel_2d(x_data: torch.Tensor, x_offsets: torch.Tensor, x_feature_counts: torch.Tensor, max_M_tensor: torch.Tensor):
+    """
+    Compute the mean of each row in a 2D jagged tensor with variable features per row.
+
+    Args
+    ----
+    x_data          : 2-D tensor of shape (total_elements, max_M) holding all elements.
+    x_offsets       : (num_rows + 1) tensor. Row i is the slice
+                      x_data[x_offsets[i] : x_offsets[i+1], :].
+    x_feature_counts: (num_rows) tensor. Number of valid features for each row.
+    max_M_tensor    : Dummy tensor whose numel() gives max number of features.
+
+    Returns
+    -------
+    result : 2-D tensor of shape (num_rows, max_M) containing the mean of each row.
+             Invalid features (beyond x_feature_counts[i]) are set to 0.
+    """
+    num_rows = x_offsets.size(0) - 1
+    max_M = max_M_tensor.numel()
+    out = torch.zeros([num_rows, max_M], dtype=x_data.dtype, device=x_data.device)
+    x_flat = x_data.view(-1)
+    _BLOCK_SIZE_0 = 16
+    _BLOCK_SIZE_1 = 8
+    _BLOCK_SIZE_2 = 16
+    _jagged_mean_kernel_2d_kernel[triton.cdiv(num_rows, _BLOCK_SIZE_0),](x_offsets, x_feature_counts, x_flat, out, out.stride(0), out.stride(1), x_feature_counts.stride(0), x_flat.stride(0), x_offsets.stride(0), num_rows, max_M, _BLOCK_SIZE_0, _BLOCK_SIZE_1, _BLOCK_SIZE_2, num_warps=4, num_stages=3)
+    return out
+
+def _jagged_mean_kernel_2d_make_precompiler(x_data: torch.Tensor, x_offsets: torch.Tensor, x_feature_counts: torch.Tensor, max_M_tensor: torch.Tensor):
+    """
+    Compute the mean of each row in a 2D jagged tensor with variable features per row.
+
+    Args
+    ----
+    x_data          : 2-D tensor of shape (total_elements, max_M) holding all elements.
+    x_offsets       : (num_rows + 1) tensor. Row i is the slice
+                      x_data[x_offsets[i] : x_offsets[i+1], :].
+    x_feature_counts: (num_rows) tensor. Number of valid features for each row.
+    max_M_tensor    : Dummy tensor whose numel() gives max number of features.
+
+    Returns
+    -------
+    result : 2-D tensor of shape (num_rows, max_M) containing the mean of each row.
+             Invalid features (beyond x_feature_counts[i]) are set to 0.
+    """
+    num_rows = x_offsets.size(0) - 1
+    max_M = max_M_tensor.numel()
+    out = torch.zeros([num_rows, max_M], dtype=x_data.dtype, device=x_data.device)
+    x_flat = x_data.view(-1)
+    _BLOCK_SIZE_0 = 16
+    _BLOCK_SIZE_1 = 8
+    _BLOCK_SIZE_2 = 16
+    from helion.runtime.precompile_shim import make_precompiler
+    return make_precompiler(_jagged_mean_kernel_2d_kernel)(x_offsets, x_feature_counts, x_flat, out, out.stride(0), out.stride(1), x_feature_counts.stride(0), x_flat.stride(0), x_offsets.stride(0), num_rows, max_M, _BLOCK_SIZE_0, _BLOCK_SIZE_1, _BLOCK_SIZE_2, num_warps=4, num_stages=3)
+
 --- assertExpectedJournal(TestExamples.test_matmul)
 from __future__ import annotations
 
diff --git a/test/test_examples.py b/test/test_examples.py
@@ -267,6 +267,20 @@ def test_softmax_two_pass_block_ptr(self):
             )
         )
 
+    def test_cross_entropy(self):
+        n, v = 128, 1000
+        args = (
+            torch.randn(n, v, device=DEVICE, dtype=torch.float32),
+            torch.randint(0, v, (n,), device=DEVICE, dtype=torch.long),
+        )
+        self.assertExpectedJournal(
+            check_example(
+                "cross_entropy",
+                args,
+                torch.nn.functional.cross_entropy(*args),
+            )
+        )
+
     def test_rms_norm(self):
         args = (
             torch.randn([128, 256], device=DEVICE, dtype=torch.float16),
