{"record":{"id":"8334a57f324beca4","repo":"sgl-project/sglang","slug":"cosmos3causalattention-requires-num-key-value-head","errorCode":null,"errorMessage":"Cosmos3CausalAttention requires num_key_value_heads divisible by tp_size, got {num_key_value_heads=} {self.tp_size=}.","messagePattern":"Cosmos3CausalAttention requires num_key_value_heads divisible by tp_size, got (.+?) (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py","lineNumber":590,"sourceCode":"        use_k_norm_und_for_gen: bool = False,\n        rms_norm_eps: float = 1e-6,\n        prefix: str = \"\",\n        quant_config: QuantizationConfig | None = None,\n    ):\n        super().__init__()\n        self.hidden_size = hidden_size\n        self.num_attention_heads = num_attention_heads\n        self.num_key_value_heads = num_key_value_heads\n        self.head_dim = head_dim\n        self.qk_norm = qk_norm\n        self.tp_size = get_tp_world_size()\n        if num_attention_heads % self.tp_size != 0:\n            raise ValueError(\n                \"Cosmos3CausalAttention requires num_attention_heads divisible \"\n                f\"by tp_size, got {num_attention_heads=} {self.tp_size=}.\"\n            )\n        if num_key_value_heads % self.tp_size != 0:\n            raise ValueError(\n                \"Cosmos3CausalAttention requires num_key_value_heads divisible \"\n                f\"by tp_size, got {num_key_value_heads=} {self.tp_size=}.\"\n            )\n        self.local_num_attention_heads = num_attention_heads // self.tp_size\n        self.local_num_key_value_heads = num_key_value_heads // self.tp_size\n\n        self.q_size = num_attention_heads * head_dim\n        self.kv_size = num_key_value_heads * head_dim\n        self.to_qkv = MergedColumnParallelLinear(\n            hidden_size,\n            [self.q_size, self.kv_size, self.kv_size],\n            bias=False,\n            gather_output=False,\n            quant_config=quant_config,\n            prefix=add_prefix(\"to_qkv\", prefix),\n        )\n        self.to_out = RowParallelLinear(\n            num_attention_heads * head_dim,","sourceCodeStart":572,"sourceCodeEnd":608,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py#L572-L608","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Set tp_size to a divisor of num_key_value_heads (and of num_attention_heads); the effective max TP = num_key_value_heads","Read num_key_value_heads from the HF/config file and choose --tp accordingly","If more parallelism is needed, combine a smaller TP with pipeline/sequence parallelism instead"],"exampleFix":"# before\n--tp 8   # model has 4 kv heads\n# after\n--tp 4   # 4 % 4 == 0","handlingStrategy":"validation","validationCode":"tp = get_tp_world_size()\nassert num_key_value_heads % tp == 0, f\"max tp is {num_key_value_heads}\"","typeGuard":"def valid_tp_for_gqa(num_kv_heads: int, tp: int) -> bool:\n    return num_kv_heads % tp == 0","tryCatchPattern":null,"preventionTips":["Remember max TP is capped by num_key_value_heads in GQA models","Document the model's TP-compatible values next to its launch command"],"tags":["sglang","cosmos3","tensor-parallel","gqa","kv-heads"],"backgroundTag":"attention-heads-not-divisible-by-tp","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}