[Qwen3.8 CP 1/4] Context parallelism for QSA attention and the sparse indexer

python/sglang/srt/layers/attention/qsa/metadata.py

@@ -54,6 +54,9 @@ class QSAIndexerMetadata(msgspec.Struct, frozen=True):
     - ``token_to_batch_idx`` maps every query/token row handled by the indexer
       onto a row of ``sequence_lengths``/``token_slot_table``; DP attention
       token padding adds physical rows beyond this mapping, never inside it.
+      For CP extend, it covers the global packed *new* tokens used by the
+      compression write plan. Query selection uses a zigzag-sharded copy;
+      cached prefix tokens are K/V context, not rows in this mapping.
     - For the paged modes the mapping is the identity
       (``arange(num_query_rows)``), so page-table/MQA inputs built per
       ``sequence_lengths`` row line up with per-query sparse-attention rows.
@@ -123,8 +126,13 @@ def get_prefill_mqa_inputs(
         self,
         layer_id: int,
         positions: torch.Tensor,
+        query_sequence_ids: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Gather packed compressed K and ragged ranges for prefill MQA."""
+        """Gather packed compressed K and ragged ranges for prefill MQA.
+
+        ``query_sequence_ids`` overrides ``token_to_batch_idx`` for callers that
+        score a subset of the batch's rows (prefill CP: this rank's zigzag rows);
+        ``positions`` must then be those rows' logical positions."""
 
         pool = self.token_to_kv_pool
         ratio = self.compress_ratio
@@ -152,7 +160,9 @@ def get_prefill_mqa_inputs(
                 (0, pool.qsa_index_kv_heads, pool.qsa_index_head_dim)
             )
         )
