sgl-project/sglang · error · ValueError

Cosmos3CausalAttention requires num_key_value_heads divisibl

Error message

Cosmos3CausalAttention requires num_key_value_heads divisible by tp_size, got {num_key_value_heads=} {self.tp_size=}.

What it means

Cosmos3CausalAttention also shards key/value heads for GQA. num_key_value_heads must be divisible by tp_size or the per-rank KV head count would be unequal/zero, so construction fails with this ValueError.

Source

Thrown at python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py:590

        use_k_norm_und_for_gen: bool = False,
        rms_norm_eps: float = 1e-6,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.qk_norm = qk_norm
        self.tp_size = get_tp_world_size()
        if num_attention_heads % self.tp_size != 0:
            raise ValueError(
                "Cosmos3CausalAttention requires num_attention_heads divisible "
                f"by tp_size, got {num_attention_heads=} {self.tp_size=}."
            )
        if num_key_value_heads % self.tp_size != 0:
            raise ValueError(
                "Cosmos3CausalAttention requires num_key_value_heads divisible "
                f"by tp_size, got {num_key_value_heads=} {self.tp_size=}."
            )
        self.local_num_attention_heads = num_attention_heads // self.tp_size
        self.local_num_key_value_heads = num_key_value_heads // self.tp_size

        self.q_size = num_attention_heads * head_dim
        self.kv_size = num_key_value_heads * head_dim
        self.to_qkv = MergedColumnParallelLinear(
            hidden_size,
            [self.q_size, self.kv_size, self.kv_size],
            bias=False,
            gather_output=False,
            quant_config=quant_config,
            prefix=add_prefix("to_qkv", prefix),
        )
        self.to_out = RowParallelLinear(
            num_attention_heads * head_dim,

View on GitHub (pinned to 0132848349)

Solutions

1. Set tp_size to a divisor of num_key_value_heads (and of num_attention_heads); the effective max TP = num_key_value_heads
2. Read num_key_value_heads from the HF/config file and choose --tp accordingly
3. If more parallelism is needed, combine a smaller TP with pipeline/sequence parallelism instead

Example fix

# before
--tp 8   # model has 4 kv heads
# after
--tp 4   # 4 % 4 == 0

tp = get_tp_world_size()
assert num_key_value_heads % tp == 0, f"max tp is {num_key_value_heads}"

def valid_tp_for_gqa(num_kv_heads: int, tp: int) -> bool:
    return num_kv_heads % tp == 0

When it happens

Trigger: Running Cosmos3 with --tp N where N exceeds or does not divide num_key_value_heads — e.g. tp=8 with only 4 KV heads (GQA models often have few KV heads, capping maximum TP).

Common situations: GQA architectures with small KV-head counts (4 or 8) launched with aggressive TP; reusing TP settings from the query-head check without considering KV heads. Max valid tp_size is bounded by num_key_value_heads.

Related errors

• Cosmos3CrossAttention requires num_key_value_heads divisible
• Cosmos3CausalAttention requires num_attention_heads divisibl
• Cosmos3CrossAttention requires num_attention_heads divisible
• num_key_value_heads must be divisible by attention TP
• attention TP must be divisible by num_key_value_heads

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