| 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(): |
| 90 | if not current_platform.is_cuda(): |
| 91 | pytest.skip("CUTLASS FP8 kernels require CUDA.") |
| 92 | torch.set_default_device("cuda") |
| 93 | |
| 94 | |
| 95 | @requires_fp8 |
| 96 | @pytest.mark.parametrize("weight_shape", [(1024, 2048), (4608, 4096)]) |
| 97 | @pytest.mark.parametrize("batch_size", [1, 16, 17, 32, 64, 65, 256, 257]) |
| 98 | @torch.inference_mode() |
| 99 | def test_cutlass_fp8_batch_invariant_fixed_config( |
| 100 | weight_shape: tuple[int, int], |
| 101 | batch_size: int, |
| 102 | default_vllm_config, |
| 103 | monkeypatch: pytest.MonkeyPatch, |
| 104 | ): |
| 105 | monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") |
| 106 | monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True) |
| 107 | |
| 108 | torch.manual_seed(0) |
| 109 | layer = TestFP8Layer( |
| 110 | weight_shape=weight_shape, |
| 111 | activation_quant_key=kFp8DynamicTokenSym, |
| 112 | weight_quant_key=kFp8StaticTensorSym, |
| 113 | input_dtype=torch.bfloat16, |
| 114 | out_dtype=torch.bfloat16, |
| 115 | device=torch.device("cuda"), |
| 116 | force_kernel=CutlassFP8ScaledMMLinearKernel, |
| 117 | ) |
| 118 | assert isinstance(layer.kernel, CutlassFP8ScaledMMLinearKernel) |
| 119 | |
| 120 | in_features = weight_shape[1] |
| 121 | needle = torch.randn((1, in_features), device="cuda", dtype=torch.bfloat16) |
| 122 | baseline = layer(needle)[0] |
| 123 | |
| 124 | filler = torch.randn( |
| 125 | (max(batch_size - 1, 0), in_features), device="cuda", dtype=torch.bfloat16 |
| 126 | ) |
| 127 | |
| 128 | front_batch = torch.cat([needle, filler], dim=0) |
| 129 | back_batch = torch.cat([filler, needle], dim=0) |
| 130 | |
| 131 | front_output = layer(front_batch)[0] |
| 132 | back_output = layer(back_batch)[-1] |
| 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 | prepare_finalize = maybe_make_prepare_finalize( |
| 225 | moe=moe_config, |
| 226 | quant_config=quant_config, |
| 227 | allow_new_interface=True, |
| 228 | use_monolithic=False, |
| 229 | ) |
| 230 | assert prepare_finalize is not None |
| 231 | kernel = mk.FusedMoEKernel( |
| 232 | prepare_finalize, |
| 233 | CutlassExpertsFp4( |
| 234 | moe_config=moe_config, |
| 235 | quant_config=quant_config, |
| 236 | ), |
| 237 | ) |
| 238 | |
| 239 | return { |
| 240 | "config": case_config, |
| 241 | "hidden_states": hidden_states, |
| 242 | "score": score, |
| 243 | "kernel": kernel, |
| 244 | "w1_q": w1_q, |
| 245 | "w2_q": w2_q, |
| 246 | "activation": activation, |
| 247 | "e": e, |
| 248 | "topk": topk, |
| 249 | "dtype": dtype, |
| 250 | } |
| 251 | |
| 252 | |
| 253 | def _run_cutlass_fp4_moe( |
| 254 | case: dict[str, Any], |
| 255 | hidden_states: torch.Tensor, |
| 256 | score: torch.Tensor, |
| 257 | ) -> torch.Tensor: |
| 258 | topk_weights, topk_ids, _ = fused_topk( |
| 259 | hidden_states, score, case["topk"], renormalize=False |
| 260 | ) |
| 261 | return case["kernel"].apply( |
| 262 | hidden_states=hidden_states, |
| 263 | w1=case["w1_q"], |
| 264 | w2=case["w2_q"], |
| 265 | topk_weights=topk_weights, |
| 266 | topk_ids=topk_ids, |
| 267 | global_num_experts=case["e"], |
| 268 | activation=case["activation"], |
| 269 | apply_router_weight_on_input=False, |
| 270 | expert_map=None, |
| 271 | ) |
| 272 | |
| 273 | |
| 274 | @_NVFP4_REQUIRES_SM100 |
| 275 | @pytest.mark.parametrize("case_config", _NVFP4_MOE_BATCH_INVARIANT_CASES) |
| 276 | @pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.SWIGLUSTEP]) |
| 277 | @pytest.mark.parametrize("e", _NVFP4_MOE_NUM_EXPERTS) |
| 278 | @pytest.mark.parametrize("topk", _NVFP4_MOE_TOPKS) |
| 279 | @pytest.mark.parametrize("dtype", _NVFP4_MOE_DTYPES) |
| 280 | @torch.inference_mode() |
| 281 | def test_cutlass_nvfp4_moe_batch_invariant( |
| 282 | case_config: dict[str, Any], |
| 283 | activation: MoEActivation, |
| 284 | e: int, |
| 285 | topk: int, |
| 286 | dtype: torch.dtype, |
| 287 | default_vllm_config, |
| 288 | workspace_init, |
| 289 | ) -> None: |
| 290 | case = _make_cutlass_fp4_moe_batch_invariant_case( |
| 291 | case_config, activation, e, topk, dtype |
| 292 | ) |
| 293 | case_id = f"{case['config']}, e={case['e']}, topk={case['topk']}" |
| 294 | # Establish the baseline output for the full mixed batch once, then |
| 295 | # compare every smaller replay against the corresponding slice here. |
| 296 | batch_output = _run_cutlass_fp4_moe(case, case["hidden_states"], case["score"]) |
| 297 | |
| 298 | assert CutlassExpertsFp4._supports_batch_invariance() |
| 299 | |
| 300 | # Re-run the whole batch with a different batch-row permutation to make |
| 301 | # sure the grouped GEMM result is invariant to how tokens are packed into |
| 302 | # expert work for scheduling. |
| 303 | indices = torch.arange( |
| 304 | case["hidden_states"].size(0), device=case["hidden_states"].device |
| 305 | ) |
| 306 | for perm_name, perm in ( |
| 307 | ("reversed", torch.flip(indices, dims=(0,))), |
| 308 | ("evens_then_odds", torch.cat((indices[::2], indices[1::2]))), |
| 309 | ): |
| 310 | permuted_output = _run_cutlass_fp4_moe( |
| 311 | case, |
| 312 | case["hidden_states"][perm], |
| 313 | case["score"][perm], |
| 314 | ) |
| 315 | torch.testing.assert_close( |
| 316 | batch_output[perm], |
| 317 | permuted_output, |
| 318 | atol=0, |
| 319 | rtol=0, |
| 320 | msg=( |
| 321 | f"{case_id}: permutation '{perm_name}' changed outputs " |
| 322 | "relative to the baseline batch order." |
| 323 | ), |
| 324 | ) |
| 325 | |
| 326 | # Re-run every batch-row as a batch-size-1 input and compare against the |
| 327 | # matching row from the full batch. |
| 328 | for idx in range(case["hidden_states"].size(0)): |
| 329 | single_idx = torch.tensor([idx], device=case["hidden_states"].device) |
| 330 | single_output = _run_cutlass_fp4_moe( |
| 331 | case, |
| 332 | case["hidden_states"][single_idx], |
| 333 | case["score"][single_idx], |
| 334 | ) |
| 335 | torch.testing.assert_close( |
| 336 | batch_output[single_idx], |
| 337 | single_output, |
| 338 | atol=0, |
| 339 | rtol=0, |
| 340 | msg=( |
| 341 | f"{case_id}: batch-row {idx} changed between full-batch " |
| 342 | "and batch-size-1 execution." |
| 343 | ), |
| 344 | ) |
| 345 | |
| 346 | # Re-run a few nontrivial sub-batches to catch interactions that only |
| 347 | # appear when multiple tokens are grouped together. |
| 348 | for subset_ids in case["config"]["subset_indices"]: |
| 349 | subset = torch.tensor(subset_ids, device=case["hidden_states"].device) |
| 350 | subset_output = _run_cutlass_fp4_moe( |
| 351 | case, |
| 352 | case["hidden_states"][subset], |
| 353 | case["score"][subset], |
| 354 | ) |
| 355 | torch.testing.assert_close( |
| 356 | batch_output[subset], |
| 357 | subset_output, |
| 358 | atol=0, |
| 359 | rtol=0, |
| 360 | msg=( |
| 361 | f"{case_id}: sub-batch {list(subset_ids)} changed " |
| 362 | "relative to full-batch execution." |
| 363 | ), |
| 364 | ) |