diff --git a/tests/distributed/test_dcp_a2a.py b/tests/distributed/test_dcp_a2a.py index 5ab0f3de97b5..7d46ddac7967 100644 --- a/tests/distributed/test_dcp_a2a.py +++ b/tests/distributed/test_dcp_a2a.py @@ -27,6 +27,7 @@ class _FakeCPGroup: def __init__(self, world_size: int, device_group: dist.ProcessGroup): self.world_size = world_size self.device_group = device_group + self.rank_in_group = dist.get_rank(device_group) def _dtype_from_name(dtype_name: str) -> torch.dtype: @@ -375,6 +376,57 @@ def test_pack_unpack_combine_matches_reference( else: _assert_packed_a2a_close(actual, expected_out, dtype) + @pytest.mark.skipif( + torch.accelerator.device_count() < 1, reason="CUDA is required." + ) + def test_pack_send_zeroes_empty_local_rows(self): + from vllm.v1.attention.ops.dcp_alltoall import ( + _dcp_a2a_lse_pack_dim, + _dcp_a2a_pack_send, + ) + + device = torch.device("cuda") + world_size, B, h_per_rank, D = 4, 5, 2, 32 + H = world_size * h_per_rank + cp_attn_out = torch.randn(B, H, D, device=device) + cp_attn_lse = torch.randn(B, H, device=device) + valid_counts = torch.tensor([3, 0, 1, 0, 2], device=device) + lse_pack_dim = _dcp_a2a_lse_pack_dim(cp_attn_out.dtype) + send_buffer = torch.empty( + (world_size, B, h_per_rank, D + lse_pack_dim), + device=device, + ) + + _dcp_a2a_pack_send( + cp_attn_out, + cp_attn_lse, + send_buffer, + world_size, + h_per_rank, + D, + lse_pack_dim, + valid_counts=valid_counts, + ) + torch.accelerator.synchronize() + + empty_rows = valid_counts == 0 + non_empty_rows = ~empty_rows + expected_out = ( + cp_attn_out.view(B, world_size, h_per_rank, D) + .permute(1, 0, 2, 3) + .contiguous() + ) + empty_payload = send_buffer[:, empty_rows, :, :D] + torch.testing.assert_close(empty_payload, torch.zeros_like(empty_payload)) + torch.testing.assert_close( + send_buffer[:, non_empty_rows, :, :D], + expected_out[:, non_empty_rows], + ) + torch.testing.assert_close( + send_buffer[:, empty_rows, :, D], + torch.full_like(send_buffer[:, empty_rows, :, D], float("-inf")), + ) + def _distributed_packed_a2a_worker(env: dict[str, str]) -> None: update_environment_variables(env) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index b8f197b31614..f7b80449c06f 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -513,6 +513,10 @@ def __init__( ) # Initialize q/k/v range constants. + # Project attention output through W_UV before the DCP merge: shrinks + # the merge payload from kv_lora_rank to v_head_dim per head. + self.W_UV_dcp: torch.Tensor | None = None + self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32) self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32) self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32) @@ -811,25 +815,7 @@ def forward_impl( assert attn_metadata.decode is not None attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined] - # correct dcp attn_out with lse. - if self.impl.dcp_world_size > 1: - if self.dcp_a2a: - attn_out = dcp_a2a_lse_reduce( - attn_out, - lse, - get_dcp_group(), - is_lse_base_on_e=self.impl.lse_base_on_e, - ) - else: - attn_out = cp_lse_ag_out_rs( - attn_out, - lse, - get_dcp_group(), - is_lse_base_on_e=self.impl.lse_base_on_e, - ) - - # v_up projection - self._v_up_proj(attn_out, out=mqa_output_slice) + self._dcp_merge_and_v_up_proj(attn_out, lse, mqa_output_slice) if quant_key is not None: quant_idx = num_mqa_tokens if mha_use_quant_output else num_actual_toks @@ -958,6 +944,11 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): else: # Convert from (L, N, V) to (N, L, V) self.W_UV = W_UV.transpose(0, 1) + if getattr(self.impl, "dcp_world_size", 1) > 1: + # all_gather_into_tensor requires a contiguous input + self.W_UV_dcp = get_dcp_group().all_gather( + self.W_UV.contiguous(), dim=0 + ) # Convert from (L, N, P) to (N, P, L) self.W_UK_T = W_UK.permute(1, 2, 0) @@ -1009,6 +1000,57 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: cache_dtype_str=vllm_config.cache_config.cache_dtype, ) + def _dcp_lse_merge( + self, + attn_out: torch.Tensor, + lse: torch.Tensor, + out: torch.Tensor | None, + ) -> torch.Tensor: + """LSE-weighted combine of the per-rank attention outputs across the + DCP group. The a2a transport can write directly into ``out``.""" + if self.dcp_a2a: + return dcp_a2a_lse_reduce( + attn_out, + lse, + get_dcp_group(), + is_lse_base_on_e=self.impl.lse_base_on_e, + out=out, + valid_counts=getattr(self.impl, "_last_dcp_valid_counts", None), + ) + return cp_lse_ag_out_rs( + attn_out, + lse, + get_dcp_group(), + is_lse_base_on_e=self.impl.lse_base_on_e, + ) + + def _dcp_merge_and_v_up_proj( + self, + attn_out: torch.Tensor, + lse: torch.Tensor, + out: torch.Tensor, + ) -> None: + """Combine the decode attention output across DCP ranks (if any) and + apply the W_UV up-projection into ``out`` (flattened v_head_dim).""" + if self.impl.dcp_world_size > 1 and self.W_UV_dcp is not None: + # Project kv_lora_rank -> v_head_dim BEFORE the merge to halve the + # DCP exchange payload; the LSE-weighted merge commutes with the + # linear W_UV projection, so this is exact. + projected = attn_out.new_empty( + attn_out.shape[0], attn_out.shape[1], self.v_head_dim + ) + self._v_up_proj_bmm(attn_out, projected, self.W_UV_dcp) + out_view = out.view(-1, self.num_heads, self.v_head_dim) + merged = self._dcp_lse_merge(projected, lse, out=out_view) + if merged is not out_view: + out.copy_(merged.reshape(out.shape)) + return + + # Project after the merge (dcp=1, or backends without a gathered W_UV). + if self.impl.dcp_world_size > 1: + attn_out = self._dcp_lse_merge(attn_out, lse, out=None) + self._v_up_proj(attn_out, out=out) + def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor): # Convert from (B, N, L) to (N, B, L) x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1) @@ -1033,6 +1075,14 @@ def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor): # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V) torch.bmm(x, self.W_UV, out=out.transpose(0, 1)) + def _v_up_proj_bmm( + self, x: torch.Tensor, out: torch.Tensor, w_uv: torch.Tensor + ) -> None: + num_heads = w_uv.shape[0] + x = x.view(-1, num_heads, self.kv_lora_rank).transpose(0, 1) + out = out.view(-1, num_heads, self.v_head_dim) + torch.bmm(x, w_uv, out=out.transpose(0, 1)) + def unified_mla_kv_cache_update( kv_c_normed: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index b66952cc74be..54443734bd2c 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """FlashInfer sparse MLA attention backend.""" +import os from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar @@ -9,13 +10,14 @@ import torch from vllm import envs -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config_or_none from vllm.config.cache import CacheDType from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( get_mla_dims, ) from vllm.platforms.interface import DeviceCapability +from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor from vllm.v1.attention.backend import ( AttentionBackend, @@ -43,6 +45,108 @@ logger = init_logger(__name__) +_FLASHINFER_SPARSE_MLA_WORKSPACE_SLOP = 64 * 1024 * 1024 + + +@triton.jit +def _zero_empty_sparse_mla_rows_kernel( + out_ptr, + lse_ptr, + seq_lens_ptr, + out_stride_t, + out_stride_h, + out_stride_d, + lse_stride_t, + lse_stride_h, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + HAS_LSE: tl.constexpr, +) -> None: + token_idx = tl.program_id(0) + head_idx = tl.program_id(1) + offsets = tl.arange(0, BLOCK_D) + + is_empty = tl.load(seq_lens_ptr + token_idx) == 0 + out_offsets = ( + token_idx * out_stride_t + head_idx * out_stride_h + offsets * out_stride_d + ) + tl.store(out_ptr + out_offsets, 0.0, mask=is_empty & (offsets < HEAD_DIM)) + + if HAS_LSE: + lse_offset = token_idx * lse_stride_t + head_idx * lse_stride_h + tl.store(lse_ptr + lse_offset, -float("inf"), mask=is_empty) + + +def _zero_empty_sparse_mla_rows( + out: torch.Tensor, + lse: torch.Tensor | None, + seq_lens: torch.Tensor, +) -> None: + block_d = triton.next_power_of_2(out.shape[2]) + _zero_empty_sparse_mla_rows_kernel[(out.shape[0], out.shape[1])]( + out, + lse, + seq_lens, + out.stride(0), + out.stride(1), + out.stride(2), + lse.stride(0) if lse is not None else 0, + lse.stride(1) if lse is not None else 0, + HEAD_DIM=out.shape[2], + BLOCK_D=block_d, + HAS_LSE=lse is not None, + ) + + +@triton.jit +def _sanitize_empty_sparse_mla_rows_kernel( + block_tables_ptr, + seq_lens_ptr, + valid_counts_ptr, + block_tables_stride_b, + block_tables_stride_k, + NUM_TOPK_TOKENS: tl.constexpr, + NUM_BLOCKS: tl.constexpr, + BLOCK_K: tl.constexpr, +) -> None: + token_idx = tl.program_id(0) + k_block = tl.program_id(1) + offsets = k_block * BLOCK_K + tl.arange(0, BLOCK_K) + mask = offsets < NUM_TOPK_TOKENS + block_offsets = token_idx * block_tables_stride_b + offsets * block_tables_stride_k + block_ids = tl.load(block_tables_ptr + block_offsets, mask=mask, other=0) + tl.store( + block_tables_ptr + block_offsets, + tl.where((block_ids < 0) | (block_ids >= NUM_BLOCKS), 0, block_ids), + mask=mask, + ) + + count = tl.load(seq_lens_ptr + token_idx) + is_first_block = k_block == 0 + tl.store(valid_counts_ptr + token_idx, count, mask=is_first_block) + tl.store(seq_lens_ptr + token_idx, 1, mask=is_first_block & (count == 0)) + + +def _sanitize_empty_sparse_mla_rows( + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + num_blocks: int, +) -> torch.Tensor: + valid_counts = torch.empty_like(seq_lens) + block_k = min(triton.next_power_of_2(block_tables.shape[1]), 1024) + grid = (seq_lens.shape[0], triton.cdiv(block_tables.shape[1], block_k)) + _sanitize_empty_sparse_mla_rows_kernel[grid]( + block_tables, + seq_lens, + valid_counts, + block_tables.stride(0), + block_tables.stride(1), + NUM_TOPK_TOKENS=block_tables.shape[1], + NUM_BLOCKS=num_blocks, + BLOCK_K=block_k, + ) + return valid_counts + class _FlashInferMLASparseBackendBase(AttentionBackend): """Common metadata for concrete FlashInfer sparse MLA backends.""" @@ -352,13 +456,66 @@ def build( _fi_sparse_workspace: torch.Tensor | None = None -def _get_workspace_buffer(device: torch.device) -> torch.Tensor: +def _round_up_to_mib(size: int) -> int: + mib = 1024 * 1024 + return ((size + mib - 1) // mib) * mib + + +def _infer_workspace_buffer_size( + num_heads: int, + topk_indices_buffer: torch.Tensor | None, + return_lse: bool, + workspace_multiplier: int = 1, + min_decode_tokens: int = 1, +) -> int: + configured_size = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + if "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE" in os.environ: + return configured_size + + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return configured_size + + topk_tokens = ( + topk_indices_buffer.shape[1] + if topk_indices_buffer is not None + else getattr(vllm_config.model_config.hf_config, "index_topk", 2048) + ) + capture_size = ( + vllm_config.compilation_config.max_cudagraph_capture_size + or vllm_config.scheduler_config.max_num_seqs + ) + max_decode_tokens = max( + min(capture_size, vllm_config.scheduler_config.max_num_seqs), + min_decode_tokens, + ) + + # TRTLLM-gen sparse MLA allocates a softmax workspace proportional to the + # captured decode batch, sparse top-k, and local query heads. + lse_multiplier = 2 if return_lse else 1 + required_size = ( + int(max_decode_tokens) + * int(topk_tokens) + * int(num_heads) + * 256 + * lse_multiplier + * int(workspace_multiplier) + + _FLASHINFER_SPARSE_MLA_WORKSPACE_SLOP + ) + return max(configured_size, _round_up_to_mib(required_size)) + + +def _get_workspace_buffer( + device: torch.device, + min_size: int | None = None, +) -> torch.Tensor: global _fi_sparse_workspace - if _fi_sparse_workspace is None: + min_size = min_size or envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + if _fi_sparse_workspace is None or _fi_sparse_workspace.numel() < min_size: # FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace; # the trtllm-gen path views it as uint8, so int8 is safe for all backends. _fi_sparse_workspace = torch.zeros( - envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE, + min_size, dtype=torch.int8, device=device, ) @@ -426,8 +583,30 @@ def __init__( ) self._workspace_buffer: torch.Tensor | None = None + # DCP all-gathers the query heads before forward_mqa. + self._workspace_buffer_size = _infer_workspace_buffer_size( + self.num_heads * max(self.dcp_world_size, 1), + self.topk_indices_buffer, + self.need_to_return_lse_for_decode, + workspace_multiplier=2 if self.dcp_world_size > 1 else 1, + min_decode_tokens=64 if self.dcp_world_size > 1 else 1, + ) + if ( + self.topk_indices_buffer is not None + and self.topk_indices_buffer.device.type == "cuda" + ): + self._workspace_buffer = _get_workspace_buffer( + self.topk_indices_buffer.device, self._workspace_buffer_size + ) self.bmm1_scale: float | None = None self.bmm2_scale: float | None = None + vllm_config = get_current_vllm_config_or_none() + self._zero_empty_in_a2a_pack = ( + self.dcp_world_size > 1 + and vllm_config is not None + and vllm_config.parallel_config.dcp_comm_backend == "a2a" + ) + self._last_dcp_valid_counts: torch.Tensor | None = None # fp8 query quantization is required when using fp8 kv_cache, # as the TRTLLM-GEN sparse MLA kernel requires matching dtypes @@ -470,9 +649,20 @@ def forward_mqa( NUM_TOPK_TOKENS=topk_indices.shape[1], return_valid_counts=True, ) + self._last_dcp_valid_counts = None + if self.dcp_world_size > 1: + valid_counts = _sanitize_empty_sparse_mla_rows( + topk_indices_physical, + seq_lens, + kv_c_and_k_pe_cache.shape[0], + ) + if self._zero_empty_in_a2a_pack: + self._last_dcp_valid_counts = valid_counts if self._workspace_buffer is None: - self._workspace_buffer = _get_workspace_buffer(q.device) + self._workspace_buffer = _get_workspace_buffer( + q.device, self._workspace_buffer_size + ) if self.bmm1_scale is None: self.bmm1_scale = self.scale @@ -520,9 +710,8 @@ def forward_mqa( out = o.view(-1, o.shape[-2], o.shape[-1]) if lse is not None: lse = self._normalize_lse(lse, out.shape[0], out.shape[1]) - empty_rows = (topk_indices_physical == -1).all(dim=-1) - out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) - lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) + if not self._zero_empty_in_a2a_pack: + _zero_empty_sparse_mla_rows(out, lse, seq_lens) return out, lse @staticmethod diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 1e12f43caacb..d9cf15181482 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -904,7 +904,7 @@ def get_dcp_local_seq_lens( ) seq_lens_tiled = seq_lens_i32.unsqueeze(-1) else: - rank_offsets = torch.tensor(dcp_rank, dtype=torch.int32, device=seq_lens.device) + rank_offsets = dcp_rank seq_lens_tiled = seq_lens_i32 base = ( seq_lens_tiled diff --git a/vllm/v1/attention/ops/dcp_alltoall.py b/vllm/v1/attention/ops/dcp_alltoall.py index e100dbc79ee2..7a0926488520 100644 --- a/vllm/v1/attention/ops/dcp_alltoall.py +++ b/vllm/v1/attention/ops/dcp_alltoall.py @@ -25,12 +25,16 @@ import torch import torch.distributed as dist +from vllm.config import get_current_vllm_config_or_none +from vllm.logger import init_logger from vllm.triton_utils import tl, triton if TYPE_CHECKING: from vllm.distributed.parallel_state import GroupCoordinator from vllm.v1.attention.ops.common import CPTritonContext +logger = init_logger(__name__) + def _lse_weighted_combine( outputs: torch.Tensor, @@ -108,31 +112,199 @@ def _dcp_a2a_lse_pack_dim(output_dtype: torch.dtype) -> int: raise ValueError(f"Cannot pack fp32 LSE into output dtype {output_dtype}.") +_dcp_a2a_buffer_cache: dict[ + tuple[torch.device, torch.dtype], + tuple[torch.Tensor, torch.Tensor], +] = {} +_dcp_a2a_fi_send_buffer_cache: dict[ + tuple[torch.device, torch.dtype], + tuple[torch.Tensor, torch.Tensor], +] = {} +_dcp_a2a_fi_workspace_cache: dict[tuple[int, torch.device], torch.Tensor] = {} +_flashinfer_dcp_a2a_supported: bool | None = None + + +def _flashinfer_dcp_a2a_available() -> bool: + global _flashinfer_dcp_a2a_supported + if _flashinfer_dcp_a2a_supported is None: + try: + import flashinfer.comm as flashinfer_comm + + _flashinfer_dcp_a2a_supported = hasattr( + flashinfer_comm, "decode_cp_a2a_alltoall" + ) + except ImportError: + _flashinfer_dcp_a2a_supported = False + return _flashinfer_dcp_a2a_supported + + +_dcp_a2a_fi_workspace_memory_cache: dict[tuple[int, torch.device], object] = {} + + def _dcp_a2a_send_recv_buffers( shape: tuple[int, ...], device: torch.device, dtype: torch.dtype, ) -> tuple[torch.Tensor, torch.Tensor]: - # Don't use the shared WorkspaceManager here. A FULL cudagraph bakes in the - # buffer address at capture, but the workspace is growable and sized only to - # the largest *captured* batch (the cudagraph capture cap). Any eager a2a - # with a bigger batch regrows it, freeing that address and poisoning every - # captured graph -> illegal memory access on replay. This bites the very - # first request: the post-capture warmup runs an eager decode at - # max_num_seqs (> the cap), so the graphs are already dangling before the - # server is ready. torch.empty buffers instead live in the graph's private - # pool and stay valid for its lifetime (as _dcp_a2a_unpack_combine and the - # AG+RS combine path already rely on). + # FULL cudagraph replay needs stable addresses, while the first eager + # prefill/decode warmup can use many more tokens than the captured decode + # graph. Allocate flat max-token buffers once so memory profiling accounts + # for the largest configured A2A payload and graph addresses stay stable. + requested_numel = 1 + for dim in shape: + requested_numel *= dim + + alloc_shape = list(shape) + vllm_config = get_current_vllm_config_or_none() + if vllm_config is not None and len(alloc_shape) >= 2: + alloc_shape[1] = max( + alloc_shape[1], + int(vllm_config.scheduler_config.max_num_batched_tokens), + ) + + alloc_numel = 1 + for dim in alloc_shape: + alloc_numel *= dim + + cache_key = (device, dtype) + buffers = _dcp_a2a_buffer_cache.get(cache_key) + if buffers is None or buffers[0].numel() < alloc_numel: + buffers = ( + torch.empty(alloc_numel, device=device, dtype=dtype), + torch.empty(alloc_numel, device=device, dtype=dtype), + ) + _dcp_a2a_buffer_cache[cache_key] = buffers + return ( - torch.empty(shape, device=device, dtype=dtype), - torch.empty(shape, device=device, dtype=dtype), + buffers[0][:requested_numel].view(shape), + buffers[1][:requested_numel].view(shape), ) +def _dcp_a2a_fi_send_buffers( + partial_o_shape: tuple[int, ...], + softmax_stats_shape: tuple[int, ...], + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + alloc_partial_o_shape = list(partial_o_shape) + alloc_softmax_stats_shape = list(softmax_stats_shape) + vllm_config = get_current_vllm_config_or_none() + if vllm_config is not None and len(alloc_partial_o_shape) >= 1: + max_tokens = int(vllm_config.scheduler_config.max_num_batched_tokens) + alloc_partial_o_shape[0] = max(alloc_partial_o_shape[0], max_tokens) + alloc_softmax_stats_shape[0] = max(alloc_softmax_stats_shape[0], max_tokens) + + requested_partial_o_numel = 1 + for dim in partial_o_shape: + requested_partial_o_numel *= dim + requested_softmax_stats_numel = 1 + for dim in softmax_stats_shape: + requested_softmax_stats_numel *= dim + + alloc_partial_o_numel = 1 + for dim in alloc_partial_o_shape: + alloc_partial_o_numel *= dim + alloc_softmax_stats_numel = 1 + for dim in alloc_softmax_stats_shape: + alloc_softmax_stats_numel *= dim + + cache_key = (device, dtype) + buffers = _dcp_a2a_fi_send_buffer_cache.get(cache_key) + if ( + buffers is None + or buffers[0].numel() < alloc_partial_o_numel + or buffers[1].numel() < alloc_softmax_stats_numel + ): + buffers = ( + torch.empty(alloc_partial_o_numel, device=device, dtype=dtype), + torch.empty( + alloc_softmax_stats_numel, + device=device, + dtype=torch.float32, + ), + ) + _dcp_a2a_fi_send_buffer_cache[cache_key] = buffers + + return ( + buffers[0][:requested_partial_o_numel].view(partial_o_shape), + buffers[1][:requested_softmax_stats_numel].view(softmax_stats_shape), + ) + + +def _dcp_a2a_fi_workspace( + cp_group: GroupCoordinator, + device: torch.device, +) -> torch.Tensor: + cache_key = (id(cp_group.device_group), device) + workspace = _dcp_a2a_fi_workspace_cache.get(cache_key) + if workspace is not None: + return workspace + + try: + import flashinfer.comm as flashinfer_comm + from flashinfer.comm.mnnvl import MnnvlConfig, MnnvlMemory, TorchDistBackend + except ImportError as err: + raise RuntimeError( + "The FlashInfer DCP A2A path requires flashinfer.comm " + "with decode_cp_a2a support." + ) from err + + if not hasattr(flashinfer_comm, "decode_cp_a2a_alltoall"): + raise RuntimeError( + "The FlashInfer DCP A2A path requires " + "flashinfer.comm.decode_cp_a2a_alltoall." + ) + + comm_mapping = flashinfer_comm.Mapping( + world_size=cp_group.world_size, + rank=cp_group.rank_in_group, + gpus_per_node=torch.accelerator.device_count(), + cp_size=1, + tp_size=cp_group.world_size, + pp_size=1, + ) + workspace_mapping = flashinfer_comm.Mapping( + world_size=cp_group.world_size, + rank=cp_group.rank_in_group, + gpus_per_node=torch.accelerator.device_count(), + cp_size=cp_group.world_size, + tp_size=1, + pp_size=1, + ) + config = MnnvlConfig( + comm_backend=TorchDistBackend(group=cp_group.device_group), + fabric_page_size=1 << 29, + allocation_granularity=0, + ) + # FlashInfer's MNNVL helper splits its communicator by CP rank and uses TP + # rank ordering. Seed it with a TP-shaped mapping so our DCP group shares + # one workspace address space, then allocate bytes sized for real DCP. + MnnvlMemory.initialize() + MnnvlMemory.set_comm_from_config(comm_mapping, config) + workspace_bytes = flashinfer_comm.decode_cp_a2a_workspace_size(cp_group.world_size) + mnnvl_memory = MnnvlMemory(workspace_mapping, workspace_bytes) + workspace = mnnvl_memory.as_torch_strided_tensor(torch.int64) + flashinfer_comm.decode_cp_a2a_init_workspace( + workspace, + cp_group.rank_in_group, + cp_group.world_size, + ) + dist.barrier(group=cp_group.device_group) + _dcp_a2a_fi_workspace_cache[cache_key] = workspace + _dcp_a2a_fi_workspace_memory_cache[cache_key] = mnnvl_memory + logger.info_once( + "Initialized FlashInfer decode CP A2A workspace for DCP size %d", + cp_group.world_size, + ) + return workspace + + @triton.jit def _dcp_a2a_pack_send_kernel( out_ptr, lse_ptr, + valid_counts_ptr, send_ptr, out_stride_B, out_stride_H, @@ -147,10 +319,14 @@ def _dcp_a2a_pack_send_kernel( HEAD_DIM: tl.constexpr, H_PER_RANK: tl.constexpr, LSE_PACK_DIM: tl.constexpr, + HAS_VALID_COUNTS: tl.constexpr, ): batch_idx = tl.program_id(0).to(tl.int64) local_head_idx = tl.program_id(1).to(tl.int64) d_offsets = tl.arange(0, HEAD_DIM) + has_values = True + if HAS_VALID_COUNTS: + has_values = tl.load(valid_counts_ptr + batch_idx) != 0 for rank_idx in tl.static_range(N): src_head_idx = rank_idx * H_PER_RANK + local_head_idx @@ -167,11 +343,13 @@ def _dcp_a2a_pack_send_kernel( ) tl.store( send_ptr + send_base + d_offsets * send_stride_D, - tl.load(out_ptr + out_offsets), + tl.where(has_values, tl.load(out_ptr + out_offsets), 0.0), ) - lse_val = tl.load( - lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H + lse_val = tl.where( + has_values, + tl.load(lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H), + -float("inf"), ) if LSE_PACK_DIM == 1: tl.store( @@ -192,6 +370,69 @@ def _dcp_a2a_pack_send_kernel( ) +@triton.jit +def _dcp_a2a_fi_pack_send_kernel( + out_ptr, + lse_ptr, + valid_counts_ptr, + partial_o_ptr, + softmax_stats_ptr, + out_stride_B, + out_stride_H, + out_stride_D, + lse_stride_B, + lse_stride_H, + partial_o_stride_B, + partial_o_stride_H, + partial_o_stride_N, + partial_o_stride_D, + stats_stride_B, + stats_stride_H, + stats_stride_N, + stats_stride_S, + N: tl.constexpr, + HEAD_DIM: tl.constexpr, + H_PER_RANK: tl.constexpr, + HAS_VALID_COUNTS: tl.constexpr, +): + batch_idx = tl.program_id(0).to(tl.int64) + local_head_idx = tl.program_id(1).to(tl.int64) + rank_idx = tl.program_id(2).to(tl.int64) + d_offsets = tl.arange(0, HEAD_DIM) + has_values = True + if HAS_VALID_COUNTS: + has_values = tl.load(valid_counts_ptr + batch_idx) != 0 + + src_head_idx = rank_idx * H_PER_RANK + local_head_idx + out_offsets = ( + batch_idx * out_stride_B + + src_head_idx * out_stride_H + + d_offsets * out_stride_D + ) + partial_o_base = ( + batch_idx * partial_o_stride_B + + local_head_idx * partial_o_stride_H + + rank_idx * partial_o_stride_N + ) + tl.store( + partial_o_ptr + partial_o_base + d_offsets * partial_o_stride_D, + tl.where(has_values, tl.load(out_ptr + out_offsets), 0.0), + ) + + lse_val = tl.where( + has_values, + tl.load(lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H), + -float("inf"), + ) + stats_base = ( + batch_idx * stats_stride_B + + local_head_idx * stats_stride_H + + rank_idx * stats_stride_N + ) + tl.store(softmax_stats_ptr + stats_base, lse_val) + tl.store(softmax_stats_ptr + stats_base + stats_stride_S, 0.0) + + @triton.jit def _dcp_a2a_unpack_combine_kernel( recv_ptr, @@ -316,6 +557,114 @@ def _dcp_a2a_unpack_combine_kernel( tl.store(out_lse_ptr + out_lse_offset, global_lse) +@triton.jit +def _dcp_a2a_fi_unpack_combine_kernel( + partial_o_ptr, + softmax_stats_ptr, + out_ptr, + out_lse_ptr, + partial_o_stride_B, + partial_o_stride_H, + partial_o_stride_N, + partial_o_stride_D, + stats_stride_B, + stats_stride_H, + stats_stride_N, + stats_stride_S, + out_stride_B, + out_stride_H, + out_stride_D, + out_lse_stride_B, + out_lse_stride_H, + N: tl.constexpr, + HEAD_DIM: tl.constexpr, + IS_BASE_E: tl.constexpr, + RETURN_LSE: tl.constexpr, +): + batch_idx = tl.program_id(0).to(tl.int64) + head_idx = tl.program_id(1).to(tl.int64) + d_offsets = tl.arange(0, HEAD_DIM) + + lse_max = -float("inf") + for rank_idx in tl.static_range(N): + stats_base = ( + batch_idx * stats_stride_B + + head_idx * stats_stride_H + + rank_idx * stats_stride_N + ) + lse_val = tl.load(softmax_stats_ptr + stats_base).to(tl.float32) + lse_val = tl.where( + (lse_val != lse_val) | (lse_val == float("inf")), + -float("inf"), + lse_val, + ) + lse_max = tl.maximum(lse_max, lse_val) + + lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max) + + lse_sum = 0.0 + for rank_idx in tl.static_range(N): + stats_base = ( + batch_idx * stats_stride_B + + head_idx * stats_stride_H + + rank_idx * stats_stride_N + ) + lse_val = tl.load(softmax_stats_ptr + stats_base).to(tl.float32) + lse_val = tl.where( + (lse_val != lse_val) | (lse_val == float("inf")), + -float("inf"), + lse_val, + ) + if IS_BASE_E: + lse_sum += tl.exp(lse_val - lse_max) + else: + lse_sum += tl.exp2(lse_val - lse_max) + + if IS_BASE_E: # noqa: SIM108 + global_lse = tl.log(lse_sum) + lse_max + else: + global_lse = tl.log2(lse_sum) + lse_max + + acc = tl.zeros([HEAD_DIM], dtype=tl.float32) + for rank_idx in tl.static_range(N): + stats_base = ( + batch_idx * stats_stride_B + + head_idx * stats_stride_H + + rank_idx * stats_stride_N + ) + lse_val = tl.load(softmax_stats_ptr + stats_base).to(tl.float32) + lse_val = tl.where( + (lse_val != lse_val) | (lse_val == float("inf")), + -float("inf"), + lse_val, + ) + if IS_BASE_E: + weight = tl.exp(lse_val - global_lse) + else: + weight = tl.exp2(lse_val - global_lse) + weight = tl.where(weight != weight, 0.0, weight) + partial_o_base = ( + batch_idx * partial_o_stride_B + + head_idx * partial_o_stride_H + + rank_idx * partial_o_stride_N + ) + acc += ( + tl.load(partial_o_ptr + partial_o_base + d_offsets * partial_o_stride_D).to( + tl.float32 + ) + * weight + ) + + final_offsets = ( + batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D + ) + tl.store(out_ptr + final_offsets, acc) + + if RETURN_LSE: + out_lse_offset = batch_idx * out_lse_stride_B + head_idx * out_lse_stride_H + tl.store(out_lse_ptr + out_lse_offset, global_lse) + + def _dcp_a2a_pack_send( cp_attn_out: torch.Tensor, cp_attn_lse: torch.Tensor, @@ -324,11 +673,13 @@ def _dcp_a2a_pack_send( h_per_rank: int, head_dim: int, lse_pack_dim: int, + valid_counts: torch.Tensor | None = None, ) -> None: grid = (cp_attn_out.shape[0], h_per_rank, 1) _dcp_a2a_pack_send_kernel[grid]( cp_attn_out, cp_attn_lse, + valid_counts, send_buffer, cp_attn_out.stride(0), cp_attn_out.stride(1), @@ -343,6 +694,44 @@ def _dcp_a2a_pack_send( HEAD_DIM=head_dim, H_PER_RANK=h_per_rank, LSE_PACK_DIM=lse_pack_dim, + HAS_VALID_COUNTS=valid_counts is not None, + ) + + +def _dcp_a2a_fi_pack_send( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + world_size: int, + h_per_rank: int, + head_dim: int, + valid_counts: torch.Tensor | None = None, +) -> None: + grid = (cp_attn_out.shape[0], h_per_rank, world_size) + _dcp_a2a_fi_pack_send_kernel[grid]( + cp_attn_out, + cp_attn_lse, + valid_counts, + partial_o, + softmax_stats, + cp_attn_out.stride(0), + cp_attn_out.stride(1), + cp_attn_out.stride(2), + cp_attn_lse.stride(0), + cp_attn_lse.stride(1), + partial_o.stride(0), + partial_o.stride(1), + partial_o.stride(2), + partial_o.stride(3), + softmax_stats.stride(0), + softmax_stats.stride(1), + softmax_stats.stride(2), + softmax_stats.stride(3), + N=world_size, + HEAD_DIM=head_dim, + H_PER_RANK=h_per_rank, + HAS_VALID_COUNTS=valid_counts is not None, ) @@ -352,13 +741,20 @@ def _dcp_a2a_unpack_combine( lse_pack_dim: int, return_lse: bool, is_lse_base_on_e: bool, + out: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: world_size, num_tokens, h_per_rank, _ = recv_buffer.shape - out = torch.empty( - (num_tokens, h_per_rank, head_dim), - device=recv_buffer.device, - dtype=recv_buffer.dtype, - ) + if out is None: + out = torch.empty( + (num_tokens, h_per_rank, head_dim), + device=recv_buffer.device, + dtype=recv_buffer.dtype, + ) + elif out.shape != (num_tokens, h_per_rank, head_dim): + raise ValueError( + "Invalid DCP A2A output shape: " + f"{tuple(out.shape)}, expected {(num_tokens, h_per_rank, head_dim)}." + ) out_lse = torch.empty( (num_tokens, h_per_rank) if return_lse else (1, 1), device=recv_buffer.device, @@ -389,42 +785,148 @@ def _dcp_a2a_unpack_combine( return out -def dcp_a2a_lse_reduce( +def _dcp_a2a_fi_unpack_combine( + partial_o: torch.Tensor, + softmax_stats: torch.Tensor, + head_dim: int, + return_lse: bool, + is_lse_base_on_e: bool, + device: torch.device, + dtype: torch.dtype, + out: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + num_tokens, h_per_rank, world_size, _ = partial_o.shape + if out is None: + out = torch.empty( + (num_tokens, h_per_rank, head_dim), + device=device, + dtype=dtype, + ) + elif out.shape != (num_tokens, h_per_rank, head_dim): + raise ValueError( + "Invalid DCP A2A output shape: " + f"{tuple(out.shape)}, expected {(num_tokens, h_per_rank, head_dim)}." + ) + out_lse = torch.empty( + (num_tokens, h_per_rank) if return_lse else (1, 1), + device=device, + dtype=torch.float32 if return_lse else dtype, + ) + grid = (num_tokens, h_per_rank, 1) + _dcp_a2a_fi_unpack_combine_kernel[grid]( + partial_o, + softmax_stats, + out, + out_lse, + partial_o.stride(0), + partial_o.stride(1), + partial_o.stride(2), + partial_o.stride(3), + softmax_stats.stride(0), + softmax_stats.stride(1), + softmax_stats.stride(2), + softmax_stats.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out_lse.stride(0), + out_lse.stride(1), + N=world_size, + HEAD_DIM=head_dim, + IS_BASE_E=is_lse_base_on_e, + RETURN_LSE=return_lse, + ) + if return_lse: + return out, out_lse + return out + + +def _as_torch_tensor(tensor: object) -> torch.Tensor: + if torch.is_tensor(tensor): + return tensor + return torch.utils.dlpack.from_dlpack(tensor) + + +def _dcp_a2a_flashinfer_lse_reduce( cp_attn_out: torch.Tensor, cp_attn_lse: torch.Tensor, cp_group: GroupCoordinator, - ctx: CPTritonContext | None = None, - return_lse: bool = False, - is_lse_base_on_e: bool = True, + return_lse: bool, + is_lse_base_on_e: bool, + out: torch.Tensor | None = None, + valid_counts: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - Combine partial attention outputs across DCP ranks using All-to-All. - - The output and fp32 LSE are packed into a single output-dtype buffer, sent - with one All-to-All, then unpacked and combined with exact LSE weighting. + if cp_attn_out.dtype not in (torch.float16, torch.bfloat16): + logger.warning_once( + "FlashInfer DCP A2A supports fp16/bf16 outputs; falling back to " + "Triton+NCCL for dtype %s.", + cp_attn_out.dtype, + ) + return _dcp_a2a_triton_lse_reduce( + cp_attn_out, + cp_attn_lse, + cp_group, + return_lse=return_lse, + is_lse_base_on_e=is_lse_base_on_e, + out=out, + valid_counts=valid_counts, + ) - Args: - cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim - cp_attn_lse: [B, H] log-sum-exp values (fp32) - cp_group: GroupCoordinator for DCP communication - ctx: CPTritonContext (unused, for signature compatibility) - return_lse: If True, also return the combined global LSE - is_lse_base_on_e: If True, LSE is base e; if False, base 2 + import flashinfer.comm as flashinfer_comm - Returns: - Combined output [B, H/N, D] (head-scattered) - If return_lse=True, also returns global_lse [B, H/N] - """ world_size = cp_group.world_size + B, H, D = cp_attn_out.shape + H_per_rank = H // world_size + workspace = _dcp_a2a_fi_workspace(cp_group, cp_attn_out.device) + partial_o, softmax_stats = _dcp_a2a_fi_send_buffers( + (B, H_per_rank, world_size, D), + (B, H_per_rank, world_size, 2), + device=cp_attn_out.device, + dtype=cp_attn_out.dtype, + ) + _dcp_a2a_fi_pack_send( + cp_attn_out, + cp_attn_lse, + partial_o, + softmax_stats, + world_size, + H_per_rank, + D, + valid_counts=valid_counts, + ) + partial_o, softmax_stats = flashinfer_comm.decode_cp_a2a_alltoall( + partial_o, + softmax_stats, + workspace, + cp_group.rank_in_group, + world_size, + ) + partial_o = _as_torch_tensor(partial_o) + softmax_stats = _as_torch_tensor(softmax_stats) + return _dcp_a2a_fi_unpack_combine( + partial_o, + softmax_stats, + D, + return_lse, + is_lse_base_on_e, + cp_attn_out.device, + cp_attn_out.dtype, + out=out, + ) - if world_size == 1: - if return_lse: - return cp_attn_out, cp_attn_lse - return cp_attn_out + +def _dcp_a2a_triton_lse_reduce( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_group: GroupCoordinator, + return_lse: bool = False, + is_lse_base_on_e: bool = True, + out: torch.Tensor | None = None, + valid_counts: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + world_size = cp_group.world_size B, H, D = cp_attn_out.shape - if H % world_size != 0: - raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.") H_per_rank = H // world_size # The pack kernel bit-casts the LSE as fp32; some MLA backends return it in # the activation dtype (bf16/fp16), so enforce the documented fp32 contract. @@ -446,6 +948,7 @@ def dcp_a2a_lse_reduce( H_per_rank, D, lse_pack_dim, + valid_counts=valid_counts, ) work = dist.all_to_all_single( @@ -457,5 +960,77 @@ def dcp_a2a_lse_reduce( work.wait() return _dcp_a2a_unpack_combine( - recv_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e + recv_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e, out=out + ) + + +def dcp_a2a_lse_reduce( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_group: GroupCoordinator, + ctx: CPTritonContext | None = None, + return_lse: bool = False, + is_lse_base_on_e: bool = True, + out: torch.Tensor | None = None, + valid_counts: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Combine partial attention outputs across DCP ranks using All-to-All. + + The output and fp32 LSE are packed into a single output-dtype buffer, sent + with one All-to-All, then unpacked and combined with exact LSE weighting. + + Args: + cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim + cp_attn_lse: [B, H] log-sum-exp values (fp32) + cp_group: GroupCoordinator for DCP communication + ctx: CPTritonContext (unused, for signature compatibility) + return_lse: If True, also return the combined global LSE + is_lse_base_on_e: If True, LSE is base e; if False, base 2 + out: Optional output tensor [B, H/N, D] to write the combined result. + valid_counts: Optional [B] local valid sparse-block count. Tokens with + count 0 are packed as zero output and -inf LSE. + + Returns: + Combined output [B, H/N, D] (head-scattered) + If return_lse=True, also returns global_lse [B, H/N] + """ + world_size = cp_group.world_size + + if world_size == 1: + if return_lse: + if out is not None: + out.copy_(cp_attn_out) + return out, cp_attn_lse + return cp_attn_out, cp_attn_lse + if out is not None: + out.copy_(cp_attn_out) + return out + return cp_attn_out + + _B, H, _D = cp_attn_out.shape + if H % world_size != 0: + raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.") + + # FlashInfer's fused decode-CP all-to-all (helix) is the fastest transport + # measured for this exchange; fall back to the Triton/NCCL path when + # flashinfer.comm is unavailable. + if _flashinfer_dcp_a2a_available(): + return _dcp_a2a_flashinfer_lse_reduce( + cp_attn_out, + cp_attn_lse, + cp_group, + return_lse=return_lse, + is_lse_base_on_e=is_lse_base_on_e, + out=out, + valid_counts=valid_counts, + ) + return _dcp_a2a_triton_lse_reduce( + cp_attn_out, + cp_attn_lse, + cp_group, + return_lse=return_lse, + is_lse_base_on_e=is_lse_base_on_e, + out=out, + valid_counts=valid_counts, )