[Bugfix][MLA] Reserve sparse prefill buffers before KV cache sizing (#57575)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Lucas Wilkinson committed
Sep 18, 2026 at 17:22 UTC
62af3df7347cf890b4b925a18ab6a2afd00c70b9
3 files changed
+59
-12
tests/v1/attention/test_sparse_mla_backends.py
+15
@@ -421,7 +421,18 @@ def test_sparse_backend_decode_correctness(
421
workspace_init,
422
q_scale: float,
423
k_scale: float,
424
+ monkeypatch,
425
):
426
+ if (
427
+ batch_name == "large_q_pure_prefill"
428
+ and backend_cls == FlashMLASparseBackend
429
+ and kv_cache_dtype == "fp8_ds_mla"
430
+ and tensor_parallel_size == 4
431
+ ):
432
+ monkeypatch.setattr(
433
+ "vllm.v1.attention.backends.mla.flashmla_sparse.split_prefill_chunks",
434
+ lambda rows, capacity: [(i, i + 1) for i in range(len(rows))],
435
+ )
436
if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:
437
pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")
438
@@ -774,6 +785,10 @@ def test_sparse_backend_decode_correctness(
785
)
786
787
with torch.inference_mode():
788
+ if backend_cls == FlashMLASparseBackend and kv_cache_dtype == "fp8_ds_mla":
789
+ from vllm.v1.worker.workspace import current_workspace_manager
790
+
791
+ current_workspace_manager().lock()
792
backend_output = mock_layer.forward_impl(
793
query_vllm,
794
kv_c_vllm,
vllm/models/hy_v4/nvidia/flashmla_sparse.py
+11
-1
@@ -34,6 +34,7 @@ from vllm.v1.attention.ops.flashmla import (
34
flash_mla_sparse_fwd,
35
flash_mla_with_kvcache,
36
)
37
+from vllm.v1.worker.workspace import current_workspace_manager
38
39
if TYPE_CHECKING:
40
from vllm.model_executor.models.deepseek_v2 import Indexer
@@ -220,6 +221,7 @@ class HYV4FlashMLASparseImpl(FlashMLASparseImpl):
221
topk_indices: torch.Tensor,
222
topk_length: torch.Tensor | None = None,
223
actual_num_heads: int | None = None,
224
+ out: torch.Tensor | None = None,
225
) -> tuple[torch.Tensor, torch.Tensor]:
226
num_tokens = q.shape[0]
227
kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
@@ -250,7 +252,14 @@ class HYV4FlashMLASparseImpl(FlashMLASparseImpl):
252
# Zero (not new_empty) the padded lanes: topk_indices is shared by
253
# all heads, so the kernel reduces across the head group and NaNs
254
# from uninitialized memory would leak into the real heads.
253
- q_padded = q.new_zeros((q.shape[0], padded_num_heads, q.shape[2]))
255
+ if out is None:
256
+ q_padded = q.new_zeros((q.shape[0], padded_num_heads, q.shape[2]))
257
+ else:
258
+ *_, q_padded, _, _ = current_workspace_manager().get_simultaneous(
259
+ *self.workspace_specs
260
+ )
261
+ q_padded = q_padded[:num_tokens]
262
+ q_padded[:, actual_num_heads:].zero_()
263
q_padded[:, :actual_num_heads, :] = q
264
q = q_padded
265
@@ -262,6 +271,7 @@ class HYV4FlashMLASparseImpl(FlashMLASparseImpl):
271
self.softmax_scale,
272
attn_sink=attn_sink,
273
topk_length=topk_length,
274
+ out=out,
275
)
276
277
output = output[:, :actual_num_heads, :]
vllm/v1/attention/backends/mla/flashmla_sparse.py
+33
-11
@@ -19,6 +19,7 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import (
19
)
20
from vllm.platforms import current_platform
21
from vllm.platforms.interface import DeviceCapability
22
+from vllm.utils.math_utils import round_up
23
from vllm.utils.platform_utils import num_compute_units
24
from vllm.utils.torch_utils import (
25
async_tensor_h2d,
@@ -745,6 +746,20 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
746
self.workspace_specs.append(
747
((prefill_workspace_size, head_size), torch.bfloat16)
748
)
749
+ prefill_query_heads = num_heads
750
+ if self.pcp_dcp_kv_gather and self.dcp_world_size > self.pcp_world_size:
751
+ prefill_query_heads *= parallel_config.tensor_parallel_size
752
+ padded_prefill_query_heads = round_up(
753
+ prefill_query_heads, self.prefill_padding
754
+ )
755
+ self.workspace_specs.extend(
756
+ (shape, torch.bfloat16)
757
+ for shape in (
758
+ (max_tokens, padded_prefill_query_heads, head_size),
759
+ (max_tokens, padded_prefill_query_heads, self.kv_lora_rank),
760
+ (max_tokens, prefill_query_heads, self.kv_lora_rank),
761
+ )
762
+ )
763
# Reserve capacity without retaining views that prevent old storage
764
# from being released when another layer grows the shared workspace.
765
current_workspace_manager().get_simultaneous(*self.workspace_specs)
@@ -845,7 +860,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
860
"""All-gather this rank's upconverted KV shard so the chunk's rows attend
861
the whole context, and map their top-k onto the rank-major result."""
862
shard_rows = int(chunk.chunk_tot_seqlen)
848
- _, _, gathered_kv_workspace = current_workspace_manager().get_simultaneous(
863
+ _, _, gathered_kv_workspace, *_ = current_workspace_manager().get_simultaneous(
864
*self.workspace_specs
865
)
866
gathered_kv = gathered_kv_workspace[: self.dcp_world_size * shard_rows]
@@ -1007,12 +1022,10 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1022
q, decode_topk if decode_topk is not None else topk_indices
1023
)
1024
else:
1010
- # Mixed or pure prefill: allocate output tensor
1011
- attn_out = q.new_empty(
1012
- (num_mqa_tokens, self.num_heads, self.kv_lora_rank),
1013
- dtype=q.dtype,
1014
- device=q.device,
1025
+ _, prefill_bf16_workspace, *_, padded_out, compact_out = (
1026
+ current_workspace_manager().get_simultaneous(*self.workspace_specs)
1027
)
1028
+ attn_out = compact_out[:num_mqa_tokens]
1029
1030
if num_decode_tokens > 0:
1031
attn_out[:num_decode_tokens] = _fp8_decode(
@@ -1023,9 +1036,6 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1036
)
1037
1038
assert fp8_metadata.prefill is not None
1026
- _, prefill_bf16_workspace, *_ = (
1027
- current_workspace_manager().get_simultaneous(*self.workspace_specs)
1028
- )
1039
for chunk_index, chunk in enumerate(fp8_metadata.prefill.chunks):
1040
chunk_workspace = prefill_bf16_workspace[: chunk.chunk_tot_seqlen]
1041
if uses_host_cache and chunk_index > 0:
@@ -1076,12 +1086,15 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1086
assert topk_length is not None
1087
chunk_topk_length = topk_length[chunk.tokens_slice]
1088
1079
- attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel(
1089
+ chunk_out, chunk_lse = self._bf16_flash_mla_kernel(
1090
chunk_q,
1091
chunk_workspace,
1092
chunk_topk_indices_workspace,
1093
chunk_topk_length,
1094
+ out=padded_out[chunk.tokens_slice],
1095
)
1096
+ attn_out[chunk.tokens_slice].copy_(chunk_out)
1097
+ del chunk_out, chunk_lse
1098
1099
if self.pcp_dcp_kv_gather and lse is None:
1100
# No decode rows: the DCP merge still expects an LSE, of no rows.
@@ -1254,6 +1267,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1267
topk_indices: torch.Tensor,
1268
topk_length: torch.Tensor | None = None,
1269
actual_num_heads: int | None = None,
1270
+ out: torch.Tensor | None = None,
1271
) -> tuple[torch.Tensor, torch.Tensor]:
1272
num_tokens = q.shape[0]
1273
kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
@@ -1275,7 +1289,14 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1289
f"Padding num_heads from {actual_num_heads} to "
1290
f"{padded_num_heads} for BF16 sparse prefill kernel"
1291
)
1278
- q_padded = q.new_empty((q.shape[0], padded_num_heads, q.shape[2]))
1292
+ if out is None:
1293
+ q_padded = q.new_empty((q.shape[0], padded_num_heads, q.shape[2]))
1294
+ else:
1295
+ *_, q_padded, _, _ = current_workspace_manager().get_simultaneous(
1296
+ *self.workspace_specs
1297
+ )
1298
+ q_padded = q_padded[:num_tokens]
1299
+ q_padded[:, actual_num_heads:].zero_()
1300
q_padded[:, :actual_num_heads, :] = q
1301
q = q_padded
1302
@@ -1286,6 +1307,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1307
topk_indices,
1308
self.softmax_scale,
1309
topk_length=topk_length,
1310
+ out=out,
1311
)
1312
1313
output = output[:, :actual_num_heads, :]