Skip to content

[Bug]: flashinfer_b12x MoE broken after PR #43328: "num_tokens exceeds max_num_tokens" with dp_size>1 and OOM in B12xMoEWrapper buffer allocation #47982

Description

@R-omk

Your current environment

The output of python collect_env.py
==============================
        System Info
==============================
OS                           : Ubuntu 22.04.5 LTS (x86_64)
GCC version                  : (Ubuntu 11.4.0-1ubuntu1~22.04.3) 11.4.0
Clang version                : Could not collect
CMake version                : Could not collect
Libc version                 : glibc-2.35

==============================
       PyTorch Info
==============================
PyTorch version              : 2.11.0+cu130
Is debug build               : False
CUDA used to build PyTorch   : 13.0
ROCM used to build PyTorch   : N/A
XPU used to build PyTorch    : N/A

==============================
      Python Environment
==============================
Python version               : 3.12.13 (main, Mar  4 2026, 09:23:07) [GCC 11.4.0] (64-bit runtime)
Python platform              : Linux-6.8.0-134-generic-x86_64-with-glibc2.35

==============================
       CUDA / GPU Info
==============================
Is CUDA available            : True
CUDA runtime version         : 13.0.88
CUDA_MODULE_LOADING set to   : 
GPU models and configuration : GPU 0: NVIDIA GeForce RTX 5090
Nvidia driver version        : Could not collect
cuDNN version                : Could not collect
HIP runtime version          : N/A
MIOpen runtime version       : N/A
Is XNNPACK available         : True

==============================
          CPU Info
==============================
Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Address sizes:                           46 bits physical, 57 bits virtual
Byte Order:                              Little Endian
CPU(s):                                  32
Vendor ID:                               AuthenticAMD
Model name:                              AMD EPYC 9115 16-Core Processor
CPU family:                              26
Model:                                   2
Thread(s) per core:                      2
Core(s) per socket:                      16
Socket(s):                               1
Stepping:                                1
Frequency boost:                         enabled
CPU max MHz:                             4115.8198
CPU min MHz:                             1500.0000
BogoMIPS:                                5191.79
Virtualization:                          AMD-V
L1d cache:                               768 KiB (16 instances)
L1i cache:                               512 KiB (16 instances)
L2 cache:                                16 MiB (16 instances)
L3 cache:                                64 MiB (2 instances)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-31

==============================
Versions of relevant libraries
==============================
[pip3] flashinfer-python==0.6.13
[pip3] numpy==2.2.6
[pip3] nvidia-cublas==13.1.0.3
[pip3] nvidia-cuda-cccl==13.3.3.4.1
[pip3] nvidia-cuda-crt==13.3.73
[pip3] nvidia-cuda-cupti==13.0.85
[pip3] nvidia-cuda-nvcc==13.2.78
[pip3] nvidia-cuda-nvrtc==13.0.88
[pip3] nvidia-cuda-runtime==13.0.96
[pip3] nvidia-cuda-tileiras==13.2.78
[pip3] nvidia-cudnn-cu13==9.19.0.56
[pip3] nvidia-cudnn-frontend==1.26.0
[pip3] nvidia-cufft==12.0.0.61
[pip3] nvidia-cufile==1.15.1.6
[pip3] nvidia-curand==10.4.0.35
[pip3] nvidia-cusolver==12.0.4.66
[pip3] nvidia-cusparse==12.6.3.3
[pip3] nvidia-cusparselt-cu13==0.8.0
[pip3] nvidia-cutlass-dsl==4.5.2
[pip3] nvidia-cutlass-dsl-libs-base==4.5.2
[pip3] nvidia-cutlass-dsl-libs-cu13==4.5.2
[pip3] nvidia-ml-py==13.610.43
[pip3] nvidia-nccl-cu13==2.28.9
[pip3] nvidia-nvjitlink==13.0.88
[pip3] nvidia-nvshmem-cu13==3.4.5
[pip3] nvidia-nvtx==13.0.85
[pip3] nvidia-nvvm==13.2.78
[pip3] pyzmq==27.1.0
[pip3] tokenspeed-triton==3.7.10.post20260531
[pip3] torch==2.11.0+cu130
[pip3] torch_c_dlpack_ext==0.1.5
[pip3] torchaudio==2.11.0+cu130
[pip3] torchcodec==0.14.0+cu130
[pip3] torchvision==0.26.0+cu130
[pip3] transformers==5.13.0
[pip3] triton==3.6.0
[conda] Could not collect

==============================
         vLLM Info
==============================
ROCM Version                 : Could not collect
vLLM Version                 : 0.23.1rc1.dev925+g2afa3f7e9 (git sha: 2afa3f7e9)
vLLM Build Flags:
  CUDA Archs: 7.5 8.0 8.6 8.9 9.0 10.0 12.0; ROCm: Disabled; XPU: Disabled
GPU Topology:
	GPU0	CPU Affinity	NUMA Affinity	GPU NUMA ID
GPU0	 X 	0-31	0		N/A

==============================
     Environment Variables
