perf(dcp): rotate destinations in the PCIe peer KV gather

.buildkite/test_areas/distributed.yaml

@@ -182,7 +182,7 @@ steps:
 
 - label: Distributed Tests (8xH100)
   key: distributed-tests-8xh100
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
   device: h100
   num_devices: 8
   working_dir: "/vllm-workspace/tests"
@@ -194,13 +194,19 @@ steps:
   - vllm/v1/engine/llm_engine.py
   - vllm/v1/executor/uniproc_executor.py
   - vllm/v1/worker/gpu_worker.py
+  - vllm/v1/attention/ops/dcp_utils.py
+  - vllm/model_executor/layers/attention/mla_attention.py
+  - vllm/model_executor/layers/attention/sparse_mla_attention.py
+  - csrc/libtorch_stable/attention/dcp_utils/
+  - tests/distributed/test_dcp_direct_a2a_lse_reduce.py
   - tests/distributed/test_mnnvl_alltoall.py
 
   commands:
   # https://github.com/NVIDIA/nccl/issues/1838
   - export NCCL_CUMEM_HOST_ENABLE=0
   # test with torchrun tp=2 and dp=4 with ep
   - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
+  - pytest -v -s distributed/test_dcp_direct_a2a_lse_reduce.py
 
 - label: Distributed Tests (4xA100)
   key: distributed-tests-4xa100

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());