Qwen3.5 cp zigzag v2 wip

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

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 ...")

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