@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,652 lines 53.6 KB
Raw
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 dataclasses import dataclass
10 from typing import TYPE_CHECKING, Any, Protocol
11
12 import torch
13 import torch.distributed as dist
14
15 import vllm.envs as envs
16 from vllm.config import VllmConfig
17 from vllm.distributed import get_dcp_group
18 from vllm.logger import init_logger
19 from vllm.model_executor.warmup.jit_warmup import (
20 WarmupIntRange,
21 kernel_launcher,
22 )
23 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
24 LaunchSpec,
25 TritonWarmupTensor,
26 VllmTritonJitKernel,
27 triton_scalar_specialization_rep,
28 )
29 from vllm.triton_utils import tl, triton
30 from vllm.v1.attention.backend import max_decode_query_len
31 from vllm.v1.attention.ops.cp_common import (
32 DirectCPWorkspace,
33 direct_cp_enabled,
34 direct_cp_multicast_enabled,
35 )
36 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
37
38 logger = init_logger(__name__)
39
40 if TYPE_CHECKING:
41 from torch.distributed import ProcessGroup
42
43 from vllm.distributed.parallel_state import GroupCoordinator
44
45
46 # LSE/output combine
47
48
49 def _validate_dcp_empty_shard_args(
50 seq_lens: torch.Tensor | None,
51 query_start_loc: torch.Tensor | None,
52 ) -> bool:
53 """Validate the empty-shard mask inputs.
54
55 Returns True when masking is requested, False when it is disabled. Raises on
56 an inconsistent pair. Performs no device work so it is safe to call from a
57 CUDA-graph-captured region.
58 """
59 if seq_lens is None and query_start_loc is None:
60 return False
61 if seq_lens is None or query_start_loc is None:
62 raise ValueError("seq_lens and query_start_loc must be provided together")
63 if (
64 seq_lens.ndim != 1
65 or query_start_loc.ndim != 1
66 or query_start_loc.shape[0] != seq_lens.shape[0] + 1
67 ):
68 raise ValueError("query_start_loc must contain one boundary per sequence")
69 return True
70
71
72 def mask_dcp_empty_shards_(
73 lse: torch.Tensor,
74 seq_lens: torch.Tensor | None,
75 query_start_loc: torch.Tensor | None,
76 ) -> None:
77 if not _validate_dcp_empty_shard_args(seq_lens, query_start_loc):
78 return
79 assert seq_lens is not None and query_start_loc is not None
80
81 # A DCP rank can receive no local sequences during CUDA graph warmup even
82 # though the padded LSE buffer still has rows. In that case every row is an
83 # empty shard; avoid indexing the empty seq_lens tensor below.
84 if seq_lens.shape[0] == 0:
85 lse.fill_(float("-inf"))
86 return
87
88 row_indices = torch.arange(
89 lse.shape[0], device=lse.device, dtype=query_start_loc.dtype
90 )
91 sequence_indices = torch.searchsorted(
92 query_start_loc[1:], row_indices, right=True
93 ).clamp_max(seq_lens.shape[0] - 1)
94 empty_rows = (row_indices >= query_start_loc[-1]) | (
95 seq_lens[sequence_indices] == 0
96 )
97 lse.masked_fill_(empty_rows[:, None], float("-inf"))
98
99
100 # AG + RS/AR implementation
101
102
103 class CorrectAttnCPOutKernel(VllmTritonJitKernel["CorrectAttnCPOutKernel.CompileKey"]):
104 @dataclass(frozen=True)
105 class CompileKey:
106 output_dtype: torch.dtype
107 lse_dtype: torch.dtype
108 outputs_stride_b: int
109 outputs_stride_h: int
110 outputs_stride_d: int
111 lses_stride_n: int
112 lses_stride_b: int
113 lses_stride_h: int
114 lse_idx: int
115 head_dim: int
116 n_rounded: int
117 is_base_e: bool
118
119 @staticmethod
120 @triton.jit
121 def kernel(
122 outputs_ptr,
123 new_output_ptr,
124 lses_ptr,
125 vlse_ptr,
126 outputs_stride_B,
127 outputs_stride_H,
128 outputs_stride_D,
129 lses_stride_N,
130 lses_stride_B,
131 lses_stride_H,
132 lse_idx,
133 HEAD_DIM: tl.constexpr,
134 N_ROUNDED: tl.constexpr,
135 IS_BASE_E: tl.constexpr,
136 ):
137 """Apply the all-gathered lses to correct each local rank's attention
138 output. we still need perform a cross-rank reduction to obtain the
139 final attention output.
140
141 Args:
142 outputs_ptr (triton.PointerType):
143 Pointer to input tensor of shape [ B, H, D ]
144 lses_ptr (triton.PointerType):
145 Pointer to input tensor of shape [ N, B, H ]
146 new_output_ptr (triton.PointerType):
147 Pointer to output tensor of shape [ B, H, D ]
148 vlse_ptr (triton.PointerType):
149 Pointer to output tensor of shape [ B, H ]
150 outputs_stride_B (int): Batch stride of ``outputs_ptr``
151 outputs_stride_H (int): Head stride of ``outputs_ptr``
152 outputs_stride_D (int): Head-dim stride of ``outputs_ptr``
153 lses_stride_N (int): Rank stride of ``lses_ptr``
154 lses_stride_B (int): Batch stride of ``lses_ptr``
155 lses_stride_H (int): Head stride of ``lses_ptr``
156 lse_idx (int): Index of this rank's lse within the all-gathered tensor
157 HEAD_DIM: Head dimension, as a constexpr
158 N_ROUNDED: Rank count rounded to a power of two, as a constexpr
159 IS_BASE_E: Whether the lses are natural-log based, as a constexpr
160
161 """
162 batch_idx = tl.program_id(axis=0).to(tl.int64)
163 head_idx = tl.program_id(axis=1).to(tl.int64)
164 d_offsets = tl.arange(0, HEAD_DIM)
165 num_n_offsets = tl.arange(0, N_ROUNDED)
166
167 # shape = [N]
168 lse_offsets = (
169 num_n_offsets * lses_stride_N
170 + batch_idx * lses_stride_B
171 + head_idx * lses_stride_H
172 )
173
174 # calc final lse
175 lse = tl.load(lses_ptr + lse_offsets).to(tl.float32)
176 lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse)
177 lse_max = tl.max(lse, axis=0)
178 lse_max = tl.where(lse_max == -float("inf"), 0, lse_max)
179 lse -= lse_max
180 if IS_BASE_E:
181 lse_exp = tl.exp(lse)
182 lse_acc = tl.sum(lse_exp, axis=0)
183 lse = tl.log(lse_acc)
184 else:
185 lse_exp = tl.exp2(lse)
186 lse_acc = tl.sum(lse_exp, axis=0)
187 lse = tl.log2(lse_acc)
188 lse += lse_max
189
190 lse_offsets = batch_idx * lses_stride_B + head_idx * lses_stride_H
191 tl.store(vlse_ptr + lse_offsets, lse)
192
193 # shape = [D]
194 output_offsets = (
195 batch_idx * outputs_stride_B
196 + head_idx * outputs_stride_H
197 + d_offsets * outputs_stride_D
198 )
199
200 # correct output
201 lse_offset = (
202 lse_idx * lses_stride_N
203 + batch_idx * lses_stride_B
204 + head_idx * lses_stride_H
205 )
206 lse_tmp = tl.load(lses_ptr + lse_offset).to(tl.float32)
207 lse_finally = lse_tmp - lse
208 lse_finally = tl.where(
209 (lse_finally != lse_finally) | (lse_finally == float("inf")),
210 -float("inf"),
211 lse_finally,
212 )
213 factor = tl.exp(lse_finally) if IS_BASE_E else tl.exp2(lse_finally)
214 output = tl.load(outputs_ptr + output_offsets)
215 output = output * factor
216 output = tl.where(factor == 0.0, 0.0, output)
217
218 tl.store(new_output_ptr + output_offsets, output)
219
220 def dispatch( # type: ignore[override]
221 self,
222 *,
223 output_dtype: torch.dtype,
224 lse_dtype: torch.dtype,
225 num_tokens: int,
226 num_heads: int,
227 head_dim: int,
228 n_rounded: int,
229 lse_idx: int,
230 is_base_e: bool,
231 ) -> CompileKey:
232 # Warm the contiguous [B, H, D] output and [N, B, H] LSE layouts used
233 # by the production collective path. Runtime still forwards real
234 # strides so unsupported layouts remain correct, but may compile on
235 # first use under the JIT monitor.
236 outputs_stride_d = 1
237 outputs_stride_h = head_dim
238 outputs_stride_b = num_heads * head_dim
239 lses_stride_h = 1
240 lses_stride_b = num_heads
241 lses_stride_n = num_tokens * num_heads
242 return self.CompileKey(
243 output_dtype=output_dtype,
244 lse_dtype=lse_dtype,
245 outputs_stride_b=triton_scalar_specialization_rep(outputs_stride_b),
246 outputs_stride_h=triton_scalar_specialization_rep(outputs_stride_h),
247 outputs_stride_d=triton_scalar_specialization_rep(outputs_stride_d),
248 lses_stride_n=triton_scalar_specialization_rep(lses_stride_n),
249 lses_stride_b=triton_scalar_specialization_rep(lses_stride_b),
250 lses_stride_h=triton_scalar_specialization_rep(lses_stride_h),
251 lse_idx=triton_scalar_specialization_rep(lse_idx),
252 head_dim=head_dim,
253 n_rounded=n_rounded,
254 is_base_e=is_base_e,
255 )
256
257 def get_warmup_keys(
258 self,
259 vllm_config: Any,
260 *,
261 output_dtype: torch.dtype | None = None,
262 num_heads: int | None = None,
263 head_dim: int | None = None,
264 is_base_e: bool | tuple[bool, ...] = (False, True),
265 ) -> list[CompileKey]:
266 from vllm.model_executor.layers.attention.mla_attention import (
267 get_mla_dims,
268 )
269
270 dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
271 max_tokens = min(
272 16,
273 vllm_config.scheduler_config.max_num_batched_tokens,
274 )
275 hf_config = vllm_config.model_config.hf_config
276 if num_heads is None:
277 num_heads = (
278 hf_config.num_attention_heads
279 // vllm_config.parallel_config.tensor_parallel_size
280 )
281 if head_dim is None:
282 head_dim = get_mla_dims(vllm_config.model_config).v_head_dim
283 if output_dtype is None:
284 output_dtype = vllm_config.model_config.dtype
285 if dcp_world_size <= 1 or max_tokens <= 0 or num_heads <= 0:
286 return []
287 return self._trace_dispatch(self.dispatch)(
288 output_dtype=output_dtype,
289 lse_dtype=torch.float32,
290 num_tokens=WarmupIntRange(1, max_tokens + 1),
291 num_heads=num_heads,
292 head_dim=head_dim,
293 n_rounded=dcp_world_size,
294 lse_idx=WarmupIntRange(0, dcp_world_size),
295 is_base_e=is_base_e,
296 )
297
298 def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
299 output_ptr = TritonWarmupTensor(
300 compile_key.output_dtype,
301 shape=(1, 1, compile_key.head_dim),
302 strides=(
303 compile_key.outputs_stride_b,
304 compile_key.outputs_stride_h,
305 compile_key.outputs_stride_d,
306 ),
307 )
308 lse_ptr = TritonWarmupTensor(
309 compile_key.lse_dtype,
310 shape=(compile_key.n_rounded, 1, 1),
311 strides=(
312 compile_key.lses_stride_n,
313 compile_key.lses_stride_b,
314 compile_key.lses_stride_h,
315 ),
316 )
317 return dict(
318 outputs=output_ptr,
319 new_output=output_ptr,
320 lses=lse_ptr,
321 vlse=lse_ptr,
322 lse_idx=compile_key.lse_idx,
323 ctx=None,
324 is_base_e=compile_key.is_base_e,
325 )
326
327 @kernel_launcher
328 def __call__(
329 self,
330 outputs: torch.Tensor,
331 new_output: torch.Tensor,
332 lses: torch.Tensor,
333 vlse: torch.Tensor,
334 lse_idx: int,
335 ctx: Any,
336 *,
337 is_base_e: bool,
338 ) -> LaunchSpec:
339 num_tokens, num_heads, head_dim = outputs.shape
340 n_rounded = lses.shape[0]
341 outputs_stride_b, outputs_stride_h, outputs_stride_d = outputs.stride()
342 lses_stride_n, lses_stride_b, lses_stride_h = lses.stride()
343 grid = (num_tokens, num_heads, 1)
344 return grid, dict(
345 outputs_stride_B=outputs_stride_b,
346 outputs_stride_H=outputs_stride_h,
347 outputs_stride_D=outputs_stride_d,
348 lses_stride_N=lses_stride_n,
349 lses_stride_B=lses_stride_b,
350 lses_stride_H=lses_stride_h,
351 HEAD_DIM=head_dim,
352 N_ROUNDED=n_rounded,
353 IS_BASE_E=is_base_e,
354 _runtime_launcher=None if self._warming else ctx.call_kernel,
355 # CPTritonContext caches the non-constexpr positional prefix; derive
356 # its length so adding/reordering a kernel arg cannot silently
357 # misalign the cached replay path.
358 _runtime_launcher_arg_count=len(self.kernel.arg_names)
359 - len(self.kernel.constexprs),
360 )
361
362
363 class CPTritonContext:
364 """The CPTritonContext is used to avoid recompilation of the Triton JIT."""
365
366 def __init__(self):
367 self.inner_kernel = None
368
369 def call_kernel(self, kernel, grid, *regular_args, **const_args):
370 if self.inner_kernel is None:
371 self.inner_kernel = kernel[grid](*regular_args, **const_args)
372 else:
373 self.inner_kernel[grid](*regular_args)
374
375
376 def correct_attn_out(
377 out: torch.Tensor,
378 lses: torch.Tensor,
379 cp_rank: int,
380 ctx: CPTritonContext,
381 is_lse_base_on_e: bool = True,
382 ) -> tuple[torch.Tensor, torch.Tensor]:
383 """Correct the attention output using the all-gathered lses.
384
385 Args:
386 out: Tensor of shape [ B, H, D ]
387 lses: Tensor of shape [ N, B, H ]
388 cp_rank: Current rank in the context-parallel group
389 ctx: Triton context to avoid recompilation
390 is_lse_base_on_e: Whether the lses use base e rather than base 2
391
392 Returns:
393 Tuple of (out, lse) with corrected attention and final log-sum-exp.
394
395 """
396 if ctx is None:
397 ctx = CPTritonContext()
398
399 # --- Normalize to 3D views ---
400 if out.ndim == 4 and out.shape[1] == 1:
401 out = out.squeeze(1)
402 assert out.ndim == 3, f"expected out [B,H,D] or [B,1,H,D], got {tuple(out.shape)}"
403
404 if lses.ndim == 4 and lses.shape[-1] == 1:
405 lses = lses.squeeze(-1)
406 if lses.ndim == 4 and lses.shape[1] == 1:
407 lses = lses.squeeze(1)
408 assert lses.ndim == 3, (
409 f"expected lses [N,B,H] (optionally with a 1-sized extra dim), "
410 f"got {tuple(lses.shape)}"
411 )
412
413 B, H, _ = out.shape
414
415 # Strides after we normalized shapes to 3-D views. The kernel computes
416 # offsets for `vlse_ptr` using lses_stride_B/H, so the output buffer must
417 # have the same B/H stride layout as a slice of `lses`.
418 _, l_sB, l_sH = lses.stride()
419
420 # Allocate LSE with the same B/H strides as `lses` so writes land correctly
421 # even when `lses` is a non-contiguous view (e.g., 4-D to 3-D squeeze).
422 lse = torch.empty_strided(
423 (B, H), (l_sB, l_sH), device=lses.device, dtype=lses.dtype
424 )
425
426 _CORRECT_ATTN_CP_OUT_KERNEL(
427 out,
428 out,
429 lses,
430 lse,
431 cp_rank,
432 ctx,
433 is_base_e=is_lse_base_on_e,
434 )
435 return out, lse
436
437
438 def _cp_lse_common(
439 cp_attn_out: torch.Tensor,
440 cp_attn_lse: torch.Tensor,
441 cp_group: GroupCoordinator,
442 ctx: CPTritonContext | None = None,
443 is_lse_base_on_e=True,
444 seq_lens: torch.Tensor | None = None,
445 query_start_loc: torch.Tensor | None = None,
446 ):
447 """cp_attn_out: [ B, H, D ]
448 cp_attn_lse: [ B, H ]
449 """
450 if cp_group.world_size == 1:
451 return cp_attn_out
452
453 if ctx is None:
454 ctx = CPTritonContext()
455
456 cp_attn_lse = cp_attn_lse.contiguous()
457 mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc)
458 lses = cp_group.all_gather(cp_attn_lse, dim=0).reshape(
459 (cp_group.world_size,) + cp_attn_lse.shape
460 )
461 out, lse = correct_attn_out(
462 cp_attn_out,
463 lses,
464 cp_group.rank_in_group,
465 ctx,
466 is_lse_base_on_e=is_lse_base_on_e,
467 )
468 return out, lse
469
470
471 def cp_lse_ag_out_rs(
472 cp_attn_out: torch.Tensor,
473 cp_attn_lse: torch.Tensor,
474 cp_group: GroupCoordinator,
475 ctx: CPTritonContext | None = None,
476 return_lse: bool = False,
477 is_lse_base_on_e=True,
478 seq_lens: torch.Tensor | None = None,
479 query_start_loc: torch.Tensor | None = None,
480 ):
481 """cp_attn_out: [ B, H, D ]
482 cp_attn_lse: [ B, H ]
483 """
484 out, lse = _cp_lse_common(
485 cp_attn_out,
486 cp_attn_lse,
487 cp_group,
488 ctx=ctx,
489 is_lse_base_on_e=is_lse_base_on_e,
490 seq_lens=seq_lens,
491 query_start_loc=query_start_loc,
492 )
493 out = cp_group.reduce_scatter(out, dim=1)
494
495 if return_lse:
496 cp_num_heads = lse.shape[1] // cp_group.world_size
497 cp_rank = cp_group.rank_in_group
498 lse = lse[:, cp_num_heads * cp_rank : cp_num_heads * (cp_rank + 1)]
499 return out, lse
500 return out
501
502
503 def cp_lse_ag_out_ar(
504 cp_attn_out: torch.Tensor,
505 cp_attn_lse: torch.Tensor,
506 cp_group: GroupCoordinator,
507 ctx: CPTritonContext | None = None,
508 return_lse: bool = False,
509 is_lse_base_on_e=True,
510 seq_lens: torch.Tensor | None = None,
511 query_start_loc: torch.Tensor | None = None,
512 ):
513 """cp_attn_out: [ B, H, D ]
514 cp_attn_lse: [ B, H ]
515 """
516 out, lse = _cp_lse_common(
517 cp_attn_out,
518 cp_attn_lse,
519 cp_group,
520 ctx=ctx,
521 is_lse_base_on_e=is_lse_base_on_e,
522 seq_lens=seq_lens,
523 query_start_loc=query_start_loc,
524 )
525 out = cp_group.all_reduce(out)
526
527 if return_lse:
528 return out, lse
529 return out
530
531
532 # Standard A2A implementation
533
534
535 def _lse_weighted_combine(
536 outputs: torch.Tensor,
537 lses: torch.Tensor,
538 return_lse: bool = False,
539 is_lse_base_on_e: bool = True,
540 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
541 """CPU reference implementation for LSE-weighted combination.
542
543 This is a pure PyTorch implementation used for testing and validation.
544
545 Args:
546 outputs: Partial attention outputs [N, B, H, D]
547 N = number of KV shards (ranks)
548 B = batch size (num_tokens)
549 H = number of heads per rank
550 D = head dimension
551 lses: Log-sum-exp values [N, B, H]
552 return_lse: If True, also return the global LSE
553 is_lse_base_on_e: If True, LSE is base e; if False, base 2
554
555 Returns:
556 Combined output [B, H, D], and optionally global LSE [B, H]
557
558 """
559 N, B, H, D = outputs.shape
560
561 # Handle NaN and inf in LSEs
562 lses = torch.where(
563 torch.isnan(lses) | torch.isinf(lses),
564 torch.tensor(float("-inf"), device=lses.device, dtype=lses.dtype),
565 lses,
566 )
567
568 # Compute max LSE for numerical stability
569 lse_max, _ = lses.max(dim=0) # [B, H]
570 lse_max = torch.where(
571 lse_max == float("-inf"),
572 torch.zeros_like(lse_max),
573 lse_max,
574 )
575
576 # Compute weights: softmax over the N dimension
577 if is_lse_base_on_e:
578 weights = torch.exp(lses - lse_max.unsqueeze(0)) # [N, B, H]
579 else:
580 weights = torch.pow(2.0, lses - lse_max.unsqueeze(0)) # [N, B, H]
581
582 # Handle NaN weights
583 weights = torch.where(torch.isnan(weights), torch.zeros_like(weights), weights)
584
585 # Normalize weights
586 weight_sum = weights.sum(dim=0, keepdim=True) # [1, B, H]
587 weights = weights / weight_sum.clamp(min=1e-10) # [N, B, H]
588
589 # Weighted combination: sum over N dimension
590 weights = weights.unsqueeze(-1)
591 outputs = torch.where(weights == 0, torch.zeros_like(outputs), outputs)
592 result = (outputs * weights).sum(dim=0) # [B, H, D]
593
594 if return_lse:
595 if is_lse_base_on_e:
596 global_lse = torch.log(weight_sum.squeeze(0)) + lse_max # [B, H]
597 else:
598 global_lse = torch.log2(weight_sum.squeeze(0)) + lse_max # [B, H]
599 return result, global_lse
600
601 return result
602
603
604 def _dcp_a2a_lse_pack_dim(output_dtype: torch.dtype) -> int:
605 bits = torch.finfo(output_dtype).bits
606 if bits == 16:
607 return 2
608 if bits == 32:
609 return 1
610 raise ValueError(f"Cannot pack fp32 LSE into output dtype {output_dtype}.")
611
612
613 def _dcp_a2a_send_recv_buffers(
614 shape: tuple[int, ...],
615 device: torch.device,
616 dtype: torch.dtype,
617 ) -> tuple[torch.Tensor, torch.Tensor]:
618 # Don't use the shared WorkspaceManager here. A FULL cudagraph bakes in the
619 # buffer address at capture, but the workspace is growable and sized only to
620 # the largest *captured* batch (the cudagraph capture cap). Any eager a2a
621 # with a bigger batch regrows it, freeing that address and poisoning every
622 # captured graph -> illegal memory access on replay. This bites the very
623 # first request: the post-capture warmup runs an eager decode at
624 # max_num_seqs (> the cap), so the graphs are already dangling before the
625 # server is ready. torch.empty buffers instead live in the graph's private
626 # pool and stay valid for its lifetime (as _dcp_a2a_unpack_combine and the
627 # AG+RS combine path already rely on).
628 return (
629 torch.empty(shape, device=device, dtype=dtype),
630 torch.empty(shape, device=device, dtype=dtype),
631 )
632
633
634 @triton.jit
635 def _dcp_a2a_pack_send_kernel(
636 out_ptr,
637 lse_ptr,
638 send_ptr,
639 seq_lens_ptr,
640 query_start_loc_ptr,
641 num_seqs,
642 out_stride_B,
643 out_stride_H,
644 out_stride_D,
645 lse_stride_B,
646 lse_stride_H,
647 send_stride_N,
648 send_stride_B,
649 send_stride_H,
650 send_stride_D,
651 N: tl.constexpr,
652 HEAD_DIM: tl.constexpr,
653 H_PER_RANK: tl.constexpr,
654 LSE_PACK_DIM: tl.constexpr,
655 HAS_MASK: tl.constexpr,
656 NUM_SEQS_POW2: tl.constexpr,
657 ):
658 batch_idx = tl.program_id(0).to(tl.int64)
659 local_head_idx = tl.program_id(1).to(tl.int64)
660 d_offsets = tl.arange(0, HEAD_DIM)
661
662 # Empty-shard masking, fused in to avoid a separate per-layer pass over the
663 # LSE. Rows that belong to a request with no local KV shard, and rows past
664 # the end of the batch (CUDA-graph padding), must not contribute to the
665 # combine, which the reduction achieves by giving them an LSE of -inf.
666 row_is_empty = False
667 if HAS_MASK:
668 # searchsorted(query_start_loc[1:], batch_idx, right=True), i.e. the
669 # number of request boundaries at or below this row.
670 seq_offsets = tl.arange(0, NUM_SEQS_POW2)
671 boundaries = tl.load(
672 query_start_loc_ptr + 1 + seq_offsets,
673 mask=seq_offsets < num_seqs,
674 other=2**31 - 1,
675 )
676 seq_idx = tl.sum((boundaries <= batch_idx).to(tl.int32), axis=0)
677 seq_idx = tl.minimum(seq_idx, num_seqs - 1)
678 num_rows = tl.load(query_start_loc_ptr + num_seqs)
679 seq_len = tl.load(seq_lens_ptr + seq_idx)
680 row_is_empty = (batch_idx >= num_rows) | (seq_len == 0)
681
682 for rank_idx in tl.static_range(N):
683 src_head_idx = rank_idx * H_PER_RANK + local_head_idx
684 send_base = (
685 rank_idx * send_stride_N
686 + batch_idx * send_stride_B
687 + local_head_idx * send_stride_H
688 )
689
690 out_offsets = (
691 batch_idx * out_stride_B
692 + src_head_idx * out_stride_H
693 + d_offsets * out_stride_D
694 )
695 tl.store(
696 send_ptr + send_base + d_offsets * send_stride_D,
697 tl.load(out_ptr + out_offsets),
698 )
699
700 lse_val = tl.load(
701 lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H
702 ).to(tl.float32)
703 if HAS_MASK:
704 lse_val = tl.where(row_is_empty, -float("inf"), lse_val)
705 if LSE_PACK_DIM == 1:
706 tl.store(
707 send_ptr + send_base + HEAD_DIM * send_stride_D,
708 lse_val.to(send_ptr.dtype.element_ty),
709 )
710 else:
711 lse_bits = lse_val.to(tl.uint32, bitcast=True)
712 lo = (lse_bits & 0xFFFF).to(tl.uint16)
713 hi = ((lse_bits >> 16) & 0xFFFF).to(tl.uint16)
714 tl.store(
715 send_ptr + send_base + HEAD_DIM * send_stride_D,
716 lo.to(send_ptr.dtype.element_ty, bitcast=True),
717 )
718 tl.store(
719 send_ptr + send_base + (HEAD_DIM + 1) * send_stride_D,
720 hi.to(send_ptr.dtype.element_ty, bitcast=True),
721 )
722
723
724 @triton.jit
725 def _dcp_a2a_unpack_combine_kernel(
726 recv_ptr,
727 out_ptr,
728 out_lse_ptr,
729 recv_stride_N,
730 recv_stride_B,
731 recv_stride_H,
732 recv_stride_D,
733 out_stride_B,
734 out_stride_H,
735 out_stride_D,
736 out_lse_stride_B,
737 out_lse_stride_H,
738 N: tl.constexpr,
739 HEAD_DIM: tl.constexpr,
740 IS_BASE_E: tl.constexpr,
741 RETURN_LSE: tl.constexpr,
742 LSE_PACK_DIM: tl.constexpr,
743 ):
744 batch_idx = tl.program_id(0).to(tl.int64)
745 head_idx = tl.program_id(1).to(tl.int64)
746 d_offsets = tl.arange(0, HEAD_DIM)
747
748 lse_max = -float("inf")
749 for rank_idx in tl.static_range(N):
750 recv_base = (
751 rank_idx * recv_stride_N
752 + batch_idx * recv_stride_B
753 + head_idx * recv_stride_H
754 )
755 if LSE_PACK_DIM == 1:
756 lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
757 tl.float32
758 )
759 else:
760 lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
761 hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
762 lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
763 hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
764 lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
765 lse_val = tl.where(
766 (lse_val != lse_val) | (lse_val == float("inf")),
767 -float("inf"),
768 lse_val,
769 )
770 lse_max = tl.maximum(lse_max, lse_val)
771
772 lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max)
773
774 lse_sum = 0.0
775 for rank_idx in tl.static_range(N):
776 recv_base = (
777 rank_idx * recv_stride_N
778 + batch_idx * recv_stride_B
779 + head_idx * recv_stride_H
780 )
781 if LSE_PACK_DIM == 1:
782 lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
783 tl.float32
784 )
785 else:
786 lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
787 hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
788 lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
789 hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
790 lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
791 lse_val = tl.where(
792 (lse_val != lse_val) | (lse_val == float("inf")),
793 -float("inf"),
794 lse_val,
795 )
796 if IS_BASE_E:
797 lse_sum += tl.exp(lse_val - lse_max)
798 else:
799 lse_sum += tl.exp2(lse_val - lse_max)
800
801 if IS_BASE_E: # noqa: SIM108
802 global_lse = tl.log(lse_sum) + lse_max
803 else:
804 global_lse = tl.log2(lse_sum) + lse_max
805
806 acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
807 for rank_idx in tl.static_range(N):
808 recv_base = (
809 rank_idx * recv_stride_N
810 + batch_idx * recv_stride_B
811 + head_idx * recv_stride_H
812 )
813 if LSE_PACK_DIM == 1:
814 lse_val = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D).to(
815 tl.float32
816 )
817 else:
818 lo_raw = tl.load(recv_ptr + recv_base + HEAD_DIM * recv_stride_D)
819 hi_raw = tl.load(recv_ptr + recv_base + (HEAD_DIM + 1) * recv_stride_D)
820 lo = lo_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
821 hi = hi_raw.to(tl.uint16, bitcast=True).to(tl.uint32)
822 lse_val = (lo | (hi << 16)).to(tl.float32, bitcast=True)
823 lse_val = tl.where(
824 (lse_val != lse_val) | (lse_val == float("inf")),
825 -float("inf"),
826 lse_val,
827 )
828 if IS_BASE_E:
829 weight = tl.exp(lse_val - global_lse)
830 else:
831 weight = tl.exp2(lse_val - global_lse)
832 weight = tl.where(weight != weight, 0.0, weight)
833 partial = tl.load(recv_ptr + recv_base + d_offsets * recv_stride_D).to(
834 tl.float32
835 )
836 partial = tl.where(weight == 0.0, 0.0, partial)
837 acc += partial * weight
838
839 final_offsets = (
840 batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D
841 )
842 tl.store(out_ptr + final_offsets, acc)
843
844 if RETURN_LSE:
845 out_lse_offset = batch_idx * out_lse_stride_B + head_idx * out_lse_stride_H
846 tl.store(out_lse_ptr + out_lse_offset, global_lse)
847
848
849 def _dcp_a2a_pack_send(
850 cp_attn_out: torch.Tensor,
851 cp_attn_lse: torch.Tensor,
852 send_buffer: torch.Tensor,
853 world_size: int,
854 h_per_rank: int,
855 head_dim: int,
856 lse_pack_dim: int,
857 seq_lens: torch.Tensor | None = None,
858 query_start_loc: torch.Tensor | None = None,
859 ) -> None:
860 # The empty-shard mask is fused into the pack kernel below rather than
861 # applied as a separate pass: the mask is a pure function of seq_lens and
862 # query_start_loc, which are identical for every MLA layer in a step, while
863 # the eager form costs eight launch-bound kernels per layer.
864 has_mask = _validate_dcp_empty_shard_args(seq_lens, query_start_loc)
865 num_seqs = seq_lens.shape[0] if has_mask and seq_lens is not None else 0
866
867 grid = (cp_attn_out.shape[0], h_per_rank, 1)
868 _dcp_a2a_pack_send_kernel[grid](
869 cp_attn_out,
870 cp_attn_lse,
871 send_buffer,
872 seq_lens,
873 query_start_loc,
874 num_seqs,
875 cp_attn_out.stride(0),
876 cp_attn_out.stride(1),
877 cp_attn_out.stride(2),
878 cp_attn_lse.stride(0),
879 cp_attn_lse.stride(1),
880 send_buffer.stride(0),
881 send_buffer.stride(1),
882 send_buffer.stride(2),
883 send_buffer.stride(3),
884 N=world_size,
885 HEAD_DIM=head_dim,
886 H_PER_RANK=h_per_rank,
887 LSE_PACK_DIM=lse_pack_dim,
888 HAS_MASK=has_mask,
889 # Only the arange bound is a constexpr: num_seqs itself is a runtime
890 # argument so the kernel specializes on at most log2(max_num_seqs)
891 # variants instead of one per distinct batch size.
892 NUM_SEQS_POW2=triton.next_power_of_2(max(num_seqs, 1)),
893 )
894
895
896 def _dcp_a2a_unpack_combine(
897 recv_buffer: torch.Tensor,
898 head_dim: int,
899 lse_pack_dim: int,
900 return_lse: bool,
901 is_lse_base_on_e: bool,
902 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
903 world_size, num_tokens, h_per_rank, _ = recv_buffer.shape
904 out = torch.empty(
905 (num_tokens, h_per_rank, head_dim),
906 device=recv_buffer.device,
907 dtype=recv_buffer.dtype,
908 )
909 out_lse = torch.empty(
910 (num_tokens, h_per_rank) if return_lse else (1, 1),
911 device=recv_buffer.device,
912 dtype=torch.float32 if return_lse else recv_buffer.dtype,
913 )
914 grid = (num_tokens, h_per_rank, 1)
915 _dcp_a2a_unpack_combine_kernel[grid](
916 recv_buffer,
917 out,
918 out_lse,
919 recv_buffer.stride(0),
920 recv_buffer.stride(1),
921 recv_buffer.stride(2),
922 recv_buffer.stride(3),
923 out.stride(0),
924 out.stride(1),
925 out.stride(2),
926 out_lse.stride(0),
927 out_lse.stride(1),
928 N=world_size,
929 HEAD_DIM=head_dim,
930 IS_BASE_E=is_lse_base_on_e,
931 RETURN_LSE=return_lse,
932 LSE_PACK_DIM=lse_pack_dim,
933 )
934 if return_lse:
935 return out, out_lse
936 return out
937
938
939 def dcp_a2a_lse_reduce(
940 cp_attn_out: torch.Tensor,
941 cp_attn_lse: torch.Tensor,
942 cp_group: GroupCoordinator,
943 ctx: CPTritonContext | None = None,
944 return_lse: bool = False,
945 is_lse_base_on_e: bool = True,
946 seq_lens: torch.Tensor | None = None,
947 query_start_loc: torch.Tensor | None = None,
948 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
949 """Combine partial attention outputs across DCP ranks using All-to-All.
950
951 The output and LSE are packed into a single output-dtype buffer, sent
952 with one All-to-All, then unpacked and combined with exact LSE weighting.
953
954 Args:
955 cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim
956 cp_attn_lse: [B, H] floating-point log-sum-exp values
957 cp_group: GroupCoordinator for DCP communication
958 ctx: CPTritonContext (unused, for signature compatibility)
959 return_lse: If True, also return the combined global LSE
960 is_lse_base_on_e: If True, LSE is base e; if False, base 2
961 seq_lens: Local KV lengths. Empty shards contribute zero weight.
962 query_start_loc: Cumulative query-token offsets for each request.
963
964 Returns:
965 Combined output [B, H/N, D] (head-scattered)
966 If return_lse=True, also returns global_lse [B, H/N]
967
968 """
969 world_size = cp_group.world_size
970
971 if world_size == 1:
972 if return_lse:
973 return cp_attn_out, cp_attn_lse
974 return cp_attn_out
975
976 B, H, D = cp_attn_out.shape
977 if H % world_size != 0:
978 raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.")
979 H_per_rank = H // world_size
980 lse_pack_dim = _dcp_a2a_lse_pack_dim(cp_attn_out.dtype)
981
982 send_buffer, recv_buffer = _dcp_a2a_send_recv_buffers(
983 (world_size, B, H_per_rank, D + lse_pack_dim),
984 device=cp_attn_out.device,
985 dtype=cp_attn_out.dtype,
986 )
987
988 _dcp_a2a_pack_send(
989 cp_attn_out,
990 cp_attn_lse,
991 send_buffer,
992 world_size,
993 H_per_rank,
994 D,
995 lse_pack_dim,
996 seq_lens=seq_lens,
997 query_start_loc=query_start_loc,
998 )
999
1000 work = dist.all_to_all_single(
1001 recv_buffer.view(-1),
1002 send_buffer.view(-1),
1003 group=cp_group.device_group,
1004 async_op=True,
1005 )
1006 work.wait()
1007
1008 return _dcp_a2a_unpack_combine(
1009 recv_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e
1010 )
1011
1012
1013 def get_dcp_workspace_max_num_tokens(vllm_config: VllmConfig) -> int:
1014 scheduler_config = vllm_config.scheduler_config
1015 # num_speculative_tokens also counts a diffusion canvas.
1016 tokens_per_seq = max(
1017 1 + vllm_config.num_speculative_tokens, max_decode_query_len(vllm_config)
1018 )
1019 return min(
1020 scheduler_config.max_num_batched_tokens,
1021 max(
1022 scheduler_config.max_num_seqs * tokens_per_seq,
1023 vllm_config.compilation_config.max_cudagraph_capture_size or 0,
1024 ),
1025 )
1026
1027
1028 def reserve_query_head_storage(
1029 query: torch.Tensor, padded_num_heads: int
1030 ) -> torch.Tensor:
1031 """Reserve backing storage for fixed-head decode kernels."""
1032 assert query.ndim == 3
1033 assert query.shape[1] <= padded_num_heads
1034 padded = query.new_empty((query.shape[0], padded_num_heads, query.shape[2]))
1035 padded.resize_(query.shape)
1036 padded.copy_(query)
1037 return padded
1038
1039
1040 # Symmetric-memory A2A implementation
1041
1042
1043 _A2A_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
1044
1045
1046 class DirectDCPA2AWorkspace(DirectCPWorkspace):
1047 """Persistent symmetric buffers for direct DCP output exchange."""
1048
1049 def __init__(
1050 self,
1051 group: ProcessGroup,
1052 device: torch.device,
1053 max_num_tokens: int,
1054 heads_per_rank: int,
1055 head_dim: int,
1056 dtype: torch.dtype = torch.bfloat16,
1057 num_ubatches: int = 1,
1058 ) -> None:
1059 if dtype not in _A2A_SUPPORTED_DTYPES:
1060 raise ValueError(f"Direct DCP A2A does not support {dtype}")
1061 if num_ubatches < 1:
1062 raise ValueError(
1063 f"Direct DCP A2A requires at least one ubatch slot, got {num_ubatches}"
1064 )
1065 super().__init__(group, device, num_ubatches)
1066 self.max_num_tokens = max_num_tokens
1067 self.heads_per_rank = heads_per_rank
1068 self.head_dim = head_dim
1069
1070 output_shape = (
1071 num_ubatches,
1072 2,
1073 self.world_size,
1074 max_num_tokens,
1075 heads_per_rank,
1076 head_dim,
1077 )
1078 lse_shape = (
1079 num_ubatches,
1080 2,
1081 self.world_size,
1082 max_num_tokens,
1083 heads_per_rank,
1084 )
1085 signal_shape = (num_ubatches, 2, self.world_size)
1086 self.received_output, self.peer_output_ptrs = self._allocate(
1087 output_shape, dtype
1088 )
1089 self.received_lse, self.peer_lse_ptrs = self._allocate(lse_shape, torch.float32)
1090 self.received_signal, self.peer_signal_ptrs = self._allocate(
1091 signal_shape, torch.int32
1092 )
1093
1094 def lse_reduce(
1095 self,
1096 partial_output: torch.Tensor,
1097 partial_lse: torch.Tensor,
1098 is_lse_base_on_e: bool,
1099 seq_lens: torch.Tensor | None = None,
1100 query_start_loc: torch.Tensor | None = None,
1101 ) -> torch.Tensor:
1102 ubatch = dbo_current_ubatch_id()
1103 num_tokens = partial_output.shape[0]
1104 output = partial_output.new_empty(
1105 (num_tokens, self.heads_per_rank, self.head_dim)
1106 )
1107 torch.ops._C.direct_dcp_a2a_lse_reduce(
1108 partial_output,
1109 partial_lse,
1110 seq_lens,
1111 query_start_loc,
1112 self.peer_output_ptrs[ubatch],
1113 self.peer_lse_ptrs[ubatch],
1114 self.peer_signal_ptrs[ubatch],
1115 self.received_output[ubatch],
1116 self.received_lse[ubatch],
1117 self.received_signal[ubatch],
1118 self.epoch[ubatch : ubatch + 1],
1119 output,
1120 self.world_size,
1121 self.rank,
1122 self.max_num_tokens,
1123 is_lse_base_on_e,
1124 )
1125 return output
1126
1127
1128 @functools.cache
1129 def get_direct_dcp_a2a_workspace(
1130 group: GroupCoordinator,
1131 device: torch.device,
1132 max_num_tokens: int,
1133 heads_per_rank: int,
1134 head_dim: int,
1135 dtype: torch.dtype,
1136 num_ubatches: int,
1137 ) -> DirectDCPA2AWorkspace | None:
1138 if not direct_cp_enabled(
1139 group, dtype, envs.VLLM_USE_DIRECT_DCP_A2A, _A2A_SUPPORTED_DTYPES
1140 ):
1141 return None
1142 return DirectDCPA2AWorkspace(
1143 group.device_group,
1144 device,
1145 max_num_tokens,
1146 heads_per_rank,
1147 head_dim,
1148 dtype,
1149 num_ubatches,
1150 )
1151
1152
1153 # Q gather
1154
1155 # Symmetric-memory implementation
1156
1157
1158 def _q_gather_layout_supported(
1159 world_size: int,
1160 heads_per_rank: int,
1161 head_dim: int,
1162 dtype: torch.dtype,
1163 padded_num_heads: int | None,
1164 ) -> bool:
1165 element_size = torch.empty((), dtype=dtype).element_size()
1166 gathered_num_heads = world_size * heads_per_rank
1167 storage_num_heads = (
1168 gathered_num_heads if padded_num_heads is None else padded_num_heads
1169 )
1170 return (
1171 heads_per_rank * head_dim * element_size % 16 == 0
1172 and storage_num_heads * head_dim * element_size % 16 == 0
1173 )
1174
1175
1176 class DirectDCPQGatherWorkspace(DirectCPWorkspace):
1177 """Publish query shards directly into the consumer-final symmetric buffer.
1178
1179 The final buffer is reusable after the downstream DCP output combine. That
1180 combine orders all ranks after attention has consumed the gathered query.
1181 """
1182
1183 def __init__(
1184 self,
1185 group: ProcessGroup,
1186 device: torch.device,
1187 max_num_tokens: int,
1188 heads_per_rank: int,
1189 head_dim: int,
1190 dtype: torch.dtype = torch.bfloat16,
1191 num_ubatches: int = 1,
1192 padded_num_heads: int | None = None,
1193 ) -> None:
1194 if num_ubatches < 1:
1195 raise ValueError(
1196 "Direct DCP q-gather requires at least one ubatch slot, "
1197 f"got {num_ubatches}"
1198 )
1199 if max_num_tokens < 1 or heads_per_rank < 1 or head_dim < 1:
1200 raise ValueError(
1201 "Direct DCP q-gather dimensions must be positive, got "
1202 f"T={max_num_tokens}, H={heads_per_rank}, D={head_dim}"
1203 )
1204 gathered_num_heads = group.size() * heads_per_rank
1205 if not _q_gather_layout_supported(
1206 group.size(), heads_per_rank, head_dim, dtype, padded_num_heads
1207 ):
1208 raise ValueError("Direct DCP q-gather requires 16-byte-aligned query rows.")
1209 super().__init__(group, device, num_ubatches)
1210 if self.world_size <= 1:
1211 raise ValueError("Direct DCP q-gather requires at least two ranks")
1212 self.max_num_tokens = max_num_tokens
1213 self.heads_per_rank = heads_per_rank
1214 self.gathered_num_heads = gathered_num_heads
1215 self.padded_num_heads = (
1216 self.gathered_num_heads if padded_num_heads is None else padded_num_heads
1217 )
1218 if self.padded_num_heads < self.gathered_num_heads:
1219 raise ValueError(
1220 "Direct DCP q-gather padded heads must cover gathered heads: "
1221 f"{self.padded_num_heads} < {self.gathered_num_heads}"
1222 )
1223 self.head_dim = head_dim
1224
1225 query_shape = (
1226 num_ubatches,
1227 max_num_tokens,
1228 self.padded_num_heads,
1229 head_dim,
1230 )
1231 signal_shape = (num_ubatches, 2, self.world_size)
1232 self.final_query, _ = self._allocate(query_shape, dtype)
1233 self.received_signal, _ = self._allocate(signal_shape, torch.int32)
1234 query_multicast_ptrs = self._multicast_ptrs(self.final_query)
1235 signal_multicast_ptrs = self._multicast_ptrs(self.received_signal)
1236 self.multicast_ptrs = list(
1237 zip(query_multicast_ptrs, signal_multicast_ptrs, strict=True)
1238 )
1239 if not all(
1240 query_ptr and signal_ptr for query_ptr, signal_ptr in self.multicast_ptrs
1241 ):
1242 raise RuntimeError(
1243 "Direct DCP q-gather requires NVLS symmetric-memory multicast."
1244 )
1245 self.completion = self.received_signal.new_zeros((num_ubatches, 1))
1246 torch.accelerator.synchronize()
1247
1248 def gather(self, local_query: torch.Tensor) -> torch.Tensor:
1249 ubatch = dbo_current_ubatch_id()
1250 if not 0 <= ubatch < self.num_ubatches:
1251 raise ValueError(
1252 f"DCP q-gather ubatch {ubatch} exceeds {self.num_ubatches} slots"
1253 )
1254 if local_query.ndim == 3 and local_query.shape[1] != self.heads_per_rank:
1255 raise ValueError(
1256 f"DCP q-gather expected {self.heads_per_rank} local query heads, "
1257 f"got {local_query.shape[1]}"
1258 )
1259
1260 num_tokens = local_query.shape[0]
1261 output = torch.as_strided(
1262 self.final_query[ubatch],
1263 size=(num_tokens, self.gathered_num_heads, self.head_dim),
1264 stride=(
1265 self.gathered_num_heads * self.head_dim,
1266 self.head_dim,
1267 1,
1268 ),
1269 )
1270 query_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch]
1271 torch.ops._C.direct_dcp_q_gather(
1272 local_query,
1273 output,
1274 self.received_signal[ubatch],
1275 self.completion[ubatch],
1276 self.epoch[ubatch : ubatch + 1],
1277 self.world_size,
1278 self.rank,
1279 self.max_num_tokens,
1280 self.padded_num_heads,
1281 query_multicast_ptr,
1282 signal_multicast_ptr,
1283 )
1284 return output
1285
1286
1287 @functools.cache
1288 def get_direct_dcp_q_gather_workspace(
1289 group: GroupCoordinator,
1290 device: torch.device,
1291 max_num_tokens: int,
1292 heads_per_rank: int,
1293 head_dim: int,
1294 dtype: torch.dtype,
1295 num_ubatches: int,
1296 padded_num_heads: int | None = None,
1297 ) -> DirectDCPQGatherWorkspace | None:
1298 if not direct_cp_multicast_enabled(group, dtype, envs.VLLM_USE_DIRECT_DCP_Q_GATHER):
1299 return None
1300 if not _q_gather_layout_supported(
1301 group.world_size, heads_per_rank, head_dim, dtype, padded_num_heads
1302 ):
1303 return None
1304 return DirectDCPQGatherWorkspace(
1305 group.device_group,
1306 device,
1307 max_num_tokens,
1308 heads_per_rank,
1309 head_dim,
1310 dtype,
1311 num_ubatches,
1312 padded_num_heads,
1313 )
1314
1315
1316 # KV gather
1317
1318 # Symmetric-memory implementation
1319
1320
1321 _KV_GATHER_SUPPORTED_DTYPES = (
1322 torch.float16,
1323 torch.bfloat16,
1324 torch.float8_e4m3fn,
1325 )
1326
1327
1328 def _kv_gather_layout_supported(token_dim: int, dtype: torch.dtype) -> bool:
1329 return token_dim * torch.empty((), dtype=dtype).element_size() % 16 == 0
1330
1331
1332 class DirectDCPKVGatherWorkspace(DirectCPWorkspace):
1333 """Persistent symmetric buffers for direct DCP KV gather."""
1334
1335 def __init__(
1336 self,
1337 group: ProcessGroup,
1338 device: torch.device,
1339 max_gathered_tokens: int,
1340 token_dim: int,
1341 dtype: torch.dtype = torch.bfloat16,
1342 num_ubatches: int = 1,
1343 ) -> None:
1344 if dtype not in _KV_GATHER_SUPPORTED_DTYPES:
1345 raise ValueError(f"Direct DCP kv-gather does not support {dtype}")
1346 if num_ubatches < 1:
1347 raise ValueError(
1348 "Direct DCP kv-gather requires at least one ubatch slot, "
1349 f"got {num_ubatches}"
1350 )
1351 if max_gathered_tokens < 1 or token_dim < 1:
1352 raise ValueError(
1353 "Direct DCP kv-gather dimensions must be positive, got "
1354 f"T={max_gathered_tokens}, D={token_dim}"
1355 )
1356 if not _kv_gather_layout_supported(token_dim, dtype):
1357 raise ValueError("Direct DCP kv-gather requires 16-byte-aligned KV rows.")
1358 super().__init__(group, device, num_ubatches)
1359 if self.world_size <= 1:
1360 raise ValueError("Direct DCP kv-gather requires at least two ranks")
1361 if max_gathered_tokens % self.world_size != 0:
1362 raise ValueError(
1363 "Direct DCP kv-gather capacity must divide evenly across "
1364 f"ranks: {max_gathered_tokens} % {self.world_size} != 0"
1365 )
1366 self.max_gathered_tokens = max_gathered_tokens
1367
1368 kv_shape = (num_ubatches, 2, max_gathered_tokens, token_dim)
1369 signal_shape = (num_ubatches, 2, self.world_size)
1370 self.received_kv, _ = self._allocate(kv_shape, dtype)
1371 self.received_signal, _ = self._allocate(signal_shape, torch.int32)
1372 kv_multicast_ptrs = self._multicast_ptrs(self.received_kv)
1373 signal_multicast_ptrs = self._multicast_ptrs(self.received_signal)
1374 self.multicast_ptrs = list(
1375 zip(kv_multicast_ptrs, signal_multicast_ptrs, strict=True)
1376 )
1377 if not all(kv_ptr and signal_ptr for kv_ptr, signal_ptr in self.multicast_ptrs):
1378 raise RuntimeError(
1379 "Direct DCP kv-gather requires NVLS symmetric-memory multicast."
1380 )
1381 self.completion = self.received_signal.new_zeros((num_ubatches, 2))
1382 torch.accelerator.synchronize()
1383
1384 def gather(self, gathered_kv: torch.Tensor, local_kv: torch.Tensor) -> None:
1385 ubatch = dbo_current_ubatch_id()
1386 if not 0 <= ubatch < self.num_ubatches:
1387 raise ValueError(
1388 f"DCP kv-gather ubatch {ubatch} exceeds {self.num_ubatches} slots"
1389 )
1390 kv_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch]
1391 torch.ops._C.direct_dcp_kv_gather(
1392 local_kv,
1393 self.received_kv[ubatch],
1394 self.received_signal[ubatch],
1395 self.completion[ubatch],
1396 self.epoch[ubatch : ubatch + 1],
1397 gathered_kv,
1398 self.world_size,
1399 self.rank,
1400 self.max_gathered_tokens,
1401 kv_multicast_ptr,
1402 signal_multicast_ptr,
1403 )
1404
1405
1406 @functools.cache
1407 def get_direct_dcp_kv_gather_workspace(
1408 group: GroupCoordinator,
1409 device: torch.device,
1410 max_gathered_tokens: int,
1411 token_dim: int,
1412 dtype: torch.dtype,
1413 num_ubatches: int,
1414 ) -> DirectDCPKVGatherWorkspace | None:
1415 if not direct_cp_multicast_enabled(
1416 group,
1417 dtype,
1418 envs.VLLM_USE_DIRECT_DCP_KV_GATHER,
1419 _KV_GATHER_SUPPORTED_DTYPES,
1420 ):
1421 return None
1422 if not _kv_gather_layout_supported(token_dim, dtype):
1423 return None
1424 return DirectDCPKVGatherWorkspace(
1425 group.device_group,
1426 device,
1427 max_gathered_tokens,
1428 token_dim,
1429 dtype,
1430 num_ubatches,
1431 )
1432
1433
1434 # MLA DCP backend selection
1435
1436
1437 class DCPCombine(Protocol):
1438 def __call__(
1439 self,
1440 partial_output: torch.Tensor,
1441 partial_lse: torch.Tensor,
1442 *,
1443 seq_lens: torch.Tensor,
1444 query_start_loc: torch.Tensor,
1445 ) -> torch.Tensor: ...
1446
1447
1448 class MLADCPManager:
1449 """Select and own layer-level collective implementations for MLA DCP."""
1450
1451 _kv_gather: Callable[[torch.Tensor, torch.Tensor], object]
1452
1453 def __init__(
1454 self,
1455 vllm_config: VllmConfig,
1456 device: torch.device,
1457 num_heads: int,
1458 query_head_dim: int,
1459 output_head_dim: int,
1460 query_dtype: torch.dtype,
1461 output_dtype: torch.dtype,
1462 padded_num_heads: int | None,
1463 is_lse_base_on_e: bool,
1464 use_pcp: bool,
1465 ) -> None:
1466 parallel_config = vllm_config.parallel_config
1467 self.group = get_dcp_group()
1468 self.device = torch.device(device)
1469 self.num_ubatches = max(parallel_config.num_ubatches, 1)
1470 self.max_num_tokens = get_dcp_workspace_max_num_tokens(vllm_config)
1471 self.use_a2a = parallel_config.dcp_comm_backend == "a2a"
1472 self.padded_num_heads = padded_num_heads
1473
1474 self.combine = self._init_combine(
1475 num_heads,
1476 output_head_dim,
1477 output_dtype,
1478 is_lse_base_on_e,
1479 use_pcp,
1480 )
1481 if vllm_config.kernel_config.enable_jit_warmup and not self.use_a2a:
1482 _CORRECT_ATTN_CP_OUT_KERNEL.register_warmup(
1483 vllm_config,
1484 output_dtype=output_dtype,
1485 num_heads=num_heads,
1486 head_dim=output_head_dim,
1487 is_base_e=is_lse_base_on_e,
1488 )
1489 self.query_gather = (
1490 None
1491 if use_pcp
1492 else self._init_query_gather(
1493 num_heads,
1494 query_head_dim,
1495 query_dtype,
1496 )
1497 )
1498
1499 def _init_combine(
1500 self,
1501 num_heads: int,
1502 head_dim: int,
1503 dtype: torch.dtype,
1504 is_lse_base_on_e: bool,
1505 use_pcp: bool,
1506 ) -> DCPCombine:
1507 direct_workspace = None
1508 if self.use_a2a:
1509 direct_workspace = get_direct_dcp_a2a_workspace(
1510 self.group,
1511 self.device,
1512 self.max_num_tokens,
1513 num_heads,
1514 head_dim,
1515 dtype,
1516 self.num_ubatches,
1517 )
1518 if direct_workspace is not None:
1519 logger.info_once("Using direct symmetric-memory DCP A2A for MLA.")
1520 return functools.partial(
1521 self._direct_workspace_combine,
1522 direct_workspace,
1523 is_lse_base_on_e=is_lse_base_on_e,
1524 )
1525
1526 combine_fn = (
1527 dcp_a2a_lse_reduce
1528 if self.use_a2a
1529 else cp_lse_ag_out_ar
1530 if use_pcp
1531 else cp_lse_ag_out_rs
1532 )
1533 return functools.partial(
1534 combine_fn,
1535 cp_group=self.group,
1536 is_lse_base_on_e=is_lse_base_on_e,
1537 )
1538
1539 def _direct_workspace_combine(
1540 self,
1541 direct_workspace: DirectDCPA2AWorkspace,
1542 partial_output: torch.Tensor,
1543 partial_lse: torch.Tensor,
1544 is_lse_base_on_e: bool,
1545 seq_lens: torch.Tensor | None = None,
1546 query_start_loc: torch.Tensor | None = None,
1547 ) -> torch.Tensor:
1548 # Forced MQA path pass all batch tokens (including prefill) into combine,
1549 # which may exceed the direct symmetric-memory workspace. Fall back to
1550 # the nccl a2a combine for those cases.
1551 if partial_output.shape[0] <= direct_workspace.max_num_tokens:
1552 return direct_workspace.lse_reduce(
1553 partial_output,
1554 partial_lse,
1555 is_lse_base_on_e,
1556 seq_lens,
1557 query_start_loc,
1558 )
1559 return dcp_a2a_lse_reduce(
1560 partial_output,
1561 partial_lse,
1562 cp_group=self.group,
1563 is_lse_base_on_e=is_lse_base_on_e,
1564 seq_lens=seq_lens,
1565 query_start_loc=query_start_loc,
1566 )
1567
1568 def _init_query_gather(
1569 self,
1570 num_heads: int,
1571 head_dim: int,
1572 dtype: torch.dtype,
1573 ) -> Callable[[torch.Tensor], torch.Tensor]:
1574 direct_workspace = get_direct_dcp_q_gather_workspace(
1575 self.group,
1576 self.device,
1577 self.max_num_tokens,
1578 num_heads,
1579 head_dim,
1580 dtype,
1581 self.num_ubatches,
1582 self.padded_num_heads,
1583 )
1584 if direct_workspace is not None:
1585 logger.info_once("Using direct symmetric-memory DCP query gather for MLA.")
1586 return functools.partial(
1587 self._direct_workspace_query_gather,
1588 direct_workspace,
1589 )
1590 return self._gather_query
1591
1592 def _gather_query(self, query: torch.Tensor) -> torch.Tensor:
1593 query = self.group.all_gather(query, dim=1)
1594 if self.padded_num_heads is not None:
1595 query = reserve_query_head_storage(query, self.padded_num_heads)
1596 return query
1597
1598 def _direct_workspace_query_gather(
1599 self,
1600 direct_workspace: DirectDCPQGatherWorkspace,
1601 query: torch.Tensor,
1602 ) -> torch.Tensor:
1603 # Forced MQA path can be taken for long sparse MLA prefills, whose batch size
1604 # may exceed the direct symmetric-memory workspace. Fall back to allgather
1605 # for those cases.
1606 if query.shape[0] <= direct_workspace.max_num_tokens:
1607 return direct_workspace.gather(query)
1608 return self._gather_query(query)
1609
1610 def init_kv_gather(
1611 self,
1612 workspace: torch.Tensor,
1613 max_gathered_tokens: int,
1614 ) -> None:
1615 world_size = self.group.world_size
1616 assert max_gathered_tokens > 0
1617 assert max_gathered_tokens % world_size == 0
1618 assert workspace.ndim == 2
1619 assert workspace.is_contiguous()
1620 assert workspace.shape[0] == (
1621 max_gathered_tokens + max_gathered_tokens // world_size
1622 )
1623 assert workspace.shape[1] > 0
1624
1625 direct_workspace = get_direct_dcp_kv_gather_workspace(
1626 self.group,
1627 workspace.device,
1628 max_gathered_tokens,
1629 workspace.shape[1],
1630 workspace.dtype,
1631 self.num_ubatches,
1632 )
1633 if direct_workspace is not None:
1634 logger.info_once(
1635 "Using direct symmetric-memory DCP chunked-context KV gather for MLA."
1636 )
1637 self._kv_gather = direct_workspace.gather
1638 else:
1639 self._kv_gather = functools.partial(
1640 torch.distributed.all_gather_into_tensor,
1641 group=self.group.device_group,
1642 )
1643
1644 def kv_gather(
1645 self,
1646 gathered_kv: torch.Tensor,
1647 local_kv: torch.Tensor,
1648 ) -> object:
1649 return self._kv_gather(gathered_kv, local_kv)
1650
1651
1652 _CORRECT_ATTN_CP_OUT_KERNEL = CorrectAttnCPOutKernel()