From 65375d2fab16197cd266bf31b50a08d5c53045e3 Mon Sep 17 00:00:00 2001
From: William Zhang <wzhang20@yahoo.com>
Date: Fri, 28 Aug 2026 20:56:59 -0400
Subject: [PATCH 1/5] fix tile kernels benchmarking+speed

Signed-off-by: William Zhang <wzhang20@yahoo.com>
---
 examples/quantization/fp8_check.py            | 155 ++++++++++++
 examples/quantization/per_token_cast.py       | 221 ++++++++++------
 .../swiglu_forward_and_per_token_cast.py      | 239 ++++++++++--------
 tests/examples/test_examples.py               |   2 +
 4 files changed, 433 insertions(+), 184 deletions(-)
 create mode 100644 examples/quantization/fp8_check.py

diff --git a/examples/quantization/fp8_check.py b/examples/quantization/fp8_check.py
new file mode 100644
index 00000000..4d4104e7
--- /dev/null
+++ b/examples/quantization/fp8_check.py
@@ -0,0 +1,155 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Shared reference and correctness checks for the per-token FP8 cast examples.
+
+FP8 e4m3 has a three-bit mantissa, so adjacent representable values are up to
+6.25% apart and the ladder tops out at 448.  Comparing *decoded* e4m3 values
+with an absolute float tolerance is therefore nearly vacuous: a tolerance loose
+enough to absorb a single rounding step near the top of the range is +/-32,
+which is also loose enough to accept a kernel that is quantizing against a
+completely wrong scale factor.
+
+These helpers compare on the code ladder instead.  e4m3 is stored
+sign-magnitude, so ranking the seven magnitude bits and re-applying the sign
+gives an integer ordinal in which *any* two adjacent representable values
+differ by exactly one, uniformly across the dynamic range.  "Agrees to within
+one rounding step" is then ``|ordinal(a) - ordinal(b)| <= 1``.
+
+One rounding step is the tightest claim that holds here.  Tilus compiles every
+kernel with ``-prec-div=false`` (see ``python/tilus/hidet/backend/build.py``),
+so the fp32 divisions that produce the scale factor and its reciprocal are
+approximate to about one ulp.  TileLang uses exact division.  With identical
+inputs and otherwise identical arithmetic the two therefore still straddle an
+e4m3 rounding boundary on a small fraction of elements; everything else is
+bit-identical.  A real bug moves the mismatch rate or the code distance well
+outside those bounds.
+"""
+
+from typing import NamedTuple
+
+import torch
+
+# Largest finite magnitude representable in e4m3 (TileKernels' ``T.max_value``).
+E4M3_MAX = 448.0
+
+# TileKernels clamps the group absmax from below before dividing, so that an
+# all-zero group yields a tiny scale rather than a division by zero.  See
+# ``CastOutputConfig.clamp_min_value`` in ``tile_kernels/quant/common.py``.
+SF_CLAMP_MIN = 1e-4
+
+
+class CodeLadderStats(NamedTuple):
+    """Result of comparing two e4m3 tensors on the code ladder."""
+
+    max_code_diff: int
+    mismatch_frac: float
+
+
+def fp8_ordinal(x: torch.Tensor) -> torch.Tensor:
+    """Rank each e4m3 value on the ladder of representable values.
+
+    e4m3 is sign-magnitude, so the seven magnitude bits are already a monotone
+    rank within one sign.  Negating them for the negative half gives a single
+    signed ordinal where consecutive representable values always differ by one.
+    """
+    assert x.dtype == torch.float8_e4m3fn
+    bits = x.view(torch.uint8).to(torch.int32)
+    magnitude = bits & 0x7F
+    return torch.where(bits & 0x80 != 0, -magnitude, magnitude)
+
+
+def check_fp8_close(
+    actual: torch.Tensor,
+    expected: torch.Tensor,
+    *,
+    label: str,
+    max_mismatch_frac: float = 0.01,
+) -> CodeLadderStats:
+    """Assert two e4m3 tensors agree to within one rounding step.
+
+    Fails if any element is more than one representable value away, or if more
+    than ``max_mismatch_frac`` of elements disagree at all.  The second bound
+    matters: a systematic error that happens to be small still shows up as a
+    mismatch rate far above the ~0.1% produced by fp32 division rounding.
+    """
+    code_diff = (fp8_ordinal(actual) - fp8_ordinal(expected)).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 "
+        f"1 (a single rounding step)"
+    )
+    assert mismatch_frac <= max_mismatch_frac, (
+        f"{label}: {mismatch_frac:.4%} of elements differ, above the "
+        f"{max_mismatch_frac:.4%} budget for fp32 division rounding"
+    )
+    return CodeLadderStats(max_code_diff, mismatch_frac)
+
+
+def check_scales_close(
+    actual: torch.Tensor,
+    expected: torch.Tensor,
+    *,
+    label: str,
+    rtol: float = 1e-6,
+) -> float:
+    """Assert two fp32 scale-factor tensors agree to a few fp32 ulps.
+
+    The scale is ``max(absmax, 1e-4) / 448`` on both sides, computed from an
+    absmax that is a max over identical inputs and so is bit-identical.  Only
+    the division differs, hence a tolerance in ulps (1 ulp ~ 1.2e-7 relative)
+    rather than the 1e-5 that would also pass a wrong reduction.
+    """
+    torch.testing.assert_close(actual, expected, rtol=rtol, atol=0.0, msg=label)
+    return float(((actual - expected).abs() / expected.abs()).max().item())
+
+
+def torch_per_token_cast(
+    values: torch.Tensor,
+    num_per_channels: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Ground-truth per-token FP8 cast in fp32 PyTorch.
+
+    ``values`` holds the activations to quantize (for SwiGLU, the post-SwiGLU
+    values).  Checking both kernels against this catches the case where they
+    agree with each other but are both wrong.
+    """
+    num_tokens, hidden = values.shape
+    grouped = values.float().reshape(num_tokens, hidden // num_per_channels, -1)
+    amax = grouped.abs().amax(dim=-1, keepdim=True).clamp_min(SF_CLAMP_MIN)
+    scale = (amax / E4M3_MAX).squeeze(-1)
+    scaled = (grouped * (E4M3_MAX / amax)).clamp(-E4M3_MAX, E4M3_MAX)
+    out = scaled.reshape(num_tokens, hidden).to(torch.float8_e4m3fn)
+    return out, scale
+
+
+def dequantize(
+    out: torch.Tensor,
+    scales: torch.Tensor,
+    num_per_channels: int,
+) -> torch.Tensor:
+    """Decode an e4m3 tensor and its per-group scales back to fp32."""
+    grouped = out.float().reshape(
+        out.shape[0],
+        out.shape[1] // num_per_channels,
+        num_per_channels,
+    )
+    return (grouped * scales[:, :, None]).reshape(out.shape)
+
+
+def quantization_snr_db(reference: torch.Tensor, dequantized: torch.Tensor) -> float:
+    """Signal-to-noise ratio of a dequantized tensor against its fp32 source.
+
+    Reported rather than asserted on: it is the sanity check that the cast is
+    carrying real information.  Round-to-nearest into a three-bit mantissa
+    lands near 32 dB for Gaussian activations, and is stable across shapes, so
+    a value well below that means the scale factors are wrong even if the code
+    ladder checks pass.
+    """
+    reference = reference.float()
+    noise = dequantized.float() - reference
+    return float(
+        10.0
+        * torch.log10(reference.square().sum() / noise.square().sum().clamp_min(1e-30))
+    )
diff --git a/examples/quantization/per_token_cast.py b/examples/quantization/per_token_cast.py
index 49b97960..01ef260f 100644
--- a/examples/quantization/per_token_cast.py
+++ b/examples/quantization/per_token_cast.py
@@ -3,22 +3,43 @@
 """Per-token FP8 cast with scale factors.
 
 This is a Tilus translation of DeepSeek TileKernels'
-``per_token_cast_kernel.py`` for the common FP16 -> FP8 e4m3 path.  Each CTA
-processes one token and one channel group, computes the absolute maximum within
-that group, stores a float32 scale factor, and writes the scaled FP8 output.
+``per_token_cast_kernel.py``.  Each CTA processes one token and one channel
+group, computes the absolute maximum within that group, stores a float32 scale
+factor, and writes the scaled FP8 e4m3 output.
+
+The input is bfloat16.  That is not a free choice: TileKernels'
+``get_cast_input_and_config`` asserts the unquantized input is bfloat16 or
+float32, so bf16 is the only 16-bit type both implementations accept.  Since
+this kernel is purely DRAM-bound on its input read, handing the reference fp32
+while Tilus reads 16-bit halves the reference's achievable bandwidth-limited
+runtime and makes the comparison meaningless.  Both sides here read the same
+bf16 tensor.
+
+Scale-factor arithmetic follows TileKernels' ``get_sf_and_inv`` exactly: the
+group absmax is clamped from below to 1e-4, the scale is ``absmax / 448`` and
+its reciprocal is ``448 / absmax``.
 """
 
 import pandas
 import tilus
 import torch
+from fp8_check import (
+    E4M3_MAX,
+    SF_CLAMP_MIN,
+    check_fp8_close,
+    check_scales_close,
+    dequantize,
+    quantization_snr_db,
+    torch_per_token_cast,
+)
 from tile_kernels.quant.per_token_cast_kernel import per_token_cast
-from tilus import float8_e4m3, float16, float32, int32
+from tilus import bfloat16, float8_e4m3, float32, int32
 from tilus.utils import benchmark_func, cdiv
 
 
-@tilus.autotune("block_m", [1, 2, 4, 8])
-@tilus.autotune("groups_per_block", [1, 2, 4, 8])
-@tilus.autotune("warps", [4, 8])
+@tilus.autotune("block_m", [1, 2, 4])
+@tilus.autotune("groups_per_block", [1, 4, 16, 64])
+@tilus.autotune("warps", [4, 8, 16])
 class PerTokenCast(tilus.Script):
     def __init__(
         self,
@@ -38,7 +59,7 @@ def __call__(
         self,
         num_tokens: int,
         hidden: int32,
-        x_ptr: ~float16,
+        x_ptr: ~bfloat16,
         out_ptr: ~float8_e4m3,
         out_sf_ptr: ~float32,
     ):
@@ -55,7 +76,7 @@ def __call__(
 
         g_x = self.global_view(
             x_ptr,
-            dtype=float16,
+            dtype=bfloat16,
             shape=[num_tokens, hidden],
         )
         g_out = self.global_view(
@@ -69,49 +90,49 @@ def __call__(
             shape=[num_tokens, cdiv(hidden, self.num_per_channels)],
         )
 
-        for gi in range(self.groups_per_block):
-            offset_n = base_offset_n + gi * self.block_n
-            sf_col = offset_n // self.num_per_channels
-
-            r_x = self.load_global(
-                g_x,
-                offsets=[offset_m, offset_n],
-                shape=[self.block_m, self.block_n],
-            ).to(float32)
-
-            r_absmax = self.max(self.abs(r_x), dim=1, keepdim=True)
-            r_fp8_max = self.register_tensor(
-                dtype=float32,
-                shape=[self.block_m, 1],
-                init=448.0,
-            )
-            r_scale = self.where(r_absmax > 0.0, x=r_absmax / 448.0, y=1.0)
-            r_inv_scale = self.where(r_absmax > 0.0, x=r_fp8_max / r_absmax, y=1.0)
-
-            self.store_global(g_out_sf, r_scale, offsets=[offset_m, sf_col])
-            self.store_global(
-                g_out,
-                (r_x * r_inv_scale).to(float8_e4m3),
-                offsets=[offset_m, offset_n],
-            )
-
-
-def tilekernels_per_token_cast_reference(
-    x: torch.Tensor,
-    num_per_channels: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    return per_token_cast(x, "e4m3", num_per_channels)
+        # One wide load of the whole tile rather than a loop of per-group
+        # loads.  The loop form does not get unrolled, so each iteration
+        # exposed the full global-load latency with nothing to overlap it;
+        # issuing every load up front and reshaping for the reduction is worth
+        # ~5% at the DRAM-bound shape.
+        r_x = self.load_global(
+            g_x,
+            offsets=[offset_m, base_offset_n],
+            shape=[self.block_m, n_step],
+        ).to(float32)
+
+        # Reshape into [block_m, groups_per_block, num_per_channels] so the
+        # per-group absmax is a single reduce on dim=2.
+        r_x_grouped = self.reshape(
+            r_x,
+            shape=[self.block_m, self.groups_per_block, self.num_per_channels],
+        )
 
+        # Clamp the absmax from below exactly as TileKernels does, so an
+        # all-zero group produces a tiny scale instead of a division by zero,
+        # and so both sides agree bit-for-bit on the clamped value.
+        r_absmax = self.max(self.abs(r_x_grouped), dim=2, keepdim=True)
+        r_amax = self.where(r_absmax > SF_CLAMP_MIN, x=r_absmax, y=SF_CLAMP_MIN)
+        r_fp8_max = self.register_tensor(
+            dtype=float32,
+            shape=[self.block_m, self.groups_per_block, 1],
+            init=E4M3_MAX,
+        )
+        r_scale = r_amax / E4M3_MAX
+        r_inv_scale = r_fp8_max / r_amax
+
+        # Store one fp32 scale per group.
+        r_scale_2d = self.reshape(r_scale, shape=[self.block_m, self.groups_per_block])
+        self.store_global(
+            g_out_sf,
+            r_scale_2d,
+            offsets=[offset_m, base_offset_n // self.num_per_channels],
+        )
 
-def dequantized_sum(
-    out: torch.Tensor, scales: torch.Tensor, num_per_channels: int
-) -> torch.Tensor:
-    grouped = out.float().reshape(
-        out.shape[0],
-        out.shape[1] // num_per_channels,
-        num_per_channels,
-    )
-    return (grouped * scales[:, :, None]).sum()
+        # Apply scaling, flatten back, cast to fp8, bulk store.
+        r_out_grouped = (r_x_grouped * r_inv_scale).to(float8_e4m3)
+        r_out = self.reshape(r_out_grouped, shape=[self.block_m, n_step])
+        self.store_global(g_out, r_out, offsets=[offset_m, base_offset_n])
 
 
 def main():
@@ -122,56 +143,84 @@ def main():
         "tilekernels (ms)",
         "tilus (ms)",
         "speedup",
-        "sum diff",
+        "code mismatch",
+        "sf rel err",
+        "snr (dB)",
     ]
 
+    # The first three shapes move only a few hundred KB, so at ~6 us they are
+    # dominated by launch and dispatch overhead rather than by the kernel.  The
+    # last one reads 128 MB and is genuinely DRAM-bound, which is the regime
+    # this kernel is written for and the only one where the speedup column
+    # says anything about code quality.
     for num_tokens, hidden in [
         (128, 1024),
         (256, 2048),
         (257, 4096),
+        (8192, 8192),
     ]:
         num_per_channels = 128
         kernel = PerTokenCast(num_per_channels=num_per_channels)
 
+        # One bf16 tensor, read by both implementations.
         x = (
             torch.randn(
                 num_tokens,
                 hidden,
                 device="cuda",
-                dtype=torch.float16,
+                dtype=torch.bfloat16,
             )
             * 2.0
         ).contiguous()
-        out = torch.empty((num_tokens, hidden), device="cuda", dtype=torch.float8_e4m3fn)
-        out_sf = torch.empty(
-            (num_tokens, hidden // num_per_channels),
-            device="cuda",
-            dtype=torch.float32,
-        )
-        x_tilekernels = x.float()
 
-        kernel(num_tokens, hidden, x, out, out_sf)
-        expected_out, expected_sf = tilekernels_per_token_cast_reference(
-            x_tilekernels,
-            num_per_channels,
-        )
+        def run_tilus():
+            # Allocate here rather than reusing preallocated buffers, because
+            # the TileKernels entry point allocates its returns on every call.
+            out = torch.empty(
+                (num_tokens, hidden), device="cuda", dtype=torch.float8_e4m3fn
+            )
+            out_sf = torch.empty(
+                (num_tokens, hidden // num_per_channels),
+                device="cuda",
+                dtype=torch.float32,
+            )
+            kernel(num_tokens, hidden, x, out, out_sf)
+            return out, out_sf
 
-        max_code_diff = (out.float() - expected_out.float()).abs().max().item()
-        assert max_code_diff <= 32.0, f"max decoded FP8 code diff is {max_code_diff}"
-        torch.testing.assert_close(out_sf, expected_sf, atol=1e-5, rtol=1e-5)
+        def run_tilekernels():
+            return per_token_cast(x, "e4m3", num_per_channels)
 
-        actual_sum = dequantized_sum(out, out_sf, num_per_channels)
-        expected_sum = dequantized_sum(expected_out, expected_sf, num_per_channels)
-        torch.testing.assert_close(actual_sum, expected_sum, atol=2.0, rtol=2e-2)
-        sum_diff = (actual_sum - expected_sum).abs().item()
+        out, out_sf = run_tilus()
+        expected_out, expected_sf = run_tilekernels()
+        torch_out, torch_sf = torch_per_token_cast(x, num_per_channels)
 
-        tilekernels_ms = benchmark_func(
-            lambda: tilekernels_per_token_cast_reference(
-                x_tilekernels,
-                num_per_channels,
-            )
+        # Tilus against TileKernels: same inputs, same arithmetic, so they may
+        # differ only by the approximate fp32 division Tilus compiles with.
+        sf_rel_err = check_scales_close(
+            out_sf, expected_sf, label=f"({num_tokens}, {hidden}) sf vs tilekernels"
+        )
+        stats = check_fp8_close(
+            out,
+            expected_out,
+            label=f"({num_tokens}, {hidden}) out vs tilekernels",
         )
-        tilus_ms = benchmark_func(lambda: kernel(num_tokens, hidden, x, out, out_sf))
+
+        # Both against an independent fp32 PyTorch cast, so that agreeing with
+        # each other is not mistaken for being correct.
+        check_scales_close(
+            out_sf, torch_sf, label=f"({num_tokens}, {hidden}) sf vs torch"
+        )
+        check_fp8_close(out, torch_out, label=f"({num_tokens}, {hidden}) out vs torch")
+        check_fp8_close(
+            expected_out,
+            torch_out,
+            label=f"({num_tokens}, {hidden}) tilekernels out vs torch",
+        )
+
+        snr_db = quantization_snr_db(x, dequantize(out, out_sf, num_per_channels))
+
+        tilekernels_ms = benchmark_func(run_tilekernels)
+        tilus_ms = benchmark_func(run_tilus)
         rows.append(
             [
                 num_tokens,
@@ -179,16 +228,26 @@ def main():
                 tilekernels_ms,
                 tilus_ms,
                 f"{tilekernels_ms / tilus_ms:.2f}x",
-                sum_diff,
+                f"{stats.mismatch_frac:.4%}",
+                f"{sf_rel_err:.2e}",
+                f"{snr_db:.1f}",
             ]
         )
         print(
             "Per-token FP8 cast matches reference for size "
-            f"({num_tokens}, {hidden}); max code diff={max_code_diff:.6g}; "
-            f"dequantized sum diff={sum_diff:.6g}"
+            f"({num_tokens}, {hidden}): every code within 1 of TileKernels and "
+            f"of fp32 torch, {stats.mismatch_frac:.4%} of codes differ at all, "
+            f"scale factors agree to {sf_rel_err:.2e} relative, "
+            f"quantization SNR {snr_db:.1f} dB"
         )
 
-    print(pandas.DataFrame(rows, columns=headers))
+    print(pandas.DataFrame(rows, columns=headers).to_string(index=False))
+    print(
+        "\nBoth implementations read the same bf16 input and allocate their own "
+        "outputs.\nThe Tilus kernel is autotuned per shape; the TileKernels "
+        "kernel picks its tiling\nanalytically from `hidden`, so it is not tuned "
+        "against this measurement."
+    )
 
 
 if __name__ == "__main__":
diff --git a/examples/quantization/swiglu_forward_and_per_token_cast.py b/examples/quantization/swiglu_forward_and_per_token_cast.py
index 153cc343..50fef78c 100644
--- a/examples/quantization/swiglu_forward_and_per_token_cast.py
+++ b/examples/quantization/swiglu_forward_and_per_token_cast.py
@@ -5,26 +5,46 @@
 This is a Tilus translation of DeepSeek TileKernels'
 ``swiglu_forward_and_per_token_cast_kernel.py``.  It computes
 
-    out = silu(x[:, :hidden]) * x[:, hidden:]
+    out = silu(clamp(x[:, :hidden])) * clamp(x[:, hidden:])
 
 optionally applies a routing weight and expert mask, then quantizes each
 ``num_per_channels`` group to FP8 e4m3 with one float32 scale factor per
 token/group.
+
+The input is bfloat16 on both sides.  The kernel is DRAM-bound on its two
+input reads, so giving the reference fp32 while Tilus reads 16-bit would double
+the reference's traffic and produce a speedup number that measures the dtype
+rather than the kernel.  bf16 is also what the companion ``per_token_cast``
+example must use, since TileKernels' cast entry point accepts only bf16 or
+fp32 for unquantized input.
+
+Scale-factor arithmetic follows TileKernels' ``get_sf_and_inv`` exactly: the
+group absmax is clamped from below to 1e-4, the scale is ``absmax / 448`` and
+its reciprocal is ``448 / absmax``.
 """
 
 import pandas
 import tilus
 import torch
+from fp8_check import (
+    E4M3_MAX,
+    SF_CLAMP_MIN,
+    check_fp8_close,
+    check_scales_close,
+    dequantize,
+    quantization_snr_db,
+    torch_per_token_cast,
+)
 from tile_kernels.quant.swiglu_forward_and_per_token_cast_kernel import (
     swiglu_forward_and_per_token_cast,
 )
-from tilus import float8_e4m3, float16, float32, int32
+from tilus import bfloat16, float8_e4m3, float32, int32
 from tilus.utils import benchmark_func, cdiv
 
 
 @tilus.autotune("block_m", [1])
-@tilus.autotune("groups_per_block", [1, 2, 4, 8, 16])
-@tilus.autotune("warps", [1, 2, 4, 8])
+@tilus.autotune("groups_per_block", [1, 4, 16, 64])
+@tilus.autotune("warps", [2, 4, 8, 16])
 class SwiGLUForwardAndPerTokenCast(tilus.Script):
     def __init__(
         self,
@@ -51,7 +71,7 @@ def __call__(
         num_expanded_tokens: int,
         hidden: int32,
         num_topk_values: int32,
-        x_ptr: ~float16,
+        x_ptr: ~bfloat16,
         out_ptr: ~float8_e4m3,
         out_sf_ptr: ~float32,
         pos_to_token_topk_ptr: ~int32,
@@ -72,7 +92,7 @@ def __call__(
 
         g_x = self.global_view(
             x_ptr,
-            dtype=float16,
+            dtype=bfloat16,
             shape=[num_expanded_tokens, hidden * 2],
         )
         g_out = self.global_view(
@@ -144,13 +164,15 @@ def __call__(
             r_absmax = self.max(
                 self.abs(r_value_grouped), dim=2, keepdim=True
             )  # [block_m, groups_per_block, 1]
+            # Clamp the absmax from below exactly as TileKernels does.
+            r_amax = self.where(r_absmax > SF_CLAMP_MIN, x=r_absmax, y=SF_CLAMP_MIN)
             r_fp8_max = self.register_tensor(
                 dtype=float32,
                 shape=[self.block_m, self.groups_per_block, 1],
-                init=448.0,
+                init=E4M3_MAX,
             )
-            r_scale = self.where(r_absmax > 0.0, x=r_absmax / 448.0, y=1.0)
-            r_inv_scale = self.where(r_absmax > 0.0, x=r_fp8_max / r_absmax, y=1.0)
+            r_scale = r_amax / E4M3_MAX
+            r_inv_scale = r_fp8_max / r_amax
 
             # Store one fp32 scale per group.
             r_scale_2d = self.reshape(
@@ -164,34 +186,23 @@ def __call__(
             self.store_global(g_out, r_out, offsets=[offset_m, base_offset_n])
 
 
-def tilekernels_swiglu_reference(
+def torch_swiglu(
     x: torch.Tensor,
     pos_to_token_topk: torch.Tensor,
     topk_weights: torch.Tensor,
-    pos_to_expert: torch.Tensor,
     clamp_value: float,
-    num_per_channels: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    return swiglu_forward_and_per_token_cast(
-        x,
-        "e4m3",
-        num_per_channels,
-        pos_to_token_topk=pos_to_token_topk,
-        topk_weights=topk_weights,
-        pos_to_expert=pos_to_expert,
-        swiglu_clamp_value=clamp_value,
-    )
-
-
-def dequantized_sum(
-    out: torch.Tensor, scales: torch.Tensor, num_per_channels: int
 ) -> torch.Tensor:
-    grouped = out.float().reshape(
-        out.shape[0],
-        out.shape[1] // num_per_channels,
-        num_per_channels,
-    )
-    return (grouped * scales[:, :, None]).sum()
+    """Reference SwiGLU activation in fp32, before quantization.
+
+    Mirrors TileKernels' operation order: clamp, silu, multiply by the gate,
+    then apply the routing weight.
+    """
+    hidden = x.shape[1] // 2
+    left = x[:, :hidden].float().clamp(max=clamp_value)
+    right = x[:, hidden:].float().clamp(min=-clamp_value, max=clamp_value)
+    value = left / (1.0 + torch.exp(-left)) * right
+    weight = topk_weights.reshape(-1)[pos_to_token_topk.long()]
+    return value * weight[:, None]
 
 
 def main():
@@ -202,24 +213,33 @@ def main():
         "tilekernels (ms)",
         "tilus (ms)",
         "speedup",
-        "sum diff",
+        "code mismatch",
+        "sf rel err",
+        "snr (dB)",
     ]
 
+    # The first four shapes are dominated by launch and dispatch overhead: they
+    # move at most a few MB and run in under 25 us.  The last one reads 128 MB
+    # and is genuinely DRAM-bound, which is the regime this kernel is written
+    # for and the only one where the speedup column says anything about code
+    # quality.
     for num_expanded_tokens, hidden, num_tokens, num_topk in [
         (128, 1024, 64, 2),
         (256, 2048, 128, 2),
         (257, 4096, 128, 2),
         (1024, 4096, 512, 2),
+        (4096, 8192, 2048, 2),
     ]:
         num_per_channels = 128
         kernel = SwiGLUForwardAndPerTokenCast(num_per_channels=num_per_channels)
 
+        # One bf16 tensor, read by both implementations.
         x = (
             torch.randn(
                 num_expanded_tokens,
                 hidden * 2,
                 device="cuda",
-                dtype=torch.float16,
+                dtype=torch.bfloat16,
             )
             * 2.0
         ).contiguous()
@@ -236,73 +256,22 @@ def main():
         )
         pos_to_expert = torch.ones(num_expanded_tokens, device="cuda", dtype=torch.int32)
         pos_to_expert[::17] = -1
-
-        out = torch.empty(
-            (num_expanded_tokens, hidden),
-            device="cuda",
-            dtype=torch.float8_e4m3fn,
-        )
-        out_sf = torch.empty(
-            (num_expanded_tokens, hidden // num_per_channels),
-            device="cuda",
-            dtype=torch.float32,
-        )
-        x_tilekernels = x.float()
-
         clamp_value = 6.0
-        kernel(
-            num_expanded_tokens,
-            hidden,
-            num_tokens * num_topk,
-            x,
-            out,
-            out_sf,
-            pos_to_token_topk,
-            topk_weights,
-            pos_to_expert,
-            clamp_value,
-        )
 
-        expected_out, expected_sf = tilekernels_swiglu_reference(
-            x_tilekernels,
-            pos_to_token_topk,
-            topk_weights,
-            pos_to_expert,
-            clamp_value,
-            num_per_channels,
-        )
-        valid = pos_to_expert >= 0
-        max_code_diff = (
-            (out[valid].float() - expected_out[valid].float()).abs().max().item()
-        )
-        assert max_code_diff <= 32.0, f"max decoded FP8 code diff is {max_code_diff}"
-        torch.testing.assert_close(
-            out_sf[valid],
-            expected_sf[valid],
-            atol=1e-5,
-            rtol=1e-5,
-        )
-        actual_sum = dequantized_sum(out[valid], out_sf[valid], num_per_channels)
-        expected_sum = dequantized_sum(
-            expected_out[valid],
-            expected_sf[valid],
-            num_per_channels,
-        )
-        torch.testing.assert_close(actual_sum, expected_sum, atol=2.0, rtol=2e-2)
-        sum_diff = (actual_sum - expected_sum).abs().item()
-
-        tilekernels_ms = benchmark_func(
-            lambda: tilekernels_swiglu_reference(
-                x_tilekernels,
-                pos_to_token_topk,
-                topk_weights,
-                pos_to_expert,
-                clamp_value,
-                num_per_channels,
+        def run_tilus():
+            # Allocate here rather than reusing preallocated buffers, because
+            # the TileKernels entry point allocates its returns on every call.
+            out = torch.empty(
+                (num_expanded_tokens, hidden),
+                device="cuda",
+                dtype=torch.float8_e4m3fn,
             )
-        )
-        tilus_ms = benchmark_func(
-            lambda: kernel(
+            out_sf = torch.empty(
+                (num_expanded_tokens, hidden // num_per_channels),
+                device="cuda",
+                dtype=torch.float32,
+            )
+            kernel(
                 num_expanded_tokens,
                 hidden,
                 num_tokens * num_topk,
@@ -314,7 +283,61 @@ def main():
                 pos_to_expert,
                 clamp_value,
             )
+            return out, out_sf
+
+        def run_tilekernels():
+            return swiglu_forward_and_per_token_cast(
+                x,
+                "e4m3",
+                num_per_channels,
+                pos_to_token_topk=pos_to_token_topk,
+                topk_weights=topk_weights,
+                pos_to_expert=pos_to_expert,
+                swiglu_clamp_value=clamp_value,
+            )
+
+        out, out_sf = run_tilus()
+        expected_out, expected_sf = run_tilekernels()
+
+        # Masked-out rows are never written by either kernel, so they hold
+        # whatever the allocator handed back; compare only the live rows.
+        valid = pos_to_expert >= 0
+        activations = torch_swiglu(
+            x[valid], pos_to_token_topk[valid], topk_weights, clamp_value
+        )
+        torch_out, torch_sf = torch_per_token_cast(activations, num_per_channels)
+
+        # Tilus against TileKernels: same inputs, same arithmetic, so they may
+        # differ only by the approximate fp32 division Tilus compiles with.
+        label = f"({num_expanded_tokens}, {hidden})"
+        sf_rel_err = check_scales_close(
+            out_sf[valid], expected_sf[valid], label=f"{label} sf vs tilekernels"
+        )
+        stats = check_fp8_close(
+            out[valid], expected_out[valid], label=f"{label} out vs tilekernels"
+        )
+
+        # Both against an independent fp32 PyTorch SwiGLU and cast, so that
+        # agreeing with each other is not mistaken for being correct.  The
+        # tolerance on the scale factors is looser than in `per_token_cast`
+        # because the absmax is taken over a transcendental (`exp`), where the
+        # device and PyTorch implementations may differ by an ulp.
+        check_scales_close(
+            out_sf[valid], torch_sf, label=f"{label} sf vs torch", rtol=1e-4
         )
+        check_fp8_close(out[valid], torch_out, label=f"{label} out vs torch")
+        check_fp8_close(
+            expected_out[valid],
+            torch_out,
+            label=f"{label} tilekernels out vs torch",
+        )
+
+        snr_db = quantization_snr_db(
+            activations, dequantize(out[valid], out_sf[valid], num_per_channels)
+        )
+
+        tilekernels_ms = benchmark_func(run_tilekernels)
+        tilus_ms = benchmark_func(run_tilus)
         rows.append(
             [
                 num_expanded_tokens,
@@ -322,16 +345,26 @@ def main():
                 tilekernels_ms,
                 tilus_ms,
                 f"{tilekernels_ms / tilus_ms:.2f}x",
-                sum_diff,
+                f"{stats.mismatch_frac:.4%}",
+                f"{sf_rel_err:.2e}",
+                f"{snr_db:.1f}",
             ]
         )
         print(
-            "SwiGLU FP8 cast matches reference for size "
-            f"({num_expanded_tokens}, {hidden}); max code diff={max_code_diff:.6g}; "
-            f"dequantized sum diff={sum_diff:.6g}"
+            f"SwiGLU FP8 cast matches reference for size {label}: every code "
+            f"within 1 of TileKernels and of fp32 torch, "
+            f"{stats.mismatch_frac:.4%} of codes differ at all, scale factors "
+            f"agree to {sf_rel_err:.2e} relative, "
+            f"quantization SNR {snr_db:.1f} dB"
         )
 
-    print(pandas.DataFrame(rows, columns=headers))
+    print(pandas.DataFrame(rows, columns=headers).to_string(index=False))
+    print(
+        "\nBoth implementations read the same bf16 input and allocate their own "
+        "outputs.\nThe Tilus kernel is autotuned per shape; the TileKernels "
+        "kernel picks its tiling\nanalytically from `hidden`, so it is not tuned "
+        "against this measurement."
+    )
 
 
 if __name__ == "__main__":
diff --git a/tests/examples/test_examples.py b/tests/examples/test_examples.py
index 1a027f38..d99898f9 100644
--- a/tests/examples/test_examples.py
+++ b/tests/examples/test_examples.py
@@ -84,6 +84,8 @@
     # Benchmark utilities
     ("blackwell_matmul", "benchmark.py"),
     ("hopper_matmul", "benchmark.py"),
+    # Shared correctness helpers for the quantization examples
+    ("quantization", "fp8_check.py"),
 ]
 
 

From 99d834075d3e74c86cc9b45eca006de0e40a1b1a Mon Sep 17 00:00:00 2001
From: William Zhang <wzhang20@yahoo.com>
Date: Wed, 9 Sep 2026 18:15:04 -0400
Subject: [PATCH 2/5] add mhc, moe, and per channel tilekernels

Signed-off-by: William Zhang <wzhang20@yahoo.com>
---
 examples/mhc/normw_merge.py                   | 63 +++++++++++++
 examples/moe/topk_gate.py                     | 89 +++++++++++++++++++
 examples/quantization/per_channel_cast.py     | 76 ++++++++++++++++
 python/tilus/backends/codegen.py              |  5 ++
 python/tilus/backends/emitters/__init__.py    |  1 +
 .../tilus/backends/emitters/fp8_epilogue.py   | 53 +++++++++++
 python/tilus/backends/emitters/reduce.py      | 13 +--
 .../tilus/hidet/ir/primitives/cuda/float8.py  | 40 +++++++++
 .../hidet/ir/primitives/cuda/math/bfloat16.py |  4 +
 python/tilus/ir/builders/stmt_builder.py      | 11 +++
 python/tilus/ir/instructions/__init__.py      |  1 +
 python/tilus/ir/instructions/generic.py       | 17 ++++
 .../inference/inference_rules/empty_rule.py   |  2 +
 .../inference/validation_rules/always_ok.py   |  2 +
 python/tilus/lang/instructions/root.py        | 12 ++-
 15 files changed, 383 insertions(+), 6 deletions(-)
 create mode 100644 examples/mhc/normw_merge.py
 create mode 100644 examples/moe/topk_gate.py
 create mode 100644 examples/quantization/per_channel_cast.py
 create mode 100644 python/tilus/backends/emitters/fp8_epilogue.py
 create mode 100644 python/tilus/hidet/ir/primitives/cuda/float8.py

diff --git a/examples/mhc/normw_merge.py b/examples/mhc/normw_merge.py
new file mode 100644
index 00000000..cf4fe397
--- /dev/null
+++ b/examples/mhc/normw_merge.py
@@ -0,0 +1,63 @@
+# 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()
diff --git a/examples/moe/topk_gate.py b/examples/moe/topk_gate.py
new file mode 100644
index 00000000..997003a5
--- /dev/null
+++ b/examples/moe/topk_gate.py
@@ -0,0 +1,89 @@
+# 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
+        )
+        negative_max = -3.402823466e38
+        # A vector load past the logical expert dimension is zero-filled by
+        # Tilus.  Only materialize a validity mask when there actually is a
+        # tail: for the standard 256-expert route it is compile-time dead work.
+        if self.aligned_experts != self.num_experts:
+            valid = self.register_tensor(
+                dtype=int32,
+                shape=[1, self.aligned_experts],
+                init=lambda _, j: j < self.num_experts,
+            )
+            values = self.where(valid != 0, x=values, y=negative_max)
+
+        # ``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(values == best_value, x=expert_ids, y=int32.max_value)
+            best_index = self.min(candidates, 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])
+            values = self.where(expert_ids == best_index, x=negative_max, 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)
+        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()
diff --git a/examples/quantization/per_channel_cast.py b/examples/quantization/per_channel_cast.py
new file mode 100644
index 00000000..6b911913
--- /dev/null
+++ b/examples/quantization/per_channel_cast.py
@@ -0,0 +1,76 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""TileKernels-compatible FP8 e4m3 cast with one scale per 128-token column."""
+
+import tilus
+import torch
+from fp8_check import check_fp8_close
+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
+
+
+class PerChannelCast(tilus.Script):
+    def __init__(self, block_n: int = 128, warps: int = 8):
+        super().__init__()
+        self.block_m, self.block_n, self.warps = 128, block_n, warps
+
+    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])
+        self.store_scaled_fp8e4m3_from_shared(out, shared_x, inv_scale, 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()
diff --git a/python/tilus/backends/codegen.py b/python/tilus/backends/codegen.py
index 4d3be984..6bcb2f0e 100644
--- a/python/tilus/backends/codegen.py
+++ b/python/tilus/backends/codegen.py
@@ -198,6 +198,11 @@ def visit_Function(self, func: Function) -> IRModule:
         current_target = get_current_target()
         if current_target.supports(nvgpu_sm90):
             cluster_blocks = self._function.metadata.cluster_blocks
+            # A 1x1x1 cluster has no cluster semantics.  Do not attach the
+            # Hopper cluster launch attribute for that default: it needlessly
+            # changes ordinary-kernel scheduling/resource configuration.
+            if cluster_blocks == (1, 1, 1):
+                cluster_blocks = None
         else:
             if self._function.metadata.cluster_blocks != (1, 1, 1):
                 raise RuntimeError(
diff --git a/python/tilus/backends/emitters/__init__.py b/python/tilus/backends/emitters/__init__.py
index 4a1863af..e2515012 100644
--- a/python/tilus/backends/emitters/__init__.py
+++ b/python/tilus/backends/emitters/__init__.py
@@ -22,6 +22,7 @@
     cuda,
     debug,
     elementwise,
+    fp8_epilogue,
     gmem,
     ldst,
     random,
diff --git a/python/tilus/backends/emitters/fp8_epilogue.py b/python/tilus/backends/emitters/fp8_epilogue.py
new file mode 100644
index 00000000..db5ecf23
--- /dev/null
+++ b/python/tilus/backends/emitters/fp8_epilogue.py
@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Fused FP8 epilogues for bandwidth-bound quantization kernels."""
+
+from tilus.backends.emitter import BaseInstEmitter, register_emitter
+from tilus.hidet.ir.dtypes import bfloat16, float8_e4m3, float32, uint32
+from tilus.hidet.ir.expr import cast
+from tilus.hidet.ir.primitives.cuda.float8 import scale_bf16x4_to_fp8e4m3x4
+from tilus.hidet.ir.type import void_p
+from tilus.ir.instructions import StoreScaledFp8E4M3FromSharedInst
+from tilus.ir.tensor import GlobalTensor, RegisterTensor, SharedTensor
+
+
+@register_emitter(StoreScaledFp8E4M3FromSharedInst)
+class StoreScaledFp8E4M3FromSharedEmitter(BaseInstEmitter):
+    """Lower the 128x128 per-channel FP8 epilogue without register tiles."""
+
+    def emit(self, inst: StoreScaledFp8E4M3FromSharedInst) -> None:
+        dst: GlobalTensor = inst.inputs[0].as_global_tensor()
+        src: SharedTensor = inst.inputs[1].as_shared_tensor()
+        inv_scale: RegisterTensor = inst.inputs[2].as_register_tensor()
+        if (
+            dst.dtype != float8_e4m3
+            or src.dtype != bfloat16
+            or inv_scale.dtype != float32
+            or tuple(src.shape) != (128, 128)
+            or tuple(inv_scale.shape) != (1, 128)
+            or inv_scale.layout.local_size != 4
+        ):
+            raise ValueError("fused FP8 epilogue requires shared bf16[128,128] and register float32[1,128]")
+
+        dst_buf = self.tensor2var[dst]
+        src_buf = self.tensor2var[src]
+        scale_buf = self.tensor2var[inv_scale]
+        offset_m, offset_n = inst.offsets
+        lane_id = self.lane_id()
+        warp_id = self.warp_id()
+        col = lane_id * 4
+        with self.for_range(16, attr="u+") as i:
+            row = warp_id * 16 + i
+            src_offset = src.layout(row, col)
+            dst_offset = dst.layout(offset_m + row, offset_n + col)
+            self.append(
+                scale_bf16x4_to_fp8e4m3x4(
+                    cast(~dst_buf[dst_offset], void_p),
+                    cast(~src_buf[src_offset], void_p),
+                    cast(~src_buf[src_offset + 2], void_p),
+                    scale_buf[0],
+                    scale_buf[1],
+                    scale_buf[2],
+                    scale_buf[3],
+                )
+            )
diff --git a/python/tilus/backends/emitters/reduce.py b/python/tilus/backends/emitters/reduce.py
index 836c51eb..9bc06674 100644
--- a/python/tilus/backends/emitters/reduce.py
+++ b/python/tilus/backends/emitters/reduce.py
@@ -21,7 +21,7 @@
 from tilus.hidet.ir import DataType
 from tilus.hidet.ir.dtypes import int32, uint32
 from tilus.hidet.ir.expr import Expr, Var, bitwise_and, bitwise_or, cast, if_then_else, logical_and
-from tilus.hidet.ir.primitives.cuda.shfl import shfl_down_sync, shfl_up_sync
+from tilus.hidet.ir.primitives.cuda.shfl import shfl_up_sync, shfl_xor_sync
 from tilus.hidet.ir.type import tensor_pointer_type
 from tilus.hidet.ir.utils.index_transform import index_deserialize, index_serialize
 from tilus.hidet.utils.py import is_power_of_two
@@ -138,10 +138,15 @@ def intra_warp_reduce(self, inst: ReduceInst) -> None:
                     indices=[dst_local],
                     value=self.scalar_reduce(
                         lhs=dst_buf[dst_local],
-                        rhs=shfl_down_sync(
+                        # XOR butterfly reduction produces the final value in
+                        # every participating lane.  The old down-tree needed
+                        # a second, reverse shuffle tree to broadcast its
+                        # lane-zero result, doubling shuffle traffic for every
+                        # warp-local reduction.
+                        rhs=shfl_xor_sync(
                             mask=uint32(0xFFFFFFFF),
                             var=dst_buf[dst_local],
-                            delta=1 << lane_bit,
+                            lane_mask=1 << lane_bit,
                             width=1 << (lane_bit + 1),
                         ),
                         op=inst.op,
@@ -350,8 +355,6 @@ def efficient_reduce(self, inst: ReduceInst) -> None:
         if self.requires_inter_warp_reduction(inst):
             # reduce between warps
             self.inter_warp_reduce(inst)
-        else:
-            self.intra_warp_broadcast(inst)
 
     def emit(self, inst: ReduceInst) -> None:
         self.efficient_reduce(inst)
diff --git a/python/tilus/hidet/ir/primitives/cuda/float8.py b/python/tilus/hidet/ir/primitives/cuda/float8.py
new file mode 100644
index 00000000..498005b3
--- /dev/null
+++ b/python/tilus/hidet/ir/primitives/cuda/float8.py
@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Small CUDA FP8 conversion primitives used by vectorized epilogues."""
+
+from typing import no_type_check
+
+from tilus.hidet.ir.expr import Expr
+from tilus.hidet.ir.func import Function
+from tilus.hidet.ir.primitives.func import call_primitive_func, register_primitive_function
+from tilus.hidet.ir.stmt import BlackBoxStmt
+from tilus.hidet.utils import initialize
+
+
+@initialize()
+def register_functions():
+    from tilus.hidet.lang import attrs, script  # pylint: disable=import-outside-toplevel
+    from tilus.hidet.lang.types import float32, void_p
+
+    bf16_template = r"""
+    float2 f0 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>({}));
+    float2 f1 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>({}));
+    __nv_fp8x2_storage_t p0 = __nv_cvt_float2_to_fp8x2(make_float2(f0.x * {}, f0.y * {}), __NV_SATFINITE, __NV_E4M3);
+    __nv_fp8x2_storage_t p1 = __nv_cvt_float2_to_fp8x2(make_float2(f1.x * {}, f1.y * {}), __NV_SATFINITE, __NV_E4M3);
+    *reinterpret_cast<uint32_t*>({}) = static_cast<uint32_t>(p0) | (static_cast<uint32_t>(p1) << 16);
+    """
+
+    @no_type_check
+    @script
+    def scale_bf16x4_to_fp8e4m3x4_(d: void_p, ab: void_p, cd: void_p, s0: float32, s1: float32, s2: float32, s3: float32):
+        attrs.func_kind = "cuda_internal"
+        attrs.func_name = "scale_bf16x4_to_fp8e4m3x4"
+        BlackBoxStmt(bf16_template, ab, cd, s0, s1, s2, s3, d)
+
+    for func in [scale_bf16x4_to_fp8e4m3x4_]:
+        assert isinstance(func, Function)
+        register_primitive_function(name=func.name, func_or_type=func)
+
+
+def scale_bf16x4_to_fp8e4m3x4(d: Expr, ab: Expr, cd: Expr, s0: Expr, s1: Expr, s2: Expr, s3: Expr) -> Expr:
+    return call_primitive_func("scale_bf16x4_to_fp8e4m3x4", args=[d, ab, cd, s0, s1, s2, s3])
diff --git a/python/tilus/hidet/ir/primitives/cuda/math/bfloat16.py b/python/tilus/hidet/ir/primitives/cuda/math/bfloat16.py
index 060ae0b7..b05e7e52 100644
--- a/python/tilus/hidet/ir/primitives/cuda/math/bfloat16.py
+++ b/python/tilus/hidet/ir/primitives/cuda/math/bfloat16.py
@@ -43,6 +43,7 @@ def register(self):
             "round": ["hrint", 1],
             "ceil": ["hceil", 1],
             "floor": ["hfloor", 1],
+            "abs": ["__habs", 1],
             "min": ["__hmin", 2],
             "max": ["__hmax", 2],
             "fma": ["__hfma", 3],
@@ -99,6 +100,9 @@ def ceil(self, a: Expr) -> Expr:
     def floor(self, a: Expr) -> Expr:
         return self.call("cuda_bf16_floor", a)
 
+    def abs(self, a: Expr) -> Expr:
+        return self.call("cuda_bf16_abs", a)
+
     def min(self, a: Expr, b: Expr) -> Expr:
         return self.call("cuda_bf16_min", a, b)
 
diff --git a/python/tilus/ir/builders/stmt_builder.py b/python/tilus/ir/builders/stmt_builder.py
index 549e818d..bb870b6c 100644
--- a/python/tilus/ir/builders/stmt_builder.py
+++ b/python/tilus/ir/builders/stmt_builder.py
@@ -119,6 +119,7 @@
     SqueezeInst,
     StoreGlobalGenericInst,
     StoreGlobalInst,
+    StoreScaledFp8E4M3FromSharedInst,
     StoreGlobalScatterInst,
     StoreSharedInst,
     StoreSharedScatterInst,
@@ -1273,6 +1274,16 @@ def store_global(
         inst = StoreGlobalInst.create(dst=dst, x=src, offsets=[as_expr(ofs) for ofs in offsets], dims=dims)
         self.append(inst)
 
+    def store_scaled_fp8e4m3_from_shared(
+        self, dst: GlobalTensor, src: SharedTensor, inv_scale: RegisterTensor, offsets: Sequence[Expr | int]
+    ) -> None:
+        self.append(
+            StoreScaledFp8E4M3FromSharedInst.create(
+                dst=dst, src=src, inv_scale=inv_scale, offsets=[as_expr(offset) for offset in offsets]
+            )
+        )
+
+
     def store_global_scatter(
         self,
         dst: GlobalTensor,
diff --git a/python/tilus/ir/instructions/__init__.py b/python/tilus/ir/instructions/__init__.py
index e71bc540..a2321d4d 100644
--- a/python/tilus/ir/instructions/__init__.py
+++ b/python/tilus/ir/instructions/__init__.py
@@ -77,6 +77,7 @@
     SqueezeInst,
     StoreGlobalGenericInst,
     StoreGlobalInst,
+    StoreScaledFp8E4M3FromSharedInst,
     StoreGlobalScatterInst,
     StoreSharedInst,
     StoreSharedScatterInst,
diff --git a/python/tilus/ir/instructions/generic.py b/python/tilus/ir/instructions/generic.py
index cf66b572..12733a6a 100644
--- a/python/tilus/ir/instructions/generic.py
+++ b/python/tilus/ir/instructions/generic.py
@@ -87,6 +87,23 @@ def create(dst: GlobalTensor, x: RegisterTensor, offsets: Sequence[Expr], dims:
         return StoreGlobalInst(output=None, inputs=(dst, x), offsets=tuple(offsets), dims=tuple(dims))
 
 
+@dataclass(frozen=True, eq=False)
+class StoreScaledFp8E4M3FromSharedInst(Instruction):
+    """Fused shared-BF16 to scaled global E4M3 store for 128x128 tiles."""
+
+    offsets: tuple[Expr, Expr]
+
+    @staticmethod
+    def create(
+        dst: GlobalTensor, src: SharedTensor, inv_scale: RegisterTensor, offsets: Sequence[Expr]
+    ) -> StoreScaledFp8E4M3FromSharedInst:
+        if len(offsets) != 2:
+            raise InstructionError("StoreScaledFp8E4M3FromSharedInst expects two offsets")
+        return StoreScaledFp8E4M3FromSharedInst(output=None, inputs=(dst, src, inv_scale), offsets=tuple(offsets))
+
+
+
+
 @dataclass(frozen=True, eq=False)
 class SliceGlobalInst(Instruction):
     offsets: tuple[Expr, ...]
diff --git a/python/tilus/ir/layout/inference/inference_rules/empty_rule.py b/python/tilus/ir/layout/inference/inference_rules/empty_rule.py
index 50f2184e..cebe9d89 100644
--- a/python/tilus/ir/layout/inference/inference_rules/empty_rule.py
+++ b/python/tilus/ir/layout/inference/inference_rules/empty_rule.py
@@ -20,6 +20,7 @@
     GlobalViewInst,
     PrintTensorInst,
     StoreGlobalInst,
+    StoreScaledFp8E4M3FromSharedInst,
 )
 from tilus.ir.instructions.cuda.cp_async_bulk import (
     CopyAsyncBulkGlobalToClusterSharedInst,
@@ -48,6 +49,7 @@
 @register_rule(FreeSharedInst)
 @register_rule(AllocateRegisterInst)
 @register_rule(StoreGlobalInst)
+@register_rule(StoreScaledFp8E4M3FromSharedInst)
 class EmptyRule(LayoutInferenceRule):
     @staticmethod
     def validate(inst: GlobalViewInst) -> bool:
diff --git a/python/tilus/ir/layout/inference/validation_rules/always_ok.py b/python/tilus/ir/layout/inference/validation_rules/always_ok.py
index 4b680f37..94ace98b 100644
--- a/python/tilus/ir/layout/inference/validation_rules/always_ok.py
+++ b/python/tilus/ir/layout/inference/validation_rules/always_ok.py
@@ -36,6 +36,7 @@
     SliceSharedInst,
     StoreGlobalGenericInst,
     StoreGlobalInst,
+    StoreScaledFp8E4M3FromSharedInst,
     StoreGlobalScatterInst,
     StoreSharedInst,
     StoreSharedScatterInst,
@@ -90,6 +91,7 @@
 @register_rule(LoadGlobalInst)
 @register_rule(LoadGlobalGenericInst)
 @register_rule(StoreGlobalInst)
+@register_rule(StoreScaledFp8E4M3FromSharedInst)
 @register_rule(SliceSharedInst)
 @register_rule(PermuteSharedInst)
 @register_rule(ReshapeSharedInst)
diff --git a/python/tilus/lang/instructions/root.py b/python/tilus/lang/instructions/root.py
index 18d01cfd..249048d9 100644
--- a/python/tilus/lang/instructions/root.py
+++ b/python/tilus/lang/instructions/root.py
@@ -17,7 +17,7 @@
 
 from tilus.hidet.ir.dtypes import boolean
 from tilus.hidet.ir.expr import Constant, Expr, Var, as_expr
-from tilus.hidet.ir.primitives.cuda.vars import blockIdx, gridDim
+from tilus.hidet.ir.primitives.cuda.vars import blockIdx, gridDim, threadIdx
 from tilus.hidet.ir.tools import infer_type
 from tilus.hidet.ir.type import DataType
 from tilus.ir.inst import InstructionError
@@ -30,6 +30,10 @@
 
 
 class RootInstructionGroup(InstructionGroup):
+    def get_thread_binding(self) -> Expr:
+        """Return the physical CUDA thread index within the thread block."""
+        return threadIdx.x
+
     @property
     def blockIdx(self) -> Dim3:
         """Get the block index of the current thread block."""
@@ -601,6 +605,12 @@ def store_global_scatter(
         """
         self._builder.store_global_scatter(dst=dst, indices=indices, values=values, dim=dim)
 
+    def store_scaled_fp8e4m3_from_shared(
+        self, dst: GlobalTensor, src: SharedTensor, inv_scale: RegisterTensor, *, offsets: Sequence[Expr | int]
+    ) -> None:
+        """Store a 128x128 shared BF16 tile as per-column scaled E4M3 FP8."""
+        self._builder.store_scaled_fp8e4m3_from_shared(dst=dst, src=src, inv_scale=inv_scale, offsets=offsets)
+
     def store_shared_scatter(
         self,
         dst: SharedTensor,

From 6021303870e5677429322c8d4bacdbfce1011066 Mon Sep 17 00:00:00 2001
From: William Zhang <wzhang20@yahoo.com>
Date: Fri, 11 Sep 2026 20:12:21 -0400
Subject: [PATCH 3/5] clean up tilekernels

Signed-off-by: William Zhang <wzhang20@yahoo.com>
---
 examples/moe/topk_gate.py                     |  37 ++++-
 examples/quantization/fp8_check.py            | 155 ------------------
 examples/quantization/per_channel_cast.py     |  27 ++-
 examples/quantization/per_token_cast.py       |  61 ++++++-
 .../swiglu_forward_and_per_token_cast.py      |  61 ++++++-
 python/tilus/backends/emitters/__init__.py    |   1 -
 .../tilus/backends/emitters/fp8_epilogue.py   |  53 ------
 .../tilus/hidet/ir/primitives/cuda/float8.py  |  40 -----
 .../hidet/ir/primitives/cuda/math/float16.py  |   9 +-
 python/tilus/ir/builders/stmt_builder.py      |  10 --
 python/tilus/ir/instructions/__init__.py      |   1 -
 python/tilus/ir/instructions/generic.py       |  17 --
 .../inference/inference_rules/empty_rule.py   |   2 -
 .../inference/validation_rules/always_ok.py   |   2 -
 python/tilus/lang/instructions/root.py        |   6 -
 tests/examples/test_examples.py               |   7 +-
 tests/instructions/test_reduce.py             | 120 +++++++++++++-
 17 files changed, 281 insertions(+), 328 deletions(-)
 delete mode 100644 examples/quantization/fp8_check.py
 delete mode 100644 python/tilus/backends/emitters/fp8_epilogue.py
 delete mode 100644 python/tilus/hidet/ir/primitives/cuda/float8.py

diff --git a/examples/moe/topk_gate.py b/examples/moe/topk_gate.py
index 997003a5..5e359630 100644
--- a/examples/moe/topk_gate.py
+++ b/examples/moe/topk_gate.py
@@ -35,32 +35,48 @@ def __call__(self, num_tokens: int32, scores_ptr: ~float32, output_ptr: ~int64):
         expert_ids = self.register_tensor(
             dtype=int32, shape=[1, self.aligned_experts], init=lambda _, j: j
         )
-        negative_max = -3.402823466e38
-        # A vector load past the logical expert dimension is zero-filled by
-        # Tilus.  Only materialize a validity mask when there actually is a
-        # tail: for the standard 256-expert route it is compile-time dead work.
+        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:
-            valid = self.register_tensor(
+            active = self.register_tensor(
                 dtype=int32,
                 shape=[1, self.aligned_experts],
                 init=lambda _, j: j < self.num_experts,
             )
-            values = self.where(valid != 0, x=values, y=negative_max)
+        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(values == best_value, x=expert_ids, y=int32.max_value)
+            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)
+            # Only a padding-sentinel maximum can leave no matching active
+            # candidate (all valid scores are -inf).  Keeping this fallback in
+            # a uniform runtime branch avoids a second warp reduction in the
+            # normal finite-score path.
+            if best_value[0, 0].item() == padding_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])
-            values = self.where(expert_ids == best_index, x=negative_max, y=values)
+            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():
@@ -74,6 +90,11 @@ def main():
         kernel(num_tokens, scores, output)
         expected = torch.sort(scores, dim=1, descending=True, stable=True).indices[:, :num_topk]
         torch.testing.assert_close(output, expected)
+        # Padding must never win when valid experts contain -inf.
+        if num_experts % 32:
+            scores.fill_(float("-inf"))
+            kernel(num_tokens, scores, output)
+            torch.testing.assert_close(output, torch.arange(num_topk, device="cuda", dtype=torch.int64)[None, :].expand(num_tokens, -1))
         def run_tilus():
             out = torch.empty(num_tokens, num_topk, device="cuda", dtype=torch.int64)
             kernel(num_tokens, scores, out)
diff --git a/examples/quantization/fp8_check.py b/examples/quantization/fp8_check.py
deleted file mode 100644
index 4d4104e7..00000000
--- a/examples/quantization/fp8_check.py
+++ /dev/null
@@ -1,155 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Shared reference and correctness checks for the per-token FP8 cast examples.
-
-FP8 e4m3 has a three-bit mantissa, so adjacent representable values are up to
-6.25% apart and the ladder tops out at 448.  Comparing *decoded* e4m3 values
-with an absolute float tolerance is therefore nearly vacuous: a tolerance loose
-enough to absorb a single rounding step near the top of the range is +/-32,
-which is also loose enough to accept a kernel that is quantizing against a
-completely wrong scale factor.
-
-These helpers compare on the code ladder instead.  e4m3 is stored
-sign-magnitude, so ranking the seven magnitude bits and re-applying the sign
-gives an integer ordinal in which *any* two adjacent representable values
-differ by exactly one, uniformly across the dynamic range.  "Agrees to within
-one rounding step" is then ``|ordinal(a) - ordinal(b)| <= 1``.
-
-One rounding step is the tightest claim that holds here.  Tilus compiles every
-kernel with ``-prec-div=false`` (see ``python/tilus/hidet/backend/build.py``),
-so the fp32 divisions that produce the scale factor and its reciprocal are
-approximate to about one ulp.  TileLang uses exact division.  With identical
-inputs and otherwise identical arithmetic the two therefore still straddle an
-e4m3 rounding boundary on a small fraction of elements; everything else is
-bit-identical.  A real bug moves the mismatch rate or the code distance well
-outside those bounds.
-"""
-
-from typing import NamedTuple
-
-import torch
-
-# Largest finite magnitude representable in e4m3 (TileKernels' ``T.max_value``).
-E4M3_MAX = 448.0
-
-# TileKernels clamps the group absmax from below before dividing, so that an
-# all-zero group yields a tiny scale rather than a division by zero.  See
-# ``CastOutputConfig.clamp_min_value`` in ``tile_kernels/quant/common.py``.
-SF_CLAMP_MIN = 1e-4
-
-
-class CodeLadderStats(NamedTuple):
-    """Result of comparing two e4m3 tensors on the code ladder."""
-
-    max_code_diff: int
-    mismatch_frac: float
-
-
-def fp8_ordinal(x: torch.Tensor) -> torch.Tensor:
-    """Rank each e4m3 value on the ladder of representable values.
-
-    e4m3 is sign-magnitude, so the seven magnitude bits are already a monotone
-    rank within one sign.  Negating them for the negative half gives a single
-    signed ordinal where consecutive representable values always differ by one.
-    """
-    assert x.dtype == torch.float8_e4m3fn
-    bits = x.view(torch.uint8).to(torch.int32)
-    magnitude = bits & 0x7F
-    return torch.where(bits & 0x80 != 0, -magnitude, magnitude)
-
-
-def check_fp8_close(
-    actual: torch.Tensor,
-    expected: torch.Tensor,
-    *,
-    label: str,
-    max_mismatch_frac: float = 0.01,
-) -> CodeLadderStats:
-    """Assert two e4m3 tensors agree to within one rounding step.
-
-    Fails if any element is more than one representable value away, or if more
-    than ``max_mismatch_frac`` of elements disagree at all.  The second bound
-    matters: a systematic error that happens to be small still shows up as a
-    mismatch rate far above the ~0.1% produced by fp32 division rounding.
-    """
-    code_diff = (fp8_ordinal(actual) - fp8_ordinal(expected)).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 "
-        f"1 (a single rounding step)"
-    )
-    assert mismatch_frac <= max_mismatch_frac, (
-        f"{label}: {mismatch_frac:.4%} of elements differ, above the "
-        f"{max_mismatch_frac:.4%} budget for fp32 division rounding"
-    )
-    return CodeLadderStats(max_code_diff, mismatch_frac)
-
-
-def check_scales_close(
-    actual: torch.Tensor,
-    expected: torch.Tensor,
-    *,
-    label: str,
-    rtol: float = 1e-6,
-) -> float:
-    """Assert two fp32 scale-factor tensors agree to a few fp32 ulps.
-
-    The scale is ``max(absmax, 1e-4) / 448`` on both sides, computed from an
-    absmax that is a max over identical inputs and so is bit-identical.  Only
-    the division differs, hence a tolerance in ulps (1 ulp ~ 1.2e-7 relative)
-    rather than the 1e-5 that would also pass a wrong reduction.
-    """
-    torch.testing.assert_close(actual, expected, rtol=rtol, atol=0.0, msg=label)
-    return float(((actual - expected).abs() / expected.abs()).max().item())
-
-
-def torch_per_token_cast(
-    values: torch.Tensor,
-    num_per_channels: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Ground-truth per-token FP8 cast in fp32 PyTorch.
-
-    ``values`` holds the activations to quantize (for SwiGLU, the post-SwiGLU
-    values).  Checking both kernels against this catches the case where they
-    agree with each other but are both wrong.
-    """
-    num_tokens, hidden = values.shape
-    grouped = values.float().reshape(num_tokens, hidden // num_per_channels, -1)
-    amax = grouped.abs().amax(dim=-1, keepdim=True).clamp_min(SF_CLAMP_MIN)
-    scale = (amax / E4M3_MAX).squeeze(-1)
-    scaled = (grouped * (E4M3_MAX / amax)).clamp(-E4M3_MAX, E4M3_MAX)
-    out = scaled.reshape(num_tokens, hidden).to(torch.float8_e4m3fn)
-    return out, scale
-
-
-def dequantize(
-    out: torch.Tensor,
-    scales: torch.Tensor,
-    num_per_channels: int,
-) -> torch.Tensor:
-    """Decode an e4m3 tensor and its per-group scales back to fp32."""
-    grouped = out.float().reshape(
-        out.shape[0],
-        out.shape[1] // num_per_channels,
-        num_per_channels,
-    )
-    return (grouped * scales[:, :, None]).reshape(out.shape)
-
-
-def quantization_snr_db(reference: torch.Tensor, dequantized: torch.Tensor) -> float:
-    """Signal-to-noise ratio of a dequantized tensor against its fp32 source.
-
-    Reported rather than asserted on: it is the sanity check that the cast is
-    carrying real information.  Round-to-nearest into a three-bit mantissa
-    lands near 32 dB for Gaussian activations, and is stable across shapes, so
-    a value well below that means the scale factors are wrong even if the code
-    ladder checks pass.
-    """
-    reference = reference.float()
-    noise = dequantized.float() - reference
-    return float(
-        10.0
-        * torch.log10(reference.square().sum() / noise.square().sum().clamp_min(1e-30))
-    )
diff --git a/examples/quantization/per_channel_cast.py b/examples/quantization/per_channel_cast.py
index 6b911913..bf601ce5 100644
--- a/examples/quantization/per_channel_cast.py
+++ b/examples/quantization/per_channel_cast.py
@@ -1,10 +1,9 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""TileKernels-compatible FP8 e4m3 cast with one scale per 128-token column."""
+"""Common BF16 128x128 path of TileKernels' per-channel E4M3 cast."""
 
 import tilus
 import torch
-from fp8_check import check_fp8_close
 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
@@ -13,11 +12,25 @@
 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):
-    def __init__(self, block_n: int = 128, warps: int = 8):
+    """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, block_n, warps
+        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))
@@ -46,7 +59,11 @@ def __call__(self, num_tokens: int32, hidden: int32, x_ptr: ~bfloat16, out_ptr:
         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])
-        self.store_scaled_fp8e4m3_from_shared(out, shared_x, inv_scale, offsets=[offset_m, 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)
 
 
diff --git a/examples/quantization/per_token_cast.py b/examples/quantization/per_token_cast.py
index 01ef260f..1dbf1685 100644
--- a/examples/quantization/per_token_cast.py
+++ b/examples/quantization/per_token_cast.py
@@ -20,22 +20,65 @@
 its reciprocal is ``448 / absmax``.
 """
 
+from typing import NamedTuple
+
 import pandas
 import tilus
 import torch
-from fp8_check import (
-    E4M3_MAX,
-    SF_CLAMP_MIN,
-    check_fp8_close,
-    check_scales_close,
-    dequantize,
-    quantization_snr_db,
-    torch_per_token_cast,
-)
 from tile_kernels.quant.per_token_cast_kernel import per_token_cast
 from tilus import bfloat16, float8_e4m3, float32, int32
 from tilus.utils import benchmark_func, cdiv
 
+E4M3_MAX = 448.0
+SF_CLAMP_MIN = 1e-4
+
+
+class CodeLadderStats(NamedTuple):
+    max_code_diff: int
+    mismatch_frac: float
+
+
+def fp8_ordinal(x: torch.Tensor) -> torch.Tensor:
+    """Map e4m3 values to consecutive signed integer codes."""
+    assert x.dtype == torch.float8_e4m3fn
+    bits = x.view(torch.uint8).to(torch.int32)
+    magnitude = bits & 0x7F
+    return torch.where(bits & 0x80 != 0, -magnitude, magnitude)
+
+
+def check_fp8_close(actual: torch.Tensor, expected: torch.Tensor, *, label: str, max_mismatch_frac: float = 0.01) -> CodeLadderStats:
+    code_diff = (fp8_ordinal(actual) - fp8_ordinal(expected)).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"
+    return CodeLadderStats(max_code_diff, mismatch_frac)
+
+
+def check_scales_close(actual: torch.Tensor, expected: torch.Tensor, *, label: str, rtol: float = 1e-6) -> float:
+    torch.testing.assert_close(actual, expected, rtol=rtol, atol=0.0, msg=label)
+    return float(((actual - expected).abs() / expected.abs()).max().item())
+
+
+def torch_per_token_cast(values: torch.Tensor, num_per_channels: int) -> tuple[torch.Tensor, torch.Tensor]:
+    num_tokens, hidden = values.shape
+    grouped = values.float().reshape(num_tokens, hidden // num_per_channels, -1)
+    amax = grouped.abs().amax(dim=-1, keepdim=True).clamp_min(SF_CLAMP_MIN)
+    scale = (amax / E4M3_MAX).squeeze(-1)
+    out = (grouped * (E4M3_MAX / amax)).clamp(-E4M3_MAX, E4M3_MAX).reshape(num_tokens, hidden).to(torch.float8_e4m3fn)
+    return out, scale
+
+
+def dequantize(out: torch.Tensor, scales: torch.Tensor, num_per_channels: int) -> torch.Tensor:
+    grouped = out.float().reshape(out.shape[0], out.shape[1] // num_per_channels, num_per_channels)
+    return (grouped * scales[:, :, None]).reshape(out.shape)
+
+
+def quantization_snr_db(reference: torch.Tensor, dequantized: torch.Tensor) -> float:
+    noise = dequantized.float() - reference.float()
+    return float(10.0 * torch.log10(reference.float().square().sum() / noise.square().sum().clamp_min(1e-30)))
+
+
 
 @tilus.autotune("block_m", [1, 2, 4])
 @tilus.autotune("groups_per_block", [1, 4, 16, 64])
diff --git a/examples/quantization/swiglu_forward_and_per_token_cast.py b/examples/quantization/swiglu_forward_and_per_token_cast.py
index 50fef78c..e9a2d12a 100644
--- a/examples/quantization/swiglu_forward_and_per_token_cast.py
+++ b/examples/quantization/swiglu_forward_and_per_token_cast.py
@@ -23,24 +23,67 @@
 its reciprocal is ``448 / absmax``.
 """
 
+from typing import NamedTuple
+
 import pandas
 import tilus
 import torch
-from fp8_check import (
-    E4M3_MAX,
-    SF_CLAMP_MIN,
-    check_fp8_close,
-    check_scales_close,
-    dequantize,
-    quantization_snr_db,
-    torch_per_token_cast,
-)
 from tile_kernels.quant.swiglu_forward_and_per_token_cast_kernel import (
     swiglu_forward_and_per_token_cast,
 )
 from tilus import bfloat16, float8_e4m3, float32, int32
 from tilus.utils import benchmark_func, cdiv
 
+E4M3_MAX = 448.0
+SF_CLAMP_MIN = 1e-4
+
+
+class CodeLadderStats(NamedTuple):
+    max_code_diff: int
+    mismatch_frac: float
+
+
+def fp8_ordinal(x: torch.Tensor) -> torch.Tensor:
+    """Map e4m3 values to consecutive signed integer codes."""
+    assert x.dtype == torch.float8_e4m3fn
+    bits = x.view(torch.uint8).to(torch.int32)
+    magnitude = bits & 0x7F
+    return torch.where(bits & 0x80 != 0, -magnitude, magnitude)
+
+
+def check_fp8_close(actual: torch.Tensor, expected: torch.Tensor, *, label: str, max_mismatch_frac: float = 0.01) -> CodeLadderStats:
+    code_diff = (fp8_ordinal(actual) - fp8_ordinal(expected)).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"
+    return CodeLadderStats(max_code_diff, mismatch_frac)
+
+
+def check_scales_close(actual: torch.Tensor, expected: torch.Tensor, *, label: str, rtol: float = 1e-6) -> float:
+    torch.testing.assert_close(actual, expected, rtol=rtol, atol=0.0, msg=label)
+    return float(((actual - expected).abs() / expected.abs()).max().item())
+
+
+def torch_per_token_cast(values: torch.Tensor, num_per_channels: int) -> tuple[torch.Tensor, torch.Tensor]:
+    num_tokens, hidden = values.shape
+    grouped = values.float().reshape(num_tokens, hidden // num_per_channels, -1)
+    amax = grouped.abs().amax(dim=-1, keepdim=True).clamp_min(SF_CLAMP_MIN)
+    scale = (amax / E4M3_MAX).squeeze(-1)
+    out = (grouped * (E4M3_MAX / amax)).clamp(-E4M3_MAX, E4M3_MAX).reshape(num_tokens, hidden).to(torch.float8_e4m3fn)
+    return out, scale
+
+
+def dequantize(out: torch.Tensor, scales: torch.Tensor, num_per_channels: int) -> torch.Tensor:
+    grouped = out.float().reshape(out.shape[0], out.shape[1] // num_per_channels, num_per_channels)
+    return (grouped * scales[:, :, None]).reshape(out.shape)
+
+
+def quantization_snr_db(reference: torch.Tensor, dequantized: torch.Tensor) -> float:
+    noise = dequantized.float() - reference.float()
+    return float(10.0 * torch.log10(reference.float().square().sum() / noise.square().sum().clamp_min(1e-30)))
+
+
 
 @tilus.autotune("block_m", [1])
 @tilus.autotune("groups_per_block", [1, 4, 16, 64])
diff --git a/python/tilus/backends/emitters/__init__.py b/python/tilus/backends/emitters/__init__.py
index e2515012..4a1863af 100644
--- a/python/tilus/backends/emitters/__init__.py
+++ b/python/tilus/backends/emitters/__init__.py
@@ -22,7 +22,6 @@
     cuda,
     debug,
     elementwise,
-    fp8_epilogue,
     gmem,
     ldst,
     random,
diff --git a/python/tilus/backends/emitters/fp8_epilogue.py b/python/tilus/backends/emitters/fp8_epilogue.py
deleted file mode 100644
index db5ecf23..00000000
--- a/python/tilus/backends/emitters/fp8_epilogue.py
+++ /dev/null
@@ -1,53 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Fused FP8 epilogues for bandwidth-bound quantization kernels."""
-
-from tilus.backends.emitter import BaseInstEmitter, register_emitter
-from tilus.hidet.ir.dtypes import bfloat16, float8_e4m3, float32, uint32
-from tilus.hidet.ir.expr import cast
-from tilus.hidet.ir.primitives.cuda.float8 import scale_bf16x4_to_fp8e4m3x4
-from tilus.hidet.ir.type import void_p
-from tilus.ir.instructions import StoreScaledFp8E4M3FromSharedInst
-from tilus.ir.tensor import GlobalTensor, RegisterTensor, SharedTensor
-
-
-@register_emitter(StoreScaledFp8E4M3FromSharedInst)
-class StoreScaledFp8E4M3FromSharedEmitter(BaseInstEmitter):
-    """Lower the 128x128 per-channel FP8 epilogue without register tiles."""
-
-    def emit(self, inst: StoreScaledFp8E4M3FromSharedInst) -> None:
-        dst: GlobalTensor = inst.inputs[0].as_global_tensor()
-        src: SharedTensor = inst.inputs[1].as_shared_tensor()
-        inv_scale: RegisterTensor = inst.inputs[2].as_register_tensor()
-        if (
-            dst.dtype != float8_e4m3
-            or src.dtype != bfloat16
-            or inv_scale.dtype != float32
-            or tuple(src.shape) != (128, 128)
-            or tuple(inv_scale.shape) != (1, 128)
-            or inv_scale.layout.local_size != 4
-        ):
-            raise ValueError("fused FP8 epilogue requires shared bf16[128,128] and register float32[1,128]")
-
-        dst_buf = self.tensor2var[dst]
-        src_buf = self.tensor2var[src]
-        scale_buf = self.tensor2var[inv_scale]
-        offset_m, offset_n = inst.offsets
-        lane_id = self.lane_id()
-        warp_id = self.warp_id()
-        col = lane_id * 4
-        with self.for_range(16, attr="u+") as i:
-            row = warp_id * 16 + i
-            src_offset = src.layout(row, col)
-            dst_offset = dst.layout(offset_m + row, offset_n + col)
-            self.append(
-                scale_bf16x4_to_fp8e4m3x4(
-                    cast(~dst_buf[dst_offset], void_p),
-                    cast(~src_buf[src_offset], void_p),
-                    cast(~src_buf[src_offset + 2], void_p),
-                    scale_buf[0],
-                    scale_buf[1],
-                    scale_buf[2],
-                    scale_buf[3],
-                )
-            )
diff --git a/python/tilus/hidet/ir/primitives/cuda/float8.py b/python/tilus/hidet/ir/primitives/cuda/float8.py
deleted file mode 100644
index 498005b3..00000000
--- a/python/tilus/hidet/ir/primitives/cuda/float8.py
+++ /dev/null
@@ -1,40 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Small CUDA FP8 conversion primitives used by vectorized epilogues."""
-
-from typing import no_type_check
-
-from tilus.hidet.ir.expr import Expr
-from tilus.hidet.ir.func import Function
-from tilus.hidet.ir.primitives.func import call_primitive_func, register_primitive_function
-from tilus.hidet.ir.stmt import BlackBoxStmt
-from tilus.hidet.utils import initialize
-
-
-@initialize()
-def register_functions():
-    from tilus.hidet.lang import attrs, script  # pylint: disable=import-outside-toplevel
-    from tilus.hidet.lang.types import float32, void_p
-
-    bf16_template = r"""
-    float2 f0 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>({}));
-    float2 f1 = __bfloat1622float2(*reinterpret_cast<const __nv_bfloat162*>({}));
-    __nv_fp8x2_storage_t p0 = __nv_cvt_float2_to_fp8x2(make_float2(f0.x * {}, f0.y * {}), __NV_SATFINITE, __NV_E4M3);
-    __nv_fp8x2_storage_t p1 = __nv_cvt_float2_to_fp8x2(make_float2(f1.x * {}, f1.y * {}), __NV_SATFINITE, __NV_E4M3);
-    *reinterpret_cast<uint32_t*>({}) = static_cast<uint32_t>(p0) | (static_cast<uint32_t>(p1) << 16);
-    """
-
-    @no_type_check
-    @script
-    def scale_bf16x4_to_fp8e4m3x4_(d: void_p, ab: void_p, cd: void_p, s0: float32, s1: float32, s2: float32, s3: float32):
-        attrs.func_kind = "cuda_internal"
-        attrs.func_name = "scale_bf16x4_to_fp8e4m3x4"
-        BlackBoxStmt(bf16_template, ab, cd, s0, s1, s2, s3, d)
-
-    for func in [scale_bf16x4_to_fp8e4m3x4_]:
-        assert isinstance(func, Function)
-        register_primitive_function(name=func.name, func_or_type=func)
-
-
-def scale_bf16x4_to_fp8e4m3x4(d: Expr, ab: Expr, cd: Expr, s0: Expr, s1: Expr, s2: Expr, s3: Expr) -> Expr:
-    return call_primitive_func("scale_bf16x4_to_fp8e4m3x4", args=[d, ab, cd, s0, s1, s2, s3])
diff --git a/python/tilus/hidet/ir/primitives/cuda/math/float16.py b/python/tilus/hidet/ir/primitives/cuda/math/float16.py
index 6a1eb57a..5a1d13ad 100644
--- a/python/tilus/hidet/ir/primitives/cuda/math/float16.py
+++ b/python/tilus/hidet/ir/primitives/cuda/math/float16.py
@@ -32,6 +32,7 @@
 from tilus.hidet.ir.primitives.math import MathFunctionSet, register_math_function_set
 from tilus.hidet.ir.type import DataType, FuncType
 from tilus.hidet.utils import initialize
+from tilus.target import get_current_target
 
 
 @initialize()
@@ -202,17 +203,13 @@ def floor(self, a: Expr) -> Expr:
         return self.call("cuda_f16_floor", a)
 
     def min(self, a: Expr, b: Expr) -> Expr:
-        arch_pair: Tuple[int, int] = hidet.option.cuda.get_arch_pair()
-
-        if arch_pair >= (8, 0):
+        if get_current_target().properties.compute_capability >= (8, 0):
             return self.call("cuda_f16_min_sm80", a, b)
         else:
             return self.call("cuda_f16_min", a, b)
 
     def max(self, a: Expr, b: Expr) -> Expr:
-        arch_pair: Tuple[int, int] = hidet.option.cuda.get_arch_pair()
-
-        if arch_pair >= (8, 0):
+        if get_current_target().properties.compute_capability >= (8, 0):
             return self.call("cuda_f16_max_sm80", a, b)
         else:
             return self.call("cuda_f16_max", a, b)
diff --git a/python/tilus/ir/builders/stmt_builder.py b/python/tilus/ir/builders/stmt_builder.py
index bb870b6c..40e48513 100644
--- a/python/tilus/ir/builders/stmt_builder.py
+++ b/python/tilus/ir/builders/stmt_builder.py
@@ -119,7 +119,6 @@
     SqueezeInst,
     StoreGlobalGenericInst,
     StoreGlobalInst,
-    StoreScaledFp8E4M3FromSharedInst,
     StoreGlobalScatterInst,
     StoreSharedInst,
     StoreSharedScatterInst,
@@ -1274,15 +1273,6 @@ def store_global(
         inst = StoreGlobalInst.create(dst=dst, x=src, offsets=[as_expr(ofs) for ofs in offsets], dims=dims)
         self.append(inst)
 
-    def store_scaled_fp8e4m3_from_shared(
-        self, dst: GlobalTensor, src: SharedTensor, inv_scale: RegisterTensor, offsets: Sequence[Expr | int]
-    ) -> None:
-        self.append(
-            StoreScaledFp8E4M3FromSharedInst.create(
-                dst=dst, src=src, inv_scale=inv_scale, offsets=[as_expr(offset) for offset in offsets]
-            )
-        )
-
 
     def store_global_scatter(
         self,
diff --git a/python/tilus/ir/instructions/__init__.py b/python/tilus/ir/instructions/__init__.py
index a2321d4d..e71bc540 100644
--- a/python/tilus/ir/instructions/__init__.py
+++ b/python/tilus/ir/instructions/__init__.py
@@ -77,7 +77,6 @@
     SqueezeInst,
     StoreGlobalGenericInst,
     StoreGlobalInst,
-    StoreScaledFp8E4M3FromSharedInst,
     StoreGlobalScatterInst,
     StoreSharedInst,
     StoreSharedScatterInst,
diff --git a/python/tilus/ir/instructions/generic.py b/python/tilus/ir/instructions/generic.py
index 12733a6a..cf66b572 100644
--- a/python/tilus/ir/instructions/generic.py
+++ b/python/tilus/ir/instructions/generic.py
@@ -87,23 +87,6 @@ def create(dst: GlobalTensor, x: RegisterTensor, offsets: Sequence[Expr], dims:
         return StoreGlobalInst(output=None, inputs=(dst, x), offsets=tuple(offsets), dims=tuple(dims))
 
 
-@dataclass(frozen=True, eq=False)
-class StoreScaledFp8E4M3FromSharedInst(Instruction):
-    """Fused shared-BF16 to scaled global E4M3 store for 128x128 tiles."""
-
-    offsets: tuple[Expr, Expr]
-
-    @staticmethod
-    def create(
-        dst: GlobalTensor, src: SharedTensor, inv_scale: RegisterTensor, offsets: Sequence[Expr]
-    ) -> StoreScaledFp8E4M3FromSharedInst:
-        if len(offsets) != 2:
-            raise InstructionError("StoreScaledFp8E4M3FromSharedInst expects two offsets")
-        return StoreScaledFp8E4M3FromSharedInst(output=None, inputs=(dst, src, inv_scale), offsets=tuple(offsets))
-
-
-
-
 @dataclass(frozen=True, eq=False)
 class SliceGlobalInst(Instruction):
     offsets: tuple[Expr, ...]
diff --git a/python/tilus/ir/layout/inference/inference_rules/empty_rule.py b/python/tilus/ir/layout/inference/inference_rules/empty_rule.py
index cebe9d89..50f2184e 100644
--- a/python/tilus/ir/layout/inference/inference_rules/empty_rule.py
+++ b/python/tilus/ir/layout/inference/inference_rules/empty_rule.py
@@ -20,7 +20,6 @@
     GlobalViewInst,
     PrintTensorInst,
     StoreGlobalInst,
-    StoreScaledFp8E4M3FromSharedInst,
 )
 from tilus.ir.instructions.cuda.cp_async_bulk import (
     CopyAsyncBulkGlobalToClusterSharedInst,
@@ -49,7 +48,6 @@
 @register_rule(FreeSharedInst)
 @register_rule(AllocateRegisterInst)
 @register_rule(StoreGlobalInst)
-@register_rule(StoreScaledFp8E4M3FromSharedInst)
 class EmptyRule(LayoutInferenceRule):
     @staticmethod
     def validate(inst: GlobalViewInst) -> bool:
diff --git a/python/tilus/ir/layout/inference/validation_rules/always_ok.py b/python/tilus/ir/layout/inference/validation_rules/always_ok.py
index 94ace98b..4b680f37 100644
--- a/python/tilus/ir/layout/inference/validation_rules/always_ok.py
+++ b/python/tilus/ir/layout/inference/validation_rules/always_ok.py
@@ -36,7 +36,6 @@
     SliceSharedInst,
     StoreGlobalGenericInst,
     StoreGlobalInst,
-    StoreScaledFp8E4M3FromSharedInst,
     StoreGlobalScatterInst,
     StoreSharedInst,
     StoreSharedScatterInst,
@@ -91,7 +90,6 @@
 @register_rule(LoadGlobalInst)
 @register_rule(LoadGlobalGenericInst)
 @register_rule(StoreGlobalInst)
-@register_rule(StoreScaledFp8E4M3FromSharedInst)
 @register_rule(SliceSharedInst)
 @register_rule(PermuteSharedInst)
 @register_rule(ReshapeSharedInst)
diff --git a/python/tilus/lang/instructions/root.py b/python/tilus/lang/instructions/root.py
index 249048d9..be92411a 100644
--- a/python/tilus/lang/instructions/root.py
+++ b/python/tilus/lang/instructions/root.py
@@ -605,12 +605,6 @@ def store_global_scatter(
         """
         self._builder.store_global_scatter(dst=dst, indices=indices, values=values, dim=dim)
 
-    def store_scaled_fp8e4m3_from_shared(
-        self, dst: GlobalTensor, src: SharedTensor, inv_scale: RegisterTensor, *, offsets: Sequence[Expr | int]
-    ) -> None:
-        """Store a 128x128 shared BF16 tile as per-column scaled E4M3 FP8."""
-        self._builder.store_scaled_fp8e4m3_from_shared(dst=dst, src=src, inv_scale=inv_scale, offsets=offsets)
-
     def store_shared_scatter(
         self,
         dst: SharedTensor,
diff --git a/tests/examples/test_examples.py b/tests/examples/test_examples.py
index d99898f9..173e55eb 100644
--- a/tests/examples/test_examples.py
+++ b/tests/examples/test_examples.py
@@ -84,8 +84,11 @@
     # Benchmark utilities
     ("blackwell_matmul", "benchmark.py"),
     ("hopper_matmul", "benchmark.py"),
-    # Shared correctness helpers for the quantization examples
-    ("quantization", "fp8_check.py"),
+    # Standalone TileKernels comparison benchmarks; exercised explicitly in
+    # the port benchmark workflow rather than the generic example smoke test.
+    ("mhc", "normw_merge.py"),
+    ("moe", "topk_gate.py"),
+    ("quantization", "per_channel_cast.py"),
 ]
 
 
diff --git a/tests/instructions/test_reduce.py b/tests/instructions/test_reduce.py
index 52b71ddf..ed9bf748 100644
--- a/tests/instructions/test_reduce.py
+++ b/tests/instructions/test_reduce.py
@@ -15,9 +15,9 @@
 import pytest
 import tilus
 import torch
-from tilus import boolean, int32
+from tilus import bfloat16, boolean, float16, float32, int32
 from tilus.ir.layout import RegisterLayout, register_layout
-from tilus.ir.layout.ops import spatial
+from tilus.ir.layout.ops import replicated, spatial
 
 
 class ReduceKernelExample(tilus.Script):
@@ -55,6 +55,82 @@ def __call__(self, x_ptr: ~int32, y_ptr: ~boolean) -> None:
         self.store_global(g_y, src=self.all(r_x != 0), offsets=[1], dims=[])
 
 
+class IntraWarpReductionMatrixExample(tilus.Script):
+    """Expose the reduction result from every lane in each warp-local group."""
+
+    def __init__(self, lane_width: int, num_warps: int, op: str, dtype):
+        super().__init__()
+        self.lane_width = lane_width
+        self.num_warps = num_warps
+        self.op = op
+        self.dtype = dtype
+        self.is_boolean = dtype == boolean
+        # The negative spatial mode replicates one logical reduction group across
+        # all other lanes.  Thus b[0] is defined in every thread and a store per
+        # physical lane verifies the broadcast part of the reduction contract.
+        self.layout = replicated(num_workers=32 * num_warps // lane_width) * spatial(lane_width)
+
+    def __call__(self, out_ptr: ~int32) -> None:
+        self.attrs.blocks = 1
+        self.attrs.warps = self.num_warps
+
+        if self.is_boolean:
+            a = self.register_tensor(dtype=self.dtype, shape=[self.lane_width], init=lambda i: (i % 2) == 0)
+        else:
+            a = self.register_tensor(dtype=self.dtype, shape=[self.lane_width], init=lambda _i: 1)
+        if self.op == "sum":
+            b = self.sum(a, dim=0, keepdim=True)
+        elif self.op == "max":
+            b = self.max(a, dim=0, keepdim=True)
+        elif self.op == "min":
+            b = self.min(a, dim=0, keepdim=True)
+        elif self.op == "any":
+            b = self.any(a, dim=0, keepdim=True)
+        elif self.op == "all":
+            b = self.all(a, dim=0, keepdim=True)
+        else:
+            raise ValueError(f"Unsupported operation: {self.op}")
+
+        g_out = self.global_view(ptr=out_ptr, dtype=int32, shape=[32 * self.num_warps])
+        self.store_global(g_out, b[0].to(int32), offsets=[self.get_thread_binding()], dims=[])
+        self.annotate_layout(a, self.layout)
+
+
+class InterWarpReductionMatrixExample(tilus.Script):
+    """Exercise the shared-memory inter-warp path in addition to XOR shuffles."""
+
+    def __init__(self, op: str, dtype):
+        super().__init__()
+        self.op = op
+        self.dtype = dtype
+        self.is_boolean = dtype == boolean
+
+    def __call__(self, out_ptr: ~int32) -> None:
+        self.attrs.blocks = 1
+        self.attrs.warps = 2
+        layout = spatial(2, 32)
+        if self.is_boolean:
+            a = self.register_tensor(dtype=self.dtype, shape=layout.shape, init=lambda i, j: (i + j) % 2 == 0)
+        else:
+            a = self.register_tensor(dtype=self.dtype, shape=layout.shape, init=lambda _i, _j: 1)
+        if self.op == "sum":
+            b = self.sum(a, dim=0, keepdim=True)
+        elif self.op == "max":
+            b = self.max(a, dim=0, keepdim=True)
+        elif self.op == "min":
+            b = self.min(a, dim=0, keepdim=True)
+        elif self.op == "any":
+            b = self.any(a, dim=0, keepdim=True)
+        elif self.op == "all":
+            b = self.all(a, dim=0, keepdim=True)
+        else:
+            raise ValueError(f"Unsupported operation: {self.op}")
+
+        g_out = self.global_view(ptr=out_ptr, dtype=int32, shape=b.shape)
+        self.store_global(g_out, b.to(int32), offsets=[0, 0], dims=[0, 1])
+        self.annotate_layout(a, layout)
+
+
 @pytest.mark.parametrize("dim", [0, 1])
 @pytest.mark.parametrize(
     "layout",
@@ -99,3 +175,43 @@ def test_any_all_reduce_instruction():
         y_actual = torch.empty_like(y)
         kernel(x, y_actual)
         assert torch.allclose(y_actual, y), f"Failed for x={x} and y={y}, y_actual={y_actual}"
+
+
+@pytest.mark.parametrize("lane_width", [2, 4, 8, 16, 32])
+@pytest.mark.parametrize("num_warps", [1, 2], ids=["single_warp", "multi_warp"])
+@pytest.mark.parametrize(
+    ("op", "dtype", "expected"),
+    [
+        ("sum", int32, lambda width: width),
+        ("sum", float32, lambda width: width),
+        ("max", float16, lambda _width: 1),
+        ("min", bfloat16, lambda _width: 1),
+        ("any", boolean, lambda _width: 1),
+        ("all", boolean, lambda _width: 0),
+    ],
+    ids=["sum_int32", "sum_fp32", "max_fp16", "min_bf16", "any", "all"],
+)
+def test_intra_warp_reduction_equivalence_matrix(lane_width: int, num_warps: int, op: str, dtype, expected):
+    """All lanes must agree for every supported XOR-reduction subgroup width."""
+    actual = torch.empty(32 * num_warps, dtype=torch.int32, device="cuda")
+    IntraWarpReductionMatrixExample(lane_width, num_warps, op, dtype)(actual)
+    torch.testing.assert_close(actual, torch.full_like(actual, expected(lane_width)))
+
+
+@pytest.mark.parametrize(
+    ("op", "dtype", "expected"),
+    [
+        ("sum", int32, 2),
+        ("sum", float32, 2),
+        ("max", float16, 1),
+        ("min", bfloat16, 1),
+        ("any", boolean, 1),
+        ("all", boolean, 0),
+    ],
+    ids=["sum_int32", "sum_fp32", "max_fp16", "min_bf16", "any", "all"],
+)
+def test_inter_warp_reduction_equivalence_matrix(op: str, dtype, expected: int):
+    """The shared-memory handoff preserves the same result for two warps."""
+    actual = torch.empty((1, 32), dtype=torch.int32, device="cuda")
+    InterWarpReductionMatrixExample(op, dtype)(actual)
+    torch.testing.assert_close(actual, torch.full_like(actual, expected))

From e4b5df8028840efed9900ebedcda5c383a051cf3 Mon Sep 17 00:00:00 2001
From: William Zhang <wzhang20@yahoo.com>
Date: Sat, 12 Sep 2026 22:10:16 -0400
Subject: [PATCH 4/5] modfiy topk kernel

Signed-off-by: William Zhang <wzhang20@yahoo.com>
---
 examples/moe/topk_gate.py | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/examples/moe/topk_gate.py b/examples/moe/topk_gate.py
index 5e359630..270d54f0 100644
--- a/examples/moe/topk_gate.py
+++ b/examples/moe/topk_gate.py
@@ -61,11 +61,11 @@ def __call__(self, num_tokens: int32, scores_ptr: ~float32, output_ptr: ~int64):
                 y=int32.max_value,
             )
             best_index = self.min(candidates, dim=1, keepdim=True)
-            # Only a padding-sentinel maximum can leave no matching active
-            # candidate (all valid scores are -inf).  Keeping this fallback in
-            # a uniform runtime branch avoids a second warp reduction in the
-            # normal finite-score path.
-            if best_value[0, 0].item() == padding_value:
+            # 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
                 )
@@ -90,15 +90,27 @@ def main():
         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:
-            scores.fill_(float("-inf"))
-            kernel(num_tokens, scores, output)
-            torch.testing.assert_close(output, torch.arange(num_topk, device="cuda", dtype=torch.int64)[None, :].expand(num_tokens, -1))
+            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))

From 58d60c52bbad5d8a83f059e6bec30ec1d07df874 Mon Sep 17 00:00:00 2001
From: Yaoyao Ding <yaoyaod@nvidia.com>
Date: Thu, 17 Sep 2026 12:02:56 -0400
Subject: [PATCH 5/5] [Fix] Resolve format-and-lint CI failures

- Apply ruff-format to the new/updated example scripts (examples/ uses
  line-length 90 via examples/pyproject.toml) and drop a stray blank line
  in stmt_builder.py.
- Fix mypy errors surfaced by the new code:
  - `RootInstructionGroup.get_thread_binding` needs the same
    `attr-defined` ignore used by the other threadIdx/blockIdx accessors.
  - `register_tensor(init=...)` was annotated `Callable[[Var, ...], ...]`,
    which is not valid typing syntax; mypy read it as a fixed 2-argument
    callable and rejected rank-1 init lambdas. It is variadic over the
    index vars, so use `Callable[..., ...]`.
  - `RegisterTensor.__getitem__`/`__setitem__` were annotated with the
    fixed-length `tuple[Expr | int | slice]`, rejecting multi-dimensional
    indexing. Use `tuple[..., ...]`, matching SharedTensor/GlobalTensor.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yaoyao Ding <yaoyaod@nvidia.com>
---
 examples/mhc/normw_merge.py                   | 12 +++-
 examples/moe/topk_gate.py                     | 40 +++++++----
 examples/quantization/per_channel_cast.py     | 66 +++++++++++++++----
 examples/quantization/per_token_cast.py       | 47 ++++++++++---
 .../swiglu_forward_and_per_token_cast.py      | 47 ++++++++++---
 python/tilus/ir/builders/stmt_builder.py      |  1 -
 python/tilus/ir/tensor.py                     |  6 +-
 python/tilus/lang/instructions/root.py        |  6 +-
 8 files changed, 172 insertions(+), 53 deletions(-)

diff --git a/examples/mhc/normw_merge.py b/examples/mhc/normw_merge.py
index cf4fe397..d4890954 100644
--- a/examples/mhc/normw_merge.py
+++ b/examples/mhc/normw_merge.py
@@ -20,7 +20,9 @@ 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):
+    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])
@@ -28,7 +30,9 @@ def __call__(self, m: int32, n: int32, fn_ptr: ~float32, normw_ptr: ~float32, ou
         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_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])
 
@@ -56,7 +60,9 @@ def run_tilekernels():
 
     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")
+    print(
+        f"mHC normw merge: Tilus {tilus_ms:.4f} ms, TileKernels {tilekernels_ms:.4f} ms"
+    )
 
 
 if __name__ == "__main__":
diff --git a/examples/moe/topk_gate.py b/examples/moe/topk_gate.py
index 270d54f0..fb4b545e 100644
--- a/examples/moe/topk_gate.py
+++ b/examples/moe/topk_gate.py
@@ -25,11 +25,17 @@ 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])
+        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])
+        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(
@@ -45,7 +51,9 @@ def __call__(self, num_tokens: int32, scores_ptr: ~float32, output_ptr: ~int64):
                 init=lambda _, j: j < self.num_experts,
             )
         else:
-            active = self.register_tensor(dtype=int32, shape=[1, self.aligned_experts], init=lambda _i, _j: 1)
+            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)
@@ -67,7 +75,9 @@ def __call__(self, num_tokens: int32, scores_ptr: ~float32, output_ptr: ~int64):
             # 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
+                    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
@@ -81,29 +91,37 @@ def __call__(self, num_tokens: int32, scores_ptr: ~float32, output_ptr: ~int64):
 
 def main():
     rows = []
-    for num_tokens, num_experts, num_topk in [(128, 72, 6), (1024, 256, 8), (8192, 256, 8)]:
+    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]
+        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
-            )
+            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]
+            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():
diff --git a/examples/quantization/per_channel_cast.py b/examples/quantization/per_channel_cast.py
index bf601ce5..cf52a966 100644
--- a/examples/quantization/per_channel_cast.py
+++ b/examples/quantization/per_channel_cast.py
@@ -12,17 +12,32 @@
 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:
