[Fix] Reject non-divisible GQA head counts

fla/ops/attn/decoding.py

@@ -10,6 +10,7 @@
 import triton.language as tl
 
 from fla.ops.utils.cumsum import chunk_global_cumsum
+from fla.ops.utils.head import get_gqa_group_size
 from fla.ops.utils.op import exp
 from fla.utils import autotune_cache_kwargs, check_shared_mem
 
@@ -160,7 +161,7 @@ def attn_decoding_one_step(
     B, T, H, K, V = *k.shape, v.shape[-1]
     N = len(cu_seqlens) - 1
     HQ = q.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     if scale is None:
         scale = K ** -0.5
     if sink_bias is not None:

fla/ops/attn/naive.py

@@ -8,6 +8,8 @@
 import torch
 import torch.nn.functional as F
 
+from fla.ops.utils.head import get_gqa_group_size
+
 
 def naive_parallel_attn(
     q: torch.Tensor,
@@ -45,7 +47,7 @@ def naive_parallel_attn(
     """
     B, T, HQ, D = q.shape
     H = k.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
 
     if scale is None:
         scale = D ** -0.5
@@ -126,7 +128,7 @@ def naive_attn_decoding(
     HQ, D = q.shape[-2], q.shape[-1]
     V = v.shape[-1]
     H = k.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     if scale is None:
         scale = D ** -0.5
     if sink_bias is not None:

fla/ops/attn/parallel.py

@@ -14,6 +14,7 @@
 from fla.ops.utils import prepare_chunk_indices
 from fla.ops.utils.constant import RCP_LN2
 from fla.ops.utils.cumsum import chunk_global_cumsum
+from fla.ops.utils.head import get_gqa_group_size
 from fla.ops.utils.op import exp2, log2
 from fla.utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, contiguous
 
@@ -519,7 +520,7 @@ def parallel_attn_fwd(
 ):
     B, T, H, K, V = *k.shape, v.shape[-1]
     HQ = q.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     BT = 128
     if check_shared_mem('hopper', q.device.index):
         BS = min(64, max(16, triton.next_power_of_2(T)))
@@ -609,7 +610,7 @@ def parallel_attn_bwd(
 ):
     B, T, H, K, V = *k.shape, v.shape[-1]
     HQ = q.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     # dq/dk are reduced over the full value dim in one program (no cross-program accumulation),
     # so BV must span all of V (NV == 1). Don't cap it here -- the forward can, the backward can't.
     if check_shared_mem('hopper'):
@@ -825,6 +826,7 @@ def parallel_attn(
         )
     if scale is None:
         scale = k.shape[-1] ** -0.5
+    get_gqa_group_size(q.shape[2], k.shape[2])
     if cu_seqlens is not None and q.shape[0] != 1:
         raise ValueError(
             f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. "

fla/ops/dsa/naive.py

@@ -11,6 +11,8 @@
 import torch.nn.functional as F
 from einops import repeat
 
+from fla.ops.utils.head import get_gqa_group_size
+
 
 def naive_dsa_indexer(
     q_idx: torch.Tensor,
@@ -155,7 +157,7 @@ def naive_dsa(
         assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided"
 
     dtype = q.dtype
-    G = q.shape[2] // k.shape[2]
+    G = get_gqa_group_size(q.shape[2], k.shape[2])
     q, k, v = (x.float() for x in (q, k, v))
     k, v = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k, v))
 

fla/ops/forgetting_attn/naive.py

@@ -9,6 +9,8 @@
 import torch.nn.functional as F
 from einops import rearrange, repeat
 
+from fla.ops.utils.head import get_gqa_group_size
+
 
 def naive_forgetting_attn(
     q: torch.Tensor,
@@ -35,7 +37,7 @@ def naive_forgetting_attn(
     """
     _, T, HQ, D = q.shape
     H = k.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
 
     if scale is None:
         scale = D ** -0.5

fla/ops/nsa/compression.py

@@ -11,6 +11,7 @@
 
 from fla.ops.attn.parallel import parallel_attn_bwd_preprocess
 from fla.ops.utils import prepare_chunk_indices, prepare_chunk_offsets, prepare_token_indices
+from fla.ops.utils.head import get_gqa_group_size
 from fla.ops.utils.op import exp, log
 from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, contiguous
 
@@ -336,7 +337,7 @@ def parallel_nsa_compression_fwd(
 ):
     B, TQ, HQ, K, V = *q.shape, v.shape[-1]
     H = k.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     BC = BS = block_size
     if check_shared_mem('hopper', q.device.index):
         BK = min(256, triton.next_power_of_2(K))
@@ -396,7 +397,7 @@ def parallel_nsa_compression_bwd(
     B, T, HQ, K, V = *q.shape, v.shape[-1]
     TC = k.shape[1]
     H = k.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     BC = BS = block_size
     BK = max(triton.next_power_of_2(K), 16)
     BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16))

fla/ops/nsa/naive.py

@@ -11,6 +11,7 @@
 from einops import repeat
 
 from fla.ops.utils import prepare_chunk_offsets
+from fla.ops.utils.head import get_gqa_group_size
 from fla.ops.utils.pooling import mean_pooling
 
 try:
@@ -68,7 +69,7 @@ def naive_nsa_selection(
         scale = k.shape[-1] ** -0.5
 
     dtype = q.dtype
-    G = q.shape[2] // k.shape[2]
+    G = get_gqa_group_size(q.shape[2], k.shape[2])
     BS = block_size
     k, v, block_indices = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k, v, block_indices))
     q, k, v = map(lambda x: x.float(), (q, k, v))
@@ -225,7 +226,7 @@ def naive_nsa_topk(
     """
     B, TQ, HQ, _ = q.shape
     H = k_cmp.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     k_cmp = repeat(k_cmp, 'b t h d -> b t (h g) d', g=G)
 
     device = q.device
@@ -364,8 +365,9 @@ def naive_nsa(
         scale = k.shape[-1] ** -0.5
     if cu_seqlens is not None:
         assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided"
-    G = q.shape[2] // k.shape[2]
-    assert G >= 16 and (G & (G - 1)) == 0, "Group size (HQ/H) must be a power of 2 and >= 16 in NSA"
+    G = get_gqa_group_size(q.shape[2], k.shape[2])
+    if G < 16 or (G & (G - 1)) != 0:
+        raise ValueError("Group size (HQ/H) must be a power of 2 and >= 16 in NSA")
 
     if cu_seqlens is not None:
         if isinstance(cu_seqlens, tuple):

fla/ops/nsa/parallel.py

@@ -15,6 +15,7 @@
 from fla.ops.nsa.compression import parallel_nsa_compression
 from fla.ops.nsa.utils import _bitonic_merge
 from fla.ops.utils import prepare_block_csr, prepare_chunk_indices, prepare_chunk_offsets, prepare_lens, prepare_token_indices
+from fla.ops.utils.head import get_gqa_group_size
 from fla.ops.utils.op import exp, log
 from fla.ops.utils.pooling import mean_pooling
 from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, contiguous
@@ -540,7 +541,7 @@ def parallel_nsa_topk(
     else:
         cu_seqlens_q = cu_seqlens_k = token_indices_q = None
 
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     # the number of selected blocks for each token
     S = block_counts if isinstance(block_counts, int) else block_counts.max().item()
     S = triton.next_power_of_2(S)
@@ -592,7 +593,7 @@ def parallel_nsa_fwd(
 ):
     B, TK, H, K, V, S = *k.shape, v.shape[-1], block_indices.shape[-1]
     _, TQ, HQ, _ = q.shape
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     BS = block_size
     if check_shared_mem('hopper', q.device.index):
         BK = min(256, triton.next_power_of_2(K))
@@ -652,7 +653,7 @@ def parallel_nsa_bwd(
 ):
     B, T, H, K, V, S = *k.shape, v.shape[-1], block_indices.shape[-1]
     HQ = q.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     BS = block_size
     BK = max(triton.next_power_of_2(K), 16)
     BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16))
@@ -875,8 +876,9 @@ def parallel_nsa(
             f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. "
             f"Please flatten variable-length inputs before processing.",
         )
-    G = q.shape[2] // k.shape[2]
-    assert G >= 16 and (G & (G - 1)) == 0, "Group size (HQ/H) must be a power of 2 and >= 16 in NSA"
+    G = get_gqa_group_size(q.shape[2], k.shape[2])
+    if G < 16 or (G & (G - 1)) != 0:
+        raise ValueError("Group size (HQ/H) must be a power of 2 and >= 16 in NSA")
 
     if cu_seqlens is not None:
         if isinstance(cu_seqlens, tuple):

fla/ops/parallax/decode.py

@@ -10,6 +10,7 @@
 import triton.language as tl
 
 from fla.ops.parallax.parallel import _block_size
+from fla.ops.utils.head import get_gqa_group_size
 from fla.ops.utils.op import exp2
 
 
@@ -162,7 +163,7 @@ def parallax_decode(
     """
     B, Sq, HQ, K = q.shape
     Skv, H = k.shape[1], k.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     if scale is None:
         scale = K ** -0.5
     window_size_left = -1 if window_size is None else window_size
@@ -315,7 +316,7 @@ def parallax_decode_one_step(
     if Sq != 1:
         raise ValueError(f"parallax_decode_one_step expects a single query (Sq=1), got Sq={Sq}")
     Skv, H = k.shape[1], k.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     if scale is None:
         scale = K ** -0.5
     window_size_left = -1 if window_size is None else window_size

fla/ops/parallax/naive.py

@@ -7,6 +7,8 @@
 
 import torch
 
+from fla.ops.utils.head import get_gqa_group_size
+
 
 def naive_parallax(
     q: torch.Tensor,
@@ -56,7 +58,7 @@ def naive_parallax(
     """
     B, T, HQ, D = q.shape
     H = k.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
 
     if scale is None:
         scale = D ** -0.5

fla/ops/parallax/parallel.py

@@ -11,6 +11,7 @@
 from einops import reduce
 
 from fla.ops.utils import prepare_chunk_indices
+from fla.ops.utils.head import get_gqa_group_size
 from fla.ops.utils.op import exp2
 from fla.utils import IS_NVIDIA_BLACKWELL, autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, contiguous
 
@@ -650,7 +651,7 @@ def parallel_parallax_fwd(q, r, k, v, scale, cu_seqlens=None, chunk_indices=None
     """
     B, T, HQ, K = q.shape
     H = k.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     BK = triton.next_power_of_2(K)
     BT = _block_size(K, q.device.index)
     o = torch.empty_like(q)
@@ -675,7 +676,7 @@ def parallel_parallax_bwd(q, r, k, v, o, barv, d1, bart, m, grad_o, scale, cu_se
     """Parallax backward (Triton). Returns grads matching `q, r, k, v`."""
     B, T, HQ, K = q.shape
     H = k.shape[2]
-    G = HQ // H
+    G = get_gqa_group_size(HQ, H)
     BK = triton.next_power_of_2(K)
     BT = _block_size(K, q.device.index)
 
@@ -796,6 +797,7 @@ def parallel_parallax(
         raise TypeError(f"parallel_parallax requires bf16 or fp16 inputs, got q.dtype={q.dtype}")
     if scale is None:
         scale = k.shape[-1] ** -0.5
+    get_gqa_group_size(q.shape[2], k.shape[2])
     if cu_seqlens is not None and q.shape[0] != 1:
         raise ValueError(
             f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. "

fla/ops/utils/head.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+
+def get_gqa_group_size(num_query_heads: int, num_kv_heads: int) -> int:
+    if num_kv_heads == 0 or num_query_heads % num_kv_heads != 0:
+        raise ValueError(
+            f"The number of query heads ({num_query_heads}) must be divisible by "
+            f"the number of key/value heads ({num_kv_heads})."
+        )
+    return num_query_heads // num_kv_heads

tests/ops/test_attn.py

@@ -15,6 +15,16 @@
 from fla.utils import assert_close, check_shared_mem, device
 
 
+@pytest.mark.parametrize("op", [naive_parallel_attn, parallel_attn], ids=["naive", "parallel"])
+def test_parallel_rejects_invalid_gqa_head_counts(op):
+    q = torch.empty(1, 1, 3, 16, dtype=torch.float16)
+    k = torch.empty(1, 1, 2, 16, dtype=torch.float16)
+    v = torch.empty(1, 1, 2, 16, dtype=torch.float16)
+
+    with pytest.raises(ValueError, match="must be divisible"):
+        op(q=q, k=k, v=v)
+
+
 @pytest.mark.parametrize(
     ('B', 'T', 'H', 'HQ', 'D', 'scale'),
     [

tests/ops/test_attn_sink.py

@@ -17,6 +17,17 @@
 from fla.utils import assert_close, device
 
 
+@pytest.mark.parametrize("op", [naive_attn_decoding, attn_decoding_one_step], ids=["naive", "triton"])
+def test_attn_decoding_rejects_invalid_gqa_head_counts(op):
+    q = torch.empty(1, 1, 3, 16, dtype=torch.float16)
+    k = torch.empty(1, 2, 2, 16, dtype=torch.float16)
+    v = torch.empty(1, 2, 2, 16, dtype=torch.float16)
+    cu_seqlens = torch.tensor([0, 2], dtype=torch.int32)
+
+    with pytest.raises(ValueError, match="must be divisible"):
+        op(q=q, k=k, v=v, cu_seqlens=cu_seqlens)
+
+
 def _repeat_kv_for_gpt_oss(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     if n_rep == 1:
         return hidden_states

tests/ops/test_forgetting_attn.py

@@ -183,3 +183,13 @@ def test_parallel_swa(
     assert_close("dk", ref_dk, tri_dk, 0.005)
     assert_close("dv", ref_dv, tri_dv, 0.005)
     assert_close("dg", ref_dg, tri_dg, 0.005)
+
+
+def test_naive_forgetting_attn_rejects_invalid_gqa_head_counts():
+    q = torch.empty(1, 1, 3, 16)
+    k = torch.empty(1, 1, 2, 16)
+    v = torch.empty_like(k)
+    g = torch.empty(1, 1, 3)
+
+    with pytest.raises(ValueError, match="must be divisible"):
+        naive_forgetting_attn(q=q, k=k, v=v, g=g)

tests/ops/test_nsa.py

@@ -23,6 +23,17 @@
 from fla.utils import assert_close, device  # noqa: E402
 
 
+@pytest.mark.parametrize("op", [naive_nsa, parallel_nsa], ids=["naive", "parallel"])
+def test_parallel_nsa_rejects_invalid_gqa_head_counts(op):
+    q = torch.empty(1, 1, 33, 16, dtype=torch.float16)
+    k = torch.empty(1, 1, 2, 16, dtype=torch.float16)
+    v = torch.empty(1, 1, 2, 16, dtype=torch.float16)
+    block_indices = torch.zeros(1, 1, 2, 1, dtype=torch.long)
+
+    with pytest.raises(ValueError, match="must be divisible"):
+        op(q=q, k=k, v=v, block_indices=block_indices, block_counts=1)
+
+
 def build_block_indices(B, T, H, S, block_size, seq_indices=None):
     block_indices = torch.full((B, T, H, S), -1, dtype=torch.long, device=device)
     for b in range(B):