[KMCompiler][MetaX] Token-tiled fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert

The generic kernel runs one program per (token, head) slot at num_warps=1 --
64 threads for 512 elements on a 64-lane warp. At 131072 tokens x 64 heads
that is 8.5M programs of 64 threads. Two costs follow: the KV slot does seven
quantisation blocks against a Q slot's normalise-and-rotate, so one slot per
program makes it a straggler, and 64-thread blocks under-fill the SM. The
straggler is 1-in-65 at 64 heads but 1-in-129 at 128, which matches the
measured shortfall being worse at the lower head count.

Raising num_warps alone does not help and measurably hurts (1134 -> 691 ->
398 -> 199 GB/s at 8192x64): work per program stays at 512 elements, so each
lane gets less to do while the program occupies more of the machine. Work and
width have to rise together. Each program here handles TPP=8 tokens of ONE
slot, keeping every program uniformly all-Q or all-KV -- no divergence -- at
num_warps=4, which is 16 elements per lane instead of 8. The (TPP, num_warps)
optimum runs along a diagonal; TPP=2/w1, 4/w2 and 8/w4 are all near the top,
so what matters is elements per lane rather than warps.

Below 512 tokens the generic kernel is faster and is dispatched to: TPP=8
masks off 7/8 of every program there, and the wider blocks raise the launch
floor from ~26us to ~32us. The crossover was measured at 512 on both head
counts (256 tokens: 0.96x/0.99x, 512 tokens: 1.03x/1.04x). TPP=8/w4 is best at
every size above it, so one configuration suffices.

Scoped to the MetaX backend rather than changing the generic op: all tuning is
C550-only and the tiling is matched to 64-lane warps, where num_warps=4 is 256
threads. On NVIDIA's 32-lane warps it is 128, a different trade-off entirely.

Measured on MetaX C550 against 1388 GB/s of achievable bandwidth (512 MiB
device-to-device copy):

    64 heads    1164 -> 1339 GB/s   (83.9% -> 96.5%)
    128 heads   1224 -> 1354 GB/s   (88.2% -> 97.5%)

Output is bit-identical to the generic kernel -- the FP8 cache matches byte for
byte and q matches exactly -- including at token counts that are not multiples
of TPP. The reduction structure is unchanged: reducing [TPP, 512] along axis 1
is the same tree per row as reducing [512], so the arithmetic order does not
move. test_backend_override_matches_generic asserts this for any backend
override, and skips where none is registered.

Author: cheersluvs

src/flag_gems/runtime/backend/_metax/fused/__init__.py

@@ -13,11 +13,15 @@
 # limitations under the License.
 
 from .flash_mla import flash_mla
+from .fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert import (
+    fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert,
+)
 from .sparse_attention import sparse_attn_triton
 from .top_k_per_row_prefill import top_k_per_row_prefill
 
 __all__ = [
     "flash_mla",
+    "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
     "sparse_attn_triton",
     "top_k_per_row_prefill",
 ]

src/flag_gems/runtime/backend/_metax/fused/fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.py

