Repository navigation
[Kernel] Fuse native BF16 ReLU2 with NVFP4 quantization - #59384
samnordmann wants to merge 1 commit into
Conversation
mergify
Bot
added
performance
9f04ae8 to
597a94c
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
d3aa128 to
f81138a
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
Fuse the native BF16 activation and NVFP4 producer without changing the CuTe-DSL GEMM. Preserve rounding, padded scales, and unsupported-input fallbacks. Keep dynamic-row launch selection outside guard-free compiled graphs and restrict automatic fusion to the validated SM10x CuTe path. Add startup warmup and targeted regression coverage. Co-authored-by: Codex Signed-off-by: samnordmann <snordmann@nvidia.com>
f81138a to
462744b
Compare
| def register_relu_squared_nvfp4_quant_warmup(module: torch.nn.Module) -> None: | ||
| """Defer a Nemotron MLP's producer selection until startup kernel warmup.""" | ||
| config = get_current_vllm_config_or_none() | ||
| if config is None or not config.kernel_config.enable_jit_warmup: | ||
| return | ||
| if ( | ||
| _FUSED_ACT_QUANT.get((ReLUSquaredActivation, kNvfp4Dynamic)) | ||
| is not relu_squared_nvfp4_quant | ||
| ): | ||
| return | ||
| max_tokens = max( | ||
| config.scheduler_config.max_num_batched_tokens, | ||
| config.compilation_config.max_cudagraph_capture_size or 0, | ||
| ) | ||
| _relu_squared_nvfp4_quant_kernel.register_warmup( | ||
| module=module, max_tokens=max_tokens | ||
| ) |
There was a problem hiding this comment.
This should not be needed imo. We do a first forward pass through the model with max_num_batched_tokens as the memory profiling run for kv cache allocation, so you should be able to spin out shape warmup on the first invocation
There was a problem hiding this comment.
You're right that the memory-profile forward runs before capture, so a first-call routine could warm the other shapes. I kept this as metadata-only registration to follow the current JIT warmup contract: materialize registered kernels through kernel_warmup() and keep runtime execution separate from startup compilation. This function does not compile or launch anything.
The max-token launch alone is insufficient: at K=5376, M=2048/4096 uses PAD_PER_CTA=128, whereas M=65 needs 256 and M=1 needs 512. The provider covers those compile keys plus pointer-alignment variants, and resolves eligibility after weight processing/LoRA replacement. The existing regression test then rejects any further producer JIT across unseen rows and graph capture.
I've clarified this in the description and left the implementation unchanged. Do you prefer a first-invocation exception to that warmup contract here?
This branch has not been deployed
2 participants
Overview
Fuse native BF16 ReLU2 with NVFP4 quantization for Nemotron shared experts on SM10x/FlashInfer CuTe-DSL. Removes one launch and the BF16 intermediate; reuses the consumer from #59612, with no GEMM changes.
Claims
Validation
Revision
462744b0, GB200, Nemotron shared-expert dimensions (K=5376, N=4096). Ordinary native compiler settings, warm CUDA graphs; eight paired samples per shape:All pairs improve; paired component speedups are 1.147–1.151× at 2K and 1.246–1.258× at 4K. Nsight confirms 3 → 2 kernels with identical GEMM kernel and grid. These are component gains, not full-model speedups.
Correctness: 325 GPU checks and pre-commit/mypy pass. Full-model oracle: 707,680 fused calls, zero value/scale/padding errors, no post-startup compilation. GSM8K: 90.75% in the current instrumented run; prior controlled A/B: 91.66% in all four runs, identical tokens (not a new-stack paired comparison).
Serving: prior-stack C50, 2K-input/1K-output median throughput is 2,359 → 2,375 output tokens/s (+0.66%, six trials/arm), with overlapping ranges—not a robust serving-speedup claim.
Prefill-heavy serving validation remains incomplete: C50, 8K-input/128-output, 2K/4K token budgets. Initial cohorts exceeded the harness deadline; full-protocol retries are queued. No speedup is claimed from partial trials.
Reproduction and scope
Current measurements use
462744b0, FlashInfer 0.7.0.post1/CuTe DSL 4.8. Run the component benchmark against this checkout, with a fresh cache and largest shape first:uv run --no-sync -- .venv/bin/python benchmark_relu2_nvfp4_quant.py \ --rows 4096 2048 128 1 --widths 5376 --output-width 4096 \ --rounds 8 --output benchmark.json uv run --no-sync -- .venv/bin/python -m pytest -q \ tests/fusion/test_quant_activation_contract.py \ tests/kernels/test_fused_quant_activation.py uv run --no-sync -- .venv/bin/python -m pytest -q \ tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py \ -k 'test_cutedsl_relu2_prequantized_bypasses_quantizer or test_flashinfer_nvfp4_preserves_quantized_input_shape'Cold-cache speedups at 2K/4K are 1.18×/1.22×; the earlier broader 20-shape sweep averaged 1.08× warm.
Prior serving/GSM8K used
f81138a8with fixed compiler choices; producer unchanged. Larger gains with pointwise tuning disabled are excluded above. GPU validation was external; hosted CI needs contributor authorization.Details
Reuses QuantizedActivation; opaque dispatch preserves dynamic rows under guard-free Dynamo. Warmup registration is metadata-only: the shared
kernel_warmup()compiles padding/alignment variants after consumer selection is resolved post-load, before CUDA graph capture, following the startup warmup contract. Only the validated native-ReLU2/SM10x/CuTe path is enabled. No scheduler, runner, communication or dependency changes.Distinct from routed-MoE #56740, GELU #59467 and normalization #61073. No DISAGG, TP>1 or default-stack serving claim.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)