==============================
CUDA_VERSION=13.0.2
VLLM_USAGE_SOURCE=production-docker-image
PYTORCH_NVML_BASED_CUDA_CHECK=1
TORCHINDUCTOR_COMPILE_THREADS=1
TORCHINDUCTOR_CACHE_DIR=/tmp/torchinductor_root

🐛 Describe the bug

When using --moe-backend flashinfer_b12x (required for NVFP4 MoE on SM120/SM121 GPUs like RTX 5090) together with --data-parallel-size 2, vLLM fails during engine initialization with:

ValueError: num_tokens (16384) exceeds max_num_tokens (8192)

The error originates in flashinfer's B12xMoEWrapper.run() (flashinfer/fused_moe/cute_dsl/b12x_moe.py:495).

Probable root cause (analysis based on source code inspection, not confirmed by debugging)

The following is our best understanding based on reading the vLLM and flashinfer source code. We have not been able to fully confirm this by running with debug logging, so these conclusions should be treated as educated hypotheses rather than definitive root cause analysis.

  1. B12xMoEWrapper appears to be initialized with max_num_tokens = max_num_batched_tokens (set at vllm/model_executor/layers/fused_moe/layer.py:345, which passes vllm_config.scheduler_config.max_num_batched_tokens).

  2. When data_parallel_size > 1, the MoE runner may use "naive dispatch/combine" (vllm/model_executor/layers/fused_moe/runner/moe_runner.py:732):

    def do_naive_dispatch_combine(self) -> bool:
        return (
            self.moe_config.dp_size > 1 or self.moe_config.is_sequence_parallel
        ) and not self._quant_method.supports_internal_mk

    If this path is taken, it performs an all-gather of tokens across all DP ranks before calling the MoE kernel, effectively multiplying num_tokens by dp_size. However, we have not confirmed whether supports_internal_mk is True or False for this configuration — if it were True, naive dispatch/combine would be skipped, and the token doubling would need to come from another source.

  3. If B12xMoEWrapper.max_num_tokens was set to max_num_batched_tokens without accounting for dp_size, then the all-gathered token count (max_num_batched_tokens * dp_size) passed to wrapper.run() would exceed the wrapper's max_num_tokens limit.

  4. The error triggers during determine_available_memory()profile_run()_dummy_run() → model forward → MoE layer → B12xMoEWrapper.run().

Why reducing --max-num-batched-tokens doesn't help

Since the all-gather (if it occurs) always doubles the token count (num_tokens = max_num_batched_tokens * dp_size), reducing max_num_batched_tokens just scales both numbers proportionally:

  • max_num_batched_tokens=16384num_tokens=32768 exceeds max_num_tokens=16384
  • max_num_batched_tokens=8192num_tokens=16384 exceeds max_num_tokens=8192

This behavior is consistent with the naive dispatch/combine hypothesis, but we cannot fully rule out other explanations.

Workaround

No reliable workaround found with the current nightly.

Setting --data-parallel-size 1 avoids the num_tokens exceeds max_num_tokens error, but reveals a second OOM issue (see "Second issue" below). Even --gpu-memory-utilization 0.96 with --max-num-batched-tokens 8192 on a single RTX 5090 (32 GiB) is not enough, as B12xMoEWrapper._allocate_buffers requires ~8.5 GiB of static workspace on top of the 21.94 GiB model weights.

Running two separate vLLM containers (each with dp_size=1, tp_size=1) on separate GPUs would not help either, as each container hits the same OOM in B12xMoEWrapper._allocate_buffers.

Reproduction

docker run --gpus '"device=0,1"' --ipc host \
  -e CUDA_DEVICE_ORDER=PCI_BUS_ID \
  -e HF_HUB_OFFLINE=1 \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  vllm/vllm-openai:nightly \
  --model nvidia/Qwen3.6-35B-A3B-NVFP4 \
  --tensor-parallel-size 1 \
  --data-parallel-size 2 \
  --moe-backend flashinfer_b12x \
  --kv-cache-dtype fp8_e4m3 \
  --enable-chunked-prefill \
  --max-num-batched-tokens 16384 \
  --speculative-config '{"method":"mtp","num_speculative_tokens":3,"moe_backend":"triton"}' \
  --trust-remote-code \
  --max-model-len auto

Error output (stack trace)

