kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-02-04 02:02:01 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None][feat] Support new Transformers RoPE configuration format (#10636)

Browse Source 
Signed-off-by: lkm2835 <lkm2835@gmail.com>
This commit is contained in:
Kyungmin Lee 2026-01-14 19:41:27 +09:00 committed by GitHub
parent e9817461ba
commit 25148d3fee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
tensorrt_llm/_torch/attention_backend/interface.py
View File

@ -24,6 +24,14 @@ from ..metadata import KVCacheParams
from ..pyexecutor.resource_manager import KVCacheManager
from ..utils import get_model_extra_attrs
try:
# Transformers v5
from transformers.configuration_utils import ALLOWED_ATTENTION_LAYER_TYPES
except ImportError:
# Transformers v4
from transformers.configuration_utils import \
ALLOWED_LAYER_TYPES as ALLOWED_ATTENTION_LAYER_TYPES
@dataclass
class AttentionRuntimeFeatures:
@ -448,6 +456,13 @@ class RopeParams:
def from_config(config) -> "RopeParams":
rope_params = RopeParams()
hf_rope_parameters = getattr(config, 'rope_parameters', None)
if hf_rope_parameters is not None:
assert not set(hf_rope_parameters.keys()).issubset(
ALLOWED_ATTENTION_LAYER_TYPES), (
"Per-layer-type RoPE configuration is not supported yet.")
config.update(hf_rope_parameters)
# get rotary parameters.
hidden_size = config.hidden_size
num_attention_heads = config.num_attention_heads