From 494aedfbb2640324848ccdd0a451d16b76763fd9 Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 18 Sep 2026 08:02:28 +0000 Subject: [PATCH 01/15] perf: optimize DCP query communication --- lmdeploy/pytorch/backends/base.py | 3 +- lmdeploy/pytorch/backends/communicator.py | 26 +++- .../pytorch/backends/cuda/attention/cp.py | 26 ++-- .../pytorch/backends/cuda/attention/mla.py | 7 ++ .../backends/cuda/comm/communicator.py | 89 ++++++++++---- .../backends/cuda/comm/symm_mem_allgather.py | 37 +++--- lmdeploy/pytorch/backends/cuda/op_backend.py | 4 +- .../pytorch/backends/dlinfer/op_backend.py | 6 +- lmdeploy/pytorch/distributed.py | 16 ++- .../pytorch/engine/executor/ray_executor.py | 3 +- lmdeploy/pytorch/envs.py | 2 +- lmdeploy/pytorch/nn/embedding.py | 5 +- tests/pytorch/kernel/test_mla_attention.py | 6 +- tests/pytorch/test_cuda_communicator.py | 114 ++++++++++++++++-- 14 files changed, 263 insertions(+), 81 deletions(-) diff --git a/lmdeploy/pytorch/backends/base.py b/lmdeploy/pytorch/backends/base.py index d03ce96ecc..b82983c9bc 100644 --- a/lmdeploy/pytorch/backends/base.py +++ b/lmdeploy/pytorch/backends/base.py @@ -92,13 +92,14 @@ def model_build_context(ctx_mgr): yield @classmethod - def build_communicator(cls, cpu_group, device_group, dist_config): + def build_communicator(cls, cpu_group, device_group, dist_config, *, group_roles=('tp', )): """Build a device communicator.""" from .communicator import build_communicator return build_communicator( cpu_group=cpu_group, device_group=device_group, dist_config=dist_config, + group_roles=group_roles, ) @staticmethod diff --git a/lmdeploy/pytorch/backends/communicator.py b/lmdeploy/pytorch/backends/communicator.py index 418458871d..48e8556fde 100644 --- a/lmdeploy/pytorch/backends/communicator.py +++ b/lmdeploy/pytorch/backends/communicator.py @@ -29,12 +29,36 @@ def all_reduce_(self, input: torch.Tensor): """All-reduce ``input`` in place.""" dist.all_reduce(input, group=self.device_group) + def prepare_query_gather(self, num_heads: int, head_size: int): + """Prepare an optional head-axis query gather before graph capture.""" + pass + + def gather_query(self, query: torch.Tensor) -> torch.Tensor: + """Gather query heads across ranks. + + Consume the result on the same stream before the next query gather; CUDA implementations may borrow an arena. + """ + world_size = dist.get_world_size(self.device_group) + if world_size == 1: + return query + if not query.is_contiguous(): + transposed = query.transpose(0, 1).contiguous() + gathered = transposed.new_empty(world_size * transposed.size(0), *transposed.shape[1:]) + dist.all_gather_into_tensor(gathered, transposed, group=self.device_group) + return gathered.transpose(0, 1).contiguous() + + # Contiguous token-major inputs need only the final head-axis packing. + gathered = query.new_empty(world_size * query.size(0), *query.shape[1:]) + dist.all_gather_into_tensor(gathered, query, group=self.device_group) + gathered = gathered.view(world_size, *query.shape) + return gathered.transpose(0, 1).reshape(query.size(0), -1, query.size(2)).contiguous() + def close(self): """Release communicator-owned resources.""" pass def build_communicator(cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup, - dist_config): + dist_config, *, group_roles: tuple[str, ...] = ('tp', )): """Build the default process-group communicator.""" return DeviceCommunicator(device_group=device_group) diff --git a/lmdeploy/pytorch/backends/cuda/attention/cp.py b/lmdeploy/pytorch/backends/cuda/attention/cp.py index fe55818372..dbae22fa95 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/cp.py +++ b/lmdeploy/pytorch/backends/cuda/attention/cp.py @@ -5,26 +5,16 @@ def gather_dcp_query(query: torch.Tensor, *, dcp_world_size: int) -> torch.Tensor: - """Gather [tokens, local_heads, dim] queries along the head axis.""" + """Gather [tokens, local_heads, dim] queries along the head axis. + + Consume the result on the same stream before the next query gather; the optimized path borrows a shared arena. + """ if dcp_world_size == 1: return query - from lmdeploy.pytorch.distributed import all_gather_into_tensor - - # Keep the existing packing path for strided queries. Direct gather - # saves a copy only when the token-major input is already contiguous. - if not query.is_contiguous(): - transposed = query.transpose(0, 1).contiguous() - gathered = transposed.new_empty( - dcp_world_size * transposed.size(0), *transposed.shape[1:]) - all_gather_into_tensor(gathered, transposed, group='dcp') - return gathered.transpose(0, 1).contiguous() - - gathered = query.new_empty(dcp_world_size * query.size(0), *query.shape[1:]) - all_gather_into_tensor(gathered, query, group='dcp') - # Gather token-major queries directly, then join rank-local heads. - # Contiguous inputs need only this final layout conversion. - gathered = gathered.view(dcp_world_size, *query.shape) - return gathered.transpose(0, 1).reshape(query.size(0), -1, query.size(2)).contiguous() + from lmdeploy.pytorch.distributed import get_dist_manager + + communicator = get_dist_manager().current_context().dcp_group.communicator + return communicator.gather_query(query) def merge_dcp_attention(local_output: torch.Tensor, diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index f73862847e..b03326c080 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -254,6 +254,13 @@ def __init__( assert num_kv_heads == 1, 'MLA requires num kv heads equal to 1' self.use_fa3 = use_fa3 + if self.dcp_world_size > 1: + from lmdeploy.pytorch.distributed import get_dist_manager + + communicator = get_dist_manager().current_context().dcp_group.communicator + if communicator is not None: + communicator.prepare_query_gather(num_heads, head_size) + def get_step_metadata_provider(self): """Describe metadata required by this selected implementation.""" return FlashMLAAttentionMetaBuilder(num_attention_heads=self.num_heads) diff --git a/lmdeploy/pytorch/backends/cuda/comm/communicator.py b/lmdeploy/pytorch/backends/cuda/comm/communicator.py index cd0a2bddf3..6147386dcb 100644 --- a/lmdeploy/pytorch/backends/cuda/comm/communicator.py +++ b/lmdeploy/pytorch/backends/cuda/comm/communicator.py @@ -12,24 +12,26 @@ class CudaCommunicator(DeviceCommunicator): """Dispatch optional CUDA collectives with a process-group fallback.""" - def __init__(self, cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup): + def __init__(self, cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup, *, + all_reduce_backend: str | None = None, symm_mem_query_gather: bool = False): super().__init__(device_group=device_group) - if _envs.enable_flashinfer_allreduce and _envs.enable_symm_mem_allreduce: - raise ValueError('FlashInfer and symmetric-memory all-reduce cannot be enabled together.') - - self._flashinfer = (FlashInferAllReduce(cpu_group) - if _envs.enable_flashinfer_allreduce else None) - self._symm_mem = (SymmetricMemoryAllReduce(cpu_group) - if _envs.enable_symm_mem_allreduce else None) + self._all_reduce_backend = all_reduce_backend + self._all_reduce = None + if all_reduce_backend == 'flashinfer': + self._all_reduce = FlashInferAllReduce(cpu_group) + elif all_reduce_backend == 'symm_mem': + self._all_reduce = SymmetricMemoryAllReduce(cpu_group) + self._symm_mem_query_gather = symm_mem_query_gather + self._query_gatherer = None + self._query_shape = None def supports_optimized_all_reduce(self) -> bool: """Whether an optimized all-reduce implementation is available.""" - return ((self._flashinfer is not None and self._flashinfer.is_available()) - or (self._symm_mem is not None and self._symm_mem.is_available())) + return self._all_reduce is not None and self._all_reduce.is_available() def supports_fused_all_reduce_residual_rms_norm(self) -> bool: """Whether fused all-reduce, residual and RMSNorm is available.""" - return self._flashinfer is not None and self._flashinfer.is_available() + return self._all_reduce_backend == 'flashinfer' and self._all_reduce.is_available() def try_fused_all_reduce_residual_rms_norm(self, input: torch.Tensor, @@ -37,9 +39,9 @@ def try_fused_all_reduce_residual_rms_norm(self, weight: torch.Tensor, eps: float): """Run fused all-reduce, residual and RMSNorm when eligible.""" - if self._flashinfer is None: + if self._all_reduce_backend != 'flashinfer': return None - return self._flashinfer.fused_all_reduce_residual_rms_norm( + return self._all_reduce.fused_all_reduce_residual_rms_norm( input=input, residual=residual, weight=weight, @@ -48,18 +50,41 @@ def try_fused_all_reduce_residual_rms_norm(self, def all_reduce_(self, input: torch.Tensor): """Dispatch all-reduce through optimized CUDA backends.""" - if self._flashinfer is not None and self._flashinfer.all_reduce_(input): - return - if self._symm_mem is not None and self._symm_mem.all_reduce_(input): + if self._all_reduce is not None and self._all_reduce.all_reduce_(input): return super().all_reduce_(input) def close(self): """Release communicator-owned workspaces.""" - if self._flashinfer is not None: - self._flashinfer.close() - if self._symm_mem is not None: - self._symm_mem.close() + if self._all_reduce is not None: + self._all_reduce.close() + if self._query_gatherer is not None: + self._query_gatherer.release() + self._query_gatherer = None + + def prepare_query_gather(self, num_heads: int, head_size: int): + """Share one bounded query arena across attention layers and MTP.""" + if self._query_gatherer is not None: + return + from .symm_mem_allgather import MultimemAllGatherer + + self._query_shape = (num_heads, head_size) + self._query_gatherer = MultimemAllGatherer( + self.device_group, dist.get_rank(self.device_group), + num_heads * head_size * dist.get_world_size(self.device_group), + torch.device('cuda', torch.cuda.current_device()), torch.bfloat16, + enabled=self._symm_mem_query_gather, capacity_bytes=64 * 1024 * 1024) + + def gather_query(self, query: torch.Tensor) -> torch.Tensor: + """Gather query heads, borrowing the symmetric arena when eligible.""" + if (self._query_gatherer is not None and query.dtype == torch.bfloat16 + and tuple(query.shape[1:]) == self._query_shape and query.is_contiguous()): + # Same-stream attention finishes before the next gather's entry + # barrier permits any rank to overwrite the borrowed arena. + output = self._query_gatherer(query.flatten(1), safe=False) + if output is not None: + return output.view(query.size(0), -1, query.size(2)) + return super().gather_query(query) def should_try_symm_mem(dist_config) -> bool: @@ -70,10 +95,24 @@ def should_try_symm_mem(dist_config) -> bool: def build_cuda_communicator(cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup, - dist_config): - """Build the optional CUDA communicator for a TP group.""" + dist_config, *, group_roles: tuple[str, ...] = ('tp', )): + """Build the optional CUDA communicator for a TP or DCP group.""" compatible = dist_config.dp == 1 and dist_config.ep == 1 and not dist_config.enable_microbatch - enabled = _envs.enable_flashinfer_allreduce or should_try_symm_mem(dist_config) - if not compatible or not enabled: + if not compatible: + return None + all_reduce_backend = None + if 'tp' in group_roles: + if _envs.enable_flashinfer_allreduce and _envs.enable_symm_mem_allreduce: + raise ValueError('FlashInfer and symmetric-memory all-reduce cannot be enabled together.') + if _envs.enable_flashinfer_allreduce: + all_reduce_backend = 'flashinfer' + elif should_try_symm_mem(dist_config): + all_reduce_backend = 'symm_mem' + query_gather = 'dcp' in group_roles and dist_config.dcp > 1 + if all_reduce_backend is None and not query_gather: return None - return CudaCommunicator(cpu_group=cpu_group, device_group=device_group) + # DCP ranks prepare together even when disabled, so admission can agree + # on NCCL fallback if the opt-in flag differs between ranks. + return CudaCommunicator(cpu_group=cpu_group, device_group=device_group, + all_reduce_backend=all_reduce_backend, + symm_mem_query_gather=query_gather and _envs.enable_symm_mem_dcp) diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py index cf95b9ed0a..04d575dee9 100644 --- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py +++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py @@ -1,5 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. -"""Setup-owned LM-head all-gather with a portable NCCL fallback. +"""Setup-owned symmetric-memory hidden-axis all-gather. The optional Triton kernels are imported only during collective preparation. Their SGLang source reference is in kernels/cuda/symm_mem_allgather.py. @@ -10,15 +10,13 @@ import torch import torch.distributed as dist -from lmdeploy.pytorch import envs as _envs - logger = logging.getLogger(__name__) class MultimemAllGatherer: - """Own configuration, TP admission and the symmetric arena for one LM-head. + """Own configuration, collective admission and one symmetric arena. - Construct on every TP rank after process-group setup. Only preparation performs allocation and consensus; forward + Construct on every group rank after process-group setup. Only preparation performs allocation and consensus; forward never initializes or rebuilds. Call reset_for_weight after coordinated, quiescent device/dtype changes. Runtime rows must match across ranks, as required by NCCL all-gather too. """ @@ -28,12 +26,13 @@ class MultimemAllGatherer: _WORLD_SIZES = {2, 4, 8} def __init__(self, group: dist.ProcessGroup, rank: int, gathered_width: int, - device: torch.device, dtype: torch.dtype): + device: torch.device, dtype: torch.dtype, *, + enabled: bool, capacity_bytes: int): self._group = group self._rank = rank self._gathered_width = gathered_width - self._enabled = _envs.enable_symm_mem_lmhead - self._capacity_bytes = _envs.symm_mem_lmhead_max_mb * 1024 * 1024 + self._enabled = enabled + self._capacity_bytes = capacity_bytes self._state = None self._kernels = None self._graph_ready_shapes = set() @@ -54,7 +53,7 @@ def _same_config(self, values: tuple[int, ...], device: torch.device) -> bool: def _disabled(self, reason: str) -> bool: if self._rank == 0: - logger.warning('symmetric-memory LM-head disabled: %s', reason) + logger.warning('symmetric-memory all-gather disabled: %s', reason) return False def prepare(self, device: torch.device, dtype: torch.dtype) -> bool: @@ -86,9 +85,9 @@ def prepare(self, device: torch.device, dtype: torch.dtype) -> bool: and self._capacity_bytes >= width * torch.bfloat16.itemsize and capability_ok) if not self._agree(valid, device): - return self._disabled('unsupported TP, dtype, width, capacity or device') + return self._disabled('unsupported group size, dtype, width, capacity or device') if not self._same_config((width, self._capacity_bytes), device): - return self._disabled('inconsistent TP arena configuration') + return self._disabled('inconsistent group arena configuration') kernels = None try: @@ -110,7 +109,7 @@ def prepare(self, device: torch.device, dtype: torch.dtype) -> bool: except Exception as exc: logger.debug('Symmetric allocation failed: %s', exc) if not self._agree(allocation_ok, device): - return self._disabled('a TP rank could not allocate its arena') + return self._disabled('a rank could not allocate its arena') # Once all ranks commit, rendezvous errors are fatal, not a rank-local # reason to fall back and strand peers inside a collective. @@ -144,7 +143,8 @@ def _valid_handle(self, state, kernels) -> bool: return False def reset_for_weight(self, weight: torch.Tensor) -> None: - """Notify the provider after a quiescent, TP-wide model transition.""" + """Notify the provider after a quiescent, group-wide model + transition.""" if (weight.device, weight.dtype) != self._weight_contract: self.prepare(weight.device, weight.dtype) @@ -161,23 +161,24 @@ def release(self) -> None: self._kernels = None self._graph_ready_shapes.clear() - def __call__(self, x: torch.Tensor) -> torch.Tensor | None: - """Gather admitted logits into an owning output; None means NCCL.""" + def __call__(self, x: torch.Tensor, *, safe: bool = True) -> torch.Tensor | None: + """Gather hidden shards; unsafe output is borrowed until the next + call.""" state = self._state if state is None: return None if not (x.dim() == 2 and x.dtype == torch.bfloat16 and x.device == state.device and x.is_contiguous() and x.data_ptr() % 16 == 0 and x.shape[-1] * state.world_size == state.hidden_dim): - # After static TP admission this is a programming error. A local + # After collective admission this is a programming error. A local # NCCL fallback could disagree with peers launching multimem. - raise RuntimeError('multimem all-gather input contract changed after TP-wide admission') + raise RuntimeError('multimem all-gather input contract changed after group-wide admission') rows = x.shape[0] if not self._MIN_TOKENS <= rows <= state.max_token_num: return None if rows not in self._graph_ready_shapes and torch.cuda.is_current_stream_capturing(): return None output = self._kernels.all_gather_inner( - state, x, tp_hidden_dim=self._gathered_width, safe=True, _validated=True) + state, x, tp_hidden_dim=self._gathered_width, safe=safe, _validated=True) self._graph_ready_shapes.add(rows) return output diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py index 5847064776..bb829b5566 100644 --- a/lmdeploy/pytorch/backends/cuda/op_backend.py +++ b/lmdeploy/pytorch/backends/cuda/op_backend.py @@ -228,13 +228,14 @@ def get_v4_attention_metadata_cls(): return CudaV4AttentionMetadata @classmethod - def build_communicator(cls, cpu_group, device_group, dist_config): + def build_communicator(cls, cpu_group, device_group, dist_config, *, group_roles=('tp', )): """Build a CUDA communicator.""" from .comm.communicator import build_cuda_communicator communicator = build_cuda_communicator( cpu_group=cpu_group, device_group=device_group, dist_config=dist_config, + group_roles=group_roles, ) if communicator is not None: return communicator @@ -242,6 +243,7 @@ def build_communicator(cls, cpu_group, device_group, dist_config): cpu_group=cpu_group, device_group=device_group, dist_config=dist_config, + group_roles=group_roles, ) @staticmethod diff --git a/lmdeploy/pytorch/backends/dlinfer/op_backend.py b/lmdeploy/pytorch/backends/dlinfer/op_backend.py index 7b4b838526..a7b5a82a94 100644 --- a/lmdeploy/pytorch/backends/dlinfer/op_backend.py +++ b/lmdeploy/pytorch/backends/dlinfer/op_backend.py @@ -127,15 +127,17 @@ def get_cache_backend(cls): return DlinferCacheBackend @classmethod - def build_communicator(cls, cpu_group, device_group, dist_config): + def build_communicator(cls, cpu_group, device_group, dist_config, *, group_roles=('tp', )): """Build a DLInfer communicator.""" from lmdeploy.pytorch import envs - cuda_communicator_enabled = envs.enable_flashinfer_allreduce or envs.enable_symm_mem_allreduce + cuda_communicator_enabled = (envs.enable_flashinfer_allreduce or envs.enable_symm_mem_allreduce + or envs.enable_symm_mem_dcp) assert not cuda_communicator_enabled, 'CUDA communicators are not supported by DLInfer.' return super().build_communicator( cpu_group=cpu_group, device_group=device_group, dist_config=dist_config, + group_roles=group_roles, ) @staticmethod diff --git a/lmdeploy/pytorch/distributed.py b/lmdeploy/pytorch/distributed.py index 525de1bf58..c7a88e05ec 100644 --- a/lmdeploy/pytorch/distributed.py +++ b/lmdeploy/pytorch/distributed.py @@ -237,17 +237,23 @@ def _build_dcp_group(context: 'DistContext', timeout: timedelta, ) -def _build_tp_communicators(context: 'DistContext'): - """Attach one communicator to each rank-local, unique TP group.""" +def _build_communicators(context: 'DistContext'): + """Attach one communicator to each rank-local, unique TP/DCP group.""" build_communicator = context.communicator_builder - groups = (context.attn_tp_group, context.mlp_tp_group, context.moe_tp_group) + tp_groups = (context.attn_tp_group, context.mlp_tp_group, context.moe_tp_group) + groups = (*tp_groups, context.dcp_group) for group in {id(group): group for group in groups}.values(): if group.gpu_group is None: continue + # Aliased groups share one communicator with both operation roles. + roles = ('tp', ) if any(group is tp_group for tp_group in tp_groups) else () + if group is context.dcp_group: + roles += ('dcp', ) group.communicator = build_communicator( cpu_group=group.cpu_group, device_group=group.gpu_group, dist_config=context.dist_config, + group_roles=roles, ) @@ -330,11 +336,13 @@ def build(cls, # tp _build_tp_group(context, timeout, cpu_backend=cpu_backend, ccl_backend=ccl_backend) - _build_tp_communicators(context) # cp _build_dcp_group(context, timeout, cpu_backend=cpu_backend, ccl_backend=ccl_backend) + # communicator + _build_communicators(context) + # ep cls._build_ep_group(context, timeout, ccl_backend=ccl_backend) diff --git a/lmdeploy/pytorch/engine/executor/ray_executor.py b/lmdeploy/pytorch/engine/executor/ray_executor.py index 3c77d91590..6939388e30 100644 --- a/lmdeploy/pytorch/engine/executor/ray_executor.py +++ b/lmdeploy/pytorch/engine/executor/ray_executor.py @@ -106,7 +106,8 @@ def _needs_symm_mem_device_setup(dist_config: DistConfig) -> bool: """ from lmdeploy.pytorch.backends.cuda.comm.communicator import should_try_symm_mem return (should_try_symm_mem(dist_config) - or (_envs.enable_symm_mem_lmhead and dist_config.attn_tp > 1)) + or (_envs.enable_symm_mem_lmhead and dist_config.attn_tp > 1) + or (_envs.enable_symm_mem_dcp and dist_config.dcp > 1)) class RemoteLogger: diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py index 8a714e1a06..3319049eb3 100644 --- a/lmdeploy/pytorch/envs.py +++ b/lmdeploy/pytorch/envs.py @@ -249,8 +249,8 @@ def _patched_get_env( # cuda communicator enable_flashinfer_allreduce = env_to_bool('LMDEPLOY_ENABLE_FLASHINFER_ALLREDUCE', False) enable_symm_mem_allreduce = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_ALLREDUCE', False) - # NCCL remains the default; launch and small-token policies are internal. enable_symm_mem_lmhead = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_LMHEAD', False) + enable_symm_mem_dcp = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_DCP', False) symm_mem_lmhead_max_mb = max(1, env_to_int('LMDEPLOY_SYMM_MEM_LMHEAD_MAX_MB', 64)) # opt-ttft diff --git a/lmdeploy/pytorch/nn/embedding.py b/lmdeploy/pytorch/nn/embedding.py index 13983fa6db..fb04f9b179 100644 --- a/lmdeploy/pytorch/nn/embedding.py +++ b/lmdeploy/pytorch/nn/embedding.py @@ -3,6 +3,7 @@ import torch.distributed as dist from torch import nn +from lmdeploy.pytorch import envs as _envs from lmdeploy.pytorch.backends import get_backend from lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather import MultimemAllGatherer from lmdeploy.pytorch.backends.embedding import EmbeddingBuildSpec @@ -157,7 +158,9 @@ def __init__( self._symm_mem_gatherer = ( MultimemAllGatherer(self.tp_group, self.tp_rank, self.tp * self.vocab_size_padded, - self.weight.device, self.weight.dtype) + self.weight.device, self.weight.dtype, + enabled=_envs.enable_symm_mem_lmhead, + capacity_bytes=_envs.symm_mem_lmhead_max_mb * 1024 * 1024) if self.all_reduce else None) def tie_weights(self, embedding: ParallelEmbedding): diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py index ac7a27b1a3..5220f24a50 100644 --- a/tests/pytorch/kernel/test_mla_attention.py +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -146,6 +146,7 @@ def test_dcp_query_all_gather_preserves_contiguous_head_order(monkeypatch, devic if device == 'cuda' and not torch.cuda.is_available(): pytest.skip('requires CUDA') from lmdeploy.pytorch import distributed + from lmdeploy.pytorch.backends.communicator import DeviceCommunicator queries = torch.arange(dcp_size * num_tokens * 3 * 8, device=device, dtype=torch.float32) queries = queries.reshape(dcp_size, num_tokens, 3, 8) @@ -162,7 +163,10 @@ def fake_all_gather(output, input_tensor, group='tp', async_op=False): source = queries[rank].transpose(0, 1) if strided else queries[rank] output[rank * source.size(0):(rank + 1) * source.size(0)].copy_(source) - monkeypatch.setattr(distributed, 'all_gather_into_tensor', fake_all_gather) + monkeypatch.setattr(torch.distributed, 'all_gather_into_tensor', fake_all_gather) + monkeypatch.setattr(torch.distributed, 'get_world_size', lambda group: dcp_size) + monkeypatch.setattr(distributed.get_dist_manager().current_context().dcp_group, + 'communicator', DeviceCommunicator('dcp')) gathered = gather_dcp_query(rank0_query, dcp_world_size=dcp_size) expected = torch.cat(list(queries), dim=1) diff --git a/tests/pytorch/test_cuda_communicator.py b/tests/pytorch/test_cuda_communicator.py index 4c98aaa146..4b841fc7a7 100644 --- a/tests/pytorch/test_cuda_communicator.py +++ b/tests/pytorch/test_cuda_communicator.py @@ -12,6 +12,80 @@ from lmdeploy.pytorch.nn import norm as norm_module +def _run_dcp_query_gather(rank, rendezvous, enabled): + from datetime import timedelta + + from torch import distributed as dist + + from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_query + from lmdeploy.pytorch.backends.cuda.op_backend import CudaOpsBackend + from lmdeploy.pytorch.config import DistConfig + from lmdeploy.pytorch.distributed import DistContext, get_dist_manager + + torch.cuda.set_device(rank) + communicator_module._envs.enable_symm_mem_dcp = enabled[rank] + dist.init_process_group('nccl', init_method=rendezvous, rank=rank, world_size=2, + timeout=timedelta(seconds=60)) + ctx = DistContext.build(rank, DistConfig(tp=2, dcp=2), + communicator_builder=CudaOpsBackend.build_communicator) + try: + with get_dist_manager().context(ctx): + comm = ctx.dcp_group.communicator + comm.prepare_query_gather(32, 576) + if all(enabled): + assert comm._query_gatherer._state is not None + else: + assert comm._query_gatherer._state is None + assert comm._all_reduce is None + query = torch.empty(384, 32, 576, device='cuda', dtype=torch.bfloat16) + # Ineligible inputs still return correct results through NCCL. + for fallback in (query[:1], query[..., ::2], query.float(), + torch.empty(1024, 32, 576, device='cuda', dtype=query.dtype)): + fallback.fill_(rank) + output = comm.gather_query(fallback) + expected = torch.arange(2, device='cuda', dtype=fallback.dtype) + expected = expected.repeat_interleave(32)[None, :, None].expand_as(output) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + rows = (2, 16, 96, 384) + for count in rows: + gather_dcp_query(query[:count], dcp_world_size=2) + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + outputs = [] + with torch.cuda.graph(graph): + for step in range(16): + query.fill_(rank + step * 8) + gathered = gather_dcp_query(query[:rows[step % len(rows)]], dcp_world_size=2) + outputs.append(gathered.clone()) + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + graph.reset() + for step, output in enumerate(outputs): + expected = (torch.arange(2, device='cuda', dtype=query.dtype) + step * 8) + expected = expected.repeat_interleave(32)[None, :, None].expand_as(output) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + finally: + ctx.close() + dist.destroy_process_group() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='requires two CUDA GPUs') +@pytest.mark.parametrize('enabled', [(False, False), (True, True), (True, False)], + ids=['nccl', 'symm_mem', 'mixed_flags']) +def test_dcp_query_gather_graph_reuses_arena(tmp_path, enabled): + if all(enabled): + from torch.distributed._symmetric_memory import DeviceType, _SymmetricMemory + if any(torch.cuda.get_device_capability(i)[0] < 9 + or not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, i) for i in range(2)): + pytest.skip('requires SM90 or newer with multicast support') + if not torch.cuda.can_device_access_peer(0, 1): + pytest.skip('requires peer access') + torch.multiprocessing.spawn(_run_dcp_query_gather, + args=((tmp_path / 'rendezvous').as_uri(), enabled), nprocs=2) + + class _Collective: def __init__(self, *, result=None, handled=False): @@ -41,7 +115,8 @@ def _build_communicator(monkeypatch, symm_mem_cls = Mock(return_value=collective) monkeypatch.setattr(communicator_module, 'FlashInferAllReduce', flashinfer_cls) monkeypatch.setattr(communicator_module, 'SymmetricMemoryAllReduce', symm_mem_cls) - communicator = communicator_module.CudaCommunicator(cpu_group='cpu', device_group='gpu') + communicator = communicator_module.CudaCommunicator(cpu_group='cpu', device_group='gpu', + all_reduce_backend=backend) enabled_cls, disabled_cls = ((flashinfer_cls, symm_mem_cls) if use_flashinfer else (symm_mem_cls, flashinfer_cls)) enabled_cls.assert_called_once_with('cpu') @@ -93,7 +168,8 @@ def test_cuda_communicator_rejects_multiple_backends(monkeypatch): monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', True) monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', True) with pytest.raises(ValueError, match='cannot be enabled together'): - communicator_module.CudaCommunicator(cpu_group='cpu', device_group='gpu') + config = SimpleNamespace(dp=1, ep=1, attn_tp=2, dcp=1, enable_microbatch=False) + communicator_module.build_cuda_communicator('cpu', 'gpu', config) def test_rms_norm_fuses_or_falls_back(monkeypatch): @@ -236,17 +312,18 @@ def test_build_cuda_communicator_gates_unsupported_parallelism(monkeypatch): monkeypatch.setattr(communicator_module, 'CudaCommunicator', communicator_cls) for config in ( - SimpleNamespace(dp=2, ep=1, attn_tp=8, enable_microbatch=False), - SimpleNamespace(dp=1, ep=2, attn_tp=8, enable_microbatch=False), - SimpleNamespace(dp=1, ep=1, attn_tp=8, enable_microbatch=True), + SimpleNamespace(dp=2, ep=1, attn_tp=8, dcp=1, enable_microbatch=False), + SimpleNamespace(dp=1, ep=2, attn_tp=8, dcp=1, enable_microbatch=False), + SimpleNamespace(dp=1, ep=1, attn_tp=8, dcp=1, enable_microbatch=True), ): assert communicator_module.build_cuda_communicator('cpu', 'device', config) is None communicator_cls.assert_not_called() - config = SimpleNamespace(dp=1, ep=1, attn_tp=4, enable_microbatch=False) + config = SimpleNamespace(dp=1, ep=1, attn_tp=4, dcp=1, enable_microbatch=False) communicator = communicator_module.build_cuda_communicator('cpu', 'device', config) assert communicator is communicator_cls.return_value - communicator_cls.assert_called_once_with(cpu_group='cpu', device_group='device') + communicator_cls.assert_called_once_with(cpu_group='cpu', device_group='device', + all_reduce_backend='flashinfer', symm_mem_query_gather=False) monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', False) monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', True) @@ -255,6 +332,29 @@ def test_build_cuda_communicator_gates_unsupported_parallelism(monkeypatch): config.attn_tp = 2 assert communicator_module.build_cuda_communicator('cpu', 'device', config) is communicator_cls.return_value + # TP all-reduce flags must not allocate all-reduce resources on DCP groups. + config.dcp = 2 + monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_dcp', True) + communicator_module.build_cuda_communicator('cpu', 'device', config, group_roles=('dcp', )) + communicator_cls.assert_called_with(cpu_group='cpu', device_group='device', + all_reduce_backend=None, symm_mem_query_gather=True) + + +@pytest.mark.parametrize('shared_dcp_group', [False, True]) +def test_communicator_group_roles(shared_dcp_group): + from lmdeploy.pytorch.distributed import DistContext, DistGroup, _build_communicators + + tp_group = DistGroup(cpu_group='tp_cpu', gpu_group='tp_gpu') + dcp_group = tp_group if shared_dcp_group else DistGroup(cpu_group='dcp_cpu', gpu_group='dcp_gpu') + builder = Mock(side_effect=lambda **kwargs: object()) + context = DistContext(attn_tp_group=tp_group, mlp_tp_group=tp_group, moe_tp_group=tp_group, + dcp_group=dcp_group, communicator_builder=builder) + _build_communicators(context) + + roles = [call.kwargs['group_roles'] for call in builder.call_args_list] + assert roles == ([('tp', 'dcp')] if shared_dcp_group else [('tp', ), ('dcp', )]) + assert (tp_group.communicator is dcp_group.communicator) == shared_dcp_group + def test_dlinfer_communicator_rejects_cuda_options(monkeypatch): monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', True) From ab9394c2d6bdb9ce355a8d7dd3f8606d66a87847 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 21 Sep 2026 09:19:28 +0000 Subject: [PATCH 02/15] feat: support GQA DCP with FP8 KV cache Add GQA decode shard merging and bounded cached-prefix prefill using shared DCP collectives. Extend paged attention to return LSE and FP8 KV fill to write rank-owned tokens. Share KV workspace sizing and update configuration, documentation, and regression tests. --- docs/en/advance/context_parallel.md | 18 ++- lmdeploy/pytorch/backends/cp_utils.py | 41 ++++-- .../backends/cuda/attention/__init__.py | 5 + .../pytorch/backends/cuda/attention/cp.py | 132 +++++++++++++++++- .../pytorch/backends/cuda/attention/mla.py | 30 +--- lmdeploy/pytorch/config.py | 6 +- lmdeploy/pytorch/engine/executor/__init__.py | 3 + lmdeploy/pytorch/engine/executor/base.py | 3 +- .../pytorch/kernels/cuda/fill_kv_cache.py | 21 +-- .../pytorch/kernels/cuda/pagedattention.py | 27 +++- .../third_party/flash_attn_interface.py | 20 ++- tests/pytorch/config/test_gqa_dcp.py | 84 +++++++++++ tests/pytorch/kernel/test_fill_kv_cache.py | 59 ++++++++ tests/pytorch/kernel/test_mla_attention.py | 31 ++-- tests/pytorch/test_gqa_dcp.py | 77 ++++++++++ 15 files changed, 497 insertions(+), 60 deletions(-) create mode 100644 tests/pytorch/config/test_gqa_dcp.py create mode 100644 tests/pytorch/test_gqa_dcp.py diff --git a/docs/en/advance/context_parallel.md b/docs/en/advance/context_parallel.md index a474ca74b9..2cb741bf5c 100644 --- a/docs/en/advance/context_parallel.md +++ b/docs/en/advance/context_parallel.md @@ -34,12 +34,16 @@ lmdeploy serve api_server Qwen/Qwen3-235B-A22B --tp 8 --cp 2 ## PyTorch decode context parallelism -PyTorch DCP distributes MLA KV cache across existing TP ranks, increasing +PyTorch DCP distributes replicated KV cache across existing TP ranks, increasing effective cache capacity without additional GPUs. It supports FlashMLA-backed -dense MLA and sparse DSA models. +dense MLA, sparse DSA, and GQA with replicated KV heads. ```bash lmdeploy serve api_server --backend pytorch --tp 4 --dcp 2 + +# Both Qwen3 models have four KV heads; TP=8 replicates each head twice. +lmdeploy serve api_server Qwen/Qwen3-30B-A3B --backend pytorch --tp 8 --dcp 2 +lmdeploy serve api_server Qwen/Qwen3-235B-A22B --backend pytorch --tp 8 --dcp 2 ``` For Python, use `PytorchEngineConfig(tp=4, dcp=2)`. The default `dcp=1` @@ -47,13 +51,21 @@ disables DCP. ### Requirements and supported features -- Requires NVIDIA Hopper/SM90 GPUs, FlashMLA, and BF16 activations. +- MLA requires NVIDIA Hopper/SM90 GPUs, FlashMLA, and BF16 activations. Sparse DSA also requires compatible DeepGEMM and TileLang top-k kernels (top-k 512 or 2048). +- GQA uses Triton for decode and FlashAttention-3, when available, for prefill, + with a Triton prefill fallback. It supports BF16/FP16 activations and + unquantized or per-tensor FP8 KV cache (`--quant-policy fp8` for E4M3, + `--quant-policy fp8_e5m2` for E5M2). Cached-prefix chunks are dequantized + before gathering. INT8, INT4, TurboQuant, ALiBi, and attention sinks are not supported. - `dcp` must divide the attention TP size and replicated KV-head count; `dp=1` and `ep=1` are required. - Prefix caching and `deepseek_mtp` are supported for compatible MLA models, including DeepSeek V3/V3.1, DeepSeek V3.2, and GLM DSA. +- GQA prefill retains local query heads and gathers cached K/V in bounded + chunks. Decode gathers query heads within each replicated-KV group and + merges shard outputs using their softmax normalization factors. - BF16 KV cache is supported. Sparse MLA also supports FP8 KV cache via `--quant-policy fp8`. - Sliding-window attention, MemDecode, prefill/decode disaggregation, diff --git a/lmdeploy/pytorch/backends/cp_utils.py b/lmdeploy/pytorch/backends/cp_utils.py index 4343290097..85f0d21b6f 100644 --- a/lmdeploy/pytorch/backends/cp_utils.py +++ b/lmdeploy/pytorch/backends/cp_utils.py @@ -10,29 +10,33 @@ def get_dcp_topk_workspace_size(num_rows: int, topk: int, dcp_size: int) -> int: return num_rows * topk * (16 + 8 * dcp_size) -def get_dcp_prefill_workspace_size(*, batch_size: int, head_dim: int, block_size: int, dcp_size: int) -> int: + +def get_dcp_prefill_workspace_size(*, batch_size: int, kv_width: int, block_size: int, dcp_size: int) -> int: """Per-rank prefix KV budget shared by chunk planning and cache sizing. Target 64 MiB, enlarged if one virtual block per request needs more. FP8 caches also gather dequantized BF16 KV. """ # Simultaneous BF16 buffers: one local block, D gathered, D reordered. - minimum = batch_size * block_size * (1 + 2 * dcp_size) * head_dim * 2 + minimum = batch_size * block_size * (1 + 2 * dcp_size) * kv_width * 2 return max(64 << 20, minimum) def get_dcp_workspace_size(*, num_prefill_tokens: int, num_decode_tokens: int, batch_size: int, num_heads: int, head_dim: int, dtype: torch.dtype, block_size: int, dcp_size: int, - topk: int | None, score_workspace_bytes: int) -> int: + topk: int | None, score_workspace_bytes: int, + kv_width: int) -> int: """Estimate per-rank peak temporary bytes for enabled DCP. ``num_heads`` is the rank-local query-head count before DCP gathering. Reserve max(indexer, attention) plus final - indices shared by both phases. Flattened indexer KV relies on memory headroom. + indices shared by both phases. + ``kv_width`` counts both K and V elements per cached token for GQA, or the shared + latent width for MLA. Flattened indexer KV relies on memory headroom. """ # Prefix KV: local, gathered, and reordered buffers. attention_workspace = get_dcp_prefill_workspace_size( batch_size=batch_size, - head_dim=head_dim, + kv_width=kv_width, block_size=block_size, dcp_size=dcp_size, ) @@ -67,8 +71,19 @@ class DCPPrefixChunk: local_kv_seqlens: torch.Tensor # [dcp_size, requests] local_cu_seqlens: torch.Tensor # current rank + @property + def local_capacity(self) -> int: + """Padded token capacity per rank, identical across the DCP group.""" + return self.kv_seqlens.numel() * self.size // self.local_kv_seqlens.size(0) + + def local_block_slice(self, block_size: int) -> slice: + """Select local cache blocks for this globally aligned chunk.""" + virtual_block_size = block_size * self.local_kv_seqlens.size(0) + start = self.start // virtual_block_size + return slice(start, start + self.size // virtual_block_size) + -def build_dcp_prefix_chunks(*, prefix_lens: torch.Tensor, prefix_limit: int, block_size: int, head_dim: int, +def build_dcp_prefix_chunks(*, prefix_lens: torch.Tensor, prefix_limit: int, block_size: int, kv_width: int, dcp_world_rank: tuple[int, int]) -> tuple[DCPPrefixChunk, ...]: """Plan identical collective shapes on all ranks without device sync.""" if prefix_limit <= 0: @@ -76,11 +91,11 @@ def build_dcp_prefix_chunks(*, prefix_lens: torch.Tensor, prefix_limit: int, blo dcp_size, dcp_rank = dcp_world_rank batch_size = prefix_lens.numel() workspace_bytes = get_dcp_prefill_workspace_size(batch_size=batch_size, - head_dim=head_dim, + kv_width=kv_width, block_size=block_size, dcp_size=dcp_size) virtual_block_size = block_size * dcp_size - bytes_per_block = batch_size * block_size * (1 + 2 * dcp_size) * head_dim * 2 + bytes_per_block = batch_size * block_size * (1 + 2 * dcp_size) * kv_width * 2 # Longer prefixes add chunks rather than enlarge the gather buffers. chunk_size = (workspace_bytes // bytes_per_block) * virtual_block_size ranks = torch.arange(dcp_size, device=prefix_lens.device, dtype=prefix_lens.dtype)[:, None] @@ -97,6 +112,14 @@ def build_dcp_prefix_chunks(*, prefix_lens: torch.Tensor, prefix_limit: int, blo return tuple(chunks) +def get_dcp_kv_width(model_config) -> int: + """Elements per cached token: shared MLA latent or combined rank-local GQA K/V.""" + if model_config.use_flash_mla: + return model_config.head_dim + _, num_kv_heads = model_config.get_num_qkv_head_by_tp() + return num_kv_heads * (model_config.k_head_dim + model_config.v_head_dim) + + def get_dcp_local_seq_lens(seq_lens: torch.Tensor, dcp_world_rank: tuple[int, int]) -> torch.Tensor: """Return interleaved per-rank lengths for global sequence lengths.""" @@ -150,7 +173,7 @@ def update_dcp_metadata(attn_metadata, step_context) -> None: prefix_lens=attn_metadata.kv_seqlens - attn_metadata.q_seqlens, prefix_limit=prefix_limit, block_size=step_context.cache_config.block_size, - head_dim=step_context.model_config.head_dim, + kv_width=get_dcp_kv_width(step_context.model_config), dcp_world_rank=dcp_world_rank, ) topk = step_context.model_config.mla_index_topk diff --git a/lmdeploy/pytorch/backends/cuda/attention/__init__.py b/lmdeploy/pytorch/backends/cuda/attention/__init__.py index 1c09484b6b..fa4b8a17b1 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/__init__.py +++ b/lmdeploy/pytorch/backends/cuda/attention/__init__.py @@ -126,6 +126,11 @@ def _build_paged_attention(spec: PagedAttentionBuildSpec) -> TritonAttentionImpl logger.debug('Build FlashMLAImpl Attention') from .mla import FlashMLAImpl return FlashMLAImpl(use_fa3=use_fa3, **common_args) + + from lmdeploy.pytorch.distributed import get_dcp_world_rank + if get_dcp_world_rank()[0] > 1: + from .cp import DCPAttentionImpl + return DCPAttentionImpl(use_fa3=enable_fa3, **common_args) elif enable_fa3: logger.debug('Build FA3Impl Attention') from .fa3 import FA3Impl diff --git a/lmdeploy/pytorch/backends/cuda/attention/cp.py b/lmdeploy/pytorch/backends/cuda/attention/cp.py index dbae22fa95..921156934f 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/cp.py +++ b/lmdeploy/pytorch/backends/cuda/attention/cp.py @@ -1,8 +1,31 @@ # Copyright (c) OpenMMLab. All rights reserved. -"""CUDA DCP query gathering and attention-output merging.""" +"""CUDA DCP attention implementation and shared collectives.""" import torch +from lmdeploy.pytorch.backends.cp_utils import DCPPrefixChunk, get_dcp_local_causal_seq_lens + +from .default import TritonAttentionImpl, TritonAttentionMetaBuilder, TritonAttentionMetadata + + +def gather_dcp_prefix_kv(local_kv: torch.Tensor, chunk: DCPPrefixChunk, + *, zero_padding: bool = False) -> torch.Tensor: + """Gather a flattened prefix chunk and restore global token order. + + Each rank supplies [chunk.local_capacity, ...] with valid tokens packed by request. The result has + [requests * chunk.size, ...] capacity and valid tokens packed according to chunk.cu_seqlens. + Set zero_padding for consumers such as FA3 that may load masked V lanes: 0 * NaN can contaminate valid outputs. + """ + from lmdeploy.pytorch.distributed import all_gather_into_tensor + from lmdeploy.pytorch.kernels.cuda.dcp import reorder_dcp_prefill_kv + + gathered = local_kv.new_empty(chunk.kv_seqlens.numel() * chunk.size, *local_kv.shape[1:]) + all_gather_into_tensor(gathered, local_kv, group='dcp') + context = torch.zeros_like(gathered) if zero_padding else torch.empty_like(gathered) + reorder_dcp_prefill_kv(gathered, context, chunk_kv_seqlens=chunk.kv_seqlens, + kv_start_loc=chunk.cu_seqlens[:-1], local_lens=chunk.local_kv_seqlens) + return context + def gather_dcp_query(query: torch.Tensor, *, dcp_world_size: int) -> torch.Tensor: """Gather [tokens, local_heads, dim] queries along the head axis. @@ -51,3 +74,110 @@ def merge_dcp_attention(local_output: torch.Tensor, contribution.size(2)) reduce_scatter_tensor(scattered_output, contribution, group='dcp') return scattered_output.transpose(0, 1) + + +class DCPAttentionImpl(TritonAttentionImpl): + """Keep TP query heads local in prefill and gather them for decode. + + Model configuration ensures that every rank in a DCP group holds the same KV head. Cached prefill gathers bounded + K/V chunks; decode evaluates each context shard and reduce-scatters softmax-corrected outputs to the TP heads. + """ + + def __init__(self, num_heads: int, head_size: int, *, use_fa3: bool = False, **kwargs): + super().__init__(num_heads=num_heads, head_size=head_size, **kwargs) + self.use_fa3 = use_fa3 + if use_fa3: + from lmdeploy.pytorch.third_party.flash_attn_interface import flash_attn_varlen_func + self._fa3_prefill = flash_attn_varlen_func + from lmdeploy.pytorch.distributed import get_dist_manager + communicator = get_dist_manager().current_context().dcp_group.communicator + if communicator is not None: + communicator.prepare_query_gather(self.num_heads, self.head_size) + + def get_step_metadata_provider(self): + """Use common DCP lengths without a separate kernel scheduler.""" + return TritonAttentionMetaBuilder() + + def supports_piecewise_cuda_graph(self) -> bool: + return True + + def _decode(self, query: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, + metadata: TritonAttentionMetadata, k_scales_zeros: torch.Tensor = None, + v_scales_zeros: torch.Tensor = None) -> torch.Tensor: + """Evaluate interleaved shards and restore rank-local query heads.""" + query = gather_dcp_query(query, dcp_world_size=self.dcp_world_size) + query_len = query.size(0) // metadata.q_seqlens.numel() + local_lens = metadata.dcp_local_kv_seqlens + blocks = metadata.block_offsets + if query_len > 1: + # A global query step need not advance this shard's causal length. + local_lens = get_dcp_local_causal_seq_lens( + metadata.kv_seqlens, query_len, (self.dcp_world_size, self.dcp_rank)) + blocks = blocks.repeat_interleave(query_len, dim=0) + output, lse = self.paged_attention_fwd( + query, k_cache, v_cache, cache_seqlens=local_lens, + page_table=blocks, max_seqlen_q=1, softmax_scale=self.scale, + softcap=self.logit_softcapping, return_lse=True, quant_policy=metadata.quant_policy, + k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros) + return merge_dcp_attention(output, lse, valid_counts=local_lens, + dcp_world_rank=(self.dcp_world_size, self.dcp_rank)) + + def _prefill_attention(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, *, + cu_q: torch.Tensor, cu_k: torch.Tensor, max_q: int, max_k: int, + causal: bool, return_lse: bool): + """Evaluate one prefill partition with optional natural-log LSE.""" + kwargs = dict(cu_seqlens_q=cu_q, cu_seqlens_k=cu_k, + max_seqlen_q=max_q, max_seqlen_k=max_k, + softmax_scale=self.scale, softcap=max(0.0, self.logit_softcapping), causal=causal) + if self.use_fa3: + return self._fa3_prefill(query, key, value, return_lse=return_lse, **kwargs) + return self.flash_attention_fwd(query, key, value, kv_layout='shd', return_lse=return_lse, **kwargs) + + def _gather_prefix(self, k_cache: torch.Tensor, v_cache: torch.Tensor, + metadata: TritonAttentionMetadata, chunk: DCPPrefixChunk, + dtype: torch.dtype, k_scales_zeros: torch.Tensor = None, + v_scales_zeros: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]: + """Gather and reorder separate K/V caches within the chunk budget.""" + local_k, local_v = self.flatten_kv_cache( + k_cache, v_cache, chunk.local_kv_seqlens[self.dcp_rank], + metadata.block_offsets[:, chunk.local_block_slice(k_cache.size(1))], + start_loc=chunk.local_cu_seqlens[:-1], out_size=chunk.local_capacity, + out_dtype=dtype, flatten_kv_layout='shd', quant_policy=metadata.quant_policy, + k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros) + return (gather_dcp_prefix_kv(local_k, chunk, zero_padding=True), + gather_dcp_prefix_kv(local_v, chunk, zero_padding=True)) + + def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + k_cache: torch.Tensor, v_cache: torch.Tensor, attn_metadata: TritonAttentionMetadata, + k_scales_zeros: torch.Tensor = None, v_scales_zeros: torch.Tensor = None, + learnable_sink: torch.Tensor = None, inplace: bool = True, **kwargs) -> torch.Tensor: + """Write owned KV tokens, then evaluate GQA prefill or decode.""" + assert learnable_sink is None, 'GQA DCP does not support attention sinks' + max_q = (query.size(0) // attn_metadata.q_seqlens.numel() + if attn_metadata.is_decoding else self._get_max_q_seqlen(query, attn_metadata)) + if key is not None and value is not None: + self._fill_kv_cache_impl(key, value, k_cache, v_cache, attn_metadata, max_q, + k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros) + if attn_metadata.is_decoding: + return self._decode(query, k_cache, v_cache, attn_metadata, + k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros) + + assert key is not None and value is not None, 'GQA DCP prefill requires current K/V' + chunks = attn_metadata.dcp_prefix_chunks + result = self._prefill_attention( + query, key, value, cu_q=attn_metadata.cu_seqlens_q, cu_k=attn_metadata.cu_seqlens_q, + max_q=max_q, max_k=max_q, causal=True, return_lse=bool(chunks)) + if not chunks: + return result + from lmdeploy.pytorch.kernels.cuda.dcp import merge_attention_states + output, lse = result + del result + for chunk in chunks: + context_k, context_v = self._gather_prefix(k_cache, v_cache, attn_metadata, chunk, query.dtype, + k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros) + chunk_output, chunk_lse = self._prefill_attention( + query, context_k, context_v, cu_q=attn_metadata.cu_seqlens_q, cu_k=chunk.cu_seqlens, + max_q=max_q, max_k=chunk.size, causal=False, return_lse=True) + output, lse = merge_attention_states(output, lse, chunk_output, chunk_lse) + del context_k, context_v, chunk_output, chunk_lse + return output diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index b03326c080..bf2e578311 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -14,7 +14,7 @@ from lmdeploy.utils import get_logger from ..step_metadata import CudaAttentionMetaBuilder -from .cp import gather_dcp_query, merge_dcp_attention +from .cp import gather_dcp_prefix_kv, gather_dcp_query, merge_dcp_attention from .default import TritonAttentionImpl, TritonAttentionMetadata logger = get_logger('lmdeploy') @@ -525,23 +525,14 @@ def _gather_dcp_prefix_chunk( v_scales_zeros: torch.Tensor = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Gather one globally ordered cached-prefix chunk.""" - from lmdeploy.pytorch.distributed import all_gather_into_tensor - from lmdeploy.pytorch.kernels.cuda.dcp import reorder_dcp_prefill_kv - - dcp_world_size = self.dcp_world_size - block_size = k_cache.size(1) - virtual_block_size = block_size * dcp_world_size - local_capacity = chunk.kv_seqlens.numel() * chunk.size // dcp_world_size - block_start = chunk.start // virtual_block_size - num_blocks = chunk.size // virtual_block_size local_meta = replace( attn_metadata, - block_offsets=attn_metadata.block_offsets[:, block_start:block_start + num_blocks], + block_offsets=attn_metadata.block_offsets[:, chunk.local_block_slice(k_cache.size(1))], kv_start_loc=chunk.local_cu_seqlens[:-1], kv_seqlens=chunk.local_kv_seqlens[self.dcp_rank], cu_seqlens_k=chunk.local_cu_seqlens, - kv_flatten_size=local_capacity, - max_kv_seqlen=chunk.size // dcp_world_size, + kv_flatten_size=chunk.local_capacity, + max_kv_seqlen=chunk.size // self.dcp_world_size, ) local_k, _ = self._flatten_prefill_kv_cache( k_cache, @@ -552,18 +543,7 @@ def _gather_dcp_prefix_chunk( k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, ) - gathered = local_k.new_empty(dcp_world_size * local_capacity, - *local_k.shape[1:]) - all_gather_into_tensor(gathered, local_k, group='dcp') - - context_k = local_k.new_empty(chunk.kv_seqlens.numel() * chunk.size, *local_k.shape[1:]) - reorder_dcp_prefill_kv( - gathered, - context_k, - chunk_kv_seqlens=chunk.kv_seqlens, - kv_start_loc=chunk.cu_seqlens[:-1], - local_lens=chunk.local_kv_seqlens, - ) + context_k = gather_dcp_prefix_kv(local_k, chunk) return context_k, chunk.cu_seqlens def _prefill_dcp_context( diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py index 1ee96e1c92..2e38c39c58 100644 --- a/lmdeploy/pytorch/config.py +++ b/lmdeploy/pytorch/config.py @@ -620,7 +620,11 @@ def from_hf_config( model_config.dist_config = dist_config if dist_config.dcp > 1: - assert model_config.use_flash_mla, 'DCP requires FlashMLA attention' + if not model_config.use_flash_mla: + replicas = model_config.num_replicate_key_value_heads + assert replicas % dist_config.dcp == 0, ( + 'GQA DCP groups must share replicated KV heads: ' + f'KV replication factor {replicas} must be divisible by dcp {dist_config.dcp}') assert model_config.sliding_window < 0, 'DCP does not support sliding-window attention' if model_config.mla_index_topk is not None: from lmdeploy.pytorch import envs diff --git a/lmdeploy/pytorch/engine/executor/__init__.py b/lmdeploy/pytorch/engine/executor/__init__.py index eec27c22c0..a1e370c6da 100644 --- a/lmdeploy/pytorch/engine/executor/__init__.py +++ b/lmdeploy/pytorch/engine/executor/__init__.py @@ -107,6 +107,9 @@ def build_executor( # Finalize cache policy before any executor copies configs to workers or # builds backend operators. Target and memory models share CacheConfig. + if (dist_config.dcp > 1 and not model_config.use_flash_mla + and cache_config.quant_policy not in (QuantPolicy.NONE, QuantPolicy.FP8, QuantPolicy.FP8_E5M2)): + raise ValueError('GQA DCP supports only unquantized or per-tensor FP8 KV cache') shared_cache_models = [model_config] if memdecode_config := misc_config.memdecode_config: shared_cache_models.append(memdecode_config.memory_model_config) diff --git a/lmdeploy/pytorch/engine/executor/base.py b/lmdeploy/pytorch/engine/executor/base.py index 8fab1fe4f5..2f26201c13 100644 --- a/lmdeploy/pytorch/engine/executor/base.py +++ b/lmdeploy/pytorch/engine/executor/base.py @@ -221,7 +221,7 @@ def _get_dcp_workspace_size(self, num_prefill_tokens: int) -> int: if self.cache_config.dcp <= 1: return 0 - from lmdeploy.pytorch.backends.cp_utils import get_dcp_workspace_size + from lmdeploy.pytorch.backends.cp_utils import get_dcp_kv_width, get_dcp_workspace_size config = self.cache_config model = self.model_config @@ -240,6 +240,7 @@ def _get_dcp_workspace_size(self, num_prefill_tokens: int) -> int: dcp_size=config.dcp, topk=model.mla_index_topk, score_workspace_bytes=self._get_dsa_score_workspace_size(), + kv_width=get_dcp_kv_width(model), ) def _get_runtime_size(self, free_mems: list[int], cache_block_sizes: list[_WorkerCachePlanSizes], diff --git a/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py b/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py index 20cc7097ac..15c91184c8 100644 --- a/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py +++ b/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py @@ -168,6 +168,8 @@ def _fill_kv_cache_fp8_scalar_kernel( stride_vch: tl.constexpr, stride_vcd: tl.constexpr, stride_boff, + DCP_SIZE: tl.constexpr, + DCP_RANK: tl.constexpr, BLOCK: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_DV: tl.constexpr, @@ -181,23 +183,22 @@ def _fill_kv_cache_fp8_scalar_kernel( q_seqlen = tl.load(QSeqLens + batch_id) kv_seqlen = tl.load(KVSeqLens + batch_id) history_seqlen = kv_seqlen - q_seqlen - kv_block_id = history_seqlen // BLOCK + block_id + kv_block_id = history_seqlen // (BLOCK * DCP_SIZE) + block_id if kv_seqlen <= 0: return - if kv_block_id * BLOCK >= kv_seqlen: + if kv_block_id * BLOCK * DCP_SIZE >= kv_seqlen: return if is_decoding: - page_offs = tl.full((1, ), history_seqlen % BLOCK, dtype=tl.int32) - kv_mask = tl.full((1, ), 1, dtype=tl.int1) + page_offs = tl.full((1, ), (history_seqlen // DCP_SIZE) % BLOCK, dtype=tl.int32) + kv_mask = tl.full((1, ), history_seqlen % DCP_SIZE == DCP_RANK, dtype=tl.int1) q_offs = tl.full((1, ), q_startloc, dtype=tl.int32) else: page_offs = tl.arange(0, BLOCK) - kv_offs = kv_block_id * BLOCK + page_offs + kv_offs = (kv_block_id * BLOCK + page_offs) * DCP_SIZE + DCP_RANK kv_mask = (kv_offs >= history_seqlen) & (kv_offs < kv_seqlen) - token_off = q_startloc + kv_block_id * BLOCK - history_seqlen - q_offs = token_off + page_offs + q_offs = q_startloc + kv_offs - history_seqlen block_off = tl.load(BlockOffsets + batch_id * stride_boff + kv_block_id) block_off = block_off.to(tl.int64) @@ -830,8 +831,8 @@ def fill_kv_cache(k_states: Tensor, if v_states.size(-1) == 0: head_dim_v = 0 - if dcp_size > 1 and quant_policy != QuantPolicy.NONE: - raise ValueError('DCP generic KV fill currently supports only unquantized cache payloads') + if dcp_size > 1 and quant_policy not in (QuantPolicy.NONE, QuantPolicy.FP8, QuantPolicy.FP8_E5M2): + raise ValueError('DCP generic KV fill supports only unquantized or per-tensor FP8 cache payloads') if max_q_seq_length == 1: max_num_blocks = 1 else: @@ -952,6 +953,8 @@ def fill_kv_cache(k_states: Tensor, stride_vch=v_caches.stride(h_dim), stride_vcd=v_caches.stride(d_dim), stride_boff=block_offsets.stride(0), + DCP_SIZE=dcp_size, + DCP_RANK=dcp_rank, BLOCK=BLOCK, BLOCK_D=BLOCK_D, BLOCK_DV=BLOCK_DV, diff --git a/lmdeploy/pytorch/kernels/cuda/pagedattention.py b/lmdeploy/pytorch/kernels/cuda/pagedattention.py index 1b62c591ca..a9637a0929 100644 --- a/lmdeploy/pytorch/kernels/cuda/pagedattention.py +++ b/lmdeploy/pytorch/kernels/cuda/pagedattention.py @@ -576,6 +576,8 @@ def _reduce_split_kernel( acc_ptr, out_ptr, sinks_ptr, + lse_ptr, + cache_seqlens_ptr, stride_ak, stride_abs, stride_ah, @@ -587,6 +589,8 @@ def _reduce_split_kernel( SPLIT_K: tl.constexpr, BLOCK_DV: tl.constexpr, USE_PDL: tl.constexpr, + NUM_HEADS: tl.constexpr, + SEQ_LEN: tl.constexpr, ): """Second step kernel of split k attention.""" cur_batch = tl.program_id(1) @@ -609,6 +613,13 @@ def _reduce_split_kernel( if USE_PDL: tl.extra.cuda.gdc_wait() + if lse_ptr is not None: + # Empty DCP shards do not write split scratch. Never read it. + if tl.load(cache_seqlens_ptr + cur_batch // SEQ_LEN) <= 0: + tl.store(out_ptr + offs_dv * stride_od, 0.0, mask=mask_dv) + tl.store(lse_ptr + cur_batch * NUM_HEADS + cur_head, -float('inf')) + return + m_k = tl.load(acc_ptr + offs_mi) l_k = tl.load(acc_ptr + offs_mi + 1) # (m_k[:, None] > -float('inf')) produce invalid mask for triton 3.5.1 @@ -625,6 +636,10 @@ def _reduce_split_kernel( if sinks_ptr is not None: sink = tl.load(sinks_ptr + cur_head).to(l_sum.dtype) l_sum = l_sum + tl.exp2(sink * tl_log2(math.e) - m_max) + if lse_ptr is not None: + # Split accumulators use log2; the DCP merge consumes natural-log LSE. + lse = (m_max + tl_log2(l_sum)) / tl_log2(math.e) + tl.store(lse_ptr + cur_batch * NUM_HEADS + cur_head, lse) acc = acc / (l_sum + 1e-10) out_offs = offs_dv * stride_od @@ -794,6 +809,7 @@ def flash_attn_with_kvcache( quant_policy: QuantPolicy = QuantPolicy.NONE, sinks: Tensor = None, kv_layout: str = 'bshd', + return_lse: bool = False, ): """Paged Attention forward. @@ -804,8 +820,12 @@ def flash_attn_with_kvcache( per-tensor key scale for per-tensor FP8 KV cache. v_scales_zeros: Per-token scale/zero metadata for int KV cache, or per-tensor value scale for per-tensor FP8 KV cache. + return_lse: Return ``(output, lse)`` with natural-log FP32 LSE shaped + [tokens, heads]. Empty cache rows return zero output and -inf LSE. """ + if return_lse and quant_policy == QuantPolicy.TURBO_QUANT: + raise NotImplementedError('TurboQuant paged attention does not return LSE') if kv_layout == 'bshd': b_dim, s_dim, h_dim, d_dim = (0, 1, 2, 3) elif kv_layout == 'bhsd': @@ -1032,6 +1052,7 @@ def flash_attn_with_kvcache( num_warps = 2 grid = (head, num_tokens) + lse = torch.empty((num_tokens, head), device=q.device, dtype=torch.float32) if return_lse else None if quant_policy == QuantPolicy.INT4: Lv *= 2 BLOCK_DV *= 2 @@ -1063,6 +1084,8 @@ def flash_attn_with_kvcache( _reduce_split_kernel[grid](acc, o, sinks, + lse, + cache_seqlens, stride_ak=acc.stride(2), stride_abs=acc.stride(0), stride_ah=acc.stride(1), @@ -1074,7 +1097,9 @@ def flash_attn_with_kvcache( head_size_v=Lv, BLOCK_DV=BLOCK_DV, USE_PDL=use_pdl, + NUM_HEADS=head, + SEQ_LEN=seq_len, num_warps=num_warps, num_stages=1, **pdl_launch) - return o + return (o, lse) if return_lse else o diff --git a/lmdeploy/pytorch/third_party/flash_attn_interface.py b/lmdeploy/pytorch/third_party/flash_attn_interface.py index 5dbcc2fbec..bf9984ff11 100644 --- a/lmdeploy/pytorch/third_party/flash_attn_interface.py +++ b/lmdeploy/pytorch/third_party/flash_attn_interface.py @@ -1,13 +1,31 @@ # Copyright (c) OpenMMLab. All rights reserved. import functools +import torch + from flash_attn_interface import flash_attn_varlen_func as _flash_attn_varlen_func from flash_attn_interface import flash_attn_with_kvcache as _flash_attn_with_kvcache @functools.wraps(_flash_attn_varlen_func) -def flash_attn_varlen_func(*args, **kwargs): +def flash_attn_varlen_func(*args, return_lse: bool = False, **kwargs): + """Normalize FA3 returns with optional natural-log LSE. + + LSE is FP32 [tokens, heads]. Callers requesting LSE provide cu_seqlens_q and cu_seqlens_k as keyword arguments. + Empty request partitions have -inf LSE. + """ + if return_lse: + kwargs['return_attn_probs'] = True output = _flash_attn_varlen_func(*args, **kwargs) + if return_lse: + from lmdeploy.pytorch.kernels.cuda.dcp import sanitize_dcp_lse + + output, lse = output + cu_q, cu_k = kwargs['cu_seqlens_q'], kwargs['cu_seqlens_k'] + valid_counts = torch.repeat_interleave(cu_k[1:] - cu_k[:-1], cu_q[1:] - cu_q[:-1], + output_size=output.size(0)) + lse = sanitize_dcp_lse(lse.transpose(0, 1), valid_counts) + return output, lse if isinstance(output, tuple): # for old api return output[0] diff --git a/tests/pytorch/config/test_gqa_dcp.py b/tests/pytorch/config/test_gqa_dcp.py new file mode 100644 index 0000000000..390481fe41 --- /dev/null +++ b/tests/pytorch/config/test_gqa_dcp.py @@ -0,0 +1,84 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lmdeploy.messages import QuantPolicy +from lmdeploy.pytorch.config import DistConfig, ModelConfig + + +@pytest.mark.parametrize('dcp_size', [1, 2]) +@pytest.mark.parametrize('policy', list(QuantPolicy)) +def test_gqa_dcp_executor_cache_policy(monkeypatch, dcp_size, policy): + from lmdeploy.pytorch.engine import executor + from lmdeploy.pytorch.engine.executor import mp_executor + + model = SimpleNamespace(use_flash_mla=False, mla_index_topk=None) + monkeypatch.setattr(ModelConfig, 'from_pretrained', Mock(return_value=model)) + constructor = Mock() + monkeypatch.setattr(mp_executor, 'MPExecutor', constructor) + cache = SimpleNamespace(block_size=16, quant_policy=policy) + misc = SimpleNamespace(hf_overrides=None, model_format=None, memdecode_config=None, empty_init=False) + kwargs = dict(model_path='', cache_config=cache, backend_config=None, + dist_config=DistConfig(tp=2, dcp=dcp_size), misc_config=misc, + distributed_executor_backend='mp') + if dcp_size > 1 and policy in (QuantPolicy.INT4, QuantPolicy.INT8, QuantPolicy.TURBO_QUANT): + with pytest.raises(ValueError, match='unquantized or per-tensor FP8'): + executor.build_executor(**kwargs) + constructor.assert_not_called() + else: + assert executor.build_executor(**kwargs) is constructor.return_value + assert constructor.call_args.kwargs['cache_config'].quant_policy == policy + + +@pytest.mark.parametrize('num_heads', [32, 64]) +@pytest.mark.parametrize(('tp', 'dcp', 'supported'), [(4, 2, False), (8, 2, True), (8, 4, False), (16, 4, True)]) +def test_qwen_gqa_dcp_head_replication(num_heads, tp, dcp, supported): + config = SimpleNamespace( + architectures=['Qwen3MoeForCausalLM'], model_type='qwen3_moe', hidden_size=2048, + num_hidden_layers=1, num_attention_heads=num_heads, num_key_value_heads=4, + head_dim=128, bos_token_id=1, eos_token_id=2, vocab_size=151936) + if not supported: + with pytest.raises(AssertionError, match='share replicated KV heads'): + ModelConfig.from_hf_config(config, dist_config=DistConfig(tp=tp, dcp=dcp)) + else: + model = ModelConfig.from_hf_config(config, dist_config=DistConfig(tp=tp, dcp=dcp)) + assert model.get_num_qkv_head_by_tp() == (num_heads // tp, 1) + assert model.num_replicate_key_value_heads == tp // 4 + + +@pytest.mark.parametrize('fa3', [True, False]) +def test_gqa_dcp_dispatch(monkeypatch, fa3): + from lmdeploy.pytorch import distributed + from lmdeploy.pytorch.backends.attention import PagedAttentionBuildSpec + from lmdeploy.pytorch.backends.cuda import attention + from lmdeploy.pytorch.backends.cuda.attention import cp + constructor = Mock() + monkeypatch.setattr(cp, 'DCPAttentionImpl', constructor) + monkeypatch.setattr(distributed, 'get_dcp_world_rank', lambda: (2, 0)) + monkeypatch.setattr(attention, '_enable_fa3', lambda *args: fa3) + spec = PagedAttentionBuildSpec(num_heads=4, head_dim=128, num_kv_heads=1, v_head_dim=128, + scale=None, alibi=False, sliding_window=None, logit_softcapping=0.0, + causal=True, use_flash_mla=False, mla_index_topk=None, + learnable_sink=False, block_sparse_size=1) + assert attention._build_paged_attention(spec) is constructor.return_value + assert constructor.call_args.kwargs['use_fa3'] == fa3 + + +def test_gqa_prefix_workspace_accounts_for_separate_kv(monkeypatch): + from lmdeploy.pytorch import distributed + from lmdeploy.pytorch.backends import cp_utils + monkeypatch.setattr(distributed, 'get_dcp_world_rank', lambda: (2, 0)) + planner = Mock(return_value=()) + monkeypatch.setattr(cp_utils, 'build_dcp_prefix_chunks', planner) + model = ModelConfig(hidden_size=2048, num_layers=1, num_attention_heads=32, + num_key_value_heads=8, head_dim=128, k_head_dim=128, v_head_dim=128, + bos_token_id=1, eos_token_id=[2], dist_config=DistConfig(tp=8, dcp=2)) + metadata = SimpleNamespace(kv_seqlens=torch.tensor([65]), q_seqlens=torch.tensor([1]), + is_decoding=False, kv_flatten_size=65, max_kv_seqlen=65) + step = SimpleNamespace(input_ids=torch.tensor([[1]]), model_config=model, + cache_config=SimpleNamespace(block_size=16)) + cp_utils.update_dcp_metadata(metadata, step) + assert planner.call_args.kwargs['kv_width'] == 256 diff --git a/tests/pytorch/kernel/test_fill_kv_cache.py b/tests/pytorch/kernel/test_fill_kv_cache.py index 2a711f45bb..8c191858b3 100644 --- a/tests/pytorch/kernel/test_fill_kv_cache.py +++ b/tests/pytorch/kernel/test_fill_kv_cache.py @@ -855,6 +855,65 @@ def quant_policy(self): yield QuantPolicy.FP8_E5M2 +@pytest.mark.skipif(not torch.cuda.is_available(), reason='CUDA required') +@pytest.mark.parametrize(('quant_policy', 'cache_dtype'), [ + (QuantPolicy.FP8, torch.float8_e4m3fn), (QuantPolicy.FP8_E5M2, torch.float8_e5m2)]) +@pytest.mark.parametrize('dcp_size', [2, 4]) +@pytest.mark.parametrize('kv_layout', ['bshd', 'bhsd']) +@pytest.mark.parametrize('q_lengths', [[1, 1, 1, 1], [0, 1, 5, 35]]) +def test_fill_fp8_dcp_matches_unsharded_bytes(quant_policy, cache_dtype, dcp_size, kv_layout, q_lengths): + """Check ownership and untouched slots with unequal K/V widths and strided + inputs.""" + from lmdeploy.pytorch.kernels.cuda.fill_kv_cache import fill_kv_cache + + _skip_unsupported_triton_fp8_dtype(cache_dtype) + torch.manual_seed(101) + block_size, num_heads = 16, 2 + history = [0, 31, 32, 65] + kv_lengths = [h + q for h, q in zip(history, q_lengths)] + q_lens = torch.tensor(q_lengths, device='cuda', dtype=torch.int32) + kv_lens = torch.tensor(kv_lengths, device='cuda', dtype=torch.int32) + starts = torch.nn.functional.pad(q_lens.cumsum(0, dtype=torch.int32), (1, 0))[:-1] + k = torch.randn(sum(q_lengths) * 2, num_heads, 128, device='cuda', dtype=torch.bfloat16)[::2] + v = torch.randn(sum(q_lengths) * 2, num_heads, 96, device='cuda', dtype=torch.bfloat16)[::2] + # Exercise saturation as well as non-unit, independent K/V scales. + k[0, 0, :2] = torch.tensor([1e5, -1e5], device='cuda') + scales = dict(k_scales_zeros=torch.tensor(0.25, device='cuda'), + v_scales_zeros=torch.tensor(0.5, device='cuda')) + global_blocks, global_count = _dcp_block_offsets(kv_lengths, block_size, 1, k.device) + local_blocks, local_count = _dcp_block_offsets(kv_lengths, block_size, dcp_size, k.device) + + def allocate(count, width): + cache = torch.full((count, block_size, num_heads, width), 7, device='cuda', dtype=cache_dtype) + return cache if kv_layout == 'bshd' else cache.transpose(1, 2).contiguous() + + def canonical(cache): + return cache if kv_layout == 'bshd' else cache.transpose(1, 2) + + global_k, global_v = allocate(global_count, 128), allocate(global_count, 96) + fill_kv_cache(k, v, global_k, global_v, starts, q_lens, kv_lens, max(q_lengths), global_blocks, + quant_policy=quant_policy, kv_layout=kv_layout, **scales) + for rank in range(dcp_size): + local_k, local_v = allocate(local_count, 128), allocate(local_count, 96) + fill_kv_cache(k, v, local_k, local_v, starts, q_lens, kv_lens, max(q_lengths), local_blocks, + quant_policy=quant_policy, kv_layout=kv_layout, dcp_size=dcp_size, dcp_rank=rank, **scales) + for actual, global_cache in [(local_k, global_k), (local_v, global_v)]: + expected = torch.full_like(canonical(actual).view(torch.uint8), + torch.tensor(7, dtype=cache_dtype).view(torch.uint8).item()) + for request, kv_len in enumerate(kv_lengths): + num_global = _div_up(kv_len, block_size) + num_local = _div_up(kv_len, block_size * dcp_size) + if num_local == 0: + continue + global_ids = global_blocks[request, :num_global] + local_ids = local_blocks[request, :num_local] + owned = canonical(global_cache).view(torch.uint8)[global_ids].flatten(0, 1)[rank::dcp_size] + padded = expected[local_ids].flatten(0, 1) + padded[:owned.size(0)] = owned + expected[local_ids] = padded.view(num_local, block_size, num_heads, -1) + assert torch.equal(canonical(actual).view(torch.uint8), expected) + + def _dcp_owned_current_tokens(states, q_seqlens, history_lens, dcp_size, dcp_rank): """Pack current-chunk tokens owned by one interleaved DCP rank.""" diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py index 5220f24a50..99446f4206 100644 --- a/tests/pytorch/kernel/test_mla_attention.py +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -494,7 +494,8 @@ def test_dcp_cached_prefill_matches_reference_across_chunks(monkeypatch, sparse, max_q_seqlen=3, input_ids=torch.zeros(1, 5), cache_config=SimpleNamespace(block_size=64), - model_config=SimpleNamespace(head_dim=576, mla_index_topk=512 if sparse else None), + model_config=SimpleNamespace(head_dim=576, use_flash_mla=True, + mla_index_topk=512 if sparse else None), kv_quant_policy=0) monkeypatch.setattr(distributed, 'get_dcp_world_rank', lambda: (dcp_size, 0)) # One virtual block per chunk, with an empty second request in every chunk. @@ -631,8 +632,11 @@ def test_dcp_attention_correction_kernel_matches_torch(dtype): @pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA') @pytest.mark.parametrize('dcp_size', [2, 4]) -def test_reorder_dcp_prefill_kv_handles_uneven_requests(monkeypatch, dcp_size): +@pytest.mark.parametrize('gather_prefix', [False, True], ids=['kernel', 'gather']) +def test_reorder_dcp_prefill_kv_handles_uneven_requests(monkeypatch, dcp_size, gather_prefix): + from lmdeploy.pytorch import distributed from lmdeploy.pytorch.backends import cp_utils + from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_prefix_kv from lmdeploy.pytorch.kernels.cuda.dcp import reorder_dcp_prefill_kv device = 'cuda' @@ -643,7 +647,7 @@ def test_reorder_dcp_prefill_kv_handles_uneven_requests(monkeypatch, dcp_size): lambda **kwargs: len(lengths) * (1 + 2 * dcp_size) * 2) chunks = cp_utils.build_dcp_prefix_chunks( prefix_lens=prefix_lens, prefix_limit=max(lengths), block_size=1, - head_dim=1, dcp_world_rank=(dcp_size, 0)) + kv_width=1, dcp_world_rank=(dcp_size, 0)) for chunk in chunks: values = [request * 100 + torch.arange(length, device=device)[chunk.start:chunk.start + chunk.size] for request, length in enumerate(lengths)] @@ -653,15 +657,24 @@ def test_reorder_dcp_prefill_kv_handles_uneven_requests(monkeypatch, dcp_size): owned = torch.cat([value[rank::dcp_size] for value in values]) gathered[rank, :owned.numel(), 0] = owned output = torch.full((len(lengths) * chunk.size, 1), -99, dtype=torch.int32, device=device) - expected = torch.full_like(output, -99) + expected = torch.full_like(output, 0 if gather_prefix else -99) valid_values = torch.cat(values) expected[:valid_values.numel(), 0] = valid_values gathered = gathered.flatten(0, 1) - reorder_dcp_prefill_kv(gathered, - output, - chunk_kv_seqlens=chunk.kv_seqlens, - kv_start_loc=chunk.cu_seqlens[:-1], - local_lens=chunk.local_kv_seqlens) + if gather_prefix: + def all_gather(destination, local, *, group): + assert group == 'dcp' + assert local.size(0) == local_capacity + destination.copy_(gathered) + + monkeypatch.setattr(distributed, 'all_gather_into_tensor', all_gather) + output = gather_dcp_prefix_kv(gathered[:local_capacity], chunk, zero_padding=True) + else: + reorder_dcp_prefill_kv(gathered, + output, + chunk_kv_seqlens=chunk.kv_seqlens, + kv_start_loc=chunk.cu_seqlens[:-1], + local_lens=chunk.local_kv_seqlens) assert torch.equal(output, expected) diff --git a/tests/pytorch/test_gqa_dcp.py b/tests/pytorch/test_gqa_dcp.py new file mode 100644 index 0000000000..7deaddc3fc --- /dev/null +++ b/tests/pytorch/test_gqa_dcp.py @@ -0,0 +1,77 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""CPU checks for GQA DCP prefill return values and temporary lifetimes.""" +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from lmdeploy.pytorch.backends.cuda.attention.cp import DCPAttentionImpl +from lmdeploy.pytorch.backends.cuda.attention.default import TritonAttentionMetadata + + +@pytest.fixture +def prefill_impl(monkeypatch): + # Exercise Python orchestration without initializing CUDA collectives. + impl = object.__new__(DCPAttentionImpl) + monkeypatch.setattr(impl, '_fill_kv_cache_impl', lambda *args, **kwargs: None) + return impl + + +@pytest.mark.parametrize('tuple_return', [False, True], ids=['tensor_api', 'tuple_api']) +def test_fa3_fresh_prefill_returns_tensor(prefill_impl, tuple_return, monkeypatch): + pytest.importorskip('flash_attn_interface') + from lmdeploy.pytorch.third_party import flash_attn_interface + + query = torch.ones(2, 1, 4) + cu_q = torch.tensor([0, 2], dtype=torch.int32) + prefill_impl.use_fa3 = True + prefill_impl.scale = 0.5 + prefill_impl.logit_softcapping = 0.0 + + def fa3(*args, **kwargs): + assert kwargs.get('return_attn_probs', False) is False + return (query, torch.zeros(1, 2)) if tuple_return else query + + monkeypatch.setattr(flash_attn_interface, '_flash_attn_varlen_func', fa3) + prefill_impl._fa3_prefill = flash_attn_interface.flash_attn_varlen_func + metadata = TritonAttentionMetadata( + is_decoding=False, block_offsets=None, cu_seqlens_q=cu_q) + result = prefill_impl.forward(query, query, query, None, None, metadata) + assert result is query + + +def test_cached_prefill_releases_initial_state(prefill_impl, monkeypatch): + from lmdeploy.pytorch.kernels.cuda import dcp + + query = torch.ones(2, 1, 4) + cu_q = torch.tensor([0, 2], dtype=torch.int32) + chunks = tuple(SimpleNamespace(size=2, cu_seqlens=cu_q) for _ in range(2)) + metadata = TritonAttentionMetadata( + is_decoding=False, block_offsets=None, cu_seqlens_q=cu_q, dcp_prefix_chunks=chunks) + initial_refs = [] + calls = 0 + + def prefill(*args, **kwargs): + nonlocal calls + calls += 1 + output = torch.full_like(query, calls) + lse = torch.zeros(query.shape[:2]) + if calls == 1: + initial_refs.extend((weakref.ref(output), weakref.ref(lse))) + return output, lse + + def gather(k_cache, v_cache, metadata, chunk, dtype, **kwargs): + if chunk is chunks[1]: + assert all(ref() is None for ref in initial_refs), 'initial attention state is still retained' + return query, query + + # Allocate new merged states, as the real merge kernel does. The initial + # states must be released before allocating the next chunk's K/V buffers. + monkeypatch.setattr(dcp, 'merge_attention_states', + lambda output, lse, chunk_output, chunk_lse: (output + chunk_output, lse + chunk_lse)) + monkeypatch.setattr(prefill_impl, '_prefill_attention', prefill) + monkeypatch.setattr(prefill_impl, '_gather_prefix', gather) + result = prefill_impl.forward(query, query, query, None, None, metadata) + assert calls == 3 + torch.testing.assert_close(result, torch.full_like(query, 6)) From 1cfd056f59300d9edfb45d3eb4d7bdfdcce79c81 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 21 Sep 2026 09:24:35 +0000 Subject: [PATCH 03/15] test: remove standalone GQA DCP tests --- tests/pytorch/config/test_gqa_dcp.py | 84 ---------------------------- tests/pytorch/test_gqa_dcp.py | 77 ------------------------- 2 files changed, 161 deletions(-) delete mode 100644 tests/pytorch/config/test_gqa_dcp.py delete mode 100644 tests/pytorch/test_gqa_dcp.py diff --git a/tests/pytorch/config/test_gqa_dcp.py b/tests/pytorch/config/test_gqa_dcp.py deleted file mode 100644 index 390481fe41..0000000000 --- a/tests/pytorch/config/test_gqa_dcp.py +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest -import torch - -from lmdeploy.messages import QuantPolicy -from lmdeploy.pytorch.config import DistConfig, ModelConfig - - -@pytest.mark.parametrize('dcp_size', [1, 2]) -@pytest.mark.parametrize('policy', list(QuantPolicy)) -def test_gqa_dcp_executor_cache_policy(monkeypatch, dcp_size, policy): - from lmdeploy.pytorch.engine import executor - from lmdeploy.pytorch.engine.executor import mp_executor - - model = SimpleNamespace(use_flash_mla=False, mla_index_topk=None) - monkeypatch.setattr(ModelConfig, 'from_pretrained', Mock(return_value=model)) - constructor = Mock() - monkeypatch.setattr(mp_executor, 'MPExecutor', constructor) - cache = SimpleNamespace(block_size=16, quant_policy=policy) - misc = SimpleNamespace(hf_overrides=None, model_format=None, memdecode_config=None, empty_init=False) - kwargs = dict(model_path='', cache_config=cache, backend_config=None, - dist_config=DistConfig(tp=2, dcp=dcp_size), misc_config=misc, - distributed_executor_backend='mp') - if dcp_size > 1 and policy in (QuantPolicy.INT4, QuantPolicy.INT8, QuantPolicy.TURBO_QUANT): - with pytest.raises(ValueError, match='unquantized or per-tensor FP8'): - executor.build_executor(**kwargs) - constructor.assert_not_called() - else: - assert executor.build_executor(**kwargs) is constructor.return_value - assert constructor.call_args.kwargs['cache_config'].quant_policy == policy - - -@pytest.mark.parametrize('num_heads', [32, 64]) -@pytest.mark.parametrize(('tp', 'dcp', 'supported'), [(4, 2, False), (8, 2, True), (8, 4, False), (16, 4, True)]) -def test_qwen_gqa_dcp_head_replication(num_heads, tp, dcp, supported): - config = SimpleNamespace( - architectures=['Qwen3MoeForCausalLM'], model_type='qwen3_moe', hidden_size=2048, - num_hidden_layers=1, num_attention_heads=num_heads, num_key_value_heads=4, - head_dim=128, bos_token_id=1, eos_token_id=2, vocab_size=151936) - if not supported: - with pytest.raises(AssertionError, match='share replicated KV heads'): - ModelConfig.from_hf_config(config, dist_config=DistConfig(tp=tp, dcp=dcp)) - else: - model = ModelConfig.from_hf_config(config, dist_config=DistConfig(tp=tp, dcp=dcp)) - assert model.get_num_qkv_head_by_tp() == (num_heads // tp, 1) - assert model.num_replicate_key_value_heads == tp // 4 - - -@pytest.mark.parametrize('fa3', [True, False]) -def test_gqa_dcp_dispatch(monkeypatch, fa3): - from lmdeploy.pytorch import distributed - from lmdeploy.pytorch.backends.attention import PagedAttentionBuildSpec - from lmdeploy.pytorch.backends.cuda import attention - from lmdeploy.pytorch.backends.cuda.attention import cp - constructor = Mock() - monkeypatch.setattr(cp, 'DCPAttentionImpl', constructor) - monkeypatch.setattr(distributed, 'get_dcp_world_rank', lambda: (2, 0)) - monkeypatch.setattr(attention, '_enable_fa3', lambda *args: fa3) - spec = PagedAttentionBuildSpec(num_heads=4, head_dim=128, num_kv_heads=1, v_head_dim=128, - scale=None, alibi=False, sliding_window=None, logit_softcapping=0.0, - causal=True, use_flash_mla=False, mla_index_topk=None, - learnable_sink=False, block_sparse_size=1) - assert attention._build_paged_attention(spec) is constructor.return_value - assert constructor.call_args.kwargs['use_fa3'] == fa3 - - -def test_gqa_prefix_workspace_accounts_for_separate_kv(monkeypatch): - from lmdeploy.pytorch import distributed - from lmdeploy.pytorch.backends import cp_utils - monkeypatch.setattr(distributed, 'get_dcp_world_rank', lambda: (2, 0)) - planner = Mock(return_value=()) - monkeypatch.setattr(cp_utils, 'build_dcp_prefix_chunks', planner) - model = ModelConfig(hidden_size=2048, num_layers=1, num_attention_heads=32, - num_key_value_heads=8, head_dim=128, k_head_dim=128, v_head_dim=128, - bos_token_id=1, eos_token_id=[2], dist_config=DistConfig(tp=8, dcp=2)) - metadata = SimpleNamespace(kv_seqlens=torch.tensor([65]), q_seqlens=torch.tensor([1]), - is_decoding=False, kv_flatten_size=65, max_kv_seqlen=65) - step = SimpleNamespace(input_ids=torch.tensor([[1]]), model_config=model, - cache_config=SimpleNamespace(block_size=16)) - cp_utils.update_dcp_metadata(metadata, step) - assert planner.call_args.kwargs['kv_width'] == 256 diff --git a/tests/pytorch/test_gqa_dcp.py b/tests/pytorch/test_gqa_dcp.py deleted file mode 100644 index 7deaddc3fc..0000000000 --- a/tests/pytorch/test_gqa_dcp.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. -"""CPU checks for GQA DCP prefill return values and temporary lifetimes.""" -import weakref -from types import SimpleNamespace - -import pytest -import torch - -from lmdeploy.pytorch.backends.cuda.attention.cp import DCPAttentionImpl -from lmdeploy.pytorch.backends.cuda.attention.default import TritonAttentionMetadata - - -@pytest.fixture -def prefill_impl(monkeypatch): - # Exercise Python orchestration without initializing CUDA collectives. - impl = object.__new__(DCPAttentionImpl) - monkeypatch.setattr(impl, '_fill_kv_cache_impl', lambda *args, **kwargs: None) - return impl - - -@pytest.mark.parametrize('tuple_return', [False, True], ids=['tensor_api', 'tuple_api']) -def test_fa3_fresh_prefill_returns_tensor(prefill_impl, tuple_return, monkeypatch): - pytest.importorskip('flash_attn_interface') - from lmdeploy.pytorch.third_party import flash_attn_interface - - query = torch.ones(2, 1, 4) - cu_q = torch.tensor([0, 2], dtype=torch.int32) - prefill_impl.use_fa3 = True - prefill_impl.scale = 0.5 - prefill_impl.logit_softcapping = 0.0 - - def fa3(*args, **kwargs): - assert kwargs.get('return_attn_probs', False) is False - return (query, torch.zeros(1, 2)) if tuple_return else query - - monkeypatch.setattr(flash_attn_interface, '_flash_attn_varlen_func', fa3) - prefill_impl._fa3_prefill = flash_attn_interface.flash_attn_varlen_func - metadata = TritonAttentionMetadata( - is_decoding=False, block_offsets=None, cu_seqlens_q=cu_q) - result = prefill_impl.forward(query, query, query, None, None, metadata) - assert result is query - - -def test_cached_prefill_releases_initial_state(prefill_impl, monkeypatch): - from lmdeploy.pytorch.kernels.cuda import dcp - - query = torch.ones(2, 1, 4) - cu_q = torch.tensor([0, 2], dtype=torch.int32) - chunks = tuple(SimpleNamespace(size=2, cu_seqlens=cu_q) for _ in range(2)) - metadata = TritonAttentionMetadata( - is_decoding=False, block_offsets=None, cu_seqlens_q=cu_q, dcp_prefix_chunks=chunks) - initial_refs = [] - calls = 0 - - def prefill(*args, **kwargs): - nonlocal calls - calls += 1 - output = torch.full_like(query, calls) - lse = torch.zeros(query.shape[:2]) - if calls == 1: - initial_refs.extend((weakref.ref(output), weakref.ref(lse))) - return output, lse - - def gather(k_cache, v_cache, metadata, chunk, dtype, **kwargs): - if chunk is chunks[1]: - assert all(ref() is None for ref in initial_refs), 'initial attention state is still retained' - return query, query - - # Allocate new merged states, as the real merge kernel does. The initial - # states must be released before allocating the next chunk's K/V buffers. - monkeypatch.setattr(dcp, 'merge_attention_states', - lambda output, lse, chunk_output, chunk_lse: (output + chunk_output, lse + chunk_lse)) - monkeypatch.setattr(prefill_impl, '_prefill_attention', prefill) - monkeypatch.setattr(prefill_impl, '_gather_prefix', gather) - result = prefill_impl.forward(query, query, query, None, None, metadata) - assert calls == 3 - torch.testing.assert_close(result, torch.full_like(query, 6)) From cced26e204f307faa792d11429d0a717285d8d2e Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 21 Sep 2026 12:11:17 +0000 Subject: [PATCH 04/15] fix: propagate KV quantization policy to draft attention --- lmdeploy/pytorch/spec_decode/proposers/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lmdeploy/pytorch/spec_decode/proposers/base.py b/lmdeploy/pytorch/spec_decode/proposers/base.py index 42d3996773..5721938c8a 100644 --- a/lmdeploy/pytorch/spec_decode/proposers/base.py +++ b/lmdeploy/pytorch/spec_decode/proposers/base.py @@ -73,6 +73,7 @@ def draft_model_forward( model_config=model_config, cache_config=cache_engine.cache_config, kv_caches=kv_caches, + kv_quant_policy=cache_engine.cache_config.quant_policy, ) context.block_caches = cache_engine.block_caches with ctx_mgr.context(context): From a37a37a1282b8e4e246f4b46557c3264e73e4fc5 Mon Sep 17 00:00:00 2001 From: zxy Date: Wed, 23 Sep 2026 10:21:05 +0000 Subject: [PATCH 05/15] refactor: simplify CUDA communicator routing --- lmdeploy/cli/cli.py | 1 + lmdeploy/cli/serve.py | 2 + lmdeploy/cli/utils.py | 8 + lmdeploy/messages.py | 6 +- lmdeploy/pytorch/backends/base.py | 4 +- lmdeploy/pytorch/backends/communicator.py | 63 +- .../pytorch/backends/cuda/attention/cp.py | 22 +- .../pytorch/backends/cuda/attention/mla.py | 6 +- .../backends/cuda/comm/communicator.py | 164 +++-- .../cuda/comm/flashinfer_allreduce.py | 209 ------- .../backends/cuda/comm/symm_mem_allgather.py | 156 ++--- .../backends/cuda/comm/symm_mem_allreduce.py | 52 +- lmdeploy/pytorch/backends/cuda/op_backend.py | 6 +- .../pytorch/backends/default/embedding.py | 3 +- lmdeploy/pytorch/backends/default/linear.py | 3 +- .../pytorch/backends/dlinfer/op_backend.py | 14 - lmdeploy/pytorch/config.py | 6 +- lmdeploy/pytorch/distributed.py | 52 +- .../pytorch/engine/executor/ray_executor.py | 7 +- lmdeploy/pytorch/envs.py | 7 - .../kernels/cuda/symm_mem_allgather.py | 6 +- lmdeploy/pytorch/models/deepseek_v32.py | 18 +- lmdeploy/pytorch/nn/embedding.py | 48 +- lmdeploy/pytorch/nn/linear/base.py | 3 +- lmdeploy/pytorch/nn/norm.py | 26 +- lmdeploy/pytorch/spec_decode/base.py | 1 + tests/pytorch/engine/test_executor_base.py | 2 +- tests/pytorch/spec_decode/test_spec_agent.py | 12 +- tests/pytorch/test_cuda_communicator.py | 578 ++++++++++-------- 29 files changed, 652 insertions(+), 833 deletions(-) delete mode 100644 lmdeploy/pytorch/backends/cuda/comm/flashinfer_allreduce.py diff --git a/lmdeploy/cli/cli.py b/lmdeploy/cli/cli.py index 29ce6d3960..4dc719344c 100644 --- a/lmdeploy/cli/cli.py +++ b/lmdeploy/cli/cli.py @@ -56,6 +56,7 @@ def add_parser_chat(): ArgumentHelper.dllm_block_length(pt_group) ArgumentHelper.prefix_cache_state_budget(pt_group) ArgumentHelper.prefix_cache_decode_state_interval(pt_group) + ArgumentHelper.communication_backend(pt_group) # common engine args dtype_act = ArgumentHelper.dtype(pt_group) tp_act = ArgumentHelper.tp(pt_group) diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index 19707a03af..787c9d0b18 100644 --- a/lmdeploy/cli/serve.py +++ b/lmdeploy/cli/serve.py @@ -145,6 +145,7 @@ def add_parser_api_server(): ArgumentHelper.migration_backend(pt_group) ArgumentHelper.cudagraph_capture_batch_sizes(pt_group) ArgumentHelper.piecewise_cudagraph_max_tokens(pt_group) + ArgumentHelper.communication_backend(pt_group) # multi-node serving args node_rank_act = ArgumentHelper.node_rank(pt_group) num_nodes_act = ArgumentHelper.num_nodes(pt_group) @@ -287,6 +288,7 @@ def api_server(args): enable_return_routed_experts=args.enable_return_routed_experts, distributed_executor_backend=args.distributed_executor_backend, kv_transfer_config=kv_transfer_config, + communication_backend=args.communication_backend, ) else: from lmdeploy.messages import TurbomindEngineConfig diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index 7d53cc4b95..602c8ca447 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -886,6 +886,14 @@ def trust_remote_code(parser): default=False, help='Whether to trust remote code from model repositories.') + @staticmethod + def communication_backend(parser): + """Select native or automatic PyTorch communication.""" + return parser.add_argument( + '--communication-backend', choices=['nccl', 'auto'], default='nccl', + help='PyTorch communication backend. nccl uses native process-group collectives; ' + 'auto opts into eligible optimizations. Default: nccl.') + @staticmethod def kv_transfer_config(parser): """Add external KV-cache connector configuration.""" diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index c8a33f5aa3..c747afbe35 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -540,6 +540,8 @@ class PytorchEngineConfig: dllm_denoising_steps: Dllm denoising steps. dllm_confidence_threshold: dllm unmasking threshold for dynamic unmasking. + communication_backend: ``nccl`` (default) uses native process-group collectives. + ``auto`` opts into eligible communication optimizations. kv_transfer_config: External KV-cache connector configuration. This is supported only by the PyTorch engine. ``None`` disables external KV-cache transfer. @@ -599,6 +601,7 @@ class PytorchEngineConfig: migration_backend: MigrationBackend = MigrationBackend.DLSlime kv_transfer_config: KVTransferConfig | dict[str, Any] | None = None piecewise_cudagraph_max_tokens: int | None = None + communication_backend: Literal['nccl', 'auto'] = 'nccl' def __post_init__(self): """Check input validation.""" @@ -651,7 +654,8 @@ def __post_init__(self): self.kv_transfer_config = KVTransferConfig(**self.kv_transfer_config) elif self.kv_transfer_config is not None and not isinstance(self.kv_transfer_config, KVTransferConfig): raise TypeError('kv_transfer_config must be a KVTransferConfig, dict, or None') - + if self.communication_backend not in ('nccl', 'auto'): + raise ValueError('communication_backend must be nccl or auto') class ResponseType(enum.Enum): """Response type.""" diff --git a/lmdeploy/pytorch/backends/base.py b/lmdeploy/pytorch/backends/base.py index b82983c9bc..c435d72814 100644 --- a/lmdeploy/pytorch/backends/base.py +++ b/lmdeploy/pytorch/backends/base.py @@ -92,14 +92,14 @@ def model_build_context(ctx_mgr): yield @classmethod - def build_communicator(cls, cpu_group, device_group, dist_config, *, group_roles=('tp', )): + def build_communicator(cls, cpu_group, device_group, dist_config, *, group_name='tp'): """Build a device communicator.""" from .communicator import build_communicator return build_communicator( cpu_group=cpu_group, device_group=device_group, dist_config=dist_config, - group_roles=group_roles, + group_name=group_name, ) @staticmethod diff --git a/lmdeploy/pytorch/backends/communicator.py b/lmdeploy/pytorch/backends/communicator.py index 48e8556fde..a02791ef07 100644 --- a/lmdeploy/pytorch/backends/communicator.py +++ b/lmdeploy/pytorch/backends/communicator.py @@ -1,64 +1,49 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import Literal + import torch from torch import distributed as dist class DeviceCommunicator: - """Device-neutral collective operations for one process group.""" + """Device-neutral collective operations for one process group. + + Workspace creation returns a prepared, caller-owned workspace when available. Collective methods execute with native + fallback when an optimized implementation cannot handle the input. + """ def __init__(self, device_group: dist.ProcessGroup): self.device_group = device_group - def supports_optimized_all_reduce(self) -> bool: - """Whether an optimized all-reduce implementation is available.""" - return False - - def supports_fused_all_reduce_residual_rms_norm(self) -> bool: - """Whether fused all-reduce, residual and RMSNorm is available.""" - return False - - def try_fused_all_reduce_residual_rms_norm(self, - input: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float): - """Run fused all-reduce, residual and RMSNorm when eligible.""" - return None - def all_reduce_(self, input: torch.Tensor): """All-reduce ``input`` in place.""" dist.all_reduce(input, group=self.device_group) - def prepare_query_gather(self, num_heads: int, head_size: int): - """Prepare an optional head-axis query gather before graph capture.""" - pass - - def gather_query(self, query: torch.Tensor) -> torch.Tensor: - """Gather query heads across ranks. + def all_gather(self, input: torch.Tensor, *, workspace=None, copy_output: bool = True) -> torch.Tensor: + """Gather 2D inputs along the last dimension on every group rank. - Consume the result on the same stream before the next query gather; CUDA implementations may borrow an arena. + With copy_output=False, optimized output may borrow the workspace until its next use. Consume it on the same + stream. Native gathering always returns an owned output for multi-rank groups. """ world_size = dist.get_world_size(self.device_group) if world_size == 1: - return query - if not query.is_contiguous(): - transposed = query.transpose(0, 1).contiguous() - gathered = transposed.new_empty(world_size * transposed.size(0), *transposed.shape[1:]) - dist.all_gather_into_tensor(gathered, transposed, group=self.device_group) - return gathered.transpose(0, 1).contiguous() - - # Contiguous token-major inputs need only the final head-axis packing. - gathered = query.new_empty(world_size * query.size(0), *query.shape[1:]) - dist.all_gather_into_tensor(gathered, query, group=self.device_group) - gathered = gathered.view(world_size, *query.shape) - return gathered.transpose(0, 1).reshape(query.size(0), -1, query.size(2)).contiguous() + return input + input = input.contiguous() + gathered = input.new_empty(world_size * input.size(0), input.size(1)) + dist.all_gather_into_tensor(gathered, input, group=self.device_group) + return gathered.view(world_size, *input.shape).transpose(0, 1).reshape(input.size(0), -1) + + def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype): + """Create and prepare an optional workspace and return it to its + owner.""" + return None def close(self): - """Release communicator-owned resources.""" - pass + """Release group-owned resources; gather workspaces belong to + callers.""" def build_communicator(cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup, - dist_config, *, group_roles: tuple[str, ...] = ('tp', )): + dist_config, *, group_name: Literal['tp', 'dcp'] = 'tp'): """Build the default process-group communicator.""" return DeviceCommunicator(device_group=device_group) diff --git a/lmdeploy/pytorch/backends/cuda/attention/cp.py b/lmdeploy/pytorch/backends/cuda/attention/cp.py index 921156934f..dba14b1f02 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/cp.py +++ b/lmdeploy/pytorch/backends/cuda/attention/cp.py @@ -27,6 +27,18 @@ def gather_dcp_prefix_kv(local_kv: torch.Tensor, chunk: DCPPrefixChunk, return context +def init_dcp_query_gather(num_heads: int, head_size: int): + """Initialize one query workspace shared by attention layers in this DCP + group.""" + from lmdeploy.pytorch.distributed import get_dist_manager + + group = get_dist_manager().current_context().dcp_group + if group.communicator is not None and group.query_gather_workspace is None: + width = num_heads * head_size * torch.distributed.get_world_size(group.gpu_group) + group.query_gather_workspace = group.communicator.create_all_gather_workspace( + width, device=torch.device('cuda'), dtype=torch.bfloat16) + + def gather_dcp_query(query: torch.Tensor, *, dcp_world_size: int) -> torch.Tensor: """Gather [tokens, local_heads, dim] queries along the head axis. @@ -36,8 +48,9 @@ def gather_dcp_query(query: torch.Tensor, *, dcp_world_size: int) -> torch.Tenso return query from lmdeploy.pytorch.distributed import get_dist_manager - communicator = get_dist_manager().current_context().dcp_group.communicator - return communicator.gather_query(query) + group = get_dist_manager().current_context().dcp_group + output = group.communicator.all_gather(query.flatten(1), workspace=group.query_gather_workspace, copy_output=False) + return output.view(query.size(0), -1, query.size(2)) def merge_dcp_attention(local_output: torch.Tensor, @@ -89,10 +102,7 @@ def __init__(self, num_heads: int, head_size: int, *, use_fa3: bool = False, **k if use_fa3: from lmdeploy.pytorch.third_party.flash_attn_interface import flash_attn_varlen_func self._fa3_prefill = flash_attn_varlen_func - from lmdeploy.pytorch.distributed import get_dist_manager - communicator = get_dist_manager().current_context().dcp_group.communicator - if communicator is not None: - communicator.prepare_query_gather(self.num_heads, self.head_size) + init_dcp_query_gather(self.num_heads, self.head_size) def get_step_metadata_provider(self): """Use common DCP lengths without a separate kernel scheduler.""" diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index bf2e578311..8fdd344c4d 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -255,11 +255,9 @@ def __init__( self.use_fa3 = use_fa3 if self.dcp_world_size > 1: - from lmdeploy.pytorch.distributed import get_dist_manager + from .cp import init_dcp_query_gather - communicator = get_dist_manager().current_context().dcp_group.communicator - if communicator is not None: - communicator.prepare_query_gather(num_heads, head_size) + init_dcp_query_gather(num_heads, head_size) def get_step_metadata_provider(self): """Describe metadata required by this selected implementation.""" diff --git a/lmdeploy/pytorch/backends/cuda/comm/communicator.py b/lmdeploy/pytorch/backends/cuda/comm/communicator.py index 6147386dcb..79128b4aee 100644 --- a/lmdeploy/pytorch/backends/cuda/comm/communicator.py +++ b/lmdeploy/pytorch/backends/cuda/comm/communicator.py @@ -1,118 +1,102 @@ # Copyright (c) OpenMMLab. All rights reserved. +import socket +from typing import Literal + import torch from torch import distributed as dist -from lmdeploy.pytorch import envs as _envs +from lmdeploy.utils import get_logger from ...communicator import DeviceCommunicator -from .flashinfer_allreduce import FlashInferAllReduce from .symm_mem_allreduce import SymmetricMemoryAllReduce +logger = get_logger('lmdeploy') + class CudaCommunicator(DeviceCommunicator): - """Dispatch optional CUDA collectives with a process-group fallback.""" + """Prepare group-owned providers and dispatch by operation and input + shape.""" def __init__(self, cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup, *, - all_reduce_backend: str | None = None, symm_mem_query_gather: bool = False): + group_name: Literal['tp', 'dcp']): super().__init__(device_group=device_group) - self._all_reduce_backend = all_reduce_backend - self._all_reduce = None - if all_reduce_backend == 'flashinfer': - self._all_reduce = FlashInferAllReduce(cpu_group) - elif all_reduce_backend == 'symm_mem': - self._all_reduce = SymmetricMemoryAllReduce(cpu_group) - self._symm_mem_query_gather = symm_mem_query_gather - self._query_gatherer = None - self._query_shape = None - - def supports_optimized_all_reduce(self) -> bool: - """Whether an optimized all-reduce implementation is available.""" - return self._all_reduce is not None and self._all_reduce.is_available() - - def supports_fused_all_reduce_residual_rms_norm(self) -> bool: - """Whether fused all-reduce, residual and RMSNorm is available.""" - return self._all_reduce_backend == 'flashinfer' and self._all_reduce.is_available() - - def try_fused_all_reduce_residual_rms_norm(self, - input: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float): - """Run fused all-reduce, residual and RMSNorm when eligible.""" - if self._all_reduce_backend != 'flashinfer': - return None - return self._all_reduce.fused_all_reduce_residual_rms_norm( - input=input, - residual=residual, - weight=weight, - eps=eps, - ) + self._cpu_group = cpu_group + self._all_reduce_provider = SymmetricMemoryAllReduce(cpu_group) if group_name == 'tp' else None + if self._all_reduce_provider is not None: + self._check_all_reduce_provider() + if self._all_reduce_provider is not None: + self._all_reduce_provider.prepare() + self._check_all_reduce_provider() + if dist.get_rank(device_group) == 0: + logger.info('Communication backend: auto group=%s symmetric_all_reduce=%s device=%s', + group_name, self._all_reduce_provider is not None, torch.cuda.get_device_name()) + + def _check_all_reduce_provider(self): + """Keep or retire symmetric all-reduce consistently on every rank.""" + provider = self._all_reduce_provider + flags = [None] * dist.get_world_size(self._cpu_group) + dist.all_gather_object(flags, provider.is_available(), group=self._cpu_group) + if not all(flags): + provider.close() + self._all_reduce_provider = None + logger.info('Symmetric-memory all-reduce unavailable; using process-group fallback') def all_reduce_(self, input: torch.Tensor): - """Dispatch all-reduce through optimized CUDA backends.""" - if self._all_reduce is not None and self._all_reduce.all_reduce_(input): + if self._all_reduce_provider is not None and self._all_reduce_provider.all_reduce_(input): return super().all_reduce_(input) - def close(self): - """Release communicator-owned workspaces.""" - if self._all_reduce is not None: - self._all_reduce.close() - if self._query_gatherer is not None: - self._query_gatherer.release() - self._query_gatherer = None - - def prepare_query_gather(self, num_heads: int, head_size: int): - """Share one bounded query arena across attention layers and MTP.""" - if self._query_gatherer is not None: - return - from .symm_mem_allgather import MultimemAllGatherer - - self._query_shape = (num_heads, head_size) - self._query_gatherer = MultimemAllGatherer( - self.device_group, dist.get_rank(self.device_group), - num_heads * head_size * dist.get_world_size(self.device_group), - torch.device('cuda', torch.cuda.current_device()), torch.bfloat16, - enabled=self._symm_mem_query_gather, capacity_bytes=64 * 1024 * 1024) - - def gather_query(self, query: torch.Tensor) -> torch.Tensor: - """Gather query heads, borrowing the symmetric arena when eligible.""" - if (self._query_gatherer is not None and query.dtype == torch.bfloat16 - and tuple(query.shape[1:]) == self._query_shape and query.is_contiguous()): - # Same-stream attention finishes before the next gather's entry - # barrier permits any rank to overwrite the borrowed arena. - output = self._query_gatherer(query.flatten(1), safe=False) - if output is not None: - return output.view(query.size(0), -1, query.size(2)) - return super().gather_query(query) + def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype): + from .symm_mem_allgather import SymmetricMemoryAllGather + + workspace = SymmetricMemoryAllGather( + self.device_group, dist.get_rank(self.device_group), gathered_width, + device=device, dtype=dtype, capacity_bytes=64 * 1024 * 1024) + workspace.prepare() + return workspace + def all_gather(self, input: torch.Tensor, *, workspace=None, copy_output: bool = True) -> torch.Tensor: + if workspace is not None: + output = workspace.all_gather(input, copy_output=copy_output) + if output is not None: + return output + return super().all_gather(input) -def should_try_symm_mem(dist_config) -> bool: - """Whether this configuration is a symmetric-memory candidate.""" - return (_envs.enable_symm_mem_allreduce and dist_config.dp == 1 - and dist_config.ep == 1 and dist_config.attn_tp > 1 - and not dist_config.enable_microbatch) + def close(self): + if self._all_reduce_provider is not None: + self._all_reduce_provider.close() + self._all_reduce_provider = None def build_cuda_communicator(cpu_group: dist.ProcessGroup, device_group: dist.ProcessGroup, - dist_config, *, group_roles: tuple[str, ...] = ('tp', )): - """Build the optional CUDA communicator for a TP or DCP group.""" + dist_config, *, group_name: Literal['tp', 'dcp'] = 'tp'): + """Agree on configuration and topology before optional collective setup.""" + backend = dist_config.communication_backend compatible = dist_config.dp == 1 and dist_config.ep == 1 and not dist_config.enable_microbatch if not compatible: return None - all_reduce_backend = None - if 'tp' in group_roles: - if _envs.enable_flashinfer_allreduce and _envs.enable_symm_mem_allreduce: - raise ValueError('FlashInfer and symmetric-memory all-reduce cannot be enabled together.') - if _envs.enable_flashinfer_allreduce: - all_reduce_backend = 'flashinfer' - elif should_try_symm_mem(dist_config): - all_reduce_backend = 'symm_mem' - query_gather = 'dcp' in group_roles and dist_config.dcp > 1 - if all_reduce_backend is None and not query_gather: + + device = torch.cuda.current_device() + local = (backend, group_name, socket.gethostname(), device, + torch.cuda.get_device_name(device), torch.cuda.get_device_capability(device)) + world_size = dist.get_world_size(device_group) + peers = [None] * world_size + dist.all_gather_object(peers, local, group=cpu_group) + if any(peer[:2] != local[:2] for peer in peers): + raise ValueError('communication_backend and group_name must agree across ranks') + if backend == 'nccl': + return None + + same_node = all(peer[2] == local[2] for peer in peers) + same_device = all(peer[4:] == local[4:] for peer in peers) + peer_access = (same_node and same_device and len({peer[3] for peer in peers}) == world_size + and all(peer[3] == device or torch.cuda.can_device_access_peer(device, peer[3]) for peer in peers)) + flags = [None] * world_size + dist.all_gather_object(flags, peer_access, group=cpu_group) + if not all(flags): + if dist.get_rank(device_group) == 0: + logger.info('Communication policy: process-group fallback (peer topology unavailable)') return None - # DCP ranks prepare together even when disabled, so admission can agree - # on NCCL fallback if the opt-in flag differs between ranks. + return CudaCommunicator(cpu_group=cpu_group, device_group=device_group, - all_reduce_backend=all_reduce_backend, - symm_mem_query_gather=query_gather and _envs.enable_symm_mem_dcp) + group_name=group_name) diff --git a/lmdeploy/pytorch/backends/cuda/comm/flashinfer_allreduce.py b/lmdeploy/pytorch/backends/cuda/comm/flashinfer_allreduce.py deleted file mode 100644 index 9cd5a431b5..0000000000 --- a/lmdeploy/pytorch/backends/cuda/comm/flashinfer_allreduce.py +++ /dev/null @@ -1,209 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. -# Modified from: https://github.com/vllm-project/vllm -import torch -from torch import distributed as dist - -from lmdeploy.utils import get_logger - -logger = get_logger('lmdeploy') - -_MiB = 1024 * 1024 -_MAX_TOKEN_NUM = 2048 - -# Empirical FlashInfer fusion limits used by vLLM: -# https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/allreduce_rms_fusion.py -# The separate token cap bounds each process-group workspace. -_FUSION_MAX_SIZE = { - 90: { - 2: 64 * _MiB, - 4: 2 * _MiB, - 8: _MiB // 2, - }, - 100: { - 2: 64 * _MiB, - 4: 32 * _MiB, - 8: _MiB, - }, - 103: { - 2: 64 * _MiB, - 4: 64 * _MiB, - 8: 2 * _MiB, - }, -} -_ONE_SHOT_MAX_SIZE = { - 90: { - 2: 32 * _MiB, - 4: 2 * _MiB, - 8: _MiB // 2, - }, - 100: { - 2: 32 * _MiB, - 4: 4 * _MiB, - 8: _MiB, - }, - 103: { - 2: 32 * _MiB, - 4: 4 * _MiB, - 8: 2 * _MiB, - }, -} - - -class FlashInferAllReduce: - """FlashInfer all-reduce operations for one distributed group.""" - - def __init__(self, group: dist.ProcessGroup): - """Initialize lazy FlashInfer state for ``group``.""" - self.group = group - self._comm = None - self._workspace = None - self._hidden_dim = None - self._dtype = None - self._disabled = False - self._world_size = dist.get_world_size(group) - major, minor = torch.cuda.get_device_capability() - self._device_capability = major * 10 + minor - self._max_size = _FUSION_MAX_SIZE.get(self._device_capability, {}).get(self._world_size, 0) - self._one_shot_max_size = _ONE_SHOT_MAX_SIZE.get(self._device_capability, {}).get(self._world_size, 0) - - def _disable(self, reason): - """Disable this backend after a process-group-wide setup failure.""" - self._disabled = True - logger.warning( - f'Disabling FlashInfer all-reduce for this process group: {reason}') - - def _initialize(self) -> bool: - """Load and validate the optional FlashInfer communication API.""" - if self._disabled: - return False - if self._comm is not None: - return True - message = 'FlashInfer all-reduce fusion requires flashinfer-python with the unified comm API.' - try: - import flashinfer.comm as comm - from flashinfer.comm.cuda_ipc import cudart - if not (hasattr(comm, 'allreduce_fusion') and hasattr( - comm, 'create_allreduce_fusion_workspace')): - raise ImportError(message) - - # Resolve FlashInfer's CUDA runtime before TileLang loads its - # libcudart shim, which the lazy loader could otherwise select. - cudart.cudaSetDevice(torch.cuda.current_device()) - except Exception as e: - self._disable(e) - return False - self._comm = comm - return True - - def is_available(self) -> bool: - """Whether the fused collective implementation is available.""" - return self._max_size > 0 and self._initialize() - - def supports(self, dtype: torch.dtype) -> bool: - """Whether this group can handle the given dtype.""" - return dtype in (torch.float16, torch.bfloat16) and self.is_available() - - def _supports_input(self, input: torch.Tensor) -> bool: - """Whether a flattened input satisfies the FlashInfer launch limits.""" - return (input.nbytes <= self._max_size - and input.size(0) <= _MAX_TOKEN_NUM - and input.is_contiguous() - and self.supports(input.dtype)) - - def _get_workspace(self, input: torch.Tensor): - """Return the workspace for the input shape and dtype.""" - hidden_dim = input.size(-1) - if self._workspace is not None: - assert self._hidden_dim == hidden_dim and self._dtype == input.dtype - return self._workspace - - rank = dist.get_rank(self.group) - max_token_num = min(_MAX_TOKEN_NUM, self._max_size // input[0].nbytes) - try: - workspace = self._comm.create_allreduce_fusion_workspace( - backend='trtllm', - world_size=self._world_size, - rank=rank, - max_token_num=max_token_num, - hidden_dim=hidden_dim, - dtype=input.dtype, - group=self.group, - ) - if workspace is None: - raise RuntimeError('workspace creation returned None') - except Exception as e: - self._disable(f'workspace initialization failed: {e}') - return None - - self._workspace = workspace - self._hidden_dim = hidden_dim - self._dtype = input.dtype - logger.info( - f'FlashInfer all-reduce workspace initialized: rank={rank}, world_size={self._world_size}, ' - f'max_token_num={max_token_num}, hidden_dim={hidden_dim}') - return self._workspace - - def all_reduce_(self, input: torch.Tensor) -> bool: - """All-reduce ``input`` in place, returning whether it was handled.""" - if input.dim() < 2 or not input.is_contiguous(): - return False - input_2d = input.flatten(0, -2) - if not self._supports_input(input_2d): - return False - - workspace = self._get_workspace(input_2d) - if workspace is None: - return False - output = self._comm.allreduce_fusion( - input=input_2d, - workspace=workspace, - pattern=self._comm.AllReduceFusionPattern.kAllReduce, - launch_with_pdl=True, - trigger_completion_at_end=True, - use_oneshot=input_2d.nbytes <= self._one_shot_max_size, - ) - input_2d.copy_(output) - return True - - def fused_all_reduce_residual_rms_norm(self, - input: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float): - """Fuse all-reduce, residual addition and RMSNorm when supported.""" - if weight.dtype != input.dtype: - return None - input_2d = input.flatten(0, -2) - residual_2d = residual.flatten(0, -2) - if (not self._supports_input(input_2d) - or not residual_2d.is_contiguous() - or not weight.is_contiguous()): - return None - - norm_out = torch.empty_like(input_2d) - residual_out = torch.empty_like(residual_2d) - workspace = self._get_workspace(input_2d) - if workspace is None: - return None - self._comm.allreduce_fusion( - input=input_2d, - workspace=workspace, - pattern=self._comm.AllReduceFusionPattern.kARResidualRMSNorm, - launch_with_pdl=True, - # Keep completion conservative; early signaling requires separate - # validation for the one-shot and two-shot paths. - trigger_completion_at_end=True, - use_oneshot=input_2d.nbytes <= self._one_shot_max_size, - residual_in=residual_2d, - residual_out=residual_out, - norm_out=norm_out, - rms_gamma=weight, - rms_eps=eps, - ) - return norm_out.view_as(input), residual_out.view_as(residual) - - def close(self): - """Release the group-bound FlashInfer workspace.""" - if self._workspace is not None: - self._workspace.destroy() - self._workspace = None diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py index 04d575dee9..cc5a8f3826 100644 --- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py +++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py @@ -1,24 +1,25 @@ # Copyright (c) OpenMMLab. All rights reserved. -"""Setup-owned symmetric-memory hidden-axis all-gather. +"""Symmetric-memory workspace for last-dimension all-gather. The optional Triton kernels are imported only during collective preparation. Their SGLang source reference is in kernels/cuda/symm_mem_allgather.py. """ import importlib -import logging import torch import torch.distributed as dist -logger = logging.getLogger(__name__) +from lmdeploy.utils import get_logger +logger = get_logger('lmdeploy') -class MultimemAllGatherer: + +class SymmetricMemoryAllGather: """Own configuration, collective admission and one symmetric arena. Construct on every group rank after process-group setup. Only preparation performs allocation and consensus; forward - never initializes or rebuilds. Call reset_for_weight after coordinated, quiescent device/dtype changes. Runtime rows - must match across ranks, as required by NCCL all-gather too. + never initializes or rebuilds. Re-prepare after coordinated, quiescent device/dtype changes. Runtime rows must match + across ranks, as required by NCCL all-gather too. """ # Keep the conservative single-token NCCL policy internal to the provider. @@ -26,29 +27,28 @@ class MultimemAllGatherer: _WORLD_SIZES = {2, 4, 8} def __init__(self, group: dist.ProcessGroup, rank: int, gathered_width: int, - device: torch.device, dtype: torch.dtype, *, - enabled: bool, capacity_bytes: int): - self._group = group + *, device: torch.device, dtype: torch.dtype, capacity_bytes: int): + self.device = torch.device(device) + self.dtype = dtype + self._prepared = False + self.group = group self._rank = rank self._gathered_width = gathered_width - self._enabled = enabled self._capacity_bytes = capacity_bytes self._state = None self._kernels = None self._graph_ready_shapes = set() - self._weight_contract = None - self.prepare(device, dtype) def _agree(self, ready: bool, device: torch.device) -> bool: flag = torch.tensor(int(ready), dtype=torch.int32, device=device) - dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=self._group) + dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=self.group) return bool(flag.item()) def _same_config(self, values: tuple[int, ...], device: torch.device) -> bool: lower = torch.tensor(values, dtype=torch.int64, device=device) upper = lower.clone() - dist.all_reduce(lower, op=dist.ReduceOp.MIN, group=self._group) - dist.all_reduce(upper, op=dist.ReduceOp.MAX, group=self._group) + dist.all_reduce(lower, op=dist.ReduceOp.MIN, group=self.group) + dist.all_reduce(upper, op=dist.ReduceOp.MAX, group=self.group) return bool(torch.equal(lower, upper)) def _disabled(self, reason: str) -> bool: @@ -56,31 +56,49 @@ def _disabled(self, reason: str) -> bool: logger.warning('symmetric-memory all-gather disabled: %s', reason) return False - def prepare(self, device: torch.device, dtype: torch.dtype) -> bool: + def _valid_handle(self, state, kernels) -> bool: + """Validate rank metadata and both pointer and signal-pad bounds.""" + try: + handle = state.symm_mem_hdl + multicast = int(handle.multicast_ptr or 0) + signals = handle.signal_pad_ptrs_dev + signal_addr = int(signals.data_ptr() if hasattr(signals, 'data_ptr') else signals or 0) + pad_size = getattr(handle, 'signal_pad_size', None) + if pad_size is None: + pad_size = kernels.symm_mem.get_signal_pad_size() + return (int(handle.rank) == self._rank + and int(handle.world_size) == state.world_size + and multicast > 0 and multicast % 16 == 0 + and signal_addr > 0 and signal_addr % 8 == 0 + and (not hasattr(signals, 'numel') or signals.numel() >= state.world_size) + and int(pad_size) >= kernels._MAX_BLOCKS * state.world_size * 4) + except (AttributeError, TypeError, ValueError, RuntimeError, OverflowError): + return False + + def _prepare(self, device: torch.device, dtype: torch.dtype) -> bool: """Collectively rebuild and admit the arena, outside CUDA graphs.""" device = torch.device(device) if device.type == 'cuda' and device.index is None: device = torch.device('cuda', torch.cuda.current_device()) if device.type == 'cuda' and torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): raise RuntimeError('symmetric-memory preparation is not graph capturable') - self.release() - self._weight_contract = (device, dtype) + self.close() if device.type != 'cuda' or not torch.cuda.is_available(): return False - if not isinstance(self._group, dist.ProcessGroup): + if not isinstance(self.group, dist.ProcessGroup): return False # Even disabled ranks participate: mixed flags/dtypes must not leave # peers entering an optional rendezvous on their own. - if not self._agree(self._enabled and dtype == torch.bfloat16, device): + if not self._agree(dtype == torch.bfloat16, device): return False - world_size = dist.get_world_size(self._group) + world_size = dist.get_world_size(self.group) width = self._gathered_width try: capability_ok = torch.cuda.get_device_capability(device) >= (9, 0) except (RuntimeError, AssertionError): capability_ok = False valid = (world_size in self._WORLD_SIZES - and self._rank == dist.get_rank(self._group) + and self._rank == dist.get_rank(self.group) and width > 0 and width % (world_size * 8) == 0 and self._capacity_bytes >= width * torch.bfloat16.itemsize and capability_ok) @@ -101,7 +119,7 @@ def prepare(self, device: torch.device, dtype: torch.dtype) -> bool: buffer = None allocation_ok = False try: - buffer = kernels._allocate_symmetric_buffer(self._group, max_tokens, width, device) + buffer = kernels._allocate_symmetric_buffer(self.group, max_tokens, width, device) allocation_ok = (buffer.shape == (max_tokens, width) and buffer.dtype == torch.bfloat16 and buffer.device == device and buffer.is_contiguous() and buffer.storage_offset() == 0 @@ -113,7 +131,7 @@ def prepare(self, device: torch.device, dtype: torch.dtype) -> bool: # Once all ranks commit, rendezvous errors are fatal, not a rank-local # reason to fall back and strand peers inside a collective. - state = kernels.create_state(self._group, self._rank, max_tokens, width, device, comm_buff=buffer) + state = kernels.create_state(self.group, self._rank, max_tokens, width, device, comm_buff=buffer) if not self._agree(self._valid_handle(state, kernels), device): return self._disabled('invalid symmetric-memory handle or unsupported multicast') self._kernels = kernels @@ -123,32 +141,51 @@ def prepare(self, device: torch.device, dtype: torch.dtype) -> bool: world_size, width, max_tokens) return True - def _valid_handle(self, state, kernels) -> bool: - """Validate rank metadata and both pointer and signal-pad bounds.""" - try: - handle = state.symm_mem_hdl - multicast = int(handle.multicast_ptr or 0) - signals = handle.signal_pad_ptrs_dev - signal_addr = int(signals.data_ptr() if hasattr(signals, 'data_ptr') else signals or 0) - pad_size = getattr(handle, 'signal_pad_size', None) - if pad_size is None: - pad_size = kernels.symm_mem.get_signal_pad_size() - return (int(handle.rank) == self._rank - and int(handle.world_size) == state.world_size - and multicast > 0 and multicast % 16 == 0 - and signal_addr > 0 and signal_addr % 8 == 0 - and (not hasattr(signals, 'numel') or signals.numel() >= state.world_size) - and int(pad_size) >= kernels._MAX_BLOCKS * state.world_size * 4) - except (AttributeError, TypeError, ValueError, RuntimeError, OverflowError): - return False + def is_available(self) -> bool: + """Whether the symmetric-memory workspace is ready for gathering.""" + return self._state is not None - def reset_for_weight(self, weight: torch.Tensor) -> None: - """Notify the provider after a quiescent, group-wide model - transition.""" - if (weight.device, weight.dtype) != self._weight_contract: - self.prepare(weight.device, weight.dtype) + def prepare(self): + """Allocate once during setup, never during forward.""" + if not self._prepared: + self._prepare(self.device, self.dtype) + self._prepared = True - def release(self) -> None: + def reset(self, device: torch.device, dtype: torch.dtype): + """Rebuild after a coordinated, quiescent device/dtype transition.""" + device = torch.device(device) + if (device, dtype) == (self.device, self.dtype): + return + self.device, self.dtype = device, dtype + if self._prepared: + self._prepared = False + self.prepare() + + def all_gather(self, input: torch.Tensor, *, copy_output: bool = True) -> torch.Tensor | None: + """Return gathered shards, or None when the input is unsupported. + + With copy_output=False, the output borrows the workspace until its next use. + """ + state = self._state + if state is None: + return None + if not (input.dim() == 2 and input.dtype == torch.bfloat16 and input.is_contiguous() + and input.shape[-1] * state.world_size == state.hidden_dim): + # Shape/dtype/layout contracts must match across ranks, as for NCCL. + return None + if input.device != state.device or input.data_ptr() % 16 != 0: + raise RuntimeError('multimem all-gather device or alignment changed after group-wide admission') + rows = input.shape[0] + if not self._MIN_TOKENS <= rows <= state.max_token_num: + return None + if rows not in self._graph_ready_shapes and torch.cuda.is_current_stream_capturing(): + return None + output = self._kernels.all_gather_inner( + state, input, tp_hidden_dim=self._gathered_width, copy_output=copy_output, _validated=True) + self._graph_ready_shapes.add(rows) + return output + + def close(self) -> None: """Release only after all consumers/graphs have been retired.""" state = self._state if state is not None: @@ -160,25 +197,4 @@ def release(self) -> None: self._state = None self._kernels = None self._graph_ready_shapes.clear() - - def __call__(self, x: torch.Tensor, *, safe: bool = True) -> torch.Tensor | None: - """Gather hidden shards; unsafe output is borrowed until the next - call.""" - state = self._state - if state is None: - return None - if not (x.dim() == 2 and x.dtype == torch.bfloat16 and x.device == state.device - and x.is_contiguous() and x.data_ptr() % 16 == 0 - and x.shape[-1] * state.world_size == state.hidden_dim): - # After collective admission this is a programming error. A local - # NCCL fallback could disagree with peers launching multimem. - raise RuntimeError('multimem all-gather input contract changed after group-wide admission') - rows = x.shape[0] - if not self._MIN_TOKENS <= rows <= state.max_token_num: - return None - if rows not in self._graph_ready_shapes and torch.cuda.is_current_stream_capturing(): - return None - output = self._kernels.all_gather_inner( - state, x, tp_hidden_dim=self._gathered_width, safe=safe, _validated=True) - self._graph_ready_shapes.add(rows) - return output + self._prepared = False diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allreduce.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allreduce.py index 7f1af61c21..1a9e328e8c 100644 --- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allreduce.py +++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allreduce.py @@ -43,8 +43,7 @@ class SymmetricMemoryAllReduce: def __init__(self, group: dist.ProcessGroup): self.group = group self._buffer = None - self._enabled = False - self._max_size = 0 + self._available = False world_size = dist.get_world_size(group) major, minor = torch.cuda.get_device_capability() @@ -59,38 +58,53 @@ def __init__(self, group: dist.ProcessGroup): op_name = 'multimem_all_reduce_' if self._use_multimem else 'two_shot_all_reduce_' if not (hasattr(symm_mem, 'empty') and hasattr(symm_mem, 'rendezvous') and hasattr(torch.ops.symm_mem, op_name)): - logger.warning('PyTorch symmetric-memory all-reduce is unavailable in this PyTorch build.') - return + raise RuntimeError('PyTorch symmetric-memory all-reduce is unavailable in this PyTorch build') + self._symm_mem = symm_mem + self._available = True + except (ImportError, RuntimeError) as e: + logger.warning(f'PyTorch symmetric-memory all-reduce is unavailable: {e}') + + def is_available(self) -> bool: + """Whether the optimized collective implementation is available.""" + return self._available + + def prepare(self): + """Allocate the group buffer during communicator construction.""" + if self._buffer is not None or not self._available: + return + + group = self.group + world_size = dist.get_world_size(group) + symm_mem = self._symm_mem + try: self._buffer = symm_mem.empty( self._max_size // torch.bfloat16.itemsize, dtype=torch.bfloat16, device=torch.device('cuda', torch.cuda.current_device()), ) - handle = symm_mem.rendezvous(self._buffer, group.group_name) - except (ImportError, RuntimeError) as e: - logger.warning(f'PyTorch symmetric-memory all-reduce is unavailable: {e}') + except RuntimeError as e: + logger.warning(f'PyTorch symmetric-memory allocation failed: {e}') self._buffer = None + + ready = [None] * world_size + dist.all_gather_object(ready, self._buffer is not None, group=group) + if not all(ready): + self.close() return + # Once ranks commit, rendezvous failures are fatal rather than a local fallback. + handle = symm_mem.rendezvous(self._buffer, group.group_name) if getattr(handle, 'multicast_ptr', 0) == 0: logger.warning('PyTorch symmetric-memory all-reduce is unavailable: multicast is not supported.') - self._buffer = None + self.close() return - self._enabled = True + if dist.get_rank(group) == 0: logger.info(f'Using PyTorch symmetric-memory all-reduce for TP{world_size}.') - def supports(self, dtype: torch.dtype) -> bool: - """Whether this group can all-reduce the given dtype.""" - return self._enabled and dtype == torch.bfloat16 - - def is_available(self) -> bool: - """Whether the optimized collective implementation is available.""" - return self._enabled - def all_reduce_(self, input: torch.Tensor) -> bool: """All-reduce ``input`` in place, returning whether it was handled.""" - if (not self.supports(input.dtype) or not input.is_contiguous() + if (self._buffer is None or input.dtype != torch.bfloat16 or not input.is_contiguous() or input.nbytes > self._max_size or input.nbytes % 4 != 0): return False @@ -106,4 +120,4 @@ def all_reduce_(self, input: torch.Tensor) -> bool: def close(self): """Release the symmetric buffer.""" self._buffer = None - self._enabled = False + self._available = False diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py index bb829b5566..8a9a691220 100644 --- a/lmdeploy/pytorch/backends/cuda/op_backend.py +++ b/lmdeploy/pytorch/backends/cuda/op_backend.py @@ -228,14 +228,14 @@ def get_v4_attention_metadata_cls(): return CudaV4AttentionMetadata @classmethod - def build_communicator(cls, cpu_group, device_group, dist_config, *, group_roles=('tp', )): + def build_communicator(cls, cpu_group, device_group, dist_config, *, group_name='tp'): """Build a CUDA communicator.""" from .comm.communicator import build_cuda_communicator communicator = build_cuda_communicator( cpu_group=cpu_group, device_group=device_group, dist_config=dist_config, - group_roles=group_roles, + group_name=group_name, ) if communicator is not None: return communicator @@ -243,7 +243,7 @@ def build_communicator(cls, cpu_group, device_group, dist_config, *, group_roles cpu_group=cpu_group, device_group=device_group, dist_config=dist_config, - group_roles=group_roles, + group_name=group_name, ) @staticmethod diff --git a/lmdeploy/pytorch/backends/default/embedding.py b/lmdeploy/pytorch/backends/default/embedding.py index c5596d71f9..49a94b7e4a 100644 --- a/lmdeploy/pytorch/backends/default/embedding.py +++ b/lmdeploy/pytorch/backends/default/embedding.py @@ -1,8 +1,9 @@ # Copyright (c) OpenMMLab. All rights reserved. import torch -import torch.distributed as dist import torch.nn.functional as F +import lmdeploy.pytorch.distributed as dist + from ..embedding import EmbeddingImpl diff --git a/lmdeploy/pytorch/backends/default/linear.py b/lmdeploy/pytorch/backends/default/linear.py index e7eba90485..589694605e 100644 --- a/lmdeploy/pytorch/backends/default/linear.py +++ b/lmdeploy/pytorch/backends/default/linear.py @@ -1,9 +1,10 @@ # Copyright (c) OpenMMLab. All rights reserved. import torch -import torch.distributed as dist import torch.nn.functional as F +import lmdeploy.pytorch.distributed as dist + from ..linear import LinearImpl diff --git a/lmdeploy/pytorch/backends/dlinfer/op_backend.py b/lmdeploy/pytorch/backends/dlinfer/op_backend.py index a7b5a82a94..a825df8a74 100644 --- a/lmdeploy/pytorch/backends/dlinfer/op_backend.py +++ b/lmdeploy/pytorch/backends/dlinfer/op_backend.py @@ -126,20 +126,6 @@ def get_cache_backend(cls): from .cache import DlinferCacheBackend return DlinferCacheBackend - @classmethod - def build_communicator(cls, cpu_group, device_group, dist_config, *, group_roles=('tp', )): - """Build a DLInfer communicator.""" - from lmdeploy.pytorch import envs - cuda_communicator_enabled = (envs.enable_flashinfer_allreduce or envs.enable_symm_mem_allreduce - or envs.enable_symm_mem_dcp) - assert not cuda_communicator_enabled, 'CUDA communicators are not supported by DLInfer.' - return super().build_communicator( - cpu_group=cpu_group, - device_group=device_group, - dist_config=dist_config, - group_roles=group_roles, - ) - @staticmethod def get_attention_metadata_cls(): from .attention import DlinferAttentionMetadata diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py index 2e38c39c58..5f5133fbc7 100644 --- a/lmdeploy/pytorch/config.py +++ b/lmdeploy/pytorch/config.py @@ -3,7 +3,7 @@ import re from collections.abc import Callable, Mapping from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal import torch @@ -224,6 +224,9 @@ class DistConfig: mlp_tp_mode: TPMode = TPMode.DEFAULT moe_tp_mode: TPMode = TPMode.DEFAULT + # communication + communication_backend: Literal['nccl', 'auto'] = 'nccl' + def __post_init__(self): """Post init.""" assert self.dp_rank < self.dp @@ -297,6 +300,7 @@ def from_engine_config(cls, engine_config: PytorchEngineConfig): dp_rank=engine_config.dp_rank, enable_microbatch=engine_config.enable_microbatch, enable_eplb=engine_config.enable_eplb, + communication_backend=engine_config.communication_backend, tp=engine_config.tp, dcp=engine_config.dcp, attn_tp=engine_config.attn_tp_size, diff --git a/lmdeploy/pytorch/distributed.py b/lmdeploy/pytorch/distributed.py index c7a88e05ec..f76ca7145a 100644 --- a/lmdeploy/pytorch/distributed.py +++ b/lmdeploy/pytorch/distributed.py @@ -26,38 +26,16 @@ class DistGroup: gpu_groups: list[dist.ProcessGroup] = None gpu_gather_group: dist.ProcessGroup = None communicator: 'DeviceCommunicator' = None - - def supports_optimized_all_reduce(self) -> bool: - """Whether this group has an optimized all-reduce implementation.""" - return self.communicator is not None and self.communicator.supports_optimized_all_reduce() - - def supports_fused_all_reduce_residual_rms_norm(self) -> bool: - """Whether this group can fuse all-reduce, residual and RMSNorm.""" - return (self.communicator is not None - and self.communicator.supports_fused_all_reduce_residual_rms_norm()) - - def try_fused_all_reduce_residual_rms_norm(self, - input: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - eps: float): - """Run fused all-reduce, residual and RMSNorm, or return ``None``.""" - if self.communicator is None: - return None - return self.communicator.try_fused_all_reduce_residual_rms_norm( - input=input, residual=residual, weight=weight, eps=eps) - - def all_reduce_(self, tensor: torch.Tensor): - """All-reduce a tensor in place on this group.""" - if self.communicator is not None: - return self.communicator.all_reduce_(tensor) - return dist.all_reduce(tensor, group=self.gpu_group) + query_gather_workspace: object = None def close(self): """Close groups.""" if self.communicator is not None: self.communicator.close() self.communicator = None + if self.query_gather_workspace is not None: + self.query_gather_workspace.close() + self.query_gather_workspace = None if not dist.is_initialized(): return if self.cpu_groups is not None: @@ -241,19 +219,16 @@ def _build_communicators(context: 'DistContext'): """Attach one communicator to each rank-local, unique TP/DCP group.""" build_communicator = context.communicator_builder tp_groups = (context.attn_tp_group, context.mlp_tp_group, context.moe_tp_group) - groups = (*tp_groups, context.dcp_group) - for group in {id(group): group for group in groups}.values(): + groups = [('tp', group) for group in {id(group): group for group in tp_groups}.values()] + groups.append(('dcp', context.dcp_group)) + for group_name, group in groups: if group.gpu_group is None: continue - # Aliased groups share one communicator with both operation roles. - roles = ('tp', ) if any(group is tp_group for tp_group in tp_groups) else () - if group is context.dcp_group: - roles += ('dcp', ) group.communicator = build_communicator( cpu_group=group.cpu_group, device_group=group.gpu_group, dist_config=context.dist_config, - group_roles=roles, + group_name=group_name, ) @@ -491,6 +466,17 @@ def all_reduce(tensor, op=ReduceOp.SUM, group='tp', async_op=False): """All reduce.""" if isinstance(group, str): group = get_group(group, 'gpu') + + # Optimize synchronous CUDA SUM on explicit TP groups; keep other calls native. + if tensor.is_cuda and op == ReduceOp.SUM and not async_op and group is not None: + context = get_dist_manager().current_context() + # Resolve the raw ProcessGroup to its communicator. + for tp_group in (context.attn_tp_group, context.mlp_tp_group, context.moe_tp_group): + if tp_group is not None and tp_group.gpu_group is group: + if tp_group.communicator is not None: + return tp_group.communicator.all_reduce_(tensor) + break + return dist.all_reduce(tensor, op, group, async_op) diff --git a/lmdeploy/pytorch/engine/executor/ray_executor.py b/lmdeploy/pytorch/engine/executor/ray_executor.py index 6939388e30..0462eeab02 100644 --- a/lmdeploy/pytorch/engine/executor/ray_executor.py +++ b/lmdeploy/pytorch/engine/executor/ray_executor.py @@ -104,10 +104,9 @@ def _needs_symm_mem_device_setup(dist_config: DistConfig) -> bool: memory rendezvous identifies allocations by the process-local device ordinal, so TP actors on one host must instead inherit the full visibility and select their assigned GPU. """ - from lmdeploy.pytorch.backends.cuda.comm.communicator import should_try_symm_mem - return (should_try_symm_mem(dist_config) - or (_envs.enable_symm_mem_lmhead and dist_config.attn_tp > 1) - or (_envs.enable_symm_mem_dcp and dist_config.dcp > 1)) + if dist_config.dp != 1 or dist_config.ep != 1 or dist_config.enable_microbatch: + return False + return dist_config.communication_backend == 'auto' and (dist_config.attn_tp > 1 or dist_config.dcp > 1) class RemoteLogger: diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py index 3319049eb3..82059829c3 100644 --- a/lmdeploy/pytorch/envs.py +++ b/lmdeploy/pytorch/envs.py @@ -246,13 +246,6 @@ def _patched_get_env( fake_capture = env_to_bool('LMDEPLOY_FAKE_CUDA_GRAPH_CAPTURE', False) enable_decode_torch_compile = env_to_bool('LMDEPLOY_ENABLE_DECODE_TORCH_COMPILE', False) - # cuda communicator - enable_flashinfer_allreduce = env_to_bool('LMDEPLOY_ENABLE_FLASHINFER_ALLREDUCE', False) - enable_symm_mem_allreduce = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_ALLREDUCE', False) - enable_symm_mem_lmhead = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_LMHEAD', False) - enable_symm_mem_dcp = env_to_bool('LMDEPLOY_ENABLE_SYMM_MEM_DCP', False) - symm_mem_lmhead_max_mb = max(1, env_to_int('LMDEPLOY_SYMM_MEM_LMHEAD_MAX_MB', 64)) - # opt-ttft opt_ttft_policy = env_to_choice('LMDEPLOY_PT_TTFT_POLICY', 'size', {'fifo', 'size'}) opt_ttft_short_turns = max(1, env_to_int('LMDEPLOY_PT_TTFT_SHORT_TURNS', 3)) diff --git a/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py b/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py index 486f21396a..06bd0d2e7e 100644 --- a/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py +++ b/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py @@ -451,13 +451,13 @@ def all_gather_inner( hidden_states: torch.Tensor, tp_hidden_dim: int, skip_entry_sync: bool = False, - safe: bool = True, + copy_output: bool = True, *, _validated: bool = False, ) -> torch.Tensor: """Gather ``[T, H/TP]`` shards into ``[T, H]`` along the hidden dim. - ``tp_hidden_dim`` is the gathered width ``H``. Returns a clone when ``safe``, + ``tp_hidden_dim`` is the gathered width ``H``. Returns a clone when ``copy_output``, else a view into the symmetric buffer (valid until the next collective). ``_validated`` is reserved for the admitted provider, whose admission check already enforces the immutable dtype/layout/width @@ -528,7 +528,7 @@ def all_gather_inner( # publish the whole grid before the output view is returned. barrier_inner(state, slot=1, release=True) output = state.comm_buff[:total_tokens, :tp_hidden_dim] - return output.clone() if safe else output + return output.clone() if copy_output else output def barrier_inner(state: MultimemAllGatherState, *, slot: int, diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index fd93567af0..810ec1ed7a 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -492,22 +492,12 @@ def __init__(self, nn.Module.__init__(self) self.layer_idx = layer_idx quantization_config = None - # Row-wise TP outputs normally reduce inside their projections. An - # optimized communicator lets the following RMSNorm consume that - # reduction instead. Attention is consumed in this layer. - defer_attn_all_reduce = RMSNorm.can_handle_all_reduce('attn') - # MLP is consumed by the next target layer, so terminal and MTP blocks - # must still reduce their outputs. - defer_mlp_all_reduce = (layer_idx < config.num_hidden_layers - 1 - and RMSNorm.can_handle_all_reduce('mlp')) - # build attention layer self.self_attn = self.attention_cls( config, layer_idx, dtype=dtype, device=device, - all_reduce=not defer_attn_all_reduce, prefix=add_prefix('self_attn', prefix)) # mlp @@ -517,14 +507,12 @@ def __init__(self, layer_idx, dtype=dtype, device=device, - all_reduce=not defer_mlp_all_reduce, prefix=mlp_prefix) if (config.n_routed_experts is not None and layer_idx >= config.first_k_dense_replace and layer_idx % config.moe_layer_freq == 0) else DeepseekV2MLP( config, dtype=dtype, device=device, - all_reduce=not defer_mlp_all_reduce, prefix=mlp_prefix)) # build input layer norm @@ -532,15 +520,13 @@ def __init__(self, config.rms_norm_eps, quant_config=quantization_config, dtype=torch.float32, - device=device, - all_reduce_group='mlp') + device=device) # build attention layer norm self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps, dtype=torch.float32, - device=device, - all_reduce_group='attn') + device=device) class DeepseekV32Model(DeepseekV2Model): decoder_layer_cls = DeepseekV32DecoderLayer diff --git a/lmdeploy/pytorch/nn/embedding.py b/lmdeploy/pytorch/nn/embedding.py index fb04f9b179..812a5ed1ce 100644 --- a/lmdeploy/pytorch/nn/embedding.py +++ b/lmdeploy/pytorch/nn/embedding.py @@ -1,11 +1,8 @@ # Copyright (c) OpenMMLab. All rights reserved. import torch -import torch.distributed as dist from torch import nn -from lmdeploy.pytorch import envs as _envs from lmdeploy.pytorch.backends import get_backend -from lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather import MultimemAllGatherer from lmdeploy.pytorch.backends.embedding import EmbeddingBuildSpec from lmdeploy.pytorch.backends.linear import LinearBuildSpec from lmdeploy.pytorch.distributed import get_dist_group, get_dist_manager, get_tp_world_rank @@ -155,23 +152,28 @@ def __init__( enable_deterministic=get_build_model_context().enable_deterministic, ) - self._symm_mem_gatherer = ( - MultimemAllGatherer(self.tp_group, self.tp_rank, - self.tp * self.vocab_size_padded, - self.weight.device, self.weight.dtype, - enabled=_envs.enable_symm_mem_lmhead, - capacity_bytes=_envs.symm_mem_lmhead_max_mb * 1024 * 1024) - if self.all_reduce else None) + communicator = get_dist_group(layer_type=layer_type).communicator + self._logits_gather_workspace = None + if self.all_reduce: + if communicator is None: + from lmdeploy.pytorch.backends.communicator import DeviceCommunicator + communicator = DeviceCommunicator(self.tp_group) + self._logits_gather_workspace = communicator.create_all_gather_workspace( + self.tp * self.vocab_size_padded, device=self.weight.device, dtype=self.weight.dtype) + + self._communicator = communicator def tie_weights(self, embedding: ParallelEmbedding): """Tie the local LM-head shard to a parallel embedding shard.""" self.weight = embedding.weight + if self._logits_gather_workspace is not None: + self._logits_gather_workspace.reset(self.weight.device, self.weight.dtype) def _apply(self, fn, recurse=True): - """Notify the provider after coordinated model device/dtype moves.""" + """Update the workspace after coordinated model device/dtype moves.""" result = super()._apply(fn, recurse=recurse) - if self._symm_mem_gatherer is not None: - self._symm_mem_gatherer.reset_for_weight(self.weight) + if self._logits_gather_workspace is not None: + self._logits_gather_workspace.reset(self.weight.device, self.weight.dtype) return result def get_local_logits(self, hidden_states: torch.Tensor): @@ -185,22 +187,10 @@ def all_gather_logits(self, local_logits: torch.Tensor) -> torch.Tensor: if not self.all_reduce: return local_logits[..., :self.vocab_size] - if self._symm_mem_gatherer is not None: - gathered = self._symm_mem_gatherer(local_logits.reshape(-1, local_logits.shape[-1])) - if gathered is not None: - output_shape = local_logits.shape[:-1] + (self.tp * local_logits.shape[-1], ) - return gathered.reshape(output_shape)[..., :self.vocab_size] - - input_size = local_logits.size() - output_size = (input_size[0] * self.tp, ) + input_size[1:] - logits = local_logits.new_empty(output_size) - dist.all_gather_into_tensor(logits, local_logits, group=self.tp_group) - # The collective concatenates dim 0. Move its rank dimension beside - # the vocabulary shard before reconstructing the full last dimension. - logits = logits.reshape((self.tp, ) + input_size) - logits = logits.movedim(0, local_logits.dim() - 1) - logits = logits.reshape(input_size[:-1] + (self.tp * input_size[-1], )) - return logits[..., :self.vocab_size] + gathered = self._communicator.all_gather( + local_logits.reshape(-1, local_logits.shape[-1]), workspace=self._logits_gather_workspace) + output_shape = local_logits.shape[:-1] + (self.tp * local_logits.shape[-1], ) + return gathered.reshape(output_shape)[..., :self.vocab_size] def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """Compute TP-local logits and all-gather them on every rank.""" diff --git a/lmdeploy/pytorch/nn/linear/base.py b/lmdeploy/pytorch/nn/linear/base.py index 03ec93bb0f..6cebbe8b82 100644 --- a/lmdeploy/pytorch/nn/linear/base.py +++ b/lmdeploy/pytorch/nn/linear/base.py @@ -7,6 +7,7 @@ from lmdeploy.pytorch.config import TPMode from lmdeploy.pytorch.distributed import ( + all_reduce, gather_by_tp_sizes, get_dist_group, get_dist_manager, @@ -207,7 +208,7 @@ def _forward_lora(self, x, tp_sizes: list[int] = None): if self.tp_mode == TPMode.DP_TP: out = reduce_scatter_by_tp_sizes(out, self.tp_rank, tp_sizes, group=self.tp_group) else: - dist.all_reduce(out, group=self.tp_group) + all_reduce(out, group=self.tp_group) return out def _forward_dp_tp(self, x): diff --git a/lmdeploy/pytorch/nn/norm.py b/lmdeploy/pytorch/nn/norm.py index 362c2ef6e5..3a17326d2a 100644 --- a/lmdeploy/pytorch/nn/norm.py +++ b/lmdeploy/pytorch/nn/norm.py @@ -3,7 +3,7 @@ import torch from torch import nn -from lmdeploy.pytorch.distributed import get_dist_group, get_tp_world_rank +from lmdeploy.pytorch.distributed import get_tp_world_rank from lmdeploy.pytorch.models.patch import get_build_model_context from ..backends import get_backend @@ -40,7 +40,6 @@ def __init__( tp: bool = False, align: int = 1, prefix: str = '', - all_reduce_group: str | None = None, ): super().__init__() backend = get_backend() @@ -69,18 +68,6 @@ def __init__( if tp: self.weight.weight_loader = self.weight_loader self.align = align - self.eps = eps - self._all_reduce_group = None - self._fuse_all_reduce = False - if all_reduce_group is not None and self.can_handle_all_reduce(all_reduce_group): - group = get_dist_group(all_reduce_group) - self._all_reduce_group = group - self._fuse_all_reduce = group.supports_fused_all_reduce_residual_rms_norm() - - @staticmethod - def can_handle_all_reduce(layer_type: str) -> bool: - """Whether this norm can consume a pending all-reduce.""" - return get_dist_group(layer_type).supports_optimized_all_reduce() def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): """Weight loader.""" @@ -100,17 +87,6 @@ def create_weight(hidden_size: int, dtype: torch.dtype | None = None, device: to def forward(self, x: torch.Tensor, residual: torch.Tensor = None): """forward.""" - if residual is not None and self._all_reduce_group is not None: - if self._fuse_all_reduce: - output = self._all_reduce_group.try_fused_all_reduce_residual_rms_norm( - input=x, - residual=residual, - weight=self.weight, - eps=self.eps, - ) - if output is not None: - return output - self._all_reduce_group.all_reduce_(x) return self.impl.forward(x, self.weight, residual) diff --git a/lmdeploy/pytorch/spec_decode/base.py b/lmdeploy/pytorch/spec_decode/base.py index 58f1d2f4e9..4f5a053c3d 100644 --- a/lmdeploy/pytorch/spec_decode/base.py +++ b/lmdeploy/pytorch/spec_decode/base.py @@ -18,6 +18,7 @@ def _build_draft_dist_ctx(dist_ctx: DistContext, specdecode_config: SpecDecodeCo return None draft_dist_config = specdecode_config.dist_config + draft_dist_config.communication_backend = dist_ctx.dist_config.communication_backend if specdecode_config.method in ('qwen3_5_mtp', 'hy3_mtp') or draft_dist_config == dist_ctx.dist_config: return dist_ctx diff --git a/tests/pytorch/engine/test_executor_base.py b/tests/pytorch/engine/test_executor_base.py index 0f426327e8..777358ba43 100644 --- a/tests/pytorch/engine/test_executor_base.py +++ b/tests/pytorch/engine/test_executor_base.py @@ -193,7 +193,7 @@ def test_get_min_num_gpu_blocks_rejects_worker_count_mismatch(): def test_runtime_size_reserves_dcp_peak_phase(monkeypatch, draft_tokens, topk, score_mb, dtype): monkeypatch.setattr(executor_base._envs, 'dsa_indexer_max_logits_mb', score_mb) executor = object.__new__(ExecutorBase) - executor.model_config = SimpleNamespace(mla_index_topk=topk, head_dim=576, v_head_dim=0, + executor.model_config = SimpleNamespace(use_flash_mla=True, mla_index_topk=topk, head_dim=576, v_head_dim=0, num_attention_heads=64, dtype=dtype) executor.specdecode_config = SimpleNamespace(num_speculative_tokens=draft_tokens) if draft_tokens else None executor.dist_config = SimpleNamespace(attn_tp=8) diff --git a/tests/pytorch/spec_decode/test_spec_agent.py b/tests/pytorch/spec_decode/test_spec_agent.py index 09580d3eaf..381bc6544b 100644 --- a/tests/pytorch/spec_decode/test_spec_agent.py +++ b/tests/pytorch/spec_decode/test_spec_agent.py @@ -478,13 +478,14 @@ def test_spec_model_agent_method_when_enabled(): assert agent.method == specdecode_config.method -def test_qwen35_mtp_reuses_main_dist_context(monkeypatch): +@pytest.mark.parametrize('backend', ['auto', 'nccl']) +def test_qwen35_mtp_reuses_main_dist_context(monkeypatch, backend): """Qwen3.5 MTP mirrors the target topology, so it should share groups.""" from lmdeploy.pytorch.config import DistConfig, SpecDecodeConfig from lmdeploy.pytorch.distributed import DistContext from lmdeploy.pytorch.spec_decode import base as base_mod - dist_config = DistConfig(dp=2, ep=2) + dist_config = DistConfig(dp=2, ep=2, communication_backend=backend) dist_ctx = DistContext(rank=1, dp_rank=1, dist_config=dist_config, ep_gpu_group=object()) specdecode_config = SpecDecodeConfig(model='draft-model', method='qwen3_5_mtp', @@ -497,16 +498,18 @@ def fail_build(*args, **kwargs): monkeypatch.setattr(base_mod.DistContext, 'build', staticmethod(fail_build)) assert base_mod._build_draft_dist_ctx(dist_ctx, specdecode_config) is dist_ctx + assert specdecode_config.dist_config.communication_backend == backend -def test_non_qwen35_mtp_builds_draft_dist_context(monkeypatch): +@pytest.mark.parametrize('backend', ['auto', 'nccl']) +def test_non_qwen35_mtp_builds_draft_dist_context(monkeypatch, backend): """Other speculative methods keep their separate draft distribution path.""" from lmdeploy.pytorch.config import DistConfig, SpecDecodeConfig from lmdeploy.pytorch.distributed import DistContext from lmdeploy.pytorch.spec_decode import base as base_mod - main_dist_config = DistConfig(dp=2, ep=2) + main_dist_config = DistConfig(dp=2, ep=2, communication_backend=backend) draft_dist_config = DistConfig() dist_ctx = DistContext(rank=1, dp_rank=1, dist_config=main_dist_config) specdecode_config = SpecDecodeConfig(model='draft-model', @@ -524,6 +527,7 @@ def fake_build(*, rank, dist_config): assert base_mod._build_draft_dist_ctx(dist_ctx, specdecode_config) is draft_dist_ctx assert build_calls == [(dist_ctx.rank, draft_dist_config)] + assert draft_dist_config.communication_backend == backend def test_async_model_forward_dp1_non_last_chunk_skips_remaining_spec_forwards(): diff --git a/tests/pytorch/test_cuda_communicator.py b/tests/pytorch/test_cuda_communicator.py index 4b841fc7a7..aa26b0f22e 100644 --- a/tests/pytorch/test_cuda_communicator.py +++ b/tests/pytorch/test_cuda_communicator.py @@ -1,5 +1,4 @@ # Copyright (c) OpenMMLab. All rights reserved. -import sys from types import SimpleNamespace from unittest.mock import Mock @@ -8,8 +7,6 @@ from lmdeploy.pytorch.backends import communicator as base_communicator_module from lmdeploy.pytorch.backends.cuda.comm import communicator as communicator_module -from lmdeploy.pytorch.backends.dlinfer.op_backend import DlinferOpsBackend -from lmdeploy.pytorch.nn import norm as norm_module def _run_dcp_query_gather(rank, rendezvous, enabled): @@ -17,32 +14,41 @@ def _run_dcp_query_gather(rank, rendezvous, enabled): from torch import distributed as dist - from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_query + from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_query, init_dcp_query_gather from lmdeploy.pytorch.backends.cuda.op_backend import CudaOpsBackend from lmdeploy.pytorch.config import DistConfig from lmdeploy.pytorch.distributed import DistContext, get_dist_manager torch.cuda.set_device(rank) - communicator_module._envs.enable_symm_mem_dcp = enabled[rank] dist.init_process_group('nccl', init_method=rendezvous, rank=rank, world_size=2, timeout=timedelta(seconds=60)) - ctx = DistContext.build(rank, DistConfig(tp=2, dcp=2), + backend = 'auto' if enabled[rank] else 'nccl' + if enabled[0] != enabled[1]: + try: + with pytest.raises(ValueError, match='agree across ranks'): + DistContext.build(rank, DistConfig(tp=2, dcp=2, communication_backend=backend), + communicator_builder=CudaOpsBackend.build_communicator) + finally: + dist.destroy_process_group() + return + ctx = DistContext.build(rank, DistConfig(tp=2, dcp=2, communication_backend=backend), communicator_builder=CudaOpsBackend.build_communicator) try: with get_dist_manager().context(ctx): - comm = ctx.dcp_group.communicator - comm.prepare_query_gather(32, 576) + init_dcp_query_gather(32, 576) + workspace = ctx.dcp_group.query_gather_workspace if all(enabled): - assert comm._query_gatherer._state is not None + assert workspace.is_available() else: - assert comm._query_gatherer._state is None - assert comm._all_reduce is None + assert workspace is None + init_dcp_query_gather(32, 576) + assert ctx.dcp_group.query_gather_workspace is workspace query = torch.empty(384, 32, 576, device='cuda', dtype=torch.bfloat16) # Ineligible inputs still return correct results through NCCL. for fallback in (query[:1], query[..., ::2], query.float(), torch.empty(1024, 32, 576, device='cuda', dtype=query.dtype)): fallback.fill_(rank) - output = comm.gather_query(fallback) + output = gather_dcp_query(fallback, dcp_world_size=2) expected = torch.arange(2, device='cuda', dtype=fallback.dtype) expected = expected.repeat_interleave(32)[None, :, None].expand_as(output) torch.testing.assert_close(output, expected, rtol=0, atol=0) @@ -66,14 +72,100 @@ def _run_dcp_query_gather(rank, rendezvous, enabled): expected = (torch.arange(2, device='cuda', dtype=query.dtype) + step * 8) expected = expected.repeat_interleave(32)[None, :, None].expand_as(output) torch.testing.assert_close(output, expected, rtol=0, atol=0) + _check_lm_head_lifecycle(rank, all(enabled)) finally: ctx.close() dist.destroy_process_group() +def _run_auto_all_reduce(rank, rendezvous): + from datetime import timedelta + + from torch import distributed as dist + + from lmdeploy.pytorch.config import DistConfig + + torch.cuda.set_device(rank) + dist.init_process_group('nccl', init_method=rendezvous, rank=rank, world_size=2, + timeout=timedelta(seconds=60)) + cpu_group = dist.new_group(backend='gloo') + communicator = communicator_module.build_cuda_communicator( + cpu_group, dist.group.WORLD, DistConfig(tp=2, communication_backend='auto')) + try: + provider = communicator._all_reduce_provider + assert provider is not None + provider.all_reduce_ = Mock(wraps=provider.all_reduce_) + for dtype in (torch.bfloat16, torch.float16, torch.float32): + tensor = torch.empty(7, 4096, device='cuda', dtype=dtype) + tensor.fill_(rank + 1) + communicator.all_reduce_(tensor) + provider.all_reduce_.assert_called_with(tensor) + torch.testing.assert_close(tensor, torch.full_like(tensor, 3), rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + tensor.fill_(rank + 1) + communicator.all_reduce_(tensor) + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(tensor, torch.full_like(tensor, 3), rtol=0, atol=0) + graph.reset() + # Strided input retains the native NCCL error. + strided = torch.empty(7, 8192, device='cuda', dtype=torch.bfloat16)[:, ::2] + assert not provider.all_reduce_(strided) + with pytest.raises(ValueError, match='contiguous'): + communicator.all_reduce_(strided) + finally: + communicator.close() + dist.destroy_process_group() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='requires two CUDA GPUs') +def test_auto_all_reduce_eager_and_graph(tmp_path): + from torch.distributed._symmetric_memory import DeviceType, _SymmetricMemory + + if any(torch.cuda.get_device_capability(i) != (9, 0) + or not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, i) for i in range(2)): + pytest.skip('requires SM90 with multicast support') + torch.multiprocessing.spawn(_run_auto_all_reduce, + args=(f'file://{tmp_path}/auto_all_reduce',), nprocs=2, join=True) + + +def _check_lm_head_lifecycle(rank, direct): + from lmdeploy.pytorch.nn import ParallelEmbedding, ParallelLMHead + + head = ParallelLMHead(250, 128, dtype=torch.bfloat16, device='cuda') + if direct: + assert head._logits_gather_workspace.is_available() + head.weight.data.fill_(rank + 1) + hidden = torch.ones(2, 128, dtype=torch.bfloat16, device='cuda') + expected = torch.arange(1, 3, device='cuda', dtype=hidden.dtype).repeat_interleave(128)[:250] * 128 + logits = head(hidden) + head(hidden * 2) + torch.testing.assert_close(logits, expected.expand_as(logits), atol=0, rtol=0) + if direct: + assert head._logits_gather_workspace.is_available() + # A coordinated dtype move retires the BF16 arena and uses native gathering. + head.to(dtype=torch.float16) + if direct: + assert not head._logits_gather_workspace.is_available() + logits = head(hidden) + torch.testing.assert_close(logits, expected.to(logits.dtype).expand_as(logits), atol=0, rtol=0) + # Tying BF16 weights must re-admit the arena before the next forward. + embedding = ParallelEmbedding(250, 128, None, dtype=torch.bfloat16, device='cuda', is_tp=True) + embedding.weight.data.fill_(rank + 1) + head.tie_weights(embedding) + if direct: + assert head._logits_gather_workspace.is_available() + logits = head(hidden) + torch.testing.assert_close(logits, expected.expand_as(logits), atol=0, rtol=0) + if direct: + head._logits_gather_workspace.close() + + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason='requires two CUDA GPUs') @pytest.mark.parametrize('enabled', [(False, False), (True, True), (True, False)], - ids=['nccl', 'symm_mem', 'mixed_flags']) + ids=['nccl', 'auto', 'mixed_config']) def test_dcp_query_gather_graph_reuses_arena(tmp_path, enabled): if all(enabled): from torch.distributed._symmetric_memory import DeviceType, _SymmetricMemory @@ -86,192 +178,182 @@ def test_dcp_query_gather_graph_reuses_arena(tmp_path, enabled): args=((tmp_path / 'rendezvous').as_uri(), enabled), nprocs=2) -class _Collective: - - def __init__(self, *, result=None, handled=False): - self._result = result - self._handled = handled - - def is_available(self): - return True - - def fused_all_reduce_residual_rms_norm(self, **kwargs): - return self._result - - def all_reduce_(self, input): - return self._handled - - -def _build_communicator(monkeypatch, - *, - backend, - fused_result=None, - handled=False): - collective = _Collective(result=fused_result, handled=handled) - use_flashinfer = backend == 'flashinfer' - monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', use_flashinfer) - monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', not use_flashinfer) - flashinfer_cls = Mock(return_value=collective) - symm_mem_cls = Mock(return_value=collective) - monkeypatch.setattr(communicator_module, 'FlashInferAllReduce', flashinfer_cls) - monkeypatch.setattr(communicator_module, 'SymmetricMemoryAllReduce', symm_mem_cls) - communicator = communicator_module.CudaCommunicator(cpu_group='cpu', device_group='gpu', - all_reduce_backend=backend) - enabled_cls, disabled_cls = ((flashinfer_cls, symm_mem_cls) - if use_flashinfer else (symm_mem_cls, flashinfer_cls)) - enabled_cls.assert_called_once_with('cpu') - disabled_cls.assert_not_called() - return communicator, collective - - -def _build_norm(monkeypatch, *, fused=False): - impl = Mock() - backend = Mock() - backend.build_op.return_value = impl - group = Mock() - group.supports_optimized_all_reduce.return_value = True - group.supports_fused_all_reduce_residual_rms_norm.return_value = fused - monkeypatch.setattr(norm_module, 'get_backend', lambda: backend) - monkeypatch.setattr(norm_module, 'get_dist_group', lambda layer_type: group) - norm = norm_module.RMSNorm(4, dtype=torch.float32, device='cpu', all_reduce_group='attn') - return norm, impl, group - - -def test_cuda_communicator_dispatch(monkeypatch): - fused_result = object() - communicator, flashinfer = _build_communicator( - monkeypatch, backend='flashinfer', fused_result=fused_result) - nccl_all_reduce = Mock() - monkeypatch.setattr(communicator_module.dist, 'all_reduce', nccl_all_reduce) - - assert communicator.supports_optimized_all_reduce() - assert communicator.supports_fused_all_reduce_residual_rms_norm() - assert communicator.try_fused_all_reduce_residual_rms_norm( - input=torch.ones(1), residual=torch.ones(1), weight=torch.ones(1), eps=1e-6) is fused_result - - input = torch.ones(1) +@pytest.fixture +def comm_env(monkeypatch): + """Two ranks with available providers; no CUDA or distributed runtime.""" + monkeypatch.setattr(torch.cuda, 'get_device_name', lambda *args: 'GPU') + monkeypatch.setattr(torch.cuda, 'current_device', lambda: 0) + monkeypatch.setattr(communicator_module.dist, 'get_world_size', lambda group: 2) + monkeypatch.setattr(communicator_module.dist, 'get_rank', lambda group: 0) + monkeypatch.setattr(communicator_module.dist, 'all_gather_object', + lambda output, value, group: output.__setitem__(slice(None), [value, value])) + provider = Mock() + provider.is_available.return_value = True + monkeypatch.setattr(communicator_module, 'SymmetricMemoryAllReduce', Mock(return_value=provider)) + return provider + + +def test_gather_registration_and_native_fallback(monkeypatch): + from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_query, init_dcp_query_gather + from lmdeploy.pytorch.distributed import DistContext, DistGroup, get_dist_manager + + module = base_communicator_module + monkeypatch.setattr(module.dist, 'get_world_size', lambda group: 2) + + def gather(output, input, group): + output.copy_(torch.cat((input, input + 10), dim=0)) + + monkeypatch.setattr(module.dist, 'all_gather_into_tensor', gather) + communicator = module.DeviceCommunicator('group') + group = DistGroup(gpu_group='group', communicator=communicator) + context = DistContext(dcp_group=group) + with get_dist_manager().context(context): + init_dcp_query_gather(2, 4) + assert group.query_gather_workspace is None + logits = communicator.create_all_gather_workspace(6, device=torch.device('cpu'), dtype=torch.float32) + assert logits is None + query = torch.arange(48).view(3, 2, 8)[..., ::2] + with get_dist_manager().context(context): + torch.testing.assert_close(gather_dcp_query(query, dcp_world_size=2), torch.cat((query, query + 10), dim=1)) + input = torch.ones(3, 3) + torch.testing.assert_close(communicator.all_gather(input, workspace=logits), torch.cat((input, input + 10), dim=-1)) + + +def test_gather_preparation_and_weight_transition(monkeypatch): + from lmdeploy.pytorch.backends.cuda.comm.symm_mem_allgather import SymmetricMemoryAllGather + + workspace = SymmetricMemoryAllGather('group', 0, 16, device='cpu', dtype=torch.bfloat16, capacity_bytes=1024) + prepare = Mock(return_value=False) + monkeypatch.setattr(workspace, '_prepare', prepare) + workspace.reset('cpu', torch.float16) + prepare.assert_not_called() + workspace.prepare() + workspace.prepare() + prepare.assert_called_once_with(torch.device('cpu'), torch.float16) + workspace.reset('cpu', torch.float16) + prepare.assert_called_once() + workspace.reset('cpu', torch.bfloat16) + assert prepare.call_count == 2 + prepare.assert_called_with(torch.device('cpu'), torch.bfloat16) + assert not workspace.is_available() + workspace.close() + workspace.prepare() + assert prepare.call_count == 3 + + +def test_cuda_gather_dispatch_and_workspace_ownership(monkeypatch, comm_env): + from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_query, init_dcp_query_gather + from lmdeploy.pytorch.backends.cuda.comm import symm_mem_allgather + from lmdeploy.pytorch.distributed import DistContext, DistGroup, get_dist_manager + + # The group shares one query workspace; logits workspaces remain caller-owned. + factory = Mock(side_effect=lambda *args, **kwargs: Mock()) + monkeypatch.setattr(symm_mem_allgather, 'SymmetricMemoryAllGather', factory) + communicator = communicator_module.CudaCommunicator('cpu', 'gpu', group_name='dcp') + group = DistGroup(gpu_group='gpu', communicator=communicator) + context = DistContext(dcp_group=group) + with get_dist_manager().context(context): + init_dcp_query_gather(2, 4) + init_dcp_query_gather(2, 4) + query_workspace = group.query_gather_workspace + logits_workspace = communicator.create_all_gather_workspace(6, device=torch.device('cpu'), dtype=torch.float32) + assert factory.call_count == 2 + query_workspace.prepare.assert_called_once() + logits_workspace.prepare.assert_called_once() + + query = torch.arange(24).view(3, 2, 4) + gathered_query = torch.cat((query, query + 10), dim=1) + query_workspace.all_gather.return_value = gathered_query.flatten(1) + with get_dist_manager().context(context): + torch.testing.assert_close(gather_dcp_query(query, dcp_world_size=2), gathered_query) + assert query_workspace.all_gather.call_args.kwargs == {'copy_output': False} + + native = Mock(side_effect=lambda output, input, group: output.copy_(torch.cat((input, input + 10)))) + monkeypatch.setattr(base_communicator_module.dist, 'all_gather_into_tensor', native) + input = torch.ones(3, 3) + logits_workspace.all_gather.return_value = torch.ones(3, 6) + torch.testing.assert_close(communicator.all_gather(input, workspace=logits_workspace), torch.ones(3, 6)) + logits_workspace.all_gather.assert_called_once_with(input, copy_output=True) + native.assert_not_called() + logits_workspace.all_gather.return_value = None + torch.testing.assert_close(communicator.all_gather(input, workspace=logits_workspace), + torch.cat((input, input + 10), dim=-1)) + native.assert_called_once() + + group.close() + group.close() + query_workspace.close.assert_called_once() + assert group.query_gather_workspace is None + logits_workspace.close.assert_not_called() + + +@pytest.mark.parametrize('native', ['none', 'async', 'max', 'cpu', 'other_group', 'no_communicator']) +def test_public_all_reduce_routes_only_synchronous_cuda_tp_sum(monkeypatch, native): + import lmdeploy.pytorch.distributed as distributed + + communicator = Mock() + group = object() + context = distributed.DistContext( + attn_tp_group=distributed.DistGroup(gpu_group=group, communicator=communicator)) + if native == 'no_communicator': + context.attn_tp_group.communicator = None + tensor = SimpleNamespace(is_cuda=native != 'cpu') + op = distributed.ReduceOp.MAX if native == 'max' else distributed.ReduceOp.SUM + async_op = native == 'async' + selected_group = object() if native == 'other_group' else group + native_reduce = Mock() + monkeypatch.setattr(distributed.dist, 'all_reduce', native_reduce) + with distributed.get_dist_manager().context(context): + result = distributed.all_reduce(tensor, op=op, group=selected_group, async_op=async_op) + if native == 'none': + communicator.all_reduce_.assert_called_once_with(tensor) + native_reduce.assert_not_called() + else: + native_reduce.assert_called_once_with(tensor, op, selected_group, async_op) + communicator.all_reduce_.assert_not_called() + assert result is native_reduce.return_value + + +def test_all_reduce_workspace_failure_agrees_across_ranks(monkeypatch, comm_env): + def peer_failure(flags, ready, group): + flags[:] = [ready, not comm_env.prepare.called] + + monkeypatch.setattr(communicator_module.dist, 'all_gather_object', peer_failure) + communicator = communicator_module.CudaCommunicator('cpu', 'gpu', group_name='tp') + comm_env.close.assert_called_once() + comm_env.prepare.assert_called_once() + assert communicator._all_reduce_provider is None + native = Mock() + monkeypatch.setattr(communicator_module.dist, 'all_reduce', native) + input = torch.ones(2, 4) communicator.all_reduce_(input) - nccl_all_reduce.assert_called_once_with(input, group='gpu') + native.assert_called_once_with(input, group='gpu') - flashinfer._handled = True - communicator.all_reduce_(input) - nccl_all_reduce.assert_called_once() - communicator, symm_mem = _build_communicator(monkeypatch, backend='symm_mem', handled=True) +@pytest.mark.parametrize('handled', [True, False]) +def test_auto_all_reduce_falls_back_for_unsupported_inputs(monkeypatch, comm_env, handled): + communicator = communicator_module.CudaCommunicator('cpu', 'gpu', group_name='tp') + comm_env.all_reduce_.return_value = handled + native = Mock() + monkeypatch.setattr(communicator_module.dist, 'all_reduce', native) + comm_env.prepare.assert_called_once() + input = torch.ones(7, 128, dtype=torch.bfloat16) communicator.all_reduce_(input) - assert communicator.supports_optimized_all_reduce() - assert not communicator.supports_fused_all_reduce_residual_rms_norm() - nccl_all_reduce.assert_called_once() - - -def test_cuda_communicator_rejects_multiple_backends(monkeypatch): - monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', True) - monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', True) - with pytest.raises(ValueError, match='cannot be enabled together'): - config = SimpleNamespace(dp=1, ep=1, attn_tp=2, dcp=1, enable_microbatch=False) - communicator_module.build_cuda_communicator('cpu', 'gpu', config) - - -def test_rms_norm_fuses_or_falls_back(monkeypatch): - norm, impl, group = _build_norm(monkeypatch, fused=True) - fused_output = (torch.full((1, 4), 2.0), torch.full((1, 4), 3.0)) - group.try_fused_all_reduce_residual_rms_norm.return_value = fused_output - input = torch.ones((1, 4), dtype=torch.bfloat16) - residual = torch.ones_like(input) - - assert norm(input, residual) is fused_output - group.try_fused_all_reduce_residual_rms_norm.assert_called_once_with( - input=input, - residual=residual, - weight=norm.weight, - eps=norm.eps, - ) - group.all_reduce_.assert_not_called() - impl.forward.assert_not_called() - - group.reset_mock() - impl.reset_mock() - group.try_fused_all_reduce_residual_rms_norm.return_value = None - impl.forward.return_value = object() - output = norm(input, residual) - - assert output is impl.forward.return_value - group.try_fused_all_reduce_residual_rms_norm.assert_called_once() - group.all_reduce_.assert_called_once_with(input) - impl.forward.assert_called_once_with(input, norm.weight, residual) - - -def test_flashinfer_allreduce_in_place_and_dtype_guards(monkeypatch): - from lmdeploy.pytorch.backends.cuda.comm import flashinfer_allreduce as flashinfer_module - - FlashInferAllReduce = flashinfer_module.FlashInferAllReduce - - flashinfer = FlashInferAllReduce.__new__(FlashInferAllReduce) - flashinfer.is_available = Mock(return_value=True) - assert flashinfer.supports(torch.float16) - assert flashinfer.supports(torch.bfloat16) - assert not flashinfer.supports(torch.float32) - - flashinfer._max_size = 1024 - flashinfer._one_shot_max_size = 1024 - flashinfer._comm = SimpleNamespace( - AllReduceFusionPattern=SimpleNamespace(kAllReduce=0), - allreduce_fusion=Mock( - return_value=torch.full((2, 4), 2.0, dtype=torch.bfloat16)), - ) - flashinfer._get_workspace = Mock(return_value='workspace') - flashinfer.supports = Mock(return_value=True) - - input = torch.ones(1, 2, 4, dtype=torch.bfloat16) - assert flashinfer.all_reduce_(input) - torch.testing.assert_close(input, torch.full_like(input, 2.0)) - call_kwargs = flashinfer._comm.allreduce_fusion.call_args.kwargs - assert call_kwargs['input'].shape == (2, 4) - assert call_kwargs['trigger_completion_at_end'] - assert call_kwargs['use_oneshot'] - - fused_calls = flashinfer._comm.allreduce_fusion.call_count - bf16_input = torch.ones(2, 4, dtype=torch.bfloat16) - output = flashinfer.fused_all_reduce_residual_rms_norm( - input=bf16_input, - residual=torch.ones_like(bf16_input), - weight=torch.ones(4, dtype=torch.float32), - eps=1e-6, - ) - assert output is None - assert flashinfer._comm.allreduce_fusion.call_count == fused_calls - - unavailable = FlashInferAllReduce.__new__(FlashInferAllReduce) - unavailable._comm = None - unavailable._disabled = False - unavailable._max_size = 1024 - monkeypatch.setitem(sys.modules, 'flashinfer.comm', None) - assert not unavailable.is_available() - assert unavailable._disabled - - workspace_error = FlashInferAllReduce.__new__(FlashInferAllReduce) - workspace_error.group = 'cpu' - workspace_error._world_size = 2 - workspace_error._max_size = 1024 - workspace_error._one_shot_max_size = 1024 - workspace_error._workspace = None - workspace_error._hidden_dim = None - workspace_error._dtype = None - workspace_error._disabled = False - create_workspace = Mock(side_effect=RuntimeError('unsupported topology')) - workspace_error._comm = SimpleNamespace( - AllReduceFusionPattern=SimpleNamespace(kAllReduce=0), - create_allreduce_fusion_workspace=create_workspace, - allreduce_fusion=Mock(), - ) - monkeypatch.setattr(flashinfer_module.dist, 'get_rank', lambda group: 0) + comm_env.all_reduce_.assert_called_once_with(input) + if handled: + native.assert_not_called() + else: + native.assert_called_once_with(input, group='gpu') + - input = torch.ones(2, 4, dtype=torch.bfloat16) - assert not workspace_error.all_reduce_(input) - assert workspace_error._disabled - assert not workspace_error.all_reduce_(input) - create_workspace.assert_called_once() - workspace_error._comm.allreduce_fusion.assert_not_called() +@pytest.mark.parametrize('group_name,available', [('tp', True), ('tp', False), ('dcp', True)]) +def test_auto_initializes_all_reduce_by_availability(comm_env, group_name, available): + comm_env.is_available.return_value = available + communicator = communicator_module.CudaCommunicator('cpu', 'gpu', group_name=group_name) + assert communicator._all_reduce_provider is (comm_env if available and group_name == 'tp' else None) + assert communicator_module.SymmetricMemoryAllReduce.call_count == (group_name == 'tp') + assert comm_env.close.call_count == (group_name == 'tp' and not available) + assert comm_env.prepare.call_count == (group_name == 'tp' and available) def test_symm_mem_allreduce_selects_group_algorithm(monkeypatch): @@ -291,7 +373,6 @@ def test_symm_mem_allreduce_selects_group_algorithm(monkeypatch): communicator = SymmetricMemoryAllReduce.__new__(SymmetricMemoryAllReduce) communicator.group = SimpleNamespace(group_name='group') communicator._buffer = torch.empty(8, dtype=torch.bfloat16) - communicator._enabled = True communicator._max_size = communicator._buffer.nbytes input = torch.ones(4, dtype=torch.bfloat16) @@ -305,64 +386,61 @@ def test_symm_mem_allreduce_selects_group_algorithm(monkeypatch): two_shot.assert_called_once() -def test_build_cuda_communicator_gates_unsupported_parallelism(monkeypatch): - monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', True) - monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', False) - communicator_cls = Mock(return_value=object()) - monkeypatch.setattr(communicator_module, 'CudaCommunicator', communicator_cls) - - for config in ( - SimpleNamespace(dp=2, ep=1, attn_tp=8, dcp=1, enable_microbatch=False), - SimpleNamespace(dp=1, ep=2, attn_tp=8, dcp=1, enable_microbatch=False), - SimpleNamespace(dp=1, ep=1, attn_tp=8, dcp=1, enable_microbatch=True), - ): - assert communicator_module.build_cuda_communicator('cpu', 'device', config) is None - communicator_cls.assert_not_called() - - config = SimpleNamespace(dp=1, ep=1, attn_tp=4, dcp=1, enable_microbatch=False) - communicator = communicator_module.build_cuda_communicator('cpu', 'device', config) - assert communicator is communicator_cls.return_value - communicator_cls.assert_called_once_with(cpu_group='cpu', device_group='device', - all_reduce_backend='flashinfer', symm_mem_query_gather=False) - - monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', False) - monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', True) - config.attn_tp = 1 - assert communicator_module.build_cuda_communicator('cpu', 'device', config) is None - config.attn_tp = 2 - assert communicator_module.build_cuda_communicator('cpu', 'device', config) is communicator_cls.return_value - - # TP all-reduce flags must not allocate all-reduce resources on DCP groups. - config.dcp = 2 - monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_dcp', True) - communicator_module.build_cuda_communicator('cpu', 'device', config, group_roles=('dcp', )) - communicator_cls.assert_called_with(cpu_group='cpu', device_group='device', - all_reduce_backend=None, symm_mem_query_gather=True) - - -@pytest.mark.parametrize('shared_dcp_group', [False, True]) -def test_communicator_group_roles(shared_dcp_group): +def test_symm_mem_allreduce_peer_allocation_failure(monkeypatch): + import torch.distributed._symmetric_memory as symm_mem + + from lmdeploy.pytorch.backends.cuda.comm import symm_mem_allreduce as module + + monkeypatch.setattr(module.dist, 'get_world_size', lambda group: 2) + monkeypatch.setattr(torch.cuda, 'get_device_capability', lambda: (9, 0)) + monkeypatch.setattr(torch.cuda, 'current_device', lambda: 0) + monkeypatch.setattr(torch.ops, 'symm_mem', SimpleNamespace(two_shot_all_reduce_=Mock())) + monkeypatch.setattr(symm_mem, 'empty', Mock(return_value=object())) + rendezvous = Mock() + monkeypatch.setattr(symm_mem, 'rendezvous', rendezvous) + + def peer_failure(flags, ready, group): + flags[:] = [ready, False] + + monkeypatch.setattr(module.dist, 'all_gather_object', peer_failure) + provider = module.SymmetricMemoryAllReduce(SimpleNamespace(group_name='group')) + symm_mem.empty.assert_not_called() + provider.prepare() + assert not provider.is_available() + assert provider._buffer is None + rendezvous.assert_not_called() + + +def test_communicator_rejects_rank_backend_mismatch(monkeypatch, comm_env): + from lmdeploy.pytorch.backends.cuda.op_backend import CudaOpsBackend + from lmdeploy.pytorch.config import DistConfig + + monkeypatch.setattr(torch.cuda, 'get_device_capability', lambda device: (9, 0)) + + def disagree(output, local, group): + output[:] = [local, ('auto', *local[1:])] + + monkeypatch.setattr(communicator_module.dist, 'all_gather_object', disagree) + optimized = Mock() + monkeypatch.setattr(communicator_module, 'CudaCommunicator', optimized) + with pytest.raises(ValueError, match='agree across ranks'): + CudaOpsBackend.build_communicator('cpu', 'gpu', DistConfig(tp=2, communication_backend='nccl')) + optimized.assert_not_called() + + +@pytest.mark.parametrize('shared_tp_group', [False, True]) +def test_communicator_group_names_and_tp_deduplication(shared_tp_group): from lmdeploy.pytorch.distributed import DistContext, DistGroup, _build_communicators tp_group = DistGroup(cpu_group='tp_cpu', gpu_group='tp_gpu') - dcp_group = tp_group if shared_dcp_group else DistGroup(cpu_group='dcp_cpu', gpu_group='dcp_gpu') - builder = Mock(side_effect=lambda **kwargs: object()) - context = DistContext(attn_tp_group=tp_group, mlp_tp_group=tp_group, moe_tp_group=tp_group, + mlp_group = tp_group if shared_tp_group else DistGroup(cpu_group='mlp_cpu', gpu_group='mlp_gpu') + dcp_group = DistGroup(cpu_group='dcp_cpu', gpu_group='dcp_gpu') + builder = Mock(side_effect=lambda **kwargs: Mock()) + context = DistContext(attn_tp_group=tp_group, mlp_tp_group=mlp_group, moe_tp_group=mlp_group, dcp_group=dcp_group, communicator_builder=builder) _build_communicators(context) - roles = [call.kwargs['group_roles'] for call in builder.call_args_list] - assert roles == ([('tp', 'dcp')] if shared_dcp_group else [('tp', ), ('dcp', )]) - assert (tp_group.communicator is dcp_group.communicator) == shared_dcp_group - - -def test_dlinfer_communicator_rejects_cuda_options(monkeypatch): - monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', True) - monkeypatch.setattr(communicator_module._envs, 'enable_symm_mem_allreduce', False) - with pytest.raises(AssertionError, match='not supported by DLInfer'): - DlinferOpsBackend.build_communicator('cpu', 'device', SimpleNamespace()) - - monkeypatch.setattr(communicator_module._envs, 'enable_flashinfer_allreduce', False) - communicator = DlinferOpsBackend.build_communicator('cpu', 'device', SimpleNamespace()) - assert isinstance(communicator, base_communicator_module.DeviceCommunicator) - assert communicator.device_group == 'device' + names = [call.kwargs['group_name'] for call in builder.call_args_list] + assert names == (['tp', 'dcp'] if shared_tp_group else ['tp', 'tp', 'dcp']) + assert (tp_group.communicator is mlp_group.communicator) == shared_tp_group + assert tp_group.communicator is not dcp_group.communicator From 525073facd769767f9cc7fc2b3f15f1e0b31848e Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 24 Sep 2026 04:09:29 +0000 Subject: [PATCH 06/15] refactor: centralize DCP workspace and collective ownership --- .../pytorch/backends/cuda/attention/cp.py | 145 ++++++++++-------- .../pytorch/backends/cuda/attention/mla.py | 14 +- .../backends/cuda/attention/sparse_mla.py | 8 +- lmdeploy/pytorch/backends/cuda/nsa.py | 14 +- lmdeploy/pytorch/distributed.py | 11 +- tests/pytorch/kernel/test_mla_attention.py | 14 +- tests/pytorch/test_cuda_communicator.py | 105 ++++++++++--- 7 files changed, 194 insertions(+), 117 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/attention/cp.py b/lmdeploy/pytorch/backends/cuda/attention/cp.py index dba14b1f02..14669ef1ac 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/cp.py +++ b/lmdeploy/pytorch/backends/cuda/attention/cp.py @@ -27,66 +27,91 @@ def gather_dcp_prefix_kv(local_kv: torch.Tensor, chunk: DCPPrefixChunk, return context -def init_dcp_query_gather(num_heads: int, head_size: int): - """Initialize one query workspace shared by attention layers in this DCP - group.""" - from lmdeploy.pytorch.distributed import get_dist_manager - - group = get_dist_manager().current_context().dcp_group - if group.communicator is not None and group.query_gather_workspace is None: - width = num_heads * head_size * torch.distributed.get_world_size(group.gpu_group) - group.query_gather_workspace = group.communicator.create_all_gather_workspace( - width, device=torch.device('cuda'), dtype=torch.bfloat16) - +class DCPManager: + """Own shared DCP gathers and the output combine that orders query reuse. -def gather_dcp_query(query: torch.Tensor, *, dcp_world_size: int) -> torch.Tensor: - """Gather [tokens, local_heads, dim] queries along the head axis. - - Consume the result on the same stream before the next query gather; the optimized path borrows a shared arena. + Layers in one distributed context share this manager. Prepare workspaces during model construction, before CUDA + graph capture, and close only after all consumers and graphs have retired. """ - if dcp_world_size == 1: - return query - from lmdeploy.pytorch.distributed import get_dist_manager - group = get_dist_manager().current_context().dcp_group - output = group.communicator.all_gather(query.flatten(1), workspace=group.query_gather_workspace, copy_output=False) - return output.view(query.size(0), -1, query.size(2)) - - -def merge_dcp_attention(local_output: torch.Tensor, - local_lse: torch.Tensor, - valid_counts: torch.Tensor, - *, - dcp_world_rank: tuple[int, int]) -> torch.Tensor: - """Merge normalized CUDA shard outputs and scatter heads back to each rank. + def __init__(self, group): + self.group = group.gpu_group + self.communicator = group.communicator + self.world_size = torch.distributed.get_world_size(self.group) + self.rank = group.rank + self._query_workspace = None + + def prepare_query_gather(self, num_heads: int, head_size: int): + """Prepare one query arena shared across attention layers.""" + if self._query_workspace is None: + width = num_heads * head_size * self.world_size + self._query_workspace = self.communicator.create_all_gather_workspace( + width, device=torch.device('cuda'), dtype=torch.bfloat16) + + def gather_query(self, query: torch.Tensor) -> torch.Tensor: + """Gather query heads; consume and combine on the same stream before + reuse.""" + if self.world_size == 1: + return query + output = self.communicator.all_gather( + query.flatten(1), workspace=self._query_workspace, copy_output=False) + return output.view(query.size(0), -1, query.size(2)) + + def gather_candidates(self, packed: torch.Tensor) -> torch.Tensor: + """Gather [rows, topk, score/ID] payloads in rank-major order.""" + from lmdeploy.pytorch.distributed import all_gather_into_tensor + + output = packed.new_empty(self.world_size * packed.size(0), *packed.shape[1:]) + all_gather_into_tensor(output, packed, group=self.group) + return output.view(self.world_size, *packed.shape) + + def combine(self, local_output: torch.Tensor, local_lse: torch.Tensor, + valid_counts: torch.Tensor) -> torch.Tensor: + """Merge shard outputs and scatter heads, ordering the next query + gather. + + Inputs are [tokens, gathered_heads, value_dim] outputs and natural-log [tokens, gathered_heads] LSE. Empty rows + contribute zero. LSE and correction arithmetic use FP32; reduce-scatter uses the input output dtype. + """ + if self.world_size == 1: + return local_output + from lmdeploy.pytorch.distributed import all_gather_into_tensor, reduce_scatter_tensor + from lmdeploy.pytorch.kernels.cuda.dcp import correct_dcp_attention_output, sanitize_dcp_lse + + local_lse = sanitize_dcp_lse(local_lse, valid_counts) + gathered_lse = local_lse.new_empty( + self.world_size * local_lse.size(0), local_lse.size(1)) + all_gather_into_tensor(gathered_lse, local_lse, group=self.group) + gathered_lse = gathered_lse.view(self.world_size, *local_lse.shape) + contribution = correct_dcp_attention_output( + local_output, gathered_lse, dcp_rank=self.rank) + + num_heads = contribution.size(0) + # The correction kernel writes [heads, tokens, dim] directly so the + # head-sharded reduce-scatter needs no separate transpose/copy. + assert num_heads % self.world_size == 0 + local_heads = num_heads // self.world_size + scattered_output = contribution.new_empty(local_heads, + contribution.size(1), + contribution.size(2)) + reduce_scatter_tensor(scattered_output, contribution, group=self.group) + return scattered_output.transpose(0, 1) + + def close(self): + """Release shared arenas before their process group is destroyed.""" + if self._query_workspace is not None: + self._query_workspace.close() + self._query_workspace = None + + +def get_dcp_manager() -> DCPManager: + """Get the context-owned manager during layer initialization.""" + from lmdeploy.pytorch.distributed import get_dist_manager - Inputs are [tokens, gathered_heads, value_dim] outputs and natural-log [tokens, gathered_heads] LSE. Empty rows - contribute zero. LSE and correction arithmetic use FP32; reduce-scatter uses the input output dtype. - """ - dcp_world_size, dcp_rank = dcp_world_rank - if dcp_world_size == 1: - return local_output - from lmdeploy.pytorch.distributed import all_gather_into_tensor, reduce_scatter_tensor - from lmdeploy.pytorch.kernels.cuda.dcp import correct_dcp_attention_output, sanitize_dcp_lse - - local_lse = sanitize_dcp_lse(local_lse, valid_counts) - gathered_lse = local_lse.new_empty( - dcp_world_size * local_lse.size(0), local_lse.size(1)) - all_gather_into_tensor(gathered_lse, local_lse, group='dcp') - gathered_lse = gathered_lse.view(dcp_world_size, *local_lse.shape) - contribution = correct_dcp_attention_output( - local_output, gathered_lse, dcp_rank=dcp_rank) - - num_heads = contribution.size(0) - # The correction kernel writes [heads, tokens, dim] directly so the - # head-sharded reduce-scatter needs no separate transpose/copy. - assert num_heads % dcp_world_size == 0 - local_heads = num_heads // dcp_world_size - scattered_output = contribution.new_empty(local_heads, - contribution.size(1), - contribution.size(2)) - reduce_scatter_tensor(scattered_output, contribution, group='dcp') - return scattered_output.transpose(0, 1) + context = get_dist_manager().current_context() + if context.dcp_manager is None: + context.dcp_manager = DCPManager(context.dcp_group) + return context.dcp_manager class DCPAttentionImpl(TritonAttentionImpl): @@ -102,7 +127,8 @@ def __init__(self, num_heads: int, head_size: int, *, use_fa3: bool = False, **k if use_fa3: from lmdeploy.pytorch.third_party.flash_attn_interface import flash_attn_varlen_func self._fa3_prefill = flash_attn_varlen_func - init_dcp_query_gather(self.num_heads, self.head_size) + self.dcp_manager = get_dcp_manager() + self.dcp_manager.prepare_query_gather(self.num_heads, self.head_size) def get_step_metadata_provider(self): """Use common DCP lengths without a separate kernel scheduler.""" @@ -115,7 +141,7 @@ def _decode(self, query: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Ten metadata: TritonAttentionMetadata, k_scales_zeros: torch.Tensor = None, v_scales_zeros: torch.Tensor = None) -> torch.Tensor: """Evaluate interleaved shards and restore rank-local query heads.""" - query = gather_dcp_query(query, dcp_world_size=self.dcp_world_size) + query = self.dcp_manager.gather_query(query) query_len = query.size(0) // metadata.q_seqlens.numel() local_lens = metadata.dcp_local_kv_seqlens blocks = metadata.block_offsets @@ -129,8 +155,7 @@ def _decode(self, query: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Ten page_table=blocks, max_seqlen_q=1, softmax_scale=self.scale, softcap=self.logit_softcapping, return_lse=True, quant_policy=metadata.quant_policy, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros) - return merge_dcp_attention(output, lse, valid_counts=local_lens, - dcp_world_rank=(self.dcp_world_size, self.dcp_rank)) + return self.dcp_manager.combine(output, lse, valid_counts=local_lens) def _prefill_attention(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, *, cu_q: torch.Tensor, cu_k: torch.Tensor, max_q: int, max_k: int, diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index 8fdd344c4d..f87f78d69a 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -14,7 +14,7 @@ from lmdeploy.utils import get_logger from ..step_metadata import CudaAttentionMetaBuilder -from .cp import gather_dcp_prefix_kv, gather_dcp_query, merge_dcp_attention +from .cp import gather_dcp_prefix_kv, get_dcp_manager from .default import TritonAttentionImpl, TritonAttentionMetadata logger = get_logger('lmdeploy') @@ -255,9 +255,8 @@ def __init__( self.use_fa3 = use_fa3 if self.dcp_world_size > 1: - from .cp import init_dcp_query_gather - - init_dcp_query_gather(num_heads, head_size) + self.dcp_manager = get_dcp_manager() + self.dcp_manager.prepare_query_gather(num_heads, head_size) def get_step_metadata_provider(self): """Describe metadata required by this selected implementation.""" @@ -710,7 +709,7 @@ def _forward_decoding( if self.dcp_world_size == 1: return self._decoding_paged(query, k_cache, attn_metadata) - query = gather_dcp_query(query, dcp_world_size=self.dcp_world_size) + query = self.dcp_manager.gather_query(query) query_len = query.size(0) // attn_metadata.q_seqlens.numel() if query_len > 1: # Consecutive global queries do not advance each interleaved @@ -725,11 +724,10 @@ def _forward_decoding( ) local_output, local_lse = self._decoding_paged( query, k_cache, attn_metadata, causal=False, return_lse=True) - return merge_dcp_attention( + return self.dcp_manager.combine( local_output, local_lse, - valid_counts=attn_metadata.dcp_local_kv_seqlens, - dcp_world_rank=(self.dcp_world_size, self.dcp_rank)) + valid_counts=attn_metadata.dcp_local_kv_seqlens) def _forward_prefill( self, diff --git a/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py b/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py index db9966ef59..9033a2c3b8 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py @@ -6,7 +6,6 @@ from lmdeploy.utils import get_logger -from .cp import gather_dcp_query, merge_dcp_attention from .default import TritonAttentionMetadata from .mla import FlashMLAAttentionMetaBuilder, FlashMLAImpl @@ -326,14 +325,11 @@ def _forward_decoding(self, dcp_world_rank = self.dcp_world_size, self.dcp_rank local_indices, local_counts = filter_and_compact_dcp_indices( nsa_indices, dcp_world_rank=dcp_world_rank) - query = gather_dcp_query(query, dcp_world_size=self.dcp_world_size) + query = self.dcp_manager.gather_query(query) local_output, local_lse = self._decoding_sparse( query, k_cache, local_indices, attn_metadata, return_lse=True, topk_length=local_counts) - return merge_dcp_attention(local_output, - local_lse, - valid_counts=local_counts, - dcp_world_rank=dcp_world_rank) + return self.dcp_manager.combine(local_output, local_lse, valid_counts=local_counts) return self._decoding_sparse(query, k_cache, nsa_indices, attn_metadata) def _forward_prefill(self, diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index 1122d7003e..6c708dd89c 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -365,6 +365,8 @@ def __init__(self, topk: int, softmax_scale: float, block_size: int, raise RuntimeError('DCP requires the TileLang sparse index top-k kernel.') if _get_deep_gemm() is None: raise RuntimeError('DCP DSA scoring requires a compatible DeepGEMM installation.') + from .attention.cp import get_dcp_manager + self.dcp_manager = get_dcp_manager() self._step_meta_group: int | None = None self._piecewise_forward: Callable[..., Tensor] | None = None self._piecewise_forward_fused: Callable[..., Tensor] | None = None @@ -476,26 +478,16 @@ def _merge_dcp_topk(self, scores: Tensor, local_indices: Tensor) -> Tensor: if self.dcp_world_size == 1: return local_indices - from lmdeploy.pytorch.distributed import all_gather_into_tensor from lmdeploy.pytorch.kernels.cuda.sparse_index_dcp_topk import ( pack_dcp_topk_candidates, sparse_dcp_global_topk, ) - num_rows = scores.size(0) packed = pack_dcp_topk_candidates( scores, local_indices, dcp_world_rank=(self.dcp_world_size, self.dcp_rank)) - gathered = torch.empty( - self.dcp_world_size * num_rows, - self.topk, - 2, - dtype=packed.dtype, - device=scores.device, - ) - all_gather_into_tensor(gathered, packed, group='dcp') - gathered = gathered.view(self.dcp_world_size, num_rows, self.topk, 2) + gathered = self.dcp_manager.gather_candidates(packed) return sparse_dcp_global_topk(gathered, k=self.topk, fill=self.fill) def _score_and_select_prefill( diff --git a/lmdeploy/pytorch/distributed.py b/lmdeploy/pytorch/distributed.py index f76ca7145a..b556687ebb 100644 --- a/lmdeploy/pytorch/distributed.py +++ b/lmdeploy/pytorch/distributed.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: from .backends.communicator import DeviceCommunicator + from .backends.cuda.attention.cp import DCPManager @dataclass @@ -26,16 +27,12 @@ class DistGroup: gpu_groups: list[dist.ProcessGroup] = None gpu_gather_group: dist.ProcessGroup = None communicator: 'DeviceCommunicator' = None - query_gather_workspace: object = None def close(self): """Close groups.""" if self.communicator is not None: self.communicator.close() self.communicator = None - if self.query_gather_workspace is not None: - self.query_gather_workspace.close() - self.query_gather_workspace = None if not dist.is_initialized(): return if self.cpu_groups is not None: @@ -243,6 +240,7 @@ class DistContext: mlp_tp_group: DistGroup = None moe_tp_group: DistGroup = None dcp_group: DistGroup = None + dcp_manager: 'DCPManager' = None cpu_group: dist.ProcessGroup = None ep_gpu_group: dist.ProcessGroup = None @@ -324,7 +322,10 @@ def build(cls, return context def close(self): - """Close groups.""" + """Close DCP resources before their process groups.""" + if self.dcp_manager is not None: + self.dcp_manager.close() + self.dcp_manager = None if not dist.is_initialized(): return if self.dcp_group is not None: diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py index 99446f4206..7faf7624ea 100644 --- a/tests/pytorch/kernel/test_mla_attention.py +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -9,7 +9,7 @@ from lmdeploy.pytorch.backends.cuda import attention as attention_module from lmdeploy.pytorch.backends.cuda.attention import mla as mla_module from lmdeploy.pytorch.backends.cuda.attention import sparse_mla as sparse_mla_module -from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_query +from lmdeploy.pytorch.backends.cuda.attention.cp import DCPManager from lmdeploy.pytorch.backends.cuda.attention.sparse_mla import ( FlashMLAIndexMapper, FlashMLASparseImpl, @@ -160,14 +160,14 @@ def fake_all_gather(output, input_tensor, group='tp', async_op=False): if not strided: assert input_tensor.data_ptr() == rank0_query.data_ptr() for rank in range(dcp_size): - source = queries[rank].transpose(0, 1) if strided else queries[rank] + source = queries[rank].flatten(1) output[rank * source.size(0):(rank + 1) * source.size(0)].copy_(source) monkeypatch.setattr(torch.distributed, 'all_gather_into_tensor', fake_all_gather) monkeypatch.setattr(torch.distributed, 'get_world_size', lambda group: dcp_size) - monkeypatch.setattr(distributed.get_dist_manager().current_context().dcp_group, - 'communicator', DeviceCommunicator('dcp')) - gathered = gather_dcp_query(rank0_query, dcp_world_size=dcp_size) + group = distributed.DistGroup(gpu_group='dcp', communicator=DeviceCommunicator('dcp')) + manager = DCPManager(group) + gathered = manager.gather_query(rank0_query) expected = torch.cat(list(queries), dim=1) assert torch.equal(gathered, expected) @@ -179,7 +179,7 @@ def fake_all_gather(output, input_tensor, group='tp', async_op=False): torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - graph_output = gather_dcp_query(rank0_query, dcp_world_size=dcp_size) + graph_output = manager.gather_query(rank0_query) queries.add_(1) graph.replay() torch.cuda.synchronize() @@ -498,6 +498,8 @@ def test_dcp_cached_prefill_matches_reference_across_chunks(monkeypatch, sparse, mla_index_topk=512 if sparse else None), kv_quant_policy=0) monkeypatch.setattr(distributed, 'get_dcp_world_rank', lambda: (dcp_size, 0)) + # This prefill test simulates ranks without creating process groups. + monkeypatch.setattr(mla_module, 'get_dcp_manager', Mock()) # One virtual block per chunk, with an empty second request in every chunk. monkeypatch.setattr(cp_utils, 'get_dcp_prefill_workspace_size', lambda **kwargs: 2 * 64 * (1 + 2 * dcp_size) * 576 * 2) diff --git a/tests/pytorch/test_cuda_communicator.py b/tests/pytorch/test_cuda_communicator.py index aa26b0f22e..158c24e3b6 100644 --- a/tests/pytorch/test_cuda_communicator.py +++ b/tests/pytorch/test_cuda_communicator.py @@ -14,7 +14,7 @@ def _run_dcp_query_gather(rank, rendezvous, enabled): from torch import distributed as dist - from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_query, init_dcp_query_gather + from lmdeploy.pytorch.backends.cuda.attention.cp import get_dcp_manager from lmdeploy.pytorch.backends.cuda.op_backend import CudaOpsBackend from lmdeploy.pytorch.config import DistConfig from lmdeploy.pytorch.distributed import DistContext, get_dist_manager @@ -35,26 +35,35 @@ def _run_dcp_query_gather(rank, rendezvous, enabled): communicator_builder=CudaOpsBackend.build_communicator) try: with get_dist_manager().context(ctx): - init_dcp_query_gather(32, 576) - workspace = ctx.dcp_group.query_gather_workspace + manager = get_dcp_manager() + manager.prepare_query_gather(32, 576) + workspace = manager._query_workspace if all(enabled): assert workspace.is_available() else: assert workspace is None - init_dcp_query_gather(32, 576) - assert ctx.dcp_group.query_gather_workspace is workspace + manager = get_dcp_manager() + manager.prepare_query_gather(32, 576) + assert manager._query_workspace is workspace query = torch.empty(384, 32, 576, device='cuda', dtype=torch.bfloat16) # Ineligible inputs still return correct results through NCCL. for fallback in (query[:1], query[..., ::2], query.float(), torch.empty(1024, 32, 576, device='cuda', dtype=query.dtype)): fallback.fill_(rank) - output = gather_dcp_query(fallback, dcp_world_size=2) + output = manager.gather_query(fallback) expected = torch.arange(2, device='cuda', dtype=fallback.dtype) expected = expected.repeat_interleave(32)[None, :, None].expand_as(output) torch.testing.assert_close(output, expected, rtol=0, atol=0) rows = (2, 16, 96, 384) + + def merge(query): + lse = torch.zeros(query.shape[:2], device='cuda') + counts = torch.ones(query.size(0), device='cuda', dtype=torch.int32) + return manager.combine(query, lse, counts) + + query.fill_(rank) for count in rows: - gather_dcp_query(query[:count], dcp_world_size=2) + torch.testing.assert_close(merge(manager.gather_query(query[:count])), query[:count], rtol=0, atol=0) torch.cuda.synchronize() dist.barrier() graph = torch.cuda.CUDAGraph() @@ -62,8 +71,12 @@ def _run_dcp_query_gather(rank, rendezvous, enabled): with torch.cuda.graph(graph): for step in range(16): query.fill_(rank + step * 8) - gathered = gather_dcp_query(query[:rows[step % len(rows)]], dcp_world_size=2) + gathered = manager.gather_query(query[:rows[step % len(rows)]]) + # A slow reader must finish before any rank reuses the arena. + if rank == step % 2: + torch.cuda._sleep(20000) outputs.append(gathered.clone()) + merge(gathered) for _ in range(3): graph.replay() torch.cuda.synchronize() @@ -73,11 +86,52 @@ def _run_dcp_query_gather(rank, rendezvous, enabled): expected = expected.repeat_interleave(32)[None, :, None].expand_as(output) torch.testing.assert_close(output, expected, rtol=0, atol=0) _check_lm_head_lifecycle(rank, all(enabled)) + _check_dcp_candidate_gather(rank, ctx.dcp_group) finally: ctx.close() dist.destroy_process_group() +def _check_dcp_candidate_gather(rank, group): + from lmdeploy.pytorch.backends.cuda.nsa import TritonNSAIndexFP8Impl + from lmdeploy.pytorch.kernels.cuda.sparse_index_dcp_topk import pack_dcp_topk_candidates, sparse_dcp_global_topk + + k = 512 + impl = TritonNSAIndexFP8Impl(k, softmax_scale=1.0, block_size=128, fill=-1) + scores = torch.randn(16, k * 2, device='cuda') + indices = torch.arange(k, device='cuda', dtype=torch.int32).expand(16, -1).clone() + + def reference(): + packed = pack_dcp_topk_candidates(scores, indices, dcp_world_rank=(2, rank)) + gathered = packed.new_empty(32, k, 2) + torch.distributed.all_gather_into_tensor(gathered, packed, group=group.gpu_group) + gathered = gathered.view(2, 16, k, 2) + return sparse_dcp_global_topk(gathered, k=k) + + impl._merge_dcp_topk(scores, indices) + graph = torch.cuda.CUDAGraph() + outputs = [] + with torch.cuda.graph(graph): + for step in range(4): + if rank == step % 2: + torch.cuda._sleep(20000) + outputs.append(impl._merge_dcp_topk(scores, indices)) + for case in range(3): + scores.normal_() + if case == 1: + scores.zero_() + if case == 2: + indices[:, k // 2:] = -1 + indices[0] = -1 + expected = reference() + torch.testing.assert_close(impl._merge_dcp_topk(scores, indices), expected, rtol=0, atol=0) + graph.replay() + torch.cuda.synchronize() + for output in outputs: + torch.testing.assert_close(output, expected, rtol=0, atol=0) + graph.reset() + + def _run_auto_all_reduce(rank, rendezvous): from datetime import timedelta @@ -194,7 +248,7 @@ def comm_env(monkeypatch): def test_gather_registration_and_native_fallback(monkeypatch): - from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_query, init_dcp_query_gather + from lmdeploy.pytorch.backends.cuda.attention.cp import get_dcp_manager from lmdeploy.pytorch.distributed import DistContext, DistGroup, get_dist_manager module = base_communicator_module @@ -208,13 +262,14 @@ def gather(output, input, group): group = DistGroup(gpu_group='group', communicator=communicator) context = DistContext(dcp_group=group) with get_dist_manager().context(context): - init_dcp_query_gather(2, 4) - assert group.query_gather_workspace is None + manager = get_dcp_manager() + manager.prepare_query_gather(2, 4) + assert manager._query_workspace is None logits = communicator.create_all_gather_workspace(6, device=torch.device('cpu'), dtype=torch.float32) assert logits is None query = torch.arange(48).view(3, 2, 8)[..., ::2] with get_dist_manager().context(context): - torch.testing.assert_close(gather_dcp_query(query, dcp_world_size=2), torch.cat((query, query + 10), dim=1)) + torch.testing.assert_close(manager.gather_query(query), torch.cat((query, query + 10), dim=1)) input = torch.ones(3, 3) torch.testing.assert_close(communicator.all_gather(input, workspace=logits), torch.cat((input, input + 10), dim=-1)) @@ -242,20 +297,22 @@ def test_gather_preparation_and_weight_transition(monkeypatch): def test_cuda_gather_dispatch_and_workspace_ownership(monkeypatch, comm_env): - from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_query, init_dcp_query_gather + from lmdeploy.pytorch.backends.cuda.attention.cp import get_dcp_manager from lmdeploy.pytorch.backends.cuda.comm import symm_mem_allgather from lmdeploy.pytorch.distributed import DistContext, DistGroup, get_dist_manager - # The group shares one query workspace; logits workspaces remain caller-owned. + # The context shares a DCP manager; logits workspaces remain caller-owned. factory = Mock(side_effect=lambda *args, **kwargs: Mock()) monkeypatch.setattr(symm_mem_allgather, 'SymmetricMemoryAllGather', factory) communicator = communicator_module.CudaCommunicator('cpu', 'gpu', group_name='dcp') group = DistGroup(gpu_group='gpu', communicator=communicator) context = DistContext(dcp_group=group) with get_dist_manager().context(context): - init_dcp_query_gather(2, 4) - init_dcp_query_gather(2, 4) - query_workspace = group.query_gather_workspace + manager = get_dcp_manager() + manager.prepare_query_gather(2, 4) + manager = get_dcp_manager() + manager.prepare_query_gather(2, 4) + query_workspace = manager._query_workspace logits_workspace = communicator.create_all_gather_workspace(6, device=torch.device('cpu'), dtype=torch.float32) assert factory.call_count == 2 query_workspace.prepare.assert_called_once() @@ -265,7 +322,7 @@ def test_cuda_gather_dispatch_and_workspace_ownership(monkeypatch, comm_env): gathered_query = torch.cat((query, query + 10), dim=1) query_workspace.all_gather.return_value = gathered_query.flatten(1) with get_dist_manager().context(context): - torch.testing.assert_close(gather_dcp_query(query, dcp_world_size=2), gathered_query) + torch.testing.assert_close(manager.gather_query(query), gathered_query) assert query_workspace.all_gather.call_args.kwargs == {'copy_output': False} native = Mock(side_effect=lambda output, input, group: output.copy_(torch.cat((input, input + 10)))) @@ -280,10 +337,16 @@ def test_cuda_gather_dispatch_and_workspace_ownership(monkeypatch, comm_env): torch.cat((input, input + 10), dim=-1)) native.assert_called_once() - group.close() - group.close() + assert context.dcp_manager is manager + monkeypatch.setattr(torch.distributed, 'is_initialized', lambda: True) + group_close = Mock(wraps=group.close) + monkeypatch.setattr(group, 'close', group_close) + query_workspace.close.side_effect = lambda: group_close.assert_not_called() + context.close() + context.close() query_workspace.close.assert_called_once() - assert group.query_gather_workspace is None + assert context.dcp_manager is None + assert manager._query_workspace is None logits_workspace.close.assert_not_called() From cdc3e4bfb0cee9f5a746a6d64e13552523aca413 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 24 Sep 2026 04:24:43 +0000 Subject: [PATCH 07/15] perf: accelerate DCP top-k candidate gathering with symmetric memory --- lmdeploy/pytorch/backends/communicator.py | 15 ++++--- .../pytorch/backends/cuda/attention/cp.py | 20 +++++++-- .../backends/cuda/comm/communicator.py | 12 +++--- .../backends/cuda/comm/symm_mem_allgather.py | 35 +++++++++------ lmdeploy/pytorch/backends/cuda/nsa.py | 1 + .../kernels/cuda/symm_mem_allgather.py | 27 ++++++++---- tests/pytorch/test_cuda_communicator.py | 43 +++++++++++++++---- 7 files changed, 110 insertions(+), 43 deletions(-) diff --git a/lmdeploy/pytorch/backends/communicator.py b/lmdeploy/pytorch/backends/communicator.py index a02791ef07..8c73dc9f57 100644 --- a/lmdeploy/pytorch/backends/communicator.py +++ b/lmdeploy/pytorch/backends/communicator.py @@ -19,23 +19,28 @@ def all_reduce_(self, input: torch.Tensor): """All-reduce ``input`` in place.""" dist.all_reduce(input, group=self.device_group) - def all_gather(self, input: torch.Tensor, *, workspace=None, copy_output: bool = True) -> torch.Tensor: - """Gather 2D inputs along the last dimension on every group rank. + def all_gather(self, input: torch.Tensor, *, dim: int = -1, + workspace=None, copy_output: bool = True) -> torch.Tensor: + """Gather 2D inputs along dimension 0 or -1 on every group rank. With copy_output=False, optimized output may borrow the workspace until its next use. Consume it on the same stream. Native gathering always returns an owned output for multi-rank groups. """ + if input.dim() != 2 or dim not in (0, -1): + raise ValueError('all_gather requires a 2D input and dim=0 or dim=-1') world_size = dist.get_world_size(self.device_group) if world_size == 1: return input input = input.contiguous() gathered = input.new_empty(world_size * input.size(0), input.size(1)) dist.all_gather_into_tensor(gathered, input, group=self.device_group) + if dim == 0: + return gathered return gathered.view(world_size, *input.shape).transpose(0, 1).reshape(input.size(0), -1) - def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype): - """Create and prepare an optional workspace and return it to its - owner.""" + def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype, + *, dim: int = -1): + """Create a caller-owned workspace.""" return None def close(self): diff --git a/lmdeploy/pytorch/backends/cuda/attention/cp.py b/lmdeploy/pytorch/backends/cuda/attention/cp.py index 14669ef1ac..1872bc714f 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/cp.py +++ b/lmdeploy/pytorch/backends/cuda/attention/cp.py @@ -40,6 +40,7 @@ def __init__(self, group): self.world_size = torch.distributed.get_world_size(self.group) self.rank = group.rank self._query_workspace = None + self._candidate_workspace = None def prepare_query_gather(self, num_heads: int, head_size: int): """Prepare one query arena shared across attention layers.""" @@ -48,6 +49,12 @@ def prepare_query_gather(self, num_heads: int, head_size: int): self._query_workspace = self.communicator.create_all_gather_workspace( width, device=torch.device('cuda'), dtype=torch.bfloat16) + def prepare_candidate_gather(self, topk: int): + """Prepare one candidate arena shared across indexer layers.""" + if self._candidate_workspace is None: + self._candidate_workspace = self.communicator.create_all_gather_workspace( + topk * 2, device=torch.device('cuda'), dtype=torch.float32, dim=0) + def gather_query(self, query: torch.Tensor) -> torch.Tensor: """Gather query heads; consume and combine on the same stream before reuse.""" @@ -58,11 +65,13 @@ def gather_query(self, query: torch.Tensor) -> torch.Tensor: return output.view(query.size(0), -1, query.size(2)) def gather_candidates(self, packed: torch.Tensor) -> torch.Tensor: - """Gather [rows, topk, score/ID] payloads in rank-major order.""" - from lmdeploy.pytorch.distributed import all_gather_into_tensor + """Gather [rows, topk, score/ID] payloads without changing their bits. - output = packed.new_empty(self.world_size * packed.size(0), *packed.shape[1:]) - all_gather_into_tensor(output, packed, group=self.group) + The entry barrier protects the previous top-k consumer before arena reuse. Consume the borrowed output on the + same stream before the next candidate gather. + """ + output = self.communicator.all_gather( + packed.flatten(1), dim=0, workspace=self._candidate_workspace, copy_output=False) return output.view(self.world_size, *packed.shape) def combine(self, local_output: torch.Tensor, local_lse: torch.Tensor, @@ -102,6 +111,9 @@ def close(self): if self._query_workspace is not None: self._query_workspace.close() self._query_workspace = None + if self._candidate_workspace is not None: + self._candidate_workspace.close() + self._candidate_workspace = None def get_dcp_manager() -> DCPManager: diff --git a/lmdeploy/pytorch/backends/cuda/comm/communicator.py b/lmdeploy/pytorch/backends/cuda/comm/communicator.py index 79128b4aee..2165a106f4 100644 --- a/lmdeploy/pytorch/backends/cuda/comm/communicator.py +++ b/lmdeploy/pytorch/backends/cuda/comm/communicator.py @@ -46,21 +46,23 @@ def all_reduce_(self, input: torch.Tensor): return super().all_reduce_(input) - def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype): + def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype, + *, dim: int = -1): from .symm_mem_allgather import SymmetricMemoryAllGather workspace = SymmetricMemoryAllGather( self.device_group, dist.get_rank(self.device_group), gathered_width, - device=device, dtype=dtype, capacity_bytes=64 * 1024 * 1024) + device=device, dtype=dtype, capacity_bytes=64 * 1024 * 1024, dim=dim) workspace.prepare() return workspace - def all_gather(self, input: torch.Tensor, *, workspace=None, copy_output: bool = True) -> torch.Tensor: + def all_gather(self, input: torch.Tensor, *, dim: int = -1, + workspace=None, copy_output: bool = True) -> torch.Tensor: if workspace is not None: - output = workspace.all_gather(input, copy_output=copy_output) + output = workspace.all_gather(input, dim=dim, copy_output=copy_output) if output is not None: return output - return super().all_gather(input) + return super().all_gather(input, dim=dim) def close(self): if self._all_reduce_provider is not None: diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py index cc5a8f3826..9f7333435a 100644 --- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py +++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py @@ -1,5 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. -"""Symmetric-memory workspace for last-dimension all-gather. +"""Symmetric-memory workspace for 2D all-gather. The optional Triton kernels are imported only during collective preparation. Their SGLang source reference is in kernels/cuda/symm_mem_allgather.py. @@ -27,13 +27,15 @@ class SymmetricMemoryAllGather: _WORLD_SIZES = {2, 4, 8} def __init__(self, group: dist.ProcessGroup, rank: int, gathered_width: int, - *, device: torch.device, dtype: torch.dtype, capacity_bytes: int): + *, device: torch.device, dtype: torch.dtype, capacity_bytes: int, + dim: int = -1): self.device = torch.device(device) self.dtype = dtype self._prepared = False self.group = group self._rank = rank self._gathered_width = gathered_width + self._dim = dim self._capacity_bytes = capacity_bytes self._state = None self._kernels = None @@ -89,22 +91,26 @@ def _prepare(self, device: torch.device, dtype: torch.dtype) -> bool: return False # Even disabled ranks participate: mixed flags/dtypes must not leave # peers entering an optional rendezvous on their own. - if not self._agree(dtype == torch.bfloat16, device): + supported = dtype == torch.bfloat16 or (self._dim == 0 and dtype == torch.float32) + if not self._agree(supported, device): return False world_size = dist.get_world_size(self.group) - width = self._gathered_width + # The copy kernel addresses BF16 words, but never converts payloads. + width = self._gathered_width * dtype.itemsize // torch.bfloat16.itemsize + alignment = 8 if self._dim == 0 else world_size * 8 try: capability_ok = torch.cuda.get_device_capability(device) >= (9, 0) except (RuntimeError, AssertionError): capability_ok = False valid = (world_size in self._WORLD_SIZES and self._rank == dist.get_rank(self.group) - and width > 0 and width % (world_size * 8) == 0 + and self._dim in (0, -1) + and width > 0 and width % alignment == 0 and self._capacity_bytes >= width * torch.bfloat16.itemsize and capability_ok) if not self._agree(valid, device): return self._disabled('unsupported group size, dtype, width, capacity or device') - if not self._same_config((width, self._capacity_bytes), device): + if not self._same_config((width, dtype.itemsize, self._dim, self._capacity_bytes), device): return self._disabled('inconsistent group arena configuration') kernels = None @@ -161,7 +167,8 @@ def reset(self, device: torch.device, dtype: torch.dtype): self._prepared = False self.prepare() - def all_gather(self, input: torch.Tensor, *, copy_output: bool = True) -> torch.Tensor | None: + def all_gather(self, input: torch.Tensor, *, dim: int = -1, + copy_output: bool = True) -> torch.Tensor | None: """Return gathered shards, or None when the input is unsupported. With copy_output=False, the output borrows the workspace until its next use. @@ -169,21 +176,25 @@ def all_gather(self, input: torch.Tensor, *, copy_output: bool = True) -> torch. state = self._state if state is None: return None - if not (input.dim() == 2 and input.dtype == torch.bfloat16 and input.is_contiguous() - and input.shape[-1] * state.world_size == state.hidden_dim): + width = self._gathered_width if dim == 0 else self._gathered_width // state.world_size + if not (dim == self._dim and input.dim() == 2 and input.dtype == self.dtype + and input.is_contiguous() and input.shape[-1] == width): # Shape/dtype/layout contracts must match across ranks, as for NCCL. return None if input.device != state.device or input.data_ptr() % 16 != 0: raise RuntimeError('multimem all-gather device or alignment changed after group-wide admission') rows = input.shape[0] - if not self._MIN_TOKENS <= rows <= state.max_token_num: + min_rows = 1 if dim == 0 else self._MIN_TOKENS + max_rows = state.max_token_num // state.world_size if dim == 0 else state.max_token_num + if not min_rows <= rows <= max_rows: return None if rows not in self._graph_ready_shapes and torch.cuda.is_current_stream_capturing(): return None output = self._kernels.all_gather_inner( - state, input, tp_hidden_dim=self._gathered_width, copy_output=copy_output, _validated=True) + state, input.view(torch.bfloat16), tp_hidden_dim=state.hidden_dim, dim=dim, + copy_output=copy_output, _validated=True) self._graph_ready_shapes.add(rows) - return output + return output.view(self.dtype) def close(self) -> None: """Release only after all consumers/graphs have been retired.""" diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index 6c708dd89c..5ff734ebb8 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -367,6 +367,7 @@ def __init__(self, topk: int, softmax_scale: float, block_size: int, raise RuntimeError('DCP DSA scoring requires a compatible DeepGEMM installation.') from .attention.cp import get_dcp_manager self.dcp_manager = get_dcp_manager() + self.dcp_manager.prepare_candidate_gather(topk) self._step_meta_group: int | None = None self._piecewise_forward: Callable[..., Tensor] | None = None self._piecewise_forward_fused: Callable[..., Tensor] | None = None diff --git a/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py b/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py index 06bd0d2e7e..8d8af29df3 100644 --- a/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py +++ b/lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py @@ -453,11 +453,12 @@ def all_gather_inner( skip_entry_sync: bool = False, copy_output: bool = True, *, + dim: int = -1, _validated: bool = False, ) -> torch.Tensor: - """Gather ``[T, H/TP]`` shards into ``[T, H]`` along the hidden dim. + """Gather 2D BF16 words along dimension 0 or -1 without arithmetic. - ``tp_hidden_dim`` is the gathered width ``H``. Returns a clone when ``copy_output``, + ``tp_hidden_dim`` is the output row width. Returns a clone when ``copy_output``, else a view into the symmetric buffer (valid until the next collective). ``_validated`` is reserved for the admitted provider, whose admission check already enforces the immutable dtype/layout/width @@ -471,10 +472,9 @@ def all_gather_inner( f"hidden_states.data_ptr()={hex(hidden_states.data_ptr())} must be " f"16-byte aligned for 128-bit multimem.st" ) - assert ( - tp_hidden_dim % world_size == 0 - ), f"tp_hidden_dim={tp_hidden_dim} must be divisible by world_size={world_size}" - local_hidden = tp_hidden_dim // world_size + assert dim in (0, -1) + assert dim == 0 or tp_hidden_dim % world_size == 0 + local_hidden = tp_hidden_dim if dim == 0 else tp_hidden_dim // world_size total_tokens, in_hidden = hidden_states.shape if not _validated: assert local_hidden % _NUMEL_PER_THREAD == 0, ( @@ -489,14 +489,20 @@ def all_gather_inner( in_hidden == local_hidden ), f"input hidden ({in_hidden}) != this rank's shard ({local_hidden})" assert ( - total_tokens <= state.max_token_num + total_tokens * (world_size if dim == 0 else 1) <= state.max_token_num ), f"total_tokens={total_tokens} exceeds max_token_num={state.max_token_num}" + logical_rows = total_tokens + if dim == 0: + # Copy each rank's entire payload as one contiguous shard. Preserve the + # logical row count for launch tuning, without a rank/token transpose. + local_hidden *= total_tokens + total_tokens = 1 hidden_offset = local_hidden * state.rank_in_group symm_mem_hdl = state.symm_mem_hdl num_blocks, block_size, num_warps, numel_per_thread = _launch_config( total_tokens * local_hidden, - total_tokens=total_tokens, + total_tokens=logical_rows, world_size=world_size, ) split_barrier = _use_single_barrier(num_blocks, world_size) @@ -527,7 +533,10 @@ def all_gather_inner( # The completion barrier can therefore be a single CTA and still # publish the whole grid before the output view is returned. barrier_inner(state, slot=1, release=True) - output = state.comm_buff[:total_tokens, :tp_hidden_dim] + if dim == 0: + output = state.comm_buff.view(-1)[:world_size * local_hidden].view(world_size * logical_rows, tp_hidden_dim) + else: + output = state.comm_buff[:total_tokens, :tp_hidden_dim] return output.clone() if copy_output else output diff --git a/tests/pytorch/test_cuda_communicator.py b/tests/pytorch/test_cuda_communicator.py index 158c24e3b6..adcecd4c10 100644 --- a/tests/pytorch/test_cuda_communicator.py +++ b/tests/pytorch/test_cuda_communicator.py @@ -86,26 +86,43 @@ def merge(query): expected = expected.repeat_interleave(32)[None, :, None].expand_as(output) torch.testing.assert_close(output, expected, rtol=0, atol=0) _check_lm_head_lifecycle(rank, all(enabled)) - _check_dcp_candidate_gather(rank, ctx.dcp_group) + _check_dcp_candidate_gather(rank, ctx.dcp_group, all(enabled)) finally: ctx.close() dist.destroy_process_group() -def _check_dcp_candidate_gather(rank, group): +def _check_dcp_candidate_gather(rank, group, direct): from lmdeploy.pytorch.backends.cuda.nsa import TritonNSAIndexFP8Impl from lmdeploy.pytorch.kernels.cuda.sparse_index_dcp_topk import pack_dcp_topk_candidates, sparse_dcp_global_topk k = 512 impl = TritonNSAIndexFP8Impl(k, softmax_scale=1.0, block_size=128, fill=-1) + workspace = impl.dcp_manager._candidate_workspace + assert (workspace is not None and workspace.is_available()) == direct + TritonNSAIndexFP8Impl(k, softmax_scale=1.0, block_size=128, fill=-1) + assert impl.dcp_manager._candidate_workspace is workspace + native = base_communicator_module.DeviceCommunicator(group.gpu_group) + # IDs are bitcast, not converted to FP32; include large IDs and NaN encodings. + bits = (torch.arange(16 * k * 2, device='cuda', dtype=torch.int32) * 1234567 + rank).view(16, k * 2) + payload = bits.view(torch.float32) + for value in (payload, payload[:1], payload[:, ::2], payload.double()): + actual = group.communicator.all_gather(value, dim=0, workspace=workspace) + expected = native.all_gather(value, dim=0) + torch.testing.assert_close(actual.view(torch.int32), expected.view(torch.int32), rtol=0, atol=0) + if direct: + # Capacity rejection must use NCCL consistently, including a captured + # shape that was not warmed before capture. + oversized = payload[:1].expand(workspace._state.max_token_num // 2 + 1, -1).contiguous() + torch.testing.assert_close( + group.communicator.all_gather(oversized, dim=0, workspace=workspace).view(torch.int32), + native.all_gather(oversized, dim=0).view(torch.int32), rtol=0, atol=0) scores = torch.randn(16, k * 2, device='cuda') indices = torch.arange(k, device='cuda', dtype=torch.int32).expand(16, -1).clone() def reference(): packed = pack_dcp_topk_candidates(scores, indices, dcp_world_rank=(2, rank)) - gathered = packed.new_empty(32, k, 2) - torch.distributed.all_gather_into_tensor(gathered, packed, group=group.gpu_group) - gathered = gathered.view(2, 16, k, 2) + gathered = native.all_gather(packed.flatten(1), dim=0).view(2, 16, k, 2) return sparse_dcp_global_topk(gathered, k=k) impl._merge_dcp_topk(scores, indices) @@ -116,6 +133,7 @@ def reference(): if rank == step % 2: torch.cuda._sleep(20000) outputs.append(impl._merge_dcp_topk(scores, indices)) + unwarmed = group.communicator.all_gather(payload[:3], dim=0, workspace=workspace) for case in range(3): scores.normal_() if case == 1: @@ -129,6 +147,8 @@ def reference(): torch.cuda.synchronize() for output in outputs: torch.testing.assert_close(output, expected, rtol=0, atol=0) + torch.testing.assert_close(unwarmed.view(torch.int32), + native.all_gather(payload[:3], dim=0).view(torch.int32), rtol=0, atol=0) graph.reset() @@ -272,6 +292,7 @@ def gather(output, input, group): torch.testing.assert_close(manager.gather_query(query), torch.cat((query, query + 10), dim=1)) input = torch.ones(3, 3) torch.testing.assert_close(communicator.all_gather(input, workspace=logits), torch.cat((input, input + 10), dim=-1)) + torch.testing.assert_close(communicator.all_gather(input, dim=0), torch.cat((input, input + 10), dim=0)) def test_gather_preparation_and_weight_transition(monkeypatch): @@ -323,30 +344,36 @@ def test_cuda_gather_dispatch_and_workspace_ownership(monkeypatch, comm_env): query_workspace.all_gather.return_value = gathered_query.flatten(1) with get_dist_manager().context(context): torch.testing.assert_close(manager.gather_query(query), gathered_query) - assert query_workspace.all_gather.call_args.kwargs == {'copy_output': False} + assert query_workspace.all_gather.call_args.kwargs == {'dim': -1, 'copy_output': False} native = Mock(side_effect=lambda output, input, group: output.copy_(torch.cat((input, input + 10)))) monkeypatch.setattr(base_communicator_module.dist, 'all_gather_into_tensor', native) input = torch.ones(3, 3) logits_workspace.all_gather.return_value = torch.ones(3, 6) torch.testing.assert_close(communicator.all_gather(input, workspace=logits_workspace), torch.ones(3, 6)) - logits_workspace.all_gather.assert_called_once_with(input, copy_output=True) + logits_workspace.all_gather.assert_called_once_with(input, dim=-1, copy_output=True) native.assert_not_called() logits_workspace.all_gather.return_value = None torch.testing.assert_close(communicator.all_gather(input, workspace=logits_workspace), torch.cat((input, input + 10), dim=-1)) native.assert_called_once() + manager.prepare_candidate_gather(512) + candidate_workspace = manager._candidate_workspace + manager.prepare_candidate_gather(512) + assert manager._candidate_workspace is candidate_workspace assert context.dcp_manager is manager monkeypatch.setattr(torch.distributed, 'is_initialized', lambda: True) group_close = Mock(wraps=group.close) monkeypatch.setattr(group, 'close', group_close) query_workspace.close.side_effect = lambda: group_close.assert_not_called() + candidate_workspace.close.side_effect = lambda: group_close.assert_not_called() context.close() context.close() query_workspace.close.assert_called_once() + candidate_workspace.close.assert_called_once() assert context.dcp_manager is None - assert manager._query_workspace is None + assert manager._query_workspace is manager._candidate_workspace is None logits_workspace.close.assert_not_called() From 4fc799aa7d5c332fb3fc3035188dbf8b872c1953 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 24 Sep 2026 04:25:52 +0000 Subject: [PATCH 08/15] perf: skip redundant DCP query-gather entry synchronization --- lmdeploy/pytorch/backends/communicator.py | 7 +++++-- lmdeploy/pytorch/backends/cuda/attention/cp.py | 4 +++- lmdeploy/pytorch/backends/cuda/comm/communicator.py | 4 ++-- .../pytorch/backends/cuda/comm/symm_mem_allgather.py | 11 +++++++---- 4 files changed, 17 insertions(+), 9 deletions(-) diff --git a/lmdeploy/pytorch/backends/communicator.py b/lmdeploy/pytorch/backends/communicator.py index 8c73dc9f57..9223ffeb80 100644 --- a/lmdeploy/pytorch/backends/communicator.py +++ b/lmdeploy/pytorch/backends/communicator.py @@ -39,8 +39,11 @@ def all_gather(self, input: torch.Tensor, *, dim: int = -1, return gathered.view(world_size, *input.shape).transpose(0, 1).reshape(input.size(0), -1) def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype, - *, dim: int = -1): - """Create a caller-owned workspace.""" + *, dim: int = -1, reuse_sync: bool = True): + """Create a caller-owned workspace. + + Disable reuse_sync only when a subsequent same-group collective orders all consumers before the next gather. + """ return None def close(self): diff --git a/lmdeploy/pytorch/backends/cuda/attention/cp.py b/lmdeploy/pytorch/backends/cuda/attention/cp.py index 1872bc714f..12bca74d9e 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/cp.py +++ b/lmdeploy/pytorch/backends/cuda/attention/cp.py @@ -46,8 +46,10 @@ def prepare_query_gather(self, num_heads: int, head_size: int): """Prepare one query arena shared across attention layers.""" if self._query_workspace is None: width = num_heads * head_size * self.world_size + # combine() orders all ranks after attention consumes these queries, + # so the next query gather does not need another entry barrier. self._query_workspace = self.communicator.create_all_gather_workspace( - width, device=torch.device('cuda'), dtype=torch.bfloat16) + width, device=torch.device('cuda'), dtype=torch.bfloat16, reuse_sync=False) def prepare_candidate_gather(self, topk: int): """Prepare one candidate arena shared across indexer layers.""" diff --git a/lmdeploy/pytorch/backends/cuda/comm/communicator.py b/lmdeploy/pytorch/backends/cuda/comm/communicator.py index 2165a106f4..9408f67f10 100644 --- a/lmdeploy/pytorch/backends/cuda/comm/communicator.py +++ b/lmdeploy/pytorch/backends/cuda/comm/communicator.py @@ -47,12 +47,12 @@ def all_reduce_(self, input: torch.Tensor): super().all_reduce_(input) def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype, - *, dim: int = -1): + *, dim: int = -1, reuse_sync: bool = True): from .symm_mem_allgather import SymmetricMemoryAllGather workspace = SymmetricMemoryAllGather( self.device_group, dist.get_rank(self.device_group), gathered_width, - device=device, dtype=dtype, capacity_bytes=64 * 1024 * 1024, dim=dim) + device=device, dtype=dtype, capacity_bytes=64 * 1024 * 1024, dim=dim, reuse_sync=reuse_sync) workspace.prepare() return workspace diff --git a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py index 9f7333435a..87e3c966b1 100644 --- a/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py +++ b/lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py @@ -19,7 +19,8 @@ class SymmetricMemoryAllGather: Construct on every group rank after process-group setup. Only preparation performs allocation and consensus; forward never initializes or rebuilds. Re-prepare after coordinated, quiescent device/dtype changes. Runtime rows must match - across ranks, as required by NCCL all-gather too. + across ranks, as required by NCCL all-gather too. With reuse_sync=False, callers must order all ranks' consumers + through a subsequent same-group collective before reusing the arena. """ # Keep the conservative single-token NCCL policy internal to the provider. @@ -28,7 +29,7 @@ class SymmetricMemoryAllGather: def __init__(self, group: dist.ProcessGroup, rank: int, gathered_width: int, *, device: torch.device, dtype: torch.dtype, capacity_bytes: int, - dim: int = -1): + dim: int = -1, reuse_sync: bool = True): self.device = torch.device(device) self.dtype = dtype self._prepared = False @@ -37,6 +38,7 @@ def __init__(self, group: dist.ProcessGroup, rank: int, gathered_width: int, self._gathered_width = gathered_width self._dim = dim self._capacity_bytes = capacity_bytes + self._reuse_sync = reuse_sync self._state = None self._kernels = None self._graph_ready_shapes = set() @@ -110,7 +112,8 @@ def _prepare(self, device: torch.device, dtype: torch.dtype) -> bool: and capability_ok) if not self._agree(valid, device): return self._disabled('unsupported group size, dtype, width, capacity or device') - if not self._same_config((width, dtype.itemsize, self._dim, self._capacity_bytes), device): + if not self._same_config((width, dtype.itemsize, self._dim, + self._capacity_bytes, int(self._reuse_sync)), device): return self._disabled('inconsistent group arena configuration') kernels = None @@ -192,7 +195,7 @@ def all_gather(self, input: torch.Tensor, *, dim: int = -1, return None output = self._kernels.all_gather_inner( state, input.view(torch.bfloat16), tp_hidden_dim=state.hidden_dim, dim=dim, - copy_output=copy_output, _validated=True) + skip_entry_sync=not self._reuse_sync, copy_output=copy_output, _validated=True) self._graph_ready_shapes.add(rows) return output.view(self.dtype) From f460e044d6a9bcd087aae4c5c2e6032de6dea1bc Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 24 Sep 2026 05:53:24 +0000 Subject: [PATCH 09/15] test: decouple DCP candidate gathering from DeepGEMM --- tests/pytorch/test_cuda_communicator.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_cuda_communicator.py b/tests/pytorch/test_cuda_communicator.py index adcecd4c10..53288083bd 100644 --- a/tests/pytorch/test_cuda_communicator.py +++ b/tests/pytorch/test_cuda_communicator.py @@ -93,14 +93,20 @@ def merge(query): def _check_dcp_candidate_gather(rank, group, direct): + from lmdeploy.pytorch.backends.cuda.attention.cp import get_dcp_manager from lmdeploy.pytorch.backends.cuda.nsa import TritonNSAIndexFP8Impl from lmdeploy.pytorch.kernels.cuda.sparse_index_dcp_topk import pack_dcp_topk_candidates, sparse_dcp_global_topk k = 512 - impl = TritonNSAIndexFP8Impl(k, softmax_scale=1.0, block_size=128, fill=-1) + # Candidate merging consumes existing scores and does not require DeepGEMM. + impl = object.__new__(TritonNSAIndexFP8Impl) + impl.topk, impl.fill = k, -1 + impl.dcp_world_size, impl.dcp_rank = 2, rank + impl.dcp_manager = get_dcp_manager() + impl.dcp_manager.prepare_candidate_gather(k) workspace = impl.dcp_manager._candidate_workspace assert (workspace is not None and workspace.is_available()) == direct - TritonNSAIndexFP8Impl(k, softmax_scale=1.0, block_size=128, fill=-1) + get_dcp_manager().prepare_candidate_gather(k) assert impl.dcp_manager._candidate_workspace is workspace native = base_communicator_module.DeviceCommunicator(group.gpu_group) # IDs are bitcast, not converted to FP32; include large IDs and NaN encodings. From 9161e644b0e0ba3bb7811b6ee9fe9d3620bf51b6 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 24 Sep 2026 11:44:25 +0000 Subject: [PATCH 10/15] perf: accelerate DCP LSE gathering with symmetric memory --- lmdeploy/pytorch/backends/communicator.py | 2 +- .../pytorch/backends/cuda/attention/cp.py | 23 +++++--- .../pytorch/backends/cuda/attention/mla.py | 2 +- .../backends/cuda/comm/communicator.py | 4 +- tests/pytorch/test_cuda_communicator.py | 56 ++++++++++++++++--- 5 files changed, 69 insertions(+), 18 deletions(-) diff --git a/lmdeploy/pytorch/backends/communicator.py b/lmdeploy/pytorch/backends/communicator.py index 9223ffeb80..c17e01d451 100644 --- a/lmdeploy/pytorch/backends/communicator.py +++ b/lmdeploy/pytorch/backends/communicator.py @@ -39,7 +39,7 @@ def all_gather(self, input: torch.Tensor, *, dim: int = -1, return gathered.view(world_size, *input.shape).transpose(0, 1).reshape(input.size(0), -1) def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype, - *, dim: int = -1, reuse_sync: bool = True): + *, dim: int = -1, reuse_sync: bool = True, capacity_bytes: int = 64 * 1024 * 1024): """Create a caller-owned workspace. Disable reuse_sync only when a subsequent same-group collective orders all consumers before the next gather. diff --git a/lmdeploy/pytorch/backends/cuda/attention/cp.py b/lmdeploy/pytorch/backends/cuda/attention/cp.py index 12bca74d9e..435c278a56 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/cp.py +++ b/lmdeploy/pytorch/backends/cuda/attention/cp.py @@ -40,16 +40,23 @@ def __init__(self, group): self.world_size = torch.distributed.get_world_size(self.group) self.rank = group.rank self._query_workspace = None + self._lse_workspace = None self._candidate_workspace = None - def prepare_query_gather(self, num_heads: int, head_size: int): - """Prepare one query arena shared across attention layers.""" + def prepare_attention(self, num_heads: int, head_size: int): + """Prepare query and LSE arenas shared across attention layers.""" if self._query_workspace is None: width = num_heads * head_size * self.world_size # combine() orders all ranks after attention consumes these queries, # so the next query gather does not need another entry barrier. self._query_workspace = self.communicator.create_all_gather_workspace( width, device=torch.device('cuda'), dtype=torch.bfloat16, reuse_sync=False) + if self._lse_workspace is None: + # Output reduce-scatter orders LSE readers before the next reuse. + # LSE is small; larger inputs fall back to NCCL beyond this budget. + self._lse_workspace = self.communicator.create_all_gather_workspace( + num_heads * self.world_size, device=torch.device('cuda'), dtype=torch.float32, + dim=0, reuse_sync=False, capacity_bytes=1024 * 1024) def prepare_candidate_gather(self, topk: int): """Prepare one candidate arena shared across indexer layers.""" @@ -86,13 +93,12 @@ def combine(self, local_output: torch.Tensor, local_lse: torch.Tensor, """ if self.world_size == 1: return local_output - from lmdeploy.pytorch.distributed import all_gather_into_tensor, reduce_scatter_tensor + from lmdeploy.pytorch.distributed import reduce_scatter_tensor from lmdeploy.pytorch.kernels.cuda.dcp import correct_dcp_attention_output, sanitize_dcp_lse local_lse = sanitize_dcp_lse(local_lse, valid_counts) - gathered_lse = local_lse.new_empty( - self.world_size * local_lse.size(0), local_lse.size(1)) - all_gather_into_tensor(gathered_lse, local_lse, group=self.group) + gathered_lse = self.communicator.all_gather( + local_lse, dim=0, workspace=self._lse_workspace, copy_output=False) gathered_lse = gathered_lse.view(self.world_size, *local_lse.shape) contribution = correct_dcp_attention_output( local_output, gathered_lse, dcp_rank=self.rank) @@ -113,6 +119,9 @@ def close(self): if self._query_workspace is not None: self._query_workspace.close() self._query_workspace = None + if self._lse_workspace is not None: + self._lse_workspace.close() + self._lse_workspace = None if self._candidate_workspace is not None: self._candidate_workspace.close() self._candidate_workspace = None @@ -142,7 +151,7 @@ def __init__(self, num_heads: int, head_size: int, *, use_fa3: bool = False, **k from lmdeploy.pytorch.third_party.flash_attn_interface import flash_attn_varlen_func self._fa3_prefill = flash_attn_varlen_func self.dcp_manager = get_dcp_manager() - self.dcp_manager.prepare_query_gather(self.num_heads, self.head_size) + self.dcp_manager.prepare_attention(self.num_heads, self.head_size) def get_step_metadata_provider(self): """Use common DCP lengths without a separate kernel scheduler.""" diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index f87f78d69a..e55659b614 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -256,7 +256,7 @@ def __init__( if self.dcp_world_size > 1: self.dcp_manager = get_dcp_manager() - self.dcp_manager.prepare_query_gather(num_heads, head_size) + self.dcp_manager.prepare_attention(num_heads, head_size) def get_step_metadata_provider(self): """Describe metadata required by this selected implementation.""" diff --git a/lmdeploy/pytorch/backends/cuda/comm/communicator.py b/lmdeploy/pytorch/backends/cuda/comm/communicator.py index 9408f67f10..1e2569267d 100644 --- a/lmdeploy/pytorch/backends/cuda/comm/communicator.py +++ b/lmdeploy/pytorch/backends/cuda/comm/communicator.py @@ -47,12 +47,12 @@ def all_reduce_(self, input: torch.Tensor): super().all_reduce_(input) def create_all_gather_workspace(self, gathered_width: int, device: torch.device, dtype: torch.dtype, - *, dim: int = -1, reuse_sync: bool = True): + *, dim: int = -1, reuse_sync: bool = True, capacity_bytes: int = 64 * 1024 * 1024): from .symm_mem_allgather import SymmetricMemoryAllGather workspace = SymmetricMemoryAllGather( self.device_group, dist.get_rank(self.device_group), gathered_width, - device=device, dtype=dtype, capacity_bytes=64 * 1024 * 1024, dim=dim, reuse_sync=reuse_sync) + device=device, dtype=dtype, capacity_bytes=capacity_bytes, dim=dim, reuse_sync=reuse_sync) workspace.prepare() return workspace diff --git a/tests/pytorch/test_cuda_communicator.py b/tests/pytorch/test_cuda_communicator.py index 53288083bd..50a13ea14b 100644 --- a/tests/pytorch/test_cuda_communicator.py +++ b/tests/pytorch/test_cuda_communicator.py @@ -36,14 +36,14 @@ def _run_dcp_query_gather(rank, rendezvous, enabled): try: with get_dist_manager().context(ctx): manager = get_dcp_manager() - manager.prepare_query_gather(32, 576) + manager.prepare_attention(32, 576) workspace = manager._query_workspace if all(enabled): assert workspace.is_available() else: assert workspace is None manager = get_dcp_manager() - manager.prepare_query_gather(32, 576) + manager.prepare_attention(32, 576) assert manager._query_workspace is workspace query = torch.empty(384, 32, 576, device='cuda', dtype=torch.bfloat16) # Ineligible inputs still return correct results through NCCL. @@ -85,6 +85,7 @@ def merge(query): expected = (torch.arange(2, device='cuda', dtype=query.dtype) + step * 8) expected = expected.repeat_interleave(32)[None, :, None].expand_as(output) torch.testing.assert_close(output, expected, rtol=0, atol=0) + _check_dcp_merge(rank, ctx.dcp_group, manager, all(enabled)) _check_lm_head_lifecycle(rank, all(enabled)) _check_dcp_candidate_gather(rank, ctx.dcp_group, all(enabled)) finally: @@ -92,6 +93,43 @@ def merge(query): dist.destroy_process_group() +def _check_dcp_merge(rank, group, manager, direct): + from lmdeploy.pytorch.backends.cuda.attention.cp import DCPManager + from lmdeploy.pytorch.distributed import DistGroup + + native = DCPManager(DistGroup(rank=rank, gpu_group=group.gpu_group, + communicator=base_communicator_module.DeviceCommunicator(group.gpu_group))) + workspace = manager._lse_workspace + assert (workspace is not None and workspace.is_available()) == direct + for rows in (1, 16, 96): + output = torch.randn(rows, 64, 128, device='cuda', dtype=torch.bfloat16) + lse = torch.randn(rows, 68, device='cuda')[:, :64] + counts = torch.ones(rows, device='cuda', dtype=torch.int32) + counts[-1] = rank + lse[0, 0] = float('inf') + expected = native.combine(output, lse, counts) + for _ in range(3): + actual = manager.combine(output, lse, counts) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + outputs = [] + with torch.cuda.graph(graph): + for step in range(4): + if rank == step % 2: + torch.cuda._sleep(20000) + outputs.append(manager.combine(output, lse, counts)) + for _ in range(3): + output.normal_() + lse.normal_() + lse[0, 0] = float('nan') + expected = native.combine(output, lse, counts) + graph.replay() + torch.cuda.synchronize() + for actual in outputs: + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + graph.reset() + + def _check_dcp_candidate_gather(rank, group, direct): from lmdeploy.pytorch.backends.cuda.attention.cp import get_dcp_manager from lmdeploy.pytorch.backends.cuda.nsa import TritonNSAIndexFP8Impl @@ -289,7 +327,7 @@ def gather(output, input, group): context = DistContext(dcp_group=group) with get_dist_manager().context(context): manager = get_dcp_manager() - manager.prepare_query_gather(2, 4) + manager.prepare_attention(2, 4) assert manager._query_workspace is None logits = communicator.create_all_gather_workspace(6, device=torch.device('cpu'), dtype=torch.float32) assert logits is None @@ -336,12 +374,14 @@ def test_cuda_gather_dispatch_and_workspace_ownership(monkeypatch, comm_env): context = DistContext(dcp_group=group) with get_dist_manager().context(context): manager = get_dcp_manager() - manager.prepare_query_gather(2, 4) + manager.prepare_attention(2, 4) manager = get_dcp_manager() - manager.prepare_query_gather(2, 4) + manager.prepare_attention(2, 4) query_workspace = manager._query_workspace logits_workspace = communicator.create_all_gather_workspace(6, device=torch.device('cpu'), dtype=torch.float32) - assert factory.call_count == 2 + assert factory.call_count == 3 + lse_workspace = manager._lse_workspace + lse_workspace.prepare.assert_called_once() query_workspace.prepare.assert_called_once() logits_workspace.prepare.assert_called_once() @@ -373,13 +413,15 @@ def test_cuda_gather_dispatch_and_workspace_ownership(monkeypatch, comm_env): group_close = Mock(wraps=group.close) monkeypatch.setattr(group, 'close', group_close) query_workspace.close.side_effect = lambda: group_close.assert_not_called() + lse_workspace.close.side_effect = lambda: group_close.assert_not_called() candidate_workspace.close.side_effect = lambda: group_close.assert_not_called() context.close() context.close() query_workspace.close.assert_called_once() + lse_workspace.close.assert_called_once() candidate_workspace.close.assert_called_once() assert context.dcp_manager is None - assert manager._query_workspace is manager._candidate_workspace is None + assert manager._query_workspace is manager._lse_workspace is manager._candidate_workspace is None logits_workspace.close.assert_not_called() From 24ab1c7f0b0dd009e022ee781bce2ef541bbc2eb Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 24 Sep 2026 11:44:47 +0000 Subject: [PATCH 11/15] perf: accelerate DCP prefix gathering and reordering --- .../pytorch/backends/cuda/attention/cp.py | 55 ++++++++++++------- .../pytorch/backends/cuda/attention/mla.py | 5 +- lmdeploy/pytorch/kernels/cuda/dcp.py | 31 ++++++----- tests/pytorch/kernel/test_mla_attention.py | 27 +++++---- tests/pytorch/test_cuda_communicator.py | 49 +++++++++++++++++ 5 files changed, 120 insertions(+), 47 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/attention/cp.py b/lmdeploy/pytorch/backends/cuda/attention/cp.py index 435c278a56..f4573cb10e 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/cp.py +++ b/lmdeploy/pytorch/backends/cuda/attention/cp.py @@ -8,25 +8,6 @@ from .default import TritonAttentionImpl, TritonAttentionMetaBuilder, TritonAttentionMetadata -def gather_dcp_prefix_kv(local_kv: torch.Tensor, chunk: DCPPrefixChunk, - *, zero_padding: bool = False) -> torch.Tensor: - """Gather a flattened prefix chunk and restore global token order. - - Each rank supplies [chunk.local_capacity, ...] with valid tokens packed by request. The result has - [requests * chunk.size, ...] capacity and valid tokens packed according to chunk.cu_seqlens. - Set zero_padding for consumers such as FA3 that may load masked V lanes: 0 * NaN can contaminate valid outputs. - """ - from lmdeploy.pytorch.distributed import all_gather_into_tensor - from lmdeploy.pytorch.kernels.cuda.dcp import reorder_dcp_prefill_kv - - gathered = local_kv.new_empty(chunk.kv_seqlens.numel() * chunk.size, *local_kv.shape[1:]) - all_gather_into_tensor(gathered, local_kv, group='dcp') - context = torch.zeros_like(gathered) if zero_padding else torch.empty_like(gathered) - reorder_dcp_prefill_kv(gathered, context, chunk_kv_seqlens=chunk.kv_seqlens, - kv_start_loc=chunk.cu_seqlens[:-1], local_lens=chunk.local_kv_seqlens) - return context - - class DCPManager: """Own shared DCP gathers and the output combine that orders query reuse. @@ -42,6 +23,7 @@ def __init__(self, group): self._query_workspace = None self._lse_workspace = None self._candidate_workspace = None + self._prefix_workspaces = {} def prepare_attention(self, num_heads: int, head_size: int): """Prepare query and LSE arenas shared across attention layers.""" @@ -64,6 +46,31 @@ def prepare_candidate_gather(self, topk: int): self._candidate_workspace = self.communicator.create_all_gather_workspace( topk * 2, device=torch.device('cuda'), dtype=torch.float32, dim=0) + def prepare_prefix_gather(self, *kv_widths: int): + """Share prefix arenas across layers and equal-width K/V payloads.""" + for width in kv_widths: + if width not in self._prefix_workspaces: + self._prefix_workspaces[width] = self.communicator.create_all_gather_workspace( + width, device=torch.device('cuda'), dtype=torch.bfloat16, dim=0) + + def gather_prefix(self, local_kv: torch.Tensor, chunk: DCPPrefixChunk, + *, zero_padding: bool = False) -> torch.Tensor: + """Gather cached KV and restore sequence order into an owned output. + + The entry barrier protects the previous reorder before arena reuse. Zero padding for consumers such as FA3 + that may load masked V lanes: 0 * NaN can contaminate valid outputs. + """ + from lmdeploy.pytorch.kernels.cuda.dcp import reorder_dcp_prefill_kv + + flat_kv = local_kv.flatten(1) + gathered = self.communicator.all_gather( + flat_kv, dim=0, workspace=self._prefix_workspaces.get(flat_kv.size(1)), copy_output=False) + gathered = gathered.view(-1, *local_kv.shape[1:]) + context = torch.zeros_like(gathered) if zero_padding else torch.empty_like(gathered) + reorder_dcp_prefill_kv(gathered, context, chunk_kv_seqlens=chunk.kv_seqlens, + kv_start_loc=chunk.cu_seqlens[:-1], local_lens=chunk.local_kv_seqlens) + return context + def gather_query(self, query: torch.Tensor) -> torch.Tensor: """Gather query heads; consume and combine on the same stream before reuse.""" @@ -125,6 +132,10 @@ def close(self): if self._candidate_workspace is not None: self._candidate_workspace.close() self._candidate_workspace = None + for workspace in self._prefix_workspaces.values(): + if workspace is not None: + workspace.close() + self._prefix_workspaces.clear() def get_dcp_manager() -> DCPManager: @@ -152,6 +163,8 @@ def __init__(self, num_heads: int, head_size: int, *, use_fa3: bool = False, **k self._fa3_prefill = flash_attn_varlen_func self.dcp_manager = get_dcp_manager() self.dcp_manager.prepare_attention(self.num_heads, self.head_size) + self.dcp_manager.prepare_prefix_gather(self.num_kv_heads * self.head_size, + self.num_kv_heads * self.v_head_size) def get_step_metadata_provider(self): """Use common DCP lengths without a separate kernel scheduler.""" @@ -202,8 +215,8 @@ def _gather_prefix(self, k_cache: torch.Tensor, v_cache: torch.Tensor, start_loc=chunk.local_cu_seqlens[:-1], out_size=chunk.local_capacity, out_dtype=dtype, flatten_kv_layout='shd', quant_policy=metadata.quant_policy, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros) - return (gather_dcp_prefix_kv(local_k, chunk, zero_padding=True), - gather_dcp_prefix_kv(local_v, chunk, zero_padding=True)) + return (self.dcp_manager.gather_prefix(local_k, chunk, zero_padding=True), + self.dcp_manager.gather_prefix(local_v, chunk, zero_padding=True)) def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, attn_metadata: TritonAttentionMetadata, diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index e55659b614..1022f32910 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -14,7 +14,7 @@ from lmdeploy.utils import get_logger from ..step_metadata import CudaAttentionMetaBuilder -from .cp import gather_dcp_prefix_kv, get_dcp_manager +from .cp import get_dcp_manager from .default import TritonAttentionImpl, TritonAttentionMetadata logger = get_logger('lmdeploy') @@ -257,6 +257,7 @@ def __init__( if self.dcp_world_size > 1: self.dcp_manager = get_dcp_manager() self.dcp_manager.prepare_attention(num_heads, head_size) + self.dcp_manager.prepare_prefix_gather(head_size) def get_step_metadata_provider(self): """Describe metadata required by this selected implementation.""" @@ -540,7 +541,7 @@ def _gather_dcp_prefix_chunk( k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, ) - context_k = gather_dcp_prefix_kv(local_k, chunk) + context_k = self.dcp_manager.gather_prefix(local_k, chunk) return context_k, chunk.cu_seqlens def _prefill_dcp_context( diff --git a/lmdeploy/pytorch/kernels/cuda/dcp.py b/lmdeploy/pytorch/kernels/cuda/dcp.py index ac005b7d47..9952a6e403 100644 --- a/lmdeploy/pytorch/kernels/cuda/dcp.py +++ b/lmdeploy/pytorch/kernels/cuda/dcp.py @@ -263,21 +263,23 @@ def _reorder_dcp_prefill_kv_kernel( stride_lr, stride_ls, DCP_SIZE: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_D: tl.constexpr, ): - source_row = tl.program_id(0) + source_row = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + row_mask = source_row < local_capacity * DCP_SIZE # Each rank contributes the same capacity, including trailing padding. rank = source_row // local_capacity local_position = source_row % local_capacity request = 0 - request_start = 0 - request_id = 0 - position_in_request = 0 - found = False + request_start = tl.full((BLOCK_M, ), 0, tl.int32) + request_id = tl.full((BLOCK_M, ), 0, tl.int32) + position_in_request = tl.full((BLOCK_M, ), 0, tl.int32) + found = tl.full((BLOCK_M, ), False, tl.int1) while request < num_sequences: request_len = tl.load(LocalLens + rank * stride_lr + - request * stride_ls).to(tl.int32) + request * stride_ls, mask=row_mask, other=0).to(tl.int32) owns_position = ((local_position >= request_start) & (local_position < request_start + request_len)) request_id = tl.where(owns_position, request, request_id) @@ -291,18 +293,18 @@ def _reorder_dcp_prefill_kv_kernel( global_position = position_in_request * DCP_SIZE + rank chunk_kv_seqlen = tl.load(ChunkKvSeqLens + request_id, mask=found, other=0) output_start = tl.load(KvStartLoc + request_id, mask=found, other=0) - valid_row = found & (global_position < chunk_kv_seqlen) + valid_row = row_mask & found & (global_position < chunk_kv_seqlen) dim_offsets = tl.arange(0, BLOCK_D) dim_mask = dim_offsets < row_width - value = tl.load(Gathered + source_row * stride_gs + - dim_offsets * stride_gd, - mask=valid_row & dim_mask, + value = tl.load(Gathered + source_row[:, None] * stride_gs + + dim_offsets[None, :] * stride_gd, + mask=valid_row[:, None] & dim_mask[None, :], other=0.0) output_row = output_start + global_position - tl.store(Output + output_row * stride_os + dim_offsets * stride_od, + tl.store(Output + output_row[:, None] * stride_os + dim_offsets[None, :] * stride_od, value, - mask=valid_row & dim_mask) + mask=valid_row[:, None] & dim_mask[None, :]) def sanitize_dcp_lse(local_lse: torch.Tensor, @@ -429,7 +431,9 @@ def reorder_dcp_prefill_kv(gathered: torch.Tensor, output: torch.Tensor, *, row_width = gathered_rows.size(1) block_d = triton.next_power_of_2(row_width) - _reorder_dcp_prefill_kv_kernel[(gathered_rows.size(0), )]( + # Process four rows per CTA to amortize scheduling and request-scan control overhead. + block_m = 4 + _reorder_dcp_prefill_kv_kernel[(triton.cdiv(gathered_rows.size(0), block_m), )]( gathered_rows, output_rows, chunk_kv_seqlens, @@ -442,6 +446,7 @@ def reorder_dcp_prefill_kv(gathered: torch.Tensor, output: torch.Tensor, *, *output_rows.stride(), *local_lens.stride(), DCP_SIZE=dcp_size, + BLOCK_M=block_m, BLOCK_D=block_d, num_warps=8, ) diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py index 7faf7624ea..1daf92a90b 100644 --- a/tests/pytorch/kernel/test_mla_attention.py +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -499,7 +499,11 @@ def test_dcp_cached_prefill_matches_reference_across_chunks(monkeypatch, sparse, kv_quant_policy=0) monkeypatch.setattr(distributed, 'get_dcp_world_rank', lambda: (dcp_size, 0)) # This prefill test simulates ranks without creating process groups. - monkeypatch.setattr(mla_module, 'get_dcp_manager', Mock()) + manager = object.__new__(DCPManager) + manager._prefix_workspaces = {} + manager.prepare_attention = Mock() + manager.prepare_prefix_gather = Mock() + monkeypatch.setattr(mla_module, 'get_dcp_manager', lambda: manager) # One virtual block per chunk, with an empty second request in every chunk. monkeypatch.setattr(cp_utils, 'get_dcp_prefill_workspace_size', lambda **kwargs: 2 * 64 * (1 + 2 * dcp_size) * 576 * 2) @@ -537,9 +541,10 @@ def test_dcp_cached_prefill_matches_reference_across_chunks(monkeypatch, sparse, reference_keys[:prefix_length, :, :512] = latent.flatten(-2).to(torch.bfloat16) gather_calls = 0 - def gather(output, local, group='tp'): + def gather(local, **kwargs): nonlocal gather_calls - assert group == 'dcp' + local = local.view(-1, 1, 576) + output = local.new_empty(dcp_size * local.size(0), 1, 576) chunk = metadata.dcp_prefix_chunks[gather_calls] prefix = reference_keys[chunk.start:min(chunk.start + chunk.size, prefix_length)] torch.testing.assert_close(local[:prefix[::dcp_size].size(0)], prefix[::dcp_size], atol=0, rtol=0) @@ -549,8 +554,9 @@ def gather(output, local, group='tp'): remote.zero_() remote[:prefix[rank::dcp_size].size(0)].copy_(prefix[rank::dcp_size]) gather_calls += 1 + return output.flatten(1) - monkeypatch.setattr(distributed, 'all_gather_into_tensor', gather) + manager.communicator = SimpleNamespace(all_gather=gather) indices = None if use_sparse: # Non-contiguous selections span cached chunks and current tokens. @@ -636,9 +642,7 @@ def test_dcp_attention_correction_kernel_matches_torch(dtype): @pytest.mark.parametrize('dcp_size', [2, 4]) @pytest.mark.parametrize('gather_prefix', [False, True], ids=['kernel', 'gather']) def test_reorder_dcp_prefill_kv_handles_uneven_requests(monkeypatch, dcp_size, gather_prefix): - from lmdeploy.pytorch import distributed from lmdeploy.pytorch.backends import cp_utils - from lmdeploy.pytorch.backends.cuda.attention.cp import gather_dcp_prefix_kv from lmdeploy.pytorch.kernels.cuda.dcp import reorder_dcp_prefill_kv device = 'cuda' @@ -664,13 +668,14 @@ def test_reorder_dcp_prefill_kv_handles_uneven_requests(monkeypatch, dcp_size, g expected[:valid_values.numel(), 0] = valid_values gathered = gathered.flatten(0, 1) if gather_prefix: - def all_gather(destination, local, *, group): - assert group == 'dcp' + def all_gather(local, **kwargs): assert local.size(0) == local_capacity - destination.copy_(gathered) + return gathered - monkeypatch.setattr(distributed, 'all_gather_into_tensor', all_gather) - output = gather_dcp_prefix_kv(gathered[:local_capacity], chunk, zero_padding=True) + manager = object.__new__(DCPManager) + manager.communicator = SimpleNamespace(all_gather=all_gather) + manager._prefix_workspaces = {} + output = manager.gather_prefix(gathered[:local_capacity], chunk, zero_padding=True) else: reorder_dcp_prefill_kv(gathered, output, diff --git a/tests/pytorch/test_cuda_communicator.py b/tests/pytorch/test_cuda_communicator.py index 50a13ea14b..04e20005f7 100644 --- a/tests/pytorch/test_cuda_communicator.py +++ b/tests/pytorch/test_cuda_communicator.py @@ -88,6 +88,7 @@ def merge(query): _check_dcp_merge(rank, ctx.dcp_group, manager, all(enabled)) _check_lm_head_lifecycle(rank, all(enabled)) _check_dcp_candidate_gather(rank, ctx.dcp_group, all(enabled)) + _check_dcp_prefix_gather(rank, ctx.dcp_group, all(enabled)) finally: ctx.close() dist.destroy_process_group() @@ -130,6 +131,46 @@ def _check_dcp_merge(rank, group, manager, direct): graph.reset() +def _check_dcp_prefix_gather(rank, group, direct): + from lmdeploy.pytorch.backends.cp_utils import build_dcp_prefix_chunks + from lmdeploy.pytorch.backends.cuda.attention.cp import DCPManager, get_dcp_manager + + manager = get_dcp_manager() + native = DCPManager(SimpleNamespace(gpu_group=group.gpu_group, rank=rank, + communicator=base_communicator_module.DeviceCommunicator(group.gpu_group))) + manager.prepare_prefix_gather(576, 256, 384, 576) + assert len(manager._prefix_workspaces) == 3 + assert all((ws is not None and ws.is_available()) == direct for ws in manager._prefix_workspaces.values()) + lengths = torch.tensor([0, 1, 3, 127, 256, 259], device='cuda', dtype=torch.int32) + chunk, = build_dcp_prefix_chunks(prefix_lens=lengths, prefix_limit=259, block_size=64, + kv_width=576, dcp_world_rank=(2, rank)) + for heads, dim, dtype in ((1, 576, torch.bfloat16), (2, 128, torch.bfloat16), + (2, 192, torch.bfloat16), (1, 576, torch.float16)): + local = torch.randn(chunk.local_capacity, heads, dim, device='cuda', dtype=dtype) + expected = native.gather_prefix(local, chunk, zero_padding=True) + actual = manager.gather_prefix(local, chunk, zero_padding=True) + torch.testing.assert_close(actual.view(torch.int16), expected.view(torch.int16), rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + outputs = [] + with torch.cuda.graph(graph): + for step in range(4): + if rank == step % 2: + torch.cuda._sleep(20000) + outputs.append(manager.gather_prefix(local, chunk, zero_padding=True)) + for _ in range(3): + local.normal_() + expected = native.gather_prefix(local, chunk, zero_padding=True) + graph.replay() + torch.cuda.synchronize() + for actual in outputs: + torch.testing.assert_close(actual.view(torch.int16), expected.view(torch.int16), rtol=0, atol=0) + graph.reset() + # A strided payload with an unregistered width retains native behavior. + local = local[..., ::2] + torch.testing.assert_close(manager.gather_prefix(local, chunk, zero_padding=True), + native.gather_prefix(local, chunk, zero_padding=True), rtol=0, atol=0) + + def _check_dcp_candidate_gather(rank, group, direct): from lmdeploy.pytorch.backends.cuda.attention.cp import get_dcp_manager from lmdeploy.pytorch.backends.cuda.nsa import TritonNSAIndexFP8Impl @@ -408,6 +449,9 @@ def test_cuda_gather_dispatch_and_workspace_ownership(monkeypatch, comm_env): candidate_workspace = manager._candidate_workspace manager.prepare_candidate_gather(512) assert manager._candidate_workspace is candidate_workspace + manager.prepare_prefix_gather(576, 128, 576) + prefix_workspaces = list(manager._prefix_workspaces.values()) + assert len(prefix_workspaces) == 2 assert context.dcp_manager is manager monkeypatch.setattr(torch.distributed, 'is_initialized', lambda: True) group_close = Mock(wraps=group.close) @@ -415,11 +459,16 @@ def test_cuda_gather_dispatch_and_workspace_ownership(monkeypatch, comm_env): query_workspace.close.side_effect = lambda: group_close.assert_not_called() lse_workspace.close.side_effect = lambda: group_close.assert_not_called() candidate_workspace.close.side_effect = lambda: group_close.assert_not_called() + for workspace in prefix_workspaces: + workspace.close.side_effect = lambda: group_close.assert_not_called() context.close() context.close() query_workspace.close.assert_called_once() lse_workspace.close.assert_called_once() candidate_workspace.close.assert_called_once() + for workspace in prefix_workspaces: + workspace.close.assert_called_once() + assert not manager._prefix_workspaces assert context.dcp_manager is None assert manager._query_workspace is manager._lse_workspace is manager._candidate_workspace is None logits_workspace.close.assert_not_called() From 8ea44bf223ac320d016ae4d1ccc8229ee792771b Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 9 Oct 2026 12:05:05 +0000 Subject: [PATCH 12/15] test: simplify CUDA communicator tests --- tests/pytorch/test_cuda_communicator.py | 320 +++--------------------- 1 file changed, 34 insertions(+), 286 deletions(-) diff --git a/tests/pytorch/test_cuda_communicator.py b/tests/pytorch/test_cuda_communicator.py index 04e20005f7..c3ee73257d 100644 --- a/tests/pytorch/test_cuda_communicator.py +++ b/tests/pytorch/test_cuda_communicator.py @@ -9,232 +9,49 @@ from lmdeploy.pytorch.backends.cuda.comm import communicator as communicator_module -def _run_dcp_query_gather(rank, rendezvous, enabled): +def _run_auto_all_gather(rank, rendezvous): from datetime import timedelta from torch import distributed as dist - from lmdeploy.pytorch.backends.cuda.attention.cp import get_dcp_manager from lmdeploy.pytorch.backends.cuda.op_backend import CudaOpsBackend from lmdeploy.pytorch.config import DistConfig - from lmdeploy.pytorch.distributed import DistContext, get_dist_manager torch.cuda.set_device(rank) dist.init_process_group('nccl', init_method=rendezvous, rank=rank, world_size=2, timeout=timedelta(seconds=60)) - backend = 'auto' if enabled[rank] else 'nccl' - if enabled[0] != enabled[1]: - try: - with pytest.raises(ValueError, match='agree across ranks'): - DistContext.build(rank, DistConfig(tp=2, dcp=2, communication_backend=backend), - communicator_builder=CudaOpsBackend.build_communicator) - finally: - dist.destroy_process_group() - return - ctx = DistContext.build(rank, DistConfig(tp=2, dcp=2, communication_backend=backend), - communicator_builder=CudaOpsBackend.build_communicator) + cpu_group = dist.new_group(backend='gloo') + communicator = communicator_module.build_cuda_communicator( + cpu_group, + dist.group.WORLD, + DistConfig(tp=2, dcp=2, communication_backend='auto'), + group_name='dcp') try: - with get_dist_manager().context(ctx): - manager = get_dcp_manager() - manager.prepare_attention(32, 576) - workspace = manager._query_workspace - if all(enabled): - assert workspace.is_available() - else: - assert workspace is None - manager = get_dcp_manager() - manager.prepare_attention(32, 576) - assert manager._query_workspace is workspace - query = torch.empty(384, 32, 576, device='cuda', dtype=torch.bfloat16) - # Ineligible inputs still return correct results through NCCL. - for fallback in (query[:1], query[..., ::2], query.float(), - torch.empty(1024, 32, 576, device='cuda', dtype=query.dtype)): - fallback.fill_(rank) - output = manager.gather_query(fallback) - expected = torch.arange(2, device='cuda', dtype=fallback.dtype) - expected = expected.repeat_interleave(32)[None, :, None].expand_as(output) - torch.testing.assert_close(output, expected, rtol=0, atol=0) - rows = (2, 16, 96, 384) - - def merge(query): - lse = torch.zeros(query.shape[:2], device='cuda') - counts = torch.ones(query.size(0), device='cuda', dtype=torch.int32) - return manager.combine(query, lse, counts) + workspace = communicator.create_all_gather_workspace( + 256, device=torch.device('cuda'), dtype=torch.bfloat16) + assert workspace.is_available() - query.fill_(rank) - for count in rows: - torch.testing.assert_close(merge(manager.gather_query(query[:count])), query[:count], rtol=0, atol=0) - torch.cuda.synchronize() - dist.barrier() - graph = torch.cuda.CUDAGraph() - outputs = [] - with torch.cuda.graph(graph): - for step in range(16): - query.fill_(rank + step * 8) - gathered = manager.gather_query(query[:rows[step % len(rows)]]) - # A slow reader must finish before any rank reuses the arena. - if rank == step % 2: - torch.cuda._sleep(20000) - outputs.append(gathered.clone()) - merge(gathered) - for _ in range(3): - graph.replay() - torch.cuda.synchronize() - graph.reset() - for step, output in enumerate(outputs): - expected = (torch.arange(2, device='cuda', dtype=query.dtype) + step * 8) - expected = expected.repeat_interleave(32)[None, :, None].expand_as(output) - torch.testing.assert_close(output, expected, rtol=0, atol=0) - _check_dcp_merge(rank, ctx.dcp_group, manager, all(enabled)) - _check_lm_head_lifecycle(rank, all(enabled)) - _check_dcp_candidate_gather(rank, ctx.dcp_group, all(enabled)) - _check_dcp_prefix_gather(rank, ctx.dcp_group, all(enabled)) - finally: - ctx.close() - dist.destroy_process_group() - - -def _check_dcp_merge(rank, group, manager, direct): - from lmdeploy.pytorch.backends.cuda.attention.cp import DCPManager - from lmdeploy.pytorch.distributed import DistGroup - - native = DCPManager(DistGroup(rank=rank, gpu_group=group.gpu_group, - communicator=base_communicator_module.DeviceCommunicator(group.gpu_group))) - workspace = manager._lse_workspace - assert (workspace is not None and workspace.is_available()) == direct - for rows in (1, 16, 96): - output = torch.randn(rows, 64, 128, device='cuda', dtype=torch.bfloat16) - lse = torch.randn(rows, 68, device='cuda')[:, :64] - counts = torch.ones(rows, device='cuda', dtype=torch.int32) - counts[-1] = rank - lse[0, 0] = float('inf') - expected = native.combine(output, lse, counts) - for _ in range(3): - actual = manager.combine(output, lse, counts) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) + rows = (2, 16, 96, 384) + inputs = [torch.empty(row, 128, device='cuda', dtype=torch.bfloat16) for row in rows] graph = torch.cuda.CUDAGraph() outputs = [] with torch.cuda.graph(graph): - for step in range(4): - if rank == step % 2: - torch.cuda._sleep(20000) - outputs.append(manager.combine(output, lse, counts)) - for _ in range(3): - output.normal_() - lse.normal_() - lse[0, 0] = float('nan') - expected = native.combine(output, lse, counts) - graph.replay() - torch.cuda.synchronize() - for actual in outputs: - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - graph.reset() - + for input in inputs: + outputs.append(communicator.all_gather(input, workspace=workspace, copy_output=False)) -def _check_dcp_prefix_gather(rank, group, direct): - from lmdeploy.pytorch.backends.cp_utils import build_dcp_prefix_chunks - from lmdeploy.pytorch.backends.cuda.attention.cp import DCPManager, get_dcp_manager - - manager = get_dcp_manager() - native = DCPManager(SimpleNamespace(gpu_group=group.gpu_group, rank=rank, - communicator=base_communicator_module.DeviceCommunicator(group.gpu_group))) - manager.prepare_prefix_gather(576, 256, 384, 576) - assert len(manager._prefix_workspaces) == 3 - assert all((ws is not None and ws.is_available()) == direct for ws in manager._prefix_workspaces.values()) - lengths = torch.tensor([0, 1, 3, 127, 256, 259], device='cuda', dtype=torch.int32) - chunk, = build_dcp_prefix_chunks(prefix_lens=lengths, prefix_limit=259, block_size=64, - kv_width=576, dcp_world_rank=(2, rank)) - for heads, dim, dtype in ((1, 576, torch.bfloat16), (2, 128, torch.bfloat16), - (2, 192, torch.bfloat16), (1, 576, torch.float16)): - local = torch.randn(chunk.local_capacity, heads, dim, device='cuda', dtype=dtype) - expected = native.gather_prefix(local, chunk, zero_padding=True) - actual = manager.gather_prefix(local, chunk, zero_padding=True) - torch.testing.assert_close(actual.view(torch.int16), expected.view(torch.int16), rtol=0, atol=0) - graph = torch.cuda.CUDAGraph() - outputs = [] - with torch.cuda.graph(graph): - for step in range(4): - if rank == step % 2: - torch.cuda._sleep(20000) - outputs.append(manager.gather_prefix(local, chunk, zero_padding=True)) - for _ in range(3): - local.normal_() - expected = native.gather_prefix(local, chunk, zero_padding=True) + for step in range(3): + for input in inputs: + input.fill_(rank + step) graph.replay() torch.cuda.synchronize() - for actual in outputs: - torch.testing.assert_close(actual.view(torch.int16), expected.view(torch.int16), rtol=0, atol=0) + for input, output in zip(inputs, outputs): + expected = torch.cat((input, input + 1), dim=-1) + torch.testing.assert_close(output, expected, rtol=0, atol=0) graph.reset() - # A strided payload with an unregistered width retains native behavior. - local = local[..., ::2] - torch.testing.assert_close(manager.gather_prefix(local, chunk, zero_padding=True), - native.gather_prefix(local, chunk, zero_padding=True), rtol=0, atol=0) - - -def _check_dcp_candidate_gather(rank, group, direct): - from lmdeploy.pytorch.backends.cuda.attention.cp import get_dcp_manager - from lmdeploy.pytorch.backends.cuda.nsa import TritonNSAIndexFP8Impl - from lmdeploy.pytorch.kernels.cuda.sparse_index_dcp_topk import pack_dcp_topk_candidates, sparse_dcp_global_topk - - k = 512 - # Candidate merging consumes existing scores and does not require DeepGEMM. - impl = object.__new__(TritonNSAIndexFP8Impl) - impl.topk, impl.fill = k, -1 - impl.dcp_world_size, impl.dcp_rank = 2, rank - impl.dcp_manager = get_dcp_manager() - impl.dcp_manager.prepare_candidate_gather(k) - workspace = impl.dcp_manager._candidate_workspace - assert (workspace is not None and workspace.is_available()) == direct - get_dcp_manager().prepare_candidate_gather(k) - assert impl.dcp_manager._candidate_workspace is workspace - native = base_communicator_module.DeviceCommunicator(group.gpu_group) - # IDs are bitcast, not converted to FP32; include large IDs and NaN encodings. - bits = (torch.arange(16 * k * 2, device='cuda', dtype=torch.int32) * 1234567 + rank).view(16, k * 2) - payload = bits.view(torch.float32) - for value in (payload, payload[:1], payload[:, ::2], payload.double()): - actual = group.communicator.all_gather(value, dim=0, workspace=workspace) - expected = native.all_gather(value, dim=0) - torch.testing.assert_close(actual.view(torch.int32), expected.view(torch.int32), rtol=0, atol=0) - if direct: - # Capacity rejection must use NCCL consistently, including a captured - # shape that was not warmed before capture. - oversized = payload[:1].expand(workspace._state.max_token_num // 2 + 1, -1).contiguous() - torch.testing.assert_close( - group.communicator.all_gather(oversized, dim=0, workspace=workspace).view(torch.int32), - native.all_gather(oversized, dim=0).view(torch.int32), rtol=0, atol=0) - scores = torch.randn(16, k * 2, device='cuda') - indices = torch.arange(k, device='cuda', dtype=torch.int32).expand(16, -1).clone() - - def reference(): - packed = pack_dcp_topk_candidates(scores, indices, dcp_world_rank=(2, rank)) - gathered = native.all_gather(packed.flatten(1), dim=0).view(2, 16, k, 2) - return sparse_dcp_global_topk(gathered, k=k) - - impl._merge_dcp_topk(scores, indices) - graph = torch.cuda.CUDAGraph() - outputs = [] - with torch.cuda.graph(graph): - for step in range(4): - if rank == step % 2: - torch.cuda._sleep(20000) - outputs.append(impl._merge_dcp_topk(scores, indices)) - unwarmed = group.communicator.all_gather(payload[:3], dim=0, workspace=workspace) - for case in range(3): - scores.normal_() - if case == 1: - scores.zero_() - if case == 2: - indices[:, k // 2:] = -1 - indices[0] = -1 - expected = reference() - torch.testing.assert_close(impl._merge_dcp_topk(scores, indices), expected, rtol=0, atol=0) - graph.replay() - torch.cuda.synchronize() - for output in outputs: - torch.testing.assert_close(output, expected, rtol=0, atol=0) - torch.testing.assert_close(unwarmed.view(torch.int32), - native.all_gather(payload[:3], dim=0).view(torch.int32), rtol=0, atol=0) - graph.reset() + workspace.close() + finally: + communicator.close() + dist.destroy_process_group() def _run_auto_all_reduce(rank, rendezvous): @@ -290,51 +107,17 @@ def test_auto_all_reduce_eager_and_graph(tmp_path): args=(f'file://{tmp_path}/auto_all_reduce',), nprocs=2, join=True) -def _check_lm_head_lifecycle(rank, direct): - from lmdeploy.pytorch.nn import ParallelEmbedding, ParallelLMHead - - head = ParallelLMHead(250, 128, dtype=torch.bfloat16, device='cuda') - if direct: - assert head._logits_gather_workspace.is_available() - head.weight.data.fill_(rank + 1) - hidden = torch.ones(2, 128, dtype=torch.bfloat16, device='cuda') - expected = torch.arange(1, 3, device='cuda', dtype=hidden.dtype).repeat_interleave(128)[:250] * 128 - logits = head(hidden) - head(hidden * 2) - torch.testing.assert_close(logits, expected.expand_as(logits), atol=0, rtol=0) - if direct: - assert head._logits_gather_workspace.is_available() - # A coordinated dtype move retires the BF16 arena and uses native gathering. - head.to(dtype=torch.float16) - if direct: - assert not head._logits_gather_workspace.is_available() - logits = head(hidden) - torch.testing.assert_close(logits, expected.to(logits.dtype).expand_as(logits), atol=0, rtol=0) - # Tying BF16 weights must re-admit the arena before the next forward. - embedding = ParallelEmbedding(250, 128, None, dtype=torch.bfloat16, device='cuda', is_tp=True) - embedding.weight.data.fill_(rank + 1) - head.tie_weights(embedding) - if direct: - assert head._logits_gather_workspace.is_available() - logits = head(hidden) - torch.testing.assert_close(logits, expected.expand_as(logits), atol=0, rtol=0) - if direct: - head._logits_gather_workspace.close() - - @pytest.mark.skipif(torch.cuda.device_count() < 2, reason='requires two CUDA GPUs') -@pytest.mark.parametrize('enabled', [(False, False), (True, True), (True, False)], - ids=['nccl', 'auto', 'mixed_config']) -def test_dcp_query_gather_graph_reuses_arena(tmp_path, enabled): - if all(enabled): - from torch.distributed._symmetric_memory import DeviceType, _SymmetricMemory - if any(torch.cuda.get_device_capability(i)[0] < 9 - or not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, i) for i in range(2)): - pytest.skip('requires SM90 or newer with multicast support') - if not torch.cuda.can_device_access_peer(0, 1): - pytest.skip('requires peer access') - torch.multiprocessing.spawn(_run_dcp_query_gather, - args=((tmp_path / 'rendezvous').as_uri(), enabled), nprocs=2) +def test_auto_all_gather_reuses_workspace(tmp_path): + from torch.distributed._symmetric_memory import DeviceType, _SymmetricMemory + + if any(torch.cuda.get_device_capability(i)[0] < 9 + or not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, i) for i in range(2)): + pytest.skip('requires SM90 or newer with multicast support') + if not torch.cuda.can_device_access_peer(0, 1): + pytest.skip('requires peer access') + torch.multiprocessing.spawn(_run_auto_all_gather, + args=(f'file://{tmp_path}/auto_all_gather',), nprocs=2, join=True) @pytest.fixture @@ -596,38 +379,3 @@ def peer_failure(flags, ready, group): assert not provider.is_available() assert provider._buffer is None rendezvous.assert_not_called() - - -def test_communicator_rejects_rank_backend_mismatch(monkeypatch, comm_env): - from lmdeploy.pytorch.backends.cuda.op_backend import CudaOpsBackend - from lmdeploy.pytorch.config import DistConfig - - monkeypatch.setattr(torch.cuda, 'get_device_capability', lambda device: (9, 0)) - - def disagree(output, local, group): - output[:] = [local, ('auto', *local[1:])] - - monkeypatch.setattr(communicator_module.dist, 'all_gather_object', disagree) - optimized = Mock() - monkeypatch.setattr(communicator_module, 'CudaCommunicator', optimized) - with pytest.raises(ValueError, match='agree across ranks'): - CudaOpsBackend.build_communicator('cpu', 'gpu', DistConfig(tp=2, communication_backend='nccl')) - optimized.assert_not_called() - - -@pytest.mark.parametrize('shared_tp_group', [False, True]) -def test_communicator_group_names_and_tp_deduplication(shared_tp_group): - from lmdeploy.pytorch.distributed import DistContext, DistGroup, _build_communicators - - tp_group = DistGroup(cpu_group='tp_cpu', gpu_group='tp_gpu') - mlp_group = tp_group if shared_tp_group else DistGroup(cpu_group='mlp_cpu', gpu_group='mlp_gpu') - dcp_group = DistGroup(cpu_group='dcp_cpu', gpu_group='dcp_gpu') - builder = Mock(side_effect=lambda **kwargs: Mock()) - context = DistContext(attn_tp_group=tp_group, mlp_tp_group=mlp_group, moe_tp_group=mlp_group, - dcp_group=dcp_group, communicator_builder=builder) - _build_communicators(context) - - names = [call.kwargs['group_name'] for call in builder.call_args_list] - assert names == (['tp', 'dcp'] if shared_tp_group else ['tp', 'tp', 'dcp']) - assert (tp_group.communicator is mlp_group.communicator) == shared_tp_group - assert tp_group.communicator is not dcp_group.communicator From ce91b1f99cf89ad39167f72cd53762df55f5a214 Mon Sep 17 00:00:00 2001 From: zxy Date: Sat, 10 Oct 2026 09:29:19 +0000 Subject: [PATCH 13/15] refactor: clarify optimized all-reduce routing --- lmdeploy/pytorch/distributed.py | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/lmdeploy/pytorch/distributed.py b/lmdeploy/pytorch/distributed.py index b556687ebb..0a87b87e81 100644 --- a/lmdeploy/pytorch/distributed.py +++ b/lmdeploy/pytorch/distributed.py @@ -463,20 +463,29 @@ def get_group(group_type: str, device: str): raise RuntimeError(f'Unknown group type: {group_type}') +def _try_optimized_all_reduce(tensor, op, group, async_op): + """Run an optimized synchronous CUDA SUM when eligible.""" + if not (tensor.is_cuda and op == ReduceOp.SUM and not async_op and group is not None): + return False + + context = get_dist_manager().current_context() + # The public API exposes raw process groups, so find the wrapper that owns this one. + for tp_group in (context.attn_tp_group, context.mlp_tp_group, context.moe_tp_group): + if tp_group is not None and tp_group.gpu_group is group: + if tp_group.communicator is not None: + tp_group.communicator.all_reduce_(tensor) + return True + break + return False + + def all_reduce(tensor, op=ReduceOp.SUM, group='tp', async_op=False): """All reduce.""" if isinstance(group, str): group = get_group(group, 'gpu') - # Optimize synchronous CUDA SUM on explicit TP groups; keep other calls native. - if tensor.is_cuda and op == ReduceOp.SUM and not async_op and group is not None: - context = get_dist_manager().current_context() - # Resolve the raw ProcessGroup to its communicator. - for tp_group in (context.attn_tp_group, context.mlp_tp_group, context.moe_tp_group): - if tp_group is not None and tp_group.gpu_group is group: - if tp_group.communicator is not None: - return tp_group.communicator.all_reduce_(tensor) - break + if _try_optimized_all_reduce(tensor, op, group, async_op): + return None return dist.all_reduce(tensor, op, group, async_op) From d6f7c1d78ead6de35149d98bafb456cba8b667d8 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 12 Oct 2026 03:06:00 +0000 Subject: [PATCH 14/15] style: fix lint issues --- lmdeploy/pytorch/third_party/flash_attn_interface.py | 1 - tests/pytorch/test_cuda_communicator.py | 1 - 2 files changed, 2 deletions(-) diff --git a/lmdeploy/pytorch/third_party/flash_attn_interface.py b/lmdeploy/pytorch/third_party/flash_attn_interface.py index bf9984ff11..3d7b11f4b7 100644 --- a/lmdeploy/pytorch/third_party/flash_attn_interface.py +++ b/lmdeploy/pytorch/third_party/flash_attn_interface.py @@ -2,7 +2,6 @@ import functools import torch - from flash_attn_interface import flash_attn_varlen_func as _flash_attn_varlen_func from flash_attn_interface import flash_attn_with_kvcache as _flash_attn_with_kvcache diff --git a/tests/pytorch/test_cuda_communicator.py b/tests/pytorch/test_cuda_communicator.py index c3ee73257d..31b4e9e0a1 100644 --- a/tests/pytorch/test_cuda_communicator.py +++ b/tests/pytorch/test_cuda_communicator.py @@ -14,7 +14,6 @@ def _run_auto_all_gather(rank, rendezvous): from torch import distributed as dist - from lmdeploy.pytorch.backends.cuda.op_backend import CudaOpsBackend from lmdeploy.pytorch.config import DistConfig torch.cuda.set_device(rank) From 61bf5d6a19579d16c231e5d78580f9ba791762bd Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 12 Oct 2026 09:51:42 +0000 Subject: [PATCH 15/15] fix: reject ALiBi for GQA DCP --- lmdeploy/pytorch/backends/cuda/attention/cp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lmdeploy/pytorch/backends/cuda/attention/cp.py b/lmdeploy/pytorch/backends/cuda/attention/cp.py index f4573cb10e..15dbbcdd47 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/cp.py +++ b/lmdeploy/pytorch/backends/cuda/attention/cp.py @@ -158,6 +158,7 @@ class DCPAttentionImpl(TritonAttentionImpl): def __init__(self, num_heads: int, head_size: int, *, use_fa3: bool = False, **kwargs): super().__init__(num_heads=num_heads, head_size=head_size, **kwargs) self.use_fa3 = use_fa3 + assert not self.alibi, 'GQA DCP does not support ALiBi' if use_fa3: from lmdeploy.pytorch.third_party.flash_attn_interface import flash_attn_varlen_func self._fa3_prefill = flash_attn_varlen_func