[Bugfix][Kernel] Fix batch invariance in RMSNorm kernels by pinning block size (#48391)

Signed-off-by: oops-oom <[email protected]>
Signed-off-by: oops-oom <[email protected]>
Co-authored-by: oops-oom <[email protected]>
Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
Co-authored-by: Wentao Ye <[email protected]>
Co-authored-by: Shengqi Chen <[email protected]>

.buildkite/test_areas/misc.yaml

@@ -398,7 +398,7 @@ steps:
 
 - label: Batch Invariance (A100)
   key: batch-invariance-a100
-  timeout_in_minutes: 40
+  timeout_in_minutes: 60
   device: a100
   source_file_dependencies:
     - vllm/v1/attention
@@ -408,11 +408,11 @@ steps:
     - 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[TRITON_MLA]
+    - 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: Batch Invariance (H100)
   key: batch-invariance-h100
-  timeout_in_minutes: 40
+  timeout_in_minutes: 60
   device: h100
   source_file_dependencies:
     - vllm/v1/attention
@@ -423,8 +423,8 @@ 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
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
-    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
+    - 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 FLASH_ATTN
 
 - label: Batch Invariance (B200)
   key: batch-invariance-b200
@@ -439,8 +439,8 @@ 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
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
-    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
+    - 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 FLASH_ATTN
     - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
     - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
     - pytest -v -s v1/determinism/test_matmul_batch_invariant.py

csrc/libtorch_stable/layernorm_kernels.cu

@@ -249,7 +249,9 @@ void rms_norm(torch::stable::Tensor& out,    // [..., hidden_size]
   int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;
 
   // For large num_tokens, use smaller blocks to increase SM concurrency.
-  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
+  const int max_block_size =
+      batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
   dim3 grid(num_tokens);
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
@@ -325,8 +327,13 @@ void fused_add_rms_norm(torch::stable::Tensor& input,     // [..., hidden_size]
   /* This kernel is memory-latency bound in many scenarios.
      When num_tokens is large, a smaller block size allows
      for increased block occupancy on CUs and better latency
-     hiding on global mem ops. */
-  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+     hiding on global mem ops. In batch-invariant mode the block size must
+     not depend on num_tokens, otherwise the same token would use a different
+     reduction width (and thus a different floating-point summation order)
+     across batches; lock it to 1024 to keep results bit-exact. */
+  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
+  const int max_block_size =
+      batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
   dim3 block(std::min(hidden_size, max_block_size));
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
@@ -337,7 +344,6 @@ void fused_add_rms_norm(torch::stable::Tensor& input,     // [..., hidden_size]
   auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
   bool offsets_are_multiple_of_vector_width =
       hidden_size % vector_width == 0 && input_stride % vector_width == 0;
-  bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
   const bool has_weight = weight.has_value();
   if (has_weight) {
     auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight->data_ptr());

csrc/libtorch_stable/layernorm_quant_kernels.cu

@@ -215,7 +215,9 @@ void rms_norm_static_fp8_quant(
   int num_tokens = input.numel() / hidden_size;
 
   // For large num_tokens, use smaller blocks to increase SM concurrency.
-  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
+  const int max_block_size =
+      batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
   dim3 grid(num_tokens);
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
@@ -279,7 +281,9 @@ void fused_add_rms_norm_static_fp8_quant(
      When num_tokens is large, a smaller block size allows
      for increased block occupancy on CUs and better latency
      hiding on global mem ops. */
-  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
+  const int max_block_size =
+      batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256);
   dim3 block(std::min(hidden_size, max_block_size));
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
@@ -296,7 +300,6 @@ void fused_add_rms_norm_static_fp8_quant(
   auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
   bool ptrs_are_aligned =
       inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;
-  bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
   if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0 &&
       !batch_invariant_launch) {
     LAUNCH_FUSED_ADD_RMS_NORM(8);

csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

@@ -2,6 +2,7 @@
 #include "../../torch_utils.h"
 
 #include "../../dispatch_utils.h"
+#include "../../../core/batch_invariant.hpp"
 #include "layernorm_utils.cuh"
 #include "quant_conversions.cuh"
 
@@ -231,7 +232,9 @@ void rms_norm_per_block_quant_dispatch(
   auto num_tokens = input.numel() / hidden_size;
 
   dim3 grid(num_tokens);
-  const int max_block_size = (num_tokens <= 256) ? 512 : 256;
+  const bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
+  const int max_block_size =
+      batch_invariant_launch ? 512 : ((num_tokens <= 256) ? 512 : 256);
   dim3 block(std::min(hidden_size, max_block_size));
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());

tests/v1/determinism/test_batch_invariance.py

@@ -27,8 +27,10 @@ from vllm.platforms import current_platform
     "backend",
     BACKENDS,
 )
+@pytest.mark.parametrize("rms_norm_impl", ["default", "vllm_c"])
 def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
     backend,
+    rms_norm_impl,
 ):
     """
     Ensures that the same request (the 'needle' prompt) yields identical output
@@ -60,6 +62,16 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
     random.seed(seed)
 
     attention_config = {"backend": backend}
+    # Force the C++ RMSNorm implementation so we actually exercise the
+    # num_tokens-dependent block-size branches.
+    kernel_config = None
+    if rms_norm_impl == "vllm_c":
+        kernel_config = {
+            "ir_op_priority": {
+                "rms_norm": ["vllm_c"],
+                "fused_add_rms_norm": ["vllm_c"],
+            }
+        }
     # Allow overrides from environment (useful for CI tuning)
     # "facebook/opt-125m" is too small, doesn't reliably test determinism
     model = TEST_MODEL
@@ -96,6 +108,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
             gpu_memory_utilization=gpu_mem_util,
             max_model_len=max_model_len,
             attention_config=attention_config,
+            kernel_config=kernel_config,
         )
 
         # Baseline generation for the needle prompt alone.
@@ -923,11 +936,15 @@ def LLM_with_max_seqs(
     gpu_memory_utilization: float,
     max_model_len: int,
     attention_config: dict | None = None,
+    kernel_config: dict | None = None,
 ) -> LLM:
     """
     Helper to construct an LLM with a specific max_num_seqs (batch-size limit)
     using the high-level v1 LLM API, while constraining memory usage.
     """
+    extra_kwargs: dict = {}
+    if kernel_config is not None:
+        extra_kwargs["kernel_config"] = kernel_config
     return LLM(
         model=model,
         max_num_seqs=max_num_seqs,
@@ -939,4 +956,5 @@ def LLM_with_max_seqs(
         attention_config=attention_config,
         # Enable for MOE models
         # enable_expert_parallel=True,
+        **extra_kwargs,
     )

tests/v1/determinism/test_rms_norm_batch_invariant.py

@@ -28,16 +28,18 @@ def _rms_norm_reference(
 
 
 @skip_if_not_cuda
-@pytest.mark.parametrize("batch_size", [1, 4, 16, 64])
+@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])
 @pytest.mark.parametrize("eps", [1e-6, 1e-5])
+@pytest.mark.parametrize("seed", list(range(4)))
 def test_rms_norm_batch_invariant_vs_reference(
     default_vllm_config,
     batch_size: int,
     hidden_size: int,
     dtype: torch.dtype,
     eps: float,
+    seed: int,
 ):
     """
     Compare batch-invariant Triton RMS norm against a PyTorch reference.
@@ -48,7 +50,7 @@ def test_rms_norm_batch_invariant_vs_reference(
     device = torch.device(DEVICE_TYPE)
 
     # Create test input and weight
-    torch.manual_seed(42)
+    torch.manual_seed(seed)
     input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
     weight = torch.randn(hidden_size, dtype=dtype, device=device)
 
@@ -71,7 +73,7 @@ def test_rms_norm_batch_invariant_vs_reference(
         atol=atol,
         msg=f"RMS norm mismatch for batch_size={batch_size}, "
         f"hidden_size={hidden_size}, "
-        f"dtype={dtype}, eps={eps}",
+        f"dtype={dtype}, eps={eps}, seed={seed}",
     )
 
 
@@ -79,17 +81,21 @@ def test_rms_norm_batch_invariant_vs_reference(
 @pytest.mark.parametrize("hidden_size", [512, 4096])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 @pytest.mark.parametrize("eps", [1e-6])
+@pytest.mark.parametrize("n_extra", [3, 299])
+@pytest.mark.parametrize("seed", list(range(16)))
 def test_fused_add_rms_norm_batch_invariant_residual_path(
     hidden_size: int,
     dtype: torch.dtype,
     eps: float,
+    n_extra: int,
+    seed: int,
 ):
     """
     Test the batch-invariant fused residual-add + RMSNorm helper directly.
     """
     device = torch.device(DEVICE_TYPE)
 
-    torch.manual_seed(42)
+    torch.manual_seed(seed)
     x_single = torch.randn(1, hidden_size, dtype=dtype, device=device)
     residual_single = torch.randn(1, hidden_size, dtype=dtype, device=device)
     weight = torch.randn(hidden_size, dtype=dtype, device=device)
@@ -97,14 +103,14 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(
     x_batch = torch.cat(
         [
             x_single,
-            torch.randn(3, hidden_size, dtype=dtype, device=device),
+            torch.randn(n_extra, hidden_size, dtype=dtype, device=device),
         ],
         dim=0,
     )
     residual_batch = torch.cat(
         [
             residual_single,
-            torch.randn(3, hidden_size, dtype=dtype, device=device),
+            torch.randn(n_extra, hidden_size, dtype=dtype, device=device),
         ],
         dim=0,
     )
@@ -168,6 +174,138 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(
     )
 
 
+FP8_DTYPE = current_platform.fp8_dtype()
+
+# The large launch (num_tokens=300 >= 256) drops an un-pinned kernel to block
+# 256, while the small launch (255 rows) stays under the threshold and keeps the
+# larger block (1024, or 512 for per-block quant). Under the pin the two launches
+# use the same block, so the shared first 255 rows must match bit-for-bit; 255 is
+# the most rows a single small launch can hold (< 256, and <= 256 for per-block).
+_LARGE_TOKENS = 300
+_SMALL_TOKENS = 255
+
+
+def _assert_rows_bit_identical(small, large, msg):
+    if small.dtype == FP8_DTYPE:
+        assert torch.equal(small.view(torch.uint8), large.view(torch.uint8)), msg
+    else:
+        torch.testing.assert_close(small, large, rtol=0.0, atol=0.0, msg=msg)
+
+
+@skip_if_not_cuda
+@pytest.mark.parametrize("hidden_size", [512, 4096])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("seed", list(range(4)))
+def test_rms_norm_batch_invariant_nonresidual_kernel(
+    hidden_size: int, dtype: torch.dtype, seed: int
+):
+    """C++ ``rms_norm`` (no residual) must be batch invariant across the block
+    threshold. Reached in compiled mode with ``ir_op_priority.rms_norm=["vllm_c"]``
+    (default priority is ``native``/inductor codegen when compiling).
+    """
+    import vllm._custom_ops as ops
+
+    device = torch.device(DEVICE_TYPE)
+    torch.manual_seed(seed)
+    rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device)
+    weight = torch.randn(hidden_size, dtype=dtype, device=device)
+
+    def rms_norm(x):
+        out = torch.empty_like(x)
+        ops.rms_norm(out, x, weight, 1e-6)
+        return out
+
+    large = rms_norm(rows.clone())
+    small = rms_norm(rows[:_SMALL_TOKENS].clone())
+    _assert_rows_bit_identical(
+        small,
+        large[:_SMALL_TOKENS],
+        "rms_norm output depends on num_tokens (block size)",
+    )
+
+
+@skip_if_not_cuda
+@pytest.mark.parametrize("hidden_size", [512, 4096])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("seed", list(range(4)))
+@pytest.mark.parametrize("add_residual", [False, True])
+def test_rms_norm_static_fp8_quant_batch_invariant(
+    hidden_size: int, dtype: torch.dtype, seed: int, add_residual: bool
+):
+    """C++ static per-tensor fp8-quant RMSNorm must be batch invariant across
+    the block threshold. Covers ``rms_norm_static_fp8_quant`` and, with
+    ``add_residual``, ``fused_add_rms_norm_static_fp8_quant`` (the compiled fp8
+    path where ``RMSNormQuantFusionPass`` rewrites norm + quant into them).
+    """
+    device = torch.device(DEVICE_TYPE)
+    torch.manual_seed(seed)
+    rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device)
+    residual = (
+        torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device)
+        if add_residual
+        else None
+    )
+    weight = torch.randn(hidden_size, dtype=dtype, device=device)
+    quant_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
+
+    def quant(x, res):
+        out = torch.empty_like(x, dtype=FP8_DTYPE)
+        if add_residual:
+            torch.ops._C.fused_add_rms_norm_static_fp8_quant(
+                out, x, res, weight, quant_scale, 1e-6
+            )
+        else:
+            torch.ops._C.rms_norm_static_fp8_quant(out, x, weight, quant_scale, 1e-6)
+        return out
+
+    large = quant(rows.clone(), residual.clone() if residual is not None else None)
+    small = quant(
+        rows[:_SMALL_TOKENS].clone(),
+        residual[:_SMALL_TOKENS].clone() if residual is not None else None,
+    )
+    _assert_rows_bit_identical(
+        small,
+        large[:_SMALL_TOKENS],
+        "static-fp8-quant RMSNorm output depends on num_tokens (block size)",
+    )
+
+
+@skip_if_not_cuda
+@pytest.mark.parametrize("hidden_size", [512, 4096])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("seed", list(range(4)))
+def test_rms_norm_per_block_quant_batch_invariant(
+    hidden_size: int, dtype: torch.dtype, seed: int
+):
+    """C++ ``rms_norm_per_block_quant`` must be batch invariant across the
+    block threshold (compiled fp8 block-quant path; block pinned to 512)."""
+    import vllm._custom_ops as ops
+
+    device = torch.device(DEVICE_TYPE)
+    torch.manual_seed(seed)
+    rows = torch.randn(_LARGE_TOKENS, hidden_size, dtype=dtype, device=device)
+    weight = torch.randn(hidden_size, dtype=dtype, device=device)
+    group_size = [1, 128]
+
+    def per_block_quant(x):
+        return ops.rms_norm_per_block_quant(x, weight, 1e-6, FP8_DTYPE, group_size)
+
+    out_large, scale_large = per_block_quant(rows.clone())
+    out_small, scale_small = per_block_quant(rows[:_SMALL_TOKENS].clone())
+    _assert_rows_bit_identical(
+        out_small,
+        out_large[:_SMALL_TOKENS],
+        "rms_norm_per_block_quant output depends on num_tokens (block size)",
+    )
+    torch.testing.assert_close(
+        scale_small,
+        scale_large[:_SMALL_TOKENS],
+        rtol=0.0,
+        atol=0.0,
+        msg="rms_norm_per_block_quant scales depend on num_tokens (block size)",
+    )
+
+
 @skip_if_not_cuda
 @pytest.mark.parametrize("batch_size", [1, 16, 128])
 @pytest.mark.parametrize("seq_len", [1, 32, 512])