feat(inference): per-token NVFP4 squared-ReLU and SwiGLU MoE via FlashInfer CuTe DSL

[YigongQin]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds a per-token NVFP4 x NVFP4 routed-MoE path to Megatron inference for squared-ReLU (Nemotron-family) and SwiGLU experts, running on FlashInfer's CuTe DSL block-scaled grouped GEMMs. It is CUDA-graph capturable and batch invariant, and comes with the same test coverage the torch MXFP8 path has: a bitwise quantizer reference, a BF16 accuracy check, CUDA-graph replay, batch invariance, and a train-gen parity test against TE NVFP4 training. The NVFP4 weight path, CUDA-graph-stable refresh, batch-invariant finalize and tests are shared by FlashInfer's other NVFP4-weight recipes (per-tensor NVFP4 x NVFP4 and BF16 x NVFP4), so adding those later only changes the activation side; see "Strength: one integration for every NVFP4-weight recipe" and "How to extend".

Background

Recipe. Weights are NVFP4 with 1x16 E4M3 block scales and one FP32 global scale per expert. Activations are quantized online with one FP32 scale per token instead of a static, calibrated per-tensor input scale. That removes the calibration step and the activation-outlier sensitivity of a per-tensor scale, which matters in RL because the policy weights change every step.

Kernel level (FlashInfer). All kernels already exist upstream; this PR only integrates them.

• flashinfer#3027 adds per-token NVFP4 quantization (nvfp4_quantize(..., per_token_activation=True)) and the per-token TRT-LLM-gen MoE.
• flashinfer#3645 adds per-token NVFP4 activations to the CuTe DSL MoE (CuteDslMoEWrapper.run(..., per_token_scale=...)), the backend this PR uses.
• flashinfer#3976 adds the deterministic two-stage finalize (use_fused_finalize=False) to the CuTe DSL NVFP4 MoE, which batch-invariant mode uses.
• flashinfer#4618 adds the ReLU2 (squared-ReLU) activation for per-token NVFP4 MoE, which Nemotron needs.
• FlashInfer 0.7.0 ships per-token NVFP4 with ReLU2 and SwiGLU on both the TRT-LLM-gen and the CuTe DSL (CuteDslMoEWrapper, quant_mode="w4a4") MoE backends.

Engine level (peers of Megatron inference).

• sglang#22918 adds SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION, which replaces the checkpoint's static input scales with 1 so FlashInfer computes per-token scales at runtime.
• sglang#39199 routes ReLU2 per-token NVFP4 MoE to TRT-LLM-gen, padding hidden to 512 and intermediate to 128.
• Both are validated with GSM8K accuracy thresholds. Neither tests bitwise behavior, batch invariance, CUDA graphs, or train-gen parity.

RL level.

• NeMo-RL#3566 trains with TE NVFP4 (per-token activation scaling on routed experts), refits BF16 weights into vLLM, and quantizes the routed-expert weights to NVFP4 inside the rollout worker.
• miles#1261 does NVFP4 rollout/train with TE's NVFP4Quantizer producing the checkpoint and SGLang (FlashInfer TRT-LLM routed MoE) serving it.
• Both keep BF16 as the refit format and quantize on the inference side. This PR brings the same design to Megatron inference, so Megatron RL can roll out in per-token NVFP4 without an external engine.

What Transformer Engine already has. TE's NVFP4BlockScaling recipe supports row_scaled_activation=True (one scale per token), and with disable_2d_quantization=True, disable_rht=True and disable_stochastic_rounding=True it is the training-side form of this recipe. That is what NeMo-RL#3566 trains with, and it is the reference for the train-gen parity test below, in the same way TE MXFP8 is the reference for the MXFP8 inference path (#8195). TE has no inference MoE kernel for this recipe, so inference uses FlashInfer.

What is added

