A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 01:15 UTC
[Bug] Fix `'CutlassMLAImpl' object has no attribute '_workspace_buffer'` (#31173)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Dec 22, 2025 at 17:24 UTC
5312a7284e58e47bd10f7acbab52f703714027c2
1 file changed
+7
-3
vllm/v1/attention/backends/mla/common.py
+7
-3
@@ -355,6 +355,7 @@ class MLACommonPrefillMetadata:
355
max_query_len: int
356
chunked_context: ChunkedContextMetadata | None = None
357
query_seq_lens: torch.Tensor | None = None
358
+ workspace_buffer: torch.Tensor | None = None
359
q_data_type: torch.dtype | None = None
360
361
@@ -986,6 +987,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
987
prefill_metadata.query_seq_lens = (
988
prefill_query_start_loc[1:] - prefill_query_start_loc[:-1]
989
)
990
+ prefill_metadata.workspace_buffer = self._workspace_buffer
991
992
decode_metadata = None
993
if num_decodes > 0:
@@ -1567,6 +1569,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1569
from flashinfer.prefill import trtllm_ragged_attention_deepseek
1570
1571
assert prefill.query_seq_lens is not None
1572
+ assert prefill.workspace_buffer is not None
1573
1574
if fp8_attention:
1575
logger.debug_once("Running TRT-LLM ragged prefill in FP8")
@@ -1579,7 +1582,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1582
query=q,
1583
key=k,
1584
value=v,
1582
- workspace_buffer=self._workspace_buffer,
1585
+ workspace_buffer=prefill.workspace_buffer,
1586
seq_lens=prefill.query_seq_lens,
1587
max_q_len=prefill.max_query_len,
1588
max_kv_len=prefill.max_query_len,
@@ -1615,6 +1618,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1618
1619
assert prefill.chunked_context is not None
1620
assert prefill.chunked_context.seq_lens[chunk_idx] is not None
1621
+ assert prefill.workspace_buffer is not None
1622
1623
out = torch.zeros(
1624
q.shape[0],
@@ -1623,7 +1627,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1627
device=q.device,
1628
dtype=q.dtype,
1629
)
1626
- self._workspace_buffer.fill_(0)
1630
+ prefill.workspace_buffer.fill_(0)
1631
1632
if fp8_attention:
1633
logger.debug_once("Running TRT-LLM ragged prefill context chunk in FP8")
@@ -1636,7 +1640,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1640
query=q,
1641
key=k,
1642
value=v,
1639
- workspace_buffer=self._workspace_buffer,
1643
+ workspace_buffer=prefill.workspace_buffer,
1644
seq_lens=prefill.chunked_context.seq_lens[chunk_idx],
1645
max_q_len=prefill.max_query_len,
1646
max_kv_len=prefill.chunked_context.max_seq_lens[chunk_idx],