Skip to content

Commit b0a9c07

Browse files
appleinskyappleinsky
authored
Merge pull request #14 from 1190201122/qwen36_dense_moe-pto-path
Qwen36 dense moe change pto path
2 parents 8a3df63 + d488559 commit b0a9c07

13 files changed

Lines changed: 68 additions & 36 deletions

File tree

‎.gitignore‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,7 @@ csrc/output/
218218
vllm_plugin_fl.egg-info/
219219
/vllm_fl/_cann_ops_custom/*
220220
!/vllm_fl/_cann_ops_custom/.gitkeep
221-
vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/*.so
221+
vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/kernels/compiled_lib/*.so
222222

223223
# CANN framework operator build directories
224224
csrc/ascend/build/

‎csrc/ascend/pto_chunk_gdn/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ if(NOT DEFINED Python_SITELIB)
7575
)
7676
endif()
7777
set(PTO_COMPILED_INSTALL_DIR
78-
"${Python_SITELIB}/vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib"
78+
"${Python_SITELIB}/vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/kernels/compiled_lib"
7979
)
8080

8181
# ── Common compiler flags ──────────────────────────────────────────────────────

‎docs/qwen3_6_ascend_optimization.md‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ FL 仓库已在 `csrc/ascend` 编译注册到 `torch.ops._C_ascend`（`TORCH_LIB
9999

100100
### 4.1 接入
101101

102-
参照 vllm-ascend PR #8872（`/workspace/patch/pto_gdn.patch`）。FL 已内置 `vllm_fl/ops/pto_chunk_gdn/`（6 阶段融合：cumsum → scaled_dot_kkt → solve_tril → wy_fast → chunk_h → chunk_o，单次 launch）。
102+
参照 vllm-ascend PR #8872（`/workspace/patch/pto_gdn.patch`）。FL 已内置 `vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/`（6 阶段融合：cumsum → scaled_dot_kkt → solve_tril → wy_fast → chunk_h → chunk_o，单次 launch）。
103103

104104
适配（`chunk_gated_delta_wrapper.py`）：l2norm/run_mega_kernel 导入改 FL 路径；`_triton()` 回退调用去掉 FL chunk 不支持的 `prebuilt_meta`。`compile.py` 增加 `.so` 已存在则跳过（避免 4 个 TP worker 每次启动并发重编）。
105105

@@ -124,7 +124,7 @@ Qwen3.6-27B TP4 → 每 rank `H=NV/TP=12, Hg=NK/TP=4`，而 `mega_transpose_TH_t
124124
- 数值：PTO vs Triton chunk，输出 o 与 final_state 一致（fp16 计算、bf16 往返）。
125125
- eager 自测：**TTFT +15%~95%**；A/B 输出逐 token 全等。
126126
- aclgraph 自测：**TTFT +18%~84%**，6 case 全过。
127-
- kernel `.so` 按 (H,Hg,D,C) 缓存于 `vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/`，新配置首用编译 ~9s。
127+
- kernel `.so` 按 (H,Hg,D,C) 缓存于 `vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/kernels/compiled_lib/`，新配置首用编译 ~9s。
128128

129129
---
130130

@@ -250,9 +250,9 @@ A/B 正确性：AscendC vs Triton 输出逐 token 全等；PTO vs 无 PTO 输出
250250
| `vllm_fl/dispatch/backends/vendor/ascend/patch.py` | 注册 `patch_qwen3_6_gdn` |
251251
| `vllm_fl/dispatch/backends/vendor/ascend/patches/README.md` | patch 条目 |
252252
| `vllm_fl/__init__.py` | `_bootstrap_cann_custom_op_env`（OPP 路径 + RTLD_LOCAL 预载） |
253-
| `vllm_fl/ops/pto_chunk_gdn/chunk_gated_delta_wrapper.py` | FL 导入适配、`_triton` 参数适配 |
254-
| `vllm_fl/ops/pto_chunk_gdn/compile.py` | `.so` 存在则跳过编译 |
255-
| `vllm_fl/ops/pto_chunk_gdn/mega_kernel.py` | `total_chunks` 可选参数 |
253+
| `vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/chunk_gated_delta_wrapper.py` | FL 导入适配、`_triton` 参数适配 |
254+
| `vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/compile.py` | `.so` 存在则跳过编译 |
255+
| `vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/mega_kernel.py` | `total_chunks` 可选参数 |
256256
| `csrc/ascend/pto_chunk_gdn/mega_kernel.cpp` | 头维 32B 补齐 |
257257
| `csrc/ascend/pto_chunk_gdn/chunk_h.cpp` | fs UB 写竞争排干 |
258258
| `vllm_fl/worker/model_runner.py` | `_update_hybrid_attention_mamba_layout` 跳过 re-stride |

‎tests/custom_ops_tests/README.md‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@
2626
- **依赖头库**：`csrc/ascend/third_party/pto-isa/`。
2727
- **产物**：多个 `mega_kernel_H*_Hg*_D*_C*.so`。
2828
- **安装位置**：
29-
- 预编译模式：安装到 Python site-packages 下的 `vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/`；
30-
- JIT 模式：首次调用时由 `vllm_fl/ops/pto_chunk_gdn/compile.py` 自动编译并缓存到同一目录。
29+
- 预编译模式：安装到 Python site-packages 下的 `vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/kernels/compiled_lib/`；
30+
- JIT 模式：首次调用时由 `vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/compile.py` 自动编译并缓存到同一目录。
3131
- **运行时加载**：Python 代码通过 `ctypes.CDLL` / `torch.ops.load_library` 直接加载 `.so`，不经过 CANN `opp/vendors` 路径。
3232

3333
## 2. 环境准备
@@ -171,28 +171,28 @@ cmake --build build/temp.linux-aarch64-cpython-311 \
171171
产物会安装到当前 Python 环境 site-packages 下的：
172172

173173
```text
174-
vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/
174+
vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/kernels/compiled_lib/
175175
├── mega_kernel_H16_Hg8_D128_C128.so
176176
├── mega_kernel_H16_Hg16_D128_C128.so
177177
└── ...
178178
```
179179

180-
> 在 editable install（`pip install -e .`）下，`_PACKAGE_ROOT` 等于仓库根目录，因此也会写到仓库内的 `vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/`。
180+
> 在 editable install（`pip install -e .`）下，`_PACKAGE_ROOT` 等于仓库根目录，因此也会写到仓库内的 `vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/kernels/compiled_lib/`。
181181
182182
### 方式 B：JIT 首次编译（开发调试用）
183183

184-
不预编译，直接运行 `tests/custom_ops_tests/test_pto_chunk_gdn.py`。`vllm_fl/ops/pto_chunk_gdn/compile.py` 会：
184+
不预编译，直接运行 `tests/custom_ops_tests/test_pto_chunk_gdn.py`。`vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/compile.py` 会：
185185

186186
1. 自动查找 `csrc/ascend/third_party/pto-isa`；
187187
2. 调用系统 `bisheng` 编译对应配置的 `mega_kernel_*.so`；
188-
3. 缓存到 `vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/`；
188+
3. 缓存到 `vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/kernels/compiled_lib/`；
189189
4. 后续调用直接复用缓存。
190190

191191
## 7. 目录结构总览
192192

193193
```text
194194
csrc/ascend/
195-
├── CMakeLists.txt # 构建 _C_ascend + 分发 pto_chunk_gdn
195+
├── CMakeLists.txt # 构建 _C_ascend + PTO GDN 预编译
196196
├── torch_binding.cpp # torch.ops._C_ascend 注册
197197
├── torch_binding_meta.cpp # meta kernel 注册
198198
├── camem_allocator.cpp # NPU 显存分配器

‎tests/custom_ops_tests/test_pto_chunk_gdn.py‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,39 @@
1616
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
1717
sys.path.insert(0, ROOT)
1818

19+
# ---------------------------------------------------------------------------
20+
# Load mega_kernel via importlib, bypassing
21+
# vllm_fl.dispatch.backends.vendor.ascend.impl.__init__ (which eagerly
22+
# imports modules that require a fully-initialized vLLM runtime).
23+
#
24+
# We register the module under a dotted package name so that the relative
25+
# import "from .compile import ..." inside mega_kernel.py resolves correctly
26+
# against sibling modules registered in sys.modules under the same package.
27+
# ---------------------------------------------------------------------------
28+
import importlib.util as _ilu
29+
30+
_PTO_PKG = "vllm_fl.dispatch.backends.vendor.ascend.impl.pto_chunk_gdn"
31+
_PTO_DIR = os.path.join(
32+
ROOT,
33+
"vllm_fl", "dispatch", "backends", "vendor", "ascend", "impl", "pto_chunk_gdn",
34+
)
35+
36+
37+
def _load_sibling(name, path):
38+
"""Load a Python file as a module of package _PTO_PKG."""
39+
spec = _ilu.spec_from_file_location(f"{_PTO_PKG}.{name}", path)
40+
mod = _ilu.module_from_spec(spec)
41+
mod.__package__ = _PTO_PKG
42+
sys.modules[f"{_PTO_PKG}.{name}"] = mod
43+
spec.loader.exec_module(mod)
44+
return mod
45+
46+
47+
# Load compile.py first (dependency of mega_kernel.py), then mega_kernel.py.
48+
_load_sibling("compile", os.path.join(_PTO_DIR, "compile.py"))
49+
_mega = _load_sibling("mega_kernel", os.path.join(_PTO_DIR, "mega_kernel.py"))
50+
run_mega_kernel = _mega.run_mega_kernel
51+
1952

2053
def main() -> int:
2154
torch.npu.config.allow_internal_format = True
@@ -24,8 +57,6 @@ def main() -> int:
2457

2558
enable_custom_op()
2659

27-
from vllm_fl.ops.pto_chunk_gdn.mega_kernel import run_mega_kernel
28-
2960
device = torch.device("npu:0")
3061
B, T, Hg, H, D = 1, 128, 8, 16, 128
3162

@@ -65,4 +96,4 @@ def main() -> int:
6596

6697

6798
if __name__ == "__main__":
68-
raise SystemExit(main())
99+
raise SystemExit(main())

‎vllm_fl/_version.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@
33
from __future__ import annotations
44
import re as _re
55

6-
__version__: str = "0.1.0+vllm0.13.0.g95f9bca10.d20260723"
6+
__version__: str = "0.0.0+g8a3df6354.d20260807"
77
version: str = __version__
8-
_node_date: str = "2026-07-23"
8+
_node_date: str = "2026-08-06"
99

1010
_sha_m = _re.search(r"(?:^|[.+-])g(?P<sha>[0-9a-f]{7,40})(?:[.+-]|$)", __version__)
1111
_date_m = _re.search(r"(?:^|[.+-])d(?P<ymd>\d{8})(?:[.+-]|$)", __version__)

vllm_fl/ops/pto_chunk_gdn/__init__.py renamed to vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/__init__.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
# See the License for the specific language governing permissions and
1515
# limitations under the License.
1616
#
17-
"""PTO chunk-GDN megakernel module for vLLM-Ascend.
17+
"""PTO chunk-GDN megakernel module — Ascend dispatch impl.
1818
1919
Provides a Bisheng-JIT-compiled fused Ascend NPU megakernel for the chunk
2020
GatedDeltaNet (GDN) recurrent layer used in Qwen3.5 / Qwen3.6 models,
@@ -25,9 +25,9 @@
2525
2626
The decode phase always uses the original Triton implementation.
2727
"""
28-
from vllm_fl.ops.pto_chunk_gdn.compile import BLOCK_DIM, PTO_LIB_PATH
29-
from vllm_fl.ops.pto_chunk_gdn.mega_kernel import run_mega_kernel
30-
from vllm_fl.ops.pto_chunk_gdn.worker_hook import apply_pto_gdn_patch, is_pto_gdn_patch_active
28+
from .compile import BLOCK_DIM, PTO_LIB_PATH
29+
from .mega_kernel import run_mega_kernel
30+
from .worker_hook import apply_pto_gdn_patch, is_pto_gdn_patch_active
3131

3232
__all__ = [
3333
"BLOCK_DIM",

vllm_fl/ops/pto_chunk_gdn/chunk_gated_delta_wrapper.py renamed to vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/chunk_gated_delta_wrapper.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ def _triton():
105105
except Exception:
106106
pass
107107

108-
from vllm_fl.ops.pto_chunk_gdn.mega_kernel import run_mega_kernel
108+
from .mega_kernel import run_mega_kernel
109109

110110
kh = q.shape[2]
111111
cu32 = cu_seqlens.to(torch.int32).contiguous()

vllm_fl/ops/pto_chunk_gdn/compile.py renamed to vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/compile.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
"""Bisheng JIT compilation for the PTO GDN megakernel on Ascend NPU.
1818
1919
The megakernel is compiled on first use and cached under
20-
``vllm_fl/ops/pto_chunk_gdn/kernels/compiled_lib/``.
20+
``vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/kernels/compiled_lib/``.
2121
Re-compilation is triggered when the C++ source mtime changes.
2222
2323
Environment variables:
@@ -37,13 +37,14 @@
3737
from pathlib import Path
3838

3939
import torch
40+
import vllm_fl
4041

4142
# ---------------------------------------------------------------------------
42-
# Paths — resolved relative to this file's installed location
43+
# Paths — resolved via vllm_fl package location
4344
# ---------------------------------------------------------------------------
4445
_THIS_DIR = Path(__file__).resolve().parent
45-
_VLLM_FL_DIR = _THIS_DIR.parent.parent # vllm_fl/
46-
_PACKAGE_ROOT = _VLLM_FL_DIR.parent # site-packages root
46+
_VLLM_FL_DIR = Path(vllm_fl.__file__).resolve().parent # vllm_fl/
47+
_PACKAGE_ROOT = _VLLM_FL_DIR.parent # repo root (editable) or site-packages
4748

4849
# C++ sources live in csrc/ascend/pto_chunk_gdn/ inside the source tree.
4950
# In editable installs (pip install -e .) _PACKAGE_ROOT == repo root.

vllm_fl/ops/pto_chunk_gdn/mega_kernel.py renamed to vllm_fl/dispatch/backends/vendor/ascend/impl/pto_chunk_gdn/mega_kernel.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
2626
Usage::
2727
28-
from vllm_ascend.ops.pto_chunk_gdn.mega_kernel import run_mega_kernel
28+
from vllm_fl.dispatch.backends.vendor.ascend.impl.pto_chunk_gdn.mega_kernel import run_mega_kernel
2929
3030
o = run_mega_kernel(q, k, v, g, beta, cu_seqlens, stream=stream,
3131
chunk_size=128, scale=head_dim**-0.5, key_heads=Hg)
@@ -38,7 +38,7 @@
3838

3939
import torch
4040

41-
from vllm_fl.ops.pto_chunk_gdn.compile import BLOCK_DIM, KERNELS_PTO, compile_mega_kernel
41+
from .compile import BLOCK_DIM, KERNELS_PTO, compile_mega_kernel
4242

4343

4444
# ---------------------------------------------------------------------------

0 commit comments

Comments
 (0)