sgl-project/sglang · error · ValueError

Cosmos3CrossAttention requires num_key_value_heads divisible

Error message

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

What it means

Cosmos3CrossAttention shards key/value heads for the text-encoding side across TP ranks. num_key_value_heads of the cross-attention block must be divisible by tp_size; otherwise per-rank KV heads cannot be computed and __init__ fails.

Source

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

        num_key_value_heads: int,
        head_dim: int,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
        supported_attention_backends: set | 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.tp_size = get_tp_world_size()
        if num_attention_heads % self.tp_size != 0:
            raise ValueError(
                "Cosmos3CrossAttention 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(
                "Cosmos3CrossAttention 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

Lower --tp to a divisor of the cross-attention num_key_value_heads
Inspect the config for all of num_attention_heads / num_key_value_heads in both causal and cross attention; choose TP dividing every one
When max TP is too small, shard other dims or use more replicas instead

Example fix

# before
--tp 8   # cross-attn kv heads = 4
# after
--tp 4

Defensive patterns

Strategy: validation

Validate before calling

assert cross_attn_num_kv_heads % tp == 0

Type guard

def tp_divides(n: int, tp: int) -> bool:
    return n % tp == 0

Prevention

Treat cross-attention KV heads as the TP bound in DiT models
Pre-flight check configs before multi-GPU launches

When it happens

Trigger: TP degree larger than, or not a divisor of, the cross-attention layer's num_key_value_heads — e.g. tp=8 with 4 cross-attn KV heads.

Common situations: Same class of mistake as the other head-divisibility errors but on cross-attention KV heads, which are often the smallest head count in the model and therefore the binding TP constraint.

Related errors

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