THE GQA MEMORY ADVANTAGE

Difficulty: Hard Tags: Transformers, Grouped-Query Attention, KV Cache

Story

A serving team switches from multi-head attention to grouped-query attention (GQA) to reduce KV-cache memory. A capacity calculator incorrectly uses the number of query heads for key/value storage. Compute the KV-cache memory under the stated configuration.

The Math

For batch B, layers L, cached tokens T, query heads Hq, key/value heads Hkv, head dimension D, and bytes per scalar b:

M = 2 * B * L * T * Hkv * D * b

The configuration is valid only if Hq % Hkv == 0.

Input Format

b, l, t, hq, hkv, d, bytes_per_scalar (all positive integers).

Output Format

None if Hq is not divisible by Hkv; otherwise the cache size in bytes (int).

Constraints

• All dimensions positive integers at most 10^5, b in {1,2,4,8}.
• Result fits in a signed 64-bit integer.
• Time limit: 1.0 second.

Example

Input

B=1, L=32, T=4096, Hq=32, Hkv=8, D=128, bytes_per_scalar=2

Output

536870912

Explanation: The cache stores keys and values for 8 KV heads, not 32 query heads: 2*1*32*4096*8*128*2 = 536870912 bytes.