From b6cbba8bc893c61e412a205533aafbee1ae6be31 Mon Sep 17 00:00:00 2001 From: oops-oom Date: Tue, 28 Jul 2026 22:24:02 +0800 Subject: [PATCH] [Bugfix][Kernel] Fix batch invariance in RMSNorm kernels by pinning block size (#48391) Signed-off-by: oops-oom <73481342@qq.com> Signed-off-by: oops-oom Co-authored-by: oops-oom <73481342@qq.com> Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Shengqi Chen --- .buildkite/test_areas/misc.yaml | 14 +- csrc/libtorch_stable/layernorm_kernels.cu | 14 +- .../layernorm_quant_kernels.cu | 9 +- ...fused_layernorm_dynamic_per_token_quant.cu | 5 +- tests/v1/determinism/test_batch_invariance.py | 18 +++ .../test_rms_norm_batch_invariant.py | 150 +++++++++++++++++- 6 files changed, 189 insertions(+), 21 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 763fbcbfec2..caa56c21b37 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.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 diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index 878b44df936..7a1051d2c00 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/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(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(weight->data_ptr()); diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/libtorch_stable/layernorm_quant_kernels.cu index f3bf8882e77..f43be531de0 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/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(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); diff --git a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu b/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu index 2152e64dc96..56dd4703872 100644 --- a/csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu +++ b/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()); diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index b2706ed89b7..37fd5cba6a5 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/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, ) diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index dfd08351277..232a43b1f98 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/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])