+
+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)
+    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"
+    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):
@@ -32,20 +47,34 @@ 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):
+    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
+        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])
+        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()
@@ -53,11 +82,16 @@ def __call__(self, num_tokens: int32, hidden: int32, x_ptr: ~bfloat16, out_ptr:
         # 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))
+        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
+        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))
@@ -79,14 +113,20 @@ def main():
     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")
+    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__":
diff --git a/examples/quantization/per_token_cast.py b/examples/quantization/per_token_cast.py
index 1dbf1685..c5bfaa1d 100644
--- a/examples/quantization/per_token_cast.py
+++ b/examples/quantization/per_token_cast.py
@@ -46,38 +46,65 @@ def fp8_ordinal(x: torch.Tensor) -> torch.Tensor:
     return torch.where(bits & 0x80 != 0, -magnitude, magnitude)
 
 
-def check_fp8_close(actual: torch.Tensor, expected: torch.Tensor, *, label: str, max_mismatch_frac: float = 0.01) -> CodeLadderStats:
+def check_fp8_close(
+    actual: torch.Tensor,
+    expected: torch.Tensor,
+    *,
+    label: str,
+    max_mismatch_frac: float = 0.01,
+) -> CodeLadderStats:
     code_diff = (fp8_ordinal(actual) - fp8_ordinal(expected)).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"
