[Bug] mega_chunk_gdn unsupported head config for Qwen3.6-27B (28 heads GQA)

[JosephNew]

Description

mega_chunk_gdn operator crashes on Qwen3.6-27B because its head configuration is not in the allowed list.

Error

RuntimeError: unsupported mega_chunk_gdn (NumValueHeads, NumKeyHeads) pair

Root Cause

File: fla/mega_chunk_gdn.py lines 5-7, 59-66

GLOBAL_VALUE_HEADS = (16, 32, 48, 64)
GLOBAL_KEY_HEADS = 16
SUPPORTED_TP_DEGREES = (1, 2, 4, 8)

def _head_pair_supported(num_value_heads: int, num_key_heads: int) -> bool:
    if num_key_heads <= 0 or GLOBAL_KEY_HEADS % num_key_heads != 0:
        return False
    tp_degree = GLOBAL_KEY_HEADS // num_key_heads
    if tp_degree not in SUPPORTED_TP_DEGREES:
        return False
    return num_value_heads * tp_degree in GLOBAL_VALUE_HEADS

Qwen3.6-27B head config (TP=2):

• num_attention_heads = 28 (per card: 28)
• num_key_value_heads = 4 (GQA, per card: 2)
• tp_degree = GLOBAL_KEY_HEADS / num_key_heads = 16 / 2 = 8 ✓
• Check: num_value_heads * tp_degree = 28 * 8 = 224 ✗
• 224 is NOT in GLOBAL_VALUE_HEADS = (16, 32, 48, 64)

Fix Suggestion

Add 224 to GLOBAL_VALUE_HEADS, or make the check dynamic based on actual head count rather than a hardcoded tuple.

Environment

• sgl-kernel-npu: 2026.5.1
• Hardware: Ascend 910B4-1
• torch: 2.10.0, torch_npu: 2.10.0
• Model: Qwen3.6-27B, TP=2