From accee8c89f7cbf05e9a1eff54724c90c9664f141 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:47:39 -0700 Subject: [PATCH 1/6] [CP] Context parallelism for QSA attention and the sparse indexer Adds the prefill-CP path to the Qwen sparse attention backend and its indexer, plus the shared block-layout helpers the CP collectives need. QwenSparseAttnBackend._forward_extend_cp all-gathers the local K/V into global token order, writes the complete sequence into every rank's pool (out_cache_loc is trimmed to the full extend count by the runner), and issues one sparse_gqa call per zigzag block. QSAIndexer.forward_cuda_cp projects q/k on local rows with local positions, all-gathers the raw index-K into global order for the unchanged ring store and block compression, then scores and selects top-k on local rows -- so the output rows already line up with the local q rows and need no gather after. layers/cp/collocated.py holds the row-layout helpers shared by any collocated-CP model: the rank-major block gather/scatter and the reduce-scatter that lands each rank's zigzag rows. The CPU unit test pins the invariant they rest on -- RS(cp_rank_major_blocks(x))[rank] must equal what ZigzagCPStrategy.shard_hidden_states gives that rank. Depends only on the existing prefill-CP infrastructure (is_cp_active, attn_cp_metadata, ZigzagCPStrategy); no model is wired to it yet. Co-Authored-By: Claude Opus 5 --- .../srt/layers/attention/qsa/metadata.py | 13 +- .../srt/layers/attention/qsa/qsa_indexer.py | 97 ++++++++++- .../attention/qwen_sparse_attn_backend.py | 76 ++++++++ python/sglang/srt/layers/cp/collocated.py | 164 ++++++++++++++++++ .../cp/test_qwen4_exp_collocated_cp_unit.py | 61 +++++++ 5 files changed, 407 insertions(+), 4 deletions(-) create mode 100644 python/sglang/srt/layers/cp/collocated.py create mode 100644 test/registered/unit/cp/test_qwen4_exp_collocated_cp_unit.py diff --git a/python/sglang/srt/layers/attention/qsa/metadata.py b/python/sglang/srt/layers/attention/qsa/metadata.py index a9e3270f4ef2..8c161efa9823 100644 --- a/python/sglang/srt/layers/attention/qsa/metadata.py +++ b/python/sglang/srt/layers/attention/qsa/metadata.py @@ -122,8 +122,13 @@ def get_prefill_mqa_inputs( self, layer_id: int, positions: torch.Tensor, + query_sequence_ids: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Gather packed compressed K and ragged ranges for prefill MQA.""" + """Gather packed compressed K and ragged ranges for prefill MQA. + + ``query_sequence_ids`` overrides ``token_to_batch_idx`` for callers that + score a subset of the batch's rows (prefill CP: this rank's zigzag rows); + ``positions`` must then be those rows' logical positions.""" pool = self.token_to_kv_pool ratio = self.compress_ratio @@ -151,7 +156,9 @@ def get_prefill_mqa_inputs( (0, pool.qsa_index_kv_heads, pool.qsa_index_head_dim) ) ) - num_valid_tokens = self.token_to_batch_idx.numel() + if query_sequence_ids is None: + query_sequence_ids = self.token_to_batch_idx + num_valid_tokens = query_sequence_ids.numel() if positions.numel() < num_valid_tokens: raise ValueError( "QSA prefill positions are shorter than the request mapping: " @@ -161,7 +168,7 @@ def get_prefill_mqa_inputs( row_starts, row_ends, _ = build_qsa_row_ranges( sequence_lengths, positions.to(sequence_lengths.device), - self.token_to_batch_idx.to(sequence_lengths.device), + query_sequence_ids.to(sequence_lengths.device), self.compress_ratio, ) return compressed_keys, row_starts, row_ends, sequence_lengths diff --git a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py index a308419db642..9a3e01a4c2a5 100644 --- a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py +++ b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py @@ -17,6 +17,12 @@ build_rope_position_matrix, ) from sglang.srt.layers.attention.qsa.mqa import qsa_mqa_decode, qsa_mqa_prefill +from sglang.srt.layers.cp.collocated import cp_symmetric_memory +from sglang.srt.layers.cp.utils import ( + cp_materialize_global_token_order, + cp_shard_hidden_states, + cp_shard_position_ids, +) from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb @@ -519,12 +525,86 @@ def select_decode_tokens( token_topk=self.token_topk, ) + def forward_cuda_cp( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch, + indexer_metadata, + global_rope_positions: torch.Tensor, + ) -> torch.Tensor: + """Prefill CP: score only this rank's zigzag query rows. + + ``hidden_states``/``positions`` are the local (padded) rows and their + RoPE positions. Index q/k are projected for the local rows only; the raw + keys are all-gathered into global token order so the pending-group ring + store and the block compression run on the batch's global write plan + exactly as without CP (every rank ends up with the complete compressed-K + cache). Selection then scores the local rows against all compressed + blocks of their request, position-driven, so the output rows line up + with the local query rows the sparse attention consumes.""" + meta = forward_batch.attn_cp_metadata + num_local = int(meta.total_q_prev_tokens + meta.total_q_next_tokens) + global_logical = getattr(forward_batch, "positions", None) + global_logical = ( + global_logical[0] if global_logical.ndim == 2 else global_logical + ).flatten() + num_global = indexer_metadata.get_token_to_batch_idx().numel() + global_logical = global_logical[:num_global] + local_logical = cp_shard_position_ids(global_logical, forward_batch)[:num_local] + local_sequence_ids = cp_shard_hidden_states( + indexer_metadata.get_token_to_batch_idx(), forward_batch + )[:num_local] + hidden_states = hidden_states[:num_local] + positions = ( + positions[:, :num_local] if positions.ndim == 2 else positions[:num_local] + ) + # Local q (normed + roped) and raw local k; no ring store here. + q, token_k, _ = self.project_qk(hidden_states, positions) + # Raw keys of the whole sequence in global token order: the ring store + # and the compression use the global write plan unchanged. + with cp_symmetric_memory(): + token_k_local = token_k.reshape(num_local, -1).contiguous() + token_k_full = cp_materialize_global_token_order( + token_k_local, forward_batch + ).reshape(-1, self.index_kv_heads, self.index_head_dim) + global_rope = ( + global_rope_positions[:, :num_global] + if global_rope_positions.ndim == 2 + else global_rope_positions[:num_global] + ) + self.update_key_state_and_compress( + token_k_full, + global_logical, + global_rope, + indexer_metadata, + state_slots=indexer_metadata.pending_ring_slots, + state_stored=False, + ) + compressed_keys, row_starts, row_ends, sequence_lengths = ( + indexer_metadata.get_prefill_mqa_inputs( + self.layer_id, local_logical, query_sequence_ids=local_sequence_ids + ) + ) + row_sequence_lengths = sequence_lengths.index_select( + 0, local_sequence_ids.long() + ) + return self.select_prefill_tokens( + q, + compressed_keys, + row_starts, + row_ends, + local_logical, + row_sequence_lengths, + ) + def _forward_impl( self, hidden_states: torch.Tensor, positions: torch.Tensor, forward_batch, indexer_metadata, + cp_global_rope_positions: torch.Tensor | None = None, ) -> torch.Tensor: """Portable orchestration shared by different platforms. Fast paths are gated per platforms inside kernel calls. @@ -532,6 +612,16 @@ def _forward_impl( forward_mode = forward_batch.forward_mode is_target_verify = getattr(forward_mode, "is_target_verify", lambda: False)() is_draft_extend = getattr(forward_mode, "is_draft_extend_v2", lambda: False)() + if cp_global_rope_positions is not None and not ( + forward_mode.is_decode() or is_target_verify or is_draft_extend + ): + return self.forward_cuda_cp( + hidden_states, + positions, + forward_batch, + indexer_metadata, + cp_global_rope_positions, + ) if forward_mode.is_decode() or is_target_verify or is_draft_extend: # EAGLE/MTP may advance the model's RoPE coordinate independently # from the physical paged-KV position. Compression and sparse @@ -637,9 +727,14 @@ def forward_cuda( positions: torch.Tensor, forward_batch, indexer_metadata, + cp_global_rope_positions: torch.Tensor | None = None, ) -> torch.Tensor: return self._forward_impl( - hidden_states, positions, forward_batch, indexer_metadata + hidden_states, + positions, + forward_batch, + indexer_metadata, + cp_global_rope_positions, ) def forward_xpu( diff --git a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py index b6b032dc06f1..9e4545354a08 100644 --- a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py +++ b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py @@ -37,6 +37,11 @@ sparse_gqa_fwd_interface_triton_ck, sparse_gqa_packed_decode_triton, ) +from sglang.srt.layers.cp.collocated import cp_symmetric_memory +from sglang.srt.layers.cp.utils import ( + cp_materialize_global_token_order, + is_cp_active, +) from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.utils import is_hip @@ -1273,6 +1278,12 @@ def forward_extend( ) -> torch.Tensor: if topk_indices is None: raise ValueError("QSA sparse attention requires topk_indices") + if getattr( + forward_batch, "attn_cp_metadata", None + ) is not None and is_cp_active(forward_batch): + return self._forward_extend_cp( + q, k, v, layer, forward_batch, topk_indices, save_kv_cache + ) if save_kv_cache: self.token_to_kv_pool.set_kv_buffer( layer, forward_batch.out_cache_loc, k, v @@ -1363,6 +1374,71 @@ def forward_extend( ) return self._pad_extend_output(output, num_output_rows) + def _forward_extend_cp( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer, + forward_batch, + topk_indices: torch.Tensor, + save_kv_cache: bool, + ) -> torch.Tensor: + """Prefill CP: this rank holds the query rows of its two zigzag blocks + (all heads) and the K/V of the same rows. All-gather K/V into global + token order (also written to the KV pool so decode sees every token), + then run the chunk-prefill sparse kernel with one entry per zigzag + block: q rows are contiguous per block, ``kv_lens`` is the block's + absolute end position (bottom-right causal), ``cu_k`` is the request's + base row in the packed full K/V, and ``topk_indices`` are per-request + logical positions exactly as in the non-CP path.""" + meta = forward_batch.attn_cp_metadata + prefix_lens = getattr(forward_batch, "extend_prefix_lens_cpu", None) + if prefix_lens is not None and any(int(x) for x in prefix_lens): + raise NotImplementedError( + "QSA prefill CP requires zero prefix (disable the radix cache and " + "chunked prefill)." + ) + q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) + num_output_rows = q.shape[0] + k = k.reshape(-1, layer.tp_k_head_num, layer.qk_head_dim) + v = v.reshape(-1, layer.tp_v_head_num, layer.v_head_dim) + k_width = k.shape[1] * k.shape[2] + v_width = v.shape[1] * v.shape[2] + with cp_symmetric_memory(): + kv_local = torch.cat([k.flatten(1), v.flatten(1)], dim=-1) + kv_full = cp_materialize_global_token_order(kv_local, forward_batch) + k_full, v_full = kv_full.split([k_width, v_width], dim=-1) + k_full = k_full.reshape(-1, k.shape[1], k.shape[2]).contiguous() + v_full = v_full.reshape(-1, v.shape[1], v.shape[2]).contiguous() + if save_kv_cache: + # out_cache_loc was trimmed to the full extend token count by + # prepare_cp_forward; write the complete sequence on every rank. + self.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k_full, v_full + ) + num_local_rows = int(meta.total_q_prev_tokens + meta.total_q_next_tokens) + q = q[:num_local_rows].contiguous() + indices = topk_indices[:num_local_rows].to(torch.int32).contiguous() + seq_lens = [int(x) for x in forward_batch.seq_lens_cpu[: meta.bs]] + base = [0] + for length in seq_lens[:-1]: + base.append(base[-1] + length) + # one entry per (request, prev block) then per (request, next block); + # the kernel reads cu_k[e] as the request's K base row only. + cu_k = torch.tensor(base + base, dtype=torch.int32, device=q.device) + output = sparse_gqa_fwd_interface_triton_ck( + q, + k_full, + v_full, + indices, + meta.cu_seqlens_q_combined_tensor, + cu_k, + meta.kv_len_combined_tensor, + layer.scaling, + ) + return self._pad_extend_output(output, num_output_rows) + @staticmethod def _pad_extend_output(output: torch.Tensor, num_rows: int) -> torch.Tensor: output = output.reshape(output.shape[0], -1) diff --git a/python/sglang/srt/layers/cp/collocated.py b/python/sglang/srt/layers/cp/collocated.py new file mode 100644 index 000000000000..e7e3b3865caf --- /dev/null +++ b/python/sglang/srt/layers/cp/collocated.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Collocated prefill CP helpers: the residual stream stays CP-sharded. + +Shared by hybrid linear-attention models that keep the residual stream +sharded over the CP group for the whole layer stack when the CP group is the +TP group (attn_cp_size == tp_size), so every collective (MoE all-gather / +reduce-scatter, GDN gather, attention K/V gather) runs over that one group. +The functions follow the Kimi-K3 prefill-CP implementation (branch +linear-attn-prefill-cp-ulysses) so both models share one row-order contract: + +- a rank's local rows are exactly ``ZigzagCPStrategy.shard_hidden_states``: + [block r of every request] ++ [block 2C-1-r of every request], zero-padded + to the uniform physical per-rank row count; +- ``cp_all_gather_blocks`` / ``cp_reduce_scatter_blocks`` move equal rank-major + row blocks over the CP(=TP) group with no reorder (valid for token-wise work); + ``cp_all_gather_blocks_multi`` gathers several such tensors (e.g. MoE rows + plus their top-k weights and ids) in one NCCL group launch; +- ``cp_reduce_scatter_global_rows`` first permutes FULL logical rows (global + token order) into the rank-major zigzag layout so a reduce-scatter lands + every rank its own local rows. +""" + +from __future__ import annotations + +from contextlib import nullcontext + +import torch + +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + is_symmetric_memory_enabled, + use_symmetric_memory, +) +from sglang.srt.layers.dp_attention import is_allocation_symmetric +from sglang.srt.runtime_context import get_parallel + + +def cp_symmetric_memory(): + """Allocation context for tensors that feed a collective over the CP group. + + NCCL's symmetric-memory kernels need every buffer of a collective, inputs + included, in a window registered with the group's communicator; the pool + behind ``use_symmetric_memory`` registers its segments with a communicator + when a context for that group exits, so producers of collective inputs + allocate under this context and the collective helpers below allocate + their outputs under it right before the call. A no-op unless symmetric + memory is enabled.""" + if not is_symmetric_memory_enabled() or not is_allocation_symmetric(): + # Keep the disabled path free of group lookups: CPU unit tests build + # the CP metadata without a communicator. + return nullcontext() + return use_symmetric_memory(get_parallel().attn_cp_group) + + +def cp_all_gather_blocks(x: torch.Tensor) -> torch.Tensor: + """Rank-major all-gather of equal-length local row blocks over the CP group.""" + group = get_parallel().attn_cp_group + with cp_symmetric_memory(): + x = x.contiguous() + out = torch.empty( + (x.shape[0] * group.world_size, *x.shape[1:]), + dtype=x.dtype, + device=x.device, + ) + group.all_gather_into_tensor(out, x) + return out + + +def cp_all_gather_blocks_multi(*tensors: torch.Tensor) -> list[torch.Tensor]: + """Rank-major all-gather of several equal-row tensors over the CP group. + + On the torch.distributed path the collectives are coalesced into one NCCL + group (one kernel launch for all of them); on the pynccl / symmetric-memory + path they are issued one after another.""" + group = get_parallel().attn_cp_group + with cp_symmetric_memory(): + inputs = [x.contiguous() for x in tensors] + outputs = [ + torch.empty( + (x.shape[0] * group.world_size, *x.shape[1:]), + dtype=x.dtype, + device=x.device, + ) + for x in inputs + ] + pynccl = group.pynccl_comm + torch_dist_path = pynccl is None or ( + pynccl.disabled and not group.is_symmetric_memory_enabled() + ) + coalesce = getattr(torch.distributed, "_coalescing_manager", None) + if torch_dist_path and coalesce is not None and len(inputs) > 1: + with coalesce(group=group.device_group, device=inputs[0].device): + for out, x in zip(outputs, inputs): + torch.distributed.all_gather_into_tensor( + out, x, group=group.device_group + ) + else: + for out, x in zip(outputs, inputs): + group.all_gather_into_tensor(out, x) + return outputs + + +def cp_rank_major_blocks(x: torch.Tensor, metadata, cp_size: int) -> torch.Tensor: + """Permute full logical rows into the concatenation of every CP rank's local + zigzag block (rank r: block r of each request, then block 2C-1-r of each + request), each block zero-padded to the physical per-rank length. + + One flat chunk list and one ``torch.cat`` (a single copy of ``x``), with a + shared zero pad slab; no per-rank intermediate concatenations.""" + segments = 2 * cp_size + bs = len(metadata.split_list) // segments + chunks = torch.split(x[: metadata.total_seq_lens], metadata.split_list, dim=0) + phys = max(metadata.per_rank_actual_token) + logical = metadata.per_rank_logical_token or metadata.per_rank_actual_token + max_pad = max(phys - n for n in logical) + pad_slab = x.new_zeros((max_pad, *x.shape[1:])) if max_pad > 0 else None + parts = [] + for rank in range(cp_size): + parts.extend(chunks[i] for i in range(rank, bs * segments, segments)) + parts.extend( + chunks[i] for i in range(segments - 1 - rank, bs * segments, segments) + ) + pad = phys - logical[rank] + assert pad >= 0, (phys, logical[rank]) + if pad: + parts.append(pad_slab[:pad]) + # The permuted rows are the reduce-scatter input: allocate them in the pool. + with cp_symmetric_memory(): + out = x.new_empty((cp_size * phys, *x.shape[1:])) + return torch.cat(parts, dim=0, out=out) + + +def cp_reduce_scatter_blocks(x: torch.Tensor) -> torch.Tensor: + """Sum rank-major row blocks over the CP group; return this rank's block.""" + group = get_parallel().attn_cp_group + with cp_symmetric_memory(): + x = x.contiguous() + out = torch.empty( + (x.shape[0] // group.world_size, *x.shape[1:]), + dtype=x.dtype, + device=x.device, + ) + group.reduce_scatter_tensor(out, x) + return out + + +def cp_reduce_scatter_global_rows(partial: torch.Tensor, forward_batch) -> torch.Tensor: + """Sum TP-partial full-row outputs (global token order) over the CP(=TP) + group and keep this rank's zigzag rows: one reduce-scatter instead of + all-reduce + re-shard.""" + parallel = get_parallel() + blocks = cp_rank_major_blocks( + partial, forward_batch.attn_cp_metadata, parallel.attn_cp_size + ) + return cp_reduce_scatter_blocks(blocks) + + +__all__ = [ + "cp_all_gather_blocks", + "cp_symmetric_memory", + "cp_all_gather_blocks_multi", + "cp_rank_major_blocks", + "cp_reduce_scatter_blocks", + "cp_reduce_scatter_global_rows", +] diff --git a/test/registered/unit/cp/test_qwen4_exp_collocated_cp_unit.py b/test/registered/unit/cp/test_qwen4_exp_collocated_cp_unit.py new file mode 100644 index 000000000000..983d934ec400 --- /dev/null +++ b/test/registered/unit/cp/test_qwen4_exp_collocated_cp_unit.py @@ -0,0 +1,61 @@ +"""CPU checks for the shared (Qwen4-Exp, Kimi-K3) collocated prefill CP row layout. + +The rank-major zigzag block layout that `cp_rank_major_blocks` produces for a +reduce-scatter must equal, per rank, what `ZigzagCPStrategy.shard_hidden_states` +gives that rank (plus its zero padding), so that RS(blocks)[rank] lands exactly +the rows the rank's residual stream holds.""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.layers.cp.collocated import cp_rank_major_blocks +from sglang.srt.layers.cp.padding import pad_local_rows, pad_logical_token_to_physical +from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy +from sglang.srt.runtime_context import get_parallel +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestCollocatedCPRankMajorBlocks(CustomTestCase): + def _check(self, cp_size, extend_lens): + num_tokens = sum(extend_lens) + hidden = torch.arange(num_tokens, dtype=torch.float32).unsqueeze( + 1 + ) * torch.ones(1, 3) + with get_parallel().override(attn_cp_rank=0, attn_cp_size=cp_size): + strategy = ZigzagCPStrategy(cp_size=cp_size) + metadata = strategy.build_metadata(num_tokens, extend_lens) + pad_logical_token_to_physical(metadata) # as prepare_cp_forward does + blocks = cp_rank_major_blocks(hidden, metadata, cp_size) + phys = metadata.per_rank_actual_token[0] + self.assertEqual(tuple(blocks.shape), (phys * cp_size, 3)) + for rank in range(cp_size): + with get_parallel().override(attn_cp_rank=rank, attn_cp_size=cp_size): + strategy_r = ZigzagCPStrategy(cp_size=cp_size) + meta_r = strategy_r.build_metadata(num_tokens, extend_lens) + pad_logical_token_to_physical(meta_r) + forward_batch = SimpleNamespace(attn_cp_metadata=meta_r) + local = strategy_r.shard_hidden_states(hidden, forward_batch) + local = pad_local_rows(local, meta_r, dim=0) + with self.subTest(rank=rank): + torch.testing.assert_close( + blocks[rank * phys : (rank + 1) * phys], local, rtol=0, atol=0 + ) + + def test_single_balanced_request(self): + self._check(8, [32768 // 64]) + + def test_ragged_requests_with_padding(self): + self._check(8, [100, 33]) + self._check(4, [257, 16, 4096 // 16]) + + def test_cp2(self): + self._check(2, [7, 9, 24]) + + +if __name__ == "__main__": + unittest.main() From 6d9aeea459477002d54b1c9b3721c281d6b001b9 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Mon, 28 Sep 2026 22:45:16 -0700 Subject: [PATCH 2/6] Refactor QSA context parallelism to reuse zigzag collectives --- .../srt/layers/attention/qsa/qsa_indexer.py | 13 +- .../attention/qwen_sparse_attn_backend.py | 13 +- python/sglang/srt/layers/cp/collocated.py | 164 ----------------- python/sglang/srt/layers/cp/zigzag.py | 25 ++- test/registered/unit/cp/test_qsa_cp_unit.py | 71 +++++++ .../cp/test_qwen4_exp_collocated_cp_unit.py | 61 ------ .../unit/cp/test_zigzag_collectives_unit.py | 173 ++++++++++++++++++ 7 files changed, 281 insertions(+), 239 deletions(-) delete mode 100644 python/sglang/srt/layers/cp/collocated.py create mode 100644 test/registered/unit/cp/test_qsa_cp_unit.py delete mode 100644 test/registered/unit/cp/test_qwen4_exp_collocated_cp_unit.py create mode 100644 test/registered/unit/cp/test_zigzag_collectives_unit.py diff --git a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py index 9a3e01a4c2a5..464fbec36c0d 100644 --- a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py +++ b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py @@ -17,11 +17,12 @@ build_rope_position_matrix, ) from sglang.srt.layers.attention.qsa.mqa import qsa_mqa_decode, qsa_mqa_prefill -from sglang.srt.layers.cp.collocated import cp_symmetric_memory from sglang.srt.layers.cp.utils import ( + ContextParallelStrategyKind, cp_materialize_global_token_order, cp_shard_hidden_states, cp_shard_position_ids, + get_cp_strategy, ) from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.linear import ReplicatedLinear @@ -543,6 +544,13 @@ def forward_cuda_cp( cache). Selection then scores the local rows against all compressed blocks of their request, position-driven, so the output rows line up with the local query rows the sparse attention consumes.""" + strategy = get_cp_strategy() + if strategy is None or strategy.kind != ContextParallelStrategyKind.ZIGZAG: + strategy_name = strategy.name if strategy is not None else "none" + raise NotImplementedError( + "QSA prefill CP only supports the zigzag strategy; " + f"got {strategy_name}." + ) meta = forward_batch.attn_cp_metadata num_local = int(meta.total_q_prev_tokens + meta.total_q_next_tokens) global_logical = getattr(forward_batch, "positions", None) @@ -563,8 +571,7 @@ def forward_cuda_cp( q, token_k, _ = self.project_qk(hidden_states, positions) # Raw keys of the whole sequence in global token order: the ring store # and the compression use the global write plan unchanged. - with cp_symmetric_memory(): - token_k_local = token_k.reshape(num_local, -1).contiguous() + token_k_local = token_k.reshape(num_local, -1).contiguous() token_k_full = cp_materialize_global_token_order( token_k_local, forward_batch ).reshape(-1, self.index_kv_heads, self.index_head_dim) diff --git a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py index 9e4545354a08..e7f3602d5162 100644 --- a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py +++ b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py @@ -37,9 +37,10 @@ sparse_gqa_fwd_interface_triton_ck, sparse_gqa_packed_decode_triton, ) -from sglang.srt.layers.cp.collocated import cp_symmetric_memory from sglang.srt.layers.cp.utils import ( + ContextParallelStrategyKind, cp_materialize_global_token_order, + get_cp_strategy, is_cp_active, ) from sglang.srt.model_executor.forward_batch_info import ForwardMode @@ -1392,6 +1393,13 @@ def _forward_extend_cp( absolute end position (bottom-right causal), ``cu_k`` is the request's base row in the packed full K/V, and ``topk_indices`` are per-request logical positions exactly as in the non-CP path.""" + strategy = get_cp_strategy() + if strategy is None or strategy.kind != ContextParallelStrategyKind.ZIGZAG: + strategy_name = strategy.name if strategy is not None else "none" + raise NotImplementedError( + "QSA prefill CP only supports the zigzag strategy; " + f"got {strategy_name}." + ) meta = forward_batch.attn_cp_metadata prefix_lens = getattr(forward_batch, "extend_prefix_lens_cpu", None) if prefix_lens is not None and any(int(x) for x in prefix_lens): @@ -1405,8 +1413,7 @@ def _forward_extend_cp( v = v.reshape(-1, layer.tp_v_head_num, layer.v_head_dim) k_width = k.shape[1] * k.shape[2] v_width = v.shape[1] * v.shape[2] - with cp_symmetric_memory(): - kv_local = torch.cat([k.flatten(1), v.flatten(1)], dim=-1) + kv_local = torch.cat([k.flatten(1), v.flatten(1)], dim=-1) kv_full = cp_materialize_global_token_order(kv_local, forward_batch) k_full, v_full = kv_full.split([k_width, v_width], dim=-1) k_full = k_full.reshape(-1, k.shape[1], k.shape[2]).contiguous() diff --git a/python/sglang/srt/layers/cp/collocated.py b/python/sglang/srt/layers/cp/collocated.py deleted file mode 100644 index e7e3b3865caf..000000000000 --- a/python/sglang/srt/layers/cp/collocated.py +++ /dev/null @@ -1,164 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Collocated prefill CP helpers: the residual stream stays CP-sharded. - -Shared by hybrid linear-attention models that keep the residual stream -sharded over the CP group for the whole layer stack when the CP group is the -TP group (attn_cp_size == tp_size), so every collective (MoE all-gather / -reduce-scatter, GDN gather, attention K/V gather) runs over that one group. -The functions follow the Kimi-K3 prefill-CP implementation (branch -linear-attn-prefill-cp-ulysses) so both models share one row-order contract: - -- a rank's local rows are exactly ``ZigzagCPStrategy.shard_hidden_states``: - [block r of every request] ++ [block 2C-1-r of every request], zero-padded - to the uniform physical per-rank row count; -- ``cp_all_gather_blocks`` / ``cp_reduce_scatter_blocks`` move equal rank-major - row blocks over the CP(=TP) group with no reorder (valid for token-wise work); - ``cp_all_gather_blocks_multi`` gathers several such tensors (e.g. MoE rows - plus their top-k weights and ids) in one NCCL group launch; -- ``cp_reduce_scatter_global_rows`` first permutes FULL logical rows (global - token order) into the rank-major zigzag layout so a reduce-scatter lands - every rank its own local rows. -""" - -from __future__ import annotations - -from contextlib import nullcontext - -import torch - -from sglang.srt.distributed.device_communicators.pynccl_allocator import ( - is_symmetric_memory_enabled, - use_symmetric_memory, -) -from sglang.srt.layers.dp_attention import is_allocation_symmetric -from sglang.srt.runtime_context import get_parallel - - -def cp_symmetric_memory(): - """Allocation context for tensors that feed a collective over the CP group. - - NCCL's symmetric-memory kernels need every buffer of a collective, inputs - included, in a window registered with the group's communicator; the pool - behind ``use_symmetric_memory`` registers its segments with a communicator - when a context for that group exits, so producers of collective inputs - allocate under this context and the collective helpers below allocate - their outputs under it right before the call. A no-op unless symmetric - memory is enabled.""" - if not is_symmetric_memory_enabled() or not is_allocation_symmetric(): - # Keep the disabled path free of group lookups: CPU unit tests build - # the CP metadata without a communicator. - return nullcontext() - return use_symmetric_memory(get_parallel().attn_cp_group) - - -def cp_all_gather_blocks(x: torch.Tensor) -> torch.Tensor: - """Rank-major all-gather of equal-length local row blocks over the CP group.""" - group = get_parallel().attn_cp_group - with cp_symmetric_memory(): - x = x.contiguous() - out = torch.empty( - (x.shape[0] * group.world_size, *x.shape[1:]), - dtype=x.dtype, - device=x.device, - ) - group.all_gather_into_tensor(out, x) - return out - - -def cp_all_gather_blocks_multi(*tensors: torch.Tensor) -> list[torch.Tensor]: - """Rank-major all-gather of several equal-row tensors over the CP group. - - On the torch.distributed path the collectives are coalesced into one NCCL - group (one kernel launch for all of them); on the pynccl / symmetric-memory - path they are issued one after another.""" - group = get_parallel().attn_cp_group - with cp_symmetric_memory(): - inputs = [x.contiguous() for x in tensors] - outputs = [ - torch.empty( - (x.shape[0] * group.world_size, *x.shape[1:]), - dtype=x.dtype, - device=x.device, - ) - for x in inputs - ] - pynccl = group.pynccl_comm - torch_dist_path = pynccl is None or ( - pynccl.disabled and not group.is_symmetric_memory_enabled() - ) - coalesce = getattr(torch.distributed, "_coalescing_manager", None) - if torch_dist_path and coalesce is not None and len(inputs) > 1: - with coalesce(group=group.device_group, device=inputs[0].device): - for out, x in zip(outputs, inputs): - torch.distributed.all_gather_into_tensor( - out, x, group=group.device_group - ) - else: - for out, x in zip(outputs, inputs): - group.all_gather_into_tensor(out, x) - return outputs - - -def cp_rank_major_blocks(x: torch.Tensor, metadata, cp_size: int) -> torch.Tensor: - """Permute full logical rows into the concatenation of every CP rank's local - zigzag block (rank r: block r of each request, then block 2C-1-r of each - request), each block zero-padded to the physical per-rank length. - - One flat chunk list and one ``torch.cat`` (a single copy of ``x``), with a - shared zero pad slab; no per-rank intermediate concatenations.""" - segments = 2 * cp_size - bs = len(metadata.split_list) // segments - chunks = torch.split(x[: metadata.total_seq_lens], metadata.split_list, dim=0) - phys = max(metadata.per_rank_actual_token) - logical = metadata.per_rank_logical_token or metadata.per_rank_actual_token - max_pad = max(phys - n for n in logical) - pad_slab = x.new_zeros((max_pad, *x.shape[1:])) if max_pad > 0 else None - parts = [] - for rank in range(cp_size): - parts.extend(chunks[i] for i in range(rank, bs * segments, segments)) - parts.extend( - chunks[i] for i in range(segments - 1 - rank, bs * segments, segments) - ) - pad = phys - logical[rank] - assert pad >= 0, (phys, logical[rank]) - if pad: - parts.append(pad_slab[:pad]) - # The permuted rows are the reduce-scatter input: allocate them in the pool. - with cp_symmetric_memory(): - out = x.new_empty((cp_size * phys, *x.shape[1:])) - return torch.cat(parts, dim=0, out=out) - - -def cp_reduce_scatter_blocks(x: torch.Tensor) -> torch.Tensor: - """Sum rank-major row blocks over the CP group; return this rank's block.""" - group = get_parallel().attn_cp_group - with cp_symmetric_memory(): - x = x.contiguous() - out = torch.empty( - (x.shape[0] // group.world_size, *x.shape[1:]), - dtype=x.dtype, - device=x.device, - ) - group.reduce_scatter_tensor(out, x) - return out - - -def cp_reduce_scatter_global_rows(partial: torch.Tensor, forward_batch) -> torch.Tensor: - """Sum TP-partial full-row outputs (global token order) over the CP(=TP) - group and keep this rank's zigzag rows: one reduce-scatter instead of - all-reduce + re-shard.""" - parallel = get_parallel() - blocks = cp_rank_major_blocks( - partial, forward_batch.attn_cp_metadata, parallel.attn_cp_size - ) - return cp_reduce_scatter_blocks(blocks) - - -__all__ = [ - "cp_all_gather_blocks", - "cp_symmetric_memory", - "cp_all_gather_blocks_multi", - "cp_rank_major_blocks", - "cp_reduce_scatter_blocks", - "cp_reduce_scatter_global_rows", -] diff --git a/python/sglang/srt/layers/cp/zigzag.py b/python/sglang/srt/layers/cp/zigzag.py index 877b55e23bff..748635bb0525 100644 --- a/python/sglang/srt/layers/cp/zigzag.py +++ b/python/sglang/srt/layers/cp/zigzag.py @@ -39,6 +39,7 @@ import torch.nn.functional as F from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + is_symmetric_memory_enabled, use_symmetric_memory, ) from sglang.srt.layers.cp.base import ( @@ -468,17 +469,25 @@ def _all_gather_reorganized(self, x: torch.Tensor, forward_batch): assert x.shape[0] >= local_len x = x[:local_len] pad_size = max_len - x.shape[0] - if pad_size > 0: - padding = [0, 0] * (x.ndim - 1) + [0, pad_size] - x = F.pad(x, padding, mode="constant", value=0) - group = get_parallel().attn_cp_group - ctx = ( - use_symmetric_memory(group, disabled=not is_allocation_symmetric()) - if x.is_cuda - else nullcontext() + symmetric = ( + x.is_cuda and is_symmetric_memory_enabled() and is_allocation_symmetric() ) + ctx = use_symmetric_memory(group) if symmetric else nullcontext() with ctx: + if symmetric: + # Allocate the same send/receive shapes on every rank, including + # ranks without padding. Both buffers must belong to the pool; + # contiguous() can return a producer's existing allocation. + local = x.new_empty((max_len, *x.shape[1:])) + local[:local_len].copy_(x) + local[local_len:].zero_() + x = local + elif pad_size > 0: + padding = [0, 0] * (x.ndim - 1) + [0, pad_size] + x = F.pad(x, padding, mode="constant", value=0) + else: + x = x.contiguous() gathered = torch.empty( max_len * self.cp_size, *x.shape[1:], diff --git a/test/registered/unit/cp/test_qsa_cp_unit.py b/test/registered/unit/cp/test_qsa_cp_unit.py new file mode 100644 index 000000000000..c257fbe1237c --- /dev/null +++ b/test/registered/unit/cp/test_qsa_cp_unit.py @@ -0,0 +1,71 @@ +"""CPU checks for QSA's zigzag-only context parallel entrypoints.""" + +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from sglang.srt.layers.attention import qwen_sparse_attn_backend +from sglang.srt.layers.attention.qsa import qsa_indexer +from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer +from sglang.srt.layers.attention.qwen_sparse_attn_backend import QwenSparseAttnBackend +from sglang.srt.layers.cp.base import ContextParallelStrategyKind +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +@pytest.mark.parametrize( + "strategy", + [ + None, + SimpleNamespace(name="interleave", kind=ContextParallelStrategyKind.INTERLEAVE), + ], +) +@pytest.mark.parametrize("entrypoint", ["indexer", "attention"]) +def test_qsa_cp_rejects_unsupported_strategy_before_metadata_or_collectives( + monkeypatch, strategy, entrypoint +): + module = qsa_indexer if entrypoint == "indexer" else qwen_sparse_attn_backend + monkeypatch.setattr(module, "get_cp_strategy", lambda: strategy) + gather = Mock(side_effect=AssertionError("unsupported CP must not communicate")) + monkeypatch.setattr(module, "cp_materialize_global_token_order", gather) + # None metadata and tensors ensure rejection precedes zigzag-specific access. + with pytest.raises( + NotImplementedError, match="QSA prefill CP only supports the zigzag strategy" + ): + if entrypoint == "indexer": + QSAIndexer.forward_cuda_cp(None, None, None, None, None, None) + else: + QwenSparseAttnBackend._forward_extend_cp( + None, None, None, None, None, None, None, False + ) + gather.assert_not_called() + + +def test_cuda_indexer_forwards_global_rope_to_cp_path(): + hidden, local_rope, metadata, global_rope = (object() for _ in range(4)) + batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND) + expected = object() + indexer = SimpleNamespace(forward_cuda_cp=Mock(return_value=expected)) + indexer._forward_impl = lambda *args: QSAIndexer._forward_impl(indexer, *args) + + actual = QSAIndexer.forward_cuda( + indexer, + hidden, + local_rope, + batch, + metadata, + cp_global_rope_positions=global_rope, + ) + + assert actual is expected + indexer.forward_cuda_cp.assert_called_once_with( + hidden, local_rope, batch, metadata, global_rope + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/cp/test_qwen4_exp_collocated_cp_unit.py b/test/registered/unit/cp/test_qwen4_exp_collocated_cp_unit.py deleted file mode 100644 index 983d934ec400..000000000000 --- a/test/registered/unit/cp/test_qwen4_exp_collocated_cp_unit.py +++ /dev/null @@ -1,61 +0,0 @@ -"""CPU checks for the shared (Qwen4-Exp, Kimi-K3) collocated prefill CP row layout. - -The rank-major zigzag block layout that `cp_rank_major_blocks` produces for a -reduce-scatter must equal, per rank, what `ZigzagCPStrategy.shard_hidden_states` -gives that rank (plus its zero padding), so that RS(blocks)[rank] lands exactly -the rows the rank's residual stream holds.""" - -import unittest -from types import SimpleNamespace - -import torch - -from sglang.srt.layers.cp.collocated import cp_rank_major_blocks -from sglang.srt.layers.cp.padding import pad_local_rows, pad_logical_token_to_physical -from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy -from sglang.srt.runtime_context import get_parallel -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=5, suite="base-a-test-cpu") - - -class TestCollocatedCPRankMajorBlocks(CustomTestCase): - def _check(self, cp_size, extend_lens): - num_tokens = sum(extend_lens) - hidden = torch.arange(num_tokens, dtype=torch.float32).unsqueeze( - 1 - ) * torch.ones(1, 3) - with get_parallel().override(attn_cp_rank=0, attn_cp_size=cp_size): - strategy = ZigzagCPStrategy(cp_size=cp_size) - metadata = strategy.build_metadata(num_tokens, extend_lens) - pad_logical_token_to_physical(metadata) # as prepare_cp_forward does - blocks = cp_rank_major_blocks(hidden, metadata, cp_size) - phys = metadata.per_rank_actual_token[0] - self.assertEqual(tuple(blocks.shape), (phys * cp_size, 3)) - for rank in range(cp_size): - with get_parallel().override(attn_cp_rank=rank, attn_cp_size=cp_size): - strategy_r = ZigzagCPStrategy(cp_size=cp_size) - meta_r = strategy_r.build_metadata(num_tokens, extend_lens) - pad_logical_token_to_physical(meta_r) - forward_batch = SimpleNamespace(attn_cp_metadata=meta_r) - local = strategy_r.shard_hidden_states(hidden, forward_batch) - local = pad_local_rows(local, meta_r, dim=0) - with self.subTest(rank=rank): - torch.testing.assert_close( - blocks[rank * phys : (rank + 1) * phys], local, rtol=0, atol=0 - ) - - def test_single_balanced_request(self): - self._check(8, [32768 // 64]) - - def test_ragged_requests_with_padding(self): - self._check(8, [100, 33]) - self._check(4, [257, 16, 4096 // 16]) - - def test_cp2(self): - self._check(2, [7, 9, 24]) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/cp/test_zigzag_collectives_unit.py b/test/registered/unit/cp/test_zigzag_collectives_unit.py new file mode 100644 index 000000000000..73ac39747aa8 --- /dev/null +++ b/test/registered/unit/cp/test_zigzag_collectives_unit.py @@ -0,0 +1,173 @@ +"""CPU layout and allocation checks for the shared zigzag collectives.""" + +import sys +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import PropertyMock, patch + +import pytest +import torch + +from sglang.srt.layers.cp import base, padding, zigzag +from sglang.srt.layers.cp.padding import pad_logical_token_to_physical +from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +@contextmanager +def _parallel(cp_size, rank, group=None): + parallel = SimpleNamespace( + attn_cp_rank=rank, attn_cp_size=cp_size, attn_cp_group=group + ) + with ( + patch.object(base, "get_parallel", return_value=parallel), + patch.object(zigzag, "get_parallel", return_value=parallel), + patch.object(padding, "get_parallel", return_value=parallel), + ): + yield + + +def _make_shards(payload, cp_size, lengths, physical_padding): + strategy = ZigzagCPStrategy(cp_size) + shards, batches = [], [] + with ( + patch.object(zigzag, "get_device", return_value=SimpleNamespace(device="cpu")), + patch.object(base, "_STRATEGY", strategy), + ): + for rank in range(cp_size): + with _parallel(cp_size, rank): + metadata = strategy.build_metadata(sum(lengths), lengths) + if physical_padding: + pad_logical_token_to_physical(metadata) + batch = SimpleNamespace(attn_cp_metadata=metadata) + shards.append(strategy.shard_hidden_states(payload, batch)) + batches.append(batch) + return strategy, shards, batches + + +def _noncontiguous(tensor): + backing = tensor.new_empty((*tensor.shape[:-1], tensor.shape[-1] * 2)) + view = backing[..., ::2] + view.copy_(tensor) + assert not view.is_contiguous() + return view + + +def _gathered_shards(shards, metadata): + lengths = metadata.per_rank_logical_token or metadata.per_rank_actual_token + max_len = max(lengths) + padded = shards[0].new_zeros((len(shards), max_len, *shards[0].shape[1:])) + for rank, length in enumerate(lengths): + padded[rank, :length].copy_(shards[rank][:length]) + return padded.flatten(0, 1) + + +@pytest.mark.parametrize("method", ["gather_hidden_states", "gather_kv_cache"]) +@pytest.mark.parametrize("physical_padding", [False, True]) +@pytest.mark.parametrize( + "cp_size,lengths,shape", + [ + (8, [512], (3,)), + (8, [100, 33], (2, 3)), + (4, [257, 16, 256], (3,)), + (2, [7, 9, 24], (2, 3)), + ], +) +def test_gather_restores_global_token_order( + method, physical_padding, cp_size, lengths, shape +): + num_tokens = sum(lengths) + payload = torch.arange(num_tokens * torch.Size(shape).numel()).reshape( + num_tokens, *shape + ) + strategy, shards, batches = _make_shards( + payload, cp_size, lengths, physical_padding + ) + expected_gathered = _gathered_shards(shards, batches[0].attn_cp_metadata) + + for rank, (shard, batch) in enumerate(zip(shards, batches)): + # Nonzero physical padding must not leak into reconstructed token order. + logical = batch.attn_cp_metadata.per_rank_logical_token + if logical is not None: + shard[logical[rank] :].fill_(-1) + local = _noncontiguous(shard) + + def gather(output, input): + assert input.is_contiguous() + torch.testing.assert_close(input, expected_gathered.chunk(cp_size)[rank]) + output.copy_(expected_gathered) + + group = SimpleNamespace(all_gather_into_tensor=gather) + with _parallel(cp_size, rank, group): + actual = getattr(strategy, method)(local, batch) + torch.testing.assert_close(actual, payload, rtol=0, atol=0) + + +@pytest.mark.parametrize("physical_padding", [False, True]) +@pytest.mark.parametrize("strided", [False, True]) +def test_symmetric_gather_allocates_equal_send_and_receive_buffers_in_pool( + physical_padding, strided +): + cp_size, lengths = 4, [17, 19] + payload = torch.arange(sum(lengths) * 6).reshape(-1, 2, 3) + strategy, shards, batches = _make_shards( + payload, cp_size, lengths, physical_padding + ) + gathered = _gathered_shards(shards, batches[0].attn_cp_metadata) + send_shape = tuple(gathered.chunk(cp_size)[0].shape) + original_new_empty, original_empty = torch.Tensor.new_empty, torch.empty + in_pool = False + allocations = {} + + @contextmanager + def pool(group): + nonlocal in_pool + in_pool = True + try: + yield + finally: + in_pool = False + + def record_new_empty(tensor, *args, **kwargs): + result = original_new_empty(tensor, *args, **kwargs) + allocations[result.data_ptr()] = (in_pool, tuple(result.shape)) + return result + + def record_empty(*args, **kwargs): + result = original_empty(*args, **kwargs) + allocations[result.data_ptr()] = (in_pool, tuple(result.shape)) + return result + + for rank, (shard, batch) in enumerate(zip(shards, batches)): + local = _noncontiguous(shard) if strided else shard + allocations.clear() + + def gather(output, input): + assert input.is_contiguous() + assert not in_pool + assert input.data_ptr() != local.data_ptr() + assert allocations[input.data_ptr()] == (True, send_shape) + assert allocations[output.data_ptr()] == (True, tuple(gathered.shape)) + assert torch.equal(input, gathered.chunk(cp_size)[rank]) + output.copy_(gathered) + + group = SimpleNamespace(all_gather_into_tensor=gather) + with ( + _parallel(cp_size, rank, group), + patch.object( + torch.Tensor, "is_cuda", new_callable=PropertyMock, return_value=True + ), + patch.object(zigzag, "is_symmetric_memory_enabled", return_value=True), + patch.object(zigzag, "is_allocation_symmetric", return_value=True), + patch.object(zigzag, "use_symmetric_memory", side_effect=pool), + patch.object(torch.Tensor, "new_empty", record_new_empty), + patch.object(torch, "empty", record_empty), + ): + actual = strategy.gather_kv_cache(local, batch) + torch.testing.assert_close(actual, payload, rtol=0, atol=0) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) From 24c5d200f68ccef13c8d6e33af91f80c1a01609e Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Mon, 28 Sep 2026 23:24:11 -0700 Subject: [PATCH 3/6] Reuse layer boundaries for Qwen4 attention and FFN communication --- .../srt/layers/layer_boundary/construction.py | 4 + .../srt/layers/layer_boundary/factories.py | 6 +- .../srt/layers/layer_boundary/prepare.py | 4 +- .../layer_boundary/residual/__init__.py | 5 + .../layers/layer_boundary/residual/gated.py | 148 ++++++ python/sglang/srt/models/qwen4_exp.py | 119 ++--- .../layer_boundary/test_gated_residual.py | 142 ++++++ .../test_moe_output_reduction.py | 7 +- .../test_qwen4_gated_boundary.py | 457 ++++++++++++++++++ 9 files changed, 814 insertions(+), 78 deletions(-) create mode 100644 python/sglang/srt/layers/layer_boundary/residual/gated.py create mode 100644 test/registered/unit/layer_boundary/test_gated_residual.py create mode 100644 test/registered/unit/layer_boundary/test_qwen4_gated_boundary.py diff --git a/python/sglang/srt/layers/layer_boundary/construction.py b/python/sglang/srt/layers/layer_boundary/construction.py index 61907661abab..9b11e21bb2de 100644 --- a/python/sglang/srt/layers/layer_boundary/construction.py +++ b/python/sglang/srt/layers/layer_boundary/construction.py @@ -292,6 +292,10 @@ def _bind_stage(declaration, norm, incoming, outgoing, **options): declaration.update.at_producer and TokenAxis.ATTN_CP in edge.produced.layout.sharded - edge.need.layout.sharded + and not ( + getattr(declaration.update, "supports_moe_cp_gather", False) + and not _gathers_over_attention_cp() + ) ): raise NotImplementedError("MHC with a gather over attention CP") handoff = None diff --git a/python/sglang/srt/layers/layer_boundary/factories.py b/python/sglang/srt/layers/layer_boundary/factories.py index 5efd7ce9dff0..ac1e21b65096 100644 --- a/python/sglang/srt/layers/layer_boundary/factories.py +++ b/python/sglang/srt/layers/layer_boundary/factories.py @@ -357,7 +357,11 @@ def declare_ffn( def _resolve(stage, variant, following=None): axes, attention, local, full = _rows(variant) if stage.update.at_producer: - if stage.sparse and sparse_moe_gathers_over_moe_cp(): + if ( + stage.sparse + and sparse_moe_gathers_over_moe_cp() + and not getattr(stage.update, "supports_moe_cp_gather", False) + ): raise NotImplementedError( "MHC does not support a MoE gathered over the MoE-CP group" ) diff --git a/python/sglang/srt/layers/layer_boundary/prepare.py b/python/sglang/srt/layers/layer_boundary/prepare.py index e923bcd6a872..23b9b9038073 100644 --- a/python/sglang/srt/layers/layer_boundary/prepare.py +++ b/python/sglang/srt/layers/layer_boundary/prepare.py @@ -443,9 +443,11 @@ def _mlp_input_gather_moe_cp( # so no rank hits this path while others proceed to the allgather. # - During decode: moe_cp allgather is skipped (guarded by is_context_parallel_extend). # - CUDA graph warmup: not applicable when --cuda-graph-backend-prefill=disabled is used. - if hidden_states.shape[0] == 0: + if hidden_states.shape[0] == 0 and update.adds_plainly: return hidden_states, residual + # Stateful nonlinear reads must still initialize their FFN residual state + # on empty batches, even though there are no rows to communicate. hidden_states, residual = gather( hidden_states, residual, forward_batch, layernorm, update=update, cache=cache ) diff --git a/python/sglang/srt/layers/layer_boundary/residual/__init__.py b/python/sglang/srt/layers/layer_boundary/residual/__init__.py index 19c12d9882d0..4c632e59cd96 100644 --- a/python/sglang/srt/layers/layer_boundary/residual/__init__.py +++ b/python/sglang/srt/layers/layer_boundary/residual/__init__.py @@ -34,6 +34,11 @@ class StageUpdate(Protocol): The shard conversion methods must move any update-associated state together with the residual. Nonlinear updates must not claim adds_plainly. + + Producer-written updates may optionally set supports_moe_cp_gather=True + when their auxiliary state stays on local residual rows while the FFN + input is gathered. The boundary returns the FFN output to those rows + before applying the update. Other producer-written updates remain gated. """ # A plain add, which one rank may run before the sum it adds into diff --git a/python/sglang/srt/layers/layer_boundary/residual/gated.py b/python/sglang/srt/layers/layer_boundary/residual/gated.py new file mode 100644 index 000000000000..a2df02b6575d --- /dev/null +++ b/python/sglang/srt/layers/layer_boundary/residual/gated.py @@ -0,0 +1,148 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Gated hyper-connections expressed as stage reads and residual updates.""" + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.layers.layer_boundary.adapters.attention import ( + _redistribute_to_attn_tp_shards, +) +from sglang.srt.layers.layer_boundary.ops import gather_attention_tp +from sglang.srt.layers.layer_boundary.residual import LayerResidual + +if TYPE_CHECKING: + from sglang.srt.layers.hyperconnection import GatedResidual + + +class _BranchState: + def __init__(self, connection: "GatedResidual"): + self.connection = connection + self.normalized: Optional[torch.Tensor] = None + + def mix(self, residual): + hidden_states, (residual, self.normalized) = self.connection.mix(residual) + return hidden_states, residual + + def combine(self, hidden_states, residual): + if self.normalized is None: + raise RuntimeError("a gated residual update requires its stage's read") + output = self.connection.combine(hidden_states, (residual, self.normalized)) + self.normalized = None + return output + + def to_attn_tp_shard(self, residual): + if self.normalized is None: + raise RuntimeError("a gated residual shard requires its stage's read") + self.normalized = _redistribute_to_attn_tp_shards(self.normalized) + return _redistribute_to_attn_tp_shards(residual) + + def from_attn_tp_shards(self, residual): + if self.normalized is None: + raise RuntimeError("a gated residual gather requires its stage's read") + # The tuple path allocates separate outputs. Gathering both through the + # shared DP scratch buffer would overwrite the first gathered tensor. + residual, self.normalized = gather_attention_tp( + (residual, self.normalized), None + ) + return residual + + +class _GatedRead: + norms_plainly = False + before_gather = True + + def __init__(self, state: _BranchState, *, enters_stack: bool): + self.state = state + self.enters_stack = enters_stack + + def enter(self, hidden_states): + if not self.enters_stack: + raise RuntimeError("a gated FFN does not enter the layer stack") + connection = self.state.connection + width = hidden_states.shape[-1] + if width == connection.hc_count * connection.hidden_size: + return hidden_states + if width != connection.hidden_size: + raise ValueError(f"unexpected gated residual input width: {width}") + return torch.cat([hidden_states] * connection.hc_count, dim=-1) + + def read(self, residual, norm, quant_format="", post_residual_addition=None): + if norm is not None or quant_format: + raise NotImplementedError("gated residual mixing owns its normalization") + if post_residual_addition is not None: + residual = residual + post_residual_addition + return self.state.mix(residual) + + def update_and_read( + self, + update, + hidden_states, + residual, + norm, + quant_format="", + post_residual_addition=None, + ): + return self.read( + update.update(hidden_states, residual), + norm, + quant_format, + post_residual_addition, + ) + + +class _GatedUpdate: + adds_plainly = False + can_defer_across_layers = False + # MoE CP gathers only the compute input. The boundary returns the output + # to the local residual rows before applying this update. + supports_moe_cp_gather = True + + def __init__(self, state: _BranchState, *, at_producer: bool): + self.state = state + self.at_producer = at_producer + + def update(self, hidden_states, residual): + return self.state.combine(hidden_states, residual) + + def residual_to_attn_tp_shard(self, residual): + return self.state.to_attn_tp_shard(residual) + + def residual_from_attn_tp_shards(self, residual): + return self.state.from_attn_tp_shards(residual) + + +class GatedResidualState: + """Adapt attention/FFN GatedResidual modules to the layer boundaries. + + The batch's residual stream retains the unnormalized hyper-input tensor. + Each branch retains the normalized input needed by its next combine, and + moves that auxiliary tensor whenever its residual changes token rows. + FFN updates run at the producer so no layer carries another layer's state. + The owning model continues to register and load the connection modules. + """ + + def __init__(self, attention: "GatedResidual", ffn: "GatedResidual"): + attention_state = _BranchState(attention) + ffn_state = _BranchState(ffn) + self._residual = LayerResidual( + attention_read=_GatedRead(attention_state, enters_stack=True), + attention_update=_GatedUpdate(attention_state, at_producer=False), + ffn_read=_GatedRead(ffn_state, enters_stack=False), + ffn_update=_GatedUpdate(ffn_state, at_producer=True), + ) + + def layer_residual(self) -> LayerResidual: + return self._residual diff --git a/python/sglang/srt/models/qwen4_exp.py b/python/sglang/srt/models/qwen4_exp.py index 7bcc07ca5aed..759a9168cc16 100644 --- a/python/sglang/srt/models/qwen4_exp.py +++ b/python/sglang/srt/models/qwen4_exp.py @@ -23,13 +23,9 @@ from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.layers.aux_hidden_states import AuxHiddenStateList from sglang.srt.layers.dp_attention import ( - attn_tp_all_gather, attn_tp_all_reduce, dp_gather_replicate, dp_scatter, - get_dp_global_num_tokens, - get_global_dp_buffer, - get_local_dp_buffer, is_allocation_symmetric, is_dp_attention_enabled, ) @@ -37,10 +33,18 @@ GatedResidual, HyperConnectionConfig, ) -from sglang.srt.layers.layer_boundary import get_attn_tp_context +from sglang.srt.layers.layer_boundary import ( + HandoffRows, + declare_attn, + declare_ffn, + get_attn_tp_context, + make_stages, +) +from sglang.srt.layers.layer_boundary.residual import batch as residual_batch +from sglang.srt.layers.layer_boundary.residual.gated import GatedResidualState from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.logits_processor import LogitsProcessorOutput -from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv +from sglang.srt.layers.moe import get_moe_a2a_backend from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.modelopt_quant import ( @@ -72,7 +76,7 @@ make_ple_file_prefetcher, make_ple_file_rss_trimmer, ) -from sglang.srt.runtime_context import get_forward, get_parallel +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import get_bool_env_var, is_hip, logger _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() @@ -1320,8 +1324,6 @@ def _init_qwen4_exp_layer_extensions( for attr_name in ( "input_layernorm", "post_attention_layernorm", - "attn_boundary", - "ffn_boundary", ): if hasattr(self, attr_name): delattr(self, attr_name) @@ -1360,6 +1362,29 @@ def _init_qwen4_exp_layer_extensions( use_mix=True, use_combine=True, ) + residual = GatedResidualState( + self.attn_hyper_connection, self.mlp_hyper_connection + ).layer_residual() + sparse = bool(config.num_experts) + ffn = declare_ffn( + sparse=sparse, + read=residual.ffn_read, + update=residual.ffn_update, + dense_tp_size=None if sparse else get_parallel().tp_size, + handoff_rows=HandoffRows.ATTENTION, + ) + self.attn_boundary, self.ffn_boundary = make_stages( + ( + declare_attn( + read=residual.attention_read, + update=residual.attention_update, + ), + None, + ), + (ffn, None), + previous=ffn if layer_id != 0 else None, + terminal=layer_id == config.num_hidden_layers - 1, + ) def _prepare_qwen4_exp_attn( self, @@ -1391,8 +1416,10 @@ def _prepare_qwen4_exp_attn( ple_query, forward_batch, ple_batch ) - hidden_states, residual = self.attn_hyper_connection.mix(hidden_states) - return hidden_states, residual + # Qwen4 hands complete hyper-connection streams between layers. PLE + # updates that tensor before the attention boundary reads its input. + hidden_states = residual_batch.written(hidden_states, forward_batch) + return self.attn_boundary.prepare(hidden_states, forward_batch), None def _prepare_qwen4_exp_mlp( self, @@ -1400,70 +1427,21 @@ def _prepare_qwen4_exp_mlp( residual: Optional[torch.Tensor], forward_batch: ForwardBatch, ): - if not forward_batch.forward_mode.is_idle(): - hidden_states = attn_tp_all_reduce(hidden_states) - hidden_states = self.attn_hyper_connection.combine(hidden_states, residual) - hidden_states, residual = self.mlp_hyper_connection.mix(hidden_states) - return hidden_states, residual - - def _qwen4_exp_use_dp_moe_gather(self) -> bool: - return get_parallel().attn_dp_size > 1 and get_moe_a2a_backend().is_none() - - def _qwen4_exp_use_attn_tp_a2a_scatter(self) -> bool: - return get_parallel().attn_tp_size > 1 and not get_moe_a2a_backend().is_none() + hidden_states = self.attn_boundary.finish(hidden_states, forward_batch) + return self.ffn_boundary.prepare(hidden_states, forward_batch), None def _run_qwen4_exp_mlp( self, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - if not self.config.num_experts: - return self.mlp(hidden_states) - - use_dp_moe_gather = self._qwen4_exp_use_dp_moe_gather() - use_attn_tp_a2a_scatter = self._qwen4_exp_use_attn_tp_a2a_scatter() - - if use_dp_moe_gather: - hidden_states, local_hidden_states = ( - get_global_dp_buffer(get_parallel().tp_group), - hidden_states, - ) - dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) - elif hidden_states.shape[0] == 0 and get_moe_a2a_backend().is_none(): - # Only safe to short-circuit an empty batch when the MoE holds no collective; - # under deepep an idle DP rank must still join dispatch/combine or peers hang. - return hidden_states - - attn_tp_chunks = None - if use_attn_tp_a2a_scatter: - attn_tp_size = get_parallel().attn_tp_size - attn_tp_chunks = list(hidden_states.tensor_split(attn_tp_size)) - hidden_states = attn_tp_chunks[get_parallel().attn_tp_rank].contiguous() - - use_reduce_scatterv = use_dp_moe_gather and should_use_dp_reduce_scatterv() - with get_forward().scoped(mlp_reduce_scatter=use_reduce_scatterv): - hidden_states = self.mlp(hidden_states, forward_batch) - - if use_dp_moe_gather: - hidden_states, global_hidden_states = ( - get_local_dp_buffer(get_parallel().tp_group), - hidden_states, - ) - if use_reduce_scatterv: - get_parallel().tp_group.reduce_scatterv( - global_hidden_states, - output=hidden_states, - sizes=get_dp_global_num_tokens(), - ) - else: - dp_scatter(hidden_states, global_hidden_states, forward_batch) - elif use_attn_tp_a2a_scatter: - assert attn_tp_chunks is not None - gathered = [torch.empty_like(t) for t in attn_tp_chunks] - attn_tp_all_gather(gathered, hidden_states.contiguous()) - hidden_states = torch.cat(gathered) - - return hidden_states + with self.ffn_boundary.exit(forward_batch) as ffn_exit: + if not self.config.num_experts: + hidden_states = self.mlp(hidden_states) + elif hidden_states.shape[0] or not get_moe_a2a_backend().is_none(): + # An idle all-to-all rank must still participate in MoE dispatch. + hidden_states = self.mlp(hidden_states, forward_batch) + return ffn_exit.finish(hidden_states) def _postprocess_qwen4_exp_layer( self, @@ -1471,7 +1449,6 @@ def _postprocess_qwen4_exp_layer( residual: Optional[torch.Tensor], forward_batch: ForwardBatch, ): - hidden_states = self.mlp_hyper_connection.combine(hidden_states, residual) return hidden_states, None @@ -1775,6 +1752,8 @@ def forward( _commit_ple_batch(ple_batch, forward_batch) + hidden_states = residual_batch.take_output(hidden_states, forward_batch) + if not self.pp_group.is_last_rank: proxy_tensors = {"hidden_states": hidden_states} if forward_batch.mm_input_embeds is not None: diff --git a/test/registered/unit/layer_boundary/test_gated_residual.py b/test/registered/unit/layer_boundary/test_gated_residual.py new file mode 100644 index 000000000000..9f5952831be9 --- /dev/null +++ b/test/registered/unit/layer_boundary/test_gated_residual.py @@ -0,0 +1,142 @@ +"""Gated hyper-connection state follows boundary residual rows and ordering.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.layers.layer_boundary.adapters import attention +from sglang.srt.layers.layer_boundary.residual import gated +from sglang.srt.layers.layer_boundary.residual.gated import GatedResidualState +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class ToyGatedResidual: + """A nonlinear mix/combine that consumes both raw and normalized input.""" + + hc_count = 2 + hidden_size = 3 + + def __init__(self, scale): + self.scale = scale + + def mix(self, hyper_input): + normalized = hyper_input / (1 + hyper_input.square().mean(-1, keepdim=True)) + mixed = normalized.reshape(-1, self.hc_count, self.hidden_size).sum(1) + return mixed * self.scale, (hyper_input, normalized) + + def combine(self, hidden_states, residuals): + residual, normalized = residuals + injected = hidden_states.repeat(1, self.hc_count) + return residual + self.scale * injected * torch.sigmoid(normalized) + + +class TestGatedResidual(CustomTestCase): + def setUp(self): + self.attention = ToyGatedResidual(0.5) + self.ffn = ToyGatedResidual(0.75) + self.state = GatedResidualState(self.attention, self.ffn) + self.residual = self.state.layer_residual() + self.hidden = torch.arange(24, dtype=torch.float32).reshape(4, 6) / 10 + + def test_preserves_attention_then_ffn_mix_combine_order(self): + expected_attn, expected_residual = self.attention.mix(self.hidden) + expected_after_attn = self.attention.combine( + expected_attn.square(), expected_residual + ) + expected_ffn, expected_residual = self.ffn.mix(expected_after_attn) + expected = self.ffn.combine(expected_ffn.sin(), expected_residual) + + attn, residual = self.residual.attention_read.read(self.hidden, None) + ffn, residual = self.residual.ffn_read.update_and_read( + self.residual.attention_update, attn.square(), residual, None + ) + actual = self.residual.ffn_update.update(ffn.sin(), residual) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + self.assertIsNone(self.residual.attention_update.state.normalized) + self.assertIsNone(self.residual.ffn_update.state.normalized) + self.assertFalse(self.residual.attention_update.adds_plainly) + self.assertFalse(self.residual.attention_update.at_producer) + self.assertTrue(self.residual.ffn_update.at_producer) + self.assertFalse(self.residual.ffn_update.can_defer_across_layers) + + def test_first_attention_expands_embeddings_only_once(self): + embedding = self.hidden[:, : self.attention.hidden_size] + expanded = self.residual.attention_read.enter(embedding) + torch.testing.assert_close(expanded, embedding.repeat(1, 2)) + self.assertIs(self.residual.attention_read.enter(expanded), expanded) + with self.assertRaisesRegex(ValueError, "width"): + self.residual.attention_read.enter(torch.empty(4, 5)) + + def test_shard_moves_normalized_state_with_raw_residual(self): + mixed, residual = self.residual.attention_read.read(self.hidden, None) + expected = self.attention.combine(mixed, self.attention.mix(self.hidden)[1]) + update = self.residual.attention_update + with patch.object( + attention, + "get_parallel", + lambda: SimpleNamespace(attn_tp_size=2, attn_tp_rank=1), + ): + shard = update.residual_to_attn_tp_shard(residual) + actual = update.update(mixed[2:], shard) + torch.testing.assert_close(actual, expected[2:], rtol=0, atol=0) + + def test_gather_moves_both_states_without_shared_scratch_aliasing(self): + full_mixed, full_residuals = self.attention.mix(self.hidden) + _, local_residual = self.residual.attention_read.read(self.hidden[2:], None) + update = self.residual.attention_update + + def gather_pair(pair, forward_batch): + self.assertIsNone(forward_batch) + torch.testing.assert_close(pair[0], self.hidden[2:]) + torch.testing.assert_close(pair[1], full_residuals[1][2:]) + return tuple(t.clone() for t in full_residuals) + + with patch.object(gated, "gather_attention_tp", side_effect=gather_pair) as ag: + gathered = update.residual_from_attn_tp_shards(local_residual) + ag.assert_called_once() + actual = update.update(full_mixed, gathered) + expected = self.attention.combine(full_mixed, full_residuals) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_empty_stage_preserves_widened_residual_shape(self): + hidden = torch.empty(0, self.attention.hidden_size) + expanded = self.residual.attention_read.enter(hidden) + attn, residual = self.residual.attention_read.read(expanded, None) + ffn, residual = self.residual.ffn_read.update_and_read( + self.residual.attention_update, attn, residual, None + ) + output = self.residual.ffn_update.update(ffn, residual) + self.assertEqual(output.shape, (0, 6)) + self.assertIsNone(self.residual.ffn_update.state.normalized) + + def test_update_uses_actual_producer_and_addition_precedes_mix(self): + updated = self.hidden + 2 + producer = SimpleNamespace(update=Mock(return_value=updated)) + addition = torch.full_like(self.hidden, 0.5) + actual, residual = self.residual.ffn_read.update_and_read( + producer, + torch.empty(4, 3), + self.hidden, + None, + post_residual_addition=addition, + ) + producer.update.assert_called_once() + expected, _ = self.ffn.mix(updated + addition) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(residual, updated + addition) + + def test_update_cannot_reuse_auxiliary_state_after_consumption(self): + mixed, residual = self.residual.attention_read.read(self.hidden, None) + self.residual.attention_update.update(mixed, residual) + with self.assertRaisesRegex(RuntimeError, "requires its stage's read"): + self.residual.attention_update.update(mixed, residual) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layer_boundary/test_moe_output_reduction.py b/test/registered/unit/layer_boundary/test_moe_output_reduction.py index b2e9c7807026..f3d68505cf7b 100644 --- a/test/registered/unit/layer_boundary/test_moe_output_reduction.py +++ b/test/registered/unit/layer_boundary/test_moe_output_reduction.py @@ -169,10 +169,7 @@ def test_direct_dp_exits_publish_their_selected_sum(self): # These models own their DP exit instead of using StageBoundary. import __future__ - for filename, method in ( - ("deepseek_v4.py", "_run_moe_ffn_dp_sync"), - ("qwen4_exp.py", "_run_qwen4_exp_mlp"), - ): + for filename, method in (("deepseek_v4.py", "_run_moe_ffn_dp_sync"),): path = MODELS_DIR / filename node = next( n @@ -229,8 +226,6 @@ def scatter(output, value, batch): model = types.SimpleNamespace( dsa_enable_prefill_cp=False, config=types.SimpleNamespace(num_experts=4), - _qwen4_exp_use_dp_moe_gather=lambda: True, - _qwen4_exp_use_attn_tp_a2a_scatter=lambda: False, mlp=lambda value, batch, **kwargs: reduce_moe_output(value), ) batch = types.SimpleNamespace( diff --git a/test/registered/unit/layer_boundary/test_qwen4_gated_boundary.py b/test/registered/unit/layer_boundary/test_qwen4_gated_boundary.py new file mode 100644 index 000000000000..1c75317b9668 --- /dev/null +++ b/test/registered/unit/layer_boundary/test_qwen4_gated_boundary.py @@ -0,0 +1,457 @@ +"""Qwen4's actual layer orchestration through production stage boundaries. + +Only the hyper-connection kernels and distributed collectives are replaced: +the nonlinear reference catches mixing before reduction or combining before +the FFN output returns to the residual's rows. +""" + +import __future__ + +import ast +import itertools +import unittest +from contextlib import ExitStack, contextmanager, nullcontext +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +import sglang +from sglang.srt.layers import layernorm_sp +from sglang.srt.layers.dp_attention import DpPaddingMode +from sglang.srt.layers.layer_boundary import ( + HandoffRows, + declare_attn, + declare_ffn, + make_stages, +) +from sglang.srt.layers.layer_boundary.residual import batch as residual_batch +from sglang.srt.layers.layer_boundary.residual.gated import GatedResidualState +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.runtime_context import get_forward +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.communicator_patch import patch_communicator +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +HIDDEN = 3 + + +def mix_reference(value): + first, second = value.chunk(2, dim=-1) + return torch.tanh(first + second / 2) + + +def combine_reference(value, residual): + return residual + value.repeat(1, 2) * torch.sigmoid(residual) + + +class Connection: + hc_count = 2 + hidden_size = HIDDEN + + def __init__(self): + self.mixed = [] + self.combined = [] + + def mix(self, value): + self.mixed.append(value.clone()) + return mix_reference(value), (value, torch.sigmoid(value)) + + def combine(self, value, residuals): + residual, normalized = residuals + self.combined.append(value.clone()) + torch.testing.assert_close(normalized, torch.sigmoid(residual)) + return residual + value.repeat(1, 2) * normalized + + +def load_mixin(namespace): + # Avoid importing GPU-only attention kernels just to execute the model's + # orchestration. The methods compiled here are the production source. + path = Path(sglang.__file__).parent / "srt/models/qwen4_exp.py" + node = next( + node + for node in ast.parse(path.read_text()).body + if isinstance(node, ast.ClassDef) and node.name == "Qwen4ExpLayerExtensionMixin" + ) + exec( + compile( + ast.Module(body=[node], type_ignores=[]), + str(path), + "exec", + flags=__future__.annotations.compiler_flag, + ), + namespace, + ) + return namespace[node.name] + + +class Harness: + def __init__( + self, + *, + dp=1, + tp=1, + a2a=False, + rsv=False, + rows=2, + rank=0, + cp_rows=None, + cp_rank=0, + ): + self.dp, self.tp, self.a2a, self.rsv = dp, tp, a2a, rsv + self.rows, self.rank = rows, rank + self.cp_size = len(cp_rows) if cp_rows is not None else 1 + self.cp_rank = cp_rank + self.events = [] + self.compute_inputs = [] + self.compute_flags = [] + self.expected_gathered = None + self.expected_cp_inputs = None + self.parallel = SimpleNamespace( + tp_size=dp * tp * self.cp_size, + tp_rank=rank, + attn_dp_size=dp, + attn_dp_rank=0, + attn_tp_size=tp, + attn_tp_rank=rank, + attn_cp_size=self.cp_size, + attn_cp_rank=cp_rank, + moe_dp_size=1, + moe_ep_size=1, + moe_tp_size=dp * tp * self.cp_size, + moe_dense_tp_size=None, + dwdp_size=1, + enable_dp_attention=dp > 1, + enable_prefill_cp=self.cp_size > 1, + enable_attn_tp_input_scattered=False, + tp_group=SimpleNamespace(reduce_scatterv=self.reduce_scatterv), + attn_tp_group=SimpleNamespace(), + ) + self.batch = SimpleNamespace( + forward_mode=ForwardMode.DECODE if rows else ForwardMode.IDLE, + residual_stream=None, + dp_padding_mode=DpPaddingMode.SUM_LEN, + attn_cp_metadata=( + SimpleNamespace(per_rank_actual_token=cp_rows) + if cp_rows is not None + else None + ), + global_num_tokens_cpu=[rows] * dp, + input_ids=torch.zeros(rows, dtype=torch.int64), + ) + if cp_rows is not None: + self.batch.forward_mode = ForwardMode.EXTEND + + def attention_sum(self, value): + self.events.append("attention_sum") + return value * self.tp + + def attention_scatter(self, output, value): + self.events.append("attention_reduce_scatter") + output.copy_((value * self.tp).tensor_split(self.tp)[self.rank]) + + def attention_gather(self, output, value): + self.events.append("attention_gather") + torch.testing.assert_close( + value, self.expected_gathered.tensor_split(self.tp)[self.rank] + ) + output.copy_(self.expected_gathered) + + def gather_dp(self, output, value, batch, cp_shard_counts=None): + self.events.append("dp_gather") + output.copy_(torch.cat([value + i * 10 for i in range(self.dp)])) + + def gather_cp(self, output, value): + self.events.append("cp_gather") + torch.testing.assert_close( + value, self.expected_cp_inputs.tensor_split(self.cp_size)[self.cp_rank] + ) + output.copy_(self.expected_cp_inputs) + + def scatter_dp(self, output, value, batch): + self.events.append("dp_scatter") + output.copy_(value[: self.rows]) + + def reduce_scatterv(self, value, *, output, sizes): + self.events.append("dp_reduce_scatterv") + output.copy_(value[: self.rows] * self.parallel.tp_size) + + def mlp(self, value, *args): + self.compute_inputs.append(value.clone()) + skipped = get_forward().mlp_reduce_scatter + self.compute_flags.append(skipped) + return value * (3 / self.parallel.tp_size if skipped else 3) + + @contextmanager + def running(self): + backend = SimpleNamespace(is_none=lambda: not self.a2a) + with ExitStack() as stack: + overrides = { + "get_parallel": lambda: self.parallel, + "get_moe_a2a_backend": lambda: backend, + "is_moe_input_scattered_across_dp_ranks": lambda: self.a2a, + "is_enable_moe_cp_allgather": lambda: self.cp_size > 1, + "get_moe_cp_size": lambda: self.cp_size, + "get_moe_cp_rank": lambda: self.cp_rank, + "moe_cp_all_gather_into_tensor": self.gather_cp, + "is_dsa_enable_prefill_cp": lambda: False, + "is_mla_cp_enabled": lambda: False, + "get_spec": lambda: SimpleNamespace(speculative_algorithm=None), + "get_lora": lambda: SimpleNamespace(enable_lora=False), + "get_exec": lambda: SimpleNamespace( + comm=SimpleNamespace( + boundary_reduction="rs+rsv", enable_quant_communications=False + ), + overlap=SimpleNamespace(enable_two_batch_overlap=False), + ), + "get_attn_tp_context": lambda: SimpleNamespace(input_scattered=False), + "post_experts_reduction_group": lambda: self.parallel.tp_group, + "should_use_dp_reduce_scatterv": lambda: self.rsv, + "can_use_dp_reduce_scatter": lambda: False, + "use_symmetric_memory": lambda *a, **k: nullcontext(), + "is_allocation_symmetric": lambda: False, + "attention_tensor_model_parallel_all_reduce": self.attention_sum, + "attn_tp_reduce_scatter_tensor": self.attention_scatter, + "attn_tp_all_gather_into_tensor": self.attention_gather, + "get_global_dp_buffer": lambda group: torch.empty( + self.rows * self.dp, HIDDEN + ), + "get_local_dp_buffer": lambda group, hidden_size=None: torch.empty( + self.rows, hidden_size or HIDDEN + ), + "get_dp_global_num_tokens": lambda: [self.rows] * self.dp, + "dp_gather_replicate": self.gather_dp, + "dp_scatter": self.scatter_dp, + } + for name, replacement in overrides.items(): + stack.enter_context(patch_communicator(name, replacement)) + stack.enter_context( + patch.object(layernorm_sp, "layernorm_sp_enabled", lambda: False) + ) + stack.enter_context( + get_forward().scoped( + mlp_reduce_scatter=False, + fuse_mlp_allreduce=False, + sp_active=False, + attn_input_scattered=False, + ) + ) + self.namespace = dict( + torch=torch, + residual_batch=residual_batch, + GatedResidualState=GatedResidualState, + GatedResidual=lambda *a, **k: Connection(), + HyperConnectionConfig=SimpleNamespace, + HandoffRows=HandoffRows, + declare_attn=declare_attn, + declare_ffn=declare_ffn, + make_stages=make_stages, + get_parallel=lambda: self.parallel, + get_moe_a2a_backend=lambda: backend, + _get_ple_forward_mode=lambda batch: batch.forward_mode, + ) + yield + + def model(self, *, sparse, layer_id=0, num_layers=1): + model = load_mixin(self.namespace)() + model.config = SimpleNamespace( + num_experts=4 if sparse else 0, + num_hidden_layers=num_layers, + hc_count=2, + hidden_size=HIDDEN, + ple_layer_ids=[], + hc_lowrank=2, + rms_norm_eps=1e-5, + ) + model._init_qwen4_exp_layer_extensions(model.config, layer_id) + model.mlp = self.mlp + return model + + def forward(self, model, value, *, ple_batch=None): + hidden, residual = model._prepare_qwen4_exp_attn( + value, None, self.batch, ple_batch=ple_batch + ) + if not self.batch.forward_mode.is_idle(): + hidden = hidden * (2 / self.tp) + hidden, residual = model._prepare_qwen4_exp_mlp(hidden, residual, self.batch) + hidden = model._run_qwen4_exp_mlp(hidden, self.batch) + return model._postprocess_qwen4_exp_layer(hidden, residual, self.batch) + + +def reference(value): + if value.shape[-1] == HIDDEN: + value = value.repeat(1, 2) + after_attention = combine_reference(mix_reference(value) * 2, value) + mlp_input = mix_reference(after_attention) + return combine_reference(mlp_input * 3, after_attention), mlp_input + + +class TestQwen4Boundaries(CustomTestCase): + def test_dense_and_moe_complete_attention_sum_before_nonlinear_mix(self): + value = torch.arange(6, dtype=torch.float32).reshape(2, HIDDEN) / 10 + expected, mlp_input = reference(value) + for sparse in (False, True): + with self.subTest(sparse=sparse): + harness = Harness(tp=2) + with harness.running(): + model = harness.model(sparse=sparse) + output, residual = harness.forward(model, value) + torch.testing.assert_close(output, expected) + torch.testing.assert_close(harness.compute_inputs[0], mlp_input) + self.assertIsNone(residual) + self.assertEqual(harness.events, ["attention_sum"]) + self.assertIsNone(harness.batch.residual_stream.pending) + + def test_dp_ffn_gathers_mixed_input_and_combines_after_return(self): + value = torch.arange(6, dtype=torch.float32).reshape(2, HIDDEN) / 10 + expected, local_input = reference(value) + for sparse, rsv in itertools.product((False, True), repeat=2): + with self.subTest(sparse=sparse, reduce_scatterv=rsv): + harness = Harness(dp=2, tp=2, rsv=rsv) + with harness.running(): + output, residual = harness.forward( + harness.model(sparse=sparse), value + ) + self.assertFalse(get_forward().mlp_reduce_scatter) + torch.testing.assert_close(output, expected) + torch.testing.assert_close( + harness.compute_inputs[0], + torch.cat([local_input, local_input + 10]), + ) + self.assertEqual(harness.compute_flags, [rsv]) + self.assertEqual( + harness.events, + [ + "attention_sum", + "dp_gather", + "dp_reduce_scatterv" if rsv else "dp_scatter", + ], + ) + self.assertIsNone(residual) + + def test_a2a_moves_gated_state_with_attention_tp_slice(self): + value = torch.arange(12, dtype=torch.float32).reshape(4, HIDDEN) / 10 + expected, mlp_input = reference(value) + for rank in (0, 1): + with self.subTest(rank=rank): + harness = Harness(tp=2, a2a=True, rows=4, rank=rank) + harness.expected_gathered = expected + with harness.running(): + output, residual = harness.forward( + harness.model(sparse=True), value + ) + torch.testing.assert_close(output, expected) + torch.testing.assert_close( + harness.compute_inputs[0], mlp_input.tensor_split(2)[rank] + ) + self.assertEqual( + harness.events, ["attention_reduce_scatter", "attention_gather"] + ) + self.assertIsNone(residual) + + def test_idle_a2a_still_participates_in_expert_dispatch(self): + value = torch.empty(0, HIDDEN) + for a2a in (False, True): + with self.subTest(a2a=a2a): + harness = Harness(a2a=a2a, rows=0) + with harness.running(): + output, residual = harness.forward( + harness.model(sparse=True), value + ) + self.assertEqual(len(harness.compute_inputs), int(a2a)) + self.assertEqual(output.shape, (0, 2 * HIDDEN)) + self.assertIsNone(residual) + + def test_moe_cp_takes_back_uneven_rows_before_gated_combine(self): + values = [ + torch.arange(6, dtype=torch.float32).reshape(2, HIDDEN) / 10, + torch.arange(3, dtype=torch.float32).reshape(1, HIDDEN) / 4, + ] + references = [reference(value) for value in values] + gathered_inputs = torch.cat( + [references[0][1], references[1][1], torch.zeros(1, HIDDEN)] + ) + for rank, value in enumerate(values): + with self.subTest(cp_rank=rank): + harness = Harness(rows=len(value), cp_rows=[2, 1], cp_rank=rank) + harness.expected_cp_inputs = gathered_inputs + with harness.running(): + output, residual = harness.forward( + harness.model(sparse=True), value + ) + torch.testing.assert_close(output, references[rank][0]) + torch.testing.assert_close(harness.compute_inputs[0], gathered_inputs) + self.assertEqual(harness.events, ["cp_gather"]) + self.assertIsNone(residual) + + def test_idle_with_moe_cp_still_consumes_gated_residual_reads(self): + harness = Harness(rows=0, cp_rows=[0, 0]) + harness.batch.forward_mode = ForwardMode.IDLE + harness.batch.attn_cp_metadata = None + with harness.running(): + output, residual = harness.forward( + harness.model(sparse=True), torch.empty(0, HIDDEN) + ) + self.assertEqual(output.shape, (0, 2 * HIDDEN)) + self.assertIsNone(residual) + self.assertEqual(harness.compute_inputs, []) + self.assertEqual(harness.events, []) + + def test_intermediate_layers_complete_stream_and_reuse_auxiliary_state(self): + value = torch.arange(6, dtype=torch.float32).reshape(2, HIDDEN) / 10 + harness = Harness(tp=2) + with harness.running(): + layers = [ + harness.model(sparse=True, layer_id=i, num_layers=2) for i in range(2) + ] + for repetition in range(2): + hidden = value + repetition + expected = hidden + for model in layers: + hidden, residual = harness.forward(model, hidden) + expected = reference(expected)[0] + torch.testing.assert_close(hidden, expected) + self.assertIsNone(residual) + self.assertIsNone(harness.batch.residual_stream.pending) + self.assertIs(harness.batch.residual_stream.residual, hidden) + + def test_ple_reads_expanded_input_before_attention_mix(self): + value = torch.ones(2, HIDDEN) / 4 + calls = [] + + def ple(hidden, batch, ple_batch): + calls.append(hidden.clone()) + return hidden / 2 + + harness = Harness() + with harness.running(): + model = harness.model(sparse=False) + model.ple = ple + output, _ = harness.forward(model, value, ple_batch=object()) + torch.testing.assert_close(calls[0], value.repeat(1, 2)) + torch.testing.assert_close(output, reference(value * 1.5)[0]) + + def test_ple_idle_enters_collective_path_and_missing_active_batch_fails(self): + from unittest.mock import Mock + + for rows in (0, 2): + with self.subTest(rows=rows): + harness = Harness(rows=rows) + with harness.running(): + model = harness.model(sparse=False) + model.ple = Mock() + value = torch.empty(rows, HIDDEN) + if rows: + with self.assertRaisesRegex(RuntimeError, "missing its batch"): + harness.forward(model, value) + else: + output, _ = harness.forward(model, value) + model.ple.forward_idle.assert_called_once_with(harness.batch) + self.assertEqual(output.shape, (0, 2 * HIDDEN)) + + +if __name__ == "__main__": + unittest.main() From 4478d77240772fd89b9f32bff6aa7fc0210360a9 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Mon, 28 Sep 2026 23:40:19 -0700 Subject: [PATCH 4/6] Keep QSA CP refactor scoped to the backend and zigzag strategy --- .../srt/layers/layer_boundary/construction.py | 4 - .../srt/layers/layer_boundary/factories.py | 6 +- .../srt/layers/layer_boundary/prepare.py | 4 +- .../layer_boundary/residual/__init__.py | 5 - .../layers/layer_boundary/residual/gated.py | 148 ------ python/sglang/srt/models/qwen4_exp.py | 119 +++-- .../layer_boundary/test_gated_residual.py | 142 ------ .../test_moe_output_reduction.py | 7 +- .../test_qwen4_gated_boundary.py | 457 ------------------ 9 files changed, 78 insertions(+), 814 deletions(-) delete mode 100644 python/sglang/srt/layers/layer_boundary/residual/gated.py delete mode 100644 test/registered/unit/layer_boundary/test_gated_residual.py delete mode 100644 test/registered/unit/layer_boundary/test_qwen4_gated_boundary.py diff --git a/python/sglang/srt/layers/layer_boundary/construction.py b/python/sglang/srt/layers/layer_boundary/construction.py index 9b11e21bb2de..61907661abab 100644 --- a/python/sglang/srt/layers/layer_boundary/construction.py +++ b/python/sglang/srt/layers/layer_boundary/construction.py @@ -292,10 +292,6 @@ def _bind_stage(declaration, norm, incoming, outgoing, **options): declaration.update.at_producer and TokenAxis.ATTN_CP in edge.produced.layout.sharded - edge.need.layout.sharded - and not ( - getattr(declaration.update, "supports_moe_cp_gather", False) - and not _gathers_over_attention_cp() - ) ): raise NotImplementedError("MHC with a gather over attention CP") handoff = None diff --git a/python/sglang/srt/layers/layer_boundary/factories.py b/python/sglang/srt/layers/layer_boundary/factories.py index ac1e21b65096..5efd7ce9dff0 100644 --- a/python/sglang/srt/layers/layer_boundary/factories.py +++ b/python/sglang/srt/layers/layer_boundary/factories.py @@ -357,11 +357,7 @@ def declare_ffn( def _resolve(stage, variant, following=None): axes, attention, local, full = _rows(variant) if stage.update.at_producer: - if ( - stage.sparse - and sparse_moe_gathers_over_moe_cp() - and not getattr(stage.update, "supports_moe_cp_gather", False) - ): + if stage.sparse and sparse_moe_gathers_over_moe_cp(): raise NotImplementedError( "MHC does not support a MoE gathered over the MoE-CP group" ) diff --git a/python/sglang/srt/layers/layer_boundary/prepare.py b/python/sglang/srt/layers/layer_boundary/prepare.py index 23b9b9038073..e923bcd6a872 100644 --- a/python/sglang/srt/layers/layer_boundary/prepare.py +++ b/python/sglang/srt/layers/layer_boundary/prepare.py @@ -443,11 +443,9 @@ def _mlp_input_gather_moe_cp( # so no rank hits this path while others proceed to the allgather. # - During decode: moe_cp allgather is skipped (guarded by is_context_parallel_extend). # - CUDA graph warmup: not applicable when --cuda-graph-backend-prefill=disabled is used. - if hidden_states.shape[0] == 0 and update.adds_plainly: + if hidden_states.shape[0] == 0: return hidden_states, residual - # Stateful nonlinear reads must still initialize their FFN residual state - # on empty batches, even though there are no rows to communicate. hidden_states, residual = gather( hidden_states, residual, forward_batch, layernorm, update=update, cache=cache ) diff --git a/python/sglang/srt/layers/layer_boundary/residual/__init__.py b/python/sglang/srt/layers/layer_boundary/residual/__init__.py index 4c632e59cd96..19c12d9882d0 100644 --- a/python/sglang/srt/layers/layer_boundary/residual/__init__.py +++ b/python/sglang/srt/layers/layer_boundary/residual/__init__.py @@ -34,11 +34,6 @@ class StageUpdate(Protocol): The shard conversion methods must move any update-associated state together with the residual. Nonlinear updates must not claim adds_plainly. - - Producer-written updates may optionally set supports_moe_cp_gather=True - when their auxiliary state stays on local residual rows while the FFN - input is gathered. The boundary returns the FFN output to those rows - before applying the update. Other producer-written updates remain gated. """ # A plain add, which one rank may run before the sum it adds into diff --git a/python/sglang/srt/layers/layer_boundary/residual/gated.py b/python/sglang/srt/layers/layer_boundary/residual/gated.py deleted file mode 100644 index a2df02b6575d..000000000000 --- a/python/sglang/srt/layers/layer_boundary/residual/gated.py +++ /dev/null @@ -1,148 +0,0 @@ -# Copyright 2023-2024 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Gated hyper-connections expressed as stage reads and residual updates.""" - -from typing import TYPE_CHECKING, Optional - -import torch - -from sglang.srt.layers.layer_boundary.adapters.attention import ( - _redistribute_to_attn_tp_shards, -) -from sglang.srt.layers.layer_boundary.ops import gather_attention_tp -from sglang.srt.layers.layer_boundary.residual import LayerResidual - -if TYPE_CHECKING: - from sglang.srt.layers.hyperconnection import GatedResidual - - -class _BranchState: - def __init__(self, connection: "GatedResidual"): - self.connection = connection - self.normalized: Optional[torch.Tensor] = None - - def mix(self, residual): - hidden_states, (residual, self.normalized) = self.connection.mix(residual) - return hidden_states, residual - - def combine(self, hidden_states, residual): - if self.normalized is None: - raise RuntimeError("a gated residual update requires its stage's read") - output = self.connection.combine(hidden_states, (residual, self.normalized)) - self.normalized = None - return output - - def to_attn_tp_shard(self, residual): - if self.normalized is None: - raise RuntimeError("a gated residual shard requires its stage's read") - self.normalized = _redistribute_to_attn_tp_shards(self.normalized) - return _redistribute_to_attn_tp_shards(residual) - - def from_attn_tp_shards(self, residual): - if self.normalized is None: - raise RuntimeError("a gated residual gather requires its stage's read") - # The tuple path allocates separate outputs. Gathering both through the - # shared DP scratch buffer would overwrite the first gathered tensor. - residual, self.normalized = gather_attention_tp( - (residual, self.normalized), None - ) - return residual - - -class _GatedRead: - norms_plainly = False - before_gather = True - - def __init__(self, state: _BranchState, *, enters_stack: bool): - self.state = state - self.enters_stack = enters_stack - - def enter(self, hidden_states): - if not self.enters_stack: - raise RuntimeError("a gated FFN does not enter the layer stack") - connection = self.state.connection - width = hidden_states.shape[-1] - if width == connection.hc_count * connection.hidden_size: - return hidden_states - if width != connection.hidden_size: - raise ValueError(f"unexpected gated residual input width: {width}") - return torch.cat([hidden_states] * connection.hc_count, dim=-1) - - def read(self, residual, norm, quant_format="", post_residual_addition=None): - if norm is not None or quant_format: - raise NotImplementedError("gated residual mixing owns its normalization") - if post_residual_addition is not None: - residual = residual + post_residual_addition - return self.state.mix(residual) - - def update_and_read( - self, - update, - hidden_states, - residual, - norm, - quant_format="", - post_residual_addition=None, - ): - return self.read( - update.update(hidden_states, residual), - norm, - quant_format, - post_residual_addition, - ) - - -class _GatedUpdate: - adds_plainly = False - can_defer_across_layers = False - # MoE CP gathers only the compute input. The boundary returns the output - # to the local residual rows before applying this update. - supports_moe_cp_gather = True - - def __init__(self, state: _BranchState, *, at_producer: bool): - self.state = state - self.at_producer = at_producer - - def update(self, hidden_states, residual): - return self.state.combine(hidden_states, residual) - - def residual_to_attn_tp_shard(self, residual): - return self.state.to_attn_tp_shard(residual) - - def residual_from_attn_tp_shards(self, residual): - return self.state.from_attn_tp_shards(residual) - - -class GatedResidualState: - """Adapt attention/FFN GatedResidual modules to the layer boundaries. - - The batch's residual stream retains the unnormalized hyper-input tensor. - Each branch retains the normalized input needed by its next combine, and - moves that auxiliary tensor whenever its residual changes token rows. - FFN updates run at the producer so no layer carries another layer's state. - The owning model continues to register and load the connection modules. - """ - - def __init__(self, attention: "GatedResidual", ffn: "GatedResidual"): - attention_state = _BranchState(attention) - ffn_state = _BranchState(ffn) - self._residual = LayerResidual( - attention_read=_GatedRead(attention_state, enters_stack=True), - attention_update=_GatedUpdate(attention_state, at_producer=False), - ffn_read=_GatedRead(ffn_state, enters_stack=False), - ffn_update=_GatedUpdate(ffn_state, at_producer=True), - ) - - def layer_residual(self) -> LayerResidual: - return self._residual diff --git a/python/sglang/srt/models/qwen4_exp.py b/python/sglang/srt/models/qwen4_exp.py index 759a9168cc16..7bcc07ca5aed 100644 --- a/python/sglang/srt/models/qwen4_exp.py +++ b/python/sglang/srt/models/qwen4_exp.py @@ -23,9 +23,13 @@ from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.layers.aux_hidden_states import AuxHiddenStateList from sglang.srt.layers.dp_attention import ( + attn_tp_all_gather, attn_tp_all_reduce, dp_gather_replicate, dp_scatter, + get_dp_global_num_tokens, + get_global_dp_buffer, + get_local_dp_buffer, is_allocation_symmetric, is_dp_attention_enabled, ) @@ -33,18 +37,10 @@ GatedResidual, HyperConnectionConfig, ) -from sglang.srt.layers.layer_boundary import ( - HandoffRows, - declare_attn, - declare_ffn, - get_attn_tp_context, - make_stages, -) -from sglang.srt.layers.layer_boundary.residual import batch as residual_batch -from sglang.srt.layers.layer_boundary.residual.gated import GatedResidualState +from sglang.srt.layers.layer_boundary import get_attn_tp_context from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.logits_processor import LogitsProcessorOutput -from sglang.srt.layers.moe import get_moe_a2a_backend +from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.modelopt_quant import ( @@ -76,7 +72,7 @@ make_ple_file_prefetcher, make_ple_file_rss_trimmer, ) -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_forward, get_parallel from sglang.srt.utils import get_bool_env_var, is_hip, logger _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() @@ -1324,6 +1320,8 @@ def _init_qwen4_exp_layer_extensions( for attr_name in ( "input_layernorm", "post_attention_layernorm", + "attn_boundary", + "ffn_boundary", ): if hasattr(self, attr_name): delattr(self, attr_name) @@ -1362,29 +1360,6 @@ def _init_qwen4_exp_layer_extensions( use_mix=True, use_combine=True, ) - residual = GatedResidualState( - self.attn_hyper_connection, self.mlp_hyper_connection - ).layer_residual() - sparse = bool(config.num_experts) - ffn = declare_ffn( - sparse=sparse, - read=residual.ffn_read, - update=residual.ffn_update, - dense_tp_size=None if sparse else get_parallel().tp_size, - handoff_rows=HandoffRows.ATTENTION, - ) - self.attn_boundary, self.ffn_boundary = make_stages( - ( - declare_attn( - read=residual.attention_read, - update=residual.attention_update, - ), - None, - ), - (ffn, None), - previous=ffn if layer_id != 0 else None, - terminal=layer_id == config.num_hidden_layers - 1, - ) def _prepare_qwen4_exp_attn( self, @@ -1416,10 +1391,8 @@ def _prepare_qwen4_exp_attn( ple_query, forward_batch, ple_batch ) - # Qwen4 hands complete hyper-connection streams between layers. PLE - # updates that tensor before the attention boundary reads its input. - hidden_states = residual_batch.written(hidden_states, forward_batch) - return self.attn_boundary.prepare(hidden_states, forward_batch), None + hidden_states, residual = self.attn_hyper_connection.mix(hidden_states) + return hidden_states, residual def _prepare_qwen4_exp_mlp( self, @@ -1427,21 +1400,70 @@ def _prepare_qwen4_exp_mlp( residual: Optional[torch.Tensor], forward_batch: ForwardBatch, ): - hidden_states = self.attn_boundary.finish(hidden_states, forward_batch) - return self.ffn_boundary.prepare(hidden_states, forward_batch), None + if not forward_batch.forward_mode.is_idle(): + hidden_states = attn_tp_all_reduce(hidden_states) + hidden_states = self.attn_hyper_connection.combine(hidden_states, residual) + hidden_states, residual = self.mlp_hyper_connection.mix(hidden_states) + return hidden_states, residual + + def _qwen4_exp_use_dp_moe_gather(self) -> bool: + return get_parallel().attn_dp_size > 1 and get_moe_a2a_backend().is_none() + + def _qwen4_exp_use_attn_tp_a2a_scatter(self) -> bool: + return get_parallel().attn_tp_size > 1 and not get_moe_a2a_backend().is_none() def _run_qwen4_exp_mlp( self, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - with self.ffn_boundary.exit(forward_batch) as ffn_exit: - if not self.config.num_experts: - hidden_states = self.mlp(hidden_states) - elif hidden_states.shape[0] or not get_moe_a2a_backend().is_none(): - # An idle all-to-all rank must still participate in MoE dispatch. - hidden_states = self.mlp(hidden_states, forward_batch) - return ffn_exit.finish(hidden_states) + if not self.config.num_experts: + return self.mlp(hidden_states) + + use_dp_moe_gather = self._qwen4_exp_use_dp_moe_gather() + use_attn_tp_a2a_scatter = self._qwen4_exp_use_attn_tp_a2a_scatter() + + if use_dp_moe_gather: + hidden_states, local_hidden_states = ( + get_global_dp_buffer(get_parallel().tp_group), + hidden_states, + ) + dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) + elif hidden_states.shape[0] == 0 and get_moe_a2a_backend().is_none(): + # Only safe to short-circuit an empty batch when the MoE holds no collective; + # under deepep an idle DP rank must still join dispatch/combine or peers hang. + return hidden_states + + attn_tp_chunks = None + if use_attn_tp_a2a_scatter: + attn_tp_size = get_parallel().attn_tp_size + attn_tp_chunks = list(hidden_states.tensor_split(attn_tp_size)) + hidden_states = attn_tp_chunks[get_parallel().attn_tp_rank].contiguous() + + use_reduce_scatterv = use_dp_moe_gather and should_use_dp_reduce_scatterv() + with get_forward().scoped(mlp_reduce_scatter=use_reduce_scatterv): + hidden_states = self.mlp(hidden_states, forward_batch) + + if use_dp_moe_gather: + hidden_states, global_hidden_states = ( + get_local_dp_buffer(get_parallel().tp_group), + hidden_states, + ) + if use_reduce_scatterv: + get_parallel().tp_group.reduce_scatterv( + global_hidden_states, + output=hidden_states, + sizes=get_dp_global_num_tokens(), + ) + else: + dp_scatter(hidden_states, global_hidden_states, forward_batch) + elif use_attn_tp_a2a_scatter: + assert attn_tp_chunks is not None + gathered = [torch.empty_like(t) for t in attn_tp_chunks] + attn_tp_all_gather(gathered, hidden_states.contiguous()) + hidden_states = torch.cat(gathered) + + return hidden_states def _postprocess_qwen4_exp_layer( self, @@ -1449,6 +1471,7 @@ def _postprocess_qwen4_exp_layer( residual: Optional[torch.Tensor], forward_batch: ForwardBatch, ): + hidden_states = self.mlp_hyper_connection.combine(hidden_states, residual) return hidden_states, None @@ -1752,8 +1775,6 @@ def forward( _commit_ple_batch(ple_batch, forward_batch) - hidden_states = residual_batch.take_output(hidden_states, forward_batch) - if not self.pp_group.is_last_rank: proxy_tensors = {"hidden_states": hidden_states} if forward_batch.mm_input_embeds is not None: diff --git a/test/registered/unit/layer_boundary/test_gated_residual.py b/test/registered/unit/layer_boundary/test_gated_residual.py deleted file mode 100644 index 9f5952831be9..000000000000 --- a/test/registered/unit/layer_boundary/test_gated_residual.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Gated hyper-connection state follows boundary residual rows and ordering.""" - -import unittest -from types import SimpleNamespace -from unittest.mock import Mock, patch - -import torch - -from sglang.srt.layers.layer_boundary.adapters import attention -from sglang.srt.layers.layer_boundary.residual import gated -from sglang.srt.layers.layer_boundary.residual.gated import GatedResidualState -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=2, suite="base-a-test-cpu") - - -class ToyGatedResidual: - """A nonlinear mix/combine that consumes both raw and normalized input.""" - - hc_count = 2 - hidden_size = 3 - - def __init__(self, scale): - self.scale = scale - - def mix(self, hyper_input): - normalized = hyper_input / (1 + hyper_input.square().mean(-1, keepdim=True)) - mixed = normalized.reshape(-1, self.hc_count, self.hidden_size).sum(1) - return mixed * self.scale, (hyper_input, normalized) - - def combine(self, hidden_states, residuals): - residual, normalized = residuals - injected = hidden_states.repeat(1, self.hc_count) - return residual + self.scale * injected * torch.sigmoid(normalized) - - -class TestGatedResidual(CustomTestCase): - def setUp(self): - self.attention = ToyGatedResidual(0.5) - self.ffn = ToyGatedResidual(0.75) - self.state = GatedResidualState(self.attention, self.ffn) - self.residual = self.state.layer_residual() - self.hidden = torch.arange(24, dtype=torch.float32).reshape(4, 6) / 10 - - def test_preserves_attention_then_ffn_mix_combine_order(self): - expected_attn, expected_residual = self.attention.mix(self.hidden) - expected_after_attn = self.attention.combine( - expected_attn.square(), expected_residual - ) - expected_ffn, expected_residual = self.ffn.mix(expected_after_attn) - expected = self.ffn.combine(expected_ffn.sin(), expected_residual) - - attn, residual = self.residual.attention_read.read(self.hidden, None) - ffn, residual = self.residual.ffn_read.update_and_read( - self.residual.attention_update, attn.square(), residual, None - ) - actual = self.residual.ffn_update.update(ffn.sin(), residual) - - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - self.assertIsNone(self.residual.attention_update.state.normalized) - self.assertIsNone(self.residual.ffn_update.state.normalized) - self.assertFalse(self.residual.attention_update.adds_plainly) - self.assertFalse(self.residual.attention_update.at_producer) - self.assertTrue(self.residual.ffn_update.at_producer) - self.assertFalse(self.residual.ffn_update.can_defer_across_layers) - - def test_first_attention_expands_embeddings_only_once(self): - embedding = self.hidden[:, : self.attention.hidden_size] - expanded = self.residual.attention_read.enter(embedding) - torch.testing.assert_close(expanded, embedding.repeat(1, 2)) - self.assertIs(self.residual.attention_read.enter(expanded), expanded) - with self.assertRaisesRegex(ValueError, "width"): - self.residual.attention_read.enter(torch.empty(4, 5)) - - def test_shard_moves_normalized_state_with_raw_residual(self): - mixed, residual = self.residual.attention_read.read(self.hidden, None) - expected = self.attention.combine(mixed, self.attention.mix(self.hidden)[1]) - update = self.residual.attention_update - with patch.object( - attention, - "get_parallel", - lambda: SimpleNamespace(attn_tp_size=2, attn_tp_rank=1), - ): - shard = update.residual_to_attn_tp_shard(residual) - actual = update.update(mixed[2:], shard) - torch.testing.assert_close(actual, expected[2:], rtol=0, atol=0) - - def test_gather_moves_both_states_without_shared_scratch_aliasing(self): - full_mixed, full_residuals = self.attention.mix(self.hidden) - _, local_residual = self.residual.attention_read.read(self.hidden[2:], None) - update = self.residual.attention_update - - def gather_pair(pair, forward_batch): - self.assertIsNone(forward_batch) - torch.testing.assert_close(pair[0], self.hidden[2:]) - torch.testing.assert_close(pair[1], full_residuals[1][2:]) - return tuple(t.clone() for t in full_residuals) - - with patch.object(gated, "gather_attention_tp", side_effect=gather_pair) as ag: - gathered = update.residual_from_attn_tp_shards(local_residual) - ag.assert_called_once() - actual = update.update(full_mixed, gathered) - expected = self.attention.combine(full_mixed, full_residuals) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - def test_empty_stage_preserves_widened_residual_shape(self): - hidden = torch.empty(0, self.attention.hidden_size) - expanded = self.residual.attention_read.enter(hidden) - attn, residual = self.residual.attention_read.read(expanded, None) - ffn, residual = self.residual.ffn_read.update_and_read( - self.residual.attention_update, attn, residual, None - ) - output = self.residual.ffn_update.update(ffn, residual) - self.assertEqual(output.shape, (0, 6)) - self.assertIsNone(self.residual.ffn_update.state.normalized) - - def test_update_uses_actual_producer_and_addition_precedes_mix(self): - updated = self.hidden + 2 - producer = SimpleNamespace(update=Mock(return_value=updated)) - addition = torch.full_like(self.hidden, 0.5) - actual, residual = self.residual.ffn_read.update_and_read( - producer, - torch.empty(4, 3), - self.hidden, - None, - post_residual_addition=addition, - ) - producer.update.assert_called_once() - expected, _ = self.ffn.mix(updated + addition) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - torch.testing.assert_close(residual, updated + addition) - - def test_update_cannot_reuse_auxiliary_state_after_consumption(self): - mixed, residual = self.residual.attention_read.read(self.hidden, None) - self.residual.attention_update.update(mixed, residual) - with self.assertRaisesRegex(RuntimeError, "requires its stage's read"): - self.residual.attention_update.update(mixed, residual) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/layer_boundary/test_moe_output_reduction.py b/test/registered/unit/layer_boundary/test_moe_output_reduction.py index f3d68505cf7b..b2e9c7807026 100644 --- a/test/registered/unit/layer_boundary/test_moe_output_reduction.py +++ b/test/registered/unit/layer_boundary/test_moe_output_reduction.py @@ -169,7 +169,10 @@ def test_direct_dp_exits_publish_their_selected_sum(self): # These models own their DP exit instead of using StageBoundary. import __future__ - for filename, method in (("deepseek_v4.py", "_run_moe_ffn_dp_sync"),): + for filename, method in ( + ("deepseek_v4.py", "_run_moe_ffn_dp_sync"), + ("qwen4_exp.py", "_run_qwen4_exp_mlp"), + ): path = MODELS_DIR / filename node = next( n @@ -226,6 +229,8 @@ def scatter(output, value, batch): model = types.SimpleNamespace( dsa_enable_prefill_cp=False, config=types.SimpleNamespace(num_experts=4), + _qwen4_exp_use_dp_moe_gather=lambda: True, + _qwen4_exp_use_attn_tp_a2a_scatter=lambda: False, mlp=lambda value, batch, **kwargs: reduce_moe_output(value), ) batch = types.SimpleNamespace( diff --git a/test/registered/unit/layer_boundary/test_qwen4_gated_boundary.py b/test/registered/unit/layer_boundary/test_qwen4_gated_boundary.py deleted file mode 100644 index 1c75317b9668..000000000000 --- a/test/registered/unit/layer_boundary/test_qwen4_gated_boundary.py +++ /dev/null @@ -1,457 +0,0 @@ -"""Qwen4's actual layer orchestration through production stage boundaries. - -Only the hyper-connection kernels and distributed collectives are replaced: -the nonlinear reference catches mixing before reduction or combining before -the FFN output returns to the residual's rows. -""" - -import __future__ - -import ast -import itertools -import unittest -from contextlib import ExitStack, contextmanager, nullcontext -from pathlib import Path -from types import SimpleNamespace -from unittest.mock import patch - -import torch - -import sglang -from sglang.srt.layers import layernorm_sp -from sglang.srt.layers.dp_attention import DpPaddingMode -from sglang.srt.layers.layer_boundary import ( - HandoffRows, - declare_attn, - declare_ffn, - make_stages, -) -from sglang.srt.layers.layer_boundary.residual import batch as residual_batch -from sglang.srt.layers.layer_boundary.residual.gated import GatedResidualState -from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.srt.runtime_context import get_forward -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.communicator_patch import patch_communicator -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=5, suite="base-a-test-cpu") - -HIDDEN = 3 - - -def mix_reference(value): - first, second = value.chunk(2, dim=-1) - return torch.tanh(first + second / 2) - - -def combine_reference(value, residual): - return residual + value.repeat(1, 2) * torch.sigmoid(residual) - - -class Connection: - hc_count = 2 - hidden_size = HIDDEN - - def __init__(self): - self.mixed = [] - self.combined = [] - - def mix(self, value): - self.mixed.append(value.clone()) - return mix_reference(value), (value, torch.sigmoid(value)) - - def combine(self, value, residuals): - residual, normalized = residuals - self.combined.append(value.clone()) - torch.testing.assert_close(normalized, torch.sigmoid(residual)) - return residual + value.repeat(1, 2) * normalized - - -def load_mixin(namespace): - # Avoid importing GPU-only attention kernels just to execute the model's - # orchestration. The methods compiled here are the production source. - path = Path(sglang.__file__).parent / "srt/models/qwen4_exp.py" - node = next( - node - for node in ast.parse(path.read_text()).body - if isinstance(node, ast.ClassDef) and node.name == "Qwen4ExpLayerExtensionMixin" - ) - exec( - compile( - ast.Module(body=[node], type_ignores=[]), - str(path), - "exec", - flags=__future__.annotations.compiler_flag, - ), - namespace, - ) - return namespace[node.name] - - -class Harness: - def __init__( - self, - *, - dp=1, - tp=1, - a2a=False, - rsv=False, - rows=2, - rank=0, - cp_rows=None, - cp_rank=0, - ): - self.dp, self.tp, self.a2a, self.rsv = dp, tp, a2a, rsv - self.rows, self.rank = rows, rank - self.cp_size = len(cp_rows) if cp_rows is not None else 1 - self.cp_rank = cp_rank - self.events = [] - self.compute_inputs = [] - self.compute_flags = [] - self.expected_gathered = None - self.expected_cp_inputs = None - self.parallel = SimpleNamespace( - tp_size=dp * tp * self.cp_size, - tp_rank=rank, - attn_dp_size=dp, - attn_dp_rank=0, - attn_tp_size=tp, - attn_tp_rank=rank, - attn_cp_size=self.cp_size, - attn_cp_rank=cp_rank, - moe_dp_size=1, - moe_ep_size=1, - moe_tp_size=dp * tp * self.cp_size, - moe_dense_tp_size=None, - dwdp_size=1, - enable_dp_attention=dp > 1, - enable_prefill_cp=self.cp_size > 1, - enable_attn_tp_input_scattered=False, - tp_group=SimpleNamespace(reduce_scatterv=self.reduce_scatterv), - attn_tp_group=SimpleNamespace(), - ) - self.batch = SimpleNamespace( - forward_mode=ForwardMode.DECODE if rows else ForwardMode.IDLE, - residual_stream=None, - dp_padding_mode=DpPaddingMode.SUM_LEN, - attn_cp_metadata=( - SimpleNamespace(per_rank_actual_token=cp_rows) - if cp_rows is not None - else None - ), - global_num_tokens_cpu=[rows] * dp, - input_ids=torch.zeros(rows, dtype=torch.int64), - ) - if cp_rows is not None: - self.batch.forward_mode = ForwardMode.EXTEND - - def attention_sum(self, value): - self.events.append("attention_sum") - return value * self.tp - - def attention_scatter(self, output, value): - self.events.append("attention_reduce_scatter") - output.copy_((value * self.tp).tensor_split(self.tp)[self.rank]) - - def attention_gather(self, output, value): - self.events.append("attention_gather") - torch.testing.assert_close( - value, self.expected_gathered.tensor_split(self.tp)[self.rank] - ) - output.copy_(self.expected_gathered) - - def gather_dp(self, output, value, batch, cp_shard_counts=None): - self.events.append("dp_gather") - output.copy_(torch.cat([value + i * 10 for i in range(self.dp)])) - - def gather_cp(self, output, value): - self.events.append("cp_gather") - torch.testing.assert_close( - value, self.expected_cp_inputs.tensor_split(self.cp_size)[self.cp_rank] - ) - output.copy_(self.expected_cp_inputs) - - def scatter_dp(self, output, value, batch): - self.events.append("dp_scatter") - output.copy_(value[: self.rows]) - - def reduce_scatterv(self, value, *, output, sizes): - self.events.append("dp_reduce_scatterv") - output.copy_(value[: self.rows] * self.parallel.tp_size) - - def mlp(self, value, *args): - self.compute_inputs.append(value.clone()) - skipped = get_forward().mlp_reduce_scatter - self.compute_flags.append(skipped) - return value * (3 / self.parallel.tp_size if skipped else 3) - - @contextmanager - def running(self): - backend = SimpleNamespace(is_none=lambda: not self.a2a) - with ExitStack() as stack: - overrides = { - "get_parallel": lambda: self.parallel, - "get_moe_a2a_backend": lambda: backend, - "is_moe_input_scattered_across_dp_ranks": lambda: self.a2a, - "is_enable_moe_cp_allgather": lambda: self.cp_size > 1, - "get_moe_cp_size": lambda: self.cp_size, - "get_moe_cp_rank": lambda: self.cp_rank, - "moe_cp_all_gather_into_tensor": self.gather_cp, - "is_dsa_enable_prefill_cp": lambda: False, - "is_mla_cp_enabled": lambda: False, - "get_spec": lambda: SimpleNamespace(speculative_algorithm=None), - "get_lora": lambda: SimpleNamespace(enable_lora=False), - "get_exec": lambda: SimpleNamespace( - comm=SimpleNamespace( - boundary_reduction="rs+rsv", enable_quant_communications=False - ), - overlap=SimpleNamespace(enable_two_batch_overlap=False), - ), - "get_attn_tp_context": lambda: SimpleNamespace(input_scattered=False), - "post_experts_reduction_group": lambda: self.parallel.tp_group, - "should_use_dp_reduce_scatterv": lambda: self.rsv, - "can_use_dp_reduce_scatter": lambda: False, - "use_symmetric_memory": lambda *a, **k: nullcontext(), - "is_allocation_symmetric": lambda: False, - "attention_tensor_model_parallel_all_reduce": self.attention_sum, - "attn_tp_reduce_scatter_tensor": self.attention_scatter, - "attn_tp_all_gather_into_tensor": self.attention_gather, - "get_global_dp_buffer": lambda group: torch.empty( - self.rows * self.dp, HIDDEN - ), - "get_local_dp_buffer": lambda group, hidden_size=None: torch.empty( - self.rows, hidden_size or HIDDEN - ), - "get_dp_global_num_tokens": lambda: [self.rows] * self.dp, - "dp_gather_replicate": self.gather_dp, - "dp_scatter": self.scatter_dp, - } - for name, replacement in overrides.items(): - stack.enter_context(patch_communicator(name, replacement)) - stack.enter_context( - patch.object(layernorm_sp, "layernorm_sp_enabled", lambda: False) - ) - stack.enter_context( - get_forward().scoped( - mlp_reduce_scatter=False, - fuse_mlp_allreduce=False, - sp_active=False, - attn_input_scattered=False, - ) - ) - self.namespace = dict( - torch=torch, - residual_batch=residual_batch, - GatedResidualState=GatedResidualState, - GatedResidual=lambda *a, **k: Connection(), - HyperConnectionConfig=SimpleNamespace, - HandoffRows=HandoffRows, - declare_attn=declare_attn, - declare_ffn=declare_ffn, - make_stages=make_stages, - get_parallel=lambda: self.parallel, - get_moe_a2a_backend=lambda: backend, - _get_ple_forward_mode=lambda batch: batch.forward_mode, - ) - yield - - def model(self, *, sparse, layer_id=0, num_layers=1): - model = load_mixin(self.namespace)() - model.config = SimpleNamespace( - num_experts=4 if sparse else 0, - num_hidden_layers=num_layers, - hc_count=2, - hidden_size=HIDDEN, - ple_layer_ids=[], - hc_lowrank=2, - rms_norm_eps=1e-5, - ) - model._init_qwen4_exp_layer_extensions(model.config, layer_id) - model.mlp = self.mlp - return model - - def forward(self, model, value, *, ple_batch=None): - hidden, residual = model._prepare_qwen4_exp_attn( - value, None, self.batch, ple_batch=ple_batch - ) - if not self.batch.forward_mode.is_idle(): - hidden = hidden * (2 / self.tp) - hidden, residual = model._prepare_qwen4_exp_mlp(hidden, residual, self.batch) - hidden = model._run_qwen4_exp_mlp(hidden, self.batch) - return model._postprocess_qwen4_exp_layer(hidden, residual, self.batch) - - -def reference(value): - if value.shape[-1] == HIDDEN: - value = value.repeat(1, 2) - after_attention = combine_reference(mix_reference(value) * 2, value) - mlp_input = mix_reference(after_attention) - return combine_reference(mlp_input * 3, after_attention), mlp_input - - -class TestQwen4Boundaries(CustomTestCase): - def test_dense_and_moe_complete_attention_sum_before_nonlinear_mix(self): - value = torch.arange(6, dtype=torch.float32).reshape(2, HIDDEN) / 10 - expected, mlp_input = reference(value) - for sparse in (False, True): - with self.subTest(sparse=sparse): - harness = Harness(tp=2) - with harness.running(): - model = harness.model(sparse=sparse) - output, residual = harness.forward(model, value) - torch.testing.assert_close(output, expected) - torch.testing.assert_close(harness.compute_inputs[0], mlp_input) - self.assertIsNone(residual) - self.assertEqual(harness.events, ["attention_sum"]) - self.assertIsNone(harness.batch.residual_stream.pending) - - def test_dp_ffn_gathers_mixed_input_and_combines_after_return(self): - value = torch.arange(6, dtype=torch.float32).reshape(2, HIDDEN) / 10 - expected, local_input = reference(value) - for sparse, rsv in itertools.product((False, True), repeat=2): - with self.subTest(sparse=sparse, reduce_scatterv=rsv): - harness = Harness(dp=2, tp=2, rsv=rsv) - with harness.running(): - output, residual = harness.forward( - harness.model(sparse=sparse), value - ) - self.assertFalse(get_forward().mlp_reduce_scatter) - torch.testing.assert_close(output, expected) - torch.testing.assert_close( - harness.compute_inputs[0], - torch.cat([local_input, local_input + 10]), - ) - self.assertEqual(harness.compute_flags, [rsv]) - self.assertEqual( - harness.events, - [ - "attention_sum", - "dp_gather", - "dp_reduce_scatterv" if rsv else "dp_scatter", - ], - ) - self.assertIsNone(residual) - - def test_a2a_moves_gated_state_with_attention_tp_slice(self): - value = torch.arange(12, dtype=torch.float32).reshape(4, HIDDEN) / 10 - expected, mlp_input = reference(value) - for rank in (0, 1): - with self.subTest(rank=rank): - harness = Harness(tp=2, a2a=True, rows=4, rank=rank) - harness.expected_gathered = expected - with harness.running(): - output, residual = harness.forward( - harness.model(sparse=True), value - ) - torch.testing.assert_close(output, expected) - torch.testing.assert_close( - harness.compute_inputs[0], mlp_input.tensor_split(2)[rank] - ) - self.assertEqual( - harness.events, ["attention_reduce_scatter", "attention_gather"] - ) - self.assertIsNone(residual) - - def test_idle_a2a_still_participates_in_expert_dispatch(self): - value = torch.empty(0, HIDDEN) - for a2a in (False, True): - with self.subTest(a2a=a2a): - harness = Harness(a2a=a2a, rows=0) - with harness.running(): - output, residual = harness.forward( - harness.model(sparse=True), value - ) - self.assertEqual(len(harness.compute_inputs), int(a2a)) - self.assertEqual(output.shape, (0, 2 * HIDDEN)) - self.assertIsNone(residual) - - def test_moe_cp_takes_back_uneven_rows_before_gated_combine(self): - values = [ - torch.arange(6, dtype=torch.float32).reshape(2, HIDDEN) / 10, - torch.arange(3, dtype=torch.float32).reshape(1, HIDDEN) / 4, - ] - references = [reference(value) for value in values] - gathered_inputs = torch.cat( - [references[0][1], references[1][1], torch.zeros(1, HIDDEN)] - ) - for rank, value in enumerate(values): - with self.subTest(cp_rank=rank): - harness = Harness(rows=len(value), cp_rows=[2, 1], cp_rank=rank) - harness.expected_cp_inputs = gathered_inputs - with harness.running(): - output, residual = harness.forward( - harness.model(sparse=True), value - ) - torch.testing.assert_close(output, references[rank][0]) - torch.testing.assert_close(harness.compute_inputs[0], gathered_inputs) - self.assertEqual(harness.events, ["cp_gather"]) - self.assertIsNone(residual) - - def test_idle_with_moe_cp_still_consumes_gated_residual_reads(self): - harness = Harness(rows=0, cp_rows=[0, 0]) - harness.batch.forward_mode = ForwardMode.IDLE - harness.batch.attn_cp_metadata = None - with harness.running(): - output, residual = harness.forward( - harness.model(sparse=True), torch.empty(0, HIDDEN) - ) - self.assertEqual(output.shape, (0, 2 * HIDDEN)) - self.assertIsNone(residual) - self.assertEqual(harness.compute_inputs, []) - self.assertEqual(harness.events, []) - - def test_intermediate_layers_complete_stream_and_reuse_auxiliary_state(self): - value = torch.arange(6, dtype=torch.float32).reshape(2, HIDDEN) / 10 - harness = Harness(tp=2) - with harness.running(): - layers = [ - harness.model(sparse=True, layer_id=i, num_layers=2) for i in range(2) - ] - for repetition in range(2): - hidden = value + repetition - expected = hidden - for model in layers: - hidden, residual = harness.forward(model, hidden) - expected = reference(expected)[0] - torch.testing.assert_close(hidden, expected) - self.assertIsNone(residual) - self.assertIsNone(harness.batch.residual_stream.pending) - self.assertIs(harness.batch.residual_stream.residual, hidden) - - def test_ple_reads_expanded_input_before_attention_mix(self): - value = torch.ones(2, HIDDEN) / 4 - calls = [] - - def ple(hidden, batch, ple_batch): - calls.append(hidden.clone()) - return hidden / 2 - - harness = Harness() - with harness.running(): - model = harness.model(sparse=False) - model.ple = ple - output, _ = harness.forward(model, value, ple_batch=object()) - torch.testing.assert_close(calls[0], value.repeat(1, 2)) - torch.testing.assert_close(output, reference(value * 1.5)[0]) - - def test_ple_idle_enters_collective_path_and_missing_active_batch_fails(self): - from unittest.mock import Mock - - for rows in (0, 2): - with self.subTest(rows=rows): - harness = Harness(rows=rows) - with harness.running(): - model = harness.model(sparse=False) - model.ple = Mock() - value = torch.empty(rows, HIDDEN) - if rows: - with self.assertRaisesRegex(RuntimeError, "missing its batch"): - harness.forward(model, value) - else: - output, _ = harness.forward(model, value) - model.ple.forward_idle.assert_called_once_with(harness.batch) - self.assertEqual(output.shape, (0, 2 * HIDDEN)) - - -if __name__ == "__main__": - unittest.main() From 4173eb92f475bc76548cd94536e7ade2399ac0be Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Tue, 29 Sep 2026 22:49:21 -0700 Subject: [PATCH 5/6] Support cached prefixes in QSA context-parallel prefill --- .../srt/layers/attention/qsa/metadata.py | 3 + .../srt/layers/attention/qsa/qsa_indexer.py | 5 +- .../attention/qwen_sparse_attn_backend.py | 40 +++- test/registered/unit/cp/test_qsa_cp_unit.py | 188 ++++++++++++++++++ 4 files changed, 226 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/layers/attention/qsa/metadata.py b/python/sglang/srt/layers/attention/qsa/metadata.py index 8c161efa9823..4dd3705d2a88 100644 --- a/python/sglang/srt/layers/attention/qsa/metadata.py +++ b/python/sglang/srt/layers/attention/qsa/metadata.py @@ -54,6 +54,9 @@ class QSAIndexerMetadata(msgspec.Struct, frozen=True): - ``token_to_batch_idx`` maps every query/token row handled by the indexer onto a row of ``sequence_lengths``/``token_slot_table``; DP attention token padding adds physical rows beyond this mapping, never inside it. + For CP extend, it covers the global packed *new* tokens used by the + compression write plan. Query selection uses a zigzag-sharded copy; + cached prefix tokens are K/V context, not rows in this mapping. - For the paged modes the mapping is the identity (``arange(num_query_rows)``), so page-table/MQA inputs built per ``sequence_lengths`` row line up with per-query sparse-attention rows. diff --git a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py index 464fbec36c0d..3c1923bf2ce7 100644 --- a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py +++ b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py @@ -569,8 +569,9 @@ def forward_cuda_cp( ) # Local q (normed + roped) and raw local k; no ring store here. q, token_k, _ = self.project_qk(hidden_states, positions) - # Raw keys of the whole sequence in global token order: the ring store - # and the compression use the global write plan unchanged. + # Raw keys of all new tokens in global packed order: the ring store + # and compression use the global write plan unchanged. Cached prefix + # groups already reside in the compressed-K pool on every rank. token_k_local = token_k.reshape(num_local, -1).contiguous() token_k_full = cp_materialize_global_token_order( token_k_local, forward_batch diff --git a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py index e7f3602d5162..b37b0665fe49 100644 --- a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py +++ b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py @@ -1386,8 +1386,9 @@ def _forward_extend_cp( save_kv_cache: bool, ) -> torch.Tensor: """Prefill CP: this rank holds the query rows of its two zigzag blocks - (all heads) and the K/V of the same rows. All-gather K/V into global + (all heads) and the K/V of the same rows. All-gather new K/V into global token order (also written to the KV pool so decode sees every token), + prepend each request's cached prefix from the replicated local KV pool, then run the chunk-prefill sparse kernel with one entry per zigzag block: q rows are contiguous per block, ``kv_lens`` is the block's absolute end position (bottom-right causal), ``cu_k`` is the request's @@ -1401,12 +1402,6 @@ def _forward_extend_cp( f"got {strategy_name}." ) meta = forward_batch.attn_cp_metadata - prefix_lens = getattr(forward_batch, "extend_prefix_lens_cpu", None) - if prefix_lens is not None and any(int(x) for x in prefix_lens): - raise NotImplementedError( - "QSA prefill CP requires zero prefix (disable the radix cache and " - "chunked prefill)." - ) q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) num_output_rows = q.shape[0] k = k.reshape(-1, layer.tp_k_head_num, layer.qk_head_dim) @@ -1420,7 +1415,7 @@ def _forward_extend_cp( v_full = v_full.reshape(-1, v.shape[1], v.shape[2]).contiguous() if save_kv_cache: # out_cache_loc was trimmed to the full extend token count by - # prepare_cp_forward; write the complete sequence on every rank. + # prepare_cp_forward; write all new tokens on every rank. self.token_to_kv_pool.set_kv_buffer( layer, forward_batch.out_cache_loc, k_full, v_full ) @@ -1428,6 +1423,35 @@ def _forward_extend_cp( q = q[:num_local_rows].contiguous() indices = topk_indices[:num_local_rows].to(torch.int32).contiguous() seq_lens = [int(x) for x in forward_batch.seq_lens_cpu[: meta.bs]] + extend_lens = [int(x) for x in forward_batch.extend_seq_lens_cpu[: meta.bs]] + prefix_lens = [seq - extend for seq, extend in zip(seq_lens, extend_lens)] + if any(prefix_lens): + # The all-gather contains only new tokens. Prefix K/V is already + # replicated on each CP rank; pack prefix + new tokens per request + # so absolute logical top-k indices address the full context. + pool = self.token_to_kv_pool + k_buffer = pool.get_key_buffer(layer.layer_id) + v_buffer = pool.get_value_buffer(layer.layer_id) + req_to_token = self.req_to_token_pool.req_to_token + req_indices = forward_batch.req_pool_indices[: meta.bs].tolist() + k_parts, v_parts = [], [] + extend_start = 0 + for req_idx, prefix_len, extend_len in zip( + req_indices, prefix_lens, extend_lens + ): + if prefix_len: + slots = req_to_token[req_idx, :prefix_len].long() + # QSA stores K/V without scales; casting also supports an + # FP8 prefix alongside the unquantized new tokens. + k_parts.append(k_buffer.index_select(0, slots).to(k_full.dtype)) + v_parts.append(v_buffer.index_select(0, slots).to(v_full.dtype)) + extend_end = extend_start + extend_len + k_parts.append(k_full[extend_start:extend_end]) + v_parts.append(v_full[extend_start:extend_end]) + extend_start = extend_end + # Use gathered new K/V even when save_kv_cache=False; those rows + # need not have been written to the cache by the caller. + k_full, v_full = torch.cat(k_parts), torch.cat(v_parts) base = [0] for length in seq_lens[:-1]: base.append(base[-1] + length) diff --git a/test/registered/unit/cp/test_qsa_cp_unit.py b/test/registered/unit/cp/test_qsa_cp_unit.py index c257fbe1237c..eb0763d883c6 100644 --- a/test/registered/unit/cp/test_qsa_cp_unit.py +++ b/test/registered/unit/cp/test_qsa_cp_unit.py @@ -5,11 +5,14 @@ from unittest.mock import Mock import pytest +import torch from sglang.srt.layers.attention import qwen_sparse_attn_backend from sglang.srt.layers.attention.qsa import qsa_indexer from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer from sglang.srt.layers.attention.qwen_sparse_attn_backend import QwenSparseAttnBackend +from sglang.srt.layers.cp import base as cp_base +from sglang.srt.layers.cp import zigzag from sglang.srt.layers.cp.base import ContextParallelStrategyKind from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.test.ci.ci_register import register_cpu_ci @@ -67,5 +70,190 @@ def test_cuda_indexer_forwards_global_rope_to_cp_path(): ) +def _cpu_sparse_chunk_attention(q, k, v, indices, cu_q, cu_k, kv_lens, scale): + """CPU oracle for the chunk kernel's packed K/V and causal conventions.""" + output = torch.empty_like(q) + for entry in range(kv_lens.numel()): + start, end = int(cu_q[entry]), int(cu_q[entry + 1]) + for row in range(start, end): + visible = int(kv_lens[entry]) - (end - start) + row - start + 1 + selected = indices[row] + selected = selected[selected >= 0].long() + assert (selected < visible).all() + selected = selected + int(cu_k[entry]) + keys = k.index_select(0, selected).repeat_interleave( + q.shape[1] // k.shape[1], dim=1 + ) + values = v.index_select(0, selected).repeat_interleave( + q.shape[1] // v.shape[1], dim=1 + ) + scores = torch.einsum("hd,nhd->hn", q[row], keys) * scale + output[row] = torch.einsum("hn,nhd->hd", scores.softmax(-1), values) + return output + + +@pytest.mark.parametrize("prefix_lens", [(4, 8), (0, 4), (0, 0)]) +@pytest.mark.parametrize("rank", [0, 1]) +@pytest.mark.parametrize("save_kv_cache", [True, False]) +@pytest.mark.parametrize( + "query_dtype,cache_dtype", + [(torch.float32, torch.float32), (torch.bfloat16, torch.float8_e4m3fn)], +) +def test_qsa_cp_cached_prefix_matches_full_context_attention( + monkeypatch, prefix_lens, rank, save_kv_cache, query_dtype, cache_dtype +): + # Unequal request lengths exercise both zigzag halves and uneven rank sizes. + extend_lens = [8, 5] + seq_lens = [prefix + length for prefix, length in zip(prefix_lens, extend_lens)] + monkeypatch.setattr( + cp_base, "get_parallel", lambda: SimpleNamespace(attn_cp_rank=rank) + ) + monkeypatch.setattr(zigzag, "get_device", lambda: SimpleNamespace(device="cpu")) + strategy = zigzag.ZigzagCPStrategy(cp_size=2) + meta = strategy.build_metadata(sum(extend_lens), seq_lens, extend_lens) + meta.per_rank_logical_token = meta.per_rank_actual_token + meta.per_rank_actual_token = [8, 8] + meta.max_rank_len = [8, 8] + monkeypatch.setattr(qwen_sparse_attn_backend, "get_cp_strategy", lambda: strategy) + + generator = torch.Generator().manual_seed(42) + keys = [ + torch.randn(length, 1, 4, generator=generator).to(query_dtype) + for length in seq_lens + ] + values = [ + torch.randn(length, 1, 4, generator=generator).to(query_dtype) + for length in seq_lens + ] + # Cached prefixes have already incurred the KV pool's storage rounding. + for key, value, prefix in zip(keys, values, prefix_lens): + key[:prefix] = key[:prefix].to(cache_dtype).to(query_dtype) + value[:prefix] = value[:prefix].to(cache_dtype).to(query_dtype) + queries = torch.randn(sum(extend_lens), 2, 4, generator=generator).to(query_dtype) + new_k = torch.cat([key[prefix:] for key, prefix in zip(keys, prefix_lens)]) + new_v = torch.cat([value[prefix:] for value, prefix in zip(values, prefix_lens)]) + + # Request IDs and physical token slots deliberately differ from packed order. + req_indices = [4, 1] + req_to_token = torch.zeros(5, max(seq_lens), dtype=torch.int32) + slots = torch.randperm(80, generator=generator)[: sum(seq_lens)] + 1 + k_cache = torch.full((81, 1, 4), float("nan"), dtype=cache_dtype) + v_cache = torch.full_like(k_cache, float("nan")) + out_cache_locs = [] + offset = 0 + for i, (length, prefix) in enumerate(zip(seq_lens, prefix_lens)): + request_slots = slots[offset : offset + length] + offset += length + req_to_token[req_indices[i], :length] = request_slots.to(torch.int32) + k_cache[request_slots[:prefix]] = keys[i][:prefix].to(cache_dtype) + v_cache[request_slots[:prefix]] = values[i][:prefix].to(cache_dtype) + out_cache_locs.append(request_slots[prefix:]) + out_cache_loc = torch.cat(out_cache_locs) + initial_k, initial_v = k_cache.clone(), v_cache.clone() + + def store_kv(layer, locations, k, v): + k_cache[locations.long()] = k.to(cache_dtype) + v_cache[locations.long()] = v.to(cache_dtype) + + pool = SimpleNamespace( + set_kv_buffer=Mock(side_effect=store_kv), + get_key_buffer=Mock(return_value=k_cache), + get_value_buffer=Mock(return_value=v_cache), + ) + if not any(prefix_lens): + pool.get_key_buffer.side_effect = AssertionError( + "zero prefix needs no cache read" + ) + pool.get_value_buffer.side_effect = AssertionError( + "zero prefix needs no cache read" + ) + backend = SimpleNamespace( + token_to_kv_pool=pool, + req_to_token_pool=SimpleNamespace(req_to_token=req_to_token), + _pad_extend_output=QwenSparseAttnBackend._pad_extend_output, + ) + batch = SimpleNamespace( + attn_cp_metadata=meta, + seq_lens_cpu=seq_lens, + extend_seq_lens_cpu=extend_lens, + extend_prefix_lens_cpu=list(prefix_lens), + req_pool_indices=torch.tensor(req_indices), + out_cache_loc=out_cache_loc, + ) + local_q = strategy.shard_hidden_states(queries, batch) + local_k = strategy.shard_hidden_states(new_k, batch) + local_v = strategy.shard_hidden_states(new_v, batch) + num_local = meta.total_q_prev_tokens + meta.total_q_next_tokens + local_rows = strategy.shard_hidden_states(torch.arange(sum(extend_lens)), batch)[ + :num_local + ] + logical_positions = torch.cat( + [torch.arange(prefix, length) for prefix, length in zip(prefix_lens, seq_lens)] + ) + local_positions = logical_positions.index_select(0, local_rows) + + def gather_new_kv(local_kv, forward_batch): + assert forward_batch is batch + torch.testing.assert_close( + local_kv, torch.cat([local_k.flatten(1), local_v.flatten(1)], dim=-1) + ) + return torch.cat([new_k.flatten(1), new_v.flatten(1)], dim=-1) + + monkeypatch.setattr( + qwen_sparse_attn_backend, "cp_materialize_global_token_order", gather_new_kv + ) + + def sparse_chunk_attention(q, k, v, indices, cu_q, cu_k, kv_lens, scale): + # The last local query in each block identifies its absolute causal end. + expected_ends = local_positions[cu_q[1:].long() - 1] + 1 + torch.testing.assert_close(kv_lens.long(), expected_ends) + assert k.dtype == v.dtype == q.dtype + return _cpu_sparse_chunk_attention(q, k, v, indices, cu_q, cu_k, kv_lens, scale) + + monkeypatch.setattr( + qwen_sparse_attn_backend, + "sparse_gqa_fwd_interface_triton_ck", + sparse_chunk_attention, + ) + layer = SimpleNamespace( + layer_id=0, + tp_q_head_num=2, + tp_k_head_num=1, + tp_v_head_num=1, + head_dim=4, + qk_head_dim=4, + v_head_dim=4, + scaling=0.5, + ) + # The indexer supplies causal logical positions and marks other slots invalid. + candidates = torch.arange(max(seq_lens), dtype=torch.int32).expand(num_local, -1) + topk = torch.where(candidates <= local_positions[:, None], candidates, -1) + actual = QwenSparseAttnBackend._forward_extend_cp( + backend, local_q, local_k, local_v, layer, batch, topk, save_kv_cache + ) + + expected = [] + row = 0 + for i, (length, prefix) in enumerate(zip(extend_lens, prefix_lens)): + for position in range(prefix, prefix + length): + visible_keys = keys[i][: position + 1, 0] + visible_values = values[i][: position + 1, 0] + scores = queries[row] @ visible_keys.T * layer.scaling + expected.append(scores.softmax(-1) @ visible_values) + row += 1 + expected = torch.stack(expected).index_select(0, local_rows).flatten(1) + torch.testing.assert_close(actual[:num_local], expected) + assert actual.shape == (8, 8) + torch.testing.assert_close(actual[num_local:], torch.zeros_like(actual[num_local:])) + if save_kv_cache: + pool.set_kv_buffer.assert_called_once() + initial_k[out_cache_loc] = new_k.to(cache_dtype) + initial_v[out_cache_loc] = new_v.to(cache_dtype) + else: + pool.set_kv_buffer.assert_not_called() + torch.testing.assert_close(k_cache.float(), initial_k.float(), equal_nan=True) + torch.testing.assert_close(v_cache.float(), initial_v.float(), equal_nan=True) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) From 8453e4ca76220784dbe2a26c4c1be9373d794f5b Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Tue, 29 Sep 2026 23:05:50 -0700 Subject: [PATCH 6/6] Keep QSA CP support independent of zigzag allocation changes --- python/sglang/srt/layers/cp/zigzag.py | 25 +-- .../unit/cp/test_zigzag_collectives_unit.py | 173 ------------------ 2 files changed, 8 insertions(+), 190 deletions(-) delete mode 100644 test/registered/unit/cp/test_zigzag_collectives_unit.py diff --git a/python/sglang/srt/layers/cp/zigzag.py b/python/sglang/srt/layers/cp/zigzag.py index 748635bb0525..877b55e23bff 100644 --- a/python/sglang/srt/layers/cp/zigzag.py +++ b/python/sglang/srt/layers/cp/zigzag.py @@ -39,7 +39,6 @@ import torch.nn.functional as F from sglang.srt.distributed.device_communicators.pynccl_allocator import ( - is_symmetric_memory_enabled, use_symmetric_memory, ) from sglang.srt.layers.cp.base import ( @@ -469,25 +468,17 @@ def _all_gather_reorganized(self, x: torch.Tensor, forward_batch): assert x.shape[0] >= local_len x = x[:local_len] pad_size = max_len - x.shape[0] + if pad_size > 0: + padding = [0, 0] * (x.ndim - 1) + [0, pad_size] + x = F.pad(x, padding, mode="constant", value=0) + group = get_parallel().attn_cp_group - symmetric = ( - x.is_cuda and is_symmetric_memory_enabled() and is_allocation_symmetric() + ctx = ( + use_symmetric_memory(group, disabled=not is_allocation_symmetric()) + if x.is_cuda + else nullcontext() ) - ctx = use_symmetric_memory(group) if symmetric else nullcontext() with ctx: - if symmetric: - # Allocate the same send/receive shapes on every rank, including - # ranks without padding. Both buffers must belong to the pool; - # contiguous() can return a producer's existing allocation. - local = x.new_empty((max_len, *x.shape[1:])) - local[:local_len].copy_(x) - local[local_len:].zero_() - x = local - elif pad_size > 0: - padding = [0, 0] * (x.ndim - 1) + [0, pad_size] - x = F.pad(x, padding, mode="constant", value=0) - else: - x = x.contiguous() gathered = torch.empty( max_len * self.cp_size, *x.shape[1:], diff --git a/test/registered/unit/cp/test_zigzag_collectives_unit.py b/test/registered/unit/cp/test_zigzag_collectives_unit.py deleted file mode 100644 index 73ac39747aa8..000000000000 --- a/test/registered/unit/cp/test_zigzag_collectives_unit.py +++ /dev/null @@ -1,173 +0,0 @@ -"""CPU layout and allocation checks for the shared zigzag collectives.""" - -import sys -from contextlib import contextmanager -from types import SimpleNamespace -from unittest.mock import PropertyMock, patch - -import pytest -import torch - -from sglang.srt.layers.cp import base, padding, zigzag -from sglang.srt.layers.cp.padding import pad_logical_token_to_physical -from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy -from sglang.test.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=5, suite="base-a-test-cpu") - - -@contextmanager -def _parallel(cp_size, rank, group=None): - parallel = SimpleNamespace( - attn_cp_rank=rank, attn_cp_size=cp_size, attn_cp_group=group - ) - with ( - patch.object(base, "get_parallel", return_value=parallel), - patch.object(zigzag, "get_parallel", return_value=parallel), - patch.object(padding, "get_parallel", return_value=parallel), - ): - yield - - -def _make_shards(payload, cp_size, lengths, physical_padding): - strategy = ZigzagCPStrategy(cp_size) - shards, batches = [], [] - with ( - patch.object(zigzag, "get_device", return_value=SimpleNamespace(device="cpu")), - patch.object(base, "_STRATEGY", strategy), - ): - for rank in range(cp_size): - with _parallel(cp_size, rank): - metadata = strategy.build_metadata(sum(lengths), lengths) - if physical_padding: - pad_logical_token_to_physical(metadata) - batch = SimpleNamespace(attn_cp_metadata=metadata) - shards.append(strategy.shard_hidden_states(payload, batch)) - batches.append(batch) - return strategy, shards, batches - - -def _noncontiguous(tensor): - backing = tensor.new_empty((*tensor.shape[:-1], tensor.shape[-1] * 2)) - view = backing[..., ::2] - view.copy_(tensor) - assert not view.is_contiguous() - return view - - -def _gathered_shards(shards, metadata): - lengths = metadata.per_rank_logical_token or metadata.per_rank_actual_token - max_len = max(lengths) - padded = shards[0].new_zeros((len(shards), max_len, *shards[0].shape[1:])) - for rank, length in enumerate(lengths): - padded[rank, :length].copy_(shards[rank][:length]) - return padded.flatten(0, 1) - - -@pytest.mark.parametrize("method", ["gather_hidden_states", "gather_kv_cache"]) -@pytest.mark.parametrize("physical_padding", [False, True]) -@pytest.mark.parametrize( - "cp_size,lengths,shape", - [ - (8, [512], (3,)), - (8, [100, 33], (2, 3)), - (4, [257, 16, 256], (3,)), - (2, [7, 9, 24], (2, 3)), - ], -) -def test_gather_restores_global_token_order( - method, physical_padding, cp_size, lengths, shape -): - num_tokens = sum(lengths) - payload = torch.arange(num_tokens * torch.Size(shape).numel()).reshape( - num_tokens, *shape - ) - strategy, shards, batches = _make_shards( - payload, cp_size, lengths, physical_padding - ) - expected_gathered = _gathered_shards(shards, batches[0].attn_cp_metadata) - - for rank, (shard, batch) in enumerate(zip(shards, batches)): - # Nonzero physical padding must not leak into reconstructed token order. - logical = batch.attn_cp_metadata.per_rank_logical_token - if logical is not None: - shard[logical[rank] :].fill_(-1) - local = _noncontiguous(shard) - - def gather(output, input): - assert input.is_contiguous() - torch.testing.assert_close(input, expected_gathered.chunk(cp_size)[rank]) - output.copy_(expected_gathered) - - group = SimpleNamespace(all_gather_into_tensor=gather) - with _parallel(cp_size, rank, group): - actual = getattr(strategy, method)(local, batch) - torch.testing.assert_close(actual, payload, rtol=0, atol=0) - - -@pytest.mark.parametrize("physical_padding", [False, True]) -@pytest.mark.parametrize("strided", [False, True]) -def test_symmetric_gather_allocates_equal_send_and_receive_buffers_in_pool( - physical_padding, strided -): - cp_size, lengths = 4, [17, 19] - payload = torch.arange(sum(lengths) * 6).reshape(-1, 2, 3) - strategy, shards, batches = _make_shards( - payload, cp_size, lengths, physical_padding - ) - gathered = _gathered_shards(shards, batches[0].attn_cp_metadata) - send_shape = tuple(gathered.chunk(cp_size)[0].shape) - original_new_empty, original_empty = torch.Tensor.new_empty, torch.empty - in_pool = False - allocations = {} - - @contextmanager - def pool(group): - nonlocal in_pool - in_pool = True - try: - yield - finally: - in_pool = False - - def record_new_empty(tensor, *args, **kwargs): - result = original_new_empty(tensor, *args, **kwargs) - allocations[result.data_ptr()] = (in_pool, tuple(result.shape)) - return result - - def record_empty(*args, **kwargs): - result = original_empty(*args, **kwargs) - allocations[result.data_ptr()] = (in_pool, tuple(result.shape)) - return result - - for rank, (shard, batch) in enumerate(zip(shards, batches)): - local = _noncontiguous(shard) if strided else shard - allocations.clear() - - def gather(output, input): - assert input.is_contiguous() - assert not in_pool - assert input.data_ptr() != local.data_ptr() - assert allocations[input.data_ptr()] == (True, send_shape) - assert allocations[output.data_ptr()] == (True, tuple(gathered.shape)) - assert torch.equal(input, gathered.chunk(cp_size)[rank]) - output.copy_(gathered) - - group = SimpleNamespace(all_gather_into_tensor=gather) - with ( - _parallel(cp_size, rank, group), - patch.object( - torch.Tensor, "is_cuda", new_callable=PropertyMock, return_value=True - ), - patch.object(zigzag, "is_symmetric_memory_enabled", return_value=True), - patch.object(zigzag, "is_allocation_symmetric", return_value=True), - patch.object(zigzag, "use_symmetric_memory", side_effect=pool), - patch.object(torch.Tensor, "new_empty", record_new_empty), - patch.object(torch, "empty", record_empty), - ): - actual = strategy.gather_kv_cache(local, batch) - torch.testing.assert_close(actual, payload, rtol=0, atol=0) - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v"]))