[Kernel][LoRA] Add deterministic split-K=8 LoRA shrink for batch invariance

.buildkite/test_areas/misc.yaml

@@ -655,12 +655,14 @@ steps:
   source_file_dependencies:
     - vllm/v1/attention
     - vllm/model_executor/layers
+    - vllm/lora/ops/triton_ops
     - 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
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+    - pytest -v -s v1/determinism/test_lora_shrink_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 FLASH_ATTN
   mirror:
@@ -690,12 +692,14 @@ steps:
   source_file_dependencies:
     - vllm/v1/attention
     - vllm/model_executor/layers
+    - vllm/lora/ops/triton_ops
     - 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
     - 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 FLASH_ATTN

docs/features/batch_invariance.md

@@ -152,6 +152,7 @@ When batch invariance is enabled, vLLM:
 5. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell),
    runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out
    (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes).
+6. Splits the LoRA shrink reduction into 8 fixed K chunks that are summed in a fixed order instead of with atomics
 
 !!! warning
     Batch invariance under tensor parallelism is not yet supported for all-reduces whose size isn't a multiple of 16 bytes, for example a hidden size that isn't a multiple of 8 in fp16/bf16. Such tensors can switch between NCCL and custom all-reduce depending on batch size. All validated models meet this requirement.

tests/v1/determinism/test_lora_shrink_batch_invariant.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Batch-invariant LoRA shrink: fixed split-K partials, fixed-order reduction."""
+
+import pytest
+import torch
+from utils import skip_if_not_cuda
+
+import vllm.lora.ops.triton_ops.lora_shrink_op as shrink_op
+import vllm.lora.ops.triton_ops.utils as lora_utils
+from vllm.lora.ops.triton_ops import LoRAKernelMeta, lora_shrink
+from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT
+from vllm.utils.torch_utils import set_random_seed
+
+DEVICE = "cuda"
+
+
+@pytest.fixture(autouse=True)
+def batch_invariant(monkeypatch):
+    monkeypatch.setattr(lora_utils, "is_batch_invariant", True)
+    monkeypatch.setattr(shrink_op, "is_batch_invariant", True)
+    lora_utils.get_lora_op_configs.cache_clear()
+    yield
+    lora_utils.get_lora_op_configs.cache_clear()
+
+
+def _shrink(inputs, weights, mapping, scaling):
+    num_tokens = inputs.size(0)
+    meta = LoRAKernelMeta.make(weights[0].size(0), num_tokens, device=DEVICE)
+    meta.prepare_tensors(mapping)
+    out = torch.zeros(
+        (len(weights), num_tokens, weights[0].size(1)),
+        dtype=torch.float32,
+        device=DEVICE,
+    )
+    _LORA_A_PTR_DICT.clear()
+    args = meta.meta_args(token_nums=num_tokens, specialize_active_lora=False)
+    lora_shrink(inputs, weights, out, *args, scaling)
+    return out
+
+
+@skip_if_not_cuda
+@pytest.mark.parametrize(
+    "hidden_size,rank,nslices,dtype",
+    [
+        (2560, 8, 3, torch.bfloat16),
+        (9728, 16, 1, torch.float16),
+        (4097, 64, 1, torch.bfloat16),  # K tail block
+        (1024, 8, 1, torch.float16),  # K < 8 * BLOCK_K: some splits are empty
+    ],
+)
+def test_lora_shrink_matches_reference(hidden_size, rank, nslices, dtype):
+    set_random_seed(0)
+    inputs = torch.randn((33, hidden_size), dtype=dtype, device=DEVICE)
+    weights = [
+        torch.randn((2, rank, hidden_size), dtype=dtype, device=DEVICE)
+        for _ in range(nslices)
+    ]
+    mapping = torch.randint(-1, 2, (33,), dtype=torch.int32, device=DEVICE)
+    out = _shrink(inputs, weights, mapping, 0.5).double().cpu()
+
+    ids, x = mapping.cpu(), inputs.double().cpu()
+    ref = torch.zeros_like(out)
+    for s, weight in enumerate(weights):
+        for lora_id in (0, 1):
+            rows = ids == lora_id
+            ref[s, rows] = 0.5 * x[rows] @ weight[lora_id].double().cpu().T
+    torch.testing.assert_close(out, ref, rtol=5e-3, atol=5e-3)
+
+
+@skip_if_not_cuda
+@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, []
+
+    class RecordReduce:
+        def __getitem__(self, grid):
+            reduce_grids.append(grid)
+            return reduce_kernel[grid]
+
+    monkeypatch.setattr(shrink_op, "_lora_shrink_reduce_kernel", RecordReduce())
+    set_random_seed(0)
+    # K = 2560 is ten BLOCK_K=256 blocks, so all 8 splits contribute.
+    dtype, hidden_size = torch.bfloat16, 2560
+    weights = [
+        torch.randn((3, 16, hidden_size), dtype=dtype, device=DEVICE)
+        for _ in range(nslices)
+    ]
+    target = torch.randn((1, hidden_size), dtype=dtype, device=DEVICE)
+    target_id = torch.zeros(1, dtype=torch.int32, device=DEVICE)
+    expected = _shrink(target, weights, target_id, 0.7)[:, 0]
+    for num_tokens in (31, 32, 33, 127, 128, 129):
+        inputs = torch.randn((num_tokens, hidden_size), dtype=dtype, device=DEVICE)
+        mapping = torch.randint(-1, 3, (num_tokens,), dtype=torch.int32, device=DEVICE)
+        pos = num_tokens // 2
+        inputs[pos], mapping[pos] = target[0], 0
+        out = _shrink(inputs, weights, mapping, 0.7)
+        assert torch.equal(out[:, pos], expected), num_tokens
+    assert len(reduce_grids) == 7

vllm/lora/ops/triton_ops/kernel_utils.py

@@ -257,6 +257,7 @@ def do_shrink_kernel(
     output_d0_stride,
     output_d1_stride,
     output_d2_stride,
+    output_split_stride,
     scaling,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
@@ -265,6 +266,7 @@ def do_shrink_kernel(
     SPLIT_K: tl.constexpr,
     SLICE_NUM: tl.constexpr,
     USE_GDC: tl.constexpr,
+    STORE_PARTIALS: tl.constexpr,
 ):
     """Given an array of integers that identifies the rows of A, ram,
     a lora index that identifies which LoRA to use from lora_ptr, lora_index,
