diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index cbd14f618660..f60c89a5b6cf 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -25,6 +25,10 @@ get_attention_tp_size, is_dp_attention_enabled, ) +from sglang.srt.layers.utils.cp_utils import ( + cp_allgather_and_save_kv_cache, + cp_attn_forward_extend, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import is_gfx95_supported @@ -148,6 +152,8 @@ def __init__( ) self.kv_cache_dtype = model_runner.kv_cache_dtype + self.attn_cp_size = getattr(model_runner, "attn_cp_size", 1) + self.req_to_token = model_runner.req_to_token_pool.req_to_token self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA @@ -2328,9 +2334,15 @@ def forward_extend( k_descale = layer.k_scale if layer.k_scale is not None else self.k_scale v_descale = layer.v_scale if layer.v_scale is not None else self.k_scale + is_cp_mode = ( + forward_batch.forward_mode.is_context_parallel_extend() + and forward_batch.attn_cp_metadata is not None + and self.attn_cp_size > 1 + ) + if k is not None: assert v is not None - if save_kv_cache: + if save_kv_cache and not is_cp_mode: # Only use SWA-specific kv cache write (reshape_and_cache_flash) when # both unified attention and sliding window kv pool are active. # Non-SWA models (e.g. Qwen3-VL) enabled via SGLANG_USE_AITER_UNIFIED_ATTN @@ -2370,6 +2382,10 @@ def forward_extend( forward_batch.token_to_kv_pool.set_kv_buffer( layer, cache_loc, k, v, k_descale, v_descale ) + if is_cp_mode: + cp_allgather_and_save_kv_cache( + forward_batch, layer, k, v, self.attn_cp_size + ) if self.use_mla: max_q_len = self.forward_metadata.max_q_len @@ -2729,26 +2745,65 @@ def forward_extend( if self.forward_metadata.swa_page_table is not None: page_table = self.forward_metadata.swa_page_table - o = mha_batch_prefill_func( - q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), - k_cache, - v_cache, - self.qo_indptr[:bs0], - self.forward_metadata.kv_indptr[:bs0], - page_table, - self.forward_metadata.max_q_len, - self.forward_metadata.max_kv_len, - causal=True, - logits_soft_cap=self.logits_soft_cap, - alibi_slopes=None, - return_lse=False, - return_attn_probs=False, - window_size=window_size, - sink_ptr=sinks, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, - ) + if is_cp_mode: + cp_meta = forward_batch.attn_cp_metadata + + def _aiter_cp_attn( + q_chunk, cu_seqlens_q_cp, cache_seqlens_cp, max_seqlen_q_cp + ): + kv_len_cp = int(cache_seqlens_cp.item()) + kv_indptr_cp = torch.tensor( + [0, kv_len_cp], device=self.device, dtype=torch.int32 + ) + page_table_cp = page_table[:kv_len_cp] + return mha_batch_prefill_func( + q_chunk, + k_cache, + v_cache, + cu_seqlens_q_cp, + kv_indptr_cp, + page_table_cp, + max_seqlen_q_cp, + kv_len_cp, + causal=True, + logits_soft_cap=self.logits_soft_cap, + alibi_slopes=None, + return_lse=False, + return_attn_probs=False, + window_size=window_size, + sink_ptr=sinks, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + + o = cp_attn_forward_extend( + forward_batch, + q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + self.device, + _aiter_cp_attn, + ) + else: + o = mha_batch_prefill_func( + q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k_cache, + v_cache, + self.qo_indptr[:bs0], + self.forward_metadata.kv_indptr[:bs0], + page_table, + self.forward_metadata.max_q_len, + self.forward_metadata.max_kv_len, + causal=True, + logits_soft_cap=self.logits_soft_cap, + alibi_slopes=None, + return_lse=False, + return_attn_probs=False, + window_size=window_size, + sink_ptr=sinks, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) # The fp8bf16 aiter prefill kernel returns bf16 even when the # model computes in fp16. Cast back so the attention output keeps diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py index 23d8d58feca2..d87e685a7e26 100644 --- a/python/sglang/srt/layers/attention/fla/chunk.py +++ b/python/sglang/srt/layers/attention/fla/chunk.py @@ -35,6 +35,7 @@ def chunk_gated_delta_rule_fwd( initial_state_indices: torch.Tensor, cu_seqlens: Optional[torch.LongTensor] = None, chunk_indices: torch.LongTensor | None = None, + M_out: Optional[torch.Tensor] = None, ): g = chunk_local_cumsum( g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices @@ -59,6 +60,7 @@ def chunk_gated_delta_rule_fwd( initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + M_out=M_out, ) o = chunk_fwd_o( q=q, @@ -92,6 +94,7 @@ def forward( initial_state_indices: torch.Tensor, cu_seqlens: Optional[torch.LongTensor] = None, use_qk_l2norm_in_kernel: bool = False, + M_out: Optional[torch.Tensor] = None, ): q_orig = q k_orig = k @@ -116,6 +119,7 @@ def forward( initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + M_out=M_out, ) return o.to(q.dtype), h @@ -133,6 +137,7 @@ def chunk_gated_delta_rule( cu_seqlens: Optional[torch.LongTensor] = None, head_first: bool = False, use_qk_l2norm_in_kernel: bool = False, + M_out: Optional[torch.Tensor] = None, ): r""" Args: @@ -247,6 +252,7 @@ def chunk_gated_delta_rule( initial_state_indices, cu_seqlens, use_qk_l2norm_in_kernel, + M_out, ) if head_first: o = rearrange(o, "b t h ... -> b h t ...") diff --git a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py index 0c7f80f42f67..f3177c35c14e 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py +++ b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py @@ -271,6 +271,97 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) +def compute_M_total_pytorch_ref( + k: torch.Tensor, # [B=1, T_total, Hg, K] bf16 + w: torch.Tensor, # [B=1, T_total, H, K] bf16 + g_cumsum: torch.Tensor, # [B=1, T_total, H] fp32 (chunk-local cumsum, same as fwd_h's g input) + cu_seqlens: torch.LongTensor, + M_out: torch.Tensor, # [N, H, K, K] fp32 — pre-allocated output, filled in place + chunk_size: int = CHUNK_SIZE, +) -> None: + """Compute the segment-end ``M_total`` matrix used by the (b, M)-merge + Context Parallel scheme for GDN linear-attention. + + For each (sequence n, head h), ``M[n, h]`` is defined by the affine + recurrence over the segment's NT chunks:: + + b_h_after = b_h_initial @ M[n, h] + b_seg[n, h] + + Per chunk t (NT chunks total, applied in causal order):: + + N_chunk[h, k_in, k_out] = γ^C[h] · I[k_in, k_out] + - Σ_t W[t, h, k_in] · K_arrow[t, h, k_out] + + where ``K_arrow[t, h, k] = exp(g_last[h] - g_chunk[t, h]) · K[t, h, k]``; + the γ-to-end factor mirrors ``fwd_h``'s ``b_v`` scaling above. Then + ``M_seg = N_chunk_1 @ N_chunk_2 @ ... @ N_chunk_NT``. + + This is a PyTorch reference and the source of truth for the math; a + fused Triton kernel can be substituted later for perf without changing + the semantics. + """ + assert k.shape[0] == 1, "varlen path expects B=1" + assert M_out.dtype == torch.float32, "M_out must be fp32" + assert g_cumsum is not None, "compute_M_total_pytorch_ref requires g (USE_G path)" + + _, T_total, Hg, K = k.shape + H = w.shape[2] + group_size = H // Hg + N = len(cu_seqlens) - 1 + assert M_out.shape == (N, H, K, K) + + device = k.device + eye_K = torch.eye(K, dtype=torch.float32, device=device) + + cu = cu_seqlens.tolist() if cu_seqlens.is_cuda else cu_seqlens.cpu().tolist() + + for n in range(N): + bos, eos = int(cu[n]), int(cu[n + 1]) + T_seg = eos - bos + NT = (T_seg + chunk_size - 1) // chunk_size + + # Initialize M = identity for each head: [H, K, K] + M = eye_K.unsqueeze(0).expand(H, K, K).contiguous() + + for i_t in range(NT): + t0 = i_t * chunk_size + t1 = min(t0 + chunk_size, T_seg) + BT_actual = t1 - t0 + + # Slice this chunk + k_chunk = k[0, bos + t0 : bos + t1] # [BT, Hg, K] bf16 + w_chunk = w[0, bos + t0 : bos + t1] # [BT, H, K] bf16 + g_chunk = g_cumsum[0, bos + t0 : bos + t1] # [BT, H] fp32 (chunk-local cumsum) + + # Expand k from Hg head-groups to H full heads (GVA: H/Hg = group_size) + # k_chunk[t, hg, k] -> k_expanded[t, hg*group_size + g, k] + k_expanded = ( + k_chunk[:, :, None, :] + .expand(BT_actual, Hg, group_size, K) + .reshape(BT_actual, H, K) + .float() + ) # [BT, H, K] fp32 + + # γ-to-end factor for each (token, head): exp(g_last - g_t) + g_last = g_chunk[-1] # [H] fp32 + gamma_to_end = torch.exp(g_last[None, :] - g_chunk) # [BT, H] + gamma_C = torch.exp(g_last) # [H] + + # K_arrow[t, h, k] = gamma_to_end[t, h] * k_expanded[t, h, k] + K_arrow = gamma_to_end[:, :, None] * k_expanded # [BT, H, K] + + # WK[h, k_in, k_out] = Σ_t w_chunk[t, h, k_in] · K_arrow[t, h, k_out] + WK = torch.einsum("thi,thj->hij", w_chunk.float(), K_arrow) # [H, K, K] + + # N_chunk[h] = γ^C[h] · I - WK[h] + N_chunk = gamma_C[:, None, None] * eye_K[None, :, :] - WK # [H, K, K] + + # Update M ← M @ N_chunk per head + M = torch.bmm(M, N_chunk) # [H, K, K] + + M_out[n].copy_(M) + + def chunk_gated_delta_rule_fwd_h( k: torch.Tensor, w: torch.Tensor, @@ -282,7 +373,8 @@ def chunk_gated_delta_rule_fwd_h( save_new_value: bool = True, cu_seqlens: Optional[torch.LongTensor] = None, chunk_indices: Optional[torch.LongTensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + M_out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: B, T, Hg, K, V = *k.shape, u.shape[-1] H = u.shape[-2] BT = CHUNK_SIZE @@ -335,4 +427,14 @@ def grid(meta): num_warps=4, num_stages=2, ) + + # Optionally fill the segment-end M matrix used by the GDN Context + # Parallel (b, M)-merge path. Default M_out=None skips this entirely + # (zero numeric or perf impact for non-CP callers). + if M_out is not None: + compute_M_total_pytorch_ref( + k=k, w=w, g_cumsum=g, cu_seqlens=cu_seqlens, + M_out=M_out, chunk_size=BT, + ) + return h, v_new diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 1f463430e472..17e54f7cbd45 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -1,10 +1,17 @@ from typing import Optional, Tuple, Union import torch +import torch.nn.functional as F +from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel +from sglang.srt.layers.dp_attention import ( + attn_cp_all_gather_into_tensor, + get_attention_cp_group, + get_attention_cp_rank, +) from sglang.srt.layers.attention.linear.utils import ( LinearAttnKernelBackend, get_linear_attn_decode_backend, @@ -259,6 +266,13 @@ def __init__(self, model_runner: ModelRunner): self.req_to_token_pool.size, dtype=torch.int32, device=model_runner.device ) + # CP attributes (mirror FA backend / aiter backend pattern). + # Used by _forward_extend_cp dispatch and conv1d D1a boundary exchange. + self.attn_cp_size = getattr(model_runner, "attn_cp_size", 1) + self.attn_cp_rank = ( + get_attention_cp_rank() if self.attn_cp_size > 1 else 0 + ) + def init_forward_metadata(self, forward_batch: ForwardBatch): super().init_forward_metadata(forward_batch) if self.forward_metadata.has_mamba_track_mask: @@ -358,6 +372,32 @@ def forward_extend( assert isinstance(mixed_qkv, torch.Tensor) seq_len = mixed_qkv.shape[0] + # CP-aware prefill path. Activated when the model layer set + # forward_batch.attn_cp_metadata via cp_split_and_rebuild_data, so + # mixed_qkv arrives ALREADY local-T per rank (= rank's 2 owned + # zigzag segments concatenated). The full-sequence equal-segment + # algorithm requires the FULL kv_len divisible by 2*cp_size with + # at least one CHUNK_SIZE=64 chunk per segment; this is gated + # upstream via can_cp_split() inside the model. Warmup / short + # prompts have attn_cp_metadata == None and fall through here. + cp_meta = getattr(forward_batch, "attn_cp_metadata", None) + cp_size = self.attn_cp_size + if ( + cp_meta is not None + and cp_size > 1 + and not forward_batch.forward_mode.is_target_verify() + and forward_batch.forward_mode.is_context_parallel_extend() + ): + return self._forward_extend_cp( + layer, + forward_batch, + mixed_qkv, + a, + b, + cp_size, + self.attn_cp_rank, + ) + is_target_verify = forward_batch.forward_mode.is_target_verify() forward_metadata = self.forward_metadata @@ -475,3 +515,277 @@ def forward_extend( ) return core_attn_out + + # CP-aware prefill for the GDN linear-attention path. + # + # Input: local-T mixed_qkv [seg_a_len + seg_b_len, dim] per rank + # (model splits at embed, gathers at exit before lm_head). + # + # Algorithm (zigzag two-pass + (b, M) chain reduce): + # 1. conv1d boundary exchange: all_gather (K-1) tail tokens from + # adjacent segments, run F.conv1d per segment locally. + # 2. Pass 1: run each owned segment with init_state=0, capture the + # affine recurrence maps (b, M) per segment. + # 3. Chain reduce: all_gather (b, M), chain S_init[i+1] = M[i]*S[i] + b[i] + # in fp32 to derive correct initial states. + # 4. Pass 2: rerun owned segments with correct S_init, emit output. + # 5. Writeback conv_state and ssm_state to cache for decode. + def _forward_extend_cp( + self, + layer: RadixLinearAttention, + forward_batch: ForwardBatch, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + cp_size: int, + cp_rank: int, + ): + SEG_PER_RANK = 2 # zigzag invariant: each rank holds {head, tail} pair + cp_meta = forward_batch.attn_cp_metadata + device = mixed_qkv.device + dtype = mixed_qkv.dtype + dim = mixed_qkv.shape[-1] + + # Owned segment lengths from standard cp metadata (split_list is the + # 2*cp_size segment lengths in causal order; remainder distributed + # to the first segments — seg_a / seg_b can differ by 1 in pathological + # cases, hence torch.split with explicit lengths instead of chunk(2)). + num_segs = SEG_PER_RANK * cp_size + seg_a_len = cp_meta.split_list[cp_rank] # = seg_(cp_rank) + seg_b_len = cp_meta.split_list[num_segs - 1 - cp_rank] # = seg_(2cp-1-cp_rank) + local_T = seg_a_len + seg_b_len + assert mixed_qkv.shape[0] == local_T, ( + f"GDN backend CP path expects local-T input (= seg_a + seg_b = " + f"{seg_a_len} + {seg_b_len}), got mixed_qkv.shape[0]={mixed_qkv.shape[0]}" + ) + + # === conv1d boundary exchange === + K = layer.conv_weights.shape[-1] + Kp = K - 1 + def _tail(seg_input): + if seg_input.shape[0] >= Kp: + return seg_input[-Kp:].contiguous() + pad = torch.zeros( + Kp - seg_input.shape[0], dim, dtype=dtype, device=device + ) + return torch.cat([pad, seg_input], dim=0).contiguous() + + seg_a_input = mixed_qkv[:seg_a_len] + seg_b_input = mixed_qkv[seg_a_len:] + own_seg_a_tail = _tail(seg_a_input) + own_seg_b_tail = _tail(seg_b_input) + + # === Boundary exchange via all_gather === + # Pack each rank's segment tails, all_gather, then locally pick the + # correct left-context feeds for each owned segment. + cp_group = get_attention_cp_group() + zeros_tail = torch.zeros(Kp, dim, dtype=dtype, device=device) + + # Pack [seg_a_tail, seg_b_tail] per rank → shape [2, Kp, dim] + local_tails = torch.stack([own_seg_a_tail, own_seg_b_tail], dim=0) + # Gathered shape: [cp_size, 2, Kp, dim] + all_tails = torch.empty( + cp_size, 2, Kp, dim, dtype=dtype, device=device + ) + attn_cp_all_gather_into_tensor(all_tails, local_tails.contiguous()) + + if cp_rank == 0: + seg_a_left = zeros_tail + else: + seg_a_left = all_tails[cp_rank - 1, 0] + + prev_b_idx = 2 * cp_size - 2 - cp_rank + if prev_b_idx < cp_size: + seg_b_left = all_tails[prev_b_idx, 0] + else: + seg_b_left = all_tails[cp_rank + 1, 1] + + # === Per-segment causal conv1d via F.conv1d (depthwise) === + # No kernel flip: causal_conv1d_fn and F.conv1d both use cross-correlation. + conv_w = layer.conv_weights.unsqueeze(1) # [dim, 1, K] + conv_b = layer.bias # [dim] or None + + def _causal_conv1d_local(seg_input, left_context): + """seg_input [T, C], left_context [Kp, C]. Returns [T, C] post-act.""" + x = torch.cat([left_context, seg_input], dim=0) # [T+Kp, C] + # F.conv1d expects (N, C, L) with depthwise weight (C, 1, K). + x_t = x.transpose(0, 1).unsqueeze(0) # [1, C, T+Kp] + y = F.conv1d(x_t, conv_w, bias=conv_b, padding=0, groups=conv_w.shape[0]) + # y shape: [1, C, T+Kp - K + 1] = [1, C, T] + y = y.squeeze(0).transpose(0, 1) # [T, C] + # Activation (mamba/qwen3.5 conv uses silu/swish; both equivalent) + if layer.activation in ("silu", "swish"): + y = F.silu(y) + elif layer.activation in (None, "identity"): + pass + else: + raise NotImplementedError( + f"GDN CP conv1d activation {layer.activation!r} not handled" + ) + return y + + seg_a_conv = _causal_conv1d_local(seg_a_input, seg_a_left) # [seg_a_len, dim] + seg_b_conv = _causal_conv1d_local(seg_b_input, seg_b_left) # [seg_b_len, dim] + mixed_qkv = torch.cat([seg_a_conv, seg_b_conv], dim=0) # [local_T, dim] + + # === conv_state writeback === + # Full-sequence conv_state = last Kp INPUT tokens of seg_(2cp-1), + # which is rank 0's seg_b. All-gather and use rank 0's slot. + mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) + conv_states = mamba_cache_params.conv[0] + ssm_states = mamba_cache_params.temporal + cache_indices = self.forward_metadata.mamba_cache_indices + + all_seg_b_tails = torch.empty( + cp_size, Kp, dim, dtype=dtype, device=device + ) + attn_cp_all_gather_into_tensor(all_seg_b_tails, own_seg_b_tail.contiguous()) + conv_states[cache_indices[0]] = all_seg_b_tails[0].transpose(0, 1).to( + conv_states.dtype + ) + + query, key, value = torch.split( + mixed_qkv, + [layer.q_dim, layer.k_dim, layer.v_dim], + dim=-1, + ) + Hq, Hk, Hv = layer.num_q_heads, layer.num_k_heads, layer.num_v_heads + Kdim, Vdim = layer.head_k_dim, layer.head_v_dim + query = query.view(1, local_T, Hq, layer.head_q_dim) + key = key.view(1, local_T, Hk, Kdim) + value = value.view(1, local_T, Hv, Vdim) + g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) + + state_idx = torch.tensor([0], dtype=torch.int64, device=device) + def _slice_local(t, start, length): + return t.narrow(1, start, length).contiguous() + + owned_starts = [0, seg_a_len] + owned_lens = [seg_a_len, seg_b_len] + owned_seg_idx = [cp_rank, num_segs - 1 - cp_rank] + + # === Pass 1: capture (b_seg, M_seg) per owned segment === + b_pair = torch.zeros( + SEG_PER_RANK, Hv, Vdim, Kdim, dtype=torch.float32, device=device + ) + M_pair = torch.zeros( + SEG_PER_RANK, Hv, Kdim, Kdim, dtype=torch.float32, device=device + ) + for slot in range(SEG_PER_RANK): + seg_len = owned_lens[slot] + start = owned_starts[slot] + cu_seg = torch.tensor([0, seg_len], dtype=torch.int32, device=device) + q_seg = _slice_local(query, start, seg_len) + k_seg = _slice_local(key, start, seg_len) + v_seg = _slice_local(value, start, seg_len) + g_seg = _slice_local(g, start, seg_len) + beta_seg = _slice_local(beta, start, seg_len) + state = torch.zeros( + 1, Hv, Vdim, Kdim, dtype=torch.float32, device=device + ) + M_buf = torch.zeros( + 1, Hv, Kdim, Kdim, dtype=torch.float32, device=device + ) + chunk_gated_delta_rule( + q=q_seg, + k=k_seg, + v=v_seg, + g=g_seg, + beta=beta_seg, + initial_state=state, + initial_state_indices=state_idx, + cu_seqlens=cu_seg, + use_qk_l2norm_in_kernel=True, + M_out=M_buf, + ) + b_pair[slot] = state[0] + M_pair[slot] = M_buf[0] + + # === All-gather (b, M) across cp_group === + b_gathered = torch.empty( + num_segs, Hv, Vdim, Kdim, dtype=torch.float32, device=device + ) + M_gathered = torch.empty( + num_segs, Hv, Kdim, Kdim, dtype=torch.float32, device=device + ) + attn_cp_all_gather_into_tensor(b_gathered, b_pair.contiguous()) + attn_cp_all_gather_into_tensor(M_gathered, M_pair.contiguous()) + + # Reorder from rank-major to causal segment order. + perm = [] + for i in range(num_segs): + owner = min(i, num_segs - 1 - i) + slot = 0 if i == owner else 1 + perm.append(owner * SEG_PER_RANK + slot) + perm_t = torch.tensor(perm, dtype=torch.long, device=device) + b_causal = b_gathered[perm_t] + M_causal = M_gathered[perm_t] + + # === fp32 deterministic chain reduce === + S_init = [None] * num_segs + S_init[0] = torch.zeros( + 1, Hv, Vdim, Kdim, dtype=torch.float32, device=device + ) + for i in range(1, num_segs): + S_init[i] = ( + torch.einsum("nhvk,hkj->nhvj", S_init[i - 1], M_causal[i - 1]) + + b_causal[i - 1].unsqueeze(0) + ) + + # === Pass 2: rerun owned segments with correct S_init === + # Output buffer: 2 contiguous segments concatenated in (prev,next) layout + # to match local-T input ordering. + seg_a_out = torch.empty( + seg_a_len, Hv, Vdim, dtype=value.dtype, device=device + ) + seg_b_out = torch.empty( + seg_b_len, Hv, Vdim, dtype=value.dtype, device=device + ) + out_slots = [seg_a_out, seg_b_out] + captured_final_state = None + last_seg_idx = num_segs - 1 + for slot in range(SEG_PER_RANK): + seg_len = owned_lens[slot] + start = owned_starts[slot] + seg_idx = owned_seg_idx[slot] + cu_seg = torch.tensor([0, seg_len], dtype=torch.int32, device=device) + q_seg = _slice_local(query, start, seg_len) + k_seg = _slice_local(key, start, seg_len) + v_seg = _slice_local(value, start, seg_len) + g_seg = _slice_local(g, start, seg_len) + beta_seg = _slice_local(beta, start, seg_len) + S_init_seg = S_init[seg_idx].clone().contiguous() + o, _, _ = chunk_gated_delta_rule( + q=q_seg, + k=k_seg, + v=v_seg, + g=g_seg, + beta=beta_seg, + initial_state=S_init_seg, + initial_state_indices=state_idx, + cu_seqlens=cu_seg, + use_qk_l2norm_in_kernel=True, + M_out=None, + ) + out_slots[slot].copy_(o[0]) # [seg_len, Hv, Vdim] + if seg_idx == last_seg_idx: + captured_final_state = S_init_seg[0].clone().contiguous() + + # Assemble local-T output matching input layout. + core_attn_out = torch.cat([seg_a_out, seg_b_out], dim=0).reshape( + 1, local_T, Hv, Vdim + ) + + # === ssm_state writeback === + # Rank 0 owns the last causal segment; all-gather its final state. + if captured_final_state is None: + captured_final_state = torch.empty( + Hv, Vdim, Kdim, dtype=torch.float32, device=device + ) + gathered_states = torch.empty( + cp_size, Hv, Vdim, Kdim, dtype=torch.float32, device=device + ) + attn_cp_all_gather_into_tensor(gathered_states, captured_final_state) + ssm_states[cache_indices[0]] = gathered_states[0].to(ssm_states.dtype) + + return core_attn_out diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 2c6fd4da71fe..29b36a1c6080 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -29,8 +29,11 @@ from sglang.srt.distributed import ( get_moe_data_parallel_world_size, get_moe_expert_parallel_world_size, + get_moe_tensor_parallel_world_size, get_pp_group, get_tensor_model_parallel_world_size, + moe_expert_parallel_all_reduce, + moe_tensor_model_parallel_all_reduce, tensor_model_parallel_all_reduce, ) from sglang.srt.distributed.parallel_state import ( @@ -205,6 +208,8 @@ def __init__( ): super().__init__() self.tp_size = get_tensor_model_parallel_world_size() + self.moe_tp_size = get_moe_tensor_parallel_world_size() + self.moe_ep_size = get_moe_expert_parallel_world_size() self.layer_id = layer_id self.alt_stream = alt_stream if self.tp_size > config.num_experts: @@ -466,12 +471,20 @@ def forward( # An out-of-place add would allocate a new tensor outside symm # memory, breaking subsequent symmetric collective operations. final_hidden_states += shared_output - if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + if self.moe_ep_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=False, + use_reduce_scatter=use_reduce_scatter, + should_allreduce_fusion=should_allreduce_fusion, + ): + final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states) + if self.moe_tp_size > 1 and not should_skip_post_experts_all_reduce( is_tp_path=True, use_reduce_scatter=use_reduce_scatter, should_allreduce_fusion=should_allreduce_fusion, ): - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + final_hidden_states = moe_tensor_model_parallel_all_reduce( + final_hidden_states + ) return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 2b3315fbf975..d2c8dc6284d6 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -43,11 +43,25 @@ from sglang.srt.layers.attention.mamba.mamba import mamba_v2_sharded_weight_loader from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes from sglang.srt.layers.dp_attention import ( + get_attention_cp_rank, + get_attention_cp_size, get_attention_tp_rank, get_attention_tp_size, is_dp_attention_enabled, ) +from sglang.srt.layers.utils.cp_utils import ( + can_cp_split, + cp_all_gather_rerange_output, + cp_split_and_rebuild_data, + cp_split_and_rebuild_position, + is_prefill_context_parallel_enabled, + prepare_context_parallel_metadata, +) +from sglang.srt.distributed.parallel_state import ( + get_moe_tensor_parallel_world_size, +) + # Layers - Others from sglang.srt.layers.layernorm import GemmaRMSNorm @@ -1059,6 +1073,17 @@ def get_layer(idx: int, prefix: str): self.layers_to_capture = [] + self.attn_cp_size = get_attention_cp_size() + self.attn_cp_rank = get_attention_cp_rank() + + if self.attn_cp_size > 1: + moe_tp = get_moe_tensor_parallel_world_size() + assert moe_tp == 1, ( + f"Qwen3.5 + prefill CP requires moe_tp_size == 1. " + f"Got moe_tp_size={moe_tp}. Set --moe-dp-size and " + f"--ep-size so tp_size / (moe_dp_size * ep_size) == 1." + ) + def get_input_embeddings(self): return self.embed_tokens @@ -1085,6 +1110,29 @@ def forward( pp_proxy_tensors: Optional[PPProxyTensors] = None, input_deepstack_embeds: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, PPProxyTensors]: + # Zigzag CP: prepare metadata for sequence splitting. + if is_prefill_context_parallel_enabled(): + cp_size = self.attn_cp_size + # input_ids is None when this forward is invoked from + # general_mm_embed_routine; fall back to input_embeds in that case. + if input_ids is not None: + seq_len_for_cp = input_ids.shape[0] + elif input_embeds is not None: + seq_len_for_cp = input_embeds.shape[0] + else: + seq_len_for_cp = None + if ( + cp_size > 1 + and seq_len_for_cp is not None + and can_cp_split(seq_len_for_cp, cp_size, forward_batch) + ): + forward_batch.attn_cp_metadata = prepare_context_parallel_metadata( + seq_len_for_cp, + self.attn_cp_rank, + cp_size, + forward_batch.seq_lens_cpu.tolist(), + ) + # Initialize hidden states if self.pp_group.is_first_rank: if input_embeds is None: @@ -1097,6 +1145,16 @@ def forward( hidden_states = pp_proxy_tensors["hidden_states"] residual = pp_proxy_tensors["residual"] + # Split to local-T per rank (zigzag layout). + if ( + is_prefill_context_parallel_enabled() + and forward_batch.forward_mode.is_context_parallel_extend() + and forward_batch.attn_cp_metadata is not None + ): + if self.pp_group.is_first_rank: + hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) + positions = cp_split_and_rebuild_position(forward_batch, positions) + aux_hidden_states = [] # Pass through decoder layers for layer_idx in range(self.start_layer, self.end_layer): @@ -1143,6 +1201,20 @@ def forward( else: hidden_states, _ = self.norm(hidden_states, residual) + # Gather local-T back to full-T before lm_head. + if ( + self.pp_group.is_last_rank + and is_prefill_context_parallel_enabled() + and forward_batch.forward_mode.is_context_parallel_extend() + and forward_batch.attn_cp_metadata is not None + ): + hidden_states = cp_all_gather_rerange_output( + hidden_states, + self.attn_cp_size, + forward_batch, + torch.cuda.current_stream(), + ) + if len(aux_hidden_states) == 0: return hidden_states