[Bug] mega_chunk_gdn unsupported head config for Qwen3.6-27B GDN attention

[JosephNew]

Description

mega_chunk_gdn operator does not support the GDN attention head configuration used by Qwen3.6-27B, causing inference crash on Ascend NPU.

Error

RuntimeError: unsupported mega_chunk_gdn (NumValueHeads, NumKeyHeads) pair

Root Cause

fla/mega_chunk_gdn.py defines strict head config constraints:

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, num_key_heads):
    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's GDN head config (NumValueHeads, NumKeyHeads) does not satisfy these constraints.

Reproduction

On Ascend 910B4-1, torch 2.10 + torch_npu 2.10, sgl-kernel-npu 2026.5.1:

sglang serve --model-path /path/to/Qwen3.6-27B --tp 2 --disable-cuda-graph --disable-radix-cache --trust-remote-code

curl http://127.0.0.1:30000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"default","messages":[{"role":"user","content":"hi"}],"max_tokens":16}'

Full call chain: chunk_gated_delta_rule_fwd -> run_mega_chunk_gdn -> torch.ops.npu.mega_chunk_gdn

Request

Please add support for Qwen3.6-27B's GDN head configuration, or provide guidance on how to extend GLOBAL_VALUE_HEADS to cover this model.