+    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"
+    )
     return CodeLadderStats(max_code_diff, mismatch_frac)
 
 
-def check_scales_close(actual: torch.Tensor, expected: torch.Tensor, *, label: str, rtol: float = 1e-6) -> float:
+def check_scales_close(
+    actual: torch.Tensor, expected: torch.Tensor, *, label: str, rtol: float = 1e-6
+) -> float:
     torch.testing.assert_close(actual, expected, rtol=rtol, atol=0.0, msg=label)
     return float(((actual - expected).abs() / expected.abs()).max().item())
 
 
-def torch_per_token_cast(values: torch.Tensor, num_per_channels: int) -> tuple[torch.Tensor, torch.Tensor]:
+def torch_per_token_cast(
+    values: torch.Tensor, num_per_channels: int
+) -> tuple[torch.Tensor, torch.Tensor]:
     num_tokens, hidden = values.shape
     grouped = values.float().reshape(num_tokens, hidden // num_per_channels, -1)
     amax = grouped.abs().amax(dim=-1, keepdim=True).clamp_min(SF_CLAMP_MIN)
     scale = (amax / E4M3_MAX).squeeze(-1)
-    out = (grouped * (E4M3_MAX / amax)).clamp(-E4M3_MAX, E4M3_MAX).reshape(num_tokens, hidden).to(torch.float8_e4m3fn)
+    out = (
+        (grouped * (E4M3_MAX / amax))
+        .clamp(-E4M3_MAX, E4M3_MAX)
+        .reshape(num_tokens, hidden)
+        .to(torch.float8_e4m3fn)
+    )
     return out, scale
 
 
-def dequantize(out: torch.Tensor, scales: torch.Tensor, num_per_channels: int) -> torch.Tensor:
-    grouped = out.float().reshape(out.shape[0], out.shape[1] // num_per_channels, num_per_channels)
+def dequantize(
+    out: torch.Tensor, scales: torch.Tensor, num_per_channels: int
+) -> torch.Tensor:
+    grouped = out.float().reshape(
+        out.shape[0], out.shape[1] // num_per_channels, num_per_channels
+    )
     return (grouped * scales[:, :, None]).reshape(out.shape)
 
 
 def quantization_snr_db(reference: torch.Tensor, dequantized: torch.Tensor) -> float:
     noise = dequantized.float() - reference.float()
-    return float(10.0 * torch.log10(reference.float().square().sum() / noise.square().sum().clamp_min(1e-30)))
-
+    return float(
+        10.0
+        * torch.log10(
+            reference.float().square().sum() / noise.square().sum().clamp_min(1e-30)
+        )
+    )
 
 
 @tilus.autotune("block_m", [1, 2, 4])
diff --git a/examples/quantization/swiglu_forward_and_per_token_cast.py b/examples/quantization/swiglu_forward_and_per_token_cast.py
index e9a2d12a..121e9668 100644
--- a/examples/quantization/swiglu_forward_and_per_token_cast.py
+++ b/examples/quantization/swiglu_forward_and_per_token_cast.py
@@ -51,38 +51,65 @@ def fp8_ordinal(x: torch.Tensor) -> torch.Tensor:
     return torch.where(bits & 0x80 != 0, -magnitude, magnitude)
 
 
-def check_fp8_close(actual: torch.Tensor, expected: torch.Tensor, *, label: str, max_mismatch_frac: float = 0.01) -> CodeLadderStats:
+def check_fp8_close(
+    actual: torch.Tensor,
+    expected: torch.Tensor,
+    *,
+    label: str,
+    max_mismatch_frac: float = 0.01,
+) -> CodeLadderStats:
     code_diff = (fp8_ordinal(actual) - fp8_ordinal(expected)).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"
+    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"
+    )
     return CodeLadderStats(max_code_diff, mismatch_frac)
 
 
-def check_scales_close(actual: torch.Tensor, expected: torch.Tensor, *, label: str, rtol: float = 1e-6) -> float:
+def check_scales_close(
+    actual: torch.Tensor, expected: torch.Tensor, *, label: str, rtol: float = 1e-6
+) -> float:
     torch.testing.assert_close(actual, expected, rtol=rtol, atol=0.0, msg=label)
     return float(((actual - expected).abs() / expected.abs()).max().item())
 
 
-def torch_per_token_cast(values: torch.Tensor, num_per_channels: int) -> tuple[torch.Tensor, torch.Tensor]:
+def torch_per_token_cast(
+    values: torch.Tensor, num_per_channels: int
+) -> tuple[torch.Tensor, torch.Tensor]:
     num_tokens, hidden = values.shape
     grouped = values.float().reshape(num_tokens, hidden // num_per_channels, -1)
     amax = grouped.abs().amax(dim=-1, keepdim=True).clamp_min(SF_CLAMP_MIN)
     scale = (amax / E4M3_MAX).squeeze(-1)
-    out = (grouped * (E4M3_MAX / amax)).clamp(-E4M3_MAX, E4M3_MAX).reshape(num_tokens, hidden).to(torch.float8_e4m3fn)
+    out = (
+        (grouped * (E4M3_MAX / amax))
+        .clamp(-E4M3_MAX, E4M3_MAX)
+        .reshape(num_tokens, hidden)
+        .to(torch.float8_e4m3fn)
+    )
     return out, scale
 
 
-def dequantize(out: torch.Tensor, scales: torch.Tensor, num_per_channels: int) -> torch.Tensor:
-    grouped = out.float().reshape(out.shape[0], out.shape[1] // num_per_channels, num_per_channels)
+def dequantize(
+    out: torch.Tensor, scales: torch.Tensor, num_per_channels: int
+) -> torch.Tensor:
+    grouped = out.float().reshape(
+        out.shape[0], out.shape[1] // num_per_channels, num_per_channels
+    )
     return (grouped * scales[:, :, None]).reshape(out.shape)
 
 
 def quantization_snr_db(reference: torch.Tensor, dequantized: torch.Tensor) -> float:
     noise = dequantized.float() - reference.float()
-    return float(10.0 * torch.log10(reference.float().square().sum() / noise.square().sum().clamp_min(1e-30)))
-
+    return float(
+        10.0
+        * torch.log10(
+            reference.float().square().sum() / noise.square().sum().clamp_min(1e-30)
+        )
+    )
 
 
 @tilus.autotune("block_m", [1])
diff --git a/python/tilus/ir/builders/stmt_builder.py b/python/tilus/ir/builders/stmt_builder.py
index 40e48513..549e818d 100644
--- a/python/tilus/ir/builders/stmt_builder.py
+++ b/python/tilus/ir/builders/stmt_builder.py
@@ -1273,7 +1273,6 @@ def store_global(
         inst = StoreGlobalInst.create(dst=dst, x=src, offsets=[as_expr(ofs) for ofs in offsets], dims=dims)
         self.append(inst)
 
-
     def store_global_scatter(
         self,
         dst: GlobalTensor,
diff --git a/python/tilus/ir/tensor.py b/python/tilus/ir/tensor.py
index 2e33294b..f89d33c3 100644
--- a/python/tilus/ir/tensor.py
+++ b/python/tilus/ir/tensor.py
@@ -96,11 +96,13 @@ class RegisterTensor(Tensor):
     shape: tuple[int, ...]
     optional_layout: Optional[RegisterLayout] = None
 
-    def __getitem__(self, indices: tuple[Expr | int | slice] | Expr | int | slice) -> RegisterTensor:
+    def __getitem__(self, indices: tuple[Expr | int | slice, ...] | Expr | int | slice) -> RegisterTensor:
         raise RuntimeError("register_tensor[...] could only be used in Tilus Script.")
 
     def __setitem__(
-        self, indices: tuple[Expr | int | slice] | Expr | int | slice, value: RegisterTensor | Expr | int | float | None
+        self,
+        indices: tuple[Expr | int | slice, ...] | Expr | int | slice,
+        value: RegisterTensor | Expr | int | float | None,
     ) -> None:
         raise RuntimeError("register_tensor[...] = value could only be used in Tilus Script.")
 
diff --git a/python/tilus/lang/instructions/root.py b/python/tilus/lang/instructions/root.py
index be92411a..e17cb5d1 100644
--- a/python/tilus/lang/instructions/root.py
+++ b/python/tilus/lang/instructions/root.py
@@ -32,7 +32,7 @@
 class RootInstructionGroup(InstructionGroup):
     def get_thread_binding(self) -> Expr:
         """Return the physical CUDA thread index within the thread block."""
-        return threadIdx.x
+        return threadIdx.x  # type: ignore[attr-defined]
 
     @property
     def blockIdx(self) -> Dim3:
@@ -300,7 +300,7 @@ def register_tensor(
         *,
         dtype: DataType,
         shape: Sequence[int],
-        init: Optional[Callable[[Var, ...], Expr | int | float | bool] | Expr | int | float] = None,  # type: ignore [misc]
+        init: Optional[Callable[..., Expr | int | float | bool] | Expr | int | float] = None,
     ) -> RegisterTensor:
         """Create a register tensor.
 
@@ -320,7 +320,7 @@ def register_tensor(
             The data type of the tensor elements.
         shape: Sequence[int]
             The shape of the tensor.
-        init: Callable[[Var, ...], Expr | int | float | bool] | Expr | int | float, optional
+        init: Callable[..., Expr | int | float | bool] | Expr | int | float, optional
             The initialization value or function to initialize the tensor elements.
 
         Returns