70c898ac8b
真机验证发现的问题: 1. qwen3_5.py 做 'from ex_engine.python.ix_bridge' 但包结构是 ex_engine.ix_bridge → 创建 python/ 子目录 + symlinks 2. flash_qla_sm70 只部署到 /workspace 没有到 vllm models 目录 → 显式 cp -r 到 VLLM/model_executor/models/ 3. 从 vllm v0.5.5 搬 MoE CUDA kernels (torch::Tensor API): - topk_softmax_kernels.cu (506行, CUB BlockReduce) - moe_align_block_size_kernels.cu (134行) - moe_pybind.cpp (pybind11 入口) 真机验证结果: ✓ ix_bridge import OK, available=True ✓ topk_softmax (64 experts, top8) OK — CUDA kernel 命中 ✓ ix_full_bridge silu_and_mul OK ✓ qwen3_5.py import OK ✓ ex_engine build 2/2 factors ✓ moe_topk_softmax_v3.so 编译成功 ✓ flash_qla_sm70_gdn_strided.so 编译成功 ✗ 单卡 32GB OOM (正常, 竞赛 4卡 tp=4)
96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
precompile_moe_kernels.py — JIT compile vllm v0.5.5 MoE CUDA kernels for BI-V100.
|
|
|
|
Produces: moe_kernels.so with:
|
|
- topk_softmax(topk_weights, topk_indices, token_expert_indices, gating_output)
|
|
- moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids, expert_ids, num_tokens_post_pad)
|
|
|
|
Usage:
|
|
python3 precompile_moe_kernels.py # JIT compile
|
|
python3 precompile_moe_kernels.py --test # compile + smoke test
|
|
"""
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
def compile_moe_kernels():
|
|
"""JIT compile MoE CUDA kernels via torch.utils.cpp_extension."""
|
|
import torch
|
|
from torch.utils.cpp_extension import load
|
|
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
moe_dir = os.path.join(script_dir, 'csrc', 'moe_v055')
|
|
|
|
sources = [
|
|
os.path.join(moe_dir, 'moe_pybind.cpp'),
|
|
os.path.join(moe_dir, 'topk_softmax_kernels.cu'),
|
|
os.path.join(moe_dir, 'moe_align_block_size_kernels.cu'),
|
|
]
|
|
|
|
for s in sources:
|
|
if not os.path.isfile(s):
|
|
raise FileNotFoundError(f"Missing: {s}")
|
|
|
|
print(f"[moe_kernels] Compiling from {moe_dir}")
|
|
t0 = time.time()
|
|
|
|
mod = load(
|
|
name='moe_kernels',
|
|
sources=sources,
|
|
extra_include_paths=[moe_dir],
|
|
extra_cflags=['-O2', '-std=c++17'],
|
|
extra_cuda_cflags=['-O2', '--expt-relaxed-constexpr'],
|
|
verbose=True,
|
|
)
|
|
|
|
dt = time.time() - t0
|
|
funcs = [x for x in dir(mod) if not x.startswith('_')]
|
|
print(f"[moe_kernels] Compiled in {dt:.1f}s — functions: {funcs}")
|
|
return mod
|
|
|
|
|
|
def smoke_test(mod):
|
|
"""Quick functional test of compiled kernels."""
|
|
import torch
|
|
|
|
print("\n=== Smoke test ===")
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
if device == 'cpu':
|
|
print(" SKIP: no CUDA device")
|
|
return
|
|
|
|
# Test topk_softmax
|
|
num_tokens, num_experts, topk = 4, 8, 2
|
|
gating = torch.randn(num_tokens, num_experts, device=device, dtype=torch.float32)
|
|
topk_weights = torch.empty(num_tokens, topk, device=device, dtype=torch.float32)
|
|
topk_indices = torch.empty(num_tokens, topk, device=device, dtype=torch.int32)
|
|
token_expert_indices = torch.empty(num_tokens, topk, device=device, dtype=torch.int32)
|
|
|
|
mod.topk_softmax(topk_weights, topk_indices, token_expert_indices, gating)
|
|
|
|
print(f" topk_softmax: weights={topk_weights.shape}, NaN={topk_weights.isnan().any()}")
|
|
print(f" weights[0] = {topk_weights[0].tolist()}")
|
|
print(f" indices[0] = {topk_indices[0].tolist()}")
|
|
|
|
# Test moe_align_block_size
|
|
block_size = 4
|
|
max_num_tokens_padded = (num_tokens * topk + num_experts * block_size)
|
|
sorted_ids = torch.empty(max_num_tokens_padded, device=device, dtype=torch.int32)
|
|
expert_ids = torch.empty(max_num_tokens_padded // block_size, device=device, dtype=torch.int32)
|
|
num_tokens_post_pad = torch.empty(1, device=device, dtype=torch.int32)
|
|
|
|
mod.moe_align_block_size(topk_indices, num_experts, block_size,
|
|
sorted_ids, expert_ids, num_tokens_post_pad)
|
|
|
|
print(f" moe_align: sorted_ids[:8]={sorted_ids[:8].tolist()}, "
|
|
f"num_post_pad={num_tokens_post_pad.item()}")
|
|
|
|
print("\n ✓ All smoke tests passed")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
mod = compile_moe_kernels()
|
|
if '--test' in sys.argv:
|
|
smoke_test(mod)
|