diff --git a/benchmarks/_benchmark.py b/benchmarks/_benchmark.py
new file mode 100644
index 0000000..71ae179
--- /dev/null
+++ b/benchmarks/_benchmark.py
@@ -0,0 +1,47 @@
+"""Shared Triton benchmark and autotune reporting helpers."""
+
+from collections.abc import Mapping
+
+import torch
+from triton.testing import do_bench
+
+
+def benchmark_mean(function):
+    """Return mean latency in microseconds using Triton's benchmark policy."""
+    with torch.inference_mode():
+        return (
+            do_bench(
+                function,
+                warmup=25,
+                rep=100,
+                return_mode="mean",
+            )
+            * 1000.0
+        )
+
+
+def selected_config(kernel, fallback=None):
+    """Return the most recently cached ``(num_warps, num_stages)`` pair."""
+    # The SSA-first NineToothed compiler records the selected candidate on
+    # the public handle instead of exposing Triton's Autotuner object.
+    candidate = getattr(kernel, "_selected_tuning_candidate", None)
+    if candidate is not None:
+        if isinstance(candidate, Mapping):
+            return int(candidate["num_warps"]), int(candidate["num_stages"])
+        return int(candidate.num_warps), int(candidate.num_stages)
+
+    # Keep compatibility with the legacy frontend, which emits a decorated
+    # Triton kernel and stores the winner in the decorator's cache.
+    globals_ = getattr(getattr(kernel, "_kernel", None), "__globals__", {})
+    autotuners = (
+        value
+        for name, value in globals_.items()
+        if name.endswith("_with_auto_tuning") and hasattr(value, "cache")
+    )
+    autotuner = next(autotuners, None)
+    if autotuner is None or not autotuner.cache:
+        if fallback is not None:
+            return tuple(int(value) for value in fallback)
+        raise RuntimeError("the kernel has not completed autotuning")
+    config = next(reversed(autotuner.cache.values()))
+    return int(config.num_warps), int(config.num_stages)
diff --git a/benchmarks/bench_block_scaled_fp8_mm.py b/benchmarks/bench_block_scaled_fp8_mm.py
new file mode 100644
index 0000000..e857e5a
--- /dev/null
+++ b/benchmarks/bench_block_scaled_fp8_mm.py
@@ -0,0 +1,139 @@
+"""Benchmark BlockWise1x128 x BlockWise128x128 FP8 projection shapes."""
+
+import importlib
+
+import torch
+import torch.nn.functional as F
+from _benchmark import benchmark_mean, selected_config
+
+import ntops
+
+_block_scaled_fp8_mm_module = importlib.import_module("ntops.torch.block_scaled_fp8_mm")
+
+
+def _column_major(value):
+    return value.t().contiguous().t()
+
+
+def _make_inputs(m, n, k):
+    mat_a = torch.randn((m, k), device="cuda").clamp(-3, 3).to(torch.float8_e4m3fn)
+    weight = torch.randn((n, k), device="cuda").clamp(-3, 3).to(torch.float8_e4m3fn)
+    k_blocks = k // 128
+    padded_k_blocks = ((k_blocks + 3) // 4) * 4
+    scale_a = _column_major(
+        torch.ones((m, k_blocks), device="cuda", dtype=torch.float32)
+    )
+    scale_b = _column_major(
+        torch.ones(
+            (padded_k_blocks, n // 128),
+            device="cuda",
+            dtype=torch.float32,
+        )
+    )
+    return mat_a, weight.t(), scale_a, scale_b
+
+
+def _torch_dtype_reference(mat_a, mat_b, scale_a, scale_b):
+    """Dequantize through PyTorch's FP8 dtype conversion, then run FP32 MM."""
+    m, k = mat_a.shape
+    n = mat_b.shape[1]
+    k_blocks = k // 128
+
+    dequant_a = (
+        mat_a.float().reshape(m, k_blocks, 128) * scale_a.float().unsqueeze(-1)
+    ).reshape(m, k)
+    expanded_scale_b = scale_b[:k_blocks].float().repeat_interleave(128, dim=1)
+    dequant_b = (
+        mat_b.float().reshape(k_blocks, 128, n) * expanded_scale_b.unsqueeze(1)
+    ).reshape(k, n)
+    return (dequant_a @ dequant_b).to(torch.bfloat16)
+
+
+def _select_reference(mat_a, mat_b, scale_a, scale_b):
+    native = getattr(F, "scaled_mm", None)
+    if callable(native):
+
+        def run_native():
+            return native(
+                mat_a,
+                mat_b,
+                scale_a,
+                ntops.torch.ScalingType.BlockWise1x128,
+                scale_b,
+                ntops.torch.ScalingType.BlockWise128x128,
+            )
+
+        try:
+            return "torch.nn.functional.scaled_mm", run_native, run_native()
+        except (NotImplementedError, RuntimeError, TypeError, ValueError):
+            pass
+
+    def run_dtype_reference():
+        return _torch_dtype_reference(mat_a, mat_b, scale_a, scale_b)
+
+    return "torch_dtype_dequant_mm", run_dtype_reference, run_dtype_reference()
+
+
+def main():
+    if not torch.cuda.is_available():
+        raise RuntimeError("a CUDA-compatible accelerator is required")
+
+    print(f"torch={torch.__version__}")
+    print(f"hip={torch.version.hip}")
+    print(f"device={torch.cuda.get_device_name()}")
+    print(
+        "scenario,M,N,K,reference_provider,reference_mean_us,ntops_mean_us,"
+        "best_num_warps,best_num_stages,speedup_vs_reference,tflops,"
+        "max_abs_error"
+    )
+    scenarios = (
+        ("attention_decode", 1, 4096, 4096),
+        ("moe_expert", 32, 14336, 4096),
+        ("linear_prefill", 128, 4096, 4096),
+    )
+    for name, m, n, k in scenarios:
+        mat_a, mat_b, scale_a, scale_b = _make_inputs(m, n, k)
+
+        def run_ntops():
+            return ntops.torch.block_scaled_fp8_mm(
+                mat_a,
+                mat_b,
+                scale_a,
+                ntops.torch.ScalingType.BlockWise1x128,
+                scale_b,
+                ntops.torch.ScalingType.BlockWise128x128,
+            )
+
+        with torch.inference_mode():
+            output = run_ntops()
+            reference_provider, run_reference, reference = _select_reference(
+                mat_a, mat_b, scale_a, scale_b
+            )
+            torch.testing.assert_close(output, reference, rtol=0.03, atol=0.03)
+            max_abs_error = (output.float() - reference.float()).abs().max().item()
+            torch.cuda.synchronize()
+
+        kernel = _block_scaled_fp8_mm_module._make_kernel(
+            mat_a.dtype,
+            torch.bfloat16,
+            None,
+            mat_a.device,
+        )
+        _, fixed_warps, fixed_stages, _ = (
+            _block_scaled_fp8_mm_module._kernel_tuning_config(mat_a.device)
+        )
+        fallback_warps = fixed_warps if isinstance(fixed_warps, int) else fixed_warps[0]
+        warps, stages = selected_config(kernel, fallback=(fallback_warps, fixed_stages))
+
+        reference_us = benchmark_mean(run_reference)
+        ntops_us = benchmark_mean(run_ntops)
+        tflops = 2.0 * m * n * k / (ntops_us * 1.0e6)
+        print(
+            f"{name},{m},{n},{k},{reference_provider},{reference_us:.3f},"
+            f"{ntops_us:.3f},{warps},{stages},{reference_us / ntops_us:.3f},"
+            f"{tflops:.3f},{max_abs_error:.6f}"
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/bench_fused_mla_rope_cache_write.py b/benchmarks/bench_fused_mla_rope_cache_write.py
new file mode 100644
index 0000000..71da821
--- /dev/null
+++ b/benchmarks/bench_fused_mla_rope_cache_write.py
@@ -0,0 +1,151 @@
+"""Benchmark MLA RoPE + compressed KV-cache write fusion."""
+
+import torch
+from _benchmark import benchmark_mean, selected_config
+
+import ntops
+from ntops.torch.utils import _cached_make
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+NUM_HEADS = 128  # documented MLA geometry; cache writer itself is head-shared
+KV_LORA_RANK = 512
+ROPE_DIM = 64
+CACHE_BLOCK_SIZE = 16
+TILE_SIZE = 128
+AUTOTUNE_WARPS = (1, 2, 4, 8)
+AUTOTUNE_STAGES = (1, 2)
+MAX_NUM_CONFIGS = 8
+
+
+def _make_inputs(tokens):
+    kv_c = torch.randn(tokens, KV_LORA_RANK, device=DEVICE, dtype=DTYPE)
+    k_pe = torch.randn(tokens, ROPE_DIM, device=DEVICE, dtype=DTYPE)
+    num_blocks = (tokens + CACHE_BLOCK_SIZE - 1) // CACHE_BLOCK_SIZE + 4
+    kv_cache = torch.empty(
+        num_blocks,
+        CACHE_BLOCK_SIZE,
+        KV_LORA_RANK + ROPE_DIM,
+        device=DEVICE,
+        dtype=DTYPE,
+    )
+    slot_mapping = torch.arange(tokens, device=DEVICE, dtype=torch.int64)
+    positions = torch.arange(tokens, device=DEVICE, dtype=torch.int64)
+    half = ROPE_DIM // 2
+    theta = 10000 ** (-2 * torch.arange(half, device=DEVICE) / ROPE_DIM)
+    phase = positions.to(torch.float32)[:, None] * theta[None, :]
+    cos_sin_cache = torch.cat((phase.cos(), phase.sin()), dim=-1)
+    return kv_c, k_pe, kv_cache, slot_mapping, positions, cos_sin_cache
+
+
+def _torch_rope(k_pe, positions, cos_sin_cache):
+    half = ROPE_DIM // 2
+    table = cos_sin_cache.index_select(0, positions)
+    cos = table[:, :half].to(torch.float32)
+    sin = table[:, half:].to(torch.float32)
+    x0 = k_pe[:, 0::2].to(torch.float32)
+    x1 = k_pe[:, 1::2].to(torch.float32)
+    output = torch.empty_like(k_pe, dtype=torch.float32)
+    output[:, 0::2] = x0 * cos - x1 * sin
+    output[:, 1::2] = x0 * sin + x1 * cos
+    return output.to(k_pe.dtype)
+
+
+def _torch_unfused(inputs):
+    kv_c, k_pe, cache, slots, positions, table = inputs
+    k_rot = _torch_rope(k_pe, positions, table)
+    block = slots // CACHE_BLOCK_SIZE
+    offset = slots % CACHE_BLOCK_SIZE
+    cache[block, offset, :KV_LORA_RANK] = kv_c
+    cache[block, offset, KV_LORA_RANK:] = k_rot
+
+
+def _torch_rope_only(inputs):
+    _, k_pe, _, slots, positions, table = inputs
+    del slots
+    return _torch_rope(k_pe, positions, table)
+
+
+def _torch_cache_only(inputs, k_rot):
+    kv_c, _, cache, slots, _, _ = inputs
+    block = slots // CACHE_BLOCK_SIZE
+    offset = slots % CACHE_BLOCK_SIZE
+    cache[block, offset, :KV_LORA_RANK] = kv_c
+    cache[block, offset, KV_LORA_RANK:] = k_rot
+
+
+def _kernel_handle(inputs):
+    entry_dim = KV_LORA_RANK + ROPE_DIM
+    tile_size = max(
+        1 << (TILE_SIZE - 1).bit_length(),
+        1 << (entry_dim - 1).bit_length(),
+    )
+    return _cached_make(
+        ntops.kernels.fused_mla_rope_cache_write.premake,
+        KV_LORA_RANK,
+        ROPE_DIM,
+        dtype=DTYPE,
+        block_size=tile_size,
+        cache_block_size=CACHE_BLOCK_SIZE,
+        cos_dtype=inputs[-1].dtype,
+        num_warps=AUTOTUNE_WARPS,
+        num_stages=AUTOTUNE_STAGES,
+        max_num_configs=MAX_NUM_CONFIGS,
+    )
+
+
+def main():
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for this benchmark")
+    print(f"torch={torch.__version__}")
+    print(f"device={torch.cuda.get_device_name()}")
+    print(
+        "scenario,tokens,heads,latent,rope,torch_rope_mean_us,"
+        "torch_cache_mean_us,torch_unfused_mean_us,best_num_warps,"
+        "best_num_stages,autotuned_mean_us,"
+        "speedup_vs_pytorch"
+    )
+
+    for scenario, tokens in (
+        ("decode", 1),
+        ("concurrent_decode", 10),
+        ("long_context", 2048),
+    ):
+        inputs = _make_inputs(tokens)
+
+        def rope_only():
+            return _torch_rope_only(inputs)
+
+        rotated = _torch_rope_only(inputs)
+
+        def cache_only():
+            return _torch_cache_only(inputs, rotated)
+
+        def unfused():
+            return _torch_unfused(inputs)
+
+        def autotuned():
+            return ntops.torch.fused_mla_rope_cache_write(
+                *inputs,
+                block_size=TILE_SIZE,
+                num_warps=AUTOTUNE_WARPS,
+                num_stages=AUTOTUNE_STAGES,
+                max_num_configs=MAX_NUM_CONFIGS,
+            )
+
+        autotuned()
+        torch.cuda.synchronize()
+        warps, stages = selected_config(_kernel_handle(inputs))
+        rope_us = benchmark_mean(rope_only)
+        cache_us = benchmark_mean(cache_only)
+        unfused_us = benchmark_mean(unfused)
+        auto_us = benchmark_mean(autotuned)
+        print(
+            f"{scenario},{tokens},{NUM_HEADS},{KV_LORA_RANK},{ROPE_DIM},"
+            f"{rope_us:.3f},{cache_us:.3f},{unfused_us:.3f},{warps},{stages},"
+            f"{auto_us:.3f},{unfused_us / auto_us:.3f}"
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/bench_mxfp4_w4a16_grouped_mm.py b/benchmarks/bench_mxfp4_w4a16_grouped_mm.py
new file mode 100644
index 0000000..75f6013
--- /dev/null
+++ b/benchmarks/bench_mxfp4_w4a16_grouped_mm.py
@@ -0,0 +1,233 @@
+"""Benchmark W4A16 MXFP4 scaled grouped matrix multiplication."""
+
+import importlib
+
+import torch
+import torch.nn.functional as F
+from _benchmark import benchmark_mean, selected_config
+
+import ntops
+
+_mxfp4_w4a16_grouped_mm_module = importlib.import_module(
+    "ntops.torch.mxfp4_w4a16_grouped_mm"
+)
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+GROUP_COUNT = 8
+K = 4096
+N = 4096
+
+SCENARIOS = (
+    ("uniform_decode", 1),
+    ("uniform_concurrent", 16),
+    ("routed_uneven", (1, 0, 3, 8, 16, 24, 32, 44)),
+)
+
+
+def _native_packed_dtypes():
+    return (
+        getattr(torch, "float4_e2m1fn_x2", None),
+        getattr(torch, "float8_e8m0fnu", None),
+    )
+
+
+def _decode_nibble(code):
+    magnitude_code = code & 0x7
+    exponent = (magnitude_code >> 1).to(torch.int32)
+    mantissa = (magnitude_code & 1).to(torch.float32)
+    normal = (1.0 + 0.5 * mantissa) * torch.exp2(exponent.to(torch.float32) - 1.0)
+    magnitude = torch.where(exponent == 0, 0.5 * mantissa, normal)
+    sign = torch.where((code & 0x8) == 0, 1.0, -1.0)
+    return sign * magnitude
+
+
+def _manual_decode_mxfp4(packed, scales):
+    if packed.dtype != torch.uint8:
+        packed = packed.view(torch.uint8)
+    if scales.dtype != torch.uint8:
+        scales = scales.view(torch.uint8)
+
+    group_count, packed_k, n = packed.shape
+    weight = torch.empty(
+        group_count,
+        packed_k * 2,
+        n,
+        dtype=torch.bfloat16,
+        device=packed.device,
+    )
+    block_scales = torch.exp2(scales.to(torch.float32) - 127.0)
+    block_scales = block_scales.repeat_interleave(32, dim=1)
+    weight[:, 0::2] = (_decode_nibble(packed & 0xF) * block_scales[:, 0::2]).to(
+        torch.bfloat16
+    )
+    weight[:, 1::2] = (_decode_nibble(packed >> 4) * block_scales[:, 1::2]).to(
+        torch.bfloat16
+    )
+    return weight
+
+
+def _make_inputs(rows):
+    packed = torch.randint(
+        0,
+        256,
+        (GROUP_COUNT, K // 2, N),
+        dtype=torch.uint8,
+        device=DEVICE,
+    )
+    scales = torch.randint(
+        124,
+        128,
+        (GROUP_COUNT, K // 32, N),
+        dtype=torch.uint8,
+        device=DEVICE,
+    )
+    packed_dtype, scale_dtype = _native_packed_dtypes()
+    native_dtypes = packed_dtype is not None and scale_dtype is not None
+    mat_b = packed.view(packed_dtype) if native_dtypes else packed
+    scale_b = scales.view(scale_dtype) if native_dtypes else scales
+
+    if isinstance(rows, int):
+        mat_a = torch.randn(
+            GROUP_COUNT, rows, K, dtype=DTYPE, device=DEVICE
+        ).contiguous()
+        return mat_a, mat_b, scale_b, None, None, native_dtypes
+
+    row_ends = tuple(torch.tensor(rows).cumsum(0).tolist())
+    mat_a = torch.randn(sum(rows), K, dtype=DTYPE, device=DEVICE).contiguous()
+    offs = torch.tensor(row_ends, dtype=torch.int32, device=DEVICE)
+    return mat_a, mat_b, scale_b, offs, row_ends, native_dtypes
+
+
+def _torch_matmul(mat_a, weight, row_ends):
+    if row_ends is None:
+        return torch.bmm(mat_a.float(), weight.float()).to(torch.bfloat16)
+
+    outputs = []
+    start = 0
+    for expert, end in enumerate(row_ends):
+        outputs.append(
+            (mat_a[start:end].float() @ weight[expert].float()).to(torch.bfloat16)
+        )
+        start = end
+    return torch.cat(outputs)
+
+
+def _ntops_mxfp4_w4a16_grouped_mm(mat_a, mat_b, scale_b, offs):
+    def run():
+        return ntops.torch.mxfp4_w4a16_grouped_mm(
+            mat_a,
+            mat_b,
+            None,
+            None,
+            scale_b,
+            ntops.torch.ScalingType.BlockWise1x32,
+            offs=offs,
+        )
+
+    if offs is None:
+        return run()
+
+    # PyTorch's jagged nested-tensor constructor requires a version counter.
+    with torch.inference_mode(False):
+        return run()
+
+
+def _select_reference(mat_a, mat_b, scale_b, offs, row_ends, native_dtypes):
+    native = getattr(F, "scaled_grouped_mm", None)
+    fallback_reason = "native_dtype_unavailable"
+    if native_dtypes and callable(native):
+
+        def run_native():
+            return native(
+                mat_a,
+                mat_b,
+                None,
+                None,
+                scale_b,
+                ntops.torch.ScalingType.BlockWise1x32,
+                offs=offs,
+            )
+
+        try:
+            return "torch.nn.functional.scaled_grouped_mm", run_native, run_native()
+        except (NotImplementedError, RuntimeError, TypeError, ValueError):
+            fallback_reason = "native_operator_unavailable"
+    elif native_dtypes:
+        fallback_reason = "native_operator_unavailable"
+
+    def run_manual_reference():
+        weight = _manual_decode_mxfp4(mat_b, scale_b)
+        return _torch_matmul(mat_a, weight, row_ends)
+
+    provider = f"manual_dequant_mm_{fallback_reason}"
+    return provider, run_manual_reference, run_manual_reference()
+
+
+def main():
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA-compatible accelerator is required for this benchmark")
+
+    torch.manual_seed(0)
+    backend = "hip_reduction" if torch.version.hip is not None else "block_dot"
+    print(f"torch={torch.__version__}")
+    print(f"device={torch.cuda.get_device_name()}")
+    print(
+        "scenario,mode,groups,total_rows,max_rows,k,n,backend,input_storage,"
+        "reference_provider,reference_mean_us,ntops_mean_us,"
+        "best_num_warps,best_num_stages,speedup_vs_reference,effective_tflops,"
+        "max_abs_error"
+    )
+
+    for scenario, rows in SCENARIOS:
+        mat_a, mat_b, scale_b, offs, row_ends, native_dtypes = _make_inputs(rows)
+
+        def run_ntops():
+            return _ntops_mxfp4_w4a16_grouped_mm(mat_a, mat_b, scale_b, offs)
+
+        with torch.inference_mode():
+            output = run_ntops()
+            reference_provider, run_reference, reference = _select_reference(
+                mat_a, mat_b, scale_b, offs, row_ends, native_dtypes
+            )
+            if reference is not None:
+                torch.testing.assert_close(output, reference, rtol=0.03, atol=0.03)
+                max_abs_error = (output.float() - reference.float()).abs().max().item()
+            else:
+                max_abs_error = float("nan")
+            torch.cuda.synchronize()
+
+        kernel = _mxfp4_w4a16_grouped_mm_module._make_kernel(offs is not None)
+        _, _, _, fixed_warps, fixed_stages, _ = (
+            _mxfp4_w4a16_grouped_mm_module._kernel_launch_config(offs is not None)
+        )
+        fallback_warps = fixed_warps if isinstance(fixed_warps, int) else fixed_warps[0]
+        warps, stages = selected_config(kernel, fallback=(fallback_warps, fixed_stages))
+
+        reference_us = (
+            benchmark_mean(run_reference) if run_reference is not None else float("nan")
+        )
+        ntops_us = benchmark_mean(run_ntops)
+
+        if isinstance(rows, int):
+            mode = "uniform"
+            total_rows = GROUP_COUNT * rows
+            max_rows = rows
+        else:
+            mode = "routed"
+            total_rows = sum(rows)
+            max_rows = max(rows)
+        input_storage = "native_packed" if native_dtypes else "raw_uint8"
+        speedup = reference_us / ntops_us
+        effective_tflops = 2 * total_rows * K * N / (ntops_us * 1e6)
+
+        print(
+            f"{scenario},{mode},{GROUP_COUNT},{total_rows},{max_rows},{K},{N},"
+            f"{backend},{input_storage},{reference_provider},{reference_us:.3f},"
+            f"{ntops_us:.3f},{warps},{stages},{speedup:.3f},"
+            f"{effective_tflops:.3f},{max_abs_error:.6f}"
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/bench_rms_norm_gated.py b/benchmarks/bench_rms_norm_gated.py
new file mode 100644
index 0000000..50b3d21
--- /dev/null
+++ b/benchmarks/bench_rms_norm_gated.py
@@ -0,0 +1,103 @@
+"""Benchmark the autotuned RMSNormGated implementation."""
+
+import torch
+import torch.nn.functional as F
+from _benchmark import benchmark_mean, selected_config
+
+import ntops
+from ntops.torch.utils import _cached_make
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+HIDDEN_SIZE = 128
+LOCAL_VALUE_HEADS = 8  # Qwen3-Next: 32 value heads with TP=4.
+EPS = 1e-5
+BLOCK_SIZE = 128
+AUTOTUNE_WARPS = (1, 2, 4, 8)
+AUTOTUNE_STAGES = (1, 2)
+MAX_NUM_CONFIGS = 8
+
+
+def pytorch_reference(input, z, weight):
+    x = input.float()
+    gate = F.silu(z.float())
+    output = x * torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + EPS)
+    return (output * weight.float() * gate).to(input.dtype)
+
+
+def _kernel_handle(input, z, weight):
+    return _cached_make(
+        ntops.kernels.rms_norm_gated.premake,
+        input.ndim,
+        HIDDEN_SIZE,
+        None,
+        True,
+        "silu",
+        input_dtype=input.dtype,
+        gate_dtype=z.dtype,
+        weight_dtype=weight.dtype,
+        output_dtype=input.dtype,
+        block_size=BLOCK_SIZE,
+        has_gate=True,
+        num_warps=AUTOTUNE_WARPS,
+        num_stages=AUTOTUNE_STAGES,
+        max_num_configs=MAX_NUM_CONFIGS,
+    )
+
+
+def main():
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for this benchmark")
+
+    print(f"torch={torch.__version__}")
+    print(f"device={torch.cuda.get_device_name()}")
+    print(
+        "scenario,tokens,rows,hidden,dtype,best_num_warps,best_num_stages,"
+        "autotuned_mean_us,pytorch_mean_us,"
+        "speedup_vs_pytorch,max_abs_error"
+    )
+
+    scenarios = (
+        ("decode", 1),
+        ("concurrent_decode", 10),
+        ("prefill_2048", 2048),
+    )
+    for name, num_tokens in scenarios:
+        rows = num_tokens * LOCAL_VALUE_HEADS
+        input = torch.randn((rows, HIDDEN_SIZE), device=DEVICE, dtype=DTYPE)
+        z = torch.randn_like(input)
+        weight = torch.randn(HIDDEN_SIZE, device=DEVICE, dtype=torch.float32)
+
+        def autotuned_function():
+            return ntops.torch.rms_norm_gated(
+                input,
+                z,
+                weight,
+                eps=EPS,
+                norm_before_gate=True,
+                activation="silu",
+            )
+
+        def pytorch_function():
+            return pytorch_reference(input, z, weight)
+
+        with torch.inference_mode():
+            autotuned_output = autotuned_function()
+            pytorch_output = pytorch_function()
+            torch.cuda.synchronize()
+            max_abs_error = (
+                (autotuned_output.float() - pytorch_output.float()).abs().max().item()
+            )
+
+        warps, stages = selected_config(_kernel_handle(input, z, weight))
+        autotuned_mean = benchmark_mean(autotuned_function)
+        pytorch_mean = benchmark_mean(pytorch_function)
+        print(
+            f"{name},{num_tokens},{rows},{HIDDEN_SIZE},{DTYPE},"
+            f"{warps},{stages},{autotuned_mean:.3f},{pytorch_mean:.3f},"
+            f"{pytorch_mean / autotuned_mean:.3f},{max_abs_error:.6f}"
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/docs/build_and_evaluate.md b/docs/build_and_evaluate.md
new file mode 100644
index 0000000..b6fe9ee
--- /dev/null
+++ b/docs/build_and_evaluate.md
@@ -0,0 +1,87 @@
+# 一键构建与评测
+
+## 1. 环境要求
+
+- Linux、Python 3.10 或更高版本；
+- 已适配目标设备的 PyTorch，以及可用的 CUDA、HIP 或 CoreX runtime；
+- `ninetoothed>=0.16.0`；
+- 海光 DCU 或天数智芯 MR-V100。仅执行构建时不要求加速卡。
+
+## 2. 一键执行
+
+构建、定向正确性测试和四个算子 benchmark 使用同一个入口：
+
+```bash
+./scripts/build_and_evaluate.sh all
+```
+
+脚本默认执行 editable install，打印 Python、PyTorch、CUDA/HIP、NineToothed 实际导入路径和设备信息，然后运行测试及 benchmark。日志写入 UTC 时间戳目录：
+
+```text
+artifacts/YYYYmmddTHHMMSSZ/
+  build.log
+  environment.txt
+  pytest.log
+  bench_*.csv
+```
+
+### 2.1 Build 实际命令
+
+脚本的 build 步骤执行：
+
+```bash
+python -m pip install -e ".[testing]"
+```
+
+这一步安装当前源码及测试依赖，不执行 NineToothed kernel 的 AOT 编译。kernel 在测试或 benchmark 第一次调用对应算子时 JIT 编译；启用多个候选时，自动调优也在该阶段完成。`build` 模式在安装并记录环境信息后退出：
+
+```bash
+./scripts/build_and_evaluate.sh build
+```
+
+可按阶段运行：
+
+```bash
+./scripts/build_and_evaluate.sh build
+./scripts/build_and_evaluate.sh test
+./scripts/build_and_evaluate.sh benchmark
+```
+
+已经安装依赖时可跳过安装；也可指定解释器和结果目录：
+
+```bash
+SKIP_INSTALL=1 PYTHON_BIN=/path/to/python \
+RESULT_ROOT=/path/to/results ./scripts/build_and_evaluate.sh all
+```
+
+## 3. 手工等价命令
+
+```bash
+python -m pip install -e ".[testing]"
+PYTHONPATH=src python -m pytest -q \
+  tests/test_block_scaled_fp8_mm.py \
+  tests/test_mxfp4_w4a16_grouped_mm.py \
+  tests/test_rms_norm_gated.py \
+  tests/test_fused_mla_rope_cache_write.py
+```
+
+```bash
+PYTHONPATH=src python benchmarks/bench_block_scaled_fp8_mm.py
+PYTHONPATH=src python benchmarks/bench_mxfp4_w4a16_grouped_mm.py
+PYTHONPATH=src python benchmarks/bench_rms_norm_gated.py
+PYTHONPATH=src python benchmarks/bench_fused_mla_rope_cache_write.py
+```
+
+benchmark 使用 `triton.testing.do_bench(warmup=25, rep=100, return_mode="mean")`，结果单位为微秒。首次编译与自动调优应在正式计时前完成。量化算子的 CSV 中必须同时保留 `reference_provider` 和输入存储类型；软件反量化 reference 的加速比不能解释为相对平台原生量化算子的收益。
+
+## 4. 两平台复验
+
+海光和天数智芯必须分别在同一提交上执行 `all`，不得共用另一平台的调优配置或性能结论。归档时保留：
+
+- `git rev-parse HEAD` 与工作区是否干净；
+- 脚本生成的环境、测试和 benchmark 日志；
+- 设备型号、驱动/runtime 版本；
+- 量化 benchmark 实际采用的 reference provider；
+- 自动调优选出的 `num_warps`、`num_stages`。
+
+HIP 原生编译器错误可能结束整个 Python 进程。引入新的 tile、dot 或多 wave 候选前，应在独立进程逐个编译和校验，不能直接加入进程内自动调优集合。
diff --git a/docs/technical_report.md b/docs/technical_report.md
new file mode 100644
index 0000000..efd542d
--- /dev/null
+++ b/docs/technical_report.md
@@ -0,0 +1,276 @@
+# KernelSwift Competition 技术报告
+
+## 1. 修改内容与动机
+
+本分支交付四个面向 LLM 推理的 NineToothed 算子。所有改动都属于 A 部分的算子实现与优化；B 部分没有修改，没有新增编译 pass，也没有修改 NineToothed、Triton、CoreX 或 HIP 后端源码。
+
+| 范围 | 正式入口 | 目标问题 | 主要场景 |
+| --- | --- | --- | --- |
+| A 部分 | `block_scaled_fp8_mm` | 1x128/128x128 block scale 的 FP8 GEMM | Attention、线性层、单专家 MoE |
+| A 部分 | `mxfp4_w4a16_grouped_mm` | MXFP4 权重解码、缩放和 grouped GEMM 融合 | uniform/routed MoE |
+| A 部分 | `rms_norm_gated` | RMSNorm、门控与激活融合 | decode、prefill |
+| A 部分 | `fused_mla_rope_cache_write` | MLA RoPE 与压缩 paged KV cache 写入融合 | MLA decode、长上下文 |
+| B 部分 | 无 | 不修改编译器和平台后端 | 不适用 |
+
+## 2. 方案与实现
+
+### 2.1 公共实现原则
+
+四个算子先冻结 shape、dtype、layout、原地语义和拒绝路径，再建立独立 PyTorch reference。kernel 只使用 `source[...]`、masked load/store、基础算术、公开 cast、`ntl.sum` 和普通 `ntl.dot` 等公共 lowering 能力。
+
+平台差异在 PyTorch wrapper 或 premake 参数处显式分派。通用数学语义和公开 API 保持一致，平台特化只改变安全的 tile、dot/reduction 路径与调优集合。修改没有侵入 NineToothed 或 Triton 安装目录，也没有进程级 patch。
+
+### 2.2 Block-scaled FP8 MM
+
+输入为 `A[M,K] @ B[K,N]`。A scale 使用 `BlockWise1x128`；B scale 支持 `BlockWise1x128` 和 `BlockWise128x128`。A 是 row-major FP8，B 是由权重转置得到的 column-major dense FP8；scale 同时接受 row-major、column-major 及 PyTorch 使用的 K-block padding。
+
+每个 128-wide K block 的 FP32 contribution 为：
+
+```text
+acc[m,n] += dot(A[m,p], B[p,n]) * scale_a[m,p] * scale_b[p,n]
+```
+
+HIP 和支持 FP8 Tensor Core 的 NVIDIA CUDA 使用普通 FP8 `ntl.dot`。CoreX 与不支持原生 FP8 dot 的 CUDA 路径将有限 E4M3FN 值精确转换为 BF16，再执行 BF16 dot；这是一条正确性回退，不改变公开 FP8 输入。HIP 固定 `BLOCK_M=16, BLOCK_N=16, BLOCK_K=32, num_warps=1, num_stages=1`，避免 gfx936 在多 wave FP8 specialization 上发生 AMD LLVM 进程级崩溃。
+
+权重按 128-N block 建立无分配 view；A 与 A scale 使用 stride-0 view 共享，wrapper 不物化完整反量化矩阵。bias 在 FP32 accumulator 上融合，最后一次性转换到 FP16、BF16 或 FP32。
+
+### 2.3 MXFP4 W4A16 Grouped MM
+
+支持两种布局：
+
+| 模式 | activation | packed weight | scale | output |
+| --- | --- | --- | --- | --- |
+| uniform | `[G,M,K]` | `[G,K/2,N]` | `[G,K/32,N]` | `[G,M,N]` |
+| routed | `[total_M,K]` | `[G,K/2,N]` | `[G,K/32,N]` | `[total_M,N]` |
+
+每 byte 的低、高 nibble 分别保存偶数和奇数 K 的 E2M1 code；E8M0 byte 解码为 `2^(s-127)`。每个 32-wide scale block 拆成两组 16-wide K。CoreX 使用两次 BF16 dot，HIP 使用无分支整数解码加显式 FP32 reduction，避免 gfx936 的 BF16 MMAC codegen 崩溃。两条路径都不物化 `[G,K,N]` 权重。
+
+uniform HIP 在已验证安全的 1/4-wave 中有界调优；routed HIP 固定单 wave。`offs[g]` 保持运行时值，防止按最大序列长度 specialization 后由 padding program 覆盖后续专家输出。零 token 专家由相邻相同 offset 表达。
+
+### 2.4 Gated RMSNorm
+
+数学定义为：
+
+```text
+RMSNorm(x) = x * rsqrt(mean(x^2) + eps) * weight
+
+norm_before_gate=True:  y = RMSNorm(x) * activation(z)
+norm_before_gate=False: y = RMSNorm(x * activation(z))
+```
+
+平方和、激活和归一化因子使用 FP32，最后转换回输入 dtype。一个 program 负责一行或一个完整 normalization group，归约轴不跨 program。不同门控顺序、激活、分组和可选输入在 premake 阶段选择 application，热路径没有对应的运行时分支。
+
+默认搜索 `num_warps=(1,2,4,8)`、`num_stages=(1,2)`，winner 按 shape、dtype、stride 和静态语义缓存。若新增 HIP 候选，必须先在独立进程验证编译稳定性。
+
+### 2.5 Fused MLA RoPE Cache Write
+
+接口为：
+
+```python
+ntops.torch.fused_mla_rope_cache_write(
+    kv_c,             # [T_source, L]
+    k_pe,             # [T_source, R] or [T_source, 1, R]
+    kv_cache,         # [num_blocks, cache_block_size, L + R]
+    slot_mapping,     # [T] int64; -1 skips
+    positions,        # [T] int32/int64
+    cos_sin_cache,    # [max_position, R], packed [cos | sin]
+    block_size=128,
+    num_warps=(1, 2, 4, 8),
+    num_stages=(1, 2),
+    max_num_configs=8,
+)
+```
+
+`kv_c` 是共享低秩 K/V latent，不按 head 展开。典型 `L=512, R=64` 时，每 token 只写 576 个元素。cache 地址为：
+
+```text
+block_idx = slot // cache_block_size
+offset    = slot % cache_block_size
+entry     = [kv_c | RoPE(k_pe)]
+```
+
+一个 program 负责一个 `[token, entry tile]`。token/feature 坐标来自 `output_anchor.offsets()`，latent、RoPE table 与 paged cache 通过 `source[...]` 间接访问。tile 至少覆盖完整 entry 的下一个 2 的幂，避免长上下文产生多余 program。
+
+NineToothed SSA 图需要真实输出根，部分 DCU backend 不能可靠 lower 只有间接副作用的 launch。当前唯一 kernel 因此生成 `[T,L+R]` 的压缩 entry 输出锚定，同时在同一 application 中写 paged cache；wrapper 丢弃锚定输出。它不生成 query 或按 head 展开，比借用其他融合语义更直接，也使 CoreX/HIP 共用同一 kernel。非连续 source 先连续化，非连续 cache 通过临时连续 cache 执行后 copy back；连续生产输入没有该额外复制。
+
+`slot=-1` 被映射到负 block 并由 store mask 抑制。`T_source` 可以大于有效 `T`，以适配 CUDA Graph padding。FP8 cache 与 `fp8_ds_mla` 量化布局尚未实现，明确拒绝而不是静默套用未量化语义。
+
+### 2.6 编译 pass、后端与官方 NineToothed 对比
+
+本分支没有 B 部分改动，也没有自定义 pass。因此不存在“修改后编译后端相对官方 NineToothed 后端”的独立收益数据；所有表格均是官方 NineToothed 安装编译本仓库算子后，相对 PyTorch reference 的端到端结果。实际收益来自 program 映射、融合、在线解码、无分配 view 和平台安全分派，不应归因于后端修改。
+
+## 3. 实验环境与方法
+
+### 3.1 环境
+
+| 平台 | 设备 | PyTorch | Triton/runtime |
+| --- | --- | --- | --- |
+| 天数智芯 | Iluvatar MR-V100 | 2.7.1+corex.4.4.0 | Triton 3.1.0+corex.4.4.0、CUDA compatibility 10.2 |
+| 海光 | BW/gfx936 | 2.9.0+das.opt1.dtk2604 | Triton 3.5.1+das.opt1.dtk2604.torch290、HIP 6.3.26093 |
+
+benchmark 使用 `do_bench(warmup=25, rep=100, return_mode="mean")`，正式采样前完成首次编译与自动调优，结果为平均微秒。量化算子优先使用 PyTorch 原生 dtype/operator；缺少原生 operator 时回退到软件反量化 reference，并在结果中单独标识。
+
+### 3.2 量化 dtype、输入与 reference
+
+两个平台的 PyTorch dtype 和原生 operator 支持情况如下：
+
+| 平台 | Block FP8 dtype | `torch.nn.functional.scaled_mm` | MXFP4 dtype | `torch.nn.functional.scaled_grouped_mm` |
+| --- | --- | --- | --- | --- |
+| 天数 MR-V100 | `float8_e4m3fn` 可用 | 不可用 | `float8_e8m0fnu` 可用，`float4_e2m1fn_x2` 不可用 | 不可用 |
+| 海光 BW | `float8_e4m3fn` 可用 | 没有可运行实现 | packed E2M1 与 E8M0 dtype 可用 | 没有可运行实现 |
+
+dtype 转换必须区分数值转换与位模式解释：
+
+- `.to(torch.float8_e4m3fn)` 是数值转换，会把浮点输入量化并舍入到 E4M3FN；
+- FP8 tensor 的 `.float()` 解码 E4M3FN 自身表示的数值，但外部 block scale 仍需
+  单独相乘，二者共同构成完整反量化；
+- `.view(torch.float4_e2m1fn_x2)`、`.view(torch.float8_e8m0fnu)` 或
+  `.view(torch.uint8)` 只重新解释相同位模式，不执行量化或反量化；
+- 普通 BF16/FP32 matmul 不能直接消费 packed code，必须先解码 code 并应用 scale。
+
+Block-scaled FP8 benchmark 在两个平台都从随机浮点数据出发，通过 `.to(torch.float8_e4m3fn)` 生成真实 FP8 tensor。由于没有可运行的原生 `scaled_mm`，两端都使用 `reference_provider=torch_dtype_dequant_mm`：先用 `.float()` 解码 FP8，再按 128-wide block 应用 A/B scale，最后执行 FP32 matmul。性能场景的 scale 初始化为 1，但 reference 仍执行完整 scale 路径；非单位 scale 由正确性测试覆盖。当前 benchmark 不提供 FP8 `uint8` 存储回退；若 PyTorch 不提供 E4M3FN dtype，该 benchmark 不能运行。
+
+MXFP4 benchmark 直接随机生成 packed E2M1 byte 和 E8M0 scale byte，它们是合法的量化编码，不是普通 `uint8` 权重，也不是从一份浮点模型权重经 calibration 得到。海光将这些 byte 以原生 packed dtype 暴露，天数因缺少 packed E2M1 dtype而保持 `uint8` 存储。两端原生 grouped operator 都不可用，因此 reference 都显式拆分高低 nibble、解码 E2M1、解码并应用 E8M0 scale，再执行 FP32 grouped matmul。对应 provider 分别为 `manual_dequant_mm_native_operator_unavailable` 和 `manual_dequant_mm_native_dtype_unavailable`。
+
+软件 reference 的计时包含解码、scale 展开、反量化临时张量和 FP32 matmul；因此相关加速比表示融合 ntops kernel 相对软件端到端 reference 的收益，不能解释为相对平台原生量化 operator 的收益。ntops kernel 不物化完整反量化权重：Block FP8 在 tile 内执行 dot 与 scale，MXFP4 在 tile 内在线解码并计算。
+
+### 3.3 正确性范围
+
+天数侧在当前分支运行四个定向测试文件，结果为 `78 passed, 2 skipped`。两个 skip 仅因 PyTorch 2.7.1 没有原生 MXFP4 dtype；raw `uint8` 数值路径仍实际执行。测试覆盖 FP8 两种 recipe、MXFP4 全部 16 个 code、Gated RMSNorm 全部分支、MLA 融合算子十项用例及生成源码检查。统一命令见 `docs/build_and_evaluate.md`。
+
+海光结果包括 Block-scaled FP8 三个代表 shape 的数值复验、MXFP4 三个 benchmark 场景、Gated RMSNorm 三个场景以及 MLA 融合算子定向测试；各项正确性检查均通过。
+
+## 4. 实验结果
+
+### 4.1 Block-scaled FP8 MM
+
+两个平台均使用 `torch_dtype_dequant_mm` 软件 reference。天数 MR-V100：
+
+| 场景 `(M,N,K)` | reference | reference (us) | ntops (us) | 加速比 | 最大绝对误差 |
+| --- | --- | ---: | ---: | ---: | ---: |
+| decode `(1,4096,4096)` | software dequant MM | 609.584 | 228.115 | 2.672x | 0.000000 |
+| MoE `(32,14336,4096)` | software dequant MM | 1910.558 | 1531.907 | 1.247x | 0.250000 |
+| prefill `(128,4096,4096)` | software dequant MM | 677.933 | 1726.408 | 0.393x | 0.000004 |
+
+海光 BW，固定 `(warps,stages)=(1,1)`：
+
+| 场景 `(M,N,K)` | reference (us) | ntops (us) | 加速比 | 最大绝对误差 |
+| --- | ---: | ---: | ---: | ---: |
+| decode `(1,4096,4096)` | 391.849 | 902.313 | 0.434x | 0.000000 |
+| MoE `(32,14336,4096)` | 1252.682 | 3853.278 | 0.325x | 0.250000 |
+| prefill `(128,4096,4096)` | 439.770 | 1733.113 | 0.254x | 0.125000 |
+
+两端都没有可运行的原生 public block-scaled operator，因此加速比只代表相对完整 FP8 解码、block scale 应用和 FP32 matmul 软件 reference。decode 在天数侧获益，prefill 与海光三项仍有优化空间。
+
+### 4.2 MXFP4 W4A16 Grouped MM
+
+天数 MR-V100 使用 raw `uint8` packed code，reference provider 为 `manual_dequant_mm_native_dtype_unavailable`：
+
+| 场景 | 模式 | ntops (us) | 加速比 | TFLOP/s | 最大绝对误差 |
+| --- | --- | ---: | ---: | ---: | ---: |
+| uniform decode | uniform | 892.727 | 42.268x | 0.301 | 0.000000 |
+| uniform concurrent | uniform | 1373.795 | 27.453x | 3.126 | 0.500000 |
+| routed uneven | routed | 5837.419 | 6.474x | 0.736 | 2.000000 |
+
+海光 BW 的原生 packed dtype 可用，但原生 grouped operator 不可用；reference provider 为 `manual_dequant_mm_native_operator_unavailable`：
+
+| 场景 | 模式 | reference (us) | ntops (us) | 配置 | 加速比 | 最大绝对误差 |
+| --- | --- | ---: | ---: | --- | ---: | ---: |
+| uniform decode | uniform | 18681.152 | 2219.319 | `(4,1)` | 8.418x | 0.000000 |
+| uniform concurrent | uniform | 18947.936 | 14143.086 | `(4,1)` | 1.340x | 0.031250 |
+| routed uneven | routed | 18885.632 | 53806.240 | `(1,1)` | 0.351x | 0.500000 |
+
+routed HIP 固定单 wave 后，相对错误选择 4-wave 时的 1296.786 ms 降到 53.806 ms，约 24.1 倍改善，但仍慢于软件 reference。
+
+### 4.3 Gated RMSNorm
+
+天数 MR-V100，BF16、hidden size 128：
+
+| 场景 | 行数 | 配置 | ntops (us) | PyTorch (us) | 加速比 | 最大绝对误差 |
+| --- | ---: | --- | ---: | ---: | ---: | ---: |
+| decode | 8 | `(8,2)` | 8.047 | 42.421 | 5.271x | 0.000000 |
+| concurrent decode | 80 | `(2,2)` | 8.888 | 53.066 | 5.971x | 0.000000 |
+| prefill 2048 | 16384 | `(1,1)` | 54.793 | 348.886 | 6.367x | 0.003906 |
+
+海光 BW：
+
+| 场景 | 行数 | 配置 | ntops (us) | PyTorch (us) | 加速比 | 最大绝对误差 |
+| --- | ---: | --- | ---: | ---: | ---: | ---: |
+| decode | 8 | `(8,2)` | 663.288 | 54.758 | 0.083x | 0.000000 |
+| concurrent decode | 80 | `(8,2)` | 645.244 | 51.970 | 0.081x | 0.000000 |
+| prefill 2048 | 16384 | `(4,2)` | 692.300 | 179.012 | 0.259x | 0.003906 |
+
+融合在天数侧有稳定收益；海光侧数值正确但 launch/lowering 开销占主导，继续调整候选本身不能替代后端与 program 映射分析。
+
+### 4.4 Fused MLA RoPE Cache Write
+
+shape 为 `L=512, R=64, cache_block_size=16, BF16`。天数数据由当前分支重新测得：
+
+| 场景 | T | PyTorch 未融合 (us) | 配置 | ntops (us) | 加速比 |
+| --- | ---: | ---: | --- | ---: | ---: |
+| decode | 1 | 63.273 | `(4,1)` | 9.103 | 6.951x |
+| concurrent decode | 10 | 81.243 | `(4,1)` | 10.497 | 7.739x |
+| long context | 2048 | 170.413 | `(2,2)` | 22.808 | 7.472x |
+
+相对未融合 PyTorch，三个场景均获得 6.95x 至 7.74x 加速。
+
+海光 BW：
+
+| 场景 | T | PyTorch 未融合 (us) | 配置 | ntops (us) | 加速比 |
+| --- | ---: | ---: | --- | ---: | ---: |
+| decode | 1 | 193.815 | `(8,1)` | 1289.147 | 0.150x |
+| concurrent decode | 10 | 195.311 | `(8,1)` | 1323.701 | 0.148x |
+| long context | 2048 | 180.899 | `(1,1)` | 1319.881 | 0.137x |
+
+海光三个场景的数值检查通过，但当前延迟高于 PyTorch 未融合 reference，主要开销来自该平台的输出锚定与 kernel launch/lowering。两平台结果分别报告，不跨平台外推配置或加速比。
+
+## 5. 对照与消融
+
+- Block-scaled FP8：CoreX 的 BF16 dot 是 FP8 有限值的精确回退。HIP 多 wave
+  specialization 会触发 AMD LLVM 崩溃，恢复单 wave 后正确性与服务稳定性恢复，
+  代价是当前性能低于软件 reference。
+- MXFP4：查表解码曾因 tensor-valued 索引坐标丢失产生约 98% 元素错误；改为纯
+  integer bitwise/arithmetic 后可穷举通过全部 code。routed 4-wave 到 1-wave 的
+  24.1x 延迟改善证明 wave 策略是独立关键因素。
+- Gated RMSNorm：相对 PyTorch eager 的 5.27x 至 6.37x 天数收益来自归约、门控、
+  激活和权重乘法单 kernel 融合；海光反例表明融合不能自动抵消后端 launch 开销。
+- Fused MLA RoPE Cache Write：直接对照 eager RoPE 加两次 cache scatter，长上下文
+  获得 7.47x 加速。
+
+## 6. 工程质量与适用边界
+
+### 6.1 测试与兼容性
+
+测试覆盖公开 dtype、主要分支、尾块、空/零 token、padding、非连续 view、原地写入、非法参数、平台分派和生成源码。量化 reference 独立解码；MLA reference 使用 PyTorch FP32 RoPE 和显式 paged scatter。执行命令见 `docs/build_and_evaluate.md`。
+
+API 改动控制在新增的任务入口；未修改 PyTorch 或 NineToothed 全局行为。MLA 融合算子只暴露一个正式入口。其他三个算子只接受明确列出的 recipe/layout，不把不支持的组合静默解释为相近语义。
+
+### 6.2 通用能力与平台特化
+
+通用层包含数学语义、layout contract、reference、公开 DSL application 和参数校验。平台特化仅包含 wrapper 选择的 dot/reduction、tile、wave 和候选集合。平台路径由 HIP/CoreX/CUDA capability 决定，不把设备名称散落进数学实现。
+
+### 6.3 已知限制
+
+- Block-scaled FP8 暂不支持 A 侧 128x128 scale、swizzle、多级 scale、batch、
+  grouped launch、N 尾块和 fast accumulation；
+- MXFP4 只支持 W4A16 `BlockWise1x32`，不支持 activation scale、bias、swizzle、
+  grouped-K 和 `[G,N,K/2]` 直接输入；
+- Gated RMSNorm 的新 HIP 配置需先做独立进程编译筛选；
+- Fused MLA RoPE Cache Write 支持 FP16/BF16/FP32 未量化 cache，不支持 FP8 cache layout；
+- 非连续 cache 会产生临时连续副本；
+- 软件 reference 的性能不能当作平台原生算子基线。
+
+### 6.4 第三方来源
+
+实现代码为本仓库 NineToothed DSL 代码，没有复制第三方 kernel。接口与数学语义参考以下上游公开资料：
+
+- PyTorch `torch.nn.functional.scaled_mm` 与 block-scaled tests；
+- PyTorch `torch.nn.functional.scaled_grouped_mm`；
+- vLLM MXFP4 quantization utilities；
+- vLLM `concat_and_cache_mla` 与 fusion design。
+
+依赖许可证与仓库本身保持一致。引用只用于接口、布局和基线语义，不引入私有 backend patch。
+
+## 7. 结论
+
+分支完成四个 A 部分算子的跨 CoreX/HIP 实现与验证，B 部分保持官方编译后端不变。天数侧 Gated RMSNorm、MXFP4 和 Fused MLA RoPE Cache Write 获得明确收益；Block-scaled FP8 的收益集中在 decode。海光侧 MXFP4 uniform 有收益，其余路径暴露出 codegen、wave 与 launch 开销边界，报告没有隐藏负收益。MLA 算子采用命名明确、接口唯一、自包含的融合 kernel；两平台正确性与性能结果均已记录。
diff --git a/scripts/build_and_evaluate.sh b/scripts/build_and_evaluate.sh
new file mode 100755
index 0000000..3e0a2ab
--- /dev/null
+++ b/scripts/build_and_evaluate.sh
@@ -0,0 +1,87 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+MODE="${1:-all}"
+PROJECT_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+PYTHON_BIN="${PYTHON_BIN:-python}"
+RESULT_ROOT="${RESULT_ROOT:-${PROJECT_ROOT}/artifacts}"
+RUN_ID="$(date -u +%Y%m%dT%H%M%SZ)"
+RESULT_DIR="${RESULT_ROOT}/${RUN_ID}"
+
+TEST_FILES=(
+  tests/test_block_scaled_fp8_mm.py
+  tests/test_mxfp4_w4a16_grouped_mm.py
+  tests/test_rms_norm_gated.py
+  tests/test_fused_mla_rope_cache_write.py
+)
+
+BENCHMARK_FILES=(
+  benchmarks/bench_block_scaled_fp8_mm.py
+  benchmarks/bench_mxfp4_w4a16_grouped_mm.py
+  benchmarks/bench_rms_norm_gated.py
+  benchmarks/bench_fused_mla_rope_cache_write.py
+)
+
+usage() {
+  echo "usage: $0 [all|build|test|benchmark]" >&2
+}
+
+case "${MODE}" in
+  all|build|test|benchmark) ;;
+  *)
+    usage
+    exit 2
+    ;;
+esac
+
+cd "${PROJECT_ROOT}"
+mkdir -p "${RESULT_DIR}"
+
+if [[ "${SKIP_INSTALL:-0}" != "1" ]]; then
+  "${PYTHON_BIN}" -m pip install -e ".[testing]" 2>&1 \
+    | tee "${RESULT_DIR}/build.log"
+fi
+
+"${PYTHON_BIN}" - <<'PY' | tee "${RESULT_DIR}/environment.txt"
+import inspect
+import platform
+
+import ninetoothed
+import torch
+
+print("python", platform.python_version())
+print("torch", torch.__version__)
+print("cuda", torch.version.cuda)
+print("hip", torch.version.hip)
+print("ninetoothed", inspect.getfile(ninetoothed))
+print("accelerator_available", torch.cuda.is_available())
+if torch.cuda.is_available():
+    print("device", torch.cuda.get_device_name())
+PY
+
+if [[ "${MODE}" == "build" ]]; then
+  echo "build artifacts: ${RESULT_DIR}"
+  exit 0
+fi
+
+"${PYTHON_BIN}" - <<'PY'
+import torch
+
+if not torch.cuda.is_available():
+    raise SystemExit("a CUDA/HIP/CoreX-compatible accelerator is required")
+PY
+
+if [[ "${MODE}" == "all" || "${MODE}" == "test" ]]; then
+  PYTHONPATH=src "${PYTHON_BIN}" -m pytest -q "${TEST_FILES[@]}" 2>&1 \
+    | tee "${RESULT_DIR}/pytest.log"
+fi
+
+if [[ "${MODE}" == "all" || "${MODE}" == "benchmark" ]]; then
+  for benchmark_file in "${BENCHMARK_FILES[@]}"; do
+    benchmark_name="$(basename "${benchmark_file}" .py)"
+    PYTHONPATH=src "${PYTHON_BIN}" "${benchmark_file}" 2>&1 \
+      | tee "${RESULT_DIR}/${benchmark_name}.csv"
+  done
+fi
+
+echo "evaluation artifacts: ${RESULT_DIR}"
diff --git a/src/ntops/kernels/__init__.py b/src/ntops/kernels/__init__.py
index 12d337b..74ee0a2 100644
--- a/src/ntops/kernels/__init__.py
+++ b/src/ntops/kernels/__init__.py
@@ -1,12 +1,21 @@
 from ntops.kernels import (
     abs,
+    acosh,
+    adaptive_avg_pool2d,
+    adaptive_max_pool2d,
     add,
     addmm,
+    addmv,
     alpha_dropout,
+    argsort,
+    atan,
     avg_pool2d,
+    batch_norm,
+    bincount,
     bitwise_and,
     bitwise_not,
     bitwise_or,
+    block_scaled_fp8_mm,
     bmm,
     celu,
     clamp,
@@ -18,6 +27,8 @@
     dropout,
     eq,
     exp,
+    fmax,
+    fused_mla_rope_cache_write,
     ge,
     gelu,
     gt,
@@ -26,17 +37,29 @@
     isnan,
     layer_norm,
     le,
+    logsumexp,
+    lp_pool1d,
+    lp_pool2d,
+    lp_pool3d,
     lt,
+    max,
+    max_pool1d,
     max_pool2d,
+    max_pool3d,
+    maximum,
+    mean,
+    median,
     mm,
     msort,
     mul,
+    mxfp4_w4a16_grouped_mm,
     ne,
     neg,
     pow,
     quantile,
     relu,
     rms_norm,
+    rms_norm_gated,
     rot90,
     rotary_position_embedding,
     round,
@@ -51,29 +74,10 @@
     sin,
     softmax,
     sort,
+    stack,
     sub,
     tanh,
     threshold,
-    max_pool1d,
-    max_pool3d,
-    stack,  
-    mean,
-    median,
-    maximum,
-    atan,
-    batch_norm,
-    bincount,
-    adaptive_max_pool2d,  
-    acosh,
-    adaptive_avg_pool2d,
-    addmv,
-    argsort,
-    fmax,
-    logsumexp,
-    lp_pool1d,
-    lp_pool2d,
-    lp_pool3d,
-    max,
 )
 
 __all__ = [
@@ -115,11 +119,14 @@
     "quantile",
     "relu",
     "rms_norm",
+    "rms_norm_gated",
     "rot90",
     "rotary_position_embedding",
     "round",
     "rsqrt",
     "scaled_dot_product_attention",
+    "mxfp4_w4a16_grouped_mm",
+    "block_scaled_fp8_mm",
     "select_copy",
     "sgn",
     "sigmoid",
@@ -152,4 +159,5 @@
     "lp_pool2d",
     "lp_pool3d",
     "max",
+    "fused_mla_rope_cache_write",
 ]
diff --git a/src/ntops/kernels/block_scaled_fp8_mm.py b/src/ntops/kernels/block_scaled_fp8_mm.py
new file mode 100644
index 0000000..52ecac3
--- /dev/null
+++ b/src/ntops/kernels/block_scaled_fp8_mm.py
@@ -0,0 +1,189 @@
+import functools
+
+import ninetoothed.language as ntl
+import torch
+from ninetoothed import Tensor
+
+BLOCK_SIZE_M = 16
+BLOCK_SIZE_N = 16
+BLOCK_SIZE_K = 32
+
+
+def _arrange_mat_a(mat_a, output_arranged, block_size_m, block_size_k):
+    arranged = mat_a.tile((1, block_size_m, block_size_k))
+    arranged = arranged.tile((1, 1, -1))
+    arranged = arranged.expand((-1, -1, output_arranged.shape[-1]))
+    arranged.dtype = arranged.dtype.squeeze((0, 1))
+    arranged.dtype.dtype = arranged.dtype.dtype.squeeze(0)
+    return arranged
+
+
+def _arrange_mat_b(mat_b, output_arranged, block_size_n, block_size_k):
+    arranged = mat_b.tile((1, block_size_k, block_size_n))
+    arranged = arranged.tile((1, -1, 1))
+    arranged = arranged.expand((-1, output_arranged.shape[-2], -1))
+    arranged.dtype = arranged.dtype.squeeze((0, 2))
+    arranged.dtype.dtype = arranged.dtype.dtype.squeeze(0)
+    return arranged
+
+
+def _arrange_scale_a(scale_a, output_arranged, block_size_m):
+    arranged = scale_a.tile((1, block_size_m, 1))
+    arranged = arranged.tile((1, 1, -1))
+    arranged = arranged.expand((-1, -1, output_arranged.shape[-1]))
+    arranged.dtype = arranged.dtype.squeeze((0, 1))
+    arranged.dtype.dtype = arranged.dtype.dtype.squeeze(0)
+    return arranged
+
+
+def _arrange_scale_b(scale_b, output_arranged, block_size_n):
+    arranged = scale_b.tile((1, 1, block_size_n))
+    arranged = arranged.tile((1, -1, 1))
+    arranged = arranged.expand((-1, output_arranged.shape[-2], -1))
+    arranged.dtype = arranged.dtype.squeeze((0, 2))
+    arranged.dtype.dtype = arranged.dtype.dtype.squeeze(0)
+    return arranged
+
+
+def arrangement(
+    mat_a,
+    mat_b,
+    scale_a,
+    scale_b,
+    output,
+    block_size_m=BLOCK_SIZE_M,
+    block_size_n=BLOCK_SIZE_N,
+    block_size_k=BLOCK_SIZE_K,
+):
+    output_arranged = output.tile((1, block_size_m, block_size_n))
+    output_arranged.dtype = output_arranged.dtype.squeeze(0)
+
+    return (
+        _arrange_mat_a(mat_a, output_arranged, block_size_m, block_size_k),
+        _arrange_mat_b(mat_b, output_arranged, block_size_n, block_size_k),
+        _arrange_scale_a(scale_a, output_arranged, block_size_m),
+        _arrange_scale_b(scale_b, output_arranged, block_size_n),
+        output_arranged,
+    )
+
+
+def arrangement_with_bias(
+    mat_a,
+    mat_b,
+    scale_a,
+    scale_b,
+    bias,
+    output,
+    block_size_m=BLOCK_SIZE_M,
+    block_size_n=BLOCK_SIZE_N,
+    block_size_k=BLOCK_SIZE_K,
+):
+    arranged = arrangement(
+        mat_a,
+        mat_b,
+        scale_a,
+        scale_b,
+        output,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        block_size_k=block_size_k,
+    )
+    output_arranged = arranged[-1]
+    bias_arranged = bias.tile((1, 1, block_size_n))
+    bias_arranged = bias_arranged.expand((-1, output_arranged.shape[-2], -1))
+    bias_arranged.dtype = bias_arranged.dtype.squeeze((0, 1))
+    return (*arranged[:-1], bias_arranged, output_arranged)
+
+
+def _scaled_dot(mat_a, mat_b, scale_a, scale_b, output, dots_per_scale):
+    accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
+
+    for scale_index in range(scale_a.shape[0]):
+        block_accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
+        for k_offset in range(dots_per_scale):
+            k = scale_index * dots_per_scale + k_offset
+            block_accumulator += ntl.dot(mat_a[k], mat_b[k])
+        scale = (scale_a[scale_index] + 0).to(ntl.float32) * (
+            scale_b[scale_index] + 0
+        ).to(ntl.float32)
+        accumulator += block_accumulator * scale
+
+    return accumulator
+
+
+def _scaled_dot_bf16(mat_a, mat_b, scale_a, scale_b, output, dots_per_scale):
+    accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
+
+    for scale_index in range(scale_a.shape[0]):
+        block_accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
+        for k_offset in range(dots_per_scale):
+            k = scale_index * dots_per_scale + k_offset
+            activation = mat_a[k].to(ntl.bfloat16)
+            weight = mat_b[k].to(ntl.bfloat16)
+            block_accumulator += ntl.dot(activation, weight)
+        scale = (scale_a[scale_index] + 0).to(ntl.float32) * (
+            scale_b[scale_index] + 0
+        ).to(ntl.float32)
+        accumulator += block_accumulator * scale
+
+    return accumulator
+
+
+def block_scaled_fp8_mm_k16(mat_a, mat_b, scale_a, scale_b, output):
+    output = _scaled_dot_bf16(  # noqa: F841
+        mat_a, mat_b, scale_a, scale_b, output, 8
+    )
+
+
+def block_scaled_fp8_mm_k32(mat_a, mat_b, scale_a, scale_b, output):
+    output = _scaled_dot(mat_a, mat_b, scale_a, scale_b, output, 4)  # noqa: F841
+
+
+def block_scaled_fp8_mm_with_bias_k16(mat_a, mat_b, scale_a, scale_b, bias, output):
+    accumulator = _scaled_dot_bf16(mat_a, mat_b, scale_a, scale_b, output, 8)
+    output = accumulator + (bias + 0).to(ntl.float32)  # noqa: F841
+
+
+def block_scaled_fp8_mm_with_bias_k32(mat_a, mat_b, scale_a, scale_b, bias, output):
+    accumulator = _scaled_dot(mat_a, mat_b, scale_a, scale_b, output, 4)
+    output = accumulator + (bias + 0).to(ntl.float32)  # noqa: F841
+
+
+def premake(
+    input_dtype,
+    output_dtype,
+    bias_dtype=None,
+    block_size_m=BLOCK_SIZE_M,
+    block_size_n=BLOCK_SIZE_N,
+    block_size_k=BLOCK_SIZE_K,
+):
+    if block_size_k not in (16, 32):
+        raise ValueError("block_size_k must be 16 or 32")
+
+    has_bias = bias_dtype is not None
+    arrangement_function = arrangement_with_bias if has_bias else arrangement
+    arrangement_ = functools.partial(
+        arrangement_function,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        block_size_k=block_size_k,
+    )
+    tensors = (
+        Tensor(3, dtype=input_dtype, other=0.0),
+        Tensor(3, dtype=input_dtype, other=0.0),
+        Tensor(3, dtype=torch.float32, other=0.0),
+        Tensor(3, dtype=torch.float32, other=0.0),
+    )
+
+    if has_bias:
+        tensors += (Tensor(3, dtype=bias_dtype, other=0.0),)
+
+    tensors += (Tensor(3, dtype=output_dtype),)
+    applications = {
+        (16, False): block_scaled_fp8_mm_k16,
+        (16, True): block_scaled_fp8_mm_with_bias_k16,
+        (32, False): block_scaled_fp8_mm_k32,
+        (32, True): block_scaled_fp8_mm_with_bias_k32,
+    }
+    application_function = applications[(block_size_k, has_bias)]
+    return arrangement_, application_function, tensors
diff --git a/src/ntops/kernels/fused_mla_rope_cache_write.py b/src/ntops/kernels/fused_mla_rope_cache_write.py
new file mode 100644
index 0000000..4c6f053
--- /dev/null
+++ b/src/ntops/kernels/fused_mla_rope_cache_write.py
@@ -0,0 +1,159 @@
+"""NineToothed fused MLA RoPE + compressed KV-cache write kernel.
+
+The cache entry follows vLLM's MLA layout:
+
+    [compressed kv_c | RoPE(k_pe)]
+
+``kv_c`` is the shared low-rank K/V latent.  It is not expanded into
+per-head K/V in this kernel.  ``k_pe`` is the shared rotary key component and
+is rotated in-place in the write path.
+"""
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Tensor
+
+
+def _source_view(tensor):
+    return tensor.tile((1,) * tensor.ndim)
+
+
+def arrangement(
+    output_anchor,
+    kv_c,
+    k_pe,
+    kv_cache,
+    slot_mapping,
+    positions,
+    cos_sin_cache,
+    entry_dim,
+    tile_size,
+    kv_lora_rank,
+    rope_dim,
+    cache_block_size,
+):
+    # A real output anchor keeps the side-effecting cache store in the SSA graph
+    # on backends that cannot lower a cache-only launch. The wrapper discards
+    # this duplicate compressed entry after the kernel completes.
+    output_anchor = output_anchor.tile((1, tile_size.value))
+    output_anchor.dtype = output_anchor.dtype.squeeze(0)
+    kv_c, k_pe, kv_cache, slot_mapping, positions, cos_sin_cache = (
+        _source_view(tensor)
+        for tensor in (
+            kv_c,
+            k_pe,
+            kv_cache,
+            slot_mapping,
+            positions,
+            cos_sin_cache,
+        )
+    )
+    return (
+        output_anchor,
+        kv_c,
+        k_pe,
+        kv_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        entry_dim,
+        tile_size,
+        kv_lora_rank,
+        rope_dim,
+        cache_block_size,
+    )
+
+
+def application(
+    output_anchor,
+    kv_c,
+    k_pe,
+    kv_cache,
+    slot_mapping,
+    positions,
+    cos_sin_cache,
+    entry_dim,
+    tile_size,
+    kv_lora_rank,
+    rope_dim,
+    cache_block_size,
+):
+    token_idx = output_anchor.offsets(0)
+    feature = output_anchor.offsets(1)
+    feature_valid = feature < entry_dim
+    nope_mask = feature_valid & (feature < kv_lora_rank)
+    rope_mask = feature_valid & (feature >= kv_lora_rank)
+    nope_feature = ntl.where(nope_mask, feature, 0)
+    rope_feature = ntl.where(rope_mask, feature - kv_lora_rank, 0)
+    pair = rope_feature // 2
+    even = (rope_feature & 1) == 0
+
+    latent_value = kv_c.source[token_idx, nope_feature]
+    kpe0 = k_pe.source[token_idx, pair * 2]
+    kpe1 = k_pe.source[token_idx, pair * 2 + 1]
+    position = positions.source[token_idx].to(ntl.int64)
+    cos = cos_sin_cache.source[position, pair]
+    sin = cos_sin_cache.source[position, rope_dim // 2 + pair]
+    rotated = ntl.where(
+        even,
+        kpe0 * cos - kpe1 * sin,
+        kpe0 * sin + kpe1 * cos,
+    )
+    cache_value = ntl.where(nope_mask, latent_value, rotated).to(kv_c.dtype)
+
+    output_anchor = cache_value  # noqa: F841
+
+    # Negative slots are padding. Indexed stores include source-shape bounds
+    # masks, so mapping padding to a negative block suppresses the write.
+    slot = slot_mapping.source[token_idx].to(ntl.int64)
+    write_slot = ntl.where(slot >= 0, slot, -1)
+    block_idx = write_slot // cache_block_size
+    block_offset = write_slot % cache_block_size
+    kv_cache.source[block_idx, block_offset, feature] = cache_value
+
+
+def premake(
+    kv_lora_rank,
+    rope_dim,
+    dtype=None,
+    block_size=128,
+    cache_block_size=16,
+    cos_dtype=None,
+):
+    entry_dim = kv_lora_rank + rope_dim
+    dynamic = {"constexpr": True, "upper_bound": 2**20}
+    tensors = (
+        Tensor(
+            shape=(None, entry_dim),
+            dtype=dtype,
+            shape_options=(dynamic, {"constexpr": True}),
+        ),
+        Tensor(
+            shape=(None, kv_lora_rank),
+            dtype=dtype,
+            shape_options=(dynamic, {"constexpr": True}),
+        ),
+        Tensor(
+            shape=(None, rope_dim),
+            dtype=dtype,
+            shape_options=(dynamic, {"constexpr": True}),
+        ),
+        Tensor(
+            shape=(None, None, entry_dim),
+            dtype=dtype,
+            shape_options=(dynamic, dynamic, {"constexpr": True}),
+        ),
+        Tensor(shape=(None,), dtype="int64", shape_options=(dynamic,)),
+        Tensor(shape=(None,), dtype="int64", shape_options=(dynamic,)),
+        Tensor(
+            shape=(None, rope_dim),
+            dtype=cos_dtype or ninetoothed.float32,
+            shape_options=(dynamic, {"constexpr": True}),
+        ),
+        Tensor(0, constexpr=True, value=entry_dim),
+        Tensor(0, constexpr=True, value=block_size),
+        Tensor(0, constexpr=True, value=kv_lora_rank),
+        Tensor(0, constexpr=True, value=rope_dim),
+        Tensor(0, constexpr=True, value=cache_block_size),
+    )
+    return arrangement, application, tensors
diff --git a/src/ntops/kernels/mxfp4_w4a16_grouped_mm.py b/src/ntops/kernels/mxfp4_w4a16_grouped_mm.py
new file mode 100644
index 0000000..e2746b1
--- /dev/null
+++ b/src/ntops/kernels/mxfp4_w4a16_grouped_mm.py
@@ -0,0 +1,195 @@
+import functools
+
+import ninetoothed
+import ninetoothed.language as ntl
+import torch
+from ninetoothed import Tensor
+
+BLOCK_SIZE_M = ninetoothed.block_size(lower_bound=16)
+BLOCK_SIZE_N = ninetoothed.block_size(lower_bound=16)
+MICROSCALE_K = 32
+PACKED_MICROSCALE_K = MICROSCALE_K // 2
+
+
+def _arrange_activation(mat_a, output_arranged, block_size_m):
+    arranged = mat_a.tile(
+        (1, block_size_m, PACKED_MICROSCALE_K),
+        strides=(1, block_size_m, MICROSCALE_K),
+        dilation=(1, 1, 2),
+    )
+    arranged = arranged.tile((1, 1, -1))
+    arranged = arranged.expand((-1, -1, output_arranged.shape[-1]))
+    arranged.dtype = arranged.dtype.squeeze((0, 1))
+    arranged.dtype.dtype = arranged.dtype.dtype.squeeze(0)
+    return arranged
+
+
+def arrangement(
+    mat_a_even,
+    mat_a_odd,
+    mat_b,
+    scale_b,
+    output,
+    block_size_m=None,
+    block_size_n=None,
+):
+    if block_size_m is None:
+        block_size_m = BLOCK_SIZE_M
+
+    if block_size_n is None:
+        block_size_n = BLOCK_SIZE_N
+
+    output_arranged = output.tile((1, block_size_m, block_size_n))
+    output_arranged.dtype = output_arranged.dtype.squeeze(0)
+
+    mat_a_even_arranged = _arrange_activation(mat_a_even, output_arranged, block_size_m)
+    mat_a_odd_arranged = _arrange_activation(mat_a_odd, output_arranged, block_size_m)
+
+    mat_b_arranged = mat_b.tile((1, PACKED_MICROSCALE_K, block_size_n))
+    mat_b_arranged = mat_b_arranged.tile((1, -1, 1))
+    mat_b_arranged = mat_b_arranged.expand((-1, output_arranged.shape[-2], -1))
+    mat_b_arranged.dtype = mat_b_arranged.dtype.squeeze((0, 2))
+    mat_b_arranged.dtype.dtype = mat_b_arranged.dtype.dtype.squeeze(0)
+
+    scale_b_arranged = scale_b.tile((1, 1, block_size_n))
+    scale_b_arranged = scale_b_arranged.tile((1, -1, 1))
+    scale_b_arranged = scale_b_arranged.expand((-1, output_arranged.shape[-2], -1))
+    scale_b_arranged.dtype = scale_b_arranged.dtype.squeeze((0, 2))
+    scale_b_arranged.dtype.dtype = scale_b_arranged.dtype.dtype.squeeze(0)
+
+    return (
+        mat_a_even_arranged,
+        mat_a_odd_arranged,
+        mat_b_arranged,
+        scale_b_arranged,
+        output_arranged,
+    )
+
+
+def mxfp4_w4a16_grouped_mm_dot(mat_a_even, mat_a_odd, mat_b, scale_b, output):
+    accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
+
+    for k in range(mat_a_even.shape[0]):
+        packed = (mat_b[k] + 0).to(ntl.int32)
+        scale = ntl.exp2((scale_b[k] + 0).to(ntl.float32) - 127.0)
+
+        even_code = packed & 0xF
+        even_magnitude_code = even_code & 0x7
+        even_exponent = even_magnitude_code >> 1
+        even_mantissa = ((even_magnitude_code & 1) + 0).to(ntl.float32)
+        even_normal = (1.0 + 0.5 * even_mantissa) * ntl.exp2(
+            (even_exponent + 0).to(ntl.float32) - 1.0
+        )
+        even_magnitude = ntl.where(even_exponent == 0, 0.5 * even_mantissa, even_normal)
+        even_sign = ntl.where((even_code & 0x8) == 0, 1.0, -1.0)
+        weight_even = (even_sign * even_magnitude * scale).to(ntl.bfloat16)
+
+        odd_code = (packed >> 4) & 0xF
+        odd_magnitude_code = odd_code & 0x7
+        odd_exponent = odd_magnitude_code >> 1
+        odd_mantissa = ((odd_magnitude_code & 1) + 0).to(ntl.float32)
+        odd_normal = (1.0 + 0.5 * odd_mantissa) * ntl.exp2(
+            (odd_exponent + 0).to(ntl.float32) - 1.0
+        )
+        odd_magnitude = ntl.where(odd_exponent == 0, 0.5 * odd_mantissa, odd_normal)
+        odd_sign = ntl.where((odd_code & 0x8) == 0, 1.0, -1.0)
+        weight_odd = (odd_sign * odd_magnitude * scale).to(ntl.bfloat16)
+
+        accumulator += ntl.dot(mat_a_even[k], weight_even)
+        accumulator += ntl.dot(mat_a_odd[k], weight_odd)
+
+    output = accumulator
+
+
+def mxfp4_w4a16_grouped_mm_reduction(mat_a_even, mat_a_odd, mat_b, scale_b, output):
+    accumulator = ntl.zeros(output.shape, dtype=ntl.float32)
+
+    for k in range(mat_a_even.shape[0]):
+        packed = (mat_b[k] + 0).to(ntl.int32)
+        scale = ntl.exp2((scale_b[k] + 0).to(ntl.float32) - 127.0)
+
+        even_code = packed & 0xF
+        even_bit_0 = even_code & 1
+        even_bit_1 = (even_code >> 1) & 1
+        even_bit_2 = (even_code >> 2) & 1
+        even_low_magnitude = even_bit_0 + (even_bit_1 << 1)
+        even_twice_magnitude = even_low_magnitude + even_bit_2 * (
+            4 + even_bit_0 + (even_bit_1 << 1) + ((even_bit_0 * even_bit_1) << 1)
+        )
+        even_sign = 1 - (((even_code >> 3) & 1) << 1)
+        decoded_even = ((even_sign * even_twice_magnitude) + 0).to(ntl.float32) * 0.5
+        weight_even = (decoded_even * scale).to(ntl.bfloat16)
+
+        odd_code = (packed >> 4) & 0xF
+        odd_bit_0 = odd_code & 1
+        odd_bit_1 = (odd_code >> 1) & 1
+        odd_bit_2 = (odd_code >> 2) & 1
+        odd_low_magnitude = odd_bit_0 + (odd_bit_1 << 1)
+        odd_twice_magnitude = odd_low_magnitude + odd_bit_2 * (
+            4 + odd_bit_0 + (odd_bit_1 << 1) + ((odd_bit_0 * odd_bit_1) << 1)
+        )
+        odd_sign = 1 - (((odd_code >> 3) & 1) << 1)
+        decoded_odd = ((odd_sign * odd_twice_magnitude) + 0).to(ntl.float32) * 0.5
+        weight_odd = (decoded_odd * scale).to(ntl.bfloat16)
+        activation_even = (mat_a_even[k] + 0).to(ntl.float32)
+        activation_odd = (mat_a_odd[k] + 0).to(ntl.float32)
+        weight_even = (weight_even + 0).to(ntl.float32)
+        weight_odd = (weight_odd + 0).to(ntl.float32)
+
+        accumulator += ntl.sum(
+            activation_even[:, :, None] * weight_even[None, :, :], axis=1
+        )
+        accumulator += ntl.sum(
+            activation_odd[:, :, None] * weight_odd[None, :, :], axis=1
+        )
+
+    output = accumulator
+
+
+def premake(jagged=False, block_size_m=None, block_size_n=None, reduction=False):
+    arrangement_ = functools.partial(
+        arrangement,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+    )
+    jagged_dim = 1 if jagged else None
+    # HIP uses the reduction path only. Specializing its matrix dimensions lets
+    # Triton fold the repeated shape predicates before AMD LLVM codegen.
+    dense_shape_options = {"constexpr": True} if reduction else None
+    activation_shape_options = (
+        ({"constexpr": True}, None, {"constexpr": True})
+        if reduction and jagged
+        else dense_shape_options
+    )
+    common_tensors = (
+        Tensor(
+            3,
+            dtype=torch.bfloat16,
+            jagged_dim=jagged_dim,
+            other=0,
+            shape_options=activation_shape_options,
+        ),
+        Tensor(
+            3,
+            dtype=torch.bfloat16,
+            jagged_dim=jagged_dim,
+            other=0,
+            shape_options=activation_shape_options,
+        ),
+        Tensor(3, dtype=torch.uint8, other=0, shape_options=dense_shape_options),
+        Tensor(3, dtype=torch.uint8, other=127, shape_options=dense_shape_options),
+    )
+    output_tensor = (
+        Tensor(
+            3,
+            dtype=torch.bfloat16,
+            jagged_dim=jagged_dim,
+            shape_options=activation_shape_options,
+        ),
+    )
+    tensors = common_tensors + output_tensor
+
+    application = (
+        mxfp4_w4a16_grouped_mm_reduction if reduction else mxfp4_w4a16_grouped_mm_dot
+    )
+    return arrangement_, application, tensors
diff --git a/src/ntops/kernels/rms_norm_gated.py b/src/ntops/kernels/rms_norm_gated.py
new file mode 100644
index 0000000..bbcfec9
--- /dev/null
+++ b/src/ntops/kernels/rms_norm_gated.py
@@ -0,0 +1,185 @@
+import functools
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Tensor
+
+
+def _silu(gate):
+    return gate / (1 + ntl.exp(-gate))
+
+
+def _sigmoid(gate):
+    return 1 / (1 + ntl.exp(-gate))
+
+
+def _arrange_last_dim(tensor, block_size):
+    # Keep each row in one program so the RMS reduction is local.
+    tile_shape = (1,) * (tensor.ndim - 1) + (block_size,)
+    return tensor.tile(tile_shape)
+
+
+def _arrange_group(tensor, group_size):
+    return tensor.flatten().tile((group_size,))
+
+
+def arrangement(
+    input,
+    gate,
+    weight,
+    eps,
+    output,
+    num_normalized_elements,
+    group_size=None,
+    block_size=128,
+):
+    if group_size is None:
+        input, gate, weight, output = (
+            _arrange_last_dim(tensor, block_size)
+            for tensor in (input, gate, weight, output)
+        )
+    else:
+        input, gate, weight, output = (
+            _arrange_group(tensor, group_size)
+            for tensor in (input, gate, weight, output)
+        )
+
+    return input, gate, weight, eps, output, num_normalized_elements
+
+
+def arrangement_no_gate(
+    input,
+    weight,
+    eps,
+    output,
+    num_normalized_elements,
+    group_size=None,
+    block_size=128,
+):
+    if group_size is None:
+        input, weight, output = (
+            _arrange_last_dim(tensor, block_size) for tensor in (input, weight, output)
+        )
+    else:
+        input, weight, output = (
+            _arrange_group(tensor, group_size) for tensor in (input, weight, output)
+        )
+
+    return input, weight, eps, output, num_normalized_elements
+
+
+def application_silu_before(input, gate, weight, eps, output, num_normalized_elements):
+    input_f32 = (input + 0).to(ntl.float32)
+    gate_f32 = (gate + 0).to(ntl.float32)
+    weight_f32 = (weight + 0).to(ntl.float32)
+    rms_value = ntl.sqrt(ntl.sum(input_f32 * input_f32) / num_normalized_elements + eps)
+    output = input_f32 / rms_value * weight_f32 * _silu(gate_f32)  # noqa: F841
+
+
+def application_silu_after(input, gate, weight, eps, output, num_normalized_elements):
+    input_f32 = (input + 0).to(ntl.float32)
+    gate_f32 = (gate + 0).to(ntl.float32)
+    weight_f32 = (weight + 0).to(ntl.float32)
+    gated = input_f32 * _silu(gate_f32)
+    rms_value = ntl.sqrt(ntl.sum(gated * gated) / num_normalized_elements + eps)
+    output = gated / rms_value * weight_f32  # noqa: F841
+
+
+def application_sigmoid_before(
+    input, gate, weight, eps, output, num_normalized_elements
+):
+    input_f32 = (input + 0).to(ntl.float32)
+    gate_f32 = (gate + 0).to(ntl.float32)
+    weight_f32 = (weight + 0).to(ntl.float32)
+    rms_value = ntl.sqrt(ntl.sum(input_f32 * input_f32) / num_normalized_elements + eps)
+    output = (  # noqa: F841
+        input_f32 / rms_value * weight_f32 * _sigmoid(gate_f32)
+    )
+
+
+def application_sigmoid_after(
+    input, gate, weight, eps, output, num_normalized_elements
+):
+    input_f32 = (input + 0).to(ntl.float32)
+    gate_f32 = (gate + 0).to(ntl.float32)
+    weight_f32 = (weight + 0).to(ntl.float32)
+    gated = input_f32 * _sigmoid(gate_f32)
+    rms_value = ntl.sqrt(ntl.sum(gated * gated) / num_normalized_elements + eps)
+    output = gated / rms_value * weight_f32  # noqa: F841
+
+
+def application_no_gate(input, weight, eps, output, num_normalized_elements):
+    input_f32 = (input + 0).to(ntl.float32)
+    weight_f32 = (weight + 0).to(ntl.float32)
+    rms_value = ntl.sqrt(ntl.sum(input_f32 * input_f32) / num_normalized_elements + eps)
+    output = input_f32 / rms_value * weight_f32  # noqa: F841
+
+
+def premake(
+    ndim,
+    hidden_size,
+    group_size=None,
+    norm_before_gate=True,
+    activation="silu",
+    dtype=None,
+    input_dtype=None,
+    gate_dtype=None,
+    weight_dtype=None,
+    output_dtype=None,
+    block_size=128,
+    has_gate=True,
+):
+    if activation not in ("silu", "sigmoid", "swish"):
+        raise ValueError("activation must be one of 'silu', 'sigmoid', or 'swish'")
+    if dtype is not None:
+        input_dtype = input_dtype or dtype
+        gate_dtype = gate_dtype or dtype
+        weight_dtype = weight_dtype or dtype
+        output_dtype = output_dtype or dtype
+    if block_size is None:
+        block_size = 128
+
+    tensor_shape = (None,) * (ndim - 1) + (hidden_size,)
+
+    if has_gate:
+        arrangement_ = functools.partial(
+            arrangement,
+            group_size=group_size,
+            block_size=max(block_size, hidden_size),
+        )
+        tensors = (
+            Tensor(shape=tensor_shape, other=0, dtype=input_dtype),
+            Tensor(shape=tensor_shape, other=0, dtype=gate_dtype),
+            Tensor(shape=tensor_shape, dtype=weight_dtype),
+            Tensor(0, dtype=ninetoothed.float64),
+            Tensor(shape=tensor_shape, dtype=output_dtype),
+            Tensor(0, dtype=ninetoothed.float64),
+        )
+    else:
+        arrangement_ = functools.partial(
+            arrangement_no_gate,
+            group_size=group_size,
+            block_size=max(block_size, hidden_size),
+        )
+        tensors = (
+            Tensor(shape=tensor_shape, other=0, dtype=input_dtype),
+            Tensor(shape=tensor_shape, dtype=weight_dtype),
+            Tensor(0, dtype=ninetoothed.float64),
+            Tensor(shape=tensor_shape, dtype=output_dtype),
+            Tensor(0, dtype=ninetoothed.float64),
+        )
+
+    if not has_gate:
+        application = application_no_gate
+    elif activation == "sigmoid":
+        application = (
+            application_sigmoid_before
+            if norm_before_gate
+            else application_sigmoid_after
+        )
+    else:
+        application = (
+            application_silu_before if norm_before_gate else application_silu_after
+        )
+
+    return arrangement_, application, tensors
diff --git a/src/ntops/torch/__init__.py b/src/ntops/torch/__init__.py
index ad6fd4c..8bc0400 100644
--- a/src/ntops/torch/__init__.py
+++ b/src/ntops/torch/__init__.py
@@ -1,11 +1,20 @@
 from ntops.torch.abs import abs
+from ntops.torch.acosh import acosh
+from ntops.torch.adaptive_avg_pool2d import adaptive_avg_pool2d
+from ntops.torch.adaptive_max_pool2d import adaptive_max_pool2d
 from ntops.torch.add import add
 from ntops.torch.addmm import addmm
+from ntops.torch.addmv import addmv
 from ntops.torch.alpha_dropout import alpha_dropout
+from ntops.torch.argsort import argsort
+from ntops.torch.atan import atan
 from ntops.torch.avg_pool2d import avg_pool2d
+from ntops.torch.batch_norm import batch_norm
+from ntops.torch.bincount import bincount
 from ntops.torch.bitwise_and import bitwise_and
 from ntops.torch.bitwise_not import bitwise_not
 from ntops.torch.bitwise_or import bitwise_or
+from ntops.torch.block_scaled_fp8_mm import block_scaled_fp8_mm
 from ntops.torch.bmm import bmm
 from ntops.torch.celu import celu
 from ntops.torch.clamp import clamp
@@ -17,6 +26,11 @@
 from ntops.torch.dropout import dropout
 from ntops.torch.eq import eq
 from ntops.torch.exp import exp
+from ntops.torch.fmax import fmax
+from ntops.torch.fused_mla_rope_cache_write import (
+    fused_mla_rope_cache_write,
+    fused_mla_rope_cache_write_reference,
+)
 from ntops.torch.ge import ge
 from ntops.torch.gelu import gelu
 from ntops.torch.gt import gt
@@ -25,18 +39,34 @@
 from ntops.torch.isnan import isnan
 from ntops.torch.layer_norm import layer_norm
 from ntops.torch.le import le
+from ntops.torch.logsumexp import logsumexp
+from ntops.torch.lp_pool1d import lp_pool1d
+from ntops.torch.lp_pool2d import lp_pool2d
+from ntops.torch.lp_pool3d import lp_pool3d
 from ntops.torch.lt import lt
 from ntops.torch.matmul import matmul
+from ntops.torch.max import max
+from ntops.torch.max_pool1d import max_pool1d
 from ntops.torch.max_pool2d import max_pool2d
+from ntops.torch.max_pool3d import max_pool3d
+from ntops.torch.maximum import maximum
+from ntops.torch.mean import mean
+from ntops.torch.median import median
 from ntops.torch.mm import mm
 from ntops.torch.msort import msort
 from ntops.torch.mul import mul
+from ntops.torch.mxfp4_w4a16_grouped_mm import (
+    ScalingType,
+    SwizzleType,
+    mxfp4_w4a16_grouped_mm,
+)
 from ntops.torch.ne import ne
 from ntops.torch.neg import neg
 from ntops.torch.pow import pow
 from ntops.torch.quantile import quantile
 from ntops.torch.relu import relu
 from ntops.torch.rms_norm import rms_norm
+from ntops.torch.rms_norm_gated import rms_norm_gated
 from ntops.torch.rot90 import rot90
 from ntops.torch.rotary_position_embedding import rotary_position_embedding
 from ntops.torch.round import round
@@ -51,29 +81,10 @@
 from ntops.torch.sin import sin
 from ntops.torch.softmax import softmax
 from ntops.torch.sort import sort
+from ntops.torch.stack import stack
 from ntops.torch.sub import sub
 from ntops.torch.tanh import tanh
 from ntops.torch.threshold import threshold
-from ntops.torch.max_pool1d import max_pool1d
-from ntops.torch.max_pool3d import max_pool3d
-from ntops.torch.stack import stack
-from ntops.torch.mean import mean
-from ntops.torch.median import median
-from ntops.torch.maximum import maximum
-from ntops.torch.atan import atan
-from ntops.torch.batch_norm import batch_norm
-from ntops.torch.bincount import bincount
-from ntops.torch.adaptive_max_pool2d import adaptive_max_pool2d
-from ntops.torch.acosh import acosh
-from ntops.torch.adaptive_avg_pool2d import adaptive_avg_pool2d
-from ntops.torch.addmv import addmv
-from ntops.torch.argsort import argsort
-from ntops.torch.fmax import fmax
-from ntops.torch.logsumexp import logsumexp
-from ntops.torch.lp_pool1d import lp_pool1d
-from ntops.torch.lp_pool2d import lp_pool2d
-from ntops.torch.lp_pool3d import lp_pool3d
-from ntops.torch.max import max
 
 __all__ = [
     "abs",
@@ -115,11 +126,16 @@
     "quantile",
     "relu",
     "rms_norm",
+    "rms_norm_gated",
     "rot90",
     "rotary_position_embedding",
     "round",
     "rsqrt",
     "scaled_dot_product_attention",
+    "ScalingType",
+    "SwizzleType",
+    "mxfp4_w4a16_grouped_mm",
+    "block_scaled_fp8_mm",
     "select_copy",
     "sgn",
     "sigmoid",
@@ -130,7 +146,7 @@
     "softmax",
     "sort",
     "sub",
-    "tanh",    
+    "tanh",
     "max_pool1d",
     "max_pool3d",
     "stack",
@@ -153,4 +169,6 @@
     "lp_pool2d",
     "lp_pool3d",
     "max",
+    "fused_mla_rope_cache_write",
+    "fused_mla_rope_cache_write_reference",
 ]
diff --git a/src/ntops/torch/block_scaled_fp8_mm.py b/src/ntops/torch/block_scaled_fp8_mm.py
new file mode 100644
index 0000000..9e6105d
--- /dev/null
+++ b/src/ntops/torch/block_scaled_fp8_mm.py
@@ -0,0 +1,299 @@
+import torch
+
+import ntops
+from ntops.torch.mxfp4_w4a16_grouped_mm import ScalingType, SwizzleType
+from ntops.torch.utils import _cached_make
+
+_SCALE_BLOCK_SIZE = 128
+_SUPPORTED_INPUT_DTYPE = torch.float8_e4m3fn
+
+
+def _kernel_launch_config(device):
+    if torch.version.hip is not None:
+        return 32, 1
+
+    is_corex = bool(getattr(torch, "corex", False))
+    has_cuda = torch.version.cuda is not None and torch.cuda.is_available()
+    if not is_corex and has_cuda:
+        capability = torch.cuda.get_device_capability(device)
+        # NVIDIA's native E4M3 Tensor Core path starts at SM89.
+        if capability >= (8, 9):
+            return 32, 4
+
+    return 16, 4
+
+
+def _kernel_tuning_config(device):
+    block_size_k, num_warps = _kernel_launch_config(device)
+    # gfx936 only compiles the one-wave specialization reliably. Multi-wave
+    # candidates crash AMD make_amdgcn, which an in-process tuner cannot catch.
+    return block_size_k, num_warps, 1, 1
+
+
+def _make_kernel(input_dtype, output_dtype, bias_dtype, device):
+    block_size_k, num_warps, num_stages, max_num_configs = _kernel_tuning_config(device)
+    return _cached_make(
+        ntops.kernels.block_scaled_fp8_mm.premake,
+        input_dtype,
+        output_dtype,
+        bias_dtype,
+        block_size_m=16,
+        block_size_n=16,
+        block_size_k=block_size_k,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        max_num_configs=max_num_configs,
+    )
+
+
+def _enum_matches(value, enum_type, member_name):
+    member = getattr(enum_type, member_name)
+    return value == member or getattr(value, "name", None) == member_name
+
+
+def _require_tensor(name, value):
+    if not isinstance(value, torch.Tensor):
+        raise TypeError(f"{name} must be a torch.Tensor")
+
+
+def _unwrap_single_level(name, value):
+    if not isinstance(value, (tuple, list)):
+        return value
+    if len(value) != 1:
+        raise NotImplementedError(f"{name} must contain exactly one level")
+    return value[0]
+
+
+def _is_dense_matrix(value):
+    return value.is_contiguous() or value.t().is_contiguous()
+
+
+def _validate_options(
+    scale_recipe_a,
+    scale_recipe_b,
+    swizzle_a,
+    swizzle_b,
+    output_dtype,
+    contraction_dim,
+    use_fast_accum,
+):
+    if not _enum_matches(scale_recipe_a, ScalingType, "BlockWise1x128"):
+        raise NotImplementedError("scale_recipe_a must be ScalingType.BlockWise1x128")
+    if not (
+        _enum_matches(scale_recipe_b, ScalingType, "BlockWise1x128")
+        or _enum_matches(scale_recipe_b, ScalingType, "BlockWise128x128")
+    ):
+        raise NotImplementedError(
+            "scale_recipe_b must be ScalingType.BlockWise1x128 or "
+            "ScalingType.BlockWise128x128"
+        )
+    for name, value in (("swizzle_a", swizzle_a), ("swizzle_b", swizzle_b)):
+        if value is not None and not _enum_matches(value, SwizzleType, "NO_SWIZZLE"):
+            raise NotImplementedError(f"{name} must be SwizzleType.NO_SWIZZLE")
+    if output_dtype is None:
+        output_dtype = torch.bfloat16
+    if output_dtype not in (torch.bfloat16, torch.float16, torch.float32):
+        raise ValueError(
+            "output_dtype must be torch.bfloat16, torch.float16, or torch.float32"
+        )
+    if contraction_dim not in (None, (), []):
+        raise NotImplementedError("contraction_dim is not supported")
+    if use_fast_accum:
+        raise NotImplementedError("use_fast_accum=True is not supported")
+    return output_dtype
+
+
+def _validate_inputs(mat_a, mat_b, scale_a, scale_b, scale_recipe_b, bias):
+    for name, value in (
+        ("mat_a", mat_a),
+        ("mat_b", mat_b),
+        ("scale_a", scale_a),
+        ("scale_b", scale_b),
+    ):
+        _require_tensor(name, value)
+
+    if mat_a.dtype != _SUPPORTED_INPUT_DTYPE or mat_b.dtype != _SUPPORTED_INPUT_DTYPE:
+        raise TypeError("mat_a and mat_b must have dtype torch.float8_e4m3fn")
+    if scale_a.dtype != torch.float32 or scale_b.dtype != torch.float32:
+        raise TypeError("scale_a and scale_b must have dtype torch.float32")
+    if mat_a.ndim != 2 or mat_b.ndim != 2:
+        raise ValueError("mat_a and mat_b must both be 2D tensors")
+    if scale_a.ndim != 2 or scale_b.ndim != 2:
+        raise ValueError("scale_a and scale_b must both be 2D tensors")
+    if not mat_a.is_contiguous():
+        raise ValueError("mat_a must be row-major contiguous")
+    if not mat_b.t().is_contiguous():
+        raise ValueError("mat_b must be column-major (a transposed contiguous matrix)")
+    if not _is_dense_matrix(scale_a) or not _is_dense_matrix(scale_b):
+        raise ValueError(
+            "scale_a and scale_b must have a dense row- or column-major layout"
+        )
+
+    m, k = mat_a.shape
+    mat_b_k, n = mat_b.shape
+    if mat_b_k != k:
+        raise ValueError(
+            f"mat_a.shape[1] must equal mat_b.shape[0], got {k} and {mat_b_k}"
+        )
+    if k == 0 or k % _SCALE_BLOCK_SIZE != 0:
+        raise ValueError("K must be a positive multiple of 128")
+    if n % _SCALE_BLOCK_SIZE != 0:
+        raise ValueError("N must be a multiple of 128")
+
+    k_blocks = k // _SCALE_BLOCK_SIZE
+    n_blocks = n // _SCALE_BLOCK_SIZE
+    if scale_a.shape != (m, k_blocks):
+        raise ValueError(
+            f"scale_a must have shape {(m, k_blocks)}, got {tuple(scale_a.shape)}"
+        )
+
+    b_uses_1x128 = _enum_matches(scale_recipe_b, ScalingType, "BlockWise1x128")
+    if b_uses_1x128:
+        expected_scale_b_shape = (n, k_blocks)
+        if scale_b.shape != expected_scale_b_shape:
+            raise ValueError(
+                f"scale_b must have shape {expected_scale_b_shape}, "
+                f"got {tuple(scale_b.shape)}"
+            )
+    else:
+        padded_k_blocks = ((k_blocks + 3) // 4) * 4
+        valid_shapes = ((k_blocks, n_blocks), (padded_k_blocks, n_blocks))
+        if scale_b.shape not in valid_shapes:
+            raise ValueError(
+                "scale_b must have shape "
+                f"{valid_shapes[0]} or padded shape {valid_shapes[1]}, "
+                f"got {tuple(scale_b.shape)}"
+            )
+
+    device = mat_a.device
+    if any(value.device != device for value in (mat_b, scale_a, scale_b)):
+        raise ValueError(
+            "mat_a, mat_b, scale_a, and scale_b must be on the same device"
+        )
+
+    if bias is not None:
+        _require_tensor("bias", bias)
+        if bias.device != device:
+            raise ValueError("bias must be on the same device as the inputs")
+        if bias.dtype not in (torch.bfloat16, torch.float16, torch.float32):
+            raise TypeError(
+                "bias must have dtype torch.bfloat16, torch.float16, or torch.float32"
+            )
+        if bias.shape != (n,):
+            raise ValueError(f"bias must have shape {(n,)}, got {tuple(bias.shape)}")
+        if not bias.is_contiguous():
+            raise ValueError("bias must be contiguous")
+
+    return m, n, k_blocks, n_blocks, b_uses_1x128
+
+
+def _group_views(
+    mat_a,
+    mat_b,
+    scale_a,
+    scale_b,
+    output,
+    bias,
+    k_blocks,
+    n_blocks,
+    b_uses_1x128,
+):
+    m, k = mat_a.shape
+    mat_a_groups = mat_a.unsqueeze(0).expand(n_blocks, -1, -1)
+    mat_b_groups = mat_b.t().reshape(n_blocks, _SCALE_BLOCK_SIZE, k)
+    mat_b_groups = mat_b_groups.transpose(-2, -1)
+    scale_a_groups = scale_a.unsqueeze(0).expand(n_blocks, -1, -1)
+
+    if b_uses_1x128:
+        scale_b_groups = scale_b.reshape(n_blocks, _SCALE_BLOCK_SIZE, k_blocks).permute(
+            0, 2, 1
+        )
+    else:
+        scale_b_groups = scale_b[:k_blocks].t().unsqueeze(-1)
+        scale_b_groups = scale_b_groups.expand(-1, -1, _SCALE_BLOCK_SIZE)
+
+    output_groups = output.view(m, n_blocks, _SCALE_BLOCK_SIZE).permute(1, 0, 2)
+    bias_groups = None
+    if bias is not None:
+        bias_groups = bias.view(n_blocks, 1, _SCALE_BLOCK_SIZE)
+
+    return (
+        mat_a_groups,
+        mat_b_groups,
+        scale_a_groups,
+        scale_b_groups,
+        bias_groups,
+        output_groups,
+    )
+
+
+def block_scaled_fp8_mm(
+    mat_a,
+    mat_b,
+    scale_a,
+    scale_recipe_a,
+    scale_b,
+    scale_recipe_b,
+    swizzle_a=None,
+    swizzle_b=None,
+    bias=None,
+    output_dtype=torch.bfloat16,
+    contraction_dim=(),
+    use_fast_accum=False,
+):
+    """Block-scaled FP8 matrix multiplication for LLM projections."""
+    scale_a = _unwrap_single_level("scale_a", scale_a)
+    scale_recipe_a = _unwrap_single_level("scale_recipe_a", scale_recipe_a)
+    scale_b = _unwrap_single_level("scale_b", scale_b)
+    scale_recipe_b = _unwrap_single_level("scale_recipe_b", scale_recipe_b)
+    swizzle_a = _unwrap_single_level("swizzle_a", swizzle_a)
+    swizzle_b = _unwrap_single_level("swizzle_b", swizzle_b)
+
+    output_dtype = _validate_options(
+        scale_recipe_a,
+        scale_recipe_b,
+        swizzle_a,
+        swizzle_b,
+        output_dtype,
+        contraction_dim,
+        use_fast_accum,
+    )
+    m, n, k_blocks, n_blocks, b_uses_1x128 = _validate_inputs(
+        mat_a, mat_b, scale_a, scale_b, scale_recipe_b, bias
+    )
+    output = torch.empty((m, n), dtype=output_dtype, device=mat_a.device)
+    if output.numel() == 0:
+        return output
+
+    grouped = _group_views(
+        mat_a,
+        mat_b,
+        scale_a,
+        scale_b,
+        output,
+        bias,
+        k_blocks,
+        n_blocks,
+        b_uses_1x128,
+    )
+    (
+        mat_a_groups,
+        mat_b_groups,
+        scale_a_groups,
+        scale_b_groups,
+        bias_groups,
+        output_groups,
+    ) = grouped
+
+    kernel = _make_kernel(
+        mat_a.dtype,
+        output_dtype,
+        bias.dtype if bias is not None else None,
+        mat_a.device,
+    )
+    arguments = (mat_a_groups, mat_b_groups, scale_a_groups, scale_b_groups)
+    if bias_groups is None:
+        kernel(*arguments, output_groups)
+    else:
+        kernel(*arguments, bias_groups, output_groups)
+    return output
diff --git a/src/ntops/torch/fused_mla_rope_cache_write.py b/src/ntops/torch/fused_mla_rope_cache_write.py
new file mode 100644
index 0000000..f071d62
--- /dev/null
+++ b/src/ntops/torch/fused_mla_rope_cache_write.py
@@ -0,0 +1,204 @@
+"""PyTorch interface for fused MLA RoPE + compressed cache writes."""
+
+import torch
+
+import ntops
+from ntops.torch.utils import _cached_make
+
+
+def _validate_inputs(
+    kv_c,
+    k_pe,
+    kv_cache,
+    slot_mapping,
+    positions,
+    cos_sin_cache,
+):
+    if kv_c.ndim != 2:
+        raise ValueError("kv_c must have shape [num_tokens, kv_lora_rank]")
+    if k_pe.ndim not in (2, 3):
+        raise ValueError(
+            "k_pe must have shape [num_tokens, rope_dim] or [num_tokens, 1, rope_dim]"
+        )
+    if kv_cache.ndim != 3:
+        raise ValueError("kv_cache must have shape [num_blocks, block_size, entry_dim]")
+    if slot_mapping.ndim != 1 or positions.ndim != 1:
+        raise ValueError("slot_mapping and positions must be one-dimensional")
+    if cos_sin_cache.ndim != 2:
+        raise ValueError("cos_sin_cache must have shape [max_position, rope_dim]")
+
+    rope_dim = k_pe.shape[-1]
+    num_tokens = slot_mapping.shape[0]
+    if rope_dim <= 0 or rope_dim % 2:
+        raise ValueError("rope_dim must be a positive even number")
+    if kv_c.shape[1] <= 0:
+        raise ValueError("kv_c must have a positive latent width")
+    if kv_c.shape[0] < num_tokens or k_pe.shape[0] < num_tokens:
+        raise ValueError(
+            "source tensors must contain at least slot_mapping.size(0) tokens"
+        )
+    if positions.shape[0] != num_tokens:
+        raise ValueError("positions must have one value per slot_mapping entry")
+    if k_pe.ndim == 3 and k_pe.shape[1] != 1:
+        raise ValueError("rank-3 k_pe must have shape [num_tokens, 1, rope_dim]")
+    if k_pe.numel() < num_tokens * rope_dim:
+        raise ValueError("k_pe does not contain enough token values")
+    if cos_sin_cache.shape[1] != rope_dim:
+        raise ValueError(
+            "cos_sin_cache width must equal packed rope_dim (cos[0:R/2] | sin[0:R/2])"
+        )
+    if kv_cache.shape[2] != kv_c.shape[1] + rope_dim:
+        raise ValueError(
+            "kv_cache entry dimension must equal latent width + rope width"
+        )
+    if kv_cache.shape[0] <= 0 or kv_cache.shape[1] <= 0:
+        raise ValueError("kv_cache must have positive block and block-size dimensions")
+    if kv_c.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise TypeError("MLA inputs and cache must use float16, bfloat16, or float32")
+    if slot_mapping.dtype != torch.int64:
+        raise TypeError("slot_mapping must be torch.int64")
+    if positions.dtype not in (torch.int32, torch.int64):
+        raise TypeError("positions must be torch.int32 or torch.int64")
+    tensors = (kv_c, k_pe, kv_cache, slot_mapping, positions, cos_sin_cache)
+    if any(t.device != kv_c.device for t in tensors):
+        raise ValueError("all tensors must be on the same device")
+    if any(t.dtype != kv_c.dtype for t in (k_pe, kv_cache)):
+        raise TypeError("kv_c, k_pe, and kv_cache must share dtype")
+    if cos_sin_cache.dtype not in (kv_c.dtype, torch.float32):
+        raise TypeError("cos_sin_cache must use the input dtype or torch.float32")
+
+
+def fused_mla_rope_cache_write(
+    kv_c,
+    k_pe,
+    kv_cache,
+    slot_mapping,
+    positions,
+    cos_sin_cache,
+    *,
+    block_size=128,
+    num_warps=(1, 2, 4, 8),
+    num_stages=(1, 2),
+    max_num_configs=8,
+):
+    """Fuse RoPE(k_pe) with an in-place compressed MLA KV-cache write.
+
+    The cache entry written for each non-padding slot is
+    ``[kv_c[token] | RoPE(k_pe[token])]``.  ``kv_c`` remains the shared
+    low-rank K/V latent; it is never expanded to per-head K/V.  This operation
+    returns ``None`` like vLLM's ``concat_and_cache_mla``.  By default Triton
+    autotunes the launch configuration and caches the winner by input key.
+    """
+    _validate_inputs(kv_c, k_pe, kv_cache, slot_mapping, positions, cos_sin_cache)
+    if not kv_c.is_cuda:
+        raise RuntimeError("fused_mla_rope_cache_write requires a CUDA device")
+    if (
+        not isinstance(block_size, int)
+        or isinstance(block_size, bool)
+        or block_size <= 0
+    ):
+        raise ValueError("block_size must be a positive integer")
+
+    if k_pe.ndim == 3:
+        k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1])
+    if positions.dtype == torch.int32:
+        positions = positions.to(torch.int64)
+
+    num_tokens = slot_mapping.shape[0]
+    if num_tokens == 0:
+        return None
+    # Keep indirect source accesses on one contiguous specialization. Some
+    # Triton backends cannot compile the non-unit-stride indexed-store variant.
+    cache_output = kv_cache
+    kv_c, k_pe, slot_mapping, positions, cos_sin_cache = (
+        tensor if tensor.is_contiguous() else tensor.contiguous()
+        for tensor in (kv_c, k_pe, slot_mapping, positions, cos_sin_cache)
+    )
+    if not kv_cache.is_contiguous():
+        kv_cache = kv_cache.contiguous()
+
+    latent = kv_c.shape[1]
+    rope_dim = k_pe.shape[-1]
+    entry_dim = latent + rope_dim
+    tile_size = max(
+        1 << (block_size - 1).bit_length(),
+        1 << (entry_dim - 1).bit_length(),
+    )
+    # A real output root is required by the NineToothed SSA path on DCU. It is
+    # one compressed entry per token, not a per-head query tensor, and is
+    # discarded after anchoring the cache side effect.
+    output_anchor = torch.empty(
+        (num_tokens, entry_dim), dtype=kv_c.dtype, device=kv_c.device
+    )
+    kernel = _cached_make(
+        ntops.kernels.fused_mla_rope_cache_write.premake,
+        latent,
+        rope_dim,
+        dtype=kv_c.dtype,
+        block_size=tile_size,
+        cache_block_size=kv_cache.shape[1],
+        cos_dtype=cos_sin_cache.dtype,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        max_num_configs=max_num_configs,
+    )
+    kernel(
+        output_anchor,
+        kv_c,
+        k_pe,
+        kv_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        entry_dim,
+        tile_size,
+        latent,
+        rope_dim,
+        kv_cache.shape[1],
+    )
+    if kv_cache is not cache_output:
+        cache_output.copy_(kv_cache)
+    return None
+
+
+def fused_mla_rope_cache_write_reference(
+    kv_c,
+    k_pe,
+    kv_cache,
+    slot_mapping,
+    positions,
+    cos_sin_cache,
+):
+    """PyTorch reference for the fused cache-only operation."""
+    _validate_inputs(kv_c, k_pe, kv_cache, slot_mapping, positions, cos_sin_cache)
+    if k_pe.ndim == 3:
+        k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1])
+    num_tokens = slot_mapping.shape[0]
+    if num_tokens == 0:
+        return None
+    rope_dim = k_pe.shape[-1]
+    half = rope_dim // 2
+    table = cos_sin_cache.index_select(0, positions.to(torch.long))
+    cos = table[:, :half].to(torch.float32)
+    sin = table[:, half:].to(torch.float32)
+    x0 = k_pe[:num_tokens, 0::2].to(torch.float32)
+    x1 = k_pe[:num_tokens, 1::2].to(torch.float32)
+    rotated = torch.empty_like(k_pe[:num_tokens], dtype=torch.float32)
+    rotated[:, 0::2] = x0 * cos - x1 * sin
+    rotated[:, 1::2] = x0 * sin + x1 * cos
+    rotated = rotated.to(kv_c.dtype)
+
+    for token_idx, slot in enumerate(slot_mapping.tolist()):
+        if slot < 0:
+            continue
+        block_idx, block_offset = divmod(slot, kv_cache.shape[1])
+        cache_base = kv_cache[block_idx, block_offset]
+        cache_base[: kv_c.shape[1]] = kv_c[token_idx]
+        cache_base[kv_c.shape[1] :] = rotated[token_idx]
+    return None
+
+
+__all__ = [
+    "fused_mla_rope_cache_write",
+    "fused_mla_rope_cache_write_reference",
+]
diff --git a/src/ntops/torch/mxfp4_w4a16_grouped_mm.py b/src/ntops/torch/mxfp4_w4a16_grouped_mm.py
new file mode 100644
index 0000000..b167635
--- /dev/null
+++ b/src/ntops/torch/mxfp4_w4a16_grouped_mm.py
@@ -0,0 +1,292 @@
+import enum
+
+import torch
+import torch.nn.functional as F
+
+import ntops
+from ntops.torch.utils import _cached_make
+
+
+class _ScalingType(enum.IntEnum):
+    TensorWise = 0
+    RowWise = 1
+    BlockWise1x16 = 2
+    BlockWise1x32 = 3
+    BlockWise1x128 = 4
+    BlockWise128x128 = 5
+
+
+class _SwizzleType(enum.IntEnum):
+    NO_SWIZZLE = 0
+    SWIZZLE_32_4_4 = 1
+
+
+ScalingType = getattr(F, "ScalingType", _ScalingType)
+SwizzleType = getattr(F, "SwizzleType", _SwizzleType)
+# Multi-wave launches regress the routed gfx936 specialization by more than an
+# order of magnitude. Four waves still wins the uniform cases.
+_HIP_AUTOTUNE_WARPS = (1, 4)
+
+
+def _kernel_launch_config(jagged=False):
+    if torch.version.hip is not None:
+        if jagged:
+            return True, 1, 16, 1, 1, 1
+        return True, 1, 16, _HIP_AUTOTUNE_WARPS, 1, len(_HIP_AUTOTUNE_WARPS)
+    return False, 16, 64, 4, 1, 1
+
+
+def _make_kernel(jagged):
+    reduction, block_size_m, block_size_n, num_warps, num_stages, limit = (
+        _kernel_launch_config(jagged)
+    )
+    return _cached_make(
+        ntops.kernels.mxfp4_w4a16_grouped_mm.premake,
+        jagged,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        reduction=reduction,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        max_num_configs=limit,
+    )
+
+
+def _dtype_if_available(name):
+    return getattr(torch, name, None)
+
+
+def _is_dtype(dtype, expected):
+    return dtype is expected or (expected is not None and dtype == expected)
+
+
+def _enum_matches(value, enum_type, member_name):
+    member = getattr(enum_type, member_name)
+
+    if value == member:
+        return True
+
+    return getattr(value, "name", None) == member_name
+
+
+def _require_tensor(name, value):
+    if not isinstance(value, torch.Tensor):
+        raise TypeError(f"{name} must be a torch.Tensor")
+
+
+def _require_contiguous(name, value):
+    if not value.is_contiguous():
+        raise ValueError(f"{name} must be contiguous")
+
+
+def _unwrap_single_level(name, value):
+    if not isinstance(value, (tuple, list)):
+        return value
+
+    if len(value) != 1:
+        raise NotImplementedError(f"{name} must contain exactly one level")
+
+    return value[0]
+
+
+def _as_uint8(value):
+    if value.dtype == torch.uint8:
+        return value
+
+    return value.view(torch.uint8)
+
+
+def _validate_optional_arguments(
+    scale_a,
+    scale_recipe_a,
+    scale_recipe_b,
+    swizzle_a,
+    swizzle_b,
+    bias,
+    output_dtype,
+    contraction_dim,
+    use_fast_accum,
+):
+    if scale_a is not None:
+        raise NotImplementedError("scale_a is not supported for W4A16")
+
+    if scale_recipe_a is not None:
+        raise NotImplementedError("scale_recipe_a is not supported for W4A16")
+
+    if not _enum_matches(scale_recipe_b, ScalingType, "BlockWise1x32"):
+        raise ValueError("scale_recipe_b must be ScalingType.BlockWise1x32")
+
+    if swizzle_a is not None and not _enum_matches(
+        swizzle_a, SwizzleType, "NO_SWIZZLE"
+    ):
+        raise NotImplementedError("swizzle_a is not supported")
+
+    if swizzle_b is not None and not _enum_matches(
+        swizzle_b, SwizzleType, "NO_SWIZZLE"
+    ):
+        raise NotImplementedError("swizzled inputs are not supported")
+
+    if bias is not None:
+        raise NotImplementedError("bias is not supported")
+
+    if output_dtype not in (None, torch.bfloat16):
+        raise ValueError("output_dtype must be torch.bfloat16")
+
+    if contraction_dim not in (None, (), []):
+        raise NotImplementedError("contraction_dim is not supported")
+
+    if use_fast_accum:
+        raise NotImplementedError("use_fast_accum=True is not supported")
+
+
+def _validate_inputs(mat_a, mat_b, scale_b, offs):
+    _require_tensor("mat_a", mat_a)
+    _require_tensor("mat_b", mat_b)
+    _require_tensor("scale_b", scale_b)
+
+    if mat_a.dtype != torch.bfloat16:
+        raise TypeError("mat_a must have dtype torch.bfloat16")
+
+    packed_dtype = _dtype_if_available("float4_e2m1fn_x2")
+    if mat_b.dtype != torch.uint8 and not _is_dtype(mat_b.dtype, packed_dtype):
+        raise TypeError("mat_b must have dtype torch.uint8 or torch.float4_e2m1fn_x2")
+
+    scale_dtype = _dtype_if_available("float8_e8m0fnu")
+    if scale_b.dtype != torch.uint8 and not _is_dtype(scale_b.dtype, scale_dtype):
+        raise TypeError("scale_b must have dtype torch.uint8 or torch.float8_e8m0fnu")
+
+    if mat_b.ndim != 3:
+        raise ValueError("mat_b must have shape (G, K // 2, N)")
+
+    if scale_b.ndim != 3:
+        raise ValueError("scale_b must have shape (G, K // 32, N)")
+
+    _require_contiguous("mat_a", mat_a)
+    _require_contiguous("mat_b", mat_b)
+    _require_contiguous("scale_b", scale_b)
+
+    group_count, packed_k, n = mat_b.shape
+    if group_count == 0:
+        raise ValueError("mat_b must contain at least one group")
+
+    k = packed_k * 2
+    if k == 0 or k % 32 != 0:
+        raise ValueError("the logical K dimension must be a positive multiple of 32")
+
+    if scale_b.shape != (group_count, k // 32, n):
+        raise ValueError(
+            f"scale_b must have shape {(group_count, k // 32, n)}, "
+            f"but got {tuple(scale_b.shape)}"
+        )
+
+    device = mat_a.device
+    if mat_b.device != device or scale_b.device != device:
+        raise ValueError("mat_a, mat_b, and scale_b must be on the same device")
+
+    if offs is None:
+        if mat_a.ndim != 3:
+            raise ValueError("mat_a must have shape (G, M, K) when offs is None")
+
+        if mat_a.shape[0] != group_count or mat_a.shape[2] != k:
+            raise ValueError(
+                f"mat_a must have shape (G, M, K) with G={group_count} and K={k}"
+            )
+
+        return False, group_count, mat_a.shape[1], k, n
+
+    _require_tensor("offs", offs)
+    if mat_a.ndim != 2:
+        raise ValueError("mat_a must have shape (total_M, K) when offs is provided")
+
+    if mat_a.shape[1] != k:
+        raise ValueError(f"mat_a must have K={k}, but got K={mat_a.shape[1]}")
+
+    if offs.ndim != 1 or offs.numel() != group_count:
+        raise ValueError(f"offs must have shape ({group_count},)")
+
+    if offs.dtype != torch.int32:
+        raise TypeError("offs must have dtype torch.int32")
+
+    if offs.device != device:
+        raise ValueError("offs must be on the same device as the inputs")
+
+    _require_contiguous("offs", offs)
+
+    if bool(torch.any(offs < 0).item()):
+        raise ValueError("offs must contain non-negative cumulative row counts")
+
+    if group_count > 1 and bool(torch.any(offs[1:] < offs[:-1]).item()):
+        raise ValueError("offs must be nondecreasing")
+
+    if int(offs[-1].item()) != mat_a.shape[0]:
+        raise ValueError("offs[-1] must equal mat_a.shape[0]")
+
+    return True, group_count, mat_a.shape[0], k, n
+
+
+def mxfp4_w4a16_grouped_mm(
+    mat_a,
+    mat_b,
+    scale_a,
+    scale_recipe_a,
+    scale_b,
+    scale_recipe_b,
+    swizzle_a=None,
+    swizzle_b=None,
+    bias=None,
+    offs=None,
+    output_dtype=torch.bfloat16,
+    contraction_dim=(),
+    use_fast_accum=False,
+):
+    """Compute grouped MXFP4 W4A16 expert matrix multiplication."""
+    scale_b = _unwrap_single_level("scale_b", scale_b)
+    scale_recipe_b = _unwrap_single_level("scale_recipe_b", scale_recipe_b)
+    swizzle_a = _unwrap_single_level("swizzle_a", swizzle_a)
+    swizzle_b = _unwrap_single_level("swizzle_b", swizzle_b)
+
+    _validate_optional_arguments(
+        scale_a,
+        scale_recipe_a,
+        scale_recipe_b,
+        swizzle_a,
+        swizzle_b,
+        bias,
+        output_dtype,
+        contraction_dim,
+        use_fast_accum,
+    )
+    jagged, group_count, m, _, n = _validate_inputs(mat_a, mat_b, scale_b, offs)
+
+    output_shape = (group_count, m, n) if not jagged else (m, n)
+    output = torch.empty(output_shape, dtype=torch.bfloat16, device=mat_a.device)
+    if output.numel() == 0:
+        return output
+
+    # The block-dot decoder lowers to BF16 MMAC on AMD and crashes gfx936
+    # codegen. Use a branchless decoder with an ordinary FP32 reduction on HIP.
+    kernel = _make_kernel(jagged)
+    mat_b_uint8 = _as_uint8(mat_b)
+    scale_b_uint8 = _as_uint8(scale_b)
+    mat_a_even = mat_a[..., :-1]
+    mat_a_odd = mat_a[..., 1:]
+
+    common_args = (mat_a_even, mat_a_odd, mat_b_uint8, scale_b_uint8)
+
+    if not jagged:
+        kernel(*common_args, output)
+        return output
+
+    offsets = torch.cat((offs.new_zeros(1), offs))
+    mat_a_even_jagged = torch.nested.nested_tensor_from_jagged(mat_a_even, offsets)
+    mat_a_odd_jagged = torch.nested.nested_tensor_from_jagged(mat_a_odd, offsets)
+    output_jagged = torch.nested.nested_tensor_from_jagged(output, offsets)
+    jagged_args = (
+        mat_a_even_jagged,
+        mat_a_odd_jagged,
+        *common_args[2:],
+        output_jagged,
+    )
+    kernel(*jagged_args)
+
+    return output
diff --git a/src/ntops/torch/rms_norm_gated.py b/src/ntops/torch/rms_norm_gated.py
new file mode 100644
index 0000000..be8e0f7
--- /dev/null
+++ b/src/ntops/torch/rms_norm_gated.py
@@ -0,0 +1,81 @@
+import torch
+
+import ntops
+from ntops.torch.utils import _cached_make
+
+
+def rms_norm_gated(
+    input,
+    z=None,
+    weight=None,
+    eps=1e-5,
+    group_size=None,
+    norm_before_gate=False,
+    activation="swish",
+    block_size=128,
+    num_warps=(1, 2, 4, 8),
+    num_stages=(1, 2),
+    max_num_configs=8,
+):
+    """RMSNormGated with a specialized layout and shape-keyed autotuning.
+
+    The default candidates target the vLLM GDN hot path and cache the best
+    launch configuration for each input shape, dtype, and stride signature.
+    """
+    if input.ndim == 0:
+        raise ValueError("input must have at least one dimension")
+
+    if activation not in ("silu", "sigmoid", "swish"):
+        raise ValueError("activation must be one of 'silu', 'sigmoid', or 'swish'")
+
+    hidden_size = input.shape[-1]
+    if group_size is not None:
+        if not isinstance(group_size, int) or isinstance(group_size, bool):
+            raise TypeError("group_size must be a positive integer or None")
+        if group_size <= 0:
+            raise ValueError("group_size must be a positive integer")
+        if hidden_size % group_size != 0:
+            raise ValueError("hidden size must be divisible by group_size")
+        num_normalized_elements = group_size
+    else:
+        num_normalized_elements = hidden_size
+
+    if weight is None:
+        weight = torch.ones(
+            hidden_size,
+            device=input.device,
+            dtype=input.dtype,
+        )
+    weight = weight.expand_as(input)
+
+    has_gate = z is not None
+    gate_dtype = input.dtype
+    if has_gate:
+        z = z.expand_as(input)
+        gate_dtype = z.dtype
+
+    output = torch.empty_like(input)
+    kernel = _cached_make(
+        ntops.kernels.rms_norm_gated.premake,
+        input.ndim,
+        hidden_size,
+        group_size,
+        norm_before_gate,
+        activation,
+        input_dtype=input.dtype,
+        gate_dtype=gate_dtype,
+        weight_dtype=weight.dtype,
+        output_dtype=output.dtype,
+        block_size=block_size,
+        has_gate=has_gate,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        max_num_configs=max_num_configs,
+    )
+
+    if has_gate:
+        kernel(input, z, weight, eps, output, num_normalized_elements)
+    else:
+        kernel(input, weight, eps, output, num_normalized_elements)
+
+    return output
diff --git a/tests/test_block_scaled_fp8_mm.py b/tests/test_block_scaled_fp8_mm.py
new file mode 100644
index 0000000..74c2526
--- /dev/null
+++ b/tests/test_block_scaled_fp8_mm.py
@@ -0,0 +1,322 @@
+import importlib
+import inspect
+import pathlib
+
+import ninetoothed
+import pytest
+import torch
+import torch.nn.functional as F
+
+import ntops
+
+skip_if_cuda_is_unavailable = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="CUDA-compatible accelerator is unavailable",
+)
+
+
+def _column_major(value):
+    return value.t().contiguous().t()
+
+
+def _make_inputs(m, n, k, recipe_b, output_dtype=torch.bfloat16, bias=False):
+    device = "cuda"
+    k_blocks = k // 128
+    n_blocks = n // 128
+    mat_a = torch.randn((m, k), device=device).clamp(-3, 3).to(torch.float8_e4m3fn)
+    weight = torch.randn((n, k), device=device).clamp(-3, 3).to(torch.float8_e4m3fn)
+    mat_b = weight.t()
+    scale_a_logical = torch.rand((m, k_blocks), device=device) + 0.25
+    scale_a = _column_major(scale_a_logical)
+
+    if recipe_b == ntops.torch.ScalingType.BlockWise1x128:
+        scale_b_logical = torch.rand((n, k_blocks), device=device) + 0.25
+        scale_b = _column_major(scale_b_logical)
+        weight_dequantized = (
+            weight.float().reshape(n, k_blocks, 128) * scale_b_logical[..., None]
+        ).reshape(n, k)
+    else:
+        scale_b_logical = torch.rand((n_blocks, k_blocks), device=device) + 0.25
+        padded_k_blocks = ((k_blocks + 3) // 4) * 4
+        scale_b_padded = F.pad(scale_b_logical, (0, padded_k_blocks - k_blocks))
+        scale_b = scale_b_padded.t()
+        weight_dequantized = (
+            weight.float().reshape(n_blocks, 128, k_blocks, 128)
+            * scale_b_logical[:, None, :, None]
+        ).reshape(n, k)
+
+    mat_a_dequantized = (
+        mat_a.float().reshape(m, k_blocks, 128) * scale_a_logical[..., None]
+    ).reshape(m, k)
+    bias_value = None
+    if bias:
+        bias_value = torch.randn((n,), device=device, dtype=output_dtype)
+    expected = mat_a_dequantized @ weight_dequantized.t()
+    if bias_value is not None:
+        expected = expected + bias_value.float()
+    return mat_a, mat_b, scale_a, scale_b, bias_value, expected.to(output_dtype)
+
+
+def _block_scaled_fp8_mm(mat_a, mat_b, scale_a, scale_b, recipe_b, **kwargs):
+    return ntops.torch.block_scaled_fp8_mm(
+        mat_a,
+        mat_b,
+        scale_a,
+        ntops.torch.ScalingType.BlockWise1x128,
+        scale_b,
+        recipe_b,
+        **kwargs,
+    )
+
+
+@pytest.mark.parametrize(
+    ("hip_version", "is_corex", "cuda_version", "capability", "expected"),
+    (
+        ("6.3", False, None, None, (32, 1)),
+        (None, True, "10.2", (9, 0), (16, 4)),
+        (None, False, "12.4", (8, 9), (32, 4)),
+        (None, False, "12.4", (9, 0), (32, 4)),
+        (None, False, "12.4", (8, 6), (16, 4)),
+        (None, False, None, None, (16, 4)),
+    ),
+)
+def test_block_scaled_fp8_mm_selects_dot_path_by_backend_capability(
+    monkeypatch,
+    hip_version,
+    is_corex,
+    cuda_version,
+    capability,
+    expected,
+):
+    module = importlib.import_module("ntops.torch.block_scaled_fp8_mm")
+    monkeypatch.setattr(torch.version, "hip", hip_version)
+    monkeypatch.setattr(torch.version, "cuda", cuda_version)
+    monkeypatch.setattr(torch, "corex", is_corex, raising=False)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: cuda_version is not None)
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_capability",
+        lambda device: capability,
+    )
+
+    assert module._kernel_launch_config(torch.device("cuda")) == expected
+
+
+def test_block_scaled_fp8_mm_keeps_fixed_hip_config(monkeypatch):
+    module = importlib.import_module("ntops.torch.block_scaled_fp8_mm")
+    monkeypatch.setattr(torch.version, "hip", "6.3")
+
+    assert module._kernel_tuning_config(torch.device("cuda")) == (32, 1, 1, 1)
+
+
+def test_block_scaled_fp8_mm_keeps_fixed_non_hip_config(monkeypatch):
+    module = importlib.import_module("ntops.torch.block_scaled_fp8_mm")
+    monkeypatch.setattr(torch.version, "hip", None)
+    monkeypatch.setattr(torch.version, "cuda", None)
+    monkeypatch.setattr(torch, "corex", True, raising=False)
+
+    assert module._kernel_tuning_config(torch.device("cuda")) == (16, 4, 1, 1)
+
+
+@skip_if_cuda_is_unavailable
+@pytest.mark.parametrize(
+    "recipe_b",
+    (
+        ntops.torch.ScalingType.BlockWise1x128,
+        ntops.torch.ScalingType.BlockWise128x128,
+    ),
+)
+def test_block_scaled_fp8_mm_block_fp8_numerics(recipe_b):
+    torch.manual_seed(0)
+    has_bias = recipe_b == ntops.torch.ScalingType.BlockWise128x128
+    inputs = _make_inputs(17, 256, 256, recipe_b, bias=has_bias)
+    mat_a, mat_b, scale_a, scale_b, bias, expected = inputs
+    output = _block_scaled_fp8_mm(mat_a, mat_b, scale_a, scale_b, recipe_b, bias=bias)
+
+    assert output.shape == expected.shape
+    assert output.dtype == torch.bfloat16
+    torch.testing.assert_close(output, expected, rtol=0.03, atol=0.125)
+
+
+@skip_if_cuda_is_unavailable
+def test_block_scaled_fp8_mm_attention_decode_with_bias_and_fp32_output():
+    torch.manual_seed(1)
+    recipe_b = ntops.torch.ScalingType.BlockWise128x128
+    inputs = _make_inputs(1, 128, 128, recipe_b, torch.float32, bias=True)
+    mat_a, mat_b, scale_a, scale_b, bias, expected = inputs
+    output = _block_scaled_fp8_mm(
+        mat_a,
+        mat_b,
+        [scale_a],
+        [scale_b],
+        recipe_b,
+        bias=bias,
+        output_dtype=torch.float32,
+        swizzle_a=[ntops.torch.SwizzleType.NO_SWIZZLE],
+        swizzle_b=ntops.torch.SwizzleType.NO_SWIZZLE,
+    )
+
+    assert output.dtype == torch.float32
+    torch.testing.assert_close(output, expected, rtol=0.01, atol=0.1)
+
+
+@skip_if_cuda_is_unavailable
+def test_block_scaled_fp8_mm_covers_the_full_128_element_scale_block():
+    recipe_a = ntops.torch.ScalingType.BlockWise1x128
+    recipe_b = ntops.torch.ScalingType.BlockWise128x128
+    mat_a = torch.zeros((1, 4, 32), device="cuda")
+    mat_a[:, :, 16:] = 1
+    mat_a = mat_a.reshape(1, 128).to(torch.float8_e4m3fn)
+    column_values = torch.arange(128, device="cuda") % 7 - 3
+    weight = column_values[:, None].expand(-1, 128)
+    mat_b = weight.to(torch.float8_e4m3fn).t()
+    scale_a = torch.ones((1, 1), device="cuda")
+    scale_b = torch.ones((1, 1), device="cuda")
+
+    output = ntops.torch.block_scaled_fp8_mm(
+        mat_a,
+        mat_b,
+        scale_a,
+        recipe_a,
+        scale_b,
+        recipe_b,
+        output_dtype=torch.float32,
+    )
+
+    expected = (column_values * 64).to(torch.float32).unsqueeze(0)
+    torch.testing.assert_close(output, expected)
+
+
+def _cpu_inputs(m=2, n=128, k=128):
+    mat_a = torch.empty((m, k), dtype=torch.float8_e4m3fn)
+    mat_b = torch.empty((n, k), dtype=torch.float8_e4m3fn).t()
+    scale_a = _column_major(torch.ones((m, k // 128)))
+    scale_b = torch.ones((k // 128, n // 128))
+    return mat_a, mat_b, scale_a, scale_b
+
+
+def test_block_scaled_fp8_mm_public_name_and_signature():
+    expected = (
+        "mat_a",
+        "mat_b",
+        "scale_a",
+        "scale_recipe_a",
+        "scale_b",
+        "scale_recipe_b",
+        "swizzle_a",
+        "swizzle_b",
+        "bias",
+        "output_dtype",
+        "contraction_dim",
+        "use_fast_accum",
+    )
+    assert not hasattr(ntops.torch, "scaled_mm")
+    assert (
+        tuple(inspect.signature(ntops.torch.block_scaled_fp8_mm).parameters) == expected
+    )
+
+
+def test_block_scaled_fp8_mm_accepts_list_api_and_empty_m():
+    mat_a, mat_b, scale_a, scale_b = _cpu_inputs(m=0)
+    output = ntops.torch.block_scaled_fp8_mm(
+        mat_a,
+        mat_b,
+        [scale_a],
+        [ntops.torch.ScalingType.BlockWise1x128],
+        [scale_b],
+        [ntops.torch.ScalingType.BlockWise128x128],
+        swizzle_a=[ntops.torch.SwizzleType.NO_SWIZZLE],
+        swizzle_b=[ntops.torch.SwizzleType.NO_SWIZZLE],
+        output_dtype=None,
+        contraction_dim=[],
+    )
+    assert output.shape == (0, mat_b.shape[1])
+    assert output.dtype == torch.bfloat16
+
+
+@pytest.mark.parametrize(
+    ("argument", "value", "error"),
+    (
+        (
+            "scale_recipe_a",
+            ntops.torch.ScalingType.TensorWise,
+            NotImplementedError,
+        ),
+        (
+            "scale_recipe_b",
+            ntops.torch.ScalingType.RowWise,
+            NotImplementedError,
+        ),
+        (
+            "swizzle_a",
+            ntops.torch.SwizzleType.SWIZZLE_32_4_4,
+            NotImplementedError,
+        ),
+        (
+            "scale_recipe_b",
+            [
+                ntops.torch.ScalingType.BlockWise128x128,
+                ntops.torch.ScalingType.BlockWise128x128,
+            ],
+            NotImplementedError,
+        ),
+        ("output_dtype", torch.int32, ValueError),
+        ("contraction_dim", (1,), NotImplementedError),
+        ("use_fast_accum", True, NotImplementedError),
+    ),
+)
+def test_block_scaled_fp8_mm_rejects_unsupported_options(argument, value, error):
+    mat_a, mat_b, scale_a, scale_b = _cpu_inputs()
+    arguments = {
+        "scale_recipe_a": ntops.torch.ScalingType.BlockWise1x128,
+        "scale_recipe_b": ntops.torch.ScalingType.BlockWise128x128,
+    }
+    arguments[argument] = value
+    with pytest.raises(error):
+        ntops.torch.block_scaled_fp8_mm(
+            mat_a, mat_b, scale_a, scale_b=scale_b, **arguments
+        )
+
+
+def test_block_scaled_fp8_mm_validates_shapes_dtypes_and_layouts():
+    mat_a, mat_b, scale_a, scale_b = _cpu_inputs()
+    recipe_b = ntops.torch.ScalingType.BlockWise128x128
+
+    with pytest.raises(TypeError, match="float8_e4m3fn"):
+        _block_scaled_fp8_mm(mat_a.float(), mat_b, scale_a, scale_b, recipe_b)
+    with pytest.raises(ValueError, match="scale_a must have shape"):
+        _block_scaled_fp8_mm(mat_a, mat_b, scale_a[:, :0], scale_b, recipe_b)
+    with pytest.raises(ValueError, match="column-major"):
+        _block_scaled_fp8_mm(mat_a, mat_b.contiguous(), scale_a, scale_b, recipe_b)
+    bad_n = torch.empty((129, 128), dtype=torch.float8_e4m3fn).t()
+    with pytest.raises(ValueError, match="N must be a multiple"):
+        _block_scaled_fp8_mm(mat_a, bad_n, scale_a, scale_b, recipe_b)
+
+
+def _generated_kernel_source():
+    kernel = ninetoothed.make(
+        *ntops.kernels.block_scaled_fp8_mm.premake(torch.float8_e4m3fn, torch.bfloat16),
+        num_warps=1,
+        num_stages=1,
+        max_num_configs=1,
+    )
+    if hasattr(kernel, "_compilation"):
+        sources = (
+            value
+            for name, value in kernel._compilation.artifact.sources.items()
+            if name.endswith(".py")
+        )
+        return "\n".join(str(value) for value in sources)
+    return pathlib.Path(kernel._source).read_text()
+
+
+def test_block_scaled_fp8_mm_lowers_to_portable_fp8_dot():
+    source = _generated_kernel_source()
+    dot_count = source.count("tl.dot(") + source.count("triton.language.dot(")
+    assert dot_count == 1
+    assert (
+        "def block_scaled_fp8_mm_k32(" in source
+        or "def block_scaled_fp8_mm_k32_kernel(" in source
+    )
+    assert "dot_scaled" not in source
+    assert "ntl." not in source
diff --git a/tests/test_fused_mla_rope_cache_write.py b/tests/test_fused_mla_rope_cache_write.py
new file mode 100644
index 0000000..471187a
--- /dev/null
+++ b/tests/test_fused_mla_rope_cache_write.py
@@ -0,0 +1,198 @@
+import pytest
+import torch
+
+import ntops
+from tests.skippers import skip_if_cuda_not_available
+
+
+def _tables(max_position, rope_dim, dtype, device):
+    return torch.randn(max_position, rope_dim, dtype=dtype, device=device)
+
+
+def test_fused_mla_rope_cache_write_validates_contract():
+    kv_c = torch.randn(2, 8)
+    k_pe = torch.randn(2, 8)
+    cache = torch.randn(2, 4, 16)
+    slots = torch.tensor((0, 1), dtype=torch.int64)
+    positions = torch.tensor((0, 1), dtype=torch.int64)
+    table = _tables(4, 8, torch.float32, "cpu")
+
+    with pytest.raises(TypeError, match="float16, bfloat16, or float32"):
+        ntops.torch.fused_mla_rope_cache_write(
+            kv_c.double(),
+            k_pe.double(),
+            cache.double(),
+            slots,
+            positions,
+            table,
+        )
+
+    with pytest.raises(ValueError, match="positive even"):
+        ntops.torch.fused_mla_rope_cache_write(
+            kv_c,
+            torch.randn(2, 7),
+            torch.randn(2, 4, 15),
+            slots,
+            positions,
+            torch.randn(4, 7),
+        )
+
+    with pytest.raises(TypeError, match="slot_mapping"):
+        ntops.torch.fused_mla_rope_cache_write(
+            kv_c,
+            k_pe,
+            cache,
+            slots.to(torch.int32),
+            positions,
+            table,
+        )
+
+
+@skip_if_cuda_not_available
+@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16, torch.float32))
+@pytest.mark.parametrize("kpe_rank3", (False, True))
+def test_fused_mla_rope_cache_write(dtype, kpe_rank3):
+    device = "cuda"
+    tokens, latent, rope = 3, 8, 8
+    kv_c = torch.randn(tokens, latent, device=device, dtype=dtype)
+    k_pe = torch.randn(tokens, rope, device=device, dtype=dtype)
+    if kpe_rank3:
+        k_pe = k_pe.unsqueeze(1)
+    kv_cache = torch.randn(6, 4, latent + rope, device=device, dtype=dtype)
+    reference_cache = kv_cache.clone()
+    slot_mapping = torch.tensor((0, -1, 7), device=device, dtype=torch.int64)
+    positions = torch.tensor((0, 2, 5), device=device, dtype=torch.int64)
+    cos_sin_cache = _tables(16, rope, torch.float32, device)
+
+    result = ntops.torch.fused_mla_rope_cache_write(
+        kv_c,
+        k_pe,
+        kv_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        block_size=8,
+        num_warps=1,
+        num_stages=1,
+    )
+    reference = ntops.torch.fused_mla_rope_cache_write_reference(
+        kv_c,
+        k_pe,
+        reference_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+    )
+
+    assert result is None
+    assert reference is None
+    atol = 2e-3 if dtype != torch.bfloat16 else 2e-2
+    assert torch.allclose(kv_cache, reference_cache, rtol=atol, atol=atol)
+
+
+@skip_if_cuda_not_available
+def test_fused_mla_rope_cache_write_supports_padding_source_rows():
+    device = "cuda"
+    source_tokens, latent, rope = 4, 8, 8
+    kv_c = torch.randn(source_tokens, latent, device=device, dtype=torch.float16)
+    k_pe = torch.randn(source_tokens, 1, rope, device=device, dtype=torch.float16)
+    cache = torch.zeros(4, 4, latent + rope, device=device, dtype=torch.float16)
+    reference_cache = cache.clone()
+    slots = torch.tensor((0, 5), device=device, dtype=torch.int64)
+    positions = torch.tensor((1, 3), device=device, dtype=torch.int64)
+    table = _tables(8, rope, torch.float32, device)
+
+    ntops.torch.fused_mla_rope_cache_write(
+        kv_c,
+        k_pe,
+        cache,
+        slots,
+        positions,
+        table,
+        block_size=8,
+        num_warps=1,
+        num_stages=1,
+    )
+    ntops.torch.fused_mla_rope_cache_write_reference(
+        kv_c,
+        k_pe,
+        reference_cache,
+        slots,
+        positions,
+        table,
+    )
+    assert torch.allclose(cache, reference_cache, rtol=2e-3, atol=2e-3)
+
+
+@skip_if_cuda_not_available
+def test_fused_mla_rope_cache_write_strided_inputs():
+    device = "cuda"
+    tokens, latent, rope = 2, 8, 8
+    kv_c = torch.randn(tokens, latent * 2, device=device, dtype=torch.float16)[:, ::2]
+    k_pe = torch.randn(tokens, rope * 2, device=device, dtype=torch.float16)[:, ::2]
+    cache = torch.randn(4, 4, (latent + rope) * 2, device=device, dtype=torch.float16)[
+        ..., ::2
+    ]
+    reference_cache = cache.clone()
+    slots = torch.tensor((0, 5), device=device, dtype=torch.int64)
+    positions = torch.tensor((0, 1), device=device, dtype=torch.int64)
+    table = _tables(16, rope, torch.float32, device)
+
+    result = ntops.torch.fused_mla_rope_cache_write(
+        kv_c,
+        k_pe,
+        cache,
+        slots,
+        positions,
+        table,
+        block_size=128,
+        num_warps=(1, 2, 4, 8),
+        num_stages=(1, 2),
+        max_num_configs=8,
+    )
+    reference = ntops.torch.fused_mla_rope_cache_write_reference(
+        kv_c,
+        k_pe,
+        reference_cache,
+        slots,
+        positions,
+        table,
+    )
+    assert result is None
+    assert reference is None
+    assert torch.allclose(cache, reference_cache, rtol=2e-3, atol=2e-3)
+
+
+@skip_if_cuda_not_available
+def test_fused_mla_rope_cache_write_output_anchor():
+    source_tokens, latent, rope = 4, 8, 8
+    kv_c = torch.randn(source_tokens, latent, device="cuda", dtype=torch.float16)
+    k_pe = torch.randn(source_tokens, rope, device="cuda", dtype=torch.float16)
+    cache = torch.randn(4, 4, latent + rope, device="cuda", dtype=torch.float16)
+    reference_cache = cache.clone()
+    slots = torch.tensor((0, -1, 5), device="cuda", dtype=torch.int64)
+    positions = torch.tensor((0, 1, 2), device="cuda", dtype=torch.int64)
+    table = _tables(8, rope, torch.float32, "cuda")
+
+    result = ntops.torch.fused_mla_rope_cache_write(
+        kv_c,
+        k_pe,
+        cache,
+        slots,
+        positions,
+        table,
+        block_size=8,
+        num_warps=1,
+        num_stages=1,
+    )
+    ntops.torch.fused_mla_rope_cache_write_reference(
+        kv_c,
+        k_pe,
+        reference_cache,
+        slots,
+        positions,
+        table,
+    )
+
+    assert result is None
+    assert torch.allclose(cache, reference_cache, rtol=2e-3, atol=2e-3)
diff --git a/tests/test_mxfp4_w4a16_grouped_mm.py b/tests/test_mxfp4_w4a16_grouped_mm.py
new file mode 100644
index 0000000..6a57e84
--- /dev/null
+++ b/tests/test_mxfp4_w4a16_grouped_mm.py
@@ -0,0 +1,312 @@
+import importlib
+import pathlib
+
+import ninetoothed
+import pytest
+import torch
+
+import ntops
+
+skip_if_cuda_is_unavailable = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="CUDA-compatible accelerator is unavailable",
+)
+native_dtypes = pytest.param(
+    True,
+    marks=pytest.mark.skipif(
+        not hasattr(torch, "float4_e2m1fn_x2") or not hasattr(torch, "float8_e8m0fnu"),
+        reason="native MXFP4 dtypes are unavailable",
+    ),
+)
+
+
+def test_mxfp4_w4a16_grouped_mm_enables_bounded_hip_autotuning(monkeypatch):
+    module = importlib.import_module("ntops.torch.mxfp4_w4a16_grouped_mm")
+    monkeypatch.setattr(torch.version, "hip", "6.3")
+
+    assert module._kernel_launch_config() == (
+        True,
+        1,
+        16,
+        (1, 4),
+        1,
+        2,
+    )
+
+
+def test_mxfp4_w4a16_grouped_mm_keeps_fixed_hip_jagged_config(monkeypatch):
+    module = importlib.import_module("ntops.torch.mxfp4_w4a16_grouped_mm")
+    monkeypatch.setattr(torch.version, "hip", "6.3")
+
+    assert module._kernel_launch_config(jagged=True) == (True, 1, 16, 1, 1, 1)
+
+
+def test_mxfp4_w4a16_grouped_mm_keeps_fixed_non_hip_config(monkeypatch):
+    module = importlib.import_module("ntops.torch.mxfp4_w4a16_grouped_mm")
+    monkeypatch.setattr(torch.version, "hip", None)
+
+    assert module._kernel_launch_config() == (False, 16, 64, 4, 1, 1)
+
+
+def _make_mxfp4_weight(group_count, k, n, device):
+    codes = torch.randint(0, 16, (group_count, k, n), dtype=torch.uint8, device=device)
+    packed = codes[:, 0::2] | (codes[:, 1::2] << 4)
+    scales = torch.randint(
+        124,
+        128,
+        (group_count, k // 32, n),
+        dtype=torch.uint8,
+        device=device,
+    )
+
+    return codes, packed.contiguous(), scales.contiguous()
+
+
+def _decode_mxfp4(codes, scales):
+    magnitude_code = codes & 0x7
+    exponent = (magnitude_code >> 1).to(torch.int32)
+    mantissa = (magnitude_code & 1).to(torch.float32)
+    normal = (1.0 + 0.5 * mantissa) * torch.exp2(exponent.to(torch.float32) - 1.0)
+    magnitude = torch.where(exponent == 0, 0.5 * mantissa, normal)
+    sign = torch.where((codes & 0x8) == 0, 1.0, -1.0)
+    block_scales = torch.exp2(scales.to(torch.float32) - 127.0)
+    block_scales = block_scales.repeat_interleave(32, dim=1)
+
+    return (sign * magnitude * block_scales).to(torch.bfloat16)
+
+
+def _mxfp4_w4a16_grouped_mm(mat_a, mat_b, scale_b, offs=None):
+    return ntops.torch.mxfp4_w4a16_grouped_mm(
+        mat_a,
+        mat_b,
+        None,
+        None,
+        scale_b,
+        ntops.torch.ScalingType.BlockWise1x32,
+        offs=offs,
+    )
+
+
+@skip_if_cuda_is_unavailable
+@pytest.mark.parametrize("native_dtypes", (False, native_dtypes))
+def test_mxfp4_w4a16_grouped_mm_uniform(native_dtypes):
+    torch.manual_seed(0)
+    group_count, m, k, n = 2, 17, 96, 19
+    mat_a = (
+        0.25 * torch.randn((group_count, m, k), dtype=torch.bfloat16, device="cuda")
+    ).contiguous()
+    codes, mat_b, scale_b = _make_mxfp4_weight(group_count, k, n, mat_a.device)
+
+    if native_dtypes:
+        mat_b = mat_b.view(torch.float4_e2m1fn_x2)
+        scale_b = scale_b.view(torch.float8_e8m0fnu)
+
+    output = _mxfp4_w4a16_grouped_mm(mat_a, mat_b, scale_b)
+    weight = _decode_mxfp4(codes, scale_b.view(torch.uint8))
+    expected = torch.bmm(mat_a.float(), weight.float()).to(torch.bfloat16)
+
+    assert output.dtype == torch.bfloat16
+    assert output.shape == (group_count, m, n)
+    torch.testing.assert_close(output, expected, rtol=0.03, atol=0.03)
+
+
+@skip_if_cuda_is_unavailable
+@pytest.mark.parametrize("native_dtypes", (False, native_dtypes))
+def test_mxfp4_w4a16_grouped_mm_jagged_with_zero_token_expert(native_dtypes):
+    torch.manual_seed(1)
+    group_count, k, n = 3, 96, 23
+    expert_rows = (4, 0, 7)
+    total_m = sum(expert_rows)
+    mat_a = (
+        0.25 * torch.randn((total_m, k), dtype=torch.bfloat16, device="cuda")
+    ).contiguous()
+    codes, mat_b, scale_b = _make_mxfp4_weight(group_count, k, n, mat_a.device)
+    offs = torch.tensor((4, 4, 11), dtype=torch.int32, device=mat_a.device)
+
+    if native_dtypes:
+        mat_b = mat_b.view(torch.float4_e2m1fn_x2)
+        scale_b = scale_b.view(torch.float8_e8m0fnu)
+
+    output = _mxfp4_w4a16_grouped_mm(mat_a, mat_b, scale_b, offs=offs)
+    weight = _decode_mxfp4(codes, scale_b.view(torch.uint8))
+    expected_parts = []
+    start = 0
+    for expert, end in enumerate(offs.cpu().tolist()):
+        expected_parts.append(
+            (mat_a[start:end].float() @ weight[expert].float()).to(torch.bfloat16)
+        )
+        start = end
+    expected = torch.cat(expected_parts)
+
+    assert output.dtype == torch.bfloat16
+    assert output.shape == (total_m, n)
+    torch.testing.assert_close(output, expected, rtol=0.03, atol=0.03)
+
+
+@skip_if_cuda_is_unavailable
+def test_mxfp4_w4a16_grouped_mm_decodes_every_e2m1_code_and_nibble_position():
+    codes = torch.arange(16, dtype=torch.uint8, device="cuda").repeat(2)
+    codes = codes.reshape(1, 32, 1)
+    mat_b = (codes[:, 0::2] | (codes[:, 1::2] << 4)).contiguous()
+    scale_b = torch.full((1, 1, 1), 127, dtype=torch.uint8, device="cuda")
+    mat_a = torch.eye(32, dtype=torch.bfloat16, device="cuda").unsqueeze(0)
+
+    output = _mxfp4_w4a16_grouped_mm(mat_a, mat_b, scale_b)
+    expected = _decode_mxfp4(codes, scale_b)
+
+    torch.testing.assert_close(output, expected, rtol=0, atol=0)
+
+
+def _generated_kernel_source(reduction=False):
+    kernel = ninetoothed.make(
+        *ntops.kernels.mxfp4_w4a16_grouped_mm.premake(False, reduction=reduction),
+        max_num_configs=1,
+    )
+
+    if hasattr(kernel, "_compilation"):
+        sources = (
+            value
+            for name, value in kernel._compilation.artifact.sources.items()
+            if name.endswith(".py")
+        )
+        source = "\n".join(str(value) for value in sources)
+    else:
+        source = pathlib.Path(kernel._source).read_text()
+
+    return source
+
+
+def test_mxfp4_w4a16_grouped_mm_lowers_to_portable_dot_pair():
+    source = _generated_kernel_source()
+
+    dot_count = source.count("tl.dot(") + source.count("triton.language.dot(")
+    assert dot_count == 2
+    assert (
+        "def mxfp4_w4a16_grouped_mm_dot(" in source
+        or "def mxfp4_w4a16_grouped_mm_dot_kernel(" in source
+    )
+    assert "dot_scaled" not in source
+    assert "ntl." not in source
+
+
+def test_mxfp4_w4a16_grouped_mm_hip_reduction_lowers_without_mmac_or_branches():
+    source = _generated_kernel_source(reduction=True)
+
+    dot_count = source.count("tl.dot(") + source.count("triton.language.dot(")
+    sum_count = source.count("tl.sum(") + source.count("triton.language.sum(")
+    reduction_loop_count = source.count("_body_i in range(0, 16, 1):")
+    assert dot_count == 0
+    assert sum_count == 2 or reduction_loop_count == 2
+    assert (
+        "def mxfp4_w4a16_grouped_mm_reduction(" in source
+        or "def mxfp4_w4a16_grouped_mm_reduction_kernel(" in source
+    )
+    assert "dot_scaled" not in source
+    assert "ntl." not in source
+    assert "tl.where(" not in source
+    assert "decode_even" not in source
+    assert "decode_odd" not in source
+
+
+def _cpu_inputs(group_count=3, total_m=None):
+    k, n = 32, 5
+    if total_m is None:
+        mat_a = torch.zeros((group_count, 2, k), dtype=torch.bfloat16)
+    else:
+        mat_a = torch.zeros((total_m, k), dtype=torch.bfloat16)
+    mat_b = torch.zeros((group_count, k // 2, n), dtype=torch.uint8)
+    scale_b = torch.full((group_count, k // 32, n), 127, dtype=torch.uint8)
+
+    return mat_a, mat_b, scale_b
+
+
+@pytest.mark.parametrize(
+    ("argument", "value", "error"),
+    (
+        ("scale_a", torch.ones(1), NotImplementedError),
+        ("scale_recipe_a", ntops.torch.ScalingType.TensorWise, NotImplementedError),
+        (
+            "scale_recipe_b",
+            [
+                ntops.torch.ScalingType.BlockWise1x32,
+                ntops.torch.ScalingType.BlockWise1x32,
+            ],
+            NotImplementedError,
+        ),
+        ("swizzle_a", ntops.torch.SwizzleType.SWIZZLE_32_4_4, NotImplementedError),
+        ("bias", torch.ones(1), NotImplementedError),
+        ("output_dtype", torch.float16, ValueError),
+        ("contraction_dim", (0,), NotImplementedError),
+        ("use_fast_accum", True, NotImplementedError),
+    ),
+)
+def test_mxfp4_w4a16_grouped_mm_rejects_unsupported_options(argument, value, error):
+    mat_a, mat_b, scale_b = _cpu_inputs()
+    arguments = {
+        "scale_a": None,
+        "scale_recipe_a": None,
+        "scale_recipe_b": ntops.torch.ScalingType.BlockWise1x32,
+    }
+    arguments[argument] = value
+
+    with pytest.raises(error):
+        ntops.torch.mxfp4_w4a16_grouped_mm(
+            mat_a,
+            mat_b,
+            scale_b=scale_b,
+            **arguments,
+        )
+
+
+def test_mxfp4_w4a16_grouped_mm_validates_shapes_and_offsets():
+    mat_a, mat_b, scale_b = _cpu_inputs()
+
+    with pytest.raises(TypeError, match="mat_a must have dtype"):
+        _mxfp4_w4a16_grouped_mm(mat_a.float(), mat_b, scale_b)
+
+    with pytest.raises(ValueError, match="scale_b must have shape"):
+        _mxfp4_w4a16_grouped_mm(mat_a, mat_b, scale_b[:, :, :-1].contiguous())
+
+    jagged_a, mat_b, scale_b = _cpu_inputs(total_m=5)
+
+    with pytest.raises(TypeError, match="offs must have dtype"):
+        _mxfp4_w4a16_grouped_mm(jagged_a, mat_b, scale_b, torch.tensor((2, 2, 5)))
+
+    with pytest.raises(ValueError, match="offs must be nondecreasing"):
+        _mxfp4_w4a16_grouped_mm(
+            jagged_a,
+            mat_b,
+            scale_b,
+            torch.tensor((3, 2, 5), dtype=torch.int32),
+        )
+
+    with pytest.raises(ValueError, match="offs\\[-1\\]"):
+        _mxfp4_w4a16_grouped_mm(
+            jagged_a,
+            mat_b,
+            scale_b,
+            torch.tensor((2, 2, 4), dtype=torch.int32),
+        )
+
+
+def test_mxfp4_w4a16_grouped_mm_accepts_supported_option_forms():
+    mat_a, mat_b, scale_b = _cpu_inputs(group_count=1, total_m=0)
+    offs = torch.tensor((0,), dtype=torch.int32)
+
+    output = ntops.torch.mxfp4_w4a16_grouped_mm(
+        mat_a,
+        mat_b,
+        None,
+        None,
+        [scale_b],
+        [ntops.torch.ScalingType.BlockWise1x32],
+        swizzle_a=[ntops.torch.SwizzleType.NO_SWIZZLE],
+        swizzle_b=ntops.torch.SwizzleType.NO_SWIZZLE,
+        offs=offs,
+        output_dtype=None,
+        contraction_dim=[],
+    )
+
+    assert output.shape == (0, mat_b.shape[-1])
+    assert output.dtype == torch.bfloat16
+    assert not hasattr(ntops.torch, "scaled_grouped_mm")
diff --git a/tests/test_rms_norm_gated.py b/tests/test_rms_norm_gated.py
new file mode 100644
index 0000000..b95e956
--- /dev/null
+++ b/tests/test_rms_norm_gated.py
@@ -0,0 +1,111 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+import ntops
+from tests.skippers import skip_if_cuda_not_available
+
+
+def _reference_rms_norm_gated(
+    input,
+    z,
+    weight,
+    eps,
+    group_size,
+    norm_before_gate,
+    activation,
+):
+    x = input.float()
+    weight = weight.float()
+    z = None if z is None else z.float()
+    act_fn = torch.sigmoid if activation == "sigmoid" else F.silu
+
+    if z is not None and not norm_before_gate:
+        x = x * act_fn(z)
+
+    if group_size is None:
+        variance = x.square().mean(dim=-1, keepdim=True)
+        output = x * torch.rsqrt(variance + eps) * weight
+    else:
+        x_group = x.reshape(*x.shape[:-1], -1, group_size)
+        variance = x_group.square().mean(dim=-1, keepdim=True)
+        output = (x_group * torch.rsqrt(variance + eps)).reshape_as(x) * weight
+
+    if z is not None and norm_before_gate:
+        output = output * act_fn(z)
+
+    return output.to(input.dtype)
+
+
+@skip_if_cuda_not_available
+@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16))
+@pytest.mark.parametrize("group_size", (None, 4))
+@pytest.mark.parametrize("norm_before_gate", (False, True))
+@pytest.mark.parametrize("activation", ("swish", "sigmoid"))
+def test_rms_norm_gated(dtype, group_size, norm_before_gate, activation):
+    input = torch.randn((2, 3, 8), dtype=dtype, device="cuda")
+    z = torch.randn_like(input)
+    weight = torch.randn((8,), dtype=dtype, device="cuda")
+
+    output = ntops.torch.rms_norm_gated(
+        input,
+        z,
+        weight,
+        group_size=group_size,
+        norm_before_gate=norm_before_gate,
+        activation=activation,
+    )
+    reference = _reference_rms_norm_gated(
+        input,
+        z,
+        weight,
+        1e-5,
+        group_size,
+        norm_before_gate,
+        activation,
+    )
+
+    torch.testing.assert_close(output, reference, rtol=2e-3, atol=2e-3)
+
+
+@skip_if_cuda_not_available
+def test_rms_norm_gated_without_gate_or_weight():
+    input = torch.randn((2, 8), dtype=torch.float32, device="cuda")
+
+    output = ntops.torch.rms_norm_gated(input)
+    reference = torch.rsqrt(input.square().mean(dim=-1, keepdim=True) + 1e-5) * input
+
+    torch.testing.assert_close(output, reference, rtol=2e-3, atol=2e-3)
+
+
+@skip_if_cuda_not_available
+def test_rms_norm_gated_hidden_size_larger_than_block_size():
+    input = torch.randn((2, 257), dtype=torch.float32, device="cuda")
+    z = torch.randn_like(input)
+    weight = torch.randn((257,), dtype=torch.float32, device="cuda")
+
+    output = ntops.torch.rms_norm_gated(input, z, weight, block_size=128)
+    reference = _reference_rms_norm_gated(
+        input,
+        z,
+        weight,
+        1e-5,
+        None,
+        False,
+        "swish",
+    )
+
+    torch.testing.assert_close(output, reference, rtol=2e-3, atol=2e-3)
+
+
+def test_rms_norm_gated_validates_arguments():
+    input = torch.randn((2, 8))
+
+    with pytest.raises(ValueError, match="activation"):
+        ntops.torch.rms_norm_gated(input, activation="gelu")
+
+    with pytest.raises(ValueError, match="positive integer"):
+        ntops.torch.rms_norm_gated(input, group_size=0)
+
+    with pytest.raises(ValueError, match="divisible"):
+        ntops.torch.rms_norm_gated(input, group_size=3)