[None][feat] Upgrade transformers dependency to 5.5.3

requirements.txt

@@ -30,7 +30,7 @@ nvidia-modelopt[torch]~=0.37.0
 # torch 2.10.0+cu130 depends on nvidia-nccl-cu13==2.28.9
 nvidia-nccl-cu13>=2.28.9,<=2.29.2
 nvidia-cuda-nvrtc
-transformers==5.3.0
+transformers==5.5.3
 prometheus_client
 prometheus_fastapi_instrumentator
 pydantic>=2.9.1
@@ -78,7 +78,7 @@ partial_json_parser
 mcp
 apache-tvm-ffi==0.1.6 # used for reduce nvidia-cutlass-dsl host overhead
 torch-c-dlpack-ext==0.1.3 # used for reduce nvidia-cutlass-dsl host overhead, optional package for improved torch tensor calling perf
-mistral-common==1.9.1
+mistral-common>=1.10.0
 torchao>=0.14.1,<0.16.0
 cuda-core
 llist

tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py

@@ -417,10 +417,11 @@ def __init__(self, batch_info_host: Optional[torch.Tensor] = None):
         # torch.compile metadata tracing without requiring a real .numpy() view.
         self._batch_info = batch_info_host
         # Cached zero-copy numpy view for fast scalar/list writes on the host path.
-        # `.numpy()` raises RuntimeError on fake tensors, so guard the call.
+        # `.numpy()` raises RuntimeError on fake tensors and TypeError on cuda tensors;
+        # both fall back to the slow tensor-write path.
         try:
             self._batch_info_np = batch_info_host.numpy()
-        except RuntimeError:
+        except (RuntimeError, TypeError):
             self._batch_info_np = None
 
     def serialize(self) -> torch.Tensor:

tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/torch_quant.py

@@ -513,6 +513,151 @@ def _safe_act_quant(x: torch.Tensor, block_size: int = 128) -> tuple:
     return y, s
 
 
