diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index b4c4b59fdf..49bc3c32b6 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -273,7 +273,7 @@ steps: #------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------# - label: ":amd: (MI250) Batch Invariance" # TBD - timeout_in_minutes: 35 + timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] dind: false agent_pool: mi250_1 @@ -282,14 +282,18 @@ steps: source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: + - export VLLM_GPU_MEMORY_UTILIZATION=0.8 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - 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 - label: ":amd: (MI250) CUDAGraph" # TBD timeout_in_minutes: 55 @@ -2443,6 +2447,56 @@ steps: #------------------------------------------------------------ mi300 · v1 -------------------------------------------------------------# +- label: ":amd: (MI300) Batch Invariance" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - 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 + - 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 TRITON_ATTN + +- label: ":amd: (MI300) Batch Invariance (Distributed)" # TBD + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/device_communicators/ + - vllm/model_executor/determinism + - vllm/model_executor/layers + - vllm/v1/attention + - tests/v1/determinism/ + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_TP_SIZE=4 pytest -v -s v1/determinism/test_batch_invariance.py -k "bs1_vs_bsN or without_batch_invariance_should_fail" + - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN + - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py + - label: ":amd: (MI300) Spec Decode Speculators + MTP" # TBD timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -5215,6 +5269,37 @@ steps: #------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# +- label: ":amd: (MI355) Batch Invariance" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_batch_invariance_vlm.py + - 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 + - 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 TRITON_ATTN + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py + - label: ":amd: (MI355) Spec Decode AL DFlash2 Nightly" timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 20b48cb98f..02f78e280c 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -644,12 +644,33 @@ steps: source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - 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 + mirror: + amd: + label: ":amd: (MI250) Batch Invariance" + dind: false + device: mi250_1 + timeout_in_minutes: 60 + env: + VLLM_GPU_MEMORY_UTILIZATION: "0.8" + VLLM_TEST_CLEAN_GPU_MEMORY: "1" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (H100) Batch Invariance" key: batch-invariance-h100 @@ -660,6 +681,8 @@ steps: - vllm/v1/attention - vllm/model_executor/layers - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn @@ -671,15 +694,18 @@ steps: - 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 mirror: amd: - label: ":amd: (MI250) Batch Invariance" + label: ":amd: (MI300) Batch Invariance" dind: false - device: mi250_1 - timeout_in_minutes: 35 + device: mi300_1 + timeout_in_minutes: 60 depends_on: - image-build-amd source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py @@ -688,6 +714,11 @@ steps: - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - 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 + - 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 TRITON_ATTN - label: ":nvidia: (B200) Batch Invariance" key: batch-invariance-b200 @@ -697,6 +728,8 @@ steps: - vllm/v1/attention - vllm/model_executor/layers - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn @@ -711,6 +744,79 @@ steps: - pytest -v -s v1/determinism/test_matmul_batch_invariant.py - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py - pytest -v -s v1/determinism/test_online_batch_invariance.py + mirror: + amd: + label: ":amd: (MI355) Batch Invariance" + dind: false + device: mi355_1 + timeout_in_minutes: 60 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pip install pytest-timeout pytest-forked + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_batch_invariance_vlm.py + - 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 + - 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 TRITON_ATTN + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py + +- label: ":nvidia: (H100) Batch Invariance (Distributed)" + key: batch-invariance-distributed-4xh100 + timeout_in_minutes: 45 + device: h100 + optional: true + num_devices: 4 + source_file_dependencies: + - vllm/distributed/device_communicators/ + - vllm/model_executor/determinism + - vllm/model_executor/layers + - vllm/v1/attention + - tests/v1/determinism/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_TP_SIZE=4 pytest -v -s v1/determinism/test_batch_invariance.py -k "bs1_vs_bsN or without_batch_invariance_should_fail" + - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN + - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN + - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py + mirror: + amd: + label: ":amd: (MI300) Batch Invariance (Distributed)" + dind: false + device: mi300_4 + timeout_in_minutes: 45 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/device_communicators/ + - vllm/model_executor/determinism + - vllm/model_executor/layers + - vllm/v1/attention + - tests/v1/determinism/ + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - VLLM_TEST_TP_SIZE=4 pytest -v -s v1/determinism/test_batch_invariance.py -k "bs1_vs_bsN or without_batch_invariance_should_fail" + - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN + - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py - label: ":nvidia: (H200 MIG 35GB) Acceptance Length (Large Models)" device: h200_35gb diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 03d07896cc..01dc7116fe 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -20,6 +20,7 @@ Batch invariance is crucial for several use cases: Batch invariance is supported on the following platforms: - NVIDIA GPUs with compute capability 8.0 or higher. +- AMD Instinct GPUs (gfx90a, gfx942, gfx950). - Intel XPUs with Triton support. ### Attention Backend Selection for XPU diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py index 55d1f76b2d..b0a2c215ab 100644 --- a/tests/kernels/quantization/test_triton_scaled_mm.py +++ b/tests/kernels/quantization/test_triton_scaled_mm.py @@ -32,11 +32,9 @@ def torch_scaled_mm( out = torch.mm(a.to(torch.float32), b.to(torch.float32)) out = scale_a * out out = scale_b.T * out - out = out.to(out_dtype) if bias is not None: - out = out + bias - - return out + out = out + bias.float() + return out.to(out_dtype) def get_8bit_types(): @@ -125,6 +123,25 @@ def test_scaled_mm( torch.testing.assert_close(c_check, c_actual, rtol=1e-1, atol=1e-1) +@pytest.mark.parametrize("in_dtype", get_8bit_types()) +@pytest.mark.parametrize("use_heuristic", [True, False]) +def test_scaled_mm_adds_bias_before_rounding(in_dtype, use_heuristic): + # The accumulator is 1 + 3 * 2**-10 and the bias 2**-8. Rounding to bf16 + # first gives 1.0 (1 + 2**-8 is a tie, to even); adding first, 1 + 2**-7. + a = torch.ones((1, 32), device=device).to(in_dtype) + col = torch.tensor([64.0] * 16 + [2.0, 1.0] + [0.0] * 14, device=device) + b = col[:, None].expand(32, 16).to(in_dtype).contiguous() + scale_a = torch.ones((1, 1), device=device) + scale_b = torch.full((1, 1), 2**-10, device=device) + bias = torch.full((16,), 2**-8, device=device, dtype=torch.bfloat16) + args = (a, b, scale_a, scale_b, torch.bfloat16, bias) + + out = triton_scaled_mm(*args, use_heuristic=use_heuristic) + expected = torch_scaled_mm(*args) + assert (expected == 1 + 2**-7).all() + torch.testing.assert_close(out, expected, rtol=0, atol=0) + + # TD operand loads must be bit-exact vs the plain masked-load path. @pytest.mark.skipif( not (current_platform.is_cuda_alike() or current_platform.is_xpu()), diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index 4523627151..139c68ed2b 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -180,6 +180,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( llm = LLM( model=TEST_MODEL, tensor_parallel_size=tp_size, + disable_custom_all_reduce=os.getenv("VLLM_TEST_DISABLE_CUSTOM_AR") == "1", max_num_seqs=128, max_model_len=8192, dtype="auto", # not everything is supported diff --git a/tests/v1/determinism/test_batch_invariance_vlm.py b/tests/v1/determinism/test_batch_invariance_vlm.py index fe95b03c61..5c557a1364 100644 --- a/tests/v1/determinism/test_batch_invariance_vlm.py +++ b/tests/v1/determinism/test_batch_invariance_vlm.py @@ -9,7 +9,7 @@ import numpy as np import pytest import torch from PIL import Image -from utils import _extract_step_logprobs, skip_if_not_cuda +from utils import _extract_step_logprobs, skip_if_not_cuda_alike from vllm import LLM, SamplingParams from vllm.inputs import TextPrompt @@ -89,7 +89,7 @@ def _assert_batch_invariant(llm: LLM, inputs: list[TextPrompt], sampling) -> Non ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("input_type", ["image", "video"]) @pytest.mark.parametrize("mm_encoder_attn_backend", ["FLASH_ATTN", "TORCH_SDPA"]) def test_vlm_batch_invariance_bs1_vs_bsN( @@ -101,7 +101,7 @@ def test_vlm_batch_invariance_bs1_vs_bsN( _run_vlm_batch_invariance(input_type, mm_encoder_attn_backend, enforce_eager=True) -@skip_if_not_cuda +@skip_if_not_cuda_alike def test_vlm_batch_invariance_default_execution(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") _run_vlm_batch_invariance("image", "FLASH_ATTN", enforce_eager=False) diff --git a/tests/v1/determinism/test_cutlass_batch_invariance.py b/tests/v1/determinism/test_cutlass_batch_invariance.py index c3384a62f6..b8f8381706 100644 --- a/tests/v1/determinism/test_cutlass_batch_invariance.py +++ b/tests/v1/determinism/test_cutlass_batch_invariance.py @@ -19,8 +19,10 @@ from tests.kernels.moe.utils import make_dummy_moe_config, make_test_quant_confi from tests.kernels.quantization.nvfp4_utils import get_nvfp4_global_scale from tests.utils import TestFP8Layer, requires_fp8 from vllm import _custom_ops as ops -from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( +from vllm.model_executor.kernels.linear import fp8_scaled_mm_is_batch_invariant +from vllm.model_executor.kernels.linear.scaled_mm import ( CutlassFP8ScaledMMLinearKernel, + TritonFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.fused_moe import fused_topk from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -31,8 +33,11 @@ from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( CutlassExpertsFp4, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTensorSym, kFp8DynamicTokenSym, + kFp8StaticChannelSym, kFp8StaticTensorSym, + kFp8StaticTokenSym, ) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -87,16 +92,42 @@ _NVFP4_MOE_BATCH_INVARIANT_CASES = ( @pytest.fixture(autouse=True) def setup_cuda(): - if not current_platform.is_cuda(): - pytest.skip("CUTLASS FP8 kernels require CUDA.") + if not current_platform.is_cuda_alike(): + pytest.skip("Requires CUDA or ROCm.") torch.set_default_device("cuda") @requires_fp8 +@pytest.mark.parametrize( + "kernel", + [ + pytest.param( + CutlassFP8ScaledMMLinearKernel, + marks=pytest.mark.skipif( + not current_platform.is_cuda(), reason="CUTLASS requires CUDA." + ), + ), + TritonFP8ScaledMMLinearKernel, + ], +) +@pytest.mark.parametrize( + "act_key,weight_key", + [ + (kFp8DynamicTokenSym, kFp8StaticTensorSym), + (kFp8DynamicTokenSym, kFp8StaticChannelSym), + (kFp8DynamicTokenSym, kFp8StaticTokenSym), + (kFp8StaticTensorSym, kFp8StaticTensorSym), + # An amax over the whole batch: the GEMM is invariant, the layer is not. + (kFp8DynamicTensorSym, kFp8StaticTensorSym), + ], +) @pytest.mark.parametrize("weight_shape", [(1024, 2048), (4608, 4096)]) @pytest.mark.parametrize("batch_size", [1, 16, 17, 32, 64, 65, 256, 257]) @torch.inference_mode() -def test_cutlass_fp8_batch_invariant_fixed_config( +def test_fp8_scaled_mm_batch_invariant_fixed_config( + kernel: type, + act_key, + weight_key, weight_shape: tuple[int, int], batch_size: int, default_vllm_config, @@ -108,14 +139,16 @@ def test_cutlass_fp8_batch_invariant_fixed_config( torch.manual_seed(0) layer = TestFP8Layer( weight_shape=weight_shape, - activation_quant_key=kFp8DynamicTokenSym, - weight_quant_key=kFp8StaticTensorSym, + activation_quant_key=act_key, + weight_quant_key=weight_key, input_dtype=torch.bfloat16, out_dtype=torch.bfloat16, device=torch.device("cuda"), - force_kernel=CutlassFP8ScaledMMLinearKernel, + force_kernel=kernel, ) - assert isinstance(layer.kernel, CutlassFP8ScaledMMLinearKernel) + assert isinstance(layer.kernel, kernel) + dynamic_per_tensor = act_key == kFp8DynamicTensorSym + assert fp8_scaled_mm_is_batch_invariant(layer.kernel) != dynamic_per_tensor in_features = weight_shape[1] needle = torch.randn((1, in_features), device="cuda", dtype=torch.bfloat16) @@ -124,6 +157,12 @@ def test_cutlass_fp8_batch_invariant_fixed_config( filler = torch.randn( (max(batch_size - 1, 0), in_features), device="cuda", dtype=torch.bfloat16 ) + if dynamic_per_tensor: + if batch_size == 1: + pytest.skip("no other rows to move the scale") + front_output = layer(torch.cat([needle, 100 * filler], dim=0))[0] + assert not torch.equal(front_output, baseline) + return front_batch = torch.cat([needle, filler], dim=0) back_batch = torch.cat([filler, needle], dim=0) @@ -135,6 +174,69 @@ def test_cutlass_fp8_batch_invariant_fixed_config( torch.testing.assert_close(back_output, baseline, rtol=0, atol=0) +@requires_fp8 +@pytest.mark.parametrize("scheme", ["fp8_per_tensor", "fp8_per_channel"]) +@torch.inference_mode() +def test_online_fp8_linear_batch_invariant( + scheme, monkeypatch, default_vllm_config, dist_init +): + from vllm.config.model import ModelConfig + from vllm.config.quantization import resolve_quantization_config + from vllm.model_executor.layers.linear import ReplicatedLinear + from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, + ) + from vllm.utils.torch_utils import set_default_torch_dtype + + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True) + default_vllm_config.model_config = ModelConfig(dtype="bfloat16") + quant_config = OnlineQuantizationConfig(resolve_quantization_config(scheme, None)) + set_random_seed(0) + with set_default_torch_dtype(torch.bfloat16), torch.device("cuda"): + layer = ReplicatedLinear( + 1024, 768, bias=False, quant_config=quant_config, prefix="mlp.down_proj" + ) + layer.weight.weight_loader(layer.weight, torch.randn(768, 1024)) + + needle = torch.randn(1, 1024) + filler = 100 * torch.randn(63, 1024) + baseline = layer(needle)[0] + assert torch.equal(layer(torch.cat([needle, filler]))[0][:1], baseline) + assert torch.equal(layer(torch.cat([filler, needle]))[0][-1:], baseline) + + +@requires_fp8 +@torch.inference_mode() +def test_block_fp8_batch_invariant_across_tuned_block_size_k(monkeypatch): + """The block-FP8 kernel applies the block scales once per K tile, so a tuned + BLOCK_SIZE_K below the scale group size reorders the sum.""" + from vllm.model_executor.layers.quantization.utils import fp8_utils + + base = dict( + BLOCK_SIZE_M=64, BLOCK_SIZE_N=128, GROUP_SIZE_M=1, num_warps=4, num_stages=2 + ) + table = {256: {**base, "BLOCK_SIZE_K": 128}, 512: {**base, "BLOCK_SIZE_K": 64}} + monkeypatch.setattr(fp8_utils, "get_w8a8_block_fp8_configs", lambda *_: table) + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True) + N, K, block = 12288, 2048, [128, 128] + + set_random_seed(0) + fp8 = current_platform.fp8_dtype() + A = torch.randn(512, K).to(fp8) + B = torch.randn(N, K).to(fp8) + As = torch.rand(512, K // 128) / 100 + Bs = torch.rand(N // 128, K // 128) / 100 + + def run(m): + return fp8_utils.w8a8_triton_block_scaled_mm( + A[:m], B, As[:m], Bs, block, torch.bfloat16 + )[:256] + + assert torch.equal(run(256), run(512)) + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False) + assert not torch.equal(run(256), run(512)) + + @_NVFP4_REQUIRES_SM100 @pytest.mark.parametrize("dtype", _NVFP4_LINEAR_DTYPES) @pytest.mark.parametrize("shape", _NVFP4_LINEAR_CONSISTENCY_SHAPES) diff --git a/tests/v1/determinism/test_lora_shrink_batch_invariant.py b/tests/v1/determinism/test_lora_shrink_batch_invariant.py index 2a7029a603..0c4c595c1d 100644 --- a/tests/v1/determinism/test_lora_shrink_batch_invariant.py +++ b/tests/v1/determinism/test_lora_shrink_batch_invariant.py @@ -4,7 +4,7 @@ import pytest import torch -from utils import skip_if_not_cuda +from utils import skip_if_not_cuda_alike import vllm.lora.ops.triton_ops.lora_shrink_op as shrink_op import vllm.lora.ops.triton_ops.utils as lora_utils @@ -39,7 +39,7 @@ def _shrink(inputs, weights, mapping, scaling): return out -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize( "hidden_size,rank,nslices,dtype", [ @@ -68,7 +68,7 @@ def test_lora_shrink_matches_reference(hidden_size, rank, nslices, dtype): torch.testing.assert_close(out, ref, rtol=5e-3, atol=5e-3) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("nslices", [1, 3]) def test_lora_shrink_invariant_across_batch_sizes(monkeypatch, nslices): reduce_kernel, reduce_grids = shrink_op._lora_shrink_reduce_kernel, [] diff --git a/tests/v1/determinism/test_matmul_batch_invariant.py b/tests/v1/determinism/test_matmul_batch_invariant.py index af395e74e7..b6ea1648f3 100644 --- a/tests/v1/determinism/test_matmul_batch_invariant.py +++ b/tests/v1/determinism/test_matmul_batch_invariant.py @@ -10,7 +10,11 @@ import pytest import torch from utils import skip_unsupported -from vllm.model_executor.determinism.batch_invariant import matmul_batch_invariant +import vllm.model_executor.determinism.batch_invariant as bi +from vllm.model_executor.determinism.batch_invariant import ( + enable_batch_invariant_mode, + matmul_batch_invariant, +) from vllm.model_executor.determinism.batch_invariant_configs import ( _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS, _get_tuned_matmul_arch_family, @@ -129,3 +133,49 @@ def test_matmul_batch_invariance_across_tuned_m_buckets(m, transpose_b): batch_output = matmul_batch_invariant(a, b) assert torch.equal(single_output[0], batch_output[0]) + + +@pytest.fixture +def batch_invariant_mode(): + # The overrides are process-global: remove them so later tests compare + # against the stock aten kernels again. + enable_batch_invariant_mode() + yield + assert bi._batch_invariant_LIB is not None + bi._batch_invariant_LIB._destroy() + bi._batch_invariant_MODE = False + + +@skip_unsupported +@pytest.mark.usefixtures("batch_invariant_mode") +@pytest.mark.parametrize("op", ["mm", "addmm", "bmm"]) +def test_compiled_matmul_reaches_batch_invariant_out_overloads(op): + """Inductor lowers mm/addmm/bmm to ``extern_kernels.(..., out=buf)``, + i.e. the ``.out`` overload, which eager never dispatches.""" + if ( + op != "bmm" + and current_platform.is_cuda() + and not current_platform.is_device_capability_family(80) + ): + pytest.skip("mm/addmm are not overridden on SM90+") + torch.manual_seed(0) + dtype, k, n, needle = torch.bfloat16, 1024, 768, 37 + w = torch.randn(4, k, n, device=DEVICE_TYPE, dtype=dtype) + bias = torch.randn(n, device=DEVICE_TYPE, dtype=dtype) + fn = { + "mm": lambda x: torch.mm(x[0], w[0]), + # Bare addmm: a following pointwise op would make Inductor unfuse the bias. + "addmm": lambda x: torch.addmm(bias, x[0], w[0]), + "bmm": lambda x: torch.bmm(x, w), + }[op] + compiled = torch.compile(fn, fullgraph=True) + x = torch.randn(4, 256, k, device=DEVICE_TYPE, dtype=dtype) + row = None + for m in (64, 256, 128): + xs = x[:, :m].clone() + # With dynamic=True, a specialization on M would silently recompile. + torch._dynamo.mark_dynamic(xs, 1) + out = compiled(xs) + assert torch.equal(out, fn(xs)), f"compiled {op} bypassed the BI kernel (M={m})" + row = out[..., needle, :] if row is None else row + assert torch.equal(out[..., needle, :], row) diff --git a/tests/v1/determinism/test_mx_linear_batch_invariant.py b/tests/v1/determinism/test_mx_linear_batch_invariant.py new file mode 100644 index 0000000000..af31302ac9 --- /dev/null +++ b/tests/v1/determinism/test_mx_linear_batch_invariant.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The ROCm native MX linear kernels must not depend on the number of rows.""" + +import pytest +import torch +from utils import requires_mx + +from vllm.utils.torch_utils import set_random_seed + +SEED = 0 + +# E8M0 exponent half-width. With 4-bit operands and a narrow spread the fp32 +# partial sums are exact, so a reordering of them would be invisible. +SCALE_SPREAD = 15 + +# (N, K, use_asm_gemm) +MXFP4_CASES = [ + (4096, 2048, False), + (2048, 6144, False), + (4608, 7168, False), + (8192, 2048, True), + (8192, 1024, True), + (8192, 14336, True), +] + +MXFP8_CASES = [(4096, 2048), (2048, 6144), (1024, 768), (1536, 1024)] + +# The untuned AITER MXFP4 configs switch at M = 9/33/65/129/257/513. +MXFP4_TOKEN_COUNTS = [ + 1, + 2, + 8, + 9, + 16, + 32, + 33, + 64, + 65, + 128, + 129, + 256, + 257, + 512, + 513, + 1024, + 2048, +] + +MXFP8_TOKEN_COUNTS = [1, 32, 64, 65, 128, 129, 256, 257, 512, 1024, 1025, 2048] + +# Row 0 lands in the first tile of every decomposition, so it can stay +# invariant when the rest of the output does not. +CHECK_ROWS = [0, 1, 2, 3, 7, 15, 31] + + +def _probe(out: torch.Tensor) -> torch.Tensor: + """The leading rows of an output, kept so the sweep does not hold every GEMM.""" + return out[: CHECK_ROWS[-1] + 1].clone() + + +def _variant_rows(probes: dict[int, torch.Tensor]) -> list[str]: + """Rows whose contents changed with the number of rows in the launch.""" + failures = [] + for row in CHECK_ROWS: + counts = [n for n in sorted(probes) if n > row] + if not counts: + continue + reference = probes[counts[0]][row] + variant = [n for n in counts if not torch.equal(probes[n][row], reference)] + if variant: + failures.append(f"row {row} changed at row counts {variant}") + return failures + + +def _mxfp4_weights(n: int, k: int, use_asm_gemm: bool): + """Weights laid out as AiterMxfp4LinearKernel.process_weights_after_loading.""" + weight = torch.randint(0, 255, (n, k // 2), dtype=torch.uint8, device="cuda") + weight_scale = torch.randint( + 127 - SCALE_SPREAD, + 128 + SCALE_SPREAD, + (n, k // 32), + dtype=torch.uint8, + device="cuda", + ) + if not use_asm_gemm: + return weight, weight_scale.T.contiguous() + + from aiter.ops.shuffle import shuffle_weight + + sm, sn = weight_scale.shape + weight_scale = ( + weight_scale.view(sm // 32, 2, 16, sn // 8, 2, 4, 1) + .permute(0, 3, 5, 2, 4, 1, 6) + .contiguous() + .view(sm, sn) + ) + return shuffle_weight(weight, layout=(16, 16)), weight_scale + + +def _mxfp8_weights(n: int, k: int): + """E4M3 weights with E8M0 block scales, as the MXFP8 layer stores them.""" + weight = (torch.randn(n, k, device="cuda") / 8).to(torch.float8_e4m3fn) + weight_scale = torch.randint( + 127 - SCALE_SPREAD, + 128 + SCALE_SPREAD, + (n, k // 32), + dtype=torch.uint8, + device="cuda", + ) + return weight, weight_scale + + +@requires_mx +@pytest.mark.parametrize("n,k,use_asm_gemm", MXFP4_CASES) +def test_mxfp4_linear_is_batch_invariant(n: int, k: int, use_asm_gemm: bool): + pytest.importorskip("aiter") + + # Importing the module registers torch.ops.vllm.gemm_with_dynamic_quant. + import vllm.model_executor.kernels.linear.mxfp4.aiter # noqa: F401 + + set_random_seed(SEED) + weight, weight_scale = _mxfp4_weights(n, k, use_asm_gemm) + x = torch.randn(max(MXFP4_TOKEN_COUNTS), k, device="cuda", dtype=torch.bfloat16) + + probes = { + num_tokens: _probe( + torch.ops.vllm.gemm_with_dynamic_quant( + x[:num_tokens], weight, weight_scale, use_asm_gemm, torch.bfloat16 + ) + ) + for num_tokens in MXFP4_TOKEN_COUNTS + } + + failures = _variant_rows(probes) + assert not failures, ( + f"MXFP4 linear depends on the row count " + f"(N={n}, K={k}, asm={use_asm_gemm}):\n " + "\n ".join(failures) + ) + + +@requires_mx +@pytest.mark.parametrize("n,k", MXFP8_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_mxfp8_linear_is_batch_invariant(n: int, k: int, dtype: torch.dtype): + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + _mxfp8_dot_scaled_linear, + ) + + set_random_seed(SEED) + weight, weight_scale = _mxfp8_weights(n, k) + x = torch.randn(max(MXFP8_TOKEN_COUNTS), k, device="cuda", dtype=dtype) + + probes = { + num_tokens: _probe( + _mxfp8_dot_scaled_linear(x[:num_tokens], weight, weight_scale) + ) + for num_tokens in MXFP8_TOKEN_COUNTS + } + + failures = _variant_rows(probes) + assert not failures, ( + f"MXFP8 linear depends on the row count (N={n}, K={k}, {dtype}):\n " + + "\n ".join(failures) + ) diff --git a/tests/v1/determinism/test_nvfp4_batch_invariant.py b/tests/v1/determinism/test_nvfp4_batch_invariant.py index 1e9b0cc8fc..162e6b36c3 100644 --- a/tests/v1/determinism/test_nvfp4_batch_invariant.py +++ b/tests/v1/determinism/test_nvfp4_batch_invariant.py @@ -9,10 +9,14 @@ import torch from utils import ( _extract_step_logprobs, _random_prompt, - skip_if_not_cuda, + skip_if_not_cuda_alike, ) from vllm import LLM, SamplingParams +from vllm.platforms import current_platform + +# FLASH_ATTN is CUDA-only. +E2E_BACKENDS = ["TRITON_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"] pytestmark = pytest.mark.skipif( not hasattr(torch, "float8_e4m3fn"), @@ -40,8 +44,8 @@ def _make_llm(max_num_seqs: int, backend: str) -> LLM: ) -@skip_if_not_cuda -@pytest.mark.parametrize("backend", ["FLASH_ATTN"]) +@skip_if_not_cuda_alike +@pytest.mark.parametrize("backend", E2E_BACKENDS) def test_dense_nvfp4_generation_is_deterministic_across_batch_sizes_e2e(backend): seed = int(os.getenv("VLLM_TEST_SEED", "12345")) random.seed(seed) diff --git a/tests/v1/determinism/test_online_batch_invariance.py b/tests/v1/determinism/test_online_batch_invariance.py index dcba87dc49..fb5db254ba 100644 --- a/tests/v1/determinism/test_online_batch_invariance.py +++ b/tests/v1/determinism/test_online_batch_invariance.py @@ -6,6 +6,7 @@ vLLM server and compare BS=1 vs BS=N results (tokens and per-step logprobs). Environment variables: - VLLM_TEST_MODEL: served model name (e.g., Qwen/Qwen3-1.7B / DeepSeek-R1) - VLLM_TP_SIZE: tensor parallelism size (e.g., 4) + - VLLM_TEST_SERVER_ARGS: extra server args (e.g., "--data-parallel-size 4") """ @@ -16,7 +17,7 @@ from typing import Any import openai import pytest -from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_if_not_cuda +from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_if_not_cuda_alike from tests.utils import RemoteOpenAIServer @@ -132,7 +133,7 @@ def _compare_bs1_vs_bsn_single_process( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("backend", BACKENDS) def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( backend: str, @@ -156,6 +157,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( ] if tp_size: server_args += ["-tp", tp_size] + server_args += os.getenv("VLLM_TEST_SERVER_ARGS", "").split() with RemoteOpenAIServer(TEST_MODEL, server_args) as server: client = server.get_client() diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index 1feb5e6741..f26217bd74 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -4,7 +4,7 @@ import pytest import torch -from utils import skip_if_not_cuda, skip_unsupported +from utils import skip_if_not_cuda_alike, skip_unsupported from vllm.model_executor.determinism.batch_invariant import ( rms_norm_batch_invariant, @@ -27,7 +27,7 @@ def _rms_norm_reference( return (output * weight.float()).to(input_tensor.dtype) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("batch_size", [1, 4, 64, 300]) @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -76,7 +76,7 @@ def test_rms_norm_batch_invariant_vs_reference( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("eps", [1e-6]) @@ -189,7 +189,7 @@ def _assert_rows_bit_identical(small, large, msg): torch.testing.assert_close(small, large, rtol=0.0, atol=0.0, msg=msg) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("seed", list(range(4))) @@ -221,7 +221,7 @@ def test_rms_norm_batch_invariant_nonresidual_kernel( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("seed", list(range(4))) @@ -267,7 +267,7 @@ def test_rms_norm_static_fp8_quant_batch_invariant( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [512, 4096]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("seed", list(range(4))) @@ -303,7 +303,7 @@ def test_rms_norm_per_block_quant_batch_invariant( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("batch_size", [1, 16, 128]) @pytest.mark.parametrize("seq_len", [1, 32, 512]) @pytest.mark.parametrize("hidden_size", [2048, 4096]) @@ -343,7 +343,7 @@ def test_rms_norm_3d_input( ) -@skip_if_not_cuda +@skip_if_not_cuda_alike def test_rms_norm_numerical_stability(default_vllm_config): """Test RMS norm numerical stability with extreme values. @@ -431,7 +431,7 @@ def test_rms_norm_formula(default_vllm_config): ) -@skip_if_not_cuda +@skip_if_not_cuda_alike @pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384]) def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int): """Test RMS norm with various hidden sizes to ensure block size handling. diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py index f03ea05b43..405b5ac3d8 100644 --- a/tests/v1/determinism/utils.py +++ b/tests/v1/determinism/utils.py @@ -34,12 +34,16 @@ DEVICE_BACKENDS: dict[str, DeviceConfig] = { available=current_platform.is_xpu() and HAS_TRITON, backends=["TRITON_ATTN"], ), + "rocm": DeviceConfig( + available=current_platform.is_rocm() and HAS_TRITON, + backends=["TRITON_ATTN"], + ), } DEFAULT_MODEL = "Qwen/Qwen3-1.7B" TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL) -# Override backends for MLA models (MLA only supported on CUDA). +# Override backends for MLA models. if os.getenv("VLLM_TEST_MODEL"): config = get_config(TEST_MODEL, trust_remote_code=False) if ModelArchConfigConvertorBase(config, config.get_text_config()).is_deepseek_mla(): @@ -52,6 +56,10 @@ if os.getenv("VLLM_TEST_MODEL"): available=DEVICE_BACKENDS["xpu"].available, backends=[], ) + DEVICE_BACKENDS["rocm"] = DeviceConfig( + available=DEVICE_BACKENDS["rocm"].available, + backends=["TRITON_MLA"], + ) # Only include backends for devices that are actually available. BACKENDS: list[str] = sorted( @@ -60,7 +68,7 @@ BACKENDS: list[str] = sorted( skip_unsupported = pytest.mark.skipif( not any(cfg.available for cfg in DEVICE_BACKENDS.values()), - reason="Requires CUDA >= Ampere (SM80) or Intel XPU with Triton", + reason="Requires CUDA >= Ampere (SM80), ROCm, or Intel XPU with Triton", ) skip_if_not_cuda = pytest.mark.skipif( @@ -68,6 +76,16 @@ skip_if_not_cuda = pytest.mark.skipif( reason="Requires CUDA >= Ampere (SM80)", ) +skip_if_not_cuda_alike = pytest.mark.skipif( + not (DEVICE_BACKENDS["cuda"].available or DEVICE_BACKENDS["rocm"].available), + reason="Requires CUDA >= Ampere (SM80) or ROCm", +) + +requires_mx = pytest.mark.skipif( + not (current_platform.is_rocm() and current_platform.supports_mx()), + reason="requires a ROCm device with native MX support (gfx95x)", +) + def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str: # Generate more realistic prompts that will actually produce varied tokens diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 29d5bae89d..1bc6182288 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -413,6 +413,17 @@ class CudaCommunicator(DeviceCommunicatorBase): out = ca_comm.custom_all_reduce(input_) assert out is not None return out + if ( + envs.VLLM_BATCH_INVARIANT + and current_platform.is_rocm() + and input_.is_floating_point() + ): + # Integer sums are exact in any order and stay on RCCL. + from vllm.model_executor.determinism.batch_invariant import ( + all_reduce_batch_invariant, + ) + + return all_reduce_batch_invariant(input_, self.device_group) symm_mem_comm = self.symm_mem_comm if symm_mem_comm is not None and symm_mem_comm.should_use_symm_mem(input_): out = symm_mem_comm.all_reduce(input_) @@ -498,7 +509,17 @@ class CudaCommunicator(DeviceCommunicatorBase): chunk_size = input_tensor.shape[0] // world_size output_shape = (chunk_size,) + input_tensor.shape[1:] - if should_nccl_symm_mem_ag_rs(): + if ( + envs.VLLM_BATCH_INVARIANT + and current_platform.is_rocm() + and input_.is_floating_point() + ): + from vllm.model_executor.determinism.batch_invariant import ( + reduce_scatter_batch_invariant, + ) + + output = reduce_scatter_batch_invariant(input_tensor, self.device_group) + elif should_nccl_symm_mem_ag_rs(): output = self._reduce_scatter_symm_mem(input_tensor) else: output = torch.empty( @@ -563,6 +584,22 @@ class CudaCommunicator(DeviceCommunicatorBase): assert output.device == input_tensor.device assert output.is_contiguous() + if ( + envs.VLLM_BATCH_INVARIANT + and current_platform.is_rocm() + and input_.is_floating_point() + ): + from vllm.model_executor.determinism.batch_invariant import ( + reduce_scatter_batch_invariant, + ) + + result = reduce_scatter_batch_invariant( + input_tensor, self.device_group, sizes=sizes + ) + if output is not None: + return output.copy_(result) + return result.movedim(0, dim).contiguous() + if self._can_use_aiter_ag_rs(sizes): aiter_comm = self.aiter_ar_comm assert aiter_comm is not None diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index 403f475702..e4b04793fe 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -6,6 +6,7 @@ from collections.abc import Callable from typing import Any import torch +import torch.distributed as dist import vllm.envs as envs from vllm.model_executor.determinism.batch_invariant_configs import ( @@ -59,6 +60,7 @@ def matmul_kernel_persistent( stride_bn, stride_cm, stride_cn, + stride_bias, BLOCK_SIZE_M: tl.constexpr, # BLOCK_SIZE_N: tl.constexpr, # BLOCK_SIZE_K: tl.constexpr, # @@ -128,7 +130,7 @@ def matmul_kernel_persistent( c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) if HAS_BIAS: - bias_ptrs = bias_ptr + offs_cn + bias_ptrs = bias_ptr + offs_cn * stride_bias bias = tl.load(bias_ptrs, mask=offs_cn < N, other=0.0).to(tl.float32) accumulator += bias c = accumulator.to(c_ptr.dtype.element_ty) @@ -305,6 +307,8 @@ def matmul_persistent( ), ) + if not _BLOCK_SIZES_RESOLVED: + _resolve_block_sizes() configs = { torch.bfloat16: { "BLOCK_SIZE_M": 128, @@ -345,6 +349,7 @@ def matmul_persistent( b.stride(1), # c.stride(0), c.stride(1), # + bias.stride(0) if bias is not None else 1, NUM_SMS=NUM_SMS, # A_LARGE=a.numel() > 2**31, B_LARGE=b.numel() > 2**31, @@ -489,16 +494,19 @@ def bmm_kernel( @triton.jit -def _log_softmax_kernel( +def _softmax_kernel( input_ptr, output_ptr, input_row_stride, + input_col_stride, output_row_stride, n_cols, BLOCK_SIZE: tl.constexpr, + LOG: tl.constexpr = False, ): - """Compute log_softmax along the last dimension of a 2D tensor. - Each block handles one row of the input tensor. + """Compute softmax, or log_softmax when ``LOG``, along the last dimension of a + 2D tensor. Each block handles one row of the input tensor, in a fixed block + order, so the result never depends on the row count. """ # Get the row index for this block row_idx = tl.program_id(0).to(tl.int64) @@ -514,7 +522,9 @@ def _log_softmax_kernel( mask = col_idx < n_cols # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask, other=-float("inf")) + vals = tl.load( + row_start_ptr + col_idx * input_col_stride, mask=mask, other=-float("inf") + ) # Update maximum max_val = tl.max(tl.maximum(vals, max_val)) @@ -526,25 +536,26 @@ def _log_softmax_kernel( mask = col_idx < n_cols # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + vals = tl.load(row_start_ptr + col_idx * input_col_stride, mask=mask, other=0.0) # Compute exp(x - max_val) and accumulate exp_vals = tl.exp(vals - max_val) sum_exp += tl.sum(tl.where(mask, exp_vals, 0.0)) - # Compute log(sum_exp) - log_sum_exp = tl.log(sum_exp) + log_sum_exp = tl.log(sum_exp) if LOG else 0.0 - # Step 3: Compute final log_softmax values: x - max_val - log_sum_exp + # Step 3: normalise for col_offset in range(0, n_cols, BLOCK_SIZE): col_idx = col_offset + tl.arange(0, BLOCK_SIZE) mask = col_idx < n_cols # Load values - vals = tl.load(row_start_ptr + col_idx, mask=mask) + vals = tl.load(row_start_ptr + col_idx * input_col_stride, mask=mask) - # Compute log_softmax - output = vals - max_val - log_sum_exp + if LOG: + output = vals - max_val - log_sum_exp + else: + output = tl.exp(vals - max_val) / sum_exp # Store results tl.store(output_row_start_ptr + col_idx, output, mask=mask) @@ -570,30 +581,55 @@ def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: # Flatten all dimensions except the last one original_shape = input.shape input_2d = input.reshape(-1, input.shape[-1]) - input_2d = input_2d.contiguous() n_rows, n_cols = input_2d.shape # Allocate output tensor - output = torch.empty_like(input_2d) + output = torch.empty((n_rows, n_cols), dtype=input_2d.dtype, device=input_2d.device) # Choose block size based on the number of columns BLOCK_SIZE = 1024 # Launch kernel with one block per row grid = (n_rows,) - _log_softmax_kernel[grid]( + _softmax_kernel[grid]( input_2d, output, input_2d.stride(0), + input_2d.stride(1), output.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE, + LOG=True, ) # Reshape output back to original shape return output.reshape(original_shape) +def softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: + """Compute softmax along the last dimension using a Triton kernel.""" + if dim != -1 and dim != input.ndim - 1: + raise ValueError( + "This implementation only supports softmax along the last dimension" + ) + + original_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + n_rows, n_cols = input_2d.shape + output = torch.empty((n_rows, n_cols), dtype=input_2d.dtype, device=input_2d.device) + + _softmax_kernel[(n_rows,)]( + input_2d, + output, + input_2d.stride(0), + input_2d.stride(1), + output.stride(0), + n_cols, + BLOCK_SIZE=1024, + ) + return output.reshape(original_shape) + + @triton.jit def mean_kernel( input_ptr, @@ -735,6 +771,9 @@ def mean_dim( def mm_batch_invariant(a, b, bias=None): if a.device.type == "xpu": + # The descriptor kernel reads the bias at unit stride. + if bias is not None: + bias = bias.contiguous() return matmul_descriptor_persistent(a, b, bias=bias) return matmul_persistent(a, b, bias=bias) @@ -820,6 +859,8 @@ def bmm_batch_invariant(a, b, *, out=None): assert out.dtype == dtype and out.device == a.device, "out tensor mismatch" c = out + if not _BLOCK_SIZES_RESOLVED: + _resolve_block_sizes() configs = { torch.bfloat16: { "BLOCK_SIZE_M": 128, @@ -877,8 +918,50 @@ def bmm_batch_invariant(a, b, *, out=None): return c -def addmm_batch_invariant(bias, a, b): - return mm_batch_invariant(a, b, bias=bias) +def _addmm_fused_bias(bias, beta, alpha, n): + # The kernel adds bias into its fp32 accumulator indexed by output column, + # so it can only absorb a row vector of width n, unscaled by alpha. + if beta == 0 or alpha != 1: + return None + if bias.dim() != 1 or bias.shape[0] != n: + return None + # fp32 keeps the scaling exact; the kernel casts the bias to fp32 anyway. + return bias if beta == 1 else bias.float() * beta + + +def _addmm_impl(bias, a, b, beta, alpha, out): + fused_bias = _addmm_fused_bias(bias, beta, alpha, b.shape[1]) + result = mm_batch_invariant(a, b, bias=fused_bias) + if fused_bias is None: + if alpha != 1: + result = result * alpha + if beta != 0: + result = result + (bias if beta == 1 else beta * bias) + if out is None: + return result + return out.copy_(result) + + +def addmm_batch_invariant(bias, a, b, *, beta=1, alpha=1): + return _addmm_impl(bias, a, b, beta, alpha, None) + + +# Inductor lowers mm/addmm/bmm to ``extern_kernels.(..., out=buf)``, which +# dispatches the ``.out`` overload rather than the default one. Registering only +# the default leaves the compiled path on the vendor GEMM, so a compiled +# ``torch.addmm`` is batch variant while the eager one is not. + + +def mm_out_batch_invariant(a, b, *, out): + return matmul_batch_invariant(a, b, out=out) + + +def addmm_out_batch_invariant(bias, a, b, *, beta=1, alpha=1, out): + return _addmm_impl(bias, a, b, beta, alpha, out) + + +def bmm_out_batch_invariant(a, b, *, out): + return bmm_batch_invariant(a, b, out=out) def _log_softmax_batch_invariant(input, dim, _half_to_float): @@ -887,10 +970,16 @@ def _log_softmax_batch_invariant(input, dim, _half_to_float): return log_softmax(input, dim=dim) -def softmax_batch_invariant(input, dim, dtype=None): - # Compute softmax in a deterministic way - # First subtract max for numerical stability (standard practice) +def softmax_batch_invariant(input, dim, half_to_float=False): + if half_to_float: + input = input.float() + if dim == -1 or dim == input.ndim - 1: + return softmax(input, dim=-1) + + # Off the last dimension this falls back to torch ops, whose reduction split + # depends on the tensor shape. input_max = torch.amax(input, dim=dim, keepdim=True) + # First subtract max for numerical stability (standard practice) input = input - input_max exp_x = torch.exp(input) sum_exp_x = torch.sum(exp_x, dim=dim, keepdim=True) @@ -1035,11 +1124,150 @@ def rms_norm_batch_invariant( def linear_batch_invariant(input, weight, bias=None): - output = matmul_batch_invariant(input, weight.t()) - if bias is not None: - output = output + bias - return output + bias = bias.expand(weight.shape[0]) + output = mm_batch_invariant( + input.reshape(-1, input.shape[-1]), weight.t(), bias=bias + ) + return output.view(*input.shape[:-1], output.shape[-1]) + + +@triton.jit +def _fixed_order_sum_kernel( + src_ptr, # (WORLD_SIZE, numel) gathered contributions, rank-major + out_ptr, # (numel,) + numel, + stride_rank, + WORLD_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + LARGE: tl.constexpr, +): + """Sum ``WORLD_SIZE`` contributions in ascending rank order, fp32 accumulator. + + An output element sums exactly ``WORLD_SIZE`` values in a compile-time + constant order and rounds once, so the result depends only on the rank + ordering and never on how many elements are being reduced. + """ + pid = tl.program_id(0) + if LARGE: + pid = pid.to(tl.int64) + stride_rank = stride_rank.to(tl.int64) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs < numel + + acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for r in tl.static_range(WORLD_SIZE): + vals = tl.load(src_ptr + r * stride_rank + offs, mask=mask, other=0.0) + acc += vals.to(tl.float32) + + tl.store(out_ptr + offs, acc.to(out_ptr.dtype.element_ty), mask=mask) + + +def all_reduce_batch_invariant( + input_: torch.Tensor, group: "torch.distributed.ProcessGroup | None" = None +) -> torch.Tensor: + """Sum all-reduce whose result does not depend on the message size. + + Library all-reduces (NCCL/RCCL) pick their algorithm, channel count and chunk + boundaries from the message size, so the order in which a given element's + contributions are summed changes with the number of tokens in the batch. + + Instead, all-gather the contributions -- pure data movement, so bitwise + reproducible at any size -- and reduce them with ``_fixed_order_sum_kernel``. + """ + world_size = dist.get_world_size(group) + if world_size == 1: + return input_ + + x = input_.contiguous() + numel = x.numel() + gathered = torch.empty((world_size, numel), dtype=x.dtype, device=x.device) + dist.all_gather_into_tensor(gathered, x.view(-1), group=group) + + out = torch.empty_like(x) + BLOCK_SIZE = 1024 + grid = (triton.cdiv(numel, BLOCK_SIZE),) + _fixed_order_sum_kernel[grid]( + gathered, + out.view(-1), + numel, + numel, + WORLD_SIZE=world_size, + BLOCK_SIZE=BLOCK_SIZE, + LARGE=gathered.numel() > 2**31, + ) + return out + + +def reduce_scatter_batch_invariant( + input_: torch.Tensor, + group: "torch.distributed.ProcessGroup | None" = None, + sizes: list[int] | None = None, +) -> torch.Tensor: + """Sum reduce-scatter over dim 0 whose result does not depend on the size. + + ``ncclReduceScatter`` picks its chunking from the message size just as + ``ncclAllReduce`` does. Route around it the same way + ``all_reduce_batch_invariant`` does: an all-to-all sends rank ``d`` exactly + the rows it owns from every rank -- pure data movement, so bitwise + reproducible however the library chunks it -- and lands them rank-major, + which is the layout ``_fixed_order_sum_kernel`` already consumes. + + ``sizes`` (the reduce-scatterv case) changes only *which* rows a rank + receives, never how the received contributions are summed, so variable + shard sizes are invariant for the same reason. + + Args: + input_: Contiguous full-size input; dim 0 is the scattered axis. + group: Process group, defaults to the world group. + sizes: Per-rank row counts. Uniform split when ``None``. + + Returns: + This rank's shard of the sum. + + """ + world_size = dist.get_world_size(group) + if world_size == 1: + return input_ + + x = input_.contiguous() + if sizes is None: + assert x.shape[0] % world_size == 0 + sizes = [x.shape[0] // world_size] * world_size + else: + assert len(sizes) == world_size + assert x.shape[0] == sum(sizes) + + rows = sizes[dist.get_rank(group)] + out = torch.empty((rows, *x.shape[1:]), dtype=x.dtype, device=x.device) + recv = torch.empty( + (world_size * rows, *x.shape[1:]), dtype=x.dtype, device=x.device + ) + # Every rank takes part even when its own shard is empty: it still has rows + # to send to the others. + dist.all_to_all_single( + recv, + x, + output_split_sizes=[rows] * world_size, + input_split_sizes=sizes, + group=group, + ) + + numel = out.numel() + if numel == 0: + return out + BLOCK_SIZE = 1024 + grid = (triton.cdiv(numel, BLOCK_SIZE),) + _fixed_order_sum_kernel[grid]( + recv.view(-1), + out.view(-1), + numel, + numel, + WORLD_SIZE=world_size, + BLOCK_SIZE=BLOCK_SIZE, + LARGE=recv.numel() > 2**31, + ) + return out _batch_invariant_MODE = False @@ -1047,9 +1275,33 @@ _batch_invariant_LIB = None _fp16_block_size_n = 256 _fp32_block_size_n = 128 _fp32_num_stages = 3 +_BLOCK_SIZES_RESOLVED = False _NUM_SMS: int = 0 +def _resolve_block_sizes(): + """Size the tiles to the shared memory, to avoid triton OutOfResources. + + Runs when the mode is enabled, or on first use by a caller that never + enables it, such as a kernel unit test. + """ + global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages + global _BLOCK_SIZES_RESOLVED + _BLOCK_SIZES_RESOLVED = True + if not current_platform.is_cuda_alike(): + return + if get_max_shared_memory_bytes() > 106496: + _fp16_block_size_n = 256 + _fp32_block_size_n = 128 + _fp32_num_stages = 3 + else: + _fp16_block_size_n = 128 + # SM89 N=1 fp32 128/stages=3 compiles to 131072B (>101376). + # Only N=1 OORs; keep 128/stages=3 for every other N. + _fp32_block_size_n = 32 + _fp32_num_stages = 2 + + def enable_batch_invariant_mode(): global _batch_invariant_MODE, _batch_invariant_LIB global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages @@ -1062,12 +1314,17 @@ def enable_batch_invariant_mode(): _batch_invariant_LIB = torch.library.Library("aten", "IMPL") key = current_platform.dispatch_key - if current_platform.is_cuda(): - if current_platform.is_device_capability_family(80): - # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the - # triton persistent matmul overrides for mm/addmm/matmul/linear. + if current_platform.is_cuda_alike(): + if current_platform.is_rocm() or current_platform.is_device_capability_family( + 80 + ): + # Neither cuBLASLt on SM80 (Ampere) nor hipBLASLt can be made batch + # invariant; install the triton persistent matmul overrides for + # mm/addmm/matmul/linear. _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key) _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key) + _batch_invariant_LIB.impl("aten::mm.out", mm_out_batch_invariant, key) + _batch_invariant_LIB.impl("aten::addmm.out", addmm_out_batch_invariant, key) _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key) _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key) else: @@ -1077,18 +1334,7 @@ def enable_batch_invariant_mode(): os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1" - # Query the shared memory size and set block size - # accordingly to avoid triton OutOfResources. - if get_max_shared_memory_bytes() > 106496: - _fp16_block_size_n = 256 - _fp32_block_size_n = 128 - _fp32_num_stages = 3 - else: - _fp16_block_size_n = 128 - # SM89 N=1 fp32 128/stages=3 compiles to 131072B (>101376). - # Only N=1 OORs; keep 128/stages=3 for every other N. - _fp32_block_size_n = 32 - _fp32_num_stages = 2 + _resolve_block_sizes() elif current_platform.is_xpu(): # Tensor descriptors need a global-memory allocator; must be set outside # torch.compile regions (triton.set_allocator modifies global state). @@ -1100,6 +1346,8 @@ def enable_batch_invariant_mode(): _NUM_SMS = num_compute_units(0) _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key) _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key) + _batch_invariant_LIB.impl("aten::mm.out", mm_out_batch_invariant, key) + _batch_invariant_LIB.impl("aten::addmm.out", addmm_out_batch_invariant, key) _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key) _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key) @@ -1123,7 +1371,9 @@ def enable_batch_invariant_mode(): _batch_invariant_LIB.impl( "aten::bmm", bmm_batch_invariant, key, allow_override=True ) - torch.bmm = bmm_batch_invariant + _batch_invariant_LIB.impl( + "aten::bmm.out", bmm_out_batch_invariant, key, allow_override=True + ) reduced_precision_val = ( (False, False) if is_torch_equal_or_newer("2.10.0") else False @@ -1145,28 +1395,32 @@ def override_envs_for_invariance(): os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - # NCCL determinism settings. NCCL keeps the launch mode process-wide, so, - # unlike the pins below, it also applies to communicators shared with other - # processes. - os.environ["NCCL_LAUNCH_MODE"] = "GROUP" - pin_nccl_env( - { - "NCCL_COLLNET_ENABLE": "0", - "NCCL_NVLS_ENABLE": "0", - "NCCL_P2P_NET_DISABLE": "1", - "NCCL_MIN_NCHANNELS": "1", - "NCCL_MAX_NCHANNELS": "1", - "NCCL_PROTO": "Simple", - # NCCL >= 2.31 zero-fills the algorithm table of every collective - # when NCCL_ALGO is set and re-enables only the named ones; together - # with the NCCL_PROTO above, collectives not named here end up with - # no algorithm and fail with ncclInvalidUsage. Re-enable Ring and - # Tree for all collectives, then pin AllReduce to Tree. - "NCCL_ALGO": "ring,tree;allreduce:tree", - "NCCL_NTHREADS": "1", - "NCCL_SOCKET_NTHREADS": "1", - } - ) + if current_platform.is_rocm(): + # The skinny GEMMs pick their split of K from the batch size. + os.environ["VLLM_ROCM_USE_SKINNY_GEMM"] = "0" + else: + # NCCL determinism settings. NCCL keeps the launch mode process-wide, so, + # unlike the pins below, it also applies to communicators shared with other + # processes. + os.environ["NCCL_LAUNCH_MODE"] = "GROUP" + pin_nccl_env( + { + "NCCL_COLLNET_ENABLE": "0", + "NCCL_NVLS_ENABLE": "0", + "NCCL_P2P_NET_DISABLE": "1", + "NCCL_MIN_NCHANNELS": "1", + "NCCL_MAX_NCHANNELS": "1", + "NCCL_PROTO": "Simple", + # NCCL >= 2.31 zero-fills the algorithm table of every collective + # when NCCL_ALGO is set and re-enables only the named ones; together + # with the NCCL_PROTO above, collectives not named here end up with + # no algorithm and fail with ncclInvalidUsage. Re-enable Ring and + # Tree for all collectives, then pin AllReduce to Tree. + "NCCL_ALGO": "ring,tree;allreduce:tree", + "NCCL_NTHREADS": "1", + "NCCL_SOCKET_NTHREADS": "1", + } + ) # torch.compile settings os.environ["VLLM_USE_AOT_COMPILE"] = "0" diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 484143ffc1..0ce22f477a 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -218,6 +218,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.rocm import ( ) from vllm.model_executor.kernels.linear.scaled_mm.triton import ( TritonFp8BlockScaledMMKernel, + TritonFP8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.xpu import ( @@ -253,6 +254,20 @@ def _get_linear_backend(*, quantization: str) -> str: return config.kernel_config.linear_backend +def _batch_invariant_fp8_kernel(kernel, config, quantization: str): + can_implement, reason = is_supported_and_can_implement_kernel(kernel, config, None) + if not can_implement: + raise ValueError(f"Batch invariance on ROCm: {reason}") + linear_backend = _get_linear_backend(quantization=quantization) + if linear_backend not in ("auto", "triton"): + logger.warning_once( + "VLLM_BATCH_INVARIANT overrides --linear-backend=%s; " + "using the Triton backend for deterministic execution.", + linear_backend, + ) + return kernel + + # Kernel classes covered by each --linear-backend value. _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { "b12x": { @@ -723,6 +738,13 @@ def init_fp8_linear_kernel( ) if activation_quant_key.scale.group_shape.is_per_group(): + if envs.VLLM_BATCH_INVARIANT and current_platform.is_rocm(): + force_kernel = _batch_invariant_fp8_kernel( + TritonFp8BlockScaledMMKernel, + scaled_mm_linear_kernel_config, + "fp8_block_w8a8", + ) + kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS, # type: ignore[misc] @@ -755,6 +777,13 @@ def init_fp8_linear_kernel( ) else: + if envs.VLLM_BATCH_INVARIANT and current_platform.is_rocm(): + force_kernel = _batch_invariant_fp8_kernel( + TritonFP8ScaledMMLinearKernel, + scaled_mm_linear_kernel_config, + "fp8_w8a8", + ) + kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type] @@ -780,6 +809,19 @@ def init_fp8_linear_kernel( ) +def fp8_scaled_mm_is_batch_invariant( + kernel: FP8ScaledMMLinearKernel | Fp8BlockScaledMMLinearKernel, +) -> bool: + """Whether a row's output from the kernel is independent of the other rows.""" + if not isinstance( + kernel, (CutlassFP8ScaledMMLinearKernel, TritonFP8ScaledMMLinearKernel) + ): + return False + # A dynamic per-tensor activation scale is an amax over every row. + scale = kernel.config.activation_quant_key.scale + return scale.static or not scale.group_shape.is_per_tensor() + + def init_int8_linear_kernel( is_channelwise: bool, is_static_input_scheme: bool, @@ -1250,6 +1292,7 @@ def register_linear_kernel( __all__ = [ + "fp8_scaled_mm_is_batch_invariant", "init_fp8_linear_kernel", "init_int8_linear_kernel", "init_nvfp4_linear_kernel", diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py index 420e7f9493..4cd09c1f90 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/aiter.py +++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py @@ -62,7 +62,12 @@ if is_aiter_found_and_supported(): N = weight.shape[0] K = weight.shape[1] if rocm_use_aiter_fp4_asm_gemm: - if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K): + # The preshuffle kernel's tuned configs split K in some M buckets. + if ( + M <= 64 + and not envs.VLLM_BATCH_INVARIANT + and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K) + ): if x_scales is None: # use hip quant kernel for performance if M >= 32: @@ -116,7 +121,16 @@ if is_aiter_found_and_supported(): x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype ) - gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y) + config = None + if envs.VLLM_BATCH_INVARIANT: + # The configs split K in some M buckets (up to M <= 256 for + # tuned shapes). Keep the tile, drop the split. + from aiter.ops.triton.gemm_afp4wfp4 import _get_config + + config, _ = _get_config(M, N, K) + config = dict(config, NUM_KSPLIT=1) + + gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y, config) return y def gemm_with_dynamic_quant_fake( diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index faf1e180fa..98920f2df1 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -12,6 +12,7 @@ instead, since ``can_implement`` does not filter on K. import torch from torch.nn.parameter import Parameter +import vllm.envs as envs from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_BLOCK_SIZE, MXFP8_SCALE_DTYPE, @@ -246,7 +247,10 @@ class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel): N, K = weight.shape scale_k = K // MXFP8_BLOCK_SIZE weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() - block_scale = _as_block32_scale(weight_scale) + # rocm_mxfp8_block32_gemm's split-K varies with M. + block_scale = ( + None if envs.VLLM_BATCH_INVARIANT else _as_block32_scale(weight_scale) + ) if block_scale is not None: # Checkpoints with 32x32 scale blocks (DeepSeek V4/V4.1) get them # back from the loader expanded per row; keep the blocks and use diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index 75db4bcaa9..ac096b47b9 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py @@ -39,6 +39,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.triton import ( + TritonFP8ScaledMMLinearKernel, TritonInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.xpu import ( @@ -65,6 +66,7 @@ __all__ = [ "PerTensorTorchFP8ScaledMMLinearKernel", "RowWiseTorchFP8ScaledMMLinearKernel", "ROCmFP8ScaledMMLinearKernel", + "TritonFP8ScaledMMLinearKernel", "TritonInt8ScaledMMLinearKernel", "ZentorchInt8ScaledMMLinearKernel", "Fp8BlockScaledMMLinearKernel", diff --git a/vllm/model_executor/kernels/linear/scaled_mm/triton.py b/vllm/model_executor/kernels/linear/scaled_mm/triton.py index d6169d2cfc..fed471e474 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/triton.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/triton.py @@ -20,6 +20,8 @@ from .BlockScaledMMLinearKernel import ( ) from .cutlass import CutlassInt8ScaledMMLinearKernel from .ScaledMMLinearKernel import ( + FP8ScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, Int8ScaledMMLinearLayerConfig, ) @@ -156,6 +158,46 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel): return out +class TritonFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda_alike(): + return False, "requires ROCm or CUDA." + return True, None + + @classmethod + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: + act_group = c.activation_quant_key.scale.group_shape + weight_group = c.weight_quant_key.scale.group_shape + if not (act_group.is_per_tensor() or act_group.is_per_token()): + return False, "requires per-tensor or per-token activation scales." + if not ( + weight_group.is_per_tensor() + or weight_group.is_per_channel() + or weight_group.is_per_token() + ): + return False, "requires per-tensor, per-channel or per-token weight scales." + return True, None + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + out = triton_scaled_mm( + A, B, scale_a=As, scale_b=Bs, out_dtype=out_dtype, bias=bias + ) + return out.view(*output_shape) + + class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): @classmethod def is_supported(cls, compute_capability=None): diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py index c15a1a13d4..1efb058ef0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py @@ -51,6 +51,10 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase): def _supports_parallel_config(moe_parallel_config) -> bool: return True + @staticmethod + def _supports_batch_invariance() -> bool: + return False + @staticmethod def is_supported_config( cls, moe_config, weight_key, activation_key, activation_format diff --git a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py index da0bca4977..ed44bf4cb1 100644 --- a/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py @@ -218,7 +218,7 @@ class OCP_MXQuantizationEmulationTritonExperts(TritonExperts): global_num_experts=global_num_experts, expert_map=expert_map, a1q_scale=None, - a2_scale=None, + a2_scale=a2_scale, workspace13=workspace13, workspace2=workspace2, expert_tokens_meta=expert_tokens_meta, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py b/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py index e0d1fdf271..8adb41151c 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/triton_scaled_mm.py @@ -150,16 +150,16 @@ def scaled_mm_kernel( scale_b = scale_b.broadcast_to((BLOCK_SIZE_N, 1)) accumulator = scale_b.T * accumulator.to(tl.float32) - # Convert to output format. - c = accumulator.to(c_ptr.type.element_ty) - - # Add bias, it's already in output format, so add it after conversion. + # Add bias in accumulator dtype for less rounding error. if bias_ptr: offsets_bias = offsets_bn bias_ptrs = bias_ptr + offsets_bias bias_mask = offsets_bias < N bias = tl.load(bias_ptrs, bias_mask) - c += bias + accumulator += bias + + # Convert to output format. + c = accumulator.to(c_ptr.type.element_ty) # Save output offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) @@ -224,6 +224,8 @@ def triton_scaled_mm( tile_shape = (64, 128, 128) else: tile_shape = (128, 128, 128) + else: + tile_shape = (block_size_m, block_size_n, block_size_k) block_size_m, block_size_n, block_size_k = tile_shape diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 1c1d18614e..b18ec0a0aa 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -9,10 +9,10 @@ import vllm.envs as envs from vllm.config import get_current_vllm_config from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.kernels.linear import ( + fp8_scaled_mm_is_batch_invariant, init_fp8_linear_kernel, ) from vllm.model_executor.kernels.linear.scaled_mm import ( - CutlassFP8ScaledMMLinearKernel, MarlinFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.attention import Attention @@ -432,7 +432,7 @@ class Fp8LinearMethod(LinearMethodBase): bias, ) else: - if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel): + if fp8_scaled_mm_is_batch_invariant(self.fp8_linear): return self.fp8_linear.apply_weights(layer, x, bias) # per-tensor/channel: dequant to BF16 and run GEMM diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index f4662ad5a6..8d8a7eca60 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -15,9 +15,11 @@ if TYPE_CHECKING: import vllm.envs as envs from vllm import _custom_ops as ops from vllm.config import get_current_vllm_config -from vllm.model_executor.kernels.linear import init_fp8_linear_kernel +from vllm.model_executor.kernels.linear import ( + fp8_scaled_mm_is_batch_invariant, + init_fp8_linear_kernel, +) from vllm.model_executor.kernels.linear.scaled_mm import ( - CutlassFP8ScaledMMLinearKernel, MarlinFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.fused_moe import RoutedExperts @@ -278,7 +280,7 @@ class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): ) -> torch.Tensor: # if batch invariant mode is enabled, use BF16 dequant if envs.VLLM_BATCH_INVARIANT: - if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel): + if fp8_scaled_mm_is_batch_invariant(self.fp8_linear): return self.fp8_linear.apply_weights(layer, x, bias) weight_fp8 = layer.weight.to(torch.bfloat16) @@ -466,7 +468,7 @@ class Fp8PtpcOnlineLinearMethod(OnlineLinearBase): # if batch invariant mode is enabled dequant if ( envs.VLLM_BATCH_INVARIANT - and not isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel) + and not fp8_scaled_mm_is_batch_invariant(self.fp8_linear) and not isinstance(x, QuantizedActivation) ): weight_dequant = ( diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index fb87c3a614..00f9e5b92e 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -1041,7 +1041,10 @@ def w8a8_triton_block_scaled_mm( C_shape = A.shape[:-1] + (N,) C = A.new_empty(C_shape, dtype=output_dtype) - configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1]) + if envs.VLLM_BATCH_INVARIANT: + configs = None # use the default config + else: + configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1]) if configs: # Get the optimal config if there is one config = configs[min(configs.keys(), key=lambda x: abs(x - M))] diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 80eef7c75a..8d285dda75 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -288,6 +288,13 @@ def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, boo def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: + if envs.VLLM_BATCH_INVARIANT: + from vllm.model_executor.determinism.batch_invariant import ( + linear_batch_invariant, + ) + + return linear_batch_invariant(x, weight, bias) + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250 n = x.numel() // x.size(-1) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 5840ffe3b7..807e571002 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -1068,6 +1068,15 @@ class RocmPlatform(Platform): if use_aiter_fp8_linear and "-quant_fp8" not in compilation_config.custom_ops: compilation_config.custom_ops.append("+quant_fp8") + # Only RMSNorm's custom op reaches rms_norm_batch_invariant. Inductor + # lowers the native one to a reduction whose block size it picks by + # benchmarking, so the summation order varies with the batch and run. + if ( + envs.VLLM_BATCH_INVARIANT + and "-rms_norm" not in compilation_config.custom_ops + ): + compilation_config.custom_ops.append("+rms_norm") + if use_aiter_fused_se and "-grouped_topk" in compilation_config.custom_ops: logger.warning_once( "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled, which " diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py index 8b0e337bb2..ecc2a50135 100644 --- a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING import torch +import vllm.envs as envs from vllm.platforms import current_platform from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend @@ -47,6 +48,15 @@ class AiterFlashAttnPrefillBackend(MLAPrefillBackend): return rocm_aiter_ops.is_enabled() + @classmethod + def validate_configuration(cls, device_capability, selector_config): + invalid_reasons = super().validate_configuration( + device_capability, selector_config + ) + if envs.VLLM_BATCH_INVARIANT: + invalid_reasons.append("batch invariance not supported") + return invalid_reasons + def __init__( self, num_heads: int, diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 3af623b1f1..b6cf96e5b7 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -427,13 +427,13 @@ class FlashAttnPrefillBackend(MLAPrefillBackend): kwargs["return_softmax_lse"] = return_softmax_lse kwargs["out"] = out kwargs["output_scale"] = output_scale + if envs.VLLM_BATCH_INVARIANT: + kwargs["num_splits"] = 1 else: # ROCm leverages the upstream flash_attn, which takes a parameter # called "return_attn_probs" instead of return_softmax_lse kwargs["return_attn_probs"] = return_softmax_lse assert out is None and output_scale is None - if envs.VLLM_BATCH_INVARIANT: - kwargs["num_splits"] = 1 attn_out = _FA4_MLA_PREFILL_KERNEL( q=q,