dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

build(SM70): precompile GDN CUDA kernel to .so during docker build

Browse Source 
precompile_gdn.py: calls torch.utils.cpp_extension.load with build_directory
to produce .so at build time. If build env has no GPU/compiler, fails
gracefully — kernel JIT compiles at runtime instead.

fused_fwd.py: _load_ext() now checks build/ dir for precompiled .so first,
skips 2-minute JIT compilation if found.
This commit is contained in:
Claude
2026-08-10 01:08:38 +00:00
parent 8cf73ad39c
commit 20cd2d8904
3 changed files with 75 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
qwen3_6_scripts/patch_ops.sh
View File

@@ -188,9 +188,12 @@ if [ -d "./flash_qla_sm70" ]; then
rm -rf "$FLASH_QLA_DST" 2>/dev/null
cp -r ./flash_qla_sm70 "$FLASH_QLA_DST" 2>/dev/null && \
echo "[patch_ops] flash_qla_sm70 deployed to $FLASH_QLA_DST" || true
# Pre-compile CUDA kernel → .so (skipped if no GPU/compiler at build time)
python3 ./precompile_gdn.py "$FLASH_QLA_DST" 2>&1 || \
echo "[patch_ops] WARNING: precompile failed — kernel will JIT at runtime"
# Also deploy to VLLM2 if present
if [ -n "$VLLM2" ]; then
rm -rf "$VLLM2/model_executor/models/flash_qla_sm70" 2>/dev/null
cp -r ./flash_qla_sm70 "$VLLM2/model_executor/models/flash_qla_sm70" 2>/dev/null || true
cp -r "$FLASH_QLA_DST" "$VLLM2/model_executor/models/flash_qla_sm70" 2>/dev/null || true
fi
fi