[Bugfix][SM120] DSv4: pass contiguous C128A decode topk indices on SM120 (#53574)
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Gabriel Wu committed
Aug 31, 2026 at 17:44 UTC
699e180df48d78b5275d528124641c67b4b1664c
2 files changed
+31
-7
tests/kernels/attention/test_flashmla_sparse.py
+22
-5
@@ -4,8 +4,19 @@ import pytest
4
import torch
5
6
7
-def test_deepseek_v4_c128a_adaptive_width_has_capture_stable_stride():
8
- from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata
7
+@pytest.mark.parametrize("sm120", [False, True])
8
+def test_deepseek_v4_c128a_adaptive_width_has_capture_stable_stride(
9
+ monkeypatch: pytest.MonkeyPatch,
10
+ sm120: bool,
11
+):
12
+ from vllm.models.deepseek_v4 import sparse_mla
13
+ from vllm.platforms.interface import DeviceCapability
14
+
15
+ monkeypatch.setattr(
16
+ sparse_mla.current_platform,
17
+ "get_device_capability",
18
+ lambda: DeviceCapability(12, 0) if sm120 else DeviceCapability(10, 0),
19
+ )
20
21
device = torch.device("cuda")
22
capacity_width = 512
@@ -27,12 +38,18 @@ def test_deepseek_v4_c128a_adaptive_width_has_capture_stable_stride():
38
decode_lens_buffer=torch.empty(2, dtype=torch.int32, device=device),
39
prefill_buffer=prefill_buffer,
40
)
30
- captured_decode, _, captured_prefill = build_c128a_topk_metadata(
41
+ captured_decode, _, captured_prefill = sparse_mla.build_c128a_topk_metadata(
42
max_compressed_tokens=256,
43
**kwargs,
44
)
34
- assert captured_decode.shape == captured_prefill.shape == (2, 256)
45
+ # SM120 keeps the decode view contiguous across the full buffer width;
46
+ # other backends get the active-width slice. The prefill view is always
47
+ # narrowed.
48
+ expected_decode_width = capacity_width if sm120 else 256
49
+ assert captured_decode.shape == (2, expected_decode_width)
50
+ assert captured_prefill.shape == (2, 256)
51
assert captured_decode.stride(0) == captured_prefill.stride(0) == capacity_width
52
+ assert captured_decode.is_contiguous() == sm120
53
54
captured_rows = torch.empty((4, 4), dtype=torch.int32, device=device)
55
captured_rows[:2].copy_(captured_decode[:, :4])
@@ -45,7 +62,7 @@ def test_deepseek_v4_c128a_adaptive_width_has_capture_stable_stride():
62
63
global_decode_buffer.fill_(-99)
64
prefill_buffer.fill_(-99)
48
- build_c128a_topk_metadata(
65
+ sparse_mla.build_c128a_topk_metadata(
66
max_compressed_tokens=128,
67
**kwargs,
68
)
vllm/models/deepseek_v4/sparse_mla.py
+9
-2
@@ -9,6 +9,7 @@ import torch
9
10
from vllm.config import VllmConfig
11
from vllm.config.cache import CacheDType
12
+from vllm.platforms import current_platform
13
from vllm.platforms.interface import DeviceCapability
14
from vllm.triton_utils import tl, triton
15
from vllm.utils.math_utils import cdiv
@@ -310,7 +311,7 @@ def build_c128a_topk_metadata(
311
Prefill tokens: position → local indices [0, ..., n-1, -1, ...].
312
313
Writes into pre-allocated buffers for CUDA graph address stability.
313
- Returns slices of the buffers.
314
+ Returns views of the buffers.
315
"""
316
num_tokens = positions.shape[0]
317
num_prefill_tokens = num_tokens - num_decode_tokens
@@ -322,7 +323,13 @@ def build_c128a_topk_metadata(
323
)
324
assert global_decode_buffer.stride(-1) == prefill_buffer.stride(-1) == 1
325
325
- global_decode = global_decode_buffer[:num_decode_tokens, :max_compressed_tokens]
326
+ # TODO: support adaptive-width decode on SM120 (needs the FlashInfer
327
+ # SM120 kernel to accept a real row stride for eidx).
328
+ capability = current_platform.get_device_capability()
329
+ if capability is not None and capability.major == 12:
330
+ global_decode = global_decode_buffer[:num_decode_tokens]
331
+ else:
332
+ global_decode = global_decode_buffer[:num_decode_tokens, :max_compressed_tokens]
333
decode_lens = decode_lens_buffer[:num_decode_tokens]
334
prefill_local = prefill_buffer[:num_prefill_tokens, :max_compressed_tokens]
335
assert global_decode.stride(0) == global_decode_buffer.stride(0)