perf(kimi-k3): DMA-ring TP all-reduces and deferred draft ingest for chunked prefill

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_b12x_fused_all_reduce.py

@@ -231,6 +231,7 @@ def test_b12x_oneshot_defaults_to_stream_isolation(
     ),
     [
         (2, "oneshot", True, False, 2),
+        (2, "oneshot", None, False, 2),
         (16, "hierarchical", False, False, 1),
         (16, "hierarchical", False, True, 1),
     ],
@@ -239,14 +240,17 @@ def test_b12x_dispatcher_prepares_single_stable_eager_owner(
     monkeypatch: pytest.MonkeyPatch,
     world_size: int,
     algorithm: str,
-    supports_all_peer_auxiliary: bool,
+    supports_all_peer_auxiliary: bool | None,
     single_channel: bool,
     expected_channel_capacity: int,
 ) -> None:
     captured = {}
     runtime = MagicMock()
     runtime.algorithm = algorithm
-    runtime.supports_all_peer_auxiliary = supports_all_peer_auxiliary
+    if supports_all_peer_auxiliary is not None:
+        runtime.supports_all_peer_auxiliary = supports_all_peer_auxiliary
+    else:
+        del runtime.supports_all_peer_auxiliary
     dma_min_bytes = MagicMock(return_value=None)
 
     class FakeDispatcher:
@@ -367,7 +371,7 @@ def fake_all_gather(gather_list, tensor, *, group):
     assert built.backend_name() == (
         "B12X_PCIE_ONESHOT" if algorithm == "oneshot" else "B12X_PCIE_HIERARCHICAL"
     )
-    assert dma_min_bytes.call_count == int(supports_all_peer_auxiliary)
+    assert dma_min_bytes.call_count == int(supports_all_peer_auxiliary is True)
 
 
 def test_b12x_channel_checkpoint_delegates_to_runtime() -> None:

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
 
@@ -84,6 +88,17 @@ def test_each_semantic_channel_gets_an_independent_workspace() -> None:
     assert eager.destroyed
 
 
+def test_all_peer_capability_requires_a_prepared_workspace() -> None:
+    pool = make_pool()
+    assert not pool.supports_all_peer_auxiliary
+
+    pool.prepare_channels(("vllm:target:production",))
+    assert pool.supports_all_peer_auxiliary
+
+    pool.close()
+    assert not pool.supports_all_peer_auxiliary
+
+
 def test_capture_routes_graph_calls_without_reusing_eager_state() -> None:
     pool = make_pool()
     inp = torch.arange(4, dtype=torch.float32)
@@ -101,6 +116,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()
 
 
@@ -168,6 +184,8 @@ def test_custom_allreduce_constructs_flashinfer_pool(
 ) -> None:
     captured: dict[str, Any] = {}
     runtime = MagicMock()
+    runtime.supports_all_peer_auxiliary = True
+    dma_min_bytes = MagicMock(return_value=None)
 
     class FakePool:
         @classmethod
@@ -211,6 +229,7 @@ def fake_all_gather(gather_list, tensor, *, group) -> None:
     monkeypatch.setattr(
         custom_all_reduce, "_load_flashinfer_pcie_oneshot_pool", lambda: FakePool
     )
+    monkeypatch.setattr(custom_all_reduce, "_b12x_pcie_dma_min_bytes", dma_min_bytes)
     monkeypatch.setattr(
         custom_all_reduce.current_platform, "get_device_capability", lambda: None
     )
@@ -255,3 +274,4 @@ def fake_all_gather(gather_list, tensor, *, group) -> None:
     assert captured["max_size"] > 0
     runtime.prepare_channels.assert_called_once()
     runtime.for_stream.assert_called_once()
+    dma_min_bytes.assert_called_once_with()

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 = (