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 3, 16:30 UTC
[Bugfix] Fix TRTLLM ragged MLA prefill workspace warmup (#42112)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Mohammad Miadh Angkad committed
May 14, 2026 at 21:48 UTC
2317682f95114c4088b9223473a8c50836deb994
2 files changed
+9
-15
vllm/v1/attention/backends/mla/prefill/flashinfer.py
+6
-7
@@ -77,6 +77,9 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
77
self._prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None
78
self._prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []
79
self._global_hyperparameters: PerLayerParameters | None = None
80
+ (self._workspace_buffer,) = current_workspace_manager().get_simultaneous(
81
+ ((envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), torch.uint8),
82
+ )
83
84
def _ensure_chunks(
85
self,
@@ -123,21 +126,17 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
126
global_hyperparameters = self._resolve_global_hyperparameters()
127
qo_indptr = prefill_metadata.query_start_loc
128
has_context = prefill_metadata.chunked_context is not None
126
- (workspace_buffer,) = current_workspace_manager().get_simultaneous(
127
- ((envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), torch.uint8),
128
- )
129
-
129
if self._prefill_main is None:
130
self._prefill_main = BatchPrefillWithRaggedKVCacheWrapper(
132
- workspace_buffer, "NHD", backend="cutlass"
131
+ self._workspace_buffer, "NHD", backend="cutlass"
132
)
134
- self._ensure_chunks(_DEFAULT_NUM_CHUNKS, workspace_buffer)
133
+ self._ensure_chunks(_DEFAULT_NUM_CHUNKS, self._workspace_buffer)
134
135
if has_context:
136
chunked_context = prefill_metadata.chunked_context
137
assert chunked_context is not None
138
num_chunks = chunked_context.cu_seq_lens.shape[0]
140
- self._ensure_chunks(num_chunks, workspace_buffer)
139
+ self._ensure_chunks(num_chunks, self._workspace_buffer)
140
141
num_qo_heads = self.num_heads
142
num_kv_heads = num_qo_heads
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
+3
-8
@@ -61,15 +61,12 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
61
v_head_dim=v_head_dim,
62
vllm_config=vllm_config,
63
)
64
-
65
- def _get_workspace_buffer(self) -> torch.Tensor:
66
- (workspace_buffer,) = current_workspace_manager().get_simultaneous(
64
+ (self._workspace_buffer,) = current_workspace_manager().get_simultaneous(
65
(
66
(envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,),
67
torch.uint8,
68
),
69
)
72
- return workspace_buffer
70
71
def prepare_metadata(
72
self,
@@ -89,7 +86,6 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
86
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
87
from flashinfer.prefill import trtllm_ragged_attention_deepseek
88
92
- workspace_buffer = self._get_workspace_buffer()
89
out = torch.empty(
90
q.shape[0],
91
q.shape[1],
@@ -102,7 +98,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
98
query=q,
99
key=k,
100
value=v,
105
- workspace_buffer=workspace_buffer,
101
+ workspace_buffer=self._workspace_buffer,
102
seq_lens=self._query_seq_lens,
103
max_q_len=self._prefill_metadata.max_query_len,
104
max_kv_len=self._prefill_metadata.max_query_len,
@@ -135,7 +131,6 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
131
132
assert self._prefill_metadata.chunked_context is not None
133
assert self._prefill_metadata.chunked_context.seq_lens[chunk_idx] is not None
138
- workspace_buffer = self._get_workspace_buffer()
134
135
out = torch.empty(
136
q.shape[0],
@@ -149,7 +144,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
144
query=q,
145
key=k,
146
value=v,
152
- workspace_buffer=workspace_buffer,
147
+ workspace_buffer=self._workspace_buffer,
148
seq_lens=self._prefill_metadata.chunked_context.seq_lens[chunk_idx],
149
max_q_len=self._prefill_metadata.max_query_len,
150
max_kv_len=self._prefill_metadata.chunked_context.max_seq_lens[chunk_idx],