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..fee856cb4646 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -2,9 +2,11 @@ import torch +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 from sglang.srt.layers.attention.linear.utils import ( LinearAttnKernelBackend, get_linear_attn_decode_backend, @@ -358,6 +360,32 @@ def forward_extend( assert isinstance(mixed_qkv, torch.Tensor) seq_len = mixed_qkv.shape[0] + # CP-aware prefill path. The model layer sets the private fields + # below only when CP is active; we deliberately do not consume the + # standard forward_batch.attn_cp_metadata so that other downstream + # paths (MoE communicator etc.) keep seeing full-T hidden_states. + # The equal-segment zigzag algorithm requires seq_len divisible by + # 2 * cp_size with at least one CHUNK_SIZE=64 chunk per segment; + # warmup / short prompts fall through to the non-CP path. + cp_meta = getattr(forward_batch, "_gdn_cp_metadata", None) + cp_size = getattr(forward_batch, "_gdn_cp_size", 1) + if ( + cp_meta is not None + and cp_size > 1 + and not forward_batch.forward_mode.is_target_verify() + and seq_len % (2 * cp_size) == 0 + and seq_len >= 2 * cp_size * 64 + ): + return self._forward_extend_cp( + layer, + forward_batch, + mixed_qkv, + a, + b, + cp_size, + forward_batch._gdn_cp_rank, + ) + is_target_verify = forward_batch.forward_mode.is_target_verify() forward_metadata = self.forward_metadata @@ -475,3 +503,256 @@ def forward_extend( ) return core_attn_out + + # CP-aware prefill for the GDN linear-attention path. + # + # mixed_qkv arrives full-T because the model layer does NOT call + # cp_split_and_rebuild_data; instead it sets a private + # `_gdn_cp_metadata` field that only this backend reads. The reasons + # for routing CP through a custom field rather than the standard + # `attn_cp_metadata`: + # + # - aiter full-attn has no CP support, so it must run full-T per + # rank (redundant compute, correct output). + # - Setting `attn_cp_metadata` would activate sglang's standard CP + # wrapper chain on the MoE path (cp_split_and_rebuild_data at + # model entry, moe_cp_all_gather_into_tensor in the layer + # communicator, cp_all_gather_rerange_output before lm_head), + # all of which assume hidden_states have already been split to + # local-T. MoE compute itself is per-token and CP-orthogonal — + # the incompatibility is in the wrapper wiring, not in MoE. + # - This backend therefore receives full-T and is the only place + # that splits internally; everything else keeps working as if + # CP were off. + # + # Algorithm (zigzag two-pass + (b, M) chain reduce): + # + # Pass 1: each rank runs its 2 owned segments with initial_state=0 + # and captures (b_seg, M_seg) — the segment's affine output + # on the recurrent state. + # Reduce: all_gather (b, M); each rank deterministically chains the + # per-segment affine maps in fp32 to derive S_init for every + # segment. + # Pass 2: each rank reruns its 2 owned segments with the correct + # S_init and emits the per-segment outputs. + # Gather: all_gather Pass-2 outputs and reorder to causal full-T. + # + # The decode-after-CP ssm_state writeback is handled at the bottom of + # this method; conv_state writeback happens via causal_conv1d_fn which + # every rank runs over the full sequence with identical inputs. + # + # Constraints (see report.md §1.4): prefix=0 only, no cuda graph, + # batch=1, attn_tp ∈ {1, 2}. + 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 + seq_len = mixed_qkv.shape[0] + device = mixed_qkv.device + + forward_metadata = self.forward_metadata + query_start_loc = forward_metadata.query_start_loc + cache_indices = forward_metadata.mamba_cache_indices + + 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 + has_initial_states = forward_batch.extend_prefix_lens > 0 + + # conv1d on the full sequence — every CP rank computes the same result + # and writes the same conv_state into its local cache slot. + mixed_qkv_t = mixed_qkv.transpose(0, 1) + if forward_metadata.has_mamba_track_mask: + mixed_qkv_to_track = mixed_qkv_t[ + :, forward_metadata.track_conv_indices + ].transpose(0, 1) + conv_states[forward_metadata.conv_states_mask_indices] = mixed_qkv_to_track + mixed_qkv = causal_conv1d_fn( + mixed_qkv_t, + layer.conv_weights, + layer.bias, + activation=layer.activation, + conv_states=conv_states, + has_initial_state=has_initial_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + seq_lens_cpu=forward_batch.extend_seq_lens_cpu, + ).transpose(0, 1)[:seq_len] + + # Split q/k/v and run gating on the full sequence. + query, key, value = torch.split( + mixed_qkv, + [layer.q_dim, layer.k_dim, layer.v_dim], + dim=-1, + ) + Hq = layer.num_q_heads + Hk = layer.num_k_heads + Hv = layer.num_v_heads + Kdim = layer.head_k_dim + Vdim = layer.head_v_dim + query = query.view(1, seq_len, Hq, layer.head_q_dim) + key = key.view(1, seq_len, Hk, Kdim) + value = value.view(1, seq_len, Hv, Vdim) + g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) + + # Zigzag segment plan: 2 * cp_size equal segments. Each rank owns + # the {rank, num_segs - 1 - rank} pair so rank 0 always owns the + # last causal segment (used below for ssm_state writeback). + num_segs = SEG_PER_RANK * cp_size + assert seq_len % num_segs == 0, ( + f"CP requires seq_len divisible by 2*cp_size; " + f"got seq_len={seq_len} num_segs={num_segs}" + ) + seg_len = seq_len // num_segs + cu_seg = torch.tensor([0, seg_len], dtype=torch.int32, device=device) + state_idx = torch.tensor([0], dtype=torch.int64, device=device) + owned = [cp_rank, num_segs - 1 - cp_rank] + + def _slice_seg(t, seg_idx): + s = seg_idx * seg_len + return t.narrow(1, s, seg_len).contiguous() + + # Pass 1: run owned segments with initial_state=0 and capture + # (b_seg, M_seg) — the per-segment affine map on the recurrent state. + 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, seg_idx in enumerate(owned): + q_seg = _slice_seg(query, seg_idx) + k_seg = _slice_seg(key, seg_idx) + v_seg = _slice_seg(value, seg_idx) + g_seg = _slice_seg(g, seg_idx) + beta_seg = _slice_seg(beta, seg_idx) + 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 the CP group so every rank has the full + # per-segment affine maps. + 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 gathered tensors from rank-major (rank 0 head, rank 0 tail, + # rank 1 head, ...) 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] + + # Deterministic fp32 chain reduce: S_init[i+1] = S_init[i] @ M[i] + b[i]. + # Same op order on every rank → bit-identical S_init across ranks. + 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 the correct S_init. The rank + # that owns the last causal segment also captures the kernel's + # in-place final-state writeback, which is the authoritative + # full-sequence ssm_state used below for decode-after-CP. + out_pair = torch.empty( + SEG_PER_RANK, seg_len, Hv, Vdim, dtype=value.dtype, device=device + ) + captured_final_state = None + last_seg_idx = num_segs - 1 + for slot, seg_idx in enumerate(owned): + q_seg = _slice_seg(query, seg_idx) + k_seg = _slice_seg(key, seg_idx) + v_seg = _slice_seg(value, seg_idx) + g_seg = _slice_seg(g, seg_idx) + beta_seg = _slice_seg(beta, seg_idx) + 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_pair[slot] = o[0] # [seg_len, Hv, Vdim] + if seg_idx == last_seg_idx: + captured_final_state = S_init_seg[0].clone().contiguous() + + # All-gather Pass-2 outputs and reorder to causal full-T. + out_gathered = torch.empty( + num_segs, seg_len, Hv, Vdim, dtype=value.dtype, device=device + ) + attn_cp_all_gather_into_tensor(out_gathered, out_pair.contiguous()) + out_causal = out_gathered[perm_t] + core_attn_out = out_causal.reshape(1, num_segs * seg_len, Hv, Vdim) + + # ssm_state writeback for decode-after-CP-prefill. + # + # The scheduler clears CP metadata on the extend→decode transition, + # so decode runs the standard non-CP path which reads + # ssm_states[cache_indices[0]] on each rank. Every CP rank therefore + # needs the full-sequence final state in its local cache slot. + # + # Rank 0 owns the last causal segment (zigzag pairing puts seg_0 and + # seg_{num_segs-1} on rank 0) and has already captured the kernel's + # in-place final-state writeback for that segment. We all-gather and + # use rank 0's slot — using the kernel's actual writeback (rather + # than reconstructing it via the affine chain) keeps the non-CP and + # CP paths bit-identical for decode. + # + # conv_state writeback is already done by causal_conv1d_fn above: + # every rank ran full-T conv1d on identical inputs, so the per-rank + # writes are idempotent. + 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/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 2b3315fbf975..90c3a8358bc5 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -43,11 +43,22 @@ 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, ) +# Context Parallel for the GDN linear-attention path. Only the GDN backend +# consumes this metadata; the standard attn_cp_metadata is intentionally left +# unset (see Qwen3_5ForCausalLM.forward). +from sglang.srt.layers.utils.cp_utils import ( + can_cp_split, + is_prefill_context_parallel_enabled, + prepare_context_parallel_metadata, +) + # Layers - Others from sglang.srt.layers.layernorm import GemmaRMSNorm @@ -1085,6 +1096,42 @@ def forward( pp_proxy_tensors: Optional[PPProxyTensors] = None, input_deepstack_embeds: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, PPProxyTensors]: + # CP wiring for the GDN backend only. + # + # We deliberately do NOT set forward_batch.attn_cp_metadata. The standard + # field activates downstream paths (MoE communicator, KV cache writeback, + # lm_head gather) that assume hidden_states has been split to local-T, + # which is incompatible with the AMD aiter full-attn backend (no CP + # support) used for the non-GDN layers in this hybrid model. + # + # Instead we attach private fields read only by GDNAttnBackend; every + # other layer sees attn_cp_metadata is None and runs full-T redundantly + # per CP rank. The GDN backend internally splits, runs the zigzag + # two-pass + (b, M) chain reduce, and returns full-T output. + if is_prefill_context_parallel_enabled(): + cp_size = get_attention_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._gdn_cp_metadata = prepare_context_parallel_metadata( + seq_len_for_cp, + get_attention_cp_rank(), + cp_size, + forward_batch.seq_lens_cpu.tolist(), + ) + forward_batch._gdn_cp_size = cp_size + forward_batch._gdn_cp_rank = get_attention_cp_rank() + # Initialize hidden states if self.pp_group.is_first_rank: if input_embeds is None: