[Kernel] Batch invariant NVFP4 MoE using cutlass (#40372)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Jakub Zakrzewski committed
Aug 5, 2026 at 06:31 UTC
c25093a30580a5e81604f36147500cfdb739e2a6
5 files changed
+316
-102
.buildkite/test_areas/misc.yaml
-1
@@ -415,7 +415,6 @@ steps:
415
- 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 -k TRITON_MLA
416
- 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 -k FLASH_ATTN
417
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
418
- - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
418
- pytest -v -s v1/determinism/test_matmul_batch_invariant.py
419
- pytest -v -s v1/determinism/test_cutlass_batch_invariance.py
420
- pytest -v -s v1/determinism/test_online_batch_invariance.py
csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu
+18
@@ -274,6 +274,15 @@ void run_fp4_blockwise_scaled_group_mm_sm100(
274
275
using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
276
using Gemm = Gemm1SM;
277
+ static_assert(
278
+ cute::is_same_v<
279
+ typename Gemm::GemmKernel::TileScheduler,
280
+ cutlass::gemm::kernel::detail::PersistentTileSchedulerSm100Group<
281
+ ProblemShape, Gemm::GemmKernel::SchedulerPipelineStageCount>>,
282
+ "SM100 NVFP4 grouped GEMM must use PersistentTileSchedulerSm100Group "
283
+ "for batch invariance (VLLM_BATCH_INVARIANT=1). "
284
+ "If you want to change this, add a dedicated config/code path "
285
+ "for batch invariant mode, instead of relaxing this check.");
286
using StrideA = typename Gemm::GemmKernel::InternalStrideA;
287
using StrideB = typename Gemm::GemmKernel::InternalStrideB;
288
using StrideC = typename Gemm::GemmKernel::InternalStrideC;
@@ -473,6 +482,15 @@ void run_fp4_blockwise_scaled_group_mm_sm120(
482
CollectiveEpilogue>;
483
484
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
485
+ static_assert(
486
+ cute::is_base_of_v<
487
+ cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative,
488
+ typename Gemm::GemmKernel::CollectiveMainloop::DispatchPolicy::
489
+ Schedule>,
490
+ "SM120 NVFP4 grouped GEMM must use a cooperative ptr-array mainloop "
491
+ "schedule for batch invariance (VLLM_BATCH_INVARIANT=1). "
492
+ "If you want to change this, add a dedicated config/code path "
493
+ "for batch invariant mode, instead of relaxing this check.");
494
using StrideA = typename Gemm::GemmKernel::InternalStrideA;
495
using StrideB = typename Gemm::GemmKernel::InternalStrideB;
496
using StrideC = typename Gemm::GemmKernel::InternalStrideC;
tests/v1/determinism/test_cutlass_batch_invariance.py
+294
@@ -1,22 +1,89 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CUTLASS batch-invariance tests.
4
+
5
+The NVFP4 CUTLASS tests in this file must run with ``VLLM_BATCH_INVARIANT=1``
6
+before the first relevant native kernel call. Do not run them in the same pytest
7
+process after tests that intentionally exercise non-batch-invariant NVFP4
8
+scaled-mm kernels.
9
+"""
10
+
11
+from typing import Any
12
13
import pytest
14
import torch
15
16
import vllm.envs as envs
17
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
18
+from tests.kernels.moe.utils import make_dummy_moe_config, make_test_quant_config
19
+from tests.kernels.quantization.nvfp4_utils import get_nvfp4_global_scale
20
from tests.utils import TestFP8Layer, requires_fp8
21
+from vllm import _custom_ops as ops
22
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
23
CutlassFP8ScaledMMLinearKernel,
24
)
25
+from vllm.model_executor.layers.fused_moe import fused_topk
26
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
27
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
28
+ maybe_make_prepare_finalize,
29
+)
30
+from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
31
+ CutlassExpertsFp4,
32
+)
33
from vllm.model_executor.layers.quantization.utils.quant_utils import (
34
kFp8DynamicTokenSym,
35
kFp8StaticTensorSym,
36
)
37
from vllm.platforms import current_platform
38
+from vllm.utils.torch_utils import set_random_seed
39
40
pytest.importorskip("torch.cuda")
41
42
+_NVFP4_REQUIRES_SM100 = pytest.mark.skipif(
43
+ not current_platform.has_device_capability(100),
44
+ reason="Nvfp4 Requires compute capability of 10 or above.",
45
+)
46
+
47
+_NVFP4_LINEAR_DTYPES = [torch.float16, torch.bfloat16]
48
+_NVFP4_LINEAR_CONSISTENCY_SHAPES = [
49
+ (256, 128, 4096),
50
+ (512, 256, 4096),
51
+ (256, 256, 2048),
52
+ (241, 160, 2048),
53
+ (401, 352, 1984),
54
+ (333, 320, 1008),
55
+ (287, 96, 4096),
56
+]
57
+
58
+_NVFP4_MOE_DTYPES = (torch.bfloat16,)
59
+_NVFP4_MOE_NUM_EXPERTS = (40, 64)
60
+_NVFP4_MOE_TOPKS = (1, 4)
61
+
62
+_NVFP4_MOE_BATCH_INVARIANT_CASES = (
63
+ {
64
+ "m": 12,
65
+ "n": 1024,
66
+ "k": 1024,
67
+ "subset_indices": ((0, 3, 7), (1, 4, 8, 10), (11, 5, 2, 9, 6)),
68
+ },
69
+ {
70
+ "m": 129,
71
+ "n": 1024,
72
+ "k": 1024,
73
+ "subset_indices": (
74
+ (0, 64, 128),
75
+ (1, 63, 64, 127),
76
+ (32, 33, 95, 96, 128),
77
+ ),
78
+ },
79
+ {
80
+ "m": 73,
81
+ "n": 1472,
82
+ "k": 1536,
83
+ "subset_indices": ((0, 17, 35), (8, 36, 72), (1, 9, 33, 48, 64)),
84
+ },
85
+)
86
+
87
88
@pytest.fixture(autouse=True)
89
def setup_cuda():
@@ -66,3 +133,230 @@ def test_cutlass_fp8_batch_invariant_fixed_config(
133
134
torch.testing.assert_close(front_output, baseline, rtol=0, atol=0)
135
torch.testing.assert_close(back_output, baseline, rtol=0, atol=0)
136
+
137
+
138
+@_NVFP4_REQUIRES_SM100
139
+@pytest.mark.parametrize("dtype", _NVFP4_LINEAR_DTYPES)
140
+@pytest.mark.parametrize("shape", _NVFP4_LINEAR_CONSISTENCY_SHAPES)
141
+@torch.inference_mode()
142
+def test_cutlass_nvfp4_scaled_mm_batch_invariant(
143
+ dtype: torch.dtype,
144
+ shape: tuple[int, int, int],
145
+) -> None:
146
+ """Each row of a full-M GEMM must match its M=1 counterpart."""
147
+ set_random_seed(12345)
148
+ m, n, packed_k = shape
149
+ k = packed_k * 2 # real K (FP4 elements)
150
+
151
+ a_dtype = torch.randn((m, k), dtype=dtype, device="cuda")
152
+ b_dtype = torch.randn((n, k), dtype=dtype, device="cuda")
153
+
154
+ a_global_scale = get_nvfp4_global_scale(a_dtype)
155
+ b_global_scale = get_nvfp4_global_scale(b_dtype)
156
+ alpha = 1.0 / (a_global_scale * b_global_scale)
157
+
158
+ b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(b_dtype, b_global_scale)
159
+
160
+ a_fp4_full, a_sf_full = ops.scaled_fp4_quant(a_dtype, a_global_scale)
161
+ out_full = ops.cutlass_scaled_fp4_mm(
162
+ a_fp4_full,
163
+ b_fp4,
164
+ a_sf_full,
165
+ b_scale_interleaved,
166
+ alpha,
167
+ dtype,
168
+ )
169
+
170
+ for i in range(m):
171
+ a_row = a_dtype[i : i + 1]
172
+ a_fp4_row, a_sf_row = ops.scaled_fp4_quant(a_row, a_global_scale)
173
+ out_row = ops.cutlass_scaled_fp4_mm(
174
+ a_fp4_row,
175
+ b_fp4,
176
+ a_sf_row,
177
+ b_scale_interleaved,
178
+ alpha,
179
+ dtype,
180
+ )
181
+
182
+ assert torch.equal(out_full[i], out_row[0]), (
183
+ f"VLLM_BATCH_INVARIANT: row {i} differs between M={m} and M=1: "
184
+ f"max_abs_diff={(out_full[i] - out_row[0]).abs().max().item()}"
185
+ )
186
+
187
+
188
+def _make_cutlass_fp4_moe_batch_invariant_case(
189
+ case_config: dict[str, Any],
190
+ activation: MoEActivation,
191
+ e: int,
192
+ topk: int,
193
+ dtype: torch.dtype,
194
+) -> dict[str, Any]:
195
+ set_random_seed(7)
196
+ assert topk <= e
197
+
198
+ hidden_states = (
199
+ torch.randn((case_config["m"], case_config["k"]), device="cuda", dtype=dtype)
200
+ / 10
201
+ )
202
+
203
+ w1_q, w2_q, quant_config = make_test_quant_config(
204
+ e,
205
+ case_config["n"],
206
+ case_config["k"],
207
+ in_dtype=dtype,
208
+ quant_dtype="nvfp4",
209
+ block_shape=None,
210
+ per_act_token_quant=False,
211
+ make_gate=activation.is_gated,
212
+ )
213
+
214
+ score = torch.randn((case_config["m"], e), device="cuda", dtype=dtype)
215
+
216
+ moe_config = make_dummy_moe_config(
217
+ num_experts=e,
218
+ experts_per_token=topk,
219
+ hidden_dim=case_config["k"],
220
+ intermediate_size=case_config["n"],
221
+ in_dtype=dtype,
222
+ activation=activation,
223
+ )
224
+ kernel = mk.FusedMoEKernel(
225
+ maybe_make_prepare_finalize(
226
+ moe=moe_config,
227
+ quant_config=quant_config,
228
+ allow_new_interface=True,
229
+ use_monolithic=False,
230
+ ),
231
+ CutlassExpertsFp4(
232
+ moe_config=moe_config,
233
+ quant_config=quant_config,
234
+ ),
235
+ )
236
+
237
+ return {
238
+ "config": case_config,
239
+ "hidden_states": hidden_states,
240
+ "score": score,
241
+ "kernel": kernel,
242
+ "w1_q": w1_q,
243
+ "w2_q": w2_q,
244
+ "activation": activation,
245
+ "e": e,
246
+ "topk": topk,
247
+ "dtype": dtype,
248
+ }
249
+
250
+
251
+def _run_cutlass_fp4_moe(
252
+ case: dict[str, Any],
253
+ hidden_states: torch.Tensor,
254
+ score: torch.Tensor,
255
+) -> torch.Tensor:
256
+ topk_weights, topk_ids, _ = fused_topk(
257
+ hidden_states, score, case["topk"], renormalize=False
258
+ )
259
+ return case["kernel"].apply(
260
+ hidden_states=hidden_states,
261
+ w1=case["w1_q"],
262
+ w2=case["w2_q"],
263
+ topk_weights=topk_weights,
264
+ topk_ids=topk_ids,
265
+ global_num_experts=case["e"],
266
+ activation=case["activation"],
267
+ apply_router_weight_on_input=False,
268
+ expert_map=None,
269
+ )
270
+
271
+
272
+@_NVFP4_REQUIRES_SM100
273
+@pytest.mark.parametrize("case_config", _NVFP4_MOE_BATCH_INVARIANT_CASES)
274
+@pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.SWIGLUSTEP])
275
+@pytest.mark.parametrize("e", _NVFP4_MOE_NUM_EXPERTS)
276
+@pytest.mark.parametrize("topk", _NVFP4_MOE_TOPKS)
277
+@pytest.mark.parametrize("dtype", _NVFP4_MOE_DTYPES)
278
+@torch.inference_mode()
279
+def test_cutlass_nvfp4_moe_batch_invariant(
280
+ case_config: dict[str, Any],
281
+ activation: MoEActivation,
282
+ e: int,
283
+ topk: int,
284
+ dtype: torch.dtype,
285
+ default_vllm_config,
286
+ workspace_init,
287
+) -> None:
288
+ case = _make_cutlass_fp4_moe_batch_invariant_case(
289
+ case_config, activation, e, topk, dtype
290
+ )
291
+ case_id = f"{case['config']}, e={case['e']}, topk={case['topk']}"
292
+ # Establish the baseline output for the full mixed batch once, then
293
+ # compare every smaller replay against the corresponding slice here.
294
+ batch_output = _run_cutlass_fp4_moe(case, case["hidden_states"], case["score"])
295
+
296
+ assert CutlassExpertsFp4._supports_batch_invariance()
297
+
298
+ # Re-run the whole batch with a different batch-row permutation to make
299
+ # sure the grouped GEMM result is invariant to how tokens are packed into
300
+ # expert work for scheduling.
301
+ indices = torch.arange(
302
+ case["hidden_states"].size(0), device=case["hidden_states"].device
303
+ )
304
+ for perm_name, perm in (
305
+ ("reversed", torch.flip(indices, dims=(0,))),
306
+ ("evens_then_odds", torch.cat((indices[::2], indices[1::2]))),
307
+ ):
308
+ permuted_output = _run_cutlass_fp4_moe(
309
+ case,
310
+ case["hidden_states"][perm],
311
+ case["score"][perm],
312
+ )
313
+ torch.testing.assert_close(
314
+ batch_output[perm],
315
+ permuted_output,
316
+ atol=0,
317
+ rtol=0,
318
+ msg=(
319
+ f"{case_id}: permutation '{perm_name}' changed outputs "
320
+ "relative to the baseline batch order."
321
+ ),
322
+ )
323
+
324
+ # Re-run every batch-row as a batch-size-1 input and compare against the
325
+ # matching row from the full batch.
326
+ for idx in range(case["hidden_states"].size(0)):
327
+ single_idx = torch.tensor([idx], device=case["hidden_states"].device)
328
+ single_output = _run_cutlass_fp4_moe(
329
+ case,
330
+ case["hidden_states"][single_idx],
331
+ case["score"][single_idx],
332
+ )
333
+ torch.testing.assert_close(
334
+ batch_output[single_idx],
335
+ single_output,
336
+ atol=0,
337
+ rtol=0,
338
+ msg=(
339
+ f"{case_id}: batch-row {idx} changed between full-batch "
340
+ "and batch-size-1 execution."
341
+ ),
342
+ )
343
+
344
+ # Re-run a few nontrivial sub-batches to catch interactions that only
345
+ # appear when multiple tokens are grouped together.
346
+ for subset_ids in case["config"]["subset_indices"]:
347
+ subset = torch.tensor(subset_ids, device=case["hidden_states"].device)
348
+ subset_output = _run_cutlass_fp4_moe(
349
+ case,
350
+ case["hidden_states"][subset],
351
+ case["score"][subset],
352
+ )
353
+ torch.testing.assert_close(
354
+ batch_output[subset],
355
+ subset_output,
356
+ atol=0,
357
+ rtol=0,
358
+ msg=(
359
+ f"{case_id}: sub-batch {list(subset_ids)} changed "
360
+ "relative to full-batch execution."
361
+ ),
362
+ )
tests/v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
deleted
-101
@@ -1,101 +0,0 @@
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/layers/fused_moe/experts/cutlass_moe.py
+4
@@ -734,6 +734,10 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
734
# needed for STANDARD activation format kernels in EP mode.
735
return moe_parallel_config.ep_size == 1
736
737
+ @staticmethod
738
+ def _supports_batch_invariance() -> bool:
739
+ return True
740
+
741
@staticmethod
742
def activation_format() -> mk.FusedMoEActivationFormat:
743
return mk.FusedMoEActivationFormat.Standard