--- baked-ours/model_executor/layers/sparse_attn_indexer.py 2026-09-10 20:13:20.000000000 +0200 +++ baked/model_executor/layers/sparse_attn_indexer.py 2026-09-11 10:13:47.000000000 +0200 @@ -17,6 +17,9 @@ apply_candidate_mask as _apply_candidate_mask, ) from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + _PrefillTimer as _DSV41PrefillTimer, +) +from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( select_candidate_blocks as _select_candidate_blocks, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -325,6 +328,7 @@ candidate_blocks: torch.Tensor | None = None, candidate_block_size: int = 0, candidate_write: bool = False, + compress_ratio: int = 1, ) -> torch.Tensor: # careful! this will be None in dummy run forward_context = get_forward_context() @@ -332,13 +336,25 @@ fp8_dtype = current_platform.fp8_dtype() k_cache_prefix = _resolve_layer_name(k_cache_prefix) + # v4.1 two-level candidate filtering under DCP: the logits columns are this + # rank's local compressed states; with interleave 1 and world | block, local + # block b (size block // world) is exactly this rank's shard of global block + # b, so candidate ids need no translation. The producer MAX-reduces block + # scores across ranks and pins the global newest block (candidate_blocks.py). + local_cand_bs = candidate_block_size + dcp_group = None if candidate_blocks is not None: - # Candidate blocks are request-local; the DCP-sharded logits layout - # would need per-rank translation that is not implemented. - assert dcp_world_size == 1, ( - "v4.1 two-level candidate filtering is not supported with DCP." - ) assert candidate_block_size > 0 + if dcp_world_size > 1: + assert cp_kv_cache_interleave_size == 1, ( + "v4.1 candidate filtering under DCP needs cp_kv_cache_interleave_size=1" + ) + assert candidate_block_size % dcp_world_size == 0, ( + f"candidate_block_size ({candidate_block_size}) must be a multiple " + f"of dcp_world_size ({dcp_world_size})" + ) + local_cand_bs = candidate_block_size // dcp_world_size + dcp_group = get_dcp_group() # assert isinstance(attn_metadata, dict) if not isinstance(attn_metadata, dict): @@ -494,6 +510,23 @@ if chunk.local_total_seq_lens == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) topk_indices.fill_(-1) + if candidate_blocks is not None and candidate_write and dcp_group is not None: + # Other ranks may hold states for this chunk: join the MAX + # reduce with an empty local contribution. + _select_candidate_blocks( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + candidate_blocks.shape[1], + local_cand_bs, + candidate_blocks[chunk.token_start : chunk.token_end], + dcp_group=dcp_group, + global_row_ke=chunk.global_ctx_lens, + global_block_size=candidate_block_size, + max_global_blocks=triton.cdiv( + chunk.max_global_ctx, candidate_block_size + ), + ) else: # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. @@ -539,8 +572,16 @@ cu_seqlen_ks, cu_seqlen_ke, chunk_candidates.shape[1], - candidate_block_size, + local_cand_bs, chunk_candidates, + dcp_group=dcp_group, + global_row_ke=chunk.global_ctx_lens, + global_block_size=candidate_block_size, + max_global_blocks=( + triton.cdiv(chunk.max_global_ctx, candidate_block_size) + if dcp_group is not None + else None + ), ) else: _apply_candidate_mask( @@ -548,7 +589,7 @@ cu_seqlen_ks, cu_seqlen_ke, chunk_candidates, - candidate_block_size, + local_cand_bs, ) ops.top_k_per_row_prefill( logits, @@ -666,14 +707,44 @@ vis = vis[:num_rows] decode_candidates = candidate_blocks[:num_rows] if candidate_write: + global_row_ke = None + max_global_blocks = None + if dcp_group is not None: + # Global (un-sharded) per-row context in STATES for the + # newest-block pin: per token when flattened / next_n == 1, + # per request (row_repeat) or expanded per row otherwise. + gl = decode_metadata.global_seq_lens + assert gl is not None + gl = gl.reshape(-1) + if row_repeat > 1: + gl = gl[:batch_size] + elif gl.numel() != num_rows: + # native spec decode: request lens -> per-row lens + gl = ( + gl[:batch_size].unsqueeze(1) + - next_n + + 1 + + torch.arange(next_n, device=gl.device, dtype=gl.dtype) + ).clamp_(min=0).reshape(-1) + else: + gl = gl[:num_rows] + global_row_ke = gl // compress_ratio + max_global_blocks = triton.cdiv( + attn_metadata_narrowed.max_seq_len // compress_ratio, + candidate_block_size, + ) _select_candidate_blocks( logits, None, vis, decode_candidates.shape[1], - candidate_block_size, + local_cand_bs, decode_candidates, row_repeat, + dcp_group=dcp_group, + global_row_ke=global_row_ke, + global_block_size=candidate_block_size, + max_global_blocks=max_global_blocks, ) else: _apply_candidate_mask( @@ -681,7 +752,7 @@ None, vis, decode_candidates, - candidate_block_size, + local_cand_bs, row_repeat, ) topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] @@ -790,6 +861,7 @@ candidate_blocks: torch.Tensor | None = None, candidate_block_size: int = 0, candidate_write: bool = False, + compress_ratio: int = 1, ) -> torch.Tensor: return topk_indices_buffer @@ -904,6 +976,17 @@ k: torch.Tensor | None, weights: torch.Tensor, ): + # Stage 7: phase "indexer" of the prefill timing (inert unless enabled). + with _DSV41PrefillTimer("indexer", int(hidden_states.shape[0])): + return self._forward_cuda_impl(hidden_states, q_quant, k, weights) + + def _forward_cuda_impl( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ): # FP8 path: single tensor (per-token scale is folded into `weights`). # FP4 path: (values, scales) tuple with scales required by the kernel. if isinstance(q_quant, tuple): @@ -935,6 +1018,7 @@ candidate_blocks=self.candidate_blocks, candidate_block_size=self.candidate_block_size, candidate_write=self.candidate_write, + compress_ratio=self.compress_ratio, ) def forward_xpu(