[Bugfix] Use larger workspace size for Flashinfer MLA LSE (#47074)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Jun 30, 2026 at 01:11 UTC
b153dd3f2811d76dadac9d563180a6dffd0b10fb
1 file changed
+18
-9
vllm/v1/attention/backends/mla/flashinfer_mla.py
+18
-9
@@ -28,6 +28,22 @@ from vllm.v1.attention.backends.utils import KVCacheLayoutType
28
logger = init_logger(__name__)
29
30
FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
31
+FLASHINFER_MLA_LSE_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
32
+
33
+_fi_workspace: torch.Tensor | None = None
34
+
35
+
36
+def _get_workspace_buffer(return_lse: bool) -> torch.Tensor:
37
+ global _fi_workspace
38
+
39
+ buffer_size = (
40
+ FLASHINFER_MLA_LSE_WORKSPACE_BUFFER_SIZE
41
+ if return_lse
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")
46
+ return _fi_workspace
47
48
49
class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
@@ -105,13 +121,6 @@ class FlashInferMLABackend(MLACommonBackend):
121
return "HND"
122
123
108
-g_fi_workspace = torch.zeros(
109
- FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE,
110
- dtype=torch.uint8,
111
- device="cuda",
112
-)
113
-
114
-
124
class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
125
can_return_lse_for_decode: bool = True
126
# trtllm-gen MLA decode emits LSE in log2 (per flashinfer's own
@@ -166,7 +175,6 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
175
"FlashInferMLAImpl"
176
)
177
169
- self._workspace_buffer = g_fi_workspace
178
self.bmm1_scale: float | None = None
179
self.bmm2_scale: float | None = None
180
@@ -206,10 +214,11 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
214
self.bmm2_scale *= layer._k_scale_float
215
216
return_lse = self.need_to_return_lse_for_decode
217
+ workspace_buffer = _get_workspace_buffer(return_lse)
218
kernel_out = trtllm_batch_decode_with_kv_cache_mla(
219
query=q,
220
kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
212
- workspace_buffer=self._workspace_buffer,
221
+ workspace_buffer=workspace_buffer,
222
qk_nope_head_dim=self.qk_nope_head_dim,
223
kv_lora_rank=self.kv_lora_rank,
224
qk_rope_head_dim=self.qk_rope_head_dim,