[NPU/Ascend] RuntimeError: unsupported mega_chunk_gdn (NumValueHeads, NumKeyHeads) pair on Qwen3-6-27B with GDN backend

[AndyLi429]

Environment

Model: Qwen3-6-27B (1P config)
CANN:9.0.0
Device: Ascend NPU
Attention backend: ascend (GDN triton path)
Speculative decoding: NEXTN

Launch Command

PYTORCH_NPU_ALLOC_CONF=expandable_segments:True \
STREAMS_PER_DEVICE=32 \
HCCL_SOCKET_IFNAME=lo \
GLOO_SOCKET_IFNAME=lo \
HCCL_OP_EXPANSION_MODE=AIV \
SGLANG_SET_CPU_AFFINITY=1 \
SGLANG_ENABLE_SPEC_V2=1 \
SGLANG_ENABLE_OVERLAP_PLAN_STREAM=0 \
SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE=1 \
SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES=130 \
python -m sglang.launch_server \
    --tp-size 2 \
    --nnodes 1 \
    --attention-backend ascend \
    --device npu \
    --chunked-prefill-size -1 \
    --max-prefill-tokens 60000 \
    --disable-radix-cache \
    --trust-remote-code \
    --max-running-requests 54 \
    --max-mamba-cache-size 64 \
    --mem-fraction-static 0.7 \
    --cuda-graph-bs 2 8 16 32 45 54 \
    --enable-multimodal \
    --quantization modelslim \
    --mm-attention-backend ascend_attn \
    --dtype bfloat16 \
    --mamba-ssm-dtype bfloat16 \
    --speculative-algorithm NEXTN \
    --speculative-num-steps 3 \
    --speculative-eagle-topk 1 \
    --speculative-num-draft-tokens 4

Error
Server crashes during prefill/extend with:
RuntimeError: unsupported mega_chunk_gdn (NumValueHeads, NumKeyHeads) pair
Full Stack Trace

