From 61e57d253b827187e9609349d3bce28889ce2534 Mon Sep 17 00:00:00 2001 From: dabuliu123 <270334047@qq.com> Date: Wed, 9 Sep 2026 10:08:10 +0800 Subject: [PATCH 1/3] feat(NPU): add latest sglang optimize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch 'wqw_base_lw_dev_0808' into ascend-dev-0808 feat(docker): upgrade NPU env to CANN 9.0.0 with FLA and MindSpeed-Ops # โญ Feature ## Upgrade CANN base image and toolchain - Upgrade CANN base image from 8.5.1-a3 to 9.0.0-a3 - Upgrade torch_npu from v2.9.0-7.3.0 to v26.0.1-pytorch2.9.0 - Upgrade triton-ascend from 3.2.0 to 3.2.1 ## Add AscendC Flash Linear Attention (FLA) support - Clone and build fla_npu from flash-linear-attention-npu v26.1.0 - Compile causal_conv1d and gated_delta_rule ops for ascend910_93 ## Add MindSpeed-Ops support - Clone, checkout and install MindSpeed-Ops - Add mindspeed-ops.patch for gated_delta_rule autotune key fix --- # โ™ป๏ธ Refactor ## Restructure Dockerfile build order - Move torch/torch_npu install before repo clone - Delay COPY . /root/Relax to just before patching - Consolidate Megatron-Bridge into Megatron-LM via cp instead of separate path - Install MindSpeed-Bridge with --no-deps to avoid circular dependency ## Migrate mindspeed-bridge to FLA ops - Replace causal_conv1d from causal_conv1d with FLA implementation - Replace mindspeed_ops l2norm with naive_l2norm fallback - Switch to flash_gated_delta_rule when FLA is available --- # ๐Ÿ› Bug Fix ## Fix various compatibility issues - Fix autotune key in gated_delta_rule kernel by removing USE_G and IS_VARLEN - Fix lambda closure bug in transformer_config_init_subclass (use default arg binding) - Fix attention_mask dtype from int32 to bool in model preprocess - Fix vision model config with MoE permute fusion disabled - Handle OmegaConf DictConfig/ListConfig in remove_non_pickleables - Guard apex MixedFusedLayerNorm import with is_npu_available check ## Update sgl-kernel-npu - Upgrade sgl-kernel-npu checkout from 2026.04.15.rc3 to 2026.7.2 - Remove obsolete cherry-pick workaround feat(docker): upgrade NPU env to CANN 9.0.0 with FLA and MindSpeed-Ops feat(docker): upgrade NPU env to CANN 9.0.0 with FLA and MindSpeed-Ops # โญ Feature ## Upgrade CANN base image and toolchain - Upgrade CANN base image from 8.5.1-a3 to 9.0.0-a3 - Upgrade torch_npu from v2.9.0-7.3.0 to v26.0.1-pytorch2.9.0 - Upgrade triton-ascend from 3.2.0 to 3.2.1 ## Add AscendC Flash Linear Attention (FLA) support - Clone and build fla_npu from flash-linear-attention-npu v26.1.0 - Compile causal_conv1d and gated_delta_rule ops for ascend910_93 ## Add MindSpeed-Ops support - Clone, checkout and install MindSpeed-Ops - Add mindspeed-ops.patch for gated_delta_rule autotune key fix --- # โ™ป๏ธ Refactor ## Restructure Dockerfile build order - Move torch/torch_npu install before repo clone - Delay COPY . /root/Relax to just before patching - Consolidate Megatron-Bridge into Megatron-LM via cp instead of separate path - Install MindSpeed-Bridge with --no-deps to avoid circular dependency ## Migrate mindspeed-bridge to FLA ops - Replace causal_conv1d from causal_conv1d with FLA implementation - Replace mindspeed_ops l2norm with naive_l2norm fallback - Switch to flash_gated_delta_rule when FLA is available --- # ๐Ÿ› Bug Fix ## Fix various compatibility issues - Fix autotune key in gated_delta_rule kernel by removing USE_G and IS_VARLEN - Fix lambda closure bug in transformer_config_init_subclass (use default arg binding) - Fix attention_mask dtype from int32 to bool in model preprocess - Fix vision model config with MoE permute fusion disabled - Handle OmegaConf DictConfig/ListConfig in remove_non_pickleables - Guard apex MixedFusedLayerNorm import with is_npu_available check ## Update sgl-kernel-npu - Upgrade sgl-kernel-npu checkout from 2026.04.15.rc3 to 2026.7.2 - Remove obsolete cherry-pick workaround [NPU] update sglang fix(megatron): attach _hf_config for MTP bridge when pp>1 # ๐Ÿ› Bug Fix ## Fix bridge mapping registry failure with PP > 1 - Set model_bridge._hf_config from bridge.hf_pretrained.config when missing chore(docker): refresh NPU patch set # ๐Ÿ”ฉ Chore ## Sync NPU patches - Add sgl-kernel-npu.patch (1847 lines) - Update sglang-npu.patch (+1070 lines) - Update mindspeed.patch / mindspeed-bridge.patch / megatron patches build(docker): upgrade NPU image to CANN 9.0.0 # ๐Ÿ”ฉ Chore ## Upgrade NPU Dockerfile dependencies - Bump base image CANN 8.5.1 โ†’ 9.0.0 - Upgrade torch_npu to v26.0.1 (PyTorch 2.9.0) - Upgrade triton-ascend 3.2.0 โ†’ 3.2.1 feat(npu): add Qwen3.5 MTP SFT/training scripts feat(npu): add Qwen3.5 MTP training scripts for 0821 # โญ Feature ## Add Qwen3.5 MTP training launch scripts - run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh: Qwen3.5-35B-A3B MTP SFT on pokemon-gpt4o-captions, 8xNPU single-node, ray-submit launch - run_qwen35-35B-A3B-16xnpu-colocate-thd.sh: Qwen3.5-35B-A3B colocate THD training on 16xNPU - run_qwen35_9B_mtp_8xnpu_thd.sh: Qwen3.5-9B MTP THD training on 8xNPU --- # ๐Ÿ› Bug Fix ## Fix EXP_DIR silently overridden by MODEL_DIR default - EXP_DIR now uses its own default with MODEL_DIR following EXP_DIR, matching the 9B THD script pattern --- # ๐Ÿ”ฉ Chore ## Align script naming and comments - Rename scripts to the 8xnpu naming convention and drop duplicated .sh suffix - Remove commented-out --qkv-format bshd / --micro-batch-size 1 lines - Sync Usage comments with actual script names Revert "fix(megatron): attach _hf_config for MTP bridge when pp>1" This reverts commit 57809b43772529b339dc5ba2a96f25594a711a71. fix: update qwen35 mtp mapping in mindspeed-bridge.patch for relax bridge_converter.py fix(npu): rename Qwen3.5-9B CP script to 16xnpu feat(npu): add Qwen3.5-9B CP training script # โญ Feature ## Add Qwen3.5-9B CP colocate training script - Add run-qwen35-9B-8xnpu-cp.sh for DAPO math training with TP4/CP4 on 16 NPUs - Set MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV=1 to use torch conv fallback for GDN ## Support GDN causal conv on NPU via fla_npu - Replace causal_conv1d with fla_npu in mindspeed-bridge.patch GDN layer - Add MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV env switch to fall back to torch conv --- # ๐Ÿ› Bug Fix ## Fix repatch ordering before Megatron init - Move repatch(args) before init(args) in MegatronTrainRayActor so bridge patches apply during initialization fix(npu): rename Qwen3.5-9B CP script to 16xnpu # ๐Ÿ› Bug Fix ## Fix Qwen3.5-9B CP script naming for 16 NPUs - Rename run-qwen35-9B-8xnpu-cp.sh to run-qwen35-9B-16xnpu-cp.sh to match the actual 16-NPU setup - Update header comment to 16xNPU and log filename to qwen35-9B-GRPO-npu16 docs(npu): add feature support table # ๐Ÿ“ Documentation ## Add feature support table to NPU training guide - Add ็‰นๆ€งๆ”ฏๆŒ table covering Qwen3.5-9B CP, Qwen3.5-9B MTP, and Qwen3.5-35B-A3B SFT MTP with reference scripts - Update ไธ‹ไธ€ๆญฅ checklist from feature support to performance optimization - Fix MTP row minimum card count to 4ๅก to match the script's 8 NPUs fix(gitleaks): drop /mnt/ defaults in NPU scripts # ๐Ÿ”’ Security ## Remove hardcoded internal paths from Qwen3.5 NPU scripts - Replace /mnt/tidalfs-hwwh01 EXP_DIR defaults with the repo-standard `${SCRIPT_DIR}/../../../../exps` in run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh, run_qwen35-35B-A3B-16xnpu-colocate-thd.sh and run_qwen35_9B_mtp_8xnpu_thd.sh - Replace bare /mnt/ placeholder in run-qwen35-9B-16xnpu-cp.sh with the same repo-standard default - Switch the SFT script DATA_DIR default to the `${DATA_DIR:-${EXP_DIR}}` convention used by all other training scripts --- # ๐ŸŽจ Style ## End-of-file fixer - Remove trailing blank line from docker/Dockerfile.npu Co-Authored-By: Claude update Dockerfile.npu optimizer npu update_weights Add optimized 9B&35B NPU colocate script Add optimized 9B&35B NPU sglang patch Add optimized 9B&35B NPU sgl-kernel-npu patch FIX dockerfile build error Fix sgl-kernel-npu patch [clean code] sglang&kernel patch --- docker/Dockerfile.npu | 12 +- docker/npu_patch/sgl-kernel-npu.patch | 2914 ++++++++++++++- docker/npu_patch/sglang-npu.patch | 3138 ++++++++++++++++- relax/utils/reloadable_process_group.py | 2 + .../run-qwen35-35B-A3B-16xnpu-colocate.sh | 230 +- .../text/run-qwen35-9B-8xnpu-colocate.sh | 211 ++ 6 files changed, 6265 insertions(+), 242 deletions(-) create mode 100644 scripts/training/text/run-qwen35-9B-8xnpu-colocate.sh diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index ac5f95c8d..7d6d2e3bf 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -12,6 +12,7 @@ ARG NO_PROXY ARG http_proxy ARG https_proxy ARG no_proxy +ARG SOC_VERSION_FOR_SGL_BUILD="Ascend910_9382" ENV http_proxy=${HTTP_PROXY} \ https_proxy=${HTTPS_PROXY} \ @@ -136,15 +137,14 @@ RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org git add . && git commit -m "install info" && \ git fetch && \ git cherry-pick ece02ffc9cc32e94382d4f1b553b2c755f83f722 && \ - patch -p1 < /root/Relax/docker/npu_patch/sglang-npu.patch && \ - git add . && git commit -m "sglang-npu.patch" + git am /root/Relax/docker/npu_patch/sglang-npu.patch # Install SGLang NPU kernels. RUN git clone https://github.com/sgl-project/sgl-kernel-npu /root/sgl-kernel-npu && \ cd /root/sgl-kernel-npu && git checkout 2026.7.2 && \ - patch -p1 < /root/Relax/docker/npu_patch/sgl-kernel-npu.patch && \ - git add . && git commit -m "sgl-kernel-npu.patch" && \ - bash build.sh && \ + git submodule update --init third_party/catlass third_party/pto-isa && \ + git am /root/Relax/docker/npu_patch/sgl-kernel-npu.patch && \ + bash build.sh -a kernels ${SOC_VERSION_FOR_SGL_BUILD} && bash build.sh -a memory-saver && bash build.sh -a deepep && \ pip install output/*.whl && \ cd /root && rm -rf /root/sgl-kernel-npu @@ -168,7 +168,7 @@ WORKDIR /root COPY requirements.txt /tmp/requirements.txt RUN pip install -r /tmp/requirements.txt --no-cache-dir && \ - pip install --index-url https://pypi.org/simple ray==2.55.1 protobuf==6.33.6 + pip install --index-url https://pypi.org/simple ray==2.55.1 protobuf==6.33.6 fastokens==0.3.1 # Clean caches. RUN pip cache purge && \ diff --git a/docker/npu_patch/sgl-kernel-npu.patch b/docker/npu_patch/sgl-kernel-npu.patch index 3ef4d0a96..3b597a8ed 100644 --- a/docker/npu_patch/sgl-kernel-npu.patch +++ b/docker/npu_patch/sgl-kernel-npu.patch @@ -1,3 +1,86 @@ +From 9b321f1e34e7ae4adcf283a845b14852cce0a783 Mon Sep 17 00:00:00 2001 +From: liyongwen <1310439159@qq.com> +Date: Sat, 29 Aug 2026 08:44:22 +0000 +Subject: [PATCH] sgl-kernel-npu optimize for qwen3.5 9B&35B + +--- + build.sh | 6 +- + .../python/torch_memory_saver/entrypoint.py | 8 +- + csrc/CMakeLists.txt | 56 ++ + .../op_host/cache_loc_assign.cpp | 36 +- + .../op_host/tiling/cache_loc_assign.h | 5 + + .../op_kernel/cache_loc_assign_kernel.cpp | 10 +- + csrc/causal_conv1d/op_host/causal_conv1d.cpp | 6 +- + csrc/causal_conv1d/op_kernel/causal_conv1d.h | 14 +- + .../op_kernel/causal_conv1d_fn_tasks.h | 7 +- + .../op_kernel/causal_conv1d_tiling_data.h | 8 + + .../op_host/fused_qkvzba_conv1d.cpp | 408 ++++++++++ + .../op_host/fused_qkvzba_conv1d.h | 37 + + .../stub/aclrtlaunch_fused_qkvzba_conv1d.h | 15 + + .../op_kernel/fused_qkvzba_conv1d_kernel.cpp | 168 ++++ + .../fused_sigmoid_gating_recurrent.cpp | 167 ++++ + .../op_host/fused_sigmoid_gating_recurrent.h | 45 ++ + ...d_sigmoid_gating_recurrent_bf16_kernel.cpp | 16 + + ...d_sigmoid_gating_recurrent_fp32_kernel.cpp | 16 + + ...used_sigmoid_gating_recurrent_kernel_lib.h | 720 ++++++++++++++++++ + csrc/pytorch_extensions.cpp | 46 +- + include/sgl_kenel_npu_ops.h | 6 +- + .../fla/fused_sigmoid_gating_recurrent.py | 36 +- + ...gmoid_gating_recurrent_decode_optimized.py | 328 ++++++++ + .../sgl_kernel_npu/fla/utils.py | 13 +- + .../sgl_kernel_npu/moe/moe_front_routing.py | 182 +++++ + .../sgl_kernel_npu/moe/persistent_gmm.py | 218 ++++++ + .../sgl_kernel_npu/sample/__init__.py | 13 + + .../sample/chain_speculative_sampling.py | 354 +++++++++ + .../sgl_kernel_npu/sample/probability.py | 40 + + .../tree_speculative_sampling_target_only.py | 383 ++++++++++ + .../test_chain_speculative_sampling.py | 138 ++++ + .../test_speculative_probability.py | 31 + + ...t_tree_speculative_sampling_target_only.py | 595 +++++++++++++++ + 33 files changed, 4093 insertions(+), 38 deletions(-) + create mode 100644 csrc/fused_qkvzba_conv1d/op_host/fused_qkvzba_conv1d.cpp + create mode 100644 csrc/fused_qkvzba_conv1d/op_host/fused_qkvzba_conv1d.h + create mode 100644 csrc/fused_qkvzba_conv1d/op_host/stub/aclrtlaunch_fused_qkvzba_conv1d.h + create mode 100644 csrc/fused_qkvzba_conv1d/op_kernel/fused_qkvzba_conv1d_kernel.cpp + create mode 100644 csrc/fused_sigmoid_gating_recurrent/op_host/fused_sigmoid_gating_recurrent.cpp + create mode 100644 csrc/fused_sigmoid_gating_recurrent/op_host/fused_sigmoid_gating_recurrent.h + create mode 100644 csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_bf16_kernel.cpp + create mode 100644 csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_fp32_kernel.cpp + create mode 100644 csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_kernel_lib.h + create mode 100644 python/sgl_kernel_npu/sgl_kernel_npu/fla/fused_sigmoid_gating_recurrent_decode_optimized.py + create mode 100644 python/sgl_kernel_npu/sgl_kernel_npu/moe/moe_front_routing.py + create mode 100644 python/sgl_kernel_npu/sgl_kernel_npu/moe/persistent_gmm.py + create mode 100644 python/sgl_kernel_npu/sgl_kernel_npu/sample/chain_speculative_sampling.py + create mode 100644 python/sgl_kernel_npu/sgl_kernel_npu/sample/probability.py + create mode 100644 python/sgl_kernel_npu/sgl_kernel_npu/sample/tree_speculative_sampling_target_only.py + create mode 100644 tests/python/sgl_kernel_npu/test_chain_speculative_sampling.py + create mode 100644 tests/python/sgl_kernel_npu/test_speculative_probability.py + create mode 100644 tests/python/sgl_kernel_npu/test_tree_speculative_sampling_target_only.py + +diff --git a/build.sh b/build.sh +index a7fc35b..32f4cb9 100755 +--- a/build.sh ++++ b/build.sh +@@ -93,6 +93,10 @@ else + fi + + echo "Use SOC_VERSION: $SOC_VERSION" ++if [[ "$BUILD_KERNELS_MODULE" == "ON" && ! "$SOC_VERSION" =~ ^Ascend910_93[0-9]{2}$ ]]; then ++ echo "Error: tp_ascendc_fusion kernels require Ascend910_93xx, got: $SOC_VERSION" ++ exit 1 ++fi + + echo "=== Fixing ASCConfig for CANN 8.3 / A2 ===" + +@@ -168,7 +172,7 @@ function build_kernels() + -DASCEND_INCLUDE_DIR=$ASCEND_INCLUDE_DIR \ + -DCMAKE_PREFIX_PATH="$ASC_CMAKE_DIR" \ + -DASC_DIR="$ASC_CMAKE_DIR" \ +- -DSOC_VERSION=Ascend910_9382 \ ++ -DSOC_VERSION="$SOC_VERSION" \ + -DDEEPEP_IS_A5_BUILD=$([[ "$SOC_VERSION" == "Ascend950" ]] && echo "ON" || echo "OFF") \ + -DBUILD_DEEPEP_MODULE=$BUILD_DEEPEP_MODULE \ + -DBUILD_KERNELS_MODULE=$BUILD_KERNELS_MODULE \ diff --git a/contrib/torch_memory_saver/python/torch_memory_saver/entrypoint.py b/contrib/torch_memory_saver/python/torch_memory_saver/entrypoint.py index b14c561..7b0508d 100644 --- a/contrib/torch_memory_saver/python/torch_memory_saver/entrypoint.py @@ -48,6 +131,94 @@ index b14c561..7b0508d 100644 - ) + ) \ No newline at end of file +diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt +index e8ec532..86801a9 100644 +--- a/csrc/CMakeLists.txt ++++ b/csrc/CMakeLists.txt +@@ -29,6 +29,9 @@ FILE(GLOB OP_SRCS + ${PROJECT_OP_SRC_BASE}/recurrent_gated_delta_rule/op_host/recurrent_gated_delta_rule.cpp + ${PROJECT_OP_SRC_BASE}/causal_conv1d/op_host/causal_conv1d.cpp + ${PROJECT_OP_SRC_BASE}/apply_token_bitmask/op_host/apply_token_bitmask.cpp ++ # TP_FUSION: fused_qkvzba_conv1d + fused_sigmoid_gating_recurrent ++ ${PROJECT_OP_SRC_BASE}/fused_qkvzba_conv1d/op_host/fused_qkvzba_conv1d.cpp ++ ${PROJECT_OP_SRC_BASE}/fused_sigmoid_gating_recurrent/op_host/fused_sigmoid_gating_recurrent.cpp + ) + if(BUILD_CATLASS_MODULE) + list(APPEND OP_SRCS +@@ -114,6 +117,55 @@ ascendc_compile_definitions(workspace_kernel PRIVATE + -DHAVE_TILING + ) + ++# The two TP_FUSION kernels live in dedicated libraries with --cce-auto-sync ++# pinned explicitly so the global default cannot drift. Both kernels use fully ++# manual synchronization (explicit PipeBarrier/SetFlag) and are correct either ++# way. ++# ++# The two kernels must be in separate libraries: under ++# HAVE_WORKSPACE/HAVE_TILING the toolchain auto_gen launch wrapper expands the ++# signature per the "last two params = workspace/tiling" convention. ++# fused_qkvzba_conv1d ends in GM_ADDR workspace/tiling and matches; ++# fused_sigmoid_gating_recurrent has no workspace/tiling and ends in 8 ++# by-value scalar params, so in the same library auto_gen would mistake the ++# second-to-last float softplusThreshold for a workspace pointer. Hence ++# recurrent goes alone into tp_fusion_recurrent_kernel, without ++# HAVE_WORKSPACE/HAVE_TILING. ++ascendc_library(tp_fusion_kernel STATIC ++ ${PROJECT_OP_SRC_BASE}/fused_qkvzba_conv1d/op_kernel/fused_qkvzba_conv1d_kernel.cpp ++) ++ascendc_include_directories(tp_fusion_kernel PRIVATE ++ ${PROJECT_OP_SRC_BASE}/utils/kernel ++ ${PROJECT_OP_SRC_BASE}/causal_conv1d_update/op_kernel ++ ${PROJECT_OP_SRC_BASE}/causal_conv1d/op_kernel ++) ++ascendc_compile_definitions(tp_fusion_kernel PRIVATE ++ -DHAVE_WORKSPACE ++ -DHAVE_TILING ++) ++ascendc_compile_options(tp_fusion_kernel PRIVATE ++ --cce-auto-sync=on ++) ++ ++# The bf16/fp32 recurrent entries must be one kernel per file: ++# KERNEL_TASK_TYPE_DEFAULT does not coexist with multiple kernels in one file ++# โ€” only the first kernel registers; the rest fail with RegisterAscendBinary ++# aiv ret 107000 and launch hangs with the kernel not found. ++# Shared classes live in op_kernel/fused_sigmoid_gating_recurrent_kernel_lib.h. ++ascendc_library(tp_fusion_recurrent_kernel STATIC ++ ${PROJECT_OP_SRC_BASE}/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_bf16_kernel.cpp ++ ${PROJECT_OP_SRC_BASE}/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_fp32_kernel.cpp ++) ++# Pin =off: this kernel is a double-buffered pipeline with fully manual event ++# discipline (AllocEventID events + DrainEvents draining at the end) and does ++# not use the TQue programming model โ€” auto-sync only applies to the canonical ++# model (EnQue/DeQue/AllocTensor/FreeTensor). =on gains nothing here and only ++# inserts extra compiler-owned events at scalar GetValue/SetValue (same ++# hardware pool as the manual events). ++ascendc_compile_options(tp_fusion_recurrent_kernel PRIVATE ++ --cce-auto-sync=off ++) ++ + set(MEGA_CHUNK_GDN_KERNEL_TARGETS) + set(MEGA_CHUNK_GDN_KERNEL_TARGETS mega_chunk_gdn_kernel) + ascendc_library(mega_chunk_gdn_kernel STATIC +@@ -139,6 +191,8 @@ target_link_libraries(${OP_PLUGIN_NAME} PRIVATE + ${MEGA_CHUNK_GDN_KERNEL_TARGETS} + workspace_kernel + no_workspace_kernel ++ tp_fusion_kernel ++ tp_fusion_recurrent_kernel + torch_npu + ascendcl + tiling_api +@@ -161,6 +215,8 @@ target_include_directories(${OP_PLUGIN_NAME} PRIVATE + ${PROJECT_OP_SRC_BASE}/causal_conv1d_update/op_host + ${PROJECT_OP_SRC_BASE}/causal_conv1d_update/op_host/stub + ${PROJECT_OP_SRC_BASE}/causal_conv1d/op_host ++ ${PROJECT_OP_SRC_BASE}/fused_qkvzba_conv1d/op_host ++ ${PROJECT_OP_SRC_BASE}/fused_sigmoid_gating_recurrent/op_host + ${TORCH_DIR}/include + ${TORCH_DIR}/include/torch/csrc/api/include + ${TORCH_NPU_DIR}/include diff --git a/csrc/cache_location_assign/op_host/cache_loc_assign.cpp b/csrc/cache_location_assign/op_host/cache_loc_assign.cpp index 5fbecbe..b42f5bc 100644 --- a/csrc/cache_location_assign/op_host/cache_loc_assign.cpp @@ -210,11 +381,1834 @@ index 109f604..c78f303 100644 uint64_t cacheLocSize; int64_t cacheIdxStart{0}; +diff --git a/csrc/causal_conv1d/op_host/causal_conv1d.cpp b/csrc/causal_conv1d/op_host/causal_conv1d.cpp +index cf76fbd..66a44b3 100644 +--- a/csrc/causal_conv1d/op_host/causal_conv1d.cpp ++++ b/csrc/causal_conv1d/op_host/causal_conv1d.cpp +@@ -70,6 +70,7 @@ struct CausalConv1dTilingKey { + int64_t hasCacheIndices; + int64_t hasInitialState; + int64_t hasNumAccept; ++ int64_t xRowStride; // x physical row stride goes into the hash; always = dim for the standalone op + }; + + struct CausalConv1dTilingKeyHash { +@@ -97,6 +98,7 @@ struct CausalConv1dTilingKeyHash { + h = HashCombine(h, static_cast(k.hasCacheIndices)); + h = HashCombine(h, static_cast(k.hasInitialState)); + h = HashCombine(h, static_cast(k.hasNumAccept)); ++ h = HashCombine(h, static_cast(k.xRowStride)); + return h; + } + }; +@@ -176,6 +178,7 @@ void ComputeTilingData(int64_t dim, int64_t cuSeqlen, int64_t seqLen, int64_t ba + std::memset(&td, 0, sizeof(td)); + + td.dim = dim; ++ td.xRowStride = dim; // the standalone op's row stride is always dim (numVHeads stays 0 via memset) + td.cuSeqlen = cuSeqlen; + td.seqLen = seqLen; + td.inputMode = inputMode; +@@ -362,7 +365,8 @@ HOST_API at::Tensor causal_conv1d_impl(const at::Tensor &x, const at::Tensor &we + hasBias ? 1 : 0, + hasCacheIndices ? 1 : 0, + hasInitialState ? 1 : 0, +- hasNumAccept ? 1 : 0}; ++ hasNumAccept ? 1 : 0, ++ dim}; + uint64_t hashValue = CausalConv1dTilingKeyHash{}(key); + + static auto globalTilingBuffer = at::empty({tilingSize * static_cast(MAX_CAPTURE_NUM)}, +diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d.h b/csrc/causal_conv1d/op_kernel/causal_conv1d.h +index 2689ed4..a7d52d8 100644 +--- a/csrc/causal_conv1d/op_kernel/causal_conv1d.h ++++ b/csrc/causal_conv1d/op_kernel/causal_conv1d.h +@@ -345,6 +345,7 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::InitRing(int32_t cacheIdx, bool hasI + const int32_t stateLen = tilingData_->stateLen; + const int32_t width = static_cast(tilingData_->width); + const int32_t ringStart = MAX_WIDTH - width; ++ const int64_t xRowStride = tilingData_->xRowStride; // x physical row stride (elements); always = dim for the standalone op + LocalTensor ring = inBuf.Get(); + + for (int32_t i = 0; i < ringStart; ++i) { +@@ -372,7 +373,7 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::InitRing(int32_t cacheIdx, bool hasI + + if (len > 0) { + const int32_t slot0 = SlotCurr(0); +- const int64_t xOffset = static_cast(start) * dim + channelStart; ++ const int64_t xOffset = static_cast(start) * xRowStride + channelStart; + DataCopy(ring[slot0 * MAX_BLOCK_DIM], xGm[xOffset], baseDim); + SetFlag(inputMte2ToVEvent_[slot0]); + } +@@ -402,6 +403,7 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::RunSeq(int32_t start, int32_t len, i + LocalTensor outT = outBuf.Get(); + const bool hasBias = HasBias(); + const bool hasActivation = HasActivation(); ++ const int64_t xRowStride = tilingData_->xRowStride; // x physical row stride (elements); always = dim for the standalone op + for (int32_t t = 0; t < len; ++t) { + const int32_t slotCurr = SlotCurr(t); + +@@ -409,7 +411,7 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::RunSeq(int32_t start, int32_t len, i + + if (t + 1 < len) { + const int32_t slotNext = SlotPrefetch(t); +- const int64_t xOffsetNext = static_cast(start + t + 1) * dim + channelStart; ++ const int64_t xOffsetNext = static_cast(start + t + 1) * xRowStride + channelStart; + WaitFlag(inputVToMte2Event_); + DataCopy(ring[slotNext * MAX_BLOCK_DIM], xGm[xOffsetNext], baseDim); + SetFlag(inputMte2ToVEvent_[slotNext]); +@@ -642,6 +644,7 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::RunSeqFnRolling(int32_t start, int32 + LocalTensor ring = inBuf.Get(); + LocalTensor outT = outBuf.Get(); + const bool hasActivation = HasActivation(); ++ const int64_t xRowStride = tilingData_->xRowStride; // x physical row stride (elements); always = dim for the standalone op + RestoreFnLocalPartials(baseDim); + + for (int32_t t = 0; t < len; ++t) { +@@ -651,7 +654,7 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::RunSeqFnRolling(int32_t start, int32 + + if (t + 1 < len) { + const int32_t slotNext = SlotPrefetch(t); +- const int64_t xOffsetNext = static_cast(start + t + 1) * dim + channelStart; ++ const int64_t xOffsetNext = static_cast(start + t + 1) * xRowStride + channelStart; + WaitFlag(inputVToMte2Event_); + DataCopy(ring[slotNext * MAX_BLOCK_DIM], xGm[xOffsetNext], baseDim); + SetFlag(inputMte2ToVEvent_[slotNext]); +@@ -735,6 +738,7 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::WriteBackStateSpec(int32_t cacheIdx, + { + const int32_t width = static_cast(tilingData_->width); + const int32_t stateLen = tilingData_->stateLen; ++ const int64_t xRowStride = tilingData_->xRowStride; // x physical row stride (elements); always = dim for the standalone op + if (len <= 0) { + return; + } +@@ -788,7 +792,7 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::WriteBackStateSpec(int32_t cacheIdx, + WaitFlag(stateShiftMte3ToMte2Event_); + } + +- const int64_t xOffset0 = static_cast(start) * dim + channelStart; ++ const int64_t xOffset0 = static_cast(start) * xRowStride + channelStart; + DataCopy(buf0, xGm[xOffset0], baseDim); + SetFlag(specWritebackMte2ToMte3Event_[0]); + +@@ -801,7 +805,7 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::WriteBackStateSpec(int32_t cacheIdx, + WaitFlag(specWritebackMte2ToMte3Event_[curr]); + + if (t + 1 < len) { +- const int64_t xOffsetNext = static_cast(start + t + 1) * dim + channelStart; ++ const int64_t xOffsetNext = static_cast(start + t + 1) * xRowStride + channelStart; + if (t > 0) { + WaitFlag(specWritebackMte3ToMte2Event_[next]); + } +diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d_fn_tasks.h b/csrc/causal_conv1d/op_kernel/causal_conv1d_fn_tasks.h +index d467dc4..fee7984 100644 +--- a/csrc/causal_conv1d/op_kernel/causal_conv1d_fn_tasks.h ++++ b/csrc/causal_conv1d/op_kernel/causal_conv1d_fn_tasks.h +@@ -108,14 +108,15 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::InitRingSeqSplit(int32_t seq, int32_ + bool hasGmHistoryCopy = false; + bool hasVectorInit = false; + const int64_t stateBaseOffset = static_cast(cacheIdx) * stateLen * dim + channelStart; +- int64_t xHistoryOffset = static_cast(historyStartTok) * dim + channelStart; ++ const int64_t xRowStride = tilingData_->xRowStride; // x physical row stride (elements); always = dim for the standalone op ++ int64_t xHistoryOffset = static_cast(historyStartTok) * xRowStride + channelStart; + + for (int32_t i = 0; i < ringStart; ++i) { + Duplicate(ring[i * MAX_BLOCK_DIM], static_cast(0), baseDim); + hasVectorInit = true; + } + +- for (int32_t i = 0, srcTok = historyStartTok; i < historyCount; ++i, ++srcTok, xHistoryOffset += dim) { ++ for (int32_t i = 0, srcTok = historyStartTok; i < historyCount; ++i, ++srcTok, xHistoryOffset += xRowStride) { + LocalTensor histSlot = ring[(ringStart + i) * MAX_BLOCK_DIM]; + if (srcTok >= seqStart) { + DataCopy(histSlot, xGm[xHistoryOffset], baseDim); +@@ -147,7 +148,7 @@ __aicore__ inline void CAUSAL_CONV1D_CLASS::InitRingSeqSplit(int32_t seq, int32_ + + if (tileLen > 0) { + const int32_t slot0 = SlotCurr(0); +- const int64_t xOffset = static_cast(tileStart) * dim + channelStart; ++ const int64_t xOffset = static_cast(tileStart) * xRowStride + channelStart; + DataCopy(ring[slot0 * MAX_BLOCK_DIM], xGm[xOffset], baseDim); + SetFlag(inputMte2ToVEvent_[slot0]); + } +diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d_tiling_data.h b/csrc/causal_conv1d/op_kernel/causal_conv1d_tiling_data.h +index 6360deb..4610e34 100644 +--- a/csrc/causal_conv1d/op_kernel/causal_conv1d_tiling_data.h ++++ b/csrc/causal_conv1d/op_kernel/causal_conv1d_tiling_data.h +@@ -62,5 +62,13 @@ struct CausalConv1dTilingData { + int64_t runModeKey; + int64_t widthKey; + int64_t fnPlanKey; ++ ++ int64_t xRowStride; // x physical row stride (elements); the standalone op's host always fills = dim ++ int64_t numVHeads; // for fused_qkvzba_conv1d's b/a copy; always 0 (memset) in the standalone op, unread ++ // fused_qkvzba_conv1d accepts row-stride view inputs (direct read of packed GEMM non-contiguous output): ++ // zWidth is carried explicitly, not derived as xRowStride - dim; baRowStride replaces hardcoded 2*numVHeads. ++ // Always 0 (memset) in the standalone op, unread. ++ int64_t zWidth; ++ int64_t baRowStride; + }; + #endif // CUSTOM_CAUSAL_CONV1D_TILING_DATA_H_ +diff --git a/csrc/fused_qkvzba_conv1d/op_host/fused_qkvzba_conv1d.cpp b/csrc/fused_qkvzba_conv1d/op_host/fused_qkvzba_conv1d.cpp +new file mode 100644 +index 0000000..828de1f +--- /dev/null ++++ b/csrc/fused_qkvzba_conv1d/op_host/fused_qkvzba_conv1d.cpp +@@ -0,0 +1,408 @@ ++// fused_qkvzba_conv1d (GDN decode: fused split + causal_conv1d) host implementation ++/** ++ * This program is free software, you can redistribute it and/or modify it. ++ * Copyright (c) 2025 Huawei Technologies Co., Ltd. ++ * This file is a part of the CANN Open Software. ++ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). ++ * Please refer to the License for details. You may not use this file except in compliance with the License. ++ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR ++ * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See ++ * LICENSE in the root of the software repository for the full text of the License. ++ */ ++ ++/*! ++ * \file fused_qkvzba_conv1d.cpp ++ * \brief fused_qkvzba_conv1d host-side implementation ++ * ++ * Fused semantics (matching fla/utils.py fused_qkvzba_split_reshape_cat_contiguous + torch.ops.npu.causal_conv1d): ++ * y[B, qkvWidth] = causal_conv1d(qkvz[:, :qkvWidth], ...) -- x rows are read with stride xRowStride = qkvz.stride(0) ++ * z[B, nv, head_v] = qkvz[:, qkvWidth : qkvWidth+nv*head_v] -- copied by the same kernel ++ * b[B, nv] = mixed_ba[:, :nv] ++ * a[B, nv] = mixed_ba[:, nv : 2nv] ++ * The conv part is bit-exact with causal_conv1d run_mode=1 (UPDATE): same tiling struct, same device class, ++ * only dim = qkvWidth and xRowStride = physical row stride of qkvz. Decode empty-call convention: ++ * has_initial_state / num_accepted_tokens are always treated as empty. ++ * ++ * qkvz / mixed_ba accept row-stride views (stride(1)==1, row stride >= logical width), e.g. non-contiguous ++ * slices of packed GEMM output: xRowStride = stride(0) (== size(1) when contiguous); zWidth is carried ++ * explicitly in tiling (it cannot be derived as xRowStride - dim when the row stride includes pack padding); ++ * the ba row stride is carried as baRowStride. Alignment contract: xRowStride must still be a multiple of 16 ++ * (32B alignment of z row starts); the ba copy uses DataCopyPad and requires no alignment. ++ */ ++ ++#include ++#include ++#include ++#include ++#include "acl/acl.h" ++#include ++#include ++#include ++#include "torch_npu/csrc/core/npu/NPUStream.h" ++#include "torch_npu/csrc/core/npu/DeviceUtils.h" ++#include "tiling/platform/platform_ascendc.h" ++#include "stub/aclrtlaunch_fused_qkvzba_conv1d.h" ++#include "defines.h" ++#include "torch_helper.h" ++#include "common.h" ++#include "fused_qkvzba_conv1d.h" ++#include "../../causal_conv1d/op_kernel/causal_conv1d_tiling_data.h" // shared tiling struct (with xRowStride/numVHeads) ++ ++namespace sglang { ++namespace npu_kernel { ++ ++constexpr uint32_t PADDING_BYTE = 32U; ++constexpr int64_t MAX_DIM_TILE = 4096; ++constexpr int32_t MAX_WIDTH = 4; ++constexpr int32_t MIN_WIDTH = 2; ++constexpr uint32_t MAX_CAPTURE_NUM = 1024; ++constexpr int64_t Z_COPY_MAX_BYTES = 65535; // single-block DataCopy blockLen limit (bytes); z row copy is never split ++ ++constexpr uint32_t CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE = 1; ++ ++// Op-local tiling cache (same field-hash + global device buffer scheme as the causal_conv1d host, ++// with xRowStride / numVHeads added so entries never collide with the standalone op or other shapes) ++static uint32_t g_fusedQkvzbaConv1dCaptureNum = 0; ++static std::unordered_map g_fusedQkvzbaConv1dCaptureMap; ++ ++struct FusedQkvzbaConv1dTilingKey { ++ int64_t dim; ++ int64_t cuSeqlen; ++ int64_t seqLen; ++ int64_t batch; ++ int64_t inputMode; ++ int64_t width; ++ int64_t stateLen; ++ int64_t numCacheLines; ++ int64_t activationMode; ++ int64_t padSlotId; ++ int64_t runMode; ++ int64_t hasBias; ++ int64_t hasCacheIndices; ++ int64_t hasInitialState; ++ int64_t hasNumAccept; ++ int64_t xRowStride; ++ int64_t numVHeads; ++ int64_t zWidth; // hashed explicitly (not derived as xRowStride - dim) ++ int64_t baRowStride; // ba physical row stride (== 2*numVHeads when contiguous) ++}; ++ ++struct FusedQkvzbaConv1dTilingKeyHash { ++ static inline std::size_t HashCombine(std::size_t seed, std::size_t value) ++ { ++ seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); ++ return seed; ++ } ++ ++ std::size_t operator()(const FusedQkvzbaConv1dTilingKey &k) const ++ { ++ std::size_t h = 0; ++ h = HashCombine(h, static_cast(k.dim)); ++ h = HashCombine(h, static_cast(k.cuSeqlen)); ++ h = HashCombine(h, static_cast(k.seqLen)); ++ h = HashCombine(h, static_cast(k.batch)); ++ h = HashCombine(h, static_cast(k.inputMode)); ++ h = HashCombine(h, static_cast(k.width)); ++ h = HashCombine(h, static_cast(k.stateLen)); ++ h = HashCombine(h, static_cast(k.numCacheLines)); ++ h = HashCombine(h, static_cast(k.activationMode)); ++ h = HashCombine(h, static_cast(k.padSlotId)); ++ h = HashCombine(h, static_cast(k.runMode)); ++ h = HashCombine(h, static_cast(k.hasBias)); ++ h = HashCombine(h, static_cast(k.hasCacheIndices)); ++ h = HashCombine(h, static_cast(k.hasInitialState)); ++ h = HashCombine(h, static_cast(k.hasNumAccept)); ++ h = HashCombine(h, static_cast(k.xRowStride)); ++ h = HashCombine(h, static_cast(k.numVHeads)); ++ h = HashCombine(h, static_cast(k.zWidth)); ++ h = HashCombine(h, static_cast(k.baRowStride)); ++ return h; ++ } ++}; ++ ++namespace { ++ ++inline int64_t CeilDiv(int64_t x, int64_t y) ++{ ++ return (x + y - 1) / y; ++} ++ ++struct UpdateDimTileChoice { ++ int64_t baseDim = 0; ++ int64_t baseDimCnt = 0; ++ int64_t gridSize = 0; ++}; ++ ++// Mirrors the causal_conv1d host's ChooseUpdateBaseDimChoice (GE update-mode tiling policy); ++// here dim = qkvWidth participates in the choice, identical to the old mixed_qkv path. ++UpdateDimTileChoice ChooseUpdateBaseDimChoice(int64_t batch, int64_t dim, int32_t numCores) ++{ ++ const int64_t candidates[] = {4096, 2048, 1024, 512, 384, 192}; ++ const int64_t coreNum = (numCores > 0) ? static_cast(numCores) : 1; ++ ++ auto chooseOnce = [&](bool requireExactDiv) -> UpdateDimTileChoice { ++ UpdateDimTileChoice bestOver; ++ int64_t bestOverGap = std::numeric_limits::max(); ++ UpdateDimTileChoice bestUnder; ++ ++ for (int64_t candBaseDim : candidates) { ++ if (candBaseDim <= 0) { ++ continue; ++ } ++ if (requireExactDiv && (dim % candBaseDim != 0)) { ++ continue; ++ } ++ const int64_t baseDimCnt = requireExactDiv ? (dim / candBaseDim) : CeilDiv(dim, candBaseDim); ++ const int64_t gridSize = batch * baseDimCnt; ++ if (gridSize <= 0) { ++ continue; ++ } ++ if (gridSize >= coreNum) { ++ const int64_t gap = gridSize - coreNum; ++ if (gap < bestOverGap) { ++ bestOver = {candBaseDim, baseDimCnt, gridSize}; ++ bestOverGap = gap; ++ } ++ } else if (gridSize > bestUnder.gridSize || ++ (gridSize == bestUnder.gridSize && candBaseDim < bestUnder.baseDim)) { ++ bestUnder = {candBaseDim, baseDimCnt, gridSize}; ++ } ++ } ++ return (bestOver.baseDim != 0) ? bestOver : bestUnder; ++ }; ++ ++ UpdateDimTileChoice result = chooseOnce(true); ++ if (result.baseDim == 0) { ++ result = chooseOnce(false); ++ } ++ return result; ++} ++ ++// Mirrors the causal_conv1d host's ComputeTilingData, keeping only the decode (run_mode=1 UPDATE) path, ++// with hasInitialState=false / hasNumAccept=false per the empty-call convention. ++// zWidth / baRowStride go into tiling explicitly (cannot be derived from xRowStride for row-stride views). ++void ComputeTilingData(int64_t dim, int64_t cuSeqlen, int64_t seqLen, int64_t batch, int64_t inputMode, int64_t width, ++ int64_t stateLen, int64_t numCacheLines, int64_t activationMode, int64_t padSlotId, bool hasBias, ++ bool hasCacheIndices, bool isBf16, int32_t numCores, int64_t xRowStride, int64_t numVHeads, ++ int64_t zWidth, int64_t baRowStride, CausalConv1dTilingData &td) ++{ ++ (void)padSlotId; ++ std::memset(&td, 0, sizeof(td)); ++ ++ td.dim = dim; ++ td.xRowStride = xRowStride; // x rows are read with this stride (= qkvz.stride(0)) ++ td.numVHeads = numVHeads; // head count for the b/a copy ++ td.zWidth = zWidth; // z column width carried explicitly (= numVHeads * head_v_dim) ++ td.baRowStride = baRowStride; // mixed_ba physical row stride (== 2*numVHeads when contiguous) ++ td.cuSeqlen = cuSeqlen; ++ td.seqLen = seqLen; ++ td.inputMode = inputMode; ++ td.width = width; ++ td.stateLen = stateLen; ++ td.numCacheLines = numCacheLines; ++ td.batch = batch; ++ td.activationMode = activationMode; ++ td.padSlotId = padSlotId; ++ td.hasBias = hasBias ? 1 : 0; ++ td.hasCacheIndices = hasCacheIndices ? 1 : 0; ++ td.hasInitialStateMode = 0; // decode empty call: no has_initial_state ++ td.hasInitStateWorkspace = 0; ++ td.hasNumAcceptedTokens = 0; // decode empty call: no num_accepted_tokens (non-speculative path) ++ ++ td.dtypeKey = isBf16 ? 0 : 1; ++ td.runModeKey = CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE; ++ td.widthKey = (width == 2) ? 1 : (width == 3) ? 2 : 3; // same WIDTH_2/3/4 encoding as the old host ++ ++ UpdateDimTileChoice choice = ChooseUpdateBaseDimChoice(batch, dim, numCores); ++ if (choice.baseDim <= 0 || choice.baseDimCnt <= 0) { ++ choice.baseDim = (dim > 0 && dim <= MAX_DIM_TILE) ? dim : MAX_DIM_TILE; ++ choice.baseDimCnt = (choice.baseDim > 0) ? CeilDiv(dim, choice.baseDim) : 1; ++ if (choice.baseDimCnt <= 0) { ++ choice.baseDimCnt = 1; ++ } ++ } ++ td.baseDim = choice.baseDim; ++ td.baseDimCnt = choice.baseDimCnt; ++ td.fnPlanKey = 0; // CAUSAL_CONV1D_TPL_FN_PLAN_INVALID ++ td.tokenBlockSize = 0; ++ td.tokenBlockCnt = 0; ++ ++ td.hasExplicitTokenSeqRanges = 0; ++ td.explicitTokenSeqRangeCount = 0; ++} ++ ++} // namespace ++ ++HOST_API std::tuple fused_qkvzba_conv1d_impl( ++ const at::Tensor &qkvz, const at::Tensor &weight, const at::Tensor &conv_states, const at::Tensor &mixed_ba, ++ int64_t num_k_heads, int64_t num_v_heads, int64_t head_k_dim, int64_t head_v_dim, const at::Tensor &bias, ++ const at::Tensor &query_start_loc, const at::Tensor &cache_indices, int64_t activation_mode, int64_t pad_slot_id) ++{ ++ TORCH_CHECK(qkvz.defined(), "qkvz tensor must be defined"); ++ TORCH_CHECK(weight.defined(), "weight tensor must be defined"); ++ TORCH_CHECK(conv_states.defined(), "conv_states tensor must be defined"); ++ TORCH_CHECK(mixed_ba.defined(), "mixed_ba tensor must be defined"); ++ ++ TORCH_CHECK(qkvz.dim() == 2, "qkvz must be 2D tensor"); ++ TORCH_CHECK(weight.dim() == 2, "weight must be 2D tensor"); ++ TORCH_CHECK(conv_states.dim() == 3, "conv_states must be 3D tensor"); ++ TORCH_CHECK(mixed_ba.dim() == 2, "mixed_ba must be 2D tensor"); ++ ++ const at::ScalarType dtype = qkvz.scalar_type(); ++ TORCH_CHECK(dtype == at::kBFloat16 || dtype == at::kHalf, "Only BF16 and FP16 are supported"); ++ TORCH_CHECK(weight.scalar_type() == dtype, "weight dtype must match qkvz dtype"); ++ TORCH_CHECK(conv_states.scalar_type() == dtype, "conv_states dtype must match qkvz dtype"); ++ TORCH_CHECK(mixed_ba.scalar_type() == dtype, "mixed_ba dtype must match qkvz dtype"); ++ ++ // Row-stride views (non-contiguous slices of packed GEMM output) are accepted: contiguous within a row ++ // (stride(1)==1) and row stride >= logical width; contiguous tensors satisfy this trivially. ++ TORCH_CHECK(qkvz.stride(1) == 1 && qkvz.stride(0) >= qkvz.size(1), ++ "qkvz must be row-contiguous (stride(1)==1 and stride(0)>=size(1)); contiguous or pack row-stride view"); ++ TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous"); ++ TORCH_CHECK(conv_states.is_contiguous(), "conv_states must be contiguous"); ++ TORCH_CHECK(mixed_ba.stride(1) == 1 && mixed_ba.stride(0) >= mixed_ba.size(1), ++ "mixed_ba must be row-contiguous (stride(1)==1 and stride(0)>=size(1))"); ++ ++ TORCH_CHECK(num_k_heads > 0 && num_v_heads > 0 && head_k_dim > 0 && head_v_dim > 0, ++ "num_k_heads/num_v_heads/head_k_dim/head_v_dim must be positive"); ++ ++ const int64_t qkvWidth = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim; ++ const int64_t zWidth = num_v_heads * head_v_dim; ++ TORCH_CHECK(qkvz.size(1) == qkvWidth + zWidth, "qkvz.size(1) must equal qkvWidth + zWidth = ", ++ qkvWidth + zWidth); ++ TORCH_CHECK(mixed_ba.size(1) == 2 * num_v_heads, "mixed_ba.size(1) must equal 2 * num_v_heads"); ++ TORCH_CHECK(mixed_ba.size(0) == qkvz.size(0), "mixed_ba and qkvz must have the same row count"); ++ ++ const int64_t dim = qkvWidth; // conv consumes only the first qkvWidth columns of each row (= the old path's split-out mixed_qkv) ++ const int64_t xRowStride = qkvz.stride(0); // physical row stride (== size(1) when contiguous) ++ const int64_t baRowStride = mixed_ba.stride(0); // physical row stride (== 2*numVHeads when contiguous) ++ ++ const int64_t width = weight.size(0); ++ TORCH_CHECK(width >= MIN_WIDTH && width <= MAX_WIDTH, "Only support width in [2,4]"); ++ TORCH_CHECK(weight.size(1) == dim, "weight must be [width, qkvWidth]"); ++ ++ // Alignment contract for the z/b/a copy: z row start byte offset = element_size * (t * xRowStride + qkvWidth) ++ // and z row bytes = element_size * zWidth must both be 32B aligned (the kernel uses plain DataCopy); ++ // for bf16/fp16 this means qkvWidth / xRowStride must be multiples of 16 (the pack side pads the packed N ++ // to a multiple of 16; when contiguous, xRowStride == size(1) == qkvWidth + zWidth). ++ TORCH_CHECK(qkvWidth % 16 == 0 && xRowStride % 16 == 0, ++ "qkvWidth and qkvz.stride(0) must be multiples of 16 (32B alignment for z DataCopy)"); ++ TORCH_CHECK(zWidth * static_cast(qkvz.element_size()) <= Z_COPY_MAX_BYTES, ++ "z row bytes exceed single-block DataCopy limit"); ++ ++ // decode is fixed to 2D varlen semantics (the old host's inputMode=0 path): batch is derived from query_start_loc ++ const int64_t inputMode = 0; ++ const int64_t seqLen = 0; ++ const int64_t cuSeqlen = qkvz.size(0); ++ TORCH_CHECK(query_start_loc.defined() && query_start_loc.numel() >= 2, ++ "query_start_loc must have at least 2 elements"); ++ const int64_t batch = query_start_loc.size(0) - 1; ++ TORCH_CHECK(batch == cuSeqlen, "decode requires qkvz.size(0) == query_start_loc.size(0) - 1"); ++ ++ const int64_t numCacheLines = conv_states.size(0); ++ const int64_t stateLen = conv_states.size(1); ++ TORCH_CHECK(conv_states.size(2) == dim, "conv_states must be [slots, stateLen, qkvWidth]"); ++ ++ bool hasBias = bias.defined() && bias.numel() > 0; ++ bool hasCacheIndices = cache_indices.defined() && cache_indices.numel() > 0; ++ bool isBf16 = (dtype == at::kBFloat16); ++ ++ at::Tensor y = at::empty({batch, qkvWidth}, qkvz.options()); ++ at::Tensor z = at::empty({batch, num_v_heads, head_v_dim}, qkvz.options()); ++ at::Tensor b = at::empty({batch, num_v_heads}, qkvz.options()); ++ at::Tensor a = at::empty({batch, num_v_heads}, qkvz.options()); ++ ++ // None -> empty conversion for optional args, matching the old host (causal_conv1d.cpp) ++ at::Tensor bias_tensor = hasBias ? bias : at::empty({0}, qkvz.options()); ++ at::Tensor query_start_loc_tensor = query_start_loc.to(at::kLong); ++ at::Tensor cache_indices_tensor = ++ hasCacheIndices ? cache_indices.to(at::kLong) : at::empty({0}, qkvz.options().dtype(at::kLong)); ++ at::Tensor has_initial_state_tensor = at::empty({0}, qkvz.options().dtype(at::kLong)); // decode empty call ++ at::Tensor num_accepted_tokens_tensor = at::empty({0}, qkvz.options().dtype(at::kInt)); // decode empty call ++ ++ auto ascendc_platform = platform_ascendc::PlatformAscendCManager::GetInstance(); ++ int32_t maxAivCore = static_cast(ascendc_platform->GetCoreNumAiv()); ++ ++ CausalConv1dTilingData tilingData; ++ ComputeTilingData(dim, cuSeqlen, seqLen, batch, inputMode, width, stateLen, numCacheLines, activation_mode, ++ pad_slot_id, hasBias, hasCacheIndices, isBf16, maxAivCore, xRowStride, num_v_heads, ++ zWidth, baRowStride, tilingData); ++ ++ // run_mode=1 (UPDATE): totalBlocks = batch * baseDimCnt (matching the old host) ++ int64_t totalBlocks = tilingData.batch * tilingData.baseDimCnt; ++ int32_t blockDim = std::min(maxAivCore, static_cast(totalBlocks)); ++ if (blockDim <= 0) { ++ blockDim = 1; ++ } ++ ++ int32_t libApiWorkspaceSize = static_cast(ascendc_platform->GetLibApiWorkSpaceSize()); ++ // decode empty call has no hasInitialState workspace: totalWorkspace = libApiWorkspaceSize (matching the old host) ++ int64_t totalWorkspace = std::max(static_cast(libApiWorkspaceSize), static_cast(0)); ++ if (totalWorkspace <= 0) { ++ totalWorkspace = libApiWorkspaceSize; ++ } ++ ++ int32_t tilingSize = ++ (static_cast(sizeof(CausalConv1dTilingData)) + PADDING_BYTE - 1) / PADDING_BYTE * PADDING_BYTE; ++ ++ FusedQkvzbaConv1dTilingKey key{dim, ++ cuSeqlen, ++ seqLen, ++ batch, ++ inputMode, ++ width, ++ stateLen, ++ numCacheLines, ++ activation_mode, ++ pad_slot_id, ++ static_cast(CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE), ++ hasBias ? 1 : 0, ++ hasCacheIndices ? 1 : 0, ++ 0, // hasInitialState: decode empty call ++ 0, // hasNumAccept: decode empty call ++ xRowStride, ++ num_v_heads, ++ zWidth, ++ baRowStride}; ++ uint64_t hashValue = FusedQkvzbaConv1dTilingKeyHash{}(key); ++ ++ static auto globalTilingBuffer = at::empty({tilingSize * static_cast(MAX_CAPTURE_NUM)}, ++ at::TensorOptions().dtype(at::kByte).device(qkvz.options().device())); ++ ++ auto copyTilingToDevice = [&]() { ++ auto cpuTiling = at::empty({tilingSize}, at::kByte); ++ std::memcpy(cpuTiling.data_ptr(), &tilingData, sizeof(CausalConv1dTilingData)); ++ return TorchNpuHelper::CopyTensorHostToDevice(cpuTiling); ++ }; ++ ++ at::Tensor tilingTensor; ++ if (g_fusedQkvzbaConv1dCaptureMap.find(hashValue) != g_fusedQkvzbaConv1dCaptureMap.end()) { ++ tilingTensor = at::from_blob( ++ globalTilingBuffer.data_ptr() + (tilingSize * g_fusedQkvzbaConv1dCaptureMap[hashValue]), ++ tilingSize, at::kByte); ++ } else if (g_fusedQkvzbaConv1dCaptureNum >= MAX_CAPTURE_NUM) { ++ tilingTensor = copyTilingToDevice(); ++ } else { ++ g_fusedQkvzbaConv1dCaptureMap[hashValue] = g_fusedQkvzbaConv1dCaptureNum; ++ auto deviceTiling = copyTilingToDevice(); ++ globalTilingBuffer ++ .slice(0, g_fusedQkvzbaConv1dCaptureNum * tilingSize, g_fusedQkvzbaConv1dCaptureNum * tilingSize + tilingSize) ++ .copy_(deviceTiling); ++ g_fusedQkvzbaConv1dCaptureNum++; ++ tilingTensor = at::from_blob( ++ globalTilingBuffer.data_ptr() + (tilingSize * g_fusedQkvzbaConv1dCaptureMap[hashValue]), ++ tilingSize, at::kByte); ++ } ++ ++ auto workspaceTensor = ++ at::empty({totalWorkspace}, at::TensorOptions().dtype(at::kByte).device(qkvz.options().device())); ++ ++ EXEC_KERNEL_CMD(fused_qkvzba_conv1d, blockDim, qkvz, weight, conv_states, mixed_ba, bias_tensor, ++ query_start_loc_tensor, cache_indices_tensor, has_initial_state_tensor, ++ num_accepted_tokens_tensor, y, z, b, a, workspaceTensor, tilingTensor); ++ ++ return std::make_tuple(y, z, b, a); ++} ++ ++} // namespace npu_kernel ++} // namespace sglang +diff --git a/csrc/fused_qkvzba_conv1d/op_host/fused_qkvzba_conv1d.h b/csrc/fused_qkvzba_conv1d/op_host/fused_qkvzba_conv1d.h +new file mode 100644 +index 0000000..b72bf7b +--- /dev/null ++++ b/csrc/fused_qkvzba_conv1d/op_host/fused_qkvzba_conv1d.h +@@ -0,0 +1,37 @@ ++// fused_qkvzba_conv1d (GDN decode: fused split + causal_conv1d) host declaration ++/** ++ * This program is free software, you can redistribute it and/or modify it. ++ * Copyright (c) 2025 Huawei Technologies Co., Ltd. ++ * This file is a part of the CANN Open Software. ++ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). ++ * Please refer to the License for details. You may not use this file except in compliance with the License. ++ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR ++ * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See ++ * LICENSE in the root of the software repository for the full text of the License. ++ */ ++ ++/*! ++ * \file fused_qkvzba_conv1d.h ++ * \brief fused_qkvzba_conv1d host-side function declaration ++ */ ++ ++#ifndef CUSTOM_FUSED_QKVZBA_CONV1D_HOST_H_ ++#define CUSTOM_FUSED_QKVZBA_CONV1D_HOST_H_ ++ ++#include ++#include ++#include "defines.h" ++ ++namespace sglang { ++namespace npu_kernel { ++ ++// Returns (y, z, b, a): y=[B, qkvWidth] (conv output), z=[B, num_v_heads, head_v_dim], b/a=[B, num_v_heads] ++HOST_API std::tuple fused_qkvzba_conv1d_impl( ++ const at::Tensor &qkvz, const at::Tensor &weight, const at::Tensor &conv_states, const at::Tensor &mixed_ba, ++ int64_t num_k_heads, int64_t num_v_heads, int64_t head_k_dim, int64_t head_v_dim, const at::Tensor &bias, ++ const at::Tensor &query_start_loc, const at::Tensor &cache_indices, int64_t activation_mode, int64_t pad_slot_id); ++ ++} // namespace npu_kernel ++} // namespace sglang ++ ++#endif // CUSTOM_FUSED_QKVZBA_CONV1D_HOST_H_ +diff --git a/csrc/fused_qkvzba_conv1d/op_host/stub/aclrtlaunch_fused_qkvzba_conv1d.h b/csrc/fused_qkvzba_conv1d/op_host/stub/aclrtlaunch_fused_qkvzba_conv1d.h +new file mode 100644 +index 0000000..e87ddc9 +--- /dev/null ++++ b/csrc/fused_qkvzba_conv1d/op_host/stub/aclrtlaunch_fused_qkvzba_conv1d.h +@@ -0,0 +1,15 @@ ++// aclrtlaunch declaration for fused_qkvzba_conv1d (handwritten per the causal_conv1d stub convention) ++#ifndef HEADER_ACLRTLAUNCH_CUSTOM_FUSED_QKVZBA_CONV1D_H ++#define HEADER_ACLRTLAUNCH_CUSTOM_FUSED_QKVZBA_CONV1D_H ++#include "acl/acl_base.h" ++ ++#ifndef ACLRT_LAUNCH_KERNEL ++#define ACLRT_LAUNCH_KERNEL(kernel_func) aclrtlaunch_##kernel_func ++#endif ++ ++extern "C" uint32_t aclrtlaunch_fused_qkvzba_conv1d(uint32_t numBlocks, aclrtStream stream, void *x, void *weight, ++ void *convStates, void *ba, void *bias, void *queryStartLoc, ++ void *cacheIndices, void *hasInitialState, ++ void *numAcceptedTokens, void *y, void *z, void *b, void *a, ++ void *workspace, void *tiling); ++#endif +diff --git a/csrc/fused_qkvzba_conv1d/op_kernel/fused_qkvzba_conv1d_kernel.cpp b/csrc/fused_qkvzba_conv1d/op_kernel/fused_qkvzba_conv1d_kernel.cpp +new file mode 100644 +index 0000000..3709aaf +--- /dev/null ++++ b/csrc/fused_qkvzba_conv1d/op_kernel/fused_qkvzba_conv1d_kernel.cpp +@@ -0,0 +1,168 @@ ++// fused_qkvzba_conv1d kernel entry (GDN decode: fused split + causal_conv1d) ++/** ++ * This program is free software, you can redistribute it and/or modify it. ++ * Copyright (c) 2025 Huawei Technologies Co., Ltd. ++ * This file is a part of the CANN Open Software. ++ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). ++ * Please refer to the License for details. You may not use this file except in compliance with the License. ++ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR ++ * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See ++ * LICENSE in the root of the software repository for the full text of the License. ++ */ ++ ++/*! ++ * \file fused_qkvzba_conv1d_kernel.cpp ++ * \brief fused_qkvzba_conv1d kernel entry: task 1 = causal_conv1d UPDATE (decode), task 2 = z/b/a split copy ++ * ++ * Bit-exactness rationale: the conv part reuses NsCausalConv1d::RunCausalConv1dUpdate (same device class, ++ * same tiling struct); the only difference is that x GM reads use row stride tilingData->xRowStride ++ * (= qkvz.stride(0), == size(1) when contiguous) instead of dim, and conv consumes only the first ++ * dim = qkvWidth columns of each row, identical in content to the old path's split-out mixed_qkv. ++ * y writeback and conv_states read/write strides remain dim, unchanged. ++ * The z/b/a copy is pure data movement like the old Triton split kernel: a z row = zWidth elements of the ++ * qkvz row starting at offset qkvWidth (zWidth carried explicitly in tiling, ba row stride carried as ++ * baRowStride -- row-stride views of packed GEMM output are supported; padding in the row stride does not ++ * affect z width or b/a content); a b/a row = the first / second numVHeads elements of a mixed_ba row. ++ * The two tasks touch disjoint GM regions, are independent, and need no synchronization. ++ */ ++ ++// Relative include by contract: the csrc tree has two headers named causal_conv1d_update.h ++// (causal_conv1d_update/op_kernel/ belongs to the old standalone op and lacks NsCausalConv1d::RunCausalConv1dUpdate; ++// causal_conv1d/op_kernel/ is the intended one). workspace_kernel's -I order places the old op's directory ++// first, so a bare include would resolve to the wrong header; a relative include resolves from this file's ++// directory and avoids the -I ambiguity. ++#include "../../causal_conv1d/op_kernel/causal_conv1d_update.h" // NsCausalConv1d::RunCausalConv1dUpdate (UPDATE/decode subclass) ++ ++using namespace AscendC; ++using namespace NsCausalConv1d; ++ ++namespace { ++ ++// Task 2: z/b/a split copy. Note: the conv object's TPipe is not reused -- the CausalConv1dUpdate object is ++// locally constructed inside RunCausalConv1dUpdate and destroyed when the call returns (TPipe lifetimes are ++// scope-isolated), so this task builds its own TPipe/TBuf. ++template ++__aicore__ inline void FusedCopyZba(GM_ADDR x, GM_ADDR ba, GM_ADDR z, GM_ADDR b, GM_ADDR a, ++ const __gm__ CausalConv1dTilingData *tilingData) ++{ ++ const int64_t batch = tilingData->batch; ++ const int64_t dim = tilingData->dim; // = qkvWidth, i.e. the column offset of z within a row ++ const int64_t xRowStride = tilingData->xRowStride; // qkvz physical row stride (may include pack padding) ++ const int64_t nv = tilingData->numVHeads; ++ // zWidth is carried explicitly in tiling (= nv * head_v_dim); it cannot be derived as xRowStride - dim ++ // because row-stride view inputs include pack padding in the stride, which would leak into z. ++ // When contiguous, xRowStride == dim + zWidth. ++ const int64_t zWidth = tilingData->zWidth; ++ const int64_t baRowStride = tilingData->baRowStride; // ba physical row stride (== 2*nv when contiguous) ++ if (batch <= 0) { ++ return; ++ } ++ ++ const int64_t blockIdx = static_cast(GetBlockIdx()); ++ const int64_t blockNum = static_cast(GetBlockNum()); ++ ++ // z/b/a share one TPipe and one pair of MTE2/MTE3 events (sequential within the task, no concurrent reuse). ++ TPipe copyPipe; ++ TEventID mte2ToMte3Event = GetTPipePtr()->AllocEventID(); ++ TEventID mte3ToMte2Event = GetTPipePtr()->AllocEventID(); ++ ++ if (zWidth > 0) { ++ GlobalTensor xGm; ++ xGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(x)); ++ GlobalTensor zGm; ++ zGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(z)); ++ ++ // z: zWidth elements per row. The host has verified that qkvWidth / xRowStride are multiples of 16 ++ // (i.e. 32B-aligned row starts and row bytes a multiple of 32 for bf16/fp16) and that a single row ++ // fits the single-block DataCopy limit, so plain DataCopy is used (same as the DataCopy inside conv). ++ TBuf zBuf; ++ copyPipe.InitBuffer(zBuf, static_cast(zWidth * sizeof(T))); ++ LocalTensor zLocal = zBuf.Get(); ++ ++ for (int64_t t = blockIdx; t < batch; t += blockNum) { ++ DataCopy(zLocal, xGm[t * xRowStride + dim], static_cast(zWidth)); ++ SetFlag(mte2ToMte3Event); ++ WaitFlag(mte2ToMte3Event); ++ DataCopy(zGm[t * zWidth], zLocal, static_cast(zWidth)); ++ // wait for this MTE3 write-out before zLocal is reused by the next row ++ // (rows are serialized; at most ceil(batch/blockNum) rows per core, negligible cost) ++ SetFlag(mte3ToMte2Event); ++ WaitFlag(mte3ToMte2Event); ++ } ++ } ++ ++ if (nv > 0) { ++ // b/a: nv elements per row (8B for bf16), below 32B alignment, so plain DataCopy is not usable; ++ // bf16 GM scalar access (GlobalTensor GetValue/SetValue) risks wrong values -- do not use it. ++ // Use the DataCopyPad path: GM->UB with padding, UB->GM exact (the padding is discarded automatically). ++ GlobalTensor baGm; ++ baGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(ba)); ++ GlobalTensor bGm; ++ bGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(b)); ++ GlobalTensor aGm; ++ aGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(a)); ++ ++ const uint32_t rowBytes = static_cast(nv) * static_cast(sizeof(T)); ++ const uint32_t bufBytes = (rowBytes + 31U) / 32U * 32U; ++ TBuf bBuf; ++ TBuf aBuf; ++ copyPipe.InitBuffer(bBuf, bufBytes); ++ copyPipe.InitBuffer(aBuf, bufBytes); ++ LocalTensor bLocal = bBuf.Get(); ++ LocalTensor aLocal = aBuf.Get(); ++ ++ DataCopyExtParams copyParams{1, rowBytes, 0, 0, 0}; ++ DataCopyPadExtParams padParams{false, 0, 0, 0}; ++ ++ for (int64_t t = blockIdx; t < batch; t += blockNum) { ++ const int64_t baRow = t * baRowStride; // ba physical row stride (== 2*nv when contiguous) ++ const int64_t outRow = t * nv; ++ DataCopyPad(bLocal, baGm[baRow], copyParams, padParams); // b = ba[:, :nv] ++ DataCopyPad(aLocal, baGm[baRow + nv], copyParams, padParams); // a = ba[:, nv:2nv] ++ SetFlag(mte2ToMte3Event); ++ WaitFlag(mte2ToMte3Event); ++ DataCopyPad(bGm[outRow], bLocal, copyParams); ++ DataCopyPad(aGm[outRow], aLocal, copyParams); ++ // wait for this MTE3 write-out before bLocal/aLocal are reused by the next row ++ SetFlag(mte3ToMte2Event); ++ WaitFlag(mte3ToMte2Event); ++ } ++ } ++ ++ GetTPipePtr()->ReleaseEventID(mte2ToMte3Event); ++ GetTPipePtr()->ReleaseEventID(mte3ToMte2Event); ++} ++ ++} // namespace ++ ++extern "C" __global__ __aicore__ void fused_qkvzba_conv1d(GM_ADDR x, GM_ADDR weight, GM_ADDR convStates, GM_ADDR ba, ++ GM_ADDR bias, GM_ADDR queryStartLoc, GM_ADDR cacheIndices, ++ GM_ADDR hasInitialState, GM_ADDR numAcceptedTokens, ++ GM_ADDR y, GM_ADDR z, GM_ADDR b, GM_ADDR a, ++ GM_ADDR workspace, GM_ADDR tiling) ++{ ++ REGISTER_TILING_DEFAULT(CausalConv1dTilingData); ++ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0); ++ GM_ADDR userWorkspace = workspace; ++ if (workspace != nullptr) { ++ userWorkspace = AscendC::GetUserWorkspace(workspace); ++ } ++ ++ auto tilingData = reinterpret_cast<__gm__ CausalConv1dTilingData *>(tiling); ++ auto dtypeKey = static_cast(tilingData->dtypeKey); ++ ++ // The host fixes run_mode=1 (UPDATE/decode); this dispatches by dtypeKey exactly like the UPDATE branch ++ // of the causal_conv1d entry. ++ if (dtypeKey == 0) { ++ // Task 1 (conv): identical to the original causal_conv1d UPDATE path (same device class); the object's ++ // scope is this call and its TPipe is destroyed on return, isolated from the copy task's TPipe below. ++ RunCausalConv1dUpdate(x, weight, bias, convStates, queryStartLoc, cacheIndices, hasInitialState, ++ numAcceptedTokens, y, userWorkspace, tilingData); ++ // Task 2 (z/b/a copy): data-independent from the conv task ++ FusedCopyZba(x, ba, z, b, a, tilingData); ++ } else { ++ RunCausalConv1dUpdate(x, weight, bias, convStates, queryStartLoc, cacheIndices, hasInitialState, ++ numAcceptedTokens, y, userWorkspace, tilingData); ++ FusedCopyZba(x, ba, z, b, a, tilingData); ++ } ++} +diff --git a/csrc/fused_sigmoid_gating_recurrent/op_host/fused_sigmoid_gating_recurrent.cpp b/csrc/fused_sigmoid_gating_recurrent/op_host/fused_sigmoid_gating_recurrent.cpp +new file mode 100644 +index 0000000..953dd77 +--- /dev/null ++++ b/csrc/fused_sigmoid_gating_recurrent/op_host/fused_sigmoid_gating_recurrent.cpp +@@ -0,0 +1,167 @@ ++// fused_sigmoid_gating_recurrent (AscendC version of GDN decode recurrent) host implementation ++/** ++ * This program is free software, you can redistribute it and/or modify it. ++ * Copyright (c) 2025 Huawei Technologies Co., Ltd. ++ * This file is a part of the CANN Open Software. ++ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). ++ * You may not use this file except in compliance with the License. ++ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR ++ * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See ++ * LICENSE in the root of the software repository for the full text of the License. ++ */ ++ ++/*! ++ * \file fused_sigmoid_gating_recurrent.cpp ++ * \brief fused_sigmoid_gating_recurrent host-side implementation ++ * ++ * Semantics: drop-in replacement for the production Triton kernel (sgl_kernel_npu ++ * fla fused_sigmoid_gating_recurrent, strided form): sigmoid gating + recurrent ++ * delta rule update, decode only (one token per sequence, T == N, varlen requires ++ * cu_seqlens). No tiling/workspace: EXEC_KERNEL_CMD direct launch with all kernel ++ * args baked by value/pointer; during graph capture the at::empty output lands in ++ * the graph-private pool and is reused on replay (same pattern as the Triton ++ * wrapper's q.new_empty). ++ * ++ * Hard specializations (the kernel is written for K == V == 128; out-of-range ++ * inputs fail here, the wrapper falls back to stock Triton first โ€” this ++ * TORCH_CHECK is a second line of defense): K == V == 128, HV <= 8, HV % H == 0, ++ * N <= 256, q/k/v bf16, a/b bf16, A_log/dt_bias fp32, pool bf16 or fp32. ++ */ ++ ++#include ++#include ++#include "acl/acl.h" ++#include ++#include ++#include ++#include "torch_npu/csrc/core/npu/NPUStream.h" ++#include "torch_npu/csrc/core/npu/DeviceUtils.h" ++#include "tiling/platform/platform_ascendc.h" ++#include "aclrtlaunch_fused_sigmoid_gating_recurrent_bf16.h" ++#include "aclrtlaunch_fused_sigmoid_gating_recurrent_fp32.h" ++#include "defines.h" ++#include "torch_helper.h" ++#include "fused_sigmoid_gating_recurrent.h" ++ ++namespace sglang { ++namespace npu_kernel { ++ ++constexpr int64_t FGR_HEAD_DIM = 128; // kernel-specialized K/V ++constexpr int64_t FGR_MAX_HV = 8; // kernel gating [8] pad cap ++constexpr int64_t FGR_MAX_N = 256; // kernel cu/idx UB residency cap ++ ++HOST_API at::Tensor fused_sigmoid_gating_recurrent_impl( ++ const at::Tensor &A_log, const at::Tensor &a, const at::Tensor &dt_bias, double softplus_beta, ++ double softplus_threshold, const at::Tensor &q, const at::Tensor &k, const at::Tensor &v, const at::Tensor &b, ++ const at::Tensor &initial_state_source, const at::Tensor &initial_state_indices, double scale, ++ const at::Tensor &cu_seqlens, bool use_qk_l2norm, int64_t q_row_stride, int64_t k_row_stride, ++ int64_t v_row_stride) ++{ ++ TORCH_CHECK(q.defined() && k.defined() && v.defined(), "q/k/v must be defined"); ++ TORCH_CHECK(a.defined() && b.defined(), "a/b must be defined"); ++ TORCH_CHECK(A_log.defined() && dt_bias.defined(), "A_log/dt_bias must be defined"); ++ TORCH_CHECK(initial_state_source.defined() && initial_state_indices.defined() && cu_seqlens.defined(), ++ "initial_state_source/initial_state_indices/cu_seqlens must be defined " ++ "(decode-only op: cu_seqlens is required)"); ++ ++ TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, "q/k/v must be 4D [1, T, H, K] / [1, T, HV, V]"); ++ const int64_t T = q.size(1); ++ const int64_t H = q.size(2); ++ const int64_t K = q.size(3); ++ const int64_t HV = v.size(2); ++ const int64_t V = v.size(3); ++ const int64_t N = cu_seqlens.size(0) - 1; ++ ++ TORCH_CHECK(q.size(0) == 1 && v.size(0) == 1, "decode requires flattened batch dim B == 1"); ++ TORCH_CHECK(k.size(1) == T && v.size(1) == T, "q/k/v token dim mismatch"); ++ TORCH_CHECK(k.size(2) == H && k.size(3) == K, "q/k head dims mismatch"); ++ TORCH_CHECK(T == N && N >= 1, ++ "decode-only: total tokens must equal sequence count (one token per sequence), got T=", T, " N=", N); ++ TORCH_CHECK(N <= FGR_MAX_N, "N exceeds kernel UB residency cap (256): ", N); ++ TORCH_CHECK(K == FGR_HEAD_DIM && V == FGR_HEAD_DIM, "kernel specialized for K == V == 128, got K=", K, " V=", V); ++ TORCH_CHECK(H >= 1 && HV >= 1 && HV <= FGR_MAX_HV && HV % H == 0, ++ "require 1 <= H <= HV <= 8 and HV % H == 0, got H=", H, " HV=", HV); ++ ++ TORCH_CHECK(q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 && ++ v.scalar_type() == at::kBFloat16, ++ "q/k/v must be bf16"); ++ TORCH_CHECK(a.scalar_type() == at::kBFloat16 && b.scalar_type() == at::kBFloat16, "a/b must be bf16"); ++ TORCH_CHECK(A_log.scalar_type() == at::kFloat && dt_bias.scalar_type() == at::kFloat, ++ "A_log/dt_bias must be fp32"); ++ const at::ScalarType poolDtype = initial_state_source.scalar_type(); ++ TORCH_CHECK(poolDtype == at::kBFloat16 || poolDtype == at::kFloat, "ssm pool must be bf16 or fp32"); ++ ++ // q/k/v accept strided views; only the last dim must be contiguous and the head dim packed by K/V ++ TORCH_CHECK(q.stride(3) == 1 && k.stride(3) == 1 && v.stride(3) == 1, "q/k/v last dim must be contiguous"); ++ TORCH_CHECK(q.stride(2) == K && k.stride(2) == K, "q/k head dim must be packed (stride(2) == K)"); ++ TORCH_CHECK(v.stride(2) == V, "v head dim must be packed (stride(2) == V)"); ++ TORCH_CHECK(q_row_stride > 0 && k_row_stride > 0 && v_row_stride > 0, "row strides must be positive"); ++ ++ TORCH_CHECK(a.dim() == 2 && a.size(0) == T && a.size(1) == HV && a.is_contiguous(), ++ "a must be contiguous [T, HV]"); ++ TORCH_CHECK(b.dim() == 2 && b.size(0) == T && b.size(1) == HV && b.is_contiguous(), ++ "b must be contiguous [T, HV]"); ++ TORCH_CHECK(A_log.numel() == HV && A_log.is_contiguous(), "A_log must be contiguous [HV]"); ++ TORCH_CHECK(dt_bias.numel() == HV && dt_bias.is_contiguous(), "dt_bias must be contiguous [HV]"); ++ ++ TORCH_CHECK(initial_state_source.dim() == 4 && initial_state_source.size(1) == HV && ++ initial_state_source.size(2) == K && initial_state_source.size(3) == V && ++ initial_state_source.is_contiguous(), ++ "ssm pool must be contiguous [slots, HV, K, V] (copying a non-contiguous pool would silently drop updates; hard assert)"); ++ TORCH_CHECK(initial_state_indices.numel() >= N, "initial_state_indices must have at least N elements"); ++ TORCH_CHECK(cu_seqlens.numel() == N + 1, "cu_seqlens must have N+1 elements"); ++ ++ // conditionally convert int64 indices to int32 (production is int32; skipping the conversion risks a silent copy, keep it explicit) ++ at::Tensor idxI32 = initial_state_indices.scalar_type() == at::kInt ++ ? initial_state_indices ++ : initial_state_indices.to(at::kInt); ++ at::Tensor cuI32 = cu_seqlens.scalar_type() == at::kInt ? cu_seqlens : cu_seqlens.to(at::kInt); ++ TORCH_CHECK(idxI32.is_contiguous() && cuI32.is_contiguous(), "indices/cu_seqlens must be contiguous"); ++ ++ TORCH_CHECK(scale > 0, "scale must be positive"); ++ TORCH_CHECK(softplus_beta != 0, "softplus_beta must be non-zero"); ++ ++ // output same shape/dtype as the Triton wrapper (o = q.new_empty(N,HV,V).view(v.shape): ++ // contiguous [1, T, HV, V], T==N one-to-one) ++ at::Tensor o = at::empty(v.sizes(), q.options()); ++ ++ auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(); ++ const int64_t totalItems = N * HV; ++ uint32_t blockDim = static_cast( ++ std::min(totalItems, static_cast(ascendcPlatform->GetCoreNumAiv()))); ++ if (blockDim == 0) { ++ blockDim = 1; ++ } ++ ++ int devidx = q.device().index(); ++ c10_npu::set_device(devidx); ++ ++ const float scaleF = static_cast(scale); ++ const float spbF = static_cast(softplus_beta); ++ const float invSpbF = 1.0f / spbF; // same fp32 division as the Triton kernel's 1.0/softplus_beta ++ const float thrF = static_cast(softplus_threshold); ++ const uint32_t useL2 = use_qk_l2norm ? 1U : 0U; ++ // EXEC_KERNEL_CMD's ConvertTypes(Ts&...) takes lvalues only; all scalars must land in named ++ // locals first (passing static_cast rvalues fails to compile) ++ const uint32_t nU32 = static_cast(N); ++ const uint32_t hU32 = static_cast(H); ++ const uint32_t hvU32 = static_cast(HV); ++ const uint32_t qRowStrideU32 = static_cast(q_row_stride); ++ const uint32_t kRowStrideU32 = static_cast(k_row_stride); ++ const uint32_t vRowStrideU32 = static_cast(v_row_stride); ++ ++ if (poolDtype == at::kBFloat16) { ++ EXEC_KERNEL_CMD(fused_sigmoid_gating_recurrent_bf16, blockDim, A_log, a, dt_bias, q, k, v, b, o, ++ initial_state_source, idxI32, cuI32, nU32, hU32, hvU32, qRowStrideU32, kRowStrideU32, ++ vRowStrideU32, scaleF, spbF, invSpbF, thrF, useL2); ++ } else { ++ EXEC_KERNEL_CMD(fused_sigmoid_gating_recurrent_fp32, blockDim, A_log, a, dt_bias, q, k, v, b, o, ++ initial_state_source, idxI32, cuI32, nU32, hU32, hvU32, qRowStrideU32, kRowStrideU32, ++ vRowStrideU32, scaleF, spbF, invSpbF, thrF, useL2); ++ } ++ ++ return o; ++} ++ ++} // namespace npu_kernel ++} // namespace sglang +diff --git a/csrc/fused_sigmoid_gating_recurrent/op_host/fused_sigmoid_gating_recurrent.h b/csrc/fused_sigmoid_gating_recurrent/op_host/fused_sigmoid_gating_recurrent.h +new file mode 100644 +index 0000000..914471e +--- /dev/null ++++ b/csrc/fused_sigmoid_gating_recurrent/op_host/fused_sigmoid_gating_recurrent.h +@@ -0,0 +1,45 @@ ++// fused_sigmoid_gating_recurrent (AscendC version of GDN decode recurrent) host declaration ++/** ++ * This program is free software, you can redistribute it and/or modify it. ++ * Copyright (c) 2025 Huawei Technologies Co., Ltd. ++ * This file is a part of the CANN Open Software. ++ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). ++ * You may not use this file except in compliance with the License. ++ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR ++ * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See ++ * LICENSE in the root of the software repository for the full text of the License. ++ */ ++ ++/*! ++ * \file fused_sigmoid_gating_recurrent.h ++ * \brief fused_sigmoid_gating_recurrent host-side function declaration ++ */ ++ ++#ifndef CUSTOM_FUSED_SIGMOID_GATING_RECURRENT_HOST_H_ ++#define CUSTOM_FUSED_SIGMOID_GATING_RECURRENT_HOST_H_ ++ ++#include ++#include "defines.h" ++ ++namespace sglang { ++namespace npu_kernel { ++ ++// GDN decode recurrent (sigmoid gating + delta rule update), AscendC AIV version. ++// Drop-in replacement for the production Triton kernel (sgl_kernel_npu fla ++// fused_sigmoid_gating_recurrent, strided form); decode only (T == N, one token ++// per sequence, varlen requires cu_seqlens). initial_state_source = ssm state pool ++// [slots, HV, K, V], updated in place. q/k/v accept strided views with contiguous ++// last dim (row strides passed explicitly in elements; contiguous case: ++// q_row_stride == H*K, v_row_stride == HV*V). Returns o, same shape as v ++// ([1, T, HV, V]). ++HOST_API at::Tensor fused_sigmoid_gating_recurrent_impl( ++ const at::Tensor &A_log, const at::Tensor &a, const at::Tensor &dt_bias, double softplus_beta, ++ double softplus_threshold, const at::Tensor &q, const at::Tensor &k, const at::Tensor &v, const at::Tensor &b, ++ const at::Tensor &initial_state_source, const at::Tensor &initial_state_indices, double scale, ++ const at::Tensor &cu_seqlens, bool use_qk_l2norm, int64_t q_row_stride, int64_t k_row_stride, ++ int64_t v_row_stride); ++ ++} // namespace npu_kernel ++} // namespace sglang ++ ++#endif // CUSTOM_FUSED_SIGMOID_GATING_RECURRENT_HOST_H_ +diff --git a/csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_bf16_kernel.cpp b/csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_bf16_kernel.cpp +new file mode 100644 +index 0000000..56f4a7e +--- /dev/null ++++ b/csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_bf16_kernel.cpp +@@ -0,0 +1,16 @@ ++// fused_sigmoid_gating_recurrent bf16 entry (bf16 pool, production arm). ++// One kernel per file: KERNEL_TASK_TYPE_DEFAULT does not coexist with multiple ++// kernels in one file (only the first registers; the rest fail with ++// RegisterAscendBinary ret 107000 โ€” see the structural constraint in the lib header). ++ ++#include "fused_sigmoid_gating_recurrent_kernel_lib.h" ++ ++extern "C" __global__ __aicore__ void fused_sigmoid_gating_recurrent_bf16(FGR_KERNEL_ARGS) ++{ ++ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); ++ TPipe pipe; ++ FusedSigmoidGatingRecurrent op(n, h, hv, qRowStride, kRowStride, vRowStride, scale, softplusBeta, ++ invSoftplusBeta, softplusThreshold, useQkL2norm); ++ op.Init(A_log, a, dt_bias, q, k, v, b, o, pool, cache_indices, cu_seqlens, &pipe); ++ op.Process(); ++} +diff --git a/csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_fp32_kernel.cpp b/csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_fp32_kernel.cpp +new file mode 100644 +index 0000000..0657207 +--- /dev/null ++++ b/csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_fp32_kernel.cpp +@@ -0,0 +1,16 @@ ++// fused_sigmoid_gating_recurrent fp32 entry (fp32 pool, accuracy-regression arm). ++// One kernel per file: KERNEL_TASK_TYPE_DEFAULT does not coexist with multiple ++// kernels in one file (only the first registers; the rest fail with ++// RegisterAscendBinary ret 107000 โ€” see the structural constraint in the lib header). ++ ++#include "fused_sigmoid_gating_recurrent_kernel_lib.h" ++ ++extern "C" __global__ __aicore__ void fused_sigmoid_gating_recurrent_fp32(FGR_KERNEL_ARGS) ++{ ++ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); ++ TPipe pipe; ++ FusedSigmoidGatingRecurrent op(n, h, hv, qRowStride, kRowStride, vRowStride, scale, softplusBeta, ++ invSoftplusBeta, softplusThreshold, useQkL2norm); ++ op.Init(A_log, a, dt_bias, q, k, v, b, o, pool, cache_indices, cu_seqlens, &pipe); ++ op.Process(); ++} +diff --git a/csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_kernel_lib.h b/csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_kernel_lib.h +new file mode 100644 +index 0000000..8facc60 +--- /dev/null ++++ b/csrc/fused_sigmoid_gating_recurrent/op_kernel/fused_sigmoid_gating_recurrent_kernel_lib.h +@@ -0,0 +1,720 @@ ++// fused_sigmoid_gating_recurrent ++// AscendC AIV kernel for GDN decode recurrent (sigmoid gating + delta rule update). ++// ++// Semantics match the production Triton kernel (sgl_kernel_npu fla ++// fused_sigmoid_gating_recurrent, strided form) statement by statement. Decode ++// only: exactly 1 token per sequence (host enforces T == N), K == V == 128 (host ++// enforced), HV <= 8 (gating vectors padded to [8]). Pool layout [slot, HV, K, V]; ++// each (slot, hv) state is one contiguous 32KB (bf16) block โ€” single DataCopyPad ++// per head in/out. ++// ++// Precision design (aligned item by item with the Triton reference): ++// 1. softplus threshold branch verbatim: softplus(x) = (beta*x <= thr) ? ++// log(1+exp(beta*x))/beta : x โ€” both branches computed (same semantics as ++// tl.where; the overflowing exp lane is discarded by the scalar select); ++// 2. decay = exp(g), g = -exp(A_log)*softplus, beta = 1/(1+exp(-b)) verbatim ++// (Div chains; no AscendC Sigmoid shortcut); ++// 3. L2norm keeps the division form q/(sqrt(sum(q^2))+1e-6): ReduceSum + scalar ++// sqrt (sqrt before +1e-6, same order as Triton) + full-width Duplicate + ++// vector Div โ€” do NOT switch to vector Rsqrt/Reciprocal (their float ++// precision fails the 2e-4 requirement, per CANN 9.0.0 docs); ++// 4. h += k x v' uses separate Mul+Add (no MulAddDst/FMA contraction, matching ++// Triton's two roundings); ++// 5. both K-dim reductions = 2 K slabs (64 rows each) x 5-pass pairwise tree ++// within a slab + sequential accumulation across slabs; ++// 6. store with CAST_RINT (fp32->bf16 RNE, same as Triton .to(bf16)); ++// bf16->fp32 CAST_NONE is exact; ++// 7. all intermediates fp32 (state cast to fp32 on UB entry, cast back to pool ++// dtype on store). ++// ++// Two points are not a-priori bitwise: (1) K-dim reduction addition order; ++// (2) Exp/Ln polynomial implementations vs triton-ascend lowering (A3 Div in ++// default INTRINSIC mode is at most 1 ulp, same bucket). ++// ++// Parallelism/pipelining: grid = min(N*HV, AIV core count) (host-side ++// GetCoreNumAiv); each core loops over work items idx = blockIdx, ++// blockIdx+blockDim, ...; double-buffered input prefetch (item i+1's MTE2 ++// overlaps item i's V compute). All cross-pipe events explicit: ++// MTE2_V[p] DMA-in(i) ready -> V reads (cast stage) ++// V_MTE2[p] V done reading qkvBuf_[p] (cast stage ends) -> next-next DMA-in may overwrite ++// V_MTE3[p] store-stage cast done -> DMA-out may read stateBuf_/oBuf_ ++// MTE3_V[p] DMA-out(i-2) done -> V may rewrite stateBuf_[p]/oBuf_[p] (start of ++// ComputeItem, before any V write; the fp32 arm rewrites stateBuf_ in ++// place as its work tile, so it must be earliest) ++// MTE3_MTE2[p] DMA-out(i-2) done -> DMA-in(i) may rewrite stateBuf_[p]/oBuf_[p] ++// (V-side writes to stateBuf_ are covered transitively via the V_MTE3 ++// chain; qkvBuf_ is never touched by MTE3, hence the separate V_MTE2; ++// must be a parity pair โ€” a single event would tail with ++// Set(M-2)->Set(M-1) and no Wait in between, i.e. same-ID consecutive ++// SetFlag with >=2 work items per core = documented hang UB, see ++// constraint 3) ++// The first two iterations use iterCnt_ conditions to skip waits whose buffers ++// have no prior accessor (waiting on a never-set flag hangs). KERNEL_TASK_TYPE ++// pinned AIV_ONLY; the kernel uses fully manual sync (explicit ++// PipeBarrier/SetFlag) and the tp_fusion_recurrent_kernel library is built with ++// --cce-auto-sync=off (auto-sync only covers the TQue programming model, unused ++// here; =on would only insert redundant compiler events). ++// ++// Structural constraint 1 (do not merge back): the bf16/fp32 entries must be one ++// kernel per file. KERNEL_TASK_TYPE_DEFAULT does not coexist with multiple ++// kernels in one file โ€” only the first registers; the rest fail with ++// RegisterAscendBinary ret 107000 and the launch hangs on a missing kernel. Hence ++// fused_sigmoid_gating_recurrent_{bf16,fp32}_kernel.cpp, with the class and arg ++// macro centralized here. ++// ++// Structural constraint 2 (do not switch back to Fetch): all events must use ++// AllocEventID (paired with ReleaseEvents at the end). FetchEventID does not ++// occupy an ID (CANN 9.0.0 docs: "this interface does not allocate a TEventID, ++// it only provides an available one"); two calls for the same HardEvent return ++// the SAME TEventID โ€” the parity double-buffer events here (MTE2_V/V_MTE2/V_MTE3/ ++// MTE3_V x2 each) would all collapse onto single hardware flags, and with >=2 ++// work items per core a collapsed ID sees two consecutive SetFlag with no Wait ++// in between = documented hang UB: kernel hangs, host sync never returns. ++// ++// Structural constraint 3 (do not exit the kernel with unconsumed SetFlag): ++// "SetFlag/WaitFlag must appear in pairs" per the official docs; all production ++// kernels in the tree drain at exit (recurrent_gated_delta_rule's ++// SEvent.release() wait loop, causal_conv1d.h's explicit trailing WaitFlag, ++// fused_qkvzba_conv1d's Set immediately followed by Wait). An unconsumed SetFlag ++// leaves the hardware flag set across launches: the next launch on the same core ++// deterministically gets the same IDs from AllocEventID, and the first SetFlag on ++// a leftover channel is "two consecutive SetFlag on the same ID with no Wait" = ++// documented hang UB โ€” shape-independent, guaranteed on the 2nd launch in the ++// process. DrainEvents waits out all leftover flags with iterCnt_ conditions; ++// MTE3_MTE2 is split into a parity pair (strict Set/Wait alternation per ID). ++// ++// UB budget: bf16 ~155KB / fp32 ~160KB (within the 192KB limit). ++ ++#ifndef __FUSED_SIGMOID_GATING_RECURRENT_KERNEL_LIB_H_ ++#define __FUSED_SIGMOID_GATING_RECURRENT_KERNEL_LIB_H_ ++ ++#include "kernel_operator.h" ++ ++// event_t / GetTPipePtr() are global-namespace symbols in CANN 9.0.0 ++// (cce_aicore_intrinsics.h / kernel_tpipe.h; causal_conv1d.h in-tree does the ++// same) โ€” do not add the AscendC:: qualifier. ++ ++using namespace AscendC; ++ ++namespace fused_sigmoid_gating_recurrent { ++ ++constexpr uint32_t FGR_K = 128; // head_k_dim (host TORCH_CHECK enforced) ++constexpr uint32_t FGR_V = 128; // head_v_dim (host TORCH_CHECK enforced) ++constexpr uint32_t FGR_V_SLAB = 64; // V-dim slab (prod buffer UB budget) ++constexpr uint32_t FGR_K_SLAB = 64; // K-dim reduction slab (prod buffer UB budget) ++constexpr uint32_t FGR_MAX_HV = 8; // gating vectors padded to [8] (host enforces HV<=8) ++constexpr uint32_t FGR_STATE = FGR_K * FGR_V; // 16384 state elements per (slot, head) ++constexpr uint32_t FGR_STATE_HALF = FGR_STATE / 2; // 8192 = within the 255-repeat x 64 single-instruction cap ++constexpr uint32_t FGR_MAX_N = 256; // cu_seqlens/cache_indices UB residency cap (production bs<=128) ++constexpr float FGR_L2_EPS = 1e-6f; // verbatim with the Triton kernel's +1e-6 (added after sqrt) ++ ++// element offsets of the qkv parity buffer in UB (bf16: q/k/v 128 elements each, a/b 16-element pad slots each) ++constexpr uint32_t OFF_Q = 0; ++constexpr uint32_t OFF_K = FGR_K; // 128 ++constexpr uint32_t OFF_V = FGR_K + FGR_K; // 256 ++constexpr uint32_t OFF_A = FGR_K + FGR_K + FGR_V; // 384 (bf16 16 elements = 32B slot) ++constexpr uint32_t OFF_B = OFF_A + 16; // 400 ++constexpr uint32_t QKV_PARITY_ELEMS = OFF_B + 16; // 416 elements = 832B ++ ++// gating group [8] slots (element offsets, in float) ++enum GateSlot { ++ GS_DT_BIAS = 0, // dt_bias padded to [8] ++ GS_EXP_ALOG = 8, // exp(A_log) padded to [8] (computed in the prologue) ++ GS_A_F32 = 16, // current item's a row (after cast) ++ GS_B_F32 = 24, // current item's b row (after cast) ++ GS_XA = 32, // x = a + dt_bias ++ GS_BETAX = 40, // beta*x ++ GS_E1 = 48, // exp(beta*x)+1 ++ GS_SPV = 56, // softplus branch value (1/beta)*log(1+exp(beta*x)) ++ GS_BETAV = 64, // sigmoid(b), all lanes ++ GS_G = 72, // [8] slot of scalar g (lane0 valid) ++ GS_DECAY = 80, // exp(g) (lane0 valid) ++ GS_ONES = 88, // all ones (numerator of the beta Div) ++ GS_TMP = 96, // scratch (exp(-b) chain) ++ GATE_ELEMS = 104 // total elements (104 * 4B = 416B, 32B aligned) ++}; ++ ++template // ssm pool dtype: bfloat16_t (production) / float (accuracy-regression arm) ++class FusedSigmoidGatingRecurrent { ++public: ++ __aicore__ inline FusedSigmoidGatingRecurrent(uint32_t n, uint32_t h, uint32_t hv, uint32_t qRowStride, ++ uint32_t kRowStride, uint32_t vRowStride, float scale, ++ float softplusBeta, float invSoftplusBeta, float softplusThreshold, ++ uint32_t useQkL2norm) ++ { ++ n_ = n; ++ h_ = h; ++ hv_ = hv; ++ hvPerH_ = (h > 0) ? (hv / h) : 1; // host enforces hv%h==0 ++ qRowStride_ = qRowStride; ++ kRowStride_ = kRowStride; ++ vRowStride_ = vRowStride; ++ scale_ = scale; ++ spb_ = softplusBeta; ++ invSpb_ = invSoftplusBeta; ++ thr_ = softplusThreshold; ++ useL2_ = useQkL2norm; ++ } ++ ++ __aicore__ inline void Init(GM_ADDR A_log, GM_ADDR a, GM_ADDR dt_bias, GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR b, ++ GM_ADDR o, GM_ADDR pool, GM_ADDR cache_indices, GM_ADDR cu_seqlens, TPipe *pipe) ++ { ++ pipe_ = pipe; ++ blockIdx_ = GetBlockIdx(); ++ blockDim_ = GetBlockNum(); ++ iterCnt_ = 0; ++ ++ aLogGm_.SetGlobalBuffer((__gm__ float *)A_log); ++ aGm_.SetGlobalBuffer((__gm__ bfloat16_t *)a); ++ dtBiasGm_.SetGlobalBuffer((__gm__ float *)dt_bias); ++ qGm_.SetGlobalBuffer((__gm__ bfloat16_t *)q); ++ kGm_.SetGlobalBuffer((__gm__ bfloat16_t *)k); ++ vGm_.SetGlobalBuffer((__gm__ bfloat16_t *)v); ++ bGm_.SetGlobalBuffer((__gm__ bfloat16_t *)b); ++ oGm_.SetGlobalBuffer((__gm__ bfloat16_t *)o); ++ poolGm_.SetGlobalBuffer((__gm__ T *)pool); ++ idxGm_.SetGlobalBuffer((__gm__ int32_t *)cache_indices); ++ cuGm_.SetGlobalBuffer((__gm__ int32_t *)cu_seqlens); ++ ++ // ---- UB allocation (bf16 arm ~155KB / fp32 arm ~160KB, within the 192KB limit) ---- ++ pipe_->InitBuffer(stateBuf_[0], FGR_STATE * sizeof(T)); ++ pipe_->InitBuffer(stateBuf_[1], FGR_STATE * sizeof(T)); ++ if (IsBf16()) { ++ // fp32 work tile (bf16 arm: state is cast here on UB entry; fp32 arm uses stateBuf directly as the work tile) ++ pipe_->InitBuffer(workBuf_, FGR_STATE * sizeof(float)); ++ } ++ pipe_->InitBuffer(prodBuf_, FGR_K_SLAB * FGR_V_SLAB * sizeof(float)); ++ pipe_->InitBuffer(brcbBuf_, FGR_K * 8 * sizeof(float)); ++ pipe_->InitBuffer(qkvBuf_[0], QKV_PARITY_ELEMS * sizeof(bfloat16_t)); ++ pipe_->InitBuffer(qkvBuf_[1], QKV_PARITY_ELEMS * sizeof(bfloat16_t)); ++ pipe_->InitBuffer(oBuf_[0], FGR_V * sizeof(bfloat16_t)); ++ pipe_->InitBuffer(oBuf_[1], FGR_V * sizeof(bfloat16_t)); ++ pipe_->InitBuffer(qF32Buf_, FGR_K * sizeof(float)); ++ pipe_->InitBuffer(kF32Buf_, FGR_K * sizeof(float)); ++ pipe_->InitBuffer(vF32Buf_, FGR_V * sizeof(float)); ++ pipe_->InitBuffer(normTmpBuf_, FGR_K * sizeof(float)); ++ pipe_->InitBuffer(reduceWorkBuf_, FGR_K * sizeof(float)); // ReduceSum workspace allocated separately ++ pipe_->InitBuffer(sumBuf_, 2 * 8 * sizeof(float)); // [0]=q sum of squares, [8]=k sum of squares ++ pipe_->InitBuffer(denomBuf_, FGR_K * sizeof(float)); ++ pipe_->InitBuffer(deltaBuf_, FGR_V * sizeof(float)); ++ pipe_->InitBuffer(oF32Buf_, FGR_V * sizeof(float)); ++ pipe_->InitBuffer(cuBuf_, (FGR_MAX_N + 8) * sizeof(int32_t)); // 32B aligned ++ pipe_->InitBuffer(idxBuf_, FGR_MAX_N * sizeof(int32_t)); ++ pipe_->InitBuffer(gateBuf_, GATE_ELEMS * sizeof(float)); ++ ++ // ---- events (TPipe forbids self-chosen IDs; must AllocEventID โ€” ++ // FetchEventID does not occupy an ID: repeated Fetch on the same ++ // HardEvent returns the same TEventID, so the parity double-buffer ++ // events would collapse onto a single hardware flag; two consecutive ++ // SetFlag on the same ID with no Wait in between hits the ++ // CANN-documented hang/UB) ---- ++ evtVS_ = GetTPipePtr()->AllocEventID(); ++ evtSV_ = GetTPipePtr()->AllocEventID(); ++ evtMte2S_ = GetTPipePtr()->AllocEventID(); ++ evtMte2V_[0] = GetTPipePtr()->AllocEventID(); ++ evtMte2V_[1] = GetTPipePtr()->AllocEventID(); ++ evtVMte2_[0] = GetTPipePtr()->AllocEventID(); ++ evtVMte2_[1] = GetTPipePtr()->AllocEventID(); ++ evtVMte3_[0] = GetTPipePtr()->AllocEventID(); ++ evtVMte3_[1] = GetTPipePtr()->AllocEventID(); ++ evtMte3V_[0] = GetTPipePtr()->AllocEventID(); ++ evtMte3V_[1] = GetTPipePtr()->AllocEventID(); ++ evtMte3Mte2_[0] = GetTPipePtr()->AllocEventID(); ++ evtMte3Mte2_[1] = GetTPipePtr()->AllocEventID(); ++ } ++ ++ __aicore__ inline void Process() ++ { ++ const uint32_t totalItems = n_ * hv_; ++ if (blockIdx_ >= totalItems) { ++ return; // extra cores have no work items (not triggered when blockDim=min(N*HV, cores); defensive) ++ } ++ Prologue(); ++ // software pipeline (iteration handles work item idx, prefetches idx+blockDim; wait ++ // conditions per the event table in the file header, iterCnt_ conditions skip ++ // nonexistent prior accesses in the first two rounds): ++ CopyIn(blockIdx_, 0); ++ uint32_t parity = 0; ++ for (uint32_t idx = blockIdx_; idx < totalItems; idx += blockDim_) { ++ const uint32_t next = idx + blockDim_; ++ const uint32_t p = parity; ++ ComputeItem(idx, p); ++ if (next < totalItems) { ++ CopyIn(next, 1 - p); ++ } ++ CopyOut(idx, p); ++ parity = 1 - parity; ++ iterCnt_++; ++ } ++ DrainEvents(); ++ ReleaseEvents(); ++ } ++ ++ // Exit drain (official "SetFlag/WaitFlag must appear in pairs" constraint, same ++ // discipline as in-tree RGDR SEvent.release() / causal_conv1d trailing WaitFlag): ++ // the SetFlags of the last two iterations have no consumer within this launch and ++ // must each be waited before exit โ€” otherwise set hardware flags persist across ++ // launches, and the next launch's first SetFlag on the same ID is "two consecutive ++ // SetFlag on the same ID with no Wait in between" hang UB (shape-independent, ++ // guaranteed on the 2nd launch in the process). ++ // Conditions are complementary to the skip conditions in CopyIn/ComputeItem: with ++ // iterCnt_ == M, the V_MTE2/MTE3_V/MTE3_MTE2 of parities (M-1)&1 and M&1 are ++ // exactly the set-but-unconsumed flags. ++ __aicore__ inline void DrainEvents() ++ { ++ if (iterCnt_ >= 1) { ++ const uint32_t lastP = (iterCnt_ - 1) & 1; ++ WaitFlag(evtVMte2_[lastP]); // set by ComputeItem(M-1) ++ WaitFlag(evtMte3V_[lastP]); // set by CopyOut(M-1) ++ WaitFlag(evtMte3Mte2_[lastP]); // set by CopyOut(M-1) ++ if (iterCnt_ >= 2) { ++ const uint32_t prevP = iterCnt_ & 1; ++ WaitFlag(evtVMte2_[prevP]); // set by ComputeItem(M-2) ++ WaitFlag(evtMte3V_[prevP]); // set by CopyOut(M-2) ++ WaitFlag(evtMte3Mte2_[prevP]); // set by CopyOut(M-2) ++ } ++ } ++ } ++ ++ // paired with AllocEventID (mirrored order, same as causal_conv1d.h; early-exit cores ++ // have no hardware side effects, TPipe is destroyed on kernel exit, pool bookkeeping ++ // only matters within this launch) ++ __aicore__ inline void ReleaseEvents() ++ { ++ GetTPipePtr()->ReleaseEventID(evtVS_); ++ GetTPipePtr()->ReleaseEventID(evtSV_); ++ GetTPipePtr()->ReleaseEventID(evtMte2S_); ++ GetTPipePtr()->ReleaseEventID(evtMte2V_[0]); ++ GetTPipePtr()->ReleaseEventID(evtMte2V_[1]); ++ GetTPipePtr()->ReleaseEventID(evtVMte2_[0]); ++ GetTPipePtr()->ReleaseEventID(evtVMte2_[1]); ++ GetTPipePtr()->ReleaseEventID(evtVMte3_[0]); ++ GetTPipePtr()->ReleaseEventID(evtVMte3_[1]); ++ GetTPipePtr()->ReleaseEventID(evtMte3V_[0]); ++ GetTPipePtr()->ReleaseEventID(evtMte3V_[1]); ++ GetTPipePtr()->ReleaseEventID(evtMte3Mte2_[0]); ++ GetTPipePtr()->ReleaseEventID(evtMte3Mte2_[1]); ++ } ++ ++private: ++ __aicore__ inline bool IsBf16() const ++ { ++ return IsSameType::value; ++ } ++ ++ __aicore__ inline LocalTensor Gate(uint32_t slot) ++ { ++ return gateBuf_.Get()[slot]; ++ } ++ ++ __aicore__ inline void Prologue() ++ { ++ // cu_seqlens / cache_indices / A_log / dt_bias enter UB once (subsequent GetValue does zero GM reads) ++ LocalTensor cuLocal = cuBuf_.Get(); ++ LocalTensor idxLocal = idxBuf_.Get(); ++ DataCopyPadParams padParams; ++ DataCopyParams cuParams{1, static_cast((n_ + 1) * sizeof(int32_t)), 0, 0}; ++ DataCopyPad(cuLocal, cuGm_, cuParams, padParams); ++ DataCopyParams idxParams{1, static_cast(n_ * sizeof(int32_t)), 0, 0}; ++ DataCopyPad(idxLocal, idxGm_, idxParams, padParams); ++ ++ LocalTensor dtPad = Gate(GS_DT_BIAS); ++ LocalTensor expAlogPad = Gate(GS_EXP_ALOG); ++ DataCopyParams hvParams{1, static_cast(hv_ * sizeof(float)), 0, 0}; ++ DataCopyPad(dtPad, dtBiasGm_, hvParams, padParams); ++ DataCopyPad(expAlogPad, aLogGm_, hvParams, padParams); ++ ++ LocalTensor ones = Gate(GS_ONES); ++ Duplicate(ones, 1.0f, FGR_MAX_HV); // all-ones numerator for the beta Div ++ // gating inputs ready after MTE2->V ++ SetFlag(evtMte2V_[0]); ++ WaitFlag(evtMte2V_[0]); ++ PipeBarrier(); ++ Exp(expAlogPad, expAlogPad, FGR_MAX_HV); // exp(A_log); pad lanes get 1.0, harmless, never read ++ PipeBarrier(); ++ // MTE2->S for cu/idx (the only cross-pipe sync of this kind in the kernel; afterwards UB GetValue reads directly) ++ SetFlag(evtMte2S_); ++ WaitFlag(evtMte2S_); ++ } ++ ++ // DMA-in all inputs of work item idx = i_n * HV + i_hv into the parity buffer ++ __aicore__ inline void CopyIn(uint32_t idx, uint32_t p) ++ { ++ // buffer-reuse waits (MTE2 pipe only; skipped while iterCnt_<1 as the parity ++ // buffers have no prior accessor): prior reader of stateBuf_[p]/oBuf_[p] is ++ // DMA-out (MTE3); prior reader of qkvBuf_[p] is the cast stage (V, only needs ++ // waiting before DMA-in reuse) ++ if (iterCnt_ >= 1) { ++ WaitFlag(evtMte3Mte2_[p]); ++ WaitFlag(evtVMte2_[p]); ++ } ++ const uint32_t iN = idx / hv_; ++ const uint32_t iHv = idx - iN * hv_; ++ LocalTensor cuLocal = cuBuf_.Get(); ++ LocalTensor idxLocal = idxBuf_.Get(); ++ const int32_t bos = cuLocal.GetValue(iN); ++ const int32_t slot = idxLocal.GetValue(iN); ++ const uint32_t iH = iHv / hvPerH_; ++ ++ LocalTensor qkvLocal = qkvBuf_[p].Get(); ++ DataCopyPadParams padParams; ++ // q/k rows (strided views addressed by explicit row stride; contiguous case stride == H*K, identical addressing) ++ DataCopyParams rowQKParams{1, static_cast(FGR_K * sizeof(bfloat16_t)), 0, 0}; ++ DataCopyPad(qkvLocal[OFF_Q], qGm_[static_cast(bos) * qRowStride_ + iH * FGR_K], rowQKParams, ++ padParams); ++ DataCopyPad(qkvLocal[OFF_K], kGm_[static_cast(bos) * kRowStride_ + iH * FGR_K], rowQKParams, ++ padParams); ++ // v row ++ DataCopyParams rowVParams{1, static_cast(FGR_V * sizeof(bfloat16_t)), 0, 0}; ++ DataCopyPad(qkvLocal[OFF_V], vGm_[static_cast(bos) * vRowStride_ + iHv * FGR_V], rowVParams, ++ padParams); ++ // a/b rows ([HV] bf16, 8B/16B small transfers go through the pad path) ++ DataCopyParams rowABParams{1, static_cast(hv_ * sizeof(bfloat16_t)), 0, 0}; ++ DataCopyPad(qkvLocal[OFF_A], aGm_[static_cast(bos) * hv_], rowABParams, padParams); ++ DataCopyPad(qkvLocal[OFF_B], bGm_[static_cast(bos) * hv_], rowABParams, padParams); ++ ++ // state 32KB(bf16)/64KB(fp32) contiguous block; skipped when slot<0 (zero-initialized ++ // in Compute, same semantics as Triton's idx>=0 guard) ++ if (slot >= 0) { ++ LocalTensor stateLocal = stateBuf_[p].Get(); ++ DataCopyParams stateParams{FGR_K, static_cast(FGR_V * sizeof(T)), 0, 0}; ++ DataCopyPad(stateLocal, poolGm_[(static_cast(slot) * hv_ + iHv) * FGR_STATE], stateParams, ++ padParams); ++ } ++ SetFlag(evtMte2V_[p]); ++ } ++ ++ // V-side reduction for L2 normalization (shared by q/k); scalar sqrt+1e-6 and the ++ // division are done in ComputeItem's scalar stage / V stage B (division verbatim ++ // with Triton b_x/(sqrt(sum(x^2))+1e-6)) ++ __aicore__ inline void L2Norm(LocalTensor xF32, uint32_t sumSlot) ++ { ++ LocalTensor normTmp = normTmpBuf_.Get(); ++ LocalTensor sumLocal = sumBuf_.Get(); ++ LocalTensor reduceWork = reduceWorkBuf_.Get(); ++ Mul(normTmp, xF32, xF32, FGR_K); ++ PipeBarrier(); ++ ReduceSum(sumLocal[sumSlot], normTmp, reduceWork, FGR_K); ++ PipeBarrier(); ++ } ++ ++ // V-side chain of the softplus dual branch + beta (all [8] small vectors, one head per lane; pad lanes harmless, never read) ++ __aicore__ inline void GatingVecPart() ++ { ++ LocalTensor aF32 = Gate(GS_A_F32); ++ LocalTensor bF32 = Gate(GS_B_F32); ++ LocalTensor dtPad = Gate(GS_DT_BIAS); ++ LocalTensor xa = Gate(GS_XA); ++ LocalTensor betax = Gate(GS_BETAX); ++ LocalTensor e1 = Gate(GS_E1); ++ LocalTensor spv = Gate(GS_SPV); ++ LocalTensor betav = Gate(GS_BETAV); ++ LocalTensor ones = Gate(GS_ONES); ++ LocalTensor tmp = Gate(GS_TMP); ++ ++ Add(xa, aF32, dtPad, FGR_MAX_HV); // x = a + dt_bias ++ PipeBarrier(); ++ Muls(betax, xa, spb_, FGR_MAX_HV); // beta*x ++ PipeBarrier(); ++ Exp(e1, betax, FGR_MAX_HV); // exp(beta*x) (lanes with bx>thr may overflow; ++ PipeBarrier(); // same two-branch semantics as tl.where, discarded by the scalar select) ++ Adds(e1, e1, 1.0f, FGR_MAX_HV); // 1+exp ++ PipeBarrier(); ++ Ln(spv, e1, FGR_MAX_HV); // log(1+exp) ++ PipeBarrier(); ++ Muls(spv, spv, invSpb_, FGR_MAX_HV); // (1/beta)*log(1+exp) ++ PipeBarrier(); ++ // beta = 1/(1+exp(-b)): verbatim formula (no AscendC Sigmoid shortcut) ++ Muls(tmp, bF32, -1.0f, FGR_MAX_HV); // -b (multiply by -1.0 flips the sign exactly, no rounding) ++ PipeBarrier(); ++ Exp(tmp, tmp, FGR_MAX_HV); ++ PipeBarrier(); ++ Adds(tmp, tmp, 1.0f, FGR_MAX_HV); ++ PipeBarrier(); ++ Div(betav, ones, tmp, FGR_MAX_HV); ++ PipeBarrier(); ++ } ++ ++ __aicore__ inline void ComputeItem(uint32_t idx, uint32_t p) ++ { ++ const uint32_t iN = idx / hv_; ++ const uint32_t iHv = idx - iN * hv_; ++ LocalTensor idxLocal = idxBuf_.Get(); ++ const int32_t slot = idxLocal.GetValue(iN); ++ ++ LocalTensor qkvLocal = qkvBuf_[p].Get(); ++ LocalTensor qF32 = qF32Buf_.Get(); ++ LocalTensor kF32 = kF32Buf_.Get(); ++ LocalTensor vF32 = vF32Buf_.Get(); ++ LocalTensor work = IsBf16() ? workBuf_.Get() : stateBuf_[p].Get(); ++ ++ WaitFlag(evtMte2V_[p]); // current item's q/k/v/a/b/state ready ++ // out-buffer/work-tile rewrite wait: V may write stateBuf_[p]/oBuf_[p] only ++ // after DMA-out(i-2) (MTE3) completes โ€” the fp32 arm uses stateBuf_[p] as its ++ // work tile and rewrites it in place from V stage C, so the wait must precede ++ // any V write (the bf16 arm's first V write is much later in the store stage; ++ // waiting early costs nothing). Skipped while iterCnt_<2 (no prior DMA-out on ++ // this parity). ++ if (iterCnt_ >= 2) { ++ WaitFlag(evtMte3V_[p]); ++ } ++ // ---- V stage A: cast + L2 reduction + gating dual branch ---- ++ Cast(qF32, qkvLocal[OFF_Q], RoundMode::CAST_NONE, FGR_K); ++ Cast(kF32, qkvLocal[OFF_K], RoundMode::CAST_NONE, FGR_K); ++ Cast(vF32, qkvLocal[OFF_V], RoundMode::CAST_NONE, FGR_V); ++ Cast(Gate(GS_A_F32), qkvLocal[OFF_A], RoundMode::CAST_NONE, FGR_MAX_HV); ++ Cast(Gate(GS_B_F32), qkvLocal[OFF_B], RoundMode::CAST_NONE, FGR_MAX_HV); ++ PipeBarrier(); ++ if (IsBf16()) { ++ // state into the work tile (bf16->fp32 exact); zero-init when slot<0 (same semantics as Triton's idx>=0 guard) ++ if (slot >= 0) { ++ Cast(work, stateBuf_[p].Get(), RoundMode::CAST_NONE, FGR_STATE_HALF); ++ Cast(work[FGR_STATE_HALF], stateBuf_[p].Get()[FGR_STATE_HALF], RoundMode::CAST_NONE, ++ FGR_STATE_HALF); ++ } else { ++ Duplicate(work, 0.0f, FGR_STATE_HALF); ++ Duplicate(work[FGR_STATE_HALF], 0.0f, FGR_STATE_HALF); ++ } ++ PipeBarrier(); ++ } else if (slot < 0) { ++ Duplicate(work, 0.0f, FGR_STATE_HALF); ++ Duplicate(work[FGR_STATE_HALF], 0.0f, FGR_STATE_HALF); ++ PipeBarrier(); ++ } ++ // V-side reads of qkvBuf_[p] end here (stateBuf_'s last V-side write is in the ++ // store stage, covered transitively via the V_MTE3->MTE3->MTE3_MTE2 chain, no ++ // separate event needed) ++ SetFlag(evtVMte2_[p]); ++ if (useL2_ != 0) { ++ L2Norm(qF32, 0); // sumBuf_[0] ++ L2Norm(kF32, 8); // sumBuf_[8] ++ } ++ GatingVecPart(); ++ // ---- scalar stage (all scalars within one V_S + one S_V) ---- ++ SetFlag(evtVS_); ++ WaitFlag(evtVS_); ++ LocalTensor sumLocal = sumBuf_.Get(); ++ float denomQ = 1.0f; ++ float denomK = 1.0f; ++ if (useL2_ != 0) { ++ // same order as Triton: sqrt first, then +1e-6 (scalar sqrt; ++ // do NOT switch to vector Rsqrt/Reciprocal โ€” float precision fails 2e-4) ++ denomQ = sqrt(sumLocal.GetValue(0)) + FGR_L2_EPS; ++ denomK = sqrt(sumLocal.GetValue(8)) + FGR_L2_EPS; ++ } ++ const float xHv = Gate(GS_XA).GetValue(iHv); ++ const float bxHv = Gate(GS_BETAX).GetValue(iHv); ++ const float spHv = Gate(GS_SPV).GetValue(iHv); ++ const float betaS = Gate(GS_BETAV).GetValue(iHv); ++ const float expAlogS = Gate(GS_EXP_ALOG).GetValue(iHv); ++ // softplus threshold branch (same per-lane semantics as tl.where; each lane here is exactly one head) ++ const float softplusS = (bxHv <= thr_) ? spHv : xHv; ++ const float gS = (-expAlogS) * softplusS; // g = -exp(A_log)*softplus (sign flip exact) ++ SetFlag(evtSV_); ++ WaitFlag(evtSV_); ++ // ---- V stage B: division normalization + decay exp ---- ++ LocalTensor denom = denomBuf_.Get(); ++ if (useL2_ != 0) { ++ Duplicate(denom, denomQ, FGR_K); ++ PipeBarrier(); ++ Div(qF32, qF32, denom, FGR_K); // true division (A3 Div in default INTRINSIC mode ++ PipeBarrier(); // is at most 1 ulp, counted in the instruction-level difference list) ++ Duplicate(denom, denomK, FGR_K); ++ PipeBarrier(); ++ Div(kF32, kF32, denom, FGR_K); ++ PipeBarrier(); ++ } ++ Muls(qF32, qF32, scale_, FGR_K); // scale after norm (Triton order) ++ PipeBarrier(); ++ Duplicate(Gate(GS_G), gS, FGR_MAX_HV); ++ PipeBarrier(); ++ Exp(Gate(GS_DECAY), Gate(GS_G), FGR_MAX_HV); // decay = exp(g) (lane0 valid) ++ PipeBarrier(); ++ SetFlag(evtVS_); ++ WaitFlag(evtVS_); ++ const float decayS = Gate(GS_DECAY).GetValue(0); ++ SetFlag(evtSV_); // S->V sync so decayS is visible as the ++ WaitFlag(evtSV_); // scalar operand of the Muls in V stage C ++ // ---- V stage C: delta rule main chain ---- ++ // h *= decay (Triton: b_h *= b_decay, before delta) ++ Muls(work, work, decayS, FGR_STATE_HALF); ++ Muls(work[FGR_STATE_HALF], work[FGR_STATE_HALF], decayS, FGR_STATE_HALF); ++ PipeBarrier(); ++ // k broadcast tile (Brcb: each element widened 8x; with src1BlkStride=0 gives full row x k[k]) ++ LocalTensor brcb = brcbBuf_.Get(); ++ Brcb(brcb, kF32, FGR_K / 8, {1, 8}); ++ PipeBarrier(); ++ LocalTensor prod = prodBuf_.Get(); ++ LocalTensor delta = deltaBuf_.Get(); ++ // delta[v] = sum_k h[k,v]*k[k] (2 K slabs: in-slab 5-pass tree leaves 2 rows, merged then accumulated across slabs) ++ for (uint32_t s = 0; s < FGR_V / FGR_V_SLAB; s++) { ++ ReduceKSlab(prod, work, brcb, s, 0); ++ Add(delta[s * FGR_V_SLAB], prod, prod[FGR_V_SLAB], FGR_V_SLAB); // merge slab-0's two rows into the slot ++ PipeBarrier(); ++ ReduceKSlab(prod, work, brcb, s, 1); ++ Add(prod, prod, prod[FGR_V_SLAB], FGR_V_SLAB); // merge slab-1's two rows into prod row 0 ++ PipeBarrier(); ++ Add(delta[s * FGR_V_SLAB], delta[s * FGR_V_SLAB], prod, FGR_V_SLAB); // cross-slab accumulation ++ PipeBarrier(); ++ } ++ // v' = (v - delta) * beta (Sub then Muls, two roundings, same as Triton) ++ Sub(vF32, vF32, delta, FGR_V); ++ PipeBarrier(); ++ Muls(vF32, vF32, betaS, FGR_V); ++ PipeBarrier(); ++ // h += k x v' (separate Mul+Add, two roundings; no MulAddDst/FMA contraction) ++ for (uint32_t s = 0; s < FGR_V / FGR_V_SLAB; s++) { ++ for (uint32_t ks = 0; ks < FGR_K / FGR_K_SLAB; ks++) { ++ // prod[r, :] = v'[s slot] * k[ks*64+r] (src0RepStride=0: the whole slab reuses the same v' slot) ++ Mul(prod, vF32[s * FGR_V_SLAB], brcb[ks * FGR_K_SLAB * 8], FGR_V_SLAB, ++ static_cast(FGR_K_SLAB), {1, 1, 0, 8, 0, 1}); ++ PipeBarrier(); ++ const uint32_t wOff = ks * FGR_K_SLAB * FGR_V + s * FGR_V_SLAB; ++ Add(work[wOff], work[wOff], prod, FGR_V_SLAB, static_cast(FGR_K_SLAB), ++ {1, 1, 1, 16, 16, 8}); ++ PipeBarrier(); ++ } ++ } ++ // q broadcast tile (overwrites the k broadcast tile: k unused from here on) ++ Brcb(brcb, qF32, FGR_K / 8, {1, 8}); ++ PipeBarrier(); ++ // o[v] = sum_k h_new[k,v]*q[k] (same shape as delta; uses the rank-1-updated h) ++ LocalTensor oF32 = oF32Buf_.Get(); ++ for (uint32_t s = 0; s < FGR_V / FGR_V_SLAB; s++) { ++ ReduceKSlab(prod, work, brcb, s, 0); ++ Add(oF32[s * FGR_V_SLAB], prod, prod[FGR_V_SLAB], FGR_V_SLAB); ++ PipeBarrier(); ++ ReduceKSlab(prod, work, brcb, s, 1); ++ Add(prod, prod, prod[FGR_V_SLAB], FGR_V_SLAB); ++ PipeBarrier(); ++ Add(oF32[s * FGR_V_SLAB], oF32[s * FGR_V_SLAB], prod, FGR_V_SLAB); ++ PipeBarrier(); ++ } ++ // ---- store stage ---- ++ LocalTensor oLocal = oBuf_[p].Get(); ++ Cast(oLocal, oF32, RoundMode::CAST_RINT, FGR_V); // RNE, same as Triton .to(bf16) ++ if (IsBf16()) { ++ Cast(stateBuf_[p].Get(), work, RoundMode::CAST_RINT, FGR_STATE_HALF); ++ Cast(stateBuf_[p].Get()[FGR_STATE_HALF], work[FGR_STATE_HALF], RoundMode::CAST_RINT, FGR_STATE_HALF); ++ } ++ PipeBarrier(); ++ SetFlag(evtVMte3_[p]); ++ } ++ ++ // prod[r, :] = work[(ks*64+r)*128 + s slot] * brcbVec[ks*64+r], then a 5-pass ++ // pairwise tree reduction (in place, read addresses never below write addresses, ++ // no self-overwrite) leaving 2 rows at prod[0..63] and prod[64..127] (the caller ++ // merges/places them with one Add). ++ __aicore__ inline void ReduceKSlab(LocalTensor prod, LocalTensor work, LocalTensor brcb, ++ uint32_t vSlab, uint32_t kSlab) ++ { ++ const uint32_t wOff = kSlab * FGR_K_SLAB * FGR_V + vSlab * FGR_V_SLAB; ++ // src1BlkStride=0: all 8 blocks within a repeat read the same 8-element brcb ++ // block (= scalar x8); src1RepStride=1: advance block by block across repeats ++ // (k/q row by row) ++ Mul(prod, work[wOff], brcb[kSlab * FGR_K_SLAB * 8], FGR_V_SLAB, static_cast(FGR_K_SLAB), ++ {1, 1, 0, 8, 16, 1}); ++ PipeBarrier(); ++ // 5-pass tree: 64->32->16->8->4->2 rows (the final two-row merge is done by the caller per slot placement) ++ for (uint32_t repeats = FGR_K_SLAB / 2; repeats >= 2; repeats = repeats / 2) { ++ Add(prod, prod, prod[FGR_V_SLAB], FGR_V_SLAB, static_cast(repeats), {1, 1, 1, 8, 16, 16}); ++ PipeBarrier(); ++ } ++ } ++ ++ __aicore__ inline void CopyOut(uint32_t idx, uint32_t p) ++ { ++ const uint32_t iN = idx / hv_; ++ const uint32_t iHv = idx - iN * hv_; ++ LocalTensor cuLocal = cuBuf_.Get(); ++ LocalTensor idxLocal = idxBuf_.Get(); ++ const int32_t bos = cuLocal.GetValue(iN); ++ const int32_t slot = idxLocal.GetValue(iN); ++ ++ WaitFlag(evtVMte3_[p]); // store-stage cast done ++ // o row store (same addressing as Triton: (bos*HV+i_hv)*V) ++ DataCopyParams oParams{1, static_cast(FGR_V * sizeof(bfloat16_t)), 0, 0}; ++ DataCopyPad(oGm_[(static_cast(bos) * hv_ + iHv) * FGR_V], oBuf_[p].Get(), oParams); ++ // state writeback (not written when slot<0, same guard as Triton) ++ if (slot >= 0) { ++ DataCopyParams stateParams{FGR_K, static_cast(FGR_V * sizeof(T)), 0, 0}; ++ DataCopyPad(poolGm_[(static_cast(slot) * hv_ + iHv) * FGR_STATE], stateBuf_[p].Get(), ++ stateParams); ++ } ++ // two sets: next-round rewrite gates for stateBuf_/oBuf_ (MTE3_V for the V side, ++ // MTE3_MTE2 for the MTE2 side; both per parity โ€” a single MTE3_MTE2 event would ++ // tail with two Sets and no Wait in between = hang UB) ++ SetFlag(evtMte3V_[p]); ++ SetFlag(evtMte3Mte2_[p]); ++ } ++ ++private: ++ GlobalTensor aLogGm_; ++ GlobalTensor aGm_; ++ GlobalTensor dtBiasGm_; ++ GlobalTensor qGm_; ++ GlobalTensor kGm_; ++ GlobalTensor vGm_; ++ GlobalTensor bGm_; ++ GlobalTensor oGm_; ++ GlobalTensor poolGm_; ++ GlobalTensor idxGm_; ++ GlobalTensor cuGm_; ++ ++ TPipe *pipe_; ++ TBuf stateBuf_[2]; ++ TBuf workBuf_; ++ TBuf prodBuf_; ++ TBuf brcbBuf_; ++ TBuf qkvBuf_[2]; ++ TBuf oBuf_[2]; ++ TBuf qF32Buf_; ++ TBuf kF32Buf_; ++ TBuf vF32Buf_; ++ TBuf normTmpBuf_; ++ TBuf reduceWorkBuf_; ++ TBuf sumBuf_; ++ TBuf denomBuf_; ++ TBuf deltaBuf_; ++ TBuf oF32Buf_; ++ TBuf cuBuf_; ++ TBuf idxBuf_; ++ TBuf gateBuf_; ++ ++ TEventID evtVS_; ++ TEventID evtSV_; ++ TEventID evtMte2S_; ++ TEventID evtMte2V_[2]; ++ TEventID evtVMte2_[2]; ++ TEventID evtVMte3_[2]; ++ TEventID evtMte3V_[2]; ++ TEventID evtMte3Mte2_[2]; ++ ++ uint32_t n_; ++ uint32_t h_; ++ uint32_t hv_; ++ uint32_t hvPerH_; ++ uint32_t qRowStride_; ++ uint32_t kRowStride_; ++ uint32_t vRowStride_; ++ float scale_; ++ float spb_; ++ float invSpb_; ++ float thr_; ++ uint32_t useL2_; ++ uint32_t blockIdx_; ++ uint32_t blockDim_; ++ uint32_t iterCnt_; ++}; ++ ++} // namespace fused_sigmoid_gating_recurrent ++ ++using fused_sigmoid_gating_recurrent::FusedSigmoidGatingRecurrent; ++ ++#define FGR_KERNEL_ARGS \ ++ GM_ADDR A_log, GM_ADDR a, GM_ADDR dt_bias, GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR b, GM_ADDR o, GM_ADDR pool, \ ++ GM_ADDR cache_indices, GM_ADDR cu_seqlens, uint32_t n, uint32_t h, uint32_t hv, uint32_t qRowStride, \ ++ uint32_t kRowStride, uint32_t vRowStride, float scale, float softplusBeta, float invSoftplusBeta, \ ++ float softplusThreshold, uint32_t useQkL2norm ++ ++#endif // __FUSED_SIGMOID_GATING_RECURRENT_KERNEL_LIB_H_ diff --git a/csrc/pytorch_extensions.cpp b/csrc/pytorch_extensions.cpp -index 1cfccb0..8e4ad91 100644 +index 1cfccb0..fb2d90a 100644 --- a/csrc/pytorch_extensions.cpp +++ b/csrc/pytorch_extensions.cpp -@@ -32,11 +32,11 @@ TORCH_LIBRARY_FRAGMENT(npu, m) +@@ -17,6 +17,8 @@ + #include "sgl_kenel_npu_ops.h" + #include "causal_conv1d_update/op_host/causal_conv1d_update.h" + #include "causal_conv1d/op_host/causal_conv1d.h" ++#include "fused_qkvzba_conv1d/op_host/fused_qkvzba_conv1d.h" ++#include "fused_sigmoid_gating_recurrent/op_host/fused_sigmoid_gating_recurrent.h" + + namespace { + TORCH_LIBRARY_FRAGMENT(npu, m) +@@ -32,11 +34,11 @@ TORCH_LIBRARY_FRAGMENT(npu, m) m.def( "cache_loc_assign(Tensor req_indices, Tensor token_pool, Tensor start_offset, Tensor end_offset, Tensor " @@ -228,6 +2222,59 @@ index 1cfccb0..8e4ad91 100644 m.def( "assign_cache_op(Tensor! out, Tensor src, Tensor dst_start_idx, Tensor dst_end_idx, Tensor src_start_idx, " +@@ -147,6 +149,26 @@ TORCH_LIBRARY_FRAGMENT(npu, m) + "Tensor? query_start_loc=None, Tensor? cache_indices=None, Tensor? has_initial_state=None, " + "Tensor? num_accepted_tokens=None, int activation_mode=0, int pad_slot_id=-1, " + "int run_mode=0) -> Tensor"); ++ ++ // GDN decode split + causal_conv1d in a single kernel ++ m.def( ++ "fused_qkvzba_conv1d(Tensor qkvz, Tensor weight, Tensor conv_states, Tensor mixed_ba, " ++ "int num_k_heads, int num_v_heads, int head_k_dim, int head_v_dim, Tensor? bias=None, " ++ "Tensor? query_start_loc=None, Tensor? cache_indices=None, int activation_mode=0, " ++ "int pad_slot_id=-1) -> (Tensor, Tensor, Tensor, Tensor)"); ++ ++ // AscendC AIV version of the GDN decode recurrent (sigmoid gating + delta ++ // rule update), decode only (T==N, one token per sequence); ++ // initial_state_source is the ssm state pool and is updated in place; ++ // q/k/v may be strided views with a contiguous last dim (row strides are ++ // passed explicitly). ++ m.def( ++ "fused_sigmoid_gating_recurrent(Tensor A_log, Tensor a, Tensor dt_bias, " ++ "float softplus_beta, float softplus_threshold, " ++ "Tensor q, Tensor k, Tensor v, Tensor b, " ++ "Tensor(a!) initial_state_source, Tensor initial_state_indices, " ++ "float scale, Tensor cu_seqlens, bool use_qk_l2norm, " ++ "int q_row_stride, int k_row_stride, int v_row_stride) -> Tensor"); + } + } // namespace + +@@ -246,5 +268,25 @@ TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) + x, weight, bias_or_empty, conv_states, query_start_loc_or_empty, cache_indices_or_empty, + has_initial_state_or_empty, num_accepted_tokens_or_empty, activation_mode, pad_slot_id, run_mode); + }); ++ ++ m.impl("fused_qkvzba_conv1d", ++ [](const at::Tensor &qkvz, const at::Tensor &weight, const at::Tensor &conv_states, ++ const at::Tensor &mixed_ba, int64_t num_k_heads, int64_t num_v_heads, int64_t head_k_dim, ++ int64_t head_v_dim, const c10::optional &bias, ++ const c10::optional &query_start_loc, const c10::optional &cache_indices, ++ int64_t activation_mode, int64_t pad_slot_id) { ++ // Handle optional parameters - convert None to empty tensors ++ auto bias_or_empty = bias.has_value() ? *bias : at::empty({0}, qkvz.options()); ++ auto query_start_loc_or_empty = ++ query_start_loc.has_value() ? *query_start_loc : at::empty({0}, qkvz.options().dtype(at::kLong)); ++ auto cache_indices_or_empty = ++ cache_indices.has_value() ? *cache_indices : at::empty({0}, qkvz.options().dtype(at::kLong)); ++ ++ return sglang::npu_kernel::fused_qkvzba_conv1d_impl( ++ qkvz, weight, conv_states, mixed_ba, num_k_heads, num_v_heads, head_k_dim, head_v_dim, ++ bias_or_empty, query_start_loc_or_empty, cache_indices_or_empty, activation_mode, pad_slot_id); ++ }); ++ ++ m.impl("fused_sigmoid_gating_recurrent", TORCH_FN(sglang::npu_kernel::fused_sigmoid_gating_recurrent_impl)); + } + } // namespace diff --git a/include/sgl_kenel_npu_ops.h b/include/sgl_kenel_npu_ops.h index c6a4fba..ad51c12 100644 --- a/include/sgl_kenel_npu_ops.h @@ -250,6 +2297,866 @@ index c6a4fba..ad51c12 100644 bool assign_cache_op(at::Tensor &dst_tensor, const at::Tensor &src_tensor, const at::Tensor &dst_start_idx, +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/fla/fused_sigmoid_gating_recurrent.py b/python/sgl_kernel_npu/sgl_kernel_npu/fla/fused_sigmoid_gating_recurrent.py +index 2c6288e..770bd67 100644 +--- a/python/sgl_kernel_npu/sgl_kernel_npu/fla/fused_sigmoid_gating_recurrent.py ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/fla/fused_sigmoid_gating_recurrent.py +@@ -5,7 +5,6 @@ import torch + import torch.nn.functional as F + import triton + import triton.language as tl +-from sgl_kernel_npu.fla.utils import input_guard + + + @triton.heuristics( +@@ -29,6 +28,9 @@ def fused_sigmoid_gating_delta_rule_update_npu_kernel( + h0_source, + h0_indices, + cu_seqlens, ++ q_row_stride, ++ k_row_stride, ++ v_row_stride, + scale, + T, + B: tl.constexpr, +@@ -70,9 +72,9 @@ def fused_sigmoid_gating_delta_rule_update_npu_kernel( + for i_bhv in range(BHV): + i_hv = i_nhv * BHV + i_bhv + i_h = i_hv // (HV // H) +- p_q = q + (bos * H + i_h) * K + o_k # q + 0 ~ 127 +- p_k = k + (bos * H + i_h) * K + o_k # k + 0 ~ 127 +- p_v = v + (bos * HV + i_hv) * V + o_v # v + i_nh * 128 + 0 ~ 15 * 8 + 0 ~ 7 ++ p_q = q + bos * q_row_stride + i_h * K + o_k # q + 0 ~ 127 ++ p_k = k + bos * k_row_stride + i_h * K + o_k # k + 0 ~ 127 ++ p_v = v + bos * v_row_stride + i_hv * V + o_v # v + i_nh * 128 + 0 ~ 15 * 8 + 0 ~ 7 + p_b = b + bos * HV + i_hv + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + +@@ -88,9 +90,9 @@ def fused_sigmoid_gating_delta_rule_update_npu_kernel( + + for i in range(T): + # Load inputs +- b_q = tl.load(p_q + i * H * K, mask=mask_k).to(tl.float32) # 128 * float32 +- b_k = tl.load(p_k + i * H * K, mask=mask_k).to(tl.float32) # 128 * float32 +- b_v = tl.load(p_v + i * HV * V, mask=mask_v).to(tl.float32) # 64 * float32 ++ b_q = tl.load(p_q + i * q_row_stride, mask=mask_k).to(tl.float32) # 128 * float32 ++ b_k = tl.load(p_k + i * k_row_stride, mask=mask_k).to(tl.float32) # 128 * float32 ++ b_v = tl.load(p_v + i * v_row_stride, mask=mask_v).to(tl.float32) # 64 * float32 + b_b = tl.load(p_b + i * HV).to(tl.float32) + + # Compute sigmoid gating +@@ -166,7 +168,6 @@ def fused_sigmoid_gating_delta_rule_update_npu_kernel( + tl.store(p_o + i * HV * V, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + +-@input_guard + def fused_sigmoid_gating_delta_rule_update_npu( + A_log: torch.Tensor, + a: torch.Tensor, +@@ -188,7 +189,23 @@ def fused_sigmoid_gating_delta_rule_update_npu( + Fused triton implementation of sigmoid gating delta rule update. + This function uses a single fused kernel that combines both sigmoid gating computation + and the recurrent delta rule update for better performance. ++ ++ v1 ๅ˜ๆ›ด๏ผˆGDN decode ไผ˜ๅŒ–ๆ–นๆกˆB๏ผ‰๏ผšๅŽปๆŽ‰ @input_guard ็š„ๅ…จ้‡ contiguousโ€”โ€” ++ q/k/v ๅ…่ฎธ้ž่ฟž็ปญ่ง†ๅ›พ๏ผˆไป…่ฆๆฑ‚ๆœซ็ปด stride==1๏ผ‰๏ผŒkernel ๆŒ‰ไผ ๅ…ฅ็š„ token ่กŒ stride ++ ๅฏปๅ€๏ผŒๆถˆๅŽป decode ่ทฏๅพ„ๆฏๅฑ‚ 3 ไปฝ q/k/v ๆ‹ท่ด๏ผ›่ฟž็ปญ่พ“ๅ…ฅๆ—ถๅฏปๅ€ไธŽๅŽŸ็‰ˆๅฎŒๅ…จไธ€่‡ด ++ ๏ผˆ้€ bit ็ญ‰ไปท๏ผŒๅทฒ็”ฑ resource/code/gdn_mamba/test_b ้ชŒ่ฏ๏ผ‰ใ€‚ + """ ++ assert q.stride(-1) == 1 and k.stride(-1) == 1 and v.stride(-1) == 1 ++ assert initial_state_source.is_contiguous() ++ a = a.contiguous() ++ b = b.contiguous() ++ A_log = A_log.contiguous() ++ dt_bias = dt_bias.contiguous() ++ if initial_state_indices is not None: ++ initial_state_indices = initial_state_indices.contiguous() ++ if cu_seqlens is not None: ++ cu_seqlens = cu_seqlens.contiguous() ++ + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 +@@ -222,6 +239,9 @@ def fused_sigmoid_gating_delta_rule_update_npu( + h0_source=initial_state_source, + h0_indices=initial_state_indices, + cu_seqlens=cu_seqlens, ++ q_row_stride=q.stride(1), ++ k_row_stride=k.stride(1), ++ v_row_stride=v.stride(1), + scale=scale, + T=T, + B=B, +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/fla/fused_sigmoid_gating_recurrent_decode_optimized.py b/python/sgl_kernel_npu/sgl_kernel_npu/fla/fused_sigmoid_gating_recurrent_decode_optimized.py +new file mode 100644 +index 0000000..a9057c3 +--- /dev/null ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/fla/fused_sigmoid_gating_recurrent_decode_optimized.py +@@ -0,0 +1,328 @@ ++# Decode-optimized GDN recurrent state update for Ascend NPU. ++# ++# Based on sgl-kernel-npu PR #740's ++# ``fla/fused_sigmoid_gating_recurrent_decode_optimized.py``, with strided ++# q/k/v support added (``q/k/v_row_stride`` addressing + contiguous-free ++# wrapper). The PR original addresses the packed layout ``(bos*H+i_h)*K`` and ++# its wrapper calls ``.contiguous()`` on non-contiguous inputs โ€” with fused ++# split producing strided views, that costs 3 extra copies per layer. ++# ++# Optimizations (from PR #740): 1-D grid capped at the AIV vector core count, ++# per-program loop over (sequence, value-head) tiles, BV=64 blocking over V ++# with gating computed once per (seq, v-head), num_warps=4, and a lightweight ++# contiguity check instead of the generic input_guard. ++# ++# Limitations: decode only (T=1 per sequence; the wrapper asserts ++# total_tokens == N). NK>1 is unsupported. The target_verify extras ++# (disable_state_update etc.) are unsupported; do not use for verify. ++ ++from typing import Optional ++ ++import torch ++import triton ++import triton.language as tl ++ ++from sgl_kernel_npu.utils.triton_utils import get_device_properties ++ ++ ++def _maybe_contiguous(x): ++ """Return ``x`` unchanged when already contiguous, else a contiguous copy.""" ++ return x if x.is_contiguous() else x.contiguous() ++ ++ ++@triton.heuristics( ++ { ++ "USE_INITIAL_STATE": lambda args: args["h0_source"] is not None, ++ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, ++ } ++) ++@triton.jit(do_not_specialize=["T", "N", "NHV"]) ++def _fused_sigmoid_gating_delta_rule_update_decode_kernel( ++ A_log, ++ a, ++ dt_bias, ++ softplus_beta, ++ softplus_threshold, ++ q, ++ k, ++ v, ++ b, ++ o, ++ h0_source, ++ h0_indices, ++ cu_seqlens, ++ q_row_stride, ++ k_row_stride, ++ v_row_stride, ++ scale, ++ T, ++ N, ++ NHV, ++ B: tl.constexpr, ++ H: tl.constexpr, ++ HV: tl.constexpr, ++ K: tl.constexpr, ++ V: tl.constexpr, ++ BK: tl.constexpr, ++ BV: tl.constexpr, ++ BHV: tl.constexpr, ++ OVERSUB: tl.constexpr, ++ USE_INITIAL_STATE: tl.constexpr, ++ IS_VARLEN: tl.constexpr, ++ USE_QK_L2NORM_IN_KERNEL: tl.constexpr, ++): ++ """ ++ Decode-optimized fused gating + recurrent delta rule update. ++ ++ Grid: (num_programs,) where num_programs = min(N*NHV, OVERSUB*num_vectorcore). ++ Each program loops over its assigned (sequence, value-head) tiles with ++ stride num_programs, processing the value dimension in BV=64 blocks and ++ reusing the gating computation across those blocks. ++ ++ q/k/v are addressed by token row stride: for contiguous inputs ++ q_row_stride == H*K / v_row_stride == HV*V, identical to the PR original. ++ """ ++ pid = tl.program_id(0) ++ num_programs = tl.num_programs(0) ++ total_tiles = N * NHV ++ ++ o_k = tl.arange(0, BK) ++ o_base_v = tl.arange(0, BV) ++ mask_k = o_k < K ++ ++ NV = tl.cdiv(V, BV) ++ ++ for g_idx in tl.range(pid, total_tiles, num_programs): ++ i_n = g_idx // NHV ++ i_nhv = g_idx % NHV ++ ++ if IS_VARLEN: ++ bos = tl.load(cu_seqlens + i_n).to(tl.int64) ++ eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) ++ t_len = (eos - bos).to(tl.int32) ++ else: ++ bos = i_n * T ++ t_len = T ++ ++ for i_bhv in tl.static_range(0, BHV): ++ i_hv = i_nhv * BHV + i_bhv ++ i_h = i_hv // (HV // H) ++ ++ # Gating: compute once per (sequence, value-head). ++ b_A_log = tl.load(A_log + i_hv).to(tl.float32) ++ b_a = tl.load(a + bos * HV + i_hv).to(tl.float32) ++ b_dt_bias = tl.load(dt_bias + i_hv).to(tl.float32) ++ b_b = tl.load(b + bos * HV + i_hv).to(tl.float32) ++ ++ x = b_a + b_dt_bias ++ beta_x = softplus_beta * x ++ softplus_x = tl.where( ++ beta_x <= softplus_threshold, ++ (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)), ++ x, ++ ) ++ b_g = -tl.exp(b_A_log) * softplus_x ++ b_decay = tl.exp(b_g) ++ b_beta = 1.0 / (1.0 + tl.exp(-b_b)) ++ ++ # q/k are shared across the value dimension; set pointers once. ++ p_q = q + bos * q_row_stride + i_h * K + o_k ++ p_k = k + bos * k_row_stride + i_h * K + o_k ++ ++ for i_v in tl.range(0, NV): ++ o_v = i_v * BV + o_base_v ++ mask_v = o_v < V ++ mask_h = mask_k[:, None] & mask_v[None, :] ++ ++ # v uses v_row_stride likewise; o is a contiguous buffer ++ # owned by this wrapper and keeps the PR-original addressing. ++ p_v = v + bos * v_row_stride + i_hv * V + o_v ++ p_o = o + (bos * HV + i_hv) * V + o_v ++ ++ if USE_INITIAL_STATE: ++ idx = tl.load(h0_indices + i_n) ++ p_h0 = ( ++ h0_source ++ + idx * HV * K * V ++ + i_hv * K * V ++ + o_k[:, None] * V ++ + o_v[None, :] ++ ) ++ ++ for i in tl.range(0, t_len): ++ b_q = tl.load(p_q + i * q_row_stride, mask=mask_k).to(tl.float32) ++ b_k = tl.load(p_k + i * k_row_stride, mask=mask_k).to(tl.float32) ++ b_v = tl.load(p_v + i * v_row_stride, mask=mask_v).to(tl.float32) ++ ++ if USE_INITIAL_STATE: ++ if idx >= 0: ++ b_h = tl.load(p_h0 + i * HV * K * V, mask=mask_h).to( ++ tl.float32 ++ ) ++ else: ++ b_h = tl.zeros([BK, BV], dtype=tl.float32) ++ else: ++ b_h = tl.zeros([BK, BV], dtype=tl.float32) ++ ++ if USE_QK_L2NORM_IN_KERNEL: ++ b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6) ++ b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6) ++ ++ b_q = b_q * scale ++ ++ # Recurrent delta rule update. ++ b_h *= b_decay ++ b_v -= tl.sum(b_h * b_k[:, None], 0) ++ b_v *= b_beta ++ b_h += b_k[:, None] * b_v[None, :] ++ b_o = tl.sum(b_h * b_q[:, None], 0) ++ ++ if USE_INITIAL_STATE: ++ if idx >= 0: ++ tl.store( ++ p_h0 + i * HV * K * V, ++ b_h.to(h0_source.dtype.element_ty), ++ mask=mask_h, ++ ) ++ ++ tl.store(p_o + i * HV * V, b_o.to(o.dtype.element_ty), mask=mask_v) ++ ++ ++def fused_sigmoid_gating_delta_rule_update_decode_npu( ++ A_log: torch.Tensor, ++ a: torch.Tensor, ++ dt_bias: torch.Tensor, ++ softplus_beta: float, ++ softplus_threshold: float, ++ q: torch.Tensor, ++ k: torch.Tensor, ++ v: torch.Tensor, ++ b: torch.Tensor, ++ initial_state_source: torch.Tensor, ++ initial_state_indices: torch.Tensor, ++ scale: Optional[float] = None, ++ use_qk_l2norm_in_kernel: bool = False, ++ cu_seqlens: Optional[torch.Tensor] = None, ++): ++ """ ++ Decode-optimized recurrent delta rule update for Ascend NPU. ++ ++ Drop-in replacement for the generic ++ ``fused_sigmoid_gating_delta_rule_update_npu`` on the decode path, with ++ the same algorithm. It launches with num_warps=4 for better small-batch ++ throughput on the AIV vector cores. ++ ++ q/k/v may be non-contiguous views (only last-dim stride==1 is required); ++ the kernel addresses them by the passed token row strides, eliminating 3 ++ per-layer q/k/v copies on the decode path. Contiguous inputs address ++ identically to the PR original. ++ ++ Args / Returns: same as ``fused_sigmoid_gating_delta_rule_update_npu`` ++ (return shape matches v; the PR original returns (N, HV, V) and this ++ version ``view``s back to v's shape โ€” under decode the token count == N, ++ so the elements correspond one-to-one). ++ """ ++ # q/k/v stay as views; only last-dim stride==1 is asserted, row strides go ++ # to the kernel. ++ assert q.stride(-1) == 1 and k.stride(-1) == 1 and v.stride(-1) == 1 ++ # The pool must be contiguous: copying a non-contiguous pool via ++ # _maybe_contiguous would write state updates back into the copy ++ # (silently losing them), hence the hard assert. ++ assert initial_state_source.is_contiguous() ++ A_log = _maybe_contiguous(A_log) ++ a = _maybe_contiguous(a) ++ dt_bias = _maybe_contiguous(dt_bias) ++ b = _maybe_contiguous(b) ++ if initial_state_indices is not None: ++ initial_state_indices = _maybe_contiguous(initial_state_indices) ++ if cu_seqlens is not None: ++ cu_seqlens = _maybe_contiguous(cu_seqlens) ++ ++ with torch.npu.device(q.device.index): ++ B, T, H, K = q.shape ++ HV = v.shape[2] ++ V = v.shape[-1] ++ N = B if cu_seqlens is None else len(cu_seqlens) - 1 ++ ++ # Decode-only guard: o is addressed by token index (bos) but has only N ++ # rows, valid only with exactly 1 token per sequence โ€” varlen requires ++ # total_tokens == N (production decode has query_start_loc step 1), ++ # non-varlen requires T == 1. verify/prefill must not use this kernel. ++ if cu_seqlens is not None: ++ assert T == N, ( ++ f"decode-optimized kernel requires one token per sequence, " ++ f"got total_tokens={T}, sequences={N}" ++ ) ++ else: ++ assert T == 1, ( ++ f"decode-optimized kernel requires T == 1 without cu_seqlens, got T={T}" ++ ) ++ ++ BK = triton.next_power_of_2(K) ++ BV = min(triton.next_power_of_2(V), 64) ++ NK = triton.cdiv(K, BK) ++ NV = triton.cdiv(V, BV) ++ assert NK == 1, "NK > 1 is not supported in the decode-optimized kernel" ++ ++ if scale is None: ++ scale = K**-0.5 ++ else: ++ assert scale > 0, "scale must be positive" ++ ++ # Each program handles exactly one value head to keep the UB working set ++ # small and avoid cross-head synchronization. ++ BHV = 1 ++ NHV = HV ++ ++ o = q.new_empty(N, HV, V) ++ ++ # 1-D grid sized to the AIV vector core count; num_warps=4 was tuned for ++ # small-batch decode. Oversubscription does not help because the loop ++ # trip count is already small. ++ num_aicore, num_vectorcore = get_device_properties() ++ OVERSUB = 1 ++ num_programs = min(N * NHV, num_vectorcore * OVERSUB) ++ num_programs = max(1, num_programs) ++ grid = (num_programs,) ++ ++ num_warps = 4 ++ num_stages = 1 ++ ++ _fused_sigmoid_gating_delta_rule_update_decode_kernel[grid]( ++ A_log=A_log, ++ a=a, ++ dt_bias=dt_bias, ++ softplus_beta=softplus_beta, ++ softplus_threshold=softplus_threshold, ++ q=q, ++ k=k, ++ v=v, ++ b=b, ++ o=o, ++ h0_source=initial_state_source, ++ h0_indices=initial_state_indices, ++ cu_seqlens=cu_seqlens, ++ q_row_stride=q.stride(1), ++ k_row_stride=k.stride(1), ++ v_row_stride=v.stride(1), ++ scale=scale, ++ T=T, ++ N=N, ++ NHV=NHV, ++ B=B, ++ H=H, ++ HV=HV, ++ K=K, ++ V=V, ++ BK=BK, ++ BV=BV, ++ BHV=BHV, ++ OVERSUB=OVERSUB, ++ USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, ++ num_warps=num_warps, ++ num_stages=num_stages, ++ multibuffer=False, ++ ) ++ # Match the generic wrapper's return shape (*v.shape); under decode the ++ # token count == N, so the elements correspond one-to-one. ++ return o.view(v.shape) +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/fla/utils.py b/python/sgl_kernel_npu/sgl_kernel_npu/fla/utils.py +index f3028e2..04f32d0 100644 +--- a/python/sgl_kernel_npu/sgl_kernel_npu/fla/utils.py ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/fla/utils.py +@@ -603,8 +603,17 @@ def fused_qkvzba_split_reshape_cat_contiguous( + total_k = num_heads_qk * head_qk + total_v = num_heads_v * head_v + +- qkvz_row_stride = total_q + total_k + total_v + total_v +- ba_row_stride = num_heads_v * 2 ++ # GDN_QKVZBA_PACK: take the row stride from the physical stride(0) โ€” a row ++ # slice view of the packed GEMM output has row stride N_qkvz+N_ba (larger ++ # than the logical width); for contiguous inputs it equals the logical ++ # width, so behavior is unchanged. Non row-major inputs (stride(1) != 1) ++ # fall back to a materialized contiguous copy. ++ if mixed_qkvz.stride(1) != 1: ++ mixed_qkvz = mixed_qkvz.contiguous() ++ if mixed_ba.stride(1) != 1: ++ mixed_ba = mixed_ba.contiguous() ++ qkvz_row_stride = mixed_qkvz.stride(0) ++ ba_row_stride = mixed_ba.stride(0) + qkv_row_stride = total_q + total_k + total_v + z_row_stride = num_heads_v * head_v + ba_out_row_stride = num_heads_v +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/moe/moe_front_routing.py b/python/sgl_kernel_npu/sgl_kernel_npu/moe/moe_front_routing.py +new file mode 100644 +index 0000000..e628291 +--- /dev/null ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/moe/moe_front_routing.py +@@ -0,0 +1,182 @@ ++# -*- coding: utf-8 -*- ++"""MoE ๅ‰ๆฎต v2.2 ่‡ชๅ†™ init_routing๏ผˆTriton / triton-ascend๏ผŒCANN 9.0.0๏ผ‰ใ€‚ ++ ++ๆœๅŠกๅ™จ่ฝ็‚น๏ผšsgl_kernel_npu ๅŒ…ๅ†… ``sgl_kernel_npu/moe/moe_front_routing.py`` ++๏ผˆnamespace ๅญๅŒ…๏ผŒๆ— ้œ€ __init__.py๏ผ›้›†ๆˆๆ–นไปฅ ++``from sgl_kernel_npu.moe.moe_front_routing import moe_init_routing_v22`` ๅฏผๅ…ฅ๏ผ‰ใ€‚ ++ ++่ฏญไน‰ๅฅ‘็บฆ๏ผšไธŽ ``torch.ops.npu.npu_moe_init_routing_v2``๏ผˆexpert_tokens_num_type=1๏ผ‰ ++**้€ไฝไธ€่‡ด**๏ผˆRL ็ฒพๅบฆ็กฌ็บฆๆŸ๏ผŒ0 ๅฎนๅทฎ้ชŒๆ”ถ๏ผ‰๏ผš ++ - expanded_x๏ผšx ่กŒๆŒ‰ใ€Œexpert id ็จณๅฎš่ฎกๆ•ฐๆŽ’ๅบใ€็š„็ฝฎๆข๏ผˆๅŒ expert ๅ†…ไฟๆŒ ++ flat slot ๅŽŸๅบ๏ผŒflat slot ไธบ t-major๏ผšf = t*top_k + k๏ผ‰๏ผ› ++ - eri[f] = flat slot f ๅฏนๅบ”็š„ expanded ่กŒๅท๏ผˆ= stock expanded_row_idx ้€ไฝ๏ผ‰๏ผ› ++ - counts = per-expert ่ฎกๆ•ฐ๏ผˆint64๏ผŒstock type=1 ๅฝขๆ€๏ผŒๅ–‚ stock GMM ++ ``group_list_type=1``๏ผ‰ใ€‚ ++้กบๅธฆๅŽŸ็”Ÿ่พ“ๅ‡บไธคไธชๅ‰็ผ€ๅ’Œ๏ผˆๅ…ไปปไฝ•ๅ‰็ฝฎ kernel๏ผ‰๏ผš ++ - excl๏ผˆint32๏ผ‰๏ผšexclusive ๅ‰็ผ€ๅ’Œ๏ผŒ็›ดๅ–‚ persistent_gmm ็š„ ``offsets=`` ++ ๏ผˆ็œ _gmm2_offsets_kernel๏ผŒ~4.4ยตs/ๅฑ‚๏ผ‰๏ผ› ++ - incl๏ผˆint64๏ผ‰๏ผšinclusive ๅ‰็ผ€ๅ’Œ๏ผŒstock GMM ``group_list_type=0`` ๅฝขๆ€๏ผˆๅค‡็”จ๏ผ‰ใ€‚ ++ ++formulation๏ผˆv2.2๏ผŒcode/moe_front_fusion ๅ…ญ่ฝฎๅ•ๆต‹ๅฎšๆกˆ๏ผ› ++็ป“ๆžœ src/moe_front_fusion_result{,_2.._6}.csv๏ผŒ็ป“่ฎบ่ง ++docs/practice/npu-moe-forward-region-probe.md๏ผ‰๏ผš ++ ++ - rank๏ผˆeri๏ผ‰๏ผš(e,slot) ๅญ—ๅ…ธๅบๆ‰“ๅŒ…ๅ•ๆฏ”่พƒ๏ผŒkey = e*PM + slot๏ผŒใ€Œj ๆŽ’ๅœจ i ๅ‰ใ€ ++ โŸบ key_j < key_i๏ผ›ๆฏ”่พƒ็Ÿฉ้˜ต่ฝฌ็ฝฎ [PM, BM]ใ€ๅฝ’็บฆ่ตฐ axis=0๏ผˆๆœฌๆœบๅฎž่ฏๅฟซๆ–นๅ‘๏ผ‰๏ผ› ++ ๆฏ program ้กบๅธฆไบง**้ƒจๅˆ†็›ดๆ–นๅ›พ**๏ผˆ[BM,PE] axis=0 ไธ€ๆฌกๅฝ’็บฆ๏ผ‰๏ผŒ่ทจ program ++ ๅนถ่กŒๅ–ไปฃไธฒ่กŒๆ‰ซๆใ€‚ ++ - counts/excl/incl๏ผš็ฌฌไบŒไธช kernel ๅฏน partials [NP,PE] ๅš axis=0 ไบŒๆฌกๅฝ’็บฆ + ++ 1D tl.cumsum๏ผˆ็”Ÿไบง้ชŒ่ฏๅŽŸ่ฏญ๏ผ‰็›ดๅ‡บ๏ผŒi64 ็›ดๅ†™็œ cast๏ผ›ๆ—  atomicใ€ๆ— ้›ถๅˆๅง‹ๅŒ–ใ€‚ ++ - gather๏ผˆexpanded_x๏ผ‰๏ผš็‹ฌ็ซ‹ kernel ่กŒๅ—ๅ‘้‡ๅŒ– [RPP,H]๏ผˆRPP=8 ไธบ 32KB๏ผ‰๏ผŒ ++ **ไธไธŽ rank ๅˆๅนถ**โ€”โ€”ไบ”่ฝฎๅฎž่ฏ่žๅˆ็‰ˆ 18.2ยตs ่ดฅ๏ผˆๅฝ’็บฆ้“พๅŽไธฒๆŽฅๅคง tile ่ฎฟๅญ˜ ++ ้•ฟ้“พ๏ผ‰๏ผŒๆ‹†ๅˆ†็‰ˆ้ƒจไปถๅ’Œไป… 6.2ยตsใ€‚ ++ - ็ฒพๅบฆ้ชŒๆ”ถ๏ผšไนๅˆ†ๅธƒ๏ผˆuniform/random/zipf/skew/first/last ๅซๆž็ซฏ๏ผ‰ร— ๅคš seed ++ ไธŽ stock ้€ไฝไธ€่‡ด + finalize ้—ญ็Žฏ 0 diff + ๅŒ่ท‘็กฎๅฎš๏ผˆไธคๅฝขๆ€ๅ„ไธ€่ฝฎๅ…จ็ปฟ๏ผ‰ใ€‚ ++ - ๆ€ง่ƒฝ๏ผˆgraph-loop๏ผŒT=32/topk=8/E=256/H=2048๏ผŒยตs/ๅฑ‚๏ผ‰๏ผšๅ…จ็จ‹ 10.9 vs ++ stock 17.0~18.2๏ผ›้“พ่ทฏๅฃๅพ„ chain_renorm_v22 24.0 vs stock ้“พ 37.4 ++ ๏ผˆๅŒ run -13.4๏ผ‰ใ€‚ ++ ++ๅฝขๆ€้—จ๏ผˆwrapper ๅ†…่‡ช้€‚ๅบ”๏ผ›่ถ…ๅ‡บๅฝขๆ€็”ฑ้›†ๆˆๆ–นๅ›ž้€€ stock๏ผ‰๏ผš ++ M = T*top_k โ‰ค 512 ไธ” M % 8 == 0๏ผŒH ไธบ 2 ็š„ๅน‚๏ผˆgather ่กŒๅ—่ฆๆฑ‚๏ผ‰๏ผŒ ++ x ไธบ bf16ใ€‚bm/rpp ๆŒ‰ M ๆ•ด้™ค่‡ช้€‚ๅบ”๏ผˆ32/16/8 ไธŽ 8/4/2/1๏ผ‰ใ€‚ ++""" ++ ++import torch ++import triton ++import triton.language as tl ++ ++try: ++ from sgl_kernel_npu.fla.utils import input_guard ++except ImportError: # ็‹ฌ็ซ‹่ฐƒ่ฏ•๏ผˆๆœชๅฎ‰่ฃ… sgl_kernel_npu๏ผ‰ๆ—ถ้€€ๅŒ–ไธบๆ’็ญ‰่ฃ…้ฅฐๅ™จ ++ def input_guard(fn): ++ return fn ++ ++ ++def _pow2_ceil(n: int) -> int: ++ p = 1 ++ while p < n: ++ p <<= 1 ++ return p ++ ++ ++@triton.jit ++def _rank_hist_kernel( ++ ids_ptr, eri_ptr, part_ptr, ++ M: tl.constexpr, PM: tl.constexpr, PE: tl.constexpr, BM: tl.constexpr, ++): ++ """eri + ้ƒจๅˆ†็›ดๆ–นๅ›พ๏ผˆgrid = M//BM๏ผ‰ใ€‚ ++ ++ lt[j, i] = key_j < key_i๏ผˆ[PM, BM]๏ผŒaxis=0 ๅฝ’็บฆโ€”โ€”ๆœฌๆœบๅฎž่ฏๅฟซๆ–นๅ‘๏ผ‰๏ผ› ++ part[pid] = ๆœฌ program BM ไธช slot ็š„ expert ็›ดๆ–นๅ›พ๏ผˆ[BM, PE] axis=0๏ผ‰ใ€‚ ++ padding slot ็š„ key=32768*PM+... ๅคงไบŽไธ€ๅˆ‡็œŸๅฎž key๏ผŒไธๆ”นๅ˜็œŸๅฎžๅๆฌกใ€‚ ++ """ ++ pid = tl.program_id(0) ++ rows = pid * BM + tl.arange(0, BM) ++ jj = tl.arange(0, PM) ++ e_all = tl.load(ids_ptr + jj, mask=jj < M, other=32768).to(tl.int32) ++ e_rows = tl.load(ids_ptr + rows, mask=rows < M, other=-1).to(tl.int32) ++ key_all = e_all * PM + jj ++ key_rows = e_rows * PM + rows ++ lt = key_all[:, None] < key_rows[None, :] # [PM, BM] ++ rank = tl.sum(lt.to(tl.int16), axis=0) # [BM] ++ tl.store(eri_ptr + rows, rank.to(tl.int32), mask=rows < M) ++ rngE = tl.arange(0, PE) ++ eq = e_rows[:, None] == rngE[None, :] # [BM, PE]๏ผˆ-1 ไธๅ‘ฝไธญไปปไฝ• expert๏ผ‰ ++ part = tl.sum(eq.to(tl.int16), axis=0).to(tl.int32) ++ tl.store(part_ptr + pid * PE + rngE, part) ++ ++ ++@triton.jit ++def _partials_cumsum_kernel( ++ part_ptr, counts_ptr, excl_ptr, incl_ptr, ++ NP: tl.constexpr, NPP: tl.constexpr, PE: tl.constexpr, E: tl.constexpr, ++): ++ """partials [NP, PE] -> counts(i64) / excl(i32) / incl(i64)๏ผŒๅ• programใ€‚ ++ ++ ๅ…จ็™ฝๅๅ•ๅŽŸ่ฏญ๏ผšaxis=0 ไบŒๆฌกๅฝ’็บฆ + 1D tl.cumsum๏ผˆGDN ็”Ÿไบง่ทฏๅพ„ๅŒๆฌพ๏ผ‰ใ€‚ ++ """ ++ np_ = tl.arange(0, NPP) ++ rngE = tl.arange(0, PE) ++ part = tl.load(part_ptr + np_[:, None] * PE + rngE[None, :], ++ mask=(np_ < NP)[:, None], other=0) ++ counts = tl.sum(part, axis=0) # [PE] i32 ++ incl = tl.cumsum(counts, axis=0) ++ excl = incl - counts ++ e_valid = rngE < E ++ tl.store(counts_ptr + rngE, counts.to(tl.int64), mask=e_valid) ++ tl.store(excl_ptr + rngE, excl, mask=e_valid) ++ tl.store(incl_ptr + rngE, incl.to(tl.int64), mask=e_valid) ++ ++ ++@triton.jit ++def _routing_gather_kernel( ++ x_ptr, eri_ptr, out_ptr, ++ TOPK: tl.constexpr, H: tl.constexpr, RPP: tl.constexpr, ++): ++ """่กŒๅ—็ฝฎๆข๏ผšout[eri[f]] = x[f // TOPK]๏ผŒๆฏ program ๅค„็† RPP ไธช flat slotใ€‚ ++ ++ [RPP, H] ไบŒ็ปด tile ไธ€ๆฌก่ฃ…่ฝฝ/ๅ†™ๅ‡บ๏ผˆRPP=8 ไธบ 32KB๏ผ‰๏ผŒ่กŒๅ†… 4KB ่ฟž็ปญใ€‚ ++ grid = (M // RPP,)๏ผˆwrapper ไฟ่ฏ M % RPP == 0๏ผŒๆ•…ไธๅธฆ mask๏ผ‰ใ€‚ ++ """ ++ pid = tl.program_id(0) ++ fs = pid * RPP + tl.arange(0, RPP) ++ dest = tl.load(eri_ptr + fs) # [RPP] int32 ++ tok = fs // TOPK ++ hh = tl.arange(0, H) ++ rows = tl.load(x_ptr + tok[:, None] * H + hh[None, :]) ++ tl.store(out_ptr + dest[:, None] * H + hh[None, :], rows) ++ ++ ++def _pick_tile(m: int, candidates) -> int: ++ for c in candidates: ++ if m % c == 0: ++ return c ++ return candidates[-1] ++ ++ ++@input_guard ++def moe_init_routing_v22( ++ x: torch.Tensor, ++ topk_ids: torch.Tensor, ++ num_experts: int, ++ top_k: int, ++): ++ """v2.2 ่‡ชๅ†™ init_routing๏ผˆrank_hist + partials_cumsum + gather ไธ‰ launch๏ผ‰ใ€‚ ++ ++ ๅฝขๆ€้—จ๏ผˆ่ฐƒ็”จๆ–น่ดŸ่ดฃๅˆคๅฎš๏ผŒ่ถ…ๅ‡บ้—จ่ฏทๅ›ž้€€ stock๏ผ‰๏ผšM = T*top_k โ‰ค 512 ไธ” ++ M % 8 == 0๏ผ›H ไธบ 2 ็š„ๅน‚๏ผ›x ไธบ bf16 contiguousใ€‚ ++ ++ ๅ‚ๆ•ฐ๏ผš ++ x: [T, H] bf16๏ผŒdecode ้š่—็Šถๆ€ใ€‚ ++ topk_ids: [T, top_k] int32 contiguous๏ผˆๆฅ่‡ช gating ็ฎ—ๅญ๏ผ‰ใ€‚ ++ num_experts / top_k๏ผš่ทฏ็”ฑ้…็ฝฎใ€‚ ++ ++ ่ฟ”ๅ›ž (expanded_x, eri, counts_i64, excl_i32, incl_i64)๏ผš ++ expanded_x [M, H] bf16 / eri [M] int32 โ€”โ€” ไธŽ stock ้€ไฝไธ€่‡ด๏ผˆ0 ๅฎนๅทฎ๏ผ‰๏ผ› ++ counts_i64 [E] โ€”โ€” stock type=1 ๅฝขๆ€๏ผˆGMM group_list_type=1๏ผ‰๏ผ› ++ excl_i32 [E] โ€”โ€” persistent_gmm ็š„ offsets= ็›ด็”จ๏ผˆ็œ 4.4ยตs ๅ‰็ฝฎ kernel๏ผ‰๏ผ› ++ incl_i64 [E] โ€”โ€” stock GMM group_list_type=0 ๅฝขๆ€๏ผˆๅค‡็”จ๏ผ‰ใ€‚ ++ """ ++ T, H = x.shape ++ M = T * top_k ++ ids_flat = topk_ids.reshape(-1).to(torch.int32) # ๅŒ dtype ๆ—ถ .to ไธๆ‹ท่ด ++ dev = x.device ++ PM = _pow2_ceil(M) ++ PE = _pow2_ceil(num_experts) ++ bm = _pick_tile(M, (32, 16, 8)) ++ rpp = _pick_tile(M, (8, 4, 2, 1)) ++ npart = M // bm ++ eri = torch.empty(M, dtype=torch.int32, device=dev) ++ partials = torch.empty((npart, PE), dtype=torch.int32, device=dev) ++ expanded = torch.empty((M, H), dtype=x.dtype, device=dev) ++ counts = torch.empty(num_experts, dtype=torch.int64, device=dev) ++ excl = torch.empty(num_experts, dtype=torch.int32, device=dev) ++ incl = torch.empty(num_experts, dtype=torch.int64, device=dev) ++ _rank_hist_kernel[(npart,)](ids_flat, eri, partials, ++ M=M, PM=PM, PE=PE, BM=bm, num_stages=1) ++ _routing_gather_kernel[(M // rpp,)](x, eri, expanded, ++ TOPK=top_k, H=H, RPP=rpp) ++ _partials_cumsum_kernel[(1,)](partials, counts, excl, incl, ++ NP=npart, NPP=_pow2_ceil(npart), PE=PE, ++ E=num_experts, num_stages=1) ++ return expanded, eri, counts, excl, incl +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/moe/persistent_gmm.py b/python/sgl_kernel_npu/sgl_kernel_npu/moe/persistent_gmm.py +new file mode 100644 +index 0000000..8e28e82 +--- /dev/null ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/moe/persistent_gmm.py +@@ -0,0 +1,218 @@ ++# -*- coding: utf-8 -*- ++"""Persistent grouped matmul๏ผˆTriton / triton-ascend๏ผ‰๏ผŒ้ขๅ‘ EP1 decode ็š„ skinny GMM2ใ€‚ ++ ++ๆœๅŠกๅ™จ่ฝ็‚น๏ผšsgl_kernel_npu ๅŒ…ๅ†… ``sgl_kernel_npu/moe/persistent_gmm.py`` ++๏ผˆnamespace ๅญๅŒ…๏ผŒๆ— ้œ€ __init__.py๏ผ›้›†ๆˆๆ–นไปฅ ++``from sgl_kernel_npu.moe.persistent_gmm import persistent_grouped_matmul`` ๅฏผๅ…ฅ๏ผ‰ใ€‚ ++ ++่ฟ็งป่ฏดๆ˜Ž๏ผˆgmm/v1๏ผŒๅŸบ็บฟ sglang main@bd3f6a793 + sgl_kernel_npu 3.2.1๏ผ‰๏ผš ++ ้ขๅ‘ 8 ๅก 16 die / 4 engine๏ผˆMoE-TPใ€EP1๏ผ‰decode๏ผšE=256, N=2048๏ผŒ ++ K=64(TP8)/128(TP4)๏ผŒๅนณๅ‡ ~1 rows/expert๏ผŒstock GMM2 ๅฎžๆต‹ ~65us ++ ๏ผˆmac_ratio 6.9%๏ผŒscalar/fixpipe ๆŽฅ่ฟ‘้ฅฑๅ’Œ๏ผŒAIV ๅ…จ้—ฒ๏ผ‰ใ€‚็“ถ้ขˆ็ป“ๆž„๏ผš ++ ้€ group ่ฐƒๅบฆไธŽ็ป„้—ดๆตๆฐด็บฟๆŽ’็ฉบ๏ผŒ่€Œ้ž FLOPsใ€‚ ++ ++v1.1 ไฟฎๅค๏ผˆ2026-08-10๏ผŒbisect ๅฎž่ฏ่ง src/bisect_result.txt๏ผ‰๏ผš ++ v1 ็š„ masked-sum in-kernel offsets๏ผˆrow0 = sum(where(e_rng < e, c_vec, 0))๏ผŒ ++ ็ปงๆ‰ฟ่‡ช old_env v1.2.1๏ผŒไธ”ๅฝ“ๅนดไฟฎๅคๅŽๆ—  rollout ๅคๆต‹่ฎฐๅฝ•๏ผ‰ๅœจๆœฌๆœบๅทฅๅ…ท้“พ ++ ๏ผˆCANN 9.0.0 / triton-ascend๏ผŒ้ป˜่ฎค --enable-auto-multi-buffer๏ผ‰ไธ‹่ขซ่ฏฏ็ผ–่ฏ‘๏ผš ++ persistent ๆ•ฐๆฎไพ่ต–ๆŽงๅˆถๆตไธญ็š„ scalar-masked reduction ่ฎก็ฎ—ๅ‡บ้”™่ฏฏ็š„ row0๏ผŒ ++ ่พ“ๅ‡บๆ•ดไฝ“้”™ไฝ๏ผˆstages=3 rel_diffโ‰ˆ1.3๏ผ›stages=1 ็›ดๆŽฅ 1.7e38 ๅžƒๅœพๅ€ผ๏ผ‰๏ผŒ ++ rollout ไนฑ็ ใ€‚ไบŒๅˆ†่ฏๅฎž๏ผšๅŒไธ€ kernel ๆ”น็”จๅค–้ƒจ offsets ๆ ‡้‡็›ด่ฏปๅŽไธŽ stock ++ GMM ้€ไฝไธ€่‡ด๏ผˆmax_abs=0.000000๏ผŒtl.trans / num_stages / base_n ๅ‡ๆ— ่พœ๏ผ‰ใ€‚ ++ ๆœฌ็‰ˆ offsets ๆ”นไธบใ€Œๅ‰็ฝฎ offsets kernel + ๆ ‡้‡็›ด่ฏปใ€๏ผš ++ - _gmm2_offsets_kernel๏ผšๅ• program ๅ‰็ผ€ๅ’Œ๏ผ›ๆ›ฟไปฃ torch.cumsum ้“พ ++ ๏ผˆ~58us/ๅฑ‚๏ผŒไธ”ๅ…ถ torch ็ฎ—ๅญๅฏน graph capture ไธๅ‹ๅฅฝ๏ผ‰ใ€‚ ++ - _persistent_gmm_kernel๏ผšrow0 = tl.load(offsets_ptr + e) ๆ ‡้‡็›ด่ฏป๏ผŒ ++ ไธŽ bisect ไธญ้€ไฝๆญฃ็กฎ็š„ ext_off ่ทฏๅพ„ๅฎŒๅ…จไธ€่‡ดใ€‚ ++ ่‹ฅๅ‰็ฝฎ kernel ๅœจๆœฌๆœบไป้ชŒ่ฏไธ่ฟ‡๏ผŒๅ›ž้€€ๆ–นๆกˆ๏ผšwrapper ็”จ ++ ``offsets = torch.cumsum(group_list, 0) - group_list``๏ผˆ2 ไธชๅฐ็ฎ—ๅญ๏ผŒ ++ ่ง persistent_grouped_matmul ๆณจ้‡Š๏ผ‰๏ผŒkernel ๆ— ้œ€ๆ”นๅŠจใ€‚ ++ ++v1.2 ไฟฎๅค๏ผˆ2026-08-11๏ผŒrollout ๅฎž่ฏ๏ผšoffsets kernel 37us ๅžๆŽ‰ๅ…จ้ƒจๆ”ถ็›Š๏ผ‰๏ผš ++ v1.1 ็š„ไธค็บงๅˆ†็ป„ๆ‰ซๆ็”จไบ† [G,G,G] ไธ‰็ปด tl.where+tl.sum(axis=2)๏ผŒๅœจๆœฌๆœบ ++ triton-ascend ไธŠๆœช่ขซๅ‘้‡ๅŒ–ใ€scalarize ๆˆ้€ๅ…ƒ็ด ๅพช็Žฏ๏ผˆ4096 ๆฌกๆ ‡้‡ไน˜ๅŠ  ร— ++ ๅ• program๏ผ‰๏ผŒๅฎžๆต‹ ~37us/ๅฑ‚๏ผŒๅžๆŽ‰ persistent kernel ็œไธ‹็š„ ~24us/ๅฑ‚ใ€‚ ++ ๏ผˆrollout ๅฎžๆต‹๏ผšpersistent kernel ๆœฌไฝ“ 41us vs stock 65us ๆญฃ็กฎไธ”็จณๅฎšใ€‚๏ผ‰ ++ ++v1.3 ไฟฎๅค๏ผˆ2026-08-11๏ผ‰๏ผšv1.2 ็š„ Hillis-Steele๏ผˆๅ…จๅฑ€ scratch ๅฐฑๅœฐ ++ store->็งปไฝ load๏ผ‰ๅœจ ConvertTritonIRToLinalgIR ้˜ถๆฎต PassManager ็ผ–่ฏ‘ๅคฑ่ดฅใ€‚ ++ ๆ”น็”จ tl.cumsum ไธ€็ปดๆ‰ซๆโ€”โ€”่ฏฅๅŽŸ่ฏญๅœจๆœฌๅทฅๅ…ท้“พๅทฒ่ขซ็”Ÿไบง้ชŒ่ฏ ++ ๏ผˆsgl_kernel_npu.fla.cumsum ็š„ chunk_local_cumsum_scalar_kernel๏ผŒGDN ่ทฏๅพ„๏ผ‰๏ผŒ ++ ๆ‰ซๆๅ‰ๅ…ˆ่ฝฌ int32๏ผˆint64 reduction ้ฃŽ้™ฉ็š„่€ๆ•™่ฎญ๏ผ‰ใ€‚็›ฎๆ ‡ <=5-7us/ๅฑ‚ใ€‚ ++ ++v1.4 ่ฐƒไผ˜๏ผˆ2026-08-11๏ผŒcode/moe_whole_process ๆŽข้’ˆๅฎžๆต‹๏ผ‰๏ผš ++ num_progs 40โ†’24ใ€‚Ascend910_9382 ไธบ 24 AIC/die๏ผŒ้™ๆ€ round-robin ไธ‹ ++ p40 ๆœ‰ 16 ๆ ธ่ƒŒ 2 ไธช program ๅฝขๆˆๆ…ขๆณขใ€p32 ๆ›ด็พ้šพ๏ผˆ8 ๆ ธๅŒไปฝ๏ผ‰๏ผ› ++ uniform counts๏ผˆๅ…จๆดป่ทƒ๏ผ‰ๅ•ๆต‹ 64.3โ†’54.1us๏ผˆ1.03โ†’1.24TB/s๏ผŒ่พพ ++ read_only_sum ๅฎžๆต‹ไธŠ้™ 1.33TB/s ็š„ ~93%๏ผŒๅŸบๆœฌ่ดดๅข™๏ผ‰ใ€‚ ++ ๅŒๆฌกๆŽข้’ˆ่ฏไผชไธคไธชๅ‡่ฎพ๏ผš ++ - ๆƒ้‡ๆ— ๆกไปถ่ฃ…่ฝฝ๏ผˆvarU๏ผŒ่ฃ…่ฝฝ็งปๅ‡บ if m_i>0๏ผ‰ๆŒ‰ๅฎž้™…่ฏปๅ–ๅญ—่Š‚ๆŠ˜็ฎ—้€Ÿ็Ž‡ ++ ไธŽๆกไปถ่ฃ…่ฝฝๅฎŒๅ…จ็›ธๅŒ๏ผˆ~1.1TB/s๏ผ‰โ€”โ€”็“ถ้ขˆไธๆ˜ฏ if ้˜ปๆ–ญ่ทจ่ฟญไปฃๆตๆฐด๏ผ› ++ - num_stages 2/3/4 ่ฎกๆ—ถๆ— ๅทฎๅผ‚ใ€‚ ++ HBM ๅฏ่พพๅธฆๅฎฝๆ ‡ๅฎš๏ผšread 1328 / copy r+w 1275 GB/s๏ผˆไธŽๆ—ง็Žฏๅขƒ ++ 1311/1214 ไธ€่‡ด๏ผ‰๏ผ›GMM1 ็”Ÿไบง profile 61us๏ผˆimplied 2.2TB/s๏ผ‰ไธŽ่ฏฅไธŠ้™ ++ ็Ÿ›็›พ๏ผŒๅ†ทๆ€ floor ๅบ” ~101us๏ผŒๅฃๅพ„ๅพ…ๆพ„ๆธ…๏ผˆL2 ็ƒญๆˆ– shape ๆททๅ…ฅ๏ผ‰ใ€‚ ++ ++kernel ่ฎพ่ฎก๏ผˆ็ปงๆ‰ฟ v1.2.1 ็ป“่ฎบ๏ผ‰๏ผš ++ ++ - persistent grid๏ผšnum_progs ไธชๅธธ้ฉป program๏ผŒ้™ๆ€ round-robin ้ข†ๅ– (expert, n_tile) ++ ไปปๅŠก๏ผ›ไปปๅŠกๆˆๆœฌ็”ฑๆƒ้‡ๅญ—่Š‚ไธปๅฏผไธ”ๅ„ไปปๅŠกๅ‡ๅŒ€๏ผˆไธŽ M_i ๆ— ๅ…ณ๏ผ‰๏ผŒ้™ๆ€่ฐƒๅบฆๅณๅคŸ๏ผŒ ++ ๆ— ้œ€ atomic ้˜Ÿๅˆ—ใ€‚ ++ - ็ฉบ expert๏ผˆcount=0๏ผ‰ๆ•ดไปปๅŠก่ทณ่ฟ‡๏ผŒๆญป expert ๆƒ้‡้›ถๆต้‡ใ€‚ ++ - ๆƒ้‡ tile ๅœจไปปๅŠกๅ†…ๅฏนๆ‰€ๆœ‰ m_tile ๅค็”จ๏ผ›BASE_M=16 + ๅ†…ๅฑ‚ m_tile ๅพช็Žฏ่ฆ†็›–็ƒญ็‚น expertใ€‚ ++ - fp32 ็ดฏๅŠ ๏ผˆtl.dot ้ป˜่ฎค๏ผ‰๏ผŒไธŽ torch.ops.npu.npu_grouped_matmul ๆ•ฐๅ€ผ่ฏญไน‰ไธ€่‡ดใ€‚ ++ - layout="kn" ่ทฏๅพ„ๆณจๆ„๏ผšๆœฌๆœบๅทฅๅ…ท้“พ base_n=512 ็ผ–่ฏ‘ๆœŸ cbuf overflow๏ผŒ ++ ๅฆ‚็กฎ้œ€ kn ๅธƒๅฑ€ๅ…ˆ่ฏ• base_n<=256๏ผˆnk ไธบ็”Ÿไบง่ทฏๅพ„๏ผŒไธๅ—ๅฝฑๅ“๏ผ‰ใ€‚ ++ ++่พ“ๅ…ฅๅฅ‘็บฆ๏ผš ++ x: [total_M, K] bf16๏ผŒ่กŒๆŒ‰ expert ๅˆ†็ป„่ฟž็ปญ๏ผˆnpu_moe_init_routing_v2 ่พ“ๅ‡บๅบ๏ผ‰ ++ w: [E, N, K] bf16๏ผˆlayout="nk"๏ผŒ็”Ÿไบงๅญ˜ๅ‚จๅธƒๅฑ€๏ผ‰ๆˆ– [E, K, N] ่ฟž็ปญ๏ผˆlayout="kn"๏ผ‰ใ€‚ ++ ๆณจๆ„๏ผšw ๅฟ…้กปๆ˜ฏ ND ่ฟž็ปญๅญ˜ๅ‚จ๏ผ›้›†ๆˆๆ–นๅœจ _gmm2_triton ๅผ€ๅฏๆ—ถๅทฒ่ทณ่ฟ‡ ++ w2 ็š„ FRACTAL_NZ cast๏ผŒNZ ้‡ๆŽ’ๅญ˜ๅ‚จไผš่ฎฉ่ฃธๆŒ‡้’ˆ่ฏปๅ–ๅ‡บ้”™ใ€‚ ++ group_list: [E] per-expert counts๏ผˆcount ๆจกๅผ๏ผŒint32/int64๏ผŒdevice๏ผ‰ ++""" ++ ++import torch ++import triton ++import triton.language as tl ++ ++try: ++ from sgl_kernel_npu.fla.utils import input_guard ++except ImportError: # ็‹ฌ็ซ‹่ฐƒ่ฏ•๏ผˆๆœชๅฎ‰่ฃ… sgl_kernel_npu๏ผ‰ๆ—ถ้€€ๅŒ–ไธบๆ’็ญ‰่ฃ…้ฅฐๅ™จ ++ def input_guard(fn): ++ return fn ++ ++ ++@triton.jit ++def _gmm2_offsets_kernel( ++ counts_ptr, offsets_ptr, ++ E: tl.constexpr, P: tl.constexpr, ++): ++ """exclusive prefix offsets[e] = sum(counts[=E ็š„ๆœ€ๅฐ 2 ็š„ๅน‚๏ผ‰ใ€‚formulation ้€‰ๆ‹ฉไพๆฎ ++ ๏ผˆๆœฌๆœบ CANN 9.0.0 / triton-ascend ๅฎž่ฏ๏ผ‰๏ผš ++ - tl.cumsum ๅœจๆœฌๅทฅๅ…ท้“พๅฏ็”จ๏ผšsgl_kernel_npu.fla.cumsum ็š„ ++ chunk_local_cumsum_scalar_kernel๏ผˆGDN ็”Ÿไบง่ทฏๅพ„๏ผ‰ๅณ็”จ tl.cumsum๏ผ› ++ - ๅ…ˆๅš int32 ่ฝฌๆขๅ†ๆ‰ซๆ๏ผš่ง„้ฟ triton-ascend ็š„ int64 reduction ้ฃŽ้™ฉ ++ ๏ผˆv1.2.1 ่ตท็š„่€ๆ•™่ฎญ๏ผ‰๏ผ› ++ - ็ฆๆญข formulation๏ผšpersistent ๅŠจๆ€ๆŽงๅˆถๆตๅ†… scalar-masked reduction ++ ๏ผˆv1๏ผŒ่ฏฏ็ผ–่ฏ‘ๆ•ฐๅ€ผ้”™๏ผ‰๏ผ›ไธ‰็ปด where/sum tile๏ผˆv1.1๏ผŒscalarize ~37us๏ผ‰๏ผ› ++ ๅ…จๅฑ€ scratch ๅฐฑๅœฐ store->็งปไฝ load ็š„ Hillis-Steele๏ผˆv1.2๏ผŒ ++ ConvertTritonIRToLinalgIR PassManager ็ผ–่ฏ‘ๅคฑ่ดฅ๏ผ‰ใ€‚ ++ """ ++ rng = tl.arange(0, P) ++ e_mask = rng < E ++ c = tl.load(counts_ptr + rng, mask=e_mask, other=0).to(tl.int32) ++ incl = tl.cumsum(c, axis=0) ++ tl.store(offsets_ptr + rng, incl - c, mask=e_mask) ++ ++ ++@triton.jit ++def _persistent_gmm_kernel( ++ x_ptr, w_ptr, counts_ptr, offsets_ptr, out_ptr, ++ num_progs, ++ E: tl.constexpr, K: tl.constexpr, N: tl.constexpr, ++ BASE_M: tl.constexpr, BASE_N: tl.constexpr, W_KN: tl.constexpr, ++): ++ pid = tl.program_id(0) ++ N_TILES: tl.constexpr = N // BASE_N ++ TOTAL_TASKS: tl.constexpr = E * N_TILES ++ k_rng = tl.arange(0, K) ++ bm_rng = tl.arange(0, BASE_M) ++ bn_rng = tl.arange(0, BASE_N) ++ ++ for t in range(pid, TOTAL_TASKS, num_progs): ++ e = t // N_TILES ++ j = t % N_TILES ++ m_i = tl.load(counts_ptr + e).to(tl.int32) ++ if m_i > 0: ++ # bisect ้€ไฝ้ชŒ่ฏ่ฟ‡็š„ๅฝขๆ€๏ผšๅค–้ƒจ offsets + ๆ ‡้‡็›ด่ฏปใ€‚ ++ # ไธฅ็ฆๆ”นๅ›ž kernel ๅ†… masked-sum ๆŽจๅฏผ๏ผˆๆœฌๆœบๅทฅๅ…ท้“พ่ฏฏ็ผ–่ฏ‘๏ผŒrollout ไนฑ็ ๏ผ‰ใ€‚ ++ row0 = tl.load(offsets_ptr + e) ++ if W_KN: ++ # [E,K,N] ่ฟž็ปญ๏ผš็›ดๆŽฅ่ฏป [K, BASE_N] tile๏ผˆk ่กŒๅ„ 2*BASE_N ๅญ—่Š‚่ฟž็ปญ็ชๅ‘๏ผ‰ ++ w_t = tl.load( ++ w_ptr + e.to(tl.int64) * N * K ++ + k_rng[:, None] * N + (j * BASE_N + bn_rng)[None, :] ++ ) ++ else: ++ # [E,N,K] ่ฟž็ปญ๏ผˆ็”Ÿไบงๅญ˜ๅ‚จๅธƒๅฑ€๏ผ‰๏ผš่ฏป [BASE_N, K] ๅŽๅฏ„ๅญ˜ๅ™จ่ฝฌ็ฝฎ ++ w_tile = tl.load( ++ w_ptr + e.to(tl.int64) * N * K ++ + (j * BASE_N + bn_rng)[:, None] * K + k_rng[None, :] ++ ) ++ w_t = tl.trans(w_tile) ++ for m0 in range(0, m_i, BASE_M): ++ rows = row0 + m0 + bm_rng ++ m_mask = (m0 + bm_rng) < m_i ++ x_tile = tl.load( ++ x_ptr + rows[:, None] * K + k_rng[None, :], ++ mask=m_mask[:, None], other=0.0, ++ ) ++ acc = tl.dot(x_tile, w_t) # fp32 [BASE_M, BASE_N] ++ tl.store( ++ out_ptr + rows[:, None] * N + (j * BASE_N + bn_rng)[None, :], ++ acc.to(tl.bfloat16), ++ mask=m_mask[:, None], ++ ) ++ ++ ++def _pow2_ceil(n: int) -> int: ++ p = 1 ++ while p < n: ++ p <<= 1 ++ return p ++ ++ ++@input_guard ++def persistent_grouped_matmul( ++ x: torch.Tensor, ++ w: torch.Tensor, ++ group_list: torch.Tensor, ++ base_n: int = 512, ++ num_progs: int = 24, ++ num_stages: int = 3, ++ layout: str = "nk", ++ out: torch.Tensor = None, ++ offsets: torch.Tensor = None, ++) -> torch.Tensor: ++ """GMM2 ไธ“็”จ persistent grouped matmul๏ผŒ่ฟ”ๅ›ž [total_M, N] bf16ใ€‚ ++ ++ base_n: N ๆ–นๅ‘ tile๏ผˆ{256,512,1024} ไน‹ไธ€๏ผŒ้กปๆ•ด้™ค N๏ผ‰ ++ num_progs: ๅธธ้ฉป program ๆ•ฐใ€‚้ป˜่ฎค 24๏ผšAscend910_9382 ไธบ 24 AIC/die๏ผŒ ++ ้™ๆ€ round-robin ไธ‹ไธŽๆ ธๆ•ฐ็›ธ็ญ‰ๆ‰่ดŸ่ฝฝๅ‡่กก๏ผˆv1.4 ๆŽข้’ˆๅฎžๆต‹๏ผš ++ p40 ๆ…ข ~19%๏ผŒp32 ๆœ€็พ้šพ๏ผ›ๆ‰ซๅ‚่Œƒๅ›ด {20,24,32,40,48}๏ผ‰ใ€‚ ++ num_stages: triton ๆตๆฐด็บงๆ•ฐ๏ผˆๆŽข้’ˆๅฎžๆต‹ 2/3/4 ๆ— ๅทฎๅผ‚๏ผŒไฟๆŒ้ป˜่ฎค 3๏ผ‰ ++ layout: "nk" = w ไธบ [E,N,K]๏ผˆ็”Ÿไบงๅญ˜ๅ‚จๅธƒๅฑ€๏ผ‰๏ผ›"kn" = w ไธบ [E,K,N] ่ฟž็ปญ ++ offsets: ๅฏ้€‰ๅค–้ƒจ exclusive offsets๏ผˆint32 [E]๏ผ‰๏ผ›็ผบ็œๆ—ถ็”ฑๆœฌๅ‡ฝๆ•ฐๅ†…็ฝฎ็š„ ++ offsets ๅ‰็ฝฎ kernel ๆŽจๅฏผ๏ผˆๆ—  torch ่ฎก็ฎ—๏ผŒgraph capture ๅ‹ๅฅฝ๏ผ‰ใ€‚ ++ ๅ›ž้€€ๅฝขๆ€๏ผšoffsets = torch.cumsum(group_list, 0) - group_list ++ ๏ผˆto(int32) ๅŽไผ ๅ…ฅ๏ผ‰๏ผŒkernel ๆ— ้œ€ไปปไฝ•ๆ”นๅŠจใ€‚ ++ """ ++ E, N, K = w.shape ++ total_M = x.shape[0] ++ if out is None: ++ out = torch.empty((total_M, N), dtype=torch.bfloat16, device=x.device) ++ if offsets is None: ++ offsets = torch.empty(E, dtype=torch.int32, device=x.device) ++ P = _pow2_ceil(E) ++ _gmm2_offsets_kernel[(1,)]( ++ group_list, offsets, E=E, P=P ++ ) ++ _persistent_gmm_kernel[(num_progs,)]( ++ x_ptr=x, ++ w_ptr=w, ++ counts_ptr=group_list, ++ offsets_ptr=offsets, ++ out_ptr=out, ++ num_progs=num_progs, ++ E=E, ++ K=K, ++ N=N, ++ BASE_M=16, ++ BASE_N=base_n, ++ W_KN=(layout == "kn"), ++ num_stages=num_stages, ++ ) ++ return out diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/sample/__init__.py b/python/sgl_kernel_npu/sgl_kernel_npu/sample/__init__.py index e69de29..6e7bbd6 100644 --- a/python/sgl_kernel_npu/sgl_kernel_npu/sample/__init__.py @@ -1845,3 +4752,6 @@ index 0000000..3c12d53 + ) + else: + raise SystemExit(pytest.main([__file__])) +-- +2.54.0.windows.1 + diff --git a/docker/npu_patch/sglang-npu.patch b/docker/npu_patch/sglang-npu.patch index f7bc22585..beb8590f6 100644 --- a/docker/npu_patch/sglang-npu.patch +++ b/docker/npu_patch/sglang-npu.patch @@ -1,3 +1,48 @@ +From 28aef8685f471c589d2d85ec954e6c2a786a30c7 Mon Sep 17 00:00:00 2001 +From: liyongwen <1310439159@qq.com> +Date: Sat, 29 Aug 2026 03:49:49 +0000 +Subject: [PATCH] sglang optimize for qwen3.5 9B&35B + +--- + python/sglang/srt/entrypoints/http_server.py | 2 +- + python/sglang/srt/environ.py | 42 ++ + .../npu/attention/ascend_gdn_backend.py | 190 +++++-- + .../ascend_hybrid_linear_attn_backend.py | 61 +- + .../attention/full_attention_fusion_npu.py | 468 ++++++++++++++++ + .../npu/graph_runner/npu_cudagraph_backend.py | 18 +- + .../srt/hardware_backend/npu/moe/topk.py | 53 +- + .../npu/moe_weight_prefetch.py | 363 ++++++++++++ + .../npu/tp_ascendc_fusion_npu.py | 347 ++++++++++++ + .../sglang/srt/layers/attention/fla/cumsum.py | 7 +- + .../attention/linear/kernels/gdn_triton.py | 26 +- + .../srt/layers/exp_race_fused_triton.py | 239 ++++++++ + python/sglang/srt/layers/layernorm.py | 21 +- + .../srt/layers/moe/token_dispatcher/deepep.py | 6 + + .../sglang/srt/layers/quantization/unquant.py | 100 +++- + .../srt/layers/radix_linear_attention.py | 21 + + python/sglang/srt/layers/sampler.py | 354 +++++++++++- + python/sglang/srt/managers/scheduler.py | 2 +- + python/sglang/srt/managers/tp_worker.py | 30 + + .../sglang/srt/mem_cache/mamba_radix_cache.py | 9 + + .../srt/model_executor/forward_batch_info.py | 25 +- + .../sglang/srt/model_executor/model_runner.py | 26 + + python/sglang/srt/models/qwen2_moe.py | 48 +- + python/sglang/srt/models/qwen3_5.py | 521 +++++++++++++++--- + .../srt/sampling/sampling_batch_info.py | 10 + + python/sglang/srt/speculative/eagle_utils.py | 98 ++-- + .../srt/speculative/triton_ops/cache_locs.py | 9 + + python/sglang/srt/utils/common.py | 4 +- + .../srt/utils/torch_memory_saver_adapter.py | 9 +- + .../layers/test_conv_weight_cache_refresh.py | 109 ++++ + .../layers/test_sampler_attach_logprobs.py | 156 ++++++ + 31 files changed, 3134 insertions(+), 240 deletions(-) + create mode 100644 python/sglang/srt/hardware_backend/npu/attention/full_attention_fusion_npu.py + create mode 100644 python/sglang/srt/hardware_backend/npu/moe_weight_prefetch.py + create mode 100644 python/sglang/srt/hardware_backend/npu/tp_ascendc_fusion_npu.py + create mode 100644 python/sglang/srt/layers/exp_race_fused_triton.py + create mode 100644 test/registered/unit/layers/test_conv_weight_cache_refresh.py + create mode 100644 test/registered/unit/layers/test_sampler_attach_logprobs.py + diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index bd9d7eafa1..1740283de2 100644 --- a/python/sglang/srt/entrypoints/http_server.py @@ -11,11 +56,65 @@ index bd9d7eafa1..1740283de2 100644 # uuid keeps rids unique across tokenizer workers (a bare time.time() can # collide and crash the shared DetokenizerManager decode_status). rid = f"{HEALTH_CHECK_RID_PREFIX}_{uuid.uuid4().hex}" +diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py +index 790acce902..b97ee18eeb 100644 +--- a/python/sglang/srt/environ.py ++++ b/python/sglang/srt/environ.py +@@ -526,7 +526,49 @@ class Envs: + # NPU + SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False) + SGLANG_NPU_USE_MULTI_STREAM = EnvBool(False) ++ # MoE expert weight L2 prefetch (decode aclgraph; see ++ # hardware_backend/npu/moe_weight_prefetch.py) ++ SGLANG_NPU_MOE_PREFETCH = EnvBool(False) ++ # Comma-separated op list: gmm1 (w13, default) / gmm2 (w2) ++ SGLANG_NPU_MOE_PREFETCH_OPS = EnvStr("gmm1") ++ # auto / full / active; auto == full, active falls back to full with warning ++ SGLANG_NPU_MOE_PREFETCH_MODE = EnvStr("auto") ++ SGLANG_NPU_MOE_PREFETCH_CHUNK_MIB = EnvInt(16) ++ # Per-tensor capacity cap in MiB; 0 = auto (0.8 * queried L2 size) ++ SGLANG_NPU_MOE_PREFETCH_BUDGET_MIB = EnvInt(0) ++ # full_attention decode fusion master switch (A1b split+KV scatter / ++ # A2 sigmoid_mul / A3a add_gemma_rms_norm, hardware_backend/npu/attention/ ++ # full_attention_fusion_npu.py); DEBUG=1 logs per-layer guard misses ++ SGLANG_NPU_FULL_ATTN_FUSION = EnvBool(True) ++ SGLANG_NPU_FULL_ATTN_FUSION_DEBUG = EnvBool(False) ++ # TP-line AscendC fusion master switch (only op1 fused_qkvzba_conv1d ++ # follows it); QKVZBA is op1's own switch (unset = follow master, explicit ++ # 0 disables op1 alone); GDN_RECURRENT_ASCENDC is op2 recurrent's switch ++ # (off by default, resolved at import time โ€” set before graph capture / ++ # server start); DEBUG=1 logs guard misses ++ SGLANG_NPU_TP_ASCENDC_FUSION = EnvBool(False) ++ SGLANG_NPU_TP_ASCENDC_FUSION_QKVZBA = EnvBool(False) ++ SGLANG_NPU_GDN_RECURRENT_ASCENDC = EnvBool(False) ++ SGLANG_NPU_TP_ASCENDC_FUSION_DEBUG = EnvBool(False) ++ # GDN decode recurrent PR#740 Triton variant switch ++ SGLANG_NPU_GDN_UPDATE_FUSED = EnvBool(False) ++ # GDN input-projection qkvz+ba weight pack into a single GEMM; MAX_M gates ++ # small M โ€” M above the threshold falls back to two GEMMs ++ SGLANG_NPU_GDN_QKVZBA_PACK = EnvBool(True) ++ SGLANG_NPU_GDN_QKVZBA_PACK_MAX_M = EnvInt(256) ++ # post_sample: async exponential-race sampling (910C only); ++ # EXP_RACE_TRITON enables its fused Triton consumer kernel ++ SGLANG_NPU_ASYNC_EXPONENTIAL = EnvBool(False) ++ SGLANG_NPU_EXP_RACE_TRITON = EnvBool(True) + SGLANG_NPU_USE_MLAPO = EnvBool(False) ++ # MoE ๅ‰ๆฎต่žๅˆๅŒ…๏ผˆmoe_front_fusion/v1๏ผŒๅˆๅ…ฅ่‡ช daikang ๅˆ†ๆ”ฏ๏ผ‰ๆ€ปๅผ€ๅ…ณ๏ผš ++ # renorm=1 ๅ•็ฎ—ๅญ่ทฏ็”ฑ + v2.2 ่‡ชๅ†™ init_routing๏ผŒไป…ไฝœ็”จไบŽๆ—  group/ๆ—  bias ++ # ็š„ BF16 ่ทฏๅพ„ใ€‚้ป˜่ฎคๅ…ณ๏ผŒๅผ€ๅฏ๏ผšSGLANG_MOE_FRONT_FUSION=1ใ€‚ ++ SGLANG_MOE_FRONT_FUSION = EnvBool(False) ++ # GMM2๏ผˆw2 down_proj๏ผ‰่ตฐ sgl_kernel_npu.moe.persistent_gmm ็š„ Triton ++ # persistent kernel๏ผˆ้ป˜่ฎค stock npu_grouped_matmul๏ผ‰ใ€‚ไป…ไฝœ็”จไบŽ BF16 ๆ— ้‡ๅŒ– ++ # ๆ—  bias ่ทฏๅพ„๏ผŒๅผ€ๅฏ๏ผšSGLANG_GMM2_TRITON=1ใ€‚ ++ SGLANG_GMM2_TRITON = EnvBool(False) + # Forward native implementation for activation gelu tanh for model Skywork-Reward-Gemma-2-27B-v0.2 + SGLANG_NPU_FORWARD_NATIVE_GELUTANH = EnvBool(False) + # Forward native implementation for gemma rms norm for model Skywork-Reward-Gemma-2-27B-v0.2 diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py -index 7f506e7f1a..204b9382e8 100644 +index 7f506e7f1a..dad52ee943 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py -@@ -5,11 +5,6 @@ from sgl_kernel_npu.fla.fused_gdn_gating import ( +@@ -5,12 +5,16 @@ from sgl_kernel_npu.fla.fused_gdn_gating import ( fused_gdn_gating_kernel_without_sigmoid, fused_gdn_gating_npu, ) @@ -23,11 +122,20 @@ index 7f506e7f1a..204b9382e8 100644 - causal_conv1d_fn_npu, - causal_conv1d_update_npu, - causal_conv1d_update_v2, --) ++from sgl_kernel_npu.fla.utils import ( ++ fused_qkvzba_split_reshape_cat_contiguous, + ) ++from sglang.srt.hardware_backend.npu.tp_ascendc_fusion_npu import ( ++ tp_debug_log, ++ tp_fused_qkvzba_conv1d_inputs_ok, ++ tp_fusion_qkvzba_enabled, ++ tp_op_available, ++) from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import ( AscendMambaAttnBackendBase, -@@ -26,8 +21,6 @@ from sglang.srt.model_executor.model_runner import ModelRunner + ) +@@ -26,8 +30,6 @@ from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput fused_gdn_gating = fused_gdn_gating_npu @@ -36,11 +144,24 @@ index 7f506e7f1a..204b9382e8 100644 class AscendGDNAttnBackend(AscendMambaAttnBackendBase): -@@ -109,6 +102,13 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): +@@ -100,6 +102,10 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): + def init_forward_metadata(self, forward_batch: ForwardBatch): + if forward_batch.forward_mode.is_draft_extend_v2(): + return ++ # In eager (non-graph) forward the base class builds int32 metadata, so ++ # no cast is needed for recurrent โ€” drop any graph-flow int32 shadow to ++ # avoid reusing a stale buffer. ++ self._decode_recurrent_shadow_i32 = None + super().init_forward_metadata(forward_batch) + self.prepare_gdn_inputs( + forward_batch.batch_size, +@@ -109,6 +115,15 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): self._prepare_mamba_track_metadata(forward_batch) self.graph_mode = False + def _get_conv_weights_t(self, layer: RadixLinearAttention) -> torch.Tensor: ++ # The cached copy is refreshed in place by wrap_conv1d_weight_loader ++ # on every (re)load of conv1d.weight; it stays valid until then. + w = getattr(layer, "_conv_weights_t", None) + if w is None: + w = layer.conv_weights.transpose(0, 1).contiguous() @@ -50,33 +171,133 @@ index 7f506e7f1a..204b9382e8 100644 def forward_decode( self, layer: RadixLinearAttention, -@@ -125,16 +125,17 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): +@@ -123,18 +138,94 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): + ssm_states = layer_cache.temporal + query_start_loc = self.forward_metadata.query_start_loc cache_indices = self.forward_metadata.mamba_cache_indices ++ # Recurrent side uses int32 shadow buffers (synced outside the decode ++ # graph by _capture/_replay_metadata, identical in content to the int64 ++ # primary buffers); when the shadow is None (eager/verify) the original ++ # metadata is used. The conv side keeps the int64 primary buffers. ++ recurrent_query_start_loc = query_start_loc ++ recurrent_cache_indices = cache_indices ++ if self._decode_recurrent_shadow_i32 is not None: ++ recurrent_query_start_loc, recurrent_cache_indices = ( ++ self._decode_recurrent_shadow_i32 ++ ) - assert isinstance(mixed_qkv, torch.Tensor) +- assert isinstance(mixed_qkv, torch.Tensor) - conv_states_tmp = conv_states.transpose(1, 2).clone() - mixed_qkv = causal_conv1d_update( -+ mixed_qkv = torch.ops.npu.causal_conv1d( - mixed_qkv, +- mixed_qkv, - conv_states_tmp, - layer.conv_weights, - layer.bias, - layer.activation, - conv_state_indices=cache_indices, -+ self._get_conv_weights_t(layer), -+ conv_states=conv_states, -+ bias=layer.bias, -+ query_start_loc=query_start_loc, -+ cache_indices=cache_indices, -+ activation_mode=1, -+ pad_slot_id=-1, -+ run_mode=1, - ) +- ) - conv_states[:] = conv_states_tmp.transpose(1, 2) ++ # When the qwen3_5.py guard hits, the raw projections arrive as a tuple ++ # (projected_states_qkvz, projected_states_ba) before split โ€” split + ++ # causal_conv1d(run_mode=1) are done by the single kernel ++ # torch.ops.npu.fused_qkvzba_conv1d (conv_states in-place semantics ++ # unchanged), and z flows back to the model side as (core_attn_out, z). ++ # If the runtime check fails (op unregistered / tensor attrs mismatch), ++ # fall back to the local split kernel + stock conv, still returning a ++ # tuple so the caller is unaffected (silent fallback; the branch is ++ # baked into the graph at capture time). ++ fused_z = None ++ if isinstance(mixed_qkv, tuple): ++ qkvz, mixed_ba = mixed_qkv ++ if tp_fusion_qkvzba_enabled() and tp_op_available( ++ "fused_qkvzba_conv1d" ++ ) and tp_fused_qkvzba_conv1d_inputs_ok( ++ qkvz, ++ mixed_ba, ++ layer.num_k_heads, ++ layer.num_v_heads, ++ layer.head_k_dim, ++ layer.head_v_dim, ++ ): ++ mixed_qkv, fused_z, b, a = torch.ops.npu.fused_qkvzba_conv1d( ++ qkvz, ++ self._get_conv_weights_t(layer), ++ conv_states, ++ mixed_ba, ++ layer.num_k_heads, ++ layer.num_v_heads, ++ layer.head_k_dim, ++ layer.head_v_dim, ++ bias=layer.bias, ++ query_start_loc=query_start_loc, ++ cache_indices=cache_indices, ++ activation_mode=1, ++ pad_slot_id=-1, ++ ) ++ else: ++ tp_debug_log( ++ ("qkvzba_rt", layer.layer_id), ++ f"layer {layer.layer_id}: fused_qkvzba_conv1d runtime check " ++ f"missed (enabled={tp_fusion_qkvzba_enabled()} " ++ f"op={tp_op_available('fused_qkvzba_conv1d')}), " ++ "falling back to local split + stock conv", ++ ) ++ mixed_qkv, fused_z, b, a = fused_qkvzba_split_reshape_cat_contiguous( ++ qkvz, ++ mixed_ba, ++ layer.num_k_heads, ++ layer.num_v_heads, ++ layer.head_k_dim, ++ layer.head_v_dim, ++ ) ++ mixed_qkv = torch.ops.npu.causal_conv1d( ++ mixed_qkv, ++ self._get_conv_weights_t(layer), ++ conv_states=conv_states, ++ bias=layer.bias, ++ query_start_loc=query_start_loc, ++ cache_indices=cache_indices, ++ activation_mode=1, ++ pad_slot_id=-1, ++ run_mode=1, ++ ) ++ else: ++ assert isinstance(mixed_qkv, torch.Tensor) ++ mixed_qkv = torch.ops.npu.causal_conv1d( ++ mixed_qkv, ++ self._get_conv_weights_t(layer), ++ conv_states=conv_states, ++ bias=layer.bias, ++ query_start_loc=query_start_loc, ++ cache_indices=cache_indices, ++ activation_mode=1, ++ pad_slot_id=-1, ++ run_mode=1, ++ ) query, key, value = torch.split( mixed_qkv, -@@ -219,44 +220,41 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): +@@ -155,13 +246,17 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): + A_log=layer.A_log, + dt_bias=layer.dt_bias, + ssm_states=ssm_states, +- cache_indices=cache_indices, +- query_start_loc=query_start_loc, ++ cache_indices=recurrent_cache_indices, ++ query_start_loc=recurrent_query_start_loc, + ) + + self._track_mamba_state_decode( + forward_batch, conv_states, ssm_states, cache_indices + ) ++ # Tuple-input path returns (core_attn_out, z) to match the qwen3_5.py ++ # unpacking convention; the tensor path return matches stock exactly. ++ if fused_z is not None: ++ return core_attn_out, fused_z + return core_attn_out + + def forward_extend( +@@ -219,44 +314,41 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): dtype=torch.int32, device=mixed_qkv.device, ) @@ -148,6 +369,601 @@ index 7f506e7f1a..204b9382e8 100644 if is_target_verify: g, beta = fused_gdn_gating_kernel_without_sigmoid( layer.A_log, a, b, layer.dt_bias +diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py +index 2a628c286d..cc3a655ad9 100644 +--- a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py ++++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py +@@ -27,16 +27,30 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase): + def __init__(self, model_runner: ModelRunner): + super().__init__(model_runner) + self.state_indices_list_gdn = [] ++ # decode int32 shadow buffers: the AscendC recurrent host reads ++ # cache_indices/cu_seqlens as int32_t*; synced once per step outside ++ # the decode graph (zero cast inside the graph); always None on the ++ # eager/verify paths. ++ self.decode_state_indices_i32_list = [] ++ self.decode_query_start_loc_i32_list = [] ++ self.cached_cuda_graph_decode_query_start_loc_i32 = None ++ self._decode_recurrent_shadow_i32 = None + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + assert ( + max_num_tokens % max_bs == 0 + ), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}" + draft_token_num = max_num_tokens // max_bs ++ # state_indices / query_start_loc primary buffers use int64: the conv ++ # AscendC ops (causal_conv1d / fused_qkvzba_conv1d) read them as ++ # int64_t* on host, so feeding directly avoids an in-graph cast. ++ # Exception: state_indices_list_gdn stays int32 (recurrent_gated_delta_rule ++ # reads it as int32_t* without validation โ€” changing it silently ++ # corrupts results). + for i in range(max_bs): + self.state_indices_list.append( + torch.full( +- (i + 1,), self.pad_slot_id, dtype=torch.int32, device=self.device ++ (i + 1,), self.pad_slot_id, dtype=torch.int64, device=self.device + ) + ) + self.state_indices_list_gdn.append( +@@ -48,7 +62,7 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase): + ) + ) + self.query_start_loc_list.append( +- torch.zeros((i + 2,), dtype=torch.int32, device=self.device) ++ torch.zeros((i + 2,), dtype=torch.int64, device=self.device) + ) + self.retrieve_next_token_list.append( + torch.zeros( +@@ -65,17 +79,54 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase): + (i + 1, draft_token_num), dtype=torch.int32, device=self.device + ) + ) ++ # decode shadow buffers (recurrent side), synced outside the graph ++ # by _capture/_replay_metadata each step; not used in verify. ++ self.decode_state_indices_i32_list.append( ++ torch.full( ++ (i + 1,), self.pad_slot_id, dtype=torch.int32, device=self.device ++ ) ++ ) ++ self.decode_query_start_loc_i32_list.append( ++ torch.zeros((i + 2,), dtype=torch.int32, device=self.device) ++ ) + self.cached_cuda_graph_decode_query_start_loc = torch.arange( ++ 0, max_bs + 1, dtype=torch.int64, device=self.device ++ ) ++ # decode query_start_loc is a static arange โ€” the int32 copy is built ++ # once at init, so shadow sync is just a same-dtype copy. ++ self.cached_cuda_graph_decode_query_start_loc_i32 = torch.arange( + 0, max_bs + 1, dtype=torch.int32, device=self.device + ) + self.cached_cuda_graph_verify_query_start_loc = torch.arange( + 0, + max_bs * draft_token_num + 1, + step=draft_token_num, +- dtype=torch.int32, ++ dtype=torch.int64, + device=self.device, + ) + ++ def _update_decode_recurrent_shadow_i32( ++ self, bs: int, mamba_indices: torch.Tensor, num_padding: int = 0 ++ ): ++ """Sync the decode int32 shadow buffers (outside the graph, once per step). ++ ++ Content is identical to the int64 primary buffers, for direct reads by ++ the AscendC recurrent (host reads cache_indices/cu_seqlens as int32_t*). ++ Called only for decode/idle; on the verify/eager paths the caller sets ++ _decode_recurrent_shadow_i32 to None. ++ """ ++ idx_i32 = self.decode_state_indices_i32_list[bs - 1] ++ qsl_i32 = self.decode_query_start_loc_i32_list[bs - 1] ++ idx_i32[: len(mamba_indices)].copy_(mamba_indices) ++ if num_padding == 0: ++ qsl_i32.copy_(self.cached_cuda_graph_decode_query_start_loc_i32[: bs + 1]) ++ else: ++ qsl_i32[: bs - num_padding].copy_( ++ self.cached_cuda_graph_decode_query_start_loc_i32[: bs - num_padding] ++ ) ++ qsl_i32[bs - num_padding :].fill_(bs - num_padding) ++ self._decode_recurrent_shadow_i32 = (qsl_i32, idx_i32) ++ + def _capture_metadata( + self, + bs: int, +@@ -89,7 +140,9 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase): + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_decode_query_start_loc[: bs + 1] + ) ++ self._update_decode_recurrent_shadow_i32(bs, mamba_indices) + elif forward_mode.is_target_verify(): ++ self._decode_recurrent_shadow_i32 = None + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_verify_query_start_loc[: bs + 1] + ) +@@ -163,7 +216,9 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase): + self.query_start_loc_list[bs - 1][bs - num_padding :].fill_( + bs - num_padding + ) ++ self._update_decode_recurrent_shadow_i32(bs, mamba_indices, num_padding) + elif forward_mode.is_target_verify(): ++ self._decode_recurrent_shadow_i32 = None + ssm_state_indices = torch.arange( + bs * spec_info.draft_token_num, + dtype=torch.int32, +diff --git a/python/sglang/srt/hardware_backend/npu/attention/full_attention_fusion_npu.py b/python/sglang/srt/hardware_backend/npu/attention/full_attention_fusion_npu.py +new file mode 100644 +index 0000000000..4751217a5e +--- /dev/null ++++ b/python/sglang/srt/hardware_backend/npu/attention/full_attention_fusion_npu.py +@@ -0,0 +1,468 @@ ++# Qwen3.5-35B-A3B NPU decode ๅ…จๆณจๆ„ๅŠ›ๆฎต่žๅˆ kernel๏ผˆfull_attention ไผ˜ๅŒ–็‚น A1b/A2/A3a๏ผ‰ใ€‚ ++# ++# A1b split_qkvgate_v4_scatter๏ผšsplit q/gate/k/v + q/k gemma rmsnorm + rope ++# single kernel, and k/v are scattered by loc directly into the KV ++# cache; returned k/v are None (cache already written โ€” the caller must ++# skip set_kv_buffer). ++# A2 fa_sigmoid_mul๏ผšsigmoid(gate)*attn ่žๅˆๅ• kernel๏ผˆๆ›ฟไปฃ sigmoid+mul_ ไธค ++# kernel๏ผ‰๏ผ›round to bf16 before the multiply, bitwise-identical to ++# stock's two-op rounding path. ++# A3a fa_add_gemma_rms_norm_v2๏ผšgrid=(batch,) row-parallel; math identical ++# to stock add_gemma_rms_norm. ++# ++# ๆ•ฐๅ€ผ็บฆๅฎšไธŽ sgl_kernel_npu stock ้€ๅญ—ไธ€่‡ด๏ผšfp32 ไธญ้—ด็ฒพๅบฆใ€gemma (w+1)ใ€ ++# neox rotate_half rope๏ผˆcat=[-x2,x1]; out=cat*sin+rot*cos๏ผ‰ใ€‚ ++# ++# ๅฎ‰ๅ…จ็ฝ‘๏ผšๆฏไธช่ฐƒ็”จ็‚นๅ…ˆ่ฟ‡ๅฎˆๅซ๏ผˆ็ฒพ็กฎๅ‘ฝไธญๅทฒ้ชŒ่ฏๅฝข็Šถๆ‰่ตฐ่‡ช็ ” kernel๏ผŒๅฆๅˆ™ๅ›ž้€€ ++# stock๏ผ‰๏ผ›็Žฏๅขƒๅ˜้‡ SGLANG_NPU_FULL_ATTN_FUSION=0 ๅฏๆ•ดไฝ“ๅ…ณๅœ๏ผˆๅ›พๆจกๅผไธ‹ๅฎˆๅซ ++# ๅˆคๅฎšๅœจ capture ๆ—ถ็ƒ˜่ฟ›ๅ›พ๏ผŒ้œ€ๅœจๅปบๅ›พๅ‰่ฎพ็ฝฎ๏ผ‰ใ€‚ ++# ++# ๅฎšไฝๆ—ฅๅฟ—๏ผšSGLANG_NPU_FULL_ATTN_FUSION_DEBUG=1 ๆ—ถ๏ผŒA1b ๅฎˆๅซๆŒ‰ๅฑ‚ไธ€ๆฌกๆ€งๆ‰“ๅฐ ++# "guard miss reason" or "activated" confirmation. ++# ++# โš ๏ธ triton-ascend constraints: no runtime dispatch on program_id, no ++# scalar-conditional mask store, no 2D masked store (miscompiles). All kernels ++# in this file use only: ++# ็บฏ็บฟๆ€ง็ดขๅผ• + ๆฏ program ๅ›บๅฎš่Œ่ดฃ + 1D ๆ— ๆŽฉ็  store + constexpr ifใ€‚ ++ ++import logging ++import os ++ ++import torch ++import triton ++import triton.language as tl ++import triton.language.extra.cann.extension as al ++ ++logger = logging.getLogger(__name__) ++ ++_FA_FUSION_ENV = "SGLANG_NPU_FULL_ATTN_FUSION" ++_FA_DEBUG_ENV = "SGLANG_NPU_FULL_ATTN_FUSION_DEBUG" ++_fa_debug_logged = set() ++ ++ ++def fa_fusion_enabled() -> bool: ++ """full_attention ่žๅˆๆ€ปๅผ€ๅ…ณ๏ผˆ้ป˜่ฎคๅผ€๏ผ›=0 ๆ—ถๅ…จ้ƒจๅ›ž้€€ stock๏ผ‰ใ€‚""" ++ return os.environ.get(_FA_FUSION_ENV, "1") == "1" ++ ++ ++def _fa_debug_log(key, msg): ++ """debug ๅผ€ๅ…ณๆ‰“ๅผ€ๆ—ถๆŒ‰ key ไธ€ๆฌกๆ€งๆ‰“ๆ—ฅๅฟ—๏ผˆๅฎšไฝๅฎˆๅซๆœชๅ‘ฝไธญ็Žฏ่Š‚็”จ๏ผ‰ใ€‚""" ++ if os.environ.get(_FA_DEBUG_ENV, "0") != "1": ++ return ++ if key in _fa_debug_logged: ++ return ++ _fa_debug_logged.add(key) ++ logger.warning("[full_attention_fusion] %s", msg) ++ ++ ++# --------------------------------------------------------------------------- ++# A3a๏ผšadd_gemma_rms_norm ่กŒๅนถ่กŒ็‰ˆ๏ผˆgrid=(batch,)๏ผŒๆฏ program ไธ€ๆ•ด่กŒ๏ผ‰ ++# ๆ•ฐๅญฆไธŽ stock add_gemma_rms_norm ๅฎŒๅ…จไธ€่‡ด๏ผˆbf16 ๅŠ ๆฎ‹ๅทฎโ†’fp32 ๅฝ’็บฆโ†’fp32 (w+1)๏ผ‰ใ€‚ ++# --------------------------------------------------------------------------- ++@triton.jit ++def add_gemma_rms_norm_v2_kernel( ++ x_ptr, ++ res_ptr, ++ w_ptr, ++ add_out_ptr, ++ norm_out_ptr, ++ eps, ++ DIM: tl.constexpr, ++): ++ row = tl.program_id(0) ++ cols = tl.arange(0, DIM) ++ x = tl.load(x_ptr + row * DIM + cols) ++ r = tl.load(res_ptr + row * DIM + cols) ++ add = x + r ++ tl.store(add_out_ptr + row * DIM + cols, add) ++ xf = add.to(tl.float32) ++ w = tl.load(w_ptr + cols).to(tl.float32) + 1.0 ++ var = tl.sum(xf * xf) / DIM ++ out = xf * tl.rsqrt(var + eps) * w ++ tl.store(norm_out_ptr + row * DIM + cols, out.to(x_ptr.dtype.element_ty)) ++ ++ ++def fa_add_gemma_rms_norm_v2_supported(x, residual) -> bool: ++ """A3a shape guard: 2D, bf16, contiguous, hidden a power of 2 ++ (tl.arange is invalid for non-power-of-2 hidden).""" ++ if not fa_fusion_enabled(): ++ return False ++ if residual is None or x.dim() != 2: ++ return False ++ if x.dtype != torch.bfloat16 or residual.dtype != torch.bfloat16: ++ return False ++ if not (x.is_contiguous() and residual.is_contiguous()): ++ return False ++ dim = x.shape[-1] ++ return dim > 0 and (dim & (dim - 1)) == 0 ++ ++ ++def fa_add_gemma_rms_norm_v2(x, weight, residual, eps): ++ """่ฟ”ๅ›ž (norm_out, add_out)๏ผŒ่ฏญไน‰ๅŒ sgl_kernel_npu add_gemma_rms_normใ€‚ ++ ่ฐƒ็”จๅ‰้กป่ฟ‡ fa_add_gemma_rms_norm_v2_supportedใ€‚""" ++ add_out = torch.empty_like(x) ++ norm_out = torch.empty_like(x) ++ B, D = x.shape ++ add_gemma_rms_norm_v2_kernel[(B,)]( ++ x, residual, weight, add_out, norm_out, eps, DIM=D ++ ) ++ return norm_out, add_out ++ ++ ++# --------------------------------------------------------------------------- ++# A2๏ผšsigmoid(gate) * attn ่žๅˆ๏ผˆๆ›ฟไปฃ torch.sigmoid + mul_ ไธคไธช kernel๏ผ‰ ++# --------------------------------------------------------------------------- ++@triton.jit ++def sigmoid_mul_kernel( ++ attn_ptr, ++ gate_ptr, ++ out_ptr, ++ numel, ++ BLOCK: tl.constexpr, ++ BITWISE: tl.constexpr, ++): ++ pid = tl.program_id(0) ++ offs = pid * BLOCK + tl.arange(0, BLOCK) ++ mask = offs < numel ++ a = tl.load(attn_ptr + offs, mask=mask).to(tl.float32) ++ g = tl.load(gate_ptr + offs, mask=mask).to(tl.float32) ++ s = tl.sigmoid(g) ++ if BITWISE: ++ # ๆจกๆ‹Ÿ stock ไธคไธช op ไน‹้—ด็š„ bf16 ไธญ้—ด่ˆๅ…ฅ๏ผšsigmoid ๅ…ˆ่ˆๅˆฐ bf16 ๅ†ไน˜๏ผŒ ++ # bitwise-identical to the attn.mul_(torch.sigmoid(gate)) target. ++ s = s.to(tl.bfloat16).to(tl.float32) ++ out = (a * s).to(tl.bfloat16) ++ tl.store(out_ptr + offs, out, mask=mask) ++ ++ ++def fa_sigmoid_mul_supported(attn_output, gate) -> bool: ++ """A2 ๅฝข็Šถๅฎˆๅซ๏ผš็ญ‰ๅฝขใ€bf16ใ€่ฟž็ปญ๏ผˆelementwise ่žๅˆ๏ผŒๆ— ๅฝข็Šถ็‰นๅŒ–๏ผ‰ใ€‚""" ++ if not fa_fusion_enabled(): ++ return False ++ if attn_output.dtype != torch.bfloat16 or gate.dtype != torch.bfloat16: ++ return False ++ if attn_output.shape != gate.shape: ++ return False ++ return attn_output.is_contiguous() and gate.is_contiguous() ++ ++ ++def fa_sigmoid_mul(attn, gate): ++ """out = attn * sigmoid(gate)๏ผŒไธŽ stock ๅŒ op ่ˆๅ…ฅ่ทฏๅพ„้€ไฝไธ€่‡ดใ€‚ ++ ่ฐƒ็”จๅ‰้กป่ฟ‡ fa_sigmoid_mul_supportedใ€‚""" ++ out = torch.empty_like(attn) ++ numel = attn.numel() ++ BLOCK = 512 ++ sigmoid_mul_kernel[(triton.cdiv(numel, BLOCK),)]( ++ attn, gate, out, numel, BLOCK=BLOCK, BITWISE=True ++ ) ++ return out ++ ++ ++# --------------------------------------------------------------------------- ++# A1b (row-parallel straight-line version): split q/kv/gate + q/k gemma ++# rmsnorm + rope, ++# ไธ” k/v ๆŒ‰ loc ็›ดๆŽฅ scatter ่ฟ› KV cache๏ผˆๅ…ไธญ้—ดๅผ ้‡ไธŽไธคๆฌก scatter_nd_update๏ผ‰ใ€‚ ++# grid=(batch,)๏ผšไธ€ program ไธ€่กŒ๏ผŒq+gate / k / v ไธ‰ๆฎตๅœจไธ€ๆก็›ด็บฟ้‡Œ้กบๅบๅšๅฎŒใ€‚ ++# Design constraints (triton-ascend): pid is used only for linear indexing ++# (no div/mod), no inter-segment runtime ++# ้€‰ๆ‹ฉใ€store ไธๅธฆๆ ‡้‡ๆกไปถ mask๏ผ›SCATTER_KV ไธบ constexpr ๅˆ†ๆ”ฏ๏ผŒๆœฌๆ–‡ไปถ็”Ÿไบง ++# wrapper ๆ’ =1๏ผŒk/v ไธญ้—ดๅผ ้‡ๅ†™ๅ‡บๆ•ดๆฎต็ผ–่ฏ‘ๆœŸๆถˆ้™คใ€‚ ++# TP4/TP8/TP16 shape specialization: NUM_Q_HEADS โˆˆ {1,2,4} (TP16=1, TP8=2, ++# 9B-TP4=4, passed per runtime shape), NUM_KV_HEADS=1. ++# --------------------------------------------------------------------------- ++@triton.jit ++def split_qkvgate_v4_kernel( ++ input_ptr, ++ sin_ptr, ++ cos_ptr, ++ q_ptr, ++ gate_ptr, ++ k_ptr, ++ v_ptr, ++ kbuf_ptr, ++ vbuf_ptr, ++ loc_ptr, ++ q_weight_ptr, ++ k_weight_ptr, ++ q_hidden_size: tl.constexpr, ++ kv_hidden_size: tl.constexpr, ++ total_hidden_size: tl.constexpr, ++ eps: tl.constexpr, ++ NUM_Q_HEADS: tl.constexpr, # โˆˆ{1,2,4} ++ HEAD_DIM: tl.constexpr, ++ ROPE_DIM: tl.constexpr, ++ HALF_ROPE_DIM: tl.constexpr, ++ SCATTER_KV: tl.constexpr, ++): ++ row = tl.program_id(0) ++ base = input_ptr + row * total_hidden_size ++ d = tl.arange(0, HEAD_DIM) ++ ++ # ---- q+gate๏ผšqkv ่กŒๅ‰ 2*q_hidden_size ่ฟž็ปญ่ฏปๅ…ฅ๏ผŒreshape ๆˆ (ๅคดๆ•ฐ, 512) ๅŽ ++ # ้€ๅคดๅˆ‡ๅ‡บ [q(256) | gate(256)]๏ผˆไธŽ stock ็š„ extract_slice ๅˆ‡ๆณ•ๅฎŒๅ…จไธ€่‡ด๏ผ‰ ++ qg_cols = tl.arange(0, 2 * NUM_Q_HEADS * HEAD_DIM) ++ qg = tl.load(base + qg_cols).to(tl.float32).reshape( ++ NUM_Q_HEADS, 2 * HEAD_DIM ++ ) ++ q = al.extract_slice( ++ qg, offsets=(0, 0), sizes=(NUM_Q_HEADS, HEAD_DIM), strides=(1, 1) ++ ) ++ gate = al.extract_slice( ++ qg, offsets=(0, HEAD_DIM), sizes=(NUM_Q_HEADS, HEAD_DIM), strides=(1, 1) ++ ) ++ ++ # q norm๏ผˆfp32 ไธญ้—ดใ€gemma w+1๏ผ›่ฟ็ฎ—้กบๅบไธŽ stock ้€ๅญ—ไธ€่‡ด๏ผ‰ ++ qw = tl.load(q_weight_ptr + d).to(tl.float32) + 1.0 ++ qvar = tl.sum(q * q, axis=1) / HEAD_DIM ++ qn = q * tl.rsqrt(qvar + eps)[:, None] * qw[None, :] ++ ++ # rope (neox rotate_half, first ROPE_DIM dims only; op order matches stock ++ # exactly) ++ sc = row * ROPE_DIM + tl.arange(0, ROPE_DIM) ++ sin = tl.load(sin_ptr + sc).to(tl.float32).reshape(1, ROPE_DIM) ++ cos = tl.load(cos_ptr + sc).to(tl.float32).reshape(1, ROPE_DIM) ++ rot = al.extract_slice( ++ qn, offsets=(0, 0), sizes=(NUM_Q_HEADS, ROPE_DIM), strides=(1, 1) ++ ) ++ x1 = al.extract_slice( ++ rot, offsets=(0, 0), sizes=(NUM_Q_HEADS, HALF_ROPE_DIM), strides=(1, 1) ++ ) ++ x2 = al.extract_slice( ++ rot, offsets=(0, HALF_ROPE_DIM), sizes=(NUM_Q_HEADS, HALF_ROPE_DIM), strides=(1, 1) ++ ) ++ cat = tl.zeros((NUM_Q_HEADS, ROPE_DIM), dtype=tl.float32) ++ cat = al.insert_slice( ++ cat, -x2, offsets=(0, 0), sizes=(NUM_Q_HEADS, HALF_ROPE_DIM), strides=(1, 1) ++ ) ++ cat = al.insert_slice( ++ cat, x1, offsets=(0, HALF_ROPE_DIM), sizes=(NUM_Q_HEADS, HALF_ROPE_DIM), strides=(1, 1) ++ ) ++ roped = cat * sin + rot * cos ++ qn = al.insert_slice( ++ qn, roped, offsets=(0, 0), sizes=(NUM_Q_HEADS, ROPE_DIM), strides=(1, 1) ++ ) ++ ++ # q/gate store: 1D unmasked store (same as stock reshape(Q_BLOCK); avoids ++ # 2D masked store โ€” see the triton-ascend constraints at the top) ++ qg_out = row * q_hidden_size + tl.arange(0, NUM_Q_HEADS * HEAD_DIM) ++ tl.store( ++ q_ptr + qg_out, ++ qn.reshape(NUM_Q_HEADS * HEAD_DIM).to(input_ptr.dtype.element_ty), ++ ) ++ tl.store( ++ gate_ptr + qg_out, ++ gate.reshape(NUM_Q_HEADS * HEAD_DIM).to(input_ptr.dtype.element_ty), ++ ) ++ ++ # ---- k๏ผˆ1 ๅคด๏ผ‰๏ผšnorm + rope๏ผŒไธŽ stock k ๆฎต็š„ (1, HEAD_DIM) ๅฝขๆ€ไธ€่‡ด ++ kt = tl.load(base + 2 * q_hidden_size + d).to(tl.float32).reshape(1, HEAD_DIM) ++ kw = tl.load(k_weight_ptr + d).to(tl.float32) + 1.0 ++ kvar = tl.sum(kt * kt, axis=1) / HEAD_DIM ++ kn = kt * tl.rsqrt(kvar + eps)[:, None] * kw[None, :] ++ krot = al.extract_slice(kn, offsets=(0, 0), sizes=(1, ROPE_DIM), strides=(1, 1)) ++ kx1 = al.extract_slice( ++ krot, offsets=(0, 0), sizes=(1, HALF_ROPE_DIM), strides=(1, 1) ++ ) ++ kx2 = al.extract_slice( ++ krot, offsets=(0, HALF_ROPE_DIM), sizes=(1, HALF_ROPE_DIM), strides=(1, 1) ++ ) ++ kcat = tl.zeros((1, ROPE_DIM), dtype=tl.float32) ++ kcat = al.insert_slice( ++ kcat, -kx2, offsets=(0, 0), sizes=(1, HALF_ROPE_DIM), strides=(1, 1) ++ ) ++ kcat = al.insert_slice( ++ kcat, kx1, offsets=(0, HALF_ROPE_DIM), sizes=(1, HALF_ROPE_DIM), strides=(1, 1) ++ ) ++ kroped = kcat * sin + krot * cos ++ kn = al.insert_slice(kn, kroped, offsets=(0, 0), sizes=(1, ROPE_DIM), strides=(1, 1)) ++ ++ # ---- v๏ผšbf16 ๅŽŸๆ ทๆ‹ท่ด๏ผˆstock v ๆฎตไธ่ฟ‡ fp32๏ผ‰ ++ vt = tl.load(base + 2 * q_hidden_size + kv_hidden_size + d) ++ ++ # constexpr branch: with SCATTER_KV=1 the k/v intermediate-tensor stores ++ # are eliminated at compile time, and the k_ptr/v_ptr placeholder pointers ++ # are never dereferenced ++ if SCATTER_KV: ++ slot = tl.load(loc_ptr + row).to(tl.int64) ++ tl.store( ++ kbuf_ptr + slot * kv_hidden_size + d, ++ kn.reshape(HEAD_DIM).to(input_ptr.dtype.element_ty), ++ ) ++ tl.store(vbuf_ptr + slot * kv_hidden_size + d, vt) ++ else: ++ tl.store( ++ k_ptr + row * kv_hidden_size + d, ++ kn.reshape(HEAD_DIM).to(input_ptr.dtype.element_ty), ++ ) ++ tl.store(v_ptr + row * kv_hidden_size + d, vt) ++ ++ ++def fa_split_qkvgate_scatter_supported( ++ num_heads, num_kv_heads, head_dim, rope_dim, attn_output_gate ++) -> bool: ++ """A1b shape guard (TP layout specialization): per-rank q โˆˆ {1,2,4} heads ++ (TP16=1/TP8=2/9B-TP4=4), kv=1 head, head_dim=256, rope=64, with ++ attn_output_gate. All other TP/model configs fall back to stock. ++ Note: after enabling q=4 (9B TP4 form), run the in-graph UT ++ bench_fa_graph.py --num-q 4 on the server before production use.""" ++ ok = ( ++ fa_fusion_enabled() ++ and num_heads in (1, 2, 4) ++ and num_kv_heads == 1 ++ and head_dim == 256 ++ and rope_dim == 64 ++ and attn_output_gate ++ ) ++ if not ok: ++ _fa_debug_log( ++ ("shape_guard",), ++ "A1b ๅฝข็Šถๅฎˆๅซๆœชๅ‘ฝไธญ๏ผˆๅ›ž้€€ stock๏ผ‰๏ผš" ++ f"num_heads={num_heads}(need 1/2/4) num_kv_heads={num_kv_heads}(need 1) " ++ f"head_dim={head_dim}(้œ€256) rope_dim={rope_dim}(้œ€64) " ++ f"gate={attn_output_gate}(้œ€True) " ++ f"{_FA_FUSION_ENV}={os.environ.get(_FA_FUSION_ENV, '1')!r}", ++ ) ++ return ok ++ ++ ++def fa_kv_pool_buffers(layer_id): ++ """ไปŽๅฝ“ๅ‰ NPU attention backend ็š„ KV pool ๅ–่ฏฅๅฑ‚ k/v bufferใ€‚ ++ ไป…ๅœจ use_fia ๅธƒๅฑ€๏ผˆ[pages*page_size, 1, head_num, head_dim]๏ผŒไธŽ token ็บง ++ slot=loc ็›ดๅ†™ๅฃๅพ„ไธ€่‡ด๏ผ‰ไธ”้ž hybrid SWA ๆ—ถ่ฟ”ๅ›ž (kbuf, vbuf)๏ผŒๅฆๅˆ™่ฟ”ๅ›ž Noneใ€‚ ++ ++ GDN ๆททๅˆๆจกๅž‹็š„ pool ๆ˜ฏ HybridLinearKVPool ๅŒ…่ฃ…๏ผšuse_fia ๅˆคๅฎš่ฆ่ฏปๅ†…ๅฑ‚ ++ full_kv_pool๏ผˆNPUMHATokenToKVPool๏ผŒๅค–ๅฑ‚ๆ— ๆญคๅฑžๆ€ง๏ผ‰๏ผ›ๅ– buffer ่ตฐๅค–ๅฑ‚ ++ get_key/value_bufferโ€”โ€”ๅค–ๅฑ‚ๆŒ‰ full_attention_layer_id_mapping ๆŠŠๅ…จๅฑ€ ++ layer_id ็ฟป่ฏ‘ๆˆๅ…จๆณจๆ„ๅŠ›ๅฑ€้ƒจๅบๅทๅŽๅ† delegate๏ผˆๅ†…ๅฑ‚ pool ๅชๆœ‰ๅ…จๆณจๆ„ๅŠ›ๅฑ‚๏ผŒ ++ ็›ดๆŽฅ็”จๅ…จๅฑ€ id ไผšๅ–้”™ๅฑ‚๏ผ‰ใ€‚ ++ """ ++ from sglang.srt.model_executor.forward_context import get_attn_backend ++ ++ backend = get_attn_backend() ++ if getattr(backend, "is_hybrid_swa", False): ++ _fa_debug_log( ++ ("pool_swa", layer_id), ++ f"layer {layer_id}: backend is_hybrid_swa=True๏ผŒA1b ๅ›ž้€€ stock", ++ ) ++ return None ++ pool = getattr(backend, "token_to_kv_pool", None) ++ if pool is None: ++ _fa_debug_log( ++ ("pool_none", layer_id), ++ f"layer {layer_id}: backend {type(backend).__name__} ๆ—  " ++ "token_to_kv_pool๏ผŒA1b ๅ›ž้€€ stock", ++ ) ++ return None ++ inner = getattr(pool, "full_kv_pool", pool) ++ if not getattr(inner, "use_fia", False): ++ _fa_debug_log( ++ ("pool_fia", layer_id), ++ f"layer {layer_id}: pool {type(pool).__name__} ๅ†…ๅฑ‚ " ++ f"{type(inner).__name__}.use_fia=False๏ผˆASCEND_USE_FIA ๆœชๅผ€๏ผŸ๏ผ‰๏ผŒ" ++ "A1b ๅ›ž้€€ stock", ++ ) ++ return None ++ return pool.get_key_buffer(layer_id), pool.get_value_buffer(layer_id) ++ ++ ++def fa_v4_scatter_context(layer, forward_batch, qkv): ++ """A1b ่ฟ่กŒๆ—ถๅฎˆๅซๆ€ปๅ…ฅๅฃ๏ผšๅ…จ้ƒจๆกไปถๅ‘ฝไธญ่ฟ”ๅ›ž (kbuf, vbuf, loc)๏ผŒไปปไธ€ๆœชๅ‘ฝไธญ ++ ่ฟ”ๅ›ž None๏ผˆ่ฐƒ็”จๆ–นๅ›ž้€€ stock split + stock set_kv_buffer๏ผŒ่กŒไธบไธŽ origin ++ ๅฎŒๅ…จไธ€่‡ด๏ผ‰ใ€‚DEBUG=1 ๆ—ถๆŒ‰ๅฑ‚ไธ€ๆฌกๆ€งๆ‰“ๅฐๆœชๅ‘ฝไธญๅŽŸๅ› /ๆฟ€ๆดป็กฎ่ฎคใ€‚""" ++ layer_id = layer.attn.layer_id ++ if not getattr(layer, "_fa_v4_shape_ok", False): ++ _fa_debug_log( ++ ("shape", layer_id), ++ f"layer {layer_id}: _fa_v4_shape_ok=False๏ผˆ่ง init ๆ—ถๅฝข็Šถๅฎˆๅซๆ—ฅๅฟ—๏ผ‰๏ผŒ" ++ "A1b ๅ›ž้€€ stock", ++ ) ++ return None ++ if not forward_batch.forward_mode.is_decode(): ++ _fa_debug_log( ++ ("mode", layer_id), ++ f"layer {layer_id}: forward_mode={forward_batch.forward_mode} ้ž " ++ "DECODE๏ผˆMTP verify / mixed / prefill ๅ‡ไธ่ตฐ A1b๏ผ‰๏ผŒๅ›ž้€€ stock", ++ ) ++ return None ++ if not torch.is_tensor(qkv) or qkv.dtype != torch.bfloat16 or not qkv.is_contiguous(): ++ _fa_debug_log( ++ ("qkv", layer_id), ++ f"layer {layer_id}: qkv ไธๆปก่ถณ๏ผˆtensor={torch.is_tensor(qkv)} " ++ f"dtype={getattr(qkv, 'dtype', None)} " ++ f"contiguous={qkv.is_contiguous() if torch.is_tensor(qkv) else None}๏ผ‰๏ผŒ" ++ "A1b ๅ›ž้€€ stock", ++ ) ++ return None ++ kv_bufs = fa_kv_pool_buffers(layer_id) ++ if kv_bufs is None: ++ return None # ๅ…ทไฝ“ๅŽŸๅ› ๅทฒๅœจ fa_kv_pool_buffers ๅ†…ๆ‰“ๆ—ฅๅฟ— ++ kbuf, vbuf = kv_bufs ++ if kbuf.dtype != qkv.dtype: ++ _fa_debug_log( ++ ("kvdtype", layer_id), ++ f"layer {layer_id}: KV cache dtype={kbuf.dtype} != qkv " ++ f"dtype={qkv.dtype}๏ผˆๅฆ‚ fp8 KV cache๏ผ‰๏ผŒA1b ๅ›ž้€€ stock", ++ ) ++ return None ++ _fa_debug_log( ++ ("active", layer_id), ++ f"layer {layer_id}: A1b ่žๅˆ split+KV scatter ๅทฒๆฟ€ๆดป", ++ ) ++ return kbuf, vbuf, forward_batch.out_cache_loc ++ ++ ++def fa_split_qkvgate_scatter( ++ input, ++ sin, ++ cos, ++ q_hidden_size, ++ kv_hidden_size, ++ head_dim, ++ rope_dim, ++ eps, ++ q_weight, ++ k_weight, ++ kbuf, ++ vbuf, ++ loc, ++): ++ """stock split_qkvgate_gemma_rmsnorm_rope ็š„่กŒๅนถ่กŒ็›ด็บฟ็‰ˆ๏ผˆgrid=(batch,)๏ผ‰๏ผŒ ++ k/v ไธๅ†™ไธญ้—ดๅผ ้‡ใ€ๆŒ‰ loc ็›ดๆŽฅ scatter ่ฟ› KV cacheใ€‚ ++ ่ฟ”ๅ›ž (q, None, None, gate)๏ผšk/v ไธบ None ่กจ็คบ cache ๅทฒๅ†™๏ผŒ ++ ่ฐƒ็”จๆ–น้กปไปฅ save_kv_cache=False ่ตฐ attention backendใ€‚ ++ Only q โˆˆ {1,2,4} heads and kv=1 head are supported (TP16=1, TP8=2, ++ 9B-TP4=4 head specialization); the fa_v4_scatter_context guard must pass ++ before calling. ++ fa_v4_scatter_context ๅฎˆๅซใ€‚""" ++ batch = input.shape[0] ++ num_q = q_hidden_size // head_dim ++ num_kv = kv_hidden_size // head_dim ++ assert num_q in (1, 2, 4) and num_kv == 1, "TP4/TP8/TP16 shape specialization (q โˆˆ {1,2,4} heads, kv=1 head)" ++ assert kbuf.dtype == input.dtype and vbuf.dtype == input.dtype ++ q_out = torch.empty(batch, q_hidden_size, device=input.device, dtype=input.dtype) ++ gate_out = torch.empty(batch, q_hidden_size, device=input.device, dtype=input.dtype) ++ # SCATTER_KV=1๏ผšk_ptr/v_ptr ไธบๅ ไฝๆŒ‡้’ˆ๏ผˆq_out/gate_out๏ผ‰๏ผŒ็ผ–่ฏ‘ๆœŸๆถˆ้™คไธไผš่ขซๅ†™ ++ split_qkvgate_v4_kernel[(batch,)]( ++ input, ++ sin, ++ cos, ++ q_out, ++ gate_out, ++ q_out, ++ gate_out, ++ kbuf, ++ vbuf, ++ loc, ++ q_weight, ++ k_weight, ++ q_hidden_size, ++ kv_hidden_size, ++ q_hidden_size * 2 + kv_hidden_size * 2, ++ eps, ++ NUM_Q_HEADS=num_q, ++ HEAD_DIM=head_dim, ++ ROPE_DIM=rope_dim, ++ HALF_ROPE_DIM=rope_dim // 2, ++ SCATTER_KV=True, ++ ) ++ return q_out, None, None, gate_out diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py index 919f46619a..c8a185388a 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py @@ -156,39 +972,1228 @@ index 919f46619a..c8a185388a 100644 BaseCudaGraphRunner, ) -+@contextmanager -+def _disable_tms_during_graph_capture(): -+ try: -+ from torch_memory_saver import torch_memory_saver -+ _impl = torch_memory_saver._impl -+ except Exception: -+ _impl = None ++@contextmanager ++def _disable_tms_during_graph_capture(): ++ try: ++ from torch_memory_saver import torch_memory_saver ++ _impl = torch_memory_saver._impl ++ except Exception: ++ _impl = None ++ ++ if _impl is not None: ++ _impl._binary_wrapper.cdll.tms_set_interesting_region(False) ++ try: ++ yield ++ finally: ++ if _impl is not None: ++ _impl._binary_wrapper.cdll.tms_set_interesting_region(True) + + class NPUCudaGraphBackend(BaseCudaGraphBackend): + """One torch.npu.NPUGraph per shape; attention metadata captured +@@ -117,7 +132,8 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend): + stream=self._capture_stream, + auto_dispatch_capture=True, + ): +- out = forward_fn() ++ with _disable_tms_during_graph_capture(): ++ out = forward_fn() + + self._graphs[shape_key] = graph + self._outputs[shape_key] = out +diff --git a/python/sglang/srt/hardware_backend/npu/moe/topk.py b/python/sglang/srt/hardware_backend/npu/moe/topk.py +index 5a51ef4e96..7fbe869fad 100644 +--- a/python/sglang/srt/hardware_backend/npu/moe/topk.py ++++ b/python/sglang/srt/hardware_backend/npu/moe/topk.py +@@ -5,6 +5,7 @@ from sgl_kernel_npu.norm.l1_norm import l1_norm + + from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder + from sglang.srt.eplb.expert_location_dispatch import topk_ids_logical_to_physical ++from sglang.srt.environ import envs + from sglang.srt.layers.moe.topk import ( + StandardTopKOutput, + capture_routed_experts_if_allowed, +@@ -15,6 +16,11 @@ if TYPE_CHECKING: + from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo + from sglang.srt.layers.moe.topk import TopKConfig, TopKOutput + ++# MoE ๅ‰ๆฎต่žๅˆๅŒ…๏ผˆmoe_front_fusion/v1๏ผŒๅˆๅ…ฅ่‡ช daikang ๅˆ†ๆ”ฏ๏ผ‰ๆ€ปๅผ€ๅ…ณ๏ผš ++# renorm=1 ๅ•็ฎ—ๅญ่ทฏ็”ฑ + v2.2 ่‡ชๅ†™ init_routingใ€‚้ป˜่ฎคๅ…ณ๏ผŒๅผ€ๅฏ๏ผš ++# SGLANG_MOE_FRONT_FUSION=1ใ€‚ไป…ไฝœ็”จไบŽไธ‹ๆ–น fast path๏ผˆๆ—  group/ๆ—  bias๏ผ‰ใ€‚ ++_moe_front_fusion = envs.SGLANG_MOE_FRONT_FUSION.get() ++ + + def fused_topk_npu( + hidden_states: torch.Tensor, +@@ -31,18 +37,43 @@ def fused_topk_npu( + + # Fast path: simple top-k without grouped routing and bias + if not use_grouped_topk and correction_bias is None: +- topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k_softmax( +- router_logits, +- k=topk_config.top_k, +- ) +- +- if renormalize: +- topk_weights = l1_norm( +- topk_weights +- if topk_config.num_fused_shared_experts == 0 +- else topk_weights[:, :-1] ++ if ( ++ _moe_front_fusion ++ and renormalize ++ and topk_config.num_fused_shared_experts == 0 ++ ): ++ # renorm=1 single op = gating_top_k_softmax + l1_norm. ++ # Feed bf16 directly: aclnnMoeGatingTopK supports BF16 input and ++ # computes softmax in fp32 after exact widening, so ids are ++ # bitwise unchanged; no .to(fp32) on the output (with bf16 output ++ # that would be a real cast). The deepep path's fp32 contract for ++ # topk_weights is restored by the defensive .to(fp32) before the ++ # two dispatch_a call sites in deepep.py. ++ topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k( ++ router_logits, ++ k=topk_config.top_k, ++ bias=None, ++ k_group=1, ++ group_count=1, ++ group_select_mode=0, ++ renorm=1, ++ norm_type=0, ++ routed_scaling_factor=1.0, ++ eps=float(1e-20), + ) +- topk_weights = topk_weights.to(torch.float32) ++ else: ++ topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k_softmax( ++ router_logits, ++ k=topk_config.top_k, ++ ) ++ ++ if renormalize: ++ topk_weights = l1_norm( ++ topk_weights ++ if topk_config.num_fused_shared_experts == 0 ++ else topk_weights[:, :-1] ++ ) ++ topk_weights = topk_weights.to(torch.float32) + + # sqrtsoftplus (DSV4 noaux_tc): the NPU op only scores sigmoid/softmax, so use + # a torch path. top-k over (scores + bias); weights from un-biased scores. +diff --git a/python/sglang/srt/hardware_backend/npu/moe_weight_prefetch.py b/python/sglang/srt/hardware_backend/npu/moe_weight_prefetch.py +new file mode 100644 +index 0000000000..5c05bffbaa +--- /dev/null ++++ b/python/sglang/srt/hardware_backend/npu/moe_weight_prefetch.py +@@ -0,0 +1,363 @@ ++"""MoE expert weight L2 prefetch for NPU decode (aclgraph replay). ++ ++Design: ++- Self-targeted single launch point: after this layer's prepare_attn, before ++ attention starts; the layer's own weight prefetch is launched there. EP ++ inter-layer communication (reduceScatter + allGather) completes inside the next ++ layer's prepare_attn, and TP's layer-end all_reduce inside the previous layer's ++ postprocess_layer โ€” so this point sits after all inter-layer communication on ++ both paths and does not contend with it for bandwidth; self-targeting also gives ++ layer 0 an attention window. ++- GMM1 (w13) full prefetch: launches w13(j), GDN layers only; on by default. ++- GMM2 (w2) full prefetch: merged into the same launch point, w2(j) chunks appended ++ after w13(j) (consumption order: w2 is consumed by GMM2(j), one GMM1+SwiGLU later ++ than w13); off by default. Merged-launch gate: w13+w2 <= 0.8*L2. ++- Best-effort: the main stream does not wait for the CMO stream before GMM (cache ++ warm-up has no RAW hazard); drained at step end (prevents cross-step backlog, and ++ capture legality requires the side stream to return to the main stream via event). ++- Launched only during graph capture (eager / prefill never activate); CMO is an ++ SDMA task and occupies no AIV/AIC cores. ++- Only layer objects registered for the target model trigger launches ++ (`_moe_prefetch_emit` marker); layers of the MTP draft model never launch even if ++ they reuse this forward from the same file. ++ ++Note: torch_npu is always lazily imported inside functions, so this module is safe ++to import on CUDA and other platforms. ++""" ++ ++from __future__ import annotations ++ ++import ctypes ++import glob ++import logging ++import os ++from typing import Dict, Optional ++ ++import torch ++ ++from sglang.srt.environ import envs ++ ++logger = logging.getLogger(__name__) ++ ++# ACL_DEV_ATTR_L2_CACHE_SIZE = RT_DEV_ATTR_L2_CACHE_SIZE = 302 ++_L2_CACHE_ATTR = 302 ++ ++_VALID_OPS = ("gmm1", "gmm2") ++_VALID_MODES = ("auto", "full", "active") ++ ++ ++# --------------------------------------------------------------------------- ++# L2 capacity query (ctypes, no compilation needed) ++# --------------------------------------------------------------------------- ++def _dlopen_first(candidates): ++ for path in candidates: ++ try: ++ return ctypes.CDLL(path) ++ except OSError: ++ continue ++ return None ++ ++ ++def _try_get_info(lib, func_name) -> Optional[int]: ++ fn = getattr(lib, func_name, None) ++ if fn is None: ++ return None ++ fn.restype = ctypes.c_int ++ fn.argtypes = [ctypes.c_uint32, ctypes.c_int, ctypes.POINTER(ctypes.c_int64)] ++ val = ctypes.c_int64(-1) ++ try: ++ rc = fn(0, _L2_CACHE_ATTR, ctypes.byref(val)) ++ except Exception: ++ return None ++ if rc == 0 and val.value > 0: ++ return int(val.value) ++ return None ++ ++ ++def _query_l2_size_bytes() -> Optional[int]: ++ """Return the L2 cache size in bytes; None on query failure (caller disables prefetch).""" ++ home = os.environ.get("ASCEND_TOOLKIT_HOME") or os.environ.get("ASCEND_HOME_PATH") ++ acl_paths = ["libascendcl.so", "libascendcl.so.1"] ++ rt_paths = ["libruntime.so"] ++ if home: ++ acl_paths.append(os.path.join(home, "lib64", "libascendcl.so")) ++ rt_paths.append(os.path.join(home, "lib64", "libruntime.so")) ++ acl_paths += sorted( ++ glob.glob("/usr/local/Ascend/ascend-toolkit/latest/lib64/libascendcl.so*") ++ ) ++ rt_paths += sorted( ++ glob.glob("/usr/local/Ascend/ascend-toolkit/latest/lib64/libruntime.so*") ++ ) ++ ++ lib = _dlopen_first(acl_paths) ++ if lib is not None: ++ for name in ("aclrtGetDeviceInfo", "aclrtGetDeviceAttr"): ++ size = _try_get_info(lib, name) ++ if size is not None: ++ return size ++ lib = _dlopen_first(rt_paths) ++ if lib is not None: ++ size = _try_get_info(lib, "rtsDeviceGetInfo") ++ if size is not None: ++ return size ++ return None ++ ++ ++def _is_capture_mode() -> bool: ++ """Lazy import to avoid a module-level dependency on the model_executor chain.""" ++ from sglang.srt.model_executor.runner_utils.capture_mode import ( ++ get_is_capture_mode, ++ ) ++ ++ return get_is_capture_mode() ++ ++ ++# --------------------------------------------------------------------------- ++# Manager (process-level singleton) ++# --------------------------------------------------------------------------- ++class _MoeWeightPrefetchManager: ++ def __init__(self): ++ self.configured = False ++ self.enabled = False ++ self.prefetch_gmm1 = False ++ self.prefetch_gmm2 = False ++ self.mode = "full" ++ self.chunk_bytes = 16 << 20 ++ self.l2_size: Optional[int] = None ++ self.budget_bytes = 0 ++ self.stream = None ++ # layer_id -> {"w13": Tensor, "w2": Tensor, "prefetch_target": bool} ++ self.layers: Dict[int, dict] = {} ++ self._capacity_checked = False ++ self._emitted_in_pass = False ++ ++ # ---- Configuration (parsed once at first registration) ---- ++ def configure(self): ++ if self.configured: ++ return ++ self.configured = True ++ ++ if not envs.SGLANG_NPU_MOE_PREFETCH.get(): ++ return # enabled stays False; everything below is a no-op ++ ++ ops_raw = str(envs.SGLANG_NPU_MOE_PREFETCH_OPS.get()).lower() ++ ops = [t.strip() for t in ops_raw.split(",") if t.strip()] ++ unknown = [t for t in ops if t not in _VALID_OPS] ++ if unknown: ++ logger.warning( ++ f"[MOE_PREFETCH] unknown ops {unknown} ignored, valid: {_VALID_OPS}" ++ ) ++ self.prefetch_gmm1 = "gmm1" in ops ++ self.prefetch_gmm2 = "gmm2" in ops ++ ++ mode = str(envs.SGLANG_NPU_MOE_PREFETCH_MODE.get()).lower() ++ if mode not in _VALID_MODES: ++ logger.warning( ++ f"[MOE_PREFETCH] unknown mode {mode!r}, fallback to 'auto'" ++ ) ++ mode = "auto" ++ if mode == "active": ++ # In-graph npu_prefetch offset/max_size are capture-time constants and ++ # nodes execute unconditionally, so runtime active indices cannot be fed ++ # into the baked CMO params; 'active' would need a copy-kernel + ++ # write-allocate design. Fall back to 'full'. ++ logger.warning( ++ "[MOE_PREFETCH] mode='active' is not supported " ++ "(in-graph CMO params are capture-time constants); " ++ "falling back to 'full'" ++ ) ++ mode = "full" ++ self.mode = "full" if mode == "auto" else mode # auto == full ++ ++ chunk_mib = int(envs.SGLANG_NPU_MOE_PREFETCH_CHUNK_MIB.get()) ++ if chunk_mib <= 0: ++ logger.warning( ++ f"[MOE_PREFETCH] invalid CHUNK_MIB={chunk_mib}, fallback to 16" ++ ) ++ chunk_mib = 16 ++ self.chunk_bytes = chunk_mib << 20 ++ ++ self.l2_size = _query_l2_size_bytes() ++ if self.l2_size is None: ++ logger.warning( ++ "[MOE_PREFETCH] failed to query L2 cache size " ++ "(aclrtGetDeviceInfo(302)); prefetch disabled" ++ ) ++ return ++ ++ budget_mib = int(envs.SGLANG_NPU_MOE_PREFETCH_BUDGET_MIB.get()) ++ self.budget_bytes = ( ++ (budget_mib << 20) if budget_mib > 0 else int(0.8 * self.l2_size) ++ ) ++ ++ self.enabled = self.prefetch_gmm1 or self.prefetch_gmm2 ++ logger.info( ++ f"[MOE_PREFETCH] configured: enabled={self.enabled} " ++ f"gmm1={self.prefetch_gmm1} gmm2={self.prefetch_gmm2} " ++ f"mode={self.mode} chunk={chunk_mib}MiB " ++ f"l2={self.l2_size >> 20}MiB budget={self.budget_bytes >> 20}MiB" ++ ) ++ ++ # ---- Capacity check (checked once at first MoE layer registration; layers are homogeneous) ---- ++ def _check_capacity(self, w13_bytes: int, w2_bytes: int): ++ if self._capacity_checked: ++ return ++ self._capacity_checked = True ++ mib = 1 << 20 ++ if self.prefetch_gmm1 and w13_bytes > self.budget_bytes: ++ logger.warning( ++ f"[MOE_PREFETCH] w13 size {w13_bytes / mib:.1f}MiB > budget " ++ f"{self.budget_bytes / mib:.1f}MiB (0.8*L2 or BUDGET_MIB); " ++ f"gmm1 prefetch disabled for this config " ++ f"(e.g. TP4 w13=256MiB is a known unsupported case)" ++ ) ++ self.prefetch_gmm1 = False ++ if self.prefetch_gmm2: ++ # Merged-launch gate: w13 (if gmm1 is on) + w2 stay resident in the same ++ # cross-layer window; gmm2 is allowed only if the total fits the budget. ++ resident = (w13_bytes if self.prefetch_gmm1 else 0) + w2_bytes ++ if resident > self.budget_bytes: ++ logger.warning( ++ f"[MOE_PREFETCH] resident w13+w2 size {resident / mib:.1f}MiB " ++ f"> budget {self.budget_bytes / mib:.1f}MiB; " ++ f"merged gmm2 prefetch disabled for this config " ++ f"(w13 and w2 launch together, so both must fit; " ++ f"TP8 w13+w2=192MiB is a known unsupported case)" ++ ) ++ self.prefetch_gmm2 = False ++ self.enabled = self.prefetch_gmm1 or self.prefetch_gmm2 ++ if self.enabled: ++ logger.info( ++ f"[MOE_PREFETCH] capacity check passed: w13={w13_bytes / mib:.1f}MiB " ++ f"(gmm1={self.prefetch_gmm1}) w2={w2_bytes / mib:.1f}MiB " ++ f"(gmm2={self.prefetch_gmm2})" ++ ) ++ ++ # ---- Model registration (called during model init, outside capture) ---- ++ def register_model(self, model): ++ self.configure() ++ if not self.enabled: ++ return ++ ++ if self.stream is None: ++ import torch_npu # noqa: F401 ++ ++ # Dedicated CMO stream: created outside capture (driver call), reused as a singleton. ++ self.stream = torch.npu.Stream() ++ ++ block_types = None ++ config = getattr(model, "config", None) ++ if config is not None: ++ try: ++ block_types = config.layers_block_type ++ except Exception: ++ block_types = None ++ ++ n_reg = 0 ++ layers = getattr(model, "layers", None) ++ if layers is None: ++ return ++ for layer in layers: ++ mlp = getattr(layer, "mlp", None) ++ experts = getattr(mlp, "experts", None) ++ w13 = getattr(experts, "w13_weight", None) ++ w2 = getattr(experts, "w2_weight", None) ++ layer_id = getattr(layer, "layer_id", None) ++ if w13 is None or w2 is None or layer_id is None: ++ continue # dense MLP / PPMissingLayer / non-MoE layer ++ ++ self._check_capacity( ++ w13.numel() * w13.element_size(), ++ w2.numel() * w2.element_size(), ++ ) ++ # Prefetch target: GDN layers only (an FA layer's FIA KV stream ++ # self-evicts and steals bandwidth); w13 and w2 share the target set. ++ is_gdn = ( ++ block_types is not None ++ and 0 <= layer_id < len(block_types) ++ and block_types[layer_id] == "linear_attention" ++ ) ++ self.layers[layer_id] = { ++ "w13": w13, ++ "w2": w2, ++ "prefetch_target": bool(is_gdn), ++ } ++ # The launch gate is marked per layer object: unregistered layers (e.g. ++ # MTP draft model layers, whose layer_id may collide with the registry) ++ # never launch. ++ layer._moe_prefetch_emit = True ++ n_reg += 1 ++ ++ if n_reg: ++ n_target = sum(1 for v in self.layers.values() if v["prefetch_target"]) ++ logger.info( ++ f"[MOE_PREFETCH] registered {n_reg} MoE layers, " ++ f"{n_target} GDN prefetch targets" ++ ) ++ ++ # ---- Launch (capture only; best-effort, no wait) ---- ++ def _emit_chunked(self, weights, anchor: torch.Tensor): ++ import torch_npu ++ ++ cur = torch.npu.current_stream() ++ self.stream.wait_stream(cur) # fork edge: the penalty is paid by the prefetch stream ++ with torch.npu.stream(self.stream): ++ for weight in weights: ++ nbytes = weight.numel() * weight.element_size() ++ offset = 0 ++ while offset < nbytes: ++ size = min(self.chunk_bytes, nbytes - offset) ++ torch_npu.npu_prefetch(weight, anchor, size, offset) ++ offset += size ++ self._emitted_in_pass = True ++ ++ def emit(self, layer, anchor: torch.Tensor): ++ """Launch this layer's w13(+w2) prefetch before its attention starts (after all inter-layer communication). ++ ++ Launch point = after prepare_attn in decoder layer forward. EP inter-layer ++ communication (reduceScatter + allGather) completes inside the next layer's ++ prepare_attn and TP's layer-end AR inside the previous layer's ++ postprocess_layer, so this point is after all inter-layer communication on ++ both paths. Chunk order: w13 first (consumed by GMM1), w2 after (consumed by ++ GMM2, with one extra GMM1+SwiGLU of slack). ++ """ ++ if not self.enabled or not _is_capture_mode(): ++ return ++ if not getattr(layer, "_moe_prefetch_emit", False): ++ return # unregistered layers (MTP draft etc.) never launch ++ target = self.layers.get(layer.layer_id) ++ if target is None or not target["prefetch_target"]: ++ return ++ weights = [] ++ if self.prefetch_gmm1: ++ weights.append(target["w13"]) ++ if self.prefetch_gmm2: ++ weights.append(target["w2"]) ++ if not weights: ++ return ++ self._emit_chunked(weights, anchor) ++ ++ # ---- Step-end drain (called after the model forward layer loop) ---- ++ def drain(self): ++ if not self.enabled or not self._emitted_in_pass or not _is_capture_mode(): ++ return ++ torch.npu.current_stream().wait_stream(self.stream) ++ self._emitted_in_pass = False ++ ++ ++_MANAGER = _MoeWeightPrefetchManager() ++ ++ ++# --------------------------------------------------------------------------- ++# Public interface (called by qwen3_5.py) ++# --------------------------------------------------------------------------- ++def moe_prefetch_register_model(model): ++ _MANAGER.register_model(model) ++ ++ ++def moe_prefetch_emit(layer, anchor: torch.Tensor): ++ _MANAGER.emit(layer, anchor) ++ ++ ++def moe_prefetch_step_drain(): ++ _MANAGER.drain() +diff --git a/python/sglang/srt/hardware_backend/npu/tp_ascendc_fusion_npu.py b/python/sglang/srt/hardware_backend/npu/tp_ascendc_fusion_npu.py +new file mode 100644 +index 0000000000..ea67cd2293 +--- /dev/null ++++ b/python/sglang/srt/hardware_backend/npu/tp_ascendc_fusion_npu.py +@@ -0,0 +1,347 @@ ++# Switch/guard module for Qwen3.5 NPU decode TP-path AscendC fusion. ++# ++# Two ops (registered by sgl-kernel-npu csrc as torch.ops.npu.*): ++# op1 fused_qkvzba_conv1d GDN decode: qkvzba split + causal_conv1d(UPDATE), ++# single kernel, returns (y, z, b, a); conv_states updated in place. ++# op2 fused_sigmoid_gating_recurrent GDN decode: sigmoid gating + delta rule update, ++# single AIV_ONLY kernel; drop-in replacement for the ++# production Triton kernel; ssm pool updated in place. ++# ++# fused_norm_qkv_proj_scatter / fused_sigmoid_mul_mm are not built in this version ++# (dropped); their guard functions are kept as always-disabled stubs so that an old ++# qwen3_5.py importing them still works. ++# ++# Switches: ++# SGLANG_NPU_TP_ASCENDC_FUSION master switch (default "0"), op1 only ++# SGLANG_NPU_TP_ASCENDC_FUSION_QKVZBA op1 switch: follows master when unset, explicit "0" disables op1 alone ++# SGLANG_NPU_GDN_RECURRENT_ASCENDC op2 switch (default "0"); evaluated at gdn_triton.py ++# import time, must be set before graph capture / server start ++# SGLANG_NPU_TP_ASCENDC_FUSION_DEBUG=1 log guard misses once per key ++# Guard decisions are baked into the graph at capture time; changing env afterwards has no effect. ++ ++import logging ++import os ++ ++import torch ++ ++logger = logging.getLogger(__name__) ++ ++_TP_FUSION_ENV = "SGLANG_NPU_TP_ASCENDC_FUSION" ++_TP_FUSION_QKVZBA_ENV = "SGLANG_NPU_TP_ASCENDC_FUSION_QKVZBA" ++_GDN_RECURRENT_ASCENDC_ENV = "SGLANG_NPU_GDN_RECURRENT_ASCENDC" ++_TP_DEBUG_ENV = "SGLANG_NPU_TP_ASCENDC_FUSION_DEBUG" ++_tp_debug_logged = set() ++ ++# Host-side hard constraints of the ops (mirror the csrc host TORCH_CHECKs; see each REGISTRATION.md): ++_TP_C0_ALIGN = 16 # bf16 C0 alignment (op1 row width must be a multiple of 16) ++_RECURRENT_HEAD_DIM = 128 # fused_sigmoid_gating_recurrent kernel specializes K/V to 128 ++_RECURRENT_MAX_HV = 8 # gating [8] pad limit of the same kernel ++_RECURRENT_MAX_N = 256 # cu/idx UB residency limit of the same kernel ++ ++_TP_OP_CACHE = {} ++ ++ ++def tp_fusion_enabled() -> bool: ++ """Master switch for TP-path AscendC fusion (default off; "1" enables, op1 only).""" ++ return os.environ.get(_TP_FUSION_ENV, "0") == "1" ++ ++ ++def tp_fusion_qkvzba_enabled() -> bool: ++ """op1 (fused_qkvzba_conv1d) switch: follows the master switch when unset; explicit "0" disables op1 alone.""" ++ return os.environ.get( ++ _TP_FUSION_QKVZBA_ENV, os.environ.get(_TP_FUSION_ENV, "0") ++ ) == "1" ++ ++ ++def gdn_recurrent_ascendc_enabled() -> bool: ++ """op2 (fused_sigmoid_gating_recurrent, AscendC recurrent) switch: default "0", enabled only by explicit "1".""" ++ return os.environ.get(_GDN_RECURRENT_ASCENDC_ENV, "0") == "1" ++ ++ ++def tp_debug_log(key, msg): ++ """When the debug switch is on, log once per key (for locating guard misses).""" ++ if os.environ.get(_TP_DEBUG_ENV, "0") != "1": ++ return ++ if key in _tp_debug_logged: ++ return ++ _tp_debug_logged.add(key) ++ logger.warning("[tp_ascendc_fusion] %s", msg) ++ ++ ++def tp_op_available(name: str) -> bool: ++ """Whether the op is registered on torch.ops.npu (i.e. sgl-kernel-npu was built with these ops). ++ ++ Registration happens when the sgl_kernel_npu extension loads (process start), so ++ the result is process-static and cached. ++ """ ++ got = _TP_OP_CACHE.get(name) ++ if got is None: ++ got = hasattr(torch.ops.npu, name) ++ _TP_OP_CACHE[name] = got ++ return got ++ ++ ++# --------------------------------------------------------------------------- ++# op1 (GDN decode): fused_qkvzba_conv1d ++# = fused_qkvzba_split_reshape_cat_contiguous + causal_conv1d(run_mode=1) ++# --------------------------------------------------------------------------- ++def tp_fused_qkvzba_conv1d_shape_supported( ++ num_k_heads_tp, num_v_heads_tp, head_k_dim, head_v_dim, conv_kernel_size ++) -> bool: ++ """Shape guard baked at init (Qwen3_5GatedDeltaNet level; head counts are per-rank after TP split). ++ ++ Mirrors the op host TORCH_CHECKs: num_v % num_k == 0 with ratio in {1,2,4} (same ++ as the replaced qwen3_5.py split-branch condition), conv width in [2,4], qkvWidth ++ and qkvz row width 16-aligned (32B DataCopy alignment of the z copy), z row ++ bytes <= 65535 (single-block DataCopy limit). ++ """ ++ ok = ( ++ tp_fusion_qkvzba_enabled() ++ and tp_op_available("fused_qkvzba_conv1d") ++ and num_k_heads_tp > 0 ++ and num_v_heads_tp > 0 ++ and head_k_dim > 0 ++ and head_v_dim > 0 ++ and num_v_heads_tp % num_k_heads_tp == 0 ++ and (num_v_heads_tp // num_k_heads_tp) in (1, 2, 4) ++ and 2 <= conv_kernel_size <= 4 ++ ) ++ if ok: ++ qkv_width = 2 * num_k_heads_tp * head_k_dim + num_v_heads_tp * head_v_dim ++ x_row_stride = qkv_width + num_v_heads_tp * head_v_dim ++ ok = ( ++ qkv_width % _TP_C0_ALIGN == 0 ++ and x_row_stride % _TP_C0_ALIGN == 0 ++ and num_v_heads_tp * head_v_dim * 2 <= 65535 ++ ) ++ if not ok: ++ tp_debug_log( ++ ("qkvzba_shape",), ++ "fused_qkvzba_conv1d init shape guard missed (falling back to stock split+causal_conv1d): " ++ f"nk_tp={num_k_heads_tp} nv_tp={num_v_heads_tp} dk={head_k_dim} " ++ f"dv={head_v_dim} width={conv_kernel_size} " ++ f"op_registered={tp_op_available('fused_qkvzba_conv1d')} " ++ f"{_TP_FUSION_QKVZBA_ENV}={os.environ.get(_TP_FUSION_QKVZBA_ENV, '')!r} " ++ f"{_TP_FUSION_ENV}={os.environ.get(_TP_FUSION_ENV, '0')!r}", ++ ) ++ return ok ++ ++ ++def tp_fused_qkvzba_conv1d_inputs_ok( ++ qkvz, mixed_ba, num_k_heads_tp, num_v_heads_tp, head_k_dim, head_v_dim ++) -> bool: ++ """Lightweight runtime check on the backend side (tuple-input tensor attributes; the main shape set is baked at init). ++ ++ Contiguity contract is a "row-stride view": column-contiguous (stride(1)==1), row ++ stride >= logical width, and qkvz row stride a multiple of 16 (mirrors the host ++ TORCH_CHECK for 32B alignment of the z copy). Contiguous inputs satisfy this ++ naturally; row-stride views (slices of the packed GEMM output) are read directly, ++ avoiding two .contiguous() calls. Must be deployed together with the matching ++ csrc host version (older hosts still TORCH_CHECK is_contiguous). ++ """ ++ qkv_width = 2 * num_k_heads_tp * head_k_dim + num_v_heads_tp * head_v_dim ++ return ( ++ torch.is_tensor(qkvz) ++ and torch.is_tensor(mixed_ba) ++ and qkvz.dim() == 2 ++ and mixed_ba.dim() == 2 ++ and qkvz.dtype == torch.bfloat16 ++ and mixed_ba.dtype == torch.bfloat16 ++ and qkvz.stride(1) == 1 ++ and qkvz.stride(0) >= qkvz.shape[1] ++ and qkvz.stride(0) % _TP_C0_ALIGN == 0 ++ and mixed_ba.stride(1) == 1 ++ and mixed_ba.stride(0) >= mixed_ba.shape[1] ++ and qkvz.shape[1] == qkv_width + num_v_heads_tp * head_v_dim ++ and mixed_ba.shape[0] == qkvz.shape[0] ++ and mixed_ba.shape[1] == 2 * num_v_heads_tp ++ ) ++ ++ ++# --------------------------------------------------------------------------- ++# op2 (GDN decode): fused_sigmoid_gating_recurrent (AscendC recurrent) ++# ++# Not bit-exact with the production Triton kernel (K-dim reduction order + ++# instruction-level Exp/Ln/Div differences); acceptance follows bounded error + ++# long-run drift + e2e A/B (see the op package README "precision gate" section). ++# --------------------------------------------------------------------------- ++def _ascendc_recurrent_supported( ++ q, k, v, a, b, initial_state_source, initial_state_indices, cu_seqlens, ++ A_log, dt_bias, ++) -> bool: ++ """Runtime guard (mirrors host TORCH_CHECKs; on miss the wrapper falls back to stock Triton). ++ ++ Reads tensor attributes only (shape/stride/dtype/contiguity), never device data, ++ so it is graph-capture safe. fp32 dtype of A_log/dt_bias is not enforced here โ€” ++ the wrapper widens losslessly (see call site); only numel is checked to avoid a ++ host TORCH_CHECK failure. ++ """ ++ ok = ( ++ tp_op_available("fused_sigmoid_gating_recurrent") ++ and torch.is_tensor(q) ++ and q.dim() == 4 ++ and q.size(0) == 1 ++ and cu_seqlens is not None ++ and torch.is_tensor(cu_seqlens) ++ and q.size(1) == cu_seqlens.numel() - 1 # T == N: 1 token per sequence ++ and 1 <= q.size(1) <= _RECURRENT_MAX_N ++ and q.size(3) == _RECURRENT_HEAD_DIM ++ and v.size(3) == _RECURRENT_HEAD_DIM ++ and q.size(2) >= 1 ++ and v.size(2) >= 1 ++ and v.size(2) <= _RECURRENT_MAX_HV ++ and v.size(2) % q.size(2) == 0 ++ and q.dtype == torch.bfloat16 ++ and k.dtype == torch.bfloat16 ++ and v.dtype == torch.bfloat16 ++ and a.dtype == torch.bfloat16 ++ and b.dtype == torch.bfloat16 ++ and initial_state_source.dtype in (torch.bfloat16, torch.float32) ++ and q.stride(3) == 1 ++ and k.stride(3) == 1 ++ and v.stride(3) == 1 ++ and q.stride(2) == q.size(3) ++ and k.stride(2) == k.size(3) ++ and v.stride(2) == v.size(3) ++ and a.is_contiguous() ++ and b.is_contiguous() ++ and initial_state_source.is_contiguous() ++ and initial_state_source.dim() == 4 ++ and initial_state_source.size(1) == v.size(2) ++ and initial_state_source.size(2) == _RECURRENT_HEAD_DIM ++ and initial_state_source.size(3) == _RECURRENT_HEAD_DIM ++ and initial_state_indices is not None ++ and initial_state_indices.numel() >= q.size(1) ++ and A_log.numel() == v.size(2) # host: A_log/dt_bias numel == HV ++ and dt_bias.numel() == v.size(2) ++ ) ++ if not ok: ++ tp_debug_log( ++ ("recurrent_shape",), ++ "fused_sigmoid_gating_recurrent runtime guard missed (falling back to stock Triton): " ++ f"q={tuple(q.shape)}/{q.dtype} v={tuple(v.shape)}/{v.dtype} " ++ f"pool={tuple(initial_state_source.shape)}/{initial_state_source.dtype} " ++ f"T_vs_N={q.size(1)}/{cu_seqlens.numel() - 1 if torch.is_tensor(cu_seqlens) else None} " ++ f"op_registered={tp_op_available('fused_sigmoid_gating_recurrent')}", ++ ) ++ return ok ++ ++ ++def fused_sigmoid_gating_delta_rule_update_ascendc( ++ A_log, ++ a, ++ dt_bias, ++ softplus_beta, ++ softplus_threshold, ++ q, ++ k, ++ v, ++ b, ++ initial_state_source, ++ initial_state_indices, ++ scale=None, ++ use_qk_l2norm_in_kernel=False, ++ cu_seqlens=None, ++): ++ """AscendC drop-in replacement for the production Triton wrapper ++ (fused_sigmoid_gating_delta_rule_update_npu), same signature and semantics; ++ decode only (T==N). ++ ++ On guard miss, falls back to the stock Triton wrapper (silent; with ++ SGLANG_NPU_TP_ASCENDC_FUSION_DEBUG=1 the reason is logged once). ++ """ ++ if not _ascendc_recurrent_supported( ++ q, k, v, a, b, initial_state_source, initial_state_indices, cu_seqlens, ++ A_log, dt_bias, ++ ): ++ from sgl_kernel_npu.fla.fused_sigmoid_gating_recurrent import ( ++ fused_sigmoid_gating_delta_rule_update_npu, ++ ) ++ ++ return fused_sigmoid_gating_delta_rule_update_npu( ++ A_log=A_log, ++ a=a, ++ dt_bias=dt_bias, ++ softplus_beta=softplus_beta, ++ softplus_threshold=softplus_threshold, ++ q=q, ++ k=k, ++ v=v, ++ b=b, ++ initial_state_source=initial_state_source, ++ initial_state_indices=initial_state_indices, ++ scale=scale, ++ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ++ cu_seqlens=cu_seqlens, ++ ) ++ # scale default matches the Triton wrapper (python-side K**-0.5, same double->float path) ++ if scale is None: ++ scale = k.shape[-1] ** -0.5 ++ else: ++ assert scale > 0, "scale must be positive" ++ # A_log/dt_bias: the host only accepts fp32, but they may be bf16 after checkpoint ++ # load. Widen losslessly here (bit-exact with the in-kernel .to(tl.float32) of the ++ # Triton kernel), mirroring the Triton wrapper's .contiguous(). ++ if A_log.dtype != torch.float32: ++ A_log = A_log.float() ++ if dt_bias.dtype != torch.float32: ++ dt_bias = dt_bias.float() ++ A_log = A_log.contiguous() ++ dt_bias = dt_bias.contiguous() ++ return torch.ops.npu.fused_sigmoid_gating_recurrent( ++ A_log, ++ a, ++ dt_bias, ++ float(softplus_beta), ++ float(softplus_threshold), ++ q, ++ k, ++ v, ++ b, ++ initial_state_source, ++ initial_state_indices, ++ float(scale), ++ cu_seqlens, ++ bool(use_qk_l2norm_in_kernel), ++ q.stride(1), ++ k.stride(1), ++ v.stride(1), ++ ) ++ ++ ++# --------------------------------------------------------------------------- ++# Legacy stubs (always disabled): fused_norm_qkv_proj_scatter / fused_sigmoid_mul_mm ++# These two ops were dropped and are not built in this version. If an old qwen3_5.py ++# imports the 5 functions below, the stubs keep it working and never activate ++# (callers fall back to the original path). ++# --------------------------------------------------------------------------- ++def tp_fusion_nqps_enabled() -> bool: ++ """fused_norm_qkv_proj_scatter switch stub: always False (op not built).""" ++ return False ++ ++ ++def tp_fusion_sigmm_enabled() -> bool: ++ """fused_sigmoid_mul_mm switch stub: always False (op not built).""" ++ return False ++ ++ ++def tp_norm_qkv_scatter_shape_supported(layer) -> bool: ++ """fused_norm_qkv_proj_scatter init guard stub: always False.""" ++ return False ++ ++ ++def tp_norm_qkv_scatter_context( ++ layer, hidden_states, residual, forward_batch, captured_last_layer_outputs ++): ++ """fused_norm_qkv_proj_scatter runtime guard stub: always None (caller falls back to the original path).""" ++ return None ++ ++ ++def tp_sigmoid_mul_mm_shape_supported(o_proj) -> bool: ++ """fused_sigmoid_mul_mm init guard stub: always False.""" ++ return False ++ ++ ++def tp_sigmoid_mul_mm_runtime_ok(attn_output, gate) -> bool: ++ """fused_sigmoid_mul_mm runtime guard stub: always False.""" ++ return False +diff --git a/python/sglang/srt/layers/attention/fla/cumsum.py b/python/sglang/srt/layers/attention/fla/cumsum.py +index 911b78e004..ff78179f56 100644 +--- a/python/sglang/srt/layers/attention/fla/cumsum.py ++++ b/python/sglang/srt/layers/attention/fla/cumsum.py +@@ -75,7 +75,12 @@ def chunk_local_cumsum_scalar_kernel( + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], +- key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE", "HAS_SCALE"], ++ # Keep key at 6 entries: triton-ascend's AutoTilingTuner (which replaces ++ # triton.autotune after the first Triton compile in the process) rejects ++ # longer key lists at decoration time, crashing engine boot if this module ++ # is imported late. The key only controls autotune cache granularity; the ++ # tuned configs are unchanged. ++ key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE"], + ) + @triton.jit(do_not_specialize=["T"]) + def chunk_local_cumsum_vector_kernel( +diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py +index abdfa5e376..26b9d0a4fd 100644 +--- a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py ++++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py +@@ -1,3 +1,5 @@ ++import os ++ + import torch + + from sglang.srt.layers.attention.linear.kernels.kernel_backend import ( +@@ -25,6 +27,25 @@ if is_npu(): + + chunk_gated_delta_rule = chunk_gated_delta_rule_npu + fused_sigmoid_gating_delta_rule_update = fused_sigmoid_gating_delta_rule_update_npu ++ ++ # Three-way choice of the NPU decode recurrent kernel (resolved at import ++ # time โ€” set the env before server start; changing it after capture has no ++ # effect): ++ # SGLANG_NPU_GDN_RECURRENT_ASCENDC=1 โ†’ AscendC version (falls back to ++ # stock internally on guard miss) ++ # SGLANG_NPU_GDN_UPDATE_FUSED=1 โ†’ decode-optimized Triton version ++ # neither set โ†’ stock ++ # If both envs are set, ASCENDC wins. ++ if os.environ.get("SGLANG_NPU_GDN_RECURRENT_ASCENDC", "0") == "1": ++ from sglang.srt.hardware_backend.npu.tp_ascendc_fusion_npu import ( ++ fused_sigmoid_gating_delta_rule_update_ascendc as _gdn_decode_update, ++ ) ++ elif os.environ.get("SGLANG_NPU_GDN_UPDATE_FUSED", "0") == "1": ++ from sgl_kernel_npu.fla.fused_sigmoid_gating_recurrent_decode_optimized import ( ++ fused_sigmoid_gating_delta_rule_update_decode_npu as _gdn_decode_update, ++ ) ++ else: ++ _gdn_decode_update = fused_sigmoid_gating_delta_rule_update + elif is_cpu(): + from sgl_kernel.mamba import chunk_gated_delta_rule_cpu + +@@ -150,7 +171,10 @@ class TritonGDNKernel(LinearAttnKernelBase): + query_start_loc: torch.Tensor, + **kwargs, + ) -> torch.Tensor: +- return fused_sigmoid_gating_delta_rule_update( ++ # _gdn_decode_update is bound by env at import time (three-way); the ++ # kwargs signatures are aligned across all three. target_verify passes ++ # extended args and bypasses this method, so it is unaffected. ++ return _gdn_decode_update( + A_log=A_log, + dt_bias=dt_bias, + q=q, +diff --git a/python/sglang/srt/layers/exp_race_fused_triton.py b/python/sglang/srt/layers/exp_race_fused_triton.py +new file mode 100644 +index 0000000000..f6ee69e480 +--- /dev/null ++++ b/python/sglang/srt/layers/exp_race_fused_triton.py +@@ -0,0 +1,239 @@ ++# Fused exponential-race sampling consumption for Ascend NPU. ++# ++# Replaces the main-stream consumption chain ++# u.neg_().add_(1.0); torch.minimum(u, bound, out=u); u.log_().neg_() ++# sampled = torch.div(probs, u).argmax(dim=-1) ++# with fused Triton kernels that read probs and u once and write one int32 ++# index per row. ++# ++# Math is identical to the unfused chain (stock aten::exponential_ NPU ++# composite semantics): ++# x = min(1 - u, 1 - eps/2); q = -log(x); out = argmax_rows(probs / q) ++# Tie-breaking matches torch.argmax: first occurrence wins, within a tile ++# (lowest lane holding the tile max) and across tiles (partials carry the ++# global index; the reducer takes the lowest index among value ties). ++# ++# Two-kernel "partials" structure: kernel 1 writes one (max, arg) slot per ++# tile with NO loop-carried state, kernel 2 reduces the <=64 partials in one ++# tile. Kernel 1 uses only compiler-verified idioms on triton-ascend: ++# - loop form `for v0 in tl.range(0, V, BLOCK_V)`, loop variable IS the ++# offset; ++# - `lane = tl.arange(0, BLOCK_V)` is the only index vector used as a VALUE ++# (never an arithmetically-derived offs); ++# - global index assembled from scalars only: `v0 + tile_lane`, ++# `v0 // BLOCK_V`. ++# ++# Fail-closed: if triton is unavailable the module still imports and ++# exp_race_argmax_available() returns False; the sampler keeps the unfused ++# chain. ++ ++import torch ++ ++try: ++ import triton ++ import triton.language as tl ++ ++ _HAS_TRITON = True ++except Exception: # pragma: no cover - depends on server env ++ _HAS_TRITON = False ++ ++# BLOCK_V=4096 keeps the tile buffers within the per-AIV-core UB budget on ++# 910C; 8192 overflows UB at compile time. If a future triton-ascend version ++# still overflows, drop to 2048. ++_BLOCK_V = 4096 ++ ++# Partials reducer tile width: 64 covers any V <= 262144 with one tile. ++_BLOCK_NT = 64 ++ ++if _HAS_TRITON: ++ ++ @triton.jit ++ def _exp_race_partial_kernel( ++ probs_ptr, ++ u_ptr, ++ pval_ptr, ++ parg_ptr, ++ V, ++ NT, ++ bound, ++ BLOCK_V: tl.constexpr, ++ ): ++ row = tl.program_id(0).to(tl.int64) ++ base = row * V ++ pbase = row * NT ++ lane = tl.arange(0, BLOCK_V) ++ ++ # No loop-carried state: every iteration writes only its own slot. ++ for v0 in tl.range(0, V, BLOCK_V): ++ offs = v0 + lane ++ mask = offs < V ++ # x = min(1 - u, 1 - eps/2): exact in fp32 (1-u is exact on the ++ # uniform grid, min is exact). q = -log(x) stays finite and ++ # strictly positive, so masked lanes pinned at -inf can never win. ++ u = tl.load(u_ptr + base + offs, mask=mask, other=0.5) ++ x = 1.0 - u ++ x = tl.minimum(x, bound) ++ q = -tl.log(x) ++ p = tl.load(probs_ptr + base + offs, mask=mask, other=0.0) ++ s = p / q ++ s = tl.where(mask, s, float("-inf")) ++ ++ tile_max = tl.max(s, axis=0) ++ # First-occurrence tie-break inside the tile: lowest LANE holding ++ # the tile max (pure arange as the only index value; global index ++ # reassembled from scalars below). Avoids tl.argmax whose tie ++ # semantics on triton-ascend are unverified. ++ cand = tl.where(s == tile_max, lane, BLOCK_V) ++ tile_lane = tl.min(cand, axis=0) ++ tile_arg = v0 + tile_lane ++ ++ slot = v0 // BLOCK_V ++ tl.store(pval_ptr + pbase + slot, tile_max) ++ tl.store(parg_ptr + pbase + slot, tile_arg) ++ ++ @triton.jit ++ def _exp_race_partial_argmax_kernel( ++ probs_ptr, ++ u_ptr, ++ pval_ptr, ++ parg_ptr, ++ V, ++ NT, ++ bound, ++ BLOCK_V: tl.constexpr, ++ ): ++ # Diagnostic arm: identical to _exp_race_partial_kernel but takes the ++ # per-tile lane via tl.argmax (tie behavior unverified). ++ row = tl.program_id(0).to(tl.int64) ++ base = row * V ++ pbase = row * NT ++ lane = tl.arange(0, BLOCK_V) ++ ++ for v0 in tl.range(0, V, BLOCK_V): ++ offs = v0 + lane ++ mask = offs < V ++ u = tl.load(u_ptr + base + offs, mask=mask, other=0.5) ++ x = 1.0 - u ++ x = tl.minimum(x, bound) ++ q = -tl.log(x) ++ p = tl.load(probs_ptr + base + offs, mask=mask, other=0.0) ++ s = p / q ++ s = tl.where(mask, s, float("-inf")) ++ ++ tile_max = tl.max(s, axis=0) ++ tile_lane = tl.argmax(s, axis=0) ++ tile_arg = v0 + tile_lane ++ ++ slot = v0 // BLOCK_V ++ tl.store(pval_ptr + pbase + slot, tile_max) ++ tl.store(parg_ptr + pbase + slot, tile_arg) ++ ++ @triton.jit ++ def _exp_race_reduce_kernel( ++ pval_ptr, ++ parg_ptr, ++ out_ptr, ++ NT, ++ BLOCK_NT: tl.constexpr, ++ ): ++ row = tl.program_id(0).to(tl.int64) ++ pbase = row * NT ++ ++ offs = tl.arange(0, BLOCK_NT) ++ mask = offs < NT ++ vals = tl.load(pval_ptr + pbase + offs, mask=mask, other=float("-inf")) ++ args = tl.load(parg_ptr + pbase + offs, mask=mask, other=2147483647) ++ ++ total = tl.max(vals, axis=0) ++ # First occurrence across tiles: lowest global index among value ties ++ # (partials store global indices, so min over matches == torch.argmax). ++ cand = tl.where(vals == total, args, 2147483647) ++ best = tl.min(cand, axis=0) ++ ++ tl.store(out_ptr + row, best) ++ ++ @triton.jit ++ def _exp_race_q_kernel( ++ u_ptr, ++ q_ptr, ++ V, ++ bound, ++ BLOCK_V: tl.constexpr, ++ ): ++ # Validation twin (UT only): materializes q = -log(min(1-u, bound)) ++ # elementwise so the UT can diff it against the unfused torch chain. ++ row = tl.program_id(0).to(tl.int64) ++ base = row * V ++ for v0 in tl.range(0, V, BLOCK_V): ++ offs = v0 + tl.arange(0, BLOCK_V) ++ mask = offs < V ++ u = tl.load(u_ptr + base + offs, mask=mask, other=0.5) ++ x = 1.0 - u ++ x = tl.minimum(x, bound) ++ q = -tl.log(x) ++ tl.store(q_ptr + base + offs, q, mask=mask) ++ ++ ++def exp_race_argmax_available() -> bool: ++ return _HAS_TRITON ++ ++ ++def _bound_for(dtype: torch.dtype) -> float: ++ # Same bound as sampler._async_exp_min_bound: 1 - finfo(dtype).eps / 2 ++ # (fp32 -> 1 - 2^-24, exactly representable; passed as an fp32 kernel arg). ++ return 1.0 - torch.finfo(dtype).eps / 2.0 ++ ++ ++def _check_inputs(probs: torch.Tensor, u: torch.Tensor): ++ assert probs.dtype == torch.float32 and u.dtype == torch.float32 ++ assert probs.shape == u.shape and probs.dim() == 2 ++ ++ ++def _run_partials(partial_kernel, probs: torch.Tensor, u: torch.Tensor) -> torch.Tensor: ++ B, V = probs.shape ++ NT = (V + _BLOCK_V - 1) // _BLOCK_V ++ assert NT <= _BLOCK_NT, "vocab too large for the partials reducer tile" ++ device = probs.device ++ pval = torch.empty((B, NT), dtype=torch.float32, device=device) ++ parg = torch.empty((B, NT), dtype=torch.int32, device=device) ++ bound = _bound_for(probs.dtype) ++ partial_kernel[(B,)]( ++ probs, u, pval, parg, V, NT, bound, BLOCK_V=_BLOCK_V, num_warps=4 ++ ) ++ out = torch.empty(B, dtype=torch.int32, device=device) ++ _exp_race_reduce_kernel[(B,)]( ++ pval, parg, out, NT, BLOCK_NT=_BLOCK_NT, num_warps=4 ++ ) ++ return out ++ ++ ++def exp_race_argmax(probs: torch.Tensor, u: torch.Tensor) -> torch.Tensor: ++ """Fused consumption: row argmax of probs / q(u), int32 [B]. ++ ++ Two-kernel partials path (production, exact first-occurrence tie-break). ++ Non-destructive: neither probs nor u is modified. ++ """ ++ _check_inputs(probs, u) ++ probs = probs.contiguous() ++ u = u.contiguous() ++ return _run_partials(_exp_race_partial_kernel, probs, u) ++ ++ ++def exp_race_argmax_argmaxarm(probs: torch.Tensor, u: torch.Tensor) -> torch.Tensor: ++ """Diagnostic arm (UT only): per-tile lane via tl.argmax.""" ++ _check_inputs(probs, u) ++ probs = probs.contiguous() ++ u = u.contiguous() ++ return _run_partials(_exp_race_partial_argmax_kernel, probs, u) ++ ++ ++def exp_race_q_values(u: torch.Tensor) -> torch.Tensor: ++ """Validation twin: materialize q for the ULP/bitwise diff (UT only).""" ++ assert u.dtype == torch.float32 and u.dim() == 2 ++ u = u.contiguous() ++ B, V = u.shape ++ q = torch.empty_like(u) ++ _exp_race_q_kernel[(B,)]( ++ u, q, V, _bound_for(u.dtype), BLOCK_V=_BLOCK_V, num_warps=4 ++ ) ++ return q +diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py +index 23c46113a1..492931c979 100644 +--- a/python/sglang/srt/layers/layernorm.py ++++ b/python/sglang/srt/layers/layernorm.py +@@ -142,6 +142,13 @@ if _is_npu: + import torch_npu + from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm + ++ # full_attention A3a๏ผš่กŒๅนถ่กŒ add_gemma_rms_norm๏ผˆgrid=(batch,)๏ผ‰๏ผŒไธŽ stock ++ # ้€ไฝไธ€่‡ด๏ผ›่ฐƒ็”จ็‚นๅธฆๅฝข็Šถๅฎˆๅซ๏ผŒๆœชๅ‘ฝไธญ่‡ชๅŠจๅ›ž้€€ stock ++ from sglang.srt.hardware_backend.npu.attention.full_attention_fusion_npu import ( ++ fa_add_gemma_rms_norm_v2, ++ fa_add_gemma_rms_norm_v2_supported, ++ ) + -+ if _impl is not None: -+ _impl._binary_wrapper.cdll.tms_set_interesting_region(False) -+ try: -+ yield -+ finally: -+ if _impl is not None: -+ _impl._binary_wrapper.cdll.tms_set_interesting_region(True) - class NPUCudaGraphBackend(BaseCudaGraphBackend): - """One torch.npu.NPUGraph per shape; attention metadata captured -@@ -117,7 +132,8 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend): - stream=self._capture_stream, - auto_dispatch_capture=True, - ): -- out = forward_fn() -+ with _disable_tms_during_graph_capture(): -+ out = forward_fn() + def _forward_with_allreduce_fusion( + norm_module, +@@ -818,9 +825,17 @@ class GemmaRMSNorm(MultiPlatformOp): + if residual is not None: + if post_residual_addition is not None: + residual = residual + post_residual_addition +- norm_out, residual = add_gemma_rms_norm( +- x, self.weight, residual, self.variance_epsilon +- ) ++ # full_attention A3a๏ผš่กŒๅนถ่กŒ็‰ˆ๏ผˆgrid=(batch,)๏ผ‰๏ผŒๆ•ฐๅญฆไธŽ stock ๅฎŒๅ…จ ++ # ไธ€่‡ด๏ผˆfp32 ไธญ้—ดใ€gemma w+1๏ผ‰๏ผ›hidden ้ž 2 ็š„ๅน‚/้ž bf16/้ž่ฟž็ปญๆ—ถ ++ # ๅ›ž้€€ stock ++ if fa_add_gemma_rms_norm_v2_supported(x, residual): ++ norm_out, residual = fa_add_gemma_rms_norm_v2( ++ x, self.weight, residual, self.variance_epsilon ++ ) ++ else: ++ norm_out, residual = add_gemma_rms_norm( ++ x, self.weight, residual, self.variance_epsilon ++ ) + return norm_out, residual - self._graphs[shape_key] = graph - self._outputs[shape_key] = out + x, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.variance_epsilon) +diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +index 33a51a89da..9f0bc41b78 100644 +--- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py ++++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +@@ -487,6 +487,9 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase): + topk_output: TopKOutput, + ): + topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids ++ # Restore the deepep path's fp32 contract for topk_weights when ++ # front_fusion emits bf16 directly (no-op when the output is fp32). ++ topk_weights = topk_weights.to(torch.float32) + topk_ids = topk_ids.to(torch.int64) + if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and self.use_fp8: + # TODO hard code 128 block quant,use fp8 communication +@@ -652,6 +655,9 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): + ): + buffer = self._get_buffer() + topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids ++ # Restore the deepep path's fp32 contract for topk_weights when ++ # front_fusion emits bf16 directly (no-op when the output is fp32). ++ topk_weights = topk_weights.to(torch.float32) + topk_ids = topk_ids.to(torch.int64) + expected_m = ( + hidden_states.shape[0] * buffer.group_size * topk_ids.shape[1] diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py -index 82a3d77f05..25d4c21406 100644 +index 82a3d77f05..f42e166ce7 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py -@@ -402,10 +402,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): +@@ -55,6 +55,13 @@ _is_hip = is_hip() + _is_cpu = is_cpu() + _is_npu = is_npu() + _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip ++# MoE ๅ‰ๆฎต่žๅˆๅŒ…๏ผˆmoe_front_fusion/v1๏ผŒๅˆๅ…ฅ่‡ช daikang ๅˆ†ๆ”ฏ๏ผ‰ๆ€ปๅผ€ๅ…ณ๏ผš ++# v2.2 ่‡ชๅ†™ init_routing๏ผˆ0 ๅฎนๅทฎ้€ไฝ้ชŒๆ”ถ๏ผ‰๏ผŒไป…ไฝœ็”จไบŽ BF16 ๆ— ้‡ๅŒ–่ทฏๅพ„ใ€‚ ++_moe_front_fusion = envs.SGLANG_MOE_FRONT_FUSION.get() ++# GMM2๏ผˆw2 down_proj๏ผ‰่ตฐ sgl_kernel_npu.moe.persistent_gmm ็š„ Triton ++# persistent kernel๏ผˆ้ป˜่ฎค stock๏ผŒๅผ€ๅฏ๏ผšSGLANG_GMM2_TRITON=1๏ผ‰ใ€‚ไป…ไฝœ็”จไบŽ ++# BF16 ๆ— ้‡ๅŒ–ๆ—  bias ่ทฏๅพ„ใ€‚ ++_gmm2_triton = envs.SGLANG_GMM2_TRITON.get() + + if _use_aiter: + from aiter.ops.shuffle import shuffle_weight +@@ -402,10 +409,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): layer.w2_weight.data = layer.w2_weight.data.reshape( layer.num_local_experts, *new_shape_w2 ) @@ -196,6 +2201,11 @@ index 82a3d77f05..25d4c21406 100644 - for weight_name in ["w13_weight", "w2_weight"]: - weight = getattr(layer, weight_name) - weight.data = npu_format_cast(weight) ++ if _is_npu and _gmm2_triton: ++ # Triton GMM2 kernel ไปฅ่ฃธๆŒ‡้’ˆๆŒ‰ [E, N, K] ND ่ฟž็ปญๅญ˜ๅ‚จ่ฏปๅ– w2๏ผŒ ++ # ๅฟ…้กปไฟ่ฏ ND ่ฟž็ปญ๏ผšFRACTAL_NZ ๅˆ†ๅฝข้‡ๆŽ’ไผš่ฎฉ่ฃธๆŒ‡้’ˆ่ฏปๅ–ๅ‡บ้”™ ++ # ๏ผˆๆœฌๅˆ†ๆ”ฏ unquant NPU ่ทฏๅพ„ๆœฌ่บซไธๅš format cast๏ผŒๆญคๅค„ไป…ๅ…œๅบ•๏ผ‰ใ€‚ ++ layer.w2_weight.data = layer.w2_weight.data.contiguous() + # if _is_npu: + # for weight_name in ["w13_weight", "w2_weight"]: + # weight = getattr(layer, weight_name) @@ -203,8 +2213,103 @@ index 82a3d77f05..25d4c21406 100644 return +@@ -739,19 +751,40 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): + num_experts = layer.num_experts + top_k = layer.top_k or topk_ids.shape[1] # in case layer.top_k is not set + +- hidden_states, expanded_row_idx, expert_tokens, _ = ( +- torch.ops.npu.npu_moe_init_routing_v2( +- x, +- topk_ids, +- active_num=num_tokens * top_k, +- expert_num=num_experts, +- expert_tokens_num_type=1, +- expert_tokens_num_flag=True, +- active_expert_range=[0, num_experts], +- quant_mode=-1, ++ m = num_tokens * top_k ++ h = x.shape[1] ++ expert_offsets = None ++ if ( ++ _moe_front_fusion ++ and m <= 512 ++ and m % 8 == 0 ++ and (h & (h - 1)) == 0 ++ and x.dtype == torch.bfloat16 ++ ): ++ # v2.2 ่‡ชๅ†™ init routing๏ผˆmoe_front_fusion/v1๏ผŒๅˆๅ…ฅ่‡ช daikang ๅˆ†ๆ”ฏ๏ผŒ ++ # ๅ…ญ่ฝฎๅ•ๆต‹ 0 ๅฎนๅทฎ้€ไฝ้ชŒๆ”ถ๏ผ‰๏ผš่ฏญไน‰ไธŽ npu_moe_init_routing_v2(type=1) ++ # ้€ไฝไธ€่‡ด๏ผŒๅนถๅŽŸ็”Ÿ่พ“ๅ‡บ exclusive offsets๏ผˆint32 [E]๏ผ‰ไพ› persistent ++ # GMM2 offsets= ็›ด็”จใ€‚ๅฝขๆ€้—จๅค–ๅ›ž้€€ stock v2ใ€‚ ++ from sgl_kernel_npu.moe.moe_front_routing import moe_init_routing_v22 ++ ++ hidden_states, expanded_row_idx, expert_tokens, excl, _incl = ( ++ moe_init_routing_v22(x, topk_ids, num_experts, top_k) + ) +- ) +- expert_tokens = expert_tokens.to(torch.int64) ++ expert_offsets = excl ++ else: ++ hidden_states, expanded_row_idx, expert_tokens, _ = ( ++ torch.ops.npu.npu_moe_init_routing_v2( ++ x, ++ topk_ids, ++ active_num=m, ++ expert_num=num_experts, ++ expert_tokens_num_type=1, ++ expert_tokens_num_flag=True, ++ active_expert_range=[0, num_experts], ++ quant_mode=-1, ++ ) ++ ) ++ expert_tokens = expert_tokens.to(torch.int64) + w13_bias = [layer.w13_weight_bias] if self.with_bias else None + w2_bias = [layer.w2_weight_bias] if self.with_bias else None + +@@ -792,16 +825,31 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): + hidden_states = GeluAndMul()(hidden_states) + + # gmm2: down_proj +- hidden_states = torch.ops.npu.npu_grouped_matmul( +- x=[hidden_states], +- weight=[layer.w2_weight.transpose(1, 2)], +- bias=w2_bias, +- split_item=2, +- group_list_type=1, +- group_type=0, +- group_list=expert_tokens, +- output_dtype=original_dtype, +- )[0] ++ if _gmm2_triton and w2_bias is None: ++ # Triton persistent GMM2๏ผˆgmm/v1๏ผŒๅˆๅ…ฅ่‡ช daikang ๅˆ†ๆ”ฏ๏ผ‰๏ผšw2 ไธบ ++ # [E, N, K] ND ่ฟž็ปญๅญ˜ๅ‚จ๏ผˆ่ทณ่ฟ‡ FRACTAL_NZ cast๏ผ‰๏ผŒexpert_tokens ไธบ ++ # per-expert counts๏ผˆgroup_list_type=1๏ผŒint64๏ผ‰ใ€‚offsets ็ผบ็œๆ—ถ็”ฑ ++ # ๅ†…็ฝฎๅ‰็ฝฎ kernel ๆŽจๅฏผ๏ผ›SGLANG_MOE_FRONT_FUSION=1 ๆ—ถ v2.2 init ++ # ๅŽŸ็”Ÿ excl ็›ดๅ–‚๏ผŒ็œ _gmm2_offsets_kernel ~4.4ยตs/ๅฑ‚ใ€‚ ++ from sgl_kernel_npu.moe.persistent_gmm import persistent_grouped_matmul ++ ++ hidden_states = persistent_grouped_matmul( ++ hidden_states, ++ layer.w2_weight, ++ expert_tokens, ++ offsets=expert_offsets, ++ ) ++ else: ++ hidden_states = torch.ops.npu.npu_grouped_matmul( ++ x=[hidden_states], ++ weight=[layer.w2_weight.transpose(1, 2)], ++ bias=w2_bias, ++ split_item=2, ++ group_list_type=1, ++ group_type=0, ++ group_list=expert_tokens, ++ output_dtype=original_dtype, ++ )[0] + + final_hidden_states = torch.ops.npu.npu_moe_finalize_routing( + hidden_states, diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py -index 6696945..7854f7a 100644 +index 6696945d2c..7854f7a5f3 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/python/sglang/srt/layers/radix_linear_attention.py @@ -35,6 +35,27 @@ if TYPE_CHECKING: @@ -236,10 +2341,36 @@ index 6696945..7854f7a 100644 """ The Linear Attention Layer Implementation. diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py -index e83e4157fa..19ed3f244b 100644 +index e83e4157fa..1a9f98f12c 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py -@@ -81,6 +81,170 @@ class Sampler(nn.Module): +@@ -58,6 +58,25 @@ _disable_aiter_greedy_sample = get_bool_env_var("SGLANG_DISABLE_AITER_GREEDY_SAM + if is_npu(): + import torch_npu + ++ ++def _is_ascend_910c() -> bool: ++ """Return whether the current NPU is an Ascend 910C-series chip. ++ ++ torch.npu.get_device_name() reports names like "Ascend910_9382" on 910C. ++ The async exponential-race path relies on uniform_() decomposing to ++ DSARandomUniform on the DSA_SQE core, which only exists on the 910C ++ series; other NPUs route uniform_() to AIV/AICore RNG kernels that would ++ contend with the model forward. Fail closed on any query error so ++ sampling stays on the stock torch.multinomial path. ++ """ ++ if not is_npu(): ++ return False ++ try: ++ return torch.npu.get_device_name().startswith("Ascend910_93") ++ except Exception: ++ return False ++ ++ + logger = logging.getLogger(__name__) + + SYNC_TOKEN_IDS_ACROSS_TP = get_bool_env_var("SYNC_TOKEN_IDS_ACROSS_TP") +@@ -81,6 +100,253 @@ class Sampler(nn.Module): # In RL on-policy mode, we use log_softmax to compute logprobs to match the trainer. self.use_log_softmax_logprob = self.rl_on_policy_target is not None self.use_ascend_backend = get_flags().sampling_backend == "ascend" @@ -253,16 +2384,94 @@ index e83e4157fa..19ed3f244b 100644 + # even when it overlaps the next forward (v1 ran the full exponential_() + # chain here and its AIV tail contended with forward). + # This is opt-in because it changes the random-number sequence, although it -+ # preserves the categorical sampling distribution. -+ self.enable_async_exponential = is_npu() and get_bool_env_var( ++ # preserves the categorical sampling distribution. The gate is ++ # additionally restricted to the Ascend 910C series: the side stream is ++ # AIV-free only if uniform_() runs on the DSA_SQE core (910C-only). ++ # On any other device the flag stays off and every batch falls back to ++ # the stock torch.multinomial path. ++ async_exponential_requested = is_npu() and get_bool_env_var( + "SGLANG_NPU_ASYNC_EXPONENTIAL" + ) ++ self.enable_async_exponential = ( ++ async_exponential_requested and _is_ascend_910c() ++ ) ++ if async_exponential_requested and not self.enable_async_exponential: ++ logger.warning( ++ "SGLANG_NPU_ASYNC_EXPONENTIAL is set but the current device is " ++ "not an Ascend 910C series NPU; async exponential-race sampling " ++ "is disabled, falling back to torch.multinomial" ++ ) + self._async_exponential_stream = None + self._async_exponential_event = None + self._async_exponential_u = None + self._async_exponential_pending = False + self._async_exp_min_bound = None + ++ # Optional fused Triton consumption kernel, replacing the multi-pass ++ # stock-exponential transform + RealDiv + ArgMaxV2 chain with one ++ # kernel that reads probs and u once. Same math, same torch.argmax ++ # first-occurrence tie-breaking. Fail-closed: any import/compile/smoke ++ # failure keeps the unfused chain. ++ self._exp_race_fused = None ++ if self.enable_async_exponential and get_bool_env_var( ++ "SGLANG_NPU_EXP_RACE_TRITON", "true" ++ ): ++ self._exp_race_fused = self._init_exp_race_fused() ++ ++ @staticmethod ++ def _init_exp_race_fused(): ++ """Resolve the fused kernel and smoke-test it against the unfused chain. ++ ++ The smoke shape (4, 10000) deliberately exceeds BLOCK_V=4096 so the ++ masked tail path is exercised, and plants boundary u values (0 and ++ 1-2^-24, the exact cap grid points) plus an exact score tie to check ++ first-occurrence tie-breaking. Any exception or mismatch disables the ++ fused path (fail-closed). ++ """ ++ try: ++ from sglang.srt.layers.exp_race_fused_triton import exp_race_argmax ++ ++ device = torch.device("npu") ++ u = torch.rand((4, 10000), dtype=torch.float32, device=device) ++ u[0, 0] = 0.0 ++ u[1, 0] = 1.0 - 2.0**-24 ++ probs = torch.rand((4, 10000), dtype=torch.float32, device=device) ++ probs = probs / probs.sum(dim=-1, keepdim=True) ++ # Exact score tie at the maximal score: identical (p, u) pair at ++ # two columns with u=0; the lower index must win, matching ++ # torch.argmax. ++ u[:, 50] = 0.0 ++ u[:, 60] = 0.0 ++ probs[:, 50] = 0.5 ++ probs[:, 60] = 0.5 ++ fused = exp_race_argmax(probs, u) ++ q = u.neg_().add_(1.0) ++ bound = torch.full( ++ (), ++ 1.0 - torch.finfo(torch.float32).eps / 2.0, ++ dtype=torch.float32, ++ device=device, ++ ) ++ torch.minimum(q, bound, out=q) ++ q.log_().neg_() ++ reference = torch.div(probs, q).argmax(dim=-1).view(-1).to(torch.int32) ++ if not torch.equal(fused, reference): ++ raise RuntimeError( ++ f"fused exp-race smoke test mismatch: {fused} vs {reference}" ++ ) ++ logger.info( ++ "Enabled fused Triton exp-race consumption kernel " ++ "(SGLANG_NPU_EXP_RACE_TRITON)" ++ ) ++ return exp_race_argmax ++ except Exception as e: ++ logger.warning( ++ "Fused Triton exp-race kernel unavailable (%s); falling back " ++ "to the unfused stock-exponential chain", ++ e, ++ ) ++ return None ++ + def can_prepare_async_exponential( + self, sampling_info: SamplingBatchInfo + ) -> bool: @@ -393,6 +2602,11 @@ index e83e4157fa..19ed3f244b 100644 + # inf/NaN edge. + # All transform ops run here on the main stream; the side stream stays + # pure DSA uniform with zero AIV kernels. ++ if self._exp_race_fused is not None: ++ # Same math as the chain below in one Triton kernel (single pass ++ # over probs and u). Non-destructive on u; returns int32 [B]. ++ return self._exp_race_fused(probs, u) ++ + bound = self._async_exp_min_bound + if bound is None or bound.dtype != u.dtype or bound.device != u.device: + bound = torch.full( @@ -410,7 +2624,7 @@ index e83e4157fa..19ed3f244b 100644 def _preprocess_logits( self, logits: torch.Tensor, sampling_info: SamplingBatchInfo -@@ -115,6 +279,9 @@ class Sampler(nn.Module): +@@ -115,6 +381,9 @@ class Sampler(nn.Module): to get the unique seed for each position. """ logits = logits_output.next_token_logits @@ -420,13 +2634,22 @@ index e83e4157fa..19ed3f244b 100644 # Preprocess logits (custom processors and NaN handling) logits = self._preprocess_logits(logits, sampling_info) -@@ -180,9 +347,13 @@ class Sampler(nn.Module): +@@ -178,11 +447,20 @@ class Sampler(nn.Module): + logprobs = logprobs_via_logsoftmax_kernel + else: # Standard path: do softmax and sample from probs. - logits.div_(sampling_info.temperatures) - +- logits.div_(sampling_info.temperatures) +- - # In-place op to save memory - logits[:] = torch.softmax(logits, dim=-1) - probs = logits ++ # All temperature == 1.0 makes div_ an IEEE bitwise identity ++ # (x/1.0 == x), so skip the RealDiv. The flag is maintained ++ # host-side by SamplingBatchInfo; older versions without the ++ # field default to False (no skip) via getattr. ++ if not getattr(sampling_info, "temperatures_all_one", False): ++ logits.div_(sampling_info.temperatures) ++ + # Do not write the softmax output back into logits: the + # write-back costs a full-matrix TensorMove pass (~0.2 ms on + # NPU at bs=128 x vocab=248320 fp32). logits (now x/T) is not @@ -437,7 +2660,7 @@ index e83e4157fa..19ed3f244b 100644 batch_next_token_ids = self._sample_from_probs( probs, sampling_info, positions, simple_sampling_case -@@ -191,8 +362,9 @@ class Sampler(nn.Module): +@@ -191,8 +469,9 @@ class Sampler(nn.Module): logprobs = ( logprobs_via_logsoftmax_kernel if logprobs_via_logsoftmax_kernel is not None @@ -448,7 +2671,7 @@ index e83e4157fa..19ed3f244b 100644 del probs # Attach logprobs to logits_output (in-place modification) -@@ -206,6 +378,7 @@ class Sampler(nn.Module): +@@ -206,6 +485,7 @@ class Sampler(nn.Module): token_ids_logprobs, sampling_info, batch_next_token_ids, @@ -456,7 +2679,7 @@ index e83e4157fa..19ed3f244b 100644 ) self._sync_token_ids_across_tp(batch_next_token_ids, sampling_info) -@@ -225,11 +398,13 @@ class Sampler(nn.Module): +@@ -225,11 +505,13 @@ class Sampler(nn.Module): Handles both simple (direct multinomial) and complex (top-k/top-p/min-p) cases. """ if simple_sampling_case: @@ -475,7 +2698,7 @@ index e83e4157fa..19ed3f244b 100644 else: backend = get_flags().sampling_backend if backend == "flashinfer": -@@ -302,7 +477,11 @@ class Sampler(nn.Module): +@@ -302,7 +584,11 @@ class Sampler(nn.Module): probabilities, sampling_info.sampling_seed, positions ).view(-1) else: @@ -488,7 +2711,7 @@ index e83e4157fa..19ed3f244b 100644 return batch_next_token_ids.to(torch.int32) else: assert ( -@@ -353,9 +532,14 @@ class Sampler(nn.Module): +@@ -353,16 +639,31 @@ class Sampler(nn.Module): token_ids_logprobs: List[List[int]], sampling_info: SamplingBatchInfo, batch_next_token_ids: torch.Tensor, @@ -505,7 +2728,6 @@ index e83e4157fa..19ed3f244b 100644 # Attach logprobs to logits_output (in-place modification) if any(x > 0 for x in top_logprobs_nums): -@@ -362,7 +546,17 @@ class Sampler(nn.Module): - ( - logits_output.next_token_top_logprobs_val, - logits_output.next_token_top_logprobs_idx, @@ -527,7 +2749,7 @@ index e83e4157fa..19ed3f244b 100644 if any(x is not None for x in token_ids_logprobs): ( -@@ -372,10 +567,19 @@ class Sampler(nn.Module): +@@ -372,10 +673,19 @@ class Sampler(nn.Module): logprobs, token_ids_logprobs, no_copy_to_cpu=True ) @@ -551,6 +2773,19 @@ index e83e4157fa..19ed3f244b 100644 def _sync_token_ids_across_tp( self, batch_next_token_ids: torch.Tensor, sampling_info: SamplingBatchInfo +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index bfdd72bd54..5f3840e359 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -3740,7 +3740,7 @@ class Scheduler( + self.draft_worker.clear_cache_pool() + + if empty_cache: +- current_platform.empty_cache() ++ torch.npu.empty_cache() + # Per-DP-group leader logs once: ranks within a DP group are + # state-synchronous, but DP groups may diverge. + if self.metrics_reporter.is_stats_logging_rank: diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 3db49daf82..fdc5f40d81 100644 --- a/python/sglang/srt/managers/tp_worker.py @@ -687,8 +2922,52 @@ index 031ab4fe87..93e14ff483 100644 if self.forward_mode.is_decode(): # 3 * N if ( +diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py +index ea17a25587..22586658da 100644 +--- a/python/sglang/srt/model_executor/model_runner.py ++++ b/python/sglang/srt/model_executor/model_runner.py +@@ -610,6 +610,28 @@ class ModelRunner(ModelRunnerKVCacheMixin): + self._model_update_group = {} + self._weights_send_group = {} + ++ # profile้‡‡้›† ++ import os ++ import torch_npu ++ if os.environ.get('ROLLOUT_PROFILE', "false") == "true": ++ # Initialize profiler ++ import torch_npu ++ experimental_config = torch_npu.profiler._ExperimentalConfig( ++ profiler_level=torch_npu.profiler.ProfilerLevel.Level1, ++ ) ++ self.profiler_npu = torch_npu.profiler.profile( ++ activities=[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU], ++ with_modules=os.environ.get('WITH_MODULES', "false") == "true", ++ profile_memory=os.environ.get('WITH_MEMORY', "false") == "true", ++ record_shapes=os.environ.get('WITH_SHAPE', "false") == "true", ++ with_stack=os.environ.get('WITH_STACK', "false") == "true", ++ experimental_config=experimental_config, ++ # ่ทณ่ฟ‡ๅ‰29ๆญฅ๏ผŒwarmupไธ€ๆญฅ๏ผŒ้‡‡้›†30ๆญฅ๏ผŒ้‡ๅค1ๆฌกใ€‚ ++ schedule=torch_npu.profiler.schedule(wait=29, warmup=0, active=10, repeat=1), ++ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(os.environ.get('ROLLOUT_PROFILE_PATH'), analyse_flag=False) # ้‡‡้›†ๆ•ฐๆฎไฟๅญ˜่ทฏๅพ„๏ผŒๆ˜ฏๅฆๅœจ็บฟ่งฃๆž ++ ) ++ self.profiler_npu.start() ++ + def _build_model_config( + self, server_args, model_path=None, model_revision=None, is_draft_model=False + ): +@@ -3090,6 +3112,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): + if self.server_args.elastic_ep_backend is not None: + self.maybe_recover_ep_ranks() + ++ import os ++ if os.environ.get('ROLLOUT_PROFILE', "false") == "true": ++ self.profiler_npu.step() # ้ฉฑๅŠจ schedule๏ผŒๅฏน้ƒจๅˆ†decode step่ฟ›่กŒ้‡‡้›† ++ + return output + + def _maybe_execute_deferred_mamba_cow_and_clear( diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py -index d5aac381dc..e710d83b05 100644 +index d5aac381dc..c1b2966e84 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -37,6 +37,7 @@ from sglang.srt.distributed import ( @@ -753,7 +3032,33 @@ index d5aac381dc..e710d83b05 100644 if shared_output is not None: final_hidden_states.add_(shared_output) -@@ -991,7 +990,14 @@ class Qwen2MoeForCausalLM(nn.Module): +@@ -494,6 +493,15 @@ class Qwen2MoeSparseMoeBlock(nn.Module): + ) -> torch.Tensor: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) ++ # [moe_whole_process/v1] enqueue ๆขๅบ๏ผšๅ…ˆ alt ๆต้•ฟ้“พ๏ผˆrouter+routed ++ # experts๏ผ‰๏ผŒๅ†ไธปๆต shared expertใ€‚L0 profiling ่ฏๅฎžๅŽŸ้กบๅบไธ‹ shared ๆฎต ++ # ๏ผˆๅŒบๅŸŸ้ ๅ‰้ƒจๅˆ†๏ผ‰ๅœจ replay ๆ—ถ็‹ฌๅ ่ฎพๅค‡ๅ…ˆ่ท‘ใ€ๆœช่ขซ routed ้“พๆŽฉ็›–๏ผ›ๆขๅบๅŽ ++ # ้•ฟ้“พๅ…ˆ่ตทๆญฅ๏ผŒshared ็š„ๅฐ GEMM ๅฏๅกซ routed ๅ‰ๆฎตๅฐ kernel ็š„็ฉบ็ช—ใ€‚ ++ # ไพ่ต–่พนไธๅ˜๏ผˆalt ไปๅช็ญ‰ pre-MoE ไธปๆต็‚น๏ผŒๆฑ‡ๅˆไป main ็ญ‰ alt๏ผ‰๏ผŒ ++ # ไธคๆ”ฏๅฏน hidden_states ๅ‡ๅช่ฏปใ€clone ไปๅœจไธปๆตโ€”โ€”ๆ•ฐๅ€ผ้€ไฝ็ญ‰ไปทใ€‚ ++ with torch.cuda.stream(self.alt_stream): ++ router_output = self._forward_router_experts(hidden_states) ++ + shared_output = ( + self._forward_shared_experts( + hidden_states.clone(), apply_gate=not use_fused_gate +@@ -518,9 +526,6 @@ class Qwen2MoeSparseMoeBlock(nn.Module): + staged = True + # ===== END TO BE REFACTORED ==== + +- with torch.cuda.stream(self.alt_stream): +- router_output = self._forward_router_experts(hidden_states) +- + current_stream.wait_stream(self.alt_stream) + + if staged and unstage_shared_expert_add() is None: +@@ -991,7 +996,14 @@ class Qwen2MoeForCausalLM(nn.Module): self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config @@ -770,18 +3075,34 @@ index d5aac381dc..e710d83b05 100644 config, quant_config, diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py -index fed78cf..11cbadd 100644 +index fed78cf888..806eec9ffb 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py -@@ -35,6 +35,7 @@ from sglang.srt.configs.qwen3_5 import ( +@@ -15,6 +15,7 @@ + """Inference-only Qwen3.5 model and Qwen3.5 MoE model compatible with HuggingFace weights.""" + + import logging ++import os + from functools import lru_cache + from typing import Iterable, Optional, Set, Tuple, Union + +@@ -35,6 +36,15 @@ from sglang.srt.configs.qwen3_5 import ( # Distributed from sglang.srt.distributed import get_pp_group +from sglang.srt.environ import envs ++ ++# MoE weight L2 prefetch (SGLANG_NPU_MOE_PREFETCH): launched after each layer's ++# prepare_attn, drained at step end; no torch_npu import at module top level, no-op when disabled. ++from sglang.srt.hardware_backend.npu.moe_weight_prefetch import ( ++ moe_prefetch_emit, ++ moe_prefetch_register_model, ++ moe_prefetch_step_drain, ++) from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation -@@ -64,7 +65,10 @@ from sglang.srt.layers.parameter import ( +@@ -64,7 +74,10 @@ from sglang.srt.layers.parameter import ( ) from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention @@ -793,7 +3114,116 @@ index fed78cf..11cbadd 100644 from sglang.srt.layers.rotary_embedding import get_rope from sglang.srt.layers.utils import PPMissingLayer, get_layer_id from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding -@@ -229,23 +233,6 @@ class Qwen3_5GatedDeltaNet(nn.Module): +@@ -127,6 +140,16 @@ _qknorm_use_alt_stream = _is_cuda or ( + ) + _is_amx_available = cpu_has_amx_support() + ++# GDN_QKVZBA_PACK: pack the in_proj_qkvz / in_proj_ba weights of GDN layers along N ++# into a single [N_pack, K] block (N_pack zero-padded up to a multiple of 16 for the ++# row-stride alignment contract of downstream fused ops), replacing two GEMMs with one ++# stock linear; column-wise equivalent (pad output columns are never consumed). ++# On by default; =0 reverts. NPU + unquantized bf16/fp16 only, guarded at runtime. ++_gdn_qkvzba_pack_env = get_bool_env_var("SGLANG_NPU_GDN_QKVZBA_PACK", "true") ++# Pack only for small M: at large M the ragged-tiling penalty of the wide N ++# (not a multiple of the 256 tile) outweighs the saved ba GEMM. ++_gdn_qkvzba_pack_max_m = int(os.getenv("SGLANG_NPU_GDN_QKVZBA_PACK_MAX_M", "256")) ++ + cached_get_processor = lru_cache(get_processor) + + +@@ -153,6 +176,91 @@ if _is_npu: + split_qkvgate_gemma_rmsnorm_rope, + ) + ++ # NPU ๅฏ็”จ่žๅˆ qkvzba split kernel๏ผšๆ”น็”จ sgl_kernel_npu ็‰ˆ ++ # ๏ผˆgraph ๅฃๅพ„ๆฏ” sglang ไพง triton_gdn_fused_proj ็‰ˆๅฟซ๏ผŒไธ”้€ bit ็ญ‰ไปท๏ผ‰ ++ from sgl_kernel_npu.fla.utils import ( ++ fused_qkvzba_split_reshape_cat_contiguous, ++ ) ++ ++ # ๅ…จๆณจๆ„ๅŠ›ๆฎต decode ่žๅˆ kernel๏ผˆfull_attention A1b/A2๏ผ‰โ€”โ€” ++ # A1b split+norm+rope+KV scatter ๅˆไธ€ใ€A2 sigmoid_mul ่žๅˆ๏ผ›่ฐƒ็”จ็‚นๅธฆๅฝข็Šถ ++ # ๅฎˆๅซ๏ผŒๆœชๅ‘ฝไธญๅทฒ้ชŒ่ฏๅฝข็Šถ่‡ชๅŠจๅ›ž้€€ stockใ€‚ ++ from sglang.srt.hardware_backend.npu.attention.full_attention_fusion_npu import ( ++ fa_sigmoid_mul, ++ fa_sigmoid_mul_supported, ++ fa_split_qkvgate_scatter, ++ fa_split_qkvgate_scatter_supported, ++ fa_v4_scatter_context, ++ ) ++ ++ # TP AscendC fusion wiring (master switch SGLANG_NPU_TP_ASCENDC_FUSION, default ++ # off; candidates 2/3 are retired, their interfaces are permanently disabled ++ # stubs kept only for import compatibility) ++ from sglang.srt.hardware_backend.npu.tp_ascendc_fusion_npu import ( ++ tp_fused_qkvzba_conv1d_shape_supported, ++ tp_norm_qkv_scatter_context, ++ tp_norm_qkv_scatter_shape_supported, ++ tp_sigmoid_mul_mm_runtime_ok, ++ tp_sigmoid_mul_mm_shape_supported, ++ ) ++ ++ ++def _gdn_qkvzba_packable(proj_qkvz, proj_ba) -> bool: ++ """Runtime guard for GDN_QKVZBA_PACK (any failure -> fall back to two GEMMs). ++ ++ Only plain unquantized 2D weights are allowed: quantized formats such as FP8 ++ have scale layouts incompatible with the packed GEMM and are excluded ++ (detected via the quant_method class name). ++ """ ++ wq = proj_qkvz.weight ++ wb = proj_ba.weight ++ if wq.dim() != 2 or wb.dim() != 2 or wq.shape[1] != wb.shape[1]: ++ return False ++ if wq.dtype != wb.dtype or wq.dtype not in (torch.bfloat16, torch.float16): ++ return False ++ if proj_qkvz.bias is not None or proj_ba.bias is not None: ++ return False ++ for proj in (proj_qkvz, proj_ba): ++ qm = getattr(proj, "quant_method", None) ++ if qm is not None and qm.__class__.__name__ != "UnquantizedLinearMethod": ++ return False ++ return True ++ ++ ++def _gdn_qkvzba_pack_materialize(proj_qkvz, proj_ba) -> torch.Tensor: ++ """Materialize the packed [N_pack, K] weight and rebind both modules' ++ weight.data to row slices of it. ++ ++ N_pack = (N_qkvz+N_ba) rounded up to a multiple of 16, so the packed GEMM ++ output keeps a 16-aligned physical row stride as required by downstream fused ++ ops consuming row-stride views; pad rows are zero and pad output columns have ++ no consumer (the ba slice excludes the pad tail rows). ++ ++ After rebinding, in-place narrow+copy_ writes from weight loaders / RL weight ++ sync land directly on the packed buffer; if some path replaces weight.data ++ wholesale (pointer drift), the caller's data_ptr check re-materializes on the ++ next forward. ++ """ ++ wq = proj_qkvz.weight ++ wb = proj_ba.weight ++ n_qkvz = wq.shape[0] ++ n_ba = wb.shape[0] ++ n_pack = (n_qkvz + n_ba + 15) // 16 * 16 ++ if n_pack == n_qkvz + n_ba: ++ packed = torch.cat([wq.detach(), wb.detach()], dim=0) ++ else: ++ packed = torch.cat( ++ [ ++ wq.detach(), ++ wb.detach(), ++ wq.detach().new_zeros(n_pack - n_qkvz - n_ba, wq.shape[1]), ++ ], ++ dim=0, ++ ) ++ wq.data = packed[:n_qkvz] ++ wb.data = packed[n_qkvz : n_qkvz + n_ba] ++ return packed ++ + + class Qwen3_5GatedDeltaNet(nn.Module): + def __init__( +@@ -229,26 +337,19 @@ class Qwen3_5GatedDeltaNet(nn.Module): self._bind_packed_weight_loaders(self.in_proj_qkvz) self._bind_packed_weight_loaders(self.in_proj_ba) @@ -813,11 +3243,24 @@ index fed78cf..11cbadd 100644 - self.attn_tp_rank, - ), - ) -- ++ # GDN_QKVZBA_PACK: enable flag and lazily materialized pack cache. Packing ++ # only reorganizes compute โ€” parameters, names, loaders and checkpoint ++ # mappings of the two modules are unchanged. ++ self._qkvzba_pack_enabled = _is_npu and _gdn_qkvzba_pack_env ++ self._packed_qkvzba_weight: Optional[torch.Tensor] = None + # State parameters ++ # dt_bias is materialized in fp32 like A_log: gating/recurrent kernels widen ++ # to fp32 losslessly anyway and the AscendC recurrent host requires fp32, so ++ # with fp32 at the source the wrapper's .float() becomes a no-op. Loading and ++ # RL weight sync are copy_-based (implicit cast), same as A_log. self.dt_bias = nn.Parameter( - torch.ones(self.num_v_heads // self.attn_tp_size), -@@ -275,6 +262,28 @@ class Qwen3_5GatedDeltaNet(nn.Module): +- torch.ones(self.num_v_heads // self.attn_tp_size), ++ torch.ones(self.num_v_heads // self.attn_tp_size, dtype=torch.float32), + ) + self.A_log = nn.Parameter( + torch.empty(self.num_v_heads // self.attn_tp_size, dtype=torch.float32), +@@ -275,6 +376,28 @@ class Qwen3_5GatedDeltaNet(nn.Module): dt_bias=self.dt_bias, ) @@ -846,7 +3289,180 @@ index fed78cf..11cbadd 100644 self.norm = RMSNormGated( self.head_v_dim, eps=self.layer_norm_epsilon, -@@ -608,7 +617,11 @@ class Qwen3_5LinearDecoderLayer(nn.Module): +@@ -301,6 +424,18 @@ class Qwen3_5GatedDeltaNet(nn.Module): + prefix=add_prefix("out_proj", prefix), + ) + ++ # TP_FUSION candidate 1: bake the init-time shape guard for ++ # fused_qkvzba_conv1d (split+causal_conv1d in one kernel); head counts are ++ # per-rank (post-TP) values. On guard miss or master switch off, forward ++ # uses the original split+causal_conv1d path. ++ self._tp_fused_qkvzba_ok = _is_npu and tp_fused_qkvzba_conv1d_shape_supported( ++ triton.cdiv(self.num_k_heads, self.attn_tp_size), ++ triton.cdiv(self.num_v_heads, self.attn_tp_size), ++ self.head_k_dim, ++ self.head_v_dim, ++ self.conv_kernel_size, ++ ) ++ + @staticmethod + def _override_weight_loader(param, loader): + """Robustly override loader for: +@@ -493,10 +628,49 @@ class Qwen3_5GatedDeltaNet(nn.Module): + projected_states_ba, _ = self.in_proj_ba(hidden_states) + current_stream.wait_stream(self.alt_stream) + else: +- projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states) +- projected_states_ba, _ = self.in_proj_ba(hidden_states) ++ if ( ++ self._qkvzba_pack_enabled ++ and seq_len <= _gdn_qkvzba_pack_max_m ++ and _gdn_qkvzba_packable(self.in_proj_qkvz, self.in_proj_ba) ++ ): ++ # GDN_QKVZBA_PACK: one packed GEMM whose [M, N_pack] output is ++ # sliced into two row-stride views (contiguous within each row); ++ # all consumers take them with zero extra copies โ€” the split kernel ++ # reads the row stride via stride(0), the op1 fused op accepts ++ # row-stride views natively, and the fix_query_key_value_ordering ++ # fallback copies anyway. The ba slice must explicitly exclude the ++ # pad tail columns (N_pack may exceed N_qkvz+N_ba). Graph capture ++ # bakes this branch per fixed batch size. ++ projected_states_qkvzba = nn.functional.linear( ++ hidden_states, self._get_packed_qkvzba_weight() ++ ) ++ n_qkvz = self.in_proj_qkvz.weight.shape[0] ++ n_ba = self.in_proj_ba.weight.shape[0] ++ projected_states_qkvz = projected_states_qkvzba[:, :n_qkvz] ++ projected_states_ba = projected_states_qkvzba[:, n_qkvz : n_qkvz + n_ba] ++ else: ++ projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states) ++ projected_states_ba, _ = self.in_proj_ba(hidden_states) + return projected_states_qkvz, projected_states_ba + ++ def _get_packed_qkvzba_weight(self) -> torch.Tensor: ++ """GDN_QKVZBA_PACK: return the packed [N_qkvz+N_ba, K] weight. ++ ++ Materialized lazily on the first forward after weights are loaded; every ++ later forward revalidates via data_ptr โ€” in-place writes (loader / RL weight ++ sync) keep the pointer so the buffer stays in sync, while wholesale .data ++ replacement moves the pointer and triggers re-materialization. ++ """ ++ packed = self._packed_qkvzba_weight ++ if ( ++ packed is not None ++ and self.in_proj_qkvz.weight.data_ptr() == packed.data_ptr() ++ ): ++ return packed ++ packed = _gdn_qkvzba_pack_materialize(self.in_proj_qkvz, self.in_proj_ba) ++ self._packed_qkvzba_weight = packed ++ return packed ++ + def forward( + self, + hidden_states: torch.Tensor, +@@ -512,48 +686,67 @@ class Qwen3_5GatedDeltaNet(nn.Module): + hidden_states + ) + +- if ( +- self.num_v_heads // self.num_k_heads in [1, 2, 4] +- and not _is_cpu +- and not _is_npu +- ): +- mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat_contiguous( +- projected_states_qkvz, +- projected_states_ba, +- triton.cdiv(self.num_k_heads, self.attn_tp_size), +- triton.cdiv(self.num_v_heads, self.attn_tp_size), +- self.head_k_dim, +- self.head_v_dim, ++ # TP_FUSION candidate 1: on decode with the init guard hit, pass the qkvz/ba ++ # projections to the GDN backend as a tuple โ€” split+causal_conv1d run inside ++ # the single torch.ops.npu.fused_qkvzba_conv1d kernel, returning ++ # (core_attn_out, z); on runtime guard miss the backend falls back to local ++ # split+stock conv, still returning a tuple. Row-stride views from the ++ # packed path can be passed directly (host contract: contiguous within a row ++ # + 16-aligned row stride; the packed N is already padded). ++ core_attn_out = None ++ z = None ++ if self._tp_fused_qkvzba_ok and forward_batch.forward_mode.is_decode(): ++ core_attn_out, z = self.attn( ++ forward_batch, ++ mixed_qkv=(projected_states_qkvz, projected_states_ba), ++ a=None, ++ b=None, + ) +- elif _is_cpu and _is_amx_available: +- mixed_qkv, z, b, a = ( +- torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_contiguous_cpu( ++ ++ if core_attn_out is None: ++ # NPU uses the fused split kernel too (the import above is rebound per ++ # platform to the sgl_kernel_npu version) ++ if ( ++ self.num_v_heads // self.num_k_heads in [1, 2, 4] ++ and not _is_cpu ++ ): ++ mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat_contiguous( + projected_states_qkvz, + projected_states_ba, +- self.num_k_heads // self.attn_tp_size, +- self.num_v_heads // self.attn_tp_size, ++ triton.cdiv(self.num_k_heads, self.attn_tp_size), ++ triton.cdiv(self.num_v_heads, self.attn_tp_size), + self.head_k_dim, + self.head_v_dim, + ) +- ) +- else: +- query, key, value, z, b, a = self.fix_query_key_value_ordering( +- projected_states_qkvz, projected_states_ba +- ) +- b = b.contiguous() +- a = a.contiguous() ++ elif _is_cpu and _is_amx_available: ++ mixed_qkv, z, b, a = ( ++ torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_contiguous_cpu( ++ projected_states_qkvz, ++ projected_states_ba, ++ self.num_k_heads // self.attn_tp_size, ++ self.num_v_heads // self.attn_tp_size, ++ self.head_k_dim, ++ self.head_v_dim, ++ ) ++ ) ++ else: ++ query, key, value, z, b, a = self.fix_query_key_value_ordering( ++ projected_states_qkvz, projected_states_ba ++ ) ++ b = b.contiguous() ++ a = a.contiguous() + +- query, key, value = map( +- lambda x: x.reshape(x.shape[0], -1), (query, key, value) +- ) +- mixed_qkv = torch.cat((query, key, value), dim=-1) ++ query, key, value = map( ++ lambda x: x.reshape(x.shape[0], -1), (query, key, value) ++ ) ++ mixed_qkv = torch.cat((query, key, value), dim=-1) + +- core_attn_out = self.attn( +- forward_batch, +- mixed_qkv=mixed_qkv, +- a=a, +- b=b, +- ) ++ core_attn_out = self.attn( ++ forward_batch, ++ mixed_qkv=mixed_qkv, ++ a=a, ++ b=b, ++ ) + + z_shape_og = z.shape + # reshape input data into 2D tensor +@@ -608,7 +801,11 @@ class Qwen3_5LinearDecoderLayer(nn.Module): quant_config=quant_config, alt_stream=( alt_stream @@ -859,7 +3475,20 @@ index fed78cf..11cbadd 100644 else None ), prefix=add_prefix("mlp", prefix.replace(".linear_attn", "")), -@@ -824,7 +837,11 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): +@@ -671,6 +868,12 @@ class Qwen3_5LinearDecoderLayer(nn.Module): + ) + ) + ++ # MoE weight L2 prefetch: launched from the target itself โ€” after this ++ # layer's prepare_attn and before attention, a point after all inter-layer ++ # communication on both TP and EP paths. GDN layers only; the main stream ++ # does not wait, everything is drained at step end. ++ moe_prefetch_emit(self, hidden_states) ++ + if not forward_batch.forward_mode.is_idle(): + hidden_states = self.linear_attn( + hidden_states, +@@ -824,7 +1027,11 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): quant_config=quant_config, alt_stream=( alt_stream @@ -872,7 +3501,237 @@ index fed78cf..11cbadd 100644 else None ), prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), -@@ -1207,6 +1224,11 @@ class Qwen3_5ForCausalLM(nn.Module): +@@ -863,6 +1070,26 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): + + self.alt_stream = alt_stream + ++ # full_attention A1b๏ผšTP ๅธƒๅฑ€ๅฝข็Šถๅฎˆๅซๅœจ init ๆ—ถไธ€ๆฌก็ฎ—ๅฎš ++ # (per-rank q heads in {1,2,4} / 1 kv head / head_dim=256 / rope=64 / gated) ++ self._fa_v4_shape_ok = _is_npu and fa_split_qkvgate_scatter_supported( ++ self.num_heads, ++ self.num_kv_heads, ++ self.head_dim, ++ int(self.head_dim * self.partial_rotary_factor), ++ self.attn_output_gate, ++ ) ++ ++ # TP_FUSION candidates 2/3: bake init-time shape guards for the ++ # norm+qkv_proj+KV scatter triple fusion and the sigmoid_mul+o_proj fusion ++ # (guard miss or master switch off -> original path at runtime) ++ self._tp_fa_norm_qkv_scatter_ok = _is_npu and tp_norm_qkv_scatter_shape_supported( ++ self ++ ) ++ self._tp_sigmoid_mul_mm_ok = _is_npu and tp_sigmoid_mul_mm_shape_supported( ++ self.o_proj ++ ) ++ + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: +@@ -979,6 +1206,30 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): + if self.attn.layer_id == (self.config.full_attention_interval - 1): + self.rotary_emb.get_cos_sin_with_position(positions) + ++ # full_attention A1b๏ผšdecode ไธ”ๅ‘ฝไธญๅทฒ้ชŒ่ฏๅฝข็Šถๆ—ถ๏ผŒsplit+norm+rope ไธŽ KV ++ # scatter ๅˆไธ€โ€”โ€”k/v ็”ฑ kernel ๆŒ‰ out_cache_loc ็›ดๆŽฅๅ†™ๅ…ฅ KV cache๏ผŒ่ฟ”ๅ›ž ++ # k/v=None๏ผŒself_attention ๆฎๆญคไปฅ save_kv_cache=False ่ตฐ backend๏ผˆ่ทณ่ฟ‡ ++ # stock ็š„ไธคๆฌก npu_scatter_nd_update_๏ผ‰ใ€‚ๅฎˆๅซไปปไธ€็Žฏ่Š‚ๆœชๅ‘ฝไธญๅˆ™ๅ›ž้€€ stock ++ # split๏ผ›SGLANG_NPU_FULL_ATTN_FUSION_DEBUG=1 ๅฏๆŒ‰ๅฑ‚ๆ‰“ๅฐๆœชๅ‘ฝไธญๅŽŸๅ› ใ€‚ ++ fa_ctx = fa_v4_scatter_context(self, forward_batch, qkv) ++ if fa_ctx is not None: ++ q, _, _, gate = fa_split_qkvgate_scatter( ++ qkv, ++ self.rotary_emb.position_sin, ++ self.rotary_emb.position_cos, ++ self.q_size, ++ self.kv_size, ++ self.head_dim, ++ int(self.head_dim * self.partial_rotary_factor), ++ eps=self.q_norm.variance_epsilon, ++ q_weight=self.q_norm.weight, ++ k_weight=self.k_norm.weight, ++ kbuf=fa_ctx[0], ++ vbuf=fa_ctx[1], ++ loc=fa_ctx[2], ++ ) ++ return q, None, None, gate ++ + q, k, v, gate = split_qkvgate_gemma_rmsnorm_rope( + qkv, + self.rotary_emb.position_sin, +@@ -1026,18 +1277,108 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): + forward_batch=forward_batch, + ) + +- attn_output = self.attn(q, k, v, forward_batch) ++ # full_attention A1b๏ผšk is None โŸบ KV ๅทฒๅœจ่žๅˆ kernel ๅ†… scatter ่ฟ› cache๏ผŒ ++ # backend ไพง่ทณ่ฟ‡ set_kv_buffer๏ผ›ๅ…ถไฝ™่ทฏๅพ„ k ๆ’้ž None๏ผŒsave_kv_cache=True ++ # ไธŽ stock ้ป˜่ฎคๅ€ผไธ€่‡ด ++ attn_output = self.attn(q, k, v, forward_batch, save_kv_cache=k is not None) + + if self.attn_output_gate: + if not _is_npu: + attn_output = fused_sigmoid_mul(attn_output, gate, inplace=True) + else: + gate_val = gate.reshape(gate.shape[0], -1) if gate.ndim == 3 else gate +- attn_output.mul_(torch.sigmoid(gate_val)) ++ # TP_FUSION candidate 3: sigmoid_mul + o_proj GEMM in one kernel ++ # (outputs a partial sum, no bias no all-reduce, same return ++ # semantics as o_proj(reduce_results=False), so it is returned ++ # directly as this layer's output); guard miss falls back to the ++ # A2/stock path ++ if self._tp_sigmoid_mul_mm_ok and tp_sigmoid_mul_mm_runtime_ok( ++ attn_output, gate_val ++ ): ++ return torch.ops.npu.fused_sigmoid_mul_mm( ++ attn_output, gate_val, self.o_proj.weight ++ ) ++ # full_attention A2๏ผšsigmoid+mul ่žๅˆๅ• kernel๏ผŒไธŽ stock ๅŒ op ++ # ่ˆๅ…ฅ่ทฏๅพ„้€ไฝไธ€่‡ด๏ผ›ไธๆปก่ถณๅฎˆๅซ๏ผˆ้ž bf16/้ž่ฟž็ปญ๏ผ‰ๅ›ž้€€ stock ++ if fa_sigmoid_mul_supported(attn_output, gate_val): ++ attn_output = fa_sigmoid_mul(attn_output, gate_val) ++ else: ++ attn_output.mul_(torch.sigmoid(gate_val)) + + output, _ = self.o_proj(attn_output) + return output + ++ def _self_attention_tp_fused( ++ self, ++ positions: torch.Tensor, ++ hidden_states: torch.Tensor, ++ residual: torch.Tensor, ++ forward_batch: ForwardBatch, ++ tp_ctx, ++ ) -> Tuple[torch.Tensor, torch.Tensor]: ++ """FA decode body of TP_FUSION candidate 2 (prepare's input_layernorm is ++ folded into the kernel). ++ ++ hidden_states/residual are the raw pre-prepare tensors (no add+norm ++ applied); returns (o_proj partial output, add_out โ€” the new residual ++ produced by the kernel) so forward can go straight to prepare_mlp. Caller ++ must have passed tp_norm_qkv_scatter_context (tp_ctx is not None). ++ """ ++ # Same as forward_prepare_npu: the first full-attention layer refreshes the ++ # sin/cos cache first ++ if self.attn.layer_id == (self.config.full_attention_interval - 1): ++ self.rotary_emb.get_cos_sin_with_position(positions) ++ ++ rope_dim = int(self.head_dim * self.partial_rotary_factor) ++ # The kernel requires contiguous [M,64] fp32; the NPU rope cache defaults to ++ # bf16, so widen-cast to fp32 (bf16 -> fp32 is exact) ++ sin = self.rotary_emb.position_sin.view(-1, rope_dim) ++ cos = self.rotary_emb.position_cos.view(-1, rope_dim) ++ if sin.dtype != torch.float32: ++ sin = sin.float() ++ cos = cos.float() ++ ++ kbuf, vbuf, loc = tp_ctx ++ add_out, q, gate = torch.ops.npu.fused_norm_qkv_proj_scatter( ++ hidden_states, ++ residual, ++ self.input_layernorm.weight, ++ self.qkv_proj.weight, ++ self.q_norm.weight, ++ self.k_norm.weight, ++ sin, ++ cos, ++ loc, ++ kbuf, ++ vbuf, ++ self.num_heads, ++ self.num_kv_heads, ++ self.head_dim, ++ rope_dim, ++ self.input_layernorm.variance_epsilon, ++ self.q_norm.variance_epsilon, ++ ) ++ # k/v were already written into the cache by the kernel at out_cache_loc ++ # (same contract as A1b: save_kv_cache=False) ++ attn_output = self.attn(q, None, None, forward_batch, save_kv_cache=False) ++ ++ # Candidate 3 (same guards as the self_attention tail): on hit, A2+o_proj ++ # are fused together; on miss, fall back to the A2/stock tail (gate is a ++ # contiguous 2D tensor straight from the kernel) ++ if self._tp_sigmoid_mul_mm_ok and tp_sigmoid_mul_mm_runtime_ok( ++ attn_output, gate ++ ): ++ output = torch.ops.npu.fused_sigmoid_mul_mm( ++ attn_output, gate, self.o_proj.weight ++ ) ++ elif fa_sigmoid_mul_supported(attn_output, gate): ++ attn_output = fa_sigmoid_mul(attn_output, gate) ++ output, _ = self.o_proj(attn_output) ++ else: ++ attn_output.mul_(torch.sigmoid(gate)) ++ output, _ = self.o_proj(attn_output) ++ return output, add_out ++ + def forward( + self, + positions: torch.Tensor, +@@ -1047,21 +1388,52 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): + captured_last_layer_outputs: Optional[list[torch.Tensor]] = None, + **kwargs, + ): +- hidden_states, residual = ( +- self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( ++ # TP_FUSION candidate 2: on decode with all runtime guards hit, skip the ++ # input_layernorm inside prepare โ€” add+norm is folded into ++ # fused_norm_qkv_proj_scatter, the raw hidden_states/residual go straight ++ # into the kernel, and add_out becomes the new residual for prepare_mlp. ++ # In this configuration (TP, no dp-attention, not the first layer, no ++ # pending AR fusion) prepare_attn is a no-op apart from the norm, so ++ # skipping it loses nothing; if any premise breaks, tp_ctx is None and the ++ # original path below runs. ++ tp_ctx = ( ++ tp_norm_qkv_scatter_context( ++ self, + hidden_states, + residual, + forward_batch, +- captured_last_layer_outputs=captured_last_layer_outputs, ++ captured_last_layer_outputs, + ) ++ if (_is_npu and self._tp_fa_norm_qkv_scatter_ok) ++ else None + ) ++ if tp_ctx is None: ++ hidden_states, residual = ( ++ self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( ++ hidden_states, ++ residual, ++ forward_batch, ++ captured_last_layer_outputs=captured_last_layer_outputs, ++ ) ++ ) + +- if not forward_batch.forward_mode.is_idle(): +- hidden_states = self.self_attention( +- positions=positions, +- hidden_states=hidden_states, +- forward_batch=forward_batch, ++ # MoE weight L2 prefetch: keep this after all inter-layer communication of ++ # the layer and before attention. When TP candidate 2 hits, prepare_attn is ++ # replaced by the equivalent fused path and the inter-layer AR was already ++ # done by the previous layer. ++ moe_prefetch_emit(self, hidden_states) ++ ++ if tp_ctx is not None: ++ hidden_states, residual = self._self_attention_tp_fused( ++ positions, hidden_states, residual, forward_batch, tp_ctx + ) ++ else: ++ if not forward_batch.forward_mode.is_idle(): ++ hidden_states = self.self_attention( ++ positions=positions, ++ hidden_states=hidden_states, ++ forward_batch=forward_batch, ++ ) + + # Fully Connected + hidden_states, residual = self.layer_communicator.prepare_mlp( +@@ -1207,6 +1579,11 @@ class Qwen3_5ForCausalLM(nn.Module): self._maybe_autodisable_shared_experts_fusion(config, quant_config) alt_stream = torch.cuda.Stream() if _is_cuda or _hip_use_alt_stream else None @@ -884,6 +3743,67 @@ index fed78cf..11cbadd 100644 # Embedding layer if self.pp_group.is_first_rank: +@@ -1252,6 +1629,13 @@ class Qwen3_5ForCausalLM(nn.Module): + + self.layers_to_capture = [] + ++ # MoE weight L2 prefetch: register each layer's MoE weights (layer_id -> ++ # w13/w2 plus the GDN target-layer flag) and create a dedicated prefetch ++ # stream outside capture; the is_nextn MTP draft model is not registered ++ # (its layers carry no launch flag, so draft capture never launches prefetch). ++ if not is_nextn: ++ moe_prefetch_register_model(self) ++ + def get_input_embeddings(self): + return self.embed_tokens + +@@ -1320,6 +1704,11 @@ class Qwen3_5ForCausalLM(nn.Module): + input_deepstack_embeds[:, sep : sep + self.hidden_size] + ) + ++ # MoE weight L2 prefetch: drain at step end โ€” the main stream joins the ++ # prefetch stream (capture legality + no cross-step backlog); no-op when ++ # disabled or when this pass launched nothing. ++ moe_prefetch_step_drain() ++ + # Return intermediate tensors for pipeline parallelism + if not self.pp_group.is_last_rank: + return PPProxyTensors( +diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py +index cfb22d8419..4f478f3539 100644 +--- a/python/sglang/srt/sampling/sampling_batch_info.py ++++ b/python/sglang/srt/sampling/sampling_batch_info.py +@@ -42,6 +42,12 @@ class SamplingBatchInfo: + + # Masking tensors for grammar-guided structured outputs + vocab_size: int ++ # Whether all requests use temperature == 1.0 (computed host-side in ++ # from_schedule_batch; merge_batch takes the AND, filter_batch subsetting ++ # preserves it). True means logits.div_(temperatures) is a bitwise ++ # identity and the sampler skips it. Must sit after the non-default ++ # fields above (dataclass field ordering). ++ temperatures_all_one: bool = False + grammars: Optional[List] = None + rids_int: Optional[torch.Tensor] = None + bootstrap_room_ids_int: Optional[torch.Tensor] = None +@@ -191,6 +197,9 @@ class SamplingBatchInfo: + need_top_p_sampling=any(r.sampling_params.top_p != 1.0 for r in reqs), + need_top_k_sampling=any(r.sampling_params.top_k != TOP_K_ALL for r in reqs), + need_min_p_sampling=any(r.sampling_params.min_p > 0 for r in reqs), ++ temperatures_all_one=all( ++ r.sampling_params.temperature == 1.0 for r in reqs ++ ), + vocab_size=vocab_size, + penalizer_orchestrator=penalizer_orchestrator, + has_custom_logit_processor=has_custom_logit_processor, +@@ -412,6 +421,7 @@ class SamplingBatchInfo: + self.need_top_p_sampling |= other.need_top_p_sampling + self.need_top_k_sampling |= other.need_top_k_sampling + self.need_min_p_sampling |= other.need_min_p_sampling ++ self.temperatures_all_one &= other.temperatures_all_one + + self.adjusted_merge_batch(other) + diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index daf9fa82e2..f00f718fd2 100644 --- a/python/sglang/srt/speculative/eagle_utils.py @@ -1109,7 +4029,7 @@ index ad98e59283..a5b2cdfd32 100644 @contextmanager diff --git a/test/registered/unit/layers/test_conv_weight_cache_refresh.py b/test/registered/unit/layers/test_conv_weight_cache_refresh.py new file mode 100644 -index 0000000000..ed2574df66 100644 +index 0000000000..ed2574df66 --- /dev/null +++ b/test/registered/unit/layers/test_conv_weight_cache_refresh.py @@ -0,0 +1,109 @@ @@ -1222,72 +4142,9 @@ index 0000000000..ed2574df66 100644 + +if __name__ == "__main__": + unittest.main() -diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py -index 9ad577b42e..390a3721c5 100755 ---- a/test/registered/unit/mem_cache/test_mamba_unittest.py -+++ b/test/registered/unit/mem_cache/test_mamba_unittest.py -@@ -419,6 +419,58 @@ class TestMamba(unittest.TestCase): - self.assertEqual(list(second_insert_events[0].token_ids), [5]) - self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash) - -+ def test_mamba_radix_cache_limited_partial_page_match_does_not_split(self): -+ page_size = 64 -+ tree = self._setup_minimal_mamba_radix_cache(page_size) -+ token_ids = array("q", range(page_size)) -+ -+ tree.insert( -+ InsertParams( -+ key=RadixKey(token_ids, None), -+ value=torch.arange(page_size), -+ mamba_value=torch.tensor([0]), -+ ) -+ ) -+ -+ match = tree.match_prefix( -+ MatchPrefixParams(key=RadixKey(token_ids, None, limit=page_size - 1)) -+ ) -+ -+ self.assertEqual(len(match.device_indices), 0) -+ self.assertEqual(self._non_root_key_lengths(tree), [page_size]) -+ -+ def _setup_minimal_mamba_radix_cache(self, page_size: int) -> MambaRadixCache: -+ tree = MambaRadixCache.__new__(MambaRadixCache) -+ tree.page_size = page_size -+ tree.mamba_cache_chunk_size = page_size -+ tree.disable = False -+ tree.device = torch.device("cpu") -+ tree.enable_kv_cache_events = False -+ tree.kv_event_queue = [] -+ tree.full_evictable_size_ = 0 -+ tree.mamba_evictable_size_ = 0 -+ tree.full_protected_size_ = 0 -+ tree.mamba_protected_size_ = 0 -+ -+ tree.root_node = TreeNode() -+ tree.root_node.key = RadixKey(array("q"), None) -+ tree.root_node.value = [] -+ tree.root_node.hash_value = [] -+ tree.root_node.full_lock_ref = 1 -+ tree.root_node.mamba_lock_ref = 1 -+ tree.full_lru_list = LRUList(mamba=False) -+ tree.mamba_lru_list = LRUList(mamba=True) -+ return tree -+ -+ def _non_root_key_lengths(self, tree: MambaRadixCache) -> list[int]: -+ lengths = [] -+ stack = list(tree.root_node.children.values()) -+ while stack: -+ node = stack.pop() -+ lengths.append(len(node.key)) -+ stack.extend(node.children.values()) -+ return lengths -+ - def _setup_tree_and_allocator(self, enable_kv_cache_events=False): - """Helper to create a MambaRadixCache with allocator for testing.""" - server_args = ServerArgs(model_path="dummy", page_size=1) diff --git a/test/registered/unit/layers/test_sampler_attach_logprobs.py b/test/registered/unit/layers/test_sampler_attach_logprobs.py new file mode 100644 -index 0000000000..9dfc930392 100644 +index 0000000000..9dfc930392 --- /dev/null +++ b/test/registered/unit/layers/test_sampler_attach_logprobs.py @@ -0,0 +1,156 @@ @@ -1447,3 +4304,6 @@ index 0000000000..9dfc930392 100644 + +if __name__ == "__main__": + unittest.main() +-- +2.54.0.windows.1 + diff --git a/relax/utils/reloadable_process_group.py b/relax/utils/reloadable_process_group.py index 4c29427ac..3dd9ce884 100644 --- a/relax/utils/reloadable_process_group.py +++ b/relax/utils/reloadable_process_group.py @@ -338,6 +338,8 @@ def destroy_process_groups(post_destroy_delay: float = 2.0): """Destroy all reloadable process groups.""" if _should_skip_reload_and_destroy(): return + if device_utils.is_npu_available: + post_destroy_delay = 0.0 ReloadableProcessGroup.destroy_process_groups(post_destroy_delay=post_destroy_delay) diff --git a/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh b/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh index ae208cf74..d5801af6c 100644 --- a/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh +++ b/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh @@ -2,7 +2,7 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. # -# Qwen3.5-35B-A3B 16xNPU colocate training script. +# Qwen35-35B 16xNPU colocate training script. # # Usage: # bash scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh @@ -12,7 +12,7 @@ set -o pipefail unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY ulimit -n 65535 - +export MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV=1 export HCCL_SOCKET_IFNAME="${HCCL_SOCKET_IFNAME:-enp23s0f3}" export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-enp23s0f3}" export TP_SOCKET_IFNAME="${TP_SOCKET_IFNAME:-enp23s0f3}" @@ -29,6 +29,34 @@ export HCCL_NPU_SOCKET_PORT_RANGE=64000-64050 export TMS_HOOK_MODE="preload" export HYDRA_FULL_ERROR=1 +# optimize +export CPU_AFFINITY_CONF=1 +export TORCH_HCCL_ZERO_COPY=1 +export MULTI_STREAM_MEMORY_REUSE=1 +export HCCL_OP_EXPANSION_MODE="AIV" + +export SGLANG_NPU_USE_MULTI_STREAM=1 + +export SGLANG_NPU_ASYNC_EXPONENTIAL=1 +export SGLANG_GMM2_TRITON=1 +export SGLANG_MOE_FRONT_FUSION=1 +export SGLANG_NPU_FULL_ATTN_FUSION=1 + +export ASCEND_USE_FIA=1 +export GDN_ATTN_BACKEND_TRITON=0 +export STREAMS_PER_DEVICE=32 +export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 +export SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES=30 +export TASK_QUEUE_ENABLE=1 +export SGLANG_NPU_GDN_UPDATE_FUSED=1 + +export SGLANG_NPU_TP_ASCENDC_FUSION=1 +export SGLANG_NPU_GDN_RECURRENT_ASCENDC=1 + +export SGLANG_NPU_EXP_RACE_TRITON=1 +export SGLANG_NPU_GDN_QKVZBA_PACK=1 +export SGLANG_NPU_GDN_QKVZBA_PACK_MAX_M=256 +export SGLANG_NPU_MOE_PREFETCH=0 SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" # Auto-source local environment when not launched via an external entrypoint @@ -52,117 +80,129 @@ CKPT_ARGS=( PROMPT_SET=${EXP_DIR}/dapo-math-17k/dapo-math-17k.jsonl ROLLOUT_ARGS=( - --prompt-data ${PROMPT_SET} - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type dapo - --reward-key score - --num-rollout ${NUM_ROLLOUT} - --rollout-batch-size 32 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 1 - --global-batch-size 256 - --use-fault-tolerance + --prompt-data ${PROMPT_SET} + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type dapo + --reward-key score + --num-rollout ${NUM_ROLLOUT} + --rollout-batch-size 16 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + --global-batch-size 128 + --use-fault-tolerance ) EVAL_ARGS=( - --log-passrate - --eval-interval 20 - --skip-eval-before-train - --eval-prompt-data aime ${EXP_DIR}/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 8 - --eval-max-response-len 8192 - #--eval-top-p 0.7 + --log-passrate + --eval-interval 20000 + --skip-eval-before-train + --eval-prompt-data aime aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 8192 + #--eval-top-p 0.7 ) PERF_ARGS=( - --tensor-model-parallel-size 4 - --sequence-parallel - --pipeline-model-parallel-size 2 - --context-parallel-size 1 - --expert-model-parallel-size 8 - --expert-tensor-parallel-size 1 - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - #--use-dynamic-batch-size - # Packing is not supported for GDN currently - --qkv-format bshd - --micro-batch-size 1 - --max-tokens-per-gpu 10240 - --no-rope-fusion - --no-gradient-accumulation-fusion + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 2 + --context-parallel-size 1 + --expert-model-parallel-size 16 + --expert-tensor-parallel-size 1 + # --recompute-granularity full + # --recompute-method uniform + # --recompute-num-layers 2 + --use-dynamic-batch-size + # Packing is not supported for GDN currently + --qkv-format thd + # --micro-batch-size 1 + --max-tokens-per-gpu 20480 + --no-rope-fusion + --no-gradient-accumulation-fusion + --balance-data ) GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 - --use-tis + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis ) OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer - --use-distributed-optimizer + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer + --use-distributed-optimizer + --overlap-grad-reduce + --overlap-param-gather ) SGLANG_ARGS=( - --rollout-num-gpus-per-engine 8 - --sglang-mem-fraction-static 0.6 - --sglang-max-running-requests 256 - --sglang-cuda-graph-bs 4 8 16 32 64 128 192 256 - --sglang-device npu - --sglang-disable-radix-cache - --sglang-chunked-prefill-size 8192 - --sglang-max-prefill-tokens 8192 - --sglang-enable-dp-attention - --sglang-enable-dp-lm-head - --sglang-attention-backend ascend + --rollout-num-gpus-per-engine 8 + --sglang-mem-fraction-static 0.85 + --sglang-max-running-requests 132 + --sglang-cuda-graph-bs 4 8 16 24 32 40 48 64 128 + --sglang-device npu + # --sglang-disable-radix-cache + --mamba-scheduler-strategy extra_buffer + --sglang-chunked-prefill-size 8192 + --sglang-max-prefill-tokens 8192 + --sglang-enable-dp-attention + --sglang-enable-dp-lm-head + --sglang-attention-backend ascend + --sglang-pp-size 1 + --sglang-dp-size 1 + --sglang-ep-size 1 + --sglang-max-mamba-cache-size 192 + --sglang-router-policy round_robin + --sglang-mamba-ssm-dtype bfloat16 + --sglang-tokenizer-backend fastokens ) MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - --use-flash-attn + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + --use-flash-attn ) mkdir -p log - ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://${MASTER_ADDR}:8265" \ - ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 -m relax.entrypoints.train \ - --resource '{"actor": [1, 16], "rollout": [1, 16]}'\ - --max-staleness 0 \ - --colocate \ - --num-gpus-per-node 16 \ - --use-health-check \ - "${MODEL_ARGS[@]}" \ - "${CKPT_ARGS[@]}" \ - "${ROLLOUT_ARGS[@]}" \ - "${OPTIMIZER_ARGS[@]}" \ - "${GRPO_ARGS[@]}" \ - "${PERF_ARGS[@]}" \ - "${EVAL_ARGS[@]}" \ - "${SGLANG_ARGS[@]}" \ - "${MISC_ARGS[@]}" 2>&1 | tee log/qwen35-35B-MATH-gpu16-sync-${now}.log + ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://${MASTER_ADDR}:8265" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"actor": [1, 32], "rollout": [1, 32]}'\ + --max-staleness 0 \ + --colocate \ + --nnodes 2 \ + --num-gpus-per-node 16 \ + --use-health-check \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" 2>&1 | tee log/qwen35-35B-MATH-gpu16-sync-${now}.log diff --git a/scripts/training/text/run-qwen35-9B-8xnpu-colocate.sh b/scripts/training/text/run-qwen35-9B-8xnpu-colocate.sh new file mode 100644 index 000000000..42663b5c7 --- /dev/null +++ b/scripts/training/text/run-qwen35-9B-8xnpu-colocate.sh @@ -0,0 +1,211 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# Qwen35-9B 8xNPU colocate training script. +# +# Usage: +# bash scripts/training/text/run-qwen35-9B-8xnpu-colocate.sh + +set -ex +set -o pipefail +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY + +ulimit -n 65535 +export MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV=1 +export HCCL_SOCKET_IFNAME="${HCCL_SOCKET_IFNAME:-enp23s0f3}" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-enp23s0f3}" +export TP_SOCKET_IFNAME="${TP_SOCKET_IFNAME:-enp23s0f3}" +export HCCL_CONNECT_TIMEOUT=1200 +export RAY_DEDUP_LOGS=0 +export PYTHONBUFFERED=1 + +now=$(date "+%Y-%m-%d-%H:%M:%S") +echo "ๅฝ“ๅ‰ๆ—ถ้—ด: $now" +export ASCEND_COREDUMP_SIGNAL=none +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 +export HCCL_HOST_SOCKET_PORT_RANGE=63000-63050 +export HCCL_NPU_SOCKET_PORT_RANGE=64000-64050 +export TMS_HOOK_MODE="preload" +export HYDRA_FULL_ERROR=1 + +# optimize +export CPU_AFFINITY_CONF=1 +export TORCH_HCCL_ZERO_COPY=1 +export MULTI_STREAM_MEMORY_REUSE=1 +export HCCL_OP_EXPANSION_MODE="AIV" + +export SGLANG_NPU_USE_MULTI_STREAM=1 + +export SGLANG_NPU_ASYNC_EXPONENTIAL=1 +export SGLANG_GMM2_TRITON=0 +export SGLANG_MOE_FRONT_FUSION=1 +export SGLANG_NPU_FULL_ATTN_FUSION=1 + +export ASCEND_USE_FIA=1 +export GDN_ATTN_BACKEND_TRITON=0 +export STREAMS_PER_DEVICE=32 +export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1 +export SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES=30 +export TASK_QUEUE_ENABLE=1 + +export SGLANG_NPU_GDN_UPDATE_FUSED=1 + +export SGLANG_NPU_TP_ASCENDC_FUSION=1 +export SGLANG_NPU_GDN_RECURRENT_ASCENDC=1 + +export SGLANG_NPU_EXP_RACE_TRITON=1 +export SGLANG_NPU_GDN_QKVZBA_PACK=1 +export SGLANG_NPU_GDN_QKVZBA_PACK_MAX_M=256 +export SGLANG_NPU_MOE_PREFETCH=0 + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +# Auto-source local environment when not launched via an external entrypoint +if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then + source "${SCRIPT_DIR}/../../entrypoint/local-npu.sh" +fi + +source "${MODEL_CONFIG_DIR}/qwen35-9B.sh" +EXP_DIR="${EXP_DIR:-/root/exps}" +PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" +NUM_ROLLOUT="${NUM_ROLLOUT:=3000}" + +CKPT_ARGS=( + --hf-checkpoint ${MODEL_DIR}/Qwen3.5-9B/ + --ref-load ${MODEL_DIR}/Qwen3.5-9B/ + --megatron-to-hf-mode bridge + --load ${EXP_DIR}/Qwen3.5-9B_mcore_8xnpu/ + --save ${EXP_DIR}/Qwen3.5-9B_mcore_8xnpu/ + --save-interval 50 + --max-actor-ckpt-to-keep 1 +) + +PROMPT_SET=${DATA_DIR}/dapo-math-17k/dapo-math-17k.jsonl +ROLLOUT_ARGS=( + --prompt-data ${PROMPT_SET} + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type dapo + --reward-key score + --num-rollout ${NUM_ROLLOUT} + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + --global-batch-size 256 + --use-fault-tolerance +) + +EVAL_ARGS=( + --log-passrate + --eval-interval 20000 + --skip-eval-before-train + --eval-prompt-data aime ${EXP_DIR}/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 8192 + #--eval-top-p 0.7 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + # --recompute-granularity full + # --recompute-method uniform + # --recompute-num-layers 2 + --use-dynamic-batch-size + # Packing is not supported for GDN currently + --qkv-format thd + # --micro-batch-size 1 + --max-tokens-per-gpu 10240 + --log-probs-max-tokens-per-gpu 40960 + --no-rope-fusion + --no-gradient-accumulation-fusion + --balance-data +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + # --optimizer-cpu-offload + # --overlap-cpu-optimizer-d2h-h2d + # --use-precision-aware-optimizer + --use-distributed-optimizer + --overlap-grad-reduce + --overlap-param-gather +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 4 + --sglang-mem-fraction-static 0.85 + --sglang-max-running-requests 272 + --sglang-cuda-graph-bs 2 8 16 24 32 64 128 192 256 + --sglang-device npu + --sglang-mamba-scheduler-strategy extra_buffer + --sglang-chunked-prefill-size 16384 + --sglang-max-prefill-tokens 16384 + --sglang-enable-dp-attention + --sglang-enable-dp-lm-head + --sglang-attention-backend ascend + --sglang-pp-size 1 + --sglang-dp-size 1 + --sglang-ep-size 1 + --sglang-max-mamba-cache-size 640 + --sglang-router-policy round_robin + --sglang-mamba-ssm-dtype bfloat16 + --sglang-tokenizer-backend fastokens +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + --use-flash-attn +) +# --debug-rollout-only \ +mkdir -p log + ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://${MASTER_ADDR}:8265" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"actor": [1, 16], "rollout": [1, 16]}'\ + --max-staleness 0 \ + --num-data-storage-units 1 \ + --colocate \ + --nnodes 1 \ + --num-gpus-per-node 16 \ + --use-health-check \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" 2>&1 | tee log/qwen35-9B-GRPO-npu8-colocate-${now}.log From c6d7b235376f98b46f7343e8a7c1474afdeb114f Mon Sep 17 00:00:00 2001 From: dabuliu123 <270334047@qq.com> Date: Thu, 10 Sep 2026 16:55:12 +0800 Subject: [PATCH 2/3] feat(npu): support multi-node FLA training MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit feat(npu): support multi-node FLA training # โญ Feature ## Add multi-node NPU entrypoint - Add `scripts/entrypoint/local-npu-multinode.sh` for Ray head/worker cluster setup - Head node starts Ray head and waits until all worker nodes join - Worker nodes poll head GCS port and join the cluster with retries ## Update Qwen3.5-35B colocate training script - Switch entrypoint to local-npu-multinode.sh and branch head/worker on `MASTER_ADDR` - Enable AscendC GDN backend via `MINDSPEED_BRIDGE_GDN_BACKEND=ascendc` - Default `MODEL_DIR`/`DATA_DIR` to `EXP_DIR` - Disable `--sglang-mamba-scheduler-strategy extra_buffer` ## Upgrade NPU docker image - Bump `SOC_VERSION_FOR_SGL_BUILD` to Ascend910_9392 - Upgrade flash-linear-attention-npu v26.1.0 โ†’ v26.6.0 - Replace build.sh flow with pip wheel build flow for fla_npu --- # ๐ŸŽจ Style ## Normalize patch file endings - Remove trailing blank lines in docker/npu_patch (pre-commit end-of-file-fixer) fix(npu): scope colocate script to master node # ๐Ÿ› Bug Fix ## Scope master-node training block - Indent config and ray job submit under the master-node branch - Remove dead worker-node else branch ## Disable SGLang radix cache - Enable --sglang-disable-radix-cache in SGLANG_ARGS --- # ๐ŸŽจ Style ## Clean up training script - Remove stale packing and micro-batch comments - Normalize trailing whitespace and file endings (pre-commit) --- docker/Dockerfile.npu | 12 +- docker/npu_patch/sgl-kernel-npu.patch | 1 - docker/npu_patch/sglang-npu.patch | 1 - scripts/entrypoint/local-npu-multinode.sh | 145 +++++++++ .../run-qwen35-35B-A3B-16xnpu-colocate.sh | 291 +++++++++--------- 5 files changed, 296 insertions(+), 154 deletions(-) create mode 100644 scripts/entrypoint/local-npu-multinode.sh diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 7d6d2e3bf..1e369795b 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -12,7 +12,7 @@ ARG NO_PROXY ARG http_proxy ARG https_proxy ARG no_proxy -ARG SOC_VERSION_FOR_SGL_BUILD="Ascend910_9382" +ARG SOC_VERSION_FOR_SGL_BUILD="Ascend910_9392" ENV http_proxy=${HTTP_PROXY} \ https_proxy=${HTTPS_PROXY} \ @@ -150,17 +150,13 @@ RUN git clone https://github.com/sgl-project/sgl-kernel-npu /root/sgl-kernel-npu # Install AscendC FLA RUN git clone https://github.com/flashserve/flash-linear-attention-npu.git /root/flash-linear-attention-npu && \ - cd /root/flash-linear-attention-npu && git checkout v26.1.0 && \ + cd /root/flash-linear-attention-npu && git checkout v26.6.0 && \ apt update && apt-get install -y --no-install-recommends gawk && \ - # ็ผ–่ฏ‘ๅ‘ฝไปค๏ผŒๆณจๆ„--soc=${soc_version}้œ€่ฆๆŒ‡ๅฎšไธบๅฝ“ๅ‰ๆœบๅ™จ็š„่Šฏ็‰‡็ฑปๅž‹{ascend910b/ascend910_93/ascend950} - bash build.sh --soc=ascend910_93 --pkg --ops=causal_conv1d,chunk_bwd_dv_local,chunk_bwd_dqkwg,chunk_gated_delta_rule_bwd_dhu,prepare_wy_repr_bwd_da,prepare_wy_repr_bwd_full,chunk_fwd_o,chunk_gated_delta_rule_fwd_h,recurrent_gated_delta_rule,recompute_wu_fwd && \ - # ๅฎ‰่ฃ…runๅŒ… - ./build_out/cann-*.run && \ source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ source /usr/local/Ascend/nnal/atb/set_env.sh && \ source /usr/local/Ascend/cann-9.0.0/share/info/ascendnpu-ir/bin/set_env.sh && \ - # ไธ€้”ฎ็ผ–่ฏ‘ๅฎ‰่ฃ…่„šๆœฌ๏ผŒๅ…ˆ่ฐƒ็”จtorchnpugen่‡ชๅŠจๆŽฅๅ…ฅ็ฎ—ๅญ๏ผŒๅ†่ฟ่กŒsetup็ผ–whlๅŒ…๏ผŒๆœ€ๅŽๅฎ‰่ฃ…whlๅŒ… - cd torch_custom/fla_npu && bash build.sh + FLA_NPU_SOC=ascend910_93 python -m pip wheel --no-build-isolation --no-deps . -w dist && \ + python -m pip install --force-reinstall --no-deps dist/flash_linear_attention_npu-*.whl FROM train as relax diff --git a/docker/npu_patch/sgl-kernel-npu.patch b/docker/npu_patch/sgl-kernel-npu.patch index 3b597a8ed..a4b074977 100644 --- a/docker/npu_patch/sgl-kernel-npu.patch +++ b/docker/npu_patch/sgl-kernel-npu.patch @@ -4754,4 +4754,3 @@ index 0000000..3c12d53 + raise SystemExit(pytest.main([__file__])) -- 2.54.0.windows.1 - diff --git a/docker/npu_patch/sglang-npu.patch b/docker/npu_patch/sglang-npu.patch index beb8590f6..085bb7240 100644 --- a/docker/npu_patch/sglang-npu.patch +++ b/docker/npu_patch/sglang-npu.patch @@ -4306,4 +4306,3 @@ index 0000000000..9dfc930392 + unittest.main() -- 2.54.0.windows.1 - diff --git a/scripts/entrypoint/local-npu-multinode.sh b/scripts/entrypoint/local-npu-multinode.sh new file mode 100644 index 000000000..3b35920b9 --- /dev/null +++ b/scripts/entrypoint/local-npu-multinode.sh @@ -0,0 +1,145 @@ +if [ -n "${RELAX_ENTRYPOINT_MODE:-}" ]; then + return 0 2>/dev/null || exit 0 +fi + +_LOCAL_SH_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" + +# โ”€โ”€ delegate to ray-job.sh when inside an existing Ray cluster โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +# When RAY_ADDRESS is set AND `ray status` succeeds, we're already part of an +# externally-managed Ray cluster. Skip local Ray startup / process cleanup and +# fall through to ray-job.sh (source mode) for env setup. +if [ -n "${RAY_ADDRESS:-}" ] && timeout 5 ray status >/dev/null 2>&1; then + echo "=== Detected existing Ray cluster (RAY_ADDRESS=$RAY_ADDRESS); delegating to ray-job.sh ===" + source "${_LOCAL_SH_DIR}/ray-job-npu.sh" + return 0 2>/dev/null || exit 0 +fi + +set -eo pipefail + +# โ”€โ”€ process cleanup โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +echo "=== Cleaning up stale processes ===" +pkill -9 sglang 2>/dev/null || true +sleep 3 +ray stop --force 2>/dev/null || true +pkill -9 ray 2>/dev/null || true +pkill -9 python 2>/dev/null || true +sleep 3 +pkill -9 ray 2>/dev/null || true +pkill -9 python 2>/dev/null || true + +set -x + +# โ”€โ”€ environment setup โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +export PYTHONUNBUFFERED=1 +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export MEGATRON=${MEGATRON:-/root/Megatron-LM/} +export MEGATRON_BRIDGE_SRC=${MEGATRON_BRIDGE_SRC:-/root/Megatron-Bridge/src/} +export MINDSPEED=${MINDSPEED:-/root/MindSpeed/} +export RELAX=${RELAX:-${_LOCAL_SH_DIR}/../../} +export PYTHONPATH=${RELAX}:${MEGATRON_BRIDGE_SRC}:${MINDSPEED}:$MEGATRON:$RELAX:${PYTHONPATH:-} +export MODEL_CONFIG_DIR="${_LOCAL_SH_DIR}/../models" + +# โ”€โ”€ Ray cluster startup (multi node) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +export MASTER_ADDR_IP=$(ping -c 1 $MASTER_ADDR | head -n 1 | awk -F'[()]' '{print $2}') +# โ”€โ”€ multi-node parameters โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +NUM_NPUS="${NUM_NPUS:-16}" +NNODES="${WORLD_SIZE:-2}" + +if [ "$MASTER_ADDR" = "$POD_NAME" ]; then + # โ”€โ”€ HEAD NODE โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + echo "=== Head node: starting Ray cluster ===" + ray start --head \ + --node-ip-address "${HOST_IP}" \ + --resources="{\"NPU\": ${NUM_NPUS}}" \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 + + sleep 5 + + # Wait for all worker nodes to join + while true; do + ray_status_output=$(ray status) + npu_count=$(echo "$ray_status_output" | grep -oP '(?<=/)\d+\.\d+(?=\s*NPU)' | head -n 1) + echo "Current NPU count: $npu_count" + npu_count_int=$(echo "$npu_count" | awk '{print int($1)}') + device_count=$((npu_count_int / ${NUM_NPUS})) + + if [ "$device_count" -eq "$NNODES" ]; then + echo "Ray cluster is ready with $device_count devices (from $npu_count NPU resources)." + ray status + break + else + echo "Waiting for Ray to allocate $NNODES devices. Current device count: $device_count" + sleep 5 + fi + done + + # โ”€โ”€ set entrypoint mode โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + export RELAX_ENTRYPOINT_MODE="npu-multinode" + + # Runtime env for multi-node + export RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONUNBUFFERED\": \"1\", + \"PYTHONPATH\": \"${PYTHONPATH}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"RAY_OVERRIDE_JOB_RUNTIME_ENV\": \"1\", + \"MASTER_ADDR\": \"${HOST_IP}\", + \"RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES\": \"1\" + } + }" + + echo "=== Head node environment ready ===" +else + # โ”€โ”€ WORKER NODE โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + # NOTE: `set -e` is active, so each retry loop below must keep the + # potentially-failing command in a condition position (if/until/||), + # otherwise the first failure kills the script and there is no retry. + GCS_PORT="${GCS_PORT:-6379}" + echo "=== Worker node: waiting for head GCS at ${MASTER_ADDR_IP}:${GCS_PORT} ===" + for i in $(seq 1 120); do + if timeout 2 bash -c "/dev/null; then + echo "Head GCS reachable after ${i} attempt(s)" + break + fi + if [ "$i" -eq 120 ]; then + echo "ERROR: head GCS at ${MASTER_ADDR_IP}:${GCS_PORT} unreachable after 10min" >&2 + exit 1 + fi + sleep 5 + done + + echo "=== Worker node: joining Ray cluster at ${MASTER_ADDR_IP}:${GCS_PORT} ===" + joined=0 + for i in $(seq 1 30); do + ray stop --force >/dev/null 2>&1 || true + if ray start \ + --address="${MASTER_ADDR_IP}:${GCS_PORT}" \ + --resources="{\"NPU\": ${NUM_NPUS}}" \ + --node-ip-address "${HOST_IP}" \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265; then + echo "Joined Ray cluster on attempt ${i}" + joined=1 + break + fi + echo "ray start failed on attempt ${i}, retrying in 5s..." + sleep 5 + done + if [ "$joined" -ne 1 ]; then + echo "ERROR: worker failed to join Ray cluster after 30 attempts" >&2 + exit 1 + fi + + if ! ray status >/dev/null 2>&1; then + echo "ERROR: ray status failed after join" >&2 + exit 1 + fi + echo "Successfully connected to the Ray cluster!" + + # Worker nodes block indefinitely (training runs on head node) + echo "=== Worker node ready, waiting for training to complete ===" + sleep inf +fi diff --git a/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh b/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh index d5801af6c..561f91e32 100644 --- a/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh +++ b/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh @@ -13,6 +13,7 @@ unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY ulimit -n 65535 export MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV=1 +export MINDSPEED_BRIDGE_GDN_BACKEND=ascendc export HCCL_SOCKET_IFNAME="${HCCL_SOCKET_IFNAME:-enp23s0f3}" export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-enp23s0f3}" export TP_SOCKET_IFNAME="${TP_SOCKET_IFNAME:-enp23s0f3}" @@ -61,148 +62,150 @@ export SGLANG_NPU_MOE_PREFETCH=0 SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" # Auto-source local environment when not launched via an external entrypoint if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then - source "${SCRIPT_DIR}/../../entrypoint/local-npu.sh" + source "${SCRIPT_DIR}/../../entrypoint/local-npu-multinode.sh" +fi +if [ "$MASTER_ADDR" = "$POD_NAME" ]; then + source "${MODEL_CONFIG_DIR}/qwen35-35B-A3B.sh" + EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" + MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" + DATA_DIR="${DATA_DIR:-${EXP_DIR}}" + PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" + NUM_ROLLOUT="${NUM_ROLLOUT:=3000}" + + CKPT_ARGS=( + --hf-checkpoint ${EXP_DIR}/Qwen3.5-35B-A3B + --ref-load ${EXP_DIR}/Qwen3.5-35B-A3B + --megatron-to-hf-mode bridge + # --load ${EXP_DIR}/Qwen3.5-35B-A3B-save + --save ${EXP_DIR}/Qwen3.5-35B-A3B-save + --save-interval 100 + ) + + PROMPT_SET=${EXP_DIR}/dapo-math-17k/dapo-math-17k.jsonl + + ROLLOUT_ARGS=( + --prompt-data ${PROMPT_SET} + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type dapo + --reward-key score + --num-rollout ${NUM_ROLLOUT} + --rollout-batch-size 16 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + --global-batch-size 128 + --use-fault-tolerance + ) + + EVAL_ARGS=( + --log-passrate + --eval-interval 20000 + --skip-eval-before-train + --eval-prompt-data aime aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 8192 + #--eval-top-p 0.7 + ) + + PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 2 + --context-parallel-size 1 + --expert-model-parallel-size 16 + --expert-tensor-parallel-size 1 + # --recompute-granularity full + # --recompute-method uniform + # --recompute-num-layers 2 + --use-dynamic-batch-size + --qkv-format thd + --max-tokens-per-gpu 20480 + --no-rope-fusion + --no-gradient-accumulation-fusion + --balance-data + ) + + GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis + ) + + OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer + --use-distributed-optimizer + --overlap-grad-reduce + --overlap-param-gather + ) + + SGLANG_ARGS=( + --rollout-num-gpus-per-engine 8 + --sglang-mem-fraction-static 0.85 + --sglang-max-running-requests 132 + --sglang-cuda-graph-bs 4 8 16 24 32 40 48 64 128 + --sglang-device npu + --sglang-disable-radix-cache + # --sglang-mamba-scheduler-strategy extra_buffer + --sglang-chunked-prefill-size 8192 + --sglang-max-prefill-tokens 8192 + --sglang-enable-dp-attention + --sglang-enable-dp-lm-head + --sglang-attention-backend ascend + --sglang-pp-size 1 + --sglang-dp-size 1 + --sglang-ep-size 1 + --sglang-max-mamba-cache-size 192 + --sglang-router-policy round_robin + --sglang-mamba-ssm-dtype bfloat16 + --sglang-tokenizer-backend fastokens + ) + + MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + --use-flash-attn + ) + + mkdir -p log + ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://${MASTER_ADDR}:8265" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"actor": [1, 32], "rollout": [1, 32]}'\ + --max-staleness 0 \ + --colocate \ + --nnodes 2 \ + --num-gpus-per-node 16 \ + --use-health-check \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" 2>&1 | tee log/qwen35-35B-MATH-gpu16-sync-${now}.log fi -source "${MODEL_CONFIG_DIR}/qwen35-35B-A3B.sh" -EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" -PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" -NUM_ROLLOUT="${NUM_ROLLOUT:=3000}" - -CKPT_ARGS=( - --hf-checkpoint ${EXP_DIR}/Qwen3.5-35B-A3B - --ref-load ${EXP_DIR}/Qwen3.5-35B-A3B - --megatron-to-hf-mode bridge - # --load ${EXP_DIR}/Qwen3.5-35B-A3B-save - --save ${EXP_DIR}/Qwen3.5-35B-A3B-save - --save-interval 100 -) - -PROMPT_SET=${EXP_DIR}/dapo-math-17k/dapo-math-17k.jsonl - -ROLLOUT_ARGS=( - --prompt-data ${PROMPT_SET} - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type dapo - --reward-key score - --num-rollout ${NUM_ROLLOUT} - --rollout-batch-size 16 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 1 - --global-batch-size 128 - --use-fault-tolerance -) - -EVAL_ARGS=( - --log-passrate - --eval-interval 20000 - --skip-eval-before-train - --eval-prompt-data aime aime-2024.jsonl - --n-samples-per-eval-prompt 8 - --eval-max-response-len 8192 - #--eval-top-p 0.7 -) - -PERF_ARGS=( - --tensor-model-parallel-size 8 - --sequence-parallel - --pipeline-model-parallel-size 2 - --context-parallel-size 1 - --expert-model-parallel-size 16 - --expert-tensor-parallel-size 1 - # --recompute-granularity full - # --recompute-method uniform - # --recompute-num-layers 2 - --use-dynamic-batch-size - # Packing is not supported for GDN currently - --qkv-format thd - # --micro-batch-size 1 - --max-tokens-per-gpu 20480 - --no-rope-fusion - --no-gradient-accumulation-fusion - --balance-data -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 - --use-tis -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer - --use-distributed-optimizer - --overlap-grad-reduce - --overlap-param-gather -) - -SGLANG_ARGS=( - --rollout-num-gpus-per-engine 8 - --sglang-mem-fraction-static 0.85 - --sglang-max-running-requests 132 - --sglang-cuda-graph-bs 4 8 16 24 32 40 48 64 128 - --sglang-device npu - # --sglang-disable-radix-cache - --mamba-scheduler-strategy extra_buffer - --sglang-chunked-prefill-size 8192 - --sglang-max-prefill-tokens 8192 - --sglang-enable-dp-attention - --sglang-enable-dp-lm-head - --sglang-attention-backend ascend - --sglang-pp-size 1 - --sglang-dp-size 1 - --sglang-ep-size 1 - --sglang-max-mamba-cache-size 192 - --sglang-router-policy round_robin - --sglang-mamba-ssm-dtype bfloat16 - --sglang-tokenizer-backend fastokens -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - --use-flash-attn -) - -mkdir -p log - ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://${MASTER_ADDR}:8265" \ - ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 -m relax.entrypoints.train \ - --resource '{"actor": [1, 32], "rollout": [1, 32]}'\ - --max-staleness 0 \ - --colocate \ - --nnodes 2 \ - --num-gpus-per-node 16 \ - --use-health-check \ - "${MODEL_ARGS[@]}" \ - "${CKPT_ARGS[@]}" \ - "${ROLLOUT_ARGS[@]}" \ - "${OPTIMIZER_ARGS[@]}" \ - "${GRPO_ARGS[@]}" \ - "${PERF_ARGS[@]}" \ - "${EVAL_ARGS[@]}" \ - "${SGLANG_ARGS[@]}" \ - "${MISC_ARGS[@]}" 2>&1 | tee log/qwen35-35B-MATH-gpu16-sync-${now}.log From 7b7997246cd9d3b32e7d5e49aaff67c4bcc49b52 Mon Sep 17 00:00:00 2001 From: liyongwen <1310439159@qq.com> Date: Fri, 11 Sep 2026 17:50:43 +0800 Subject: [PATCH 3/3] fix npu script --- .../run-qwen35-35B-A3B-16xnpu-colocate.sh | 292 +++++++++--------- .../text/run-qwen35-9B-8xnpu-colocate.sh | 2 + 2 files changed, 149 insertions(+), 145 deletions(-) diff --git a/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh b/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh index 561f91e32..dce78b24d 100644 --- a/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh +++ b/scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh @@ -64,148 +64,150 @@ SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then source "${SCRIPT_DIR}/../../entrypoint/local-npu-multinode.sh" fi -if [ "$MASTER_ADDR" = "$POD_NAME" ]; then - source "${MODEL_CONFIG_DIR}/qwen35-35B-A3B.sh" - EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" - MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" - DATA_DIR="${DATA_DIR:-${EXP_DIR}}" - PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" - NUM_ROLLOUT="${NUM_ROLLOUT:=3000}" - - CKPT_ARGS=( - --hf-checkpoint ${EXP_DIR}/Qwen3.5-35B-A3B - --ref-load ${EXP_DIR}/Qwen3.5-35B-A3B - --megatron-to-hf-mode bridge - # --load ${EXP_DIR}/Qwen3.5-35B-A3B-save - --save ${EXP_DIR}/Qwen3.5-35B-A3B-save - --save-interval 100 - ) - - PROMPT_SET=${EXP_DIR}/dapo-math-17k/dapo-math-17k.jsonl - - ROLLOUT_ARGS=( - --prompt-data ${PROMPT_SET} - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type dapo - --reward-key score - --num-rollout ${NUM_ROLLOUT} - --rollout-batch-size 16 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 1 - --global-batch-size 128 - --use-fault-tolerance - ) - - EVAL_ARGS=( - --log-passrate - --eval-interval 20000 - --skip-eval-before-train - --eval-prompt-data aime aime-2024.jsonl - --n-samples-per-eval-prompt 8 - --eval-max-response-len 8192 - #--eval-top-p 0.7 - ) - - PERF_ARGS=( - --tensor-model-parallel-size 8 - --sequence-parallel - --pipeline-model-parallel-size 2 - --context-parallel-size 1 - --expert-model-parallel-size 16 - --expert-tensor-parallel-size 1 - # --recompute-granularity full - # --recompute-method uniform - # --recompute-num-layers 2 - --use-dynamic-batch-size - --qkv-format thd - --max-tokens-per-gpu 20480 - --no-rope-fusion - --no-gradient-accumulation-fusion - --balance-data - ) - - GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 - --use-tis - ) - - OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer - --use-distributed-optimizer - --overlap-grad-reduce - --overlap-param-gather - ) - - SGLANG_ARGS=( - --rollout-num-gpus-per-engine 8 - --sglang-mem-fraction-static 0.85 - --sglang-max-running-requests 132 - --sglang-cuda-graph-bs 4 8 16 24 32 40 48 64 128 - --sglang-device npu - --sglang-disable-radix-cache - # --sglang-mamba-scheduler-strategy extra_buffer - --sglang-chunked-prefill-size 8192 - --sglang-max-prefill-tokens 8192 - --sglang-enable-dp-attention - --sglang-enable-dp-lm-head - --sglang-attention-backend ascend - --sglang-pp-size 1 - --sglang-dp-size 1 - --sglang-ep-size 1 - --sglang-max-mamba-cache-size 192 - --sglang-router-policy round_robin - --sglang-mamba-ssm-dtype bfloat16 - --sglang-tokenizer-backend fastokens - ) - - MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - --use-flash-attn - ) - - mkdir -p log - ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://${MASTER_ADDR}:8265" \ - ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 -m relax.entrypoints.train \ - --resource '{"actor": [1, 32], "rollout": [1, 32]}'\ - --max-staleness 0 \ - --colocate \ - --nnodes 2 \ - --num-gpus-per-node 16 \ - --use-health-check \ - "${MODEL_ARGS[@]}" \ - "${CKPT_ARGS[@]}" \ - "${ROLLOUT_ARGS[@]}" \ - "${OPTIMIZER_ARGS[@]}" \ - "${GRPO_ARGS[@]}" \ - "${PERF_ARGS[@]}" \ - "${EVAL_ARGS[@]}" \ - "${SGLANG_ARGS[@]}" \ - "${MISC_ARGS[@]}" 2>&1 | tee log/qwen35-35B-MATH-gpu16-sync-${now}.log -fi +# NOTE: no node-role check here. When self-launching a multi-node cluster, +# local-npu-multinode.sh already handles head/worker split (worker nodes block +# there and never return). In existing-cluster mode (ray-job-npu.sh) we submit +# directly. +source "${MODEL_CONFIG_DIR}/qwen35-35B-A3B.sh" +EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" +MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" +DATA_DIR="${DATA_DIR:-${EXP_DIR}}" +PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" +NUM_ROLLOUT="${NUM_ROLLOUT:=3000}" + +CKPT_ARGS=( + --hf-checkpoint ${EXP_DIR}/Qwen3.5-35B-A3B + --ref-load ${EXP_DIR}/Qwen3.5-35B-A3B + --megatron-to-hf-mode bridge + # --load ${EXP_DIR}/Qwen3.5-35B-A3B-save + --save ${EXP_DIR}/Qwen3.5-35B-A3B-save + --save-interval 100 +) + +PROMPT_SET=${EXP_DIR}/dapo-math-17k/dapo-math-17k.jsonl + +ROLLOUT_ARGS=( + --prompt-data ${PROMPT_SET} + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type dapo + --reward-key score + --num-rollout ${NUM_ROLLOUT} + --rollout-batch-size 16 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + --global-batch-size 128 + --use-fault-tolerance +) + +EVAL_ARGS=( + --log-passrate + --eval-interval 20000 + --skip-eval-before-train + --eval-prompt-data aime aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 8192 + #--eval-top-p 0.7 +) + +PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 2 + --context-parallel-size 1 + --expert-model-parallel-size 16 + --expert-tensor-parallel-size 1 + # --recompute-granularity full + # --recompute-method uniform + # --recompute-num-layers 2 + --use-dynamic-batch-size + --qkv-format thd + --max-tokens-per-gpu 20480 + --no-rope-fusion + --no-gradient-accumulation-fusion + --balance-data +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer + --use-distributed-optimizer + --overlap-grad-reduce + --overlap-param-gather +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 8 + --sglang-mem-fraction-static 0.85 + --sglang-max-running-requests 132 + --sglang-cuda-graph-bs 4 8 16 24 32 40 48 64 128 + --sglang-device npu + --sglang-disable-radix-cache + # --sglang-mamba-scheduler-strategy extra_buffer + --sglang-chunked-prefill-size 8192 + --sglang-max-prefill-tokens 8192 + --sglang-enable-dp-attention + --sglang-enable-dp-lm-head + --sglang-attention-backend ascend + --sglang-pp-size 1 + --sglang-dp-size 1 + --sglang-ep-size 1 + --sglang-max-mamba-cache-size 192 + --sglang-router-policy round_robin + --sglang-mamba-ssm-dtype bfloat16 + --sglang-tokenizer-backend fastokens +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash + --use-flash-attn +) + +mkdir -p log +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://${MASTER_ADDR}:8265" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"actor": [1, 32], "rollout": [1, 32]}'\ + --max-staleness 0 \ + --colocate \ + --nnodes 2 \ + --num-gpus-per-node 16 \ + --use-health-check \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" 2>&1 | tee log/qwen35-35B-MATH-gpu16-sync-${now}.log diff --git a/scripts/training/text/run-qwen35-9B-8xnpu-colocate.sh b/scripts/training/text/run-qwen35-9B-8xnpu-colocate.sh index 42663b5c7..de0be3114 100644 --- a/scripts/training/text/run-qwen35-9B-8xnpu-colocate.sh +++ b/scripts/training/text/run-qwen35-9B-8xnpu-colocate.sh @@ -67,6 +67,8 @@ fi source "${MODEL_CONFIG_DIR}/qwen35-9B.sh" EXP_DIR="${EXP_DIR:-/root/exps}" +MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" +DATA_DIR="${DATA_DIR:-${EXP_DIR}}" PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" NUM_ROLLOUT="${NUM_ROLLOUT:=3000}"