[https://nvbugspro.nvidia.com/bug/5329655] [feat] Pytorch path add spec dec param to attention op (NVIDIA#5146)

Signed-off-by: Jhao-Ting Chen <[email protected]>

Author: jhaotingc

cpp/tensorrt_llm/thop/attentionOp.cpp

@@ -78,7 +78,8 @@ class RunnerBase
         torch::optional<torch::Tensor> mrope_rotary_cos_sin, torch::optional<torch::Tensor> mrope_position_deltas,
         torch::optional<torch::Tensor> mla_context_paged_kv,
         torch::optional<torch::Tensor> mla_context_kv_cache_block_offsets,
-        torch::optional<torch::Tensor> softmax_stats_tensor) const
+        torch::optional<torch::Tensor> softmax_stats_tensor,
+        c10::ArrayRef<std::optional<torch::Tensor>> spec_decoding_tensor_params) const
         = 0;
 };
 
@@ -129,7 +130,8 @@ class Runner : public RunnerBase
         torch::optional<torch::Tensor> mrope_rotary_cos_sin, torch::optional<torch::Tensor> mrope_position_deltas,
         torch::optional<torch::Tensor> mla_context_paged_kv,
         torch::optional<torch::Tensor> mla_context_kv_cache_block_offsets,
-        torch::optional<torch::Tensor> softmax_stats_tensor) const override
+        torch::optional<torch::Tensor> softmax_stats_tensor,
+        c10::ArrayRef<std::optional<torch::Tensor>> spec_decoding_tensor_params) const override
     {
         auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
         T* attention_input = static_cast<T*>(qkv.slice(0, token_offset).data_ptr());
@@ -322,6 +324,27 @@ class Runner : public RunnerBase
             {
                 enqueue_params.mrope_position_deltas = mrope_position_deltas.value().data_ptr<int32_t>();
             }
+            if (op.mIsSpecDecodingEnabled && op.mUseSpecDecoding)
+            {
+                TORCH_CHECK(spec_decoding_tensor_params.size() == 3,
+                    "Expecting 3 tensors for spec-dec mode, spec_decoding_generation_lengths, "
+                    "spec_decoding_position_offsets and spec_decoding_packed_mask.");
+                TORCH_CHECK(spec_decoding_tensor_params[0].has_value(),
+                    "Expecting spec_decoding_generation_lengths spec-dec mode.");
+                TORCH_CHECK(spec_decoding_tensor_params[1].has_value(),
+                    "Expecting spec_decoding_position_offsets spec-dec mode.");
+                TORCH_CHECK(
+                    spec_decoding_tensor_params[2].has_value(), "Expecting spec_decoding_packed_mask spec-dec mode.");
+
+                enqueue_params.spec_decoding_generation_lengths
+                    = spec_decoding_tensor_params[0].value().data_ptr<int32_t>();
+                enqueue_params.spec_decoding_position_offsets
+                    = spec_decoding_tensor_params[1].value().data_ptr<int32_t>();
+                enqueue_params.spec_decoding_packed_mask = spec_decoding_tensor_params[2].value().data_ptr<int32_t>();
+                enqueue_params.spec_decoding_is_generation_length_variable = true;
+                enqueue_params.spec_decoding_max_generation_length = input_seq_length + 1;
+            }
+
             // Current mlaGeneration will using fmha to do attention, so we don't go into enqueueGeneration
             if (op.isMLAEnabled())
             {
@@ -384,15 +407,14 @@ void attention_inplace(torch::Tensor q, torch::optional<torch::Tensor> k, torch:
     int64_t const attention_window_size, int64_t const sink_token_length, int64_t const beam_width,
     int64_t const mask_type, int64_t const quant_mode, double const q_scaling, int64_t const position_embedding_type,
     int64_t const rotary_embedding_dim, double const rotary_embedding_base, int64_t const rotary_embedding_scale_type,
-    double const rotary_embedding_scale, double const rotary_embedding_short_m_scale,
-    double const rotary_embedding_long_m_scale, int64_t const rotary_embedding_max_positions,
-    int64_t const rotary_embedding_original_max_positions, bool const use_paged_context_fmha,
-    std::optional<int64_t> attention_input_type, bool is_mla_enable, std::optional<int64_t> q_lora_rank,
-    std::optional<int64_t> kv_lora_rank, std::optional<int64_t> qk_nope_head_dim,
+    c10::ArrayRef<double> rotary_embedding_scales, c10::ArrayRef<int64_t> rotary_embedding_max_position_info,
+    bool const use_paged_context_fmha, std::optional<int64_t> attention_input_type, bool is_mla_enable,
+    std::optional<int64_t> q_lora_rank, std::optional<int64_t> kv_lora_rank, std::optional<int64_t> qk_nope_head_dim,
     std::optional<int64_t> qk_rope_head_dim, std::optional<int64_t> v_head_dim,
     torch::optional<torch::Tensor> mrope_rotary_cos_sin, torch::optional<torch::Tensor> mrope_position_deltas,
     std::optional<torch::Tensor> mla_context_paged_kv, std::optional<torch::Tensor> mla_context_kv_cache_block_offsets,
-    std::optional<int64_t> attention_chunk_size, std::optional<torch::Tensor> softmax_stats_tensor)
+    std::optional<int64_t> attention_chunk_size, std::optional<torch::Tensor> softmax_stats_tensor,
+    c10::List<bool> spec_decoding_bool_params, c10::ArrayRef<std::optional<torch::Tensor>> spec_decoding_tensor_params)
 {
     TLLM_LOG_TRACE("Attention op starts at layer %d", layer_idx);
     // Use these tensors to infer if the attention is using KV cache
@@ -462,6 +484,12 @@ void attention_inplace(torch::Tensor q, torch::optional<torch::Tensor> k, torch:
     runner->attention_window_size = attention_window_size;
     runner->sink_token_length = sink_token_length;
 
+    double const rotary_embedding_scale = rotary_embedding_scales[0];
+    double const rotary_embedding_short_m_scale = rotary_embedding_scales[1];
+    double const rotary_embedding_long_m_scale = rotary_embedding_scales[2];
+    int64_t const rotary_embedding_max_positions = rotary_embedding_max_position_info[0];
+    int64_t const rotary_embedding_original_max_positions = rotary_embedding_max_position_info[1];
+
     auto op = std::make_shared<AttentionOp>();
     op->mType = dtype;
     op->mFMHAForceFP32Acc = dtype == nvinfer1::DataType::kBF16;
@@ -494,6 +522,12 @@ void attention_inplace(torch::Tensor q, torch::optional<torch::Tensor> k, torch:
 
     op->mAttentionChunkSize = attention_chunk_size;
 
+    TORCH_CHECK(spec_decoding_bool_params.size() == 2,
+        "Expecting 2 bools for spec-dec mode, is_spec_decoding_enabled and use_spec_decoding.");
+    op->mIsSpecDecodingEnabled = spec_decoding_bool_params[0]; // is_spec_decoding_enabled
+    op->mUseSpecDecoding = spec_decoding_bool_params[1];       // use_spec_decoding
+    op->mMultiBlockMode = op->mIsSpecDecodingEnabled ? false : true;
+
     if (is_mla_enable)
     {
         // MLA does not support NVFP4 output yet.
@@ -610,7 +644,7 @@ void attention_inplace(torch::Tensor q, torch::optional<torch::Tensor> k, torch:
             host_kv_cache_block_offsets, host_kv_cache_pool_pointers, host_kv_cache_pool_mapping, cache_indirection,
             kv_scale_orig_quant, kv_scale_quant_orig, out_scale, rotary_inv_freq, rotary_cos_sin, latent_cache, q_pe,
             block_ids_per_seq, mrope_rotary_cos_sin, mrope_position_deltas, mla_context_paged_kv,
-            mla_context_kv_cache_block_offsets, softmax_stats_tensor);
+            mla_context_kv_cache_block_offsets, softmax_stats_tensor, spec_decoding_tensor_params);
     }
 
     if ((num_generations > 0) && (attn_input_type != AttentionInputType::ContextOnly))
@@ -626,7 +660,7 @@ void attention_inplace(torch::Tensor q, torch::optional<torch::Tensor> k, torch:
             host_kv_cache_block_offsets, host_kv_cache_pool_pointers, host_kv_cache_pool_mapping, cache_indirection,
             kv_scale_orig_quant, kv_scale_quant_orig, out_scale, rotary_inv_freq, rotary_cos_sin, latent_cache, q_pe,
             block_ids_per_seq, mrope_rotary_cos_sin, mrope_position_deltas, mla_context_paged_kv,
-            mla_context_kv_cache_block_offsets, softmax_stats_tensor);
+            mla_context_kv_cache_block_offsets, softmax_stats_tensor, spec_decoding_tensor_params);
     }
 
     TLLM_LOG_TRACE("Attention op stops at layer %d", layer_idx);
@@ -731,11 +765,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         ", int rotary_embedding_dim"
         ", float rotary_embedding_base"
         ", int rotary_embedding_scale_type"
-        ", float rotary_embedding_scale"
-        ", float rotary_embedding_short_m_scale"
-        ", float rotary_embedding_long_m_scale"
-        ", int rotary_embedding_max_positions"
-        ", int rotary_embedding_original_max_positions"
+        ", float[] rotary_embedding_scales"
+        ", int[] rotary_embedding_max_position_info"
         ", bool use_paged_context_fmha"
         ", int? attention_input_type"
         ", bool is_mla_enable"
@@ -750,6 +781,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         ", Tensor? mla_context_kv_cache_block_offsets"
         ", int? attention_chunk_size"
         ", Tensor? softmax_stats_tensor"
+        ", bool[] spec_decoding_bool_params"
+        ", Tensor?[] spec_decoding_tensor_params"
         ") -> ()");
 
     m.def("attention_supports_nvfp4_output", &torch_ext::attention_supports_nvfp4_output);

