fix GQA

[XZman]

No description provided.

[Copilot]

Pull request overview

This PR aims to fix/enable GQA (grouped-query attention) modeling by threading a num_kv_heads parameter through the LLM attention op generators, updating attention block tensor shapes to reflect Q-heads vs KV-heads, and relaxing FlashAttention shape assumptions to permit Q_heads != KV_heads.

Changes:

• Add num_kv_heads plumbing across attention entrypoints (create_multi_head_attention*) and propagate it from LLMOpsGenerator.
• Update normal/self-attention (fwd/bwd) and flash-attention blocks to use KV-head shapes where applicable.
• Relax FlashAttention shape assertions to only require K_heads == V_heads (allowing GQA/MQA).

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.

File	Description
neusim/npusim/frontend/llm_ops_lib.py	Adds num_kv_heads throughout attention construction and introduces GQA-shaped einsums/softmax for attention blocks.
neusim/npusim/frontend/llm_ops_generator.py	Passes self.num_kv_heads into attention creation for prefill/decode and fwd/bwd generation paths.
neusim/npusim/backend/npusim_lib.py	Relaxes FlashAttention head-count assertion to allow Q-heads to differ from KV-heads.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.