sgl-project/sglang · error · ValueError

Cosmos3CrossAttention requires num_attention_heads divisible

Error message

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

What it means

The cross-attention block in Cosmos3 shards text-conditioning query heads across tensor-parallel ranks. If num_attention_heads of the cross-attention layer is not divisible by tp_size, __init__ raises this ValueError during model build.

Source

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

    def __init__(
        self,
        hidden_size: int,
        num_attention_heads: int,
        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,

View on GitHub (pinned to 0132848349)

Solutions

1. Check the cross-attention num_attention_heads in the model config and pick a --tp that divides it (as well as the causal-layer head counts)
2. Use a smaller TP that satisfies all head-count divisibility constraints
3. Add a startup assertion script that validates all three head counts against your chosen tp before launching

Example fix

# before
--tp 6   # cross-attn heads = 8
# after
--tp 4   # 8 % 4 == 0

for name in ("num_attention_heads", "num_key_value_heads"):
    for sub in ("self_attn", "cross_attn"):
        assert cfg[sub][name] % tp == 0, f"{sub}.{name} not divisible by tp={tp}"

def tp_valid_for_config(cfg, tp: int) -> bool:
    return all(h % tp == 0 for h in all_head_counts(cfg))

When it happens

Trigger: Launching Cosmos3 with a TP degree that divides the self-attention head count but not the cross-attention head count (some configs use different head counts for cross-attention), e.g. tp=8 with 12 cross-attn heads.

Common situations: Configs where cross_attention layers have distinct num_attention_heads from causal layers; validating TP only against the main attention config and missing the cross-attention head count.

Related errors

• Cosmos3CrossAttention requires num_key_value_heads divisible
• Cosmos3CausalAttention requires num_attention_heads divisibl
• Cosmos3CausalAttention requires num_key_value_heads divisibl
• Cosmos3 requires text_ids and text_mask to be passed
• Cosmos3 action generation does not support sequence parallel

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