From 0018f015404ea2c0488243a35b44309644b3de94 Mon Sep 17 00:00:00 2001 From: alexsun07 Date: Sat, 9 May 2026 08:09:53 +0000 Subject: [PATCH 1/3] Add Context Parallel support for Qwen3.5 hybrid GDN+full-attn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements CP for Qwen3.5 hybrid models (GDN linear-attn + full-attn MoE). Production-validated at cp ∈ {2, 4, 8} on AMD MI300X via verbatim-copy semantic test (1072-token Chinese prompt). Architecture: model layer sets a custom forward_batch._gdn_cp_metadata field (NOT the standard attn_cp_metadata) to avoid triggering MoE communicator and other downstream paths that assume CP-split hidden_states. GDN backend slices full-T input internally and runs a per-rank zigzag two-pass + (b,M) chain reduce with all_gather. Decode-after-CP works via broadcasting final ssm_state from rank holding the last causal segment (rank 0 in zigzag) so all CP ranks have correct local cache for the standard non-CP decode path. Patches: - fla/chunk.py + fla/chunk_delta_h.py: add output_M_total option to chunk_gated_delta_rule_fwd_h with PyTorch reference compute_M_total_pytorch_ref. Default-off, zero-impact for non-CP callers. - models/qwen3_5.py: in Qwen3_5MoeForCausalLM.forward, populate forward_batch._gdn_cp_metadata when CP enabled (without calling cp_split_and_rebuild_data so hidden_states stays full-T). - layers/attention/linear/gdn_backend.py: GDNAttnBackend.forward_extend adds _forward_extend_cp branch implementing zigzag two-pass + fp32 chain reduce + all_gather output + ssm_state writeback for decode-after-CP. Known limitation: cp=2 + attn_tp=4 + Qwen3.5-397B-FP8 + AMD aiter triggers an upstream sglang bug (V11 stock baseline also fails this combo). --- .../sglang/srt/layers/attention/fla/chunk.py | 6 + .../srt/layers/attention/fla/chunk_delta_h.py | 101 ++++- .../layers/attention/linear/gdn_backend.py | 346 ++++++++++++++++++ python/sglang/srt/models/qwen3_5.py | 61 +++ 4 files changed, 513 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py index 23d8d58feca2..79d5bd2a00eb 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, # NEW: Option-A CP segment-M out-param ): 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, # NEW ) 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, # NEW: Option-A CP segment-M ): 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, # NEW ) 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, # NEW: Option-A CP segment-M out-param ): r""" Args: @@ -247,6 +252,7 @@ def chunk_gated_delta_rule( initial_state_indices, cu_seqlens, use_qk_l2norm_in_kernel, + M_out, # NEW ) 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..98228114ae4d 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,95 @@ 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: + """PyTorch reference for the segment-end M_total matrix used by Option-A + Context Parallel (zigzag + (b, M) merge) 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] + (γ-to-end factor mirrors fwd_h's b_v scaling at chunk_delta_h.py:185) + + Then M_seg = N_chunk_1 @ N_chunk_2 @ ... @ N_chunk_NT (matrix product). + + NOTE: This is the PoC v0 implementation — pure PyTorch, no fused Triton + kernel. v1 will add a fused Triton kernel for perf parity. The semantics + here are the source of truth that the future Triton kernel must match. + """ + 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 +371,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, # NEW: out-param for Option-A CP +) -> Tuple[torch.Tensor, torch.Tensor]: B, T, Hg, K, V = *k.shape, u.shape[-1] H = u.shape[-2] BT = CHUNK_SIZE @@ -335,4 +425,13 @@ def grid(meta): num_warps=4, num_stages=2, ) + + # Option-A CP path: optionally fill the segment-end M matrix. + # Default M_out=None ⇒ skipped entirely (zero numeric or perf change). + 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..23fce4579004 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -2,9 +2,15 @@ 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, + get_attention_cp_rank, + get_attention_cp_size, +) from sglang.srt.layers.attention.linear.utils import ( LinearAttnKernelBackend, get_linear_attn_decode_backend, @@ -358,6 +364,36 @@ def forward_extend( assert isinstance(mixed_qkv, torch.Tensor) seq_len = mixed_qkv.shape[0] + # === V15 (B3) — CP-aware GDN forward === + # Read custom _gdn_cp_metadata (NOT attn_cp_metadata, which would activate + # other downstream CP code paths in the MoE communicator that we can't + # satisfy without splitting hidden_states). + # Fall through to non-CP path when seq_len doesn't divide evenly (warmup, + # short prompts) — CP code path requires seq_len % (2*cp_size) == 0 for + # the equal-segment PoC algorithm. + 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 # need at least one chunk per segment (CHUNK_SIZE=64) + ): + cp_rank = getattr(forward_batch, "_gdn_cp_rank", get_attention_cp_rank()) + # one-shot print so we can confirm CP path fired in production + if not getattr(self.__class__, "_v15_cp_path_printed", False): + print( + f"[V15 CP FIRE] layer={layer.layer_id} seq_len={seq_len} " + f"cp_size={cp_size} cp_rank={cp_rank} " + f"split_list={cp_meta.split_list}", + flush=True, + ) + self.__class__._v15_cp_path_printed = True + return self._forward_extend_cp( + layer, forward_batch, mixed_qkv, a, b, cp_meta, cp_size, cp_rank + ) + is_target_verify = forward_batch.forward_mode.is_target_verify() forward_metadata = self.forward_metadata @@ -474,4 +510,314 @@ def forward_extend( forward_batch, h, ssm_states, forward_metadata ) + # === V20: dump ssm_state after non-CP prefill (gated by DUMP_SSM=1 env) === + # Captures the kernel's in-place writeback for cp=1 baseline + noise floor. + # One-shot per (layer_id, world_rank) per server boot to avoid filesystem spam. + import os as _os + if _os.getenv("DUMP_SSM") == "1" and forward_batch.forward_mode.is_extend(): + try: + import torch.distributed as _dist + _world_rank = _dist.get_rank() if _dist.is_initialized() else 0 + _seq_len = mixed_qkv.shape[0] if isinstance(mixed_qkv, torch.Tensor) else seq_len + _path = ( + f"/tmp/v20_ssm_cp1_layer{layer.layer_id}_rank{_world_rank}" + f"_T{_seq_len}.pt" + ) + if not _os.path.exists(_path): + _slot = int(cache_indices[0].item()) + torch.save( + ssm_states[_slot].detach().float().cpu(), + _path, + ) + if layer.layer_id == 0: + print( + f"[V20 SSM DUMP cp=1] layer={layer.layer_id} slot={_slot} " + f"path={_path}", + flush=True, + ) + except Exception as _e: + print(f"[V20 SSM DUMP cp=1 ERROR] {_e}", flush=True) + return core_attn_out + + # =========================================================================== + # V15 (B3): CP-aware GDN forward. + # + # Contract: + # - mixed_qkv arrives FULL-T (model file did NOT split). This avoids the + # unbounded OOB chain on aiter_backend (which is not CP-aware). + # - cp_meta is set by Qwen3_5ForCausalLM.forward when CP is active. + # - We do real CP work in Pass 1 + (b,M) all-gather + chain reduce + Pass 2, + # mirroring poc-a/cp_gdn_zigzag.py:cp_zigzag_two_pass. + # - Returns FULL-T core_attn_out so downstream layers (full-attn / MLP) + # keep operating on full sequence. + # + # Caveats (B3 v1, prefill-correctness only): + # - conv1d runs on the FULL sequence on every rank (redundant). No conv1d + # boundary error (ranks compute the SAME conv1d). + # - SSM state writeback is SKIPPED. Decode after this prefill is broken; + # prefill-correctness test (V5/V6 hidden-state dump methodology) is + # unaffected. + # - Prefix cache (extend_prefix_lens > 0) not supported. PoC scope. + # =========================================================================== + def _forward_extend_cp( + self, + layer: RadixLinearAttention, + forward_batch: ForwardBatch, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + cp_meta, + 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 + + # === V17 DIAG: check if ssm_states[cache_indices[0]] is non-zero for first prefill + # If non-zero → cp=1's "initial_state=ssm_states" reads stale data ≠ my zeros. + if not getattr(self.__class__, "_v17_ssm_diag_printed", False): + try: + slot = int(cache_indices[0].item()) + ssm_slice = ssm_states[slot] + print( + f"[V17 SSM DIAG] layer={layer.layer_id} slot={slot} " + f"ssm_states[slot] abs max={ssm_slice.abs().max().item():.3e} " + f"mean abs={ssm_slice.abs().mean().item():.3e} " + f"shape={tuple(ssm_slice.shape)}", + flush=True, + ) + if layer.layer_id >= 5: + self.__class__._v17_ssm_diag_printed = True + except Exception as _e: + print(f"[V17 SSM DIAG ERROR] {_e}", flush=True) + + # ----- conv1d on FULL mixed_qkv (every rank computes same result) ----- + 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 + gating (FULL-T) ----- + 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) + + # ----- CP segment plan (zigzag) ----- + num_segs = SEG_PER_RANK * cp_size + assert seq_len % num_segs == 0, ( + f"V15 B3 PoC 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] # zigzag rule + + def _slice_seg(t, seg_idx): + s = seg_idx * seg_len + return t.narrow(1, s, seg_len).contiguous() + + # ----- Pass 1: only owned segments, init=0, capture (b, M) ----- + 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) over cp_group (use sglang's GroupCoordinator wrapper) ----- + 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 gather → causal index (PoC: causal_to_gather_perm) ----- + 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 chain reduce (LOCAL): S_init[i+1] = S_init[i] @ M[i] + b[i] ----- + 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): + # PoC: torch.einsum("nhvk,hkj->nhvj", S, M) — affine fp32 update. + S_init[i] = ( + torch.einsum("nhvk,hkj->nhvj", S_init[i - 1], M_causal[i - 1]) + + b_causal[i - 1].unsqueeze(0) + ) + + # ----- Pass 2: only owned segments with correct S_init ----- + # V19.1: capture in-place writeback of S_init_seg after Pass 2 of seg_(N-1) + # (the rank that owns the last causal segment = rank 0 by zigzag rule). + # This is the kernel's true final state for that segment, ensuring + # bit-exact match with cp=1's writeback (which uses the same kernel). + out_pair = torch.empty( + SEG_PER_RANK, seg_len, Hv, Vdim, dtype=value.dtype, device=device + ) + captured_final_state = None # only rank owning seg_(N-1) sets this + 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] + # On the rank owning seg_(N-1), capture the kernel's in-place state writeback. + 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] # [num_segs, seg_len, Hv, Vdim] + core_attn_out = out_causal.reshape(1, num_segs * seg_len, Hv, Vdim) + + # =========================================================================== + # V19.1: ssm_state writeback for decode-after-CP-prefill support. + # + # Mirrors sglang's full-attn pattern (cp_allgather_and_save_kv_cache): + # every CP rank ends prefill with the request's full state in its local + # pool. Decode then runs through the standard non-CP path (scheduler + # clears attn_cp_metadata on extend→decode transition; GDN forward_decode + # reads ssm_states[cache_indices[0]] symmetrically per CP rank). + # + # The rank owning seg_(N-1) (= rank 0 by zigzag rule, since rank 0 owns + # {0, 2*cp_size-1}) ran Pass 2 on that segment and captured the kernel's + # in-place final-state writeback into `captured_final_state`. We broadcast + # this to all CP ranks and write to local ssm_states[slot]. Using the + # kernel's actual writeback (rather than reconstructing via einsum + # extension of the chain reduce) ensures bit-exact equivalence with + # cp=1's automatic in-place writeback in the non-CP path. + # + # Conv_state writeback: already handled by causal_conv1d_fn above (every + # CP rank runs full-T conv1d on identical inputs → identical writes, + # idempotent across ranks). + # =========================================================================== + if captured_final_state is None: + # ranks not owning seg_(N-1) allocate buffer to receive broadcast + captured_final_state = torch.empty( + Hv, Vdim, Kdim, dtype=torch.float32, device=device + ) + # Broadcast from the rank that captured (= rank 0 in the cp group). + # attn_cp_all_gather_into_tensor with output sized cp_size × input gives + # us all ranks' values; rank 0's slot is at offset 0. Each rank then + # uses gathered[0] as the authoritative S_final. + 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) + S_final = gathered_states[0] # rank 0 owns seg_(N-1) → its captured state is the truth + ssm_states[cache_indices[0]] = S_final.to(ssm_states.dtype) + + # === V20: dump ssm_state after V19.1 CP writeback (gated by DUMP_SSM=1) === + import os as _os + if _os.getenv("DUMP_SSM") == "1": + try: + import torch.distributed as _dist + _world_rank = _dist.get_rank() if _dist.is_initialized() else 0 + _path = ( + f"/tmp/v20_ssm_cp{cp_size}_layer{layer.layer_id}" + f"_rank{_world_rank}_T{seq_len}.pt" + ) + if not _os.path.exists(_path): + _slot = int(cache_indices[0].item()) + torch.save( + ssm_states[_slot].detach().float().cpu(), + _path, + ) + if layer.layer_id == 0: + print( + f"[V20 SSM DUMP cp={cp_size}] layer={layer.layer_id} " + f"slot={_slot} path={_path}", + flush=True, + ) + except Exception as _e: + print(f"[V20 SSM DUMP cp ERROR] {_e}", flush=True) return core_attn_out diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 2b3315fbf975..77c474c9dd5e 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -43,11 +43,20 @@ 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, ) +# === V15 (B3): minimal CP wiring (set attn_cp_metadata only; backend handles the rest) === +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 +1094,36 @@ def forward( pp_proxy_tensors: Optional[PPProxyTensors] = None, input_deepstack_embeds: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, PPProxyTensors]: + # === V15 (B3) — minimal CP wiring === + # GDN-backend-only CP. We set a CUSTOM attribute (_gdn_cp_metadata) that the + # GDN backend reads, leaving forward_batch.attn_cp_metadata = None so that + # the MoE communicator's CP path (which expects local-T hidden_states) does + # NOT activate. Hidden_states stays full-T everywhere. Only GDN backend + # internally splits → does (b,M) merge → returns full-T. + # Full-attn (aiter) layers run full-T redundantly across CP ranks. + if is_prefill_context_parallel_enabled(): + cp_size = get_attention_cp_size() + cp_rank = get_attention_cp_rank() + # gate fix from V8: input_ids may be None when called via general_mm_embed_routine + seq_len_for_cp = None + if input_ids is not None: + seq_len_for_cp = len(input_ids) + elif input_embeds is not None: + seq_len_for_cp = input_embeds.shape[0] + 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, + cp_rank, + cp_size, + forward_batch.seq_lens_cpu.tolist(), + ) + forward_batch._gdn_cp_size = cp_size + forward_batch._gdn_cp_rank = cp_rank + # Initialize hidden states if self.pp_group.is_first_rank: if input_embeds is None: @@ -1143,6 +1182,28 @@ def forward( else: hidden_states, _ = self.norm(hidden_states, residual) + # === V16: hidden state dump for cp=1 vs cp=2 validation === + # Gated by env var DUMP_HS to avoid runtime cost in production. Dumps + # post-norm hidden_states (the input to lm_head) on each prefill. + # One file per (cp_size, world_rank) per call. + import os as _os + if ( + _os.getenv("DUMP_HS") == "1" + and forward_batch.forward_mode.is_extend() + and hidden_states.shape[0] >= 64 + ): + try: + import torch.distributed as _dist + _world_rank = _dist.get_rank() if _dist.is_initialized() else 0 + _cp_size = getattr(forward_batch, "_gdn_cp_size", get_attention_cp_size()) + _path = f"/tmp/v16_cp{_cp_size}_rank{_world_rank}_T{hidden_states.shape[0]}_hs.pt" + if not _os.path.exists(_path): # one-shot per file + torch.save(hidden_states.detach().float().cpu(), _path) + print(f"[V16 DUMP] saved {_path} shape={tuple(hidden_states.shape)}", + flush=True) + except Exception as _e: + print(f"[V16 DUMP ERROR] {_e}", flush=True) + if len(aux_hidden_states) == 0: return hidden_states From f7a3acfcc20a27b15088538431ffac1a3cd1c525 Mon Sep 17 00:00:00 2001 From: alexsun07 Date: Sat, 9 May 2026 08:33:51 +0000 Subject: [PATCH 2/3] Simplify Qwen3.5 CP: strip debug prints, dump blocks, V-numbered comments Remove instrumentation that was added during V1->V27 development but is no longer needed in production: - DUMP_HS hidden_state dump in qwen3_5.py - DUMP_SSM ssm_state dumps in gdn_backend.py (both non-CP and CP paths) - V17 SSM DIAG block at top of _forward_extend_cp - [V15 CP FIRE] one-shot print + _v15_cp_path_printed class attr - Defensive try/except + import os/torch.distributed only used by dumps Also rewrite comments in why-style instead of citing V-N debug sessions: - Architectural rationale for the custom _gdn_cp_metadata field stays (decoupling from MoE communicator's CP path is the key design point) - Algorithm walkthrough on _forward_extend_cp now describes the math, not the iteration history - compute_M_total_pytorch_ref docstring trimmed; PoC v0/v1 wording replaced with a fused-Triton-can-replace-this pointer - 'NEW' / 'Option-A' annotations on M_out parameters removed No algorithm changes; the cp=2 + Qwen3.5-397B-FP8 verbatim echo test (see report.md \xc2\xa71.3) still passes with the same output fidelity as the original V21 PASS run. Net: -92 LOC across the 4 CP files (-213 / +121, mostly reflowed comments). --- .../sglang/srt/layers/attention/fla/chunk.py | 12 +- .../srt/layers/attention/fla/chunk_delta_h.py | 31 +-- .../layers/attention/linear/gdn_backend.py | 233 ++++++------------ python/sglang/srt/models/qwen3_5.py | 58 ++--- 4 files changed, 121 insertions(+), 213 deletions(-) diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py index 79d5bd2a00eb..d87e685a7e26 100644 --- a/python/sglang/srt/layers/attention/fla/chunk.py +++ b/python/sglang/srt/layers/attention/fla/chunk.py @@ -35,7 +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, # NEW: Option-A CP segment-M out-param + M_out: Optional[torch.Tensor] = None, ): g = chunk_local_cumsum( g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices @@ -60,7 +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, # NEW + M_out=M_out, ) o = chunk_fwd_o( q=q, @@ -94,7 +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, # NEW: Option-A CP segment-M + M_out: Optional[torch.Tensor] = None, ): q_orig = q k_orig = k @@ -119,7 +119,7 @@ def forward( initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, - M_out=M_out, # NEW + M_out=M_out, ) return o.to(q.dtype), h @@ -137,7 +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, # NEW: Option-A CP segment-M out-param + M_out: Optional[torch.Tensor] = None, ): r""" Args: @@ -252,7 +252,7 @@ def chunk_gated_delta_rule( initial_state_indices, cu_seqlens, use_qk_l2norm_in_kernel, - M_out, # NEW + 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 98228114ae4d..f3177c35c14e 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py +++ b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py @@ -279,24 +279,26 @@ def compute_M_total_pytorch_ref( M_out: torch.Tensor, # [N, H, K, K] fp32 — pre-allocated output, filled in place chunk_size: int = CHUNK_SIZE, ) -> None: - """PyTorch reference for the segment-end M_total matrix used by Option-A - Context Parallel (zigzag + (b, M) merge) for GDN linear-attention. + """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:: - 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): + 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] - (γ-to-end factor mirrors fwd_h's b_v scaling at chunk_delta_h.py:185) - Then M_seg = N_chunk_1 @ N_chunk_2 @ ... @ N_chunk_NT (matrix product). + 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``. - NOTE: This is the PoC v0 implementation — pure PyTorch, no fused Triton - kernel. v1 will add a fused Triton kernel for perf parity. The semantics - here are the source of truth that the future Triton kernel must match. + 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" @@ -371,7 +373,7 @@ def chunk_gated_delta_rule_fwd_h( save_new_value: bool = True, cu_seqlens: Optional[torch.LongTensor] = None, chunk_indices: Optional[torch.LongTensor] = None, - M_out: Optional[torch.Tensor] = None, # NEW: out-param for Option-A CP + 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] @@ -426,8 +428,9 @@ def grid(meta): num_stages=2, ) - # Option-A CP path: optionally fill the segment-end M matrix. - # Default M_out=None ⇒ skipped entirely (zero numeric or perf change). + # 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, diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 23fce4579004..56afe52f2ee4 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -6,11 +6,7 @@ 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_rank, - get_attention_cp_size, -) +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, @@ -364,13 +360,13 @@ def forward_extend( assert isinstance(mixed_qkv, torch.Tensor) seq_len = mixed_qkv.shape[0] - # === V15 (B3) — CP-aware GDN forward === - # Read custom _gdn_cp_metadata (NOT attn_cp_metadata, which would activate - # other downstream CP code paths in the MoE communicator that we can't - # satisfy without splitting hidden_states). - # Fall through to non-CP path when seq_len doesn't divide evenly (warmup, - # short prompts) — CP code path requires seq_len % (2*cp_size) == 0 for - # the equal-segment PoC algorithm. + # 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 ( @@ -378,20 +374,16 @@ def forward_extend( 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 # need at least one chunk per segment (CHUNK_SIZE=64) + and seq_len >= 2 * cp_size * 64 ): - cp_rank = getattr(forward_batch, "_gdn_cp_rank", get_attention_cp_rank()) - # one-shot print so we can confirm CP path fired in production - if not getattr(self.__class__, "_v15_cp_path_printed", False): - print( - f"[V15 CP FIRE] layer={layer.layer_id} seq_len={seq_len} " - f"cp_size={cp_size} cp_rank={cp_rank} " - f"split_list={cp_meta.split_list}", - flush=True, - ) - self.__class__._v15_cp_path_printed = True return self._forward_extend_cp( - layer, forward_batch, mixed_qkv, a, b, cp_meta, cp_size, cp_rank + layer, + forward_batch, + mixed_qkv, + a, + b, + cp_size, + forward_batch._gdn_cp_rank, ) is_target_verify = forward_batch.forward_mode.is_target_verify() @@ -510,55 +502,31 @@ def forward_extend( forward_batch, h, ssm_states, forward_metadata ) - # === V20: dump ssm_state after non-CP prefill (gated by DUMP_SSM=1 env) === - # Captures the kernel's in-place writeback for cp=1 baseline + noise floor. - # One-shot per (layer_id, world_rank) per server boot to avoid filesystem spam. - import os as _os - if _os.getenv("DUMP_SSM") == "1" and forward_batch.forward_mode.is_extend(): - try: - import torch.distributed as _dist - _world_rank = _dist.get_rank() if _dist.is_initialized() else 0 - _seq_len = mixed_qkv.shape[0] if isinstance(mixed_qkv, torch.Tensor) else seq_len - _path = ( - f"/tmp/v20_ssm_cp1_layer{layer.layer_id}_rank{_world_rank}" - f"_T{_seq_len}.pt" - ) - if not _os.path.exists(_path): - _slot = int(cache_indices[0].item()) - torch.save( - ssm_states[_slot].detach().float().cpu(), - _path, - ) - if layer.layer_id == 0: - print( - f"[V20 SSM DUMP cp=1] layer={layer.layer_id} slot={_slot} " - f"path={_path}", - flush=True, - ) - except Exception as _e: - print(f"[V20 SSM DUMP cp=1 ERROR] {_e}", flush=True) return core_attn_out - # =========================================================================== - # V15 (B3): CP-aware GDN forward. + # CP-aware prefill for the GDN linear-attention path. + # + # mixed_qkv arrives full-T (the model layer does not pre-split it; the + # MoE / aiter full-attn layers in this hybrid model are not CP-aware, + # so all non-GDN layers must keep operating on full sequences). The + # algorithm is zigzag two-pass + (b, M) chain reduce: # - # Contract: - # - mixed_qkv arrives FULL-T (model file did NOT split). This avoids the - # unbounded OOB chain on aiter_backend (which is not CP-aware). - # - cp_meta is set by Qwen3_5ForCausalLM.forward when CP is active. - # - We do real CP work in Pass 1 + (b,M) all-gather + chain reduce + Pass 2, - # mirroring poc-a/cp_gdn_zigzag.py:cp_zigzag_two_pass. - # - Returns FULL-T core_attn_out so downstream layers (full-attn / MLP) - # keep operating on full sequence. + # 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. # - # Caveats (B3 v1, prefill-correctness only): - # - conv1d runs on the FULL sequence on every rank (redundant). No conv1d - # boundary error (ranks compute the SAME conv1d). - # - SSM state writeback is SKIPPED. Decode after this prefill is broken; - # prefill-correctness test (V5/V6 hidden-state dump methodology) is - # unaffected. - # - Prefix cache (extend_prefix_lens > 0) not supported. PoC scope. - # =========================================================================== + # 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, @@ -566,7 +534,6 @@ def _forward_extend_cp( mixed_qkv: torch.Tensor, a: torch.Tensor, b: torch.Tensor, - cp_meta, cp_size: int, cp_rank: int, ): @@ -583,25 +550,8 @@ def _forward_extend_cp( ssm_states = mamba_cache_params.temporal has_initial_states = forward_batch.extend_prefix_lens > 0 - # === V17 DIAG: check if ssm_states[cache_indices[0]] is non-zero for first prefill - # If non-zero → cp=1's "initial_state=ssm_states" reads stale data ≠ my zeros. - if not getattr(self.__class__, "_v17_ssm_diag_printed", False): - try: - slot = int(cache_indices[0].item()) - ssm_slice = ssm_states[slot] - print( - f"[V17 SSM DIAG] layer={layer.layer_id} slot={slot} " - f"ssm_states[slot] abs max={ssm_slice.abs().max().item():.3e} " - f"mean abs={ssm_slice.abs().mean().item():.3e} " - f"shape={tuple(ssm_slice.shape)}", - flush=True, - ) - if layer.layer_id >= 5: - self.__class__._v17_ssm_diag_printed = True - except Exception as _e: - print(f"[V17 SSM DIAG ERROR] {_e}", flush=True) - - # ----- conv1d on FULL mixed_qkv (every rank computes same result) ----- + # 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[ @@ -620,7 +570,7 @@ def _forward_extend_cp( seq_lens_cpu=forward_batch.extend_seq_lens_cpu, ).transpose(0, 1)[:seq_len] - # ----- Split q/k/v + gating (FULL-T) ----- + # 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], @@ -636,22 +586,25 @@ def _forward_extend_cp( value = value.view(1, seq_len, Hv, Vdim) g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) - # ----- CP segment plan (zigzag) ----- + # 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"V15 B3 PoC requires seq_len divisible by 2*cp_size; " + 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] # zigzag rule + 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: only owned segments, init=0, capture (b, M) ----- + # 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 ) @@ -685,7 +638,8 @@ def _slice_seg(t, seg_idx): b_pair[slot] = state[0] M_pair[slot] = M_buf[0] - # ----- All-gather (b, M) over cp_group (use sglang's GroupCoordinator wrapper) ----- + # 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 ) @@ -695,7 +649,8 @@ def _slice_seg(t, seg_idx): attn_cp_all_gather_into_tensor(b_gathered, b_pair.contiguous()) attn_cp_all_gather_into_tensor(M_gathered, M_pair.contiguous()) - # ----- Reorder gather → causal index (PoC: causal_to_gather_perm) ----- + # 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) @@ -705,27 +660,26 @@ def _slice_seg(t, seg_idx): b_causal = b_gathered[perm_t] M_causal = M_gathered[perm_t] - # ----- fp32 chain reduce (LOCAL): S_init[i+1] = S_init[i] @ M[i] + b[i] ----- + # 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): - # PoC: torch.einsum("nhvk,hkj->nhvj", S, M) — affine fp32 update. S_init[i] = ( torch.einsum("nhvk,hkj->nhvj", S_init[i - 1], M_causal[i - 1]) + b_causal[i - 1].unsqueeze(0) ) - # ----- Pass 2: only owned segments with correct S_init ----- - # V19.1: capture in-place writeback of S_init_seg after Pass 2 of seg_(N-1) - # (the rank that owns the last causal segment = rank 0 by zigzag rule). - # This is the kernel's true final state for that segment, ensuring - # bit-exact match with cp=1's writeback (which uses the same kernel). + # 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 # only rank owning seg_(N-1) sets this + captured_final_state = None last_seg_idx = num_segs - 1 for slot, seg_idx in enumerate(owned): q_seg = _slice_seg(query, seg_idx) @@ -747,77 +701,42 @@ def _slice_seg(t, seg_idx): M_out=None, ) out_pair[slot] = o[0] # [seg_len, Hv, Vdim] - # On the rank owning seg_(N-1), capture the kernel's in-place state writeback. 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 ----- + # 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] # [num_segs, seg_len, Hv, Vdim] + out_causal = out_gathered[perm_t] core_attn_out = out_causal.reshape(1, num_segs * seg_len, Hv, Vdim) - # =========================================================================== - # V19.1: ssm_state writeback for decode-after-CP-prefill support. + # ssm_state writeback for decode-after-CP-prefill. # - # Mirrors sglang's full-attn pattern (cp_allgather_and_save_kv_cache): - # every CP rank ends prefill with the request's full state in its local - # pool. Decode then runs through the standard non-CP path (scheduler - # clears attn_cp_metadata on extend→decode transition; GDN forward_decode - # reads ssm_states[cache_indices[0]] symmetrically per CP rank). + # 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. # - # The rank owning seg_(N-1) (= rank 0 by zigzag rule, since rank 0 owns - # {0, 2*cp_size-1}) ran Pass 2 on that segment and captured the kernel's - # in-place final-state writeback into `captured_final_state`. We broadcast - # this to all CP ranks and write to local ssm_states[slot]. Using the - # kernel's actual writeback (rather than reconstructing via einsum - # extension of the chain reduce) ensures bit-exact equivalence with - # cp=1's automatic in-place writeback in the non-CP path. + # 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: already handled by causal_conv1d_fn above (every - # CP rank runs full-T conv1d on identical inputs → identical writes, - # idempotent across ranks). - # =========================================================================== + # 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: - # ranks not owning seg_(N-1) allocate buffer to receive broadcast captured_final_state = torch.empty( Hv, Vdim, Kdim, dtype=torch.float32, device=device ) - # Broadcast from the rank that captured (= rank 0 in the cp group). - # attn_cp_all_gather_into_tensor with output sized cp_size × input gives - # us all ranks' values; rank 0's slot is at offset 0. Each rank then - # uses gathered[0] as the authoritative S_final. 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) - S_final = gathered_states[0] # rank 0 owns seg_(N-1) → its captured state is the truth - ssm_states[cache_indices[0]] = S_final.to(ssm_states.dtype) - - # === V20: dump ssm_state after V19.1 CP writeback (gated by DUMP_SSM=1) === - import os as _os - if _os.getenv("DUMP_SSM") == "1": - try: - import torch.distributed as _dist - _world_rank = _dist.get_rank() if _dist.is_initialized() else 0 - _path = ( - f"/tmp/v20_ssm_cp{cp_size}_layer{layer.layer_id}" - f"_rank{_world_rank}_T{seq_len}.pt" - ) - if not _os.path.exists(_path): - _slot = int(cache_indices[0].item()) - torch.save( - ssm_states[_slot].detach().float().cpu(), - _path, - ) - if layer.layer_id == 0: - print( - f"[V20 SSM DUMP cp={cp_size}] layer={layer.layer_id} " - f"slot={_slot} path={_path}", - flush=True, - ) - except Exception as _e: - print(f"[V20 SSM DUMP cp ERROR] {_e}", flush=True) + 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 77c474c9dd5e..90c3a8358bc5 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -50,7 +50,9 @@ is_dp_attention_enabled, ) -# === V15 (B3): minimal CP wiring (set attn_cp_metadata only; backend handles the rest) === +# 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, @@ -1094,22 +1096,28 @@ def forward( pp_proxy_tensors: Optional[PPProxyTensors] = None, input_deepstack_embeds: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, PPProxyTensors]: - # === V15 (B3) — minimal CP wiring === - # GDN-backend-only CP. We set a CUSTOM attribute (_gdn_cp_metadata) that the - # GDN backend reads, leaving forward_batch.attn_cp_metadata = None so that - # the MoE communicator's CP path (which expects local-T hidden_states) does - # NOT activate. Hidden_states stays full-T everywhere. Only GDN backend - # internally splits → does (b,M) merge → returns full-T. - # Full-attn (aiter) layers run full-T redundantly across CP ranks. + # 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() - cp_rank = get_attention_cp_rank() - # gate fix from V8: input_ids may be None when called via general_mm_embed_routine - seq_len_for_cp = None + # 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 = len(input_ids) + 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 @@ -1117,12 +1125,12 @@ def forward( ): forward_batch._gdn_cp_metadata = prepare_context_parallel_metadata( seq_len_for_cp, - cp_rank, + get_attention_cp_rank(), cp_size, forward_batch.seq_lens_cpu.tolist(), ) forward_batch._gdn_cp_size = cp_size - forward_batch._gdn_cp_rank = cp_rank + forward_batch._gdn_cp_rank = get_attention_cp_rank() # Initialize hidden states if self.pp_group.is_first_rank: @@ -1182,28 +1190,6 @@ def forward( else: hidden_states, _ = self.norm(hidden_states, residual) - # === V16: hidden state dump for cp=1 vs cp=2 validation === - # Gated by env var DUMP_HS to avoid runtime cost in production. Dumps - # post-norm hidden_states (the input to lm_head) on each prefill. - # One file per (cp_size, world_rank) per call. - import os as _os - if ( - _os.getenv("DUMP_HS") == "1" - and forward_batch.forward_mode.is_extend() - and hidden_states.shape[0] >= 64 - ): - try: - import torch.distributed as _dist - _world_rank = _dist.get_rank() if _dist.is_initialized() else 0 - _cp_size = getattr(forward_batch, "_gdn_cp_size", get_attention_cp_size()) - _path = f"/tmp/v16_cp{_cp_size}_rank{_world_rank}_T{hidden_states.shape[0]}_hs.pt" - if not _os.path.exists(_path): # one-shot per file - torch.save(hidden_states.detach().float().cpu(), _path) - print(f"[V16 DUMP] saved {_path} shape={tuple(hidden_states.shape)}", - flush=True) - except Exception as _e: - print(f"[V16 DUMP ERROR] {_e}", flush=True) - if len(aux_hidden_states) == 0: return hidden_states From a26c79310f043275e59ae43fba61ce29ff173283 Mon Sep 17 00:00:00 2001 From: alexsun07 Date: Sat, 9 May 2026 09:18:37 +0000 Subject: [PATCH 3/3] Clarify _forward_extend_cp comment: distinguish MoE vs aiter reasons The previous comment said 'MoE / aiter full-attn layers are not CP-aware, so all non-GDN layers must keep operating on full sequences', which lumped two distinct reasons together and was misleading about MoE. MoE compute itself is per-token and CP-orthogonal. What forces full-T on the MoE path is sglang's standard CP wrapper chain that activates on forward_batch.attn_cp_metadata is not None: - cp_split_and_rebuild_data at the model entry (qwen2_moe.py:778) - moe_cp_all_gather_into_tensor in the layer communicator (communicator.py:1106) - cp_all_gather_rerange_output before lm_head (qwen2_moe.py:831) All three assume hidden_states are already split to local-T. By using a private _gdn_cp_metadata field instead, we avoid triggering any of them and keep the MoE path running unmodified on full-T. aiter full-attn separately has no CP code at all and would always need full-T regardless. Comment now states each reason precisely. No code change. --- .../layers/attention/linear/gdn_backend.py | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 56afe52f2ee4..fee856cb4646 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -506,10 +506,26 @@ def forward_extend( # CP-aware prefill for the GDN linear-attention path. # - # mixed_qkv arrives full-T (the model layer does not pre-split it; the - # MoE / aiter full-attn layers in this hybrid model are not CP-aware, - # so all non-GDN layers must keep operating on full sequences). The - # algorithm is zigzag two-pass + (b, M) chain reduce: + # 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