buffer size insuffient Dspark sd for FlashInfer MNNVL allreduce

tests/distributed/test_flashinfer_all_reduce.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Admission-gate tests for :class:`FlashInferAllReduce`.
+
+These exercise the capacity accounting only, so they run on CPU without CUDA or
+flashinfer installed: both the workspace and the input tensor are stubbed.
+"""
+
+import pytest
+import torch
+
+from vllm.distributed.device_communicators import flashinfer_all_reduce
+from vllm.distributed.device_communicators.flashinfer_all_reduce import (
+    FlashInferAllReduce,
+)
+from vllm.utils.math_utils import round_up
+
+# The repro from the bug report: DeepSeek-V4 hidden size, TP8, bf16, and the 2MB
+# per-rank budget FI_ALLREDUCE_FUSION_MAX_SIZE_MB gives TP8 on sm103.
+HIDDEN_DIM = 7168
+WORLD_SIZE = 8
+DTYPE = torch.bfloat16
+WORKSPACE_BUDGET = 2 * 1024 * 1024
+
+# Byte counts quoted in the traceback, reproduced by _FakeMnnvlWorkspace below.
+REPORTED_BUFFER_BYTES = 698368
+REPORTED_REQUIRED_BYTES = 802816
+
+# Tokens whose bf16 payload exceeds a single Lamport buffer but still fits the
+# whole workspace budget -- the window that used to reach the kernel and abort.
+OVERSIZED_TOKENS = 56
+# Largest token count that genuinely fits one Lamport buffer here.
+FITTING_TOKENS = 48
+
+
+class _FakeMnnvlWorkspace:
+    """Mimics flashinfer's MNNVL workspace capacity accounting.
+
+    flashinfer sizes one Lamport buffer at roughly ``payload / 3`` and allocates
+    three of them, so only about a third of the requested budget backs any single
+    all-reduce. ``is_buffer_size_sufficient`` is the real API that the production
+    code consults instead of reimplementing this arithmetic.
+    """
+
+    NUM_LAMPORT_BUFFERS = 3
+    backend = "mnnvl"
+
+    def __init__(self, max_token_num: int, hidden_dim: int, dtype: torch.dtype):
+        payload = max_token_num * hidden_dim * dtype.itemsize
+        self.buffer_size_bytes = round_up(payload // self.NUM_LAMPORT_BUFFERS, 1024)
+
+    def is_buffer_size_sufficient(
+        self,
+        tp_size: int,
+        num_tokens: int,
+        hidden_dim: int,
+        dtype: torch.dtype,
+        use_oneshot=None,
+    ) -> bool:
+        return num_tokens * hidden_dim * dtype.itemsize <= self.buffer_size_bytes
+
+
+class _FakeTensor:
+    """Just the surface ``should_use_fi_ar`` inspects, so no GPU is needed."""
+
+    def __init__(self, num_tokens: int, hidden_dim: int, dtype: torch.dtype):
+        self.shape = (num_tokens, hidden_dim)
+        self.dtype = dtype
+        self.is_cuda = True
+
+    def is_contiguous(self) -> bool:
+        return True
+
+
+@pytest.fixture
+def comm(monkeypatch):
+    """A FlashInferAllReduce wired to a stub MNNVL workspace.
+
+    ``__init__`` is bypassed because it needs a live process group and a CUDA
+    platform; the attributes it would set are filled in directly. The stub is
+    cached like the real process-wide singleton, so the first caller's shape
+    fixes the capacity for every later caller.
+    """
+    singleton: list[_FakeMnnvlWorkspace] = []
+
+    def fake_get_workspace(world_size, rank, max_token_num, hidden_dim, dtype, group):
+        if not singleton:
+            singleton.append(_FakeMnnvlWorkspace(max_token_num, hidden_dim, dtype))
+        return singleton[0]
+
+    monkeypatch.setattr(
+        flashinfer_all_reduce, "get_fi_ar_workspace", fake_get_workspace
+    )
+
+    obj = FlashInferAllReduce.__new__(FlashInferAllReduce)
+    obj.disabled = False
+    obj.group = None
+    obj.world_size = WORLD_SIZE
+    obj.rank = 0
+    obj.device = "cuda:0"
+    obj.max_workspace_size = WORKSPACE_BUDGET
+    obj.max_num_tokens = 0
+    obj.workspace = None
+    return obj
+
+
+def test_rejects_tensor_larger_than_one_lamport_buffer(comm):
+    """The reported crash: 56 tokens fit the budget but not a Lamport buffer."""
+    tensor = _FakeTensor(OVERSIZED_TOKENS, HIDDEN_DIM, DTYPE)
+
+    assert comm.should_use_fi_ar(tensor) is False
+
+    # The cheap budget bound admits it, so only the authoritative workspace check
+    # can be what rejected it. Without that check this tensor reached the kernel.
+    assert comm.max_num_tokens >= OVERSIZED_TOKENS
+
+
+def test_accepts_tensor_that_fits_one_lamport_buffer(comm):
+    tensor = _FakeTensor(FITTING_TOKENS, HIDDEN_DIM, DTYPE)
+
+    assert comm.should_use_fi_ar(tensor) is True
+    assert comm.workspace is not None
+
+
+def test_reproduces_the_byte_counts_from_the_traceback(comm):
+    oversized = _FakeTensor(OVERSIZED_TOKENS, HIDDEN_DIM, DTYPE)
+    assert comm.should_use_fi_ar(oversized) is False
+
+    # "Buffer: 698368 bytes, Required: 802816 bytes."
+    assert comm.workspace.buffer_size_bytes == REPORTED_BUFFER_BYTES
+    required = OVERSIZED_TOKENS * HIDDEN_DIM * DTYPE.itemsize
+    assert required == REPORTED_REQUIRED_BYTES
+
+
+def test_still_rejects_tensors_beyond_the_whole_budget(comm):
+    """The cheap pre-check must keep short-circuiting huge prefill tensors."""
+    huge = _FakeTensor(8192, HIDDEN_DIM, DTYPE)
+
+    assert comm.should_use_fi_ar(huge) is False
+    # Rejected before any workspace was created.
+    assert comm.workspace is None
+
+
+def test_gate_follows_payload_not_the_first_shape_seen(comm):
+    """A drafter with a different hidden size must not inherit a stale limit."""
+    assert comm.should_use_fi_ar(_FakeTensor(FITTING_TOKENS, HIDDEN_DIM, DTYPE)) is True
+
+    # Same token count, twice the hidden size: twice the payload, so it no longer
+    # fits, even though the cached token-count bound would still admit it.
+    wide = _FakeTensor(FITTING_TOKENS, HIDDEN_DIM * 2, DTYPE)
+    assert comm.max_num_tokens >= FITTING_TOKENS
+    assert comm.should_use_fi_ar(wide) is False

vllm/distributed/device_communicators/flashinfer_all_reduce.py

@@ -359,15 +359,20 @@ def __init__(
                 self.world_size,
             )
             return
-        self.max_workspace_size = max_workspace_size * MiB
+        self.max_workspace_size = int(max_workspace_size * MiB)
         self.max_num_tokens = 0
+        self.workspace: Any = None
         self.disabled = False
 
     def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool:
         """Ensure the all reduce workspace is initialized."""
         if self.max_num_tokens == 0:
-            element_size = torch.tensor([], dtype=dtype, device="cpu").element_size()
-            self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size)
+            self.max_num_tokens = self.max_workspace_size // (
+                hidden_dim * dtype.itemsize
+            )
+        # Always go through the accessor rather than caching only on first use:
+        # the workspaces are process-wide singletons that destroy_fi_ar_workspace()
+        # can clear underneath us.
         workspace = get_fi_ar_workspace(
             world_size=self.world_size,
             rank=self.rank,
@@ -379,6 +384,7 @@ def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool:
         if workspace is None:
             self.disabled = True
             return False
+        self.workspace = workspace
         return True
 
     def should_use_fi_ar(self, input_tensor: torch.Tensor) -> bool:
@@ -396,27 +402,39 @@ def should_use_fi_ar(self, input_tensor: torch.Tensor) -> bool:
 
         num_tokens, hidden_dim = input_tensor.shape
         if not self.max_num_tokens:
-            element_size = torch.tensor([], dtype=input_tensor.dtype).element_size()
-            self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size)
+            self.max_num_tokens = self.max_workspace_size // (
+                hidden_dim * input_tensor.dtype.itemsize
+            )
 
+        # Cheap upper bound: no workspace can hold more than the whole size
+        # budget, so reject obviously over-sized tensors before paying for
+        # workspace creation. Not authoritative -- see below.
         if num_tokens > self.max_num_tokens:
             return False
 
-        return self._ensure_workspace(hidden_dim, input_tensor.dtype)
+        if not self._ensure_workspace(hidden_dim, input_tensor.dtype):
+            return False
 
-    def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor:
-        num_tokens, hidden_dim = input_tensor.shape
-        workspace = get_fi_ar_workspace(
-            world_size=self.world_size,
-            rank=self.rank,
-            max_token_num=self.max_num_tokens,
+        # Authoritative capacity check. max_workspace_size budgets the whole
+        # *allocation*, but a backend may only devote a fraction of it to any one
+        # all-reduce: mnnvl is Lamport-based and splits its allocation into three
+        # buffers, so only about a third of the budget is usable per call. Ask the
+        # workspace instead of reimplementing that arithmetic -- otherwise tensors
+        # sized between the real capacity and the budget pass this gate and then
+        # abort the engine inside flashinfer with "The buffer size in the given
+        # workspace is insufficient for the given problem size".
+        return self.workspace.is_buffer_size_sufficient(
+            tp_size=self.world_size,
+            num_tokens=num_tokens,
             hidden_dim=hidden_dim,
             dtype=input_tensor.dtype,
-            group=self.group,
         )
+
+    def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        num_tokens = input_tensor.shape[0]
         return flashinfer_comm.allreduce_fusion(
             input=input_tensor,
-            workspace=workspace,
+            workspace=self.workspace,
             pattern=flashinfer_comm.AllReduceFusionPattern.kAllReduce,
             launch_with_pdl=True,
             trigger_completion_at_end=num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
@@ -425,3 +443,4 @@ def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor:
     def destroy(self):
         if not self.disabled:
             destroy_fi_ar_workspace()
+        self.workspace = None

vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py

@@ -82,6 +82,7 @@ def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool
 
     num_tokens, hidden_size = hidden_states.shape
     max_token_num = _max_token_num(tp_size, hidden_size, hidden_states.dtype)
+    # Cheap upper bound only; the authoritative check needs the workspace.
     if max_token_num is None or num_tokens > max_token_num:
         return False, 0
 
@@ -97,6 +98,19 @@ def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool
     )
     if workspace is None:
         return False, 0
+    # max_token_num budgets the whole *allocation*, but a backend may only devote
+    # a fraction of it to any one call -- mnnvl is Lamport-based and splits its
+    # allocation into three buffers. Ask the workspace rather than trusting the
+    # budget, otherwise tensors sized between the real capacity and the budget
+    # reach the kernel and abort with "The buffer size in the given workspace is
+    # insufficient for the given problem size".
+    if not workspace.is_buffer_size_sufficient(
+        tp_size=tp_size,
+        num_tokens=num_tokens,
+        hidden_dim=hidden_size,
+        dtype=hidden_states.dtype,
+    ):
+        return False, 0
     return True, max_token_num