[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>
Toby Mao committed
Aug 15, 2026 at 18:39 UTC
edd4c8176cfd98ece8a29beda574378c42971967
2 files changed
+7
-56
tests/kernels/attention/test_flashmla_sparse.py
-37
@@ -4,43 +4,6 @@ import pytest
4
import torch
5
6
7
-def test_deepseek_v4_c128a_dynamic_topk_packed_buffers():
8
- from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata
9
-
10
- device = torch.device("cuda")
11
- capacity_width = 256
12
- active_width = 128
13
- global_decode_buffer = torch.empty(
14
- (2, capacity_width), dtype=torch.int32, device=device
15
- )
16
- decode_lens_buffer = torch.empty(2, dtype=torch.int32, device=device)
17
- prefill_buffer = torch.empty((2, capacity_width), dtype=torch.int32, device=device)
18
-
19
- global_decode, decode_lens, prefill_local = build_c128a_topk_metadata(
20
- positions=torch.tensor([255, 511], dtype=torch.int64, device=device),
21
- compress_ratio=128,
22
- num_decode_tokens=1,
23
- token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32, device=device),
24
- block_table=torch.tensor([[3]], dtype=torch.int32, device=device),
25
- block_size=capacity_width,
26
- slot_mapping=torch.tensor([0, 1], dtype=torch.int64, device=device),
27
- global_decode_buffer=global_decode_buffer,
28
- decode_lens_buffer=decode_lens_buffer,
29
- prefill_buffer=prefill_buffer,
30
- max_compressed_tokens=active_width,
31
- )
32
-
33
- assert global_decode.shape == (1, active_width)
34
- assert prefill_local.shape == (1, active_width)
35
- assert global_decode.stride() == (active_width, 1)
36
- assert prefill_local.stride() == (active_width, 1)
37
- assert global_decode[0, :2].cpu().tolist() == [768, 769]
38
- assert decode_lens.cpu().tolist() == [2]
39
- assert prefill_local[0, :4].cpu().tolist() == list(range(4))
40
- assert torch.all(global_decode[0, 2:] == -1)
41
- assert torch.all(prefill_local[0, 4:] == -1)
42
-
43
-
7
def test_sparse_flashmla_metadata_smoke():
8
import vllm.v1.attention.ops.flashmla as fm
9
vllm/models/deepseek_v4/sparse_mla.py
+7
-19
@@ -257,13 +257,6 @@ class DeepseekV4SparseMLAMetadataBuilder(
257
assert cm.positions is not None, (
258
"positions is required for C128A metadata build"
259
)
260
- active_topk_width = min(
261
- max(
262
- triton.next_power_of_2(max(cm.max_seq_len // self.compress_ratio, 1)),
263
- _C128A_TOPK_ALIGNMENT,
264
- ),
265
- self.c128a_max_compressed,
266
- )
260
block_size = self.kv_cache_spec.block_size // self.compress_ratio
261
global_decode, decode_lens, prefill_local = build_c128a_topk_metadata(
262
cm.positions[:num_total],
@@ -276,7 +269,7 @@ class DeepseekV4SparseMLAMetadataBuilder(
269
self.c128a_global_decode_buffer,
270
self.c128a_decode_lens_buffer,
271
self.c128a_prefill_buffer,
279
- max_compressed_tokens=active_topk_width,
272
+ max_compressed_tokens=self.c128a_max_compressed,
273
)
274
275
result: dict[str, torch.Tensor | None] = {}
@@ -322,30 +315,25 @@ def build_c128a_topk_metadata(
315
Decode tokens: position → block_table lookup → global slot ids + topk_lens.
316
Prefill tokens: position → local indices [0, ..., n-1, -1, ...].
317
325
- Writes into packed views of pre-allocated buffers for CUDA graph stability.
318
+ Writes into pre-allocated buffers for CUDA graph address stability.
319
+ Returns slices of the buffers.
320
"""
321
num_tokens = positions.shape[0]
322
num_prefill_tokens = num_tokens - num_decode_tokens
323
330
- # view(-1) as 1-d array and then expanded to
331
- # [num_decode_tokens, max_compressed_tokens]
332
- global_decode = global_decode_buffer.view(-1)[
333
- : num_decode_tokens * max_compressed_tokens
334
- ].view(num_decode_tokens, max_compressed_tokens)
324
+ global_decode = global_decode_buffer[:num_decode_tokens]
325
decode_lens = decode_lens_buffer[:num_decode_tokens]
336
- prefill_local = prefill_buffer.view(-1)[
337
- : num_prefill_tokens * max_compressed_tokens
338
- ].view(num_prefill_tokens, max_compressed_tokens)
326
+ prefill_local = prefill_buffer[:num_prefill_tokens]
327
328
if num_tokens == 0:
329
return global_decode, decode_lens, prefill_local
330
331
_build_c128a_topk_metadata_kernel[(num_tokens,)](
332
global_decode_buffer,
345
- max_compressed_tokens,
333
+ global_decode_buffer.stride(0),
334
decode_lens_buffer,
335
prefill_buffer,
348
- max_compressed_tokens,
336
+ prefill_buffer.stride(0),
337
positions,
338
compress_ratio,
339
max_compressed_tokens,