[Example] Fix TileKernels Benchmarking

examples/mhc/normw_merge.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""mHC norm-weight merge (forward).
+
+This is the load-bearing elementwise mHC primitive: ``out = fn * normw``.
+"""
+
+import tilus
+import torch
+from tile_kernels.mhc.norm_fn_kernel import _mhc_fn_normw_merge_fwd
+from tilus import float32, int32
+from tilus.utils import benchmark_func, cdiv
+
+
+class MHCNormWMerge(tilus.Script):
+    # One 256-wide row per CTA and two values per lane. Coalescing rows leaves
+    # this small mHC launch with too few CTAs, while 256 threads makes each
+    # lane handle only one value.
+    def __init__(self, block_m: int = 1, block_n: int = 256, warps: int = 4):
+        super().__init__()
+        self.block_m, self.block_n, self.warps = block_m, block_n, warps
+
+    def __call__(
+        self, m: int32, n: int32, fn_ptr: ~float32, normw_ptr: ~float32, out_ptr: ~float32
+    ):
+        self.attrs.blocks = (cdiv(m, self.block_m), cdiv(n, self.block_n))
+        self.attrs.warps = self.warps
+        fn = self.global_view(fn_ptr, dtype=float32, shape=[m, n])
+        normw = self.global_view(normw_ptr, dtype=float32, shape=[n])
+        out = self.global_view(out_ptr, dtype=float32, shape=[m, n])
+        rows = self.blockIdx.x * self.block_m
+        cols = self.blockIdx.y * self.block_n
+        r_fn = self.load_global(
+            fn, offsets=[rows, cols], shape=[self.block_m, self.block_n]
+        )
+        r_w = self.load_global(normw, offsets=[cols], shape=[self.block_n])
+        self.store_global(out, r_fn * r_w, offsets=[rows, cols])
+
+
+def main():
+    # These are the mHC shapes used by the TileKernels norm-fn path: fn is
+    # [mhc_mult ** 3, mhc_hidden_size], normw is [mhc_hidden_size].
+    m, n = 24, 4096
+    fn = torch.randn((m, n), device="cuda", dtype=torch.float32)
+    normw = torch.randn((n,), device="cuda", dtype=torch.float32)
+    out = torch.empty_like(fn)
+    kernel = MHCNormWMerge()
+    kernel(m, n, fn, normw, out)
+    tile_out = torch.empty_like(fn)
+    tile_kernel = _mhc_fn_normw_merge_fwd(m, n)
+    tile_kernel(fn, normw, tile_out)
+    torch.testing.assert_close(out, fn * normw, rtol=0, atol=0)
+    torch.testing.assert_close(out, tile_out, rtol=0, atol=0)
+
+    def run_tilus():
+        kernel(m, n, fn, normw, out)
+
+    def run_tilekernels():
+        tile_kernel(fn, normw, tile_out)
+
+    tilus_ms = benchmark_func(run_tilus, warmup=10, repeat=100)
+    tilekernels_ms = benchmark_func(run_tilekernels, warmup=10, repeat=100)
+    print(
+        f"mHC normw merge: Tilus {tilus_ms:.4f} ms, TileKernels {tilekernels_ms:.4f} ms"
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/moe/topk_gate.py

@@ -0,0 +1,140 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Stable MoE top-k routing, ported from TileKernels' ``topk_gate``.
+
+Each CTA owns one token.  Repeated max/min reductions retain TileKernels'
+tie rule: when equal scores occur, the lower expert index wins.
+"""
+
+import tilus
+import torch
+from tile_kernels.moe.topk_gate_kernel import topk_gate
+from tilus import float32, int32, int64
+from tilus.utils import benchmark_func, cdiv
+
+
+class TopKGate(tilus.Script):
+    def __init__(self, num_experts: int, num_topk: int):
+        super().__init__()
+        assert 1 <= num_topk <= num_experts
+        self.num_experts = num_experts
+        self.num_topk = num_topk
+        self.aligned_experts = cdiv(num_experts, 32) * 32
+
+    def __call__(self, num_tokens: int32, scores_ptr: ~float32, output_ptr: ~int64):
+        self.attrs.blocks = (num_tokens,)
+        self.attrs.warps = 1
+
+        scores = self.global_view(
+            scores_ptr, dtype=float32, shape=[num_tokens, self.num_experts]
+        )
+        output = self.global_view(
+            output_ptr, dtype=int64, shape=[num_tokens, self.num_topk]
+        )
+        token = self.blockIdx.x
+
+        values = self.load_global(
+            scores, offsets=[token, 0], shape=[1, self.aligned_experts]
+        )
+        # TileKernels performs the stable reducer in int32 and widens only at
+        # the required int64 output boundary.
+        expert_ids = self.register_tensor(
+            dtype=int32, shape=[1, self.aligned_experts], init=lambda _, j: j
+        )
+        padding_value = -3.402823466e38
+        # Keep selection state separate from score values: neither -inf nor
+        # any finite sentinel is safe to use as an in-band removed marker.
+        if self.aligned_experts != self.num_experts:
+            active = self.register_tensor(
+                dtype=int32,
+                shape=[1, self.aligned_experts],
+                init=lambda _, j: j < self.num_experts,
+            )
+        else:
+            active = self.register_tensor(
+                dtype=int32, shape=[1, self.aligned_experts], init=lambda _i, _j: 1
+            )
+        # Out-of-bounds vector-load lanes are not necessarily initialized to a
+        # value below every valid score.  Mask them before the first max.
+        values = self.where(active != 0, x=values, y=padding_value)
+
+        # ``num_topk`` is a specialization constant.  TileKernels unrolls
+        # this selection loop; retain that property in the generated CUDA.
+        for rank in self.range(0, self.num_topk, 1, unroll="all"):
+            best_value = self.max(values, dim=1, keepdim=True)
+            # ``min`` over matching candidates is the stable tie breaker.
+            candidates = self.where(
+                active != 0,
+                x=self.where(values == best_value, x=expert_ids, y=int32.max_value),
+                y=int32.max_value,
+            )
+            best_index = self.min(candidates, dim=1, keepdim=True)
+            # A sentinel maximum may also match a valid finite score.  Fall
+            # back only when no active candidate matches the maximum (the
+            # remaining scores are -inf).  This uniform branch avoids a second
+            # warp reduction in the normal finite-score path.
+            if best_index[0, 0].item() == int32.max_value:
+                best_index = self.min(
+                    self.where(active != 0, x=expert_ids, y=int32.max_value),
+                    dim=1,
+                    keepdim=True,
+                )
+            # The reduction result is replicated across the warp.  A direct
+            # store from every lane creates 32 identical global writes; only
+            # lane 0 owns the scalar output (as in TileKernels' shared-output
+            # epilogue).
+            if self.get_thread_binding() == 0:
+                self.store_global(output, best_index.to(int64), offsets=[token, rank])
+            active = self.where(expert_ids == best_index, x=0, y=active)
+            values = self.where(expert_ids == best_index, x=padding_value, y=values)
+
+
+def main():
+    rows = []
+    for num_tokens, num_experts, num_topk in [
+        (128, 72, 6),
+        (1024, 256, 8),
+        (8192, 256, 8),
+    ]:
+        scores = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.float32)
+        # Make ties observable: the implementation must choose the lower index.
+        scores[:, 0] = scores[:, 1]
+        kernel = TopKGate(num_experts, num_topk)
+        output = torch.empty(num_tokens, num_topk, device="cuda", dtype=torch.int64)
+        kernel(num_tokens, scores, output)
+        expected = torch.sort(scores, dim=1, descending=True, stable=True).indices[
+            :, :num_topk
+        ]
+        torch.testing.assert_close(output, expected)
+        torch.testing.assert_close(topk_gate(scores, num_topk), expected)
+        # Padding must never win when valid experts contain -inf.
+        if num_experts % 32:
+            edge_scores = torch.full_like(scores, float("-inf"))
+            kernel(num_tokens, edge_scores, output)
+            edge_expected = torch.arange(num_topk, device="cuda", dtype=torch.int64)[
+                None, :
+            ].expand(num_tokens, -1)
+            torch.testing.assert_close(output, edge_expected)
+            # The minimum finite score must beat -inf even though it equals
+            # the padding sentinel.
+            edge_scores[:, 1] = torch.finfo(torch.float32).min
+            kernel(num_tokens, edge_scores, output)
+            edge_expected = torch.sort(
+                edge_scores, dim=1, descending=True, stable=True
+            ).indices[:, :num_topk]
+            torch.testing.assert_close(output, edge_expected)
+
+        def run_tilus():
+            out = torch.empty(num_tokens, num_topk, device="cuda", dtype=torch.int64)
+            kernel(num_tokens, scores, out)
+            return out
+
+        tilus_ms = benchmark_func(run_tilus)
+        tilekernels_ms = benchmark_func(lambda: topk_gate(scores, num_topk))
+        rows.append((num_tokens, num_experts, num_topk, tilus_ms, tilekernels_ms))
+    for row in rows:
+        print("tokens=%d experts=%d topk=%d: Tilus %.4f ms, TileKernels %.4f ms" % row)
+
+
+if __name__ == "__main__":
+    main()

