How a faster MXFP8 kernel can keep its own weight layout and still take new weights every step

NVIDIA-NeMo/RL #3638 — links pin head 83c655cb and vLLM v0.29.0.

Before
Asking vLLM for its CuTeDSL dense kernel was silently replaced with CUTLASS, and the TRTLLM dense kernel failed with Unsupported MXFP8 linear kernel for refit. Only CUTLASS could take new weights.
This PR
Keeps an explicitly chosen CuTeDSL or TRTLLM kernel, and copies each step's new weights into it through vLLM's layerwise reload. The default (auto) still uses CUTLASS.
Status
New opt-in feature; works as described. Both kernels pass the GB200 refit test with CUDA graphs on. A side change in this PR has a separate bug: see page 2.

Two pages on this PR — this one: how the new path compares with the default one.

Background. In RL training the rollout engine (vLLM) gets new weights from the trainer every step; that copy is the refit. vLLM records CUDA graphs that read the weights at fixed GPU addresses, so a refit must write new values into the same tensors, never swap in new ones. Dotted words have a short explanation: hover or tap them.

After loading, every MXFP8 dense kernel may rearrange the weight into the layout it reads fastest. CUTLASS keeps the trainer's layout, so a refit can write straight into its tensors. The two kernels this PR enables do not: CuTeDSL stores the weight transposed, and TRTLLM pads the rows to a multiple of 128, shuffles them, and flattens the scale. So the PR routes these layers through vLLM's layerwise reload: load into a temporary trainer-layout copy, let the kernel rearrange it, then copy the result into the tensors the graph reads.

One dense layer, one refit, both paths

Example layer: weight [4, 256] = [output rows, inputs], one FP8 value each; scale [4, 8] = one 8-bit scale per 32 inputs (256 / 32 = 8).

1 · trainer sends 2 · where the bytes land 3 · kernel layout step 4 · graph reads default auto → CUTLASS weight, fp8 [4,256] …_scale_from_checkpoint [4,8] written into the live weight [4,256] in place same layout as trainer swizzle scale into live weight_scale same addresses ✓ fp8.py L751 picks the scale name fp8.py L1243 in-place scale the paths split here: _weight_update_lifecycle → _uses_native_layerwise_refit (vllm_backend.py L1752) native linear_backend: flashinfer_cutedsl weight, fp8 [4,256] …_scale [4,8] same bytes, other name live tensors set aside weight [256,4] (transposed) bytes go to a temp copy weight [4,256], scale [4,8] kernel rearranges the temp copy: weight.t() → [256,4] scale swizzled vLLM copies into the set-aside tensors same addresses layerwise.py L111 initialize flashinfer.py L121 CuTeDSL layerwise.py L392 TRTLLM is the same path; only step 3 differs: rows padded 4 → 128 and shuffled, scale flattened to one dimension. the graph keeps reading the same GPU tensors the one place the two paths differ orange = data at that step Both paths get the same quantized bytes from the trainer and end with the graph reading unchanged addresses. What differs: the scale's name, and whether the kernel's rearranging step runs on a temporary copy first.

The split is one check in _uses_native_layerwise_refit: it is true when the model has layers whose kernel uses_native_mxfp8_linear_refit recognizes. Those layers join the same reload as the BF16 TRTLLM MoE experts that already used it, in one pass, so neither group undoes the other's layout. If a reload fails halfway, _abort_layerwise_reload puts the set-aside tensors back and the worker refuses further refits.

Why keep both

default: auto → CUTLASSnative: flashinfer_cutedsl / flashinfer_trtllm
How to get itnothing to setvllm_kwargs.linear_backend with precision: fp8, is_mx: true
Generation speedbaselineCuTeDSL: +0.9% to +14.5% in the author's GB200 A/B over several MoE models; one async case was slower because of different overlap with refit pauses. No TRTLLM numbers posted.
Refit transportsallIPC and collective only. NCCL reshard, sparse, NIXL and checkpoint-engine refit are rejected at startup.
Other limitsnone newno FP8 KV cache, no co-trained MTP drafter
So what. The default does not change: auto still swaps CuTeDSL for CUTLASS. If you set linear_backend to a native kernel, refit now keeps it, at the cost of a reload pass per step and the transport limits above. The feature itself checks out. The separate bug the review found is in the cache reset this PR added to every refit: page 2.