Skip to content

[WIP]SGLang NPU optimizations for Qwen3.5 9B & 35B - #308

Open
Lw135 wants to merge 2 commits into
redai-studio:mainfrom
Lw135:ascend-dev-0908
Open

[WIP]SGLang NPU optimizations for Qwen3.5 9B & 35B#308
Lw135 wants to merge 2 commits into
redai-studio:mainfrom
Lw135:ascend-dev-0908

Conversation

@Lw135

@Lw135 Lw135 commented Sep 9, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

This merge (!70) brings the latest SGLang NPU optimizations for Qwen3.5 9B & 35B into ascend-dev-0908: refreshed sglang-npu.patch / sgl-kernel-npu.patch patch sets with fused GDN & MoE kernels, an updated NPU Dockerfile build flow (Ascend910_9382 SOC, git am patch application, fastokens), a new Qwen3.5-9B colocate training script, and a reworked 35B-A3B script scaled to 2×16 NPUs.

Key Changes

1. SGLang NPU Patch Refresh (sglang-npu.patch, ~3.1k lines)

  • GDN / hybrid linear attention NPU backend optimizations: GDN recurrent AscendC fusion, QKVZBA packing, exp-race fused Triton kernel, full-attention fusion, TP AscendC fusion
  • MoE: front-routing fusion, weight prefetch, DeepEP token dispatcher updates
  • Mamba: radix cache with extra_buffer scheduling, bf16 SSM dtype
  • New optimization switches: multi-stream, async exponential, prefill delayer, task queue
  • fastokens tokenizer backend and EAGLE speculative-sampling cache-loc Triton op

2. sgl-kernel-npu Patch Refresh (sgl-kernel-npu.patch, ~2.9k lines)

  • New fused kernels: fused_qkvzba_conv1d, fused_sigmoid_gating_recurrent (bf16/fp32), updated causal_conv1d tiling, cache-location assignment
  • MoE front routing + persistent GMM; chain/tree speculative-sampling kernels
  • Build: init catlass/pto-isa submodules, per-target builds (kernels, memory-saver, deepep) for Ascend910_9382

3. NPU Dockerfile Updates

  • New SOC_VERSION_FOR_SGL_BUILD=Ascend910_9382 build arg; patches applied via git am
  • Install fastokens 0.3.1

4. Training Scripts

  • New run-qwen35-9B-8xnpu-colocate.sh: Qwen3.5-9B DAPO-math colocate training (TP4, THD, dynamic batch size, mamba extra_buffer, fastokens tokenizer)
  • Reworked run-qwen35-35B-A3B-16xnpu-colocate.sh: 2×16 NPUs, TP8/PP2/EP16, THD with 20480 tokens/GPU, mem fraction 0.85, --overlap-grad-reduce / --overlap-param-gather, plus the full set of NPU optimization env flags (SGLANG_NPU_*, ASCEND_USE_FIA, etc.)
  • loss:
image
  • performance:
Env Job perf/train_wait_time (s) perf/train_time (s) perf/step_time (s)
910C-hdk26.0.rc1 qwen35-35B-A3B(16x910C gbs128) step 1~35 AVG 192.14 81.72 273.86
H800 qwen35-35B-A3B(16xH800 gbs128) step 1~35 AVG 141.71 88.01 229.71
compare - 0.74 1.08 0.84

5. Runtime Fix

  • reloadable_process_group.py: skip post-destroy delay on NPU (post_destroy_delay = 0.0)

Major Commits

Commit Description
61e57d2 feat(NPU): add latest sglang optimize (merge of wqw_base_lw_dev_0808)
597b703 !70 merge ascend-dev-0908-pr into ascend-dev-0908

Lw135 and others added 2 commits September 9, 2026 10:08
@Lw135 @Li-Yongwen
feat(NPU): add latest sglang optimize
61e57d2 
Merge branch 'wqw_base_lw_dev_0808' into ascend-dev-0808

feat(docker): upgrade NPU env to CANN 9.0.0 with FLA and MindSpeed-Ops

# ⭐ Feature

## Upgrade CANN base image and toolchain

- Upgrade CANN base image from 8.5.1-a3 to 9.0.0-a3
- Upgrade torch_npu from v2.9.0-7.3.0 to v26.0.1-pytorch2.9.0
- Upgrade triton-ascend from 3.2.0 to 3.2.1

## Add AscendC Flash Linear Attention (FLA) support

- Clone and build fla_npu from flash-linear-attention-npu v26.1.0
- Compile causal_conv1d and gated_delta_rule ops for ascend910_93

## Add MindSpeed-Ops support

- Clone, checkout and install MindSpeed-Ops
- Add mindspeed-ops.patch for gated_delta_rule autotune key fix

---
# ♻️ Refactor

## Restructure Dockerfile build order

