Summary

In the GQA (Grouped Query Attention) reshape operations across multiple modules, int(kg / self.num_kv_heads) uses float division (/) instead of integer division (//). When kg is not exactly divisible by num_kv_heads, Python's int() silently truncates the result, producing an incorrect reshape dimension with no error.

Affected files

Example (from gm/nn/_modules.py)

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

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

Impact

If kg % num_kv_heads != 0, the truncated dimension causes the reshape to produce a tensor whose total element count does not match the original. This leads to:

• A ValueError at runtime from JAX/NumPy (reshape size mismatch), or
• Silent incorrect computation if the truncation happens to produce a valid (but wrong) shape by coincidence

In practice, standard model configurations keep num_query_heads a multiple of num_kv_heads, so this bug is latent. However, custom or experimental configurations will hit it unexpectedly.

Fix

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