From 7bd9ccae5405432ac79b389380c33a40280bfdda Mon Sep 17 00:00:00 2001 From: Andreas Hassellof Date: Mon, 3 Aug 2026 18:06:45 +0200 Subject: [PATCH 1/4] Fix DSPARK SM120 decode dispatch for non-instantiated topk widths --- .../kernels/ops/attention/flash_mla_sm120.py | 88 +++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py index ec43c55089dd..be9c41663b2a 100644 --- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py +++ b/python/sglang/kernels/ops/attention/flash_mla_sm120.py @@ -470,6 +470,15 @@ def _split_kv_pages_to_64( ) +# CUTLASS SM120 sparse-MLA kernels are instantiated only for these topk +# widths (decode: (num_heads, topk) table with topk in {128, 512, 1024}; +# prefill orchestrator: {128, 512, 1024, 2048}). +_SUPPORTED_TOPK_WIDTHS = (128, 512, 1024, 2048) + +_noted_bucket_pad = False +_warned_triton_fb = False + + def _flash_mla_flashinfer( q, k_cache, @@ -526,6 +535,85 @@ def _flash_mla_flashinfer( else extra_indices ) + # --- Bucket alignment for non-instantiated topk widths --- + # The CUTLASS SM120 sparse-MLA kernels are instantiated only for a + # fixed set of topk widths (decode: topk in {128, 512, 1024}; prefill + # orchestrator: {128, 512, 1024, 2048}), and the prefill kernel + # additionally asserts num_tokens > 64. DSPARK's draft indexer emits + # topk=192, which is in no bucket on either path, so both the warmup + # draft pass and the draft CUDA-graph capture crash the server at boot + # ("num_tokens > 64" check fail / "Unsupported sparse-MLA prefill + # configuration ... topk=192"). Same failure family as + # sgl-project/sglang#33134 (DGX Spark, sm_121). + # + # Fix: right-pad the index tensor with -1 (the kernels' documented + # "skip" sentinel, see flashinfer csrc/sparse_mla_sm120.cu) up to the + # next instantiated bucket, and cap the scan via topk_length so the + # padding is never read. If a decode-sized batch is still not + # dispatchable (e.g. draft head_dim != 512), fall back to the Triton + # sparse-decode kernel for that call instead of crashing. + global _noted_bucket_pad, _warned_triton_fb + _topk = idx.shape[-1] + _d_qk = q.shape[-1] + if _d_qk == 512 and _topk not in _SUPPORTED_TOPK_WIDTHS: + _next_w = next((t for t in _SUPPORTED_TOPK_WIDTHS if t >= _topk), None) + if _next_w is not None: + if topk_length is None: + # Cap the scan at the true width so the -1 padding is + # never even read. + topk_length = torch.full((B,), _topk, dtype=torch.int32, device=dev) + idx = torch.nn.functional.pad(idx, (0, _next_w - _topk), value=-1) + if not _noted_bucket_pad: + _noted_bucket_pad = True + logger.info( + "SM120 sparse-MLA: padding topk %d -> %d (next " + "instantiated bucket, -1 skip sentinel; scan capped " + "via topk_length).", + _topk, + _next_w, + ) + + if B <= _FI_DECODE_MAX_TOKENS: + from flashinfer.mla._sparse_mla_sm120 import _decode_dsv4_dispatchable + + _extra_topk = extra_idx.shape[-1] if extra_idx is not None else 0 + if not _decode_dsv4_dispatchable( + B, H, idx.shape[-1], _d_qk, _PBS_DST, _extra_topk + ): + # Decode-sized but not coverable by the CUTLASS decode kernel + # even after bucket padding (e.g. head_dim != 512 or an + # uninstantiated head count) — the prefill kernel would reject + # num_tokens <= 64, so route to Triton instead of crashing. + if not _warned_triton_fb: + _warned_triton_fb = True + logger.warning( + "SM120 sparse-MLA: decode-sized batch not dispatchable " + "to CUTLASS decode kernel (num_tokens=%d heads=%d " + "topk=%d d_qk=%d) — using Triton fallback for these " + "calls.", + B, + H, + idx.shape[-1], + _d_qk, + ) + from sglang.kernels.ops.attention.flash_mla_sm120_triton import ( + flash_mla_sparse_decode_triton, + ) + + out, lse = flash_mla_sparse_decode_triton( + q, + k_cache, + indices, + topk_length, + attn_sink, + head_dim_v, + softmax_scale, + extra_k_cache, + extra_indices, + extra_topk_length, + ) + return (out, lse) + output = torch.empty(B, H, head_dim_v, dtype=torch.bfloat16, device=dev) out_lse = torch.empty(B, H, dtype=torch.float32, device=dev) From 21065dffe6cc3927e36bd5541339b346077bf29c Mon Sep 17 00:00:00 2001 From: Andreas Hassellof Date: Tue, 4 Aug 2026 07:59:00 +0200 Subject: [PATCH 2/4] Add SM120 topk-bucket dispatch tests; extract _next_topk_bucket helper Adds hermetic (CPU-only, CUDA_VISIBLE_DEVICES-independent) coverage for the topk-bucket padding + Triton-fallback dispatch this PR introduces: bucket arithmetic (192 -> 512, instantiated widths self-bucket, >2048 has no bucket, widths match the installed flashinfer decode table) and dispatch behaviour via recorders around the CUTLASS kernel and the Triton fallback (padded-to-512 call, -1 skip sentinel, topk_length capped to the true width, split-K scratch sized to the padded width, and the three fall-back-to-Triton geometries). To make the pad target testable, the inline "next instantiated width" search is extracted into a small pure helper, _next_topk_bucket(). Test authored by @efschu (github.com/efschu/htsglang); import paths adapted to this tree's sglang.kernels.ops.attention layout. Co-authored-by: efschu --- .../kernels/ops/attention/flash_mla_sm120.py | 9 +- .../test_flash_mla_sm120_topk_buckets.py | 229 ++++++++++++++++++ 2 files changed, 237 insertions(+), 1 deletion(-) create mode 100644 test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py index be9c41663b2a..00f61f920c1c 100644 --- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py +++ b/python/sglang/kernels/ops/attention/flash_mla_sm120.py @@ -479,6 +479,13 @@ def _split_kv_pages_to_64( _warned_triton_fb = False +def _next_topk_bucket(topk: int) -> Optional[int]: + """Smallest instantiated topk width >= ``topk``, or ``None`` if wider than + every kernel. The pad target for a request whose indexer topk (e.g. DSpark's + 192) is not itself an instantiated width.""" + return next((w for w in _SUPPORTED_TOPK_WIDTHS if w >= topk), None) + + def _flash_mla_flashinfer( q, k_cache, @@ -556,7 +563,7 @@ def _flash_mla_flashinfer( _topk = idx.shape[-1] _d_qk = q.shape[-1] if _d_qk == 512 and _topk not in _SUPPORTED_TOPK_WIDTHS: - _next_w = next((t for t in _SUPPORTED_TOPK_WIDTHS if t >= _topk), None) + _next_w = _next_topk_bucket(_topk) if _next_w is not None: if topk_length is None: # Cap the scan at the true width so the -1 padding is diff --git a/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py b/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py new file mode 100644 index 000000000000..613863b53d5d --- /dev/null +++ b/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SM120 sparse-MLA topk bucket alignment (coverage for #33407). + +Test authored by @efschu (github.com/efschu/htsglang), contributed here with +their attribution; import paths adapted from that fork's layout to this tree +(sglang.kernels.ops.attention) and _next_topk_bucket exposed as a named helper. + + +The CUTLASS SM120 sparse-MLA decode kernels are instantiated for a fixed set +of ``(num_heads, topk)`` pairs -- ``topk in {128, 512, 1024}`` -- and DSpark's +draft indexer emits ``topk=192``, which is in no bucket. Before the port, +``_flash_mla_flashinfer`` handed that width straight to +``sparse_mla_sm120_decode_dsv4`` and the server died at boot on the first +DSpark call. + +Everything here is hermetic (``CUDA_VISIBLE_DEVICES=99``, CPU tensors, no +kernel launch): the fix is index arithmetic plus a dispatch decision, and both +are observable without a card. The kernel and the Triton fallback are replaced +by recorders, so "which path was taken, with which index width" is an +observation rather than an inference. + +The KV cache in these tests already has ``page_size == 64``, which is the +production short-circuit that skips the page-split (``_PBS_DST``); that keeps +the test free of any Triton launch. +""" + +from __future__ import annotations + +import unittest + +import torch + +import flashinfer.mla._sparse_mla_sm120 as fi +from sglang.kernels.ops.attention import flash_mla_sm120 as fmod +from sglang.kernels.ops.attention import flash_mla_sm120_triton as tmod +from sglang.kernels.ops.attention.flash_mla_sm120 import ( + _NOPE_ROPE_STRIDE, + _PBS_DST, + _SCALE_STRIDE, + _SUPPORTED_TOPK_WIDTHS, + _next_topk_bucket, +) + +_BYTES_PER_TOKEN = _NOPE_ROPE_STRIDE + _SCALE_STRIDE +_HEAD_DIM_QK = 512 +_HEAD_DIM_V = 448 + + +class TestBucketArithmetic(unittest.TestCase): + """The whole correctness argument of the pad, without a kernel.""" + + def test_dspark_192_pads_to_the_next_instantiated_width(self): + self.assertEqual(_next_topk_bucket(192), 512) + + def test_an_instantiated_width_is_its_own_bucket(self): + for width in _SUPPORTED_TOPK_WIDTHS: + self.assertEqual(_next_topk_bucket(width), width) + + def test_the_bucket_is_never_narrower_than_the_request(self): + for topk in range(1, 2049): + bucket = _next_topk_bucket(topk) + self.assertIsNotNone(bucket, topk) + self.assertGreaterEqual(bucket, topk, topk) + + def test_the_bucket_is_the_smallest_that_fits(self): + for topk in (1, 127, 128, 129, 192, 511, 513, 1025, 2048): + bucket = _next_topk_bucket(topk) + narrower = [w for w in _SUPPORTED_TOPK_WIDTHS if w >= topk and w < bucket] + self.assertEqual(narrower, [], topk) + + def test_wider_than_every_kernel_has_no_bucket(self): + """Padding cannot rescue this; the dispatch check must catch it.""" + self.assertIsNone(_next_topk_bucket(2049)) + + def test_the_widths_match_the_installed_flashinfer_decode_table(self): + """Spread precondition: the table this port targets is really this one. + + A constant copied out of an upstream comment is worth nothing if the + installed kernels were built for other widths. + """ + table_widths = {topk for _, topk in fi._DECODE_DSV4_DISPATCH} + self.assertTrue(table_widths) + self.assertNotIn(192, table_widths, "the defect this port fixes is gone") + self.assertTrue( + table_widths.issubset(set(_SUPPORTED_TOPK_WIDTHS)), + f"instantiated decode widths {sorted(table_widths)} are not covered " + f"by _SUPPORTED_TOPK_WIDTHS {_SUPPORTED_TOPK_WIDTHS}", + ) + + +class _Recorders: + """Replace the CUTLASS kernel and the Triton fallback with recorders.""" + + def __enter__(self): + self.decode_calls = [] + self.triton_calls = [] + self._fi = fi.sparse_mla_sm120_decode_dsv4 + self._tr = tmod.flash_mla_sparse_decode_triton + + def decode(**kwargs): + self.decode_calls.append(kwargs) + return None + + def triton(q, k_cache, indices, topk_length, *args, **kwargs): + self.triton_calls.append((indices.shape, topk_length)) + b = q.shape[0] + h = q.shape[2] if q.ndim == 4 else q.shape[1] + return ( + torch.zeros(b, 1, h, _HEAD_DIM_V, dtype=torch.bfloat16), + torch.zeros(b, h, dtype=torch.float32), + ) + + fi.sparse_mla_sm120_decode_dsv4 = decode + tmod.flash_mla_sparse_decode_triton = triton + return self + + def __exit__(self, *exc): + fi.sparse_mla_sm120_decode_dsv4 = self._fi + tmod.flash_mla_sparse_decode_triton = self._tr + return False + + +def _call(topk: int, heads: int = 128, batch: int = 1, d_qk: int = _HEAD_DIM_QK): + """Drive ``_flash_mla_flashinfer`` on CPU and report what it reached.""" + q = torch.zeros(batch, 1, heads, d_qk, dtype=torch.bfloat16) + # page_size == _PBS_DST short-circuits the page-split, so no Triton kernel + # is launched anywhere on this path. + k_cache = torch.zeros(4, _PBS_DST, 1, _BYTES_PER_TOKEN, dtype=torch.uint8) + indices = torch.zeros(batch, topk, dtype=torch.int32) + with _Recorders() as rec: + out = fmod._flash_mla_flashinfer( + q, + k_cache, + indices, + None, # topk_length + None, # attn_sink + _HEAD_DIM_V, + d_qk ** (-0.5), + None, # extra_k_cache + None, # extra_indices + None, # extra_topk_length + ) + return rec, out + + +class TestDispatchOnCall(unittest.TestCase): + def setUp(self): + # The two "log once" latches are module state; reset them so a test + # order cannot decide whether a branch logs. + fmod._noted_bucket_pad = False + fmod._warned_triton_fb = False + + def test_topk_192_reaches_the_kernel_padded_to_512(self): + rec, _ = _call(192) + self.assertEqual(len(rec.decode_calls), 1) + self.assertEqual(rec.triton_calls, []) + kwargs = rec.decode_calls[0] + self.assertEqual(kwargs["indices"].shape[-1], 512) + self.assertTrue( + fi._decode_dsv4_dispatchable( + 1, 128, kwargs["indices"].shape[-1], _HEAD_DIM_QK, _PBS_DST, 0 + ), + "the padded width must be dispatchable -- that is the whole fix", + ) + + def test_the_padding_is_the_minus_one_skip_sentinel(self): + rec, _ = _call(192) + idx = rec.decode_calls[0]["indices"] + self.assertTrue(bool((idx[:, :192] == 0).all()), "real indices were altered") + self.assertTrue(bool((idx[:, 192:] == -1).all()), "pad is not the sentinel") + + def test_the_scan_is_capped_at_the_true_width(self): + """Without topk_length the kernel would read the -1 padding.""" + rec, _ = _call(192) + capped = rec.decode_calls[0]["topk_length"] + self.assertIsNotNone(capped, "topk_length must be synthesised by the pad") + self.assertEqual(capped.dtype, torch.int32) + self.assertEqual(capped.tolist(), [192]) + + def test_an_instantiated_width_is_passed_through_untouched(self): + """Neutrality: the pre-port behaviour for every width that worked.""" + rec, _ = _call(512) + self.assertEqual(len(rec.decode_calls), 1) + self.assertEqual(rec.decode_calls[0]["indices"].shape[-1], 512) + self.assertIsNone( + rec.decode_calls[0]["topk_length"], + "an already-instantiated width must not grow a synthetic cap", + ) + + def test_the_split_k_scratch_covers_the_padded_width(self): + """mid_out/mid_lse are sized from the width the kernel actually scans.""" + rec, _ = _call(192) + kwargs = rec.decode_calls[0] + self.assertEqual(kwargs["mid_out"].shape[2], 512 // 64) + self.assertEqual(kwargs["mid_lse"].shape[2], 512 // 64) + + def test_an_undispatchable_geometry_falls_back_to_triton(self): + """d_qk != 512 is out of the CUTLASS table at every topk width.""" + rec, out = _call(128, d_qk=256) + self.assertEqual(rec.decode_calls, []) + self.assertEqual(len(rec.triton_calls), 1) + self.assertIsNotNone(out[0]) + + def test_a_batch_above_the_decode_maximum_falls_back_to_triton(self): + rec, _ = _call(128, batch=fi._DECODE_MAX_TOKENS + 1) + self.assertEqual(rec.decode_calls, []) + self.assertEqual(len(rec.triton_calls), 1) + + def test_an_uninstantiated_head_count_falls_back_to_triton(self): + heads = next( + h + for h in range(1, 4096) + if (h, 512) not in fi._DECODE_DSV4_DISPATCH + and (h, 128) not in fi._DECODE_DSV4_DISPATCH + ) + rec, _ = _call(128, heads=heads) + self.assertEqual(rec.decode_calls, []) + self.assertEqual(len(rec.triton_calls), 1) + + def test_the_fallback_gets_the_unpadded_indices(self): + """Triton reads the real width; handing it the -1 pad would be wrong.""" + rec, _ = _call(192, d_qk=256) + self.assertEqual(len(rec.triton_calls), 1) + shape, _ = rec.triton_calls[0] + self.assertEqual(shape[-1], 192) + + +if __name__ == "__main__": + unittest.main() From bb72537d5f70df69a2f077012ecdafe737b58928 Mon Sep 17 00:00:00 2001 From: Andreas Hassellof Date: Tue, 4 Aug 2026 15:30:59 +0200 Subject: [PATCH 3/4] Fix lint: isort import order + register CI suite for SM120 topk-bucket test isort grouped flashinfer/torch third-party imports and the new test/registered/ file needs a CI registry call (check-registered-tests hook). Register it on base-b/1-gpu-small, matching the sibling SM120 kernel test. --- .../ops/attention/test_flash_mla_sm120_topk_buckets.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py b/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py index 613863b53d5d..6c1c1dc0e014 100644 --- a/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py +++ b/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py @@ -28,9 +28,9 @@ import unittest +import flashinfer.mla._sparse_mla_sm120 as fi import torch -import flashinfer.mla._sparse_mla_sm120 as fi from sglang.kernels.ops.attention import flash_mla_sm120 as fmod from sglang.kernels.ops.attention import flash_mla_sm120_triton as tmod from sglang.kernels.ops.attention.flash_mla_sm120 import ( @@ -40,6 +40,9 @@ _SUPPORTED_TOPK_WIDTHS, _next_topk_bucket, ) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small") _BYTES_PER_TOKEN = _NOPE_ROPE_STRIDE + _SCALE_STRIDE _HEAD_DIM_QK = 512 From 1a4cba439d56767bee98f290f2dad696d9ccf5aa Mon Sep 17 00:00:00 2001 From: Andreas Hassellof Date: Thu, 6 Aug 2026 05:21:51 +0200 Subject: [PATCH 4/4] Fix SM120 topk-bucket test port: intercept at the paged-attention entry The port from the htsglang fork kept that fork's recorder seam and geometry, and never ran green on this tree: - The recorder patched fi.sparse_mla_sm120_decode_dsv4, but this tree's wrapper calls _sparse_mla_sm120_paged_attention, which JIT-builds the SM120 module before any Python-level dispatch -- the old seam is only reachable after a successful build on an SM12x card. Intercepting the paged-attention entry (imported per-call by _flash_mla_flashinfer) keeps the suite genuinely hermetic. - _HEAD_DIM_V was 448; flashinfer hard-requires d_v == 512 on this path (_require_d_v_512), so every dispatch test died in the entry gate. - The fork's decode-only wrapper sent batch > 64 to Triton; this tree hands prefill-sized batches to the CUTLASS prefill orchestrator via the same entry. The expectation now pins that (no decode scratch, no Triton detour). All 15 cases pass on a card-less host (CUDA_VISIBLE_DEVICES=99, stock flashinfer 0.6.15.post1, no JIT). --- .../test_flash_mla_sm120_topk_buckets.py | 108 +++++++++++++----- 1 file changed, 77 insertions(+), 31 deletions(-) diff --git a/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py b/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py index 6c1c1dc0e014..a5dd24510a8f 100644 --- a/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py +++ b/test/registered/kernels/ops/attention/test_flash_mla_sm120_topk_buckets.py @@ -2,22 +2,28 @@ """SM120 sparse-MLA topk bucket alignment (coverage for #33407). Test authored by @efschu (github.com/efschu/htsglang), contributed here with -their attribution; import paths adapted from that fork's layout to this tree -(sglang.kernels.ops.attention) and _next_topk_bucket exposed as a named helper. - +their attribution; adapted from that fork's layout to this tree: import paths +(sglang.kernels.ops.attention), _next_topk_bucket exposed as a named helper, +the recorder seam (this tree calls flashinfer's +``_sparse_mla_sm120_paged_attention`` entry, which JIT-builds the SM120 module +before any Python-level dispatch, so interception must happen at that entry -- +not at ``sparse_mla_sm120_decode_dsv4``), the geometry (flashinfer requires +``d_v == 512`` on this path), and the batch > 64 expectation (the fork's +decode-only wrapper sent those to Triton; this tree hands them to the CUTLASS +prefill orchestrator via the same entry). The CUTLASS SM120 sparse-MLA decode kernels are instantiated for a fixed set of ``(num_heads, topk)`` pairs -- ``topk in {128, 512, 1024}`` -- and DSpark's -draft indexer emits ``topk=192``, which is in no bucket. Before the port, +draft indexer emits ``topk=192``, which is in no bucket. Before the fix, ``_flash_mla_flashinfer`` handed that width straight to -``sparse_mla_sm120_decode_dsv4`` and the server died at boot on the first -DSpark call. +``_sparse_mla_sm120_paged_attention`` and the server died at boot on the +first DSpark call. Everything here is hermetic (``CUDA_VISIBLE_DEVICES=99``, CPU tensors, no -kernel launch): the fix is index arithmetic plus a dispatch decision, and both -are observable without a card. The kernel and the Triton fallback are replaced -by recorders, so "which path was taken, with which index width" is an -observation rather than an inference. +kernel launch or JIT build): the fix is index arithmetic plus a dispatch +decision, and both are observable without a card. The flashinfer entry point +and the Triton fallback are replaced by recorders, so "which path was taken, +with which index width" is an observation rather than an inference. The KV cache in these tests already has ``page_size == 64``, which is the production short-circuit that skips the page-split (``_PBS_DST``); that keeps @@ -46,7 +52,9 @@ _BYTES_PER_TOKEN = _NOPE_ROPE_STRIDE + _SCALE_STRIDE _HEAD_DIM_QK = 512 -_HEAD_DIM_V = 448 +# flashinfer's SM120 sparse-MLA path hard-requires d_v == 512 +# (_require_d_v_512); this is the only geometry the wrapper can pass on. +_HEAD_DIM_V = 512 class TestBucketArithmetic(unittest.TestCase): @@ -92,16 +100,46 @@ def test_the_widths_match_the_installed_flashinfer_decode_table(self): class _Recorders: - """Replace the CUTLASS kernel and the Triton fallback with recorders.""" + """Replace the flashinfer entry point and the Triton fallback with recorders. + + ``_flash_mla_flashinfer`` imports ``_sparse_mla_sm120_paged_attention`` + from ``flashinfer.mla._sparse_mla_sm120`` inside the call, so patching the + module attribute intercepts every call -- before flashinfer's JIT build, + which is what keeps the suite runnable on a card-less host. + """ def __enter__(self): - self.decode_calls = [] + self.paged_calls = [] self.triton_calls = [] - self._fi = fi.sparse_mla_sm120_decode_dsv4 + self._fi = fi._sparse_mla_sm120_paged_attention self._tr = tmod.flash_mla_sparse_decode_triton - def decode(**kwargs): - self.decode_calls.append(kwargs) + def paged_attention( + q, + kv_cache, + indices, + output, + out_lse, + sm_scale, + *, + d_v, + topk_length=None, + attn_sink=None, + extra_kv_cache=None, + extra_indices=None, + extra_topk_length=None, + mid_out=None, + mid_lse=None, + ): + self.paged_calls.append( + dict( + indices=indices, + topk_length=topk_length, + d_v=d_v, + mid_out=mid_out, + mid_lse=mid_lse, + ) + ) return None def triton(q, k_cache, indices, topk_length, *args, **kwargs): @@ -113,12 +151,12 @@ def triton(q, k_cache, indices, topk_length, *args, **kwargs): torch.zeros(b, h, dtype=torch.float32), ) - fi.sparse_mla_sm120_decode_dsv4 = decode + fi._sparse_mla_sm120_paged_attention = paged_attention tmod.flash_mla_sparse_decode_triton = triton return self def __exit__(self, *exc): - fi.sparse_mla_sm120_decode_dsv4 = self._fi + fi._sparse_mla_sm120_paged_attention = self._fi tmod.flash_mla_sparse_decode_triton = self._tr return False @@ -155,9 +193,9 @@ def setUp(self): def test_topk_192_reaches_the_kernel_padded_to_512(self): rec, _ = _call(192) - self.assertEqual(len(rec.decode_calls), 1) + self.assertEqual(len(rec.paged_calls), 1) self.assertEqual(rec.triton_calls, []) - kwargs = rec.decode_calls[0] + kwargs = rec.paged_calls[0] self.assertEqual(kwargs["indices"].shape[-1], 512) self.assertTrue( fi._decode_dsv4_dispatchable( @@ -168,14 +206,14 @@ def test_topk_192_reaches_the_kernel_padded_to_512(self): def test_the_padding_is_the_minus_one_skip_sentinel(self): rec, _ = _call(192) - idx = rec.decode_calls[0]["indices"] + idx = rec.paged_calls[0]["indices"] self.assertTrue(bool((idx[:, :192] == 0).all()), "real indices were altered") self.assertTrue(bool((idx[:, 192:] == -1).all()), "pad is not the sentinel") def test_the_scan_is_capped_at_the_true_width(self): """Without topk_length the kernel would read the -1 padding.""" rec, _ = _call(192) - capped = rec.decode_calls[0]["topk_length"] + capped = rec.paged_calls[0]["topk_length"] self.assertIsNotNone(capped, "topk_length must be synthesised by the pad") self.assertEqual(capped.dtype, torch.int32) self.assertEqual(capped.tolist(), [192]) @@ -183,31 +221,39 @@ def test_the_scan_is_capped_at_the_true_width(self): def test_an_instantiated_width_is_passed_through_untouched(self): """Neutrality: the pre-port behaviour for every width that worked.""" rec, _ = _call(512) - self.assertEqual(len(rec.decode_calls), 1) - self.assertEqual(rec.decode_calls[0]["indices"].shape[-1], 512) + self.assertEqual(len(rec.paged_calls), 1) + self.assertEqual(rec.paged_calls[0]["indices"].shape[-1], 512) self.assertIsNone( - rec.decode_calls[0]["topk_length"], + rec.paged_calls[0]["topk_length"], "an already-instantiated width must not grow a synthetic cap", ) def test_the_split_k_scratch_covers_the_padded_width(self): """mid_out/mid_lse are sized from the width the kernel actually scans.""" rec, _ = _call(192) - kwargs = rec.decode_calls[0] + kwargs = rec.paged_calls[0] self.assertEqual(kwargs["mid_out"].shape[2], 512 // 64) self.assertEqual(kwargs["mid_lse"].shape[2], 512 // 64) def test_an_undispatchable_geometry_falls_back_to_triton(self): """d_qk != 512 is out of the CUTLASS table at every topk width.""" rec, out = _call(128, d_qk=256) - self.assertEqual(rec.decode_calls, []) + self.assertEqual(rec.paged_calls, []) self.assertEqual(len(rec.triton_calls), 1) self.assertIsNotNone(out[0]) - def test_a_batch_above_the_decode_maximum_falls_back_to_triton(self): + def test_a_batch_above_the_decode_maximum_goes_to_the_prefill_path(self): + """num_tokens > 64 is exactly what the prefill orchestrator accepts. + + The pre-fix crash was the *decode*-sized batch falling through to the + prefill kernel; a genuinely prefill-sized batch must keep using it -- + with no decode split-K scratch and no Triton detour. + """ rec, _ = _call(128, batch=fi._DECODE_MAX_TOKENS + 1) - self.assertEqual(rec.decode_calls, []) - self.assertEqual(len(rec.triton_calls), 1) + self.assertEqual(len(rec.paged_calls), 1) + self.assertEqual(rec.triton_calls, []) + self.assertIsNone(rec.paged_calls[0]["mid_out"]) + self.assertIsNone(rec.paged_calls[0]["mid_lse"]) def test_an_uninstantiated_head_count_falls_back_to_triton(self): heads = next( @@ -217,7 +263,7 @@ def test_an_uninstantiated_head_count_falls_back_to_triton(self): and (h, 128) not in fi._DECODE_DSV4_DISPATCH ) rec, _ = _call(128, heads=heads) - self.assertEqual(rec.decode_calls, []) + self.assertEqual(rec.paged_calls, []) self.assertEqual(len(rec.triton_calls), 1) def test_the_fallback_gets_the_unpadded_indices(self):