[NPU] Add TeleChat4 support with fused mHC backends

[momaekar1]

Motivation

TeleChat4-29B uses four mHC residual streams with hidden_size=3584. On Ascend NPU, the previous Torch decomposition generated many eager operators, synchronization points, and intermediate tensors. The generic CUDA mHC path is not an NPU implementation.

This PR adds native TeleChat4 registration and explicit NPU mHC backends while preserving the existing CUDA path. The AscendC backend depends on sgl-project/sgl-kernel-npu#670. If that package is unavailable, auto falls back to Triton and then Torch.

Modifications

• Register TeleChat4Config, TeleChat4ForCausalLM, the TeleChat4 tool parser, and reasoning parser.
• Accept released checkpoints with native model_type: telechat4 and convert legacy rope_scaling fields to the Transformers v5 rope_parameters representation.
• Add SGLANG_TELECHAT4_MHC_BACKEND=auto|ascendc|triton|torch; auto resolves as AscendC -> Triton -> Torch on NPU.
• Keep the NPU Triton prefill implementation as two kernels (split-K GEMM and finalize). Finalize writes the HcPost physical layout directly, removing an extra transpose copy.
• Append shared-expert indices after routed top-k indices on NPU and add focused coverage for the fused path.
• Leave CUDA dispatch on the existing sglang.kernels.ops.layernorm.mhc implementation.

Checkpoint configuration

The checkpoint config.json must identify the native class:

{
  "architectures": ["TeleChat4ForCausalLM"],
  "model_type": "telechat4"
}

The released legacy rope_scaling dictionary is supported. No --json-model-override-args is required with this configuration.

Build and source setup

This setup reuses the existing Torch, torch-npu, and CANN installation; it does not reinstall them.

export SGLANG_SRC=/path/to/sglang
export SGL_KERNEL_NPU_SRC=/path/to/sgl-kernel-npu

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

cd "$SGL_KERNEL_NPU_SRC"
./build.sh -a kernels

source "$SGL_KERNEL_NPU_SRC/python/sgl_kernel_npu/sgl_kernel_npu/vendors/aie_ascendc/bin/set_env.bash"
export PYTHONPATH="$SGLANG_SRC/python:$SGL_KERNEL_NPU_SRC/python/sgl_kernel_npu:${PYTHONPATH:-}"

python3 - <<'PY'
import torch
import torch_npu
import sgl_kernel_npu

print("NPU available:", torch.npu.is_available())
print("hc_pre:", hasattr(torch.ops.npu, "hc_pre"))
print("hc_post:", hasattr(torch.ops.npu, "hc_post"))
PY

Minimal TP=2 smoke launch

export ASCEND_RT_VISIBLE_DEVICES=0,1
export SGLANG_TELECHAT4_MHC_BACKEND=auto

python3 -m sglang.launch_server \
  --model-path /path/to/TeleChat4-29B \
  --tp-size 2 \
  --device npu \
  --attention-backend ascend \
  --dtype bfloat16 \
  --moe-a2a-backend none \
  --cuda-graph-backend-decode disabled \
  --cuda-graph-backend-prefill disabled \
  --disable-overlap-schedule \
  --disable-radix-cache \
  --disable-shared-experts-fusion \
  --context-length 2048 \
  --chunked-prefill-size 1024 \
  --max-prefill-tokens 2048 \
  --max-running-requests 1 \
  --mem-fraction-static 0.7 \
  --trust-remote-code \
  --host 0.0.0.0 \
  --port 30000

NEXTN throughput launch

Graph capture, RadixCache, and overlap scheduling remain enabled by omitting their disable flags.

export ASCEND_RT_VISIBLE_DEVICES=0,1
export SGLANG_TELECHAT4_MHC_BACKEND=auto

python3 -m sglang.launch_server \
  --model-path /path/to/TeleChat4-29B \
  --tp-size 2 \
  --device npu \
  --attention-backend ascend \
  --dtype bfloat16 \
  --moe-a2a-backend none \
  --speculative-algorithm NEXTN \
  --speculative-num-steps 3 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 4 \
  --disable-shared-experts-fusion \
  --context-length 131072 \
  --chunked-prefill-size 1024 \
  --max-prefill-tokens 8192 \
  --max-running-requests 16 \
  --mem-fraction-static 0.7 \
  --trust-remote-code \
  --host 0.0.0.0 \
  --port 30000

Accuracy Tests

• Native checkpoint config loaded as TeleChat4Config / TeleChat4ForCausalLM with json_model_override_args='{}'; TP=2 BF16 weight loading and API dialogue smoke completed.
• pytest -q test/registered/unit/configs/test_telechat4_config.py: 1 passed.
• pytest -q test/srt/models/test_telechat4_mhc_triton_npu.py test/registered/unit/npu/test_npu_topk_shared_expert.py: 19 passed on Ascend 910B3.
• Coordinated kernel tests in sgl-kernel-npu: 4 passed at 1, 8, and 129 tokens, including 3584 shape validation.

Full GSM8K: 1319 questions, dataset SHA256 fb581f0270b25988e071316835842a1c8449f4e27af6fc0f539ef270b987f9ff, 5-shot, temperature 0, TP=2, BF16, concurrency 16, NEXTN steps=3, four draft tokens, Graph/Radix/overlap enabled.

Backend	Accuracy	Invalid
AscendC	85.2%	0.2%
Triton split-direct	85.0%	0.2%
Torch fallback	84.9%	0.3%

Speed Tests and Profiling

Backend	End-to-end latency	Output throughput
AscendC	497.391 s	291.248 token/s
Triton split-direct	525.785 s	277.691 token/s
Torch fallback	1358.826 s	107.135 token/s

Triton and Torch were a same-container/card A/B and had nearly identical NEXTN acceptance (2.9299 / 64.33% vs. 2.9241 / 64.16%). Triton is 2.592x the Torch fallback throughput. The AscendC result came from the existing 0721 integration container, so its 4.882% lead over Triton also includes container/framework differences and is not attributed solely to the mHC kernel.

At 512 tokens, msprof op measured AscendC HcPre at 135.1-136.3 us and HcPost at 48.8-49.3 us. Reproducible kernel microbenchmarks are included in the dependent kernel PR.

Checklist

• Format Python/C++ and pass all applicable pre-commit checks.
• Add CPU and NPU unit tests and register them with SGLang CI.
• Document checkpoint, build, environment, and launch requirements in this PR.
• Provide accuracy, end-to-end speed, kernel correctness, and msprof results.
• Follow the existing SGLang CUDA/NPU backend boundaries.