sgl-project/sglang · error · ValueError

AITer backend requires num_heads ({num_heads}) to be a multi

Error message

AITer backend requires num_heads ({num_heads}) to be a multiple of num_kv_heads ({num_kv_heads}).

What it means

AITerImpl.__init__ validates that num_heads divides evenly by num_kv_heads (GQA constraint). aiter's MHA kernels broadcast each KV head across its query-head group, so an uneven split cannot be mapped. num_kv_heads=None skips the check.

Source

Thrown at python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py:136

    """

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        softmax_scale: float,
        causal: bool = False,
        num_kv_heads: int | None = None,
        prefix: str = "",
        dropout_p: float = 0.0,
        **extra_impl_args,
    ) -> None:
        # aiter's mha entry points take GQA/MQA K/V directly (they broadcast
        # each KV head across its group of query heads), so the only
        # requirement is an even split. The FP8 ASM path is MHA-only and
        # already routes grouped shapes back to BF16 below.
        if num_kv_heads is not None and num_heads % num_kv_heads != 0:
            raise ValueError(
                f"AITer backend requires num_heads ({num_heads}) to be a "
                f"multiple of num_kv_heads ({num_kv_heads})."
            )
        self.causal = causal
        self.dropout_p = dropout_p
        self.softmax_scale = softmax_scale

    @torch.compiler.disable
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        """
        Performs attention using one of:
          - _fmha_fp8_prefill_attention (FP8, SGLANG_DIFFUSION_AITER_FP8_ATTN=1 when eligible)

View on GitHub (pinned to 0132848349)

Solutions

  1. Use a model/config where num_attention_heads is a multiple of num_key_value_heads (standard GQA shapes)
  2. Fix the num_heads/num_kv_heads values being passed if they were misread from the model config
  3. Switch to a backend that supports uneven GQA grouping (e.g. flash attention) for this model

Example fix

// before: model config has 28 query heads, 6 kv heads -> 28 % 6 != 0
AITerImpl(num_heads=28, head_size=128, softmax_scale=s, num_kv_heads=6)
// after: use a valid GQA shape or another backend
AITerImpl(num_heads=32, head_size=128, softmax_scale=s, num_kv_heads=8)  # 32 % 8 == 0
Defensive patterns

Strategy: validation

Validate before calling

if num_kv_heads is not None:
    assert num_heads % num_kv_heads == 0, (
        f"invalid GQA shape: {num_heads} heads / {num_kv_heads} kv heads"
    )

Type guard

def is_valid_gqa(num_heads: int, num_kv_heads: int | None) -> bool:
    return num_kv_heads is None or num_heads % num_kv_heads == 0

Prevention

When it happens

Trigger: Constructing AITerImpl with num_kv_heads not None and num_heads % num_kv_heads != 0, e.g. num_heads=24, num_kv_heads=8 (24%8==0 passes) vs num_heads=28, num_kv_heads=6 (fails). Typically the values come from the model's attention config (num_attention_heads / num_key_value_heads).

Common situations: Selecting the aiter attention backend on AMD GPUs for a model whose GQA head grouping is fractional (query heads not a multiple of KV heads), or a config typo in num_key_value_heads.

Related errors


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