Repository: NVIDIA/TensorRT-LLM
Pull request #3517: feat: Support cos_sin_cache in all cases.
Author: login: yuxianq
Created: 2025-04-14T04:17:51Z | Merged: 2025-04-16T05:48:44Z

This MR contains the following updates:

1. Handle `fuse_pos_embd=True/False` and create `RotaryEmbedding` inside attention module, so that the users don't need to handle it in the modeling files.
2. Cache `cos_sin` for unfused rope implementation. If flashinfer is available, use `apply_rope_with_cos_sin_cache_inplace` instead of `apply_rope_inplace`. Otherwise, we fallback to pure pytorch implementation, which can support any rope now.



Diff:
diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py
index 74f097b952ca..28557d8f15f7 100644
--- a/tensorrt_llm/_torch/attention_backend/interface.py
+++ b/tensorrt_llm/_torch/attention_backend/interface.py
@@ -364,7 +364,7 @@ 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
 
@@ -372,7 +372,7 @@ def create_rope_const_params(self):
         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():
diff --git a/tensorrt_llm/_torch/custom_ops/__init__.py b/tensorrt_llm/_torch/custom_ops/__init__.py
index 652f3b65b0f1..f9a0f69f29ce 100644
--- a/tensorrt_llm/_torch/custom_ops/__init__.py
+++ b/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',
     ]
diff --git a/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py b/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py
index acdc93d56187..5aa2f91b6dfb 100644
--- a/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py
+++ b/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
diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py
index ebd68954c54f..091eb3463503 100644
--- a/tensorrt_llm/_torch/model_config.py
+++ b/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':
diff --git a/tensorrt_llm/_torch/models/modeling_bert.py b/tensorrt_llm/_torch/models/modeling_bert.py
index 19d61f6f0e2f..e8218a53f707 100644
--- a/tensorrt_llm/_torch/models/modeling_bert.py
+++ b/tensorrt_llm/_torch/models/modeling_bert.py
@@ -81,7 +81,6 @@ 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,
@@ -89,7 +88,6 @@ def __init__(self,
             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,
diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
index f2b9e64b0883..42109529105b 100644
--- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py
+++ b/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,
diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py
index 621e44c08a1c..0c76013ab0f2 100644
--- a/tensorrt_llm/_torch/models/modeling_llama.py
+++ b/tensorrt_llm/_torch/models/modeling_llama.py
@@ -22,7 +22,6 @@
 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,
@@ -30,22 +29,6 @@
                              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,
diff --git a/tensorrt_llm/_torch/models/modeling_mamba_hybrid.py b/tensorrt_llm/_torch/models/modeling_mamba_hybrid.py
index 543c13f6b87e..3b8540d6ffd4 100644
--- a/tensorrt_llm/_torch/models/modeling_mamba_hybrid.py
+++ b/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,
diff --git a/tensorrt_llm/_torch/models/modeling_mixtral.py b/tensorrt_llm/_torch/models/modeling_mixtral.py
index 8d2beed94f7b..bc297bc01dba 100644
--- a/tensorrt_llm/_torch/models/modeling_mixtral.py
+++ b/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)
diff --git a/tensorrt_llm/_torch/models/modeling_nemotron.py b/tensorrt_llm/_torch/models/modeling_nemotron.py
index b5a40ffba6d9..e75c8c65571c 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron.py
@@ -14,27 +14,10 @@
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.mlp import MLP
-from ..modules.rotary_embedding import RotaryEmbedding
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM,
                              register_auto_model)
 
 
-class NemotronRotaryEmbedding(RotaryEmbedding):
-
-    def __init__(
-        self,
-        config: NemotronConfig,
-        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")
-
-
 class NemotronAttention(Attention):
 
     def __init__(
@@ -43,21 +26,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=NemotronRotaryEmbedding(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,
diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_nas.py b/tensorrt_llm/_torch/models/modeling_nemotron_nas.py
index 6d02765d8c52..62e749608379 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron_nas.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron_nas.py
@@ -15,27 +15,10 @@
 from ..modules.gated_mlp import GatedMLP
 from ..modules.linear import Linear, TensorParallelMode
 from ..modules.rms_norm import RMSNorm
-from ..modules.rotary_embedding import RotaryEmbedding
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM,
                              register_auto_model)
 
 
-class NemotronNASRotaryEmbedding(RotaryEmbedding):
-
-    def __init__(self, config: PretrainedConfig, layer_idx: int):
-        if config.rope_scaling is not None:
-            rope_type = config.rope_scaling.get("rope_type",
-                                                config.rope_scaling.get("type"))
-        else:
-            rope_type = "default"
-        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,
-                         rope_type=rope_type)
-
-
 def _ffn_mult_to_intermediate_size(ffn_mult: float, n_embd: int) -> int:
     intermediate_size = int(2 * ffn_mult * n_embd / 3)
     return _find_multiple(intermediate_size, 256)
@@ -67,21 +50,16 @@ class NemotronNASAttention(Attention):
     def __init__(self, model_config: ModelConfig[PretrainedConfig],
                  layer_idx: int):
         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[layer_idx],
             max_position_embeddings=config.max_position_embeddings,
             bias=False,
-            pos_embd_params=pos_embd_params,
-            rotary_emb=NemotronNASRotaryEmbedding(config, layer_idx=layer_idx),
+            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)
diff --git a/tensorrt_llm/_torch/models/modeling_qwen.py b/tensorrt_llm/_torch/models/modeling_qwen.py
index 558bd06b03a3..c0c76b301c76 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen.py
@@ -15,28 +15,11 @@
 from ..modules.gated_mlp import GatedMLP
 from ..modules.linear import Linear, TensorParallelMode
 from ..modules.rms_norm import RMSNorm
-from ..modules.rotary_embedding import RotaryEmbedding
 from ..pipeline_interface import PipelineInterface
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM,
                              register_auto_model)
 
 
