Skip to content

Commit 15f8231

Browse files
vbaddiquic-rishinr
vbaddi
and
quic-rishinr
authored
feat(rope_fix): Hoist layer-invariant RoPE indexing out of decoder subfunctions for cached text models (quic#928)
## Summary This change moves layer-invariant RoPE cos/sin indexing out of repeated decoder-layer subfunctions and into model-level forward paths. For cached decoder models, we were repeatedly doing: ``` cos = cos[position_ids].unsqueeze(1) sin = sin[position_ids].unsqueeze(1) ``` inside each decoder attention block. With ONNX subfunctions enabled, that indexing becomes part of the exported repeated subfunction body and contributes to the on-device regression we observed after the single-subfunction Rope Fix work quic#880 . This patch hoists that work once per forward pass and passes the already-shaped cos/sin tensors into each decoder layer. ## What changed Applied the refactor to the applicable QEff model families that thread static cached RoPE tensors through repeated decoder layers, including: - Llama - Llama SwiftKV - Gemma - Gemma2 - Mistral - Falcon - GPT-OSS - Granite - GraniteMoE - Mllama text path - Mixtral - Olmo2 - Phi3 - Qwen2 - Qwen3 - Qwen3 MoE - Qwen2.5 VL text path - Qwen3 VL text path - Qwen3 VL MoE text path For the Qwen VL text towers, the same idea is applied to the indexed/interleaved MRoPE preparation: the already-indexed cos/sin tensors are prepared once before the decoder-layer loop and reused across layers. ## Tests Added a TinyLlama regression test to assert that export with subfunctions still produces a single decoder-layer ONNX function. Verified: `python -m pytest -q tests/unit_test/models/test_model_quickcheck.py -n auto` --------- Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com> Signed-off-by: Rishin Raj <rishinr@qti.qualcomm.com> Co-authored-by: Rishin Raj <rishinr@qti.qualcomm.com>
1 parent 5231e91 commit 15f8231

20 files changed

Lines changed: 183 additions & 399 deletions

File tree

‎QEfficient/transformers/models/falcon/modeling_falcon.py‎

Lines changed: 6 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -60,30 +60,17 @@ def _set_cos_sin_cache(self, seq_len, device, dtype):
6060
self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
6161

6262

63-
def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
63+
def qeff_apply_rotary_pos_emb(q, k, cos, sin):
6464
"""Applies Rotary Position Embedding to the query and key tensors.
6565
6666
Args:
6767
q (`torch.Tensor`): The query tensor.
6868
k (`torch.Tensor`): The key tensor.
6969
cos (`torch.Tensor`): The cosine part of the rotary embedding.
7070
sin (`torch.Tensor`): The sine part of the rotary embedding.
71-
position_ids (`torch.Tensor`):
72-
The position indices of the tokens corresponding to the query and key tensors. For example, this can be
73-
used to pass offsetted position ids when working with a KV-cache.
74-
unsqueeze_dim (`int`, *optional*, defaults to 1):
75-
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
76-
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
77-
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
78-
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
79-
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
80-
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
8171
Returns:
8272
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
8373
"""
84-
cos = cos[position_ids].unsqueeze(unsqueeze_dim)
85-
sin = sin[position_ids].unsqueeze(unsqueeze_dim)
86-
8774
# Apply rotation
8875
q_embed = (q * cos) + (rotate_half(q) * sin)
8976
k_embed = (k * cos) + (rotate_half(k) * sin)
@@ -127,7 +114,7 @@ def forward(
127114
value_layer = value_layer.transpose(1, 2).reshape(batch_size, num_kv_heads, query_length, self.head_dim)
128115

129116
# kv_seq_len = past_key_value.get_seq_length(self.layer_idx, cache_position)
130-
query_layer, key_layer = qeff_apply_rotary_pos_emb(query_layer, key_layer, cos_cached, sin_cached, position_ids)
117+
query_layer, key_layer = qeff_apply_rotary_pos_emb(query_layer, key_layer, cos_cached, sin_cached)
131118

132119
if layer_past is not None:
133120
past_seen_tokens = layer_past.get_seq_length()
@@ -308,6 +295,8 @@ def forward(
308295

309296
all_self_attentions = () if output_attentions else None
310297
all_hidden_states = () if output_hidden_states else None
298+
sin = self.sin_cached[position_ids].unsqueeze(1)
299+
cos = self.cos_cached[position_ids].unsqueeze(1)
311300

312301
for i, block in enumerate(self.h):
313302
if output_hidden_states:
@@ -326,8 +315,8 @@ def forward(
326315
output_attentions=output_attentions,
327316
alibi=alibi,
328317
cache_position=cache_position,
329-
sin_cached=self.sin_cached,
330-
cos_cached=self.cos_cached,
318+
sin_cached=sin,
319+
cos_cached=cos,
331320
)
332321

333322
hidden_states = outputs[0]

‎QEfficient/transformers/models/gemma/modeling_gemma.py‎

Lines changed: 6 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -62,30 +62,17 @@ def _set_cos_sin_cache(self, seq_len, device, dtype):
6262
self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
6363

6464

65-
def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
65+
def qeff_apply_rotary_pos_emb(q, k, cos, sin):
6666
"""Applies Rotary Position Embedding to the query and key tensors.
6767
6868
Args:
6969
q (`torch.Tensor`): The query tensor.
7070
k (`torch.Tensor`): The key tensor.
7171
cos (`torch.Tensor`): The cosine part of the rotary embedding.
7272
sin (`torch.Tensor`): The sine part of the rotary embedding.
73-
position_ids (`torch.Tensor`):
74-
The position indices of the tokens corresponding to the query and key tensors. For example, this can be
75-
used to pass offsetted position ids when working with a KV-cache.
76-
unsqueeze_dim (`int`, *optional*, defaults to 1):
77-
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
78-
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
79-
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
80-
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
81-
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
82-
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
8373
Returns:
8474
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
8575
"""
86-
cos = cos[position_ids].unsqueeze(unsqueeze_dim)
87-
sin = sin[position_ids].unsqueeze(unsqueeze_dim)
88-
8976
# Apply rotation
9077
q_embed = (q * cos) + (rotate_half(q) * sin)
9178
k_embed = (k * cos) + (rotate_half(k) * sin)
@@ -144,9 +131,7 @@ def forward(
144131
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
145132
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
146133

147-
query_states, key_states = qeff_apply_rotary_pos_emb(
148-
query_states, key_states, cos_cached, sin_cached, position_ids
149-
)
134+
query_states, key_states = qeff_apply_rotary_pos_emb(query_states, key_states, cos_cached, sin_cached)
150135

151136
past_seen_tokens = past_key_values.get_seq_length(self.layer_idx) if past_key_values is not None else 0
152137
blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig())
@@ -321,6 +306,8 @@ def forward(
321306

322307
# decoder layers
323308
all_hidden_states = () if output_hidden_states else None
309+
sin = self.sin_cached[position_ids].unsqueeze(1)
310+
cos = self.cos_cached[position_ids].unsqueeze(1)
324311

325312
for decoder_layer in self.layers[: self.config.num_hidden_layers]:
326313
if output_hidden_states:
@@ -335,8 +322,8 @@ def forward(
335322
batch_index=batch_index,
336323
use_cache=use_cache,
337324
cache_position=cache_position,
338-
sin_cached=self.sin_cached,
339-
cos_cached=self.cos_cached,
325+
sin_cached=sin,
326+
cos_cached=cos,
340327
**kwargs,
341328
)
342329

‎QEfficient/transformers/models/gemma2/modeling_gemma2.py‎

Lines changed: 6 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -65,30 +65,17 @@ def _set_cos_sin_cache(self, seq_len, device, dtype):
6565
self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
6666

6767

68-
def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
68+
def qeff_apply_rotary_pos_emb(q, k, cos, sin):
6969
"""Applies Rotary Position Embedding to the query and key tensors.
7070
7171
Args:
7272
q (`torch.Tensor`): The query tensor.
7373
k (`torch.Tensor`): The key tensor.
7474
cos (`torch.Tensor`): The cosine part of the rotary embedding.
7575
sin (`torch.Tensor`): The sine part of the rotary embedding.
76-
position_ids (`torch.Tensor`):
77-
The position indices of the tokens corresponding to the query and key tensors. For example, this can be
78-
used to pass offsetted position ids when working with a KV-cache.
79-
unsqueeze_dim (`int`, *optional*, defaults to 1):
80-
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
81-
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
82-
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
83-
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
84-
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
85-
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
8676
Returns:
8777
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
8878
"""
89-
cos = cos[position_ids].unsqueeze(unsqueeze_dim)
90-
sin = sin[position_ids].unsqueeze(unsqueeze_dim)
91-
9279
# Apply rotation
9380
q_embed = (q * cos) + (rotate_half(q) * sin)
9481
k_embed = (k * cos) + (rotate_half(k) * sin)
@@ -151,9 +138,7 @@ def forward(
151138
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
152139
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
153140

154-
query_states, key_states = qeff_apply_rotary_pos_emb(
155-
query_states, key_states, cos_cached, sin_cached, position_ids
156-
)
141+
query_states, key_states = qeff_apply_rotary_pos_emb(query_states, key_states, cos_cached, sin_cached)
157142

158143
past_seen_tokens = past_key_values.get_seq_length(self.layer_idx) if past_key_values is not None else 0
159144
blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig())
@@ -358,6 +343,8 @@ def forward(
358343
# decoder layers
359344
all_hidden_states = () if output_hidden_states else None
360345
all_self_attns = () if output_attentions else None
346+
sin = self.sin_cached[position_ids].unsqueeze(1)
347+
cos = self.cos_cached[position_ids].unsqueeze(1)
361348

362349
for decoder_layer in self.layers:
363350
if output_hidden_states:
@@ -373,8 +360,8 @@ def forward(
373360
output_attentions=output_attentions,
374361
use_cache=use_cache,
375362
cache_position=cache_position,
376-
sin_cached=self.sin_cached,
377-
cos_cached=self.cos_cached,
363+
sin_cached=sin,
364+
cos_cached=cos,
378365
**kwargs,
379366
)
380367

‎QEfficient/transformers/models/gpt_oss/modeling_gpt_oss.py‎

Lines changed: 14 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -541,7 +541,7 @@ def rotate_half(x):
541541
return torch.cat((-x2, x1), dim=-1)
542542

543543

544-
def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
544+
def qeff_apply_rotary_pos_emb(q, k, cos, sin):
545545
"""Applies Rotary Position Embedding with Multimodal Sections to the query and key tensors (https://qwenlm.github.io/blog/qwen2-vl/).
546546
547547
Explanation:
@@ -558,25 +558,10 @@ def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
558558
k (`torch.Tensor`): The key tensor.
559559
cos (`torch.Tensor`): The cosine part of the rotary embedding.
560560
sin (`torch.Tensor`): The sine part of the rotary embedding.
561-
position_ids (`torch.Tensor`):
562-
The position indices of the tokens corresponding to the query and key tensors. For example, this can be
563-
used to pass offsetted position ids when working with a KV-cache.
564-
mrope_section(`List(int)`):
565-
Multimodal rope section is for channel dimension of temporal, height and width in rope calculation.
566-
unsqueeze_dim (`int`, *optional*, defaults to 1):
567-
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
568-
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
569-
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
570-
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
571-
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
572-
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
573561
Returns:
574562
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
575563
"""
576564

577-
cos = cos[position_ids].unsqueeze(unsqueeze_dim)
578-
sin = sin[position_ids].unsqueeze(unsqueeze_dim)
579-
580565
q_embed = (q * cos) + (rotate_half(q) * sin)
581566
k_embed = (k * cos) + (rotate_half(k) * sin)
582567

@@ -703,9 +688,7 @@ def forward(
703688
hidden_shape = (*input_shape, -1, self.head_dim)
704689
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
705690
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
706-
query_states, key_states = qeff_apply_rotary_pos_emb(
707-
query_states, key_states, cos_cached, sin_cached, position_ids
708-
)
691+
query_states, key_states = qeff_apply_rotary_pos_emb(query_states, key_states, cos_cached, sin_cached)
709692

710693
if past_key_values is not None:
711694
# sin and cos are specific to RoPE models; cache_position needed for the static cache
@@ -787,9 +770,7 @@ def forward(
787770
hidden_shape = (*input_shape, -1, self.head_dim)
788771
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
789772
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
790-
query_states, key_states = qeff_apply_rotary_pos_emb(
791-
query_states, key_states, cos_cached, sin_cached, position_ids
792-
)
773+
query_states, key_states = qeff_apply_rotary_pos_emb(query_states, key_states, cos_cached, sin_cached)
793774

794775
if past_key_values is not None:
795776
# sin and cos are specific to RoPE models; cache_position needed for the static cache
@@ -868,9 +849,7 @@ def forward(
868849
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
869850
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
870851
past_seen_tokens = past_key_values.get_seq_length(self.layer_idx) if past_key_values is not None else 0
871-
query_states, key_states = qeff_apply_rotary_pos_emb(
872-
query_states, key_states, cos_cached, sin_cached, position_ids
873-
)
852+
query_states, key_states = qeff_apply_rotary_pos_emb(query_states, key_states, cos_cached, sin_cached)
874853

875854
if self.sliding_window is not None:
876855
attention_mask = sliding_mask
@@ -1044,6 +1023,8 @@ def forward(
10441023
# decoder layers
10451024
all_hidden_states = () if output_hidden_states else None
10461025
all_self_attns = () if output_attentions else None
1026+
sin = self.sin_cached[position_ids].unsqueeze(1)
1027+
cos = self.cos_cached[position_ids].unsqueeze(1)
10471028

10481029
for decoder_layer in self.layers:
10491030
if output_hidden_states:
@@ -1059,8 +1040,8 @@ def forward(
10591040
output_attentions=output_attentions,
10601041
cache_position=cache_position,
10611042
sliding_mask=sliding_mask,
1062-
sin_cached=self.sin_cached,
1063-
cos_cached=self.cos_cached,
1043+
sin_cached=sin,
1044+
cos_cached=cos,
10641045
**kwargs,
10651046
)
10661047
hidden_states = layer_outputs[0]
@@ -1085,8 +1066,8 @@ def forward(
10851066
class QEffGptOssModel(GptOssModel):
10861067
def __qeff_init__(self):
10871068
self.rotary_emb = QEffGptOssRotaryEmbedding(config=self.config)
1088-
self.sin_cached = torch.nn.Parameter(self.rotary_emb.sin_cached)
1089-
self.cos_cached = torch.nn.Parameter(self.rotary_emb.cos_cached)
1069+
self.sin_cached = torch.nn.Parameter(self.rotary_emb.sin_cached * self.rotary_emb.attention_scaling)
1070+
self.cos_cached = torch.nn.Parameter(self.rotary_emb.cos_cached * self.rotary_emb.attention_scaling)
10901071

10911072
def forward(
10921073
self,
@@ -1144,6 +1125,8 @@ def forward(
11441125
# decoder layers
11451126
all_hidden_states = () if output_hidden_states else None
11461127
all_self_attns = () if output_attentions else None
1128+
sin = self.sin_cached[position_ids].unsqueeze(1)
1129+
cos = self.cos_cached[position_ids].unsqueeze(1)
11471130

11481131
for decoder_layer in self.layers:
11491132
if output_hidden_states:
@@ -1160,8 +1143,8 @@ def forward(
11601143
output_attentions=output_attentions,
11611144
cache_position=cache_position,
11621145
sliding_mask=sliding_mask,
1163-
sin_cached=self.sin_cached,
1164-
cos_cached=self.cos_cached,
1146+
sin_cached=sin,
1147+
cos_cached=cos,
11651148
**kwargs,
11661149
)
11671150
hidden_states = layer_outputs[0]

‎QEfficient/transformers/models/granite/modeling_granite.py‎

Lines changed: 6 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -60,30 +60,17 @@ def _set_cos_sin_cache(self, seq_len, device, dtype):
6060
self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
6161

6262

63-
def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
63+
def qeff_apply_rotary_pos_emb(q, k, cos, sin):
6464
"""Applies Rotary Position Embedding to the query and key tensors.
6565
6666
Args:
6767
q (`torch.Tensor`): The query tensor.
6868
k (`torch.Tensor`): The key tensor.
6969
cos (`torch.Tensor`): The cosine part of the rotary embedding.
7070
sin (`torch.Tensor`): The sine part of the rotary embedding.
71-
position_ids (`torch.Tensor`):
72-
The position indices of the tokens corresponding to the query and key tensors. For example, this can be
73-
used to pass offsetted position ids when working with a KV-cache.
74-
unsqueeze_dim (`int`, *optional*, defaults to 1):
75-
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
76-
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
77-
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
78-
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
79-
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
80-
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
8171
Returns:
8272
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
8373
"""
84-
cos = cos[position_ids].unsqueeze(unsqueeze_dim)
85-
sin = sin[position_ids].unsqueeze(unsqueeze_dim)
86-
8774
# Apply rotation
8875
q_embed = (q * cos) + (rotate_half(q) * sin)
8976
k_embed = (k * cos) + (rotate_half(k) * sin)
@@ -137,9 +124,7 @@ def forward(
137124
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
138125
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
139126

140-
query_states, key_states = qeff_apply_rotary_pos_emb(
141-
query_states, key_states, cos_cached, sin_cached, position_ids
142-
)
127+
query_states, key_states = qeff_apply_rotary_pos_emb(query_states, key_states, cos_cached, sin_cached)
143128

144129
past_seen_tokens = past_key_values.get_seq_length(self.layer_idx) if past_key_values is not None else 0
145130
blocking_config = getattr(self, "attn_blocking_config", AttentionBlockingConfig())
@@ -320,6 +305,8 @@ def forward(
320305
# decoder layers
321306
all_hidden_states = () if output_hidden_states else None
322307
all_self_attns = () if output_attentions else None
308+
sin = self.sin_cached[position_ids].unsqueeze(1)
309+
cos = self.cos_cached[position_ids].unsqueeze(1)
323310

324311
for decoder_layer in self.layers[: self.config.num_hidden_layers]:
325312
if output_hidden_states:
@@ -335,8 +322,8 @@ def forward(
335322
output_attentions=output_attentions,
336323
use_cache=use_cache,
337324
cache_position=cache_position,
338-
sin_cached=self.sin_cached,
339-
cos_cached=self.cos_cached,
325+
sin_cached=sin,
326+
cos_cached=cos,
340327
**kwargs,
341328
)
342329

0 commit comments

Comments
 (0)