[Perf][DSA] Unify q8kv8 sparse-prefill topk_length source with the bf16 path; delete backscan kernel

python/sglang/kernels/ops/attention/utils.py

@@ -30,9 +30,6 @@
 from sglang.kernels.ops.kvcache.cache_ops import (
     launch_reshape_and_cache_flash as launch_reshape_and_cache_flash,
 )
-from sglang.kernels.ops.kvcache.cache_ops import (
-    q8kv8_topk_length_from_indices as q8kv8_topk_length_from_indices,
-)
 from sglang.kernels.ops.kvcache.cache_ops import (
     reshape_and_cache_flash as reshape_and_cache_flash,
 )

python/sglang/kernels/ops/kvcache/cache_ops.py

@@ -736,54 +736,6 @@ def absorbed_bmm_concat_cast_q_fp8(
     )
 
 
-@triton.jit
-def q8kv8_topk_length_backscan_kernel(
-    indices_ptr,
-    out_ptr,
-    stride_row,
-    topk,
-    BLOCK: tl.constexpr,
-):
-    row = tl.program_id(0).to(tl.int64)
-    base = indices_ptr + row * stride_row
-    off = topk
-    length = 1
-    found = 0
-    while (found == 0) & (off > 0):
-        off -= BLOCK
-        idx = off + tl.arange(0, BLOCK)
-        vals = tl.load(base + idx)
-        pos = tl.max(tl.where(vals >= 0, idx, -1), axis=0)
-        found = tl.where(pos >= 0, 1, found)
-        length = tl.where(pos >= 0, pos + 1, length)
-    tl.store(out_ptr + row, length)
-
-
-def q8kv8_topk_length_from_indices(indices: torch.Tensor) -> torch.Tensor:
-    """Per-row valid-topk count = last non-negative position + 1 (min 1).
-
-    ``indices``: [s_q, topk] int32 topk output whose pad slots are -1.
-    Backward block scan per row: the loop exits at the first block holding a
-    valid entry, so the cost is proportional to the trailing pad run — one
-    block (~topk/4 elements) for rows with a full topk, which dominate long
-    contexts. Semantics match the unfused ``(indices >= 0) * ramp).amax``
-    derivation exactly, including all-pad rows (length 1: one pad-only block
-    keeps the kernel on its clamp+mask path, contributing zero).
-    """
-    s_q, topk = indices.shape
-    assert indices.dtype == torch.int32 and indices.stride(1) == 1
-    out = torch.empty(s_q, dtype=torch.int32, device=indices.device)
-    block = 512 if topk % 512 == 0 else (256 if topk % 256 == 0 else 128)
-    q8kv8_topk_length_backscan_kernel[(s_q,)](
-        indices,
-        out,
-        indices.stride(0),
-        topk,
-        BLOCK=block,
-    )
-    return out
-
-
 # ---------------------------------------------------------------------------
 # Decode Context Parallel (DCP) helpers.
 #

python/sglang/srt/environ.py

@@ -1406,13 +1406,6 @@ class Envs:
     # different GEMM accumulation order), hence default OFF until accuracy-
     # gated (oracle + full-set gsm8k).
     SGLANG_ENABLE_DSA_Q8KV8_BORN_FP8_Q = EnvBool(False)
-    # Opt-in perf path for --dsa-prefill-backend flashmla_sparse_q8: pass a
-    # per-row valid-topk count (derived from the trailing -1 pad run of the
-    # topk indices) so the kernel skips whole pad-only topk blocks instead of
-    # computing masked zero contributions.  Bit-exact by construction: skipped
-    # blocks contain only -1 pads, and -1 entries inside the consumed range
-    # still take the in-kernel clamp+mask path.
-    SGLANG_ENABLE_DSA_Q8KV8_TOPK_LENGTH = EnvBool(False)
     # Opt-in: run the born-fp8 q-prep (absorbed bmm + concat + fp8 cast,
     # ~173us/layer-call) on alt_stream underneath the DSA indexer — the two
     # chains fork independently from the q_a_layernorm output.  Requires

python/sglang/srt/layers/attention/dsa_backend.py