+# Adapted from sgl-project/sglang fp8 block matmul kernel, vendored here to
+# decouple from transformers.integrations.finegrained_fp8 (which removed
+# w8a8_block_fp8_matmul_triton in transformers 5.5.x).
+@triton.jit
+def _w8a8_block_fp8_matmul_kernel(
+    A,
+    B,
+    C,
+    As,
+    Bs,
+    M,
+    N,
+    K,
+    group_n,
+    group_k,
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_As_m,
+    stride_As_k,
+    stride_Bs_k,
+    stride_Bs_n,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + (pid % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
+    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+
+    As_ptrs = As + offs_am * stride_As_m
+    offs_bsn = offs_bn // group_n
+    Bs_ptrs = Bs + offs_bsn * stride_Bs_n
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
+        b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
+
+        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * stride_bk
+
+    if C.dtype.element_ty == tl.bfloat16:
+        c = accumulator.to(tl.bfloat16)
+    elif C.dtype.element_ty == tl.float16:
+        c = accumulator.to(tl.float16)
+    else:
+        c = accumulator.to(tl.float32)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+
+def _w8a8_block_fp8_matmul_triton(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor,
+    Bs: torch.Tensor,
+    block_size: List[int],
+    output_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    if block_size is None:
+        block_n, block_k = 128, 128
+    else:
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+
+    assert A.shape[-1] == B.shape[-1]
+    assert As.numel() != 1, "per-tensor scales unsupported in vendored path"
+    assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
+    assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
+
+    M = A.numel() // A.shape[-1]
+    N, K = B.shape
+    assert B.ndim == 2 and B.is_contiguous()
+    assert Bs.ndim == 2
+    assert triton.cdiv(N, block_n) == Bs.shape[0]
+    assert triton.cdiv(K, block_k) == Bs.shape[1]
+
+    C_shape = A.shape[:-1] + (N,)
+    C = A.new_empty(C_shape, dtype=output_dtype)
+
+    BLOCK_SIZE_M = 128
+    if M < BLOCK_SIZE_M:
+        BLOCK_SIZE_M = max(triton.next_power_of_2(M), 16)
+    BLOCK_SIZE_K = block_k
+    BLOCK_SIZE_N = block_n
+
+    def grid(META):
+        return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
+
+    _w8a8_block_fp8_matmul_kernel[grid](
+        A,
+        B,
+        C,
+        As,
+        Bs,
+        M,
+        N,
+        K,
+        block_n,
+        block_k,
+        A.stride(-2),
+        A.stride(-1),
+        B.stride(1),
+        B.stride(0),
+        C.stride(-2),
+        C.stride(-1),
+        As.stride(-2),
+        As.stride(-1),
+        Bs.stride(1),
+        Bs.stride(0),
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=BLOCK_SIZE_N,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+        GROUP_SIZE_M=8,
+    )
+    return C
+
+
 @torch.library.custom_op("auto_deploy::torch_fake_quant_finegrained_fp8_linear", mutates_args=())
 def torch_fake_quant_finegrained_fp8_linear(
     input: torch.Tensor,  # [..., K]
@@ -532,8 +677,6 @@ def torch_fake_quant_finegrained_fp8_linear(
     - input_scale, input_zp, weight_zp are unused
     - block_size is inferred from weight and weight_scale_inv shapes
     """
-    from transformers.integrations.finegrained_fp8 import w8a8_block_fp8_matmul_triton
-
     weight_scale_inv = weight_scale[0]
 
     # Infer block_size from weight and weight_scale_inv shapes
@@ -545,7 +688,7 @@ def torch_fake_quant_finegrained_fp8_linear(
     block_size = [block_n, block_k]
 
     qinput, scale = _safe_act_quant(input, block_size[1])
-    output = w8a8_block_fp8_matmul_triton(
+    output = _w8a8_block_fp8_matmul_triton(
         qinput,
         weight_quantized,
         scale,

tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py

@@ -29,7 +29,7 @@
 
 from dataclasses import dataclass
 from types import SimpleNamespace
-from typing import Any, Dict, Optional, Union
+from typing import Any, ClassVar, Dict, Optional, Union
 
 import torch
 import torch.nn as nn
@@ -104,9 +104,11 @@ class EagleConfig(PretrainedConfig):
         model_type: The base model type (e.g., "llama", "nemotron_h") used to look up defaults.
     """
 
-    # Map model_type -> default Eagle config values
-    # Includes _checkpoint_conversion_mapping for model-specific weight key transformations
-    _drafter_defaults: Dict[str, Dict[str, Any]] = {
+    # Map model_type -> default Eagle config values.
+    # Annotated as ClassVar so transformers 5.5.x's @dataclass(kw_only=True)
+    # treatment of PretrainedConfig subclasses skips this (otherwise the
+    # mutable-dict default is rejected at class-creation time).
+    _drafter_defaults: ClassVar[Dict[str, Dict[str, Any]]] = {
         "llama": {
             "load_embedding_from_target": True,
             "load_lm_head_from_target": False,
@@ -140,24 +142,31 @@ class EagleConfig(PretrainedConfig):
     # Some custom HF config classes expose backward-compatibility fields as properties instead of
     # storing them directly in __dict__. Those values do not survive config.to_dict(), so carry
     # them over explicitly before rebuilding a generic EagleConfig.
-    _preserved_config_attrs: Dict[str, tuple[str, ...]] = {
+    _preserved_config_attrs: ClassVar[Dict[str, tuple[str, ...]]] = {
         "nemotron_h": ("mtp_hybrid_override_pattern",),
     }
 
-    def __init__(
-        self,
+    @classmethod
+    def from_base_config(
+        cls,
         config: PretrainedConfig,
         model_type: str,
     ):
-        if model_type not in self._drafter_defaults:
+        """Build an EagleConfig by merging a base model's config with type-specific defaults.
+
+        Use this factory instead of constructing ``EagleConfig`` directly: transformers>=5.5
+        applies ``@dataclass(kw_only=True)`` to ``PretrainedConfig`` subclasses, which
+        overrides any manually defined ``__init__``.
+        """
+        if model_type not in cls._drafter_defaults:
             raise ValueError(
                 f"Unsupported model_type '{model_type}' for EagleConfig. "
-                f"Supported types: {list(self._drafter_defaults.keys())}"
+                f"Supported types: {list(cls._drafter_defaults.keys())}"
             )
 
-        defaults = self._drafter_defaults[model_type]
+        defaults = cls._drafter_defaults[model_type]
         config_dict = config.to_dict()
-        for key in self._preserved_config_attrs.get(model_type, ()):
+        for key in cls._preserved_config_attrs.get(model_type, ()):
             if key not in config_dict and hasattr(config, key):
                 config_dict[key] = getattr(config, key)
 
@@ -170,7 +179,7 @@ def __init__(
                 )
 
         merged = deep_merge_dicts(defaults, config_dict)
-        super().__init__(**merged)
+        return cls(**merged)
 
 
 class LlamaRotaryEmbedding(nn.Module):

tensorrt_llm/_torch/auto_deploy/models/custom/modeling_nemotron_flash.py

@@ -878,7 +878,7 @@ class NemotronFlashPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["NemotronFlashAttentionDecoderLayer", "NemotronFlashMambaDecoderLayer"]
     _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn_2 = True
+    _supports_flash_attn = True
     _supports_sdpa = True
     _supports_cache_class = True
 

tensorrt_llm/_torch/auto_deploy/models/eagle.py

@@ -66,7 +66,7 @@ def _build_model(self, device: DeviceLikeType) -> nn.Module:
 
         # Convert base config to EagleConfig, preserving existing values
         # and applying model-specific defaults based on model_type
-        model_config = EagleConfig(model_config, model_type)
+        model_config = EagleConfig.from_base_config(model_config, model_type)
 
         with (init_empty_weights if device == "meta" else nullcontext)():
             model = EagleDrafterForCausalLM._from_config(model_config, **unused_kwargs)