Hotfix axk2 indexer norm

[kmswin1]

There is currently a mismatch between the HF and vLLM implementations of the AXK2 indexer input.
HF passes the pre-normalization q_compressed, directly after q_a_proj, to the indexer.
vLLM passes the post-normalization q_c, after q_a_layernorm, to the indexer.
This difference may change the indexer scores and top-k selections. For consistency, we are considering updating the HF implementation to match the current vLLM serving behavior by passing the post-normalization query to the indexer.
We are planning to submit the following change as a PR. Would it be possible to handle this as a hotfix and merge it quickly?
topk_indices = self.indexer(
hidden_states,

• q_compressed,

• q_resid, # post-normalization query
  position_embeddings,
  indexer_mask,
  position_ids,
  past_key_values=past_key_values,
  )
  This change only affects the indexer input. The main query and attention output-gate paths remain unchanged.!

[Rocketknight1]

cc @vasqu since you added the original model!

[vasqu]

Just did a small change to inherit completely from the indexer, merging now

[vasqu]

Ah no let me update the expectations of the integration tests

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: axk2

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 31081375519:2
Result: success | Jobs: 8 | Tests: 377 | Failures: 0 | Duration: 2m 20s