- Move torch/torch_npu install before repo clone
- Delay COPY . /root/Relax to just before patching
- Consolidate Megatron-Bridge into Megatron-LM via cp instead of separate path
- Install MindSpeed-Bridge with --no-deps to avoid circular dependency

## Migrate mindspeed-bridge to FLA ops

- Replace causal_conv1d from causal_conv1d with FLA implementation
- Replace mindspeed_ops l2norm with naive_l2norm fallback
- Switch to flash_gated_delta_rule when FLA is available

---
# 🐛 Bug Fix

## Fix various compatibility issues

- Fix autotune key in gated_delta_rule kernel by removing USE_G and IS_VARLEN
- Fix lambda closure bug in transformer_config_init_subclass (use default arg binding)
- Fix attention_mask dtype from int32 to bool in model preprocess
- Fix vision model config with MoE permute fusion disabled
- Handle OmegaConf DictConfig/ListConfig in remove_non_pickleables
- Guard apex MixedFusedLayerNorm import with is_npu_available check

## Update sgl-kernel-npu

- Upgrade sgl-kernel-npu checkout from 2026.04.15.rc3 to 2026.7.2
- Remove obsolete cherry-pick workaround

feat(docker): upgrade NPU env to CANN 9.0.0 with FLA and MindSpeed-Ops

feat(docker): upgrade NPU env to CANN 9.0.0 with FLA and MindSpeed-Ops

# ⭐ Feature

## Upgrade CANN base image and toolchain

- Upgrade CANN base image from 8.5.1-a3 to 9.0.0-a3
- Upgrade torch_npu from v2.9.0-7.3.0 to v26.0.1-pytorch2.9.0
- Upgrade triton-ascend from 3.2.0 to 3.2.1

## Add AscendC Flash Linear Attention (FLA) support

- Clone and build fla_npu from flash-linear-attention-npu v26.1.0
- Compile causal_conv1d and gated_delta_rule ops for ascend910_93

## Add MindSpeed-Ops support

- Clone, checkout and install MindSpeed-Ops
- Add mindspeed-ops.patch for gated_delta_rule autotune key fix

---
# ♻️ Refactor

## Restructure Dockerfile build order

- Move torch/torch_npu install before repo clone
- Delay COPY . /root/Relax to just before patching
- Consolidate Megatron-Bridge into Megatron-LM via cp instead of separate path
- Install MindSpeed-Bridge with --no-deps to avoid circular dependency

## Migrate mindspeed-bridge to FLA ops

- Replace causal_conv1d from causal_conv1d with FLA implementation
- Replace mindspeed_ops l2norm with naive_l2norm fallback
- Switch to flash_gated_delta_rule when FLA is available

---
# 🐛 Bug Fix

## Fix various compatibility issues

- Fix autotune key in gated_delta_rule kernel by removing USE_G and IS_VARLEN
- Fix lambda closure bug in transformer_config_init_subclass (use default arg binding)
- Fix attention_mask dtype from int32 to bool in model preprocess
- Fix vision model config with MoE permute fusion disabled
- Handle OmegaConf DictConfig/ListConfig in remove_non_pickleables
- Guard apex MixedFusedLayerNorm import with is_npu_available check

## Update sgl-kernel-npu

- Upgrade sgl-kernel-npu checkout from 2026.04.15.rc3 to 2026.7.2
- Remove obsolete cherry-pick workaround


[NPU] update sglang


fix(megatron): attach _hf_config for MTP bridge when pp>1

# 🐛 Bug Fix

## Fix bridge mapping registry failure with PP > 1

- Set model_bridge._hf_config from bridge.hf_pretrained.config when missing

chore(docker): refresh NPU patch set

# 🔩 Chore

## Sync NPU patches

- Add sgl-kernel-npu.patch (1847 lines)
- Update sglang-npu.patch (+1070 lines)
- Update mindspeed.patch / mindspeed-bridge.patch / megatron patches

build(docker): upgrade NPU image to CANN 9.0.0

# 🔩 Chore

## Upgrade NPU Dockerfile dependencies

- Bump base image CANN 8.5.1 → 9.0.0
- Upgrade torch_npu to v26.0.1 (PyTorch 2.9.0)
- Upgrade triton-ascend 3.2.0 → 3.2.1

feat(npu): add Qwen3.5 MTP SFT/training scripts

feat(npu): add Qwen3.5 MTP training scripts for 0821

# ⭐ Feature

## Add Qwen3.5 MTP training launch scripts

- run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh: Qwen3.5-35B-A3B MTP SFT on pokemon-gpt4o-captions, 8xNPU single-node, ray-submit launch
- run_qwen35-35B-A3B-16xnpu-colocate-thd.sh: Qwen3.5-35B-A3B colocate THD training on 16xNPU
- run_qwen35_9B_mtp_8xnpu_thd.sh: Qwen3.5-9B MTP THD training on 8xNPU

