From fe7c4b968080837a55509f4f6f8c276628d1530d Mon Sep 17 00:00:00 2001
From: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com>
Date: Tue, 29 Sep 2026 21:51:20 -0700
Subject: [PATCH 1/3] [Kernel][LoRA] Add deterministic split-K=8 LoRA shrink
 for batch invariance

Under VLLM_BATCH_INVARIANT=1 the LoRA shrink kernel runs with SPLIT_K=1
because split-K normally reduces with float atomics. The opt-in
VLLM_LORA_DETERMINISTIC_SPLIT_K=8 keeps K parallelism without atomics:
_lora_shrink_kernel stores unscaled FP32 partials for 8 fixed K splits
(BLOCK_K=256 for every M) through a STORE_PARTIALS branch in
do_shrink_kernel, and a reducer adds the splits in a fixed order and
applies scaling once. Partials are allocated per call with torch.empty.

The default path (split-K 0, BI on or off) launches exactly as before.
S=8 is rejected at import without VLLM_BATCH_INVARIANT=1 or together
with VLLM_LORA_ENABLE_DUAL_STREAM=1.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com>
---
 docs/features/batch_invariance.md          |  20 ++
 tests/lora/test_lora_shrink_splitk.py      | 338 +++++++++++++++++++++
 tests/test_envs.py                         |  15 +
 vllm/envs.py                               |  21 ++
 vllm/lora/ops/triton_ops/kernel_utils.py   |  26 +-
 vllm/lora/ops/triton_ops/lora_shrink_op.py | 159 +++++++++-
 6 files changed, 570 insertions(+), 9 deletions(-)
 create mode 100644 tests/lora/test_lora_shrink_splitk.py

diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md
index 789878f48a96..5b159c0a9e68 100644
--- a/docs/features/batch_invariance.md
+++ b/docs/features/batch_invariance.md
@@ -138,6 +138,26 @@ Batch invariance has been tested and verified on the following models:
 
 Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose).
 
+## LoRA Deterministic Split-K (Experimental)
+
+In batch-invariant mode the LoRA shrink kernel does not split the K
+dimension. `VLLM_LORA_DETERMINISTIC_SPLIT_K=8` instead writes FP32 partial
+sums for 8 fixed K splits and adds them in a fixed order, without atomics.
+
+```bash
+VLLM_BATCH_INVARIANT=1 VLLM_LORA_DETERMINISTIC_SPLIT_K=8 vllm serve <model> \
+    --enable-lora --lora-modules ...
+```
+
+- Requires `VLLM_BATCH_INVARIANT=1`; accepted values are `0` (default) and `8`.
+- Cannot be combined with `VLLM_LORA_ENABLE_DUAL_STREAM=1`.
+- Tested on CUDA dense LoRA with BF16/FP16 inputs on SM89, SM90, SM100 and SM120 GPUs.
+
+!!! warning
+    Split-K=8 is batch-invariant on its own but uses a different reduction
+    order than the default path, so generated tokens and logprobs can differ
+    between the two settings.
+
 ## Implementation Details
 
 When batch invariance is enabled, vLLM:
diff --git a/tests/lora/test_lora_shrink_splitk.py b/tests/lora/test_lora_shrink_splitk.py
new file mode 100644
index 000000000000..eba2546abd0f
--- /dev/null
+++ b/tests/lora/test_lora_shrink_splitk.py
@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the LoRA shrink deterministic split-K path.
+
+The shrink op reads VLLM_LORA_DETERMINISTIC_SPLIT_K and VLLM_BATCH_INVARIANT
+once at import, so every case runs in a fresh interpreter whose environment
+is set before start-up.
+"""
+
+import os
+import subprocess
+import sys
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="LoRA shrink kernel requires CUDA"
+)
+
+S8 = {
+    "VLLM_BATCH_INVARIANT": "1",
+    "VLLM_LORA_DETERMINISTIC_SPLIT_K": "8",
+    "VLLM_LORA_ENABLE_DUAL_STREAM": "0",
+}
+S1 = {**S8, "VLLM_LORA_DETERMINISTIC_SPLIT_K": "0"}
+BI_OFF = {**S1, "VLLM_BATCH_INVARIANT": "0"}
+DEVICE = "cuda:0"
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+
+
+@pytest.fixture(autouse=True)
+def cleanup_fixture():
+    """GPU work runs in subprocesses; skip conftest teardown (needs ray)."""
+    yield
+
+
+@pytest.fixture(autouse=True)
+def dynamo_reset():
+    yield
+
+
+def _subprocess(env: dict[str, str], code: str) -> subprocess.CompletedProcess:
+    return subprocess.run(
+        [sys.executable, "-c", code],
+        env={**os.environ, **env},
+        cwd=REPO_ROOT,
+        capture_output=True,
+        text=True,
+        timeout=600,
+    )
+
+
+def _run_case(env: dict[str, str], case: str, *args) -> None:
+    code = f"import runpy; runpy.run_path({__file__!r})[{case!r}](*{args!r})"
+    proc = _subprocess(env, code)
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+
+
+def _lora_data(m, k, rank, nslices, num_loras, dtype, seed):
+    g = torch.Generator(device=DEVICE).manual_seed(seed)
+    inputs = torch.randn((m, k), dtype=dtype, device=DEVICE, generator=g)
+    weights = [
+        torch.randn((num_loras, rank, k), dtype=dtype, device=DEVICE, generator=g)
+        for _ in range(nslices)
+    ]
+    mapping = torch.randint(
+        0, num_loras, (m,), dtype=torch.int32, device=DEVICE, generator=g
+    )
+    return inputs, weights, mapping
+
+
+def _meta(mapping, max_loras, meta=None):
+    from vllm.lora.ops.triton_ops.lora_kernel_metadata import LoRAKernelMeta
+
+    if meta is None:
+        meta = LoRAKernelMeta.make(max_loras, mapping.numel(), device=DEVICE)
+    meta.prepare_tensors(mapping)
+    return meta
+
+
+def _new_out(weights, m, fill=0.0):
+    shape = (len(weights), m, weights[0].size(1))
+    return torch.full(shape, fill, dtype=torch.float32, device=DEVICE)
+
+
+def _shrink(inputs, weights, meta, out, scaling):
+    from vllm.lora.ops.triton_ops import lora_shrink
+    from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT
+
+    _LORA_A_PTR_DICT.clear()
+    args = meta.meta_args(token_nums=inputs.size(0), specialize_active_lora=False)
+    lora_shrink(inputs, weights, out, *args, scaling)
+    return out
+
+
+def _assert_matches_fp64(out, inputs, weights, mapping, scaling):
+    x, ids = inputs.double().cpu(), mapping.cpu()
+    ref = torch.zeros(out.shape, dtype=torch.float64)
+    for s, w in enumerate(weights):
+        for lid in ids[ids >= 0].unique().tolist():
+            rows = ids == lid
+            ref[s, rows] = scaling * (x[rows] @ w[lid].double().cpu().T)
+    out = out.double().cpu()
+    assert torch.isfinite(out).all()
+    assert torch.all(out[:, ids < 0] == 0)
+    torch.testing.assert_close(out, ref, rtol=0.005, atol=0.005)
+
+
+class _LaunchRecorder:
+    def __init__(self, name, kernel, sink):
+        self.name, self.kernel, self.sink = name, kernel, sink
+
+    def __getitem__(self, grid):
+        launch = self.kernel[grid]
+
+        def record(*args, **kwargs):
+            bound = dict(zip(self.kernel.arg_names, args))
+            self.sink.append((self.name, {**bound, **kwargs}))
+            return launch(*args, **kwargs)
+
+        return record
+
+
+def _case_numerics(k, rank, nslices, dtype, scaling):
+    dtype = getattr(torch, dtype)
+    inputs, weights, mapping = _lora_data(17, k, rank, nslices, 2, dtype, seed=0)
+    out = _shrink(inputs, weights, _meta(mapping, 2), _new_out(weights, 17), scaling)
+    _assert_matches_fp64(out, inputs, weights, mapping, scaling)
+
+
+def _case_dispatch(split_k):
+    import vllm.lora.ops.triton_ops.lora_shrink_op as op
+
+    launches: list = []
+    for name in ("_lora_shrink_kernel", "_lora_shrink_reduce_kernel"):
+        setattr(op, name, _LaunchRecorder(name, getattr(op, name), launches))
+
+    inputs, weights, mapping = _lora_data(16, 2049, 32, 1, 4, torch.bfloat16, 1)
+    out = _shrink(inputs, weights, _meta(mapping, 4), _new_out(weights, 16), 0.5)
+    _assert_matches_fp64(out, inputs, weights, mapping, 0.5)
+
+    names = [name for name, _ in launches]
+    if not split_k:
+        assert names == ["_lora_shrink_kernel"], names
+        assert not launches[0][1].get("STORE_PARTIALS", False)
+        return
+    assert names == ["_lora_shrink_kernel", "_lora_shrink_reduce_kernel"], names
+    (_, first), (_, reduce) = launches
+    assert first["STORE_PARTIALS"] and first["out_ptr"].dtype == torch.float32
+    launch = dict(BLOCK_M=32, BLOCK_N=16, SPLIT_K=8, num_warps=4, num_stages=2)
+    assert {key: first[key] for key in launch} == launch
+    assert (first["BLOCK_K"], first["num_ctas"]) == (256, 1)
+    assert {key: reduce[key] for key in launch} == launch
+
+
+def _case_batch_invariance(nslices):
+    from vllm.lora.ops.triton_ops import lora_expand
+    from vllm.lora.ops.triton_ops.utils import _LORA_B_PTR_DICT
+
+    # K=2560 is ten BK=256 blocks, so every one of the eight splits is used.
+    hidden, rank, num_loras, scaling = 2560, 16, 3, 0.7
+    g = torch.Generator(device=DEVICE).manual_seed(1234 + nslices)
+
+    def randn(*shape):
+        return torch.randn(shape, dtype=torch.bfloat16, device=DEVICE, generator=g)
+
+    lora_a = [randn(num_loras, rank, hidden) for _ in range(nslices)]
+    lora_b = [randn(num_loras, hidden, rank) for _ in range(nslices)]
+    target, target_residual = randn(hidden), randn(hidden * nslices)
+
+    def run(m, pos, others):
+        ids = torch.tensor(others[:pos] + [0] + others[pos:], device=DEVICE)
+        inputs, residual = randn(m, hidden), randn(m, hidden * nslices)
+        inputs[pos], residual[pos] = target, target_residual
+        meta = _meta(ids.to(torch.int32), num_loras)
+        shrunk = _shrink(inputs, lora_a, meta, _new_out(lora_a, m), scaling)
+        _LORA_B_PTR_DICT.clear()
+        args = meta.meta_args(token_nums=m, specialize_active_lora=False)
+        lora_expand(shrunk, lora_b, residual, *args, add_inputs=True)
+        return shrunk[:, pos].clone(), residual[pos].clone()
+
+    ref_shrunk, ref_expanded = run(1, 0, [])
+    for m, pos, others in [  # (M, target row, adapters of the other rows)
+        (127, 0, [0] * 126),
+        (127, 126, [1, 2] * 63),
+        (128, 0, [1] * 127),
+        (129, 128, [2] * 128),
+        (31, 15, [0, 1, 2] * 10),
+        (32, 16, [0, 1, 2] * 10 + [1]),
+        (33, 16, [0, 1, 2] * 10 + [1, 2]),
+        (64, 40, ([0, 1, 2] * 21)[:63]),
+    ]:
+        for _ in range(20):
+            shrunk, expanded = run(m, pos, others)
+            assert torch.equal(shrunk, ref_shrunk), (m, pos)
+            assert torch.equal(expanded, ref_expanded), (m, pos)
+
+
+def _case_empty_and_metadata_swap():
+    rank, scaling = 8, 1.0
+    inputs, weights, _ = _lora_data(16, 128, rank, 1, 2, torch.bfloat16, seed=3)
+
+    no_rows = torch.empty((0,), dtype=torch.int32, device=DEVICE)
+    _shrink(inputs[:0], weights, _meta(no_rows, 2), _new_out(weights, 0), scaling)
+
+    untouched = _new_out(weights, 16, fill=-1.0)
+    no_lora = torch.full((16,), -1, dtype=torch.int32, device=DEVICE)
+    _shrink(inputs, weights, _meta(no_lora, 2), untouched, scaling)
+    assert torch.all(untouched == -1.0)
+
+    mappings = [
+        torch.tensor(ids * 2, dtype=torch.int32, device=DEVICE)
+        for ids in ([0, -1, 1, -1, 0, 1, -1, 0], [0, 0, -1, -1] * 2, [-1, 1] * 4)
+    ]
+    out, meta, first = _new_out(weights, 16, fill=-1.0), None, None
+    for mapping in mappings + mappings[:1]:
+        meta = _meta(mapping, 2, meta)
+        _shrink(inputs, weights, meta, out, scaling)
+        _assert_matches_fp64(out, inputs, weights, mapping, scaling)
+        first = out.clone() if first is None else first
+    assert torch.equal(out, first)
+
+
+def _case_cuda_graph_replay():
+    hidden, m, scaling, bf16 = 2560, 96, 0.6, torch.bfloat16
+    inputs, weights, mapping = _lora_data(m, hidden, 16, 1, 3, bf16, 5)
+
+    def eager(x, ids):
+        return _shrink(x, weights, _meta(ids, 3), _new_out(weights, x.size(0)), scaling)
+
+    reference = eager(inputs, mapping)
+    static_inputs, static_out = inputs.clone(), _new_out(weights, m)
+    meta = _meta(mapping, 3)
+    stream = torch.cuda.Stream(device=DEVICE)
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            _shrink(static_inputs, weights, meta, static_out, scaling)
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        _shrink(static_inputs, weights, meta, static_out, scaling)
+    graph.replay()
+    assert torch.equal(static_out, reference)
+
+    # A larger eager call between replays must not disturb the graph's scratch.
+    with torch.cuda.stream(stream):
+        large_inputs, _, large_mapping = _lora_data(256, hidden, 16, 1, 3, bf16, 6)
+        eager(large_inputs, large_mapping)
+    torch.cuda.current_stream().wait_stream(stream)
+
+    new_inputs, _, new_mapping = _lora_data(m, hidden, 16, 1, 3, bf16, 7)
+    for x, ids, expected in (
+        (new_inputs, new_mapping, eager(new_inputs, new_mapping)),
+        (inputs, mapping, reference),
+    ):
+        static_inputs.copy_(x)
+        _meta(ids, 3, meta)
+        graph.replay()
+        assert torch.equal(static_out, expected)
+
+
+def _case_concurrent_streams():
+    # Same lane on two streams, as MoE shared-experts overlap schedules it.
+    aux_stream = torch.cuda.Stream(device=DEVICE)
+    for i in range(50):
+        calls = [
+            _lora_data(32, 2560, 16, 1, 2, torch.bfloat16, seed=2 * i + j)
+            for j in range(2)
+        ]
+        refs = [
+            _shrink(x, w, _meta(ids, 2), _new_out(w, 32), 0.9) for x, w, ids in calls
+        ]
+        metas = [_meta(ids, 2) for _, _, ids in calls]
+        outs = [_new_out(w, 32) for _, w, _ in calls]
+        inputs_ready = torch.cuda.Event()
+        inputs_ready.record()
+        (x0, w0, _), (x1, w1, _) = calls
+        _shrink(x0, w0, metas[0], outs[0], 0.9)
+        with torch.cuda.stream(aux_stream):
+            inputs_ready.wait(aux_stream)
+            _shrink(x1, w1, metas[1], outs[1], 0.9)
+            aux_done = torch.cuda.Event()
+            aux_done.record(aux_stream)
+        aux_done.wait(torch.cuda.current_stream())
+        assert torch.equal(outs[0], refs[0]) and torch.equal(outs[1], refs[1]), i
+
+
+@pytest.mark.parametrize(
+    "k,rank,nslices,dtype,scaling",
+    [
+        (2560, 8, 3, "bfloat16", 0.73),
+        (9728, 16, 1, "float16", 1.0),
+        (4097, 64, 1, "bfloat16", 0.5),
+        (1024, 8, 1, "float16", 1.0),
+    ],
+    ids=["K2560_r8_3slices", "K9728_r16", "K4097_r64_tail", "K1024_empty_split"],
+)
+def test_splitk_matches_fp64_reference(k, rank, nslices, dtype, scaling):
+    _run_case(S8, "_case_numerics", k, rank, nslices, dtype, scaling)
+
+
+@pytest.mark.parametrize(
+    "env,split_k", [(S8, 8), (S1, 0), (BI_OFF, 0)], ids=["S8", "BI_S1", "BI_off"]
+)
+def test_shrink_dispatch(env, split_k):
+    _run_case(env, "_case_dispatch", split_k)
+
+
+@pytest.mark.parametrize("nslices", [1, 3])
+def test_splitk_batch_invariance(nslices):
+    _run_case(S8, "_case_batch_invariance", nslices)
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        "_case_empty_and_metadata_swap",
+        "_case_cuda_graph_replay",
+        "_case_concurrent_streams",
+    ],
+)
+def test_splitk(case):
+    _run_case(S8, case)
+
+
+@pytest.mark.parametrize(
+    "env,message",
+    [
+        ({**S8, "VLLM_BATCH_INVARIANT": "0"}, "VLLM_BATCH_INVARIANT"),
+        ({**S8, "VLLM_LORA_ENABLE_DUAL_STREAM": "1"}, "VLLM_LORA_ENABLE_DUAL_STREAM"),
+    ],
+    ids=["requires_batch_invariant", "rejects_dual_stream"],
+)
+def test_splitk_rejects_unsupported_env(env, message):
+    proc = _subprocess(env, "import vllm.lora.ops.triton_ops.lora_shrink_op")
+    assert proc.returncode != 0 and message in proc.stderr, proc.stderr
diff --git a/tests/test_envs.py b/tests/test_envs.py
index 1ef40cde0cc0..29cf2d3906f1 100644
--- a/tests/test_envs.py
+++ b/tests/test_envs.py
@@ -272,6 +272,21 @@ def test_gdn_decode_kernel_env(monkeypatch: pytest.MonkeyPatch):
             env_func()
 
 
+def test_lora_deterministic_split_k_env(monkeypatch: pytest.MonkeyPatch):
+    env_func = environment_variables["VLLM_LORA_DETERMINISTIC_SPLIT_K"]
+    monkeypatch.delenv("VLLM_LORA_DETERMINISTIC_SPLIT_K", raising=False)
+    assert env_func() == 0
+
+    for value in ("0", "8"):
+        monkeypatch.setenv("VLLM_LORA_DETERMINISTIC_SPLIT_K", value)
+        assert env_func() == int(value)
+
+    for value in ("4", "16", "eight"):
+        monkeypatch.setenv("VLLM_LORA_DETERMINISTIC_SPLIT_K", value)
+        with pytest.raises(ValueError, match="VLLM_LORA_DETERMINISTIC_SPLIT_K"):
+            env_func()
+
+
 class TestEnvListWithChoices:
     """Test cases for env_list_with_choices function."""
 
diff --git a/vllm/envs.py b/vllm/envs.py
index 8d836ffc4e71..79ef3deb1ff9 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -320,6 +320,9 @@
     VLLM_XPU_USE_SAMPLER_KERNEL: bool = True
     VLLM_XPU_INC_WNA16_BACKEND: Literal["auto", "ark", "w4a16", "w4a8"] = "auto"
     VLLM_LORA_ENABLE_DUAL_STREAM: bool = False
+    VLLM_LORA_DETERMINISTIC_SPLIT_K: int = 0
+    """``8`` runs the batch-invariant LoRA shrink as a fixed split-K=8 two-pass
+    kernel (requires ``VLLM_BATCH_INVARIANT=1``); ``0`` (default) disables it."""
     VLLM_GPU_NIC_PCIE_MAPPING: str = ""
     VLLM_NIC_SELECTION_VARS: str = ""
     VLLM_ENABLE_HPC_OPS: bool = False
@@ -503,6 +506,21 @@ def _get_validated_env_set() -> set[str]:
     return _get_validated_env_set
 
 
+def _parse_lora_deterministic_split_k(value: str) -> int:
+    try:
+        split_k = int(value)
+    except ValueError as err:
+        raise ValueError(
+            f"VLLM_LORA_DETERMINISTIC_SPLIT_K '{value}' must be an integer"
+        ) from err
+    if split_k not in (0, 8):
+        raise ValueError(
+            f"VLLM_LORA_DETERMINISTIC_SPLIT_K '{split_k}' is not supported; "
+            "only 0 (disabled) or 8 (fixed split-K=8) are allowed"
+        )
+    return split_k
+
+
 def get_vllm_port() -> int | None:
     """Get the port from VLLM_PORT environment variable.
 
@@ -2156,6 +2174,9 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_LORA_ENABLE_DUAL_STREAM": lambda: bool(
         int(os.getenv("VLLM_LORA_ENABLE_DUAL_STREAM", "0"))
     ),
+    "VLLM_LORA_DETERMINISTIC_SPLIT_K": lambda: _parse_lora_deterministic_split_k(
+        os.getenv("VLLM_LORA_DETERMINISTIC_SPLIT_K", "0")
+    ),
     # If set to 1, use Python spinloop extension to poll in a more efficient
     # way when using the mp backend.
     "VLLM_USE_SPINLOOP_EXT": lambda: bool(int(os.getenv("VLLM_USE_SPINLOOP_EXT", "0"))),
diff --git a/vllm/lora/ops/triton_ops/kernel_utils.py b/vllm/lora/ops/triton_ops/kernel_utils.py
index fd6ae4442869..bc07817ae19c 100644
--- a/vllm/lora/ops/triton_ops/kernel_utils.py
+++ b/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")
diff --git a/vllm/lora/ops/triton_ops/lora_shrink_op.py b/vllm/lora/ops/triton_ops/lora_shrink_op.py
index ec29065b9acf..97136aa69b99 100644
--- a/vllm/lora/ops/triton_ops/lora_shrink_op.py
+++ b/vllm/lora/ops/triton_ops/lora_shrink_op.py
@@ -18,6 +18,95 @@
 from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import direct_register_custom_op
 
