[Bugfix][DSv4] Revert adaptive C128A metadata packing (#51318)

Signed-off-by: tobymao <toby.mao@gmail.com>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

tests/kernels/attention/test_flashmla_sparse.py

@@ -4,43 +4,6 @@ import pytest
 import torch
 
 
-def test_deepseek_v4_c128a_dynamic_topk_packed_buffers():
-    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
-    )
-    decode_lens_buffer = torch.empty(2, dtype=torch.int32, device=device)
-    prefill_buffer = torch.empty((2, capacity_width), 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),
-        compress_ratio=128,
-        num_decode_tokens=1,
-        token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32, device=device),
-        block_table=torch.tensor([[3]], dtype=torch.int32, device=device),
-        block_size=capacity_width,
-        slot_mapping=torch.tensor([0, 1], dtype=torch.int64, device=device),
-        global_decode_buffer=global_decode_buffer,
-        decode_lens_buffer=decode_lens_buffer,
-        prefill_buffer=prefill_buffer,
-        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[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)
-
-
 def test_sparse_flashmla_metadata_smoke():
     import vllm.v1.attention.ops.flashmla as fm
 

vllm/models/deepseek_v4/sparse_mla.py

@@ -257,13 +257,6 @@ class DeepseekV4SparseMLAMetadataBuilder(
         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],
@@ -276,7 +269,7 @@ class DeepseekV4SparseMLAMetadataBuilder(
             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,
         )
 
         result: dict[str, torch.Tensor | None] = {}
@@ -322,30 +315,25 @@ 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 pre-allocated buffers for CUDA graph address stability.
+    Returns slices of the buffers.
     """
     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,