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

docs/features/batch_invariance.md

@@ -150,6 +150,7 @@ When batch invariance is enabled, vLLM:
 4. 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).
+5. 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.

tests/lora/test_punica_ops.py

@@ -639,3 +639,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

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

vllm/lora/ops/triton_ops/lora_shrink_op.py

@@ -13,11 +13,90 @@
 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
 
+# Batch-invariant shrink splits K into this many fixed partials (no atomics).
+_BI_SPLIT_K = 8
+
+
+@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 +128,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 +203,7 @@ def _lora_shrink_kernel(
         output_d0_stride,
         output_d1_stride,
         output_d2_stride,
+        output_split_stride,
         scaling,
         BLOCK_M,
         BLOCK_N,
@@ -124,6 +212,7 @@ def _lora_shrink_kernel(
         SPLIT_K,
         SLICE_NUM,
         USE_GDC,
+        STORE_PARTIALS,
     )
 
 
@@ -167,6 +256,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 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),
+                _BI_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 +319,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 partials is not None:
+        BLOCK_K = 256
+        SPLIT_K = _BI_SPLIT_K
     EVEN_K = K % (BLOCK_K * SPLIT_K) == 0  # type: ignore
 
     # TODO (varun): This grid formulation maximizes parallelization at the
@@ -225,11 +336,24 @@ 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:
+        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 +367,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 +382,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