(Worker_DP0 pid=339) ERROR 07-08 08:15:11 [multiproc_executor.py:901]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/modular_kernel.py", line 1283, in _fused_experts
(Worker_DP0 pid=339) ERROR 07-08 08:15:11 [multiproc_executor.py:901]     self.fused_experts.apply(
(Worker_DP0 pid=339) ERROR 07-08 08:15:11 [multiproc_executor.py:901]   File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py", line 283, in apply
(Worker_DP0 pid=339) ERROR 07-08 08:15:11 [multiproc_executor.py:901]     wrapper_output = wrapper.run(
(Worker_DP0 pid=339) ERROR 07-08 08:15:11 [multiproc_executor.py:901]   File "/usr/local/lib/python3.12/dist-packages/flashinfer/fused_moe/cute_dsl/b12x_moe.py", line 495, in run
(Worker_DP0 pid=339) ERROR 07-08 08:15:11 [multiproc_executor.py:901]     raise ValueError(
(Worker_DP0 pid=339) ERROR 07-08 08:15:11 [multiproc_executor.py:901] ValueError: num_tokens (32768) exceeds max_num_tokens (16384)

Suggested fix (unverified)

Based on the above analysis, a potential fix would be to multiply max_num_tokens by dp_size in vllm/model_executor/layers/fused_moe/layer.py when constructing FusedMoEConfig. However, this has not been tested and may not be the correct approach — the actual fix might need to be in flashinfer, or the issue might be that supports_internal_mk should be True for b12x (which would bypass naive dispatch/combine entirely).

# layer.py, around line 345 — UNVERIFIED, for discussion only
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
dp_size_ = dp_size if dp_size is not None else get_dp_group().world_size
moe_max_num_tokens = max_num_batched_tokens * max(dp_size_, 1)

moe_config = FusedMoEConfig(
    ...
    max_num_tokens=moe_max_num_tokens,
    ...
)

Additional context

  • This only affects flashinfer_b12x MoE backend (used for NVFP4 on SM120/SM121 consumer Blackwell GPUs).
  • Other MoE backends (e.g. flashinfer_cutlass) may handle this differently via supports_internal_mk which bypasses naive dispatch/combine.
  • The num_tokens exceeds max_num_tokens issue does not occur with data_parallel_size=1, but a second OOM issue appears instead (see below).
  • No existing issue or PR was found on GitHub for this specific combination.

Second issue: OOM in B12xMoEWrapper._allocate_buffers (affects dp_size=1 too)

After switching to --data-parallel-size 1 to avoid the first error, the engine still fails to start — now with a CUDA OOM during B12xMoEWrapper buffer allocation. This affects all flashinfer_b12x configurations, not just DP.

Model weights occupy 21.94 GiB on a single RTX 5090 (32 GiB total). The new B12xMoEWrapper (introduced by PR #43328) pre-allocates a large static workspace via allocate_sm120_static_workspace, bringing total PyTorch allocation to ~30.47 GiB — exceeding the 31.40 GiB GPU capacity even with --gpu-memory-utilization 0.96.

The previous stateless flashinfer_b12x_fused_moe() function did not pre-allocate static buffers, so this OOM did not occur.

This was confirmed by testing the vllm-qwen3.6-35b-nvfp4-b12x-1gpu service (single GPU, dp_size=1, gpu_memory_utilization=0.96, max_num_batched_tokens=8192) — it also crashes with the same OOM:

(EngineCore pid=135) ERROR 07-08 09:40:55 [core.py:1231]     self._wrapper = B12xMoEWrapper(
(EngineCore pid=135) ERROR 07-08 09:40:55 [core.py:1231]     self._allocate_buffers()
(EngineCore pid=135) ERROR 07-08 09:40:55 [core.py:1231]   File "/usr/local/lib/python3.12/dist-packages/flashinfer/fused_moe/cute_dsl/b12x_moe.py", line 410, in _allocate_buffers
(EngineCore pid=135) ERROR 07-08 09:40:55 [core.py:1231]     self._static_workspace = allocate_sm120_moe_workspace(
(EngineCore pid=135) ERROR 07-08 09:40:55 [core.py:1231]   File "/usr/local/lib/python3.12/dist-packages/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py", line 2280, in allocate_sm120_moe_workspace
(EngineCore pid=135) ERROR 07-08 09:40:55 [core.py:1231]     return allocate_sm120_static_workspace(
(EngineCore pid=135) ERROR 07-08 09:40:55 [core.py:1231]   File "/usr/local/lib/python3.12/dist-packages/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py", line 307, in allocate_sm120_static_workspace
(EngineCore pid=135) ERROR 07-08 09:40:55 [core.py:1231] torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 160.00 MiB. GPU 0 has a total capacity of 31.40 GiB of which 84.56 MiB is free. Including non-PyTorch memory, this process has 31.30 GiB memory in use. Of the allocated memory 30.47 GiB is allocated by PyTorch, and 160.58 MiB is reserved by PyTorch but unallocated.

Probable regression introduced by PR #43328

This is likely a regression, though we have not bisected it by running older nightlies. Based on source code history:

Before PR #43328 (commit b1384f5ec, merged July 6, 2026), the b12x MoE backend used a stateless function call:

flashinfer_b12x_fused_moe(x=hidden_states, ...)

PR #43328 appears to have switched to a stateful B12xMoEWrapper with pre-allocated buffers and a max_num_tokens validation check:

self._wrapper = B12xMoEWrapper(
    use_cuda_graph=True,
    max_num_tokens=self.max_num_tokens,  # = max_num_batched_tokens (no dp_size factor)
)

When use_cuda_graph=True, B12xMoEWrapper.run() raises ValueError if num_tokens > max_num_tokens. With dp_size > 1, if naive dispatch/combine all-gathers tokens across DP ranks (doubling the count), but max_num_tokens was not adjusted to account for dp_size, this would cause the crash.

This configuration (dp_size=2 + flashinfer_b12x) was working a few days prior to this report, which is consistent with — but not proof of — the PR #43328 regression hypothesis.

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions