From 574c5184989ecbf0f8d31badd3006e8a5dee0426 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 25 May 2026 03:27:55 +0000 Subject: [PATCH 01/12] [None][refactor] Add derived attention properties for the thop.attention call site Squashed snapshot of PR #14279 pre-rebase. Combines all 17 commits from derived-attn-props onto its old merge-base. The final state on rebase target (origin/main) is reconstructed in the next commit if any of this got displaced. Major themes folded in: - Add derived @property accessors on AttentionForwardArgs and TrtllmAttention / TrtllmAttentionMetadata so the ~75-arg thop.attention call site becomes an explicit-kwarg call sourced from rich objects. - Add AttentionSparseArgs dataclass nested in AttentionForwardArgs. - Adapt to main's class-based trtllm_gen dispatch. - Make MLA preprocess thop ops accept None for kv_scale_* (kernels default to 1.0). - Drop unused kv_scale param per MLA preprocess op (each only needs one direction). - Null-safe descale ptrs in FlashMLA fp8 kernel so Python can pass None end-to-end. - Drop the now-unused ``out_scale_sf`` field and ``_get_attention_output_orig_quant`` helper. - Add test_attention_op_sync.py: AST sync test between thop call site and C++ binding schema. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../kernels/flashMLA/flash_fwd_mla_kernel.h | 4 +- cpp/tensorrt_llm/thop/attentionOp.cpp | 14 +- cpp/tensorrt_llm/thop/attentionOp.h | 2 +- cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp | 29 +- .../_torch/attention_backend/interface.py | 39 +- .../_torch/attention_backend/sparse/dsa.py | 3 +- .../_torch/attention_backend/trtllm.py | 430 +++++++++------ .../_torch/attention_backend/trtllm_gen.py | 18 +- .../auto_deploy/custom_ops/mla/trtllm_mla.py | 2 - tensorrt_llm/_torch/modules/attention.py | 30 +- .../test_attention_op_sync.py | 519 ++++++++++++++++++ 11 files changed, 866 insertions(+), 224 deletions(-) create mode 100644 tests/unittest/_torch/attention_backend/test_attention_op_sync.py diff --git a/cpp/tensorrt_llm/kernels/flashMLA/flash_fwd_mla_kernel.h b/cpp/tensorrt_llm/kernels/flashMLA/flash_fwd_mla_kernel.h index d34104cb845b..23772e585f94 100644 --- a/cpp/tensorrt_llm/kernels/flashMLA/flash_fwd_mla_kernel.h +++ b/cpp/tensorrt_llm/kernels/flashMLA/flash_fwd_mla_kernel.h @@ -686,8 +686,8 @@ __global__ void __launch_bounds__(Kernel_traits::kNThreads, 1, 1) float scale_softmax_log2 = params.scale_softmax_log2; if constexpr (Kernel_traits::Is_FP8) { - float descale_q = __ldg(params.descale_q_ptr); - descale_k = __ldg(params.descale_k_ptr); + float descale_q = params.descale_q_ptr ? __ldg(params.descale_q_ptr) : 1.f; + descale_k = params.descale_k_ptr ? __ldg(params.descale_k_ptr) : 1.f; scale_softmax = scale_softmax * descale_q * descale_k; scale_softmax_log2 = scale_softmax_log2 * descale_q * descale_k; } diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 086e999e8b1d..79447cc93172 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -927,7 +927,7 @@ void attention(torch::Tensor q, std::optional k, std::optional rotary_cos_sin, std::optional latent_cache, std::optional q_pe, std::optional block_ids_per_seq, std::optional attention_sinks, bool const is_fused_qkv, bool const update_kv_cache, - int64_t const predicted_tokens_per_seq, int64_t const layer_idx, int64_t const num_heads, + int64_t const predicted_tokens_per_seq, int64_t const local_layer_idx, int64_t const num_heads, int64_t const num_kv_heads, int64_t const head_size, std::optional const tokens_per_block, int64_t const max_num_requests, int64_t const max_context_length, int64_t const attention_window_size, int64_t const beam_width, int64_t const mask_type, int64_t const quant_mode, double const q_scaling, @@ -956,7 +956,7 @@ void attention(torch::Tensor q, std::optional k, std::optional compressed_kv_cache_pool_ptr) { - TLLM_LOG_TRACE("Attention op starts at layer %d", layer_idx); + TLLM_LOG_TRACE("Attention op starts at layer %d", local_layer_idx); // Use these tensors to infer if the attention is using KV cache bool const use_kv_cache = kv_cache_block_offsets.has_value() && host_kv_cache_pool_pointers.has_value() && host_kv_cache_pool_mapping.has_value(); @@ -1044,7 +1044,7 @@ void attention(torch::Tensor q, std::optional k, std::optional(); op->mType = dtype; op->mFMHAForceFP32Acc = dtype == nvinfer1::DataType::kBF16; - op->mLayerIdx = layer_idx; + op->mLayerIdx = local_layer_idx; op->mNumHeads = num_heads; op->mNumKVHeads = num_kv_heads; op->mHeadSize = head_size; @@ -1162,13 +1162,13 @@ void attention(torch::Tensor q, std::optional k, std::optional, hash> op_cache; if (auto it = op_cache.find(cache_key); it != op_cache.end()) { - TLLM_LOG_TRACE("Attention op for layer %d is cached", layer_idx); + TLLM_LOG_TRACE("Attention op for layer %d is cached", local_layer_idx); op = it->second; } else { - TLLM_LOG_TRACE( - "Preparing new attention op for layer %d with cache key: %s", layer_idx, to_string(cache_key).c_str()); + TLLM_LOG_TRACE("Preparing new attention op for layer %d with cache key: %s", local_layer_idx, + to_string(cache_key).c_str()); op->initialize(); runner->prepare(*op); op_cache[cache_key] = op; @@ -1257,7 +1257,7 @@ void attention(torch::Tensor q, std::optional k, std::optional k, std::optional rotary_cos_sin, std::optional latent_cache, std::optional q_pe, std::optional block_ids_per_seq, std::optional attention_sinks, bool const is_fused_qkv, bool const update_kv_cache, - int64_t const predicted_tokens_per_seq, int64_t const layer_idx, int64_t const num_heads, + int64_t const predicted_tokens_per_seq, int64_t const local_layer_idx, int64_t const num_heads, int64_t const num_kv_heads, int64_t const head_size, std::optional const tokens_per_block, int64_t const max_num_requests, int64_t const max_context_length, int64_t const attention_window_size, int64_t const beam_width, int64_t const mask_type, int64_t const quant_mode, double const q_scaling, diff --git a/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp b/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp index f6c5eba6bc68..69cceb64de65 100644 --- a/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp +++ b/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp @@ -129,17 +129,14 @@ std::vector loadPagedKVCacheForMLA(torch::ScalarType out_dtype, i attention_window_size, beam_width, 0 /*seq_offset*/, true /*is_mla_enable*/, torch::elementSize(out_dtype)) .kvCacheBuffer; - float const* kv_scale_orig_quant_ptr = nullptr; float const* kv_scale_quant_orig_ptr = nullptr; if (kv_cache_quant_mode.hasKvCacheQuant()) { TLLM_CHECK_WITH_INFO(kv_cache_quant_mode.hasFp8KvCache(), "Only FP8 KV cache is supported for now"); - TORCH_CHECK(kv_scale_orig_quant.has_value()); - TORCH_CHECK(kv_scale_quant_orig.has_value()); - kv_scale_orig_quant_ptr = kv_scale_orig_quant.value().data_ptr(); + } + if (kv_scale_quant_orig.has_value()) + { kv_scale_quant_orig_ptr = kv_scale_quant_orig.value().data_ptr(); - TLLM_CHECK(kv_scale_orig_quant_ptr != nullptr); - TLLM_CHECK(kv_scale_quant_orig_ptr != nullptr); } std::vector outputs; @@ -217,16 +214,10 @@ std::vector loadChunkedKVCacheForMLA(torch::ScalarType out_dtype, attention_window_size, beam_width, 0 /*seq_offset*/, true /*is_mla_enable*/, torch::elementSize(out_dtype)) .kvCacheBuffer; - float const* kv_scale_orig_quant_ptr = nullptr; float const* kv_scale_quant_orig_ptr = nullptr; - if (kv_cache_quant_mode.hasKvCacheQuant()) + if (kv_scale_quant_orig.has_value()) { - TORCH_CHECK(kv_scale_orig_quant.has_value()); - TORCH_CHECK(kv_scale_quant_orig.has_value()); - kv_scale_orig_quant_ptr = kv_scale_orig_quant.value().data_ptr(); kv_scale_quant_orig_ptr = kv_scale_quant_orig.value().data_ptr(); - TLLM_CHECK(kv_scale_orig_quant_ptr != nullptr); - TLLM_CHECK(kv_scale_quant_orig_ptr != nullptr); } std::vector outputs; @@ -321,16 +312,13 @@ void MLARopeAppendPagedKVAssignQ(torch::Tensor& q, torch::Tensor& latent_cache, .kvCacheBuffer; float const* kv_scale_orig_quant_ptr = nullptr; - float const* kv_scale_quant_orig_ptr = nullptr; if (kv_cache_quant_mode.hasKvCacheQuant()) { TLLM_CHECK_WITH_INFO(kv_cache_quant_mode.hasFp8KvCache(), "Only FP8 KV cache is supported for now"); - TORCH_CHECK(kv_scale_orig_quant.has_value()); - TORCH_CHECK(kv_scale_quant_orig.has_value()); + } + if (kv_scale_orig_quant.has_value()) + { kv_scale_orig_quant_ptr = kv_scale_orig_quant.value().data_ptr(); - kv_scale_quant_orig_ptr = kv_scale_quant_orig.value().data_ptr(); - TLLM_CHECK(kv_scale_orig_quant_ptr != nullptr); - TLLM_CHECK(kv_scale_quant_orig_ptr != nullptr); } if (input_dtype == torch::kFloat16) @@ -425,7 +413,6 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) ", Tensor kv_cache_block_offsets" ", Tensor host_kv_cache_pool_pointers" ", Tensor host_kv_cache_pool_mapping" - ", Tensor? kv_scale_orig_quant" ", Tensor? kv_scale_quant_orig" ", int layer_idx" ", int lora_size" @@ -454,7 +441,6 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) ", Tensor kv_cache_block_offsets" ", Tensor host_kv_cache_pool_pointers" ", Tensor host_kv_cache_pool_mapping" - ", Tensor? kv_scale_orig_quant" ", Tensor? kv_scale_quant_orig" ", int layer_idx" ", int lora_size" @@ -491,7 +477,6 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) ", Tensor host_kv_cache_pool_pointers" ", Tensor host_kv_cache_pool_mapping" ", Tensor? kv_scale_orig_quant" - ", Tensor? kv_scale_quant_orig" ", int layer_idx" ", int tokens_per_block" ", int attention_window_size" diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index bdcdbe7dc6fc..2921672d3f81 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -16,8 +16,8 @@ from ..speculative.spec_tree_manager import SpecTreeManager from tensorrt_llm._utils import get_hf_rope_theta, maybe_pin_memory -from tensorrt_llm.functional import (PositionEmbeddingType, RopeEmbeddingUtils, - RotaryScalingType) +from tensorrt_llm.functional import (AttentionMaskType, PositionEmbeddingType, + RopeEmbeddingUtils, RotaryScalingType) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig @@ -694,6 +694,20 @@ class CustomAttentionMask(str, Enum): AttentionMask = Union[PredefinedAttentionMask, CustomAttentionMask] +@dataclass(kw_only=True, slots=True) +class AttentionSparseArgs: + """Sparse-attention inputs passed to the attention op. + + Backends without sparse attention leave ``AttentionForwardArgs.sparse`` + at its default-constructed value (all-``None`` / ``0`` fields). + """ + sparse_kv_indices: Optional[torch.Tensor] = None + sparse_kv_offsets: Optional[torch.Tensor] = None + sparse_attn_indices: Optional[torch.Tensor] = None + sparse_attn_offsets: Optional[torch.Tensor] = None + sparse_attn_indices_block_size: int = 0 + + @dataclass(kw_only=True, slots=True) class AttentionForwardArgs: """Per-forward optional arguments for attention backends.""" @@ -702,9 +716,8 @@ class AttentionForwardArgs: output_sf: Optional[torch.Tensor] = None out_scale: Optional[torch.Tensor] = None - out_scale_sf: Optional[torch.Tensor] = None - kv_scales_sf: Optional[torch.Tensor] = None - kv_scales_sf_inv: Optional[torch.Tensor] = None + kv_scale_orig_quant: Optional[torch.Tensor] = None + kv_scale_quant_orig: Optional[torch.Tensor] = None attention_mask: AttentionMask = PredefinedAttentionMask.CAUSAL attention_input_type: AttentionInputType = AttentionInputType.mixed @@ -735,6 +748,22 @@ class AttentionForwardArgs: topk_indices: Optional[torch.Tensor] = None + is_fused_qkv: bool = False + update_kv_cache: bool = True + + sparse: AttentionSparseArgs = field(default_factory=AttentionSparseArgs) + + @property + def mask_type(self) -> int: + """Integer mask type accepted by the C++ attention op + (``causal`` or ``padding``).""" + if self.attention_mask == PredefinedAttentionMask.CAUSAL: + return int(AttentionMaskType.causal) + if self.attention_mask == PredefinedAttentionMask.FULL: + return int(AttentionMaskType.padding) + raise ValueError( + f"Unexpected attention mask type: {self.attention_mask!r}") + _ATTENTION_FORWARD_ARGS_FIELDS = frozenset( AttentionForwardArgs.__dataclass_fields__) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 48b2b2b3d8de..204b81d95a58 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -2532,8 +2532,7 @@ def mla_rope_append_paged_kv_assign_q( block_offsets, metadata.kv_cache_manager.kv_cache_pool_pointers, metadata.kv_cache_manager.kv_cache_pool_mapping, - self.kv_scale_orig_quant, - self.kv_scale_quant_orig, + None, # kv_scale_orig_quant self.get_local_layer_idx(metadata), metadata.kv_cache_manager.tokens_per_block, metadata.kv_cache_manager.max_seq_len, diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 66899f3d5b26..9a4eb2deff6d 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -2,7 +2,7 @@ import math import os import weakref -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -22,14 +22,33 @@ get_model_extra_attrs) from .interface import (AttentionBackend, AttentionForwardArgs, AttentionInputType, AttentionMask, AttentionMetadata, - KVCacheParams, MLAParams, PositionalEmbeddingParams, - PredefinedAttentionMask, RopeParams, - merge_attention_forward_args) + AttentionSparseArgs, KVCacheParams, MLAParams, + PositionalEmbeddingParams, PredefinedAttentionMask, + RopeParams, merge_attention_forward_args) # Enable TRTLLM-Gen attention backend via environment variable (default: on). _TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION = (os.environ.get( "TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", "0") == "1") +# ``AttentionForwardArgs`` fields that this backend does not consume. +# Sync test (test_attention_op_sync.py) requires every other field to map to a +# kwarg name, a @property on the dataclass, or a field that some @property +# transitively reads; entries here are exempt. +_THOP_EXCLUDED_FIELDS: frozenset = frozenset({ + "topk_indices", # DSA-only + "is_generation", # trtllm_gen dispatch flag + "attention_mask_data", # custom-mask code path +}) + +# ``thop.attention`` kwargs hard-wired to a literal at the call site (no +# rich object owns them). Sync test enforces both the kwarg name and the +# literal value. +_THOP_LITERALS: dict = { + "sink_token_length": 0, + "sparse_mla_topk_lens": None, + "compressed_kv_cache_pool_ptr": None, +} + @functools.cache def generate_spec_decoding_position_offsets(max_num_requests: int, @@ -129,6 +148,48 @@ class TrtllmAttentionMetadata(AttentionMetadata): init=False, repr=False) + use_paged_context_fmha: bool = field(init=False, default=False, repr=False) + + # ``DSAtrtllmAttentionMetadata`` overrides this; the dense path keeps 0. + num_sparse_topk: int = 0 + + @property + def effective_workspace(self) -> Optional[torch.Tensor]: + """Attention-kernel workspace, switching to the CUDA-graph copy under capture.""" + return self.cuda_graph_workspace if self.is_cuda_graph else self.workspace + + @property + def helix_tensor_params(self) -> List[Optional[torch.Tensor]]: + """``[helix_position_offsets, helix_is_inactive_rank]`` — the positional + helix tensor list expected by the C++ attention op.""" + return [self.helix_position_offsets, self.helix_is_inactive_rank] + + @property + def spec_decoding_bool_params(self) -> List[bool]: + """``[is_spec_decoding_enabled, use_spec_decoding, is_spec_dec_tree]`` — + the positional bool list expected by the C++ attention op.""" + return [ + self.is_spec_decoding_enabled, + self.use_spec_decoding, + self.is_spec_dec_tree, + ] + + @property + def spec_decoding_position_offsets_for_cpp(self) -> Optional[torch.Tensor]: + """``spec_decoding_position_offsets`` reshaped to the 2D layout the C++ + kernel expects. 1D inputs (dynamic-tree shorthand) are viewed as + ``(max_num_requests, -1)``.""" + offsets = self.spec_decoding_position_offsets + if offsets is not None and offsets.dim() == 1: + return offsets.view(self.max_num_requests, -1) + return offsets + + @property + def max_context_length(self) -> int: + """``min(max_seq_len - 1, max_num_tokens)`` — upper bound for a single + context window.""" + return min(self.max_seq_len - 1, self.max_num_tokens) + @property def max_seq_len(self) -> int: """ @@ -174,6 +235,11 @@ def __post_init__(self) -> None: super().__post_init__() self.enable_helix = self.mapping.has_cp_helix( ) if self.mapping is not None else False + self.use_paged_context_fmha = ( + self.runtime_features.chunked_prefill + or self.runtime_features.cache_reuse + or self.runtime_features.has_speculative_draft_tokens + ) if self.runtime_features is not None else False self._post_init_with_buffers(self.cuda_graph_buffers) def update_position_offsets_for_cpp(self, query_len: int) -> None: @@ -986,6 +1052,22 @@ def generate_spec_decoding_generation_length(self, runtime_draft_len): def is_sm_version_trtllm_gen_kernel(self, sm): return not (sm < 100 or sm in [120, 121]) + @property + def spec_decoding_tensor_params(self) -> List[Optional[torch.Tensor]]: + """Positional spec-decoding tensor list for the C++ attention op. + Includes three Blackwell-tree mask tensors on SM versions that take + the trtllm-gen kernel.""" + params = [ + self.spec_decoding_generation_lengths, + self.spec_decoding_position_offsets_for_cpp, + self.spec_decoding_packed_mask, + ] + if self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()): + params.append(self.spec_decoding_bl_tree_mask_offset) + params.append(self.spec_decoding_bl_tree_mask) + params.append(self.spec_bl_tree_first_sparse_mask_offset_kv) + return params + class TrtllmAttention(AttentionBackend[TrtllmAttentionMetadata]): @@ -1062,11 +1144,7 @@ def __init__( self.print_skip_softmax_stat = os.environ.get( "TRTLLM_PRINT_SKIP_SOFTMAX_STAT", "0") == "1" - self.kv_cache_scaling_factor = torch.ones(1, - dtype=torch.float32, - device='cuda') - self.kv_scale_quant_orig = self.kv_cache_scaling_factor - self.kv_scale_orig_quant = 1.0 / self.kv_cache_scaling_factor + self.local_layer_idx: Optional[int] = None if not skip_create_weights_in_init: self.update_quant_config(self.quant_config) @@ -1235,13 +1313,52 @@ def create_output(self, q, *, is_quantize_output: bool, dtype=out_dtype) ] - def _get_mask_type(self, - attention_mask: AttentionMask) -> AttentionMaskType: - if attention_mask == PredefinedAttentionMask.CAUSAL: - return AttentionMaskType.causal - if attention_mask == PredefinedAttentionMask.FULL: - return AttentionMaskType.padding - raise ValueError("Unexpected attention mask type") + @property + def rotary_embedding_dim(self) -> int: + return self.rope_params.dim + + @property + def rotary_embedding_base(self) -> float: + return self.rope_params.theta + + @property + def rotary_embedding_scale_type(self) -> int: + return int(self.rope_params.scale_type) + + @property + def rotary_embedding_scales(self) -> List[float]: + """``[scale, short_m_scale, long_m_scale]`` — the positional RoPE-scale + list expected by the C++ attention op.""" + return [ + self.rope_params.scale, + self.rope_params.short_m_scale, + self.rope_params.long_m_scale, + ] + + @property + def rotary_embedding_max_position_info(self) -> List[int]: + """``[max_positions, original_max_positions]`` — the positional + RoPE-positions list expected by the C++ attention op.""" + return [ + self.rope_params.max_positions, + self.rope_params.original_max_positions, + ] + + @property + def skip_softmax_threshold_scale_factor_prefill(self) -> Optional[float]: + """Prefill skip-softmax threshold; ``None`` unless + ``sparse_attention_config`` is a ``SkipSoftmaxAttentionConfig``.""" + if isinstance(self.sparse_attention_config, SkipSoftmaxAttentionConfig): + return self.sparse_attention_config.threshold_scale_factor_prefill + return None + + @property + def skip_softmax_threshold_scale_factor_decode(self) -> Optional[float]: + """Decode skip-softmax threshold; ``None`` unless + ``sparse_attention_config`` is a ``SkipSoftmaxAttentionConfig``.""" + if isinstance(self.sparse_attention_config, SkipSoftmaxAttentionConfig): + return self.sparse_attention_config.threshold_scale_factor_decode + return None def _get_trtllm_gen_backend( self) -> trtllm_gen.FlashInferTrtllmGenAttention: @@ -1257,38 +1374,19 @@ def _run( q: torch.Tensor, k: Optional[torch.Tensor], v: Optional[torch.Tensor], - output: torch.Tensor, - output_sf: Optional[torch.Tensor], metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, - use_paged_context_fmha: bool, - sparse_kv_indices: Optional[torch.Tensor], - sparse_kv_offsets: Optional[torch.Tensor], - sparse_attn_indices: Optional[torch.Tensor], - sparse_attn_offsets: Optional[torch.Tensor], - sparse_attn_indices_block_size: int, - num_sparse_topk: int, - sparse_mla_topk_lens: Optional[torch.Tensor], - compressed_kv_cache_pool_ptr: Optional[int], - skip_softmax_threshold_scale_factor_prefill: Optional[float], - skip_softmax_threshold_scale_factor_decode: Optional[float], ) -> None: - is_fused_qkv = not metadata.is_cross and k is None - update_kv_cache = not metadata.is_cross or k is not None - assert (is_fused_qkv and k is None - and v is None) or (not is_fused_qkv and k is not None - and v is not None) - attention_input_type = forward_args.attention_input_type if not self.is_mla_enable: - if is_fused_qkv: + if forward_args.is_fused_qkv: qkv_hidden_size = (self.num_heads + 2 * self.num_kv_heads) * self.head_dim assert q.shape[1] == qkv_hidden_size else: q_hidden_size = self.num_heads * self.head_dim assert q.shape[1] == q_hidden_size - if update_kv_cache: + if forward_args.update_kv_cache: kv_hidden_size = self.num_kv_heads * self.head_dim assert k.shape[1] == kv_hidden_size assert v.shape[1] == kv_hidden_size @@ -1297,18 +1395,18 @@ def _run( assert k.shape[0] == num_tokens assert v.shape[0] == num_tokens else: - is_sparse_attn = sparse_attn_indices is not None and sparse_attn_indices.numel( + is_sparse_attn = forward_args.sparse.sparse_attn_indices is not None and forward_args.sparse.sparse_attn_indices.numel( ) > 0 if attention_input_type == AttentionInputType.context_only and is_sparse_attn: - assert is_fused_qkv + assert forward_args.is_fused_qkv qkv_hidden_size = self.num_heads * (self.kv_lora_rank + self.qk_rope_head_dim) elif attention_input_type == AttentionInputType.context_only: - assert not is_fused_qkv + assert not forward_args.is_fused_qkv qkv_hidden_size = self.num_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim) elif attention_input_type == AttentionInputType.generation_only: - assert is_fused_qkv + assert forward_args.is_fused_qkv qkv_hidden_size = self.num_heads * (self.kv_lora_rank + self.qk_rope_head_dim) else: @@ -1324,7 +1422,6 @@ def _run( assert metadata.prompt_lens_cpu_runtime.shape[0] == batch_size assert metadata.host_request_types_runtime.shape[0] == batch_size - mask_type = self._get_mask_type(forward_args.attention_mask) self._ensure_rope_table_size(metadata.max_seq_len) rotary_embedding_dim = self.rope_params.dim @@ -1365,26 +1462,8 @@ def _run( if self.print_skip_softmax_stat: self.skip_softmax_stat.zero_() - use_nvfp4_output = output_sf is not None - out_scale = (forward_args.out_scale_sf - if use_nvfp4_output else forward_args.out_scale) - kv_scale_orig_quant = (self.kv_scale_orig_quant - if forward_args.kv_scales_sf_inv is None else - forward_args.kv_scales_sf_inv) - kv_scale_quant_orig = (self.kv_scale_quant_orig - if forward_args.kv_scales_sf is None else - forward_args.kv_scales_sf) - mrope_rotary_cos_sin = forward_args.mrope_rotary_cos_sin - mrope_position_deltas = forward_args.mrope_position_deltas - workspace = metadata.workspace if not metadata.is_cuda_graph else metadata.cuda_graph_workspace - flash_mla_tile_scheduler_metadata = ( - metadata.flash_mla_tile_scheduler_metadata - if metadata.enable_flash_mla else None) - flash_mla_num_splits = metadata.flash_mla_num_splits if metadata.enable_flash_mla else None - attention_window_size = (forward_args.attention_window_size - or metadata.max_seq_len) - max_context_length = min(metadata.max_seq_len - 1, - metadata.max_num_tokens) + if forward_args.attention_window_size is None: + forward_args.attention_window_size = metadata.max_seq_len helix_active = metadata.helix_position_offsets is not None use_sage_attn = (forward_args.sage_attn_num_elts_per_blk_q > 0 @@ -1394,14 +1473,11 @@ def _run( use_trtllm_gen = False if _TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION: trtllm_gen_backend = self._get_trtllm_gen_backend() - trtllm_gen_forward_args = replace(forward_args, - output=output, - output_sf=output_sf) use_trtllm_gen = trtllm_gen_backend.is_supported( q, metadata=metadata, - forward_args=trtllm_gen_forward_args, - mask_type=int(mask_type), + forward_args=forward_args, + mask_type=int(forward_args.mask_type), active_helix=helix_active, use_sage_attn=use_sage_attn, )[0] @@ -1410,11 +1486,15 @@ def _run( trtllm_gen_backend.attention( q, metadata=metadata, - forward_args=trtllm_gen_forward_args, - mask_type=int(mask_type), - use_paged_context_fmha=use_paged_context_fmha, + forward_args=forward_args, + mask_type=int(forward_args.mask_type), + use_paged_context_fmha=metadata.use_paged_context_fmha, ) else: + # Every kwarg sources from ``self`` / ``metadata`` / + # ``forward_args`` (with ``forward_args.sparse`` for sparse-attn + # inputs) or a literal allowlisted in ``_THOP_LITERALS``. + # ``test_attention_op_sync.py`` enforces this statically. thop.attention( q, k, @@ -1479,31 +1559,90 @@ def _run( forward_args.softmax_stats_tensor, spec_decoding_bool_params, spec_decoding_tensor_params, - sparse_kv_indices, - sparse_kv_offsets, - sparse_attn_indices, - sparse_attn_offsets, - sparse_attn_indices_block_size, - num_sparse_topk, - sparse_mla_topk_lens, - skip_softmax_threshold_scale_factor_prefill, - skip_softmax_threshold_scale_factor_decode, - self.skip_softmax_stat, - forward_args.cu_q_seqlens, - forward_args.cu_kv_seqlens, - forward_args.fmha_scheduler_counter, - forward_args.mla_bmm1_scale, - forward_args.mla_bmm2_scale, - forward_args.quant_q_buffer, + num_sparse_topk=metadata.num_sparse_topk, + flash_mla_tile_scheduler_metadata=metadata. flash_mla_tile_scheduler_metadata, - flash_mla_num_splits, - forward_args.sage_attn_num_elts_per_blk_q, - forward_args.sage_attn_num_elts_per_blk_k, - forward_args.sage_attn_num_elts_per_blk_v, - forward_args.sage_attn_qk_int8, + flash_mla_num_splits=metadata.flash_mla_num_splits, num_contexts=metadata.num_contexts, num_ctx_tokens=metadata.num_ctx_tokens, - compressed_kv_cache_pool_ptr=compressed_kv_cache_pool_ptr, + max_context_length=metadata.max_context_length, + + # --- Per-call (AttentionForwardArgs) --- + out_scale=forward_args.out_scale, + kv_scale_orig_quant=forward_args.kv_scale_orig_quant, + kv_scale_quant_orig=forward_args.kv_scale_quant_orig, + latent_cache=forward_args.latent_cache, + q_pe=forward_args.q_pe, + attention_sinks=forward_args.attention_sinks, + mask_type=forward_args.mask_type, + attention_input_type=int(forward_args.attention_input_type), + attention_window_size=forward_args.attention_window_size, + chunked_prefill_buffer_batch_size=forward_args. + chunked_prefill_buffer_batch_size, + mrope_rotary_cos_sin=forward_args.mrope_rotary_cos_sin, + mrope_position_deltas=forward_args.mrope_position_deltas, + softmax_stats_tensor=forward_args.softmax_stats_tensor, + cu_q_seqlens=forward_args.cu_q_seqlens, + cu_kv_seqlens=forward_args.cu_kv_seqlens, + fmha_scheduler_counter=forward_args.fmha_scheduler_counter, + mla_bmm1_scale=forward_args.mla_bmm1_scale, + mla_bmm2_scale=forward_args.mla_bmm2_scale, + quant_q_buffer=forward_args.quant_q_buffer, + sage_attn_num_elts_per_blk_q=forward_args. + sage_attn_num_elts_per_blk_q, + sage_attn_num_elts_per_blk_k=forward_args. + sage_attn_num_elts_per_blk_k, + sage_attn_num_elts_per_blk_v=forward_args. + sage_attn_num_elts_per_blk_v, + sage_attn_qk_int8=forward_args.sage_attn_qk_int8, + is_fused_qkv=forward_args.is_fused_qkv, + update_kv_cache=forward_args.update_kv_cache, + + # --- Module config (TrtllmAttention) --- + rotary_inv_freq=self.rotary_inv_freq, + rotary_cos_sin=self.rotary_cos_sin, + predicted_tokens_per_seq=self.predicted_tokens_per_seq, + local_layer_idx=self.local_layer_idx, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + quant_mode=self.quant_mode, + q_scaling=self.q_scaling, + position_embedding_type=self.position_embedding_type, + rotary_embedding_dim=self.rotary_embedding_dim, + rotary_embedding_base=self.rotary_embedding_base, + rotary_embedding_scale_type=self.rotary_embedding_scale_type, + rotary_embedding_scales=self.rotary_embedding_scales, + rotary_embedding_max_position_info=self. + rotary_embedding_max_position_info, + is_mla_enable=self.is_mla_enable, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + rope_append=self.rope_append, + attention_chunk_size=self.attention_chunk_size, + skip_softmax_threshold_scale_factor_prefill=self. + skip_softmax_threshold_scale_factor_prefill, + skip_softmax_threshold_scale_factor_decode=self. + skip_softmax_threshold_scale_factor_decode, + skip_softmax_stat=self.skip_softmax_stat, + + # --- Sparse-specific (AttentionForwardArgs.sparse) --- + sparse_kv_indices=forward_args.sparse.sparse_kv_indices, + sparse_kv_offsets=forward_args.sparse.sparse_kv_offsets, + sparse_attn_indices=forward_args.sparse.sparse_attn_indices, + sparse_attn_offsets=forward_args.sparse.sparse_attn_offsets, + sparse_attn_indices_block_size=forward_args.sparse. + sparse_attn_indices_block_size, + + # --- Literals intentionally None / 0 (see _THOP_LITERALS) --- + # ``sparse_mla_topk_lens`` and ``compressed_kv_cache_pool_ptr`` + # stay as literal ``None`` until DeepSeek V4 sparse-MLA lands. + sink_token_length=0, + sparse_mla_topk_lens=None, + compressed_kv_cache_pool_ptr=None, ) if self.print_skip_softmax_stat: @@ -1530,31 +1669,22 @@ def forward( ) assert not metadata.is_cross, "TRT-LLM Attention does not support cross attention yet." - use_paged_context_fmha = ( - metadata.runtime_features.chunked_prefill - or metadata.runtime_features.cache_reuse - or metadata.runtime_features.has_speculative_draft_tokens - ) if metadata.runtime_features else False - - # This is a workaround for https://nvbugs/5624818 - # Paged context FMHA is forced on SM90 for correctness + # SM90 forces ``use_paged_context_fmha`` on for correctness + # (https://nvbugs/5624818). if get_sm_version() == 90: - use_paged_context_fmha = True + metadata.use_paged_context_fmha = True # Sparse mqa/gqa attention uses generation kernel which reads Q from qPtr (separate buffer). # Force paged context FMHA so QKV preprocessing writes Q to q_buf_2_. if (self.sparse_attention_config is not None and getattr( self.sparse_attention_config, 'algorithm', None) == 'mqa_gqa'): - use_paged_context_fmha = True + metadata.use_paged_context_fmha = True if self.is_mla_enable: # Context MLA uses separate qkv instead of paged_context_fmha - use_paged_context_fmha = False + metadata.use_paged_context_fmha = False - output = forward_args.output - output_sf = forward_args.output_sf - if output is None: - # Output is not provided. + if forward_args.output is None: is_gen_only = (forward_args.attention_input_type == AttentionInputType.generation_only) outputs = self.create_output( @@ -1562,34 +1692,36 @@ def forward( is_quantize_output=forward_args.out_scale is not None, metadata=metadata, attention_mask=forward_args.attention_mask, - use_paged_context_fmha=use_paged_context_fmha, + use_paged_context_fmha=metadata.use_paged_context_fmha, is_mla_enable=self.is_mla_enable, is_gen_only=is_gen_only, ) - - output = outputs[0] - output_sf = outputs[1] if len(outputs) == 2 else None - - sparse_kv_indices, sparse_kv_offsets, sparse_attn_indices, sparse_attn_offsets = None, None, None, None - sparse_attn_indices_block_size = 1 - skip_softmax_threshold_scale_factor_prefill = None - skip_softmax_threshold_scale_factor_decode = None - num_sparse_topk = getattr(metadata, 'num_sparse_topk', 0) - sparse_mla_topk_lens = None - compressed_kv_cache_pool_ptr = None - if self.sparse_attention_config is not None: - if isinstance(self.sparse_attention_config, - SkipSoftmaxAttentionConfig): - skip_softmax_threshold_scale_factor_prefill = self.sparse_attention_config.threshold_scale_factor_prefill - skip_softmax_threshold_scale_factor_decode = self.sparse_attention_config.threshold_scale_factor_decode - - else: - sparse_kv_indices, sparse_kv_offsets = self.sparse_kv_predict( - q, k, metadata, forward_args) - sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict( - q, k, metadata, forward_args) - sparse_attn_indices_block_size = self.sparse_attention_config.get_indices_block_size( - ) + forward_args.output = outputs[0] + forward_args.output_sf = outputs[1] if len(outputs) == 2 else None + + forward_args.is_fused_qkv = not metadata.is_cross and k is None + forward_args.update_kv_cache = not metadata.is_cross or k is not None + assert (forward_args.is_fused_qkv and k is None + and v is None) or (not forward_args.is_fused_qkv + and k is not None and v is not None) + + # ``SkipSoftmax`` configs contribute nothing here — their thresholds + # are read via the ``skip_softmax_threshold_scale_factor_*`` + # @property accessors directly on ``self``. + if (self.sparse_attention_config is not None and not isinstance( + self.sparse_attention_config, SkipSoftmaxAttentionConfig)): + kv_idx, kv_off = self.sparse_kv_predict(q, k, metadata, + forward_args) + at_idx, at_off = self.sparse_attn_predict(q, k, metadata, + forward_args) + forward_args.sparse = AttentionSparseArgs( + sparse_kv_indices=kv_idx, + sparse_kv_offsets=kv_off, + sparse_attn_indices=at_idx, + sparse_attn_offsets=at_off, + sparse_attn_indices_block_size=self.sparse_attention_config. + get_indices_block_size(), + ) # Compute FlashMLA tile-scheduler metadata once per forward pass. # The flag is reset in prepare_flash_mla() and update_for_spec_dec() to trigger @@ -1603,24 +1735,17 @@ def forward( metadata._flash_mla_metadata_valid = True # Blackwell first_sparse: refresh at layer 0 before kernel launch. - layer_idx = self.get_local_layer_idx(metadata) - if layer_idx == 0 and (metadata.spec_bl_tree_first_sparse_mask_offset_kv - is not None - and metadata._seq_lens_cuda is not None): + if self.get_local_layer_idx(metadata) == 0 and ( + metadata.spec_bl_tree_first_sparse_mask_offset_kv is not None + and metadata._seq_lens_cuda is not None): metadata.update_blackwell_first_sparse_mask_offset() - self._run(q, k, v, output, output_sf, metadata, forward_args, - use_paged_context_fmha, sparse_kv_indices, sparse_kv_offsets, - sparse_attn_indices, sparse_attn_offsets, - sparse_attn_indices_block_size, num_sparse_topk, - sparse_mla_topk_lens, compressed_kv_cache_pool_ptr, - skip_softmax_threshold_scale_factor_prefill, - skip_softmax_threshold_scale_factor_decode) + self._run(q, k, v, metadata, forward_args) - if output_sf is None: - return output + if forward_args.output_sf is None: + return forward_args.output else: - return output, output_sf + return forward_args.output, forward_args.output_sf @classmethod def support_fused_rope(cls) -> bool: @@ -1684,8 +1809,7 @@ def load_paged_kv_cache_for_mla( metadata.kv_cache_block_offsets, metadata.kv_cache_manager.kv_cache_pool_pointers, metadata.kv_cache_manager.kv_cache_pool_mapping, - self.kv_scale_orig_quant, - self.kv_scale_quant_orig, + None, # kv_scale_quant_orig self.get_local_layer_idx(metadata), self.mla_params.kv_lora_rank, self.mla_params.qk_rope_head_dim, @@ -1730,8 +1854,7 @@ def load_chunked_kv_cache_for_mla( metadata.kv_cache_block_offsets, metadata.kv_cache_manager.kv_cache_pool_pointers, metadata.kv_cache_manager.kv_cache_pool_mapping, - self.kv_scale_orig_quant, - self.kv_scale_quant_orig, + None, # kv_scale_quant_orig self.get_local_layer_idx(metadata), self.mla_params.kv_lora_rank, self.mla_params.qk_rope_head_dim, @@ -1774,8 +1897,7 @@ def mla_rope_append_paged_kv_assign_q( metadata.kv_cache_block_offsets, metadata.kv_cache_manager.kv_cache_pool_pointers, metadata.kv_cache_manager.kv_cache_pool_mapping, - self.kv_scale_orig_quant, - self.kv_scale_quant_orig, + None, # kv_scale_orig_quant self.get_local_layer_idx(metadata), metadata.kv_cache_manager.tokens_per_block, metadata.kv_cache_manager.max_seq_len, @@ -1889,8 +2011,8 @@ def mla_rope_generation( metadata.kv_cache_block_offsets, metadata.kv_cache_manager.kv_cache_pool_pointers, metadata.kv_cache_manager.kv_cache_pool_mapping, - self.kv_scale_orig_quant, - self.kv_scale_quant_orig, + None, # kv_scale_orig_quant + None, # kv_scale_quant_orig out_scale, metadata.block_ids_per_seq, helix_tensor_params, diff --git a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py index 4fd3bc0d0143..4d09adc046bd 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py @@ -619,23 +619,11 @@ def _get_kv_scale_params( return kv_scale_orig_quant, kv_scale_quant_orig - @staticmethod - def _get_attention_output_orig_quant( - forward_args: AttentionForwardArgs, - ) -> Optional[torch.Tensor]: - return ( - forward_args.out_scale_sf - if forward_args.output_sf is not None - else forward_args.out_scale - ) - @staticmethod def _get_mrope_rotary_cos_sin( forward_args: AttentionForwardArgs, ) -> Optional[torch.Tensor]: - if forward_args.mrope_config is None: - return None - return forward_args.mrope_config.get("mrope_rotary_cos_sin") + return forward_args.mrope_rotary_cos_sin def is_supported( self, @@ -989,7 +977,7 @@ def run_context( kv_scale_orig_quant, kv_scale_quant_orig = self._get_kv_scale_params( params.forward, params.kv_cache_quant_mode ) - attention_output_orig_quant = self._get_attention_output_orig_quant(params.forward) + attention_output_orig_quant = params.forward.out_scale mrope_rotary_cos_sin = self._get_mrope_rotary_cos_sin(params.forward) ( @@ -1098,7 +1086,7 @@ def run_generation( kv_scale_orig_quant, kv_scale_quant_orig = self._get_kv_scale_params( params.forward, params.kv_cache_quant_mode ) - attention_output_orig_quant = self._get_attention_output_orig_quant(params.forward) + attention_output_orig_quant = params.forward.out_scale ( q_processed, kv_pool, diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py index 37017dd5e946..d36d030cb402 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py @@ -1208,7 +1208,6 @@ def _handle_prefill_thop_cached_kv( host_kv_cache_pool_pointers, planner.host_pool_mapping, planner.kv_scale_orig_quant, - planner.kv_scale_quant_orig, _CONTEXT_LAYER_OFFSET, tokens_per_block, max_context_length, @@ -1277,7 +1276,6 @@ def _handle_prefill_thop_cached_kv( kv_cache_block_offsets, host_kv_cache_pool_pointers, planner.host_pool_mapping, - planner.kv_scale_orig_quant, planner.kv_scale_quant_orig, _CONTEXT_LAYER_OFFSET, kv_lora_rank, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 8f44c008ee9c..f7a9c8991f75 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -712,7 +712,7 @@ def _attn_impl( # then exchange and combine across CP ranks. # NOTE: The helix post-process combine step works on unquantized # (BF16/FP16) partial outputs and softmax stats from each rank. - # We intentionally skip passing out_scale/out_scale_sf to FMHA here + # We intentionally skip passing out_scale to FMHA here # so it produces BF16 output. After combining, the downstream o_proj # linear layer handles quantization (FP8/NVFP4) in its apply() method. if self.mapping.has_cp_helix() and attn_metadata.num_contexts == 0: @@ -741,21 +741,24 @@ def _attn_impl( attn_output = self._helix_post_process(attn_output, softmax_stats) return attn_output, None + # Don't set out_scale if o_proj has pre_quant_scale — this prevents + # FP8/FP4 output and keeps attention output in BF16 for better + # precision when applying pre_quant_scale. Also don't set out_scale + # if LoRA is active — LoRA grouped_gemm doesn't support FP8. + # When ``output_sf`` is set the kernel writes NVFP4 output and reads + # ``input_scale`` as its scaling factor; otherwise it reads + # ``inv_input_scale``. out_scale = None - out_scale_sf = None - # Don't set out_scale if o_proj has pre_quant_scale - this prevents FP8/FP4 output - # and keeps attention output in BF16 for better precision when applying pre_quant_scale - # Also don't set out_scale if LoRA is active - LoRA grouped_gemm doesn't support FP8 if self._use_quantize_output() and not has_lora: - out_scale = self.o_proj.inv_input_scale - out_scale_sf = self.o_proj.input_scale + out_scale = (self.o_proj.input_scale if output_sf is not None else + self.o_proj.inv_input_scale) - kv_scales_sf = None - kv_scales_sf_inv = None + kv_scale_orig_quant = None + kv_scale_quant_orig = None if self.quant_config is not None and self.quant_config.layer_quant_mode.has_fp4_kv_cache( ): - kv_scales_sf = self.qkv_proj.kv_scales - kv_scales_sf_inv = self.qkv_proj.inv_kv_scales + kv_scale_orig_quant = self.qkv_proj.inv_kv_scales + kv_scale_quant_orig = self.qkv_proj.kv_scales attn_output = self.attn.forward( q, @@ -764,9 +767,8 @@ def _attn_impl( attn_metadata, forward_args=AttentionForwardArgs( out_scale=out_scale, - out_scale_sf=out_scale_sf, - kv_scales_sf=kv_scales_sf, - kv_scales_sf_inv=kv_scales_sf_inv, + kv_scale_orig_quant=kv_scale_orig_quant, + kv_scale_quant_orig=kv_scale_quant_orig, attention_mask=attention_mask, mrope_rotary_cos_sin=mrope_rotary_cos_sin, mrope_position_deltas=mrope_position_deltas, diff --git a/tests/unittest/_torch/attention_backend/test_attention_op_sync.py b/tests/unittest/_torch/attention_backend/test_attention_op_sync.py new file mode 100644 index 000000000000..91aee99d4c78 --- /dev/null +++ b/tests/unittest/_torch/attention_backend/test_attention_op_sync.py @@ -0,0 +1,519 @@ +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Static sync test for the inline ``thop.attention(...)`` call in +``TrtllmAttention._run``. + +That call is the single explicit-kwarg call site for the C++ ``thop.attention`` +binding. This test parses both the call site (Python AST) and the C++ +function declaration in ``attentionOp.h`` (text/regex), and enforces: + +1. Every C++ parameter name appears at the call site (and nothing extra). +2. Every call-site kwarg sourced as ``root.attr[.attr...]`` resolves on + exactly one of ``self`` / ``metadata`` / ``forward_args``, and its + declared C++ type matches the source attribute's Python type at a + coarse-category level (tensor / int / bool / float / list-of-X). +3. Every dataclass field reachable from ``AttentionForwardArgs`` (including + nested dataclass sub-bags like ``AttentionSparseArgs``) is consumed at + the call site — directly, transitively via a @property of the + containing class, or listed in ``_THOP_EXCLUDED_FIELDS``. +4. Every kwarg passed as a literal constant matches an entry in + ``_THOP_LITERALS`` (both name and value). + +The test is AST-only (no kernel run) so it fails fast and runs without a GPU. +""" + +import ast +import dataclasses +import inspect +import pathlib +import re +import textwrap +import typing +from dataclasses import fields + +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention_backend.trtllm import ( + _THOP_EXCLUDED_FIELDS, + _THOP_LITERALS, + TrtllmAttention, + TrtllmAttentionMetadata, +) + +# Roots used as the LHS of attribute chains at the call site. Match the +# parameter names inside ``TrtllmAttention._run``. +_SOURCE_CLASSES = { + "self": TrtllmAttention, + "metadata": TrtllmAttentionMetadata, + "forward_args": AttentionForwardArgs, +} + +# The C++ attention() declaration is the single source of truth for kwarg +# names, ordering, and types. +_HEADER = pathlib.Path(__file__).resolve().parents[4] / ("cpp/tensorrt_llm/thop/attentionOp.h") + + +# ---- C++ declaration parser ------------------------------------------------- + + +_TENSOR_RE = re.compile(r"\btorch::Tensor\b") +_QUANT_MODE_RE = re.compile(r"\bcommon::QuantMode\b") +_INT64_RE = re.compile(r"\bint64_t\b") +_DOUBLE_RE = re.compile(r"\bdouble\b") +_BOOL_RE = re.compile(r"\bbool\b") + + +def _split_top_level(s: str, sep: str = ",") -> list[str]: + """Split ``s`` by ``sep`` at angle-bracket depth 0.""" + out: list[str] = [] + depth = 0 + buf: list[str] = [] + for ch in s: + if ch == "<": + depth += 1 + buf.append(ch) + elif ch == ">": + depth -= 1 + buf.append(ch) + elif ch == sep and depth == 0: + out.append("".join(buf).strip()) + buf = [] + else: + buf.append(ch) + if buf: + out.append("".join(buf).strip()) + return out + + +def _strip_inner_type(cpp_type: str) -> str: + """Strip ``std::optional<...>`` and ``const`` / references to expose the + inner element type. Idempotent.""" + t = cpp_type.replace("const", "").replace("&", "").strip() + m = re.fullmatch(r"std::optional<\s*(.+)\s*>", t) + if m: + t = m.group(1).strip() + return t + + +def _cpp_category(cpp_type: str) -> str: + """Coarse Python-side category for a C++ parameter type. + + Returns one of ``tensor`` / ``int`` / ``bool`` / ``float`` / + ``list_tensor`` / ``list_int`` / ``list_bool`` / ``list_float`` / + ``unknown``. Optional/const wrappers and references are ignored. + """ + bare = _strip_inner_type(cpp_type) + if bare.startswith("std::vector<"): + inner_match = re.fullmatch(r"std::vector<\s*(.+)\s*>", bare) + if inner_match: + inner_cat = _cpp_category(inner_match.group(1)) + return f"list_{inner_cat}" + if _TENSOR_RE.search(bare): + return "tensor" + if _BOOL_RE.fullmatch(bare): + return "bool" + if _INT64_RE.fullmatch(bare) or _QUANT_MODE_RE.fullmatch(bare): + return "int" + if _DOUBLE_RE.fullmatch(bare): + return "float" + return "unknown" + + +def _parse_attention_decl() -> list[tuple[str, str]]: + """Parse the ``void attention(...)`` declaration in ``attentionOp.h`` + and return ``[(name, cpp_type), ...]`` in declaration order.""" + src = _HEADER.read_text() + # Strip line comments to keep the regex tidy. + src = re.sub(r"//[^\n]*", "", src) + m = re.search(r"void\s+attention\s*\(", src) + if not m: + raise AssertionError(f"Could not find void attention(...) in {_HEADER}") + # Walk forward, matching the opening paren at m.end()-1 to its close. + start = m.end() - 1 + depth = 0 + end = None + for i in range(start, len(src)): + if src[i] == "(": + depth += 1 + elif src[i] == ")": + depth -= 1 + if depth == 0: + end = i + break + if end is None: + raise AssertionError("Unbalanced parens in attention() declaration") + body = src[start + 1 : end] + + params: list[tuple[str, str]] = [] + for raw in _split_top_level(body): + # Drop ``= default`` suffix at top level only. + param = _split_top_level(raw, sep="=")[0].strip() + # The name is the trailing identifier. Strip trailing '&' or '*' + # attached to the type, not the name. + m = re.match(r"(.*?)([A-Za-z_]\w*)\s*$", param, re.DOTALL) + if not m: + raise AssertionError(f"Could not parse param: {param!r}") + cpp_type = m.group(1).strip() + name = m.group(2) + params.append((name, cpp_type)) + return params + + +def _binding_kwargs() -> set[str]: + """Set of parameter names declared on ``void attention(...)``.""" + return {name for name, _ in _parse_attention_decl()} + + +def _binding_types() -> dict[str, str]: + """Map parameter name → declared C++ type.""" + return dict(_parse_attention_decl()) + + +# ---- Call-site AST helpers -------------------------------------------------- + + +def _parse_thop_attention_call() -> ast.Call: + """Locate the single ``thop.attention(...)`` call inside ``_run``.""" + src = textwrap.dedent(inspect.getsource(TrtllmAttention._run)) + tree = ast.parse(src) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "attention" + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "thop" + ): + return node + raise AssertionError("Could not find thop.attention(...) call in TrtllmAttention._run") + + +def _attribute_path(node: ast.AST) -> tuple[str, tuple[str, ...]] | None: + """If ``node`` is a pure attribute chain ``Name.attr1.attr2...``, return + ``(root_name_id, (attr1, attr2, ...))``. Otherwise return ``None``. + """ + if not isinstance(node, ast.Attribute): + return None + attrs: list[str] = [] + current: ast.AST = node + while isinstance(current, ast.Attribute): + attrs.append(current.attr) + current = current.value + if not isinstance(current, ast.Name): + return None + return current.id, tuple(reversed(attrs)) + + +def _classify_kwargs() -> tuple[ + dict[str, tuple[str, tuple[str, ...]]], dict[str, object], set[str] +]: + """Split the call site's kwargs into three buckets: + + - ``attr_kwargs``: ``kwarg=source.attr[...]`` (or + ``kwarg=int(source.attr)``) → ``{kwarg: (root, path)}``. + - ``literal_kwargs``: ``kwarg=`` → ``{kwarg: value}``. + - ``other_kwargs``: kwargs whose value is anything else (e.g. a bare + Name like ``q``). + """ + call = _parse_thop_attention_call() + attr_kwargs: dict[str, tuple[str, tuple[str, ...]]] = {} + literal_kwargs: dict[str, object] = {} + other_kwargs: set[str] = set() + for kw in call.keywords: + v = kw.value + # int(source.attr) → equivalent to source.attr for sync purposes. + if ( + isinstance(v, ast.Call) + and isinstance(v.func, ast.Name) + and v.func.id == "int" + and len(v.args) == 1 + ): + path = _attribute_path(v.args[0]) + if path is not None: + attr_kwargs[kw.arg] = path + continue + other_kwargs.add(kw.arg) + continue + if isinstance(v, ast.Constant): + literal_kwargs[kw.arg] = v.value + continue + path = _attribute_path(v) + if path is not None: + attr_kwargs[kw.arg] = path + continue + other_kwargs.add(kw.arg) + return attr_kwargs, literal_kwargs, other_kwargs + + +# ---- Python-side attribute & type resolution -------------------------------- + + +def _runtime_instance_attrs(cls) -> set[str]: + """Names assigned as ``self. = ...`` anywhere in ``cls`` or any of + its base classes.""" + cache = _runtime_instance_attrs._cache # type: ignore[attr-defined] + if cls in cache: + return cache[cls] + + def _walk_target(tgt: ast.AST, names: set[str]) -> None: + if isinstance(tgt, (ast.Tuple, ast.List)): + for elt in tgt.elts: + _walk_target(elt, names) + elif ( + isinstance(tgt, ast.Attribute) + and isinstance(tgt.value, ast.Name) + and tgt.value.id == "self" + ): + names.add(tgt.attr) + + names: set[str] = set() + for base in cls.__mro__: + if base is object: + continue + try: + src = textwrap.dedent(inspect.getsource(base)) + except (OSError, TypeError): + continue + for node in ast.walk(ast.parse(src)): + if isinstance(node, ast.Assign): + for tgt in node.targets: + _walk_target(tgt, names) + elif isinstance(node, (ast.AnnAssign, ast.AugAssign)): + _walk_target(node.target, names) + cache[cls] = names + return names + + +_runtime_instance_attrs._cache = {} # type: ignore[attr-defined] + + +def _has_attr_or_field(cls, name: str) -> bool: + if hasattr(cls, name): + return True + try: + if name in {f.name for f in fields(cls)}: + return True + except TypeError: + pass + return name in _runtime_instance_attrs(cls) + + +def _dataclass_field_type(cls, name: str): + try: + f = cls.__dataclass_fields__.get(name) # type: ignore[attr-defined] + except AttributeError: + return None + if f is None: + return None + return f.type if not isinstance(f.type, str) else None + + +def _resolve_path(root_cls, path: tuple[str, ...]): + """Walk ``path[:-1]`` from ``root_cls`` and return ``(leaf_cls, + leaf_attr)``. Returns ``(None, leaf_attr)`` if an intermediate step is + not a dataclass field.""" + cls = root_cls + for step in path[:-1]: + nxt = _dataclass_field_type(cls, step) + if nxt is None: + return None, path[-1] + cls = nxt + return cls, path[-1] + + +def _python_category(py_type) -> str: + """Coarse classification of a Python annotation, mirroring + ``_cpp_category``. Returns ``unknown`` for anything we can't classify + confidently (the type check is then skipped for that kwarg).""" + # Unwrap Optional[X] / Union[X, None]. + origin = typing.get_origin(py_type) + if origin is typing.Union: + args = [a for a in typing.get_args(py_type) if a is not type(None)] + if len(args) == 1: + return _python_category(args[0]) + return "unknown" + if origin in (list, typing.List): + args = typing.get_args(py_type) + if args: + return f"list_{_python_category(args[0])}" + return "list_unknown" + if py_type is bool: + return "bool" + if py_type is int: + return "int" + if py_type is float: + return "float" + if isinstance(py_type, type): + if py_type.__name__ == "Tensor": + return "tensor" + return "unknown" + + +# ---- Tests ------------------------------------------------------------------ + + +def test_call_site_kwargs_match_binding_kwargs(): + """The call site's keyword set must equal the C++ binding's keyword set.""" + cpp_kwargs = _binding_kwargs() + attr_kwargs, literal_kwargs, other_kwargs = _classify_kwargs() + call_site_kwargs = set(attr_kwargs) | set(literal_kwargs) | other_kwargs + assert call_site_kwargs == cpp_kwargs, ( + f"missing in call site: {cpp_kwargs - call_site_kwargs}, " + f"unknown to C++ binding: {call_site_kwargs - cpp_kwargs}" + ) + + +def test_each_source_attr_kwarg_resolves_uniquely(): + """``root.attr[.attr...]`` kwargs must resolve along the chain to an + attribute that exists on exactly one of the source classes, and the + leaf's Python-category must match the C++ kwarg's category (where + both are unambiguously known).""" + attr_kwargs, _, _ = _classify_kwargs() + cpp_types = _binding_types() + for thop_kwarg, (root, path) in attr_kwargs.items(): + assert root in _SOURCE_CLASSES, ( + f"thop kwarg `{thop_kwarg}` sources from unknown root `{root}` — " + f"call site must use one of {set(_SOURCE_CLASSES)}." + ) + leaf_cls, leaf_attr = _resolve_path(_SOURCE_CLASSES[root], path) + chain = ".".join((root, *path)) + assert leaf_cls is not None, ( + f"thop kwarg `{thop_kwarg}` reads `{chain}` but an intermediate " + f"step is not a dataclass field of its parent." + ) + assert _has_attr_or_field(leaf_cls, leaf_attr), ( + f"thop kwarg `{thop_kwarg}` reads `{chain}` but `{leaf_attr}` " + f"is not exposed on {leaf_cls.__name__}." + ) + # Coarse type cross-check (skipped when either side is unknown). + cpp_cat = _cpp_category(cpp_types[thop_kwarg]) + leaf_py_type = _dataclass_field_type(leaf_cls, leaf_attr) + if cpp_cat != "unknown" and leaf_py_type is not None: + py_cat = _python_category(leaf_py_type) + if py_cat != "unknown": + assert cpp_cat == py_cat, ( + f"thop kwarg `{thop_kwarg}` (C++ {cpp_types[thop_kwarg]!r}" + f" → {cpp_cat}) bound to `{chain}` " + f"({leaf_py_type!r} → {py_cat})." + ) + + +def test_literal_kwargs_match_allowlist(): + """Every literal-constant kwarg at the call site must appear in + ``_THOP_LITERALS`` with the matching value, and every entry in + ``_THOP_LITERALS`` must be used at the call site.""" + _, literal_kwargs, _ = _classify_kwargs() + unknown = set(literal_kwargs) - set(_THOP_LITERALS) + assert not unknown, ( + f"kwargs passed as literals but not in _THOP_LITERALS: " + f"{sorted(unknown)}. Source from one of {set(_SOURCE_CLASSES)} or " + f"add to the allowlist." + ) + stale = set(_THOP_LITERALS) - set(literal_kwargs) + assert not stale, ( + f"_THOP_LITERALS entries no longer passed as literals: " + f"{sorted(stale)}. Drop them or restore the call-site literal." + ) + for kwarg, value in literal_kwargs.items(): + assert value == _THOP_LITERALS[kwarg], ( + f"thop kwarg `{kwarg}` passed as literal {value!r} but " + f"_THOP_LITERALS expects {_THOP_LITERALS[kwarg]!r}." + ) + + +def _self_attrs_in_property(prop: property) -> set[str]: + """``self.`` names read inside ``prop``'s getter body.""" + src = textwrap.dedent(inspect.getsource(prop.fget)) + return { + node.attr + for node in ast.walk(ast.parse(src)) + if isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "self" + } + + +def _collect_chains(root: str) -> set[tuple[str, ...]]: + """All attribute paths in ``_run`` that start with ``Name(root).``.""" + src = textwrap.dedent(inspect.getsource(TrtllmAttention._run)) + chains: set[tuple[str, ...]] = set() + for node in ast.walk(ast.parse(src)): + if not isinstance(node, ast.Attribute): + continue + path: list[str] = [] + cur: ast.AST = node + while isinstance(cur, ast.Attribute): + path.insert(0, cur.attr) + cur = cur.value + if isinstance(cur, ast.Name) and cur.id == root: + chains.add(tuple(path)) + return chains + + +def _verify_consumed(cls, chains: set[tuple[str, ...]], excluded=frozenset()): + """Recursively assert that every field on ``cls`` is consumed by some + chain in ``chains`` (or excluded). For nested dataclass fields, recurse + into the sub-bag with chain tails. Properties on ``cls`` accessed at + the call site transitively consume the fields they read on ``self``. + """ + direct = {p[0] for p in chains if p} + transitive: set[str] = set() + for name in direct: + obj = vars(cls).get(name) + if isinstance(obj, property): + transitive |= _self_attrs_in_property(obj) + consumed = direct | transitive + for f in fields(cls): + if f.name in excluded: + continue + ftype = f.type if not isinstance(f.type, str) else None + if ftype is not None and dataclasses.is_dataclass(ftype): + sub = {p[1:] for p in chains if len(p) >= 2 and p[0] == f.name} + assert sub, ( + f"Nested dataclass field `{f.name}` on {cls.__name__} is " + f"declared but `{f.name}.` is never read at the " + f"call site." + ) + _verify_consumed(ftype, sub) + else: + assert f.name in consumed, ( + f"Field `{f.name}` on {cls.__name__} not consumed by the " + f"thop call site (directly or via @property). Source it at " + f"the call site or add it to _THOP_EXCLUDED_FIELDS." + ) + + +def test_every_forward_args_field_is_consumed(): + """Recursively check that every dataclass field reachable from + ``AttentionForwardArgs`` (including nested sub-bags such as + ``AttentionSparseArgs``) is consumed at the call site, transitively + via @property where applicable, or listed in ``_THOP_EXCLUDED_FIELDS``. + """ + _verify_consumed( + AttentionForwardArgs, + _collect_chains("forward_args"), + excluded=_THOP_EXCLUDED_FIELDS, + ) + + +def test_no_unexpected_other_kwargs(): + """The only call-site kwargs that aren't ``source.attr`` chains or + allowlisted literals are the ``_run`` positional parameters.""" + _, _, other_kwargs = _classify_kwargs() + expected = {"q", "k", "v"} + unexpected = other_kwargs - expected + assert not unexpected, ( + f"thop kwargs with unexpected expression form: {sorted(unexpected)}. " + f"Allowed: source.attr[.attr...], int(source.attr), literal " + f"constant, or one of {sorted(expected)}." + ) From 3e3e5950e680b2984ed7e939f9234d925e4b33a2 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 25 May 2026 03:49:05 +0000 Subject: [PATCH 02/12] [Rebase] re-apply derived-attn-props intent atop main MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step 2 of the squash + 2-step rebase workflow. Step 1 used ``git rebase -X ours`` so the tree silently took main's lines wherever it could; this commit re-applies the parts of the derived-attn-props refactor that step 1 dropped on the floor. Two displaced regions: 1. ``tensorrt_llm/_torch/attention_backend/trtllm.py`` — the ``thop.attention(...)`` call inside ``TrtllmAttention._run`` was left as a syntactically-valid hybrid: main's positional args followed by my kwarg form, with the positional args referencing locals (output, workspace, layer_idx, ...) that the refactor moved into ``forward_args.X`` properties. Restore the kwarg-only form, drop the dead prelude that built the now-unused locals, and remove ``sink_token_length`` from both the call site and ``_THOP_LITERALS`` (main dropped it from the C++ signature in #14275). 2. ``cpp/tensorrt_llm/nanobind/thop/bindings.cpp`` — the ``nb::arg("layer_idx")`` for ``thop.attention`` reverted to main's spelling even though the C++ function signature in attentionOp.h/.cpp already renamed the parameter to ``local_layer_idx``. Restore the ``nb::arg("local_layer_idx")`` so the binding name matches the underlying parameter and the Python kwarg ``local_layer_idx=self.local_layer_idx`` resolves. Verification: - pre-commit clean on all 10 touched files (autoflake + clang-format ran) - ``test_attention_op_sync.py``: 5/5 pass (call-site kwargs match the ``nb::arg`` list and ``_THOP_LITERALS`` allowlist) - All 6 touched Python modules import cleanly Reviewability: ``git diff rebase-step1..HEAD`` shows every line of manual judgement in this rebase, isolated from main's churn. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/thop/bindings.cpp | 15 +- .../_torch/attention_backend/trtllm.py | 130 +++++------------- 2 files changed, 40 insertions(+), 105 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index f19a8799ea8f..76c95037e945 100644 --- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp @@ -140,13 +140,14 @@ void initBindings(nb::module_& m) nb::arg("kv_scale_quant_orig").none(), nb::arg("out_scale").none(), nb::arg("rotary_inv_freq").none(), nb::arg("rotary_cos_sin").none(), nb::arg("latent_cache").none(), nb::arg("q_pe").none(), nb::arg("block_ids_per_seq").none(), nb::arg("attention_sinks").none(), nb::arg("is_fused_qkv"), - nb::arg("update_kv_cache"), nb::arg("predicted_tokens_per_seq"), nb::arg("layer_idx"), nb::arg("num_heads"), - nb::arg("num_kv_heads"), nb::arg("head_size"), nb::arg("tokens_per_block").none(), nb::arg("max_num_requests"), - nb::arg("max_context_length"), nb::arg("attention_window_size"), nb::arg("beam_width"), nb::arg("mask_type"), - nb::arg("quant_mode"), nb::arg("q_scaling"), nb::arg("position_embedding_type"), - nb::arg("rotary_embedding_dim"), nb::arg("rotary_embedding_base"), nb::arg("rotary_embedding_scale_type"), - nb::arg("rotary_embedding_scales"), nb::arg("rotary_embedding_max_position_info"), - nb::arg("use_paged_context_fmha"), nb::arg("attention_input_type").none(), nb::arg("is_mla_enable"), + nb::arg("update_kv_cache"), nb::arg("predicted_tokens_per_seq"), nb::arg("local_layer_idx"), + nb::arg("num_heads"), nb::arg("num_kv_heads"), nb::arg("head_size"), nb::arg("tokens_per_block").none(), + nb::arg("max_num_requests"), nb::arg("max_context_length"), nb::arg("attention_window_size"), + nb::arg("beam_width"), nb::arg("mask_type"), nb::arg("quant_mode"), nb::arg("q_scaling"), + nb::arg("position_embedding_type"), nb::arg("rotary_embedding_dim"), nb::arg("rotary_embedding_base"), + nb::arg("rotary_embedding_scale_type"), nb::arg("rotary_embedding_scales"), + nb::arg("rotary_embedding_max_position_info"), nb::arg("use_paged_context_fmha"), + nb::arg("attention_input_type").none(), nb::arg("is_mla_enable"), nb::arg("chunked_prefill_buffer_batch_size").none(), nb::arg("q_lora_rank").none(), nb::arg("kv_lora_rank").none(), nb::arg("qk_nope_head_dim").none(), nb::arg("qk_rope_head_dim").none(), nb::arg("v_head_dim").none(), nb::arg("rope_append").none(), nb::arg("mrope_rotary_cos_sin").none(), diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 9a4eb2deff6d..59e4e14882e6 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -44,7 +44,6 @@ # rich object owns them). Sync test enforces both the kwarg name and the # literal value. _THOP_LITERALS: dict = { - "sink_token_length": 0, "sparse_mla_topk_lens": None, "compressed_kv_cache_pool_ptr": None, } @@ -1424,39 +1423,10 @@ def _run( self._ensure_rope_table_size(metadata.max_seq_len) - rotary_embedding_dim = self.rope_params.dim - rotary_embedding_base = self.rope_params.theta - rotary_embedding_scale_type = int(self.rope_params.scale_type) - rotary_embedding_scales = [ - self.rope_params.scale, self.rope_params.short_m_scale, - self.rope_params.long_m_scale - ] - rotary_embedding_max_position_info = [ - self.rope_params.max_positions, - self.rope_params.original_max_positions - ] - spec_decoding_bool_params = [ - metadata.is_spec_decoding_enabled, metadata.use_spec_decoding, - metadata.is_spec_dec_tree - ] - spec_decoding_tensor_params = [ - metadata.spec_decoding_generation_lengths, - metadata.spec_decoding_position_offsets_cpp, - metadata.spec_decoding_packed_mask - ] - if self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()): - spec_decoding_tensor_params.append( - metadata.spec_decoding_bl_tree_mask_offset) - spec_decoding_tensor_params.append( - metadata.spec_decoding_bl_tree_mask) - spec_decoding_tensor_params.append( - metadata.spec_bl_tree_first_sparse_mask_offset_kv) - helix_tensor_params = [ - metadata.helix_position_offsets, metadata.helix_is_inactive_rank - ] - - layer_idx = self.get_local_layer_idx(metadata) - if metadata.spec_decoding_bl_tree_mask is not None and layer_idx == 0: + # Refresh every call: the cache-manager mapping can change between + # forward passes. + self.local_layer_idx = self.get_local_layer_idx(metadata) + if metadata.spec_decoding_bl_tree_mask is not None and self.local_layer_idx == 0: metadata.spec_decoding_bl_tree_mask.zero_() if self.print_skip_softmax_stat: @@ -1496,68 +1466,33 @@ def _run( # inputs) or a literal allowlisted in ``_THOP_LITERALS``. # ``test_attention_op_sync.py`` enforces this statically. thop.attention( - q, - k, - v, - output, - output_sf, - workspace, - metadata.kv_lens_cuda_runtime, - metadata.kv_lens_runtime, - metadata.host_total_kv_lens, - metadata.prompt_lens_cuda_runtime, - metadata.prompt_lens_cpu_runtime, - metadata.host_request_types_runtime, - metadata.kv_cache_block_offsets, - metadata.host_kv_cache_pool_pointers, - metadata.host_kv_cache_pool_mapping, - metadata.cache_indirection, - kv_scale_orig_quant, - kv_scale_quant_orig, - out_scale, - self.rotary_inv_freq, - self.rotary_cos_sin, - forward_args.latent_cache, - forward_args.q_pe, - metadata.block_ids_per_seq, - forward_args.attention_sinks, - is_fused_qkv, - update_kv_cache, - self.predicted_tokens_per_seq, - layer_idx, - self.num_heads, - self.num_kv_heads, - self.head_dim, - metadata.tokens_per_block, - metadata.max_num_requests, - max_context_length, - attention_window_size, - metadata.beam_width, - int(mask_type), - self.quant_mode, - self.q_scaling, - self.position_embedding_type, - rotary_embedding_dim, - rotary_embedding_base, - rotary_embedding_scale_type, - rotary_embedding_scales, - rotary_embedding_max_position_info, - use_paged_context_fmha, - int(attention_input_type), - self.is_mla_enable, - forward_args.chunked_prefill_buffer_batch_size, - self.q_lora_rank, - self.kv_lora_rank, - self.qk_nope_head_dim, - self.qk_rope_head_dim, - self.v_head_dim, - self.rope_append, - mrope_rotary_cos_sin, - mrope_position_deltas, - helix_tensor_params, - self.attention_chunk_size, - forward_args.softmax_stats_tensor, - spec_decoding_bool_params, + q=q, + k=k, + v=v, + output=forward_args.output, + output_sf=forward_args.output_sf, + workspace_=metadata.effective_workspace, + + # --- Per-step batch state (TrtllmAttentionMetadata) --- + sequence_length=metadata.kv_lens_cuda_runtime, + host_past_key_value_lengths=metadata.kv_lens_runtime, + host_total_kv_lens=metadata.host_total_kv_lens, + context_lengths=metadata.prompt_lens_cuda_runtime, + host_context_lengths=metadata.prompt_lens_cpu_runtime, + host_request_types=metadata.host_request_types_runtime, + kv_cache_block_offsets=metadata.kv_cache_block_offsets, + host_kv_cache_pool_pointers=metadata. + host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=metadata.host_kv_cache_pool_mapping, + cache_indirection=metadata.cache_indirection, + block_ids_per_seq=metadata.block_ids_per_seq, + tokens_per_block=metadata.tokens_per_block, + max_num_requests=metadata.max_num_requests, + beam_width=metadata.beam_width, + use_paged_context_fmha=metadata.use_paged_context_fmha, + helix_tensor_params=metadata.helix_tensor_params, + spec_decoding_bool_params=metadata.spec_decoding_bool_params, + spec_decoding_tensor_params=metadata. spec_decoding_tensor_params, num_sparse_topk=metadata.num_sparse_topk, flash_mla_tile_scheduler_metadata=metadata. @@ -1637,10 +1572,9 @@ def _run( sparse_attn_indices_block_size=forward_args.sparse. sparse_attn_indices_block_size, - # --- Literals intentionally None / 0 (see _THOP_LITERALS) --- + # --- Literals intentionally None (see _THOP_LITERALS) --- # ``sparse_mla_topk_lens`` and ``compressed_kv_cache_pool_ptr`` # stay as literal ``None`` until DeepSeek V4 sparse-MLA lands. - sink_token_length=0, sparse_mla_topk_lens=None, compressed_kv_cache_pool_ptr=None, ) From bb17fa91a05da9ccd532fbe2243643dcb5c3fc5b Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 25 May 2026 06:48:55 +0000 Subject: [PATCH 03/12] [None][fix] Drop unused kv_scale_* param from each MLA preprocess op MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI build #39698 failed in nanobind.stubgen with a schema mismatch on ``trtllm::load_paged_kv_cache_for_mla`` (16 vs 17 args). Root cause: main commit 58311b2345a ("Remove unused params in attn") added a ``kv_scale_orig_quant`` parameter to three MLA preprocess C++ functions without a corresponding schema update — and ironic to its title, the added parameter is dead in two of the three functions. Inspection of each function body shows only ONE direction of scale is actually consumed: | Function | Used | Dead | |-----------------------------------|---------------------|---------------------| | loadPagedKVCacheForMLA | kv_scale_quant_orig | kv_scale_orig_quant | | loadChunkedKVCacheForMLA | kv_scale_quant_orig | kv_scale_orig_quant | | MLARopeAppendPagedKVAssignQ | kv_scale_orig_quant | kv_scale_quant_orig | Fix: drop the dead direction from each C++ function signature. The ``TORCH_LIBRARY`` schemas and Python callers (trtllm.py, sparse/dsa.py, auto_deploy/.../trtllm_mla.py) already pass only the used direction, so no changes are needed on the Python side. This diverges from main's signatures but drops dead code rather than preserving it. Build-stage failure: validated by source inspection; schema and function now textually match. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp b/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp index 69cceb64de65..a52c86a46917 100644 --- a/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp +++ b/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp @@ -107,10 +107,9 @@ void mergeChunkedAttentionForMLAHelper(torch::Tensor& merged_attn, torch::Tensor std::vector loadPagedKVCacheForMLA(torch::ScalarType out_dtype, int64_t const num_contexts, int64_t const num_ctx_cached_tokens, int64_t const max_ctx_cached_kv_len, torch::Tensor& cu_ctx_cached_kv_lens, torch::Tensor const& kv_cache_block_offsets, torch::Tensor const& host_kv_cache_pool_pointers, - torch::Tensor const& host_kv_cache_pool_mapping, torch::optional kv_scale_orig_quant, - torch::optional kv_scale_quant_orig, int64_t const layer_idx, int64_t const lora_size, - int64_t const rope_size, int64_t const tokens_per_block, int64_t const attention_window_size, - int64_t const beam_width, int64_t const quant_mode) + torch::Tensor const& host_kv_cache_pool_mapping, torch::optional kv_scale_quant_orig, + int64_t const layer_idx, int64_t const lora_size, int64_t const rope_size, int64_t const tokens_per_block, + int64_t const attention_window_size, int64_t const beam_width, int64_t const quant_mode) { TORCH_CHECK(out_dtype == torch::kFloat16 || out_dtype == torch::kFloat32 || out_dtype == torch::kBFloat16, "out_dtype only support float16, float32, bfloat16"); @@ -196,9 +195,9 @@ std::vector loadChunkedKVCacheForMLA(torch::ScalarType out_dtype, int64_t const num_ctx_cached_tokens, torch::Tensor const& cu_ctx_chunked_kv_lens, torch::Tensor const& chunked_ld_global_offset, torch::Tensor const& kv_cache_block_offsets, torch::Tensor const& host_kv_cache_pool_pointers, torch::Tensor const& host_kv_cache_pool_mapping, - torch::optional kv_scale_orig_quant, torch::optional kv_scale_quant_orig, - int64_t const layer_idx, int64_t const lora_size, int64_t const rope_size, int64_t const tokens_per_block, - int64_t const max_seq_len, int64_t const attention_window_size, int64_t const beam_width, int64_t const quant_mode) + torch::optional kv_scale_quant_orig, int64_t const layer_idx, int64_t const lora_size, + int64_t const rope_size, int64_t const tokens_per_block, int64_t const max_seq_len, + int64_t const attention_window_size, int64_t const beam_width, int64_t const quant_mode) { TORCH_CHECK(out_dtype == torch::kFloat16 || out_dtype == torch::kFloat32 || out_dtype == torch::kBFloat16, "out_dtype only support float16, float32, bfloat16"); @@ -284,8 +283,8 @@ void MLARopeAppendPagedKVAssignQ(torch::Tensor& q, torch::Tensor& latent_cache, int64_t const nope_size, int64_t const rope_size, int64_t const lora_size, torch::Tensor const& kv_cache_block_offsets, torch::Tensor const& host_kv_cache_pool_pointers, torch::Tensor const& host_kv_cache_pool_mapping, torch::optional kv_scale_orig_quant, - torch::optional kv_scale_quant_orig, int64_t const layer_idx, int64_t const tokens_per_block, - int64_t const attention_window_size, int64_t const beam_width, int64_t const quant_mode) + int64_t const layer_idx, int64_t const tokens_per_block, int64_t const attention_window_size, + int64_t const beam_width, int64_t const quant_mode) { auto input_dtype = q.scalar_type(); TORCH_CHECK(input_dtype == torch::kFloat16 || input_dtype == torch::kFloat32 || input_dtype == torch::kBFloat16); From c701220b2e7d98158f0d951a6c69b0bed4dfd6ff Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 25 May 2026 08:58:22 +0000 Subject: [PATCH 04/12] [None][refactor] Lazy-memoize TrtllmAttention.local_layer_idx ``get_local_layer_idx`` previously re-read ``metadata.kv_cache_manager. layer_offsets[self.layer_idx]`` on every call. The dict and its contents are immutable after the cache manager's ``__init__`` (``resource_manager.py:562``), so the value can be cached on the attention instance after the first call that has a cache manager. Drop the per-call write in ``_run`` and let ``get_local_layer_idx`` memoize internally. The 5 other call sites already go through ``get_local_layer_idx`` and now hit the cached path after the first forward. The ``kv_cache_manager is None`` branch stays uncached so a manager attached on a later forward is picked up correctly. No behavior change; ``test_attention_op_sync.py`` (5 tests) still passes. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/trtllm.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 59e4e14882e6..74a518e3f03c 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1168,10 +1168,14 @@ def update_quant_config(self, new_quant_config: Optional[QuantConfig]): ) def get_local_layer_idx(self, metadata: TrtllmAttentionMetadata) -> int: + if self.local_layer_idx is not None: + return self.local_layer_idx if metadata.kv_cache_manager is None: + # Uncached: recomputed each call until a cache manager appears. return self.layer_idx - else: - return metadata.kv_cache_manager.layer_offsets[self.layer_idx] + self.local_layer_idx = metadata.kv_cache_manager.layer_offsets[ + self.layer_idx] + return self.local_layer_idx def use_nvfp4_output( self, @@ -1423,10 +1427,8 @@ def _run( self._ensure_rope_table_size(metadata.max_seq_len) - # Refresh every call: the cache-manager mapping can change between - # forward passes. - self.local_layer_idx = self.get_local_layer_idx(metadata) - if metadata.spec_decoding_bl_tree_mask is not None and self.local_layer_idx == 0: + if metadata.spec_decoding_bl_tree_mask is not None and self.get_local_layer_idx( + metadata) == 0: metadata.spec_decoding_bl_tree_mask.zero_() if self.print_skip_softmax_stat: From 6dc82efdeba4a0d33d4304392ddffbb9912f7f34 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Tue, 26 May 2026 02:22:11 +0000 Subject: [PATCH 05/12] [None][fix] Prime self.local_layer_idx before thop.attention call CI build #39731 surfaced 896 ``TypeError: attention(): incompatible function arguments`` (plus ~200 downstream illegal-memory-access / setup failures). The thop.attention kwarg ``local_layer_idx=self.local_layer_idx`` reads the attribute directly; my previous lazy-memoize commit (c701220b2e) removed the unconditional priming write at the top of ``_run``, so when ``spec_decoding_bl_tree_mask`` is None (the common case) ``get_local_layer_idx`` was never called before thop.attention and the attribute stayed at its ``None`` placeholder. C++ binding then rejects ``NoneType`` against ``local_layer_idx: int``. Restore the priming line. The internal lazy memoize in ``get_local_layer_idx`` still elides the per-call ``layer_offsets`` dict lookup after the first forward; this priming call hits the fast-return path on every subsequent call. 5/5 ``test_attention_op_sync.py`` still pass. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/trtllm.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 74a518e3f03c..d02ded49164b 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1427,8 +1427,11 @@ def _run( self._ensure_rope_table_size(metadata.max_seq_len) - if metadata.spec_decoding_bl_tree_mask is not None and self.get_local_layer_idx( - metadata) == 0: + # Prime ``self.local_layer_idx`` so the ``thop.attention`` kwarg + # below reads a populated int rather than the ``None`` placeholder. + # The call is a fast cache hit after the first forward. + self.local_layer_idx = self.get_local_layer_idx(metadata) + if metadata.spec_decoding_bl_tree_mask is not None and self.local_layer_idx == 0: metadata.spec_decoding_bl_tree_mask.zero_() if self.print_skip_softmax_stat: From 319ec1c7241c3363da74883d222920628caae7f6 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Tue, 26 May 2026 13:44:37 +0000 Subject: [PATCH 06/12] [None][fix] Null-safe scale-pointer derefs in MMHA + unfused KV-cache write kernels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI build #39787 on PR #14279 surfaced 162 illegal-memory-access failures (1 root + 161 cascade from a corrupted CUDA context) in ``test_attention_mla[v2_kv_cache-...float8_e4m3fn-context_sequence_lengths=[10,12,5]...]`` on RTX5090 (sm_120) and RTXPro6000D. Reproduced locally on RTX5090. The crash is a WARP_MMU_FAULT in ``kernel_mha`` during the generation phase. Root cause: ``decoderMaskedMultiheadAttentionTemplate.h`` dereferences ``params.kv_scale_quant_orig[0]`` and ``params.kv_scale_orig_quant[0]`` gated only on the compile-time template flags ``ENABLE_8BITS_K_CACHE`` / ``ENABLE_8BITS_KV_CACHE``, NOT on a runtime nullptr check. When the caller passes ``None`` for the scales (e.g. unit tests, non-fp8 paths exercising the same template), the pointer is nullptr and the kernel crashes. This mirrors the FlashMLA ``descale_q_ptr`` / ``descale_k_ptr`` nullptr-safety fix already in this branch — same root cause, sister kernel that the FlashMLA fix didn't cover. Apply the same defensive default-to-1.0 pattern in: - ``decoderMaskedMultiheadAttentionTemplate.h`` (the ``kernel_mha`` template: k_scale_quant_orig, kv_scale_quant_orig, kv_scale_orig_quant) - ``unfusedAttentionKernels_2_template.h`` line ~1034 (qkv_scale_orig_quant in the 8-bit cache write path — the only outlier in that file; peer reads at lines 579/989 already guarded the same way) Validated on RTX5090 (sm_120) inside the tekit-agent container after a fresh sm_120 build: the previously-crashing ``v2_kv_cache + float8_e4m3fn + [10,12,5]`` scenario passes. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../decoderMaskedMultiheadAttentionTemplate.h | 15 +++++++++++---- .../unfusedAttentionKernels_2_template.h | 8 ++++++-- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h index 5bb632465dd9..cde3a765ce82 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h @@ -1522,13 +1522,20 @@ __global__ void __launch_bounds__(MAX_THEADS_PER_BLOCK, MIN_BLOCKS_PER_SM) maske bool const load_qkv_quant = params.qkv_scale_quant_orig != nullptr; bool const write_attention_quant = params.attention_out_scale_orig_quant != nullptr; - // Quant/Dequant scales for 8bits kv cache. + // Quant/Dequant scales for 8bits kv cache. The ENABLE_8BITS_* guards are + // compile-time template flags; the runtime pointer can still be nullptr + // when the caller omits scales (e.g. unit tests, non-fp8 paths exercising + // the same template). Default to 1.0 in that case to mirror the FlashMLA + // descale_*_ptr nullptr-safety fix. using T_scale = typename kv_cache_scale_type_t::Type; T_scale kv_scale_orig_quant, k_scale_quant_orig; - float const k_scale_quant_orig_f = (ENABLE_8BITS_K_CACHE ? params.kv_scale_quant_orig[0] : 1.0f); - float const kv_scale_quant_orig_f = (ENABLE_8BITS_KV_CACHE ? params.kv_scale_quant_orig[0] : 1.0f); + float const k_scale_quant_orig_f + = (ENABLE_8BITS_K_CACHE && params.kv_scale_quant_orig != nullptr ? params.kv_scale_quant_orig[0] : 1.0f); + float const kv_scale_quant_orig_f + = (ENABLE_8BITS_KV_CACHE && params.kv_scale_quant_orig != nullptr ? params.kv_scale_quant_orig[0] : 1.0f); convert_from_float(&k_scale_quant_orig, k_scale_quant_orig_f); - convert_from_float(&kv_scale_orig_quant, (ENABLE_8BITS_KV_CACHE ? params.kv_scale_orig_quant[0] : 1.0f)); + convert_from_float(&kv_scale_orig_quant, + (ENABLE_8BITS_KV_CACHE && params.kv_scale_orig_quant != nullptr ? params.kv_scale_orig_quant[0] : 1.0f)); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index 78fad6b6a0fd..f8993fed6146 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -1028,10 +1028,14 @@ __global__ void applyBiasRopeUpdateKVCacheV2(QKVPreprocessingParams::Type; TScale scaleOrigQuant; - mmha::convert_from_float(&scaleOrigQuant, params.qkv_scale_orig_quant[0]); + mmha::convert_from_float(&scaleOrigQuant, + params.qkv_scale_orig_quant != nullptr ? params.qkv_scale_orig_quant[0] : 1.0f); // Store 8bits kv cache. mmha::store_8bits_vec(kDst, k, inBlockIdx, scaleOrigQuant); mmha::store_8bits_vec(vDst, v, inBlockIdx, scaleOrigQuant); From 5040cc7194791cf341b3d186f8a1d72ef32fd999 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Tue, 26 May 2026 14:49:06 +0000 Subject: [PATCH 07/12] [None][fix] Null-safe kvCacheScale deref in MLA sm_120 XQA JIT kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Followup to 319ec1c724. Validation on RTX5090 (sm_120) showed the v2_kv_cache + fp8 + short-context test still crashed with WARP_MMU_FAULT in kernel_mha, this time in the XQA JIT path. On sm_120, MLA generation routes through cpp/kernels/xqa/mla_sm120.cu (compiled at runtime via nvrtc), not the prebuilt MMHA template fixed in 319ec1c724. Two unguarded args.kvCacheScale[0] dereferences in mla_sm120.cu: - line ~539 (Producer ctor): smem.qkScaleLog2e factor - line ~1447 (epilogue): xvScale for output rescale Both gate fp8 path correctness but crash the kernel when the caller passes nullptr — same pattern as the MMHA template fix and the FlashMLA descale_*_ptr fix. Apply the same default-to-1.0-when-nullptr pattern. With this fix, the full test_attention_mla.py suite (84 cases, both v1 and v2 kv_cache, bf16 and fp8 KV cache) passes 80/80 (4 skipped, 0 failed) on RTX5090 — same scenario that produced 162 cascade failures in CI #39787. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- cpp/kernels/xqa/mla_sm120.cu | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/cpp/kernels/xqa/mla_sm120.cu b/cpp/kernels/xqa/mla_sm120.cu index 84ee4f7ccba4..f77e83732fbc 100644 --- a/cpp/kernels/xqa/mla_sm120.cu +++ b/cpp/kernels/xqa/mla_sm120.cu @@ -536,7 +536,12 @@ struct Producer #endif if (threadIdx.x == 0) { - smem.qkScaleLog2e = args.qScale * args.kvCacheScale[0] * log2e; + // ``kvCacheScale`` is the device-memory fp8 dequant scale. Callers + // may pass nullptr when scales are unset (e.g. unit tests, or + // scale=1.0 paths); default to 1.0 in that case to keep the kernel + // crash-safe. + float const kvScale = (args.kvCacheScale != nullptr) ? args.kvCacheScale[0] : 1.0f; + smem.qkScaleLog2e = args.qScale * kvScale * log2e; } if (threadIdx.x < headGrpSize) @@ -1444,7 +1449,8 @@ __device__ inline void Consumer::compute() smem.mathWarpsBar.arrive(); ThrdRegRowMax const accRowSum = loadShmRowMax(smem.accRowSum[tileIdx.x], tileBase.y, lane); - float const xvScale = computeRowSumFromF8 ? args.kvCacheScale[0] : args.kvCacheScale[0] * xScale; + float const kvScale = (args.kvCacheScale != nullptr) ? args.kvCacheScale[0] : 1.0f; + float const xvScale = computeRowSumFromF8 ? kvScale : kvScale * xScale; WarpOutputTile const output = finalize(acc, accRowSum, xvScale, lane); bool const isMultiBlockMode = (nbSubSeq != 1); From 9c76d6f2bb347c12765e9065c12677202bd485be Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 27 May 2026 03:34:04 +0000 Subject: [PATCH 08/12] [None][fix] Update trtllm_gen.py to renamed kv_scale forward_args fields MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After commit 574c518498 renamed AttentionForwardArgs.kv_scales_sf and kv_scales_sf_inv to kv_scale_quant_orig and kv_scale_orig_quant, two references in FlashInferTrtllmGenAttention._get_kv_scale_params were left pointing at the old (now-nonexistent) field names. Because AttentionForwardArgs uses slots=True, hitting these references at runtime raises AttributeError. The path is gated by TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION=1, so default-config tests do not exercise it; the failure surfaces only when the trtllm-gen kernel is enabled (e.g. test_attention_mla with the env override). Static-only fix: pure rename to the current field names. No GPU validation possible without rebuilding the .so on a Blackwell node, but the change cannot regress anything else — these are the only references to the old names in tensorrt_llm/. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/trtllm_gen.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py index 4d09adc046bd..211e7718ec46 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py @@ -594,13 +594,13 @@ def _get_kv_scale_params( kv_cache_quant_mode = QuantMode(quant_mode) kv_scale_orig_quant = ( attention_layer.kv_scale_orig_quant - if forward_args.kv_scales_sf_inv is None - else forward_args.kv_scales_sf_inv + if forward_args.kv_scale_orig_quant is None + else forward_args.kv_scale_orig_quant ) kv_scale_quant_orig = ( attention_layer.kv_scale_quant_orig - if forward_args.kv_scales_sf is None - else forward_args.kv_scales_sf + if forward_args.kv_scale_quant_orig is None + else forward_args.kv_scale_quant_orig ) if ( not kv_cache_quant_mode.has_kv_cache_quant() From 9897b55da754137d2a770aa779c89e5c61c10547 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 27 May 2026 04:24:34 +0000 Subject: [PATCH 09/12] [None][fix] Defer out_scale/out_scale_sf selection to backend _run CI build #39880 surfaced NVFP4 accuracy regressions on DGX_B200: - TestLlama3_1_8B.test_nvfp4: 5.517 vs threshold 23.479 - TestLlama3_1_8B.test_nvfp4_streaming: same - TestQwen3_30B_A3B.test_nvfp4[latency_moe_*]: same root cause Root cause: refactor commit 574c518498 consolidated main's out_scale/out_scale_sf pair into a single forward_args.out_scale field, picking at the modules/attention.py call site based on ``output_sf is not None``. But ``output_sf`` at that point is the ``_attn_impl`` function argument, which forward_impl() never populates (the buffer is created later inside the backend's forward by ``create_output``). So the predicate was always False and the NVFP4 path silently picked ``inv_input_scale`` instead of ``input_scale`` -> kernel-side rescale wrong -> ~5x accuracy drop. Main avoided this by passing both scales through to trtllm.py and deciding inside ``_run`` after ``output_sf`` is allocated. Fix: * interface.py: re-introduce ``out_scale_sf`` field on AttentionForwardArgs. * modules/attention.py: populate both ``out_scale`` (= ``inv_input_scale``) and ``out_scale_sf`` (= ``input_scale``) whenever quantized output is enabled, matching main. * trtllm.py: in ``_run``, after the fp8-KV fallback, promote ``out_scale_sf`` -> ``out_scale`` when ``forward_args.output_sf`` is allocated. The static sync test's ``_THOP_EXCLUDED_FIELDS`` gains ``out_scale_sf`` since the kernel binding still takes only a single ``out_scale``. Validation on B200 (umbriel-b200-081, sm_100): * TestLlama3_1_8B::test_nvfp4 PASSED (MMLU 62.622 vs threshold 61.343) * TestLlama3_1_8B::test_nvfp4_streaming[stream_interval_4,64] PASSED * test_attention_op_sync.py (5 tests) PASSED Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/interface.py | 1 + tensorrt_llm/_torch/attention_backend/trtllm.py | 10 ++++++++++ tensorrt_llm/_torch/modules/attention.py | 14 +++++++++----- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 2921672d3f81..af4fd26ee0ca 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -716,6 +716,7 @@ class AttentionForwardArgs: output_sf: Optional[torch.Tensor] = None out_scale: Optional[torch.Tensor] = None + out_scale_sf: Optional[torch.Tensor] = None kv_scale_orig_quant: Optional[torch.Tensor] = None kv_scale_quant_orig: Optional[torch.Tensor] = None diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index d02ded49164b..8468f964d6f1 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -38,6 +38,7 @@ "topk_indices", # DSA-only "is_generation", # trtllm_gen dispatch flag "attention_mask_data", # custom-mask code path + "out_scale_sf", # promoted into ``out_scale`` in ``_run`` for NVFP4 path }) # ``thop.attention`` kwargs hard-wired to a literal at the call site (no @@ -1440,6 +1441,15 @@ def _run( if forward_args.attention_window_size is None: forward_args.attention_window_size = metadata.max_seq_len + # Promote ``out_scale_sf`` -> ``out_scale`` for the NVFP4-output path + # (kernel reads a single ``out_scale`` and interprets it as the SF + # quant scale when ``output_sf`` is allocated). ``output_sf`` is + # populated by ``create_output`` in ``forward`` above, so the + # decision is correct only here, not at the modules/attention.py + # call site where ``output_sf`` is always ``None``. + if forward_args.output_sf is not None and forward_args.out_scale_sf is not None: + forward_args.out_scale = forward_args.out_scale_sf + helix_active = metadata.helix_position_offsets is not None use_sage_attn = (forward_args.sage_attn_num_elts_per_blk_q > 0 or forward_args.sage_attn_num_elts_per_blk_k > 0 diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index f7a9c8991f75..616724550497 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -745,13 +745,16 @@ def _attn_impl( # FP8/FP4 output and keeps attention output in BF16 for better # precision when applying pre_quant_scale. Also don't set out_scale # if LoRA is active — LoRA grouped_gemm doesn't support FP8. - # When ``output_sf`` is set the kernel writes NVFP4 output and reads - # ``input_scale`` as its scaling factor; otherwise it reads - # ``inv_input_scale``. + # Pass both scales; the backend selects ``out_scale_sf`` when the + # kernel writes NVFP4 output (``forward_args.output_sf`` is + # allocated downstream by ``create_output``) and ``out_scale`` + # otherwise. Deciding here would be premature — ``output_sf`` is + # not populated yet at this call site. out_scale = None + out_scale_sf = None if self._use_quantize_output() and not has_lora: - out_scale = (self.o_proj.input_scale if output_sf is not None else - self.o_proj.inv_input_scale) + out_scale = self.o_proj.inv_input_scale + out_scale_sf = self.o_proj.input_scale kv_scale_orig_quant = None kv_scale_quant_orig = None @@ -767,6 +770,7 @@ def _attn_impl( attn_metadata, forward_args=AttentionForwardArgs( out_scale=out_scale, + out_scale_sf=out_scale_sf, kv_scale_orig_quant=kv_scale_orig_quant, kv_scale_quant_orig=kv_scale_quant_orig, attention_mask=attention_mask, From 5e5334ffc902e95c81b4092cf456e3d57b709533 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 27 May 2026 06:43:48 +0000 Subject: [PATCH 10/12] [None][refactor] Source kv_scale from forward_args in trtllm_gen.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After the AttentionForwardArgs field rename in 574c518498, the trtllm-gen backend's _get_kv_scale_params kept a two-tier lookup (per-call forward_args.kv_scale_* OR layer-level attention_layer.kv_scale_*). The layer-level fallback is dead for the PyTorch backend: * The PyTorch model loader never writes to TrtllmAttention. kv_cache_scaling_factor — it stays at the torch.ones(1) initial value, so attention_layer.kv_scale_quant_orig is always 1.0. * The attention-kernel side (mlaPreprocessOp / MMHA / unfused / sm_120 XQA) now defaults to 1.0 when the scale pointer is null (commits 319ec1c724, 5040cc7194), reproducing the exact same effective value. Sourcing only from forward_args also removes the implicit dependency on attention_layer's kv_scale_* attributes, allowing the layer-level mirrors to be dropped from TrtllmAttention.__init__ in a follow-up. Validation on B200 (umb-b200-236): * TestLlama3_1_8B::test_nvfp4 with TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION=1 PASSED (exercises the simplified _get_kv_scale_params). Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/trtllm_gen.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py index 211e7718ec46..16ce98ef876a 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm_gen.py @@ -590,18 +590,9 @@ def _get_kv_scale_params( forward_args: AttentionForwardArgs, quant_mode: int, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - attention_layer = self._get_attention_layer() kv_cache_quant_mode = QuantMode(quant_mode) - kv_scale_orig_quant = ( - attention_layer.kv_scale_orig_quant - if forward_args.kv_scale_orig_quant is None - else forward_args.kv_scale_orig_quant - ) - kv_scale_quant_orig = ( - attention_layer.kv_scale_quant_orig - if forward_args.kv_scale_quant_orig is None - else forward_args.kv_scale_quant_orig - ) + kv_scale_orig_quant = forward_args.kv_scale_orig_quant + kv_scale_quant_orig = forward_args.kv_scale_quant_orig if ( not kv_cache_quant_mode.has_kv_cache_quant() or kv_scale_orig_quant is None From 15b0fdce4b8f976b0d55d14200cbe279e5c11d1c Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 27 May 2026 07:02:31 +0000 Subject: [PATCH 11/12] [None][test] Catch stale _THOP_EXCLUDED_FIELDS entries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add ``test_excluded_fields_match_real_fields`` to the AST sync test. It walks ``AttentionForwardArgs`` (recursively into nested dataclass sub-bags) and asserts every entry in ``_THOP_EXCLUDED_FIELDS`` names a real field. A dead entry would silently allow a newly-added real field to slip past ``test_every_forward_args_field_is_consumed``. While at it, drop the dead ``is_generation`` entry — it was added in 574c518498 but never matched a field on the renamed ``AttentionForwardArgs`` (the matching field was removed earlier in the same refactor iteration). The companion sanity check on ``_THOP_LITERALS`` already exists in ``test_literal_kwargs_match_allowlist`` (both directions: unknowns and stales), so this commit covers only the gap on the exclude set. Validation: ran the full sync test locally (no GPU); all 6 tests pass, including the new check. Temporarily re-adding ``is_generation`` makes the new check fail with a clear message, as expected. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention_backend/trtllm.py | 1 - .../test_attention_op_sync.py | 29 +++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 8468f964d6f1..9768cef00a9d 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -36,7 +36,6 @@ # transitively reads; entries here are exempt. _THOP_EXCLUDED_FIELDS: frozenset = frozenset({ "topk_indices", # DSA-only - "is_generation", # trtllm_gen dispatch flag "attention_mask_data", # custom-mask code path "out_scale_sf", # promoted into ``out_scale`` in ``_run`` for NVFP4 path }) diff --git a/tests/unittest/_torch/attention_backend/test_attention_op_sync.py b/tests/unittest/_torch/attention_backend/test_attention_op_sync.py index 91aee99d4c78..b664eda11ab5 100644 --- a/tests/unittest/_torch/attention_backend/test_attention_op_sync.py +++ b/tests/unittest/_torch/attention_backend/test_attention_op_sync.py @@ -517,3 +517,32 @@ def test_no_unexpected_other_kwargs(): f"Allowed: source.attr[.attr...], int(source.attr), literal " f"constant, or one of {sorted(expected)}." ) + + +def _all_forward_args_field_names() -> set[str]: + """All dataclass field names reachable from ``AttentionForwardArgs``, + recursively descending into nested dataclass sub-bags.""" + seen: set[str] = set() + + def _walk(cls) -> None: + for f in fields(cls): + seen.add(f.name) + ftype = f.type if not isinstance(f.type, str) else None + if ftype is not None and dataclasses.is_dataclass(ftype): + _walk(ftype) + + _walk(AttentionForwardArgs) + return seen + + +def test_excluded_fields_match_real_fields(): + """Every entry in ``_THOP_EXCLUDED_FIELDS`` must name a real field on + ``AttentionForwardArgs`` (or a nested sub-bag). Dead entries (typos, + fields removed in a refactor) silently allow newly-added real fields + to slip past ``test_every_forward_args_field_is_consumed``.""" + stale = set(_THOP_EXCLUDED_FIELDS) - _all_forward_args_field_names() + assert not stale, ( + f"_THOP_EXCLUDED_FIELDS entries that don't match any " + f"AttentionForwardArgs field: {sorted(stale)}. Drop them or " + f"restore the field." + ) From 2d4798ecb5d82e163bab2be2f9c3cd3b30c93d91 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 27 May 2026 13:03:50 +0000 Subject: [PATCH 12/12] [None][fix] Restore layer-level kv_scale_* mirrors and fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit L0 build #40005 surfaced 5 illegal-memory-access failures on RTXPro6000D and 1 CUBLAS_STATUS_EXECUTION_FAILED on DGX_H100: - TestLagunaXS.{test_fp8,test_nvfp4} (RTXPro6000D, sm_120) - test_ptp_quickstart_advanced[{Llama3.1-8B-NVFP4,Qwen3-30B-A3B_{fp8,nvfp4}}_hf] (RTXPro6000D, sm_120) - TestNemotronV3Nano.test_fp8 (DGX_H100, sm_90) Cross-PR scan: 0 other PRs hit these tests; PR #14279 alone since commit 5e5334ffc9 (which dropped the older ecf3fa5d03 "kv_scale fallback" commit during the rebase that simplified trtllm_gen.py). Root cause: dropping ecf3fa5d03 removed two things that turn out to be load-bearing for fp8 KV cache models on the default-path (non-trtllm-gen) attention kernels: 1. Layer-level mirrors ``self.kv_cache_scaling_factor`` / ``self.kv_scale_quant_orig`` / ``self.kv_scale_orig_quant`` on ``TrtllmAttention.__init__``. 2. The ``_run``-side fallback that defaults ``forward_args.kv_scale_*`` to those mirrors when the caller leaves them ``None``. I earlier (correctly) noted the PyTorch backend never writes real values to ``kv_cache_scaling_factor`` and concluded the mirrors were 1.0-no-ops the kernel could re-derive from a nullptr default. That conclusion was incomplete: several XQA kernels still deref ``kvCacheScale[0]`` whenever ``isKVCacheQuantized`` is true and do NOT check for nullptr — * cpp/kernels/xqa/mha.cu (lines 1744, 2495) - sm_100/sm_120 * cpp/kernels/xqa/mha_sm90.cu (lines 885, 1127, 1530, 1620) - sm_90 The earlier null-safe deref fixes (commits 319ec1c724, 5040cc7194) only covered the MMHA template, the unfused KV-cache write path, and ``mla_sm120.cu`` — the non-MLA XQA kernels were missed. Passing a nullptr to them on fp8-KV models reads invalid memory at the kernel entry → illegal memory access at executor init → cascades to CUBLAS_STATUS_EXECUTION_FAILED on subsequent calls. Fix: restore the layer-level mirrors and the ``_run`` fallback. This is verbatim equivalent to the trtllm.py state at the previously-passing build #2656 (commit 2aea9fc525) modulo this commit message's comment edits. ``modules/attention.py`` continues to populate the per-call fields only for fp4 KV cache; fp8-KV models fall through to the layer-level mirrors and the kernel receives a valid 1.0 pointer. Local validation: my single-arch / multi-arch sm_120 builds cannot reproduce the CI's complete FMHA cubin set (e4m3 + sm_120 + head_dim 128 + ReturnSoftmaxStats variants are not generated by setup.py / fmha_cubin.cpp in the local build) — TestLagunaXS::test_fp8 fails at executor init with "FMHA kernels are not found" before reaching the attention runtime path. This is a build-environment gap, not a code regression. The fix is byte-equivalent (modulo comments) to commit 2aea9fc525 which DID pass this test on RTXPro6000D in CI build #2656. A follow-up should add null-safe ``kvCacheScale[0]`` guards in ``mha.cu`` and ``mha_sm90.cu`` so the layer-level mirrors are no longer load-bearing, mirroring the pattern already in ``mla_sm120.cu``. Tracked separately. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention_backend/trtllm.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 9768cef00a9d..aedd4c839d03 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1143,6 +1143,21 @@ def __init__( self.print_skip_softmax_stat = os.environ.get( "TRTLLM_PRINT_SKIP_SOFTMAX_STAT", "0") == "1" + # Layer-level fp8 KV-cache scales. Stay at 1.0 (no-op) for the + # PyTorch backend, which never overrides them. They guarantee the + # kernel always receives a valid pointer, since several non-MLA + # XQA kernels (cpp/kernels/xqa/mha.cu, mha_sm90.cu) deref + # ``kvCacheScale[0]`` whenever ``isKVCacheQuantized`` is true and + # do not check for nullptr. ``modules/attention.py`` only assigns + # ``forward_args.kv_scale_*`` for fp4 KV cache, so without this + # fallback the kernel takes nullptr on fp8-KV models → illegal + # memory access. + self.kv_cache_scaling_factor = torch.ones(1, + dtype=torch.float32, + device='cuda') + self.kv_scale_quant_orig = self.kv_cache_scaling_factor + self.kv_scale_orig_quant = 1.0 / self.kv_cache_scaling_factor + self.local_layer_idx: Optional[int] = None if not skip_create_weights_in_init: self.update_quant_config(self.quant_config) @@ -1449,6 +1464,17 @@ def _run( if forward_args.output_sf is not None and forward_args.out_scale_sf is not None: forward_args.out_scale = forward_args.out_scale_sf + # Default ``forward_args.kv_scale_*`` to the layer-level mirrors when + # the caller didn't populate them. ``modules/attention.py`` only sets + # these for fp4 KV cache; fp8-KV models leave them ``None``. Several + # XQA kernels (mha.cu, mha_sm90.cu) deref ``kvCacheScale[0]`` when + # ``isKVCacheQuantized`` is true and don't check for nullptr, so + # passing ``None`` crashes with illegal memory access. + if forward_args.kv_scale_orig_quant is None: + forward_args.kv_scale_orig_quant = self.kv_scale_orig_quant + if forward_args.kv_scale_quant_orig is None: + forward_args.kv_scale_quant_orig = self.kv_scale_quant_orig + helix_active = metadata.helix_position_offsets is not None use_sage_attn = (forward_args.sage_attn_num_elts_per_blk_q > 0 or forward_args.sage_attn_num_elts_per_blk_k > 0