@@ -35,7 +35,6 @@
 from sglang.kernels.ops.attention.utils import (
     concat_mla_absorb_q_general,
     mla_quantize_and_rope_for_fp8,
-    q8kv8_topk_length_from_indices,
     seqlens_expand_triton,
 )
 from sglang.kernels.ops.kvcache.cache_ops import concat_and_cast_q_fp8_pad
@@ -446,12 +445,6 @@ def __init__(
         # the gather kernel fuses that in.  Same single-stream reuse
         # argument as `_q8kv8_qpad_buf`.
         self._q8kv8_kv_buf: Optional[torch.Tensor] = None
-        # Per-row valid-topk early-exit (SGLANG_ENABLE_DSA_Q8KV8_TOPK_LENGTH):
-        # rows whose topk indices end in a -1 pad run skip whole topk blocks
-        # in-kernel.
-        self._q8kv8_topk_length_enabled: bool = (
-            envs.SGLANG_ENABLE_DSA_Q8KV8_TOPK_LENGTH.get()
-        )
         # Persistent (grow-only) kernel-output buffers (out/max_logits/lse).
         self._q8kv8_out_bufs: Optional[tuple] = None
         # Fused non-prefix KV prep (cast-concat k/k_rope directly into the
@@ -2075,6 +2068,7 @@ def forward_extend(
                             sm_scale=layer.scaling,
                             v_head_dim=layer.v_head_dim,
                             layer_id=layer.layer_id,
+                            topk_length=metadata.dsa_cache_seqlens_int32,
                         )
                     if self._q8kv8_kv_cat_fusion:
                         # Fused path: no bf16 concat materialization — k and
@@ -2092,6 +2086,7 @@ def forward_extend(
                             sm_scale=layer.scaling,
                             v_head_dim=layer.v_head_dim,
                             layer_id=layer.layer_id,
+                            topk_length=metadata.dsa_cache_seqlens_int32,
                         )
                     kv_cache = _cat([k, k_rope], dim=-1)
                     return self._forward_flashmla_sparse_q8kv8(
@@ -2104,6 +2099,7 @@ def forward_extend(
                         sm_scale=layer.scaling,
                         v_head_dim=layer.v_head_dim,
                         layer_id=layer.layer_id,
+                        topk_length=metadata.dsa_cache_seqlens_int32,
                     )
 
                 # bf16 path (dsa_impl == "flashmla_sparse").
@@ -2558,17 +2554,19 @@ def _forward_flashmla_sparse_q8kv8(
         layer_id: Optional[int] = None,
         kv_k: Optional[torch.Tensor] = None,
         kv_k_rope: Optional[torch.Tensor] = None,
+        topk_length: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """Native FP8 (q8 x kv8) sparse-prefill attention (SM90 JIT kernel).
 
         Same contract as ``_forward_flashmla_sparse`` but executed through the
         FP8 ``sparse_mla_q8kv8_prefill_fwd`` kernel.  Identity per-tensor
         scales (scalar 1.0) are used: a raw bf16->fp8 cast of q/kv is accurate
         on real DeepSeek-V3 magnitudes, so no dynamic rescaling is applied.
-        The kernel runs via its fixed full-topk entry (``attn_sink`` /
-        ``topk_length`` left None), keeping control flow identical across DP
-        ranks; -1 topk sentinels are clamped to distinct zero pad rows inside
-        the kernel.
+        ``topk_length`` (``dsa_cache_seqlens_int32``, the same metadata the
+        bf16 ``_forward_flashmla_sparse`` path consumes) is a device tensor,
+        so host-side control flow stays identical across DP ranks; -1 topk
+        sentinels are masked out in-kernel (their KV loads are predicated
+        off).
 
         Two KV paths:
           * non-prefix extend: ``kv_bf16`` (the gathered bf16 KV) is cast into
@@ -2715,15 +2713,20 @@ def _forward_flashmla_sparse_q8kv8(
             kv_padded[num_kv_tokens:].zero_()
             kv_padded = kv_padded.view(-1, 1, head_dim)
 
-        # Per-row valid-topk count = last non-pad position + 1.  Bit-exact
-        # vs topk_length=None: the skipped tail blocks contain only -1 pads
-        # (masked to zero contribution today), and -1 entries inside the
-        # consumed range still take the kernel's clamp+mask path.  The
-        # backscan's cost is proportional to the trailing pad run, so rows
-        # with a full topk (all rows at long context) pay ~one block read.
-        topk_length = None
-        if self._q8kv8_topk_length_enabled:
-            topk_length = q8kv8_topk_length_from_indices(page_table_1)
+        # Per-row valid-topk early-exit: topk_length is the per-row count of
+        # valid indices (`dsa_cache_seqlens_int32` = seqlens clipped to
+        # `index_topk`), the same metadata source the bf16
+        # `_forward_flashmla_sparse` path uses. Rows whose context is shorter
+        # than `index_topk` have their indices tail-padded with -1; the valid
+        # length lets the kernel skip whole pad-only topk blocks instead of
+        # computing masked zero contributions (-1 entries inside the consumed
+        # range are still masked out in-kernel, and the metadata
+        # value can only exceed the position of the last valid index).
+        if topk_length is not None and topk_length.shape[0] != num_tokens:
+            # Metadata rows are expected to match q rows (the DP/CP padding
+            # helpers keep them aligned); fall back to full-width compute if
+            # they ever diverge.
+            topk_length = None
 
         # Persistent kernel-output buffers (out / max_logits / lse): the
         # wrapper otherwise torch.empty's all three per layer-call.  The

test/registered/kernels/ops/attention/test_sparse_mla_q8kv8_prefill_sm90.py

@@ -192,9 +192,11 @@ def test_sparse_mla_q8kv8_prefill_corner_cases(
     _run_and_check(d_qk, with_sink, s_q=s_q, topk=topk, s_kv=s_kv)
 
 
-# topk_length WITHOUT attn_sink (the production early-exit path for
-# SGLANG_ENABLE_DSA_Q8KV8_TOPK_LENGTH): rows with a trailing -1 pad run must
-# be BITWISE identical to the full-topk dispatch that masks those pads, and
+# topk_length WITHOUT attn_sink (the production early-exit path, fed by the
+# metadata-derived per-row valid-topk count): rows with a trailing -1 pad run
+# must be BITWISE identical to the full-topk dispatch that masks those pads
+# (which holds for these specific length patterns; arbitrary lengths that
+# change the block-iteration count may perturb `out` by 1 bf16 ULP), and
 # must match the fp32 reference on the truncated index range.
 @pytest.mark.skipif(
     not _sm90_available(), reason="Q8KV8 sparse prefill requires SM90 CUDA"
@@ -681,43 +683,93 @@ def test_sparse_mla_q8kv8_prefill_large_skv():
     assert cos > 0.99, f"cos {cos:.4f} <= 0.99"
 
 
-# Backend-side topk_length derivation (backscan Triton kernel): must equal the
-# reference "last non-negative position + 1 (min 1)" on every pad pattern the
-# production topk output can produce (trailing runs), plus adversarial ones
-# (interleaved -1s, all-pad, full rows) where the trailing-run semantics still
-# define the correct consumed range.
+# Metadata-derived topk_length: production passes dsa_cache_seqlens_int32
+# (seqlens clipped to topk) instead of a count derived from the indices
+# themselves. The metadata value can only meet or exceed the exact
+# last-valid-position count (in-range -1 sentinels are clamp+masked either
+# way), so both sources must be BITWISE interchangeable. All-pad rows are the
+# one place the two sources disagree numerically (metadata 0 -- e.g. DP/CP
+# padding rows, which pad_dsa_cache_seqlens fills with zeros -- vs the
+# index-derived floor of 1): both must yield the empty-row sentinel outputs
+# (zero out, max_logits=-inf, lse=+inf) without wedging the kernel's
+# producer/consumer handshake.
 @pytest.mark.skipif(
     not _sm90_available(), reason="Q8KV8 sparse prefill requires SM90 CUDA"
 )
-@pytest.mark.parametrize("s_q,topk", [(437, 2048), (7, 128), (65, 256), (4096, 2048)])
-def test_q8kv8_topk_length_backscan(s_q: int, topk: int):
-    from sglang.kernels.ops.kvcache.cache_ops import (
-        q8kv8_topk_length_from_indices,
+@pytest.mark.parametrize(
+    "s_q,topk,s_kv", [(437, 2048, 4096), (65, 256, 592), (128, 128, 512)]
+)
+def test_q8kv8_topk_length_metadata_equivalence(s_q: int, topk: int, s_kv: int):
+    from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
+        sparse_mla_q8kv8_prefill_fwd,
     )
 
+    d_qk = 576
     generator = torch.Generator(device="cuda")
     generator.manual_seed(4000 + s_q + topk)
+    q = _make_fp8_tensor((s_q, H_Q, d_qk), seed=s_q)
+    kv = _make_fp8_tensor((s_kv, H_KV, d_qk), seed=s_kv)
+    one = torch.ones(1, dtype=torch.float32, device="cuda")
+    sm_scale = 1.0 / math.sqrt(d_qk)
+
+    # Metadata analogue: per-row seqlens clipped to topk, with all-pad rows
+    # (value 0) mixed in like the DP/CP padding helpers produce.
+    seqlens = torch.randint(
+        1, topk + 1, (s_q,), dtype=torch.int32, device="cuda", generator=generator
+    )
+    seqlens[::5] = 0
     indices = torch.randint(
-        0, 1 << 20, (s_q, topk), dtype=torch.int32, device="cuda", generator=generator
-    )
-    # Row patterns: full, trailing pad runs of every length, all-pad,
-    # interleaved -1s inside the valid range.
-    for i in range(s_q):
-        mode = i % 5
-        if mode == 1:
-            indices[i, max(1, i % topk) :] = -1
-        elif mode == 2:
-            indices[i, :] = -1
-        elif mode == 3:
-            indices[i, i % topk :: 7] = -1  # interleaved + trailing mix
-        elif mode == 4:
-            indices[i, topk - 1 :] = -1
-
-    got = q8kv8_topk_length_from_indices(indices)
-
-    ramp = torch.arange(1, topk + 1, dtype=torch.int32, device="cuda")
-    ref = ((indices >= 0).int() * ramp).amax(dim=-1).clamp_(min=1)
-    assert torch.equal(got, ref)
+        0,
+        s_kv,
+        (s_q, H_KV, topk),
+        dtype=torch.int32,
+        device="cuda",
+        generator=generator,
+    )
+    ramp = torch.arange(topk, dtype=torch.int32, device="cuda").view(1, 1, topk)
+    indices = torch.where(
+        ramp < seqlens.view(-1, 1, 1), indices, torch.full_like(indices, -1)
+    )
+    # In-range -1 sentinels: exercise the clamp+mask path inside the consumed
+    # range (makes the index-derived count strictly smaller than the metadata
+    # value on some rows).
+    sprinkle = (
+        torch.rand((s_q, 1, topk), device="cuda", generator=generator) < 0.02
+    ) & (ramp < seqlens.view(-1, 1, 1))
+    indices = torch.where(sprinkle, torch.full_like(indices, -1), indices)
+
+    # Exact index-derived count = last non-negative position + 1 (min 1), the
+    # semantics of the deleted backend-side backscan derivation.
+    ramp1 = torch.arange(1, topk + 1, dtype=torch.int32, device="cuda")
+    exact = ((indices.squeeze(1) >= 0).int() * ramp1).amax(dim=-1).clamp_(min=1)
+
+    outs = {}
+    for name, topk_length in (("metadata", seqlens), ("exact", exact)):
+        outs[name] = sparse_mla_q8kv8_prefill_fwd(
+            q=q,
+            kv=kv,
+            indices=indices,
+            sm_scale=sm_scale,
+            q_scale=one,
+            kv_scale=one,
+            d_v=D_V,
+            attn_sink=None,
+            topk_length=topk_length,
+        )
+    torch.cuda.synchronize()
+
+    for (om, mlm, lsem), (oe, mle, lsee) in ((outs["metadata"], outs["exact"]),):
+        assert torch.equal(om, oe)
+        assert torch.equal(mlm, mle)
+        assert torch.equal(lsem, lsee)
+
+    # All-pad rows: empty-row sentinels from both sources.
+    allpad = seqlens == 0
+    assert bool(allpad.any())
+    out_m, ml_m, lse_m = outs["metadata"]
+    assert bool((out_m[allpad] == 0).all())
+    assert bool((ml_m[allpad] == float("-inf")).all())
+    assert bool((lse_m[allpad] == float("inf")).all())
 
 
 if __name__ == "__main__":