A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 22:00 UTC
[Kernel] Batch invariant NVFP4 linear using cutlass (#39912)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Jakub Zakrzewski committed
May 23, 2026 at 15:41 UTC
5bb8d2767a2829b56e58c68fa8f380e9e4e2bd3e
5 files changed
+210
-49
.buildkite/test_areas/misc.yaml
+1
@@ -367,6 +367,7 @@ steps:
367
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
368
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
369
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
370
+ - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
371
372
- label: Acceptance Length Test (Large Models) # optional
373
device: h200_35gb
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu
+29
-5
@@ -22,6 +22,8 @@
22
23
#include "cutlass/cutlass.h"
24
25
+#include <type_traits>
26
+
27
#include "cutlass/gemm/collective/collective_builder.hpp"
28
#include "cutlass/epilogue/collective/collective_builder.hpp"
29
#include "cutlass/gemm/device/gemm_universal_adapter.h"
@@ -30,15 +32,21 @@
32
#include "cutlass/util/packed_stride.hpp"
33
34
#include "core/math.hpp"
35
+#include "core/batch_invariant.hpp"
36
37
using namespace cute;
38
39
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
40
38
-// Configuration for M in (256, inf)
41
+// Configuration for M in (256, inf), also reused for batch-invariant mode
42
+// to keep a fixed large-M tiling across all batch sizes.
43
+// Do not change the tile K or tile scheduler here unless you are also
44
+// updating the batch-invariant behavior; if batch-invariant mode needs a
45
+// different schedule, add a dedicated batch-invariant config/path instead.
46
struct sm100_fp4_config_default {
47
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
48
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
49
+ using TileScheduler = cutlass::gemm::PersistentScheduler;
50
using TileShape = Shape<_256, _256, _256>;
51
using ClusterShape = Shape<_2, _1, _1>;
52
using PerSmTileShape_MNK = Shape<_128, _256, _256>;
@@ -48,6 +56,7 @@ struct sm100_fp4_config_default {
56
struct sm100_fp4_config_M256 {
57
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
58
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
59
+ using TileScheduler = void;
60
using TileShape = Shape<_256, _128, _256>;
61
using ClusterShape = Shape<_2, _1, _1>;
62
using PerSmTileShape_MNK = Shape<_128, _128, _256>;
@@ -57,6 +66,7 @@ struct sm100_fp4_config_M256 {
66
struct sm100_fp4_config_M16 {
67
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
68
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
69
+ using TileScheduler = void;
70
using TileShape = Shape<_128, _128, _256>;
71
using ClusterShape = Shape<_1, _1, _1>;
72
using PerSmTileShape_MNK = Shape<_128, _128, _256>;
@@ -98,7 +108,7 @@ struct Fp4GemmSm100 {
108
cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
109
ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD,
110
LayoutDTag, AlignmentD,
101
- cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
111
+ typename Config::EpilogueSchedule>::CollectiveOp;
112
113
using CollectiveMainloop =
114
typename cutlass::gemm::collective::CollectiveBuilder<
@@ -107,10 +117,13 @@ struct Fp4GemmSm100 {
117
ClusterShape,
118
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
119
sizeof(typename CollectiveEpilogue::SharedStorage))>,
110
- cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
120
+ typename Config::KernelSchedule>::CollectiveOp;
121
112
- using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
113
- Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
122
+ using TileScheduler = typename Config::TileScheduler;
123
+ using GemmKernel =
124
+ cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>,
125
+ CollectiveMainloop,
126
+ CollectiveEpilogue, TileScheduler>;
127
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
128
using StrideA = typename Gemm::GemmKernel::StrideA;
129
using LayoutA = decltype(cute::make_layout(make_shape(0, 0, 0), StrideA{}));
@@ -205,6 +218,17 @@ void cutlass_fp4_gemm_dispatch(torch::stable::Tensor& D,
218
torch::stable::Tensor const& B_sf,
219
torch::stable::Tensor const& alpha, int64_t m,
220
int64_t n, int64_t k, cudaStream_t stream) {
221
+ if (vllm::vllm_is_batch_invariant()) {
222
+ using BiGemm = Fp4GemmSm100<sm100_fp4_config_default, OutType>;
223
+ static_assert(
224
+ cute::is_same_v<typename BiGemm::TileScheduler,
225
+ cutlass::gemm::PersistentScheduler>,
226
+ "batch_invariant requires a persistent tile scheduler; stream-K or "
227
+ "split-K would break numerical invariance");
228
+ runGemm<BiGemm>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
229
+ return;
230
+ }
231
+
232
uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
233
234
if (mp2 <= 16) {
csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu
+50
-32
@@ -22,6 +22,8 @@
22
23
#include "cutlass/cutlass.h"
24
25
+#include <type_traits>
26
+
27
#include "cutlass/gemm/collective/collective_builder.hpp"
28
#include "cutlass/epilogue/collective/collective_builder.hpp"
29
#include "cutlass/gemm/device/gemm_universal_adapter.h"
@@ -30,6 +32,7 @@
32
#include "cutlass/util/packed_stride.hpp"
33
34
#include "core/math.hpp"
35
+#include "core/batch_invariant.hpp"
36
37
using namespace cute;
38
@@ -49,12 +52,22 @@ constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
52
constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn;
53
54
struct sm120_fp4_config_M256 {
55
+ using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
56
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
57
+ using TileScheduler = void;
58
using ClusterShape = Shape<_1, _1, _1>;
59
using MmaTileShape = Shape<_128, _128, _128>;
60
using PerSmTileShape_MNK = Shape<_128, _128, _128>;
61
};
62
63
struct sm120_fp4_config_default {
64
+ // Also used for batch-invariant mode.
65
+ // Do not change the tile K or tile scheduler here unless you are also
66
+ // updating the batch-invariant behavior; if batch-invariant mode needs a
67
+ // different schedule, add a dedicated batch-invariant config/path instead.
68
+ using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
69
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
70
+ using TileScheduler = cutlass::gemm::PersistentScheduler;
71
using ClusterShape = Shape<_1, _1, _1>;
72
using MmaTileShape = Shape<_256, _128, _128>;
73
using PerSmTileShape_MNK = Shape<_256, _128, _128>;
@@ -91,7 +104,7 @@ struct Fp4GemmSm120 {
104
cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
105
ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD,
106
LayoutDTag, AlignmentD,
94
- cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
107
+ typename Config::EpilogueSchedule>::CollectiveOp;
108
109
using CollectiveMainloop =
110
typename cutlass::gemm::collective::CollectiveBuilder<
@@ -100,10 +113,13 @@ struct Fp4GemmSm120 {
113
ClusterShape,
114
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
115
sizeof(typename CollectiveEpilogue::SharedStorage))>,
103
- cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
116
+ typename Config::KernelSchedule>::CollectiveOp;
117
105
- using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
106
- Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
118
+ using TileScheduler = typename Config::TileScheduler;
119
+ using GemmKernel =
120
+ cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>,
121
+ CollectiveMainloop,
122
+ CollectiveEpilogue, TileScheduler>;
123
124
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
125
};
@@ -180,40 +196,42 @@ void runGemm(torch::stable::Tensor& D, torch::stable::Tensor const& A,
196
CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
197
}
198
183
-void cutlass_fp4_bf16_gemm_dispatch(torch::stable::Tensor& D,
184
- torch::stable::Tensor const& A,
185
- torch::stable::Tensor const& B,
186
- torch::stable::Tensor const& A_sf,
187
- torch::stable::Tensor const& B_sf,
188
- torch::stable::Tensor const& alpha, int m,
189
- int n, int k, cudaStream_t stream) {
190
- uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
191
- if (mp2 <= 256) {
192
- runGemm<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::bfloat16_t>::Gemm>(
193
- D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
194
- } else {
195
- runGemm<Fp4GemmSm120<sm120_fp4_config_default, cutlass::bfloat16_t>::Gemm>(
196
- D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
199
+namespace {
200
+
201
+// Dispatch function to select appropriate config based on M (file-local;
202
+// internal linkage avoids clashing with SM100's cutlass_fp4_gemm_dispatch in
203
+// nvfp4_scaled_mm_kernels.cu).
204
+template <typename OutType>
205
+void cutlass_fp4_gemm_dispatch(torch::stable::Tensor& D,
206
+ torch::stable::Tensor const& A,
207
+ torch::stable::Tensor const& B,
208
+ torch::stable::Tensor const& A_sf,
209
+ torch::stable::Tensor const& B_sf,
210
+ torch::stable::Tensor const& alpha, int m, int n,
211
+ int k, cudaStream_t stream) {
212
+ if (vllm::vllm_is_batch_invariant()) {
213
+ using BiGemm = Fp4GemmSm120<sm120_fp4_config_default, OutType>;
214
+ static_assert(
215
+ cute::is_same_v<typename BiGemm::TileScheduler,
216
+ cutlass::gemm::PersistentScheduler>,
217
+ "batch_invariant requires a persistent tile scheduler; stream-K or "
218
+ "split-K would break numerical invariance");
219
+ runGemm<typename BiGemm::Gemm>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
220
+ return;
221
}
198
-}
222
200
-void cutlass_fp4_f16_gemm_dispatch(torch::stable::Tensor& D,
201
- torch::stable::Tensor const& A,
202
- torch::stable::Tensor const& B,
203
- torch::stable::Tensor const& A_sf,
204
- torch::stable::Tensor const& B_sf,
205
- torch::stable::Tensor const& alpha, int m,
206
- int n, int k, cudaStream_t stream) {
223
uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
224
if (mp2 <= 256) {
209
- runGemm<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::half_t>::Gemm>(
225
+ runGemm<typename Fp4GemmSm120<sm120_fp4_config_M256, OutType>::Gemm>(
226
D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
227
} else {
212
- runGemm<Fp4GemmSm120<sm120_fp4_config_default, cutlass::half_t>::Gemm>(
228
+ runGemm<typename Fp4GemmSm120<sm120_fp4_config_default, OutType>::Gemm>(
229
D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
230
}
231
}
232
233
+} // namespace
234
+
235
void cutlass_scaled_fp4_mm_sm120a(torch::stable::Tensor& D,
236
torch::stable::Tensor const& A,
237
torch::stable::Tensor const& B,
@@ -275,11 +293,11 @@ void cutlass_scaled_fp4_mm_sm120a(torch::stable::Tensor& D,
293
const cudaStream_t stream = get_current_cuda_stream(A.get_device_index());
294
295
if (out_dtype == torch::headeronly::ScalarType::BFloat16) {
278
- return cutlass_fp4_bf16_gemm_dispatch(D, A, B, A_sf, B_sf, alpha, m, n, k,
279
- stream);
296
+ return cutlass_fp4_gemm_dispatch<cutlass::bfloat16_t>(
297
+ D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
298
} else if (out_dtype == torch::headeronly::ScalarType::Half) {
281
- return cutlass_fp4_f16_gemm_dispatch(D, A, B, A_sf, B_sf, alpha, m, n, k,
282
- stream);
299
+ return cutlass_fp4_gemm_dispatch<cutlass::half_t>(D, A, B, A_sf, B_sf,
300
+ alpha, m, n, k, stream);
301
} else {
302
STD_TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm sm120 (",
303
out_dtype, ")");
tests/v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
new
+101
@@ -0,0 +1,101 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NVFP4 CUTLASS GEMM tests that require ``VLLM_BATCH_INVARIANT=1``.
4
+
5
+Must run in a **fresh** pytest process:
6
+
7
+ pytest tests/v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py -v
8
+
9
+Do not share a session with ``tests/kernels/quantization/test_nvfp4_scaled_mm.py``:
10
+the native code caches whether batch invariance is enabled on the first GEMM, and
11
+if ``VLLM_BATCH_INVARIANT`` was not set at that moment, it stays disabled for the
12
+rest of the process.
13
+"""
14
+
15
+import os
16
+
17
+import pytest
18
+import torch
19
+
20
+from tests.kernels.quantization.nvfp4_utils import get_nvfp4_global_scale
21
+from vllm import _custom_ops as ops
22
+from vllm.platforms import current_platform
23
+from vllm.utils.torch_utils import set_random_seed
24
+
25
+if not current_platform.has_device_capability(100):
26
+ pytest.skip(
27
+ reason="Nvfp4 Requires compute capability of 10 or above.",
28
+ allow_module_level=True,
29
+ )
30
+
31
+DTYPES = [torch.float16, torch.bfloat16]
32
+SHAPES = [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]
33
+PAD_SHAPES = [(150, 128, 64), (128, 128, 96)]
34
+SHAPES.extend(PAD_SHAPES)
35
+
36
+
37
+CONSISTENCY_SHAPES = [
38
+ (256, 128, 4096),
39
+ (512, 256, 4096),
40
+ (256, 256, 2048),
41
+ (241, 160, 2048),
42
+ (401, 352, 1984),
43
+ (333, 320, 1008),
44
+ (287, 96, 4096),
45
+]
46
+
47
+
48
+@pytest.mark.parametrize("dtype", DTYPES)
49
+@pytest.mark.parametrize("shape", CONSISTENCY_SHAPES)
50
+@torch.inference_mode()
51
+def test_nvfp4_gemm_batch_invariance(
52
+ dtype: torch.dtype,
53
+ shape: tuple[int, int, int],
54
+) -> None:
55
+ """Batch invariance: each row of a full-``M`` GEMM matches its ``M=1`` counterpart.
56
+
57
+ For row ``i``, compares ``cutlass_scaled_fp4_mm`` run once over all ``M``
58
+ rows against a separate call with ``A`` sliced to ``a_dtype[i : i+1]``.
59
+ Catches kernels whose reduction or scheduling depends on ``M`` or adjacent
60
+ rows.
61
+ """
62
+ seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
63
+ set_random_seed(seed)
64
+ m, n, packed_k = shape
65
+ k = packed_k * 2 # real K (FP4 elements)
66
+
67
+ a_dtype = torch.randn((m, k), dtype=dtype, device="cuda")
68
+ b_dtype = torch.randn((n, k), dtype=dtype, device="cuda")
69
+
70
+ a_global_scale = get_nvfp4_global_scale(a_dtype)
71
+ b_global_scale = get_nvfp4_global_scale(b_dtype)
72
+ alpha = 1.0 / (a_global_scale * b_global_scale)
73
+
74
+ b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(b_dtype, b_global_scale)
75
+
76
+ a_fp4_full, a_sf_full = ops.scaled_fp4_quant(a_dtype, a_global_scale)
77
+ out_full = ops.cutlass_scaled_fp4_mm(
78
+ a_fp4_full,
79
+ b_fp4,
80
+ a_sf_full,
81
+ b_scale_interleaved,
82
+ alpha,
83
+ dtype,
84
+ )
85
+
86
+ for i in range(m):
87
+ a_row = a_dtype[i : i + 1]
88
+ a_fp4_row, a_sf_row = ops.scaled_fp4_quant(a_row, a_global_scale)
89
+ out_row = ops.cutlass_scaled_fp4_mm(
90
+ a_fp4_row,
91
+ b_fp4,
92
+ a_sf_row,
93
+ b_scale_interleaved,
94
+ alpha,
95
+ dtype,
96
+ )
97
+
98
+ assert torch.equal(out_full[i], out_row[0]), (
99
+ f"VLLM_BATCH_INVARIANT: row {i} differs between M={m} and M=1: "
100
+ f"max_abs_diff={(out_full[i] - out_row[0]).abs().max().item()}"
101
+ )
vllm/model_executor/kernels/linear/__init__.py
+29
-12
@@ -833,24 +833,41 @@ def init_nvfp4_linear_kernel() -> NvFp4LinearKernel:
833
current platform."""
834
config = NvFp4LinearLayerConfig()
835
836
- # VLLM_BATCH_INVARIANT unconditionally forces emulation for deterministic
837
- # execution. It overrides both --linear-backend and the deprecated env
838
- # vars below.
836
+ # VLLM_BATCH_INVARIANT forces deterministic execution. Prefer the
837
+ # batch-invariant CUTLASS implementation when available, otherwise fall
838
+ # back to emulation. It overrides both --linear-backend and the deprecated
839
+ # env vars below.
840
force_kernel: type[NvFp4LinearKernel] | None = None
841
linear_backend = _get_linear_backend()
842
if envs.VLLM_BATCH_INVARIANT:
842
- if linear_backend not in ("auto", "emulation"):
843
- logger.warning_once(
844
- "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; using "
845
- "the emulation backend for deterministic execution.",
846
- linear_backend,
847
- )
843
+ bi_supported, reason = CutlassNvFp4LinearKernel.is_supported()
844
+ if bi_supported:
845
+ if linear_backend not in ("auto", "cutlass"):
846
+ logger.warning_once(
847
+ "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; "
848
+ "using the CUTLASS backend for deterministic execution.",
849
+ linear_backend,
850
+ )
851
+ else:
852
+ logger.info_once(
853
+ "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the "
854
+ "CUTLASS backend for deterministic execution."
855
+ )
856
+ force_kernel = CutlassNvFp4LinearKernel
857
else:
858
+ if linear_backend not in ("auto", "emulation"):
859
+ logger.warning_once(
860
+ "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; "
861
+ "using the emulation backend for deterministic execution.",
862
+ linear_backend,
863
+ )
864
logger.info_once(
850
- "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the "
851
- "emulation backend for deterministic execution."
865
+ "VLLM_BATCH_INVARIANT is set but the batch-invariant NVFP4 "
866
+ "kernel is not supported on this platform; falling back to "
867
+ "emulation for deterministic execution. Reason: %s",
868
+ reason,
869
)
853
- force_kernel = EmulationNvFp4LinearKernel
870
+ force_kernel = EmulationNvFp4LinearKernel
871
elif linear_backend == "auto":
872
# Deprecated env-var overrides — only honoured when --linear-backend
873
# is "auto". Deprecation warnings are emitted from vllm/envs.py.