cf1b701afe
唯一改动: computility-run.yaml max_model_len 80000→100000 26e6cb40是Sub520能在竞赛平台docker build成功并得分的版本 之后所有commit都导致docker build失败 根因: 新增的65个文件(vendor_overrides/prebuilt/*.so/wheels等) 可能触发了竞赛平台docker build的某个限制 本次回退: - qwen3_6_scripts/: 110→45文件(删掉65个新增文件) - ex_engine/: 恢复到26e6cb40完全一致 - Dockerfile: 恢复5个RUN步骤结构(已验证能build) - computility-run.yaml: max_model_len=100000(避免replay 400拒绝)
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)
|