[II] Fix DS4 C128A CUDA-graph metadata stride

tests/kernels/attention/test_flashmla_sparse.py

@@ -4,17 +4,20 @@
 import torch
 
 
-def test_deepseek_v4_c128a_dynamic_topk_packed_buffers():
+def test_deepseek_v4_c128a_metadata_preserves_capacity_stride():
     from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata
 
     device = torch.device("cuda")
     capacity_width = 256
     active_width = 128
-    global_decode_buffer = torch.empty(
-        (2, capacity_width), dtype=torch.int32, device=device
+    untouched = -77
+    global_decode_buffer = torch.full(
+        (2, capacity_width), untouched, dtype=torch.int32, device=device
     )
     decode_lens_buffer = torch.empty(2, dtype=torch.int32, device=device)
-    prefill_buffer = torch.empty((2, capacity_width), dtype=torch.int32, device=device)
+    prefill_buffer = torch.full(
+        (2, capacity_width), untouched, dtype=torch.int32, device=device
+    )
 
     global_decode, decode_lens, prefill_local = build_c128a_topk_metadata(
         positions=torch.tensor([255, 511], dtype=torch.int64, device=device),
@@ -31,15 +34,17 @@ def test_deepseek_v4_c128a_dynamic_topk_packed_buffers():
         max_compressed_tokens=active_width,
     )
 
-    assert global_decode.shape == (1, active_width)
-    assert prefill_local.shape == (1, active_width)
-    assert global_decode.stride() == (active_width, 1)
-    assert prefill_local.stride() == (active_width, 1)
+    assert global_decode.shape == (1, capacity_width)
+    assert prefill_local.shape == (1, capacity_width)
+    assert global_decode.stride() == (capacity_width, 1)
+    assert prefill_local.stride() == (capacity_width, 1)
     assert global_decode[0, :2].cpu().tolist() == [768, 769]
     assert decode_lens.cpu().tolist() == [2]
     assert prefill_local[0, :4].cpu().tolist() == list(range(4))
-    assert torch.all(global_decode[0, 2:] == -1)
-    assert torch.all(prefill_local[0, 4:] == -1)
+    assert torch.all(global_decode[0, 2:active_width] == -1)
+    assert torch.all(prefill_local[0, 4:active_width] == -1)
+    assert torch.all(global_decode[0, active_width:] == untouched)
+    assert torch.all(prefill_local[0, active_width:] == untouched)
 
 
 def test_sparse_flashmla_metadata_smoke():

tests/v1/attention/test_deepseek_v4_sparse_mla_metadata.py

@@ -14,6 +14,63 @@
 from vllm.v1.kv_cache_interface import MLAAttentionSpec
 
 
+def test_c128a_builder_uses_capture_time_row_capacity(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    import vllm.models.deepseek_v4.sparse_mla as sparse_mla
+
+    capacity_width = 256
+    builder = object.__new__(DeepseekV4FlashMLAMetadataBuilder)
+    builder.deepseek_v4_decode_threshold = 1
+    builder.compress_ratio = 128
+    builder.c128a_max_compressed = capacity_width
+    builder.kv_cache_spec = SimpleNamespace(block_size=capacity_width * 128)
+    builder.c128a_global_decode_buffer = torch.empty(2, capacity_width)
+    builder.c128a_decode_lens_buffer = torch.empty(2, dtype=torch.int32)
+    builder.c128a_prefill_buffer = torch.empty(2, capacity_width)
+
+    monkeypatch.setattr(
+        sparse_mla,
+        "split_decodes_and_prefills",
+        lambda *_args, **_kwargs: (2, 0, 2, 0),
+    )
+    observed: dict[str, int] = {}
+
+    def record_width(*_args, max_compressed_tokens: int, **_kwargs):
+        observed["width"] = max_compressed_tokens
+        return (
+            torch.empty(2, capacity_width, dtype=torch.int32),
+            torch.empty(2, dtype=torch.int32),
+            torch.empty(0, capacity_width, dtype=torch.int32),
+        )
+
+    monkeypatch.setattr(sparse_mla, "build_c128a_topk_metadata", record_width)
+    common_metadata = SimpleNamespace(
+        # This active context would select a 128-wide packed row in the
+        # incompatible implementation.
+        max_seq_len=512,
+        positions=torch.arange(2),
+        block_table_tensor=torch.empty(2, 1, dtype=torch.int32),
+        slot_mapping=torch.arange(2),
+    )
+
+    metadata = builder._build_c128a_metadata(
+        common_metadata,
+        req_id_per_token=torch.zeros(2, dtype=torch.int32),
+        actual_num_query_tokens=2,
+        dcp_world_size=1,
+        dcp_rank=0,
+        cp_kv_cache_interleave_size=1,
+    )
+
+    assert observed["width"] == capacity_width
+    assert metadata["c128a_global_decode_topk_indices"].shape == (
+        2,
+        1,
+        capacity_width,
+    )
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 @pytest.mark.parametrize(
     ("dcp_rank", "valid_offset", "expected_slot"),

vllm/models/deepseek_v4/sparse_mla.py

@@ -311,13 +311,6 @@ def _build_c128a_metadata(
         assert cm.positions is not None, (
             "positions is required for C128A metadata build"
         )
-        active_topk_width = min(
-            max(
-                triton.next_power_of_2(max(cm.max_seq_len // self.compress_ratio, 1)),
-                _C128A_TOPK_ALIGNMENT,
-            ),
-            self.c128a_max_compressed,
-        )
         block_size = self.kv_cache_spec.block_size // self.compress_ratio
         global_decode, decode_lens, prefill_local = build_c128a_topk_metadata(
             cm.positions[:num_total],
@@ -331,7 +324,7 @@ def _build_c128a_metadata(
             self.c128a_global_decode_buffer,
             self.c128a_decode_lens_buffer,
             self.c128a_prefill_buffer,
-            max_compressed_tokens=active_topk_width,
+            max_compressed_tokens=self.c128a_max_compressed,
             dcp_world_size=dcp_world_size,
             dcp_rank=dcp_rank,
             cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
@@ -370,30 +363,26 @@ def build_c128a_topk_metadata(
     Decode tokens: position → block_table lookup → global slot ids + topk_lens.
     Prefill tokens: position → local indices [0, ..., n-1, -1, ...].
 
-    Writes into packed views of pre-allocated buffers for CUDA graph stability.
+    Writes into capacity-strided views of pre-allocated buffers. FULL CUDA
+    graphs capture consumers with this row stride, so runtime metadata must
+    preserve it independently of the active sequence length.
     """
     num_tokens = positions.shape[0]
     num_prefill_tokens = num_tokens - num_decode_tokens
 
-    # view(-1) as 1-d array and then expanded to
-    # [num_decode_tokens, max_compressed_tokens]
-    global_decode = global_decode_buffer.view(-1)[
-        : num_decode_tokens * max_compressed_tokens
-    ].view(num_decode_tokens, max_compressed_tokens)
+    global_decode = global_decode_buffer[:num_decode_tokens]
     decode_lens = decode_lens_buffer[:num_decode_tokens]
-    prefill_local = prefill_buffer.view(-1)[
-        : num_prefill_tokens * max_compressed_tokens
-    ].view(num_prefill_tokens, max_compressed_tokens)
+    prefill_local = prefill_buffer[:num_prefill_tokens]
 
     if num_tokens == 0:
         return global_decode, decode_lens, prefill_local
 
     _build_c128a_topk_metadata_kernel[(num_tokens,)](
         global_decode_buffer,
-        max_compressed_tokens,
+        global_decode_buffer.stride(0),
         decode_lens_buffer,
         prefill_buffer,
-        max_compressed_tokens,
+        prefill_buffer.stride(0),
         positions,
         compress_ratio,
         max_compressed_tokens,