New config field. inference_moe_nvfp4_recipe: Optional[Literal['per_token']] = None in TransformerConfig (--inference-moe-nvfp4-recipe per_token). It is a recipe choice rather than a bool so the other NVFP4 recipes (see "How to extend") can be added as new values. It requires transformer_impl='inference_optimized', inference_grouped_gemm_backend='flashinfer', non-gated squared_relu experts or gated F.silu (SwiGLU, gated_linear_unit=True) experts, without activation_func_tanh_clamp_scale and, for SwiGLU, without activation_func_clamp_value, glu_linear_offset or moe_mlp_glu_interleave_size, BF16 expert parameters (no fp8 / fp4_param), and FlashInfer >= 0.7.0 on SM100/SM103. Validation rejects every other combination with a specific error. batch_invariant_mode is allowed with this recipe (with the NVLS dispatcher when EP > 1), as it is for FlashInfer MXFP8.

megatron/core/inference/moe/flashinfer_nvfp4.py (new).

• prepare_routed_nvfp4_weights quantizes the BF16 FC1/FC2 expert stacks into the CuTe DSL layout: packed E2M1, 128x4-swizzled E4M3 block scales plus their MMA view, and per-expert FP32 alphas. SwiGLU is detected from the shapes (FC1 has twice FC2's input width): MCore's [gate; up] FC1 is padded per half and reordered into FlashInfer's alternating 64-row [up, gate] blocks. With out= it refreshes the existing buffers in place, so addresses captured by CUDA graphs stay valid across refits.
• quantize_routed_nvfp4_input pads and quantizes activations per token. The global scale is the recipe constant 1 / (448 * 6) passed as a Python float, so there is no host sync and it is graph-safe.
• flashinfer_routed_nvfp4_moe runs CuteDslMoEWrapper. Wrappers are cached per shape and created outside graph capture. Under torch.use_deterministic_algorithms or batch-invariant mode it uses the two-stage finalize, which reduces each token's top-k routes in a fixed order, instead of the fused atomic finalize. With the NVLS dispatcher it writes straight into the reduce-scatter buffer, so ordered_reduce_scatter_v keeps the cross-rank reduction order fixed.

InferenceGroupedMLP (experts.py). The BF16 expert stacks stay resident as the canonical weights and refit destination. The NVFP4 copy is built on the first inference forward, and the FlashInfer forward dispatches to it. refresh_flashinfer_nvfp4_weights() requantizes it in place.

Not in this PR: refit and RL. Nothing calls refresh_flashinfer_nvfp4_weights() after a refit or an optimizer step yet, so don't set inference_moe_nvfp4_recipe for RL rollouts until the follow-up PR (TODO: link) lands. That PR wires the refresh into megatron/core/resharding/refit.py and megatron/rl/rl_utils.py; it is split out to keep the refit and RL reviewer groups off this PR. Static inference from a loaded checkpoint is unaffected.

moe_layer.py. The NVFP4 path checks for the FlashInfer per-token APIs instead of flashinfer-jit-cache, because CuTe DSL kernels are JIT-compiled through nvidia-cutlass-dsl rather than shipped in the jit-cache wheel.

Dependencies. flashinfer-python is relaxed from <0.7.0 to <0.8.0 and uv.lock moves to 0.7.0.post1, which adds nccl-extensions. nvidia-cutlass-dsl is already a dependency of flashinfer-python.

Why CuTe DSL and not TRT-LLM-gen

Both FlashInfer backends support per-token NVFP4 with ReLU2 and SwiGLU on SM100/SM103. An earlier revision of this PR supported both; we kept only CuTe DSL:

• No runtime downloads. TRT-LLM-gen loads its cubins from an artifactory at first use (into ~/.cache/flashinfer/cubins unless flashinfer-cubin is installed). On our test node that download failed with "Disk quota exceeded" and the GEMM then failed at 128 and 256 tokens. CuTe DSL compiles locally through nvidia-cutlass-dsl.
• Controllable determinism. CuTe DSL exposes use_fused_finalize (flashinfer#3976), which gives a fixed-order top-k reduction for batch-invariant mode. Make args and kwargs optional positional arguments for the Module hooks. #3976 measures its cost at about 1.5-14% over the atomic finalize, highest at 1-4 tokens.
• Less padding and simpler weight prep. CuTe DSL pads to 256 (Nemotron 2688 -> 2816 hidden) where TRT-LLM-gen pads to 512 (2688 -> 3072), and it needs no TRT-LLM weight shuffling. A full requantization of 32 local experts (what each refit will cost) takes 16.7 ms against 21.5 ms for TRT-LLM-gen.
• No speed reason to keep TRT-LLM-gen. FlashInfer's own per-token comparison in flashinfer#3645 (bench_moe_deepseek.py --use-per-token-activation, DeepSeek-V3 shapes, B200, CUDA graphs) has TRT-LLM-gen about 12% faster at 128-256 tokens and CuTe DSL 1.12-1.29x faster from 512 to 4096 tokens, a 1.08x geomean in CuTe DSL's favor. Our earlier Nemotron-shape runs on GB300 showed the same pattern.

One backend also halves the test matrix and removes a config field.

Performance

MoE layer forward (router, dispatcher, experts and combine) under CUDA-graph replay, on one GB300 with EP = 1. The shapes are the DeepSeek-V3 shapes that FlashInfer uses for its per-token CuTe DSL vs TRT-LLM-gen comparison in flashinfer#3645: hidden 7168, intermediate 2048, 256 experts (all local) and top-8. There are two differences from #3645. The experts are squared-ReLU, Nemotron's activation, where #3645 uses SwiGLU. TODO: add the SwiGLU numbers (bench_moe_layer.py --activation swiglu). Routing is plain top-8 without group-limited routing, because Megatron's inference router does not support it. All three backends load the same BF16 weights. Times are in µs, the mean over 50 replays.

Tokens	Default: vLLM BF16	Default: torch MXFP8	Default: NVFP4 (this PR)	Batch-inv.: vLLM BF16	Batch-inv.: torch MXFP8	Batch-inv.: NVFP4 (this PR)
1	140.4	138.2	56.6	179.7	237.4	82.3
8	507.1	526.4	187.2	575.3	488.1	215.2
32	1389.8	1604.4	488.7	1502.3	1188.0	522.7
64	1878.8	2074.0	626.9	1990.7	1519.1	662.7
128	2162.9	2352.7	704.1	2216.6	1695.1	739.8
256	2247.1	2407.9	723.7	2339.4	1739.6	758.1
512	2336.9	2433.8	739.0	2498.8	1792.4	772.5
1024	2732.7	2486.9	767.4	2994.2	1889.0	810.5
2048	3003.8	2650.0	822.5	3502.6	2081.8	879.0
4096	4054.7	3474.3	1015.6	5006.8	2560.6	1107.4

Batch-invariant mode uses the te_native backend for the non-MoE GEMMs (router) in all three columns.

• Default mode: NVFP4 is 2.5x faster than BF16 at 1 token, 3.2x at 512 and 4.0x at 4096, and 3.2-3.4x faster than MXFP8 from 64 tokens up.
• Batch-invariant mode: NVFP4 is 2.2x faster than BF16 at 1 token, 3.2x at 512 and 4.5x at 4096, and about 2.3x faster than MXFP8 from 8 tokens up. The deterministic two-stage finalize costs 26 µs at 1 token and 4-15% over default mode from 8 tokens up, in line with the 1.4-14% that flashinfer#3976 reports.
• Where the time goes at 512 tokens (default mode, profiled): The expert GEMMs take 2180 µs in BF16, 1433 µs in MXFP8 and 656 µs in NVFP4, roughly in proportion to weight bytes. MXFP8 is no faster than BF16 overall because its unpermute kernel takes another 842 µs; in batch-invariant mode, MXFP8's expert-order unpermute takes 61 µs, which is why MXFP8 is faster there than in its own default mode.
• Accuracy: cosine similarity to vLLM BF16 is about 0.998 for MXFP8 and 0.98 for NVFP4 in both modes, which is the expected precision gap.
• Memory: layer memory after the first forward is 18.4 GB for NVFP4 against 14.4 GB for BF16, because the BF16 expert stacks stay resident as the refit source.

Known gap with EP > 1. With the NVLS dispatcher, the expert input buffer is sized for the largest batch across all ranks. The vLLM and MXFP8 paths process only the valid rows (vLLM reads a device-side valid_tokens count; FlashInfer MXFP8 takes inference_flashinfer_mxfp8_token_capacity), but the NVFP4 path quantizes and routes the whole buffer. On Nemotron-3 Super shapes at EP = 4 with an 8192-token buffer, this gives NVFP4 a floor of about 820 µs per layer, against 80 µs for BF16 at 1 token, although NVFP4 is still 1.4x faster than BF16 at 8192 tokens per rank. TODO: bound the NVFP4 rows the same way before marking this PR ready.

Requantizing a layer's expert weights for a refit takes 16.7 ms for 32 Nemotron-3 Super experts.

Strength: one integration for every NVFP4-weight recipe

FlashInfer's CuteDslMoEWrapper runs three NVFP4-weight recipes on the same expert-weight layout: per-token NVFP4 x NVFP4 (this PR), per-tensor NVFP4 x NVFP4, and BF16 x NVFP4. FlashInfer's own tests (tests/moe/test_cute_dsl_fused_moe.py, cases w4a4-per-token, w4a4-per-tensor, w4a16 and w4a16-relu2) feed identical weight tensors to all three. Everything this PR builds is therefore shared, and the other two recipes only differ in how activations enter the kernel:

• Weights: prepare_routed_nvfp4_weights (BF16 -> packed E2M1 with swizzled E4M3 block scales, padded to 256) and its in-place, CUDA-graph-stable refresh after refit.
• Kernel plumbing: the cached wrapper, the two-stage finalize for batch invariance, and writing into the NVLS reduce-scatter buffer.
• Model integration: the build, refresh and dispatch in InferenceGroupedMLP, and the config validation.
• Tests: the BF16 reference, CUDA-graph replay, batch invariance, and the layer and NVLS tests all take the recipe as a parameter.

inference_moe_nvfp4_recipe is a recipe choice rather than a bool for this reason: each variant below is a new value, not a new flag.

How to extend

	Per-token NVFP4 x NVFP4 (this PR)	Per-tensor NVFP4 x NVFP4	BF16 x NVFP4
Config value	'per_token'	'per_tensor'	'weight_only'
quant_mode	w4a4	w4a4	w4a16
FC1 input	nvfp4_quantize(x, 1/(448*6), per_token_activation=True)	nvfp4_quantize(x, a1_scale) with a fixed scale tensor	BF16 hidden states, padded; not quantized
x_sf	block scales	block scales	None
per_token_scale	from the quantizer	None	None
fc1_alpha / fc2_alpha	1 / w_scale	1 / (w_scale * a1_scale) / 1 / (w_scale * a2_scale)	1 / w_scale
fc2_input_scale	1/(448*6)	a2_scale	None
Use case	RL rollout, no calibration	Serving calibrated (ModelOpt-style) checkpoints	4x smaller expert weights at BF16 activation accuracy

w_scale is the per-expert weight global scale that prepare_routed_nvfp4_weights already computes. a1_scale / a2_scale are the static FC1 / FC2 input scales.

Per-tensor NVFP4 x NVFP4 ('per_tensor')

1. Config: add 'per_tensor' to the Literal, plus a way to supply the calibrated FC1 and FC2 input scales per layer (from a ModelOpt checkpoint's input_scale, or a calibration pass).
2. flashinfer_nvfp4.py:
   • Make quant_mode an argument of _get_cute_dsl_wrapper and part of its cache key; it stays w4a4 here.
   • Add an input path that calls nvfp4_quantize(padded_hidden, a1_scale) without per_token_activation, and pass per_token_scale=None.
   • In prepare_routed_nvfp4_weights, fold the activation scales into fc1_alpha / fc2_alpha and store a2_scale in fc2_input_scale. The packed weights and block scales don't change.
3. Graph safety: keep a1_scale / a2_scale in fixed device tensors, and confirm the tensor-scale quantizer path triggers no host sync under CUDA graphs.
4. Tests: add 'per_tensor' to the recipe parametrization. The quantizer reference test gains a per-tensor case. The train-gen parity test doesn't apply: TE's per-tensor NVFP4 computes the activation scale from each tensor's amax at runtime, so it has no static-scale counterpart.

BF16 x NVFP4 ('weight_only')

1. Config: add 'weight_only' to the Literal. No calibration and no new inputs are needed.
2. flashinfer_nvfp4.py:
   • Create the wrapper with quant_mode="w4a16".
   • Skip activation quantization: pad the BF16 hidden states to padded_hidden_size and pass them as x, with x_sf, fc2_input_scale and per_token_scale all None.
   • Weights, fc1_alpha, fc2_alpha and the refresh don't change.
3. Tests: add 'weight_only' to the recipe parametrization. The quantizer reference test doesn't apply because there's no activation quantizer. The BF16 reference tolerance can be tighter, since only weights are quantized.

Both variants inherit the rest unchanged: batch invariance through the two-stage finalize (w4a16 supports it too), CUDA-graph capture, NVLS, refit refresh, and support for squared-ReLU and SwiGLU experts.

What is tested

The tests mirror the torch MXFP8 coverage from #7302 (batch invariance) and #8195 (train-gen parity) and sit next to those tests. GPU tests skip unless FlashInfer >= 0.7.0 and SM100/SM103 are available; kernel tests are marked launch_on_gb200. Results below are from a GB300 4-GPU node (torch 2.11, FlashInfer 0.7.0, TE 2.18).

Run all commands from the repository root. Every NVFP4 kernel, layer, batch-invariance, NVLS and parity test is parametrized over squared-ReLU and SwiGLU.

1. Correctness

tests/unit_tests/inference/test_flashinfer_nvfp4.py

• Quantizer bitwise reference (test_per_token_nvfp4_quantize_matches_reference): FlashInfer's per-token quantizer matches a pure-torch reference of the recipe exactly (per-token scales, block-scale bytes and E2M1 codes), including an all-zero row and logspace-spread magnitudes. The reference mirrors TE's NVFP4 quantizer. It runs with FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1; production keeps FlashInfer's default fast math.
• Gated weight layout (test_gated_fc1_is_padded_per_half_and_interleaved_up_gate, CPU): the SwiGLU FC1 reordering matches FlashInfer's interleaved [up, gate] layout exactly, including per-half padding.
• MoE against BF16 (test_flashinfer_routed_nvfp4_moe_matches_bf16_reference): the fused MoE stays within the expected NVFP4 error of a BF16 squared-ReLU or SwiGLU MoE (per-token cosine mean > 0.97, minimum > 0.95), including non-256-aligned shapes that exercise padding. Like MXFP8, the MoE itself is checked against BF16, not an emulated reference.
• Layer level (TestNVFP4MoELayer, NCCL and NVLS dispatchers): an inference_optimized MoE layer matches the BF16 training layer; refreshing after a BF16 weight update tracks the update; CUDA-graph replay follows new inputs and refreshed weights.
• Config, dispatch and refresh plumbing: accepted and rejected configurations, the missing-capability error message, in-place requantization, and dispatch to the NVFP4 path.

uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q \
  tests/unit_tests/inference/test_flashinfer_nvfp4.py

Result: 57 passed (squared ReLU and SwiGLU).

2. CUDA-graph capture and bitwise replay

tests/unit_tests/determinism/kernels/test_inference_kernels.py (next to test_mxfp8_swiglu_moe_replays): test_per_token_nvfp4_quantize_replays and test_flashinfer_nvfp4_moe_replays.

tests/unit_tests/determinism/kernels/test_moe_kernels.py: TestMoEModules::test_inference_nvfp4_moe_layer_replays replays the whole MoE layer.

The new kernels are registered in tests/unit_tests/determinism/kernels/manifest.py as inference_flashinfer_nvfp4_moe, and test_manifest.py checks the registration.

uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q \
  tests/unit_tests/determinism/kernels/test_inference_kernels.py -k "nvfp4 or mxfp8"
uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q \
  tests/unit_tests/determinism/kernels/test_moe_kernels.py -k nvfp4
uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q \
  tests/unit_tests/determinism/kernels/test_manifest.py

Results: 9 passed (NVFP4 together with the existing MXFP8 cases), 2 passed, 269 passed.

3. Batch invariance

• Kernel (test_inference_kernels.py::test_flashinfer_nvfp4_moe_is_batch_invariant, next to test_torch_mxfp8_moe_is_batch_invariant): a target batch gives bitwise identical output alone and inside batches of 193 and 1024 tokens, where the co-batched tokens are 100x larger and routed to the same experts.
• Local EP (test_flashinfer_nvfp4.py::test_flashinfer_routed_nvfp4_is_batch_invariant): the same check across local expert offsets, in deterministic and batch-invariant modes.
• NVLS layer (test_moe_dispatching_and_routing.py::TestNVLSAllGatherVDispatcher::test_cuda_graph_batch_invariant_nvfp4_moe_layer_uses_ordered_rsv, next to the MXFP8 equivalent): with EP = 4 the layer writes into the reduce-scatter buffer, reduces with ordered_reduce_scatter_v, and CUDA-graph replay equals eager execution bitwise.

These run with the commands in sections 1 and 2, plus:

uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q \
  tests/unit_tests/inference/test_moe_dispatching_and_routing.py \
  -k "nvfp4 or test_batch_invariant_moe_matches_training"

4. Train-gen parity (expected failure for now)

test_moe_dispatching_and_routing.py::TestNVLSAllGatherVDispatcher::test_batch_invariant_nvfp4_moe_matches_training, next to the MXFP8 test_batch_invariant_moe_matches_training from #8195. It loads the same weights into a TE training MoE layer running NVFP4 with row-scaled activations (fp4="e2m1" with the recipe from the Background section) and into the batch-invariant inference layer, then compares outputs.

TE does not match bitwise yet. The test asserts that generation is much closer to TE NVFP4 training than to BF16 (relative error less than half of the error against BF16), then calls pytest.xfail when the outputs are not bitwise equal. It measures about 0.03 relative error against TE NVFP4 training versus about 0.20 against BF16 for squared ReLU (maximum absolute difference about 0.0018), and about 0.035 versus 0.23 for SwiGLU (about 0.0008). The test reports XFAIL, which does not fail CI, and will start passing once the kernels are aligned. The known gaps are:

• Training applies the routing probabilities before quantizing the FC2 input; the FlashInfer kernel applies them in the finalize, after FC2.
• The CuTe DSL kernel quantizes the FC2 input inside the FC1 epilogue with approximate reciprocals. FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH only affects the standalone quantizer, so this step can't be made IEEE-exact today.
• The GEMMs differ in how the FC1 output is rounded and in accumulation order.

Command: the test_moe_dispatching_and_routing.py command in section 3.

Result: 6 passed, 4 skipped, 2 xfailed (the parity test for squared ReLU and SwiGLU). The skips are pre-existing MXFP8 parity cases that need DeepGEMM BF16 or cuBLASLt 13.5.1.

5. Regressions

uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q \
  tests/unit_tests/inference/test_flashinfer_mxfp8.py
uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q \
  tests/unit_tests/inference/test_moe_dispatching_and_routing.py

Results: 28 passed for test_flashinfer_mxfp8.py; 61 passed, 5 skipped and 2 xfailed (the parity test above) for the dispatching file. The skips are pre-existing torch and MXFP8 cases that need DeepGEMM BF16 or cuBLASLt 13.5.1.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[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.

[YigongQin]

/ok to test 466c868

[YigongQin]

the only CI failure in 466c868 was addressed in b0a26ec

[zhongbozhu]

What's the version in CI?

[YigongQin]

CI will use the updated version

[zhongbozhu]

have we added launch on gb200 flag?

[YigongQin]

added gb200 flags in 2a82c00

[wdykas]

At EP>1 we quantize and run the whole max-size NVLS buffer? I think the bounded work is useful here?

[YigongQin]

trying to reuse self._flashinfer_mxfp8_token_capacity and bounded rows

[wdykas]

can we add a contiguous call for safety here. will be no op if aligned

[YigongQin]

added

[wdykas]

was stuck on this line for a while. It'ssafe because config validation blocks gated + FlashInfer without NVFP4? if so could we just put like a small note

[YigongQin]

added one comment on gated with non-nvfp4 flashinfer are rejected

[wdykas]

nice work! can we have the system error if it tries to refit something with fp4?

[YigongQin]

nice work! can we have the system error if it tries to refit something with fp4?

added not implement error in refit.py, thank you