feat: Support cos_sin_cache in all cases.

tensorrt_llm/_torch/attention_backend/interface.py

@@ -364,15 +364,15 @@ def from_config(config) -> "RopeParams":
 
         return rope_params
 
-    def create_rope_const_params(self):
+    def create_rope_const_params(self, interleave: bool = True):
         if self.dim == 0:
             return None, None
 
         RopeConstParams = namedtuple("RopeConstParams", ["inv_freq", "cos_sin"])
         extra_attrs = get_model_extra_attrs()
         if extra_attrs is not None:
             cache = extra_attrs.setdefault("rope_const_params", {})
-            rope_const_params = cache.get(self, None)
+            rope_const_params = cache.get((self, interleave), None)
             if rope_const_params is not None and rope_const_params.cos_sin(
             ) is not None:
                 return (
@@ -416,13 +416,17 @@ def create_rope_const_params(self):
                 dtype=torch.float32,
                 device='cuda',
             )
+        if not interleave:
+            rope_cos_sin = rope_cos_sin.reshape(
+                self.max_positions, -1,
+                2)[:, :self.dim // 2, :].transpose(0, 2, 1).reshape(1, -1)
         rope_cos_sin = torch.torch.tensor(
             rope_cos_sin,
             dtype=torch.float32,
             device='cuda',
         )
         if extra_attrs is not None:
-            cache[self] = RopeConstParams(
+            cache[(self, interleave)] = RopeConstParams(
                 weakref.ref(rope_inv_freq)
                 if rope_inv_freq is not None else None,
                 weakref.ref(rope_cos_sin),
@@ -437,6 +441,7 @@ class PositionalEmbeddingParams:
 
     # RoPE params
     rope: Optional[RopeParams] = None
+    is_neox: bool = True
 
     def __post_init__(self) -> None:
         if self.type.is_deferred():

tensorrt_llm/_torch/custom_ops/__init__.py

@@ -13,13 +13,13 @@
 ]
 
 if IS_FLASHINFER_AVAIABLE:
-    from .flashinfer_custom_ops import (flashinfer_apply_rope_inplace,
-                                        flashinfer_fused_add_rmsnorm,
-                                        flashinfer_rmsnorm,
-                                        flashinfer_silu_and_mul)
+    from .flashinfer_custom_ops import (
+        flashinfer_apply_rope_with_cos_sin_cache_inplace,
+        flashinfer_fused_add_rmsnorm, flashinfer_rmsnorm,
+        flashinfer_silu_and_mul)
     __all__ += [
         'flashinfer_silu_and_mul',
         'flashinfer_rmsnorm',
         'flashinfer_fused_add_rmsnorm',
-        'flashinfer_apply_rope_inplace',
+        'flashinfer_apply_rope_with_cos_sin_cache_inplace',
     ]

tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py

@@ -1,7 +1,6 @@
 import os
 import platform
 import traceback
-from typing import Optional
 
 import torch
 
@@ -48,8 +47,8 @@ def flashinfer_rmsnorm(input: torch.Tensor, weight: torch.Tensor,
         return rmsnorm(input, weight, eps, enable_pdl=ENABLE_PDL)
 
     @flashinfer_rmsnorm.register_fake
-    def rmsnorm_fake(input: torch.Tensor, weight: torch.Tensor,
-                     eps: float) -> torch.Tensor:
+    def _(input: torch.Tensor, weight: torch.Tensor,
+          eps: float) -> torch.Tensor:
         return torch.empty_like(input)
 
     @torch.library.custom_op("trtllm::flashinfer_fused_add_rmsnorm",
@@ -59,30 +58,33 @@ def flashinfer_fused_add_rmsnorm(input: torch.Tensor,
                                      weight: torch.Tensor, eps: float) -> None:
         fused_add_rmsnorm(input, residual, weight, eps, enable_pdl=ENABLE_PDL)
 
-    @torch.library.custom_op("trtllm::flashinfer_apply_rope_inplace",
-                             mutates_args=("q", "k"))
-    def flashinfer_apply_rope_inplace(
-        q: torch.Tensor,
-        k: torch.Tensor,
-        indptr: torch.Tensor,
-        offsets: torch.Tensor,
-        rotary_dim: Optional[int] = None,
-        interleave: bool = False,
-        rope_scale: float = 1,
-        rope_theta: float = 1e4,
+    @torch.library.custom_op(
+        "trtllm::flashinfer_apply_rope_with_cos_sin_cache_inplace",
+        mutates_args=("query", "key"))
+    def flashinfer_apply_rope_with_cos_sin_cache_inplace(
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        head_size: int,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool = True,
     ) -> None:
-        flashinfer.apply_rope_inplace(q, k, indptr, offsets, rotary_dim,
-                                      interleave, rope_scale, rope_theta)
-
-    @flashinfer_apply_rope_inplace.register_fake
-    def apply_rope_inplace_fake(
-        q: torch.Tensor,
-        k: torch.Tensor,
-        indptr: torch.Tensor,
-        offsets: torch.Tensor,
-        rotary_dim: Optional[int] = None,
-        interleave: bool = False,
-        rope_scale: float = 1,
-        rope_theta: float = 1e4,
+        flashinfer.apply_rope_with_cos_sin_cache_inplace(
+            positions,
+            query,
+            key,
+            head_size,
+            cos_sin_cache,
+            is_neox,
+        )
+
+    @flashinfer_apply_rope_with_cos_sin_cache_inplace.register_fake
+    def _(
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        head_size: int,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool = True,
     ):
         return

tensorrt_llm/_torch/model_config.py

@@ -35,14 +35,6 @@ def __post_init__(self):
             self.is_generation = self.is_generation_model(
                 self.pretrained_config.architectures)
 
-    @property
-    def fuse_pos_embd(self):
-        if self.attn_backend == 'TRTLLM':
-            return True
-        elif self.attn_backend == 'FLASHINFER':
-            return False
-        return False
-
     @property
     def enable_flash_mla(self):
         if self.attn_backend == 'TRTLLM':

tensorrt_llm/_torch/models/modeling_bert.py

@@ -81,15 +81,13 @@ def __init__(self,
                  layer_idx: Optional[int] = None):
         config = model_config.pretrained_config
         pos_embd_params = None
-        rotary_emb = None
         bias = True
         super().__init__(
             hidden_size=config.hidden_size,
             num_attention_heads=config.num_attention_heads,
             num_key_value_heads=config.num_attention_heads,
             max_position_embeddings=config.max_position_embeddings,
             bias=bias,
-            rotary_emb=rotary_emb,
             pos_embd_params=pos_embd_params,
             layer_idx=layer_idx,
             dtype=config.torch_dtype,

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -25,7 +25,6 @@
 from ..modules.linear import Linear
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
 from ..modules.rms_norm import RMSNorm
-from ..modules.rotary_embedding import RotaryEmbedding
 from ..pipeline_interface import PipelineInterface
 from ..speculative import MTPEagleWorker, MTPSpecMetadata, MTPWorker
 from ..utils import (AuxStreamType, EventType, Fp4QuantizedTensor,
@@ -60,20 +59,6 @@ def forward(self, hidden_states: torch.Tensor, lm_head: Linear,
         return logits
 
 
-class DeepseekV3RotaryEmbedding(RotaryEmbedding):
-
-    def __init__(self,
-                 config: PretrainedConfig,
-                 device: Optional[torch.device] = None):
-        head_dim = config.hidden_size // config.num_attention_heads
-        super().__init__(config,
-                         head_dim=head_dim,
-                         num_attention_heads=config.num_attention_heads,
-                         max_position_embeddings=config.max_position_embeddings,
-                         device=device,
-                         rope_type="default")
-
-
 class DeepseekV3Attention(MLA):
 
     def __init__(
@@ -83,14 +68,6 @@ def __init__(
         aux_stream: Optional[torch.cuda.Stream] = None,
     ):
         config = model_config.pretrained_config
-        if model_config.fuse_pos_embd:
-            pos_embd_params = PositionalEmbeddingParams(
-                type=PositionEmbeddingType.yarn,
-                rope=RopeParams.from_config(config),
-            )
-        else:
-            pos_embd_params = None
-
         predicted_tokens_per_seq = model_config.spec_config.num_nextn_predict_layers + 1 if model_config.spec_config is not None else 1
         super().__init__(hidden_size=config.hidden_size,
                          num_attention_heads=config.num_attention_heads,
@@ -103,8 +80,11 @@ def __init__(
                          predicted_tokens_per_seq=predicted_tokens_per_seq,
                          max_position_embeddings=config.max_position_embeddings,
                          bias=False,
-                         rotary_emb=DeepseekV3RotaryEmbedding(config),
-                         pos_embd_params=pos_embd_params,
+                         pos_embd_params=PositionalEmbeddingParams(
+                             type=PositionEmbeddingType.yarn,
+                             rope=RopeParams.from_config(config),
+                             is_neox=False,
+                         ),
                          layer_idx=layer_idx,
                          dtype=config.torch_dtype,
                          config=model_config,

tensorrt_llm/_torch/models/modeling_llama.py

@@ -22,30 +22,13 @@
 from ..modules.linear import (Linear, TensorParallelMode, WeightMode,
                               WeightsLoadingConfig)
 from ..modules.rms_norm import RMSNorm
-from ..modules.rotary_embedding import RotaryEmbedding
 from ..speculative import Eagle3SpecMetadata, SpecMetadata
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM,
                              EagerFusionConfig, MissingLayer,
                              register_auto_model, support_pp,
                              unpack_hidden_states)
 
 
-class LlamaRotaryEmbedding(RotaryEmbedding):
-
-    def __init__(
-        self,
-        config: LlamaConfig,
-        device: Optional[torch.device] = None,
-    ):
-        super().__init__(
-            config,
-            head_dim=config.hidden_size // config.num_attention_heads,
-            num_attention_heads=config.num_attention_heads,
-            max_position_embeddings=config.max_position_embeddings,
-            device=device,
-            rope_type="default" if config.rope_scaling is None else "llama3")
-
-
 class Llama4Attention(Attention):
 
     def __init__(
@@ -60,11 +43,11 @@ def __init__(
         # Note the convention no_rope_layers[layer_idx] == 0 means nope_layer
         is_nope_layer = config.no_rope_layers[layer_idx] == 0
 
-        use_rope = not is_nope_layer
-        if use_rope and model_config.fuse_pos_embd:
+        if not is_nope_layer:
             pos_embd_params = PositionalEmbeddingParams(
                 type=PositionEmbeddingType.rope_gptj,
                 rope=RopeParams.from_config(config),
+                is_neox=False,
             )
         else:
             pos_embd_params = None
@@ -75,7 +58,6 @@ def __init__(
             num_key_value_heads=config.num_key_value_heads,
             max_position_embeddings=config.max_position_embeddings,
             bias=config.attention_bias,
-            rotary_emb=LlamaRotaryEmbedding(config),
             pos_embd_params=pos_embd_params,
             layer_idx=layer_idx,
             dtype=config.torch_dtype,
@@ -93,21 +75,16 @@ def __init__(
         layer_idx: Optional[int] = None,
     ):
         config = model_config.pretrained_config
-        if model_config.fuse_pos_embd:
-            pos_embd_params = PositionalEmbeddingParams(
-                type=PositionEmbeddingType.rope_gpt_neox,
-                rope=RopeParams.from_config(config),
-            )
-        else:
-            pos_embd_params = None
         super().__init__(
             hidden_size=config.hidden_size,
             num_attention_heads=config.num_attention_heads,
             num_key_value_heads=config.num_key_value_heads,
             max_position_embeddings=config.max_position_embeddings,
             bias=config.attention_bias,
-            rotary_emb=LlamaRotaryEmbedding(config),
-            pos_embd_params=pos_embd_params,
+            pos_embd_params=PositionalEmbeddingParams(
+                type=PositionEmbeddingType.rope_gpt_neox,
+                rope=RopeParams.from_config(config),
+            ),
             layer_idx=layer_idx,
             dtype=config.torch_dtype,
             config=model_config,

tensorrt_llm/_torch/models/modeling_mamba_hybrid.py

@@ -76,7 +76,6 @@ def __init__(
             num_key_value_heads=config.num_key_value_heads,
             max_position_embeddings=config.max_position_embeddings,
             bias=config.attention_bias,
-            rotary_emb=None,
             pos_embd_params=None,
             layer_idx=layer_idx,
             dtype=config.torch_dtype,

tensorrt_llm/_torch/models/modeling_mixtral.py

@@ -15,7 +15,6 @@
 from ..modules.fused_moe import FusedMoE, RenormalizeMoeRoutingMethod
 from ..modules.linear import Linear
 from ..modules.rms_norm import RMSNorm
-from ..modules.rotary_embedding import RotaryEmbedding
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM,
                              register_auto_model)
 
@@ -71,20 +70,6 @@ def forward(
         return final_hidden_states
 
 
-class MixtralRotaryEmbedding(RotaryEmbedding):
-
-    def __init__(self,
-                 config: PretrainedConfig,
-                 device: Optional[torch.device] = None):
-        head_dim = config.hidden_size // config.num_attention_heads
-        super().__init__(config,
-                         head_dim=head_dim,
-                         num_attention_heads=config.num_attention_heads,
-                         max_position_embeddings=config.max_position_embeddings,
-                         device=device,
-                         rope_type="default")
-
-
 class MixtralAttention(Attention):
 
     def __init__(
@@ -93,21 +78,15 @@ def __init__(
         layer_idx: Optional[int] = None,
     ):
         config = model_config.pretrained_config
-        if model_config.fuse_pos_embd:
-            pos_embd_params = PositionalEmbeddingParams(
-                type=PositionEmbeddingType.rope_gpt_neox,
-                rope=RopeParams.from_config(config),
-            )
-        else:
-            pos_embd_params = None
-
         super().__init__(hidden_size=config.hidden_size,
                          num_attention_heads=config.num_attention_heads,
                          num_key_value_heads=config.num_key_value_heads,
                          max_position_embeddings=config.max_position_embeddings,
                          bias=False,
-                         rotary_emb=MixtralRotaryEmbedding(config),
-                         pos_embd_params=pos_embd_params,
+                         pos_embd_params=PositionalEmbeddingParams(
+                             type=PositionEmbeddingType.rope_gpt_neox,
+                             rope=RopeParams.from_config(config),
+                         ),
                          layer_idx=layer_idx,
                          dtype=config.torch_dtype,
                          config=model_config)