[McZyWu]

Summary

Avoid running the CUDA-oriented Triton autotune search for the Kimi-K3 FLA cumsum/KDA kernels on Ascend. The NPU path launches the raw JIT kernels with the A3 configurations validated below; CUDA, ROCm, Intel XPU, and other non-NPU backends keep the existing autotuned wrappers.

This is intentionally limited to the operator layer:

• chunk_local_cumsum_vector: BS=64, num_warps=4, num_stages=3
• KDA inter-chunk kernel: BK=64, num_warps=4, num_stages=3
• KDA intra-chunk kernel: num_warps=4, num_stages=3
• backend selection is guarded by sglang.srt.utils.is_npu()

Why

The K3 CI failure has two different Triton Ascend behaviors that should not be conflated:

1. With the older Triton Ascend runtime used by the failing job, scheduler stacks stopped in
   chunk_kda_scaled_dot_kkt_fwd -> triton.autotuner.do_bench -> torch_npu.synchronize.
2. After force-installing triton_ascend==3.2.2, the autotuner/JIT can finish. Fresh-cache stack samples progressed through cumsum, KDA, recompute/state kernels, and later operators; the cache file count kept increasing. However, the first request still spent a long time compiling/benchmarking CUDA-style candidates independently on each of the four nodes. While one node was in linalg_to_bin_enable_npu_compile_A2_A3, peers waited in model forward / DP collectives, the detokenizer heartbeat could become stale, and /health temporarily returned 503.

The public kernels currently enumerate 18 cumsum configs, 24 KDA inter configs, and 4 KDA intra configs for each new autotune key. This PR preserves that search for non-NPU backends and avoids it only on Ascend.

Four-node validation

Environment:

Item	Value
Base	sgl-project/sglang@5f216fc33f315ade378bc7cd295ee88f5322d4d3
Relevant merged PRs in base	#35314 (2d8484740d5e...), #36603 (69a49fede863...)
Topology	4 nodes x 16 Ascend NPUs, TP=64, DP=4, EP=64
Container	CANN 9.1, sglwmc-813
Triton wheel	triton_ascend-3.2.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl
Wheel SHA256	bb02d9eb2181172f783aefdb2b4505fb7c52bf93c506d43290bac58b2b3f4e7e on all four nodes
Imported runtime	triton.__version__ == 3.2.0, Python 3.12
Target	/home/weights/Kimi-K3-w4a8-int-moe, ModelSlim W4A8
Draft	/home/weights/Kimi-K3-DSpark, DSpark
Important launch options	--attention-backend ascend --linear-attn-verify-backend triton --tp-size 64 --enable-dp-attention --dp-size 4 --moe-a2a-backend deepep --speculative-algorithm DSPARK --cuda-graph-bs 11 16

Each run used a new per-node TRITON_CACHE_DIR; no compiled Triton cache was shared across nodes.

No-code baseline on Triton Ascend 3.2.2

Check	Result
Target graph capture	completed, 82.09-84.77 s across rank-0 workers
Draft graph capture	completed, 6.62-6.92 s
Built-in 448x448 K3 VLM warmup	HTTP 200, 04:07:31 -> 04:11:07 (~216 s)
Follow-up health	HTTP 503 at 04:14:27
Stack evidence	one node was still compiling in linalg_to_bin_enable_npu_compile_A2_A3; other nodes were in model forward / DP synchronization. A later sample showed DP ranks entering the two maybe_prepare_mlp_sync_batch all-gathers from different scheduler sites.

This proves that 3.2.2 can compile the operators; changing only the package version did not make the fresh-cache service reliably ready within the observed window.

Patched fresh-cache run

Check	Result
Target weights	all 351 shards loaded; 218.27-231.27 s on rank-0 workers
DSpark weights	loaded; 0.72-0.99 s
Target graph capture	completed; max 79.45 s
Draft graph capture	completed; max 7.03 s
Built-in 448x448 K3 VLM warmup	HTTP 200, 04:33:37 -> 04:34:53 (~76 s)
First external health while other first-use JIT was active	HTTP 503 at 04:37:12
First external text request	HTTP 200 at 04:38:58, 16 generated tokens
Follow-up health	HTTP 200 at 04:38:58; a separately timed check was HTTP 200 in 10.013 s
Second external text request	HTTP 200, 8 generated tokens, 28.329 s
Post-GSM8K health	HTTP 200 at 04:47:52

The temporary 503 is a stale-heartbeat symptom while another node is doing first-use JIT; it is not evidence that the 3.2.2 compiler or the patched kernel is permanently stuck. The external request completed once the lagging node finished compilation.

Operator numerical smoke

Single-NPU tests compared the patched operators against Torch references using a new Triton cache:

Operator path	Max absolute error
cumsum	0.00000000
KDA triangular component	0.00000002
KDA query-key component	0.00000000

Result: operator_smoke=PASS.

GSM8K-200

Command parameters: 5-shot, 200 questions, max_new_tokens=512, parallel=128, temperature=0.

Metric	Patched fresh-cache run	Previous #36603 comparison run
Completed requests	200/200	200/200
Exit status	0	0
Accuracy	0.975	0.970
Invalid	0.005	0.005
Eval latency	338.068 s	156.517 s
Output throughput	59.003 token/s	126.421 token/s
Client wall time	349 s	not recorded separately

The patched run intentionally started from fresh per-node caches. Its first GSM8K completion took 87.81 s while a lagging node compiled additional input-shape kernels, so this end-to-end throughput is not a hot-cache performance comparison. Correctness did not regress, all 200 requests completed, and the service remained healthy afterward.

Compatibility with #35314

#35314 is an ancestor of this branch. There is no changed-file overlap: #35314 adds K3/DeepSeek SSD expert-pack and GGUF loader/model paths, whereas this PR changes only:

• python/sglang/kernels/ops/attention/fla/cumsum.py
• python/sglang/kernels/ops/attention/fla/kda.py

No weight loading, SSD, expert-pack, GGUF, quantization, or model code is modified. The four-node latest-main run successfully loaded the K3 W4A8 target and DSpark draft. The exact SSD/GGUF artifact from #35314 was not part of this hardware run, so that compatibility statement is based on the disjoint code paths plus successful latest-main loading, not a separate SSD throughput benchmark.

Other checks

• python -m py_compile for both modified files: PASS on all four nodes
• identical modified-file SHA256 on all four nodes: PASS
• git diff --check: PASS
• final branch contains only the two operator files above

---

CI States

Latest PR Test (Base): ❌ Run #33234977390
Latest PR Test (Extra): ❌ Run #33234977218
Latest PR Test (AMD ROCm 7.2): ❌ Run #33234977328

[McZyWu]

/tag-and-rerun-ci