sgl-project/sglang · error · ValueError

GQA/MQA requires query heads to be a multiple of KV heads, g

Error message

GQA/MQA requires query heads to be a multiple of KV heads, got q_heads={query.shape[1]} and kv_heads={key.shape[1]}

What it means

SageAttention3's Blackwell kernel assumes MHA (equal query and KV head counts); for GQA/MQA it falls back to torch SDPA, but that fallback requires query heads to be an integer multiple of KV heads. Otherwise attention grouping is undefined and forward raises ValueError.

Source

Thrown at python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py:71

        self.causal = causal
        self.softmax_scale = softmax_scale
        self.dropout = extra_impl_args.get("dropout_p", 0.0)

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata,
    ) -> torch.Tensor:
        query = query.transpose(1, 2)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)
        # SageAttention3's Blackwell kernel assumes MHA (Hq == Hkv). For GQA/MQA
        # (Hq != Hkv), fall back to torch SDPA which supports GQA.
        if key.shape[1] != query.shape[1]:
            if query.shape[1] % key.shape[1] != 0:
                raise ValueError(
                    "GQA/MQA requires query heads to be a multiple of KV heads, "
                    f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}"
                )
            if not type(self)._warned_gqa_fallback_global:
                logger.warning(
                    "SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA."
                )
                type(self)._warned_gqa_fallback_global = True
            output = F.scaled_dot_product_attention(
                query,
                key,
                value,
                is_causal=self.causal,
                dropout_p=self.dropout,
                scale=self.softmax_scale,
                enable_gqa=True,
            )
        else:

View on GitHub (pinned to 0132848349)

Solutions

  1. Fix the model/config head counts so q_heads is a multiple of kv_heads
  2. Use a backend with native GQA support (flash_attn) for this model
  3. If heads are genuinely non-divisible, select a backend supporting arbitrary grouped attention

Example fix

# before
impl = SageAttention3Impl(head_size=128, num_heads=12, num_kv_heads=8, ...)  # 12 % 8 != 0
# after
impl = FlashAttentionImpl(head_size=128, num_heads=12, num_kv_heads=8, ...)
Defensive patterns

Strategy: validation

Validate before calling

assert query.shape[1] % key.shape[1] == 0, "q_heads must be a multiple of kv_heads for sage_attn3 fallback"

Type guard

def sage3_compatible(q: torch.Tensor, k: torch.Tensor) -> bool:
    return q.shape[1] == k.shape[1] or q.shape[1] % k.shape[1] == 0

Try / catch

try:
    out = impl.forward(q, k, v, meta)
except ValueError as e:
    if "GQA/MQA" in str(e):
        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)

Prevention

When it happens

Trigger: Calling forward on the sage_attn3 backend with q.shape[1] % k.shape[1] != 0, e.g. 12 query heads and 8 KV heads.

Common situations: Loading a model with unusual head-group sizes not supported by SageAttention3; quantized/remapped attention with mismatched heads; new model configs hitting the GQA fallback path.

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/73751e210631c0b2. Report an issue: GitHub.