[Bugfix] Use int8 workspace for FlashInfer MLA decode (#48046)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Nick Hill committed
Jul 9, 2026 at 00:39 UTC
95d6d6f4bba8234088e62124ff20a482acd98714
3 files changed
+93
-47
tests/kernels/attention/test_flashinfer_mla_decode.py
+87
-45
@@ -17,6 +17,43 @@ if not current_platform.has_device_capability(100):
17
else:
18
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
19
20
+# Deepseek R1 MLA config.
21
+NUM_HEADS = 128
22
+KV_LORA_RANK = 512
23
+QK_NOPE_HEAD_DIM = 128
24
+QK_ROPE_HEAD_DIM = 64
25
+QK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM
26
+SCALE = (QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM) ** -0.5
27
+
28
+
29
+def _make_decode_inputs(bs: int, block_size: int, dtype: torch.dtype):
30
+ """Build valid trtllm MLA decode inputs on the current CUDA device."""
31
+ max_seq_len_cap = 1024
32
+ seq_lens = [torch.randint(2, max_seq_len_cap, (1,)).item() for _ in range(bs)]
33
+ seq_lens[-1] = max_seq_len_cap
34
+ max_seq_len = max(seq_lens)
35
+ seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
36
+
37
+ # Generate block tables with random but unique block IDs
38
+ # From https://github.com/flashinfer-ai/flashinfer/pull/1222
39
+ blocks_per_seq = (seq_lens_tensor + block_size - 1) // block_size
40
+ max_num_blocks_per_seq = max(blocks_per_seq.max().item(), 4)
41
+ total_blocks_needed = int(sum(blocks_per_seq))
42
+ all_block_ids = torch.randperm(total_blocks_needed)
43
+
44
+ block_tables = torch.zeros((bs, max_num_blocks_per_seq), dtype=torch.int32)
45
+ block_id = 0
46
+ for i in range(bs):
47
+ num_blocks_needed = blocks_per_seq[i]
48
+ block_tables[i, :num_blocks_needed] = all_block_ids[
49
+ block_id : block_id + num_blocks_needed
50
+ ]
51
+ block_id += num_blocks_needed
52
+
53
+ kv_cache = torch.randn(block_tables.numel(), block_size, QK_HEAD_DIM).to(dtype)
54
+ q = torch.randn(bs, NUM_HEADS, QK_HEAD_DIM).to(dtype)
55
+ return q, kv_cache, block_tables, seq_lens_tensor, max_seq_len
56
+
57
58
def ref_mla(
59
out: Tensor, # (bs, num_heads, v_head_dim)
@@ -49,49 +86,12 @@ def test_flashinfer_mla_decode(dtype: torch.dtype, bs: int, block_size: int):
86
torch.set_default_device("cuda")
87
torch.manual_seed(42)
88
52
- # Deepseek R1 config
53
- num_heads = 128
54
- kv_lora_rank = 512
55
- qk_nope_head_dim = 128
56
- qk_rope_head_dim = 64
57
- qk_head_dim = kv_lora_rank + qk_rope_head_dim
58
- scale = (qk_nope_head_dim + qk_rope_head_dim) ** -0.5
59
-
60
- MAX_SEQ_LEN = 1024
61
-
62
- seq_lens = [torch.randint(2, MAX_SEQ_LEN, (1,)).item() for _ in range(bs)]
63
- seq_lens[-1] = MAX_SEQ_LEN
64
- max_seq_len = max(seq_lens)
65
- seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
66
-
67
- # Generate block tables with random but unique block IDs
68
- # From https://github.com/flashinfer-ai/flashinfer/pull/1222
69
- blocks_per_seq = (seq_lens_tensor + block_size - 1) // block_size
70
- max_num_blocks_per_seq = max(blocks_per_seq.max().item(), 4)
71
- total_blocks_needed = sum(blocks_per_seq)
72
- # Get random unique IDs for all blocks
73
- all_block_ids = torch.randperm(total_blocks_needed)
74
-
75
- block_id = 0
76
- block_tables = torch.zeros(
77
- (bs, max_num_blocks_per_seq),
78
- dtype=torch.int32,
89
+ q, kv_cache, block_tables, seq_lens_tensor, max_seq_len = _make_decode_inputs(
90
+ bs, block_size, dtype
91
)
92
81
- # Populate block tables and track block assignments
82
- block_id = 0
83
- for i in range(bs):
84
- num_blocks_needed = blocks_per_seq[i]
85
- block_tables[i, :num_blocks_needed] = all_block_ids[
86
- block_id : block_id + num_blocks_needed
87
- ]
88
- block_id += num_blocks_needed
89
-
90
- kv_cache = torch.randn(block_tables.numel(), block_size, qk_head_dim).to(dtype)
91
- q = torch.randn(bs, num_heads, qk_head_dim).to(dtype)
92
-
93
- out_ref = q.new_zeros(bs, num_heads, kv_lora_rank)
94
- ref_mla(out_ref, q, kv_cache, scale, block_tables, seq_lens_tensor)
93
+ out_ref = q.new_zeros(bs, NUM_HEADS, KV_LORA_RANK)
94
+ ref_mla(out_ref, q, kv_cache, SCALE, block_tables, seq_lens_tensor)
95
96
workspace_buffer = torch.zeros(
97
FLASHINFER_WORKSPACE_BUFFER_SIZE,
@@ -107,13 +107,55 @@ def test_flashinfer_mla_decode(dtype: torch.dtype, bs: int, block_size: int):
107
query=q,
108
kv_cache=kv_cache.unsqueeze(1),
109
workspace_buffer=workspace_buffer,
110
- qk_nope_head_dim=qk_nope_head_dim,
111
- kv_lora_rank=kv_lora_rank,
112
- qk_rope_head_dim=qk_rope_head_dim,
110
+ qk_nope_head_dim=QK_NOPE_HEAD_DIM,
111
+ kv_lora_rank=KV_LORA_RANK,
112
+ qk_rope_head_dim=QK_ROPE_HEAD_DIM,
113
block_tables=block_tables,
114
seq_lens=seq_lens_tensor,
115
max_seq_len=max_seq_len,
116
- bmm1_scale=scale,
116
+ bmm1_scale=SCALE,
117
)
118
out_ans = out_ans.squeeze(1)
119
torch.testing.assert_close(out_ans, out_ref, atol=1e-2, rtol=1e-2)
120
+
121
+
122
+def test_flashinfer_mla_decode_workspace_supports_autotune():
123
+ """vLLM's FlashInfer MLA decode workspace must be int8 for autotuning.
124
+
125
+ Model Runner V2's warmup autotunes ``trtllm_batch_decode_mla``, which makes
126
+ the FlashInfer autotuner enumerate the CuteDSL tactic. That tactic asserts
127
+ ``workspace_buffer.dtype == torch.int8``; the trtllm-gen path (used for
128
+ normal, non-autotuned inference) instead views the buffer as uint8, so a
129
+ uint8 workspace only fails once the autotuner tries CuteDSL. That regressed
130
+ every DeepSeek MLA test on Blackwell under V2 with
131
+ ``workspace_buffer must be torch.int8`` (vllm-project/vllm#46646).
132
+ """
133
+ from flashinfer.autotuner import autotune
134
+
135
+ from vllm.v1.attention.backends.mla.flashinfer_mla import _get_workspace_buffer
136
+
137
+ torch.set_default_device("cuda")
138
+ torch.manual_seed(0)
139
+
140
+ workspace_buffer = _get_workspace_buffer(return_lse=False)
141
+ assert workspace_buffer.dtype == torch.int8
142
+
143
+ q, kv_cache, block_tables, seq_lens_tensor, max_seq_len = _make_decode_inputs(
144
+ bs=1, block_size=64, dtype=torch.bfloat16
145
+ )
146
+
147
+ # Under the autotuner the CuteDSL tactic is instantiated with our workspace;
148
+ # a uint8 buffer raises AssertionError here, an int8 buffer succeeds.
149
+ with torch.inference_mode(), autotune(True):
150
+ trtllm_batch_decode_with_kv_cache_mla(
151
+ query=q.unsqueeze(1),
152
+ kv_cache=kv_cache.unsqueeze(1),
153
+ workspace_buffer=workspace_buffer,
154
+ qk_nope_head_dim=QK_NOPE_HEAD_DIM,
155
+ kv_lora_rank=KV_LORA_RANK,
156
+ qk_rope_head_dim=QK_ROPE_HEAD_DIM,
157
+ block_tables=block_tables,
158
+ seq_lens=seq_lens_tensor,
159
+ max_seq_len=max_seq_len,
160
+ bmm1_scale=SCALE,
161
+ )
vllm/v1/attention/backends/mla/flashinfer_mla.py
+3
-1
@@ -42,7 +42,9 @@ def _get_workspace_buffer(return_lse: bool) -> torch.Tensor:
42
else FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE
43
)
44
if _fi_workspace is None or _fi_workspace.numel() < buffer_size:
45
- _fi_workspace = torch.zeros(buffer_size, dtype=torch.uint8, device="cuda")
45
+ # FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace;
46
+ # the trtllm-gen path views it as uint8, so int8 is safe for all backends.
47
+ _fi_workspace = torch.zeros(buffer_size, dtype=torch.int8, device="cuda")
48
return _fi_workspace
49
50
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+3
-1
@@ -355,9 +355,11 @@ _fi_sparse_workspace: torch.Tensor | None = None
355
def _get_workspace_buffer(device: torch.device) -> torch.Tensor:
356
global _fi_sparse_workspace
357
if _fi_sparse_workspace is None:
358
+ # FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace;
359
+ # the trtllm-gen path views it as uint8, so int8 is safe for all backends.
360
_fi_sparse_workspace = torch.zeros(
361
envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
360
- dtype=torch.uint8,
362
+ dtype=torch.int8,
363
device=device,
364
)
365
return _fi_sparse_workspace