From adcdfb17ec108321b7f0f132568e6e3e3bdff4f3 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 5 Aug 2026 08:09:28 +0000 Subject: [PATCH 01/18] [None][refactor] modularize sparse top-k selection Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/__init__.py | 2 - .../attention_backend/sparse/dsa/indexer.py | 287 ++++---------- .../attention_backend/sparse/dsa/metadata.py | 20 +- .../triattention/triattention.py | 11 +- tensorrt_llm/_torch/modules/top_k.py | 374 ++++++++++++++++++ .../attention/sparse/dsa/test_dsa_indexer.py | 58 ++- tests/unittest/_torch/modules/test_top_k.py | 218 ++++++++++ 7 files changed, 733 insertions(+), 237 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/top_k.py create mode 100644 tests/unittest/_torch/modules/test_top_k.py diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py index 80908e08b0ac..d08340286dbf 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py @@ -20,7 +20,6 @@ rotate_activation, split_prefill_chunks, transform_local_topk_and_prepare_pool_view, - warmup_heuristic_topk_decode, ) from .metadata import DSAtrtllmAttentionMetadata, build_req_idx_per_token from .params import DSABackendForwardArgs, DSAMetadataParams, DSAParams @@ -49,5 +48,4 @@ "rotate_activation", "split_prefill_chunks", "transform_local_topk_and_prepare_pool_view", - "warmup_heuristic_topk_decode", ] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 3e67594b1690..8760affd4042 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -5,10 +5,9 @@ from __future__ import annotations import os -import threading from contextlib import contextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING, List, Optional, Set, Tuple +from typing import TYPE_CHECKING, List, Optional, Tuple import torch import torch.nn as nn @@ -26,6 +25,12 @@ maybe_execute_in_parallel, ) from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding +from tensorrt_llm._torch.modules.top_k import ( + DecodeTopK, + DecodeTopKPolicy, + PrefillTopK, + PrefillTopKImplementation, +) from tensorrt_llm._torch.utils import Fp4QuantizedTensor, maybe_compile from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned from tensorrt_llm.deep_gemm import ( @@ -67,71 +72,9 @@ hadamard_transform = None HAS_FAST_HADAMARD = False -# Idempotency guard for warmup_heuristic_topk_decode — keyed by -# (device_index, top_k, hint_size, num_cols). Prevents repeated allocations -# and synchronizations when multiple Indexer modules invoke the warmup with -# the same parameters during model construction. -_HEURISTIC_TOPK_WARMUP_DONE: Set[Tuple[int, int, int, int]] = set() -_HEURISTIC_TOPK_WARMUP_LOCK = threading.Lock() _DG_SCHEDULE_BLOCK_KV = 64 -def warmup_heuristic_topk_decode( - top_k: int = 2048, hint_size: int = 2048, num_cols: int = 4096 -) -> None: - """Pre-initialize cached hardware attributes in the C++ Scheme X dispatcher. - - The dispatcher inside ``invokeIndexerTopKDecode`` lazily queries - ``cudaDeviceGetAttribute`` for ``MultiProcessorCount`` and - ``L2CacheSize`` on its first call. Those host-side queries must not - be issued during ``cudaStreamBeginCapture / EndCapture``: the values - captured there become frozen into the graph and cannot be refreshed - across replays on a different device. - - This warmup issues one small heuristic decode call so the static - caches are populated before any CUDA Graph capture begins. Must be - called from the Indexer setup hook (``layer_idx == 0``) when - ``enable_heuristic_topk`` is true. - - Repeated invocations with the same ``(device, top_k, hint_size, - num_cols)`` key are short-circuited so that constructing many Indexer - modules in the same process does not re-allocate scratch tensors or - issue redundant synchronizations. - """ - key = (torch.cuda.current_device(), top_k, hint_size, num_cols) - with _HEURISTIC_TOPK_WARMUP_LOCK: - if key in _HEURISTIC_TOPK_WARMUP_DONE: - return - _HEURISTIC_TOPK_WARMUP_DONE.add(key) - - device = torch.device("cuda") - logits = torch.zeros((1, num_cols), dtype=torch.float32, device=device) - seq_lens = torch.tensor([num_cols], dtype=torch.int32, device=device) - indices = torch.empty((1, top_k), dtype=torch.int32, device=device) - pre_idx = torch.zeros((1, hint_size), dtype=torch.int32, device=device) - scratch = torch.empty((top_k,), dtype=torch.float32, device=device) - # The default warmup geometry (num_cols=4096) falls below kSeqSmall=12288 - # and routes to the Radix path with blocks_per_row=2 (num_rows=1 sweeps - # bp ∈ [2, maxByCols=2]). The cpp op rejects blocks_per_row > 1 without - # caller-owned radix aux scratch, so supply worst-case (kMaxBlocksPerRowDecode=10) - # buffers here. Cost is negligible (~80 KB) and the warmup is a one-shot. - _radix_max_bp = 10 - radix_aux_indices = torch.empty((1, _radix_max_bp, top_k), dtype=torch.int32, device=device) - radix_aux_logits = torch.empty((1, _radix_max_bp, top_k), dtype=torch.float32, device=device) - torch.ops.trtllm.indexer_topk_decode( - logits, - seq_lens, - indices, - 1, - top_k, - pre_idx=pre_idx, - heuristic_scratch=scratch, - radix_aux_indices=radix_aux_indices, - radix_aux_logits=radix_aux_logits, - ) - torch.cuda.synchronize() - - def _pick_dsl_expand( next_n: int, num_sms: int, @@ -699,11 +642,35 @@ def __init__( ) self.mtp_index_share = sparse_params.mtp_index_share - if self._enable_heuristic_topk and layer_idx == 0: - # Populate static caches (sm_count, L2 cache size) inside the C++ - # Scheme X dispatcher before any CUDA Graph capture so the host - # attribute queries do not end up frozen into a captured graph. - warmup_heuristic_topk_decode(top_k=self.index_topk) + self.prefill_top_k = PrefillTopK( + self.index_topk, + PrefillTopKImplementation.TRTLLM, + ) + if self.use_cute_dsl_topk: + decode_top_k_policy = ( + DecodeTopKPolicy.CUTE_DSL_GVR + if self._enable_heuristic_topk + else DecodeTopKPolicy.CUTE_DSL_PREFERRED + ) + elif self._enable_heuristic_topk: + decode_top_k_policy = DecodeTopKPolicy.TRTLLM_HEURISTIC + else: + decode_top_k_policy = DecodeTopKPolicy.TRTLLM + self.decode_top_k = DecodeTopK( + self.index_topk, + decode_top_k_policy, + compress_ratio=self.compress_ratio, + ) + + if decode_top_k_policy == DecodeTopKPolicy.TRTLLM_HEURISTIC and layer_idx == 0: + # Populate static caches inside the C++ dispatcher before CUDA + # Graph capture. DecodeTopK globally de-duplicates this warmup. + self.decode_top_k.prepare( + device=torch.device("cuda", torch.cuda.current_device()), + max_num_columns=4096, + next_n=1, + input_dtype=torch.float32, + ) # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM # (populated in cache_derived_state; maps to TF32 tensor cores on Ampere+) @@ -1099,11 +1066,11 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): """ num_contexts = metadata.num_contexts num_generations = metadata.num_generations + gen_seq_lens = metadata.get_indexer_kv_lens( + metadata.kv_lens_cuda_runtime[num_contexts : num_contexts + num_generations] + ) + metadata.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens if not metadata.use_expanded_buffers_for_mtp: - gen_seq_lens = metadata.get_indexer_kv_lens( - metadata.kv_lens_cuda_runtime[num_contexts : num_contexts + num_generations] - ) - metadata.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens next_n_cap = metadata.kv_lens_cuda_2d.shape[1] metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_( gen_seq_lens.unsqueeze(-1).expand(-1, next_n_cap) @@ -1379,7 +1346,6 @@ def sparse_attn_indexer( k_fp8: torch.Tensor, k_scale: torch.Tensor, weights: torch.Tensor, - use_custom_topk: bool = True, q_scale: Optional[torch.Tensor] = None, is_generation: Optional[bool] = None, ) -> torch.Tensor: @@ -1431,8 +1397,6 @@ def sparse_attn_indexer( dtype=torch.int32, capture_graph=metadata.is_cuda_graph, ) - if not use_custom_topk: - topk_indices_buffer[: hidden_states.shape[0]] = -1 if has_prefill and not metadata.skip_indexer_for_ctx_reqs: # Use chunked prefill to reduce memory footprint @@ -1517,36 +1481,14 @@ def sparse_attn_indexer( chunk.cu_seqlen_ks[c0:c1], chunk.cu_seqlen_ke[c0:c1], tile_q_scale, - clean_logits=not use_custom_topk, + clean_logits=False, + ) + self.prefill_top_k( + logits, + chunk.cu_seqlen_ks[c0:c1], + chunk.cu_seqlen_ke[c0:c1], + topk_indices_buffer[g0:g1, :], ) - if use_custom_topk: - torch.ops.trtllm.indexer_topk_prefill( - logits, - chunk.cu_seqlen_ks[c0:c1], - chunk.cu_seqlen_ke[c0:c1], - topk_indices_buffer[g0:g1, :], - self.index_topk, - ) - else: - topk_indices = logits.topk( - min(self.index_topk, logits.shape[-1]), dim=-1 - )[1] - topk_indices -= chunk.cu_seqlen_ks[c0:c1][:, None] - - mask_lo = topk_indices >= 0 - mask_hi = ( - topk_indices - - (chunk.cu_seqlen_ke[c0:c1] - chunk.cu_seqlen_ks[c0:c1])[:, None] - < 0 - ) - mask = mask_lo & mask_hi - - # local indices per sequence - topk_indices = topk_indices.masked_fill(~mask, -1) - - topk_indices_buffer[g0:g1, : topk_indices.shape[-1]] = topk_indices.to( - dtype=torch.int32 - ) if apply_q_split: q_sizes = [ @@ -1577,28 +1519,14 @@ def sparse_attn_indexer( cu_seqlen_ks, cu_seqlen_ke, ctx_q_scale, - clean_logits=not use_custom_topk, + clean_logits=False, + ) + self.prefill_top_k( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + topk_indices_buffer[:num_ctx_tokens, :], ) - if use_custom_topk: - torch.ops.trtllm.indexer_topk_prefill( - logits, - cu_seqlen_ks, - cu_seqlen_ke, - topk_indices_buffer[:num_ctx_tokens, :], - self.index_topk, - ) - else: - topk_indices = logits.topk(min(self.index_topk, logits.shape[-1]), dim=-1)[1] - topk_indices -= cu_seqlen_ks[:, None] - mask_lo = topk_indices >= 0 - mask_hi = topk_indices - (cu_seqlen_ke - cu_seqlen_ks)[:, None] < 0 - mask = mask_lo & mask_hi - - # local indices per sequence - topk_indices = topk_indices.masked_fill(~mask, -1) - topk_indices_buffer[:num_ctx_tokens, : topk_indices.shape[-1]] = ( - topk_indices.to(dtype=torch.int32) - ) elif has_prefill and metadata.skip_indexer_for_ctx_reqs: # Fill topk_indices_buffer with pre-defined dense topk indices topk_indices_buffer[:num_ctx_tokens, :] = metadata.topk_indices_buffer[ @@ -1802,84 +1730,37 @@ def sparse_attn_indexer( decode_q_scale, ) - if use_custom_topk: - # Kernel expects kv_lens (total cache length), not seq_lens (new tokens) - # This is because rowEnd = seq_len - next_n + offset + 1 - gen_kv_lens_cuda = metadata.kv_lens_cuda_runtime[ - num_contexts : num_contexts + num_generations - ] + # The native TRTLLM and GVR paths consume logical KV lengths and + # apply compress_ratio internally. CuTe radix and the Torch oracle + # consume lengths in the score-column coordinate system. + gen_kv_lens_cuda = metadata.kv_lens_cuda_runtime[ + num_contexts : num_contexts + num_generations + ] + scan_lengths = metadata.gen_indexer_kv_lens_cuda_runtime + assert scan_lengths is not None - pre_idx = None - heuristic_scratch = None - if self._enable_heuristic_topk: - local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] - # Pass prev_topk directly; the +1 temporal offset is - # handled inside the C++ kernel (preIdxOffset += 1). - pre_idx = metadata.heuristic_prev_topk[local_layer, :num_generations] - if not metadata.use_cute_dsl_topk: - heuristic_scratch = metadata.heuristic_scratch_values[:num_gen_tokens] - - if self.use_cute_dsl_topk and self._enable_heuristic_topk: - torch.ops.trtllm.cute_dsl_gvr_topk_decode( - logits_decode, - pre_idx, - gen_kv_lens_cuda, - topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], - self.index_topk, - next_n=next_n, - compress_ratio=self.compress_ratio, - max_seq_len=indexer_max_seq_len, - order_row=metadata.kv_lens_row_reorder, - ) - elif self.use_cute_dsl_topk and (self.compress_ratio == 1 or next_n == 1): - torch.ops.trtllm.cute_dsl_indexer_topk_decode( - logits_decode, - context_lens if self.compress_ratio > 1 else gen_kv_lens_cuda, - topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], - self.index_topk, - next_n, - ) - else: - torch.ops.trtllm.indexer_topk_decode( - logits_decode, - gen_kv_lens_cuda, - topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], - next_n, - self.index_topk, - pre_idx=pre_idx, - heuristic_scratch=heuristic_scratch, - compress_ratio=self.compress_ratio, - radix_aux_indices=metadata.radix_aux_indices, - radix_aux_logits=metadata.radix_aux_logits, - ) - else: - # padded - positions = ( - torch.arange(logits_decode.shape[-1], device=q_decode.device) - .unsqueeze(0) - .expand(num_gen_tokens, -1) - ) - row_indices = torch.arange(num_gen_tokens, device=q_decode.device) // next_n - next_n_offset = torch.arange(num_gen_tokens, device=q_decode.device) % next_n - index_end_pos = (context_lens[row_indices] - next_n + next_n_offset).unsqueeze(1) - # index_end_pos: [B * N, 1] - mask = positions <= index_end_pos - # mask: [B * N, L] - logits_decode = logits_decode.masked_fill(~mask, float("-inf")) - topk_indices_decode = logits_decode.topk( - min(self.index_topk, logits_decode.shape[-1]), dim=-1 - )[1].to(torch.int32) # [B * N, K] - # ensure we don't set indices for the top k - # that is out of range(masked already) - # this will happen if context length is shorter than K - mask_decode = topk_indices_decode <= index_end_pos - - # local indices per sequence - topk_indices_decode = topk_indices_decode.masked_fill(~mask_decode, -1) - # Store in buffer - topk_indices_buffer[ - token_offset : token_offset + num_gen_tokens, : topk_indices_decode.shape[-1] - ] = topk_indices_decode.to(dtype=torch.int32) + prior_indices = None + heuristic_values = None + if self._enable_heuristic_topk: + local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] + # The +1 temporal offset is handled by the native kernels. + prior_indices = metadata.heuristic_prev_topk[local_layer, :num_generations] + if not metadata.use_cute_dsl_topk: + heuristic_values = metadata.heuristic_scratch_values[:num_gen_tokens] + + self.decode_top_k( + logits_decode, + gen_kv_lens_cuda, + scan_lengths, + topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], + next_n=next_n, + prior_indices=prior_indices, + heuristic_values=heuristic_values, + radix_indices=metadata.radix_aux_indices, + radix_values=metadata.radix_aux_logits, + max_num_columns=indexer_max_seq_len, + row_order=getattr(metadata, "kv_lens_row_reorder", None), + ) if self._enable_heuristic_topk: local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 998366944daf..d2815ae8514b 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -13,6 +13,7 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.modules.top_k import DecodeTopK, DecodeTopKPolicy from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import get_sm_version, prefer_pinned from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata @@ -347,17 +348,16 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: # use it there, drop this guard so the case is pre-compiled too. if self._indexer_compress_ratio > 1 and next_n > 1: return - try: - from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( - warmup_cute_dsl_radix_topk_decode, - ) - except ImportError: - return - warmup_cute_dsl_radix_topk_decode( - top_k=int(top_k), - num_cols=int(self.get_indexer_max_seq_len()), + decode_top_k = DecodeTopK( + int(top_k), + DecodeTopKPolicy.CUTE_DSL_PREFERRED, + compress_ratio=self._indexer_compress_ratio, + ) + decode_top_k.prepare( + device=self.kv_lens_cuda.device, + max_num_columns=int(self.get_indexer_max_seq_len()), next_n=next_n, - dtype=_INDEXER_LOGITS_DTYPE, + input_dtype=_INDEXER_LOGITS_DTYPE, num_sms=self.num_sms, ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 06ce761b28e1..6b8799085ff5 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -36,6 +36,7 @@ from tensorrt_llm.logger import logger from ...distributed import allgather +from ...modules.top_k import DecodeTopK, DecodeTopKPolicy from ...pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from ...pyexecutor.llm_request import LlmRequestState from ...pyexecutor.resource_manager import KVCacheCompressionManager @@ -632,12 +633,12 @@ def _select_kept_ordinals(self, request_count: int) -> None: """Select top-k tokens and settle score ties into kept-ordinal rows.""" rows = request_count * self._selection_rows_per_request # The trailing 1 is next_n: decode scores one query token per request. - torch.ops.trtllm.cute_dsl_indexer_topk_decode( + self._selection_top_k( self._selection_scores_rows[:rows], self._selection_row_lengths[:rows], + self._selection_row_lengths[:rows], self._provisional_rows[:rows], - self.budget, - 1, + next_n=1, ) settle_ties( self._selection_scores_rows, @@ -799,6 +800,10 @@ def _allocate_metadata_buffers( def _allocate_selection_buffers(self, device: torch.device, *, tp_size: int) -> None: """Allocate fixed manager-lifetime TopK inputs and outputs.""" + self._selection_top_k = DecodeTopK( + self.budget, + DecodeTopKPolicy.CUTE_DSL_PREFERRED, + ) request_capacity = self._request_capacity selection_width = self._selection_width_capacity union = self.eviction_mode == "union" diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py new file mode 100644 index 000000000000..f2d685ebe30d --- /dev/null +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -0,0 +1,374 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Reusable index-selection Top-K modules for sparse inference paths.""" + +from __future__ import annotations + +import threading +from enum import Enum + +import torch +import torch.nn as nn + + +class PrefillTopKImplementation(str, Enum): + """Available segmented prefill Top-K implementations.""" + + TORCH = "torch" + TRTLLM = "trtllm" + + +class DecodeTopKPolicy(str, Enum): + """Decode Top-K dispatch policies. + + ``CUTE_DSL_PREFERRED`` uses the CuTe DSL radix implementation when it + supports the runtime shape and falls back to TRTLLM otherwise. + """ + + TORCH = "torch" + TRTLLM = "trtllm" + TRTLLM_HEURISTIC = "trtllm_heuristic" + CUTE_DSL_PREFERRED = "cute_dsl_preferred" + CUTE_DSL_GVR = "cute_dsl_gvr" + + +_PREPARE_LOCK = threading.Lock() +_PREPARED_DECODE_TOP_K: set[tuple[object, ...]] = set() +_HEURISTIC_WARMUP_COLUMNS = 4096 +_RADIX_MAX_BLOCKS_PER_ROW = 10 + + +def _cuda_device(device: torch.device) -> torch.device: + device = torch.device(device) + if device.type != "cuda": + raise ValueError(f"Top-K preparation requires a CUDA device, got {device}") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + return device + + +def _validate_output( + scores: torch.Tensor, + output_indices: torch.Tensor, + top_k: int, +) -> None: + if scores.ndim != 2: + raise ValueError(f"scores must be rank 2, got shape {tuple(scores.shape)}") + if output_indices.ndim != 2: + raise ValueError(f"output_indices must be rank 2, got shape {tuple(output_indices.shape)}") + expected_shape = (scores.shape[0], top_k) + if output_indices.shape != expected_shape: + raise ValueError( + f"output_indices must have shape {expected_shape}, got {tuple(output_indices.shape)}" + ) + if output_indices.dtype != torch.int32: + raise TypeError(f"output_indices must have dtype torch.int32, got {output_indices.dtype}") + if output_indices.device != scores.device: + raise ValueError( + "scores and output_indices must be on the same device, got " + f"{scores.device} and {output_indices.device}" + ) + + +def _validate_lengths( + scores: torch.Tensor, + lengths: torch.Tensor, + name: str, + expected_size: int, +) -> None: + if lengths.ndim != 1 or lengths.shape[0] != expected_size: + raise ValueError(f"{name} must have shape ({expected_size},), got {tuple(lengths.shape)}") + if lengths.dtype != torch.int32: + raise TypeError(f"{name} must have dtype torch.int32, got {lengths.dtype}") + if lengths.device != scores.device: + raise ValueError( + f"scores and {name} must be on the same device, got {scores.device} and {lengths.device}" + ) + + +class PrefillTopK(nn.Module): + """Select row-local Top-K indices from segmented prefill scores.""" + + def __init__( + self, + top_k: int, + implementation: PrefillTopKImplementation, + ) -> None: + super().__init__() + if top_k <= 0: + raise ValueError(f"top_k must be positive, got {top_k}") + self.top_k = top_k + self.implementation = PrefillTopKImplementation(implementation) + + def prepare(self) -> None: + """Prepare the implementation for execution. + + Segmented prefill implementations currently need no explicit warmup. + """ + + def forward( + self, + scores: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + output_indices: torch.Tensor, + ) -> torch.Tensor: + """Write row-local Top-K indices into ``output_indices``. + + Args: + scores: Score matrix with shape ``[num_rows, num_columns]``. + row_starts: Inclusive valid-column starts with shape ``[num_rows]``. + row_ends: Exclusive valid-column ends with shape ``[num_rows]``. + output_indices: Caller-owned int32 output with shape + ``[num_rows, top_k]``. + """ + _validate_output(scores, output_indices, self.top_k) + _validate_lengths(scores, row_starts, "row_starts", scores.shape[0]) + _validate_lengths(scores, row_ends, "row_ends", scores.shape[0]) + + if self.implementation == PrefillTopKImplementation.TRTLLM: + torch.ops.trtllm.indexer_topk_prefill( + scores, + row_starts, + row_ends, + output_indices, + self.top_k, + ) + return output_indices + + output_indices.fill_(-1) + selected_count = min(self.top_k, scores.shape[1]) + if selected_count == 0: + return output_indices + columns = torch.arange(scores.shape[1], device=scores.device).unsqueeze(0) + valid = (columns >= row_starts.unsqueeze(1)) & (columns < row_ends.unsqueeze(1)) + selected = scores.masked_fill(~valid, float("-inf")).topk(selected_count, dim=-1).indices + selected_valid = torch.gather(valid, 1, selected) + selected = selected - row_starts.unsqueeze(1) + selected = selected.masked_fill(~selected_valid, -1) + output_indices[:, :selected_count].copy_(selected.to(torch.int32)) + return output_indices + + +class DecodeTopK(nn.Module): + """Select Top-K indices from decode scores using a fixed dispatch policy.""" + + def __init__( + self, + top_k: int, + policy: DecodeTopKPolicy, + *, + compress_ratio: int = 1, + ) -> None: + super().__init__() + if top_k <= 0: + raise ValueError(f"top_k must be positive, got {top_k}") + if compress_ratio <= 0: + raise ValueError(f"compress_ratio must be positive, got {compress_ratio}") + self.top_k = top_k + self.policy = DecodeTopKPolicy(policy) + self.compress_ratio = compress_ratio + + def prepare( + self, + *, + device: torch.device, + max_num_columns: int, + next_n: int, + input_dtype: torch.dtype, + num_sms: int | None = None, + ) -> None: + """Warm up static implementation state for one deployment shape.""" + if max_num_columns <= 0 or next_n <= 0: + return + if self.policy in ( + DecodeTopKPolicy.TORCH, + DecodeTopKPolicy.TRTLLM, + DecodeTopKPolicy.CUTE_DSL_GVR, + ): + return + if ( + self.policy == DecodeTopKPolicy.CUTE_DSL_PREFERRED + and self.compress_ratio > 1 + and next_n > 1 + ): + return + + device = _cuda_device(device) + key = ( + self.policy, + device.index, + input_dtype, + self.top_k, + max_num_columns, + next_n, + num_sms, + self.compress_ratio, + ) + with _PREPARE_LOCK: + if key in _PREPARED_DECODE_TOP_K: + return + with torch.cuda.device(device): + if self.policy == DecodeTopKPolicy.TRTLLM_HEURISTIC: + self._warmup_trtllm_heuristic(input_dtype) + else: + try: + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + warmup_cute_dsl_radix_topk_decode, + ) + except ImportError: + return + warmup_cute_dsl_radix_topk_decode( + top_k=self.top_k, + num_cols=max_num_columns, + next_n=next_n, + dtype=input_dtype, + num_sms=num_sms, + ) + _PREPARED_DECODE_TOP_K.add(key) + + def _warmup_trtllm_heuristic(self, input_dtype: torch.dtype) -> None: + num_columns = max(_HEURISTIC_WARMUP_COLUMNS, self.top_k) + device = torch.device("cuda") + scores = torch.zeros((1, num_columns), dtype=input_dtype, device=device) + sequence_lengths = torch.tensor([num_columns], dtype=torch.int32, device=device) + output_indices = torch.empty((1, self.top_k), dtype=torch.int32, device=device) + prior_indices = torch.zeros((1, self.top_k), dtype=torch.int32, device=device) + heuristic_values = torch.empty((1, self.top_k), dtype=input_dtype, device=device) + radix_indices = torch.empty( + (1, _RADIX_MAX_BLOCKS_PER_ROW, self.top_k), + dtype=torch.int32, + device=device, + ) + radix_values = torch.empty( + (1, _RADIX_MAX_BLOCKS_PER_ROW, self.top_k), + dtype=torch.float32, + device=device, + ) + torch.ops.trtllm.indexer_topk_decode( + scores, + sequence_lengths, + output_indices, + 1, + self.top_k, + pre_idx=prior_indices, + heuristic_scratch=heuristic_values, + radix_aux_indices=radix_indices, + radix_aux_logits=radix_values, + ) + torch.cuda.synchronize() + + def forward( + self, + scores: torch.Tensor, + sequence_lengths: torch.Tensor, + scan_lengths: torch.Tensor, + output_indices: torch.Tensor, + *, + next_n: int, + prior_indices: torch.Tensor | None = None, + heuristic_values: torch.Tensor | None = None, + radix_indices: torch.Tensor | None = None, + radix_values: torch.Tensor | None = None, + max_num_columns: int | None = None, + row_order: torch.Tensor | None = None, + ) -> torch.Tensor: + """Write decode Top-K indices into ``output_indices``. + + ``sequence_lengths`` are logical request KV lengths. ``scan_lengths`` + are request lengths in the score-column coordinate system. + """ + _validate_output(scores, output_indices, self.top_k) + if next_n <= 0: + raise ValueError(f"next_n must be positive, got {next_n}") + if scores.shape[0] % next_n != 0: + raise ValueError( + f"score rows ({scores.shape[0]}) must be divisible by next_n ({next_n})" + ) + num_requests = scores.shape[0] // next_n + for name, lengths in ( + ("sequence_lengths", sequence_lengths), + ("scan_lengths", scan_lengths), + ): + _validate_lengths(scores, lengths, name, num_requests) + + if (radix_indices is None) != (radix_values is None): + raise ValueError("radix_indices and radix_values must be provided together") + + if self.policy == DecodeTopKPolicy.TORCH: + return self._forward_torch(scores, scan_lengths, output_indices, next_n) + + use_trtllm = self.policy in ( + DecodeTopKPolicy.TRTLLM, + DecodeTopKPolicy.TRTLLM_HEURISTIC, + ) or ( + self.policy == DecodeTopKPolicy.CUTE_DSL_PREFERRED + and self.compress_ratio > 1 + and next_n > 1 + ) + if use_trtllm: + if self.policy == DecodeTopKPolicy.TRTLLM_HEURISTIC: + if prior_indices is None or heuristic_values is None: + raise ValueError("TRTLLM_HEURISTIC requires prior_indices and heuristic_values") + torch.ops.trtllm.indexer_topk_decode( + scores, + sequence_lengths, + output_indices, + next_n, + self.top_k, + pre_idx=prior_indices, + heuristic_scratch=heuristic_values, + compress_ratio=self.compress_ratio, + radix_aux_indices=radix_indices, + radix_aux_logits=radix_values, + ) + return output_indices + + if self.policy == DecodeTopKPolicy.CUTE_DSL_PREFERRED: + torch.ops.trtllm.cute_dsl_indexer_topk_decode( + scores, + scan_lengths, + output_indices, + self.top_k, + next_n, + ) + return output_indices + + if prior_indices is None: + raise ValueError("CUTE_DSL_GVR requires prior_indices") + if max_num_columns is None: + raise ValueError("CUTE_DSL_GVR requires max_num_columns") + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + scores, + prior_indices, + sequence_lengths, + output_indices, + self.top_k, + next_n=next_n, + compress_ratio=self.compress_ratio, + max_seq_len=max_num_columns, + order_row=row_order, + ) + return output_indices + + def _forward_torch( + self, + scores: torch.Tensor, + scan_lengths: torch.Tensor, + output_indices: torch.Tensor, + next_n: int, + ) -> torch.Tensor: + output_indices.fill_(-1) + selected_count = min(self.top_k, scores.shape[1]) + if selected_count == 0: + return output_indices + positions = torch.arange(scores.shape[1], device=scores.device).unsqueeze(0) + row_indices = torch.arange(scores.shape[0], device=scores.device) // next_n + next_n_offsets = torch.arange(scores.shape[0], device=scores.device) % next_n + row_ends = scan_lengths[row_indices] - next_n + next_n_offsets + 1 + valid = positions < row_ends.unsqueeze(1) + selected = scores.masked_fill(~valid, float("-inf")).topk(selected_count, dim=-1).indices + selected_valid = torch.gather(valid, 1, selected) + selected = selected.masked_fill(~selected_valid, -1) + output_indices[:, :selected_count].copy_(selected.to(torch.int32)) + return output_indices diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 76552749651f..ca48c801463f 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -56,6 +56,12 @@ from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.pyexecutor._util import get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.modules.top_k import ( + DecodeTopK, + DecodeTopKPolicy, + PrefillTopK, + PrefillTopKImplementation, +) from tensorrt_llm._torch.speculative.interface import ( prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, @@ -83,6 +89,18 @@ def has_deep_gemm(): return False +def _set_torch_top_k(indexer: Indexer) -> None: + indexer.prefill_top_k = PrefillTopK( + indexer.index_topk, + PrefillTopKImplementation.TORCH, + ) + indexer.decode_top_k = DecodeTopK( + indexer.index_topk, + DecodeTopKPolicy.TORCH, + compress_ratio=indexer.compress_ratio, + ) + + def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): sparse_config = DeepSeekV4SparseAttentionConfig( compress_ratios=[1, 4, 128], @@ -2660,7 +2678,7 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, compr @pytest.mark.parametrize("seq_len_range", [(2048, 8192), (512, 1024)]) def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_len_range): """ - Test that use_custom_topk=True and use_custom_topk=False produce identical results + Test that the production and Torch Top-K modules produce identical results in the decode phase of sparse_attn_indexer. This test validates: @@ -2806,7 +2824,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l try: topk_indices_custom = indexer.sparse_attn_indexer( - metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights ) except Exception as e: pytest.skip(f"Custom topk not available: {e}") @@ -2830,8 +2848,9 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) + _set_torch_top_k(indexer) topk_indices_fallback = indexer.sparse_attn_indexer( - metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False + metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights ) # Test with indexer skip enabled @@ -2858,7 +2877,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l try: topk_indices_skip = indexer.sparse_attn_indexer( - metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights ) except Exception as e: raise RuntimeError(f"Error when testing indexer skip: {e}") @@ -3009,9 +3028,7 @@ def make_inputs(n_tokens): h0, q0, k0_fp8, k0_scale, w0 = make_inputs(step0_tokens) indexer._update_k_cache(k0_fp8, k0_scale, meta0) try: - topk0 = indexer.sparse_attn_indexer( - meta0, h0, q0, k0_fp8, k0_scale, w0, use_custom_topk=True - ) + topk0 = indexer.sparse_attn_indexer(meta0, h0, q0, k0_fp8, k0_scale, w0) except Exception as e: pytest.skip(f"Custom topk not available: {e}") @@ -3049,7 +3066,7 @@ def make_inputs(n_tokens): meta.indexer_skip_topk = True hs, qs, ks_fp8, ks_scale, ws = make_inputs(batch_size) indexer._update_k_cache(ks_fp8, ks_scale, meta) - topk = indexer.sparse_attn_indexer(meta, hs, qs, ks_fp8, ks_scale, ws, use_custom_topk=True) + topk = indexer.sparse_attn_indexer(meta, hs, qs, ks_fp8, ks_scale, ws) assert torch.equal(topk, stash[:batch_size, :]), ( f"{step0_mode} draft reuse step {step} should copy the stash 1:1 (next_n=1)" ) @@ -3062,7 +3079,7 @@ def make_inputs(n_tokens): @pytest.mark.parametrize("chunk_size", [1024, 2048]) def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chunk_size): """ - Test chunked prefill: use_custom_topk=True vs use_custom_topk=False + Test chunked prefill: production Top-K vs Torch Top-K with metadata.indexer_prefill_chunks != None. This test validates: @@ -3135,7 +3152,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun try: topk_indices_custom = indexer.sparse_attn_indexer( - metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights ) except Exception as e: pytest.skip(f"Custom topk not available: {e}") @@ -3158,8 +3175,9 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) + _set_torch_top_k(indexer) topk_indices_fallback = indexer.sparse_attn_indexer( - metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False + metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights ) # Validation @@ -3179,7 +3197,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun @pytest.mark.parametrize("seq_len_range", [(1, 512)]) def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, seq_len_range): """ - Test single-pass prefill: use_custom_topk=True vs use_custom_topk=False + Test single-pass prefill: production Top-K vs Torch Top-K with metadata.indexer_prefill_chunks == None (else branch). """ torch.manual_seed(42) @@ -3243,7 +3261,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, try: topk_indices_custom = indexer.sparse_attn_indexer( - metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights ) except Exception as e: pytest.skip(f"Custom topk not available: {e}") @@ -3269,8 +3287,9 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, # Force single-pass path by setting indexer_prefill_chunks to None metadata_fallback.indexer_prefill_chunks = None + _set_torch_top_k(indexer) topk_indices_fallback = indexer.sparse_attn_indexer( - metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False + metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights ) # Test with indexer skip enabled @@ -3295,7 +3314,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, try: topk_indices_skip = indexer.sparse_attn_indexer( - metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights ) except Exception as e: raise RuntimeError(f"Indexer skip not available: {e}") @@ -3399,12 +3418,13 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): # Test custom kernel topk_custom = indexer.sparse_attn_indexer( - metadata, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True - ) + metadata, hidden_states, q_fp8, k_fp8, k_scale, weights + ).clone() # Test fallback + _set_torch_top_k(indexer) topk_fallback = indexer.sparse_attn_indexer( - metadata, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False + metadata, hidden_states, q_fp8, k_fp8, k_scale, weights ) # Test with indexer skip enabled @@ -3428,7 +3448,7 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) topk_indices_skip = indexer.sparse_attn_indexer( - metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights ) # Validate: custom and fallback should match diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py new file mode 100644 index 000000000000..7ab7a11150db --- /dev/null +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the reusable sparse index-selection Top-K modules.""" + +from contextlib import nullcontext +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.modules import top_k as top_k_module +from tensorrt_llm._torch.modules.top_k import ( + DecodeTopK, + DecodeTopKPolicy, + PrefillTopK, + PrefillTopKImplementation, +) + + +def test_prefill_torch_masks_dirty_scores_and_pads_output() -> None: + scores = torch.tensor( + [ + [1000.0, 9.0, 8.0, 7.0, 1000.0, 1000.0], + [1000.0, 1000.0, 3.0, 1000.0, 1000.0, 1000.0], + [1000.0, 1000.0, 1000.0, 1000.0, 1000.0, 1000.0], + ] + ) + row_starts = torch.tensor([1, 2, 4], dtype=torch.int32) + row_ends = torch.tensor([4, 3, 4], dtype=torch.int32) + output = torch.full((3, 4), 77, dtype=torch.int32) + + result = PrefillTopK(4, PrefillTopKImplementation.TORCH)( + scores, + row_starts, + row_ends, + output, + ) + + assert result is output + assert output.tolist() == [[0, 1, 2, -1], [0, -1, -1, -1], [-1, -1, -1, -1]] + + +def test_decode_torch_uses_scan_lengths() -> None: + scores = torch.tensor( + [ + [1.0, 2.0, 3.0, 1000.0, 1000.0, 1000.0], + [1.0, 2.0, 3.0, 4.0, 1000.0, 1000.0], + ] + ) + logical_lengths = torch.tensor([16], dtype=torch.int32) + scan_lengths = torch.tensor([3], dtype=torch.int32) + output = torch.full((2, 4), 77, dtype=torch.int32) + + result = DecodeTopK(4, DecodeTopKPolicy.TORCH, compress_ratio=4)( + scores, + logical_lengths, + scan_lengths, + output, + next_n=2, + ) + + assert result is output + assert output.tolist() == [[1, 0, -1, -1], [2, 1, 0, -1]] + + +def test_cute_dsl_preferred_preserves_compressed_mtp_fallback(monkeypatch) -> None: + cute_dsl = Mock() + trtllm = Mock() + monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_indexer_topk_decode", cute_dsl) + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", trtllm) + + top_k = DecodeTopK( + 2, + DecodeTopKPolicy.CUTE_DSL_PREFERRED, + compress_ratio=4, + ) + logical_lengths = torch.tensor([16], dtype=torch.int32) + scan_lengths = torch.tensor([4], dtype=torch.int32) + + scores = torch.randn(1, 4) + output = torch.empty(1, 2, dtype=torch.int32) + top_k(scores, logical_lengths, scan_lengths, output, next_n=1) + cute_dsl.assert_called_once_with(scores, scan_lengths, output, 2, 1) + trtllm.assert_not_called() + + cute_dsl.reset_mock() + scores = torch.randn(2, 4) + output = torch.empty(2, 2, dtype=torch.int32) + radix_indices = torch.empty(2, 10, 2, dtype=torch.int32) + radix_values = torch.empty(2, 10, 2) + top_k( + scores, + logical_lengths, + scan_lengths, + output, + next_n=2, + radix_indices=radix_indices, + radix_values=radix_values, + ) + cute_dsl.assert_not_called() + trtllm.assert_called_once_with( + scores, + logical_lengths, + output, + 2, + 2, + pre_idx=None, + heuristic_scratch=None, + compress_ratio=4, + radix_aux_indices=radix_indices, + radix_aux_logits=radix_values, + ) + + +def test_gvr_routes_logical_lengths_and_workspace(monkeypatch) -> None: + gvr = Mock() + monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) + top_k = DecodeTopK(2, DecodeTopKPolicy.CUTE_DSL_GVR, compress_ratio=4) + scores = torch.randn(1, 8) + logical_lengths = torch.tensor([32], dtype=torch.int32) + scan_lengths = torch.tensor([8], dtype=torch.int32) + output = torch.empty(1, 2, dtype=torch.int32) + prior = torch.zeros(1, 2, dtype=torch.int32) + row_order = torch.zeros(1, dtype=torch.int32) + + top_k( + scores, + logical_lengths, + scan_lengths, + output, + next_n=1, + prior_indices=prior, + max_num_columns=8, + row_order=row_order, + ) + + gvr.assert_called_once_with( + scores, + prior, + logical_lengths, + output, + 2, + next_n=1, + compress_ratio=4, + max_seq_len=8, + order_row=row_order, + ) + + +@pytest.mark.parametrize("top_k", [0, -1]) +def test_top_k_must_be_positive(top_k: int) -> None: + with pytest.raises(ValueError, match="top_k must be positive"): + PrefillTopK(top_k, PrefillTopKImplementation.TORCH) + with pytest.raises(ValueError, match="top_k must be positive"): + DecodeTopK(top_k, DecodeTopKPolicy.TORCH) + + +def test_decode_validates_workspace_pairs() -> None: + top_k = DecodeTopK(2, DecodeTopKPolicy.TRTLLM) + with pytest.raises(ValueError, match="must be provided together"): + top_k( + torch.randn(1, 4), + torch.tensor([4], dtype=torch.int32), + torch.tensor([4], dtype=torch.int32), + torch.empty(1, 2, dtype=torch.int32), + next_n=1, + radix_indices=torch.empty(1, 10, 2, dtype=torch.int32), + ) + + +def test_prepare_deduplicates_success_for_full_key(monkeypatch) -> None: + prepared: set[tuple[object, ...]] = set() + monkeypatch.setattr(top_k_module, "_PREPARED_DECODE_TOP_K", prepared) + monkeypatch.setattr(top_k_module, "_cuda_device", lambda _: torch.device("cuda:3")) + monkeypatch.setattr(torch.cuda, "device", lambda _: nullcontext()) + + top_k = DecodeTopK(32, DecodeTopKPolicy.TRTLLM_HEURISTIC) + warmup = Mock() + monkeypatch.setattr(top_k, "_warmup_trtllm_heuristic", warmup) + + prepare_args = dict( + device=torch.device("cuda:3"), + max_num_columns=4096, + next_n=1, + input_dtype=torch.bfloat16, + num_sms=148, + ) + top_k.prepare(**prepare_args) + top_k.prepare(**prepare_args) + + warmup.assert_called_once_with(torch.bfloat16) + assert len(prepared) == 1 + + +def test_prepare_does_not_cache_failure(monkeypatch) -> None: + prepared: set[tuple[object, ...]] = set() + monkeypatch.setattr(top_k_module, "_PREPARED_DECODE_TOP_K", prepared) + monkeypatch.setattr(top_k_module, "_cuda_device", lambda _: torch.device("cuda:0")) + monkeypatch.setattr(torch.cuda, "device", lambda _: nullcontext()) + + top_k = DecodeTopK(16, DecodeTopKPolicy.TRTLLM_HEURISTIC) + warmup = Mock(side_effect=RuntimeError("warmup failed")) + monkeypatch.setattr(top_k, "_warmup_trtllm_heuristic", warmup) + prepare_args = dict( + device=torch.device("cuda:0"), + max_num_columns=1024, + next_n=1, + input_dtype=torch.float32, + ) + + with pytest.raises(RuntimeError, match="warmup failed"): + top_k.prepare(**prepare_args) + assert not prepared + + warmup.side_effect = None + top_k.prepare(**prepare_args) + assert warmup.call_count == 2 + assert len(prepared) == 1 From db56dd271431d68e6dc0a6ddeefb2de10fe3a523 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:31:55 +0000 Subject: [PATCH 02/18] [None][refactor] unify sparse top-k module Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 53 ++-- .../attention_backend/sparse/dsa/metadata.py | 8 +- .../triattention/triattention.py | 11 +- tensorrt_llm/_torch/modules/top_k.py | 245 +++++++++++------- .../attention/sparse/dsa/test_dsa_indexer.py | 33 ++- tests/unittest/_torch/modules/test_top_k.py | 104 +++++--- 6 files changed, 279 insertions(+), 175 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 8760affd4042..82a6e90d39db 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -25,12 +25,7 @@ maybe_execute_in_parallel, ) from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding -from tensorrt_llm._torch.modules.top_k import ( - DecodeTopK, - DecodeTopKPolicy, - PrefillTopK, - PrefillTopKImplementation, -) +from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation from tensorrt_llm._torch.utils import Fp4QuantizedTensor, maybe_compile from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned from tensorrt_llm.deep_gemm import ( @@ -642,30 +637,27 @@ def __init__( ) self.mtp_index_share = sparse_params.mtp_index_share - self.prefill_top_k = PrefillTopK( - self.index_topk, - PrefillTopKImplementation.TRTLLM, - ) if self.use_cute_dsl_topk: - decode_top_k_policy = ( - DecodeTopKPolicy.CUTE_DSL_GVR + decode_top_k_implementation = ( + TopKImplementation.CUTE_DSL_GVR if self._enable_heuristic_topk - else DecodeTopKPolicy.CUTE_DSL_PREFERRED + else TopKImplementation.CUTE_DSL_PREFERRED ) elif self._enable_heuristic_topk: - decode_top_k_policy = DecodeTopKPolicy.TRTLLM_HEURISTIC + decode_top_k_implementation = TopKImplementation.TRTLLM_HEURISTIC else: - decode_top_k_policy = DecodeTopKPolicy.TRTLLM - self.decode_top_k = DecodeTopK( + decode_top_k_implementation = TopKImplementation.TRTLLM + self.top_k = TopK( self.index_topk, - decode_top_k_policy, + prefill_implementation=TopKImplementation.TRTLLM, + decode_implementation=decode_top_k_implementation, compress_ratio=self.compress_ratio, ) - if decode_top_k_policy == DecodeTopKPolicy.TRTLLM_HEURISTIC and layer_idx == 0: + if decode_top_k_implementation == TopKImplementation.TRTLLM_HEURISTIC and layer_idx == 0: # Populate static caches inside the C++ dispatcher before CUDA - # Graph capture. DecodeTopK globally de-duplicates this warmup. - self.decode_top_k.prepare( + # Graph capture. TopK globally de-duplicates this warmup. + self.top_k.prepare( device=torch.device("cuda", torch.cuda.current_device()), max_num_columns=4096, next_n=1, @@ -1483,11 +1475,12 @@ def sparse_attn_indexer( tile_q_scale, clean_logits=False, ) - self.prefill_top_k( + self.top_k( logits, - chunk.cu_seqlen_ks[c0:c1], - chunk.cu_seqlen_ke[c0:c1], topk_indices_buffer[g0:g1, :], + is_prefill=True, + row_starts=chunk.cu_seqlen_ks[c0:c1], + row_ends=chunk.cu_seqlen_ke[c0:c1], ) if apply_q_split: @@ -1521,11 +1514,12 @@ def sparse_attn_indexer( ctx_q_scale, clean_logits=False, ) - self.prefill_top_k( + self.top_k( logits, - cu_seqlen_ks, - cu_seqlen_ke, topk_indices_buffer[:num_ctx_tokens, :], + is_prefill=True, + row_starts=cu_seqlen_ks, + row_ends=cu_seqlen_ke, ) elif has_prefill and metadata.skip_indexer_for_ctx_reqs: # Fill topk_indices_buffer with pre-defined dense topk indices @@ -1748,11 +1742,12 @@ def sparse_attn_indexer( if not metadata.use_cute_dsl_topk: heuristic_values = metadata.heuristic_scratch_values[:num_gen_tokens] - self.decode_top_k( + self.top_k( logits_decode, - gen_kv_lens_cuda, - scan_lengths, topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], + is_prefill=False, + sequence_lengths=gen_kv_lens_cuda, + scan_lengths=scan_lengths, next_n=next_n, prior_indices=prior_indices, heuristic_values=heuristic_values, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index d2815ae8514b..88ab09cfa325 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -13,7 +13,7 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from tensorrt_llm._torch.modules.top_k import DecodeTopK, DecodeTopKPolicy +from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import get_sm_version, prefer_pinned from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata @@ -348,12 +348,12 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: # use it there, drop this guard so the case is pre-compiled too. if self._indexer_compress_ratio > 1 and next_n > 1: return - decode_top_k = DecodeTopK( + top_k_module = TopK( int(top_k), - DecodeTopKPolicy.CUTE_DSL_PREFERRED, + decode_implementation=TopKImplementation.CUTE_DSL_PREFERRED, compress_ratio=self._indexer_compress_ratio, ) - decode_top_k.prepare( + top_k_module.prepare( device=self.kv_lens_cuda.device, max_num_columns=int(self.get_indexer_max_seq_len()), next_n=next_n, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 6b8799085ff5..0c9e65e28628 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -36,7 +36,7 @@ from tensorrt_llm.logger import logger from ...distributed import allgather -from ...modules.top_k import DecodeTopK, DecodeTopKPolicy +from ...modules.top_k import TopK, TopKImplementation from ...pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from ...pyexecutor.llm_request import LlmRequestState from ...pyexecutor.resource_manager import KVCacheCompressionManager @@ -635,9 +635,10 @@ def _select_kept_ordinals(self, request_count: int) -> None: # The trailing 1 is next_n: decode scores one query token per request. self._selection_top_k( self._selection_scores_rows[:rows], - self._selection_row_lengths[:rows], - self._selection_row_lengths[:rows], self._provisional_rows[:rows], + is_prefill=False, + sequence_lengths=self._selection_row_lengths[:rows], + scan_lengths=self._selection_row_lengths[:rows], next_n=1, ) settle_ties( @@ -800,9 +801,9 @@ def _allocate_metadata_buffers( def _allocate_selection_buffers(self, device: torch.device, *, tp_size: int) -> None: """Allocate fixed manager-lifetime TopK inputs and outputs.""" - self._selection_top_k = DecodeTopK( + self._selection_top_k = TopK( self.budget, - DecodeTopKPolicy.CUTE_DSL_PREFERRED, + decode_implementation=TopKImplementation.CUTE_DSL_PREFERRED, ) request_capacity = self._request_capacity selection_width = self._selection_width_capacity diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index f2d685ebe30d..358e57eb5267 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Reusable index-selection Top-K modules for sparse inference paths.""" +"""Reusable index-selection Top-K module for sparse inference paths.""" from __future__ import annotations @@ -11,15 +11,8 @@ import torch.nn as nn -class PrefillTopKImplementation(str, Enum): - """Available segmented prefill Top-K implementations.""" - - TORCH = "torch" - TRTLLM = "trtllm" - - -class DecodeTopKPolicy(str, Enum): - """Decode Top-K dispatch policies. +class TopKImplementation(str, Enum): + """Top-K implementations for prefill and decode. ``CUTE_DSL_PREFERRED`` uses the CuTe DSL radix implementation when it supports the runtime shape and falls back to TRTLLM otherwise. @@ -32,6 +25,12 @@ class DecodeTopKPolicy(str, Enum): CUTE_DSL_GVR = "cute_dsl_gvr" +_PREFILL_IMPLEMENTATIONS = { + TopKImplementation.TORCH, + TopKImplementation.TRTLLM, +} + + _PREPARE_LOCK = threading.Lock() _PREPARED_DECODE_TOP_K: set[tuple[object, ...]] = set() _HEURISTIC_WARMUP_COLUMNS = 4096 @@ -86,78 +85,42 @@ def _validate_lengths( ) -class PrefillTopK(nn.Module): - """Select row-local Top-K indices from segmented prefill scores.""" - - def __init__( - self, - top_k: int, - implementation: PrefillTopKImplementation, - ) -> None: - super().__init__() - if top_k <= 0: - raise ValueError(f"top_k must be positive, got {top_k}") - self.top_k = top_k - self.implementation = PrefillTopKImplementation(implementation) - - def prepare(self) -> None: - """Prepare the implementation for execution. +def _require_tensor(tensor: torch.Tensor | None, name: str) -> torch.Tensor: + if tensor is None: + raise ValueError(f"{name} is required") + return tensor - Segmented prefill implementations currently need no explicit warmup. - """ - def forward( - self, - scores: torch.Tensor, - row_starts: torch.Tensor, - row_ends: torch.Tensor, - output_indices: torch.Tensor, - ) -> torch.Tensor: - """Write row-local Top-K indices into ``output_indices``. - - Args: - scores: Score matrix with shape ``[num_rows, num_columns]``. - row_starts: Inclusive valid-column starts with shape ``[num_rows]``. - row_ends: Exclusive valid-column ends with shape ``[num_rows]``. - output_indices: Caller-owned int32 output with shape - ``[num_rows, top_k]``. - """ - _validate_output(scores, output_indices, self.top_k) - _validate_lengths(scores, row_starts, "row_starts", scores.shape[0]) - _validate_lengths(scores, row_ends, "row_ends", scores.shape[0]) - - if self.implementation == PrefillTopKImplementation.TRTLLM: - torch.ops.trtllm.indexer_topk_prefill( - scores, - row_starts, - row_ends, - output_indices, - self.top_k, - ) - return output_indices - - output_indices.fill_(-1) - selected_count = min(self.top_k, scores.shape[1]) - if selected_count == 0: - return output_indices - columns = torch.arange(scores.shape[1], device=scores.device).unsqueeze(0) - valid = (columns >= row_starts.unsqueeze(1)) & (columns < row_ends.unsqueeze(1)) - selected = scores.masked_fill(~valid, float("-inf")).topk(selected_count, dim=-1).indices - selected_valid = torch.gather(valid, 1, selected) - selected = selected - row_starts.unsqueeze(1) - selected = selected.masked_fill(~selected_valid, -1) - output_indices[:, :selected_count].copy_(selected.to(torch.int32)) +def _forward_prefill_torch( + scores: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + output_indices: torch.Tensor, + top_k: int, +) -> torch.Tensor: + output_indices.fill_(-1) + selected_count = min(top_k, scores.shape[1]) + if selected_count == 0: return output_indices + columns = torch.arange(scores.shape[1], device=scores.device).unsqueeze(0) + valid = (columns >= row_starts.unsqueeze(1)) & (columns < row_ends.unsqueeze(1)) + selected = scores.masked_fill(~valid, float("-inf")).topk(selected_count, dim=-1).indices + selected_valid = torch.gather(valid, 1, selected) + selected = selected - row_starts.unsqueeze(1) + selected = selected.masked_fill(~selected_valid, -1) + output_indices[:, :selected_count].copy_(selected.to(torch.int32)) + return output_indices -class DecodeTopK(nn.Module): - """Select Top-K indices from decode scores using a fixed dispatch policy.""" +class TopK(nn.Module): + """Select Top-K indices for sparse prefill and decode paths.""" def __init__( self, top_k: int, - policy: DecodeTopKPolicy, *, + prefill_implementation: TopKImplementation | None = None, + decode_implementation: TopKImplementation | None = None, compress_ratio: int = 1, ) -> None: super().__init__() @@ -165,8 +128,24 @@ def __init__( raise ValueError(f"top_k must be positive, got {top_k}") if compress_ratio <= 0: raise ValueError(f"compress_ratio must be positive, got {compress_ratio}") + if prefill_implementation is None and decode_implementation is None: + raise ValueError("at least one Top-K implementation must be configured") self.top_k = top_k - self.policy = DecodeTopKPolicy(policy) + self.prefill_implementation = ( + TopKImplementation(prefill_implementation) + if prefill_implementation is not None + else None + ) + if ( + self.prefill_implementation is not None + and self.prefill_implementation not in _PREFILL_IMPLEMENTATIONS + ): + raise ValueError( + f"{self.prefill_implementation.value} is not supported for prefill Top-K" + ) + self.decode_implementation = ( + TopKImplementation(decode_implementation) if decode_implementation is not None else None + ) self.compress_ratio = compress_ratio def prepare( @@ -178,17 +157,20 @@ def prepare( input_dtype: torch.dtype, num_sms: int | None = None, ) -> None: - """Warm up static implementation state for one deployment shape.""" + """Warm up decode implementation state for one deployment shape.""" + implementation = self.decode_implementation + if implementation is None: + raise ValueError("decode Top-K is not configured") if max_num_columns <= 0 or next_n <= 0: return - if self.policy in ( - DecodeTopKPolicy.TORCH, - DecodeTopKPolicy.TRTLLM, - DecodeTopKPolicy.CUTE_DSL_GVR, + if implementation in ( + TopKImplementation.TORCH, + TopKImplementation.TRTLLM, + TopKImplementation.CUTE_DSL_GVR, ): return if ( - self.policy == DecodeTopKPolicy.CUTE_DSL_PREFERRED + implementation == TopKImplementation.CUTE_DSL_PREFERRED and self.compress_ratio > 1 and next_n > 1 ): @@ -196,7 +178,7 @@ def prepare( device = _cuda_device(device) key = ( - self.policy, + implementation, device.index, input_dtype, self.top_k, @@ -209,7 +191,7 @@ def prepare( if key in _PREPARED_DECODE_TOP_K: return with torch.cuda.device(device): - if self.policy == DecodeTopKPolicy.TRTLLM_HEURISTIC: + if implementation == TopKImplementation.TRTLLM_HEURISTIC: self._warmup_trtllm_heuristic(input_dtype) else: try: @@ -261,11 +243,14 @@ def _warmup_trtllm_heuristic(self, input_dtype: torch.dtype) -> None: def forward( self, scores: torch.Tensor, - sequence_lengths: torch.Tensor, - scan_lengths: torch.Tensor, output_indices: torch.Tensor, *, - next_n: int, + is_prefill: bool, + row_starts: torch.Tensor | None = None, + row_ends: torch.Tensor | None = None, + sequence_lengths: torch.Tensor | None = None, + scan_lengths: torch.Tensor | None = None, + next_n: int = 1, prior_indices: torch.Tensor | None = None, heuristic_values: torch.Tensor | None = None, radix_indices: torch.Tensor | None = None, @@ -273,12 +258,80 @@ def forward( max_num_columns: int | None = None, row_order: torch.Tensor | None = None, ) -> torch.Tensor: - """Write decode Top-K indices into ``output_indices``. + """Write prefill or decode Top-K indices into ``output_indices``. - ``sequence_lengths`` are logical request KV lengths. ``scan_lengths`` - are request lengths in the score-column coordinate system. + Prefill uses ``row_starts`` and ``row_ends``. Decode uses logical + ``sequence_lengths`` plus ``scan_lengths`` in score-column coordinates. """ _validate_output(scores, output_indices, self.top_k) + if is_prefill: + return self._forward_prefill( + scores, + _require_tensor(row_starts, "row_starts"), + _require_tensor(row_ends, "row_ends"), + output_indices, + ) + return self._forward_decode( + scores, + _require_tensor(sequence_lengths, "sequence_lengths"), + _require_tensor(scan_lengths, "scan_lengths"), + output_indices, + next_n=next_n, + prior_indices=prior_indices, + heuristic_values=heuristic_values, + radix_indices=radix_indices, + radix_values=radix_values, + max_num_columns=max_num_columns, + row_order=row_order, + ) + + def _forward_prefill( + self, + scores: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + output_indices: torch.Tensor, + ) -> torch.Tensor: + implementation = self.prefill_implementation + if implementation is None: + raise ValueError("prefill Top-K is not configured") + _validate_lengths(scores, row_starts, "row_starts", scores.shape[0]) + _validate_lengths(scores, row_ends, "row_ends", scores.shape[0]) + if implementation == TopKImplementation.TORCH: + return _forward_prefill_torch( + scores, + row_starts, + row_ends, + output_indices, + self.top_k, + ) + torch.ops.trtllm.indexer_topk_prefill( + scores, + row_starts, + row_ends, + output_indices, + self.top_k, + ) + return output_indices + + def _forward_decode( + self, + scores: torch.Tensor, + sequence_lengths: torch.Tensor, + scan_lengths: torch.Tensor, + output_indices: torch.Tensor, + *, + next_n: int, + prior_indices: torch.Tensor | None, + heuristic_values: torch.Tensor | None, + radix_indices: torch.Tensor | None, + radix_values: torch.Tensor | None, + max_num_columns: int | None, + row_order: torch.Tensor | None, + ) -> torch.Tensor: + implementation = self.decode_implementation + if implementation is None: + raise ValueError("decode Top-K is not configured") if next_n <= 0: raise ValueError(f"next_n must be positive, got {next_n}") if scores.shape[0] % next_n != 0: @@ -295,19 +348,19 @@ def forward( if (radix_indices is None) != (radix_values is None): raise ValueError("radix_indices and radix_values must be provided together") - if self.policy == DecodeTopKPolicy.TORCH: - return self._forward_torch(scores, scan_lengths, output_indices, next_n) + if implementation == TopKImplementation.TORCH: + return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n) - use_trtllm = self.policy in ( - DecodeTopKPolicy.TRTLLM, - DecodeTopKPolicy.TRTLLM_HEURISTIC, + use_trtllm = implementation in ( + TopKImplementation.TRTLLM, + TopKImplementation.TRTLLM_HEURISTIC, ) or ( - self.policy == DecodeTopKPolicy.CUTE_DSL_PREFERRED + implementation == TopKImplementation.CUTE_DSL_PREFERRED and self.compress_ratio > 1 and next_n > 1 ) if use_trtllm: - if self.policy == DecodeTopKPolicy.TRTLLM_HEURISTIC: + if implementation == TopKImplementation.TRTLLM_HEURISTIC: if prior_indices is None or heuristic_values is None: raise ValueError("TRTLLM_HEURISTIC requires prior_indices and heuristic_values") torch.ops.trtllm.indexer_topk_decode( @@ -324,7 +377,7 @@ def forward( ) return output_indices - if self.policy == DecodeTopKPolicy.CUTE_DSL_PREFERRED: + if implementation == TopKImplementation.CUTE_DSL_PREFERRED: torch.ops.trtllm.cute_dsl_indexer_topk_decode( scores, scan_lengths, @@ -351,7 +404,7 @@ def forward( ) return output_indices - def _forward_torch( + def _forward_decode_torch( self, scores: torch.Tensor, scan_lengths: torch.Tensor, diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index ca48c801463f..650194a455a2 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -56,12 +56,7 @@ from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.pyexecutor._util import get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role -from tensorrt_llm._torch.modules.top_k import ( - DecodeTopK, - DecodeTopKPolicy, - PrefillTopK, - PrefillTopKImplementation, -) +from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation from tensorrt_llm._torch.speculative.interface import ( prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, @@ -90,13 +85,10 @@ def has_deep_gemm(): def _set_torch_top_k(indexer: Indexer) -> None: - indexer.prefill_top_k = PrefillTopK( - indexer.index_topk, - PrefillTopKImplementation.TORCH, - ) - indexer.decode_top_k = DecodeTopK( + indexer.top_k = TopK( indexer.index_topk, - DecodeTopKPolicy.TORCH, + prefill_implementation=TopKImplementation.TORCH, + decode_implementation=TopKImplementation.TORCH, compress_ratio=indexer.compress_ratio, ) @@ -247,6 +239,23 @@ def test_indexer_post_load_weights_caches_fused_weight(): assert not hasattr(indexer, "_weights_transformed") +@skip_pre_hopper +def test_indexer_configures_one_top_k_module(): + sparse_config = DeepSeekSparseAttentionConfig( + index_head_dim=128, + index_n_heads=32, + index_topk=128, + ) + + indexer = create_indexer(sparse_config) + + assert isinstance(indexer.top_k, TopK) + assert indexer.top_k.prefill_implementation == TopKImplementation.TRTLLM + assert indexer.top_k.decode_implementation == TopKImplementation.TRTLLM + assert not hasattr(indexer, "prefill_top_k") + assert not hasattr(indexer, "decode_top_k") + + def _ceil_to_ue8m0(x: torch.Tensor): """Round tensor values up to the nearest power of two (UE8M0 format).""" return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 7ab7a11150db..3eabb6f46261 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for the reusable sparse index-selection Top-K modules.""" +"""Tests for the reusable sparse index-selection Top-K module.""" from contextlib import nullcontext from unittest.mock import Mock @@ -9,12 +9,7 @@ import torch from tensorrt_llm._torch.modules import top_k as top_k_module -from tensorrt_llm._torch.modules.top_k import ( - DecodeTopK, - DecodeTopKPolicy, - PrefillTopK, - PrefillTopKImplementation, -) +from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation def test_prefill_torch_masks_dirty_scores_and_pads_output() -> None: @@ -29,11 +24,12 @@ def test_prefill_torch_masks_dirty_scores_and_pads_output() -> None: row_ends = torch.tensor([4, 3, 4], dtype=torch.int32) output = torch.full((3, 4), 77, dtype=torch.int32) - result = PrefillTopK(4, PrefillTopKImplementation.TORCH)( + result = TopK(4, prefill_implementation=TopKImplementation.TORCH)( scores, - row_starts, - row_ends, output, + is_prefill=True, + row_starts=row_starts, + row_ends=row_ends, ) assert result is output @@ -51,11 +47,16 @@ def test_decode_torch_uses_scan_lengths() -> None: scan_lengths = torch.tensor([3], dtype=torch.int32) output = torch.full((2, 4), 77, dtype=torch.int32) - result = DecodeTopK(4, DecodeTopKPolicy.TORCH, compress_ratio=4)( + result = TopK( + 4, + decode_implementation=TopKImplementation.TORCH, + compress_ratio=4, + )( scores, - logical_lengths, - scan_lengths, output, + is_prefill=False, + sequence_lengths=logical_lengths, + scan_lengths=scan_lengths, next_n=2, ) @@ -63,15 +64,43 @@ def test_decode_torch_uses_scan_lengths() -> None: assert output.tolist() == [[1, 0, -1, -1], [2, 1, 0, -1]] +def test_one_module_dispatches_prefill_and_decode() -> None: + top_k = TopK( + 1, + prefill_implementation=TopKImplementation.TORCH, + decode_implementation=TopKImplementation.TORCH, + ) + scores = torch.tensor([[1.0, 3.0, 2.0]]) + output = torch.empty(1, 1, dtype=torch.int32) + + top_k( + scores, + output, + is_prefill=True, + row_starts=torch.tensor([1], dtype=torch.int32), + row_ends=torch.tensor([3], dtype=torch.int32), + ) + assert output.item() == 0 + + top_k( + scores, + output, + is_prefill=False, + sequence_lengths=torch.tensor([3], dtype=torch.int32), + scan_lengths=torch.tensor([3], dtype=torch.int32), + ) + assert output.item() == 1 + + def test_cute_dsl_preferred_preserves_compressed_mtp_fallback(monkeypatch) -> None: cute_dsl = Mock() trtllm = Mock() monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_indexer_topk_decode", cute_dsl) monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", trtllm) - top_k = DecodeTopK( + top_k = TopK( 2, - DecodeTopKPolicy.CUTE_DSL_PREFERRED, + decode_implementation=TopKImplementation.CUTE_DSL_PREFERRED, compress_ratio=4, ) logical_lengths = torch.tensor([16], dtype=torch.int32) @@ -79,7 +108,14 @@ def test_cute_dsl_preferred_preserves_compressed_mtp_fallback(monkeypatch) -> No scores = torch.randn(1, 4) output = torch.empty(1, 2, dtype=torch.int32) - top_k(scores, logical_lengths, scan_lengths, output, next_n=1) + top_k( + scores, + output, + is_prefill=False, + sequence_lengths=logical_lengths, + scan_lengths=scan_lengths, + next_n=1, + ) cute_dsl.assert_called_once_with(scores, scan_lengths, output, 2, 1) trtllm.assert_not_called() @@ -90,9 +126,10 @@ def test_cute_dsl_preferred_preserves_compressed_mtp_fallback(monkeypatch) -> No radix_values = torch.empty(2, 10, 2) top_k( scores, - logical_lengths, - scan_lengths, output, + is_prefill=False, + sequence_lengths=logical_lengths, + scan_lengths=scan_lengths, next_n=2, radix_indices=radix_indices, radix_values=radix_values, @@ -115,7 +152,11 @@ def test_cute_dsl_preferred_preserves_compressed_mtp_fallback(monkeypatch) -> No def test_gvr_routes_logical_lengths_and_workspace(monkeypatch) -> None: gvr = Mock() monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) - top_k = DecodeTopK(2, DecodeTopKPolicy.CUTE_DSL_GVR, compress_ratio=4) + top_k = TopK( + 2, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + compress_ratio=4, + ) scores = torch.randn(1, 8) logical_lengths = torch.tensor([32], dtype=torch.int32) scan_lengths = torch.tensor([8], dtype=torch.int32) @@ -125,9 +166,10 @@ def test_gvr_routes_logical_lengths_and_workspace(monkeypatch) -> None: top_k( scores, - logical_lengths, - scan_lengths, output, + is_prefill=False, + sequence_lengths=logical_lengths, + scan_lengths=scan_lengths, next_n=1, prior_indices=prior, max_num_columns=8, @@ -150,19 +192,23 @@ def test_gvr_routes_logical_lengths_and_workspace(monkeypatch) -> None: @pytest.mark.parametrize("top_k", [0, -1]) def test_top_k_must_be_positive(top_k: int) -> None: with pytest.raises(ValueError, match="top_k must be positive"): - PrefillTopK(top_k, PrefillTopKImplementation.TORCH) - with pytest.raises(ValueError, match="top_k must be positive"): - DecodeTopK(top_k, DecodeTopKPolicy.TORCH) + TopK(top_k, prefill_implementation=TopKImplementation.TORCH) + + +def test_top_k_requires_an_implementation() -> None: + with pytest.raises(ValueError, match="at least one Top-K implementation"): + TopK(1) def test_decode_validates_workspace_pairs() -> None: - top_k = DecodeTopK(2, DecodeTopKPolicy.TRTLLM) + top_k = TopK(2, decode_implementation=TopKImplementation.TRTLLM) with pytest.raises(ValueError, match="must be provided together"): top_k( torch.randn(1, 4), - torch.tensor([4], dtype=torch.int32), - torch.tensor([4], dtype=torch.int32), torch.empty(1, 2, dtype=torch.int32), + is_prefill=False, + sequence_lengths=torch.tensor([4], dtype=torch.int32), + scan_lengths=torch.tensor([4], dtype=torch.int32), next_n=1, radix_indices=torch.empty(1, 10, 2, dtype=torch.int32), ) @@ -174,7 +220,7 @@ def test_prepare_deduplicates_success_for_full_key(monkeypatch) -> None: monkeypatch.setattr(top_k_module, "_cuda_device", lambda _: torch.device("cuda:3")) monkeypatch.setattr(torch.cuda, "device", lambda _: nullcontext()) - top_k = DecodeTopK(32, DecodeTopKPolicy.TRTLLM_HEURISTIC) + top_k = TopK(32, decode_implementation=TopKImplementation.TRTLLM_HEURISTIC) warmup = Mock() monkeypatch.setattr(top_k, "_warmup_trtllm_heuristic", warmup) @@ -198,7 +244,7 @@ def test_prepare_does_not_cache_failure(monkeypatch) -> None: monkeypatch.setattr(top_k_module, "_cuda_device", lambda _: torch.device("cuda:0")) monkeypatch.setattr(torch.cuda, "device", lambda _: nullcontext()) - top_k = DecodeTopK(16, DecodeTopKPolicy.TRTLLM_HEURISTIC) + top_k = TopK(16, decode_implementation=TopKImplementation.TRTLLM_HEURISTIC) warmup = Mock(side_effect=RuntimeError("warmup failed")) monkeypatch.setattr(top_k, "_warmup_trtllm_heuristic", warmup) prepare_args = dict( From c4f3c1fcb6457c71e08f8e9ee9aeecc18109cbb6 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 5 Aug 2026 12:59:56 +0000 Subject: [PATCH 03/18] [None][refactor] move GVR state into top-k module Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 77 +++-------- .../attention_backend/sparse/dsa/metadata.py | 57 +------- tensorrt_llm/_torch/modules/top_k.py | 130 ++++++++++++++++-- .../attention/sparse/dsa/test_dsa_indexer.py | 52 ------- tests/unittest/_torch/modules/test_top_k.py | 89 ++++++++++-- 5 files changed, 222 insertions(+), 183 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 82a6e90d39db..07cde017bfa6 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -654,16 +654,6 @@ def __init__( compress_ratio=self.compress_ratio, ) - if decode_top_k_implementation == TopKImplementation.TRTLLM_HEURISTIC and layer_idx == 0: - # Populate static caches inside the C++ dispatcher before CUDA - # Graph capture. TopK globally de-duplicates this warmup. - self.top_k.prepare( - device=torch.device("cuda", torch.cuda.current_device()), - max_num_columns=4096, - next_n=1, - input_dtype=torch.float32, - ) - # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM # (populated in cache_derived_state; maps to TF32 tensor cores on Ampere+) self._fused_wk_wp_weight: Optional[torch.Tensor] = None @@ -1364,6 +1354,15 @@ def sparse_attn_indexer( num_tokens = metadata.num_tokens num_gen_tokens = num_tokens - num_ctx_tokens + if self._enable_heuristic_topk and not self.top_k.is_state_prepared: + self.top_k.prepare( + device=hidden_states.device, + max_num_columns=metadata.get_indexer_max_seq_len(), + next_n=1 + metadata.max_draft_tokens, + input_dtype=torch.float32, + num_sms=metadata.num_sms, + max_num_requests=metadata.max_num_sequences, + ) if is_generation is None: has_prefill = num_contexts > 0 has_decode = num_generations > 0 @@ -1527,27 +1526,13 @@ def sparse_attn_indexer( :num_ctx_tokens, : ] - # Prefill→decode GVR handoff: seed each finishing-prefill sequence's - # heuristic_prev_topk slot with its own last-context-token top-K, so - # the FIRST decode step of that sequence gets a warm-started preIdx - # (~60-75% set-overlap with the eventual decode top-K on this - # workload) instead of the all-zero / all-(-1) cold start that the - # default `heuristic_prev_topk.zero_()` initialization leaves behind. - # Without this, GVR P2 secant on decode step 0 runs from a benign - # but uninformative seed (kernel +1 offset on zeros → all indices - # point at compressed-token position 1), wasting iterations. - # Slot convention (mirrors the existing decode write-back at the - # bottom of the decode block): new gens from finishing prefill - # append after currently-active gens, i.e., slots - # [num_generations : num_generations + num_contexts]. + # Seed finishing prefill requests after the active generation slots. if self._enable_heuristic_topk and has_prefill and not metadata.skip_indexer_for_ctx_reqs: - local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] - ctx_seq_lens = metadata.seq_lens[:num_contexts] - # Per-sequence last context-token offset (exclusive cumsum minus 1). - last_ctx_idx = (torch.cumsum(ctx_seq_lens, dim=0) - 1).to(dtype=torch.long) - metadata.heuristic_prev_topk[ - local_layer, num_generations : num_generations + num_contexts - ].copy_(topk_indices_buffer[last_ctx_idx, :]) + self.top_k.seed_from_prefill( + topk_indices_buffer[:num_ctx_tokens], + metadata.seq_lens[:num_contexts], + request_offset=num_generations, + ) reuse_topk = ( self.mtp_index_share @@ -1724,23 +1709,16 @@ def sparse_attn_indexer( decode_q_scale, ) - # The native TRTLLM and GVR paths consume logical KV lengths and - # apply compress_ratio internally. CuTe radix and the Torch oracle - # consume lengths in the score-column coordinate system. + # Native/GVR use logical lengths; radix uses score-column lengths. gen_kv_lens_cuda = metadata.kv_lens_cuda_runtime[ num_contexts : num_contexts + num_generations ] scan_lengths = metadata.gen_indexer_kv_lens_cuda_runtime assert scan_lengths is not None - prior_indices = None heuristic_values = None - if self._enable_heuristic_topk: - local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] - # The +1 temporal offset is handled by the native kernels. - prior_indices = metadata.heuristic_prev_topk[local_layer, :num_generations] - if not metadata.use_cute_dsl_topk: - heuristic_values = metadata.heuristic_scratch_values[:num_gen_tokens] + if self._enable_heuristic_topk and not metadata.use_cute_dsl_topk: + heuristic_values = metadata.heuristic_scratch_values[:num_gen_tokens] self.top_k( logits_decode, @@ -1749,31 +1727,12 @@ def sparse_attn_indexer( sequence_lengths=gen_kv_lens_cuda, scan_lengths=scan_lengths, next_n=next_n, - prior_indices=prior_indices, heuristic_values=heuristic_values, radix_indices=metadata.radix_aux_indices, radix_values=metadata.radix_aux_logits, max_num_columns=indexer_max_seq_len, - row_order=getattr(metadata, "kv_lens_row_reorder", None), ) - if self._enable_heuristic_topk: - local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] - decode_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens] - last_mtp_topk = decode_topk[next_n - 1 :: next_n] - prev_topk_dst = metadata.heuristic_prev_topk[local_layer, :num_generations] - if do_multi_stream() and self.aux_stream is not None: - # Fork the write-back onto the aux stream to overlap with this - # layer's core sparse attention; joined by maybe_join_prev_topk_copy(). - self.prev_topk_copy_events[0].record() - with torch.cuda.stream(self.aux_stream): - self.prev_topk_copy_events[0].wait() - prev_topk_dst.copy_(last_mtp_topk) - self.prev_topk_copy_events[1].record() - self._prev_topk_copy_pending = True - else: - prev_topk_dst.copy_(last_mtp_topk) - elif has_decode and metadata.skip_indexer_for_gen_reqs: # Fill topk_indices_buffer with pre-defined dense topk indices topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = ( diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 88ab09cfa325..34bbe820c37b 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -161,7 +161,6 @@ def __post_init__(self): self.use_cute_dsl_topk = ( sparse_metadata_params.use_cute_dsl_topk and IS_CUTLASS_DSL_AVAILABLE ) - self.kv_lens_row_reorder = None capture_graph = self.is_cuda_graph # Plain DSA has no compression and uses the default [1]. DeepSeek-V4's # metadata params carry the model-specific compression ratios. @@ -477,24 +476,8 @@ def on_update_kv_lens(self): self.scheduler_metadata_buffer_expanded.copy_( scheduler_metadata_buffer_expanded, non_blocking=True ) - self._compute_kv_lens_row_reorder() self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True) - def _compute_kv_lens_row_reorder(self): - """Prepare longest-job-first row order for GVR top-k.""" - next_n = 1 + self.max_draft_tokens - if ( - self.enable_heuristic_topk - and self.use_cute_dsl_topk - and self.num_generations * next_n >= 2 * self.num_sms - ): - gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs] - order = torch.argsort(gen_kv_lens, descending=True).to(torch.int32) - self.kv_lens_row_reorder_buffer[: self.num_generations].copy_(order) - self.kv_lens_row_reorder = self.kv_lens_row_reorder_buffer[: self.num_generations] - else: - self.kv_lens_row_reorder = None - def update_for_spec_dec(self): super().update_for_spec_dec() # host @@ -786,45 +769,19 @@ def create_buffers_for_indexer(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - # Per-layer persistent buffers for heuristic TopK pre_idx. - # Indexed by [local_layer_idx, generation_position, :]. - # The graph captures reads/writes on these stable-address buffers; - # each replay's write becomes the next replay's read (feedback loop). self.enable_heuristic_topk = ( sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100 ) - if self.enable_heuristic_topk: - num_local_layers = self.kv_cache_manager.num_local_layers - self.heuristic_prev_topk = self.get_empty( + if self.enable_heuristic_topk and not self.use_cute_dsl_topk: + # Shared C++ heuristic scratch; per-layer history lives in TopK. + max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) + self.heuristic_scratch_values = self.get_empty( self.cuda_graph_buffers, - (num_local_layers, self.max_num_sequences, self.num_sparse_topk), - cache_name="heuristic_prev_topk", - dtype=torch.int32, + (max_gen_tokens, self.num_sparse_topk), + cache_name="heuristic_scratch_values", + dtype=torch.float32, capture_graph=capture_graph, ) - # Zero-initialize so the first decode step's pre_idx (kernel - # adds +1 offset) points to index 1 — a valid but benign candidate. - # Without this, uninitialized memory produces random hint indices. - self.heuristic_prev_topk.zero_() - # The C++ top-k path needs a stable scratch address. - if not self.use_cute_dsl_topk: - max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) - self.heuristic_scratch_values = self.get_empty( - self.cuda_graph_buffers, - (max_gen_tokens, self.num_sparse_topk), - cache_name="heuristic_scratch_values", - dtype=torch.float32, - capture_graph=capture_graph, - ) - # GVR row order also needs a stable address for CUDA graphs. - if self.use_cute_dsl_topk: - self.kv_lens_row_reorder_buffer = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), - cache_name="kv_lens_row_reorder_buffer", - dtype=torch.int32, - capture_graph=capture_graph, - ) # Persistent scratch for the Radix-split-work indexer path. Re-created # in update_spec_dec_param when max_draft_tokens changes so it stays diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 358e57eb5267..f3f55b1c35b2 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -147,6 +147,14 @@ def __init__( TopKImplementation(decode_implementation) if decode_implementation is not None else None ) self.compress_ratio = compress_ratio + self.register_buffer("_prior_indices", None, persistent=False) + self.register_buffer("_row_order_buffer", None, persistent=False) + self._num_sms: int | None = None + + @property + def is_state_prepared(self) -> bool: + """Whether stateful heuristic decode has persistent buffers.""" + return self._prior_indices is not None def prepare( self, @@ -156,11 +164,14 @@ def prepare( next_n: int, input_dtype: torch.dtype, num_sms: int | None = None, + max_num_requests: int | None = None, ) -> None: - """Warm up decode implementation state for one deployment shape.""" + """Prepare persistent state and warm up one decode deployment shape.""" implementation = self.decode_implementation if implementation is None: raise ValueError("decode Top-K is not configured") + if max_num_requests is not None: + self._prepare_state(device, max_num_requests, num_sms) if max_num_columns <= 0 or next_n <= 0: return if implementation in ( @@ -209,6 +220,53 @@ def prepare( ) _PREPARED_DECODE_TOP_K.add(key) + def _prepare_state( + self, + device: torch.device, + max_num_requests: int, + num_sms: int | None, + ) -> None: + implementation = self.decode_implementation + if implementation not in ( + TopKImplementation.TRTLLM_HEURISTIC, + TopKImplementation.CUTE_DSL_GVR, + ): + return + if max_num_requests <= 0: + raise ValueError(f"max_num_requests must be positive, got {max_num_requests}") + + device = torch.device(device) + prior_shape = (max_num_requests, self.top_k) + if self._prior_indices is None: + self._prior_indices = torch.zeros(prior_shape, dtype=torch.int32, device=device) + elif ( + self._prior_indices.device != device or self._prior_indices.shape[0] < max_num_requests + ): + raise RuntimeError( + "Top-K state was already prepared with an incompatible device or capacity: " + f"got {self._prior_indices.device} {tuple(self._prior_indices.shape)}, " + f"requested {device} {prior_shape}" + ) + + if implementation == TopKImplementation.CUTE_DSL_GVR: + if num_sms is None or num_sms <= 0: + raise ValueError("num_sms must be positive when preparing CUTE_DSL_GVR") + self._num_sms = num_sms + if self._row_order_buffer is None: + self._row_order_buffer = torch.empty( + (max_num_requests,), dtype=torch.int32, device=device + ) + elif ( + self._row_order_buffer.device != device + or self._row_order_buffer.shape[0] < max_num_requests + ): + raise RuntimeError( + "GVR row-order state was already prepared with an incompatible device or " + f"capacity: got {self._row_order_buffer.device} " + f"{tuple(self._row_order_buffer.shape)}, requested {device} " + f"({max_num_requests},)" + ) + def _warmup_trtllm_heuristic(self, input_dtype: torch.dtype) -> None: num_columns = max(_HEURISTIC_WARMUP_COLUMNS, self.top_k) device = torch.device("cuda") @@ -251,12 +309,10 @@ def forward( sequence_lengths: torch.Tensor | None = None, scan_lengths: torch.Tensor | None = None, next_n: int = 1, - prior_indices: torch.Tensor | None = None, heuristic_values: torch.Tensor | None = None, radix_indices: torch.Tensor | None = None, radix_values: torch.Tensor | None = None, max_num_columns: int | None = None, - row_order: torch.Tensor | None = None, ) -> torch.Tensor: """Write prefill or decode Top-K indices into ``output_indices``. @@ -277,12 +333,35 @@ def forward( _require_tensor(scan_lengths, "scan_lengths"), output_indices, next_n=next_n, - prior_indices=prior_indices, heuristic_values=heuristic_values, radix_indices=radix_indices, radix_values=radix_values, max_num_columns=max_num_columns, - row_order=row_order, + ) + + def seed_from_prefill( + self, + output_indices: torch.Tensor, + request_lengths: torch.Tensor, + *, + request_offset: int = 0, + ) -> None: + """Seed decode hints from the last selected row of each prefill request.""" + if self._prior_indices is None: + return + if request_lengths.ndim != 1: + raise ValueError( + f"request_lengths must be rank 1, got shape {tuple(request_lengths.shape)}" + ) + num_requests = request_lengths.shape[0] + if request_offset < 0 or request_offset + num_requests > self._prior_indices.shape[0]: + raise ValueError( + f"request range [{request_offset}, {request_offset + num_requests}) exceeds " + f"prepared capacity {self._prior_indices.shape[0]}" + ) + last_rows = (torch.cumsum(request_lengths, dim=0) - 1).to(dtype=torch.long) + self._prior_indices[request_offset : request_offset + num_requests].copy_( + output_indices[last_rows] ) def _forward_prefill( @@ -322,12 +401,10 @@ def _forward_decode( output_indices: torch.Tensor, *, next_n: int, - prior_indices: torch.Tensor | None, heuristic_values: torch.Tensor | None, radix_indices: torch.Tensor | None, radix_values: torch.Tensor | None, max_num_columns: int | None, - row_order: torch.Tensor | None, ) -> torch.Tensor: implementation = self.decode_implementation if implementation is None: @@ -348,6 +425,15 @@ def _forward_decode( if (radix_indices is None) != (radix_values is None): raise ValueError("radix_indices and radix_values must be provided together") + prior_indices = None + if implementation in ( + TopKImplementation.TRTLLM_HEURISTIC, + TopKImplementation.CUTE_DSL_GVR, + ): + if self._prior_indices is None: + raise RuntimeError("Top-K state must be prepared before heuristic decode") + prior_indices = self._prior_indices[:num_requests] + if implementation == TopKImplementation.TORCH: return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n) @@ -375,6 +461,8 @@ def _forward_decode( radix_aux_indices=radix_indices, radix_aux_logits=radix_values, ) + if implementation == TopKImplementation.TRTLLM_HEURISTIC: + self._update_prior_indices(output_indices, num_requests, next_n) return output_indices if implementation == TopKImplementation.CUTE_DSL_PREFERRED: @@ -387,10 +475,9 @@ def _forward_decode( ) return output_indices - if prior_indices is None: - raise ValueError("CUTE_DSL_GVR requires prior_indices") if max_num_columns is None: raise ValueError("CUTE_DSL_GVR requires max_num_columns") + row_order = self._prepare_row_order(sequence_lengths, next_n) torch.ops.trtllm.cute_dsl_gvr_topk_decode( scores, prior_indices, @@ -402,8 +489,33 @@ def _forward_decode( max_seq_len=max_num_columns, order_row=row_order, ) + self._update_prior_indices(output_indices, num_requests, next_n) return output_indices + def _prepare_row_order( + self, + sequence_lengths: torch.Tensor, + next_n: int, + ) -> torch.Tensor | None: + if self._num_sms is None or self._row_order_buffer is None: + raise RuntimeError("GVR state must be prepared before decode") + num_requests = sequence_lengths.shape[0] + if num_requests * next_n < 2 * self._num_sms: + return None + order = torch.argsort(sequence_lengths, descending=True).to(torch.int32) + row_order = self._row_order_buffer[:num_requests] + row_order.copy_(order) + return row_order + + def _update_prior_indices( + self, + output_indices: torch.Tensor, + num_requests: int, + next_n: int, + ) -> None: + assert self._prior_indices is not None + self._prior_indices[:num_requests].copy_(output_indices[next_n - 1 :: next_n]) + def _forward_decode_torch( self, scores: torch.Tensor, diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 650194a455a2..3436019633be 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -3812,55 +3812,3 @@ def _run_indexer(): assert "indexer_topk_out_buffer" in metadata.cuda_graph_buffers.buffers, ( "indexer topk-output buffer must be drawn from the cuda_graph_buffers arena" ) - - -def test_kv_lens_row_reorder_threshold(): - """_compute_kv_lens_row_reorder engages iff num_generations * next_n >= 2 * num_sms, - and produces a descending argsort of gen_kv_lens when active.""" - num_sms = 16 # small synthetic value; threshold = 2 * 16 = 32 rows - next_n = 2 # max_draft_tokens=1 → next_n = 1 + 1 = 2 - - def make_mock(num_generations, kv_lens_list): - kv_cuda = torch.tensor(kv_lens_list, dtype=torch.int32, device="cuda") - buf = torch.zeros(64, dtype=torch.int32, device="cuda") - ns = SimpleNamespace( - enable_heuristic_topk=True, - use_cute_dsl_topk=True, - num_generations=num_generations, - num_sms=num_sms, - max_draft_tokens=next_n - 1, - num_contexts=0, - num_seqs=num_generations, - kv_lens_cuda=kv_cuda, - kv_lens_row_reorder_buffer=buf, - kv_lens_row_reorder=None, - ) - ns._compute_kv_lens_row_reorder = ( - lambda: DSAtrtllmAttentionMetadata._compute_kv_lens_row_reorder(ns) - ) - return ns - - # Fixed unsorted sequence for deterministic sort verification (len == num_sms) - kv_vals = [4, 1, 8, 2, 16, 3, 12, 6, 7, 9, 5, 11, 13, 10, 14, 15] - - # Below threshold: 1 * 2 = 2 < 32 → None - md_below = make_mock(1, [1000]) - md_below._compute_kv_lens_row_reorder() - assert md_below.kv_lens_row_reorder is None - - # At threshold: num_sms * 2 = 32 → engages, verify descending argsort - md_at = make_mock(num_sms, kv_vals) - md_at._compute_kv_lens_row_reorder() - assert md_at.kv_lens_row_reorder is not None - reorder = md_at.kv_lens_row_reorder.cpu().tolist() - assert [kv_vals[i] for i in reorder] == sorted(kv_vals, reverse=True), ( - "order_row must be a descending argsort of gen_kv_lens" - ) - - # Above threshold: (num_sms + 1) * 2 = 34 > 32 → also engages with correct sort - kv_vals2 = kv_vals + [100] - md_above = make_mock(num_sms + 1, kv_vals2) - md_above._compute_kv_lens_row_reorder() - assert md_above.kv_lens_row_reorder is not None - reorder2 = md_above.kv_lens_row_reorder.cpu().tolist() - assert [kv_vals2[i] for i in reorder2] == sorted(kv_vals2, reverse=True) diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 3eabb6f46261..33e1c570878b 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -149,7 +149,7 @@ def test_cute_dsl_preferred_preserves_compressed_mtp_fallback(monkeypatch) -> No ) -def test_gvr_routes_logical_lengths_and_workspace(monkeypatch) -> None: +def test_gvr_owns_prior_state_and_updates_it(monkeypatch) -> None: gvr = Mock() monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) top_k = TopK( @@ -162,7 +162,17 @@ def test_gvr_routes_logical_lengths_and_workspace(monkeypatch) -> None: scan_lengths = torch.tensor([8], dtype=torch.int32) output = torch.empty(1, 2, dtype=torch.int32) prior = torch.zeros(1, 2, dtype=torch.int32) - row_order = torch.zeros(1, dtype=torch.int32) + top_k.prepare( + device=torch.device("cpu"), + max_num_columns=8, + next_n=1, + input_dtype=torch.float32, + num_sms=16, + max_num_requests=1, + ) + assert top_k._prior_indices is not None + top_k._prior_indices.copy_(prior) + gvr.side_effect = lambda *args, **kwargs: output.copy_(torch.tensor([[5, 3]])) top_k( scores, @@ -171,22 +181,75 @@ def test_gvr_routes_logical_lengths_and_workspace(monkeypatch) -> None: sequence_lengths=logical_lengths, scan_lengths=scan_lengths, next_n=1, - prior_indices=prior, max_num_columns=8, - row_order=row_order, ) - gvr.assert_called_once_with( - scores, - prior, - logical_lengths, - output, - 2, + args, kwargs = gvr.call_args + assert args[0] is scores + assert args[1].data_ptr() == top_k._prior_indices.data_ptr() + assert args[2] is logical_lengths + assert args[3] is output + assert args[4] == 2 + assert kwargs == { + "next_n": 1, + "compress_ratio": 4, + "max_seq_len": 8, + "order_row": None, + } + assert top_k._prior_indices.tolist() == [[5, 3]] + + +def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: + gvr = Mock(side_effect=lambda *args, **kwargs: args[3].zero_()) + monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) + top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) + num_sms = 4 + next_n = 2 + lengths = torch.tensor([4, 1, 8, 2], dtype=torch.int32) + top_k.prepare( + device=torch.device("cpu"), + max_num_columns=8, + next_n=next_n, + input_dtype=torch.float32, + num_sms=num_sms, + max_num_requests=lengths.shape[0], + ) + + top_k( + torch.randn(lengths.shape[0] * next_n, 8), + torch.empty(lengths.shape[0] * next_n, 2, dtype=torch.int32), + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths, + next_n=next_n, + max_num_columns=8, + ) + + row_order = gvr.call_args.kwargs["order_row"] + assert row_order is not None + assert row_order.tolist() == [2, 0, 3, 1] + + +def test_seed_from_prefill_uses_last_request_rows() -> None: + top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) + top_k.prepare( + device=torch.device("cpu"), + max_num_columns=8, next_n=1, - compress_ratio=4, - max_seq_len=8, - order_row=row_order, + input_dtype=torch.float32, + num_sms=4, + max_num_requests=4, ) + prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32) + + top_k.seed_from_prefill( + prefill_indices, + torch.tensor([2, 1], dtype=torch.int32), + request_offset=1, + ) + + assert top_k._prior_indices is not None + assert top_k._prior_indices.tolist() == [[0, 0], [2, 3], [4, 5], [0, 0]] @pytest.mark.parametrize("top_k", [0, -1]) From 77051d1be3d2fd5c7d3fa110b0e29824bb2b350c Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:48:15 +0000 Subject: [PATCH 04/18] [None][refactor] prepare sparse top-k before model forward Separate per-step indexer metadata preparation from per-module TopK preparation, and route the latter through sparse MLA lifecycle hooks before model execution. Remove forward-time lazy preparation and metadata-owned radix warmup plumbing. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../sparse/deepseek_v4/metadata.py | 2 +- .../sparse/deepseek_v4/module.py | 4 + .../attention_backend/sparse/dsa/indexer.py | 28 ++++--- .../attention_backend/sparse/dsa/metadata.py | 66 +--------------- .../attention_backend/sparse/dsa/module.py | 4 + .../_torch/attention_backend/sparse/hooks.py | 3 + tensorrt_llm/_torch/modules/mla.py | 5 ++ tensorrt_llm/_torch/modules/top_k.py | 5 -- .../_torch/pyexecutor/model_engine.py | 22 +++++- .../attention/sparse/dsa/test_dsa_indexer.py | 78 ++++++++++++------- .../attention/sparse/test_dsa_fp4_indexer.py | 2 +- .../attention/sparse/test_sparse_attention.py | 26 +++++++ .../executor/test_pytorch_model_engine.py | 14 ++++ 13 files changed, 147 insertions(+), 112 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py index cfd22a4da152..3cf478436904 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py @@ -576,7 +576,7 @@ def prepare(self): # Prepare metadata for indexer (only needed when sparse layers exist) if has_sparse_layers: - DeepseekV4Indexer.prepare(metadata=self) + DeepseekV4Indexer.prepare_metadata(metadata=self) # --- Per-ratio metadata --- # 1) CPU-side: compute scalar metadata (num_total_compressed_tokens, etc.) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py index 2c09926c9835..f904312ae0b3 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -1243,6 +1243,10 @@ def initialize(self, mla: MLA) -> None: def create_weights(self, mla: MLA) -> None: create_sparse_attn_weights(mla) + def prepare(self, mla: MLA, attn_metadata: AttentionMetadata) -> None: + if mla.indexer is not None: + mla.indexer.prepare(attn_metadata) + def prepare_outputs( self, mla: MLA, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 07cde017bfa6..3f2b11fd0796 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -136,7 +136,7 @@ def _compute_slot_mappings( ) -> Tuple[torch.Tensor, torch.Tensor]: """Compute flat byte indices for FP8/FP4 data and scales from global token positions. - Shared by Indexer.prepare() (CPU) and on_update_kv_lens() (GPU) to avoid + Shared by Indexer.prepare_metadata() (CPU) and on_update_kv_lens() (GPU) to avoid duplicating the slot mapping arithmetic. Args: @@ -669,6 +669,19 @@ def cache_derived_state(self) -> None: def post_load_weights(self) -> None: self.cache_derived_state() + def prepare(self, metadata: DSAtrtllmAttentionMetadata) -> None: + """Prepare this indexer's Top-K state before model forward.""" + if metadata.kv_cache_manager is None: + return + self.top_k.prepare( + device=metadata.kv_lens_cuda.device, + max_num_columns=metadata.get_indexer_max_seq_len(), + next_n=1 + metadata.max_draft_tokens, + input_dtype=torch.float32, + num_sms=metadata.num_sms, + max_num_requests=metadata.max_num_sequences, + ) + @staticmethod def prepare_one_prefill_chunk( metadata: DSAtrtllmAttentionMetadata, @@ -1084,7 +1097,7 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): ) @staticmethod - def prepare(metadata: DSAtrtllmAttentionMetadata): + def prepare_metadata(metadata: DSAtrtllmAttentionMetadata): """ Prepare indexer for the forward pass. This should be called during metadata.prepare() stage. @@ -1354,15 +1367,6 @@ def sparse_attn_indexer( num_tokens = metadata.num_tokens num_gen_tokens = num_tokens - num_ctx_tokens - if self._enable_heuristic_topk and not self.top_k.is_state_prepared: - self.top_k.prepare( - device=hidden_states.device, - max_num_columns=metadata.get_indexer_max_seq_len(), - next_n=1 + metadata.max_draft_tokens, - input_dtype=torch.float32, - num_sms=metadata.num_sms, - max_num_requests=metadata.max_num_sequences, - ) if is_generation is None: has_prefill = num_contexts > 0 has_decode = num_generations > 0 @@ -1601,7 +1605,7 @@ def sparse_attn_indexer( if self.use_cute_dsl_paged_mqa_logits: # DSL kernel design: 1 atom per q (atom = real next_n positions), # kNumNextNAtoms = 1 for any real next_n. The matching schedule - # is `scheduler_metadata_buffer` — built in `Indexer.prepare()` + # is `scheduler_metadata_buffer` — built in `Indexer.prepare_metadata()` # with a (num_gen, 1) input shape, which makes DeepGEMM's wrapper # compute `num_next_n_atoms = 1`. (DeepGEMM uses the same buffer # for its own next_n=1 kernel; DSL piggy-backs on it for all diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 34bbe820c37b..110891f65936 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -13,7 +13,6 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import get_sm_version, prefer_pinned from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata @@ -32,16 +31,6 @@ ModelConfig = tensorrt_llm.bindings.ModelConfig -# dtype of the indexer MQA-logits that feed the top-k. All paged_mqa_logits -# paths produce fp32 today (DSL fp8/fp4 default output_dtype=fp32; DeepGEMM -# fp8 hardcodes kFloat; DeepGEMM fp4 defaults logits_dtype=kFloat32 and is not -# overridden here), and the decode forward feeds logits to the top-k without a -# cast. dtype is a top-k compile-key dimension, so the warmup pre-compiles for -# exactly this value. If a paged_mqa_logits caller ever emits a different dtype -# (e.g. overriding the DeepGEMM fp4 logits_dtype to bf16), update this constant -# or the warmup silently compiles the wrong variant. -_INDEXER_LOGITS_DTYPE = torch.float32 - if TYPE_CHECKING: from tensorrt_llm._torch.speculative.interface import SpecMetadata from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager @@ -59,8 +48,6 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" sparse_metadata_params: Optional[DSAMetadataParams] = None - # Store reference to indexer for preparation stage - indexer: Optional["Indexer"] = None # Chunked prefill metadata for indexer (prefill-only, no CUDA graph needed) indexer_prefill_chunks: Optional[List[IndexerPrefillChunkMetadata]] = None # Max chunk size for two-level chunking: @@ -216,7 +203,7 @@ def prepare(self): self.prepare_for_spec_decode(kv_lens) # Prepare metadata for indexer - Indexer.prepare(metadata=self) + Indexer.prepare_metadata(metadata=self) def prepare_for_draft_forward(self) -> dict | None: """Select native DSA indexer metadata for a draft forward.""" @@ -312,54 +299,6 @@ def get_indexer_max_seq_len(self) -> int: return self.kv_cache_manager.max_seq_len return max(1, self.kv_cache_manager.max_seq_len // self._indexer_compress_ratio) - def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: - """Pre-compile the radix-filter CuTe DSL decode top-k during warmup. - - Eager decode iters (mixed prefill+decode batch, or ``cuda_graph`` - disabled) whose ``num_rows`` lands in a ``cluster_size`` band that - graph capture did not exercise otherwise pay a first-touch JIT stall - on a live request. ``num_cols`` is fixed at ``indexer_max_seq_len``, - so only the ``cluster_size`` dimension needs sweeping; delegate to the - custom-op warmup helper, which owns the band enumeration. - - ``next_n`` (a compile-key dimension) is supplied by the caller from - the engine's static spec-decode config. - - No-op unless decode actually routes to - ``cute_dsl_indexer_topk_decode``: heuristic top-k uses the GVR kernel - and plain (no cute_dsl_topk) decode uses the C++ op. Called once from - ``ModelEngine.warmup``. - """ - if not self.use_cute_dsl_topk or self.enable_heuristic_topk: - return - if self.kv_cache_manager is None: - return - top_k = getattr(self.sparse_metadata_params, "index_topk", None) - if not top_k: - return - # The radix-filter DSL kernel does not support a compressed indexer - # combined with multi-row MTP: decode dispatches to it only when - # compress_ratio == 1 or next_n == 1. The compress_ratio > 1 && - # next_n > 1 case routes to the C++ op (or GVR when heuristic top-k is - # on), so there is nothing to pre-compile here. - # TODO: extending the radix-filter path to compress_ratio > 1 && - # next_n > 1 is straightforward; once the dispatch above is relaxed to - # use it there, drop this guard so the case is pre-compiled too. - if self._indexer_compress_ratio > 1 and next_n > 1: - return - top_k_module = TopK( - int(top_k), - decode_implementation=TopKImplementation.CUTE_DSL_PREFERRED, - compress_ratio=self._indexer_compress_ratio, - ) - top_k_module.prepare( - device=self.kv_lens_cuda.device, - max_num_columns=int(self.get_indexer_max_seq_len()), - next_n=next_n, - input_dtype=_INDEXER_LOGITS_DTYPE, - num_sms=self.num_sms, - ) - def on_update_kv_lens(self): # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas. # Especially for the changes in the _preprocess_inputs() of model_engine.py. @@ -678,7 +617,8 @@ def create_buffers_for_indexer(self, capture_graph=False): pin_memory=prefer_pinned(), ) # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation. - # These buffers will be allocated dynamically in Indexer.prepare() based on actual total_kv_len to save memory. + # Allocate these buffers dynamically in Indexer.prepare_metadata() + # based on the actual total_kv_len to save memory. if self.enable_context_mla_with_cached_kv: self.slot_mapping_fp8_fullkv = None self.slot_mapping_scale_fullkv = None diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py index bf36a14655ed..80498f65b438 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py @@ -595,6 +595,10 @@ def initialize(self, mla: MLA) -> None: self.need_dense_mha and mla.mapping.cp_size == 1 and mla.mqa.support_fused_rope() ) + def prepare(self, mla: MLA, attn_metadata: AttentionMetadata) -> None: + if mla.indexer is not None: + mla.indexer.prepare(attn_metadata) + def forward( self, mla: MLA, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/hooks.py b/tensorrt_llm/_torch/attention_backend/sparse/hooks.py index 6bc1177323aa..12edceabf237 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/hooks.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/hooks.py @@ -56,6 +56,9 @@ def create_weights(self, mla: "MLA") -> None: def transform_weights(self, mla: "MLA") -> None: """Transform algorithm-specific weights.""" + def prepare(self, mla: "MLA", attn_metadata: "AttentionMetadata") -> None: + """Prepare algorithm-specific module state before model forward.""" + def prepare_outputs( self, mla: "MLA", diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 646509bb25da..f486e36dfba3 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -795,6 +795,11 @@ def create_output(self, hidden_states: torch.Tensor, num_contexts: int): [num_tokens, self.attention_output_hidden_size], dtype=hidden_states.dtype ) + def prepare_sparse_attn(self, attn_metadata: AttentionMetadata) -> None: + """Prepare sparse module state before the model forward starts.""" + if self.sparse_attn_hooks is not None: + self.sparse_attn_hooks.prepare(self, attn_metadata) + def _create_outputs( self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata ) -> list[torch.Tensor]: diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index f3f55b1c35b2..cc31f04a7016 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -151,11 +151,6 @@ def __init__( self.register_buffer("_row_order_buffer", None, persistent=False) self._num_sms: int | None = None - @property - def is_state_prepared(self) -> bool: - """Whether stateful heuristic decode has persistent buffers.""" - return self._prior_indices is not None - def prepare( self, *, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 3da7a73a7e9e..4e702ebd5083 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1492,8 +1492,6 @@ def warmup(self, resource_manager: ResourceManager) -> None: # warmup. No-op on non-DSA models. self._warmup_dg_paged_mqa_logits_metadata() log_mem_snapshot("warmup/after_dg_paged_mqa_logits_metadata") - self._warmup_cute_dsl_radix_topk() - log_mem_snapshot("warmup/after_cute_dsl_radix_topk") if can_run_general_warmup: # Pre-populate the memory pool with max-shape allocations to reduce # fragmentation at runtime. @@ -3158,6 +3156,23 @@ def _update_draft_inference_state_for_warmup( req.py_is_first_draft = True req.py_draft_tokens = [] + def _prepare_sparse_attention_modules( + self, attn_metadata: AttentionMetadata) -> None: + """Prepare sparse module state before any model forward uses metadata.""" + prepared_metadata_ids = getattr( + self, "_prepared_sparse_attention_metadata_ids", None) + if prepared_metadata_ids is None: + prepared_metadata_ids = set() + self._prepared_sparse_attention_metadata_ids = prepared_metadata_ids + metadata_id = id(attn_metadata) + if metadata_id in prepared_metadata_ids: + return + for module in self.model.modules(): + prepare = getattr(module, "prepare_sparse_attn", None) + if callable(prepare): + prepare(attn_metadata) + prepared_metadata_ids.add(metadata_id) + def _set_up_attn_metadata( self, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], @@ -3232,7 +3247,6 @@ def _set_up_attn_metadata( num_heads_per_kv=num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, ) - return self.attn_metadata @property @@ -7504,6 +7518,8 @@ def forward(self, spec_resource_manager = None spec_metadata = None + self._prepare_sparse_attention_modules(attn_metadata) + moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', None) if kv_cache_manager is None: diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 3436019633be..b3a3d5220e50 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -256,6 +256,30 @@ def test_indexer_configures_one_top_k_module(): assert not hasattr(indexer, "decode_top_k") +def test_indexer_prepare_delegates_to_top_k_before_forward(): + top_k = Mock() + indexer = SimpleNamespace(top_k=top_k) + metadata = SimpleNamespace( + kv_cache_manager=SimpleNamespace(), + kv_lens_cuda=torch.empty(0), + get_indexer_max_seq_len=Mock(return_value=4096), + max_draft_tokens=3, + num_sms=148, + max_num_sequences=32, + ) + + Indexer.prepare(indexer, metadata) + + top_k.prepare.assert_called_once_with( + device=metadata.kv_lens_cuda.device, + max_num_columns=4096, + next_n=4, + input_dtype=torch.float32, + num_sms=148, + max_num_requests=32, + ) + + def _ceil_to_ue8m0(x: torch.Tensor): """Round tensor values up to the nearest power of two (UE8M0 format).""" return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) @@ -793,7 +817,7 @@ def __init__(self): # so allocate a separate buffer for the full-next_n schedule. # DeepGEMM expects the full-next_n schedule in # `scheduler_metadata_buffer` itself (the alias makes - # `Indexer.prepare()`'s second populate overwrite the first). + # `Indexer.prepare_metadata()`'s second populate overwrite the first). if use_cute_dsl_paged_mqa_logits: self.scheduler_metadata_buffer_full_next_n = torch.zeros( (self.num_sms + 1, 2), device="cuda", dtype=torch.int32 @@ -1026,7 +1050,7 @@ def validate_topk_indices(topk_indices_0, topk_indices_1, total_tokens): @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper def test_recompute_slot_mappings_matches_prepare_with_cached_tokens(): - """Recompute slot mappings without re-running full Indexer.prepare().""" + """Recompute slot mappings without re-running full Indexer.prepare_metadata().""" head_dim = 128 block_size = 64 request_ids = [0, 1] @@ -1060,7 +1084,7 @@ def test_recompute_slot_mappings_matches_prepare_with_cached_tokens(): indexer_head_dim=head_dim, ) - Indexer.prepare(metadata) + Indexer.prepare_metadata(metadata) expected_fp8 = metadata.slot_mapping_fp8[:num_tokens].clone() expected_scale = metadata.slot_mapping_scale[:num_tokens].clone() @@ -1129,7 +1153,7 @@ def test_indexer_k_cache_scatter_custom_op(): from tensorrt_llm._torch.attention_backend.sparse.dsa import Indexer - Indexer.prepare(metadata) + Indexer.prepare_metadata(metadata) # Generate test data k_original = torch.randn((num_tokens, head_dim), device="cuda", dtype=torch.bfloat16) @@ -1290,7 +1314,7 @@ def test_fp8_k_cache_roundtrip(): num_tokens=total_tokens, indexer_head_dim=head_dim, ) - Indexer.prepare(metadata) + Indexer.prepare_metadata(metadata) # Generate unique patterns for each request and quantize k_original = torch.randn((total_tokens, head_dim), device="cuda", dtype=torch.bfloat16) @@ -1447,7 +1471,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, compres compress_ratio=compress_ratio, indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_context) + Indexer.prepare_metadata(metadata_context) k_context_fp8, k_context_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_context_bf16) @@ -1473,7 +1497,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, compres compress_ratio=compress_ratio, indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_gen) + Indexer.prepare_metadata(metadata_gen) k_gen_fp8, k_gen_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_gen_bf16) indexer._update_k_cache(k_gen_fp8, k_gen_scale, metadata_gen) @@ -1809,7 +1833,7 @@ def _force_dsl_expand_setup(meta): ) if not use_dsl: _force_direct_path(metadata_context) - Indexer.prepare(metadata_context) + Indexer.prepare_metadata(metadata_context) # Real path: split K at head_dim//2 + fused_cat_fp4 (mirrors # Indexer._prep_q_or_k at dsa.py:2046-2050). @@ -1843,7 +1867,7 @@ def _force_dsl_expand_setup(meta): # >1 = atom-split). Mirrors dsa.py's `if expand_for_dsl and # num_generations > 0` block which runs for any next_n ≥ 2. _force_dsl_expand_setup(metadata_gen) - Indexer.prepare(metadata_gen) + Indexer.prepare_metadata(metadata_gen) k_gen_fp4, k_gen_scale = torch.ops.trtllm.fused_cat_fp4( k_gen_bf16[:, :pe_dim].contiguous(), @@ -2529,7 +2553,7 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, compr compress_ratio=compress_ratio, ) - Indexer.prepare(metadata_chunked) + Indexer.prepare_metadata(metadata_chunked) assert metadata_chunked.indexer_prefill_chunks is not None num_chunks = len(metadata_chunked.indexer_prefill_chunks) @@ -2570,7 +2594,7 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, compr compress_ratio=compress_ratio, ) - Indexer.prepare(metadata_baseline) + Indexer.prepare_metadata(metadata_baseline) if metadata_baseline.indexer_prefill_chunks is not None: num_baseline_chunks = len(metadata_baseline.indexer_prefill_chunks) @@ -2782,7 +2806,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l max_draft_tokens=next_n - 1, indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_context) + Indexer.prepare_metadata(metadata_context) indexer._update_k_cache(k_context_fp8, k_context_scale, metadata_context) # Generate decode phase test data @@ -2808,7 +2832,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l max_draft_tokens=next_n - 1, indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_gen_write) + Indexer.prepare_metadata(metadata_gen_write) indexer._update_k_cache(k_fp8, k_scale, metadata_gen_write) # Test with custom CUDA kernel @@ -2828,7 +2852,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_custom) + Indexer.prepare_metadata(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) try: @@ -2855,7 +2879,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_fallback) + Indexer.prepare_metadata(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) _set_torch_top_k(indexer) topk_indices_fallback = indexer.sparse_attn_indexer( @@ -2881,7 +2905,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_skip) + Indexer.prepare_metadata(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) try: @@ -2984,7 +3008,7 @@ def make_inputs(n_tokens): kv_lens.sum().item(), max_draft_tokens=md, ) - Indexer.prepare(meta_ctx) + Indexer.prepare_metadata(meta_ctx) indexer._update_k_cache(ctx_k_fp8, ctx_k_scale, meta_ctx) step0_tokens = batch_size * step0_next_n @@ -3002,7 +3026,7 @@ def make_inputs(n_tokens): max_model_len, max_draft_tokens=md, ) - Indexer.prepare(meta0) + Indexer.prepare_metadata(meta0) # indexer_topk_decode needs caller-owned radix aux buffers for small gen batches. _radix_bp = 10 meta0.radix_aux_indices = torch.zeros( @@ -3027,7 +3051,7 @@ def make_inputs(n_tokens): max_model_len, max_draft_tokens=md, ) - Indexer.prepare(meta0) + Indexer.prepare_metadata(meta0) # context stash branch reads seq_lens_cuda (a read-only property); set its backing field. meta0._seq_lens_cuda = kv_lens.clone().cuda() @@ -3069,7 +3093,7 @@ def make_inputs(n_tokens): max_model_len, max_draft_tokens=md, ) - Indexer.prepare(meta) + Indexer.prepare_metadata(meta) meta.in_mtp_draft_loop = True meta.shared_topk_indices = stash meta.indexer_skip_topk = True @@ -3154,7 +3178,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_custom) + Indexer.prepare_metadata(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) assert metadata_custom.indexer_prefill_chunks is not None @@ -3182,7 +3206,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_fallback) + Indexer.prepare_metadata(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) _set_torch_top_k(indexer) topk_indices_fallback = indexer.sparse_attn_indexer( @@ -3263,7 +3287,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_custom) + Indexer.prepare_metadata(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) # Force single-pass path by setting indexer_prefill_chunks to None metadata_custom.indexer_prefill_chunks = None @@ -3291,7 +3315,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_fallback) + Indexer.prepare_metadata(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) # Force single-pass path by setting indexer_prefill_chunks to None metadata_fallback.indexer_prefill_chunks = None @@ -3317,7 +3341,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, enable_indexer_skip=True, indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_skip) + Indexer.prepare_metadata(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) metadata_skip.indexer_prefill_chunks = None @@ -3454,7 +3478,7 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): enable_indexer_skip=True, indexer_head_dim=head_dim, ) - Indexer.prepare(metadata_skip) + Indexer.prepare_metadata(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) topk_indices_skip = indexer.sparse_attn_indexer( metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights @@ -3702,7 +3726,7 @@ def test_cutedsl_mqa_logits_output_buffer_persistent(): index_topk=index_topk, use_cute_dsl_paged_mqa_logits=True, ) - Indexer.prepare(metadata) + Indexer.prepare_metadata(metadata) kv_cache = cache_manager.get_indexer_k_cache_buffers(0) q = torch.randn((batch_size, next_n, heads, head_dim), device="cuda", dtype=torch.bfloat16).to( diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py b/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py index 65d02bf07fbc..67b8b88831ca 100644 --- a/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py +++ b/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py @@ -368,7 +368,7 @@ def test_indexer_k_cache_scatter_custom_op_fp4(): from tensorrt_llm._torch.attention_backend.sparse.dsa import Indexer - Indexer.prepare(metadata) + Indexer.prepare_metadata(metadata) # FP4 packed data: [num_tokens, 64] int8; scale: [num_tokens, 1] int32 k_fp4 = torch.randint(-128, 127, (num_tokens, fp4_data_dim), device="cuda", dtype=torch.int8) diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py index 40135270736c..0696d72ebab3 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py @@ -260,6 +260,32 @@ def test_mla_backend_only_forward() -> None: ) +def test_mla_sparse_prepare_delegates_to_hooks() -> None: + mla = MLA.__new__(MLA) + torch.nn.Module.__init__(mla) + mla.sparse_attn_hooks = Mock() + attn_metadata = Mock() + + MLA.prepare_sparse_attn(mla, attn_metadata) + + mla.sparse_attn_hooks.prepare.assert_called_once_with(mla, attn_metadata) + + +@pytest.mark.parametrize("algorithm", ["dsa", "deepseek_v4"]) +def test_sparse_mla_hooks_prepare_the_indexer(algorithm: str) -> None: + hook_module = ModuleType(f"{algorithm}_prepare") + hook_module.sparse_params = MockSparseParams() + hook_module.sparse_params.algorithm = algorithm + hooks = get_sparse_mla_hooks(hook_module) + indexer = Mock() + mla = Mock(indexer=indexer) + attn_metadata = Mock() + + hooks.prepare(mla, attn_metadata) + + indexer.prepare.assert_called_once_with(attn_metadata) + + def test_sparse_runtime_params_without_prediction() -> None: attention = TrtllmAttention.__new__(TrtllmAttention) attention.sparse_params = MockSparseParams() diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 7d44477e520e..bf3b531994de 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -91,6 +91,19 @@ def forward(self, *args, **kwargs) -> torch.Tensor: return {"logits": torch.randn((batch_size, 10), device='cuda')} +def test_sparse_attention_modules_are_prepared_once_per_metadata(): + sparse_module = Mock() + sparse_module.prepare_sparse_attn = Mock() + engine = SimpleNamespace(model=Mock()) + engine.model.modules.return_value = [object(), sparse_module] + attn_metadata = SimpleNamespace() + + PyTorchModelEngine._prepare_sparse_attention_modules(engine, attn_metadata) + PyTorchModelEngine._prepare_sparse_attention_modules(engine, attn_metadata) + + sparse_module.prepare_sparse_attn.assert_called_once_with(attn_metadata) + + class DummyMultimodalIndexModel(torch.nn.Module): class Config: @@ -200,6 +213,7 @@ def _make_forward_only_engine( engine = object.__new__(PyTorchModelEngine) engine.model = SimpleNamespace( extra_attrs={}, + modules=lambda: [], model_config=SimpleNamespace(pretrained_config=SimpleNamespace( rope_scaling=None))) engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER From a14c0b14b013525e944256d47b7f0b3308086066 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Fri, 7 Aug 2026 09:31:47 +0000 Subject: [PATCH 05/18] [None][refactor] simplify sparse top-k lifecycle Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../sparse/deepseek_v4/metadata.py | 2 +- .../sparse/deepseek_v4/module.py | 4 - .../attention_backend/sparse/dsa/indexer.py | 47 +- .../attention_backend/sparse/dsa/metadata.py | 32 +- .../attention_backend/sparse/dsa/module.py | 4 - .../_torch/attention_backend/sparse/hooks.py | 3 - .../triattention/triattention.py | 2 +- tensorrt_llm/_torch/modules/mla.py | 5 - tensorrt_llm/_torch/modules/top_k.py | 468 ++++++------------ .../_torch/pyexecutor/model_engine.py | 26 +- .../attention/sparse/dsa/test_dsa_indexer.py | 105 ++-- .../attention/sparse/test_dsa_fp4_indexer.py | 2 +- .../attention/sparse/test_sparse_attention.py | 26 - .../executor/test_pytorch_model_engine.py | 50 +- tests/unittest/_torch/modules/test_top_k.py | 125 ++--- 15 files changed, 358 insertions(+), 543 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py index 3cf478436904..cfd22a4da152 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py @@ -576,7 +576,7 @@ def prepare(self): # Prepare metadata for indexer (only needed when sparse layers exist) if has_sparse_layers: - DeepseekV4Indexer.prepare_metadata(metadata=self) + DeepseekV4Indexer.prepare(metadata=self) # --- Per-ratio metadata --- # 1) CPU-side: compute scalar metadata (num_total_compressed_tokens, etc.) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py index f904312ae0b3..2c09926c9835 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -1243,10 +1243,6 @@ def initialize(self, mla: MLA) -> None: def create_weights(self, mla: MLA) -> None: create_sparse_attn_weights(mla) - def prepare(self, mla: MLA, attn_metadata: AttentionMetadata) -> None: - if mla.indexer is not None: - mla.indexer.prepare(attn_metadata) - def prepare_outputs( self, mla: MLA, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 3f2b11fd0796..2369d9db389a 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -136,7 +136,7 @@ def _compute_slot_mappings( ) -> Tuple[torch.Tensor, torch.Tensor]: """Compute flat byte indices for FP8/FP4 data and scales from global token positions. - Shared by Indexer.prepare_metadata() (CPU) and on_update_kv_lens() (GPU) to avoid + Shared by Indexer.prepare() (CPU) and on_update_kv_lens() (GPU) to avoid duplicating the slot mapping arithmetic. Args: @@ -641,15 +641,15 @@ def __init__( decode_top_k_implementation = ( TopKImplementation.CUTE_DSL_GVR if self._enable_heuristic_topk - else TopKImplementation.CUTE_DSL_PREFERRED + else TopKImplementation.CUTE_DSL_RADIX ) elif self._enable_heuristic_topk: - decode_top_k_implementation = TopKImplementation.TRTLLM_HEURISTIC + decode_top_k_implementation = TopKImplementation.CUDA_GVR else: - decode_top_k_implementation = TopKImplementation.TRTLLM + decode_top_k_implementation = TopKImplementation.CUDA_RADIX self.top_k = TopK( self.index_topk, - prefill_implementation=TopKImplementation.TRTLLM, + prefill_implementation=TopKImplementation.CUDA_RADIX, decode_implementation=decode_top_k_implementation, compress_ratio=self.compress_ratio, ) @@ -669,19 +669,6 @@ def cache_derived_state(self) -> None: def post_load_weights(self) -> None: self.cache_derived_state() - def prepare(self, metadata: DSAtrtllmAttentionMetadata) -> None: - """Prepare this indexer's Top-K state before model forward.""" - if metadata.kv_cache_manager is None: - return - self.top_k.prepare( - device=metadata.kv_lens_cuda.device, - max_num_columns=metadata.get_indexer_max_seq_len(), - next_n=1 + metadata.max_draft_tokens, - input_dtype=torch.float32, - num_sms=metadata.num_sms, - max_num_requests=metadata.max_num_sequences, - ) - @staticmethod def prepare_one_prefill_chunk( metadata: DSAtrtllmAttentionMetadata, @@ -1097,7 +1084,7 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): ) @staticmethod - def prepare_metadata(metadata: DSAtrtllmAttentionMetadata): + def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: """ Prepare indexer for the forward pass. This should be called during metadata.prepare() stage. @@ -1159,6 +1146,16 @@ def prepare_metadata(metadata: DSAtrtllmAttentionMetadata): # This is a preprocessing step that computes scheduling information for the kernel Indexer.prepare_scheduler_metadata(metadata) + for indexer in getattr(metadata, "indexers", ()): + indexer.top_k.prepare( + device=metadata.kv_lens_cuda.device, + max_num_columns=metadata.get_indexer_max_seq_len(), + next_n=1 + metadata.max_draft_tokens, + input_dtype=torch.float32, + num_sms=metadata.num_sms, + max_num_requests=metadata.max_num_sequences, + ) + def _update_k_cache( self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: DSAtrtllmAttentionMetadata ) -> None: @@ -1530,8 +1527,8 @@ def sparse_attn_indexer( :num_ctx_tokens, : ] - # Seed finishing prefill requests after the active generation slots. - if self._enable_heuristic_topk and has_prefill and not metadata.skip_indexer_for_ctx_reqs: + # Seed GVR state after the active generation slots. + if has_prefill and not metadata.skip_indexer_for_ctx_reqs: self.top_k.seed_from_prefill( topk_indices_buffer[:num_ctx_tokens], metadata.seq_lens[:num_contexts], @@ -1605,7 +1602,7 @@ def sparse_attn_indexer( if self.use_cute_dsl_paged_mqa_logits: # DSL kernel design: 1 atom per q (atom = real next_n positions), # kNumNextNAtoms = 1 for any real next_n. The matching schedule - # is `scheduler_metadata_buffer` — built in `Indexer.prepare_metadata()` + # is `scheduler_metadata_buffer` — built in `Indexer.prepare()` # with a (num_gen, 1) input shape, which makes DeepGEMM's wrapper # compute `num_next_n_atoms = 1`. (DeepGEMM uses the same buffer # for its own next_n=1 kernel; DSL piggy-backs on it for all @@ -1720,10 +1717,6 @@ def sparse_attn_indexer( scan_lengths = metadata.gen_indexer_kv_lens_cuda_runtime assert scan_lengths is not None - heuristic_values = None - if self._enable_heuristic_topk and not metadata.use_cute_dsl_topk: - heuristic_values = metadata.heuristic_scratch_values[:num_gen_tokens] - self.top_k( logits_decode, topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], @@ -1731,10 +1724,8 @@ def sparse_attn_indexer( sequence_lengths=gen_kv_lens_cuda, scan_lengths=scan_lengths, next_n=next_n, - heuristic_values=heuristic_values, radix_indices=metadata.radix_aux_indices, radix_values=metadata.radix_aux_logits, - max_num_columns=indexer_max_seq_len, ) elif has_decode and metadata.skip_indexer_for_gen_reqs: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 110891f65936..f3af752fd521 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -48,6 +48,7 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" sparse_metadata_params: Optional[DSAMetadataParams] = None + indexers: tuple[Indexer, ...] = () # Chunked prefill metadata for indexer (prefill-only, no CUDA graph needed) indexer_prefill_chunks: Optional[List[IndexerPrefillChunkMetadata]] = None # Max chunk size for two-level chunking: @@ -203,7 +204,7 @@ def prepare(self): self.prepare_for_spec_decode(kv_lens) # Prepare metadata for indexer - Indexer.prepare_metadata(metadata=self) + Indexer.prepare(metadata=self) def prepare_for_draft_forward(self) -> dict | None: """Select native DSA indexer metadata for a draft forward.""" @@ -617,7 +618,7 @@ def create_buffers_for_indexer(self, capture_graph=False): pin_memory=prefer_pinned(), ) # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation. - # Allocate these buffers dynamically in Indexer.prepare_metadata() + # Allocate these buffers dynamically in Indexer.prepare() # based on the actual total_kv_len to save memory. if self.enable_context_mla_with_cached_kv: self.slot_mapping_fp8_fullkv = None @@ -709,20 +710,6 @@ def create_buffers_for_indexer(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - self.enable_heuristic_topk = ( - sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100 - ) - if self.enable_heuristic_topk and not self.use_cute_dsl_topk: - # Shared C++ heuristic scratch; per-layer history lives in TopK. - max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) - self.heuristic_scratch_values = self.get_empty( - self.cuda_graph_buffers, - (max_gen_tokens, self.num_sparse_topk), - cache_name="heuristic_scratch_values", - dtype=torch.float32, - capture_graph=capture_graph, - ) - # Persistent scratch for the Radix-split-work indexer path. Re-created # in update_spec_dec_param when max_draft_tokens changes so it stays # large enough for the MTP generation-row count. @@ -816,6 +803,7 @@ def update_spec_dec_param( num_contexts: int = 0, ): """Update speculative decoding parameters and create expanded buffers.""" + previous_max_draft_tokens = self.max_draft_tokens super().update_spec_dec_param( batch_size, is_spec_decoding_enabled, @@ -835,16 +823,6 @@ def update_spec_dec_param( init_shape = self.kv_lens_expanded_host.shape[0] if self.max_num_sequences * (1 + self.max_draft_tokens) != init_shape: self.create_expanded_buffers(capture_graph=capture_graph) - # Resize heuristic scratch buffer for new max_draft_tokens. - if self.enable_heuristic_topk and not self.use_cute_dsl_topk: - max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) - self.heuristic_scratch_values = self.get_empty( - self.cuda_graph_buffers, - (max_gen_tokens, self.num_sparse_topk), - cache_name="heuristic_scratch_values", - dtype=torch.float32, - capture_graph=capture_graph, - ) # The Radix-split-work scratch (radix_aux_*) is sized the same way # (num_seqs * (1 + max_draft_tokens) rows) and is allocated # unconditionally, so it must be resized here too -- otherwise the @@ -852,6 +830,8 @@ def update_spec_dec_param( # ("radix_aux_* must hold at least num_rows*blocks_per_row*index_topk # elements"). self._create_radix_aux_buffers(capture_graph=capture_graph) + if self.max_draft_tokens != previous_max_draft_tokens: + Indexer.prepare(metadata=self) def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor: """Refresh INDEX_KEY offsets and return their physical pool slots.""" diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py index 80498f65b438..bf36a14655ed 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py @@ -595,10 +595,6 @@ def initialize(self, mla: MLA) -> None: self.need_dense_mha and mla.mapping.cp_size == 1 and mla.mqa.support_fused_rope() ) - def prepare(self, mla: MLA, attn_metadata: AttentionMetadata) -> None: - if mla.indexer is not None: - mla.indexer.prepare(attn_metadata) - def forward( self, mla: MLA, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/hooks.py b/tensorrt_llm/_torch/attention_backend/sparse/hooks.py index 12edceabf237..6bc1177323aa 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/hooks.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/hooks.py @@ -56,9 +56,6 @@ def create_weights(self, mla: "MLA") -> None: def transform_weights(self, mla: "MLA") -> None: """Transform algorithm-specific weights.""" - def prepare(self, mla: "MLA", attn_metadata: "AttentionMetadata") -> None: - """Prepare algorithm-specific module state before model forward.""" - def prepare_outputs( self, mla: "MLA", diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 0c9e65e28628..8462a4bd8cc7 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -803,7 +803,7 @@ def _allocate_selection_buffers(self, device: torch.device, *, tp_size: int) -> """Allocate fixed manager-lifetime TopK inputs and outputs.""" self._selection_top_k = TopK( self.budget, - decode_implementation=TopKImplementation.CUTE_DSL_PREFERRED, + decode_implementation=TopKImplementation.CUTE_DSL_RADIX, ) request_capacity = self._request_capacity selection_width = self._selection_width_capacity diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index f486e36dfba3..646509bb25da 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -795,11 +795,6 @@ def create_output(self, hidden_states: torch.Tensor, num_contexts: int): [num_tokens, self.attention_output_hidden_size], dtype=hidden_states.dtype ) - def prepare_sparse_attn(self, attn_metadata: AttentionMetadata) -> None: - """Prepare sparse module state before the model forward starts.""" - if self.sparse_attn_hooks is not None: - self.sparse_attn_hooks.prepare(self, attn_metadata) - def _create_outputs( self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata ) -> list[torch.Tensor]: diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index cc31f04a7016..fe74572c6231 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -4,112 +4,80 @@ from __future__ import annotations -import threading from enum import Enum +from functools import cache import torch import torch.nn as nn class TopKImplementation(str, Enum): - """Top-K implementations for prefill and decode. - - ``CUTE_DSL_PREFERRED`` uses the CuTe DSL radix implementation when it - supports the runtime shape and falls back to TRTLLM otherwise. - """ + """Top-K implementations grouped by backend and algorithm.""" TORCH = "torch" - TRTLLM = "trtllm" - TRTLLM_HEURISTIC = "trtllm_heuristic" - CUTE_DSL_PREFERRED = "cute_dsl_preferred" + CUDA_RADIX = "cuda_radix" + CUTE_DSL_RADIX = "cute_dsl_radix" + CUDA_GVR = "cuda_gvr" CUTE_DSL_GVR = "cute_dsl_gvr" -_PREFILL_IMPLEMENTATIONS = { - TopKImplementation.TORCH, - TopKImplementation.TRTLLM, +_GVR_IMPLEMENTATIONS = { + TopKImplementation.CUDA_GVR, + TopKImplementation.CUTE_DSL_GVR, } - - -_PREPARE_LOCK = threading.Lock() -_PREPARED_DECODE_TOP_K: set[tuple[object, ...]] = set() -_HEURISTIC_WARMUP_COLUMNS = 4096 +_GVR_WARMUP_COLUMNS = 4096 _RADIX_MAX_BLOCKS_PER_ROW = 10 -def _cuda_device(device: torch.device) -> torch.device: - device = torch.device(device) - if device.type != "cuda": - raise ValueError(f"Top-K preparation requires a CUDA device, got {device}") - if device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) - return device - - -def _validate_output( - scores: torch.Tensor, - output_indices: torch.Tensor, +@cache +def _warmup_decode_top_k( + implementation: TopKImplementation, + device: torch.device, + input_dtype: torch.dtype, top_k: int, + max_num_columns: int, + next_n: int, + num_sms: int | None, + compress_ratio: int, ) -> None: - if scores.ndim != 2: - raise ValueError(f"scores must be rank 2, got shape {tuple(scores.shape)}") - if output_indices.ndim != 2: - raise ValueError(f"output_indices must be rank 2, got shape {tuple(output_indices.shape)}") - expected_shape = (scores.shape[0], top_k) - if output_indices.shape != expected_shape: - raise ValueError( - f"output_indices must have shape {expected_shape}, got {tuple(output_indices.shape)}" + if implementation == TopKImplementation.CUDA_GVR: + num_columns = max(_GVR_WARMUP_COLUMNS, top_k) + scores = torch.zeros((1, num_columns), dtype=input_dtype, device=device) + sequence_lengths = torch.tensor([num_columns], dtype=torch.int32, device=device) + output_indices = torch.empty((1, top_k), dtype=torch.int32, device=device) + prior_indices = torch.zeros((1, top_k), dtype=torch.int32, device=device) + heuristic_values = torch.empty((1, top_k), dtype=input_dtype, device=device) + radix_indices = torch.empty( + (1, _RADIX_MAX_BLOCKS_PER_ROW, top_k), dtype=torch.int32, device=device ) - if output_indices.dtype != torch.int32: - raise TypeError(f"output_indices must have dtype torch.int32, got {output_indices.dtype}") - if output_indices.device != scores.device: - raise ValueError( - "scores and output_indices must be on the same device, got " - f"{scores.device} and {output_indices.device}" + radix_values = torch.empty( + (1, _RADIX_MAX_BLOCKS_PER_ROW, top_k), dtype=torch.float32, device=device ) - - -def _validate_lengths( - scores: torch.Tensor, - lengths: torch.Tensor, - name: str, - expected_size: int, -) -> None: - if lengths.ndim != 1 or lengths.shape[0] != expected_size: - raise ValueError(f"{name} must have shape ({expected_size},), got {tuple(lengths.shape)}") - if lengths.dtype != torch.int32: - raise TypeError(f"{name} must have dtype torch.int32, got {lengths.dtype}") - if lengths.device != scores.device: - raise ValueError( - f"scores and {name} must be on the same device, got {scores.device} and {lengths.device}" + torch.ops.trtllm.indexer_topk_decode( + scores, + sequence_lengths, + output_indices, + 1, + top_k, + pre_idx=prior_indices, + heuristic_scratch=heuristic_values, + radix_aux_indices=radix_indices, + radix_aux_logits=radix_values, + ) + elif implementation == TopKImplementation.CUTE_DSL_RADIX and not ( + compress_ratio > 1 and next_n > 1 + ): + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + warmup_cute_dsl_radix_topk_decode, ) - -def _require_tensor(tensor: torch.Tensor | None, name: str) -> torch.Tensor: - if tensor is None: - raise ValueError(f"{name} is required") - return tensor - - -def _forward_prefill_torch( - scores: torch.Tensor, - row_starts: torch.Tensor, - row_ends: torch.Tensor, - output_indices: torch.Tensor, - top_k: int, -) -> torch.Tensor: - output_indices.fill_(-1) - selected_count = min(top_k, scores.shape[1]) - if selected_count == 0: - return output_indices - columns = torch.arange(scores.shape[1], device=scores.device).unsqueeze(0) - valid = (columns >= row_starts.unsqueeze(1)) & (columns < row_ends.unsqueeze(1)) - selected = scores.masked_fill(~valid, float("-inf")).topk(selected_count, dim=-1).indices - selected_valid = torch.gather(valid, 1, selected) - selected = selected - row_starts.unsqueeze(1) - selected = selected.masked_fill(~selected_valid, -1) - output_indices[:, :selected_count].copy_(selected.to(torch.int32)) - return output_indices + warmup_cute_dsl_radix_topk_decode( + top_k=top_k, + num_cols=max_num_columns, + next_n=next_n, + dtype=input_dtype, + num_sms=num_sms, + ) class TopK(nn.Module): @@ -124,31 +92,17 @@ def __init__( compress_ratio: int = 1, ) -> None: super().__init__() - if top_k <= 0: - raise ValueError(f"top_k must be positive, got {top_k}") - if compress_ratio <= 0: - raise ValueError(f"compress_ratio must be positive, got {compress_ratio}") - if prefill_implementation is None and decode_implementation is None: - raise ValueError("at least one Top-K implementation must be configured") self.top_k = top_k - self.prefill_implementation = ( - TopKImplementation(prefill_implementation) - if prefill_implementation is not None - else None + self.prefill_implementation = TopKImplementation( + prefill_implementation or TopKImplementation.CUDA_RADIX ) - if ( - self.prefill_implementation is not None - and self.prefill_implementation not in _PREFILL_IMPLEMENTATIONS - ): - raise ValueError( - f"{self.prefill_implementation.value} is not supported for prefill Top-K" - ) - self.decode_implementation = ( - TopKImplementation(decode_implementation) if decode_implementation is not None else None + self.decode_implementation = TopKImplementation( + decode_implementation or TopKImplementation.CUDA_RADIX ) self.compress_ratio = compress_ratio - self.register_buffer("_prior_indices", None, persistent=False) - self.register_buffer("_row_order_buffer", None, persistent=False) + self.register_buffer("_gvr_prior_indices", None, persistent=False) + self.register_buffer("_cuda_gvr_scratch", None, persistent=False) + self.register_buffer("_gvr_row_order", None, persistent=False) self._num_sms: int | None = None def prepare( @@ -158,34 +112,33 @@ def prepare( max_num_columns: int, next_n: int, input_dtype: torch.dtype, + max_num_requests: int, num_sms: int | None = None, - max_num_requests: int | None = None, ) -> None: - """Prepare persistent state and warm up one decode deployment shape.""" + """Allocate GVR state and warm up the configured decode implementation.""" implementation = self.decode_implementation - if implementation is None: - raise ValueError("decode Top-K is not configured") - if max_num_requests is not None: - self._prepare_state(device, max_num_requests, num_sms) - if max_num_columns <= 0 or next_n <= 0: - return - if implementation in ( - TopKImplementation.TORCH, - TopKImplementation.TRTLLM, - TopKImplementation.CUTE_DSL_GVR, - ): - return - if ( - implementation == TopKImplementation.CUTE_DSL_PREFERRED - and self.compress_ratio > 1 - and next_n > 1 - ): - return + if implementation in _GVR_IMPLEMENTATIONS: + if self._gvr_prior_indices is None: + self._gvr_prior_indices = torch.zeros( + (max_num_requests, self.top_k), dtype=torch.int32, device=device + ) + + if implementation == TopKImplementation.CUDA_GVR: + scratch_shape = (max_num_requests * next_n, self.top_k) + if self._cuda_gvr_scratch is None or self._cuda_gvr_scratch.shape != scratch_shape: + self._cuda_gvr_scratch = torch.empty( + scratch_shape, dtype=input_dtype, device=device + ) + else: + self._num_sms = num_sms + if self._gvr_row_order is None: + self._gvr_row_order = torch.empty( + (max_num_requests,), dtype=torch.int32, device=device + ) - device = _cuda_device(device) - key = ( + _warmup_decode_top_k( implementation, - device.index, + device, input_dtype, self.top_k, max_num_columns, @@ -193,105 +146,6 @@ def prepare( num_sms, self.compress_ratio, ) - with _PREPARE_LOCK: - if key in _PREPARED_DECODE_TOP_K: - return - with torch.cuda.device(device): - if implementation == TopKImplementation.TRTLLM_HEURISTIC: - self._warmup_trtllm_heuristic(input_dtype) - else: - try: - from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( - warmup_cute_dsl_radix_topk_decode, - ) - except ImportError: - return - warmup_cute_dsl_radix_topk_decode( - top_k=self.top_k, - num_cols=max_num_columns, - next_n=next_n, - dtype=input_dtype, - num_sms=num_sms, - ) - _PREPARED_DECODE_TOP_K.add(key) - - def _prepare_state( - self, - device: torch.device, - max_num_requests: int, - num_sms: int | None, - ) -> None: - implementation = self.decode_implementation - if implementation not in ( - TopKImplementation.TRTLLM_HEURISTIC, - TopKImplementation.CUTE_DSL_GVR, - ): - return - if max_num_requests <= 0: - raise ValueError(f"max_num_requests must be positive, got {max_num_requests}") - - device = torch.device(device) - prior_shape = (max_num_requests, self.top_k) - if self._prior_indices is None: - self._prior_indices = torch.zeros(prior_shape, dtype=torch.int32, device=device) - elif ( - self._prior_indices.device != device or self._prior_indices.shape[0] < max_num_requests - ): - raise RuntimeError( - "Top-K state was already prepared with an incompatible device or capacity: " - f"got {self._prior_indices.device} {tuple(self._prior_indices.shape)}, " - f"requested {device} {prior_shape}" - ) - - if implementation == TopKImplementation.CUTE_DSL_GVR: - if num_sms is None or num_sms <= 0: - raise ValueError("num_sms must be positive when preparing CUTE_DSL_GVR") - self._num_sms = num_sms - if self._row_order_buffer is None: - self._row_order_buffer = torch.empty( - (max_num_requests,), dtype=torch.int32, device=device - ) - elif ( - self._row_order_buffer.device != device - or self._row_order_buffer.shape[0] < max_num_requests - ): - raise RuntimeError( - "GVR row-order state was already prepared with an incompatible device or " - f"capacity: got {self._row_order_buffer.device} " - f"{tuple(self._row_order_buffer.shape)}, requested {device} " - f"({max_num_requests},)" - ) - - def _warmup_trtllm_heuristic(self, input_dtype: torch.dtype) -> None: - num_columns = max(_HEURISTIC_WARMUP_COLUMNS, self.top_k) - device = torch.device("cuda") - scores = torch.zeros((1, num_columns), dtype=input_dtype, device=device) - sequence_lengths = torch.tensor([num_columns], dtype=torch.int32, device=device) - output_indices = torch.empty((1, self.top_k), dtype=torch.int32, device=device) - prior_indices = torch.zeros((1, self.top_k), dtype=torch.int32, device=device) - heuristic_values = torch.empty((1, self.top_k), dtype=input_dtype, device=device) - radix_indices = torch.empty( - (1, _RADIX_MAX_BLOCKS_PER_ROW, self.top_k), - dtype=torch.int32, - device=device, - ) - radix_values = torch.empty( - (1, _RADIX_MAX_BLOCKS_PER_ROW, self.top_k), - dtype=torch.float32, - device=device, - ) - torch.ops.trtllm.indexer_topk_decode( - scores, - sequence_lengths, - output_indices, - 1, - self.top_k, - pre_idx=prior_indices, - heuristic_scratch=heuristic_values, - radix_aux_indices=radix_indices, - radix_aux_logits=radix_values, - ) - torch.cuda.synchronize() def forward( self, @@ -304,34 +158,23 @@ def forward( sequence_lengths: torch.Tensor | None = None, scan_lengths: torch.Tensor | None = None, next_n: int = 1, - heuristic_values: torch.Tensor | None = None, radix_indices: torch.Tensor | None = None, radix_values: torch.Tensor | None = None, - max_num_columns: int | None = None, ) -> torch.Tensor: - """Write prefill or decode Top-K indices into ``output_indices``. - - Prefill uses ``row_starts`` and ``row_ends``. Decode uses logical - ``sequence_lengths`` plus ``scan_lengths`` in score-column coordinates. - """ - _validate_output(scores, output_indices, self.top_k) + """Write prefill or decode Top-K indices into ``output_indices``.""" if is_prefill: - return self._forward_prefill( - scores, - _require_tensor(row_starts, "row_starts"), - _require_tensor(row_ends, "row_ends"), - output_indices, - ) + assert row_starts is not None and row_ends is not None + return self._forward_prefill(scores, row_starts, row_ends, output_indices) + + assert sequence_lengths is not None and scan_lengths is not None return self._forward_decode( scores, - _require_tensor(sequence_lengths, "sequence_lengths"), - _require_tensor(scan_lengths, "scan_lengths"), + sequence_lengths, + scan_lengths, output_indices, - next_n=next_n, - heuristic_values=heuristic_values, - radix_indices=radix_indices, - radix_values=radix_values, - max_num_columns=max_num_columns, + next_n, + radix_indices, + radix_values, ) def seed_from_prefill( @@ -341,21 +184,12 @@ def seed_from_prefill( *, request_offset: int = 0, ) -> None: - """Seed decode hints from the last selected row of each prefill request.""" - if self._prior_indices is None: + """Seed GVR from the last selected row of each prefill request.""" + if self._gvr_prior_indices is None: return - if request_lengths.ndim != 1: - raise ValueError( - f"request_lengths must be rank 1, got shape {tuple(request_lengths.shape)}" - ) - num_requests = request_lengths.shape[0] - if request_offset < 0 or request_offset + num_requests > self._prior_indices.shape[0]: - raise ValueError( - f"request range [{request_offset}, {request_offset + num_requests}) exceeds " - f"prepared capacity {self._prior_indices.shape[0]}" - ) last_rows = (torch.cumsum(request_lengths, dim=0) - 1).to(dtype=torch.long) - self._prior_indices[request_offset : request_offset + num_requests].copy_( + num_requests = request_lengths.shape[0] + self._gvr_prior_indices[request_offset : request_offset + num_requests].copy_( output_indices[last_rows] ) @@ -366,18 +200,16 @@ def _forward_prefill( row_ends: torch.Tensor, output_indices: torch.Tensor, ) -> torch.Tensor: - implementation = self.prefill_implementation - if implementation is None: - raise ValueError("prefill Top-K is not configured") - _validate_lengths(scores, row_starts, "row_starts", scores.shape[0]) - _validate_lengths(scores, row_ends, "row_ends", scores.shape[0]) - if implementation == TopKImplementation.TORCH: - return _forward_prefill_torch( + if self.prefill_implementation == TopKImplementation.TORCH: + return self._forward_prefill_torch( scores, row_starts, row_ends, output_indices, - self.top_k, + ) + if self.prefill_implementation != TopKImplementation.CUDA_RADIX: + raise NotImplementedError( + f"{self.prefill_implementation.value} does not support prefill Top-K" ) torch.ops.trtllm.indexer_topk_prefill( scores, @@ -394,56 +226,30 @@ def _forward_decode( sequence_lengths: torch.Tensor, scan_lengths: torch.Tensor, output_indices: torch.Tensor, - *, next_n: int, - heuristic_values: torch.Tensor | None, radix_indices: torch.Tensor | None, radix_values: torch.Tensor | None, - max_num_columns: int | None, ) -> torch.Tensor: implementation = self.decode_implementation - if implementation is None: - raise ValueError("decode Top-K is not configured") - if next_n <= 0: - raise ValueError(f"next_n must be positive, got {next_n}") - if scores.shape[0] % next_n != 0: - raise ValueError( - f"score rows ({scores.shape[0]}) must be divisible by next_n ({next_n})" - ) - num_requests = scores.shape[0] // next_n - for name, lengths in ( - ("sequence_lengths", sequence_lengths), - ("scan_lengths", scan_lengths), - ): - _validate_lengths(scores, lengths, name, num_requests) - - if (radix_indices is None) != (radix_values is None): - raise ValueError("radix_indices and radix_values must be provided together") - - prior_indices = None - if implementation in ( - TopKImplementation.TRTLLM_HEURISTIC, - TopKImplementation.CUTE_DSL_GVR, - ): - if self._prior_indices is None: - raise RuntimeError("Top-K state must be prepared before heuristic decode") - prior_indices = self._prior_indices[:num_requests] - if implementation == TopKImplementation.TORCH: return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n) - use_trtllm = implementation in ( - TopKImplementation.TRTLLM, - TopKImplementation.TRTLLM_HEURISTIC, + use_cuda = implementation in ( + TopKImplementation.CUDA_RADIX, + TopKImplementation.CUDA_GVR, ) or ( - implementation == TopKImplementation.CUTE_DSL_PREFERRED + implementation == TopKImplementation.CUTE_DSL_RADIX and self.compress_ratio > 1 and next_n > 1 ) - if use_trtllm: - if implementation == TopKImplementation.TRTLLM_HEURISTIC: - if prior_indices is None or heuristic_values is None: - raise ValueError("TRTLLM_HEURISTIC requires prior_indices and heuristic_values") + if use_cuda: + prior_indices = None + heuristic_values = None + if implementation == TopKImplementation.CUDA_GVR: + assert self._gvr_prior_indices is not None + assert self._cuda_gvr_scratch is not None + prior_indices = self._gvr_prior_indices[: sequence_lengths.shape[0]] + heuristic_values = self._cuda_gvr_scratch[: scores.shape[0]] torch.ops.trtllm.indexer_topk_decode( scores, sequence_lengths, @@ -456,11 +262,11 @@ def _forward_decode( radix_aux_indices=radix_indices, radix_aux_logits=radix_values, ) - if implementation == TopKImplementation.TRTLLM_HEURISTIC: - self._update_prior_indices(output_indices, num_requests, next_n) + if prior_indices is not None: + prior_indices.copy_(output_indices[next_n - 1 :: next_n]) return output_indices - if implementation == TopKImplementation.CUTE_DSL_PREFERRED: + if implementation == TopKImplementation.CUTE_DSL_RADIX: torch.ops.trtllm.cute_dsl_indexer_topk_decode( scores, scan_lengths, @@ -470,9 +276,8 @@ def _forward_decode( ) return output_indices - if max_num_columns is None: - raise ValueError("CUTE_DSL_GVR requires max_num_columns") - row_order = self._prepare_row_order(sequence_lengths, next_n) + assert self._gvr_prior_indices is not None + prior_indices = self._gvr_prior_indices[: sequence_lengths.shape[0]] torch.ops.trtllm.cute_dsl_gvr_topk_decode( scores, prior_indices, @@ -481,10 +286,10 @@ def _forward_decode( self.top_k, next_n=next_n, compress_ratio=self.compress_ratio, - max_seq_len=max_num_columns, - order_row=row_order, + max_seq_len=scores.shape[1], + order_row=self._prepare_row_order(sequence_lengths, next_n), ) - self._update_prior_indices(output_indices, num_requests, next_n) + prior_indices.copy_(output_indices[next_n - 1 :: next_n]) return output_indices def _prepare_row_order( @@ -492,24 +297,33 @@ def _prepare_row_order( sequence_lengths: torch.Tensor, next_n: int, ) -> torch.Tensor | None: - if self._num_sms is None or self._row_order_buffer is None: - raise RuntimeError("GVR state must be prepared before decode") + assert self._num_sms is not None and self._gvr_row_order is not None num_requests = sequence_lengths.shape[0] if num_requests * next_n < 2 * self._num_sms: return None - order = torch.argsort(sequence_lengths, descending=True).to(torch.int32) - row_order = self._row_order_buffer[:num_requests] - row_order.copy_(order) + row_order = self._gvr_row_order[:num_requests] + row_order.copy_(torch.argsort(sequence_lengths, descending=True).to(torch.int32)) return row_order - def _update_prior_indices( + def _forward_prefill_torch( self, + scores: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, output_indices: torch.Tensor, - num_requests: int, - next_n: int, - ) -> None: - assert self._prior_indices is not None - self._prior_indices[:num_requests].copy_(output_indices[next_n - 1 :: next_n]) + ) -> torch.Tensor: + output_indices.fill_(-1) + selected_count = min(self.top_k, scores.shape[1]) + if selected_count == 0: + return output_indices + columns = torch.arange(scores.shape[1], device=scores.device).unsqueeze(0) + valid = (columns >= row_starts.unsqueeze(1)) & (columns < row_ends.unsqueeze(1)) + selected = scores.masked_fill(~valid, float("-inf")).topk(selected_count, dim=-1).indices + selected_valid = torch.gather(valid, 1, selected) + selected = selected - row_starts.unsqueeze(1) + selected = selected.masked_fill(~selected_valid, -1) + output_indices[:, :selected_count].copy_(selected.to(torch.int32)) + return output_indices def _forward_decode_torch( self, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 4e702ebd5083..516f8b4f8ae6 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3156,23 +3156,6 @@ def _update_draft_inference_state_for_warmup( req.py_is_first_draft = True req.py_draft_tokens = [] - def _prepare_sparse_attention_modules( - self, attn_metadata: AttentionMetadata) -> None: - """Prepare sparse module state before any model forward uses metadata.""" - prepared_metadata_ids = getattr( - self, "_prepared_sparse_attention_metadata_ids", None) - if prepared_metadata_ids is None: - prepared_metadata_ids = set() - self._prepared_sparse_attention_metadata_ids = prepared_metadata_ids - metadata_id = id(attn_metadata) - if metadata_id in prepared_metadata_ids: - return - for module in self.model.modules(): - prepare = getattr(module, "prepare_sparse_attn", None) - if callable(prepare): - prepare(attn_metadata) - prepared_metadata_ids.add(metadata_id) - def _set_up_attn_metadata( self, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], @@ -3247,6 +3230,13 @@ def _set_up_attn_metadata( num_heads_per_kv=num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, ) + if hasattr(self.attn_metadata, "indexers"): + indexers = {} + for module in self.model.modules(): + indexer = getattr(module, "indexer", None) + if indexer is not None and hasattr(indexer, "top_k"): + indexers[id(indexer)] = indexer + self.attn_metadata.indexers = tuple(indexers.values()) return self.attn_metadata @property @@ -7518,8 +7508,6 @@ def forward(self, spec_resource_manager = None spec_metadata = None - self._prepare_sparse_attention_modules(attn_metadata) - moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', None) if kv_cache_manager is None: diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index b3a3d5220e50..5806ae23dce4 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -250,26 +250,37 @@ def test_indexer_configures_one_top_k_module(): indexer = create_indexer(sparse_config) assert isinstance(indexer.top_k, TopK) - assert indexer.top_k.prefill_implementation == TopKImplementation.TRTLLM - assert indexer.top_k.decode_implementation == TopKImplementation.TRTLLM + assert indexer.top_k.prefill_implementation == TopKImplementation.CUDA_RADIX + assert indexer.top_k.decode_implementation == TopKImplementation.CUDA_RADIX assert not hasattr(indexer, "prefill_top_k") assert not hasattr(indexer, "decode_top_k") -def test_indexer_prepare_delegates_to_top_k_before_forward(): +def test_indexer_prepare_keeps_metadata_and_top_k_preparation_together(): top_k = Mock() indexer = SimpleNamespace(top_k=top_k) metadata = SimpleNamespace( - kv_cache_manager=SimpleNamespace(), + indexers=(indexer,), + num_contexts=0, + num_generations=0, + num_ctx_tokens=0, + seq_lens=torch.empty(0, dtype=torch.int32), + compress_ratios=[1], kv_lens_cuda=torch.empty(0), get_indexer_max_seq_len=Mock(return_value=4096), max_draft_tokens=3, num_sms=148, max_num_sequences=32, ) + indexer_params = SimpleNamespace(new_kv_tokens=torch.empty(0, dtype=torch.int32)) - Indexer.prepare(indexer, metadata) + with ( + patch.object(Indexer, "build_indexer_params", return_value=indexer_params), + patch.object(Indexer, "prepare_for_update_k_cache") as prepare_metadata, + ): + Indexer.prepare(metadata) + prepare_metadata.assert_called_once_with(metadata, indexer_params) top_k.prepare.assert_called_once_with( device=metadata.kv_lens_cuda.device, max_num_columns=4096, @@ -280,6 +291,36 @@ def test_indexer_prepare_delegates_to_top_k_before_forward(): ) +def test_draft_width_change_reprepares_indexer(): + metadata = object.__new__(DSAtrtllmAttentionMetadata) + metadata.max_draft_tokens = 0 + metadata.max_num_sequences = 2 + metadata.is_cuda_graph = False + metadata.kv_lens_cuda_2d = torch.empty((2, 1)) + metadata.kv_lens_expanded_host = torch.empty(2) + metadata._create_kv_lens_2d_buffer = Mock() + metadata.create_expanded_buffers = Mock() + metadata._create_radix_aux_buffers = Mock() + + with ( + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata." + "TrtllmAttentionMetadata.update_spec_dec_param" + ), + patch.object(Indexer, "prepare") as prepare, + ): + metadata.update_spec_dec_param( + batch_size=2, + is_spec_decoding_enabled=True, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + max_draft_len=3, + max_total_draft_tokens=3, + ) + + prepare.assert_called_once_with(metadata=metadata) + + def _ceil_to_ue8m0(x: torch.Tensor): """Round tensor values up to the nearest power of two (UE8M0 format).""" return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) @@ -817,7 +858,7 @@ def __init__(self): # so allocate a separate buffer for the full-next_n schedule. # DeepGEMM expects the full-next_n schedule in # `scheduler_metadata_buffer` itself (the alias makes - # `Indexer.prepare_metadata()`'s second populate overwrite the first). + # `Indexer.prepare()`'s second populate overwrite the first). if use_cute_dsl_paged_mqa_logits: self.scheduler_metadata_buffer_full_next_n = torch.zeros( (self.num_sms + 1, 2), device="cuda", dtype=torch.int32 @@ -1050,7 +1091,7 @@ def validate_topk_indices(topk_indices_0, topk_indices_1, total_tokens): @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper def test_recompute_slot_mappings_matches_prepare_with_cached_tokens(): - """Recompute slot mappings without re-running full Indexer.prepare_metadata().""" + """Recompute slot mappings without re-running full Indexer.prepare().""" head_dim = 128 block_size = 64 request_ids = [0, 1] @@ -1084,7 +1125,7 @@ def test_recompute_slot_mappings_matches_prepare_with_cached_tokens(): indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata) + Indexer.prepare(metadata) expected_fp8 = metadata.slot_mapping_fp8[:num_tokens].clone() expected_scale = metadata.slot_mapping_scale[:num_tokens].clone() @@ -1153,7 +1194,7 @@ def test_indexer_k_cache_scatter_custom_op(): from tensorrt_llm._torch.attention_backend.sparse.dsa import Indexer - Indexer.prepare_metadata(metadata) + Indexer.prepare(metadata) # Generate test data k_original = torch.randn((num_tokens, head_dim), device="cuda", dtype=torch.bfloat16) @@ -1314,7 +1355,7 @@ def test_fp8_k_cache_roundtrip(): num_tokens=total_tokens, indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata) + Indexer.prepare(metadata) # Generate unique patterns for each request and quantize k_original = torch.randn((total_tokens, head_dim), device="cuda", dtype=torch.bfloat16) @@ -1471,7 +1512,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, compres compress_ratio=compress_ratio, indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_context) + Indexer.prepare(metadata_context) k_context_fp8, k_context_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_context_bf16) @@ -1497,7 +1538,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, compres compress_ratio=compress_ratio, indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_gen) + Indexer.prepare(metadata_gen) k_gen_fp8, k_gen_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_gen_bf16) indexer._update_k_cache(k_gen_fp8, k_gen_scale, metadata_gen) @@ -1833,7 +1874,7 @@ def _force_dsl_expand_setup(meta): ) if not use_dsl: _force_direct_path(metadata_context) - Indexer.prepare_metadata(metadata_context) + Indexer.prepare(metadata_context) # Real path: split K at head_dim//2 + fused_cat_fp4 (mirrors # Indexer._prep_q_or_k at dsa.py:2046-2050). @@ -1867,7 +1908,7 @@ def _force_dsl_expand_setup(meta): # >1 = atom-split). Mirrors dsa.py's `if expand_for_dsl and # num_generations > 0` block which runs for any next_n ≥ 2. _force_dsl_expand_setup(metadata_gen) - Indexer.prepare_metadata(metadata_gen) + Indexer.prepare(metadata_gen) k_gen_fp4, k_gen_scale = torch.ops.trtllm.fused_cat_fp4( k_gen_bf16[:, :pe_dim].contiguous(), @@ -2553,7 +2594,7 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, compr compress_ratio=compress_ratio, ) - Indexer.prepare_metadata(metadata_chunked) + Indexer.prepare(metadata_chunked) assert metadata_chunked.indexer_prefill_chunks is not None num_chunks = len(metadata_chunked.indexer_prefill_chunks) @@ -2594,7 +2635,7 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, compr compress_ratio=compress_ratio, ) - Indexer.prepare_metadata(metadata_baseline) + Indexer.prepare(metadata_baseline) if metadata_baseline.indexer_prefill_chunks is not None: num_baseline_chunks = len(metadata_baseline.indexer_prefill_chunks) @@ -2806,7 +2847,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l max_draft_tokens=next_n - 1, indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_context) + Indexer.prepare(metadata_context) indexer._update_k_cache(k_context_fp8, k_context_scale, metadata_context) # Generate decode phase test data @@ -2832,7 +2873,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l max_draft_tokens=next_n - 1, indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_gen_write) + Indexer.prepare(metadata_gen_write) indexer._update_k_cache(k_fp8, k_scale, metadata_gen_write) # Test with custom CUDA kernel @@ -2852,7 +2893,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_custom) + Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) try: @@ -2879,7 +2920,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_fallback) + Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) _set_torch_top_k(indexer) topk_indices_fallback = indexer.sparse_attn_indexer( @@ -2905,7 +2946,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_skip) + Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) try: @@ -3008,7 +3049,7 @@ def make_inputs(n_tokens): kv_lens.sum().item(), max_draft_tokens=md, ) - Indexer.prepare_metadata(meta_ctx) + Indexer.prepare(meta_ctx) indexer._update_k_cache(ctx_k_fp8, ctx_k_scale, meta_ctx) step0_tokens = batch_size * step0_next_n @@ -3026,7 +3067,7 @@ def make_inputs(n_tokens): max_model_len, max_draft_tokens=md, ) - Indexer.prepare_metadata(meta0) + Indexer.prepare(meta0) # indexer_topk_decode needs caller-owned radix aux buffers for small gen batches. _radix_bp = 10 meta0.radix_aux_indices = torch.zeros( @@ -3051,7 +3092,7 @@ def make_inputs(n_tokens): max_model_len, max_draft_tokens=md, ) - Indexer.prepare_metadata(meta0) + Indexer.prepare(meta0) # context stash branch reads seq_lens_cuda (a read-only property); set its backing field. meta0._seq_lens_cuda = kv_lens.clone().cuda() @@ -3093,7 +3134,7 @@ def make_inputs(n_tokens): max_model_len, max_draft_tokens=md, ) - Indexer.prepare_metadata(meta) + Indexer.prepare(meta) meta.in_mtp_draft_loop = True meta.shared_topk_indices = stash meta.indexer_skip_topk = True @@ -3178,7 +3219,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_custom) + Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) assert metadata_custom.indexer_prefill_chunks is not None @@ -3206,7 +3247,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_fallback) + Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) _set_torch_top_k(indexer) topk_indices_fallback = indexer.sparse_attn_indexer( @@ -3287,7 +3328,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_custom) + Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) # Force single-pass path by setting indexer_prefill_chunks to None metadata_custom.indexer_prefill_chunks = None @@ -3315,7 +3356,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_fallback) + Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) # Force single-pass path by setting indexer_prefill_chunks to None metadata_fallback.indexer_prefill_chunks = None @@ -3341,7 +3382,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, enable_indexer_skip=True, indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_skip) + Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) metadata_skip.indexer_prefill_chunks = None @@ -3478,7 +3519,7 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): enable_indexer_skip=True, indexer_head_dim=head_dim, ) - Indexer.prepare_metadata(metadata_skip) + Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) topk_indices_skip = indexer.sparse_attn_indexer( metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights @@ -3726,7 +3767,7 @@ def test_cutedsl_mqa_logits_output_buffer_persistent(): index_topk=index_topk, use_cute_dsl_paged_mqa_logits=True, ) - Indexer.prepare_metadata(metadata) + Indexer.prepare(metadata) kv_cache = cache_manager.get_indexer_k_cache_buffers(0) q = torch.randn((batch_size, next_n, heads, head_dim), device="cuda", dtype=torch.bfloat16).to( diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py b/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py index 67b8b88831ca..65d02bf07fbc 100644 --- a/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py +++ b/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py @@ -368,7 +368,7 @@ def test_indexer_k_cache_scatter_custom_op_fp4(): from tensorrt_llm._torch.attention_backend.sparse.dsa import Indexer - Indexer.prepare_metadata(metadata) + Indexer.prepare(metadata) # FP4 packed data: [num_tokens, 64] int8; scale: [num_tokens, 1] int32 k_fp4 = torch.randint(-128, 127, (num_tokens, fp4_data_dim), device="cuda", dtype=torch.int8) diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py index 0696d72ebab3..40135270736c 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py @@ -260,32 +260,6 @@ def test_mla_backend_only_forward() -> None: ) -def test_mla_sparse_prepare_delegates_to_hooks() -> None: - mla = MLA.__new__(MLA) - torch.nn.Module.__init__(mla) - mla.sparse_attn_hooks = Mock() - attn_metadata = Mock() - - MLA.prepare_sparse_attn(mla, attn_metadata) - - mla.sparse_attn_hooks.prepare.assert_called_once_with(mla, attn_metadata) - - -@pytest.mark.parametrize("algorithm", ["dsa", "deepseek_v4"]) -def test_sparse_mla_hooks_prepare_the_indexer(algorithm: str) -> None: - hook_module = ModuleType(f"{algorithm}_prepare") - hook_module.sparse_params = MockSparseParams() - hook_module.sparse_params.algorithm = algorithm - hooks = get_sparse_mla_hooks(hook_module) - indexer = Mock() - mla = Mock(indexer=indexer) - attn_metadata = Mock() - - hooks.prepare(mla, attn_metadata) - - indexer.prepare.assert_called_once_with(attn_metadata) - - def test_sparse_runtime_params_without_prediction() -> None: attention = TrtllmAttention.__new__(TrtllmAttention) attention.sparse_params = MockSparseParams() diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index bf3b531994de..3cae8cbb186e 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -91,17 +91,49 @@ def forward(self, *args, **kwargs) -> torch.Tensor: return {"logits": torch.randn((batch_size, 10), device='cuda')} -def test_sparse_attention_modules_are_prepared_once_per_metadata(): - sparse_module = Mock() - sparse_module.prepare_sparse_attn = Mock() - engine = SimpleNamespace(model=Mock()) - engine.model.modules.return_value = [object(), sparse_module] - attn_metadata = SimpleNamespace() +def test_setup_attn_metadata_registers_unique_topk_indexers(): + + class SparseMetadata: + indexers = () + + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + indexer = SimpleNamespace(top_k=object()) + model = SimpleNamespace( + model_config=SimpleNamespace( + pretrained_config=SimpleNamespace( + architectures=["LlamaForCausalLM"], + num_attention_heads=8, + num_key_value_heads=2, + ), + enable_flash_mla=False, + ), + modules=lambda: [ + SimpleNamespace(indexer=indexer), + SimpleNamespace(indexer=indexer), + SimpleNamespace(indexer=None), + ], + ) + engine = SimpleNamespace( + model=model, + attn_runtime_features=SimpleNamespace(cache_reuse=False, + chunked_prefill=False), + cache_indirection_attention=None, + attn_backend=SimpleNamespace(Metadata=SparseMetadata), + sparse_attention_config=None, + encoder_attn_metadata=None, + attn_metadata=None, + batch_size=4, + max_num_tokens=32, + max_beam_width=1, + mapping=Mock(), + ) + cache_manager = Mock() - PyTorchModelEngine._prepare_sparse_attention_modules(engine, attn_metadata) - PyTorchModelEngine._prepare_sparse_attention_modules(engine, attn_metadata) + metadata = PyTorchModelEngine._set_up_attn_metadata(engine, cache_manager) - sparse_module.prepare_sparse_attn.assert_called_once_with(attn_metadata) + assert metadata.indexers == (indexer, ) class DummyMultimodalIndexModel(torch.nn.Module): diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 33e1c570878b..f621da005a4c 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for the reusable sparse index-selection Top-K module.""" -from contextlib import nullcontext from unittest.mock import Mock import pytest @@ -92,7 +91,7 @@ def test_one_module_dispatches_prefill_and_decode() -> None: assert output.item() == 1 -def test_cute_dsl_preferred_preserves_compressed_mtp_fallback(monkeypatch) -> None: +def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: cute_dsl = Mock() trtllm = Mock() monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_indexer_topk_decode", cute_dsl) @@ -100,7 +99,7 @@ def test_cute_dsl_preferred_preserves_compressed_mtp_fallback(monkeypatch) -> No top_k = TopK( 2, - decode_implementation=TopKImplementation.CUTE_DSL_PREFERRED, + decode_implementation=TopKImplementation.CUTE_DSL_RADIX, compress_ratio=4, ) logical_lengths = torch.tensor([16], dtype=torch.int32) @@ -170,8 +169,8 @@ def test_gvr_owns_prior_state_and_updates_it(monkeypatch) -> None: num_sms=16, max_num_requests=1, ) - assert top_k._prior_indices is not None - top_k._prior_indices.copy_(prior) + assert top_k._gvr_prior_indices is not None + top_k._gvr_prior_indices.copy_(prior) gvr.side_effect = lambda *args, **kwargs: output.copy_(torch.tensor([[5, 3]])) top_k( @@ -181,12 +180,11 @@ def test_gvr_owns_prior_state_and_updates_it(monkeypatch) -> None: sequence_lengths=logical_lengths, scan_lengths=scan_lengths, next_n=1, - max_num_columns=8, ) args, kwargs = gvr.call_args assert args[0] is scores - assert args[1].data_ptr() == top_k._prior_indices.data_ptr() + assert args[1].data_ptr() == top_k._gvr_prior_indices.data_ptr() assert args[2] is logical_lengths assert args[3] is output assert args[4] == 2 @@ -196,7 +194,7 @@ def test_gvr_owns_prior_state_and_updates_it(monkeypatch) -> None: "max_seq_len": 8, "order_row": None, } - assert top_k._prior_indices.tolist() == [[5, 3]] + assert top_k._gvr_prior_indices.tolist() == [[5, 3]] def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: @@ -222,7 +220,6 @@ def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: sequence_lengths=lengths, scan_lengths=lengths, next_n=next_n, - max_num_columns=8, ) row_order = gvr.call_args.kwargs["order_row"] @@ -248,80 +245,94 @@ def test_seed_from_prefill_uses_last_request_rows() -> None: request_offset=1, ) - assert top_k._prior_indices is not None - assert top_k._prior_indices.tolist() == [[0, 0], [2, 3], [4, 5], [0, 0]] + assert top_k._gvr_prior_indices is not None + assert top_k._gvr_prior_indices.tolist() == [[0, 0], [2, 3], [4, 5], [0, 0]] -@pytest.mark.parametrize("top_k", [0, -1]) -def test_top_k_must_be_positive(top_k: int) -> None: - with pytest.raises(ValueError, match="top_k must be positive"): - TopK(top_k, prefill_implementation=TopKImplementation.TORCH) +def test_implementations_are_named_by_backend_and_algorithm() -> None: + assert {implementation.value for implementation in TopKImplementation} == { + "torch", + "cuda_radix", + "cute_dsl_radix", + "cuda_gvr", + "cute_dsl_gvr", + } + +def test_none_implementations_use_cuda_radix_defaults() -> None: + top_k = TopK(1) -def test_top_k_requires_an_implementation() -> None: - with pytest.raises(ValueError, match="at least one Top-K implementation"): - TopK(1) + assert top_k.prefill_implementation == TopKImplementation.CUDA_RADIX + assert top_k.decode_implementation == TopKImplementation.CUDA_RADIX -def test_decode_validates_workspace_pairs() -> None: - top_k = TopK(2, decode_implementation=TopKImplementation.TRTLLM) - with pytest.raises(ValueError, match="must be provided together"): - top_k( - torch.randn(1, 4), - torch.empty(1, 2, dtype=torch.int32), - is_prefill=False, - sequence_lengths=torch.tensor([4], dtype=torch.int32), - scan_lengths=torch.tensor([4], dtype=torch.int32), - next_n=1, - radix_indices=torch.empty(1, 10, 2, dtype=torch.int32), - ) +def test_cuda_gvr_owns_scratch_and_updates_prior(monkeypatch) -> None: + top_k_module._warmup_decode_top_k.cache_clear() + decode = Mock(side_effect=lambda *args, **kwargs: args[2].copy_(torch.tensor([[3, 1]]))) + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) + + top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) + top_k.prepare( + device=torch.device("cpu"), + max_num_columns=8, + next_n=1, + input_dtype=torch.float32, + max_num_requests=2, + num_sms=4, + ) + scores = torch.randn(1, 8) + lengths = torch.tensor([8], dtype=torch.int32) + output = torch.empty(1, 2, dtype=torch.int32) + + top_k( + scores, + output, + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths, + ) + runtime_call = decode.call_args_list[-1] + assert runtime_call.kwargs["pre_idx"].data_ptr() == top_k._gvr_prior_indices.data_ptr() + assert runtime_call.kwargs["heuristic_scratch"].data_ptr() == top_k._cuda_gvr_scratch.data_ptr() + assert top_k._gvr_prior_indices.tolist() == [[3, 1], [0, 0]] -def test_prepare_deduplicates_success_for_full_key(monkeypatch) -> None: - prepared: set[tuple[object, ...]] = set() - monkeypatch.setattr(top_k_module, "_PREPARED_DECODE_TOP_K", prepared) - monkeypatch.setattr(top_k_module, "_cuda_device", lambda _: torch.device("cuda:3")) - monkeypatch.setattr(torch.cuda, "device", lambda _: nullcontext()) - top_k = TopK(32, decode_implementation=TopKImplementation.TRTLLM_HEURISTIC) - warmup = Mock() - monkeypatch.setattr(top_k, "_warmup_trtllm_heuristic", warmup) +def test_prepare_deduplicates_warmup(monkeypatch) -> None: + top_k_module._warmup_decode_top_k.cache_clear() + decode = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) prepare_args = dict( - device=torch.device("cuda:3"), + device=torch.device("cpu"), max_num_columns=4096, next_n=1, - input_dtype=torch.bfloat16, + input_dtype=torch.float32, + max_num_requests=2, num_sms=148, ) - top_k.prepare(**prepare_args) - top_k.prepare(**prepare_args) + TopK(32, decode_implementation=TopKImplementation.CUDA_GVR).prepare(**prepare_args) + TopK(32, decode_implementation=TopKImplementation.CUDA_GVR).prepare(**prepare_args) - warmup.assert_called_once_with(torch.bfloat16) - assert len(prepared) == 1 + decode.assert_called_once() def test_prepare_does_not_cache_failure(monkeypatch) -> None: - prepared: set[tuple[object, ...]] = set() - monkeypatch.setattr(top_k_module, "_PREPARED_DECODE_TOP_K", prepared) - monkeypatch.setattr(top_k_module, "_cuda_device", lambda _: torch.device("cuda:0")) - monkeypatch.setattr(torch.cuda, "device", lambda _: nullcontext()) - - top_k = TopK(16, decode_implementation=TopKImplementation.TRTLLM_HEURISTIC) - warmup = Mock(side_effect=RuntimeError("warmup failed")) - monkeypatch.setattr(top_k, "_warmup_trtllm_heuristic", warmup) + top_k_module._warmup_decode_top_k.cache_clear() + decode = Mock(side_effect=RuntimeError("warmup failed")) + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) + top_k = TopK(16, decode_implementation=TopKImplementation.CUDA_GVR) prepare_args = dict( - device=torch.device("cuda:0"), + device=torch.device("cpu"), max_num_columns=1024, next_n=1, input_dtype=torch.float32, + max_num_requests=1, ) with pytest.raises(RuntimeError, match="warmup failed"): top_k.prepare(**prepare_args) - assert not prepared - warmup.side_effect = None + decode.side_effect = None top_k.prepare(**prepare_args) - assert warmup.call_count == 2 - assert len(prepared) == 1 + assert decode.call_count == 2 From 6aa45a5a25fcba764258984368b8df36fda2d166 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Fri, 7 Aug 2026 09:42:13 +0000 Subject: [PATCH 06/18] [None][fix] keep singular sparse indexer metadata Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 4 +-- .../attention_backend/sparse/dsa/metadata.py | 5 +-- .../_torch/pyexecutor/model_engine.py | 22 +++++++++++-- .../attention/sparse/dsa/test_dsa_indexer.py | 32 +------------------ .../executor/test_pytorch_model_engine.py | 32 +++++++++++++++++-- 5 files changed, 53 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 2369d9db389a..e8f39527e100 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1146,8 +1146,8 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: # This is a preprocessing step that computes scheduling information for the kernel Indexer.prepare_scheduler_metadata(metadata) - for indexer in getattr(metadata, "indexers", ()): - indexer.top_k.prepare( + if metadata.indexer is not None: + metadata.indexer.top_k.prepare( device=metadata.kv_lens_cuda.device, max_num_columns=metadata.get_indexer_max_seq_len(), next_n=1 + metadata.max_draft_tokens, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index f3af752fd521..e4951b13a57a 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -48,7 +48,7 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" sparse_metadata_params: Optional[DSAMetadataParams] = None - indexers: tuple[Indexer, ...] = () + indexer: Optional["Indexer"] = None # Chunked prefill metadata for indexer (prefill-only, no CUDA graph needed) indexer_prefill_chunks: Optional[List[IndexerPrefillChunkMetadata]] = None # Max chunk size for two-level chunking: @@ -803,7 +803,6 @@ def update_spec_dec_param( num_contexts: int = 0, ): """Update speculative decoding parameters and create expanded buffers.""" - previous_max_draft_tokens = self.max_draft_tokens super().update_spec_dec_param( batch_size, is_spec_decoding_enabled, @@ -830,8 +829,6 @@ def update_spec_dec_param( # ("radix_aux_* must hold at least num_rows*blocks_per_row*index_topk # elements"). self._create_radix_aux_buffers(capture_graph=capture_graph) - if self.max_draft_tokens != previous_max_draft_tokens: - Indexer.prepare(metadata=self) def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor: """Refresh INDEX_KEY offsets and return their physical pool slots.""" diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 516f8b4f8ae6..64fd4950a442 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3230,15 +3230,30 @@ def _set_up_attn_metadata( num_heads_per_kv=num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, ) - if hasattr(self.attn_metadata, "indexers"): + if hasattr(self.attn_metadata, "indexer"): indexers = {} for module in self.model.modules(): indexer = getattr(module, "indexer", None) if indexer is not None and hasattr(indexer, "top_k"): indexers[id(indexer)] = indexer - self.attn_metadata.indexers = tuple(indexers.values()) + self._top_k_indexers = tuple(indexers.values()) + self.attn_metadata.indexer = (self._top_k_indexers[0] + if self._top_k_indexers else None) return self.attn_metadata + def _prepare_sparse_top_k_modules(self, + attn_metadata: AttentionMetadata) -> None: + """Prepare layer-local Top-K state before model forward.""" + for indexer in getattr(self, "_top_k_indexers", ()): + indexer.top_k.prepare( + device=attn_metadata.kv_lens_cuda.device, + max_num_columns=attn_metadata.get_indexer_max_seq_len(), + next_n=1 + attn_metadata.max_draft_tokens, + input_dtype=torch.float32, + num_sms=attn_metadata.num_sms, + max_num_requests=attn_metadata.max_num_sequences, + ) + @property def is_multimodal(self) -> bool: """True iff this engine drives a multimodal model. @@ -7508,6 +7523,9 @@ def forward(self, spec_resource_manager = None spec_metadata = None + if kv_cache_manager is not None: + self._prepare_sparse_top_k_modules(attn_metadata) + moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', None) if kv_cache_manager is None: diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 5806ae23dce4..74544b28a773 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -260,7 +260,7 @@ def test_indexer_prepare_keeps_metadata_and_top_k_preparation_together(): top_k = Mock() indexer = SimpleNamespace(top_k=top_k) metadata = SimpleNamespace( - indexers=(indexer,), + indexer=indexer, num_contexts=0, num_generations=0, num_ctx_tokens=0, @@ -291,36 +291,6 @@ def test_indexer_prepare_keeps_metadata_and_top_k_preparation_together(): ) -def test_draft_width_change_reprepares_indexer(): - metadata = object.__new__(DSAtrtllmAttentionMetadata) - metadata.max_draft_tokens = 0 - metadata.max_num_sequences = 2 - metadata.is_cuda_graph = False - metadata.kv_lens_cuda_2d = torch.empty((2, 1)) - metadata.kv_lens_expanded_host = torch.empty(2) - metadata._create_kv_lens_2d_buffer = Mock() - metadata.create_expanded_buffers = Mock() - metadata._create_radix_aux_buffers = Mock() - - with ( - patch( - "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata." - "TrtllmAttentionMetadata.update_spec_dec_param" - ), - patch.object(Indexer, "prepare") as prepare, - ): - metadata.update_spec_dec_param( - batch_size=2, - is_spec_decoding_enabled=True, - is_spec_dec_tree=False, - is_spec_dec_dynamic_tree=False, - max_draft_len=3, - max_total_draft_tokens=3, - ) - - prepare.assert_called_once_with(metadata=metadata) - - def _ceil_to_ue8m0(x: torch.Tensor): """Round tensor values up to the nearest power of two (UE8M0 format).""" return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 3cae8cbb186e..2c7b1511c5e3 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -91,10 +91,10 @@ def forward(self, *args, **kwargs) -> torch.Tensor: return {"logits": torch.randn((batch_size, 10), device='cuda')} -def test_setup_attn_metadata_registers_unique_topk_indexers(): +def test_setup_attn_metadata_registers_single_indexer(): class SparseMetadata: - indexers = () + indexer = None def __init__(self, **kwargs): self.__dict__.update(kwargs) @@ -133,7 +133,33 @@ def __init__(self, **kwargs): metadata = PyTorchModelEngine._set_up_attn_metadata(engine, cache_manager) - assert metadata.indexers == (indexer, ) + assert metadata.indexer is indexer + assert engine._top_k_indexers == (indexer, ) + + +def test_prepare_sparse_top_k_modules_prepares_every_indexer(): + top_ks = (Mock(), Mock()) + engine = SimpleNamespace(_top_k_indexers=tuple( + SimpleNamespace(top_k=top_k) for top_k in top_ks)) + metadata = SimpleNamespace( + kv_lens_cuda=torch.empty(0), + get_indexer_max_seq_len=Mock(return_value=4096), + max_draft_tokens=3, + num_sms=148, + max_num_sequences=32, + ) + + PyTorchModelEngine._prepare_sparse_top_k_modules(engine, metadata) + + for top_k in top_ks: + top_k.prepare.assert_called_once_with( + device=metadata.kv_lens_cuda.device, + max_num_columns=4096, + next_n=4, + input_dtype=torch.float32, + num_sms=148, + max_num_requests=32, + ) class DummyMultimodalIndexModel(torch.nn.Module): From 2356231a92da8db9cbe75cb3bc15e222121274d6 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:25:52 +0000 Subject: [PATCH 07/18] [None][refactor] remove sparse top-k prepare lifecycle Keep GVR state within each TopK module and materialize capacity-dependent buffers from prefill seeding or decode forward. Remove engine and shared-metadata coupling, along with redundant kernel warmup calls. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 10 - .../attention_backend/sparse/dsa/metadata.py | 1 - tensorrt_llm/_torch/modules/top_k.py | 172 +++++++----------- .../_torch/pyexecutor/model_engine.py | 25 --- .../attention/sparse/dsa/test_dsa_indexer.py | 18 +- .../executor/test_pytorch_model_engine.py | 71 -------- tests/unittest/_torch/modules/test_top_k.py | 98 ++-------- 7 files changed, 86 insertions(+), 309 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index e8f39527e100..66e9f347b9d6 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1146,16 +1146,6 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: # This is a preprocessing step that computes scheduling information for the kernel Indexer.prepare_scheduler_metadata(metadata) - if metadata.indexer is not None: - metadata.indexer.top_k.prepare( - device=metadata.kv_lens_cuda.device, - max_num_columns=metadata.get_indexer_max_seq_len(), - next_n=1 + metadata.max_draft_tokens, - input_dtype=torch.float32, - num_sms=metadata.num_sms, - max_num_requests=metadata.max_num_sequences, - ) - def _update_k_cache( self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: DSAtrtllmAttentionMetadata ) -> None: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index e4951b13a57a..db51ac877c1e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -48,7 +48,6 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" sparse_metadata_params: Optional[DSAMetadataParams] = None - indexer: Optional["Indexer"] = None # Chunked prefill metadata for indexer (prefill-only, no CUDA graph needed) indexer_prefill_chunks: Optional[List[IndexerPrefillChunkMetadata]] = None # Max chunk size for two-level chunking: diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index fe74572c6231..3f83ae22da3c 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -5,7 +5,6 @@ from __future__ import annotations from enum import Enum -from functools import cache import torch import torch.nn as nn @@ -25,59 +24,6 @@ class TopKImplementation(str, Enum): TopKImplementation.CUDA_GVR, TopKImplementation.CUTE_DSL_GVR, } -_GVR_WARMUP_COLUMNS = 4096 -_RADIX_MAX_BLOCKS_PER_ROW = 10 - - -@cache -def _warmup_decode_top_k( - implementation: TopKImplementation, - device: torch.device, - input_dtype: torch.dtype, - top_k: int, - max_num_columns: int, - next_n: int, - num_sms: int | None, - compress_ratio: int, -) -> None: - if implementation == TopKImplementation.CUDA_GVR: - num_columns = max(_GVR_WARMUP_COLUMNS, top_k) - scores = torch.zeros((1, num_columns), dtype=input_dtype, device=device) - sequence_lengths = torch.tensor([num_columns], dtype=torch.int32, device=device) - output_indices = torch.empty((1, top_k), dtype=torch.int32, device=device) - prior_indices = torch.zeros((1, top_k), dtype=torch.int32, device=device) - heuristic_values = torch.empty((1, top_k), dtype=input_dtype, device=device) - radix_indices = torch.empty( - (1, _RADIX_MAX_BLOCKS_PER_ROW, top_k), dtype=torch.int32, device=device - ) - radix_values = torch.empty( - (1, _RADIX_MAX_BLOCKS_PER_ROW, top_k), dtype=torch.float32, device=device - ) - torch.ops.trtllm.indexer_topk_decode( - scores, - sequence_lengths, - output_indices, - 1, - top_k, - pre_idx=prior_indices, - heuristic_scratch=heuristic_values, - radix_aux_indices=radix_indices, - radix_aux_logits=radix_values, - ) - elif implementation == TopKImplementation.CUTE_DSL_RADIX and not ( - compress_ratio > 1 and next_n > 1 - ): - from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( - warmup_cute_dsl_radix_topk_decode, - ) - - warmup_cute_dsl_radix_topk_decode( - top_k=top_k, - num_cols=max_num_columns, - next_n=next_n, - dtype=input_dtype, - num_sms=num_sms, - ) class TopK(nn.Module): @@ -100,52 +46,22 @@ def __init__( decode_implementation or TopKImplementation.CUDA_RADIX ) self.compress_ratio = compress_ratio - self.register_buffer("_gvr_prior_indices", None, persistent=False) - self.register_buffer("_cuda_gvr_scratch", None, persistent=False) - self.register_buffer("_gvr_row_order", None, persistent=False) - self._num_sms: int | None = None - - def prepare( - self, - *, - device: torch.device, - max_num_columns: int, - next_n: int, - input_dtype: torch.dtype, - max_num_requests: int, - num_sms: int | None = None, - ) -> None: - """Allocate GVR state and warm up the configured decode implementation.""" - implementation = self.decode_implementation - if implementation in _GVR_IMPLEMENTATIONS: - if self._gvr_prior_indices is None: - self._gvr_prior_indices = torch.zeros( - (max_num_requests, self.top_k), dtype=torch.int32, device=device - ) - - if implementation == TopKImplementation.CUDA_GVR: - scratch_shape = (max_num_requests * next_n, self.top_k) - if self._cuda_gvr_scratch is None or self._cuda_gvr_scratch.shape != scratch_shape: - self._cuda_gvr_scratch = torch.empty( - scratch_shape, dtype=input_dtype, device=device - ) - else: - self._num_sms = num_sms - if self._gvr_row_order is None: - self._gvr_row_order = torch.empty( - (max_num_requests,), dtype=torch.int32, device=device - ) - - _warmup_decode_top_k( - implementation, - device, - input_dtype, - self.top_k, - max_num_columns, - next_n, - num_sms, - self.compress_ratio, + self.register_buffer( + "_gvr_prior_indices", + torch.empty((0, top_k), dtype=torch.int32), + persistent=False, + ) + self.register_buffer( + "_cuda_gvr_scratch", + torch.empty((0, top_k)), + persistent=False, ) + self.register_buffer( + "_gvr_row_order", + torch.empty((0,), dtype=torch.int32), + persistent=False, + ) + self._num_sms = 0 def forward( self, @@ -185,10 +101,19 @@ def seed_from_prefill( request_offset: int = 0, ) -> None: """Seed GVR from the last selected row of each prefill request.""" - if self._gvr_prior_indices is None: + if self.decode_implementation not in _GVR_IMPLEMENTATIONS: return last_rows = (torch.cumsum(request_lengths, dim=0) - 1).to(dtype=torch.long) num_requests = request_lengths.shape[0] + required_rows = request_offset + num_requests + if ( + self._gvr_prior_indices.device != output_indices.device + or self._gvr_prior_indices.shape[0] < required_rows + ): + prior_indices = output_indices.new_zeros((required_rows, self.top_k), dtype=torch.int32) + if self._gvr_prior_indices.device == output_indices.device: + prior_indices[: self._gvr_prior_indices.shape[0]].copy_(self._gvr_prior_indices) + self._gvr_prior_indices = prior_indices self._gvr_prior_indices[request_offset : request_offset + num_requests].copy_( output_indices[last_rows] ) @@ -234,6 +159,9 @@ def _forward_decode( if implementation == TopKImplementation.TORCH: return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n) + if implementation in _GVR_IMPLEMENTATIONS: + self._ensure_gvr_state(scores, sequence_lengths, next_n, radix_indices) + use_cuda = implementation in ( TopKImplementation.CUDA_RADIX, TopKImplementation.CUDA_GVR, @@ -246,8 +174,6 @@ def _forward_decode( prior_indices = None heuristic_values = None if implementation == TopKImplementation.CUDA_GVR: - assert self._gvr_prior_indices is not None - assert self._cuda_gvr_scratch is not None prior_indices = self._gvr_prior_indices[: sequence_lengths.shape[0]] heuristic_values = self._cuda_gvr_scratch[: scores.shape[0]] torch.ops.trtllm.indexer_topk_decode( @@ -276,7 +202,6 @@ def _forward_decode( ) return output_indices - assert self._gvr_prior_indices is not None prior_indices = self._gvr_prior_indices[: sequence_lengths.shape[0]] torch.ops.trtllm.cute_dsl_gvr_topk_decode( scores, @@ -292,12 +217,51 @@ def _forward_decode( prior_indices.copy_(output_indices[next_n - 1 :: next_n]) return output_indices + def _ensure_gvr_state( + self, + scores: torch.Tensor, + sequence_lengths: torch.Tensor, + next_n: int, + radix_indices: torch.Tensor | None, + ) -> None: + row_capacity = radix_indices.shape[0] if radix_indices is not None else scores.shape[0] + request_capacity = max(sequence_lengths.shape[0], row_capacity // next_n) + + if ( + self._gvr_prior_indices.device != scores.device + or self._gvr_prior_indices.shape[0] < request_capacity + ): + prior_indices = scores.new_zeros((request_capacity, self.top_k), dtype=torch.int32) + if self._gvr_prior_indices.device == scores.device: + prior_indices[: self._gvr_prior_indices.shape[0]].copy_(self._gvr_prior_indices) + self._gvr_prior_indices = prior_indices + + if self.decode_implementation == TopKImplementation.CUDA_GVR: + if ( + self._cuda_gvr_scratch.device != scores.device + or self._cuda_gvr_scratch.dtype != scores.dtype + or self._cuda_gvr_scratch.shape[0] < row_capacity + ): + self._cuda_gvr_scratch = scores.new_empty((row_capacity, self.top_k)) + return + + if ( + self._gvr_row_order.device != scores.device + or self._gvr_row_order.shape[0] < request_capacity + ): + self._gvr_row_order = scores.new_empty((request_capacity,), dtype=torch.int32) + if self._num_sms == 0: + self._num_sms = ( + torch.cuda.get_device_properties(scores.device).multi_processor_count + if scores.is_cuda + else 1 + ) + def _prepare_row_order( self, sequence_lengths: torch.Tensor, next_n: int, ) -> torch.Tensor | None: - assert self._num_sms is not None and self._gvr_row_order is not None num_requests = sequence_lengths.shape[0] if num_requests * next_n < 2 * self._num_sms: return None diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 64fd4950a442..f3b90b487a9a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3230,30 +3230,8 @@ def _set_up_attn_metadata( num_heads_per_kv=num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, ) - if hasattr(self.attn_metadata, "indexer"): - indexers = {} - for module in self.model.modules(): - indexer = getattr(module, "indexer", None) - if indexer is not None and hasattr(indexer, "top_k"): - indexers[id(indexer)] = indexer - self._top_k_indexers = tuple(indexers.values()) - self.attn_metadata.indexer = (self._top_k_indexers[0] - if self._top_k_indexers else None) return self.attn_metadata - def _prepare_sparse_top_k_modules(self, - attn_metadata: AttentionMetadata) -> None: - """Prepare layer-local Top-K state before model forward.""" - for indexer in getattr(self, "_top_k_indexers", ()): - indexer.top_k.prepare( - device=attn_metadata.kv_lens_cuda.device, - max_num_columns=attn_metadata.get_indexer_max_seq_len(), - next_n=1 + attn_metadata.max_draft_tokens, - input_dtype=torch.float32, - num_sms=attn_metadata.num_sms, - max_num_requests=attn_metadata.max_num_sequences, - ) - @property def is_multimodal(self) -> bool: """True iff this engine drives a multimodal model. @@ -7523,9 +7501,6 @@ def forward(self, spec_resource_manager = None spec_metadata = None - if kv_cache_manager is not None: - self._prepare_sparse_top_k_modules(attn_metadata) - moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', None) if kv_cache_manager is None: diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 74544b28a773..6a95fbd986b7 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -256,21 +256,13 @@ def test_indexer_configures_one_top_k_module(): assert not hasattr(indexer, "decode_top_k") -def test_indexer_prepare_keeps_metadata_and_top_k_preparation_together(): - top_k = Mock() - indexer = SimpleNamespace(top_k=top_k) +def test_indexer_prepare_updates_metadata_without_layer_state(): metadata = SimpleNamespace( - indexer=indexer, num_contexts=0, num_generations=0, num_ctx_tokens=0, seq_lens=torch.empty(0, dtype=torch.int32), compress_ratios=[1], - kv_lens_cuda=torch.empty(0), - get_indexer_max_seq_len=Mock(return_value=4096), - max_draft_tokens=3, - num_sms=148, - max_num_sequences=32, ) indexer_params = SimpleNamespace(new_kv_tokens=torch.empty(0, dtype=torch.int32)) @@ -281,14 +273,6 @@ def test_indexer_prepare_keeps_metadata_and_top_k_preparation_together(): Indexer.prepare(metadata) prepare_metadata.assert_called_once_with(metadata, indexer_params) - top_k.prepare.assert_called_once_with( - device=metadata.kv_lens_cuda.device, - max_num_columns=4096, - next_n=4, - input_dtype=torch.float32, - num_sms=148, - max_num_requests=32, - ) def _ceil_to_ue8m0(x: torch.Tensor): diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 2c7b1511c5e3..a323f89f3158 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -91,77 +91,6 @@ def forward(self, *args, **kwargs) -> torch.Tensor: return {"logits": torch.randn((batch_size, 10), device='cuda')} -def test_setup_attn_metadata_registers_single_indexer(): - - class SparseMetadata: - indexer = None - - def __init__(self, **kwargs): - self.__dict__.update(kwargs) - - indexer = SimpleNamespace(top_k=object()) - model = SimpleNamespace( - model_config=SimpleNamespace( - pretrained_config=SimpleNamespace( - architectures=["LlamaForCausalLM"], - num_attention_heads=8, - num_key_value_heads=2, - ), - enable_flash_mla=False, - ), - modules=lambda: [ - SimpleNamespace(indexer=indexer), - SimpleNamespace(indexer=indexer), - SimpleNamespace(indexer=None), - ], - ) - engine = SimpleNamespace( - model=model, - attn_runtime_features=SimpleNamespace(cache_reuse=False, - chunked_prefill=False), - cache_indirection_attention=None, - attn_backend=SimpleNamespace(Metadata=SparseMetadata), - sparse_attention_config=None, - encoder_attn_metadata=None, - attn_metadata=None, - batch_size=4, - max_num_tokens=32, - max_beam_width=1, - mapping=Mock(), - ) - cache_manager = Mock() - - metadata = PyTorchModelEngine._set_up_attn_metadata(engine, cache_manager) - - assert metadata.indexer is indexer - assert engine._top_k_indexers == (indexer, ) - - -def test_prepare_sparse_top_k_modules_prepares_every_indexer(): - top_ks = (Mock(), Mock()) - engine = SimpleNamespace(_top_k_indexers=tuple( - SimpleNamespace(top_k=top_k) for top_k in top_ks)) - metadata = SimpleNamespace( - kv_lens_cuda=torch.empty(0), - get_indexer_max_seq_len=Mock(return_value=4096), - max_draft_tokens=3, - num_sms=148, - max_num_sequences=32, - ) - - PyTorchModelEngine._prepare_sparse_top_k_modules(engine, metadata) - - for top_k in top_ks: - top_k.prepare.assert_called_once_with( - device=metadata.kv_lens_cuda.device, - max_num_columns=4096, - next_n=4, - input_dtype=torch.float32, - num_sms=148, - max_num_requests=32, - ) - - class DummyMultimodalIndexModel(torch.nn.Module): class Config: diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index f621da005a4c..514884cf7a78 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -4,10 +4,8 @@ from unittest.mock import Mock -import pytest import torch -from tensorrt_llm._torch.modules import top_k as top_k_module from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation @@ -160,17 +158,6 @@ def test_gvr_owns_prior_state_and_updates_it(monkeypatch) -> None: logical_lengths = torch.tensor([32], dtype=torch.int32) scan_lengths = torch.tensor([8], dtype=torch.int32) output = torch.empty(1, 2, dtype=torch.int32) - prior = torch.zeros(1, 2, dtype=torch.int32) - top_k.prepare( - device=torch.device("cpu"), - max_num_columns=8, - next_n=1, - input_dtype=torch.float32, - num_sms=16, - max_num_requests=1, - ) - assert top_k._gvr_prior_indices is not None - top_k._gvr_prior_indices.copy_(prior) gvr.side_effect = lambda *args, **kwargs: output.copy_(torch.tensor([[5, 3]])) top_k( @@ -201,17 +188,8 @@ def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: gvr = Mock(side_effect=lambda *args, **kwargs: args[3].zero_()) monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) - num_sms = 4 next_n = 2 lengths = torch.tensor([4, 1, 8, 2], dtype=torch.int32) - top_k.prepare( - device=torch.device("cpu"), - max_num_columns=8, - next_n=next_n, - input_dtype=torch.float32, - num_sms=num_sms, - max_num_requests=lengths.shape[0], - ) top_k( torch.randn(lengths.shape[0] * next_n, 8), @@ -229,14 +207,6 @@ def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: def test_seed_from_prefill_uses_last_request_rows() -> None: top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) - top_k.prepare( - device=torch.device("cpu"), - max_num_columns=8, - next_n=1, - input_dtype=torch.float32, - num_sms=4, - max_num_requests=4, - ) prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32) top_k.seed_from_prefill( @@ -245,8 +215,19 @@ def test_seed_from_prefill_uses_last_request_rows() -> None: request_offset=1, ) - assert top_k._gvr_prior_indices is not None - assert top_k._gvr_prior_indices.tolist() == [[0, 0], [2, 3], [4, 5], [0, 0]] + assert top_k._gvr_prior_indices.tolist() == [[0, 0], [2, 3], [4, 5]] + + +def test_gvr_state_buffers_are_registered_during_init() -> None: + top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) + + buffers = dict(top_k.named_buffers()) + assert set(buffers) == { + "_gvr_prior_indices", + "_cuda_gvr_scratch", + "_gvr_row_order", + } + assert all(buffer.numel() == 0 for buffer in buffers.values()) def test_implementations_are_named_by_backend_and_algorithm() -> None: @@ -267,22 +248,15 @@ def test_none_implementations_use_cuda_radix_defaults() -> None: def test_cuda_gvr_owns_scratch_and_updates_prior(monkeypatch) -> None: - top_k_module._warmup_decode_top_k.cache_clear() decode = Mock(side_effect=lambda *args, **kwargs: args[2].copy_(torch.tensor([[3, 1]]))) monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) - top_k.prepare( - device=torch.device("cpu"), - max_num_columns=8, - next_n=1, - input_dtype=torch.float32, - max_num_requests=2, - num_sms=4, - ) scores = torch.randn(1, 8) lengths = torch.tensor([8], dtype=torch.int32) output = torch.empty(1, 2, dtype=torch.int32) + radix_indices = torch.empty(2, 10, 2, dtype=torch.int32) + radix_values = torch.empty(2, 10, 2) top_k( scores, @@ -290,49 +264,11 @@ def test_cuda_gvr_owns_scratch_and_updates_prior(monkeypatch) -> None: is_prefill=False, sequence_lengths=lengths, scan_lengths=lengths, + radix_indices=radix_indices, + radix_values=radix_values, ) runtime_call = decode.call_args_list[-1] assert runtime_call.kwargs["pre_idx"].data_ptr() == top_k._gvr_prior_indices.data_ptr() assert runtime_call.kwargs["heuristic_scratch"].data_ptr() == top_k._cuda_gvr_scratch.data_ptr() assert top_k._gvr_prior_indices.tolist() == [[3, 1], [0, 0]] - - -def test_prepare_deduplicates_warmup(monkeypatch) -> None: - top_k_module._warmup_decode_top_k.cache_clear() - decode = Mock() - monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) - - prepare_args = dict( - device=torch.device("cpu"), - max_num_columns=4096, - next_n=1, - input_dtype=torch.float32, - max_num_requests=2, - num_sms=148, - ) - TopK(32, decode_implementation=TopKImplementation.CUDA_GVR).prepare(**prepare_args) - TopK(32, decode_implementation=TopKImplementation.CUDA_GVR).prepare(**prepare_args) - - decode.assert_called_once() - - -def test_prepare_does_not_cache_failure(monkeypatch) -> None: - top_k_module._warmup_decode_top_k.cache_clear() - decode = Mock(side_effect=RuntimeError("warmup failed")) - monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) - top_k = TopK(16, decode_implementation=TopKImplementation.CUDA_GVR) - prepare_args = dict( - device=torch.device("cpu"), - max_num_columns=1024, - next_n=1, - input_dtype=torch.float32, - max_num_requests=1, - ) - - with pytest.raises(RuntimeError, match="warmup failed"): - top_k.prepare(**prepare_args) - - decode.side_effect = None - top_k.prepare(**prepare_args) - assert decode.call_count == 2 From ac38ccb5b376b4b7b8aa4c7d8a5798c9307930ab Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:37:37 +0000 Subject: [PATCH 08/18] [None][fix] warm up sparse top-k before graph capture Initialize CUDA GVR dispatcher hardware caches through the C++ custom-op helper before model warmup. Route sparse Top-K warmup through shared DSA metadata, retain CuTe radix precompilation, and let eager model warmup compile CuTe DSL GVR. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/metadata.py | 27 +++++++++ .../_torch/custom_ops/cpp_custom_ops.py | 38 ++++++++++++ .../_torch/pyexecutor/model_engine.py | 7 +++ .../attention/sparse/dsa/test_dsa_indexer.py | 58 +++++++++++++++++++ .../executor/test_pytorch_model_engine.py | 13 +++++ tests/unittest/_torch/modules/test_top_k.py | 39 +++++++++++++ 6 files changed, 182 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index db51ac877c1e..f4d6d71f366f 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -299,6 +299,33 @@ def get_indexer_max_seq_len(self) -> int: return self.kv_cache_manager.max_seq_len return max(1, self.kv_cache_manager.max_seq_len // self._indexer_compress_ratio) + def warmup_top_k(self, next_n: int) -> None: + """Warm up the configured decode Top-K implementation.""" + sparse_params = self.sparse_metadata_params + use_gvr = sparse_params.enable_heuristic_topk and get_sm_version() >= 100 + if use_gvr: + if self.use_cute_dsl_topk: + # The regular eager attention warmup compiles CuTe DSL GVR. + return + from tensorrt_llm._torch.custom_ops.cpp_custom_ops import warmup_cuda_gvr_topk_decode + + warmup_cuda_gvr_topk_decode(top_k=self.num_sparse_topk) + return + + if not self.use_cute_dsl_topk or (self._indexer_compress_ratio > 1 and next_n > 1): + return + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + warmup_cute_dsl_radix_topk_decode, + ) + + warmup_cute_dsl_radix_topk_decode( + top_k=self.num_sparse_topk, + num_cols=self.get_indexer_max_seq_len(), + next_n=next_n, + dtype=torch.float32, + num_sms=self.num_sms, + ) + def on_update_kv_lens(self): # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas. # Especially for the changes in the _preprocess_inputs() of model_engine.py. diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 3feac32cee17..46d4137197ad 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from typing import List, Optional, Tuple import torch @@ -1606,3 +1609,38 @@ def _(like: torch.Tensor, out_shape = shape if shape is not None else list(like.shape) dtype = out_dtype if out_dtype is not None else like.dtype return like.new_empty(out_shape, dtype=dtype), output_buffer_kind + + +def warmup_cuda_gvr_topk_decode(top_k: int = 2048) -> None: + """Initialize CUDA GVR before CUDA Graph capture. + + The first dispatcher call queries and caches device SM and L2 attributes. + """ + num_columns = 4096 + max_blocks_per_row = 10 + device = torch.device("cuda") + logits = torch.zeros((1, num_columns), dtype=torch.float32, device=device) + sequence_lengths = torch.tensor([num_columns], + dtype=torch.int32, + device=device) + output_indices = torch.empty((1, top_k), dtype=torch.int32, device=device) + prior_indices = torch.zeros((1, top_k), dtype=torch.int32, device=device) + scratch_values = torch.empty((1, top_k), dtype=torch.float32, device=device) + radix_indices = torch.empty((1, max_blocks_per_row, top_k), + dtype=torch.int32, + device=device) + radix_values = torch.empty((1, max_blocks_per_row, top_k), + dtype=torch.float32, + device=device) + torch.ops.trtllm.indexer_topk_decode( + logits, + sequence_lengths, + output_indices, + 1, + top_k, + pre_idx=prior_indices, + heuristic_scratch=scratch_values, + radix_aux_indices=radix_indices, + radix_aux_logits=radix_values, + ) + torch.cuda.synchronize() diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index f3b90b487a9a..8c1904716ee5 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1507,6 +1507,13 @@ def warmup(self, resource_manager: ResourceManager) -> None: self.cuda_graph_runner.preallocate_padding_dummies(resource_manager) log_mem_snapshot("warmup/after_preallocate_padding_dummies") + def _warmup_sparse_top_k(self) -> None: + """Warm up DSA Top-K kernels before any model forward or graph capture.""" + from ..attention_backend.sparse.dsa import DSAtrtllmAttentionMetadata + + if isinstance(self.attn_metadata, DSAtrtllmAttentionMetadata): + self.attn_metadata.warmup_top_k(1 + self.original_max_draft_len) + def _warmup_dg_paged_mqa_logits_metadata(self) -> None: """Pre-compile DeepGEMM's `get_paged_mqa_logits_metadata` helper for every 32-aligned batch bucket the runtime can produce. diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 6a95fbd986b7..6db490516635 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -122,6 +122,64 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): assert metadata._tokens_per_block == 64 +@pytest.mark.parametrize( + "enable_heuristic,use_cute_dsl,sm_version,compress_ratio,next_n,expected", + [ + (True, False, 100, 1, 1, "cuda_gvr"), + (True, True, 100, 1, 1, None), + (False, True, 100, 1, 1, "cute_dsl_radix"), + (True, True, 90, 1, 1, "cute_dsl_radix"), + (False, True, 100, 4, 2, None), + (False, False, 100, 1, 1, None), + ], +) +def test_metadata_warmup_top_k_dispatches_configured_implementation( + enable_heuristic, + use_cute_dsl, + sm_version, + compress_ratio, + next_n, + expected, +): + metadata = SimpleNamespace( + sparse_metadata_params=SimpleNamespace(enable_heuristic_topk=enable_heuristic), + use_cute_dsl_topk=use_cute_dsl, + num_sparse_topk=512, + _indexer_compress_ratio=compress_ratio, + get_indexer_max_seq_len=Mock(return_value=32768), + num_sms=148, + ) + + with ( + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.get_sm_version", + return_value=sm_version, + ), + patch( + "tensorrt_llm._torch.custom_ops.cpp_custom_ops.warmup_cuda_gvr_topk_decode" + ) as cuda_gvr, + patch( + "tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops.warmup_cute_dsl_radix_topk_decode" + ) as cute_dsl_radix, + ): + DSAtrtllmAttentionMetadata.warmup_top_k(metadata, next_n) + + if expected == "cuda_gvr": + cuda_gvr.assert_called_once_with(top_k=512) + else: + cuda_gvr.assert_not_called() + if expected == "cute_dsl_radix": + cute_dsl_radix.assert_called_once_with( + top_k=512, + num_cols=32768, + next_n=next_n, + dtype=torch.float32, + num_sms=148, + ) + else: + cute_dsl_radix.assert_not_called() + + def test_shared_topk_lifecycle(): sparse_config = DeepSeekSparseAttentionConfig( index_n_heads=1, diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index a323f89f3158..7aed76d3ee70 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -91,6 +91,19 @@ def forward(self, *args, **kwargs) -> torch.Tensor: return {"logits": torch.randn((batch_size, 10), device='cuda')} +def test_sparse_top_k_warmup_uses_shared_metadata(): + from tensorrt_llm._torch.attention_backend.sparse.dsa import \ + DSAtrtllmAttentionMetadata + + metadata = object.__new__(DSAtrtllmAttentionMetadata) + metadata.warmup_top_k = Mock() + engine = SimpleNamespace(attn_metadata=metadata, original_max_draft_len=3) + + PyTorchModelEngine._warmup_sparse_top_k(engine) + + metadata.warmup_top_k.assert_called_once_with(4) + + class DummyMultimodalIndexModel(torch.nn.Module): class Config: diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 514884cf7a78..75eb263dc265 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -230,6 +230,45 @@ def test_gvr_state_buffers_are_registered_during_init() -> None: assert all(buffer.numel() == 0 for buffer in buffers.values()) +def test_cuda_gvr_warmup_calls_cpp_op(monkeypatch) -> None: + from tensorrt_llm._torch.custom_ops import cpp_custom_ops + + decode = Mock() + synchronize = Mock() + torch_empty = torch.empty + torch_tensor = torch.tensor + torch_zeros = torch.zeros + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) + monkeypatch.setattr(torch.cuda, "synchronize", synchronize) + monkeypatch.setattr( + torch, + "empty", + lambda shape, *, dtype, device: torch_empty(shape, dtype=dtype), + ) + monkeypatch.setattr( + torch, + "tensor", + lambda data, *, dtype, device: torch_tensor(data, dtype=dtype), + ) + monkeypatch.setattr( + torch, + "zeros", + lambda shape, *, dtype, device: torch_zeros(shape, dtype=dtype), + ) + + cpp_custom_ops.warmup_cuda_gvr_topk_decode(top_k=512) + + args = decode.call_args.args + kwargs = decode.call_args.kwargs + assert args[0].shape == (1, 4096) + assert args[2].shape == (1, 512) + assert kwargs["pre_idx"].shape == (1, 512) + assert kwargs["heuristic_scratch"].shape == (1, 512) + assert kwargs["radix_aux_indices"].shape == (1, 10, 512) + assert kwargs["radix_aux_logits"].shape == (1, 10, 512) + synchronize.assert_called_once_with() + + def test_implementations_are_named_by_backend_and_algorithm() -> None: assert {implementation.value for implementation in TopKImplementation} == { "torch", From c456b10090d868ee80a7fb555878784c9a08de95 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 12 Aug 2026 21:52:29 -0700 Subject: [PATCH 09/18] [None][refactor] simplify sparse top-k GVR lifecycle Keep GVR state management inside the TopK module, separate radix and GVR decode paths, and rely on model-engine execution to initialize CUDA GVR. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 4 +- .../attention_backend/sparse/dsa/metadata.py | 18 +- .../_torch/custom_ops/cpp_custom_ops.py | 35 ---- tensorrt_llm/_torch/modules/top_k.py | 198 ++++++++++-------- .../_torch/pyexecutor/model_engine.py | 2 +- .../attention/sparse/dsa/test_dsa_indexer.py | 27 +-- tests/unittest/_torch/modules/test_top_k.py | 43 +--- 7 files changed, 129 insertions(+), 198 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 66e9f347b9d6..37f98f66b8c1 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1517,9 +1517,9 @@ def sparse_attn_indexer( :num_ctx_tokens, : ] - # Seed GVR state after the active generation slots. + # Chunked prefill is final only after its TP all-gathers; update GVR prior once here. if has_prefill and not metadata.skip_indexer_for_ctx_reqs: - self.top_k.seed_from_prefill( + self.top_k.update_gvr_prior_from_prefill( topk_indices_buffer[:num_ctx_tokens], metadata.seq_lens[:num_contexts], request_offset=num_generations, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index f4d6d71f366f..346e74142fd4 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -300,19 +300,13 @@ def get_indexer_max_seq_len(self) -> int: return max(1, self.kv_cache_manager.max_seq_len // self._indexer_compress_ratio) def warmup_top_k(self, next_n: int) -> None: - """Warm up the configured decode Top-K implementation.""" + """Pre-compile CuTe DSL radix variants not covered by engine warmup.""" sparse_params = self.sparse_metadata_params - use_gvr = sparse_params.enable_heuristic_topk and get_sm_version() >= 100 - if use_gvr: - if self.use_cute_dsl_topk: - # The regular eager attention warmup compiles CuTe DSL GVR. - return - from tensorrt_llm._torch.custom_ops.cpp_custom_ops import warmup_cuda_gvr_topk_decode - - warmup_cuda_gvr_topk_decode(top_k=self.num_sparse_topk) - return - - if not self.use_cute_dsl_topk or (self._indexer_compress_ratio > 1 and next_n > 1): + if ( + not self.use_cute_dsl_topk + or (self._indexer_compress_ratio > 1 and next_n > 1) + or (sparse_params.enable_heuristic_topk and get_sm_version() >= 100) + ): return from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( warmup_cute_dsl_radix_topk_decode, diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 46d4137197ad..0b1eb3b038a7 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1609,38 +1609,3 @@ def _(like: torch.Tensor, out_shape = shape if shape is not None else list(like.shape) dtype = out_dtype if out_dtype is not None else like.dtype return like.new_empty(out_shape, dtype=dtype), output_buffer_kind - - -def warmup_cuda_gvr_topk_decode(top_k: int = 2048) -> None: - """Initialize CUDA GVR before CUDA Graph capture. - - The first dispatcher call queries and caches device SM and L2 attributes. - """ - num_columns = 4096 - max_blocks_per_row = 10 - device = torch.device("cuda") - logits = torch.zeros((1, num_columns), dtype=torch.float32, device=device) - sequence_lengths = torch.tensor([num_columns], - dtype=torch.int32, - device=device) - output_indices = torch.empty((1, top_k), dtype=torch.int32, device=device) - prior_indices = torch.zeros((1, top_k), dtype=torch.int32, device=device) - scratch_values = torch.empty((1, top_k), dtype=torch.float32, device=device) - radix_indices = torch.empty((1, max_blocks_per_row, top_k), - dtype=torch.int32, - device=device) - radix_values = torch.empty((1, max_blocks_per_row, top_k), - dtype=torch.float32, - device=device) - torch.ops.trtllm.indexer_topk_decode( - logits, - sequence_lengths, - output_indices, - 1, - top_k, - pre_idx=prior_indices, - heuristic_scratch=scratch_values, - radix_aux_indices=radix_indices, - radix_aux_logits=radix_values, - ) - torch.cuda.synchronize() diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 3f83ae22da3c..e1c4edbe4b03 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -93,31 +93,6 @@ def forward( radix_values, ) - def seed_from_prefill( - self, - output_indices: torch.Tensor, - request_lengths: torch.Tensor, - *, - request_offset: int = 0, - ) -> None: - """Seed GVR from the last selected row of each prefill request.""" - if self.decode_implementation not in _GVR_IMPLEMENTATIONS: - return - last_rows = (torch.cumsum(request_lengths, dim=0) - 1).to(dtype=torch.long) - num_requests = request_lengths.shape[0] - required_rows = request_offset + num_requests - if ( - self._gvr_prior_indices.device != output_indices.device - or self._gvr_prior_indices.shape[0] < required_rows - ): - prior_indices = output_indices.new_zeros((required_rows, self.top_k), dtype=torch.int32) - if self._gvr_prior_indices.device == output_indices.device: - prior_indices[: self._gvr_prior_indices.shape[0]].copy_(self._gvr_prior_indices) - self._gvr_prior_indices = prior_indices - self._gvr_prior_indices[request_offset : request_offset + num_requests].copy_( - output_indices[last_rows] - ) - def _forward_prefill( self, scores: torch.Tensor, @@ -155,44 +130,43 @@ def _forward_decode( radix_indices: torch.Tensor | None, radix_values: torch.Tensor | None, ) -> torch.Tensor: - implementation = self.decode_implementation - if implementation == TopKImplementation.TORCH: + if self.decode_implementation == TopKImplementation.TORCH: return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n) - if implementation in _GVR_IMPLEMENTATIONS: - self._ensure_gvr_state(scores, sequence_lengths, next_n, radix_indices) - - use_cuda = implementation in ( - TopKImplementation.CUDA_RADIX, - TopKImplementation.CUDA_GVR, - ) or ( - implementation == TopKImplementation.CUTE_DSL_RADIX - and self.compress_ratio > 1 - and next_n > 1 - ) - if use_cuda: - prior_indices = None - heuristic_values = None - if implementation == TopKImplementation.CUDA_GVR: - prior_indices = self._gvr_prior_indices[: sequence_lengths.shape[0]] - heuristic_values = self._cuda_gvr_scratch[: scores.shape[0]] - torch.ops.trtllm.indexer_topk_decode( + if self.decode_implementation in _GVR_IMPLEMENTATIONS: + return self._forward_decode_gvr( scores, sequence_lengths, output_indices, next_n, - self.top_k, - pre_idx=prior_indices, - heuristic_scratch=heuristic_values, - compress_ratio=self.compress_ratio, - radix_aux_indices=radix_indices, - radix_aux_logits=radix_values, + radix_indices, + radix_values, ) - if prior_indices is not None: - prior_indices.copy_(output_indices[next_n - 1 :: next_n]) - return output_indices - if implementation == TopKImplementation.CUTE_DSL_RADIX: + return self._forward_decode_radix( + scores, + sequence_lengths, + scan_lengths, + output_indices, + next_n, + radix_indices, + radix_values, + ) + + def _forward_decode_radix( + self, + scores: torch.Tensor, + sequence_lengths: torch.Tensor, + scan_lengths: torch.Tensor, + output_indices: torch.Tensor, + next_n: int, + radix_indices: torch.Tensor | None, + radix_values: torch.Tensor | None, + ) -> torch.Tensor: + use_cute_dsl = self.decode_implementation == TopKImplementation.CUTE_DSL_RADIX and not ( + self.compress_ratio > 1 and next_n > 1 + ) + if use_cute_dsl: torch.ops.trtllm.cute_dsl_indexer_topk_decode( scores, scan_lengths, @@ -202,54 +176,110 @@ def _forward_decode( ) return output_indices - prior_indices = self._gvr_prior_indices[: sequence_lengths.shape[0]] - torch.ops.trtllm.cute_dsl_gvr_topk_decode( + torch.ops.trtllm.indexer_topk_decode( scores, - prior_indices, sequence_lengths, output_indices, + next_n, self.top_k, - next_n=next_n, + pre_idx=None, + heuristic_scratch=None, compress_ratio=self.compress_ratio, - max_seq_len=scores.shape[1], - order_row=self._prepare_row_order(sequence_lengths, next_n), + radix_aux_indices=radix_indices, + radix_aux_logits=radix_values, ) - prior_indices.copy_(output_indices[next_n - 1 :: next_n]) return output_indices - def _ensure_gvr_state( + def _forward_decode_gvr( self, scores: torch.Tensor, sequence_lengths: torch.Tensor, + output_indices: torch.Tensor, next_n: int, radix_indices: torch.Tensor | None, - ) -> None: + radix_values: torch.Tensor | None, + ) -> torch.Tensor: row_capacity = radix_indices.shape[0] if radix_indices is not None else scores.shape[0] request_capacity = max(sequence_lengths.shape[0], row_capacity // next_n) + self._ensure_gvr_buffers(scores, request_capacity, row_capacity) + + num_requests = sequence_lengths.shape[0] + prior_indices = self._gvr_prior_indices[:num_requests] + if self.decode_implementation == TopKImplementation.CUDA_GVR: + torch.ops.trtllm.indexer_topk_decode( + scores, + sequence_lengths, + output_indices, + next_n, + self.top_k, + pre_idx=prior_indices, + heuristic_scratch=self._cuda_gvr_scratch[: scores.shape[0]], + compress_ratio=self.compress_ratio, + radix_aux_indices=radix_indices, + radix_aux_logits=radix_values, + ) + else: + row_order = None + if num_requests * next_n >= 2 * self._num_sms: + row_order = self._gvr_row_order[:num_requests] + row_order.copy_(torch.argsort(sequence_lengths, descending=True).to(torch.int32)) + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + scores, + prior_indices, + sequence_lengths, + output_indices, + self.top_k, + next_n=next_n, + compress_ratio=self.compress_ratio, + max_seq_len=scores.shape[1], + order_row=row_order, + ) + prior_indices.copy_(output_indices[next_n - 1 :: next_n]) + return output_indices - if ( - self._gvr_prior_indices.device != scores.device - or self._gvr_prior_indices.shape[0] < request_capacity - ): - prior_indices = scores.new_zeros((request_capacity, self.top_k), dtype=torch.int32) - if self._gvr_prior_indices.device == scores.device: - prior_indices[: self._gvr_prior_indices.shape[0]].copy_(self._gvr_prior_indices) + def update_gvr_prior_from_prefill( + self, + output_indices: torch.Tensor, + request_lengths: torch.Tensor, + *, + request_offset: int = 0, + ) -> None: + """Update GVR prior indices from each prefill request's last selected row.""" + if self.decode_implementation not in _GVR_IMPLEMENTATIONS: + return + last_rows = (torch.cumsum(request_lengths, dim=0) - 1).to(dtype=torch.long) + num_requests = request_lengths.shape[0] + required_rows = request_offset + num_requests + if self._gvr_prior_indices.shape[0] < required_rows: + prior_indices = self._gvr_prior_indices.new_zeros((required_rows, self.top_k)) + prior_indices[: self._gvr_prior_indices.shape[0]].copy_(self._gvr_prior_indices) + self._gvr_prior_indices = prior_indices + self._gvr_prior_indices[request_offset : request_offset + num_requests].copy_( + output_indices[last_rows] + ) + + def _ensure_gvr_buffers( + self, + scores: torch.Tensor, + request_capacity: int, + row_capacity: int, + ) -> None: + """Grow persistent GVR buffers before CUDA Graph capture.""" + if self._gvr_prior_indices.shape[0] < request_capacity: + prior_indices = self._gvr_prior_indices.new_zeros((request_capacity, self.top_k)) + prior_indices[: self._gvr_prior_indices.shape[0]].copy_(self._gvr_prior_indices) self._gvr_prior_indices = prior_indices if self.decode_implementation == TopKImplementation.CUDA_GVR: if ( - self._cuda_gvr_scratch.device != scores.device - or self._cuda_gvr_scratch.dtype != scores.dtype + self._cuda_gvr_scratch.dtype != scores.dtype or self._cuda_gvr_scratch.shape[0] < row_capacity ): self._cuda_gvr_scratch = scores.new_empty((row_capacity, self.top_k)) return - if ( - self._gvr_row_order.device != scores.device - or self._gvr_row_order.shape[0] < request_capacity - ): - self._gvr_row_order = scores.new_empty((request_capacity,), dtype=torch.int32) + if self._gvr_row_order.shape[0] < request_capacity: + self._gvr_row_order = self._gvr_row_order.new_empty((request_capacity,)) if self._num_sms == 0: self._num_sms = ( torch.cuda.get_device_properties(scores.device).multi_processor_count @@ -257,18 +287,6 @@ def _ensure_gvr_state( else 1 ) - def _prepare_row_order( - self, - sequence_lengths: torch.Tensor, - next_n: int, - ) -> torch.Tensor | None: - num_requests = sequence_lengths.shape[0] - if num_requests * next_n < 2 * self._num_sms: - return None - row_order = self._gvr_row_order[:num_requests] - row_order.copy_(torch.argsort(sequence_lengths, descending=True).to(torch.int32)) - return row_order - def _forward_prefill_torch( self, scores: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 8c1904716ee5..db8b58d259fd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1508,7 +1508,7 @@ def warmup(self, resource_manager: ResourceManager) -> None: log_mem_snapshot("warmup/after_preallocate_padding_dummies") def _warmup_sparse_top_k(self) -> None: - """Warm up DSA Top-K kernels before any model forward or graph capture.""" + """Pre-compile CuTe DSL radix variants before model warmup.""" from ..attention_backend.sparse.dsa import DSAtrtllmAttentionMetadata if isinstance(self.attn_metadata, DSAtrtllmAttentionMetadata): diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 6db490516635..522f3b126016 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -123,23 +123,23 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): @pytest.mark.parametrize( - "enable_heuristic,use_cute_dsl,sm_version,compress_ratio,next_n,expected", + "enable_heuristic,use_cute_dsl,sm_version,compress_ratio,next_n,should_warmup", [ - (True, False, 100, 1, 1, "cuda_gvr"), - (True, True, 100, 1, 1, None), - (False, True, 100, 1, 1, "cute_dsl_radix"), - (True, True, 90, 1, 1, "cute_dsl_radix"), - (False, True, 100, 4, 2, None), - (False, False, 100, 1, 1, None), + (True, False, 100, 1, 1, False), + (True, True, 100, 1, 1, False), + (False, True, 100, 1, 1, True), + (True, True, 90, 1, 1, True), + (False, True, 100, 4, 2, False), + (False, False, 100, 1, 1, False), ], ) -def test_metadata_warmup_top_k_dispatches_configured_implementation( +def test_metadata_warmup_top_k_only_precompiles_cute_dsl_radix( enable_heuristic, use_cute_dsl, sm_version, compress_ratio, next_n, - expected, + should_warmup, ): metadata = SimpleNamespace( sparse_metadata_params=SimpleNamespace(enable_heuristic_topk=enable_heuristic), @@ -155,20 +155,13 @@ def test_metadata_warmup_top_k_dispatches_configured_implementation( "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.get_sm_version", return_value=sm_version, ), - patch( - "tensorrt_llm._torch.custom_ops.cpp_custom_ops.warmup_cuda_gvr_topk_decode" - ) as cuda_gvr, patch( "tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops.warmup_cute_dsl_radix_topk_decode" ) as cute_dsl_radix, ): DSAtrtllmAttentionMetadata.warmup_top_k(metadata, next_n) - if expected == "cuda_gvr": - cuda_gvr.assert_called_once_with(top_k=512) - else: - cuda_gvr.assert_not_called() - if expected == "cute_dsl_radix": + if should_warmup: cute_dsl_radix.assert_called_once_with( top_k=512, num_cols=32768, diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 75eb263dc265..3c054c5d30b3 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -205,11 +205,11 @@ def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: assert row_order.tolist() == [2, 0, 3, 1] -def test_seed_from_prefill_uses_last_request_rows() -> None: +def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None: top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32) - top_k.seed_from_prefill( + top_k.update_gvr_prior_from_prefill( prefill_indices, torch.tensor([2, 1], dtype=torch.int32), request_offset=1, @@ -230,45 +230,6 @@ def test_gvr_state_buffers_are_registered_during_init() -> None: assert all(buffer.numel() == 0 for buffer in buffers.values()) -def test_cuda_gvr_warmup_calls_cpp_op(monkeypatch) -> None: - from tensorrt_llm._torch.custom_ops import cpp_custom_ops - - decode = Mock() - synchronize = Mock() - torch_empty = torch.empty - torch_tensor = torch.tensor - torch_zeros = torch.zeros - monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) - monkeypatch.setattr(torch.cuda, "synchronize", synchronize) - monkeypatch.setattr( - torch, - "empty", - lambda shape, *, dtype, device: torch_empty(shape, dtype=dtype), - ) - monkeypatch.setattr( - torch, - "tensor", - lambda data, *, dtype, device: torch_tensor(data, dtype=dtype), - ) - monkeypatch.setattr( - torch, - "zeros", - lambda shape, *, dtype, device: torch_zeros(shape, dtype=dtype), - ) - - cpp_custom_ops.warmup_cuda_gvr_topk_decode(top_k=512) - - args = decode.call_args.args - kwargs = decode.call_args.kwargs - assert args[0].shape == (1, 4096) - assert args[2].shape == (1, 512) - assert kwargs["pre_idx"].shape == (1, 512) - assert kwargs["heuristic_scratch"].shape == (1, 512) - assert kwargs["radix_aux_indices"].shape == (1, 10, 512) - assert kwargs["radix_aux_logits"].shape == (1, 10, 512) - synchronize.assert_called_once_with() - - def test_implementations_are_named_by_backend_and_algorithm() -> None: assert {implementation.value for implementation in TopKImplementation} == { "torch", From ed54529fad72466fc64328afd5944ef7528609de Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 12 Aug 2026 22:04:17 -0700 Subject: [PATCH 10/18] [None][refactor] retain main sparse top-k warmup flow Keep CuTe DSL radix warmup after CUDA graph warmup so it only fills eager batch variants not covered by normal engine warmup. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/metadata.py | 31 ++++++++++++------- .../_torch/custom_ops/cpp_custom_ops.py | 3 -- .../_torch/pyexecutor/model_engine.py | 9 ++---- .../attention/sparse/dsa/test_dsa_indexer.py | 5 +-- .../executor/test_pytorch_model_engine.py | 8 ++--- 5 files changed, 29 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 346e74142fd4..e6e9dc3220a0 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -31,6 +31,10 @@ ModelConfig = tensorrt_llm.bindings.ModelConfig +# Indexer MQA-logits are currently always fp32. The dtype is part of the +# CuTe DSL Top-K compile key, so warmup must use the runtime dtype. +_INDEXER_LOGITS_DTYPE = torch.float32 + if TYPE_CHECKING: from tensorrt_llm._torch.speculative.interface import SpecMetadata from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager @@ -299,24 +303,29 @@ def get_indexer_max_seq_len(self) -> int: return self.kv_cache_manager.max_seq_len return max(1, self.kv_cache_manager.max_seq_len // self._indexer_compress_ratio) - def warmup_top_k(self, next_n: int) -> None: + def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: """Pre-compile CuTe DSL radix variants not covered by engine warmup.""" sparse_params = self.sparse_metadata_params - if ( - not self.use_cute_dsl_topk - or (self._indexer_compress_ratio > 1 and next_n > 1) - or (sparse_params.enable_heuristic_topk and get_sm_version() >= 100) + if not self.use_cute_dsl_topk or ( + sparse_params.enable_heuristic_topk and get_sm_version() >= 100 ): return - from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( - warmup_cute_dsl_radix_topk_decode, - ) + if self.kv_cache_manager is None or not self.num_sparse_topk: + return + if self._indexer_compress_ratio > 1 and next_n > 1: + return + try: + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + warmup_cute_dsl_radix_topk_decode, + ) + except ImportError: + return warmup_cute_dsl_radix_topk_decode( - top_k=self.num_sparse_topk, - num_cols=self.get_indexer_max_seq_len(), + top_k=int(self.num_sparse_topk), + num_cols=int(self.get_indexer_max_seq_len()), next_n=next_n, - dtype=torch.float32, + dtype=_INDEXER_LOGITS_DTYPE, num_sms=self.num_sms, ) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 0b1eb3b038a7..3feac32cee17 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - from typing import List, Optional, Tuple import torch diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index db8b58d259fd..e7d8df92ab9f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1492,6 +1492,8 @@ def warmup(self, resource_manager: ResourceManager) -> None: # warmup. No-op on non-DSA models. self._warmup_dg_paged_mqa_logits_metadata() log_mem_snapshot("warmup/after_dg_paged_mqa_logits_metadata") + self._warmup_cute_dsl_radix_topk() + log_mem_snapshot("warmup/after_cute_dsl_radix_topk") if can_run_general_warmup: # Pre-populate the memory pool with max-shape allocations to reduce # fragmentation at runtime. @@ -1507,13 +1509,6 @@ def warmup(self, resource_manager: ResourceManager) -> None: self.cuda_graph_runner.preallocate_padding_dummies(resource_manager) log_mem_snapshot("warmup/after_preallocate_padding_dummies") - def _warmup_sparse_top_k(self) -> None: - """Pre-compile CuTe DSL radix variants before model warmup.""" - from ..attention_backend.sparse.dsa import DSAtrtllmAttentionMetadata - - if isinstance(self.attn_metadata, DSAtrtllmAttentionMetadata): - self.attn_metadata.warmup_top_k(1 + self.original_max_draft_len) - def _warmup_dg_paged_mqa_logits_metadata(self) -> None: """Pre-compile DeepGEMM's `get_paged_mqa_logits_metadata` helper for every 32-aligned batch bucket the runtime can produce. diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 522f3b126016..6e824bdc1fef 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -133,7 +133,7 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): (False, False, 100, 1, 1, False), ], ) -def test_metadata_warmup_top_k_only_precompiles_cute_dsl_radix( +def test_metadata_warmup_cute_dsl_radix_topk_dispatch( enable_heuristic, use_cute_dsl, sm_version, @@ -145,6 +145,7 @@ def test_metadata_warmup_top_k_only_precompiles_cute_dsl_radix( sparse_metadata_params=SimpleNamespace(enable_heuristic_topk=enable_heuristic), use_cute_dsl_topk=use_cute_dsl, num_sparse_topk=512, + kv_cache_manager=SimpleNamespace(), _indexer_compress_ratio=compress_ratio, get_indexer_max_seq_len=Mock(return_value=32768), num_sms=148, @@ -159,7 +160,7 @@ def test_metadata_warmup_top_k_only_precompiles_cute_dsl_radix( "tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops.warmup_cute_dsl_radix_topk_decode" ) as cute_dsl_radix, ): - DSAtrtllmAttentionMetadata.warmup_top_k(metadata, next_n) + DSAtrtllmAttentionMetadata.warmup_cute_dsl_radix_topk(metadata, next_n) if should_warmup: cute_dsl_radix.assert_called_once_with( diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 7aed76d3ee70..2629dc3c2c8b 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -91,17 +91,17 @@ def forward(self, *args, **kwargs) -> torch.Tensor: return {"logits": torch.randn((batch_size, 10), device='cuda')} -def test_sparse_top_k_warmup_uses_shared_metadata(): +def test_cute_dsl_radix_top_k_warmup_uses_shared_metadata(): from tensorrt_llm._torch.attention_backend.sparse.dsa import \ DSAtrtllmAttentionMetadata metadata = object.__new__(DSAtrtllmAttentionMetadata) - metadata.warmup_top_k = Mock() + metadata.warmup_cute_dsl_radix_topk = Mock() engine = SimpleNamespace(attn_metadata=metadata, original_max_draft_len=3) - PyTorchModelEngine._warmup_sparse_top_k(engine) + PyTorchModelEngine._warmup_cute_dsl_radix_topk(engine) - metadata.warmup_top_k.assert_called_once_with(4) + metadata.warmup_cute_dsl_radix_topk.assert_called_once_with(4) class DummyMultimodalIndexModel(torch.nn.Module): From a6dfd4089c0dcc535b4f6b12aa67c91bd983b66b Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 12 Aug 2026 23:30:22 -0700 Subject: [PATCH 11/18] [None][refactor] preserve model engine top-k warmup Keep the existing model-engine CuTe DSL radix warmup wrapper unchanged and limit the TopK refactor to sparse metadata and module ownership. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 1 + .../_torch/executor/test_pytorch_model_engine.py | 13 ------------- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e7d8df92ab9f..3da7a73a7e9e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3232,6 +3232,7 @@ def _set_up_attn_metadata( num_heads_per_kv=num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, ) + return self.attn_metadata @property diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 2629dc3c2c8b..a323f89f3158 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -91,19 +91,6 @@ def forward(self, *args, **kwargs) -> torch.Tensor: return {"logits": torch.randn((batch_size, 10), device='cuda')} -def test_cute_dsl_radix_top_k_warmup_uses_shared_metadata(): - from tensorrt_llm._torch.attention_backend.sparse.dsa import \ - DSAtrtllmAttentionMetadata - - metadata = object.__new__(DSAtrtllmAttentionMetadata) - metadata.warmup_cute_dsl_radix_topk = Mock() - engine = SimpleNamespace(attn_metadata=metadata, original_max_draft_len=3) - - PyTorchModelEngine._warmup_cute_dsl_radix_topk(engine) - - metadata.warmup_cute_dsl_radix_topk.assert_called_once_with(4) - - class DummyMultimodalIndexModel(torch.nn.Module): class Config: From 446fb8d62d5a3549749a7154612b048352ec3a03 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 12 Aug 2026 23:53:13 -0700 Subject: [PATCH 12/18] [None][test] streamline sparse top-k coverage Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention/sparse/dsa/test_dsa_indexer.py | 24 +----- .../executor/test_pytorch_model_engine.py | 1 - tests/unittest/_torch/modules/test_top_k.py | 78 +++++++------------ 3 files changed, 28 insertions(+), 75 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 6e824bdc1fef..94dae740eac5 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -54,9 +54,9 @@ transform_local_topk_and_prepare_pool_view, ) from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation from tensorrt_llm._torch.pyexecutor._util import get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role -from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation from tensorrt_llm._torch.speculative.interface import ( prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, @@ -125,7 +125,6 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): @pytest.mark.parametrize( "enable_heuristic,use_cute_dsl,sm_version,compress_ratio,next_n,should_warmup", [ - (True, False, 100, 1, 1, False), (True, True, 100, 1, 1, False), (False, True, 100, 1, 1, True), (True, True, 90, 1, 1, True), @@ -304,27 +303,6 @@ def test_indexer_configures_one_top_k_module(): assert isinstance(indexer.top_k, TopK) assert indexer.top_k.prefill_implementation == TopKImplementation.CUDA_RADIX assert indexer.top_k.decode_implementation == TopKImplementation.CUDA_RADIX - assert not hasattr(indexer, "prefill_top_k") - assert not hasattr(indexer, "decode_top_k") - - -def test_indexer_prepare_updates_metadata_without_layer_state(): - metadata = SimpleNamespace( - num_contexts=0, - num_generations=0, - num_ctx_tokens=0, - seq_lens=torch.empty(0, dtype=torch.int32), - compress_ratios=[1], - ) - indexer_params = SimpleNamespace(new_kv_tokens=torch.empty(0, dtype=torch.int32)) - - with ( - patch.object(Indexer, "build_indexer_params", return_value=indexer_params), - patch.object(Indexer, "prepare_for_update_k_cache") as prepare_metadata, - ): - Indexer.prepare(metadata) - - prepare_metadata.assert_called_once_with(metadata, indexer_params) def _ceil_to_ue8m0(x: torch.Tensor): diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index a323f89f3158..7d44477e520e 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -200,7 +200,6 @@ def _make_forward_only_engine( engine = object.__new__(PyTorchModelEngine) engine.model = SimpleNamespace( extra_attrs={}, - modules=lambda: [], model_config=SimpleNamespace(pretrained_config=SimpleNamespace( rope_scaling=None))) engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 3c054c5d30b3..1b01d7dead9f 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -61,34 +61,6 @@ def test_decode_torch_uses_scan_lengths() -> None: assert output.tolist() == [[1, 0, -1, -1], [2, 1, 0, -1]] -def test_one_module_dispatches_prefill_and_decode() -> None: - top_k = TopK( - 1, - prefill_implementation=TopKImplementation.TORCH, - decode_implementation=TopKImplementation.TORCH, - ) - scores = torch.tensor([[1.0, 3.0, 2.0]]) - output = torch.empty(1, 1, dtype=torch.int32) - - top_k( - scores, - output, - is_prefill=True, - row_starts=torch.tensor([1], dtype=torch.int32), - row_ends=torch.tensor([3], dtype=torch.int32), - ) - assert output.item() == 0 - - top_k( - scores, - output, - is_prefill=False, - sequence_lengths=torch.tensor([3], dtype=torch.int32), - scan_lengths=torch.tensor([3], dtype=torch.int32), - ) - assert output.item() == 1 - - def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: cute_dsl = Mock() trtllm = Mock() @@ -218,33 +190,37 @@ def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None: assert top_k._gvr_prior_indices.tolist() == [[0, 0], [2, 3], [4, 5]] -def test_gvr_state_buffers_are_registered_during_init() -> None: - top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) - - buffers = dict(top_k.named_buffers()) - assert set(buffers) == { - "_gvr_prior_indices", - "_cuda_gvr_scratch", - "_gvr_row_order", - } - assert all(buffer.numel() == 0 for buffer in buffers.values()) - - -def test_implementations_are_named_by_backend_and_algorithm() -> None: - assert {implementation.value for implementation in TopKImplementation} == { - "torch", - "cuda_radix", - "cute_dsl_radix", - "cuda_gvr", - "cute_dsl_gvr", - } - - -def test_none_implementations_use_cuda_radix_defaults() -> None: +def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: + decode = Mock() + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) top_k = TopK(1) + scores = torch.randn(1, 8) + lengths = torch.tensor([8], dtype=torch.int32) + output = torch.empty((1, 1), dtype=torch.int32) + + result = top_k( + scores, + output, + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths, + ) assert top_k.prefill_implementation == TopKImplementation.CUDA_RADIX assert top_k.decode_implementation == TopKImplementation.CUDA_RADIX + assert result is output + decode.assert_called_once_with( + scores, + lengths, + output, + 1, + 1, + pre_idx=None, + heuristic_scratch=None, + compress_ratio=1, + radix_aux_indices=None, + radix_aux_logits=None, + ) def test_cuda_gvr_owns_scratch_and_updates_prior(monkeypatch) -> None: From ed01ab2916dc0ccef961e904fd05dafd026538b0 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Thu, 13 Aug 2026 04:34:49 -0700 Subject: [PATCH 13/18] fix: stabilize GVR Top-K state Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 5 +- .../attention_backend/sparse/dsa/metadata.py | 7 +- tensorrt_llm/_torch/modules/top_k.py | 81 ++++++++++++++++--- .../attention/sparse/dsa/test_dsa_indexer.py | 50 ++++++++++-- tests/unittest/_torch/modules/test_top_k.py | 76 +++++++++++++++++ 5 files changed, 198 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 37f98f66b8c1..597c4660f369 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1517,8 +1517,8 @@ def sparse_attn_indexer( :num_ctx_tokens, : ] - # Chunked prefill is final only after its TP all-gathers; update GVR prior once here. - if has_prefill and not metadata.skip_indexer_for_ctx_reqs: + # Update GVR state after chunk all-gathers or the dense-index copy. + if has_prefill: self.top_k.update_gvr_prior_from_prefill( topk_indices_buffer[:num_ctx_tokens], metadata.seq_lens[:num_contexts], @@ -1716,6 +1716,7 @@ def sparse_attn_indexer( next_n=next_n, radix_indices=metadata.radix_aux_indices, radix_values=metadata.radix_aux_logits, + request_capacity=metadata.max_num_sequences, ) elif has_decode and metadata.skip_indexer_for_gen_reqs: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index e6e9dc3220a0..34ba0e2f4fd4 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -310,7 +310,10 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: sparse_params.enable_heuristic_topk and get_sm_version() >= 100 ): return - if self.kv_cache_manager is None or not self.num_sparse_topk: + if self.kv_cache_manager is None: + return + top_k = getattr(sparse_params, "index_topk", None) + if not top_k: return if self._indexer_compress_ratio > 1 and next_n > 1: return @@ -322,7 +325,7 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: return warmup_cute_dsl_radix_topk_decode( - top_k=int(self.num_sparse_topk), + top_k=int(top_k), num_cols=int(self.get_indexer_max_seq_len()), next_n=next_n, dtype=_INDEXER_LOGITS_DTYPE, diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index e1c4edbe4b03..ad3ead2bea45 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -76,6 +76,7 @@ def forward( next_n: int = 1, radix_indices: torch.Tensor | None = None, radix_values: torch.Tensor | None = None, + request_capacity: int | None = None, ) -> torch.Tensor: """Write prefill or decode Top-K indices into ``output_indices``.""" if is_prefill: @@ -91,6 +92,7 @@ def forward( next_n, radix_indices, radix_values, + request_capacity, ) def _forward_prefill( @@ -129,6 +131,7 @@ def _forward_decode( next_n: int, radix_indices: torch.Tensor | None, radix_values: torch.Tensor | None, + request_capacity: int | None, ) -> torch.Tensor: if self.decode_implementation == TopKImplementation.TORCH: return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n) @@ -141,6 +144,7 @@ def _forward_decode( next_n, radix_indices, radix_values, + request_capacity, ) return self._forward_decode_radix( @@ -198,9 +202,12 @@ def _forward_decode_gvr( next_n: int, radix_indices: torch.Tensor | None, radix_values: torch.Tensor | None, + request_capacity: int | None, ) -> torch.Tensor: row_capacity = radix_indices.shape[0] if radix_indices is not None else scores.shape[0] - request_capacity = max(sequence_lengths.shape[0], row_capacity // next_n) + request_capacity = request_capacity or max( + sequence_lengths.shape[0], row_capacity // next_n + ) self._ensure_gvr_buffers(scores, request_capacity, row_capacity) num_requests = sequence_lengths.shape[0] @@ -250,8 +257,28 @@ def update_gvr_prior_from_prefill( last_rows = (torch.cumsum(request_lengths, dim=0) - 1).to(dtype=torch.long) num_requests = request_lengths.shape[0] required_rows = request_offset + num_requests - if self._gvr_prior_indices.shape[0] < required_rows: - prior_indices = self._gvr_prior_indices.new_zeros((required_rows, self.top_k)) + needs_resize = ( + self._gvr_prior_indices.shape[0] < required_rows + or self._gvr_prior_indices.device != output_indices.device + ) + if needs_resize: + decode_initialized = ( + self._cuda_gvr_scratch.numel() > 0 + if self.decode_implementation == TopKImplementation.CUDA_GVR + else self._gvr_row_order.numel() > 0 + ) + if decode_initialized or ( + output_indices.is_cuda and torch.cuda.is_current_stream_capturing() + ): + raise RuntimeError( + "GVR prior indices cannot be resized after decode initialization" + ) + prior_capacity = max(required_rows, self._gvr_prior_indices.shape[0]) + prior_indices = torch.zeros( + (prior_capacity, self.top_k), + dtype=torch.int32, + device=output_indices.device, + ) prior_indices[: self._gvr_prior_indices.shape[0]].copy_(self._gvr_prior_indices) self._gvr_prior_indices = prior_indices self._gvr_prior_indices[request_offset : request_offset + num_requests].copy_( @@ -264,22 +291,52 @@ def _ensure_gvr_buffers( request_capacity: int, row_capacity: int, ) -> None: - """Grow persistent GVR buffers before CUDA Graph capture.""" - if self._gvr_prior_indices.shape[0] < request_capacity: - prior_indices = self._gvr_prior_indices.new_zeros((request_capacity, self.top_k)) + """Initialize fixed-address GVR buffers before CUDA Graph capture.""" + needs_prior = ( + self._gvr_prior_indices.shape[0] < request_capacity + or self._gvr_prior_indices.device != scores.device + ) + needs_scratch = self.decode_implementation == TopKImplementation.CUDA_GVR and ( + self._cuda_gvr_scratch.dtype != scores.dtype + or self._cuda_gvr_scratch.device != scores.device + or self._cuda_gvr_scratch.shape[0] < row_capacity + ) + needs_row_order = self.decode_implementation == TopKImplementation.CUTE_DSL_GVR and ( + self._gvr_row_order.device != scores.device + or self._gvr_row_order.shape[0] < request_capacity + ) + needs_resize = needs_prior or needs_scratch or needs_row_order + decode_initialized = ( + self._cuda_gvr_scratch.numel() > 0 + if self.decode_implementation == TopKImplementation.CUDA_GVR + else self._gvr_row_order.numel() > 0 + ) + if needs_resize and ( + decode_initialized or (scores.is_cuda and torch.cuda.is_current_stream_capturing()) + ): + raise RuntimeError("GVR buffers cannot be resized after decode initialization") + + if needs_prior: + prior_capacity = max(request_capacity, self._gvr_prior_indices.shape[0]) + prior_indices = torch.zeros( + (prior_capacity, self.top_k), + dtype=torch.int32, + device=scores.device, + ) prior_indices[: self._gvr_prior_indices.shape[0]].copy_(self._gvr_prior_indices) self._gvr_prior_indices = prior_indices if self.decode_implementation == TopKImplementation.CUDA_GVR: - if ( - self._cuda_gvr_scratch.dtype != scores.dtype - or self._cuda_gvr_scratch.shape[0] < row_capacity - ): + if needs_scratch: self._cuda_gvr_scratch = scores.new_empty((row_capacity, self.top_k)) return - if self._gvr_row_order.shape[0] < request_capacity: - self._gvr_row_order = self._gvr_row_order.new_empty((request_capacity,)) + if needs_row_order: + self._gvr_row_order = torch.empty( + (request_capacity,), + dtype=torch.int32, + device=scores.device, + ) if self._num_sms == 0: self._num_sms = ( torch.cuda.get_device_properties(scores.device).multi_processor_count diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 94dae740eac5..4d4664d1d249 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -141,7 +141,10 @@ def test_metadata_warmup_cute_dsl_radix_topk_dispatch( should_warmup, ): metadata = SimpleNamespace( - sparse_metadata_params=SimpleNamespace(enable_heuristic_topk=enable_heuristic), + sparse_metadata_params=SimpleNamespace( + enable_heuristic_topk=enable_heuristic, + index_topk=384, + ), use_cute_dsl_topk=use_cute_dsl, num_sparse_topk=512, kv_cache_manager=SimpleNamespace(), @@ -163,7 +166,7 @@ def test_metadata_warmup_cute_dsl_radix_topk_dispatch( if should_warmup: cute_dsl_radix.assert_called_once_with( - top_k=512, + top_k=384, num_cols=32768, next_n=next_n, dtype=torch.float32, @@ -291,18 +294,43 @@ def test_indexer_post_load_weights_caches_fused_weight(): @skip_pre_hopper -def test_indexer_configures_one_top_k_module(): +@pytest.mark.parametrize( + "use_cute_dsl,enable_heuristic,expected_decode", + [ + (False, False, TopKImplementation.CUDA_RADIX), + (True, False, TopKImplementation.CUTE_DSL_RADIX), + (False, True, TopKImplementation.CUDA_GVR), + (True, True, TopKImplementation.CUTE_DSL_GVR), + ], +) +def test_indexer_configures_one_top_k_module( + use_cute_dsl, + enable_heuristic, + expected_decode, +): sparse_config = DeepSeekSparseAttentionConfig( index_head_dim=128, index_n_heads=32, index_topk=128, + use_cute_dsl_topk=use_cute_dsl, + enable_heuristic_topk=enable_heuristic, ) - indexer = create_indexer(sparse_config) + with ( + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.IS_CUTLASS_DSL_AVAILABLE", + True, + ), + patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.get_sm_version", + return_value=100, + ), + ): + indexer = create_indexer(sparse_config) assert isinstance(indexer.top_k, TopK) assert indexer.top_k.prefill_implementation == TopKImplementation.CUDA_RADIX - assert indexer.top_k.decode_implementation == TopKImplementation.CUDA_RADIX + assert indexer.top_k.decode_implementation == expected_decode def _ceil_to_ue8m0(x: torch.Tensor): @@ -794,6 +822,7 @@ def __init__(self): self.num_contexts = num_contexts self.num_generations = num_generations self._num_seqs = num_contexts + num_generations + self.max_num_sequences = batch_size self.max_draft_tokens = max_draft_tokens self.num_sparse_topk = index_topk self.enable_indexer_skip = enable_indexer_skip @@ -3369,6 +3398,11 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) metadata_skip.indexer_prefill_chunks = None + indexer.top_k = TopK( + index_topk, + prefill_implementation=TopKImplementation.CUDA_RADIX, + decode_implementation=TopKImplementation.CUTE_DSL_GVR, + ) try: topk_indices_skip = indexer.sparse_attn_indexer( @@ -3377,6 +3411,12 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, except Exception as e: raise RuntimeError(f"Indexer skip not available: {e}") + last_rows = torch.cumsum(metadata_skip.seq_lens[:batch_size], dim=0) - 1 + torch.testing.assert_close( + indexer.top_k._gvr_prior_indices[:batch_size], + topk_indices_skip[last_rows], + ) + # Validation ## Custom vs fallback num_exact_matches, total_similarity, _ = validate_topk_indices( diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 1b01d7dead9f..6c5d0355a17b 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -4,6 +4,7 @@ from unittest.mock import Mock +import pytest import torch from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation @@ -248,3 +249,78 @@ def test_cuda_gvr_owns_scratch_and_updates_prior(monkeypatch) -> None: assert runtime_call.kwargs["pre_idx"].data_ptr() == top_k._gvr_prior_indices.data_ptr() assert runtime_call.kwargs["heuristic_scratch"].data_ptr() == top_k._cuda_gvr_scratch.data_ptr() assert top_k._gvr_prior_indices.tolist() == [[3, 1], [0, 0]] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cuda_gvr_buffers_use_scores_device_and_keep_addresses(monkeypatch) -> None: + decode = Mock(side_effect=lambda *args, **kwargs: args[2].zero_()) + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) + top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) + scores = torch.randn(2, 8, device="cuda") + lengths = torch.tensor([8, 8], dtype=torch.int32, device="cuda") + output = torch.empty(2, 2, dtype=torch.int32, device="cuda") + radix_indices = torch.empty(2, 10, 2, dtype=torch.int32, device="cuda") + radix_values = torch.empty(2, 10, 2, device="cuda") + + pointers = None + for _ in range(2): + top_k( + scores, + output, + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths, + radix_indices=radix_indices, + radix_values=radix_values, + request_capacity=2, + ) + current_pointers = ( + top_k._gvr_prior_indices.data_ptr(), + top_k._cuda_gvr_scratch.data_ptr(), + ) + if pointers is None: + pointers = current_pointers + else: + assert current_pointers == pointers + assert top_k._gvr_prior_indices.device == scores.device + assert top_k._cuda_gvr_scratch.device == scores.device + + +def test_unsupported_prefill_implementation_raises() -> None: + top_k = TopK(1, prefill_implementation=TopKImplementation.CUTE_DSL_RADIX) + + with pytest.raises(NotImplementedError, match="does not support prefill Top-K"): + top_k( + torch.ones(1, 1), + torch.empty(1, 1, dtype=torch.int32), + is_prefill=True, + row_starts=torch.zeros(1, dtype=torch.int32), + row_ends=torch.ones(1, dtype=torch.int32), + ) + + +def test_gvr_buffers_cannot_grow_after_decode_initialization(monkeypatch) -> None: + decode = Mock(side_effect=lambda *args, **kwargs: args[2].zero_()) + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) + top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) + scores = torch.randn(1, 8) + lengths = torch.tensor([8], dtype=torch.int32) + + top_k( + scores, + torch.empty(1, 2, dtype=torch.int32), + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths, + request_capacity=1, + ) + + with pytest.raises(RuntimeError, match="cannot be resized"): + top_k( + scores, + torch.empty(1, 2, dtype=torch.int32), + is_prefill=False, + sequence_lengths=lengths, + scan_lengths=lengths, + request_capacity=2, + ) From 748fd25ec4adbc47fd73adb14314467b4556b264 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Thu, 13 Aug 2026 07:26:22 -0700 Subject: [PATCH 14/18] [None][fix] stabilize sparse top-k workspaces Keep GVR prior indices in per-layer DSA metadata while allocating GVR and radix scratch from the reusable memory-buffer arena. Preserve the prior state across bypassed decode paths and use the indexer Top-K value for warmup. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- cpp/tensorrt_llm/thop/IndexerTopKOp.cpp | 5 +- .../attention_backend/sparse/dsa/indexer.py | 24 ++- .../attention_backend/sparse/dsa/metadata.py | 61 ++---- tensorrt_llm/_torch/modules/top_k.py | 192 ++++++------------ .../attention/sparse/dsa/test_dsa_indexer.py | 21 +- tests/unittest/_torch/modules/test_top_k.py | 161 +++++++-------- 6 files changed, 192 insertions(+), 272 deletions(-) diff --git a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp index fc29041bbf93..7e9f3bd7070d 100644 --- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp @@ -127,9 +127,8 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t if (radix_aux_indices.has_value() && radix_aux_logits.has_value()) { // Caller-owned scratch with stable address (CUDA Graph safe; - // matches the heuristic_scratch convention noted above). All - // in-tree callers under CUDA Graph capture go through dsa.py's - // DSAtrtllmAttentionMetadata which always pre-allocates these. + // matches the heuristic_scratch convention noted above). The + // Python TopK module supplies these from its reusable buffer arena. auto const& ai = radix_aux_indices.value(); auto const& al = radix_aux_logits.value(); TORCH_CHECK(ai.is_cuda() && al.is_cuda(), "radix_aux_{indices,logits} must be CUDA tensors"); diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 597c4660f369..eb6110609a9d 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1354,6 +1354,10 @@ def sparse_attn_indexer( num_tokens = metadata.num_tokens num_gen_tokens = num_tokens - num_ctx_tokens + gvr_prior_indices = None + if self._enable_heuristic_topk: + local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] + gvr_prior_indices = metadata.gvr_prior_indices[local_layer] if is_generation is None: has_prefill = num_contexts > 0 has_decode = num_generations > 0 @@ -1522,6 +1526,7 @@ def sparse_attn_indexer( self.top_k.update_gvr_prior_from_prefill( topk_indices_buffer[:num_ctx_tokens], metadata.seq_lens[:num_contexts], + gvr_prior_indices, request_offset=num_generations, ) @@ -1714,9 +1719,9 @@ def sparse_attn_indexer( sequence_lengths=gen_kv_lens_cuda, scan_lengths=scan_lengths, next_n=next_n, - radix_indices=metadata.radix_aux_indices, - radix_values=metadata.radix_aux_logits, - request_capacity=metadata.max_num_sequences, + gvr_prior_indices=( + gvr_prior_indices[:num_generations] if gvr_prior_indices is not None else None + ), ) elif has_decode and metadata.skip_indexer_for_gen_reqs: @@ -1725,6 +1730,19 @@ def sparse_attn_indexer( metadata.topk_indices_buffer[num_ctx_tokens:num_tokens, :] ) + if ( + gvr_prior_indices is not None + and has_decode + and (reuse_topk or metadata.skip_indexer_for_gen_reqs) + ): + next_n = num_gen_tokens // num_generations + # Keep bypassed decode results as the next GVR prior; MTP reuse holds the accepted row. + gvr_prior_indices[:num_generations].copy_( + topk_indices_buffer[token_offset : token_offset + num_gen_tokens][ + next_n - 1 :: next_n + ] + ) + if self.mtp_index_share and metadata.in_mtp_draft_loop and not reuse_topk: rows = None if has_decode: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 34ba0e2f4fd4..64f6396ba273 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -152,6 +152,9 @@ def __post_init__(self): self.use_cute_dsl_topk = ( sparse_metadata_params.use_cute_dsl_topk and IS_CUTLASS_DSL_AVAILABLE ) + self.enable_gvr_topk = ( + sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100 + ) capture_graph = self.is_cuda_graph # Plain DSA has no compression and uses the default [1]. DeepSeek-V4's # metadata params carry the model-specific compression ratios. @@ -312,7 +315,7 @@ def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: return if self.kv_cache_manager is None: return - top_k = getattr(sparse_params, "index_topk", None) + top_k = self.sparse_mla_topk if not top_k: return if self._indexer_compress_ratio > 1 and next_n > 1: @@ -515,37 +518,6 @@ def create_buffers_for_mla_rope_append(self, capture_graph=False): pin_memory=prefer_pinned(), ) - def _create_radix_aux_buffers(self, capture_graph=False): - # Persistent scratch for Radix-split-work indexer path (blocks_per_row > 1). - # Mirrors the fix the Heuristic path applied: per-call th::empty inside - # indexer_topk_decode produces stale pointers under CUDA Graph replay when - # the caching allocator is perturbed by chunked prefill at high CONC. - # Sized to the worst case kMaxBlocksPerRowDecode=10 from - # cpp/tensorrt_llm/kernels/indexerTopK.cu, times the max number of - # generation rows (num_seqs * (1 + max_draft_tokens)); the cpp op aborts - # if this is smaller than num_rows*blocks_per_row*index_topk. Allocated - # unconditionally: even with enable_heuristic_topk=True the dispatcher can - # fall back to Radix when canUseHeuristic returns False (small numColumns, - # etc.). MUST be re-created whenever max_draft_tokens changes (see - # update_spec_dec_param) or it is left too small once MTP raises the - # generation-row count. - _radix_max_blocks_per_row = 10 - _radix_max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) - self.radix_aux_indices = self.get_empty( - self.cuda_graph_buffers, - (_radix_max_gen_tokens, _radix_max_blocks_per_row, self.num_sparse_topk), - cache_name="radix_aux_indices", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.radix_aux_logits = self.get_empty( - self.cuda_graph_buffers, - (_radix_max_gen_tokens, _radix_max_blocks_per_row, self.num_sparse_topk), - cache_name="radix_aux_logits", - dtype=torch.float32, - capture_graph=capture_graph, - ) - def create_buffers_for_indexer(self, capture_graph=False): sparse_metadata_params = self.sparse_metadata_params if not isinstance(sparse_metadata_params, DSAMetadataParams): @@ -742,11 +714,19 @@ def create_buffers_for_indexer(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - # Persistent scratch for the Radix-split-work indexer path. Re-created - # in update_spec_dec_param when max_draft_tokens changes so it stays - # large enough for the MTP generation-row count. - self._create_radix_aux_buffers(capture_graph=capture_graph) - + if self.enable_gvr_topk: + self.gvr_prior_indices = self.get_empty( + self.cuda_graph_buffers, + ( + self.kv_cache_manager.num_local_layers, + self.max_num_sequences, + self.num_sparse_topk, + ), + cache_name="gvr_prior_indices", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.gvr_prior_indices.zero_() # Create expanded buffers for MTP support self.create_expanded_buffers(capture_graph=capture_graph) @@ -854,13 +834,6 @@ def update_spec_dec_param( init_shape = self.kv_lens_expanded_host.shape[0] if self.max_num_sequences * (1 + self.max_draft_tokens) != init_shape: self.create_expanded_buffers(capture_graph=capture_graph) - # The Radix-split-work scratch (radix_aux_*) is sized the same way - # (num_seqs * (1 + max_draft_tokens) rows) and is allocated - # unconditionally, so it must be resized here too -- otherwise the - # cpp indexer_topk_decode op aborts once MTP raises max_draft_tokens - # ("radix_aux_* must hold at least num_rows*blocks_per_row*index_topk - # elements"). - self._create_radix_aux_buffers(capture_graph=capture_graph) def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor: """Refresh INDEX_KEY offsets and return their physical pool slots.""" diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index ad3ead2bea45..2fab70b92dca 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -9,6 +9,8 @@ import torch import torch.nn as nn +from ..memory_buffer_utils import get_memory_buffers + class TopKImplementation(str, Enum): """Top-K implementations grouped by backend and algorithm.""" @@ -24,11 +26,14 @@ class TopKImplementation(str, Enum): TopKImplementation.CUDA_GVR, TopKImplementation.CUTE_DSL_GVR, } +_MAX_RADIX_BLOCKS_PER_ROW = 10 class TopK(nn.Module): """Select Top-K indices for sparse prefill and decode paths.""" + _memory_buffers = get_memory_buffers() + def __init__( self, top_k: int, @@ -46,21 +51,6 @@ def __init__( decode_implementation or TopKImplementation.CUDA_RADIX ) self.compress_ratio = compress_ratio - self.register_buffer( - "_gvr_prior_indices", - torch.empty((0, top_k), dtype=torch.int32), - persistent=False, - ) - self.register_buffer( - "_cuda_gvr_scratch", - torch.empty((0, top_k)), - persistent=False, - ) - self.register_buffer( - "_gvr_row_order", - torch.empty((0,), dtype=torch.int32), - persistent=False, - ) self._num_sms = 0 def forward( @@ -74,9 +64,7 @@ def forward( sequence_lengths: torch.Tensor | None = None, scan_lengths: torch.Tensor | None = None, next_n: int = 1, - radix_indices: torch.Tensor | None = None, - radix_values: torch.Tensor | None = None, - request_capacity: int | None = None, + gvr_prior_indices: torch.Tensor | None = None, ) -> torch.Tensor: """Write prefill or decode Top-K indices into ``output_indices``.""" if is_prefill: @@ -90,9 +78,7 @@ def forward( scan_lengths, output_indices, next_n, - radix_indices, - radix_values, - request_capacity, + gvr_prior_indices, ) def _forward_prefill( @@ -129,9 +115,7 @@ def _forward_decode( scan_lengths: torch.Tensor, output_indices: torch.Tensor, next_n: int, - radix_indices: torch.Tensor | None, - radix_values: torch.Tensor | None, - request_capacity: int | None, + gvr_prior_indices: torch.Tensor | None, ) -> torch.Tensor: if self.decode_implementation == TopKImplementation.TORCH: return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n) @@ -142,9 +126,7 @@ def _forward_decode( sequence_lengths, output_indices, next_n, - radix_indices, - radix_values, - request_capacity, + gvr_prior_indices, ) return self._forward_decode_radix( @@ -153,8 +135,6 @@ def _forward_decode( scan_lengths, output_indices, next_n, - radix_indices, - radix_values, ) def _forward_decode_radix( @@ -164,8 +144,6 @@ def _forward_decode_radix( scan_lengths: torch.Tensor, output_indices: torch.Tensor, next_n: int, - radix_indices: torch.Tensor | None, - radix_values: torch.Tensor | None, ) -> torch.Tensor: use_cute_dsl = self.decode_implementation == TopKImplementation.CUTE_DSL_RADIX and not ( self.compress_ratio > 1 and next_n > 1 @@ -180,6 +158,7 @@ def _forward_decode_radix( ) return output_indices + radix_indices, radix_values = self._get_radix_workspace(scores) torch.ops.trtllm.indexer_topk_decode( scores, sequence_lengths, @@ -194,45 +173,78 @@ def _forward_decode_radix( ) return output_indices + def _get_radix_workspace( + self, scores: torch.Tensor + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if scores.dtype != torch.float32: + return None, None + + shape = (scores.shape[0], _MAX_RADIX_BLOCKS_PER_ROW, self.top_k) + capture_graph = scores.is_cuda and torch.cuda.is_current_stream_capturing() + radix_indices = self._memory_buffers.get_buffer( + shape, + dtype=torch.int32, + buffer_name="top_k_radix_indices_workspace", + reserve_buffer=capture_graph, + ) + radix_values = self._memory_buffers.get_buffer( + shape, + dtype=torch.float32, + buffer_name="top_k_radix_values_workspace", + reserve_buffer=capture_graph, + ) + return radix_indices, radix_values + def _forward_decode_gvr( self, scores: torch.Tensor, sequence_lengths: torch.Tensor, output_indices: torch.Tensor, next_n: int, - radix_indices: torch.Tensor | None, - radix_values: torch.Tensor | None, - request_capacity: int | None, + gvr_prior_indices: torch.Tensor | None, ) -> torch.Tensor: - row_capacity = radix_indices.shape[0] if radix_indices is not None else scores.shape[0] - request_capacity = request_capacity or max( - sequence_lengths.shape[0], row_capacity // next_n - ) - self._ensure_gvr_buffers(scores, request_capacity, row_capacity) - + assert gvr_prior_indices is not None + capture_graph = scores.is_cuda and torch.cuda.is_current_stream_capturing() num_requests = sequence_lengths.shape[0] - prior_indices = self._gvr_prior_indices[:num_requests] if self.decode_implementation == TopKImplementation.CUDA_GVR: + workspace = self._memory_buffers.get_buffer( + (scores.shape[0], self.top_k), + dtype=scores.dtype, + buffer_name="top_k_cuda_gvr_workspace", + reserve_buffer=capture_graph, + ) + radix_indices, radix_values = self._get_radix_workspace(scores) torch.ops.trtllm.indexer_topk_decode( scores, sequence_lengths, output_indices, next_n, self.top_k, - pre_idx=prior_indices, - heuristic_scratch=self._cuda_gvr_scratch[: scores.shape[0]], + pre_idx=gvr_prior_indices, + heuristic_scratch=workspace, compress_ratio=self.compress_ratio, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, ) else: + if self._num_sms == 0: + self._num_sms = ( + torch.cuda.get_device_properties(scores.device).multi_processor_count + if scores.is_cuda + else 1 + ) row_order = None if num_requests * next_n >= 2 * self._num_sms: - row_order = self._gvr_row_order[:num_requests] + row_order = self._memory_buffers.get_buffer( + (num_requests,), + dtype=torch.int32, + buffer_name="top_k_cute_dsl_gvr_row_order", + reserve_buffer=capture_graph, + ) row_order.copy_(torch.argsort(sequence_lengths, descending=True).to(torch.int32)) torch.ops.trtllm.cute_dsl_gvr_topk_decode( scores, - prior_indices, + gvr_prior_indices, sequence_lengths, output_indices, self.top_k, @@ -241,109 +253,27 @@ def _forward_decode_gvr( max_seq_len=scores.shape[1], order_row=row_order, ) - prior_indices.copy_(output_indices[next_n - 1 :: next_n]) + gvr_prior_indices.copy_(output_indices[next_n - 1 :: next_n]) return output_indices def update_gvr_prior_from_prefill( self, output_indices: torch.Tensor, request_lengths: torch.Tensor, + gvr_prior_indices: torch.Tensor | None, *, request_offset: int = 0, ) -> None: """Update GVR prior indices from each prefill request's last selected row.""" if self.decode_implementation not in _GVR_IMPLEMENTATIONS: return + assert gvr_prior_indices is not None last_rows = (torch.cumsum(request_lengths, dim=0) - 1).to(dtype=torch.long) num_requests = request_lengths.shape[0] - required_rows = request_offset + num_requests - needs_resize = ( - self._gvr_prior_indices.shape[0] < required_rows - or self._gvr_prior_indices.device != output_indices.device - ) - if needs_resize: - decode_initialized = ( - self._cuda_gvr_scratch.numel() > 0 - if self.decode_implementation == TopKImplementation.CUDA_GVR - else self._gvr_row_order.numel() > 0 - ) - if decode_initialized or ( - output_indices.is_cuda and torch.cuda.is_current_stream_capturing() - ): - raise RuntimeError( - "GVR prior indices cannot be resized after decode initialization" - ) - prior_capacity = max(required_rows, self._gvr_prior_indices.shape[0]) - prior_indices = torch.zeros( - (prior_capacity, self.top_k), - dtype=torch.int32, - device=output_indices.device, - ) - prior_indices[: self._gvr_prior_indices.shape[0]].copy_(self._gvr_prior_indices) - self._gvr_prior_indices = prior_indices - self._gvr_prior_indices[request_offset : request_offset + num_requests].copy_( + gvr_prior_indices[request_offset : request_offset + num_requests].copy_( output_indices[last_rows] ) - def _ensure_gvr_buffers( - self, - scores: torch.Tensor, - request_capacity: int, - row_capacity: int, - ) -> None: - """Initialize fixed-address GVR buffers before CUDA Graph capture.""" - needs_prior = ( - self._gvr_prior_indices.shape[0] < request_capacity - or self._gvr_prior_indices.device != scores.device - ) - needs_scratch = self.decode_implementation == TopKImplementation.CUDA_GVR and ( - self._cuda_gvr_scratch.dtype != scores.dtype - or self._cuda_gvr_scratch.device != scores.device - or self._cuda_gvr_scratch.shape[0] < row_capacity - ) - needs_row_order = self.decode_implementation == TopKImplementation.CUTE_DSL_GVR and ( - self._gvr_row_order.device != scores.device - or self._gvr_row_order.shape[0] < request_capacity - ) - needs_resize = needs_prior or needs_scratch or needs_row_order - decode_initialized = ( - self._cuda_gvr_scratch.numel() > 0 - if self.decode_implementation == TopKImplementation.CUDA_GVR - else self._gvr_row_order.numel() > 0 - ) - if needs_resize and ( - decode_initialized or (scores.is_cuda and torch.cuda.is_current_stream_capturing()) - ): - raise RuntimeError("GVR buffers cannot be resized after decode initialization") - - if needs_prior: - prior_capacity = max(request_capacity, self._gvr_prior_indices.shape[0]) - prior_indices = torch.zeros( - (prior_capacity, self.top_k), - dtype=torch.int32, - device=scores.device, - ) - prior_indices[: self._gvr_prior_indices.shape[0]].copy_(self._gvr_prior_indices) - self._gvr_prior_indices = prior_indices - - if self.decode_implementation == TopKImplementation.CUDA_GVR: - if needs_scratch: - self._cuda_gvr_scratch = scores.new_empty((row_capacity, self.top_k)) - return - - if needs_row_order: - self._gvr_row_order = torch.empty( - (request_capacity,), - dtype=torch.int32, - device=scores.device, - ) - if self._num_sms == 0: - self._num_sms = ( - torch.cuda.get_device_properties(scores.device).multi_processor_count - if scores.is_cuda - else 1 - ) - def _forward_prefill_torch( self, scores: torch.Tensor, diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 4d4664d1d249..999625bf6142 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -143,10 +143,10 @@ def test_metadata_warmup_cute_dsl_radix_topk_dispatch( metadata = SimpleNamespace( sparse_metadata_params=SimpleNamespace( enable_heuristic_topk=enable_heuristic, - index_topk=384, ), use_cute_dsl_topk=use_cute_dsl, num_sparse_topk=512, + sparse_mla_topk=384, kv_cache_manager=SimpleNamespace(), _indexer_compress_ratio=compress_ratio, get_indexer_max_seq_len=Mock(return_value=32768), @@ -202,11 +202,11 @@ def test_shared_topk_lifecycle(): metadata.kv_cache_manager = SimpleNamespace(max_blocks_per_seq=2) metadata.enable_context_mla_with_cached_kv = False metadata.enable_indexer_skip = False + metadata.enable_gvr_topk = False metadata.get_empty = Mock( side_effect=lambda _, shape, **kwargs: torch.empty(tuple(shape), dtype=kwargs["dtype"]) ) metadata._create_kv_lens_2d_buffer = Mock() - metadata._create_radix_aux_buffers = Mock() metadata.create_expanded_buffers = Mock() with patch( @@ -3081,14 +3081,6 @@ def make_inputs(n_tokens): max_draft_tokens=md, ) Indexer.prepare(meta0) - # indexer_topk_decode needs caller-owned radix aux buffers for small gen batches. - _radix_bp = 10 - meta0.radix_aux_indices = torch.zeros( - (step0_tokens, _radix_bp, index_topk), device="cuda", dtype=torch.int32 - ) - meta0.radix_aux_logits = torch.zeros( - (step0_tokens, _radix_bp, index_topk), device="cuda", dtype=torch.float32 - ) else: # context: first gen round -- step 0 runs the context/prefill path. step0_tokens = kv_lens.sum().item() meta0 = _create_mock_metadata( @@ -3403,6 +3395,12 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, prefill_implementation=TopKImplementation.CUDA_RADIX, decode_implementation=TopKImplementation.CUTE_DSL_GVR, ) + indexer._enable_heuristic_topk = True + metadata_skip.gvr_prior_indices = torch.zeros( + (cache_manager.num_local_layers, batch_size, index_topk), + device="cuda", + dtype=torch.int32, + ) try: topk_indices_skip = indexer.sparse_attn_indexer( @@ -3412,8 +3410,9 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, raise RuntimeError(f"Indexer skip not available: {e}") last_rows = torch.cumsum(metadata_skip.seq_lens[:batch_size], dim=0) - 1 + local_layer = cache_manager.layer_offsets[layer_idx] torch.testing.assert_close( - indexer.top_k._gvr_prior_indices[:batch_size], + metadata_skip.gvr_prior_indices[local_layer, :batch_size], topk_indices_skip[last_rows], ) diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index 6c5d0355a17b..e5437174b5f2 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for the reusable sparse index-selection Top-K module.""" -from unittest.mock import Mock +from unittest.mock import Mock, call import pytest import torch @@ -94,6 +94,9 @@ def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: output = torch.empty(2, 2, dtype=torch.int32) radix_indices = torch.empty(2, 10, 2, dtype=torch.int32) radix_values = torch.empty(2, 10, 2) + buffers = Mock() + buffers.get_buffer.side_effect = [radix_indices, radix_values] + monkeypatch.setattr(TopK, "_memory_buffers", buffers) top_k( scores, output, @@ -101,10 +104,22 @@ def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: sequence_lengths=logical_lengths, scan_lengths=scan_lengths, next_n=2, - radix_indices=radix_indices, - radix_values=radix_values, ) cute_dsl.assert_not_called() + assert buffers.get_buffer.call_args_list == [ + call( + (2, 10, 2), + dtype=torch.int32, + buffer_name="top_k_radix_indices_workspace", + reserve_buffer=False, + ), + call( + (2, 10, 2), + dtype=torch.float32, + buffer_name="top_k_radix_values_workspace", + reserve_buffer=False, + ), + ] trtllm.assert_called_once_with( scores, logical_lengths, @@ -119,7 +134,7 @@ def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: ) -def test_gvr_owns_prior_state_and_updates_it(monkeypatch) -> None: +def test_gvr_uses_prior_state_and_updates_it(monkeypatch) -> None: gvr = Mock() monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) top_k = TopK( @@ -131,6 +146,7 @@ def test_gvr_owns_prior_state_and_updates_it(monkeypatch) -> None: logical_lengths = torch.tensor([32], dtype=torch.int32) scan_lengths = torch.tensor([8], dtype=torch.int32) output = torch.empty(1, 2, dtype=torch.int32) + prior_indices = torch.zeros(1, 2, dtype=torch.int32) gvr.side_effect = lambda *args, **kwargs: output.copy_(torch.tensor([[5, 3]])) top_k( @@ -140,11 +156,12 @@ def test_gvr_owns_prior_state_and_updates_it(monkeypatch) -> None: sequence_lengths=logical_lengths, scan_lengths=scan_lengths, next_n=1, + gvr_prior_indices=prior_indices, ) args, kwargs = gvr.call_args assert args[0] is scores - assert args[1].data_ptr() == top_k._gvr_prior_indices.data_ptr() + assert args[1] is prior_indices assert args[2] is logical_lengths assert args[3] is output assert args[4] == 2 @@ -154,7 +171,7 @@ def test_gvr_owns_prior_state_and_updates_it(monkeypatch) -> None: "max_seq_len": 8, "order_row": None, } - assert top_k._gvr_prior_indices.tolist() == [[5, 3]] + assert prior_indices.tolist() == [[5, 3]] def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: @@ -163,6 +180,10 @@ def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) next_n = 2 lengths = torch.tensor([4, 1, 8, 2], dtype=torch.int32) + workspace = torch.empty(lengths.shape[0], dtype=torch.int32) + buffers = Mock() + buffers.get_buffer.return_value = workspace + monkeypatch.setattr(TopK, "_memory_buffers", buffers) top_k( torch.randn(lengths.shape[0] * next_n, 8), @@ -171,8 +192,15 @@ def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: sequence_lengths=lengths, scan_lengths=lengths, next_n=next_n, + gvr_prior_indices=torch.zeros(lengths.shape[0], 2, dtype=torch.int32), ) + buffers.get_buffer.assert_called_once_with( + (lengths.shape[0],), + dtype=torch.int32, + buffer_name="top_k_cute_dsl_gvr_row_order", + reserve_buffer=False, + ) row_order = gvr.call_args.kwargs["order_row"] assert row_order is not None assert row_order.tolist() == [2, 0, 3, 1] @@ -181,14 +209,16 @@ def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None: top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32) + prior_indices = torch.zeros(3, 2, dtype=torch.int32) top_k.update_gvr_prior_from_prefill( prefill_indices, torch.tensor([2, 1], dtype=torch.int32), + prior_indices, request_offset=1, ) - assert top_k._gvr_prior_indices.tolist() == [[0, 0], [2, 3], [4, 5]] + assert prior_indices.tolist() == [[0, 0], [2, 3], [4, 5]] def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: @@ -198,6 +228,11 @@ def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: scores = torch.randn(1, 8) lengths = torch.tensor([8], dtype=torch.int32) output = torch.empty((1, 1), dtype=torch.int32) + radix_indices = torch.empty(1, 10, 1, dtype=torch.int32) + radix_values = torch.empty(1, 10, 1) + buffers = Mock() + buffers.get_buffer.side_effect = [radix_indices, radix_values] + monkeypatch.setattr(TopK, "_memory_buffers", buffers) result = top_k( scores, @@ -210,6 +245,7 @@ def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: assert top_k.prefill_implementation == TopKImplementation.CUDA_RADIX assert top_k.decode_implementation == TopKImplementation.CUDA_RADIX assert result is output + assert buffers.get_buffer.call_count == 2 decode.assert_called_once_with( scores, lengths, @@ -219,21 +255,27 @@ def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: pre_idx=None, heuristic_scratch=None, compress_ratio=1, - radix_aux_indices=None, - radix_aux_logits=None, + radix_aux_indices=radix_indices, + radix_aux_logits=radix_values, ) -def test_cuda_gvr_owns_scratch_and_updates_prior(monkeypatch) -> None: +def test_cuda_gvr_reserves_workspace_during_capture_and_updates_prior(monkeypatch) -> None: decode = Mock(side_effect=lambda *args, **kwargs: args[2].copy_(torch.tensor([[3, 1]]))) monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True)) top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) - scores = torch.randn(1, 8) + scores = Mock(shape=(1, 8), dtype=torch.float32, is_cuda=True) lengths = torch.tensor([8], dtype=torch.int32) output = torch.empty(1, 2, dtype=torch.int32) - radix_indices = torch.empty(2, 10, 2, dtype=torch.int32) - radix_values = torch.empty(2, 10, 2) + radix_indices = torch.empty(1, 10, 2, dtype=torch.int32) + radix_values = torch.empty(1, 10, 2) + workspace = torch.empty(1, 2) + prior_indices = torch.zeros(1, 2, dtype=torch.int32) + buffers = Mock() + buffers.get_buffer.side_effect = [workspace, radix_indices, radix_values] + monkeypatch.setattr(TopK, "_memory_buffers", buffers) top_k( scores, @@ -241,49 +283,35 @@ def test_cuda_gvr_owns_scratch_and_updates_prior(monkeypatch) -> None: is_prefill=False, sequence_lengths=lengths, scan_lengths=lengths, - radix_indices=radix_indices, - radix_values=radix_values, + gvr_prior_indices=prior_indices, ) + assert buffers.get_buffer.call_args_list == [ + call( + (scores.shape[0], 2), + dtype=scores.dtype, + buffer_name="top_k_cuda_gvr_workspace", + reserve_buffer=True, + ), + call( + (scores.shape[0], 10, 2), + dtype=torch.int32, + buffer_name="top_k_radix_indices_workspace", + reserve_buffer=True, + ), + call( + (scores.shape[0], 10, 2), + dtype=torch.float32, + buffer_name="top_k_radix_values_workspace", + reserve_buffer=True, + ), + ] runtime_call = decode.call_args_list[-1] - assert runtime_call.kwargs["pre_idx"].data_ptr() == top_k._gvr_prior_indices.data_ptr() - assert runtime_call.kwargs["heuristic_scratch"].data_ptr() == top_k._cuda_gvr_scratch.data_ptr() - assert top_k._gvr_prior_indices.tolist() == [[3, 1], [0, 0]] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_cuda_gvr_buffers_use_scores_device_and_keep_addresses(monkeypatch) -> None: - decode = Mock(side_effect=lambda *args, **kwargs: args[2].zero_()) - monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) - top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) - scores = torch.randn(2, 8, device="cuda") - lengths = torch.tensor([8, 8], dtype=torch.int32, device="cuda") - output = torch.empty(2, 2, dtype=torch.int32, device="cuda") - radix_indices = torch.empty(2, 10, 2, dtype=torch.int32, device="cuda") - radix_values = torch.empty(2, 10, 2, device="cuda") - - pointers = None - for _ in range(2): - top_k( - scores, - output, - is_prefill=False, - sequence_lengths=lengths, - scan_lengths=lengths, - radix_indices=radix_indices, - radix_values=radix_values, - request_capacity=2, - ) - current_pointers = ( - top_k._gvr_prior_indices.data_ptr(), - top_k._cuda_gvr_scratch.data_ptr(), - ) - if pointers is None: - pointers = current_pointers - else: - assert current_pointers == pointers - assert top_k._gvr_prior_indices.device == scores.device - assert top_k._cuda_gvr_scratch.device == scores.device + assert runtime_call.kwargs["pre_idx"] is prior_indices + assert runtime_call.kwargs["heuristic_scratch"].data_ptr() == workspace.data_ptr() + assert runtime_call.kwargs["radix_aux_indices"] is radix_indices + assert runtime_call.kwargs["radix_aux_logits"] is radix_values + assert prior_indices.tolist() == [[3, 1]] def test_unsupported_prefill_implementation_raises() -> None: @@ -297,30 +325,3 @@ def test_unsupported_prefill_implementation_raises() -> None: row_starts=torch.zeros(1, dtype=torch.int32), row_ends=torch.ones(1, dtype=torch.int32), ) - - -def test_gvr_buffers_cannot_grow_after_decode_initialization(monkeypatch) -> None: - decode = Mock(side_effect=lambda *args, **kwargs: args[2].zero_()) - monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) - top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) - scores = torch.randn(1, 8) - lengths = torch.tensor([8], dtype=torch.int32) - - top_k( - scores, - torch.empty(1, 2, dtype=torch.int32), - is_prefill=False, - sequence_lengths=lengths, - scan_lengths=lengths, - request_capacity=1, - ) - - with pytest.raises(RuntimeError, match="cannot be resized"): - top_k( - scores, - torch.empty(1, 2, dtype=torch.int32), - is_prefill=False, - sequence_lengths=lengths, - scan_lengths=lengths, - request_capacity=2, - ) From a41ece9383d29d1ec3d827b85f944c88c7e47d50 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:15:18 -0700 Subject: [PATCH 15/18] [None][fix] address sparse top-k review feedback Scope reusable Top-K workspaces by the scores device, restore once-per-step GVR row ordering in DSA metadata, and document the external GVR state contract. Update focused tests for device-aware workspace keys and metadata-owned row order. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 1 + .../attention_backend/sparse/dsa/metadata.py | 25 ++++ tensorrt_llm/_torch/modules/top_k.py | 115 ++++++++++++------ .../attention/sparse/dsa/test_dsa_indexer.py | 33 +++++ tests/unittest/_torch/modules/test_top_k.py | 39 +++--- 5 files changed, 158 insertions(+), 55 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index eb6110609a9d..78d464fba7ca 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1722,6 +1722,7 @@ def sparse_attn_indexer( gvr_prior_indices=( gvr_prior_indices[:num_generations] if gvr_prior_indices is not None else None ), + gvr_row_order=metadata.kv_lens_row_reorder, ) elif has_decode and metadata.skip_indexer_for_gen_reqs: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py index 64f6396ba273..2854c71b6f96 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -155,6 +155,7 @@ def __post_init__(self): self.enable_gvr_topk = ( sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100 ) + self.kv_lens_row_reorder = None capture_graph = self.is_cuda_graph # Plain DSA has no compression and uses the default [1]. DeepSeek-V4's # metadata params carry the model-specific compression ratios. @@ -451,8 +452,24 @@ def on_update_kv_lens(self): self.scheduler_metadata_buffer_expanded.copy_( scheduler_metadata_buffer_expanded, non_blocking=True ) + self._compute_kv_lens_row_reorder() self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True) + def _compute_kv_lens_row_reorder(self) -> None: + """Prepare the longest-job-first GVR row order once per forward step.""" + next_n = 1 + self.max_draft_tokens + if ( + self.enable_gvr_topk + and self.use_cute_dsl_topk + and self.num_generations * next_n >= 2 * self.num_sms + ): + gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs] + order = torch.argsort(gen_kv_lens, descending=True).to(torch.int32) + self.kv_lens_row_reorder_buffer[: self.num_generations].copy_(order) + self.kv_lens_row_reorder = self.kv_lens_row_reorder_buffer[: self.num_generations] + else: + self.kv_lens_row_reorder = None + def update_for_spec_dec(self): super().update_for_spec_dec() # host @@ -727,6 +744,14 @@ def create_buffers_for_indexer(self, capture_graph=False): capture_graph=capture_graph, ) self.gvr_prior_indices.zero_() + if self.use_cute_dsl_topk: + self.kv_lens_row_reorder_buffer = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences,), + cache_name="kv_lens_row_reorder_buffer", + dtype=torch.int32, + capture_graph=capture_graph, + ) # Create expanded buffers for MTP support self.create_expanded_buffers(capture_graph=capture_graph) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 2fab70b92dca..12c21529de8f 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -30,7 +30,11 @@ class TopKImplementation(str, Enum): class TopK(nn.Module): - """Select Top-K indices for sparse prefill and decode paths.""" + """Select Top-K indices for sparse prefill and decode paths. + + GVR decode state is owned by the caller so it can be shared with the + request metadata and retain a stable address across CUDA Graph replays. + """ _memory_buffers = get_memory_buffers() @@ -51,7 +55,6 @@ def __init__( decode_implementation or TopKImplementation.CUDA_RADIX ) self.compress_ratio = compress_ratio - self._num_sms = 0 def forward( self, @@ -65,8 +68,30 @@ def forward( scan_lengths: torch.Tensor | None = None, next_n: int = 1, gvr_prior_indices: torch.Tensor | None = None, + gvr_row_order: torch.Tensor | None = None, ) -> torch.Tensor: - """Write prefill or decode Top-K indices into ``output_indices``.""" + """Write prefill or decode Top-K indices into ``output_indices``. + + Args: + scores: Top-K input scores with shape ``[num_rows, num_columns]``. + output_indices: Int32 output with shape ``[num_rows, top_k]``. + is_prefill: Whether to run the prefill implementation. + row_starts: Per-row inclusive starts for prefill. + row_ends: Per-row exclusive ends for prefill. + sequence_lengths: Per-request logical KV lengths for decode. + scan_lengths: Per-request score-column lengths for decode. + next_n: Number of decode rows per request. + gvr_prior_indices: Required for GVR decode. Int32 tensor with shape + ``[num_requests, top_k]`` on ``scores.device``. Each call reads + the previous selection and updates it in place with the last + decode row for every request. + gvr_row_order: Optional int32 tensor with shape ``[num_requests]`` + on ``scores.device``. It contains a reusable request ordering + prepared once for the current forward step. + + Returns: + ``output_indices`` after the selected implementation writes it. + """ if is_prefill: assert row_starts is not None and row_ends is not None return self._forward_prefill(scores, row_starts, row_ends, output_indices) @@ -79,6 +104,7 @@ def forward( output_indices, next_n, gvr_prior_indices, + gvr_row_order, ) def _forward_prefill( @@ -116,6 +142,7 @@ def _forward_decode( output_indices: torch.Tensor, next_n: int, gvr_prior_indices: torch.Tensor | None, + gvr_row_order: torch.Tensor | None, ) -> torch.Tensor: if self.decode_implementation == TopKImplementation.TORCH: return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n) @@ -127,6 +154,7 @@ def _forward_decode( output_indices, next_n, gvr_prior_indices, + gvr_row_order, ) return self._forward_decode_radix( @@ -173,6 +201,30 @@ def _forward_decode_radix( ) return output_indices + def _get_workspace( + self, + scores: torch.Tensor, + shape: tuple[int, ...], + dtype: torch.dtype, + buffer_name: str, + ) -> torch.Tensor: + device_buffer_name = f"{buffer_name}_{scores.device}" + if scores.is_cuda: + with torch.cuda.device(scores.device): + capture_graph = torch.cuda.is_current_stream_capturing() + return self._memory_buffers.get_buffer( + shape, + dtype=dtype, + buffer_name=device_buffer_name, + reserve_buffer=capture_graph, + ) + return self._memory_buffers.get_buffer( + shape, + dtype=dtype, + buffer_name=device_buffer_name, + reserve_buffer=False, + ) + def _get_radix_workspace( self, scores: torch.Tensor ) -> tuple[torch.Tensor | None, torch.Tensor | None]: @@ -180,18 +232,17 @@ def _get_radix_workspace( return None, None shape = (scores.shape[0], _MAX_RADIX_BLOCKS_PER_ROW, self.top_k) - capture_graph = scores.is_cuda and torch.cuda.is_current_stream_capturing() - radix_indices = self._memory_buffers.get_buffer( + radix_indices = self._get_workspace( + scores, shape, - dtype=torch.int32, - buffer_name="top_k_radix_indices_workspace", - reserve_buffer=capture_graph, + torch.int32, + "top_k_radix_indices_workspace", ) - radix_values = self._memory_buffers.get_buffer( + radix_values = self._get_workspace( + scores, shape, - dtype=torch.float32, - buffer_name="top_k_radix_values_workspace", - reserve_buffer=capture_graph, + torch.float32, + "top_k_radix_values_workspace", ) return radix_indices, radix_values @@ -202,16 +253,15 @@ def _forward_decode_gvr( output_indices: torch.Tensor, next_n: int, gvr_prior_indices: torch.Tensor | None, + gvr_row_order: torch.Tensor | None, ) -> torch.Tensor: assert gvr_prior_indices is not None - capture_graph = scores.is_cuda and torch.cuda.is_current_stream_capturing() - num_requests = sequence_lengths.shape[0] if self.decode_implementation == TopKImplementation.CUDA_GVR: - workspace = self._memory_buffers.get_buffer( + workspace = self._get_workspace( + scores, (scores.shape[0], self.top_k), - dtype=scores.dtype, - buffer_name="top_k_cuda_gvr_workspace", - reserve_buffer=capture_graph, + scores.dtype, + "top_k_cuda_gvr_workspace", ) radix_indices, radix_values = self._get_radix_workspace(scores) torch.ops.trtllm.indexer_topk_decode( @@ -227,21 +277,6 @@ def _forward_decode_gvr( radix_aux_logits=radix_values, ) else: - if self._num_sms == 0: - self._num_sms = ( - torch.cuda.get_device_properties(scores.device).multi_processor_count - if scores.is_cuda - else 1 - ) - row_order = None - if num_requests * next_n >= 2 * self._num_sms: - row_order = self._memory_buffers.get_buffer( - (num_requests,), - dtype=torch.int32, - buffer_name="top_k_cute_dsl_gvr_row_order", - reserve_buffer=capture_graph, - ) - row_order.copy_(torch.argsort(sequence_lengths, descending=True).to(torch.int32)) torch.ops.trtllm.cute_dsl_gvr_topk_decode( scores, gvr_prior_indices, @@ -251,7 +286,7 @@ def _forward_decode_gvr( next_n=next_n, compress_ratio=self.compress_ratio, max_seq_len=scores.shape[1], - order_row=row_order, + order_row=gvr_row_order, ) gvr_prior_indices.copy_(output_indices[next_n - 1 :: next_n]) return output_indices @@ -264,7 +299,17 @@ def update_gvr_prior_from_prefill( *, request_offset: int = 0, ) -> None: - """Update GVR prior indices from each prefill request's last selected row.""" + """Update GVR prior indices from each prefill request's last row. + + Args: + output_indices: Int32 prefill selections with shape + ``[num_prefill_rows, top_k]``. + request_lengths: Per-request prefill row counts. + gvr_prior_indices: Int32 caller-owned state on + ``output_indices.device`` with shape ``[capacity, top_k]``. + The slice starting at ``request_offset`` is updated in place. + request_offset: First request row to update in the prior state. + """ if self.decode_implementation not in _GVR_IMPLEMENTATIONS: return assert gvr_prior_indices is not None diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 999625bf6142..aa0b82018074 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -176,6 +176,39 @@ def test_metadata_warmup_cute_dsl_radix_topk_dispatch( cute_dsl_radix.assert_not_called() +def test_kv_lens_row_reorder_threshold(): + """Prepare row order only when CuTe DSL GVR has enough decode rows.""" + num_sms = 16 + next_n = 2 + + def make_mock(num_generations, kv_lens_list): + kv_lens_cuda = torch.tensor(kv_lens_list, dtype=torch.int32, device="cuda") + row_order_buffer = torch.zeros(64, dtype=torch.int32, device="cuda") + return SimpleNamespace( + enable_gvr_topk=True, + use_cute_dsl_topk=True, + num_generations=num_generations, + num_sms=num_sms, + max_draft_tokens=next_n - 1, + num_contexts=0, + num_seqs=num_generations, + kv_lens_cuda=kv_lens_cuda, + kv_lens_row_reorder_buffer=row_order_buffer, + kv_lens_row_reorder=None, + ) + + kv_lens = [4, 1, 8, 2, 16, 3, 12, 6, 7, 9, 5, 11, 13, 10, 14, 15] + + metadata_below = make_mock(1, [1000]) + DSAtrtllmAttentionMetadata._compute_kv_lens_row_reorder(metadata_below) + assert metadata_below.kv_lens_row_reorder is None + + metadata_at = make_mock(num_sms, kv_lens) + DSAtrtllmAttentionMetadata._compute_kv_lens_row_reorder(metadata_at) + row_order = metadata_at.kv_lens_row_reorder.cpu().tolist() + assert [kv_lens[i] for i in row_order] == sorted(kv_lens, reverse=True) + + def test_shared_topk_lifecycle(): sparse_config = DeepSeekSparseAttentionConfig( index_n_heads=1, diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index e5437174b5f2..b47320d14e27 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for the reusable sparse index-selection Top-K module.""" +from contextlib import nullcontext from unittest.mock import Mock, call import pytest @@ -110,13 +111,13 @@ def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: call( (2, 10, 2), dtype=torch.int32, - buffer_name="top_k_radix_indices_workspace", + buffer_name="top_k_radix_indices_workspace_cpu", reserve_buffer=False, ), call( (2, 10, 2), dtype=torch.float32, - buffer_name="top_k_radix_values_workspace", + buffer_name="top_k_radix_values_workspace_cpu", reserve_buffer=False, ), ] @@ -174,16 +175,13 @@ def test_gvr_uses_prior_state_and_updates_it(monkeypatch) -> None: assert prior_indices.tolist() == [[5, 3]] -def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: +def test_gvr_uses_caller_prepared_row_order(monkeypatch) -> None: gvr = Mock(side_effect=lambda *args, **kwargs: args[3].zero_()) monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR) next_n = 2 lengths = torch.tensor([4, 1, 8, 2], dtype=torch.int32) - workspace = torch.empty(lengths.shape[0], dtype=torch.int32) - buffers = Mock() - buffers.get_buffer.return_value = workspace - monkeypatch.setattr(TopK, "_memory_buffers", buffers) + row_order = torch.tensor([2, 0, 3, 1], dtype=torch.int32) top_k( torch.randn(lengths.shape[0] * next_n, 8), @@ -193,17 +191,10 @@ def test_gvr_prepares_row_order_at_threshold(monkeypatch) -> None: scan_lengths=lengths, next_n=next_n, gvr_prior_indices=torch.zeros(lengths.shape[0], 2, dtype=torch.int32), + gvr_row_order=row_order, ) - buffers.get_buffer.assert_called_once_with( - (lengths.shape[0],), - dtype=torch.int32, - buffer_name="top_k_cute_dsl_gvr_row_order", - reserve_buffer=False, - ) - row_order = gvr.call_args.kwargs["order_row"] - assert row_order is not None - assert row_order.tolist() == [2, 0, 3, 1] + assert gvr.call_args.kwargs["order_row"] is row_order def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None: @@ -264,9 +255,16 @@ def test_cuda_gvr_reserves_workspace_during_capture_and_updates_prior(monkeypatc decode = Mock(side_effect=lambda *args, **kwargs: args[2].copy_(torch.tensor([[3, 1]]))) monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True)) + device_context = Mock(side_effect=lambda _: nullcontext()) + monkeypatch.setattr(torch.cuda, "device", device_context) top_k = TopK(2, decode_implementation=TopKImplementation.CUDA_GVR) - scores = Mock(shape=(1, 8), dtype=torch.float32, is_cuda=True) + scores = Mock( + shape=(1, 8), + dtype=torch.float32, + is_cuda=True, + device=torch.device("cuda", 3), + ) lengths = torch.tensor([8], dtype=torch.int32) output = torch.empty(1, 2, dtype=torch.int32) radix_indices = torch.empty(1, 10, 2, dtype=torch.int32) @@ -290,22 +288,23 @@ def test_cuda_gvr_reserves_workspace_during_capture_and_updates_prior(monkeypatc call( (scores.shape[0], 2), dtype=scores.dtype, - buffer_name="top_k_cuda_gvr_workspace", + buffer_name="top_k_cuda_gvr_workspace_cuda:3", reserve_buffer=True, ), call( (scores.shape[0], 10, 2), dtype=torch.int32, - buffer_name="top_k_radix_indices_workspace", + buffer_name="top_k_radix_indices_workspace_cuda:3", reserve_buffer=True, ), call( (scores.shape[0], 10, 2), dtype=torch.float32, - buffer_name="top_k_radix_values_workspace", + buffer_name="top_k_radix_values_workspace_cuda:3", reserve_buffer=True, ), ] + assert device_context.call_args_list == [call(scores.device)] * 3 runtime_call = decode.call_args_list[-1] assert runtime_call.kwargs["pre_idx"] is prior_indices assert runtime_call.kwargs["heuristic_scratch"].data_ptr() == workspace.data_ptr() From ab968f5311a52afdc9f898f3957d1ea2f0e72c20 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:21:43 -0700 Subject: [PATCH 16/18] [None][fix] preserve multistream GVR write-back Keep GVR prior write-back in the Indexer so it can overlap with sparse attention on the auxiliary stream and join within the same layer. Apply the lifecycle uniformly to computed, reused, and skipped decode Top-K results, and add focused coverage. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 27 ++++---- tensorrt_llm/_torch/modules/top_k.py | 6 +- .../attention/sparse/dsa/test_dsa_indexer.py | 69 +++++++++++++++++++ tests/unittest/_torch/modules/test_top_k.py | 8 +-- 4 files changed, 90 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 78d464fba7ca..4c429c8ece29 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1731,18 +1731,21 @@ def sparse_attn_indexer( metadata.topk_indices_buffer[num_ctx_tokens:num_tokens, :] ) - if ( - gvr_prior_indices is not None - and has_decode - and (reuse_topk or metadata.skip_indexer_for_gen_reqs) - ): + if gvr_prior_indices is not None and has_decode: next_n = num_gen_tokens // num_generations - # Keep bypassed decode results as the next GVR prior; MTP reuse holds the accepted row. - gvr_prior_indices[:num_generations].copy_( - topk_indices_buffer[token_offset : token_offset + num_gen_tokens][ - next_n - 1 :: next_n - ] - ) + decode_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens] + last_mtp_topk = decode_topk[next_n - 1 :: next_n] + prev_topk_dst = gvr_prior_indices[:num_generations] + if do_multi_stream() and self.aux_stream is not None: + # Overlap the GVR prior write-back with this layer's sparse attention. + self.prev_topk_copy_events[0].record() + with torch.cuda.stream(self.aux_stream): + self.prev_topk_copy_events[0].wait() + prev_topk_dst.copy_(last_mtp_topk) + self.prev_topk_copy_events[1].record() + self._prev_topk_copy_pending = True + else: + prev_topk_dst.copy_(last_mtp_topk) if self.mtp_index_share and metadata.in_mtp_draft_loop and not reuse_topk: rows = None @@ -1787,7 +1790,7 @@ def _mtp_last_accepted_rows( return gen_topk[base + offset] def maybe_join_prev_topk_copy(self) -> None: - """Join the aux-stream heuristic prev_topk write-back, if forked.""" + """Join the aux-stream GVR prior write-back, if forked.""" if self._prev_topk_copy_pending: self.prev_topk_copy_events[1].wait() self._prev_topk_copy_pending = False diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 12c21529de8f..b0c90dc3022c 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -82,9 +82,8 @@ def forward( scan_lengths: Per-request score-column lengths for decode. next_n: Number of decode rows per request. gvr_prior_indices: Required for GVR decode. Int32 tensor with shape - ``[num_requests, top_k]`` on ``scores.device``. Each call reads - the previous selection and updates it in place with the last - decode row for every request. + ``[num_requests, top_k]`` on ``scores.device`` containing the + previous selection. The caller owns its write-back lifecycle. gvr_row_order: Optional int32 tensor with shape ``[num_requests]`` on ``scores.device``. It contains a reusable request ordering prepared once for the current forward step. @@ -288,7 +287,6 @@ def _forward_decode_gvr( max_seq_len=scores.shape[1], order_row=gvr_row_order, ) - gvr_prior_indices.copy_(output_indices[next_n - 1 :: next_n]) return output_indices def update_gvr_prior_from_prefill( diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index aa0b82018074..afbfefea9993 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -54,6 +54,7 @@ transform_local_topk_and_prepare_pool_view, ) from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream from tensorrt_llm._torch.modules.top_k import TopK, TopKImplementation from tensorrt_llm._torch.pyexecutor._util import get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role @@ -209,6 +210,74 @@ def make_mock(num_generations, kv_lens_list): assert [kv_lens[i] for i in row_order] == sorted(kv_lens, reverse=True) +@skip_pre_hopper +def test_gvr_prior_writeback_uses_aux_stream(): + batch_size = 2 + index_topk = 4 + cache_manager, sparse_config = create_dsa_cache_manager( + batch_size=batch_size, + head_dim=128, + tokens_per_block=64, + max_seq_len=64, + num_layers=1, + index_topk=index_topk, + ) + try: + request_ids = list(range(batch_size)) + kv_lens = torch.full((batch_size,), index_topk, dtype=torch.int32) + cache_manager.add_dummy_requests( + request_ids, + kv_lens.tolist(), + is_gen=False, + prepare_resource=True, + ) + metadata = _create_mock_metadata( + request_ids, + batch_size, + num_contexts=0, + num_generations=batch_size, + seq_lens=torch.ones(batch_size, dtype=torch.int32), + kv_lens=kv_lens, + num_cached_tokens=[index_topk - 1] * batch_size, + cache_manager=cache_manager, + num_ctx_tokens=0, + num_tokens=batch_size, + index_topk=index_topk, + enable_indexer_skip=True, + ) + indexer = create_indexer(sparse_config) + indexer._enable_heuristic_topk = True + indexer.aux_stream = torch.cuda.Stream() + metadata.gvr_prior_indices = torch.zeros( + (cache_manager.num_local_layers, batch_size, index_topk), + device="cuda", + dtype=torch.int32, + ) + hidden_states = torch.empty((batch_size, 1), device="cuda") + unused = torch.empty((batch_size, 1), device="cuda") + + with with_multi_stream(True): + topk_indices = indexer.sparse_attn_indexer( + metadata, + hidden_states, + unused, + unused, + unused, + unused, + ) + assert indexer._prev_topk_copy_pending + indexer.maybe_join_prev_topk_copy() + + local_layer = cache_manager.layer_offsets[indexer.layer_idx] + torch.testing.assert_close( + metadata.gvr_prior_indices[local_layer, :batch_size], + topk_indices, + ) + assert not indexer._prev_topk_copy_pending + finally: + cache_manager.shutdown() + + def test_shared_topk_lifecycle(): sparse_config = DeepSeekSparseAttentionConfig( index_n_heads=1, diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index b47320d14e27..e2d27290d28e 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -135,7 +135,7 @@ def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: ) -def test_gvr_uses_prior_state_and_updates_it(monkeypatch) -> None: +def test_gvr_uses_caller_prior_state(monkeypatch) -> None: gvr = Mock() monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) top_k = TopK( @@ -172,7 +172,7 @@ def test_gvr_uses_prior_state_and_updates_it(monkeypatch) -> None: "max_seq_len": 8, "order_row": None, } - assert prior_indices.tolist() == [[5, 3]] + assert prior_indices.tolist() == [[0, 0]] def test_gvr_uses_caller_prepared_row_order(monkeypatch) -> None: @@ -251,7 +251,7 @@ def test_cuda_radix_defaults_dispatch_to_cpp(monkeypatch) -> None: ) -def test_cuda_gvr_reserves_workspace_during_capture_and_updates_prior(monkeypatch) -> None: +def test_cuda_gvr_reserves_workspace_during_capture(monkeypatch) -> None: decode = Mock(side_effect=lambda *args, **kwargs: args[2].copy_(torch.tensor([[3, 1]]))) monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", decode) monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True)) @@ -310,7 +310,7 @@ def test_cuda_gvr_reserves_workspace_during_capture_and_updates_prior(monkeypatc assert runtime_call.kwargs["heuristic_scratch"].data_ptr() == workspace.data_ptr() assert runtime_call.kwargs["radix_aux_indices"] is radix_indices assert runtime_call.kwargs["radix_aux_logits"] is radix_values - assert prior_indices.tolist() == [[3, 1]] + assert prior_indices.tolist() == [[0, 0]] def test_unsupported_prefill_implementation_raises() -> None: From 806644b7660e6f12deb0c125fa4bd08b71185f7a Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Thu, 20 Aug 2026 04:05:37 -0700 Subject: [PATCH 17/18] [None][docs] clarify GVR Top-K contracts Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py | 5 ++++- tensorrt_llm/_torch/modules/top_k.py | 2 ++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 4c429c8ece29..8cc1b7287c3b 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1521,7 +1521,7 @@ def sparse_attn_indexer( :num_ctx_tokens, : ] - # Update GVR state after chunk all-gathers or the dense-index copy. + # Dense skip outputs remain valid priors if a later step enters GVR. if has_prefill: self.top_k.update_gvr_prior_from_prefill( topk_indices_buffer[:num_ctx_tokens], @@ -1712,6 +1712,8 @@ def sparse_attn_indexer( scan_lengths = metadata.gen_indexer_kv_lens_cuda_runtime assert scan_lengths is not None + # Paged MQA logits allocate their score width at + # indexer_max_seq_len, so TopK's CuTe GVR tuning key is stable. self.top_k( logits_decode, topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], @@ -1731,6 +1733,7 @@ def sparse_attn_indexer( metadata.topk_indices_buffer[num_ctx_tokens:num_tokens, :] ) + # Keep the GVR prior current for computed, reused, and dense-skip TopK. if gvr_prior_indices is not None and has_decode: next_n = num_gen_tokens // num_generations decode_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens] diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index b0c90dc3022c..83ce03a1d737 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -228,6 +228,8 @@ def _get_radix_workspace( self, scores: torch.Tensor ) -> tuple[torch.Tensor | None, torch.Tensor | None]: if scores.dtype != torch.float32: + # The C++ bf16/fp16 entry has no split-work tier or aux-buffer + # arguments and rejects widths that would require split work. return None, None shape = (scores.shape[0], _MAX_RADIX_BLOCKS_PER_ROW, self.top_k) From 7a95a0a8ec947e000b30244ff6b82ef70b03f580 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Fri, 21 Aug 2026 00:01:04 -0700 Subject: [PATCH 18/18] [None][refactor] refine GVR Top-K interface Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../attention_backend/sparse/dsa/indexer.py | 16 +++++---- tensorrt_llm/_torch/modules/top_k.py | 36 ++++++++++--------- tests/unittest/_torch/modules/test_top_k.py | 14 +++++--- 3 files changed, 38 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 8cc1b7287c3b..e32e806edb0e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -1712,8 +1712,14 @@ def sparse_attn_indexer( scan_lengths = metadata.gen_indexer_kv_lens_cuda_runtime assert scan_lengths is not None - # Paged MQA logits allocate their score width at - # indexer_max_seq_len, so TopK's CuTe GVR tuning key is stable. + gvr_ext_kwargs = ( + { + "gvr_prior_indices": gvr_prior_indices[:num_generations], + "gvr_row_order": metadata.kv_lens_row_reorder, + } + if gvr_prior_indices is not None + else None + ) self.top_k( logits_decode, topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], @@ -1721,10 +1727,8 @@ def sparse_attn_indexer( sequence_lengths=gen_kv_lens_cuda, scan_lengths=scan_lengths, next_n=next_n, - gvr_prior_indices=( - gvr_prior_indices[:num_generations] if gvr_prior_indices is not None else None - ), - gvr_row_order=metadata.kv_lens_row_reorder, + max_seq_len=indexer_max_seq_len, + gvr_ext_kwargs=gvr_ext_kwargs, ) elif has_decode and metadata.skip_indexer_for_gen_reqs: diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 83ce03a1d737..704ce11283bc 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -67,8 +67,8 @@ def forward( sequence_lengths: torch.Tensor | None = None, scan_lengths: torch.Tensor | None = None, next_n: int = 1, - gvr_prior_indices: torch.Tensor | None = None, - gvr_row_order: torch.Tensor | None = None, + max_seq_len: int | None = None, + gvr_ext_kwargs: dict[str, torch.Tensor | None] | None = None, ) -> torch.Tensor: """Write prefill or decode Top-K indices into ``output_indices``. @@ -81,12 +81,12 @@ def forward( sequence_lengths: Per-request logical KV lengths for decode. scan_lengths: Per-request score-column lengths for decode. next_n: Number of decode rows per request. - gvr_prior_indices: Required for GVR decode. Int32 tensor with shape - ``[num_requests, top_k]`` on ``scores.device`` containing the - previous selection. The caller owns its write-back lifecycle. - gvr_row_order: Optional int32 tensor with shape ``[num_requests]`` - on ``scores.device``. It contains a reusable request ordering - prepared once for the current forward step. + max_seq_len: Maximum decode score width used for GVR kernel tuning. + gvr_ext_kwargs: GVR-only keyword arguments. ``gvr_prior_indices`` + is the required caller-owned int32 previous selection with + shape ``[num_requests, top_k]`` on ``scores.device``. + ``gvr_row_order`` is an optional int32 request ordering with + shape ``[num_requests]`` on the same device. Returns: ``output_indices`` after the selected implementation writes it. @@ -102,8 +102,8 @@ def forward( scan_lengths, output_indices, next_n, - gvr_prior_indices, - gvr_row_order, + max_seq_len, + gvr_ext_kwargs, ) def _forward_prefill( @@ -140,8 +140,8 @@ def _forward_decode( scan_lengths: torch.Tensor, output_indices: torch.Tensor, next_n: int, - gvr_prior_indices: torch.Tensor | None, - gvr_row_order: torch.Tensor | None, + max_seq_len: int | None, + gvr_ext_kwargs: dict[str, torch.Tensor | None] | None, ) -> torch.Tensor: if self.decode_implementation == TopKImplementation.TORCH: return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n) @@ -152,8 +152,8 @@ def _forward_decode( sequence_lengths, output_indices, next_n, - gvr_prior_indices, - gvr_row_order, + max_seq_len=max_seq_len, + **(gvr_ext_kwargs or {}), ) return self._forward_decode_radix( @@ -253,8 +253,9 @@ def _forward_decode_gvr( sequence_lengths: torch.Tensor, output_indices: torch.Tensor, next_n: int, - gvr_prior_indices: torch.Tensor | None, - gvr_row_order: torch.Tensor | None, + max_seq_len: int | None, + gvr_prior_indices: torch.Tensor | None = None, + gvr_row_order: torch.Tensor | None = None, ) -> torch.Tensor: assert gvr_prior_indices is not None if self.decode_implementation == TopKImplementation.CUDA_GVR: @@ -278,6 +279,7 @@ def _forward_decode_gvr( radix_aux_logits=radix_values, ) else: + assert max_seq_len is not None torch.ops.trtllm.cute_dsl_gvr_topk_decode( scores, gvr_prior_indices, @@ -286,7 +288,7 @@ def _forward_decode_gvr( self.top_k, next_n=next_n, compress_ratio=self.compress_ratio, - max_seq_len=scores.shape[1], + max_seq_len=max_seq_len, order_row=gvr_row_order, ) return output_indices diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index e2d27290d28e..9682d7981ab5 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -157,7 +157,8 @@ def test_gvr_uses_caller_prior_state(monkeypatch) -> None: sequence_lengths=logical_lengths, scan_lengths=scan_lengths, next_n=1, - gvr_prior_indices=prior_indices, + max_seq_len=16, + gvr_ext_kwargs={"gvr_prior_indices": prior_indices}, ) args, kwargs = gvr.call_args @@ -169,7 +170,7 @@ def test_gvr_uses_caller_prior_state(monkeypatch) -> None: assert kwargs == { "next_n": 1, "compress_ratio": 4, - "max_seq_len": 8, + "max_seq_len": 16, "order_row": None, } assert prior_indices.tolist() == [[0, 0]] @@ -190,8 +191,11 @@ def test_gvr_uses_caller_prepared_row_order(monkeypatch) -> None: sequence_lengths=lengths, scan_lengths=lengths, next_n=next_n, - gvr_prior_indices=torch.zeros(lengths.shape[0], 2, dtype=torch.int32), - gvr_row_order=row_order, + max_seq_len=8, + gvr_ext_kwargs={ + "gvr_prior_indices": torch.zeros(lengths.shape[0], 2, dtype=torch.int32), + "gvr_row_order": row_order, + }, ) assert gvr.call_args.kwargs["order_row"] is row_order @@ -281,7 +285,7 @@ def test_cuda_gvr_reserves_workspace_during_capture(monkeypatch) -> None: is_prefill=False, sequence_lengths=lengths, scan_lengths=lengths, - gvr_prior_indices=prior_indices, + gvr_ext_kwargs={"gvr_prior_indices": prior_indices}, ) assert buffers.get_buffer.call_args_list == [