diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 7caa97acb..ac5f95c8d 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -3,7 +3,7 @@ ARG HTTP_PROXY ARG HTTPS_PROXY ARG NO_PROXY -ARG BASE_IMAGE=quay.io/ascend/cann:8.5.1-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 @@ -45,11 +45,13 @@ 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 # 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 && \ + 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" @@ -71,13 +73,13 @@ RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org 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 + pip install dist/torch_npu*.whl +RUN pip install triton-ascend==3.2.1 --extra-index-url=https://triton-ascend.osinfra.cn/pypi/simple && \ + cd /root && rm -rf /root/pytorch RUN pip_install() { pip install "$@" || pip install --index-url https://pypi.org/simple "$@"; }; \ - cd /root && rm -rf /root/pytorch && \ - pip_install triton-ascend==3.2.0 && \ 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 && \ @@ -89,57 +91,77 @@ RUN git clone https://github.com/redai-studio/TransferQueue.git /root/TransferQu # 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 pip_install() { pip install "$@" || pip install --index-url https://pypi.org/simple "$@"; }; \ - cd /root/MindSpeed && git checkout core_r0.16.0 && pip_install -r requirements.txt && pip_install -e . --no-build-isolation && \ - cd /root/Megatron-LM && git checkout core_v0.16.1 && pip_install -e . --no-build-isolation && \ + 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-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 # Keep Relax code changes out of the dependency cache. COPY docker/npu_patch /root/Relax/docker/npu_patch # Apply NPU patches. -RUN cd /root/MindSpeed && \ - patch -p1 < /root/Relax/docker/npu_patch/mindspeed.patch && \ - cd /root/Megatron-Bridge && \ +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 && \ + 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" -# Merge bridge packages 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. +# Install 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.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 && \ - echo "torch==2.9.0" > /tmp/sglang-constraints.txt && \ - pip_install -e "python[srt_npu]" --constraint /tmp/sglang-constraints.txt && \ - python -c 'import importlib.metadata as m; v = m.version("torch"); assert v.startswith("2.9.0"), v' && \ - rm /tmp/sglang-constraints.txt && \ - git stash && \ - git fetch origin pull/23815/head:pr-23815 && \ - git checkout pr-23815 && \ - patch -p1 < /root/Relax/docker/npu_patch/sglang-npu.patch + pip install -e "python[srt_npu]" --constraint <(echo "torch==2.9.0") && \ + 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 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.04.15.rc3 && \ - 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 && \ + cd /root/sgl-kernel-npu && git checkout 2026.7.2 && \ + patch -p1 < /root/Relax/docker/npu_patch/sgl-kernel-npu.patch && \ + git add . && git commit -m "sgl-kernel-npu.patch" && \ + bash build.sh && \ pip install output/*.whl && \ 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 && \ + cd /root/flash-linear-attention-npu && git checkout v26.1.0 && \ + apt update && apt-get install -y --no-install-recommends gawk && \ + # 编译命令，注意--soc=${soc_version}需要指定为当前机器的芯片类型{ascend910b/ascend910_93/ascend950} + bash build.sh --soc=ascend910_93 --pkg --ops=causal_conv1d,chunk_bwd_dv_local,chunk_bwd_dqkwg,chunk_gated_delta_rule_bwd_dhu,prepare_wy_repr_bwd_da,prepare_wy_repr_bwd_full,chunk_fwd_o,chunk_gated_delta_rule_fwd_h,recurrent_gated_delta_rule,recompute_wu_fwd && \ + # 安装run包 + ./build_out/cann-*.run && \ + source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + source /usr/local/Ascend/cann-9.0.0/share/info/ascendnpu-ir/bin/set_env.sh && \ + # 一键编译安装脚本，先调用torchnpugen自动接入算子，再运行setup编whl包，最后安装whl包 + cd torch_custom/fla_npu && bash build.sh + FROM train as relax WORKDIR /root diff --git a/docker/npu-training.md b/docker/npu-training.md index 4650af73e..f5012126f 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` | + ## 环境准备 ### 前置准备 @@ -19,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` ### 环境检查 @@ -49,7 +57,7 @@ REGISTRY= make docker-ascend REGISTRY= ASCEND_DOCKER_BUILDKIT=0 make docker-ascend ``` -可配置变量：`BASE_IMAGE`（默认 `quay.io/ascend/cann:8.5.1-a3-ubuntu22.04-py3.11`）、`SOC_VERSION`（默认 `ascend910_9391`）、`REGISTRY`、`DO_PUSH`（默认 `1`，设为 `0` 时不推送并检查本地镜像）、`ASCEND_DOCKER_BUILDKIT`（默认 `1`）。远端已存在同名镜像时会跳过构建。 +可配置变量：`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`： @@ -146,6 +154,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` 的值时脚本会报错退出 + ## 下一步 -- [ ] 特性支持：多模态、CP长序列、MTP 等 +- [ ] 性能优化：Qwen3.5-35B-A3B 等 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..6d574e221 100644 --- a/docker/npu_patch/mindspeed-bridge.patch +++ b/docker/npu_patch/mindspeed-bridge.patch @@ -19,64 +19,304 @@ 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..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 -@@ -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 +@@ -36,13 +36,15 @@ from megatron.core.ssm.gated_delta_net import ( + _split_tensor_factory, + ) + +-try: +- from causal_conv1d import causal_conv1d +-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 +61,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 +199,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__() - -- def finalize(self) -> None: -- self.validate_parallelism() -- super().finalize() + else: + self.gated_delta_rule = torch_chunk_gated_delta_rule +@@ -474,7 +483,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 +564,7 @@ class GatedDeltaNet(MegatronModule): + + # 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()) + + # 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, # ← 新增：同上 + ) + + # 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) - -- def validate_parallelism(self): -- """Validate that parallelism settings are compatible with this model's architecture. +- 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)) - -- 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." -- ) +- 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) - - 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__() - -- def finalize(self) -> None: -- self.validate_parallelism() -- super().finalize() +- hf_config = self._hf_config +- num_mtp_layers = getattr(hf_config.text_config, "mtp_num_hidden_layers", None) - -- def validate_parallelism(self): -- """Validate that parallelism settings are compatible with this model's architecture. +- 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)) - -- 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 model instance with vision and language components. +- 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) + 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..f7bc22585 100644 --- a/docker/npu_patch/sglang-npu.patch +++ b/docker/npu_patch/sglang-npu.patch @@ -1,56 +1,1449 @@ +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/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py +index 6696945..7854f7a 100644 +--- a/python/sglang/srt/layers/radix_linear_attention.py ++++ b/python/sglang/srt/layers/radix_linear_attention.py +@@ -35,6 +35,27 @@ if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + ++def wrap_conv1d_weight_loader(attn: RadixLinearAttention, base_loader): ++ """Wrap conv1d.weight's loader so every (re)load also refreshes the ++ transposed copy cached on the RadixLinearAttention wrapper. ++ ++ mamba_v2_sharded_weight_loader writes through ``param.data[...] = ...``; ++ ``Tensor.data`` carries a version counter independent of the Parameter's, ++ so ``_version``-based invalidation cannot observe online weight updates ++ (e.g. RL Actor->SGLang sync via ``update_weights_from_tensor``). The ++ refresh is explicit and in place so captured graphs keep referencing the ++ same storage. ++ """ ++ ++ def loader(param, loaded_weight): ++ base_loader(param, loaded_weight) ++ cached = getattr(attn, "_conv_weights_t", None) ++ if cached is not None: ++ cached.copy_(attn.conv_weights.transpose(0, 1)) ++ ++ return loader ++ ++ + class RadixLinearAttention(nn.Module): + """ + The Linear Attention Layer Implementation. +diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py +index e83e4157fa..19ed3f244b 100644 +--- 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): +@@ -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. Runs topk exactly once. ++ 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 +567,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.next_token_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) +@@ -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 fed78cf..11cbadd 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 + +@@ -64,7 +65,10 @@ from sglang.srt.layers.parameter import ( + ) + from sglang.srt.layers.quantization.base_config import QuantizationConfig + from sglang.srt.layers.radix_attention import RadixAttention +-from sglang.srt.layers.radix_linear_attention import RadixLinearAttention ++from sglang.srt.layers.radix_linear_attention import ( ++ RadixLinearAttention, ++ wrap_conv1d_weight_loader, ++) + from sglang.srt.layers.rotary_embedding import get_rope + from sglang.srt.layers.utils import PPMissingLayer, get_layer_id + from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +@@ -229,23 +233,6 @@ class Qwen3_5GatedDeltaNet(nn.Module): + self._bind_packed_weight_loaders(self.in_proj_qkvz) + self._bind_packed_weight_loaders(self.in_proj_ba) + +- # Conv1d weight loader setup +- query_key_settings = (self.key_dim, 0, False) +- value_settings = (self.value_dim, 0, False) +- +- self._override_weight_loader( +- self.conv1d.weight, +- mamba_v2_sharded_weight_loader( +- [ +- query_key_settings, +- query_key_settings, +- value_settings, +- ], +- self.attn_tp_size, +- self.attn_tp_rank, +- ), +- ) +- + # State parameters + self.dt_bias = nn.Parameter( + torch.ones(self.num_v_heads // self.attn_tp_size), +@@ -275,6 +262,28 @@ class Qwen3_5GatedDeltaNet(nn.Module): + dt_bias=self.dt_bias, + ) + ++ # Conv1d weight loader setup. The wrapper refreshes the transposed ++ # copy cached on self.attn by the attention backend after every ++ # (re)load, so it must be bound after self.attn exists. ++ query_key_settings = (self.key_dim, 0, False) ++ value_settings = (self.value_dim, 0, False) ++ ++ self._override_weight_loader( ++ self.conv1d.weight, ++ wrap_conv1d_weight_loader( ++ self.attn, ++ mamba_v2_sharded_weight_loader( ++ [ ++ query_key_settings, ++ query_key_settings, ++ value_settings, ++ ], ++ self.attn_tp_size, ++ self.attn_tp_rank, ++ ), ++ ), ++ ) ++ + self.norm = RMSNormGated( + self.head_v_dim, + eps=self.layer_norm_epsilon, +@@ -608,7 +617,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 +837,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 +1224,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/layers/test_conv_weight_cache_refresh.py b/test/registered/unit/layers/test_conv_weight_cache_refresh.py +new file mode 100644 +index 0000000000..ed2574df66 100644 +--- /dev/null ++++ b/test/registered/unit/layers/test_conv_weight_cache_refresh.py +@@ -0,0 +1,109 @@ ++"""Regression tests for the GDN conv1d weight cache refresh. ++ ++The Ascend GDN backend caches a transposed copy of conv1d.weight on the ++RadixLinearAttention wrapper (AscendGDNAttnBackend._get_conv_weights_t). ++mamba_v2_sharded_weight_loader writes through ``param.data[...] = ...``; ++because ``Tensor.data`` carries a version counter independent of the ++Parameter's, online weight updates (e.g. RL Actor->SGLang sync via ++``update_weights_from_tensor``) are invisible to ``_version``-based ++invalidation. conv1d.weight's loader is therefore wrapped with ++``wrap_conv1d_weight_loader`` to refresh the cached copy explicitly, in ++place, so captured graphs keep referencing the same storage. ++""" ++ ++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.attention.mamba.mamba import mamba_v2_sharded_weight_loader ++from sglang.srt.layers.radix_linear_attention import wrap_conv1d_weight_loader ++ ++register_cpu_ci(est_time=10, suite="base-a-test-cpu") ++ ++_CHANNELS = 16 # q(4) + k(4) + v(8) on a single TP rank ++_KERNEL = 4 ++ ++ ++def _make_conv_param() -> torch.nn.Parameter: ++ # conv1d.weight layout: [channels, 1, kernel] ++ return torch.nn.Parameter( ++ torch.randn(_CHANNELS, 1, _KERNEL), requires_grad=False ++ ) ++ ++ ++def _make_base_loader(): ++ # Same shard layout as Qwen3_5: q/k shards then the v shard, TP size 1. ++ return mamba_v2_sharded_weight_loader( ++ [(4, 0, False), (4, 0, False), (8, 0, False)], 1, 0 ++ ) ++ ++ ++class _FakeAttn: ++ """Stands in for RadixLinearAttention: holds conv_weights and the cache.""" ++ ++ def __init__(self, conv_param: torch.nn.Parameter): ++ self.conv_weights = conv_param.view(_CHANNELS, _KERNEL) ++ self._conv_weights_t = self.conv_weights.transpose(0, 1).contiguous() ++ ++ ++class TestConvWeightCacheRefresh(CustomTestCase): ++ ++ def test_mamba_loader_write_bypasses_param_version(self): ++ """Root-cause guard: the real loader must not bump ``_version``. ++ ++ If a future PyTorch/SGLang change makes param.data writes bump the ++ shared version counter, this test fails and the explicit wrapper may ++ be revisited. ++ """ ++ param = _make_conv_param() ++ view = param.view(_CHANNELS, _KERNEL) ++ param_version = param._version ++ view_version = view._version ++ ++ _make_base_loader()(param, torch.randn_like(param)) ++ ++ self.assertEqual(param._version, param_version) ++ self.assertEqual(view._version, view_version) ++ ++ def test_wrapped_loader_refreshes_cache_in_place(self): ++ """An online reload must refresh the cached transpose in place.""" ++ param = _make_conv_param() ++ attn = _FakeAttn(param) ++ cached = attn._conv_weights_t ++ cached_ptr = cached.data_ptr() ++ ++ new_weight = torch.randn_like(param) ++ wrapped = wrap_conv1d_weight_loader(attn, _make_base_loader()) ++ wrapped(param, new_weight) ++ ++ # The loader actually wrote the new weights. ++ torch.testing.assert_close(param.detach(), new_weight, rtol=0, atol=0) ++ # The cache is refreshed in place: same tensor object, same storage ++ # (captured graphs keep referencing it), content matches the ++ # transpose of the new weights. ++ self.assertIs(attn._conv_weights_t, cached) ++ self.assertEqual(attn._conv_weights_t.data_ptr(), cached_ptr) ++ reference = param.view(_CHANNELS, _KERNEL).transpose(0, 1) ++ torch.testing.assert_close(attn._conv_weights_t, reference, rtol=0, atol=0) ++ ++ def test_wrapped_loader_without_cache_is_noop(self): ++ """Initial load (before the first forward) must not create a cache.""" ++ param = _make_conv_param() ++ attn = _FakeAttn(param) ++ del attn._conv_weights_t ++ ++ new_weight = torch.randn_like(param) ++ wrapped = wrap_conv1d_weight_loader(attn, _make_base_loader()) ++ wrapped(param, new_weight) ++ ++ self.assertFalse(hasattr(attn, "_conv_weights_t")) ++ torch.testing.assert_close(param.detach(), new_weight, rtol=0, atol=0) ++ ++ ++if __name__ == "__main__": ++ unittest.main() +diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py +index 9ad577b42e..390a3721c5 100755 +--- a/test/registered/unit/mem_cache/test_mamba_unittest.py ++++ b/test/registered/unit/mem_cache/test_mamba_unittest.py +@@ -419,6 +419,58 @@ class TestMamba(unittest.TestCase): + self.assertEqual(list(second_insert_events[0].token_ids), [5]) + self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash) + ++ def test_mamba_radix_cache_limited_partial_page_match_does_not_split(self): ++ page_size = 64 ++ tree = self._setup_minimal_mamba_radix_cache(page_size) ++ token_ids = array("q", range(page_size)) ++ ++ tree.insert( ++ InsertParams( ++ key=RadixKey(token_ids, None), ++ value=torch.arange(page_size), ++ mamba_value=torch.tensor([0]), ++ ) ++ ) ++ ++ match = tree.match_prefix( ++ MatchPrefixParams(key=RadixKey(token_ids, None, limit=page_size - 1)) ++ ) ++ ++ self.assertEqual(len(match.device_indices), 0) ++ self.assertEqual(self._non_root_key_lengths(tree), [page_size]) ++ ++ def _setup_minimal_mamba_radix_cache(self, page_size: int) -> MambaRadixCache: ++ tree = MambaRadixCache.__new__(MambaRadixCache) ++ tree.page_size = page_size ++ tree.mamba_cache_chunk_size = page_size ++ tree.disable = False ++ tree.device = torch.device("cpu") ++ tree.enable_kv_cache_events = False ++ tree.kv_event_queue = [] ++ tree.full_evictable_size_ = 0 ++ tree.mamba_evictable_size_ = 0 ++ tree.full_protected_size_ = 0 ++ tree.mamba_protected_size_ = 0 ++ ++ tree.root_node = TreeNode() ++ tree.root_node.key = RadixKey(array("q"), None) ++ tree.root_node.value = [] ++ tree.root_node.hash_value = [] ++ tree.root_node.full_lock_ref = 1 ++ tree.root_node.mamba_lock_ref = 1 ++ tree.full_lru_list = LRUList(mamba=False) ++ tree.mamba_lru_list = LRUList(mamba=True) ++ return tree ++ ++ def _non_root_key_lengths(self, tree: MambaRadixCache) -> list[int]: ++ lengths = [] ++ stack = list(tree.root_node.children.values()) ++ while stack: ++ node = stack.pop() ++ lengths.append(len(node.key)) ++ stack.extend(node.children.values()) ++ return lengths ++ + def _setup_tree_and_allocator(self, enable_kv_cache_events=False): + """Helper to create a MambaRadixCache with allocator for testing.""" + server_args = ServerArgs(model_path="dummy", page_size=1) +diff --git a/test/registered/unit/layers/test_sampler_attach_logprobs.py b/test/registered/unit/layers/test_sampler_attach_logprobs.py +new file mode 100644 +index 0000000000..9dfc930392 100644 +--- /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() diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 38df4258a..4a7364031 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -214,9 +214,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/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..606b50240 --- /dev/null +++ b/scripts/training/sft/run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh @@ -0,0 +1,164 @@ +#!/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 + +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:-${SCRIPT_DIR}/../../../../exps}" +MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}" +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'" +) +PROMPT_DATA="[$(IFS=,; echo "${TRAIN_FILES[*]}")]" +SAVE_DIR="${SAVE_DIR:=${MODEL_DIR}/checkpoint/checkpoints/qwen3.5-35B-A3B-mtp-pokemon-sft-0821}" + +CKPT_ARGS=( + --hf-checkpoint ${MODEL_DIR}/Qwen3.5-35B-A3B + --ref-load ${MODEL_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 +) + +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 + --mtp-loss-scaling-factor ${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-9B-16xnpu-cp.sh b/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh new file mode 100755 index 000000000..5bbdb075d --- /dev/null +++ b/scripts/training/text/run-qwen35-9B-16xnpu-cp.sh @@ -0,0 +1,179 @@ +#!/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 + +source "${MODEL_CONFIG_DIR}/qwen35-9B.sh" + +PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" +EXP_DIR="${EXP_DIR:-${SCRIPT_DIR}/../../../../exps}" +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 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 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 +) + +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 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..3d82bc949 --- /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:-${SCRIPT_DIR}/../../../../exps}" +PROJECT_NAME="${PROJECT_NAME:=Relax/dev/dapo-math}" +NUM_ROLLOUT="${NUM_ROLLOUT:=3000}" + +CKPT_ARGS=( + --hf-checkpoint ${EXP_DIR}/Qwen3.5-35B-A3B + --ref-load ${EXP_DIR}/Qwen3.5-35B-A3B + --megatron-to-hf-mode bridge + # --load ${EXP_DIR}/Qwen3.5-35B-A3B-save-0821 + --save ${EXP_DIR}/Qwen3.5-35B-A3B-save-0821 + --save-interval 100 + --max-actor-ckpt-to-keep 1 +) + +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..dbf4142de --- /dev/null +++ b/scripts/training/text/run_qwen35_9B_mtp_8xnpu_thd.sh @@ -0,0 +1,211 @@ +#!/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:-${SCRIPT_DIR}/../../../../exps}" +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 +) + +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 + --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