---

# 🐛 Bug Fix

## Fix EXP_DIR silently overridden by MODEL_DIR default

- EXP_DIR now uses its own default with MODEL_DIR following EXP_DIR, matching the 9B THD script pattern

---

# 🔩 Chore

## Align script naming and comments

- Rename scripts to the 8xnpu naming convention and drop duplicated .sh suffix
- Remove commented-out --qkv-format bshd / --micro-batch-size 1 lines
- Sync Usage comments with actual script names


Revert "fix(megatron): attach _hf_config for MTP bridge when pp>1"

This reverts commit 57809b4.

fix: update qwen35 mtp mapping in mindspeed-bridge.patch for relax bridge_converter.py

fix(npu): rename Qwen3.5-9B CP script to 16xnpu

feat(npu): add Qwen3.5-9B CP training script

# ⭐ Feature

## Add Qwen3.5-9B CP colocate training script

- Add run-qwen35-9B-8xnpu-cp.sh for DAPO math training with TP4/CP4 on 16 NPUs
- Set MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV=1 to use torch conv fallback for GDN

## Support GDN causal conv on NPU via fla_npu

- Replace causal_conv1d with fla_npu in mindspeed-bridge.patch GDN layer
- Add MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV env switch to fall back to torch conv

---

# 🐛 Bug Fix

## Fix repatch ordering before Megatron init

- Move repatch(args) before init(args) in MegatronTrainRayActor so bridge patches apply during initialization

fix(npu): rename Qwen3.5-9B CP script to 16xnpu

# 🐛 Bug Fix

## Fix Qwen3.5-9B CP script naming for 16 NPUs

- Rename run-qwen35-9B-8xnpu-cp.sh to run-qwen35-9B-16xnpu-cp.sh to match the actual 16-NPU setup
- Update header comment to 16xNPU and log filename to qwen35-9B-GRPO-npu16

docs(npu): add feature support table

# 📝 Documentation

## Add feature support table to NPU training guide

- Add 特性支持 table covering Qwen3.5-9B CP, Qwen3.5-9B MTP, and Qwen3.5-35B-A3B SFT MTP with reference scripts
- Update 下一步 checklist from feature support to performance optimization
- Fix MTP row minimum card count to 4卡 to match the script's 8 NPUs


fix(gitleaks): drop /mnt/ defaults in NPU scripts

# 🔒 Security

## Remove hardcoded internal paths from Qwen3.5 NPU scripts

- Replace /mnt/tidalfs-hwwh01 EXP_DIR defaults with the repo-standard `${SCRIPT_DIR}/../../../../exps` in run_qwen35-35B-pokemon-sft-mtp-8xnpu.sh, run_qwen35-35B-A3B-16xnpu-colocate-thd.sh and run_qwen35_9B_mtp_8xnpu_thd.sh
- Replace bare /mnt/ placeholder in run-qwen35-9B-16xnpu-cp.sh with the same repo-standard default
- Switch the SFT script DATA_DIR default to the `${DATA_DIR:-${EXP_DIR}}` convention used by all other training scripts

---

# 🎨 Style

## End-of-file fixer

- Remove trailing blank line from docker/Dockerfile.npu

Co-Authored-By: Claude <noreply@anthropic.com>

update Dockerfile.npu

optimizer npu update_weights

Add optimized 9B&35B NPU colocate script

Add optimized 9B&35B NPU sglang patch

Add optimized 9B&35B NPU sgl-kernel-npu patch

FIX dockerfile build error

Fix sgl-kernel-npu patch

[clean code] sglang&kernel patch
@Li-Yongwen
!70 merge ascend-dev-0908-pr into ascend-dev-0908
597b703 
feat(NPU): add latest sglang optimize

Created-by: llxll
Commit-by: liyongwen;jufanbo;wuqiwei;dabuliu123;lixionglong;Tgz27
Merged-by: llxll
Description: ## What

add latest sglang optimize &  clean code for sglang

## Why

<!-- Why are these changes needed? Link related issues with "Fixes redai-studio#123" or "Relates to #456". -->

## How

<!-- How do the changes work? Describe the technical approach. -->

## Testing

<!-- How were the changes tested? Include commands, test results, or screenshots. -->

- [ ] `pre-commit run --all-files` passes
- [ ] Tests pass (`pytest tests/`)
- [ ] New tests added (if applicable)
- [ ] Documentation updated (if applicable)

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation update
- [ ] Refactoring (no functional changes)
- [ ] Performance improvement
- [ ] CI/CD or build changes

## Screenshots / Logs

<!-- If applicable, add screenshots or log output to help explain the changes. -->


See merge request: hw-pbclouds/Relax!70
@Lw135 Lw135 changed the title Ascend dev 0908 [WIP]SGLang NPU optimizations for Qwen3.5 9B & 35B Sep 9, 2026
rai-studio-bot
rai-studio-bot reviewed Sep 9, 2026
View reviewed changes