tensorrt_llm/_torch/attention_backend/interface.py

@@ -316,6 +316,12 @@ def create_cuda_graph_metadata(self,
         cuda_graph_metadata.__post_init__()
         return cuda_graph_metadata
 
+    def update_spec_dec_param(self, is_spec_decoding_enabled, is_spec_dec_tree,
+                              is_spec_dec_dynamic_tree, max_draft_tokens):
+        """
+        Hook to be called when using TRTLLM attention backend in spec-dec mode.
+        """
+
 
 class PositionalEmbedder(Protocol):
     """

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -1,10 +1,12 @@
+import math
 import os
 import weakref
 from dataclasses import dataclass, field
 from typing import Optional
 
 import torch
 
+from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.functional import AttentionMaskType
 from tensorrt_llm.logger import logger
 from tensorrt_llm.models.modeling_utils import QuantConfig
@@ -64,6 +66,10 @@ class TrtllmAttentionWrapper:
     qk_nope_head_dim: Optional[int]
     v_head_dim: Optional[int]
     attention_chunk_size: Optional[int]
+    use_spec_decoding: bool
+    spec_decoding_position_offsets: Optional[torch.Tensor]
+    spec_decoding_packed_mask: Optional[torch.Tensor]
+    spec_decoding_generation_lengths: Optional[torch.Tensor]
     kwargs: dict
 
     def __init__(
@@ -169,6 +175,11 @@ def plan(
         mla_context_paged_kv: Optional[torch.Tensor] = None,
         mla_context_kv_cache_block_offsets: Optional[torch.Tensor] = None,
         softmax_stats_tensor: Optional[torch.Tensor] = None,
+        is_spec_decoding_enabled: bool = False,
+        use_spec_decoding: bool = False,
+        spec_decoding_position_offsets: Optional[torch.Tensor] = None,
+        spec_decoding_packed_mask: Optional[torch.Tensor] = None,
+        spec_decoding_generation_lengths: Optional[torch.Tensor] = None,
         **kwargs,
     ):
         """
@@ -245,7 +256,11 @@ def plan(
             self.rope_params.max_positions = max_sequence_length
             self.rotary_inv_freq, self.rotary_cos_sin = self.rope_params.create_rope_const_params(
             )
-
+        self.is_spec_decoding_enabled = is_spec_decoding_enabled
+        self.use_spec_decoding = use_spec_decoding
+        self.spec_decoding_position_offsets = spec_decoding_position_offsets
+        self.spec_decoding_packed_mask = spec_decoding_packed_mask
+        self.spec_decoding_generation_lengths = spec_decoding_generation_lengths
         self.kwargs.update(kwargs)
 
     def run(
@@ -374,6 +389,23 @@ def run(
             # output is provided, expect output_sf be provided as well if has NVFP4 output.
             assert out_dtype is None or out_dtype != torch.uint8 or output_sf is not None
 
+        # packing parameters to avoid maxing out 64 arguments
+        rotary_embedding_scales = [
+            self.rotary_embedding_scale, self.rotary_embedding_short_m_scale,
+            self.rotary_embedding_long_m_scale
+        ]
+        rotary_embedding_max_position_info = [
+            self.rotary_embedding_max_positions,
+            self.rotary_embedding_original_max_positions
+        ]
+        spec_decoding_bool_params = [
+            self.is_spec_decoding_enabled, self.use_spec_decoding
+        ]
+        spec_decoding_tensor_params = [
+            self.spec_decoding_generation_lengths,
+            self.spec_decoding_position_offsets, self.spec_decoding_packed_mask
+        ]
+
         torch.ops.trtllm.attention_inplace(
             q,
             k,
@@ -420,11 +452,8 @@ def run(
             self.rotary_embedding_dim,
             self.rotary_embedding_base,
             self.rotary_embedding_scale_type,
-            self.rotary_embedding_scale,
-            self.rotary_embedding_short_m_scale,
-            self.rotary_embedding_long_m_scale,
-            self.rotary_embedding_max_positions,
-            self.rotary_embedding_original_max_positions,
+            rotary_embedding_scales,
+            rotary_embedding_max_position_info,
             self.use_paged_context_fmha,
             self.attention_input_type,
             self.is_mla_enable,
@@ -439,6 +468,8 @@ def run(
             self.mla_context_kv_cache_block_offsets,
             self.attention_chunk_size,
             self.softmax_stats_tensor,
+            spec_decoding_bool_params,
+            spec_decoding_tensor_params,
         )
 
         # reset the planned states (especially tensors) to avoid memory leak
@@ -495,6 +526,23 @@ class TrtllmAttentionMetadata(AttentionMetadata):
                                                 init=True,
                                                 repr=False)
 
+    # Flags to enable spec-dec mode (multi-query mode) in TRTLLM XQA Kernels
+    # spec decoding mode can be enabled for non-TRTLLM-gen kernels (pre-Blackwell XQA kernels)
+    # is_spec_decoding_enabled specifies if spec-dec mode is supported for the entire runtime.
+    is_spec_decoding_enabled: bool = False
+    # use_spec_decoding determines if the attention layer should be run in spec-dec mode at the specific step / layer.
+    use_spec_decoding: bool = False
+
+    # if spec-dec tree is a tree or a chain (linear tree)
+    is_spec_dec_tree: bool = False
+    # if spec-dec tree wouldn't be changed at all, the mask won't be computed every step.
+    is_spec_dec_dynamic_tree: bool = False
+
+    # parameters required for spec-dec mode
+    spec_decoding_position_offsets: Optional[torch.Tensor] = None
+    spec_decoding_packed_mask: Optional[torch.Tensor] = None
+    spec_decoding_generation_lengths: Optional[torch.Tensor] = None
+
     @property
     def max_seq_len(self) -> int:
         """
@@ -849,6 +897,76 @@ def prepare_paged_context_mla(self, cached_token_lens: torch.Tensor,
         self.ctx_kv_indptr[:self.num_contexts + 1].copy_(
             self.host_ctx_kv_indptr[:self.num_contexts + 1], non_blocking=True)
 
+    def update_spec_dec_param(self, is_spec_decoding_enabled, is_spec_dec_tree,
+                              is_spec_dec_dynamic_tree, max_draft_tokens):
+        # spec_dec mode should only be enabled for pre-Blackwell machines and when there's a spec-dec tree.
+        self.is_spec_decoding_enabled = is_spec_decoding_enabled and get_sm_version(
+        ) < 100
+
+        # use_spec_decoding is default to true by default, change in runtime by layers / requests
+        self.use_spec_decoding = self.is_spec_decoding_enabled
+
+        self.is_spec_dec_tree = is_spec_dec_tree
+        self.is_spec_dec_dynamic_tree = is_spec_dec_dynamic_tree
+
+        # Parameters can be fixed and not changed during runtime if the
+        if self.is_spec_decoding_enabled:
+            self.spec_decoding_position_offsets = torch.empty(
+                [self.max_num_requests, max_draft_tokens + 1],
+                dtype=torch.int,
+                device='cuda',
+            )
+
+            self.spec_decoding_packed_mask = torch.empty(
+                [
+                    self.max_num_requests, max_draft_tokens + 1,
+                    math.ceil(max_draft_tokens / 32)
+                ],
+                dtype=torch.int,
+                device='cuda',
+            )
+
+            self.spec_decoding_generation_lengths = torch.empty(
+                [self.max_num_requests],
+                dtype=torch.int,
+                device='cuda',
+            )
+
+            if self.is_spec_dec_dynamic_tree:
+                assert False, "currently dynamic tree is not supported"
+            else:
+                # Populate the mask that won't change during inference phase.
+                self.generate_spec_decoding_position_offsets(
+                    max_draft_tokens=max_draft_tokens)
+                self.generate_spec_decoding_packed_mask(
+                    max_draft_tokens=max_draft_tokens)
+                self.generate_spec_decoding_generation_length(
+                    max_draft_tokens=max_draft_tokens)
+
+    def generate_spec_decoding_position_offsets(self, max_draft_tokens):
+        assert not self.is_spec_dec_tree, "only chained/linear tree is supported now"
+        position_offset = torch.arange(max_draft_tokens + 1,
+                                       dtype=torch.int,
+                                       device='cpu',
+                                       pin_memory=True)
+
+        # fill all the batches with same position offset
+        self.spec_decoding_position_offsets.copy_(position_offset,
+                                                  non_blocking=True)
+
+    def generate_spec_decoding_packed_mask(self, max_draft_tokens):
+        assert not self.is_spec_dec_tree, "only chained/linear tree is supported now"
+        dummy_idx = torch.arange(max_draft_tokens + 1)
+        spec_decoding_packed_mask = torch.pow(2, dummy_idx + 1) - 1
+        self.spec_decoding_packed_mask[:, :, 0].copy_(spec_decoding_packed_mask,
+                                                      non_blocking=True)
+
+    def generate_spec_decoding_generation_length(self, max_draft_tokens):
+        spec_decoding_generation_length = torch.full((self.max_num_requests, ),
+                                                     max_draft_tokens + 1)
+        self.spec_decoding_generation_lengths[:self.max_num_requests].copy_(
+            spec_decoding_generation_length, non_blocking=True)
+
 
 class TrtllmAttention(AttentionBackend[TrtllmAttentionMetadata]):
 
@@ -984,7 +1102,6 @@ def forward(
                 use_paged_context_fmha=use_paged_context_fmha,
                 is_mla_enable=self.is_mla_enable,
             )
-
         self.wrapper.plan(
             layer_idx=self.get_local_layer_idx(metadata),
             tokens_per_block=metadata.tokens_per_block,
@@ -1021,6 +1138,13 @@ def forward(
             mla_context_kv_cache_block_offsets=
             mla_context_kv_cache_block_offsets,
             softmax_stats_tensor=softmax_stats_tensor,
+            is_spec_decoding_enabled=metadata.is_spec_decoding_enabled,
+            use_spec_decoding=metadata.use_spec_decoding,
+            spec_decoding_position_offsets=metadata.
+            spec_decoding_position_offsets,
+            spec_decoding_packed_mask=metadata.spec_decoding_packed_mask,
+            spec_decoding_generation_lengths=metadata.
+            spec_decoding_generation_lengths,
         )
         out_dtype = None
         if out_scale is not None: