Fuse the DSA (V3.2, GLM-5.x) indexer Q/K paths into single kernels (#27705)

Co-authored-by: Brayden Zhong <brayden@radixark.ai>
Co-authored-by: Kaixi <kaiximatteoc@nvidia.com>

test/registered/8-gpu-models/test_deepseek_v32_indexcache.py

@@ -63,7 +63,7 @@ class TestDeepseekV32IndexTopkPattern(CustomTestCase):
             write_github_step_summary(
                 f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n'
             )
-            self.assertGreater(metrics["accuracy"], 0.935)
+            self.assertGreater(metrics["accuracy"], 0.93)
 
 
 class TestDeepseekV32IndexFreq(CustomTestCase):

test/registered/jit/test_dsv32_indexer_fusion.py

@@ -0,0 +1,206 @@
+"""Correctness tests for the DeepSeek-V3.2 DSA indexer fused kernels.
+
+Covers:
+  - fused_q_indexer_rope_first_quant      (Q: rope-first + fp8 quant + head-gate fold)
+  - fused_k_indexer_norm_rope             (K: LayerNorm + rope-first -> bf16)
+  - fused_k_indexer_norm_rope_store       (K: the above + fp8 quant + paged index-k cache write)
+
+The store kernel is checked for byte-exact equivalence against the un-fused path
+(bf16 K kernel + standalone fused_store_index_k_cache), so it needs no fp8
+reference. The Q/K math kernels are checked against torch references. Strided
+inputs (the non-contiguous wk_weights_proj slices) are checked to match
+contiguous inputs (the no-copy path).
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from sglang.jit_kernel.dsv4 import fused_q_indexer_rope_first_quant
+from sglang.jit_kernel.dsv32 import (
+    fused_k_indexer_norm_rope,
+    fused_k_indexer_norm_rope_store,
+)
+from sglang.jit_kernel.fused_store_index_cache import (
+    can_use_dsa_fused_store,
+    fused_store_index_k_cache,
+)
+from sglang.srt.utils import is_hip
+from sglang.test.ci.ci_register import register_cuda_ci
+
+_is_hip = is_hip()
+
+register_cuda_ci(est_time=45, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=90, suite="nightly-kernel-1-gpu", nightly=True)
+
+HEAD_DIM = 128
+ROPE_DIM = 64
+HALF = ROPE_DIM // 2
+FP8_MAX = 448.0
+PAGE_SIZE = 64
+BYTES_PER_TOKEN = HEAD_DIM + 4  # 128 fp8 + 4-byte fp32 scale
+EPS = 1e-6
+MAX_POS = 8192
+
+
+def _skip_if_unavailable():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    if _is_hip:
+        pytest.skip("Indexer fused kernels are CUDA-specific")
+
+
+def _make_inputs(B, seed=0, pos_dtype=torch.int32):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    dev = "cuda"
+    cos = torch.randn(MAX_POS, HALF, device=dev, generator=g)
+    sin = torch.randn(MAX_POS, HALF, device=dev, generator=g)
+    freqs_cis = torch.complex(cos, sin)
+    positions = torch.randint(0, 4096, (B,), device=dev, dtype=pos_dtype, generator=g)
+    return cos, sin, freqs_cis, positions
+
+
+def _rope_first(x, cos_p, sin_p):
+    """Interleaved complex rope on the leading ROPE_DIM dims (kRopeFirst)."""
+    x = x.clone()
+    xr = x[..., 0:ROPE_DIM:2].clone()
+    xi = x[..., 1:ROPE_DIM:2].clone()
+    x[..., 0:ROPE_DIM:2] = xr * cos_p - xi * sin_p
+    x[..., 1:ROPE_DIM:2] = xr * sin_p + xi * cos_p
+    return x
+
+
+# ----------------------------------------------------------------------------
+# K kernel (-> bf16): LayerNorm + rope-first
+# ----------------------------------------------------------------------------
+def test_k_norm_rope_matches_reference():
+    _skip_if_unavailable()
+    dev = "cuda"
+    B = 37
+    cos, sin, freqs_cis, positions = _make_inputs(B)
+    key = torch.randn(B, HEAD_DIM, dtype=torch.bfloat16, device=dev)
+    weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
+    bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
+
+    out = fused_k_indexer_norm_rope(key, weight, bias, EPS, freqs_cis, positions)
+    torch.cuda.synchronize()
+
+    normed = torch.nn.functional.layer_norm(
+        key.float(), (HEAD_DIM,), weight=weight, bias=bias, eps=EPS
+    )
+    cp, sp = cos[positions.long()], sin[positions.long()]
+    ref = _rope_first(normed, cp, sp)
+
+    torch.testing.assert_close(out.float(), ref, atol=0.06, rtol=0.0)
+
+
+# ----------------------------------------------------------------------------
+# K kernel + fused store == bf16 K kernel + standalone store (byte-exact).
+# Also covers the strided (non-contiguous wk slice) no-copy input path.
+# ----------------------------------------------------------------------------
+@pytest.mark.parametrize("strided", [False, True])
+def test_k_store_matches_unfused(strided):
+    _skip_if_unavailable()
+    if not can_use_dsa_fused_store(torch.bfloat16, torch.int64, PAGE_SIZE):
+        pytest.skip("fused store JIT unavailable")
+    dev = "cuda"
+    B, n_heads = 41, 64
+    cos, sin, freqs_cis, positions = _make_inputs(B)
+    weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
+    bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
+
+    if strided:
+        # Mimic the fused wk_weights_proj GEMM output: key is kw[:, :head_dim].
+        kw = torch.randn(B, HEAD_DIM + n_heads, dtype=torch.bfloat16, device=dev)
+        key = kw[:, :HEAD_DIM]
+        assert not key.is_contiguous()
+    else:
+        key = torch.randn(B, HEAD_DIM, dtype=torch.bfloat16, device=dev)
+
+    loc = torch.randperm(B * 4, device=dev)[:B].to(torch.int64)
+    num_pages = int(loc.max().item()) // PAGE_SIZE + 2
+    buf_ref = torch.zeros(
+        num_pages, BYTES_PER_TOKEN * PAGE_SIZE, dtype=torch.uint8, device=dev
+    )
+    buf_fused = torch.zeros_like(buf_ref)
+
+    key_bf16 = fused_k_indexer_norm_rope(key, weight, bias, EPS, freqs_cis, positions)
+    fused_store_index_k_cache(key_bf16, buf_ref, loc, PAGE_SIZE)
+
+    fused_k_indexer_norm_rope_store(
+        key, buf_fused, loc, weight, bias, EPS, freqs_cis, positions, PAGE_SIZE
+    )
+    torch.cuda.synchronize()
+
+    assert torch.equal(buf_ref, buf_fused)
+
+
+# ----------------------------------------------------------------------------
+# Q kernel: rope-first + fp8 quant + head-gate fold
+# ----------------------------------------------------------------------------
+@pytest.mark.parametrize("pos_dtype", [torch.int32, torch.int64])
+def test_q_rope_quant_matches_reference(pos_dtype):
+    _skip_if_unavailable()
+    dev = "cuda"
+    B, n_heads = 37, 64
+    cos, sin, freqs_cis, positions = _make_inputs(B, pos_dtype=pos_dtype)
+    q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev)
+    weight = torch.randn(B, n_heads, dtype=torch.bfloat16, device=dev)
+    weight_scale = 0.137
+
+    q_fp8, weights_out = fused_q_indexer_rope_first_quant(
+        q, weight, weight_scale, freqs_cis, positions
+    )
+    torch.cuda.synchronize()
+
+    cp = cos[positions.long()][:, None, :]
+    sp = sin[positions.long()][:, None, :]
+    ref = _rope_first(q.float(), cp, sp)  # [B, n_heads, 128]
+    amax = ref.abs().amax(dim=-1, keepdim=True)
+    scale = torch.clamp(amax, min=1e-4) / FP8_MAX
+
+    # weights_out[b,h] = weight * weight_scale * scale
+    w_ref = weight.float() * weight_scale * scale.squeeze(-1)
+    torch.testing.assert_close(weights_out.squeeze(-1), w_ref, atol=1e-3, rtol=1e-3)
+
+    # dequantized q should match the rope result within fp8-e4m3 precision:
+    # round-to-nearest with 3 mantissa bits => <= 1/16 relative error, plus one
+    # scale step at the bottom of the range.
+    deq = q_fp8.float() * scale
+    err = (deq - ref).abs()
+    assert (
+        err <= 0.0625 * ref.abs() + scale
+    ).all(), f"max fp8 dequant error {err.max().item()}"
+
+
+# ----------------------------------------------------------------------------
+# Strided weight (the wk_weights_proj slice) matches contiguous for the Q kernel
+# ----------------------------------------------------------------------------
+def test_q_strided_weight_matches_contiguous():
+    _skip_if_unavailable()
+    dev = "cuda"
+    B, n_heads = 29, 64
+    cos, sin, freqs_cis, positions = _make_inputs(B)
+    q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev)
+    # weights_raw = kw[:, head_dim:] is a non-contiguous slice.
+    kw = torch.randn(B, HEAD_DIM + n_heads, dtype=torch.bfloat16, device=dev)
+    w_strided = kw[:, HEAD_DIM:]
+    w_contig = w_strided.contiguous()
+    assert not w_strided.is_contiguous()
+
+    a_fp8, a_w = fused_q_indexer_rope_first_quant(
+        q, w_strided, 0.137, freqs_cis, positions
+    )
+    b_fp8, b_w = fused_q_indexer_rope_first_quant(
+        q, w_contig, 0.137, freqs_cis, positions
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(a_fp8, b_fp8)
+    assert torch.equal(a_w, b_w)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/models_e2e/test_dsa_glm5_dp_mtp.py

@@ -19,7 +19,7 @@ class TestGLM5DPMTP(
     DsaMtpServerBase, DsaMtpEvalConfigDefaults, GSM8KMixin, SpecDecodingMixin
 ):
     model = "zai-org/GLM-5-FP8"
-    mem_fraction_static = 0.8
+    mem_fraction_static = 0.88
     enable_dp_attention = True
     bs_1_speed_thres = 70