[Bugfix] Release stale FlashMLA workspace views after growth (#56902)
Signed-off-by: Maroon Ayoub <mayoub@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Maroon Ayoub committed
Sep 17, 2026 at 23:23 UTC
9612f77077e09acbae9cc1d1b2ddf14a23e91597
1 file changed
+26
-19
vllm/v1/attention/backends/mla/flashmla_sparse.py
+26
-19
@@ -720,7 +720,9 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
720
)
721
722
self.pcp_dcp_kv_gather = False
723
- self.gathered_kv_workspace: torch.Tensor | None = None
723
+ self.workspace_specs: list[tuple[tuple[int, ...], torch.dtype]] = [
724
+ (q_concat_shape, torch.bfloat16)
725
+ ]
726
if kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS:
727
# Reserve workspace during initialization
728
assert vllm_config is not None and vllm_config.model_config is not None
@@ -738,20 +740,14 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
740
# shards into a workspace of the full prefill size.
741
shard_rows //= parallel_config.decode_context_parallel_size
742
self.prefill_workspace_shape = (shard_rows, head_size)
741
- specs = [
742
- (q_concat_shape, torch.bfloat16),
743
- (self.prefill_workspace_shape, torch.bfloat16),
744
- ]
743
+ self.workspace_specs.append((self.prefill_workspace_shape, torch.bfloat16))
744
if self.pcp_dcp_kv_gather:
746
- specs.append(((prefill_workspace_size, head_size), torch.bfloat16))
747
- buffers = current_workspace_manager().get_simultaneous(*specs)
748
- self.q_concat_buffer, self.prefill_bf16_workspace = buffers[:2]
749
- if self.pcp_dcp_kv_gather:
750
- self.gathered_kv_workspace = buffers[2]
751
- else:
752
- (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous(
753
- (q_concat_shape, torch.bfloat16),
754
- )
745
+ self.workspace_specs.append(
746
+ ((prefill_workspace_size, head_size), torch.bfloat16)
747
+ )
748
+ # Reserve capacity without retaining views that prevent old storage
749
+ # from being released when another layer grows the shared workspace.
750
+ current_workspace_manager().get_simultaneous(*self.workspace_specs)
751
752
def _forward_bf16_kv(
753
self,
@@ -849,8 +845,10 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
845
"""All-gather this rank's upconverted KV shard so the chunk's rows attend
846
the whole context, and map their top-k onto the rank-major result."""
847
shard_rows = int(chunk.chunk_tot_seqlen)
852
- assert self.gathered_kv_workspace is not None
853
- gathered_kv = self.gathered_kv_workspace[: self.dcp_world_size * shard_rows]
848
+ _, _, gathered_kv_workspace = current_workspace_manager().get_simultaneous(
849
+ *self.workspace_specs
850
+ )
851
+ gathered_kv = gathered_kv_workspace[: self.dcp_world_size * shard_rows]
852
dist.all_gather_into_tensor(
853
gathered_kv, shard, group=get_dcp_group().device_group
854
)
@@ -902,10 +900,13 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
900
assert fp8_metadata.prefill is not None
901
first_chunk = fp8_metadata.prefill.chunks[0]
902
assert isinstance(index_group, HiSparseMLAIndexGroup)
903
+ _, prefill_bf16_workspace, *_ = (
904
+ current_workspace_manager().get_simultaneous(*self.workspace_specs)
905
+ )
906
prefill_ready = index_group.gather_fp8_prefill(
907
self.index_group_index,
908
kv_c_and_k_pe_cache,
908
- self.prefill_bf16_workspace[: first_chunk.chunk_tot_seqlen],
909
+ prefill_bf16_workspace[: first_chunk.chunk_tot_seqlen],
910
first_chunk.block_table,
911
first_chunk.workspace_starts,
912
len(first_chunk.block_table),
@@ -1022,8 +1023,11 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1023
)
1024
1025
assert fp8_metadata.prefill is not None
1026
+ _, prefill_bf16_workspace, *_ = (
1027
+ current_workspace_manager().get_simultaneous(*self.workspace_specs)
1028
+ )
1029
for chunk_index, chunk in enumerate(fp8_metadata.prefill.chunks):
1026
- chunk_workspace = self.prefill_bf16_workspace[: chunk.chunk_tot_seqlen]
1030
+ chunk_workspace = prefill_bf16_workspace[: chunk.chunk_tot_seqlen]
1031
if uses_host_cache and chunk_index > 0:
1032
assert isinstance(index_group, HiSparseMLAIndexGroup)
1033
prefill_ready = index_group.gather_fp8_prefill(
@@ -1302,7 +1306,10 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1306
actual_num_heads = self.num_heads
1307
if isinstance(q, tuple):
1308
ql_nope, q_pe = q
1305
- q = self.q_concat_buffer[: ql_nope.shape[0]]
1309
+ q_concat_buffer, *_ = current_workspace_manager().get_simultaneous(
1310
+ *self.workspace_specs
1311
+ )
1312
+ q = q_concat_buffer[: ql_nope.shape[0]]
1313
ops.concat_mla_q(ql_nope, q_pe, q)
1314
else:
1315
actual_num_heads = q.shape[1]