[https://nvbugs/5123103][fix] Fix torch compile for DeepSeekV3

cpp/tensorrt_llm/thop/mtpOp.cpp

@@ -93,8 +93,8 @@ std::tuple<th::Tensor, th::Tensor> mtp_prepare_drafter_inputs_op(th::Tensor& inp
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////
 std::tuple<th::Tensor, th::Tensor> mtp_sampling_and_accepted_draft_tokens_op(th::Tensor& logits,
-    th::Tensor& draftTokens, th::Tensor& targetTokens, th::Tensor& acceptedTokens, th::Tensor& numAcceptedTokens,
-    int64_t numMTPModules, int64_t batchSize, int64_t numContextRequest, int64_t vocabSize)
+    th::Tensor& draftTokens, th::Tensor& targetTokens, int64_t numMTPModules, int64_t batchSize,
+    int64_t numContextRequest, int64_t vocabSize)
 {
     int const numGenerationRequest = batchSize - numContextRequest;
     auto dataType = logits.scalar_type();
@@ -109,6 +109,9 @@ std::tuple<th::Tensor, th::Tensor> mtp_sampling_and_accepted_draft_tokens_op(th:
     TLLM_CHECK(draftTokensSizes[0] == (numGenerationRequest * numMTPModules));
 
     auto stream = at::cuda::getCurrentCUDAStream(logits.get_device());
+    auto acceptedTokens = torch::empty(
+        {batchSize, numMTPModules + 1}, at::TensorOptions().dtype(torch::kInt32).device(logits.device()));
+    auto numAcceptedTokens = torch::ones({batchSize}, at::TensorOptions().dtype(torch::kInt32).device(logits.device()));
 
     // Fill params
     tk::MTPSampleAndAcceptDraftTokensParam params;
@@ -288,7 +291,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
     m.def(
         "mtp_sampling_and_accepted_draft_tokens_op(Tensor logits, Tensor draftTokens, Tensor "
-        "targetTokens, Tensor acceptedTokens, Tensor numAcceptedTokens, int numMTPModules, "
+        "targetTokens, int numMTPModules, "
         "int batchSize, int numContextRequest, int vocabSize) -> (Tensor, Tensor)");
 }
 

tensorrt_llm/_torch/attention_backend/flashinfer.py

@@ -6,6 +6,7 @@
 import flashinfer
 import torch
 from flashinfer.jit.core import check_cuda_arch
+from typing_extensions import Self
 
 from tensorrt_llm.functional import AttentionMaskType
 from tensorrt_llm.models.modeling_utils import QuantConfig
@@ -163,6 +164,19 @@ def __post_init__(self) -> None:
                                                  device='cuda',
                                                  dtype=torch.int)
 
+    def create_cuda_graph_metadata(self,
+                                   max_batch_size: int,
+                                   sub_cross_metadata: bool = False,
+                                   max_draft_tokens: int = 0) -> Self:
+        metadata = super().create_cuda_graph_metadata(max_batch_size,
+                                                      sub_cross_metadata,
+                                                      max_draft_tokens)
+        metadata.max_num_requests = max_batch_size
+        metadata.max_num_tokens = max_batch_size * (1 + max_draft_tokens)
+        # Post init again to make sure all tensors are allocated
+        metadata.__post_init__()
+        return metadata
+
     @property
     def page_size(self) -> int:
         """
@@ -172,9 +186,7 @@ def page_size(self) -> int:
 
     def prepare(self) -> None:
         extra_attrs = get_model_extra_attrs()
-        if extra_attrs is not None:
-            extra_attrs["attention_metadata"] = weakref.ref(self)
-        else:
+        if extra_attrs is None:
             get_global_attrs().attention_metadata = weakref.ref(self)
         # start and end indices of each sequence in the ragged query
         torch.cumsum(self.seq_lens_cuda,

tensorrt_llm/_torch/attention_backend/interface.py

@@ -144,9 +144,9 @@ def on_update(self):
             ).item()
             self._num_generations = self._seq_lens.shape[0] - self.num_contexts
         if self._seq_lens_kv is not None:
-            self._num_tokens = int(self._seq_lens_kv.sum())
+            self._num_tokens = self._seq_lens_kv.sum().item()
         elif self._seq_lens is not None:
-            self._num_tokens = int(self._seq_lens.sum())
+            self._num_tokens = self._seq_lens.sum().item()
 
     @property
     def seq_lens(self) -> Optional[torch.Tensor]:
@@ -296,9 +296,6 @@ def create_cuda_graph_metadata(self,
                 )
 
         cuda_graph_metadata.num_contexts = 0
-        cuda_graph_metadata.max_num_requests = max_batch_size
-        cuda_graph_metadata.max_num_tokens = max_batch_size * (1 +
-                                                               max_draft_tokens)
         cuda_graph_metadata.__post_init__()
         return cuda_graph_metadata
 

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -1,3 +1,4 @@
+import weakref
 from dataclasses import dataclass, field
 from typing import Optional
 
@@ -7,6 +8,7 @@
 from tensorrt_llm.logger import logger
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
+from ..utils import get_global_attrs, get_model_extra_attrs
 from .interface import (AttentionBackend, AttentionInputType, AttentionMask,
                         AttentionMetadata, KVCacheParams, MLAParams,
                         PositionalEmbeddingParams, PredefinedAttentionMask,
@@ -531,7 +533,10 @@ def __post_init__(self) -> None:
                 )
 
     def prepare(self) -> None:
-
+        extra_attrs = get_model_extra_attrs()
+        # If model extra attrs is set, attention_metadata is setup in executor.
+        if extra_attrs is None:
+            get_global_attrs().attention_metadata = weakref.ref(self)
         if self.kv_cache_manager is None:
             # Convert the attention metadata to a TRT-LLM no cache attention metadata.
             assert self.kv_cache_manager is None, "no cache attention should not have KV cache manager"
@@ -726,7 +731,6 @@ def update_quant_config(self, new_quant_config: Optional[QuantConfig]):
             self.has_fp8_block_wise = self.quant_config.layer_quant_mode.has_fp8_block_scales(
             )
             self.has_nvfp4 = self.quant_config.layer_quant_mode.has_nvfp4()
-            self.has_nvfp4 = self.quant_config.layer_quant_mode.has_nvfp4()
 
     def get_local_layer_idx(self, metadata: TrtllmAttentionMetadata) -> int:
         if metadata.kv_cache_manager is None:

tensorrt_llm/_torch/compilation/patterns/ub_allreduce.py

@@ -191,15 +191,16 @@ def register_convert_supported_ar_to_ub(custom_pass: PatternMatcherPass):
         fusion = KeywordArg('fusion_op')
         trtllm_allreduce_default = CallFunction(
             torch.ops.trtllm.allreduce.default, input_node,
-            KeywordArg('residual_in'), KeywordArg('gamma'), Ignored(),
-            Ignored(), Ignored(), mapping.tp_group, strategy, fusion,
+            KeywordArg('residual_in'), KeywordArg('gamma'), KeywordArg('scale'),
+            None, Ignored(), mapping.tp_group, strategy, fusion,
             KeywordArg('eps'))
         convert_pattern = MultiOutputPattern([trtllm_allreduce_default])
 
         def empty_convert_supported_ar_to_ub(
             input: torch.Tensor,
             residual_in: torch.Tensor,
             gamma: torch.Tensor,
+            scale: torch.Tensor,
             fusion_op: int,
             eps: float,
         ):
@@ -209,12 +210,13 @@ def target_convert_supported_ar_to_ub(
             input: torch.Tensor,
             residual_in: torch.Tensor,
             gamma: torch.Tensor,
+            scale: torch.Tensor,
             fusion_op: int,
             eps: float,
         ):
             input = torch.ops.trtllm.copy_to_userbuffers(input)
             all_reduce_output = torch.ops.trtllm.allreduce(
-                input, residual_in, gamma, None, None, None, mapping.tp_group,
+                input, residual_in, gamma, scale, None, None, mapping.tp_group,
                 int(AllReduceStrategy.UB), fusion_op, eps)
             finalize_output = torch.ops.trtllm.userbuffers_allreduce_finalize(
                 all_reduce_output[-1], False)

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -4,6 +4,8 @@
 
 import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils
 
+from ..._utils import get_sm_version
+
 
 def _register_fake():
 
@@ -123,7 +125,7 @@ def _(input, force_applying_finalize):
     def _(a, b, a_scale, b_scale):
         m = a.shape[0]
         n = b.shape[0]
-        return a.new_empty((m, n))
+        return a.new_empty((m, n), dtype=torch.bfloat16)
 
     @torch.library.register_fake(
         "tensorrt_llm::static_quantize_e4m3_per_tensor")
@@ -284,3 +286,24 @@ def _(
         weight_bias: float,
     ) -> List[torch.Tensor]:
         return outputs
+
+    @torch.library.register_fake(
+        "trtllm::mtp_sampling_and_accepted_draft_tokens_op")
+    def _(logits: torch.Tensor, draft_tokens: torch.Tensor,
+          target_tokens: torch.Tensor, num_mtp_modules: int, batch_size: int,
+          num_context_request: int, vocab_size: int):
+        return logits.new_empty((batch_size, num_mtp_modules + 1),
+                                dtype=torch.int32), logits.new_empty(
+                                    (batch_size, ), dtype=torch.int32)
+
+    @torch.library.register_fake("trtllm::fp8_quantize_1x128")
+    def _(input: torch.Tensor):
+        pad_m = fp4_utils.pad_up(input.shape[0], 4)
+        blocked_n = (input.shape[1] + 127) // 128
+        if get_sm_version() >= 100:
+            sz = (blocked_n, input.shape[0])
+        else:
+            sz = (fp4_utils.pad_up(pad_m * blocked_n * 4, 128) // 4, )
+        return torch.empty_like(input,
+                                dtype=torch.float8_e4m3fn), input.new_empty(
+                                    sz, dtype=torch.float)

tensorrt_llm/_torch/model_config.py

@@ -32,6 +32,8 @@ class ModelConfig(Generic[TConfig]):
     attn_backend: str = 'TRTLLM'
     moe_backend: str = 'CUTLASS'  # options can be CUTLASS, TRTLLM
 
+    extra_attrs: Dict = field(default_factory=dict, repr=False, init=False)
+
     def __post_init__(self):
         if self.pretrained_config and hasattr(self.pretrained_config,
                                               "architectures"):

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -633,19 +633,19 @@ def _compute_mlp_tp_size(self, intermediate_size: int,
         """
 
         assert intermediate_size % block_size == 0, "intermediate_size must be divisible by block_size."
-
         if self.enable_attention_dp:
             # If using attention DP, the MLP also uses DP instead of TP.
             mlp_tp_size = 1
         else:
             # The two math.gcd operations ensure that mlp_tp_size falls in the candidate TP sizes.
-            mlp_tp_size = math.gcd(
-                math.gcd(
-                    intermediate_size // block_size,
-                    self.mapping.tp_size,
-                ),
-                self.mapping.gpus_per_node,  # Avoid costly inter-node TP
+            tp = math.gcd(
+                intermediate_size // block_size,
+                self.mapping.tp_size,
             )
+            mlp_tp_size = math.gcd(
+                tp,
+                self.mapping.gpus_per_node,
+            ) if tp > self.mapping.gpus_per_node else tp  # Avoid costly inter-node TP
         return mlp_tp_size
 
     def _enable_min_latency_mode(self, num_tokens: int):

tensorrt_llm/_torch/modules/attention.py

@@ -1,4 +1,5 @@
 import math
+import weakref
 from enum import IntEnum
 from typing import Optional, cast
 
@@ -15,6 +16,7 @@
 from ..distributed import AllReduceParams
 from ..model_config import ModelConfig
 from ..peft.lora.layer import LoraLayer, LoraModuleType
+from ..utils import get_model_extra_attrs
 from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
 from .multi_stream_utils import maybe_execute_in_parallel
 from .rms_norm import RMSNorm
@@ -277,6 +279,47 @@ def apply_qk_norm(self, q, k):
             "Please override the `apply_qk_norm` method in the subclass.")
 
 
+def extract_extra_attrs(layer_idx: str):
+    extra_attrs = get_model_extra_attrs()
+    assert extra_attrs is not None, "Model extra attrs is not set"
+
+    metadata_ref = extra_attrs.get("attention_metadata", None)
+    assert metadata_ref is not None, "Attention metadata is not set"
+    metadata = metadata_ref()
+    assert isinstance(
+        metadata,
+        TrtllmAttentionMetadata,
+    )
+
+    mla_layers = extra_attrs.get("mla_layers", None)
+    assert mla_layers is not None, "MLA layers is not registered"
+    mla_layer_ref = mla_layers.get(layer_idx, None)
+    assert mla_layer_ref is not None, f"Cannot find MLA layer for layer {layer_idx}"
+    mla_layer = mla_layer_ref()
+    assert isinstance(
+        mla_layer,
+        MLA), "MLA layer must be a subclass of MLA or an instance of MLA"
+
+    return metadata, mla_layer
+
+
+@torch.library.custom_op("trtllm::mla_custom_op", mutates_args=())
+def mla_custom_op(
+    position_ids: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    layer_idx: str,
+) -> torch.Tensor:
+    metadata, mla_layer = extract_extra_attrs(layer_idx)
+
+    return mla_layer.forward_impl(position_ids, hidden_states, metadata)
+
+
+@mla_custom_op.register_fake
+def _(position_ids, hidden_states, layer_idx):
+    _, mla_layer = extract_extra_attrs(layer_idx)
+    return mla_layer.forward_impl_fake(hidden_states)
+
+
 class MLA(nn.Module):
 
     def __init__(
@@ -324,6 +367,7 @@ def __init__(
         """
         super().__init__()
         self.layer_idx = layer_idx
+        self.layer_idx_str = str(layer_idx)
         self.dtype = dtype
 
         self.hidden_size = hidden_size
@@ -351,6 +395,14 @@ def __init__(
 
         assert pos_embd_params is not None, "pos_embd_params must be provided in MLA"
 
+        self.register_to_config = False
+        if config is not None:
+            if "mla_layers" not in config.extra_attrs:
+                config.extra_attrs["mla_layers"] = {}
+            config.extra_attrs["mla_layers"][self.layer_idx_str] = weakref.ref(
+                self)
+            self.register_to_config = True
+
         # tensor parallel
         config = config or ModelConfig()
         tp_size = config.mapping.tp_size
@@ -584,12 +636,17 @@ def apply_rope(
                                                    self.qk_rope_head_dim)
         return k_pe
 
-    def forward(
+    def forward_impl_fake(self, hidden_states: torch.Tensor):
+        num_tokens = hidden_states.shape[0]
+        hidden_size = self.o_proj.in_features
+        return hidden_states.new_empty([num_tokens, hidden_size],
+                                       dtype=hidden_states.dtype)
+
+    def forward_impl(
         self,
-        position_ids: Optional[torch.LongTensor],
+        position_ids: Optional[torch.Tensor],
         hidden_states: torch.Tensor,
         attn_metadata: AttentionMetadata,
-        all_reduce_params: Optional[AllReduceParams] = None,
     ) -> torch.Tensor:
         """
         Forward pass for the MLA module.
@@ -692,8 +749,6 @@ def forward(
         else:
             attn_output = attn_output_gen
 
-        attn_output = self.o_proj(attn_output,
-                                  all_reduce_params=all_reduce_params)
         return attn_output
 
     def _maybe_concat_qkv(self, q, k, v):
@@ -987,3 +1042,20 @@ def forward_generation(
 
         # [seq, num_heads * v_head_dim]
         return attn_output.flatten(1, 2)
+
+    def forward(
+        self,
+        position_ids: Optional[torch.Tensor],
+        hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        all_reduce_params: Optional[AllReduceParams] = None,
+    ) -> torch.Tensor:
+        if self.register_to_config:
+            attn_output = torch.ops.trtllm.mla_custom_op(
+                position_ids, hidden_states, self.layer_idx_str)
+        else:
+            attn_output = self.forward_impl(position_ids, hidden_states,
+                                            attn_metadata)
+        attn_output = self.o_proj(attn_output,
+                                  all_reduce_params=all_reduce_params)
+        return attn_output