buffer size insuffient Dspark sd for FlashInfer MNNVL allreduce (#50932)
Khushali Desai committed
Aug 27, 2026 at 08:01 UTC
c7b0467c2524f85d619a8143bc7cb0d94f968396
2 files changed
+32
-2
vllm/distributed/device_communicators/flashinfer_all_reduce.py
+21
-2
@@ -386,7 +386,6 @@ class FlashInferAllReduce:
386
self.world_size,
387
)
388
return
389
-
389
backend, _ = _resolve_fi_ar_backend()
390
tuned_max_size = _get_tuned_standalone_max_size(
391
self.world_size,
@@ -450,7 +449,27 @@ class FlashInferAllReduce:
449
if num_tokens > self.max_num_tokens:
450
return False
451
453
- return self._ensure_workspace(hidden_dim, input_tensor.dtype)
452
+ if not self._ensure_workspace(hidden_dim, input_tensor.dtype):
453
+ return False
454
+
455
+ workspace = get_fi_ar_workspace(
456
+ world_size=self.world_size,
457
+ rank=self.rank,
458
+ max_token_num=self.max_num_tokens,
459
+ hidden_dim=hidden_dim,
460
+ dtype=input_tensor.dtype,
461
+ group=self.group,
462
+ )
463
+ assert workspace is not None
464
+ # The token bound above uses the full allocation budget, but mnnvl's
465
+ # Lamport buffers rotate through three slots, so only ~1/3 is usable per
466
+ # call. Asking the workspace directly to reject sizes the kernel can't fit.
467
+ return workspace.is_buffer_size_sufficient(
468
+ tp_size=self.world_size,
469
+ num_tokens=num_tokens,
470
+ hidden_dim=hidden_dim,
471
+ dtype=input_tensor.dtype,
472
+ )
473
474
def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor:
475
num_tokens, hidden_dim = input_tensor.shape
vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
+11
@@ -97,6 +97,17 @@ def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool
97
)
98
if workspace is None:
99
return False, 0
100
+ # The token-count bound above uses the whole workspace budget, but a backend
101
+ # may use only a fraction of it per call (mnnvl rotates through three Lamport
102
+ # buffers). Ask the workspace so we don't admit tensors that the kernel will
103
+ # reject with "The buffer size in the given workspace is insufficient".
104
+ if not workspace.is_buffer_size_sufficient(
105
+ tp_size=tp_size,
106
+ num_tokens=num_tokens,
107
+ hidden_dim=hidden_size,
108
+ dtype=hidden_states.dtype,
109
+ ):
110
+ return False, 0
111
return True, max_token_num
112
113