# # Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. # Copyright 2024 The vLLM team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # This file is a part of the vllm-ascend project. # import math import sys from collections.abc import Callable from unittest.mock import MagicMock import pytest import torch if "torch_npu._inductor" not in sys.modules: sys.modules["torch_npu._inductor"] = MagicMock() from vllm.config import VllmConfig # noqa: E402 from vllm.forward_context import set_forward_context # noqa: E402 from tests.ut.attention.utils import ( # noqa: E402 BatchSpec, create_common_attn_metadata, create_vllm_config, ) from vllm_ascend.attention.context_parallel.sfa_cp import AscendSFACPImpl # noqa: E402 from vllm_ascend.utils import enable_custom_op enable_custom_op() SPARSE_COUNT = 2048 _BLOCK_SIZE = 128 _TEST_NUM_HEADS = 8 DEFAULT_RTOL = 1e-2 DEFAULT_ATOL = 1e-2 _MAX_SIG_REL_ERR = 1e-2 # max |out-ref| / peak |ref| _MAX_MEAN_SIG_ERR = 5e-3 # mean |out-ref| / mean |ref| _MAX_REL_ERR = 1e-2 # max per-element rel err where |ref| >= floor _SIG_FLOOR_FRAC = 5e-1 # floor = this fraction of peak |ref| def _validate_spec(spec: BatchSpec) -> None: """Require ``seq_len <= SPARSE_COUNT`` so sparse matches dense reference.""" for s, q in zip(spec.seq_lens, spec.query_lens): assert q <= s, f"query_len ({q}) must not exceed seq_len ({s})" assert s <= SPARSE_COUNT, ( f"seq_len ({s}) must be <= SPARSE_COUNT ({SPARSE_COUNT}) so the " "sparse attention degenerates into dense attention for the " "reference comparison." ) _VLLM_CONFIG_CACHE: dict = {} def _get_vllm_config( model: str, dtype: torch.dtype, *, max_model_len: int = 4096, tensor_parallel_size: int = 1, ) -> VllmConfig: key = (model, dtype, tensor_parallel_size) cfg = _VLLM_CONFIG_CACHE.get(key) if cfg is not None: return cfg dtype_str = "bfloat16" if dtype == torch.bfloat16 else "float16" sim_num_heads = max(1, _TEST_NUM_HEADS // tensor_parallel_size) cfg = create_vllm_config( model_name=model, tensor_parallel_size=1, # always TP=1; head split is simulated max_model_len=max_model_len, dtype=dtype_str, block_size=_BLOCK_SIZE, num_gpu_blocks=4096, max_num_seqs=64, max_num_batched_tokens=max(8192, max_model_len * 2), enable_chunked_prefill=True, hf_overrides={"quantization_config": None}, hf_config_override={ "num_attention_heads": sim_num_heads, "num_key_value_heads": 1, }, ) _VLLM_CONFIG_CACHE[key] = cfg return cfg # Spec name prefixes drive the SFA-CP branch under test. BATCH_SPECS: dict[str, BatchSpec] = { "decode_single": BatchSpec(seq_lens=[1024], query_lens=[1], name="decode_single"), "decode_small_batch": BatchSpec( seq_lens=[512, 1024, 1536, 2048], query_lens=[1, 1, 1, 1], name="decode_small_batch" ), "decode_large_batch": BatchSpec(seq_lens=[2048] * 8, query_lens=[1] * 8, name="decode_large_batch"), "mtp_1_plus_1": BatchSpec(seq_lens=[512, 1024, 1536], query_lens=[2, 2, 2], name="mtp_1_plus_1"), "mtp_1_plus_3": BatchSpec(seq_lens=[1024, 1536, 2048, 2048], query_lens=[4, 4, 4, 4], name="mtp_1_plus_3"), "mtp_1_plus_7": BatchSpec(seq_lens=[1024, 1536, 2048], query_lens=[8, 8, 8], name="mtp_1_plus_7"), "prefill_single": BatchSpec(seq_lens=[256], query_lens=[256], name="prefill_single"), "prefill_small_batch": BatchSpec(seq_lens=[256, 512, 384], query_lens=[256, 512, 384], name="prefill_small_batch"), "prefill_with_context": BatchSpec(seq_lens=[512, 1024], query_lens=[128, 256], name="prefill_with_context"), "mixed_small": BatchSpec(seq_lens=[512, 1024, 256, 512], query_lens=[1, 1, 64, 128], name="mixed_small"), "mixed_medium": BatchSpec( seq_lens=[1024, 1536, 2048, 256, 512], query_lens=[1, 1, 1, 64, 128], name="mixed_medium" ), } def _infer_mode(spec: BatchSpec) -> str: """Return one of: ``decode``, ``mtp``, ``prefill``, ``mixed``.""" name = spec.name for prefix in ("decode_", "mtp_", "prefill_", "mixed_"): if name.startswith(prefix): return prefix.rstrip("_") raise ValueError( f"BatchSpec name {name!r} does not start with a known mode prefix ('decode_', 'mtp_', 'prefill_', 'mixed_')" ) def _build_topk_indices( seq_lens: list[int], query_lens: list[int], sparse_count: int, device: torch.device, ) -> torch.Tensor: """Build causal topk indices with INVALID_IDX (-1) padding.""" num_tokens = sum(query_lens) topk = torch.full((num_tokens, 1, sparse_count), -1, dtype=torch.int32, device=device) cum_q = 0 for b, s_len in enumerate(seq_lens): q_len = query_lens[b] ctx_len = s_len - q_len for j in range(q_len): valid_end = ctx_len + j + 1 topk[cum_q + j, 0, :valid_end] = torch.arange(valid_end, dtype=torch.int32, device=device) cum_q += q_len return topk def _reference_sparse_attention( ql_nope: torch.Tensor, q_pe: torch.Tensor, k_nope_full_per_req: list[torch.Tensor], k_rope_full_per_req: list[torch.Tensor], seq_lens: list[int], query_lens: list[int], scale: float, out_dtype: torch.dtype, ) -> torch.Tensor: """Pure-PyTorch dense MQA baseline in fp32.""" outputs: list[torch.Tensor] = [] cum_q = 0 for b, s_len in enumerate(seq_lens): q_len = query_lens[b] ctx_len = s_len - q_len K_nope = k_nope_full_per_req[b][:s_len].float() K_rope = k_rope_full_per_req[b][:s_len].float() K = torch.cat([K_nope, K_rope], dim=-1) V = K_nope for j in range(q_len): t = cum_q + j valid_end = ctx_len + j + 1 q_n = ql_nope[t].float() q_p = q_pe[t].float() Q = torch.cat([q_n, q_p], dim=-1) K_b = K[:valid_end] V_b = V[:valid_end] scores = (Q @ K_b.transpose(0, 1)) * scale attn = torch.softmax(scores, dim=-1) out = attn @ V_b outputs.append(out.to(out_dtype)) cum_q += q_len return torch.stack(outputs, dim=0) def _build_cp_paged_kv_cache( seq_lens: list[int], k_nope_contexts: list[torch.Tensor], k_rope_contexts: list[torch.Tensor], block_size: int, kv_lora_rank: int, qk_rope_head_dim: int, cp_size: int, dtype: torch.dtype, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: """Allocate a paged KV cache that matches the *decode* CP gather/block_table layout. Each request is padded so its block count is exactly ``cp_size * L`` (i.e. divisible by the total CP world size). The local rank is treated as "rank 0" and owns the first ``L`` of each request's logical blocks (the actual values in ``local_block_table`` only matter via shape because ``gather_kv_cross_cp`` is mocked). The ``gathered_*`` tensors mirror what the impl would see after both DCP and PCP all-gathers; layout matches the indexing produced by ``gather_block_table``: ``gathered[i*N + r*L + b]`` holds K data for logical block ``b*cp_size + i`` of request ``r`` (where ``N == batch_size * L``). """ batch_size = len(seq_lens) raw_blocks_per_req = [(s + block_size - 1) // block_size for s in seq_lens] max_raw = max(raw_blocks_per_req) total_blocks_per_req = ((max_raw + cp_size - 1) // cp_size) * cp_size L = total_blocks_per_req // cp_size total_blocks = batch_size * total_blocks_per_req + 1 # +1 reserves block 0 full_k_nope_cache = torch.zeros(total_blocks, block_size, 1, kv_lora_rank, dtype=dtype, device=device) full_k_rope_cache = torch.zeros(total_blocks, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device) full_block_table = torch.zeros(batch_size, total_blocks_per_req, dtype=torch.int32, device=device) next_block = 1 for r in range(batch_size): s_len = seq_lens[r] for p in range(total_blocks_per_req): full_block_table[r, p] = next_block tok_start = p * block_size tok_end = min(tok_start + block_size, s_len) length = max(tok_end - tok_start, 0) if length > 0: full_k_nope_cache[next_block, :length, 0, :] = k_nope_contexts[r][tok_start:tok_end] full_k_rope_cache[next_block, :length, 0, :] = k_rope_contexts[r][tok_start:tok_end] next_block += 1 local_block_table = full_block_table[:, :L].contiguous() N = batch_size * L gathered_k_nope = torch.zeros(cp_size * N, block_size, 1, kv_lora_rank, dtype=dtype, device=device) gathered_k_rope = torch.zeros(cp_size * N, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device) for i in range(cp_size): for r in range(batch_size): for b in range(L): p_logical = b * cp_size + i phys_block = int(full_block_table[r, p_logical].item()) dst = i * N + r * L + b gathered_k_nope[dst] = full_k_nope_cache[phys_block] gathered_k_rope[dst] = full_k_rope_cache[phys_block] return ( full_k_nope_cache, full_k_rope_cache, full_block_table, local_block_table, gathered_k_nope, gathered_k_rope, L, ) def _build_cp_prefill_compact_metadata( prefill_full_block_table: torch.Tensor, full_k_nope_cache: torch.Tensor, full_k_rope_cache: torch.Tensor, cp_size: int, block_size: int, seq_lens: list[int], kv_lora_rank: int, qk_rope_head_dim: int, dtype: torch.dtype, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Build the prefill compact metadata + matching gathered KV view. Reproduces the production ``AscendSFACPMetadataBuilder.build_prefill_compact_block_metadata`` formula: ``block_table_cp[r, b*cp_size + i] = new_block_table[r, b] + i*M`` where ``M = num_unique_local_blocks``. Then materialises the ``gathered_compact`` tensor that ``gather_kv_cross_cp_compact`` would produce so that the kernel, when walking ``block_table_cp[r, :]``, sees the contiguous K context of request ``r`` block-by-block. Inputs ------ prefill_full_block_table: Shape ``(num_prefill_reqs, L * cp_size)``. The "full" (logical) block table for prefill requests as built by ``_build_cp_paged_kv_cache``. Entry ``[r, p]`` holds the physical block id that owns logical position ``p`` of request ``r``. The "local" (rank-0) prefill block table -- which the production code actually receives in ``attn_metadata.block_table[num_decodes:]`` -- is derived as ``prefill_full_block_table[:, :L]``. This matches the convention used by ``_build_cp_paged_kv_cache`` for the decode path. """ num_prefill_reqs, total_blocks_per_req = prefill_full_block_table.shape assert total_blocks_per_req % cp_size == 0, ( f"prefill_full_block_table has {total_blocks_per_req} columns, which is not divisible by cp_size={cp_size}" ) L = total_blocks_per_req // cp_size prefill_local_block_table = prefill_full_block_table[:, :L].contiguous() block_arange = torch.arange(cp_size, dtype=prefill_local_block_table.dtype, device=device) valid_block_ids, new_block_table_flat = prefill_local_block_table.flatten().unique(return_inverse=True) num_blocks = valid_block_ids.shape[0] block_table_cp = ( new_block_table_flat.unsqueeze(-1).to(prefill_local_block_table) + (block_arange * num_blocks).view(1, 1, -1).to(prefill_local_block_table) ).reshape(prefill_local_block_table.shape[0], -1) new_block_table_2d = new_block_table_flat.view(prefill_local_block_table.shape) M = int(num_blocks) gathered_compact_nope = torch.zeros(cp_size * M, block_size, 1, kv_lora_rank, dtype=dtype, device=device) gathered_compact_rope = torch.zeros(cp_size * M, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device) for r in range(num_prefill_reqs): s_len = seq_lens[r] for b_local in range(L): for i in range(cp_size): p_logical = b_local * cp_size + i tok_start = p_logical * block_size tok_end = min(tok_start + block_size, s_len) length = max(tok_end - tok_start, 0) if length <= 0: continue src_block = int(prefill_full_block_table[r, p_logical].item()) dst = int(new_block_table_2d[r, b_local].item()) + i * M gathered_compact_nope[dst, :length, 0, :] = full_k_nope_cache[src_block, :length, 0, :] gathered_compact_rope[dst, :length, 0, :] = full_k_rope_cache[src_block, :length, 0, :] return valid_block_ids, block_table_cp, gathered_compact_nope, gathered_compact_rope def _make_fake_self( *, scale: float, pcp_size: int, dcp_size: int, gather_kv_cross_cp_fn: Callable | None = None, gather_kv_cross_cp_compact_fn: Callable | None = None, ) -> MagicMock: """Construct a ``MagicMock`` matching the slice of ``AscendSFACPImpl`` that ``_execute_sparse_flash_attention_process`` reads. Bound methods that don't depend on init-time state (``gather_block_table``, ``_execute_sparse_flash_attention``) are delegated to the real (unbound) implementations so the production kernel call stays under test. The collective gathers (``gather_kv_cross_cp`` / ``gather_kv_cross_cp_compact``) are mocked via ``side_effect`` because they require an initialised PCP / DCP ``ProcessGroupHCCL``, which can't be created on a single rank. """ fake_self = MagicMock() fake_self.scale = scale fake_self.pcp_size = pcp_size fake_self.dcp_size = dcp_size if gather_kv_cross_cp_fn is not None: fake_self.gather_kv_cross_cp = MagicMock(side_effect=gather_kv_cross_cp_fn) if gather_kv_cross_cp_compact_fn is not None: fake_self.gather_kv_cross_cp_compact = MagicMock(side_effect=gather_kv_cross_cp_compact_fn) fake_self.gather_block_table = lambda block_num, block_tables, block_arange: AscendSFACPImpl.gather_block_table( fake_self, block_num, block_tables, block_arange ) fake_self._execute_sparse_flash_attention = lambda *args, **kwargs: AscendSFACPImpl._execute_sparse_flash_attention( fake_self, *args, **kwargs ) fake_self._align_to_graph_bucket_tokens = lambda x, m: x return fake_self def _run_sfa_cp_kernel( *, ql_nope: torch.Tensor, q_pe: torch.Tensor, full_k_nope_cache: torch.Tensor, full_k_rope_cache: torch.Tensor, local_block_table: torch.Tensor, topk_indices: torch.Tensor, cum_query_lens: torch.Tensor, seq_lens_tensor: torch.Tensor, scale: float, pcp_size: int, dcp_size: int, num_decodes: int, num_decode_tokens: int, num_prefills: int, device: torch.device, gathered_k_nope: torch.Tensor | None = None, gathered_k_rope: torch.Tensor | None = None, valid_block_ids: torch.Tensor | None = None, block_table_cp: torch.Tensor | None = None, gathered_compact_nope: torch.Tensor | None = None, gathered_compact_rope: torch.Tensor | None = None, prefill_q_cum_seqlens: torch.Tensor | None = None, ) -> torch.Tensor: """Run the SFA-CP kernel with mocked CP collectives.""" cp_size = pcp_size * dcp_size gather_kv_cross_cp_fn = None gather_kv_cross_cp_compact_fn = None if num_decode_tokens > 0: assert gathered_k_nope is not None and gathered_k_rope is not None, ( "decode branch requires gathered_k_nope/gathered_k_rope" ) gathered_lookup = { id(full_k_nope_cache): gathered_k_nope, id(full_k_rope_cache): gathered_k_rope, } def gather_kv_cross_cp_fn(kv: torch.Tensor, block_tables: torch.Tensor): gathered = gathered_lookup.get(id(kv)) assert gathered is not None, "gather_kv_cross_cp called with an unexpected kv tensor" return gathered, block_tables.numel() if num_prefills > 0: assert ( valid_block_ids is not None and block_table_cp is not None and gathered_compact_nope is not None and gathered_compact_rope is not None and prefill_q_cum_seqlens is not None ), ( "prefill branch requires valid_block_ids / block_table_cp / " "gathered_compact_{nope,rope} / prefill_q_cum_seqlens" ) gathered_compact_lookup = { id(full_k_nope_cache): gathered_compact_nope, id(full_k_rope_cache): gathered_compact_rope, } def gather_kv_cross_cp_compact_fn(kv: torch.Tensor, vbid: torch.Tensor): gathered = gathered_compact_lookup.get(id(kv)) assert gathered is not None, "gather_kv_cross_cp_compact called with an unexpected kv tensor" return gathered fake_self = _make_fake_self( scale=scale, pcp_size=pcp_size, dcp_size=dcp_size, gather_kv_cross_cp_fn=gather_kv_cross_cp_fn, gather_kv_cross_cp_compact_fn=gather_kv_cross_cp_compact_fn, ) fake_attn_metadata = MagicMock() fake_attn_metadata.block_table = local_block_table fake_attn_metadata.num_decodes = num_decodes fake_attn_metadata.num_decode_tokens = num_decode_tokens fake_attn_metadata.num_prefills = num_prefills fake_sfa_cp_metadata = MagicMock() fake_sfa_cp_metadata.block_arange = torch.arange(cp_size, dtype=torch.int32, device=device) if num_prefills > 0: fake_sfa_cp_metadata.valid_block_ids = valid_block_ids fake_sfa_cp_metadata.block_table_cp = block_table_cp fake_sfa_cp_metadata.prefill_q_cum_seqlens = prefill_q_cum_seqlens fake_attn_metadata.sfa_cp_metadata = fake_sfa_cp_metadata cum_lens_arg = cum_query_lens if num_decode_tokens > 0 else prefill_q_cum_seqlens return AscendSFACPImpl._execute_sparse_flash_attention_process( fake_self, ql_nope, q_pe, (full_k_nope_cache, full_k_rope_cache), topk_indices, fake_attn_metadata, cum_lens_arg, seq_lens_tensor, ) def _record_and_assert( backend_output: torch.Tensor, reference_output: torch.Tensor, tag: str, *, dtype: torch.dtype, atol: float = DEFAULT_ATOL, rtol: float = DEFAULT_RTOL, ) -> tuple[float, float]: """Assert numerical closeness and record signal-relative metrics.""" assert backend_output.shape == reference_output.shape, ( f"[{tag}] backend shape {tuple(backend_output.shape)} != reference shape {tuple(reference_output.shape)}" ) assert backend_output.dtype == reference_output.dtype, ( f"[{tag}] backend dtype {backend_output.dtype} != reference dtype {reference_output.dtype}" ) assert torch.isfinite(backend_output).all(), f"[{tag}] sparse flash attention produced non-finite values" torch.testing.assert_close( backend_output, reference_output, rtol=rtol, atol=atol, msg=lambda m: f"[SFA-CP:{tag}] kernel output diverges from baseline. {m}", ) ref_f32 = reference_output.float() out_f32 = backend_output.float() diff = (out_f32 - ref_f32).abs() ref_abs = ref_f32.abs() peak = float(ref_abs.max()) mean_ref_abs = float(ref_abs.mean()) sig_floor = peak * _SIG_FLOOR_FRAC max_abs_err = float(diff.max()) mean_abs_err = float(diff.mean()) max_sig_rel_err = max_abs_err / peak if peak > 0 else 0.0 mean_sig_rel_err = mean_abs_err / mean_ref_abs if mean_ref_abs > 0 else 0.0 significant_mask = ref_abs >= sig_floor if significant_mask.any(): per_elem_rel = diff[significant_mask] / ref_abs[significant_mask] max_rel_err_sig = float(per_elem_rel.max()) else: max_rel_err_sig = 0.0 assert max_sig_rel_err < _MAX_SIG_REL_ERR, ( f"[SFA-CP:{tag}] dtype={dtype} signal-relative max error " f"{max_sig_rel_err * 100:.4f}% exceeds 1% budget " f"(peak={peak:.4e}, max_abs_err={max_abs_err:.4e})" ) assert mean_sig_rel_err < _MAX_MEAN_SIG_ERR, ( f"[SFA-CP:{tag}] dtype={dtype} signal-relative mean error " f"{mean_sig_rel_err * 100:.4f}% exceeds 0.5% drift budget " f"(mean_ref_abs={mean_ref_abs:.4e}, mean_abs_err={mean_abs_err:.4e})" ) assert max_rel_err_sig < _MAX_REL_ERR, ( f"[SFA-CP:{tag}] dtype={dtype} per-element relative error on " f">={int(_SIG_FLOOR_FRAC * 100)}%-of-peak elements " f"{max_rel_err_sig * 100:.4f}% exceeds 1% budget " f"(peak={peak:.4e}, max_abs_err={max_abs_err:.4e})" ) return max_abs_err, max_rel_err_sig def _make_synthetic_kv_contexts( seq_lens: list[int], kv_lora_rank: int, qk_rope_head_dim: int, dtype: torch.dtype, device: torch.device, ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: """Per-request K context generated independently so values across requests are uncorrelated.""" k_nope = [torch.randn(s, kv_lora_rank, dtype=dtype, device=device) * 0.1 for s in seq_lens] k_rope = [torch.randn(s, qk_rope_head_dim, dtype=dtype, device=device) * 0.1 for s in seq_lens] return k_nope, k_rope def _test_sfa_cp_correctness( batch_spec: BatchSpec, model: str, *, pcp_size: int, dcp_size: int, dtype: torch.dtype = torch.bfloat16, atol: float = 1e-2, rtol: float = 1e-2, tensor_parallel_size: int = 1, ) -> None: """Test ``AscendSFACPImpl`` against a fp32 dense MQA reference.""" mode = _infer_mode(batch_spec) assert not (pcp_size > 1 and mode in ("prefill", "mixed")), ( f"PCP>1 {mode} is out of scope; the parametrize whitelist should not have generated this combination" ) torch.manual_seed(2026) _validate_spec(batch_spec) vllm_config = _get_vllm_config( model, dtype, tensor_parallel_size=tensor_parallel_size, ) device = torch.device("npu") seq_lens = list(batch_spec.seq_lens) query_lens = list(batch_spec.query_lens) batch_size = batch_spec.batch_size num_tokens = batch_spec.compute_num_tokens() cp_size = pcp_size * dcp_size cache_config = vllm_config.cache_config hf_text = vllm_config.model_config.hf_text_config block_size = cache_config.block_size kv_lora_rank = hf_text.kv_lora_rank qk_rope_head_dim = hf_text.qk_rope_head_dim num_heads = hf_text.num_attention_heads head_dim = kv_lora_rank + qk_rope_head_dim scale = 1.0 / math.sqrt(head_dim) common_attn_metadata = create_common_attn_metadata( batch_spec, block_size=block_size, device=device, ) k_nope_contexts, k_rope_contexts = _make_synthetic_kv_contexts( seq_lens, kv_lora_rank, qk_rope_head_dim, dtype, device, ) ( full_k_nope_cache, full_k_rope_cache, full_block_table, local_block_table, gathered_k_nope, gathered_k_rope, _L, ) = _build_cp_paged_kv_cache( seq_lens=seq_lens, k_nope_contexts=k_nope_contexts, k_rope_contexts=k_rope_contexts, block_size=block_size, kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, cp_size=cp_size, dtype=dtype, device=device, ) ql_nope = ( torch.randn( num_tokens, num_heads, kv_lora_rank, dtype=dtype, device=device, ) * 0.1 ) q_pe = ( torch.randn( num_tokens, num_heads, qk_rope_head_dim, dtype=dtype, device=device, ) * 0.1 ) topk_indices = _build_topk_indices(seq_lens, query_lens, SPARSE_COUNT, device) cum_query_lens = common_attn_metadata.query_start_loc[1:].to(torch.int32) seq_lens_tensor = common_attn_metadata.seq_lens.to(torch.int32) decode_gathered_k_nope: torch.Tensor | None = None decode_gathered_k_rope: torch.Tensor | None = None prefill_full_block_table: torch.Tensor | None = None if mode in ("decode", "mtp"): num_decodes = batch_size num_decode_tokens = num_tokens num_prefills = 0 decode_gathered_k_nope = gathered_k_nope decode_gathered_k_rope = gathered_k_rope n_decode_reqs = batch_size elif mode == "prefill": num_decodes = 0 num_decode_tokens = 0 num_prefills = batch_size prefill_full_block_table = full_block_table n_decode_reqs = 0 else: # mixed (chunked-prefill) n_decode_reqs = sum(1 for q in query_lens if q == 1) assert all(q == 1 for q in query_lens[:n_decode_reqs]), ( f"mixed spec {batch_spec.name}: decode requests must come first" ) assert all(q > 1 for q in query_lens[n_decode_reqs:]), ( f"mixed spec {batch_spec.name}: prefill requests must come after decode" ) num_decodes = n_decode_reqs num_decode_tokens = n_decode_reqs num_prefills = batch_size - n_decode_reqs prefill_full_block_table = full_block_table[n_decode_reqs:] L_decode = local_block_table[:n_decode_reqs].shape[1] N_decode = n_decode_reqs * L_decode decode_gathered_k_nope = torch.zeros( cp_size * N_decode, block_size, 1, kv_lora_rank, dtype=dtype, device=device, ) decode_gathered_k_rope = torch.zeros( cp_size * N_decode, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device, ) for i in range(cp_size): for r in range(n_decode_reqs): for b in range(L_decode): p_logical = b * cp_size + i phys_block = int(full_block_table[r, p_logical].item()) dst = i * N_decode + r * L_decode + b decode_gathered_k_nope[dst] = full_k_nope_cache[phys_block] decode_gathered_k_rope[dst] = full_k_rope_cache[phys_block] valid_block_ids = None block_table_cp = None gathered_compact_nope = None gathered_compact_rope = None prefill_q_cum_seqlens = None if num_prefills > 0: valid_block_ids, block_table_cp, gathered_compact_nope, gathered_compact_rope = ( _build_cp_prefill_compact_metadata( prefill_full_block_table=prefill_full_block_table, full_k_nope_cache=full_k_nope_cache, full_k_rope_cache=full_k_rope_cache, cp_size=cp_size, block_size=block_size, seq_lens=seq_lens[n_decode_reqs:], kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, dtype=dtype, device=device, ) ) if n_decode_reqs > 0: prefill_q_cum_seqlens = cum_query_lens[n_decode_reqs:] - cum_query_lens[n_decode_reqs - 1] else: prefill_q_cum_seqlens = cum_query_lens with set_forward_context(attn_metadata=None, vllm_config=vllm_config): backend_output = _run_sfa_cp_kernel( ql_nope=ql_nope, q_pe=q_pe, full_k_nope_cache=full_k_nope_cache, full_k_rope_cache=full_k_rope_cache, local_block_table=local_block_table, topk_indices=topk_indices, cum_query_lens=cum_query_lens, seq_lens_tensor=seq_lens_tensor, scale=scale, pcp_size=pcp_size, dcp_size=dcp_size, num_decodes=num_decodes, num_decode_tokens=num_decode_tokens, num_prefills=num_prefills, device=device, gathered_k_nope=decode_gathered_k_nope, gathered_k_rope=decode_gathered_k_rope, valid_block_ids=valid_block_ids, block_table_cp=block_table_cp, gathered_compact_nope=gathered_compact_nope, gathered_compact_rope=gathered_compact_rope, prefill_q_cum_seqlens=prefill_q_cum_seqlens, ) reference_output = _reference_sparse_attention( ql_nope=ql_nope, q_pe=q_pe, k_nope_full_per_req=k_nope_contexts, k_rope_full_per_req=k_rope_contexts, seq_lens=seq_lens, query_lens=query_lens, scale=scale, out_dtype=dtype, ) dt = "bf16" if dtype == torch.bfloat16 else "fp16" tag = f"{mode}|{batch_spec.name}|pcp={pcp_size}|dcp={dcp_size}|tp={tensor_parallel_size}|{dt}" _record_and_assert( backend_output, reference_output, tag, dtype=dtype, atol=atol, rtol=rtol, ) # Whitelist of ``(batch_spec_name, pcp_size, dcp_size)`` combinations that are # in scope for the single-rank precision suite. ``PCP > 1`` prefill / mixed # requires per-rank PCP scheduler metadata (``q_head_idx``, ``q_tail_idx``, # ``q_full_idx``, ...) that only a multi-rank job can produce faithfully -- so # those combos are intentionally absent from the matrix (covered structurally # by ``tests/ut/attention/test_sfa_cp.py``). _TOPOLOGIES_ALL = [(2, 2), (2, 1), (1, 2), (1, 4)] _TOPOLOGIES_PCP1 = [(1, 2), (1, 4)] _TEST_CASES: list[tuple[str, int, int]] = [ (name, pcp, dcp) for name in BATCH_SPECS for (pcp, dcp) in (_TOPOLOGIES_ALL if _infer_mode(BATCH_SPECS[name]) in ("decode", "mtp") else _TOPOLOGIES_PCP1) ] @pytest.mark.parametrize( "batch_spec_name,pcp_size,dcp_size", _TEST_CASES, ids=[f"{n}-pcp{p}-dcp{d}" for (n, p, d) in _TEST_CASES], ) @pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-V3.2-Exp"]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4]) def test_sfa_cp_correctness( batch_spec_name: str, model: str, pcp_size: int, dcp_size: int, dtype: torch.dtype, tensor_parallel_size: int, ) -> None: """Test SFA-CP correctness across workload, topology, dtype, and TP size.""" _test_sfa_cp_correctness( BATCH_SPECS[batch_spec_name], model, pcp_size=pcp_size, dcp_size=dcp_size, dtype=dtype, tensor_parallel_size=tensor_parallel_size, )