fix(gqa): replace float division with integer division in GQA reshape

[Sumu004]

Problem

In the GQA (Grouped Query Attention) reshape operations, the number of query heads per KV head group is computed as:

int(kg / self.num_kv_heads)   # float division → int()

This is semantically wrong for two reasons:

1. Silent truncation: When kg % num_kv_heads != 0, float division produces a non-integer (e.g. 7/4 = 1.75), and int() silently truncates it to 1. The reshape proceeds with a wrong dimension — either silently corrupting the tensor shape or crashing with a confusing size-mismatch error.

2. Semantic mismatch: Tensor dimensions are integers by definition. Using float division here obscures the intent and masks invalid configurations.

Standard Gemma model configs (where num_query_heads is always a multiple of num_kv_heads) happen to produce exact divisions, so this bug is latent in production — but it surfaces immediately with custom or experimental head configurations.

Fix

Replace all 8 occurrences of int(kg / self.num_kv_heads) with kg // self.num_kv_heads:

# Before
query_scaled = query_scaled.reshape(
    (b, t, self.num_kv_heads, int(kg / self.num_kv_heads), h)
)

# After
query_scaled = query_scaled.reshape(
    (b, t, self.num_kv_heads, kg // self.num_kv_heads, h)
)

Affected files

File	Lines
gemma/gm/nn/_modules.py	244, 288
gemma/gm/nn/gemma4/_modules.py	326, 365
gemma/gm/nn/gemma3n/_modules.py	342, 387
gemma/research/t5gemma/modules.py	240, 272

Verification

Confirmed all 8 occurrences replaced, zero remaining instances of int(kg / self.num_kv_heads) in the codebase.

Fixes #641

[Sumu004]

Friendly bump on this one — it's a small, self-contained fix (8 occurrences across 4 files, same pattern each time) so should be quick to review. Happy to address any feedback or split it differently if that's easier to land. Thanks!

[Sumu004]

Following up again, it's been about two weeks since my last note. Still happy to address any feedback or split this differently if that helps it land. Thanks!