diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py index b0c378923bc5..ba5550f0adb6 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py @@ -1,6 +1,7 @@ -"""Candidate block selection of the two-level indexer: the per-row block counts -and sparse-row lengths, the block keys and the JIT block top-k, and the torch -block ids and masks.""" +"""Candidate blocks of the two-level indexer: a source keeps whole blocks of +``block_size`` compressed positions, its newest block always, and a consumer +selects among them. Block counts, block keys, the block top-k (JIT / torch), and +the torch top-k among chosen blocks.""" from typing import Optional, Union @@ -142,83 +143,64 @@ def amax_topk_blocks( return blocks -def mask_topk_scores( - scores: torch.Tensor, - indices: torch.Tensor, - offsets: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """Keep masked indexer scores out of attention even when top-k underfills.""" - columns = indices.to(torch.int64) - if offsets is not None: - columns = columns - offsets[:, None] - selected_scores = scores.gather(1, columns.clamp(0, scores.shape[1] - 1)) - valid = ( - (columns >= 0) & (columns < scores.shape[1]) & (selected_scores > -torch.inf) - ) - return indices.masked_fill(~valid, -1) - - -def _candidate_block_topk( +def select_candidate_block_ids( logits: torch.Tensor, compress_lens: Union[torch.Tensor, int], topk_blocks: int, block_size: int, -) -> torch.return_types.topk: +) -> torch.Tensor: + """Per row the ids of the ``topk_blocks`` blocks with the largest score, the + block of position ``compress_lens - 1`` always kept: int32 + ``[rows, min(topk_blocks, blocks)]``, unordered, ``-1`` past the finite blocks. + ``logits`` must be ``-inf`` past each row's causal length.""" width = logits.size(-1) padding = -width % block_size scores = F.pad(logits, (0, padding), value=-torch.inf) if padding else logits scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1) num_blocks = scores.size(-1) - last = (compress_lens - 1) // block_size scores = scores.masked_fill( torch.arange(num_blocks, device=logits.device) == last, torch.inf ) - - return scores.topk(min(topk_blocks, num_blocks), dim=-1) - - -def select_candidate_block_ids( - logits: torch.Tensor, - compress_lens: Union[torch.Tensor, int], - topk_blocks: int, - block_size: int, -) -> torch.Tensor: - top = _candidate_block_topk( - logits=logits, - compress_lens=compress_lens, - topk_blocks=topk_blocks, - block_size=block_size, - ) + top = scores.topk(min(topk_blocks, num_blocks), dim=-1) return top.indices.to(torch.int32).masked_fill_(~(top.values > -torch.inf), -1) -def candidate_block_mask( - blocks: torch.Tensor, width: int, block_size: int -) -> torch.Tensor: - num_blocks = (width + block_size - 1) // block_size - keep = torch.zeros( - (*blocks.shape[:-1], num_blocks + 1), dtype=torch.bool, device=blocks.device - ) - keep.scatter_(-1, blocks.to(torch.int64).masked_fill(blocks < 0, num_blocks), True) - return keep[..., :num_blocks].repeat_interleave(block_size, dim=-1)[..., :width] - - -def select_candidate_blocks( - logits: torch.Tensor, - compress_lens: Union[torch.Tensor, int], - topk_blocks: int, +def topk_among_blocks( + scores: torch.Tensor, + lens: torch.Tensor, + blocks: torch.Tensor, + k: int, block_size: int, ) -> torch.Tensor: - top = _candidate_block_topk( - logits=logits, - compress_lens=compress_lens, - topk_blocks=topk_blocks, - block_size=block_size, + """Top-``k`` of each row of ``scores`` ``[rows, width]`` within its ``blocks`` + ``[rows, n]`` (``-1`` padded) and its causal ``lens``: int64 ``[rows, k]`` + positions, unordered, ``-1`` where fewer than ``k`` candidates are finite.""" + rows, width = scores.shape + n = blocks.shape[1] + if width == 0 or n == 0: + return torch.full((rows, k), -1, dtype=torch.int64, device=scores.device) + if width % block_size: + # the DeepGEMM callers slice their tile block-aligned, so this is a no-op there + scores = F.pad(scores, (0, -width % block_size), value=-torch.inf) + num_blocks = scores.shape[1] // block_size + by_block = scores.unflatten(1, (num_blocks, block_size)) + blocks64 = blocks.to(torch.int64) + ids = blocks64.clamp(0, num_blocks - 1) + # [rows, n, block_size], the only scratch of the row count's size + candidates = by_block.gather(1, ids[:, :, None].expand(rows, n, block_size)) + # column j of block b is position b * block_size + j, kept if j < lens - b * block_size + offsets = torch.arange(block_size, device=scores.device, dtype=lens.dtype) + room = lens[:, None] - blocks64 * block_size + drop = (offsets[None, None, :] >= room[:, :, None]) | (blocks64 < 0)[:, :, None] + candidates.masked_fill_(drop, -torch.inf) + flat = candidates.flatten(1) + top = flat.topk(min(k, flat.shape[1]), dim=-1, sorted=False) + block_of_pick = top.indices // block_size + picked = blocks64.gather(1, block_of_pick) * block_size + ( + top.indices - block_of_pick * block_size ) - width = logits.shape[-1] - num_blocks = (width + block_size - 1) // block_size - keep = torch.zeros( - (*logits.shape[:-1], num_blocks), dtype=torch.bool, device=logits.device - ).scatter_(-1, top.indices, top.values > -torch.inf) - return keep.repeat_interleave(block_size, dim=-1)[..., :width] + picked = picked.masked_fill(top.values == -torch.inf, -1) + if picked.shape[1] < k: + picked = F.pad(picked, (0, k - picked.shape[1]), value=-1) + return picked diff --git a/python/sglang/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index 1a57689230be..a68cb99bb9cd 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/python/sglang/kernels/ops/attention/dsv4/topk.py @@ -331,13 +331,13 @@ def topk_transform_packed_v2( def topk_transform_sparse( logits: torch.Tensor, valid_lens: torch.Tensor, - phys_blocks: torch.Tensor, - page_indices: torch.Tensor, + blocks: torch.Tensor, + out_indices: torch.Tensor, ) -> None: - """Top-``k`` (``k = page_indices.shape[1]``) of each row of the bf16 sparse - ``logits`` within its first ``valid_lens[b]`` columns, as pool slots through - ``phys_blocks`` (the published blocks as pool slots / 8); ``-1`` padded, - unordered.""" + """Top-``k`` (``k = out_indices.shape[1]``) of each row of the bf16 sparse + ``logits`` within its first ``valid_lens[b]`` columns, ``-1`` padded, unordered; + column ``j`` is written as ``blocks[b, j // 8] * 8 + j % 8``: pool slots for the + published blocks as pool slots / 8, compressed positions for logical ids.""" topk_transform_bf16_small( - logits, valid_lens, phys_blocks, page_indices, CANDIDATE_BLOCK_SIZE + logits, valid_lens, blocks, out_indices, CANDIDATE_BLOCK_SIZE ) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 4a7255c494dd..ce8d926134ec 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -19,10 +19,6 @@ import torch import torch.nn.functional as F -from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( - mask_topk_scores, - select_candidate_blocks, -) from sglang.kernels.ops.attention.dsv4.decode_attention_sm100 import ( can_use_swapab_attention, ) @@ -33,10 +29,6 @@ gather_dequant_requant_fp8_paged, q8kv8_padded_num_heads, ) -from sglang.kernels.ops.attention.dsv4.fp4_indexer import fp4_index_logits_decode -from sglang.kernels.ops.attention.dsv4.index_logits import ( - deep_gemm_fp4_paged_mqa_logits, -) from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout from sglang.kernels.ops.attention.dsv4.metadata_kernel import ( fill_all_compressed_indices, @@ -45,7 +37,6 @@ init_compression_metadata as _init_compression_metadata_triton, ) from sglang.kernels.ops.attention.dsv4.online_c128_mtp import OnlineC128MTPController -from sglang.kernels.ops.attention.dsv4.topk import topk_transform_paged_torch from sglang.kernels.ops.attention.dsv4_attn_metadata_kernels import ( BuildCausalSwaPageIndices, BuildPageTablePositions, @@ -64,34 +55,21 @@ ) from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp -from sglang.srt.layers.attention.dsv4.candidate_indexer import ( - CandidateMasks, - CandidateMetadata, - IndexerInputs, - PrefillIndexerInputs, - cut_request_masks, - expand_index_page_table, - make_candidate_indexer, - published_masks, -) from sglang.srt.layers.attention.dsv4.compressor_v2 import ( CompressorBackendMixin, FusedCompressMetadata, create_paged_compressor_data, ) -from sglang.srt.layers.attention.dsv4.dense_prefill_indexer import plain_prefill_topk from sglang.srt.layers.attention.dsv4.dsv41_sparse import ( _rope_fq4, token_req_indices, ) -from sglang.srt.layers.attention.dsv4.indexer import ( - C4IndexerBackendMixin, - topk_transform_paged_from_metadata, -) +from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin from sglang.srt.layers.attention.dsv4.metadata import ( _LARGE_INDEXER_QUERY_THRESHOLD, PagedIndexerMetadata, copy_metadata, + expand_index_page_table, maybe_copy_inplace, ) from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( @@ -99,6 +77,17 @@ SparsePrefillWorkspace, use_dsv4_q8kv8_sparse_prefill, ) +from sglang.srt.layers.attention.dsv4.v41_indexer import ( + CandidateMetadata, + CapturedPrefillInputs, + DecodeInputs, + PrefillInputs, + Selection, + has_dense_fp4_indexer, + is_sm100_or_newer, + make_candidate_indexer, + make_full_topk_indexer, +) from sglang.srt.layers.attention.verify_mask import ( VerifyMask, maybe_create_verify_mask, @@ -151,12 +140,6 @@ PAGE_INDEX_ALIGNED_SIZE = 64 -@functools.lru_cache(maxsize=None) -def _is_sm100_or_newer() -> bool: - # DeepGEMM's fp8_fp4 mqa-logits kernels need SM100+; Hopper takes the torch indexer. - return torch.cuda.get_device_capability()[0] >= 10 - - def _get_logical_forward_mode(forward_batch: ForwardBatch) -> ForwardMode: # IDLE is a real per-DP-rank mode. Do not let a stale _original_forward_mode # from a reused/padded ForwardBatch turn an empty rank into TARGET_VERIFY. @@ -265,35 +248,6 @@ def _maybe_precompute_flashmla_sched_meta( flashmla_metadata.num_splits = num_splits -# Arbitrary cap on one bf16 [rows, heads, lc] score chunk; transients run ~3x this. -_TORCH_INDEXER_SCORE_BUDGET_BYTES = 1 << 30 - - -def _every_request_fits() -> bool: - from sglang.srt.model_executor.runner_utils.capture_mode import ( - get_capture_attention_variant, - ) - - # Captured only for batches where every request fits the candidate budget, so - # the plain top-k is the whole selection. - return get_capture_attention_variant() in ( - "candidate_all", - "candidate_c2_all", - "candidate_unfiltered", - ) - - -@functools.cache -def _has_dense_fp4_indexer() -> bool: - if not torch.cuda.is_available() or torch.version.cuda is None: - return False - try: - import deep_gemm - except ImportError: - return False - return hasattr(deep_gemm, "fp8_fp4_mqa_logits") - - def _low_ratio_source_projections(layer, x, q_lora, positions, bufs): from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( get_tc_piecewise_forward_context, @@ -1169,12 +1123,17 @@ def __init__( ) self.has_c4: bool = 4 in self.present_ratios self.has_c128: bool = 128 in self.present_ratios - # Two-level low-ratio indexer (dsv4/candidate_indexer.py). cfg = model_runner.model_config.hf_text_config self.is_dsv41: bool = getattr(cfg, "model_type", None) == "deepseek_v41" - self.candidate_indexer = make_candidate_indexer( - getattr(cfg, "candidate_topk_blocks", 0), - getattr(cfg, "candidate_block_size", 0), + self.prefill_candidates, self.decode_candidates = make_candidate_indexer( + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + page_size=self.page_size, + candidate_topk_blocks=getattr(cfg, "candidate_topk_blocks", 0), + candidate_block_size=getattr(cfg, "candidate_block_size", 0), + ) + self.full_topk_indexer = make_full_topk_indexer( + token_to_kv_pool=self.token_to_kv_pool, req_to_token=self.req_to_token ) self.MAX_SEQ_LEN_FOR_CAPTURE = self.req_to_token.shape[1] @@ -1540,9 +1499,7 @@ def _low_ratio_prefill_indexer_metadata( @property def low_ratio_prefill_graph(self) -> bool: - return ( - bool(self.low_ratios) and _has_dense_fp4_indexer() and _is_sm100_or_newer() - ) + return bool(self.low_ratios) and has_dense_fp4_indexer() and is_sm100_or_newer() def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool: max_seq_len = _prefill_graph_max_seq_len() @@ -2772,7 +2729,6 @@ def forward_low_ratio_sources( req, pos = hoisted_req, hoisted_pos else: req = token_req_indices(forward_batch, num_tokens=positions.shape[0]) - # Every consumer takes int32 or int64 positions; keep the caller's. pos = positions if ( forward_batch.forward_mode.is_extend() @@ -2785,9 +2741,7 @@ def forward_low_ratio_sources( layer, x, req, pos, projected=(bufs["kv"], bufs.get("score")) ) if run_indexer and layer.indexer is not None: - self._low_ratio_index_topk_prefill_graph( - layer, pos, bufs["q"], bufs["w"] - ) + self._low_ratio_index_topk_captured(layer, bufs["q"], bufs["w"]) return if run_compressor and layer.compressor is not None: self._low_ratio_compress(layer, x, req, pos, forward_batch) @@ -2819,14 +2773,15 @@ def _forward_low_ratio_sources_cp( )[:total] self._low_ratio_compress_torch(layer, x_global, req_global, pos_global) if run_indexer and layer.indexer is not None: - self._low_ratio_index_topk_dense( + # A rank's local rows are not the batch's, so no req_rows to pass. + self._low_ratio_index_topk( layer, x[:num_local], q_lora[:num_local], + None, positions[:num_local].to(torch.int64), forward_batch, - torch.tensor(q_lens_cpu, dtype=torch.int32, device=x.device), - q_lens_cpu, + rows_per_request=q_lens_cpu, ) def _low_ratio_compress(self, layer, x, req, pos, forward_batch) -> None: @@ -2884,9 +2839,7 @@ def _low_ratio_compress_decode(self, layer, x, req, pos) -> None: slots, pos, fuse_index_store=( - x.is_cuda - and torch.version.cuda is not None - and _is_sm100_or_newer() + x.is_cuda and torch.version.cuda is not None and is_sm100_or_newer() ), ) return @@ -2916,7 +2869,7 @@ def _low_ratio_compress_decode(self, layer, x, req, pos) -> None: pooled, slots, group_pos, - fuse_index_store=_is_sm100_or_newer(), + fuse_index_store=is_sm100_or_newer(), ) def _low_ratio_compress_fused(self, layer, x, req, pos, *, draft_len=1) -> None: @@ -3117,201 +3070,57 @@ def _low_ratio_write_group( layer_id=layer.layer_id, loc=slots, cache_k=latent ) - def _low_ratio_index_topk(self, layer, x, q_lora, req, pos, forward_batch) -> None: - is_decode_or_verify = ( - forward_batch.forward_mode.is_decode() - or forward_batch.forward_mode.is_target_verify() - ) - if is_decode_or_verify: - if _is_sm100_or_newer(): - # DeepGEMM pairs verify rows by request id; decode has one row each. - req_ids = None if forward_batch.forward_mode.is_decode() else req - self._low_ratio_index_topk_decode(layer, x, q_lora, pos, req_ids) - else: - self._low_ratio_index_topk_sm90_decode(layer, x, q_lora, req, pos) - elif ( - self._use_dense_fp4_prefill_indexer(forward_batch) and _is_sm100_or_newer() - ): - self._low_ratio_index_topk_extend(layer, x, q_lora, pos, forward_batch) - else: - self._low_ratio_index_topk_torch(layer, x, q_lora, req, pos) - - @staticmethod - def _use_dense_fp4_prefill_indexer(forward_batch) -> bool: - return ( - not envs.SGLANG_DSV41_TORCH_PREFILL_INDEXER.get() - and _has_dense_fp4_indexer() - and forward_batch.forward_mode.is_extend() - and forward_batch.seq_lens_cpu is not None - and forward_batch.extend_seq_lens_cpu is not None - ) - - def _low_ratio_index_topk_extend( - self, layer, x, q_lora, pos, forward_batch - ) -> None: - tail = self.forward_metadata.late_layer_tail - if tail is not None: - q_lens, q_lens_cpu = tail.extend_seq_lens, tail.extend_seq_lens_cpu - else: - q_lens = forward_batch.extend_seq_lens - q_lens_cpu = _as_int_list(forward_batch.extend_seq_lens_cpu) - assert q_lens_cpu is not None - self._low_ratio_index_topk_dense( - layer, x, q_lora, pos, forward_batch, q_lens, q_lens_cpu - ) - - def _low_ratio_index_topk_dense( - self, layer, x, q_lora, pos, forward_batch, q_lens, q_lens_cpu + def _low_ratio_index_topk( + self, + layer, + x, + q_lora, + req, + pos, + forward_batch, + *, + rows_per_request=None, # NOTE: only used in CP ) -> None: - from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( - quantize_fp4_indexer_tensor, - ) - - pool = self.token_to_kv_pool - core = self.forward_metadata.core_metadata + is_source = layer.indexer.is_candidate_source + is_consumer = layer.indexer.uses_candidates ratio = layer.compress_ratio - indexer = layer.indexer - page_indices = core.sparse_page_indices(ratio) - raw_indices = core.sparse_raw_indices(ratio) - page_indices.fill_(-1) - if raw_indices is not None: - raw_indices.fill_(-1) - - seq_lens_cpu = _as_int_list(forward_batch.seq_lens_cpu) - assert seq_lens_cpu is not None - device = pos.device - # Visible compressed positions per request at its newest token; the - # per-token count (pos + 1) // ratio bounds each row below. - lc_per_req = [s // ratio for s in seq_lens_cpu] - req_pool_indices = forward_batch.req_pool_indices.to(torch.int64) - slot_chunks, starts, start = [], [], 0 - for r, lc in enumerate(lc_per_req): - starts.append(start) - if lc == 0: - continue - j = torch.arange(lc, device=device) - slot_chunks.append( - self.req_to_token[req_pool_indices[r], j * ratio].to(torch.int64) - // ratio - ) - start += lc - num_tokens = pos.shape[0] - if not slot_chunks or num_tokens == 0: - if indexer.is_candidate_source: - # no row to publish for; the consumers return here as well - self.forward_metadata.candidate_metadata = None - return - k_slots = torch.cat(slot_chunks) - k_fp4, k_sf = pool.get_low_ratio_index_k_fp4(layer.layer_id, k_slots) - - q = indexer.queries(q_lora, layer.freqs_cis[pos]) # [T, H, 128] fp4 grid - num_heads = q.shape[1] - q_fp4, q_sf = quantize_fp4_indexer_tensor(q.flatten(0, 1), rne=True) - q_fp4 = q_fp4.view(num_tokens, num_heads, 64) - q_sf = q_sf.view(num_tokens, num_heads) - weights = indexer.head_weights(x).float() - compress_lens = ((pos + 1) // ratio).to(torch.int32) - ks = torch.repeat_interleave( - torch.tensor(starts, dtype=torch.int32, device=device), - q_lens.to(torch.int64), - output_size=num_tokens, - ) - inputs = self._prefill_indexer_inputs( - layer, - (q_fp4, q_sf), - (k_fp4, k_sf), - weights, - compress_lens, - ks, - lc_per_req, - list(q_lens_cpu), - ) - topk = indexer.index_topk - selected = torch.full((num_tokens, topk), -1, dtype=torch.int32, device=device) - if indexer.is_candidate_source: - self._publish_prefill( - self.candidate_indexer.publish_prefill(inputs, selected) - ) - elif indexer.uses_candidates: - self.candidate_indexer.select_prefill( - self.forward_metadata.candidate_metadata, inputs, selected - ) - else: - plain_prefill_topk(inputs, selected) - # ascending positions, padding last: the layout the consumers expect - unselected = torch.iinfo(torch.int32).max - selected = selected.masked_fill(selected < 0, unselected).sort(dim=-1).values - chosen = selected != unselected - page_indices[:num_tokens, :topk] = torch.where( - chosen, k_slots[selected.clamp_max(k_slots.shape[0] - 1)], -1 - ).to(torch.int32) - if raw_indices is not None: - raw_indices[:num_tokens, :topk] = torch.where( - chosen, selected - ks[:, None], -1 + out = self._get_low_ratio_selection(ratio) + published = self.forward_metadata.candidate_metadata + mode = forward_batch.forward_mode + if mode.is_decode() or mode.is_target_verify(): + from sglang.srt.model_executor.runner_utils.capture_mode import ( + skip_candidate_filtering, + skip_low_ratio_indexer, ) - def _prefill_indexer_inputs( - self, layer, q, kv, weights, compress_lens, ks, lc_per_req, q_lens_cpu - ) -> PrefillIndexerInputs: - pool = self.token_to_kv_pool - ratio = layer.compress_ratio - num_tokens = q[0].shape[0] - index_page_size = pool.get_index_k_page_size(ratio) - k_cache = pool.get_index_k_with_scale_buffer(layer.layer_id) - assert k_cache.dim() == 2 - return PrefillIndexerInputs( - q_fp4=q[0], - q_sf=q[1], - weights=weights, - compress_lens=compress_lens, - request_starts=ks, - lens_per_request=lc_per_req, - rows_per_request=q_lens_cpu, - kv=kv, - k_cache=k_cache.view(k_cache.shape[0], index_page_size, 1, 68), - page_size=index_page_size, - kv_page_table=self.forward_metadata.core_metadata.page_table[:num_tokens], - kv_page_size=self.page_size, - compress_ratio=ratio, - ) - - def _tail_lens_to_publish(self) -> Optional[List[int]]: - """Rows per request of the late-layer tail, when the layers after the - switch will consume what is published now.""" - tail_metadata = self.tail_forward_metadata - if tail_metadata is None or tail_metadata is self.forward_metadata: - return None - tail = tail_metadata.late_layer_tail - return ( - tail.local_lens_cpu - if tail.cp_metadata is not None - else tail.extend_seq_lens_cpu - ) + if skip_low_ratio_indexer(ratio): + return + if skip_candidate_filtering(): + is_source = is_consumer = False - def _publish_prefill(self, published: CandidateMetadata) -> None: - self.forward_metadata.candidate_metadata = published - tail_lens = self._tail_lens_to_publish() - if tail_lens is not None: - self.tail_forward_metadata.candidate_metadata = ( - self.candidate_indexer.prefill_tail(published, tail_lens) + inputs = self._make_low_ratio_decode_indexer_inputs( + layer, x, q_lora, req, pos, mode ) - - def _publish_prefill_masks(self, masks: CandidateMasks) -> None: - """The torch prefill path's inline masks (see the TODO on CandidateMasks).""" - self.forward_metadata.candidate_metadata = masks - tail_lens = self._tail_lens_to_publish() - if tail_lens is not None: - self.tail_forward_metadata.candidate_metadata = cut_request_masks( - masks, tail_lens + if is_source: + published = self.decode_candidates.publish_decode(inputs, out) + self._publish_candidate_metadata(published) + elif is_consumer: + self.decode_candidates.consume_decode(inputs, published, out) + else: + self.full_topk_indexer.topk_decode(inputs, out) + else: + inputs = self._make_low_ratio_prefill_indexer_inputs( + layer, x, q_lora, req, pos, forward_batch, rows_per_request ) + if is_source: + published = self.prefill_candidates.publish_prefill(inputs, out) + self._publish_candidate_metadata(published) + elif is_consumer: + self.prefill_candidates.consume_prefill(inputs, published, out) + else: + self.full_topk_indexer.topk_prefill(inputs, out) - def _low_ratio_index_topk_prefill_graph(self, layer, pos, q, w) -> None: - from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( - quantize_fp4_indexer_tensor, - ) - - pool = self.token_to_kv_pool - core = self.forward_metadata.core_metadata + def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> None: ratio = layer.compress_ratio indexer = layer.indexer metadata = ( @@ -3320,325 +3129,100 @@ def _low_ratio_index_topk_prefill_graph(self, layer, pos, q, w) -> None: else self.forward_metadata.c2_indexer_metadata ) assert metadata is not None, f"no prefill graph indexer metadata for {ratio = }" - assert indexer.n_local_heads == indexer.n_heads - width = metadata.max_compressed_seq_len if indexer.uses_candidates or indexer.is_candidate_source: - # Every reachable block is a candidate inside the window, so the - # two-level selection collapses to the plain top-k below. + # Inside the window every block is a candidate: the full top-k applies. + width = metadata.max_compressed_seq_len assert ( width <= indexer.candidate_topk_blocks * indexer.candidate_block_size ), f"prefill graph indexer width {width} exceeds the candidate window" - - num_tokens, num_heads = q.shape[0], q.shape[1] - q_fp4, q_sf = quantize_fp4_indexer_tensor(q.flatten(0, 1), rne=True) - q_fp4 = q_fp4.view(num_tokens, 1, num_heads, 64) - q_sf = q_sf.view(num_tokens, 1, num_heads) - weights = w.float() - - k_cache = pool.get_index_k_with_scale_buffer(layer.layer_id) - assert k_cache.dim() == 2 - page_size = metadata.compressed_page_size - k_cache = k_cache.view(k_cache.shape[0], page_size, 1, 68) - - lens = metadata.compressed_seq_lens - page_table = metadata.page_table - page_indices = core.sparse_page_indices(ratio) - raw_indices = core.sparse_raw_indices(ratio) - topk = min(indexer.index_topk, width) - for rows, plan in metadata.row_chunks(): - logits = deep_gemm_fp4_paged_mqa_logits( - (q_fp4[rows], q_sf[rows]), - k_cache, - weights[rows], - lens[rows], - page_table[rows], - plan, - width, - ) - topk_transform_paged_torch( - logits, - lens[rows], - page_table[rows], - page_indices[rows, :topk], - page_size, - raw_indices[rows, :topk] if raw_indices is not None else None, - ) - - def _low_ratio_index_topk_decode(self, layer, x, q_lora, pos, req=None) -> None: - from sglang.srt.model_executor.runner_utils.capture_mode import ( - skip_low_ratio_indexer, + inputs = CapturedPrefillInputs( + indexer=indexer, + layer_id=layer.layer_id, + q=projected_q, + weights=projected_w, + paged_metadata=metadata, ) + out = self._get_low_ratio_selection(ratio) + self.full_topk_indexer.topk_prefill_captured(inputs, out) - if skip_low_ratio_indexer(layer.compress_ratio): - # The compressor still writes index K for later, longer contexts. - return - - from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( - quantize_fp4_indexer_tensor, + def _get_low_ratio_selection(self, compress_ratio: int) -> Selection: + core = self.forward_metadata.core_metadata + return Selection( + page_indices=core.sparse_page_indices(compress_ratio), + raw_indices=core.sparse_raw_indices(compress_ratio), ) - pool = self.token_to_kv_pool - core = self.forward_metadata.core_metadata + def _make_low_ratio_decode_indexer_inputs(self, layer, x, q_lora, req, pos, mode): ratio = layer.compress_ratio - indexer = layer.indexer metadata = ( self.forward_metadata.c1_indexer_metadata if ratio == 1 else self.forward_metadata.c2_indexer_metadata ) assert metadata is not None, f"no decode indexer metadata for {ratio = }" + return DecodeInputs( + indexer=layer.indexer, + layer_id=layer.layer_id, + compress_ratio=ratio, + freqs_cis=layer.freqs_cis, + x=x, + q_lora=q_lora, + positions=pos, + req_rows=req, + paged_metadata=metadata, + is_verify=mode.is_target_verify(), + ) - # fp4 query as (payload, scale), kernel layout [bs, 1, n_heads, dim]. The - # kernel sums head scores locally, so the indexer heads must be replicated. - assert indexer.n_local_heads == indexer.n_heads - if ( - x.is_cuda - and torch.version.cuda is not None - and x.dtype == torch.bfloat16 - and indexer.index_head_dim == 128 - ): - from sglang.kernels.ops.attention.dsv4.fp4_indexer_rope import ( - index_q_rope_pack_weights, - ) - - q, _ = indexer.wq_b(q_lora) - q = q.view(q.shape[0], indexer.n_local_heads, indexer.index_head_dim) - # The fused pack also computes head_weights(x).float(), same rounding. - q_fp4, q_sf, weights = index_q_rope_pack_weights( - q, - torch.view_as_real(layer.freqs_cis).flatten(-2), - pos, - indexer.head_weights_raw(x), # [bs, n_local] bf16, n32k5120 - indexer.head_weight_scale, + def _make_low_ratio_prefill_indexer_inputs( + self, + layer, + x, + q_lora, + req, + pos, + forward_batch, + rows_per_request, + ) -> PrefillInputs: + tail = self.forward_metadata.late_layer_tail + if rows_per_request is not None: + rows_per_request_device = torch.tensor( + rows_per_request, dtype=torch.int32, device=x.device ) + elif tail is not None: + rows_per_request = tail.extend_seq_lens_cpu + rows_per_request_device = tail.extend_seq_lens else: - q = indexer.queries(q_lora, layer.freqs_cis[pos]) - q_fp4, q_sf = quantize_fp4_indexer_tensor(q.flatten(0, 1), rne=True) - weights = indexer.head_weights(x).float() # [bs, n_local] - bs = q.shape[0] - q_fp4 = q_fp4.view(bs, 1, indexer.n_local_heads, 64) - q_sf = q_sf.view(bs, 1, indexer.n_local_heads) - - k_cache = pool.get_index_k_with_scale_buffer(layer.layer_id) - assert k_cache.dim() == 2 - # Index pool page (64 slots); metadata.page_table is expanded to match. - page_size = metadata.compressed_page_size - k_cache = k_cache.view( - k_cache.shape[0], page_size, 1, 68 - ) # fp4: 64 payload + 4 scale - - page_indices = core.sparse_page_indices(ratio) - raw_indices = core.sparse_raw_indices(ratio) - inputs = IndexerInputs( - q_fp4, - q_sf, - k_cache, - weights, - metadata, - request_ids=req, # one per query row; verify rows of a request share one + rows_per_request = _as_int_list(forward_batch.extend_seq_lens_cpu) + rows_per_request_device = forward_batch.extend_seq_lens + return PrefillInputs( + indexer=layer.indexer, + layer_id=layer.layer_id, + compress_ratio=layer.compress_ratio, + freqs_cis=layer.freqs_cis, + x=x, + q_lora=q_lora, + positions=pos, + req_rows=req, + req_pool_indices=forward_batch.req_pool_indices, + kv_page_table=self.forward_metadata.core_metadata.page_table, + seq_lens_cpu=_as_int_list(forward_batch.seq_lens_cpu), + rows_per_request=rows_per_request, + rows_per_request_device=rows_per_request_device, ) - candidate_layer = not _every_request_fits() - # use special selection for candidate layers - if indexer.uses_candidates and candidate_layer: - return self.candidate_indexer.select_decode( - self.forward_metadata.candidate_metadata, - inputs, - page_indices, - raw_indices, - ) - if indexer.is_candidate_source and candidate_layer: - self.forward_metadata.candidate_metadata = ( - self.candidate_indexer.publish_decode(inputs, page_indices, raw_indices) - ) - return - if isinstance(metadata.deep_gemm_metadata, list): - topk_plans = metadata.topk_metadata_chunks - assert not metadata.use_topk_v2 or topk_plans is not None - for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()): - logits = deep_gemm_fp4_paged_mqa_logits( - (q_fp4[rows], q_sf[rows]), - k_cache, - weights[rows], - metadata.compressed_seq_lens[rows], - metadata.page_table[rows], - plan, - metadata.max_compressed_seq_len, - ) - # TODO(dark): add bf16 topk - topk_transform_paged_from_metadata( - logits, - metadata, - page_indices, - raw_indices, - rows=rows, - topk_metadata=( - topk_plans[chunk_idx] if topk_plans is not None else None - ), - ) - else: - logits = deep_gemm_fp4_paged_mqa_logits( - (q_fp4, q_sf), - k_cache, - weights, - metadata.compressed_seq_lens, - metadata.page_table, - metadata.deep_gemm_metadata, - metadata.max_compressed_seq_len, - ) - # TODO(dark): add bf16 topk - topk_transform_paged_from_metadata( - logits, metadata, page_indices, raw_indices - ) - # TODO(candidate): Hopper decode still publishes / consumes masks inline (torch - # top-k); move into the candidate indexer with the prefill paths. - def _low_ratio_index_topk_sm90_decode(self, layer, x, q_lora, req, pos) -> None: - pool = self.token_to_kv_pool - core = self.forward_metadata.core_metadata - ratio = layer.compress_ratio - indexer = layer.indexer - page_indices = core.sparse_page_indices(ratio) - raw_indices = core.sparse_raw_indices(ratio) - page_indices.fill_(-1) - if raw_indices is not None: - raw_indices.fill_(-1) - bs = req.shape[0] - assert pos.shape[0] == bs, ( - f"decode expects one token per request, {pos.shape=} {bs=}" - ) - if bs == 0: + def _publish_candidate_metadata(self, published: Optional[CandidateMetadata]): + if published is None: return - lens = (pos + 1) // ratio - metadata = ( - self.forward_metadata.c1_indexer_metadata - if ratio == 1 - else self.forward_metadata.c2_indexer_metadata - ) - assert metadata is not None - # V4 reserves the replay bound in metadata; visibility stays on device. - # A capture-time length read would both synchronize and truncate replay. - lmax = min(metadata.max_compressed_seq_len, self.req_to_token.shape[1] // ratio) - if lmax == 0: + self.forward_metadata.candidate_metadata = published + tail_metadata = self.tail_forward_metadata + if tail_metadata is None or tail_metadata is self.forward_metadata: return - q = indexer.queries(q_lora, layer.freqs_cis[pos]) - weights = indexer.head_weights(x) - j = torch.arange(lmax, device=pos.device) - valid = j[None, :] < lens[:, None] - slots = ( - self.req_to_token[req[:, None], (j * ratio)[None, :]].to(torch.int64) - // ratio - ) - slots = slots.masked_fill(~valid, 0) - table = pool.get_index_k_with_scale_buffer(layer.layer_id) - s = fp4_index_logits_decode( - q, weights, slots, lens, table, table.shape[1] // 68 - ) - if indexer.is_candidate_source: - mask = select_candidate_blocks( - s, - lens[:, None], - topk_blocks=indexer.candidate_topk_blocks, - block_size=indexer.candidate_block_size, - ) - self.forward_metadata.candidate_metadata = CandidateMasks(mask=mask) - elif indexer.uses_candidates: - # Published this step by the candidate-source layer's decode pass above. - consume = published_masks(self.forward_metadata.candidate_metadata).mask - assert torch.is_tensor(consume) and consume.shape[0] == bs, ( - "candidate mask missing for decode" - ) - s = s.masked_fill(~consume[:, :lmax], -torch.inf) - k = min(indexer.index_topk, lmax) - idx = s.topk(k, dim=-1, sorted=False).indices - if indexer.uses_candidates and not indexer.is_candidate_source: - idx = mask_topk_scores(s, idx) - idx = idx.masked_fill(idx < 0, lmax) - idx = idx.sort(dim=-1).values - reach = idx < lens[:, None] - page_indices[:bs, :k] = torch.where( - reach, slots.gather(1, idx.clamp_max(lmax - 1)), -1 - ).to(torch.int32) - if raw_indices is not None: - raw_indices[:bs, :k] = torch.where(reach, idx, -1).to(torch.int32) - - # TODO(candidate): torch prefill still publishes / consumes masks inline; same - # move as above. - def _low_ratio_index_topk_torch(self, layer, x, q_lora, req, pos) -> None: - pool = self.token_to_kv_pool - core = self.forward_metadata.core_metadata - ratio = layer.compress_ratio - indexer = layer.indexer - # Attention scans sparse_topk_lengths slots and skips -1 entries. - page_indices = core.sparse_page_indices(ratio) - raw_indices = core.sparse_raw_indices(ratio) - page_indices.fill_(-1) - if raw_indices is not None: - raw_indices.fill_(-1) - q = indexer.queries(q_lora, layer.freqs_cis[pos]) - weights = indexer.head_weights(x) - # A compressed position is visible once the query has passed its last token. - compress_lens = (pos + 1) // ratio - topk = indexer.index_topk - publish = [] if indexer.is_candidate_source else None - consume = ( - published_masks(self.forward_metadata.candidate_metadata).request_masks - if indexer.uses_candidates - else None + tail = tail_metadata.late_layer_tail + tail_metadata.candidate_metadata = published.tail( + tail.local_lens_cpu + if tail.cp_metadata is not None + else tail.extend_seq_lens_cpu ) - for b, r in enumerate(torch.unique_consecutive(req).tolist()): - tok = (req == r).nonzero().squeeze(1) - lens = compress_lens[tok] - lc = int(lens.max().item()) - if lc == 0: - # Consumers address masks by request position, including empty requests. - if publish is not None: - publish.append( - torch.zeros(0, 0, dtype=torch.bool, device=pos.device) - ) - continue - j = torch.arange(lc, device=pos.device) - slots_j = self.req_to_token[r, j * ratio].to(torch.int64) // ratio - # Dequantize only this request's visible K rows; the table is pool-sized. - index_k = pool.get_low_ratio_index_k_dequant(layer.layer_id, slots_j) - k = min(topk, lc) - # Every step below is per query row; chunk rows so the [rows, heads, lc] - # bf16 scores stay under the budget (16 GiB at once for a 16k-token prompt). - rows_per_chunk = max( - 1, - _TORCH_INDEXER_SCORE_BUDGET_BYTES // (q.shape[1] * lc * 2), - ) - masks = [] if publish is not None else None - for start in range(0, tok.numel(), rows_per_chunk): - rows = slice(start, start + rows_per_chunk) - tok_c, lens_c = tok[rows], lens[rows] - s = indexer.scores(q[tok_c], index_k, weights[tok_c]) - s = s.masked_fill(j[None, :] >= lens_c[:, None], -torch.inf) - if masks is not None: - masks.append( - select_candidate_blocks( - s, - lens_c[:, None], - topk_blocks=indexer.candidate_topk_blocks, - block_size=indexer.candidate_block_size, - ) - ) - elif consume is not None: - s = s.masked_fill(~consume[b][rows], -torch.inf) - idx = s.topk(k, dim=-1, sorted=False).indices - if consume is not None and masks is None: - idx = mask_topk_scores(s, idx) - idx = idx.masked_fill(idx < 0, lc) - idx = idx.sort(dim=-1).values - reach = idx < lens_c[:, None] - page_indices[tok_c, :k] = torch.where( - reach, slots_j[idx.clamp_max(lc - 1)], -1 - ).to(torch.int32) - if raw_indices is not None: - raw_indices[tok_c, :k] = torch.where(reach, idx, -1).to(torch.int32) - if masks is not None: - publish.append(torch.cat(masks) if len(masks) > 1 else masks[0]) - if publish is not None: - self._publish_prefill_masks(CandidateMasks(request_masks=publish)) def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor: """Idle always re-translates at store time: its metadata may be stale, and diff --git a/python/sglang/srt/layers/attention/dsv4/candidate_indexer.py b/python/sglang/srt/layers/attention/dsv4/candidate_indexer.py deleted file mode 100644 index 9a78e2618c68..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/candidate_indexer.py +++ /dev/null @@ -1,194 +0,0 @@ -from __future__ import annotations - -from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import List, Optional, Tuple - -import msgspec -import torch - -from sglang.srt.layers.attention.dsv4.metadata import PagedIndexerMetadata -from sglang.srt.runtime_context import get_parallel, get_platform - - -class CandidateMetadata: - """Base of an implementation's published state on - ``DSV4Metadata.candidate_metadata``.""" - - -@dataclass(frozen=True) -class IndexerInputs: - """One index-source layer's operands on the paged fp4 decode path (one query - row per request, or per draft token under verify).""" - - q_fp4: torch.Tensor # [rows, 1, heads, 64] int8, packed fp4 - q_sf: torch.Tensor # [rows, 1, heads] int32, packed ue8m0 - k_cache: torch.Tensor # [pages, page_size, 1, 68] uint8, the layer's index-K pool - weights: torch.Tensor # [rows, heads] bf16/fp32 head weights - metadata: PagedIndexerMetadata # this ratio's lengths, page table and plans - # [rows] int, one request id per query row, the rows of one request - # consecutive (verify: its draft tokens); None = every row its own request - request_ids: Optional[torch.Tensor] = None - - @property - def num_rows(self) -> int: - return self.q_fp4.shape[0] - - -class PrefillIndexerInputs(msgspec.Struct, frozen=True): - """One index layer's operands on the dense fp4 prefill path: the chunk's query - rows, one request's rows consecutive.""" - - q_fp4: torch.Tensor # [rows, heads, 64] int8, packed fp4 - q_sf: torch.Tensor # [rows, heads] int32, packed ue8m0 - weights: torch.Tensor # [rows, heads] fp32 head weights - compress_lens: torch.Tensor # [rows] int32, compressed positions the row sees - # [rows] int32: first flattened-K column of the row's request; selections - # are flattened-K columns, request start + compressed position - request_starts: torch.Tensor - lens_per_request: List[int] # compressed positions at each request's newest token - rows_per_request: List[int] # query rows of each request - # the flattened index K of the batch (int8 [n, 64], int32 [n]), column j of - # a row's dense scores against kv[request_starts[t] + j] - kv: Tuple[torch.Tensor, torch.Tensor] - k_cache: torch.Tensor # [pages, page_size, 1, 68] uint8, the layer's index-K pool - page_size: int # index-K pool page size, in slots - # [rows, kv pages] int32: the row's request in the KV pool, whose pages hold - # kv_page_size tokens; expand_index_page_table gives the index-K pool's - kv_page_table: torch.Tensor - kv_page_size: int - compress_ratio: int - - @property - def num_rows(self) -> int: - return self.q_fp4.shape[0] - - -def expand_index_page_table( - page_table: torch.Tensor, - *, - full_page_size: int, - compress_ratio: int, - index_page_size: int, -) -> torch.Tensor: - """Block table of a low-ratio indexer-K pool, which pages at `index_page_size` - slots: [bs, n] -> [bs, n * blocks_per_page] int32. The kernel reads compressed - slot j at page_table[b, j // index_page_size] * index_page_size + j % - index_page_size, which after this expansion is the c1/c2 pool slot of the same - position.""" - slots_per_page = full_page_size // compress_ratio - assert slots_per_page % index_page_size == 0, ( - f"{full_page_size = } / {compress_ratio = } must be a multiple of " - f"{index_page_size = }" - ) - blocks_per_page = slots_per_page // index_page_size - if blocks_per_page == 1: - return page_table - bs, n = page_table.shape - base = page_table.to(torch.int64) * blocks_per_page - offsets = torch.arange(blocks_per_page, device=page_table.device, dtype=torch.int64) - expanded = base.unsqueeze(-1) + offsets # [bs, n, blocks_per_page] - return expanded.reshape(bs, n * blocks_per_page).to(torch.int32) - - -class CandidateIndexer(ABC): - """The two-level low-ratio indexer. The candidate source publishes what its - consumers need to select among its top candidate blocks, on the forward - metadata; a consumer selects its top-k from that. - - TODO(candidate): Hopper decode and the torch prefill path still publish and - consume masks inline in the backend; their decode side is not behind this - protocol yet.""" - - @abstractmethod - def publish_prefill( - self, inputs: PrefillIndexerInputs, out_positions: torch.Tensor - ) -> CandidateMetadata: - """The source layer's own top-``k`` (``k = out_positions.shape[1]``) as - flattened-K columns, ``-1`` padded, unordered, plus its candidates for - the chunk: both come from one pass over its dense scores.""" - - @abstractmethod - def select_prefill( - self, - published: CandidateMetadata, - inputs: PrefillIndexerInputs, - out_positions: torch.Tensor, - ) -> None: - """A consumer's top-``k`` over its published candidates, in the same - layout as ``publish_prefill`` writes.""" - - @abstractmethod - def prefill_tail( - self, published: CandidateMetadata, tail_lens: List[int] - ) -> CandidateMetadata: - """The candidates of the last ``tail_lens[b]`` rows of each request, for - the late layers that run on the tail only.""" - - -def make_candidate_indexer( - topk_blocks: int, block_size: int -) -> Optional[CandidateIndexer]: - """DeepGEMM's paged sparse indexer on SM100, with the dense-score - implementation for prefill under context parallelism (a rank's local rows - are not the page table's rows); None on Hopper, whose decode and prefill - indexers select through masks inline.""" - if topk_blocks <= 0 or get_platform().device_sm < 100: - return None - from sglang.srt.layers.deep_gemm_wrapper.configurer import ( - DEEPGEMM_PAGED_SPARSE_MQA_LOGITS, - ) - - if not DEEPGEMM_PAGED_SPARSE_MQA_LOGITS: - raise RuntimeError( - "the candidate indexer needs DeepGEMM's paged sparse MQA logits " - "(sgl-deep-gemm >= 0.2.0 with SGLANG_ENABLE_JIT_DEEPGEMM on)" - ) - from sglang.srt.layers.attention.dsv4.candidate_indexer_deep_gemm import ( - DeepGemmCandidateIndexer, - ) - from sglang.srt.layers.attention.dsv4.dense_prefill_indexer import ( - DenseCandidateIndexer, - ) - - prefill_dense = None - if get_parallel().attn_cp_size > 1: - prefill_dense = DenseCandidateIndexer(topk_blocks, block_size) - return DeepGemmCandidateIndexer( - topk_blocks, block_size, prefill_dense=prefill_dense - ) - - -# TODO(candidate): Hopper decode and the torch prefill path still select through -# these masks inline in the backend. -@dataclass -class CandidateMasks(CandidateMetadata): - mask: Optional[torch.Tensor] = None # decode: [rows, width] bool - request_masks: Optional[List[torch.Tensor]] = None # prefill: [rows_b, lc_b] each - - -def cut_request_masks(masks: CandidateMasks, tail_lens: List[int]) -> CandidateMasks: - """The last ``tail_lens[b]`` rows of each request's mask.""" - assert masks.request_masks is not None, "prefill masks missing" - return CandidateMasks( - request_masks=[ - mask[mask.shape[0] - t :] for mask, t in zip(masks.request_masks, tail_lens) - ] - ) - - -class PrefillCandidateBlocks(CandidateMetadata, msgspec.Struct): - request_blocks: List[torch.Tensor] - - def tail(self, lengths: List[int]) -> PrefillCandidateBlocks: - return PrefillCandidateBlocks( - request_blocks=[ - blocks[blocks.shape[0] - length :] - for blocks, length in zip(self.request_blocks, lengths) - ] - ) - - -def published_masks(candidate) -> CandidateMasks: - assert isinstance(candidate, CandidateMasks), "candidate masks missing" - return candidate diff --git a/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py b/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py deleted file mode 100644 index 2cb85a5bb93a..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py +++ /dev/null @@ -1,439 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import List, Optional - -import torch - -from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( - amax8_varlen, - amax_topk_blocks, - candidate_row_lens, -) -from sglang.kernels.ops.attention.dsv4.candidate_table import ( - CANDIDATE_BLOCK_SIZE, - build_sparse_indexer_schedule, - sort_candidate_blocks, -) -from sglang.kernels.ops.attention.dsv4.index_logits import ( - deep_gemm_fp4_paged_mqa_logits, - flat_index_logits_tiles, - sparse_logits, -) -from sglang.kernels.ops.attention.dsv4.topk import ( - topk_transform_bf16_small, - topk_transform_paged_v2, - topk_transform_ragged_v2, - topk_transform_sparse, -) -from sglang.srt.layers.attention.dsv4.candidate_indexer import ( - CandidateIndexer, - CandidateMetadata, - IndexerInputs, - PrefillIndexerInputs, - expand_index_page_table, -) -from sglang.srt.layers.attention.dsv4.dense_prefill_indexer import ( - _SCORE_BUDGET_BYTES, - DenseCandidateIndexer, -) -from sglang.srt.layers.attention.dsv4.indexer import ( - topk_transform_paged_from_metadata, -) - - -@dataclass -class SparseBlockTable(CandidateMetadata): - # [rows, topk_blocks] int32: ascending logical block ids, valid for the first - # min(topk_blocks, ceil(seq_len / 8)) entries of a row; DeepGEMM reads only those - blocks: torch.Tensor - # DeepGEMM's schedule metadata (uint8) for them - schedule: torch.Tensor - # [rows, topk_blocks] int32: the same blocks as pool slots / 8, so a consumer's - # top-k maps column j of the sparse row to slot phys_blocks[b, j // 8] * 8 + j % 8 - # with the plain page-table transform at page size 8 - phys_blocks: torch.Tensor - # [rows] int32: length of each row of the sparse logits: the published blocks - # laid out block by block, the newest possibly partial (`candidate_row_lens`) - valid_lens: torch.Tensor - # recorded on the side stream once the fields above are complete - ready: torch.cuda.Event - - -@dataclass -class PrefillSparseBlockTable(SparseBlockTable): - """Published by a prefill chunk's candidate source: one row per query token. - The inputs stay attached so the late-layer tail can rebuild the table for - its rows (the DeepGEMM schedule is per row set, it cannot be sliced).""" - - compress_lens: torch.Tensor # [rows] int32 - page_table: torch.Tensor # [rows, index pages] int32 - request_ids: torch.Tensor # [rows] int32 - q_dtype: torch.dtype - page_size: int # index-K pool page size (slots) - rows_per_request: List[int] # rows of each request, in row order - - -# TODO(dark): support fusion of publish + topk of publish layer -class DeepGemmCandidateIndexer(CandidateIndexer): - def __init__( - self, - topk_blocks: int, - block_size: int, - prefill_dense: Optional[DenseCandidateIndexer] = None, - ): - assert block_size == CANDIDATE_BLOCK_SIZE, block_size - self.topk_blocks = topk_blocks - self.block_size = block_size - self.alt_stream = torch.cuda.Stream() - self._row_ids: Optional[torch.Tensor] = None - self._retired_row_ids: list = [] # captured graphs keep reading the buffers they saw - # set for the forwards whose prefill rows the page table does not describe - self._prefill_dense = prefill_dense - - def _request_ids( - self, request_ids: Optional[torch.Tensor], rows: int, device: torch.device - ) -> torch.Tensor: - """int32 ``[rows]``; None means every row is its own request, served from a - cached ``arange`` (grown in steps of 8192) so that case costs no launch.""" - if request_ids is not None: - # the scheduler keeps request indices as int64; one small cast per publish - return request_ids[:rows].to(torch.int32).contiguous() - buf = self._row_ids - if buf is None or buf.numel() < rows: - assert not torch.cuda.is_current_stream_capturing(), ( - f"row-id buffer grows to {rows} rows inside a CUDA graph capture; " - "warm up with the largest row count first" - ) - if buf is not None: - self._retired_row_ids.append(buf) - size = max(8192, -(-rows // 8192) * 8192) - buf = self._row_ids = torch.arange(size, dtype=torch.int32, device=device) - return buf[:rows] - - def publish_decode( - self, - inputs: IndexerInputs, - page_indices: torch.Tensor, - raw_indices: Optional[torch.Tensor] = None, - ) -> SparseBlockTable: - """The publishing layer's own plain top-k into ``page_indices``, plus the - block table for the consumers; the backend stores it on the forward - metadata.""" - metadata = inputs.metadata - seq_lens = metadata.compressed_seq_lens.reshape(-1) - if isinstance(metadata.deep_gemm_metadata, list): - return self._publish_decode_chunked( - inputs, page_indices, raw_indices, seq_lens - ) - logits = deep_gemm_fp4_paged_mqa_logits( - (inputs.q_fp4, inputs.q_sf), - inputs.k_cache, - inputs.weights, - metadata.compressed_seq_lens, - metadata.page_table, - metadata.deep_gemm_metadata, - metadata.max_compressed_seq_len, - ) - main_stream = torch.cuda.current_stream() - self.alt_stream.wait_stream(main_stream) - # The block-selection chain reads logits after the main stream moves on. - logits.record_stream(self.alt_stream) - # TODO(candidate): one kernel for both selections below (dense logits read once) - topk_transform_paged_v2( - logits, - seq_lens, - metadata.page_table, - page_indices, - metadata.compressed_page_size, - metadata.topk_metadata, - out_raw_indices=raw_indices, - ) - # per row: block count for the block top-k, sparse-row length for the consumers - with torch.cuda.stream(self.alt_stream): - nblocks, row_valid_lens = candidate_row_lens(seq_lens, self.topk_blocks) - blocks = amax_topk_blocks(logits, seq_lens, nblocks, self.topk_blocks) - # in place: ascending, INT32_MAX padded, plus the blocks as pool slots / 8 - phys_blocks = sort_candidate_blocks( - blocks, - seq_lens, - metadata.page_table, - metadata.compressed_page_size, - ) - schedule = build_sparse_indexer_schedule( - blocks, - seq_lens, - metadata.page_table, - metadata.compressed_page_size, - inputs.q_fp4.dtype, - self._request_ids(inputs.request_ids, inputs.num_rows, blocks.device), - ) - # select_decode reads these on the main stream - for t in (blocks, schedule, phys_blocks, row_valid_lens): - t.record_stream(main_stream) - ready = torch.cuda.Event() - ready.record(self.alt_stream) - return SparseBlockTable( - blocks=blocks, - schedule=schedule, - phys_blocks=phys_blocks, - valid_lens=row_valid_lens, - ready=ready, - ) - - def _publish_decode_chunked( - self, - inputs: IndexerInputs, - page_indices: torch.Tensor, - raw_indices: Optional[torch.Tensor], - seq_lens: torch.Tensor, - ) -> SparseBlockTable: - """Publish an eager forward whose dense logits are bounded by row chunks. - - CUDA-graph metadata always carries one tensor schedule and keeps using the - asynchronous fast path above. The exceptional eager path stays on the - current stream so each chunk's full logits can be released before the next. - """ - metadata = inputs.metadata - block_chunks = [] - phys_block_chunks = [] - valid_len_chunks = [] - topk_plans = metadata.topk_metadata_chunks - assert not metadata.use_topk_v2 or topk_plans is not None - - for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()): - logits = deep_gemm_fp4_paged_mqa_logits( - (inputs.q_fp4[rows], inputs.q_sf[rows]), - inputs.k_cache, - inputs.weights[rows], - metadata.compressed_seq_lens[rows], - metadata.page_table[rows], - plan, - metadata.max_compressed_seq_len, - ) - topk_transform_paged_from_metadata( - logits, - metadata, - page_indices, - raw_indices, - rows=rows, - topk_metadata=( - topk_plans[chunk_idx] if topk_plans is not None else None - ), - ) - - chunk_seq_lens = seq_lens[rows] - nblocks, row_valid_lens = candidate_row_lens( - chunk_seq_lens, self.topk_blocks - ) - blocks = amax_topk_blocks(logits, chunk_seq_lens, nblocks, self.topk_blocks) - phys_blocks = sort_candidate_blocks( - blocks, - chunk_seq_lens, - metadata.page_table[rows], - metadata.compressed_page_size, - ) - block_chunks.append(blocks) - phys_block_chunks.append(phys_blocks) - valid_len_chunks.append(row_valid_lens) - - blocks = torch.cat(block_chunks) - phys_blocks = torch.cat(phys_block_chunks) - row_valid_lens = torch.cat(valid_len_chunks) - schedule = build_sparse_indexer_schedule( - blocks, - seq_lens, - metadata.page_table, - metadata.compressed_page_size, - inputs.q_fp4.dtype, - self._request_ids(inputs.request_ids, inputs.num_rows, blocks.device), - ) - ready = torch.cuda.Event() - ready.record(torch.cuda.current_stream()) - return SparseBlockTable( - blocks=blocks, - schedule=schedule, - phys_blocks=phys_blocks, - valid_lens=row_valid_lens, - ready=ready, - ) - - def _scores(self, table: SparseBlockTable, inputs: IndexerInputs) -> torch.Tensor: - return sparse_logits( - inputs.q_fp4, - inputs.q_sf, - inputs.k_cache, - inputs.weights.to(torch.bfloat16), - table.schedule, - table.blocks.shape[1], - ) - - def select_decode( - self, - candidate_metadata: SparseBlockTable, - inputs: IndexerInputs, - page_indices: torch.Tensor, - raw_indices: Optional[torch.Tensor] = None, - ) -> None: - assert raw_indices is None - table = candidate_metadata - torch.cuda.current_stream().wait_event(table.ready) - logits = self._scores(table, inputs) - # decode carries no raw_indices; the kernel writes slots only - topk_transform_sparse(logits, table.valid_lens, table.phys_blocks, page_indices) - - def publish_prefill( - self, inputs: PrefillIndexerInputs, out_positions: torch.Tensor - ) -> CandidateMetadata: - """The source layer's own top-k and the block table of the chunk from - one pass over its dense scores, tile by tile: per tile the plain top-k - into ``out_positions`` and one read for the block keys.""" - if self._prefill_dense is not None: - return self._prefill_dense.publish_prefill(inputs, out_positions) - rows = inputs.num_rows - device = inputs.q_fp4.device - nblocks, valid_lens = candidate_row_lens(inputs.compress_lens, self.topk_blocks) - blocks = torch.empty(rows, self.topk_blocks, dtype=torch.int32, device=device) - # the block keys read the score rows through 32-byte vectors - for tile, logits in flat_index_logits_tiles( - q=(inputs.q_fp4, inputs.q_sf), - kv=inputs.kv, - weights=inputs.weights, - starts=inputs.request_starts, - lengths=inputs.compress_lens, - context_lengths=inputs.lens_per_request, - budget_bytes=_SCORE_BUDGET_BYTES, - width_align=8, - ): - lens = inputs.compress_lens[tile] - topk_transform_ragged_v2( - logits, - lens, - out_offsets=inputs.request_starts[tile], - out_indices=out_positions[tile], - ) - # keys past a row's block count stay unset; the top-k reads a row - # up to its block count only (v2 wants the stride a multiple of 4) - width_blocks = ( - logits.shape[1] + CANDIDATE_BLOCK_SIZE - 1 - ) // CANDIDATE_BLOCK_SIZE - keys = logits.new_empty(logits.shape[0], (width_blocks + 3) // 4 * 4) - amax8_varlen(logits, lens, out=keys) - topk_transform_ragged_v2( - keys, - nblocks[tile], - out_offsets=torch.zeros( - logits.shape[0], dtype=torch.int32, device=device - ), - out_indices=blocks[tile], - ) - request_ids = torch.repeat_interleave( - torch.arange( - len(inputs.rows_per_request), dtype=torch.int32, device=device - ), - torch.tensor(inputs.rows_per_request, dtype=torch.int64, device=device), - output_size=rows, - ) - # index-K pool pages of each row's request - page_table = expand_index_page_table( - inputs.kv_page_table, - full_page_size=inputs.kv_page_size, - compress_ratio=inputs.compress_ratio, - index_page_size=inputs.page_size, - ).contiguous() - return self._prefill_table( - blocks, - inputs.compress_lens, - page_table, - inputs.page_size, - request_ids, - inputs.rows_per_request, - inputs.q_fp4.dtype, - valid_lens, - ) - - def _prefill_table( - self, - blocks, - compress_lens, - page_table, - page_size, - request_ids, - rows_per_request, - q_dtype, - valid_lens, - ) -> PrefillSparseBlockTable: - # in place: ascending, INT32_MAX padded, plus the blocks as pool slots / 8 - phys_blocks = sort_candidate_blocks( - blocks, compress_lens, page_table, page_size - ) - schedule = build_sparse_indexer_schedule( - blocks, compress_lens, page_table, page_size, q_dtype, request_ids - ) - ready = torch.cuda.Event() - ready.record() - return PrefillSparseBlockTable( - blocks=blocks, - schedule=schedule, - phys_blocks=phys_blocks, - valid_lens=valid_lens, - ready=ready, - compress_lens=compress_lens, - page_table=page_table, - request_ids=request_ids, - q_dtype=q_dtype, - page_size=page_size, - rows_per_request=rows_per_request, - ) - - def prefill_tail( - self, published: CandidateMetadata, tail_lens: List[int] - ) -> CandidateMetadata: - if self._prefill_dense is not None: - return self._prefill_dense.prefill_tail(published, tail_lens) - table = published - assert isinstance(table, PrefillSparseBlockTable), "prefill block table missing" - rows, start = [], 0 - for n, t in zip(table.rows_per_request, tail_lens): - rows.append(torch.arange(start + n - t, start + n)) - start += n - rows = torch.cat(rows).to(table.blocks.device) - compress_lens = table.compress_lens[rows].contiguous() - _, valid_lens = candidate_row_lens(compress_lens, self.topk_blocks) - return self._prefill_table( - table.blocks[rows].contiguous(), - compress_lens, - table.page_table[rows].contiguous(), - table.page_size, - table.request_ids[rows].contiguous(), - tail_lens, - table.q_dtype, - valid_lens, - ) - - def select_prefill( - self, - published: CandidateMetadata, - inputs: PrefillIndexerInputs, - out_positions: torch.Tensor, - ) -> None: - if self._prefill_dense is not None: - return self._prefill_dense.select_prefill(published, inputs, out_positions) - table = published - assert isinstance(table, PrefillSparseBlockTable), "prefill block table missing" - rows, heads = inputs.q_sf.shape - logits = sparse_logits( - inputs.q_fp4.view(rows, 1, heads, 64), - inputs.q_sf.view(rows, 1, heads), - inputs.k_cache, - inputs.weights.to(torch.bfloat16), - table.schedule, - table.blocks.shape[1], - ) - # request-relative compressed positions, -1 padded - topk_transform_bf16_small( - logits, table.valid_lens, table.blocks, out_positions, CANDIDATE_BLOCK_SIZE - ) - out_positions.add_( - torch.where(out_positions >= 0, inputs.request_starts[:, None], 0) - ) diff --git a/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py b/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py deleted file mode 100644 index bbde5ce9a4f8..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py +++ /dev/null @@ -1,234 +0,0 @@ -from __future__ import annotations - -from typing import List, Optional - -import torch - -from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( - candidate_block_mask, - mask_topk_scores, - select_candidate_block_ids, -) -from sglang.kernels.ops.attention.dsv4.index_logits import ( - flat_index_logits_rows_per_tile, -) -from sglang.srt.layers.attention.dsv4.candidate_indexer import ( - CandidateIndexer, - CandidateMetadata, - PrefillCandidateBlocks, - PrefillIndexerInputs, -) -from sglang.srt.utils.common import ceil_align - -# TODO: use a per-forward mqa_logits_budget_bytes() budget that also -# leaves room for candidate masks and block-selection scratch. -_SCORE_BUDGET_BYTES = 2 << 30 - - -def dense_prefill_topk( - *, - q: tuple[torch.Tensor, torch.Tensor], - kv: tuple[torch.Tensor, torch.Tensor], - weights: torch.Tensor, - starts: torch.Tensor, - lengths: torch.Tensor, - request_lengths: list[tuple[int, int]], - topk: int, - candidate_topk_blocks: int, - candidate_block_size: int, - publish_candidates: bool, - candidates: PrefillCandidateBlocks | None, -) -> tuple[torch.Tensor, PrefillCandidateBlocks | None]: - selected = torch.full( - (q[0].shape[0], topk), -1, dtype=torch.int32, device=weights.device - ) - published = ( - PrefillCandidateBlocks(request_blocks=[]) if publish_candidates else None - ) - request_ranges = [] - row = 0 - for query_length, context_length in request_lengths: - request_ranges.append((row, row + query_length, context_length)) - if published is not None: - num_blocks = ( - context_length + candidate_block_size - 1 - ) // candidate_block_size - published.request_blocks.append( - torch.empty( - (query_length, min(candidate_topk_blocks, num_blocks)), - dtype=torch.int32, - device=weights.device, - ) - ) - row += query_length - width = ceil_align(max((n for _, n in request_lengths), default=0), 4) - if row == 0 or width == 0: - return selected, published - rows_per_chunk = flat_index_logits_rows_per_tile( - row, width, heads=q[0].shape[1], budget_bytes=_SCORE_BUDGET_BYTES - ) - for offset in range(0, row, rows_per_chunk): - rows = slice(offset, min(offset + rows_per_chunk, row)) - _select_tile( - q=(q[0][rows], q[1][rows]), - kv=kv, - weights=weights[rows], - starts=starts[rows], - lengths=lengths[rows], - width=width, - selected=selected[rows], - block_size=candidate_block_size, - row_offset=offset, - request_ranges=request_ranges, - publish=published, - consume=candidates, - ) - return selected, published - - -def _dense_topk( - inputs: PrefillIndexerInputs, - out_positions: torch.Tensor, - *, - topk_blocks: int, - block_size: int, - publish: bool, - candidates: Optional[PrefillCandidateBlocks], -) -> Optional[PrefillCandidateBlocks]: - selected, published = dense_prefill_topk( - q=(inputs.q_fp4, inputs.q_sf), - kv=inputs.kv, - weights=inputs.weights, - starts=inputs.request_starts, - lengths=inputs.compress_lens, - request_lengths=list(zip(inputs.rows_per_request, inputs.lens_per_request)), - topk=out_positions.shape[1], - candidate_topk_blocks=topk_blocks, - candidate_block_size=block_size, - publish_candidates=publish, - candidates=candidates, - ) - out_positions.copy_(selected) - return published - - -def plain_prefill_topk( - inputs: PrefillIndexerInputs, out_positions: torch.Tensor -) -> None: - """The top-k of an index layer outside the candidate scheme, over its dense - scores tile by tile.""" - _dense_topk( - inputs, - out_positions, - topk_blocks=0, - block_size=1, - publish=False, - candidates=None, - ) - - -class DenseCandidateIndexer(CandidateIndexer): - """Candidates as block ids per request (``PrefillCandidateBlocks``): the - source keeps its best blocks from its dense scores, a consumer masks its own - dense scores to -inf outside them and runs the plain top-k, tile by tile. - The prefill implementation for the CP layout; Hopper still runs the same - selection inline in the backend.""" - - def __init__(self, topk_blocks: int, block_size: int): - self.topk_blocks = topk_blocks - self.block_size = block_size - - def publish_prefill( - self, inputs: PrefillIndexerInputs, out_positions: torch.Tensor - ) -> PrefillCandidateBlocks: - published = _dense_topk( - inputs, - out_positions, - topk_blocks=self.topk_blocks, - block_size=self.block_size, - publish=True, - candidates=None, - ) - assert published is not None - return published - - def select_prefill( - self, - published: CandidateMetadata, - inputs: PrefillIndexerInputs, - out_positions: torch.Tensor, - ) -> None: - assert isinstance(published, PrefillCandidateBlocks), "candidate blocks missing" - _dense_topk( - inputs, - out_positions, - topk_blocks=self.topk_blocks, - block_size=self.block_size, - publish=False, - candidates=published, - ) - - def prefill_tail( - self, published: CandidateMetadata, tail_lens: List[int] - ) -> PrefillCandidateBlocks: - assert isinstance(published, PrefillCandidateBlocks), "candidate blocks missing" - return published.tail(tail_lens) - - -def _select_tile( - *, - q: tuple[torch.Tensor, torch.Tensor], - kv: tuple[torch.Tensor, torch.Tensor], - weights: torch.Tensor, - starts: torch.Tensor, - lengths: torch.Tensor, - width: int, - selected: torch.Tensor, - block_size: int, - row_offset: int, - request_ranges: list[tuple[int, int, int]], - publish: PrefillCandidateBlocks | None, - consume: PrefillCandidateBlocks | None, -) -> None: - from deep_gemm import fp8_fp4_mqa_logits - - from sglang.kernels.ops.attention.dsv4.topk import topk_transform_ragged_v2 - - logits = fp8_fp4_mqa_logits(q, kv, weights, starts, starts + lengths, False, width) - if publish is not None or consume is not None: - for request, (start, end, context_length) in enumerate(request_ranges): - begin, stop = max(start, row_offset), min(end, row_offset + logits.shape[0]) - if begin >= stop or context_length == 0: - continue - rows = slice(begin - row_offset, stop - row_offset) - request_rows = slice(begin - start, stop - start) - scores = logits[rows, :context_length] - if publish is not None: - lens = lengths[rows, None] - scores.masked_fill_( - torch.arange(context_length, device=logits.device)[None, :] >= lens, - -torch.inf, - ) - blocks = publish.request_blocks[request][request_rows] - blocks.copy_( - select_candidate_block_ids( - logits=scores, - compress_lens=lens, - topk_blocks=blocks.shape[1], - block_size=block_size, - ) - ) - else: - scores.masked_fill_( - ~candidate_block_mask( - blocks=consume.request_blocks[request][request_rows], - width=context_length, - block_size=block_size, - ), - -torch.inf, - ) - topk_transform_ragged_v2(logits, lengths, out_offsets=starts, out_indices=selected) - if consume is not None: - selected.copy_( - mask_topk_scores(scores=logits, indices=selected, offsets=starts) - ) diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index 019850d05756..bd850d38b342 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -388,3 +388,30 @@ def maybe_copy_inplace(dst, *, src) -> None: assert type(src) == type(dst) if dst is not None: dst.copy_(src) + + +def expand_index_page_table( + page_table: torch.Tensor, + *, + full_page_size: int, + compress_ratio: int, + index_page_size: int, +) -> torch.Tensor: + """Block table of a low-ratio indexer-K pool, which pages at `index_page_size` + slots: [bs, n] -> [bs, n * blocks_per_page] int32. The kernel reads compressed + slot j at page_table[b, j // index_page_size] * index_page_size + j % + index_page_size, which after this expansion is the c1/c2 pool slot of the same + position.""" + slots_per_page = full_page_size // compress_ratio + assert slots_per_page % index_page_size == 0, ( + f"{full_page_size = } / {compress_ratio = } must be a multiple of " + f"{index_page_size = }" + ) + blocks_per_page = slots_per_page // index_page_size + if blocks_per_page == 1: + return page_table + bs, n = page_table.shape + base = page_table.to(torch.int64) * blocks_per_page + offsets = torch.arange(blocks_per_page, device=page_table.device, dtype=torch.int64) + expanded = base.unsqueeze(-1) + offsets # [bs, n, blocks_per_page] + return expanded.reshape(bs, n * blocks_per_page).to(torch.int32) diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py new file mode 100644 index 000000000000..60f6fc690561 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py @@ -0,0 +1,134 @@ +"""The top-k selection of the DeepSeek V4.1 ratio-1/2 index layers: the full top-k +(``full_topk``), and the candidate schemes split by what a source publishes, a +DeepGEMM sparse table (``sparse_table``) or dense block ids (``dense_blocks``).""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Tuple + +import torch + +from .full_topk import FullTopKIndexer +from .types import ( + CandidateMetadata, + CapturedPrefillInputs, + DecodeCandidates, + DecodeInputs, + PrefillCandidates, + PrefillInputs, + Selection, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + +__all__ = [ + "CandidateMetadata", + "CapturedPrefillInputs", + "DecodeCandidates", + "DecodeInputs", + "FullTopKIndexer", + "PrefillCandidates", + "PrefillInputs", + "Selection", + "has_dense_fp4_indexer", + "is_sm100_or_newer", + "make_candidate_indexer", + "make_full_topk_indexer", +] + + +@functools.cache +def is_sm100_or_newer() -> bool: + # DeepGEMM's fp8_fp4 mqa-logits kernels need SM100+; Hopper takes the torch indexer. + return torch.cuda.get_device_capability()[0] >= 10 + + +@functools.cache +def has_dense_fp4_indexer() -> bool: + if not torch.cuda.is_available() or torch.version.cuda is None: + return False + try: + import deep_gemm + except ImportError: + return False + return hasattr(deep_gemm, "fp8_fp4_mqa_logits") + + +@functools.cache +def _use_deep_gemm_prefill() -> bool: + # A CP rank scores rank-local rows, which only DeepGEMM serves, so CP ignores the env. + from sglang.srt.environ import envs + from sglang.srt.runtime_context import get_parallel + + return ( + is_sm100_or_newer() + and has_dense_fp4_indexer() + and ( + get_parallel().attn_cp_size > 1 + or not envs.SGLANG_DSV41_TORCH_PREFILL_INDEXER.get() + ) + ) + + +def make_full_topk_indexer( + *, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, +) -> FullTopKIndexer: + return FullTopKIndexer( + token_to_kv_pool=token_to_kv_pool, + req_to_token=req_to_token, + use_deep_gemm_prefill=_use_deep_gemm_prefill(), + use_deep_gemm_decode=is_sm100_or_newer(), + ) + + +def make_candidate_indexer( + *, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, + page_size: int, + candidate_topk_blocks: int, + candidate_block_size: int, +) -> Tuple[PrefillCandidates, DecodeCandidates]: + from sglang.srt.runtime_context import get_parallel + + from .dense_blocks import DenseBlocksBackend + + use_deep_gemm_prefill = _use_deep_gemm_prefill() + dense_blocks = DenseBlocksBackend( + token_to_kv_pool=token_to_kv_pool, + req_to_token=req_to_token, + candidate_topk_blocks=candidate_topk_blocks, + candidate_block_size=candidate_block_size, + use_deep_gemm_prefill=use_deep_gemm_prefill, + ) + # Without candidate blocks no layer is a candidate source or consumer. + if not is_sm100_or_newer() or candidate_topk_blocks <= 0: + return dense_blocks, dense_blocks + + from sglang.srt.layers.deep_gemm_wrapper.configurer import ( + DEEPGEMM_PAGED_SPARSE_MQA_LOGITS, + ) + + if not DEEPGEMM_PAGED_SPARSE_MQA_LOGITS: + raise RuntimeError( + "the candidate indexer needs DeepGEMM's paged sparse MQA logits " + "(sgl-deep-gemm >= 0.2.0 with SGLANG_ENABLE_JIT_DEEPGEMM on)" + ) + from .sparse_table import SparseTableBackend + + sparse_table = SparseTableBackend( + token_to_kv_pool=token_to_kv_pool, + req_to_token=req_to_token, + page_size=page_size, + candidate_topk_blocks=candidate_topk_blocks, + candidate_block_size=candidate_block_size, + ) + # A CP rank's rows are an interleaved subset of the batch; the torch prefill + # env keeps the sparse table on decode only. + if not use_deep_gemm_prefill or get_parallel().attn_cp_size > 1: + return dense_blocks, sparse_table + return sparse_table, sparse_table diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py new file mode 100644 index 000000000000..3cc44e2bfd25 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py @@ -0,0 +1,345 @@ +"""The candidate scheme carried by dense block ids: a source publishes the ids of +its best blocks, a consumer takes its top-k among them.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Optional, Tuple + +import msgspec +import torch + +from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( + select_candidate_block_ids, + topk_among_blocks, +) + +from .scoring import ( + DeepGEMMPrefillData, + decode_scores, + get_deep_gemm_prefill_data, + prefill_requests, + score_tiles, + write_decode, + write_prefill, +) +from .types import ( + CandidateMetadata, + DecodeInputs, + PrefillInputs, + Selection, + get_tail_row_indices, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + + +class BlockIds(CandidateMetadata, msgspec.Struct): + # [rows, n] int32 block ids of each query row, unordered, -1 padded; n is the + # largest kept count of the batch, at most topk_blocks + blocks: torch.Tensor + # prefill: query rows of each request in row order, for the late-layer tail + rows_per_request: Optional[List[int]] = None + + def tail(self, rows_per_request: List[int]) -> BlockIds: + assert self.rows_per_request is not None, "prefill block ids missing" + rows = get_tail_row_indices( + full_rows_per_request=self.rows_per_request, + tail_rows_per_request=rows_per_request, + device=self.blocks.device, + ) + return BlockIds(blocks=self.blocks[rows], rows_per_request=rows_per_request) + + +class DenseBlocksBackend: + def __init__( + self, + *, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, + candidate_topk_blocks: int, + candidate_block_size: int, + use_deep_gemm_prefill: bool, + ) -> None: + self.token_to_kv_pool = token_to_kv_pool + self.req_to_token = req_to_token + self.topk_blocks = candidate_topk_blocks + self.block_size = candidate_block_size + self.use_deep_gemm_prefill = use_deep_gemm_prefill + + def publish_prefill(self, inputs: PrefillInputs, out: Selection): + if self.use_deep_gemm_prefill: + return self._deep_gemm_publish_prefill(inputs, out) + return self._torch_publish_prefill(inputs, out) + + def consume_prefill( + self, + inputs: PrefillInputs, + published: Optional[BlockIds], + out: Selection, + ) -> None: + if self.use_deep_gemm_prefill: + self._deep_gemm_consume_prefill(inputs, published, out) + else: + self._torch_consume_prefill(inputs, published, out) + + def publish_decode(self, inputs: DecodeInputs, out: Selection): + d = decode_scores( + inputs=inputs, + out=out, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + ) + if d is None: + return None + published = BlockIds( + blocks=select_candidate_block_ids( + d.scores, + d.lens[:, None], + topk_blocks=self.topk_blocks, + block_size=self.block_size, + ) + ) + k = min(inputs.indexer.index_topk, d.lmax) + idx = d.scores.topk(k, dim=-1, sorted=False).indices + write_decode(out, d, idx) + return published + + def consume_decode( + self, + inputs: DecodeInputs, + published: Optional[BlockIds], + out: Selection, + ) -> None: + d = decode_scores( + inputs=inputs, + out=out, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + ) + if d is None: + return + assert published is not None and published.blocks.shape[0] == d.bs + k = min(inputs.indexer.index_topk, d.lmax) + idx = topk_among_blocks( + d.scores, d.lens, published.blocks, k, block_size=self.block_size + ) + write_decode(out, d, idx.masked_fill(idx < 0, d.lmax)) + + # ---------- DeepGEMM: flattened-K scores ---------- + + def _deep_gemm_publish_prefill(self, inputs: PrefillInputs, out: Selection): + out.reset() + data = get_deep_gemm_prefill_data(inputs, self.req_to_token) + if data is None: + return None + selected, blocks = _publish_prefill_blocks( + data=data, + kv=self.token_to_kv_pool.get_low_ratio_index_k_fp4( + inputs.layer_id, data.k_slots + ), + topk=inputs.indexer.index_topk, + topk_blocks=self.topk_blocks, + block_size=self.block_size, + ) + data.write_selection(selected=selected, out=out) + return BlockIds(blocks=blocks, rows_per_request=data.rows_per_request) + + def _deep_gemm_consume_prefill( + self, + inputs: PrefillInputs, + published: Optional[BlockIds], + out: Selection, + ) -> None: + out.reset() + data = get_deep_gemm_prefill_data(inputs, self.req_to_token) + if data is None: + return + assert published is not None + selected = _consume_prefill_blocks( + data=data, + kv=self.token_to_kv_pool.get_low_ratio_index_k_fp4( + inputs.layer_id, data.k_slots + ), + topk=inputs.indexer.index_topk, + blocks=published.blocks, + block_size=self.block_size, + ) + data.write_selection(selected=selected, out=out) + + # ---------- torch: per-request scores ---------- + + def _torch_publish_prefill(self, inputs: PrefillInputs, out: Selection): + picked = [] # (query rows, their block ids) per scored chunk + for request, chunks in prefill_requests( + inputs=inputs, + out=out, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + ): + for chunk in chunks: + picked.append( + ( + chunk.tok, + select_candidate_block_ids( + chunk.scores, + chunk.lens[:, None], + topk_blocks=self.topk_blocks, + block_size=self.block_size, + ), + ) + ) + idx = chunk.scores.topk(request.k, dim=-1, sorted=False).indices + write_prefill(out, request, chunk, idx) + width = max((ids.shape[1] for _, ids in picked), default=0) + blocks = torch.full( + (inputs.positions.shape[0], width), + -1, + dtype=torch.int32, + device=inputs.positions.device, + ) + for rows, ids in picked: + blocks[rows, : ids.shape[1]] = ids + return BlockIds(blocks=blocks, rows_per_request=inputs.rows_per_request) + + def _torch_consume_prefill( + self, + inputs: PrefillInputs, + published: Optional[BlockIds], + out: Selection, + ) -> None: + assert published is not None + for request, chunks in prefill_requests( + inputs=inputs, + out=out, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + ): + for chunk in chunks: + idx = topk_among_blocks( + chunk.scores, + chunk.lens, + published.blocks[chunk.tok], + request.k, + block_size=self.block_size, + ) + write_prefill(out, request, chunk, idx.masked_fill(idx < 0, request.lc)) + + +def _publish_prefill_blocks( + *, + data: DeepGEMMPrefillData, + kv: Tuple[torch.Tensor, torch.Tensor], + topk: int, + topk_blocks: int, + block_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + from sglang.kernels.ops.attention.dsv4 import topk_transform_ragged_v2 + + selected = data.empty_selection(topk) + width = min(topk_blocks, -(-max(data.lens_per_request, default=0) // block_size)) + blocks = torch.full( + (data.num_rows, width), -1, dtype=torch.int32, device=selected.device + ) + for tile, logits in score_tiles(data, kv, width_align=4): + _publish_tile_blocks( + data=data, + tile=tile, + logits=logits, + blocks=blocks[tile], + topk_blocks=topk_blocks, + block_size=block_size, + ) + topk_transform_ragged_v2( + logits, + data.compress_lens[tile], + out_offsets=data.request_starts[tile], + out_indices=selected[tile], + ) + # Free this tile's logits before the generator scores the next one. + del logits + return selected, blocks + + +def _publish_tile_blocks( + *, + data: DeepGEMMPrefillData, + tile: slice, + logits: torch.Tensor, + blocks: torch.Tensor, + topk_blocks: int, + block_size: int, +) -> None: + lens = data.compress_lens[tile] + for rows, lc in _requests_in_tile(data, tile): + scores = logits[rows, :lc] + scores.masked_fill_( + torch.arange(lc, device=logits.device)[None, :] >= lens[rows, None], + -torch.inf, + ) + ids = select_candidate_block_ids( + logits=scores, + compress_lens=lens[rows, None], + topk_blocks=topk_blocks, + block_size=block_size, + ) + blocks[rows, : ids.shape[1]] = ids + + +def _consume_prefill_blocks( + *, + data: DeepGEMMPrefillData, + kv: Tuple[torch.Tensor, torch.Tensor], + topk: int, + blocks: torch.Tensor, + block_size: int, +) -> torch.Tensor: + selected = data.empty_selection(topk) + for tile, logits in score_tiles(data, kv, width_align=block_size): + _consume_tile_blocks( + data=data, + tile=tile, + logits=logits, + blocks=blocks[tile], + block_size=block_size, + out=selected[tile], + ) + # Free this tile's logits before the generator scores the next one. + del logits + return selected + + +def _consume_tile_blocks( + *, + data: DeepGEMMPrefillData, + tile: slice, + logits: torch.Tensor, + blocks: torch.Tensor, + block_size: int, + out: torch.Tensor, +) -> None: + lens = data.compress_lens[tile] + starts = data.request_starts[tile] + for rows, lc in _requests_in_tile(data, tile): + # Columns past a row's length hold garbage; topk_among_blocks drops them. + # The slice is block-aligned (the tile is), so the op gathers in place. + positions = topk_among_blocks( + logits[rows, : -(-lc // block_size) * block_size], + lens[rows], + blocks[rows], + out.shape[1], + block_size=block_size, + ) + out[rows] = torch.where(positions >= 0, positions + starts[rows, None], -1).to( + torch.int32 + ) + + +def _requests_in_tile(data: DeepGEMMPrefillData, tile: slice): + start = 0 + for num_rows, lc in zip(data.rows_per_request, data.lens_per_request): + end = start + num_rows + begin, stop = max(start, tile.start), min(end, tile.stop) + if begin < stop and lc > 0: + yield slice(begin - tile.start, stop - tile.start), lc + start = end diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py new file mode 100644 index 000000000000..db9dadfcd3e8 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py @@ -0,0 +1,184 @@ +"""The full top-k of the ratio-1/2 index layers outside the candidate scheme: +DeepGEMM on SM100, torch elsewhere.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.kernels.ops.attention.dsv4.index_logits import ( + deep_gemm_fp4_paged_mqa_logits, +) +from sglang.kernels.ops.attention.dsv4.topk import topk_transform_paged_torch +from sglang.srt.layers.attention.dsv4.indexer import topk_transform_paged_from_metadata + +from .scoring import ( + decode_scores, + dense_prefill_topk, + get_deep_gemm_decode_data, + get_deep_gemm_prefill_data, + get_index_k_cache, + prefill_requests, + quantize_index_q, + write_decode, + write_prefill, +) +from .types import ( + CapturedPrefillInputs, + DecodeInputs, + PrefillInputs, + Selection, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + + +class FullTopKIndexer: + def __init__( + self, + *, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, + use_deep_gemm_prefill: bool, + use_deep_gemm_decode: bool, + ): + self.token_to_kv_pool = token_to_kv_pool + self.req_to_token = req_to_token + self.use_deep_gemm_prefill = use_deep_gemm_prefill + self.use_deep_gemm_decode = use_deep_gemm_decode + + def topk_prefill(self, inputs: PrefillInputs, out: Selection) -> None: + if self.use_deep_gemm_prefill: + self._deep_gemm_prefill(inputs, out) + else: + self._torch_prefill(inputs, out) + + def topk_prefill_captured(self, inputs: CapturedPrefillInputs, out: Selection): + self._deep_gemm_prefill_captured(inputs, out) + + def topk_decode(self, inputs: DecodeInputs, out: Selection) -> None: + if self.use_deep_gemm_decode: + self._deep_gemm_decode(inputs, out) + else: + self._torch_decode(inputs, out) + + def _deep_gemm_prefill(self, inputs: PrefillInputs, out: Selection) -> None: + out.reset() + data = get_deep_gemm_prefill_data(inputs, self.req_to_token) + if data is None: + return + kv = self.token_to_kv_pool.get_low_ratio_index_k_fp4( + inputs.layer_id, data.k_slots + ) + selected = dense_prefill_topk(data, kv, topk=inputs.indexer.index_topk) + data.write_selection(selected=selected, out=out) + + def _torch_prefill(self, inputs: PrefillInputs, out: Selection) -> None: + requests = prefill_requests( + inputs=inputs, + out=out, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + ) + for request, chunks in requests: + for chunk in chunks: + idx = chunk.scores.topk(request.k, dim=-1, sorted=False).indices + write_prefill(out, request, chunk, idx) + + def _deep_gemm_decode(self, inputs: DecodeInputs, out: Selection) -> None: + data = get_deep_gemm_decode_data(inputs, self.token_to_kv_pool) + metadata = inputs.paged_metadata + if isinstance(metadata.deep_gemm_metadata, list): + topk_plans = metadata.topk_metadata_chunks + assert not metadata.use_topk_v2 or topk_plans is not None + for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()): + logits = deep_gemm_fp4_paged_mqa_logits( + (data.q_fp4[rows], data.q_sf[rows]), + data.k_cache, + data.weights[rows], + metadata.compressed_seq_lens[rows], + metadata.page_table[rows], + plan, + metadata.max_compressed_seq_len, + ) + # TODO(dark): add bf16 topk + topk_transform_paged_from_metadata( + logits, + metadata, + out.page_indices, + out.raw_indices, + rows=rows, + topk_metadata=( + topk_plans[chunk_idx] if topk_plans is not None else None + ), + ) + return + logits = deep_gemm_fp4_paged_mqa_logits( + (data.q_fp4, data.q_sf), + data.k_cache, + data.weights, + metadata.compressed_seq_lens, + metadata.page_table, + metadata.deep_gemm_metadata, + metadata.max_compressed_seq_len, + ) + # TODO(dark): add bf16 topk + topk_transform_paged_from_metadata( + logits, metadata, out.page_indices, out.raw_indices + ) + + def _torch_decode(self, inputs: DecodeInputs, out: Selection) -> None: + d = decode_scores( + inputs=inputs, + out=out, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + ) + if d is None: + return + k = min(inputs.indexer.index_topk, d.lmax) + idx = d.scores.topk(k, dim=-1, sorted=False).indices + write_decode(out, d, idx) + + def _deep_gemm_prefill_captured( + self, inputs: CapturedPrefillInputs, out: Selection + ) -> None: + indexer = inputs.indexer + metadata = inputs.paged_metadata + assert indexer.n_local_heads == indexer.n_heads + q_fp4, q_sf = quantize_index_q(inputs.q) + num_tokens, num_heads = q_fp4.shape[0], q_fp4.shape[1] + q_fp4 = q_fp4.view(num_tokens, 1, num_heads, 64) + q_sf = q_sf.view(num_tokens, 1, num_heads) + weights = inputs.weights.float() + page_size = metadata.compressed_page_size + k_cache = get_index_k_cache( + token_to_kv_pool=self.token_to_kv_pool, + layer_id=inputs.layer_id, + page_size=page_size, + ) + + width = metadata.max_compressed_seq_len + lens = metadata.compressed_seq_lens + page_table = metadata.page_table + topk = min(indexer.index_topk, width) + for rows, plan in metadata.row_chunks(): + logits = deep_gemm_fp4_paged_mqa_logits( + (q_fp4[rows], q_sf[rows]), + k_cache, + weights[rows], + lens[rows], + page_table[rows], + plan, + width, + ) + topk_transform_paged_torch( + logits, + lens[rows], + page_table[rows], + out.page_indices[rows, :topk], + page_size, + out.raw_indices[rows, :topk] if out.raw_indices is not None else None, + ) diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py new file mode 100644 index 000000000000..1cd46108fa1d --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py @@ -0,0 +1,382 @@ +"""Index scores shared by the full top-k and both candidate schemes: the DeepGEMM +operands and flattened-K scores, the torch per-request scores, and the writes of +a selection back into pool slots.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Generator, Iterator, List, Optional, Tuple + +import msgspec +import torch + +from sglang.kernels.ops.attention.dsv4.fp4_indexer import fp4_index_logits_decode +from sglang.kernels.ops.attention.dsv4.index_logits import flat_index_logits_tiles + +from .types import ( + DecodeInputs, + PrefillInputs, + Selection, +) + +if TYPE_CHECKING: + from sglang.srt.layers.attention.dsv4.dsv41_sparse import DeepseekV41Indexer + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + +# TODO: use a per-forward mqa_logits_budget_bytes() budget that also +# leaves room for candidate block ids and block-selection scratch. +_DEEP_GEMM_SCORE_BUDGET_BYTES = 2 << 30 + + +class DeepGEMMDecodeData(msgspec.Struct, frozen=True): + q_fp4: torch.Tensor # [rows, 1, heads, 64] int8, packed fp4 + q_sf: torch.Tensor # [rows, 1, heads] int32, packed ue8m0 + weights: torch.Tensor # [rows, heads] bf16/fp32 head weights + k_cache: torch.Tensor # [pages, page_size, 1, 68] uint8, the layer's index-K pool + + +class DeepGEMMPrefillData(msgspec.Struct, frozen=True): + """The rows of one request are consecutive; a selection is a flattened-K + column, request start plus compressed position.""" + + k_slots: torch.Tensor # [columns] index-K pool slot of each flattened-K column + request_starts: torch.Tensor # [rows] int32, first column of the row's request + lens_per_request: List[int] # compressed positions at each request's newest token + rows_per_request: List[int] # query rows of each request + compress_lens: torch.Tensor # [rows] int32, compressed positions the row sees + q_fp4: torch.Tensor # [rows, heads, 64] int8, packed fp4 + q_sf: torch.Tensor # [rows, heads] int32, packed ue8m0 + weights: torch.Tensor # [rows, heads] fp32 head weights + + @property + def num_rows(self) -> int: + return self.q_fp4.shape[0] + + def empty_selection(self, topk: int) -> torch.Tensor: + return self.request_starts.new_full( + (self.num_rows, topk), -1, dtype=torch.int32 + ) + + def write_selection(self, selected: torch.Tensor, out: Selection) -> None: + num_tokens, topk = selected.shape + unselected = torch.iinfo(torch.int32).max + selected = selected.masked_fill(selected < 0, unselected).sort(dim=-1).values + chosen = selected != unselected + out.page_indices[:num_tokens, :topk] = torch.where( + chosen, self.k_slots[selected.clamp_max(self.k_slots.shape[0] - 1)], -1 + ).to(torch.int32) + if out.raw_indices is not None: + out.raw_indices[:num_tokens, :topk] = torch.where( + chosen, selected - self.request_starts[:, None], -1 + ) + + +def quantize_index_q(q: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( + quantize_fp4_indexer_tensor, + ) + + rows, heads = q.shape[0], q.shape[1] + q_fp4, q_sf = quantize_fp4_indexer_tensor(q.flatten(0, 1), rne=True) + return q_fp4.view(rows, heads, 64), q_sf.view(rows, heads) + + +def get_index_k_cache( + *, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + layer_id: int, + page_size: int, +) -> torch.Tensor: + """The layer's index-K pool as ``[pages, page_size, 1, 68]`` uint8.""" + k_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id) + assert k_cache.dim() == 2 + # fp4: 64 payload + 4 scale per slot + return k_cache.view(k_cache.shape[0], page_size, 1, 68) + + +def get_deep_gemm_decode_data( + inputs: DecodeInputs, token_to_kv_pool: DeepSeekV4TokenToKVPool +) -> DeepGEMMDecodeData: + indexer = inputs.indexer + x, q_lora, pos = inputs.x, inputs.q_lora, inputs.positions + # The kernel sums head scores locally, so the indexer heads must be replicated. + assert indexer.n_local_heads == indexer.n_heads + if ( + x.is_cuda + and torch.version.cuda is not None + and x.dtype == torch.bfloat16 + and indexer.index_head_dim == 128 + ): + from sglang.kernels.ops.attention.dsv4.fp4_indexer_rope import ( + index_q_rope_pack_weights, + ) + + q, _ = indexer.wq_b(q_lora) + q = q.view(q.shape[0], indexer.n_local_heads, indexer.index_head_dim) + # The fused pack also computes head_weights(x).float(), same rounding. + q_fp4, q_sf, weights = index_q_rope_pack_weights( + q, + torch.view_as_real(inputs.freqs_cis).flatten(-2), + pos, + indexer.head_weights_raw(x), # [bs, n_local] bf16, n32k5120 + indexer.head_weight_scale, + ) + else: + q = indexer.queries(q_lora, inputs.freqs_cis[pos]) + q_fp4, q_sf = quantize_index_q(q) + weights = indexer.head_weights(x).float() # [bs, n_local] + bs = q.shape[0] + q_fp4 = q_fp4.view(bs, 1, indexer.n_local_heads, 64) + q_sf = q_sf.view(bs, 1, indexer.n_local_heads) + + # Index pool page (64 slots); metadata.page_table is expanded to match. + k_cache = get_index_k_cache( + token_to_kv_pool=token_to_kv_pool, + layer_id=inputs.layer_id, + page_size=inputs.paged_metadata.compressed_page_size, + ) + return DeepGEMMDecodeData(q_fp4, q_sf, weights, k_cache) + + +def get_deep_gemm_prefill_data( + inputs: PrefillInputs, req_to_token: torch.Tensor +) -> Optional[DeepGEMMPrefillData]: + """None when the chunk has no row or no visible compressed position.""" + ratio = inputs.compress_ratio + indexer = inputs.indexer + pos = inputs.positions + seq_lens_cpu = inputs.seq_lens_cpu + assert ( + seq_lens_cpu is not None + and inputs.rows_per_request is not None + and inputs.rows_per_request_device is not None + ), "the DeepGEMM prefill indexer needs the CPU length vectors" + device = pos.device + # Visible compressed positions per request; (pos + 1) // ratio bounds each row. + lc_per_req = [s // ratio for s in seq_lens_cpu] + req_pool_indices = inputs.req_pool_indices.to(torch.int64) + slot_chunks, starts, start = [], [], 0 + for r, lc in enumerate(lc_per_req): + starts.append(start) + if lc == 0: + continue + j = torch.arange(lc, device=device) + slot_chunks.append( + req_to_token[req_pool_indices[r], j * ratio].to(torch.int64) // ratio + ) + start += lc + num_tokens = pos.shape[0] + if not slot_chunks or num_tokens == 0: + return None + k_slots = torch.cat(slot_chunks) + q = indexer.queries(inputs.q_lora, inputs.freqs_cis[pos]) + q_fp4, q_sf = quantize_index_q(q) + weights = indexer.head_weights(inputs.x).float() + compress_lens = ((pos + 1) // ratio).to(torch.int32) + request_starts = torch.repeat_interleave( + torch.tensor(starts, dtype=torch.int32, device=device), + inputs.rows_per_request_device.to(torch.int64), + output_size=num_tokens, + ) + return DeepGEMMPrefillData( + k_slots=k_slots, + request_starts=request_starts, + lens_per_request=lc_per_req, + rows_per_request=inputs.rows_per_request, + compress_lens=compress_lens, + q_fp4=q_fp4, + q_sf=q_sf, + weights=weights, + ) + + +def score_tiles( + data: DeepGEMMPrefillData, + kv: Tuple[torch.Tensor, torch.Tensor], + *, + width_align: int, +) -> Generator[Tuple[slice, torch.Tensor], None, None]: + yield from flat_index_logits_tiles( + q=(data.q_fp4, data.q_sf), + kv=kv, + weights=data.weights, + starts=data.request_starts, + lengths=data.compress_lens, + context_lengths=data.lens_per_request, + budget_bytes=_DEEP_GEMM_SCORE_BUDGET_BYTES, + width_align=width_align, + ) + + +def dense_prefill_topk( + data: DeepGEMMPrefillData, + kv: Tuple[torch.Tensor, torch.Tensor], + *, + topk: int, +) -> torch.Tensor: + """Flattened-K columns, ``-1`` padded, unordered.""" + from sglang.kernels.ops.attention.dsv4 import topk_transform_ragged_v2 + + selected = data.empty_selection(topk) + for tile, logits in score_tiles(data, kv, width_align=4): + topk_transform_ragged_v2( + logits, + data.compress_lens[tile], + out_offsets=data.request_starts[tile], + out_indices=selected[tile], + ) + # Free this tile's logits before the generator scores the next one. + del logits + return selected + + +# ---------- torch ---------- + + +# Arbitrary cap on one bf16 [rows, heads, lc] score chunk; transients run ~3x this. +_TORCH_SCORE_BUDGET_BYTES = 1 << 30 + + +class RequestScores(msgspec.Struct, frozen=True): + lc: int # compressed positions visible at its newest row + k: int + columns: torch.Tensor # arange(lc) + tok: torch.Tensor # [rows_b] its query rows in the chunk + lens: torch.Tensor # [rows_b] compressed positions each row sees + slots: torch.Tensor # [lc] index-K pool slot of each position + + +class ChunkScores(msgspec.Struct, frozen=True): + rows: slice # into the request's rows + tok: torch.Tensor + lens: torch.Tensor + scores: torch.Tensor # [rows, lc] bf16, -inf past each row's length + + +class DecodeScores(msgspec.Struct, frozen=True): + bs: int + lmax: int + lens: torch.Tensor # [bs] + slots: torch.Tensor # [bs, lmax] + scores: torch.Tensor # [bs, lmax] + + +def prefill_requests( + *, + inputs: PrefillInputs, + out: Selection, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, +) -> Iterator[tuple[RequestScores, Iterator[ChunkScores]]]: + """Requests with nothing visible yet are skipped.""" + pool = token_to_kv_pool + ratio = inputs.compress_ratio + indexer = inputs.indexer + req, pos = inputs.req_rows, inputs.positions + assert req is not None, "the torch indexer needs the batch's rows" + out.reset() + q = indexer.queries(inputs.q_lora, inputs.freqs_cis[pos]) + weights = indexer.head_weights(inputs.x) + # A compressed position is visible once the query has passed its last token. + compress_lens = (pos + 1) // ratio + topk = indexer.index_topk + for r in torch.unique_consecutive(req).tolist(): + tok = (req == r).nonzero().squeeze(1) + lens = compress_lens[tok] + lc = int(lens.max().item()) + if lc == 0: + continue + j = torch.arange(lc, device=pos.device) + slots = req_to_token[r, j * ratio].to(torch.int64) // ratio + # Dequantize only this request's visible K rows; the table is pool-sized. + index_k = pool.get_low_ratio_index_k_dequant(inputs.layer_id, slots) + request = RequestScores( + lc=lc, k=min(topk, lc), columns=j, tok=tok, lens=lens, slots=slots + ) + yield ( + request, + _score_chunks( + indexer=indexer, q=q, weights=weights, index_k=index_k, request=request + ), + ) + + +def _score_chunks( + *, + indexer: DeepseekV41Indexer, + q: torch.Tensor, + weights: torch.Tensor, + index_k: torch.Tensor, + request: RequestScores, +) -> Iterator[ChunkScores]: + lc, j = request.lc, request.columns + # Chunk rows so the [rows, heads, lc] bf16 scores stay under the budget. + rows_per_chunk = max(1, _TORCH_SCORE_BUDGET_BYTES // (q.shape[1] * lc * 2)) + for start in range(0, request.tok.numel(), rows_per_chunk): + rows = slice(start, start + rows_per_chunk) + tok_c, lens_c = request.tok[rows], request.lens[rows] + s = indexer.scores(q[tok_c], index_k, weights[tok_c]) + s = s.masked_fill(j[None, :] >= lens_c[:, None], -torch.inf) + yield ChunkScores(rows=rows, tok=tok_c, lens=lens_c, scores=s) + + +def write_prefill( + out: Selection, request: RequestScores, chunk: ChunkScores, idx: torch.Tensor +) -> None: + k, lc = request.k, request.lc + idx = idx.sort(dim=-1).values + reach = idx < chunk.lens[:, None] + out.page_indices[chunk.tok, :k] = torch.where( + reach, request.slots[idx.clamp_max(lc - 1)], -1 + ).to(torch.int32) + if out.raw_indices is not None: + out.raw_indices[chunk.tok, :k] = torch.where(reach, idx, -1).to(torch.int32) + + +def decode_scores( + *, + inputs: DecodeInputs, + out: Selection, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, +) -> Optional[DecodeScores]: + """None when there is no row, or nothing visible yet.""" + pool = token_to_kv_pool + ratio = inputs.compress_ratio + indexer = inputs.indexer + req, pos = inputs.req_rows, inputs.positions + out.reset() + bs = req.shape[0] + assert pos.shape[0] == bs, ( + f"decode expects one token per request, {pos.shape=} {bs=}" + ) + if bs == 0: + return None + lens = (pos + 1) // ratio + metadata = inputs.paged_metadata + # V4 reserves the replay bound in metadata; visibility stays on device. + # A capture-time length read would both synchronize and truncate replay. + lmax = min(metadata.max_compressed_seq_len, req_to_token.shape[1] // ratio) + if lmax == 0: + return None + q = indexer.queries(inputs.q_lora, inputs.freqs_cis[pos]) + weights = indexer.head_weights(inputs.x) + j = torch.arange(lmax, device=pos.device) + valid = j[None, :] < lens[:, None] + slots = req_to_token[req[:, None], (j * ratio)[None, :]].to(torch.int64) // ratio + slots = slots.masked_fill(~valid, 0) + table = pool.get_index_k_with_scale_buffer(inputs.layer_id) + scores = fp4_index_logits_decode( + q, weights, slots, lens, table, table.shape[1] // 68 + ) + return DecodeScores(bs=bs, lmax=lmax, lens=lens, slots=slots, scores=scores) + + +def write_decode(out: Selection, d: DecodeScores, idx: torch.Tensor) -> None: + k = idx.shape[1] + idx = idx.sort(dim=-1).values + reach = idx < d.lens[:, None] + out.page_indices[: d.bs, :k] = torch.where( + reach, d.slots.gather(1, idx.clamp_max(d.lmax - 1)), -1 + ).to(torch.int32) + if out.raw_indices is not None: + out.raw_indices[: d.bs, :k] = torch.where(reach, idx, -1).to(torch.int32) diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py new file mode 100644 index 000000000000..c01d76f7be94 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py @@ -0,0 +1,449 @@ +"""The candidate scheme carried by a DeepGEMM sparse table (SM100): a source +publishes its blocks with the schedule ``fp8_fp4_paged_sparse_mqa_logits`` needs, +a consumer scores only those blocks out of the index-K pool.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Optional, Tuple + +import msgspec +import torch + +from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( + amax8_varlen, + amax_topk_blocks, + candidate_row_lens, +) +from sglang.kernels.ops.attention.dsv4.candidate_table import ( + CANDIDATE_BLOCK_SIZE, + build_sparse_indexer_schedule, + sort_candidate_blocks, +) +from sglang.kernels.ops.attention.dsv4.index_logits import ( + deep_gemm_fp4_paged_mqa_logits, + sparse_logits, +) +from sglang.kernels.ops.attention.dsv4.topk import ( + topk_transform_paged_v2, + topk_transform_ragged_v2, + topk_transform_sparse, +) +from sglang.srt.layers.attention.dsv4.indexer import topk_transform_paged_from_metadata +from sglang.srt.layers.attention.dsv4.metadata import expand_index_page_table + +from .scoring import ( + DeepGEMMDecodeData, + DeepGEMMPrefillData, + get_deep_gemm_decode_data, + get_deep_gemm_prefill_data, + get_index_k_cache, + score_tiles, +) +from .types import ( + CandidateMetadata, + DecodeInputs, + PrefillInputs, + Selection, + get_tail_row_indices, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + + +class _SparseTable(CandidateMetadata, msgspec.Struct): + # [rows, topk_blocks] int32: ascending logical block ids, valid for the first + # min(topk_blocks, ceil(seq_len / 8)) entries of a row; DeepGEMM reads only those + blocks: torch.Tensor + # DeepGEMM's schedule metadata (uint8) for them + schedule: torch.Tensor + # [rows, topk_blocks] int32: the same blocks as pool slots / 8; column j of the + # sparse row maps to slot phys_blocks[b, j // 8] * 8 + j % 8 + phys_blocks: torch.Tensor + # [rows] int32: length of each row of the sparse logits: the published blocks + # laid out block by block, the newest possibly partial (`candidate_row_lens`) + valid_lens: torch.Tensor + # recorded on the side stream once the fields above are complete + ready: torch.cuda.Event + + +class _SparsePrefillTable(_SparseTable): + """The inputs stay attached so the late-layer tail can rebuild the table for + its rows: the DeepGEMM schedule is per row set, it cannot be sliced.""" + + compress_lens: torch.Tensor # [rows] int32 + page_table: torch.Tensor # [rows, index pages] int32 + request_ids: torch.Tensor # [rows] int32 + q_dtype: torch.dtype + page_size: int # index-K pool page size (slots) + rows_per_request: List[int] # rows of each request, in row order + + def tail(self, rows_per_request: List[int]) -> _SparsePrefillTable: + rows = get_tail_row_indices( + full_rows_per_request=self.rows_per_request, + tail_rows_per_request=rows_per_request, + device=self.blocks.device, + ) + compress_lens = self.compress_lens[rows].contiguous() + _, valid_lens = candidate_row_lens(compress_lens, self.blocks.shape[1]) + return _build_prefill_table( + blocks=self.blocks[rows].contiguous(), + compress_lens=compress_lens, + page_table=self.page_table[rows].contiguous(), + page_size=self.page_size, + request_ids=self.request_ids[rows].contiguous(), + rows_per_request=rows_per_request, + q_dtype=self.q_dtype, + valid_lens=valid_lens, + ) + + +class SparseTableBackend: + def __init__( + self, + *, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, + page_size: int, + candidate_topk_blocks: int, + candidate_block_size: int, + ) -> None: + assert candidate_block_size == CANDIDATE_BLOCK_SIZE + self.token_to_kv_pool = token_to_kv_pool + self.req_to_token = req_to_token + self.page_size = page_size + self.topk_blocks = candidate_topk_blocks + self.alt_stream = torch.cuda.Stream() + self._cached_row_ids: Optional[torch.Tensor] = None + # captured graphs keep reading the buffers they saw + self._retired_cached_row_ids: list = [] + + def publish_prefill(self, inputs: PrefillInputs, out: Selection): + out.reset() + data = get_deep_gemm_prefill_data(inputs, self.req_to_token) + if data is None: + return None + selected = data.empty_selection(inputs.indexer.index_topk) + index_page_size = self.token_to_kv_pool.get_index_k_page_size( + inputs.compress_ratio + ) + table = publish_prefill_table( + data=data, + kv=self.token_to_kv_pool.get_low_ratio_index_k_fp4( + inputs.layer_id, data.k_slots + ), + index_page_table=expand_index_page_table( + inputs.kv_page_table[: data.num_rows], + full_page_size=self.page_size, + compress_ratio=inputs.compress_ratio, + index_page_size=index_page_size, + ).contiguous(), + index_page_size=index_page_size, + topk_blocks=self.topk_blocks, + out_positions=selected, + ) + data.write_selection(selected=selected, out=out) + return table + + def consume_prefill( + self, + inputs: PrefillInputs, + published: Optional[_SparsePrefillTable], + out: Selection, + ) -> None: + out.reset() + data = get_deep_gemm_prefill_data(inputs, self.req_to_token) + if data is None: + return + assert published is not None + selected = data.empty_selection(inputs.indexer.index_topk) + select_prefill_table( + table=published, + data=data, + k_cache=get_index_k_cache( + token_to_kv_pool=self.token_to_kv_pool, + layer_id=inputs.layer_id, + page_size=self.token_to_kv_pool.get_index_k_page_size( + inputs.compress_ratio + ), + ), + out_positions=selected, + ) + data.write_selection(selected=selected, out=out) + + def publish_decode(self, inputs: DecodeInputs, out: Selection): + data = get_deep_gemm_decode_data(inputs, self.token_to_kv_pool) + metadata = inputs.paged_metadata + seq_lens = metadata.compressed_seq_lens.reshape(-1) + if isinstance(metadata.deep_gemm_metadata, list): + return self._publish_decode_chunked(inputs, data, seq_lens, out) + logits = deep_gemm_fp4_paged_mqa_logits( + (data.q_fp4, data.q_sf), + data.k_cache, + data.weights, + metadata.compressed_seq_lens, + metadata.page_table, + metadata.deep_gemm_metadata, + metadata.max_compressed_seq_len, + ) + main_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(main_stream) + # The block-selection chain reads logits after the main stream moves on. + logits.record_stream(self.alt_stream) + # TODO(candidate): one kernel for both selections below (dense logits read once) + topk_transform_paged_v2( + logits, + seq_lens, + metadata.page_table, + out.page_indices, + metadata.compressed_page_size, + metadata.topk_metadata, + out_raw_indices=out.raw_indices, + ) + # per row: block count for the block top-k, sparse-row length for the consumers + with torch.cuda.stream(self.alt_stream): + nblocks, row_valid_lens = candidate_row_lens(seq_lens, self.topk_blocks) + blocks = amax_topk_blocks(logits, seq_lens, nblocks, self.topk_blocks) + # in place: ascending, INT32_MAX padded, plus the blocks as pool slots / 8 + phys_blocks = sort_candidate_blocks( + blocks, + seq_lens, + metadata.page_table, + metadata.compressed_page_size, + ) + schedule = build_sparse_indexer_schedule( + blocks, + seq_lens, + metadata.page_table, + metadata.compressed_page_size, + data.q_fp4.dtype, + self._get_request_ids(inputs, data.q_fp4.shape[0], blocks.device), + ) + # consume_decode reads these on the main stream + for t in (blocks, schedule, phys_blocks, row_valid_lens): + t.record_stream(main_stream) + ready = torch.cuda.Event() + ready.record(self.alt_stream) + return _SparseTable( + blocks=blocks, + schedule=schedule, + phys_blocks=phys_blocks, + valid_lens=row_valid_lens, + ready=ready, + ) + + def _publish_decode_chunked( + self, + inputs: DecodeInputs, + data: DeepGEMMDecodeData, + seq_lens: torch.Tensor, + out: Selection, + ) -> _SparseTable: + """On the current stream, so each chunk's logits are released before the next.""" + metadata = inputs.paged_metadata + block_chunks = [] + phys_block_chunks = [] + valid_len_chunks = [] + topk_plans = metadata.topk_metadata_chunks + assert not metadata.use_topk_v2 or topk_plans is not None + for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()): + logits = deep_gemm_fp4_paged_mqa_logits( + (data.q_fp4[rows], data.q_sf[rows]), + data.k_cache, + data.weights[rows], + metadata.compressed_seq_lens[rows], + metadata.page_table[rows], + plan, + metadata.max_compressed_seq_len, + ) + topk_transform_paged_from_metadata( + logits, + metadata, + out.page_indices, + out.raw_indices, + rows=rows, + topk_metadata=( + topk_plans[chunk_idx] if topk_plans is not None else None + ), + ) + chunk_seq_lens = seq_lens[rows] + nblocks, row_valid_lens = candidate_row_lens( + chunk_seq_lens, self.topk_blocks + ) + blocks = amax_topk_blocks(logits, chunk_seq_lens, nblocks, self.topk_blocks) + phys_blocks = sort_candidate_blocks( + blocks, + chunk_seq_lens, + metadata.page_table[rows], + metadata.compressed_page_size, + ) + block_chunks.append(blocks) + phys_block_chunks.append(phys_blocks) + valid_len_chunks.append(row_valid_lens) + + blocks = torch.cat(block_chunks) + schedule = build_sparse_indexer_schedule( + blocks, + seq_lens, + metadata.page_table, + metadata.compressed_page_size, + data.q_fp4.dtype, + self._get_request_ids(inputs, data.q_fp4.shape[0], blocks.device), + ) + ready = torch.cuda.Event() + ready.record(torch.cuda.current_stream()) + return _SparseTable( + blocks=blocks, + schedule=schedule, + phys_blocks=torch.cat(phys_block_chunks), + valid_lens=torch.cat(valid_len_chunks), + ready=ready, + ) + + def consume_decode( + self, + inputs: DecodeInputs, + published: Optional[_SparseTable], + out: Selection, + ) -> None: + assert published is not None + data = get_deep_gemm_decode_data(inputs, self.token_to_kv_pool) + torch.cuda.current_stream().wait_event(published.ready) + logits = sparse_logits( + data.q_fp4, + data.q_sf, + data.k_cache, + data.weights.to(torch.bfloat16), + published.schedule, + published.blocks.shape[1], + ) + assert out.raw_indices is None + topk_transform_sparse( + logits, published.valid_lens, published.phys_blocks, out.page_indices + ) + + def _get_request_ids(self, inputs: DecodeInputs, rows: int, device: torch.device): + if inputs.is_verify: + return inputs.req_rows[:rows].to(torch.int32).contiguous() + buf = self._cached_row_ids + if buf is None or buf.numel() < rows: + assert not torch.cuda.is_current_stream_capturing(), ( + f"row-id buffer grows to {rows} rows inside a CUDA graph capture; " + "warm up with the largest row count first" + ) + if buf is not None: + self._retired_cached_row_ids.append(buf) + size = max(8192, -(-rows // 8192) * 8192) + buf = torch.arange(size, dtype=torch.int32, device=device) + self._cached_row_ids = buf + return buf[:rows] + + +def _build_prefill_table( + *, + blocks: torch.Tensor, + compress_lens: torch.Tensor, + page_table: torch.Tensor, + page_size: int, + request_ids: torch.Tensor, + rows_per_request: List[int], + q_dtype: torch.dtype, + valid_lens: torch.Tensor, +) -> _SparsePrefillTable: + # in place: ascending, INT32_MAX padded, plus the blocks as pool slots / 8 + phys_blocks = sort_candidate_blocks(blocks, compress_lens, page_table, page_size) + schedule = build_sparse_indexer_schedule( + blocks, compress_lens, page_table, page_size, q_dtype, request_ids + ) + ready = torch.cuda.Event() + ready.record() + return _SparsePrefillTable( + blocks=blocks, + schedule=schedule, + phys_blocks=phys_blocks, + valid_lens=valid_lens, + ready=ready, + compress_lens=compress_lens, + page_table=page_table, + request_ids=request_ids, + q_dtype=q_dtype, + page_size=page_size, + rows_per_request=rows_per_request, + ) + + +# TODO(dark): support fusion of publish + topk of publish layer +def publish_prefill_table( + *, + data: DeepGEMMPrefillData, + kv: Tuple[torch.Tensor, torch.Tensor], + index_page_table: torch.Tensor, + index_page_size: int, + topk_blocks: int, + out_positions: torch.Tensor, +) -> _SparsePrefillTable: + """The source's own top-k into ``out_positions`` and the chunk's table, from one + pass over its dense scores; ``index_page_table`` is at ``index_page_size`` slots.""" + rows = data.num_rows + device = data.q_fp4.device + nblocks, valid_lens = candidate_row_lens(data.compress_lens, topk_blocks) + blocks = torch.empty(rows, topk_blocks, dtype=torch.int32, device=device) + zero_offsets = torch.zeros(rows, dtype=torch.int32, device=device) + # the block keys read the score rows through 32-byte vectors + for tile, logits in score_tiles(data, kv, width_align=8): + lens = data.compress_lens[tile] + topk_transform_ragged_v2( + logits, + lens, + out_offsets=data.request_starts[tile], + out_indices=out_positions[tile], + ) + # keys past a row's block count stay unset; the top-k reads a row + # up to its block count only (v2 wants the stride a multiple of 4) + width_blocks = ( + logits.shape[1] + CANDIDATE_BLOCK_SIZE - 1 + ) // CANDIDATE_BLOCK_SIZE + keys = logits.new_empty(logits.shape[0], (width_blocks + 3) // 4 * 4) + amax8_varlen(logits, lens, out=keys) + topk_transform_ragged_v2( + keys, + nblocks[tile], + out_offsets=zero_offsets[: logits.shape[0]], + out_indices=blocks[tile], + ) + request_ids = torch.repeat_interleave( + torch.arange(len(data.rows_per_request), dtype=torch.int32, device=device), + torch.tensor(data.rows_per_request, dtype=torch.int64, device=device), + output_size=rows, + ) + return _build_prefill_table( + blocks=blocks, + compress_lens=data.compress_lens, + page_table=index_page_table, + page_size=index_page_size, + request_ids=request_ids, + rows_per_request=data.rows_per_request, + q_dtype=data.q_fp4.dtype, + valid_lens=valid_lens, + ) + + +def select_prefill_table( + *, + table: _SparsePrefillTable, + data: DeepGEMMPrefillData, + k_cache: torch.Tensor, + out_positions: torch.Tensor, +) -> None: + rows, heads = data.q_sf.shape + logits = sparse_logits( + data.q_fp4.view(rows, 1, heads, 64), + data.q_sf.view(rows, 1, heads), + k_cache, + data.weights.to(torch.bfloat16), + table.schedule, + table.blocks.shape[1], + ) + # request-relative compressed positions, -1 padded + topk_transform_sparse(logits, table.valid_lens, table.blocks, out_positions) + out_positions.add_(torch.where(out_positions >= 0, data.request_starts[:, None], 0)) diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py new file mode 100644 index 000000000000..8dee611dc624 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py @@ -0,0 +1,129 @@ +"""The inputs and outputs of the V4.1 indexer backends.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Optional, Protocol, TypeVar + +import msgspec +import torch + +if TYPE_CHECKING: + from sglang.srt.layers.attention.dsv4.dsv41_sparse import DeepseekV41Indexer + from sglang.srt.layers.attention.dsv4.metadata import PagedIndexerMetadata + + +def get_tail_row_indices( + full_rows_per_request: List[int], + tail_rows_per_request: List[int], + device: torch.device, +) -> torch.Tensor: + """Row indices of each request's tail, copied without a host sync.""" + rows, start = [], 0 + for n, t in zip(full_rows_per_request, tail_rows_per_request): + rows.extend(range(start + n - t, start + n)) + start += n + device = torch.device(device) + if device.type != "cuda": + return torch.tensor(rows, dtype=torch.int64, device=device) + staged = torch.tensor(rows, dtype=torch.int64, pin_memory=True) + return staged.to(device, non_blocking=True) + + +class CandidateMetadata: + def tail(self, rows_per_request: List[int]) -> CandidateMetadata: + raise NotImplementedError(f"{type(self).__name__} is not a prefill publish") + + +class Selection(msgspec.Struct, frozen=True): + page_indices: torch.Tensor # index-K pool slots, ascending + raw_indices: Optional[torch.Tensor] # the same picks as compressed positions + + def reset(self) -> None: + self.page_indices.fill_(-1) + if self.raw_indices is not None: + self.raw_indices.fill_(-1) + + +class PrefillInputs(msgspec.Struct, frozen=True, kw_only=True): + """An eager extend: the backend projects the queries and gathers the index K.""" + + indexer: DeepseekV41Indexer + layer_id: int + compress_ratio: int + freqs_cis: torch.Tensor + x: torch.Tensor # [rows, hidden] bf16 + q_lora: torch.Tensor # [rows, q_lora_rank] + positions: torch.Tensor # [rows] absolute position of each row + + # [rows] req_to_token row of each query row; None under prefill CP, where a + # rank's rows are an interleaved subset of the batch's. + req_rows: Optional[torch.Tensor] + req_pool_indices: torch.Tensor # [requests] req_to_token row of each request + # [rows, pages] int32 at the KV page size; the index-K pool pages it with + # expand_index_page_table. + kv_page_table: torch.Tensor + # Visible full-length context of each request at its newest row, and the + # query rows of each request in row order. + seq_lens_cpu: Optional[List[int]] + rows_per_request: Optional[List[int]] + rows_per_request_device: Optional[torch.Tensor] + + +class DecodeInputs(msgspec.Struct, frozen=True, kw_only=True): + """Decode (one row per request) or verify (one row per draft token).""" + + indexer: DeepseekV41Indexer + layer_id: int + compress_ratio: int + freqs_cis: torch.Tensor + x: torch.Tensor # [rows, hidden] bf16 + q_lora: torch.Tensor # [rows, q_lora_rank] + positions: torch.Tensor # [rows] absolute position of each row + req_rows: torch.Tensor # [rows] req_to_token row of each query row + paged_metadata: PagedIndexerMetadata + is_verify: bool + + +class CapturedPrefillInputs(msgspec.Struct, frozen=True, kw_only=True): + """An extend under the prefill graph: projected queries, paged index K.""" + + indexer: DeepseekV41Indexer + layer_id: int + q: torch.Tensor # [rows, heads, 128] bf16, roped + weights: torch.Tensor # [rows, heads] bf16 head weights + paged_metadata: PagedIndexerMetadata + + +Metadata = TypeVar("Metadata", bound=CandidateMetadata) + + +class PrefillCandidates(Protocol[Metadata]): + """The attention backend keeps a publish alive and hands it back unread.""" + + def publish_prefill( + self, + inputs: PrefillInputs, + out: Selection, + ) -> Optional[Metadata]: ... + + def consume_prefill( + self, + inputs: PrefillInputs, + published: Optional[Metadata], + out: Selection, + ) -> None: ... + + +class DecodeCandidates(Protocol[Metadata]): + def publish_decode( + self, + inputs: DecodeInputs, + out: Selection, + ) -> Optional[Metadata]: ... + + def consume_decode( + self, + inputs: DecodeInputs, + published: Optional[Metadata], + out: Selection, + ) -> None: ... diff --git a/python/sglang/srt/model_executor/runner_utils/capture_mode.py b/python/sglang/srt/model_executor/runner_utils/capture_mode.py index 38770ea4015e..ce4f08f1fb3b 100644 --- a/python/sglang/srt/model_executor/runner_utils/capture_mode.py +++ b/python/sglang/srt/model_executor/runner_utils/capture_mode.py @@ -79,6 +79,16 @@ def skip_low_ratio_indexer(compress_ratio: int) -> bool: ) +def skip_candidate_filtering() -> bool: + """Whether the captured variant stays inside the candidate window, where every + block is a candidate.""" + return _capture_attention_variant in ( + "candidate_all", + "candidate_c2_all", + "candidate_unfiltered", + ) + + def _set_capture_attention_variant(variant: Optional[str]) -> None: global _capture_attention_variant _capture_attention_variant = variant diff --git a/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py b/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py index 62c1ff6ff06b..34682a239ea9 100644 --- a/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py +++ b/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py @@ -4,7 +4,7 @@ import torch from sglang.kernels.ops.attention.dsv4.fp4_indexer import quantize_fp4_indexer_tensor -from sglang.srt.layers.attention.dsv4 import dense_prefill_indexer +from sglang.srt.layers.attention.dsv4.v41_indexer import dense_blocks, scoring from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -48,6 +48,54 @@ def make_inputs(request_lengths, ratio=1, zero_queries=False, seed=17): ) +def run_dense(inputs, *, publish_candidates, candidates): + """The plain top-k, or the CP source's own top-k plus its candidate blocks + (`publish_candidates`), or a CP consumer's top-k among `candidates`.""" + data = scoring.DeepGEMMPrefillData( + k_slots=torch.arange(inputs["kv"][0].shape[0], device="cuda"), + request_starts=inputs["starts"], + lens_per_request=[n for _, n in inputs["request_lengths"]], + rows_per_request=[q for q, _ in inputs["request_lengths"]], + compress_lens=inputs["lengths"], + q_fp4=inputs["q"][0], + q_sf=inputs["q"][1], + weights=inputs["weights"], + ) + if publish_candidates: + selected, blocks = dense_blocks._publish_prefill_blocks( + data=data, + kv=inputs["kv"], + topk=inputs["topk"], + topk_blocks=inputs["candidate_topk_blocks"], + block_size=inputs["candidate_block_size"], + ) + return selected, dense_blocks.BlockIds( + blocks=blocks, rows_per_request=data.rows_per_request + ) + if candidates is not None: + selected = dense_blocks._consume_prefill_blocks( + data=data, + kv=inputs["kv"], + topk=inputs["topk"], + blocks=candidates.blocks, + block_size=inputs["candidate_block_size"], + ) + return selected, None + return scoring.dense_prefill_topk(data, inputs["kv"], topk=inputs["topk"]), None + + +def request_blocks(inputs, candidates): + """Each request's rows of the published block ids, cut to its kept count + ``min(topk_blocks, ceil(context / block))``.""" + out, row = [], 0 + block = inputs["candidate_block_size"] + for queries, context in inputs["request_lengths"]: + kept = min(inputs["candidate_topk_blocks"], -(-context // block)) + out.append(candidates.blocks[row : row + queries, :kept]) + row += queries + return out + + def dense_scores(inputs): from deep_gemm import fp8_fp4_mqa_logits @@ -121,15 +169,16 @@ def test_ragged_source_consumer_and_replay(self): consumer_inputs["candidate_topk_blocks"] = 128 consumer_scores = dense_scores(consumer_inputs) with patch.object( - dense_prefill_indexer, "_SCORE_BUDGET_BYTES", 128 << 10 + scoring, "_DEEP_GEMM_SCORE_BUDGET_BYTES", 128 << 10 ): - selected, candidates = dense_prefill_indexer.dense_prefill_topk( - **inputs, publish_candidates=True, candidates=None + selected, candidates = run_dense( + inputs, publish_candidates=True, candidates=None ) self.assert_topk(inputs, selected, scores) row = 0 for (queries, context), blocks in zip( - inputs["request_lengths"], candidates.request_blocks + inputs["request_lengths"], + request_blocks(inputs, candidates), ): local = scores[row : row + queries, :context] if queries and context: @@ -173,8 +222,8 @@ def test_ragged_source_consumer_and_replay(self): torch.isfinite(consumer_scores[-1]).sum().item(), inputs["lengths"][-1].item(), ) - selected, published = dense_prefill_indexer.dense_prefill_topk( - **consumer_inputs, + selected, published = run_dense( + consumer_inputs, publish_candidates=False, candidates=candidates, ) @@ -201,8 +250,8 @@ def test_ragged_source_consumer_and_replay(self): ) ), ) - selected, _ = dense_prefill_indexer.dense_prefill_topk( - **tail_inputs, + selected, _ = run_dense( + tail_inputs, publish_candidates=False, candidates=candidates.tail(tail_lengths), ) @@ -213,34 +262,33 @@ def test_unfiltered_and_zero_length_requests(self): for publish in (False, True): with self.subTest(request_lengths=request_lengths, publish=publish): inputs = make_inputs(request_lengths) - selected, candidates = dense_prefill_indexer.dense_prefill_topk( - **inputs, publish_candidates=publish, candidates=None + selected, candidates = run_dense( + inputs, publish_candidates=publish, candidates=None ) self.assert_topk(inputs, selected, dense_scores(inputs)) if publish: self.assertEqual( - [b.shape[0] for b in candidates.request_blocks], + [b.shape[0] for b in request_blocks(inputs, candidates)], [q for q, _ in request_lengths], ) def test_empty_queries_or_context(self): - for request_lengths, shape, block_shapes in ( - ([], (0, 512), []), - ([(0, 0)], (0, 512), [(0, 0)]), - ([(0, 0), (0, 17)], (0, 512), [(0, 0), (0, 2)]), - ([(1, 0)], (1, 512), [(1, 0)]), + # block ids are [rows, min(topk_blocks, ceil(max context / 8))] + for request_lengths, shape, block_shape in ( + ([], (0, 512), (0, 0)), + ([(0, 0)], (0, 512), (0, 0)), + ([(0, 0), (0, 17)], (0, 512), (0, 2)), + ([(1, 0)], (1, 512), (1, 0)), ): with self.subTest(request_lengths=request_lengths): inputs = make_inputs(request_lengths) - selected, candidates = dense_prefill_indexer.dense_prefill_topk( - **inputs, publish_candidates=True, candidates=None + selected, candidates = run_dense( + inputs, publish_candidates=True, candidates=None ) torch.testing.assert_close( selected, torch.full(shape, -1, dtype=torch.int32, device="cuda") ) - self.assertEqual( - [tuple(b.shape) for b in candidates.request_blocks], block_shapes - ) + self.assertEqual(tuple(candidates.blocks.shape), block_shape) def test_score_budget_includes_allocation_padding(self): from deep_gemm import fp8_fp4_mqa_logits @@ -258,11 +306,11 @@ def checked_logits(*args, **kwargs): return logits with ( - patch.object(dense_prefill_indexer, "_SCORE_BUDGET_BYTES", budget), + patch.object(scoring, "_DEEP_GEMM_SCORE_BUDGET_BYTES", budget), patch("deep_gemm.fp8_fp4_mqa_logits", new=checked_logits), ): - selected, _ = dense_prefill_indexer.dense_prefill_topk( - **inputs, publish_candidates=True, candidates=None + selected, _ = run_dense( + inputs, publish_candidates=True, candidates=None ) self.assertTrue(allocations) self.assertLessEqual(max(allocations), budget) @@ -276,20 +324,18 @@ def test_score_memory_is_bounded(self): torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() baseline = torch.cuda.memory_allocated() - selected, candidates = dense_prefill_indexer.dense_prefill_topk( - **inputs, publish_candidates=True, candidates=None + selected, candidates = run_dense( + inputs, publish_candidates=True, candidates=None ) torch.cuda.synchronize() self.assertLess( torch.cuda.max_memory_allocated() - baseline, limit_gib << 30 ) self.assertEqual(tuple(selected.shape), (16384, 512)) - self.assertEqual( - tuple(candidates.request_blocks[0].shape), (16384, 2048) - ) + self.assertEqual(tuple(candidates.blocks.shape), (16384, 2048)) del selected - selected, published = dense_prefill_indexer.dense_prefill_topk( - **inputs, publish_candidates=False, candidates=candidates + selected, published = run_dense( + inputs, publish_candidates=False, candidates=candidates ) self.assertIsNone(published) del selected @@ -297,8 +343,8 @@ def test_score_memory_is_bounded(self): baseline = torch.cuda.memory_allocated() for _ in range(3): torch.cuda.reset_peak_memory_stats() - selected, published = dense_prefill_indexer.dense_prefill_topk( - **inputs, publish_candidates=False, candidates=candidates + selected, published = run_dense( + inputs, publish_candidates=False, candidates=candidates ) torch.cuda.synchronize() self.assertIsNone(published) diff --git a/test/registered/kernels/ops/attention/test_dsv41_prefill_sparse_indexer.py b/test/registered/kernels/ops/attention/test_dsv41_prefill_sparse_indexer.py index e5d62892ecff..c6ce00b12261 100644 --- a/test/registered/kernels/ops/attention/test_dsv41_prefill_sparse_indexer.py +++ b/test/registered/kernels/ops/attention/test_dsv41_prefill_sparse_indexer.py @@ -7,11 +7,12 @@ import msgspec import torch -from sglang.kernels.ops.attention.dsv4.candidate_blocks import select_candidate_blocks +from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( + select_candidate_block_ids, +) from sglang.kernels.ops.attention.dsv4.fp4_indexer import quantize_fp4_indexer_tensor -from sglang.srt.layers.attention.dsv4.candidate_indexer import PrefillIndexerInputs -from sglang.srt.layers.attention.dsv4.dense_prefill_indexer import ( - DenseCandidateIndexer, +from sglang.srt.layers.attention.dsv4.v41_indexer.scoring import ( + DeepGEMMPrefillData, ) from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -35,7 +36,10 @@ class Case(NamedTuple): dense: torch.Tensor # fp32 [rows, width] scores of every row against its context lens: torch.Tensor # [rows] int32 causal lengths - inputs: PrefillIndexerInputs + data: DeepGEMMPrefillData + kv: tuple # the flattened index K (int8 [n, 64], int32 [n]) + k_cache: torch.Tensor # [pages, PAGE, 1, 68] uint8, the index-K pool + page_table: torch.Tensor # [rows, pages] int32 at PAGE slots def make_case(rows, ctx, seed) -> Case: @@ -75,38 +79,35 @@ def make_case(rows, ctx, seed) -> Case: False, width, ) - inputs = PrefillIndexerInputs( - q_fp4=q_fp4, - q_sf=q_sf, - weights=weights, - compress_lens=lens, + data = DeepGEMMPrefillData( + k_slots=torch.arange(width, device=dev), request_starts=starts, lens_per_request=[ctx], rows_per_request=[rows], - kv=(k_fp4[:width], k_sf[:width]), - k_cache=k_cache, - page_size=PAGE, - # KV pages of PAGE tokens at ratio 1: the index page table is the KV one - kv_page_table=torch.arange(n_pages, device=dev, dtype=torch.int32) + compress_lens=lens, + q_fp4=q_fp4, + q_sf=q_sf, + weights=weights, + ) + # KV pages of PAGE tokens at ratio 1: the index page table is the KV one + page_table = ( + torch.arange(n_pages, device=dev, dtype=torch.int32) .expand(rows, -1) - .contiguous(), - kv_page_size=PAGE, - compress_ratio=1, + .contiguous() ) - return Case(dense, lens, inputs) + return Case(dense, lens, data, (k_fp4[:width], k_sf[:width]), k_cache, page_table) -def rows_of(inputs: PrefillIndexerInputs, idx: torch.Tensor) -> PrefillIndexerInputs: - """The inputs of a subset of rows (one request).""" +def rows_of(data: DeepGEMMPrefillData, idx: torch.Tensor) -> DeepGEMMPrefillData: + """The operands of a subset of rows (one request).""" return msgspec.structs.replace( - inputs, - q_fp4=inputs.q_fp4[idx], - q_sf=inputs.q_sf[idx], - weights=inputs.weights[idx], - compress_lens=inputs.compress_lens[idx], - request_starts=inputs.request_starts[idx], + data, + q_fp4=data.q_fp4[idx], + q_sf=data.q_sf[idx], + weights=data.weights[idx], + compress_lens=data.compress_lens[idx], + request_starts=data.request_starts[idx], rows_per_request=[idx.numel()], - kv_page_table=inputs.kv_page_table[idx], ) @@ -114,26 +115,77 @@ def reference_blocks(case: Case) -> torch.Tensor: """[rows, blocks] bool: the torch block selection of the dense scores.""" j = torch.arange(case.dense.shape[1], device=case.dense.device) scores = case.dense.masked_fill(j[None, :] >= case.lens[:, None], -torch.inf) - keep = select_candidate_blocks( + ids = select_candidate_block_ids( scores, case.lens[:, None], topk_blocks=TOPK_BLOCKS, block_size=BLOCK + ).to(torch.int64) + num_blocks = -(-scores.shape[1] // BLOCK) + keep = torch.zeros( + ids.shape[0], num_blocks + 1, dtype=torch.bool, device=ids.device + ) + keep.scatter_(1, ids.masked_fill(ids < 0, num_blocks), True) + return keep[:, :num_blocks] + + +def publish_sparse(case: Case): + """(the DeepGEMM sparse table, the source layer's own top-k).""" + from sglang.srt.layers.attention.dsv4.v41_indexer.sparse_table import ( + publish_prefill_table, ) - return keep.unflatten(1, (-1, BLOCK)).any(-1) + own = case.data.empty_selection(TOPK) + table = publish_prefill_table( + data=case.data, + kv=case.kv, + index_page_table=case.page_table, + index_page_size=PAGE, + topk_blocks=TOPK_BLOCKS, + out_positions=own, + ) + return table, own -def publish(indexer, inputs): - """(published metadata, the source layer's own top-k).""" - own = torch.full((inputs.num_rows, TOPK), -1, dtype=torch.int32, device="cuda") - return indexer.publish_prefill(inputs, own), own +def select_sparse(table, data: DeepGEMMPrefillData, k_cache: torch.Tensor): + from sglang.srt.layers.attention.dsv4.v41_indexer.sparse_table import ( + select_prefill_table, + ) -def select(indexer, published, inputs): - positions = torch.full( - (inputs.num_rows, TOPK), -1, dtype=torch.int32, device="cuda" + positions = data.empty_selection(TOPK) + select_prefill_table( + table=table, data=data, k_cache=k_cache, out_positions=positions ) - indexer.select_prefill(published, inputs, positions) return positions +def publish_dense(case: Case): + """(the CP implementation's block ids, the source layer's own top-k).""" + from sglang.srt.layers.attention.dsv4.v41_indexer.dense_blocks import ( + _publish_prefill_blocks, + ) + + own, blocks = _publish_prefill_blocks( + data=case.data, + kv=case.kv, + topk=TOPK, + topk_blocks=TOPK_BLOCKS, + block_size=BLOCK, + ) + return blocks, own + + +def select_dense(blocks, case: Case): + from sglang.srt.layers.attention.dsv4.v41_indexer.dense_blocks import ( + _consume_prefill_blocks, + ) + + return _consume_prefill_blocks( + data=case.data, + kv=case.kv, + topk=TOPK, + blocks=blocks, + block_size=BLOCK, + ) + + def picks(positions: torch.Tensor, row: int) -> set: return set(positions[row].tolist()) - {-1} @@ -143,23 +195,15 @@ def picks(positions: torch.Tensor, row: int) -> set: "DeepGEMM's paged sparse MQA logits need SM100", ) class TestPrefillSparseIndexer(CustomTestCase): - def setUp(self): - from sglang.srt.layers.attention.dsv4.candidate_indexer_deep_gemm import ( - DeepGemmCandidateIndexer, - ) - - self.sparse = DeepGemmCandidateIndexer(TOPK_BLOCKS, BLOCK) - self.dense = DenseCandidateIndexer(TOPK_BLOCKS, BLOCK) - @torch.inference_mode() def test_publish_prefill_is_the_torch_block_selection(self): - """The published blocks equal `select_candidate_blocks` block for block + """The published blocks equal `select_candidate_block_ids` block for block (ragged lengths, an empty row, block counts above and below 2048), and the source's own top-k from the same tiled pass is the dense one.""" for rows, ctx in CASES: with self.subTest(rows=rows, ctx=ctx): case = make_case(rows, ctx, seed=rows + ctx) - table, own = publish(self.sparse, case.inputs) + table, own = publish_sparse(case) expected = reference_blocks(case) nb = expected.shape[1] # INT32_MAX padding lands in a spare column instead of a block @@ -167,7 +211,7 @@ def test_publish_prefill_is_the_torch_block_selection(self): got.scatter_(1, table.blocks.clamp_max(nb).long(), True) self.assertTrue(torch.equal(got[:, :nb], expected)) self.assertFalse(got[0, :nb].any(), "an empty row keeps no block") - _, own_dense = publish(self.dense, case.inputs) + _, own_dense = publish_dense(case) for r in range(rows): self.assertEqual(picks(own, r), picks(own_dense, r), r) @@ -179,12 +223,8 @@ def test_select_prefill_matches_the_dense_implementation(self): for rows, ctx in CASES: with self.subTest(rows=rows, ctx=ctx): case = make_case(rows, ctx, seed=rows * 3 + ctx) - got = select( - self.sparse, publish(self.sparse, case.inputs)[0], case.inputs - ) - want = select( - self.dense, publish(self.dense, case.inputs)[0], case.inputs - ) + got = select_sparse(publish_sparse(case)[0], case.data, case.k_cache) + want = select_dense(publish_dense(case)[0], case) keep = reference_blocks(case).repeat_interleave(BLOCK, dim=1) self.assertEqual(picks(got, 0), set(), "an empty row selects nothing") for r in range(1, rows): @@ -209,19 +249,39 @@ def test_prefill_tail_rebuilds_the_last_rows(self): what the full table selects for those rows.""" rows, tail = 128, 16 case = make_case(rows, 20000, seed=7) - table, _ = publish(self.sparse, case.inputs) + table, _ = publish_sparse(case) idx = torch.arange(rows - tail, rows, device="cuda") - sub = self.sparse.prefill_tail(table, [tail]) + sub = table.tail([tail]) self.assertTrue(torch.equal(sub.blocks, table.blocks[idx])) self.assertTrue(torch.equal(sub.valid_lens, table.valid_lens[idx])) self.assertTrue(torch.equal(sub.compress_lens, table.compress_lens[idx])) - full = select(self.sparse, table, case.inputs) - part = select(self.sparse, sub, rows_of(case.inputs, idx)) + full = select_sparse(table, case.data, case.k_cache) + part = select_sparse(sub, rows_of(case.data, idx), case.k_cache) for r in range(tail): a, b = picks(full, rows - tail + r), picks(part, r) self.assertGreaterEqual(len(a & b), MIN_TAIL_OVERLAP * len(a), r) + def test_block_ids_tail_does_not_sync_the_host(self): + """The late-layer tail of published block ids is cut on the device + without a host sync while earlier work is still queued.""" + from sglang.srt.layers.attention.dsv4.v41_indexer.dense_blocks import ( + BlockIds, + ) + + busy = torch.randn(4096, 4096, device="cuda") + for _ in range(8): + busy = busy @ busy + blocks = torch.arange(40, dtype=torch.int32, device="cuda").view(20, 2) + ids = BlockIds(blocks=blocks, rows_per_request=[5, 0, 10, 5]) + torch.cuda.set_sync_debug_mode("error") + try: + tail = ids.tail([2, 0, 3, 5]) + finally: + torch.cuda.set_sync_debug_mode("default") + rows = [3, 4, 12, 13, 14, 15, 16, 17, 18, 19] + torch.testing.assert_close(tail.blocks, blocks[rows]) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/layers/test_dsv41_candidate_blocks.py b/test/registered/unit/layers/test_dsv41_candidate_blocks.py index ff33378fe5b3..39d023876f88 100644 --- a/test/registered/unit/layers/test_dsv41_candidate_blocks.py +++ b/test/registered/unit/layers/test_dsv41_candidate_blocks.py @@ -3,82 +3,116 @@ import torch from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( - candidate_block_mask, select_candidate_block_ids, - select_candidate_blocks, + topk_among_blocks, ) -from sglang.srt.layers.attention.dsv4.candidate_indexer import PrefillCandidateBlocks +from sglang.srt.layers.attention.dsv4.v41_indexer.dense_blocks import BlockIds +from sglang.srt.layers.attention.dsv4.v41_indexer.types import get_tail_row_indices from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=3, suite="base-a-test-cpu") -class TestPrefillCandidateBlocks(CustomTestCase): +def id_sets(blocks: torch.Tensor) -> list: + return [set(row.tolist()) - {-1} for row in blocks] + + +def masked_topk(scores, lens, blocks, k, block_size): + """The reference: -inf outside the chosen blocks and past each row's length, + then a plain top-k; -1 where fewer than k candidates are finite.""" + rows, width = scores.shape + keep = torch.zeros(rows, width, dtype=torch.bool) + for r in range(rows): + for b in blocks[r].tolist(): + if b >= 0: + keep[r, b * block_size : (b + 1) * block_size] = True + keep &= torch.arange(width)[None, :] < lens[:, None] + masked = scores.masked_fill(~keep, -torch.inf) + top = masked.topk(min(k, width), dim=-1) + return [ + {c for c, v in zip(ci.tolist(), vi.tolist()) if v > -torch.inf} + for ci, vi in zip(top.indices, top.values) + ] + + +class TestCandidateBlockIds(CustomTestCase): def test_causal_partial_blocks_and_forced_newest_block(self): scores = torch.tensor([[100.0] * 8 + [50.0] * 8 + [-10.0] * 3] * 4) lengths = torch.tensor([[0], [1], [9], [19]]) scores.masked_fill_(torch.arange(19)[None, :] >= lengths, -torch.inf) original_scores = scores.clone() - expected = torch.tensor( - [ - [False] * 19, - [True] * 8 + [False] * 11, - [True] * 16 + [False] * 3, - [True] * 8 + [False] * 8 + [True] * 3, - ] - ) blocks = select_candidate_block_ids( logits=scores, compress_lens=lengths, topk_blocks=2, block_size=8 ) self.assertEqual(blocks.dtype, torch.int32) self.assertEqual(tuple(blocks.shape), (4, 2)) - torch.testing.assert_close(blocks[0], torch.tensor([-1, -1], dtype=torch.int32)) - torch.testing.assert_close( - candidate_block_mask(blocks=blocks, width=19, block_size=8), expected - ) - torch.testing.assert_close( - select_candidate_blocks( - logits=scores, compress_lens=lengths, topk_blocks=2, block_size=8 - ), - expected, - ) + # nothing visible; one partial block; two blocks; the newest (partial) + # block is forced in ahead of the 50-score one + self.assertEqual(id_sets(blocks), [set(), {0}, {0, 1}, {0, 2}]) torch.testing.assert_close(scores, original_scores) - def test_underfilled_and_empty_candidates(self): + def test_underfilled_candidates_keep_every_block(self): for width in (0, 1, 7, 8, 9): with self.subTest(width=width): - scores = torch.zeros((2, width)) blocks = select_candidate_block_ids( - logits=scores, compress_lens=width, topk_blocks=2048, block_size=8 + logits=torch.zeros((2, width)), + compress_lens=width, + topk_blocks=2048, + block_size=8, + ) + expected = set(range((width + 7) // 8)) + self.assertEqual(id_sets(blocks), [expected, expected]) + + +class TestTopkAmongBlocks(CustomTestCase): + def test_matches_masked_topk(self): + """Gathering the chosen blocks and taking the top-k picks the same + positions as masking everything else to -inf, including causal cuts, + -1 padded blocks and rows with fewer than k candidates.""" + g = torch.Generator().manual_seed(0) + for rows, width, topk_blocks, k in ( + (7, 100, 3, 16), + (5, 1000, 20, 64), + (3, 13, 4, 32), + (2, 8, 1, 4), + ): + with self.subTest(width=width, k=k): + scores = torch.randn(rows, width, generator=g) + lens = torch.randint(0, width + 1, (rows,), generator=g) + visible = scores.masked_fill( + torch.arange(width)[None, :] >= lens[:, None], -torch.inf + ) + blocks = select_candidate_block_ids( + visible, lens[:, None], topk_blocks=topk_blocks, block_size=8 ) - torch.testing.assert_close( - candidate_block_mask(blocks=blocks, width=width, block_size=8), - torch.ones_like(scores, dtype=torch.bool), + picked = topk_among_blocks(scores, lens, blocks, k, block_size=8) + self.assertEqual(tuple(picked.shape), (rows, k)) + self.assertEqual( + id_sets(picked), masked_topk(scores, lens, blocks, k, 8) ) - blocks = torch.full((3, 2), -1, dtype=torch.int32) - self.assertFalse( - candidate_block_mask(blocks=blocks, width=19, block_size=8).any() - ) - def test_replay_tail_keeps_request_boundaries_and_empty_tails(self): - requests = [torch.arange(n * 2).reshape(n, 2) for n in (5, 0, 3)] - candidates = PrefillCandidateBlocks(request_blocks=requests) - tail = candidates.tail([2, 0, 0]) + def test_no_candidate_selects_nothing(self): + scores = torch.randn(2, 16) + blocks = torch.full((2, 2), -1, dtype=torch.int32) + picked = topk_among_blocks(scores, torch.tensor([16, 0]), blocks, 4, 8) + self.assertTrue((picked == -1).all()) + + +class TestBlockIdsTail(CustomTestCase): + def test_tail_keeps_each_requests_last_rows(self): + blocks = torch.arange(16, dtype=torch.int32).view(8, 2) + ids = BlockIds(blocks=blocks, rows_per_request=[5, 0, 3]) + tail = ids.tail([2, 0, 0]) + torch.testing.assert_close(tail.blocks, blocks[3:5]) + self.assertEqual(tail.rows_per_request, [2, 0, 0]) self.assertEqual( - [tuple(b.shape) for b in tail.request_blocks], [(2, 2), (0, 2), (0, 2)] - ) - torch.testing.assert_close(tail.request_blocks[0], requests[0][3:]) - self.assertEqual(tail.request_blocks[0].data_ptr(), requests[0][3:].data_ptr()) - self.assertEqual([b.shape[0] for b in candidates.request_blocks], [5, 0, 3]) - - def test_nonfinite_blocks_match_mask_selection(self): - logits = torch.tensor([[float("nan")] * 8 + [1.0] * 8 + [-torch.inf] * 8]) - kwargs = dict(logits=logits, compress_lens=24, topk_blocks=3, block_size=8) - blocks = select_candidate_block_ids(**kwargs) - torch.testing.assert_close( - candidate_block_mask(blocks=blocks, width=24, block_size=8), - select_candidate_blocks(**kwargs), + get_tail_row_indices( + full_rows_per_request=[5, 0, 3], + tail_rows_per_request=[2, 0, 3], + device="cpu", + ).tolist(), + [3, 4, 5, 6, 7], ) diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py index 3dbb65fab2b4..d9399f2195c8 100644 --- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py +++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py @@ -908,47 +908,91 @@ def run(rows_per_chunk): self.assertTrue(torch.equal(run(rows_per_chunk), expected)) -class TestChunkedCandidatePublisher(CustomTestCase): - def test_each_deep_gemm_call_receives_one_tensor_schedule(self): - from sglang.srt.layers.attention.dsv4 import candidate_indexer_deep_gemm as mod - - num_rows, width = 5, 16 - chunks = [slice(0, 2), slice(2, 4), slice(4, 5)] - plans = [torch.tensor([i], dtype=torch.uint8) for i in range(len(chunks))] - topk_plans = [torch.tensor([i], dtype=torch.int32) for i in range(len(chunks))] +class TestChunkedPagedDecode(CustomTestCase): + """An eager decode whose paged-MQA metadata is split into row chunks (#40637): + every DeepGEMM call must get its own chunk's schedule and every top-k its own + rows and plan, in the candidate publisher and in the dense indexer alike.""" + + num_rows, width = 5, 16 + chunks = [slice(0, 2), slice(2, 4), slice(4, 5)] + + def _inputs(self): + from sglang.srt.layers.attention.dsv4.v41_indexer import DecodeInputs + from sglang.srt.layers.attention.dsv4.v41_indexer.scoring import ( + DeepGEMMDecodeData, + ) + + num_rows, width = self.num_rows, self.width + self.plans = [torch.tensor([i], dtype=torch.uint8) for i in range(3)] + self.topk_plans = [torch.tensor([i], dtype=torch.int32) for i in range(3)] metadata = SimpleNamespace( compressed_seq_lens=torch.full((num_rows, 1), width, dtype=torch.int32), page_table=torch.zeros((num_rows, 1), dtype=torch.int32), - deep_gemm_metadata=plans, + deep_gemm_metadata=self.plans, max_compressed_seq_len=width, compressed_page_size=64, - topk_metadata_chunks=topk_plans, + topk_metadata_chunks=self.topk_plans, use_topk_v2=True, - row_chunks=lambda: list(zip(chunks, plans)), + row_chunks=lambda: list(zip(self.chunks, self.plans)), ) - inputs = SimpleNamespace( + inputs = DecodeInputs( + indexer=None, + layer_id=0, + compress_ratio=1, + freqs_cis=None, + x=None, + q_lora=None, + positions=None, + req_rows=torch.arange(num_rows), + paged_metadata=metadata, + is_verify=True, + ) + data = DeepGEMMDecodeData( q_fp4=torch.zeros((num_rows, 1, 2, 64), dtype=torch.int8), q_sf=torch.zeros((num_rows, 1, 2), dtype=torch.int32), - k_cache=torch.zeros((1, 64, 1, 68), dtype=torch.uint8), weights=torch.zeros((num_rows, 2), dtype=torch.float32), - metadata=metadata, - request_ids=torch.arange(num_rows), - num_rows=num_rows, + k_cache=torch.zeros((1, 64, 1, 68), dtype=torch.uint8), ) - page_indices = torch.full((num_rows, 4), -1, dtype=torch.int32) - raw_indices = torch.full_like(page_indices, -1) - indexer = object.__new__(mod.DeepGemmCandidateIndexer) - indexer.topk_blocks = 2 + return inputs, data + + def _selection(self): + from sglang.srt.layers.attention.dsv4.v41_indexer import Selection + page_indices = torch.full((self.num_rows, 4), -1, dtype=torch.int32) + return Selection(page_indices=page_indices, raw_indices=page_indices.clone()) + + def _mocks(self): + width = self.width deep_gemm = MagicMock( side_effect=lambda q, *_args: torch.zeros( (q[0].shape[0], width), dtype=torch.float32 ) ) - topk = MagicMock() + return deep_gemm, MagicMock() + + def _assert_chunked_calls(self, deep_gemm, topk): + self.assertEqual(deep_gemm.call_count, len(self.chunks)) + for call, plan in zip(deep_gemm.call_args_list, self.plans): + self.assertIs(call.args[5], plan) + self.assertIsInstance(call.args[5], torch.Tensor) + self.assertEqual( + [call.kwargs["rows"] for call in topk.call_args_list], self.chunks + ) + for call, plan in zip(topk.call_args_list, self.topk_plans): + self.assertIs(call.kwargs["topk_metadata"], plan) + + def test_candidate_publisher_takes_one_schedule_per_call(self): + from sglang.srt.layers.attention.dsv4.v41_indexer import sparse_table as mod + + inputs, data = self._inputs() + backend = object.__new__(mod.SparseTableBackend) + backend.token_to_kv_pool = None + backend.topk_blocks = 2 + deep_gemm, topk = self._mocks() event = MagicMock() stream = MagicMock() with ( + patch.object(mod, "get_deep_gemm_decode_data", return_value=data), patch.object(mod, "deep_gemm_fp4_paged_mqa_logits", deep_gemm), patch.object(mod, "topk_transform_paged_from_metadata", topk), patch.object( @@ -979,40 +1023,67 @@ def test_each_deep_gemm_call_receives_one_tensor_schedule(self): patch.object(mod.torch.cuda, "Event", return_value=event), patch.object(mod.torch.cuda, "current_stream", return_value=stream), ): - table = indexer.publish_decode(inputs, page_indices, raw_indices) + table = backend.publish_decode(inputs, self._selection()) - self.assertEqual(deep_gemm.call_count, len(chunks)) - for call, plan in zip(deep_gemm.call_args_list, plans): - self.assertIs(call.args[5], plan) - self.assertIsInstance(call.args[5], torch.Tensor) - self.assertEqual([call.kwargs["rows"] for call in topk.call_args_list], chunks) - for call, plan in zip(topk.call_args_list, topk_plans): - self.assertIs(call.kwargs["topk_metadata"], plan) - self.assertEqual(table.blocks.shape, (num_rows, indexer.topk_blocks)) + self._assert_chunked_calls(deep_gemm, topk) + self.assertEqual(table.blocks.shape, (self.num_rows, backend.topk_blocks)) self.assertEqual(table.phys_blocks.shape, table.blocks.shape) - self.assertEqual(table.valid_lens.shape, (num_rows,)) + self.assertEqual(table.valid_lens.shape, (self.num_rows,)) event.record.assert_called_once_with(stream) + def test_dense_decode_takes_one_schedule_per_call(self): + from sglang.srt.layers.attention.dsv4.v41_indexer import full_topk as mod + + inputs, data = self._inputs() + indexer = object.__new__(mod.FullTopKIndexer) + indexer.token_to_kv_pool = None + indexer.use_deep_gemm_decode = True + deep_gemm, topk = self._mocks() + with ( + patch.object(mod, "get_deep_gemm_decode_data", return_value=data), + patch.object(mod, "deep_gemm_fp4_paged_mqa_logits", deep_gemm), + patch.object(mod, "topk_transform_paged_from_metadata", topk), + ): + indexer.topk_decode(inputs, self._selection()) + + self._assert_chunked_calls(deep_gemm, topk) + class TestCandidateIndexerGating(CustomTestCase): def test_candidate_indexer_gating(self): - from sglang.srt.layers.attention.dsv4 import candidate_indexer + from sglang.srt.layers.attention.dsv4 import v41_indexer + from sglang.srt.layers.attention.dsv4.v41_indexer.dense_blocks import ( + DenseBlocksBackend, + ) - def platform(sm): - return patch.object( - candidate_indexer, "get_platform", lambda: SimpleNamespace(device_sm=sm) - ) + def sm100(value): + return patch.object(v41_indexer, "is_sm100_or_newer", lambda: value) + + def make(topk_blocks, block_size): + # The prefill choice reads the real DeepGEMM build and is cached. + with patch.object(v41_indexer, "_use_deep_gemm_prefill", lambda: False): + return v41_indexer.make_candidate_indexer( + token_to_kv_pool=None, + req_to_token=torch.zeros((1, 1), dtype=torch.int32), + page_size=256, + candidate_topk_blocks=topk_blocks, + candidate_block_size=block_size, + ) flag = "sglang.srt.layers.deep_gemm_wrapper.configurer.DEEPGEMM_PAGED_SPARSE_MQA_LOGITS" - # V4 models have no candidate source; Hopper selects through masks inline. - with platform(100), patch(flag, True): - self.assertIsNone(candidate_indexer.make_candidate_indexer(0, 8)) - with platform(90), patch(flag, False): - self.assertIsNone(candidate_indexer.make_candidate_indexer(2048, 8)) + # V4 configs carry no candidate fields (0 / 0): no layer is a candidate + # source, so Blackwell must not build the DeepGEMM backend for them. + # Hopper serves candidates on dense block ids. + for is_sm100, topk_blocks, block_size in ((True, 0, 0), (False, 2048, 8)): + with self.subTest(sm100=is_sm100, topk_blocks=topk_blocks): + with sm100(is_sm100), patch(flag, False): + prefill, decode = make(topk_blocks, block_size) + self.assertIsInstance(prefill, DenseBlocksBackend) + self.assertIsInstance(decode, DenseBlocksBackend) # Blackwell without DeepGEMM's sparse logits fails instead of falling back. - with platform(100), patch(flag, False): + with sm100(True), patch(flag, False): with self.assertRaises(RuntimeError): - candidate_indexer.make_candidate_indexer(2048, 8) + make(2048, 8) if __name__ == "__main__":