[Kernel] Batch invariance: pin FlashInfer FP8 bmm and W4A16 NVFP4 kernels - #62192

Overview

Under VLLM_BATCH_INVARIANT=1, pins two GEMM paths so that a row's result does not depend on the batch:

• FlashInfer's FP8 bmm (flashinfer_scaled_fp8_mm) is pinned to the cutlass backend.
• W4A16 NVFP4 linears use the Humming kernel.

This was split from #62017 by @aoshen02. Without BI, behavior is unchanged.

Changes

• vllm/utils/flashinfer.py: flashinfer_scaled_fp8_mm passes backend="cutlass" to bmm_fp8 under BI and "auto" otherwise. With FlashInfer autotune (on by default at -O1 and above), "auto" tunes over cutlass, cuBLAS and cuDNN per M bucket, so different batch sizes can run different backends.
• vllm/model_executor/kernels/linear/__init__.py: under BI, a W4A16 NVFP4 linear selects HummingNvFp4LinearKernel, which already reads use_batch_invariant from VLLM_BATCH_INVARIANT. Before this change, the BI branch forced the W4A4-only CutlassNvFp4LinearKernel, and kernel init raised ValueError: CutlassNvFp4LinearKernel does not support W4A16. With --linear-backend set to something other than Humming, a warning is logged and the setting is overridden.
• Tests:
  • test_flashinfer_fp8_mm_batch_invariant_after_autotune in tests/v1/determinism/test_cutlass_batch_invariance.py (SM100+);
  • test_batch_invariant_w4a16_selects_humming in tests/kernels/quantization/test_nvfp4_kernel_selection.py (CPU).

The router FP32-logits branch (with the out_dtype argument of mm_batch_invariant) and the two sm100 tuned M buckets for Nemotron-H shapes were moved out of this PR. They stay in #62017 for the Nemotron-H wiring.

Validation

Hardware: GB200 (SM100), vLLM nightly image 0.32.0rc1.dev281+g49cb2bbf7, FlashInfer 0.7.2. The touched files are identical between that image and this PR's base fa75d3456d, so "base" is the image as shipped and "PR" is the image with this PR's two files overlaid.

FP8 bmm. Static per-tensor FP8 inputs. Each M slice x[:M] was compared bitwise against the first M rows of the M=2048 output, for M in {1, 2, 7, 16, 64, 128, 512, 2048}, on four (N, K) shapes. The table counts rows that differ.

Max abs diff of the differing rows is 3e-5 to 6e-5. Without autotune, "auto" falls back to its first runner (cutlass), so it is invariant; the variance only appears after tuning. As a diagnostic, cuBLAS alone and cuDNN alone were each invariant after autotune on these shapes: the variance comes from switching backends between buckets. The new test failed in 1 of 2 runs at base (the tuner's pick is timing dependent) and passed in 2 of 2 runs with the PR.

W4A16 NVFP4. A synthetic CompressedTensorsW4A4Fp4(use_a16=True) layer, with the same M sweep, on shapes 4096x4096, 2688x4096 and 4608x2688.

So the W4A16 change fixes a crash at init under BI rather than a numerical difference.

Tests.

pytest tests/kernels/quantization/test_nvfp4_kernel_selection.py
# 4 passed; the new test fails at base with the ValueError above
pytest tests/v1/determinism/test_cutlass_batch_invariance.py -k test_flashinfer_fp8_mm_batch_invariant_after_autotune
# GB200: PR 2/2 passed; base 1/2 failed

pre-commit (ruff, mypy) passes on the changed files.

Not tested: end-to-end serving of an FP8 or W4A16 NVFP4 model under BI, and throughput. Humming and pinned cutlass are chosen for invariance, not speed.