[DSA] Q8KV8 FP8 Sparse Prefill on GLM-5.2 & DeepSeek-V3.2: Q8-Path & Shared-Path Optimizations (#31888)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

test/manual/test_moe_quant_once.py

@@ -0,0 +1,274 @@
+"""Standalone GPU test for SGLANG_OPT_MOE_QUANT_ONCE (quantize the MoE input
+once, feed both the fused shared-expert GEMM and the routed triton runner).
+
+
+    CUDA_VISIBLE_DEVICES=0 python test/manual/test_moe_quant_once.py
+
+Verifies, against the double-quant baseline:
+  (1) quant equivalence: the row-padded quantize-once kernel produces the
+      same q bits / scale values as the routed path's default row-major quant
+      (JIT v2 kernel) on the valid rows;
+  (2) shared consumer: cutlass_w8a8_block_fp8_linear_with_fallback with a
+      pre-quantized (q, s) tuple vs its own internal quant -- expected BITWISE
+      (baseline uses the identical row-padded quant + identical GEMM);
+  (2b) shared consumer under SGLANG_ENABLE_JIT_DEEPGEMM=1 (the recommended JIT-DeepGEMM config):
+      deepgemm_w8a8_block_fp8_linear_with_fallback with the same (q, s) tuple
+      -- expected BITWISE (DG's own quant layout, column-major TMA-aligned
+      fp32 scales, is byte-identical to the row-padded quantize-once layout);
+      skipped cleanly when deep_gemm is unavailable or UE8M0 (Blackwell);
+  (3) routed consumer: fused_experts(a1_q=..., a1_scale=...) vs the in-kernel
+      quant baseline -- expected BITWISE if (1) is bitwise (the fused kernel
+      reads A_scale through explicit strides, so the column-major scale view
+      feeds identical values).
+
+If (1) is not bitwise (AOT v2 vs JIT v2 quant kernels round differently),
+(3) falls back to an allclose check at atol=1e-2 and the discrepancy is
+reported --.
+"""
+
+import sys
+
+import torch
+
+from sglang.kernels.ops.quantization.fp8_kernel import (
+    sglang_per_token_group_quant_fp8,
+    sglang_per_token_group_quant_fp8_row_padded,
+)
+from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
+from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_experts
+from sglang.srt.layers.moe.topk import StandardTopKOutput
+from sglang.srt.layers.quantization.fp8_utils import (
+    cutlass_w8a8_block_fp8_linear_with_fallback,
+)
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+
+GROUP = 128
+FAILURES = []
+
+
+def _report(name, ok, detail=""):
+    status = "PASS" if ok else "FAIL"
+    print(f"[{status}] {name} {detail}")
+    if not ok:
+        FAILURES.append(name)
+
+
+def _quant_weight_blockwise(w_bf16, block=128):
+    """Per-[128,128]-block fp8 weight quant (reference, fp32 math)."""
+    n, k = w_bf16.shape
+    w = w_bf16.float().view(n // block, block, k // block, block)
+    amax = w.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-4)
+    scale = amax / torch.finfo(torch.float8_e4m3fn).max
+    q = (w / scale).clamp(-448, 448).to(torch.float8_e4m3fn)
+    return (
+        q.view(n, k),
+        scale.squeeze(1).squeeze(-1).to(torch.float32),  # [n/128, k/128]
+    )
+
+
+def test_quant_equivalence(T, K, device):
+    x = torch.randn(T, K, device=device, dtype=torch.bfloat16) * 3
+    q_ref, s_ref = sglang_per_token_group_quant_fp8(x, GROUP)  # routed baseline
+    q_pad, s_pad = sglang_per_token_group_quant_fp8_row_padded(x, GROUP)
+
+    bitwise_q = torch.equal(q_pad[:T].view(torch.uint8), q_ref.view(torch.uint8))
+    bitwise_s = torch.equal(s_pad[:T].contiguous(), s_ref)
+    _report(
+        f"quant-equivalence T={T} K={K}",
+        bitwise_q and bitwise_s,
+        f"(q bitwise={bitwise_q}, s bitwise={bitwise_s})",
+    )
+    return bitwise_q and bitwise_s
+
+
+def test_shared_consumer(T, K, N, device):
+    torch.manual_seed(T + K)
+    x = torch.randn(T, K, device=device, dtype=torch.bfloat16)
+    w_bf16 = torch.randn(N, K, device=device, dtype=torch.bfloat16) / K**0.5
+    w, ws = _quant_weight_blockwise(w_bf16)
+
+    ref = cutlass_w8a8_block_fp8_linear_with_fallback(x, w, [128, 128], ws)
+
+    q_pad, s_pad = sglang_per_token_group_quant_fp8_row_padded(x, GROUP)
+    out = cutlass_w8a8_block_fp8_linear_with_fallback(
+        q_pad, w, [128, 128], ws, input_scale=s_pad
+    )[:T]
+
+    bitwise = torch.equal(out, ref)
+    close = torch.allclose(out.float(), ref.float(), atol=1e-2, rtol=1e-2)
+    _report(
+        f"shared-consumer T={T} K={K} N={N}",
+        close,
+        f"(bitwise={bitwise}, max|d|={(out.float() - ref.float()).abs().max().item():.3e})",
+    )
+    return bitwise
+
+
+def test_shared_consumer_deepgemm(T, K, N, device):
+    """DG branch (SGLANG_ENABLE_JIT_DEEPGEMM=1 recommended JIT-DeepGEMM config): the shared-expert
+    linear resolves to deepgemm_w8a8_block_fp8_linear_with_fallback. Its own
+    quant (column-major + TMA-aligned fp32 scales) has the same buffer layout
+    as the row-padded quantize-once kernel, so this is expected BITWISE."""
+    from sglang.srt.layers.quantization.fp8_utils import (
+        deepgemm_w8a8_block_fp8_linear_with_fallback,
+    )
+
+    torch.manual_seed(T + K + 1)
+    x = torch.randn(T, K, device=device, dtype=torch.bfloat16)
+    w_bf16 = torch.randn(N, K, device=device, dtype=torch.bfloat16) / K**0.5
+    w, ws = _quant_weight_blockwise_n64(w_bf16)
+
+    ref = deepgemm_w8a8_block_fp8_linear_with_fallback(x, w, [128, 128], ws)
+
+    q_pad, s_pad = sglang_per_token_group_quant_fp8_row_padded(x, GROUP)
+    out = deepgemm_w8a8_block_fp8_linear_with_fallback(
+        q_pad, w, [128, 128], ws, input_scale=s_pad
+    )[:T]
+
+    bitwise = torch.equal(out, ref)
+    close = torch.allclose(out.float(), ref.float(), atol=1e-2, rtol=1e-2)
+    _report(
+        f"shared-consumer-deepgemm T={T} K={K} N={N}",
+        close,
+        f"(bitwise={bitwise}, max|d|={(out.float() - ref.float()).abs().max().item():.3e})",
+    )
+    return bitwise
+
+
+def _quant_weight_blockwise_n64(w_bf16, block=128):
+    """Like _quant_weight_blockwise but supports N % 64 == 0 (DeepGEMM's
+    minimum): the last (partial) N-block reuses ceil-division block indexing."""
+    n, k = w_bf16.shape
+    if n % block == 0:
+        return _quant_weight_blockwise(w_bf16, block)
+    import math
+
+    n_blocks = math.ceil(n / block)
+    w = w_bf16.float()
+    q = torch.empty(n, k, device=w.device, dtype=torch.float8_e4m3fn)
+    scale = torch.empty(n_blocks, k // block, device=w.device, dtype=torch.float32)
+    for bn in range(n_blocks):
+        rows = slice(bn * block, min((bn + 1) * block, n))
+        wb = w[rows].view(rows.stop - rows.start, k // block, block)
+        amax = wb.abs().amax(dim=(0, 2)).clamp(min=1e-4)
+        s = amax / torch.finfo(torch.float8_e4m3fn).max
+        q[rows] = (
+            (wb / s[None, :, None]).clamp(-448, 448).to(torch.float8_e4m3fn).view(-1, k)
+        )
+        scale[bn] = s
+    return q, scale
+
+
+def test_routed_consumer(T, K, E, I, topk, device):
+    torch.manual_seed(T * 7 + K)
+    x = torch.randn(T, K, device=device, dtype=torch.bfloat16)
+    w1 = torch.empty(E, 2 * I, K, device=device, dtype=torch.float8_e4m3fn)
+    w1s = torch.empty(E, 2 * I // 128, K // 128, device=device)
+    w2 = torch.empty(E, K, I, device=device, dtype=torch.float8_e4m3fn)
+    w2s = torch.empty(E, K // 128, I // 128, device=device)
+    for e in range(E):
+        w1[e], w1s[e] = _quant_weight_blockwise(
+            torch.randn(2 * I, K, device=device, dtype=torch.bfloat16) / K**0.5
+        )
+        w2[e], w2s[e] = _quant_weight_blockwise(
+            torch.randn(K, I, device=device, dtype=torch.bfloat16) / I**0.5
+        )
+
+    topk_weights = torch.rand(T, topk, device=device)
+    topk_weights = (topk_weights / topk_weights.sum(-1, keepdim=True)).to(torch.float32)
+    topk_ids = torch.stack(
+        [torch.randperm(E, device=device)[:topk] for _ in range(T)]
+    ).to(torch.int32)
+    topk_output = StandardTopKOutput(
+        topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
+    )
+    # num_experts == num_local_experts => filter_expert=False (pure TP layout)
+    cfg = MoeRunnerConfig(
+        num_experts=E,
+        num_local_experts=E,
+        top_k=topk,
+        inplace=False,
+        activation="silu",
+        is_gated=True,
+    )
+
+    kwargs = dict(
+        w1=w1,
+        w2=w2,
+        topk_output=topk_output,
+        moe_runner_config=cfg,
+        use_fp8_w8a8=True,
+        w1_scale=w1s,
+        w2_scale=w2s,
+        block_shape=[128, 128],
+    )
+    ref = fused_experts(hidden_states=x, **kwargs)
+
+    q_pad, s_pad = sglang_per_token_group_quant_fp8_row_padded(x, GROUP)
+    out = fused_experts(hidden_states=x, a1_q=q_pad, a1_scale=s_pad, **kwargs)
+
+    bitwise = torch.equal(out, ref)
+    close = torch.allclose(out.float(), ref.float(), atol=1e-2, rtol=1e-2)
+    _report(
+        f"routed-consumer T={T} K={K} E={E} topk={topk}",
+        close,
+        f"(bitwise={bitwise}, max|d|={(out.float() - ref.float()).abs().max().item():.3e})",
+    )
+    return bitwise
+
+
+def main():
+    assert torch.cuda.is_available(), "CUDA required"
+    set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+    device = "cuda"
+    torch.manual_seed(0)
+
+    print("== (1) quantize-once vs routed-baseline quant equivalence ==")
+    all_bitwise_q = True
+    for T in (1, 3, 4093, 4096):
+        for K in (6144, 7168):
+            all_bitwise_q &= test_quant_equivalence(T, K, device)
+
+    print("== (2) shared consumer (cutlass w8a8 linear) ==")
+    # N=512 mirrors a tp8 shared expert gate_up (2*2048/8); must be %128==0.
+    for T in (4093, 4096):
+        for K in (6144, 7168):
+            test_shared_consumer(T, K, 512, device)
+
+    print(
+        "== (2b) shared consumer (deepgemm w8a8 linear, JIT DG recommended JIT-DeepGEMM config) =="
+    )
+    from sglang.srt.layers import deep_gemm_wrapper
+
+    if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
+        print("SKIP: deep_gemm unavailable or SGLANG_ENABLE_JIT_DEEPGEMM=0")
+    elif deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+        print("SKIP: Blackwell UE8M0 scale layout (gated ineligible by design)")
+    else:
+        for T in (4093, 4096):
+            for K in (6144, 7168):
+                test_shared_consumer_deepgemm(T, K, 512, device)
+        # DG accepts N % 64 (cutlass needs % 128) -- exercise the DG-only shape.
+        test_shared_consumer_deepgemm(4096, 7168, 320, device)
+
+    print("== (3) routed consumer (triton fused_experts) ==")
+    # Identical in both cutlass and JIT-DG configs: the MoE runner stays
+    # triton with a2a=none (is_deepgemm_moe_runner_backend_enabled() is False
+    # for auto + a2a=none even when SGLANG_ENABLE_JIT_DEEPGEMM=1).
+    for T in (61, 4093, 4096):
+        test_routed_consumer(T, 7168, E=32, I=256, topk=8, device=device)
+
+    if not all_bitwise_q:
+        print(
+            "NOTE: quantize-once q/s not bitwise vs the routed baseline quant "
+            "(AOT v2 vs JIT v2 kernel rounding) -- routed consumer is then "
+            "allclose-only; document this in the PR."
+        )
+    if FAILURES:
+        print(f"FAILED: {FAILURES}")
+        sys.exit(1)
+    print("ALL PASS")
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,120 @@
+"""Tests for the SM90 Q8KV8 born-fp8 q-prep JIT kernel.
+
+Gates (mirroring benchmark/kernels/deepseek/benchmark_q8kv8_q_prep.py conventions):
+  (a) vs the Triton absorbed_bmm_concat_cast_q_fp8 "two_dot" variant with
+      atol/rtol=2e-2 on the fp32 view (the accumulation order matches, so the
+      output is empirically bitwise identical on SM90, but only the tolerance
+      is contractual);
+  (b) vs an fp64 bmm reference: mean |err| must match two_dot's;
+  rope half must be bit-exact (identical bf16 -> fp8 conversion chain).
+"""
+
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=240, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+N_LORA = 512  # kv_lora_rank
+ROPE = 64  # qk_rope_head_dim
+
+
+def _is_sm90() -> bool:
+    return torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0)
+
+
+requires_sm90 = pytest.mark.skipif(not _is_sm90(), reason="requires SM90 (Hopper)")
+
+
+def _make_inputs(T: int, H: int, K: int, seed: int = 1234, magnitude: float = 1.0):
+    # Production layout: q_nope/q_rope are strided views of one [T, H, K+R]
+    # q_b_proj output; w_kc is the N-major absorbed weight [H, K, N] with
+    # strides (K*N, 1, K).
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    q = (
+        torch.randn((T, H, K + ROPE), generator=g, device="cuda", dtype=torch.float32)
+        * magnitude
+    ).to(torch.bfloat16)
+    w = (
+        torch.randn((H, N_LORA, K), generator=g, device="cuda", dtype=torch.float32)
+        / K**0.5
+    ).to(torch.bfloat16)
+    return q[..., :K], q[..., K:], w.transpose(1, 2)
+
+
+@requires_sm90
+@pytest.mark.parametrize("T", [1, 437, 1024])
+@pytest.mark.parametrize(
+    "h_k", [(64, 192), (128, 128)], ids=["glm_h64_k192", "ds_h128_k128"]
+)
+@pytest.mark.parametrize("pad_heads_extra", [0, 2])
+def test_qprep_vs_triton_two_dot(T, h_k, pad_heads_extra):
+    from sglang.kernels.ops.attention.qprep_bf16_fp8_sm90 import q8kv8_qprep_fwd
+    from sglang.kernels.ops.kvcache.cache_ops import (
+        absorbed_bmm_concat_cast_q_fp8,
+    )
+
+    H, K = h_k
+    q_nope, q_rope, w_kc = _make_inputs(T, H, K)
+    ph = H + pad_heads_extra
+    ref = torch.zeros((T, ph, N_LORA + ROPE), dtype=torch.float8_e4m3fn, device="cuda")
+    out = torch.zeros_like(ref)
+    absorbed_bmm_concat_cast_q_fp8(ref, q_nope, w_kc, q_rope, H, variant="two_dot")
+    q8kv8_qprep_fwd(out, q_nope, w_kc, q_rope, H)
+    torch.cuda.synchronize()
+
+    # rope half: identical conversion chain -> bit-exact.
+    assert torch.equal(
+        ref[:, :H, N_LORA:].contiguous().view(torch.uint8),
+        out[:, :H, N_LORA:].contiguous().view(torch.uint8),
+    ), "rope half must be bit-exact vs the Triton kernel"
+
+    # gate (a): nope half within tolerance on the fp32 view.
+    torch.testing.assert_close(
+        out[:, :H].to(torch.float32),
+        ref[:, :H].to(torch.float32),
+        atol=2e-2,
+        rtol=2e-2,
+    )
+
+    # padded head slice must stay untouched.
+    if pad_heads_extra:
+        assert out[:, H:].view(torch.uint8).max().item() == 0
+
+
+@requires_sm90
+@pytest.mark.parametrize(
+    "h_k", [(64, 192), (128, 128)], ids=["glm_h64_k192", "ds_h128_k128"]
+)
+def test_qprep_fp64_reference_parity(h_k):
+    from sglang.kernels.ops.attention.qprep_bf16_fp8_sm90 import q8kv8_qprep_fwd
+    from sglang.kernels.ops.kvcache.cache_ops import (
+        absorbed_bmm_concat_cast_q_fp8,
+    )
+
+    H, K = h_k
+    T = 437
+    q_nope, q_rope, w_kc = _make_inputs(T, H, K, seed=5678)
+    ref8 = torch.zeros((T, H, N_LORA + ROPE), dtype=torch.float8_e4m3fn, device="cuda")
+    out8 = torch.zeros_like(ref8)
+    absorbed_bmm_concat_cast_q_fp8(ref8, q_nope, w_kc, q_rope, H, variant="two_dot")
+    q8kv8_qprep_fwd(out8, q_nope, w_kc, q_rope, H)
+    torch.cuda.synchronize()
+
+    # gate (b): the CUDA kernel's fp8 must land as close to the exact bmm as
+    # the Triton kernel's (same quantization noise floor, ~1.8e-2 mean).
+    ref64 = torch.bmm(
+        q_nope.transpose(0, 1).to(torch.float64), w_kc.to(torch.float64)
+    ).transpose(0, 1)
+    err_tri = (ref8[..., :N_LORA].to(torch.float64) - ref64).abs().mean().item()
+    err_cuda = (out8[..., :N_LORA].to(torch.float64) - ref64).abs().mean().item()
+    assert (
+        err_cuda <= 1.05 * err_tri
+    ), f"CUDA fp64-ref mean |err| {err_cuda:.4e} exceeds Triton's {err_tri:.4e}"
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

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

@@ -192,6 +192,85 @@ 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
+# must match the fp32 reference on the truncated index range.
+@pytest.mark.skipif(
+    not _sm90_available(), reason="Q8KV8 sparse prefill requires SM90 CUDA"
+)
+@pytest.mark.parametrize(
+    "d_qk,s_q,topk,s_kv",
+    [
+        (576, 8, TOPK, S_KV),
+        (576, 65, 256, 592),
+        (512, 8, TOPK, S_KV),
+    ],
+)
+def test_sparse_mla_q8kv8_prefill_topk_length_only(
+    d_qk: int, s_q: int, topk: int, s_kv: int
+):
+    from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
+        sparse_mla_q8kv8_prefill_fwd,
+    )
+
+    q, kv, indices, sm_scale, q_scale, kv_scale, _, _ = _make_case(
+        d_qk, False, s_q=s_q, topk=topk, s_kv=s_kv
+    )
+    # Trailing pad runs of varying size, including a 1-valid-entry row (the
+    # production clamp(min=1) floor) and full rows.
+    lengths = [
+        topk if i % 3 == 0 else (1 if i % 3 == 1 else max(topk - 32, topk // 2))
+        for i in range(s_q)
+    ]
+    topk_length = torch.tensor(lengths, dtype=torch.int32, device="cuda")
+    for q_idx, valid_topk in enumerate(lengths):
+        if valid_topk < topk:
+            indices[q_idx, 0, valid_topk:] = -1
+
+    out, max_logits, lse = sparse_mla_q8kv8_prefill_fwd(
+        q=q,
+        kv=kv,
+        indices=indices,
+        sm_scale=sm_scale,
+        q_scale=q_scale,
+        kv_scale=kv_scale,
+        d_v=D_V,
+        attn_sink=None,
+        topk_length=topk_length,
+    )
+    out_full, max_logits_full, lse_full = sparse_mla_q8kv8_prefill_fwd(
+        q=q,
+        kv=kv,
+        indices=indices,
+        sm_scale=sm_scale,
+        q_scale=q_scale,
+        kv_scale=kv_scale,
+        d_v=D_V,
+        attn_sink=None,
+        topk_length=None,
+    )
+    torch.cuda.synchronize()
+
+    assert torch.equal(out, out_full)
+    assert torch.equal(max_logits, max_logits_full)
+    assert torch.equal(lse, lse_full)
+
+    ref, ref_max_logits, ref_lse = _torch_sparse_attention_ref(
+        q=q,
+        kv=kv,
+        indices=indices,
+        sm_scale=sm_scale,
+        q_scale=q_scale,
+        kv_scale=kv_scale,
+        attn_sink=None,
+        topk_length=topk_length,
+    )
+    torch.testing.assert_close(out.float(), ref, atol=8e-2, rtol=8e-2)
+    torch.testing.assert_close(max_logits.float(), ref_max_logits, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(lse.float(), ref_lse, atol=2e-3, rtol=2e-3)
+
+
 # Precision / accuracy: no-sink only because these metrics are intended to
 # approximate the current DeepSeek NSA E2E path. Sink behavior is still covered
 # above as kernel feature coverage, but sink-enabled precision numbers should
@@ -602,5 +681,44 @@ 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.
+@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,
+    )
+
+    generator = torch.Generator(device="cuda")
+    generator.manual_seed(4000 + s_q + topk)
+    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)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))