-        num_valid_tokens = self.token_to_batch_idx.numel()
+        if query_sequence_ids is None:
+            query_sequence_ids = self.token_to_batch_idx
+        num_valid_tokens = query_sequence_ids.numel()
         if positions.numel() < num_valid_tokens:
             raise ValueError(
                 "QSA prefill positions are shorter than the request mapping: "
@@ -162,7 +172,7 @@ def get_prefill_mqa_inputs(
         row_starts, row_ends, _ = build_qsa_row_ranges(
             sequence_lengths,
             positions.to(sequence_lengths.device),
-            self.token_to_batch_idx.to(sequence_lengths.device),
+            query_sequence_ids.to(sequence_lengths.device),
             self.compress_ratio,
         )
         return compressed_keys, row_starts, row_ends, sequence_lengths

python/sglang/srt/layers/attention/qsa/qsa_indexer.py

@@ -17,6 +17,13 @@
     build_rope_position_matrix,
 )
 from sglang.srt.layers.attention.qsa.mqa import qsa_mqa_decode, qsa_mqa_prefill
+from sglang.srt.layers.cp.utils import (
+    ContextParallelStrategyKind,
+    cp_materialize_global_token_order,
+    cp_shard_hidden_states,
+    cp_shard_position_ids,
+    get_cp_strategy,
+)
 from sglang.srt.layers.layernorm import GemmaRMSNorm
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
@@ -525,19 +532,113 @@ def select_decode_tokens(
             token_topk=self.token_topk,
         )
 
+    def forward_cuda_cp(
+        self,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch,
+        indexer_metadata,
+        global_rope_positions: torch.Tensor,
+    ) -> torch.Tensor:
+        """Prefill CP: score only this rank's zigzag query rows.
+
+        ``hidden_states``/``positions`` are the local (padded) rows and their
+        RoPE positions. Index q/k are projected for the local rows only; the raw
+        keys are all-gathered into global token order so the pending-group ring
+        store and the block compression run on the batch's global write plan
+        exactly as without CP (every rank ends up with the complete compressed-K
+        cache). Selection then scores the local rows against all compressed
+        blocks of their request, position-driven, so the output rows line up
+        with the local query rows the sparse attention consumes."""
+        strategy = get_cp_strategy()
+        if strategy is None or strategy.kind != ContextParallelStrategyKind.ZIGZAG:
+            strategy_name = strategy.name if strategy is not None else "none"
+            raise NotImplementedError(
+                "QSA prefill CP only supports the zigzag strategy; "
+                f"got {strategy_name}."
+            )
+        meta = forward_batch.attn_cp_metadata
+        num_local = int(meta.total_q_prev_tokens + meta.total_q_next_tokens)
+        global_logical = getattr(forward_batch, "positions", None)
+        global_logical = (
+            global_logical[0] if global_logical.ndim == 2 else global_logical
+        ).flatten()
+        num_global = indexer_metadata.get_token_to_batch_idx().numel()
+        global_logical = global_logical[:num_global]
+        local_logical = cp_shard_position_ids(global_logical, forward_batch)[:num_local]
+        local_sequence_ids = cp_shard_hidden_states(
+            indexer_metadata.get_token_to_batch_idx(), forward_batch
+        )[:num_local]
+        hidden_states = hidden_states[:num_local]
+        positions = (
+            positions[:, :num_local] if positions.ndim == 2 else positions[:num_local]
+        )
+        # Use the compressed-cache dtype for Q. Leaving cache_loc unset defers
+        # the ring store until raw local K has been gathered into global order.
+        q, token_k, _ = self.project_qk(
+            hidden_states, positions, pool=indexer_metadata.token_to_kv_pool
+        )
+        # Raw keys of all new tokens in global packed order: the ring store
+        # and compression use the global write plan unchanged. Cached prefix
+        # groups already reside in the compressed-K pool on every rank.
+        token_k_local = token_k.reshape(num_local, -1).contiguous()
+        token_k_full = cp_materialize_global_token_order(
+            token_k_local, forward_batch
+        ).reshape(-1, self.index_kv_heads, self.index_head_dim)
+        global_rope = (
+            global_rope_positions[:, :num_global]
+            if global_rope_positions.ndim == 2
+            else global_rope_positions[:num_global]
+        )
+        self.update_key_state_and_compress(
+            token_k_full,
+            global_logical,
+            global_rope,
+            indexer_metadata,
+            state_slots=indexer_metadata.pending_ring_slots,
+            state_stored=False,
+        )
+        compressed_keys, row_starts, row_ends, sequence_lengths = (
+            indexer_metadata.get_prefill_mqa_inputs(
+                self.layer_id, local_logical, query_sequence_ids=local_sequence_ids
+            )
+        )
+        row_sequence_lengths = sequence_lengths.index_select(
+            0, local_sequence_ids.long()
+        )
+        return self.select_prefill_tokens(
+            q,
+            compressed_keys,
+            row_starts,
+            row_ends,
+            local_logical,
+            row_sequence_lengths,
+        )
+
     def _forward_impl(
         self,
         hidden_states: torch.Tensor,
         positions: torch.Tensor,
         forward_batch,
         indexer_metadata,
+        cp_global_rope_positions: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Portable orchestration shared by different platforms.
         Fast paths are gated per platforms inside kernel calls.
         """
         forward_mode = forward_batch.forward_mode
         is_target_verify = getattr(forward_mode, "is_target_verify", lambda: False)()
         is_draft_extend = getattr(forward_mode, "is_draft_extend_v2", lambda: False)()
+        if cp_global_rope_positions is not None and not (
+            forward_mode.is_decode() or is_target_verify or is_draft_extend
+        ):
+            return self.forward_cuda_cp(
+                hidden_states,
+                positions,
+                forward_batch,
+                indexer_metadata,
+                cp_global_rope_positions,
+            )
         if forward_mode.is_decode() or is_target_verify or is_draft_extend:
             # EAGLE/MTP may advance the model's RoPE coordinate independently
             # from the physical paged-KV position.  Compression and sparse
@@ -644,9 +745,14 @@ def forward_cuda(
         positions: torch.Tensor,
         forward_batch,
         indexer_metadata,
+        cp_global_rope_positions: torch.Tensor | None = None,
     ) -> torch.Tensor:
         return self._forward_impl(
-            hidden_states, positions, forward_batch, indexer_metadata
+            hidden_states,
+            positions,
+            forward_batch,
+            indexer_metadata,
+            cp_global_rope_positions,
         )
 
     def forward_xpu(

python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py

@@ -41,6 +41,12 @@
     sparse_gqa_fwd_interface_triton_ck,
     sparse_gqa_packed_decode_triton,
 )
+from sglang.srt.layers.cp.utils import (
+    ContextParallelStrategyKind,
+    cp_materialize_global_token_order,
+    get_cp_strategy,
+    is_cp_active,
+)
 from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, MHATokenToKVPool
 from sglang.srt.mem_cache.qsa_kv_pool import QSATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
@@ -1401,6 +1407,12 @@ def forward_extend(
     ) -> torch.Tensor:
         if topk_indices is None:
             raise ValueError("QSA sparse attention requires topk_indices")
+        if getattr(
+            forward_batch, "attn_cp_metadata", None
+        ) is not None and is_cp_active(forward_batch):
+            return self._forward_extend_cp(
+                q, k, v, layer, forward_batch, topk_indices, save_kv_cache
+            )
         if save_kv_cache:
             fused_output = self._try_fused_kv_attention(
                 q, k, v, layer, forward_batch, topk_indices
@@ -1496,6 +1508,101 @@ def forward_extend(
         )
         return self._pad_extend_output(output, num_output_rows)
 
+    def _forward_extend_cp(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer,
+        forward_batch,
+        topk_indices: torch.Tensor,
+        save_kv_cache: bool,
+    ) -> torch.Tensor:
+        """Prefill CP: this rank holds the query rows of its two zigzag blocks
+        (all heads) and the K/V of the same rows. All-gather new K/V into global
+        token order (also written to the KV pool so decode sees every token),
+        prepend each request's cached prefix from the replicated local KV pool,
+        then run the chunk-prefill sparse kernel with one entry per zigzag
+        block: q rows are contiguous per block, ``kv_lens`` is the block's
+        absolute end position (bottom-right causal), ``cu_k`` is the request's
+        base row in the packed full K/V, and ``topk_indices`` are per-request
+        logical positions exactly as in the non-CP path."""
+        strategy = get_cp_strategy()
+        if strategy is None or strategy.kind != ContextParallelStrategyKind.ZIGZAG:
+            strategy_name = strategy.name if strategy is not None else "none"
+            raise NotImplementedError(
+                "QSA prefill CP only supports the zigzag strategy; "
+                f"got {strategy_name}."
+            )
+        meta = forward_batch.attn_cp_metadata
+        q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim)
+        num_output_rows = q.shape[0]
+        k = k.reshape(-1, layer.tp_k_head_num, layer.qk_head_dim)
+        v = v.reshape(-1, layer.tp_v_head_num, layer.v_head_dim)
+        k_width = k.shape[1] * k.shape[2]
+        v_width = v.shape[1] * v.shape[2]
+        kv_local = torch.cat([k.flatten(1), v.flatten(1)], dim=-1)
+        kv_full = cp_materialize_global_token_order(kv_local, forward_batch)
+        k_full, v_full = kv_full.split([k_width, v_width], dim=-1)
+        k_full = k_full.reshape(-1, k.shape[1], k.shape[2]).contiguous()
+        v_full = v_full.reshape(-1, v.shape[1], v.shape[2]).contiguous()
+        if save_kv_cache:
+            # out_cache_loc was trimmed to the full extend token count by
+            # prepare_cp_forward; write all new tokens on every rank.
+            self.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k_full, v_full
+            )
+        num_local_rows = int(meta.total_q_prev_tokens + meta.total_q_next_tokens)
+        q = q[:num_local_rows].contiguous()
+        indices = topk_indices[:num_local_rows].to(torch.int32).contiguous()
+        seq_lens = [int(x) for x in forward_batch.seq_lens_cpu[: meta.bs]]
+        extend_lens = [int(x) for x in forward_batch.extend_seq_lens_cpu[: meta.bs]]
+        prefix_lens = [seq - extend for seq, extend in zip(seq_lens, extend_lens)]
+        if any(prefix_lens):
+            # The all-gather contains only new tokens. Prefix K/V is already
+            # replicated on each CP rank; pack prefix + new tokens per request
+            # so absolute logical top-k indices address the full context.
+            pool = self.token_to_kv_pool
+            k_buffer = pool.get_key_buffer(layer.layer_id)
+            v_buffer = pool.get_value_buffer(layer.layer_id)
+            req_to_token = self.req_to_token_pool.req_to_token
+            req_indices = forward_batch.req_pool_indices[: meta.bs].tolist()
+            k_parts, v_parts = [], []
+            extend_start = 0
+            for req_idx, prefix_len, extend_len in zip(
+                req_indices, prefix_lens, extend_lens
+            ):
+                if prefix_len:
+                    slots = req_to_token[req_idx, :prefix_len].long()
+                    # QSA stores K/V without scales; casting also supports an
+                    # FP8 prefix alongside the unquantized new tokens.
+                    k_parts.append(k_buffer.index_select(0, slots).to(k_full.dtype))
+                    v_parts.append(v_buffer.index_select(0, slots).to(v_full.dtype))
+                extend_end = extend_start + extend_len
+                k_parts.append(k_full[extend_start:extend_end])
+                v_parts.append(v_full[extend_start:extend_end])
+                extend_start = extend_end
+            # Use gathered new K/V even when save_kv_cache=False; those rows
+            # need not have been written to the cache by the caller.
+            k_full, v_full = torch.cat(k_parts), torch.cat(v_parts)
+        base = [0]
+        for length in seq_lens[:-1]:
+            base.append(base[-1] + length)
+        # one entry per (request, prev block) then per (request, next block);
+        # the kernel reads cu_k[e] as the request's K base row only.
+        cu_k = torch.tensor(base + base, dtype=torch.int32, device=q.device)
+        output = sparse_gqa_fwd_interface_triton_ck(
+            q,
+            k_full,
+            v_full,
+            indices,
+            meta.cu_seqlens_q_combined_tensor,
+            cu_k,
+            meta.kv_len_combined_tensor,
+            layer.scaling,
+        )
+        return self._pad_extend_output(output, num_output_rows)
+
     @staticmethod
     def _pad_extend_output(output: torch.Tensor, num_rows: int) -> torch.Tensor:
         output = output.reshape(output.shape[0], -1)