[NPU] Add AscendC mHC operators for TeleChat4 - #670
Draft
momaekar1 wants to merge 4 commits into
Draft
Conversation
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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'scann-recipes-inferalready 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
cann-recipes-infer@3fdb034, retaining their original copyright and CANN Open Software License Agreement notices.csrc/attentions/build/build_ascendc_ops.shfrom the main./build.sh -a kernelsentrypoint.torch.ops.npu.hc_preandtorch.ops.npu.hc_post, including Meta implementations.hidden_size=3584beside the existing 4096/7168 binding validation.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.
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.shandgit diff --check: passed../build.sh -a kernels: passed; the normal wheel containslibsgl_kernel_npu.soand theaie_ascendcvendor payload.pytest -q tests/python/sgl_kernel_npu/test_mhc.py: 4 passed on Ascend 910B3.Performance
Steady-state Ascend 910B3 latency, 10 warmups and 50 measured iterations:
msprof opat 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.