[SM120] Split the DSv4 extra KV cache to pbs=64 for sparse-MLA prefill

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

@@ -247,6 +247,24 @@ def _flash_mla_sm120_prefill(
         if extra_indices is not None and extra_indices.dim() == 3
         else extra_indices
     )
+    # The FlashInfer SM120 prefill dispatcher (dispatch_dsv4_dual) only
+    # instantiates extra_page_block_size of 64 or 2. The DSv4 secondary
+    # (C4A/C128A) cache is paged at 128, so -- unlike the decode path, which
+    # accepts 128 -- prefill must split it to 64 as well. Uses a distinct
+    # buffer key so it does not clobber the main cache's split buffer.
+    if extra_kv_u8 is not None and extra_k_cache.ndim >= 3:
+        extra_src_pbs = (
+            extra_k_cache.shape[2]
+            if extra_k_cache.ndim == 4 and extra_k_cache.shape[1] == 1
+            else extra_k_cache.shape[1]
+        )
+        if extra_src_pbs != _PBS_DST:
+            extra_kv_u8 = _split_kv_pages_to_64(
+                extra_kv_u8,
+                extra_src_pbs,
+                touched_indices=extra_idx,
+                key_suffix=":extra",
+            )
     output = q2.new_empty((num_tokens, num_heads, head_dim_v), dtype=torch.bfloat16)
     out_lse = torch.empty((num_tokens, num_heads), dtype=torch.float32, device=dev)
     _sparse_mla_sm120_paged_attention(
@@ -478,6 +496,7 @@ def _split_kv_pages_to_64(
     kv_u8: torch.Tensor,
     src_pbs: int,
     touched_indices: Optional[torch.Tensor] = None,
+    key_suffix: str = "",
 ) -> torch.Tensor:
     """Split pbs=N footer-format pages into pbs=64 footer-format pages.
 
@@ -501,7 +520,7 @@ def _split_kv_pages_to_64(
     # Pre-allocated grow-only buffer for page-split output per device.
     dev = kv_u8.device
     buffers = get_resources().buffers
-    key = f"flash_mla_sm120_split:{dev}"
+    key = f"flash_mla_sm120_split{key_suffix}:{dev}"
     buf = buffers.get(key)
     if buf is None or buf.shape[0] < num_dst_pages:
         # The first allocation can happen under inference mode (autotune), but
@@ -530,7 +549,7 @@ def _split_kv_pages_to_64(
     if use_mask:
         # Persistent per-device int8 mask, zeroed each call (cheap memset,
         # captured cleanly by CUDA graph). 1 = page is referenced this step.
-        mkey = f"flash_mla_sm120_mask:{dev}"
+        mkey = f"flash_mla_sm120_mask{key_suffix}:{dev}"
         mbuf = buffers.get(mkey)
         if mbuf is None or mbuf.shape[0] < N:
             # The first allocation can happen under inference mode (autotune),

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

@@ -0,0 +1,214 @@
+"""SM120 FlashMLA page-split tests: footer round-trip and buffer namespacing.
+
+`_split_kv_pages_to_64` rewrites footer-format KV pages from an arbitrary source
+page_block_size (pbs) down to pbs=64, which is the only main-cache page size the
+FlashInfer SM120 sparse-MLA kernels instantiate.
+
+DSv4 drives this helper from two different caches in a single forward:
+  - the main SWA cache at pbs=256, and
+  - the secondary (C4A/C128A) cache at pbs=128, which the prefill dispatcher
+    (`dispatch_dsv4_dual`) also requires at pbs=64.
+Both share one per-device persistent output buffer, so the call sites must be
+namespaced via `key_suffix` or the second split silently clobbers the first.
+
+Coverage:
+- Byte-exact round-trip of the data + scale footer regions for pbs=256 and
+  pbs=128 (the extra cache's page size)
+- `key_suffix` isolates the output and mask buffers, so a pbs=128 extra-cache
+  split does not corrupt a preceding pbs=256 main-cache split
+- The default `key_suffix=""` keeps the historic buffer key names
+- `touched_indices` masking still restricts copies to referenced pages
+
+Pure-Triton page copies on any CUDA GPU -- no SM120 hardware required.
+"""
+
+from __future__ import annotations
+
+import math
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.flash_mla_sm120 import (
+    _BYTES_PER_DST_PAGE_PADDED,
+    _NOPE_ROPE_STRIDE,
+    _PBS_DST,
+    _SCALE_STRIDE,
+    _split_kv_pages_to_64,
+)
+from sglang.srt.runtime_context import get_resources
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=8, stage="base-b", runner_config="1-gpu-small")
+
+
+def _page_bytes(pbs: int) -> int:
+    """Footer-format bytes per page, padded to _NOPE_ROPE_STRIDE alignment."""
+    raw = pbs * _NOPE_ROPE_STRIDE + pbs * _SCALE_STRIDE
+    return math.ceil(raw / _NOPE_ROPE_STRIDE) * _NOPE_ROPE_STRIDE
+
+
+def _build_cache(num_pages: int, pbs: int, device, seed: int = 0):
+    """Build a random footer-format cache.
+
+    Returns ``(raw, view)`` where ``raw`` is the (num_pages, page_bytes) uint8
+    backing store and ``view`` is the (num_pages, pbs, 1, bytes_per_token) NHD
+    view that production passes in (the last dim is a fiction for downstream
+    consumers; only stride(0) is load-bearing).
+    """
+    pb = _page_bytes(pbs)
+    gen = torch.Generator(device="cpu").manual_seed(seed)
+    raw = torch.randint(0, 256, (num_pages, pb), dtype=torch.uint8, generator=gen).to(
+        device
+    )
+    bpt = _NOPE_ROPE_STRIDE + _SCALE_STRIDE
+    view = raw.as_strided((num_pages, pbs, 1, bpt), (pb, bpt, bpt, 1))
+    return raw, view
+
+
+def _as_raw(split_out: torch.Tensor) -> torch.Tensor:
+    """Recover the flat (num_dst_pages, page_bytes) view of a split result."""
+    return split_out.as_strided(
+        (split_out.shape[0], _BYTES_PER_DST_PAGE_PADDED),
+        (_BYTES_PER_DST_PAGE_PADDED, 1),
+    )
+
+
+def _token_data(raw: torch.Tensor, pbs: int, token: int) -> torch.Tensor:
+    page, off = divmod(token, pbs)
+    return raw[page, off * _NOPE_ROPE_STRIDE : (off + 1) * _NOPE_ROPE_STRIDE]
+
+
+def _token_scale(raw: torch.Tensor, pbs: int, token: int) -> torch.Tensor:
+    page, off = divmod(token, pbs)
+    base = pbs * _NOPE_ROPE_STRIDE + off * _SCALE_STRIDE
+    return raw[page, base : base + _SCALE_STRIDE]
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+class TestFlashMLASM120PageSplit(CustomTestCase):
+    def setUp(self):
+        # The split buffers are grow-only and process-persistent; clear them so
+        # each test observes allocation behaviour from a known state.
+        get_resources().buffers.clear()
+        self.device = torch.device("cuda")
+
+    def _assert_round_trip(self, num_pages: int, src_pbs: int, key_suffix: str = ""):
+        src_raw, src_view = _build_cache(num_pages, src_pbs, self.device)
+        out = _split_kv_pages_to_64(src_view, src_pbs, key_suffix=key_suffix)
+
+        ratio = src_pbs // _PBS_DST
+        self.assertEqual(out.shape[0], num_pages * ratio)
+        dst_raw = _as_raw(out)
+
+        # Token indices are invariant under the split: a token's absolute index
+        # addresses the same logical slot before and after.
+        for token in range(num_pages * src_pbs):
+            self.assertTrue(
+                torch.equal(
+                    _token_data(src_raw, src_pbs, token),
+                    _token_data(dst_raw, _PBS_DST, token),
+                ),
+                f"data mismatch at token {token} (src_pbs={src_pbs})",
+            )
+            self.assertTrue(
+                torch.equal(
+                    _token_scale(src_raw, src_pbs, token),
+                    _token_scale(dst_raw, _PBS_DST, token),
+                ),
+                f"scale mismatch at token {token} (src_pbs={src_pbs})",
+            )
+
+    def test_round_trip_main_cache_pbs256(self):
+        """Regression guard for the pre-existing main-cache (pbs=256) path."""
+        self._assert_round_trip(num_pages=3, src_pbs=256)
+
+    def test_round_trip_extra_cache_pbs128(self):
+        """The DSv4 secondary cache is paged at 128 and must split to 64."""
+        self._assert_round_trip(num_pages=3, src_pbs=128)
+
+    def test_key_suffix_isolates_buffers(self):
+        """An extra-cache split must not clobber a preceding main-cache split.
+
+        Sized so the pbs=128 result (8 dst pages) fits inside the pbs=256 result
+        (16 dst pages): with a shared buffer key the second call would reuse and
+        overwrite the first call's pages rather than reallocating.
+        """
+        main_raw, main_view = _build_cache(4, 256, self.device, seed=1)
+        extra_raw, extra_view = _build_cache(4, 128, self.device, seed=2)
+
+        main_out = _split_kv_pages_to_64(main_view, 256)
+        extra_out = _split_kv_pages_to_64(extra_view, 128, key_suffix=":extra")
+
+        self.assertEqual(main_out.shape[0], 16)
+        self.assertEqual(extra_out.shape[0], 8)
+        self.assertNotEqual(
+            main_out.data_ptr(),
+            extra_out.data_ptr(),
+            "main and extra splits must not share an output buffer",
+        )
+
+        # The main split must still read back correctly after the extra split.
+        main_dst = _as_raw(main_out)
+        for token in range(4 * 256):
+            self.assertTrue(
+                torch.equal(
+                    _token_data(main_raw, 256, token),
+                    _token_data(main_dst, _PBS_DST, token),
+                ),
+                f"main cache corrupted at token {token} by the extra-cache split",
+            )
+
+        # ... and the extra split must be correct in its own right.
+        extra_dst = _as_raw(extra_out)
+        for token in range(4 * 128):
+            self.assertTrue(
+                torch.equal(
+                    _token_data(extra_raw, 128, token),
+                    _token_data(extra_dst, _PBS_DST, token),
+                ),
+                f"extra cache mismatch at token {token}",
+            )
+
+    def test_buffer_key_names(self):
+        """Default key_suffix preserves the historic key; ':extra' adds one."""
+        _, main_view = _build_cache(2, 256, self.device)
+        _, extra_view = _build_cache(2, 128, self.device)
+        idx = torch.arange(64, dtype=torch.int32, device=self.device)
+
+        _split_kv_pages_to_64(main_view, 256, touched_indices=idx)
+        _split_kv_pages_to_64(extra_view, 128, touched_indices=idx, key_suffix=":extra")
+
+        keys = set(get_resources().buffers)
+        # Keys are built from the tensor's device, which is always indexed
+        # (e.g. "cuda:0") regardless of how the test spelled it.
+        dev = main_view.device
+        self.assertIn(f"flash_mla_sm120_split:{dev}", keys)
+        self.assertIn(f"flash_mla_sm120_mask:{dev}", keys)
+        self.assertIn(f"flash_mla_sm120_split:extra:{dev}", keys)
+        self.assertIn(f"flash_mla_sm120_mask:extra:{dev}", keys)
+
+    def test_touched_indices_limits_copies(self):
+        """Masked splits copy only referenced source pages."""
+        src_raw, src_view = _build_cache(4, 128, self.device, seed=3)
+
+        # Reference only page 2 (tokens 256..383).
+        idx = torch.tensor([256, 300, 383], dtype=torch.int32, device=self.device)
+        out = _split_kv_pages_to_64(
+            src_view, 128, touched_indices=idx, key_suffix=":extra"
+        )
+        dst_raw = _as_raw(out)
+
+        for token in range(256, 384):
+            self.assertTrue(
+                torch.equal(
+                    _token_data(src_raw, 128, token),
+                    _token_data(dst_raw, _PBS_DST, token),
+                ),
+                f"touched page not copied at token {token}",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()