-class QwenRotaryEmbedding(RotaryEmbedding):
-
-    def __init__(
-        self,
-        config: Qwen2Config,
-        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")
-
-
 class QwenAttention(Attention):
 
     def __init__(
@@ -45,27 +28,23 @@ def __init__(
         layer_idx: Optional[int] = None,
     ):
         config = model_config.pretrained_config
-        if model_config.fuse_pos_embd:
-            if getattr(config, "rope_scaling", None) is not None:
-                pos_embd_params = PositionalEmbeddingParams(
-                    type=PositionEmbeddingType.from_string(
-                        config.rope_scaling["type"]),
-                    rope=RopeParams.from_config(config),
-                )
-            else:
-                pos_embd_params = PositionalEmbeddingParams(
-                    type=PositionEmbeddingType.rope_gpt_neox,
-                    rope=RopeParams.from_config(config),
-                )
+        if getattr(config, "rope_scaling", None) is not None:
+            pos_embd_params = PositionalEmbeddingParams(
+                type=PositionEmbeddingType.from_string(
+                    config.rope_scaling["type"]),
+                rope=RopeParams.from_config(config),
+            )
         else:
-            pos_embd_params = None
+            pos_embd_params = PositionalEmbeddingParams(
+                type=PositionEmbeddingType.rope_gpt_neox,
+                rope=RopeParams.from_config(config),
+            )
         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=True,
-            rotary_emb=QwenRotaryEmbedding(config),
             pos_embd_params=pos_embd_params,
             layer_idx=layer_idx,
             dtype=config.torch_dtype,
diff --git a/tensorrt_llm/_torch/models/modeling_qwen_moe.py b/tensorrt_llm/_torch/models/modeling_qwen_moe.py
index 9cdecd9ceb73..6286f6b1550d 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen_moe.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen_moe.py
@@ -18,7 +18,6 @@
 from ..modules.gated_mlp import GatedMLP
 from ..modules.linear import Linear, TensorParallelMode
 from ..modules.rms_norm import RMSNorm
-from ..modules.rotary_embedding import RotaryEmbedding
 from .modeling_utils import (DecoderModel, DecoderModelForCausalLM,
                              duplicate_kv_weight, register_auto_model)
 
@@ -101,22 +100,6 @@ def forward(
         return final_hidden_states.view(orig_shape)
 
 
-class QwenMoeRotaryEmbedding(RotaryEmbedding):
-
-    def __init__(
-        self,
-        config: Qwen2MoeConfig,
-        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 QwenMoeAttention(Attention):
 
     def __init__(
@@ -125,27 +108,23 @@ def __init__(
         layer_idx: Optional[int] = None,
     ):
         config = model_config.pretrained_config
-        if model_config.fuse_pos_embd:
-            if getattr(config, "rope_scaling", None) is not None:
-                pos_embd_params = PositionalEmbeddingParams(
-                    type=PositionEmbeddingType.from_string(
-                        config.rope_scaling["type"]),
-                    rope=RopeParams.from_config(config),
-                )
-            else:
-                pos_embd_params = PositionalEmbeddingParams(
-                    type=PositionEmbeddingType.rope_gpt_neox,
-                    rope=RopeParams.from_config(config),
-                )
+        if getattr(config, "rope_scaling", None) is not None:
+            pos_embd_params = PositionalEmbeddingParams(
+                type=PositionEmbeddingType.from_string(
+                    config.rope_scaling["type"]),
+                rope=RopeParams.from_config(config),
+            )
         else:
-            pos_embd_params = None
+            pos_embd_params = PositionalEmbeddingParams(
+                type=PositionEmbeddingType.rope_gpt_neox,
+                rope=RopeParams.from_config(config),
+            )
         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=True,
-            rotary_emb=QwenMoeRotaryEmbedding(config),
             pos_embd_params=pos_embd_params,
             layer_idx=layer_idx,
             dtype=config.torch_dtype,
diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index 3b0b51a10aa8..1870f62b04a4 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -6,8 +6,7 @@
 
 from tensorrt_llm.mapping import Mapping
 
-from ..attention_backend import (AttentionInputType, AttentionMetadata,
-                                 TrtllmAttention)
+from ..attention_backend import AttentionInputType, AttentionMetadata
 from ..attention_backend.interface import (PositionalEmbeddingParams,
                                            PredefinedAttentionMask)
 from ..attention_backend.utils import create_attention
@@ -49,7 +48,6 @@ def __init__(
         max_position_embeddings: int,
         bias: bool,
         pos_embd_params: Optional[PositionalEmbeddingParams] = None,
-        rotary_emb: Optional[RotaryEmbedding] = None,
         layer_idx: Optional[int] = None,
         dtype: torch.dtype = None,
         dense_bias: Optional[bool] = None,
@@ -132,7 +130,17 @@ def __init__(
         self.quant_config = config.get_quant_config()
         self.attn_backend = config.attn_backend
         self.pos_embd_params = pos_embd_params
-        self.rotary_emb = rotary_emb
+
+        self.enable_rope_fusion = self.attn_backend == "TRTLLM"
+        self.support_fused_qkv = self.attn_backend == "TRTLLM"
+        self.support_unfused_qkv = self.attn_backend != "TRTLLM"
+        self.rotary_emb = None
+        if not self.enable_rope_fusion and pos_embd_params is not None:
+            self.rotary_emb = RotaryEmbedding(
+                pos_embd_params.rope,
+                head_dim=self.head_dim,
+                is_neox=pos_embd_params.is_neox,
+            )
 
         # These two modules are mutually exclusive - either splitted_qkv_lora or fused_qkv_lora will be used,
         # but never both at the same time. splitted_qkv_lora handles Q,K,V separately while fused_qkv_lora
@@ -161,6 +169,14 @@ def create_weights(self):
             quant_config=self.quant_config,
         )
 
+    def convert_qkv(self, q, k, v):
+        if k is None and v is None and not self.support_fused_qkv:
+            q, k, v = q.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        elif k is not None and v is not None and not self.support_unfused_qkv:
+            qkv = torch.concat([q, k, v], dim=-1)
+            q, k, v = qkv, None, None
+        return q, k, v
+
     def forward(
         self,
         position_ids: Optional[torch.LongTensor],
@@ -174,9 +190,6 @@ def forward(
         **kwargs,
     ) -> torch.Tensor:
         qkv = self.qkv_proj(hidden_states)
-        is_fused_qkv = False
-        if isinstance(self.attn, TrtllmAttention):
-            is_fused_qkv = True
 
         if self.qk_norm is not None:
             # TODO: make this more efficient.
@@ -208,36 +221,25 @@ def forward(
             if qkv_lora is not None:
                 qkv = qkv + qkv_lora
 
-        if is_fused_qkv:
-            if self.pos_embd_params is None and position_ids is not None:
-                q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size],
-                                    dim=-1)
-                q, k = self.rotary_emb(position_ids, [q, k], attn_metadata)
-                qkv = torch.concat([q, k, v], dim=-1)
-
-            out_scale = None
-            if self.o_proj.has_fp8_qdq or self.o_proj.has_nv_fp4 or self.o_proj.has_fp8_block_scales:
-                out_scale = self.o_proj.inv_input_scale
-            attn_output = self.attn.forward(qkv,
-                                            None,
-                                            None,
-                                            attn_metadata,
-                                            out_scale=out_scale,
-                                            attention_mask=attention_mask,
-                                            mrope_config=mrope_config)
-        else:
+        q, k, v = qkv, None, None
+
+        if self.rotary_emb is not None and position_ids is not None:
             q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size],
                                 dim=-1)
+            q, k = self.rotary_emb(position_ids, [q, k])
 
-            if self.pos_embd_params is None and position_ids is not None:
-                q, k = self.rotary_emb(position_ids, [q, k], attn_metadata)
+        out_scale = None
+        if self.o_proj.has_fp8_qdq or self.o_proj.has_nv_fp4 or self.o_proj.has_fp8_block_scales:
+            out_scale = self.o_proj.inv_input_scale
 
-            attn_output = self.attn.forward(q.contiguous(),
-                                            k.contiguous(),
-                                            v.contiguous(),
-                                            attn_metadata,
-                                            attention_mask=attention_mask,
-                                            mrope_config=mrope_config)
+        q, k, v = self.convert_qkv(q, k, v)
+        attn_output = self.attn.forward(q,
+                                        k,
+                                        v,
+                                        attn_metadata,
+                                        out_scale=out_scale,
+                                        attention_mask=attention_mask,
+                                        mrope_config=mrope_config)
 
         attn_output = self.o_proj(attn_output,
                                   all_reduce_params=all_reduce_params)
@@ -268,7 +270,6 @@ def __init__(
         bias: bool,
         aux_stream: Optional[torch.cuda.Stream] = None,
         pos_embd_params: Optional[PositionalEmbeddingParams] = None,
-        rotary_emb: Optional[RotaryEmbedding] = None,
         layer_idx: Optional[int] = None,
         dtype: torch.dtype = None,
         dense_bias: Optional[bool] = None,
@@ -305,6 +306,7 @@ def __init__(
             raise ValueError(
                 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                 f" and `num_heads`: {self.num_heads}).")
+        assert pos_embd_params is not None, "pos_embd_params must be provided in MLA"
 
         # tensor parallel
         config = config or ModelConfig()
@@ -498,7 +500,7 @@ def yarn_get_mscale(scale=1, mscale=1):
             v_head_dim=self.kv_lora_rank,
             predicted_tokens_per_seq=self.predicted_tokens_per_seq,
         )
-        self.rotary_emb = rotary_emb
+
         self.aux_stream = aux_stream
         self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
 
diff --git a/tensorrt_llm/_torch/modules/rotary_embedding.py b/tensorrt_llm/_torch/modules/rotary_embedding.py
index 2101e0732b36..c600e34128cf 100644
--- a/tensorrt_llm/_torch/modules/rotary_embedding.py
+++ b/tensorrt_llm/_torch/modules/rotary_embedding.py
@@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
-                    else:
-                        # TODO(qijun): support apply_llama31_rope_inplace
-                        raise ValueError(
-                            f'unsupported rope type {self.rope_type}')
-                        # rope_scale = self.config.rope_scaling
-                        # flashinfer.apply_llama31_rope_inplace(
-                        #     q,
-                        #     k,
-                        #     attn_metadata.qo_indptr,
-                        #     attn_metadata.cached_token_lens,
-                        #     rope_theta=rope_theta,
-                        #     rope_scale=rope_scale["factor"],
-                        #     low_freq_factor=rope_scale["low_freq_factor"],
-                        #     high_freq_factor=rope_scale["high_freq_factor"],
-                        #     old_context_len=rope_scale[
-                        #         "original_max_position_embeddings"],
-                        # )
-
-                    q = q.view(seq_len, -1)
-                    k = k.view(seq_len, -1)
-                    return [q, k]
-
-        if use_gptj_style_rope:
-            raise ValueError(
-                "gptj style RoPE has to go through flashinfer route for correct results."
+        if IS_FLASHINFER_AVAIABLE and len(targets) == 2:
+            from ..custom_ops import \
+                flashinfer_apply_rope_with_cos_sin_cache_inplace
+            q = targets[0]
+            k = targets[1]
+            flashinfer_apply_rope_with_cos_sin_cache_inplace(
+                position_ids.view(-1),
+                q,
+                k,
+                self.head_dim,
+                self.rotary_cos_sin.view(self.max_positions, -1),
+                self.is_neox,
             )
+            return [q, k]
+
         # it is assumed all targets are of the same rank
         q_or_k = targets[0]
         remove_input_padding = (len(q_or_k.size()) == 2)
 
-        # Core RoPE block
-        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(
-            position_ids.shape[0], -1, 1)
-        position_ids_expanded = position_ids[:, None, :].float()
-        # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
-        device_type = q_or_k.device.type
-        device_type = device_type if isinstance(
-            device_type, str) and device_type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):
-            freqs = (inv_freq_expanded.float()
-                     @ position_ids_expanded.float()).transpose(1, 2)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos()
-            sin = emb.sin()
-
-        # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
-        cos = cos * self.attention_scaling
-        sin = sin * self.attention_scaling
-
-        cos = cos.to(dtype=q_or_k.dtype)
-        sin = sin.to(dtype=q_or_k.dtype)
+        cos_sin = self.rotary_cos_sin[position_ids.view(-1)]
+        cos, sin = cos_sin[:, 0, :], cos_sin[:, 1, :]
+        cos = cos.to(dtype=q_or_k.dtype).unsqueeze(0)
+        sin = sin.to(dtype=q_or_k.dtype).unsqueeze(0)
 
         if remove_input_padding:
             bsz = 1
@@ -218,7 +67,10 @@ def forward(
         def rope_target(target):
             target = target.view(bsz, seq_len, -1,
                                  self.head_dim).transpose(1, 2)
-            target = RotaryEmbedding.apply_rotary_pos_emb(target, cos, sin)
+            target = RotaryEmbedding.apply_rotary_pos_emb(target,
+                                                          cos,
+                                                          sin,
+                                                          is_neox=self.is_neox)
             target = target.transpose(1, 2).contiguous()
             if remove_input_padding:
                 target = target.view(seq_len, -1)
@@ -232,7 +84,8 @@ def rope_target(target):
     def apply_rotary_pos_emb(q_or_k: torch.Tensor,
                              cos: torch.Tensor,
                              sin: torch.Tensor,
-                             unsqueeze_dim: int = 1) -> torch.Tensor:
+                             unsqueeze_dim: int = 1,
+                             is_neox: bool = True) -> torch.Tensor:
         """Applies Rotary Position Embedding to the query and key tensors.
 
         Args:
@@ -246,22 +99,29 @@ def apply_rotary_pos_emb(q_or_k: torch.Tensor,
                 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
                 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
                 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+            is_neox (bool): Whether to use Neox style RoPE, True by default.
         Returns:
             `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
         """
         cos = cos.unsqueeze(unsqueeze_dim)
         sin = sin.unsqueeze(unsqueeze_dim)
 
-        rot_dim = cos.shape[-1]
+        rot_dim = cos.shape[-1] * 2
         # If q_or_k_pass is empty, rotary pos embedding is applied to all tensor
         q_or_k, q_or_k_pass = q_or_k[..., :rot_dim], q_or_k[..., rot_dim:]
 
-        embed = (q_or_k * cos) + (RotaryEmbedding.rotate_half(q_or_k) * sin)
-        return torch.cat((embed, q_or_k_pass), dim=-1)
+        if is_neox:
+            x1 = q_or_k[..., :q_or_k.shape[-1] // 2]
+            x2 = q_or_k[..., q_or_k.shape[-1] // 2:]
+        else:
+            x1 = q_or_k[..., ::2]
+            x2 = q_or_k[..., 1::2]
 
-    @staticmethod
-    def rotate_half(x: torch.Tensor) -> torch.Tensor:
-        """Rotates half the hidden dims of the input."""
-        x1 = x[..., :x.shape[-1] // 2]
-        x2 = x[..., x.shape[-1] // 2:]
-        return torch.cat((-x2, x1), dim=-1)
+        o1 = x1 * cos - x2 * sin
+        o2 = x2 * cos + x1 * sin
+
+        if is_neox:
+            return torch.cat((o1, o2, q_or_k_pass), dim=-1)
+        else:
+            embed = torch.stack((o1, o2), dim=-1).flatten(-2)
+            return torch.cat((embed, q_or_k_pass), dim=-1)
diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py
index 8dbe585beca0..19282837eea6 100644
--- a/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py
+++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_nas.py
@@ -247,7 +247,9 @@ def test_nemotron_nas_sanity(self):
         if config_dict["num_hidden_layers"] <= 0:
             self.skipTest("Insufficient memory for a single NemotronNAS layer")
         nemotron_nas_config = AutoConfig.from_pretrained(
-            "nvidia/Llama-3_1-Nemotron-51B-Instruct", trust_remote_code=True)
+            llm_models_root() / "nemotron-nas/Llama-3_1-Nemotron-51B-Instruct",
+            trust_remote_code=True,
+        )
         nemotron_nas_config = nemotron_nas_config.from_dict(config_dict)
 
         dtype = nemotron_nas_config.torch_dtype
@@ -370,7 +372,9 @@ def test_nemotron_nas_allclose_to_hf(self, scenario: Scenario) -> None:
         if config_dict["num_hidden_layers"] <= 0:
             self.skipTest("Insufficient memory for a single NemotronNAS layer")
         nemotron_nas_config = AutoConfig.from_pretrained(
-            "nvidia/Llama-3_1-Nemotron-51B-Instruct", trust_remote_code=True)
+            llm_models_root() / "nemotron-nas/Llama-3_1-Nemotron-51B-Instruct",
+            trust_remote_code=True,
+        )
         nemotron_nas_config = nemotron_nas_config.from_dict(config_dict)
         dtype = nemotron_nas_config.torch_dtype
         device = torch.device('cuda')
diff --git a/tests/unittest/_torch/multi_gpu/test_star_attention.py b/tests/unittest/_torch/multi_gpu/test_star_attention.py
index fdab1067199c..083185ec7979 100644
--- a/tests/unittest/_torch/multi_gpu/test_star_attention.py
+++ b/tests/unittest/_torch/multi_gpu/test_star_attention.py
@@ -1,139 +1,20 @@
 import json
 import os
 
-import modelopt.torch.quantization as mtq
 import pytest
 import torch
-from modelopt.torch.export import export_hf_checkpoint
-from transformers import AutoModelForCausalLM, AutoTokenizer
 from utils.llm_data import llm_models_root
 
-from tensorrt_llm import SamplingParams, logger
+from tensorrt_llm import SamplingParams
 from tensorrt_llm._torch import LLM
 from tensorrt_llm._torch.pyexecutor.config import PyTorchConfig
 from tensorrt_llm.llmapi import KvCacheConfig
 from tensorrt_llm.llmapi.utils import get_total_gpu_memory
 from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
-from tensorrt_llm.quantization.quantize_by_modelopt import get_calib_dataloader
 
 MAX_SEQ_LEN = 4096 + 1024
 
 
-def is_model_on_gpu(model) -> bool:
-    """Check if the model is fully loaded on GPUs."""
-    return all("cuda" in str(param.device) for param in model.parameters())
-
-
-def get_model(ckpt_path, device="cuda"):
-    """Load and return a model from the checkpoint path."""
-    logger.info(f"Initializing model from {ckpt_path}")
-
-    # Forcibly converting the model precision between bf16 and fp16 may introduce accuracy drop
-    model_kwargs = {"torch_dtype": "auto"}
-
-    model = AutoModelForCausalLM.from_pretrained(
-        ckpt_path,
-        device_map=device,
-        **model_kwargs,
-        trust_remote_code=True,
-    )
-    model.eval()
-
-    return model
-
-
-def get_tokenizer(ckpt_path, max_seq_len=MAX_SEQ_LEN, model_type=None):
-    """Load and return a tokenizer from the checkpoint path."""
-    logger.info(f"Initializing tokenizer from {ckpt_path}")
-    tokenizer = AutoTokenizer.from_pretrained(
-        ckpt_path,
-        model_max_length=max_seq_len,
-        padding_side="left",
-        trust_remote_code=True,
-    )
-
-    if tokenizer.pad_token != "<unk>":
-        tokenizer.pad_token = tokenizer.eos_token
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    assert tokenizer.pad_token is not None, f"Pad token for {model_type} cannot be set!"
-
-    return tokenizer
-
-
-def calibrate_loop(model, calib_dataloader):
-    """Calibrate the model with the calibration dataloader."""
-    if calib_dataloader is None:
-        return
-
-    for idx, data in enumerate(calib_dataloader):
-        logger.debug(f"Calibrating batch {idx}")
-        data = data["input_ids"].to(model.device)
-        model(data)
-
-
-# Come from internal/examples/llm_ptq/hf_model_export.py in modelopt repo
-def quantize(model_path,
-             output_dir,
-             quant_format,
-             enable_quant_kv_cache,
-             calib_samples=512):
-    model = get_model(model_path, device="cuda")
-    tokenizer = get_tokenizer(model_path)
-
-    device = model.device if not hasattr(model, "model") else model.model.device
-
-    # Calibrate the model
-    calib_dataloader = get_calib_dataloader(
-        dataset_name_or_dir=str(llm_models_root() / "datasets/cnn_dailymail"),
-        tokenizer=tokenizer,
-        batch_size=1,
-        calib_size=calib_samples,
-        device=device,
-    )
-
-    if quant_format == "fp8":
-        quant_cfg = mtq.FP8_DEFAULT_CFG
-    elif quant_format == "fp8_pc":
-        quant_cfg = mtq.FP8_PER_CHANNEL_CFG
-    elif quant_format == "int4_awq":
-        quant_cfg = mtq.INT4_AWQ_CFG
-    elif quant_format == "nvfp4":
-        quant_cfg = mtq.NVFP4_DEFAULT_CFG
-    else:
-        raise ValueError(f"Unsupported quantization format: {quant_format}")
-
-    quant_cfg["quant_cfg"]["*output_quantizer"] = {  # type: ignore[index]
-        "num_bits": (4, 3),
-        "axis": None,
-        "enable": enable_quant_kv_cache,
-    }
-
-    model = mtq.quantize(
-        model,
-        quant_cfg,
-        forward_loop=lambda: calibrate_loop(model, calib_dataloader),
-    )
-    model.to("cpu")
-
-    # Post-process: export the quantized checkpoint with fused and packed weights, and save the quantization config
-    post_state_dict, hf_quant_config, _ = export_hf_checkpoint(
-        model,
-        dtype=model.config.torch_dtype,
-        export_dir=output_dir,
-    )
-
-    with open(f"{output_dir}/hf_quant_config.json", "w") as file:
-        json.dump(hf_quant_config, file, indent=4)
-
-    # Save the post-processed model and configs with the HF API
-    model.save_pretrained(output_dir, state_dict=post_state_dict)
-
-    # Save the tokenizer
-    tokenizer.save_pretrained(output_dir)
-
-
 @pytest.mark.parametrize("backend", ["pytorch"])
 @pytest.mark.parametrize("model_name",
                          ["llama-models-v3/Llama-3-8B-Instruct-Gradient-1048k"],


Reviews:
user: login: litaotju
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-16T05:18:03Z
user: login: chang-l
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T23:59:28Z
user: login: yuxianq
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T03:03:58Z
user: login: chang-l
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T03:23:51Z
user: login: vanshilshah97
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T03:39:54Z
user: login: chang-l
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T03:44:57Z
user: login: yuxianq
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T04:08:51Z
user: login: yuxianq
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T04:16:38Z
user: login: chang-l
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T04:23:32Z
user: login: yuxianq
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T04:38:42Z
user: login: yuxianq
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T04:40:29Z
user: login: chang-l
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T04:47:27Z
user: login: chang-l
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T23:30:22Z
user: login: hlu1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T23:40:01Z
user: login: chang-l
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-18T01:18:24Z

Review comments:
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: chang-l
body: Just to clarify — here when `interleave=True` (for llama4), can we still use `flashinfer_apply_rope_with_cos_sin_cache_inplace`?
It looks like in L25, you enforced `interleave=False`. Are these referring to the same interleave definition? @yuxianq 
created_at: 2025-04-16T23:59:28Z
author_association: COLLABORATOR
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: yuxianq
body: llama4 uses `is_neox=False` while llama3 uses `is_neox=True`. `interleave=True/False` is for `rope_cos_sin` layout. For fused rope in trtllm attention backend, we always use `interleave=True`. For unfused rope implemented by flashinfer or native PyTorch ops, we always use `interleave=False`.
created_at: 2025-04-17T03:03:58Z
author_association: COLLABORATOR
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: chang-l
body: Thanks @yuxianq.
Can I understand it as: whether interleave is True or not would not impact the attention output itself — as long as we follow the expected format per backend (i.e., use interleaved format for TRTLLM and non-interleaved for FlashInfer)?

Also, for my own education, what does `is_neox` represent for?
created_at: 2025-04-17T03:23:50Z
author_association: COLLABORATOR
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: vanshilshah97
body: Hi @yuxianq 
Thanks for the commit. I think i was looking for this for our use case.
I was discussing this with @chang-l today afternoon regarding this.

@chang-l I think `is_neox` refers to the type of RoPE embeddings that TRT-LLM supports that were found in the respective  [model implementation](https://nvidia.github.io/TensorRT-LLM/advanced/gpt-attention.html#rotary-positional-embedding-rope).
@yuxianq one question
 For our use-case to support trt-llm attention for llama4 scout . We want to apply qk_norm on Q and K after RoPE is applied.
I was planning to use the pytorch implementation for generating unfused rope then apply qk norm on pytorch level and pass these things to trtllm attention kernels thorugh [here](https://github.com/NVIDIA/TensorRT-LLM/blob/b2fb0fe843ca6ed09be691b9219747da4167d3c2/tensorrt_llm/_torch/attention_backend/trtllm.py#L315) with `pos_embed_params=None`, according to my understanding as we have already done RoPE. Do you forsee any issue for that ? Accuracy wise or format wise ?
created_at: 2025-04-17T03:39:54Z
author_association: COLLABORATOR
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: chang-l
body: @vanshilshah97  It seems w/ this PR, unfused rope (using flashinfer kernel) can work outside together with trtllm attention backend. I will update the code path in this PR: https://github.com/NVIDIA/TensorRT-LLM/pull/3491 later tonight, then you can cherry-pick it to your llama4 branch.
created_at: 2025-04-17T03:44:57Z
author_association: COLLABORATOR
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: yuxianq
body: @vanshilshah97 The main issue is that the current MLA kernel cannot skip rope, see https://github.com/NVIDIA/TensorRT-LLM/blob/515dd0d78fe8f5f86f8a655fc7bac36a6f8da7ac/cpp/tensorrt_llm/kernels/mlaKernels.cu#L451.
@hlu1 Can you help to estimate the effort needed to skip rope in our mla kernels?
created_at: 2025-04-17T04:08:51Z
author_association: COLLABORATOR
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: yuxianq
body: @chang-l `interleave=True` means the `rope_cos_sin` layout is [cos0, sin0, cos1, sin1, ...], cos and sin are interleaved; `interleave=False` uses [cos0, cos1, ..., sin0, sin1, ...] layout.
`is_neox` is about q/k layout. When calculating rope, we need to split q/k into two tensors and use them to multiple with cos/sin. `is_neox=True` means we split q/k into the first half part (`q[..., :head_dim//2]`) and the second half part (`q[..., head_dim//2:]`), while `is_neox=False` means we split q/k into the even part (`q[..., ::2]`) and the odd part (`q[..., 1::2]`).
created_at: 2025-04-17T04:16:37Z
author_association: COLLABORATOR
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: chang-l
body: @yuxianq  Got it. Thanks Yuxian.  
So `interleave=True` must come with `is_neox=False`; and `interleave=False` must come with `is_neox=True`

If so, like you mentioned:
> For unfused rope implemented by flashinfer or native PyTorch ops, we always use interleave=False

then we cannot use the existing unfused rope from flashinfer or native pytorch op for llama4 due to format issue (interleave = True/is_neox=False for llama4)? Am i missing something?
created_at: 2025-04-17T04:23:32Z
author_association: COLLABORATOR
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: yuxianq
body: No, `is_neox` is related to algorithm, it changes the rope output; `interleave` is related to implementation, it does not change the rope output. I list all the combinations:
- `is_neox=False` and `interleave=False`: for rope_gptj + flashinfer/native op
- `is_neox=False` and `interleave=True`: for rope_gptj + trtllm fused rope
- `is_neox=True` and `interleave=False`: for rope_gpt_neox + flashinfer/native op
- `is_neox=True` and `interleave=True`: for rope_gpt_neox + trtllm fused rope

created_at: 2025-04-17T04:38:42Z
author_association: COLLABORATOR
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: yuxianq
body: @hlu1 Sorry for misunderstanding, this issue is related to llama4, not deepseek, so MLA without rope fusion is unnecessary here. 
created_at: 2025-04-17T04:40:28Z
author_association: COLLABORATOR
diff_hunk: @@ -1,213 +1,62 @@
-import math
-from typing import List, Literal, Optional, Tuple
+from typing import List
 
 import torch
 from torch import nn
-from transformers import LlamaConfig, PretrainedConfig
 
-from ..attention_backend import IS_FLASHINFER_AVAIABLE, AttentionMetadata
-
-
-def compute_default_rope_parameters(
-    config: PretrainedConfig,
-    device: Optional[torch.device] = None,
-) -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies according to the original RoPE implementation
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-        seq_len (`int`, *optional*):
-            The current sequence length. Unused for this type of RoPE.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
-    """
-    base = config.rope_theta
-    partial_rotary_factor = config.partial_rotary_factor if hasattr(
-        config, "partial_rotary_factor") else 1.0
-    dim = int((config.hidden_size // config.num_attention_heads) *
-              partial_rotary_factor)
-
-    attention_factor = 1.0  # Unused in this type of RoPE
-
-    # Compute the inverse frequencies
-    inv_freq = 1.0 / (base**(
-        torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
-    return inv_freq, attention_factor
-
-
-def compute_llama3_parameters(
-        config: LlamaConfig,
-        device: "torch.device") -> "Tuple[torch.Tensor, float]":
-    """
-    Computes the inverse frequencies for llama 3.1.
-
-    Args:
-        config ([`~transformers.PretrainedConfig`]):
-            The model configuration.
-        device (`torch.device`):
-            The device to use for initialization of the inverse frequencies.
-    Returns:
-        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
-        post-processing scaling factor applied to the computed cos/sin.
-    """
-    # Gets the default RoPE parameters
-    inv_freq, attention_factor = compute_default_rope_parameters(config, device)
-
-    factor: int = config.rope_scaling[
-        "factor"]  # `8` in the original implementation
-    low_freq_factor: int = config.rope_scaling[
-        "low_freq_factor"]  # `1` in the original implementation
-    high_freq_factor: int = config.rope_scaling[
-        "high_freq_factor"]  # `4` in the original implementation
-    old_context_len: int = config.rope_scaling[
-        "original_max_position_embeddings"]  # `8192` in the original implementation
-
-    low_freq_wavelen = old_context_len / low_freq_factor
-    high_freq_wavelen = old_context_len / high_freq_factor
-
-    wavelen = 2 * math.pi / inv_freq
-    # wavelen < high_freq_wavelen: do nothing
-    # wavelen > low_freq_wavelen: divide by factor
-    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor,
-                                 inv_freq)
-    # otherwise: interpolate between the two, using a smooth factor
-    smooth_factor = (old_context_len / wavelen -
-                     low_freq_factor) / (high_freq_factor - low_freq_factor)
-    smoothed_inv_freq = (
-        1 - smooth_factor
-    ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
-    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen
-                                                        > low_freq_wavelen)
-    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
-                                 inv_freq_llama)
-
-    return inv_freq_llama, attention_factor
+from ..attention_backend import IS_FLASHINFER_AVAIABLE
+from ..attention_backend.interface import RopeParams
 
 
 class RotaryEmbedding(nn.Module):
 
     def __init__(
         self,
-        config: PretrainedConfig,
+        rope_params: RopeParams,
         *,
         head_dim: int,
-        num_attention_heads: int,
-        max_position_embeddings: int,
-        device: Optional[torch.device] = None,
-        rope_type: Literal['default', 'llama3'] = "default",
+        is_neox: bool = True,
     ):
         super().__init__()
-        self.config = config
+        self.rope_params = rope_params
         self.head_dim = head_dim
-        self.num_attention_heads = num_attention_heads
-        self.rope_type = rope_type
-        self.max_seq_len_cached = max_position_embeddings
-        self.original_max_seq_len = max_position_embeddings
-
-        rope_init_fn = compute_llama3_parameters if rope_type == "llama3" else compute_default_rope_parameters
-
-        inv_freq, self.attention_scaling = rope_init_fn(config, device)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-        rope_percentage = (getattr(config, 'rotary_pct', None)
-                           or getattr(config, 'partial_rotary_factor', None)
-                           or 1.0)
-        self.rotary_dim = (getattr(config, 'rotary_dim', None)
-                           or getattr(config, 'rotary_emb_base', None)
-                           or int(head_dim * rope_percentage))
+        self.is_neox = is_neox
+        self.max_positions = rope_params.max_positions
+        self.rotary_cos_sin = rope_params.create_rope_const_params(
+            interleave=False)[1].reshape(rope_params.max_positions, 2, -1)
 
     def forward(
         self,
         position_ids: torch.Tensor,
         targets: List[torch.Tensor],
-        attn_metadata: Optional[AttentionMetadata] = None
     ) -> List[torch.Tensor]:
         """
         Apply RoPE to any number of target tensors with the same position_ids.
         This is useful if q_len = k_len, in which case we may RoPE q and k with the same cos and sin values.
         However, if k is cached without positional embedding, we need to apply rope to q and k with different values, so we need a separate call for each.
         """
-        use_gptj_style_rope = self.config.model_type == "llama4_text"
-        if IS_FLASHINFER_AVAIABLE:
-            from ..attention_backend import FlashInferAttentionMetadata
-            if attn_metadata is not None:
-                from ..attention_backend import (FlashInferAttentionMetadata,
-                                                 StarAttentionMetadata)
-                if isinstance(attn_metadata, StarAttentionMetadata):
-                    pass
-                elif isinstance(attn_metadata, FlashInferAttentionMetadata):
-                    from ..custom_ops import flashinfer_apply_rope_inplace
-                    assert len(targets) == 2
-                    q = targets[0]
-                    seq_len = q.size()[0]
-                    q = targets[0].view(seq_len, -1, self.head_dim)
-                    k = targets[1].view(seq_len, -1, self.head_dim)
-                    rope_theta = self.config.rope_theta
-                    if self.rope_type in ['default', 'llama3']:
-                        flashinfer_apply_rope_inplace(
-                            q,
-                            k,
-                            attn_metadata.qo_indptr,
-                            attn_metadata.cached_token_lens,
-                            rope_theta=rope_theta,
-                            rotary_dim=self.rotary_dim,
-                            interleave=use_gptj_style_rope)
path: tensorrt_llm/_torch/modules/rotary_embedding.py
user: login: chang-l
body: @yuxianq  Thanks! Thats very helpful
created_at: 2025-04-17T04:47:27Z
author_association: COLLABORATOR
diff_hunk: @@ -208,36 +221,25 @@ def forward(
             if qkv_lora is not None:
                 qkv = qkv + qkv_lora
 
-        if is_fused_qkv:
-            if self.pos_embd_params is None and position_ids is not None:
-                q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size],
-                                    dim=-1)
-                q, k = self.rotary_emb(position_ids, [q, k], attn_metadata)
-                qkv = torch.concat([q, k, v], dim=-1)
-
-            out_scale = None
-            if self.o_proj.has_fp8_qdq or self.o_proj.has_nv_fp4 or self.o_proj.has_fp8_block_scales:
-                out_scale = self.o_proj.inv_input_scale
-            attn_output = self.attn.forward(qkv,
-                                            None,
-                                            None,
-                                            attn_metadata,
-                                            out_scale=out_scale,
-                                            attention_mask=attention_mask,
-                                            mrope_config=mrope_config)
-        else:
+        q, k, v = qkv, None, None
+
+        if self.rotary_emb is not None and position_ids is not None:
             q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size],
                                 dim=-1)
+            q, k = self.rotary_emb(position_ids, [q, k])
 
-            if self.pos_embd_params is None and position_ids is not None:
-                q, k = self.rotary_emb(position_ids, [q, k], attn_metadata)
+        out_scale = None
+        if self.o_proj.has_fp8_qdq or self.o_proj.has_nv_fp4 or self.o_proj.has_fp8_block_scales:
+            out_scale = self.o_proj.inv_input_scale
 
-            attn_output = self.attn.forward(q.contiguous(),
-                                            k.contiguous(),
-                                            v.contiguous(),
-                                            attn_metadata,
-                                            attention_mask=attention_mask,
-                                            mrope_config=mrope_config)
+        q, k, v = self.convert_qkv(q, k, v)
+        attn_output = self.attn.forward(q,
path: tensorrt_llm/_torch/modules/attention.py
user: login: chang-l
body: @yuxianq  
Another question: do we need to include `.contiguous()` after q, k, v when passing to `attn.forward`  after applying rope

cc. @hlu1 
created_at: 2025-04-17T23:30:22Z
author_association: COLLABORATOR
diff_hunk: @@ -208,36 +221,25 @@ def forward(
             if qkv_lora is not None:
                 qkv = qkv + qkv_lora
 
-        if is_fused_qkv:
-            if self.pos_embd_params is None and position_ids is not None:
-                q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size],
-                                    dim=-1)
-                q, k = self.rotary_emb(position_ids, [q, k], attn_metadata)
-                qkv = torch.concat([q, k, v], dim=-1)
-
-            out_scale = None
-            if self.o_proj.has_fp8_qdq or self.o_proj.has_nv_fp4 or self.o_proj.has_fp8_block_scales:
-                out_scale = self.o_proj.inv_input_scale
-            attn_output = self.attn.forward(qkv,
-                                            None,
-                                            None,
-                                            attn_metadata,
-                                            out_scale=out_scale,
-                                            attention_mask=attention_mask,
-                                            mrope_config=mrope_config)
-        else:
+        q, k, v = qkv, None, None
+
+        if self.rotary_emb is not None and position_ids is not None:
             q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size],
                                 dim=-1)
+            q, k = self.rotary_emb(position_ids, [q, k])
 
-            if self.pos_embd_params is None and position_ids is not None:
-                q, k = self.rotary_emb(position_ids, [q, k], attn_metadata)
+        out_scale = None
+        if self.o_proj.has_fp8_qdq or self.o_proj.has_nv_fp4 or self.o_proj.has_fp8_block_scales:
+            out_scale = self.o_proj.inv_input_scale
 
-            attn_output = self.attn.forward(q.contiguous(),
-                                            k.contiguous(),
-                                            v.contiguous(),
-                                            attn_metadata,
-                                            attention_mask=attention_mask,
-                                            mrope_config=mrope_config)
+        q, k, v = self.convert_qkv(q, k, v)
+        attn_output = self.attn.forward(q,
path: tensorrt_llm/_torch/modules/attention.py
user: login: hlu1
body: It depends on if the attention kernel requires contiguous inputs or strided inputs. A common practice is for the attention kernel to call `.contiguous()` inside if contiguous inputs are required.
created_at: 2025-04-17T23:40:01Z
author_association: COLLABORATOR
diff_hunk: @@ -208,36 +221,25 @@ def forward(
             if qkv_lora is not None:
                 qkv = qkv + qkv_lora
 
-        if is_fused_qkv:
-            if self.pos_embd_params is None and position_ids is not None:
-                q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size],
-                                    dim=-1)
-                q, k = self.rotary_emb(position_ids, [q, k], attn_metadata)
-                qkv = torch.concat([q, k, v], dim=-1)
-
-            out_scale = None
-            if self.o_proj.has_fp8_qdq or self.o_proj.has_nv_fp4 or self.o_proj.has_fp8_block_scales:
-                out_scale = self.o_proj.inv_input_scale
-            attn_output = self.attn.forward(qkv,
-                                            None,
-                                            None,
-                                            attn_metadata,
-                                            out_scale=out_scale,
-                                            attention_mask=attention_mask,
-                                            mrope_config=mrope_config)
-        else:
+        q, k, v = qkv, None, None
+
+        if self.rotary_emb is not None and position_ids is not None:
             q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size],
                                 dim=-1)
+            q, k = self.rotary_emb(position_ids, [q, k])
 
-            if self.pos_embd_params is None and position_ids is not None:
-                q, k = self.rotary_emb(position_ids, [q, k], attn_metadata)
+        out_scale = None
+        if self.o_proj.has_fp8_qdq or self.o_proj.has_nv_fp4 or self.o_proj.has_fp8_block_scales:
+            out_scale = self.o_proj.inv_input_scale
 
-            attn_output = self.attn.forward(q.contiguous(),
-                                            k.contiguous(),
-                                            v.contiguous(),
-                                            attn_metadata,
-                                            attention_mask=attention_mask,
-                                            mrope_config=mrope_config)
+        q, k, v = self.convert_qkv(q, k, v)
+        attn_output = self.attn.forward(q,
path: tensorrt_llm/_torch/modules/attention.py
user: login: chang-l
body: I see.. do you know whether flashinfer and trtllm attn backend support strided input? AFAIK, there seems no `contiguous()` call inside both attn kernels? (I may miss something)
created_at: 2025-04-18T01:18:24Z
author_association: COLLABORATOR

PR comments:
user: login: yuxianq
created_at: 2025-04-14T04:19:24Z
body: /bot run --disable-fail-fast
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-14T04:25:04Z
body: [PR_Github #2109](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/2109/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-14T11:15:56Z
body: [PR_Github #2109](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/2109/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1538](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1538/) completed with status: 'FAILURE'

author_association: COLLABORATOR
user: login: yuxianq
created_at: 2025-04-15T07:26:18Z
body: /bot run --disable-fail-fast
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-15T07:32:01Z
body: [PR_Github #2291](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/2291/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-15T11:24:51Z
body: [PR_Github #2291](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/2291/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1647](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1647/) completed with status: 'FAILURE'

author_association: COLLABORATOR
user: login: yuxianq
created_at: 2025-04-15T11:38:16Z
body: /bot run --disable-fail-fast
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-15T11:43:56Z
body: [PR_Github #2332](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/2332/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-15T16:36:06Z
body: [PR_Github #2332](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/2332/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #1679](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/1679/) completed with status: 'SUCCESS'

author_association: COLLABORATOR
user: login: yuxianq
created_at: 2025-04-16T05:36:34Z
body: /bot reuse-pipeline
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-16T05:42:08Z
body: [PR_Github #2422](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/2422/) [ reuse-pipeline ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-04-16T05:48:42Z
body: [PR_Github #2422](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/2422/) [ reuse-pipeline ] completed with state `SUCCESS`
Reusing [PR_Github #2332](https://prod.blsm.nvidia.com/sw-tensorrt-top-1/job/LLM/job/helpers/job/PR_Github/2332/) for commit `54a3b4d`

author_association: COLLABORATOR