[Bugfix][MoE] Snapshot max_cudagraph_capture_size into FusedMoEConfig

[aoshen02]

Purpose

Several MoE kernels/quant-methods read
get_current_vllm_config().compilation_config.max_cudagraph_capture_size
inside their __init__:

• FlashInferExperts.__init__ (fused_moe/experts/flashinfer_cutlass_moe.py)
• TrtLlmMxfp4ExpertsBase.__init__ (fused_moe/experts/trtllm_mxfp4_moe.py)
• GptOssMxfp4MoEMethod.__init__ and Mxfp4MoEMethod.__init__
  (quantization/mxfp4.py)

UnquantizedFusedMoEMethod.process_weights_after_loading (and the mxfp4 method
paths) re-instantiate these objects every time they run. On the initial
model load that happens inside a set_current_vllm_config(...) context, so the
lookup succeeds. But when process_weights_after_loading is invoked again
outside that context — e.g. an RL framework reloading rollout weights, or
vLLM's own layerwise weight-reload path — get_current_vllm_config() trips its
assertion and the reload aborts:

AssertionError: Current vLLM config is not set. This typically means
get_current_vllm_config() was called outside of a set_current_vllm_config()
context, or a CustomOp was instantiated at module import time or model forward
time when config is not set.

The FlashInfer CUTLASS unquantized (BF16/FP16) MoE backend is the most visible
case; the same crash reproduces on the mxfp4 paths.

Fix

compilation_config.max_cudagraph_capture_size is fixed for the engine's
lifetime, so snapshot it once into FusedMoEConfig.max_capture_size at
FusedMoE layer-construction time (which always runs inside the config
context, and where compilation_config is already in scope), and have the four
read sites use moe_config.max_capture_size. This removes the implicit global
dependency from kernel/quant-method construction, so rebuilding them during a
weight reload no longer requires an active set_current_vllm_config context.

This mirrors how max_num_tokens is already plumbed from the scheduler config
into FusedMoEConfig.

Test Plan

• pre-commit / ruff check + ruff format on the changed files.
• Existing MoE kernel tests:
  tests/kernels/moe/test_flashinfer.py,
  tests/kernels/moe/test_flashinfer_moe.py,
  tests/kernels/moe/test_ocp_mx_moe.py,
  tests/kernels/moe/test_trtllm_nvfp4_moe.py,
  tests/kernels/moe/test_moe.py.
• Repro: load an unquantized FlashInfer-MoE model, then call
  process_weights_after_loading(model, model_config, device) without an
  active set_current_vllm_config context (simulating an RL weight reload) and
  run a forward pass. Before: AssertionError: Current vLLM config is not set.
  After: reload completes and generations match the reloaded weights.

Test Result

Lint: pre-commit run on the changed files — ruff check, ruff format, mypy,
typos, SPDX all pass.

End-to-end RL training run (the original verl#6563 scenario). Reproduced and
fixed in a full verl GRPO run on 8×H200 in the report's image
(verlai/verl:vllm020.dev1: vLLM 0.20.2, flashinfer 0.6.8.post1). Model: an
unquantized MoE (Qwen3-30B-A3B, qwen3_moe) taking the FlashInfer CUTLASS MoE
backend (VLLM_USE_FLASHINFER_MOE_FP16=1); VeOmni FSDP actor (ulysses=2, ep=4),
vLLM rollout (tp=4). verl syncs updated actor weights into the rollout engine
each step via update_weights_from_ipc → process_weights_after_loading, which
rebuilds the FlashInfer MoE kernel outside any set_current_vllm_config
context. Same script/config on both sides; the only difference is this patch:

run	first weight reload (global_step 0)	training
unpatched 0.20.2	crash — AssertionError: Current vLLM config is not set (all TP workers)	dies immediately
with this change	reload completes	proceeds normally — reaches step:1+ with valid metrics (actor/entropy, rollout_probs_diff_valid=1, throughput)

The unpatched crash stack, captured live, is exactly the report's:

verl/.../vllm_rollout/utils.py:240   update_weights_from_ipc
  process_weights_after_loading(model, model_config, self.device)
vllm/.../model_loader/utils.py:107   process_weights_after_loading
vllm/.../fused_moe/unquantized_fused_moe_method.py:170   _setup_kernel
vllm/.../fused_moe/oracle/unquantized.py:356             make_unquantized_moe_kernel -> experts_cls(...)
vllm/.../fused_moe/flashinfer_cutlass_moe.py:98  __init__
  get_current_vllm_config().compilation_config.max_cudagraph_capture_size
vllm/config/vllm.py:1969  get_current_vllm_config -> AssertionError

Targeted object-level repro. Verified inside the
image from the originating report (verlai/verl:vllm020.dev1: vLLM 0.20.2,
flashinfer 0.6.8.post1) on an H200, which selects the FlashInfer CUTLASS
unquantized MoE backend. The repro builds the exact FlashInferExperts(moe_config, quant_config) object that process_weights_after_loading → _setup_kernel → make_unquantized_moe_kernel re-instantiates on every weight reload — once inside
set_current_vllm_config(...) (initial model load) and once outside it (the RL
weight-reload path, e.g. verl update_weights_from_ipc):

	in-context (initial load)	outside context (reload)
unpatched 0.20.2	builds OK	AssertionError: Current vLLM config is not set (verbatim from the report)
with this change (backported to 0.20.2)	builds OK	builds OK — no crash

A value-flow check confirmed the kernel sources the value from the config: with
moe_config.max_capture_size = 2048, the kernel read 2048 outside any context
(no global lookup).

Correctness of the snapshot (lifecycle invariance). Snapshotting at build
time is equivalent to the original per-reload global read only if
compilation_config.max_cudagraph_capture_size is fixed after startup. Audited
all assignment sites: they live solely in config-construction modules
(config/compilation.py, config/vllm.py, engine/arg_utils.py,
model_executor/models/config.py — all run during engine init, before the model
is built). There are no assignments in any runtime/serving/reload path
(v1/worker, v1/engine, model_loader, engine), so the value is frozen for
the engine's lifetime and the snapshot matches what the original code read at
every reload. (The pre-existing code re-read this global on every reload even
though it never changes, so it was already effectively a build-time constant.)

Notes

Test/benchmark sites that construct FusedMoEConfig directly do not pass
max_capture_size and fall back to the default 0; the kernels clamp it via
max(self.max_capture_size, 1), and it only feeds the FlashInfer autotuner's
tune_max_num_tokens hint, so behavior is unchanged for correctness.

AI assistance (Claude Code) was used to investigate and draft this change; the
submitter has reviewed every changed line.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[zyongye]

LGTM