[None][refactor] Add derived properties for the thop.attention call site

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<warpTile.y>(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);

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<T, Tcache>::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();

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;
     }

cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h

@@ -1028,10 +1028,14 @@ __global__ void applyBiasRopeUpdateKVCacheV2(QKVPreprocessingParams<T, KVCacheBu
                     if constexpr (ENABLE_8BITS_CACHE)
                     {
                         inBlockIdx = inBlockIdx * ELTS_PER_VEC;
-                        // Cast float scale to dst data type.
+                        // Cast float scale to dst data type. Default to 1.0
+                        // when the scale pointer is nullptr; the surrounding
+                        // peer reads in this file already follow this pattern
+                        // (e.g. lines 579, 989).
                         using TScale = typename mmha::kv_cache_scale_type_t<T, TCache>::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);

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(),

cpp/tensorrt_llm/thop/attentionOp.cpp

@@ -927,7 +927,7 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     std::optional<torch::Tensor> rotary_cos_sin, std::optional<torch::Tensor> latent_cache,
     std::optional<torch::Tensor> q_pe, std::optional<torch::Tensor> block_ids_per_seq,
     std::optional<torch::Tensor> 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<int64_t> 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<torch::Tensor> k, std::optional<to
     bool sage_attn_qk_int8, int64_t num_contexts, int64_t num_ctx_tokens,
     std::optional<int64_t> 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<torch::Tensor> k, std::optional<to
     auto op = std::make_shared<AttentionOp>();
     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<torch::Tensor> k, std::optional<to
     static std::unordered_map<CacheKey, std::shared_ptr<AttentionOp>, hash<CacheKey>> 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<torch::Tensor> k, std::optional<to
             flash_mla_tile_scheduler_metadata, flash_mla_num_splits, compressed_kv_cache_pool_ptr);
     }
 
-    TLLM_LOG_TRACE("Attention op stops at layer %d", layer_idx);
+    TLLM_LOG_TRACE("Attention op stops at layer %d", local_layer_idx);
 }
 
 bool attention_supports_nvfp4_output(int64_t const num_heads, int64_t const num_kv_heads, int64_t const head_size,

cpp/tensorrt_llm/thop/attentionOp.h

@@ -55,7 +55,7 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     std::optional<torch::Tensor> rotary_cos_sin, std::optional<torch::Tensor> latent_cache,
     std::optional<torch::Tensor> q_pe, std::optional<torch::Tensor> block_ids_per_seq,
     std::optional<torch::Tensor> 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<int64_t> 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,

cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp

@@ -107,10 +107,9 @@ void mergeChunkedAttentionForMLAHelper(torch::Tensor& merged_attn, torch::Tensor
 std::vector<torch::Tensor> 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<torch::Tensor> kv_scale_orig_quant,
-    torch::optional<torch::Tensor> 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<torch::Tensor> 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");
@@ -129,17 +128,14 @@ std::vector<torch::Tensor> 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<float>();
+    }
+    if (kv_scale_quant_orig.has_value())
+    {
         kv_scale_quant_orig_ptr = kv_scale_quant_orig.value().data_ptr<float>();
-        TLLM_CHECK(kv_scale_orig_quant_ptr != nullptr);
-        TLLM_CHECK(kv_scale_quant_orig_ptr != nullptr);
     }
 
     std::vector<torch::Tensor> outputs;
@@ -199,9 +195,9 @@ std::vector<torch::Tensor> 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<torch::Tensor> kv_scale_orig_quant, torch::optional<torch::Tensor> 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<torch::Tensor> 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");
@@ -217,16 +213,10 @@ std::vector<torch::Tensor> 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<float>();
         kv_scale_quant_orig_ptr = kv_scale_quant_orig.value().data_ptr<float>();
-        TLLM_CHECK(kv_scale_orig_quant_ptr != nullptr);
-        TLLM_CHECK(kv_scale_quant_orig_ptr != nullptr);
     }
 
     std::vector<torch::Tensor> outputs;
@@ -293,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<torch::Tensor> kv_scale_orig_quant,
-    torch::optional<torch::Tensor> 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);
@@ -321,16 +311,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<float>();
-        kv_scale_quant_orig_ptr = kv_scale_quant_orig.value().data_ptr<float>();
-        TLLM_CHECK(kv_scale_orig_quant_ptr != nullptr);
-        TLLM_CHECK(kv_scale_quant_orig_ptr != nullptr);
     }
 
     if (input_dtype == torch::kFloat16)
@@ -425,7 +412,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 +440,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 +476,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"