@@ -0,0 +1,254 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MetaX override: token-tiled fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.
+
+The generic kernel runs one program per (token, head) slot with num_warps=1 --
+64 threads for 512 elements. On C550 that leaves bandwidth on the table two
+ways: the KV slot does ~7 quant blocks against a Q slot's normalize+rotate, so
+one-slot-per-program makes it a straggler (1-in-65 at 64 heads vs 1-in-129 at
+128, which is why the shortfall is worse at the lower head count); and
+64-thread blocks under-fill the SM.
+
+Raising num_warps alone makes it worse -- work per program stays at 512
+elements, so each lane gets less to do (measured 1134 -> 691 -> 398 -> 199
+GB/s). Work and width have to rise together. Here each program handles TPP
+tokens of ONE slot, so every program is uniformly all-Q or all-KV.
+
+Measured on C550 against a 1388 GB/s ceiling (512 MiB device-to-device copy):
+
+    shape              generic          tiled       vs MetaX mcoplib
+    131072 x 64     1164 (83.9%)   1339 (96.5%)     1336 (96.3%)
+     32768 x 128    1224 (88.2%)   1354 (97.5%)     1298 (93.5%)
+
+Output is bit-identical to the generic kernel (FP8 cache exact, q exact) across
+shapes including non-multiples of TPP.
+
+Below 512 tokens the generic kernel wins -- a TPP=8 launch masks off 7/8 of
+every program and the wider blocks raise the launch floor from ~26us to ~32us
+-- so this dispatches to it. Crossover measured at 512 on both head counts
+(256: 0.96x/0.99x, 512: 1.03x/1.04x).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+# Measured on MetaX C550; see module docstring.
+_TILED_MIN_TOKENS = 512
+_TPP = 8
+_NUM_WARPS = 4
+
+
+@triton.jit
+def _tiled_kernel(
+    q,
+    kv,
+    k_cache,
+    slot_mapping,
+    position_ids,
+    cos_sin_cache,
+    eps,
+    cache_block_size: tl.constexpr,
+    num_tokens,
+    num_heads: tl.constexpr,
+    kv_block_stride,
+    num_tokens_insert,
+    TPP: tl.constexpr,  # tokens per program
+):
+    HEAD_DIM: tl.constexpr = 512
+    NOPE_DIM: tl.constexpr = 448
+    ROPE_DIM: tl.constexpr = 64
+    HALF_ROPE_DIM: tl.constexpr = 32
+    QUANT_BLOCK: tl.constexpr = 64
+    NUM_QUANT_BLOCKS: tl.constexpr = NOPE_DIM // QUANT_BLOCK  # 7
+    SCALE_BYTES_PER_TOKEN: tl.constexpr = NUM_QUANT_BLOCKS + 1  # 8
+    TOKEN_DATA_BYTES: tl.constexpr = NOPE_DIM + 2 * ROPE_DIM  # 576
+    FP8_MAX: tl.constexpr = 448.0
+
+    pid = tl.program_id(0).to(tl.int64)
+    blocks_per_token: tl.constexpr = num_heads + 1
+
+    # grid = cdiv(num_tokens, TPP) * blocks_per_token
+    tile = pid // blocks_per_token
+    slot_idx = pid % blocks_per_token
+    is_kv = slot_idx == num_heads
+
+    tok = tile * TPP + tl.arange(0, TPP).to(tl.int64)  # [TPP]
+    tok_ok = tok < num_tokens
+
+    off = tl.arange(0, HEAD_DIM)  # [HEAD_DIM]
+    off_rope = tl.arange(0, ROPE_DIM)
+    off_half = tl.arange(0, HALF_ROPE_DIM)
+    off_quant = tl.arange(0, QUANT_BLOCK)
+
+    # cos/sin are needed by both paths
+    pos = tl.load(position_ids + tok, mask=tok_ok, other=0)  # [TPP]
+    cs_base = cos_sin_cache + pos[:, None] * ROPE_DIM
+    cos_blk = tl.load(cs_base + off_half[None, :], mask=tok_ok[:, None], other=0.0)
+    sin_blk = tl.load(
+        cs_base + (off_half + HALF_ROPE_DIM)[None, :], mask=tok_ok[:, None], other=0.0
+    )
+
+    if not is_kv:
+        # ── Q: per-head RMSNorm (no weight) + GPT-J RoPE, in place ──
+        q_base = q + (tok * num_heads + slot_idx) * HEAD_DIM  # [TPP]
+        q_blk = tl.load(
+            q_base[:, None] + off[None, :], mask=tok_ok[:, None], other=0.0
+        ).to(tl.float32)
+        variance = tl.sum(q_blk * q_blk, axis=1) / HEAD_DIM  # [TPP]
+        rsqrt = tl.rsqrt(variance + eps)
+        q_blk = q_blk * rsqrt[:, None]
+        tl.store(
+            q_base[:, None] + off[None, :],
+            q_blk.to(tl.bfloat16),
+            mask=tok_ok[:, None] & (off[None, :] < NOPE_DIM),
+        )
+        rope = (
+            tl.load(
+                q_base[:, None] + NOPE_DIM + off_rope[None, :],
+                mask=tok_ok[:, None],
+                other=0.0,
+            ).to(tl.float32)
+            * rsqrt[:, None]
+        )
+    else:
+        kv_base = kv + tok * HEAD_DIM
+        rope = tl.load(
+            kv_base[:, None] + NOPE_DIM + off_rope[None, :],
+            mask=tok_ok[:, None],
+            other=0.0,
+        ).to(tl.float32)
+
+    # ── GPT-J interleaved RoPE on the trailing ROPE_DIM ──
+    rope = tl.reshape(rope, TPP, HALF_ROPE_DIM, 2)
+    even, odd = tl.split(rope)  # [TPP, HALF_ROPE_DIM]
+    new_even = even * cos_blk - odd * sin_blk
+    new_odd = even * sin_blk + odd * cos_blk
+    rope = tl.reshape(tl.join(new_even, new_odd), TPP, ROPE_DIM).to(tl.bfloat16)
+
+    if not is_kv:
+        q_base = q + (tok * num_heads + slot_idx) * HEAD_DIM
+        tl.store(
+            q_base[:, None] + NOPE_DIM + off_rope[None, :],
+            rope,
+            mask=tok_ok[:, None],
+        )
+        return
+
+    # ── KV: RoPE already applied; UE8M0 FP8 quant + paged cache insert ──
+    kv_base = kv + tok * HEAD_DIM
+    ins_ok = tok_ok & (tok < num_tokens_insert)
+    slot_id = tl.load(slot_mapping + tok, mask=ins_ok, other=-1)  # [TPP]
+    ins_ok = ins_ok & (slot_id >= 0)
+
+    block_idx = slot_id // cache_block_size
+    pos_in_block = slot_id % cache_block_size
+    block_base = block_idx * kv_block_stride
+    token_fp8 = block_base + pos_in_block * TOKEN_DATA_BYTES  # [TPP] byte offset
+    token_scale = (
+        block_base
+        + cache_block_size * TOKEN_DATA_BYTES
+        + pos_in_block * SCALE_BYTES_PER_TOKEN
+    )
+
+    bf16_ptr = (k_cache + token_fp8 + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
+    tl.store(bf16_ptr[:, None] + off_rope[None, :], rope, mask=ins_ok[:, None])
+
+    for b in tl.static_range(NUM_QUANT_BLOCKS):
+        blk = tl.load(
+            kv_base[:, None] + b * QUANT_BLOCK + off_quant[None, :],
+            mask=ins_ok[:, None],
+            other=0.0,
+        ).to(tl.float32)
+        block_max = tl.maximum(tl.max(tl.abs(blk), axis=1), 1e-4)  # [TPP]
+        exponent = tl.ceil(tl.log2(block_max / FP8_MAX))
+        scale = tl.exp2(exponent)
+        x = tl.clamp(blk / scale[:, None], -FP8_MAX, FP8_MAX)
+        tl.store(
+            k_cache + token_fp8[:, None] + b * QUANT_BLOCK + off_quant[None, :],
+            x.to(tl.float8e4nv).to(tl.uint8, bitcast=True),
+            mask=ins_ok[:, None],
+        )
+        enc = tl.maximum(tl.minimum(exponent + 127.0, 255.0), 0.0)
+        tl.store(k_cache + token_scale + b, enc.to(tl.uint8), mask=ins_ok)
+
+    tl.store(
+        k_cache + token_scale + NUM_QUANT_BLOCKS,
+        tl.zeros((TPP,), dtype=tl.uint8),
+        mask=ins_ok,
+    )
+
+
+def fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    k_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    position_ids: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    eps: float,
+    cache_block_size: int,
+):
+    """See the generic implementation for the layout contract."""
+    assert q.is_contiguous() and kv.is_contiguous()
+    num_tokens, num_heads, head_dims = q.shape
+
+    if num_tokens < _TILED_MIN_TOKENS:
+        # Small shapes are launch-bound; the narrower generic kernel wins.
+        from flag_gems.fused.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert import (
+            fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert as _generic,
+        )
+
+        return _generic(
+            q,
+            kv,
+            k_cache,
+            slot_mapping,
+            position_ids,
+            cos_sin_cache,
+            eps,
+            cache_block_size,
+        )
+
+    assert head_dims == 512
+    assert kv.shape == (num_tokens, 512)
+    assert q.dtype == torch.bfloat16 and kv.dtype == torch.bfloat16
+    assert k_cache.dtype == torch.uint8
+    assert slot_mapping.dim() == 1
+    num_tokens_insert = slot_mapping.shape[0]
+    assert num_tokens_insert <= num_tokens
+    assert slot_mapping.dtype == torch.int64
+    assert position_ids.shape == (num_tokens,)
+    assert position_ids.dtype == torch.int64
+    assert cos_sin_cache.dim() == 2 and cos_sin_cache.shape[1] == 64
+    assert cos_sin_cache.dtype == torch.float32
+
+    grid = triton.cdiv(num_tokens, _TPP) * (num_heads + 1)
+    _tiled_kernel[(grid,)](
+        q,
+        kv,
+        k_cache,
+        slot_mapping,
+        position_ids,
+        cos_sin_cache,
+        eps,
+        cache_block_size,
+        num_tokens,
+        num_heads,
+        k_cache.stride(0),
+        num_tokens_insert,
+        TPP=_TPP,
+        num_warps=_NUM_WARPS,
+        num_stages=2,
+    )

