From 852d6603ce03188610ca7067807b11aa240ec984 Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Mon, 21 Sep 2026 12:04:31 +0700 Subject: [PATCH 01/10] [DCP] Add fi_a2a_fused comm backend using FlashInfer's fused A2A + LSE reduce The existing a2a and fi_a2a backends exchange head partials and then merge them with a separate Triton LSE kernel, so every MLA layer round-trips the partials through global memory between the two kernels. FlashInfer's fused decode-CP op does the exchange and the LSE-weighted reduce inside one cooperative kernel over NCCL device-side LSA symmetric memory, removing that round-trip and one launch per layer per decode step. Expose it as a fourth --dcp-comm-backend value rather than changing fi_a2a in place, so the validated MNNVL path is untouched and the two can be compared directly. Resolution never selects it on its own; it is explicit-only. The fused op allows one workspace per ordered CUDA stream, so the workspaces are allocated upfront (rendezvous is collective and not graph-capturable) and handed out per stream: one for eager decode, one for the global graph-capture stream. Requires a FlashInfer build carrying the op (flashinfer-ai/flashinfer#4929), NCCL >= 2.29 with device-communicator support, and a DCP group inside one NVLink domain. fp16/bf16 only -- unlike a2a, which ships partials as opaque bytes, the fused kernel has no fp8 path. Co-Authored-By: Claude Opus 5 (1M context) --- .../sglang/srt/arg_groups/fields/parallel.py | 11 +- python/sglang/srt/arg_groups/parallel_hook.py | 16 +- python/sglang/srt/layers/dcp/__init__.py | 2 + python/sglang/srt/layers/dcp/comm.py | 191 ++++++++++++++++ .../srt/model_executor/runner/base_runner.py | 32 +++ .../attention_forward_methods/forward_mla.py | 6 +- .../forward_mla_rocm.py | 6 +- .../dcp/test_dcp_fi_a2a_fused_equivalence.py | 213 ++++++++++++++++++ .../unit/server_args/test_server_args.py | 37 +++ 9 files changed, 502 insertions(+), 12 deletions(-) create mode 100644 test/registered/dcp/test_dcp_fi_a2a_fused_equivalence.py diff --git a/python/sglang/srt/arg_groups/fields/parallel.py b/python/sglang/srt/arg_groups/fields/parallel.py index 1b4df5bdf70c..745e3cfb780d 100644 --- a/python/sglang/srt/arg_groups/fields/parallel.py +++ b/python/sglang/srt/arg_groups/fields/parallel.py @@ -126,10 +126,13 @@ class Parallel(msgspec.Struct): help="Communication backend for the decode context-parallel (DCP) " "attention reduction: 'ag_rs' (AllGather + ReduceScatter), 'a2a' " "(fused NCCL All-to-All exchange of output+LSE + local Triton LSE " - "combine), or 'fi_a2a' (FlashInfer MNNVL All-to-All kernel; requires " - "Blackwell and a DCP group within one MNNVL domain). Unset resolves " - "to 'fi_a2a' where supported, else 'a2a' on CUDA/ROCm, else 'ag_rs'.", - choices=["ag_rs", "a2a", "fi_a2a"], + "combine), 'fi_a2a' (FlashInfer MNNVL All-to-All kernel; requires " + "Blackwell and a DCP group within one MNNVL domain), or " + "'fi_a2a_fused' (FlashInfer fused All-to-All + LSE reduce in a single " + "kernel; same topology requirement as 'fi_a2a', fp16/bf16 only). Unset " + "resolves to 'fi_a2a' where supported, else 'a2a' on CUDA/ROCm, else " + "'ag_rs'.", + choices=["ag_rs", "a2a", "fi_a2a", "fi_a2a_fused"], resolvable=True, ), ] = None diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index 1c92d5791f57..b503ff7cbc40 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -170,7 +170,7 @@ def handle_decode_context_parallelism(server_args: Any): "--decode-context-parallel-size) must be >= 1, but got " f"dcp_size={cfg.dcp_size}." ) - if cfg.dcp_comm_backend in ("a2a", "fi_a2a") and cfg.dcp_size <= 1: + if cfg.dcp_comm_backend in ("a2a", "fi_a2a", "fi_a2a_fused") and cfg.dcp_size <= 1: raise ValueError( f"--dcp-comm-backend {cfg.dcp_comm_backend} only affects the " "decode context-parallel attention reduction and therefore " @@ -184,13 +184,21 @@ def handle_decode_context_parallelism(server_args: Any): "within one MNNVL domain. Use 'a2a' or 'ag_rs' elsewhere, or leave " "the flag unset to resolve it." ) + if cfg.dcp_comm_backend == "fi_a2a_fused" and not get_platform().is_cuda: + raise ValueError( + "--dcp-comm-backend fi_a2a_fused delegates both the exchange and the " + "LSE reduce to FlashInfer's fused kernel, which requires Blackwell and " + "a DCP group within one NVLink domain. Use 'a2a' or 'ag_rs' elsewhere, " + "or leave the flag unset to resolve it." + ) if cfg.dcp_replicate_q_proj: if cfg.dcp_size <= 1: raise ValueError("--dcp-replicate-q-proj requires --dcp-size > 1.") - if cfg.dcp_comm_backend not in ("a2a", "fi_a2a"): + if cfg.dcp_comm_backend not in ("a2a", "fi_a2a", "fi_a2a_fused"): raise ValueError( - "--dcp-replicate-q-proj only applies to the a2a/fi_a2a DCP " - "communication backend (it removes the head-dim Q all-gather); " + "--dcp-replicate-q-proj only applies to the a2a/fi_a2a/" + "fi_a2a_fused DCP communication backend (it removes the " + "head-dim Q all-gather); " f"got --dcp-comm-backend={cfg.dcp_comm_backend}." ) diff --git a/python/sglang/srt/layers/dcp/__init__.py b/python/sglang/srt/layers/dcp/__init__.py index 54ae19195a1e..957b9bf7baab 100644 --- a/python/sglang/srt/layers/dcp/__init__.py +++ b/python/sglang/srt/layers/dcp/__init__.py @@ -44,6 +44,7 @@ dcp_enabled, get_attention_dcp_rank, get_attention_dcp_world_size, + init_fi_a2a_fused_workspace, init_fi_a2a_workspace, ) from sglang.srt.layers.dcp.layout import ( @@ -64,6 +65,7 @@ __all__ = [ "DecodeContextParallelMetadata", "dcp_a2a_lse_reduce", + "init_fi_a2a_fused_workspace", "init_fi_a2a_workspace", "all_gather_kv_cache_for_dcp", "all_gather_kv_cache_for_mha_chunk_extend", diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index ae2d1422d38e..4ad3025b3792 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/python/sglang/srt/layers/dcp/comm.py @@ -18,8 +18,13 @@ PR #25090 vs #14194): - cp_lse_ag_out_rs_mha: torch / natural-log logsumexp / all-reduce + head slice - cp_lse_ag_out_rs_mla: Triton (log2/exp2) correction / reduce-scatter + +The a2a family exchanges head partials instead: a2a (NCCL) and fi_a2a +(FlashInfer MNNVL) both combine locally with a Triton kernel, while +fi_a2a_fused does the exchange and the merge in one FlashInfer kernel. """ +import logging import warnings from typing import Optional @@ -40,6 +45,8 @@ from sglang.srt.utils import is_hip from sglang.srt.utils.common import is_fi_a2a_supported +logger = logging.getLogger(__name__) + _is_hip = is_hip() @@ -486,6 +493,11 @@ def dcp_a2a_lse_reduce( cp_attn_out, cp_attn_lse, cp_group, is_lse_base_on_e ) + if comm_backend == "fi_a2a_fused": + return _dcp_fi_a2a_fused_lse_reduce( + cp_attn_out, cp_attn_lse, cp_group, is_lse_base_on_e + ) + N = cp_group.world_size B, H, D = cp_attn_out.shape assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})" @@ -586,3 +598,182 @@ def _dcp_fi_a2a_lse_reduce( recv_output, recv_lse, is_lse_base_on_e=is_lse_base_on_e ) return combined + + +# fi_a2a_fused (flashinfer #4929): fuses the exchange and the LSE merge into one +# kernel. NCCL LSA symmetric memory, not the MNNVL FIFO fi_a2a uses, hence a +# separate workspace. +_FI_A2A_FUSED_STATE: Optional[dict] = None + +# Default stream + the global graph-capture stream. Rendezvous is collective and +# not capturable, so per-stream workspaces cannot be created on demand. +_FI_A2A_FUSED_NUM_STREAMS = 2 + + +def init_fi_a2a_fused_workspace( + cp_group: "GroupCoordinator", + *, + max_tokens: int, + local_heads: int, + head_dim: int, + dtype: torch.dtype, +) -> None: + """Allocate the fi_a2a_fused workspaces. Must run before CUDA-graph capture.""" + global _FI_A2A_FUSED_STATE + if _FI_A2A_FUSED_STATE is not None: + return + if cp_group.world_size == 1: + return + + import torch.distributed as dist + + try: + from flashinfer.comm import decode_cp_a2a_lse_reduce_create_workspace + except ImportError as e: + raise ImportError( + "--dcp-comm-backend fi_a2a_fused requires FlashInfer with the fused " + "DCP A2A + LSE reduce op (flashinfer #4929, unreleased as of 0.6.x); " + "could not import decode_cp_a2a_lse_reduce_create_workspace from " + "flashinfer.comm. Use --dcp-comm-backend fi_a2a or a2a otherwise." + ) from e + + cp_size = cp_group.world_size + cp_rank = cp_group.rank_in_group + parallel = get_parallel() + + # The kernel asserts one NCCL LSA domain, the same topology fi_a2a gates on. + if not is_fi_a2a_supported( + dcp_size=cp_size, + tp_size=parallel.tp_size, + pp_size=parallel.pp_size, + nnodes=parallel.nnodes, + ): + raise RuntimeError( + "--dcp-comm-backend fi_a2a_fused needs a Blackwell system whose DCP " + f"group shares one NVLink domain (got dcp_size={cp_size}, " + f"tp_size={parallel.tp_size}, pp_size={parallel.pp_size}, " + f"nnodes={parallel.nnodes}). Use --dcp-comm-backend a2a or ag_rs " + "otherwise." + ) + + if dtype not in (torch.float16, torch.bfloat16): + raise ValueError( + f"--dcp-comm-backend fi_a2a_fused supports fp16/bf16 attention output, " + f"got {dtype}. Unlike a2a (which ships the partials as opaque bytes), " + "the fused kernel reduces in-place and has no fp8 path -- use " + "--dcp-comm-backend a2a for fp8 KV-cache configurations." + ) + itemsize = torch.empty((), dtype=dtype).element_size() + if head_dim * itemsize % 16 != 0: + raise ValueError( + f"--dcp-comm-backend fi_a2a_fused needs 16-byte aligned head rows, got " + f"head_dim={head_dim} x {itemsize}B = {head_dim * itemsize}B." + ) + if cp_size > 64: + raise ValueError( + f"--dcp-comm-backend fi_a2a_fused supports at most 64 ranks, got {cp_size}." + ) + + from flashinfer.comm import decode_cp_a2a_lse_reduce_workspace_size + + nbytes = decode_cp_a2a_lse_reduce_workspace_size( + max_tokens, local_heads, cp_size, head_dim, dtype + ) + logger.info( + "fi_a2a_fused: allocating %d x %.1f MiB NCCL-symmetric workspaces " + "(max_tokens=%d, local_heads=%d, head_dim=%d, cp_size=%d)", + _FI_A2A_FUSED_NUM_STREAMS, + nbytes / (1 << 20), + max_tokens, + local_heads, + head_dim, + cp_size, + ) + workspaces = [ + decode_cp_a2a_lse_reduce_create_workspace( + max_tokens=max_tokens, + local_heads=local_heads, + cp_size=cp_size, + head_dim=head_dim, + dtype=dtype, + group=cp_group.device_group, + ) + for _ in range(_FI_A2A_FUSED_NUM_STREAMS) + ] + # No rank may publish a remote readiness word before every rank has finished. + dist.barrier(group=cp_group.device_group) + _FI_A2A_FUSED_STATE = { + "pool": workspaces, + "by_stream": {}, + "cp_rank": cp_rank, + "max_tokens": max_tokens, + } + + +def _fi_a2a_fused_workspace_for_current_stream(state: dict) -> torch.Tensor: + """Bind one pre-allocated workspace per ordered CUDA stream.""" + stream_id = int(torch.cuda.current_stream().cuda_stream) + workspace = state["by_stream"].get(stream_id) + if workspace is None: + if not state["pool"]: + raise RuntimeError( + "--dcp-comm-backend fi_a2a_fused ran out of per-stream workspaces " + f"({_FI_A2A_FUSED_NUM_STREAMS} pre-allocated, a " + f"{len(state['by_stream']) + 1}th CUDA stream appeared). Workspaces " + "cannot be allocated lazily because rendezvous is collective and " + "not graph-capturable. PDMUX captures one stream per stream group " + "and is the expected cause; raise _FI_A2A_FUSED_NUM_STREAMS to " + "1 + the number of capture streams, or use --dcp-comm-backend fi_a2a." + ) + workspace = state["pool"].pop() + state["by_stream"][stream_id] = workspace + return workspace + + +def _dcp_fi_a2a_fused_lse_reduce( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_group: "GroupCoordinator", + is_lse_base_on_e: bool = True, +) -> torch.Tensor: + """Same send layout as fi_a2a, but the LSE is its own fp32 tensor rather than + slot 0 of a padded stats pair, and the op returns the combined output.""" + from flashinfer.comm import decode_cp_a2a_lse_reduce + + state = _FI_A2A_FUSED_STATE + assert state is not None, ( + "fi_a2a_fused workspace not initialized — call " + "init_fi_a2a_fused_workspace(dcp_group, ...) at model-runner init " + "(before CUDA graph capture)." + ) + + N = cp_group.world_size + B, H, D = cp_attn_out.shape + assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})" + H_per_rank = H // N + assert B <= state["max_tokens"], ( + f"batch {B} exceeds the fi_a2a_fused workspace bound " + f"({state['max_tokens']}); the workspace is sized at init and cannot grow." + ) + + partial_o = torch.empty( + B, H_per_rank, N, D, dtype=cp_attn_out.dtype, device=cp_attn_out.device + ) + partial_lse = torch.empty( + B, H_per_rank, N, dtype=torch.float32, device=cp_attn_out.device + ) + dcp_pack_a2a_send( + cp_attn_out, + cp_attn_lse, + partial_o.permute(2, 0, 1, 3), + partial_lse.permute(2, 0, 1), + ) + + return decode_cp_a2a_lse_reduce( + partial_o, + partial_lse, + _fi_a2a_fused_workspace_for_current_stream(state), + state["cp_rank"], + N, + lse_mode="basee" if is_lse_base_on_e else "base2", + ) diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 39032b131363..7b7b798981f2 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -52,6 +52,7 @@ get_exec, get_flags, get_parallel, + get_spec, ) from sglang.srt.speculative.spec_info import create_dummy_verify_input from sglang.srt.utils import ( @@ -254,6 +255,7 @@ def warmup(self) -> None: self._pre_initialize_flashinfer_allreduce_workspace() self._pre_initialize_fi_a2a_workspace() + self._pre_initialize_fi_a2a_fused_workspace() # Model-owned communication resources may depend on the resolved # request pool and must be compiled/allocated before graph capture. @@ -316,6 +318,36 @@ def _pre_initialize_fi_a2a_workspace(self): init_fi_a2a_workspace(get_parallel().dcp_group) + def _pre_initialize_fi_a2a_fused_workspace(self): + """Allocate the fi_a2a_fused workspaces; must run before CG capture. + + Sized here rather than in layers/dcp because the bound comes from the + model geometry and the request pool. + """ + parallel = get_parallel() + if not parallel.dcp_enabled or parallel.dcp_comm_backend != "fi_a2a_fused": + return + + from sglang.srt.layers.dcp import init_fi_a2a_fused_workspace + + mr = self.model_runner + # Speculative decode widens the reduce batch to requests x draft tokens. + num_draft_tokens = ( + get_spec().speculative_num_draft_tokens + if mr.spec_algorithm.is_speculative() + else 1 + ) + head_dim = mr.model_config.kv_lora_rank + local_heads = mr.model_config.num_attention_heads // parallel.attn_tp_size + + init_fi_a2a_fused_workspace( + parallel.dcp_group, + max_tokens=mr.max_running_requests * num_draft_tokens, + local_heads=local_heads, + head_dim=head_dim, + dtype=mr.dtype, + ) + def _flashinfer_autotune(self, *, buffers, batch_size): """Run flashinfer autotune. diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index ca39935abba7..68952f2cb5b2 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -786,8 +786,10 @@ def forward_absorb_core( is_lse_base_on_e = is_mla_dcp_lse_base_on_e( self.current_attention_backend ) - if dcp_comm_backend in ("a2a", "fi_a2a"): - # A2A exchange of head partials + LSE, then local Triton combine. + if dcp_comm_backend in ("a2a", "fi_a2a", "fi_a2a_fused"): + # A2A exchange of head partials + LSE, then the merge: + # locally in Triton for a2a/fi_a2a, inside the same kernel + # for fi_a2a_fused. attn_output = dcp_a2a_lse_reduce( attn_output.contiguous(), lse.contiguous(), diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index 28ee0c1f0a95..406c450ea4fa 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -865,8 +865,10 @@ def forward_absorb_rocm_core( is_lse_base_on_e = is_mla_dcp_lse_base_on_e( self.current_attention_backend ) - if dcp_comm_backend in ("a2a", "fi_a2a"): - # A2A exchange of head partials + LSE, then local Triton combine. + if dcp_comm_backend in ("a2a", "fi_a2a", "fi_a2a_fused"): + # A2A exchange of head partials + LSE, then the merge: + # locally in Triton for a2a/fi_a2a, inside the same kernel + # for fi_a2a_fused. attn_output = dcp_a2a_lse_reduce( attn_output.contiguous(), lse.contiguous(), diff --git a/test/registered/dcp/test_dcp_fi_a2a_fused_equivalence.py b/test/registered/dcp/test_dcp_fi_a2a_fused_equivalence.py new file mode 100644 index 000000000000..ebfe70acaf9b --- /dev/null +++ b/test/registered/dcp/test_dcp_fi_a2a_fused_equivalence.py @@ -0,0 +1,213 @@ +""" +Equivalence test for the ``fi_a2a_fused`` DCP comm backend. + +``fi_a2a`` exchanges head partials with FlashInfer's MNNVL all-to-all and then +merges them with SGLang's Triton LSE kernel. ``fi_a2a_fused`` hands both halves +to FlashInfer's fused kernel (flashinfer #4929). The two must agree: this test +feeds both paths identical partials and compares their outputs. + +The fused op is unreleased (targets FlashInfer 0.7.0), so the whole module skips +when it is not importable -- see ``helix/scripts/fi4929_overlay/`` in the work +tracker for the source overlay that provides it. + +Requirements: all ranks in one NVLink/LSA domain, NCCL >= 2.29 with device +communicator support, and fp16/bf16 partials. + +Usage: + python -m pytest test_dcp_fi_a2a_fused_equivalence.py -v + +This file doubles as the torchrun worker script. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +from typing import List, Tuple + +import pytest +import torch +import torch.distributed as dist + +import sglang.srt.distributed.parallel_state as ps +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=240, + stage="extra-b", + runner_config="8-gpu-b200", +) + +# (batch, heads_per_rank, head_dim). head_dim * itemsize must be 16-byte +# aligned; 512 is the DeepSeek/Kimi MLA latent rank, the shape that matters. +TEST_SHAPES: List[Tuple[int, int, int]] = [ + (1, 2, 512), + (1, 16, 512), + (8, 16, 512), + (64, 8, 512), +] +TEST_DTYPES = [torch.bfloat16, torch.float16] +TEST_LOOP = 4 + +# The fused kernel reduces in fp32 but accumulates in a different order from the +# Triton combine, so compare at the tolerance FlashInfer's own test uses. +RTOL, ATOL = 1e-2, 1e-3 + + +def multiprocess_test(file: str, nproc: int, timeout: int = 600) -> None: + cmd = ["torchrun", f"--nproc_per_node={nproc}", file] + try: + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired as e: + raise RuntimeError( + f"torchrun (nproc={nproc}) timed out after {timeout}s\n{e.stdout}" + ) from e + assert result.returncode == 0, ( + f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n{result.stdout}" + ) + + +def _fused_available() -> bool: + try: + from flashinfer.comm import ( # noqa: F401 + decode_cp_a2a_lse_reduce, + decode_cp_a2a_lse_reduce_create_workspace, + ) + except ImportError: + return False + return True + + +@pytest.mark.parametrize("nproc", [2, 4, 8]) +def test_fi_a2a_fused_matches_fi_a2a(nproc: int) -> None: + if not _fused_available(): + pytest.skip("FlashInfer fused DCP LSE reduce (#4929) not available") + device_count = torch.cuda.device_count() + if device_count < nproc: + pytest.skip( + f"Requires at least {nproc} GPUs, but only {device_count} available" + ) + multiprocess_test(__file__, nproc) + + +def init_distributed(): + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + + dist.init_process_group(backend="gloo") + ps._WORLD = coord = ps.init_world_group( + ranks=list(range(world_size)), + local_rank=local_rank, + backend="nccl", + ) + return local_rank, device, coord + + +@torch.inference_mode() +def _run_shape( + device, coord, workspace, shape, dtype, is_lse_base_on_e, seed +) -> str | None: + """Compare the fused kernel against an all-gather + local reference merge.""" + from flashinfer.comm import decode_cp_a2a_lse_reduce + + from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton + + B, H_per_rank, D = shape + N = coord.world_size + rank = coord.rank_in_group + + # Seed per (rank, iteration): ranks must differ or the merge is trivial, and + # iterations must differ or the repeat loop re-tests one input. + torch.manual_seed(seed * 1000 + rank) + partial_o = torch.randn(B, H_per_rank, N, D, dtype=dtype, device=device) + partial_lse = torch.randn(B, H_per_rank, N, dtype=torch.float32, device=device) + + fused = decode_cp_a2a_lse_reduce( + partial_o, + partial_lse, + workspace, + rank, + N, + lse_mode="basee" if is_lse_base_on_e else "base2", + ) + + # Reference: reconstruct the all-to-all with an all-gather, take the column + # addressed to this rank, then merge with the Triton kernel fi_a2a uses. + gathered_o = [torch.empty_like(partial_o) for _ in range(N)] + gathered_lse = [torch.empty_like(partial_lse) for _ in range(N)] + dist.all_gather(gathered_o, partial_o, group=coord.device_group) + dist.all_gather(gathered_lse, partial_lse, group=coord.device_group) + + # gathered_*[src][..., rank, :] is what src sent to us. + recv_o = torch.stack([g[:, :, rank, :] for g in gathered_o], dim=0) + recv_lse = torch.stack([g[:, :, rank] for g in gathered_lse], dim=0) + + ref, _ = dcp_lse_combine_triton(recv_o, recv_lse, is_lse_base_on_e=is_lse_base_on_e) + + if not torch.allclose(fused.float(), ref.float(), rtol=RTOL, atol=ATOL): + diff = (fused.float() - ref.float()).abs().max().item() + return ( + f"shape={shape} dtype={dtype} base_e={is_lse_base_on_e}: " + f"max abs diff {diff}" + ) + return None + + +def worker_main() -> None: + from flashinfer.comm import decode_cp_a2a_lse_reduce_create_workspace + + rank, device, coord = init_distributed() + N = coord.world_size + + max_b = max(s[0] for s in TEST_SHAPES) + max_h = max(s[1] for s in TEST_SHAPES) + max_d = max(s[2] for s in TEST_SHAPES) + + for dtype in TEST_DTYPES: + # A workspace is dtype-specific and bound to one stream; build a fresh + # one per dtype rather than reusing across them. + workspace = decode_cp_a2a_lse_reduce_create_workspace( + max_tokens=max_b, + local_heads=max_h, + cp_size=N, + head_dim=max_d, + dtype=dtype, + group=coord.device_group, + ) + dist.barrier(group=coord.device_group) + + for shape in TEST_SHAPES: + for is_lse_base_on_e in (False, True): + for it in range(TEST_LOOP): + error = _run_shape( + device, coord, workspace, shape, dtype, is_lse_base_on_e, it + ) + # If any rank mismatches, every rank must fail together -- + # a one-rank raise would hang the others in the next collective. + flag = torch.tensor([int(error is not None)], device="cpu") + dist.all_reduce(flag, group=coord.cpu_group) + if flag.item(): + raise RuntimeError( + f"Rank {rank} fused/reference mismatch: {error}" + ) + + del workspace + torch.cuda.synchronize(device) + + dist.destroy_process_group() + + +if __name__ == "__main__": + if "LOCAL_RANK" in os.environ: + worker_main() + else: + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index fed3493e58c6..d070fe1594e6 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -4001,6 +4001,43 @@ def test_explicit_value_wins(self): self._resolved(dcp_size=4, dcp_comm_backend="ag_rs"), "ag_rs" ) + @override_platform(is_cuda=True, is_hip=False) + def test_fi_a2a_fused_is_explicit_only(self): + """The fused backend needs an unreleased FlashInfer op, so resolution + must never pick it on its own -- only an explicit flag selects it.""" + with patch( + "sglang.srt.arg_groups.overrides.is_fi_a2a_supported", return_value=True + ): + self.assertNotEqual(self._resolved(dcp_size=4), "fi_a2a_fused") + self.assertEqual( + self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a_fused"), + "fi_a2a_fused", + ) + + @override_platform(is_cuda=True, is_hip=False) + def test_fi_a2a_fused_requires_dcp(self): + with self.assertRaisesRegex(ValueError, "requires --dcp-size"): + self._resolved(dcp_size=1, dcp_comm_backend="fi_a2a_fused") + + @override_platform(is_cuda=False, is_hip=False) + def test_fi_a2a_fused_rejected_off_cuda(self): + with self.assertRaisesRegex(ValueError, "fi_a2a_fused"): + self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a_fused") + + @override_platform(is_cuda=True, is_hip=False) + def test_fi_a2a_fused_allows_replicate_q_proj(self): + with patch( + "sglang.srt.arg_groups.overrides.is_fi_a2a_supported", return_value=True + ): + self.assertEqual( + self._resolved( + dcp_size=4, + dcp_comm_backend="fi_a2a_fused", + dcp_replicate_q_proj=True, + ), + "fi_a2a_fused", + ) + class TestParserChoices(CustomTestCase): """The choices come from dependency-free name lists, but `cli/serve.py` From dbeafbf09fffc5ae3b99b258ce891269657d0aa4 Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Wed, 30 Sep 2026 01:02:24 +0700 Subject: [PATCH 02/10] Revert "[DCP] Add fi_a2a_fused comm backend using FlashInfer's fused A2A + LSE reduce" This reverts commit 852d6603ce03188610ca7067807b11aa240ec984. --- .../sglang/srt/arg_groups/fields/parallel.py | 11 +- python/sglang/srt/arg_groups/parallel_hook.py | 16 +- python/sglang/srt/layers/dcp/__init__.py | 2 - python/sglang/srt/layers/dcp/comm.py | 191 ---------------- .../srt/model_executor/runner/base_runner.py | 32 --- .../attention_forward_methods/forward_mla.py | 6 +- .../forward_mla_rocm.py | 6 +- .../dcp/test_dcp_fi_a2a_fused_equivalence.py | 213 ------------------ .../unit/server_args/test_server_args.py | 37 --- 9 files changed, 12 insertions(+), 502 deletions(-) delete mode 100644 test/registered/dcp/test_dcp_fi_a2a_fused_equivalence.py diff --git a/python/sglang/srt/arg_groups/fields/parallel.py b/python/sglang/srt/arg_groups/fields/parallel.py index 745e3cfb780d..1b4df5bdf70c 100644 --- a/python/sglang/srt/arg_groups/fields/parallel.py +++ b/python/sglang/srt/arg_groups/fields/parallel.py @@ -126,13 +126,10 @@ class Parallel(msgspec.Struct): help="Communication backend for the decode context-parallel (DCP) " "attention reduction: 'ag_rs' (AllGather + ReduceScatter), 'a2a' " "(fused NCCL All-to-All exchange of output+LSE + local Triton LSE " - "combine), 'fi_a2a' (FlashInfer MNNVL All-to-All kernel; requires " - "Blackwell and a DCP group within one MNNVL domain), or " - "'fi_a2a_fused' (FlashInfer fused All-to-All + LSE reduce in a single " - "kernel; same topology requirement as 'fi_a2a', fp16/bf16 only). Unset " - "resolves to 'fi_a2a' where supported, else 'a2a' on CUDA/ROCm, else " - "'ag_rs'.", - choices=["ag_rs", "a2a", "fi_a2a", "fi_a2a_fused"], + "combine), or 'fi_a2a' (FlashInfer MNNVL All-to-All kernel; requires " + "Blackwell and a DCP group within one MNNVL domain). Unset resolves " + "to 'fi_a2a' where supported, else 'a2a' on CUDA/ROCm, else 'ag_rs'.", + choices=["ag_rs", "a2a", "fi_a2a"], resolvable=True, ), ] = None diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index b503ff7cbc40..1c92d5791f57 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -170,7 +170,7 @@ def handle_decode_context_parallelism(server_args: Any): "--decode-context-parallel-size) must be >= 1, but got " f"dcp_size={cfg.dcp_size}." ) - if cfg.dcp_comm_backend in ("a2a", "fi_a2a", "fi_a2a_fused") and cfg.dcp_size <= 1: + if cfg.dcp_comm_backend in ("a2a", "fi_a2a") and cfg.dcp_size <= 1: raise ValueError( f"--dcp-comm-backend {cfg.dcp_comm_backend} only affects the " "decode context-parallel attention reduction and therefore " @@ -184,21 +184,13 @@ def handle_decode_context_parallelism(server_args: Any): "within one MNNVL domain. Use 'a2a' or 'ag_rs' elsewhere, or leave " "the flag unset to resolve it." ) - if cfg.dcp_comm_backend == "fi_a2a_fused" and not get_platform().is_cuda: - raise ValueError( - "--dcp-comm-backend fi_a2a_fused delegates both the exchange and the " - "LSE reduce to FlashInfer's fused kernel, which requires Blackwell and " - "a DCP group within one NVLink domain. Use 'a2a' or 'ag_rs' elsewhere, " - "or leave the flag unset to resolve it." - ) if cfg.dcp_replicate_q_proj: if cfg.dcp_size <= 1: raise ValueError("--dcp-replicate-q-proj requires --dcp-size > 1.") - if cfg.dcp_comm_backend not in ("a2a", "fi_a2a", "fi_a2a_fused"): + if cfg.dcp_comm_backend not in ("a2a", "fi_a2a"): raise ValueError( - "--dcp-replicate-q-proj only applies to the a2a/fi_a2a/" - "fi_a2a_fused DCP communication backend (it removes the " - "head-dim Q all-gather); " + "--dcp-replicate-q-proj only applies to the a2a/fi_a2a DCP " + "communication backend (it removes the head-dim Q all-gather); " f"got --dcp-comm-backend={cfg.dcp_comm_backend}." ) diff --git a/python/sglang/srt/layers/dcp/__init__.py b/python/sglang/srt/layers/dcp/__init__.py index 957b9bf7baab..54ae19195a1e 100644 --- a/python/sglang/srt/layers/dcp/__init__.py +++ b/python/sglang/srt/layers/dcp/__init__.py @@ -44,7 +44,6 @@ dcp_enabled, get_attention_dcp_rank, get_attention_dcp_world_size, - init_fi_a2a_fused_workspace, init_fi_a2a_workspace, ) from sglang.srt.layers.dcp.layout import ( @@ -65,7 +64,6 @@ __all__ = [ "DecodeContextParallelMetadata", "dcp_a2a_lse_reduce", - "init_fi_a2a_fused_workspace", "init_fi_a2a_workspace", "all_gather_kv_cache_for_dcp", "all_gather_kv_cache_for_mha_chunk_extend", diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index 4ad3025b3792..ae2d1422d38e 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/python/sglang/srt/layers/dcp/comm.py @@ -18,13 +18,8 @@ PR #25090 vs #14194): - cp_lse_ag_out_rs_mha: torch / natural-log logsumexp / all-reduce + head slice - cp_lse_ag_out_rs_mla: Triton (log2/exp2) correction / reduce-scatter - -The a2a family exchanges head partials instead: a2a (NCCL) and fi_a2a -(FlashInfer MNNVL) both combine locally with a Triton kernel, while -fi_a2a_fused does the exchange and the merge in one FlashInfer kernel. """ -import logging import warnings from typing import Optional @@ -45,8 +40,6 @@ from sglang.srt.utils import is_hip from sglang.srt.utils.common import is_fi_a2a_supported -logger = logging.getLogger(__name__) - _is_hip = is_hip() @@ -493,11 +486,6 @@ def dcp_a2a_lse_reduce( cp_attn_out, cp_attn_lse, cp_group, is_lse_base_on_e ) - if comm_backend == "fi_a2a_fused": - return _dcp_fi_a2a_fused_lse_reduce( - cp_attn_out, cp_attn_lse, cp_group, is_lse_base_on_e - ) - N = cp_group.world_size B, H, D = cp_attn_out.shape assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})" @@ -598,182 +586,3 @@ def _dcp_fi_a2a_lse_reduce( recv_output, recv_lse, is_lse_base_on_e=is_lse_base_on_e ) return combined - - -# fi_a2a_fused (flashinfer #4929): fuses the exchange and the LSE merge into one -# kernel. NCCL LSA symmetric memory, not the MNNVL FIFO fi_a2a uses, hence a -# separate workspace. -_FI_A2A_FUSED_STATE: Optional[dict] = None - -# Default stream + the global graph-capture stream. Rendezvous is collective and -# not capturable, so per-stream workspaces cannot be created on demand. -_FI_A2A_FUSED_NUM_STREAMS = 2 - - -def init_fi_a2a_fused_workspace( - cp_group: "GroupCoordinator", - *, - max_tokens: int, - local_heads: int, - head_dim: int, - dtype: torch.dtype, -) -> None: - """Allocate the fi_a2a_fused workspaces. Must run before CUDA-graph capture.""" - global _FI_A2A_FUSED_STATE - if _FI_A2A_FUSED_STATE is not None: - return - if cp_group.world_size == 1: - return - - import torch.distributed as dist - - try: - from flashinfer.comm import decode_cp_a2a_lse_reduce_create_workspace - except ImportError as e: - raise ImportError( - "--dcp-comm-backend fi_a2a_fused requires FlashInfer with the fused " - "DCP A2A + LSE reduce op (flashinfer #4929, unreleased as of 0.6.x); " - "could not import decode_cp_a2a_lse_reduce_create_workspace from " - "flashinfer.comm. Use --dcp-comm-backend fi_a2a or a2a otherwise." - ) from e - - cp_size = cp_group.world_size - cp_rank = cp_group.rank_in_group - parallel = get_parallel() - - # The kernel asserts one NCCL LSA domain, the same topology fi_a2a gates on. - if not is_fi_a2a_supported( - dcp_size=cp_size, - tp_size=parallel.tp_size, - pp_size=parallel.pp_size, - nnodes=parallel.nnodes, - ): - raise RuntimeError( - "--dcp-comm-backend fi_a2a_fused needs a Blackwell system whose DCP " - f"group shares one NVLink domain (got dcp_size={cp_size}, " - f"tp_size={parallel.tp_size}, pp_size={parallel.pp_size}, " - f"nnodes={parallel.nnodes}). Use --dcp-comm-backend a2a or ag_rs " - "otherwise." - ) - - if dtype not in (torch.float16, torch.bfloat16): - raise ValueError( - f"--dcp-comm-backend fi_a2a_fused supports fp16/bf16 attention output, " - f"got {dtype}. Unlike a2a (which ships the partials as opaque bytes), " - "the fused kernel reduces in-place and has no fp8 path -- use " - "--dcp-comm-backend a2a for fp8 KV-cache configurations." - ) - itemsize = torch.empty((), dtype=dtype).element_size() - if head_dim * itemsize % 16 != 0: - raise ValueError( - f"--dcp-comm-backend fi_a2a_fused needs 16-byte aligned head rows, got " - f"head_dim={head_dim} x {itemsize}B = {head_dim * itemsize}B." - ) - if cp_size > 64: - raise ValueError( - f"--dcp-comm-backend fi_a2a_fused supports at most 64 ranks, got {cp_size}." - ) - - from flashinfer.comm import decode_cp_a2a_lse_reduce_workspace_size - - nbytes = decode_cp_a2a_lse_reduce_workspace_size( - max_tokens, local_heads, cp_size, head_dim, dtype - ) - logger.info( - "fi_a2a_fused: allocating %d x %.1f MiB NCCL-symmetric workspaces " - "(max_tokens=%d, local_heads=%d, head_dim=%d, cp_size=%d)", - _FI_A2A_FUSED_NUM_STREAMS, - nbytes / (1 << 20), - max_tokens, - local_heads, - head_dim, - cp_size, - ) - workspaces = [ - decode_cp_a2a_lse_reduce_create_workspace( - max_tokens=max_tokens, - local_heads=local_heads, - cp_size=cp_size, - head_dim=head_dim, - dtype=dtype, - group=cp_group.device_group, - ) - for _ in range(_FI_A2A_FUSED_NUM_STREAMS) - ] - # No rank may publish a remote readiness word before every rank has finished. - dist.barrier(group=cp_group.device_group) - _FI_A2A_FUSED_STATE = { - "pool": workspaces, - "by_stream": {}, - "cp_rank": cp_rank, - "max_tokens": max_tokens, - } - - -def _fi_a2a_fused_workspace_for_current_stream(state: dict) -> torch.Tensor: - """Bind one pre-allocated workspace per ordered CUDA stream.""" - stream_id = int(torch.cuda.current_stream().cuda_stream) - workspace = state["by_stream"].get(stream_id) - if workspace is None: - if not state["pool"]: - raise RuntimeError( - "--dcp-comm-backend fi_a2a_fused ran out of per-stream workspaces " - f"({_FI_A2A_FUSED_NUM_STREAMS} pre-allocated, a " - f"{len(state['by_stream']) + 1}th CUDA stream appeared). Workspaces " - "cannot be allocated lazily because rendezvous is collective and " - "not graph-capturable. PDMUX captures one stream per stream group " - "and is the expected cause; raise _FI_A2A_FUSED_NUM_STREAMS to " - "1 + the number of capture streams, or use --dcp-comm-backend fi_a2a." - ) - workspace = state["pool"].pop() - state["by_stream"][stream_id] = workspace - return workspace - - -def _dcp_fi_a2a_fused_lse_reduce( - cp_attn_out: torch.Tensor, - cp_attn_lse: torch.Tensor, - cp_group: "GroupCoordinator", - is_lse_base_on_e: bool = True, -) -> torch.Tensor: - """Same send layout as fi_a2a, but the LSE is its own fp32 tensor rather than - slot 0 of a padded stats pair, and the op returns the combined output.""" - from flashinfer.comm import decode_cp_a2a_lse_reduce - - state = _FI_A2A_FUSED_STATE - assert state is not None, ( - "fi_a2a_fused workspace not initialized — call " - "init_fi_a2a_fused_workspace(dcp_group, ...) at model-runner init " - "(before CUDA graph capture)." - ) - - N = cp_group.world_size - B, H, D = cp_attn_out.shape - assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})" - H_per_rank = H // N - assert B <= state["max_tokens"], ( - f"batch {B} exceeds the fi_a2a_fused workspace bound " - f"({state['max_tokens']}); the workspace is sized at init and cannot grow." - ) - - partial_o = torch.empty( - B, H_per_rank, N, D, dtype=cp_attn_out.dtype, device=cp_attn_out.device - ) - partial_lse = torch.empty( - B, H_per_rank, N, dtype=torch.float32, device=cp_attn_out.device - ) - dcp_pack_a2a_send( - cp_attn_out, - cp_attn_lse, - partial_o.permute(2, 0, 1, 3), - partial_lse.permute(2, 0, 1), - ) - - return decode_cp_a2a_lse_reduce( - partial_o, - partial_lse, - _fi_a2a_fused_workspace_for_current_stream(state), - state["cp_rank"], - N, - lse_mode="basee" if is_lse_base_on_e else "base2", - ) diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 7b7b798981f2..39032b131363 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -52,7 +52,6 @@ get_exec, get_flags, get_parallel, - get_spec, ) from sglang.srt.speculative.spec_info import create_dummy_verify_input from sglang.srt.utils import ( @@ -255,7 +254,6 @@ def warmup(self) -> None: self._pre_initialize_flashinfer_allreduce_workspace() self._pre_initialize_fi_a2a_workspace() - self._pre_initialize_fi_a2a_fused_workspace() # Model-owned communication resources may depend on the resolved # request pool and must be compiled/allocated before graph capture. @@ -318,36 +316,6 @@ def _pre_initialize_fi_a2a_workspace(self): init_fi_a2a_workspace(get_parallel().dcp_group) - def _pre_initialize_fi_a2a_fused_workspace(self): - """Allocate the fi_a2a_fused workspaces; must run before CG capture. - - Sized here rather than in layers/dcp because the bound comes from the - model geometry and the request pool. - """ - parallel = get_parallel() - if not parallel.dcp_enabled or parallel.dcp_comm_backend != "fi_a2a_fused": - return - - from sglang.srt.layers.dcp import init_fi_a2a_fused_workspace - - mr = self.model_runner - # Speculative decode widens the reduce batch to requests x draft tokens. - num_draft_tokens = ( - get_spec().speculative_num_draft_tokens - if mr.spec_algorithm.is_speculative() - else 1 - ) - head_dim = mr.model_config.kv_lora_rank - local_heads = mr.model_config.num_attention_heads // parallel.attn_tp_size - - init_fi_a2a_fused_workspace( - parallel.dcp_group, - max_tokens=mr.max_running_requests * num_draft_tokens, - local_heads=local_heads, - head_dim=head_dim, - dtype=mr.dtype, - ) - def _flashinfer_autotune(self, *, buffers, batch_size): """Run flashinfer autotune. diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 68952f2cb5b2..ca39935abba7 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -786,10 +786,8 @@ def forward_absorb_core( is_lse_base_on_e = is_mla_dcp_lse_base_on_e( self.current_attention_backend ) - if dcp_comm_backend in ("a2a", "fi_a2a", "fi_a2a_fused"): - # A2A exchange of head partials + LSE, then the merge: - # locally in Triton for a2a/fi_a2a, inside the same kernel - # for fi_a2a_fused. + if dcp_comm_backend in ("a2a", "fi_a2a"): + # A2A exchange of head partials + LSE, then local Triton combine. attn_output = dcp_a2a_lse_reduce( attn_output.contiguous(), lse.contiguous(), diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index 406c450ea4fa..28ee0c1f0a95 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -865,10 +865,8 @@ def forward_absorb_rocm_core( is_lse_base_on_e = is_mla_dcp_lse_base_on_e( self.current_attention_backend ) - if dcp_comm_backend in ("a2a", "fi_a2a", "fi_a2a_fused"): - # A2A exchange of head partials + LSE, then the merge: - # locally in Triton for a2a/fi_a2a, inside the same kernel - # for fi_a2a_fused. + if dcp_comm_backend in ("a2a", "fi_a2a"): + # A2A exchange of head partials + LSE, then local Triton combine. attn_output = dcp_a2a_lse_reduce( attn_output.contiguous(), lse.contiguous(), diff --git a/test/registered/dcp/test_dcp_fi_a2a_fused_equivalence.py b/test/registered/dcp/test_dcp_fi_a2a_fused_equivalence.py deleted file mode 100644 index ebfe70acaf9b..000000000000 --- a/test/registered/dcp/test_dcp_fi_a2a_fused_equivalence.py +++ /dev/null @@ -1,213 +0,0 @@ -""" -Equivalence test for the ``fi_a2a_fused`` DCP comm backend. - -``fi_a2a`` exchanges head partials with FlashInfer's MNNVL all-to-all and then -merges them with SGLang's Triton LSE kernel. ``fi_a2a_fused`` hands both halves -to FlashInfer's fused kernel (flashinfer #4929). The two must agree: this test -feeds both paths identical partials and compares their outputs. - -The fused op is unreleased (targets FlashInfer 0.7.0), so the whole module skips -when it is not importable -- see ``helix/scripts/fi4929_overlay/`` in the work -tracker for the source overlay that provides it. - -Requirements: all ranks in one NVLink/LSA domain, NCCL >= 2.29 with device -communicator support, and fp16/bf16 partials. - -Usage: - python -m pytest test_dcp_fi_a2a_fused_equivalence.py -v - -This file doubles as the torchrun worker script. -""" - -from __future__ import annotations - -import os -import subprocess -import sys -from typing import List, Tuple - -import pytest -import torch -import torch.distributed as dist - -import sglang.srt.distributed.parallel_state as ps -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=240, - stage="extra-b", - runner_config="8-gpu-b200", -) - -# (batch, heads_per_rank, head_dim). head_dim * itemsize must be 16-byte -# aligned; 512 is the DeepSeek/Kimi MLA latent rank, the shape that matters. -TEST_SHAPES: List[Tuple[int, int, int]] = [ - (1, 2, 512), - (1, 16, 512), - (8, 16, 512), - (64, 8, 512), -] -TEST_DTYPES = [torch.bfloat16, torch.float16] -TEST_LOOP = 4 - -# The fused kernel reduces in fp32 but accumulates in a different order from the -# Triton combine, so compare at the tolerance FlashInfer's own test uses. -RTOL, ATOL = 1e-2, 1e-3 - - -def multiprocess_test(file: str, nproc: int, timeout: int = 600) -> None: - cmd = ["torchrun", f"--nproc_per_node={nproc}", file] - try: - result = subprocess.run( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - timeout=timeout, - ) - except subprocess.TimeoutExpired as e: - raise RuntimeError( - f"torchrun (nproc={nproc}) timed out after {timeout}s\n{e.stdout}" - ) from e - assert result.returncode == 0, ( - f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n{result.stdout}" - ) - - -def _fused_available() -> bool: - try: - from flashinfer.comm import ( # noqa: F401 - decode_cp_a2a_lse_reduce, - decode_cp_a2a_lse_reduce_create_workspace, - ) - except ImportError: - return False - return True - - -@pytest.mark.parametrize("nproc", [2, 4, 8]) -def test_fi_a2a_fused_matches_fi_a2a(nproc: int) -> None: - if not _fused_available(): - pytest.skip("FlashInfer fused DCP LSE reduce (#4929) not available") - device_count = torch.cuda.device_count() - if device_count < nproc: - pytest.skip( - f"Requires at least {nproc} GPUs, but only {device_count} available" - ) - multiprocess_test(__file__, nproc) - - -def init_distributed(): - local_rank = int(os.environ["LOCAL_RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - device = torch.device(f"cuda:{local_rank}") - torch.cuda.set_device(device) - - dist.init_process_group(backend="gloo") - ps._WORLD = coord = ps.init_world_group( - ranks=list(range(world_size)), - local_rank=local_rank, - backend="nccl", - ) - return local_rank, device, coord - - -@torch.inference_mode() -def _run_shape( - device, coord, workspace, shape, dtype, is_lse_base_on_e, seed -) -> str | None: - """Compare the fused kernel against an all-gather + local reference merge.""" - from flashinfer.comm import decode_cp_a2a_lse_reduce - - from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton - - B, H_per_rank, D = shape - N = coord.world_size - rank = coord.rank_in_group - - # Seed per (rank, iteration): ranks must differ or the merge is trivial, and - # iterations must differ or the repeat loop re-tests one input. - torch.manual_seed(seed * 1000 + rank) - partial_o = torch.randn(B, H_per_rank, N, D, dtype=dtype, device=device) - partial_lse = torch.randn(B, H_per_rank, N, dtype=torch.float32, device=device) - - fused = decode_cp_a2a_lse_reduce( - partial_o, - partial_lse, - workspace, - rank, - N, - lse_mode="basee" if is_lse_base_on_e else "base2", - ) - - # Reference: reconstruct the all-to-all with an all-gather, take the column - # addressed to this rank, then merge with the Triton kernel fi_a2a uses. - gathered_o = [torch.empty_like(partial_o) for _ in range(N)] - gathered_lse = [torch.empty_like(partial_lse) for _ in range(N)] - dist.all_gather(gathered_o, partial_o, group=coord.device_group) - dist.all_gather(gathered_lse, partial_lse, group=coord.device_group) - - # gathered_*[src][..., rank, :] is what src sent to us. - recv_o = torch.stack([g[:, :, rank, :] for g in gathered_o], dim=0) - recv_lse = torch.stack([g[:, :, rank] for g in gathered_lse], dim=0) - - ref, _ = dcp_lse_combine_triton(recv_o, recv_lse, is_lse_base_on_e=is_lse_base_on_e) - - if not torch.allclose(fused.float(), ref.float(), rtol=RTOL, atol=ATOL): - diff = (fused.float() - ref.float()).abs().max().item() - return ( - f"shape={shape} dtype={dtype} base_e={is_lse_base_on_e}: " - f"max abs diff {diff}" - ) - return None - - -def worker_main() -> None: - from flashinfer.comm import decode_cp_a2a_lse_reduce_create_workspace - - rank, device, coord = init_distributed() - N = coord.world_size - - max_b = max(s[0] for s in TEST_SHAPES) - max_h = max(s[1] for s in TEST_SHAPES) - max_d = max(s[2] for s in TEST_SHAPES) - - for dtype in TEST_DTYPES: - # A workspace is dtype-specific and bound to one stream; build a fresh - # one per dtype rather than reusing across them. - workspace = decode_cp_a2a_lse_reduce_create_workspace( - max_tokens=max_b, - local_heads=max_h, - cp_size=N, - head_dim=max_d, - dtype=dtype, - group=coord.device_group, - ) - dist.barrier(group=coord.device_group) - - for shape in TEST_SHAPES: - for is_lse_base_on_e in (False, True): - for it in range(TEST_LOOP): - error = _run_shape( - device, coord, workspace, shape, dtype, is_lse_base_on_e, it - ) - # If any rank mismatches, every rank must fail together -- - # a one-rank raise would hang the others in the next collective. - flag = torch.tensor([int(error is not None)], device="cpu") - dist.all_reduce(flag, group=coord.cpu_group) - if flag.item(): - raise RuntimeError( - f"Rank {rank} fused/reference mismatch: {error}" - ) - - del workspace - torch.cuda.synchronize(device) - - dist.destroy_process_group() - - -if __name__ == "__main__": - if "LOCAL_RANK" in os.environ: - worker_main() - else: - sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index d070fe1594e6..fed3493e58c6 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -4001,43 +4001,6 @@ def test_explicit_value_wins(self): self._resolved(dcp_size=4, dcp_comm_backend="ag_rs"), "ag_rs" ) - @override_platform(is_cuda=True, is_hip=False) - def test_fi_a2a_fused_is_explicit_only(self): - """The fused backend needs an unreleased FlashInfer op, so resolution - must never pick it on its own -- only an explicit flag selects it.""" - with patch( - "sglang.srt.arg_groups.overrides.is_fi_a2a_supported", return_value=True - ): - self.assertNotEqual(self._resolved(dcp_size=4), "fi_a2a_fused") - self.assertEqual( - self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a_fused"), - "fi_a2a_fused", - ) - - @override_platform(is_cuda=True, is_hip=False) - def test_fi_a2a_fused_requires_dcp(self): - with self.assertRaisesRegex(ValueError, "requires --dcp-size"): - self._resolved(dcp_size=1, dcp_comm_backend="fi_a2a_fused") - - @override_platform(is_cuda=False, is_hip=False) - def test_fi_a2a_fused_rejected_off_cuda(self): - with self.assertRaisesRegex(ValueError, "fi_a2a_fused"): - self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a_fused") - - @override_platform(is_cuda=True, is_hip=False) - def test_fi_a2a_fused_allows_replicate_q_proj(self): - with patch( - "sglang.srt.arg_groups.overrides.is_fi_a2a_supported", return_value=True - ): - self.assertEqual( - self._resolved( - dcp_size=4, - dcp_comm_backend="fi_a2a_fused", - dcp_replicate_q_proj=True, - ), - "fi_a2a_fused", - ) - class TestParserChoices(CustomTestCase): """The choices come from dependency-free name lists, but `cli/serve.py` From c3992a2a8b6051d87c2e030504a7a2022af6e772 Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Wed, 30 Sep 2026 01:16:12 +0700 Subject: [PATCH 03/10] [DCP] Resolve fi_a2a only where the fused FlashInfer reduce can run fi_a2a is moving to FlashInfer's decode_cp_a2a_lse_reduce, which needs torch >= 2.14, a FlashInfer build that ships the op, an MLA model and a DCP group inside one NVLink domain. It has no fallback kernel, so an unset --dcp-comm-backend must only resolve to fi_a2a where all of that holds, and otherwise to a2a. fi_a2a_platform_blocker reports the host-side reason without importing FlashInfer or creating a CUDA context; fi_a2a_blocker adds the configuration checks (MLA, PDMUX, torch symm-mem all-reduce, the PP DeepGEMM warmup stream, an exported NCCL_CUMEM_ENABLE=0). An explicit fi_a2a names the blocker, and otherwise defaults NCCL_CUMEM_ENABLE=1 because the fused workspace is an NCCL symmetric window. --- .../sglang/srt/arg_groups/fields/parallel.py | 10 ++- python/sglang/srt/arg_groups/overrides.py | 47 ++++++++++-- python/sglang/srt/arg_groups/parallel_hook.py | 18 +++-- python/sglang/srt/utils/common.py | 45 +++++++++++ .../unit/server_args/test_server_args.py | 74 +++++++++++++++---- 5 files changed, 161 insertions(+), 33 deletions(-) diff --git a/python/sglang/srt/arg_groups/fields/parallel.py b/python/sglang/srt/arg_groups/fields/parallel.py index 3e0e28e6d4c8..a3d903f6d94a 100644 --- a/python/sglang/srt/arg_groups/fields/parallel.py +++ b/python/sglang/srt/arg_groups/fields/parallel.py @@ -119,9 +119,13 @@ class Parallel(msgspec.Struct): help="Communication backend for the decode context-parallel (DCP) " "attention reduction: 'ag_rs' (AllGather + ReduceScatter), 'a2a' " "(fused NCCL All-to-All exchange of output+LSE + local Triton LSE " - "combine), or 'fi_a2a' (FlashInfer MNNVL All-to-All kernel; requires " - "Blackwell and a DCP group within one MNNVL domain). Unset resolves " - "to 'fi_a2a' where supported, else 'a2a' on CUDA/ROCm, else 'ag_rs'.", + "combine), or 'fi_a2a' (FlashInfer's decode_cp_a2a_lse_reduce: " + "exchange and LSE merge in one kernel over NCCL symmetric memory; " + "requires Blackwell, the DCP group inside one NVLink domain, an MLA " + "model, torch>=2.14 and a FlashInfer build with that op; it moves " + "torch symmetric memory to the NCCL backend, so custom all-reduce " + "uses its v1 kernel). Unset resolves to 'fi_a2a' where it can run, " + "else 'a2a' on CUDA/ROCm, else 'ag_rs'.", choices=["ag_rs", "a2a", "fi_a2a"], resolvable=True, ), diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index a2de0bc2937b..27348e6bc2fd 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -24,6 +24,7 @@ import json import logging import math +import os from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from sglang.srt.arg_groups import model_override_base @@ -66,8 +67,8 @@ get_platform, ) from sglang.srt.utils.common import ( + fi_a2a_platform_blocker, get_quantization_config, - is_fi_a2a_supported, is_gfx95_supported, xpu_has_xmx_support, ) @@ -1309,6 +1310,39 @@ def _data_parallelism_defaults(view: Any) -> dict: return {} +def fi_a2a_blocker(view: Any) -> Optional[str]: + """Why --dcp-comm-backend fi_a2a cannot run with this configuration, or None.""" + reason = fi_a2a_platform_blocker( + dcp_size=view.dcp_size, + tp_size=view.tp_size, + pp_size=view.pp_size, + nnodes=view.nnodes, + ) + if reason is not None: + return reason + if not use_mla_backend(view): + return "applies only to MLA models (the MHA DCP reduction ignores it)" + if view.enable_pdmux: + return ( + "does not support --enable-pdmux (every captured stream group would " + "need its own fused-reduce workspace)" + ) + if view.enable_torch_symm_mem: + return ( + "cannot be combined with --enable-torch-symm-mem (fi_a2a moves torch " + "symmetric memory to the NCCL backend, which that all-reduce cannot use)" + ) + if envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.get() and view.pp_size > 1: + return ( + "cannot be combined with SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP at " + "--pp-size > 1 (its decode warmup runs on a third CUDA stream)" + ) + # The fused workspace is an NCCL symmetric window, which needs cuMem memory. + if os.environ.get("NCCL_CUMEM_ENABLE") == "0": + return "needs cuMem-backed NCCL memory, but NCCL_CUMEM_ENABLE=0 is exported" + return None + + @register_post_process def _dcp_comm_backend_default(view: Any) -> dict: if view.dcp_comm_backend is not None: @@ -1316,21 +1350,18 @@ def _dcp_comm_backend_default(view: Any) -> dict: if view.dcp_size <= 1: return {"dcp_comm_backend": "ag_rs"} platform = get_platform() - if is_fi_a2a_supported( - dcp_size=view.dcp_size, - tp_size=view.tp_size, - pp_size=view.pp_size, - nnodes=view.nnodes, - ): + blocker = fi_a2a_blocker(view) + if blocker is None: backend = "fi_a2a" elif platform.is_cuda or platform.is_hip: backend = "a2a" else: backend = "ag_rs" logger.info( - "DCP (dcp_size=%d) selects communication backend %r.", + "DCP (dcp_size=%d) selects communication backend %r%s.", view.dcp_size, backend, + f" (fi_a2a {blocker})" if blocker is not None and platform.is_cuda else "", ) return {"dcp_comm_backend": backend} diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index 1a396935fbf1..53325280e3fe 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -13,6 +13,7 @@ _dp_lm_head_validation, _tp_lm_head_all_to_all_default, declare_resolution, + fi_a2a_blocker, model_config_of, resolved_view, resolving_view, @@ -205,13 +206,16 @@ def handle_decode_context_parallelism(server_args: Any): "requires --dcp-size / --decode-context-parallel-size > 1, but " f"got dcp_size={cfg.dcp_size}." ) - if cfg.dcp_comm_backend == "fi_a2a" and not get_platform().is_cuda: - raise ValueError( - "--dcp-comm-backend fi_a2a delegates the exchange to FlashInfer's " - "MNNVL All-to-All kernel, which requires Blackwell and a DCP group " - "within one MNNVL domain. Use 'a2a' or 'ag_rs' elsewhere, or leave " - "the flag unset to resolve it." - ) + if cfg.dcp_comm_backend == "fi_a2a": + blocker = fi_a2a_blocker(cfg) + if blocker is not None: + raise ValueError( + f"--dcp-comm-backend fi_a2a {blocker}. Use --dcp-comm-backend a2a, " + "or leave the flag unset to resolve a supported backend." + ) + # The fused reduce's workspace is an NCCL symmetric window, which needs + # cuMem-backed NCCL memory; set before any communicator exists. + os.environ.setdefault("NCCL_CUMEM_ENABLE", "1") if cfg.dcp_replicate_q_proj: if cfg.dcp_size <= 1: raise ValueError("--dcp-replicate-q-proj requires --dcp-size > 1.") diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 67bf51b698e4..0f06abe7d8b0 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -957,6 +957,39 @@ def is_mnnvl_fabric_device() -> bool: return any(tag in name for tag in ("GB200", "GB300")) +def fi_a2a_platform_blocker( + *, dcp_size: int, tp_size: int, pp_size: int, nnodes: int +) -> Optional[str]: + """Why this host cannot run the fused fi_a2a DCP reduce, or None if it can. + + Runs at argument resolution: no CUDA context and no flashinfer import. + """ + if not get_platform().is_sm100: + return "requires a Blackwell (SM100-family) GPU" + if not is_mnnvl_fabric_device(): + tp_size_per_node = tp_size // max(nnodes // pp_size, 1) + if tp_size_per_node % dcp_size != 0: + return ( + "requires the DCP group inside one NVLink domain (MNNVL fabric or " + f"one node); dcp_size={dcp_size} does not divide the " + f"{tp_size_per_node} TP ranks per node (tp_size={tp_size}, " + f"pp_size={pp_size}, nnodes={nnodes})" + ) + # torch 2.14 is the first release whose symmetric memory exposes the NCCL + # device communicator and window offsets the fused kernel is built on. + if torch_release < (2, 14): + return ( + "requires torch>=2.14 (NCCL device communicator in torch symmetric " + f"memory); found torch {torch.__version__}" + ) + if not _flashinfer_has_fused_dcp_reduce(): + return ( + "requires a FlashInfer build that provides " + "flashinfer.comm.decode_cp_a2a_lse_reduce" + ) + return None + + def is_fi_a2a_supported( *, dcp_size: int, tp_size: int, pp_size: int, nnodes: int ) -> bool: @@ -968,6 +1001,18 @@ def is_fi_a2a_supported( return tp_size_per_node % dcp_size == 0 +@lru_cache(maxsize=1) +def _flashinfer_has_fused_dcp_reduce() -> bool: + # A top-level find_spec locates the package without importing it. + spec = find_spec("flashinfer") + if spec is None or spec.submodule_search_locations is None: + return False + return any( + os.path.isfile(os.path.join(path, "comm", "dcp_lse_reduce.py")) + for path in spec.submodule_search_locations + ) + + @lru_cache(maxsize=1) def is_habana_available() -> bool: return find_spec("habana_frameworks") is not None diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 67010b13066f..50acb5b622e9 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -3996,40 +3996,84 @@ def test_without_all_to_all_env_is_untouched(self): class TestDcpCommBackendDefault(CustomTestCase): + def _args(self, **fields): + return ServerArgs(model_path="dummy", tp_size=8, **fields) + def _resolved(self, **fields): - args = ServerArgs(model_path="dummy", tp_size=8, **fields) + args = self._args(**fields) parallel_hook.handle_decode_context_parallelism(args) return resolution_result(args, "dcp_comm_backend") + def _fused_reduce(self, *, platform_blocker=None, mla=True): + """Patch what fi_a2a needs: the host check and the model's attention arch.""" + return ( + patch( + "sglang.srt.arg_groups.overrides.fi_a2a_platform_blocker", + return_value=platform_blocker, + ), + patch("sglang.srt.arg_groups.overrides.use_mla_backend", return_value=mla), + ) + def test_no_dcp_is_ag_rs(self): self.assertEqual(self._resolved(dcp_size=1), "ag_rs") @override_platform(is_cuda=True, is_hip=False) - def test_fi_a2a_where_supported(self): - with patch( - "sglang.srt.arg_groups.overrides.is_fi_a2a_supported", return_value=True - ): + def test_fi_a2a_where_the_fused_reduce_can_run(self): + host, mla = self._fused_reduce() + with host, mla: self.assertEqual(self._resolved(dcp_size=4), "fi_a2a") @override_platform(is_cuda=True, is_hip=False) - def test_a2a_on_cuda_without_mnnvl(self): - with patch( - "sglang.srt.arg_groups.overrides.is_fi_a2a_supported", return_value=False - ): + def test_a2a_when_the_fused_reduce_cannot_run(self): + """fi_a2a has no fallback kernel, so resolving to it on a host that + cannot run the fused reduce would fail every default DCP start.""" + host, mla = self._fused_reduce(platform_blocker="requires torch>=2.14") + with host, mla: self.assertEqual(self._resolved(dcp_size=4), "a2a") + @override_platform(is_cuda=True, is_hip=False) + def test_configs_the_fused_reduce_cannot_serve_resolve_to_a2a(self): + cases = { + "non-MLA model": ({}, {}, False), + "--enable-pdmux": ({"enable_pdmux": True}, {}, True), + "--enable-torch-symm-mem": ({"enable_torch_symm_mem": True}, {}, True), + "exported NCCL_CUMEM_ENABLE=0": ({}, {"NCCL_CUMEM_ENABLE": "0"}, True), + } + for name, (fields, env, is_mla) in cases.items(): + with self.subTest(name): + host, mla = self._fused_reduce(mla=is_mla) + with host, mla, patch.dict(os.environ, env): + self.assertEqual(self._resolved(dcp_size=4, **fields), "a2a") + with self.assertRaisesRegex(ValueError, "--dcp-comm-backend a2a"): + self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a", **fields) + + @override_platform(is_cuda=True, is_hip=False) + def test_explicit_fi_a2a_names_the_blocker(self): + host, mla = self._fused_reduce(platform_blocker="requires torch>=2.14") + with host, mla: + with self.assertRaisesRegex(ValueError, "fi_a2a requires torch>=2.14"): + self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a") + + @override_platform(is_cuda=True, is_hip=False) + def test_explicit_fi_a2a_defaults_nccl_cumem_on(self): + """The fused workspace is an NCCL symmetric window; engine startup would + otherwise force NCCL_CUMEM_ENABLE=0 for a non-symm-mem server.""" + host, mla = self._fused_reduce() + with host, mla, patch.dict(os.environ): + os.environ.pop("NCCL_CUMEM_ENABLE", None) + self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a") + self.assertEqual(os.environ["NCCL_CUMEM_ENABLE"], "1") + @override_platform(is_cuda=False, is_hip=False) def test_ag_rs_off_cuda(self): - with patch( - "sglang.srt.arg_groups.overrides.is_fi_a2a_supported", return_value=False - ): + host, mla = self._fused_reduce(platform_blocker="requires a Blackwell GPU") + with host, mla: self.assertEqual(self._resolved(dcp_size=4), "ag_rs") @override_platform(is_cuda=True, is_hip=False) def test_explicit_value_wins(self): - with patch( - "sglang.srt.arg_groups.overrides.is_fi_a2a_supported", return_value=True - ): + host, mla = self._fused_reduce() + with host, mla: self.assertEqual( self._resolved(dcp_size=4, dcp_comm_backend="ag_rs"), "ag_rs" ) From e7b03c3ca8eaa9bd037992800c6a1727302aca01 Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Wed, 30 Sep 2026 01:16:12 +0700 Subject: [PATCH 04/10] [DCP] Run fi_a2a on FlashInfer's fused all-to-all + LSE reduce fi_a2a exchanged DCP head partials with FlashInfer's MNNVL all-to-all and then merged them with a separate Triton LSE kernel, so every MLA layer round-tripped the partials through HBM between two launches. decode_cp_a2a_lse_reduce does the exchange and the LSE-weighted merge in one cooperative kernel over NCCL LSA symmetric memory; fi_a2a now uses it exclusively. Before graph capture, every DCP rank agrees over gloo that the op can run, then creates two workspaces (one per ordered CUDA stream: graph capture and serving) and checks one probe call per LSE base against an all-gather + Triton reference. That probe also makes the group's first, collective device-comm setup eager. Batches wider than the captured graphs run as row chunks on the same stream and workspace. The op is wrapped as a custom op so Dynamo does not trace the stream-to-workspace lookup. The workspace comes from torch symmetric memory on the NCCL backend, a process-global choice frozen by the first symmetric allocation, so bootstrap selects it before the groups are built. CustomAllReduceV2 cannot allocate from that backend (it passes a group name and rendezvouses on gloo), so it yields to CustomAllreduce there. --- python/sglang/srt/distributed/bootstrap.py | 21 + .../device_communicators/custom_all_reduce.py | 11 +- python/sglang/srt/layers/dcp/comm.py | 399 +++++++++++++----- .../srt/model_executor/runner/base_runner.py | 30 +- python/sglang/srt/utils/common.py | 11 - .../ops/attention/test_dcp_lse_combine.py | 16 +- .../communication/test_dcp_fi_a2a_reduce.py | 198 +++++++++ 7 files changed, 558 insertions(+), 128 deletions(-) create mode 100644 test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index 61e05396950d..c1cb401cb769 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/python/sglang/srt/distributed/bootstrap.py @@ -145,6 +145,12 @@ def init_parallel_runtime( local_omp_cpuid=local_omp_cpuid, dist_init_method=dist_init_method, ) + if ( + device == "cuda" + and parallel.dcp_enabled + and parallel.dcp_comm_backend == "fi_a2a" + ): + _select_nccl_symm_mem_backend() _init_parallel_groups( backend=backend, @@ -285,6 +291,21 @@ def _init_cpu_threads_env( ) +def _select_nccl_symm_mem_backend() -> None: + # DCP fi_a2a's fused reduce allocates torch symmetric memory on the NCCL + # backend. The choice is process-global and frozen by the first symmetric + # allocation, which custom all-reduce makes while the groups are built. + import torch.distributed._symmetric_memory as symm_mem + + try: + symm_mem.set_backend("NCCL") + except RuntimeError as e: + raise RuntimeError( + "--dcp-comm-backend fi_a2a needs torch's NCCL symmetric-memory " + f"backend, which could not be selected before process-group init: {e}" + ) from e + + def _init_parallel_groups( *, backend: str, diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index 55c292cde30c..7632a986718c 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -356,12 +356,21 @@ def dispatch_custom_allreduce( can_use_custom_all_reduce_v2); other cross-node groups fall back to NCCL. """ if _is_cuda and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get(): + import torch.distributed._symmetric_memory as torch_symm_mem + from .custom_all_reduce_v2 import ( CustomAllReduceV2, can_use_custom_all_reduce_v2, ) - if can_use_custom_all_reduce_v2(group=group, device=device): + if torch_symm_mem.get_backend(device) == "NCCL": + # CustomAllReduceV2 allocates torch symmetric memory with a group name + # and rendezvouses on the gloo group; the NCCL allocator rejects both. + logger.info( + "[AR] torch symmetric memory uses the NCCL backend (DCP fi_a2a); " + "CustomAllReduceV2 cannot allocate from it, using CustomAllreduce" + ) + elif can_use_custom_all_reduce_v2(group=group, device=device): logger.debug("[AR] Using CustomAllReduceV2 (JIT-compiled)") return CustomAllReduceV2 diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index b06a49cbe0a3..1ab4e71f89c7 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/python/sglang/srt/layers/dcp/comm.py @@ -20,9 +20,12 @@ - cp_lse_ag_out_rs_mla: Triton (log2/exp2) correction / reduce-scatter """ +import logging from typing import Optional +import msgspec import torch +import torch.distributed as dist from sglang.kernels.ops.attention.dcp_kernels import ( CPTritonContext, @@ -34,10 +37,16 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) +from sglang.srt.distributed.device_communicators.pynccl_wrapper import ( + NCCLLibrary, + find_nccl_library, +) from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import is_hip -from sglang.srt.utils.common import is_fi_a2a_supported +from sglang.srt.utils.custom_op import register_custom_op + +logger = logging.getLogger(__name__) _is_hip = is_hip() @@ -350,88 +359,296 @@ def all_gather_kv_cache_for_dcp( # --------------------------------------------------------------------------- # A2A communication backend for DCP decode (alternative to AG+RS above): exchange -# per-head partial outputs + LSEs across DCP ranks, then combine locally with the -# Triton LSE kernel. fi_a2a delegates the exchange to FlashInfer MNNVL (#2951). +# per-head partial outputs + LSEs across DCP ranks and merge them. a2a sends one +# NCCL all-to-all and merges with the Triton LSE kernel; fi_a2a hands exchange and +# merge to FlashInfer's fused decode_cp_a2a_lse_reduce (one cooperative kernel over +# NCCL LSA symmetric memory). # --------------------------------------------------------------------------- -# Per-process singleton: MNNVL workspace + this rank's cp position. Populated -# once, pre-CUDA-graph-capture, by init_fi_a2a_workspace(). -_FI_A2A_STATE: Optional[dict] = None +class _FiA2aState(msgspec.Struct, kw_only=True): + cp_rank: int + cp_size: int + # One fused call carries at most this many rows (product of the leading dims). + capacity_rows: int + # Created before graph capture; each binds to the first stream that uses it. + unbound: list + by_stream: dict + + +# Per-process singleton, populated once before CUDA graph capture by +# init_fi_a2a_workspace(); holds strong references to the workspaces. +_FI_A2A_STATE: Optional[_FiA2aState] = None -def init_fi_a2a_workspace(cp_group: "GroupCoordinator") -> None: - # Call once per process BEFORE CUDA-graph capture: the FlashInfer init syncs - # the stream and barriers cross-rank, neither of which is capturable. +# A workspace may serve only one ordered CUDA stream: the graph-capture stream +# (bound by the startup probe) and the serving stream (bound on first use). +_FI_A2A_NUM_WORKSPACES = 2 + + +def init_fi_a2a_workspace( + cp_group: GroupCoordinator, + *, + max_tokens: int, + local_heads: int, + head_dim: int, + dtype: torch.dtype, + probe_stream: torch.cuda.Stream, +) -> None: + """Create the fused fi_a2a workspaces on every DCP rank; collective, run + before graph capture (rendezvous and the first reduce are not capturable).""" global _FI_A2A_STATE - if _FI_A2A_STATE is not None: - return - if cp_group.world_size == 1: + if _FI_A2A_STATE is not None or cp_group.world_size == 1: return + # The kernel's receive spin-wait has no timeout, so every rank must agree to + # proceed before the NCCL collectives below; the verdict travels over gloo. + _fi_a2a_require_all_ranks( + cp_group=cp_group, + stage="eligibility check", + error=_fi_a2a_local_blocker(cp_group=cp_group, head_dim=head_dim, dtype=dtype), + ) + from flashinfer.comm import ( + decode_cp_a2a_lse_reduce_create_workspace, + decode_cp_a2a_lse_reduce_workspace_size, + ) + + # Rendezvous finds the group's NCCL communicator by name, and ProcessGroupNCCL + # creates it lazily: one collective on the device group makes it exist. + dist.barrier(group=cp_group.device_group) + workspaces = [ + decode_cp_a2a_lse_reduce_create_workspace( + max_tokens=max_tokens, + local_heads=local_heads, + cp_size=cp_group.world_size, + head_dim=head_dim, + dtype=dtype, + group=cp_group.device_group, + ) + for _ in range(_FI_A2A_NUM_WORKSPACES) + ] + dist.barrier(group=cp_group.device_group) + state = _FiA2aState( + cp_rank=cp_group.rank_in_group, + cp_size=cp_group.world_size, + capacity_rows=max_tokens * local_heads, + unbound=workspaces[1:], + by_stream={probe_stream.cuda_stream: workspaces[0]}, + ) + with torch.cuda.stream(probe_stream): + error = _fi_a2a_probe( + state=state, + cp_group=cp_group, + local_heads=local_heads, + head_dim=head_dim, + dtype=dtype, + ) + _fi_a2a_require_all_ranks(cp_group=cp_group, stage="startup probe", error=error) + _FI_A2A_STATE = state + workspace_mib = ( + decode_cp_a2a_lse_reduce_workspace_size( + max_tokens=max_tokens, + local_heads=local_heads, + cp_size=cp_group.world_size, + head_dim=head_dim, + dtype=dtype, + ) + / 2**20 + ) + logger.info( + "DCP fi_a2a: fused all-to-all + LSE reduce ready (rank %d/%d, %d heads/rank, " + "head_dim %d, %s, capacity %d rows, %d workspaces x %.1f MiB)", + state.cp_rank, + state.cp_size, + local_heads, + head_dim, + dtype, + state.capacity_rows, + _FI_A2A_NUM_WORKSPACES, + workspace_mib, + ) - import torch.distributed as dist +def _fi_a2a_local_blocker( + *, cp_group: GroupCoordinator, head_dim: int, dtype: torch.dtype +) -> Optional[str]: try: - from flashinfer.comm.dcp_alltoall import ( - decode_cp_a2a_allocate_mnnvl_workspace, - decode_cp_a2a_init_workspace, - ) - from flashinfer.comm.mapping import Mapping - from flashinfer.comm.mnnvl import MnnvlConfig + import flashinfer.comm.dcp_lse_reduce # noqa: F401 except ImportError as e: - raise ImportError( - "--dcp-comm-backend fi_a2a requires FlashInfer with the DCP " - "all-to-all kernel (flashinfer #2951); could not import " - "flashinfer.comm.dcp_alltoall." - ) from e - - # Reuse the MoE adapter: its Split() returns a CommBackend (what FlashInfer's - # Mapping expects); the flashinfer_comm_fusion copy has drifted to return a - # raw ProcessGroup, so don't swap without re-checking the Split() contract. - from sglang.srt.layers.moe.token_dispatcher.flashinfer_utils import ( - TorchDistributedCommBackend, - ) + return f"FlashInfer lacks flashinfer.comm.decode_cp_a2a_lse_reduce ({e})" + import torch.distributed._symmetric_memory as symm_mem + + backend = symm_mem.get_backend(torch.device("cuda", torch.cuda.current_device())) + if backend != "NCCL": + return ( + f"torch symmetric memory uses the {backend} backend; the fused reduce " + "needs NCCL, which bootstrap selects before process-group init" + ) + # The op is built against the pip NCCL headers; an older loaded libnccl would + # abort at the first call instead of raising. + nccl_version = NCCLLibrary(find_nccl_library()).ncclGetRawVersion() + if nccl_version < 22900: + return f"the loaded NCCL is {nccl_version}; the fused reduce needs >= 2.29" + if dtype not in (torch.float16, torch.bfloat16): + return f"attention partials are {dtype}; the fused reduce supports float16 and bfloat16" + if head_dim * dtype.itemsize % 16 != 0: + return f"head_dim {head_dim} rows of {dtype} are not 16-byte aligned" + if cp_group.world_size > 64: + return f"dcp_size {cp_group.world_size} exceeds the fused reduce's 64 ranks" + if dist.get_rank(cp_group.device_group) != cp_group.rank_in_group: + return "the DCP device-group rank differs from rank_in_group" + return None + + +def _fi_a2a_require_all_ranks( + *, cp_group: GroupCoordinator, stage: str, error: Optional[str] +) -> None: + errors: list = [None] * cp_group.world_size + dist.all_gather_object(errors, error, group=cp_group.cpu_group) + failed = [(rank, e) for rank, e in enumerate(errors) if e is not None] + if failed: + detail = "; ".join(f"rank {rank}: {e}" for rank, e in failed[:4]) + raise RuntimeError( + f"--dcp-comm-backend fi_a2a {stage} failed on {len(failed)}/" + f"{cp_group.world_size} DCP ranks: {detail}. " + "Use --dcp-comm-backend a2a." + ) - cp_size = cp_group.world_size - cp_rank = cp_group.rank_in_group - parallel = get_parallel() - if not is_fi_a2a_supported( - dcp_size=cp_size, - tp_size=parallel.tp_size, - pp_size=parallel.pp_size, - nnodes=parallel.nnodes, - ): +def _fi_a2a_probe( + *, + state: _FiA2aState, + cp_group: GroupCoordinator, + local_heads: int, + head_dim: int, + dtype: torch.dtype, +) -> Optional[str]: + """Check the fused reduce against all-gather + Triton merge for both LSE bases. + + Its first call also builds the group's NCCL device communicator (collective), + which must happen eagerly, before any graph capture. + """ + rank, n = cp_group.rank_in_group, cp_group.world_size + gen = torch.Generator(device="cuda").manual_seed(4929 + rank) + error = None + for is_lse_base_on_e in (False, True): + partial_o = torch.randn( + 2, local_heads, n, head_dim, dtype=dtype, device="cuda", generator=gen + ) + partial_lse = torch.randn( + 2, local_heads, n, dtype=torch.float32, device="cuda", generator=gen + ) + # The reference collectives run on every rank and every iteration, even + # after a local failure, so no rank waits in a collective the others skip. + gathered_o = cp_group.all_gather(partial_o.unsqueeze(0), dim=0) + gathered_lse = cp_group.all_gather(partial_lse.unsqueeze(0), dim=0) + expected, _ = dcp_lse_combine_triton( + gathered_o[:, :, :, rank].contiguous(), + gathered_lse[:, :, :, rank].contiguous(), + is_lse_base_on_e=is_lse_base_on_e, + ) + if error is not None: + continue + try: + got = _fi_a2a_fused_reduce_impl( + state=state, + partial_o=partial_o, + partial_lse=partial_lse, + is_lse_base_on_e=is_lse_base_on_e, + ) + torch.cuda.current_stream().synchronize() + except RuntimeError as e: + error = str(e).splitlines()[0] + continue + if not torch.allclose(got.float(), expected.float(), rtol=1e-2, atol=1e-2): + max_diff = (got.float() - expected.float()).abs().max().item() + base = "e" if is_lse_base_on_e else "2" + error = ( + f"fused reduce disagrees with the all-gather reference " + f"(LSE base {base}, max abs diff {max_diff:.3g})" + ) + return error + + +def _fi_a2a_workspace_for_current_stream(state: _FiA2aState) -> torch.Tensor: + stream = torch.cuda.current_stream().cuda_stream + workspace = state.by_stream.get(stream) + if workspace is not None: + return workspace + if not state.unbound: + bound = ", ".join(hex(s) for s in state.by_stream) raise RuntimeError( - "--dcp-comm-backend fi_a2a needs a Blackwell system whose DCP group " - "shares one MNNVL domain: either MNNVL fabric memory (GB200/GB300) " - f"or a DCP group inside one node (got dcp_size={cp_size}, " - f"tp_size={parallel.tp_size}, pp_size={parallel.pp_size}, " - f"nnodes={parallel.nnodes}). Use --dcp-comm-backend a2a or ag_rs " - "otherwise." + f"DCP fi_a2a: CUDA stream {stream:#x} is a third stream calling the " + f"fused reduce; its workspaces are bound to streams {bound}. Each " + "ordered stream needs its own workspace, created before graph capture." + ) + workspace = state.unbound.pop(0) + state.by_stream[stream] = workspace + return workspace + + +def _fi_a2a_fused_reduce_impl( + *, + state: _FiA2aState, + partial_o: torch.Tensor, + partial_lse: torch.Tensor, + is_lse_base_on_e: bool, +) -> torch.Tensor: + from flashinfer.comm import decode_cp_a2a_lse_reduce + + workspace = _fi_a2a_workspace_for_current_stream(state) + lse_mode = "basee" if is_lse_base_on_e else "base2" + batch, heads, n, head_dim = partial_o.shape + rows = batch * heads + if rows <= state.capacity_rows: + return decode_cp_a2a_lse_reduce( + partial_o=partial_o, + partial_lse=partial_lse, + workspace=workspace, + cp_rank=state.cp_rank, + cp_size=state.cp_size, + lse_mode=lse_mode, + ) + # An eager batch wider than the graph capture size: the op flattens leading + # dims into rows, so reduce capacity-sized row chunks in order on this stream. + # Every DCP rank has the same batch, hence the same chunk plan. + o_rows = partial_o.view(rows, n, head_dim) + lse_rows = partial_lse.view(rows, n) + out = partial_o.new_empty(rows, head_dim) + for start in range(0, rows, state.capacity_rows): + end = min(start + state.capacity_rows, rows) + out[start:end] = decode_cp_a2a_lse_reduce( + partial_o=o_rows[start:end], + partial_lse=lse_rows[start:end], + workspace=workspace, + cp_rank=state.cp_rank, + cp_size=state.cp_size, + lse_mode=lse_mode, ) + return out.view(batch, heads, head_dim) - mapping = Mapping( - world_size=cp_size, - rank=cp_rank, - gpus_per_node=torch.cuda.device_count(), - cp_size=cp_size, - tp_size=1, - pp_size=1, + +def _fi_a2a_fused_reduce_fake( + partial_o: torch.Tensor, partial_lse: torch.Tensor, is_lse_base_on_e: bool +) -> torch.Tensor: + return partial_o.new_empty(partial_o.shape[:-2] + partial_o.shape[-1:]) + + +# A custom op so Dynamo (decode tc-piecewise graphs, torch.compile) sees one opaque +# collective instead of tracing the Python stream-to-workspace lookup. +@register_custom_op(fake_impl=_fi_a2a_fused_reduce_fake) +def fi_a2a_fused_reduce( + partial_o: torch.Tensor, partial_lse: torch.Tensor, is_lse_base_on_e: bool +) -> torch.Tensor: + state = _FI_A2A_STATE + assert state is not None, ( + "fi_a2a workspace not initialized; BaseRunner.warmup() creates it " + "before CUDA graph capture." ) - workspace = decode_cp_a2a_allocate_mnnvl_workspace( - mapping, - mnnvl_config=MnnvlConfig( - comm_backend=TorchDistributedCommBackend(cp_group.device_group) - ), + return _fi_a2a_fused_reduce_impl( + state=state, + partial_o=partial_o, + partial_lse=partial_lse, + is_lse_base_on_e=is_lse_base_on_e, ) - decode_cp_a2a_init_workspace(workspace, cp_rank, cp_size) - # REQUIRED barrier before the first alltoall: every rank must finish init, - # else a rank writes a peer's FIFO before it is ready -> deadlock. - dist.barrier(group=cp_group.device_group) - _FI_A2A_STATE = { - "workspace": workspace, - "cp_rank": cp_rank, - } def dcp_a2a_lse_reduce( @@ -506,52 +723,34 @@ def _dcp_fi_a2a_lse_reduce( cp_group: "GroupCoordinator", is_lse_base_on_e: bool = True, ) -> torch.Tensor: - """fi_a2a: delegate only the cross-rank exchange to FlashInfer's MNNVL kernel, - then reuse the local Triton LSE combine. FlashInfer takes output + LSE as - separate tensors: partial_o [B, H_per_rank, cp_size, D] (peer axis 2nd-to-last), - softmax_stats [B, H_per_rank, cp_size, 2] fp32 (S padded 1->2). + """fi_a2a: one FlashInfer kernel exchanges the head partials and merges them. + Its layout is partial_o [B, H_per_rank, cp_size, D] (peer axis second-to-last) + plus a contiguous fp32 partial_lse [B, H_per_rank, cp_size]. """ - from flashinfer.comm.dcp_alltoall import decode_cp_a2a_alltoall - - state = _FI_A2A_STATE - assert state is not None, ( - "fi_a2a workspace not initialized — call init_fi_a2a_workspace(dcp_group) " - "at model-runner init (before CUDA graph capture)." - ) - N = cp_group.world_size B, H, D = cp_attn_out.shape assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})" H_per_rank = H // N + if cp_attn_out.dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError( + f"DCP fi_a2a: attention partials are {cp_attn_out.dtype}; the fused " + "reduce supports float16 and bfloat16." + ) + # The op rejects an empty input; B is equal on every DCP rank, so all ranks + # skip this collective together. + if B == 0: + return cp_attn_out.new_empty(0, H_per_rank, D) - # Note(kpham-sgl): empty(), not zeros() -- the pack below fills partial_o and - # stats slot 0, and slot 1 is never read by anyone. The a2a moves the stats - # field as opaque bytes and we only ever read slot 0 back off the wire. partial_o = torch.empty( B, H_per_rank, N, D, dtype=cp_attn_out.dtype, device=cp_attn_out.device ) - softmax_stats = torch.empty( - B, H_per_rank, N, 2, dtype=torch.float32, device=cp_attn_out.device + partial_lse = torch.empty( + B, H_per_rank, N, dtype=torch.float32, device=cp_attn_out.device ) dcp_pack_a2a_send( cp_attn_out, cp_attn_lse, partial_o.permute(2, 0, 1, 3), - softmax_stats[..., 0].permute(2, 0, 1), + partial_lse.permute(2, 0, 1), ) - - o_out, stats_out = decode_cp_a2a_alltoall( - partial_o, - softmax_stats, - state["workspace"], - state["cp_rank"], - N, - ) - - recv_output = o_out.permute(2, 0, 1, 3) - recv_lse = stats_out[..., 0].permute(2, 0, 1) - - combined, _ = dcp_lse_combine_triton( - recv_output, recv_lse, is_lse_base_on_e=is_lse_base_on_e - ) - return combined + return fi_a2a_fused_reduce(partial_o, partial_lse, is_lse_base_on_e) diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 643b7d0b6386..46212a434078 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -67,6 +67,10 @@ logger = logging.getLogger(__name__) +# fi_a2a workspace capacity in tokens when no decode graph is captured; +# arbitrary, since wider eager batches are reduced in chunks of this size. +_FI_A2A_EAGER_ONLY_MAX_TOKENS = 256 + def _allocate_decode_buffers( *, @@ -320,19 +324,29 @@ def _pre_initialize_flashinfer_allreduce_workspace(self): ) def _pre_initialize_fi_a2a_workspace(self): - """Allocate the FlashInfer MNNVL all-to-all workspace for the fi_a2a DCP - comm backend; must run before CG capture (it syncs the stream + barriers - cross-rank, uncapturable) and raises early on non-MNNVL platforms. + """Create the fused fi_a2a DCP reduce workspaces; must run before CG + capture (workspace rendezvous and the group's first reduce are + collective and not capturable). """ - if ( - not get_parallel().dcp_enabled - or get_parallel().dcp_comm_backend != "fi_a2a" - ): + parallel = get_parallel() + if not parallel.dcp_enabled or parallel.dcp_comm_backend != "fi_a2a": return from sglang.srt.layers.dcp import init_fi_a2a_workspace + from sglang.srt.model_executor.runner_utils.pool import ( + get_or_create_global_graph_capture_stream, + ) - init_fi_a2a_workspace(get_parallel().dcp_group) + mr = self.model_runner + init_fi_a2a_workspace( + parallel.dcp_group, + # Batches wider than the captured graphs run as several fused calls. + max_tokens=self._widest_decode_rows() or _FI_A2A_EAGER_ONLY_MAX_TOKENS, + local_heads=mr.model_config.num_attention_heads // parallel.attn_tp_size, + head_dim=mr.model_config.kv_lora_rank, + dtype=mr.dtype, + probe_stream=get_or_create_global_graph_capture_stream(), + ) def _pre_initialize_pcie_ipc_workspace(self): """Build the PCIe-IPC all-reduce workspace before graph capture. diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 0f06abe7d8b0..e00c3f70df43 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -990,17 +990,6 @@ def fi_a2a_platform_blocker( return None -def is_fi_a2a_supported( - *, dcp_size: int, tp_size: int, pp_size: int, nnodes: int -) -> bool: - if not get_platform().is_sm100: - return False - if is_mnnvl_fabric_device(): - return True - tp_size_per_node = tp_size // max(nnodes // pp_size, 1) - return tp_size_per_node % dcp_size == 0 - - @lru_cache(maxsize=1) def _flashinfer_has_fused_dcp_reduce() -> bool: # A top-level find_spec locates the package without importing it. diff --git a/test/registered/kernels/ops/attention/test_dcp_lse_combine.py b/test/registered/kernels/ops/attention/test_dcp_lse_combine.py index f177178c1f6e..9be1ad117703 100644 --- a/test/registered/kernels/ops/attention/test_dcp_lse_combine.py +++ b/test/registered/kernels/ops/attention/test_dcp_lse_combine.py @@ -398,7 +398,10 @@ def test_pack_matches_the_copy_formulation_it_replaces(self): torch.equal(lane, lse.view(B, N, H_per_rank).permute(1, 0, 2)) ) - def test_pack_serves_the_split_peer_inside_layout(self): + def test_pack_serves_the_fused_reduce_layout(self): + """fi_a2a's fused reduce reads partial_o [B, H/N, N, D] and a contiguous + partial_lse [B, H/N, N] (peer axis second-to-last); the pack writes both + through the [N, B, H/N, ...] views it is handed.""" from sglang.kernels.ops.attention.dcp_kernels import dcp_pack_a2a_send for N, B, H_per_rank, D in ((2, 4, 8, 128), (4, 1, 16, 512)): @@ -410,14 +413,14 @@ def test_pack_serves_the_split_peer_inside_layout(self): partial_o = torch.empty( B, H_per_rank, N, D, dtype=torch.bfloat16, device=self.device ) - stats = torch.zeros( - B, H_per_rank, N, 2, dtype=torch.float32, device=self.device + partial_lse = torch.empty( + B, H_per_rank, N, dtype=torch.float32, device=self.device ) dcp_pack_a2a_send( out, lse, partial_o.permute(2, 0, 1, 3), - stats[..., 0].permute(2, 0, 1), + partial_lse.permute(2, 0, 1), ) want_o = out.view(B, N, H_per_rank, D).permute(0, 2, 1, 3) @@ -428,10 +431,7 @@ def test_pack_serves_the_split_peer_inside_layout(self): want_o.contiguous().view(torch.uint8), ) ) - self.assertTrue(torch.equal(stats[..., 0], want_lse)) - self.assertTrue( - torch.equal(stats[..., 1], torch.zeros_like(stats[..., 1])) - ) + self.assertTrue(torch.equal(partial_lse, want_lse)) def test_buffers_have_fixed_data_ptrs(self): """Pre-allocated buffer data_ptr must not change -- required for graph replay.""" diff --git a/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py b/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py new file mode 100644 index 000000000000..940bc44bc278 --- /dev/null +++ b/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py @@ -0,0 +1,198 @@ +"""DCP fi_a2a: FlashInfer's fused all-to-all + LSE reduce must equal the +all-gather + Triton merge it replaces, including chunked, graph-replayed and +fully-masked rows. + +Run with ``python test_dcp_fi_a2a_reduce.py --num-gpu 2,4,8``. Skips unless the +host can run the fused reduce (Blackwell, torch>=2.14, FlashInfer with the op). +""" + +import os + +import pytest +import torch +import torch.distributed as dist + +from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton +from sglang.srt.distributed import parallel_state as ps +from sglang.srt.layers.dcp import comm +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.utils import multigpu_pytest_main +from sglang.test.test_utils import publish_build_topology + +register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +# Workspace geometry: 3 tokens x 16 heads = 48 rows, so the (H, B) cases below +# cover a single call, an exactly-full call and multi-chunk calls. +MAX_TOKENS, LOCAL_HEADS, HEAD_DIM = 3, 16, 512 + + +def _blocker(): + if not torch.cuda.is_available(): + return "needs CUDA" + from sglang.srt.utils.common import fi_a2a_platform_blocker + + world_size = int(os.environ.get("WORLD_SIZE", "1")) + return fi_a2a_platform_blocker( + dcp_size=world_size, tp_size=world_size, pp_size=1, nnodes=1 + ) + + +_reason = _blocker() +if _reason is not None: + pytest.skip(f"fused fi_a2a reduce {_reason}", allow_module_level=True) + + +@pytest.fixture(scope="module") +def dcp(): + import torch.distributed._symmetric_memory as symm_mem + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + # As bootstrap does for fi_a2a: pick the backend before any symmetric allocation. + symm_mem.set_backend("NCCL") + ps.set_custom_all_reduce(False) + ps.init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=local_rank, + distributed_init_method="env://", + ) + publish_build_topology(tp_size=world_size, world_rank=rank) + ps.initialize_model_parallel() + group = ps.get_tp_group() + capture_stream = torch.cuda.Stream() + comm.init_fi_a2a_workspace( + group, + max_tokens=MAX_TOKENS, + local_heads=LOCAL_HEADS, + head_dim=HEAD_DIM, + dtype=torch.bfloat16, + probe_stream=capture_stream, + ) + yield group, capture_stream + ps.destroy_model_parallel() + ps.destroy_distributed_environment() + + +def _partials(group, *, batch, heads, dtype, seed): + gen = torch.Generator(device="cuda").manual_seed(seed + group.rank_in_group) + total_heads = heads * group.world_size + out = torch.randn( + batch, total_heads, HEAD_DIM, dtype=dtype, device="cuda", generator=gen + ) + lse = torch.randn( + batch, total_heads, dtype=torch.float32, device="cuda", generator=gen + ) + return out, lse + + +def _reference(group, out, lse, *, is_lse_base_on_e): + """Gather every rank's partials, keep this rank's heads, merge with Triton.""" + heads = out.shape[1] // group.world_size + own = slice(group.rank_in_group * heads, (group.rank_in_group + 1) * heads) + all_out = [torch.empty_like(out) for _ in range(group.world_size)] + all_lse = [torch.empty_like(lse) for _ in range(group.world_size)] + dist.all_gather(all_out, out, group=group.device_group) + dist.all_gather(all_lse, lse, group=group.device_group) + merged, _ = dcp_lse_combine_triton( + torch.stack([o[:, own] for o in all_out]).contiguous(), + torch.stack([s[:, own] for s in all_lse]).contiguous(), + is_lse_base_on_e=is_lse_base_on_e, + ) + return merged + + +def _fused(group, out, lse, *, is_lse_base_on_e): + return comm.dcp_a2a_lse_reduce( + out, lse, group, is_lse_base_on_e=is_lse_base_on_e, comm_backend="fi_a2a" + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("is_lse_base_on_e", [False, True]) +@pytest.mark.parametrize( + "heads,batch", [(16, 1), (16, 3), (16, 7), (2, 1), (2, 24), (2, 49)] +) +def test_matches_all_gather_reference(dcp, dtype, is_lse_base_on_e, heads, batch): + """Guards the head-to-peer mapping and the row chunking (B=7 at 16 heads and + B=49 at 2 heads exceed the 48-row workspace).""" + group, _ = dcp + out, lse = _partials(group, batch=batch, heads=heads, dtype=dtype, seed=11) + expected = _reference(group, out, lse, is_lse_base_on_e=is_lse_base_on_e) + got = _fused(group, out, lse, is_lse_base_on_e=is_lse_base_on_e) + torch.testing.assert_close(got, expected, rtol=1e-2, atol=1e-2) + + +def test_rows_empty_on_every_rank_are_zero(dcp): + """A row with no KV on any rank (LSE -inf everywhere, e.g. graph padding) + merges to 0, where the Triton combine divides 0 by 0.""" + group, _ = dcp + out, lse = _partials( + group, batch=2, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=23 + ) + empty = torch.arange(0, lse.shape[1], LOCAL_HEADS, device="cuda") + lse[:, empty] = float("-inf") # local head 0 of every destination rank + expected = _reference(group, out, lse, is_lse_base_on_e=False) + got = _fused(group, out, lse, is_lse_base_on_e=False) + assert torch.equal(got[:, 0], torch.zeros_like(got[:, 0])) + torch.testing.assert_close(got[:, 1:], expected[:, 1:], rtol=1e-2, atol=1e-2) + + +def test_graph_replay_and_eager_calls_interleave(dcp): + """The captured graph and eager decode use different workspaces on one + ordered stream; replays and eager calls must not corrupt each other.""" + group, capture_stream = dcp + static_out, static_lse = _partials( + group, batch=2, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=31 + ) + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + _fused(group, static_out, static_lse, is_lse_base_on_e=False) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + static_result = _fused( + group, static_out, static_lse, is_lse_base_on_e=False + ) + torch.cuda.current_stream().wait_stream(capture_stream) + for step in range(3): + new_out, new_lse = _partials( + group, batch=2, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=40 + step + ) + static_out.copy_(new_out) + static_lse.copy_(new_lse) + graph.replay() + eager_out, eager_lse = _partials( + group, batch=3, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=50 + step + ) + eager_result = _fused(group, eager_out, eager_lse, is_lse_base_on_e=False) + torch.testing.assert_close( + static_result, + _reference(group, new_out, new_lse, is_lse_base_on_e=False), + rtol=1e-2, + atol=1e-2, + ) + torch.testing.assert_close( + eager_result, + _reference(group, eager_out, eager_lse, is_lse_base_on_e=False), + rtol=1e-2, + atol=1e-2, + ) + + +def test_third_stream_raises_before_the_collective(dcp): + """Only the capture and serving streams own workspaces; a third stream must + fail on the host on every rank, not spin in the kernel.""" + group, _ = dcp + out, lse = _partials( + group, batch=1, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=61 + ) + _fused(group, out, lse, is_lse_base_on_e=False) # binds the serving stream + with torch.cuda.stream(torch.cuda.Stream()): + with pytest.raises(RuntimeError, match="third stream"): + _fused(group, out, lse, is_lse_base_on_e=False) + + +if __name__ == "__main__": + multigpu_pytest_main(__name__, __file__, num_gpus=(2, 4)) From 4b96291c65ec2b14c451403ed7fc9795e7d134cb Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Wed, 30 Sep 2026 01:16:54 +0700 Subject: [PATCH 05/10] [DCP] Document the fused fi_a2a backend Describe what fi_a2a now requires (Blackwell, one NVLink domain, an MLA model, torch >= 2.14, a FlashInfer build with decode_cp_a2a_lse_reduce), how an unset --dcp-comm-backend resolves, and the process-wide settings it changes: NCCL symmetric-memory backend with custom all-reduce v1, NCCL_CUMEM_ENABLE=1, and PDMUX / torch symm-mem all-reduce rejected. --- .../autoregressive/Moonshotai/Kimi-K3.mdx | 2 +- docs/docs/advanced_features/dcp.mdx | 15 +++++++++++---- docs/docs/advanced_features/server_arguments.mdx | 2 +- 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx index b6dd00264aea..1b43a3f2b20e 100644 --- a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx +++ b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx @@ -203,7 +203,7 @@ Speculation: DSPARK holds block size + 1 (= 8) intermediate states per request **Blackwell DCP notes** — the DCP cells are Balanced and High-Throughput on every Blackwell platform, in both the `Unified` and `Decode` roles: - DCP is the only axis that shards the TP-replicated MLA KV; Low-Latency skips it. -- Leave `--dcp-comm-backend` unset (fabric-resolved: `fi_a2a` on GB200/GB300, `a2a` on B200/B300). +- Leave `--dcp-comm-backend` unset: it resolves to `fi_a2a` where FlashInfer's fused reduce can run (torch 2.14+ and a FlashInfer build with `decode_cp_a2a_lse_reduce`), otherwise `a2a`. - No `--enable-symm-mem` under DCP (force-disabled for decode-graph correctness). - Explicit `tokenspeed_mla` force-rewrites `--kv-cache-dtype` to fp8; the default `cutedsl_mla` serves either dtype. - Calculator ratios run well above 1 here (`r > 1` is legal): `bfloat16` state buys admission, `fp8` KV buys context. diff --git a/docs/docs/advanced_features/dcp.mdx b/docs/docs/advanced_features/dcp.mdx index 262b69c4c295..ecd3eed6fc40 100644 --- a/docs/docs/advanced_features/dcp.mdx +++ b/docs/docs/advanced_features/dcp.mdx @@ -15,7 +15,7 @@ On Kimi K3, DCP applies only to MLA layers. Request-indexed KDA state is unchang | Setting | Behavior | | --- | --- | | `--dcp-size N` | Enables DCP and widens virtual capacity and page size by `N`. Alias: `--decode-context-parallel-size`. | -| `--dcp-comm-backend` | Selects the MLA partial-output merge: `ag_rs`, `a2a`, or `fi_a2a`. Unset resolves to `fi_a2a` where the DCP group shares an MNNVL domain on Blackwell, otherwise `a2a` on CUDA/ROCm. | +| `--dcp-comm-backend` | Selects the MLA partial-output merge: `ag_rs`, `a2a`, or `fi_a2a`. Unset resolves to `fi_a2a` where FlashInfer's fused reduce can run (see [Communication backends](#communication-backends)), otherwise `a2a` on CUDA/ROCm. | | `--dcp-replicate-q-proj` | Removes the query all-gather for supported MLA weights. Use `--no-dcp-replicate-q-proj` to disable a model-specific default. | | `--enable-dp-attention` | Composes only when DCP groups nest inside attention TP. | @@ -44,7 +44,7 @@ sglang serve \ --port 30000 ``` -Kimi K3 on 32 GPUs with DPA and DCP 8. The model override enables replicated Q by default; the communication backend resolves generically to `fi_a2a` on MNNVL systems, otherwise `a2a`: +Kimi K3 on 32 GPUs with DPA and DCP 8. The model override enables replicated Q by default; the communication backend resolves generically to `fi_a2a` where the fused reduce can run, otherwise `a2a`: ```bash Command sglang serve \ @@ -133,9 +133,16 @@ Normally each layer all-gathers its absorbed and rotary query components. `--dcp | `ag_rs` | Query AG + LSE AG + FP32 output RS | Fallback for platforms without the a2a path; opt in with `--dcp-comm-backend ag_rs` | | `a2a`, gathered Q | Query AG + packed NCCL A2A | Two collectives | | `a2a`, replicated Q | One packed NCCL A2A | Canonical low-latency path | -| `fi_a2a`, replicated Q | One FlashInfer MNNVL A2A | Requires Blackwell and a DCP group inside one MNNVL domain (GB200/GB300 fabric, or a single node) | +| `fi_a2a`, replicated Q | One FlashInfer kernel: NCCL LSA all-to-all and LSE merge fused | Requires Blackwell, the DCP group inside one NVLink domain (GB200/GB300 fabric, or a single node), an MLA model, torch 2.14 or newer, and a FlashInfer build that provides `decode_cp_a2a_lse_reduce` | -When `--dcp-comm-backend` is unset, every DCP launch resolves it the same way: `fi_a2a` where the DCP group shares one MNNVL domain on Blackwell, otherwise `a2a` on CUDA/ROCm, and `ag_rs` only on platforms without the a2a path. Query replication stays model-resolved: Kimi K3 enables it by default on the a2a family, and its decode backend is `cutedsl_mla`. +When `--dcp-comm-backend` is unset, every DCP launch resolves it the same way: `fi_a2a` where the fused reduce can run, otherwise `a2a` on CUDA/ROCm, and `ag_rs` only on platforms without the a2a path. The startup log names the missing requirement when `fi_a2a` is not selected. + +`fi_a2a` exchanges and merges the partials in one cooperative kernel over NCCL symmetric memory, instead of an all-to-all followed by a separate merge kernel. Running it changes a few process-wide settings: + +- torch symmetric memory uses the NCCL backend, so custom all-reduce uses its v1 kernel instead of `CustomAllReduceV2`; +- `NCCL_CUMEM_ENABLE` defaults to `1`, because the fused workspace is an NCCL symmetric window; +- `--enable-pdmux` and `--enable-torch-symm-mem` are rejected, since they need more CUDA streams or the CUDA symmetric-memory backend; +- decode batches wider than the captured CUDA graphs run as several fused calls. Query replication stays model-resolved: Kimi K3 enables it by default on the a2a family, and its decode backend is `cutedsl_mla`. DCP adds a context-independent decode collective while cutting context-dependent KV storage and reads by about `c`. Extend is outside that decode cost model: it may gather cached prefix shards, restore token order, and append new tokens, and that work grows with prefix length. diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 15f0543c4b7c..ee003c89b766 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -668,7 +668,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--dcp-comm-backend` - Communication backend for the DCP attention reduction: AllGather + ReduceScatter, fused NCCL All-to-All, or FlashInfer MNNVL All-to-All. Unset resolves to fi_a2a where the DCP group shares an MNNVL domain on Blackwell, else a2a on CUDA/ROCm, else ag_rs. + Communication backend for the DCP attention reduction: AllGather + ReduceScatter, fused NCCL All-to-All, or FlashInfer's fused all-to-all + LSE merge (Blackwell, one NVLink domain, MLA, torch 2.14+, FlashInfer with decode_cp_a2a_lse_reduce). Unset resolves to fi_a2a where it can run, else a2a on CUDA/ROCm, else ag_rs. None ag_rs, a2a, fi_a2a From a8981f99b5af5b899715fbacad32fba9f5c420ee Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Tue, 6 Oct 2026 16:20:47 +0700 Subject: [PATCH 06/10] [DCP] Run fi_a2a on the active torch symmetric-memory backend FlashInfer's fused DCP reduce allocates its workspace on whichever torch symmetric-memory backend is active. Stop selecting NCCL at bootstrap, and drop the custom all-reduce v1 fallback, the NCCL_CUMEM_ENABLE default and the cuMem, torch 2.14, NCCL 2.29 and --enable-torch-symm-mem checks that only that backend needed. No per-rank environment check is left to agree on, so remove the startup probe and its gloo vote. Keep one workspace for graphs captured on the capture stream and one for eager calls, chosen by the current stream, and re-raise a failed workspace creation with what to change. --- python/sglang/srt/arg_groups/overrides.py | 9 - python/sglang/srt/arg_groups/parallel_hook.py | 3 - python/sglang/srt/distributed/bootstrap.py | 21 - .../device_communicators/custom_all_reduce.py | 11 +- python/sglang/srt/layers/dcp/comm.py | 401 +++++------------- .../srt/model_executor/runner/base_runner.py | 7 +- python/sglang/srt/utils/common.py | 9 +- .../communication/test_dcp_fi_a2a_reduce.py | 83 ++-- .../unit/server_args/test_server_args.py | 38 +- test/registered/unit/utils/test_common.py | 78 ++++ 10 files changed, 260 insertions(+), 400 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index cc2174d05ee0..cee730675a12 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -24,7 +24,6 @@ import json import logging import math -import os from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from sglang.srt.arg_groups import model_override_base @@ -1329,19 +1328,11 @@ def fi_a2a_blocker(view: Any) -> Optional[str]: "does not support --enable-pdmux (every captured stream group would " "need its own fused-reduce workspace)" ) - if view.enable_torch_symm_mem: - return ( - "cannot be combined with --enable-torch-symm-mem (fi_a2a moves torch " - "symmetric memory to the NCCL backend, which that all-reduce cannot use)" - ) if envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.get() and view.pp_size > 1: return ( "cannot be combined with SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP at " "--pp-size > 1 (its decode warmup runs on a third CUDA stream)" ) - # The fused workspace is an NCCL symmetric window, which needs cuMem memory. - if os.environ.get("NCCL_CUMEM_ENABLE") == "0": - return "needs cuMem-backed NCCL memory, but NCCL_CUMEM_ENABLE=0 is exported" return None diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index 404274482e52..538c316dac09 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -248,9 +248,6 @@ def handle_decode_context_parallelism(server_args: Any): f"--dcp-comm-backend fi_a2a {blocker}. Use --dcp-comm-backend a2a, " "or leave the flag unset to resolve a supported backend." ) - # The fused reduce's workspace is an NCCL symmetric window, which needs - # cuMem-backed NCCL memory; set before any communicator exists. - os.environ.setdefault("NCCL_CUMEM_ENABLE", "1") if cfg.dcp_replicate_q_proj: if cfg.dcp_size <= 1: raise ValueError("--dcp-replicate-q-proj requires --dcp-size > 1.") diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index 51016324f93f..6b3932995dc2 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/python/sglang/srt/distributed/bootstrap.py @@ -142,12 +142,6 @@ def init_parallel_runtime( local_omp_cpuid=local_omp_cpuid, dist_init_method=dist_init_method, ) - if ( - device == "cuda" - and parallel.dcp_enabled - and parallel.dcp_comm_backend == "fi_a2a" - ): - _select_nccl_symm_mem_backend() _init_parallel_groups( backend=backend, @@ -301,21 +295,6 @@ def _init_cpu_threads_env( ) -def _select_nccl_symm_mem_backend() -> None: - # DCP fi_a2a's fused reduce allocates torch symmetric memory on the NCCL - # backend. The choice is process-global and frozen by the first symmetric - # allocation, which custom all-reduce makes while the groups are built. - import torch.distributed._symmetric_memory as symm_mem - - try: - symm_mem.set_backend("NCCL") - except RuntimeError as e: - raise RuntimeError( - "--dcp-comm-backend fi_a2a needs torch's NCCL symmetric-memory " - f"backend, which could not be selected before process-group init: {e}" - ) from e - - def _init_parallel_groups( *, backend: str, diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index 7632a986718c..55c292cde30c 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -356,21 +356,12 @@ def dispatch_custom_allreduce( can_use_custom_all_reduce_v2); other cross-node groups fall back to NCCL. """ if _is_cuda and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get(): - import torch.distributed._symmetric_memory as torch_symm_mem - from .custom_all_reduce_v2 import ( CustomAllReduceV2, can_use_custom_all_reduce_v2, ) - if torch_symm_mem.get_backend(device) == "NCCL": - # CustomAllReduceV2 allocates torch symmetric memory with a group name - # and rendezvouses on the gloo group; the NCCL allocator rejects both. - logger.info( - "[AR] torch symmetric memory uses the NCCL backend (DCP fi_a2a); " - "CustomAllReduceV2 cannot allocate from it, using CustomAllreduce" - ) - elif can_use_custom_all_reduce_v2(group=group, device=device): + if can_use_custom_all_reduce_v2(group=group, device=device): logger.debug("[AR] Using CustomAllReduceV2 (JIT-compiled)") return CustomAllReduceV2 diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index 33ca71a0bcbc..b41e385965c3 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/python/sglang/srt/layers/dcp/comm.py @@ -25,7 +25,6 @@ import msgspec import torch -import torch.distributed as dist from sglang.kernels.ops.attention.dcp_kernels import ( CPTritonContext, @@ -37,13 +36,9 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) -from sglang.srt.distributed.device_communicators.pynccl_wrapper import ( - NCCLLibrary, - find_nccl_library, -) from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import is_hip +from sglang.srt.utils import is_hip, log_info_on_rank0 from sglang.srt.utils.custom_op import register_custom_op logger = logging.getLogger(__name__) @@ -396,29 +391,25 @@ def all_gather_kv_cache_for_dcp( # A2A communication backend for DCP decode (alternative to AG+RS above): exchange # per-head partial outputs + LSEs across DCP ranks and merge them. a2a sends one # NCCL all-to-all and merges with the Triton LSE kernel; fi_a2a hands exchange and -# merge to FlashInfer's fused decode_cp_a2a_lse_reduce (one cooperative kernel over -# NCCL LSA symmetric memory). +# merge to FlashInfer's fused decode_cp_a2a_lse_reduce, one kernel that writes into +# its peers' torch symmetric memory. # --------------------------------------------------------------------------- -class _FiA2aState(msgspec.Struct, kw_only=True): +class _FiA2aState(msgspec.Struct, frozen=True, kw_only=True): cp_rank: int - cp_size: int - # One fused call carries at most this many rows (product of the leading dims). + # Rows (tokens x local heads) one fused call carries. capacity_rows: int - # Created before graph capture; each binds to the first stream that uses it. - unbound: list - by_stream: dict + # A workspace serves one ordered CUDA stream: graphs captured on capture_stream + # use capture_workspace, also when replayed; eager calls use serving_workspace. + capture_stream: int + capture_workspace: torch.Tensor + serving_workspace: torch.Tensor -# Per-process singleton, populated once before CUDA graph capture by -# init_fi_a2a_workspace(); holds strong references to the workspaces. +# Set once per process, before CUDA graph capture, by init_fi_a2a_workspace(). _FI_A2A_STATE: Optional[_FiA2aState] = None -# A workspace may serve only one ordered CUDA stream: the graph-capture stream -# (bound by the startup probe) and the serving stream (bound on first use). -_FI_A2A_NUM_WORKSPACES = 2 - def init_fi_a2a_workspace( cp_group: GroupCoordinator, @@ -427,263 +418,131 @@ def init_fi_a2a_workspace( local_heads: int, head_dim: int, dtype: torch.dtype, - probe_stream: torch.cuda.Stream, + capture_stream: torch.cuda.Stream, ) -> None: - """Create the fused fi_a2a workspaces on every DCP rank; collective, run - before graph capture (rendezvous and the first reduce are not capturable).""" + """Create the fused reduce's workspaces on every DCP rank. Collective, and it + synchronizes the current stream, so it runs before graph capture.""" global _FI_A2A_STATE if _FI_A2A_STATE is not None or cp_group.world_size == 1: return - # The kernel's receive spin-wait has no timeout, so every rank must agree to - # proceed before the NCCL collectives below; the verdict travels over gloo. - _fi_a2a_require_all_ranks( - cp_group=cp_group, - stage="eligibility check", - error=_fi_a2a_local_blocker(cp_group=cp_group, head_dim=head_dim, dtype=dtype), - ) + import torch.distributed._symmetric_memory as symm_mem from flashinfer.comm import ( decode_cp_a2a_lse_reduce_create_workspace, decode_cp_a2a_lse_reduce_workspace_size, ) - # Rendezvous finds the group's NCCL communicator by name, and ProcessGroupNCCL - # creates it lazily: one collective on the device group makes it exist. - dist.barrier(group=cp_group.device_group) - workspaces = [ - decode_cp_a2a_lse_reduce_create_workspace( - max_tokens=max_tokens, - local_heads=local_heads, - cp_size=cp_group.world_size, - head_dim=head_dim, - dtype=dtype, - group=cp_group.device_group, - ) - for _ in range(_FI_A2A_NUM_WORKSPACES) - ] - dist.barrier(group=cp_group.device_group) - state = _FiA2aState( - cp_rank=cp_group.rank_in_group, + geometry = dict( + max_tokens=max_tokens, + local_heads=local_heads, cp_size=cp_group.world_size, - capacity_rows=max_tokens * local_heads, - unbound=workspaces[1:], - by_stream={probe_stream.cuda_stream: workspaces[0]}, - ) - with torch.cuda.stream(probe_stream): - error = _fi_a2a_probe( - state=state, - cp_group=cp_group, - local_heads=local_heads, - head_dim=head_dim, - dtype=dtype, - ) - _fi_a2a_require_all_ranks(cp_group=cp_group, stage="startup probe", error=error) - _FI_A2A_STATE = state - workspace_mib = ( - decode_cp_a2a_lse_reduce_workspace_size( - max_tokens=max_tokens, - local_heads=local_heads, - cp_size=cp_group.world_size, - head_dim=head_dim, - dtype=dtype, - ) - / 2**20 - ) - logger.info( - "DCP fi_a2a: fused all-to-all + LSE reduce ready (rank %d/%d, %d heads/rank, " - "head_dim %d, %s, capacity %d rows, %d workspaces x %.1f MiB)", - state.cp_rank, - state.cp_size, - local_heads, - head_dim, - dtype, - state.capacity_rows, - _FI_A2A_NUM_WORKSPACES, - workspace_mib, + head_dim=head_dim, + dtype=dtype, ) + device = torch.device("cuda", torch.cuda.current_device()) + def create_workspace() -> torch.Tensor: + return decode_cp_a2a_lse_reduce_create_workspace( + **geometry, group=cp_group.device_group + ) -def _fi_a2a_local_blocker( - *, cp_group: GroupCoordinator, head_dim: int, dtype: torch.dtype -) -> Optional[str]: try: - import flashinfer.comm.dcp_lse_reduce # noqa: F401 - except ImportError as e: - return f"FlashInfer lacks flashinfer.comm.decode_cp_a2a_lse_reduce ({e})" - import torch.distributed._symmetric_memory as symm_mem - - backend = symm_mem.get_backend(torch.device("cuda", torch.cuda.current_device())) - if backend != "NCCL": - return ( - f"torch symmetric memory uses the {backend} backend; the fused reduce " - "needs NCCL, which bootstrap selects before process-group init" - ) - # The op is built against the pip NCCL headers; an older loaded libnccl would - # abort at the first call instead of raising. - nccl_version = NCCLLibrary(find_nccl_library()).ncclGetRawVersion() - if nccl_version < 22900: - return f"the loaded NCCL is {nccl_version}; the fused reduce needs >= 2.29" - if dtype not in (torch.float16, torch.bfloat16): - return f"attention partials are {dtype}; the fused reduce supports float16 and bfloat16" - if head_dim * dtype.itemsize % 16 != 0: - return f"head_dim {head_dim} rows of {dtype} are not 16-byte aligned" - if cp_group.world_size > 64: - return f"dcp_size {cp_group.world_size} exceeds the fused reduce's 64 ranks" - if dist.get_rank(cp_group.device_group) != cp_group.rank_in_group: - return "the DCP device-group rank differs from rank_in_group" - return None - - -def _fi_a2a_require_all_ranks( - *, cp_group: GroupCoordinator, stage: str, error: Optional[str] -) -> None: - errors: list = [None] * cp_group.world_size - dist.all_gather_object(errors, error, group=cp_group.cpu_group) - failed = [(rank, e) for rank, e in enumerate(errors) if e is not None] - if failed: - detail = "; ".join(f"rank {rank}: {e}" for rank, e in failed[:4]) + capture_workspace = create_workspace() + serving_workspace = create_workspace() + except RuntimeError as e: raise RuntimeError( - f"--dcp-comm-backend fi_a2a {stage} failed on {len(failed)}/" - f"{cp_group.world_size} DCP ranks: {detail}. " - "Use --dcp-comm-backend a2a." - ) - + "--dcp-comm-backend fi_a2a: FlashInfer could not create its fused-reduce " + "workspace on this process's torch symmetric-memory backend " + f"({symm_mem.get_backend(device)}): {str(e).rstrip('.')}. Use a " + "flashinfer-python release whose decode_cp_a2a_lse_reduce does not " + "select a backend, or pass --dcp-comm-backend a2a." + ) from e + _FI_A2A_STATE = _FiA2aState( + cp_rank=cp_group.rank_in_group, + capacity_rows=max_tokens * local_heads, + capture_stream=capture_stream.cuda_stream, + capture_workspace=capture_workspace, + serving_workspace=serving_workspace, + ) + workspace_mib = decode_cp_a2a_lse_reduce_workspace_size(**geometry) / 2**20 + log_info_on_rank0( + logger, + "DCP fi_a2a: fused all-to-all + LSE reduce ready " + f"(backend={symm_mem.get_backend(device)}, capacity " + f"{max_tokens * local_heads} rows ({max_tokens} tokens x {local_heads} " + f"heads), 2 workspaces x {workspace_mib:.2f} MiB)", + ) -def _fi_a2a_probe( - *, - state: _FiA2aState, - cp_group: GroupCoordinator, - local_heads: int, - head_dim: int, - dtype: torch.dtype, -) -> Optional[str]: - """Check the fused reduce against all-gather + Triton merge for both LSE bases. - Its first call also builds the group's NCCL device communicator (collective), - which must happen eagerly, before any graph capture. - """ - rank, n = cp_group.rank_in_group, cp_group.world_size - gen = torch.Generator(device="cuda").manual_seed(4929 + rank) - error = None - for is_lse_base_on_e in (False, True): - partial_o = torch.randn( - 2, local_heads, n, head_dim, dtype=dtype, device="cuda", generator=gen - ) - partial_lse = torch.randn( - 2, local_heads, n, dtype=torch.float32, device="cuda", generator=gen - ) - # The reference collectives run on every rank and every iteration, even - # after a local failure, so no rank waits in a collective the others skip. - gathered_o = cp_group.all_gather(partial_o.unsqueeze(0), dim=0) - gathered_lse = cp_group.all_gather(partial_lse.unsqueeze(0), dim=0) - expected, _ = dcp_lse_combine_triton( - gathered_o[:, :, :, rank].contiguous(), - gathered_lse[:, :, :, rank].contiguous(), - is_lse_base_on_e=is_lse_base_on_e, - ) - if error is not None: - continue - try: - got = _fi_a2a_fused_reduce_impl( - state=state, - partial_o=partial_o, - partial_lse=partial_lse, - is_lse_base_on_e=is_lse_base_on_e, - ) - torch.cuda.current_stream().synchronize() - except RuntimeError as e: - error = str(e).splitlines()[0] - continue - if not torch.allclose(got.float(), expected.float(), rtol=1e-2, atol=1e-2): - max_diff = (got.float() - expected.float()).abs().max().item() - base = "e" if is_lse_base_on_e else "2" - error = ( - f"fused reduce disagrees with the all-gather reference " - f"(LSE base {base}, max abs diff {max_diff:.3g})" - ) - return error - - -def _fi_a2a_workspace_for_current_stream(state: _FiA2aState) -> torch.Tensor: - stream = torch.cuda.current_stream().cuda_stream - workspace = state.by_stream.get(stream) - if workspace is not None: - return workspace - if not state.unbound: - bound = ", ".join(hex(s) for s in state.by_stream) - raise RuntimeError( - f"DCP fi_a2a: CUDA stream {stream:#x} is a third stream calling the " - f"fused reduce; its workspaces are bound to streams {bound}. Each " - "ordered stream needs its own workspace, created before graph capture." - ) - workspace = state.unbound.pop(0) - state.by_stream[stream] = workspace - return workspace +def _fi_a2a_peer_partials( + cp_attn_out: torch.Tensor, cp_attn_lse: torch.Tensor, cp_size: int +) -> tuple[torch.Tensor, torch.Tensor]: + """[B, H, D] partials and [B, H] LSE as the fused reduce's [B, H/cp, cp, D] + and [B, H/cp, cp]; head h goes to peer h // (H/cp).""" + batch, heads, head_dim = cp_attn_out.shape + local_heads = heads // cp_size + partial_o = cp_attn_out.new_empty(batch, local_heads, cp_size, head_dim) + partial_lse = cp_attn_lse.new_empty(batch, local_heads, cp_size) + dcp_pack_a2a_send( + cp_attn_out, + cp_attn_lse, + partial_o.permute(2, 0, 1, 3), + partial_lse.permute(2, 0, 1), + ) + return partial_o, partial_lse -def _fi_a2a_fused_reduce_impl( - *, - state: _FiA2aState, - partial_o: torch.Tensor, - partial_lse: torch.Tensor, +def _fi_a2a_lse_reduce_fake( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_size: int, is_lse_base_on_e: bool, ) -> torch.Tensor: - from flashinfer.comm import decode_cp_a2a_lse_reduce - - workspace = _fi_a2a_workspace_for_current_stream(state) - lse_mode = "basee" if is_lse_base_on_e else "base2" - batch, heads, n, head_dim = partial_o.shape - rows = batch * heads - if rows <= state.capacity_rows: - return decode_cp_a2a_lse_reduce( - partial_o=partial_o, - partial_lse=partial_lse, - workspace=workspace, - cp_rank=state.cp_rank, - cp_size=state.cp_size, - lse_mode=lse_mode, - ) - # An eager batch wider than the graph capture size: the op flattens leading - # dims into rows, so reduce capacity-sized row chunks in order on this stream. - # Every DCP rank has the same batch, hence the same chunk plan. - o_rows = partial_o.view(rows, n, head_dim) - lse_rows = partial_lse.view(rows, n) - out = partial_o.new_empty(rows, head_dim) - for start in range(0, rows, state.capacity_rows): - end = min(start + state.capacity_rows, rows) - out[start:end] = decode_cp_a2a_lse_reduce( - partial_o=o_rows[start:end], - partial_lse=lse_rows[start:end], - workspace=workspace, - cp_rank=state.cp_rank, - cp_size=state.cp_size, - lse_mode=lse_mode, - ) - return out.view(batch, heads, head_dim) + batch, heads, head_dim = cp_attn_out.shape + return cp_attn_out.new_empty(batch, heads // cp_size, head_dim) -def _fi_a2a_fused_reduce_fake( - partial_o: torch.Tensor, partial_lse: torch.Tensor, is_lse_base_on_e: bool +# A custom op, so torch.compile sees one opaque collective rather than tracing the +# workspace choice. +@register_custom_op(fake_impl=_fi_a2a_lse_reduce_fake) +def fi_a2a_lse_reduce( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_size: int, + is_lse_base_on_e: bool, ) -> torch.Tensor: - return partial_o.new_empty(partial_o.shape[:-2] + partial_o.shape[-1:]) - + from flashinfer.comm import decode_cp_a2a_lse_reduce -# A custom op so Dynamo (decode tc-piecewise graphs, torch.compile) sees one opaque -# collective instead of tracing the Python stream-to-workspace lookup. -@register_custom_op(fake_impl=_fi_a2a_fused_reduce_fake) -def fi_a2a_fused_reduce( - partial_o: torch.Tensor, partial_lse: torch.Tensor, is_lse_base_on_e: bool -) -> torch.Tensor: state = _FI_A2A_STATE assert state is not None, ( - "fi_a2a workspace not initialized; BaseRunner.warmup() creates it " - "before CUDA graph capture." + "BaseRunner.warmup() creates the fi_a2a workspaces before graph capture" ) - return _fi_a2a_fused_reduce_impl( - state=state, - partial_o=partial_o, - partial_lse=partial_lse, - is_lse_base_on_e=is_lse_base_on_e, + batch, heads, head_dim = cp_attn_out.shape + local_heads = heads // cp_size + # The kernel rejects an empty batch; every DCP rank sees the same batch. + if batch == 0: + return cp_attn_out.new_empty(0, local_heads, head_dim) + partial_o, partial_lse = _fi_a2a_peer_partials( + cp_attn_out=cp_attn_out, cp_attn_lse=cp_attn_lse, cp_size=cp_size ) + if torch.cuda.current_stream().cuda_stream == state.capture_stream: + workspace = state.capture_workspace + else: + workspace = state.serving_workspace + # Eager batches wider than the widest captured graph run as several calls, + # in the same order on every DCP rank. + tokens_per_call = state.capacity_rows // local_heads + outputs = [ + decode_cp_a2a_lse_reduce( + partial_o=partial_o[start : start + tokens_per_call], + partial_lse=partial_lse[start : start + tokens_per_call], + workspace=workspace, + cp_rank=state.cp_rank, + cp_size=cp_size, + lse_mode="basee" if is_lse_base_on_e else "base2", + ) + for start in range(0, batch, tokens_per_call) + ] + return outputs[0] if len(outputs) == 1 else torch.cat(outputs) def dcp_a2a_lse_reduce( @@ -703,8 +562,11 @@ def dcp_a2a_lse_reduce( return cp_attn_out if comm_backend == "fi_a2a": - return _dcp_fi_a2a_lse_reduce( - cp_attn_out, cp_attn_lse, cp_group, is_lse_base_on_e + return fi_a2a_lse_reduce( + cp_attn_out=cp_attn_out, + cp_attn_lse=cp_attn_lse, + cp_size=cp_group.world_size, + is_lse_base_on_e=is_lse_base_on_e, ) N = cp_group.world_size @@ -750,42 +612,3 @@ def dcp_a2a_lse_reduce( recv_output, recv_lse, is_lse_base_on_e=is_lse_base_on_e ) return combined - - -def _dcp_fi_a2a_lse_reduce( - cp_attn_out: torch.Tensor, - cp_attn_lse: torch.Tensor, - cp_group: "GroupCoordinator", - is_lse_base_on_e: bool = True, -) -> torch.Tensor: - """fi_a2a: one FlashInfer kernel exchanges the head partials and merges them. - Its layout is partial_o [B, H_per_rank, cp_size, D] (peer axis second-to-last) - plus a contiguous fp32 partial_lse [B, H_per_rank, cp_size]. - """ - N = cp_group.world_size - B, H, D = cp_attn_out.shape - assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})" - H_per_rank = H // N - if cp_attn_out.dtype not in (torch.float16, torch.bfloat16): - raise RuntimeError( - f"DCP fi_a2a: attention partials are {cp_attn_out.dtype}; the fused " - "reduce supports float16 and bfloat16." - ) - # The op rejects an empty input; B is equal on every DCP rank, so all ranks - # skip this collective together. - if B == 0: - return cp_attn_out.new_empty(0, H_per_rank, D) - - partial_o = torch.empty( - B, H_per_rank, N, D, dtype=cp_attn_out.dtype, device=cp_attn_out.device - ) - partial_lse = torch.empty( - B, H_per_rank, N, dtype=torch.float32, device=cp_attn_out.device - ) - dcp_pack_a2a_send( - cp_attn_out, - cp_attn_lse, - partial_o.permute(2, 0, 1, 3), - partial_lse.permute(2, 0, 1), - ) - return fi_a2a_fused_reduce(partial_o, partial_lse, is_lse_base_on_e) diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 9d0773ebb466..91d30c311a77 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -325,8 +325,7 @@ def _pre_initialize_flashinfer_allreduce_workspace(self): def _pre_initialize_fi_a2a_workspace(self): """Create the fused fi_a2a DCP reduce workspaces; must run before CG - capture (workspace rendezvous and the group's first reduce are - collective and not capturable). + capture, since creation is collective and synchronizes the stream. """ parallel = get_parallel() if not parallel.dcp_enabled or parallel.dcp_comm_backend != "fi_a2a": @@ -339,13 +338,13 @@ def _pre_initialize_fi_a2a_workspace(self): mr = self.model_runner init_fi_a2a_workspace( - parallel.dcp_group, + cp_group=parallel.dcp_group, # Batches wider than the captured graphs run as several fused calls. max_tokens=self._widest_decode_rows() or _FI_A2A_EAGER_ONLY_MAX_TOKENS, local_heads=mr.model_config.num_attention_heads // parallel.attn_tp_size, head_dim=mr.model_config.kv_lora_rank, dtype=mr.dtype, - probe_stream=get_or_create_global_graph_capture_stream(), + capture_stream=get_or_create_global_graph_capture_stream(), ) def _pre_initialize_pcie_ipc_workspace(self): diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index c3d2feee7489..212aeec4eea5 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1005,7 +1005,7 @@ def fi_a2a_platform_blocker( ) -> Optional[str]: """Why this host cannot run the fused fi_a2a DCP reduce, or None if it can. - Runs at argument resolution: no CUDA context and no flashinfer import. + Runs at argument resolution, so it finds FlashInfer's op without importing it. """ if not get_platform().is_sm100: return "requires a Blackwell (SM100-family) GPU" @@ -1018,13 +1018,6 @@ def fi_a2a_platform_blocker( f"{tp_size_per_node} TP ranks per node (tp_size={tp_size}, " f"pp_size={pp_size}, nnodes={nnodes})" ) - # torch 2.14 is the first release whose symmetric memory exposes the NCCL - # device communicator and window offsets the fused kernel is built on. - if torch_release < (2, 14): - return ( - "requires torch>=2.14 (NCCL device communicator in torch symmetric " - f"memory); found torch {torch.__version__}" - ) if not _flashinfer_has_fused_dcp_reduce(): return ( "requires a FlashInfer build that provides " diff --git a/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py b/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py index 940bc44bc278..0600ed3e7b99 100644 --- a/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py +++ b/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py @@ -1,9 +1,9 @@ """DCP fi_a2a: FlashInfer's fused all-to-all + LSE reduce must equal the all-gather + Triton merge it replaces, including chunked, graph-replayed and -fully-masked rows. +fully-masked rows, next to custom all-reduce v2 on the same symmetric memory. Run with ``python test_dcp_fi_a2a_reduce.py --num-gpu 2,4,8``. Skips unless the -host can run the fused reduce (Blackwell, torch>=2.14, FlashInfer with the op). +host can run the fused reduce (Blackwell, FlashInfer with the op). """ import os @@ -14,62 +14,54 @@ from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton from sglang.srt.distributed import parallel_state as ps +from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, +) from sglang.srt.layers.dcp import comm +from sglang.srt.utils.common import fi_a2a_platform_blocker from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kernels.utils import multigpu_pytest_main from sglang.test.test_utils import publish_build_topology register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200") -# Workspace geometry: 3 tokens x 16 heads = 48 rows, so the (H, B) cases below -# cover a single call, an exactly-full call and multi-chunk calls. +# Workspace geometry: 3 tokens x 16 heads = 48 rows, so the (heads, batch) cases +# below cover a single call, an exactly-full call and multi-chunk calls. MAX_TOKENS, LOCAL_HEADS, HEAD_DIM = 3, 16, 512 - -def _blocker(): - if not torch.cuda.is_available(): - return "needs CUDA" - from sglang.srt.utils.common import fi_a2a_platform_blocker - - world_size = int(os.environ.get("WORLD_SIZE", "1")) - return fi_a2a_platform_blocker( - dcp_size=world_size, tp_size=world_size, pp_size=1, nnodes=1 - ) - - -_reason = _blocker() -if _reason is not None: - pytest.skip(f"fused fi_a2a reduce {_reason}", allow_module_level=True) +_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) +_SKIP_REASON = fi_a2a_platform_blocker( + dcp_size=_WORLD_SIZE, tp_size=_WORLD_SIZE, pp_size=1, nnodes=1 +) +pytestmark = pytest.mark.skipif( + _SKIP_REASON is not None, reason=f"fused fi_a2a reduce {_SKIP_REASON}" +) @pytest.fixture(scope="module") def dcp(): - import torch.distributed._symmetric_memory as symm_mem - rank = int(os.environ["RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) - # As bootstrap does for fi_a2a: pick the backend before any symmetric allocation. - symm_mem.set_backend("NCCL") - ps.set_custom_all_reduce(False) ps.init_distributed_environment( - world_size=world_size, + world_size=_WORLD_SIZE, rank=rank, local_rank=local_rank, distributed_init_method="env://", ) - publish_build_topology(tp_size=world_size, world_rank=rank) + publish_build_topology(tp_size=_WORLD_SIZE, world_rank=rank) + # Builds custom all-reduce v2, whose symmetric memory is allocated before the + # fused workspaces, as in serving. ps.initialize_model_parallel() group = ps.get_tp_group() capture_stream = torch.cuda.Stream() comm.init_fi_a2a_workspace( - group, + cp_group=group, max_tokens=MAX_TOKENS, local_heads=LOCAL_HEADS, head_dim=HEAD_DIM, dtype=torch.bfloat16, - probe_stream=capture_stream, + capture_stream=capture_stream, ) yield group, capture_stream ps.destroy_model_parallel() @@ -116,8 +108,8 @@ def _fused(group, out, lse, *, is_lse_base_on_e): "heads,batch", [(16, 1), (16, 3), (16, 7), (2, 1), (2, 24), (2, 49)] ) def test_matches_all_gather_reference(dcp, dtype, is_lse_base_on_e, heads, batch): - """Guards the head-to-peer mapping and the row chunking (B=7 at 16 heads and - B=49 at 2 heads exceed the 48-row workspace).""" + """Guards the head-to-peer mapping and the token chunking (B=7 at 16 heads + and B=49 at 2 heads exceed the 48-row workspace).""" group, _ = dcp out, lse = _partials(group, batch=batch, heads=heads, dtype=dtype, seed=11) expected = _reference(group, out, lse, is_lse_base_on_e=is_lse_base_on_e) @@ -182,17 +174,40 @@ def test_graph_replay_and_eager_calls_interleave(dcp): def test_third_stream_raises_before_the_collective(dcp): - """Only the capture and serving streams own workspaces; a third stream must - fail on the host on every rank, not spin in the kernel.""" + """Only the capture and serving streams own a workspace; a third stream + must fail on the host on every rank, not spin in the kernel.""" group, _ = dcp out, lse = _partials( group, batch=1, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=61 ) _fused(group, out, lse, is_lse_base_on_e=False) # binds the serving stream with torch.cuda.stream(torch.cuda.Stream()): - with pytest.raises(RuntimeError, match="third stream"): + with pytest.raises(RuntimeError, match="one ordered CUDA stream"): _fused(group, out, lse, is_lse_base_on_e=False) +def test_custom_all_reduce_v2_stays_exact_around_fused_calls(dcp): + """The fused workspaces share torch symmetric memory with custom all-reduce + v2, which allocates first; neither may disturb the other.""" + group, _ = dcp + ca_comm = group.ca_comm + if not isinstance(ca_comm, CustomAllReduceV2) or ca_comm.disabled: + pytest.skip("custom all-reduce v2 is not enabled on this group") + world_size = group.world_size + x = torch.full( + (4096,), group.rank_in_group + 1, dtype=torch.bfloat16, device="cuda" + ) + want_sum = torch.full_like(x, world_size * (world_size + 1) // 2) + out, lse = _partials( + group, batch=2, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=81 + ) + expected = _reference(group, out, lse, is_lse_base_on_e=False) + for _ in range(2): + assert torch.equal(ca_comm.custom_all_reduce(x), want_sum) + got = _fused(group, out, lse, is_lse_base_on_e=False) + torch.testing.assert_close(got, expected, rtol=1e-2, atol=1e-2) + assert torch.equal(ca_comm.custom_all_reduce(x), want_sum) + + if __name__ == "__main__": multigpu_pytest_main(__name__, __file__, num_gpus=(2, 4)) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index e6449f5bbe90..90d1923cb887 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -4004,11 +4004,13 @@ def test_without_all_to_all_env_is_untouched(self): class TestDcpCommBackendDefault(CustomTestCase): - def _args(self, **fields): - return ServerArgs(model_path="dummy", tp_size=8, **fields) + _NO_FUSED_OP = ( + "requires a FlashInfer build that provides " + "flashinfer.comm.decode_cp_a2a_lse_reduce" + ) def _resolved(self, **fields): - args = self._args(**fields) + args = ServerArgs(model_path="dummy", tp_size=8, **fields) parallel_hook.handle_decode_context_parallelism(args) return resolution_result(args, "dcp_comm_backend") @@ -4039,43 +4041,35 @@ def test_fi_a2a_where_the_fused_reduce_can_run(self): def test_a2a_when_the_fused_reduce_cannot_run(self): """fi_a2a has no fallback kernel, so resolving to it on a host that cannot run the fused reduce would fail every default DCP start.""" - host, mla = self._fused_reduce(platform_blocker="requires torch>=2.14") + host, mla = self._fused_reduce(platform_blocker=self._NO_FUSED_OP) with host, mla: self.assertEqual(self._resolved(dcp_size=4), "a2a") @override_platform(is_cuda=True, is_hip=False) def test_configs_the_fused_reduce_cannot_serve_resolve_to_a2a(self): cases = { - "non-MLA model": ({}, {}, False), - "--enable-pdmux": ({"enable_pdmux": True}, {}, True), - "--enable-torch-symm-mem": ({"enable_torch_symm_mem": True}, {}, True), - "exported NCCL_CUMEM_ENABLE=0": ({}, {"NCCL_CUMEM_ENABLE": "0"}, True), + "non-MLA model": ({}, False, False), + "--enable-pdmux": ({"enable_pdmux": True}, False, True), + "PP DeepGEMM warmup at pp_size 2": ({"pp_size": 2}, True, True), } - for name, (fields, env, is_mla) in cases.items(): + for name, (fields, pp_warmup, is_mla) in cases.items(): with self.subTest(name): host, mla = self._fused_reduce(mla=is_mla) - with host, mla, patch.dict(os.environ, env): + warmup = envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.override(pp_warmup) + with host, mla, warmup: self.assertEqual(self._resolved(dcp_size=4, **fields), "a2a") with self.assertRaisesRegex(ValueError, "--dcp-comm-backend a2a"): self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a", **fields) @override_platform(is_cuda=True, is_hip=False) def test_explicit_fi_a2a_names_the_blocker(self): - host, mla = self._fused_reduce(platform_blocker="requires torch>=2.14") + host, mla = self._fused_reduce(platform_blocker=self._NO_FUSED_OP) with host, mla: - with self.assertRaisesRegex(ValueError, "fi_a2a requires torch>=2.14"): + with self.assertRaisesRegex( + ValueError, "fi_a2a requires a FlashInfer build that provides" + ): self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a") - @override_platform(is_cuda=True, is_hip=False) - def test_explicit_fi_a2a_defaults_nccl_cumem_on(self): - """The fused workspace is an NCCL symmetric window; engine startup would - otherwise force NCCL_CUMEM_ENABLE=0 for a non-symm-mem server.""" - host, mla = self._fused_reduce() - with host, mla, patch.dict(os.environ): - os.environ.pop("NCCL_CUMEM_ENABLE", None) - self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a") - self.assertEqual(os.environ["NCCL_CUMEM_ENABLE"], "1") - @override_platform(is_cuda=False, is_hip=False) def test_ag_rs_off_cuda(self): host, mla = self._fused_reduce(platform_blocker="requires a Blackwell GPU") diff --git a/test/registered/unit/utils/test_common.py b/test/registered/unit/utils/test_common.py index 8b335a65115b..a274e7dcb4ed 100644 --- a/test/registered/unit/utils/test_common.py +++ b/test/registered/unit/utils/test_common.py @@ -1,5 +1,9 @@ +import contextlib +import importlib import io +import os import sys +import tempfile import unittest from array import array from unittest import mock @@ -7,9 +11,12 @@ import torch from PIL import Image +from sglang.srt.runtime_context import override_platform from sglang.srt.utils.common import ( + _flashinfer_has_fused_dcp_reduce, _get_device_sm_via_nvml, _load_image, + fi_a2a_platform_blocker, flatten_arrays_to_int64_tensor, get_device_sm_nvidia_smi, get_nvidia_driver_version_str, @@ -251,5 +258,76 @@ def boom(index): self.assertFalse(fake.initialized, "must not query NVML without the mapping") +class TestFiA2aPlatformBlocker(CustomTestCase): + @contextlib.contextmanager + def _flashinfer_package(self, *, with_fused_op: bool): + """Install a flashinfer package that fails when imported.""" + with tempfile.TemporaryDirectory() as root: + package = os.path.join(root, "flashinfer") + os.makedirs(os.path.join(package, "comm")) + with open(os.path.join(package, "__init__.py"), "w") as f: + f.write("raise ImportError('flashinfer was imported')\n") + if with_fused_op: + open(os.path.join(package, "comm", "dcp_lse_reduce.py"), "w").close() + with ( + mock.patch.dict(sys.modules), + mock.patch.object(sys, "path", [root, *sys.path]), + ): + for name in [n for n in sys.modules if n.split(".")[0] == "flashinfer"]: + del sys.modules[name] + importlib.invalidate_caches() + _flashinfer_has_fused_dcp_reduce.cache_clear() + try: + yield + finally: + _flashinfer_has_fused_dcp_reduce.cache_clear() + + @override_platform(is_sm100=True) + def test_finds_the_fused_op_without_importing_flashinfer(self): + """The blocker runs while the launcher resolves arguments, so it must + find FlashInfer's fused op on disk without importing flashinfer.""" + for with_fused_op in (True, False): + with ( + self.subTest(with_fused_op=with_fused_op), + self._flashinfer_package(with_fused_op=with_fused_op), + ): + reason = fi_a2a_platform_blocker( + dcp_size=8, tp_size=8, pp_size=1, nnodes=1 + ) + self.assertNotIn("flashinfer", sys.modules) + if with_fused_op: + self.assertIsNone(reason) + else: + self.assertIn("decode_cp_a2a_lse_reduce", reason) + + @override_platform(is_sm100=True) + def test_dcp_group_must_share_one_nvlink_domain(self): + # (MNNVL fabric, dcp_size, tp_size, pp_size, nnodes, blocked) + cases = ( + (False, 8, 16, 1, 2, False), + (False, 16, 16, 1, 2, True), + (False, 8, 8, 2, 2, False), + (True, 16, 16, 1, 2, False), + ) + for fabric, dcp_size, tp_size, pp_size, nnodes, blocked in cases: + with ( + self.subTest( + fabric=fabric, dcp_size=dcp_size, pp_size=pp_size, nnodes=nnodes + ), + mock.patch( + "sglang.srt.utils.common.is_mnnvl_fabric_device", + return_value=fabric, + ), + mock.patch( + "sglang.srt.utils.common._flashinfer_has_fused_dcp_reduce", + return_value=True, + ), + ): + reason = fi_a2a_platform_blocker( + dcp_size=dcp_size, tp_size=tp_size, pp_size=pp_size, nnodes=nnodes + ) + self.assertEqual(reason is not None, blocked, reason) + + if __name__ == "__main__": unittest.main() From e98cc370898cc82cd34485922ee370a1fd591b7e Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Tue, 6 Oct 2026 16:22:47 +0700 Subject: [PATCH 07/10] [DCP] Read fi_a2a partials from the attention output in place The fused reduce accepts strided [B, H/cp, cp, D] and [B, H/cp, cp] views, so view the attention output and LSE instead of packing them. Head h still goes to peer h // (H/cp). Inputs with a non-unit inner stride, or rows off an 8-byte boundary, are copied first; wide eager batches still run as token slices. Drop the pack test for the layout only fi_a2a used. --- python/sglang/srt/layers/dcp/comm.py | 33 ++++++++++------- .../ops/attention/test_dcp_lse_combine.py | 35 ------------------- .../communication/test_dcp_fi_a2a_reduce.py | 30 ++++++++++++++-- 3 files changed, 49 insertions(+), 49 deletions(-) diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index b41e385965c3..b7983e03bf47 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/python/sglang/srt/layers/dcp/comm.py @@ -473,20 +473,29 @@ def create_workspace() -> torch.Tensor: ) -def _fi_a2a_peer_partials( +def _fi_a2a_peer_views( cp_attn_out: torch.Tensor, cp_attn_lse: torch.Tensor, cp_size: int ) -> tuple[torch.Tensor, torch.Tensor]: - """[B, H, D] partials and [B, H] LSE as the fused reduce's [B, H/cp, cp, D] - and [B, H/cp, cp]; head h goes to peer h // (H/cp).""" - batch, heads, head_dim = cp_attn_out.shape + """View [B, H, D] partials and their [B, H] LSE as the fused reduce's + [B, H/cp, cp, D] and [B, H/cp, cp]; head h goes to peer h // (H/cp).""" + batch, heads, _ = cp_attn_out.shape local_heads = heads // cp_size - partial_o = cp_attn_out.new_empty(batch, local_heads, cp_size, head_dim) - partial_lse = cp_attn_lse.new_empty(batch, local_heads, cp_size) - dcp_pack_a2a_send( - cp_attn_out, - cp_attn_lse, - partial_o.permute(2, 0, 1, 3), - partial_lse.permute(2, 0, 1), + itemsize = cp_attn_out.element_size() + # The kernel reads every row in 8-byte words along a unit-stride last dim. + if cp_attn_out.stride(-1) != 1 or any( + n % 8 + for n in ( + cp_attn_out.data_ptr(), + cp_attn_out.stride(0) * itemsize, + cp_attn_out.stride(1) * itemsize, + ) + ): + cp_attn_out = cp_attn_out.clone(memory_format=torch.contiguous_format) + partial_o = cp_attn_out.unflatten(1, (cp_size, local_heads)).transpose(1, 2) + partial_lse = ( + cp_attn_lse.reshape(batch, heads) + .unflatten(1, (cp_size, local_heads)) + .transpose(1, 2) ) return partial_o, partial_lse @@ -521,7 +530,7 @@ def fi_a2a_lse_reduce( # The kernel rejects an empty batch; every DCP rank sees the same batch. if batch == 0: return cp_attn_out.new_empty(0, local_heads, head_dim) - partial_o, partial_lse = _fi_a2a_peer_partials( + partial_o, partial_lse = _fi_a2a_peer_views( cp_attn_out=cp_attn_out, cp_attn_lse=cp_attn_lse, cp_size=cp_size ) if torch.cuda.current_stream().cuda_stream == state.capture_stream: diff --git a/test/registered/kernels/ops/attention/test_dcp_lse_combine.py b/test/registered/kernels/ops/attention/test_dcp_lse_combine.py index 9be1ad117703..81fec5b61b60 100644 --- a/test/registered/kernels/ops/attention/test_dcp_lse_combine.py +++ b/test/registered/kernels/ops/attention/test_dcp_lse_combine.py @@ -398,41 +398,6 @@ def test_pack_matches_the_copy_formulation_it_replaces(self): torch.equal(lane, lse.view(B, N, H_per_rank).permute(1, 0, 2)) ) - def test_pack_serves_the_fused_reduce_layout(self): - """fi_a2a's fused reduce reads partial_o [B, H/N, N, D] and a contiguous - partial_lse [B, H/N, N] (peer axis second-to-last); the pack writes both - through the [N, B, H/N, ...] views it is handed.""" - from sglang.kernels.ops.attention.dcp_kernels import dcp_pack_a2a_send - - for N, B, H_per_rank, D in ((2, 4, 8, 128), (4, 1, 16, 512)): - with self.subTest(N=N, B=B, H_per_rank=H_per_rank, D=D): - H = H_per_rank * N - out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16) - lse = torch.randn(B, H, device=self.device, dtype=torch.float32) - - partial_o = torch.empty( - B, H_per_rank, N, D, dtype=torch.bfloat16, device=self.device - ) - partial_lse = torch.empty( - B, H_per_rank, N, dtype=torch.float32, device=self.device - ) - dcp_pack_a2a_send( - out, - lse, - partial_o.permute(2, 0, 1, 3), - partial_lse.permute(2, 0, 1), - ) - - want_o = out.view(B, N, H_per_rank, D).permute(0, 2, 1, 3) - want_lse = lse.view(B, N, H_per_rank).permute(0, 2, 1) - self.assertTrue( - torch.equal( - partial_o.view(torch.uint8), - want_o.contiguous().view(torch.uint8), - ) - ) - self.assertTrue(torch.equal(partial_lse, want_lse)) - def test_buffers_have_fixed_data_ptrs(self): """Pre-allocated buffer data_ptr must not change -- required for graph replay.""" from sglang.srt.layers.dcp import dcp_a2a_lse_reduce diff --git a/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py b/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py index 0600ed3e7b99..1cefc3bf9007 100644 --- a/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py +++ b/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py @@ -1,6 +1,6 @@ """DCP fi_a2a: FlashInfer's fused all-to-all + LSE reduce must equal the -all-gather + Triton merge it replaces, including chunked, graph-replayed and -fully-masked rows, next to custom all-reduce v2 on the same symmetric memory. +all-gather + Triton merge it replaces, including chunked, strided, graph-replayed +and fully-masked inputs, next to custom all-reduce v2 on the same symmetric memory. Run with ``python test_dcp_fi_a2a_reduce.py --num-gpu 2,4,8``. Skips unless the host can run the fused reduce (Blackwell, FlashInfer with the op). @@ -132,6 +132,32 @@ def test_rows_empty_on_every_rank_are_zero(dcp): torch.testing.assert_close(got[:, 1:], expected[:, 1:], rtol=1e-2, atol=1e-2) +@pytest.mark.parametrize( + "layout", ["strided_head_dim", "misaligned_rows", "lse_with_unit_dim"] +) +def test_layouts_the_kernel_cannot_read_are_adapted(dcp, layout): + """The kernel reads unit-stride rows that start on 8-byte boundaries, and + FlashMLA returns its LSE as [B, H, 1]; such inputs are copied or reshaped, + not rejected.""" + group, _ = dcp + out, lse = _partials( + group, batch=3, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=71 + ) + expected = _reference(group, out, lse, is_lse_base_on_e=True) + if layout == "strided_head_dim": + wide = out.new_zeros(*out.shape[:-1], 2 * HEAD_DIM) + wide[..., ::2] = out + out = wide[..., ::2] + elif layout == "misaligned_rows": + flat = out.new_empty(out.numel() + 1) + flat[1:] = out.flatten() + out = flat[1:].view(out.shape) # 2 bytes past an 8-byte boundary + else: + lse = lse.unsqueeze(-1) + got = _fused(group, out, lse, is_lse_base_on_e=True) + torch.testing.assert_close(got, expected, rtol=1e-2, atol=1e-2) + + def test_graph_replay_and_eager_calls_interleave(dcp): """The captured graph and eager decode use different workspaces on one ordered stream; replays and eager calls must not corrupt each other.""" From 3bfdabae487957aa8595386ad0d7a5980f0151e4 Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Tue, 6 Oct 2026 16:24:12 +0700 Subject: [PATCH 08/10] [DCP] Document fi_a2a's FlashInfer requirement fi_a2a needs Blackwell, an MLA model, a DCP group inside one NVLink domain and a flashinfer-python with the fused decode_cp_a2a_lse_reduce; it no longer changes the symmetric-memory backend, custom all-reduce or NCCL settings, so drop those notes. Describe the two workspaces, the chunked eager batches and the configurations that resolve to a2a. --- docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx | 2 +- docs/docs/advanced_features/dcp.mdx | 11 +++-------- docs/docs/advanced_features/server_arguments.mdx | 2 +- python/sglang/srt/arg_groups/fields/parallel.py | 10 ++++------ 4 files changed, 9 insertions(+), 16 deletions(-) diff --git a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx index 7012aa92db10..8d5d9f7a9ab4 100644 --- a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx +++ b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx @@ -205,7 +205,7 @@ Speculation: DSPARK holds block size + 1 (= 8) intermediate states per request **Blackwell DCP notes** — the DCP cells are Balanced and High-Throughput on every Blackwell platform, in both the `Unified` and `Decode` roles: - DCP is the only axis that shards the TP-replicated MLA KV; Low-Latency skips it. -- Leave `--dcp-comm-backend` unset: it resolves to `fi_a2a` where FlashInfer's fused reduce can run (torch 2.14+ and a FlashInfer build with `decode_cp_a2a_lse_reduce`), otherwise `a2a`. +- Leave `--dcp-comm-backend` unset: it resolves to `fi_a2a` (FlashInfer's fused all-to-all + LSE merge) where the DCP group shares one NVLink domain, GB200/GB300 fabric or one node, and `a2a` otherwise. - No `--enable-symm-mem` under DCP (force-disabled for decode-graph correctness). - Explicit `tokenspeed_mla` force-rewrites `--kv-cache-dtype` to fp8; the default `cutedsl_mla` serves either dtype. - Calculator ratios run well above 1 here (`r > 1` is legal): `bfloat16` state buys admission, `fp8` KV buys context. diff --git a/docs/docs/advanced_features/dcp.mdx b/docs/docs/advanced_features/dcp.mdx index 29a719a0a9ce..1a6c441d8c55 100644 --- a/docs/docs/advanced_features/dcp.mdx +++ b/docs/docs/advanced_features/dcp.mdx @@ -132,16 +132,11 @@ Normally each layer all-gathers its absorbed and rotary query components. `--dcp | `ag_rs` | Query AG + LSE AG + FP32 output RS | Fallback for platforms without the a2a path; opt in with `--dcp-comm-backend ag_rs` | | `a2a`, gathered Q | Query AG + packed NCCL A2A | Two collectives | | `a2a`, replicated Q | One packed NCCL A2A | Canonical low-latency path | -| `fi_a2a`, replicated Q | One FlashInfer kernel: NCCL LSA all-to-all and LSE merge fused | Requires Blackwell, the DCP group inside one NVLink domain (GB200/GB300 fabric, or a single node), an MLA model, torch 2.14 or newer, and a FlashInfer build that provides `decode_cp_a2a_lse_reduce` | +| `fi_a2a`, replicated Q | One FlashInfer kernel: all-to-all and LSE merge fused | Requires Blackwell, an MLA model, a DCP group inside one NVLink domain (GB200/GB300 fabric, or one node), and a `flashinfer-python` with the fused `decode_cp_a2a_lse_reduce` | -When `--dcp-comm-backend` is unset, every DCP launch resolves it the same way: `fi_a2a` where the fused reduce can run, otherwise `a2a` on CUDA/ROCm, and `ag_rs` only on platforms without the a2a path. The startup log names the missing requirement when `fi_a2a` is not selected. +When `--dcp-comm-backend` is unset, every DCP launch resolves it the same way: `fi_a2a` where the fused reduce can run, otherwise `a2a` on CUDA/ROCm, and `ag_rs` only on platforms without the a2a path. The startup log names the missing requirement when `fi_a2a` is not selected. Query replication stays model-resolved: Kimi K3 enables it by default on the a2a family, and its decode backend is `cutedsl_mla`. -`fi_a2a` exchanges and merges the partials in one cooperative kernel over NCCL symmetric memory, instead of an all-to-all followed by a separate merge kernel. Running it changes a few process-wide settings: - -- torch symmetric memory uses the NCCL backend, so custom all-reduce uses its v1 kernel instead of `CustomAllReduceV2`; -- `NCCL_CUMEM_ENABLE` defaults to `1`, because the fused workspace is an NCCL symmetric window; -- `--enable-pdmux` and `--enable-torch-symm-mem` are rejected, since they need more CUDA streams or the CUDA symmetric-memory backend; -- decode batches wider than the captured CUDA graphs run as several fused calls. Query replication stays model-resolved: Kimi K3 enables it by default on the a2a family, and its decode backend is `cutedsl_mla`. +`fi_a2a` keeps two fused-reduce workspaces, one for CUDA graph replays and one for eager decode, each sized for the widest captured decode batch; the startup log reports their size. Eager batches wider than that run as several fused calls. A workspace serves only one CUDA stream, so `--enable-pdmux`, and `SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP` with `--pp-size` above 1, resolve to `a2a`. DCP adds a context-independent decode collective while cutting context-dependent KV storage and reads by about `c`. Extend is outside that decode cost model: it may gather cached prefix shards, restore token order, and append new tokens, and that work grows with prefix length. diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 986c4d2f754e..17217b752e83 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -668,7 +668,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--dcp-comm-backend` - Communication backend for the DCP attention reduction: AllGather + ReduceScatter, fused NCCL All-to-All, or FlashInfer's fused all-to-all + LSE merge (Blackwell, one NVLink domain, MLA, torch 2.14+, FlashInfer with decode_cp_a2a_lse_reduce). Unset resolves to fi_a2a where it can run, else a2a on CUDA/ROCm, else ag_rs. + Communication backend for the DCP attention reduction: AllGather + ReduceScatter, fused NCCL All-to-All, or FlashInfer's fused all-to-all + LSE merge (Blackwell, an MLA model, one NVLink domain, and FlashInfer with decode_cp_a2a_lse_reduce). Unset resolves to fi_a2a where it can run, else a2a on CUDA/ROCm, else ag_rs. None ag_rs, a2a, fi_a2a diff --git a/python/sglang/srt/arg_groups/fields/parallel.py b/python/sglang/srt/arg_groups/fields/parallel.py index 7785ba3788c9..393886e9c0e0 100644 --- a/python/sglang/srt/arg_groups/fields/parallel.py +++ b/python/sglang/srt/arg_groups/fields/parallel.py @@ -130,12 +130,10 @@ class Parallel(msgspec.Struct): help="Communication backend for the decode context-parallel (DCP) " "attention reduction: 'ag_rs' (AllGather + ReduceScatter), 'a2a' " "(fused NCCL All-to-All exchange of output+LSE + local Triton LSE " - "combine), or 'fi_a2a' (FlashInfer's decode_cp_a2a_lse_reduce: " - "exchange and LSE merge in one kernel over NCCL symmetric memory; " - "requires Blackwell, the DCP group inside one NVLink domain, an MLA " - "model, torch>=2.14 and a FlashInfer build with that op; it moves " - "torch symmetric memory to the NCCL backend, so custom all-reduce " - "uses its v1 kernel). Unset resolves to 'fi_a2a' where it can run, " + "combine), or 'fi_a2a' (FlashInfer's decode_cp_a2a_lse_reduce, which " + "exchanges and merges the partials in one kernel; requires Blackwell, " + "an MLA model, the DCP group inside one NVLink domain and a FlashInfer " + "build with that op). Unset resolves to 'fi_a2a' where it can run, " "else 'a2a' on CUDA/ROCm, else 'ag_rs'.", choices=["ag_rs", "a2a", "fi_a2a"], resolvable=True, From 5166e20e230a382e2ee2a6a033895c7909b839fa Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Tue, 6 Oct 2026 23:09:04 +0700 Subject: [PATCH 09/10] [DCP] Document fi_a2a's workspace memory and trim its tests The DCP docs now give the size of the two fused-reduce workspaces: they follow --cuda-graph-max-bs-decode, also with decode CUDA graphs disabled, and come out of the memory left after KV-cache sizing. The Kimi K3 notes name the FlashInfer op the default backend needs. Drop the copy for strided or misaligned partials: the only caller passes contiguous ones, and FlashInfer checks its inputs. Comments that still described fi_a2a as packing and running the Triton combine now describe a2a only. Tests: keep the bf16 reference cases that cover the head mapping and chunking, and drop the cases that only exercised FlashInfer or the removed copy. Fail rather than skip when custom all-reduce v2 should be running. Add the PP case that keeps fi_a2a without the DeepGEMM warmup. --- .../autoregressive/Moonshotai/Kimi-K3.mdx | 2 +- docs/docs/advanced_features/dcp.mdx | 4 +- .../kernels/ops/attention/dcp_kernels.py | 2 +- python/sglang/srt/layers/dcp/comm.py | 23 ++--- .../srt/model_executor/runner/base_runner.py | 4 +- .../attention_forward_methods/forward_mla.py | 2 +- .../communication/test_dcp_fi_a2a_reduce.py | 94 ++++++------------- .../unit/server_args/test_server_args.py | 9 ++ 8 files changed, 54 insertions(+), 86 deletions(-) diff --git a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx index 8d5d9f7a9ab4..e593c2200823 100644 --- a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx +++ b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx @@ -205,7 +205,7 @@ Speculation: DSPARK holds block size + 1 (= 8) intermediate states per request **Blackwell DCP notes** — the DCP cells are Balanced and High-Throughput on every Blackwell platform, in both the `Unified` and `Decode` roles: - DCP is the only axis that shards the TP-replicated MLA KV; Low-Latency skips it. -- Leave `--dcp-comm-backend` unset: it resolves to `fi_a2a` (FlashInfer's fused all-to-all + LSE merge) where the DCP group shares one NVLink domain, GB200/GB300 fabric or one node, and `a2a` otherwise. +- Leave `--dcp-comm-backend` unset: it resolves to `fi_a2a` (FlashInfer's fused all-to-all + LSE merge) where the DCP group shares one NVLink domain, GB200/GB300 fabric or one node, and the installed `flashinfer-python` provides `decode_cp_a2a_lse_reduce`; otherwise to `a2a` (the startup log names the missing requirement). - No `--enable-symm-mem` under DCP (force-disabled for decode-graph correctness). - Explicit `tokenspeed_mla` force-rewrites `--kv-cache-dtype` to fp8; the default `cutedsl_mla` serves either dtype. - Calculator ratios run well above 1 here (`r > 1` is legal): `bfloat16` state buys admission, `fp8` KV buys context. diff --git a/docs/docs/advanced_features/dcp.mdx b/docs/docs/advanced_features/dcp.mdx index 1a6c441d8c55..779f2ccb1089 100644 --- a/docs/docs/advanced_features/dcp.mdx +++ b/docs/docs/advanced_features/dcp.mdx @@ -134,9 +134,9 @@ Normally each layer all-gathers its absorbed and rotary query components. `--dcp | `a2a`, replicated Q | One packed NCCL A2A | Canonical low-latency path | | `fi_a2a`, replicated Q | One FlashInfer kernel: all-to-all and LSE merge fused | Requires Blackwell, an MLA model, a DCP group inside one NVLink domain (GB200/GB300 fabric, or one node), and a `flashinfer-python` with the fused `decode_cp_a2a_lse_reduce` | -When `--dcp-comm-backend` is unset, every DCP launch resolves it the same way: `fi_a2a` where the fused reduce can run, otherwise `a2a` on CUDA/ROCm, and `ag_rs` only on platforms without the a2a path. The startup log names the missing requirement when `fi_a2a` is not selected. Query replication stays model-resolved: Kimi K3 enables it by default on the a2a family, and its decode backend is `cutedsl_mla`. +When `--dcp-comm-backend` is unset, every DCP launch resolves it the same way: `fi_a2a` where the fused reduce can run, otherwise `a2a` on CUDA/ROCm, and `ag_rs` only on platforms without the a2a path. On CUDA, the startup log names the missing requirement when `fi_a2a` is not selected. Query replication stays model-resolved: Kimi K3 enables it by default on the a2a family, and its decode backend is `cutedsl_mla`. -`fi_a2a` keeps two fused-reduce workspaces, one for CUDA graph replays and one for eager decode, each sized for the widest captured decode batch; the startup log reports their size. Eager batches wider than that run as several fused calls. A workspace serves only one CUDA stream, so `--enable-pdmux`, and `SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP` with `--pp-size` above 1, resolve to `a2a`. +`fi_a2a` keeps two fused-reduce workspaces, one for CUDA graph replays and one for eager decode, each sized for the widest decode CUDA graph batch (`--cuda-graph-max-bs-decode`, times the tokens each request verifies under speculative decoding), also with `--disable-decode-cuda-graph`; eager batches wider than that run as several fused calls. Each workspace takes about `2 × dcp_size × tokens × heads × 1.2 KB` for the 512-wide BF16 latent, where `heads` counts the attention heads per attention-TP rank: 150 MiB at 512 tokens, 16 heads and DCP 8. Both are allocated after KV-cache sizing, out of the memory `--mem-fraction-static` leaves free, and the startup log reports their size. A workspace serves only one CUDA stream, so `--enable-pdmux`, and `SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP` with `--pp-size` above 1, resolve to `a2a`. DCP adds a context-independent decode collective while cutting context-dependent KV storage and reads by about `c`. Extend is outside that decode cost model: it may gather cached prefix shards, restore token order, and append new tokens, and that work grows with prefix length. diff --git a/python/sglang/kernels/ops/attention/dcp_kernels.py b/python/sglang/kernels/ops/attention/dcp_kernels.py index 30586e886e85..75b02330f2a7 100644 --- a/python/sglang/kernels/ops/attention/dcp_kernels.py +++ b/python/sglang/kernels/ops/attention/dcp_kernels.py @@ -385,7 +385,7 @@ def correct_attn_out( # A2A DCP reduce: LSE-weighted combine of N partial attention outputs -# (used by the a2a / fi_a2a communication backends, see comm.py). +# (used by the a2a communication backend, see comm.py). def _lse_pack_dim(output_dtype: torch.dtype) -> int: diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index b7983e03bf47..60eb4968d336 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/python/sglang/srt/layers/dcp/comm.py @@ -476,21 +476,10 @@ def create_workspace() -> torch.Tensor: def _fi_a2a_peer_views( cp_attn_out: torch.Tensor, cp_attn_lse: torch.Tensor, cp_size: int ) -> tuple[torch.Tensor, torch.Tensor]: - """View [B, H, D] partials and their [B, H] LSE as the fused reduce's - [B, H/cp, cp, D] and [B, H/cp, cp]; head h goes to peer h // (H/cp).""" + """View contiguous [B, H, D] partials and their [B, H] LSE as the fused + reduce's [B, H/cp, cp, D] and [B, H/cp, cp]; head h goes to peer h // (H/cp).""" batch, heads, _ = cp_attn_out.shape local_heads = heads // cp_size - itemsize = cp_attn_out.element_size() - # The kernel reads every row in 8-byte words along a unit-stride last dim. - if cp_attn_out.stride(-1) != 1 or any( - n % 8 - for n in ( - cp_attn_out.data_ptr(), - cp_attn_out.stride(0) * itemsize, - cp_attn_out.stride(1) * itemsize, - ) - ): - cp_attn_out = cp_attn_out.clone(memory_format=torch.contiguous_format) partial_o = cp_attn_out.unflatten(1, (cp_size, local_heads)).transpose(1, 2) partial_lse = ( cp_attn_lse.reshape(batch, heads) @@ -562,9 +551,11 @@ def dcp_a2a_lse_reduce( cuda_graph_buffers: Optional[dict] = None, comm_backend: str = "a2a", ) -> torch.Tensor: - """A2A DCP reduce: all-to-all exchange of head partials, then local Triton - combine. Output + fp32 LSE are packed into ONE all_to_all (LSE reinterpreted - as output-dtype columns along D) -> 1 NCCL call/layer instead of 2. + """A2A DCP reduce: exchange head partials + LSE across DCP ranks and merge + them. fi_a2a does both in one FlashInfer kernel (fi_a2a_lse_reduce). a2a + packs output + fp32 LSE into ONE all_to_all (LSE reinterpreted as + output-dtype columns along D) -> 1 NCCL call/layer instead of 2, then merges + with the local Triton combine. is_lse_base_on_e: True=base-e (FlashAttention), False=base-2 (FlashInfer-MLA). """ if cp_group.world_size == 1: diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 91d30c311a77..2371ebd07266 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -67,8 +67,8 @@ logger = logging.getLogger(__name__) -# fi_a2a workspace capacity in tokens when no decode graph is captured; -# arbitrary, since wider eager batches are reduced in chunks of this size. +# fi_a2a workspace capacity in tokens when no decode CUDA graph batch sizes are +# configured; arbitrary, since wider eager batches are reduced in chunks of this size. _FI_A2A_EAGER_ONLY_MAX_TOKENS = 256 diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 8db5a0d51925..76185d67ee66 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -787,7 +787,7 @@ def forward_absorb_core( self.current_attention_backend ) if dcp_comm_backend in ("a2a", "fi_a2a"): - # A2A exchange of head partials + LSE, then local Triton combine. + # Exchange head partials + LSE across DCP ranks and merge them. attn_output = dcp_a2a_lse_reduce( attn_output.contiguous(), lse.contiguous(), diff --git a/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py b/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py index 1cefc3bf9007..3cb966e3e5f9 100644 --- a/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py +++ b/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py @@ -1,6 +1,6 @@ """DCP fi_a2a: FlashInfer's fused all-to-all + LSE reduce must equal the -all-gather + Triton merge it replaces, including chunked, strided, graph-replayed -and fully-masked inputs, next to custom all-reduce v2 on the same symmetric memory. +all-gather + Triton merge it replaces, including chunked and graph-replayed +calls, next to custom all-reduce v2 on the same symmetric memory. Run with ``python test_dcp_fi_a2a_reduce.py --num-gpu 2,4,8``. Skips unless the host can run the fused reduce (Blackwell, FlashInfer with the op). @@ -17,7 +17,9 @@ from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) +from sglang.srt.environ import envs from sglang.srt.layers.dcp import comm +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils.common import fi_a2a_platform_blocker from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kernels.utils import multigpu_pytest_main @@ -26,7 +28,7 @@ register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200") # Workspace geometry: 3 tokens x 16 heads = 48 rows, so the (heads, batch) cases -# below cover a single call, an exactly-full call and multi-chunk calls. +# below cover an exactly-full call and multi-chunk calls. MAX_TOKENS, LOCAL_HEADS, HEAD_DIM = 3, 16, 512 _WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) @@ -53,7 +55,7 @@ def dcp(): # Builds custom all-reduce v2, whose symmetric memory is allocated before the # fused workspaces, as in serving. ps.initialize_model_parallel() - group = ps.get_tp_group() + group = get_parallel().tp_group capture_stream = torch.cuda.Stream() comm.init_fi_a2a_workspace( cp_group=group, @@ -68,11 +70,16 @@ def dcp(): ps.destroy_distributed_environment() -def _partials(group, *, batch, heads, dtype, seed): +def _partials(group, *, batch, heads, seed): gen = torch.Generator(device="cuda").manual_seed(seed + group.rank_in_group) total_heads = heads * group.world_size out = torch.randn( - batch, total_heads, HEAD_DIM, dtype=dtype, device="cuda", generator=gen + batch, + total_heads, + HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + generator=gen, ) lse = torch.randn( batch, total_heads, dtype=torch.float32, device="cuda", generator=gen @@ -102,59 +109,24 @@ def _fused(group, out, lse, *, is_lse_base_on_e): ) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("is_lse_base_on_e", [False, True]) -@pytest.mark.parametrize( - "heads,batch", [(16, 1), (16, 3), (16, 7), (2, 1), (2, 24), (2, 49)] -) -def test_matches_all_gather_reference(dcp, dtype, is_lse_base_on_e, heads, batch): +@pytest.mark.parametrize("heads,batch", [(16, 3), (16, 7), (2, 49)]) +def test_matches_all_gather_reference(dcp, is_lse_base_on_e, heads, batch): """Guards the head-to-peer mapping and the token chunking (B=7 at 16 heads and B=49 at 2 heads exceed the 48-row workspace).""" group, _ = dcp - out, lse = _partials(group, batch=batch, heads=heads, dtype=dtype, seed=11) + out, lse = _partials(group, batch=batch, heads=heads, seed=11) expected = _reference(group, out, lse, is_lse_base_on_e=is_lse_base_on_e) got = _fused(group, out, lse, is_lse_base_on_e=is_lse_base_on_e) torch.testing.assert_close(got, expected, rtol=1e-2, atol=1e-2) -def test_rows_empty_on_every_rank_are_zero(dcp): - """A row with no KV on any rank (LSE -inf everywhere, e.g. graph padding) - merges to 0, where the Triton combine divides 0 by 0.""" - group, _ = dcp - out, lse = _partials( - group, batch=2, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=23 - ) - empty = torch.arange(0, lse.shape[1], LOCAL_HEADS, device="cuda") - lse[:, empty] = float("-inf") # local head 0 of every destination rank - expected = _reference(group, out, lse, is_lse_base_on_e=False) - got = _fused(group, out, lse, is_lse_base_on_e=False) - assert torch.equal(got[:, 0], torch.zeros_like(got[:, 0])) - torch.testing.assert_close(got[:, 1:], expected[:, 1:], rtol=1e-2, atol=1e-2) - - -@pytest.mark.parametrize( - "layout", ["strided_head_dim", "misaligned_rows", "lse_with_unit_dim"] -) -def test_layouts_the_kernel_cannot_read_are_adapted(dcp, layout): - """The kernel reads unit-stride rows that start on 8-byte boundaries, and - FlashMLA returns its LSE as [B, H, 1]; such inputs are copied or reshaped, - not rejected.""" +def test_lse_with_a_trailing_unit_dim_is_accepted(dcp): + """FlashMLA returns its LSE as [B, H, 1].""" group, _ = dcp - out, lse = _partials( - group, batch=3, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=71 - ) + out, lse = _partials(group, batch=3, heads=LOCAL_HEADS, seed=71) expected = _reference(group, out, lse, is_lse_base_on_e=True) - if layout == "strided_head_dim": - wide = out.new_zeros(*out.shape[:-1], 2 * HEAD_DIM) - wide[..., ::2] = out - out = wide[..., ::2] - elif layout == "misaligned_rows": - flat = out.new_empty(out.numel() + 1) - flat[1:] = out.flatten() - out = flat[1:].view(out.shape) # 2 bytes past an 8-byte boundary - else: - lse = lse.unsqueeze(-1) - got = _fused(group, out, lse, is_lse_base_on_e=True) + got = _fused(group, out, lse.unsqueeze(-1), is_lse_base_on_e=True) torch.testing.assert_close(got, expected, rtol=1e-2, atol=1e-2) @@ -162,9 +134,7 @@ def test_graph_replay_and_eager_calls_interleave(dcp): """The captured graph and eager decode use different workspaces on one ordered stream; replays and eager calls must not corrupt each other.""" group, capture_stream = dcp - static_out, static_lse = _partials( - group, batch=2, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=31 - ) + static_out, static_lse = _partials(group, batch=2, heads=LOCAL_HEADS, seed=31) capture_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(capture_stream): _fused(group, static_out, static_lse, is_lse_base_on_e=False) @@ -175,14 +145,12 @@ def test_graph_replay_and_eager_calls_interleave(dcp): ) torch.cuda.current_stream().wait_stream(capture_stream) for step in range(3): - new_out, new_lse = _partials( - group, batch=2, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=40 + step - ) + new_out, new_lse = _partials(group, batch=2, heads=LOCAL_HEADS, seed=40 + step) static_out.copy_(new_out) static_lse.copy_(new_lse) graph.replay() eager_out, eager_lse = _partials( - group, batch=3, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=50 + step + group, batch=3, heads=LOCAL_HEADS, seed=50 + step ) eager_result = _fused(group, eager_out, eager_lse, is_lse_base_on_e=False) torch.testing.assert_close( @@ -203,9 +171,7 @@ def test_third_stream_raises_before_the_collective(dcp): """Only the capture and serving streams own a workspace; a third stream must fail on the host on every rank, not spin in the kernel.""" group, _ = dcp - out, lse = _partials( - group, batch=1, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=61 - ) + out, lse = _partials(group, batch=1, heads=LOCAL_HEADS, seed=61) _fused(group, out, lse, is_lse_base_on_e=False) # binds the serving stream with torch.cuda.stream(torch.cuda.Stream()): with pytest.raises(RuntimeError, match="one ordered CUDA stream"): @@ -216,17 +182,19 @@ def test_custom_all_reduce_v2_stays_exact_around_fused_calls(dcp): """The fused workspaces share torch symmetric memory with custom all-reduce v2, which allocates first; neither may disturb the other.""" group, _ = dcp + if not envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get(): + pytest.skip("SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2 is off") ca_comm = group.ca_comm - if not isinstance(ca_comm, CustomAllReduceV2) or ca_comm.disabled: - pytest.skip("custom all-reduce v2 is not enabled on this group") + # GroupCoordinator logs a failed custom all-reduce setup and carries on. + assert isinstance(ca_comm, CustomAllReduceV2) and not ca_comm.disabled, ( + f"expected custom all-reduce v2 next to the fused workspaces, got {ca_comm!r}" + ) world_size = group.world_size x = torch.full( (4096,), group.rank_in_group + 1, dtype=torch.bfloat16, device="cuda" ) want_sum = torch.full_like(x, world_size * (world_size + 1) // 2) - out, lse = _partials( - group, batch=2, heads=LOCAL_HEADS, dtype=torch.bfloat16, seed=81 - ) + out, lse = _partials(group, batch=2, heads=LOCAL_HEADS, seed=81) expected = _reference(group, out, lse, is_lse_base_on_e=False) for _ in range(2): assert torch.equal(ca_comm.custom_all_reduce(x), want_sum) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 90d1923cb887..b5677c4429d6 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -4061,6 +4061,15 @@ def test_configs_the_fused_reduce_cannot_serve_resolve_to_a2a(self): with self.assertRaisesRegex(ValueError, "--dcp-comm-backend a2a"): self._resolved(dcp_size=4, dcp_comm_backend="fi_a2a", **fields) + @override_platform(is_cuda=True, is_hip=False) + def test_pp_without_the_deepgemm_warmup_keeps_fi_a2a(self): + """Only the PP DeepGEMM warmup's third stream blocks fi_a2a; PP alone + must not.""" + host, mla = self._fused_reduce() + warmup = envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.override(False) + with host, mla, warmup: + self.assertEqual(self._resolved(dcp_size=4, pp_size=2), "fi_a2a") + @override_platform(is_cuda=True, is_hip=False) def test_explicit_fi_a2a_names_the_blocker(self): host, mla = self._fused_reduce(platform_blocker=self._NO_FUSED_OP) From ffb1623ea16abc5e50860fc921177ed062436b34 Mon Sep 17 00:00:00 2001 From: Hao Phan Date: Wed, 7 Oct 2026 12:33:52 +0700 Subject: [PATCH 10/10] [DCP] Expect custom all-reduce v2 in the fi_a2a test only where it serves the group Cross-node groups use custom all-reduce v2 only when they form one NVLink clique on a VMM-backed allocator, so the coexistence case now asks the same predicate the dispatcher uses and skips otherwise. --- .../kernels/ops/communication/test_dcp_fi_a2a_reduce.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py b/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py index 3cb966e3e5f9..dc5d3cff0d24 100644 --- a/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py +++ b/test/registered/kernels/ops/communication/test_dcp_fi_a2a_reduce.py @@ -16,6 +16,7 @@ from sglang.srt.distributed import parallel_state as ps from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, + can_use_custom_all_reduce_v2, ) from sglang.srt.environ import envs from sglang.srt.layers.dcp import comm @@ -184,6 +185,8 @@ def test_custom_all_reduce_v2_stays_exact_around_fused_calls(dcp): group, _ = dcp if not envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get(): pytest.skip("SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2 is off") + if not can_use_custom_all_reduce_v2(group=group.cpu_group, device=group.device): + pytest.skip("custom all-reduce v2 does not serve this group") ca_comm = group.ca_comm # GroupCoordinator logs a failed custom all-reduce setup and carries on. assert isinstance(ca_comm, CustomAllReduceV2) and not ca_comm.disabled, (