sgl-project/sglang · error · ValueError

Cosmos3CausalAttention requires num_attention_heads divisibl

Error message

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

What it means

Cosmos3CausalAttention shards query heads across tensor-parallel ranks. num_attention_heads must be divisible by the tensor-parallel world size (tp_size from get_tp_world_size()) so each rank gets an equal head count; otherwise it raises ValueError at construction.

Source

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

        hidden_size: int,
        num_attention_heads: int,
        num_key_value_heads: int,
        head_dim: int,
        qk_norm: bool = True,
        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,

View on GitHub (pinned to 0132848349)

Solutions

1. Pick a tp_size that divides num_attention_heads (any power of 2 up to the largest power of 2 dividing it, e.g. 1/2/4/8 for 64 heads)
2. Check the model's config for num_attention_heads before choosing --tp
3. Fall back to tp=1 if unsure or for small models

Example fix

# before
python -m sglang.launch_server --model cosmos3 ... --tp 6   # 64 heads
# after
python -m sglang.launch_server --model cosmos3 ... --tp 4   # 64 % 4 == 0

tp = get_tp_world_size()
assert num_attention_heads % tp == 0, f"{num_attention_heads=} not divisible by tp={tp}"

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

When it happens

Trigger: Launching Cosmos3 video generation with --tp N where N does not divide the model's num_attention_heads (e.g. tp=4 with 30 heads). The attention module __init__ raises immediately during model build.

Common situations: Choosing a TP degree tuned for a text LLM (e.g. tp=8) and reusing it for a video DiT with a different head count; using head counts that are products of small primes incompatible with the chosen TP.

Related errors

• Cosmos3CausalAttention requires num_key_value_heads divisibl
• Cosmos3CrossAttention requires num_attention_heads divisible
• Cosmos3CrossAttention requires num_key_value_heads divisible
• {selection_error}{component_suffix}
• Cosmos3 requires text_ids and text_mask to be passed

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