From edd4c8176cfd98ece8a29beda574378c42971967 Mon Sep 17 00:00:00 2001 From: Toby Mao Date: Sat, 15 Aug 2026 18:39:23 -0700 Subject: [PATCH] [Bugfix][DSv4] Revert adaptive C128A metadata packing (#51318) Signed-off-by: tobymao Signed-off-by: Yongye Zhu Co-authored-by: Claude Co-authored-by: Roger Wang Co-authored-by: Yongye Zhu Co-authored-by: OpenAI Codex --- .../kernels/attention/test_flashmla_sparse.py | 37 ------------------- vllm/models/deepseek_v4/sparse_mla.py | 26 ++++--------- 2 files changed, 7 insertions(+), 56 deletions(-) diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index 8082a359cc..eb953faa44 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/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 diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index db8ab96e90..0475e4c96d 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/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,