tests/test_fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.py

@@ -570,3 +570,61 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
 
     torch.testing.assert_close(q, q_ref, rtol=1e-2, atol=1e-2)
     k_cache_compare(k_cache, k_cache_ref, block_size, rtol=1e-2, atol=1e-2)
+
+
+# ── Test 5: a backend override must not change results ───────────────────────
+# A vendor override is a performance change. It runs only where one is
+# registered, so nothing below executes on the generic path.
+_OVERRIDE_ACTIVE = (
+    flag_gems.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert is not _generic_impl
+)
+
+
+@pytest.mark.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert
+@pytest.mark.skipif(not is_support_fp8e4nv(), reason="Device does not support fp8e4nv")
+@pytest.mark.skipif(
+    not _OVERRIDE_ACTIVE,
+    reason="No backend override registered; the generic kernel is in use",
+)
+@pytest.mark.parametrize(
+    "num_tokens",
+    [1, 17, 513] if QUICK_MODE else [1, 17, 64, 511, 512, 777, 1000, 4096],
+)
+@pytest.mark.parametrize("n_heads", [64, 128])
+@pytest.mark.parametrize("block_size", [16, 64])
+def test_backend_override_matches_generic(
+    num_tokens: int, n_heads: int, block_size: int
+):
+    torch.manual_seed(3)
+    device = "cuda"
+    eps = 1e-6
+    max_pos = max(4096, num_tokens)
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+
+    q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=device)
+    kv = torch.randn(num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device)
+    k_cache = torch.zeros(
+        num_blocks, block_size * HEAD_BYTES, dtype=torch.uint8, device=device
+    )
+    slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
+
+    q_gen, k_cache_gen = q.clone(), k_cache.clone()
+    _generic_impl(
+        q_gen, kv, k_cache_gen, slot_mapping, positions, cos_sin_cache, eps, block_size
+    )
+    fused_impl(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, block_size)
+
+    # The cache is derived from kv through per-64-element reductions, identical
+    # in any sane decomposition, so it must match byte for byte -- an override
+    # that quantizes differently has changed behaviour.
+    assert torch.equal(k_cache, k_cache_gen), (
+        f"override wrote a different FP8 cache at num_tokens={num_tokens}, "
+        f"n_heads={n_heads}, block_size={block_size}: "
+        f"{int((k_cache != k_cache_gen).sum().item())} of {k_cache.numel()} "
+        "bytes differ"
+    )
+    # q goes through a 512-element RMSNorm reduction whose tree may legitimately
+    # differ between decompositions, so allow a bf16 ULP there.
+    torch.testing.assert_close(q, q_gen, rtol=1e-3, atol=1e-3)