Skip to content

[NPU] Add AscendC mHC operators for TeleChat4 - #670

Draft
momaekar1 wants to merge 4 commits into
sgl-project:mainfrom
momaekar1:codex/telechat4-mhc-npu
Draft

[NPU] Add AscendC mHC operators for TeleChat4#670
momaekar1 wants to merge 4 commits into
sgl-project:mainfrom
momaekar1:codex/telechat4-mhc-npu

Conversation

@momaekar1

@momaekar1 momaekar1 commented Aug 7, 2026

Copy link
Copy Markdown

Motivation

TeleChat4-29B uses four mHC residual streams with hidden_size=3584. A Torch decomposition creates many eager launches and intermediate tensors on Ascend NPU. Huawei's cann-recipes-infer already provides fused AscendC HcPre and HcPost implementations, but sgl-kernel-npu did not package or expose them and its binding validation did not include 3584.

This PR integrates those official operators into the existing sgl-kernel-npu build and wheel layout. The corresponding SGLang integration is sgl-project/sglang#33972.

Modifications

  • Import the official HcPre/HcPost sources from cann-recipes-infer@3fdb034, retaining their original copyright and CANN Open Software License Agreement notices.
  • Reuse csrc/attentions/build/build_ascendc_ops.sh from the main ./build.sh -a kernels entrypoint.
  • Register torch.ops.npu.hc_pre and torch.ops.npu.hc_post, including Meta implementations.
  • Add hidden_size=3584 beside the existing 4096/7168 binding validation.
  • Package the generated vendor tree in the wheel and initialize its custom OPP path on import.
  • Add focused correctness tests and a reproducible steady-state benchmark.

The PR does not add a second custom-ops CMake scaffold, compiled artifacts, wheels, cache directories, or a duplicate ops_common.h.

Build and use from source

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

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="$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

For the complete TeleChat4 TP=2 launch commands and SGLANG_TELECHAT4_MHC_BACKEND=auto|ascendc|triton|torch, see the dependent SGLang PR.

Validation

  • bash -n build.sh and git diff --check: passed.
  • ./build.sh -a kernels: passed; the normal wheel contains libsgl_kernel_npu.so and the aie_ascendc vendor payload.
  • pytest -q tests/python/sgl_kernel_npu/test_mhc.py: 4 passed on Ascend 910B3.
  • Real NPU HcPre/HcPost correctness checked at 1, 8, and 129 tokens against an HF32-aware Torch reference.
  • Meta dispatch accepts 3584 and rejects an unsupported 2048 hidden size.
  • Coordinated SGLang tests: 19 passed (17 mHC dispatch/layout tests and 2 shared-expert top-k tests).
  • Native TeleChat4 TP=2 BF16 weight loading and API smoke completed through the AscendC backend.

Performance

Steady-state Ascend 910B3 latency, 10 warmups and 50 measured iterations:

Tokens HcPre HcPost Combined
1 0.0545 ms 0.0322 ms 0.0891 ms
16 0.0546 ms 0.0331 ms 0.0898 ms
128 0.0688 ms 0.0313 ms 0.0968 ms
512 0.1451 ms 0.0594 ms 0.2251 ms
1024 0.2354 ms 0.1322 ms 0.3888 ms

msprof op at 512 tokens measured HcPre at 135.1-136.3 us (20 AIC + 40 AIV) and HcPost at 48.8-49.3 us (40 AIV).

The end-to-end 1319-question GSM8K run in the SGLang PR measured 291.248 output token/s for AscendC, 277.691 token/s for Triton split-direct, and 107.135 token/s for Torch fallback under TP=2 BF16, concurrency 16, and NEXTN steps=3.

@momaekar1 momaekar1 changed the title [codex] Add official AscendC mHC operators for TeleChat4 [NPU] Add AscendC mHC operators for TeleChat4 Aug 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant