support gqa

[ajassani]

This PR adds support for GQA by updating parameter naming and computation logic in performance-related functions.

Renames parameters (e.g., H → H_Q, d_k → d_h, N_K → N_KV) to align with GQA semantics.
Updates flops, bytes, and backward flops calculations and adjusts the extraction of tensor shapes in both forward and backward methods.

[Copilot]

Pull Request Overview

This PR adds support for GQA by updating parameter naming and computation logic in performance-related functions.

• Renames parameters (e.g., H → H_Q, d_k → d_h, N_K → N_KV) to align with GQA semantics.
• Updates flops, bytes, and backward flops calculations and adjusts the extraction of tensor shapes in both forward and backward methods.

Comments suppressed due to low confidence (3)

TraceLens/PerfModel/perf_model.py:873

• The dimension ordering in backward get_param_details appears inconsistent with the forward function, where Q shape is expected as (B, N_Q, H_Q, d_h). Verify that the reordering to (B, H_Q, N_Q, d_h) is intentional.

B, H_Q, N_Q, d_h = q_shape

TraceLens/PerfModel/perf_model.py:798

• [nitpick] Using integer floor division in the flops calculation may lead to unintended truncation of critical precision; please confirm that this behavior is as intended for scaling purposes in the GQA computations.

flops_vgrad +=  B * N_KV * d_h * (H_Q//H_KV -1 )

TraceLens/PerfModel/perf_model.py:809

• [nitpick] Double-check the intended effect of the integer division on H_Q relative to H_KV in the backward gradient computations, as a similar pattern appears with flops_vgrad.

flops_k_grad += B * N_KV * d_h * (H_Q//H_KV -1 )