[Feature] Support batch invariance on ROCm (#52231)

Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>

.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]

.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
+      device: mi300_1
-      timeout_in_minutes: 35
+      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

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

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
+        out = out + bias.float()
-
+    return out.to(out_dtype)
-    return out
 
 
 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()),

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

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)

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():
+    if not current_platform.is_cuda_alike():
-        pytest.skip("CUTLASS FP8 kernels require CUDA.")
+        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,
+        activation_quant_key=act_key,
-        weight_quant_key=kFp8StaticTensorSym,
+        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)

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, []

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.<op>(..., 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)

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)
+    )

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
+@skip_if_not_cuda_alike
-@pytest.mark.parametrize("backend", ["FLASH_ATTN"])
+@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)

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()

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.

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

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

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.
+    """Compute softmax, or log_softmax when ``LOG``, along the last dimension of a
-    Each block handles one row of the input tensor.
+    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) if LOG else 0.0
-    log_sum_exp = tl.log(sum_exp)
 
-    # 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
+        if LOG:
-        output = vals - max_val - log_sum_exp
+            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):
+def _addmm_fused_bias(bias, beta, alpha, n):
-    return mm_batch_invariant(a, b, bias=bias)
+    # 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.<op>(..., 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):
+def softmax_batch_invariant(input, dim, half_to_float=False):
-    # Compute softmax in a deterministic way
+    if half_to_float:
-    # First subtract max for numerical stability (standard practice)
+        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
+        bias = bias.expand(weight.shape[0])
-    return output
+    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_cuda_alike():
-        if current_platform.is_device_capability_family(80):
+        if current_platform.is_rocm() or current_platform.is_device_capability_family(
-            # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the
+            80
-            # triton persistent matmul overrides for mm/addmm/matmul/linear.
+        ):
+            # 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
+        _resolve_block_sizes()
-        # 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
     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,
+    if current_platform.is_rocm():
-    # unlike the pins below, it also applies to communicators shared with other
+        # The skinny GEMMs pick their split of K from the batch size.
-    # processes.
+        os.environ["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
-    os.environ["NCCL_LAUNCH_MODE"] = "GROUP"
+    else:
-    pin_nccl_env(
+        # NCCL determinism settings. NCCL keeps the launch mode process-wide, so,
-        {
+        # unlike the pins below, it also applies to communicators shared with other
-            "NCCL_COLLNET_ENABLE": "0",
+        # processes.
-            "NCCL_NVLS_ENABLE": "0",
+        os.environ["NCCL_LAUNCH_MODE"] = "GROUP"
-            "NCCL_P2P_NET_DISABLE": "1",
+        pin_nccl_env(
-            "NCCL_MIN_NCHANNELS": "1",
+            {
-            "NCCL_MAX_NCHANNELS": "1",
+                "NCCL_COLLNET_ENABLE": "0",
-            "NCCL_PROTO": "Simple",
+                "NCCL_NVLS_ENABLE": "0",
-            # NCCL >= 2.31 zero-fills the algorithm table of every collective
+                "NCCL_P2P_NET_DISABLE": "1",
-            # when NCCL_ALGO is set and re-enables only the named ones; together
+                "NCCL_MIN_NCHANNELS": "1",
-            # with the NCCL_PROTO above, collectives not named here end up with
+                "NCCL_MAX_NCHANNELS": "1",
-            # no algorithm and fail with ncclInvalidUsage. Re-enable Ring and
+                "NCCL_PROTO": "Simple",
-            # Tree for all collectives, then pin AllReduce to Tree.
+                # NCCL >= 2.31 zero-fills the algorithm table of every collective
-            "NCCL_ALGO": "ring,tree;allreduce:tree",
+                # when NCCL_ALGO is set and re-enables only the named ones; together
-            "NCCL_NTHREADS": "1",
+                # with the NCCL_PROTO above, collectives not named here end up with
-            "NCCL_SOCKET_NTHREADS": "1",
+                # 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"

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",

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(

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

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",

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):

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

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,

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.
+    # Add bias in accumulator dtype for less rounding error.
-    c = accumulator.to(c_ptr.type.element_ty)
-
-    # Add bias, it's already in output format, so add it after conversion.
     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
 

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

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 = (

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))]

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)

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 "

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,

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,