examples/quantization/per_channel_cast.py

@@ -0,0 +1,133 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Common BF16 128x128 path of TileKernels' per-channel E4M3 cast."""
+
+import tilus
+import torch
+from tile_kernels.quant.per_channel_cast_kernel import per_channel_cast
+from tilus import bfloat16, float8_e4m3, float32, int32
+from tilus.ir.layout.ops import spatial
+from tilus.utils import benchmark_func, cdiv
+
+E4M3_MAX = 448.0
+SF_CLAMP_MIN = 1.0e-4
+
+
+def check_fp8_close(
+    actual: torch.Tensor,
+    expected: torch.Tensor,
+    *,
+    label: str,
+    max_mismatch_frac: float = 0.01,
+) -> None:
+    """Check e4m3 values on their integer code ladder."""
+    actual_bits = actual.view(torch.uint8).to(torch.int32)
+    expected_bits = expected.view(torch.uint8).to(torch.int32)
+    actual_codes = torch.where(
+        actual_bits & 0x80 != 0, -(actual_bits & 0x7F), actual_bits & 0x7F
+    )
+    expected_codes = torch.where(
+        expected_bits & 0x80 != 0, -(expected_bits & 0x7F), expected_bits & 0x7F
+    )
+    code_diff = (actual_codes - expected_codes).abs()
+    max_code_diff = int(code_diff.max().item())
+    mismatch_frac = float((code_diff != 0).to(torch.float64).mean().item())
+    assert max_code_diff <= 1, (
+        f"{label}: {max_code_diff} e4m3 codes apart at worst; expected at most 1"
+    )
+    assert mismatch_frac <= max_mismatch_frac, (
+        f"{label}: {mismatch_frac:.4%} of elements differ"
+    )
+
+
+class PerChannelCast(tilus.Script):
+    """BF16 input, 128-token groups, and hidden sizes divisible by 128 only."""
+
+    def __init__(self):
+        super().__init__()
+        self.block_m, self.block_n, self.warps = 128, 128, 8
+
+    def __call__(
+        self,
+        num_tokens: int32,
+        hidden: int32,
+        x_ptr: ~bfloat16,
+        out_ptr: ~float8_e4m3,
+        sf_ptr: ~float32,
+    ):
+        self.attrs.blocks = (cdiv(num_tokens, self.block_m), cdiv(hidden, self.block_n))
+        self.attrs.warps = self.warps
+        self.assume(num_tokens % self.block_m == 0)
+        self.assume(hidden % self.block_n == 0)
+        x = self.global_view(x_ptr, dtype=bfloat16, shape=[num_tokens, hidden])
+        out = self.global_view(out_ptr, dtype=float8_e4m3, shape=[num_tokens, hidden])
+        sf = self.global_view(
+            sf_ptr, dtype=float32, shape=[cdiv(num_tokens, self.block_m), hidden]
+        )
+        offset_m, offset_n = (
+            self.blockIdx.x * self.block_m,
+            self.blockIdx.y * self.block_n,
+        )
+        # Keep the BF16 tile in shared memory across the reduction.  Reloading
+        # it for the FP8 epilogue avoids keeping both the complete BF16 and
+        # FP32 tiles live in registers, matching TileKernels' lifetime.
+        shared_x = self.shared_tensor(dtype=bfloat16, shape=[self.block_m, self.block_n])
+        input_values = self.load_global(
+            x, offsets=[offset_m, offset_n], shape=[self.block_m, self.block_n]
+        )
+        self.annotate_layout(input_values, spatial(8, 32).local(16, self.block_n // 32))
+        self.store_shared(shared_x, input_values)
+        self.sync()
+        values_for_reduce = self.abs(self.load_shared(shared_x)).to(float32)
+        # TileKernels maps 8 warps over [8, 32] and gives each thread a
+        # [16, 4] micro-tile.  The generic inferred layout instead made every
+        # thread retain an entire 128-value column; make this mapping explicit.
+        self.annotate_layout(
+            values_for_reduce, spatial(8, 32).local(16, self.block_n // 32)
+        )
+        amax = self.max(values_for_reduce, dim=0, keepdim=True)
+        amax = self.where(amax > SF_CLAMP_MIN, x=amax, y=SF_CLAMP_MIN)
+        scale = amax / E4M3_MAX
+        inv_scale = (
+            self.register_tensor(dtype=float32, shape=[1, self.block_n], init=E4M3_MAX)
+            / amax
+        )
+        self.store_global(sf, scale, offsets=[self.blockIdx.x, offset_n])
+        output_values = self.load_shared(shared_x)
+        self.annotate_layout(output_values, spatial(8, 32).local(16, self.block_n // 32))
+        output_values_f32 = output_values.to(float32) * inv_scale
+        output_values_fp8 = output_values_f32.to(float8_e4m3)
+        self.store_global(out, output_values_fp8, offsets=[offset_m, offset_n])
+        self.free_shared(shared_x)
+
+
+def main():
+    tokens, hidden = 8192, 8192
+    x = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16).contiguous()
+    out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    sf = torch.empty(tokens // 128, hidden, device="cuda", dtype=torch.float32)
+    kernel = PerChannelCast()
+    kernel(tokens, hidden, x, out, sf)
+    ref_amax = x.float().abs().reshape(-1, 128, hidden).amax(1).clamp_min(SF_CLAMP_MIN)
+    torch.testing.assert_close(sf, ref_amax / E4M3_MAX, rtol=1e-5, atol=1e-7)
+    tk_out, tk_sf = per_channel_cast(x, "e4m3", 128)
+    torch.testing.assert_close(sf, tk_sf, rtol=1e-5, atol=1e-7)
+    check_fp8_close(out, tk_out, label="per-channel Tilus vs TileKernels")
+
+    def run_tilus():
+        out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+        sf = torch.empty(tokens // 128, hidden, device="cuda", dtype=torch.float32)
+        kernel(tokens, hidden, x, out, sf)
+        return out, sf
+
+    tilus_ms = benchmark_func(run_tilus, warmup=10, repeat=50)
+    tilekernels_ms = benchmark_func(
+        lambda: per_channel_cast(x, "e4m3", 128), warmup=10, repeat=50
+    )
+    print(
+        f"Per-channel FP8 cast: Tilus {tilus_ms:.4f} ms, TileKernels {tilekernels_ms:.4f} ms"
+    )
+
+
+if __name__ == "__main__":
+    main()