Fix for subfn, rope issue in Grok1

[tv-karthikeya]

This PR fixes Grok1 export/runtime support by addressing two issues:

• Adds Grok1 decoder-layer subfunction export support.
• Fixes Grok1 RoPE handling to match the current shared helper API.

RoPE Compatibility Fix

Grok1 was still calling the shared RoPE helper with position_ids:

qeff_apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
The shared helper imported from the Llama wrapper now expects already-indexed RoPE tensors:

qeff_apply_rotary_pos_emb(q, k, cos, sin)
This caused Grok1 QEff execution/export to fail with:

TypeError: qeff_apply_rotary_pos_emb() takes 4 positional arguments but 5 were given

The fix applies Grok1’s position_ids indexing locally before calling the shared helper:

  cos = cos[position_ids].unsqueeze(1)
  sin = sin[position_ids].unsqueeze(1)
  query_states, key_states = qeff_apply_rotary_pos_emb(query_states, key_states, cos, sin)

This preserves Grok1’s original RoPE semantics and produces the broadcast shape expected by attention: [batch, 1, seq_len, head_dim].

Subfunction Export Support

Grok1 now declares its repeated decoder block as the ONNX subfunction boundary:

  def get_submodules_for_export(self) -> Type[nn.Module]:
      return {QEffGrok1DecoderLayer}

Grok1 support is implemented through external module mapping, so the original HF decoder layer is transformed by replacing its forward/init behavior rather than by instantiating a standalone QEff...DecoderLayer class everywhere in the model tree. Because of that, the subfunction exporter cannot rely on discovering a generic QEffDecoderLayer instance name from the original model structure.

updated :

  def get_submodules_for_export(self) -> Type[nn.Module]:
        return {self.model.layers[0].__class__}

[tv-karthikeya]

CI-Ready

[ochougul]

this logic should go insdie the qeff_apply_rotary_pos_emb method I think. Why are we needing explicit unsqueeze here only?