@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 63ac04a61e

[Kimi-K3] DCP support (#50484)

Signed-off-by: girasoley <girasoley@inferact.ai> Signed-off-by: Summer Yang <girasoleyang@gmail.com> Signed-off-by: 云挚 <ningyunxiao.nyx@antgroup.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: 云挚 <ningyunxiao.nyx@antgroup.com> Co-authored-by: foraxe <73625538+foraxe@users.noreply.github.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>

Summer Yang committed Aug 10, 2026 at 12:54 UTC 63ac04a61e6272b746d82e180e7f6276b3ee1ede
19 files changed +3529 -82
CMakeLists.txt
+5
@@ -474,6 +474,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
474 endif()
475
476 if(VLLM_GPU_LANG STREQUAL "CUDA")
477 + list(APPEND VLLM_STABLE_EXT_SRC
478 + "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu"
479 + "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu"
480 + "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu")
481 +
482 SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
483
484 # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu new
+381
@@ -0,0 +1,381 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +// Direct symmetric-memory DCP A2A LSE reduction.
4 +
5 +#include <cuda_bf16.h>
6 +#include <cuda_fp16.h>
7 +#include <math_constants.h>
8 +#include <torch/csrc/stable/library.h>
9 +#include <torch/headeronly/core/ScalarType.h>
10 +
11 +#include <cstdio>
12 +#include <optional>
13 +#include <type_traits>
14 +
15 +#include "dcp_direct_common.cuh"
16 +
17 +namespace {
18 +
19 +using vllm::direct_dcp::check_cuda_launch;
20 +using vllm::direct_dcp::get_peer_ptr;
21 +using vllm::direct_dcp::increment_epoch_kernel;
22 +using vllm::direct_dcp::store_release_system;
23 +using vllm::direct_dcp::wait_for_epoch;
24 +
25 +__device__ __forceinline__ int64_t find_sequence(const int32_t* query_start_loc,
26 + int64_t token_idx,
27 + int64_t num_seqs) {
28 + int64_t left = 0;
29 + int64_t right = num_seqs;
30 + while (left < right) {
31 + int64_t mid = (left + right) / 2;
32 + if (query_start_loc[mid] <= token_idx) {
33 + left = mid + 1;
34 + } else {
35 + right = mid;
36 + }
37 + }
38 + return left - 1;
39 +}
40 +
41 +template <typename scalar_t>
42 +__device__ __forceinline__ float to_float(scalar_t value) {
43 + if constexpr (std::is_same_v<scalar_t, float>) {
44 + return value;
45 + } else if constexpr (std::is_same_v<scalar_t, __half>) {
46 + return __half2float(value);
47 + } else {
48 + return __bfloat162float(value);
49 + }
50 +}
51 +
52 +template <typename scalar_t>
53 +__device__ __forceinline__ scalar_t from_float(float value) {
54 + if constexpr (std::is_same_v<scalar_t, __half>) {
55 + return __float2half_rn(value);
56 + } else {
57 + return __float2bfloat16_rn(value);
58 + }
59 +}
60 +
61 +// Send each destination's head slice and LSE to its symmetric buffer. Mask
62 +// undefined empty-shard rows at the source so peers give them zero weight.
63 +template <typename lse_t>
64 +__global__ void dispatch_output_lse_kernel(
65 + const uint4* partial_output, const lse_t* partial_lse,
66 + const int32_t* seq_lens, const int32_t* query_start_loc,
67 + const int64_t* peer_output_ptrs, const int64_t* peer_lse_ptrs,
68 + const int64_t* epoch_ptr, int64_t world_size, int64_t rank,
69 + int64_t num_tokens, int64_t max_num_tokens, int64_t num_seqs,
70 + int64_t heads_per_rank, int64_t head_dim, int64_t output_token_stride,
71 + int64_t lse_token_stride) {
72 + int64_t item = static_cast<int64_t>(blockIdx.x);
73 + int64_t destination_rank = item / num_tokens;
74 + int64_t token_idx = item - destination_rank * num_tokens;
75 +
76 + __shared__ bool empty_kv;
77 + if (threadIdx.x == 0) {
78 + empty_kv = false;
79 + if (seq_lens != nullptr) {
80 + int64_t seq_idx = find_sequence(query_start_loc, token_idx, num_seqs);
81 + empty_kv = seq_lens[seq_idx] == 0;
82 + }
83 + }
84 + __syncthreads();
85 +
86 + uint32_t epoch = static_cast<uint32_t>(epoch_ptr[0]);
87 + int64_t parity = static_cast<int64_t>(epoch & 1u);
88 + int64_t destination_item =
89 + ((parity * world_size + rank) * max_num_tokens + token_idx) *
90 + heads_per_rank;
91 + int64_t source_head = destination_rank * heads_per_rank;
92 +
93 + if (!empty_kv) {
94 + uint4* peer_output =
95 + get_peer_ptr<uint4>(peer_output_ptrs, destination_rank);
96 + int64_t vectors_per_item = heads_per_rank * head_dim / 8;
97 + int64_t source_vector =
98 + (token_idx * output_token_stride + source_head * head_dim) / 8;
99 + int64_t destination_vector = destination_item * head_dim / 8;
100 + for (int64_t vector_idx = threadIdx.x; vector_idx < vectors_per_item;
101 + vector_idx += blockDim.x) {
102 + peer_output[destination_vector + vector_idx] =
103 + partial_output[source_vector + vector_idx];
104 + }
105 + }
106 +
107 + float* peer_lse = get_peer_ptr<float>(peer_lse_ptrs, destination_rank);
108 + int64_t source_lse = token_idx * lse_token_stride + source_head;
109 + for (int64_t head_idx = threadIdx.x; head_idx < heads_per_rank;
110 + head_idx += blockDim.x) {
111 + peer_lse[destination_item + head_idx] =
112 + empty_kv ? -CUDART_INF_F : to_float(partial_lse[source_lse + head_idx]);
113 + }
114 +}
115 +
116 +// Publish completion after the stream-ordered payload writes.
117 +__global__ void signal_kernel(const int64_t* peer_signal_ptrs,
118 + const int64_t* epoch_ptr, int64_t world_size,
119 + int64_t rank) {
120 + int64_t destination_rank =
121 + static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
122 + if (destination_rank >= world_size) {
123 + return;
124 + }
125 +
126 + uint32_t epoch = static_cast<uint32_t>(epoch_ptr[0]);
127 + int64_t parity = static_cast<int64_t>(epoch & 1u);
128 + uint32_t* peer_signal =
129 + get_peer_ptr<uint32_t>(peer_signal_ptrs, destination_rank);
130 + int64_t signal_item = parity * world_size + rank;
131 + store_release_system(peer_signal + signal_item, epoch);
132 +}
133 +
134 +// Form stable base-2 LSE weights and combine in FP32.
135 +template <typename scalar_t>
136 +__global__ void wait_lse_combine_kernel(
137 + const scalar_t* received_output, const float* received_lse,
138 + const uint32_t* received_signal, const int64_t* epoch_ptr,
139 + scalar_t* combined_output, int64_t world_size, int64_t num_tokens,
140 + int64_t max_num_tokens, int64_t heads_per_rank, int64_t head_dim,
141 + bool is_lse_base_on_e) {
142 + extern __shared__ float weights[];
143 +
144 + int64_t item = static_cast<int64_t>(blockIdx.x);
145 + int64_t token_idx = item / heads_per_rank;
146 + int64_t head_idx = item - token_idx * heads_per_rank;
147 + uint32_t epoch = static_cast<uint32_t>(epoch_ptr[0]);
148 + int64_t parity = static_cast<int64_t>(epoch & 1u);
149 +
150 + if (threadIdx.x == 0) {
151 + float lse_max = -CUDART_INF_F;
152 + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) {
153 + int64_t source_item =
154 + ((parity * world_size + source_rank) * max_num_tokens + token_idx) *
155 + heads_per_rank +
156 + head_idx;
157 + int64_t signal_item = parity * world_size + source_rank;
158 + if (!wait_for_epoch(received_signal + signal_item, epoch)) {
159 + printf(
160 + "direct DCP A2A timeout source=%lld token=%lld head=%lld "
161 + "epoch=%u\n",
162 + static_cast<long long>(source_rank),
163 + static_cast<long long>(token_idx), static_cast<long long>(head_idx),
164 + epoch);
165 + asm volatile("trap;");
166 + }
167 + float value = received_lse[source_item];
168 + if (isnan(value) || value == CUDART_INF_F) {
169 + value = -CUDART_INF_F;
170 + }
171 + if (is_lse_base_on_e) {
172 + value *= CUDART_L2E_F;
173 + }
174 + weights[source_rank] = value;
175 + lse_max = fmaxf(lse_max, value);
176 + }
177 + if (lse_max == -CUDART_INF_F) {
178 + lse_max = 0.0f;
179 + }
180 +
181 + float lse_sum = 0.0f;
182 + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) {
183 + lse_sum += exp2f(weights[source_rank] - lse_max);
184 + }
185 + float inverse_lse_sum = lse_sum > 0.0f ? 1.0f / lse_sum : 0.0f;
186 + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) {
187 + weights[source_rank] =
188 + exp2f(weights[source_rank] - lse_max) * inverse_lse_sum;
189 + }
190 + }
191 + __syncthreads();
192 +
193 + for (int64_t dim_idx = threadIdx.x; dim_idx < head_dim;
194 + dim_idx += blockDim.x) {
195 + float accumulator = 0.0f;
196 + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) {
197 + // Empty shards leave stale payloads, so never read zero-weight sources.
198 + float weight = weights[source_rank];
199 + if (weight == 0.0f) {
200 + continue;
201 + }
202 + int64_t source_item =
203 + ((parity * world_size + source_rank) * max_num_tokens + token_idx) *
204 + heads_per_rank +
205 + head_idx;
206 + accumulator +=
207 + to_float(received_output[source_item * head_dim + dim_idx]) * weight;
208 + }
209 + combined_output[item * head_dim + dim_idx] =
210 + from_float<scalar_t>(accumulator);
211 + }
212 +}
213 +
214 +void direct_dcp_a2a_lse_reduce(
215 + const torch::stable::Tensor& partial_output,
216 + const torch::stable::Tensor& partial_lse,
217 + const std::optional<torch::stable::Tensor>& seq_lens,
218 + const std::optional<torch::stable::Tensor>& query_start_loc,
219 + const torch::stable::Tensor& peer_output_ptrs,
220 + const torch::stable::Tensor& peer_lse_ptrs,
221 + const torch::stable::Tensor& peer_signal_ptrs,
222 + torch::stable::Tensor& received_output, torch::stable::Tensor& received_lse,
223 + torch::stable::Tensor& received_signal, torch::stable::Tensor& epoch,
224 + torch::stable::Tensor& combined_output, int64_t world_size, int64_t rank,
225 + int64_t max_num_tokens, bool is_lse_base_on_e) {
226 + STD_TORCH_CHECK(partial_output.is_cuda() && partial_lse.is_cuda(),
227 + "partial output and LSE must be CUDA tensors");
228 + auto output_dtype = partial_output.scalar_type();
229 + STD_TORCH_CHECK(output_dtype == torch::headeronly::ScalarType::Half ||
230 + output_dtype == torch::headeronly::ScalarType::BFloat16,
231 + "direct DCP A2A only supports FP16 and BF16 output");
232 + auto lse_dtype = partial_lse.scalar_type();
233 + STD_TORCH_CHECK(lse_dtype == torch::headeronly::ScalarType::Float ||
234 + lse_dtype == torch::headeronly::ScalarType::Half ||
235 + lse_dtype == torch::headeronly::ScalarType::BFloat16,
236 + "partial LSE must be FP32, FP16, or BF16");
237 + STD_TORCH_CHECK(partial_output.dim() == 3 && partial_lse.dim() == 2,
238 + "expected output [T,H,D] and LSE [T,H]");
239 + STD_TORCH_CHECK(world_size > 1, "world_size must be greater than 1");
240 + STD_TORCH_CHECK(rank >= 0 && rank < world_size, "invalid rank");
241 +
242 + int64_t num_tokens = partial_output.size(0);
243 + int64_t total_heads = partial_output.size(1);
244 + int64_t head_dim = partial_output.size(2);
245 + int64_t output_token_stride = partial_output.stride(0);
246 + int64_t lse_token_stride = partial_lse.stride(0);
247 + STD_TORCH_CHECK(
248 + partial_output.stride(2) == 1 && partial_output.stride(1) == head_dim &&
249 + output_token_stride >= total_heads * head_dim &&
250 + output_token_stride % 8 == 0,
251 + "partial output must have packed heads and an aligned token stride");
252 + STD_TORCH_CHECK(partial_lse.stride(1) == 1 && lse_token_stride >= total_heads,
253 + "partial LSE must have packed heads");
254 + STD_TORCH_CHECK(num_tokens > 0 && num_tokens <= max_num_tokens,
255 + "token count exceeds symmetric buffer capacity");
256 + STD_TORCH_CHECK(total_heads % world_size == 0,
257 + "attention heads must divide evenly across DCP ranks");
258 + STD_TORCH_CHECK(
259 + partial_lse.size(0) == num_tokens && partial_lse.size(1) == total_heads,
260 + "LSE shape must match attention output");
261 + STD_TORCH_CHECK(head_dim % 8 == 0,
262 + "head_dim must be divisible by 8 for 16-byte stores");
263 + int64_t heads_per_rank = total_heads / world_size;
264 + STD_TORCH_CHECK(combined_output.scalar_type() == output_dtype &&
265 + combined_output.is_contiguous() &&
266 + combined_output.is_cuda(),
267 + "combined output must match the contiguous CUDA input");
268 + STD_TORCH_CHECK(combined_output.size(0) == num_tokens &&
269 + combined_output.size(1) == heads_per_rank &&
270 + combined_output.size(2) == head_dim,
271 + "combined output has the wrong shape");
272 + STD_TORCH_CHECK(
273 + received_output.scalar_type() == output_dtype &&
274 + received_lse.scalar_type() == torch::headeronly::ScalarType::Float &&
275 + received_signal.scalar_type() == torch::headeronly::ScalarType::Int,
276 + "invalid symmetric staging dtypes");
277 + STD_TORCH_CHECK(
278 + peer_output_ptrs.scalar_type() == torch::headeronly::ScalarType::Long &&
279 + peer_lse_ptrs.scalar_type() == torch::headeronly::ScalarType::Long &&
280 + peer_signal_ptrs.scalar_type() == torch::headeronly::ScalarType::Long,
281 + "peer pointer tables must be int64");
282 + const int32_t* seq_lens_ptr = nullptr;
283 + const int32_t* query_start_loc_ptr = nullptr;
284 + int64_t num_seqs = 0;
285 + STD_TORCH_CHECK(seq_lens.has_value() == query_start_loc.has_value(),
286 + "seq_lens and query_start_loc must be provided together");
287 + if (seq_lens.has_value() && query_start_loc.has_value()) {
288 + STD_TORCH_CHECK(
289 + seq_lens->is_cuda() && seq_lens->dim() == 1 &&
290 + seq_lens->stride(0) == 1 &&
291 + seq_lens->scalar_type() == torch::headeronly::ScalarType::Int,
292 + "seq_lens must be a contiguous 1-D int32 CUDA tensor");
293 + STD_TORCH_CHECK(
294 + query_start_loc->is_cuda() && query_start_loc->dim() == 1 &&
295 + query_start_loc->stride(0) == 1 &&
296 + query_start_loc->scalar_type() ==
297 + torch::headeronly::ScalarType::Int,
298 + "query_start_loc must be a contiguous 1-D int32 CUDA tensor");
299 + num_seqs = seq_lens->size(0);
300 + STD_TORCH_CHECK(num_seqs > 0 && query_start_loc->size(0) == num_seqs + 1,
301 + "query_start_loc must contain one boundary per sequence");
302 + seq_lens_ptr = seq_lens->const_data_ptr<int32_t>();
303 + query_start_loc_ptr = query_start_loc->const_data_ptr<int32_t>();
304 + }
305 +
306 + const torch::stable::accelerator::DeviceGuard device_guard(
307 + partial_output.get_device_index());
308 + cudaStream_t stream = get_current_cuda_stream();
309 + constexpr int kExchangeThreads = 256;
310 + constexpr int kCombineThreads = 128;
311 +
312 + increment_epoch_kernel<<<1, 1, 0, stream>>>(
313 + epoch.mutable_data_ptr<int64_t>());
314 + check_cuda_launch("direct DCP A2A");
315 + int64_t dispatch_blocks = world_size * num_tokens;
316 + auto launch_dispatch = [&]<typename lse_t>() {
317 + dispatch_output_lse_kernel<lse_t>
318 + <<<dispatch_blocks, kExchangeThreads, 0, stream>>>(
319 + reinterpret_cast<const uint4*>(partial_output.data_ptr()),
320 + reinterpret_cast<const lse_t*>(partial_lse.data_ptr()),
321 + seq_lens_ptr, query_start_loc_ptr,
322 + peer_output_ptrs.const_data_ptr<int64_t>(),
323 + peer_lse_ptrs.const_data_ptr<int64_t>(),
324 + epoch.const_data_ptr<int64_t>(), world_size, rank, num_tokens,
325 + max_num_tokens, num_seqs, heads_per_rank, head_dim,
326 + output_token_stride, lse_token_stride);
327 + };
328 + if (lse_dtype == torch::headeronly::ScalarType::Float) {
329 + launch_dispatch.operator()<float>();
330 + } else if (lse_dtype == torch::headeronly::ScalarType::BFloat16) {
331 + launch_dispatch.operator()<__nv_bfloat16>();
332 + } else {
333 + launch_dispatch.operator()<__half>();
334 + }
335 + check_cuda_launch("direct DCP A2A");
336 + int64_t signal_items = world_size;
337 + int64_t signal_blocks =
338 + (signal_items + kExchangeThreads - 1) / kExchangeThreads;
339 + signal_kernel<<<signal_blocks, kExchangeThreads, 0, stream>>>(
340 + peer_signal_ptrs.const_data_ptr<int64_t>(),
341 + epoch.const_data_ptr<int64_t>(), world_size, rank);
342 + check_cuda_launch("direct DCP A2A");
343 + int64_t combine_blocks = num_tokens * heads_per_rank;
344 + size_t shared_memory_bytes = world_size * sizeof(float);
345 + auto launch_combine = [&]<typename scalar_t>() {
346 + wait_lse_combine_kernel<scalar_t>
347 + <<<combine_blocks, kCombineThreads, shared_memory_bytes, stream>>>(
348 + reinterpret_cast<const scalar_t*>(received_output.data_ptr()),
349 + received_lse.const_data_ptr<float>(),
350 + reinterpret_cast<const uint32_t*>(
351 + received_signal.const_data_ptr<int32_t>()),
352 + epoch.const_data_ptr<int64_t>(),
353 + reinterpret_cast<scalar_t*>(combined_output.mutable_data_ptr()),
354 + world_size, num_tokens, max_num_tokens, heads_per_rank, head_dim,
355 + is_lse_base_on_e);
356 + };
357 + if (output_dtype == torch::headeronly::ScalarType::BFloat16) {
358 + launch_combine.operator()<__nv_bfloat16>();
359 + } else {
360 + launch_combine.operator()<__half>();
361 + }
362 + check_cuda_launch("direct DCP A2A");
363 +}
364 +
365 +} // namespace
366 +
367 +STABLE_TORCH_LIBRARY_FRAGMENT(_C, direct_dcp_a2a_ops) {
368 + direct_dcp_a2a_ops.def(
369 + "direct_dcp_a2a_lse_reduce("
370 + "Tensor partial_output, Tensor partial_lse, Tensor? seq_lens, "
371 + "Tensor? query_start_loc, Tensor peer_output_ptrs, "
372 + "Tensor peer_lse_ptrs, Tensor peer_signal_ptrs, Tensor! received_output, "
373 + "Tensor! received_lse, Tensor! received_signal, Tensor! epoch, "
374 + "Tensor! combined_output, int world_size, int rank, "
375 + "int max_num_tokens, bool is_lse_base_on_e) -> ()");
376 +}
377 +
378 +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, direct_dcp_a2a_ops) {
379 + direct_dcp_a2a_ops.impl("direct_dcp_a2a_lse_reduce",
380 + TORCH_BOX(&direct_dcp_a2a_lse_reduce));
381 +}
csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh new
+92
@@ -0,0 +1,92 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +#pragma once
5 +
6 +#include <cuda_runtime.h>
7 +
8 +#include <cstdint>
9 +#include <string>
10 +
11 +#include "../../torch_utils.h"
12 +
13 +namespace vllm::direct_dcp {
14 +
15 +constexpr uint64_t kSpinLimit = 100000000;
16 +
17 +// Advance the invocation ID; its low bit selects one of two staging slots.
18 +static __global__ void increment_epoch_kernel(int64_t* epoch) {
19 + if (blockIdx.x == 0 && threadIdx.x == 0) {
20 + epoch[0] += 1;
21 + }
22 +}
23 +
24 +template <typename T>
25 +__device__ __forceinline__ T* get_peer_ptr(const int64_t* peer_ptrs,
26 + int64_t peer) {
27 + return reinterpret_cast<T*>(static_cast<uintptr_t>(peer_ptrs[peer]));
28 +}
29 +
30 +// Replicate one 16-byte payload to every symmetric-buffer replica.
31 +__device__ __forceinline__ void multimem_store_16(uint4* mc_ptr, uint4 value) {
32 +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
33 + asm volatile("multimem.st.relaxed.sys.global.v4.f32 [%0], {%1,%2,%3,%4};"
34 + :
35 + : "l"(mc_ptr), "r"(value.x), "r"(value.y), "r"(value.z),
36 + "r"(value.w)
37 + : "memory");
38 +#else
39 + asm volatile("trap;");
40 +#endif
41 +}
42 +
43 +// Publish prior system-scope writes and signal every replica.
44 +__device__ __forceinline__ void multimem_store_release_system(uint32_t* mc_ptr,
45 + uint32_t value) {
46 +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
47 + asm volatile("multimem.st.release.sys.global.u32 [%0], %1;"
48 + :
49 + : "l"(mc_ptr), "r"(value)
50 + : "memory");
51 +#else
52 + asm volatile("trap;");
53 +#endif
54 +}
55 +
56 +__device__ __forceinline__ void store_release_system(uint32_t* ptr,
57 + uint32_t value) {
58 + uint64_t address = reinterpret_cast<uint64_t>(ptr);
59 + asm volatile("st.global.release.sys.u32 [%0], %1;"
60 + :
61 + : "l"(address), "r"(value)
62 + : "memory");
63 +}
64 +
65 +__device__ __forceinline__ uint32_t load_acquire_system(const uint32_t* ptr) {
66 + uint32_t value;
67 + uint64_t address = reinterpret_cast<uint64_t>(ptr);
68 + asm volatile("ld.global.acquire.sys.u32 %0, [%1];"
69 + : "=r"(value)
70 + : "l"(address)
71 + : "memory");
72 + return value;
73 +}
74 +
75 +__device__ __forceinline__ bool wait_for_epoch(const uint32_t* ptr,
76 + uint32_t epoch) {
77 + for (uint64_t spins = 0; spins < kSpinLimit; ++spins) {
78 + if (load_acquire_system(ptr) == epoch) {
79 + return true;
80 + }
81 + }
82 + return false;
83 +}
84 +
85 +inline void check_cuda_launch(const char* operation) {
86 + cudaError_t error = cudaGetLastError();
87 + STD_TORCH_CHECK(error == cudaSuccess,
88 + std::string(operation) +
89 + " kernel launch failed: " + cudaGetErrorString(error));
90 +}
91 +
92 +} // namespace vllm::direct_dcp
csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu new
+209
@@ -0,0 +1,209 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +// Direct symmetric-memory DCP KV gather.
4 +
5 +#include <torch/csrc/stable/library.h>
6 +#include <torch/headeronly/core/ScalarType.h>
7 +
8 +#include <cstdio>
9 +
10 +#include "dcp_direct_common.cuh"
11 +
12 +namespace {
13 +
14 +using vllm::direct_dcp::check_cuda_launch;
15 +using vllm::direct_dcp::increment_epoch_kernel;
16 +using vllm::direct_dcp::multimem_store_16;
17 +using vllm::direct_dcp::multimem_store_release_system;
18 +using vllm::direct_dcp::wait_for_epoch;
19 +
20 +constexpr int kThreads = 256;
21 +// KV chunks need many blocks in flight to saturate the fabric.
22 +constexpr int64_t kMaxMulticastBlocks = 128;
23 +constexpr int64_t kMaxCopyBlocks = 1024;
24 +
25 +// Multicast each rank's KV slice and completion epoch to every replica.
26 +__global__ void direct_dcp_kv_gather_multimem_kernel(
27 + const uint4* local_kv, uint4* mc_kv, uint32_t* mc_signal,
28 + const uint32_t* received_signal, const int64_t* epoch_ptr,
29 + uint32_t* completion, int64_t world_size, int64_t rank, int64_t slice_bytes,
30 + int64_t slot_stride_bytes) {
31 + uint32_t epoch = static_cast<uint32_t>(epoch_ptr[0]);
32 + int64_t buffer_slot = static_cast<int64_t>(epoch & 1u);
33 + int64_t item_count = slice_bytes / sizeof(uint4);
34 + mc_kv += buffer_slot * slot_stride_bytes / sizeof(uint4);
35 + mc_kv += rank * item_count;
36 +
37 + int64_t item_stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
38 + for (int64_t item =
39 + static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
40 + item < item_count; item += item_stride) {
41 + multimem_store_16(mc_kv + item, local_kv[item]);
42 + }
43 +
44 + // Publish all multicast writes before incrementing completion.
45 + __threadfence_system();
46 + __syncthreads();
47 + if (threadIdx.x != 0) {
48 + return;
49 + }
50 +
51 + uint32_t completed = atomicAdd(completion + buffer_slot, 1u);
52 + if (completed + 1u != gridDim.x) {
53 + return;
54 + }
55 + atomicExch(completion + buffer_slot, 0u);
56 +
57 + multimem_store_release_system(mc_signal + buffer_slot * world_size + rank,
58 + epoch);
59 +
60 + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) {
61 + int64_t signal_item = buffer_slot * world_size + source_rank;
62 + if (!wait_for_epoch(received_signal + signal_item, epoch)) {
63 + printf("direct DCP kv-gather multimem timeout source=%lld epoch=%u\n",
64 + static_cast<long long>(source_rank), epoch);
65 + asm volatile("trap;");
66 + }
67 + }
68 +}
69 +
70 +// Materialize the acquired slot into the caller's workspace.
71 +template <typename copy_t>
72 +__global__ void materialize_kv_gather_kernel(const copy_t* received_kv,
73 + const int64_t* epoch_ptr,
74 + copy_t* gathered_kv,
75 + int64_t gathered_item_count,
76 + int64_t slot_stride_items) {
77 + uint32_t epoch = static_cast<uint32_t>(epoch_ptr[0]);
78 + int64_t buffer_slot = static_cast<int64_t>(epoch & 1u);
79 + received_kv += buffer_slot * slot_stride_items;
80 + int64_t item = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
81 + int64_t item_stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
82 + for (; item < gathered_item_count; item += item_stride) {
83 + gathered_kv[item] = received_kv[item];
84 + }
85 +}
86 +
87 +void direct_dcp_kv_gather(
88 + const torch::stable::Tensor& local_kv, torch::stable::Tensor& received_kv,
89 + torch::stable::Tensor& received_signal, torch::stable::Tensor& completion,
90 + torch::stable::Tensor& epoch, torch::stable::Tensor& gathered_kv,
91 + int64_t world_size, int64_t rank, int64_t max_gathered_tokens,
92 + int64_t kv_mc_ptr, int64_t signal_mc_ptr) {
93 + using torch::headeronly::ScalarType;
94 +
95 + STD_TORCH_CHECK(local_kv.is_cuda(), "local kv must be a CUDA tensor");
96 + ScalarType dtype = local_kv.scalar_type();
97 + STD_TORCH_CHECK(dtype == ScalarType::Half || dtype == ScalarType::BFloat16 ||
98 + dtype == ScalarType::Float8_e4m3fn,
99 + "direct DCP kv-gather only supports FP16, BF16, and FP8");
100 + STD_TORCH_CHECK(local_kv.dim() == 2 && local_kv.is_contiguous(),
101 + "local kv must be a contiguous [T,D] tensor");
102 + STD_TORCH_CHECK(world_size > 1, "world_size must be greater than 1");
103 + STD_TORCH_CHECK(rank >= 0 && rank < world_size, "invalid rank");
104 +
105 + int64_t num_tokens = local_kv.size(0);
106 + int64_t token_dim = local_kv.size(1);
107 + int64_t element_size = local_kv.element_size();
108 + STD_TORCH_CHECK(num_tokens > 0 && token_dim > 0,
109 + "local kv dimensions must be positive");
110 + STD_TORCH_CHECK(num_tokens * world_size <= max_gathered_tokens,
111 + "gathered kv exceeds symmetric kv-gather buffer capacity");
112 +
113 + STD_TORCH_CHECK(received_kv.is_cuda() && received_kv.scalar_type() == dtype &&
114 + received_kv.is_contiguous() && received_kv.dim() == 3 &&
115 + received_kv.size(0) == 2 &&
116 + received_kv.size(1) == max_gathered_tokens &&
117 + received_kv.size(2) == token_dim,
118 + "received kv has the wrong symmetric buffer layout");
119 + STD_TORCH_CHECK(
120 + received_signal.is_cuda() && received_signal.is_contiguous() &&
121 + received_signal.scalar_type() == ScalarType::Int &&
122 + received_signal.dim() == 2 && received_signal.size(0) == 2 &&
123 + received_signal.size(1) == world_size,
124 + "received signal has the wrong symmetric buffer layout");
125 + STD_TORCH_CHECK(completion.is_cuda() && completion.is_contiguous() &&
126 + completion.scalar_type() == ScalarType::Int &&
127 + completion.numel() == 2,
128 + "completion counter must be a two-element CUDA int32 tensor");
129 + STD_TORCH_CHECK(epoch.is_cuda() && epoch.is_contiguous() &&
130 + epoch.scalar_type() == ScalarType::Long &&
131 + epoch.numel() == 1,
132 + "epoch must be a one-element CUDA int64 tensor");
133 + STD_TORCH_CHECK(gathered_kv.is_cuda() && gathered_kv.scalar_type() == dtype &&
134 + gathered_kv.is_contiguous() && gathered_kv.dim() == 2 &&
135 + gathered_kv.size(0) == num_tokens * world_size &&
136 + gathered_kv.size(1) == token_dim,
137 + "gathered kv must be contiguous with shape [world_size*T,D]");
138 + int64_t device_index = local_kv.get_device_index();
139 + STD_TORCH_CHECK(
140 + received_kv.get_device_index() == device_index &&
141 + received_signal.get_device_index() == device_index &&
142 + completion.get_device_index() == device_index &&
143 + epoch.get_device_index() == device_index &&
144 + gathered_kv.get_device_index() == device_index,
145 + "direct DCP kv-gather tensors must be on the same CUDA device");
146 +
147 + int64_t slice_bytes = num_tokens * token_dim * element_size;
148 + int64_t slot_stride_bytes = max_gathered_tokens * token_dim * element_size;
149 + bool vectorized =
150 + reinterpret_cast<uintptr_t>(local_kv.data_ptr()) % alignof(uint4) == 0 &&
151 + reinterpret_cast<uintptr_t>(received_kv.data_ptr()) % alignof(uint4) ==
152 + 0 &&
153 + reinterpret_cast<uintptr_t>(gathered_kv.data_ptr()) % alignof(uint4) ==
154 + 0 &&
155 + slice_bytes % sizeof(uint4) == 0 &&
156 + slot_stride_bytes % sizeof(uint4) == 0;
157 + STD_TORCH_CHECK(
158 + vectorized,
159 + "direct DCP kv-gather requires 16-byte-aligned pointers and strides");
160 + STD_TORCH_CHECK(kv_mc_ptr != 0 && signal_mc_ptr != 0,
161 + "direct DCP kv-gather requires multicast pointers");
162 +
163 + const torch::stable::accelerator::DeviceGuard device_guard(device_index);
164 + cudaStream_t stream = get_current_cuda_stream();
165 + increment_epoch_kernel<<<1, 1, 0, stream>>>(
166 + epoch.mutable_data_ptr<int64_t>());
167 + check_cuda_launch("direct DCP kv-gather");
168 +
169 + int64_t item_count = slice_bytes / sizeof(uint4);
170 + int64_t blocks = (item_count + kThreads - 1) / kThreads;
171 + blocks = blocks < kMaxMulticastBlocks ? blocks : kMaxMulticastBlocks;
172 + direct_dcp_kv_gather_multimem_kernel<<<blocks, kThreads, 0, stream>>>(
173 + reinterpret_cast<const uint4*>(local_kv.data_ptr()),
174 + reinterpret_cast<uint4*>(static_cast<uintptr_t>(kv_mc_ptr)),
175 + reinterpret_cast<uint32_t*>(static_cast<uintptr_t>(signal_mc_ptr)),
176 + reinterpret_cast<const uint32_t*>(
177 + received_signal.const_data_ptr<int32_t>()),
178 + epoch.const_data_ptr<int64_t>(),
179 + reinterpret_cast<uint32_t*>(completion.mutable_data_ptr<int32_t>()),
180 + world_size, rank, slice_bytes, slot_stride_bytes);
181 + check_cuda_launch("direct DCP kv-gather");
182 +
183 + int64_t gathered_item_count = world_size * item_count;
184 + int64_t copy_blocks = (gathered_item_count + kThreads - 1) / kThreads;
185 + copy_blocks = copy_blocks < kMaxCopyBlocks ? copy_blocks : kMaxCopyBlocks;
186 + materialize_kv_gather_kernel<uint4><<<copy_blocks, kThreads, 0, stream>>>(
187 + reinterpret_cast<const uint4*>(received_kv.data_ptr()),
188 + epoch.const_data_ptr<int64_t>(),
189 + reinterpret_cast<uint4*>(gathered_kv.mutable_data_ptr()),
190 + gathered_item_count, slot_stride_bytes / sizeof(uint4));
191 + check_cuda_launch("direct DCP kv-gather");
192 +}
193 +
194 +} // namespace
195 +
196 +STABLE_TORCH_LIBRARY_FRAGMENT(_C, direct_dcp_kv_gather_ops) {
197 + direct_dcp_kv_gather_ops.def(
198 + "direct_dcp_kv_gather("
199 + "Tensor local_kv, Tensor! received_kv, Tensor! received_signal, "
200 + "Tensor! completion, "
201 + "Tensor! epoch, Tensor! gathered_kv, "
202 + "int world_size, int rank, int max_gathered_tokens, "
203 + "int kv_mc_ptr, int signal_mc_ptr) -> ()");
204 +}
205 +
206 +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, direct_dcp_kv_gather_ops) {
207 + direct_dcp_kv_gather_ops.impl("direct_dcp_kv_gather",
208 + TORCH_BOX(&direct_dcp_kv_gather));
209 +}
csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu new
+197
@@ -0,0 +1,197 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +// Direct symmetric-memory DCP query gather into the consumer-final buffer.
4 +
5 +#include <torch/csrc/stable/library.h>
6 +#include <torch/headeronly/core/ScalarType.h>
7 +
8 +#include <cstdio>
9 +
10 +#include "dcp_direct_common.cuh"
11 +
12 +namespace {
13 +
14 +using vllm::direct_dcp::check_cuda_launch;
15 +using vllm::direct_dcp::increment_epoch_kernel;
16 +using vllm::direct_dcp::multimem_store_16;
17 +using vllm::direct_dcp::multimem_store_release_system;
18 +using vllm::direct_dcp::wait_for_epoch;
19 +
20 +// Multicast each rank's head slice directly into every consumer's final query
21 +// buffer. Reuse is ordered by the downstream DCP output synchronization.
22 +__global__ void direct_dcp_q_gather_multimem_kernel(
23 + const uint4* local_query, uint4* mc_final_query, uint32_t* mc_signal,
24 + const uint32_t* received_signal, int64_t* epoch_ptr, uint32_t* completion,
25 + int64_t world_size, int64_t rank, int64_t num_tokens,
26 + int64_t bytes_per_token, int64_t query_token_stride_bytes,
27 + int64_t destination_token_stride_bytes) {
28 + // The common one-token decode case uses one block. Fold the epoch update
29 + // into that publication kernel to avoid an otherwise separate tiny launch.
30 + if (gridDim.x == 1 && threadIdx.x == 0) {
31 + epoch_ptr[0] += 1;
32 + }
33 + __syncthreads();
34 + uint32_t epoch = static_cast<uint32_t>(epoch_ptr[0]);
35 + int64_t signal_slot = static_cast<int64_t>(epoch & 1u);
36 +
37 + int64_t items_per_token = bytes_per_token / sizeof(uint4);
38 + int64_t source_token_stride = query_token_stride_bytes / sizeof(uint4);
39 + int64_t destination_token_stride =
40 + destination_token_stride_bytes / sizeof(uint4);
41 + int64_t destination_head_offset = rank * items_per_token;
42 + for (int64_t token_idx = blockIdx.x; token_idx < num_tokens;
43 + token_idx += gridDim.x) {
44 + for (int64_t token_item = threadIdx.x; token_item < items_per_token;
45 + token_item += blockDim.x) {
46 + multimem_store_16(
47 + mc_final_query + token_idx * destination_token_stride +
48 + destination_head_offset + token_item,
49 + local_query[token_idx * source_token_stride + token_item]);
50 + }
51 + }
52 +
53 + // Publish all multicast writes before incrementing completion.
54 + __threadfence_system();
55 + __syncthreads();
56 + if (threadIdx.x != 0) {
57 + return;
58 + }
59 +
60 + if (gridDim.x > 1) {
61 + uint32_t completed = atomicAdd(completion, 1u);
62 + if (completed + 1u != gridDim.x) {
63 + return;
64 + }
65 + atomicExch(completion, 0u);
66 + }
67 +
68 + multimem_store_release_system(mc_signal + signal_slot * world_size + rank,
69 + epoch);
70 +
71 + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) {
72 + int64_t signal_item = signal_slot * world_size + source_rank;
73 + if (!wait_for_epoch(received_signal + signal_item, epoch)) {
74 + printf("direct DCP q-gather multimem timeout source=%lld epoch=%u\n",
75 + static_cast<long long>(source_rank), epoch);
76 + asm volatile("trap;");
77 + }
78 + }
79 +}
80 +
81 +void direct_dcp_q_gather(const torch::stable::Tensor& local_query,
82 + torch::stable::Tensor& final_query,
83 + torch::stable::Tensor& received_signal,
84 + torch::stable::Tensor& completion,
85 + torch::stable::Tensor& epoch, int64_t world_size,
86 + int64_t rank, int64_t max_num_tokens,
87 + int64_t padded_num_heads, int64_t query_mc_ptr,
88 + int64_t signal_mc_ptr) {
89 + using torch::headeronly::ScalarType;
90 +
91 + STD_TORCH_CHECK(local_query.is_cuda(), "local query must be a CUDA tensor");
92 + ScalarType dtype = local_query.scalar_type();
93 + STD_TORCH_CHECK(local_query.dim() == 3,
94 + "local query must have shape [T,H,D]");
95 + STD_TORCH_CHECK(world_size > 1, "world_size must be greater than 1");
96 + STD_TORCH_CHECK(rank >= 0 && rank < world_size, "invalid rank");
97 +
98 + int64_t num_tokens = local_query.size(0);
99 + int64_t heads_per_rank = local_query.size(1);
100 + int64_t head_dim = local_query.size(2);
101 + int64_t gathered_num_heads = world_size * heads_per_rank;
102 + int64_t element_size = local_query.element_size();
103 + STD_TORCH_CHECK(num_tokens > 0 && num_tokens <= max_num_tokens,
104 + "token count exceeds symmetric q-gather buffer capacity");
105 + STD_TORCH_CHECK(heads_per_rank > 0 && head_dim > 0,
106 + "query head dimensions must be positive");
107 + STD_TORCH_CHECK(padded_num_heads >= gathered_num_heads,
108 + "padded query heads must cover all gathered heads");
109 + STD_TORCH_CHECK(local_query.stride(2) == 1 &&
110 + local_query.stride(1) == head_dim &&
111 + local_query.stride(0) >= heads_per_rank * head_dim,
112 + "local query must have packed heads");
113 +
114 + STD_TORCH_CHECK(
115 + final_query.is_cuda() && final_query.scalar_type() == dtype &&
116 + final_query.is_contiguous() && final_query.dim() == 3 &&
117 + final_query.size(0) == num_tokens &&
118 + final_query.size(1) == gathered_num_heads &&
119 + final_query.size(2) == head_dim,
120 + "final query must be contiguous with shape [T,world_size*H,D]");
121 + STD_TORCH_CHECK(
122 + received_signal.is_cuda() && received_signal.is_contiguous() &&
123 + received_signal.scalar_type() == ScalarType::Int &&
124 + received_signal.dim() == 2 && received_signal.size(0) == 2 &&
125 + received_signal.size(1) == world_size,
126 + "received signal has the wrong symmetric buffer layout");
127 + STD_TORCH_CHECK(completion.is_cuda() && completion.is_contiguous() &&
128 + completion.scalar_type() == ScalarType::Int &&
129 + completion.numel() == 1,
130 + "completion counter must be one CUDA int32 tensor");
131 + STD_TORCH_CHECK(epoch.is_cuda() && epoch.is_contiguous() &&
132 + epoch.scalar_type() == ScalarType::Long &&
133 + epoch.numel() == 1,
134 + "epoch must be a one-element CUDA int64 tensor");
135 + int64_t device_index = local_query.get_device_index();
136 + STD_TORCH_CHECK(
137 + final_query.get_device_index() == device_index &&
138 + received_signal.get_device_index() == device_index &&
139 + completion.get_device_index() == device_index &&
140 + epoch.get_device_index() == device_index,
141 + "direct DCP q-gather tensors must be on the same CUDA device");
142 +
143 + int64_t query_token_stride_bytes = local_query.stride(0) * element_size;
144 + int64_t bytes_per_token = heads_per_rank * head_dim * element_size;
145 + int64_t gathered_token_stride_bytes =
146 + gathered_num_heads * head_dim * element_size;
147 + bool vectorized =
148 + reinterpret_cast<uintptr_t>(local_query.data_ptr()) % alignof(uint4) ==
149 + 0 &&
150 + reinterpret_cast<uintptr_t>(final_query.data_ptr()) % alignof(uint4) ==
151 + 0 &&
152 + query_token_stride_bytes % sizeof(uint4) == 0 &&
153 + bytes_per_token % sizeof(uint4) == 0;
154 + STD_TORCH_CHECK(
155 + vectorized,
156 + "direct DCP q-gather requires 16-byte-aligned pointers and strides");
157 + STD_TORCH_CHECK(query_mc_ptr != 0 && signal_mc_ptr != 0,
158 + "direct DCP q-gather requires multicast pointers");
159 +
160 + const torch::stable::accelerator::DeviceGuard device_guard(device_index);
161 + cudaStream_t stream = get_current_cuda_stream();
162 + constexpr int kThreads = 256;
163 + int64_t blocks = num_tokens < world_size ? num_tokens : world_size;
164 + if (blocks > 1) {
165 + increment_epoch_kernel<<<1, 1, 0, stream>>>(
166 + epoch.mutable_data_ptr<int64_t>());
167 + check_cuda_launch("direct DCP q-gather epoch");
168 + }
169 +
170 + direct_dcp_q_gather_multimem_kernel<<<blocks, kThreads, 0, stream>>>(
171 + reinterpret_cast<const uint4*>(local_query.data_ptr()),
172 + reinterpret_cast<uint4*>(static_cast<uintptr_t>(query_mc_ptr)),
173 + reinterpret_cast<uint32_t*>(static_cast<uintptr_t>(signal_mc_ptr)),
174 + reinterpret_cast<const uint32_t*>(
175 + received_signal.const_data_ptr<int32_t>()),
176 + epoch.mutable_data_ptr<int64_t>(),
177 + reinterpret_cast<uint32_t*>(completion.mutable_data_ptr<int32_t>()),
178 + world_size, rank, num_tokens, bytes_per_token, query_token_stride_bytes,
179 + gathered_token_stride_bytes);
180 + check_cuda_launch("direct DCP q-gather");
181 +}
182 +
183 +} // namespace
184 +
185 +STABLE_TORCH_LIBRARY_FRAGMENT(_C, direct_dcp_q_gather_ops) {
186 + direct_dcp_q_gather_ops.def(
187 + "direct_dcp_q_gather("
188 + "Tensor local_query, Tensor! final_query, Tensor! received_signal, "
189 + "Tensor! completion, Tensor! epoch, "
190 + "int world_size, int rank, int max_num_tokens, "
191 + "int padded_num_heads, int query_mc_ptr, int signal_mc_ptr) -> ()");
192 +}
193 +
194 +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, direct_dcp_q_gather_ops) {
195 + direct_dcp_q_gather_ops.impl("direct_dcp_q_gather",
196 + TORCH_BOX(&direct_dcp_q_gather));
197 +}
tests/distributed/test_dcp_a2a.py
+106 -2
@@ -222,6 +222,49 @@ class TestLSEWeightedCombine:
222 assert result.shape == (B, H, D)
223 torch.testing.assert_close(result, outputs[1], atol=1e-5, rtol=1e-5)
224
225 + def test_empty_shard_ignores_undefined_output(self):
226 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
227 +
228 + outputs = torch.tensor([[[[float("nan")]]], [[[3.0]]]])
229 + lses = torch.tensor([[[-float("inf")]], [[0.0]]])
230 +
231 + result = _lse_weighted_combine(outputs, lses)
232 +
233 + torch.testing.assert_close(result, outputs[1])
234 +
235 + def test_ag_rs_masks_empty_shard_and_padded_lse(self, monkeypatch):
236 + import vllm.v1.attention.ops.common as common
237 +
238 + class FakeGroup:
239 + world_size = 2
240 + rank_in_group = 0
241 +
242 + def all_gather(self, tensor, dim):
243 + assert dim == 0
244 + return torch.cat((tensor, tensor), dim=dim)
245 +
246 + monkeypatch.setattr(
247 + common,
248 + "correct_attn_out",
249 + lambda output, lses, *args, **kwargs: (output, lses[0]),
250 + )
251 + output = torch.ones(7, 1, 1)
252 + lse = torch.ones(7, 1)
253 + seq_lens = torch.tensor([0, 2], dtype=torch.int32)
254 + query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32)
255 +
256 + _, masked_lse = common._cp_lse_common(
257 + output,
258 + lse,
259 + FakeGroup(),
260 + seq_lens=seq_lens,
261 + query_start_loc=query_start_loc,
262 + )
263 +
264 + assert torch.isneginf(masked_lse[:1]).all()
265 + torch.testing.assert_close(masked_lse[1:5], torch.ones_like(masked_lse[1:5]))
266 + assert torch.isneginf(masked_lse[5:]).all()
267 +
268 def test_mathematically_correct(self):
269 """Verify mathematical correctness of LSE combination."""
270 from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
@@ -344,7 +387,7 @@ class TestPackedA2AKernels:
387 world_size, B, h_per_rank, D = 4, 7, 2, 32
388 H = world_size * h_per_rank
389 cp_attn_out = torch.randn(B, H, D, device=device, dtype=dtype)
347 - cp_attn_lse = torch.randn(B, H, device=device, dtype=torch.float32)
390 + cp_attn_lse = torch.randn(B, H, device=device, dtype=dtype)
391 lse_pack_dim = _dcp_a2a_lse_pack_dim(dtype)
392 send_buffer = torch.empty(
393 (world_size, B, h_per_rank, D + lse_pack_dim),
@@ -375,6 +418,67 @@ class TestPackedA2AKernels:
418 else:
419 _assert_packed_a2a_close(actual, expected_out, dtype)
420
421 + @pytest.mark.skipif(
422 + torch.accelerator.device_count() < 1, reason="CUDA is required."
423 + )
424 + def test_empty_seq_lens_ignore_undefined_output(self):
425 + from vllm.v1.attention.ops.dcp_alltoall import (
426 + _dcp_a2a_lse_pack_dim,
427 + _dcp_a2a_pack_send,
428 + _dcp_a2a_unpack_combine,
429 + )
430 +
431 + device = torch.device("cuda")
432 + world_size, num_tokens, h_per_rank, head_dim = 2, 5, 1, 32
433 + num_heads = world_size * h_per_rank
434 + output = torch.randn(
435 + num_tokens,
436 + num_heads,
437 + head_dim,
438 + device=device,
439 + dtype=torch.bfloat16,
440 + )
441 + output[:1] = float("nan")
442 + lse = torch.randn(num_tokens, num_heads, device=device, dtype=output.dtype)
443 + seq_lens = torch.tensor([0, 2], device=device, dtype=torch.int32)
444 + query_start_loc = torch.tensor([0, 1, 5], device=device, dtype=torch.int32)
445 + lse_pack_dim = _dcp_a2a_lse_pack_dim(output.dtype)
446 + send_buffer = torch.empty(
447 + (
448 + world_size,
449 + num_tokens,
450 + h_per_rank,
451 + head_dim + lse_pack_dim,
452 + ),
453 + device=device,
454 + dtype=output.dtype,
455 + )
456 +
457 + _dcp_a2a_pack_send(
458 + output,
459 + lse,
460 + send_buffer,
461 + world_size,
462 + h_per_rank,
463 + head_dim,
464 + lse_pack_dim,
465 + seq_lens=seq_lens,
466 + query_start_loc=query_start_loc,
467 + )
468 + actual_output, actual_lse = _dcp_a2a_unpack_combine(
469 + send_buffer,
470 + head_dim,
471 + lse_pack_dim,
472 + return_lse=True,
473 + is_lse_base_on_e=True,
474 + )
475 +
476 + torch.testing.assert_close(
477 + actual_output[:1], torch.zeros_like(actual_output[:1])
478 + )
479 + assert torch.isneginf(actual_lse[:1]).all()
480 + assert torch.isfinite(actual_output[1:]).all()
481 +
482
483 def _distributed_packed_a2a_worker(env: dict[str, str]) -> None:
484 update_environment_variables(env)
@@ -417,7 +521,7 @@ def _distributed_packed_a2a_worker(env: dict[str, str]) -> None:
521 B,
522 H,
523 device=f"cuda:{local_rank}",
420 - dtype=torch.float32,
524 + dtype=dtype,
525 generator=generator,
526 )
527 actual = dcp_a2a_lse_reduce(
tests/distributed/test_dcp_direct_a2a_lse_reduce.py new
+1077
@@ -0,0 +1,1077 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for direct symmetric-memory DCP collectives."""
4 +
5 +import functools
6 +from unittest.mock import MagicMock
7 +
8 +import multiprocess as mp
9 +import pytest
10 +import torch
11 +import torch.distributed as dist
12 +
13 +import vllm.v1.attention.ops.dcp_utils as dcp_utils
14 +from vllm.utils.network_utils import get_open_port
15 +from vllm.utils.system_utils import update_environment_variables
16 +
17 +mp.set_start_method("spawn", force=True)
18 +
19 +
20 +def _has_multicast_support() -> bool:
21 + if not torch.cuda.is_available():
22 + return False
23 + try:
24 + from torch._C._autograd import DeviceType
25 + from torch._C._distributed_c10d import _SymmetricMemory
26 +
27 + return _SymmetricMemory.has_multicast_support(DeviceType.CUDA, 0)
28 + except Exception:
29 + return False
30 +
31 +
32 +def _dtype_from_name(dtype_name: str) -> torch.dtype:
33 + return {
34 + "float16": torch.float16,
35 + "bfloat16": torch.bfloat16,
36 + "float8_e4m3fn": torch.float8_e4m3fn,
37 + "float32": torch.float32,
38 + }[dtype_name]
39 +
40 +
41 +def _assert_close(
42 + actual: torch.Tensor,
43 + expected: torch.Tensor,
44 + dtype: torch.dtype,
45 +) -> None:
46 + torch.testing.assert_close(actual.float(), expected.float(), rtol=3e-2, atol=3e-2)
47 +
48 +
49 +def _q_gather_reference(
50 + local_query: torch.Tensor,
51 + world_size: int,
52 + padded_num_heads: int | None,
53 +) -> torch.Tensor:
54 + num_tokens, heads_per_rank, head_dim = local_query.shape
55 + gathered = torch.empty(
56 + (world_size * num_tokens, heads_per_rank, head_dim),
57 + dtype=local_query.dtype,
58 + device=local_query.device,
59 + )
60 + dist.all_gather_into_tensor(gathered, local_query.contiguous())
61 + expected = (
62 + gathered.view(world_size, num_tokens, heads_per_rank, head_dim)
63 + .movedim(0, 1)
64 + .reshape(num_tokens, world_size * heads_per_rank, head_dim)
65 + )
66 + if padded_num_heads is not None:
67 + reserved = expected.new_empty((num_tokens, padded_num_heads, head_dim))
68 + reserved.resize_(expected.shape)
69 + reserved.copy_(expected)
70 + expected = reserved
71 + return expected
72 +
73 +
74 +def _assert_q_gather_matches_reference(
75 + actual: torch.Tensor,
76 + local_query: torch.Tensor,
77 + world_size: int,
78 + padded_num_heads: int | None,
79 +) -> None:
80 + expected = _q_gather_reference(local_query, world_size, padded_num_heads)
81 + assert actual.shape == expected.shape
82 + assert actual.dtype == expected.dtype
83 + assert actual.is_contiguous()
84 + assert actual.stride() == expected.stride()
85 + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
86 +
87 + storage_num_heads = padded_num_heads or world_size * local_query.shape[1]
88 + remaining_storage_bytes = (
89 + actual.untyped_storage().nbytes()
90 + - actual.storage_offset() * actual.element_size()
91 + )
92 + required_storage_bytes = (
93 + local_query.shape[0]
94 + * storage_num_heads
95 + * local_query.shape[2]
96 + * actual.element_size()
97 + )
98 + assert remaining_storage_bytes >= required_storage_bytes
99 +
100 +
101 +def _distributed_run(fn, world_size: int, extra_env: dict[str, str]) -> None:
102 + port = str(get_open_port())
103 + processes: list[mp.Process] = []
104 + for rank in range(world_size):
105 + env = {
106 + "RANK": str(rank),
107 + "LOCAL_RANK": str(rank),
108 + "WORLD_SIZE": str(world_size),
109 + "LOCAL_WORLD_SIZE": str(world_size),
110 + "MASTER_ADDR": "localhost",
111 + "MASTER_PORT": port,
112 + **extra_env,
113 + }
114 + process = mp.Process(target=fn, args=(env,))
115 + processes.append(process)
116 + process.start()
117 +
118 + for process in processes:
119 + process.join(timeout=120)
120 +
121 + for process in processes:
122 + if process.is_alive():
123 + process.kill()
124 + process.join()
125 + assert process.exitcode == 0
126 +
127 +
128 +class _FakeGroupCoordinator:
129 + device_group = None
130 + cpu_group = None
131 + world_size = 4
132 +
133 +
134 +class _FakeProcessGroup:
135 + def size(self) -> int:
136 + return 4
137 +
138 + def rank(self) -> int:
139 + return 0
140 +
141 +
142 +class TestDirectDCPGating:
143 + def test_env_disabled_returns_none(self, monkeypatch):
144 + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0")
145 + dcp_utils.get_direct_dcp_a2a_workspace.cache_clear()
146 + workspace = dcp_utils.get_direct_dcp_a2a_workspace(
147 + _FakeGroupCoordinator(), torch.device("cpu"), 16, 2, 32, torch.bfloat16, 1
148 + )
149 + assert workspace is None
150 +
151 + def test_forced_with_unsupported_dtype_raises(self, monkeypatch):
152 + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1")
153 + dcp_utils.get_direct_dcp_a2a_workspace.cache_clear()
154 + with pytest.raises(ValueError, match="does not support"):
155 + dcp_utils.get_direct_dcp_a2a_workspace(
156 + _FakeGroupCoordinator(),
157 + torch.device("cpu"),
158 + 16,
159 + 2,
160 + 32,
161 + torch.float32,
162 + 1,
163 + )
164 +
165 + def test_zero_ubatches_raises(self):
166 + with pytest.raises(ValueError, match="ubatch"):
167 + dcp_utils.DirectDCPA2AWorkspace(
168 + None, torch.device("cpu"), 16, 2, 32, torch.bfloat16, num_ubatches=0
169 + )
170 +
171 + def test_auto_with_unsupported_dtype_returns_none(self, monkeypatch):
172 + monkeypatch.delenv("VLLM_USE_DIRECT_DCP_A2A", raising=False)
173 + dcp_utils.get_direct_dcp_a2a_workspace.cache_clear()
174 + workspace = dcp_utils.get_direct_dcp_a2a_workspace(
175 + _FakeGroupCoordinator(), torch.device("cpu"), 16, 2, 32, torch.float32, 1
176 + )
177 + assert workspace is None
178 +
179 + def test_q_gather_env_disabled_returns_none(self, monkeypatch):
180 + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_Q_GATHER", "0")
181 + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1")
182 + dcp_utils.get_direct_dcp_q_gather_workspace.cache_clear()
183 + workspace = dcp_utils.get_direct_dcp_q_gather_workspace(
184 + _FakeGroupCoordinator(),
185 + torch.device("cpu"),
186 + 16,
187 + 2,
188 + 32,
189 + torch.bfloat16,
190 + 1,
191 + )
192 + assert workspace is None
193 +
194 + def test_q_gather_flag_is_independent(self, monkeypatch):
195 + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_Q_GATHER", "1")
196 + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0")
197 + monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: True)
198 + dcp_utils.get_direct_dcp_q_gather_workspace.cache_clear()
199 + workspace = object()
200 + init_workspace = MagicMock(return_value=workspace)
201 + monkeypatch.setattr(
202 + dcp_utils,
203 + "DirectDCPQGatherWorkspace",
204 + init_workspace,
205 + )
206 +
207 + result = dcp_utils.get_direct_dcp_q_gather_workspace(
208 + _FakeGroupCoordinator(),
209 + torch.device("cpu"),
210 + 16,
211 + 2,
212 + 32,
213 + torch.float32,
214 + 1,
215 + )
216 +
217 + assert result is workspace
218 + assert init_workspace.call_args.args[5] == torch.float32
219 +
220 + def test_kv_gather_env_disabled_returns_none(self, monkeypatch):
221 + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_KV_GATHER", "0")
222 + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1")
223 + dcp_utils.get_direct_dcp_kv_gather_workspace.cache_clear()
224 + workspace = dcp_utils.get_direct_dcp_kv_gather_workspace(
225 + _FakeGroupCoordinator(), torch.device("cpu"), 64, 576, torch.bfloat16, 1
226 + )
227 + assert workspace is None
228 +
229 + def test_kv_gather_flag_is_independent(self, monkeypatch):
230 + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_KV_GATHER", "1")
231 + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0")
232 + monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: True)
233 + dcp_utils.get_direct_dcp_kv_gather_workspace.cache_clear()
234 + workspace = object()
235 + init_workspace = MagicMock(return_value=workspace)
236 + monkeypatch.setattr(
237 + dcp_utils,
238 + "DirectDCPKVGatherWorkspace",
239 + init_workspace,
240 + )
241 +
242 + result = dcp_utils.get_direct_dcp_kv_gather_workspace(
243 + _FakeGroupCoordinator(), torch.device("cpu"), 64, 576, torch.bfloat16, 1
244 + )
245 +
246 + assert result is workspace
247 +
248 + @pytest.mark.parametrize(
249 + ("flag_name", "factory_name", "factory_args"),
250 + [
251 + (
252 + "VLLM_USE_DIRECT_DCP_Q_GATHER",
253 + "get_direct_dcp_q_gather_workspace",
254 + (16, 2, 32, torch.bfloat16, 1),
255 + ),
256 + (
257 + "VLLM_USE_DIRECT_DCP_KV_GATHER",
258 + "get_direct_dcp_kv_gather_workspace",
259 + (64, 576, torch.bfloat16, 1),
260 + ),
261 + ],
262 + )
263 + def test_gather_requires_multicast(
264 + self,
265 + monkeypatch,
266 + flag_name,
267 + factory_name,
268 + factory_args,
269 + ):
270 + factory = getattr(dcp_utils, factory_name)
271 + monkeypatch.setenv(flag_name, "1")
272 + monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: False)
273 + factory.cache_clear()
274 +
275 + assert (
276 + factory(
277 + _FakeGroupCoordinator(),
278 + torch.device("cpu"),
279 + *factory_args,
280 + )
281 + is None
282 + )
283 +
284 + def test_kv_gather_rejects_invalid_workspace_geometry(self):
285 + with pytest.raises(ValueError, match="ubatch"):
286 + dcp_utils.DirectDCPKVGatherWorkspace(
287 + None, torch.device("cpu"), 64, 576, num_ubatches=0
288 + )
289 + with pytest.raises(ValueError, match="divide evenly"):
290 + dcp_utils.DirectDCPKVGatherWorkspace(
291 + _FakeProcessGroup(), torch.device("cpu"), 63, 576
292 + )
293 + with pytest.raises(ValueError, match="16-byte"):
294 + dcp_utils.DirectDCPKVGatherWorkspace(
295 + _FakeProcessGroup(), torch.device("cpu"), 64, 3
296 + )
297 +
298 + def test_q_gather_rejects_invalid_workspace_geometry(self):
299 + with pytest.raises(ValueError, match="ubatch"):
300 + dcp_utils.DirectDCPQGatherWorkspace(
301 + None, torch.device("cpu"), 16, 2, 32, num_ubatches=0
302 + )
303 + with pytest.raises(ValueError, match="padded heads"):
304 + dcp_utils.DirectDCPQGatherWorkspace(
305 + _FakeProcessGroup(),
306 + torch.device("cpu"),
307 + 16,
308 + 2,
309 + 32,
310 + padded_num_heads=7,
311 + )
312 + with pytest.raises(ValueError, match="16-byte"):
313 + dcp_utils.DirectDCPQGatherWorkspace(
314 + _FakeProcessGroup(),
315 + torch.device("cpu"),
316 + 16,
317 + 1,
318 + 3,
319 + )
320 +
321 +
322 +def _manager_config(dcp_comm_backend: str = "a2a"):
323 + config = MagicMock()
324 + config.parallel_config.num_ubatches = 1
325 + config.parallel_config.dcp_comm_backend = dcp_comm_backend
326 + config.scheduler_config.max_num_batched_tokens = 16
327 + config.scheduler_config.max_num_seqs = 4
328 + config.num_speculative_tokens = 0
329 + config.speculative_config = None
330 + config.compilation_config.max_cudagraph_capture_size = 0
331 + return config
332 +
333 +
334 +def test_mla_dcp_manager_selects_direct_backends(monkeypatch):
335 + import vllm.v1.attention.ops.dcp_utils as dcp_manager
336 +
337 + group = MagicMock(world_size=2)
338 + monkeypatch.setattr(dcp_manager, "get_dcp_group", lambda: group)
339 + direct_a2a = MagicMock()
340 + direct_query = MagicMock()
341 + direct_kv = MagicMock()
342 + monkeypatch.setattr(
343 + dcp_manager, "get_direct_dcp_a2a_workspace", MagicMock(return_value=direct_a2a)
344 + )
345 + monkeypatch.setattr(
346 + dcp_manager,
347 + "get_direct_dcp_q_gather_workspace",
348 + MagicMock(return_value=direct_query),
349 + )
350 + monkeypatch.setattr(
351 + dcp_manager,
352 + "get_direct_dcp_kv_gather_workspace",
353 + MagicMock(return_value=direct_kv),
354 + )
355 +
356 + manager = dcp_manager.MLADCPManager(
357 + vllm_config=_manager_config(),
358 + device=torch.device("cpu"),
359 + num_heads=2,
360 + query_head_dim=8,
361 + output_head_dim=4,
362 + query_dtype=torch.bfloat16,
363 + output_dtype=torch.bfloat16,
364 + padded_num_heads=None,
365 + is_lse_base_on_e=False,
366 + use_pcp=False,
367 + )
368 + workspace = torch.empty(96, 8)
369 +
370 + assert manager.query_gather == direct_query.gather
371 + manager.init_kv_gather(workspace, 64)
372 + gathered_kv, local_kv = torch.empty(4, 8), torch.empty(2, 8)
373 + manager.kv_gather(gathered_kv, local_kv)
374 + direct_kv.gather.assert_called_once_with(gathered_kv, local_kv)
375 + output, lse = torch.empty(1), torch.empty(1)
376 + seq_lens = torch.ones(1, dtype=torch.int32)
377 + query_start_loc = torch.tensor([0, 1], dtype=torch.int32)
378 + manager.combine(
379 + output,
380 + lse,
381 + seq_lens=seq_lens,
382 + query_start_loc=query_start_loc,
383 + )
384 + direct_a2a.lse_reduce.assert_called_once_with(
385 + output,
386 + lse,
387 + seq_lens=seq_lens,
388 + query_start_loc=query_start_loc,
389 + is_lse_base_on_e=False,
390 + )
391 +
392 +
393 +def test_mla_dcp_manager_selects_fallback_backends(monkeypatch):
394 + import vllm.v1.attention.ops.dcp_utils as dcp_manager
395 +
396 + group = MagicMock(world_size=2)
397 + gathered_query = torch.empty(1, 4, 8)
398 + group.all_gather.return_value = gathered_query
399 + monkeypatch.setattr(dcp_manager, "get_dcp_group", lambda: group)
400 + monkeypatch.setattr(
401 + dcp_manager, "get_direct_dcp_a2a_workspace", MagicMock(return_value=None)
402 + )
403 + monkeypatch.setattr(
404 + dcp_manager, "get_direct_dcp_q_gather_workspace", MagicMock(return_value=None)
405 + )
406 + monkeypatch.setattr(
407 + dcp_manager, "get_direct_dcp_kv_gather_workspace", MagicMock(return_value=None)
408 + )
409 + fallback_combine = MagicMock(return_value=torch.empty(1))
410 + monkeypatch.setattr(dcp_manager, "dcp_a2a_lse_reduce", fallback_combine)
411 +
412 + manager = dcp_manager.MLADCPManager(
413 + vllm_config=_manager_config(),
414 + device=torch.device("cpu"),
415 + num_heads=2,
416 + query_head_dim=8,
417 + output_head_dim=4,
418 + query_dtype=torch.bfloat16,
419 + output_dtype=torch.bfloat16,
420 + padded_num_heads=None,
421 + is_lse_base_on_e=True,
422 + use_pcp=False,
423 + )
424 +
425 + all_gather = MagicMock()
426 + monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", all_gather)
427 + workspace = torch.empty(96, 8)
428 + manager.init_kv_gather(workspace, 64)
429 + output, local = torch.empty(4, 8), torch.empty(2, 8)
430 + manager.kv_gather(output, local)
431 + all_gather.assert_called_once_with(output, local, group=group.device_group)
432 +
433 + query = torch.empty(1, 2, 8)
434 + assert manager.query_gather is not None
435 + assert manager.query_gather(query) is gathered_query
436 + group.all_gather.assert_called_once_with(query, dim=1)
437 +
438 + partial_output, partial_lse = torch.empty(1), torch.empty(1)
439 + seq_lens = torch.ones(1, dtype=torch.int32)
440 + query_start_loc = torch.tensor([0, 1], dtype=torch.int32)
441 + manager.combine(
442 + partial_output,
443 + partial_lse,
444 + seq_lens=seq_lens,
445 + query_start_loc=query_start_loc,
446 + )
447 + fallback_combine.assert_called_once_with(
448 + partial_output,
449 + partial_lse,
450 + seq_lens=seq_lens,
451 + query_start_loc=query_start_loc,
452 + cp_group=group,
453 + is_lse_base_on_e=True,
454 + )
455 +
456 +
457 +def test_dcp_workspace_covers_parallel_drafting():
458 + config = _manager_config()
459 + config.scheduler_config.max_num_batched_tokens = 128
460 + config.num_speculative_tokens = 3
461 + config.speculative_config = MagicMock(parallel_drafting=True)
462 +
463 + assert dcp_utils.get_dcp_workspace_max_num_tokens(config) == 28
464 +
465 +
466 +def test_mla_dcp_manager_selects_pcp_combine(monkeypatch):
467 + import vllm.v1.attention.ops.dcp_utils as dcp_manager
468 +
469 + monkeypatch.setattr(dcp_manager, "get_dcp_group", lambda: MagicMock(world_size=2))
470 + manager = dcp_manager.MLADCPManager(
471 + vllm_config=_manager_config(dcp_comm_backend="ag_rs"),
472 + device=torch.device("cpu"),
473 + num_heads=2,
474 + query_head_dim=8,
475 + output_head_dim=4,
476 + query_dtype=torch.bfloat16,
477 + output_dtype=torch.bfloat16,
478 + padded_num_heads=None,
479 + is_lse_base_on_e=True,
480 + use_pcp=True,
481 + )
482 +
483 + assert isinstance(manager.combine, functools.partial)
484 + assert manager.combine.func is dcp_manager.cp_lse_ag_out_ar
485 + assert manager.query_gather is None
486 +
487 +
488 +def test_dcp_chunk_workspace_alignment_covers_interleave():
489 + from vllm.model_executor.layers.attention.mla_attention import (
490 + align_mla_chunked_context_workspace_size,
491 + )
492 +
493 + config = MagicMock()
494 + config.cache_config.block_size = 32
495 + config.parallel_config.decode_context_parallel_size = 8
496 + config.parallel_config.cp_kv_cache_interleave_size = 8
497 + config.scheduler_config.max_num_seqs = 3
498 +
499 + assert align_mla_chunked_context_workspace_size(config, 100) == 192
500 +
501 +
502 +def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch):
503 + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla
504 +
505 + monkeypatch.setattr(
506 + sparse_mla.AttentionMetadataBuilder,
507 + "__init__",
508 + lambda *args, **kwargs: None,
509 + )
510 + monkeypatch.setattr(
511 + sparse_mla,
512 + "get_dcp_group",
513 + lambda: MagicMock(world_size=2),
514 + )
515 + monkeypatch.setattr(
516 + sparse_mla,
517 + "get_mla_dims",
518 + lambda _: MagicMock(kv_lora_rank=8, qk_rope_head_dim=4),
519 + )
520 +
521 + manager = object.__new__(dcp_utils.MLADCPManager)
522 + manager.init_kv_gather = MagicMock()
523 + layer = MagicMock(dcp_manager=manager)
524 + config = MagicMock()
525 + config.model_config.dtype = torch.bfloat16
526 + config.model_config.max_model_len = 64
527 + config.model_config.hf_config.index_topk = 8
528 + config.scheduler_config.max_num_batched_tokens = 64
529 + config.scheduler_config.max_num_seqs = 2
530 + config.cache_config.block_size = 4
531 + config.parallel_config.prefill_context_parallel_size = 1
532 + config.parallel_config.decode_context_parallel_size = 2
533 + config.parallel_config.cp_kv_cache_interleave_size = 1
534 + config.compilation_config.static_forward_context = {"layer": layer}
535 +
536 + builder = sparse_mla.SparseMLACommonMetadataBuilder(
537 + MagicMock(),
538 + ["layer"],
539 + config,
540 + torch.device("cpu"),
541 + )
542 +
543 + assert builder.dcp_manager is manager
544 + manager.init_kv_gather.assert_called_once_with(
545 + builder.chunked_prefill_workspace,
546 + builder.chunked_prefill_workspace_size,
547 + )
548 +
549 +
550 +def test_sparse_mla_workspace_preserves_non_dcp_size():
551 + from vllm.model_executor.layers.attention.sparse_mla_attention import (
552 + SparseMLACommonMetadataBuilder,
553 + )
554 +
555 + config = MagicMock()
556 + config.model_config.max_model_len = 1
557 + config.model_config.hf_config.index_topk = 7
558 + config.scheduler_config.max_num_seqs = 3
559 + config.cache_config.block_size = 4
560 + config.parallel_config.decode_context_parallel_size = 1
561 +
562 + assert (
563 + SparseMLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(config)
564 + == 21
565 + )
566 +
567 +
568 +def _distributed_direct_q_gather_worker(env: dict[str, str]) -> None:
569 + update_environment_variables(env)
570 + local_rank = int(env["LOCAL_RANK"])
571 + device = torch.device(f"cuda:{local_rank}")
572 + torch.accelerator.set_device_index(local_rank)
573 + dist.init_process_group(backend="nccl")
574 + try:
575 + rank = dist.get_rank()
576 + world_size = dist.get_world_size()
577 + heads_per_rank, head_dim, max_num_tokens = 6, 576, 128
578 + padded_num_heads = 128 if world_size == 4 else None
579 + active_ubatch = [0]
580 + dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0]
581 + for dtype_idx, dtype_name in enumerate(
582 + ("bfloat16", "float8_e4m3fn", "float32")
583 + ):
584 + dtype = _dtype_from_name(dtype_name)
585 + workspace = dcp_utils.DirectDCPQGatherWorkspace(
586 + dist.group.WORLD,
587 + device,
588 + max_num_tokens,
589 + heads_per_rank,
590 + head_dim,
591 + dtype,
592 + num_ubatches=2,
593 + padded_num_heads=padded_num_heads,
594 + )
595 +
596 + cases = (
597 + ((1, False), (128, True), (17, True), (5, False))
598 + if dtype == torch.bfloat16
599 + else ((17, True),)
600 + )
601 + for iteration, (num_tokens, noncontiguous) in enumerate(cases):
602 + generator = torch.Generator(device=device)
603 + generator.manual_seed(9000 + rank * 101 + dtype_idx * 1009 + iteration)
604 + source_num_heads = (
605 + heads_per_rank + 2 if noncontiguous else heads_per_rank
606 + )
607 + query_storage = torch.randn(
608 + num_tokens,
609 + source_num_heads,
610 + head_dim,
611 + device=device,
612 + dtype=torch.bfloat16,
613 + generator=generator,
614 + ).to(dtype)
615 + local_query = query_storage[:, :heads_per_rank]
616 + assert local_query.is_contiguous() is not noncontiguous
617 +
618 + active_ubatch[0] = iteration % 2
619 + actual = workspace.gather(local_query)
620 + torch.accelerator.synchronize()
621 +
622 + assert (
623 + actual.data_ptr()
624 + == workspace.final_query[active_ubatch[0]].data_ptr()
625 + )
626 + assert workspace.completion[active_ubatch[0]].numel() == 1
627 + _assert_q_gather_matches_reference(
628 + actual, local_query, world_size, padded_num_heads
629 + )
630 +
631 + if env.get("TEST_CUDA_GRAPH") != "1" or dtype != torch.bfloat16:
632 + continue
633 +
634 + capture_num_tokens = 17
635 + capture_storage = torch.empty(
636 + capture_num_tokens,
637 + heads_per_rank + 2,
638 + head_dim,
639 + dtype=dtype,
640 + device=device,
641 + )
642 + captured_input = capture_storage[:, :heads_per_rank]
643 + assert not captured_input.is_contiguous()
644 + input_pattern = (
645 + torch.arange(captured_input.numel(), device=device, dtype=torch.int32)
646 + .remainder(31)
647 + .view(captured_input.shape)
648 + .to(dtype)
649 + )
650 +
651 + active_ubatch[0] = 1
652 + torch.accelerator.synchronize()
653 + graph = torch.cuda.CUDAGraph()
654 + with torch.cuda.graph(graph):
655 + captured_output = workspace.gather(captured_input)
656 + torch.accelerator.synchronize()
657 + dist.barrier()
658 +
659 + eager_input = torch.full(
660 + captured_input.shape,
661 + 96 + rank,
662 + dtype=dtype,
663 + device=device,
664 + )
665 + eager_output = workspace.gather(eager_input)
666 + torch.accelerator.synchronize()
667 + _assert_q_gather_matches_reference(
668 + eager_output, eager_input, world_size, padded_num_heads
669 + )
670 + epoch_before_replays = int(workspace.epoch[1].item())
671 +
672 + for replay in range(2):
673 + captured_input.copy_(input_pattern + rank * 32 + replay)
674 + torch.accelerator.synchronize()
675 + graph.replay()
676 + torch.accelerator.synchronize()
677 + assert int(workspace.epoch[1].item()) == (
678 + epoch_before_replays + replay + 1
679 + )
680 + _assert_q_gather_matches_reference(
681 + captured_output,
682 + captured_input,
683 + world_size,
684 + padded_num_heads,
685 + )
686 +
687 + interleaved_eager_input = torch.full(
688 + captured_input.shape,
689 + 160 + rank,
690 + dtype=dtype,
691 + device=device,
692 + )
693 + interleaved_eager_output = workspace.gather(interleaved_eager_input)
694 + torch.accelerator.synchronize()
695 + _assert_q_gather_matches_reference(
696 + interleaved_eager_output,
697 + interleaved_eager_input,
698 + world_size,
699 + padded_num_heads,
700 + )
701 +
702 + captured_input.copy_(input_pattern + rank * 32 + 2)
703 + torch.accelerator.synchronize()
704 + epoch_before_mixed_replay = int(workspace.epoch[1].item())
705 + graph.replay()
706 + torch.accelerator.synchronize()
707 + assert int(workspace.epoch[1].item()) == epoch_before_mixed_replay + 1
708 + _assert_q_gather_matches_reference(
709 + captured_output, captured_input, world_size, padded_num_heads
710 + )
711 + finally:
712 + dist.destroy_process_group()
713 +
714 +
715 +@pytest.mark.parametrize(
716 + "world_size",
717 + [
718 + pytest.param(
719 + 4,
720 + marks=pytest.mark.skipif(
721 + torch.accelerator.device_count() < 4 or not _has_multicast_support(),
722 + reason="Need 4 GPUs with symmetric-memory multicast.",
723 + ),
724 + ),
725 + ],
726 +)
727 +def test_distributed_direct_q_gather_cuda_graph_replay(world_size: int):
728 + _distributed_run(
729 + _distributed_direct_q_gather_worker,
730 + world_size=world_size,
731 + extra_env={"TEST_CUDA_GRAPH": "1"},
732 + )
733 +
734 +
735 +def _distributed_direct_kv_gather_worker(env: dict[str, str]) -> None:
736 + update_environment_variables(env)
737 + local_rank = int(env["LOCAL_RANK"])
738 + device = torch.device(f"cuda:{local_rank}")
739 + torch.accelerator.set_device_index(local_rank)
740 + dist.init_process_group(backend="nccl")
741 + try:
742 + rank = dist.get_rank()
743 + world_size = dist.get_world_size()
744 + token_dim = 576
745 + max_gathered_tokens = 128 * world_size
746 + active_ubatch = [0]
747 + dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0]
748 +
749 + for dtype_idx, dtype_name in enumerate(("bfloat16", "float16")):
750 + dtype = _dtype_from_name(dtype_name)
751 + workspace = dcp_utils.DirectDCPKVGatherWorkspace(
752 + dist.group.WORLD,
753 + device,
754 + max_gathered_tokens,
755 + token_dim,
756 + dtype,
757 + num_ubatches=2,
758 + )
759 +
760 + # Use disjoint slices of one persistent chunked-context workspace.
761 + storage = torch.zeros(
762 + (world_size + 1) * 128, token_dim, device=device, dtype=dtype
763 + )
764 + for iteration, num_tokens in enumerate((1, 128, 17)):
765 + generator = torch.Generator(device=device)
766 + generator.manual_seed(7000 + rank * 101 + dtype_idx * 977 + iteration)
767 + local_kv = storage[:num_tokens]
768 + local_kv.copy_(
769 + torch.randn(
770 + num_tokens,
771 + token_dim,
772 + device=device,
773 + dtype=torch.float32,
774 + generator=generator,
775 + ).to(dtype)
776 + )
777 + gathered = storage[128 : 128 + num_tokens * world_size]
778 + active_ubatch[0] = iteration % 2
779 + workspace.gather(gathered, local_kv)
780 + torch.accelerator.synchronize()
781 +
782 + expected = torch.empty_like(gathered)
783 + dist.all_gather_into_tensor(expected, local_kv.contiguous())
784 + assert torch.equal(
785 + gathered.view(torch.uint8), expected.view(torch.uint8)
786 + )
787 + finally:
788 + dist.destroy_process_group()
789 +
790 +
791 +@pytest.mark.parametrize(
792 + "world_size",
793 + [
794 + pytest.param(
795 + 4,
796 + marks=pytest.mark.skipif(
797 + torch.accelerator.device_count() < 4 or not _has_multicast_support(),
798 + reason="Need 4 GPUs with symmetric-memory multicast.",
799 + ),
800 + ),
801 + ],
802 +)
803 +def test_distributed_direct_kv_gather_matches_reference(world_size: int):
804 + _distributed_run(
805 + _distributed_direct_kv_gather_worker,
806 + world_size=world_size,
807 + extra_env={},
808 + )
809 +
810 +
811 +def _distributed_direct_a2a_worker(env: dict[str, str]) -> None:
812 + update_environment_variables(env)
813 + local_rank = int(env["LOCAL_RANK"])
814 + device = torch.device(f"cuda:{local_rank}")
815 + torch.accelerator.set_device_index(local_rank)
816 + dist.init_process_group(backend="nccl")
817 + try:
818 + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
819 +
820 + rank = dist.get_rank()
821 + world_size = dist.get_world_size()
822 + dtype = _dtype_from_name(env["TEST_DTYPE"])
823 + lse_dtype = _dtype_from_name(env["TEST_LSE_DTYPE"])
824 + is_lse_base_on_e = env["LSE_BASE_E"] == "1"
825 + # Match Kimi-K3's six heads per DCP rank.
826 + heads_per_rank, head_dim, max_num_tokens = 6, 512, 128
827 + total_heads = world_size * heads_per_rank
828 + active_ubatch = [0]
829 + dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0]
830 + workspace = dcp_utils.DirectDCPA2AWorkspace(
831 + dist.group.WORLD,
832 + device,
833 + max_num_tokens,
834 + heads_per_rank,
835 + head_dim,
836 + dtype,
837 + num_ubatches=2,
838 + )
839 +
840 + def check(num_tokens: int, iteration: int, padded: bool) -> None:
841 + generator = torch.Generator(device=device)
842 + generator.manual_seed(1234 + rank + iteration * world_size)
843 + storage_heads = 128 if padded else total_heads
844 + partial_output_storage = torch.randn(
845 + num_tokens,
846 + storage_heads,
847 + head_dim,
848 + device=device,
849 + dtype=dtype,
850 + generator=generator,
851 + )
852 + partial_lse_storage = torch.randn(
853 + num_tokens,
854 + storage_heads,
855 + device=device,
856 + dtype=lse_dtype,
857 + generator=generator,
858 + )
859 + partial_output = partial_output_storage[:, :total_heads, :]
860 + partial_lse = partial_lse_storage[:, :total_heads]
861 + if padded:
862 + assert not partial_output.is_contiguous()
863 + assert not partial_lse.is_contiguous()
864 + active_ubatch[0] = iteration % 2
865 + actual = workspace.lse_reduce(partial_output, partial_lse, is_lse_base_on_e)
866 + torch.accelerator.synchronize()
867 +
868 + reference_output = partial_output.contiguous()
869 + reference_lse = partial_lse.contiguous()
870 + gathered_output = [
871 + torch.empty_like(reference_output) for _ in range(world_size)
872 + ]
873 + gathered_lse = [torch.empty_like(reference_lse) for _ in range(world_size)]
874 + dist.all_gather(gathered_output, reference_output)
875 + dist.all_gather(gathered_lse, reference_lse)
876 + outputs = torch.stack(
877 + [
878 + value[
879 + :,
880 + rank * heads_per_rank : (rank + 1) * heads_per_rank,
881 + :,
882 + ]
883 + for value in gathered_output
884 + ]
885 + ).float()
886 + lses = torch.stack(
887 + [
888 + value[:, rank * heads_per_rank : (rank + 1) * heads_per_rank]
889 + for value in gathered_lse
890 + ]
891 + )
892 + expected = _lse_weighted_combine(
893 + outputs, lses, is_lse_base_on_e=is_lse_base_on_e
894 + )
895 + _assert_close(actual, expected, dtype)
896 +
897 + def check_empty_shards(query_lens: list[int], iteration: int) -> None:
898 + """Verify empty local shards contribute zero weight."""
899 + num_seqs = len(query_lens)
900 + num_tokens = sum(query_lens)
901 + generator = torch.Generator(device=device)
902 + generator.manual_seed(777 + rank + iteration * world_size)
903 + partial_output = torch.randn(
904 + num_tokens,
905 + total_heads,
906 + head_dim,
907 + device=device,
908 + dtype=dtype,
909 + generator=generator,
910 + )
911 + partial_lse = torch.randn(
912 + num_tokens,
913 + total_heads,
914 + device=device,
915 + dtype=lse_dtype,
916 + generator=generator,
917 + )
918 +
919 + # Cover globally empty, rank-local empty, and non-empty sequences.
920 + def is_empty(seq_idx: int, source_rank: int) -> bool:
921 + return seq_idx == 0 or seq_idx == 1 + source_rank
922 +
923 + seq_lens = torch.tensor(
924 + [
925 + 0 if is_empty(seq_idx, rank) else seq_idx + 3
926 + for seq_idx in range(num_seqs)
927 + ],
928 + dtype=torch.int32,
929 + device=device,
930 + )
931 + query_lens_tensor = torch.tensor(
932 + query_lens,
933 + dtype=torch.int32,
934 + device=device,
935 + )
936 + query_start_loc = torch.cat(
937 + (
938 + query_lens_tensor.new_zeros(1),
939 + query_lens_tensor.cumsum(0),
940 + )
941 + )
942 + empty_rows = torch.repeat_interleave(seq_lens == 0, query_lens_tensor)
943 + # Model undefined attention rows with NaNs.
944 + partial_output[empty_rows] = float("nan")
945 + partial_lse[empty_rows] = float("nan")
946 +
947 + active_ubatch[0] = iteration % 2
948 + actual = workspace.lse_reduce(
949 + partial_output,
950 + partial_lse,
951 + is_lse_base_on_e,
952 + seq_lens=seq_lens,
953 + query_start_loc=query_start_loc,
954 + )
955 + torch.accelerator.synchronize()
956 +
957 + gathered_output = [
958 + torch.empty_like(partial_output) for _ in range(world_size)
959 + ]
960 + gathered_lse = [torch.empty_like(partial_lse) for _ in range(world_size)]
961 + dist.all_gather(gathered_output, partial_output.contiguous())
962 + dist.all_gather(gathered_lse, partial_lse.contiguous())
963 + head_slice = slice(rank * heads_per_rank, (rank + 1) * heads_per_rank)
964 + outputs = torch.stack(
965 + [value[:, head_slice, :] for value in gathered_output]
966 + ).float()
967 + lses = torch.stack([value[:, head_slice] for value in gathered_lse])
968 + for source_rank in range(world_size):
969 + source_empty = torch.tensor(
970 + [is_empty(seq_idx, source_rank) for seq_idx in range(num_seqs)],
971 + device=device,
972 + ).repeat_interleave(query_lens_tensor)
973 + outputs[source_rank][source_empty] = 0.0
974 + lses[source_rank][source_empty] = float("-inf")
975 + expected = _lse_weighted_combine(
976 + outputs, lses, is_lse_base_on_e=is_lse_base_on_e
977 + )
978 + all_empty = torch.tensor(
979 + [
980 + all(is_empty(seq_idx, r) for r in range(world_size))
981 + for seq_idx in range(num_seqs)
982 + ],
983 + device=device,
984 + ).repeat_interleave(query_lens_tensor)
985 + assert torch.equal(actual[all_empty], torch.zeros_like(actual[all_empty]))
986 + assert not torch.isnan(actual.float()).any()
987 + _assert_close(actual, expected, dtype)
988 +
989 + cases = ((1, False), (17, True), (128, True))
990 + for iteration, (num_tokens, padded) in enumerate(cases):
991 + check(num_tokens, iteration, padded)
992 + check_empty_shards(
993 + query_lens=[1, 3, 2, *([1] * (world_size - 1))],
994 + iteration=len(cases),
995 + )
996 + generator = torch.Generator(device=device)
997 + generator.manual_seed(4321 + rank)
998 + partial_output_storage = torch.randn(
999 + 128,
1000 + 128,
1001 + head_dim,
1002 + device=device,
1003 + dtype=dtype,
1004 + generator=generator,
1005 + )
1006 + partial_lse_storage = torch.randn(
1007 + 128,
1008 + 128,
1009 + device=device,
1010 + dtype=lse_dtype,
1011 + generator=generator,
1012 + )
1013 + partial_output = partial_output_storage[:, :total_heads, :]
1014 + partial_lse = partial_lse_storage[:, :total_heads]
1015 + assert not partial_output.is_contiguous()
1016 + assert not partial_lse.is_contiguous()
1017 + torch.accelerator.synchronize()
1018 + active_ubatch[0] = 1
1019 + graph = torch.cuda.CUDAGraph()
1020 + with torch.cuda.graph(graph):
1021 + actual = workspace.lse_reduce(partial_output, partial_lse, is_lse_base_on_e)
1022 + for _ in range(3):
1023 + graph.replay()
1024 + torch.accelerator.synchronize()
1025 +
1026 + reference_output = partial_output.contiguous()
1027 + reference_lse = partial_lse.contiguous()
1028 + gathered_output = [
1029 + torch.empty_like(reference_output) for _ in range(world_size)
1030 + ]
1031 + gathered_lse = [torch.empty_like(reference_lse) for _ in range(world_size)]
1032 + dist.all_gather(gathered_output, reference_output)
1033 + dist.all_gather(gathered_lse, reference_lse)
1034 + head_slice = slice(rank * heads_per_rank, (rank + 1) * heads_per_rank)
1035 + outputs = torch.stack(
1036 + [value[:, head_slice, :] for value in gathered_output]
1037 + ).float()
1038 + lses = torch.stack([value[:, head_slice] for value in gathered_lse])
1039 + expected = _lse_weighted_combine(
1040 + outputs, lses, is_lse_base_on_e=is_lse_base_on_e
1041 + )
1042 + _assert_close(actual, expected, dtype)
1043 + finally:
1044 + dist.destroy_process_group()
1045 +
1046 +
1047 +@pytest.mark.parametrize(
1048 + "world_size",
1049 + [
1050 + pytest.param(
1051 + 2,
1052 + marks=pytest.mark.skipif(
1053 + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs."
1054 + ),
1055 + ),
1056 + pytest.param(
1057 + 4,
1058 + marks=pytest.mark.skipif(
1059 + torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs."
1060 + ),
1061 + ),
1062 + ],
1063 +)
1064 +def test_distributed_direct_a2a_matches_reference(world_size: int):
1065 + _distributed_run(
1066 + _distributed_direct_a2a_worker,
1067 + world_size=world_size,
1068 + extra_env={
1069 + "TEST_DTYPE": "bfloat16",
1070 + "TEST_LSE_DTYPE": "bfloat16",
1071 + "LSE_BASE_E": "0",
1072 + },
1073 + )
1074 +
1075 +
1076 +if __name__ == "__main__":
1077 + pytest.main([__file__, "-v"])
tests/distributed/test_kimi_linear_context_parallel.py new
+128
@@ -0,0 +1,128 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +
6 +import pytest
7 +
8 +from tests.conftest import VllmRunner
9 +from tests.utils import create_new_process_for_each_test
10 +from vllm import SamplingParams, TokensPrompt
11 +
12 +MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
13 +
14 +
15 +def _make_tiny_overrides() -> dict:
16 + linear_attn_config = {
17 + "full_attn_layers": [4],
18 + "head_dim": 128,
19 + "kda_layers": [1, 2, 3],
20 + "num_heads": 8,
21 + "short_conv_kernel_size": 4,
22 + "use_full_rank_gate": True,
23 + }
24 +
25 + return {
26 + "bos_token_id": 1,
27 + "eos_token_id": 2,
28 + "pad_token_id": 0,
29 + "vocab_size": 256,
30 + "hidden_size": 256,
31 + "head_dim": 32,
32 + "intermediate_size": 512,
33 + "num_hidden_layers": 4,
34 + "num_attention_heads": 8,
35 + "num_key_value_heads": 8,
36 + "num_experts": None,
37 + "num_experts_per_token": None,
38 + "num_shared_experts": 0,
39 + "q_lora_rank": 128,
40 + "kv_lora_rank": 512,
41 + "qk_nope_head_dim": 128,
42 + "qk_rope_head_dim": 64,
43 + "v_head_dim": 128,
44 + "mla_use_nope": True,
45 + "mla_use_output_gate": True,
46 + "linear_attn_config": linear_attn_config,
47 + "max_position_embeddings": 512,
48 + "model_max_length": 512,
49 + }
50 +
51 +
52 +def _run_tiny_model(
53 + vllm_runner: type[VllmRunner], dcp_size: int
54 +) -> list[tuple[list[int], list[float]]]:
55 + with vllm_runner(
56 + model_name=MODEL,
57 + skip_tokenizer_init=True,
58 + load_format="dummy",
59 + hf_overrides=_make_tiny_overrides(),
60 + tensor_parallel_size=2,
61 + decode_context_parallel_size=dcp_size,
62 + cp_kv_cache_interleave_size=1,
63 + distributed_executor_backend="mp",
64 + dtype="bfloat16",
65 + seed=0,
66 + enforce_eager=True,
67 + max_model_len=512,
68 + max_num_seqs=8,
69 + max_num_batched_tokens=1024,
70 + gpu_memory_utilization=0.25,
71 + enable_chunked_prefill=True,
72 + enable_prefix_caching=False,
73 + ) as runner:
74 + lengths = [1, 32, 33, 129]
75 + prompts = [
76 + TokensPrompt(
77 + prompt_token_ids=[
78 + 3 + ((request_idx + token_idx) % 251) for token_idx in range(length)
79 + ]
80 + )
81 + for request_idx, length in enumerate(lengths)
82 + ]
83 + outputs = runner.llm.generate(
84 + prompts,
85 + SamplingParams(
86 + temperature=0,
87 + max_tokens=8,
88 + seed=0,
89 + ignore_eos=True,
90 + logprobs=20,
91 + ),
92 + use_tqdm=False,
93 + )
94 +
95 + results = []
96 + for request_output in outputs:
97 + completion = request_output.outputs[0]
98 + token_ids = list(completion.token_ids)
99 + assert completion.logprobs is not None
100 + selected_logprobs = [
101 + step_logprobs[token_id].logprob
102 + for token_id, step_logprobs in zip(token_ids, completion.logprobs)
103 + ]
104 + results.append((token_ids, selected_logprobs))
105 + return results
106 +
107 +
108 +@create_new_process_for_each_test()
109 +@pytest.mark.distributed(num_gpus=2)
110 +def test_kimi_linear_dcp_tiny(
111 + vllm_runner: type[VllmRunner],
112 + num_gpus_available: int,
113 +) -> None:
114 + if num_gpus_available < 2:
115 + pytest.skip("Need at least 2 GPUs")
116 +
117 + baseline = _run_tiny_model(vllm_runner, dcp_size=1)
118 + dcp = _run_tiny_model(vllm_runner, dcp_size=2)
119 +
120 + assert [tokens for tokens, _ in dcp] == [tokens for tokens, _ in baseline]
121 + logprob_drifts = []
122 + for (_, baseline_logprobs), (_, dcp_logprobs) in zip(baseline, dcp):
123 + for baseline_logprob, dcp_logprob in zip(baseline_logprobs, dcp_logprobs):
124 + assert math.isfinite(baseline_logprob)
125 + assert math.isfinite(dcp_logprob)
126 + logprob_drifts.append(abs(baseline_logprob - dcp_logprob))
127 +
128 + assert max(logprob_drifts) <= 1e-2
tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py new
+246
@@ -0,0 +1,246 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +import vllm._custom_ops as ops
8 +from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import (
9 + fused_mla_decode_q_concat_kv_cache_insert,
10 + fused_mla_key_concat_ds_mla_insert,
11 + fused_mla_key_concat_kv_cache_insert,
12 + fused_mla_qkv_quant_kv_cache_fp8_insert,
13 +)
14 +from vllm.platforms import current_platform
15 +
16 +NUM_TOKENS = 4
17 +NUM_HEADS = 2
18 +KV_LORA_RANK = 512
19 +NOPE_HEAD_DIM = 128
20 +ROPE_HEAD_DIM = 64
21 +QK_HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM
22 +V_HEAD_DIM = 128
23 +CACHE_ENTRY = KV_LORA_RANK + ROPE_HEAD_DIM
24 +DS_MLA_CACHE_ENTRY = 656
25 +BLOCK_SIZE = 8
26 +OWNED_TOKEN_INDICES = [1, 3]
27 +
28 +pytestmark = pytest.mark.skipif(
29 + not current_platform.is_cuda(), reason="Kimi-K3 fused MLA ops require CUDA"
30 +)
31 +
32 +
33 +def _inputs() -> dict[str, torch.Tensor]:
34 + torch.manual_seed(0)
35 +
36 + def make(*shape: int) -> torch.Tensor:
37 + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16)
38 +
39 + return {
40 + "q": make(NUM_TOKENS, NUM_HEADS, QK_HEAD_DIM),
41 + "k_nope": make(NUM_TOKENS, NUM_HEADS, NOPE_HEAD_DIM),
42 + "k_pe": make(NUM_TOKENS, ROPE_HEAD_DIM),
43 + "kv_c": make(NUM_TOKENS, KV_LORA_RANK),
44 + "v": make(NUM_TOKENS, NUM_HEADS, V_HEAD_DIM),
45 + "ql_nope": make(NUM_TOKENS, NUM_HEADS, KV_LORA_RANK),
46 + "q_pe": make(NUM_TOKENS, NUM_HEADS, ROPE_HEAD_DIM),
47 + }
48 +
49 +
50 +def _slot_mapping() -> torch.Tensor:
51 + return torch.tensor([-1, 0, -1, 5], device="cuda", dtype=torch.int64)
52 +
53 +
54 +def _owned(tensor: torch.Tensor) -> torch.Tensor:
55 + return tensor[OWNED_TOKEN_INDICES]
56 +
57 +
58 +def _full_key(k_nope: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor:
59 + shared_pe = k_pe[:, None, :].expand(-1, k_nope.shape[1], -1)
60 + return torch.cat((k_nope, shared_pe), dim=-1)
61 +
62 +
63 +def _latent_query(ql_nope: torch.Tensor, q_pe: torch.Tensor) -> torch.Tensor:
64 + return torch.cat((ql_nope, q_pe), dim=-1)
65 +
66 +
67 +def _make_cache(cache_format: str) -> torch.Tensor:
68 + if cache_format == "bf16":
69 + return torch.full(
70 + (1, BLOCK_SIZE, CACHE_ENTRY),
71 + -7,
72 + device="cuda",
73 + dtype=torch.bfloat16,
74 + )
75 + if cache_format == "fp8":
76 + return torch.full(
77 + (1, BLOCK_SIZE, CACHE_ENTRY),
78 + 1.0,
79 + device="cuda",
80 + dtype=torch.float32,
81 + ).to(torch.float8_e4m3fn)
82 + assert cache_format == "fp8_ds_mla"
83 + return torch.full(
84 + (1, BLOCK_SIZE, DS_MLA_CACHE_ENTRY),
85 + 165,
86 + device="cuda",
87 + dtype=torch.uint8,
88 + )
89 +
90 +
91 +def _reference_cache(
92 + cache_format: str,
93 + inputs: dict[str, torch.Tensor],
94 + initial_cache: torch.Tensor,
95 + slots: torch.Tensor,
96 +) -> torch.Tensor:
97 + reference_cache = initial_cache.clone()
98 + valid_slots = _owned(slots)
99 + if cache_format == "fp8_ds_mla":
100 + ops.concat_and_cache_mla(
101 + _owned(inputs["kv_c"]),
102 + _owned(inputs["k_pe"]),
103 + reference_cache,
104 + valid_slots,
105 + cache_format,
106 + torch.ones(1, device="cuda", dtype=torch.float32),
107 + )
108 + return reference_cache
109 +
110 + latent = torch.cat(
111 + (_owned(inputs["kv_c"]), _owned(inputs["k_pe"])),
112 + dim=-1,
113 + )
114 + flat_reference_cache = reference_cache.view(-1, CACHE_ENTRY)
115 + for row, slot in zip(latent, valid_slots.tolist()):
116 + flat_reference_cache[slot].copy_(row.to(reference_cache.dtype))
117 + return reference_cache
118 +
119 +
120 +def _assert_cache_matches_reference(
121 + cache: torch.Tensor,
122 + reference_cache: torch.Tensor,
123 + initial_cache: torch.Tensor,
124 +) -> None:
125 + assert torch.equal(cache, reference_cache)
126 + for slot in (0, 5):
127 + assert not torch.equal(cache[0, slot], initial_cache[0, slot])
128 + for slot in (1, 2, 3, 4, 6, 7):
129 + assert torch.equal(cache[0, slot], initial_cache[0, slot])
130 +
131 +
132 +@pytest.mark.parametrize("cache_format", ["bf16", "fp8_ds_mla"])
133 +def test_prefill_concat_ignores_negative_slots_for_cache(cache_format: str) -> None:
134 + inputs = _inputs()
135 + slots = _slot_mapping()
136 + initial_cache = _make_cache(cache_format)
137 + mixed_cache = initial_cache.clone()
138 + reference_cache = _reference_cache(cache_format, inputs, initial_cache, slots)
139 +
140 + op = (
141 + fused_mla_key_concat_kv_cache_insert
142 + if cache_format == "bf16"
143 + else fused_mla_key_concat_ds_mla_insert
144 + )
145 + output = op(
146 + inputs["q"],
147 + inputs["k_nope"],
148 + inputs["k_pe"],
149 + inputs["kv_c"],
150 + mixed_cache,
151 + slots,
152 + )
153 +
154 + torch.testing.assert_close(
155 + output, _full_key(inputs["k_nope"], inputs["k_pe"]), rtol=0, atol=0
156 + )
157 + _assert_cache_matches_reference(mixed_cache, reference_cache, initial_cache)
158 +
159 +
160 +def test_fp8_prefill_outputs_ignore_negative_slots_for_cache() -> None:
161 + inputs = _inputs()
162 + slots = _slot_mapping()
163 + initial_cache = _make_cache("fp8")
164 + mixed_cache = initial_cache.clone()
165 + reference_cache = _reference_cache("fp8", inputs, initial_cache, slots)
166 + scale_inv = torch.ones(1, device="cuda", dtype=torch.float32)
167 +
168 + q, k, v = fused_mla_qkv_quant_kv_cache_fp8_insert(
169 + inputs["q"],
170 + inputs["k_nope"],
171 + inputs["k_pe"],
172 + inputs["kv_c"],
173 + inputs["v"],
174 + mixed_cache,
175 + slots,
176 + scale_inv,
177 + scale_inv,
178 + scale_inv,
179 + scale_inv,
180 + )
181 +
182 + expected = (
183 + inputs["q"].to(torch.float8_e4m3fn),
184 + _full_key(inputs["k_nope"], inputs["k_pe"]).to(torch.float8_e4m3fn),
185 + inputs["v"].to(torch.float8_e4m3fn),
186 + )
187 + for actual, reference in zip((q, k, v), expected):
188 + torch.testing.assert_close(actual.float(), reference.float(), rtol=0, atol=0)
189 + _assert_cache_matches_reference(mixed_cache, reference_cache, initial_cache)
190 +
191 +
192 +@pytest.mark.parametrize("cache_format", ["bf16", "fp8", "fp8_ds_mla"])
193 +def test_decode_concat_ignores_negative_slots_for_cache(cache_format: str) -> None:
194 + inputs = _inputs()
195 + slots = _slot_mapping()
196 + initial_cache = _make_cache(cache_format)
197 + mixed_cache = initial_cache.clone()
198 + reference_cache = _reference_cache(cache_format, inputs, initial_cache, slots)
199 + scale_inv = torch.ones(1, device="cuda", dtype=torch.float32)
200 + kwargs = {
201 + "ds_mla": cache_format == "fp8_ds_mla",
202 + "q_scale_inv": scale_inv if cache_format == "fp8" else None,
203 + "cache_scale_inv": scale_inv if cache_format == "fp8" else None,
204 + }
205 +
206 + output = fused_mla_decode_q_concat_kv_cache_insert(
207 + inputs["ql_nope"],
208 + inputs["q_pe"],
209 + inputs["kv_c"],
210 + inputs["k_pe"],
211 + mixed_cache,
212 + slots,
213 + **kwargs,
214 + )
215 +
216 + expected = _latent_query(inputs["ql_nope"], inputs["q_pe"])
217 + if cache_format == "fp8":
218 + expected = expected.to(torch.float8_e4m3fn)
219 + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0)
220 + else:
221 + torch.testing.assert_close(output, expected, rtol=0, atol=0)
222 + _assert_cache_matches_reference(mixed_cache, reference_cache, initial_cache)
223 +
224 +
225 +def test_ds_mla_cache_insert_bit_compatible_with_reference() -> None:
226 + """Fused ds_mla insertion must match the reference bit-for-bit."""
227 + torch.manual_seed(0)
228 + num_tokens, num_blocks = 33, 16
229 + dt = torch.bfloat16
230 + kv_c = torch.randn(num_tokens, KV_LORA_RANK, device="cuda", dtype=dt)
231 + k_pe = torch.randn(num_tokens, 1, ROPE_HEAD_DIM, device="cuda", dtype=dt)
232 + k_nope = torch.randn(num_tokens, NUM_HEADS, NOPE_HEAD_DIM, device="cuda", dtype=dt)
233 + q = torch.randn(num_tokens, NUM_HEADS, QK_HEAD_DIM, device="cuda", dtype=dt)
234 + slots = torch.randperm(num_blocks * BLOCK_SIZE, device="cuda")[:num_tokens].long()
235 + slots[torch.rand(num_tokens, device="cuda") < 0.3] = -1
236 + ref = torch.full(
237 + (num_blocks, BLOCK_SIZE, DS_MLA_CACHE_ENTRY),
238 + 165,
239 + device="cuda",
240 + dtype=torch.uint8,
241 + )
242 + got = ref.clone()
243 + scale = torch.ones(1, device="cuda", dtype=torch.float32)
244 + ops.concat_and_cache_mla(kv_c, k_pe.squeeze(1), ref, slots, "fp8_ds_mla", scale)
245 + fused_mla_key_concat_ds_mla_insert(q, k_nope, k_pe, kv_c, got, slots)
246 + assert torch.equal(ref, got)
tests/v1/attention/test_flashinfer_mla_dcp.py new
+80
@@ -0,0 +1,80 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from unittest.mock import MagicMock
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm.platforms import current_platform
10 +
11 +requires_flashinfer_mla = pytest.mark.skipif(
12 + not current_platform.has_device_capability(100),
13 + reason="FlashInfer MLA requires compute capability 10 or above.",
14 +)
15 +
16 +
17 +def test_mla_dcp_gathered_query_reserves_backend_head_storage():
18 + from vllm.v1.attention.ops.dcp_utils import reserve_query_head_storage
19 +
20 + query = torch.randn(3, 24, 576, dtype=torch.bfloat16)
21 +
22 + padded = reserve_query_head_storage(query, 128)
23 +
24 + assert padded.shape == query.shape
25 + assert padded.stride() == query.stride()
26 + assert padded.untyped_storage().nbytes() >= 3 * 128 * 576 * 2
27 + assert torch.equal(padded, query)
28 +
29 +
30 +@requires_flashinfer_mla
31 +def test_flashinfer_mla_selects_backend_for_gathered_heads():
32 + from vllm.v1.attention.backends.mla.flashinfer_mla import (
33 + _select_mla_decode_backend,
34 + )
35 +
36 + assert _select_mla_decode_backend(6) is None
37 + assert _select_mla_decode_backend(24) == "cute-dsl"
38 +
39 +
40 +@requires_flashinfer_mla
41 +def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch):
42 + import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla
43 +
44 + impl = MagicMock()
45 + impl.bmm1_scale = 1.0
46 + impl.bmm2_scale = 1.0
47 + impl.need_to_return_lse_for_decode = True
48 + impl.num_heads = 6
49 + impl.qk_nope_head_dim = 128
50 + impl.kv_lora_rank = 512
51 + impl.qk_rope_head_dim = 64
52 + impl.kv_cache_dtype = "auto"
53 +
54 + attn_metadata = MagicMock()
55 + attn_metadata.causal = True
56 + attn_metadata.num_decode_tokens = 2
57 + attn_metadata.num_decodes = 2
58 + attn_metadata.max_seq_len = 1
59 + attn_metadata.decode.block_table = torch.zeros(2, 1, dtype=torch.int32)
60 + attn_metadata.decode.seq_lens = torch.ones(2, dtype=torch.int32)
61 +
62 + query = torch.ones(2, 24, 576, dtype=torch.bfloat16)
63 + kv_cache = torch.ones(1, 128, 576, dtype=torch.bfloat16)
64 + kernel = MagicMock(
65 + return_value=(
66 + torch.ones(2, 1, 24, 512, dtype=torch.bfloat16),
67 + torch.ones(2, 24, dtype=torch.float32),
68 + )
69 + )
70 + monkeypatch.setattr(flashinfer_mla, "_get_workspace_buffer", MagicMock())
71 + monkeypatch.setattr(flashinfer_mla, "trtllm_batch_decode_with_kv_cache_mla", kernel)
72 + layer = MagicMock()
73 +
74 + output, lse = flashinfer_mla.FlashInferMLAImpl.forward_mqa(
75 + impl, query, kv_cache, attn_metadata, layer
76 + )
77 +
78 + assert kernel.call_args.kwargs["backend"] == "cute-dsl"
79 + assert output.shape == (2, 24, 512)
80 + assert lse is not None and lse.shape == (2, 24)
vllm/envs.py
+13
@@ -190,6 +190,9 @@ if TYPE_CHECKING:
190 VLLM_USE_DEEP_GEMM_E8M0: bool = True
191 VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True
192 VLLM_DCP_Q_REPLICATE: bool = False
193 + VLLM_USE_DIRECT_DCP_A2A: bool | None = None
194 + VLLM_USE_DIRECT_DCP_Q_GATHER: bool | None = None
195 + VLLM_USE_DIRECT_DCP_KV_GATHER: bool | None = None
196 VLLM_DEEP_GEMM_WARMUP: Literal[
197 "skip",
198 "full",
@@ -2057,6 +2060,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
2060 "VLLM_USE_SIMPLE_KV_OFFLOAD": lambda: bool(
2061 int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0"))
2062 ),
2063 + # Direct DCP ops default on when applicable; set to 1 to enforce or 0 to disable.
2064 + "VLLM_USE_DIRECT_DCP_A2A": lambda: maybe_convert_bool(
2065 + os.getenv("VLLM_USE_DIRECT_DCP_A2A")
2066 + ),
2067 + "VLLM_USE_DIRECT_DCP_Q_GATHER": lambda: maybe_convert_bool(
2068 + os.getenv("VLLM_USE_DIRECT_DCP_Q_GATHER")
2069 + ),
2070 + "VLLM_USE_DIRECT_DCP_KV_GATHER": lambda: maybe_convert_bool(
2071 + os.getenv("VLLM_USE_DIRECT_DCP_KV_GATHER")
2072 + ),
2073 # Whether to enable dual cuda streams for LoRA computation
2074 # (used by both BaseLinearLayerWithLoRA and FusedMoEWithLoRA to
2075 # overlap the base layer compute with the LoRA fast path).
vllm/model_executor/layers/attention/mla_attention.py
+87 -48
@@ -206,6 +206,7 @@ from abc import abstractmethod
206 from collections.abc import Callable
207 from dataclasses import dataclass
208 from enum import Enum
209 +from math import lcm
210 from typing import ClassVar, Generic, TypeVar, cast
211
212 import numpy as np
@@ -222,7 +223,6 @@ from vllm.config import (
223 ModelConfig,
224 VllmConfig,
225 get_current_vllm_config,
225 - get_current_vllm_config_or_none,
226 )
227 from vllm.config.cache import CacheDType
228 from vllm.distributed.parallel_state import (
@@ -264,7 +264,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
264 from vllm.model_executor.utils import replace_parameter
265 from vllm.platforms import current_platform
266 from vllm.utils.flashinfer import has_flashinfer
267 -from vllm.utils.math_utils import cdiv, round_down
267 +from vllm.utils.math_utils import cdiv, round_down, round_up
268 from vllm.utils.torch_utils import (
269 LayerNameType,
270 _encode_layer_name,
@@ -292,8 +292,7 @@ from vllm.v1.attention.backends.utils import (
292 get_num_attention_heads_from_layers,
293 split_decodes_and_prefills,
294 )
295 -from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs
296 -from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
295 +from vllm.v1.attention.ops.dcp_utils import MLADCPManager
296 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
297 from vllm.v1.attention.selector import get_attn_backend
298 from vllm.v1.kv_cache_interface import (
@@ -579,12 +578,27 @@ class MLAAttention(nn.Module, AttentionLayerBase):
578
579 self.use_sparse = use_sparse
580
582 - _vllm_config = get_current_vllm_config_or_none()
583 - self.dcp_a2a = (
584 - _vllm_config is not None
585 - and _vllm_config.parallel_config.decode_context_parallel_size > 1
586 - and _vllm_config.parallel_config.dcp_comm_backend == "a2a"
587 - )
581 + self.dcp_manager: MLADCPManager | None = None
582 + if self.impl.dcp_world_size > 1:
583 + query_dtype = (
584 + current_platform.fp8_dtype()
585 + if is_quantized_kv_cache(self.kv_cache_dtype)
586 + and self.kv_cache_dtype != "fp8_ds_mla"
587 + and self.impl.supports_quant_query_input
588 + else dtype
589 + )
590 + self.dcp_manager = MLADCPManager(
591 + vllm_config=vllm_config,
592 + device=next(kv_b_proj.parameters()).device,
593 + num_heads=self.num_heads,
594 + query_head_dim=self.kv_lora_rank + self.qk_rope_head_dim,
595 + output_head_dim=self.kv_lora_rank,
596 + query_dtype=query_dtype,
597 + output_dtype=dtype,
598 + padded_num_heads=self.q_pad_num_heads,
599 + is_lse_base_on_e=self.impl.lse_base_on_e,
600 + use_pcp=self.use_pcp,
601 + )
602
603 self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
604
@@ -753,9 +767,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
767 return quant_output.fill_(0)
768 return output.fill_(0)
769
756 - if self.impl.dcp_world_size == -1:
757 - self.impl.dcp_world_size = get_dcp_group().world_size
758 -
770 fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)
771
772 num_actual_toks = attn_metadata.num_actual_tokens
@@ -914,6 +925,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
925 mqa_q = (mqa_ql_nope, mqa_q_pe)
926 # concatenate nope + pe -> (B, N, L + P) (fp8 op above may have fused)
927 if self.impl.dcp_world_size > 1:
928 + assert self.dcp_manager is not None
929 if self.use_pcp:
930 if self.impl.dcp_world_size > self.impl.pcp_world_size:
931 if isinstance(mqa_q, tuple):
@@ -924,8 +936,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
936 # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)
937 mqa_q = torch.cat(mqa_q, dim=-1)
938 if not qrep_decode:
927 - # mqa_q do allgather in head dim.
928 - mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)
939 + assert self.dcp_manager.query_gather is not None
940 + mqa_q = self.dcp_manager.query_gather(mqa_q)
941
942 # call decode attn
943 if not self.impl.is_sparse:
@@ -935,27 +947,23 @@ class MLAAttention(nn.Module, AttentionLayerBase):
947 # correct dcp attn_out with lse.
948 if self.impl.dcp_world_size > 1:
949 assert lse is not None
938 - if self.dcp_a2a:
939 - attn_out = dcp_a2a_lse_reduce(
940 - attn_out,
941 - lse,
942 - get_dcp_group(),
943 - is_lse_base_on_e=self.impl.lse_base_on_e,
944 - )
945 - elif self.use_pcp:
946 - attn_out = cp_lse_ag_out_ar(
947 - attn_out,
948 - lse,
949 - get_dcp_group(),
950 - is_lse_base_on_e=self.impl.lse_base_on_e,
951 - )
952 - else:
953 - attn_out = cp_lse_ag_out_rs(
954 - attn_out,
955 - lse,
956 - get_dcp_group(),
957 - is_lse_base_on_e=self.impl.lse_base_on_e,
958 - )
950 + assert self.dcp_manager is not None
951 + seq_lens = (
952 + attn_metadata.decode.seq_lens
953 + if attn_metadata.decode is not None
954 + else cast(torch.Tensor, attn_metadata.seq_lens)[ # type: ignore[attr-defined]
955 + : attn_metadata.num_decodes
956 + ]
957 + )
958 + query_start_loc = attn_metadata.query_start_loc[
959 + : attn_metadata.num_decodes + 1
960 + ]
961 + attn_out = self.dcp_manager.combine(
962 + attn_out,
963 + lse,
964 + seq_lens=seq_lens,
965 + query_start_loc=query_start_loc,
966 + )
967 if self.use_pcp:
968 attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)
969
@@ -1427,6 +1435,7 @@ class MLACommonPrefillMetadata:
1435 chunks: "list[MLACommonPrefillMetadata.ContextChunk]"
1436 context_lens_list: list[int]
1437 empty_token_slices: list[slice]
1438 + dcp_manager: MLADCPManager | None = None
1439
1440 block_table: torch.Tensor
1441 query_start_loc: torch.Tensor
@@ -1684,6 +1693,25 @@ def _flat_int32(values: list[int] | np.ndarray) -> torch.Tensor:
1693 return np_to_pinned_tensor(np.asarray(values, dtype=np.int32))
1694
1695
1696 +def align_mla_chunked_context_workspace_size(
1697 + vllm_config: VllmConfig,
1698 + workspace_size: int,
1699 +) -> int:
1700 + parallel_config = vllm_config.parallel_config
1701 + alignment = vllm_config.cache_config.block_size
1702 + if parallel_config.decode_context_parallel_size > 1:
1703 + alignment = lcm(
1704 + alignment,
1705 + parallel_config.decode_context_parallel_size
1706 + * parallel_config.cp_kv_cache_interleave_size,
1707 + )
1708 + workspace_size = max(
1709 + workspace_size,
1710 + vllm_config.scheduler_config.max_num_seqs * alignment,
1711 + )
1712 + return round_up(workspace_size, alignment)
1713 +
1714 +
1715 def build_mla_chunked_context_metadata(
1716 *,
1717 context_lens_cpu: torch.Tensor,
@@ -1696,6 +1724,7 @@ def build_mla_chunked_context_metadata(
1724 dcp_world_size: int,
1725 dcp_local_block_size: int,
1726 dcp_virtual_block_size: int,
1727 + dcp_manager: MLADCPManager | None = None,
1728 ) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None":
1729 """Build chunked-context metadata for an MLA prefill.
1730
@@ -1714,6 +1743,7 @@ def build_mla_chunked_context_metadata(
1743 dcp_world_size: Decode-context-parallel world size (1 if disabled).
1744 dcp_local_block_size: Per-rank interleave block size for DCP.
1745 dcp_virtual_block_size: ``dcp_local_block_size * dcp_world_size``.
1746 + dcp_manager: Shared MLA DCP collective manager.
1747
1748 Returns:
1749 The chunked-context metadata, or None when no prefill has any context.
@@ -1902,6 +1932,7 @@ def build_mla_chunked_context_metadata(
1932 chunks=chunks,
1933 context_lens_list=context_lens,
1934 empty_token_slices=empty_token_slices,
1935 + dcp_manager=dcp_manager,
1936 )
1937
1938
@@ -1954,7 +1985,10 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1985 64 * 1024,
1986 )
1987
1957 - return max(chunked_prefill_workspace_size, cache_config.block_size)
1988 + return align_mla_chunked_context_workspace_size(
1989 + vllm_config,
1990 + chunked_prefill_workspace_size,
1991 + )
1992
1993 @staticmethod
1994 def determine_prefill_query_data_type(
@@ -2033,6 +2067,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2067 self.q_data_type = self.determine_prefill_query_data_type(
2068 vllm_config, self.model_config.dtype
2069 )
2070 + attention_layer = self.compilation_config.static_forward_context[layer_names[0]]
2071
2072 try:
2073 self.dcp_world_size = get_dcp_group().world_size
@@ -2050,6 +2085,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2085 )
2086
2087 use_packed_fp8_cache = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
2088 + self.dcp_manager: MLADCPManager | None = None
2089 if self.dcp_world_size > 1:
2090 # Note(hc): The local kvcache is incomplete when DCP is triggered,
2091 # an additional kvcache allgather across the DCP group is therefore
@@ -2067,6 +2103,12 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2103 else self.model_config.dtype,
2104 device=device,
2105 )
2106 + self.dcp_manager = getattr(attention_layer, "dcp_manager", None)
2107 + assert isinstance(self.dcp_manager, MLADCPManager)
2108 + self.dcp_manager.init_kv_gather(
2109 + self.chunked_prefill_workspace,
2110 + self.chunked_prefill_workspace_size,
2111 + )
2112 else:
2113 self.chunked_prefill_workspace = torch.empty(
2114 (
@@ -2080,9 +2122,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2122 # Metadata builders are created per ubatch when DBO is enabled. MLA
2123 # prefill backends keep the prepared metadata on the backend object, so
2124 # each builder needs its own backend instance to avoid cross-ubatch races.
2083 - self._prefill_backend = self.compilation_config.static_forward_context[
2084 - layer_names[0]
2085 - ].prefill_backend.clone()
2125 + self._prefill_backend = attention_layer.prefill_backend.clone()
2126
2127 supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY
2128 self._init_reorder_batch_threshold(
@@ -2223,6 +2263,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2263 dcp_world_size=self.dcp_world_size,
2264 dcp_local_block_size=self.dcp_local_block_size,
2265 dcp_virtual_block_size=self.dcp_virtual_block_size,
2266 + dcp_manager=self.dcp_manager,
2267 )
2268
2269 prefill_metadata = MLACommonPrefillMetadata(
@@ -2660,9 +2701,8 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2701 ]
2702 assert toks * dcp_world_size <= cur_allgather_workspace.shape[0]
2703 cur_allgather_kvcache = cur_allgather_workspace[: toks * dcp_world_size]
2663 - cur_allgather_kvcache.copy_(
2664 - get_dcp_group().all_gather(local_gathered_kvcache, dim=0)
2665 - )
2704 + dcp_manager = cast(MLADCPManager, chunked_context.dcp_manager)
2705 + dcp_manager.kv_gather(cur_allgather_kvcache, local_gathered_kvcache)
2706 assert (
2707 cur_allgather_kvcache.shape[-1]
2708 == self.kv_lora_rank + self.qk_rope_head_dim
@@ -2732,8 +2772,6 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2772 output_scale: torch.Tensor | None = None,
2773 ) -> None:
2774 assert attn_metadata.prefill is not None
2735 - assert self.dcp_world_size != -1
2736 -
2775 prefill_metadata = attn_metadata.prefill
2776 assert prefill_metadata.prefill_backend is not None
2777 use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
@@ -2877,10 +2915,11 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
2915 and (self.qk_rope_head_dim == 64)
2916 )
2917
2880 - self.dcp_world_size: int = -1
2881 -
2918 + parallel_config = get_current_vllm_config().parallel_config
2919 + # Avoid requiring an initialized DCP group in tests.
2920 + self.dcp_world_size: int = parallel_config.decode_context_parallel_size
2921 self.cp_kv_cache_interleave_size: int = (
2883 - get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
2922 + parallel_config.cp_kv_cache_interleave_size
2923 )
2924
2925 @abstractmethod
vllm/model_executor/layers/attention/sparse_mla_attention.py
+22 -3
@@ -20,6 +20,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
20 MLACommonMetadata,
21 MLACommonPrefillMetadata,
22 accumulate_mla_context_chunk,
23 + align_mla_chunked_context_workspace_size,
24 build_mla_chunked_context_metadata,
25 get_mla_dims,
26 init_mla_context_partial,
@@ -31,6 +32,7 @@ from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
32 from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder
33 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
34 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
35 +from vllm.v1.attention.ops.dcp_utils import MLADCPManager
36 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
37
38 if TYPE_CHECKING:
@@ -162,9 +164,19 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
164 dtype=torch.int32,
165 device=device,
166 )
165 - layer_prefill_backend = vllm_config.compilation_config.static_forward_context[
167 + attention_layer = vllm_config.compilation_config.static_forward_context[
168 layer_names[0]
167 - ].prefill_backend
169 + ]
170 + layer_prefill_backend = attention_layer.prefill_backend
171 + self.dcp_manager: MLADCPManager | None = None
172 + if self.dcp_world_size > 1:
173 + self.dcp_manager = getattr(attention_layer, "dcp_manager", None)
174 + assert isinstance(self.dcp_manager, MLADCPManager)
175 + if layer_prefill_backend is not None:
176 + self.dcp_manager.init_kv_gather(
177 + self.chunked_prefill_workspace,
178 + self.chunked_prefill_workspace_size,
179 + )
180 self._prefill_backend = (
181 layer_prefill_backend.clone() if layer_prefill_backend is not None else None
182 )
@@ -184,7 +196,13 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
196 64 * 1024,
197 scheduler_config.max_num_seqs * topk_tokens,
198 )
187 - return max(workspace_size, cache_config.block_size)
199 + workspace_size = max(
200 + workspace_size,
201 + scheduler_config.max_num_seqs * cache_config.block_size,
202 + )
203 + if vllm_config.parallel_config.decode_context_parallel_size > 1:
204 + return align_mla_chunked_context_workspace_size(vllm_config, workspace_size)
205 + return workspace_size
206
207 def _build_req_id_per_token(
208 self,
@@ -232,6 +250,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
250 dcp_world_size=self.dcp_world_size,
251 dcp_local_block_size=self.dcp_local_block_size,
252 dcp_virtual_block_size=self.dcp_virtual_block_size,
253 + dcp_manager=self.dcp_manager,
254 )
255
256 def build(
vllm/models/kimi_k3/nvidia/mla.py
+69 -11
@@ -24,8 +24,8 @@ KV cache, and absorbs ``kv_b_proj`` into ``W_UK_T`` / ``W_UV`` -- mirroring the
24 K3 specifics: optional rotary embedding (disabled for the target model's NoPE
25 layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``).
26
27 -Out of scope (extension points, not wired here): context parallelism (DCP/PCP),
28 -sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths.
27 +Out of scope (extension points, not wired here): prefill context parallelism
28 +(PCP), sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths.
29 """
30
31 import math
@@ -35,7 +35,11 @@ import torch
35 from torch import nn
36
37 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
38 -from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
38 +from vllm.config import (
39 + CacheConfig,
40 + VllmConfig,
41 + get_current_vllm_config,
42 +)
43 from vllm.distributed import get_tensor_model_parallel_world_size
44 from vllm.forward_context import get_forward_context
45 from vllm.logger import init_logger
@@ -78,6 +82,7 @@ from vllm.v1.attention.backend import (
82 MLAAttentionImpl,
83 )
84 from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
85 +from vllm.v1.attention.ops.dcp_utils import MLADCPManager
86 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
87 from vllm.v1.attention.selector import get_attn_backend
88 from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec, get_kv_quant_mode
@@ -310,10 +315,36 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
315
316 vllm_config = get_current_vllm_config()
317 parallel_config = vllm_config.parallel_config
313 - assert (
314 - parallel_config.decode_context_parallel_size <= 1
315 - and parallel_config.prefill_context_parallel_size <= 1
316 - ), "Kimi-K3 MultiHeadLatentAttention does not support context parallelism."
318 + assert parallel_config.prefill_context_parallel_size == 1, (
319 + "Kimi-K3 MultiHeadLatentAttention does not support prefill context "
320 + "parallelism."
321 + )
322 + self.dcp_world_size = parallel_config.decode_context_parallel_size
323 + assert self.dcp_world_size <= 1 or self.rotary_emb is None, (
324 + "Kimi-K3 MultiHeadLatentAttention does not support RoPE with decode "
325 + "context parallelism because gathered queries require gathered "
326 + "positions."
327 + )
328 + self.dcp_manager: MLADCPManager | None = None
329 + if self.dcp_world_size > 1:
330 + query_dtype = (
331 + torch.float8_e4m3fn
332 + if is_quantized_kv_cache(self.kv_cache_dtype)
333 + and self.kv_cache_dtype != "fp8_ds_mla"
334 + else dtype
335 + )
336 + self.dcp_manager = MLADCPManager(
337 + vllm_config=vllm_config,
338 + device=next(self.kv_b_proj.parameters()).device,
339 + num_heads=self.num_local_heads,
340 + query_head_dim=self.head_size,
341 + output_head_dim=self.kv_lora_rank,
342 + query_dtype=query_dtype,
343 + output_dtype=dtype,
344 + padded_num_heads=self.q_pad_num_heads,
345 + is_lse_base_on_e=self.impl.lse_base_on_e,
346 + use_pcp=False,
347 + )
348 self.prefill_backend = get_mla_prefill_backend(vllm_config)(
349 num_heads=self.num_local_heads,
350 scale=self.scale,
@@ -584,9 +615,25 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
615 cos_sin_cache,
616 slot_mapping[:num_mqa_tokens],
617 )
587 - latent_out, _lse = self.impl.forward_mqa( # type: ignore[attr-defined]
618 + if self.dcp_world_size > 1:
619 + assert self.dcp_manager is not None
620 + assert self.dcp_manager.query_gather is not None
621 + mqa_q = self.dcp_manager.query_gather(mqa_q)
622 + latent_out, lse = self.impl.forward_mqa( # type: ignore[attr-defined]
623 mqa_q, self._attn_read_kv_cache(), attn_metadata, self
624 )
625 + if self.dcp_world_size > 1:
626 + assert lse is not None
627 + assert self.dcp_manager is not None
628 + assert attn_metadata.decode is not None
629 + latent_out = self.dcp_manager.combine(
630 + latent_out,
631 + lse,
632 + seq_lens=attn_metadata.decode.seq_lens,
633 + query_start_loc=attn_metadata.query_start_loc[
634 + : attn_metadata.num_decodes + 1
635 + ],
636 + )
637 self._v_up_proj(latent_out, out=attn_out[:num_mqa_tokens])
638
639 def _decode_concat_cache(
@@ -752,9 +799,20 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
799 )
800
801 if has_context:
755 - context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined]
756 - q, self._attn_read_kv_cache(), attn_metadata, self._k_scale
757 - )
802 + if self.dcp_world_size > 1:
803 + context_output, context_lse = (
804 + self.impl._context_parallel_compute_prefill_context( # type: ignore[attr-defined]
805 + q,
806 + self._attn_read_kv_cache(),
807 + attn_metadata,
808 + k_scale=self._k_scale,
809 + dcp_world_size=self.dcp_world_size,
810 + )
811 + )
812 + else:
813 + context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined]
814 + q, self._attn_read_kv_cache(), attn_metadata, self._k_scale
815 + )
816 suffix_output, suffix_lse = output_prefill
817 out = out.view(-1, self.num_local_heads, self.v_head_dim)
818 merge_attn_states(
vllm/v1/attention/backends/mla/flashinfer_mla.py
+4 -2
@@ -309,9 +309,11 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
309
310 return_lse = self.need_to_return_lse_for_decode
311 workspace_buffer = _get_workspace_buffer(return_lse)
312 + # Parallel gathers can change the runtime Q heads from TP-local num_heads.
313 + runtime_num_heads = q.shape[-2]
314 # trtllm-gen rejects MLA head counts it can't tile (e.g. 96);
315 # fall back to cute-dsl for those.
314 - decode_backend = _select_mla_decode_backend(self.num_heads)
316 + decode_backend = _select_mla_decode_backend(runtime_num_heads)
317 extra_kwargs = {}
318 if decode_backend:
319 extra_kwargs["backend"] = decode_backend
@@ -326,7 +328,7 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
328 if self._mla_counter_bytes is None:
329 self._mla_counter_bytes = get_trtllm_gen_multi_ctas_kv_counter_bytes(
330 self._mla_counter_max_batch,
329 - self.num_heads,
331 + runtime_num_heads,
332 get_device_sm_count(q.device),
333 )
334 extra_kwargs["multi_ctas_kv_counter_buffer"] = (
vllm/v1/attention/backends/mla/tokenspeed_mla.py
+2 -1
@@ -289,8 +289,9 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
289 self.output_scale = layer._k_scale_float
290
291 if self._workspace_buffer is None:
292 + # Parallelism can change the runtime query head count.
293 self._workspace_buffer = _get_workspace(
293 - q.device, self.num_heads, self.kv_lora_rank
294 + q.device, q.shape[-2], self.kv_lora_rank
295 )
296
297 # vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size).
vllm/v1/attention/ops/common.py
+51 -4
@@ -6,6 +6,34 @@ from vllm.distributed.parallel_state import GroupCoordinator
6 from vllm.triton_utils import tl, triton
7
8
9 +def mask_dcp_empty_shards_(
10 + lse: torch.Tensor,
11 + seq_lens: torch.Tensor | None,
12 + query_start_loc: torch.Tensor | None,
13 +) -> None:
14 + if seq_lens is None and query_start_loc is None:
15 + return
16 + if seq_lens is None or query_start_loc is None:
17 + raise ValueError("seq_lens and query_start_loc must be provided together")
18 + if (
19 + seq_lens.ndim != 1
20 + or query_start_loc.ndim != 1
21 + or query_start_loc.shape[0] != seq_lens.shape[0] + 1
22 + ):
23 + raise ValueError("query_start_loc must contain one boundary per sequence")
24 +
25 + row_indices = torch.arange(
26 + lse.shape[0], device=lse.device, dtype=query_start_loc.dtype
27 + )
28 + sequence_indices = torch.searchsorted(
29 + query_start_loc[1:], row_indices, right=True
30 + ).clamp_max(seq_lens.shape[0] - 1)
31 + empty_rows = (row_indices >= query_start_loc[-1]) | (
32 + seq_lens[sequence_indices] == 0
33 + )
34 + lse.masked_fill_(empty_rows[:, None], float("-inf"))
35 +
36 +
37 @triton.jit
38 def _correct_attn_cp_out_kernel(
39 outputs_ptr,
@@ -51,7 +79,7 @@ def _correct_attn_cp_out_kernel(
79 )
80
81 # calc final lse
54 - lse = tl.load(lses_ptr + lse_offsets)
82 + lse = tl.load(lses_ptr + lse_offsets).to(tl.float32)
83 lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse)
84 lse_max = tl.max(lse, axis=0)
85 lse_max = tl.where(lse_max == -float("inf"), 0, lse_max)
@@ -80,7 +108,7 @@ def _correct_attn_cp_out_kernel(
108 lse_offset = (
109 lse_idx * lses_stride_N + batch_idx * lses_stride_B + head_idx * lses_stride_H
110 )
83 - lse_tmp = tl.load(lses_ptr + lse_offset)
111 + lse_tmp = tl.load(lses_ptr + lse_offset).to(tl.float32)
112 lse_finally = lse_tmp - lse
113 lse_finally = tl.where(
114 (lse_finally != lse_finally) | (lse_finally == float("inf")),
@@ -185,6 +213,8 @@ def _cp_lse_common(
213 cp_group: GroupCoordinator,
214 ctx: CPTritonContext | None = None,
215 is_lse_base_on_e=True,
216 + seq_lens: torch.Tensor | None = None,
217 + query_start_loc: torch.Tensor | None = None,
218 ):
219 """
220 cp_attn_out: [ B, H, D ]
@@ -197,6 +227,7 @@ def _cp_lse_common(
227 ctx = CPTritonContext()
228
229 cp_attn_lse = cp_attn_lse.contiguous()
230 + mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc)
231 lses = cp_group.all_gather(cp_attn_lse, dim=0).reshape(
232 (cp_group.world_size,) + cp_attn_lse.shape
233 )
@@ -217,13 +248,21 @@ def cp_lse_ag_out_rs(
248 ctx: CPTritonContext | None = None,
249 return_lse: bool = False,
250 is_lse_base_on_e=True,
251 + seq_lens: torch.Tensor | None = None,
252 + query_start_loc: torch.Tensor | None = None,
253 ):
254 """
255 cp_attn_out: [ B, H, D ]
256 cp_attn_lse: [ B, H ]
257 """
258 out, lse = _cp_lse_common(
226 - cp_attn_out, cp_attn_lse, cp_group, ctx=ctx, is_lse_base_on_e=is_lse_base_on_e
259 + cp_attn_out,
260 + cp_attn_lse,
261 + cp_group,
262 + ctx=ctx,
263 + is_lse_base_on_e=is_lse_base_on_e,
264 + seq_lens=seq_lens,
265 + query_start_loc=query_start_loc,
266 )
267 out = cp_group.reduce_scatter(out, dim=1)
268
@@ -242,13 +281,21 @@ def cp_lse_ag_out_ar(
281 ctx: CPTritonContext | None = None,
282 return_lse: bool = False,
283 is_lse_base_on_e=True,
284 + seq_lens: torch.Tensor | None = None,
285 + query_start_loc: torch.Tensor | None = None,
286 ):
287 """
288 cp_attn_out: [ B, H, D ]
289 cp_attn_lse: [ B, H ]
290 """
291 out, lse = _cp_lse_common(
251 - cp_attn_out, cp_attn_lse, cp_group, ctx=ctx, is_lse_base_on_e=is_lse_base_on_e
292 + cp_attn_out,
293 + cp_attn_lse,
294 + cp_group,
295 + ctx=ctx,
296 + is_lse_base_on_e=is_lse_base_on_e,
297 + seq_lens=seq_lens,
298 + query_start_loc=query_start_loc,
299 )
300 out = cp_group.all_reduce(out)
301
vllm/v1/attention/ops/dcp_alltoall.py
+20 -11
@@ -26,6 +26,7 @@ import torch
26 import torch.distributed as dist
27
28 from vllm.triton_utils import tl, triton
29 +from vllm.v1.attention.ops.common import mask_dcp_empty_shards_
30
31 if TYPE_CHECKING:
32 from vllm.distributed.parallel_state import GroupCoordinator
@@ -87,7 +88,9 @@ def _lse_weighted_combine(
88 weights = weights / weight_sum.clamp(min=1e-10) # [N, B, H]
89
90 # Weighted combination: sum over N dimension
90 - result = (outputs * weights.unsqueeze(-1)).sum(dim=0) # [B, H, D]
91 + weights = weights.unsqueeze(-1)
92 + outputs = torch.where(weights == 0, torch.zeros_like(outputs), outputs)
93 + result = (outputs * weights).sum(dim=0) # [B, H, D]
94
95 if return_lse:
96 if is_lse_base_on_e:
@@ -172,7 +175,7 @@ def _dcp_a2a_pack_send_kernel(
175
176 lse_val = tl.load(
177 lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H
175 - )
178 + ).to(tl.float32)
179 if LSE_PACK_DIM == 1:
180 tl.store(
181 send_ptr + send_base + HEAD_DIM * send_stride_D,
@@ -301,10 +304,11 @@ def _dcp_a2a_unpack_combine_kernel(
304 else:
305 weight = tl.exp2(lse_val - global_lse)
306 weight = tl.where(weight != weight, 0.0, weight)
304 - acc += (
305 - tl.load(recv_ptr + recv_base + d_offsets * recv_stride_D).to(tl.float32)
306 - * weight
307 + partial = tl.load(recv_ptr + recv_base + d_offsets * recv_stride_D).to(
308 + tl.float32
309 )
310 + partial = tl.where(weight == 0.0, 0.0, partial)
311 + acc += partial * weight
312
313 final_offsets = (
314 batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D
@@ -324,7 +328,10 @@ def _dcp_a2a_pack_send(
328 h_per_rank: int,
329 head_dim: int,
330 lse_pack_dim: int,
331 + seq_lens: torch.Tensor | None = None,
332 + query_start_loc: torch.Tensor | None = None,
333 ) -> None:
334 + mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc)
335 grid = (cp_attn_out.shape[0], h_per_rank, 1)
336 _dcp_a2a_pack_send_kernel[grid](
337 cp_attn_out,
@@ -396,20 +403,24 @@ def dcp_a2a_lse_reduce(
403 ctx: CPTritonContext | None = None,
404 return_lse: bool = False,
405 is_lse_base_on_e: bool = True,
406 + seq_lens: torch.Tensor | None = None,
407 + query_start_loc: torch.Tensor | None = None,
408 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
409 """
410 Combine partial attention outputs across DCP ranks using All-to-All.
411
403 - The output and fp32 LSE are packed into a single output-dtype buffer, sent
412 + The output and LSE are packed into a single output-dtype buffer, sent
413 with one All-to-All, then unpacked and combined with exact LSE weighting.
414
415 Args:
416 cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim
408 - cp_attn_lse: [B, H] log-sum-exp values (fp32)
417 + cp_attn_lse: [B, H] floating-point log-sum-exp values
418 cp_group: GroupCoordinator for DCP communication
419 ctx: CPTritonContext (unused, for signature compatibility)
420 return_lse: If True, also return the combined global LSE
421 is_lse_base_on_e: If True, LSE is base e; if False, base 2
422 + seq_lens: Local KV lengths. Empty shards contribute zero weight.
423 + query_start_loc: Cumulative query-token offsets for each request.
424
425 Returns:
426 Combined output [B, H/N, D] (head-scattered)
@@ -426,10 +437,6 @@ def dcp_a2a_lse_reduce(
437 if H % world_size != 0:
438 raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.")
439 H_per_rank = H // world_size
429 - # The pack kernel bit-casts the LSE as fp32; some MLA backends return it in
430 - # the activation dtype (bf16/fp16), so enforce the documented fp32 contract.
431 - if cp_attn_lse.dtype != torch.float32:
432 - cp_attn_lse = cp_attn_lse.to(torch.float32)
440 lse_pack_dim = _dcp_a2a_lse_pack_dim(cp_attn_out.dtype)
441
442 send_buffer, recv_buffer = _dcp_a2a_send_recv_buffers(
@@ -446,6 +453,8 @@ def dcp_a2a_lse_reduce(
453 H_per_rank,
454 D,
455 lse_pack_dim,
456 + seq_lens=seq_lens,
457 + query_start_loc=query_start_loc,
458 )
459
460 work = dist.all_to_all_single(
vllm/v1/attention/ops/dcp_utils.py new
+740
@@ -0,0 +1,740 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MLA DCP collective selection and direct symmetric-memory implementations."""
4 +
5 +from __future__ import annotations
6 +
7 +import functools
8 +from collections.abc import Callable
9 +from typing import TYPE_CHECKING, Any, Protocol
10 +
11 +import torch
12 +
13 +import vllm.envs as envs
14 +from vllm.config import VllmConfig
15 +from vllm.distributed import get_dcp_group
16 +from vllm.distributed.parallel_state import in_the_same_node_as
17 +from vllm.logger import init_logger
18 +from vllm.platforms import current_platform
19 +from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs
20 +from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
21 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id
22 +
23 +logger = init_logger(__name__)
24 +
25 +if TYPE_CHECKING:
26 + from torch.distributed import ProcessGroup
27 +
28 + from vllm.distributed.parallel_state import GroupCoordinator
29 +
30 +try:
31 + import torch.distributed._symmetric_memory as symm_mem
32 +
33 + symm_mem_available = True
34 +except ImportError:
35 + symm_mem = None # type: ignore[assignment]
36 + symm_mem_available = False
37 +
38 +
39 +@functools.cache
40 +def _symm_mem_spans_group(group: GroupCoordinator) -> bool:
41 + """Probe whether the group has NVLS symmetric memory."""
42 + if not symm_mem_available:
43 + return False
44 + try:
45 + from torch._C._autograd import DeviceType
46 + from torch._C._distributed_c10d import _SymmetricMemory
47 +
48 + device = torch.device("cuda", torch.accelerator.current_device_index())
49 + if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device.index):
50 + return False
51 + probe = symm_mem.empty(8, dtype=torch.uint8, device=device)
52 + probe.zero_()
53 + torch.accelerator.synchronize()
54 + handle = symm_mem.rendezvous(probe, group.device_group.group_name)
55 + spans = handle is not None and handle.multicast_ptr != 0
56 + except Exception as error:
57 + logger.debug("Direct DCP symmetric-memory probe failed: %s", error)
58 + return False
59 + logger.debug_once(
60 + "Direct DCP symmetric memory across %d ranks: %s",
61 + group.world_size,
62 + "available" if spans else "unavailable",
63 + )
64 + return spans
65 +
66 +
67 +def _direct_dcp_enabled(
68 + group: GroupCoordinator,
69 + dtype: torch.dtype,
70 + use_direct: bool | None,
71 + supported_dtypes: tuple[torch.dtype, ...] | None = None,
72 +) -> bool:
73 + if use_direct is not None:
74 + return use_direct
75 + return (
76 + symm_mem_available
77 + and current_platform.is_cuda()
78 + and (supported_dtypes is None or dtype in supported_dtypes)
79 + and (
80 + all(in_the_same_node_as(group.cpu_group, source_rank=0))
81 + or _symm_mem_spans_group(group)
82 + )
83 + )
84 +
85 +
86 +def _direct_dcp_multicast_enabled(
87 + group: GroupCoordinator,
88 + dtype: torch.dtype,
89 + use_direct: bool | None,
90 + supported_dtypes: tuple[torch.dtype, ...] | None = None,
91 +) -> bool:
92 + return _direct_dcp_enabled(
93 + group, dtype, use_direct, supported_dtypes
94 + ) and _symm_mem_spans_group(group)
95 +
96 +
97 +def get_dcp_workspace_max_num_tokens(vllm_config: VllmConfig) -> int:
98 + scheduler_config = vllm_config.scheduler_config
99 + speculative_config = vllm_config.speculative_config
100 + speculative_tokens = vllm_config.num_speculative_tokens
101 + tokens_per_seq = (
102 + 1
103 + + (
104 + 2
105 + if speculative_config is not None and speculative_config.parallel_drafting
106 + else 1
107 + )
108 + * speculative_tokens
109 + )
110 + return min(
111 + scheduler_config.max_num_batched_tokens,
112 + max(
113 + scheduler_config.max_num_seqs * tokens_per_seq,
114 + vllm_config.compilation_config.max_cudagraph_capture_size or 0,
115 + ),
116 + )
117 +
118 +
119 +class _DirectDCPWorkspace:
120 + def __init__(
121 + self,
122 + group: ProcessGroup,
123 + device: torch.device,
124 + num_ubatches: int,
125 + ) -> None:
126 + self.group = group
127 + self.world_size = group.size()
128 + self.rank = group.rank()
129 + self.device = torch.device(device)
130 + self.num_ubatches = num_ubatches
131 + self.epoch = torch.zeros(num_ubatches, dtype=torch.int64, device=self.device)
132 + self._allocations: list[tuple[torch.Tensor, Any, list[torch.Tensor]]] = []
133 +
134 + def _allocate(
135 + self, shape: tuple[int, ...], dtype: torch.dtype
136 + ) -> tuple[torch.Tensor, torch.Tensor]:
137 + storage = symm_mem.empty(shape, device=self.device, dtype=dtype)
138 + storage.zero_()
139 + torch.accelerator.synchronize()
140 + handle = symm_mem.rendezvous(storage, self.group.group_name)
141 + assert handle is not None, "DCP symmetric memory rendezvous returned None"
142 + handle.barrier()
143 + views = [
144 + handle.get_buffer(peer, list(shape), dtype, 0)
145 + for peer in range(self.world_size)
146 + ]
147 + self.device = storage.device
148 + peer_ptrs = torch.tensor(
149 + [
150 + [view[ubatch].data_ptr() for view in views]
151 + for ubatch in range(self.num_ubatches)
152 + ],
153 + dtype=torch.int64,
154 + device=self.device,
155 + )
156 + self._allocations.append((storage, handle, views))
157 + return storage, peer_ptrs
158 +
159 + def _multicast_ptrs(self, storage: torch.Tensor) -> list[int]:
160 + disabled = [0] * self.num_ubatches
161 + for allocated, handle, _ in self._allocations:
162 + if allocated is storage:
163 + break
164 + else:
165 + return disabled
166 + try:
167 + from torch._C._autograd import DeviceType
168 + from torch._C._distributed_c10d import _SymmetricMemory
169 +
170 + if not _SymmetricMemory.has_multicast_support(
171 + DeviceType.CUDA, storage.device.index
172 + ):
173 + return disabled
174 + multicast_base = handle.multicast_ptr
175 + except Exception:
176 + return disabled
177 + if not multicast_base:
178 + return disabled
179 + storage_base = storage.data_ptr()
180 + return [
181 + multicast_base + (storage[ubatch].data_ptr() - storage_base)
182 + for ubatch in range(self.num_ubatches)
183 + ]
184 +
185 +
186 +def reserve_query_head_storage(
187 + query: torch.Tensor, padded_num_heads: int
188 +) -> torch.Tensor:
189 + """Reserve backing storage for fixed-head decode kernels."""
190 + assert query.ndim == 3
191 + assert query.shape[1] <= padded_num_heads
192 + padded = query.new_empty((query.shape[0], padded_num_heads, query.shape[2]))
193 + padded.resize_(query.shape)
194 + padded.copy_(query)
195 + return padded
196 +
197 +
198 +_A2A_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
199 +
200 +
201 +class DirectDCPA2AWorkspace(_DirectDCPWorkspace):
202 + """Persistent symmetric buffers for direct DCP output exchange."""
203 +
204 + def __init__(
205 + self,
206 + group: ProcessGroup,
207 + device: torch.device,
208 + max_num_tokens: int,
209 + heads_per_rank: int,
210 + head_dim: int,
211 + dtype: torch.dtype = torch.bfloat16,
212 + num_ubatches: int = 1,
213 + ) -> None:
214 + if dtype not in _A2A_SUPPORTED_DTYPES:
215 + raise ValueError(f"Direct DCP A2A does not support {dtype}")
216 + if num_ubatches < 1:
217 + raise ValueError(
218 + f"Direct DCP A2A requires at least one ubatch slot, got {num_ubatches}"
219 + )
220 + super().__init__(group, device, num_ubatches)
221 + self.max_num_tokens = max_num_tokens
222 + self.heads_per_rank = heads_per_rank
223 + self.head_dim = head_dim
224 +
225 + output_shape = (
226 + num_ubatches,
227 + 2,
228 + self.world_size,
229 + max_num_tokens,
230 + heads_per_rank,
231 + head_dim,
232 + )
233 + lse_shape = (
234 + num_ubatches,
235 + 2,
236 + self.world_size,
237 + max_num_tokens,
238 + heads_per_rank,
239 + )
240 + signal_shape = (num_ubatches, 2, self.world_size)
241 + self.received_output, self.peer_output_ptrs = self._allocate(
242 + output_shape, dtype
243 + )
244 + self.received_lse, self.peer_lse_ptrs = self._allocate(lse_shape, torch.float32)
245 + self.received_signal, self.peer_signal_ptrs = self._allocate(
246 + signal_shape, torch.int32
247 + )
248 +
249 + def lse_reduce(
250 + self,
251 + partial_output: torch.Tensor,
252 + partial_lse: torch.Tensor,
253 + is_lse_base_on_e: bool,
254 + seq_lens: torch.Tensor | None = None,
255 + query_start_loc: torch.Tensor | None = None,
256 + ) -> torch.Tensor:
257 + ubatch = dbo_current_ubatch_id()
258 + num_tokens = partial_output.shape[0]
259 + output = partial_output.new_empty(
260 + (num_tokens, self.heads_per_rank, self.head_dim)
261 + )
262 + torch.ops._C.direct_dcp_a2a_lse_reduce(
263 + partial_output,
264 + partial_lse,
265 + seq_lens,
266 + query_start_loc,
267 + self.peer_output_ptrs[ubatch],
268 + self.peer_lse_ptrs[ubatch],
269 + self.peer_signal_ptrs[ubatch],
270 + self.received_output[ubatch],
271 + self.received_lse[ubatch],
272 + self.received_signal[ubatch],
273 + self.epoch[ubatch : ubatch + 1],
274 + output,
275 + self.world_size,
276 + self.rank,
277 + self.max_num_tokens,
278 + is_lse_base_on_e,
279 + )
280 + return output
281 +
282 +
283 +@functools.cache
284 +def get_direct_dcp_a2a_workspace(
285 + group: GroupCoordinator,
286 + device: torch.device,
287 + max_num_tokens: int,
288 + heads_per_rank: int,
289 + head_dim: int,
290 + dtype: torch.dtype,
291 + num_ubatches: int,
292 +) -> DirectDCPA2AWorkspace | None:
293 + if not _direct_dcp_enabled(
294 + group, dtype, envs.VLLM_USE_DIRECT_DCP_A2A, _A2A_SUPPORTED_DTYPES
295 + ):
296 + return None
297 + return DirectDCPA2AWorkspace(
298 + group.device_group,
299 + device,
300 + max_num_tokens,
301 + heads_per_rank,
302 + head_dim,
303 + dtype,
304 + num_ubatches,
305 + )
306 +
307 +
308 +def _q_gather_layout_supported(
309 + world_size: int,
310 + heads_per_rank: int,
311 + head_dim: int,
312 + dtype: torch.dtype,
313 + padded_num_heads: int | None,
314 +) -> bool:
315 + element_size = torch.empty((), dtype=dtype).element_size()
316 + gathered_num_heads = world_size * heads_per_rank
317 + storage_num_heads = (
318 + gathered_num_heads if padded_num_heads is None else padded_num_heads
319 + )
320 + return (
321 + heads_per_rank * head_dim * element_size % 16 == 0
322 + and storage_num_heads * head_dim * element_size % 16 == 0
323 + )
324 +
325 +
326 +class DirectDCPQGatherWorkspace(_DirectDCPWorkspace):
327 + """Publish query shards directly into the consumer-final symmetric buffer.
328 +
329 + The final buffer is reusable after the downstream DCP output combine. That
330 + combine orders all ranks after attention has consumed the gathered query.
331 + """
332 +
333 + def __init__(
334 + self,
335 + group: ProcessGroup,
336 + device: torch.device,
337 + max_num_tokens: int,
338 + heads_per_rank: int,
339 + head_dim: int,
340 + dtype: torch.dtype = torch.bfloat16,
341 + num_ubatches: int = 1,
342 + padded_num_heads: int | None = None,
343 + ) -> None:
344 + if num_ubatches < 1:
345 + raise ValueError(
346 + "Direct DCP q-gather requires at least one ubatch slot, "
347 + f"got {num_ubatches}"
348 + )
349 + if max_num_tokens < 1 or heads_per_rank < 1 or head_dim < 1:
350 + raise ValueError(
351 + "Direct DCP q-gather dimensions must be positive, got "
352 + f"T={max_num_tokens}, H={heads_per_rank}, D={head_dim}"
353 + )
354 + gathered_num_heads = group.size() * heads_per_rank
355 + if not _q_gather_layout_supported(
356 + group.size(), heads_per_rank, head_dim, dtype, padded_num_heads
357 + ):
358 + raise ValueError("Direct DCP q-gather requires 16-byte-aligned query rows.")
359 + super().__init__(group, device, num_ubatches)
360 + if self.world_size <= 1:
361 + raise ValueError("Direct DCP q-gather requires at least two ranks")
362 + self.max_num_tokens = max_num_tokens
363 + self.heads_per_rank = heads_per_rank
364 + self.gathered_num_heads = gathered_num_heads
365 + self.padded_num_heads = (
366 + self.gathered_num_heads if padded_num_heads is None else padded_num_heads
367 + )
368 + if self.padded_num_heads < self.gathered_num_heads:
369 + raise ValueError(
370 + "Direct DCP q-gather padded heads must cover gathered heads: "
371 + f"{self.padded_num_heads} < {self.gathered_num_heads}"
372 + )
373 + self.head_dim = head_dim
374 +
375 + query_shape = (
376 + num_ubatches,
377 + max_num_tokens,
378 + self.padded_num_heads,
379 + head_dim,
380 + )
381 + signal_shape = (num_ubatches, 2, self.world_size)
382 + self.final_query, _ = self._allocate(query_shape, dtype)
383 + self.received_signal, _ = self._allocate(signal_shape, torch.int32)
384 + query_multicast_ptrs = self._multicast_ptrs(self.final_query)
385 + signal_multicast_ptrs = self._multicast_ptrs(self.received_signal)
386 + self.multicast_ptrs = list(
387 + zip(query_multicast_ptrs, signal_multicast_ptrs, strict=True)
388 + )
389 + if not all(
390 + query_ptr and signal_ptr for query_ptr, signal_ptr in self.multicast_ptrs
391 + ):
392 + raise RuntimeError(
393 + "Direct DCP q-gather requires NVLS symmetric-memory multicast."
394 + )
395 + self.completion = self.received_signal.new_zeros((num_ubatches, 1))
396 + torch.accelerator.synchronize()
397 +
398 + def gather(self, local_query: torch.Tensor) -> torch.Tensor:
399 + ubatch = dbo_current_ubatch_id()
400 + if not 0 <= ubatch < self.num_ubatches:
401 + raise ValueError(
402 + f"DCP q-gather ubatch {ubatch} exceeds {self.num_ubatches} slots"
403 + )
404 + if local_query.ndim == 3 and local_query.shape[1] != self.heads_per_rank:
405 + raise ValueError(
406 + f"DCP q-gather expected {self.heads_per_rank} local query heads, "
407 + f"got {local_query.shape[1]}"
408 + )
409 +
410 + num_tokens = local_query.shape[0]
411 + output = torch.as_strided(
412 + self.final_query[ubatch],
413 + size=(num_tokens, self.gathered_num_heads, self.head_dim),
414 + stride=(
415 + self.gathered_num_heads * self.head_dim,
416 + self.head_dim,
417 + 1,
418 + ),
419 + )
420 + query_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch]
421 + torch.ops._C.direct_dcp_q_gather(
422 + local_query,
423 + output,
424 + self.received_signal[ubatch],
425 + self.completion[ubatch],
426 + self.epoch[ubatch : ubatch + 1],
427 + self.world_size,
428 + self.rank,
429 + self.max_num_tokens,
430 + self.padded_num_heads,
431 + query_multicast_ptr,
432 + signal_multicast_ptr,
433 + )
434 + return output
435 +
436 +
437 +@functools.cache
438 +def get_direct_dcp_q_gather_workspace(
439 + group: GroupCoordinator,
440 + device: torch.device,
441 + max_num_tokens: int,
442 + heads_per_rank: int,
443 + head_dim: int,
444 + dtype: torch.dtype,
445 + num_ubatches: int,
446 + padded_num_heads: int | None = None,
447 +) -> DirectDCPQGatherWorkspace | None:
448 + if not _direct_dcp_multicast_enabled(
449 + group, dtype, envs.VLLM_USE_DIRECT_DCP_Q_GATHER
450 + ):
451 + return None
452 + if not _q_gather_layout_supported(
453 + group.world_size, heads_per_rank, head_dim, dtype, padded_num_heads
454 + ):
455 + return None
456 + return DirectDCPQGatherWorkspace(
457 + group.device_group,
458 + device,
459 + max_num_tokens,
460 + heads_per_rank,
461 + head_dim,
462 + dtype,
463 + num_ubatches,
464 + padded_num_heads,
465 + )
466 +
467 +
468 +_KV_GATHER_SUPPORTED_DTYPES = (
469 + torch.float16,
470 + torch.bfloat16,
471 + torch.float8_e4m3fn,
472 +)
473 +
474 +
475 +def _kv_gather_layout_supported(token_dim: int, dtype: torch.dtype) -> bool:
476 + return token_dim * torch.empty((), dtype=dtype).element_size() % 16 == 0
477 +
478 +
479 +class DirectDCPKVGatherWorkspace(_DirectDCPWorkspace):
480 + """Persistent symmetric buffers for direct DCP KV gather."""
481 +
482 + def __init__(
483 + self,
484 + group: ProcessGroup,
485 + device: torch.device,
486 + max_gathered_tokens: int,
487 + token_dim: int,
488 + dtype: torch.dtype = torch.bfloat16,
489 + num_ubatches: int = 1,
490 + ) -> None:
491 + if dtype not in _KV_GATHER_SUPPORTED_DTYPES:
492 + raise ValueError(f"Direct DCP kv-gather does not support {dtype}")
493 + if num_ubatches < 1:
494 + raise ValueError(
495 + "Direct DCP kv-gather requires at least one ubatch slot, "
496 + f"got {num_ubatches}"
497 + )
498 + if max_gathered_tokens < 1 or token_dim < 1:
499 + raise ValueError(
500 + "Direct DCP kv-gather dimensions must be positive, got "
501 + f"T={max_gathered_tokens}, D={token_dim}"
502 + )
503 + if not _kv_gather_layout_supported(token_dim, dtype):
504 + raise ValueError("Direct DCP kv-gather requires 16-byte-aligned KV rows.")
505 + super().__init__(group, device, num_ubatches)
506 + if self.world_size <= 1:
507 + raise ValueError("Direct DCP kv-gather requires at least two ranks")
508 + if max_gathered_tokens % self.world_size != 0:
509 + raise ValueError(
510 + "Direct DCP kv-gather capacity must divide evenly across "
511 + f"ranks: {max_gathered_tokens} % {self.world_size} != 0"
512 + )
513 + self.max_gathered_tokens = max_gathered_tokens
514 +
515 + kv_shape = (num_ubatches, 2, max_gathered_tokens, token_dim)
516 + signal_shape = (num_ubatches, 2, self.world_size)
517 + self.received_kv, _ = self._allocate(kv_shape, dtype)
518 + self.received_signal, _ = self._allocate(signal_shape, torch.int32)
519 + kv_multicast_ptrs = self._multicast_ptrs(self.received_kv)
520 + signal_multicast_ptrs = self._multicast_ptrs(self.received_signal)
521 + self.multicast_ptrs = list(
522 + zip(kv_multicast_ptrs, signal_multicast_ptrs, strict=True)
523 + )
524 + if not all(kv_ptr and signal_ptr for kv_ptr, signal_ptr in self.multicast_ptrs):
525 + raise RuntimeError(
526 + "Direct DCP kv-gather requires NVLS symmetric-memory multicast."
527 + )
528 + self.completion = self.received_signal.new_zeros((num_ubatches, 2))
529 + torch.accelerator.synchronize()
530 +
531 + def gather(self, gathered_kv: torch.Tensor, local_kv: torch.Tensor) -> None:
532 + ubatch = dbo_current_ubatch_id()
533 + if not 0 <= ubatch < self.num_ubatches:
534 + raise ValueError(
535 + f"DCP kv-gather ubatch {ubatch} exceeds {self.num_ubatches} slots"
536 + )
537 + kv_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch]
538 + torch.ops._C.direct_dcp_kv_gather(
539 + local_kv,
540 + self.received_kv[ubatch],
541 + self.received_signal[ubatch],
542 + self.completion[ubatch],
543 + self.epoch[ubatch : ubatch + 1],
544 + gathered_kv,
545 + self.world_size,
546 + self.rank,
547 + self.max_gathered_tokens,
548 + kv_multicast_ptr,
549 + signal_multicast_ptr,
550 + )
551 +
552 +
553 +@functools.cache
554 +def get_direct_dcp_kv_gather_workspace(
555 + group: GroupCoordinator,
556 + device: torch.device,
557 + max_gathered_tokens: int,
558 + token_dim: int,
559 + dtype: torch.dtype,
560 + num_ubatches: int,
561 +) -> DirectDCPKVGatherWorkspace | None:
562 + if not _direct_dcp_multicast_enabled(
563 + group,
564 + dtype,
565 + envs.VLLM_USE_DIRECT_DCP_KV_GATHER,
566 + _KV_GATHER_SUPPORTED_DTYPES,
567 + ):
568 + return None
569 + if not _kv_gather_layout_supported(token_dim, dtype):
570 + return None
571 + return DirectDCPKVGatherWorkspace(
572 + group.device_group,
573 + device,
574 + max_gathered_tokens,
575 + token_dim,
576 + dtype,
577 + num_ubatches,
578 + )
579 +
580 +
581 +class DCPCombine(Protocol):
582 + def __call__(
583 + self,
584 + partial_output: torch.Tensor,
585 + partial_lse: torch.Tensor,
586 + *,
587 + seq_lens: torch.Tensor,
588 + query_start_loc: torch.Tensor,
589 + ) -> torch.Tensor: ...
590 +
591 +
592 +class MLADCPManager:
593 + """Select and own layer-level collective implementations for MLA DCP."""
594 +
595 + _kv_gather: Callable[[torch.Tensor, torch.Tensor], object]
596 +
597 + def __init__(
598 + self,
599 + vllm_config: VllmConfig,
600 + device: torch.device,
601 + num_heads: int,
602 + query_head_dim: int,
603 + output_head_dim: int,
604 + query_dtype: torch.dtype,
605 + output_dtype: torch.dtype,
606 + padded_num_heads: int | None,
607 + is_lse_base_on_e: bool,
608 + use_pcp: bool,
609 + ) -> None:
610 + parallel_config = vllm_config.parallel_config
611 + self.group = get_dcp_group()
612 + self.device = torch.device(device)
613 + self.num_ubatches = max(parallel_config.num_ubatches, 1)
614 + self.max_num_tokens = get_dcp_workspace_max_num_tokens(vllm_config)
615 + self.use_a2a = parallel_config.dcp_comm_backend == "a2a"
616 + self.padded_num_heads = padded_num_heads
617 +
618 + self.combine = self._init_combine(
619 + num_heads,
620 + output_head_dim,
621 + output_dtype,
622 + is_lse_base_on_e,
623 + use_pcp,
624 + )
625 + self.query_gather = (
626 + None
627 + if use_pcp
628 + else self._init_query_gather(
629 + num_heads,
630 + query_head_dim,
631 + query_dtype,
632 + )
633 + )
634 +
635 + def _init_combine(
636 + self,
637 + num_heads: int,
638 + head_dim: int,
639 + dtype: torch.dtype,
640 + is_lse_base_on_e: bool,
641 + use_pcp: bool,
642 + ) -> DCPCombine:
643 + direct_workspace = None
644 + if self.use_a2a:
645 + direct_workspace = get_direct_dcp_a2a_workspace(
646 + self.group,
647 + self.device,
648 + self.max_num_tokens,
649 + num_heads,
650 + head_dim,
651 + dtype,
652 + self.num_ubatches,
653 + )
654 + if direct_workspace is not None:
655 + logger.info_once("Using direct symmetric-memory DCP A2A for MLA.")
656 + return functools.partial(
657 + direct_workspace.lse_reduce,
658 + is_lse_base_on_e=is_lse_base_on_e,
659 + )
660 +
661 + combine_fn = (
662 + dcp_a2a_lse_reduce
663 + if self.use_a2a
664 + else cp_lse_ag_out_ar
665 + if use_pcp
666 + else cp_lse_ag_out_rs
667 + )
668 + return functools.partial(
669 + combine_fn,
670 + cp_group=self.group,
671 + is_lse_base_on_e=is_lse_base_on_e,
672 + )
673 +
674 + def _init_query_gather(
675 + self,
676 + num_heads: int,
677 + head_dim: int,
678 + dtype: torch.dtype,
679 + ) -> Callable[[torch.Tensor], torch.Tensor]:
680 + direct_workspace = get_direct_dcp_q_gather_workspace(
681 + self.group,
682 + self.device,
683 + self.max_num_tokens,
684 + num_heads,
685 + head_dim,
686 + dtype,
687 + self.num_ubatches,
688 + self.padded_num_heads,
689 + )
690 + if direct_workspace is not None:
691 + logger.info_once("Using direct symmetric-memory DCP query gather for MLA.")
692 + return direct_workspace.gather
693 + return self._gather_query
694 +
695 + def _gather_query(self, query: torch.Tensor) -> torch.Tensor:
696 + query = self.group.all_gather(query, dim=1)
697 + if self.padded_num_heads is not None:
698 + query = reserve_query_head_storage(query, self.padded_num_heads)
699 + return query
700 +
701 + def init_kv_gather(
702 + self,
703 + workspace: torch.Tensor,
704 + max_gathered_tokens: int,
705 + ) -> None:
706 + world_size = self.group.world_size
707 + assert max_gathered_tokens > 0
708 + assert max_gathered_tokens % world_size == 0
709 + assert workspace.ndim == 2
710 + assert workspace.is_contiguous()
711 + assert workspace.shape[0] == (
712 + max_gathered_tokens + max_gathered_tokens // world_size
713 + )
714 + assert workspace.shape[1] > 0
715 +
716 + direct_workspace = get_direct_dcp_kv_gather_workspace(
717 + self.group,
718 + workspace.device,
719 + max_gathered_tokens,
720 + workspace.shape[1],
721 + workspace.dtype,
722 + self.num_ubatches,
723 + )
724 + if direct_workspace is not None:
725 + logger.info_once(
726 + "Using direct symmetric-memory DCP chunked-context KV gather for MLA."
727 + )
728 + self._kv_gather = direct_workspace.gather
729 + else:
730 + self._kv_gather = functools.partial(
731 + torch.distributed.all_gather_into_tensor,
732 + group=self.group.device_group,
733 + )
734 +
735 + def kv_gather(
736 + self,
737 + gathered_kv: torch.Tensor,
738 + local_kv: torch.Tensor,
739 + ) -> object:
740 + return self._kv_gather(gathered_kv, local_kv)