diff --git a/python/sglang/srt/layers/attention/qsa/metadata.py b/python/sglang/srt/layers/attention/qsa/metadata.py index fe3c047dfcba..581f4bf77669 100644 --- a/python/sglang/srt/layers/attention/qsa/metadata.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py index 7939525eca94..27ebe8cd04b8 100644 --- a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py +++ b/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,12 +532,96 @@ 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. @@ -538,6 +629,16 @@ def _forward_impl( 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( diff --git a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py index 5643dbae167e..005a7717d0f2 100644 --- a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py +++ b/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) diff --git a/test/registered/unit/cp/test_qsa_cp_unit.py b/test/registered/unit/cp/test_qsa_cp_unit.py new file mode 100644 index 000000000000..8ab902c2ba65 --- /dev/null +++ b/test/registered/unit/cp/test_qsa_cp_unit.py @@ -0,0 +1,341 @@ +"""CPU checks for QSA's zigzag-only context parallel entrypoints.""" + +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from sglang.srt.layers.attention import qwen_sparse_attn_backend +from sglang.srt.layers.attention.qsa import qsa_indexer +from sglang.srt.layers.attention.qsa.mqa import _scoring_dtype, torch_qsa_mqa_prefill +from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer +from sglang.srt.layers.attention.qwen_sparse_attn_backend import QwenSparseAttnBackend +from sglang.srt.layers.cp import base as cp_base +from sglang.srt.layers.cp import zigzag +from sglang.srt.layers.cp.base import ContextParallelStrategyKind +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +@pytest.mark.parametrize( + "strategy", + [ + None, + SimpleNamespace(name="interleave", kind=ContextParallelStrategyKind.INTERLEAVE), + ], +) +@pytest.mark.parametrize("entrypoint", ["indexer", "attention"]) +def test_qsa_cp_rejects_unsupported_strategy_before_metadata_or_collectives( + monkeypatch, strategy, entrypoint +): + module = qsa_indexer if entrypoint == "indexer" else qwen_sparse_attn_backend + monkeypatch.setattr(module, "get_cp_strategy", lambda: strategy) + gather = Mock(side_effect=AssertionError("unsupported CP must not communicate")) + monkeypatch.setattr(module, "cp_materialize_global_token_order", gather) + # None metadata and tensors ensure rejection precedes zigzag-specific access. + with pytest.raises( + NotImplementedError, match="QSA prefill CP only supports the zigzag strategy" + ): + if entrypoint == "indexer": + QSAIndexer.forward_cuda_cp(None, None, None, None, None, None) + else: + QwenSparseAttnBackend._forward_extend_cp( + None, None, None, None, None, None, None, False + ) + gather.assert_not_called() + + +def test_cuda_indexer_forwards_global_rope_to_cp_path(): + hidden, local_rope, metadata, global_rope = (object() for _ in range(4)) + batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND) + expected = object() + indexer = SimpleNamespace(forward_cuda_cp=Mock(return_value=expected)) + indexer._forward_impl = lambda *args: QSAIndexer._forward_impl(indexer, *args) + + actual = QSAIndexer.forward_cuda( + indexer, + hidden, + local_rope, + batch, + metadata, + cp_global_rope_positions=global_rope, + ) + + assert actual is expected + indexer.forward_cuda_cp.assert_called_once_with( + hidden, local_rope, batch, metadata, global_rope + ) + + +@pytest.mark.parametrize("compressed_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_cp_indexer_matches_compressed_dtype_without_storing_local_keys( + monkeypatch, compressed_dtype +): + local_rows = torch.tensor([0, 1, 6, 7]) + hidden = (torch.arange(32).reshape(8, 4) / 13).to(torch.bfloat16) + positions = torch.arange(8) + compressed_keys = hidden.reshape(4, 2, 4).mean(1).unsqueeze(1).to(compressed_dtype) + row_starts = torch.zeros(4, dtype=torch.int32) + row_ends = (positions[local_rows] + 1).to(torch.int32) // 2 + pool = SimpleNamespace(qsa_compressed_dtype=compressed_dtype) + metadata = SimpleNamespace( + token_to_kv_pool=pool, + pending_ring_slots=torch.arange(8), + get_token_to_batch_idx=lambda: torch.zeros(8, dtype=torch.int32), + get_prefill_mqa_inputs=lambda *args, **kwargs: ( + compressed_keys, + row_starts, + row_ends, + torch.tensor([8], dtype=torch.int32), + ), + ) + batch = SimpleNamespace( + positions=positions, + attn_cp_metadata=SimpleNamespace(total_q_prev_tokens=2, total_q_next_tokens=2), + ) + monkeypatch.setattr( + qsa_indexer, + "get_cp_strategy", + lambda: SimpleNamespace(kind=ContextParallelStrategyKind.ZIGZAG), + ) + for name in ("cp_shard_position_ids", "cp_shard_hidden_states"): + monkeypatch.setattr(qsa_indexer, name, lambda value, batch: value[local_rows]) + + def gather_raw_keys(local, batch): + torch.testing.assert_close(local, hidden[local_rows]) + return hidden + + monkeypatch.setattr( + qsa_indexer, "cp_materialize_global_token_order", gather_raw_keys + ) + + def score_queries(q, k, starts, ends, query_positions, sequence_lengths): + # Use the same operand check as GPU scoring; the CPU reference otherwise + # accepts mixed dtypes and would conceal the FP8 compatibility failure. + assert _scoring_dtype(q, k) == compressed_dtype + return torch_qsa_mqa_prefill(q, k, starts, ends) + + indexer = SimpleNamespace( + layer_id=0, + index_n_heads=1, + index_kv_heads=1, + index_head_dim=4, + index_qk_proj=lambda value: (torch.cat([value, value], dim=-1), None), + q_layernorm=lambda value: value, + apply_rope=lambda positions, value: value, + _use_fused_prep=Mock( + side_effect=AssertionError("local keys must not be stored") + ), + update_key_state_and_compress=Mock(), + select_prefill_tokens=score_queries, + ) + indexer.project_qk = lambda *args, **kwargs: QSAIndexer.project_qk( + indexer, *args, **kwargs + ) + actual = QSAIndexer.forward_cuda_cp( + indexer, hidden[local_rows], positions[local_rows], batch, metadata, positions + ) + expected = torch_qsa_mqa_prefill( + hidden[local_rows].unsqueeze(1).to(compressed_dtype), + compressed_keys, + row_starts, + row_ends, + ) + torch.testing.assert_close(actual, expected) + indexer._use_fused_prep.assert_not_called() + update = indexer.update_key_state_and_compress.call_args + torch.testing.assert_close(update.args[0], hidden.unsqueeze(1)) + assert update.kwargs["state_stored"] is False + + +def _cpu_sparse_chunk_attention(q, k, v, indices, cu_q, cu_k, kv_lens, scale): + """CPU oracle for the chunk kernel's packed K/V and causal conventions.""" + output = torch.empty_like(q) + for entry in range(kv_lens.numel()): + start, end = int(cu_q[entry]), int(cu_q[entry + 1]) + for row in range(start, end): + visible = int(kv_lens[entry]) - (end - start) + row - start + 1 + selected = indices[row] + selected = selected[selected >= 0].long() + assert (selected < visible).all() + selected = selected + int(cu_k[entry]) + keys = k.index_select(0, selected).repeat_interleave( + q.shape[1] // k.shape[1], dim=1 + ) + values = v.index_select(0, selected).repeat_interleave( + q.shape[1] // v.shape[1], dim=1 + ) + scores = torch.einsum("hd,nhd->hn", q[row], keys) * scale + output[row] = torch.einsum("hn,nhd->hd", scores.softmax(-1), values) + return output + + +@pytest.mark.parametrize("prefix_lens", [(4, 8), (0, 4), (0, 0)]) +@pytest.mark.parametrize("rank", [0, 1]) +@pytest.mark.parametrize("save_kv_cache", [True, False]) +@pytest.mark.parametrize( + "query_dtype,cache_dtype", + [(torch.float32, torch.float32), (torch.bfloat16, torch.float8_e4m3fn)], +) +def test_qsa_cp_cached_prefix_matches_full_context_attention( + monkeypatch, prefix_lens, rank, save_kv_cache, query_dtype, cache_dtype +): + # Unequal request lengths exercise both zigzag halves and uneven rank sizes. + extend_lens = [8, 5] + seq_lens = [prefix + length for prefix, length in zip(prefix_lens, extend_lens)] + monkeypatch.setattr( + cp_base, "get_parallel", lambda: SimpleNamespace(attn_cp_rank=rank) + ) + monkeypatch.setattr(zigzag, "get_device", lambda: SimpleNamespace(device="cpu")) + strategy = zigzag.ZigzagCPStrategy(cp_size=2) + meta = strategy.build_metadata(sum(extend_lens), seq_lens, extend_lens) + meta.per_rank_logical_token = meta.per_rank_actual_token + meta.per_rank_actual_token = [8, 8] + meta.max_rank_len = [8, 8] + monkeypatch.setattr(qwen_sparse_attn_backend, "get_cp_strategy", lambda: strategy) + + generator = torch.Generator().manual_seed(42) + keys = [ + torch.randn(length, 1, 4, generator=generator).to(query_dtype) + for length in seq_lens + ] + values = [ + torch.randn(length, 1, 4, generator=generator).to(query_dtype) + for length in seq_lens + ] + # Cached prefixes have already incurred the KV pool's storage rounding. + for key, value, prefix in zip(keys, values, prefix_lens): + key[:prefix] = key[:prefix].to(cache_dtype).to(query_dtype) + value[:prefix] = value[:prefix].to(cache_dtype).to(query_dtype) + queries = torch.randn(sum(extend_lens), 2, 4, generator=generator).to(query_dtype) + new_k = torch.cat([key[prefix:] for key, prefix in zip(keys, prefix_lens)]) + new_v = torch.cat([value[prefix:] for value, prefix in zip(values, prefix_lens)]) + + # Request IDs and physical token slots deliberately differ from packed order. + req_indices = [4, 1] + req_to_token = torch.zeros(5, max(seq_lens), dtype=torch.int32) + slots = torch.randperm(80, generator=generator)[: sum(seq_lens)] + 1 + k_cache = torch.full((81, 1, 4), float("nan"), dtype=cache_dtype) + v_cache = torch.full_like(k_cache, float("nan")) + out_cache_locs = [] + offset = 0 + for i, (length, prefix) in enumerate(zip(seq_lens, prefix_lens)): + request_slots = slots[offset : offset + length] + offset += length + req_to_token[req_indices[i], :length] = request_slots.to(torch.int32) + k_cache[request_slots[:prefix]] = keys[i][:prefix].to(cache_dtype) + v_cache[request_slots[:prefix]] = values[i][:prefix].to(cache_dtype) + out_cache_locs.append(request_slots[prefix:]) + out_cache_loc = torch.cat(out_cache_locs) + initial_k, initial_v = k_cache.clone(), v_cache.clone() + + def store_kv(layer, locations, k, v): + k_cache[locations.long()] = k.to(cache_dtype) + v_cache[locations.long()] = v.to(cache_dtype) + + pool = SimpleNamespace( + set_kv_buffer=Mock(side_effect=store_kv), + get_key_buffer=Mock(return_value=k_cache), + get_value_buffer=Mock(return_value=v_cache), + ) + if not any(prefix_lens): + pool.get_key_buffer.side_effect = AssertionError( + "zero prefix needs no cache read" + ) + pool.get_value_buffer.side_effect = AssertionError( + "zero prefix needs no cache read" + ) + backend = SimpleNamespace( + token_to_kv_pool=pool, + req_to_token_pool=SimpleNamespace(req_to_token=req_to_token), + _pad_extend_output=QwenSparseAttnBackend._pad_extend_output, + ) + batch = SimpleNamespace( + attn_cp_metadata=meta, + seq_lens_cpu=seq_lens, + extend_seq_lens_cpu=extend_lens, + extend_prefix_lens_cpu=list(prefix_lens), + req_pool_indices=torch.tensor(req_indices), + out_cache_loc=out_cache_loc, + ) + local_q = strategy.shard_hidden_states(queries, batch) + local_k = strategy.shard_hidden_states(new_k, batch) + local_v = strategy.shard_hidden_states(new_v, batch) + num_local = meta.total_q_prev_tokens + meta.total_q_next_tokens + local_rows = strategy.shard_hidden_states(torch.arange(sum(extend_lens)), batch)[ + :num_local + ] + logical_positions = torch.cat( + [torch.arange(prefix, length) for prefix, length in zip(prefix_lens, seq_lens)] + ) + local_positions = logical_positions.index_select(0, local_rows) + + def gather_new_kv(local_kv, forward_batch): + assert forward_batch is batch + torch.testing.assert_close( + local_kv, torch.cat([local_k.flatten(1), local_v.flatten(1)], dim=-1) + ) + return torch.cat([new_k.flatten(1), new_v.flatten(1)], dim=-1) + + monkeypatch.setattr( + qwen_sparse_attn_backend, "cp_materialize_global_token_order", gather_new_kv + ) + + def sparse_chunk_attention(q, k, v, indices, cu_q, cu_k, kv_lens, scale): + # The last local query in each block identifies its absolute causal end. + expected_ends = local_positions[cu_q[1:].long() - 1] + 1 + torch.testing.assert_close(kv_lens.long(), expected_ends) + assert k.dtype == v.dtype == q.dtype + return _cpu_sparse_chunk_attention(q, k, v, indices, cu_q, cu_k, kv_lens, scale) + + monkeypatch.setattr( + qwen_sparse_attn_backend, + "sparse_gqa_fwd_interface_triton_ck", + sparse_chunk_attention, + ) + layer = SimpleNamespace( + layer_id=0, + tp_q_head_num=2, + tp_k_head_num=1, + tp_v_head_num=1, + head_dim=4, + qk_head_dim=4, + v_head_dim=4, + scaling=0.5, + ) + # The indexer supplies causal logical positions and marks other slots invalid. + candidates = torch.arange(max(seq_lens), dtype=torch.int32).expand(num_local, -1) + topk = torch.where(candidates <= local_positions[:, None], candidates, -1) + actual = QwenSparseAttnBackend._forward_extend_cp( + backend, local_q, local_k, local_v, layer, batch, topk, save_kv_cache + ) + + expected = [] + row = 0 + for i, (length, prefix) in enumerate(zip(extend_lens, prefix_lens)): + for position in range(prefix, prefix + length): + visible_keys = keys[i][: position + 1, 0] + visible_values = values[i][: position + 1, 0] + scores = queries[row] @ visible_keys.T * layer.scaling + expected.append(scores.softmax(-1) @ visible_values) + row += 1 + expected = torch.stack(expected).index_select(0, local_rows).flatten(1) + torch.testing.assert_close(actual[:num_local], expected) + assert actual.shape == (8, 8) + torch.testing.assert_close(actual[num_local:], torch.zeros_like(actual[num_local:])) + if save_kv_cache: + pool.set_kv_buffer.assert_called_once() + initial_k[out_cache_loc] = new_k.to(cache_dtype) + initial_v[out_cache_loc] = new_v.to(cache_dtype) + else: + pool.set_kv_buffer.assert_not_called() + torch.testing.assert_close(k_cache.float(), initial_k.float(), equal_nan=True) + torch.testing.assert_close(v_cache.float(), initial_v.float(), equal_nan=True) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))