@rai-studio-bot rai-studio-bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

未发现需要阻塞合入的确定性问题；脚本配置问题及修复建议已放在行级评论中。

两份补丁可应用到镜像固定的上游版本，Python 3.10–3.12 CI 测试均通过；pre-commit 仍因两份补丁的末尾换行失败。本次未运行算子数值及多节点训练验证，审查环境缺少所需 NPU 硬件。

Powered by Nyanpasu with gpt-6-astra medium, please check the suggestions carefully.

Comment thread scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh
--sglang-cuda-graph-bs 4 8 16 24 32 40 48 64 128
--sglang-device npu
# --sglang-disable-radix-cache
--mamba-scheduler-strategy extra_buffer

@rai-studio-bot rai-studio-bot Sep 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 优先级：P2

请为该参数补上 sglang- 前缀。relax/backends/sglang/arguments.py 的 new_add_argument_wrapper 只注册带此前缀的选项；当前写法会被 parse_known_args() 留在未知参数中，后续 Megatron 解析又允许忽略未知参数，因此指定的 extra_buffer 没有传给推理引擎。9B 脚本已使用正确写法。

Suggested change
--mamba-scheduler-strategy extra_buffer
--sglang-mamba-scheduler-strategy extra_buffer

Comment thread scripts/training/text/run-qwen35-35B-A3B-16xnpu-colocate.sh
"${MISC_ARGS[@]}" 2>&1 | tee log/qwen35-35B-MATH-gpu16-sync-${now}.log
ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://${MASTER_ADDR}:8265" \
${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \
--runtime-env-json="${RUNTIME_ENV_JSON}" \

@rai-studio-bot rai-studio-bot Sep 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 优先级：P2

请在提交前将新增优化开关合入任务的 env_vars，两份脚本都需要覆盖。连接已有 Ray 集群时，local-npu.sh 会委托 ray-job-npu.sh；后者的 RUNTIME_ENV_JSON 只包含固定名单，没有这里新增的 SGLANG_NPU_*、SGLANG_GMM2_TRITON 和 MINDSPEED_BRIDGE_GDN_USE_TORCH_CONV 等变量。提交 shell 中的 export 不会更新已运行的 raylet 环境，导致该启动路径下部分优化实际未启用。可在环境初始化完成后、ray job submit 前按以下方式合并：

RUNTIME_ENV_JSON="$(python3 - <<'PY'
import json, os
env = json.loads(os.environ["RUNTIME_ENV_JSON"])
extra = {
    "CPU_AFFINITY_CONF", "TORCH_HCCL_ZERO_COPY", "MULTI_STREAM_MEMORY_REUSE",
    "HCCL_OP_EXPANSION_MODE", "ASCEND_USE_FIA", "GDN_ATTN_BACKEND_TRITON",
    "STREAMS_PER_DEVICE", "TASK_QUEUE_ENABLE",
}
env.setdefault("env_vars", {}).update({
    k: v for k, v in os.environ.items()
    if k.startswith(("SGLANG_", "MINDSPEED_BRIDGE_")) or k in extra
})
print(json.dumps(env))
PY
)"

Comment thread scripts/training/text/run-qwen35-9B-8xnpu-colocate.sh
fi

source "${MODEL_CONFIG_DIR}/qwen35-9B.sh"
EXP_DIR="${EXP_DIR:-/root/exps}"

@rai-studio-bot rai-studio-bot Sep 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 优先级：P2

请为新引入的 MODEL_DIR 和 DATA_DIR 提供默认值。该脚本只初始化了 EXP_DIR，而 local-npu.sh、ray-job-npu.sh 和模型配置均未设置另外两个变量。按 Usage 直接运行且未额外导出它们时，checkpoint 会展开为 /Qwen3.5-9B/，训练数据会展开为 /dapo-math-17k/dapo-math-17k.jsonl；即使数据已按现有示例放在 EXP_DIR 下，也无法找到。可保留覆盖能力并默认沿用实验目录：

Suggested change
EXP_DIR="${EXP_DIR:-/root/exps}"
EXP_DIR="${EXP_DIR:-/root/exps}"
MODEL_DIR="${MODEL_DIR:-${EXP_DIR}}"
DATA_DIR="${DATA_DIR:-${EXP_DIR}}"

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Yangruipis Yangruipis Awaiting requested review from Yangruipis Yangruipis is a code owner
@NINGBENZHE NINGBENZHE Awaiting requested review from NINGBENZHE NINGBENZHE is a code owner
1 more reviewer
@rai-studio-bot rai-studio-bot rai-studio-bot left review comments
Reviewers whose approvals may not affect merge requirements

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@Lw135 @rai-studio-bot @Li-Yongwen