File ".../sglang/srt/models/qwen3_5.py", line 511, in forward
    core_attn_out = self.attn(
File ".../torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
File ".../sglang/srt/layers/radix_linear_attention.py", line 111, in forward
    return get_attn_backend().forward(
File ".../sglang/srt/layers/attention/hybrid_linear_attn_backend.py", line 1005, in forward
    return self.forward_extend(
File ".../sglang/srt/layers/attention/hybrid_linear_attn_backend.py", line 957, in forward_extend
    return self.linear_attn_backend.forward_extend(
File ".../sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py", line 338, in forward_extend
    core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend(
File ".../sglang/srt/layers/attention/linear/gdn_backend.py", line 222, in extend
    return self.extend_kernel.extend(
File ".../sglang/srt/layers/attention/linear/kernels/gdn_triton.py", line 147, in extend
    return chunk_gated_delta_rule(
File ".../sgl_kernel_npu/fla/utils.py", line 86, in wrapper
    return fn(*contiguous_args, **contiguous_kwargs)
File ".../sgl_kernel_npu/fla/chunk.py", line 384, in chunk_gated_delta_rule_npu
    _, o, _, final_state, _, h, _ = chunk_gated_delta_rule_fwd(
File ".../sgl_kernel_npu/fla/chunk.py", line 219, in chunk_gated_delta_rule_fwd
    g, o, A, final_state, w, h, v_new = run_mega_chunk_gdn(
File ".../sgl_kernel_npu/fla/mega_chunk_gdn.py", line 207, in run_mega_chunk_gdn
    torch.ops.npu.mega_chunk_gdn(
File ".../torch/_ops.py", line 1209, in __call__
    return self._op(*args, **kwargs)
RuntimeError: unsupported mega_chunk_gdn (NumValueHeads, NumKeyHeads) pair

[learning-chip]

--tp-size 2 hits the currently unsupported num_head configs in #462

Format: (num val heads, num key heads; bold values are supported)

Global $H_v$	TP=1	TP=2	TP=4	TP=8
16	(16,16)	(8,8)	(4,4)	(2,2)
32	(32,16)	(16,8)	(8,4)	(4,2)
48	(48,16)	(24,8)	(12,4)	(6,2)
64	(64,16)	(32,8)	(16,4)	(8,2)

27B model TP=2 hits the (24,8) case.

@fiskrt For unsupported num_head, it should throw an error msg tell user how to fall back to original triton backend

[fiskrt]

For a temporary fix @AndyLi429 you can run it with the environment variable GDN_ATTN_BACKEND_TRITON=1 set.

A long-term fix is waiting in #517 to be merged.

And yes, @learning-chip, agree that the error message should mention how to force the Triton backend instead.

[fiskrt]

@AndyLi429 It's now fixed.

Here you can see the speedup for prefill for 27B with TP=1 and TP=2

[AndyLi429]

@fiskrt hi,I test the lastest kernel,some problem still occurs.

# here is the launch command

echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor
sysctl -w vm.swappiness=0
sysctl -w kernel.numa_balancing=0
sysctl -w kernel.sched_migration_cost_ns=50000

export SGLANG_SET_CPU_AFFINITY=1
unset https_proxy
unset http_proxy
unset HTTPS_PROXY
unset HTTP_PROXY

source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh

export PYTHONPATH=/home/l00567497/sglang/python:$PYTHONPATH
export ASCEND_RT_VISIBLE_DEVICES=6,7,14,15
export STREAMS_PER_DEVICE=32
export HCCL_OP_EXPANSION_MODE=AIV
export HCCL_SOCKET_IFNAME=lo
export GLOO_SOCKET_IFNAME=lo
export SGLANG_NPU_FORWARD_NATIVE_GEMMA_RMS_NORM=1
export SGLANG_ENABLE_SPEC_V2=1
#export TORCHDYNAMO_VERBOSE=1
#export TORCH_LOGS="+dynamo"
export SGLANG_NPU_PROFILING=1
export SGLANG_NPU_PROFILING_STAGE="prefill"
export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1
export SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE=1
export SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES=100

python3 -m sglang.launch_server \
    --model-path /home/weight/Qwen3.6-27B \
    --attention-backend ascend \
    --device npu \
    --tp-size 4 --nnodes 1 --node-rank 0 \
    --chunked-prefill-size -1 --max-prefill-tokens 72000 \
    --disable-radix-cache \
    --trust-remote-code \
    --host 127.0.0.1 --max-running-requests 64 --max-mamba-cache-size 100 \
    --mem-fraction-static 0.7 \
    --port 6688 \
    --cuda-graph-bs 2 4 8 16 32 48 64 60 64 \
    --enable-multimodal \
    --mm-attention-backend ascend_attn \
    --dtype bfloat16 --mamba-ssm-dtype bfloat16 \
    --speculative-algorithm NEXTN \
    --speculative-num-steps 3 \
    --speculative-eagle-topk 1 \
    --speculative-num-draft-tokens 4 \

And Here is the fault stack:

[2026-06-01 03:32:29 TP3] Scheduler hit an exception: Traceback (most recent call last):
  File "/home/l00567497/sglang/python/sglang/srt/managers/scheduler.py", line 3818, in run_scheduler_process
    scheduler.run_event_loop()
  File "/home/l00567497/sglang/python/sglang/srt/managers/scheduler.py", line 1387, in run_event_loop
    dispatch_event_loop(self)
  File "/home/l00567497/sglang/python/sglang/srt/managers/scheduler.py", line 3688, in dispatch_event_loop
    scheduler.event_loop_overlap()
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/managers/scheduler.py", line 1448, in event_loop_overlap
    batch_result = self.run_batch(batch)
                   ^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/managers/scheduler.py", line 2828, in run_batch
    batch_result = self.model_worker.forward_batch_generation(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/speculative/eagle_worker_v2.py", line 701, in forward_batch_generation
    batch_output = self.target_worker.forward_batch_generation(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/managers/tp_worker.py", line 468, in forward_batch_generation
    out = self.model_runner.forward(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/model_executor/model_runner.py", line 3043, in forward
    output = self._forward_raw(
             ^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/model_executor/model_runner.py", line 3167, in _forward_raw
    ret, can_run_graph = self.forward_extend(
                         ^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/model_executor/model_runner.py", line 2966, in forward_extend
    self.model.forward(
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/models/qwen3_vl.py", line 1272, in forward
    hidden_states = general_mm_embed_routine(
                    ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/managers/mm_utils.py", line 1106, in general_mm_embed_routine
    hidden_states = language_model(
                    ^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/models/qwen3_5.py", line 1046, in forward
    hidden_states, residual = layer(
                              ^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/models/qwen3_5.py", line 617, in forward
    hidden_states = self.linear_attn(
                    ^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/models/qwen3_5.py", line 497, in forward
    core_attn_out = self.attn(
                    ^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/layers/radix_linear_attention.py", line 95, in forward
    return forward_batch.attn_backend.forward(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py", line 924, in forward
    return self.forward_extend(
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py", line 875, in forward_extend
    return self.linear_attn_backend.forward_extend(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py", line 338, in forward_extend
    core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend(
                                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/layers/attention/linear/gdn_backend.py", line 200, in extend
    return self.extend_kernel.extend(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/l00567497/sglang/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py", line 143, in extend
    return chunk_gated_delta_rule(
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/sgl_kernel_npu/fla/utils.py", line 86, in wrapper
    return fn(*contiguous_args, **contiguous_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 1181, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/sgl_kernel_npu/fla/chunk.py", line 382, in chunk_gated_delta_rule_npu
    _, o, _, final_state, _, h, _ = chunk_gated_delta_rule_fwd(
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/sgl_kernel_npu/fla/chunk.py", line 217, in chunk_gated_delta_rule_fwd
    g, o, A, final_state, w, h, v_new = run_mega_chunk_gdn(
                                        ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/sgl_kernel_npu/fla/mega_chunk_gdn.py", line 156, in run_mega_chunk_gdn
    torch.ops.npu.mega_chunk_gdn(
  File "/usr/local/python3.11.15/lib/python3.11/site-packages/torch/_ops.py", line 1209, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: unsupported mega_chunk_gdn (NumValueHeads, NumKeyHeads) pair Try running the alternative backend by setting GDN_ATTN_BACKEND_TRITON=1

I want to use this kernel,but I still fail.

[filipzrc]

TP=4 is supported in #533 . (TP=8/16 is still not supported for Qwen 27B)

Let us know @AndyLi429 if you run in to any more problems