Refactor qheads_per_kvhead calculations for clarity

[LoserCheems]

Summary

• Improve clarity and consistency in the calculation of qheads_per_kvhead in forward functions.

Root Cause

• The previous implementation had redundant calculations that could lead to confusion.

Changes

• Refactored the calculation of qheads_per_kvhead and introduced a separate variable for the packed case.

Reproduction

• No specific bug to reproduce; this is a code clarity improvement.

Tests

• Existing tests validate the functionality; no new tests required.

Compatibility

• No backward compatibility issues.

Checklist

• Linked issue provided [FEATURE REQUEST] Next-Generation Trainable Sparse Attention Mechanism #219 [FEATURE REQUEST] Triton-based efficient multi-platform, multi-variant attention #222
• Adds or updates tests
• Updates docs if needed
• No perf regressions

[Copilot]

Pull request overview

This PR refactors the qheads_per_kvhead variable calculations in the four forward dispatch functions within flash_fwd.py to improve clarity. Previously, a single variable qheads_per_kvhead was conditionally set based on pack_gqa and then also re-computed inline where the unconditional value was needed. The refactoring splits this into two clearly named variables.

Changes:

• Split the single qheads_per_kvhead variable into two: qheads_per_kvhead (always the true ratio num_heads_q // num_heads_kv) and qheads_per_kvhead_packgqa (conditionally 1 when pack_gqa is False).
• Replaced the inline num_heads_q // num_heads_kv expression at each kernel call site with the pre-computed qheads_per_kvhead variable.
• Applied the same refactoring consistently across all four forward functions: _flash_attn_base_forward, _flash_attn_varlen_base_forward, _flash_attn_sm90_forward, and _flash_attn_varlen_sm90_forward.

---

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