[seonjinn]

What does this PR do?

Allows NeMo-RL MXFP8 rollout refit to preserve an explicitly selected vLLM CuTeDSL or TRTLLM dense linear backend.

These native kernels replace checkpoint-layout parameters with backend-specific runtime layouts. Re-running their post-load hook directly during refit would replace CUDA-Graph-visible storage. This change uses vLLM's layerwise reload lifecycle for native MXFP8 linear layers, so each transformed weight is copied into the original runtime storage. Failed or incomplete reloads restore saved runtime parameters without processing partial checkpoint tensors.

The current auto behavior remains unchanged: NeMo-RL keeps FlashInfer CUTLASS as its default refit backend. Non-native refit paths now run process_weights_after_loading under set_current_vllm_config and reset the LoRA, encoder and multimodal caches after each refit. The reset calls depend on vLLM model runner method names.

NeMo-RL's vLLM 0.29 pin includes vLLM PR #52204, which provides the TRTLLM MXFP8 linear backend used by this PR.

Update (review fixes, cf947cc)

• Config time: normalize_vllm_refit_config rejects precision: fp8 + is_mx + a native linear_backend (flashinfer_cutedsl, flashinfer_trtllm) on nccl_reshard and the sparse and NIXL transports.
• Runtime: the checkpoint-engine lifecycle re-checks native MXFP8 linear kernels, so a subclass MRO cannot bypass the reject. The unreachable pending_roots branch is removed.
• DeepSeek V4 FP8 is rejected on nccl_reshard again, as on main (feat: add DeepSeek V4 Flash GRPO support - Automodel path #3595).
• process_weights_after_loading_mxfp8_linear raises when no vLLM config context is set. It no longer demotes an explicit CuTe-DSL backend to CUTLASS.
• The reload root scan imports the fp8 module only when a module has a linear kernel. The FlashInfer import is guarded.
• Native MXFP8 linear refit errors are fatal, as for TRTLLM MoE.
• Tests: nccl_reshard case restored, random-byte layout oracle with pad-row checks, fused q/k/v shards in separate batches, real-vLLM kernel name resolution, no-context raise, DSV4 nccl_reshard reject.
• Gate on OCI-HSG (vLLM extra, H100): 513 passed, 1 skipped for the affected vLLM unit files. The PR head before the fixes had 16 failures in the same files.

Usage

policy:
  generation:
    vllm_cfg:
      precision: fp8
      is_mx: true
    vllm_kwargs:
      linear_backend: flashinfer_cutedsl
      # Or: flashinfer_trtllm

Before your PR is "Ready for review"

• Read and followed the contributor guidelines.
• Added repeated-refit, layout-change, failure-cleanup, backend-default, and scale-routing tests.
• Ran the affected vLLM unit suites: 50 passed.
• Ran Ruff check and format check on all changed files.
• Ran Pyrefly: 0 errors.
• Run a vLLM 0.29 GB200 repeated-refit smoke for each explicit native backend.
  • 893a738 adds flashinfer_cutedsl and flashinfer_trtllm legs to the L1 GB200 MXFP8 shard (grpo_vllm_mxfp8_rollout_gb200.sh, Qwen3-0.6B, 2 GPUs, precision: fp8, is_mx: true, vllm_kwargs.linear_backend). Each native leg checks that the native kernel owns the refit (NRL_MXFP8_NATIVE_LINEAR_REFIT kernels=<kernel> modules=N). a7eeef4 runs every leg of the script, the default refit_prequantize legs included, for 3 GRPO steps with CUDA graphs on (enforce_eager=false). All 4 legs passed on OCI GB200 (modules=56 for both native kernels).

Additional Information

The repeated-refit tests cover CuTeDSL-style transpose and TRTLLM-style padding/layout changes while asserting stable Parameter identity and data_ptr() across two refits. Failure-path tests cover partial transfer abort, initialization failure, and one-shot finalization cleanup.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[seonjinn]

Self-review completed with independent correctness, test, and upstream-compatibility passes.

Resolved findings:

• Added an abort path that restores saved runtime tensors instead of processing incomplete checkpoint tensors.
• Protected partial initialization and made finalization one-shot while preserving the first error.
• Added CuTeDSL transpose and TRTLLM padding/layout regression cases over two refits.
• Added direct coverage for the existing CUTLASS default and backend-specific scale routing.
• Replaced class-name-only detection with dynamic isinstance checks against registered vLLM kernel types.

Validation:

• Affected vLLM unit suites: 50 passed.
• Ruff check and format check passed.
• Pyrefly: 0 errors.

Remaining dependency: the real TRTLLM GPU integration requires vLLM PR #52204 in NeMo-RL's dependency pin. This PR remains draft until that test can run.

[seonjinn]

/ok to test eb05783

[seonjinn]

MXFP8 dense backend: generation performance comparison

Tested PR3638's native-linear implementation integrated onto main b6016e678, commit 1d2852cb, plus a targeted NCCL-reshard misc-path support adjustment. The adjustment is needed for the earlier small Async runs; the performance recipes use IPC/default collective. These measurements are not exact-head validation of PR commit 32c8288a. The separate #3294 optimization stack is excluded.

Generation performance

MoE backend: flashinfer_trtllm. Dense linear: flashinfer_cutlass → flashinfer_cutedsl. MXFP8 covers eligible MoE and dense linear layers; explicit ignore list is [lm_head]. Official October 7 nightly, vLLM 0.29.0, FlashInfer 0.6.18, GB200 GPUs.

20 actual GRPO steps per arm, statistics over steps 2–20. Each pair uses the same nodes, model, seed, dataset and configuration, with only the selected dense backend changed. These are single-run comparisons. Lower generation time and higher throughput are better.

Model	Mode	GBS	Generation time: CUTLASS → CuTeDSL (s/batch)	Generation throughput: CUTLASS → CuTeDSL (output tok/s)	Throughput change
Qwen3-30B-A3B	Sync	2048	55.922 → 54.152	116,093.3 → 119,915.1	+3.3%
Qwen3-30B-A3B	Async	2048	79.066 → 78.091	82,051.5 → 83,056.8	+1.2%
Nemotron 3.5 Lightning 30B-A3B	Sync	2048	51.471 → 51.094	96,122.0 → 97,003.6	+0.9%
Nemotron 3.5 Lightning 30B-A3B	Async	2048	68.661 → 67.441	72,081.1 → 73,292.1	+1.7%
Qwen3-235B-A22B*	Async	512	138.590 → 124.827	20,853.5 → 23,199.1	+11.2%
Qwen3-235B-A22B	Sync	512	113.545 → 98.989	25,472.6 → 29,158.1	+14.5%
Nemotron 3 Super 120B-A12B†	Async	256	79.271 → 90.060	10,390.1 → 8,735.6	-15.9%

Definition: Generation time is rollout batch completion time, including reward/postprocessing, not a pure decode-kernel timer. Sync excludes preparation/refit and includes final sleep/offload. Async includes in-flight pauses inside the batch and excludes prelaunch refit waiting. Throughput is actual accepted output tokens divided by summed batch durations; overlapping Async batches mean it is not aggregate engine wall throughput. Responses use natural EOS.

Qwen235B memory-pressure qualification (*): CUTLASS logged repeated non-fatal CUDA allocator OOM warnings during initial refit and measured steps; CuTeDSL had no OOM messages. Both completed 20 steps with finite losses/gradient norms, zero reported failed generations and scheduler exit 0. The warnings' timing impact has not been isolated, so the +11.2% is the observed difference under this memory configuration.

Super Async overlap qualification (†): CuTeDSL has lower batch-normalized generation throughput in this pair, but this is not a 15.9% kernel or aggregate-engine regression. The batch-duration mix changed in this workload with in-flight refit: 9/19 CUTLASS batches versus 2/19 CuTeDSL batches completed under 70 s. Median batch duration was 97.285 → 94.716 s while the mean was 79.271 → 90.060 s. Aggregate engine wall throughput was 10,488.1 → 10,355.2 output tok/s (-1.27%), and mean output length was 3217.3 → 3073.2 tokens. The duration pattern is consistent with different overlap with refit pauses, but individual pause attribution is not proven. The effects of overlap timing and output-length variation have not been isolated. Both runs completed 20 finite steps with zero generation failures and no OOM messages.

Large-workload status (updated October 9)

• Qwen3-30B-A3B and Nemotron 3.5 Lightning GBS 2048: both Sync/Async comparisons are complete above. Qwen Async used consecutive backend arms in one allocation.
• Qwen/Qwen3-235B-A22B (original model): Sync and Async comparisons are now complete. Official GBS 512 = 16 × 32, sequence/output 8192, 64 GPUs Sync / 128 GPUs Async, generation TP4/EP1 with 16 replicas. The successful Sync pair retains optimizer state on GPU during refit, identically for both backends. It has no OOM messages; the earlier Async memory-pressure qualification applies only to the Async row.
• Nemotron 3 Super 120B-A12B: official GBS 256 = 32 × 8, sequence/output 8192, 128 GPUs for each mode (Sync colocated; Async 64 training + 64 generation), generation TP4/EP1. Async is now complete above with the overlap qualification. Sync remains incomplete: the retry again failed host OOM before any completed steps, despite retaining optimizer states on GPU for logprobs/refit; its dependent CuTeDSL arm was cancelled. No GBS 256 Sync speedup is claimed. The common overlay widens routed-MoE-only scope to eligible dense linears, sets SSM cache FP32 and disables async checkpoint workers. Earlier GBS 32 compatibility results below remain separate.

Workload, output lengths and quantization scope

• Qwen30 and Lightning: GBS 2048 = 64 prompts × 32 responses, maximum sequence/output 4096, 16 GPUs. Qwen uses the official performance recipes; Lightning adapts that workload with its Mamba settings. Sync: 16 colocated GPUs, trainer TP1/EP16, 16 generation TP1/EP1 engines, IPC refit. Async: 8 training + 8 generation GPUs, trainer TP1/EP8, 8 engines, in-flight collective updates.
• Qwen235B original model: GBS 512 = 16 × 32, maximum sequence/output 8192. Async: 64 training + 64 generation GPUs, trainer TP4/PP4/CP1/EP16/ETP1, 16 generation TP4/EP1 engines, in-flight collective updates.
• Same 4,096-problem OpenMathInstruct-2 subset, seed 23, no validation/checkpoint saving. Training and KV cache are BF16; Nemotron SSM cache is FP32. Startup/step 1 excluded; later compilation remains included.

Model	Mode	Mean output length: CUTLASS → CuTeDSL
Qwen3-30B-A3B	Sync	3170.0 → 3170.7
Qwen3-30B-A3B	Async	3167.7 → 3167.0
Nemotron 3.5 Lightning 30B-A3B	Sync	2415.8 → 2420.1
Nemotron 3.5 Lightning 30B-A3B	Async	2416.6 → 2413.5
Qwen3-235B-A22B*	Async	5644.7 → 5656.0
Qwen3-235B-A22B	Sync	5649.0 → 5637.3
Nemotron 3 Super 120B-A12B†	Async	3217.3 → 3073.2

The explicit exclusion is quantization_ignore_patterns: [lm_head]. Eligible
attention QKV/output, routed experts, shared-expert linears and Nemotron Mamba
input/output (and latent-MoE) projections use MXFP8. Model-specific exceptions
still apply: Qwen3MoE routers receive quantization; NemotronH routers explicitly
use unquantized GateLinear. Mamba convolution/state parameters, normalization,
embeddings and the excluded LM head are outside the MXFP8 linear/expert scope.
Training remains BF16 and KV cache uses BF16; Nemotron SSM cache remains FP32.

This was checked against the running nightly image: SHA256 of qwen3_moe.py,
nemotron_h.py, mamba_mixer2.py, router/gate_linear.py and quantization/modelopt.py
matches the corresponding vLLM v0.29.0 source files. See
Qwen router,
Nemotron router
and GateLinear exclusion.

Earlier Super GBS 32 compatibility workload: generation results

These reduced workloads use 8 prompts × 4 responses, maximum input/output 512/512 and total sequence 1024. They have different GPU allocations from the performance workloads and must not be used to rank backends for large batches. Each pair uses the same nodes; compare backends only within a model/mode. Super Sync's CUTLASS training arm completed 20 steps before an allocation cleanup failure; its CuTeDSL continuation used the exact same nodes.

Model	Mode	GBS	Generation time: CUTLASS → CuTeDSL (s/batch)	Generation throughput: CUTLASS → CuTeDSL (output tok/s)	Throughput change
Nemotron 3 Super 120B-A12B	Sync	32	3.999 → 3.558	3,483.3 → 3,930.1	+12.8%
Nemotron 3 Super 120B-A12B	Async	32	5.210 → 4.938	2,683.4 → 2,805.8	+4.6%

Reference: training E2E, refit and aggregate engine throughput

Generation is the main comparison above. These previously requested measurements remain supplemental. Backend weight layouts/refit and memory use can also influence the full training cycle. Training E2E uses the wall interval between logging after steps 1 and 20, including inter-step logging. Engine counters use that same window, including idle/refit and unused rollouts.

Model	Mode	GBS	Training wall s/step: CUTLASS → CuTeDSL	Training output tok/s: CUTLASS → CuTeDSL	Refit s: CUTLASS → CuTeDSL
Qwen3-30B-A3B	Sync	2048	168.893 → 166.686	38,439.3 → 38,957.3	8.887 → 9.089
Qwen3-30B-A3B	Async	2048	234.989 → 235.508	27,607.5 → 27,540.3	8.229 → 8.428
Nemotron 3.5 Lightning 30B-A3B	Sync	2048	123.375 → 118.532	40,101.6 → 41,814.3	2.835 → 2.925
Nemotron 3.5 Lightning 30B-A3B	Async	2048	124.828 → 126.396	39,647.8 → 39,106.2	2.647 → 2.691
Qwen3-235B-A22B*	Async	512	108.632 → 103.069	26,604.5 → 28,096.3	22.711 → 22.191

Model	Async aggregate engine output tok/s: CUTLASS → CuTeDSL
Qwen3-30B-A3B	27,504.5 → 27,440.8
Nemotron 3.5 Lightning 30B-A3B	39,647.5 → 39,126.4
Qwen3-235B-A22B*	26,601.1 → 28,192.9
Nemotron 3 Super 120B-A12B†	10,488.1 → 10,355.2

[seonjinn]

Follow-up: generation impact of direct CuTeDSL scale refit

This is an additional experimental optimization, separate from the CUTLASS-versus-CuTeDSL comparison above. Both arms use moe_backend=flashinfer_trtllm and linear_backend=flashinfer_cutedsl. Only NRL_MXFP8_DIRECT_SCALE_REFIT changes from 0 to 1. It combines scale padding, swizzling and final scale copying into one kernel writing existing runtime storage; weight quantization, shard loading and generation GEMMs are unchanged.

Measured source: experimental commit 348e81ac, based on the previously measured integration. This optimization is not in PR #3638's current head.

Generation performance

Qwen/Qwen3-30B-A3B, performance recipe, GBS 2048, 16 GB200 GPUs, official October 7 nightly with vLLM 0.29. Each mode ran baseline then optimized in one allocation on identical nodes, 20 actual GRPO steps per arm; statistics use steps 2–20. MXFP8 covers eligible MoE and dense linear layers with [lm_head] as the explicit ignore list; training/KV are BF16. Same model revision, source, seed and dataset. Configuration audit found only the optimization flag and logging paths differed.

Model	Mode	Generation time: baseline → optimized (s/batch)	Generation throughput: baseline → optimized (output tok/s)	Throughput change
Qwen3-30B-A3B	Sync	54.266 → 53.639	119,683.5 → 121,128.9	+1.21%
Qwen3-30B-A3B	Async	78.505 → 78.096	82,669.9 → 83,052.6	+0.46%

Interpretation: observed generation differences are small. One ordered pair does not establish a repeatable generation speedup; the complete refit timing also shows no consistent improvement. The earlier 42–48% reduction was layer postprocessing microbenchmark time, approximately 55–70 µs saved per invocation, not generation or complete refit time. The experimental flag remains disabled by default.

Generation time is rollout batch completion time, including reward/postprocessing. Sync excludes preparation/refit; Async uses the native full-batch duration and can include in-flight refit pauses, but excludes prelaunch waiting. Throughput is actual accepted output tokens divided by summed batch durations. With overlapping Async batches this is a batch-normalized rate, not aggregate engine wall throughput. Natural EOS is enabled.

All four arms completed 20 steps and exited successfully, with finite loss/gradient metrics and zero reported generation failures. Worker logs confirm the optimized path was active only in the optimized arms. GPU validation also passed 17 new tests, 260 existing tests and four repeated-refit CUDA Graph lifecycle probes.

Supplemental refit and training-wall measurements

Mode	Mean refit: baseline → optimized (s)	Median refit: baseline → optimized (s)	Training wall time: baseline → optimized (s/step)
Sync	9.355 → 9.135	9.116 → 9.145	167.883 → 167.932
Async	8.366 → 8.384	8.350 → 8.381	241.553 → 240.495

The Sync baseline's step-2 refit was 13.198 s. This accounts for its apparent mean-refit difference: diagnostic steps 3–20 averages are 9.141797 s versus 9.141367 s, effectively equal. The primary generation table retains the predefined steps 2–20 window. Async mean refit was slightly higher with the optimization. Training wall time uses the logging-boundary interval after steps 1 and 20, including logging overhead.

Mode	Mean output tokens/sample: baseline → optimized
Sync	3171.3 → 3172.5
Async	3169.0 → 3167.0

[guyueh1]

A few follow-up suggestions:

• Optional naming clarification: _is_flashinfer_mxfp8_linear_kernel() would describe the helper more precisely than _is_mxfp8_linear_kernel(), since it resolves kernel types specifically from vllm.model_executor.kernels.linear.mxfp8.flashinfer.
• Add a focused unit test for the refit_transport=null, refit_with_reload_api=true collective path using the real set_current_vllm_config context manager. The current test mocks the context manager and checks call order; it does not verify that reload sees exactly model_runner.vllm_config (including kernel_config.linear_backend), or that the previous context is restored after both success and an exception.
• Add GPU integration coverage for real CuTeDSL and TRTLLM MXFP8 kernels across two refits with changed weights and CUDA graph replay. Assert stable Parameter identity and storage pointers, and compare graph-replayed outputs with eager outputs. Cover both refit_with_reload_api values on collective. The current repeated-refit test uses synthetic transpose/padding kernels, so it does not exercise the real transforms or graph execution.
• Add a transport-level case where fused Q/K/V shards arrive in separate batches and the receive buffer is overwritten between batches. Verify the final weights/scales to exercise the native MXFP8 linear path together with deferred-loader buffer protection.

I would prioritize the config-context unit test and real-kernel GPU integration validation. These are coverage recommendations, not confirmed correctness bugs.

[seonjinn]

Thanks for the follow-up. In 8a39ac8 I renamed the helper to _is_flashinfer_mxfp8_linear_kernel and added a collective reload test using the real vLLM config context manager. It checks that reload sees the exact runner config and kernel_config.linear_backend, and that the previous context is restored after both success and a reload exception. The existing deferred-buffer-reuse test covers the general protection mechanism, but it does not cover separately delivered fused Q/K/V shards. Real CuTeDSL/TRTLLM kernels, two changed refits, and CUDA graph replay also remain unverified here

[seonjinn]

/ok to test 3e79c96

[seonjinn]

Self-review after merging latest main (head 3e79c96723). Seven threads: two medium correctness items (unreachable or inconsistent pending_roots branch, DSV4 widened to nccl_reshard), one import-scope item, one latent backend-demotion item, test gaps, and nits.

Validation: static review plus local compileall and ruff. Tests were not executed locally (no GPU, Linux-only lockfile).

[seonjinn]

/ok to test cf947cc

[seonjinn]

/ok to test 893a738

[seonjinn]

/ok to test a7eeef4

[seonjinn]

/ok to test 2870907

[seonjinn]

/ok to test 83c655c

[seonjinn]

/ok to test
4d1fb42

[terrykong]

Thanks for this — opt-in native CuTeDSL/TRTLLM dense kernels with refit that is safe under CUDA graphs, plus a careful self-review round. Every self-review item is fixed at head (checked; no replies needed).

Done well, worth keeping:

• Refit goes through vLLM's public layerwise reload, so changed weights are copied into the storage CUDA graphs captured. vLLM private code is touched only on the abort path, and that path is tested against real vLLM.
• BF16 TRTLLM experts and native MXFP8 linears reload in one pass, so one can't undo the other's layout.
• Kernel detection uses isinstance on the real vLLM classes, and the name test fails (not skips) on a rename.
• A startup reject, plus a runtime check on the kernel vLLM actually built.

Must fix before merge:

1. (C1) The new per-refit cache reset clears only the GPU worker's image cache, not the scheduler's. Non-colocated and in-flight-update VLM runs can crash with Encoder cache miss. Remove the helper.

Smaller items:
2. (C2) Run one native GB200 leg with REFIT_PREQUANTIZE=false.
3. (C3) Read linear_backend the same way vLLM does.
4. (C4) Keep the backend → kernel pairing in one place.
5. (C5) Document the new option in docs/fp8.md.
6. (C6, C7, C8) Comment and docstring fixes.

GPU CI has not run on this head yet (waiting for approval); the GB200 MXFP8 shard passed on 2870907.

Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-3638-native-refit.html
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-3638-encoder-cache-reset.html

Generated by Claude Code

[terrykong]

C1 · 1 action item — correctness bug, please fix in this PR.
TL;DR — after every refit, this helper clears the image/video encoder cache on the GPU worker but not in vLLM's scheduler. The scheduler still thinks the media is cached, skips the encoder, and the worker crashes with Encoder cache miss.

PR-introduced: on main this only ran behind refit_with_reload_api=True; this PR runs it on every refit, for every model.

1. Both finalizers now call this helper: the native one (L1804) and the default one used by every ipc / collective / checkpoint-engine / nccl_reshard refit (L1850).
2. GPUModelRunner.reset_encoder_cache only runs self.encoder_cache.clear() on the worker.
3. The scheduler's EncoderCacheManager.check_and_update_cache still has the hash in cached (finished requests leave it there as "freeable"), so it returns True and the encoder is not scheduled.
4. The worker then raises RuntimeError(f"Encoder cache miss for {mm_hash}."). Only EngineCore.reset_encoder_cache clears both sides.

Who hits it: VLM runs that refit without first pausing the engine with clear_cache=True — non-colocated runs, and async in-flight updates with recompute_kv_cache_after_weight_updates: false. The shipped Nemotron-Omni async video recipe is exactly that (in-flight updates, no KV recompute; colocated off). A later turn of a multi-turn Gym episode, a sibling rollout, or a preempted request sends the same media again after the refit and hits the stale entry. Colocated runs are safe: sleep clears both sides first.

AI-1

Delete _reset_model_runner_caches_after_refit and its two calls (L1804, L1850). In test_mxfp8_native_linear_refit_uses_vllm_layerwise_reload, drop the last two expected entries "reset_encoder_cache", "reset_mm_cache" but keep the two recording lambdas, so the test fails if the worker-only reset comes back. Not a suggestion block: three locations in two files.

Why not keep part of the helper, or swap in the engine-level reset

• reset_lora_state returns at once without a lora_config, so it does nothing on every path that reaches this helper: only native LoRA sets enable_lora, and native LoRA skips _weight_update_lifecycle.
• reset_mm_cache only clears a worker-local profiling cache.
• Clearing stale image embeddings after refit already has a safe home: the opt-in vllm_cfg.reset_encoder_cache_after_weight_update calls the engine-level reset that clears both sides. The exemplar says it is "Only safe when generation is quiesced".
• Calling the engine-level reset from this helper is not the fix either: the worker extension cannot reach the scheduler, and with in-flight updates even that call puts the two caches out of step.

Upstream vLLM has the same worker-only clear

vLLM's own GPUModelRunner.reload_weights does the same worker-only clear. That is why refit_with_reload_api=True already had this bug on main.

Test coverage

Not verified on GPU — no GPU environment. No GB200 leg covers it: all legs are text-only Qwen3-0.6B.

[terrykong]

C2 · 1 action item — test gap.
TL;DR — both native GB200 legs run with REFIT_PREQUANTIZE=true, so the new native branch on the vLLM-side quantize path (fp8.py#L751-L758) never runs on real kernels.

The script defaults REFIT_PREQUANTIZE=true. With it on, refit takes the _scale_from_checkpoint rename branch. The prequant=false branch that picks _scale for native modules is covered only by a mocked unit test (test_load_weights_uses_scale_name_for_mxfp8_linear_backend).

AI-1

Run the TRTLLM leg with prequant off. This swaps an existing leg, so it adds no runtime, and EXP_NAME already includes both knobs.

Suggested change

	run_test env LINEAR_BACKEND=flashinfer_trtllm uv run --no-sync bash \
	./tests/functional/grpo_vllm_mxfp8_rollout_gb200.sh
	run_test env REFIT_PREQUANTIZE=false LINEAR_BACKEND=flashinfer_trtllm uv run --no-sync bash \
	./tests/functional/grpo_vllm_mxfp8_rollout_gb200.sh

[terrykong]

C3 · 1 action item — config check can miss or wrongly reject a backend.
TL;DR — _requested_linear_backend reads the setting differently from vLLM, so the startup reject misses flashinfer-cutedsl and can reject a config vLLM would run as CUTLASS.

• vLLM lowercases the name and turns - into _. flashinfer-cutedsl passes this check, vLLM builds the native kernel, and the run only fails at the runtime reject, after model load.
• vLLM lets a top-level linear_backend that is not "auto" override kernel_config.linear_backend. Here kernel_config wins, so top-level flashinfer_cutlass + kernel_config.linear_backend: flashinfer_cutedsl runs CUTLASS in vLLM but is rejected here.

AI-1

Suggested change

	def _requested_linear_backend(config: VllmConfig) -> str | None:
	vllm_kwargs = config.get("vllm_kwargs") or {}
	kernel_config = vllm_kwargs.get("kernel_config") or {}
	if isinstance(kernel_config, Mapping) and kernel_config.get("linear_backend"):
	return kernel_config["linear_backend"]
	return vllm_kwargs.get("linear_backend")
	def _requested_linear_backend(config: VllmConfig) -> str | None:
	# Match vLLM: a top-level non-"auto" value overrides kernel_config, and
	# names are case-insensitive with "-" and "_" interchangeable.
	vllm_kwargs = config.get("vllm_kwargs") or {}
	backend = vllm_kwargs.get("linear_backend")
	if not backend or backend == "auto":
	kernel_config = vllm_kwargs.get("kernel_config") or {}
	if isinstance(kernel_config, Mapping):
	backend = kernel_config.get("linear_backend") or backend
	if isinstance(backend, str):
	backend = backend.lower().replace("-", "_")
	return backend

Also add the dash spelling and the both-keys-set case to test_native_mxfp8_linear_backend_rejected_on_non_native_refit_transports.

[terrykong]

C4 · 1 action item — design, low.
TL;DR — the backend → kernel pairing is written twice, backend names here and kernel class names in _NATIVE_MXFP8_LINEAR_REFIT_KERNELS, with nothing tying them together.

If a third native kernel is added to one list and not the other, the startup reject passes a config that only the runtime check on the GPU worker catches. (The class names themselves are already guarded by test_mxfp8_linear_kernel_names_resolve_on_installed_vllm; this is only about the pairing.)

AI-1

Keep one dict in config.py (it already has no vLLM import) and build the fp8.py set from it. Three edits, two files, so not a suggestion block:

1. Replace _NATIVE_MXFP8_LINEAR_BACKENDS on this line with:

# Native MXFP8 linear backends and the vLLM kernel each one builds
# (vllm/model_executor/kernels/linear/__init__.py). fp8.py derives its set from this.
NATIVE_MXFP8_LINEAR_KERNELS = {
    "flashinfer_cutedsl": "FlashInferCutedslMxfp8LinearKernel",
    "flashinfer_trtllm": "FlashInferTrtllmMxfp8LinearKernel",
}

2. In _reject_native_mxfp8_linear_refit_transport, check linear_backend not in NATIVE_MXFP8_LINEAR_KERNELS.
3. In fp8.py, replace the kernel-name set with set(NATIVE_MXFP8_LINEAR_KERNELS.values()), imported from config.py.

[terrykong]

C5 · 1 action item — missing docs.
TL;DR — the new vllm_kwargs.linear_backend: flashinfer_cutedsl | flashinfer_trtllm option and its limits are only in the PR description; nothing under docs/ mentions linear_backend.

AI-1

In docs/fp8.md, add a short subsection after "Optional trainer-side MXFP8 refit" with the PR's usage YAML and its limits:

• refit_transport must be null (IPC/collective); nccl_reshard, sparse, NIXL and checkpoint-engine are rejected.
• No FP8 KV cache.
• No co-trained MTP drafter.
• The default stays CUTLASS.

Not a suggestion block: the file is outside this PR's diff.

[terrykong]

C6 · 1 action item — nit, code comment.
TL;DR — in vLLM 0.29, auto never selects TRTLLM, so this comment describes a case that cannot happen.

On CUDA, vLLM tries Cutedsl, Cutlass, Marlin, B12x, Emulation, Humming, then Trtllm, and Emulation is always supported, so the list never reaches Trtllm.

AI-1

Suggested change

	# Under linear_backend=auto only CuTe-DSL is demoted to CUTLASS. A TRTLLM
	# kernel that vLLM selects keeps the native refit path below.
	# Under linear_backend=auto only CuTe-DSL is demoted to CUTLASS. An
	# explicitly selected TRTLLM kernel keeps the native refit path below.

[terrykong]

C7 · 1 action item — nit, missing docstring.
TL;DR — uses_native_mxfp8_linear_refit is public and called from _get_mxfp8_linear_reload_roots, but has no docstring.

AI-1

Suggested change

	def uses_native_mxfp8_linear_refit(module: torch.nn.Module) -> bool:
	def uses_native_mxfp8_linear_refit(module: torch.nn.Module) -> bool:
	"""Return True if the layer uses vLLM's FlashInfer CuTeDSL or TRTLLM MXFP8 linear kernel.
	These kernels replace the weight and scale with new tensors in their own
	layout after loading, so refit must go through vLLM's layerwise reload,
	which copies the result back into the storage CUDA graphs captured.
	"""