From bd68d89c96dfea20c39ad91861c6f7979e477c7d Mon Sep 17 00:00:00 2001 From: DarkSharpness <2040703891@qq.com> Date: Thu, 24 Sep 2026 02:27:58 -0700 Subject: [PATCH 01/16] [DSv4.1] Move the low-ratio index top-k into dsv4/low_ratio_indexer The ratio-1/2 index layers' top-k selection moves out of deepseek_v4_backend.py into its own package. The attention backend keeps only the dispatch (dense / candidate source / candidate consumer) and an opaque CandidateMetadata on the forward metadata. - DenseIndexer: the plain top-k (DeepGEMM on SM100, the portable path with the fp4 decode logits kernel elsewhere), plus the captured prefill. - CandidateIndexer: publish/consume per prefill/decode, composed from a DeepGEMM backend (sparse table), a DeepGEMM CP backend (block ids) and a torch backend (masks). - Shared scoring in deep_gemm_utils / torch_utils; the flat index K is gathered only by the dense-scoring paths. Behavior-preserving, including the chunked paged-MQA eager decode from #40637. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../layers/attention/deepseek_v4_backend.py | 700 ++++------------- .../attention/dsv4/candidate_indexer.py | 277 ------- .../dsv4/candidate_indexer_deep_gemm.py | 511 ------------ .../attention/dsv4/dense_prefill_indexer.py | 277 ------- .../dsv4/low_ratio_indexer/__init__.py | 99 +++ .../dsv4/low_ratio_indexer/block_math.py | 92 +++ .../low_ratio_indexer/candidate_indexer.py | 87 +++ .../attention/dsv4/low_ratio_indexer/caps.py | 39 + .../low_ratio_indexer/deep_gemm_backend.py | 726 ++++++++++++++++++ .../dsv4/low_ratio_indexer/deep_gemm_utils.py | 284 +++++++ .../dsv4/low_ratio_indexer/dense_indexer.py | 197 +++++ .../dsv4/low_ratio_indexer/inputs.py | 80 ++ .../dsv4/low_ratio_indexer/torch_backend.py | 152 ++++ .../dsv4/low_ratio_indexer/torch_utils.py | 178 +++++ .../srt/layers/attention/dsv4/metadata.py | 27 + .../runner_utils/capture_mode.py | 11 + 16 files changed, 2116 insertions(+), 1621 deletions(-) delete mode 100644 python/sglang/srt/layers/attention/dsv4/candidate_indexer.py delete mode 100644 python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py delete mode 100644 python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py create mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/__init__.py create mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/block_math.py create mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/candidate_indexer.py create mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/caps.py create mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_backend.py create mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py create mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/dense_indexer.py create mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/inputs.py create mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_backend.py create mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_utils.py diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 9cd3961c5f4b..c4016043f0de 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -29,7 +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.kv_layout import KVLayout from sglang.kernels.ops.attention.dsv4.metadata_kernel import ( fill_all_compressed_indices, @@ -56,37 +55,32 @@ ) 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, - mask_topk_scores, - published_masks, - select_candidate_blocks, -) 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, - deep_gemm_fp4_paged_mqa_logits, - topk_transform_paged_from_metadata, +from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin +from sglang.srt.layers.attention.dsv4.low_ratio_indexer import ( + CandidateMetadata, + CapturedPrefillInputs, + DecodeInputs, + DenseIndexer, + PrefillInputs, + Selection, + has_dense_fp4_indexer, + is_sm100_or_newer, + make_candidate_indexer, ) 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 ( @@ -146,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. @@ -260,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, @@ -1165,12 +1124,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), + 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.dense_indexer = DenseIndexer( + 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] @@ -1541,9 +1505,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() @@ -2775,7 +2737,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() @@ -2788,9 +2749,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) @@ -2822,14 +2781,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: @@ -2887,9 +2847,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 @@ -2919,7 +2877,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: @@ -3120,201 +3078,59 @@ 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 + """Select this index layer's compressed positions: a dense layer on the + dense indexer, a candidate source or consumer on the candidate indexer.""" + 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) + 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, ) - 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 - ) - - 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.candidate_indexer.publish_decode(inputs, out) + self._publish_candidate_metadata(published) + elif is_consumer: + self.candidate_indexer.consume_decode(inputs, published, out) + else: + self.dense_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.candidate_indexer.publish_prefill(inputs, out) + self._publish_candidate_metadata(published) + elif is_consumer: + self.candidate_indexer.consume_prefill(inputs, published, out) + else: + self.dense_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 = ( @@ -3323,329 +3139,101 @@ 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. + # two-level selection collapses to the plain top-k. + 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) - columns = torch.arange(width, device=lens.device) - 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, - ) - lens_c = lens[rows].unsqueeze(-1) - # Columns past a row's length hold garbage. - s = logits.masked_fill(columns[None, :] >= lens_c, -torch.inf) - idx = s.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values - reach = idx < lens_c - slots = page_table[rows].gather(-1, idx // page_size) * page_size + ( - idx % page_size - ) - page_indices[rows, :topk] = torch.where(reach, slots, -1).to(torch.int32) - if raw_indices is not None: - raw_indices[rows, :topk] = torch.where(reach, idx, -1).to(torch.int32) - - 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.dense_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 f6df650b7df8..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/candidate_indexer.py +++ /dev/null @@ -1,277 +0,0 @@ -from __future__ import annotations - -from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import List, Optional, Tuple, Union - -import msgspec -import torch -import torch.nn.functional as F - -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 - - -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( - logits: torch.Tensor, - compress_lens: Union[torch.Tensor, int], - topk_blocks: int, - block_size: int, -) -> torch.return_types.topk: - 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, - ) - 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, - block_size: int, -) -> torch.Tensor: - top = _candidate_block_topk( - logits=logits, - compress_lens=compress_lens, - topk_blocks=topk_blocks, - block_size=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] 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 26f308133897..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py +++ /dev/null @@ -1,511 +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 candidate_row_lens -from sglang.kernels.ops.attention.dsv4.candidate_table import ( - amax8_varlen, - sort_candidate_blocks, -) -from sglang.kernels.ops.attention.dsv4.topk import ( - plan_topk_v2, - topk_transform_bf16_small, - topk_transform_paged_v2, - topk_transform_ragged_v2, -) -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 ( - DenseCandidateIndexer, - score_tiles, -) -from sglang.srt.layers.attention.dsv4.indexer import ( - deep_gemm_fp4_paged_mqa_logits, - topk_transform_paged_from_metadata, -) - -CANDIDATE_BLOCK_SIZE = 8 # positions per block; DeepGEMM accepts 8 or 16 - - -@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 - - -def amax_topk_blocks( - logits: torch.Tensor, - seq_lens: torch.Tensor, - nblocks: torch.Tensor, - topk_blocks: int, - max_seq_len: Optional[int] = None, -) -> torch.Tensor: - """Per row the ``topk_blocks`` blocks of 8 positions with the largest block - maximum among its first ``seq_lens[b]`` positions, the newest block always - included: block ids in no particular order, ``-1`` past the row's count. - ``nblocks`` is ``ceil(seq_lens / 8)`` as int32.""" - rows = logits.shape[0] - block = CANDIDATE_BLOCK_SIZE - if max_seq_len is None: - max_seq_len = logits.shape[1] - # NOTE: plan cannot be the previous kernel of topk_transform_paged_v2 - plan = plan_topk_v2(nblocks) - # block maxima, the newest block +inf; the top-k reads each row up to nblocks - # only, so nothing past a row's keys is initialised (v2 needs stride % 4 == 0) - keys = logits.new_empty(rows, -(-max_seq_len // (4 * block)) * 4) - amax8_varlen(logits, seq_lens, out=keys) - blocks = torch.empty(rows, topk_blocks, dtype=torch.int32, device=logits.device) - topk_transform_paged_v2(keys, nblocks, None, blocks, 1, plan) - return blocks - - -def build_sparse_indexer_schedule( - blocks: torch.Tensor, - seq_lens: torch.Tensor, - page_table: torch.Tensor, - page_size: int, - q_dtype: torch.dtype, - request_ids: torch.Tensor, -) -> torch.Tensor: - """DeepGEMM's schedule for the published blocks: ``seq_lens`` ``[rows]`` - int32, ``page_table`` ``[rows, pages]`` int32 at the index pool's page size. - ``request_ids`` ``[rows]`` int32 lets DeepGEMM pair two rows of a request on - one KV pass; each row keeps its own block list and output layout, and paired - rows must share their page-table row.""" - import deep_gemm - - return deep_gemm.get_paged_sparse_mqa_logits_metadata( - seq_lens.contiguous(), - page_table, - request_ids, - page_size, - blocks, - q_dtype, - CANDIDATE_BLOCK_SIZE, - ) - - -def sparse_logits( - q_fp4: torch.Tensor, - q_sf: torch.Tensor, - k_cache: torch.Tensor, - weights: torch.Tensor, - table: SparseBlockTable, -) -> torch.Tensor: - """bf16 logits ``[rows, topk_blocks * 8]`` of the published blocks: ``q_fp4`` - ``[rows, 1, heads, 64]`` int8 with ``q_sf`` ``[rows, 1, heads]`` int32 (packed - ue8m0), ``k_cache`` ``[pages, page_size, 1, 68]`` uint8 whose page stride is - a multiple of 512 bytes, ``weights`` ``[rows, heads]`` bf16.""" - import deep_gemm - - return deep_gemm.fp8_fp4_paged_sparse_mqa_logits( - (q_fp4, q_sf), - k_cache, - weights, - table.schedule, - table.blocks.shape[1], - CANDIDATE_BLOCK_SIZE, - ) - - -def topk_transform_sparse( - logits: torch.Tensor, - valid_lens: torch.Tensor, - table: SparseBlockTable, - page_indices: torch.Tensor, -) -> None: - """Top-``k`` (``k = page_indices.shape[1]``) of every row of the sparse - ``logits`` (bf16 ``[rows, topk_blocks * 8]``) within its first ``valid_lens[b]`` - columns, written as pool slots, ``-1`` where a row has fewer than ``k`` valid - columns, in no particular order.""" - topk_transform_bf16_small( - logits, valid_lens, table.phys_blocks, page_indices, CANDIDATE_BLOCK_SIZE - ) - - -@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, - ) - - 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, 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 score_tiles(inputs, 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, - ) - # 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 d4a478a24de1..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py +++ /dev/null @@ -1,277 +0,0 @@ -from __future__ import annotations - -from typing import Iterator, List, Optional - -import torch - -from sglang.srt.layers.attention.dsv4.candidate_indexer import ( - CandidateIndexer, - CandidateMetadata, - PrefillCandidateBlocks, - PrefillIndexerInputs, - candidate_block_mask, - mask_topk_scores, - select_candidate_block_ids, -) -from sglang.srt.layers.attention.mqa_logits_utils import ( - mqa_logits_row_bytes, - mqa_logits_rows_per_chunk, -) -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 = _rows_per_chunk(row, width, heads=q[0].shape[1]) - 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 _rows_per_chunk(rows: int, width: int, *, heads: int) -> int: - """Query rows per logits tile so one fp32 [rows, width] tile fits the - budget; the row count stays a multiple of the kernel's row alignment.""" - row_alignment = 128 // heads - rows_per_chunk = mqa_logits_rows_per_chunk( - num_rows=ceil_align(rows, row_alignment), - row_bytes=mqa_logits_row_bytes(width), - budget_bytes=_SCORE_BUDGET_BYTES, - ) - if rows_per_chunk is None: - return rows - return max(row_alignment, rows_per_chunk // row_alignment * row_alignment) - - -def score_tiles( - inputs: PrefillIndexerInputs, *, width_align: int = 4 -) -> Iterator[tuple[slice, torch.Tensor]]: - """The dense scores of ``inputs`` row tile by row tile under the budget: - ``(rows, logits)`` with fp32 ``logits[i, j]`` the score of query row - ``rows.start + i`` against ``kv[request_starts + j]``, garbage past the row's - ``compress_lens``; the width is the batch's largest context aligned to - ``width_align`` columns.""" - from deep_gemm import fp8_fp4_mqa_logits - - rows = inputs.num_rows - width = ceil_align(max(inputs.lens_per_request, default=0), width_align) - if rows == 0 or width == 0: - return - rows_per_chunk = _rows_per_chunk(rows, width, heads=inputs.q_fp4.shape[1]) - for offset in range(0, rows, rows_per_chunk): - tile = slice(offset, min(offset + rows_per_chunk, rows)) - starts = inputs.request_starts[tile] - yield ( - tile, - fp8_fp4_mqa_logits( - (inputs.q_fp4[tile], inputs.q_sf[tile]), - inputs.kv, - inputs.weights[tile], - starts, - starts + inputs.compress_lens[tile], - False, - width, - ), - ) - - -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 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/low_ratio_indexer/__init__.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/__init__.py new file mode 100644 index 000000000000..954e5f9a89ce --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/__init__.py @@ -0,0 +1,99 @@ +"""The top-k selection of the DeepSeek V4.1 ratio-1/2 index layers. + +The dense layers go through ``DenseIndexer``; a candidate source and its +consumers go through the ``CandidateIndexer`` that ``make_candidate_indexer`` +builds. The two share only the scoring utilities. Nothing else in this package +is meant to be named from outside it. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.runtime_context import get_parallel + +from .candidate_indexer import CandidateIndexer +from .caps import ( + has_dense_fp4_indexer, + is_sm100_or_newer, + use_deep_gemm_prefill, +) +from .dense_indexer import DenseIndexer +from .inputs import ( + CandidateMetadata, + CapturedPrefillInputs, + DecodeInputs, + PrefillInputs, + Selection, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + +__all__ = [ + "CandidateIndexer", + "CandidateMetadata", + "CapturedPrefillInputs", + "DecodeInputs", + "DenseIndexer", + "PrefillInputs", + "Selection", + "has_dense_fp4_indexer", + "is_sm100_or_newer", + "make_candidate_indexer", +] + + +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, +) -> CandidateIndexer: + from .torch_backend import TorchCandidateBackend + + torch_backend = TorchCandidateBackend( + token_to_kv_pool=token_to_kv_pool, req_to_token=req_to_token + ) + # Without candidate blocks no layer is a candidate source or consumer. + if not is_sm100_or_newer() or candidate_topk_blocks <= 0: + return CandidateIndexer(prefill=torch_backend, decode=torch_backend) + + 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 .deep_gemm_backend import ( + DeepGEMMCandidateBackend, + DeepGEMMCPCandidateBackend, + ) + + deep_gemm_backend = DeepGEMMCandidateBackend( + 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, + ) + + if not use_deep_gemm_prefill(): + prefill_backend = torch_backend + elif get_parallel().attn_cp_size > 1: + prefill_backend = DeepGEMMCPCandidateBackend( + 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, + ) + else: + prefill_backend = deep_gemm_backend + return CandidateIndexer(prefill=prefill_backend, decode=deep_gemm_backend) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/block_math.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/block_math.py new file mode 100644 index 000000000000..d39ec84b986d --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/block_math.py @@ -0,0 +1,92 @@ +"""Block-selection math shared by the backends whose candidates are picked from +dense scores: a candidate source keeps whole blocks of ``block_size`` compressed +positions rather than single positions, and its newest block always.""" + +from __future__ import annotations + +from typing import Optional, Union + +import torch +import torch.nn.functional as F + + +def _candidate_block_topk( + logits: torch.Tensor, + compress_lens: Union[torch.Tensor, int], + topk_blocks: int, + block_size: int, +) -> torch.return_types.topk: + 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 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 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, + ) + 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, + block_size: int, +) -> torch.Tensor: + top = _candidate_block_topk( + logits=logits, + compress_lens=compress_lens, + topk_blocks=topk_blocks, + block_size=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] diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/candidate_indexer.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/candidate_indexer.py new file mode 100644 index 000000000000..c967bf6301ef --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/candidate_indexer.py @@ -0,0 +1,87 @@ +"""The two-level candidate scheme of the DeepSeek V4.1 ratio-1/2 index layers. + +The candidate source publishes the blocks of compressed positions it kept, and +the index layers after it select their top-k among those blocks only. What a +publish carries -- a DeepGEMM sparse table, block ids, a mask -- is the backend's +own; the attention backend keeps it alive on the forward metadata and hands it +back to the consumers unread. +""" + +from __future__ import annotations + +from typing import Optional, Protocol, TypeVar + +from .inputs import ( + CandidateMetadata, + DecodeInputs, + PrefillInputs, + Selection, +) + +T = TypeVar("T", bound=CandidateMetadata) + + +class PrefillBackend(Protocol[T]): + def publish_prefill( + self, + inputs: PrefillInputs, + out: Selection, + ) -> Optional[T]: ... + + def consume_prefill( + self, + inputs: PrefillInputs, + published: Optional[T], + out: Selection, + ) -> None: ... + + +class DecodeBackend(Protocol[T]): + def publish_decode( + self, + inputs: DecodeInputs, + out: Selection, + ) -> Optional[T]: ... + + def consume_decode( + self, + inputs: DecodeInputs, + published: Optional[T], + out: Selection, + ) -> None: ... + + +class CandidateIndexer: + def __init__(self, *, prefill: PrefillBackend, decode: DecodeBackend): + self.prefill = prefill + self.decode = decode + + def publish_prefill( + self, + inputs: PrefillInputs, + out: Selection, + ) -> Optional[CandidateMetadata]: + return self.prefill.publish_prefill(inputs, out) + + def consume_prefill( + self, + inputs: PrefillInputs, + published: Optional[CandidateMetadata], + out: Selection, + ) -> None: + self.prefill.consume_prefill(inputs, published, out) + + def publish_decode( + self, + inputs: DecodeInputs, + out: Selection, + ) -> Optional[CandidateMetadata]: + return self.decode.publish_decode(inputs, out) + + def consume_decode( + self, + inputs: DecodeInputs, + published: Optional[CandidateMetadata], + out: Selection, + ) -> None: + self.decode.consume_decode(inputs, published, out) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/caps.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/caps.py new file mode 100644 index 000000000000..92ba855c78b3 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/caps.py @@ -0,0 +1,39 @@ +"""Which low-ratio indexer kernels this build can run.""" + +from __future__ import annotations + +import functools + +import torch + + +@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: + """Whether an eager prefill scores on DeepGEMM rather than torch. 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 + + if not (is_sm100_or_newer() and has_dense_fp4_indexer()): + return False + return ( + get_parallel().attn_cp_size > 1 + or not envs.SGLANG_DSV41_TORCH_PREFILL_INDEXER.get() + ) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_backend.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_backend.py new file mode 100644 index 000000000000..6553d08a10a3 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_backend.py @@ -0,0 +1,726 @@ +"""The candidate scheme on SM100, carried by a DeepGEMM sparse table. + +A source publishes the block ids it kept plus the pool slots and the schedule +``fp8_fp4_paged_sparse_mqa_logits`` needs for them, and a consumer scores only +those blocks straight 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 candidate_row_lens +from sglang.kernels.ops.attention.dsv4.candidate_table import ( + amax8_varlen, + sort_candidate_blocks, +) +from sglang.kernels.ops.attention.dsv4.topk import ( + plan_topk_v2, + topk_transform_bf16_small, + topk_transform_paged_v2, + topk_transform_ragged_v2, +) +from sglang.srt.layers.attention.dsv4.indexer import ( + deep_gemm_fp4_paged_mqa_logits, + topk_transform_paged_from_metadata, +) +from sglang.srt.layers.attention.dsv4.metadata import expand_index_page_table + +from .block_math import ( + candidate_block_mask, + mask_topk_scores, + select_candidate_block_ids, +) +from .deep_gemm_utils import ( + DeepGEMMDecodeData, + DeepGEMMPrefillData, + get_deep_gemm_decode_data, + get_deep_gemm_prefill_data, + get_flat_index_k, + get_index_k_cache, + score_tiles, +) +from .inputs import ( + CandidateMetadata, + DecodeInputs, + PrefillInputs, + Selection, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + +CANDIDATE_BLOCK_SIZE = 8 # positions per block; DeepGEMM accepts 8 or 16 + + +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, 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 + + +class _SparsePrefillTable(_SparseTable): + """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 + topk_blocks: int + + def tail(self, rows_per_request: List[int]) -> _SparsePrefillTable: + rows, start = [], 0 + for n, t in zip(self.rows_per_request, rows_per_request): + rows.append(torch.arange(start + n - t, start + n)) + start += n + rows = torch.cat(rows).to(self.blocks.device) + compress_lens = self.compress_lens[rows].contiguous() + _, valid_lens = candidate_row_lens(compress_lens, self.topk_blocks) + 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, + topk_blocks=self.topk_blocks, + ) + + +class _DensePrefillBlocks(CandidateMetadata, msgspec.Struct): + request_blocks: List[torch.Tensor] + + def tail(self, rows_per_request: List[int]) -> _DensePrefillBlocks: + return _DensePrefillBlocks( + request_blocks=[ + blocks[blocks.shape[0] - length :] + for blocks, length in zip(self.request_blocks, rows_per_request) + ] + ) + + +class DeepGEMMCandidateBackend: + 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 = self._get_prefill_data(inputs) + 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=get_flat_index_k( + data=data, + token_to_kv_pool=self.token_to_kv_pool, + layer_id=inputs.layer_id, + ), + 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() + # Scores only the published blocks, straight out of the pool. + data = self._get_prefill_data(inputs) + 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): + """The source layer's own plain top-k, plus the table for the consumers.""" + 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_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(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: + """Publish an eager forward whose dense logits are bounded by row chunks. + It stays on the current stream so each chunk's full logits can be released + before the next; CUDA-graph metadata always carries one schedule.""" + 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_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 = deep_gemm_sparse_logits( + data.q_fp4, + data.q_sf, + data.k_cache, + data.weights.to(torch.bfloat16), + published, + ) + assert out.raw_indices is None + topk_transform_bf16_small( + logits, + published.valid_lens, + published.phys_blocks, + out.page_indices, + CANDIDATE_BLOCK_SIZE, + ) + + def _get_prefill_data(self, inputs: PrefillInputs) -> Optional[DeepGEMMPrefillData]: + return get_deep_gemm_prefill_data(inputs, self.req_to_token) + + 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] + + +class DeepGEMMCPCandidateBackend: + def __init__( + self, + *, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, + candidate_topk_blocks: int, + candidate_block_size: int, + ) -> 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 + + def publish_prefill(self, inputs: PrefillInputs, out: Selection): + out.reset() + data = self._get_prefill_data(inputs) + if data is None: + return None + selected, request_blocks = _publish_prefill_blocks( + data=data, + kv=self._get_flat_index_k(inputs, data), + topk=inputs.indexer.index_topk, + topk_blocks=self.topk_blocks, + block_size=self.block_size, + ) + data.write_selection(selected=selected, out=out) + return _DensePrefillBlocks(request_blocks) + + def consume_prefill( + self, + inputs: PrefillInputs, + published: Optional[_DensePrefillBlocks], + out: Selection, + ) -> None: + out.reset() + data = self._get_prefill_data(inputs) + if data is None: + return + assert isinstance(published, _DensePrefillBlocks) + selected = _consume_prefill_blocks( + data=data, + kv=self._get_flat_index_k(inputs, data), + topk=inputs.indexer.index_topk, + request_blocks=published.request_blocks, + block_size=self.block_size, + ) + data.write_selection(selected=selected, out=out) + + def _get_prefill_data(self, inputs: PrefillInputs) -> Optional[DeepGEMMPrefillData]: + return get_deep_gemm_prefill_data(inputs, self.req_to_token) + + def _get_flat_index_k(self, inputs: PrefillInputs, data: DeepGEMMPrefillData): + return get_flat_index_k( + data=data, token_to_kv_pool=self.token_to_kv_pool, layer_id=inputs.layer_id + ) + + +def amax_topk_blocks( + logits: torch.Tensor, + seq_lens: torch.Tensor, + nblocks: torch.Tensor, + topk_blocks: int, + max_seq_len: Optional[int] = None, +) -> torch.Tensor: + """Per row the ``topk_blocks`` blocks of 8 positions with the largest block + maximum among its first ``seq_lens[b]`` positions, the newest block always + included: block ids in no particular order, ``-1`` past the row's count. + ``nblocks`` is ``ceil(seq_lens / 8)`` as int32.""" + rows = logits.shape[0] + block = CANDIDATE_BLOCK_SIZE + if max_seq_len is None: + max_seq_len = logits.shape[1] + # NOTE: plan cannot be the previous kernel of topk_transform_paged_v2 + plan = plan_topk_v2(nblocks) + # block maxima, the newest block +inf; the top-k reads each row up to nblocks + # only, so nothing past a row's keys is initialised (v2 needs stride % 4 == 0) + keys = logits.new_empty(rows, -(-max_seq_len // (4 * block)) * 4) + amax8_varlen(logits, seq_lens, out=keys) + blocks = torch.empty(rows, topk_blocks, dtype=torch.int32, device=logits.device) + topk_transform_paged_v2(keys, nblocks, None, blocks, 1, plan) + return blocks + + +def build_sparse_schedule( + blocks: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + page_size: int, + q_dtype: torch.dtype, + request_ids: torch.Tensor, +) -> torch.Tensor: + """DeepGEMM's schedule for the published blocks: ``seq_lens`` ``[rows]`` + int32, ``page_table`` ``[rows, pages]`` int32 at the index pool's page size. + ``request_ids`` ``[rows]`` int32 groups the rows of one request onto one KV + pass; each row keeps its own block list and output layout, and rows grouped + together must share their page-table row.""" + import deep_gemm + + return deep_gemm.get_paged_sparse_mqa_logits_metadata( + seq_lens.contiguous(), + page_table, + request_ids, + page_size, + blocks, + q_dtype, + CANDIDATE_BLOCK_SIZE, + ) + + +def deep_gemm_sparse_logits( + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + k_cache: torch.Tensor, + weights: torch.Tensor, + table: _SparseTable, +) -> torch.Tensor: + """bf16 logits ``[rows, topk_blocks * 8]`` of the published blocks: ``q_fp4`` + ``[rows, 1, heads, 64]`` int8 with ``q_sf`` ``[rows, 1, heads]`` int32 (packed + ue8m0), ``k_cache`` ``[pages, page_size, 1, 68]`` uint8 whose page stride is + a multiple of 512 bytes, ``weights`` ``[rows, heads]`` bf16.""" + import deep_gemm + + return deep_gemm.fp8_fp4_paged_sparse_mqa_logits( + (q_fp4, q_sf), + k_cache, + weights, + table.schedule, + table.blocks.shape[1], + CANDIDATE_BLOCK_SIZE, + ) + + +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, + topk_blocks: int, +) -> _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_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, + topk_blocks=topk_blocks, + ) + + +# 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 layer's own top-``k`` (``k = out_positions.shape[1]``) as + flattened-K columns, ``-1`` padded, unordered, and the chunk's table: both + from one pass over its dense scores, tile by tile. ``index_page_table`` is + ``[rows, index pages]`` 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) + # 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=torch.zeros(logits.shape[0], dtype=torch.int32, device=device), + 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, + topk_blocks=topk_blocks, + ) + + +def select_prefill_table( + *, + table: _SparsePrefillTable, + data: DeepGEMMPrefillData, + k_cache: torch.Tensor, + out_positions: torch.Tensor, +) -> None: + """A consumer's top-``k`` over the published blocks, in the layout + ``publish_prefill_table`` writes; ``k_cache`` is the layer's index-K pool, + ``[pages, page_size, 1, 68]`` uint8.""" + rows, heads = data.q_sf.shape + logits = deep_gemm_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, + ) + # 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, data.request_starts[:, None], 0)) + + +def _publish_prefill_blocks( + *, + data: DeepGEMMPrefillData, + kv: Tuple[torch.Tensor, torch.Tensor], + topk: int, + topk_blocks: int, + block_size: int, +) -> Tuple[torch.Tensor, List[torch.Tensor]]: + """The source 's own plain top-k, plus the ids of its best blocks per + request: ``[rows_b, layermin(topk_blocks, ceil(lc_b / block_size))]`` int32 each.""" + from sglang.kernels.ops.attention.dsv4 import topk_transform_ragged_v2 + + selected = data.empty_selection(topk) + device = data.request_starts.device + request_blocks = [ + torch.empty( + (rows, min(topk_blocks, -(-lc // block_size))), + dtype=torch.int32, + device=device, + ) + for rows, lc in zip(data.rows_per_request, data.lens_per_request) + ] + for tile, logits in score_tiles(data, kv, width_align=4): + lens = data.compress_lens[tile] + for request, rows, request_rows, lc in _requests_in_tile(data, tile): + scores = logits[rows, :lc] + scores.masked_fill_( + torch.arange(lc, device=device)[None, :] >= lens[rows, None], + -torch.inf, + ) + blocks = request_blocks[request][request_rows] + blocks.copy_( + select_candidate_block_ids( + logits=scores, + compress_lens=lens[rows, None], + topk_blocks=blocks.shape[1], + block_size=block_size, + ) + ) + topk_transform_ragged_v2( + logits, + lens, + out_offsets=data.request_starts[tile], + out_indices=selected[tile], + ) + return selected, request_blocks + + +def _consume_prefill_blocks( + *, + data: DeepGEMMPrefillData, + kv: Tuple[torch.Tensor, torch.Tensor], + topk: int, + request_blocks: List[torch.Tensor], + block_size: int, +) -> torch.Tensor: + """A consumer's top-k over its dense scores masked to the published blocks.""" + 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): + starts = data.request_starts[tile] + for request, rows, request_rows, lc in _requests_in_tile(data, tile): + logits[rows, :lc].masked_fill_( + ~candidate_block_mask( + blocks=request_blocks[request][request_rows], + width=lc, + block_size=block_size, + ), + -torch.inf, + ) + topk_transform_ragged_v2( + logits, + data.compress_lens[tile], + out_offsets=starts, + out_indices=selected[tile], + ) + selected[tile] = mask_topk_scores( + scores=logits, indices=selected[tile], offsets=starts + ) + return selected + + +def _requests_in_tile(data: DeepGEMMPrefillData, tile: slice): + """``(request, rows, request_rows, lc)`` of each request with a visible + position and rows in ``tile``: its rows within the tile and within itself.""" + start = 0 + for request, (num_rows, lc) in enumerate( + 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 ( + request, + slice(begin - tile.start, stop - tile.start), + slice(begin - start, stop - start), + lc, + ) + start = end diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py new file mode 100644 index 000000000000..432bcff28fa6 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py @@ -0,0 +1,284 @@ +"""DeepGEMM scoring shared by the dense layers and the candidate scheme. + +On prefill: the chunk's index-K slots, fp4 queries and head weights, the tiled +``fp8_fp4_mqa_logits`` scores over the flattened index K, the plain top-k, and +the write of a selection back into the pool-slot layout attention reads. On +decode: the fp4 queries and the paged pool. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Generator, List, Optional, Tuple + +import msgspec +import torch + +from sglang.srt.layers.attention.mqa_logits_utils import ( + mqa_logits_row_bytes, + mqa_logits_rows_per_chunk, +) +from sglang.srt.utils.common import ceil_align + +from .inputs import ( + DecodeInputs, + PrefillInputs, + Selection, +) + +if TYPE_CHECKING: + 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 masks and block-selection scratch. +_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): + """One index layer's operands on a prefill chunk; the rows of one request are + consecutive, and 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, init: bool = True) -> torch.Tensor: + out = self.request_starts.new_empty((self.num_rows, topk), dtype=torch.int32) + if init: + out.fill_(-1) + return out + + def write_selection(self, selected: torch.Tensor, out: Selection) -> None: + # write_prefill_selection(out=out, data=self, selected=selected) + 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]: + """``q`` ``[rows, heads, 128]`` bf16 to fp4 ``[rows, heads, 64]`` int8 and its + scales ``[rows, heads]`` int32 (packed ue8m0).""" + 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_flat_index_k( + *, + data: DeepGEMMPrefillData, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + layer_id: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """The chunk's index K gathered in flattened-K column order: int8 + ``[columns, 64]`` and int32 ``[columns]``; a copy, for dense scoring only.""" + return token_to_kv_pool.get_low_ratio_index_k_fp4(layer_id, data.k_slots) + + +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, which is + the same for every index layer of the forward.""" + 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 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 = 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]: + """The dense scores of ``data`` row tile by row tile under the budget: + ``(rows, logits)`` with fp32 ``logits[i, j]`` the score of query row + ``rows.start + i`` against ``kv[request_starts + j]``, garbage past the row's + ``compress_lens``; the width is the batch's largest context aligned to + ``width_align`` columns.""" + from deep_gemm import fp8_fp4_mqa_logits + + rows = data.num_rows + width = ceil_align(max(data.lens_per_request, default=0), width_align) + if rows == 0 or width == 0: + return + rows_per_chunk = _rows_per_chunk(rows, width, heads=data.q_fp4.shape[1]) + for offset in range(0, rows, rows_per_chunk): + tile = slice(offset, min(offset + rows_per_chunk, rows)) + starts = data.request_starts[tile] + yield ( + tile, + fp8_fp4_mqa_logits( + (data.q_fp4[tile], data.q_sf[tile]), + kv, + data.weights[tile], + starts, + starts + data.compress_lens[tile], + False, + width, + ), + ) + + +def dense_prefill_topk( + data: DeepGEMMPrefillData, + kv: Tuple[torch.Tensor, torch.Tensor], + *, + topk: int, +) -> torch.Tensor: + """The plain top-``topk`` of each row as 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], + ) + return selected + + +# TODO(dark): make it SM aware; align to SM count to avoid wave quantization. +def _rows_per_chunk(rows: int, width: int, *, heads: int) -> int: + """Query rows per logits tile so one fp32 [rows, width] tile fits the + budget; the row count stays a multiple of the kernel's row alignment.""" + row_alignment = 128 // heads + rows_per_chunk = mqa_logits_rows_per_chunk( + num_rows=ceil_align(rows, row_alignment), + row_bytes=mqa_logits_row_bytes(width), + budget_bytes=_SCORE_BUDGET_BYTES, + ) + if rows_per_chunk is None: + return rows + return max(row_alignment, rows_per_chunk // row_alignment * row_alignment) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/dense_indexer.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/dense_indexer.py new file mode 100644 index 000000000000..d6f7d62f8780 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/dense_indexer.py @@ -0,0 +1,197 @@ +"""The plain top-k of the DeepSeek V4.1 ratio-1/2 index layers outside the +candidate scheme: score every visible compressed position, keep the best. + +DeepGEMM on SM100; elsewhere the torch path, whose decode runs the fp4 decode +logits kernel (the Hopper indexer). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.layers.attention.dsv4.indexer import ( + deep_gemm_fp4_paged_mqa_logits, + topk_transform_paged_from_metadata, +) + +from .caps import is_sm100_or_newer, use_deep_gemm_prefill +from .deep_gemm_utils import ( + dense_prefill_topk, + get_deep_gemm_decode_data, + get_deep_gemm_prefill_data, + get_flat_index_k, + get_index_k_cache, + quantize_index_q, +) +from .inputs import ( + CapturedPrefillInputs, + DecodeInputs, + PrefillInputs, + Selection, +) +from .torch_utils import ( + decode_scores, + prefill_requests, + write_decode, + write_prefill, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + + +class DenseIndexer: + def __init__( + self, + *, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, + ): + 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 = is_sm100_or_newer() + + 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): + return _deep_gemm_prefill_captured(inputs, self.token_to_kv_pool, 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 = get_flat_index_k( + data=data, token_to_kv_pool=self.token_to_kv_pool, layer_id=inputs.layer_id + ) + 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): + # An eager forward whose dense logits are bounded by row 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, + ) + # 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( + inputs: CapturedPrefillInputs, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + 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=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) + columns = torch.arange(width, device=lens.device) + 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, + ) + lens_c = lens[rows].unsqueeze(-1) + # Columns past a row's length hold garbage. + s = logits.masked_fill(columns[None, :] >= lens_c, -torch.inf) + idx = s.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values + reach = idx < lens_c + slots = page_table[rows].gather(-1, idx // page_size) * page_size + ( + idx % page_size + ) + out.page_indices[rows, :topk] = torch.where(reach, slots, -1).to(torch.int32) + if out.raw_indices is not None: + out.raw_indices[rows, :topk] = torch.where(reach, idx, -1).to(torch.int32) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/inputs.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/inputs.py new file mode 100644 index 000000000000..3877876cd858 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/inputs.py @@ -0,0 +1,80 @@ +"""What the attention backend hands a low-ratio indexer backend, and what it gets +back.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Optional + +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 + + +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 itself and gathers the + index K it needs.""" + + 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: the queries and head weights come + projected, and the index K is read paged, as on decode.""" + + indexer: DeepseekV41Indexer + layer_id: int + q: torch.Tensor # [rows, heads, 128] bf16, roped + weights: torch.Tensor # [rows, heads] bf16 head weights + paged_metadata: PagedIndexerMetadata diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_backend.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_backend.py new file mode 100644 index 000000000000..f2b8f0237471 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_backend.py @@ -0,0 +1,152 @@ +"""The candidate scheme on the portable torch path: bf16 scores materialized per +request and selected with ``torch.topk``, the candidates carried by boolean masks. +A source keeps a mask of its best blocks, a consumer scores densely and drops +everything outside it; it needs no page table, so it works wherever torch does. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Optional + +import msgspec +import torch + +from .block_math import ( + mask_topk_scores, + select_candidate_blocks, +) +from .inputs import ( + CandidateMetadata, + DecodeInputs, + PrefillInputs, + Selection, +) +from .torch_utils import ( + decode_scores, + prefill_requests, + write_decode, + write_prefill, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + + +class CandidateMasks(CandidateMetadata, msgspec.Struct): + mask: Optional[torch.Tensor] = None # decode: [rows, width] bool + request_masks: Optional[List[torch.Tensor]] = None # prefill: [rows_b, lc_b] each + + def tail(self, rows_per_request: List[int]) -> CandidateMasks: + assert self.request_masks is not None, "prefill masks missing" + return CandidateMasks( + request_masks=[ + mask[mask.shape[0] - t :] + for mask, t in zip(self.request_masks, rows_per_request) + ] + ) + + +class TorchCandidateBackend: + def __init__( + self, + *, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, + ): + self.token_to_kv_pool = token_to_kv_pool + self.req_to_token = req_to_token + + def publish_prefill(self, inputs: PrefillInputs, out: Selection): + indexer = inputs.indexer + published = [] + for request, chunks in self._prefill_requests(inputs, out): + if request.lc == 0: + # Consumers address masks by request position, including empty requests. + published.append( + torch.zeros(0, 0, dtype=torch.bool, device=inputs.positions.device) + ) + continue + masks = [] + for chunk in chunks: + masks.append( + select_candidate_blocks( + chunk.scores, + chunk.lens[:, None], + topk_blocks=indexer.candidate_topk_blocks, + block_size=indexer.candidate_block_size, + ) + ) + idx = chunk.scores.topk(request.k, dim=-1, sorted=False).indices + write_prefill(out, request, chunk, idx) + published.append(torch.cat(masks) if len(masks) > 1 else masks[0]) + return CandidateMasks(request_masks=published) + + def consume_prefill( + self, + inputs: PrefillInputs, + published: Optional[CandidateMasks], + out: Selection, + ) -> None: + assert published is not None + request_masks = published.request_masks + for request, chunks in self._prefill_requests(inputs, out): + for chunk in chunks: + keep = request_masks[request.index][chunk.rows] + s = chunk.scores.masked_fill(~keep, -torch.inf) + idx = s.topk(request.k, dim=-1, sorted=False).indices + idx = mask_topk_scores(s, idx) + idx = idx.masked_fill(idx < 0, request.lc) + write_prefill(out, request, chunk, idx) + + def publish_decode(self, inputs: DecodeInputs, out: Selection): + indexer = inputs.indexer + d = self._decode_scores(inputs, out) + if d is None: + return None + published = CandidateMasks( + mask=select_candidate_blocks( + d.scores, + d.lens[:, None], + topk_blocks=indexer.candidate_topk_blocks, + block_size=indexer.candidate_block_size, + ) + ) + k = min(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[CandidateMasks], + out: Selection, + ) -> None: + d = self._decode_scores(inputs, out) + if d is None: + return + assert published is not None + consume = published.mask + assert torch.is_tensor(consume) and consume.shape[0] == d.bs + s = d.scores.masked_fill(~consume[:, : d.lmax], -torch.inf) + k = min(inputs.indexer.index_topk, d.lmax) + idx = s.topk(k, dim=-1, sorted=False).indices + idx = mask_topk_scores(s, idx) + idx = idx.masked_fill(idx < 0, d.lmax) + write_decode(out, d, idx) + + def _prefill_requests(self, inputs: PrefillInputs, out: Selection): + return prefill_requests( + inputs=inputs, + out=out, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + ) + + def _decode_scores(self, inputs: DecodeInputs, out: Selection): + return decode_scores( + inputs=inputs, + out=out, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + ) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_utils.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_utils.py new file mode 100644 index 000000000000..22b1ae8744bb --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_utils.py @@ -0,0 +1,178 @@ +"""Torch scoring shared by the dense layers and the candidate scheme: bf16 scores +per request on prefill, the fp4 decode logits kernel on decode, and the writes of +a ``torch.topk`` back into the pool-slot layout attention reads.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Iterator, Optional + +import msgspec +import torch + +from sglang.kernels.ops.attention.dsv4.fp4_indexer import fp4_index_logits_decode + +from .inputs 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 + + +# Arbitrary cap on one bf16 [rows, heads, lc] score chunk; transients run ~3x this. +_SCORE_BUDGET_BYTES = 1 << 30 + + +class RequestScores(msgspec.Struct, frozen=True): + """One request of a prefill chunk, scored row chunk by row chunk.""" + + index: int # position among the chunk's requests, empty ones included + lc: int # compressed positions visible at its newest row; 0 scores nothing + k: int + columns: Optional[torch.Tensor] = None # arange(lc) + tok: Optional[torch.Tensor] = None # [rows_b] its query rows in the chunk + lens: Optional[torch.Tensor] = None # [rows_b] compressed positions each row sees + slots: Optional[torch.Tensor] = None # [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]]]: + """Each request of the chunk with its row chunks' scores; an empty request + comes with no chunk.""" + 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 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: + yield RequestScores(index=b, lc=0, k=0), iter(()) + 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( + index=b, 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 + # Every step 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, _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/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/model_executor/runner_utils/capture_mode.py b/python/sglang/srt/model_executor/runner_utils/capture_mode.py index 38770ea4015e..a95bba4d8962 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,17 @@ def skip_low_ratio_indexer(compress_ratio: int) -> bool: ) +def skip_candidate_filtering() -> bool: + """Whether the captured candidate variant serves only contexts inside the + candidate window, where every block is a candidate and the candidate layers + select like dense ones.""" + 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 From fb9811079ced8a51de6a2a3e8ff6d069e546e413 Mon Sep 17 00:00:00 2001 From: DarkSharpness <2040703891@qq.com> Date: Thu, 24 Sep 2026 09:10:07 -0700 Subject: [PATCH 02/16] [DSv4.1] Port the low-ratio indexer tests; free each score tile Port the tests of the removed candidate_indexer / dense_prefill_indexer / candidate_indexer_deep_gemm modules to the low_ratio_indexer package, and cover the chunked eager decode in the dense indexer and the candidate factory without candidate blocks. The dense and CP prefill loops over score_tiles kept the previous tile's logits alive while the next one was scored, doubling the peak; drop each tile before the next. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../low_ratio_indexer/deep_gemm_backend.py | 59 ++++--- .../dsv4/low_ratio_indexer/deep_gemm_utils.py | 2 + .../attention/test_dense_prefill_indexer.py | 81 ++++++--- .../test_dsv41_prefill_sparse_indexer.py | 145 ++++++++++------ .../layers/test_dsv41_candidate_blocks.py | 8 +- .../unit/layers/test_dsv4_nonpaged_indexer.py | 159 +++++++++++++----- 6 files changed, 312 insertions(+), 142 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_backend.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_backend.py index 6553d08a10a3..8de7a8b14dee 100644 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_backend.py +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_backend.py @@ -632,8 +632,8 @@ def _publish_prefill_blocks( topk_blocks: int, block_size: int, ) -> Tuple[torch.Tensor, List[torch.Tensor]]: - """The source 's own plain top-k, plus the ids of its best blocks per - request: ``[rows_b, layermin(topk_blocks, ceil(lc_b / block_size))]`` int32 each.""" + """The source layer's own plain top-k, plus the ids of its best blocks per + request: ``[rows_b, min(topk_blocks, ceil(lc_b / block_size))]`` int32 each.""" from sglang.kernels.ops.attention.dsv4 import topk_transform_ragged_v2 selected = data.empty_selection(topk) @@ -647,31 +647,50 @@ def _publish_prefill_blocks( for rows, lc in zip(data.rows_per_request, data.lens_per_request) ] for tile, logits in score_tiles(data, kv, width_align=4): - lens = data.compress_lens[tile] - for request, rows, request_rows, lc in _requests_in_tile(data, tile): - scores = logits[rows, :lc] - scores.masked_fill_( - torch.arange(lc, device=device)[None, :] >= lens[rows, None], - -torch.inf, - ) - blocks = request_blocks[request][request_rows] - blocks.copy_( - select_candidate_block_ids( - logits=scores, - compress_lens=lens[rows, None], - topk_blocks=blocks.shape[1], - block_size=block_size, - ) - ) + _publish_tile_blocks( + data=data, + tile=tile, + logits=logits, + request_blocks=request_blocks, + block_size=block_size, + ) topk_transform_ragged_v2( logits, - lens, + 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, request_blocks +def _publish_tile_blocks( + *, + data: DeepGEMMPrefillData, + tile: slice, + logits: torch.Tensor, + request_blocks: List[torch.Tensor], + block_size: int, +) -> None: + lens = data.compress_lens[tile] + for request, rows, request_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, + ) + blocks = request_blocks[request][request_rows] + blocks.copy_( + select_candidate_block_ids( + logits=scores, + compress_lens=lens[rows, None], + topk_blocks=blocks.shape[1], + block_size=block_size, + ) + ) + + def _consume_prefill_blocks( *, data: DeepGEMMPrefillData, @@ -704,6 +723,8 @@ def _consume_prefill_blocks( selected[tile] = mask_topk_scores( scores=logits, indices=selected[tile], offsets=starts ) + # Free this tile's logits before the generator scores the next one. + del logits return selected diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py index 432bcff28fa6..084e3414b955 100644 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py +++ b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py @@ -266,6 +266,8 @@ def dense_prefill_topk( 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 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..fe2962b3aea0 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,10 @@ 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.low_ratio_indexer import ( + deep_gemm_backend, + deep_gemm_utils, +) from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -48,6 +51,42 @@ 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 = deep_gemm_utils.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 = deep_gemm_backend._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, deep_gemm_backend._DensePrefillBlocks(request_blocks=blocks) + if candidates is not None: + selected = deep_gemm_backend._consume_prefill_blocks( + data=data, + kv=inputs["kv"], + topk=inputs["topk"], + request_blocks=candidates.request_blocks, + block_size=inputs["candidate_block_size"], + ) + return selected, None + return deep_gemm_utils.dense_prefill_topk( + data, inputs["kv"], topk=inputs["topk"] + ), None + + def dense_scores(inputs): from deep_gemm import fp8_fp4_mqa_logits @@ -121,10 +160,10 @@ 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 + deep_gemm_utils, "_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 @@ -173,8 +212,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 +240,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,8 +252,8 @@ 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: @@ -232,8 +271,8 @@ def test_empty_queries_or_context(self): ): 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") @@ -258,11 +297,11 @@ def checked_logits(*args, **kwargs): return logits with ( - patch.object(dense_prefill_indexer, "_SCORE_BUDGET_BYTES", budget), + patch.object(deep_gemm_utils, "_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,8 +315,8 @@ 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( @@ -288,8 +327,8 @@ def test_score_memory_is_bounded(self): tuple(candidates.request_blocks[0].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 +336,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 96d146a3cf18..fea4c25f9525 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 @@ -8,12 +8,11 @@ import torch 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.low_ratio_indexer.block_math import ( select_candidate_blocks, ) -from sglang.srt.layers.attention.dsv4.dense_prefill_indexer import ( - DenseCandidateIndexer, +from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_utils import ( + DeepGEMMPrefillData, ) from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -37,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: @@ -77,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], ) @@ -122,20 +121,66 @@ def reference_blocks(case: Case) -> torch.Tensor: return keep.unflatten(1, (-1, BLOCK)).any(-1) -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 publish_sparse(case: Case): + """(the DeepGEMM sparse table, the source layer's own top-k).""" + from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_backend import ( + publish_prefill_table, + ) + + 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 select_sparse(table, data: DeepGEMMPrefillData, k_cache: torch.Tensor): + from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_backend 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.low_ratio_indexer.deep_gemm_backend 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.low_ratio_indexer.deep_gemm_backend import ( + _consume_prefill_blocks, + ) + + return _consume_prefill_blocks( + data=case.data, + kv=case.kv, + topk=TOPK, + request_blocks=blocks, + block_size=BLOCK, + ) + + def picks(positions: torch.Tensor, row: int) -> set: return set(positions[row].tolist()) - {-1} @@ -145,14 +190,6 @@ 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 @@ -161,7 +198,7 @@ def test_publish_prefill_is_the_torch_block_selection(self): 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 @@ -169,7 +206,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) @@ -181,12 +218,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): @@ -211,15 +244,15 @@ 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) diff --git a/test/registered/unit/layers/test_dsv41_candidate_blocks.py b/test/registered/unit/layers/test_dsv41_candidate_blocks.py index 79915cf5c29c..64b27dae92ca 100644 --- a/test/registered/unit/layers/test_dsv41_candidate_blocks.py +++ b/test/registered/unit/layers/test_dsv41_candidate_blocks.py @@ -2,12 +2,14 @@ import torch -from sglang.srt.layers.attention.dsv4.candidate_indexer import ( - PrefillCandidateBlocks, +from sglang.srt.layers.attention.dsv4.low_ratio_indexer.block_math import ( candidate_block_mask, select_candidate_block_ids, select_candidate_blocks, ) +from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_backend import ( + _DensePrefillBlocks, +) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -63,7 +65,7 @@ def test_underfilled_and_empty_candidates(self): 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) + candidates = _DensePrefillBlocks(request_blocks=requests) tail = candidates.tail([2, 0, 0]) self.assertEqual( [tuple(b.shape) for b in tail.request_blocks], [(2, 2), (0, 2), (0, 2)] diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py index 3dbb65fab2b4..1865705e870c 100644 --- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py +++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py @@ -908,47 +908,93 @@ 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.low_ratio_indexer import DecodeInputs + from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_utils 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.low_ratio_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.low_ratio_indexer import ( + deep_gemm_backend as mod, + ) + + inputs, data = self._inputs() + backend = object.__new__(mod.DeepGEMMCandidateBackend) + 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( @@ -973,46 +1019,73 @@ def test_each_deep_gemm_call_receives_one_tensor_schedule(self): ), patch.object( mod, - "build_sparse_indexer_schedule", + "build_sparse_schedule", return_value=torch.tensor([7], dtype=torch.uint8), ), 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.low_ratio_indexer import ( + dense_indexer as mod, + ) + + inputs, data = self._inputs() + indexer = object.__new__(mod.DenseIndexer) + 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 low_ratio_indexer + from sglang.srt.layers.attention.dsv4.low_ratio_indexer.torch_backend import ( + TorchCandidateBackend, + ) - def platform(sm): - return patch.object( - candidate_indexer, "get_platform", lambda: SimpleNamespace(device_sm=sm) + def sm100(value): + return patch.object(low_ratio_indexer, "is_sm100_or_newer", lambda: value) + + def make(topk_blocks, block_size): + return low_ratio_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 the portable backend. + 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): + indexer = make(topk_blocks, block_size) + self.assertIsInstance(indexer.prefill, TorchCandidateBackend) + self.assertIsInstance(indexer.decode, TorchCandidateBackend) # 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__": From 22a741283f6ae2b1842ca167b6fea9f28f7eacc8 Mon Sep 17 00:00:00 2001 From: DarkSharpness <2040703891@qq.com> Date: Fri, 25 Sep 2026 06:06:31 -0700 Subject: [PATCH 03/16] [DSv4.1] Restructure the ratio-1/2 indexer into dsv4/v41_indexer Address review: rename the package to v41_indexer and split it by candidate scheme instead of by DeepGEMM / torch. - sparse_table: the DeepGEMM sparse-table candidates (SM100), unchanged. - dense_blocks: dense scores plus the chosen block ids, for torch and for the DeepGEMM prefill under CP. A publish carries one [rows, n] block-id tensor instead of per-request masks; a consumer gathers the scores of its blocks and takes the top-k among them (topk_among_blocks), which picks the same set as masking everything else. - plain: the non-candidate layers (was DenseIndexer). - scoring: the DeepGEMM and torch score helpers, one module. - make_candidate_indexer returns (prefill, decode) backends directly; the forwarding CandidateIndexer and caps.py are gone. - The torch block math and topk_among_blocks / get_tail_row_indices move to kernels/ops/attention/dsv4/candidate_blocks.py. - The late-layer tail builds its row indices without a host sync (pinned staging), which also removes the sync the sparse-table tail had. Tests are ported to the new layout, with cases for topk_among_blocks, the block-id tail and the sync-free tail. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../ops/attention/dsv4/candidate_blocks.py | 112 +++++- .../layers/attention/deepseek_v4_backend.py | 44 +-- .../dsv4/low_ratio_indexer/__init__.py | 99 ----- .../dsv4/low_ratio_indexer/block_math.py | 92 ----- .../low_ratio_indexer/candidate_indexer.py | 87 ----- .../attention/dsv4/low_ratio_indexer/caps.py | 39 -- .../dsv4/low_ratio_indexer/torch_backend.py | 152 -------- .../dsv4/low_ratio_indexer/torch_utils.py | 178 --------- .../attention/dsv4/v41_indexer/__init__.py | 142 +++++++ .../dsv4/v41_indexer/dense_blocks.py | 354 ++++++++++++++++++ .../dense_indexer.py => v41_indexer/plain.py} | 23 +- .../scoring.py} | 184 ++++++++- .../sparse_table.py} | 249 +----------- .../inputs.py => v41_indexer/types.py} | 40 +- .../attention/test_dense_prefill_indexer.py | 61 +-- .../test_dsv41_prefill_sparse_indexer.py | 49 ++- .../layers/test_dsv41_candidate_blocks.py | 136 ++++--- .../unit/layers/test_dsv4_nonpaged_indexer.py | 50 +-- 18 files changed, 1047 insertions(+), 1044 deletions(-) delete mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/__init__.py delete mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/block_math.py delete mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/candidate_indexer.py delete mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/caps.py delete mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_backend.py delete mode 100644 python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_utils.py create mode 100644 python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py create mode 100644 python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py rename python/sglang/srt/layers/attention/dsv4/{low_ratio_indexer/dense_indexer.py => v41_indexer/plain.py} (96%) rename python/sglang/srt/layers/attention/dsv4/{low_ratio_indexer/deep_gemm_utils.py => v41_indexer/scoring.py} (59%) rename python/sglang/srt/layers/attention/dsv4/{low_ratio_indexer/deep_gemm_backend.py => v41_indexer/sparse_table.py} (69%) rename python/sglang/srt/layers/attention/dsv4/{low_ratio_indexer/inputs.py => v41_indexer/types.py} (73%) diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py index b70d0f2d2c41..42fe76416fcf 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py @@ -1,6 +1,14 @@ -"""Per-row candidate block counts and sparse-row lengths for the paged indexer.""" +"""Candidate blocks of the two-level indexer: a candidate source keeps whole +blocks of ``block_size`` compressed positions, its newest block always, and a +consumer selects among the kept blocks only. + +The per-row block counts and sparse-row lengths, the block selection (the JIT +chain on SM100, torch elsewhere), and the torch top-k among chosen blocks.""" + +from typing import List, Optional, Union import torch +import torch.nn.functional as F import triton import triton.language as tl @@ -62,3 +70,105 @@ def candidate_row_lens( **pdl_kwargs, ) return nblocks, valid + + +def amax_topk_blocks( + logits: torch.Tensor, + seq_lens: torch.Tensor, + nblocks: torch.Tensor, + topk_blocks: int, + max_seq_len: Optional[int] = None, +) -> torch.Tensor: + """Per row the ``topk_blocks`` blocks of 8 positions with the largest block + maximum among its first ``seq_lens[b]`` positions, the newest block always + included: block ids in no particular order, ``-1`` past the row's count. + ``nblocks`` is ``ceil(seq_lens / 8)`` as int32.""" + from .candidate_table import amax8_varlen + from .topk import plan_topk_v2, topk_transform_paged_v2 + + rows = logits.shape[0] + block = 8 + if max_seq_len is None: + max_seq_len = logits.shape[1] + # NOTE: plan cannot be the previous kernel of topk_transform_paged_v2 + plan = plan_topk_v2(nblocks) + # block maxima, the newest block +inf; the top-k reads each row up to nblocks + # only, so nothing past a row's keys is initialised (v2 needs stride % 4 == 0) + keys = logits.new_empty(rows, -(-max_seq_len // (4 * block)) * 4) + amax8_varlen(logits, seq_lens, out=keys) + blocks = torch.empty(rows, topk_blocks, dtype=torch.int32, device=logits.device) + topk_transform_paged_v2(keys, nblocks, None, blocks, 1, plan) + return blocks + + +def select_candidate_block_ids( + logits: torch.Tensor, + compress_lens: Union[torch.Tensor, int], + topk_blocks: int, + block_size: int, +) -> torch.Tensor: + """The torch block selection: per row the ids of the ``topk_blocks`` blocks + with the largest maximum score, the block holding position + ``compress_lens - 1`` always kept; int32 ``[rows, min(topk_blocks, blocks)]``, + unordered, ``-1`` where a row has fewer finite blocks. ``logits`` must already + 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 + ) + top = scores.topk(min(topk_blocks, num_blocks), dim=-1) + return top.indices.to(torch.int32).masked_fill_(~(top.values > -torch.inf), -1) + + +def topk_among_blocks( + scores: torch.Tensor, + lens: torch.Tensor, + blocks: torch.Tensor, + k: int, + block_size: int, +) -> torch.Tensor: + """The torch top-``k`` of each row among its chosen blocks: ``scores`` + ``[rows, width]`` indexed by position, ``lens`` ``[rows]`` the causal length, + ``blocks`` ``[rows, n]`` block ids with ``-1`` padding. Returns int64 + ``[rows, k]`` positions, unordered, ``-1`` where a row has fewer than ``k`` + finite candidates; positions at or past ``lens`` never count.""" + rows, width = scores.shape + if width == 0: + return torch.full((rows, k), -1, dtype=torch.int64, device=scores.device) + offsets = torch.arange(block_size, device=scores.device) + columns = (blocks.to(torch.int64)[:, :, None] * block_size + offsets).flatten(1) + valid = ( + (blocks >= 0).repeat_interleave(block_size, dim=1) + & (columns < lens.to(torch.int64)[:, None]) + & (columns < width) + ) + candidates = scores.gather(1, columns.clamp(0, width - 1)) + candidates = candidates.masked_fill(~valid, -torch.inf) + top = candidates.topk(min(k, candidates.shape[1]), dim=-1, sorted=False) + picked = columns.gather(1, top.indices).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 + + +def get_tail_row_indices( + full_rows_per_request: List[int], + tail_rows_per_request: List[int], + device: torch.device, +) -> torch.Tensor: + """int64 indices of each request's last ``tail_rows_per_request[b]`` rows, in + row order, copied without a host sync (pinned staging, non-blocking).""" + 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) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index c4016043f0de..fe7a5c9f5d4a 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -65,17 +65,6 @@ token_req_indices, ) from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin -from sglang.srt.layers.attention.dsv4.low_ratio_indexer import ( - CandidateMetadata, - CapturedPrefillInputs, - DecodeInputs, - DenseIndexer, - PrefillInputs, - Selection, - has_dense_fp4_indexer, - is_sm100_or_newer, - make_candidate_indexer, -) from sglang.srt.layers.attention.dsv4.metadata import ( _LARGE_INDEXER_QUERY_THRESHOLD, PagedIndexerMetadata, @@ -88,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_plain_indexer, +) from sglang.srt.layers.attention.verify_mask import ( VerifyMask, maybe_create_verify_mask, @@ -1126,14 +1126,14 @@ def __init__( self.has_c128: bool = 128 in self.present_ratios 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( + 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.dense_indexer = DenseIndexer( + self.plain_indexer = make_plain_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] @@ -3089,8 +3089,8 @@ def _low_ratio_index_topk( *, rows_per_request=None, # NOTE: only used in CP ) -> None: - """Select this index layer's compressed positions: a dense layer on the - dense indexer, a candidate source or consumer on the candidate indexer.""" + """Select this index layer's compressed positions: a plain layer on the + plain indexer, a candidate source or consumer on the candidate backends.""" is_source = layer.indexer.is_candidate_source is_consumer = layer.indexer.uses_candidates ratio = layer.compress_ratio @@ -3112,23 +3112,23 @@ def _low_ratio_index_topk( layer, x, q_lora, req, pos, mode ) if is_source: - published = self.candidate_indexer.publish_decode(inputs, out) + published = self.decode_candidates.publish_decode(inputs, out) self._publish_candidate_metadata(published) elif is_consumer: - self.candidate_indexer.consume_decode(inputs, published, out) + self.decode_candidates.consume_decode(inputs, published, out) else: - self.dense_indexer.topk_decode(inputs, out) + self.plain_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.candidate_indexer.publish_prefill(inputs, out) + published = self.prefill_candidates.publish_prefill(inputs, out) self._publish_candidate_metadata(published) elif is_consumer: - self.candidate_indexer.consume_prefill(inputs, published, out) + self.prefill_candidates.consume_prefill(inputs, published, out) else: - self.dense_indexer.topk_prefill(inputs, out) + self.plain_indexer.topk_prefill(inputs, out) def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> None: ratio = layer.compress_ratio @@ -3154,7 +3154,7 @@ def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> Non paged_metadata=metadata, ) out = self._get_low_ratio_selection(ratio) - self.dense_indexer.topk_prefill_captured(inputs, out) + self.plain_indexer.topk_prefill_captured(inputs, out) def _get_low_ratio_selection(self, compress_ratio: int) -> Selection: core = self.forward_metadata.core_metadata diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/__init__.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/__init__.py deleted file mode 100644 index 954e5f9a89ce..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/__init__.py +++ /dev/null @@ -1,99 +0,0 @@ -"""The top-k selection of the DeepSeek V4.1 ratio-1/2 index layers. - -The dense layers go through ``DenseIndexer``; a candidate source and its -consumers go through the ``CandidateIndexer`` that ``make_candidate_indexer`` -builds. The two share only the scoring utilities. Nothing else in this package -is meant to be named from outside it. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch - -from sglang.srt.runtime_context import get_parallel - -from .candidate_indexer import CandidateIndexer -from .caps import ( - has_dense_fp4_indexer, - is_sm100_or_newer, - use_deep_gemm_prefill, -) -from .dense_indexer import DenseIndexer -from .inputs import ( - CandidateMetadata, - CapturedPrefillInputs, - DecodeInputs, - PrefillInputs, - Selection, -) - -if TYPE_CHECKING: - from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool - -__all__ = [ - "CandidateIndexer", - "CandidateMetadata", - "CapturedPrefillInputs", - "DecodeInputs", - "DenseIndexer", - "PrefillInputs", - "Selection", - "has_dense_fp4_indexer", - "is_sm100_or_newer", - "make_candidate_indexer", -] - - -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, -) -> CandidateIndexer: - from .torch_backend import TorchCandidateBackend - - torch_backend = TorchCandidateBackend( - token_to_kv_pool=token_to_kv_pool, req_to_token=req_to_token - ) - # Without candidate blocks no layer is a candidate source or consumer. - if not is_sm100_or_newer() or candidate_topk_blocks <= 0: - return CandidateIndexer(prefill=torch_backend, decode=torch_backend) - - 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 .deep_gemm_backend import ( - DeepGEMMCandidateBackend, - DeepGEMMCPCandidateBackend, - ) - - deep_gemm_backend = DeepGEMMCandidateBackend( - 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, - ) - - if not use_deep_gemm_prefill(): - prefill_backend = torch_backend - elif get_parallel().attn_cp_size > 1: - prefill_backend = DeepGEMMCPCandidateBackend( - 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, - ) - else: - prefill_backend = deep_gemm_backend - return CandidateIndexer(prefill=prefill_backend, decode=deep_gemm_backend) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/block_math.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/block_math.py deleted file mode 100644 index d39ec84b986d..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/block_math.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Block-selection math shared by the backends whose candidates are picked from -dense scores: a candidate source keeps whole blocks of ``block_size`` compressed -positions rather than single positions, and its newest block always.""" - -from __future__ import annotations - -from typing import Optional, Union - -import torch -import torch.nn.functional as F - - -def _candidate_block_topk( - logits: torch.Tensor, - compress_lens: Union[torch.Tensor, int], - topk_blocks: int, - block_size: int, -) -> torch.return_types.topk: - 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 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 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, - ) - 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, - block_size: int, -) -> torch.Tensor: - top = _candidate_block_topk( - logits=logits, - compress_lens=compress_lens, - topk_blocks=topk_blocks, - block_size=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] diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/candidate_indexer.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/candidate_indexer.py deleted file mode 100644 index c967bf6301ef..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/candidate_indexer.py +++ /dev/null @@ -1,87 +0,0 @@ -"""The two-level candidate scheme of the DeepSeek V4.1 ratio-1/2 index layers. - -The candidate source publishes the blocks of compressed positions it kept, and -the index layers after it select their top-k among those blocks only. What a -publish carries -- a DeepGEMM sparse table, block ids, a mask -- is the backend's -own; the attention backend keeps it alive on the forward metadata and hands it -back to the consumers unread. -""" - -from __future__ import annotations - -from typing import Optional, Protocol, TypeVar - -from .inputs import ( - CandidateMetadata, - DecodeInputs, - PrefillInputs, - Selection, -) - -T = TypeVar("T", bound=CandidateMetadata) - - -class PrefillBackend(Protocol[T]): - def publish_prefill( - self, - inputs: PrefillInputs, - out: Selection, - ) -> Optional[T]: ... - - def consume_prefill( - self, - inputs: PrefillInputs, - published: Optional[T], - out: Selection, - ) -> None: ... - - -class DecodeBackend(Protocol[T]): - def publish_decode( - self, - inputs: DecodeInputs, - out: Selection, - ) -> Optional[T]: ... - - def consume_decode( - self, - inputs: DecodeInputs, - published: Optional[T], - out: Selection, - ) -> None: ... - - -class CandidateIndexer: - def __init__(self, *, prefill: PrefillBackend, decode: DecodeBackend): - self.prefill = prefill - self.decode = decode - - def publish_prefill( - self, - inputs: PrefillInputs, - out: Selection, - ) -> Optional[CandidateMetadata]: - return self.prefill.publish_prefill(inputs, out) - - def consume_prefill( - self, - inputs: PrefillInputs, - published: Optional[CandidateMetadata], - out: Selection, - ) -> None: - self.prefill.consume_prefill(inputs, published, out) - - def publish_decode( - self, - inputs: DecodeInputs, - out: Selection, - ) -> Optional[CandidateMetadata]: - return self.decode.publish_decode(inputs, out) - - def consume_decode( - self, - inputs: DecodeInputs, - published: Optional[CandidateMetadata], - out: Selection, - ) -> None: - self.decode.consume_decode(inputs, published, out) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/caps.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/caps.py deleted file mode 100644 index 92ba855c78b3..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/caps.py +++ /dev/null @@ -1,39 +0,0 @@ -"""Which low-ratio indexer kernels this build can run.""" - -from __future__ import annotations - -import functools - -import torch - - -@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: - """Whether an eager prefill scores on DeepGEMM rather than torch. 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 - - if not (is_sm100_or_newer() and has_dense_fp4_indexer()): - return False - return ( - get_parallel().attn_cp_size > 1 - or not envs.SGLANG_DSV41_TORCH_PREFILL_INDEXER.get() - ) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_backend.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_backend.py deleted file mode 100644 index f2b8f0237471..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_backend.py +++ /dev/null @@ -1,152 +0,0 @@ -"""The candidate scheme on the portable torch path: bf16 scores materialized per -request and selected with ``torch.topk``, the candidates carried by boolean masks. -A source keeps a mask of its best blocks, a consumer scores densely and drops -everything outside it; it needs no page table, so it works wherever torch does. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, List, Optional - -import msgspec -import torch - -from .block_math import ( - mask_topk_scores, - select_candidate_blocks, -) -from .inputs import ( - CandidateMetadata, - DecodeInputs, - PrefillInputs, - Selection, -) -from .torch_utils import ( - decode_scores, - prefill_requests, - write_decode, - write_prefill, -) - -if TYPE_CHECKING: - from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool - - -class CandidateMasks(CandidateMetadata, msgspec.Struct): - mask: Optional[torch.Tensor] = None # decode: [rows, width] bool - request_masks: Optional[List[torch.Tensor]] = None # prefill: [rows_b, lc_b] each - - def tail(self, rows_per_request: List[int]) -> CandidateMasks: - assert self.request_masks is not None, "prefill masks missing" - return CandidateMasks( - request_masks=[ - mask[mask.shape[0] - t :] - for mask, t in zip(self.request_masks, rows_per_request) - ] - ) - - -class TorchCandidateBackend: - def __init__( - self, - *, - token_to_kv_pool: DeepSeekV4TokenToKVPool, - req_to_token: torch.Tensor, - ): - self.token_to_kv_pool = token_to_kv_pool - self.req_to_token = req_to_token - - def publish_prefill(self, inputs: PrefillInputs, out: Selection): - indexer = inputs.indexer - published = [] - for request, chunks in self._prefill_requests(inputs, out): - if request.lc == 0: - # Consumers address masks by request position, including empty requests. - published.append( - torch.zeros(0, 0, dtype=torch.bool, device=inputs.positions.device) - ) - continue - masks = [] - for chunk in chunks: - masks.append( - select_candidate_blocks( - chunk.scores, - chunk.lens[:, None], - topk_blocks=indexer.candidate_topk_blocks, - block_size=indexer.candidate_block_size, - ) - ) - idx = chunk.scores.topk(request.k, dim=-1, sorted=False).indices - write_prefill(out, request, chunk, idx) - published.append(torch.cat(masks) if len(masks) > 1 else masks[0]) - return CandidateMasks(request_masks=published) - - def consume_prefill( - self, - inputs: PrefillInputs, - published: Optional[CandidateMasks], - out: Selection, - ) -> None: - assert published is not None - request_masks = published.request_masks - for request, chunks in self._prefill_requests(inputs, out): - for chunk in chunks: - keep = request_masks[request.index][chunk.rows] - s = chunk.scores.masked_fill(~keep, -torch.inf) - idx = s.topk(request.k, dim=-1, sorted=False).indices - idx = mask_topk_scores(s, idx) - idx = idx.masked_fill(idx < 0, request.lc) - write_prefill(out, request, chunk, idx) - - def publish_decode(self, inputs: DecodeInputs, out: Selection): - indexer = inputs.indexer - d = self._decode_scores(inputs, out) - if d is None: - return None - published = CandidateMasks( - mask=select_candidate_blocks( - d.scores, - d.lens[:, None], - topk_blocks=indexer.candidate_topk_blocks, - block_size=indexer.candidate_block_size, - ) - ) - k = min(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[CandidateMasks], - out: Selection, - ) -> None: - d = self._decode_scores(inputs, out) - if d is None: - return - assert published is not None - consume = published.mask - assert torch.is_tensor(consume) and consume.shape[0] == d.bs - s = d.scores.masked_fill(~consume[:, : d.lmax], -torch.inf) - k = min(inputs.indexer.index_topk, d.lmax) - idx = s.topk(k, dim=-1, sorted=False).indices - idx = mask_topk_scores(s, idx) - idx = idx.masked_fill(idx < 0, d.lmax) - write_decode(out, d, idx) - - def _prefill_requests(self, inputs: PrefillInputs, out: Selection): - return prefill_requests( - inputs=inputs, - out=out, - token_to_kv_pool=self.token_to_kv_pool, - req_to_token=self.req_to_token, - ) - - def _decode_scores(self, inputs: DecodeInputs, out: Selection): - return decode_scores( - inputs=inputs, - out=out, - token_to_kv_pool=self.token_to_kv_pool, - req_to_token=self.req_to_token, - ) diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_utils.py b/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_utils.py deleted file mode 100644 index 22b1ae8744bb..000000000000 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/torch_utils.py +++ /dev/null @@ -1,178 +0,0 @@ -"""Torch scoring shared by the dense layers and the candidate scheme: bf16 scores -per request on prefill, the fp4 decode logits kernel on decode, and the writes of -a ``torch.topk`` back into the pool-slot layout attention reads.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Iterator, Optional - -import msgspec -import torch - -from sglang.kernels.ops.attention.dsv4.fp4_indexer import fp4_index_logits_decode - -from .inputs 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 - - -# Arbitrary cap on one bf16 [rows, heads, lc] score chunk; transients run ~3x this. -_SCORE_BUDGET_BYTES = 1 << 30 - - -class RequestScores(msgspec.Struct, frozen=True): - """One request of a prefill chunk, scored row chunk by row chunk.""" - - index: int # position among the chunk's requests, empty ones included - lc: int # compressed positions visible at its newest row; 0 scores nothing - k: int - columns: Optional[torch.Tensor] = None # arange(lc) - tok: Optional[torch.Tensor] = None # [rows_b] its query rows in the chunk - lens: Optional[torch.Tensor] = None # [rows_b] compressed positions each row sees - slots: Optional[torch.Tensor] = None # [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]]]: - """Each request of the chunk with its row chunks' scores; an empty request - comes with no chunk.""" - 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 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: - yield RequestScores(index=b, lc=0, k=0), iter(()) - 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( - index=b, 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 - # Every step 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, _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/__init__.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py new file mode 100644 index 000000000000..b2bd85f9f42b --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py @@ -0,0 +1,142 @@ +"""The top-k selection of the DeepSeek V4.1 ratio-1/2 index layers. + +The plain index layers go through the ``PlainIndexer`` of ``make_plain_indexer``; a +candidate source and its consumers go through the prefill / decode backends of +``make_candidate_indexer``: a DeepGEMM sparse table (``sparse_table``) or dense +block ids (``dense_blocks``). All three share only ``scoring``. Nothing else in +this package is meant to be named from outside it. +""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Tuple + +import torch + +from .plain import PlainIndexer +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", + "PlainIndexer", + "PrefillCandidates", + "PrefillInputs", + "Selection", + "has_dense_fp4_indexer", + "is_sm100_or_newer", + "make_candidate_indexer", + "make_plain_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: + """Whether an eager prefill scores on DeepGEMM rather than torch. 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_plain_indexer( + *, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + req_to_token: torch.Tensor, +) -> PlainIndexer: + return PlainIndexer( + 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]: + """(prefill, decode): the sparse table where SM100 can page the rows, dense + block ids elsewhere.""" + 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..e68e2edeb7f7 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py @@ -0,0 +1,354 @@ +"""The candidate scheme carried by dense block ids. + +A source scores its whole context, keeps its own top-k and publishes the ids of +its best blocks; a consumer scores its whole context too, gathers the scores of +the published blocks only and takes its top-k among them. Scores come from +DeepGEMM's flattened-K logits or from torch, per backend configuration; the +block math is the torch one in ``candidate_blocks``. +""" + +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 ( + get_tail_row_indices, + select_candidate_block_ids, + topk_among_blocks, +) + +from .scoring import ( + DeepGEMMPrefillData, + decode_scores, + get_deep_gemm_prefill_data, + get_flat_index_k, + prefill_requests, + score_tiles, + write_decode, + write_prefill, +) +from .types import ( + CandidateMetadata, + DecodeInputs, + PrefillInputs, + Selection, +) + +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 = self._decode_scores(inputs, out) + 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 = self._decode_scores(inputs, out) + 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._get_flat_index_k(inputs, data), + 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._get_flat_index_k(inputs, data), + topk=inputs.indexer.index_topk, + blocks=published.blocks, + block_size=self.block_size, + ) + data.write_selection(selected=selected, out=out) + + def _get_flat_index_k(self, inputs: PrefillInputs, data: DeepGEMMPrefillData): + return get_flat_index_k( + data=data, token_to_kv_pool=self.token_to_kv_pool, layer_id=inputs.layer_id + ) + + # ---------- 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 self._prefill_requests(inputs, out): + 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 self._prefill_requests(inputs, out): + 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 _prefill_requests(self, inputs: PrefillInputs, out: Selection): + return prefill_requests( + inputs=inputs, + out=out, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + ) + + def _decode_scores(self, inputs: DecodeInputs, out: Selection): + return decode_scores( + inputs=inputs, + out=out, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + ) + + +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]: + """The source layer's own plain top-k, plus the ids of each row's best blocks: + ``[rows, min(topk_blocks, ceil(max lc / block_size))]`` int32, -1 padded.""" + 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: + """A consumer's top-k among the published blocks of its dense scores, as + flattened-K columns, ``-1`` padded, unordered.""" + selected = data.empty_selection(topk) + for tile, logits in score_tiles(data, kv, width_align=4): + _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. + positions = topk_among_blocks( + logits[rows, :lc], + 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): + """``(rows, lc)`` of each request with a visible position and rows in + ``tile``: its rows within the tile and its context length.""" + 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/low_ratio_indexer/dense_indexer.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/plain.py similarity index 96% rename from python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/dense_indexer.py rename to python/sglang/srt/layers/attention/dsv4/v41_indexer/plain.py index d6f7d62f8780..82d87228563c 100644 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/dense_indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/plain.py @@ -16,43 +16,42 @@ topk_transform_paged_from_metadata, ) -from .caps import is_sm100_or_newer, use_deep_gemm_prefill -from .deep_gemm_utils import ( +from .scoring import ( + decode_scores, dense_prefill_topk, get_deep_gemm_decode_data, get_deep_gemm_prefill_data, get_flat_index_k, get_index_k_cache, + prefill_requests, quantize_index_q, + write_decode, + write_prefill, ) -from .inputs import ( +from .types import ( CapturedPrefillInputs, DecodeInputs, PrefillInputs, Selection, ) -from .torch_utils import ( - decode_scores, - prefill_requests, - write_decode, - write_prefill, -) if TYPE_CHECKING: from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool -class DenseIndexer: +class PlainIndexer: 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 = is_sm100_or_newer() + 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: diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py similarity index 59% rename from python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py rename to python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py index 084e3414b955..ce1e4d4d601c 100644 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_utils.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py @@ -1,36 +1,41 @@ -"""DeepGEMM scoring shared by the dense layers and the candidate scheme. +"""Index scores shared by the dense layers and both candidate schemes. -On prefill: the chunk's index-K slots, fp4 queries and head weights, the tiled -``fp8_fp4_mqa_logits`` scores over the flattened index K, the plain top-k, and -the write of a selection back into the pool-slot layout attention reads. On -decode: the fp4 queries and the paged pool. +DeepGEMM (SM100): on prefill the chunk's index-K slots, fp4 queries and head +weights, the tiled ``fp8_fp4_mqa_logits`` scores over the flattened index K, the +plain top-k, and the write of a selection back into the pool-slot layout +attention reads; on decode the fp4 queries and the paged pool. + +Torch (elsewhere): bf16 scores per request on prefill, the fp4 decode logits +kernel on decode, and the writes of a ``torch.topk`` back into pool slots. """ from __future__ import annotations -from typing import TYPE_CHECKING, Generator, List, Optional, Tuple +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.srt.layers.attention.mqa_logits_utils import ( mqa_logits_row_bytes, mqa_logits_rows_per_chunk, ) from sglang.srt.utils.common import ceil_align -from .inputs import ( +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 masks and block-selection scratch. -_SCORE_BUDGET_BYTES = 2 << 30 +# leaves room for candidate block ids and block-selection scratch. +_DEEP_GEMM_SCORE_BUDGET_BYTES = 2 << 30 class DeepGEMMDecodeData(msgspec.Struct, frozen=True): @@ -279,8 +284,167 @@ def _rows_per_chunk(rows: int, width: int, *, heads: int) -> int: rows_per_chunk = mqa_logits_rows_per_chunk( num_rows=ceil_align(rows, row_alignment), row_bytes=mqa_logits_row_bytes(width), - budget_bytes=_SCORE_BUDGET_BYTES, + budget_bytes=_DEEP_GEMM_SCORE_BUDGET_BYTES, ) if rows_per_chunk is None: return rows return max(row_alignment, rows_per_chunk // row_alignment * row_alignment) + + +# ---------- 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): + """One request of a prefill chunk, scored row chunk by row chunk.""" + + index: int # position among the chunk's requests, empty ones included + lc: int # compressed positions visible at its newest row; 0 scores nothing + k: int + columns: Optional[torch.Tensor] = None # arange(lc) + tok: Optional[torch.Tensor] = None # [rows_b] its query rows in the chunk + lens: Optional[torch.Tensor] = None # [rows_b] compressed positions each row sees + slots: Optional[torch.Tensor] = None # [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]]]: + """Each request of the chunk with its row chunks' scores; an empty request + comes with no chunk.""" + 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 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: + yield RequestScores(index=b, lc=0, k=0), iter(()) + 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( + index=b, 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 + # Every step 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_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/low_ratio_indexer/deep_gemm_backend.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py similarity index 69% rename from python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_backend.py rename to python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py index 8de7a8b14dee..0273ca1ba11f 100644 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/deep_gemm_backend.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py @@ -1,4 +1,4 @@ -"""The candidate scheme on SM100, carried by a DeepGEMM sparse table. +"""The candidate scheme carried by a DeepGEMM sparse table (SM100). A source publishes the block ids it kept plus the pool slots and the schedule ``fp8_fp4_paged_sparse_mqa_logits`` needs for them, and a consumer scores only @@ -12,13 +12,16 @@ import msgspec import torch -from sglang.kernels.ops.attention.dsv4.candidate_blocks import candidate_row_lens +from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( + amax_topk_blocks, + candidate_row_lens, + get_tail_row_indices, +) from sglang.kernels.ops.attention.dsv4.candidate_table import ( amax8_varlen, sort_candidate_blocks, ) from sglang.kernels.ops.attention.dsv4.topk import ( - plan_topk_v2, topk_transform_bf16_small, topk_transform_paged_v2, topk_transform_ragged_v2, @@ -29,12 +32,7 @@ ) from sglang.srt.layers.attention.dsv4.metadata import expand_index_page_table -from .block_math import ( - candidate_block_mask, - mask_topk_scores, - select_candidate_block_ids, -) -from .deep_gemm_utils import ( +from .scoring import ( DeepGEMMDecodeData, DeepGEMMPrefillData, get_deep_gemm_decode_data, @@ -43,7 +41,7 @@ get_index_k_cache, score_tiles, ) -from .inputs import ( +from .types import ( CandidateMetadata, DecodeInputs, PrefillInputs, @@ -87,11 +85,11 @@ class _SparsePrefillTable(_SparseTable): topk_blocks: int def tail(self, rows_per_request: List[int]) -> _SparsePrefillTable: - rows, start = [], 0 - for n, t in zip(self.rows_per_request, rows_per_request): - rows.append(torch.arange(start + n - t, start + n)) - start += n - rows = torch.cat(rows).to(self.blocks.device) + 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.topk_blocks) return _build_prefill_table( @@ -107,19 +105,7 @@ def tail(self, rows_per_request: List[int]) -> _SparsePrefillTable: ) -class _DensePrefillBlocks(CandidateMetadata, msgspec.Struct): - request_blocks: List[torch.Tensor] - - def tail(self, rows_per_request: List[int]) -> _DensePrefillBlocks: - return _DensePrefillBlocks( - request_blocks=[ - blocks[blocks.shape[0] - length :] - for blocks, length in zip(self.request_blocks, rows_per_request) - ] - ) - - -class DeepGEMMCandidateBackend: +class SparseTableBackend: def __init__( self, *, @@ -369,90 +355,6 @@ def _get_request_ids(self, inputs: DecodeInputs, rows: int, device: torch.device return buf[:rows] -class DeepGEMMCPCandidateBackend: - def __init__( - self, - *, - token_to_kv_pool: DeepSeekV4TokenToKVPool, - req_to_token: torch.Tensor, - candidate_topk_blocks: int, - candidate_block_size: int, - ) -> 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 - - def publish_prefill(self, inputs: PrefillInputs, out: Selection): - out.reset() - data = self._get_prefill_data(inputs) - if data is None: - return None - selected, request_blocks = _publish_prefill_blocks( - data=data, - kv=self._get_flat_index_k(inputs, data), - topk=inputs.indexer.index_topk, - topk_blocks=self.topk_blocks, - block_size=self.block_size, - ) - data.write_selection(selected=selected, out=out) - return _DensePrefillBlocks(request_blocks) - - def consume_prefill( - self, - inputs: PrefillInputs, - published: Optional[_DensePrefillBlocks], - out: Selection, - ) -> None: - out.reset() - data = self._get_prefill_data(inputs) - if data is None: - return - assert isinstance(published, _DensePrefillBlocks) - selected = _consume_prefill_blocks( - data=data, - kv=self._get_flat_index_k(inputs, data), - topk=inputs.indexer.index_topk, - request_blocks=published.request_blocks, - block_size=self.block_size, - ) - data.write_selection(selected=selected, out=out) - - def _get_prefill_data(self, inputs: PrefillInputs) -> Optional[DeepGEMMPrefillData]: - return get_deep_gemm_prefill_data(inputs, self.req_to_token) - - def _get_flat_index_k(self, inputs: PrefillInputs, data: DeepGEMMPrefillData): - return get_flat_index_k( - data=data, token_to_kv_pool=self.token_to_kv_pool, layer_id=inputs.layer_id - ) - - -def amax_topk_blocks( - logits: torch.Tensor, - seq_lens: torch.Tensor, - nblocks: torch.Tensor, - topk_blocks: int, - max_seq_len: Optional[int] = None, -) -> torch.Tensor: - """Per row the ``topk_blocks`` blocks of 8 positions with the largest block - maximum among its first ``seq_lens[b]`` positions, the newest block always - included: block ids in no particular order, ``-1`` past the row's count. - ``nblocks`` is ``ceil(seq_lens / 8)`` as int32.""" - rows = logits.shape[0] - block = CANDIDATE_BLOCK_SIZE - if max_seq_len is None: - max_seq_len = logits.shape[1] - # NOTE: plan cannot be the previous kernel of topk_transform_paged_v2 - plan = plan_topk_v2(nblocks) - # block maxima, the newest block +inf; the top-k reads each row up to nblocks - # only, so nothing past a row's keys is initialised (v2 needs stride % 4 == 0) - keys = logits.new_empty(rows, -(-max_seq_len // (4 * block)) * 4) - amax8_varlen(logits, seq_lens, out=keys) - blocks = torch.empty(rows, topk_blocks, dtype=torch.int32, device=logits.device) - topk_transform_paged_v2(keys, nblocks, None, blocks, 1, plan) - return blocks - - def build_sparse_schedule( blocks: torch.Tensor, seq_lens: torch.Tensor, @@ -622,126 +524,3 @@ def select_prefill_table( CANDIDATE_BLOCK_SIZE, ) out_positions.add_(torch.where(out_positions >= 0, data.request_starts[:, None], 0)) - - -def _publish_prefill_blocks( - *, - data: DeepGEMMPrefillData, - kv: Tuple[torch.Tensor, torch.Tensor], - topk: int, - topk_blocks: int, - block_size: int, -) -> Tuple[torch.Tensor, List[torch.Tensor]]: - """The source layer's own plain top-k, plus the ids of its best blocks per - request: ``[rows_b, min(topk_blocks, ceil(lc_b / block_size))]`` int32 each.""" - from sglang.kernels.ops.attention.dsv4 import topk_transform_ragged_v2 - - selected = data.empty_selection(topk) - device = data.request_starts.device - request_blocks = [ - torch.empty( - (rows, min(topk_blocks, -(-lc // block_size))), - dtype=torch.int32, - device=device, - ) - for rows, lc in zip(data.rows_per_request, data.lens_per_request) - ] - for tile, logits in score_tiles(data, kv, width_align=4): - _publish_tile_blocks( - data=data, - tile=tile, - logits=logits, - request_blocks=request_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, request_blocks - - -def _publish_tile_blocks( - *, - data: DeepGEMMPrefillData, - tile: slice, - logits: torch.Tensor, - request_blocks: List[torch.Tensor], - block_size: int, -) -> None: - lens = data.compress_lens[tile] - for request, rows, request_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, - ) - blocks = request_blocks[request][request_rows] - blocks.copy_( - select_candidate_block_ids( - logits=scores, - compress_lens=lens[rows, None], - topk_blocks=blocks.shape[1], - block_size=block_size, - ) - ) - - -def _consume_prefill_blocks( - *, - data: DeepGEMMPrefillData, - kv: Tuple[torch.Tensor, torch.Tensor], - topk: int, - request_blocks: List[torch.Tensor], - block_size: int, -) -> torch.Tensor: - """A consumer's top-k over its dense scores masked to the published blocks.""" - 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): - starts = data.request_starts[tile] - for request, rows, request_rows, lc in _requests_in_tile(data, tile): - logits[rows, :lc].masked_fill_( - ~candidate_block_mask( - blocks=request_blocks[request][request_rows], - width=lc, - block_size=block_size, - ), - -torch.inf, - ) - topk_transform_ragged_v2( - logits, - data.compress_lens[tile], - out_offsets=starts, - out_indices=selected[tile], - ) - selected[tile] = mask_topk_scores( - scores=logits, indices=selected[tile], offsets=starts - ) - # Free this tile's logits before the generator scores the next one. - del logits - return selected - - -def _requests_in_tile(data: DeepGEMMPrefillData, tile: slice): - """``(request, rows, request_rows, lc)`` of each request with a visible - position and rows in ``tile``: its rows within the tile and within itself.""" - start = 0 - for request, (num_rows, lc) in enumerate( - 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 ( - request, - slice(begin - tile.start, stop - tile.start), - slice(begin - start, stop - start), - lc, - ) - start = end diff --git a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/inputs.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py similarity index 73% rename from python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/inputs.py rename to python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py index 3877876cd858..e982419338af 100644 --- a/python/sglang/srt/layers/attention/dsv4/low_ratio_indexer/inputs.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py @@ -1,9 +1,9 @@ -"""What the attention backend hands a low-ratio indexer backend, and what it gets +"""What the attention backend hands the V4.1 indexer backends, and what it gets back.""" from __future__ import annotations -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, List, Optional, Protocol, TypeVar import msgspec import torch @@ -78,3 +78,39 @@ class CapturedPrefillInputs(msgspec.Struct, frozen=True, kw_only=True): 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 candidate scheme on prefill: what a publish carries is the backend's + own; the attention backend keeps it 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/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py b/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py index fe2962b3aea0..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,10 +4,7 @@ import torch from sglang.kernels.ops.attention.dsv4.fp4_indexer import quantize_fp4_indexer_tensor -from sglang.srt.layers.attention.dsv4.low_ratio_indexer import ( - deep_gemm_backend, - deep_gemm_utils, -) +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 @@ -54,7 +51,7 @@ 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 = deep_gemm_utils.DeepGEMMPrefillData( + 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"]], @@ -65,26 +62,38 @@ def run_dense(inputs, *, publish_candidates, candidates): weights=inputs["weights"], ) if publish_candidates: - selected, blocks = deep_gemm_backend._publish_prefill_blocks( + 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, deep_gemm_backend._DensePrefillBlocks(request_blocks=blocks) + return selected, dense_blocks.BlockIds( + blocks=blocks, rows_per_request=data.rows_per_request + ) if candidates is not None: - selected = deep_gemm_backend._consume_prefill_blocks( + selected = dense_blocks._consume_prefill_blocks( data=data, kv=inputs["kv"], topk=inputs["topk"], - request_blocks=candidates.request_blocks, + blocks=candidates.blocks, block_size=inputs["candidate_block_size"], ) return selected, None - return deep_gemm_utils.dense_prefill_topk( - data, inputs["kv"], topk=inputs["topk"] - ), 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): @@ -160,7 +169,7 @@ def test_ragged_source_consumer_and_replay(self): consumer_inputs["candidate_topk_blocks"] = 128 consumer_scores = dense_scores(consumer_inputs) with patch.object( - deep_gemm_utils, "_SCORE_BUDGET_BYTES", 128 << 10 + scoring, "_DEEP_GEMM_SCORE_BUDGET_BYTES", 128 << 10 ): selected, candidates = run_dense( inputs, publish_candidates=True, candidates=None @@ -168,7 +177,8 @@ def test_ragged_source_consumer_and_replay(self): 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: @@ -258,16 +268,17 @@ def test_unfiltered_and_zero_length_requests(self): 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) @@ -277,9 +288,7 @@ def test_empty_queries_or_context(self): 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 @@ -297,7 +306,7 @@ def checked_logits(*args, **kwargs): return logits with ( - patch.object(deep_gemm_utils, "_SCORE_BUDGET_BYTES", budget), + patch.object(scoring, "_DEEP_GEMM_SCORE_BUDGET_BYTES", budget), patch("deep_gemm.fp8_fp4_mqa_logits", new=checked_logits), ): selected, _ = run_dense( @@ -323,9 +332,7 @@ def test_score_memory_is_bounded(self): 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 = run_dense( inputs, publish_candidates=False, candidates=candidates 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 fea4c25f9525..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,11 @@ import msgspec import torch -from sglang.kernels.ops.attention.dsv4.fp4_indexer import quantize_fp4_indexer_tensor -from sglang.srt.layers.attention.dsv4.low_ratio_indexer.block_math import ( - select_candidate_blocks, +from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( + select_candidate_block_ids, ) -from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_utils import ( +from sglang.kernels.ops.attention.dsv4.fp4_indexer import quantize_fp4_indexer_tensor +from sglang.srt.layers.attention.dsv4.v41_indexer.scoring import ( DeepGEMMPrefillData, ) from sglang.test.ci.ci_register import register_cuda_ci @@ -115,15 +115,20 @@ 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 ) - return keep.unflatten(1, (-1, BLOCK)).any(-1) + 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.low_ratio_indexer.deep_gemm_backend import ( + from sglang.srt.layers.attention.dsv4.v41_indexer.sparse_table import ( publish_prefill_table, ) @@ -140,7 +145,7 @@ def publish_sparse(case: Case): def select_sparse(table, data: DeepGEMMPrefillData, k_cache: torch.Tensor): - from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_backend import ( + from sglang.srt.layers.attention.dsv4.v41_indexer.sparse_table import ( select_prefill_table, ) @@ -153,7 +158,7 @@ def select_sparse(table, data: DeepGEMMPrefillData, k_cache: torch.Tensor): def publish_dense(case: Case): """(the CP implementation's block ids, the source layer's own top-k).""" - from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_backend import ( + from sglang.srt.layers.attention.dsv4.v41_indexer.dense_blocks import ( _publish_prefill_blocks, ) @@ -168,7 +173,7 @@ def publish_dense(case: Case): def select_dense(blocks, case: Case): - from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_backend import ( + from sglang.srt.layers.attention.dsv4.v41_indexer.dense_blocks import ( _consume_prefill_blocks, ) @@ -176,7 +181,7 @@ def select_dense(blocks, case: Case): data=case.data, kv=case.kv, topk=TOPK, - request_blocks=blocks, + blocks=blocks, block_size=BLOCK, ) @@ -192,7 +197,7 @@ def picks(positions: torch.Tensor, row: int) -> set: class TestPrefillSparseIndexer(CustomTestCase): @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: @@ -257,6 +262,26 @@ def test_prefill_tail_rebuilds_the_last_rows(self): 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 64b27dae92ca..151bfdcea149 100644 --- a/test/registered/unit/layers/test_dsv41_candidate_blocks.py +++ b/test/registered/unit/layers/test_dsv41_candidate_blocks.py @@ -2,85 +2,117 @@ import torch -from sglang.srt.layers.attention.dsv4.low_ratio_indexer.block_math import ( - candidate_block_mask, +from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( + get_tail_row_indices, select_candidate_block_ids, - select_candidate_blocks, -) -from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_backend import ( - _DensePrefillBlocks, + topk_among_blocks, ) +from sglang.srt.layers.attention.dsv4.v41_indexer.dense_blocks import BlockIds 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, ) - torch.testing.assert_close( - candidate_block_mask(blocks=blocks, width=width, block_size=8), - torch.ones_like(scores, dtype=torch.bool), + 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 = torch.full((3, 2), -1, dtype=torch.int32) - self.assertFalse( - candidate_block_mask(blocks=blocks, width=19, block_size=8).any() - ) + blocks = select_candidate_block_ids( + visible, lens[:, None], topk_blocks=topk_blocks, block_size=8 + ) + 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) + ) + + 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()) - 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 = _DensePrefillBlocks(request_blocks=requests) - tail = candidates.tail([2, 0, 0]) + +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 1865705e870c..5e37e8e2989a 100644 --- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py +++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py @@ -917,8 +917,8 @@ class TestChunkedPagedDecode(CustomTestCase): chunks = [slice(0, 2), slice(2, 4), slice(4, 5)] def _inputs(self): - from sglang.srt.layers.attention.dsv4.low_ratio_indexer import DecodeInputs - from sglang.srt.layers.attention.dsv4.low_ratio_indexer.deep_gemm_utils import ( + from sglang.srt.layers.attention.dsv4.v41_indexer import DecodeInputs + from sglang.srt.layers.attention.dsv4.v41_indexer.scoring import ( DeepGEMMDecodeData, ) @@ -956,7 +956,7 @@ def _inputs(self): return inputs, data def _selection(self): - from sglang.srt.layers.attention.dsv4.low_ratio_indexer import Selection + 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()) @@ -982,12 +982,12 @@ def _assert_chunked_calls(self, deep_gemm, topk): self.assertIs(call.kwargs["topk_metadata"], plan) def test_candidate_publisher_takes_one_schedule_per_call(self): - from sglang.srt.layers.attention.dsv4.low_ratio_indexer import ( - deep_gemm_backend as mod, + from sglang.srt.layers.attention.dsv4.v41_indexer import ( + sparse_table as mod, ) inputs, data = self._inputs() - backend = object.__new__(mod.DeepGEMMCandidateBackend) + backend = object.__new__(mod.SparseTableBackend) backend.token_to_kv_pool = None backend.topk_blocks = 2 deep_gemm, topk = self._mocks() @@ -1034,12 +1034,12 @@ def test_candidate_publisher_takes_one_schedule_per_call(self): event.record.assert_called_once_with(stream) def test_dense_decode_takes_one_schedule_per_call(self): - from sglang.srt.layers.attention.dsv4.low_ratio_indexer import ( - dense_indexer as mod, + from sglang.srt.layers.attention.dsv4.v41_indexer import ( + plain as mod, ) inputs, data = self._inputs() - indexer = object.__new__(mod.DenseIndexer) + indexer = object.__new__(mod.PlainIndexer) indexer.token_to_kv_pool = None indexer.use_deep_gemm_decode = True deep_gemm, topk = self._mocks() @@ -1055,33 +1055,35 @@ def test_dense_decode_takes_one_schedule_per_call(self): class TestCandidateIndexerGating(CustomTestCase): def test_candidate_indexer_gating(self): - from sglang.srt.layers.attention.dsv4 import low_ratio_indexer - from sglang.srt.layers.attention.dsv4.low_ratio_indexer.torch_backend import ( - TorchCandidateBackend, + from sglang.srt.layers.attention.dsv4 import v41_indexer + from sglang.srt.layers.attention.dsv4.v41_indexer.dense_blocks import ( + DenseBlocksBackend, ) def sm100(value): - return patch.object(low_ratio_indexer, "is_sm100_or_newer", lambda: value) + return patch.object(v41_indexer, "is_sm100_or_newer", lambda: value) def make(topk_blocks, block_size): - return low_ratio_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, - ) + # 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 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 the portable backend. + # 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): - indexer = make(topk_blocks, block_size) - self.assertIsInstance(indexer.prefill, TorchCandidateBackend) - self.assertIsInstance(indexer.decode, TorchCandidateBackend) + 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 sm100(True), patch(flag, False): with self.assertRaises(RuntimeError): From 097df40dc4346e7332083e06d7530f35bcf10513 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 16:43:06 -0700 Subject: [PATCH 04/16] move dsv4.1 index ops into kernels/ops Co-Authored-By: DarkSharpness <2040703891@qq.com> --- .../ops/attention/dsv4/candidate_blocks.py | 118 ++++++++++- .../ops/attention/dsv4/candidate_table.py | 70 +++++++ .../ops/attention/dsv4/dense_prefill.py | 194 ++++++++++++++++++ .../sglang/kernels/ops/attention/dsv4/topk.py | 25 +++ .../layers/attention/deepseek_v4_backend.py | 25 ++- .../attention/dsv4/candidate_indexer.py | 85 +------- .../dsv4/candidate_indexer_deep_gemm.py | 122 +++-------- .../attention/dsv4/dense_prefill_indexer.py | 185 +---------------- .../attention/test_dense_prefill_indexer.py | 43 ++-- .../test_dsv41_prefill_sparse_indexer.py | 6 +- .../layers/test_dsv41_candidate_blocks.py | 4 +- 11 files changed, 472 insertions(+), 405 deletions(-) create mode 100644 python/sglang/kernels/ops/attention/dsv4/dense_prefill.py diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py index b70d0f2d2c41..5c422c054dee 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py @@ -1,11 +1,19 @@ -"""Per-row candidate block counts and sparse-row lengths for the paged indexer.""" +"""Candidate blocks of the two-level indexer: the per-row block counts and +sparse-row lengths, the JIT block top-k of the sparse table, and the torch +block ids and masks.""" + +from typing import Optional, Union import torch +import torch.nn.functional as F import triton import triton.language as tl from sglang.kernels.jit.utils import is_arch_support_pdl +from .candidate_table import CANDIDATE_BLOCK_SIZE, amax8_varlen +from .topk import plan_topk_v2, topk_transform_paged_v2 + @triton.jit def _candidate_row_lens_kernel( @@ -62,3 +70,111 @@ def candidate_row_lens( **pdl_kwargs, ) return nblocks, valid + + +def amax_topk_blocks( + logits: torch.Tensor, + seq_lens: torch.Tensor, + nblocks: torch.Tensor, + topk_blocks: int, + max_seq_len: Optional[int] = None, +) -> torch.Tensor: + """Per row the ``topk_blocks`` blocks of 8 positions with the largest block + maximum among its first ``seq_lens[b]`` positions, the newest block always + included: block ids in no particular order, ``-1`` past the row's count. + ``nblocks`` is ``ceil(seq_lens / 8)`` as int32.""" + rows = logits.shape[0] + block = CANDIDATE_BLOCK_SIZE + if max_seq_len is None: + max_seq_len = logits.shape[1] + # NOTE: plan cannot be the previous kernel of topk_transform_paged_v2 + plan = plan_topk_v2(nblocks) + # block maxima, the newest block +inf; the top-k reads each row up to nblocks + # only, so nothing past a row's keys is initialised (v2 needs stride % 4 == 0) + keys = logits.new_empty(rows, -(-max_seq_len // (4 * block)) * 4) + amax8_varlen(logits, seq_lens, out=keys) + blocks = torch.empty(rows, topk_blocks, dtype=torch.int32, device=logits.device) + topk_transform_paged_v2(keys, nblocks, None, blocks, 1, plan) + 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( + logits: torch.Tensor, + compress_lens: Union[torch.Tensor, int], + topk_blocks: int, + block_size: int, +) -> torch.return_types.topk: + 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, + ) + 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, + block_size: int, +) -> torch.Tensor: + top = _candidate_block_topk( + logits=logits, + compress_lens=compress_lens, + topk_blocks=topk_blocks, + block_size=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] diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_table.py b/python/sglang/kernels/ops/attention/dsv4/candidate_table.py index e956d2e6909f..bdcca9ebf3ca 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_table.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_table.py @@ -14,11 +14,14 @@ make_cpp_args, ) +from .topk import topk_transform_bf16_small from .utils import make_name if TYPE_CHECKING: pass +CANDIDATE_BLOCK_SIZE = 8 # positions per block; DeepGEMM accepts 8 or 16 + @cache_once def _jit_block_amax_module(): @@ -91,3 +94,70 @@ def sort_candidate_blocks( blocks, seq_lens, page_table, out_pages, page_size ) return out_pages + + +def build_sparse_indexer_schedule( + blocks: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + page_size: int, + q_dtype: torch.dtype, + request_ids: torch.Tensor, +) -> torch.Tensor: + """DeepGEMM's schedule for the published blocks: ``seq_lens`` ``[rows]`` + int32, ``page_table`` ``[rows, pages]`` int32 at the index pool's page size. + ``request_ids`` ``[rows]`` int32 lets DeepGEMM pair two rows of a request on + one KV pass; each row keeps its own block list and output layout, and paired + rows must share their page-table row.""" + import deep_gemm + + return deep_gemm.get_paged_sparse_mqa_logits_metadata( + seq_lens.contiguous(), + page_table, + request_ids, + page_size, + blocks, + q_dtype, + CANDIDATE_BLOCK_SIZE, + ) + + +def sparse_logits( + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + k_cache: torch.Tensor, + weights: torch.Tensor, + schedule: torch.Tensor, + topk_blocks: int, +) -> torch.Tensor: + """bf16 logits ``[rows, topk_blocks * 8]`` of the published blocks: ``q_fp4`` + ``[rows, 1, heads, 64]`` int8 with ``q_sf`` ``[rows, 1, heads]`` int32 (packed + ue8m0), ``k_cache`` ``[pages, page_size, 1, 68]`` uint8 whose page stride is + a multiple of 512 bytes, ``weights`` ``[rows, heads]`` bf16, ``schedule`` the + ``build_sparse_indexer_schedule`` of the ``topk_blocks`` published blocks.""" + import deep_gemm + + return deep_gemm.fp8_fp4_paged_sparse_mqa_logits( + (q_fp4, q_sf), + k_cache, + weights, + schedule, + topk_blocks, + CANDIDATE_BLOCK_SIZE, + ) + + +def topk_transform_sparse( + logits: torch.Tensor, + valid_lens: torch.Tensor, + phys_blocks: torch.Tensor, + page_indices: torch.Tensor, +) -> None: + """Top-``k`` (``k = page_indices.shape[1]``) of every row of the sparse + ``logits`` (bf16 ``[rows, topk_blocks * 8]``) within its first ``valid_lens[b]`` + columns, written as pool slots, ``-1`` where a row has fewer than ``k`` valid + columns, in no particular order; ``phys_blocks`` ``[rows, topk_blocks]`` int32 + holds the published blocks as pool slots / 8.""" + topk_transform_bf16_small( + logits, valid_lens, phys_blocks, page_indices, CANDIDATE_BLOCK_SIZE + ) diff --git a/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py b/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py new file mode 100644 index 000000000000..7bf6d18a36f8 --- /dev/null +++ b/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py @@ -0,0 +1,194 @@ +"""The DeepGEMM dense prefill scores of the ratio-1/2 index layers, tile by +tile under a memory budget, and the top-k over them: plain, or publishing / +consuming candidate block ids per request.""" + +from __future__ import annotations + +from typing import Iterator + +import torch + +from sglang.srt.layers.attention.mqa_logits_utils import ( + mqa_logits_row_bytes, + mqa_logits_rows_per_chunk, +) +from sglang.srt.utils.common import ceil_align + +from .candidate_blocks import ( + candidate_block_mask, + mask_topk_scores, + select_candidate_block_ids, +) +from .topk import topk_transform_ragged_v2 + +# 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: list[torch.Tensor] | None, +) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + selected = torch.full( + (q[0].shape[0], topk), -1, dtype=torch.int32, device=weights.device + ) + published = [] 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.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 = _rows_per_chunk(row, width, heads=q[0].shape[1]) + 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 _rows_per_chunk(rows: int, width: int, *, heads: int) -> int: + """Query rows per logits tile so one fp32 [rows, width] tile fits the + budget; the row count stays a multiple of the kernel's row alignment.""" + row_alignment = 128 // heads + rows_per_chunk = mqa_logits_rows_per_chunk( + num_rows=ceil_align(rows, row_alignment), + row_bytes=mqa_logits_row_bytes(width), + budget_bytes=_SCORE_BUDGET_BYTES, + ) + if rows_per_chunk is None: + return rows + return max(row_alignment, rows_per_chunk // row_alignment * row_alignment) + + +def score_tiles( + *, + q: tuple[torch.Tensor, torch.Tensor], + kv: tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + starts: torch.Tensor, + lengths: torch.Tensor, + context_lengths: list[int], + width_align: int = 4, +) -> Iterator[tuple[slice, torch.Tensor]]: + """The dense scores of the chunk row tile by row tile under the budget: + ``(rows, logits)`` with fp32 ``logits[i, j]`` the score of query row + ``rows.start + i`` against ``kv[starts + j]``, garbage past the row's + ``lengths``; the width is the largest of ``context_lengths`` (compressed + positions at each request's newest token) aligned to ``width_align`` + columns.""" + from deep_gemm import fp8_fp4_mqa_logits + + rows = q[0].shape[0] + width = ceil_align(max(context_lengths, default=0), width_align) + if rows == 0 or width == 0: + return + rows_per_chunk = _rows_per_chunk(rows, width, heads=q[0].shape[1]) + for offset in range(0, rows, rows_per_chunk): + tile = slice(offset, min(offset + rows_per_chunk, rows)) + tile_starts = starts[tile] + yield ( + tile, + fp8_fp4_mqa_logits( + (q[0][tile], q[1][tile]), + kv, + weights[tile], + tile_starts, + tile_starts + lengths[tile], + False, + width, + ), + ) + + +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: list[torch.Tensor] | None, + consume: list[torch.Tensor] | None, +) -> None: + from deep_gemm import fp8_fp4_mqa_logits + + 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][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][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/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index 74128f2b0f0e..bd26525738c5 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/python/sglang/kernels/ops/attention/dsv4/topk.py @@ -123,6 +123,31 @@ def topk_transform_paged( ) +def topk_transform_paged_torch( + scores: torch.Tensor, + seq_lens: torch.Tensor, + page_tables: torch.Tensor, + out_page_indices: torch.Tensor, + page_size: int, + out_raw_indices: Optional[torch.Tensor] = None, +) -> None: + """The torch ``topk_transform_paged``: the top-``k`` (``k = out_page_indices.shape[1]``) + of each row of ``scores`` within its first ``seq_lens[b]`` columns, ascending, as + pool slots through ``page_tables`` and, when given, as positions; ``-1`` where a + row has fewer than ``k`` columns.""" + topk = out_page_indices.shape[1] + columns = torch.arange(scores.shape[1], device=seq_lens.device) + lens_c = seq_lens.unsqueeze(-1) + # Columns past a row's length hold garbage. + s = scores.masked_fill(columns[None, :] >= lens_c, -torch.inf) + idx = s.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values + reach = idx < lens_c + slots = page_tables.gather(-1, idx // page_size) * page_size + (idx % page_size) + out_page_indices.copy_(torch.where(reach, slots, -1).to(torch.int32)) + if out_raw_indices is not None: + out_raw_indices.copy_(torch.where(reach, idx, -1).to(torch.int32)) + + # metadata is (batch+1, 2) int32: row 0 = {cluster_threshold, num_cluster_items}; # rows 1..N = {batch_id, seq_len} of items routed to the persistent cluster pool. _PLAN_METADATA_INTS_PER_BATCH = 2 diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 467110ffa1dd..1b6eda2f72c3 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -19,6 +19,10 @@ 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, ) @@ -38,6 +42,7 @@ 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,9 +69,7 @@ cut_request_masks, expand_index_page_table, make_candidate_indexer, - mask_topk_scores, published_masks, - select_candidate_blocks, ) from sglang.srt.layers.attention.dsv4.compressor_v2 import ( CompressorBackendMixin, @@ -3340,7 +3343,6 @@ def _low_ratio_index_topk_prefill_graph(self, layer, pos, q, w) -> None: page_indices = core.sparse_page_indices(ratio) raw_indices = core.sparse_raw_indices(ratio) topk = min(indexer.index_topk, width) - columns = torch.arange(width, device=lens.device) for rows, plan in metadata.row_chunks(): logits = deep_gemm_fp4_paged_mqa_logits( (q_fp4[rows], q_sf[rows]), @@ -3351,17 +3353,14 @@ def _low_ratio_index_topk_prefill_graph(self, layer, pos, q, w) -> None: plan, width, ) - lens_c = lens[rows].unsqueeze(-1) - # Columns past a row's length hold garbage. - s = logits.masked_fill(columns[None, :] >= lens_c, -torch.inf) - idx = s.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values - reach = idx < lens_c - slots = page_table[rows].gather(-1, idx // page_size) * page_size + ( - idx % page_size + 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, ) - page_indices[rows, :topk] = torch.where(reach, slots, -1).to(torch.int32) - if raw_indices is not None: - raw_indices[rows, :topk] = torch.where(reach, idx, -1).to(torch.int32) 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 ( diff --git a/python/sglang/srt/layers/attention/dsv4/candidate_indexer.py b/python/sglang/srt/layers/attention/dsv4/candidate_indexer.py index f6df650b7df8..9a78e2618c68 100644 --- a/python/sglang/srt/layers/attention/dsv4/candidate_indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/candidate_indexer.py @@ -2,11 +2,10 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import List, Optional, Tuple, Union +from typing import List, Optional, Tuple import msgspec import torch -import torch.nn.functional as F from sglang.srt.layers.attention.dsv4.metadata import PagedIndexerMetadata from sglang.srt.runtime_context import get_parallel, get_platform @@ -193,85 +192,3 @@ def tail(self, lengths: List[int]) -> PrefillCandidateBlocks: def published_masks(candidate) -> CandidateMasks: assert isinstance(candidate, CandidateMasks), "candidate masks missing" return candidate - - -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( - logits: torch.Tensor, - compress_lens: Union[torch.Tensor, int], - topk_blocks: int, - block_size: int, -) -> torch.return_types.topk: - 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, - ) - 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, - block_size: int, -) -> torch.Tensor: - top = _candidate_block_topk( - logits=logits, - compress_lens=compress_lens, - topk_blocks=topk_blocks, - block_size=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] 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 index 26f308133897..11ffee7e24b3 100644 --- a/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py +++ b/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py @@ -5,13 +5,20 @@ import torch -from sglang.kernels.ops.attention.dsv4.candidate_blocks import candidate_row_lens +from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( + amax_topk_blocks, + candidate_row_lens, +) from sglang.kernels.ops.attention.dsv4.candidate_table import ( + CANDIDATE_BLOCK_SIZE, amax8_varlen, + build_sparse_indexer_schedule, sort_candidate_blocks, + sparse_logits, + topk_transform_sparse, ) +from sglang.kernels.ops.attention.dsv4.dense_prefill import score_tiles from sglang.kernels.ops.attention.dsv4.topk import ( - plan_topk_v2, topk_transform_bf16_small, topk_transform_paged_v2, topk_transform_ragged_v2, @@ -25,15 +32,12 @@ ) from sglang.srt.layers.attention.dsv4.dense_prefill_indexer import ( DenseCandidateIndexer, - score_tiles, ) from sglang.srt.layers.attention.dsv4.indexer import ( deep_gemm_fp4_paged_mqa_logits, topk_transform_paged_from_metadata, ) -CANDIDATE_BLOCK_SIZE = 8 # positions per block; DeepGEMM accepts 8 or 16 - @dataclass class SparseBlockTable(CandidateMetadata): @@ -53,96 +57,6 @@ class SparseBlockTable(CandidateMetadata): ready: torch.cuda.Event -def amax_topk_blocks( - logits: torch.Tensor, - seq_lens: torch.Tensor, - nblocks: torch.Tensor, - topk_blocks: int, - max_seq_len: Optional[int] = None, -) -> torch.Tensor: - """Per row the ``topk_blocks`` blocks of 8 positions with the largest block - maximum among its first ``seq_lens[b]`` positions, the newest block always - included: block ids in no particular order, ``-1`` past the row's count. - ``nblocks`` is ``ceil(seq_lens / 8)`` as int32.""" - rows = logits.shape[0] - block = CANDIDATE_BLOCK_SIZE - if max_seq_len is None: - max_seq_len = logits.shape[1] - # NOTE: plan cannot be the previous kernel of topk_transform_paged_v2 - plan = plan_topk_v2(nblocks) - # block maxima, the newest block +inf; the top-k reads each row up to nblocks - # only, so nothing past a row's keys is initialised (v2 needs stride % 4 == 0) - keys = logits.new_empty(rows, -(-max_seq_len // (4 * block)) * 4) - amax8_varlen(logits, seq_lens, out=keys) - blocks = torch.empty(rows, topk_blocks, dtype=torch.int32, device=logits.device) - topk_transform_paged_v2(keys, nblocks, None, blocks, 1, plan) - return blocks - - -def build_sparse_indexer_schedule( - blocks: torch.Tensor, - seq_lens: torch.Tensor, - page_table: torch.Tensor, - page_size: int, - q_dtype: torch.dtype, - request_ids: torch.Tensor, -) -> torch.Tensor: - """DeepGEMM's schedule for the published blocks: ``seq_lens`` ``[rows]`` - int32, ``page_table`` ``[rows, pages]`` int32 at the index pool's page size. - ``request_ids`` ``[rows]`` int32 lets DeepGEMM pair two rows of a request on - one KV pass; each row keeps its own block list and output layout, and paired - rows must share their page-table row.""" - import deep_gemm - - return deep_gemm.get_paged_sparse_mqa_logits_metadata( - seq_lens.contiguous(), - page_table, - request_ids, - page_size, - blocks, - q_dtype, - CANDIDATE_BLOCK_SIZE, - ) - - -def sparse_logits( - q_fp4: torch.Tensor, - q_sf: torch.Tensor, - k_cache: torch.Tensor, - weights: torch.Tensor, - table: SparseBlockTable, -) -> torch.Tensor: - """bf16 logits ``[rows, topk_blocks * 8]`` of the published blocks: ``q_fp4`` - ``[rows, 1, heads, 64]`` int8 with ``q_sf`` ``[rows, 1, heads]`` int32 (packed - ue8m0), ``k_cache`` ``[pages, page_size, 1, 68]`` uint8 whose page stride is - a multiple of 512 bytes, ``weights`` ``[rows, heads]`` bf16.""" - import deep_gemm - - return deep_gemm.fp8_fp4_paged_sparse_mqa_logits( - (q_fp4, q_sf), - k_cache, - weights, - table.schedule, - table.blocks.shape[1], - CANDIDATE_BLOCK_SIZE, - ) - - -def topk_transform_sparse( - logits: torch.Tensor, - valid_lens: torch.Tensor, - table: SparseBlockTable, - page_indices: torch.Tensor, -) -> None: - """Top-``k`` (``k = page_indices.shape[1]``) of every row of the sparse - ``logits`` (bf16 ``[rows, topk_blocks * 8]``) within its first ``valid_lens[b]`` - columns, written as pool slots, ``-1`` where a row has fewer than ``k`` valid - columns, in no particular order.""" - topk_transform_bf16_small( - logits, valid_lens, table.phys_blocks, page_indices, CANDIDATE_BLOCK_SIZE - ) - - @dataclass class PrefillSparseBlockTable(SparseBlockTable): """Published by a prefill chunk's candidate source: one row per query token. @@ -347,7 +261,8 @@ def _scores(self, table: SparseBlockTable, inputs: IndexerInputs) -> torch.Tenso inputs.q_sf, inputs.k_cache, inputs.weights.to(torch.bfloat16), - table, + table.schedule, + table.blocks.shape[1], ) def select_decode( @@ -362,7 +277,7 @@ def select_decode( 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, page_indices) + topk_transform_sparse(logits, table.valid_lens, table.phys_blocks, page_indices) def publish_prefill( self, inputs: PrefillIndexerInputs, out_positions: torch.Tensor @@ -377,7 +292,15 @@ def publish_prefill( 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 score_tiles(inputs, width_align=8): + for tile, logits in score_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, + width_align=8, + ): lens = inputs.compress_lens[tile] topk_transform_ragged_v2( logits, @@ -500,7 +423,8 @@ def select_prefill( inputs.q_sf.view(rows, 1, heads), inputs.k_cache, inputs.weights.to(torch.bfloat16), - table, + table.schedule, + table.blocks.shape[1], ) # request-relative compressed positions, -1 padded topk_transform_bf16_small( diff --git a/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py b/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py index d4a478a24de1..d428e39c8ee4 100644 --- a/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py @@ -1,132 +1,16 @@ from __future__ import annotations -from typing import Iterator, List, Optional +from typing import List, Optional import torch +from sglang.kernels.ops.attention.dsv4.dense_prefill import dense_prefill_topk from sglang.srt.layers.attention.dsv4.candidate_indexer import ( CandidateIndexer, CandidateMetadata, PrefillCandidateBlocks, PrefillIndexerInputs, - candidate_block_mask, - mask_topk_scores, - select_candidate_block_ids, ) -from sglang.srt.layers.attention.mqa_logits_utils import ( - mqa_logits_row_bytes, - mqa_logits_rows_per_chunk, -) -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 = _rows_per_chunk(row, width, heads=q[0].shape[1]) - 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 _rows_per_chunk(rows: int, width: int, *, heads: int) -> int: - """Query rows per logits tile so one fp32 [rows, width] tile fits the - budget; the row count stays a multiple of the kernel's row alignment.""" - row_alignment = 128 // heads - rows_per_chunk = mqa_logits_rows_per_chunk( - num_rows=ceil_align(rows, row_alignment), - row_bytes=mqa_logits_row_bytes(width), - budget_bytes=_SCORE_BUDGET_BYTES, - ) - if rows_per_chunk is None: - return rows - return max(row_alignment, rows_per_chunk // row_alignment * row_alignment) - - -def score_tiles( - inputs: PrefillIndexerInputs, *, width_align: int = 4 -) -> Iterator[tuple[slice, torch.Tensor]]: - """The dense scores of ``inputs`` row tile by row tile under the budget: - ``(rows, logits)`` with fp32 ``logits[i, j]`` the score of query row - ``rows.start + i`` against ``kv[request_starts + j]``, garbage past the row's - ``compress_lens``; the width is the batch's largest context aligned to - ``width_align`` columns.""" - from deep_gemm import fp8_fp4_mqa_logits - - rows = inputs.num_rows - width = ceil_align(max(inputs.lens_per_request, default=0), width_align) - if rows == 0 or width == 0: - return - rows_per_chunk = _rows_per_chunk(rows, width, heads=inputs.q_fp4.shape[1]) - for offset in range(0, rows, rows_per_chunk): - tile = slice(offset, min(offset + rows_per_chunk, rows)) - starts = inputs.request_starts[tile] - yield ( - tile, - fp8_fp4_mqa_logits( - (inputs.q_fp4[tile], inputs.q_sf[tile]), - inputs.kv, - inputs.weights[tile], - starts, - starts + inputs.compress_lens[tile], - False, - width, - ), - ) def _dense_topk( @@ -149,10 +33,12 @@ def _dense_topk( candidate_topk_blocks=topk_blocks, candidate_block_size=block_size, publish_candidates=publish, - candidates=candidates, + candidates=candidates.request_blocks if candidates is not None else None, ) out_positions.copy_(selected) - return published + if published is None: + return None + return PrefillCandidateBlocks(request_blocks=published) def plain_prefill_topk( @@ -216,62 +102,3 @@ def prefill_tail( ) -> 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 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/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py b/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py index 62c1ff6ff06b..76dc769d7c6d 100644 --- a/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py +++ b/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py @@ -3,8 +3,8 @@ import torch +from sglang.kernels.ops.attention.dsv4 import dense_prefill 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.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -120,16 +120,14 @@ def test_ragged_source_consumer_and_replay(self): consumer_inputs["kv"] = inputs["kv"] consumer_inputs["candidate_topk_blocks"] = 128 consumer_scores = dense_scores(consumer_inputs) - with patch.object( - dense_prefill_indexer, "_SCORE_BUDGET_BYTES", 128 << 10 - ): - selected, candidates = dense_prefill_indexer.dense_prefill_topk( + with patch.object(dense_prefill, "_SCORE_BUDGET_BYTES", 128 << 10): + selected, candidates = dense_prefill.dense_prefill_topk( **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"], candidates ): local = scores[row : row + queries, :context] if queries and context: @@ -173,7 +171,7 @@ 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( + selected, published = dense_prefill.dense_prefill_topk( **consumer_inputs, publish_candidates=False, candidates=candidates, @@ -201,10 +199,13 @@ def test_ragged_source_consumer_and_replay(self): ) ), ) - selected, _ = dense_prefill_indexer.dense_prefill_topk( + selected, _ = dense_prefill.dense_prefill_topk( **tail_inputs, publish_candidates=False, - candidates=candidates.tail(tail_lengths), + candidates=[ + b[b.shape[0] - t :] + for b, t in zip(candidates, tail_lengths) + ], ) self.assert_topk(tail_inputs, selected, consumer_scores[rows]) @@ -213,13 +214,13 @@ 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( + selected, candidates = dense_prefill.dense_prefill_topk( **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 candidates], [q for q, _ in request_lengths], ) @@ -232,15 +233,13 @@ def test_empty_queries_or_context(self): ): with self.subTest(request_lengths=request_lengths): inputs = make_inputs(request_lengths) - selected, candidates = dense_prefill_indexer.dense_prefill_topk( + selected, candidates = dense_prefill.dense_prefill_topk( **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(b.shape) for b in candidates], block_shapes) def test_score_budget_includes_allocation_padding(self): from deep_gemm import fp8_fp4_mqa_logits @@ -258,10 +257,10 @@ def checked_logits(*args, **kwargs): return logits with ( - patch.object(dense_prefill_indexer, "_SCORE_BUDGET_BYTES", budget), + patch.object(dense_prefill, "_SCORE_BUDGET_BYTES", budget), patch("deep_gemm.fp8_fp4_mqa_logits", new=checked_logits), ): - selected, _ = dense_prefill_indexer.dense_prefill_topk( + selected, _ = dense_prefill.dense_prefill_topk( **inputs, publish_candidates=True, candidates=None ) self.assertTrue(allocations) @@ -276,7 +275,7 @@ 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( + selected, candidates = dense_prefill.dense_prefill_topk( **inputs, publish_candidates=True, candidates=None ) torch.cuda.synchronize() @@ -284,11 +283,9 @@ def test_score_memory_is_bounded(self): 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[0].shape), (16384, 2048)) del selected - selected, published = dense_prefill_indexer.dense_prefill_topk( + selected, published = dense_prefill.dense_prefill_topk( **inputs, publish_candidates=False, candidates=candidates ) self.assertIsNone(published) @@ -297,7 +294,7 @@ 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( + selected, published = dense_prefill.dense_prefill_topk( **inputs, publish_candidates=False, candidates=candidates ) torch.cuda.synchronize() 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 96d146a3cf18..e5d62892ecff 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,9 @@ import msgspec import torch +from sglang.kernels.ops.attention.dsv4.candidate_blocks import select_candidate_blocks from sglang.kernels.ops.attention.dsv4.fp4_indexer import quantize_fp4_indexer_tensor -from sglang.srt.layers.attention.dsv4.candidate_indexer import ( - PrefillIndexerInputs, - select_candidate_blocks, -) +from sglang.srt.layers.attention.dsv4.candidate_indexer import PrefillIndexerInputs from sglang.srt.layers.attention.dsv4.dense_prefill_indexer import ( DenseCandidateIndexer, ) diff --git a/test/registered/unit/layers/test_dsv41_candidate_blocks.py b/test/registered/unit/layers/test_dsv41_candidate_blocks.py index 79915cf5c29c..ff33378fe5b3 100644 --- a/test/registered/unit/layers/test_dsv41_candidate_blocks.py +++ b/test/registered/unit/layers/test_dsv41_candidate_blocks.py @@ -2,12 +2,12 @@ import torch -from sglang.srt.layers.attention.dsv4.candidate_indexer import ( - PrefillCandidateBlocks, +from sglang.kernels.ops.attention.dsv4.candidate_blocks import ( candidate_block_mask, select_candidate_block_ids, select_candidate_blocks, ) +from sglang.srt.layers.attention.dsv4.candidate_indexer import PrefillCandidateBlocks from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase From 0b1a636067d1b17501ddb3db0482e27cc5c8526b Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 16:53:41 -0700 Subject: [PATCH 05/16] use kernels/ops index ops in v41_indexer --- .../ops/attention/dsv4/dense_prefill.py | 126 +----------------- .../attention/dsv4/v41_indexer/plain.py | 19 ++- .../attention/dsv4/v41_indexer/scoring.py | 58 ++------ .../dsv4/v41_indexer/sparse_table.py | 70 ++-------- .../attention/test_dense_prefill_indexer.py | 7 +- .../unit/layers/test_dsv4_nonpaged_indexer.py | 10 +- 6 files changed, 40 insertions(+), 250 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py b/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py index 7bf6d18a36f8..c54222c4aa72 100644 --- a/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py +++ b/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py @@ -1,6 +1,5 @@ """The DeepGEMM dense prefill scores of the ratio-1/2 index layers, tile by -tile under a memory budget, and the top-k over them: plain, or publishing / -consuming candidate block ids per request.""" +tile under a memory budget.""" from __future__ import annotations @@ -14,75 +13,11 @@ ) from sglang.srt.utils.common import ceil_align -from .candidate_blocks import ( - candidate_block_mask, - mask_topk_scores, - select_candidate_block_ids, -) -from .topk import topk_transform_ragged_v2 - # TODO: use a per-forward mqa_logits_budget_bytes() budget that also -# leaves room for candidate masks and block-selection scratch. +# leaves room for candidate block ids 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: list[torch.Tensor] | None, -) -> tuple[torch.Tensor, list[torch.Tensor] | None]: - selected = torch.full( - (q[0].shape[0], topk), -1, dtype=torch.int32, device=weights.device - ) - published = [] 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.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 = _rows_per_chunk(row, width, heads=q[0].shape[1]) - 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 _rows_per_chunk(rows: int, width: int, *, heads: int) -> int: """Query rows per logits tile so one fp32 [rows, width] tile fits the budget; the row count stays a multiple of the kernel's row alignment.""" @@ -135,60 +70,3 @@ def score_tiles( width, ), ) - - -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: list[torch.Tensor] | None, - consume: list[torch.Tensor] | None, -) -> None: - from deep_gemm import fp8_fp4_mqa_logits - - 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][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][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/v41_indexer/plain.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/plain.py index 82d87228563c..767254827032 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/plain.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/plain.py @@ -11,6 +11,7 @@ import torch +from sglang.kernels.ops.attention.dsv4.topk import topk_transform_paged_torch from sglang.srt.layers.attention.dsv4.indexer import ( deep_gemm_fp4_paged_mqa_logits, topk_transform_paged_from_metadata, @@ -172,7 +173,6 @@ def _deep_gemm_prefill_captured( lens = metadata.compressed_seq_lens page_table = metadata.page_table topk = min(indexer.index_topk, width) - columns = torch.arange(width, device=lens.device) for rows, plan in metadata.row_chunks(): logits = deep_gemm_fp4_paged_mqa_logits( (q_fp4[rows], q_sf[rows]), @@ -183,14 +183,11 @@ def _deep_gemm_prefill_captured( plan, width, ) - lens_c = lens[rows].unsqueeze(-1) - # Columns past a row's length hold garbage. - s = logits.masked_fill(columns[None, :] >= lens_c, -torch.inf) - idx = s.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values - reach = idx < lens_c - slots = page_table[rows].gather(-1, idx // page_size) * page_size + ( - idx % page_size + 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, ) - out.page_indices[rows, :topk] = torch.where(reach, slots, -1).to(torch.int32) - if out.raw_indices is not None: - out.raw_indices[rows, :topk] = torch.where(reach, idx, -1).to(torch.int32) diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py index ce1e4d4d601c..722e14d85c07 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py @@ -16,12 +16,10 @@ import msgspec import torch -from sglang.kernels.ops.attention.dsv4.fp4_indexer import fp4_index_logits_decode -from sglang.srt.layers.attention.mqa_logits_utils import ( - mqa_logits_row_bytes, - mqa_logits_rows_per_chunk, +from sglang.kernels.ops.attention.dsv4.dense_prefill import ( + score_tiles as score_tiles_op, ) -from sglang.srt.utils.common import ceil_align +from sglang.kernels.ops.attention.dsv4.fp4_indexer import fp4_index_logits_decode from .types import ( DecodeInputs, @@ -33,10 +31,6 @@ 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 @@ -229,28 +223,15 @@ def score_tiles( ``rows.start + i`` against ``kv[request_starts + j]``, garbage past the row's ``compress_lens``; the width is the batch's largest context aligned to ``width_align`` columns.""" - from deep_gemm import fp8_fp4_mqa_logits - - rows = data.num_rows - width = ceil_align(max(data.lens_per_request, default=0), width_align) - if rows == 0 or width == 0: - return - rows_per_chunk = _rows_per_chunk(rows, width, heads=data.q_fp4.shape[1]) - for offset in range(0, rows, rows_per_chunk): - tile = slice(offset, min(offset + rows_per_chunk, rows)) - starts = data.request_starts[tile] - yield ( - tile, - fp8_fp4_mqa_logits( - (data.q_fp4[tile], data.q_sf[tile]), - kv, - data.weights[tile], - starts, - starts + data.compress_lens[tile], - False, - width, - ), - ) + yield from score_tiles_op( + 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, + width_align=width_align, + ) def dense_prefill_topk( @@ -276,21 +257,6 @@ def dense_prefill_topk( return selected -# TODO(dark): make it SM aware; align to SM count to avoid wave quantization. -def _rows_per_chunk(rows: int, width: int, *, heads: int) -> int: - """Query rows per logits tile so one fp32 [rows, width] tile fits the - budget; the row count stays a multiple of the kernel's row alignment.""" - row_alignment = 128 // heads - rows_per_chunk = mqa_logits_rows_per_chunk( - num_rows=ceil_align(rows, row_alignment), - row_bytes=mqa_logits_row_bytes(width), - budget_bytes=_DEEP_GEMM_SCORE_BUDGET_BYTES, - ) - if rows_per_chunk is None: - return rows - return max(row_alignment, rows_per_chunk // row_alignment * row_alignment) - - # ---------- torch ---------- 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 index 0273ca1ba11f..59d3a76c2dae 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py @@ -18,8 +18,11 @@ get_tail_row_indices, ) from sglang.kernels.ops.attention.dsv4.candidate_table import ( + CANDIDATE_BLOCK_SIZE, amax8_varlen, + build_sparse_indexer_schedule, sort_candidate_blocks, + sparse_logits, ) from sglang.kernels.ops.attention.dsv4.topk import ( topk_transform_bf16_small, @@ -51,8 +54,6 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool -CANDIDATE_BLOCK_SIZE = 8 # positions per block; DeepGEMM accepts 8 or 16 - class _SparseTable(CandidateMetadata, msgspec.Struct): # [rows, topk_blocks] int32: ascending logical block ids, valid for the first @@ -222,7 +223,7 @@ def publish_decode(self, inputs: DecodeInputs, out: Selection): metadata.page_table, metadata.compressed_page_size, ) - schedule = build_sparse_schedule( + schedule = build_sparse_indexer_schedule( blocks, seq_lens, metadata.page_table, @@ -292,7 +293,7 @@ def _publish_decode_chunked( valid_len_chunks.append(row_valid_lens) blocks = torch.cat(block_chunks) - schedule = build_sparse_schedule( + schedule = build_sparse_indexer_schedule( blocks, seq_lens, metadata.page_table, @@ -319,12 +320,13 @@ def consume_decode( 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 = deep_gemm_sparse_logits( + logits = sparse_logits( data.q_fp4, data.q_sf, data.k_cache, data.weights.to(torch.bfloat16), - published, + published.schedule, + published.blocks.shape[1], ) assert out.raw_indices is None topk_transform_bf16_small( @@ -355,55 +357,6 @@ def _get_request_ids(self, inputs: DecodeInputs, rows: int, device: torch.device return buf[:rows] -def build_sparse_schedule( - blocks: torch.Tensor, - seq_lens: torch.Tensor, - page_table: torch.Tensor, - page_size: int, - q_dtype: torch.dtype, - request_ids: torch.Tensor, -) -> torch.Tensor: - """DeepGEMM's schedule for the published blocks: ``seq_lens`` ``[rows]`` - int32, ``page_table`` ``[rows, pages]`` int32 at the index pool's page size. - ``request_ids`` ``[rows]`` int32 groups the rows of one request onto one KV - pass; each row keeps its own block list and output layout, and rows grouped - together must share their page-table row.""" - import deep_gemm - - return deep_gemm.get_paged_sparse_mqa_logits_metadata( - seq_lens.contiguous(), - page_table, - request_ids, - page_size, - blocks, - q_dtype, - CANDIDATE_BLOCK_SIZE, - ) - - -def deep_gemm_sparse_logits( - q_fp4: torch.Tensor, - q_sf: torch.Tensor, - k_cache: torch.Tensor, - weights: torch.Tensor, - table: _SparseTable, -) -> torch.Tensor: - """bf16 logits ``[rows, topk_blocks * 8]`` of the published blocks: ``q_fp4`` - ``[rows, 1, heads, 64]`` int8 with ``q_sf`` ``[rows, 1, heads]`` int32 (packed - ue8m0), ``k_cache`` ``[pages, page_size, 1, 68]`` uint8 whose page stride is - a multiple of 512 bytes, ``weights`` ``[rows, heads]`` bf16.""" - import deep_gemm - - return deep_gemm.fp8_fp4_paged_sparse_mqa_logits( - (q_fp4, q_sf), - k_cache, - weights, - table.schedule, - table.blocks.shape[1], - CANDIDATE_BLOCK_SIZE, - ) - - def _build_prefill_table( *, blocks: torch.Tensor, @@ -418,7 +371,7 @@ def _build_prefill_table( ) -> _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_schedule( + schedule = build_sparse_indexer_schedule( blocks, compress_lens, page_table, page_size, q_dtype, request_ids ) ready = torch.cuda.Event() @@ -508,12 +461,13 @@ def select_prefill_table( ``publish_prefill_table`` writes; ``k_cache`` is the layer's index-K pool, ``[pages, page_size, 1, 68]`` uint8.""" rows, heads = data.q_sf.shape - logits = deep_gemm_sparse_logits( + 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, + table.schedule, + table.blocks.shape[1], ) # request-relative compressed positions, -1 padded topk_transform_bf16_small( 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 9182d7b21bcf..fc4c039fb3d6 100644 --- a/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py +++ b/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py @@ -3,6 +3,7 @@ import torch +from sglang.kernels.ops.attention.dsv4 import dense_prefill from sglang.kernels.ops.attention.dsv4.fp4_indexer import quantize_fp4_indexer_tensor from sglang.srt.layers.attention.dsv4.v41_indexer import dense_blocks, scoring from sglang.test.ci.ci_register import register_cuda_ci @@ -168,9 +169,7 @@ def test_ragged_source_consumer_and_replay(self): consumer_inputs["kv"] = inputs["kv"] consumer_inputs["candidate_topk_blocks"] = 128 consumer_scores = dense_scores(consumer_inputs) - with patch.object( - scoring, "_DEEP_GEMM_SCORE_BUDGET_BYTES", 128 << 10 - ): + with patch.object(dense_prefill, "_SCORE_BUDGET_BYTES", 128 << 10): selected, candidates = run_dense( inputs, publish_candidates=True, candidates=None ) @@ -309,7 +308,7 @@ def checked_logits(*args, **kwargs): return logits with ( - patch.object(scoring, "_DEEP_GEMM_SCORE_BUDGET_BYTES", budget), + patch.object(dense_prefill, "_SCORE_BUDGET_BYTES", budget), patch("deep_gemm.fp8_fp4_mqa_logits", new=checked_logits), ): selected, _ = run_dense( diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py index 5e37e8e2989a..253140f412a9 100644 --- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py +++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py @@ -982,9 +982,7 @@ def _assert_chunked_calls(self, deep_gemm, topk): 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, - ) + from sglang.srt.layers.attention.dsv4.v41_indexer import sparse_table as mod inputs, data = self._inputs() backend = object.__new__(mod.SparseTableBackend) @@ -1019,7 +1017,7 @@ def test_candidate_publisher_takes_one_schedule_per_call(self): ), patch.object( mod, - "build_sparse_schedule", + "build_sparse_indexer_schedule", return_value=torch.tensor([7], dtype=torch.uint8), ), patch.object(mod.torch.cuda, "Event", return_value=event), @@ -1034,9 +1032,7 @@ def test_candidate_publisher_takes_one_schedule_per_call(self): 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 ( - plain as mod, - ) + from sglang.srt.layers.attention.dsv4.v41_indexer import plain as mod inputs, data = self._inputs() indexer = object.__new__(mod.PlainIndexer) From 32e7ccf5e619290f6f405b47a8a3d25791a235c5 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 17:06:44 -0700 Subject: [PATCH 06/16] score budget as a parameter Co-Authored-By: DarkSharpness <2040703891@qq.com> --- .../ops/attention/dsv4/dense_prefill.py | 30 +++++++++++-------- .../dsv4/candidate_indexer_deep_gemm.py | 2 ++ .../attention/dsv4/dense_prefill_indexer.py | 5 ++++ .../attention/test_dense_prefill_indexer.py | 8 +++-- 4 files changed, 30 insertions(+), 15 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py b/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py index 7bf6d18a36f8..9b3fb2bf0731 100644 --- a/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py +++ b/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py @@ -1,6 +1,6 @@ """The DeepGEMM dense prefill scores of the ratio-1/2 index layers, tile by -tile under a memory budget, and the top-k over them: plain, or publishing / -consuming candidate block ids per request.""" +tile under a caller-given memory budget, and the top-k over them: plain, or +publishing / consuming candidate block ids per request.""" from __future__ import annotations @@ -21,10 +21,6 @@ ) from .topk import topk_transform_ragged_v2 -# 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( *, @@ -39,7 +35,9 @@ def dense_prefill_topk( candidate_block_size: int, publish_candidates: bool, candidates: list[torch.Tensor] | None, + budget_bytes: int, ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + """``budget_bytes`` bounds one fp32 logits tile.""" selected = torch.full( (q[0].shape[0], topk), -1, dtype=torch.int32, device=weights.device ) @@ -63,7 +61,9 @@ def dense_prefill_topk( 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 = _rows_per_chunk(row, width, heads=q[0].shape[1]) + rows_per_chunk = _rows_per_chunk( + row, width, heads=q[0].shape[1], budget_bytes=budget_bytes + ) for offset in range(0, row, rows_per_chunk): rows = slice(offset, min(offset + rows_per_chunk, row)) _select_tile( @@ -83,14 +83,14 @@ def dense_prefill_topk( return selected, published -def _rows_per_chunk(rows: int, width: int, *, heads: int) -> int: - """Query rows per logits tile so one fp32 [rows, width] tile fits the - budget; the row count stays a multiple of the kernel's row alignment.""" +def _rows_per_chunk(rows: int, width: int, *, heads: int, budget_bytes: int) -> int: + """Query rows per logits tile so one fp32 [rows, width] tile fits + ``budget_bytes``; the row count stays a multiple of the kernel's row alignment.""" row_alignment = 128 // heads rows_per_chunk = mqa_logits_rows_per_chunk( num_rows=ceil_align(rows, row_alignment), row_bytes=mqa_logits_row_bytes(width), - budget_bytes=_SCORE_BUDGET_BYTES, + budget_bytes=budget_bytes, ) if rows_per_chunk is None: return rows @@ -105,9 +105,11 @@ def score_tiles( starts: torch.Tensor, lengths: torch.Tensor, context_lengths: list[int], + budget_bytes: int, width_align: int = 4, ) -> Iterator[tuple[slice, torch.Tensor]]: - """The dense scores of the chunk row tile by row tile under the budget: + """The dense scores of the chunk row tile by row tile, each fp32 tile within + ``budget_bytes``: ``(rows, logits)`` with fp32 ``logits[i, j]`` the score of query row ``rows.start + i`` against ``kv[starts + j]``, garbage past the row's ``lengths``; the width is the largest of ``context_lengths`` (compressed @@ -119,7 +121,9 @@ def score_tiles( width = ceil_align(max(context_lengths, default=0), width_align) if rows == 0 or width == 0: return - rows_per_chunk = _rows_per_chunk(rows, width, heads=q[0].shape[1]) + rows_per_chunk = _rows_per_chunk( + rows, width, heads=q[0].shape[1], budget_bytes=budget_bytes + ) for offset in range(0, rows, rows_per_chunk): tile = slice(offset, min(offset + rows_per_chunk, rows)) tile_starts = starts[tile] 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 index 11ffee7e24b3..e31a317c3487 100644 --- a/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py +++ b/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py @@ -31,6 +31,7 @@ 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 ( @@ -299,6 +300,7 @@ def publish_prefill( 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] diff --git a/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py b/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py index d428e39c8ee4..3ac4650b2ddc 100644 --- a/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py @@ -12,6 +12,10 @@ PrefillIndexerInputs, ) +# 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_topk( inputs: PrefillIndexerInputs, @@ -34,6 +38,7 @@ def _dense_topk( candidate_block_size=block_size, publish_candidates=publish, candidates=candidates.request_blocks if candidates is not None else None, + budget_bytes=_SCORE_BUDGET_BYTES, ) out_positions.copy_(selected) if published is None: 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 76dc769d7c6d..957fb65d9203 100644 --- a/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py +++ b/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py @@ -45,6 +45,7 @@ def make_inputs(request_lengths, ratio=1, zero_queries=False, seed=17): topk=512, candidate_topk_blocks=2, candidate_block_size=8, + budget_bytes=2 << 30, ) @@ -120,7 +121,10 @@ def test_ragged_source_consumer_and_replay(self): consumer_inputs["kv"] = inputs["kv"] consumer_inputs["candidate_topk_blocks"] = 128 consumer_scores = dense_scores(consumer_inputs) - with patch.object(dense_prefill, "_SCORE_BUDGET_BYTES", 128 << 10): + with ( + patch.dict(inputs, budget_bytes=128 << 10), + patch.dict(consumer_inputs, budget_bytes=128 << 10), + ): selected, candidates = dense_prefill.dense_prefill_topk( **inputs, publish_candidates=True, candidates=None ) @@ -257,7 +261,7 @@ def checked_logits(*args, **kwargs): return logits with ( - patch.object(dense_prefill, "_SCORE_BUDGET_BYTES", budget), + patch.dict(inputs, budget_bytes=budget), patch("deep_gemm.fp8_fp4_mqa_logits", new=checked_logits), ): selected, _ = dense_prefill.dense_prefill_topk( From 4930de767ee6c9555e986c260fe67e879927767f Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 17:08:49 -0700 Subject: [PATCH 07/16] use topk_transform_sparse --- .../ops/attention/dsv4/candidate_table.py | 15 ++++++++------- .../attention/dsv4/v41_indexer/sparse_table.py | 18 ++++-------------- 2 files changed, 12 insertions(+), 21 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_table.py b/python/sglang/kernels/ops/attention/dsv4/candidate_table.py index bdcca9ebf3ca..09394abf5bfa 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_table.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_table.py @@ -150,14 +150,15 @@ def sparse_logits( 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 every row of the sparse + """Top-``k`` (``k = out_indices.shape[1]``) of every row of the sparse ``logits`` (bf16 ``[rows, topk_blocks * 8]``) within its first ``valid_lens[b]`` - columns, written as pool slots, ``-1`` where a row has fewer than ``k`` valid - columns, in no particular order; ``phys_blocks`` ``[rows, topk_blocks]`` int32 - holds the published blocks as pool slots / 8.""" + columns, ``-1`` where a row has fewer than ``k`` valid columns, in no + particular order. Column ``j`` is written as ``blocks[b, j // 8] * 8 + j % 8`` + with ``blocks`` ``[rows, topk_blocks]`` int32: the published blocks as pool + slots / 8 for pool slots, or as logical block ids for compressed positions.""" 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/dsv4/v41_indexer/sparse_table.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py index 59d3a76c2dae..5634867aa0ec 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py @@ -23,9 +23,9 @@ build_sparse_indexer_schedule, sort_candidate_blocks, sparse_logits, + topk_transform_sparse, ) from sglang.kernels.ops.attention.dsv4.topk import ( - topk_transform_bf16_small, topk_transform_paged_v2, topk_transform_ragged_v2, ) @@ -329,12 +329,8 @@ def consume_decode( published.blocks.shape[1], ) assert out.raw_indices is None - topk_transform_bf16_small( - logits, - published.valid_lens, - published.phys_blocks, - out.page_indices, - CANDIDATE_BLOCK_SIZE, + topk_transform_sparse( + logits, published.valid_lens, published.phys_blocks, out.page_indices ) def _get_prefill_data(self, inputs: PrefillInputs) -> Optional[DeepGEMMPrefillData]: @@ -470,11 +466,5 @@ def select_prefill_table( 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, - ) + topk_transform_sparse(logits, table.valid_lens, table.blocks, out_positions) out_positions.add_(torch.where(out_positions >= 0, data.request_starts[:, None], 0)) From b147983edadac8c85dcb069da0c400ec72d37281 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 17:25:26 -0700 Subject: [PATCH 08/16] trim docstrings Co-Authored-By: DarkSharpness <2040703891@qq.com> --- .../kernels/ops/attention/dsv4/candidate_table.py | 9 ++++----- .../kernels/ops/attention/dsv4/dense_prefill.py | 15 ++++++--------- python/sglang/kernels/ops/attention/dsv4/topk.py | 7 +++---- 3 files changed, 13 insertions(+), 18 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_table.py b/python/sglang/kernels/ops/attention/dsv4/candidate_table.py index bdcca9ebf3ca..cad7fd9f730d 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_table.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_table.py @@ -153,11 +153,10 @@ def topk_transform_sparse( phys_blocks: torch.Tensor, page_indices: torch.Tensor, ) -> None: - """Top-``k`` (``k = page_indices.shape[1]``) of every row of the sparse - ``logits`` (bf16 ``[rows, topk_blocks * 8]``) within its first ``valid_lens[b]`` - columns, written as pool slots, ``-1`` where a row has fewer than ``k`` valid - columns, in no particular order; ``phys_blocks`` ``[rows, topk_blocks]`` int32 - holds the published blocks as pool slots / 8.""" + """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.""" topk_transform_bf16_small( logits, valid_lens, phys_blocks, page_indices, CANDIDATE_BLOCK_SIZE ) diff --git a/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py b/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py index 9b3fb2bf0731..17bfb15d264e 100644 --- a/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py +++ b/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py @@ -84,8 +84,8 @@ def dense_prefill_topk( def _rows_per_chunk(rows: int, width: int, *, heads: int, budget_bytes: int) -> int: - """Query rows per logits tile so one fp32 [rows, width] tile fits - ``budget_bytes``; the row count stays a multiple of the kernel's row alignment.""" + """Rows per fp32 logits tile within ``budget_bytes``, at the kernel's row + alignment.""" row_alignment = 128 // heads rows_per_chunk = mqa_logits_rows_per_chunk( num_rows=ceil_align(rows, row_alignment), @@ -108,13 +108,10 @@ def score_tiles( budget_bytes: int, width_align: int = 4, ) -> Iterator[tuple[slice, torch.Tensor]]: - """The dense scores of the chunk row tile by row tile, each fp32 tile within - ``budget_bytes``: - ``(rows, logits)`` with fp32 ``logits[i, j]`` the score of query row - ``rows.start + i`` against ``kv[starts + j]``, garbage past the row's - ``lengths``; the width is the largest of ``context_lengths`` (compressed - positions at each request's newest token) aligned to ``width_align`` - columns.""" + """``(rows, logits)`` per row tile, each fp32 tile within ``budget_bytes``: + ``logits[i, j]`` scores query row ``rows.start + i`` against ``kv[starts + j]``, + garbage past the row's ``lengths``; the width is ``max(context_lengths)`` + aligned to ``width_align``.""" from deep_gemm import fp8_fp4_mqa_logits rows = q[0].shape[0] diff --git a/python/sglang/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index bd26525738c5..ebacb43329aa 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/python/sglang/kernels/ops/attention/dsv4/topk.py @@ -131,10 +131,9 @@ def topk_transform_paged_torch( page_size: int, out_raw_indices: Optional[torch.Tensor] = None, ) -> None: - """The torch ``topk_transform_paged``: the top-``k`` (``k = out_page_indices.shape[1]``) - of each row of ``scores`` within its first ``seq_lens[b]`` columns, ascending, as - pool slots through ``page_tables`` and, when given, as positions; ``-1`` where a - row has fewer than ``k`` columns.""" + """The torch ``topk_transform_paged``: top-``k`` (``k = out_page_indices.shape[1]``) + of each row within its first ``seq_lens[b]`` columns, ascending, as pool slots + through ``page_tables`` and as positions when given; ``-1`` padded.""" topk = out_page_indices.shape[1] columns = torch.arange(scores.shape[1], device=seq_lens.device) lens_c = seq_lens.unsqueeze(-1) From 7f31f7bcc58cb267b3cb3975c3cf21f644d126f5 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 17:42:14 -0700 Subject: [PATCH 09/16] regroup index ops by stage Co-Authored-By: DarkSharpness <2040703891@qq.com> --- .../ops/attention/dsv4/candidate_blocks.py | 52 ++++- .../ops/attention/dsv4/candidate_table.py | 83 +------- .../ops/attention/dsv4/dense_prefill.py | 195 ------------------ .../ops/attention/dsv4/index_logits.py | 126 +++++++++++ .../sglang/kernels/ops/attention/dsv4/topk.py | 16 ++ .../layers/attention/deepseek_v4_backend.py | 4 +- .../dsv4/candidate_indexer_deep_gemm.py | 12 +- .../attention/dsv4/dense_prefill_indexer.py | 137 +++++++++++- .../srt/layers/attention/dsv4/indexer.py | 27 --- .../attention/test_dense_prefill_indexer.py | 45 ++-- 10 files changed, 355 insertions(+), 342 deletions(-) delete mode 100644 python/sglang/kernels/ops/attention/dsv4/dense_prefill.py create mode 100644 python/sglang/kernels/ops/attention/dsv4/index_logits.py diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py index 5c422c054dee..b0c378923bc5 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py @@ -1,5 +1,5 @@ -"""Candidate blocks of the two-level indexer: the per-row block counts and -sparse-row lengths, the JIT block top-k of the sparse table, and the torch +"""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.""" from typing import Optional, Union @@ -9,10 +9,16 @@ import triton import triton.language as tl -from sglang.kernels.jit.utils import is_arch_support_pdl +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) -from .candidate_table import CANDIDATE_BLOCK_SIZE, amax8_varlen +from .candidate_table import CANDIDATE_BLOCK_SIZE from .topk import plan_topk_v2, topk_transform_paged_v2 +from .utils import make_name @triton.jit @@ -72,6 +78,44 @@ def candidate_row_lens( return nblocks, valid +@cache_once +def _jit_block_amax_module(): + args = make_cpp_args(is_arch_support_pdl()) + return load_jit( + make_name("block_amax"), + *args, + cuda_files=["deepseek_v4/block_amax.cuh"], + cuda_wrappers=[("amax8_varlen", f"BlockAmaxKernel<{args}>::amax8_varlen")], + ) + + +def amax8_varlen( + scores: torch.Tensor, + seq_lens: torch.Tensor, + topk: int = 0, + *, + max_seqlen: int = 0, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Level-one keys of the two-level indexer: ``out[b, i]`` is the max of + ``scores[b, 8 i : 8 i + 8]`` for ``i < ceil(seq_lens[b] / 8)``, the last of + them ``+inf`` (the newest block is always selected), nothing written past + that count. Rows with at most ``topk`` blocks are skipped (every block is + selected anyway); ``topk=0`` never skips. ``out`` is allocated as + ``[rows, ceil(max_seqlen / 8)]`` when not given, ``max_seqlen`` defaulting to + the width of ``scores``; every ``seq_lens[b]`` must fit in ``8 * out.shape[1]``. + fp32 only for now; ``scores`` rows must be 32-byte aligned (stride a multiple + of 8). Returns ``out``. + """ + if out is None: + num_tokens, max_len = scores.shape + if max_seqlen == 0: + max_seqlen = max_len + out = scores.new_empty(num_tokens, (max_seqlen + 7) // 8) + _jit_block_amax_module().amax8_varlen(scores, seq_lens, out, topk) + return out + + def amax_topk_blocks( logits: torch.Tensor, seq_lens: torch.Tensor, diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_table.py b/python/sglang/kernels/ops/attention/dsv4/candidate_table.py index cad7fd9f730d..08871657182d 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_table.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_table.py @@ -1,5 +1,5 @@ -"""Candidate-table kernels of the two-level indexer: level-one block keys and -the sorted, page-transformed block table.""" +"""The sparse table of the two-level indexer from a row's selected blocks: the +sorted, page-transformed block table and DeepGEMM's schedule for it.""" from __future__ import annotations @@ -14,7 +14,6 @@ make_cpp_args, ) -from .topk import topk_transform_bf16_small from .utils import make_name if TYPE_CHECKING: @@ -23,44 +22,6 @@ CANDIDATE_BLOCK_SIZE = 8 # positions per block; DeepGEMM accepts 8 or 16 -@cache_once -def _jit_block_amax_module(): - args = make_cpp_args(is_arch_support_pdl()) - return load_jit( - make_name("block_amax"), - *args, - cuda_files=["deepseek_v4/block_amax.cuh"], - cuda_wrappers=[("amax8_varlen", f"BlockAmaxKernel<{args}>::amax8_varlen")], - ) - - -def amax8_varlen( - scores: torch.Tensor, - seq_lens: torch.Tensor, - topk: int = 0, - *, - max_seqlen: int = 0, - out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """Level-one keys of the two-level indexer: ``out[b, i]`` is the max of - ``scores[b, 8 i : 8 i + 8]`` for ``i < ceil(seq_lens[b] / 8)``, the last of - them ``+inf`` (the newest block is always selected), nothing written past - that count. Rows with at most ``topk`` blocks are skipped (every block is - selected anyway); ``topk=0`` never skips. ``out`` is allocated as - ``[rows, ceil(max_seqlen / 8)]`` when not given, ``max_seqlen`` defaulting to - the width of ``scores``; every ``seq_lens[b]`` must fit in ``8 * out.shape[1]``. - fp32 only for now; ``scores`` rows must be 32-byte aligned (stride a multiple - of 8). Returns ``out``. - """ - if out is None: - num_tokens, max_len = scores.shape - if max_seqlen == 0: - max_seqlen = max_len - out = scores.new_empty(num_tokens, (max_seqlen + 7) // 8) - _jit_block_amax_module().amax8_varlen(scores, seq_lens, out, topk) - return out - - @cache_once def _jit_candidate_block_table_module(): args = make_cpp_args(is_arch_support_pdl()) @@ -120,43 +81,3 @@ def build_sparse_indexer_schedule( q_dtype, CANDIDATE_BLOCK_SIZE, ) - - -def sparse_logits( - q_fp4: torch.Tensor, - q_sf: torch.Tensor, - k_cache: torch.Tensor, - weights: torch.Tensor, - schedule: torch.Tensor, - topk_blocks: int, -) -> torch.Tensor: - """bf16 logits ``[rows, topk_blocks * 8]`` of the published blocks: ``q_fp4`` - ``[rows, 1, heads, 64]`` int8 with ``q_sf`` ``[rows, 1, heads]`` int32 (packed - ue8m0), ``k_cache`` ``[pages, page_size, 1, 68]`` uint8 whose page stride is - a multiple of 512 bytes, ``weights`` ``[rows, heads]`` bf16, ``schedule`` the - ``build_sparse_indexer_schedule`` of the ``topk_blocks`` published blocks.""" - import deep_gemm - - return deep_gemm.fp8_fp4_paged_sparse_mqa_logits( - (q_fp4, q_sf), - k_cache, - weights, - schedule, - topk_blocks, - CANDIDATE_BLOCK_SIZE, - ) - - -def topk_transform_sparse( - logits: torch.Tensor, - valid_lens: torch.Tensor, - phys_blocks: torch.Tensor, - page_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.""" - topk_transform_bf16_small( - logits, valid_lens, phys_blocks, page_indices, CANDIDATE_BLOCK_SIZE - ) diff --git a/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py b/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py deleted file mode 100644 index 17bfb15d264e..000000000000 --- a/python/sglang/kernels/ops/attention/dsv4/dense_prefill.py +++ /dev/null @@ -1,195 +0,0 @@ -"""The DeepGEMM dense prefill scores of the ratio-1/2 index layers, tile by -tile under a caller-given memory budget, and the top-k over them: plain, or -publishing / consuming candidate block ids per request.""" - -from __future__ import annotations - -from typing import Iterator - -import torch - -from sglang.srt.layers.attention.mqa_logits_utils import ( - mqa_logits_row_bytes, - mqa_logits_rows_per_chunk, -) -from sglang.srt.utils.common import ceil_align - -from .candidate_blocks import ( - candidate_block_mask, - mask_topk_scores, - select_candidate_block_ids, -) -from .topk import topk_transform_ragged_v2 - - -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: list[torch.Tensor] | None, - budget_bytes: int, -) -> tuple[torch.Tensor, list[torch.Tensor] | None]: - """``budget_bytes`` bounds one fp32 logits tile.""" - selected = torch.full( - (q[0].shape[0], topk), -1, dtype=torch.int32, device=weights.device - ) - published = [] 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.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 = _rows_per_chunk( - row, width, heads=q[0].shape[1], budget_bytes=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 _rows_per_chunk(rows: int, width: int, *, heads: int, budget_bytes: int) -> int: - """Rows per fp32 logits tile within ``budget_bytes``, at the kernel's row - alignment.""" - row_alignment = 128 // heads - rows_per_chunk = mqa_logits_rows_per_chunk( - num_rows=ceil_align(rows, row_alignment), - row_bytes=mqa_logits_row_bytes(width), - budget_bytes=budget_bytes, - ) - if rows_per_chunk is None: - return rows - return max(row_alignment, rows_per_chunk // row_alignment * row_alignment) - - -def score_tiles( - *, - q: tuple[torch.Tensor, torch.Tensor], - kv: tuple[torch.Tensor, torch.Tensor], - weights: torch.Tensor, - starts: torch.Tensor, - lengths: torch.Tensor, - context_lengths: list[int], - budget_bytes: int, - width_align: int = 4, -) -> Iterator[tuple[slice, torch.Tensor]]: - """``(rows, logits)`` per row tile, each fp32 tile within ``budget_bytes``: - ``logits[i, j]`` scores query row ``rows.start + i`` against ``kv[starts + j]``, - garbage past the row's ``lengths``; the width is ``max(context_lengths)`` - aligned to ``width_align``.""" - from deep_gemm import fp8_fp4_mqa_logits - - rows = q[0].shape[0] - width = ceil_align(max(context_lengths, default=0), width_align) - if rows == 0 or width == 0: - return - rows_per_chunk = _rows_per_chunk( - rows, width, heads=q[0].shape[1], budget_bytes=budget_bytes - ) - for offset in range(0, rows, rows_per_chunk): - tile = slice(offset, min(offset + rows_per_chunk, rows)) - tile_starts = starts[tile] - yield ( - tile, - fp8_fp4_mqa_logits( - (q[0][tile], q[1][tile]), - kv, - weights[tile], - tile_starts, - tile_starts + lengths[tile], - False, - width, - ), - ) - - -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: list[torch.Tensor] | None, - consume: list[torch.Tensor] | None, -) -> None: - from deep_gemm import fp8_fp4_mqa_logits - - 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][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][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/kernels/ops/attention/dsv4/index_logits.py b/python/sglang/kernels/ops/attention/dsv4/index_logits.py new file mode 100644 index 000000000000..7832d770c569 --- /dev/null +++ b/python/sglang/kernels/ops/attention/dsv4/index_logits.py @@ -0,0 +1,126 @@ +"""The index logits of the DeepSeek V4.1 ratio-1/2 index layers, by the layout of +the index K they read: the flattened K of a prefill chunk (tile by tile under a +memory budget), the paged pool, or the published blocks of a sparse table.""" + +from __future__ import annotations + +from typing import Iterator, Tuple + +import torch + +from sglang.srt.layers.attention.mqa_logits_utils import ( + mqa_logits_row_bytes, + mqa_logits_rows_per_chunk, +) +from sglang.srt.utils.common import ceil_align + +from .candidate_table import CANDIDATE_BLOCK_SIZE + + +def flat_index_logits_rows_per_tile( + rows: int, width: int, *, heads: int, budget_bytes: int +) -> int: + """Rows per fp32 logits tile within ``budget_bytes``, at the kernel's row + alignment.""" + row_alignment = 128 // heads + rows_per_chunk = mqa_logits_rows_per_chunk( + num_rows=ceil_align(rows, row_alignment), + row_bytes=mqa_logits_row_bytes(width), + budget_bytes=budget_bytes, + ) + if rows_per_chunk is None: + return rows + return max(row_alignment, rows_per_chunk // row_alignment * row_alignment) + + +def flat_index_logits_tiles( + *, + q: tuple[torch.Tensor, torch.Tensor], + kv: tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + starts: torch.Tensor, + lengths: torch.Tensor, + context_lengths: list[int], + budget_bytes: int, + width_align: int = 4, +) -> Iterator[tuple[slice, torch.Tensor]]: + """``(rows, logits)`` per row tile, each fp32 tile within ``budget_bytes``: + ``logits[i, j]`` scores query row ``rows.start + i`` against ``kv[starts + j]``, + garbage past the row's ``lengths``; the width is ``max(context_lengths)`` + aligned to ``width_align``.""" + from deep_gemm import fp8_fp4_mqa_logits + + rows = q[0].shape[0] + width = ceil_align(max(context_lengths, default=0), width_align) + if rows == 0 or width == 0: + return + rows_per_chunk = flat_index_logits_rows_per_tile( + rows, width, heads=q[0].shape[1], budget_bytes=budget_bytes + ) + for offset in range(0, rows, rows_per_chunk): + tile = slice(offset, min(offset + rows_per_chunk, rows)) + tile_starts = starts[tile] + yield ( + tile, + fp8_fp4_mqa_logits( + (q[0][tile], q[1][tile]), + kv, + weights[tile], + tile_starts, + tile_starts + lengths[tile], + False, + width, + ), + ) + + +def deep_gemm_fp4_paged_mqa_logits( + q_fp4: Tuple[torch.Tensor, torch.Tensor], + k_cache: torch.Tensor, + weights: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + deep_gemm_metadata, + max_seq_len: int, +) -> torch.Tensor: + """DeepGEMM paged fp4 logits; no hadamard, the reference does not apply one.""" + from deep_gemm import fp8_fp4_paged_mqa_logits + + sl = seq_lens.to(torch.int32) + if sl.dim() == 1: + sl = sl.unsqueeze(-1) + return fp8_fp4_paged_mqa_logits( + q_fp4, + k_cache, + weights, + sl, + page_table, + deep_gemm_metadata, + max_seq_len, + False, + ) + + +def sparse_logits( + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + k_cache: torch.Tensor, + weights: torch.Tensor, + schedule: torch.Tensor, + topk_blocks: int, +) -> torch.Tensor: + """bf16 logits ``[rows, topk_blocks * 8]`` of the published blocks: ``q_fp4`` + ``[rows, 1, heads, 64]`` int8 with ``q_sf`` ``[rows, 1, heads]`` int32 (packed + ue8m0), ``k_cache`` ``[pages, page_size, 1, 68]`` uint8 whose page stride is + a multiple of 512 bytes, ``weights`` ``[rows, heads]`` bf16, ``schedule`` the + ``build_sparse_indexer_schedule`` of the ``topk_blocks`` published blocks.""" + import deep_gemm + + return deep_gemm.fp8_fp4_paged_sparse_mqa_logits( + (q_fp4, q_sf), + k_cache, + weights, + schedule, + topk_blocks, + CANDIDATE_BLOCK_SIZE, + ) diff --git a/python/sglang/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index ebacb43329aa..1a57689230be 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/python/sglang/kernels/ops/attention/dsv4/topk.py @@ -13,6 +13,7 @@ ) from sglang.srt.utils import is_xpu +from .candidate_table import CANDIDATE_BLOCK_SIZE from .utils import make_name @@ -325,3 +326,18 @@ def topk_transform_packed_v2( page_size, row_to_batch, ) + + +def topk_transform_sparse( + logits: torch.Tensor, + valid_lens: torch.Tensor, + phys_blocks: torch.Tensor, + page_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.""" + topk_transform_bf16_small( + logits, valid_lens, phys_blocks, page_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 1b6eda2f72c3..4a7255c494dd 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -34,6 +34,9 @@ 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, @@ -83,7 +86,6 @@ ) from sglang.srt.layers.attention.dsv4.indexer import ( C4IndexerBackendMixin, - deep_gemm_fp4_paged_mqa_logits, topk_transform_paged_from_metadata, ) from sglang.srt.layers.attention.dsv4.metadata import ( 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 index e31a317c3487..2cb85a5bb93a 100644 --- a/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py +++ b/python/sglang/srt/layers/attention/dsv4/candidate_indexer_deep_gemm.py @@ -6,22 +6,25 @@ 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, - amax8_varlen, 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, - topk_transform_sparse, ) -from sglang.kernels.ops.attention.dsv4.dense_prefill import score_tiles 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, @@ -35,7 +38,6 @@ DenseCandidateIndexer, ) from sglang.srt.layers.attention.dsv4.indexer import ( - deep_gemm_fp4_paged_mqa_logits, topk_transform_paged_from_metadata, ) @@ -293,7 +295,7 @@ def publish_prefill( 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 score_tiles( + for tile, logits in flat_index_logits_tiles( q=(inputs.q_fp4, inputs.q_sf), kv=inputs.kv, weights=inputs.weights, diff --git a/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py b/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py index 3ac4650b2ddc..bbde5ce9a4f8 100644 --- a/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/dense_prefill_indexer.py @@ -4,19 +4,88 @@ import torch -from sglang.kernels.ops.attention.dsv4.dense_prefill import dense_prefill_topk +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, @@ -37,13 +106,10 @@ def _dense_topk( candidate_topk_blocks=topk_blocks, candidate_block_size=block_size, publish_candidates=publish, - candidates=candidates.request_blocks if candidates is not None else None, - budget_bytes=_SCORE_BUDGET_BYTES, + candidates=candidates, ) out_positions.copy_(selected) - if published is None: - return None - return PrefillCandidateBlocks(request_blocks=published) + return published def plain_prefill_topk( @@ -107,3 +173,62 @@ def prefill_tail( ) -> 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/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 31b1230706ae..f0bac3d4108a 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -443,33 +443,6 @@ def topk_transform_flashinfer_fused( ) -def deep_gemm_fp4_paged_mqa_logits( - q_fp4: Tuple[torch.Tensor, torch.Tensor], - k_cache: torch.Tensor, - weights: torch.Tensor, - seq_lens: torch.Tensor, - page_table: torch.Tensor, - deep_gemm_metadata, - max_seq_len: int, -) -> torch.Tensor: - """DeepGEMM paged fp4 logits; no hadamard, the reference does not apply one.""" - from deep_gemm import fp8_fp4_paged_mqa_logits - - sl = seq_lens.to(torch.int32) - if sl.dim() == 1: - sl = sl.unsqueeze(-1) - return fp8_fp4_paged_mqa_logits( - q_fp4, - k_cache, - weights, - sl, - page_table, - deep_gemm_metadata, - max_seq_len, - False, - ) - - def topk_transform_paged_from_metadata( logits: torch.Tensor, metadata, 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 957fb65d9203..62c1ff6ff06b 100644 --- a/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py +++ b/test/registered/kernels/ops/attention/test_dense_prefill_indexer.py @@ -3,8 +3,8 @@ import torch -from sglang.kernels.ops.attention.dsv4 import dense_prefill 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.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -45,7 +45,6 @@ def make_inputs(request_lengths, ratio=1, zero_queries=False, seed=17): topk=512, candidate_topk_blocks=2, candidate_block_size=8, - budget_bytes=2 << 30, ) @@ -121,17 +120,16 @@ def test_ragged_source_consumer_and_replay(self): consumer_inputs["kv"] = inputs["kv"] consumer_inputs["candidate_topk_blocks"] = 128 consumer_scores = dense_scores(consumer_inputs) - with ( - patch.dict(inputs, budget_bytes=128 << 10), - patch.dict(consumer_inputs, budget_bytes=128 << 10), + with patch.object( + dense_prefill_indexer, "_SCORE_BUDGET_BYTES", 128 << 10 ): - selected, candidates = dense_prefill.dense_prefill_topk( + selected, candidates = dense_prefill_indexer.dense_prefill_topk( **inputs, publish_candidates=True, candidates=None ) self.assert_topk(inputs, selected, scores) row = 0 for (queries, context), blocks in zip( - inputs["request_lengths"], candidates + inputs["request_lengths"], candidates.request_blocks ): local = scores[row : row + queries, :context] if queries and context: @@ -175,7 +173,7 @@ def test_ragged_source_consumer_and_replay(self): torch.isfinite(consumer_scores[-1]).sum().item(), inputs["lengths"][-1].item(), ) - selected, published = dense_prefill.dense_prefill_topk( + selected, published = dense_prefill_indexer.dense_prefill_topk( **consumer_inputs, publish_candidates=False, candidates=candidates, @@ -203,13 +201,10 @@ def test_ragged_source_consumer_and_replay(self): ) ), ) - selected, _ = dense_prefill.dense_prefill_topk( + selected, _ = dense_prefill_indexer.dense_prefill_topk( **tail_inputs, publish_candidates=False, - candidates=[ - b[b.shape[0] - t :] - for b, t in zip(candidates, tail_lengths) - ], + candidates=candidates.tail(tail_lengths), ) self.assert_topk(tail_inputs, selected, consumer_scores[rows]) @@ -218,13 +213,13 @@ 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.dense_prefill_topk( + selected, candidates = dense_prefill_indexer.dense_prefill_topk( **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], + [b.shape[0] for b in candidates.request_blocks], [q for q, _ in request_lengths], ) @@ -237,13 +232,15 @@ def test_empty_queries_or_context(self): ): with self.subTest(request_lengths=request_lengths): inputs = make_inputs(request_lengths) - selected, candidates = dense_prefill.dense_prefill_topk( + selected, candidates = dense_prefill_indexer.dense_prefill_topk( **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], block_shapes) + self.assertEqual( + [tuple(b.shape) for b in candidates.request_blocks], block_shapes + ) def test_score_budget_includes_allocation_padding(self): from deep_gemm import fp8_fp4_mqa_logits @@ -261,10 +258,10 @@ def checked_logits(*args, **kwargs): return logits with ( - patch.dict(inputs, budget_bytes=budget), + patch.object(dense_prefill_indexer, "_SCORE_BUDGET_BYTES", budget), patch("deep_gemm.fp8_fp4_mqa_logits", new=checked_logits), ): - selected, _ = dense_prefill.dense_prefill_topk( + selected, _ = dense_prefill_indexer.dense_prefill_topk( **inputs, publish_candidates=True, candidates=None ) self.assertTrue(allocations) @@ -279,7 +276,7 @@ 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.dense_prefill_topk( + selected, candidates = dense_prefill_indexer.dense_prefill_topk( **inputs, publish_candidates=True, candidates=None ) torch.cuda.synchronize() @@ -287,9 +284,11 @@ def test_score_memory_is_bounded(self): torch.cuda.max_memory_allocated() - baseline, limit_gib << 30 ) self.assertEqual(tuple(selected.shape), (16384, 512)) - self.assertEqual(tuple(candidates[0].shape), (16384, 2048)) + self.assertEqual( + tuple(candidates.request_blocks[0].shape), (16384, 2048) + ) del selected - selected, published = dense_prefill.dense_prefill_topk( + selected, published = dense_prefill_indexer.dense_prefill_topk( **inputs, publish_candidates=False, candidates=candidates ) self.assertIsNone(published) @@ -298,7 +297,7 @@ 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.dense_prefill_topk( + selected, published = dense_prefill_indexer.dense_prefill_topk( **inputs, publish_candidates=False, candidates=candidates ) torch.cuda.synchronize() From dc4e7cb10e03a7fa1ee9a3122eb710a7ec9d70d8 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 18:40:44 -0700 Subject: [PATCH 10/16] restore record_stream for the published table --- .../srt/layers/attention/dsv4/v41_indexer/sparse_table.py | 3 +++ 1 file changed, 3 insertions(+) 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 index 360dd35d68c3..a329669c4644 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py @@ -231,6 +231,9 @@ def publish_decode(self, inputs: DecodeInputs, out: Selection): 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( From 2e650be35e655c8087b4436a75bcaf8921244051 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 18:48:10 -0700 Subject: [PATCH 11/16] rename plain indexer to full_topk --- .../srt/layers/attention/deepseek_v4_backend.py | 14 +++++++------- .../layers/attention/dsv4/v41_indexer/__init__.py | 15 ++++++++------- .../dsv4/v41_indexer/{plain.py => full_topk.py} | 4 ++-- .../unit/layers/test_dsv4_nonpaged_indexer.py | 4 ++-- 4 files changed, 19 insertions(+), 18 deletions(-) rename python/sglang/srt/layers/attention/dsv4/v41_indexer/{plain.py => full_topk.py} (98%) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 73b5134c954b..342182be7926 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -86,7 +86,7 @@ has_dense_fp4_indexer, is_sm100_or_newer, make_candidate_indexer, - make_plain_indexer, + make_full_topk_indexer, ) from sglang.srt.layers.attention.verify_mask import ( VerifyMask, @@ -1132,7 +1132,7 @@ def __init__( candidate_topk_blocks=getattr(cfg, "candidate_topk_blocks", 0), candidate_block_size=getattr(cfg, "candidate_block_size", 0), ) - self.plain_indexer = make_plain_indexer( + 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] @@ -3081,8 +3081,8 @@ def _low_ratio_index_topk( *, rows_per_request=None, # NOTE: only used in CP ) -> None: - """Select this index layer's compressed positions: a plain layer on the - plain indexer, a candidate source or consumer on the candidate backends.""" + """Select this index layer's compressed positions: a layer outside the + candidate scheme on the full top-k, a source or consumer on its backend.""" is_source = layer.indexer.is_candidate_source is_consumer = layer.indexer.uses_candidates ratio = layer.compress_ratio @@ -3109,7 +3109,7 @@ def _low_ratio_index_topk( elif is_consumer: self.decode_candidates.consume_decode(inputs, published, out) else: - self.plain_indexer.topk_decode(inputs, out) + 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 @@ -3120,7 +3120,7 @@ def _low_ratio_index_topk( elif is_consumer: self.prefill_candidates.consume_prefill(inputs, published, out) else: - self.plain_indexer.topk_prefill(inputs, out) + self.full_topk_indexer.topk_prefill(inputs, out) def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> None: ratio = layer.compress_ratio @@ -3146,7 +3146,7 @@ def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> Non paged_metadata=metadata, ) out = self._get_low_ratio_selection(ratio) - self.plain_indexer.topk_prefill_captured(inputs, out) + self.full_topk_indexer.topk_prefill_captured(inputs, out) def _get_low_ratio_selection(self, compress_ratio: int) -> Selection: core = self.forward_metadata.core_metadata diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py index b2bd85f9f42b..1b6b592ad0fc 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py @@ -1,6 +1,7 @@ """The top-k selection of the DeepSeek V4.1 ratio-1/2 index layers. -The plain index layers go through the ``PlainIndexer`` of ``make_plain_indexer``; a +The index layers outside the candidate scheme go through the ``FullTopKIndexer`` +of ``make_full_topk_indexer``; a candidate source and its consumers go through the prefill / decode backends of ``make_candidate_indexer``: a DeepGEMM sparse table (``sparse_table``) or dense block ids (``dense_blocks``). All three share only ``scoring``. Nothing else in @@ -14,7 +15,7 @@ import torch -from .plain import PlainIndexer +from .full_topk import FullTopKIndexer from .types import ( CandidateMetadata, CapturedPrefillInputs, @@ -33,14 +34,14 @@ "CapturedPrefillInputs", "DecodeCandidates", "DecodeInputs", - "PlainIndexer", + "FullTopKIndexer", "PrefillCandidates", "PrefillInputs", "Selection", "has_dense_fp4_indexer", "is_sm100_or_newer", "make_candidate_indexer", - "make_plain_indexer", + "make_full_topk_indexer", ] @@ -78,12 +79,12 @@ def _use_deep_gemm_prefill() -> bool: ) -def make_plain_indexer( +def make_full_topk_indexer( *, token_to_kv_pool: DeepSeekV4TokenToKVPool, req_to_token: torch.Tensor, -) -> PlainIndexer: - return PlainIndexer( +) -> 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(), diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/plain.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py similarity index 98% rename from python/sglang/srt/layers/attention/dsv4/v41_indexer/plain.py rename to python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py index 77efaa8a6bc8..2d6aee791e2e 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/plain.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py @@ -1,4 +1,4 @@ -"""The plain top-k of the DeepSeek V4.1 ratio-1/2 index layers outside the +"""The full top-k of the DeepSeek V4.1 ratio-1/2 index layers outside the candidate scheme: score every visible compressed position, keep the best. DeepGEMM on SM100; elsewhere the torch path, whose decode runs the fp4 decode @@ -40,7 +40,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool -class PlainIndexer: +class FullTopKIndexer: def __init__( self, *, diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py index 253140f412a9..d9399f2195c8 100644 --- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py +++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py @@ -1032,10 +1032,10 @@ def test_candidate_publisher_takes_one_schedule_per_call(self): 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 plain as mod + from sglang.srt.layers.attention.dsv4.v41_indexer import full_topk as mod inputs, data = self._inputs() - indexer = object.__new__(mod.PlainIndexer) + indexer = object.__new__(mod.FullTopKIndexer) indexer.token_to_kv_pool = None indexer.use_deep_gemm_decode = True deep_gemm, topk = self._mocks() From 81aca0f71d3202c84b5367fc418e7f1496f92579 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 19:03:45 -0700 Subject: [PATCH 12/16] trim comments --- .../ops/attention/dsv4/candidate_blocks.py | 28 +++++------- .../layers/attention/deepseek_v4_backend.py | 5 +-- .../attention/dsv4/v41_indexer/__init__.py | 17 ++----- .../dsv4/v41_indexer/dense_blocks.py | 16 +------ .../attention/dsv4/v41_indexer/full_topk.py | 9 +--- .../attention/dsv4/v41_indexer/scoring.py | 45 +++++-------------- .../dsv4/v41_indexer/sparse_table.py | 34 +++++--------- .../attention/dsv4/v41_indexer/types.py | 15 +++---- .../runner_utils/capture_mode.py | 5 +-- 9 files changed, 47 insertions(+), 127 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py index 9bb21910fd64..17112dc88403 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py @@ -1,10 +1,7 @@ -"""Candidate blocks of the two-level indexer: a candidate source keeps whole -blocks of ``block_size`` compressed positions, its newest block always, and a -consumer selects among the kept blocks only. - -The per-row block counts and sparse-row lengths, the block keys, the block -selection (the JIT chain on SM100, torch elsewhere), and the torch top-k among -chosen blocks.""" +"""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 @@ -152,11 +149,10 @@ def select_candidate_block_ids( topk_blocks: int, block_size: int, ) -> torch.Tensor: - """The torch block selection: per row the ids of the ``topk_blocks`` blocks - with the largest maximum score, the block holding position - ``compress_lens - 1`` always kept; int32 ``[rows, min(topk_blocks, blocks)]``, - unordered, ``-1`` where a row has fewer finite blocks. ``logits`` must already - be ``-inf`` past each row's causal length.""" + """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 @@ -177,11 +173,9 @@ def topk_among_blocks( k: int, block_size: int, ) -> torch.Tensor: - """The torch top-``k`` of each row among its chosen blocks: ``scores`` - ``[rows, width]`` indexed by position, ``lens`` ``[rows]`` the causal length, - ``blocks`` ``[rows, n]`` block ids with ``-1`` padding. Returns int64 - ``[rows, k]`` positions, unordered, ``-1`` where a row has fewer than ``k`` - finite candidates; positions at or past ``lens`` never count.""" + """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 if width == 0: return torch.full((rows, k), -1, dtype=torch.int64, device=scores.device) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 342182be7926..ce8d926134ec 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -3081,8 +3081,6 @@ def _low_ratio_index_topk( *, rows_per_request=None, # NOTE: only used in CP ) -> None: - """Select this index layer's compressed positions: a layer outside the - candidate scheme on the full top-k, a source or consumer on its backend.""" is_source = layer.indexer.is_candidate_source is_consumer = layer.indexer.uses_candidates ratio = layer.compress_ratio @@ -3132,8 +3130,7 @@ def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> Non ) assert metadata is not None, f"no prefill graph indexer metadata for {ratio = }" 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. + # 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 diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py index 1b6b592ad0fc..60f6fc690561 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/__init__.py @@ -1,12 +1,6 @@ -"""The top-k selection of the DeepSeek V4.1 ratio-1/2 index layers. - -The index layers outside the candidate scheme go through the ``FullTopKIndexer`` -of ``make_full_topk_indexer``; a -candidate source and its consumers go through the prefill / decode backends of -``make_candidate_indexer``: a DeepGEMM sparse table (``sparse_table``) or dense -block ids (``dense_blocks``). All three share only ``scoring``. Nothing else in -this package is meant to be named from outside it. -""" +"""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 @@ -64,8 +58,7 @@ def has_dense_fp4_indexer() -> bool: @functools.cache def _use_deep_gemm_prefill() -> bool: - """Whether an eager prefill scores on DeepGEMM rather than torch. A CP rank - scores rank-local rows, which only DeepGEMM serves, so CP ignores the env.""" + # 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 @@ -100,8 +93,6 @@ def make_candidate_indexer( candidate_topk_blocks: int, candidate_block_size: int, ) -> Tuple[PrefillCandidates, DecodeCandidates]: - """(prefill, decode): the sparse table where SM100 can page the rows, dense - block ids elsewhere.""" from sglang.srt.runtime_context import get_parallel from .dense_blocks import DenseBlocksBackend 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 index c419716a3da5..e45ebf7a0884 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py @@ -1,11 +1,5 @@ -"""The candidate scheme carried by dense block ids. - -A source scores its whole context, keeps its own top-k and publishes the ids of -its best blocks; a consumer scores its whole context too, gathers the scores of -the published blocks only and takes its top-k among them. Scores come from -DeepGEMM's flattened-K logits or from torch, per backend configuration; the -block math is the torch one in ``candidate_blocks``. -""" +"""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 @@ -238,8 +232,6 @@ def _publish_prefill_blocks( topk_blocks: int, block_size: int, ) -> Tuple[torch.Tensor, torch.Tensor]: - """The source layer's own plain top-k, plus the ids of each row's best blocks: - ``[rows, min(topk_blocks, ceil(max lc / block_size))]`` int32, -1 padded.""" from sglang.kernels.ops.attention.dsv4 import topk_transform_ragged_v2 selected = data.empty_selection(topk) @@ -300,8 +292,6 @@ def _consume_prefill_blocks( blocks: torch.Tensor, block_size: int, ) -> torch.Tensor: - """A consumer's top-k among the published blocks of its dense scores, as - flattened-K columns, ``-1`` padded, unordered.""" selected = data.empty_selection(topk) for tile, logits in score_tiles(data, kv, width_align=4): _consume_tile_blocks( @@ -343,8 +333,6 @@ def _consume_tile_blocks( def _requests_in_tile(data: DeepGEMMPrefillData, tile: slice): - """``(rows, lc)`` of each request with a visible position and rows in - ``tile``: its rows within the tile and its context length.""" start = 0 for num_rows, lc in zip(data.rows_per_request, data.lens_per_request): end = start + num_rows 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 index 2d6aee791e2e..0bcd1b45ba67 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py @@ -1,9 +1,5 @@ -"""The full top-k of the DeepSeek V4.1 ratio-1/2 index layers outside the -candidate scheme: score every visible compressed position, keep the best. - -DeepGEMM on SM100; elsewhere the torch path, whose decode runs the fp4 decode -logits kernel (the Hopper indexer). -""" +"""The full top-k of the ratio-1/2 index layers outside the candidate scheme: +DeepGEMM on SM100, torch elsewhere.""" from __future__ import annotations @@ -96,7 +92,6 @@ 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): - # An eager forward whose dense logits are bounded by row 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()): diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py index 21d232d502c9..4ae6d471fd01 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py @@ -1,13 +1,6 @@ -"""Index scores shared by the dense layers and both candidate schemes. - -DeepGEMM (SM100): on prefill the chunk's index-K slots, fp4 queries and head -weights, the tiled ``fp8_fp4_mqa_logits`` scores over the flattened index K, the -plain top-k, and the write of a selection back into the pool-slot layout -attention reads; on decode the fp4 queries and the paged pool. - -Torch (elsewhere): bf16 scores per request on prefill, the fp4 decode logits -kernel on decode, and the writes of a ``torch.topk`` back into pool slots. -""" +"""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 @@ -42,9 +35,8 @@ class DeepGEMMDecodeData(msgspec.Struct, frozen=True): class DeepGEMMPrefillData(msgspec.Struct, frozen=True): - """One index layer's operands on a prefill chunk; the rows of one request are - consecutive, and a selection is a flattened-K column, request start plus - compressed position.""" + """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 @@ -66,7 +58,6 @@ def empty_selection(self, topk: int, init: bool = True) -> torch.Tensor: return out def write_selection(self, selected: torch.Tensor, out: Selection) -> None: - # write_prefill_selection(out=out, data=self, selected=selected) num_tokens, topk = selected.shape unselected = torch.iinfo(torch.int32).max selected = selected.masked_fill(selected < 0, unselected).sort(dim=-1).values @@ -81,8 +72,6 @@ def write_selection(self, selected: torch.Tensor, out: Selection) -> None: def quantize_index_q(q: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """``q`` ``[rows, heads, 128]`` bf16 to fp4 ``[rows, heads, 64]`` int8 and its - scales ``[rows, heads]`` int32 (packed ue8m0).""" from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( quantize_fp4_indexer_tensor, ) @@ -111,8 +100,6 @@ def get_flat_index_k( token_to_kv_pool: DeepSeekV4TokenToKVPool, layer_id: int, ) -> Tuple[torch.Tensor, torch.Tensor]: - """The chunk's index K gathered in flattened-K column order: int8 - ``[columns, 64]`` and int32 ``[columns]``; a copy, for dense scoring only.""" return token_to_kv_pool.get_low_ratio_index_k_fp4(layer_id, data.k_slots) @@ -163,8 +150,7 @@ def get_deep_gemm_decode_data( 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, which is - the same for every index layer of the forward.""" + """None when the chunk has no row or no visible compressed position.""" ratio = inputs.compress_ratio indexer = inputs.indexer pos = inputs.positions @@ -175,8 +161,7 @@ def get_deep_gemm_prefill_data( 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 at its newest token; the - # per-token count (pos + 1) // ratio bounds each row below. + # 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 @@ -220,11 +205,6 @@ def score_tiles( *, width_align: int, ) -> Generator[Tuple[slice, torch.Tensor], None, None]: - """The dense scores of ``data`` row tile by row tile under the budget: - ``(rows, logits)`` with fp32 ``logits[i, j]`` the score of query row - ``rows.start + i`` against ``kv[request_starts + j]``, garbage past the row's - ``compress_lens``; the width is the batch's largest context aligned to - ``width_align`` columns.""" yield from flat_index_logits_tiles( q=(data.q_fp4, data.q_sf), kv=kv, @@ -243,8 +223,7 @@ def dense_prefill_topk( *, topk: int, ) -> torch.Tensor: - """The plain top-``topk`` of each row as flattened-K columns, ``-1`` padded, - unordered.""" + """Flattened-K columns, ``-1`` padded, unordered.""" from sglang.kernels.ops.attention.dsv4 import topk_transform_ragged_v2 selected = data.empty_selection(topk) @@ -268,8 +247,6 @@ def dense_prefill_topk( class RequestScores(msgspec.Struct, frozen=True): - """One request of a prefill chunk, scored row chunk by row chunk.""" - index: int # position among the chunk's requests, empty ones included lc: int # compressed positions visible at its newest row; 0 scores nothing k: int @@ -301,8 +278,7 @@ def prefill_requests( token_to_kv_pool: DeepSeekV4TokenToKVPool, req_to_token: torch.Tensor, ) -> Iterator[tuple[RequestScores, Iterator[ChunkScores]]]: - """Each request of the chunk with its row chunks' scores; an empty request - comes with no chunk.""" + """An empty request comes with no chunk.""" pool = token_to_kv_pool ratio = inputs.compress_ratio indexer = inputs.indexer @@ -345,8 +321,7 @@ def _score_chunks( request: RequestScores, ) -> Iterator[ChunkScores]: lc, j = request.lc, request.columns - # Every step 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). + # 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) 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 index a329669c4644..8f58fe35fa3e 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py @@ -1,9 +1,6 @@ -"""The candidate scheme carried by a DeepGEMM sparse table (SM100). - -A source publishes the block ids it kept plus the pool slots and the schedule -``fp8_fp4_paged_sparse_mqa_logits`` needs for them, and a consumer scores only -those blocks straight out of the index-K pool. -""" +"""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 @@ -61,9 +58,8 @@ class _SparseTable(CandidateMetadata, msgspec.Struct): 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 + # [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`) @@ -73,9 +69,8 @@ class _SparseTable(CandidateMetadata, msgspec.Struct): class _SparsePrefillTable(_SparseTable): - """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).""" + """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 @@ -162,7 +157,6 @@ def consume_prefill( out: Selection, ) -> None: out.reset() - # Scores only the published blocks, straight out of the pool. data = self._get_prefill_data(inputs) if data is None: return @@ -183,7 +177,6 @@ def consume_prefill( data.write_selection(selected=selected, out=out) def publish_decode(self, inputs: DecodeInputs, out: Selection): - """The source layer's own plain top-k, plus the table for the consumers.""" data = get_deep_gemm_decode_data(inputs, self.token_to_kv_pool) metadata = inputs.paged_metadata seq_lens = metadata.compressed_seq_lens.reshape(-1) @@ -251,9 +244,7 @@ def _publish_decode_chunked( seq_lens: torch.Tensor, out: Selection, ) -> _SparseTable: - """Publish an eager forward whose dense logits are bounded by row chunks. - It stays on the current stream so each chunk's full logits can be released - before the next; CUDA-graph metadata always carries one schedule.""" + """On the current stream, so each chunk's logits are released before the next.""" metadata = inputs.paged_metadata block_chunks = [] phys_block_chunks = [] @@ -401,10 +392,8 @@ def publish_prefill_table( topk_blocks: int, out_positions: torch.Tensor, ) -> _SparsePrefillTable: - """The source layer's own top-``k`` (``k = out_positions.shape[1]``) as - flattened-K columns, ``-1`` padded, unordered, and the chunk's table: both - from one pass over its dense scores, tile by tile. ``index_page_table`` is - ``[rows, index pages]`` at ``index_page_size`` slots.""" + """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) @@ -456,9 +445,6 @@ def select_prefill_table( k_cache: torch.Tensor, out_positions: torch.Tensor, ) -> None: - """A consumer's top-``k`` over the published blocks, in the layout - ``publish_prefill_table`` writes; ``k_cache`` is the layer's index-K pool, - ``[pages, page_size, 1, 68]`` uint8.""" rows, heads = data.q_sf.shape logits = sparse_logits( data.q_fp4.view(rows, 1, heads, 64), diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py index 67b21343d4de..8dee611dc624 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py @@ -1,5 +1,4 @@ -"""What the attention backend hands the V4.1 indexer backends, and what it gets -back.""" +"""The inputs and outputs of the V4.1 indexer backends.""" from __future__ import annotations @@ -18,8 +17,7 @@ def get_tail_row_indices( tail_rows_per_request: List[int], device: torch.device, ) -> torch.Tensor: - """int64 indices of each request's last ``tail_rows_per_request[b]`` rows, in - row order, copied without a host sync (pinned staging, non-blocking).""" + """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)) @@ -47,8 +45,7 @@ def reset(self) -> None: class PrefillInputs(msgspec.Struct, frozen=True, kw_only=True): - """An eager extend: the backend projects the queries itself and gathers the - index K it needs.""" + """An eager extend: the backend projects the queries and gathers the index K.""" indexer: DeepseekV41Indexer layer_id: int @@ -88,8 +85,7 @@ class DecodeInputs(msgspec.Struct, frozen=True, kw_only=True): class CapturedPrefillInputs(msgspec.Struct, frozen=True, kw_only=True): - """An extend under the prefill graph: the queries and head weights come - projected, and the index K is read paged, as on decode.""" + """An extend under the prefill graph: projected queries, paged index K.""" indexer: DeepseekV41Indexer layer_id: int @@ -102,8 +98,7 @@ class CapturedPrefillInputs(msgspec.Struct, frozen=True, kw_only=True): class PrefillCandidates(Protocol[Metadata]): - """The candidate scheme on prefill: what a publish carries is the backend's - own; the attention backend keeps it alive and hands it back unread.""" + """The attention backend keeps a publish alive and hands it back unread.""" def publish_prefill( self, 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 a95bba4d8962..ce4f08f1fb3b 100644 --- a/python/sglang/srt/model_executor/runner_utils/capture_mode.py +++ b/python/sglang/srt/model_executor/runner_utils/capture_mode.py @@ -80,9 +80,8 @@ def skip_low_ratio_indexer(compress_ratio: int) -> bool: def skip_candidate_filtering() -> bool: - """Whether the captured candidate variant serves only contexts inside the - candidate window, where every block is a candidate and the candidate layers - select like dense ones.""" + """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", From 67aa7ba55ce854d5e0e8f51cacfd9d4c7955561a Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 19:21:25 -0700 Subject: [PATCH 13/16] gather candidate blocks, not columns --- .../ops/attention/dsv4/candidate_blocks.py | 33 ++++++++++++------- .../dsv4/v41_indexer/dense_blocks.py | 5 +-- 2 files changed, 25 insertions(+), 13 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py index 17112dc88403..e0f505e55dd3 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py @@ -177,19 +177,30 @@ def topk_among_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 - if width == 0: + n = blocks.shape[1] + if width == 0 or n == 0: return torch.full((rows, k), -1, dtype=torch.int64, device=scores.device) - offsets = torch.arange(block_size, device=scores.device) - columns = (blocks.to(torch.int64)[:, :, None] * block_size + offsets).flatten(1) - valid = ( - (blocks >= 0).repeat_interleave(block_size, dim=1) - & (columns < lens.to(torch.int64)[:, None]) - & (columns < width) + if width % block_size: + # A block-aligned width lets the gather below index blocks, not columns; + # callers on the DeepGEMM path slice their tile to one so this is a no-op. + 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)) + ids = blocks.to(torch.int64).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 b is a real + # block and the position is causal: j < lens - b * block_size + offsets = torch.arange(block_size, device=scores.device, dtype=torch.int32) + room = lens.to(torch.int32)[:, None] - blocks.to(torch.int32) * block_size + valid = (blocks >= 0)[:, :, None] & (offsets[None, None, :] < room[:, :, None]) + candidates.masked_fill_(~valid, -torch.inf) + flat = candidates.flatten(1) + top = flat.topk(min(k, flat.shape[1]), dim=-1, sorted=False) + picked = blocks.to(torch.int64).gather(1, top.indices // block_size) * block_size + picked = (picked + top.indices % block_size).masked_fill( + ~(top.values > -torch.inf), -1 ) - candidates = scores.gather(1, columns.clamp(0, width - 1)) - candidates = candidates.masked_fill(~valid, -torch.inf) - top = candidates.topk(min(k, candidates.shape[1]), dim=-1, sorted=False) - picked = columns.gather(1, top.indices).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/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py index e45ebf7a0884..b85a4d757046 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py @@ -293,7 +293,7 @@ def _consume_prefill_blocks( block_size: int, ) -> torch.Tensor: selected = data.empty_selection(topk) - for tile, logits in score_tiles(data, kv, width_align=4): + for tile, logits in score_tiles(data, kv, width_align=block_size): _consume_tile_blocks( data=data, tile=tile, @@ -320,8 +320,9 @@ def _consume_tile_blocks( 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], + logits[rows, : -(-lc // block_size) * block_size], lens[rows], blocks[rows], out.shape[1], From 389900b759f3831840ff8979ee5d9772241e10b3 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 19:53:32 -0700 Subject: [PATCH 14/16] fewer ops in topk_among_blocks --- .../ops/attention/dsv4/candidate_blocks.py | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py index e0f505e55dd3..a033f9d9137c 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py @@ -186,21 +186,24 @@ def topk_among_blocks( 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)) - ids = blocks.to(torch.int64).clamp_(0, num_blocks - 1) + 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 b is a real - # block and the position is causal: j < lens - b * block_size - offsets = torch.arange(block_size, device=scores.device, dtype=torch.int32) - room = lens.to(torch.int32)[:, None] - blocks.to(torch.int32) * block_size - valid = (blocks >= 0)[:, :, None] & (offsets[None, None, :] < room[:, :, None]) - candidates.masked_fill_(~valid, -torch.inf) + # block and the position is causal: j < lens - b * block_size (a padded + # block is -1, so its room exceeds block_size and the sign test drops it) + 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) - picked = blocks.to(torch.int64).gather(1, top.indices // block_size) * block_size - picked = (picked + top.indices % block_size).masked_fill( - ~(top.values > -torch.inf), -1 + block_of_pick = top.indices // block_size + picked = blocks64.gather(1, block_of_pick) * block_size + ( + top.indices - block_of_pick * block_size ) + 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 From 0b82dfcb5fecf31ee4f0526a8c8e19c56577fc80 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 20:02:22 -0700 Subject: [PATCH 15/16] inline the backend forwarding helpers --- .../dsv4/v41_indexer/dense_blocks.py | 58 ++++++++++--------- .../attention/dsv4/v41_indexer/full_topk.py | 5 +- .../attention/dsv4/v41_indexer/scoring.py | 9 --- .../dsv4/v41_indexer/sparse_table.py | 14 ++--- 4 files changed, 36 insertions(+), 50 deletions(-) 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 index b85a4d757046..3cc44e2bfd25 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/dense_blocks.py @@ -17,7 +17,6 @@ DeepGEMMPrefillData, decode_scores, get_deep_gemm_prefill_data, - get_flat_index_k, prefill_requests, score_tiles, write_decode, @@ -85,7 +84,12 @@ def consume_prefill( self._torch_consume_prefill(inputs, published, out) def publish_decode(self, inputs: DecodeInputs, out: Selection): - d = self._decode_scores(inputs, out) + 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( @@ -107,7 +111,12 @@ def consume_decode( published: Optional[BlockIds], out: Selection, ) -> None: - d = self._decode_scores(inputs, out) + 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 @@ -126,7 +135,9 @@ def _deep_gemm_publish_prefill(self, inputs: PrefillInputs, out: Selection): return None selected, blocks = _publish_prefill_blocks( data=data, - kv=self._get_flat_index_k(inputs, 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, @@ -147,23 +158,25 @@ def _deep_gemm_consume_prefill( assert published is not None selected = _consume_prefill_blocks( data=data, - kv=self._get_flat_index_k(inputs, 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) - def _get_flat_index_k(self, inputs: PrefillInputs, data: DeepGEMMPrefillData): - return get_flat_index_k( - data=data, token_to_kv_pool=self.token_to_kv_pool, layer_id=inputs.layer_id - ) - # ---------- 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 self._prefill_requests(inputs, out): + 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( ( @@ -196,7 +209,12 @@ def _torch_consume_prefill( out: Selection, ) -> None: assert published is not None - for request, chunks in self._prefill_requests(inputs, out): + 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, @@ -207,22 +225,6 @@ def _torch_consume_prefill( ) write_prefill(out, request, chunk, idx.masked_fill(idx < 0, request.lc)) - def _prefill_requests(self, inputs: PrefillInputs, out: Selection): - return prefill_requests( - inputs=inputs, - out=out, - token_to_kv_pool=self.token_to_kv_pool, - req_to_token=self.req_to_token, - ) - - def _decode_scores(self, inputs: DecodeInputs, out: Selection): - return decode_scores( - inputs=inputs, - out=out, - token_to_kv_pool=self.token_to_kv_pool, - req_to_token=self.req_to_token, - ) - def _publish_prefill_blocks( *, 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 index 0bcd1b45ba67..a8c18b5a2c39 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py @@ -18,7 +18,6 @@ dense_prefill_topk, get_deep_gemm_decode_data, get_deep_gemm_prefill_data, - get_flat_index_k, get_index_k_cache, prefill_requests, quantize_index_q, @@ -70,8 +69,8 @@ def _deep_gemm_prefill(self, inputs: PrefillInputs, out: Selection) -> None: data = get_deep_gemm_prefill_data(inputs, self.req_to_token) if data is None: return - kv = get_flat_index_k( - data=data, token_to_kv_pool=self.token_to_kv_pool, layer_id=inputs.layer_id + 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) diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py index 4ae6d471fd01..9d43b0fb4bc6 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py @@ -94,15 +94,6 @@ def get_index_k_cache( return k_cache.view(k_cache.shape[0], page_size, 1, 68) -def get_flat_index_k( - *, - data: DeepGEMMPrefillData, - token_to_kv_pool: DeepSeekV4TokenToKVPool, - layer_id: int, -) -> Tuple[torch.Tensor, torch.Tensor]: - return token_to_kv_pool.get_low_ratio_index_k_fp4(layer_id, data.k_slots) - - def get_deep_gemm_decode_data( inputs: DecodeInputs, token_to_kv_pool: DeepSeekV4TokenToKVPool ) -> DeepGEMMDecodeData: 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 index 8f58fe35fa3e..e284db54d273 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py @@ -36,7 +36,6 @@ DeepGEMMPrefillData, get_deep_gemm_decode_data, get_deep_gemm_prefill_data, - get_flat_index_k, get_index_k_cache, score_tiles, ) @@ -123,7 +122,7 @@ def __init__( def publish_prefill(self, inputs: PrefillInputs, out: Selection): out.reset() - data = self._get_prefill_data(inputs) + 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) @@ -132,10 +131,8 @@ def publish_prefill(self, inputs: PrefillInputs, out: Selection): ) table = publish_prefill_table( data=data, - kv=get_flat_index_k( - data=data, - token_to_kv_pool=self.token_to_kv_pool, - layer_id=inputs.layer_id, + 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], @@ -157,7 +154,7 @@ def consume_prefill( out: Selection, ) -> None: out.reset() - data = self._get_prefill_data(inputs) + data = get_deep_gemm_prefill_data(inputs, self.req_to_token) if data is None: return assert published is not None @@ -327,9 +324,6 @@ def consume_decode( logits, published.valid_lens, published.phys_blocks, out.page_indices ) - def _get_prefill_data(self, inputs: PrefillInputs) -> Optional[DeepGEMMPrefillData]: - return get_deep_gemm_prefill_data(inputs, self.req_to_token) - 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() From d4a5b1c814b512e9339e1e377f0eb088c54ee701 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Fri, 25 Sep 2026 20:07:30 -0700 Subject: [PATCH 16/16] drop unused params and fields; captured prefill as a method --- .../ops/attention/dsv4/candidate_blocks.py | 7 +- .../attention/dsv4/v41_indexer/full_topk.py | 83 +++++++++---------- .../attention/dsv4/v41_indexer/scoring.py | 27 +++--- .../dsv4/v41_indexer/sparse_table.py | 10 +-- 4 files changed, 57 insertions(+), 70 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py index a033f9d9137c..ba5550f0adb6 100644 --- a/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py +++ b/python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py @@ -181,8 +181,7 @@ def topk_among_blocks( if width == 0 or n == 0: return torch.full((rows, k), -1, dtype=torch.int64, device=scores.device) if width % block_size: - # A block-aligned width lets the gather below index blocks, not columns; - # callers on the DeepGEMM path slice their tile to one so this is a no-op. + # 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)) @@ -190,9 +189,7 @@ def topk_among_blocks( 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 b is a real - # block and the position is causal: j < lens - b * block_size (a padded - # block is -1, so its room exceeds block_size and the sign test drops it) + # 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] 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 index a8c18b5a2c39..db9dadfcd3e8 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py @@ -56,7 +56,7 @@ def topk_prefill(self, inputs: PrefillInputs, out: Selection) -> None: self._torch_prefill(inputs, out) def topk_prefill_captured(self, inputs: CapturedPrefillInputs, out: Selection): - return _deep_gemm_prefill_captured(inputs, self.token_to_kv_pool, out) + self._deep_gemm_prefill_captured(inputs, out) def topk_decode(self, inputs: DecodeInputs, out: Selection) -> None: if self.use_deep_gemm_decode: @@ -142,46 +142,43 @@ def _torch_decode(self, inputs: DecodeInputs, out: Selection) -> None: idx = d.scores.topk(k, dim=-1, sorted=False).indices write_decode(out, d, idx) - -def _deep_gemm_prefill_captured( - inputs: CapturedPrefillInputs, - token_to_kv_pool: DeepSeekV4TokenToKVPool, - 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=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, + 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 index 9d43b0fb4bc6..1cd46108fa1d 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/scoring.py @@ -51,11 +51,10 @@ class DeepGEMMPrefillData(msgspec.Struct, frozen=True): def num_rows(self) -> int: return self.q_fp4.shape[0] - def empty_selection(self, topk: int, init: bool = True) -> torch.Tensor: - out = self.request_starts.new_empty((self.num_rows, topk), dtype=torch.int32) - if init: - out.fill_(-1) - return out + 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 @@ -238,13 +237,12 @@ def dense_prefill_topk( class RequestScores(msgspec.Struct, frozen=True): - index: int # position among the chunk's requests, empty ones included - lc: int # compressed positions visible at its newest row; 0 scores nothing + lc: int # compressed positions visible at its newest row k: int - columns: Optional[torch.Tensor] = None # arange(lc) - tok: Optional[torch.Tensor] = None # [rows_b] its query rows in the chunk - lens: Optional[torch.Tensor] = None # [rows_b] compressed positions each row sees - slots: Optional[torch.Tensor] = None # [lc] index-K pool slot of each position + 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): @@ -269,7 +267,7 @@ def prefill_requests( token_to_kv_pool: DeepSeekV4TokenToKVPool, req_to_token: torch.Tensor, ) -> Iterator[tuple[RequestScores, Iterator[ChunkScores]]]: - """An empty request comes with no chunk.""" + """Requests with nothing visible yet are skipped.""" pool = token_to_kv_pool ratio = inputs.compress_ratio indexer = inputs.indexer @@ -281,19 +279,18 @@ def prefill_requests( # A compressed position is visible once the query has passed its last token. compress_lens = (pos + 1) // ratio topk = indexer.index_topk - for b, r in enumerate(torch.unique_consecutive(req).tolist()): + 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: - yield RequestScores(index=b, lc=0, k=0), iter(()) 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( - index=b, lc=lc, k=min(topk, lc), columns=j, tok=tok, lens=lens, slots=slots + lc=lc, k=min(topk, lc), columns=j, tok=tok, lens=lens, slots=slots ) yield ( request, 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 index e284db54d273..c01d76f7be94 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/sparse_table.py @@ -77,7 +77,6 @@ class _SparsePrefillTable(_SparseTable): 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 - topk_blocks: int def tail(self, rows_per_request: List[int]) -> _SparsePrefillTable: rows = get_tail_row_indices( @@ -86,7 +85,7 @@ def tail(self, rows_per_request: List[int]) -> _SparsePrefillTable: device=self.blocks.device, ) compress_lens = self.compress_lens[rows].contiguous() - _, valid_lens = candidate_row_lens(compress_lens, self.topk_blocks) + _, valid_lens = candidate_row_lens(compress_lens, self.blocks.shape[1]) return _build_prefill_table( blocks=self.blocks[rows].contiguous(), compress_lens=compress_lens, @@ -96,7 +95,6 @@ def tail(self, rows_per_request: List[int]) -> _SparsePrefillTable: rows_per_request=rows_per_request, q_dtype=self.q_dtype, valid_lens=valid_lens, - topk_blocks=self.topk_blocks, ) @@ -351,7 +349,6 @@ def _build_prefill_table( rows_per_request: List[int], q_dtype: torch.dtype, valid_lens: torch.Tensor, - topk_blocks: int, ) -> _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) @@ -372,7 +369,6 @@ def _build_prefill_table( q_dtype=q_dtype, page_size=page_size, rows_per_request=rows_per_request, - topk_blocks=topk_blocks, ) @@ -392,6 +388,7 @@ def publish_prefill_table( 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] @@ -411,7 +408,7 @@ def publish_prefill_table( topk_transform_ragged_v2( keys, nblocks[tile], - out_offsets=torch.zeros(logits.shape[0], dtype=torch.int32, device=device), + out_offsets=zero_offsets[: logits.shape[0]], out_indices=blocks[tile], ) request_ids = torch.repeat_interleave( @@ -428,7 +425,6 @@ def publish_prefill_table( rows_per_request=data.rows_per_request, q_dtype=data.q_fp4.dtype, valid_lens=valid_lens, - topk_blocks=topk_blocks, )