+# 0 keeps the default shrink path; 8 (BI only) writes FP32 partials for 8
+# fixed K splits and reduces them in a fixed order.
+_TWO_PASS_SPLIT_K = envs.VLLM_LORA_DETERMINISTIC_SPLIT_K
+if _TWO_PASS_SPLIT_K:
+    if not envs.VLLM_BATCH_INVARIANT:
+        raise RuntimeError(
+            "VLLM_LORA_DETERMINISTIC_SPLIT_K=8 requires VLLM_BATCH_INVARIANT=1."
+        )
+    if envs.VLLM_LORA_ENABLE_DUAL_STREAM:
+        raise RuntimeError(
+            "VLLM_LORA_DETERMINISTIC_SPLIT_K=8 cannot be combined with "
+            "VLLM_LORA_ENABLE_DUAL_STREAM=1."
+        )
+
+
+@triton.jit
+def _lora_shrink_reduce_kernel(
+    partial_ptr,
+    out_ptr,
+    M,
+    N: tl.constexpr,
+    token_indices_sorted_by_lora_ids,
+    num_tokens_per_lora,
+    lora_token_start_loc,
+    lora_ids,
+    scaling,
+    partial_d0_stride,
+    partial_d1_stride,
+    partial_d2_stride,
+    partial_d3_stride,
+    output_d0_stride,
+    output_d1_stride,
+    output_d2_stride,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+):
+    cta_n_num = tl.cdiv(N, BLOCK_N)
+    cta_m_num = tl.cdiv(M, BLOCK_M)
+    pid_m_n = tl.program_id(axis=0)
+    num_pid_in_group = GROUP_SIZE_M * cta_n_num
+    group_id = pid_m_n // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(cta_m_num - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)
+    pid_n = (pid_m_n % num_pid_in_group) // group_size_m
+
+    slice_id = tl.program_id(axis=1)
+    lora_idx = tl.program_id(axis=2)
+    lora_id = tl.load(lora_ids + lora_idx)
+    if lora_id == -1:
+        return
+
+    lora_m_size = tl.load(num_tokens_per_lora + lora_idx)
+    cta_m_offset = pid_m * BLOCK_M
+    if cta_m_offset >= lora_m_size:
+        return
+
+    cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset)
+    lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx)
+    cta_lora_seq_indices = (
+        token_indices_sorted_by_lora_ids + lora_m_indices_start + cta_m_offset
+    )
+    offset_m = tl.arange(0, BLOCK_M) % cta_m_len
+    ram = tl.load(cta_lora_seq_indices + offset_m)
+    offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    offset_cm = tl.arange(0, BLOCK_M)
+    mask = (offset_cm[:, None] < cta_m_len) & (offset_n[None, :] < N)
+
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for split_id in range(SPLIT_K):
+        partial_in = (
+            partial_ptr
+            + slice_id.to(tl.int64) * partial_d0_stride
+            + tl.cast(split_id, tl.int64) * partial_d1_stride
+            + ram[:, None] * partial_d2_stride
+            + offset_n[None, :] * partial_d3_stride
+        )
+        accumulator += tl.load(partial_in, mask=mask, other=0.0)
+
+    output = (
+        out_ptr
+        + slice_id * output_d0_stride
+        + ram[:, None] * output_d1_stride
+        + offset_n[None, :] * output_d2_stride
+    )
+    tl.store(output, accumulator * scaling, mask=mask)
+
 
 @triton.jit
 def _lora_shrink_kernel(
@@ -49,7 +138,15 @@ def _lora_shrink_kernel(
     SLICE_NUM: tl.constexpr,
     USE_GDC: tl.constexpr,
     launch_pdl: tl.constexpr,
+    STORE_PARTIALS: tl.constexpr = False,
+    output_split_stride=0,
+    PARTIAL_N: tl.constexpr = 0,
+    PARTIAL_K: tl.constexpr = 0,
 ):
+    if STORE_PARTIALS:
+        # The split-K partial stage specializes on N and K.
+        N = PARTIAL_N
+        K = PARTIAL_K
     cta_n_num = tl.cdiv(N, BLOCK_N)
     cta_m_num = tl.cdiv(M, BLOCK_M)
 
@@ -116,6 +213,7 @@ def _lora_shrink_kernel(
         output_d0_stride,
         output_d1_stride,
         output_d2_stride,
+        output_split_stride,
         scaling,
         BLOCK_M,
         BLOCK_N,
@@ -124,6 +222,7 @@ def _lora_shrink_kernel(
         SPLIT_K,
         SLICE_NUM,
         USE_GDC,
+        STORE_PARTIALS,
     )
 
 
@@ -167,6 +266,23 @@ def _lora_shrink(
 
     """
     assert no_lora_flag_cpu.numel() == 1
+    if inputs.size(0) == 0:
+        # An empty first call has no work to do.
+        return
+    partials: torch.Tensor | None = None
+    if _TWO_PASS_SPLIT_K:
+        # Allocated before the no-LoRA return so profiling sees the scratch.
+        # Per-call allocation keeps each CUDA graph's scratch in its own pool.
+        partials = torch.empty(
+            (
+                output_tensor.size(0),
+                _TWO_PASS_SPLIT_K,
+                inputs.size(0),
+                output_tensor.size(-1),
+            ),
+            dtype=torch.float32,
+            device=output_tensor.device,
+        )
     if no_lora_flag_cpu.item():
         # None of the inputs require LoRA.
         return
@@ -213,6 +329,11 @@ def _lora_shrink(
     NUM_STAGES = kernel_config["num_stages"]
     NUM_CTAS = kernel_config["num_ctas"]
     GROUP_SIZE_M = kernel_config.get("group_size_m", 8)
+
+    # Every M must keep the same K partition for batch invariance.
+    if _TWO_PASS_SPLIT_K:
+        BLOCK_K = 256
+        SPLIT_K = _TWO_PASS_SPLIT_K
     EVEN_K = K % (BLOCK_K * SPLIT_K) == 0  # type: ignore
 
     # TODO (varun): This grid formulation maximizes parallelization at the
@@ -226,10 +347,19 @@ def _lora_shrink(
 
     # PDL only works when dual-stream is being used.
     use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM
+    first_out = output_tensor if partials is None else partials
+    partial_args = {}
+    if partials is not None:
+        partial_args = dict(
+            STORE_PARTIALS=True,
+            output_split_stride=partials.stride(1),
+            PARTIAL_N=N,
+            PARTIAL_K=K,
+        )
     _lora_shrink_kernel[grid](
         inputs,
         lora_ptr_tensor,
-        output_tensor,
+        first_out,
         M,
         N,
         K,
@@ -243,9 +373,9 @@ def _lora_shrink(
         lora_strides_d0,
         lora_strides_d1,
         lora_strides_d2,
-        output_tensor.stride(0),
-        output_tensor.stride(1),
-        output_tensor.stride(2),
+        first_out.stride(0),
+        first_out.stride(-2),
+        first_out.stride(-1),
         BLOCK_M,
         BLOCK_N,
         BLOCK_K,
@@ -258,7 +388,28 @@ def _lora_shrink(
         num_ctas=NUM_CTAS,
         num_stages=NUM_STAGES,
         launch_pdl=use_gdc,
+        **partial_args,
     )
+    if partials is not None:
+        _lora_shrink_reduce_kernel[(grid[0] // SPLIT_K, *grid[1:])](
+            partials,
+            output_tensor,
+            M,
+            N,
+            token_indices_sorted_by_lora_ids,
+            num_tokens_per_lora,
+            lora_token_start_loc,
+            lora_ids,
+            scaling,
+            *partials.stride(),
+            *output_tensor.stride(),
+            BLOCK_M,
+            BLOCK_N,
+            SPLIT_K,
+            GROUP_SIZE_M,
+            num_warps=NUM_WARPS,
+            num_stages=NUM_STAGES,
+        )
 
     return
 

From 76d4a931f2e88715dd74bf0e3d6b859d32785e64 Mon Sep 17 00:00:00 2001
From: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com>
Date: Wed, 30 Sep 2026 15:02:42 -0700
Subject: [PATCH 2/3] [Kernel][LoRA] Use deterministic split-K shrink whenever
 batch invariant

Drop VLLM_LORA_DETERMINISTIC_SPLIT_K and its import-time guards: the
fixed 8-way split-K shrink is now the batch-invariant LoRA shrink path.
It skips PDL, so LoRA dual-stream no longer needs to be rejected. Move
the tests into test_punica_ops.py as in-process checks against an FP64
reference and across batch sizes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com>
---
 docs/features/batch_invariance.md          |  21 +-
 tests/lora/test_lora_shrink_splitk.py      | 338 ---------------------
 tests/lora/test_punica_ops.py              |  91 ++++++
 tests/test_envs.py                         |  15 -
 vllm/envs.py                               |  21 --
 vllm/lora/ops/triton_ops/lora_shrink_op.py |  30 +-
 6 files changed, 104 insertions(+), 412 deletions(-)
 delete mode 100644 tests/lora/test_lora_shrink_splitk.py

diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md
index 5b159c0a9e68..fe7200430364 100644
--- a/docs/features/batch_invariance.md
+++ b/docs/features/batch_invariance.md
@@ -138,26 +138,6 @@ Batch invariance has been tested and verified on the following models:
 
 Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose).
 
-## LoRA Deterministic Split-K (Experimental)
-
-In batch-invariant mode the LoRA shrink kernel does not split the K
-dimension. `VLLM_LORA_DETERMINISTIC_SPLIT_K=8` instead writes FP32 partial
-sums for 8 fixed K splits and adds them in a fixed order, without atomics.
-
-```bash
-VLLM_BATCH_INVARIANT=1 VLLM_LORA_DETERMINISTIC_SPLIT_K=8 vllm serve <model> \
-    --enable-lora --lora-modules ...
-```
-
-- Requires `VLLM_BATCH_INVARIANT=1`; accepted values are `0` (default) and `8`.
-- Cannot be combined with `VLLM_LORA_ENABLE_DUAL_STREAM=1`.
-- Tested on CUDA dense LoRA with BF16/FP16 inputs on SM89, SM90, SM100 and SM120 GPUs.
-
-!!! warning
-    Split-K=8 is batch-invariant on its own but uses a different reduction
-    order than the default path, so generated tokens and logprobs can differ
-    between the two settings.
-
 ## Implementation Details
 
 When batch invariance is enabled, vLLM:
@@ -165,6 +145,7 @@ When batch invariance is enabled, vLLM:
 1. Uses deterministic kernel implementations for attention and other operations
 2. Ensures consistent numerical behavior across different batch sizes
 3. Disables certain optimizations that may introduce non-determinism (such as custom all-reduce operations in tensor parallel mode)
+4. Splits the LoRA shrink reduction into 8 fixed K chunks that are summed in a fixed order instead of with atomics
 
 !!! note
     Enabling batch invariance may impact performance compared to the default non-deterministic mode. This trade-off is intentional to guarantee reproducibility.
diff --git a/tests/lora/test_lora_shrink_splitk.py b/tests/lora/test_lora_shrink_splitk.py
deleted file mode 100644
index eba2546abd0f..000000000000
--- a/tests/lora/test_lora_shrink_splitk.py
+++ /dev/null
@@ -1,338 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the LoRA shrink deterministic split-K path.
-
-The shrink op reads VLLM_LORA_DETERMINISTIC_SPLIT_K and VLLM_BATCH_INVARIANT
-once at import, so every case runs in a fresh interpreter whose environment
-is set before start-up.
-"""
-
-import os
-import subprocess
-import sys
-
-import pytest
-import torch
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available(), reason="LoRA shrink kernel requires CUDA"
-)
-
-S8 = {
-    "VLLM_BATCH_INVARIANT": "1",
-    "VLLM_LORA_DETERMINISTIC_SPLIT_K": "8",
-    "VLLM_LORA_ENABLE_DUAL_STREAM": "0",
-}
-S1 = {**S8, "VLLM_LORA_DETERMINISTIC_SPLIT_K": "0"}
-BI_OFF = {**S1, "VLLM_BATCH_INVARIANT": "0"}
-DEVICE = "cuda:0"
-REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
-
-
-@pytest.fixture(autouse=True)
-def cleanup_fixture():
-    """GPU work runs in subprocesses; skip conftest teardown (needs ray)."""
-    yield
-
-
-@pytest.fixture(autouse=True)
-def dynamo_reset():
-    yield
-
-
-def _subprocess(env: dict[str, str], code: str) -> subprocess.CompletedProcess:
-    return subprocess.run(
-        [sys.executable, "-c", code],
-        env={**os.environ, **env},
-        cwd=REPO_ROOT,
-        capture_output=True,
-        text=True,
-        timeout=600,
-    )
-
-
-def _run_case(env: dict[str, str], case: str, *args) -> None:
-    code = f"import runpy; runpy.run_path({__file__!r})[{case!r}](*{args!r})"
-    proc = _subprocess(env, code)
-    assert proc.returncode == 0, proc.stdout + proc.stderr
-
-
-def _lora_data(m, k, rank, nslices, num_loras, dtype, seed):
-    g = torch.Generator(device=DEVICE).manual_seed(seed)
-    inputs = torch.randn((m, k), dtype=dtype, device=DEVICE, generator=g)
-    weights = [
-        torch.randn((num_loras, rank, k), dtype=dtype, device=DEVICE, generator=g)
-        for _ in range(nslices)
-    ]
-    mapping = torch.randint(
-        0, num_loras, (m,), dtype=torch.int32, device=DEVICE, generator=g
-    )
-    return inputs, weights, mapping
-
-
-def _meta(mapping, max_loras, meta=None):
-    from vllm.lora.ops.triton_ops.lora_kernel_metadata import LoRAKernelMeta
-
-    if meta is None:
-        meta = LoRAKernelMeta.make(max_loras, mapping.numel(), device=DEVICE)
-    meta.prepare_tensors(mapping)
-    return meta
-
-
-def _new_out(weights, m, fill=0.0):
-    shape = (len(weights), m, weights[0].size(1))
-    return torch.full(shape, fill, dtype=torch.float32, device=DEVICE)
-
-
-def _shrink(inputs, weights, meta, out, scaling):
-    from vllm.lora.ops.triton_ops import lora_shrink
-    from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT
-
-    _LORA_A_PTR_DICT.clear()
-    args = meta.meta_args(token_nums=inputs.size(0), specialize_active_lora=False)
-    lora_shrink(inputs, weights, out, *args, scaling)
-    return out
-
-
-def _assert_matches_fp64(out, inputs, weights, mapping, scaling):
-    x, ids = inputs.double().cpu(), mapping.cpu()
-    ref = torch.zeros(out.shape, dtype=torch.float64)
-    for s, w in enumerate(weights):
-        for lid in ids[ids >= 0].unique().tolist():
-            rows = ids == lid
-            ref[s, rows] = scaling * (x[rows] @ w[lid].double().cpu().T)
-    out = out.double().cpu()
-    assert torch.isfinite(out).all()
-    assert torch.all(out[:, ids < 0] == 0)
-    torch.testing.assert_close(out, ref, rtol=0.005, atol=0.005)
-
-
-class _LaunchRecorder:
-    def __init__(self, name, kernel, sink):
-        self.name, self.kernel, self.sink = name, kernel, sink
-
-    def __getitem__(self, grid):
-        launch = self.kernel[grid]
-
-        def record(*args, **kwargs):
-            bound = dict(zip(self.kernel.arg_names, args))
-            self.sink.append((self.name, {**bound, **kwargs}))
-            return launch(*args, **kwargs)
-
-        return record
-
-
-def _case_numerics(k, rank, nslices, dtype, scaling):
-    dtype = getattr(torch, dtype)
-    inputs, weights, mapping = _lora_data(17, k, rank, nslices, 2, dtype, seed=0)
-    out = _shrink(inputs, weights, _meta(mapping, 2), _new_out(weights, 17), scaling)
-    _assert_matches_fp64(out, inputs, weights, mapping, scaling)
-
-
-def _case_dispatch(split_k):
-    import vllm.lora.ops.triton_ops.lora_shrink_op as op
-
-    launches: list = []
-    for name in ("_lora_shrink_kernel", "_lora_shrink_reduce_kernel"):
-        setattr(op, name, _LaunchRecorder(name, getattr(op, name), launches))
-
-    inputs, weights, mapping = _lora_data(16, 2049, 32, 1, 4, torch.bfloat16, 1)
-    out = _shrink(inputs, weights, _meta(mapping, 4), _new_out(weights, 16), 0.5)
-    _assert_matches_fp64(out, inputs, weights, mapping, 0.5)
-
-    names = [name for name, _ in launches]
-    if not split_k:
-        assert names == ["_lora_shrink_kernel"], names
-        assert not launches[0][1].get("STORE_PARTIALS", False)
-        return
-    assert names == ["_lora_shrink_kernel", "_lora_shrink_reduce_kernel"], names
-    (_, first), (_, reduce) = launches
-    assert first["STORE_PARTIALS"] and first["out_ptr"].dtype == torch.float32
-    launch = dict(BLOCK_M=32, BLOCK_N=16, SPLIT_K=8, num_warps=4, num_stages=2)
-    assert {key: first[key] for key in launch} == launch
-    assert (first["BLOCK_K"], first["num_ctas"]) == (256, 1)
-    assert {key: reduce[key] for key in launch} == launch
-
-
-def _case_batch_invariance(nslices):
-    from vllm.lora.ops.triton_ops import lora_expand
-    from vllm.lora.ops.triton_ops.utils import _LORA_B_PTR_DICT
-
-    # K=2560 is ten BK=256 blocks, so every one of the eight splits is used.
-    hidden, rank, num_loras, scaling = 2560, 16, 3, 0.7
-    g = torch.Generator(device=DEVICE).manual_seed(1234 + nslices)
-
-    def randn(*shape):
-        return torch.randn(shape, dtype=torch.bfloat16, device=DEVICE, generator=g)
-
-    lora_a = [randn(num_loras, rank, hidden) for _ in range(nslices)]
-    lora_b = [randn(num_loras, hidden, rank) for _ in range(nslices)]
-    target, target_residual = randn(hidden), randn(hidden * nslices)
-
-    def run(m, pos, others):
-        ids = torch.tensor(others[:pos] + [0] + others[pos:], device=DEVICE)
-        inputs, residual = randn(m, hidden), randn(m, hidden * nslices)
-        inputs[pos], residual[pos] = target, target_residual
-        meta = _meta(ids.to(torch.int32), num_loras)
-        shrunk = _shrink(inputs, lora_a, meta, _new_out(lora_a, m), scaling)
-        _LORA_B_PTR_DICT.clear()
-        args = meta.meta_args(token_nums=m, specialize_active_lora=False)
-        lora_expand(shrunk, lora_b, residual, *args, add_inputs=True)
-        return shrunk[:, pos].clone(), residual[pos].clone()
-
-    ref_shrunk, ref_expanded = run(1, 0, [])
-    for m, pos, others in [  # (M, target row, adapters of the other rows)
-        (127, 0, [0] * 126),
-        (127, 126, [1, 2] * 63),
-        (128, 0, [1] * 127),
-        (129, 128, [2] * 128),
-        (31, 15, [0, 1, 2] * 10),
-        (32, 16, [0, 1, 2] * 10 + [1]),
-        (33, 16, [0, 1, 2] * 10 + [1, 2]),
-        (64, 40, ([0, 1, 2] * 21)[:63]),
-    ]:
-        for _ in range(20):
-            shrunk, expanded = run(m, pos, others)
-            assert torch.equal(shrunk, ref_shrunk), (m, pos)
-            assert torch.equal(expanded, ref_expanded), (m, pos)
-
-
-def _case_empty_and_metadata_swap():
-    rank, scaling = 8, 1.0
-    inputs, weights, _ = _lora_data(16, 128, rank, 1, 2, torch.bfloat16, seed=3)
-
-    no_rows = torch.empty((0,), dtype=torch.int32, device=DEVICE)
-    _shrink(inputs[:0], weights, _meta(no_rows, 2), _new_out(weights, 0), scaling)
-
-    untouched = _new_out(weights, 16, fill=-1.0)
-    no_lora = torch.full((16,), -1, dtype=torch.int32, device=DEVICE)
-    _shrink(inputs, weights, _meta(no_lora, 2), untouched, scaling)
-    assert torch.all(untouched == -1.0)
-
-    mappings = [
-        torch.tensor(ids * 2, dtype=torch.int32, device=DEVICE)
-        for ids in ([0, -1, 1, -1, 0, 1, -1, 0], [0, 0, -1, -1] * 2, [-1, 1] * 4)
-    ]
-    out, meta, first = _new_out(weights, 16, fill=-1.0), None, None
-    for mapping in mappings + mappings[:1]:
-        meta = _meta(mapping, 2, meta)
-        _shrink(inputs, weights, meta, out, scaling)
-        _assert_matches_fp64(out, inputs, weights, mapping, scaling)
-        first = out.clone() if first is None else first
-    assert torch.equal(out, first)
-
-
-def _case_cuda_graph_replay():
-    hidden, m, scaling, bf16 = 2560, 96, 0.6, torch.bfloat16
-    inputs, weights, mapping = _lora_data(m, hidden, 16, 1, 3, bf16, 5)
-
-    def eager(x, ids):
-        return _shrink(x, weights, _meta(ids, 3), _new_out(weights, x.size(0)), scaling)
-
-    reference = eager(inputs, mapping)
-    static_inputs, static_out = inputs.clone(), _new_out(weights, m)
-    meta = _meta(mapping, 3)
-    stream = torch.cuda.Stream(device=DEVICE)
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        for _ in range(3):
-            _shrink(static_inputs, weights, meta, static_out, scaling)
-    torch.cuda.current_stream().wait_stream(stream)
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph, stream=stream):
-        _shrink(static_inputs, weights, meta, static_out, scaling)
-    graph.replay()
-    assert torch.equal(static_out, reference)
-
-    # A larger eager call between replays must not disturb the graph's scratch.
-    with torch.cuda.stream(stream):
-        large_inputs, _, large_mapping = _lora_data(256, hidden, 16, 1, 3, bf16, 6)
-        eager(large_inputs, large_mapping)
-    torch.cuda.current_stream().wait_stream(stream)
-
-    new_inputs, _, new_mapping = _lora_data(m, hidden, 16, 1, 3, bf16, 7)
-    for x, ids, expected in (
-        (new_inputs, new_mapping, eager(new_inputs, new_mapping)),
-        (inputs, mapping, reference),
-    ):
-        static_inputs.copy_(x)
-        _meta(ids, 3, meta)
-        graph.replay()
-        assert torch.equal(static_out, expected)
-
-
-def _case_concurrent_streams():
-    # Same lane on two streams, as MoE shared-experts overlap schedules it.
-    aux_stream = torch.cuda.Stream(device=DEVICE)
-    for i in range(50):
-        calls = [
-            _lora_data(32, 2560, 16, 1, 2, torch.bfloat16, seed=2 * i + j)
-            for j in range(2)
-        ]
-        refs = [
-            _shrink(x, w, _meta(ids, 2), _new_out(w, 32), 0.9) for x, w, ids in calls
-        ]
-        metas = [_meta(ids, 2) for _, _, ids in calls]
-        outs = [_new_out(w, 32) for _, w, _ in calls]
-        inputs_ready = torch.cuda.Event()
-        inputs_ready.record()
-        (x0, w0, _), (x1, w1, _) = calls
-        _shrink(x0, w0, metas[0], outs[0], 0.9)
-        with torch.cuda.stream(aux_stream):
-            inputs_ready.wait(aux_stream)
-            _shrink(x1, w1, metas[1], outs[1], 0.9)
-            aux_done = torch.cuda.Event()
-            aux_done.record(aux_stream)
-        aux_done.wait(torch.cuda.current_stream())
-        assert torch.equal(outs[0], refs[0]) and torch.equal(outs[1], refs[1]), i
-
-
-@pytest.mark.parametrize(
-    "k,rank,nslices,dtype,scaling",
-    [
-        (2560, 8, 3, "bfloat16", 0.73),
-        (9728, 16, 1, "float16", 1.0),
-        (4097, 64, 1, "bfloat16", 0.5),
-        (1024, 8, 1, "float16", 1.0),
-    ],
-    ids=["K2560_r8_3slices", "K9728_r16", "K4097_r64_tail", "K1024_empty_split"],
-)
-def test_splitk_matches_fp64_reference(k, rank, nslices, dtype, scaling):
-    _run_case(S8, "_case_numerics", k, rank, nslices, dtype, scaling)
-
-
-@pytest.mark.parametrize(
-    "env,split_k", [(S8, 8), (S1, 0), (BI_OFF, 0)], ids=["S8", "BI_S1", "BI_off"]
-)
-def test_shrink_dispatch(env, split_k):
-    _run_case(env, "_case_dispatch", split_k)
-
-
-@pytest.mark.parametrize("nslices", [1, 3])
-def test_splitk_batch_invariance(nslices):
-    _run_case(S8, "_case_batch_invariance", nslices)
-
-
-@pytest.mark.parametrize(
-    "case",
-    [
-        "_case_empty_and_metadata_swap",
-        "_case_cuda_graph_replay",
-        "_case_concurrent_streams",
-    ],
-)
-def test_splitk(case):
-    _run_case(S8, case)
-
-
-@pytest.mark.parametrize(
-    "env,message",
-    [
-        ({**S8, "VLLM_BATCH_INVARIANT": "0"}, "VLLM_BATCH_INVARIANT"),
-        ({**S8, "VLLM_LORA_ENABLE_DUAL_STREAM": "1"}, "VLLM_LORA_ENABLE_DUAL_STREAM"),
-    ],
-    ids=["requires_batch_invariant", "rejects_dual_stream"],
-)
-def test_splitk_rejects_unsupported_env(env, message):
-    proc = _subprocess(env, "import vllm.lora.ops.triton_ops.lora_shrink_op")
-    assert proc.returncode != 0 and message in proc.stderr, proc.stderr
diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py
index f2bb18aac26e..e1d6981b76f8 100644
--- a/tests/lora/test_punica_ops.py
+++ b/tests/lora/test_punica_ops.py
@@ -641,3 +641,94 @@ def test_add_lora_fused_moe_early_exit(device):
     assert torch.equal(y, y_snapshot), (
         "add_lora_fused_moe modified output tensor despite no_lora_flag_cpu=True"
     )
+
+
+@pytest.fixture
+def batch_invariant(monkeypatch):
+    import vllm.lora.ops.triton_ops.lora_shrink_op as shrink_op
+    import vllm.lora.ops.triton_ops.utils as lora_utils
+
+    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 shrink_op
+    lora_utils.get_lora_op_configs.cache_clear()
+
+
+def _bi_shrink(inputs, weights, mapping, scaling):
+    num_tokens = inputs.size(0)
+    meta = LoRAKernelMeta.make(weights[0].size(0), num_tokens, device=DEVICE_TYPE)
+    meta.prepare_tensors(mapping)
+    out = torch.zeros(
+        (len(weights), num_tokens, weights[0].size(1)),
+        dtype=torch.float32,
+        device=DEVICE_TYPE,
+    )
+    _LORA_A_PTR_DICT.clear()
+    args = meta.meta_args(token_nums=num_tokens, specialize_active_lora=False)
+    triton_ops.lora_shrink(inputs, weights, out, *args, scaling)
+    return out
+
+
+@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_batch_invariant_matches_reference(
+    batch_invariant, hidden_size, rank, nslices, dtype
+):
+    set_random_seed(0)
+    inputs = torch.randn((33, hidden_size), dtype=dtype, device=DEVICE_TYPE)
+    weights = [
+        torch.randn((2, rank, hidden_size), dtype=dtype, device=DEVICE_TYPE)
+        for _ in range(nslices)
+    ]
+    mapping = torch.randint(-1, 2, (33,), dtype=torch.int32, device=DEVICE_TYPE)
+    out = _bi_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)
+
+
+@pytest.mark.parametrize("nslices", [1, 3])
+def test_lora_shrink_batch_invariant_across_batch_sizes(
+    batch_invariant, monkeypatch, nslices
+):
+    reduce_kernel, reduce_grids = batch_invariant._lora_shrink_reduce_kernel, []
+
+    class RecordReduce:
+        def __getitem__(self, grid):
+            reduce_grids.append(grid)
+            return reduce_kernel[grid]
+
+    monkeypatch.setattr(batch_invariant, "_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_TYPE)
+        for _ in range(nslices)
+    ]
+    target = torch.randn((1, hidden_size), dtype=dtype, device=DEVICE_TYPE)
+    target_id = torch.zeros(1, dtype=torch.int32, device=DEVICE_TYPE)
+    expected = _bi_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_TYPE)
+        mapping = torch.randint(
+            -1, 3, (num_tokens,), dtype=torch.int32, device=DEVICE_TYPE
+        )
+        pos = num_tokens // 2
+        inputs[pos], mapping[pos] = target[0], 0
+        out = _bi_shrink(inputs, weights, mapping, 0.7)
+        assert torch.equal(out[:, pos], expected), num_tokens
+    assert len(reduce_grids) == 7
diff --git a/tests/test_envs.py b/tests/test_envs.py
index 29cf2d3906f1..1ef40cde0cc0 100644
--- a/tests/test_envs.py
+++ b/tests/test_envs.py
@@ -272,21 +272,6 @@ def test_gdn_decode_kernel_env(monkeypatch: pytest.MonkeyPatch):
             env_func()
 
 