@@ -324,17 +326,31 @@ def do_shrink_kernel(
     # Identify the C output pointers to store the results of the accumulator.
     offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
     offset_cm = tl.arange(0, BLOCK_M)
-    cur_out_ptr = out_ptr if SLICE_NUM == 1 else out_ptr + slice_id * output_d0_stride
+    if STORE_PARTIALS:
+        cur_out_ptr = (
+            out_ptr
+            + slice_id.to(tl.int64) * output_d0_stride
+            + pid_sk.to(tl.int64) * output_split_stride
+        )
+    else:
+        cur_out_ptr = (
+            out_ptr if SLICE_NUM == 1 else out_ptr + slice_id * output_d0_stride
+        )
     c_ptr = (
         cur_out_ptr
         + ram[:, None] * output_d1_stride
         + offset_cn[None, :] * output_d2_stride
     )
     c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < N)
-    accumulator *= scaling
 
-    # handles write-back with reduction-splitting
-    if SPLIT_K == 1:
+    if STORE_PARTIALS:
+        # Unscaled FP32 partial per K split; the reducer adds the splits in a
+        # fixed order and applies scaling once.
         tl.store(c_ptr, accumulator, mask=c_mask)
     else:
-        tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed")
+        accumulator *= scaling
+        # handles write-back with reduction-splitting
+        if SPLIT_K == 1:
+            tl.store(c_ptr, accumulator, mask=c_mask)
+        else:
+            tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed")