perf(kquant): plan W4A16 trellis prefill on wide route blocks

csrc/libtorch_stable/attention/merge_attn_states.cu

@@ -36,18 +36,34 @@ __global__ void merge_attn_states_kernel(
   const uint pack_size = 16 / sizeof(scalar_t);
   const uint threads_per_head = head_size / pack_size;
 
-  const uint global_idx = blockIdx.x * NUM_THREADS + threadIdx.x;
+  const uint global_idx = blockIdx.x * blockDim.x + threadIdx.x;
   const uint token_head_threads = num_tokens * num_heads * threads_per_head;
 
-  if (global_idx >= token_head_threads) return;
-
-  // global_idx -> token_idx + head_idx + pack_idx
+  // Derive indices before the block barrier so every thread reaches it.
   const uint token_head_idx = global_idx / threads_per_head;
   const uint pack_idx = global_idx % threads_per_head;
-
   const uint token_idx = token_head_idx / num_heads;
   const uint head_idx = token_head_idx % num_heads;
 
+  // A running chunked-attention LSE may be both prefix_lse and output_lse.
+  // The launcher aligns block boundaries to complete head groups, allowing
+  // every group to load its LSE values before any thread overwrites them.
+  __shared__ float shared_prefix_lse[NUM_THREADS];
+  __shared__ float shared_suffix_lse[NUM_THREADS];
+  const bool is_valid = global_idx < token_head_threads;
+  const uint group_idx = threadIdx.x / threads_per_head;
+
+  if (is_valid && pack_idx == 0 && token_idx < prefix_num_tokens) {
+    shared_prefix_lse[group_idx] =
+        prefix_lse[head_idx * prefix_lse_head_stride +
+                   token_idx * prefix_lse_token_stride];
+    shared_suffix_lse[group_idx] =
+        suffix_lse[head_idx * suffix_lse_head_stride +
+                   token_idx * suffix_lse_token_stride];
+  }
+  __syncthreads();
+  if (!is_valid) return;
+
   const uint pack_offset = pack_idx * pack_size;  // (0~15)*8, etc.
   const uint src_head_offset = token_idx * num_heads * prefix_head_stride +
                                head_idx * prefix_head_stride;
@@ -95,11 +111,9 @@ __global__ void merge_attn_states_kernel(
     return;
   }
 
-  // For tokens within prefix range, merge prefix and suffix
-  float p_lse = prefix_lse[head_idx * prefix_lse_head_stride +
-                           token_idx * prefix_lse_token_stride];
-  float s_lse = suffix_lse[head_idx * suffix_lse_head_stride +
-                           token_idx * suffix_lse_token_stride];
+  // For tokens within prefix range, merge prefix and suffix.
+  float p_lse = shared_prefix_lse[group_idx];
+  float s_lse = shared_suffix_lse[group_idx];
   p_lse = std::isinf(p_lse) ? -std::numeric_limits<float>::infinity() : p_lse;
   s_lse = std::isinf(s_lse) ? -std::numeric_limits<float>::infinity() : s_lse;
 
@@ -307,10 +321,19 @@ void merge_attn_states_launcher(
   // Process one pack elements per thread. for float, the
   // pack_size is 4 for half/bf16, the pack_size is 8.
   const uint threads_per_head = head_size / pack_size;
+  STD_TORCH_CHECK(
+      threads_per_head <= NUM_THREADS,
+      "headsize requires more threads than the merge kernel block supports: ",
+      head_size);
   const uint total_threads = num_tokens * num_heads * threads_per_head;
+  // Keep each token-head group inside one block. This is required when
+  // output_lse aliases prefix_lse because the whole group must read the input
+  // LSE before its first thread writes the merged value.
+  const uint block_threads =
+      (NUM_THREADS / threads_per_head) * threads_per_head;
 
-  dim3 block(NUM_THREADS);
-  dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS);
+  dim3 block(block_threads);
+  dim3 grid((total_threads + block_threads - 1) / block_threads);
 
   const torch::stable::accelerator::DeviceGuard device_guard(
       prefix_output.get_device_index());

tests/distributed/test_flashinfer_pcie_all_reduce.py

@@ -21,6 +21,7 @@ def __init__(self, **kwargs: Any) -> None:
         self.kwargs = kwargs
         self.destroyed = False
         self.last_input: torch.Tensor | None = None
+        self.prepared: list[tuple[list[tuple[int, ...]], torch.dtype]] = []
         FakeWorkspace.instances.append(self)
 
     def supports(self, inp: torch.Tensor) -> bool:
@@ -35,6 +36,9 @@ def all_reduce(
         out.copy_(inp)
         return out
 
+    def prepare(self, shapes, *, dtype) -> None:
+        self.prepared.append((list(shapes), dtype))
+
     def destroy(self) -> None:
         self.destroyed = True
 
@@ -101,6 +105,7 @@ def test_capture_routes_graph_calls_without_reusing_eager_state() -> None:
     assert torch.equal(actual, inp)
     assert len(FakeWorkspace.instances) == 1
     assert FakeWorkspace.instances[0].last_input is inp
+    assert FakeWorkspace.instances[0].prepared == [([(1, 4)], torch.float32)]
     pool.close()
 
 

tests/distributed/test_pynccl.py

@@ -257,7 +257,9 @@ def all_gatherv_worker_fn():
     device = f"cuda:{pynccl_comm.rank}"
 
     assert world_size <= 8
-    sizes = [81, 20, 57, 52, 81, 5, 49, 49][:world_size]
+    # A zero-length rank is required when fewer multimodal inputs than TP
+    # ranks are distributed across the model-parallel group.
+    sizes = [81, 0, 57, 52, 81, 5, 49, 49][:world_size]
     num_elems = sizes[rank]
     tensor = torch.arange(num_elems, dtype=torch.float32, device=device) + rank * 100
     result = torch.zeros(sum(sizes), dtype=torch.float32, device=device)

tests/kernels/attention/test_merge_attn_states.py

@@ -99,6 +99,64 @@ def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
     assert not output.isnan().any()
 
 
+@pytest.mark.parametrize("num_tokens", [256, 4096])
+@pytest.mark.parametrize("head_size", [128, 192, 512])
+@torch.inference_mode()
+def test_merge_attn_states_cuda_inplace_accumulator(
+    num_tokens: int, head_size: int
+) -> None:
+    """The CUDA kernel supports a running partial as input and destination.
+
+    Chunked attention folds each suffix partial into one prefix allocation.
+    Both the attention output and its log-sum-exp tensor therefore alias their
+    corresponding destinations.  The result must match a merge into disjoint
+    output allocations exactly. The 192-element case is Kimi-K3's chunked
+    context-merge geometry and does not divide the kernel's 128-thread limit;
+    the 512-element case spans multiple warps per head.
+    """
+    if not current_platform.is_cuda():
+        pytest.skip("The custom merge-attention kernel requires CUDA")
+
+    torch.manual_seed(0)
+    num_heads = 6
+    shape = (num_tokens, num_heads, head_size)
+
+    prefix_output = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+    suffix_output = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+    prefix_lse = torch.randn(
+        (num_heads, num_tokens), dtype=torch.float32, device="cuda"
+    )
+    suffix_lse = torch.randn(
+        (num_heads, num_tokens), dtype=torch.float32, device="cuda"
+    )
+
+    reference_output = torch.empty_like(prefix_output)
+    reference_lse = torch.empty_like(prefix_lse)
+    merge_attn_states_cuda(
+        reference_output,
+        prefix_output,
+        prefix_lse,
+        suffix_output,
+        suffix_lse,
+        reference_lse,
+    )
+
+    inplace_output = prefix_output.clone()
+    inplace_lse = prefix_lse.clone()
+    merge_attn_states_cuda(
+        inplace_output,
+        inplace_output,
+        inplace_lse,
+        suffix_output,
+        suffix_lse,
+        inplace_lse,
+    )
+    torch.accelerator.synchronize()
+
+    torch.testing.assert_close(inplace_output, reference_output, rtol=0, atol=0)
+    torch.testing.assert_close(inplace_lse, reference_lse, rtol=0, atol=0)
+
+
 def generate_markdown_table():
     global all_case_info
     table_header = (

tests/models/kimi_k3/test_aux_attn_res_stream.py

@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Which value the DFlash drafter is fed under AttnRes.
+
+`_capture_aux_hidden_stream` picks the weights it mixes against from one of
+three places depending on where the tapped layer sits, and returns the plain
+running prefix when the feature is off. The mixture itself is the kernel's
+job and is covered by ``test_attn_res.py``; what is asserted here is the
+selection, which is the part that can silently feed the drafter the wrong
+tensor.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.models.kimi_k3.nvidia import model as k3_model
+
+END_LAYER = 4
+
+
+def _weights(tag: float) -> SimpleNamespace:
+    """A norm/projection pair that is identifiable by value."""
+    return SimpleNamespace(
+        weight=torch.full((2,), tag),
+        variance_epsilon=tag,
+    )
+
+
+def _stub_model(*, enabled: bool, use_attn_res: bool = True) -> SimpleNamespace:
+    """A stand-in carrying only what the tap reads.
+
+    Constructing the real model needs a distributed init and weights, and none
+    of it participates in the selection under test.
+    """
+    model = SimpleNamespace(
+        _aux_attn_res_stream=enabled,
+        use_attn_res=use_attn_res,
+        end_layer=END_LAYER,
+    )
+    if not use_attn_res:
+        return model
+
+    consumers = []
+    for i in range(END_LAYER):
+        consumers.append(
+            SimpleNamespace(
+                self_attention_res_norm=_weights(float(i)),
+                self_attention_res_proj=SimpleNamespace(
+                    weight=torch.full((1, 2), float(i))
+                ),
+                prev_valid_blocks=i,
+            )
+        )
+    model.layers = consumers
+    model.output_attn_res_norm = _weights(99.0)
+    model.output_attn_res_proj = SimpleNamespace(weight=torch.full((1, 2), 99.0))
+    model.num_attn_res_blocks = 99
+    return model
+
+
+@pytest.fixture
+def recorder(monkeypatch):
+    """Replace the kernel so the call it would have made is inspectable."""
+    calls = []
+
+    def _fake_attn_res(
+        prefix,
+        delta,
+        block_residual,
+        norm_weight,
+        proj_weight,
+        output_norm_weight,
+        **kwargs,
+    ):
+        calls.append(
+            SimpleNamespace(
+                prefix=prefix,
+                delta=delta,
+                block_residual=block_residual,
+                norm_weight=norm_weight,
+                proj_weight=proj_weight,
+                kwargs=kwargs,
+            )
+        )
+        return torch.full_like(prefix, -1.0)
+
+    monkeypatch.setattr(k3_model, "attn_res", _fake_attn_res)
+    return calls
+
+
+def _set_last_rank(monkeypatch, is_last: bool):
+    monkeypatch.setattr(
+        k3_model,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_last_rank=is_last),
+    )
+
+
+def _call(stub, layer_idx, prefix_sum, pending_mlp_out, block_residual):
+    return k3_model.KimiLinearModel._capture_aux_hidden_stream(
+        stub, layer_idx, prefix_sum, pending_mlp_out, block_residual
+    )
+
+
+@pytest.mark.parametrize(
+    "enabled,use_attn_res", [(False, True), (True, False), (False, False)]
+)
+def test_disabled_reproduces_the_plain_residual_sum(
+    recorder, monkeypatch, enabled, use_attn_res
+):
+    """Off, the tap must be exactly the sum it replaced.
+
+    Both conditions matter. `use_attn_res` is what constructs the norm and
+    projection weights, so without it the lookups below would raise rather
+    than fall back.
+    """
+    _set_last_rank(monkeypatch, True)
+    prefix_sum = torch.tensor([1.0, 2.0])
+    pending = torch.tensor([0.5, 0.25])
+
+    got = _call(
+        _stub_model(enabled=enabled, use_attn_res=use_attn_res),
+        0,
+        prefix_sum,
+        pending,
+        torch.zeros(2),
+    )
+
+    torch.testing.assert_close(got, prefix_sum + pending)
+    assert not recorder, "the kernel must not run when the tap is off"
+
+
+def test_taps_the_consumer_layer_when_one_follows(recorder, monkeypatch):
+    """The value the next layer reads is the mixture against *its* weights,
+    so the tap has to reach forward rather than use the current layer's."""
+    _set_last_rank(monkeypatch, True)
+
+    _call(_stub_model(enabled=True), 1, torch.zeros(2), None, torch.zeros(2))
+
+    assert len(recorder) == 1
+    call = recorder[0]
+    # Layer 2's weights, not layer 1's.
+    torch.testing.assert_close(call.norm_weight, torch.full((2,), 2.0))
+    torch.testing.assert_close(call.proj_weight, torch.full((2,), 2.0))
+    assert call.kwargs["num_blocks"] == 2
+
+
+def test_last_layer_on_the_final_rank_uses_the_output_aggregation(
+    recorder, monkeypatch
+):
+    """Nothing downstream but the model's own output-side mixture."""
+    _set_last_rank(monkeypatch, True)
+
+    _call(
+        _stub_model(enabled=True), END_LAYER - 1, torch.zeros(2), None, torch.zeros(2)
+    )
+
+    assert len(recorder) == 1
+    torch.testing.assert_close(recorder[0].norm_weight, torch.full((2,), 99.0))
+    torch.testing.assert_close(recorder[0].proj_weight, torch.full((2,), 99.0))
+    assert recorder[0].kwargs["num_blocks"] == 99
+
+
+def test_last_layer_of_a_non_final_stage_falls_back(recorder, monkeypatch):
+    """The consumer lives on the next rank and the output aggregation only
+    exists on the last one, so there is nothing here to mix against.
+
+    This is the case that would otherwise reach for weights this rank never
+    constructs. The forward guard is `layer_idx + 1 < end_layer`, where
+    `end_layer` is the rank's own exclusive bound from `get_pp_indices`, so a
+    `PPMissingLayer` is unreachable by construction -- the fallback below is
+    what makes that true rather than merely likely.
+    """
+    _set_last_rank(monkeypatch, False)
+    prefix_sum = torch.tensor([3.0, 4.0])
+
+    got = _call(
+        _stub_model(enabled=True), END_LAYER - 1, prefix_sum, None, torch.zeros(2)
+    )
+
+    torch.testing.assert_close(got, prefix_sum)
+    assert not recorder, "no weights exist on this rank to mix against"
+
+
+def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta(
+    recorder, monkeypatch
+):
+    """The kernel writes an applied delta back into the prefix in place, which
+    would double-add it into the live residual stream, so the pending output
+    has to arrive already summed into the prefix with `delta` left None."""
+    _set_last_rank(monkeypatch, True)
+    prefix_sum = torch.tensor([1.0, 2.0])
+    pending = torch.tensor([0.5, 0.25])
+
+    _call(_stub_model(enabled=True), 0, prefix_sum, pending, torch.zeros(2))
+
+    assert len(recorder) == 1
+    assert recorder[0].delta is None
+    torch.testing.assert_close(recorder[0].prefix, prefix_sum + pending)
+    # And the caller's tensor is not mutated on the way.
+    torch.testing.assert_close(prefix_sum, torch.tensor([1.0, 2.0]))