diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 2f0c73443..b933601c4 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -32,6 +32,16 @@ from sglang.srt.layers.attention.dsa.utils import ( is_graph_dsa_split_op_surface, ) from sglang.srt.layers.attention.graph_variants import DSA_DENSE +from sglang.srt.layers.attention.mqa_logits_utils import ( + MQA_LOGITS_BYTES_PER_ELEM, + MQA_LOGITS_MAX_BYTES_ROCM, + MQA_LOGITS_STATIC_SKIP_ELEMS, + MQA_LOGITS_TOTAL_MEM_FRACTION, + mqa_logits_budget_bytes, + mqa_logits_free_mem_fraction, + mqa_logits_should_chunk, + mqa_logits_static_budget_bytes, +) from sglang.srt.layers.layernorm import LayerNorm, RMSNorm from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( is_in_breakable_cuda_graph, @@ -43,7 +53,6 @@ from sglang.srt.runtime_context import ( get_device, get_exec, get_parallel, - get_schedule, ) from sglang.srt.state_capturer.indexer_topk import ( maybe_capture_indexer_topk, @@ -52,7 +61,6 @@ from sglang.srt.utils import ( add_prefix, ceil_align, get_bool_env_var, - get_device_module, is_cuda, is_gfx95_supported, is_hip, @@ -202,16 +210,16 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: class Indexer(DSANPUIndexerMixin, BaseFusedOp): - _MQA_LOGITS_BYTES_PER_ELEM = 4 - _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000 - _MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3 - # aiter's fp8_mqa_logits only compiles below 2 GiB of logits (buffer_store). - _MQA_LOGITS_MAX_BYTES_ROCM = 2**31 - 1 + _MQA_LOGITS_BYTES_PER_ELEM = MQA_LOGITS_BYTES_PER_ELEM + _MQA_LOGITS_STATIC_SKIP_ELEMS = MQA_LOGITS_STATIC_SKIP_ELEMS + _MQA_LOGITS_TOTAL_MEM_FRACTION = MQA_LOGITS_TOTAL_MEM_FRACTION + _MQA_LOGITS_MAX_BYTES_ROCM = MQA_LOGITS_MAX_BYTES_ROCM + # One measured budget per device for the process lifetime. _mqa_logits_budget_bytes: Dict[int, int] = {} @staticmethod def _mqa_logits_free_mem_fraction() -> float: - return envs.SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION.get() + return mqa_logits_free_mem_fraction() def __init__( self, @@ -988,68 +996,21 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): return topk_result def _get_mqa_logits_budget_bytes(self, device_index: int) -> int: - free_mem_fraction = self._mqa_logits_free_mem_fraction() cached_budget = self._mqa_logits_budget_bytes.get(device_index) if cached_budget is not None: return cached_budget - total_mem = get_device_module().get_device_properties(device_index).total_memory - - total_mem_budget = int(total_mem * self._MQA_LOGITS_TOTAL_MEM_FRACTION) - mem_fraction_static = get_schedule().mem_fraction_static - if mem_fraction_static is None: - static_budget = total_mem_budget - else: - static_free_mem = int(total_mem * max(0.0, 1.0 - mem_fraction_static)) - static_budget = min( - int(static_free_mem * free_mem_fraction), - total_mem_budget, - ) - static_budget = max(1, static_budget) - - # Keep the static serving-memory guard during CUDA graph capture without - # caching it. The first non-capture prefill path will cache the real - # free-memory budget below. + # Graph capture cannot sync the host; use the static guard and do not + # cache it, so the first eager prefill still measures the real budget. if get_is_capture_mode(): - return static_budget + return mqa_logits_static_budget_bytes(device_index=device_index) - # Match the original free-memory guard: logits_bytes * 2 > free_mem. - # Synchronizes the host; cache the result capped by serving-memory headroom. - if _is_xpu: - # On XPU, use total_mem budget as the free-memory estimate; - # dynamic free-memory query is not supported the same way as CUDA. - # TODO Use torch.xpu.mem_get_info() when available (planned end of 2026). - budget_bytes = static_budget - else: - free_mem, _ = torch.cuda.mem_get_info(device_index) - budget_bytes = min(int(free_mem * free_mem_fraction), static_budget) - - budget_bytes = max(1, budget_bytes) + budget_bytes = mqa_logits_budget_bytes( + device_index=device_index, allow_sync=True + ) self._mqa_logits_budget_bytes[device_index] = budget_bytes return budget_bytes - def _should_chunk_mqa_logits( - self, num_q: int, num_k: int, device_index: int - ) -> Tuple[bool, int]: - """ - Detect whether we need to chunk the MQA logits computation to avoid OOM, - and on ROCm to stay under aiter's 2 GiB logits limit - Return: (need_chunk, logits_budget_bytes) - """ - # Quick static check for normal batches - if num_q * num_k < self._MQA_LOGITS_STATIC_SKIP_ELEMS: - return False, 0 - - logits_bytes = num_q * num_k * self._MQA_LOGITS_BYTES_PER_ELEM - logits_budget_bytes = self._get_mqa_logits_budget_bytes(device_index) - if _is_hip: - logits_budget_bytes = min( - logits_budget_bytes, self._MQA_LOGITS_MAX_BYTES_ROCM - ) - - need_chunk = logits_bytes > logits_budget_bytes - return need_chunk, logits_budget_bytes - def _get_topk_ragged( self, enable_dual_stream: bool, @@ -1139,8 +1100,11 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): token_to_batch_idx = metadata.get_token_to_batch_idx() q_offset = ks.shape[0] k_offset = k_fp8.shape[0] - need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits( - q_offset, k_offset, device_index + need_chunk, logits_budget_bytes = mqa_logits_should_chunk( + num_rows=q_offset, + num_cols=k_offset, + get_budget_bytes=lambda: self._get_mqa_logits_budget_bytes(device_index), + rocm=_is_hip, ) if not need_chunk: diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 77416f4e7..28d521854 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -35,9 +35,13 @@ from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa from sglang.srt.layers.attention.dsv4.compressor import Compressor from sglang.srt.layers.attention.dsv4.metadata import ( - _SM120_INDEXER_M_CHUNK, NonPagedIndexerPlan, PagedIndexerMetadata, + iter_row_chunks, +) +from sglang.srt.layers.attention.mqa_logits_utils import ( + mqa_logits_row_bytes, + mqa_logits_rows_per_chunk, ) from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.model_executor.forward_batch_info import ForwardMode @@ -690,6 +694,14 @@ class C4IndexerBackendMixin: ke = torch.where(ke - ks > c4_indexer.index_topk, ke, ks) c4_page_size = indexer_metadata.compressed_page_size max_seqlen_k = (final_c4_len + c4_page_size - 1) // c4_page_size * c4_page_size + rows_per_chunk = None + if indexer_metadata.mqa_logits_budget_bytes is not None: + # fp8_mqa_logits allocates [query_rows, align256(max_seqlen_k)] fp32. + rows_per_chunk = mqa_logits_rows_per_chunk( + num_rows=query_rows, + row_bytes=mqa_logits_row_bytes(max_seqlen_k), + budget_bytes=indexer_metadata.mqa_logits_budget_bytes, + ) plan = NonPagedIndexerPlan( page_table=request_page_table, gather_seq_lens=gather_seq_lens, @@ -699,21 +711,18 @@ class C4IndexerBackendMixin: max_seq_len=final_c4_len, max_seqlen_k=max_seqlen_k, query_rows=query_rows, + rows_per_chunk=rows_per_chunk, ) indexer_metadata.nonpaged_plan = plan return plan @staticmethod - def _forward_nonpaged_indexer( + def _gather_nonpaged_index_k( *, - q_indexer: torch.Tensor, - weights: torch.Tensor, c4_indexer: C4Indexer, token_to_kv_pool: DeepSeekV4TokenToKVPool, plan: NonPagedIndexerPlan, - ) -> torch.Tensor: - import deep_gemm - + ) -> Tuple[torch.Tensor, torch.Tensor]: k_u8, scale_u8 = token_to_kv_pool.get_index_k_scale_buffer( layer_id=c4_indexer.layer_id, seq_len_tensor=plan.gather_seq_lens, @@ -721,14 +730,25 @@ class C4IndexerBackendMixin: seq_len_sum=plan.seq_len_sum, max_seq_len=plan.max_seq_len, ) - k_fp8 = k_u8.view(FP8_DTYPE) - k_scale = scale_u8.view(torch.float32).squeeze(-1) + return k_u8.view(FP8_DTYPE), scale_u8.view(torch.float32).squeeze(-1) + + @staticmethod + def _nonpaged_mqa_logits( + *, + q_indexer: torch.Tensor, + weights: torch.Tensor, + kv: Tuple[torch.Tensor, torch.Tensor], + plan: NonPagedIndexerPlan, + rows: slice, + ) -> torch.Tensor: + import deep_gemm + return deep_gemm.fp8_mqa_logits( - q_indexer[: plan.query_rows], - (k_fp8, k_scale), - weights[: plan.query_rows], - plan.ks, - plan.ke, + q_indexer[rows], + kv, + weights[rows], + plan.ks[rows], + plan.ke[rows], clean_logits=False, max_seqlen_k=plan.max_seqlen_k, ) @@ -904,7 +924,11 @@ class C4IndexerBackendMixin: all_rows = slice(0, _c4sl.shape[0]) - def run_topk_transform(rows: slice, logits: torch.Tensor) -> None: + def run_topk_transform( + rows: slice, + logits: torch.Tensor, + topk_plan: Optional[torch.Tensor] = None, + ) -> None: row_raw_indices = raw_indices[rows] if raw_indices is not None else None if self.dsa_topk_backend.is_torch(): topk_transform_pytorch_vectorized( @@ -925,19 +949,21 @@ class C4IndexerBackendMixin: row_raw_indices, ) elif self.dsa_topk_backend.should_use_topk_v2(): + if topk_plan is None: + # The cached plan routes rows by their index in the full + # range, so a chunk needs one built over its own rows. + topk_plan = ( + indexer_metadata.topk_metadata + if rows == all_rows + else plan_topk_v2(c4_seq_lens[rows]) + ) topk_transform_paged_v2( logits, c4_seq_lens[rows], page_table[rows], c4_sparse_page_indices[rows], indexer_metadata.compressed_page_size, - # The cached plan routes rows by their index in the full - # range, so a chunk needs one built over its own rows. - ( - indexer_metadata.topk_metadata - if rows == all_rows or not is_hip() - else plan_topk_v2(c4_seq_lens[rows]) - ), + topk_plan, row_raw_indices, ) else: @@ -952,14 +978,27 @@ class C4IndexerBackendMixin: if nonpaged_plan is not None: assert isinstance(q_indexer, torch.Tensor) - logits = self._forward_nonpaged_indexer( - q_indexer=q_indexer, - weights=weights, + # K is gathered once; each row chunk's logits are reduced to top-k + # and dropped before the next chunk allocates, so only one chunk + # of logits is live at a time. + kv = self._gather_nonpaged_index_k( c4_indexer=c4_indexer, token_to_kv_pool=token_to_kv_pool, plan=nonpaged_plan, ) - run_topk_transform(all_rows, logits) + for rows in iter_row_chunks( + num_rows=nonpaged_plan.query_rows, + rows_per_chunk=nonpaged_plan.rows_per_chunk, + ): + logits = self._nonpaged_mqa_logits( + q_indexer=q_indexer, + weights=weights, + kv=kv, + plan=nonpaged_plan, + rows=rows, + ) + run_topk_transform(rows, logits) + del logits elif use_aiter_fp4: q_fp4, q_scale = q is_decode = forward_batch.forward_mode.is_decode() @@ -1012,7 +1051,11 @@ class C4IndexerBackendMixin: c4_indexer_kv_cache.shape[0], 64, 1, head_dim_with_sf ) - def run_paged_indexer(rows: slice, metadata: torch.Tensor) -> None: + def run_paged_indexer( + rows: slice, + metadata: torch.Tensor, + topk_plan: Optional[torch.Tensor] = None, + ) -> None: row_q = (q[0][rows], q[1][rows]) if isinstance(q, tuple) else q[rows] logits = fn( row_q, @@ -1024,18 +1067,30 @@ class C4IndexerBackendMixin: indexer_metadata.max_compressed_seq_len, False, ) - run_topk_transform(rows, logits) + run_topk_transform(rows, logits, topk_plan) deep_gemm_metadata = indexer_metadata.deep_gemm_metadata if isinstance(deep_gemm_metadata, list): - # SM120 only: DeepGEMM's metadata kernel caps the row count, so - # PagedIndexerMetadata split it; run indexer + topk per chunk. + # PagedIndexerMetadata split this forward into row chunks (SM120 + # kernel cap and/or logits memory budget), one schedule each. num_rows = _c4sl.shape[0] - for chunk_idx, start in enumerate( - range(0, num_rows, _SM120_INDEXER_M_CHUNK) + assert num_rows == indexer_metadata.compressed_seq_lens.shape[0], ( + f"chunk schedules were built for " + f"{indexer_metadata.compressed_seq_lens.shape[0]} rows, " + f"got {num_rows}" + ) + topk_plans = indexer_metadata.topk_metadata_chunks + for chunk_idx, rows in enumerate( + iter_row_chunks( + num_rows=num_rows, + rows_per_chunk=indexer_metadata.rows_per_chunk, + ) ): - rows = slice(start, min(start + _SM120_INDEXER_M_CHUNK, num_rows)) - run_paged_indexer(rows, deep_gemm_metadata[chunk_idx]) + run_paged_indexer( + rows, + deep_gemm_metadata[chunk_idx], + topk_plans[chunk_idx] if topk_plans is not None else None, + ) else: run_paged_indexer(all_rows, deep_gemm_metadata) diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index 0080ad024..da0bf0d84 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -1,14 +1,30 @@ from __future__ import annotations +import logging import warnings from dataclasses import dataclass, field, fields -from typing import Any, List, Optional +from typing import Any, Iterator, List, Optional import torch from sglang.srt.environ import envs +from sglang.srt.layers.attention.mqa_logits_utils import ( + mqa_logits_budget_bytes, + mqa_logits_needs_budget_check, + mqa_logits_row_bytes, + mqa_logits_rows_per_chunk, + mqa_logits_should_chunk, +) +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, +) +from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( + is_in_tc_piecewise_cuda_graph, +) from sglang.srt.utils import is_hip, is_sm120_supported, is_xpu +logger = logging.getLogger(__name__) + _IS_SM120 = is_sm120_supported() """ @@ -48,9 +64,57 @@ Some other notes: """ _LARGE_INDEXER_QUERY_THRESHOLD = 11673 +# DeepGEMM's paged-MQA metadata kernel cannot schedule more rows than this on +# SM120 (shared-memory cap), so SM120 always splits larger batches. _SM120_INDEXER_M_CHUNK = 4096 +def plan_indexer_row_chunks( + *, + num_rows: int, + num_cols: int, + budget_bytes: Optional[int], + sm120_row_cap: Optional[int], +) -> Optional[int]: + """Query rows per paged-indexer chunk; None runs the whole batch in one call. + + The fp32 logits are [num_rows, num_cols] per layer, so the chunk is the + smaller of the SM120 kernel cap and what the memory budget allows. + """ + rows_per_chunk = None + if sm120_row_cap is not None and num_rows > sm120_row_cap: + rows_per_chunk = sm120_row_cap + if budget_bytes is not None: + need_chunk, budget_bytes = mqa_logits_should_chunk( + num_rows=num_rows, + num_cols=num_cols, + get_budget_bytes=lambda: budget_bytes, + rocm=is_hip(), + ) + by_budget = ( + mqa_logits_rows_per_chunk( + num_rows=num_rows, + row_bytes=mqa_logits_row_bytes(num_cols), + budget_bytes=budget_bytes, + ) + if need_chunk + else None + ) + if by_budget is not None: + rows_per_chunk = ( + by_budget if rows_per_chunk is None else min(rows_per_chunk, by_budget) + ) + return rows_per_chunk + + +def iter_row_chunks(*, num_rows: int, rows_per_chunk: Optional[int]) -> Iterator[slice]: + if rows_per_chunk is None or rows_per_chunk >= num_rows: + yield slice(0, num_rows) + return + for start in range(0, num_rows, rows_per_chunk): + yield slice(start, min(start + rows_per_chunk, num_rows)) + + def copy_metadata( *, src, @@ -105,6 +169,8 @@ class NonPagedIndexerPlan: max_seq_len: int max_seqlen_k: int query_rows: int + # None runs all query rows in one fp8_mqa_logits call. + rows_per_chunk: Optional[int] = None @dataclass @@ -121,11 +187,20 @@ class PagedIndexerMetadata: # Rows per logits chunk for the prefill CUDA graph low-ratio indexer; 0 plans # all rows at once. row_chunk: int = 0 + # A list when the dynamic-budget forward is row-chunked: one schedule per chunk. deep_gemm_metadata: Any = field(init=False, repr=False) topk_metadata: torch.Tensor = field(init=False, repr=False) nonpaged_plan: Optional[NonPagedIndexerPlan] = field( init=False, repr=False, default=None ) + # Decided once per forward and shared by every layer's indexer call. + rows_per_chunk: Optional[int] = field(init=False, repr=False, default=None) + mqa_logits_budget_bytes: Optional[int] = field(init=False, repr=False, default=None) + # The top-k v2 plan routes rows by index within the batch it was built for, + # so a row-chunked forward needs one plan per chunk. + topk_metadata_chunks: Optional[List[torch.Tensor]] = field( + init=False, repr=False, default=None + ) def __post_init__(self): if ( @@ -149,6 +224,23 @@ class PagedIndexerMetadata: compressed_seq_lens = self.compressed_seq_lens.to(torch.int32) if compressed_seq_lens.dim() == 1: compressed_seq_lens = compressed_seq_lens.unsqueeze(-1) + num_rows = compressed_seq_lens.shape[0] + self.mqa_logits_budget_bytes = self._mqa_logits_budget(num_rows=num_rows) + self.rows_per_chunk = plan_indexer_row_chunks( + num_rows=num_rows, + num_cols=self.max_compressed_seq_len, + budget_bytes=self.mqa_logits_budget_bytes, + sm120_row_cap=_SM120_INDEXER_M_CHUNK if _IS_SM120 else None, + ) + if self.rows_per_chunk is not None: + logger.debug( + "DSV4 indexer chunks %d query rows x %d compressed cols into " + "%d-row chunks (logits budget %s bytes)", + num_rows, + self.max_compressed_seq_len, + self.rows_per_chunk, + self.mqa_logits_budget_bytes, + ) if self.row_chunk > 0: self.deep_gemm_metadata = torch.stack( [ @@ -160,16 +252,16 @@ class PagedIndexerMetadata: for _s in range(0, compressed_seq_lens.shape[0], self.row_chunk) ] ) - elif _IS_SM120 and compressed_seq_lens.shape[0] > _SM120_INDEXER_M_CHUNK: + elif self.rows_per_chunk is not None: # Chunk metadata is shared by all indexer layers in this forward. self.deep_gemm_metadata = [ get_paged_mqa_logits_metadata( - compressed_seq_lens[_s : _s + _SM120_INDEXER_M_CHUNK], + compressed_seq_lens[rows], self.compressed_page_size, deep_gemm.get_num_sms(), ) - for _s in range( - 0, compressed_seq_lens.shape[0], _SM120_INDEXER_M_CHUNK + for rows in iter_row_chunks( + num_rows=num_rows, rows_per_chunk=self.rows_per_chunk ) ] else: @@ -185,6 +277,14 @@ class PagedIndexerMetadata: from sglang.kernels.ops.attention.dsv4 import plan_topk_v2 self.topk_metadata = plan_topk_v2(self.compressed_seq_lens) + if self.rows_per_chunk is not None: + self.topk_metadata_chunks = [ + plan_topk_v2(self.compressed_seq_lens[rows]) + for rows in iter_row_chunks( + num_rows=self.compressed_seq_lens.shape[0], + rows_per_chunk=self.rows_per_chunk, + ) + ] else: self.topk_metadata = torch.empty((0,)) @@ -193,6 +293,28 @@ class PagedIndexerMetadata: f"compress_ratio {self.compress_ratio} must divide page_size {self.page_size}" ) + def _mqa_logits_budget(self, *, num_rows: int) -> Optional[int]: + """Free-memory budget for this forward's logits; None disables chunking. + + Graph-backed forwards keep a single call: their shapes are fixed at + capture and the free-memory read would sync the host mid-capture. + """ + if self.use_prefill_cuda_graph or not self.compressed_seq_lens.is_cuda: + return None + if not mqa_logits_needs_budget_check( + num_rows=num_rows, num_cols=self.max_compressed_seq_len + ): + return None + if ( + torch.cuda.is_current_stream_capturing() + or is_in_breakable_cuda_graph() + or is_in_tc_piecewise_cuda_graph() + ): + return None + return mqa_logits_budget_bytes( + device_index=self.compressed_seq_lens.device.index, allow_sync=True + ) + @property def max_seq_len(self) -> int: return self.page_table.shape[1] * self.page_size @@ -213,13 +335,22 @@ class PagedIndexerMetadata: ] def copy_(self, other: PagedIndexerMetadata): - if is_hip(): + # A chunked schedule list has no in-place copy; rebind it instead. + chunked = isinstance(self.deep_gemm_metadata, list) or isinstance( + other.deep_gemm_metadata, list + ) + if is_hip() or chunked: copy_fields = ["page_table", "compressed_seq_lens"] assign_fields = ["deep_gemm_metadata", "nonpaged_plan"] else: copy_fields = ["page_table", "compressed_seq_lens", "deep_gemm_metadata"] assign_fields = ["nonpaged_plan"] copy_fields += ["topk_metadata"] + assign_fields += [ + "rows_per_chunk", + "mqa_logits_budget_bytes", + "topk_metadata_chunks", + ] copy_metadata( src=other, dst=self, diff --git a/python/sglang/srt/layers/attention/mqa_logits_utils.py b/python/sglang/srt/layers/attention/mqa_logits_utils.py new file mode 100644 index 000000000..84f48825e --- /dev/null +++ b/python/sglang/srt/layers/attention/mqa_logits_utils.py @@ -0,0 +1,101 @@ +"""Memory budget for the fp32 MQA-logits scratch of the DeepSeek sparse +attention indexers (DSA, DSV4, DSV4.1). + +The indexer scores every query row against every key column into one fp32 +logits matrix that no pool sized by mem_fraction_static accounts for. These +helpers bound that matrix and slice it by query rows. +""" + +from __future__ import annotations + +from typing import Callable, Optional, Tuple + +import torch + +from sglang.srt.environ import envs +from sglang.srt.runtime_context import get_schedule +from sglang.srt.utils import get_device_module, is_hip, is_xpu +from sglang.srt.utils.common import ceil_div + +MQA_LOGITS_BYTES_PER_ELEM = 4 +MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000 +MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3 +# aiter's fp8_mqa_logits only compiles below 2 GiB of logits (buffer_store). +MQA_LOGITS_MAX_BYTES_ROCM = 2**31 - 1 +# DeepGEMM pads the logits row stride to 1024 bytes, i.e. 256 fp32 columns. +MQA_LOGITS_ROW_ALIGN_ELEMS = 256 + + +def mqa_logits_free_mem_fraction() -> float: + return envs.SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION.get() + + +def mqa_logits_needs_budget_check(*, num_rows: int, num_cols: int) -> bool: + return num_rows * num_cols >= MQA_LOGITS_STATIC_SKIP_ELEMS + + +def mqa_logits_row_bytes(num_cols: int) -> int: + aligned_cols = ( + ceil_div(num_cols, MQA_LOGITS_ROW_ALIGN_ELEMS) * MQA_LOGITS_ROW_ALIGN_ELEMS + ) + return aligned_cols * MQA_LOGITS_BYTES_PER_ELEM + + +def mqa_logits_static_budget_bytes(*, device_index: int) -> int: + """Budget from configuration alone (no device query); safe during graph capture.""" + total_mem = get_device_module().get_device_properties(device_index).total_memory + total_mem_budget = int(total_mem * MQA_LOGITS_TOTAL_MEM_FRACTION) + mem_fraction_static = get_schedule().mem_fraction_static + if mem_fraction_static is None: + budget = total_mem_budget + else: + static_free_mem = int(total_mem * max(0.0, 1.0 - mem_fraction_static)) + budget = min( + int(static_free_mem * mqa_logits_free_mem_fraction()), total_mem_budget + ) + return max(1, budget) + + +def mqa_logits_budget_bytes(*, device_index: int, allow_sync: bool) -> int: + """Static budget capped by current free memory; mem_get_info syncs, so + callers pass allow_sync=False under CUDA graph capture.""" + budget = mqa_logits_static_budget_bytes(device_index=device_index) + if allow_sync and not is_xpu(): + free_mem, _ = torch.cuda.mem_get_info(device_index) + budget = min(int(free_mem * mqa_logits_free_mem_fraction()), budget) + if is_hip(): + budget = min(budget, MQA_LOGITS_MAX_BYTES_ROCM) + return max(1, budget) + + +def mqa_logits_should_chunk( + *, + num_rows: int, + num_cols: int, + get_budget_bytes: Callable[[], int], + rocm: bool, +) -> Tuple[bool, int]: + """Whether a [num_rows, num_cols] fp32 logits matrix must be row-chunked. + + Returns (need_chunk, effective_budget_bytes). Matrices below + MQA_LOGITS_STATIC_SKIP_ELEMS never chunk and the budget is not queried + (it may synchronize the host), hence the callable. On ROCm the budget is + also capped by aiter's logits ceiling. + """ + if not mqa_logits_needs_budget_check(num_rows=num_rows, num_cols=num_cols): + return False, 0 + budget_bytes = get_budget_bytes() + if rocm: + budget_bytes = min(budget_bytes, MQA_LOGITS_MAX_BYTES_ROCM) + logits_bytes = num_rows * num_cols * MQA_LOGITS_BYTES_PER_ELEM + return logits_bytes > budget_bytes, budget_bytes + + +def mqa_logits_rows_per_chunk( + *, num_rows: int, row_bytes: int, budget_bytes: int +) -> Optional[int]: + """Query rows per chunk so one logits chunk fits the budget; None if all rows fit.""" + if num_rows * row_bytes <= budget_bytes: + return None + rows = max(budget_bytes // max(row_bytes, 1), 1) + return int(rows) if rows < num_rows else None diff --git a/test/registered/unit/layers/attention/test_dsa_mqa_logits_chunking.py b/test/registered/unit/layers/attention/test_dsa_mqa_logits_chunking.py index 0500512b0..4d9e0a703 100644 --- a/test/registered/unit/layers/attention/test_dsa_mqa_logits_chunking.py +++ b/test/registered/unit/layers/attention/test_dsa_mqa_logits_chunking.py @@ -1,4 +1,5 @@ -"""Contract tests for the DSA indexer's MQA-logits chunk budget. +"""Contract tests for the MQA-logits chunk decision shared by the DSA and DSV4 +indexers. On ROCm the `[num_q x num_k]` fp32 logits tensor goes to aiter's `fp8_mqa_logits`, which only compiles below 2 GiB, so the budget that decides @@ -8,34 +9,30 @@ The measured memory budget is stubbed: it is the only input the limit has to beat, and stubbing it keeps these tests on CPU. """ -from unittest import mock - import pytest torch = pytest.importorskip("torch") -from sglang.srt.layers.attention.dsa import dsa_indexer # noqa: E402 +from sglang.srt.layers.attention.mqa_logits_utils import ( # noqa: E402 + MQA_LOGITS_MAX_BYTES_ROCM, + mqa_logits_should_chunk, +) from sglang.test.ci.ci_register import register_cpu_ci # noqa: E402 register_cpu_ci(est_time=9, suite="base-a-test-cpu") -CEILING = dsa_indexer.Indexer._MQA_LOGITS_MAX_BYTES_ROCM +CEILING = MQA_LOGITS_MAX_BYTES_ROCM # More than any single logits tensor here needs, so it never decides a case. HUGE_MEM_BUDGET = 64 * 2**30 def _decide(num_q, num_k, mem_budget=HUGE_MEM_BUDGET, is_hip=True): - # __new__ skips an __init__ that needs a model config and a device. - indexer = dsa_indexer.Indexer.__new__(dsa_indexer.Indexer) - with ( - mock.patch.object(dsa_indexer, "_is_hip", is_hip), - mock.patch.object( - dsa_indexer.Indexer, - "_get_mqa_logits_budget_bytes", - return_value=mem_budget, - ), - ): - return indexer._should_chunk_mqa_logits(num_q, num_k, 0) + return mqa_logits_should_chunk( + num_rows=num_q, + num_cols=num_k, + get_budget_bytes=lambda: mem_budget, + rocm=is_hip, + ) def test_the_ceiling_is_the_largest_logits_aiter_still_takes(): diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py index c427c7b8c..ad06b05e8 100644 --- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py +++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py @@ -11,19 +11,35 @@ from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend from sglang.srt.layers.attention.dsv4.indexer import ( FP8_DTYPE, C4IndexerBackendMixin, + topk_transform_pytorch_vectorized, ) from sglang.srt.layers.attention.dsv4.metadata import ( NonPagedIndexerPlan, PagedIndexerMetadata, + iter_row_chunks, + plan_indexer_row_chunks, +) +from sglang.srt.layers.attention.mqa_logits_utils import ( + MQA_LOGITS_MAX_BYTES_ROCM, + mqa_logits_budget_bytes, + mqa_logits_row_bytes, + mqa_logits_rows_per_chunk, + mqa_logits_should_chunk, ) from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase -register_cpu_ci(est_time=11, suite="base-a-test-cpu") +register_cpu_ci(est_time=17, suite="base-a-test-cpu") _INDEXER = "sglang.srt.layers.attention.dsv4.indexer" +_METADATA = "sglang.srt.layers.attention.dsv4.metadata" +_MQA_UTILS = "sglang.srt.layers.attention.mqa_logits_utils" + +# issue #35201: 372K raw tokens -> 92992 c4 columns, padded to 93184 by DeepGEMM. +_ISSUE_C4_COLS = 92992 +_ISSUE_ALIGNED_COLS = 93184 class TestDSV4PagedIndexerMetadata(CustomTestCase): @@ -247,8 +263,11 @@ class TestDSV4TopKDispatch(CustomTestCase): torch.empty((1, 1, 1)), ) ) - backend._get_nonpaged_indexer_plan = MagicMock(return_value=object()) - backend._forward_nonpaged_indexer = MagicMock(return_value=logits) + backend._get_nonpaged_indexer_plan = MagicMock( + return_value=SimpleNamespace(query_rows=1, rows_per_chunk=None) + ) + backend._gather_nonpaged_index_k = MagicMock(return_value=(object(), object())) + backend._nonpaged_mqa_logits = MagicMock(return_value=logits) indexer_capturer = MagicMock() with ( @@ -353,7 +372,9 @@ class TestDSV4NonPagedIndexer(CustomTestCase): extend_start_loc=torch.tensor([0], dtype=torch.int32), extend_num_tokens=query_rows, ) - metadata = SimpleNamespace(nonpaged_plan=None, compressed_page_size=64) + metadata = SimpleNamespace( + nonpaged_plan=None, compressed_page_size=64, mqa_logits_budget_bytes=None + ) page_table = torch.tensor([[3, 1]], dtype=torch.int32).repeat(query_rows, 1) c4_seq_lens = torch.tensor([62, 63, 64, 65], dtype=torch.int32) @@ -388,6 +409,50 @@ class TestDSV4NonPagedIndexer(CustomTestCase): with threshold.override(0): self.assertIsNone(build_plan()) + def test_plan_row_chunking_follows_the_forward_budget(self): + backend = SimpleNamespace(_can_use_nonpaged_indexer=lambda **_: True) + backend.dsa_topk_backend = SimpleNamespace(is_sgl_kernel=lambda: True) + c4_indexer = SimpleNamespace(use_fp4_indexer=False, index_topk=512) + query_rows = 8192 + seq_len = 372_000 + batch = SimpleNamespace( + seq_lens=torch.tensor([seq_len], dtype=torch.int32), + seq_lens_cpu=[seq_len], + extend_seq_lens_cpu=[query_rows], + extend_seq_lens=torch.tensor([query_rows], dtype=torch.int32), + extend_start_loc=torch.tensor([0], dtype=torch.int32), + extend_num_tokens=query_rows, + ) + c4_seq_lens = torch.full((query_rows,), seq_len // 4, dtype=torch.int32) + page_table = torch.zeros((query_rows, 1), dtype=torch.int32) + + def build_plan(budget): + metadata = SimpleNamespace( + nonpaged_plan=None, + compressed_page_size=64, + mqa_logits_budget_bytes=budget, + ) + return C4IndexerBackendMixin._get_nonpaged_indexer_plan( + backend, + c4_indexer=c4_indexer, + forward_batch=batch, + indexer_metadata=metadata, + page_table=page_table, + c4_seq_lens=c4_seq_lens, + query_rows=query_rows, + ) + + # No budget measured this forward (small batch or graph): one call. + self.assertIsNone(build_plan(None).rows_per_chunk) + budget = 512 << 20 + plan = build_plan(budget) + # 8192 rows x align256(93056) cols x 4 B is far over 512 MiB. + self.assertIsNotNone(plan.rows_per_chunk) + self.assertLess(plan.rows_per_chunk, query_rows) + self.assertLessEqual( + plan.rows_per_chunk * mqa_logits_row_bytes(plan.max_seqlen_k), budget + ) + def test_extreme_plan_metadata_is_bounded_and_fail_closed(self): backend = SimpleNamespace(_can_use_nonpaged_indexer=lambda **_: True) backend.dsa_topk_backend = SimpleNamespace(is_sgl_kernel=lambda: True) @@ -401,7 +466,9 @@ class TestDSV4NonPagedIndexer(CustomTestCase): extend_start_loc=torch.tensor([0], dtype=torch.int32), extend_num_tokens=query_rows, ) - metadata = SimpleNamespace(nonpaged_plan=None, compressed_page_size=64) + metadata = SimpleNamespace( + nonpaged_plan=None, compressed_page_size=64, mqa_logits_budget_bytes=None + ) page_table = torch.zeros((query_rows, 1), dtype=torch.int32) c4_seq_lens = torch.tensor( [124_997, 124_998, 124_999, 125_000], dtype=torch.int32 @@ -439,7 +506,9 @@ class TestDSV4NonPagedIndexer(CustomTestCase): backend = SimpleNamespace(_can_use_nonpaged_indexer=can_use_nonpaged_indexer) backend.dsa_topk_backend = SimpleNamespace(is_sgl_kernel=lambda: True) c4_indexer = SimpleNamespace(use_fp4_indexer=False, index_topk=512) - metadata = SimpleNamespace(nonpaged_plan=None, compressed_page_size=64) + metadata = SimpleNamespace( + nonpaged_plan=None, compressed_page_size=64, mqa_logits_budget_bytes=None + ) def build_plan(query_rows): batch = SimpleNamespace( @@ -503,13 +572,18 @@ class TestDSV4NonPagedIndexer(CustomTestCase): deep_gemm = SimpleNamespace(fp8_mqa_logits=MagicMock(return_value=expected)) with patch.dict(sys.modules, {"deep_gemm": deep_gemm}): - actual = C4IndexerBackendMixin._forward_nonpaged_indexer( - q_indexer=q_indexer, - weights=weights, + kv = C4IndexerBackendMixin._gather_nonpaged_index_k( c4_indexer=c4_indexer, token_to_kv_pool=token_to_kv_pool, plan=plan, ) + actual = C4IndexerBackendMixin._nonpaged_mqa_logits( + q_indexer=q_indexer, + weights=weights, + kv=kv, + plan=plan, + rows=slice(0, plan.query_rows), + ) self.assertIs(actual, expected) token_to_kv_pool.get_index_k_scale_buffer.assert_called_once_with( @@ -531,6 +605,276 @@ class TestDSV4NonPagedIndexer(CustomTestCase): self.assertEqual(call.kwargs, {"clean_logits": False, "max_seqlen_k": 128}) +class TestMqaLogitsBudgetArithmetic(CustomTestCase): + def test_row_bytes_follow_deepgemm_stride_alignment(self): + # DeepGEMM pads the fp32 logits row stride to 1024 B, i.e. 256 columns. + self.assertEqual(mqa_logits_row_bytes(1), 256 * 4) + self.assertEqual(mqa_logits_row_bytes(256), 256 * 4) + self.assertEqual(mqa_logits_row_bytes(257), 512 * 4) + self.assertEqual(mqa_logits_row_bytes(_ISSUE_C4_COLS), _ISSUE_ALIGNED_COLS * 4) + + def test_rows_per_chunk_keeps_one_chunk_inside_budget(self): + row_bytes = mqa_logits_row_bytes(_ISSUE_C4_COLS) + budget = 512 << 20 + # 4096 rows x 93184 cols x 4 B = 1.42 GiB > 512 MiB: must slice. + rows = mqa_logits_rows_per_chunk( + num_rows=4096, row_bytes=row_bytes, budget_bytes=budget + ) + self.assertIsNotNone(rows) + self.assertLess(rows, 4096) + self.assertLessEqual(rows * row_bytes, budget) + # The whole matrix fits: single call. + self.assertIsNone( + mqa_logits_rows_per_chunk( + num_rows=64, row_bytes=row_bytes, budget_bytes=budget + ) + ) + # A tight budget never yields a chunk larger than the budget, even when + # that means fewer rows than a full page of queries. + tight = 64 << 20 + rows = mqa_logits_rows_per_chunk( + num_rows=4096, row_bytes=row_bytes, budget_bytes=tight + ) + self.assertEqual(rows, tight // row_bytes) + self.assertLessEqual(rows * row_bytes, tight) + # Few query rows still chunk when they do not fit. + self.assertEqual( + mqa_logits_rows_per_chunk( + num_rows=100, row_bytes=row_bytes, budget_bytes=40 * row_bytes + ), + 40, + ) + # A budget below one row degrades to single-row chunks, never None. + self.assertEqual( + mqa_logits_rows_per_chunk( + num_rows=4096, row_bytes=row_bytes, budget_bytes=1 + ), + 1, + ) + + def test_should_chunk_caps_the_budget_only_on_rocm(self): + huge = 64 << 30 + # 16384 x 32768 x 4 B is exactly 2 GiB, aiter's compile-time ceiling. + self.assertEqual( + mqa_logits_should_chunk( + num_rows=16384, num_cols=32768, get_budget_bytes=lambda: huge, rocm=True + ), + (True, MQA_LOGITS_MAX_BYTES_ROCM), + ) + self.assertEqual( + mqa_logits_should_chunk( + num_rows=16384, + num_cols=32768, + get_budget_bytes=lambda: huge, + rocm=False, + ), + (False, huge), + ) + + def test_should_chunk_skips_small_matrices_without_querying_the_budget(self): + get_budget = MagicMock(return_value=1) + # 64 decode rows x 100K columns is far below the 8M-element threshold. + self.assertEqual( + mqa_logits_should_chunk( + num_rows=64, num_cols=100_000, get_budget_bytes=get_budget, rocm=False + ), + (False, 0), + ) + get_budget.assert_not_called() + + def test_plan_combines_sm120_cap_with_budget(self): + budget = 512 << 20 + by_budget = mqa_logits_rows_per_chunk( + num_rows=8192, + row_bytes=mqa_logits_row_bytes(_ISSUE_C4_COLS), + budget_bytes=budget, + ) + cases = ( + # (num_rows, num_cols, budget_bytes, sm120_row_cap) -> rows_per_chunk + ((8192, 1024, None, None), None), + ((8192, 1024, None, 4096), 4096), + ((4096, 1024, None, 4096), None), + ((8192, _ISSUE_C4_COLS, budget, None), by_budget), + ((8192, _ISSUE_C4_COLS, budget, 4096), min(4096, by_budget)), + ) + for (num_rows, num_cols, budget_bytes, cap), expected in cases: + with self.subTest(num_rows=num_rows, num_cols=num_cols, cap=cap): + self.assertEqual( + plan_indexer_row_chunks( + num_rows=num_rows, + num_cols=num_cols, + budget_bytes=budget_bytes, + sm120_row_cap=cap, + ), + expected, + ) + + def test_iter_row_chunks_covers_rows_exactly_once(self): + self.assertEqual( + list(iter_row_chunks(num_rows=10, rows_per_chunk=4)), + [slice(0, 4), slice(4, 8), slice(8, 10)], + ) + self.assertEqual( + list(iter_row_chunks(num_rows=10, rows_per_chunk=None)), [slice(0, 10)] + ) + self.assertEqual( + list(iter_row_chunks(num_rows=10, rows_per_chunk=64)), [slice(0, 10)] + ) + + def test_static_budget_never_queries_free_memory(self): + total = 80 << 30 + props = SimpleNamespace(total_memory=total) + device_module = SimpleNamespace( + get_device_properties=MagicMock(return_value=props) + ) + schedule = SimpleNamespace(mem_fraction_static=0.9) + with ( + envs.SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION.override(0.2), + patch(f"{_MQA_UTILS}.get_device_module", return_value=device_module), + patch(f"{_MQA_UTILS}.get_schedule", return_value=schedule), + patch(f"{_MQA_UTILS}.is_hip", return_value=False), + patch(f"{_MQA_UTILS}.is_xpu", return_value=False), + patch( + "torch.cuda.mem_get_info", return_value=(6 << 30, total) + ) as mem_get_info, + ): + static = mqa_logits_budget_bytes(device_index=0, allow_sync=False) + mem_get_info.assert_not_called() + live = mqa_logits_budget_bytes(device_index=0, allow_sync=True) + mem_get_info.assert_called_once_with(0) + # static: 80 GiB x (1 - 0.9) x 0.2; live is further capped by 6 GiB free x 0.2. + self.assertEqual(static, int(int(total * 0.1) * 0.2)) + self.assertEqual(live, int((6 << 30) * 0.2)) + + +class TestPagedIndexerMetadataChunking(CustomTestCase): + """The schedule list and the top-k plan list must be built over the exact + row chunks the indexer loops over; a mismatch would silently score rows + with another chunk's schedule.""" + + def _build(self, *, num_rows: int, budget, use_topk_v2: bool): + deep_gemm = SimpleNamespace( + get_num_sms=MagicMock(return_value=1), + get_paged_mqa_logits_metadata=MagicMock( + side_effect=lambda c4, *_: torch.zeros((2, 2), dtype=torch.int32) + ), + ) + c4_seq_lens = torch.arange(1, num_rows + 1, dtype=torch.int32) + page_table = torch.zeros( + (num_rows, _ISSUE_ALIGNED_COLS // 64), dtype=torch.int32 + ) + with ( + patch.dict(sys.modules, {"deep_gemm": deep_gemm}), + envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.override(False), + envs.SGLANG_OPT_USE_JIT_INDEXER_METADATA.override(False), + patch(f"{_METADATA}.is_hip", return_value=False), + patch(f"{_METADATA}.is_xpu", return_value=False), + patch(f"{_METADATA}._IS_SM120", False), + patch.object( + PagedIndexerMetadata, "_mqa_logits_budget", return_value=budget + ), + patch( + "sglang.kernels.ops.attention.dsv4.plan_topk_v2", + side_effect=lambda seq_lens: seq_lens.new_zeros( + (seq_lens.shape[0] + 1, 2) + ), + ) as plan_topk_v2, + ): + metadata = PagedIndexerMetadata( + page_size=256, + compressed_page_size=64, + page_table=page_table, + compressed_seq_lens=c4_seq_lens, + use_topk_v2=use_topk_v2, + ) + return metadata, deep_gemm, plan_topk_v2 + + def test_budget_splits_schedules_and_topk_plans_over_the_same_rows(self): + num_rows, budget = 4096, 512 << 20 + metadata, deep_gemm, plan_topk_v2 = self._build( + num_rows=num_rows, budget=budget, use_topk_v2=True + ) + expected_rows = mqa_logits_rows_per_chunk( + num_rows=num_rows, + row_bytes=mqa_logits_row_bytes(_ISSUE_ALIGNED_COLS), + budget_bytes=budget, + ) + self.assertEqual(metadata.rows_per_chunk, expected_rows) + self.assertEqual(metadata.mqa_logits_budget_bytes, budget) + chunks = list(iter_row_chunks(num_rows=num_rows, rows_per_chunk=expected_rows)) + self.assertGreater(len(chunks), 1) + + self.assertIsInstance(metadata.deep_gemm_metadata, list) + self.assertEqual(len(metadata.deep_gemm_metadata), len(chunks)) + schedule_rows = [ + call.args[0] + for call in deep_gemm.get_paged_mqa_logits_metadata.call_args_list + ] + torch.testing.assert_close( + torch.cat(schedule_rows), metadata.compressed_seq_lens.unsqueeze(-1) + ) + self.assertEqual( + [r.shape[0] for r in schedule_rows], [c.stop - c.start for c in chunks] + ) + + self.assertEqual(len(metadata.topk_metadata_chunks), len(chunks)) + # First call is the full-batch plan; the rest are one per chunk. + plan_rows = [call.args[0] for call in plan_topk_v2.call_args_list] + torch.testing.assert_close(plan_rows[0], metadata.compressed_seq_lens) + torch.testing.assert_close( + torch.cat(plan_rows[1:]), metadata.compressed_seq_lens + ) + self.assertEqual( + [r.shape[0] for r in plan_rows[1:]], [c.stop - c.start for c in chunks] + ) + + def test_no_budget_keeps_the_single_call_shape(self): + metadata, deep_gemm, plan_topk_v2 = self._build( + num_rows=4096, budget=None, use_topk_v2=True + ) + self.assertIsNone(metadata.rows_per_chunk) + self.assertIsNone(metadata.topk_metadata_chunks) + self.assertIsInstance(metadata.deep_gemm_metadata, torch.Tensor) + deep_gemm.get_paged_mqa_logits_metadata.assert_called_once() + plan_topk_v2.assert_called_once() + + +class TestChunkedTopKMatchesUnchunked(CustomTestCase): + """Each row's top-k depends only on its own logits row, sequence length and + page-table row, so scoring the batch in row chunks must select the same + pages as one pass. This is the property the chunk loop relies on.""" + + def test_row_chunks_select_the_same_pages(self): + torch.manual_seed(0) + rows, width, topk, page_size = 37, 2048, 64, 64 + logits = torch.randn(rows, width, dtype=torch.float32) + seq_lens = torch.randint(1, width, (rows,), dtype=torch.int32) + page_table = torch.randint( + 0, 4096, (rows, width // page_size), dtype=torch.int32 + ) + + def run(rows_per_chunk): + out = torch.full((rows, topk), -1, dtype=torch.int32) + for rows_slice in iter_row_chunks( + num_rows=rows, rows_per_chunk=rows_per_chunk + ): + topk_transform_pytorch_vectorized( + logits[rows_slice], + seq_lens[rows_slice], + page_table[rows_slice], + out[rows_slice], + page_size, + None, + ) + # Unsorted top-k: compare the selected sets row by row. + return out.sort(dim=1).values + + expected = run(None) + for rows_per_chunk in (1, 7, 16, rows - 1): + with self.subTest(rows_per_chunk=rows_per_chunk): + self.assertTrue(torch.equal(run(rows_per_chunk), expected)) + + class TestCandidateIndexerGating(CustomTestCase): def test_candidate_indexer_gating(self): from sglang.srt.layers.attention.dsv4 import candidate_indexer