-def test_lora_deterministic_split_k_env(monkeypatch: pytest.MonkeyPatch):
-    env_func = environment_variables["VLLM_LORA_DETERMINISTIC_SPLIT_K"]
-    monkeypatch.delenv("VLLM_LORA_DETERMINISTIC_SPLIT_K", raising=False)
-    assert env_func() == 0
-
-    for value in ("0", "8"):
-        monkeypatch.setenv("VLLM_LORA_DETERMINISTIC_SPLIT_K", value)
-        assert env_func() == int(value)
-
-    for value in ("4", "16", "eight"):
-        monkeypatch.setenv("VLLM_LORA_DETERMINISTIC_SPLIT_K", value)
-        with pytest.raises(ValueError, match="VLLM_LORA_DETERMINISTIC_SPLIT_K"):
-            env_func()
-
-
 class TestEnvListWithChoices:
     """Test cases for env_list_with_choices function."""
 
diff --git a/vllm/envs.py b/vllm/envs.py
index 79ef3deb1ff9..8d836ffc4e71 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -320,9 +320,6 @@
     VLLM_XPU_USE_SAMPLER_KERNEL: bool = True
     VLLM_XPU_INC_WNA16_BACKEND: Literal["auto", "ark", "w4a16", "w4a8"] = "auto"
     VLLM_LORA_ENABLE_DUAL_STREAM: bool = False
-    VLLM_LORA_DETERMINISTIC_SPLIT_K: int = 0
-    """``8`` runs the batch-invariant LoRA shrink as a fixed split-K=8 two-pass
-    kernel (requires ``VLLM_BATCH_INVARIANT=1``); ``0`` (default) disables it."""
     VLLM_GPU_NIC_PCIE_MAPPING: str = ""
     VLLM_NIC_SELECTION_VARS: str = ""
     VLLM_ENABLE_HPC_OPS: bool = False
@@ -506,21 +503,6 @@ def _get_validated_env_set() -> set[str]:
     return _get_validated_env_set
 
 
-def _parse_lora_deterministic_split_k(value: str) -> int:
-    try:
-        split_k = int(value)
-    except ValueError as err:
-        raise ValueError(
-            f"VLLM_LORA_DETERMINISTIC_SPLIT_K '{value}' must be an integer"
-        ) from err
-    if split_k not in (0, 8):
-        raise ValueError(
-            f"VLLM_LORA_DETERMINISTIC_SPLIT_K '{split_k}' is not supported; "
-            "only 0 (disabled) or 8 (fixed split-K=8) are allowed"
-        )
-    return split_k
-
-
 def get_vllm_port() -> int | None:
     """Get the port from VLLM_PORT environment variable.
 
@@ -2174,9 +2156,6 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_LORA_ENABLE_DUAL_STREAM": lambda: bool(
         int(os.getenv("VLLM_LORA_ENABLE_DUAL_STREAM", "0"))
     ),
-    "VLLM_LORA_DETERMINISTIC_SPLIT_K": lambda: _parse_lora_deterministic_split_k(
-        os.getenv("VLLM_LORA_DETERMINISTIC_SPLIT_K", "0")
-    ),
     # If set to 1, use Python spinloop extension to poll in a more efficient
     # way when using the mp backend.
     "VLLM_USE_SPINLOOP_EXT": lambda: bool(int(os.getenv("VLLM_USE_SPINLOOP_EXT", "0"))),
diff --git a/vllm/lora/ops/triton_ops/lora_shrink_op.py b/vllm/lora/ops/triton_ops/lora_shrink_op.py
index 97136aa69b99..abe63bd41bbf 100644
--- a/vllm/lora/ops/triton_ops/lora_shrink_op.py
+++ b/vllm/lora/ops/triton_ops/lora_shrink_op.py
@@ -13,24 +13,14 @@
 from vllm.lora.ops.triton_ops.utils import (
     _get_lora_a_ptr,
     get_lora_op_configs,
+    is_batch_invariant,
     supports_pdl,
 )
 from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import direct_register_custom_op
 
-# 0 keeps the default shrink path; 8 (BI only) writes FP32 partials for 8
-# fixed K splits and reduces them in a fixed order.
-_TWO_PASS_SPLIT_K = envs.VLLM_LORA_DETERMINISTIC_SPLIT_K
-if _TWO_PASS_SPLIT_K:
-    if not envs.VLLM_BATCH_INVARIANT:
-        raise RuntimeError(
-            "VLLM_LORA_DETERMINISTIC_SPLIT_K=8 requires VLLM_BATCH_INVARIANT=1."
-        )
-    if envs.VLLM_LORA_ENABLE_DUAL_STREAM:
-        raise RuntimeError(
-            "VLLM_LORA_DETERMINISTIC_SPLIT_K=8 cannot be combined with "
-            "VLLM_LORA_ENABLE_DUAL_STREAM=1."
-        )
+# Batch-invariant shrink splits K into this many fixed partials (no atomics).
+_BI_SPLIT_K = 8
 
 
 @triton.jit
@@ -270,13 +260,13 @@ def _lora_shrink(
         # An empty first call has no work to do.
         return
     partials: torch.Tensor | None = None
-    if _TWO_PASS_SPLIT_K:
+    if is_batch_invariant:
         # Allocated before the no-LoRA return so profiling sees the scratch.
         # Per-call allocation keeps each CUDA graph's scratch in its own pool.
         partials = torch.empty(
             (
                 output_tensor.size(0),
-                _TWO_PASS_SPLIT_K,
+                _BI_SPLIT_K,
                 inputs.size(0),
                 output_tensor.size(-1),
             ),
@@ -331,9 +321,9 @@ def _lora_shrink(
     GROUP_SIZE_M = kernel_config.get("group_size_m", 8)
 
     # Every M must keep the same K partition for batch invariance.
-    if _TWO_PASS_SPLIT_K:
+    if partials is not None:
         BLOCK_K = 256
-        SPLIT_K = _TWO_PASS_SPLIT_K
+        SPLIT_K = _BI_SPLIT_K
     EVEN_K = K % (BLOCK_K * SPLIT_K) == 0  # type: ignore
 
     # TODO (varun): This grid formulation maximizes parallelization at the
@@ -346,7 +336,11 @@ def _lora_shrink(
     )
 
     # PDL only works when dual-stream is being used.
-    use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM
+    use_gdc = (
+        supports_pdl(inputs.device)
+        and envs.VLLM_LORA_ENABLE_DUAL_STREAM
+        and partials is None
+    )
     first_out = output_tensor if partials is None else partials
     partial_args = {}
     if partials is not None:

From dbb109eddadda31dac8699b43c535b93fb835924 Mon Sep 17 00:00:00 2001
From: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com>
Date: Fri, 2 Oct 2026 14:01:49 -0700
Subject: [PATCH 3/3] [Tests][LoRA] Move batch-invariant LoRA shrink tests to
 tests/v1/determinism

Run them in the H100 and B200 batch-invariance steps and trigger those
steps on vllm/lora/ops/triton_ops changes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com>
---
 .buildkite/test_areas/misc.yaml               |  4 +
 tests/lora/test_punica_ops.py                 | 91 -----------------
 .../test_lora_shrink_batch_invariant.py       | 99 +++++++++++++++++++
 3 files changed, 103 insertions(+), 91 deletions(-)
 create mode 100644 tests/v1/determinism/test_lora_shrink_batch_invariant.py

diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml
index 602491e0e0f8..cbcfaeaea7ca 100644
--- a/.buildkite/test_areas/misc.yaml
+++ b/.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
diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py
index 971c203ec44b..34cbe061dc42 100644
--- a/tests/lora/test_punica_ops.py
+++ b/tests/lora/test_punica_ops.py
@@ -639,94 +639,3 @@ def test_add_lora_fused_moe_early_exit(device):
     assert torch.equal(y, y_snapshot), (
         "add_lora_fused_moe modified output tensor despite no_lora_flag_cpu=True"
     )
-
-
-@pytest.fixture
-def batch_invariant(monkeypatch):
-    import vllm.lora.ops.triton_ops.lora_shrink_op as shrink_op
-    import vllm.lora.ops.triton_ops.utils as lora_utils
-
-    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 shrink_op
-    lora_utils.get_lora_op_configs.cache_clear()
-
-
-def _bi_shrink(inputs, weights, mapping, scaling):
-    num_tokens = inputs.size(0)
-    meta = LoRAKernelMeta.make(weights[0].size(0), num_tokens, device=DEVICE_TYPE)
-    meta.prepare_tensors(mapping)
-    out = torch.zeros(
-        (len(weights), num_tokens, weights[0].size(1)),
-        dtype=torch.float32,
-        device=DEVICE_TYPE,
-    )
-    _LORA_A_PTR_DICT.clear()
-    args = meta.meta_args(token_nums=num_tokens, specialize_active_lora=False)
-    triton_ops.lora_shrink(inputs, weights, out, *args, scaling)
-    return out
-
-
-@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_batch_invariant_matches_reference(
-    batch_invariant, hidden_size, rank, nslices, dtype
-):
-    set_random_seed(0)
-    inputs = torch.randn((33, hidden_size), dtype=dtype, device=DEVICE_TYPE)
-    weights = [
-        torch.randn((2, rank, hidden_size), dtype=dtype, device=DEVICE_TYPE)
-        for _ in range(nslices)
-    ]
-    mapping = torch.randint(-1, 2, (33,), dtype=torch.int32, device=DEVICE_TYPE)
-    out = _bi_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)
-
-
-@pytest.mark.parametrize("nslices", [1, 3])
-def test_lora_shrink_batch_invariant_across_batch_sizes(
-    batch_invariant, monkeypatch, nslices
-):
-    reduce_kernel, reduce_grids = batch_invariant._lora_shrink_reduce_kernel, []
-
-    class RecordReduce:
-        def __getitem__(self, grid):
-            reduce_grids.append(grid)
-            return reduce_kernel[grid]
-
-    monkeypatch.setattr(batch_invariant, "_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_TYPE)
-        for _ in range(nslices)
-    ]
-    target = torch.randn((1, hidden_size), dtype=dtype, device=DEVICE_TYPE)
-    target_id = torch.zeros(1, dtype=torch.int32, device=DEVICE_TYPE)
-    expected = _bi_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_TYPE)
-        mapping = torch.randint(
-            -1, 3, (num_tokens,), dtype=torch.int32, device=DEVICE_TYPE
-        )
-        pos = num_tokens // 2
-        inputs[pos], mapping[pos] = target[0], 0
-        out = _bi_shrink(inputs, weights, mapping, 0.7)
-        assert torch.equal(out[:, pos], expected), num_tokens
-    assert len(reduce_grids) == 7
diff --git a/tests/v1/determinism/test_lora_shrink_batch_invariant.py b/tests/v1/determinism/test_lora_shrink_batch_invariant.py
new file mode 100644
index 000000000000..2a7029a60324
--- /dev/null
+++ b/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