From 6c91af95837dfd39af1cdacab99ab22ae7c3a2f4 Mon Sep 17 00:00:00 2001 From: dabuliu123 <270334047@qq.com> Date: Fri, 21 Aug 2026 11:50:54 +0800 Subject: [PATCH 01/16] Merge branch 'wqw_base_lw_dev_0808' into ascend-dev-0808 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- docker/Dockerfile.npu | 98 +- docker/npu_patch/megatron-bridge.patch | 56 +- docker/npu_patch/mindspeed-bridge.patch | 132 +- docker/npu_patch/mindspeed-ops.patch | 13 + docker/npu_patch/mindspeed.patch | 14 + docker/npu_patch/sgl-kernel-npu.patch | 1707 +++++++++++++++++++++++ docker/npu_patch/sglang-npu.patch | 433 +++++- 7 files changed, 2330 insertions(+), 123 deletions(-) create mode 100644 docker/npu_patch/mindspeed-ops.patch create mode 100644 docker/npu_patch/sgl-kernel-npu.patch diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 89ed93f0c..dc9b05197 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -2,7 +2,7 @@ ARG HTTP_PROXY ARG HTTPS_PROXY ARG NO_PROXY -FROM quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11 +FROM quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11 ARG HTTP_PROXY ARG HTTPS_PROXY @@ -43,86 +43,112 @@ RUN ARCH=$(uname -m) && \ export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ fi && \ source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ - source /usr/local/Ascend/nnal/atb/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 + + # Setting pip & git config. Global config (set once, persists across subsequent RUN layers) ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" RUN pip config set global.index-url ${PIP_INDEX_URL} && \ git config --global http.sslverify false && \ + git config --global https.sslverify false && \ git config --global http.postBuffer 2147483648 && \ git config --global user.email "temp@example.com" && \ git config --global user.name "temp" -WORKDIR /root -COPY . /root/Relax +# install torch RUN pip install --upgrade pip packaging setuptools==80.10.2 && \ pip install torch==2.9.0 && \ pip install numpy==1.26.0 - # build torch_npu + +# build torch_npu RUN pip install pyyaml && \ git clone https://gitcode.com/Ascend/pytorch.git /root/pytorch && \ cd /root/pytorch && \ - git checkout v2.9.0-7.3.0 && \ + git checkout v26.0.1-pytorch2.9.0 && \ git cherry-pick -n f495de675bce38a2fa21edbf067b73d2a5f26733 && \ bash ci/build.sh --python=3.11 && \ - pip install dist/torch_npu-2.9.0*.whl -RUN cd /root && rm -rf /root/pytorch && \ - pip install triton-ascend==3.2.0 && \ + pip install dist/torch_npu*.whl + +# install triton/TQ +RUN pip install triton-ascend==3.2.1 --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple && \ pip install tensordict==0.10.0 pyvers==0.1.0 --no-deps && \ pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps - # Clone Megatron-LM, MindSpeed, MindSpeed-Bridge, Megatron-Bridge and install +# Clone Megatron-LM, MindSpeed, MindSpeed-Bridge, Megatron-Bridge and install RUN git clone https://gitcode.com/ascend/MindSpeed.git /root/MindSpeed && \ git clone https://github.com/NVIDIA/Megatron-LM.git /root/Megatron-LM && \ + git clone https://gitcode.com/ascend/MindSpeed-Ops.git /root/MindSpeed-Ops && \ git clone https://gitcode.com/ascend/MindSpeed-Bridge.git /root/MindSpeed-Bridge && \ git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git /root/Megatron-Bridge -RUN cd /root/MindSpeed && git checkout core_r0.16.0 && pip install -r requirements.txt && pip install -e . && \ +RUN cd /root/MindSpeed && git checkout core_r0.16.0 && pip install -r requirements.txt && pip install -e . && git checkout e4772499 && \ cd /root/Megatron-LM && git checkout core_v0.16.1 && pip install -e . --no-build-isolation && \ - cd /root/Megatron-Bridge && git checkout v0.3.1 && \ - cd /root/MindSpeed-Bridge && git checkout 3655c07cbcc9 && pip install -r requirements.txt && bash tools/install_auto.sh - -# Patch Megatron-LM, MindSpeed, MindSpeed-Bridge -RUN cd /root/MindSpeed && \ - patch -p1 < /root/Relax/docker/npu_patch/mindspeed.patch && \ - cd /root/Megatron-Bridge && \ - patch -p1 < /root/Relax/docker/npu_patch/megatron-bridge.patch && \ + cd /root/Megatron-Bridge && git checkout v0.3.1 && cp -r /root/Megatron-Bridge/src/megatron/bridge /root/Megatron-LM/megatron/ && \ + cd /root/MindSpeed-Ops/ && git checkout 33ac80f7 && pip install -e . --no-build-isolation --no-deps && \ + cd /root/MindSpeed-Bridge/ && git checkout v0.3.1 && pip install -r requirements.txt && pip install -e . --no-deps + + +COPY . /root/Relax +# Patch Megatron-LM, MindSpeed, MindSpeed-Bridge, MindSpeed-Ops +RUN cd /root/Megatron-Bridge && \ + patch -p1 < /root/Relax/docker/npu_patch/megatron-bridge.patch && \ cd /root/Megatron-LM && \ - patch -p1 < /root/Relax/docker/npu_patch/megatron.patch && \ + patch -p1 < /root/Relax/docker/npu_patch/megatron.patch && \ cd /root/MindSpeed-Bridge && \ - patch -p1 < /root/Relax/docker/npu_patch/mindspeed-bridge.patch + patch -p1 < /root/Relax/docker/npu_patch/mindspeed-bridge.patch && \ + cd /root/MindSpeed-Ops && \ + patch -p1 < /root/Relax/docker/npu_patch/mindspeed-ops.patch && \ + cd /root/MindSpeed && \ + patch -p1 < /root/Relax/docker/npu_patch/mindspeed.patch -# Copy MindSpeed-Bridge, Megatron-Bridge into Megatron-LM -RUN cp -r /root/MindSpeed-Bridge/mindspeed_bridge /root/Megatron-LM/ && \ - cp -r /root/Megatron-Bridge/src/megatron/bridge /root/Megatron-LM/megatron/ && \ - cd /root && rm -rf /root/MindSpeed-Bridge && rm -rf /root/Megatron-Bridge # Install sglang RUN git clone https://github.com/sgl-project/sglang.git /root/sglang && \ - cd /root/sglang && git checkout v0.5.10 && \ + cd /root/sglang && git checkout v0.5.15 && \ mv python/pyproject.toml python/pyproject.toml.backup && \ mv python/pyproject_npu.toml python/pyproject.toml && \ pip install -e "python[srt_npu]" --constraint <(echo "torch==2.9.0") && \ - # [NPU] Fix Qwen3.5 inference acc. - git stash && \ - git fetch origin pull/23815/head:pr-23815 && \ - git checkout pr-23815 && \ - patch -p1 < /root/Relax/docker/npu_patch/sglang-npu.patch + # patch -p1 < /root/Relax/docker/npu_patch/sglang-npu.patch + git add . && git commit -m "install info" && \ + git fetch && \ + git cherry-pick ece02ffc9cc32e94382d4f1b553b2c755f83f722 && \ + git am /root/Relax/docker/npu_patch/sglang-npu.patch + # Install sgl-kernle-npu RUN git clone https://github.com/sgl-project/sgl-kernel-npu /root/sgl-kernel-npu && \ - cd /root/sgl-kernel-npu && git checkout 2026.04.15.rc3 && \ + cd /root/sgl-kernel-npu && git checkout 2026.7.2 && \ # Adapt tms for colocate train. - git cherry-pick -n 23519771d347 --no-gpg-sign && \ - patch -p1 < /root/Relax/docker/npu_patch/torch-memory-saver.patch && \ - bash build.sh -a kernels && bash build.sh -a memory-saver && \ + git am /root/Relax/docker/npu_patch/sgl-kernel-npu.patch && \ + bash build.sh && \ pip install output/*.whl && \ - cd /root && rm -rf /root/sgl-kernel-npu + cd /root + +# 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 && \ + apt update && apt install 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 + + # Install Relax +# git clone https://github.com/redai-infra/Relax.git RUN cd /root/Relax && pip install -e . + RUN pip install ray==2.55.1 && pip install protobuf==6.33.6 #Clean cache RUN pip cache purge && \ rm -rf /tmp/* + diff --git a/docker/npu_patch/megatron-bridge.patch b/docker/npu_patch/megatron-bridge.patch index 0066df63b..d49be66b3 100644 --- a/docker/npu_patch/megatron-bridge.patch +++ b/docker/npu_patch/megatron-bridge.patch @@ -1,11 +1,32 @@ diff --git a/src/megatron/bridge/models/conversion/utils.py b/src/megatron/bridge/models/conversion/utils.py -index 5a66e719..3d411c17 100644 +index 86ddf8661..115624b16 100644 --- a/src/megatron/bridge/models/conversion/utils.py +++ b/src/megatron/bridge/models/conversion/utils.py -@@ -203,6 +203,15 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0): +@@ -203,6 +203,17 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0): + ): # bound methods + return None + ++ # Convert OmegaConf containers to plain dict/list to avoid in-place ++ # mutation triggering "dictionary changed size during iteration" errors ++ # inside OmegaConf's internal _flags_cache handling. ++ try: ++ from omegaconf import DictConfig, ListConfig, OmegaConf as _OmegaConf ++ ++ if isinstance(obj, (DictConfig, ListConfig)): ++ obj = _OmegaConf.to_container(obj, resolve=True) ++ except ImportError: ++ pass ++ + # Handle dataclass/object with attributes + if hasattr(obj, "__dict__"): + # Create a copy to avoid modifying the original +@@ -213,9 +224,18 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0): # Recursively clean attribute cleaned_value = remove_non_pickleables(attr_value, max_depth, current_depth + 1) +- +- # Set the cleaned value (or None if it was removed) +- setattr(cleaned_obj, attr_name, cleaned_value) + if hasattr(cleaned_obj, '__setattr__'): + try: + setattr(cleaned_obj, attr_name, cleaned_value) @@ -15,6 +36,33 @@ index 5a66e719..3d411c17 100644 + print(f"Skipping attribute '{attr_name}' due to Union type") + continue + raise ++ else: ++ # Fallback for objects without __setattr__ override ++ setattr(cleaned_obj, attr_name, cleaned_value) + + return cleaned_obj + +diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py +index 1ca5b18bd..4797e2e42 100644 +--- a/src/megatron/bridge/peft/utils.py ++++ b/src/megatron/bridge/peft/utils.py +@@ -33,6 +33,7 @@ from megatron.core.transformer.moe.router import TopKRouter + + from megatron.bridge.utils.import_utils import safe_import_from + ++from relax.utils.device import is_npu_available + + TEColumnParallelLinear, HAVE_TE_COL_LINEAR = safe_import_from( + "megatron.core.extensions.transformer_engine", "TEColumnParallelLinear" +@@ -62,7 +63,10 @@ HAVE_TE = all( + ) + ) + +-MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm") ++if is_npu_available: ++ MixedFusedLayerNorm, HAVE_APEX = None, False ++else: ++ MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm") - # Set the cleaned value (or None if it was removed) - setattr(cleaned_obj, attr_name, cleaned_value) + TECL = (TEColumnParallelLinear, TELayerNormColumnParallelLinear, TEColumnParallelGroupedLinear) + TERL = (TERowParallelLinear, TERowParallelGroupedLinear) diff --git a/docker/npu_patch/mindspeed-bridge.patch b/docker/npu_patch/mindspeed-bridge.patch index 327ff7afa..11d7b4924 100644 --- a/docker/npu_patch/mindspeed-bridge.patch +++ b/docker/npu_patch/mindspeed-bridge.patch @@ -19,64 +19,98 @@ index 99ccba5..8fa07a1 100644 # Set the cleaned value (or None if it was removed) setattr(cleaned_obj, attr_name, cleaned_value) diff --git a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py -index 818af29..ef97078 100644 +index 19a3b9d..1417c33 100644 --- a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py +++ b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py -@@ -191,7 +191,7 @@ class GatedDeltaNet(MegatronModule): +@@ -37,12 +37,11 @@ from megatron.core.ssm.gated_delta_net import ( + ) + + try: +- from causal_conv1d import causal_conv1d ++ import fla_npu ++ from fla.modules.convolution import causal_conv1d + except ImportError: + causal_conv1d = None +- causal_conv1d_update = None + +-from mindspeed_ops.api.triton.l2norm import l2norm + from mindspeed_bridge.models.qwen_vl.modelling_qwen3_vl.chunk_gated_delta_rule import ( + torch_chunk_gated_delta_rule, + ) +@@ -59,6 +58,12 @@ except ImportError: + + from mindspeed_bridge.models.qwen_vl.modelling_qwen3_vl.flash_gated_delta_rule import flash_gated_delta_rule + ++def naive_l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): ++ """This function is intended to align with the l2norm implementation in the FLA library.""" ++ original_dtype = x.dtype ++ inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) ++ # Counteract verl's autocast promotion (bf16 -> fp32) by restoring original dtype ++ return (x * inv_norm).to(original_dtype) + + class GatedDeltaNet(MegatronModule): + """Gated Delta Net (GDN) layer class +@@ -191,8 +196,8 @@ class GatedDeltaNet(MegatronModule): ) setattr(self.A_log, "tensor_model_parallel", True) - if HAVE_FLA and self.use_triton_gdn: +- self.gated_delta_rule = chunk_gated_delta_rule + if HAVE_FLA: - self.gated_delta_rule = chunk_gated_delta_rule ++ self.gated_delta_rule = flash_gated_delta_rule elif self.use_ascend_gdn: self.gated_delta_rule = flash_gated_delta_rule -diff --git a/mindspeed_bridge/models/qwen_vl/qwen35_vl_provider.py b/mindspeed_bridge/models/qwen_vl/qwen35_vl_provider.py -index d495c47..f08ea49 100644 ---- a/mindspeed_bridge/models/qwen_vl/qwen35_vl_provider.py -+++ b/mindspeed_bridge/models/qwen_vl/qwen35_vl_provider.py -@@ -402,22 +402,6 @@ class Qwen35VLModelProvider(GPTModelProvider): - self.vision_config = Qwen3_5VisionConfig() - super().__post_init__() + else: +@@ -474,7 +479,7 @@ class GatedDeltaNet(MegatronModule): + beta=beta, + initial_state=None, + output_final_state=False, +- use_qk_l2norm_in_kernel=True, ++ use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens_q, + ) + nvtx_range_pop(suffix="gated_delta_rule") +@@ -555,7 +560,7 @@ class GatedDeltaNet(MegatronModule): -- def finalize(self) -> None: -- self.validate_parallelism() -- super().finalize() -- -- def validate_parallelism(self): -- """Validate that parallelism settings are compatible with this model's architecture. -- -- Call this after mutating parallelism attributes (e.g. tensor_model_parallel_size) -- on an already-constructed provider, since finalize() only runs once before provide(). -- """ -- if self.num_query_groups < self.tensor_model_parallel_size: -- raise ValueError( -- f"TP size {self.tensor_model_parallel_size} should be less than or equal to " -- f"num_query_groups {self.num_query_groups}. Please use a smaller TP size." -- ) -- - def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen3VLModel: - """Provide a Qwen3.5 VL dense model instance with vision and language components.""" - language_transformer_config = self -@@ -595,21 +579,6 @@ class Qwen35VLMoEModelProvider(GPTModelProvider): - self.vision_config = Qwen3_5MoeVisionConfig() - super().__post_init__() + # Apply L2 norm to query and key + if self.use_qk_l2norm: +- query_key = l2norm(query_key.contiguous()) ++ query_key = naive_l2norm(query_key.contiguous()) -- def finalize(self) -> None: -- self.validate_parallelism() -- super().finalize() -- -- def validate_parallelism(self): -- """Validate that parallelism settings are compatible with this model's architecture. -- -- Call this after mutating parallelism attributes (e.g. tensor_model_parallel_size) -- on an already-constructed provider, since finalize() only runs once before provide(). -- """ -- if self.num_query_groups < self.tensor_model_parallel_size: -- raise ValueError( -- f"TP size {self.tensor_model_parallel_size} should be less than or equal to " -- f"num_query_groups {self.num_query_groups}. Please use a smaller TP size." -- ) + # Split query and key + split_size = self.qk_dim_local_tp // self.key_head_dim // self.cp_size +diff --git a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/model.py +index a6f310a..d468fa5 100644 +--- a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/model.py ++++ b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/model.py +@@ -501,7 +501,7 @@ class Qwen3VLModel(MegatronModule): + combined_embeddings = split_data_cp_rank(combined_embeddings, cp_size, 0, cp_rank) + if packed_seq_params is not None: + if attention_mask is None: +- attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device) ++ attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device) + input_ids_thd, _ = preprocess_packed_seqs( + input_ids, + attention_mask, +@@ -560,7 +560,7 @@ class Qwen3VLModel(MegatronModule): + # convert lm_input_ids to THD format so it matches position_ids. + if packed_seq_params is not None: + if attention_mask is None: +- attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device) ++ attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device) + lm_input_ids, _ = preprocess_packed_seqs( + input_ids, + attention_mask, +diff --git a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py +index 5d03f85..2488ef7 100644 +--- a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py ++++ b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py +@@ -56,6 +56,8 @@ def get_vision_model_config(hf_config, megatron_config=None): + ffn_hidden_size=hf_config.intermediate_size, + add_bias_linear=True, + add_qkv_bias=True, ++ moe_permute_fusion=False, # ← 新增：vision model 没有 MoE，不需要 permute fusion ++ use_fused_moe_token_permute_and_unpermute=False, # ← 新增：同上 + ) - def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen3VLModel: - """Provide a Qwen3.5 VL model instance with vision and language components. + # apply text model config to vision model config diff --git a/docker/npu_patch/mindspeed-ops.patch b/docker/npu_patch/mindspeed-ops.patch new file mode 100644 index 000000000..9dae51dd4 --- /dev/null +++ b/docker/npu_patch/mindspeed-ops.patch @@ -0,0 +1,13 @@ +diff --git a/mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_bwd_dhu.py b/mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_bwd_dhu.py +index bf62b0b..1925eca 100644 +--- a/mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_bwd_dhu.py ++++ b/mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_bwd_dhu.py +@@ -21,7 +21,7 @@ from mindspeed_ops.api.triton.utils import prepare_chunk_indices, prepare_chunk_ + ) + @triton.autotune( + configs=get_autotune_config(multibuffer_list=(True, False)), +- key=['H', 'K', 'V', 'BT', 'BV', 'USE_G', 'IS_VARLEN'], ++ key=['H', 'K', 'V', 'BT', 'BV'], + ) + @triton.jit(do_not_specialize=['T']) + def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( diff --git a/docker/npu_patch/mindspeed.patch b/docker/npu_patch/mindspeed.patch index 9f5f056ae..61c0c21e9 100644 --- a/docker/npu_patch/mindspeed.patch +++ b/docker/npu_patch/mindspeed.patch @@ -37,6 +37,20 @@ index c624fd0a..cef668af 100644 return parser +diff --git a/mindspeed/core/megatron_basic/arguments_basic.py b/mindspeed/core/megatron_basic/arguments_basic.py +index 87d71f81..aa6c219b 100644 +--- a/mindspeed/core/megatron_basic/arguments_basic.py ++++ b/mindspeed/core/megatron_basic/arguments_basic.py +@@ -147,7 +147,7 @@ def transformer_config_init_subclass(cls, **kwargs): + if callable(value) and not isinstance(value, type): + value = field(default_factory=value) + elif type(value) in mutable_types: +- value = field(default_factory=lambda: value) ++ value = field(default_factory=lambda v=value: v) + else: + value = value + setattr(cls, key, value) +\ No newline at end of file diff --git a/mindspeed/features_manager/functional/profile.py b/mindspeed/features_manager/functional/profile.py index 6450b41c..86584015 100644 --- a/mindspeed/features_manager/functional/profile.py diff --git a/docker/npu_patch/sgl-kernel-npu.patch b/docker/npu_patch/sgl-kernel-npu.patch new file mode 100644 index 000000000..730580178 --- /dev/null +++ b/docker/npu_patch/sgl-kernel-npu.patch @@ -0,0 +1,1707 @@ +From 1732669a033f5fe5aa178c305ea1e4ba379fdca1 Mon Sep 17 00:00:00 2001 +From: wuqiwei +Date: Wed, 5 Aug 2026 12:32:43 +0000 +Subject: [PATCH 1/3] torch.npu.empty_cache() + +--- + .../torch_memory_saver/python/torch_memory_saver/entrypoint.py | 2 ++ + 1 file changed, 2 insertions(+) + +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..146c011 100644 +--- a/contrib/torch_memory_saver/python/torch_memory_saver/entrypoint.py ++++ b/contrib/torch_memory_saver/python/torch_memory_saver/entrypoint.py +@@ -159,6 +159,7 @@ class _TorchMemorySaverImpl: + # only be released after the memory region is resumed and empty_cache() is invoked. + torch_npu._C._npu_releasePool(torch.npu.current_device(), pool.id) + del pool ++ torch.npu.empty_cache() + finally: + self._binary_wrapper.cdll.tms_set_interesting_region(True) + +@@ -176,3 +177,4 @@ def _sanity_checks(): + raise RuntimeError( + "TorchMemorySaver is disabled for the current process because expandable_segments is not supported yet." + ) ++ +\ No newline at end of file +-- +2.34.1 + + +From 2aca54d20695004a4a8c7222d333fadd78d4d589 Mon Sep 17 00:00:00 2001 +From: cl-vv-h +Date: Sat, 25 Jul 2026 09:21:13 +0000 +Subject: [PATCH 2/3] feat: add non-greedy MTP sampling kernels + +--- + .../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 ++++++++++++++++++ + 7 files changed, 1554 insertions(+) + 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/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 ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/sample/__init__.py +@@ -0,0 +1,13 @@ ++from sgl_kernel_npu.sample.chain_speculative_sampling import ( ++ chain_speculative_sampling_rejection, ++) ++from sgl_kernel_npu.sample.probability import top_k_top_p_renorm_probs ++from sgl_kernel_npu.sample.tree_speculative_sampling_target_only import ( ++ tree_speculative_sampling_target_only, ++) ++ ++__all__ = [ ++ "chain_speculative_sampling_rejection", ++ "top_k_top_p_renorm_probs", ++ "tree_speculative_sampling_target_only", ++] +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/sample/chain_speculative_sampling.py b/python/sgl_kernel_npu/sgl_kernel_npu/sample/chain_speculative_sampling.py +new file mode 100644 +index 0000000..e8340da +--- /dev/null ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/sample/chain_speculative_sampling.py +@@ -0,0 +1,354 @@ ++import torch ++import triton ++import triton.language as tl ++ ++ ++@triton.jit ++def _chain_rejection_accept_kernel( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ target_probs, ++ draft_probs, ++ metadata, ++ num_draft_tokens: tl.constexpr, ++ num_speculative_tokens: tl.constexpr, ++ num_draft_prob_rows: tl.constexpr, ++ vocab_size: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ row_offset = req_idx * num_draft_tokens ++ ++ cur_prob_row = tl.full((), 0, tl.int64) ++ last_accepted_idx = tl.load(retrive_index + row_offset).to(tl.int64) ++ num_accepted = 0 ++ active = tl.full((), 1, tl.int32) ++ ++ tl.store(accept_index + req_idx * num_speculative_tokens, last_accepted_idx) ++ ++ # Linear Leviathan/Chen verification. Candidate 0 is the root; candidate ++ # step uses probability row step - 1 until a rejection terminates the chain. ++ for step in range(1, num_draft_tokens): ++ if active == 1: ++ draft_token = tl.load(candidates + row_offset + step).to(tl.int64) ++ target_offset = ( ++ (row_offset + cur_prob_row) * vocab_size + draft_token ++ ) ++ draft_offset = ( ++ (req_idx * num_draft_prob_rows + cur_prob_row) * vocab_size ++ + draft_token ++ ) ++ target_prob = tl.load(target_probs + target_offset).to(tl.float32) ++ draft_prob = tl.load(draft_probs + draft_offset).to(tl.float32) ++ coin = tl.load(uniform_samples + row_offset + step - 1).to( ++ tl.float32 ++ ) ++ ++ if coin * draft_prob < target_prob: ++ tl.store(predicts + last_accepted_idx, draft_token) ++ num_accepted += 1 ++ draft_idx = tl.load(retrive_index + row_offset + step).to( ++ tl.int64 ++ ) ++ tl.store( ++ accept_index ++ + req_idx * num_speculative_tokens ++ + num_accepted, ++ draft_idx, ++ ) ++ last_accepted_idx = draft_idx ++ # Keep this loop-carried value int64 across both branches. ++ # Triton infers the constexpr loop variable `step` as int32. ++ cur_prob_row = tl.full((), step, tl.int64) ++ else: ++ active = 0 ++ ++ tl.store(accept_token_num + req_idx, num_accepted) ++ ++ # metadata = [target row, output slot, all drafts accepted]. ++ metadata_offset = req_idx * 3 ++ tl.store(metadata + metadata_offset, cur_prob_row) ++ tl.store(metadata + metadata_offset + 1, last_accepted_idx) ++ tl.store(metadata + metadata_offset + 2, active) ++ ++ ++@triton.jit ++def _chain_rejection_block_sum_kernel( ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ num_draft_tokens: tl.constexpr, ++ num_draft_prob_rows: tl.constexpr, ++ vocab_size: tl.constexpr, ++ vocab_block_size: tl.constexpr, ++ num_vocab_blocks: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ block_idx = tl.program_id(1) ++ vocab_offsets = block_idx * vocab_block_size + tl.arange(0, vocab_block_size) ++ vocab_mask = vocab_offsets < vocab_size ++ ++ metadata_offset = req_idx * 3 ++ target_row = tl.load(metadata + metadata_offset).to(tl.int64) ++ all_accepted = tl.load(metadata + metadata_offset + 2).to(tl.int32) ++ target_offset = (req_idx * num_draft_tokens + target_row) * vocab_size ++ target = tl.load( ++ target_probs + target_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ ++ if all_accepted == 1: ++ residual = target ++ else: ++ draft_offset = ( ++ req_idx * num_draft_prob_rows + target_row ++ ) * vocab_size ++ draft = tl.load( ++ draft_probs + draft_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ residual = tl.maximum(target - draft, 0.0) ++ ++ tl.store( ++ block_sums + req_idx * num_vocab_blocks + block_idx, ++ tl.sum(residual, axis=0), ++ ) ++ ++ ++@triton.jit ++def _chain_rejection_sample_kernel( ++ predicts, ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ uniform_samples_for_final_sampling, ++ num_draft_tokens: tl.constexpr, ++ num_draft_prob_rows: tl.constexpr, ++ vocab_size: tl.constexpr, ++ vocab_block_size: tl.constexpr, ++ num_vocab_blocks: tl.constexpr, ++ pad_num_vocab_blocks: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ metadata_offset = req_idx * 3 ++ target_row = tl.load(metadata + metadata_offset).to(tl.int64) ++ output_idx = tl.load(metadata + metadata_offset + 1).to(tl.int64) ++ all_accepted = tl.load(metadata + metadata_offset + 2).to(tl.int32) ++ ++ block_offsets = tl.arange(0, pad_num_vocab_blocks) ++ block_mask = block_offsets < num_vocab_blocks ++ sums = tl.load( ++ block_sums + req_idx * num_vocab_blocks + block_offsets, ++ mask=block_mask, ++ other=0.0, ++ ).to(tl.float32) ++ block_cdf = tl.cumsum(sums, axis=0) ++ total = tl.sum(sums, axis=0) ++ coin = tl.load(uniform_samples_for_final_sampling + req_idx).to(tl.float32) ++ target_value = coin * total ++ ++ selected_block = tl.sum( ++ ((block_cdf <= target_value) & block_mask).to(tl.int32), axis=0 ++ ) ++ selected_block = tl.minimum(selected_block, num_vocab_blocks - 1) ++ prefix_sum = tl.sum( ++ tl.where(block_offsets < selected_block, sums, 0.0), axis=0 ++ ) ++ local_target = tl.maximum(target_value - prefix_sum, 0.0) ++ ++ local_offsets = tl.arange(0, vocab_block_size) ++ vocab_offsets = selected_block * vocab_block_size + local_offsets ++ vocab_mask = vocab_offsets < vocab_size ++ target_offset = (req_idx * num_draft_tokens + target_row) * vocab_size ++ target = tl.load( ++ target_probs + target_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ if all_accepted == 1: ++ residual = target ++ else: ++ draft_offset = ( ++ req_idx * num_draft_prob_rows + target_row ++ ) * vocab_size ++ draft = tl.load( ++ draft_probs + draft_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ residual = tl.maximum(target - draft, 0.0) ++ ++ local_cdf = tl.cumsum(residual, axis=0) ++ local_index = tl.sum( ++ ((local_cdf <= local_target) & vocab_mask).to(tl.int32), axis=0 ++ ) ++ last_valid_local = tl.max( ++ tl.where((residual > 0.0) & vocab_mask, local_offsets, -1), axis=0 ++ ) ++ valid_local_count = tl.minimum( ++ vocab_block_size, ++ vocab_size - selected_block * vocab_block_size, ++ ) ++ sampled_local = tl.where( ++ local_index < valid_local_count, ++ local_index, ++ last_valid_local, ++ ) ++ sampled_token = tl.where( ++ sampled_local >= 0, ++ selected_block * vocab_block_size + sampled_local, ++ vocab_size - 1, ++ ) ++ tl.store(predicts + output_idx, tl.minimum(sampled_token, vocab_size - 1)) ++ ++ ++def chain_speculative_sampling_rejection( ++ predicts: torch.Tensor, ++ accept_index: torch.Tensor, ++ accept_token_num: torch.Tensor, ++ candidates: torch.Tensor, ++ retrive_index: torch.Tensor, ++ retrive_next_token: torch.Tensor, ++ retrive_next_sibling: torch.Tensor, ++ uniform_samples: torch.Tensor, ++ uniform_samples_for_final_sampling: torch.Tensor, ++ target_probs: torch.Tensor, ++ draft_probs: torch.Tensor, ++ threshold_single: float = 1.0, ++ threshold_acc: float = 1.0, ++ deterministic: bool = True, ++) -> None: ++ """NPU kernel implementation of classic chain rejection sampling.""" ++ del retrive_next_token, retrive_next_sibling ++ del threshold_single, threshold_acc, deterministic ++ ++ if candidates.ndim != 2 or target_probs.ndim != 3: ++ raise ValueError("candidates must be 2-D and target_probs must be 3-D") ++ batch_size, num_draft_tokens = candidates.shape ++ if batch_size == 0: ++ return ++ if num_draft_tokens == 0: ++ raise ValueError("num_draft_tokens must be positive") ++ if target_probs.shape[:2] != (batch_size, num_draft_tokens): ++ raise ValueError( ++ "target_probs shape must be [batch, num_draft_tokens, vocab_size]" ++ ) ++ if retrive_index.shape != candidates.shape: ++ raise ValueError("retrive_index shape must match candidates") ++ if accept_index.shape != candidates.shape: ++ raise ValueError( ++ "classic rejection sampling requires a topk=1 linear chain" ++ ) ++ if accept_token_num.shape != (batch_size,): ++ raise ValueError("accept_token_num must have shape [batch]") ++ if predicts.ndim != 1: ++ raise ValueError("predicts must be 1-D") ++ if uniform_samples.shape != candidates.shape: ++ raise ValueError("uniform_samples shape must match candidates") ++ if uniform_samples_for_final_sampling.shape != (batch_size,): ++ raise ValueError( ++ "uniform_samples_for_final_sampling must have shape [batch]" ++ ) ++ if draft_probs is None or draft_probs.ndim != 3: ++ raise ValueError("draft_probs must be a 3-D tensor") ++ if draft_probs.shape[0] != batch_size: ++ raise ValueError("draft_probs batch size must match candidates") ++ if draft_probs.shape[1] < max(num_draft_tokens - 1, 1): ++ raise ValueError("draft_probs does not contain every proposal row") ++ if draft_probs.shape[-1] != target_probs.shape[-1]: ++ raise ValueError("draft_probs and target_probs vocab sizes must match") ++ if target_probs.dtype != torch.float32 or draft_probs.dtype != torch.float32: ++ raise TypeError("target_probs and draft_probs must be torch.float32") ++ if uniform_samples.dtype != torch.float32: ++ raise TypeError("uniform_samples must be torch.float32") ++ if uniform_samples_for_final_sampling.dtype != torch.float32: ++ raise TypeError("uniform_samples_for_final_sampling must be torch.float32") ++ integer_dtypes = ( ++ (predicts, torch.int32, "predicts"), ++ (accept_index, torch.int32, "accept_index"), ++ (accept_token_num, torch.int32, "accept_token_num"), ++ (candidates, torch.int64, "candidates"), ++ (retrive_index, torch.int64, "retrive_index"), ++ ) ++ for tensor, expected_dtype, name in integer_dtypes: ++ if tensor.dtype != expected_dtype: ++ raise TypeError(f"{name} must be {expected_dtype}") ++ tensors = ( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++ draft_probs, ++ ) ++ if any(tensor.device != target_probs.device for tensor in tensors): ++ raise ValueError("all tensors must be on the same NPU device") ++ if any(not tensor.is_contiguous() for tensor in tensors): ++ raise ValueError("all tensors must be contiguous") ++ ++ num_speculative_tokens = accept_index.shape[1] ++ num_draft_prob_rows = draft_probs.shape[1] ++ vocab_size = target_probs.shape[-1] ++ vocab_block_size = 2048 ++ num_vocab_blocks = triton.cdiv(vocab_size, vocab_block_size) ++ pad_num_vocab_blocks = triton.next_power_of_2(num_vocab_blocks) ++ ++ metadata = torch.empty( ++ (batch_size, 3), dtype=torch.int64, device=target_probs.device ++ ) ++ block_sums = torch.empty( ++ (batch_size, num_vocab_blocks), ++ dtype=torch.float32, ++ device=target_probs.device, ++ ) ++ ++ _chain_rejection_accept_kernel[(batch_size,)]( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ target_probs, ++ draft_probs, ++ metadata, ++ num_draft_tokens=num_draft_tokens, ++ num_speculative_tokens=num_speculative_tokens, ++ num_draft_prob_rows=num_draft_prob_rows, ++ vocab_size=vocab_size, ++ ) ++ _chain_rejection_block_sum_kernel[(batch_size, num_vocab_blocks)]( ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ num_draft_tokens=num_draft_tokens, ++ num_draft_prob_rows=num_draft_prob_rows, ++ vocab_size=vocab_size, ++ vocab_block_size=vocab_block_size, ++ num_vocab_blocks=num_vocab_blocks, ++ ) ++ _chain_rejection_sample_kernel[(batch_size,)]( ++ predicts, ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ uniform_samples_for_final_sampling, ++ num_draft_tokens=num_draft_tokens, ++ num_draft_prob_rows=num_draft_prob_rows, ++ vocab_size=vocab_size, ++ vocab_block_size=vocab_block_size, ++ num_vocab_blocks=num_vocab_blocks, ++ pad_num_vocab_blocks=pad_num_vocab_blocks, ++ ) +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/sample/probability.py b/python/sgl_kernel_npu/sgl_kernel_npu/sample/probability.py +new file mode 100644 +index 0000000..cadddd0 +--- /dev/null ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/sample/probability.py +@@ -0,0 +1,40 @@ ++import torch ++ ++ ++def top_k_top_p_renorm_probs( ++ probs: torch.Tensor, ++ top_ks: torch.Tensor, ++ top_ps: torch.Tensor, ++ need_top_k_sampling: bool, ++ need_top_p_sampling: bool, ++) -> torch.Tensor: ++ """Apply the same sequential top-k then top-p policy used by SGLang GPU.""" ++ if not need_top_k_sampling and not need_top_p_sampling: ++ return probs ++ ++ vocab_size = probs.shape[-1] ++ sorted_probs, sorted_indices = probs.sort(dim=-1, descending=True) ++ ++ if need_top_k_sampling: ++ top_ks = top_ks.to(device=probs.device, dtype=torch.long).clamp( ++ min=1, max=vocab_size ++ ) ++ positions = torch.arange(vocab_size, device=probs.device).view(1, -1) ++ sorted_probs.masked_fill_(positions >= top_ks.view(-1, 1), 0.0) ++ sorted_probs.div_( ++ sorted_probs.sum(dim=-1, keepdim=True).clamp_min_(1e-20) ++ ) ++ ++ if need_top_p_sampling: ++ top_ps = top_ps.to(device=probs.device, dtype=probs.dtype) ++ cumulative_probs = sorted_probs.cumsum(dim=-1) ++ sorted_probs.masked_fill_( ++ cumulative_probs - sorted_probs > top_ps.view(-1, 1), 0.0 ++ ) ++ sorted_probs.div_( ++ sorted_probs.sum(dim=-1, keepdim=True).clamp_min_(1e-20) ++ ) ++ ++ return torch.zeros_like(probs).scatter_( ++ dim=-1, index=sorted_indices, src=sorted_probs ++ ) +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/sample/tree_speculative_sampling_target_only.py b/python/sgl_kernel_npu/sgl_kernel_npu/sample/tree_speculative_sampling_target_only.py +new file mode 100644 +index 0000000..e34baab +--- /dev/null ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/sample/tree_speculative_sampling_target_only.py +@@ -0,0 +1,383 @@ ++import torch ++import triton ++import triton.language as tl ++ ++ ++@triton.jit ++def _tree_target_only_accept_kernel( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ retrive_next_token, ++ retrive_next_sibling, ++ uniform_samples, ++ target_probs, ++ rejected_probs, ++ metadata, ++ threshold_single, ++ threshold_acc, ++ num_draft_tokens: tl.constexpr, ++ num_speculative_tokens: tl.constexpr, ++ vocab_size: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ row_offset = req_idx * num_draft_tokens ++ ++ cur_prob_row = tl.full((), 0, tl.int64) ++ cur_node = tl.full((), 0, tl.int64) ++ last_accepted_idx = tl.load(retrive_index + row_offset).to(tl.int64) ++ coin = tl.load(uniform_samples + row_offset).to(tl.float32) ++ num_accepted = 0 ++ path_active = tl.full((), 1, tl.int32) ++ ++ tl.store(accept_index + req_idx * num_speculative_tokens, last_accepted_idx) ++ ++ # This is the same breadth-at-each-depth traversal used by the CUDA kernel: ++ # descend to the first child, then walk siblings until one is accepted. ++ for _depth in range(1, num_speculative_tokens): ++ accepted_at_depth = tl.full((), 0, tl.int32) ++ prob_acc = tl.full((), 0.0, tl.float32) ++ ++ if path_active == 1: ++ cur_node = tl.load( ++ retrive_next_token + row_offset + cur_node ++ ).to(tl.int64) ++ if cur_node == -1: ++ path_active = 0 ++ ++ # The loop is bounded by the number of tree nodes. It terminates ++ # logically when a child is accepted or the sibling list reaches -1. ++ for _sibling in range(0, num_draft_tokens): ++ if ( ++ (path_active == 1) ++ & (accepted_at_depth == 0) ++ & (cur_node != -1) ++ ): ++ draft_token = tl.load( ++ candidates + row_offset + cur_node ++ ).to(tl.int64) ++ draft_idx = tl.load( ++ retrive_index + row_offset + cur_node ++ ).to(tl.int64) ++ prob_offset = ( ++ (row_offset + cur_prob_row) * vocab_size + draft_token ++ ) ++ target_prob_single = tl.load( ++ target_probs + prob_offset ++ ).to(tl.float32) ++ prob_acc += target_prob_single ++ ++ accepted = (coin <= prob_acc / threshold_acc) | ( ++ target_prob_single >= threshold_single ++ ) ++ if accepted: ++ tl.store(predicts + last_accepted_idx, draft_token) ++ num_accepted += 1 ++ tl.store( ++ accept_index ++ + req_idx * num_speculative_tokens ++ + num_accepted, ++ draft_idx, ++ ) ++ last_accepted_idx = draft_idx ++ cur_prob_row = cur_node ++ coin = tl.load( ++ uniform_samples + row_offset + cur_node ++ ).to(tl.float32) ++ accepted_at_depth = 1 ++ else: ++ # The CUDA target-only kernel stores the rejected sibling's ++ # target probability in draft_probs and later samples from ++ # relu(target_probs - draft_probs). ++ tl.store(rejected_probs + prob_offset, target_prob_single) ++ cur_node = tl.load( ++ retrive_next_sibling + row_offset + cur_node ++ ).to(tl.int64) ++ ++ if accepted_at_depth == 0: ++ path_active = 0 ++ ++ tl.store(accept_token_num + req_idx, num_accepted) ++ ++ # metadata = [final target-probability row, final output slot]. ++ metadata_offset = req_idx * 2 ++ tl.store(metadata + metadata_offset, cur_prob_row) ++ tl.store(metadata + metadata_offset + 1, last_accepted_idx) ++ ++ ++@triton.jit ++def _tree_target_only_block_sum_kernel( ++ target_probs, ++ rejected_probs, ++ metadata, ++ block_sums, ++ num_draft_tokens: tl.constexpr, ++ vocab_size: tl.constexpr, ++ vocab_block_size: tl.constexpr, ++ num_vocab_blocks: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ block_idx = tl.program_id(1) ++ vocab_offsets = block_idx * vocab_block_size + tl.arange(0, vocab_block_size) ++ vocab_mask = vocab_offsets < vocab_size ++ ++ target_row = tl.load(metadata + req_idx * 2).to(tl.int64) ++ probs_offset = (req_idx * num_draft_tokens + target_row) * vocab_size ++ target = tl.load( ++ target_probs + probs_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ rejected = tl.load( ++ rejected_probs + probs_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ residual = tl.maximum(target - rejected, 0.0) ++ block_sum = tl.sum(residual, axis=0) ++ tl.store(block_sums + req_idx * num_vocab_blocks + block_idx, block_sum) ++ ++ ++@triton.jit ++def _tree_target_only_sample_kernel( ++ predicts, ++ target_probs, ++ rejected_probs, ++ metadata, ++ block_sums, ++ uniform_samples_for_final_sampling, ++ num_draft_tokens: tl.constexpr, ++ vocab_size: tl.constexpr, ++ vocab_block_size: tl.constexpr, ++ num_vocab_blocks: tl.constexpr, ++ pad_num_vocab_blocks: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ metadata_offset = req_idx * 2 ++ target_row = tl.load(metadata + metadata_offset).to(tl.int64) ++ output_idx = tl.load(metadata + metadata_offset + 1).to(tl.int64) ++ ++ block_offsets = tl.arange(0, pad_num_vocab_blocks) ++ block_mask = block_offsets < num_vocab_blocks ++ sums = tl.load( ++ block_sums + req_idx * num_vocab_blocks + block_offsets, ++ mask=block_mask, ++ other=0.0, ++ ).to(tl.float32) ++ block_cdf = tl.cumsum(sums, axis=0) ++ total = tl.sum(sums, axis=0) ++ coin = tl.load(uniform_samples_for_final_sampling + req_idx).to(tl.float32) ++ target = coin * total ++ ++ selected_block = tl.sum( ++ ((block_cdf <= target) & block_mask).to(tl.int32), axis=0 ++ ) ++ selected_block = tl.minimum(selected_block, num_vocab_blocks - 1) ++ prefix_sum = tl.sum( ++ tl.where(block_offsets < selected_block, sums, 0.0), axis=0 ++ ) ++ local_target = tl.maximum(target - prefix_sum, 0.0) ++ ++ local_offsets = tl.arange(0, vocab_block_size) ++ vocab_offsets = selected_block * vocab_block_size + local_offsets ++ vocab_mask = vocab_offsets < vocab_size ++ probs_offset = (req_idx * num_draft_tokens + target_row) * vocab_size ++ target_probs_block = tl.load( ++ target_probs + probs_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ rejected_probs_block = tl.load( ++ rejected_probs + probs_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ residual = tl.maximum(target_probs_block - rejected_probs_block, 0.0) ++ ++ local_cdf = tl.cumsum(residual, axis=0) ++ local_index = tl.sum( ++ ((local_cdf <= local_target) & vocab_mask).to(tl.int32), axis=0 ++ ) ++ last_valid_local = tl.max( ++ tl.where((residual > 0.0) & vocab_mask, local_offsets, -1), axis=0 ++ ) ++ valid_local_count = tl.minimum( ++ vocab_block_size, ++ vocab_size - selected_block * vocab_block_size, ++ ) ++ sampled_local = tl.where( ++ local_index < valid_local_count, ++ local_index, ++ last_valid_local, ++ ) ++ sampled_token = tl.where( ++ sampled_local >= 0, ++ selected_block * vocab_block_size + sampled_local, ++ vocab_size - 1, ++ ) ++ sampled_token = tl.minimum(sampled_token, vocab_size - 1) ++ tl.store(predicts + output_idx, sampled_token) ++ ++ ++def tree_speculative_sampling_target_only( ++ predicts: torch.Tensor, ++ accept_index: torch.Tensor, ++ accept_token_num: torch.Tensor, ++ candidates: torch.Tensor, ++ retrive_index: torch.Tensor, ++ retrive_next_token: torch.Tensor, ++ retrive_next_sibling: torch.Tensor, ++ uniform_samples: torch.Tensor, ++ uniform_samples_for_final_sampling: torch.Tensor, ++ target_probs: torch.Tensor, ++ draft_probs: torch.Tensor, ++ threshold_single: float = 1.0, ++ threshold_acc: float = 1.0, ++ deterministic: bool = True, ++) -> None: ++ """NPU port of GPU target-only tree speculative sampling. ++ ++ ``draft_probs`` is scratch storage, matching the GPU API. The function ++ clears it and records rejected sibling probabilities before sampling from ++ ``relu(target_probs - draft_probs)`` on the final selected tree row. ++ """ ++ del deterministic ++ ++ if candidates.ndim != 2 or target_probs.ndim != 3: ++ raise ValueError("candidates must be 2-D and target_probs must be 3-D") ++ ++ batch_size, num_draft_tokens = candidates.shape ++ if batch_size == 0: ++ return ++ if num_draft_tokens == 0: ++ raise ValueError("num_draft_tokens must be positive") ++ if target_probs.shape[:2] != (batch_size, num_draft_tokens): ++ raise ValueError( ++ "target_probs shape must be [batch, num_draft_tokens, vocab_size]" ++ ) ++ tree_shapes = ( ++ retrive_index.shape, ++ retrive_next_token.shape, ++ retrive_next_sibling.shape, ++ uniform_samples.shape, ++ ) ++ if any(shape != candidates.shape for shape in tree_shapes): ++ raise ValueError("all tree-index and uniform tensors must match candidates") ++ if accept_index.ndim != 2 or accept_index.shape[0] != batch_size: ++ raise ValueError("accept_index must be [batch, max_tree_depth]") ++ num_speculative_tokens = accept_index.shape[1] ++ if not 1 <= num_speculative_tokens <= num_draft_tokens: ++ raise ValueError("max_tree_depth must be in [1, num_draft_tokens]") ++ if accept_token_num.shape != (batch_size,): ++ raise ValueError("accept_token_num must have shape [batch]") ++ if predicts.ndim != 1: ++ raise ValueError("predicts must be 1-D") ++ if uniform_samples_for_final_sampling.shape != (batch_size,): ++ raise ValueError( ++ "uniform_samples_for_final_sampling must have shape [batch]" ++ ) ++ if draft_probs.shape != target_probs.shape: ++ raise ValueError("draft_probs scratch must match target_probs") ++ if draft_probs.data_ptr() == target_probs.data_ptr(): ++ raise ValueError("draft_probs must not alias target_probs") ++ if target_probs.dtype != torch.float32 or draft_probs.dtype != torch.float32: ++ raise TypeError("target_probs and draft_probs must be torch.float32") ++ if uniform_samples.dtype != torch.float32: ++ raise TypeError("uniform_samples must be torch.float32") ++ if uniform_samples_for_final_sampling.dtype != torch.float32: ++ raise TypeError("uniform_samples_for_final_sampling must be torch.float32") ++ integer_dtypes = ( ++ (predicts, torch.int32, "predicts"), ++ (accept_index, torch.int32, "accept_index"), ++ (accept_token_num, torch.int32, "accept_token_num"), ++ (candidates, torch.int64, "candidates"), ++ (retrive_index, torch.int64, "retrive_index"), ++ (retrive_next_token, torch.int64, "retrive_next_token"), ++ (retrive_next_sibling, torch.int64, "retrive_next_sibling"), ++ ) ++ for tensor, expected_dtype, name in integer_dtypes: ++ if tensor.dtype != expected_dtype: ++ raise TypeError(f"{name} must be {expected_dtype}") ++ tensors = ( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ retrive_next_token, ++ retrive_next_sibling, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++ draft_probs, ++ ) ++ if any(tensor.device != target_probs.device for tensor in tensors): ++ raise ValueError("all tensors must be on the same NPU device") ++ if any(not tensor.is_contiguous() for tensor in tensors): ++ raise ValueError("all tensors must be contiguous") ++ if not 0.0 <= threshold_single <= 1.0: ++ raise ValueError("threshold_single must be in [0, 1]") ++ if not 0.0 <= threshold_acc <= 1.0: ++ raise ValueError("threshold_acc must be in [0, 1]") ++ ++ threshold_acc = max(float(threshold_acc), 1e-9) ++ vocab_size = target_probs.shape[-1] ++ vocab_block_size = 2048 ++ num_vocab_blocks = triton.cdiv(vocab_size, vocab_block_size) ++ pad_num_vocab_blocks = triton.next_power_of_2(num_vocab_blocks) ++ ++ # The CUDA call site passes zeros_like(target_probs). Clearing in the NPU ++ # wrapper makes the scratch contract explicit and permits empty_like callers. ++ draft_probs.zero_() ++ metadata = torch.empty( ++ (batch_size, 2), dtype=torch.int64, device=target_probs.device ++ ) ++ block_sums = torch.empty( ++ (batch_size, num_vocab_blocks), ++ dtype=torch.float32, ++ device=target_probs.device, ++ ) ++ ++ _tree_target_only_accept_kernel[(batch_size,)]( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ retrive_next_token, ++ retrive_next_sibling, ++ uniform_samples, ++ target_probs, ++ draft_probs, ++ metadata, ++ float(threshold_single), ++ threshold_acc, ++ num_draft_tokens=num_draft_tokens, ++ num_speculative_tokens=num_speculative_tokens, ++ vocab_size=vocab_size, ++ ) ++ _tree_target_only_block_sum_kernel[(batch_size, num_vocab_blocks)]( ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ num_draft_tokens=num_draft_tokens, ++ vocab_size=vocab_size, ++ vocab_block_size=vocab_block_size, ++ num_vocab_blocks=num_vocab_blocks, ++ ) ++ _tree_target_only_sample_kernel[(batch_size,)]( ++ predicts, ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ uniform_samples_for_final_sampling, ++ num_draft_tokens=num_draft_tokens, ++ vocab_size=vocab_size, ++ vocab_block_size=vocab_block_size, ++ num_vocab_blocks=num_vocab_blocks, ++ pad_num_vocab_blocks=pad_num_vocab_blocks, ++ ) +diff --git a/tests/python/sgl_kernel_npu/test_chain_speculative_sampling.py b/tests/python/sgl_kernel_npu/test_chain_speculative_sampling.py +new file mode 100644 +index 0000000..40b04be +--- /dev/null ++++ b/tests/python/sgl_kernel_npu/test_chain_speculative_sampling.py +@@ -0,0 +1,138 @@ ++import torch ++import torch_npu # noqa: F401 ++ ++from sgl_kernel_npu.sample import chain_speculative_sampling_rejection ++ ++ ++def chain_rejection_reference( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++ draft_probs, ++): ++ batch_size, num_draft_tokens = candidates.shape ++ for req_idx in range(batch_size): ++ cur_prob_row = 0 ++ last_accepted_idx = int(retrive_index[req_idx, 0]) ++ accept_index[req_idx, 0] = last_accepted_idx ++ num_accepted = 0 ++ all_accepted = True ++ ++ for step in range(1, num_draft_tokens): ++ draft_token = int(candidates[req_idx, step]) ++ p = float(target_probs[req_idx, cur_prob_row, draft_token]) ++ q = float(draft_probs[req_idx, cur_prob_row, draft_token]) ++ coin = float(uniform_samples[req_idx, step - 1]) ++ if coin * q < p: ++ predicts[last_accepted_idx] = draft_token ++ num_accepted += 1 ++ last_accepted_idx = int(retrive_index[req_idx, step]) ++ accept_index[req_idx, num_accepted] = last_accepted_idx ++ cur_prob_row = step ++ else: ++ all_accepted = False ++ break ++ ++ accept_token_num[req_idx] = num_accepted ++ residual = target_probs[req_idx, cur_prob_row].clone() ++ if not all_accepted: ++ residual.sub_(draft_probs[req_idx, cur_prob_row]).clamp_min_(0.0) ++ target = float(uniform_samples_for_final_sampling[req_idx]) * float( ++ residual.sum() ++ ) ++ sampled_token = int((residual.cumsum(0) <= target).sum()) ++ if sampled_token == residual.numel(): ++ positive = torch.nonzero(residual > 0.0).flatten() ++ sampled_token = ( ++ int(positive[-1]) if positive.numel() else residual.numel() - 1 ++ ) ++ predicts[last_accepted_idx] = sampled_token ++ ++ ++def test_chain_rejection_matches_gpu_algorithm(): ++ batch_size, num_draft_tokens, vocab_size = 2, 4, 11 ++ candidates = torch.tensor([[0, 2, 3, 4], [0, 5, 6, 7]]) ++ retrive_index = torch.arange(batch_size * num_draft_tokens).view( ++ batch_size, num_draft_tokens ++ ) ++ target_probs = torch.softmax( ++ torch.tensor( ++ [ ++ [ ++ [0.1, 0.2, 2.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], ++ [0.1, 0.2, 0.1, 2.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], ++ [0.1, 0.2, 0.1, 0.1, 2.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], ++ [0.1, 0.2, 0.1, 0.1, 2.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], ++ ], ++ [ ++ [0.1, 0.1, 0.1, 0.1, 0.1, 0.2, 0.1, 2.0, 0.1, 0.1, 0.1], ++ [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 2.0, 0.2, 0.1, 0.1, 0.1], ++ [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 2.0, 0.1, 0.1, 0.1], ++ [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 2.0, 0.1, 0.1, 0.1], ++ ], ++ ], ++ dtype=torch.float32, ++ ), ++ dim=-1, ++ ) ++ draft_probs = torch.softmax( ++ torch.tensor( ++ [ ++ [[0.1] * vocab_size, [0.1] * vocab_size, [0.1] * vocab_size], ++ [[0.1] * vocab_size, [0.1] * vocab_size, [0.1] * vocab_size], ++ ], ++ dtype=torch.float32, ++ ), ++ dim=-1, ++ ) ++ draft_probs[1, 0, 5] = 0.9 ++ draft_probs[1, 0] /= draft_probs[1, 0].sum() ++ uniforms = torch.tensor([[0.1, 0.1, 0.1, 0.0], [0.99, 0.0, 0.0, 0.0]]) ++ final_uniforms = torch.tensor([0.37, 0.61]) ++ ++ expected_predicts = torch.full( ++ (batch_size * num_draft_tokens,), -1, dtype=torch.int32 ++ ) ++ expected_accept_index = torch.full( ++ (batch_size, num_draft_tokens), -1, dtype=torch.int32 ++ ) ++ expected_accept_num = torch.zeros(batch_size, dtype=torch.int32) ++ chain_rejection_reference( ++ expected_predicts, ++ expected_accept_index, ++ expected_accept_num, ++ candidates, ++ retrive_index, ++ uniforms, ++ final_uniforms, ++ target_probs, ++ draft_probs, ++ ) ++ ++ predicts = torch.full_like(expected_predicts, -1, device="npu") ++ accept_index = torch.full_like(expected_accept_index, -1, device="npu") ++ accept_num = torch.zeros_like(expected_accept_num, device="npu") ++ next_token = torch.full_like(candidates, -1, device="npu") ++ next_sibling = torch.full_like(candidates, -1, device="npu") ++ chain_speculative_sampling_rejection( ++ predicts, ++ accept_index, ++ accept_num, ++ candidates.npu(), ++ retrive_index.npu(), ++ next_token, ++ next_sibling, ++ uniforms.npu(), ++ final_uniforms.npu(), ++ target_probs.npu(), ++ draft_probs.npu(), ++ ) ++ ++ torch.testing.assert_close(predicts.cpu(), expected_predicts) ++ torch.testing.assert_close(accept_index.cpu(), expected_accept_index) ++ torch.testing.assert_close(accept_num.cpu(), expected_accept_num) +diff --git a/tests/python/sgl_kernel_npu/test_speculative_probability.py b/tests/python/sgl_kernel_npu/test_speculative_probability.py +new file mode 100644 +index 0000000..563f174 +--- /dev/null ++++ b/tests/python/sgl_kernel_npu/test_speculative_probability.py +@@ -0,0 +1,31 @@ ++import torch ++ ++from sgl_kernel_npu.sample.probability import top_k_top_p_renorm_probs ++ ++ ++def test_top_k_top_p_renorm_matches_sequential_reference(): ++ torch.manual_seed(7) ++ probs = torch.softmax(torch.randn(4, 97), dim=-1) ++ top_ks = torch.tensor([1, 7, 31, 97]) ++ top_ps = torch.tensor([0.3, 0.75, 0.95, 1.0]) ++ ++ actual = top_k_top_p_renorm_probs( ++ probs, top_ks, top_ps, True, True ++ ) ++ ++ sorted_probs, sorted_indices = probs.sort(dim=-1, descending=True) ++ positions = torch.arange(probs.shape[-1]).view(1, -1) ++ sorted_probs[positions >= top_ks.view(-1, 1)] = 0.0 ++ sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True) ++ top_k_probs = torch.zeros_like(probs).scatter( ++ -1, sorted_indices, sorted_probs ++ ) ++ sorted_probs, sorted_indices = top_k_probs.sort(dim=-1, descending=True) ++ cumulative = sorted_probs.cumsum(dim=-1) ++ sorted_probs[cumulative - sorted_probs > top_ps.view(-1, 1)] = 0.0 ++ sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True) ++ expected = torch.zeros_like(probs).scatter( ++ -1, sorted_indices, sorted_probs ++ ) ++ ++ torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-7) +diff --git a/tests/python/sgl_kernel_npu/test_tree_speculative_sampling_target_only.py b/tests/python/sgl_kernel_npu/test_tree_speculative_sampling_target_only.py +new file mode 100644 +index 0000000..3c12d53 +--- /dev/null ++++ b/tests/python/sgl_kernel_npu/test_tree_speculative_sampling_target_only.py +@@ -0,0 +1,595 @@ ++import argparse ++import time ++ ++import pytest ++import torch ++import torch_npu # noqa: F401 ++ ++from sgl_kernel_npu.sample import tree_speculative_sampling_target_only ++ ++ ++def target_only_tree_reference( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ retrive_next_token, ++ retrive_next_sibling, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++ rejected_probs, ++ threshold_single, ++ threshold_acc, ++): ++ """CPU reference translated directly from the GPU CUDA kernel.""" ++ batch_size, num_draft_tokens = candidates.shape ++ num_speculative_tokens = accept_index.shape[1] ++ threshold_acc = max(float(threshold_acc), 1e-9) ++ rejected_probs.zero_() ++ ++ for req_idx in range(batch_size): ++ cur_prob_row = 0 ++ cur_node = 0 ++ coin = float(uniform_samples[req_idx, 0]) ++ last_accepted_idx = int(retrive_index[req_idx, 0]) ++ accept_index[req_idx, 0] = last_accepted_idx ++ num_accepted = 0 ++ ++ for _ in range(1, num_speculative_tokens): ++ cur_node = int(retrive_next_token[req_idx, cur_node]) ++ prob_acc = 0.0 ++ while cur_node != -1: ++ draft_idx = int(retrive_index[req_idx, cur_node]) ++ draft_token = int(candidates[req_idx, cur_node]) ++ target_prob = float( ++ target_probs[req_idx, cur_prob_row, draft_token] ++ ) ++ prob_acc += target_prob ++ if ( ++ coin <= prob_acc / threshold_acc ++ or target_prob >= threshold_single ++ ): ++ predicts[last_accepted_idx] = draft_token ++ num_accepted += 1 ++ accept_index[req_idx, num_accepted] = draft_idx ++ last_accepted_idx = draft_idx ++ cur_prob_row = cur_node ++ coin = float(uniform_samples[req_idx, cur_node]) ++ break ++ ++ rejected_probs[req_idx, cur_prob_row, draft_token] = target_prob ++ cur_node = int(retrive_next_sibling[req_idx, cur_node]) ++ ++ if cur_node == -1: ++ break ++ ++ accept_token_num[req_idx] = num_accepted ++ residual = ( ++ target_probs[req_idx, cur_prob_row] ++ - rejected_probs[req_idx, cur_prob_row] ++ ).clamp_min(0.0) ++ target = float(uniform_samples_for_final_sampling[req_idx]) * float( ++ residual.sum() ++ ) ++ sampled_token = int((residual.cumsum(0) <= target).sum()) ++ if sampled_token == residual.numel(): ++ positive = torch.nonzero(residual > 0.0).flatten() ++ sampled_token = ( ++ int(positive[-1]) if positive.numel() else residual.numel() - 1 ++ ) ++ predicts[last_accepted_idx] = sampled_token ++ ++ return predicts, accept_index, accept_token_num, rejected_probs ++ ++ ++def target_only_chain_reference( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++ threshold_single, ++ threshold_acc, ++): ++ batch_size, num_draft_tokens = candidates.shape ++ threshold_acc = max(float(threshold_acc), 1e-9) ++ ++ for req_idx in range(batch_size): ++ last_accepted_idx = int(retrive_index[req_idx, 0]) ++ accept_index[req_idx, 0] = last_accepted_idx ++ num_accepted = 0 ++ rejected_token = -1 ++ ++ for step in range(1, num_draft_tokens): ++ draft_token = int(candidates[req_idx, step]) ++ target_prob = float(target_probs[req_idx, step - 1, draft_token]) ++ coin = float(uniform_samples[req_idx, step - 1]) ++ if ( ++ coin <= target_prob / threshold_acc ++ or target_prob >= threshold_single ++ ): ++ predicts[last_accepted_idx] = draft_token ++ num_accepted += 1 ++ last_accepted_idx = int(retrive_index[req_idx, step]) ++ accept_index[req_idx, num_accepted] = last_accepted_idx ++ else: ++ rejected_token = draft_token ++ break ++ ++ accept_token_num[req_idx] = num_accepted ++ final_probs = target_probs[req_idx, num_accepted].clone().float() ++ if rejected_token >= 0: ++ final_probs[rejected_token] = 0.0 ++ final_probs.clamp_min_(0.0) ++ target = float(uniform_samples_for_final_sampling[req_idx]) * float( ++ final_probs.sum() ++ ) ++ sampled_token = int((final_probs.cumsum(0) <= target).sum()) ++ sampled_token = min(sampled_token, final_probs.numel() - 1) ++ predicts[last_accepted_idx] = sampled_token ++ ++ return predicts, accept_index, accept_token_num ++ ++ ++@pytest.mark.parametrize( ++ "threshold_single,threshold_acc", [(1.0, 1.0), (0.0, 0.0), (0.5, 0.8)] ++) ++def test_general_tree_matches_gpu_algorithm_reference( ++ threshold_single, threshold_acc ++): ++ candidates = torch.tensor( ++ [[0, 1, 2, 3, 4, 5], [7, 8, 9, 10, 11, 12]], dtype=torch.int64 ++ ) ++ retrive_index = torch.tensor( ++ [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]], dtype=torch.int64 ++ ) ++ retrive_next_token = torch.tensor( ++ [[1, 2, -1, 4, 5, -1], [4, 2, 3, -1, 5, -1]], ++ dtype=torch.int64, ++ ) ++ retrive_next_sibling = torch.tensor( ++ [[-1, 3, -1, -1, -1, -1], [-1, -1, -1, -1, 1, -1]], ++ dtype=torch.int64, ++ ) ++ batch_size, num_draft_tokens = candidates.shape ++ vocab_size = 20 ++ target_probs = torch.full( ++ (batch_size, num_draft_tokens, vocab_size), 0.01, dtype=torch.float32 ++ ) ++ target_probs[0, 0, 1] = 0.12 ++ target_probs[0, 0, 3] = 0.72 ++ target_probs[0, 3, 4] = 0.82 ++ target_probs[0, 4, 5] = 0.75 ++ target_probs[1, 0, 11] = 0.68 ++ target_probs[1, 0, 8] = 0.14 ++ target_probs[1, 4, 12] = 0.77 ++ target_probs /= target_probs.sum(dim=-1, keepdim=True) ++ uniforms = torch.tensor( ++ [[0.55, 0.2, 0.8, 0.4, 0.3, 0.9], [0.6, 0.2, 0.8, 0.7, 0.3, 0.4]], ++ dtype=torch.float32, ++ ) ++ final_uniforms = torch.tensor([0.25, 0.75], dtype=torch.float32) ++ ++ ref_predicts = torch.full((12,), -1, dtype=torch.int32) ++ ref_accept_index = torch.full((2, 4), -1, dtype=torch.int32) ++ ref_accept_num = torch.zeros(2, dtype=torch.int32) ++ ref_rejected = torch.empty_like(target_probs) ++ target_only_tree_reference( ++ ref_predicts, ++ ref_accept_index, ++ ref_accept_num, ++ candidates, ++ retrive_index, ++ retrive_next_token, ++ retrive_next_sibling, ++ uniforms, ++ final_uniforms, ++ target_probs, ++ ref_rejected, ++ threshold_single, ++ threshold_acc, ++ ) ++ ++ npu_predicts = torch.full_like(ref_predicts, -1, device="npu") ++ npu_accept_index = torch.full_like(ref_accept_index, -1, device="npu") ++ npu_accept_num = torch.zeros_like(ref_accept_num, device="npu") ++ npu_rejected = torch.empty_like(target_probs, device="npu") ++ tree_speculative_sampling_target_only( ++ npu_predicts, ++ npu_accept_index, ++ npu_accept_num, ++ candidates.npu(), ++ retrive_index.npu(), ++ retrive_next_token.npu(), ++ retrive_next_sibling.npu(), ++ uniforms.npu(), ++ final_uniforms.npu(), ++ target_probs.npu(), ++ npu_rejected, ++ threshold_single, ++ threshold_acc, ++ True, ++ ) ++ ++ torch.testing.assert_close(npu_predicts.cpu(), ref_predicts, rtol=0, atol=0) ++ torch.testing.assert_close( ++ npu_accept_index.cpu(), ref_accept_index, rtol=0, atol=0 ++ ) ++ torch.testing.assert_close(npu_accept_num.cpu(), ref_accept_num, rtol=0, atol=0) ++ torch.testing.assert_close(npu_rejected.cpu(), ref_rejected, rtol=0, atol=0) ++ ++ ++def target_only_chain_torch( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++): ++ batch_size, num_draft_tokens = candidates.shape ++ device = candidates.device ++ draft_tokens = candidates[:, 1:].long() ++ step_probs = torch.gather( ++ target_probs[:, :-1, :], 2, draft_tokens.unsqueeze(-1) ++ ).squeeze(-1) ++ accept_steps = uniform_samples[:, : num_draft_tokens - 1] <= step_probs ++ reject_count = (~accept_steps).to(torch.int32).cumsum(dim=1) ++ num_correct = (reject_count == 0).to(torch.int32).sum(dim=1) ++ ++ accept_token_num.copy_(num_correct) ++ accept_index.fill_(-1) ++ positions = torch.arange(num_draft_tokens, device=device).view(1, -1) ++ valid_accept = positions <= num_correct.view(-1, 1) ++ accept_index.copy_( ++ torch.where( ++ valid_accept, ++ retrive_index.to(torch.int32), ++ torch.full_like(accept_index, -1), ++ ) ++ ) ++ ++ predicts.zero_() ++ parent_positions = torch.arange(num_draft_tokens - 1, device=device).view(1, -1) ++ valid_parent = parent_positions < num_correct.view(-1, 1) ++ parent_indices = retrive_index[:, :-1].reshape(-1).long() ++ parent_values = candidates[:, 1:].to(torch.int32).reshape(-1) ++ predicts[parent_indices] = torch.where( ++ valid_parent.reshape(-1), parent_values, predicts[parent_indices] ++ ) ++ ++ rows = torch.arange(batch_size, device=device) ++ final_rows = num_correct.long() ++ final_probs = target_probs[rows, final_rows].clone() ++ rejected = num_correct < num_draft_tokens - 1 ++ rejected_positions = (num_correct.long() + 1).clamp_max(num_draft_tokens - 1) ++ rejected_tokens = candidates[rows, rejected_positions].long() ++ final_probs[rejected, rejected_tokens[rejected]] = 0.0 ++ ++ probability_sums = final_probs.sum(dim=-1, keepdim=True) ++ targets = uniform_samples_for_final_sampling.view(-1, 1) * probability_sums ++ final_tokens = ( ++ (final_probs.cumsum(dim=-1) <= targets) ++ .to(torch.int32) ++ .sum(dim=-1) ++ .clamp_max(target_probs.shape[-1] - 1) ++ ) ++ final_indices = retrive_index[rows, final_rows].long() ++ predicts[final_indices] = final_tokens.to(torch.int32) ++ ++ ++def make_chain_indices(batch_size, num_draft_tokens, device): ++ retrive_index = torch.arange( ++ batch_size * num_draft_tokens, dtype=torch.int64, device=device ++ ).view(batch_size, num_draft_tokens) ++ retrive_next_token = torch.arange( ++ 1, num_draft_tokens + 1, dtype=torch.int64, device=device ++ ).repeat(batch_size, 1) ++ retrive_next_token[:, -1] = -1 ++ retrive_next_sibling = torch.full_like(retrive_next_token, -1) ++ return retrive_index, retrive_next_token, retrive_next_sibling ++ ++ ++def make_stable_chain_final_uniforms( ++ candidates, ++ uniform_samples, ++ target_probs, ++): ++ """Choose final-sampling coins away from inverse-CDF boundaries.""" ++ batch_size, num_draft_tokens = candidates.shape ++ final_uniforms = torch.empty(batch_size, dtype=torch.float32) ++ ++ for req_idx in range(batch_size): ++ num_accepted = 0 ++ rejected_token = -1 ++ for step in range(1, num_draft_tokens): ++ draft_token = int(candidates[req_idx, step]) ++ target_prob = float(target_probs[req_idx, step - 1, draft_token]) ++ if float(uniform_samples[req_idx, step - 1]) <= target_prob: ++ num_accepted += 1 ++ else: ++ rejected_token = draft_token ++ break ++ ++ final_probs = target_probs[req_idx, num_accepted].double().clone() ++ if rejected_token >= 0: ++ final_probs[rejected_token] = 0.0 ++ ++ sampled_token = int(final_probs.argmax()) ++ probability_sum = final_probs.sum() ++ cdf_before = final_probs[:sampled_token].sum() ++ cdf_midpoint = cdf_before + final_probs[sampled_token] * 0.5 ++ final_uniforms[req_idx] = (cdf_midpoint / probability_sum).float() ++ ++ return final_uniforms ++ ++ ++@pytest.mark.parametrize("batch_size", [1, 4, 17]) ++@pytest.mark.parametrize("num_draft_tokens", [2, 5]) ++@pytest.mark.parametrize("vocab_size", [20, 32000, 151552]) ++def test_target_only_chain_matches_reference( ++ batch_size, num_draft_tokens, vocab_size ++): ++ torch.manual_seed(20260717 + batch_size + num_draft_tokens + vocab_size) ++ candidates = torch.randint( ++ 0, vocab_size, (batch_size, num_draft_tokens), dtype=torch.int64 ++ ) ++ logits = torch.randn(batch_size, num_draft_tokens, vocab_size) ++ target_probs = torch.softmax(logits, dim=-1).float() ++ ++ # Give some draft tokens meaningful acceptance probability. ++ for req_idx in range(batch_size): ++ for step in range(1, num_draft_tokens): ++ token = int(candidates[req_idx, step]) ++ target_probs[req_idx, step - 1] *= 0.35 ++ target_probs[req_idx, step - 1, token] += 0.65 ++ target_probs[req_idx, step - 1] /= target_probs[ ++ req_idx, step - 1 ++ ].sum() ++ ++ uniform_samples = torch.rand(batch_size, num_draft_tokens) ++ # A random coin can land within FP32 reduction error of a CDF boundary for ++ # large vocabularies. Use the midpoint of a high-mass token's interval so ++ # exact token equality tests the algorithm instead of reduction order. ++ final_uniform_samples = make_stable_chain_final_uniforms( ++ candidates, ++ uniform_samples, ++ target_probs, ++ ) ++ retrive_index, retrive_next_token, retrive_next_sibling = make_chain_indices( ++ batch_size, num_draft_tokens, "cpu" ++ ) ++ ++ ref_predicts = torch.full( ++ (batch_size * num_draft_tokens,), -1, dtype=torch.int32 ++ ) ++ ref_accept_index = torch.full( ++ (batch_size, num_draft_tokens), -1, dtype=torch.int32 ++ ) ++ ref_accept_num = torch.zeros(batch_size, dtype=torch.int32) ++ target_only_chain_reference( ++ ref_predicts, ++ ref_accept_index, ++ ref_accept_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ final_uniform_samples, ++ target_probs, ++ 1.0, ++ 1.0, ++ ) ++ ++ npu_predicts = torch.full_like(ref_predicts, -1, device="npu") ++ npu_accept_index = torch.full_like(ref_accept_index, -1, device="npu") ++ npu_accept_num = torch.zeros_like(ref_accept_num, device="npu") ++ candidates_npu = candidates.npu() ++ retrive_index_npu = retrive_index.npu() ++ next_token_npu = retrive_next_token.npu() ++ next_sibling_npu = retrive_next_sibling.npu() ++ target_probs_npu = target_probs.npu() ++ ++ tree_speculative_sampling_target_only( ++ predicts=npu_predicts, ++ accept_index=npu_accept_index, ++ accept_token_num=npu_accept_num, ++ candidates=candidates_npu, ++ retrive_index=retrive_index_npu, ++ retrive_next_token=next_token_npu, ++ retrive_next_sibling=next_sibling_npu, ++ uniform_samples=uniform_samples.npu(), ++ uniform_samples_for_final_sampling=final_uniform_samples.npu(), ++ target_probs=target_probs_npu, ++ draft_probs=torch.empty_like(target_probs_npu), ++ threshold_single=1.0, ++ threshold_acc=1.0, ++ deterministic=True, ++ ) ++ ++ torch.testing.assert_close(npu_predicts.cpu(), ref_predicts, rtol=0, atol=0) ++ torch.testing.assert_close( ++ npu_accept_index.cpu(), ref_accept_index, rtol=0, atol=0 ++ ) ++ torch.testing.assert_close(npu_accept_num.cpu(), ref_accept_num, rtol=0, atol=0) ++ ++ ++@pytest.mark.parametrize( ++ "threshold_single,threshold_acc", ++ [(1.0, 1.0), (0.0, 0.0), (0.5, 0.8)], ++) ++def test_target_only_thresholds(threshold_single, threshold_acc): ++ batch_size, num_draft_tokens, vocab_size = 2, 4, 32 ++ candidates = torch.tensor([[0, 3, 4, 5], [0, 7, 8, 9]], dtype=torch.int64) ++ target_probs = torch.full( ++ (batch_size, num_draft_tokens, vocab_size), 1.0 / vocab_size ++ ) ++ for req_idx in range(batch_size): ++ for step in range(1, num_draft_tokens): ++ token = int(candidates[req_idx, step]) ++ target_probs[req_idx, step - 1] *= 0.2 ++ target_probs[req_idx, step - 1, token] += 0.8 ++ target_probs[req_idx, step - 1] /= target_probs[ ++ req_idx, step - 1 ++ ].sum() ++ ++ uniforms = torch.tensor([[0.1, 0.9, 0.2, 0.0], [0.7, 0.2, 0.95, 0.0]]) ++ final_uniforms = torch.tensor([0.25, 0.75]) ++ retrive_index, next_token, next_sibling = make_chain_indices( ++ batch_size, num_draft_tokens, "cpu" ++ ) ++ ++ ref_predicts = torch.full((batch_size * num_draft_tokens,), -1, dtype=torch.int32) ++ ref_accept_index = torch.full( ++ (batch_size, num_draft_tokens), -1, dtype=torch.int32 ++ ) ++ ref_accept_num = torch.zeros(batch_size, dtype=torch.int32) ++ target_only_chain_reference( ++ ref_predicts, ++ ref_accept_index, ++ ref_accept_num, ++ candidates, ++ retrive_index, ++ uniforms, ++ final_uniforms, ++ target_probs, ++ threshold_single, ++ threshold_acc, ++ ) ++ ++ npu_predicts = torch.full_like(ref_predicts, -1, device="npu") ++ npu_accept_index = torch.full_like(ref_accept_index, -1, device="npu") ++ npu_accept_num = torch.zeros_like(ref_accept_num, device="npu") ++ target_probs_npu = target_probs.npu() ++ tree_speculative_sampling_target_only( ++ npu_predicts, ++ npu_accept_index, ++ npu_accept_num, ++ candidates.npu(), ++ retrive_index.npu(), ++ next_token.npu(), ++ next_sibling.npu(), ++ uniforms.npu(), ++ final_uniforms.npu(), ++ target_probs_npu, ++ torch.empty_like(target_probs_npu), ++ threshold_single, ++ threshold_acc, ++ True, ++ ) ++ ++ torch.testing.assert_close(npu_predicts.cpu(), ref_predicts, rtol=0, atol=0) ++ torch.testing.assert_close( ++ npu_accept_index.cpu(), ref_accept_index, rtol=0, atol=0 ++ ) ++ torch.testing.assert_close(npu_accept_num.cpu(), ref_accept_num, rtol=0, atol=0) ++ ++ ++def run_benchmark(batch_size, num_draft_tokens, vocab_size, warmup, iterations): ++ candidates = torch.randint( ++ 0, ++ vocab_size, ++ (batch_size, num_draft_tokens), ++ dtype=torch.int64, ++ device="npu", ++ ) ++ target_probs = torch.softmax( ++ torch.randn( ++ batch_size, ++ num_draft_tokens, ++ vocab_size, ++ dtype=torch.float32, ++ device="npu", ++ ), ++ dim=-1, ++ ) ++ retrive_index, next_token, next_sibling = make_chain_indices( ++ batch_size, num_draft_tokens, "npu" ++ ) ++ uniforms = torch.rand( ++ batch_size, num_draft_tokens, dtype=torch.float32, device="npu" ++ ) ++ final_uniforms = torch.rand(batch_size, dtype=torch.float32, device="npu") ++ draft_probs = torch.empty_like(target_probs) ++ predicts = torch.zeros( ++ batch_size * num_draft_tokens, dtype=torch.int32, device="npu" ++ ) ++ accept_index = torch.full( ++ (batch_size, num_draft_tokens), -1, dtype=torch.int32, device="npu" ++ ) ++ accept_num = torch.zeros(batch_size, dtype=torch.int32, device="npu") ++ ++ def run_kernel(): ++ tree_speculative_sampling_target_only( ++ predicts, ++ accept_index, ++ accept_num, ++ candidates, ++ retrive_index, ++ next_token, ++ next_sibling, ++ uniforms, ++ final_uniforms, ++ target_probs, ++ draft_probs, ++ 1.0, ++ 1.0, ++ True, ++ ) ++ ++ def run_torch(): ++ target_only_chain_torch( ++ predicts, ++ accept_index, ++ accept_num, ++ candidates, ++ retrive_index, ++ uniforms, ++ final_uniforms, ++ target_probs, ++ ) ++ ++ def benchmark(fn): ++ for _ in range(warmup): ++ fn() ++ torch.npu.synchronize() ++ started = time.perf_counter() ++ for _ in range(iterations): ++ fn() ++ torch.npu.synchronize() ++ return (time.perf_counter() - started) * 1000 / iterations ++ ++ kernel_latency_ms = benchmark(run_kernel) ++ torch_latency_ms = benchmark(run_torch) ++ print( ++ f"batch={batch_size} drafts={num_draft_tokens} vocab={vocab_size} " ++ f"kernel_ms={kernel_latency_ms:.4f} torch_ms={torch_latency_ms:.4f} " ++ f"speedup={torch_latency_ms / kernel_latency_ms:.2f}x" ++ ) ++ ++ ++if __name__ == "__main__": ++ parser = argparse.ArgumentParser() ++ parser.add_argument("--perf", action="store_true") ++ parser.add_argument("--batch-size", type=int, default=16) ++ parser.add_argument("--num-draft-tokens", type=int, default=5) ++ parser.add_argument("--vocab-size", type=int, default=151552) ++ parser.add_argument("--warmup", type=int, default=10) ++ parser.add_argument("--iterations", type=int, default=100) ++ args = parser.parse_args() ++ if args.perf: ++ run_benchmark( ++ args.batch_size, ++ args.num_draft_tokens, ++ args.vocab_size, ++ args.warmup, ++ args.iterations, ++ ) ++ else: ++ raise SystemExit(pytest.main([__file__])) +-- +2.34.1 + + +From b85bb7a082f3fcf9f8edbbe12b1641f94292245d Mon Sep 17 00:00:00 2001 +From: wuqiwei +Date: Wed, 5 Aug 2026 13:52:42 +0000 +Subject: [PATCH 3/3] cuda_gragh arg patch + +--- + .../python/torch_memory_saver/entrypoint.py | 8 +++++--- + 1 file changed, 5 insertions(+), 3 deletions(-) + +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 146c011..7b0508d 100644 +--- a/contrib/torch_memory_saver/python/torch_memory_saver/entrypoint.py ++++ b/contrib/torch_memory_saver/python/torch_memory_saver/entrypoint.py +@@ -36,6 +36,7 @@ class TorchMemorySaver: + capture_error_mode="global", + tag: str = _TAG_DEFAULT, + enable_cpu_backup: bool = False, ++ **kwargs, + ): + """Similar to `torch.cuda.graph`, but ensures memory in it to be pauseable.""" + self._ensure_initialized() +@@ -46,6 +47,7 @@ class TorchMemorySaver: + capture_error_mode=capture_error_mode, + tag=tag, + enable_cpu_backup=enable_cpu_backup, ++ **kwargs + ): + yield + +@@ -115,12 +117,13 @@ class _TorchMemorySaverImpl: + capture_error_mode, + tag: str, + enable_cpu_backup: bool, ++ **kwargs, + ): + assert ( + self._hook_mode == "preload" + ), "Only hook_mode=preload supports pauseable CUDA Graph currently" + with torch.npu.graph( +- cuda_graph, pool=pool, stream=stream, capture_error_mode=capture_error_mode ++ cuda_graph, pool=pool, stream=stream, capture_error_mode=capture_error_mode, **kwargs + ): + with self._with_region_config(tag=tag, enable_cpu_backup=enable_cpu_backup): + yield +@@ -176,5 +179,4 @@ def _sanity_checks(): + if "expandable_segments:True" in os.environ.get("PYTORCH_CUDA_ALLOC_CONF", ""): + raise RuntimeError( + "TorchMemorySaver is disabled for the current process because expandable_segments is not supported yet." +- ) +- +\ No newline at end of file ++ ) +\ No newline at end of file +-- +2.34.1 + diff --git a/docker/npu_patch/sglang-npu.patch b/docker/npu_patch/sglang-npu.patch index de38657b6..fffe907a7 100644 --- a/docker/npu_patch/sglang-npu.patch +++ b/docker/npu_patch/sglang-npu.patch @@ -1,56 +1,421 @@ +From 1bf8b943fa46b025618857cf0efed7ac7b5fb9b5 Mon Sep 17 00:00:00 2001 +From: wuqiwei +Date: Thu, 23 Jul 2026 08:44:55 +0000 +Subject: [PATCH 1/5] sglang_npu_patch + +--- + python/sglang/srt/layers/quantization/unquant.py | 8 ++++---- + python/sglang/srt/utils/common.py | 4 +++- + 2 files changed, 7 insertions(+), 5 deletions(-) + diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py -index 1ade4ed9e4..0161bd398a 100644 +index 82a3d77f05..25d4c21406 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py -@@ -314,12 +314,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): +@@ -402,10 +402,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): + layer.w2_weight.data = layer.w2_weight.data.reshape( layer.num_local_experts, *new_shape_w2 ) - - if _is_npu: - for weight_name in ["w13_weight", "w2_weight"]: - weight = getattr(layer, weight_name) -- weight.data = weight.data.transpose(1, 2) -- weight.data = npu_format_cast(weight.data) -- +- weight.data = npu_format_cast(weight) ++ # if _is_npu: ++ # for weight_name in ["w13_weight", "w2_weight"]: ++ # weight = getattr(layer, weight_name) ++ # weight.data = npu_format_cast(weight) + return - def maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load( -@@ -646,7 +640,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): - # gmm1: gate_up_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], -- weight=[layer.w13_weight], -+ weight=[layer.w13_weight.transpose(1, 2)], - bias=w13_bias, - split_item=2, - group_list_type=1, -@@ -670,7 +664,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): - # gmm2: down_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], -- weight=[layer.w2_weight], -+ weight=[layer.w2_weight.transpose(1, 2)], - bias=w2_bias, - split_item=2, - group_list_type=1, diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py -index f68721f3d0..3e85ee7c10 100644 +index 82db90da3b..aabd222d7e 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py -@@ -46,6 +46,7 @@ import types - import uuid - import warnings - from collections import OrderedDict, defaultdict +@@ -73,6 +73,7 @@ from typing import ( + TypeVar, + Union, + ) +from relax.utils.device import is_npu_available - from contextlib import contextmanager - from dataclasses import dataclass - from decimal import Decimal -@@ -150,7 +151,7 @@ def is_npu() -> bool: + from unittest import SkipTest + from unittest.case import _ShouldStop + from urllib.parse import unquote, urlparse +@@ -176,7 +177,8 @@ def is_npu() -> bool: if not hasattr(torch, "npu"): return False - if not torch.npu.is_available(): ++ # if not torch.npu.is_available(): + if not is_npu_available: raise RuntimeError( "torch_npu detected, but NPU device is not available or visible." ) +-- +2.34.1 + + +From 6ba9e861e0c32c4b299ad5b7b6e7e11c8a441534 Mon Sep 17 00:00:00 2001 +From: David Wang +Date: Fri, 19 Jun 2026 00:19:40 +0000 +Subject: [PATCH 2/5] fix mamba radix partial page prefix matching + +--- + .../sglang/srt/mem_cache/mamba_radix_cache.py | 9 ++++ + .../unit/mem_cache/test_mamba_unittest.py | 52 +++++++++++++++++++ + 2 files changed, 61 insertions(+) + +diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py +index 368781d8d5..0856caae48 100644 +--- a/python/sglang/srt/mem_cache/mamba_radix_cache.py ++++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py +@@ -1089,6 +1089,10 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): + if self.disable or len(key) == 0: + return None + ++ key = key.page_aligned(self.page_size) ++ if len(key) == 0: ++ return None ++ + return key + + def _match_post_processor( +@@ -1158,6 +1162,9 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): + ) + + def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode: ++ assert ( ++ 0 < split_len < len(child.key) ++ ), f"split_len must create non-empty nodes, {split_len=}, {len(child.key)=}" + # new_node -> child + new_node = TreeNode() + new_node.children = {key[split_len:].child_key(self.page_size): child} +@@ -1166,6 +1173,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): + new_node.full_lock_ref = child.full_lock_ref + new_node.mamba_lock_ref = 0 + new_node.key = child.key[:split_len] ++ assert len(new_node.key) > 0, f"new_node.key should not be empty" + new_node.value = child.value[:split_len].clone() + + # child time should be later than parent's time for mamba tombstone +@@ -1176,6 +1184,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): + self.mamba_lru_list.remove_node(child) + child.parent = new_node + child.key = child.key[split_len:] ++ assert len(child.key) > 0, f"child.key should not be empty" + child.value = child.value[split_len:].clone() + new_node.parent.children[key.child_key(self.page_size)] = new_node + new_node.hash_value, child.hash_value = split_node_hash_value( +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) +-- +2.34.1 + + +From 9f5d5ac59b474998956f2ac2e9e494ee2c6529a6 Mon Sep 17 00:00:00 2001 +From: wuqiwei +Date: Wed, 5 Aug 2026 09:26:57 +0000 +Subject: [PATCH 3/5] [NPU] Enable non-greedy MTP sampling- #32495 #32495 + +--- + python/sglang/srt/speculative/eagle_utils.py | 98 +++++++++++++------- + 1 file changed, 67 insertions(+), 31 deletions(-) + +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 ++++ b/python/sglang/srt/speculative/eagle_utils.py +@@ -617,7 +617,7 @@ def eagle_sample( + + # Sample tokens + target_predict = None +- if sampling_info.is_all_greedy or _is_npu or _is_hip or _is_xpu: ++ if sampling_info.is_all_greedy or _is_hip or _is_xpu: + target_predict = torch.argmax(next_token_logits, dim=-1) + target_predict = target_predict.reshape(bs, verify_input.draft_token_num) + predict, accept_index, num_correct_drafts = verify_tree_greedy_func( +@@ -632,43 +632,80 @@ def eagle_sample( + topk=verify_input.tree_topk, + ) + else: +- from sgl_kernel import ( +- top_k_renorm_prob, +- top_p_renorm_prob, +- tree_speculative_sampling_target_only, +- ) +- +- from sglang.srt.speculative.reject_sampling import ( +- chain_speculative_sampling_triton, +- ) +- + use_rejection_sampling = ( + get_global_server_args().speculative_use_rejection_sampling + ) + ++ if _is_npu: ++ from sgl_kernel_npu.sample import ( ++ chain_speculative_sampling_rejection, ++ top_k_top_p_renorm_probs, ++ tree_speculative_sampling_target_only, ++ ) ++ ++ sampling_fn = ( ++ chain_speculative_sampling_rejection ++ if use_rejection_sampling ++ else tree_speculative_sampling_target_only ++ ) ++ else: ++ from sgl_kernel import ( ++ top_k_renorm_prob, ++ top_p_renorm_prob, ++ tree_speculative_sampling_target_only, ++ ) ++ ++ from sglang.srt.speculative.reject_sampling import ( ++ chain_speculative_sampling_triton, ++ ) ++ ++ sampling_fn = ( ++ chain_speculative_sampling_triton ++ if use_rejection_sampling ++ else tree_speculative_sampling_target_only ++ ) ++ + # Apply temperature and get target probs + expanded_temperature = torch.repeat_interleave( + sampling_info.temperatures, verify_input.draft_token_num, dim=0 + ) # (bs * num_draft_tokens, 1) + ++ sampling_logits = next_token_logits.float() if _is_npu else next_token_logits + target_probs = F.softmax( +- next_token_logits / expanded_temperature, dim=-1 ++ sampling_logits / expanded_temperature, dim=-1 + ) # (bs * num_draft_tokens, vocab_size) + maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax") +- target_probs = top_k_renorm_prob( +- target_probs, +- torch.repeat_interleave( +- sampling_info.top_ks, verify_input.draft_token_num, dim=0 +- ), +- ) # (bs * num_draft_tokens, vocab_size) +- maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm") +- target_probs = top_p_renorm_prob( +- target_probs, +- torch.repeat_interleave( +- sampling_info.top_ps, verify_input.draft_token_num, dim=0 +- ), +- ) +- maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm") ++ ++ if _is_npu: ++ target_probs = top_k_top_p_renorm_probs( ++ target_probs, ++ torch.repeat_interleave( ++ sampling_info.top_ks, verify_input.draft_token_num, dim=0 ++ ), ++ torch.repeat_interleave( ++ sampling_info.top_ps, verify_input.draft_token_num, dim=0 ++ ), ++ sampling_info.need_top_k_sampling, ++ sampling_info.need_top_p_sampling, ++ ) ++ maybe_detect_nan(target_probs, "v2 verify: target_probs after renorm") ++ else: ++ if sampling_info.need_top_k_sampling: ++ target_probs = top_k_renorm_prob( ++ target_probs, ++ torch.repeat_interleave( ++ sampling_info.top_ks, verify_input.draft_token_num, dim=0 ++ ), ++ ) # (bs * num_draft_tokens, vocab_size) ++ maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm") ++ if sampling_info.need_top_p_sampling: ++ target_probs = top_p_renorm_prob( ++ target_probs, ++ torch.repeat_interleave( ++ sampling_info.top_ps, verify_input.draft_token_num, dim=0 ++ ), ++ ) ++ maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm") + target_probs = target_probs.reshape(bs, verify_input.draft_token_num, -1) + draft_probs = ( + verify_input.draft_probs +@@ -687,16 +724,15 @@ def eagle_sample( + "does not produce one (draft_probs missing or vocab-mismatched)." + ) + ++ if _is_npu: ++ target_probs = target_probs.contiguous() ++ draft_probs = draft_probs.float().contiguous() ++ + # coins for rejection sampling + coins = torch.rand_like(candidates, dtype=torch.float32, device=device) + # coins for final sampling + coins_for_final_sampling = torch.rand((bs,), dtype=torch.float32, device=device) + +- sampling_fn = ( +- chain_speculative_sampling_triton +- if use_rejection_sampling +- else tree_speculative_sampling_target_only +- ) + sampling_fn( + predicts=predict, # mutable + accept_index=accept_index, # mutable +-- +2.34.1 + + +From 32ffa124368440360bc966cdaca9f44dac5aaf72 Mon Sep 17 00:00:00 2001 +From: wuqiwei +Date: Thu, 6 Aug 2026 08:10:11 +0000 +Subject: [PATCH 4/5] for tms cuda_gragh + +--- + python/sglang/srt/utils/torch_memory_saver_adapter.py | 9 ++++++--- + 1 file changed, 6 insertions(+), 3 deletions(-) + +diff --git a/python/sglang/srt/utils/torch_memory_saver_adapter.py b/python/sglang/srt/utils/torch_memory_saver_adapter.py +index ad98e59283..a5b2cdfd32 100644 +--- a/python/sglang/srt/utils/torch_memory_saver_adapter.py ++++ b/python/sglang/srt/utils/torch_memory_saver_adapter.py +@@ -41,7 +41,7 @@ class TorchMemorySaverAdapter(ABC): + def region(self, tag: str, enable_cpu_backup: bool = False): + raise NotImplementedError + +- def cuda_graph(self, **kwargs): ++ def cuda_graph(self, cuda_graph=None, **kwargs): + raise NotImplementedError + + def disable(self): +@@ -67,7 +67,10 @@ class _TorchMemorySaverAdapterReal(TorchMemorySaverAdapter): + def region(self, tag: str, enable_cpu_backup: bool = False): + return _memory_saver.region(tag=tag, enable_cpu_backup=enable_cpu_backup) + +- def cuda_graph(self, **kwargs): ++ def cuda_graph(self, cuda_graph=None, **kwargs): ++ if cuda_graph is not None: ++ kwargs["cuda_graph"] = cuda_graph ++ # kwargs.pop("auto_dispatch_capture", None) # torch_memory_saver 0.0.8 does not support this arg + return _memory_saver.cuda_graph(**kwargs) + + def disable(self): +@@ -94,7 +97,7 @@ class _TorchMemorySaverAdapterNoop(TorchMemorySaverAdapter): + yield + + @contextmanager +- def cuda_graph(self, **kwargs): ++ def cuda_graph(self, cuda_graph=None, **kwargs): + yield + + @contextmanager +-- +2.34.1 + + +From 8c023fdbdf6d0ccf062f88a26a52cb2545b0683a Mon Sep 17 00:00:00 2001 +From: wuqiwei +Date: Thu, 6 Aug 2026 08:10:48 +0000 +Subject: [PATCH 5/5] disable tms cuda_gragh + +--- + .../npu/graph_runner/npu_cudagraph_backend.py | 18 +++++++++++++++++- + 1 file changed, 17 insertions(+), 1 deletion(-) + +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 ++++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py +@@ -37,6 +37,21 @@ if TYPE_CHECKING: + 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 ++ ++ 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 +-- +2.34.1 + From 57809b43772529b339dc5ba2a96f25594a711a71 Mon Sep 17 00:00:00 2001 From: lixionglong Date: Fri, 21 Aug 2026 06:26:43 +0000 Subject: [PATCH 02/16] fix(megatron): attach _hf_config for MTP bridge when pp>1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Fix bridge mapping registry failure with PP > 1 - Set model_bridge._hf_config from bridge.hf_pretrained.config when missing --- relax/backends/megatron/weight_update/bridge_converter.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/relax/backends/megatron/weight_update/bridge_converter.py b/relax/backends/megatron/weight_update/bridge_converter.py index 608eb0ce8..6eb2c9a0a 100644 --- a/relax/backends/megatron/weight_update/bridge_converter.py +++ b/relax/backends/megatron/weight_update/bridge_converter.py @@ -82,7 +82,10 @@ def init_tasks(self) -> None: if task.param_weight is not None: self._bridge_task_map[task.global_param_name] = task - self._bridge_mapping_registry = bridge._model_bridge.mapping_registry() + model_bridge = bridge._model_bridge + if not hasattr(model_bridge, "_hf_config"): + model_bridge._hf_config = bridge.hf_pretrained.config + self._bridge_mapping_registry = model_bridge.mapping_registry() mapping_registry = self._bridge_mapping_registry for name, _param in named_params_and_buffers(self._args, self._model): global_name = strip_param_name_prefix(name) From 3f51dccd11dc33747715136780fe353d7b0b63f4 Mon Sep 17 00:00:00 2001 From: wuqiwei Date: Fri, 21 Aug 2026 06:28:27 +0000 Subject: [PATCH 03/16] chore(docker): refresh NPU patch set MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🔩 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 --- docker/npu_patch/megatron-bridge.patch | 56 +- docker/npu_patch/megatron.patch | 38 +- docker/npu_patch/mindspeed-bridge.patch | 142 +- docker/npu_patch/mindspeed-ops.patch | 13 + docker/npu_patch/mindspeed.patch | 154 ++ docker/npu_patch/sgl-kernel-npu.patch | 1847 +++++++++++++++++++++++ docker/npu_patch/sglang-npu.patch | 1070 ++++++++++++- 7 files changed, 3231 insertions(+), 89 deletions(-) create mode 100644 docker/npu_patch/mindspeed-ops.patch create mode 100644 docker/npu_patch/sgl-kernel-npu.patch diff --git a/docker/npu_patch/megatron-bridge.patch b/docker/npu_patch/megatron-bridge.patch index 0066df63b..d49be66b3 100644 --- a/docker/npu_patch/megatron-bridge.patch +++ b/docker/npu_patch/megatron-bridge.patch @@ -1,11 +1,32 @@ diff --git a/src/megatron/bridge/models/conversion/utils.py b/src/megatron/bridge/models/conversion/utils.py -index 5a66e719..3d411c17 100644 +index 86ddf8661..115624b16 100644 --- a/src/megatron/bridge/models/conversion/utils.py +++ b/src/megatron/bridge/models/conversion/utils.py -@@ -203,6 +203,15 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0): +@@ -203,6 +203,17 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0): + ): # bound methods + return None + ++ # Convert OmegaConf containers to plain dict/list to avoid in-place ++ # mutation triggering "dictionary changed size during iteration" errors ++ # inside OmegaConf's internal _flags_cache handling. ++ try: ++ from omegaconf import DictConfig, ListConfig, OmegaConf as _OmegaConf ++ ++ if isinstance(obj, (DictConfig, ListConfig)): ++ obj = _OmegaConf.to_container(obj, resolve=True) ++ except ImportError: ++ pass ++ + # Handle dataclass/object with attributes + if hasattr(obj, "__dict__"): + # Create a copy to avoid modifying the original +@@ -213,9 +224,18 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0): # Recursively clean attribute cleaned_value = remove_non_pickleables(attr_value, max_depth, current_depth + 1) +- +- # Set the cleaned value (or None if it was removed) +- setattr(cleaned_obj, attr_name, cleaned_value) + if hasattr(cleaned_obj, '__setattr__'): + try: + setattr(cleaned_obj, attr_name, cleaned_value) @@ -15,6 +36,33 @@ index 5a66e719..3d411c17 100644 + print(f"Skipping attribute '{attr_name}' due to Union type") + continue + raise ++ else: ++ # Fallback for objects without __setattr__ override ++ setattr(cleaned_obj, attr_name, cleaned_value) + + return cleaned_obj + +diff --git a/src/megatron/bridge/peft/utils.py b/src/megatron/bridge/peft/utils.py +index 1ca5b18bd..4797e2e42 100644 +--- a/src/megatron/bridge/peft/utils.py ++++ b/src/megatron/bridge/peft/utils.py +@@ -33,6 +33,7 @@ from megatron.core.transformer.moe.router import TopKRouter + + from megatron.bridge.utils.import_utils import safe_import_from + ++from relax.utils.device import is_npu_available + + TEColumnParallelLinear, HAVE_TE_COL_LINEAR = safe_import_from( + "megatron.core.extensions.transformer_engine", "TEColumnParallelLinear" +@@ -62,7 +63,10 @@ HAVE_TE = all( + ) + ) + +-MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm") ++if is_npu_available: ++ MixedFusedLayerNorm, HAVE_APEX = None, False ++else: ++ MixedFusedLayerNorm, HAVE_APEX = safe_import_from("apex.normalization.fused_layer_norm", "MixedFusedLayerNorm") - # Set the cleaned value (or None if it was removed) - setattr(cleaned_obj, attr_name, cleaned_value) + TECL = (TEColumnParallelLinear, TELayerNormColumnParallelLinear, TEColumnParallelGroupedLinear) + TERL = (TERowParallelLinear, TERowParallelGroupedLinear) diff --git a/docker/npu_patch/megatron.patch b/docker/npu_patch/megatron.patch index 0b76d6360..21de37e10 100644 --- a/docker/npu_patch/megatron.patch +++ b/docker/npu_patch/megatron.patch @@ -655,7 +655,7 @@ index 4be974017..0dfdb0928 100644 """ Maintain the expert bias in float32. diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py -index 2edb652bf..58fb89c4e 100755 +index 2edb652bf..7a068271b 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -586,6 +586,102 @@ class MTPLossAutoScaler(torch.autograd.Function): @@ -857,6 +857,42 @@ index 2edb652bf..58fb89c4e 100755 ) if self.config.recompute_method == 'uniform': +@@ -1053,7 +1192,7 @@ class MultiTokenPredictionBlock(MegatronModule): + self._build_layers(pg_collection) + assert len(self.layers) > 0, "MultiTokenPredictionBlock must have at least one layer." + self.cp_group = pg_collection.cp +- ++ self._register_mtp_grad_hooks() + def _build_layers(self, pg_collection): + def build_layer(layer_spec, layer_number): + fp8_init_context = get_fp8_context(self.config, is_init=True) +@@ -1074,6 +1213,26 @@ class MultiTokenPredictionBlock(MegatronModule): + ] + ) + ++ def _register_mtp_grad_hooks(self): ++ if len(self.layers) == 0: ++ return ++ ++ def _make_sync_hook(): ++ def _hook(grad): ++ torch.npu.synchronize() ++ return _hook ++ ++ sync_keys = [ ++ "self_attention.linear_qkv.weight", ++ "eh_proj.weight", ++ ] ++ ++ for name, param in self.layers[0].named_parameters(): ++ for key in sync_keys: ++ if key in name: ++ param.register_hook(_make_sync_hook()) ++ break ++ + def forward( + self, + input_ids: Tensor, diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index eaae58590..0f01f6bc0 100644 --- a/megatron/core/transformer/transformer_config.py diff --git a/docker/npu_patch/mindspeed-bridge.patch b/docker/npu_patch/mindspeed-bridge.patch index 327ff7afa..f658a5981 100644 --- a/docker/npu_patch/mindspeed-bridge.patch +++ b/docker/npu_patch/mindspeed-bridge.patch @@ -19,64 +19,106 @@ index 99ccba5..8fa07a1 100644 # Set the cleaned value (or None if it was removed) setattr(cleaned_obj, attr_name, cleaned_value) diff --git a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py -index 818af29..ef97078 100644 +index 19a3b9d..3dd5e25 100644 --- a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py +++ b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py -@@ -191,7 +191,7 @@ class GatedDeltaNet(MegatronModule): +@@ -1,6 +1,6 @@ + # pylint: disable=R0801 + from typing import List, Optional, Tuple +- ++import os + import torch + from torch import nn + import torch.nn.functional as F +@@ -37,12 +37,11 @@ from megatron.core.ssm.gated_delta_net import ( + ) + + try: +- from causal_conv1d import causal_conv1d ++ import fla_npu ++ from fla.modules.convolution import causal_conv1d + except ImportError: + causal_conv1d = None +- causal_conv1d_update = None + +-from mindspeed_ops.api.triton.l2norm import l2norm + from mindspeed_bridge.models.qwen_vl.modelling_qwen3_vl.chunk_gated_delta_rule import ( + torch_chunk_gated_delta_rule, + ) +@@ -59,6 +58,12 @@ except ImportError: + + from mindspeed_bridge.models.qwen_vl.modelling_qwen3_vl.flash_gated_delta_rule import flash_gated_delta_rule + ++def naive_l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): ++ """This function is intended to align with the l2norm implementation in the FLA library.""" ++ original_dtype = x.dtype ++ inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) ++ # Counteract verl's autocast promotion (bf16 -> fp32) by restoring original dtype ++ return (x * inv_norm).to(original_dtype) + + class GatedDeltaNet(MegatronModule): + """Gated Delta Net (GDN) layer class +@@ -191,9 +196,10 @@ class GatedDeltaNet(MegatronModule): ) setattr(self.A_log, "tensor_model_parallel", True) - if HAVE_FLA and self.use_triton_gdn: -+ if HAVE_FLA: ++ gdn_backend = os.environ.get("MINDSPEED_BRIDGE_GDN_BACKEND", "triton").strip().lower() ++ if gdn_backend == "triton": self.gated_delta_rule = chunk_gated_delta_rule - elif self.use_ascend_gdn: +- elif self.use_ascend_gdn: ++ elif gdn_backend=="ascendc": self.gated_delta_rule = flash_gated_delta_rule -diff --git a/mindspeed_bridge/models/qwen_vl/qwen35_vl_provider.py b/mindspeed_bridge/models/qwen_vl/qwen35_vl_provider.py -index d495c47..f08ea49 100644 ---- a/mindspeed_bridge/models/qwen_vl/qwen35_vl_provider.py -+++ b/mindspeed_bridge/models/qwen_vl/qwen35_vl_provider.py -@@ -402,22 +402,6 @@ class Qwen35VLModelProvider(GPTModelProvider): - self.vision_config = Qwen3_5VisionConfig() - super().__post_init__() + else: + self.gated_delta_rule = torch_chunk_gated_delta_rule +@@ -474,7 +480,7 @@ class GatedDeltaNet(MegatronModule): + beta=beta, + initial_state=None, + output_final_state=False, +- use_qk_l2norm_in_kernel=True, ++ use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens_q, + ) + nvtx_range_pop(suffix="gated_delta_rule") +@@ -555,7 +561,7 @@ class GatedDeltaNet(MegatronModule): -- def finalize(self) -> None: -- self.validate_parallelism() -- super().finalize() -- -- def validate_parallelism(self): -- """Validate that parallelism settings are compatible with this model's architecture. -- -- Call this after mutating parallelism attributes (e.g. tensor_model_parallel_size) -- on an already-constructed provider, since finalize() only runs once before provide(). -- """ -- if self.num_query_groups < self.tensor_model_parallel_size: -- raise ValueError( -- f"TP size {self.tensor_model_parallel_size} should be less than or equal to " -- f"num_query_groups {self.num_query_groups}. Please use a smaller TP size." -- ) -- - def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen3VLModel: - """Provide a Qwen3.5 VL dense model instance with vision and language components.""" - language_transformer_config = self -@@ -595,21 +579,6 @@ class Qwen35VLMoEModelProvider(GPTModelProvider): - self.vision_config = Qwen3_5MoeVisionConfig() - super().__post_init__() + # Apply L2 norm to query and key + if self.use_qk_l2norm: +- query_key = l2norm(query_key.contiguous()) ++ query_key = naive_l2norm(query_key.contiguous()) -- def finalize(self) -> None: -- self.validate_parallelism() -- super().finalize() -- -- def validate_parallelism(self): -- """Validate that parallelism settings are compatible with this model's architecture. -- -- Call this after mutating parallelism attributes (e.g. tensor_model_parallel_size) -- on an already-constructed provider, since finalize() only runs once before provide(). -- """ -- if self.num_query_groups < self.tensor_model_parallel_size: -- raise ValueError( -- f"TP size {self.tensor_model_parallel_size} should be less than or equal to " -- f"num_query_groups {self.num_query_groups}. Please use a smaller TP size." -- ) + # Split query and key + split_size = self.qk_dim_local_tp // self.key_head_dim // self.cp_size +diff --git a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py +index 529eaa3..1df512c 100644 +--- a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py ++++ b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py +@@ -108,6 +108,7 @@ class Qwen3VLGPTModel(GPTModel): + # args for deepstack + visual_pos_masks: Optional[torch.Tensor] = None, + deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, ++ mtp_kwargs: Optional[dict] = None, + ) -> Tensor: + """Forward function of the GPT Model This function passes the input tensors + through the embedding layer, and then the decoeder and finally into the post +@@ -195,6 +196,7 @@ class Qwen3VLGPTModel(GPTModel): + runtime_gather_output=runtime_gather_output, + extra_block_kwargs=extra_block_kwargs, + inference_context=inference_context, ++ mtp_kwargs=mtp_kwargs, + ) + + if _shadow_embedding: +diff --git a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py +index 5d03f85..2488ef7 100644 +--- a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py ++++ b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py +@@ -56,6 +56,8 @@ def get_vision_model_config(hf_config, megatron_config=None): + ffn_hidden_size=hf_config.intermediate_size, + add_bias_linear=True, + add_qkv_bias=True, ++ moe_permute_fusion=False, # ← 新增：vision model 没有 MoE，不需要 permute fusion ++ use_fused_moe_token_permute_and_unpermute=False, # ← 新增：同上 + ) - def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen3VLModel: - """Provide a Qwen3.5 VL model instance with vision and language components. + # apply text model config to vision model config diff --git a/docker/npu_patch/mindspeed-ops.patch b/docker/npu_patch/mindspeed-ops.patch new file mode 100644 index 000000000..9dae51dd4 --- /dev/null +++ b/docker/npu_patch/mindspeed-ops.patch @@ -0,0 +1,13 @@ +diff --git a/mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_bwd_dhu.py b/mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_bwd_dhu.py +index bf62b0b..1925eca 100644 +--- a/mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_bwd_dhu.py ++++ b/mindspeed_ops/arch32/triton/gdn/chunk_gated_delta_rule_bwd_dhu.py +@@ -21,7 +21,7 @@ from mindspeed_ops.api.triton.utils import prepare_chunk_indices, prepare_chunk_ + ) + @triton.autotune( + configs=get_autotune_config(multibuffer_list=(True, False)), +- key=['H', 'K', 'V', 'BT', 'BV', 'USE_G', 'IS_VARLEN'], ++ key=['H', 'K', 'V', 'BT', 'BV'], + ) + @triton.jit(do_not_specialize=['T']) + def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( diff --git a/docker/npu_patch/mindspeed.patch b/docker/npu_patch/mindspeed.patch index 9f5f056ae..1597d6bbf 100644 --- a/docker/npu_patch/mindspeed.patch +++ b/docker/npu_patch/mindspeed.patch @@ -37,6 +37,22 @@ index c624fd0a..cef668af 100644 return parser +diff --git a/mindspeed/core/megatron_basic/arguments_basic.py b/mindspeed/core/megatron_basic/arguments_basic.py +old mode 100644 +new mode 100755 +index 87d71f81..aa6c219b +--- a/mindspeed/core/megatron_basic/arguments_basic.py ++++ b/mindspeed/core/megatron_basic/arguments_basic.py +@@ -147,7 +147,7 @@ def transformer_config_init_subclass(cls, **kwargs): + if callable(value) and not isinstance(value, type): + value = field(default_factory=value) + elif type(value) in mutable_types: +- value = field(default_factory=lambda: value) ++ value = field(default_factory=lambda v=value: v) + else: + value = value + setattr(cls, key, value) +\ No newline at end of file diff --git a/mindspeed/features_manager/functional/profile.py b/mindspeed/features_manager/functional/profile.py index 6450b41c..86584015 100644 --- a/mindspeed/features_manager/functional/profile.py @@ -93,3 +109,141 @@ index f14b231d..881c7fdb 100644 full_args = get_full_args() for k, v in args.items(): setattr(full_args, k, v) +diff --git a/mindspeed/ops/triton/l2norm.py b/mindspeed/ops/triton/l2norm.py +old mode 100644 +new mode 100755 +index 0050bce3..a9791a2d +--- a/mindspeed/ops/triton/l2norm.py ++++ b/mindspeed/ops/triton/l2norm.py +@@ -5,23 +5,13 @@ + from typing import Optional + + import torch +-import torch.nn as nn ++from torch import nn + import triton + import triton.language as tl + +-from mindspeed.ops.triton.utils import input_guard, is_amd ++from mindspeed.ops.triton.utils import input_guard + +-BT_LIST = [8, 16, 32, 64, 128] +-NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if is_amd else [1, 2, 4, 8, 16, 32] + +- +-@triton.autotune( +- configs=[ +- triton.Config({}, num_warps=num_warps) +- for num_warps in NUM_WARPS_AUTOTUNE +- ], +- key=['D'] +-) + @triton.jit + def l2norm_fwd_kernel1( + x, +@@ -45,13 +35,13 @@ def l2norm_fwd_kernel1( + tl.store(rstd + i_t, b_rstd) + + +-@triton.autotune( +- configs=[ +- triton.Config({}, num_warps=num_warps) +- for num_warps in NUM_WARPS_AUTOTUNE +- ], +- key=['D'] +-) ++#@triton.autotune( ++ # configs=[ ++ # triton.Config({}, num_warps=num_warps) ++ # for num_warps in NUM_WARPS_AUTOTUNE ++ # ], ++ # key=['D'] ++#) + @triton.jit + def l2norm_bwd_kernel1( + y, +@@ -76,14 +66,14 @@ def l2norm_bwd_kernel1( + tl.store(dx + cols, b_dx, mask=mask) + + +-@triton.autotune( +- configs=[ +- triton.Config({'BT': BT}, num_warps=num_warps) +- for num_warps in [1, 2, 4, 8, 16] +- for BT in BT_LIST +- ], +- key=['D', 'NB'] +-) ++#@triton.autotune( ++# configs=[ ++# triton.Config({'BT': BT}, num_warps=num_warps) ++# for num_warps in [1, 2, 4, 8, 16] ++# for BT in BT_LIST ++# ], ++# key=['D', 'NB'] ++#) + @triton.jit + def l2norm_fwd_kernel( + x, +@@ -113,14 +103,14 @@ def l2norm_fwd_kernel( + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) + + +-@triton.autotune( +- configs=[ +- triton.Config({'BT': BT}, num_warps=num_warps) +- for num_warps in [1, 2, 4, 8, 16] +- for BT in BT_LIST +- ], +- key=['D', 'NB'] +-) ++#@triton.autotune( ++# configs=[ ++# triton.Config({'BT': BT}, num_warps=num_warps) ++# for num_warps in [1, 2, 4, 8, 16] ++# for BT in BT_LIST ++# ], ++# key=['D', 'NB'] ++#) + @triton.jit + def l2norm_bwd_kernel( + y, +@@ -188,11 +178,12 @@ def l2norm_fwd( + rstd = torch.empty((T,), dtype=torch.float32, device=x.device) + if D <= 512: + NB = triton.cdiv(T, 2048) ++ BT = 32 + bt_size = 32 + + def grid(meta): +- new_bt = meta['BT'] * bt_size +- return (triton.cdiv(T, new_bt), ) ++ new_bt = BT * bt_size ++ return (triton.cdiv(T, new_bt),) + + l2norm_fwd_kernel[grid]( + x=x, +@@ -203,6 +194,7 @@ def l2norm_fwd( + D=D, + BD=BD, + NB=NB, ++ BT=BT, + bt_size=bt_size, + ) + else: +@@ -238,6 +230,7 @@ def l2norm_bwd( + + if D <= 512: + NB = triton.cdiv(T, 2048) ++ BT=32 + bt_size = 40 + l2norm_bwd_kernel[(bt_size,)]( + y=y, +@@ -249,6 +242,7 @@ def l2norm_bwd( + D=D, + BD=BD, + NB=NB, ++ BT=BT, + bt_size=bt_size, + ) + else: diff --git a/docker/npu_patch/sgl-kernel-npu.patch b/docker/npu_patch/sgl-kernel-npu.patch new file mode 100644 index 000000000..3ef4d0a96 --- /dev/null +++ b/docker/npu_patch/sgl-kernel-npu.patch @@ -0,0 +1,1847 @@ +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 ++++ b/contrib/torch_memory_saver/python/torch_memory_saver/entrypoint.py +@@ -36,6 +36,7 @@ class TorchMemorySaver: + capture_error_mode="global", + tag: str = _TAG_DEFAULT, + enable_cpu_backup: bool = False, ++ **kwargs, + ): + """Similar to `torch.cuda.graph`, but ensures memory in it to be pauseable.""" + self._ensure_initialized() +@@ -46,6 +47,7 @@ class TorchMemorySaver: + capture_error_mode=capture_error_mode, + tag=tag, + enable_cpu_backup=enable_cpu_backup, ++ **kwargs + ): + yield + +@@ -115,12 +117,13 @@ class _TorchMemorySaverImpl: + capture_error_mode, + tag: str, + enable_cpu_backup: bool, ++ **kwargs, + ): + assert ( + self._hook_mode == "preload" + ), "Only hook_mode=preload supports pauseable CUDA Graph currently" + with torch.npu.graph( +- cuda_graph, pool=pool, stream=stream, capture_error_mode=capture_error_mode ++ cuda_graph, pool=pool, stream=stream, capture_error_mode=capture_error_mode, **kwargs + ): + with self._with_region_config(tag=tag, enable_cpu_backup=enable_cpu_backup): + yield +@@ -159,6 +162,7 @@ class _TorchMemorySaverImpl: + # only be released after the memory region is resumed and empty_cache() is invoked. + torch_npu._C._npu_releasePool(torch.npu.current_device(), pool.id) + del pool ++ torch.npu.empty_cache() + finally: + self._binary_wrapper.cdll.tms_set_interesting_region(True) + +@@ -175,4 +179,4 @@ def _sanity_checks(): + if "expandable_segments:True" in os.environ.get("PYTORCH_CUDA_ALLOC_CONF", ""): + raise RuntimeError( + "TorchMemorySaver is disabled for the current process because expandable_segments is not supported yet." +- ) ++ ) +\ No newline at end of file +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 ++++ b/csrc/cache_location_assign/op_host/cache_loc_assign.cpp +@@ -19,8 +19,8 @@ + namespace sglang { + namespace npu_kernel { + +-at::Tensor getTiling(const at::Tensor &reqPoolIndices, uint64_t rowSize, uint64_t poolSize, uint32_t &blockDim, +- bool isUpddate) ++at::Tensor getTiling(const at::Tensor &reqPoolIndices, uint64_t rowSize, uint64_t poolSize, uint64_t maxStep, ++ uint32_t &blockDim, bool isUpddate) + { + auto batchSize = reqPoolIndices.sizes()[0]; + auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(); +@@ -39,6 +39,7 @@ at::Tensor getTiling(const at::Tensor &reqPoolIndices, uint64_t rowSize, uint64_ + tillingData->rowNumNoTail = batchSize / (tillingData->vcoreNum); + tillingData->tailNum = batchSize % (tillingData->vcoreNum); + tillingData->rowSize = rowSize; ++ tillingData->maxStep = maxStep; + + if (reqPoolIndices.options().dtype() == at::kInt) { + tillingData->key = 1; +@@ -50,13 +51,13 @@ at::Tensor getTiling(const at::Tensor &reqPoolIndices, uint64_t rowSize, uint64_ + tillingData->reqInxBufferSize = tillingData->reqInxBufferCount * sizeof(int64_t); + } + +- tillingData->tokenCountAlignInt32 = host_utils::alinInt32Count(MAX_STEP); ++ tillingData->tokenCountAlignInt32 = host_utils::alinInt32Count(maxStep); + tillingData->tokenColAlignInt32 = tillingData->tokenCountAlignInt32 * sizeof(int32_t); + + tillingData->offsetCountAlignInt64 = host_utils::alinInt64Count(batchSize); + tillingData->offsetColAlignInt64 = tillingData->offsetCountAlignInt64 * sizeof(int64_t); + +- tillingData->cacheLocSize = batchSize * MAX_STEP; ++ tillingData->cacheLocSize = batchSize * maxStep; + tillingData->cacheLocCountAlignInt32 = host_utils::alinInt32Count(tillingData->cacheLocSize); + tillingData->cacheLocAlignInt32 = tillingData->cacheLocCountAlignInt32 * sizeof(int32_t); + +@@ -73,7 +74,7 @@ at::Tensor getTiling(const at::Tensor &reqPoolIndices, uint64_t rowSize, uint64_ + } + + HOST_API void checkParams(const at::Tensor &reqPoolIndices, const at::Tensor &tokenPool, const at::Tensor &startOffset, +- const at::Tensor &endOffset, const at::Tensor &outCacheLoc) ++ const at::Tensor &endOffset, const at::Tensor &outCacheLoc, int64_t maxStep) + { + auto reqIdxType = reqPoolIndices.options().dtype(); + if ((reqIdxType != at::kInt && reqIdxType != at::kLong) || tokenPool.options().dtype() != at::kInt || +@@ -83,16 +84,28 @@ HOST_API void checkParams(const at::Tensor &reqPoolIndices, const at::Tensor &to + "Only support inputTensor combo1: int64, int32, int64, int64, int32; combo2: " + "int32, int32, int64, int64, int32"); + } ++ if (maxStep < 1 || maxStep > MAX_STEP) { ++ throw std::invalid_argument("max_step must be in [1, " + std::to_string(MAX_STEP) + "], got " + ++ std::to_string(maxStep) + ++ " (the cache_loc_assign kernel handles at most MAX_STEP tokens per row)"); ++ } ++ auto batchSize = reqPoolIndices.sizes()[0]; ++ if (outCacheLoc.numel() < batchSize * static_cast(maxStep)) { ++ throw std::invalid_argument("out_cache_loc too small: needs at least batchSize * max_step = " + ++ std::to_string(batchSize * maxStep) + " elements, got " + ++ std::to_string(outCacheLoc.numel())); ++ } + } + + HOST_API at::Tensor cache_loc_assign(const at::Tensor &reqPoolIndices, const at::Tensor &tokenPool, + const at::Tensor &startOffset, const at::Tensor &endOffset, +- const at::Tensor &outCacheLoc) ++ const at::Tensor &outCacheLoc, int64_t maxStep) + { +- checkParams(reqPoolIndices, tokenPool, startOffset, endOffset, outCacheLoc); ++ checkParams(reqPoolIndices, tokenPool, startOffset, endOffset, outCacheLoc, maxStep); + uint32_t blockDim; + uint32_t cacheAssignMode = 0; +- at::Tensor tilingTensor = getTiling(reqPoolIndices, tokenPool.sizes()[1], tokenPool.sizes()[0], blockDim, false); ++ at::Tensor tilingTensor = getTiling(reqPoolIndices, tokenPool.sizes()[1], tokenPool.sizes()[0], maxStep, blockDim, ++ false); + + EXEC_KERNEL_CMD(cache_loc_assign, blockDim, reqPoolIndices, tokenPool, startOffset, endOffset, outCacheLoc, + tilingTensor, cacheAssignMode); +@@ -101,12 +114,13 @@ HOST_API at::Tensor cache_loc_assign(const at::Tensor &reqPoolIndices, const at: + + HOST_API at::Tensor cache_loc_update(const at::Tensor &reqPoolIndices, const at::Tensor &tokenPool, + const at::Tensor &startOffset, const at::Tensor &endOffset, +- const at::Tensor &outCacheLoc) ++ const at::Tensor &outCacheLoc, int64_t maxStep) + { +- checkParams(reqPoolIndices, tokenPool, startOffset, endOffset, outCacheLoc); ++ checkParams(reqPoolIndices, tokenPool, startOffset, endOffset, outCacheLoc, maxStep); + uint32_t blockDim; + uint32_t cacheAssignMode = 1; +- at::Tensor tilingTensor = getTiling(reqPoolIndices, tokenPool.sizes()[1], tokenPool.sizes()[0], blockDim, true); ++ at::Tensor tilingTensor = getTiling(reqPoolIndices, tokenPool.sizes()[1], tokenPool.sizes()[0], maxStep, blockDim, ++ true); + + EXEC_KERNEL_CMD(cache_loc_assign, blockDim, reqPoolIndices, tokenPool, startOffset, endOffset, outCacheLoc, + tilingTensor, cacheAssignMode); +diff --git a/csrc/cache_location_assign/op_host/tiling/cache_loc_assign.h b/csrc/cache_location_assign/op_host/tiling/cache_loc_assign.h +index 2f8e9bd..2d473ee 100644 +--- a/csrc/cache_location_assign/op_host/tiling/cache_loc_assign.h ++++ b/csrc/cache_location_assign/op_host/tiling/cache_loc_assign.h +@@ -26,6 +26,7 @@ struct AssignCacheTillingData { + uint64_t reqInxBufferCount{0}; + uint64_t reqInxBufferSize{0}; + ++ uint64_t maxStep{0}; + uint64_t tokenCountAlignInt32{0}; + uint64_t tokenColAlignInt32{0}; + +@@ -37,6 +38,10 @@ struct AssignCacheTillingData { + uint64_t cacheLocAlignInt32{0}; + }; + ++// Upper bound of the per-row token count this kernel handles in one call. ++// The host validates `1 <= max_step <= MAX_STEP` and derives every buffer ++// size (token transfer count, out_cache_loc tiling) from the caller-supplied ++// max_step — never from this constant. + constexpr uint32_t MAX_STEP = 16; + + #endif // CACHE_LOC_ASSIGN_TILING_H +diff --git a/csrc/cache_location_assign/op_kernel/cache_loc_assign_kernel.cpp b/csrc/cache_location_assign/op_kernel/cache_loc_assign_kernel.cpp +index 109f604..c78f303 100644 +--- a/csrc/cache_location_assign/op_kernel/cache_loc_assign_kernel.cpp ++++ b/csrc/cache_location_assign/op_kernel/cache_loc_assign_kernel.cpp +@@ -21,6 +21,12 @@ constexpr int32_t BUFFER_NUM = 2; + constexpr uint32_t ASSIGN_TO_POOL = 0; + constexpr uint32_t RETRIEVE_FROM_POOL = 1; + ++// Per-row contract: the caller passes max_step (tokens moved per row) through ++// the tiling data. The host validates `1 <= max_step <= MAX_STEP` and sizes ++// out_cache_loc as `batchSize * max_step`; every transfer below (token pool ++// copy in CopyIn, token pool write-back, cacheLoc GM read/write) is derived ++// from max_step, so no fixed-size over/under-transfer can occur. ++ + template + class CacheLocAssignKernel + { +@@ -44,6 +50,7 @@ public: + } + this->rowSize = tempTilingGM->rowSize; + this->reqInxBufferCount = tempTilingGM->reqInxBufferCount; ++ this->maxStep = tempTilingGM->maxStep; + this->tokenCountAlignInt32 = tempTilingGM->tokenCountAlignInt32; + this->offsetCountAlignInt64 = tempTilingGM->offsetCountAlignInt64; + this->cacheLocCountAlignInt32 = tempTilingGM->cacheLocCountAlignInt32; +@@ -151,7 +158,7 @@ private: + tokenPoolLocal.SetValue(j, cache); + } + +- uint32_t tokenBytes = static_cast(MAX_STEP * sizeof(int32_t)); ++ uint32_t tokenBytes = static_cast(this->maxStep * sizeof(int32_t)); + AscendC::DataCopyExtParams copyParams{1, tokenBytes, 0, 0, 0}; + AscendC::DataCopyPad(tokenPoolGM[reqIdx * this->rowSize + start], tokenPoolLocal, copyParams); + +@@ -208,6 +215,7 @@ private: + uint64_t tailOffset; + uint64_t rowOffset; + uint64_t rowSize; ++ uint64_t maxStep; + uint64_t cacheLocSize; + + int64_t cacheIdxStart{0}; +diff --git a/csrc/pytorch_extensions.cpp b/csrc/pytorch_extensions.cpp +index 1cfccb0..8e4ad91 100644 +--- a/csrc/pytorch_extensions.cpp ++++ b/csrc/pytorch_extensions.cpp +@@ -32,11 +32,11 @@ TORCH_LIBRARY_FRAGMENT(npu, m) + + m.def( + "cache_loc_assign(Tensor req_indices, Tensor token_pool, Tensor start_offset, Tensor end_offset, Tensor " +- "out_cache_loc) -> Tensor"); ++ "out_cache_loc, int max_step) -> Tensor"); + + m.def( + "cache_loc_update(Tensor req_indices, Tensor token_pool, Tensor start_offset, Tensor end_offset, Tensor " +- "out_cache_loc) -> Tensor"); ++ "out_cache_loc, int max_step) -> Tensor"); + + m.def( + "assign_cache_op(Tensor! out, Tensor src, Tensor dst_start_idx, Tensor dst_end_idx, Tensor src_start_idx, " +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 ++++ b/include/sgl_kenel_npu_ops.h +@@ -19,13 +19,15 @@ at::Tensor cache_loc_assign(const at::Tensor &req_indices, + const at::Tensor &token_pool, + const at::Tensor &start_offset, + const at::Tensor &end_offset, +- const at::Tensor &out_cache_loc); ++ const at::Tensor &out_cache_loc, ++ int64_t max_step); + + at::Tensor cache_loc_update(const at::Tensor &req_indices, + const at::Tensor &token_pool, + const at::Tensor &start_offset, + const at::Tensor &end_offset, +- const at::Tensor &out_cache_loc); ++ const at::Tensor &out_cache_loc, ++ int64_t max_step); + + 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/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 ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/sample/__init__.py +@@ -0,0 +1,13 @@ ++from sgl_kernel_npu.sample.chain_speculative_sampling import ( ++ chain_speculative_sampling_rejection, ++) ++from sgl_kernel_npu.sample.probability import top_k_top_p_renorm_probs ++from sgl_kernel_npu.sample.tree_speculative_sampling_target_only import ( ++ tree_speculative_sampling_target_only, ++) ++ ++__all__ = [ ++ "chain_speculative_sampling_rejection", ++ "top_k_top_p_renorm_probs", ++ "tree_speculative_sampling_target_only", ++] +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/sample/chain_speculative_sampling.py b/python/sgl_kernel_npu/sgl_kernel_npu/sample/chain_speculative_sampling.py +new file mode 100644 +index 0000000..e8340da +--- /dev/null ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/sample/chain_speculative_sampling.py +@@ -0,0 +1,354 @@ ++import torch ++import triton ++import triton.language as tl ++ ++ ++@triton.jit ++def _chain_rejection_accept_kernel( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ target_probs, ++ draft_probs, ++ metadata, ++ num_draft_tokens: tl.constexpr, ++ num_speculative_tokens: tl.constexpr, ++ num_draft_prob_rows: tl.constexpr, ++ vocab_size: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ row_offset = req_idx * num_draft_tokens ++ ++ cur_prob_row = tl.full((), 0, tl.int64) ++ last_accepted_idx = tl.load(retrive_index + row_offset).to(tl.int64) ++ num_accepted = 0 ++ active = tl.full((), 1, tl.int32) ++ ++ tl.store(accept_index + req_idx * num_speculative_tokens, last_accepted_idx) ++ ++ # Linear Leviathan/Chen verification. Candidate 0 is the root; candidate ++ # step uses probability row step - 1 until a rejection terminates the chain. ++ for step in range(1, num_draft_tokens): ++ if active == 1: ++ draft_token = tl.load(candidates + row_offset + step).to(tl.int64) ++ target_offset = ( ++ (row_offset + cur_prob_row) * vocab_size + draft_token ++ ) ++ draft_offset = ( ++ (req_idx * num_draft_prob_rows + cur_prob_row) * vocab_size ++ + draft_token ++ ) ++ target_prob = tl.load(target_probs + target_offset).to(tl.float32) ++ draft_prob = tl.load(draft_probs + draft_offset).to(tl.float32) ++ coin = tl.load(uniform_samples + row_offset + step - 1).to( ++ tl.float32 ++ ) ++ ++ if coin * draft_prob < target_prob: ++ tl.store(predicts + last_accepted_idx, draft_token) ++ num_accepted += 1 ++ draft_idx = tl.load(retrive_index + row_offset + step).to( ++ tl.int64 ++ ) ++ tl.store( ++ accept_index ++ + req_idx * num_speculative_tokens ++ + num_accepted, ++ draft_idx, ++ ) ++ last_accepted_idx = draft_idx ++ # Keep this loop-carried value int64 across both branches. ++ # Triton infers the constexpr loop variable `step` as int32. ++ cur_prob_row = tl.full((), step, tl.int64) ++ else: ++ active = 0 ++ ++ tl.store(accept_token_num + req_idx, num_accepted) ++ ++ # metadata = [target row, output slot, all drafts accepted]. ++ metadata_offset = req_idx * 3 ++ tl.store(metadata + metadata_offset, cur_prob_row) ++ tl.store(metadata + metadata_offset + 1, last_accepted_idx) ++ tl.store(metadata + metadata_offset + 2, active) ++ ++ ++@triton.jit ++def _chain_rejection_block_sum_kernel( ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ num_draft_tokens: tl.constexpr, ++ num_draft_prob_rows: tl.constexpr, ++ vocab_size: tl.constexpr, ++ vocab_block_size: tl.constexpr, ++ num_vocab_blocks: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ block_idx = tl.program_id(1) ++ vocab_offsets = block_idx * vocab_block_size + tl.arange(0, vocab_block_size) ++ vocab_mask = vocab_offsets < vocab_size ++ ++ metadata_offset = req_idx * 3 ++ target_row = tl.load(metadata + metadata_offset).to(tl.int64) ++ all_accepted = tl.load(metadata + metadata_offset + 2).to(tl.int32) ++ target_offset = (req_idx * num_draft_tokens + target_row) * vocab_size ++ target = tl.load( ++ target_probs + target_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ ++ if all_accepted == 1: ++ residual = target ++ else: ++ draft_offset = ( ++ req_idx * num_draft_prob_rows + target_row ++ ) * vocab_size ++ draft = tl.load( ++ draft_probs + draft_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ residual = tl.maximum(target - draft, 0.0) ++ ++ tl.store( ++ block_sums + req_idx * num_vocab_blocks + block_idx, ++ tl.sum(residual, axis=0), ++ ) ++ ++ ++@triton.jit ++def _chain_rejection_sample_kernel( ++ predicts, ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ uniform_samples_for_final_sampling, ++ num_draft_tokens: tl.constexpr, ++ num_draft_prob_rows: tl.constexpr, ++ vocab_size: tl.constexpr, ++ vocab_block_size: tl.constexpr, ++ num_vocab_blocks: tl.constexpr, ++ pad_num_vocab_blocks: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ metadata_offset = req_idx * 3 ++ target_row = tl.load(metadata + metadata_offset).to(tl.int64) ++ output_idx = tl.load(metadata + metadata_offset + 1).to(tl.int64) ++ all_accepted = tl.load(metadata + metadata_offset + 2).to(tl.int32) ++ ++ block_offsets = tl.arange(0, pad_num_vocab_blocks) ++ block_mask = block_offsets < num_vocab_blocks ++ sums = tl.load( ++ block_sums + req_idx * num_vocab_blocks + block_offsets, ++ mask=block_mask, ++ other=0.0, ++ ).to(tl.float32) ++ block_cdf = tl.cumsum(sums, axis=0) ++ total = tl.sum(sums, axis=0) ++ coin = tl.load(uniform_samples_for_final_sampling + req_idx).to(tl.float32) ++ target_value = coin * total ++ ++ selected_block = tl.sum( ++ ((block_cdf <= target_value) & block_mask).to(tl.int32), axis=0 ++ ) ++ selected_block = tl.minimum(selected_block, num_vocab_blocks - 1) ++ prefix_sum = tl.sum( ++ tl.where(block_offsets < selected_block, sums, 0.0), axis=0 ++ ) ++ local_target = tl.maximum(target_value - prefix_sum, 0.0) ++ ++ local_offsets = tl.arange(0, vocab_block_size) ++ vocab_offsets = selected_block * vocab_block_size + local_offsets ++ vocab_mask = vocab_offsets < vocab_size ++ target_offset = (req_idx * num_draft_tokens + target_row) * vocab_size ++ target = tl.load( ++ target_probs + target_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ if all_accepted == 1: ++ residual = target ++ else: ++ draft_offset = ( ++ req_idx * num_draft_prob_rows + target_row ++ ) * vocab_size ++ draft = tl.load( ++ draft_probs + draft_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ residual = tl.maximum(target - draft, 0.0) ++ ++ local_cdf = tl.cumsum(residual, axis=0) ++ local_index = tl.sum( ++ ((local_cdf <= local_target) & vocab_mask).to(tl.int32), axis=0 ++ ) ++ last_valid_local = tl.max( ++ tl.where((residual > 0.0) & vocab_mask, local_offsets, -1), axis=0 ++ ) ++ valid_local_count = tl.minimum( ++ vocab_block_size, ++ vocab_size - selected_block * vocab_block_size, ++ ) ++ sampled_local = tl.where( ++ local_index < valid_local_count, ++ local_index, ++ last_valid_local, ++ ) ++ sampled_token = tl.where( ++ sampled_local >= 0, ++ selected_block * vocab_block_size + sampled_local, ++ vocab_size - 1, ++ ) ++ tl.store(predicts + output_idx, tl.minimum(sampled_token, vocab_size - 1)) ++ ++ ++def chain_speculative_sampling_rejection( ++ predicts: torch.Tensor, ++ accept_index: torch.Tensor, ++ accept_token_num: torch.Tensor, ++ candidates: torch.Tensor, ++ retrive_index: torch.Tensor, ++ retrive_next_token: torch.Tensor, ++ retrive_next_sibling: torch.Tensor, ++ uniform_samples: torch.Tensor, ++ uniform_samples_for_final_sampling: torch.Tensor, ++ target_probs: torch.Tensor, ++ draft_probs: torch.Tensor, ++ threshold_single: float = 1.0, ++ threshold_acc: float = 1.0, ++ deterministic: bool = True, ++) -> None: ++ """NPU kernel implementation of classic chain rejection sampling.""" ++ del retrive_next_token, retrive_next_sibling ++ del threshold_single, threshold_acc, deterministic ++ ++ if candidates.ndim != 2 or target_probs.ndim != 3: ++ raise ValueError("candidates must be 2-D and target_probs must be 3-D") ++ batch_size, num_draft_tokens = candidates.shape ++ if batch_size == 0: ++ return ++ if num_draft_tokens == 0: ++ raise ValueError("num_draft_tokens must be positive") ++ if target_probs.shape[:2] != (batch_size, num_draft_tokens): ++ raise ValueError( ++ "target_probs shape must be [batch, num_draft_tokens, vocab_size]" ++ ) ++ if retrive_index.shape != candidates.shape: ++ raise ValueError("retrive_index shape must match candidates") ++ if accept_index.shape != candidates.shape: ++ raise ValueError( ++ "classic rejection sampling requires a topk=1 linear chain" ++ ) ++ if accept_token_num.shape != (batch_size,): ++ raise ValueError("accept_token_num must have shape [batch]") ++ if predicts.ndim != 1: ++ raise ValueError("predicts must be 1-D") ++ if uniform_samples.shape != candidates.shape: ++ raise ValueError("uniform_samples shape must match candidates") ++ if uniform_samples_for_final_sampling.shape != (batch_size,): ++ raise ValueError( ++ "uniform_samples_for_final_sampling must have shape [batch]" ++ ) ++ if draft_probs is None or draft_probs.ndim != 3: ++ raise ValueError("draft_probs must be a 3-D tensor") ++ if draft_probs.shape[0] != batch_size: ++ raise ValueError("draft_probs batch size must match candidates") ++ if draft_probs.shape[1] < max(num_draft_tokens - 1, 1): ++ raise ValueError("draft_probs does not contain every proposal row") ++ if draft_probs.shape[-1] != target_probs.shape[-1]: ++ raise ValueError("draft_probs and target_probs vocab sizes must match") ++ if target_probs.dtype != torch.float32 or draft_probs.dtype != torch.float32: ++ raise TypeError("target_probs and draft_probs must be torch.float32") ++ if uniform_samples.dtype != torch.float32: ++ raise TypeError("uniform_samples must be torch.float32") ++ if uniform_samples_for_final_sampling.dtype != torch.float32: ++ raise TypeError("uniform_samples_for_final_sampling must be torch.float32") ++ integer_dtypes = ( ++ (predicts, torch.int32, "predicts"), ++ (accept_index, torch.int32, "accept_index"), ++ (accept_token_num, torch.int32, "accept_token_num"), ++ (candidates, torch.int64, "candidates"), ++ (retrive_index, torch.int64, "retrive_index"), ++ ) ++ for tensor, expected_dtype, name in integer_dtypes: ++ if tensor.dtype != expected_dtype: ++ raise TypeError(f"{name} must be {expected_dtype}") ++ tensors = ( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++ draft_probs, ++ ) ++ if any(tensor.device != target_probs.device for tensor in tensors): ++ raise ValueError("all tensors must be on the same NPU device") ++ if any(not tensor.is_contiguous() for tensor in tensors): ++ raise ValueError("all tensors must be contiguous") ++ ++ num_speculative_tokens = accept_index.shape[1] ++ num_draft_prob_rows = draft_probs.shape[1] ++ vocab_size = target_probs.shape[-1] ++ vocab_block_size = 2048 ++ num_vocab_blocks = triton.cdiv(vocab_size, vocab_block_size) ++ pad_num_vocab_blocks = triton.next_power_of_2(num_vocab_blocks) ++ ++ metadata = torch.empty( ++ (batch_size, 3), dtype=torch.int64, device=target_probs.device ++ ) ++ block_sums = torch.empty( ++ (batch_size, num_vocab_blocks), ++ dtype=torch.float32, ++ device=target_probs.device, ++ ) ++ ++ _chain_rejection_accept_kernel[(batch_size,)]( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ target_probs, ++ draft_probs, ++ metadata, ++ num_draft_tokens=num_draft_tokens, ++ num_speculative_tokens=num_speculative_tokens, ++ num_draft_prob_rows=num_draft_prob_rows, ++ vocab_size=vocab_size, ++ ) ++ _chain_rejection_block_sum_kernel[(batch_size, num_vocab_blocks)]( ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ num_draft_tokens=num_draft_tokens, ++ num_draft_prob_rows=num_draft_prob_rows, ++ vocab_size=vocab_size, ++ vocab_block_size=vocab_block_size, ++ num_vocab_blocks=num_vocab_blocks, ++ ) ++ _chain_rejection_sample_kernel[(batch_size,)]( ++ predicts, ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ uniform_samples_for_final_sampling, ++ num_draft_tokens=num_draft_tokens, ++ num_draft_prob_rows=num_draft_prob_rows, ++ vocab_size=vocab_size, ++ vocab_block_size=vocab_block_size, ++ num_vocab_blocks=num_vocab_blocks, ++ pad_num_vocab_blocks=pad_num_vocab_blocks, ++ ) +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/sample/probability.py b/python/sgl_kernel_npu/sgl_kernel_npu/sample/probability.py +new file mode 100644 +index 0000000..cadddd0 +--- /dev/null ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/sample/probability.py +@@ -0,0 +1,40 @@ ++import torch ++ ++ ++def top_k_top_p_renorm_probs( ++ probs: torch.Tensor, ++ top_ks: torch.Tensor, ++ top_ps: torch.Tensor, ++ need_top_k_sampling: bool, ++ need_top_p_sampling: bool, ++) -> torch.Tensor: ++ """Apply the same sequential top-k then top-p policy used by SGLang GPU.""" ++ if not need_top_k_sampling and not need_top_p_sampling: ++ return probs ++ ++ vocab_size = probs.shape[-1] ++ sorted_probs, sorted_indices = probs.sort(dim=-1, descending=True) ++ ++ if need_top_k_sampling: ++ top_ks = top_ks.to(device=probs.device, dtype=torch.long).clamp( ++ min=1, max=vocab_size ++ ) ++ positions = torch.arange(vocab_size, device=probs.device).view(1, -1) ++ sorted_probs.masked_fill_(positions >= top_ks.view(-1, 1), 0.0) ++ sorted_probs.div_( ++ sorted_probs.sum(dim=-1, keepdim=True).clamp_min_(1e-20) ++ ) ++ ++ if need_top_p_sampling: ++ top_ps = top_ps.to(device=probs.device, dtype=probs.dtype) ++ cumulative_probs = sorted_probs.cumsum(dim=-1) ++ sorted_probs.masked_fill_( ++ cumulative_probs - sorted_probs > top_ps.view(-1, 1), 0.0 ++ ) ++ sorted_probs.div_( ++ sorted_probs.sum(dim=-1, keepdim=True).clamp_min_(1e-20) ++ ) ++ ++ return torch.zeros_like(probs).scatter_( ++ dim=-1, index=sorted_indices, src=sorted_probs ++ ) +diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/sample/tree_speculative_sampling_target_only.py b/python/sgl_kernel_npu/sgl_kernel_npu/sample/tree_speculative_sampling_target_only.py +new file mode 100644 +index 0000000..e34baab +--- /dev/null ++++ b/python/sgl_kernel_npu/sgl_kernel_npu/sample/tree_speculative_sampling_target_only.py +@@ -0,0 +1,383 @@ ++import torch ++import triton ++import triton.language as tl ++ ++ ++@triton.jit ++def _tree_target_only_accept_kernel( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ retrive_next_token, ++ retrive_next_sibling, ++ uniform_samples, ++ target_probs, ++ rejected_probs, ++ metadata, ++ threshold_single, ++ threshold_acc, ++ num_draft_tokens: tl.constexpr, ++ num_speculative_tokens: tl.constexpr, ++ vocab_size: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ row_offset = req_idx * num_draft_tokens ++ ++ cur_prob_row = tl.full((), 0, tl.int64) ++ cur_node = tl.full((), 0, tl.int64) ++ last_accepted_idx = tl.load(retrive_index + row_offset).to(tl.int64) ++ coin = tl.load(uniform_samples + row_offset).to(tl.float32) ++ num_accepted = 0 ++ path_active = tl.full((), 1, tl.int32) ++ ++ tl.store(accept_index + req_idx * num_speculative_tokens, last_accepted_idx) ++ ++ # This is the same breadth-at-each-depth traversal used by the CUDA kernel: ++ # descend to the first child, then walk siblings until one is accepted. ++ for _depth in range(1, num_speculative_tokens): ++ accepted_at_depth = tl.full((), 0, tl.int32) ++ prob_acc = tl.full((), 0.0, tl.float32) ++ ++ if path_active == 1: ++ cur_node = tl.load( ++ retrive_next_token + row_offset + cur_node ++ ).to(tl.int64) ++ if cur_node == -1: ++ path_active = 0 ++ ++ # The loop is bounded by the number of tree nodes. It terminates ++ # logically when a child is accepted or the sibling list reaches -1. ++ for _sibling in range(0, num_draft_tokens): ++ if ( ++ (path_active == 1) ++ & (accepted_at_depth == 0) ++ & (cur_node != -1) ++ ): ++ draft_token = tl.load( ++ candidates + row_offset + cur_node ++ ).to(tl.int64) ++ draft_idx = tl.load( ++ retrive_index + row_offset + cur_node ++ ).to(tl.int64) ++ prob_offset = ( ++ (row_offset + cur_prob_row) * vocab_size + draft_token ++ ) ++ target_prob_single = tl.load( ++ target_probs + prob_offset ++ ).to(tl.float32) ++ prob_acc += target_prob_single ++ ++ accepted = (coin <= prob_acc / threshold_acc) | ( ++ target_prob_single >= threshold_single ++ ) ++ if accepted: ++ tl.store(predicts + last_accepted_idx, draft_token) ++ num_accepted += 1 ++ tl.store( ++ accept_index ++ + req_idx * num_speculative_tokens ++ + num_accepted, ++ draft_idx, ++ ) ++ last_accepted_idx = draft_idx ++ cur_prob_row = cur_node ++ coin = tl.load( ++ uniform_samples + row_offset + cur_node ++ ).to(tl.float32) ++ accepted_at_depth = 1 ++ else: ++ # The CUDA target-only kernel stores the rejected sibling's ++ # target probability in draft_probs and later samples from ++ # relu(target_probs - draft_probs). ++ tl.store(rejected_probs + prob_offset, target_prob_single) ++ cur_node = tl.load( ++ retrive_next_sibling + row_offset + cur_node ++ ).to(tl.int64) ++ ++ if accepted_at_depth == 0: ++ path_active = 0 ++ ++ tl.store(accept_token_num + req_idx, num_accepted) ++ ++ # metadata = [final target-probability row, final output slot]. ++ metadata_offset = req_idx * 2 ++ tl.store(metadata + metadata_offset, cur_prob_row) ++ tl.store(metadata + metadata_offset + 1, last_accepted_idx) ++ ++ ++@triton.jit ++def _tree_target_only_block_sum_kernel( ++ target_probs, ++ rejected_probs, ++ metadata, ++ block_sums, ++ num_draft_tokens: tl.constexpr, ++ vocab_size: tl.constexpr, ++ vocab_block_size: tl.constexpr, ++ num_vocab_blocks: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ block_idx = tl.program_id(1) ++ vocab_offsets = block_idx * vocab_block_size + tl.arange(0, vocab_block_size) ++ vocab_mask = vocab_offsets < vocab_size ++ ++ target_row = tl.load(metadata + req_idx * 2).to(tl.int64) ++ probs_offset = (req_idx * num_draft_tokens + target_row) * vocab_size ++ target = tl.load( ++ target_probs + probs_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ rejected = tl.load( ++ rejected_probs + probs_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ residual = tl.maximum(target - rejected, 0.0) ++ block_sum = tl.sum(residual, axis=0) ++ tl.store(block_sums + req_idx * num_vocab_blocks + block_idx, block_sum) ++ ++ ++@triton.jit ++def _tree_target_only_sample_kernel( ++ predicts, ++ target_probs, ++ rejected_probs, ++ metadata, ++ block_sums, ++ uniform_samples_for_final_sampling, ++ num_draft_tokens: tl.constexpr, ++ vocab_size: tl.constexpr, ++ vocab_block_size: tl.constexpr, ++ num_vocab_blocks: tl.constexpr, ++ pad_num_vocab_blocks: tl.constexpr, ++): ++ req_idx = tl.program_id(0) ++ metadata_offset = req_idx * 2 ++ target_row = tl.load(metadata + metadata_offset).to(tl.int64) ++ output_idx = tl.load(metadata + metadata_offset + 1).to(tl.int64) ++ ++ block_offsets = tl.arange(0, pad_num_vocab_blocks) ++ block_mask = block_offsets < num_vocab_blocks ++ sums = tl.load( ++ block_sums + req_idx * num_vocab_blocks + block_offsets, ++ mask=block_mask, ++ other=0.0, ++ ).to(tl.float32) ++ block_cdf = tl.cumsum(sums, axis=0) ++ total = tl.sum(sums, axis=0) ++ coin = tl.load(uniform_samples_for_final_sampling + req_idx).to(tl.float32) ++ target = coin * total ++ ++ selected_block = tl.sum( ++ ((block_cdf <= target) & block_mask).to(tl.int32), axis=0 ++ ) ++ selected_block = tl.minimum(selected_block, num_vocab_blocks - 1) ++ prefix_sum = tl.sum( ++ tl.where(block_offsets < selected_block, sums, 0.0), axis=0 ++ ) ++ local_target = tl.maximum(target - prefix_sum, 0.0) ++ ++ local_offsets = tl.arange(0, vocab_block_size) ++ vocab_offsets = selected_block * vocab_block_size + local_offsets ++ vocab_mask = vocab_offsets < vocab_size ++ probs_offset = (req_idx * num_draft_tokens + target_row) * vocab_size ++ target_probs_block = tl.load( ++ target_probs + probs_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ rejected_probs_block = tl.load( ++ rejected_probs + probs_offset + vocab_offsets, ++ mask=vocab_mask, ++ other=0.0, ++ ).to(tl.float32) ++ residual = tl.maximum(target_probs_block - rejected_probs_block, 0.0) ++ ++ local_cdf = tl.cumsum(residual, axis=0) ++ local_index = tl.sum( ++ ((local_cdf <= local_target) & vocab_mask).to(tl.int32), axis=0 ++ ) ++ last_valid_local = tl.max( ++ tl.where((residual > 0.0) & vocab_mask, local_offsets, -1), axis=0 ++ ) ++ valid_local_count = tl.minimum( ++ vocab_block_size, ++ vocab_size - selected_block * vocab_block_size, ++ ) ++ sampled_local = tl.where( ++ local_index < valid_local_count, ++ local_index, ++ last_valid_local, ++ ) ++ sampled_token = tl.where( ++ sampled_local >= 0, ++ selected_block * vocab_block_size + sampled_local, ++ vocab_size - 1, ++ ) ++ sampled_token = tl.minimum(sampled_token, vocab_size - 1) ++ tl.store(predicts + output_idx, sampled_token) ++ ++ ++def tree_speculative_sampling_target_only( ++ predicts: torch.Tensor, ++ accept_index: torch.Tensor, ++ accept_token_num: torch.Tensor, ++ candidates: torch.Tensor, ++ retrive_index: torch.Tensor, ++ retrive_next_token: torch.Tensor, ++ retrive_next_sibling: torch.Tensor, ++ uniform_samples: torch.Tensor, ++ uniform_samples_for_final_sampling: torch.Tensor, ++ target_probs: torch.Tensor, ++ draft_probs: torch.Tensor, ++ threshold_single: float = 1.0, ++ threshold_acc: float = 1.0, ++ deterministic: bool = True, ++) -> None: ++ """NPU port of GPU target-only tree speculative sampling. ++ ++ ``draft_probs`` is scratch storage, matching the GPU API. The function ++ clears it and records rejected sibling probabilities before sampling from ++ ``relu(target_probs - draft_probs)`` on the final selected tree row. ++ """ ++ del deterministic ++ ++ if candidates.ndim != 2 or target_probs.ndim != 3: ++ raise ValueError("candidates must be 2-D and target_probs must be 3-D") ++ ++ batch_size, num_draft_tokens = candidates.shape ++ if batch_size == 0: ++ return ++ if num_draft_tokens == 0: ++ raise ValueError("num_draft_tokens must be positive") ++ if target_probs.shape[:2] != (batch_size, num_draft_tokens): ++ raise ValueError( ++ "target_probs shape must be [batch, num_draft_tokens, vocab_size]" ++ ) ++ tree_shapes = ( ++ retrive_index.shape, ++ retrive_next_token.shape, ++ retrive_next_sibling.shape, ++ uniform_samples.shape, ++ ) ++ if any(shape != candidates.shape for shape in tree_shapes): ++ raise ValueError("all tree-index and uniform tensors must match candidates") ++ if accept_index.ndim != 2 or accept_index.shape[0] != batch_size: ++ raise ValueError("accept_index must be [batch, max_tree_depth]") ++ num_speculative_tokens = accept_index.shape[1] ++ if not 1 <= num_speculative_tokens <= num_draft_tokens: ++ raise ValueError("max_tree_depth must be in [1, num_draft_tokens]") ++ if accept_token_num.shape != (batch_size,): ++ raise ValueError("accept_token_num must have shape [batch]") ++ if predicts.ndim != 1: ++ raise ValueError("predicts must be 1-D") ++ if uniform_samples_for_final_sampling.shape != (batch_size,): ++ raise ValueError( ++ "uniform_samples_for_final_sampling must have shape [batch]" ++ ) ++ if draft_probs.shape != target_probs.shape: ++ raise ValueError("draft_probs scratch must match target_probs") ++ if draft_probs.data_ptr() == target_probs.data_ptr(): ++ raise ValueError("draft_probs must not alias target_probs") ++ if target_probs.dtype != torch.float32 or draft_probs.dtype != torch.float32: ++ raise TypeError("target_probs and draft_probs must be torch.float32") ++ if uniform_samples.dtype != torch.float32: ++ raise TypeError("uniform_samples must be torch.float32") ++ if uniform_samples_for_final_sampling.dtype != torch.float32: ++ raise TypeError("uniform_samples_for_final_sampling must be torch.float32") ++ integer_dtypes = ( ++ (predicts, torch.int32, "predicts"), ++ (accept_index, torch.int32, "accept_index"), ++ (accept_token_num, torch.int32, "accept_token_num"), ++ (candidates, torch.int64, "candidates"), ++ (retrive_index, torch.int64, "retrive_index"), ++ (retrive_next_token, torch.int64, "retrive_next_token"), ++ (retrive_next_sibling, torch.int64, "retrive_next_sibling"), ++ ) ++ for tensor, expected_dtype, name in integer_dtypes: ++ if tensor.dtype != expected_dtype: ++ raise TypeError(f"{name} must be {expected_dtype}") ++ tensors = ( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ retrive_next_token, ++ retrive_next_sibling, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++ draft_probs, ++ ) ++ if any(tensor.device != target_probs.device for tensor in tensors): ++ raise ValueError("all tensors must be on the same NPU device") ++ if any(not tensor.is_contiguous() for tensor in tensors): ++ raise ValueError("all tensors must be contiguous") ++ if not 0.0 <= threshold_single <= 1.0: ++ raise ValueError("threshold_single must be in [0, 1]") ++ if not 0.0 <= threshold_acc <= 1.0: ++ raise ValueError("threshold_acc must be in [0, 1]") ++ ++ threshold_acc = max(float(threshold_acc), 1e-9) ++ vocab_size = target_probs.shape[-1] ++ vocab_block_size = 2048 ++ num_vocab_blocks = triton.cdiv(vocab_size, vocab_block_size) ++ pad_num_vocab_blocks = triton.next_power_of_2(num_vocab_blocks) ++ ++ # The CUDA call site passes zeros_like(target_probs). Clearing in the NPU ++ # wrapper makes the scratch contract explicit and permits empty_like callers. ++ draft_probs.zero_() ++ metadata = torch.empty( ++ (batch_size, 2), dtype=torch.int64, device=target_probs.device ++ ) ++ block_sums = torch.empty( ++ (batch_size, num_vocab_blocks), ++ dtype=torch.float32, ++ device=target_probs.device, ++ ) ++ ++ _tree_target_only_accept_kernel[(batch_size,)]( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ retrive_next_token, ++ retrive_next_sibling, ++ uniform_samples, ++ target_probs, ++ draft_probs, ++ metadata, ++ float(threshold_single), ++ threshold_acc, ++ num_draft_tokens=num_draft_tokens, ++ num_speculative_tokens=num_speculative_tokens, ++ vocab_size=vocab_size, ++ ) ++ _tree_target_only_block_sum_kernel[(batch_size, num_vocab_blocks)]( ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ num_draft_tokens=num_draft_tokens, ++ vocab_size=vocab_size, ++ vocab_block_size=vocab_block_size, ++ num_vocab_blocks=num_vocab_blocks, ++ ) ++ _tree_target_only_sample_kernel[(batch_size,)]( ++ predicts, ++ target_probs, ++ draft_probs, ++ metadata, ++ block_sums, ++ uniform_samples_for_final_sampling, ++ num_draft_tokens=num_draft_tokens, ++ vocab_size=vocab_size, ++ vocab_block_size=vocab_block_size, ++ num_vocab_blocks=num_vocab_blocks, ++ pad_num_vocab_blocks=pad_num_vocab_blocks, ++ ) +diff --git a/tests/python/sgl_kernel_npu/test_chain_speculative_sampling.py b/tests/python/sgl_kernel_npu/test_chain_speculative_sampling.py +new file mode 100644 +index 0000000..40b04be +--- /dev/null ++++ b/tests/python/sgl_kernel_npu/test_chain_speculative_sampling.py +@@ -0,0 +1,138 @@ ++import torch ++import torch_npu # noqa: F401 ++ ++from sgl_kernel_npu.sample import chain_speculative_sampling_rejection ++ ++ ++def chain_rejection_reference( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++ draft_probs, ++): ++ batch_size, num_draft_tokens = candidates.shape ++ for req_idx in range(batch_size): ++ cur_prob_row = 0 ++ last_accepted_idx = int(retrive_index[req_idx, 0]) ++ accept_index[req_idx, 0] = last_accepted_idx ++ num_accepted = 0 ++ all_accepted = True ++ ++ for step in range(1, num_draft_tokens): ++ draft_token = int(candidates[req_idx, step]) ++ p = float(target_probs[req_idx, cur_prob_row, draft_token]) ++ q = float(draft_probs[req_idx, cur_prob_row, draft_token]) ++ coin = float(uniform_samples[req_idx, step - 1]) ++ if coin * q < p: ++ predicts[last_accepted_idx] = draft_token ++ num_accepted += 1 ++ last_accepted_idx = int(retrive_index[req_idx, step]) ++ accept_index[req_idx, num_accepted] = last_accepted_idx ++ cur_prob_row = step ++ else: ++ all_accepted = False ++ break ++ ++ accept_token_num[req_idx] = num_accepted ++ residual = target_probs[req_idx, cur_prob_row].clone() ++ if not all_accepted: ++ residual.sub_(draft_probs[req_idx, cur_prob_row]).clamp_min_(0.0) ++ target = float(uniform_samples_for_final_sampling[req_idx]) * float( ++ residual.sum() ++ ) ++ sampled_token = int((residual.cumsum(0) <= target).sum()) ++ if sampled_token == residual.numel(): ++ positive = torch.nonzero(residual > 0.0).flatten() ++ sampled_token = ( ++ int(positive[-1]) if positive.numel() else residual.numel() - 1 ++ ) ++ predicts[last_accepted_idx] = sampled_token ++ ++ ++def test_chain_rejection_matches_gpu_algorithm(): ++ batch_size, num_draft_tokens, vocab_size = 2, 4, 11 ++ candidates = torch.tensor([[0, 2, 3, 4], [0, 5, 6, 7]]) ++ retrive_index = torch.arange(batch_size * num_draft_tokens).view( ++ batch_size, num_draft_tokens ++ ) ++ target_probs = torch.softmax( ++ torch.tensor( ++ [ ++ [ ++ [0.1, 0.2, 2.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], ++ [0.1, 0.2, 0.1, 2.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], ++ [0.1, 0.2, 0.1, 0.1, 2.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], ++ [0.1, 0.2, 0.1, 0.1, 2.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1], ++ ], ++ [ ++ [0.1, 0.1, 0.1, 0.1, 0.1, 0.2, 0.1, 2.0, 0.1, 0.1, 0.1], ++ [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 2.0, 0.2, 0.1, 0.1, 0.1], ++ [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 2.0, 0.1, 0.1, 0.1], ++ [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 2.0, 0.1, 0.1, 0.1], ++ ], ++ ], ++ dtype=torch.float32, ++ ), ++ dim=-1, ++ ) ++ draft_probs = torch.softmax( ++ torch.tensor( ++ [ ++ [[0.1] * vocab_size, [0.1] * vocab_size, [0.1] * vocab_size], ++ [[0.1] * vocab_size, [0.1] * vocab_size, [0.1] * vocab_size], ++ ], ++ dtype=torch.float32, ++ ), ++ dim=-1, ++ ) ++ draft_probs[1, 0, 5] = 0.9 ++ draft_probs[1, 0] /= draft_probs[1, 0].sum() ++ uniforms = torch.tensor([[0.1, 0.1, 0.1, 0.0], [0.99, 0.0, 0.0, 0.0]]) ++ final_uniforms = torch.tensor([0.37, 0.61]) ++ ++ expected_predicts = torch.full( ++ (batch_size * num_draft_tokens,), -1, dtype=torch.int32 ++ ) ++ expected_accept_index = torch.full( ++ (batch_size, num_draft_tokens), -1, dtype=torch.int32 ++ ) ++ expected_accept_num = torch.zeros(batch_size, dtype=torch.int32) ++ chain_rejection_reference( ++ expected_predicts, ++ expected_accept_index, ++ expected_accept_num, ++ candidates, ++ retrive_index, ++ uniforms, ++ final_uniforms, ++ target_probs, ++ draft_probs, ++ ) ++ ++ predicts = torch.full_like(expected_predicts, -1, device="npu") ++ accept_index = torch.full_like(expected_accept_index, -1, device="npu") ++ accept_num = torch.zeros_like(expected_accept_num, device="npu") ++ next_token = torch.full_like(candidates, -1, device="npu") ++ next_sibling = torch.full_like(candidates, -1, device="npu") ++ chain_speculative_sampling_rejection( ++ predicts, ++ accept_index, ++ accept_num, ++ candidates.npu(), ++ retrive_index.npu(), ++ next_token, ++ next_sibling, ++ uniforms.npu(), ++ final_uniforms.npu(), ++ target_probs.npu(), ++ draft_probs.npu(), ++ ) ++ ++ torch.testing.assert_close(predicts.cpu(), expected_predicts) ++ torch.testing.assert_close(accept_index.cpu(), expected_accept_index) ++ torch.testing.assert_close(accept_num.cpu(), expected_accept_num) +diff --git a/tests/python/sgl_kernel_npu/test_speculative_probability.py b/tests/python/sgl_kernel_npu/test_speculative_probability.py +new file mode 100644 +index 0000000..563f174 +--- /dev/null ++++ b/tests/python/sgl_kernel_npu/test_speculative_probability.py +@@ -0,0 +1,31 @@ ++import torch ++ ++from sgl_kernel_npu.sample.probability import top_k_top_p_renorm_probs ++ ++ ++def test_top_k_top_p_renorm_matches_sequential_reference(): ++ torch.manual_seed(7) ++ probs = torch.softmax(torch.randn(4, 97), dim=-1) ++ top_ks = torch.tensor([1, 7, 31, 97]) ++ top_ps = torch.tensor([0.3, 0.75, 0.95, 1.0]) ++ ++ actual = top_k_top_p_renorm_probs( ++ probs, top_ks, top_ps, True, True ++ ) ++ ++ sorted_probs, sorted_indices = probs.sort(dim=-1, descending=True) ++ positions = torch.arange(probs.shape[-1]).view(1, -1) ++ sorted_probs[positions >= top_ks.view(-1, 1)] = 0.0 ++ sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True) ++ top_k_probs = torch.zeros_like(probs).scatter( ++ -1, sorted_indices, sorted_probs ++ ) ++ sorted_probs, sorted_indices = top_k_probs.sort(dim=-1, descending=True) ++ cumulative = sorted_probs.cumsum(dim=-1) ++ sorted_probs[cumulative - sorted_probs > top_ps.view(-1, 1)] = 0.0 ++ sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True) ++ expected = torch.zeros_like(probs).scatter( ++ -1, sorted_indices, sorted_probs ++ ) ++ ++ torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-7) +diff --git a/tests/python/sgl_kernel_npu/test_tree_speculative_sampling_target_only.py b/tests/python/sgl_kernel_npu/test_tree_speculative_sampling_target_only.py +new file mode 100644 +index 0000000..3c12d53 +--- /dev/null ++++ b/tests/python/sgl_kernel_npu/test_tree_speculative_sampling_target_only.py +@@ -0,0 +1,595 @@ ++import argparse ++import time ++ ++import pytest ++import torch ++import torch_npu # noqa: F401 ++ ++from sgl_kernel_npu.sample import tree_speculative_sampling_target_only ++ ++ ++def target_only_tree_reference( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ retrive_next_token, ++ retrive_next_sibling, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++ rejected_probs, ++ threshold_single, ++ threshold_acc, ++): ++ """CPU reference translated directly from the GPU CUDA kernel.""" ++ batch_size, num_draft_tokens = candidates.shape ++ num_speculative_tokens = accept_index.shape[1] ++ threshold_acc = max(float(threshold_acc), 1e-9) ++ rejected_probs.zero_() ++ ++ for req_idx in range(batch_size): ++ cur_prob_row = 0 ++ cur_node = 0 ++ coin = float(uniform_samples[req_idx, 0]) ++ last_accepted_idx = int(retrive_index[req_idx, 0]) ++ accept_index[req_idx, 0] = last_accepted_idx ++ num_accepted = 0 ++ ++ for _ in range(1, num_speculative_tokens): ++ cur_node = int(retrive_next_token[req_idx, cur_node]) ++ prob_acc = 0.0 ++ while cur_node != -1: ++ draft_idx = int(retrive_index[req_idx, cur_node]) ++ draft_token = int(candidates[req_idx, cur_node]) ++ target_prob = float( ++ target_probs[req_idx, cur_prob_row, draft_token] ++ ) ++ prob_acc += target_prob ++ if ( ++ coin <= prob_acc / threshold_acc ++ or target_prob >= threshold_single ++ ): ++ predicts[last_accepted_idx] = draft_token ++ num_accepted += 1 ++ accept_index[req_idx, num_accepted] = draft_idx ++ last_accepted_idx = draft_idx ++ cur_prob_row = cur_node ++ coin = float(uniform_samples[req_idx, cur_node]) ++ break ++ ++ rejected_probs[req_idx, cur_prob_row, draft_token] = target_prob ++ cur_node = int(retrive_next_sibling[req_idx, cur_node]) ++ ++ if cur_node == -1: ++ break ++ ++ accept_token_num[req_idx] = num_accepted ++ residual = ( ++ target_probs[req_idx, cur_prob_row] ++ - rejected_probs[req_idx, cur_prob_row] ++ ).clamp_min(0.0) ++ target = float(uniform_samples_for_final_sampling[req_idx]) * float( ++ residual.sum() ++ ) ++ sampled_token = int((residual.cumsum(0) <= target).sum()) ++ if sampled_token == residual.numel(): ++ positive = torch.nonzero(residual > 0.0).flatten() ++ sampled_token = ( ++ int(positive[-1]) if positive.numel() else residual.numel() - 1 ++ ) ++ predicts[last_accepted_idx] = sampled_token ++ ++ return predicts, accept_index, accept_token_num, rejected_probs ++ ++ ++def target_only_chain_reference( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++ threshold_single, ++ threshold_acc, ++): ++ batch_size, num_draft_tokens = candidates.shape ++ threshold_acc = max(float(threshold_acc), 1e-9) ++ ++ for req_idx in range(batch_size): ++ last_accepted_idx = int(retrive_index[req_idx, 0]) ++ accept_index[req_idx, 0] = last_accepted_idx ++ num_accepted = 0 ++ rejected_token = -1 ++ ++ for step in range(1, num_draft_tokens): ++ draft_token = int(candidates[req_idx, step]) ++ target_prob = float(target_probs[req_idx, step - 1, draft_token]) ++ coin = float(uniform_samples[req_idx, step - 1]) ++ if ( ++ coin <= target_prob / threshold_acc ++ or target_prob >= threshold_single ++ ): ++ predicts[last_accepted_idx] = draft_token ++ num_accepted += 1 ++ last_accepted_idx = int(retrive_index[req_idx, step]) ++ accept_index[req_idx, num_accepted] = last_accepted_idx ++ else: ++ rejected_token = draft_token ++ break ++ ++ accept_token_num[req_idx] = num_accepted ++ final_probs = target_probs[req_idx, num_accepted].clone().float() ++ if rejected_token >= 0: ++ final_probs[rejected_token] = 0.0 ++ final_probs.clamp_min_(0.0) ++ target = float(uniform_samples_for_final_sampling[req_idx]) * float( ++ final_probs.sum() ++ ) ++ sampled_token = int((final_probs.cumsum(0) <= target).sum()) ++ sampled_token = min(sampled_token, final_probs.numel() - 1) ++ predicts[last_accepted_idx] = sampled_token ++ ++ return predicts, accept_index, accept_token_num ++ ++ ++@pytest.mark.parametrize( ++ "threshold_single,threshold_acc", [(1.0, 1.0), (0.0, 0.0), (0.5, 0.8)] ++) ++def test_general_tree_matches_gpu_algorithm_reference( ++ threshold_single, threshold_acc ++): ++ candidates = torch.tensor( ++ [[0, 1, 2, 3, 4, 5], [7, 8, 9, 10, 11, 12]], dtype=torch.int64 ++ ) ++ retrive_index = torch.tensor( ++ [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]], dtype=torch.int64 ++ ) ++ retrive_next_token = torch.tensor( ++ [[1, 2, -1, 4, 5, -1], [4, 2, 3, -1, 5, -1]], ++ dtype=torch.int64, ++ ) ++ retrive_next_sibling = torch.tensor( ++ [[-1, 3, -1, -1, -1, -1], [-1, -1, -1, -1, 1, -1]], ++ dtype=torch.int64, ++ ) ++ batch_size, num_draft_tokens = candidates.shape ++ vocab_size = 20 ++ target_probs = torch.full( ++ (batch_size, num_draft_tokens, vocab_size), 0.01, dtype=torch.float32 ++ ) ++ target_probs[0, 0, 1] = 0.12 ++ target_probs[0, 0, 3] = 0.72 ++ target_probs[0, 3, 4] = 0.82 ++ target_probs[0, 4, 5] = 0.75 ++ target_probs[1, 0, 11] = 0.68 ++ target_probs[1, 0, 8] = 0.14 ++ target_probs[1, 4, 12] = 0.77 ++ target_probs /= target_probs.sum(dim=-1, keepdim=True) ++ uniforms = torch.tensor( ++ [[0.55, 0.2, 0.8, 0.4, 0.3, 0.9], [0.6, 0.2, 0.8, 0.7, 0.3, 0.4]], ++ dtype=torch.float32, ++ ) ++ final_uniforms = torch.tensor([0.25, 0.75], dtype=torch.float32) ++ ++ ref_predicts = torch.full((12,), -1, dtype=torch.int32) ++ ref_accept_index = torch.full((2, 4), -1, dtype=torch.int32) ++ ref_accept_num = torch.zeros(2, dtype=torch.int32) ++ ref_rejected = torch.empty_like(target_probs) ++ target_only_tree_reference( ++ ref_predicts, ++ ref_accept_index, ++ ref_accept_num, ++ candidates, ++ retrive_index, ++ retrive_next_token, ++ retrive_next_sibling, ++ uniforms, ++ final_uniforms, ++ target_probs, ++ ref_rejected, ++ threshold_single, ++ threshold_acc, ++ ) ++ ++ npu_predicts = torch.full_like(ref_predicts, -1, device="npu") ++ npu_accept_index = torch.full_like(ref_accept_index, -1, device="npu") ++ npu_accept_num = torch.zeros_like(ref_accept_num, device="npu") ++ npu_rejected = torch.empty_like(target_probs, device="npu") ++ tree_speculative_sampling_target_only( ++ npu_predicts, ++ npu_accept_index, ++ npu_accept_num, ++ candidates.npu(), ++ retrive_index.npu(), ++ retrive_next_token.npu(), ++ retrive_next_sibling.npu(), ++ uniforms.npu(), ++ final_uniforms.npu(), ++ target_probs.npu(), ++ npu_rejected, ++ threshold_single, ++ threshold_acc, ++ True, ++ ) ++ ++ torch.testing.assert_close(npu_predicts.cpu(), ref_predicts, rtol=0, atol=0) ++ torch.testing.assert_close( ++ npu_accept_index.cpu(), ref_accept_index, rtol=0, atol=0 ++ ) ++ torch.testing.assert_close(npu_accept_num.cpu(), ref_accept_num, rtol=0, atol=0) ++ torch.testing.assert_close(npu_rejected.cpu(), ref_rejected, rtol=0, atol=0) ++ ++ ++def target_only_chain_torch( ++ predicts, ++ accept_index, ++ accept_token_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ uniform_samples_for_final_sampling, ++ target_probs, ++): ++ batch_size, num_draft_tokens = candidates.shape ++ device = candidates.device ++ draft_tokens = candidates[:, 1:].long() ++ step_probs = torch.gather( ++ target_probs[:, :-1, :], 2, draft_tokens.unsqueeze(-1) ++ ).squeeze(-1) ++ accept_steps = uniform_samples[:, : num_draft_tokens - 1] <= step_probs ++ reject_count = (~accept_steps).to(torch.int32).cumsum(dim=1) ++ num_correct = (reject_count == 0).to(torch.int32).sum(dim=1) ++ ++ accept_token_num.copy_(num_correct) ++ accept_index.fill_(-1) ++ positions = torch.arange(num_draft_tokens, device=device).view(1, -1) ++ valid_accept = positions <= num_correct.view(-1, 1) ++ accept_index.copy_( ++ torch.where( ++ valid_accept, ++ retrive_index.to(torch.int32), ++ torch.full_like(accept_index, -1), ++ ) ++ ) ++ ++ predicts.zero_() ++ parent_positions = torch.arange(num_draft_tokens - 1, device=device).view(1, -1) ++ valid_parent = parent_positions < num_correct.view(-1, 1) ++ parent_indices = retrive_index[:, :-1].reshape(-1).long() ++ parent_values = candidates[:, 1:].to(torch.int32).reshape(-1) ++ predicts[parent_indices] = torch.where( ++ valid_parent.reshape(-1), parent_values, predicts[parent_indices] ++ ) ++ ++ rows = torch.arange(batch_size, device=device) ++ final_rows = num_correct.long() ++ final_probs = target_probs[rows, final_rows].clone() ++ rejected = num_correct < num_draft_tokens - 1 ++ rejected_positions = (num_correct.long() + 1).clamp_max(num_draft_tokens - 1) ++ rejected_tokens = candidates[rows, rejected_positions].long() ++ final_probs[rejected, rejected_tokens[rejected]] = 0.0 ++ ++ probability_sums = final_probs.sum(dim=-1, keepdim=True) ++ targets = uniform_samples_for_final_sampling.view(-1, 1) * probability_sums ++ final_tokens = ( ++ (final_probs.cumsum(dim=-1) <= targets) ++ .to(torch.int32) ++ .sum(dim=-1) ++ .clamp_max(target_probs.shape[-1] - 1) ++ ) ++ final_indices = retrive_index[rows, final_rows].long() ++ predicts[final_indices] = final_tokens.to(torch.int32) ++ ++ ++def make_chain_indices(batch_size, num_draft_tokens, device): ++ retrive_index = torch.arange( ++ batch_size * num_draft_tokens, dtype=torch.int64, device=device ++ ).view(batch_size, num_draft_tokens) ++ retrive_next_token = torch.arange( ++ 1, num_draft_tokens + 1, dtype=torch.int64, device=device ++ ).repeat(batch_size, 1) ++ retrive_next_token[:, -1] = -1 ++ retrive_next_sibling = torch.full_like(retrive_next_token, -1) ++ return retrive_index, retrive_next_token, retrive_next_sibling ++ ++ ++def make_stable_chain_final_uniforms( ++ candidates, ++ uniform_samples, ++ target_probs, ++): ++ """Choose final-sampling coins away from inverse-CDF boundaries.""" ++ batch_size, num_draft_tokens = candidates.shape ++ final_uniforms = torch.empty(batch_size, dtype=torch.float32) ++ ++ for req_idx in range(batch_size): ++ num_accepted = 0 ++ rejected_token = -1 ++ for step in range(1, num_draft_tokens): ++ draft_token = int(candidates[req_idx, step]) ++ target_prob = float(target_probs[req_idx, step - 1, draft_token]) ++ if float(uniform_samples[req_idx, step - 1]) <= target_prob: ++ num_accepted += 1 ++ else: ++ rejected_token = draft_token ++ break ++ ++ final_probs = target_probs[req_idx, num_accepted].double().clone() ++ if rejected_token >= 0: ++ final_probs[rejected_token] = 0.0 ++ ++ sampled_token = int(final_probs.argmax()) ++ probability_sum = final_probs.sum() ++ cdf_before = final_probs[:sampled_token].sum() ++ cdf_midpoint = cdf_before + final_probs[sampled_token] * 0.5 ++ final_uniforms[req_idx] = (cdf_midpoint / probability_sum).float() ++ ++ return final_uniforms ++ ++ ++@pytest.mark.parametrize("batch_size", [1, 4, 17]) ++@pytest.mark.parametrize("num_draft_tokens", [2, 5]) ++@pytest.mark.parametrize("vocab_size", [20, 32000, 151552]) ++def test_target_only_chain_matches_reference( ++ batch_size, num_draft_tokens, vocab_size ++): ++ torch.manual_seed(20260717 + batch_size + num_draft_tokens + vocab_size) ++ candidates = torch.randint( ++ 0, vocab_size, (batch_size, num_draft_tokens), dtype=torch.int64 ++ ) ++ logits = torch.randn(batch_size, num_draft_tokens, vocab_size) ++ target_probs = torch.softmax(logits, dim=-1).float() ++ ++ # Give some draft tokens meaningful acceptance probability. ++ for req_idx in range(batch_size): ++ for step in range(1, num_draft_tokens): ++ token = int(candidates[req_idx, step]) ++ target_probs[req_idx, step - 1] *= 0.35 ++ target_probs[req_idx, step - 1, token] += 0.65 ++ target_probs[req_idx, step - 1] /= target_probs[ ++ req_idx, step - 1 ++ ].sum() ++ ++ uniform_samples = torch.rand(batch_size, num_draft_tokens) ++ # A random coin can land within FP32 reduction error of a CDF boundary for ++ # large vocabularies. Use the midpoint of a high-mass token's interval so ++ # exact token equality tests the algorithm instead of reduction order. ++ final_uniform_samples = make_stable_chain_final_uniforms( ++ candidates, ++ uniform_samples, ++ target_probs, ++ ) ++ retrive_index, retrive_next_token, retrive_next_sibling = make_chain_indices( ++ batch_size, num_draft_tokens, "cpu" ++ ) ++ ++ ref_predicts = torch.full( ++ (batch_size * num_draft_tokens,), -1, dtype=torch.int32 ++ ) ++ ref_accept_index = torch.full( ++ (batch_size, num_draft_tokens), -1, dtype=torch.int32 ++ ) ++ ref_accept_num = torch.zeros(batch_size, dtype=torch.int32) ++ target_only_chain_reference( ++ ref_predicts, ++ ref_accept_index, ++ ref_accept_num, ++ candidates, ++ retrive_index, ++ uniform_samples, ++ final_uniform_samples, ++ target_probs, ++ 1.0, ++ 1.0, ++ ) ++ ++ npu_predicts = torch.full_like(ref_predicts, -1, device="npu") ++ npu_accept_index = torch.full_like(ref_accept_index, -1, device="npu") ++ npu_accept_num = torch.zeros_like(ref_accept_num, device="npu") ++ candidates_npu = candidates.npu() ++ retrive_index_npu = retrive_index.npu() ++ next_token_npu = retrive_next_token.npu() ++ next_sibling_npu = retrive_next_sibling.npu() ++ target_probs_npu = target_probs.npu() ++ ++ tree_speculative_sampling_target_only( ++ predicts=npu_predicts, ++ accept_index=npu_accept_index, ++ accept_token_num=npu_accept_num, ++ candidates=candidates_npu, ++ retrive_index=retrive_index_npu, ++ retrive_next_token=next_token_npu, ++ retrive_next_sibling=next_sibling_npu, ++ uniform_samples=uniform_samples.npu(), ++ uniform_samples_for_final_sampling=final_uniform_samples.npu(), ++ target_probs=target_probs_npu, ++ draft_probs=torch.empty_like(target_probs_npu), ++ threshold_single=1.0, ++ threshold_acc=1.0, ++ deterministic=True, ++ ) ++ ++ torch.testing.assert_close(npu_predicts.cpu(), ref_predicts, rtol=0, atol=0) ++ torch.testing.assert_close( ++ npu_accept_index.cpu(), ref_accept_index, rtol=0, atol=0 ++ ) ++ torch.testing.assert_close(npu_accept_num.cpu(), ref_accept_num, rtol=0, atol=0) ++ ++ ++@pytest.mark.parametrize( ++ "threshold_single,threshold_acc", ++ [(1.0, 1.0), (0.0, 0.0), (0.5, 0.8)], ++) ++def test_target_only_thresholds(threshold_single, threshold_acc): ++ batch_size, num_draft_tokens, vocab_size = 2, 4, 32 ++ candidates = torch.tensor([[0, 3, 4, 5], [0, 7, 8, 9]], dtype=torch.int64) ++ target_probs = torch.full( ++ (batch_size, num_draft_tokens, vocab_size), 1.0 / vocab_size ++ ) ++ for req_idx in range(batch_size): ++ for step in range(1, num_draft_tokens): ++ token = int(candidates[req_idx, step]) ++ target_probs[req_idx, step - 1] *= 0.2 ++ target_probs[req_idx, step - 1, token] += 0.8 ++ target_probs[req_idx, step - 1] /= target_probs[ ++ req_idx, step - 1 ++ ].sum() ++ ++ uniforms = torch.tensor([[0.1, 0.9, 0.2, 0.0], [0.7, 0.2, 0.95, 0.0]]) ++ final_uniforms = torch.tensor([0.25, 0.75]) ++ retrive_index, next_token, next_sibling = make_chain_indices( ++ batch_size, num_draft_tokens, "cpu" ++ ) ++ ++ ref_predicts = torch.full((batch_size * num_draft_tokens,), -1, dtype=torch.int32) ++ ref_accept_index = torch.full( ++ (batch_size, num_draft_tokens), -1, dtype=torch.int32 ++ ) ++ ref_accept_num = torch.zeros(batch_size, dtype=torch.int32) ++ target_only_chain_reference( ++ ref_predicts, ++ ref_accept_index, ++ ref_accept_num, ++ candidates, ++ retrive_index, ++ uniforms, ++ final_uniforms, ++ target_probs, ++ threshold_single, ++ threshold_acc, ++ ) ++ ++ npu_predicts = torch.full_like(ref_predicts, -1, device="npu") ++ npu_accept_index = torch.full_like(ref_accept_index, -1, device="npu") ++ npu_accept_num = torch.zeros_like(ref_accept_num, device="npu") ++ target_probs_npu = target_probs.npu() ++ tree_speculative_sampling_target_only( ++ npu_predicts, ++ npu_accept_index, ++ npu_accept_num, ++ candidates.npu(), ++ retrive_index.npu(), ++ next_token.npu(), ++ next_sibling.npu(), ++ uniforms.npu(), ++ final_uniforms.npu(), ++ target_probs_npu, ++ torch.empty_like(target_probs_npu), ++ threshold_single, ++ threshold_acc, ++ True, ++ ) ++ ++ torch.testing.assert_close(npu_predicts.cpu(), ref_predicts, rtol=0, atol=0) ++ torch.testing.assert_close( ++ npu_accept_index.cpu(), ref_accept_index, rtol=0, atol=0 ++ ) ++ torch.testing.assert_close(npu_accept_num.cpu(), ref_accept_num, rtol=0, atol=0) ++ ++ ++def run_benchmark(batch_size, num_draft_tokens, vocab_size, warmup, iterations): ++ candidates = torch.randint( ++ 0, ++ vocab_size, ++ (batch_size, num_draft_tokens), ++ dtype=torch.int64, ++ device="npu", ++ ) ++ target_probs = torch.softmax( ++ torch.randn( ++ batch_size, ++ num_draft_tokens, ++ vocab_size, ++ dtype=torch.float32, ++ device="npu", ++ ), ++ dim=-1, ++ ) ++ retrive_index, next_token, next_sibling = make_chain_indices( ++ batch_size, num_draft_tokens, "npu" ++ ) ++ uniforms = torch.rand( ++ batch_size, num_draft_tokens, dtype=torch.float32, device="npu" ++ ) ++ final_uniforms = torch.rand(batch_size, dtype=torch.float32, device="npu") ++ draft_probs = torch.empty_like(target_probs) ++ predicts = torch.zeros( ++ batch_size * num_draft_tokens, dtype=torch.int32, device="npu" ++ ) ++ accept_index = torch.full( ++ (batch_size, num_draft_tokens), -1, dtype=torch.int32, device="npu" ++ ) ++ accept_num = torch.zeros(batch_size, dtype=torch.int32, device="npu") ++ ++ def run_kernel(): ++ tree_speculative_sampling_target_only( ++ predicts, ++ accept_index, ++ accept_num, ++ candidates, ++ retrive_index, ++ next_token, ++ next_sibling, ++ uniforms, ++ final_uniforms, ++ target_probs, ++ draft_probs, ++ 1.0, ++ 1.0, ++ True, ++ ) ++ ++ def run_torch(): ++ target_only_chain_torch( ++ predicts, ++ accept_index, ++ accept_num, ++ candidates, ++ retrive_index, ++ uniforms, ++ final_uniforms, ++ target_probs, ++ ) ++ ++ def benchmark(fn): ++ for _ in range(warmup): ++ fn() ++ torch.npu.synchronize() ++ started = time.perf_counter() ++ for _ in range(iterations): ++ fn() ++ torch.npu.synchronize() ++ return (time.perf_counter() - started) * 1000 / iterations ++ ++ kernel_latency_ms = benchmark(run_kernel) ++ torch_latency_ms = benchmark(run_torch) ++ print( ++ f"batch={batch_size} drafts={num_draft_tokens} vocab={vocab_size} " ++ f"kernel_ms={kernel_latency_ms:.4f} torch_ms={torch_latency_ms:.4f} " ++ f"speedup={torch_latency_ms / kernel_latency_ms:.2f}x" ++ ) ++ ++ ++if __name__ == "__main__": ++ parser = argparse.ArgumentParser() ++ parser.add_argument("--perf", action="store_true") ++ parser.add_argument("--batch-size", type=int, default=16) ++ parser.add_argument("--num-draft-tokens", type=int, default=5) ++ parser.add_argument("--vocab-size", type=int, default=151552) ++ parser.add_argument("--warmup", type=int, default=10) ++ parser.add_argument("--iterations", type=int, default=100) ++ args = parser.parse_args() ++ if args.perf: ++ run_benchmark( ++ args.batch_size, ++ args.num_draft_tokens, ++ args.vocab_size, ++ args.warmup, ++ args.iterations, ++ ) ++ else: ++ raise SystemExit(pytest.main([__file__])) diff --git a/docker/npu_patch/sglang-npu.patch b/docker/npu_patch/sglang-npu.patch index de38657b6..52cbe46e0 100644 --- a/docker/npu_patch/sglang-npu.patch +++ b/docker/npu_patch/sglang-npu.patch @@ -1,56 +1,1058 @@ +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 ++++ b/python/sglang/srt/entrypoints/http_server.py +@@ -598,7 +598,7 @@ async def health_generate(request: Request) -> Response: + ): + return Response(status_code=200) + +- sampling_params = {"max_new_tokens": 1, "temperature": 0.0} ++ sampling_params = {"max_new_tokens": 1, "temperature": 1.0,"top_k":-1,"top_p":1.0} + # 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/hardware_backend/npu/attention/ascend_gdn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +index 7f506e7f1a..204b9382e8 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 ( + fused_gdn_gating_kernel_without_sigmoid, + fused_gdn_gating_npu, + ) +-from sgl_kernel_npu.mamba.causal_conv1d import ( +- causal_conv1d_fn_npu, +- causal_conv1d_update_npu, +- causal_conv1d_update_v2, +-) + + 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 + from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput + + fused_gdn_gating = fused_gdn_gating_npu +-causal_conv1d_fn = causal_conv1d_fn_npu +-causal_conv1d_update = causal_conv1d_update_npu + + + class AscendGDNAttnBackend(AscendMambaAttnBackendBase): +@@ -109,6 +102,13 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): + self._prepare_mamba_track_metadata(forward_batch) + self.graph_mode = False + ++ def _get_conv_weights_t(self, layer: RadixLinearAttention) -> torch.Tensor: ++ w = getattr(layer, "_conv_weights_t", None) ++ if w is None: ++ w = layer.conv_weights.transpose(0, 1).contiguous() ++ layer._conv_weights_t = w ++ return w ++ + def forward_decode( + self, + layer: RadixLinearAttention, +@@ -125,16 +125,17 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): + cache_indices = self.forward_metadata.mamba_cache_indices + + 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, +- 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) + + query, key, value = torch.split( + mixed_qkv, +@@ -219,44 +220,41 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): + dtype=torch.int32, + device=mixed_qkv.device, + ) +- mixed_qkv = causal_conv1d_update_v2( +- x=mixed_qkv.view(batch_size, draft_token_num, -1).contiguous(), +- conv_state=conv_states.contiguous(), +- weight=layer.conv_weights.transpose(0, 1).contiguous(), ++ mixed_qkv = torch.ops.npu.causal_conv1d( ++ mixed_qkv, ++ self._get_conv_weights_t(layer), ++ conv_states=conv_states, + bias=layer.bias, +- activation=layer.activation, +- conv_state_indices=cache_indices, ++ query_start_loc=query_start_loc, ++ cache_indices=cache_indices, + num_accepted_tokens=num_accepted_tokens, ++ activation_mode=1, + pad_slot_id=-1, +- validate_data=False, +- ).view(seq_len, -1) ++ run_mode=1, ++ ) + else: +- mixed_qkv = mixed_qkv.transpose(0, 1) + if forward_metadata.has_mamba_track_mask: +- mixed_qkv_to_track = mixed_qkv[ +- :, forward_metadata.track_conv_indices +- ].transpose(0, 1) +- conv_states.transpose(1, 2)[ +- forward_metadata.conv_states_mask_indices +- ] = mixed_qkv_to_track ++ mixed_qkv_to_track = mixed_qkv[forward_metadata.track_conv_indices] ++ conv_states[forward_metadata.conv_states_mask_indices] = ( ++ mixed_qkv_to_track ++ ) + kernel_size = layer.conv_weights.shape[-1] +- conv_states_for_prefill = conv_states[:, -(kernel_size - 1) :, :] +- conv_states_tmp = conv_states_for_prefill.transpose(1, 2).contiguous() +- +- mixed_qkv = causal_conv1d_fn( ++ conv_states_for_prefill = conv_states[ ++ :, -(kernel_size - 1) :, : ++ ].contiguous() ++ mixed_qkv = torch.ops.npu.causal_conv1d( + mixed_qkv, +- layer.conv_weights, +- layer.bias, +- activation=layer.activation, +- conv_states=conv_states_tmp, +- has_initial_state=has_initial_states, +- cache_indices=cache_indices, ++ self._get_conv_weights_t(layer), ++ conv_states=conv_states_for_prefill, ++ bias=layer.bias, + query_start_loc=query_start_loc, +- seq_lens_cpu=forward_batch.extend_seq_lens_cpu, +- ).transpose(0, 1)[:seq_len] +- conv_states[:, -(kernel_size - 1) :, :] = conv_states_tmp.transpose( +- 1, 2 +- ).contiguous() ++ cache_indices=cache_indices, ++ has_initial_state=has_initial_states, ++ activation_mode=1, ++ pad_slot_id=-1, ++ run_mode=0, ++ ) ++ conv_states[:, -(kernel_size - 1) :, :] = conv_states_for_prefill + 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/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 ++++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py +@@ -37,6 +37,21 @@ if TYPE_CHECKING: + 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 ++ ++ 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/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py -index 1ade4ed9e4..0161bd398a 100644 +index 82a3d77f05..25d4c21406 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py -@@ -314,12 +314,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): +@@ -402,10 +402,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): + layer.w2_weight.data = layer.w2_weight.data.reshape( layer.num_local_experts, *new_shape_w2 ) - - if _is_npu: - for weight_name in ["w13_weight", "w2_weight"]: - weight = getattr(layer, weight_name) -- weight.data = weight.data.transpose(1, 2) -- weight.data = npu_format_cast(weight.data) -- +- weight.data = npu_format_cast(weight) ++ # if _is_npu: ++ # for weight_name in ["w13_weight", "w2_weight"]: ++ # weight = getattr(layer, weight_name) ++ # weight.data = npu_format_cast(weight) + return - def maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load( -@@ -646,7 +640,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): - # gmm1: gate_up_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], -- weight=[layer.w13_weight], -+ weight=[layer.w13_weight.transpose(1, 2)], - bias=w13_bias, - split_item=2, - group_list_type=1, -@@ -670,7 +664,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): - # gmm2: down_proj - hidden_states = torch.ops.npu.npu_grouped_matmul( - x=[hidden_states], -- weight=[layer.w2_weight], -+ weight=[layer.w2_weight.transpose(1, 2)], - bias=w2_bias, - split_item=2, - group_list_type=1, +diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py +index e83e4157fa..19ed3f244b 100644 +--- a/python/sglang/srt/layers/sampler.py ++++ b/python/sglang/srt/layers/sampler.py +@@ -81,6 +81,170 @@ 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" ++ # Generate uniform noise on a side stream while the model is running, and ++ # apply the exponential-race transform at consumption time on the main ++ # stream, exactly matching stock aten::exponential_ semantics (NPU ++ # op-plugin composite): x = min(1 - u, 1 - eps/2), q = -log(x). ++ # uniform_() decomposes to DSARandomUniform plus an async ++ # D2D copy, i.e. pure DSA-engine work with no AIV transform kernels, so ++ # the side stream cannot contend with the model's vector-core kernels ++ # 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( ++ "SGLANG_NPU_ASYNC_EXPONENTIAL" ++ ) ++ 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 ++ ++ def can_prepare_async_exponential( ++ self, sampling_info: SamplingBatchInfo ++ ) -> bool: ++ """Return whether this batch can use precomputed uniform noise.""" ++ return ( ++ self.enable_async_exponential ++ and not sampling_info.is_all_greedy ++ and sampling_info.sampling_seed is None ++ and not sampling_info.need_top_p_sampling ++ and not sampling_info.need_top_k_sampling ++ and not sampling_info.need_min_p_sampling ++ ) ++ ++ @torch.no_grad() ++ def prepare_async_exponential( ++ self, ++ batch_size: int, ++ vocab_size: int, ++ sampling_info: SamplingBatchInfo, ++ device: torch.device, ++ ) -> bool: ++ """Enqueue U(0,1) noise before model forward on a dedicated NPU stream. ++ ++ Reusing the buffer is safe because the side stream first waits for all ++ previously enqueued work on the current stream. The sampling path later ++ inserts a device-side event wait; it never synchronizes the CPU. ++ """ ++ if not self.can_prepare_async_exponential(sampling_info): ++ return False ++ ++ # A delayed sampler may still own the previous buffer. Do not overwrite it. ++ if self._async_exponential_pending: ++ logger.warning( ++ "Skip async exponential preparation because the previous batch " ++ "has not consumed its random buffer" ++ ) ++ return False ++ ++ if self._async_exponential_stream is None: ++ self._async_exponential_stream = torch.npu.Stream() ++ self._async_exponential_event = torch.npu.Event() ++ logger.info( ++ "Enabled asynchronous NPU exponential-race sampling " ++ "(uniform noise on side stream, stock-exponential argmax consumption)" ++ ) ++ ++ current_stream = torch.npu.current_stream() ++ self._async_exponential_stream.wait_stream(current_stream) ++ ++ expected_shape = (batch_size, vocab_size) ++ with torch.npu.stream(self._async_exponential_stream): ++ u = self._async_exponential_u ++ if ( ++ u is None ++ or tuple(u.shape) != expected_shape ++ or u.dtype != torch.float32 ++ or u.device.type != torch.device(device).type ++ ): ++ # SGLang converts next-token logits to FP32 before sampling; the ++ # noise grid must match CANN's fp32 uniform to preserve the ++ # incumbent sampling distribution. ++ u = torch.empty( ++ expected_shape, ++ dtype=torch.float32, ++ device=device, ++ ) ++ self._async_exponential_u = u ++ u.uniform_() ++ self._async_exponential_event.record() ++ ++ self._async_exponential_pending = True ++ return True ++ ++ def _sample_with_async_exponential( ++ self, probs: torch.Tensor ++ ) -> Optional[torch.Tensor]: ++ """Consume precomputed U(0,1) noise using the exponential-race identity. ++ ++ The full stock exponential_ transform (complement + clamp + -log) runs ++ here on the main stream (serial work that the sampling path owns ++ anyway), keeping the side stream free of AIV kernels so it cannot ++ slow down the model forward. ++ """ ++ if not self._async_exponential_pending: ++ return None ++ ++ u = self._async_exponential_u ++ self._async_exponential_pending = False ++ if ( ++ u is None ++ or u.shape != probs.shape ++ or u.dtype != probs.dtype ++ or u.device != probs.device ++ ): ++ logger.warning( ++ "Async uniform buffer does not match probs; falling back to " ++ "torch.multinomial (u=%s/%s/%s, probs=%s/%s/%s)", ++ None if u is None else tuple(u.shape), ++ None if u is None else u.dtype, ++ None if u is None else u.device, ++ tuple(probs.shape), ++ probs.dtype, ++ probs.device, ++ ) ++ return None ++ ++ current_stream = torch.npu.current_stream() ++ current_stream.wait_event(self._async_exponential_event) ++ u.record_stream(current_stream) ++ ++ # Do not modify probs in place: the standard backend reuses it for logprobs. ++ # argmax over probs / q with q ~ Exp(1): v1's production consumption ++ # form, kept because the stock argmax kernel is healthy while argmin at ++ # this shape is a slow legacy kernel (0.38ms vs ~0.13ms at bs=128). ++ # q is produced by transforming u in place with the exact stock ++ # aten::exponential_ values (NPU op-plugin composite, fp32 path): ++ # x = min(1 - u, 1 - eps/2); q = -log(x), eps = finfo(dtype).eps ++ # The guard is a single torch.minimum against a cached 0-dim bound ++ # tensor: min(x, 1-eps/2) is bitwise-identical to stock's ++ # ge+masked_fill_ pair (x <= 1 always), and unlike clamp_max_ it stays ++ # on the template-grade aclnnMinimum kernel -- clamp_max_ routes to ++ # the legacy-family ClipByValueV2 op with two scalar->device uploads ++ # per call, which costs more than the unary-template kernels used by ++ # the rest of the chain. The cap keeps x below 1 so q can never be 0: ++ # u == 0 maps to q ~= 5.96e-8 and scores huge, exactly as stock (v2.3 ++ # mapped it to +inf/+0.0 instead). For fp32 q stays within ++ # [~5.96e-8, ~16.6], finite and strictly positive, so the race has no ++ # inf/NaN edge. ++ # All transform ops run here on the main stream; the side stream stays ++ # pure DSA uniform with zero AIV kernels. ++ bound = self._async_exp_min_bound ++ if bound is None or bound.dtype != u.dtype or bound.device != u.device: ++ bound = torch.full( ++ (), ++ 1.0 - torch.finfo(u.dtype).eps / 2.0, ++ dtype=u.dtype, ++ device=u.device, ++ ) ++ self._async_exp_min_bound = bound ++ u.neg_().add_(1.0) ++ torch.minimum(u, bound, out=u) ++ u.log_().neg_() ++ sampled_index = torch.div(probs, u).argmax(dim=-1) ++ return sampled_index.view(-1).to(torch.int32) + + def _preprocess_logits( + self, logits: torch.Tensor, sampling_info: SamplingBatchInfo +@@ -115,6 +279,9 @@ class Sampler(nn.Module): + to get the unique seed for each position. + """ + logits = logits_output.next_token_logits ++ # In the plain probability path, keep the softmax output and apply log ++ # only to the values requested by the caller. ++ logprobs_are_probs = False + + # Preprocess logits (custom processors and NaN handling) + logits = self._preprocess_logits(logits, sampling_info) +@@ -180,9 +347,13 @@ class Sampler(nn.Module): + # Standard path: do softmax and sample from probs. + logits.div_(sampling_info.temperatures) + +- # In-place op to save memory +- logits[:] = torch.softmax(logits, dim=-1) +- probs = logits ++ # 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 ++ # read again on this path and is overwritten by the next ++ # forward anyway. Trade-off: one extra live [batch, vocab] ++ # fp32 tensor during sampling. (post_sample v1 的 A1 改动) ++ probs = torch.softmax(logits, dim=-1) + + batch_next_token_ids = self._sample_from_probs( + probs, sampling_info, positions, simple_sampling_case +@@ -191,8 +362,9 @@ class Sampler(nn.Module): + logprobs = ( + logprobs_via_logsoftmax_kernel + if logprobs_via_logsoftmax_kernel is not None +- else torch.log(probs) ++ else probs + ) ++ logprobs_are_probs = logprobs_via_logsoftmax_kernel is None + del probs + + # Attach logprobs to logits_output (in-place modification) +@@ -206,6 +378,7 @@ class Sampler(nn.Module): + token_ids_logprobs, + sampling_info, + batch_next_token_ids, ++ logprobs_are_probs, + ) + + self._sync_token_ids_across_tp(batch_next_token_ids, sampling_info) +@@ -225,11 +398,13 @@ class Sampler(nn.Module): + Handles both simple (direct multinomial) and complex (top-k/top-p/min-p) cases. + """ + if simple_sampling_case: +- batch_next_token_ids = sampling_from_probs_torch( +- probs, +- sampling_seed=sampling_info.sampling_seed, +- positions=positions, +- ) ++ batch_next_token_ids = self._sample_with_async_exponential(probs) ++ if batch_next_token_ids is None: ++ batch_next_token_ids = sampling_from_probs_torch( ++ probs, ++ sampling_seed=sampling_info.sampling_seed, ++ positions=positions, ++ ) + else: + backend = get_flags().sampling_backend + if backend == "flashinfer": +@@ -302,7 +477,11 @@ class Sampler(nn.Module): + probabilities, sampling_info.sampling_seed, positions + ).view(-1) + else: +- batch_next_token_ids = torch.multinomial(probs, num_samples=1).view(-1) ++ batch_next_token_ids = self._sample_with_async_exponential(probs) ++ if batch_next_token_ids is None: ++ batch_next_token_ids = torch.multinomial( ++ probs, num_samples=1 ++ ).view(-1) + return batch_next_token_ids.to(torch.int32) + else: + assert ( +@@ -353,9 +532,14 @@ class Sampler(nn.Module): + token_ids_logprobs: List[List[int]], + sampling_info: SamplingBatchInfo, + batch_next_token_ids: torch.Tensor, ++ logprobs_are_probs: bool, + ): +- # clamp to avoid -inf values +- logprobs.clamp_(min=torch.finfo(logprobs.dtype).min) ++ # Clamp the extracted values instead of the full [batch, vocab] ++ # matrix. clamp(min=const) is elementwise and monotone ++ # non-decreasing, so it commutes with topk/gather: clamping the ++ # small outputs gives identical results (-inf -> finfo.min) while ++ # skipping a full-matrix read+write pass (~0.4 ms/step on NPU). ++ clamp_min = torch.finfo(logprobs.dtype).min + + # Attach logprobs to logits_output (in-place modification) + if any(x > 0 for x in top_logprobs_nums): +@@ -363,6 +547,20 @@ class Sampler(nn.Module): + logits_output.next_token_top_logprobs_val, + logits_output.next_token_top_logprobs_idx, + ) = get_top_logprobs(logprobs, top_logprobs_nums, no_copy_to_cpu=True) ++ # Same extraction as get_top_logprobs, but clamp the ++ # [batch, max_k] topk result in a single kernel before ++ # slicing per request. ++ max_k = max(top_logprobs_nums) ++ top_vals, top_idx = logprobs.topk(max_k, dim=-1) ++ if logprobs_are_probs: ++ top_vals.log_() ++ top_vals.clamp_(min=clamp_min) ++ logits_output.next_token_top_logprobs_val = [ ++ top_vals[i][:k] for i, k in enumerate(top_logprobs_nums) ++ ] ++ logits_output.next_token_top_logprobs_idx = [ ++ top_idx[i][:k] for i, k in enumerate(top_logprobs_nums) ++ ] + + if any(x is not None for x in token_ids_logprobs): + ( +@@ -372,10 +570,19 @@ class Sampler(nn.Module): + logprobs, token_ids_logprobs, no_copy_to_cpu=True + ) + +- logits_output.next_token_logprobs = logprobs[ +- torch.arange(len(batch_next_token_ids), device=sampling_info.device), +- batch_next_token_ids, +- ] ++ for row in logits_output.token_ids_logprobs_val: ++ if torch.is_tensor(row): ++ if logprobs_are_probs: ++ row.log_() ++ row.clamp_(min=clamp_min) ++ ++ # Gather one value per row directly; this removes the temporary arange ++ # and the 2-D advanced-index operation from the hot path. ++ token_indices = batch_next_token_ids.to(dtype=torch.long).view(-1, 1) ++ next_token_logprobs = torch.gather(logprobs, dim=1, index=token_indices).view(-1) ++ if logprobs_are_probs: ++ next_token_logprobs.log_() ++ logits_output.next_token_logprobs = next_token_logprobs.clamp_(min=clamp_min) + + def _sync_token_ids_across_tp( + self, batch_next_token_ids: torch.Tensor, sampling_info: SamplingBatchInfo +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 ++++ b/python/sglang/srt/managers/tp_worker.py +@@ -486,6 +486,35 @@ class TpModelWorker(BaseTpWorker): + can_run_cuda_graph=can_run_cuda_graph, + ) + ++ def _prepare_async_exponential( ++ self, ++ forward_batch: ForwardBatch, ++ is_verify: bool, ++ ) -> None: ++ """Start simple-sampling RNG before forward so it can overlap on NPU.""" ++ if ( ++ is_verify ++ or self.enable_spec ++ or forward_batch.is_prefill_only ++ or not forward_batch.forward_mode.is_decode() ++ or forward_batch.sampling_info is None ++ or forward_batch.sampling_info.grammars is not None ++ ): ++ return ++ ++ sampler = self.model_runner.sampler ++ prepare = getattr(sampler, "prepare_async_exponential", None) ++ if prepare is None: ++ return ++ ++ prepare( ++ batch_size=forward_batch.batch_size, ++ vocab_size=self.model_runner.model_config.vocab_size, ++ sampling_info=forward_batch.sampling_info, ++ device=forward_batch.input_ids.device, ++ ) ++ ++ + def forward_batch_generation( + self, + batch: Optional[ScheduleBatch], +@@ -511,6 +540,7 @@ class TpModelWorker(BaseTpWorker): + return self._forward_batch_generation_dllm(forward_batch) + + if self.pp_group.is_last_rank: ++ self._prepare_async_exponential(forward_batch, is_verify) + out = self.model_runner.forward( + forward_batch, + pp_proxy_tensors=pp_proxy_tensors, +diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py +index 368781d8d5..0856caae48 100644 +--- a/python/sglang/srt/mem_cache/mamba_radix_cache.py ++++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py +@@ -1089,6 +1089,10 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): + if self.disable or len(key) == 0: + return None + ++ key = key.page_aligned(self.page_size) ++ if len(key) == 0: ++ return None ++ + return key + + def _match_post_processor( +@@ -1158,6 +1162,9 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): + ) + + def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode: ++ assert ( ++ 0 < split_len < len(child.key) ++ ), f"split_len must create non-empty nodes, {split_len=}, {len(child.key)=}" + # new_node -> child + new_node = TreeNode() + new_node.children = {key[split_len:].child_key(self.page_size): child} +@@ -1166,6 +1173,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): + new_node.full_lock_ref = child.full_lock_ref + new_node.mamba_lock_ref = 0 + new_node.key = child.key[:split_len] ++ assert len(new_node.key) > 0, f"new_node.key should not be empty" + new_node.value = child.value[:split_len].clone() + + # child time should be later than parent's time for mamba tombstone +@@ -1176,6 +1184,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): + self.mamba_lru_list.remove_node(child) + child.parent = new_node + child.key = child.key[split_len:] ++ assert len(child.key) > 0, f"child.key should not be empty" + child.value = child.value[split_len:].clone() + new_node.parent.children[key.child_key(self.page_size)] = new_node + new_node.hash_value, child.hash_value = split_node_hash_value( +diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py +index 031ab4fe87..93e14ff483 100644 +--- a/python/sglang/srt/model_executor/forward_batch_info.py ++++ b/python/sglang/srt/model_executor/forward_batch_info.py +@@ -49,7 +49,7 @@ from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import ( + ForwardBatchDeepSeekMHAMixin, + ) + from sglang.srt.model_executor.triton_ops.position import compute_position_triton +-from sglang.srt.runtime_context import get_parallel ++from sglang.srt.runtime_context import get_parallel, get_server_args + from sglang.srt.server_args import get_global_server_args + from sglang.srt.utils import ( + is_cuda, +@@ -1053,11 +1053,32 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): + return mrope_positions + + def _compute_mrope_positions(self, model_runner: ModelRunner, batch: ScheduleBatch): ++ mm_inputs = batch.multimodal_inputs ++ rl_on_policy_target = get_server_args().rl_on_policy_target ++ ++ if ( ++ self.spec_info is None ++ and batch.dllm_config is None ++ and ( ++ rl_on_policy_target is not None ++ or all(mm_input is None for mm_input in mm_inputs) ++ ) ++ ): ++ # Regular text generation does not need to rebuild mRoPE on the ++ # host. init_new has already produced the same flattened token ++ # positions on model_runner.device: clamp_position() for decode, ++ # or compute_position() for extend/mixed. Text mRoPE has identical ++ # temporal/height/width coordinates, so materialize the three rows ++ # directly and avoid the per-request host factories, cat, and H2D. ++ self.mrope_positions = ( ++ self.positions.to(dtype=torch.int64).unsqueeze(0).repeat(3, 1) ++ ) ++ return + # batch_size * [3 * seq_len] + batch_size = self.seq_lens_cpu.shape[0] + mrope_positions_list = [[]] * batch_size + for batch_idx in range(batch_size): +- mm_input = batch.multimodal_inputs[batch_idx] ++ mm_input = mm_inputs[batch_idx] + if self.forward_mode.is_decode(): + # 3 * N + if ( +diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py +index d5aac381dc..e710d83b05 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 ( + moe_tensor_model_parallel_all_reduce, + tensor_model_parallel_all_reduce, + ) ++from sglang.srt.environ import envs + from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder + from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation + from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +@@ -93,7 +94,10 @@ from sglang.srt.model_executor.cuda_graph_config import ( + from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors + from sglang.srt.model_executor.runner import get_is_capture_mode + from sglang.srt.model_loader.weight_utils import default_weight_loader +-from sglang.srt.runtime_context import get_flags, get_parallel ++from sglang.srt.runtime_context import ( ++ get_flags, ++ get_parallel, ++) + from sglang.srt.server_args import get_global_server_args + from sglang.srt.utils import ( + add_prefix, +@@ -106,14 +110,6 @@ from sglang.srt.utils import ( + make_layers, + use_intel_amx_backend, + ) +- +-if is_npu(): +- from sglang.srt.hardware_backend.npu.cmo import ( +- shared_expert_on_independent_stream, +- wait_share_stream, +- ) +- +-from sglang.srt.environ import envs + from sglang.srt.utils.hf_transformers_utils import get_rope_config + + _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() +@@ -446,9 +442,12 @@ class Qwen2MoeSparseMoeBlock(nn.Module): + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + if enable_dual_stream: +- shared_output = shared_expert_on_independent_stream( +- hidden_states.clone(), self._forward_shared_experts +- ) ++ current_stream = torch.npu.current_stream() ++ self.alt_stream.wait_stream(current_stream) ++ with torch.npu.stream(self.alt_stream): ++ shared_output = self._forward_shared_experts(hidden_states) ++ shared_output.record_stream(self.alt_stream) ++ shared_event = self.alt_stream.record_event() + else: + shared_output = self._forward_shared_experts(hidden_states) + topk_output = self.topk( +@@ -469,8 +468,8 @@ class Qwen2MoeSparseMoeBlock(nn.Module): + hidden_states=hidden_states, + topk_output=topk_output, + ) +- if enable_dual_stream: +- wait_share_stream() ++ if hidden_states.shape[0] > 0 and enable_dual_stream: ++ torch.npu.current_stream().wait_event(shared_event) + + if shared_output is not None: + final_hidden_states.add_(shared_output) +diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py +index fed78cf888..7b6505fa10 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 ( + + # Distributed + from sglang.srt.distributed import get_pp_group ++from sglang.srt.environ import envs + from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder + from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation + +@@ -608,7 +609,11 @@ class Qwen3_5LinearDecoderLayer(nn.Module): + quant_config=quant_config, + alt_stream=( + alt_stream +- if (_is_cuda or _disable_shared_experts_fusion()) ++ if ( ++ _is_cuda ++ or _disable_shared_experts_fusion() ++ or envs.SGLANG_NPU_USE_MULTI_STREAM.get() ++ ) + else None + ), + prefix=add_prefix("mlp", prefix.replace(".linear_attn", "")), +@@ -824,7 +829,11 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): + quant_config=quant_config, + alt_stream=( + alt_stream +- if (_is_cuda or _disable_shared_experts_fusion()) ++ if ( ++ _is_cuda ++ or _disable_shared_experts_fusion() ++ or envs.SGLANG_NPU_USE_MULTI_STREAM.get() ++ ) + else None + ), + prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), +@@ -1207,6 +1216,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 ++ alt_stream = ( ++ torch.cuda.Stream() ++ if _is_cuda or _hip_use_alt_stream or envs.SGLANG_NPU_USE_MULTI_STREAM.get() ++ else None ++ ) + + # Embedding layer + if self.pp_group.is_first_rank: +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 ++++ b/python/sglang/srt/speculative/eagle_utils.py +@@ -617,7 +617,7 @@ def eagle_sample( + + # Sample tokens + target_predict = None +- if sampling_info.is_all_greedy or _is_npu or _is_hip or _is_xpu: ++ if sampling_info.is_all_greedy or _is_hip or _is_xpu: + target_predict = torch.argmax(next_token_logits, dim=-1) + target_predict = target_predict.reshape(bs, verify_input.draft_token_num) + predict, accept_index, num_correct_drafts = verify_tree_greedy_func( +@@ -632,43 +632,80 @@ def eagle_sample( + topk=verify_input.tree_topk, + ) + else: +- from sgl_kernel import ( +- top_k_renorm_prob, +- top_p_renorm_prob, +- tree_speculative_sampling_target_only, +- ) +- +- from sglang.srt.speculative.reject_sampling import ( +- chain_speculative_sampling_triton, +- ) +- + use_rejection_sampling = ( + get_global_server_args().speculative_use_rejection_sampling + ) + ++ if _is_npu: ++ from sgl_kernel_npu.sample import ( ++ chain_speculative_sampling_rejection, ++ top_k_top_p_renorm_probs, ++ tree_speculative_sampling_target_only, ++ ) ++ ++ sampling_fn = ( ++ chain_speculative_sampling_rejection ++ if use_rejection_sampling ++ else tree_speculative_sampling_target_only ++ ) ++ else: ++ from sgl_kernel import ( ++ top_k_renorm_prob, ++ top_p_renorm_prob, ++ tree_speculative_sampling_target_only, ++ ) ++ ++ from sglang.srt.speculative.reject_sampling import ( ++ chain_speculative_sampling_triton, ++ ) ++ ++ sampling_fn = ( ++ chain_speculative_sampling_triton ++ if use_rejection_sampling ++ else tree_speculative_sampling_target_only ++ ) ++ + # Apply temperature and get target probs + expanded_temperature = torch.repeat_interleave( + sampling_info.temperatures, verify_input.draft_token_num, dim=0 + ) # (bs * num_draft_tokens, 1) + ++ sampling_logits = next_token_logits.float() if _is_npu else next_token_logits + target_probs = F.softmax( +- next_token_logits / expanded_temperature, dim=-1 ++ sampling_logits / expanded_temperature, dim=-1 + ) # (bs * num_draft_tokens, vocab_size) + maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax") +- target_probs = top_k_renorm_prob( +- target_probs, +- torch.repeat_interleave( +- sampling_info.top_ks, verify_input.draft_token_num, dim=0 +- ), +- ) # (bs * num_draft_tokens, vocab_size) +- maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm") +- target_probs = top_p_renorm_prob( +- target_probs, +- torch.repeat_interleave( +- sampling_info.top_ps, verify_input.draft_token_num, dim=0 +- ), +- ) +- maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm") ++ ++ if _is_npu: ++ target_probs = top_k_top_p_renorm_probs( ++ target_probs, ++ torch.repeat_interleave( ++ sampling_info.top_ks, verify_input.draft_token_num, dim=0 ++ ), ++ torch.repeat_interleave( ++ sampling_info.top_ps, verify_input.draft_token_num, dim=0 ++ ), ++ sampling_info.need_top_k_sampling, ++ sampling_info.need_top_p_sampling, ++ ) ++ maybe_detect_nan(target_probs, "v2 verify: target_probs after renorm") ++ else: ++ if sampling_info.need_top_k_sampling: ++ target_probs = top_k_renorm_prob( ++ target_probs, ++ torch.repeat_interleave( ++ sampling_info.top_ks, verify_input.draft_token_num, dim=0 ++ ), ++ ) # (bs * num_draft_tokens, vocab_size) ++ maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm") ++ if sampling_info.need_top_p_sampling: ++ target_probs = top_p_renorm_prob( ++ target_probs, ++ torch.repeat_interleave( ++ sampling_info.top_ps, verify_input.draft_token_num, dim=0 ++ ), ++ ) ++ maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm") + target_probs = target_probs.reshape(bs, verify_input.draft_token_num, -1) + draft_probs = ( + verify_input.draft_probs +@@ -687,16 +724,15 @@ def eagle_sample( + "does not produce one (draft_probs missing or vocab-mismatched)." + ) + ++ if _is_npu: ++ target_probs = target_probs.contiguous() ++ draft_probs = draft_probs.float().contiguous() ++ + # coins for rejection sampling + coins = torch.rand_like(candidates, dtype=torch.float32, device=device) + # coins for final sampling + coins_for_final_sampling = torch.rand((bs,), dtype=torch.float32, device=device) + +- sampling_fn = ( +- chain_speculative_sampling_triton +- if use_rejection_sampling +- else tree_speculative_sampling_target_only +- ) + sampling_fn( + predicts=predict, # mutable + accept_index=accept_index, # mutable +diff --git a/python/sglang/srt/speculative/triton_ops/cache_locs.py b/python/sglang/srt/speculative/triton_ops/cache_locs.py +index 663c238722..522fd5ee44 100644 +--- a/python/sglang/srt/speculative/triton_ops/cache_locs.py ++++ b/python/sglang/srt/speculative/triton_ops/cache_locs.py +@@ -363,6 +363,14 @@ def assign_extend_cache_locs_func( + return out_cache_loc + + elif _is_npu: ++ ''' ++ sgl-kernel-npu's cache_loc_assign / cache_loc_update operate under an explicit contract: ++ each row processes max_step tokens, and the tiling dimension of out_cache_loc is ++ cacheLocSize = batchSize * max_step, which is validated on the host side. ++ The host enforces 1 <= max_step <= MAX_STEP (16) and checks the size of out_cache_loc. ++ Therefore, here we allocate exactly batch_size * draft_token_num and pass draft_token_num ++ as max_step; no padding to 16 is needed. ++ ''' + out_cache_loc = torch.empty( + (batch_size * draft_token_num,), + dtype=torch.int32, +@@ -374,6 +382,7 @@ def assign_extend_cache_locs_func( + start_offset, + end_offset, + out_cache_loc, ++ draft_token_num, + ) + + return out_cache_loc diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py -index f68721f3d0..3e85ee7c10 100644 +index 82db90da3b..aabd222d7e 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py -@@ -46,6 +46,7 @@ import types - import uuid - import warnings - from collections import OrderedDict, defaultdict +@@ -73,6 +73,7 @@ from typing import ( + TypeVar, + Union, + ) +from relax.utils.device import is_npu_available - from contextlib import contextmanager - from dataclasses import dataclass - from decimal import Decimal -@@ -150,7 +151,7 @@ def is_npu() -> bool: + from unittest import SkipTest + from unittest.case import _ShouldStop + from urllib.parse import unquote, urlparse +@@ -176,7 +177,8 @@ def is_npu() -> bool: if not hasattr(torch, "npu"): return False - if not torch.npu.is_available(): ++ # if not torch.npu.is_available(): + if not is_npu_available: raise RuntimeError( "torch_npu detected, but NPU device is not available or visible." ) +diff --git a/python/sglang/srt/utils/torch_memory_saver_adapter.py b/python/sglang/srt/utils/torch_memory_saver_adapter.py +index ad98e59283..a5b2cdfd32 100644 +--- a/python/sglang/srt/utils/torch_memory_saver_adapter.py ++++ b/python/sglang/srt/utils/torch_memory_saver_adapter.py +@@ -41,7 +41,7 @@ class TorchMemorySaverAdapter(ABC): + def region(self, tag: str, enable_cpu_backup: bool = False): + raise NotImplementedError + +- def cuda_graph(self, **kwargs): ++ def cuda_graph(self, cuda_graph=None, **kwargs): + raise NotImplementedError + + def disable(self): +@@ -67,7 +67,10 @@ class _TorchMemorySaverAdapterReal(TorchMemorySaverAdapter): + def region(self, tag: str, enable_cpu_backup: bool = False): + return _memory_saver.region(tag=tag, enable_cpu_backup=enable_cpu_backup) + +- def cuda_graph(self, **kwargs): ++ def cuda_graph(self, cuda_graph=None, **kwargs): ++ if cuda_graph is not None: ++ kwargs["cuda_graph"] = cuda_graph ++ # kwargs.pop("auto_dispatch_capture", None) # torch_memory_saver 0.0.8 does not support this arg + return _memory_saver.cuda_graph(**kwargs) + + def disable(self): +@@ -94,7 +97,7 @@ class _TorchMemorySaverAdapterNoop(TorchMemorySaverAdapter): + yield + + @contextmanager +- def cuda_graph(self, **kwargs): ++ def cuda_graph(self, cuda_graph=None, **kwargs): + yield + + @contextmanager +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) From 25a42cc4d9156fd8618d0d042453554781002d8f Mon Sep 17 00:00:00 2001 From: wuqiwei Date: Fri, 21 Aug 2026 06:29:08 +0000 Subject: [PATCH 04/16] build(docker): upgrade NPU image to CANN 9.0.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🔩 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 --- docker/Dockerfile.npu | 104 ++++++++++++++++++++++++++++-------------- 1 file changed, 69 insertions(+), 35 deletions(-) diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 89ed93f0c..98184d709 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -2,7 +2,7 @@ ARG HTTP_PROXY ARG HTTPS_PROXY ARG NO_PROXY -FROM quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11 +FROM quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11 ARG HTTP_PROXY ARG HTTPS_PROXY @@ -43,86 +43,120 @@ RUN ARCH=$(uname -m) && \ export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \ fi && \ source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ - source /usr/local/Ascend/nnal/atb/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 + + # Setting pip & git config. Global config (set once, persists across subsequent RUN layers) ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" RUN pip config set global.index-url ${PIP_INDEX_URL} && \ git config --global http.sslverify false && \ + git config --global https.sslverify false && \ git config --global http.postBuffer 2147483648 && \ git config --global user.email "temp@example.com" && \ git config --global user.name "temp" -WORKDIR /root -COPY . /root/Relax +# install torch RUN pip install --upgrade pip packaging setuptools==80.10.2 && \ pip install torch==2.9.0 && \ pip install numpy==1.26.0 - # build torch_npu + +# build torch_npu RUN pip install pyyaml && \ git clone https://gitcode.com/Ascend/pytorch.git /root/pytorch && \ cd /root/pytorch && \ - git checkout v2.9.0-7.3.0 && \ + git checkout v26.0.1-pytorch2.9.0 && \ git cherry-pick -n f495de675bce38a2fa21edbf067b73d2a5f26733 && \ bash ci/build.sh --python=3.11 && \ - pip install dist/torch_npu-2.9.0*.whl -RUN cd /root && rm -rf /root/pytorch && \ - pip install triton-ascend==3.2.0 && \ + pip install dist/torch_npu*.whl + +# install triton/TQ +RUN pip install triton-ascend==3.2.1 --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple && \ pip install tensordict==0.10.0 pyvers==0.1.0 --no-deps && \ pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps - # Clone Megatron-LM, MindSpeed, MindSpeed-Bridge, Megatron-Bridge and install +# Clone Megatron-LM, MindSpeed, MindSpeed-Bridge, Megatron-Bridge and install RUN git clone https://gitcode.com/ascend/MindSpeed.git /root/MindSpeed && \ git clone https://github.com/NVIDIA/Megatron-LM.git /root/Megatron-LM && \ + git clone https://gitcode.com/ascend/MindSpeed-Ops.git /root/MindSpeed-Ops && \ git clone https://gitcode.com/ascend/MindSpeed-Bridge.git /root/MindSpeed-Bridge && \ git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git /root/Megatron-Bridge -RUN cd /root/MindSpeed && git checkout core_r0.16.0 && pip install -r requirements.txt && pip install -e . && \ +RUN cd /root/MindSpeed && git checkout core_r0.16.0 && pip install -r requirements.txt && pip install -e . && git checkout e4772499 && \ cd /root/Megatron-LM && git checkout core_v0.16.1 && pip install -e . --no-build-isolation && \ cd /root/Megatron-Bridge && git checkout v0.3.1 && \ - cd /root/MindSpeed-Bridge && git checkout 3655c07cbcc9 && pip install -r requirements.txt && bash tools/install_auto.sh + cd /root/MindSpeed-Ops/ && git checkout 33ac80f7 && pip install -e . --no-build-isolation --no-deps && \ + cd /root/MindSpeed-Bridge/ && git checkout v0.3.1 && pip install -r requirements.txt && pip install -e . --no-deps + -# Patch Megatron-LM, MindSpeed, MindSpeed-Bridge -RUN cd /root/MindSpeed && \ - patch -p1 < /root/Relax/docker/npu_patch/mindspeed.patch && \ - cd /root/Megatron-Bridge && \ - patch -p1 < /root/Relax/docker/npu_patch/megatron-bridge.patch && \ +COPY . /root/Relax +# Patch Megatron-LM, MindSpeed, MindSpeed-Bridge, MindSpeed-Ops +RUN cd /root/Megatron-Bridge && \ + patch -p1 < /root/Relax/docker/npu_patch/megatron-bridge.patch && \ + git add . && git commit -m "base line" && \ cd /root/Megatron-LM && \ - patch -p1 < /root/Relax/docker/npu_patch/megatron.patch && \ + patch -p1 < /root/Relax/docker/npu_patch/megatron.patch && \ + git add . && git commit -m "base line" && \ cd /root/MindSpeed-Bridge && \ - patch -p1 < /root/Relax/docker/npu_patch/mindspeed-bridge.patch + patch -p1 < /root/Relax/docker/npu_patch/mindspeed-bridge.patch && \ + git add . && git commit -m "base line" && \ + cd /root/MindSpeed-Ops && \ + patch -p1 < /root/Relax/docker/npu_patch/mindspeed-ops.patch && \ + git add . && git commit -m "base line" && \ + cd /root/MindSpeed && \ + patch -p1 < /root/Relax/docker/npu_patch/mindspeed.patch && \ + git add . && git commit -m "base line" -# Copy MindSpeed-Bridge, Megatron-Bridge into Megatron-LM -RUN cp -r /root/MindSpeed-Bridge/mindspeed_bridge /root/Megatron-LM/ && \ - cp -r /root/Megatron-Bridge/src/megatron/bridge /root/Megatron-LM/megatron/ && \ - cd /root && rm -rf /root/MindSpeed-Bridge && rm -rf /root/Megatron-Bridge # Install sglang RUN git clone https://github.com/sgl-project/sglang.git /root/sglang && \ - cd /root/sglang && git checkout v0.5.10 && \ + cd /root/sglang && git checkout v0.5.15 && \ mv python/pyproject.toml python/pyproject.toml.backup && \ mv python/pyproject_npu.toml python/pyproject.toml && \ pip install -e "python[srt_npu]" --constraint <(echo "torch==2.9.0") && \ - # [NPU] Fix Qwen3.5 inference acc. - git stash && \ - git fetch origin pull/23815/head:pr-23815 && \ - git checkout pr-23815 && \ - patch -p1 < /root/Relax/docker/npu_patch/sglang-npu.patch + # patch -p1 < /root/Relax/docker/npu_patch/sglang-npu.patch + 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" + # Install sgl-kernle-npu RUN git clone https://github.com/sgl-project/sgl-kernel-npu /root/sgl-kernel-npu && \ - cd /root/sgl-kernel-npu && git checkout 2026.04.15.rc3 && \ + cd /root/sgl-kernel-npu && git checkout 2026.7.2 && \ # Adapt tms for colocate train. - git cherry-pick -n 23519771d347 --no-gpg-sign && \ - patch -p1 < /root/Relax/docker/npu_patch/torch-memory-saver.patch && \ - bash build.sh -a kernels && bash build.sh -a memory-saver && \ + patch -p1 < /root/Relax/docker/npu_patch/sgl-kernel-npu.patch && \ + git add . && git commit -m "sgl-kernel-npu.patch" && \ + bash build.sh && \ pip install output/*.whl && \ - cd /root && rm -rf /root/sgl-kernel-npu + cd /root + +# 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 && \ + apt update && apt install 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 + + # Install Relax -RUN cd /root/Relax && pip install -e . +# git clone https://github.com/redai-infra/Relax.git +RUN cd /root/Relax && \ + pip install -e . + RUN pip install ray==2.55.1 && pip install protobuf==6.33.6 #Clean cache RUN pip cache purge && \ rm -rf /tmp/* + From 700b291ee17476ac6916b88c69f8272f6f346f8f Mon Sep 17 00:00:00 2001 From: Tgz27 <617796318@qq.com> Date: Mon, 24 Aug 2026 16:36:59 +0800 Subject: [PATCH 05/16] feat(npu): add Qwen3.5 MTP SFT/training scripts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh | 160 ++++++++++++++ .../run_qwen35-35B-A3B-16xnpu-colocate-thd.sh | 167 ++++++++++++++ .../text/run_qwen35_9B_mtp_8xnpu_thd.sh | 206 ++++++++++++++++++ 3 files changed, 533 insertions(+) create mode 100644 scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh create mode 100644 scripts/training/text/run_qwen35-35B-A3B-16xnpu-colocate-thd.sh create mode 100644 scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh diff --git a/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh b/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh new file mode 100644 index 000000000..e5415ad0c --- /dev/null +++ b/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh @@ -0,0 +1,160 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# Qwen3.5-35B-A3B MTP SFT on pokemon-gpt4o-captions, 8xNPU single-node, ray-submit launch. +# +# Usage: +# bash scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh + +set -ex +set -o pipefail + +now=$(date "+%Y-%m-%d-%H:%M:%S") +echo 当前时间: + +export ASCEND_COREDUMP_SIGNAL=none +export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export HCCL_HOST_SOCKET_PORT_RANGE=63000-63150 +export HCCL_NPU_SOCKET_PORT_RANGE=64000-64150 +export MASTER_ADDR=$(hostname -I | awk '{print $1}') + +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-35B-A3B.sh" + +PROJECT_NAME="${PROJECT_NAME:=Relax/sft/pokemon}" +EXP_NAME=qwen3.5-35B-A3B-mtp-sft-pokemon-gpu8 +EXP_DIR="${EXP_DIR:-/mnt/tidalfs-hwwh01/dataset/yuanhang/models}" +MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" +DATA_DIR="${DATA_DIR:=/mnt/tidalfs-hwwh01/dataset/yuanhang/datasets}" +TRAIN_FILES=( + "'${DATA_DIR}/sft/data/pokemon-gpt4o-captions/pokemon_gpt4o_en.parquet'" + "'${DATA_DIR}/sft/data/pokemon-gpt4o-captions/pokemon_gpt4o_zh.parquet'" +) +PROMPT_DATA="[$(IFS=,; echo "${TRAIN_FILES[*]}")]" +SAVE_DIR="${SAVE_DIR:=${EXP_DIR}/checkpoint/checkpoints/qwen3.5-35B-A3B-mtp-pokemon-sft-0821}" + +CKPT_ARGS=( + --hf-checkpoint ${EXP_DIR}/Qwen3.5-35B-A3B + --ref-load ${EXP_DIR}/Qwen3.5-35B-A3B + --megatron-to-hf-mode bridge + --save ${SAVE_DIR}/sft/${EXP_NAME} + # --load ${SAVE_DIR}/sft/${EXP_NAME} + --save-interval 100 + --num-epoch 10 +) + +SFT_ARGS=( + --loss-type sft + --prompt-data "${PROMPT_DATA}" + --input-key conversations + --multimodal-keys '{"image":"images"}' + --conversation-key-map '{"from":"role","value":"content","human":"user","gpt":"assistant"}' + --global-batch-size 64 + --use-dynamic-batch-size + --max-tokens-per-gpu 20480 + --balance-data + --per-rank-fetch + --sft-prefetch-num-workers 16 + --sft-prefetch-buffer-size 512 +) + +MTP_ARGS=( + --mtp-num-layers 1 + --enable-mtp-training + --mtp-loss-scaling-factor 0.2 + # --ci-test +) + +EVAL_ARGS=( + --eval-size 0.1 + --eval-interval 20 +) + +PREDICT_ARGS=( + # --sft-predict-interval 10 + # --eval-temperature 0.0 + # --eval-max-response-len 512 + # --rollout-num-gpus-per-engine 2 + # --sglang-mem-fraction-static 0.6 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 8 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer + --no-gradient-accumulation-fusion + + # --moe-flex-dispatcher-backend deepep + # --moe-token-dispatcher-type flex + --cross-entropy-loss-fusion + --no-rope-fusion + --sft-chunked-logits + --sft-logits-chunk-size ${SFT_LOGITS_CHUNK_SIZE:-2048} + + --colocate +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-5 + --lr-decay-style cosine + --min-lr 1e-6 + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --clip-grad 1.0 +) + +WANDB_ARGS=( + --use-clearml + --use-metrics-service + --use-tensorboard + --tb-project-name ${PROJECT_NAME} + --tb-experiment-name ${EXP_NAME}-${now} +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --use-health-check + --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 '{"sft": [1, 0], "actor": [1, 8]}' \ + --sft-max-in-flight-steps 4 \ + --num-data-storage-units 8 \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${SFT_ARGS[@]}" \ + "${MTP_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${PREDICT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${MISC_ARGS[@]}" 2>&1 | tee log/qwen3.5-35B-A3B-mtp-sft-pokemon-npu8-${now}.log diff --git a/scripts/training/text/run_qwen35-35B-A3B-16xnpu-colocate-thd.sh b/scripts/training/text/run_qwen35-35B-A3B-16xnpu-colocate-thd.sh new file mode 100644 index 000000000..c0dc81952 --- /dev/null +++ b/scripts/training/text/run_qwen35-35B-A3B-16xnpu-colocate-thd.sh @@ -0,0 +1,167 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# Qwen3.5-35B-A3B 16xNPU colocate training script. +# +# Usage: +# bash scripts/training/text/run_qwen35-35B-A3B-16xnpu-colocate-thd.sh + +set -ex +set -o pipefail +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY + +ulimit -n 65535 + +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-63150 +export HCCL_NPU_SOCKET_PORT_RANGE=64000-64150 +export TMS_HOOK_MODE="preload" +export HYDRA_FULL_ERROR=1 + + +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-35B-A3B.sh" +EXP_DIR="${EXP_DIR:-/mnt/tidalfs-hwwh01/dataset/yuanhang/models}" +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-0821 + --save ${EXP_DIR}/Qwen3.5-35B-A3B-save-0821 + --save-interval 100 + --max-actor-ckpt-to-keep 0 +) + +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 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 +) + +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 + --max-tokens-per-gpu 10240 + --no-rope-fusion + --no-gradient-accumulation-fusion +) + +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 +) + +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 +) + +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, 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-npu16-colocate-${now}.log diff --git a/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh b/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh new file mode 100644 index 000000000..5e7f7b94f --- /dev/null +++ b/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh @@ -0,0 +1,206 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# Qwen3.5-9B 4xNPU colocate (sync) GRPO + MTP joint-training script. +# +# Phase-1 RL MTP: trains the native MTP head jointly with the policy via an +# auxiliary loss (slime-style). Rollout keeps `enable_draft_weights_cpu_backup=True` +# so SGLang inference uses the base model only — no speculative decoding here. +# +# Requires the HF checkpoint to contain MTP weights (`num_nextn_predict_layers>=1`). +# +# Differences from the GPU MTP script: +# - Removes --cross-entropy-loss-fusion / --cross-entropy-fusion-impl te +# (TransformerEngine fused CE kernel is CUDA-only, not available on NPU). +# - Uses NPU-specific SGLang args (--sglang-device npu, ascend attention backend). +# - Uses NPU-specific perf args (--qkv-format bshd, --no-rope-fusion, etc.). +# +# Usage: +# bash scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh + +set -ex +set -o pipefail +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY +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 +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 + + + +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" +# Support setting env from outside +EXP_DIR="${EXP_DIR:-/mnt/tidalfs-hwwh01/dataset/yuanhang/models}" +MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" +DATA_DIR="${DATA_DIR:-${EXP_DIR}}" +PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math-mtp}" +NUM_ROLLOUT="${NUM_ROLLOUT:=200}" + + +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-mtp-save-0821 + --save ${EXP_DIR}/Qwen3.5-9B-mtp-save-0821 + --save-interval 100 + --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 + --balance-data + --use-fault-tolerance +) + +EVAL_ARGS=( + --log-passrate + --skip-eval-before-train + --eval-interval 50 + --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 1 +) + +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 1 + + --use-dynamic-batch-size + --qkv-format thd + --max-tokens-per-gpu 10240 + + --no-rope-fusion + --no-gradient-accumulation-fusion +) + +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 + + --custom-tis-function-path relax.backends.megatron.loss.icepop_function +) + +MTP_ARGS=( + --mtp-num-layers ${MTP_NUM_LAYERS:-1} + --enable-mtp-training + --mtp-loss-scaling-factor ${MTP_LOSS_SCALING_FACTOR:-0.1} + # NOTE: --cross-entropy-loss-fusion / --cross-entropy-fusion-impl te are + # intentionally omitted — the TE fused CE kernel is CUDA-only. Megatron + # will fall back to the non-fused cross-entropy path. +) + +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 +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 4 + --sglang-mem-fraction-static 0.6 + --sglang-cuda-graph-bs 1 2 4 8 16 24 32 48 64 + --sglang-max-running-requests 64 + --sglang-device npu + --sglang-chunked-prefill-size 8192 + --sglang-max-prefill-tokens 8192 + --sglang-enable-dp-attention + --sglang-enable-dp-lm-head + --sglang-attention-backend ascend + --sglang-max-mamba-cache-size 352 + --sglang-mamba-ssm-dtype bfloat16 + --sglang-mamba-scheduler-strategy extra_buffer + --sglang-speculative-algorithm NEXTN + --sglang-speculative-num-steps 2 + --sglang-speculative-eagle-topk 1 + --sglang-speculative-num-draft-tokens 3 + +) +WANDB_ARGS=( + --use-tensorboard + --use-metrics-service + --tb-project-name ${PROJECT_NAME} + --tb-experiment-name qwen35-9B-mtp-GRPO-4x-sync-${now} + # --use-wandb + # --wandb-project slime-dev + # --wandb-group qwen3-4B-test + # --wandb-key ${WANDB_KEY} +) + +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://127.0.0.1:8265" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"actor": [1, 8], "rollout": [1, 8]}' \ + --max-staleness 0 \ + --num-data-storage-units 1 \ + --colocate \ + --use-health-check \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${MTP_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" 2>&1 | tee log/qwen35-9B-MATH-npu16-colocate-mtp-${now}.log From bb7dda73430802ec2310330cc91cc9feb9d7f31e Mon Sep 17 00:00:00 2001 From: lixionglong Date: Mon, 24 Aug 2026 17:36:15 +0800 Subject: [PATCH 06/16] Revert "fix(megatron): attach _hf_config for MTP bridge when pp>1" This reverts commit 57809b43772529b339dc5ba2a96f25594a711a71. --- relax/backends/megatron/weight_update/bridge_converter.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/relax/backends/megatron/weight_update/bridge_converter.py b/relax/backends/megatron/weight_update/bridge_converter.py index 6eb2c9a0a..608eb0ce8 100644 --- a/relax/backends/megatron/weight_update/bridge_converter.py +++ b/relax/backends/megatron/weight_update/bridge_converter.py @@ -82,10 +82,7 @@ def init_tasks(self) -> None: if task.param_weight is not None: self._bridge_task_map[task.global_param_name] = task - model_bridge = bridge._model_bridge - if not hasattr(model_bridge, "_hf_config"): - model_bridge._hf_config = bridge.hf_pretrained.config - self._bridge_mapping_registry = model_bridge.mapping_registry() + self._bridge_mapping_registry = bridge._model_bridge.mapping_registry() mapping_registry = self._bridge_mapping_registry for name, _param in named_params_and_buffers(self._args, self._model): global_name = strip_param_name_prefix(name) From 41606123e6df390a9feca9c52e991e74454dc86d Mon Sep 17 00:00:00 2001 From: lixionglong Date: Mon, 24 Aug 2026 17:52:40 +0800 Subject: [PATCH 07/16] fix: update qwen35 mtp mapping in mindspeed-bridge.patch for relax bridge_converter.py --- docker/npu_patch/mindspeed-bridge.patch | 192 ++++++++++++++++++++++++ 1 file changed, 192 insertions(+) diff --git a/docker/npu_patch/mindspeed-bridge.patch b/docker/npu_patch/mindspeed-bridge.patch index f658a5981..a0e6fcfd3 100644 --- a/docker/npu_patch/mindspeed-bridge.patch +++ b/docker/npu_patch/mindspeed-bridge.patch @@ -122,3 +122,195 @@ index 5d03f85..2488ef7 100644 ) # apply text model config to vision model config +diff --git a/mindspeed_bridge/models/qwen_vl/qwen35_vl_bridge.py b/mindspeed_bridge/models/qwen_vl/qwen35_vl_bridge.py +index ebbd2a3..9accb3f 100644 +--- a/mindspeed_bridge/models/qwen_vl/qwen35_vl_bridge.py ++++ b/mindspeed_bridge/models/qwen_vl/qwen35_vl_bridge.py +@@ -458,62 +458,53 @@ class Qwen35VLMoEBridge(MegatronModelBridge): + # Megatron VL prefix: language_model.mtp.* + # HF prefix: mtp.* (top-level, not under model.language_model.) + # ================================================================= +- if not hasattr(self, "_hf_config"): +- logger.warning("No HF config found, skipping MTP mappings.") +- return MegatronMappingRegistry(*mapping_list) +- +- hf_config = self._hf_config +- num_mtp_layers = getattr(hf_config.text_config, "mtp_num_hidden_layers", None) +- +- if num_mtp_layers is not None: +- for mtp_layer in range(num_mtp_layers): +- mtp_param_mappings = { +- f"language_model.mtp.layers.{mtp_layer}.eh_proj.weight": "mtp.fc.weight", +- f"language_model.mtp.layers.{mtp_layer}.enorm.weight": "mtp.pre_fc_norm_embedding.weight", +- f"language_model.mtp.layers.{mtp_layer}.hnorm.weight": "mtp.pre_fc_norm_hidden.weight", +- f"language_model.mtp.layers.{mtp_layer}.final_layernorm.weight": "mtp.norm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.mlp.router.weight": f"mtp.layers.{mtp_layer}.mlp.gate.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.pre_mlp_layernorm.weight": f"mtp.layers.{mtp_layer}.post_attention_layernorm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_qkv.layer_norm_weight": f"mtp.layers.{mtp_layer}.input_layernorm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.self_attention.q_layernorm.weight": f"mtp.layers.{mtp_layer}.self_attn.q_norm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.self_attention.k_layernorm.weight": f"mtp.layers.{mtp_layer}.self_attn.k_norm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_proj.weight": f"mtp.layers.{mtp_layer}.self_attn.o_proj.weight", +- } +- for megatron_param, hf_param in mtp_param_mappings.items(): +- mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) +- +- mapping_list.extend( +- [ +- QKVMapping( +- megatron_param=f"language_model.mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_qkv.weight", +- q=f"mtp.layers.{mtp_layer}.self_attn.q_proj.weight", +- k=f"mtp.layers.{mtp_layer}.self_attn.k_proj.weight", +- v=f"mtp.layers.{mtp_layer}.self_attn.v_proj.weight", +- ), +- GatedMLPMapping( +- megatron_param=f"language_model.mtp.layers.{mtp_layer}.transformer_layer.mlp.experts.linear_fc1.weight*", +- gate=f"mtp.layers.{mtp_layer}.mlp.experts.*.gate_proj.weight", +- up=f"mtp.layers.{mtp_layer}.mlp.experts.*.up_proj.weight", +- ), +- AutoMapping( +- megatron_param=f"language_model.mtp.layers.{mtp_layer}.transformer_layer.mlp.experts.linear_fc2.weight*", +- hf_param=f"mtp.layers.{mtp_layer}.mlp.experts.*.down_proj.weight", +- ), +- GatedMLPMapping( +- megatron_param=f"language_model.mtp.layers.{mtp_layer}.transformer_layer.mlp.shared_experts.linear_fc1.weight", +- gate=f"mtp.layers.{mtp_layer}.mlp.shared_expert.gate_proj.weight", +- up=f"mtp.layers.{mtp_layer}.mlp.shared_expert.up_proj.weight", +- ), +- AutoMapping( +- megatron_param=f"language_model.mtp.layers.{mtp_layer}.transformer_layer.mlp.shared_experts.linear_fc2.weight", +- hf_param=f"mtp.layers.{mtp_layer}.mlp.shared_expert.down_proj.weight", +- ), +- ReplicatedMapping( +- megatron_param=f"language_model.mtp.layers.{mtp_layer}.transformer_layer.mlp.shared_experts.gate_weight", +- hf_param=f"mtp.layers.{mtp_layer}.mlp.shared_expert_gate.weight", +- ), +- ] +- ) ++ mtp_param_mappings = { ++ "language_model.mtp.layers.0.eh_proj.weight": "mtp.fc.weight", ++ "language_model.mtp.layers.0.enorm.weight": "mtp.pre_fc_norm_embedding.weight", ++ "language_model.mtp.layers.0.hnorm.weight": "mtp.pre_fc_norm_hidden.weight", ++ "language_model.mtp.layers.0.final_layernorm.weight": "mtp.norm.weight", ++ "language_model.mtp.layers.0.transformer_layer.mlp.router.weight": "mtp.layers.0.mlp.gate.weight", ++ "language_model.mtp.layers.0.transformer_layer.pre_mlp_layernorm.weight": "mtp.layers.0.post_attention_layernorm.weight", ++ "language_model.mtp.layers.0.transformer_layer.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight", ++ "language_model.mtp.layers.0.transformer_layer.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight", ++ "language_model.mtp.layers.0.transformer_layer.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight", ++ "language_model.mtp.layers.0.transformer_layer.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight", ++ } ++ for megatron_param, hf_param in mtp_param_mappings.items(): ++ mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) ++ ++ mapping_list.extend( ++ [ ++ QKVMapping( ++ megatron_param="language_model.mtp.layers.0.transformer_layer.self_attention.linear_qkv.weight", ++ q="mtp.layers.0.self_attn.q_proj.weight", ++ k="mtp.layers.0.self_attn.k_proj.weight", ++ v="mtp.layers.0.self_attn.v_proj.weight", ++ ), ++ GatedMLPMapping( ++ megatron_param="language_model.mtp.layers.0.transformer_layer.mlp.experts.linear_fc1.weight*", ++ gate="mtp.layers.0.mlp.experts.*.gate_proj.weight", ++ up="mtp.layers.0.mlp.experts.*.up_proj.weight", ++ ), ++ AutoMapping( ++ megatron_param="language_model.mtp.layers.0.transformer_layer.mlp.experts.linear_fc2.weight*", ++ hf_param="mtp.layers.0.mlp.experts.*.down_proj.weight", ++ ), ++ GatedMLPMapping( ++ megatron_param="language_model.mtp.layers.0.transformer_layer.mlp.shared_experts.linear_fc1.weight", ++ gate="mtp.layers.0.mlp.shared_expert.gate_proj.weight", ++ up="mtp.layers.0.mlp.shared_expert.up_proj.weight", ++ ), ++ AutoMapping( ++ megatron_param="language_model.mtp.layers.0.transformer_layer.mlp.shared_experts.linear_fc2.weight", ++ hf_param="mtp.layers.0.mlp.shared_expert.down_proj.weight", ++ ), ++ ReplicatedMapping( ++ megatron_param="language_model.mtp.layers.0.transformer_layer.mlp.shared_experts.gate_weight", ++ hf_param="mtp.layers.0.mlp.shared_expert_gate.weight", ++ ), ++ ] ++ ) + + return MegatronMappingRegistry(*mapping_list) + +@@ -769,46 +760,37 @@ class Qwen35VLBridge(MegatronModelBridge): + # Megatron VL prefix: language_model.mtp.* + # HF prefix: mtp.* (top-level, not under model.language_model.) + # ================================================================= +- if not hasattr(self, "_hf_config"): +- logger.warning("No HF config found, skipping MTP mappings.") +- return MegatronMappingRegistry(*mapping_list) +- +- hf_config = self._hf_config +- num_mtp_layers = getattr(hf_config.text_config, "mtp_num_hidden_layers", None) +- +- if num_mtp_layers is not None: +- for mtp_layer in range(num_mtp_layers): +- mtp_param_mappings = { +- f"language_model.mtp.layers.{mtp_layer}.eh_proj.weight": "mtp.fc.weight", +- f"language_model.mtp.layers.{mtp_layer}.enorm.weight": "mtp.pre_fc_norm_embedding.weight", +- f"language_model.mtp.layers.{mtp_layer}.hnorm.weight": "mtp.pre_fc_norm_hidden.weight", +- f"language_model.mtp.layers.{mtp_layer}.final_layernorm.weight": "mtp.norm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.mlp.linear_fc1.layer_norm_weight": f"mtp.layers.{mtp_layer}.post_attention_layernorm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.mlp.linear_fc2.weight": f"mtp.layers.{mtp_layer}.mlp.down_proj.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.pre_mlp_layernorm.weight": f"mtp.layers.{mtp_layer}.post_attention_layernorm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_qkv.layer_norm_weight": f"mtp.layers.{mtp_layer}.input_layernorm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.self_attention.q_layernorm.weight": f"mtp.layers.{mtp_layer}.self_attn.q_norm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.self_attention.k_layernorm.weight": f"mtp.layers.{mtp_layer}.self_attn.k_norm.weight", +- f"language_model.mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_proj.weight": f"mtp.layers.{mtp_layer}.self_attn.o_proj.weight", +- } +- for megatron_param, hf_param in mtp_param_mappings.items(): +- mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) +- +- mapping_list.extend( +- [ +- QKVMapping( +- megatron_param=f"language_model.mtp.layers.{mtp_layer}.transformer_layer.self_attention.linear_qkv.weight", +- q=f"mtp.layers.{mtp_layer}.self_attn.q_proj.weight", +- k=f"mtp.layers.{mtp_layer}.self_attn.k_proj.weight", +- v=f"mtp.layers.{mtp_layer}.self_attn.v_proj.weight", +- ), +- GatedMLPMapping( +- megatron_param=f"language_model.mtp.layers.{mtp_layer}.transformer_layer.mlp.linear_fc1.weight", +- gate=f"mtp.layers.{mtp_layer}.mlp.gate_proj.weight", +- up=f"mtp.layers.{mtp_layer}.mlp.up_proj.weight", +- ), +- ] +- ) ++ mtp_param_mappings = { ++ "language_model.mtp.layers.0.eh_proj.weight": "mtp.fc.weight", ++ "language_model.mtp.layers.0.enorm.weight": "mtp.pre_fc_norm_embedding.weight", ++ "language_model.mtp.layers.0.hnorm.weight": "mtp.pre_fc_norm_hidden.weight", ++ "language_model.mtp.layers.0.final_layernorm.weight": "mtp.norm.weight", ++ "language_model.mtp.layers.0.transformer_layer.mlp.linear_fc1.layer_norm_weight": "mtp.layers.0.post_attention_layernorm.weight", ++ "language_model.mtp.layers.0.transformer_layer.mlp.linear_fc2.weight": "mtp.layers.0.mlp.down_proj.weight", ++ "language_model.mtp.layers.0.transformer_layer.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight", ++ "language_model.mtp.layers.0.transformer_layer.pre_mlp_layernorm.weight": "mtp.layers.0.post_attention_layernorm.weight", ++ "language_model.mtp.layers.0.transformer_layer.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight", ++ "language_model.mtp.layers.0.transformer_layer.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight", ++ "language_model.mtp.layers.0.transformer_layer.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight", ++ } ++ for megatron_param, hf_param in mtp_param_mappings.items(): ++ mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) ++ ++ mapping_list.extend( ++ [ ++ QKVMapping( ++ megatron_param="language_model.mtp.layers.0.transformer_layer.self_attention.linear_qkv.weight", ++ q="mtp.layers.0.self_attn.q_proj.weight", ++ k="mtp.layers.0.self_attn.k_proj.weight", ++ v="mtp.layers.0.self_attn.v_proj.weight", ++ ), ++ GatedMLPMapping( ++ megatron_param="language_model.mtp.layers.0.transformer_layer.mlp.linear_fc1.weight", ++ gate="mtp.layers.0.mlp.gate_proj.weight", ++ up="mtp.layers.0.mlp.up_proj.weight", ++ ), ++ ] ++ ) + + return MegatronMappingRegistry(*mapping_list) + From 905857ade8e20bf5a4512cca476545ac5639d3cb Mon Sep 17 00:00:00 2001 From: dabuliu123 <270334047@qq.com> Date: Wed, 26 Aug 2026 16:35:08 +0800 Subject: [PATCH 08/16] fix(npu): rename Qwen3.5-9B CP script to 16xnpu MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- docker/npu-training.md | 10 +- docker/npu_patch/mindspeed-bridge.patch | 26 ++- relax/backends/megatron/actor.py | 2 +- .../training/text/run-qwen35-9B-16xnpu-cp.sh | 194 ++++++++++++++++++ 4 files changed, 220 insertions(+), 12 deletions(-) create mode 100755 scripts/training/text/run-qwen35-9B-16xnpu-cp.sh diff --git a/docker/npu-training.md b/docker/npu-training.md index d35892be1..f8909f6e4 100644 --- a/docker/npu-training.md +++ b/docker/npu-training.md @@ -12,6 +12,14 @@ | Qwen3.5-9B | DAPO | √ | √ | 910C 2卡 | `scripts/training/text/run-qwen35-9B-4xnpu-colocate.sh` | | Qwen3.5-35B-A3B | DAPO | √ | √ | 910C 8卡 | `scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh` | +## 特性支持 + +| 模型 | 训练场景 | 多模态 | MTP | CP | 训练所需最小卡数 | 参考脚本 | +| --------------- | -------- | ------ | --- | --- | ---------------- | -------------------------------------------------------------- | +| Qwen3.5-9B | DAPO | - | - | √ | 910C 8卡 | `scripts/training/text/run-qwen35-9B-16xnpu-cp.sh` | +| Qwen3.5-9B | DAPO | - | √ | - | 910C 4卡 | `scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh` | +| Qwen3.5-35B-A3B | SFT | √ | √ | - | 910C 4卡 | `scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh` | + ## 环境准备 ### 前置准备 @@ -120,4 +128,4 @@ bash scripts/training/text/run-qwen3-4B-8xgpu-async-npu.sh ## 下一步 -- [ ] 特性支持：多模态、CP长序列、MTP 等 +- [ ] 性能优化：Qwen3.5-35B-A3B 等 diff --git a/docker/npu_patch/mindspeed-bridge.patch b/docker/npu_patch/mindspeed-bridge.patch index a0e6fcfd3..6d574e221 100644 --- a/docker/npu_patch/mindspeed-bridge.patch +++ b/docker/npu_patch/mindspeed-bridge.patch @@ -19,7 +19,7 @@ index 99ccba5..8fa07a1 100644 # Set the cleaned value (or None if it was removed) setattr(cleaned_obj, attr_name, cleaned_value) diff --git a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py -index 19a3b9d..3dd5e25 100644 +index 19a3b9d..00a3a80 100644 --- a/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py +++ b/mindspeed_bridge/models/qwen_vl/modelling_qwen3_vl/gated_delta_net.py @@ -1,6 +1,6 @@ @@ -30,22 +30,28 @@ index 19a3b9d..3dd5e25 100644 import torch from torch import nn import torch.nn.functional as F -@@ -37,12 +37,11 @@ from megatron.core.ssm.gated_delta_net import ( +@@ -36,13 +36,15 @@ from megatron.core.ssm.gated_delta_net import ( + _split_tensor_factory, ) - try: +-try: - from causal_conv1d import causal_conv1d -+ import fla_npu -+ from fla.modules.convolution import causal_conv1d - except ImportError: +-except ImportError: ++if os.environ.get("MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV", "0") == "1": causal_conv1d = None - causal_conv1d_update = None ++else: ++ try: ++ import fla_npu ++ from fla.modules.convolution import causal_conv1d ++ except ImportError: ++ causal_conv1d = None -from mindspeed_ops.api.triton.l2norm import l2norm from mindspeed_bridge.models.qwen_vl.modelling_qwen3_vl.chunk_gated_delta_rule import ( torch_chunk_gated_delta_rule, ) -@@ -59,6 +58,12 @@ except ImportError: +@@ -59,6 +61,12 @@ except ImportError: from mindspeed_bridge.models.qwen_vl.modelling_qwen3_vl.flash_gated_delta_rule import flash_gated_delta_rule @@ -58,7 +64,7 @@ index 19a3b9d..3dd5e25 100644 class GatedDeltaNet(MegatronModule): """Gated Delta Net (GDN) layer class -@@ -191,9 +196,10 @@ class GatedDeltaNet(MegatronModule): +@@ -191,9 +199,10 @@ class GatedDeltaNet(MegatronModule): ) setattr(self.A_log, "tensor_model_parallel", True) @@ -71,7 +77,7 @@ index 19a3b9d..3dd5e25 100644 self.gated_delta_rule = flash_gated_delta_rule else: self.gated_delta_rule = torch_chunk_gated_delta_rule -@@ -474,7 +480,7 @@ class GatedDeltaNet(MegatronModule): +@@ -474,7 +483,7 @@ class GatedDeltaNet(MegatronModule): beta=beta, initial_state=None, output_final_state=False, @@ -80,7 +86,7 @@ index 19a3b9d..3dd5e25 100644 cu_seqlens=cu_seqlens_q, ) nvtx_range_pop(suffix="gated_delta_rule") -@@ -555,7 +561,7 @@ class GatedDeltaNet(MegatronModule): +@@ -555,7 +564,7 @@ class GatedDeltaNet(MegatronModule): # Apply L2 norm to query and key if self.use_qk_l2norm: diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..72d230001 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -184,9 +184,9 @@ def _init( self.genrm_manager = None - init(args) if repatch is not None: repatch(args) + init(args) tq.init(args.tq_config) self.data_system_client = tq.get_client() if is_megatron_main_rank(): diff --git a/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh b/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh new file mode 100755 index 000000000..4315ee447 --- /dev/null +++ b/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh @@ -0,0 +1,194 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# Qwen3.5-9B 16xNPU colocate (sync) training script for DAPO math dataset. +# + +set -ex +set -o pipefail +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY +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 +export LD_LIBRARY_PATH=/usr/local/Ascend/cann-9.0.0/opp/vendors/custom_transformer/op_api/lib/:${LD_LIBRARY_PATH} +export MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV=1 + +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 + +# Forward WANDB_API_KEY into Ray workers' runtime_env (local.sh doesn't +# propagate arbitrary env vars). No-op when the key isn't exported. +if [ -n "${WANDB_API_KEY:-}" ]; then + export RUNTIME_ENV_JSON=$(echo "$RUNTIME_ENV_JSON" | jq --arg k "$WANDB_API_KEY" '.env_vars.WANDB_API_KEY = $k') +fi +source "${MODEL_CONFIG_DIR}/qwen35-9B.sh" + +PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" +EXP_DIR="${EXP_DIR:-/mnt/}" +MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" +DATA_DIR="${DATA_DIR:-${EXP_DIR}}" + +NUM_ROLLOUT="${NUM_ROLLOUT:=1000}" + +CKPT_ARGS=( + --hf-checkpoint ${MODEL_DIR}/models/Qwen3.5-9B + --ref-load ${MODEL_DIR}/models/Qwen3.5-9B + --megatron-to-hf-mode bridge +# --warm-hf-checkpoint-page-cache + +# --load ${EXP_DIR}/Qwen3-9B_mcore_8xgpu/ + --save ${EXP_DIR}/Qwen3-9B_mcore_8xgpu/ + --save-interval 50 + --max-actor-ckpt-to-keep 0 +) + +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 + --balance-data + --use-fault-tolerance +) + +EVAL_ARGS=( + --log-passrate + --skip-eval-before-train + --eval-interval 20 + --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 4 + --calculate-per-token-loss + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity selective + # --recompute-method uniform + # --recompute-num-layers 1 + + --use-distributed-optimizer + --use-dynamic-batch-size + + --data-pad-size-multiplier 4096 + --max-tokens-per-gpu 8192 + --log-probs-max-tokens-per-gpu 8192 + --no-rope-fusion + --no-gradient-accumulation-fusion +) + +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 + # icepop: drop tokens with ratio outside [tis-clip-low, tis-clip] instead of clamping (vanilla TIS). + --custom-tis-function-path relax.backends.megatron.loss.icepop_function +) + +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 +) + +WANDB_ARGS=( + --use-tensorboard + --use-metrics-service + --tb-project-name ${PROJECT_NAME} + --tb-experiment-name qwen35-9B-8x-${now} +) +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 4 + --sglang-mem-fraction-static 0.8 + --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 +) + +# wandb: only enabled when WANDB_API_KEY is exported (see runtime_env injection above). +# wandb project names cannot contain / \ # ? % : — translate slashes to dashes. +if [ -n "${WANDB_API_KEY:-}" ]; then + WANDB_ARGS+=( + --use-wandb + --wandb-project ${PROJECT_NAME//\//-} + --wandb-group qwen35-9B-8x-${now} + ) +fi + +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://127.0.0.1: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]}' \ + --num-gpus-per-node 16 \ + --max-staleness 0 \ + --num-data-storage-units 1 \ + --colocate \ + --use-health-check \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" 2>&1 | tee log/qwen35-9B-GRPO-npu16-${now}.log From f49de709f3d026ec02198bf2b91684b60e05ce80 Mon Sep 17 00:00:00 2001 From: dabuliu123 <270334047@qq.com> Date: Wed, 26 Aug 2026 11:26:59 +0000 Subject: [PATCH 09/16] fix(gitleaks): drop /mnt/ defaults in NPU scripts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🔒 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 --- docker/Dockerfile.npu | 1 - scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh | 4 ++-- scripts/training/text/run-qwen35-9B-16xnpu-cp.sh | 2 +- .../training/text/run_qwen35-35B-A3B-16xnpu-colocate-thd.sh | 2 +- scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh | 2 +- 5 files changed, 5 insertions(+), 6 deletions(-) diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 85256d74d..ac1e37d42 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -160,4 +160,3 @@ RUN pip install ray==2.55.1 && pip install protobuf==6.33.6 #Clean cache RUN pip cache purge && \ rm -rf /tmp/* - diff --git a/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh b/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh index e5415ad0c..9cc63821b 100644 --- a/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh +++ b/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh @@ -28,9 +28,9 @@ source "${MODEL_CONFIG_DIR}/qwen35-35B-A3B.sh" PROJECT_NAME="${PROJECT_NAME:=Relax/sft/pokemon}" EXP_NAME=qwen3.5-35B-A3B-mtp-sft-pokemon-gpu8 -EXP_DIR="${EXP_DIR:-/mnt/tidalfs-hwwh01/dataset/yuanhang/models}" +EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" -DATA_DIR="${DATA_DIR:=/mnt/tidalfs-hwwh01/dataset/yuanhang/datasets}" +DATA_DIR="${DATA_DIR:-${EXP_DIR}}" TRAIN_FILES=( "'${DATA_DIR}/sft/data/pokemon-gpt4o-captions/pokemon_gpt4o_en.parquet'" "'${DATA_DIR}/sft/data/pokemon-gpt4o-captions/pokemon_gpt4o_zh.parquet'" diff --git a/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh b/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh index 4315ee447..f66519f45 100755 --- a/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh +++ b/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh @@ -33,7 +33,7 @@ fi source "${MODEL_CONFIG_DIR}/qwen35-9B.sh" PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" -EXP_DIR="${EXP_DIR:-/mnt/}" +EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" DATA_DIR="${DATA_DIR:-${EXP_DIR}}" diff --git a/scripts/training/text/run_qwen35-35B-A3B-16xnpu-colocate-thd.sh b/scripts/training/text/run_qwen35-35B-A3B-16xnpu-colocate-thd.sh index c0dc81952..4e899014b 100644 --- a/scripts/training/text/run_qwen35-35B-A3B-16xnpu-colocate-thd.sh +++ b/scripts/training/text/run_qwen35-35B-A3B-16xnpu-colocate-thd.sh @@ -36,7 +36,7 @@ if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then source "${SCRIPT_DIR}/../../entrypoint/local-npu.sh" fi source "${MODEL_CONFIG_DIR}/qwen35-35B-A3B.sh" -EXP_DIR="${EXP_DIR:-/mnt/tidalfs-hwwh01/dataset/yuanhang/models}" +EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" NUM_ROLLOUT="${NUM_ROLLOUT:=3000}" diff --git a/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh b/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh index 5e7f7b94f..4f031f57b 100644 --- a/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh +++ b/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh @@ -40,7 +40,7 @@ if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then fi source "${MODEL_CONFIG_DIR}/qwen35-9B.sh" # Support setting env from outside -EXP_DIR="${EXP_DIR:-/mnt/tidalfs-hwwh01/dataset/yuanhang/models}" +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-mtp}" From abf8a43727e404e8493ec17838f1e2983e44a85e Mon Sep 17 00:00:00 2001 From: lixionglong Date: Tue, 1 Sep 2026 15:43:47 +0800 Subject: [PATCH 10/16] fix: enforce MTP_NUM_LAYERS=1 for Qwen3.5 NPU scripts add validation in the *sft.sh script to ensure that MTP_NUM_LAYERS can only be 1 --- docker/npu-training.md | 4 ++++ .../training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh | 9 +++++++-- scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh | 5 +++++ 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/docker/npu-training.md b/docker/npu-training.md index f8909f6e4..bc4710071 100644 --- a/docker/npu-training.md +++ b/docker/npu-training.md @@ -126,6 +126,10 @@ bash scripts/training/text/run-qwen3-4B-8xgpu-async-npu.sh > MISC_ARGS，显示启用FlashAttention实现 `--use-flash-attn` +### MTP 特性说明 + +开启 MTP 训练时，`MTP_NUM_LAYERS`（MTP 头层数，脚本透传给训练参数 `--mtp-num-layers`）只能为 `1`：由于 Qwen3.5 原始 checkpoint 中仅包含 1 层 MTP 权重（`mtp_num_hidden_layers=1`），MTP 相关启动脚本（`scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh`、`scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh`）已加入校验，当该参数被设置为非 `1` 的值时脚本会报错退出 + ## 下一步 - [ ] 性能优化：Qwen3.5-35B-A3B 等 diff --git a/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh b/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh index 9cc63821b..308a8607c 100644 --- a/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh +++ b/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh @@ -63,10 +63,15 @@ SFT_ARGS=( --sft-prefetch-buffer-size 512 ) +if [[ "${MTP_NUM_LAYERS:-1}" != "1" ]]; then + echo "ERROR: MTP_NUM_LAYERS must be 1 for Qwen3.5 (checkpoint has mtp_num_hidden_layers=1)." >&2 + exit 1 +fi + MTP_ARGS=( - --mtp-num-layers 1 + --mtp-num-layers ${MTP_NUM_LAYERS:-1} --enable-mtp-training - --mtp-loss-scaling-factor 0.2 + --mtp-loss-scaling-factor ${MTP_LOSS_SCALING_FACTOR:-0.2} # --ci-test ) diff --git a/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh b/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh index 4f031f57b..dbf4142de 100644 --- a/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh +++ b/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh @@ -119,6 +119,11 @@ GRPO_ARGS=( --custom-tis-function-path relax.backends.megatron.loss.icepop_function ) +if [[ "${MTP_NUM_LAYERS:-1}" != "1" ]]; then + echo "ERROR: MTP_NUM_LAYERS must be 1 for Qwen3.5 (checkpoint has mtp_num_hidden_layers=1)." >&2 + exit 1 +fi + MTP_ARGS=( --mtp-num-layers ${MTP_NUM_LAYERS:-1} --enable-mtp-training From ebfc3648d773b30f982646d58263581438745555 Mon Sep 17 00:00:00 2001 From: dabuliu123 <270334047@qq.com> Date: Tue, 1 Sep 2026 17:53:58 +0800 Subject: [PATCH 11/16] refactor(npu): drop wandb wiring from CP script MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fix(docker): install gawk non-interactively # 🐛 Bug Fix ## Make gawk install non-interactive in NPU Dockerfile - Replace `apt install gawk` with `apt-get install -y --no-install-recommends gawk` - `apt install` prompts for confirmation and aborts in unattended Docker builds refactor(npu): drop wandb wiring from CP script # ♻️ Refactor ## Drop script-local wandb wiring from Qwen3.5-9B CP script - Remove WANDB_API_KEY injection into Ray runtime_env via jq - Remove WANDB_API_KEY-gated --use-wandb/--wandb-project/--wandb-group args - Wandb is configured centrally via --use-wandb/--wandb-key CLI args --- docker/Dockerfile.npu | 2 +- scripts/training/text/run-qwen35-9B-16xnpu-cp.sh | 15 --------------- 2 files changed, 1 insertion(+), 16 deletions(-) diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index ac1e37d42..ca4619d3e 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -135,7 +135,7 @@ 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 && \ - apt update && apt install gawk && \ + 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包 diff --git a/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh b/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh index f66519f45..6db7e65e9 100755 --- a/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh +++ b/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh @@ -25,11 +25,6 @@ if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then source "${SCRIPT_DIR}/../../entrypoint/local-npu.sh" fi -# Forward WANDB_API_KEY into Ray workers' runtime_env (local.sh doesn't -# propagate arbitrary env vars). No-op when the key isn't exported. -if [ -n "${WANDB_API_KEY:-}" ]; then - export RUNTIME_ENV_JSON=$(echo "$RUNTIME_ENV_JSON" | jq --arg k "$WANDB_API_KEY" '.env_vars.WANDB_API_KEY = $k') -fi source "${MODEL_CONFIG_DIR}/qwen35-9B.sh" PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" @@ -149,16 +144,6 @@ SGLANG_ARGS=( --sglang-attention-backend ascend ) -# wandb: only enabled when WANDB_API_KEY is exported (see runtime_env injection above). -# wandb project names cannot contain / \ # ? % : — translate slashes to dashes. -if [ -n "${WANDB_API_KEY:-}" ]; then - WANDB_ARGS+=( - --use-wandb - --wandb-project ${PROJECT_NAME//\//-} - --wandb-group qwen35-9B-8x-${now} - ) -fi - MISC_ARGS=( # default dropout in megatron is 0.1 --attention-dropout 0.0 From cda7305b739ac8095db80a93d8291d258a1abbcc Mon Sep 17 00:00:00 2001 From: Tgz27 <617796318@qq.com> Date: Thu, 3 Sep 2026 14:40:43 +0800 Subject: [PATCH 12/16] fix(sglang-npu.patch): conv weight cache invalidation, logprobs field name, qwen2_moe NPU alt_stream, duplicate topk fix(sglang-npu.patch): conv weight cache invalidation, logprobs field name, qwen2_moe NPU alt_stream, duplicate topk --- docker/npu_patch/sglang-npu.patch | 201 ++++++++++++++++++++++++++++-- 1 file changed, 191 insertions(+), 10 deletions(-) diff --git a/docker/npu_patch/sglang-npu.patch b/docker/npu_patch/sglang-npu.patch index 52cbe46e0..13d067766 100644 --- a/docker/npu_patch/sglang-npu.patch +++ b/docker/npu_patch/sglang-npu.patch @@ -36,21 +36,23 @@ index 7f506e7f1a..204b9382e8 100644 class AscendGDNAttnBackend(AscendMambaAttnBackendBase): -@@ -109,6 +102,13 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): +@@ -109,6 +102,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: + w = getattr(layer, "_conv_weights_t", None) -+ if w is None: ++ version = layer.conv_weights._version ++ if w is None or version != getattr(layer, "_conv_weights_t_version", None): + w = layer.conv_weights.transpose(0, 1).contiguous() + layer._conv_weights_t = w ++ layer._conv_weights_t_version = version + return w + def forward_decode( self, layer: RadixLinearAttention, -@@ -125,16 +125,17 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): +@@ -125,16 +127,17 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): cache_indices = self.forward_metadata.mamba_cache_indices assert isinstance(mixed_qkv, torch.Tensor) @@ -473,13 +475,14 @@ index e83e4157fa..19ed3f244b 100644 # Attach logprobs to logits_output (in-place modification) if any(x > 0 for x in top_logprobs_nums): -@@ -363,6 +547,20 @@ class Sampler(nn.Module): - logits_output.next_token_top_logprobs_val, - logits_output.next_token_top_logprobs_idx, - ) = get_top_logprobs(logprobs, top_logprobs_nums, no_copy_to_cpu=True) +@@ -362,7 +546,17 @@ class Sampler(nn.Module): +- ( +- logits_output.next_token_top_logprobs_val, +- logits_output.next_token_top_logprobs_idx, +- ) = get_top_logprobs(logprobs, top_logprobs_nums, no_copy_to_cpu=True) + # Same extraction as get_top_logprobs, but clamp the + # [batch, max_k] topk result in a single kernel before -+ # slicing per request. ++ # slicing per request. Runs topk exactly once. + max_k = max(top_logprobs_nums) + top_vals, top_idx = logprobs.topk(max_k, dim=-1) + if logprobs_are_probs: @@ -494,7 +497,7 @@ index e83e4157fa..19ed3f244b 100644 if any(x is not None for x in token_ids_logprobs): ( -@@ -372,10 +570,19 @@ class Sampler(nn.Module): +@@ -372,10 +567,19 @@ class Sampler(nn.Module): logprobs, token_ids_logprobs, no_copy_to_cpu=True ) @@ -502,7 +505,7 @@ index e83e4157fa..19ed3f244b 100644 - torch.arange(len(batch_next_token_ids), device=sampling_info.device), - batch_next_token_ids, - ] -+ for row in logits_output.token_ids_logprobs_val: ++ for row in logits_output.next_token_token_ids_logprobs_val: + if torch.is_tensor(row): + if logprobs_are_probs: + row.log_() @@ -720,6 +723,22 @@ index d5aac381dc..e710d83b05 100644 if shared_output is not None: final_hidden_states.add_(shared_output) +@@ -991,7 +990,14 @@ class Qwen2MoeForCausalLM(nn.Module): + self.pp_group = get_pp_group() + self.config = config + self.quant_config = quant_config +- alt_stream = torch.cuda.Stream() if _is_cuda else None ++ # Mirror qwen3_5.py: on NPU, SGLANG_NPU_USE_MULTI_STREAM also needs a ++ # real stream — _forward_deepep's dual-stream branch dereferences ++ # self.alt_stream unconditionally once enabled. ++ alt_stream = ( ++ torch.cuda.Stream() ++ if _is_cuda or (is_npu() and envs.SGLANG_NPU_USE_MULTI_STREAM.get()) ++ else None ++ ) + self.model = Qwen2MoeModel( + config, + quant_config, diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index fed78cf888..7b6505fa10 100644 --- a/python/sglang/srt/models/qwen3_5.py @@ -1056,3 +1075,165 @@ index 9ad577b42e..390a3721c5 100755 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 +--- /dev/null ++++ b/test/registered/unit/layers/test_sampler_attach_logprobs.py +@@ -0,0 +1,156 @@ ++"""Unit tests for Sampler._attach_logprobs_to_output. ++ ++Regression coverage for two bugs in the decode logprob attach path: ++ - iterating ``logits_output.token_ids_logprobs_val`` (a nonexistent field on ++ LogitsProcessorOutput) instead of ``next_token_token_ids_logprobs_val`` ++ raised AttributeError for any request with specific token-ID logprobs; ++ - the top-logprob branch ran a full-vocab topk twice (get_top_logprobs plus ++ an immediate re-topk that overwrote the first result). ++ ++The method under test does not touch ``self``, so the Sampler is allocated ++with ``__new__`` to skip the distributed-group init in ``__init__``. ++""" ++ ++import unittest ++ ++import torch ++ ++from sglang.test.ci.ci_register import register_cpu_ci ++from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel ++ ++maybe_stub_sgl_kernel() ++ ++from sglang.srt.layers.logits_processor import LogitsProcessorOutput ++from sglang.srt.layers.sampler import Sampler ++ ++register_cpu_ci(est_time=10, suite="base-a-test-cpu") ++ ++_VOCAB = 1000 ++ ++ ++def _make_sampler() -> Sampler: ++ # _attach_logprobs_to_output never reads instance state; bypass __init__ ++ # (which requires an initialized torch distributed group). ++ return Sampler.__new__(Sampler) ++ ++ ++def _attach( ++ logprobs: torch.Tensor, ++ top_logprobs_nums, ++ token_ids_logprobs, ++ batch_next_token_ids: torch.Tensor, ++ logprobs_are_probs: bool = False, ++) -> LogitsProcessorOutput: ++ logits_output = LogitsProcessorOutput(next_token_logits=None) ++ _make_sampler()._attach_logprobs_to_output( ++ logits_output=logits_output, ++ logprobs=logprobs, ++ top_logprobs_nums=top_logprobs_nums, ++ token_ids_logprobs=token_ids_logprobs, ++ sampling_info=None, ++ batch_next_token_ids=batch_next_token_ids, ++ logprobs_are_probs=logprobs_are_probs, ++ ) ++ return logits_output ++ ++ ++class TestAttachTokenIdsLogprobs(CustomTestCase): ++ ++ def test_token_ids_logprobs_no_attribute_error(self): ++ """Specific token-ID logprob requests must not raise AttributeError.""" ++ logprobs = torch.log_softmax(torch.randn(2, _VOCAB), dim=-1) ++ out = _attach( ++ logprobs=logprobs.clone(), ++ top_logprobs_nums=[0, 0], ++ token_ids_logprobs=[[5, 7], [42]], ++ batch_next_token_ids=torch.tensor([3, 4], dtype=torch.int32), ++ ) ++ self.assertIsNotNone(out.next_token_token_ids_logprobs_val) ++ self.assertIsNotNone(out.next_token_token_ids_logprobs_idx) ++ ++ def test_token_ids_logprobs_values(self): ++ logprobs = torch.log_softmax(torch.randn(3, _VOCAB), dim=-1) ++ reference = logprobs.clone() ++ requested = [[5, 7], None, [42]] ++ out = _attach( ++ logprobs=logprobs.clone(), ++ top_logprobs_nums=[0, 0, 0], ++ token_ids_logprobs=requested, ++ batch_next_token_ids=torch.tensor([3, 4, 9], dtype=torch.int32), ++ ) ++ vals = out.next_token_token_ids_logprobs_val ++ idxs = out.next_token_token_ids_logprobs_idx ++ self.assertEqual(len(vals), 3) ++ torch.testing.assert_close( ++ vals[0], reference[0, torch.tensor([5, 7])], rtol=0, atol=0 ++ ) ++ self.assertEqual(idxs[0], [5, 7]) ++ # None request yields empty placeholders. ++ self.assertEqual(vals[1], []) ++ self.assertEqual(idxs[1], []) ++ torch.testing.assert_close( ++ vals[2], reference[2, torch.tensor([42])], rtol=0, atol=0 ++ ) ++ self.assertEqual(idxs[2], [42]) ++ ++ def test_token_ids_logprobs_probs_input(self): ++ """logprobs_are_probs=True applies log() before clamping, per row.""" ++ probs = torch.softmax(torch.randn(2, _VOCAB), dim=-1) ++ reference = probs.clone() ++ out = _attach( ++ logprobs=probs.clone(), ++ top_logprobs_nums=[0, 0], ++ token_ids_logprobs=[[1, 2], [3]], ++ batch_next_token_ids=torch.tensor([0, 1], dtype=torch.int32), ++ logprobs_are_probs=True, ++ ) ++ torch.testing.assert_close( ++ out.next_token_token_ids_logprobs_val[0], ++ reference[0, torch.tensor([1, 2])].log(), ++ rtol=1e-6, ++ atol=1e-6, ++ ) ++ ++ ++class TestAttachTopLogprobs(CustomTestCase): ++ ++ def test_top_logprobs_values(self): ++ """Per-request k slicing stays correct with the single-topk path.""" ++ logprobs = torch.log_softmax(torch.randn(2, _VOCAB), dim=-1) ++ reference = logprobs.clone() ++ out = _attach( ++ logprobs=logprobs.clone(), ++ top_logprobs_nums=[3, 1], ++ token_ids_logprobs=[None, None], ++ batch_next_token_ids=torch.tensor([3, 4], dtype=torch.int32), ++ ) ++ ref_vals, ref_idx = reference.topk(3, dim=-1) ++ self.assertEqual(len(out.next_token_top_logprobs_val), 2) ++ torch.testing.assert_close( ++ out.next_token_top_logprobs_val[0], ref_vals[0], rtol=0, atol=0 ++ ) ++ self.assertTrue(torch.equal(out.next_token_top_logprobs_idx[0], ref_idx[0])) ++ torch.testing.assert_close( ++ out.next_token_top_logprobs_val[1], ref_vals[1][:1], rtol=0, atol=0 ++ ) ++ self.assertTrue(torch.equal(out.next_token_top_logprobs_idx[1], ref_idx[1][:1])) ++ ++ def test_next_token_logprobs_gather(self): ++ logprobs = torch.log_softmax(torch.randn(2, _VOCAB), dim=-1) ++ reference = logprobs.clone() ++ out = _attach( ++ logprobs=logprobs.clone(), ++ top_logprobs_nums=[0, 0], ++ token_ids_logprobs=[None, None], ++ batch_next_token_ids=torch.tensor([3, 4], dtype=torch.int32), ++ ) ++ torch.testing.assert_close( ++ out.next_token_logprobs, ++ torch.stack([reference[0, 3], reference[1, 4]]), ++ rtol=0, ++ atol=0, ++ ) ++ ++ ++if __name__ == "__main__": ++ unittest.main() From c53df5b134231dc53c20ad95f73ecd06d6d48f29 Mon Sep 17 00:00:00 2001 From: dabuliu123 <270334047@qq.com> Date: Fri, 4 Sep 2026 09:57:23 +0800 Subject: [PATCH 13/16] docs(npu): align BASE_IMAGE default with CANN 9.0.0 Dockerfile MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit feat(docker): add Ascend/NPU multi-stage image build with Makefile entrypoints Bring the GPU image pipeline's layered design to Ascend/910C so 云效 only selects code version, params, and scheduling while the Relax repo defines how images are built. - **`docker/Dockerfile.npu`**: refactor the monolithic NPU build into multi-stage `base → train → relax`, mirroring `docker/Dockerfile` (GPU). - `train` holds the Relax-code-independent heavy deps (torch_npu source build, MindSpeed/Megatron/MindSpeed-Bridge/Megatron-Bridge, sglang-npu, sgl-kernel-npu) and only `COPY docker/npu_patch` — a Relax code change no longer rebuilds them. - `relax` copies the full checkout (incl. `.git` for versioneer) and installs it; `ray`/`protobuf` pins stay after `pip install -e .` so they still win. - `BASE_IMAGE` / `SOC_VERSION` / `TRAIN_IMAGE` are now build args (default base `quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11`, SOC `ascend910_9391`). - A plain `docker build -f docker/Dockerfile.npu .` still works via the chained default `TRAIN_IMAGE=train`. - **`Makefile`**: add `docker-train-ascend`, `docker-dev-ascend`, and optional `docker-qs-ascend`, reusing the GPU tag/push/remote-skip logic. - Tags `ascend-train|dev|qs-YYYYMMDD-`; the `ascend-` prefix keeps aarch64 artifacts from overwriting the amd64 `train-/dev-` tags in the same repository. - `docker-qs-ascend` builds relax-ci's `Dockerfile.qs` (verified pure-python / arch-independent) on top of the dev image via `ASCEND_QS_DOCKERFILE`; cloning relax-ci and its credentials stay in CI, not the repo. - **`docker/npu-training.md`**: document the multi-stage layout, the Makefile build paths, configurable variables, and the optional QS wrapping. Co-Authored-By: Claude (cherry picked from commit 71f7a80c43dadf189c010c3ce5f5748784af5f21) (cherry picked from commit d9b9265418a24d382eaea279940a5391cec7fe38) chore: migrate references from `redai-infra` to `redai-studio` (#296) (cherry picked from commit 03cd24bcd7d739200eaa15543e2a7892cc4c7d9a) docs(npu): align BASE_IMAGE default with CANN 9.0.0 Dockerfile --- .github/ISSUE_TEMPLATE/config.yml | 4 +- .github/workflows/deploy-docs.yml | 4 +- CONTRIBUTING.md | 6 +- Makefile | 66 ++++++++++- README.md | 16 +-- README_zh.md | 16 +-- docker/Dockerfile | 4 +- docker/Dockerfile.npu | 103 ++++++++++-------- docker/npu-training.md | 32 +++++- docs/.vitepress/config.mts | 6 +- docs/.vitepress/theme/CallToAction.vue | 4 +- docs/deploy-docs.sh | 2 +- docs/en/api/actor-fwd.md | 4 +- docs/en/api/actor.md | 4 +- docs/en/api/genrm.md | 4 +- docs/en/api/rollout.md | 4 +- docs/en/guide/configuration.md | 2 +- docs/en/guide/customize-training.md | 4 +- docs/en/guide/fully-async-training.md | 2 +- docs/en/guide/how-to-contribute.md | 2 +- docs/en/guide/installation.md | 12 +- docs/en/guide/introduction.md | 4 +- .../reinforce-plus-plus-training-report.md | 2 +- docs/en/index.md | 2 +- docs/zh/api/actor-fwd.md | 4 +- docs/zh/api/actor.md | 4 +- docs/zh/api/genrm.md | 4 +- docs/zh/api/rollout.md | 4 +- docs/zh/guide/configuration.md | 2 +- docs/zh/guide/customize-training.md | 4 +- docs/zh/guide/fully-async-training.md | 2 +- docs/zh/guide/how-to-contribute.md | 2 +- docs/zh/guide/installation.md | 12 +- docs/zh/guide/introduction.md | 4 +- .../reinforce-plus-plus-training-report.md | 2 +- docs/zh/index.md | 2 +- examples/nemo_gym_agentic/README.md | 4 +- .../nemo_gym_agentic/recipes/gsm8k/README.md | 4 +- .../recipes/r2e-gym/README.md | 4 +- .../recipes/workplace-assistant/README.md | 4 +- examples/nemo_gym_agentic/service/Dockerfile | 2 +- relax/backends/megatron/model_provider.py | 2 +- relax/utils/arguments.py | 2 +- relax/utils/visualize/templates.py | 4 +- setup.py | 2 +- skills/sglang-upgrade/SKILL.md | 2 +- skills/sync-github/SKILL.md | 12 +- .../references/prompt-b-dev-to-main.md | 12 +- 48 files changed, 256 insertions(+), 153 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml index 34faf1138..e997cf03b 100644 --- a/.github/ISSUE_TEMPLATE/config.yml +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -1,8 +1,8 @@ blank_issues_enabled: true contact_links: - name: Documentation - url: https://redai-infra.github.io/Relax + url: https://redai-studio.github.io/Relax about: Check the documentation for guides and API reference - name: Questions & Discussions - url: https://github.com/redai-infra/Relax/discussions + url: https://github.com/redai-studio/Relax/discussions about: Ask questions and discuss ideas with the community diff --git a/.github/workflows/deploy-docs.yml b/.github/workflows/deploy-docs.yml index 9d8719240..bd797817a 100644 --- a/.github/workflows/deploy-docs.yml +++ b/.github/workflows/deploy-docs.yml @@ -22,7 +22,7 @@ concurrency: jobs: build: runs-on: ubuntu-latest - if: github.repository == 'redai-infra/Relax' + if: github.repository == 'redai-studio/Relax' steps: - name: Checkout uses: actions/checkout@v4 @@ -61,7 +61,7 @@ jobs: deploy: needs: build runs-on: ubuntu-latest - if: github.repository == 'redai-infra/Relax' + if: github.repository == 'redai-studio/Relax' environment: name: github-pages url: ${{ steps.deployment.outputs.page_url }} diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 062d4ebd3..b24745f0d 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -42,7 +42,7 @@ This project follows a standard code of conduct. Please be respectful, inclusive ```bash # Clone the repository -git clone https://github.com/redai-infra/Relax.git +git clone https://github.com/redai-studio/Relax.git cd Relax # Create virtual environment (recommended) @@ -155,7 +155,7 @@ feat(rollout): add streaming data consumption for async mode ## Reporting Bugs -Use the [Bug Report template](https://github.com/redai-infra/Relax/issues/new?template=bug_report.md) and include: +Use the [Bug Report template](https://github.com/redai-studio/Relax/issues/new?template=bug_report.md) and include: - **Environment** — OS, Python version, CUDA version, GPU type - **Steps to Reproduce** — Minimal commands to trigger the bug @@ -165,7 +165,7 @@ Use the [Bug Report template](https://github.com/redai-infra/Relax/issues/new?te ## Requesting Features -Use the [Feature Request template](https://github.com/redai-infra/Relax/issues/new?template=feature_request.md) and include: +Use the [Feature Request template](https://github.com/redai-studio/Relax/issues/new?template=feature_request.md) and include: - **Problem Statement** — What problem does this solve? - **Proposed Solution** — How should it work? diff --git a/Makefile b/Makefile index 79039702b..e43221b01 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,10 @@ -.PHONY: help install test lint format clean docs docs-dev docs-build docs-preview check-registry check-train-image docker-train docker-dev +.PHONY: help install test lint format clean docs docs-dev docs-build docs-preview check-registry check-train-image check-ascend-qs-base-image check-qs-dockerfile docker-train docker-dev docker-ascend docker-qs-ascend DOCKER ?= docker DOCKERFILE ?= docker/Dockerfile +ASCEND_DOCKERFILE ?= docker/Dockerfile.npu +SOC_VERSION ?= ascend910_9391 +ASCEND_DOCKER_BUILDKIT ?= 1 DOCKER_BUILD_PROGRESS ?= plain DOCKER_BUILD_ARGS ?= DO_PUSH ?= 1 @@ -13,7 +16,19 @@ IMAGE_REGISTRY := $(patsubst %/,%,$(strip $(REGISTRY))) DEFAULT_TRAIN_IMAGE := $(IMAGE_REGISTRY)/$(IMAGE_REPOSITORY):train-$(BUILD_DATE)-$(GIT_SHORT_HASH) DEV_IMAGE := $(IMAGE_REGISTRY)/$(IMAGE_REPOSITORY):dev-$(BUILD_DATE)-$(GIT_SHORT_HASH) -ifeq ($(origin TRAIN_IMAGE),undefined) +# Ascend/NPU images share the same repository as GPU; the ascend- tag prefix keeps +# aarch64 artifacts from ever overwriting the amd64 train-/dev- tags. +ASCEND_DEV_IMAGE := $(IMAGE_REGISTRY)/$(IMAGE_REPOSITORY):ascend-dev-$(BUILD_DATE)-$(GIT_SHORT_HASH) +ASCEND_QS_IMAGE := $(IMAGE_REGISTRY)/$(IMAGE_REPOSITORY):ascend-qs-$(BUILD_DATE)-$(GIT_SHORT_HASH) + +# QS wrapping reuses the external relax-ci Dockerfile.qs (verified pure-python / +# arch-independent, so ARM64-safe). CI is responsible for checking out relax-ci and +# pointing ASCEND_QS_DOCKERFILE at its docker/Dockerfile.qs; the Relax repo embeds neither +# the external repo nor its credentials. ASCEND_QS_BASE_IMAGE defaults to the dev image. +ASCEND_QS_DOCKERFILE ?= +ASCEND_QS_BASE_IMAGE ?= $(ASCEND_DEV_IMAGE) + +ifeq ($(strip $(TRAIN_IMAGE)),) TRAIN_IMAGE := $(DEFAULT_TRAIN_IMAGE) BUILD_DEFAULT_TRAIN := 1 else @@ -40,6 +55,10 @@ help: @echo " REGISTRY=... make docker-train - Build and push the Docker train stage" @echo " REGISTRY=... make docker-dev - Build and push the Docker development image" @echo " REGISTRY=... TRAIN_IMAGE=... make docker-dev - Build dev from an existing train image" + @echo " REGISTRY=... make docker-ascend - Build and push the complete Ascend/NPU image" + @echo " REGISTRY=... ASCEND_QS_DOCKERFILE=... make docker-qs-ascend - Wrap an Ascend dev image into a QS image (optional)" + @echo " Ascend targets accept BASE_IMAGE=... and SOC_VERSION=... (default ascend910_9391)" + @echo " Ascend targets use BuildKit by default; set ASCEND_DOCKER_BUILDKIT=0 for legacy DinD" @echo " Set DO_PUSH=0 before make to skip pushing Docker images" @echo " Existing remote images are skipped; DO_PUSH=0 checks local images" @@ -84,6 +103,13 @@ check-registry: check-train-image: @test -n "$(strip $(TRAIN_IMAGE))" || { echo "TRAIN_IMAGE must not be empty" >&2; exit 2; } +check-ascend-qs-base-image: + @test -n "$(strip $(ASCEND_QS_BASE_IMAGE))" || { echo "ASCEND_QS_BASE_IMAGE must not be empty" >&2; exit 2; } + +check-qs-dockerfile: + @test -n "$(strip $(ASCEND_QS_DOCKERFILE))" || { echo "ASCEND_QS_DOCKERFILE is required (path to relax-ci docker/Dockerfile.qs)" >&2; exit 2; } + @test -f "$(strip $(ASCEND_QS_DOCKERFILE))" || { echo "ASCEND_QS_DOCKERFILE not found: $(ASCEND_QS_DOCKERFILE)" >&2; exit 2; } + docker-train: check-registry @echo "[docker] output train image: $(TRAIN_IMAGE)" @set -e; \ @@ -118,3 +144,39 @@ docker-dev: check-registry check-train-image .; \ if [ "$(DO_PUSH)" != "0" ]; then $(DOCKER) push "$(DEV_IMAGE)"; fi; \ fi + +docker-ascend: check-registry + @echo "[docker] output ascend image: $(ASCEND_DEV_IMAGE)" + @set -e; \ + if $(IMAGE_INSPECT) "$(ASCEND_DEV_IMAGE)" >/dev/null 2>&1; then \ + echo "[docker] skip existing $(IMAGE_LOCATION) ascend image: $(ASCEND_DEV_IMAGE)"; \ + else \ + DOCKER_BUILDKIT=$(ASCEND_DOCKER_BUILDKIT) $(DOCKER) build --progress=$(DOCKER_BUILD_PROGRESS) \ + -f $(ASCEND_DOCKERFILE) \ + --target relax \ + -t "$(ASCEND_DEV_IMAGE)" \ + --build-arg SOC_VERSION="$(SOC_VERSION)" \ + $(PROXY_BUILD_ARGS) $(BASE_IMAGE_BUILD_ARG) $(DOCKER_BUILD_ARGS) \ + .; \ + if [ "$(DO_PUSH)" != "0" ]; then $(DOCKER) push "$(ASCEND_DEV_IMAGE)"; fi; \ + fi + +# Optional: wrap an Ascend dev image into an internal QS image using relax-ci's +# Dockerfile.qs. ASCEND_QS_DOCKERFILE must point at a relax-ci checkout; ASCEND_QS_BASE_IMAGE +# defaults to the dev image built above but can be any existing Ascend dev image. +docker-qs-ascend: check-registry check-qs-dockerfile check-ascend-qs-base-image + @echo "[docker] input ascend dev image: $(ASCEND_QS_BASE_IMAGE)" + @echo "[docker] output ascend qs image: $(ASCEND_QS_IMAGE)" + @set -e; \ + if $(IMAGE_INSPECT) "$(ASCEND_QS_IMAGE)" >/dev/null 2>&1; then \ + echo "[docker] skip existing $(IMAGE_LOCATION) ascend qs image: $(ASCEND_QS_IMAGE)"; \ + else \ + DOCKER_BUILDKIT=$(ASCEND_DOCKER_BUILDKIT) $(DOCKER) build --progress=$(DOCKER_BUILD_PROGRESS) \ + --no-cache \ + -f "$(ASCEND_QS_DOCKERFILE)" \ + -t "$(ASCEND_QS_IMAGE)" \ + --build-arg BASE_IMAGE="$(ASCEND_QS_BASE_IMAGE)" \ + $(PROXY_BUILD_ARGS) $(DOCKER_BUILD_ARGS) \ + "$(dir $(ASCEND_QS_DOCKERFILE))"; \ + if [ "$(DO_PUSH)" != "0" ]; then $(DOCKER) push "$(ASCEND_QS_IMAGE)"; fi; \ + fi diff --git a/README.md b/README.md index 03c93a2e8..79d644df8 100644 --- a/README.md +++ b/README.md @@ -16,13 +16,13 @@ arXiv - + Documentation - + WeChat QR - + Docker Image

@@ -34,7 +34,7 @@ ______________________________________________________________________ -**Relax** (**R**einforcement **E**ngine **L**everaging **A**gentic **X**-modality) is a high-performance reinforcement learning post-training framework open-sourced by the Xiaohongshu AI Infra Team for multimodal large language models. Built on Ray Serve with a service-oriented architecture, Relax uses Megatron-LM as the training backend and SGLang as the inference engine. Through the [TransferQueue](https://github.com/redai-infra/TransferQueue) data transfer system, it achieves complete decoupling of training and inference, supporting end-to-end multimodal RL training from text to images, videos, and audio. +**Relax** (**R**einforcement **E**ngine **L**everaging **A**gentic **X**-modality) is a high-performance reinforcement learning post-training framework open-sourced by the Xiaohongshu AI Infra Team for multimodal large language models. Built on Ray Serve with a service-oriented architecture, Relax uses Megatron-LM as the training backend and SGLang as the inference engine. Through the [TransferQueue](https://github.com/redai-studio/TransferQueue) data transfer system, it achieves complete decoupling of training and inference, supporting end-to-end multimodal RL training from text to images, videos, and audio. ______________________________________________________________________ @@ -129,15 +129,15 @@ The recommended way to run Relax is via the official Docker image, which ships w ```bash # Pull the official image -docker pull ghcr.io/redai-infra/relaxrl:latest +docker pull ghcr.io/redai-studio/relaxrl:latest # Launch a container with GPUs, shared memory, and your workspace mounted docker run -it --gpus all --ipc=host --network=host \ -v /path/to/your/workspace:/root \ - ghcr.io/redai-infra/relaxrl:latest bash + ghcr.io/redai-studio/relaxrl:latest bash # Inside the container -git clone https://github.com/redai-infra/Relax.git /root/Relax +git clone https://github.com/redai-studio/Relax.git /root/Relax cd /root/Relax && pip install -e . ``` @@ -246,7 +246,7 @@ ______________________________________________________________________ ## 📚 Documentation -Full bilingual documentation is available at **[redai-infra.github.io/Relax](https://redai-infra.github.io/Relax)**. +Full bilingual documentation is available at **[redai-studio.github.io/Relax](https://redai-studio.github.io/Relax)**. ______________________________________________________________________ diff --git a/README_zh.md b/README_zh.md index af13a4c36..faf13e888 100644 --- a/README_zh.md +++ b/README_zh.md @@ -16,13 +16,13 @@ arXiv - + Documentation - + WeChat QR - + Docker Image

@@ -34,7 +34,7 @@ ______________________________________________________________________ -**Relax**（**R**einforcement **E**ngine **L**everaging **A**gentic **X**-modality）是小红书 AI 平台开源的、面向多模态大模型的高性能强化学习后训练框架。Relax 基于 Ray Serve 构建面向服务的架构，以 Megatron-LM 为训练后端、SGLang 为推理引擎，通过 [TransferQueue](https://github.com/redai-infra/TransferQueue) 数据传输系统实现训练与推理的完全解耦，支持从文本到图像、视频、音频的全模态强化学习训练。 +**Relax**（**R**einforcement **E**ngine **L**everaging **A**gentic **X**-modality）是小红书 AI 平台开源的、面向多模态大模型的高性能强化学习后训练框架。Relax 基于 Ray Serve 构建面向服务的架构，以 Megatron-LM 为训练后端、SGLang 为推理引擎，通过 [TransferQueue](https://github.com/redai-studio/TransferQueue) 数据传输系统实现训练与推理的完全解耦，支持从文本到图像、视频、音频的全模态强化学习训练。 ______________________________________________________________________ @@ -129,15 +129,15 @@ ______________________________________________________________________ ```bash # 拉取官方镜像 -docker pull ghcr.io/redai-infra/relaxrl:latest +docker pull ghcr.io/redai-studio/relaxrl:latest # 启动容器，挂载 GPU、共享内存与工作目录 docker run -it --gpus all --ipc=host --network=host \ -v /path/to/your/workspace:/root \ - ghcr.io/redai-infra/relaxrl:latest bash + ghcr.io/redai-studio/relaxrl:latest bash # 容器内克隆仓库并安装 -git clone https://github.com/redai-infra/Relax.git /root/Relax +git clone https://github.com/redai-studio/Relax.git /root/Relax cd /root/Relax && pip install -e . ``` @@ -246,7 +246,7 @@ ______________________________________________________________________ ## 📚 文档 -完整的双语文档请访问 **[redai-infra.github.io/Relax](https://redai-infra.github.io/Relax)**。 +完整的双语文档请访问 **[redai-studio.github.io/Relax](https://redai-studio.github.io/Relax)**。 ______________________________________________________________________ diff --git a/docker/Dockerfile b/docker/Dockerfile index c473e89b0..70e98cd9f 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -68,7 +68,7 @@ RUN MAX_JOBS=64 \ rm -rf /opt/flash-attention/ RUN pip -v install --no-cache-dir --no-build-isolation "transformer_engine[pytorch]==2.14.1" && \ - TMS_CUDA_MAJOR=$(python -c 'import torch; print(torch.version.cuda.split(".")[0])') pip install git+https://github.com/redai-infra/torch_memory_saver.git@afc13785c50119048e2dd8ac497cc9e29ec75bd4 --no-cache-dir --force-reinstall && \ + TMS_CUDA_MAJOR=$(python -c 'import torch; print(torch.version.cuda.split(".")[0])') pip install git+https://github.com/redai-studio/torch_memory_saver.git@afc13785c50119048e2dd8ac497cc9e29ec75bd4 --no-cache-dir --force-reinstall && \ pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation --no-cache-dir && \ pip install "numpy<2" nvidia-cudnn-cu12==9.16.0.29 --no-cache-dir && \ NVCC_APPEND_FLAGS="--threads 32" \ @@ -116,7 +116,7 @@ COPY requirements.txt /tmp/requirements.txt RUN pip install --ignore-installed PyJWT && \ pip install -r /tmp/requirements.txt --no-cache-dir && \ pip install --no-cache-dir "compressed_tensors>=0.13.0" tensordict==0.10.0 pyvers==0.1.0 'nvidia-modelopt[hf]==0.44.0' --no-deps && \ - pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps + pip install "transferqueue @ git+https://github.com/redai-studio/TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps # sgl-router: override the official wheel (pulled by requirements.txt above) with # slime's r3-capable fork. The official sglang-router drops the routed_experts diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index ca4619d3e..bfc4f1a12 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -1,8 +1,10 @@ -# +# Ascend/NPU image: CANN base -> training stack -> runtime dependencies. +# CI can reuse a published train image through TRAIN_IMAGE. ARG HTTP_PROXY ARG HTTPS_PROXY ARG NO_PROXY -FROM quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11 +ARG BASE_IMAGE=quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11 +FROM ${BASE_IMAGE} as base ARG HTTP_PROXY ARG HTTPS_PROXY @@ -35,7 +37,7 @@ RUN DEBIAN_FRONTEND=noninteractive apt-get update && apt-get install -y --no-ins apt-get clean && \ rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* -# Setting env +# Configure CANN. RUN ARCH=$(uname -m) && \ if [ "$ARCH" = "aarch64" ]; then \ export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \ @@ -45,9 +47,7 @@ RUN ARCH=$(uname -m) && \ 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 - - -# Setting pip & git config. Global config (set once, persists across subsequent RUN layers) +# Configure pip and git. ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple" RUN pip config set global.index-url ${PIP_INDEX_URL} && \ git config --global http.sslverify false && \ @@ -55,82 +55,98 @@ RUN pip config set global.index-url ${PIP_INDEX_URL} && \ git config --global http.postBuffer 2147483648 && \ git config --global user.email "temp@example.com" && \ git config --global user.name "temp" +WORKDIR /root + +FROM base as train -# install torch -RUN pip install --upgrade pip packaging setuptools==80.10.2 && \ - pip install torch==2.9.0 && \ - pip install numpy==1.26.0 +# Ascend chip target, defaulting to 910C. +ARG SOC_VERSION=ascend910_9391 +ENV SOC_VERSION=${SOC_VERSION} -# build torch_npu -RUN pip install pyyaml && \ +WORKDIR /root +RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org/simple "$@"; }; \ + pip_install --upgrade pip packaging setuptools==80.10.2 && \ + pip_install torch==2.9.0 && \ + pip_install numpy==1.26.0 +# Build torch_npu. +RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org/simple "$@"; }; \ + pip_install pyyaml && \ git clone https://gitcode.com/Ascend/pytorch.git /root/pytorch && \ cd /root/pytorch && \ git checkout v26.0.1-pytorch2.9.0 && \ git cherry-pick -n f495de675bce38a2fa21edbf067b73d2a5f26733 && \ bash ci/build.sh --python=3.11 && \ pip install dist/torch_npu*.whl - -# install triton/TQ RUN pip install triton-ascend==3.2.1 --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple && \ - pip install tensordict==0.10.0 pyvers==0.1.0 --no-deps && \ - pip install "transferqueue @ git+https://github.com/redai-infra/TransferQueue.git@58054a33834aadbcf76aacd6b1e32e25c030f2c9" --no-deps - -# Clone Megatron-LM, MindSpeed, MindSpeed-Bridge, Megatron-Bridge and install + cd /root && rm -rf /root/pytorch +RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org/simple "$@"; }; \ + pip_install tensordict==0.10.0 pyvers==0.1.0 --no-deps +RUN git clone https://github.com/redai-studio/TransferQueue.git /root/TransferQueue && \ + cd /root/TransferQueue && \ + git checkout 58054a33834aadbcf76aacd6b1e32e25c030f2c9 && \ + pip install /root/TransferQueue --no-deps --no-build-isolation && \ + cd /root && \ + rm -rf /root/TransferQueue + +# Install the MindSpeed/Megatron stack. RUN git clone https://gitcode.com/ascend/MindSpeed.git /root/MindSpeed && \ git clone https://github.com/NVIDIA/Megatron-LM.git /root/Megatron-LM && \ git clone https://gitcode.com/ascend/MindSpeed-Ops.git /root/MindSpeed-Ops && \ git clone https://gitcode.com/ascend/MindSpeed-Bridge.git /root/MindSpeed-Bridge && \ git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git /root/Megatron-Bridge -RUN cd /root/MindSpeed && git checkout core_r0.16.0 && pip install -r requirements.txt && pip install -e . && git checkout e4772499 && \ +RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org/simple "$@"; }; \ + 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 && \ + cd /root/MindSpeed && git checkout core_r0.16.0 && pip install -r requirements.txt && pip install -e . && git checkout e4772499 && \ cd /root/Megatron-LM && git checkout core_v0.16.1 && pip install -e . --no-build-isolation && \ cd /root/Megatron-Bridge && git checkout v0.3.1 && \ - cd /root/MindSpeed-Ops/ && git checkout 33ac80f7 && pip install -e . --no-build-isolation --no-deps && \ - cd /root/MindSpeed-Bridge/ && git checkout v0.3.1 && pip install -r requirements.txt && pip install -e . --no-deps + cd /root/MindSpeed-Ops && git checkout 33ac80f7 && pip install -e . --no-build-isolation --no-deps && \ + cd /root/MindSpeed-Bridge && git checkout v0.3.1 && pip install -r requirements.txt && pip install -e . --no-deps +# Keep Relax code changes out of the dependency cache. +COPY docker/npu_patch /root/Relax/docker/npu_patch -COPY . /root/Relax -# Patch Megatron-LM, MindSpeed, MindSpeed-Bridge, MindSpeed-Ops +# Apply NPU patches. RUN cd /root/Megatron-Bridge && \ - patch -p1 < /root/Relax/docker/npu_patch/megatron-bridge.patch && \ + patch -p1 < /root/Relax/docker/npu_patch/megatron-bridge.patch && \ git add . && git commit -m "base line" && \ cd /root/Megatron-LM && \ - patch -p1 < /root/Relax/docker/npu_patch/megatron.patch && \ + patch -p1 < /root/Relax/docker/npu_patch/megatron.patch && \ git add . && git commit -m "base line" && \ cd /root/MindSpeed-Bridge && \ - patch -p1 < /root/Relax/docker/npu_patch/mindspeed-bridge.patch && \ + patch -p1 < /root/Relax/docker/npu_patch/mindspeed-bridge.patch && \ git add . && git commit -m "base line" && \ cd /root/MindSpeed-Ops && \ - patch -p1 < /root/Relax/docker/npu_patch/mindspeed-ops.patch && \ + patch -p1 < /root/Relax/docker/npu_patch/mindspeed-ops.patch && \ git add . && git commit -m "base line" && \ cd /root/MindSpeed && \ - patch -p1 < /root/Relax/docker/npu_patch/mindspeed.patch && \ + patch -p1 < /root/Relax/docker/npu_patch/mindspeed.patch && \ git add . && git commit -m "base line" # Install sglang -RUN git clone https://github.com/sgl-project/sglang.git /root/sglang && \ +RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org/simple "$@"; }; \ + git clone https://github.com/sgl-project/sglang.git /root/sglang && \ cd /root/sglang && git checkout v0.5.15 && \ mv python/pyproject.toml python/pyproject.toml.backup && \ mv python/pyproject_npu.toml python/pyproject.toml && \ pip install -e "python[srt_npu]" --constraint <(echo "torch==2.9.0") && \ - # patch -p1 < /root/Relax/docker/npu_patch/sglang-npu.patch 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" - -# Install sgl-kernle-npu +# 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 && \ - # Adapt tms for colocate train. patch -p1 < /root/Relax/docker/npu_patch/sgl-kernel-npu.patch && \ - git add . && git commit -m "sgl-kernel-npu.patch" && \ + git add . && git commit -m "sgl-kernel-npu.patch" && \ bash build.sh && \ pip install output/*.whl && \ - cd /root + cd /root && rm -rf /root/sgl-kernel-npu # Install AscendC FLA RUN git clone https://github.com/flashserve/flash-linear-attention-npu.git /root/flash-linear-attention-npu && \ @@ -145,18 +161,15 @@ RUN git clone https://github.com/flashserve/flash-linear-attention-npu.git /root 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 - - - -# Install Relax -# git clone https://github.com/redai-infra/Relax.git -RUN cd /root/Relax && \ - pip install -e . +FROM train as relax +WORKDIR /root -RUN pip install ray==2.55.1 && pip install protobuf==6.33.6 +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 -#Clean cache +# Clean caches. RUN pip cache purge && \ - rm -rf /tmp/* + rm -rf /tmp/* \ No newline at end of file diff --git a/docker/npu-training.md b/docker/npu-training.md index bc4710071..f5012126f 100644 --- a/docker/npu-training.md +++ b/docker/npu-training.md @@ -27,7 +27,7 @@ - 资源类型：`Ascend910 Snt9b23` - 驱动版本：`Software Version 25.5.1` - 固件版本：`Firmware Version 7.8.0.6.201` -- 基础镜像：`quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11` +- 基础镜像：`quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11` ### 环境检查 @@ -39,7 +39,35 @@ ### 安装方法 -（推荐）基于 Dockerfile 构建镜像：`docker build -f docker/Dockerfile.npu -t npu-2026 .` +`docker/Dockerfile.npu` 采用多阶段构建： + +- `base`：系统依赖 + CANN 工具链层； +- `train`：安装 torch_npu、MindSpeed、Megatron、sglang-npu 和 sgl-kernel-npu； +- `relax`：安装 Relax 的 Python 运行时依赖，不复制或安装 Relax 源码。训练任务运行时挂载目标分支代码。 + +通过 Makefile 一次构建并推送完整 Ascend 镜像，标签形如 `ascend-dev-YYYYMMDD-`： + +```bash +REGISTRY= make docker-ascend +``` + +910C DinD 使用传统 builder 时： + +```bash +REGISTRY= ASCEND_DOCKER_BUILDKIT=0 make docker-ascend +``` + +可配置变量：`BASE_IMAGE`（默认 `quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11`）、`SOC_VERSION`（默认 `ascend910_9391`）、`REGISTRY`、`DO_PUSH`（默认 `1`，设为 `0` 时不推送并检查本地镜像）、`ASCEND_DOCKER_BUILDKIT`（默认 `1`）。远端已存在同名镜像时会跳过构建。 + +（可选，内部 QS 镜像）复用 `ml-engine/tools/relax-ci` 的 `docker/Dockerfile.qs`（纯 Python 依赖，架构无关，ARM64 可直接构建）。由 CI 先 checkout relax-ci，再指向其 `Dockerfile.qs`： + +```bash +# 需先 checkout relax-ci，ASCEND_QS_DOCKERFILE 指向其 docker/Dockerfile.qs +REGISTRY= \ + ASCEND_QS_DOCKERFILE=/docker/Dockerfile.qs \ + ASCEND_QS_BASE_IMAGE=/relax:ascend-dev-YYYYMMDD- \ + make docker-qs-ascend +``` ## 启动配置 diff --git a/docs/.vitepress/config.mts b/docs/.vitepress/config.mts index c2157f054..29622c534 100644 --- a/docs/.vitepress/config.mts +++ b/docs/.vitepress/config.mts @@ -229,7 +229,7 @@ export default defineConfig({ { text: 'Resources', items: [ - { text: 'GitHub', link: 'https://github.com/redai-infra/Relax' }, + { text: 'GitHub', link: 'https://github.com/redai-studio/Relax' }, { text: 'Paper', link: 'https://arxiv.org/abs/2604.11554' } ] } @@ -341,7 +341,7 @@ export default defineConfig({ { text: '资源', items: [ - { text: 'GitHub', link: 'https://github.com/redai-infra/Relax' }, + { text: 'GitHub', link: 'https://github.com/redai-studio/Relax' }, { text: '论文', link: 'https://arxiv.org/abs/2604.11554' } ] } @@ -466,7 +466,7 @@ export default defineConfig({ themeConfig: { logo: '/rednote-logo.png', socialLinks: [ - { icon: 'github', link: 'https://github.com/redai-infra/Relax' } + { icon: 'github', link: 'https://github.com/redai-studio/Relax' } ], search: { provider: 'local' diff --git a/docs/.vitepress/theme/CallToAction.vue b/docs/.vitepress/theme/CallToAction.vue index ed957b6d7..674f503e5 100644 --- a/docs/.vitepress/theme/CallToAction.vue +++ b/docs/.vitepress/theme/CallToAction.vue @@ -26,8 +26,8 @@ const subtitle = computed(() => const primaryLabel = computed(() => isZh.value ? '参与贡献' : 'Contribute Now') const secondaryLabel = computed(() => isZh.value ? '讨论区' : 'Discussion') -const primaryLink = 'https://github.com/redai-infra/Relax/blob/main/CONTRIBUTING.md' -const secondaryLink = 'https://github.com/redai-infra/Relax/discussions' +const primaryLink = 'https://github.com/redai-studio/Relax/blob/main/CONTRIBUTING.md' +const secondaryLink = 'https://github.com/redai-studio/Relax/discussions'