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

feat(EX): corex_fa2.py — third dlopen module from comp 168 AST chain

Browse Source 
Log analysis from dockerrizhi.txt (07-23 Sub168 run) reveals THREE
corex modules, not two:

  1. corex_gdn.py — GatedDeltaNet fused kernel (already implemented)
  2. corex_moe.py — MoE routing + expert GEMM (already implemented)
  3. corex_fa2.py — FlashAttention2 dispatch (NEW)

corex_fa2.py handles 32/36 attention layers with three modes:
  :333 → FA2 packed prefill (B=2 Hq=4 Hkv=1 D=256 max_q=2048)
  :507 → FA2 paged chunked prefill (B=1 max_q=17 cache_blocks=2)
  :225 → FA2 paged decode (B=1 max_k=45455 partition=256)

Wraps ixformer.contrib.vllm_flash_attn + ixf_F.vllm_single_query_cached_kv_attention.
These .so files EXIST in the base image (libixattn.so).

Also: wired corex_fa2 import into qwen3_5.py + deploy script.
This commit is contained in:
project6-dev
2026-08-10 04:01:17 +00:00
parent f4e2264a83
commit 5efb0fcc35
3 changed files with 389 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
ex_engine/deploy_corex_modules.sh
View File

@@ -60,13 +60,24 @@ else
echo "[COREX] ✓ corex_moe.py already exists (base image or prior deploy)"
fi
# Deploy corex_fa2.py
if [[ ! -f "${VLLM_MODELS}/corex_fa2.py" ]]; then
cp "${SRC_DIR}/corex_fa2.py" "${VLLM_MODELS}/corex_fa2.py"
echo "[COREX] ✓ Deployed corex_fa2.py"
else
echo "[COREX] ✓ corex_fa2.py already exists (base image or prior deploy)"
fi
# Also deploy to ex_engine location (backup import path)
mkdir -p /workspace/ex_engine/python 2>/dev/null || true
cp "${SRC_DIR}/corex_gdn.py" /workspace/ex_engine/python/ 2>/dev/null || true
cp "${SRC_DIR}/corex_moe.py" /workspace/ex_engine/python/ 2>/dev/null || true
cp "${SRC_DIR}/corex_fa2.py" /workspace/ex_engine/python/ 2>/dev/null || true
echo "[COREX] Deploy complete"
echo "[COREX] Expected log on startup:"
echo " corex_gdn.py:NN → Loaded fused CoreX GDN decode operator ..."
echo " corex_gdn.py:NN → Using fused CoreX GDN prefill operator"
echo " corex_moe.py:NN → Using CoreX fused MoE prefill operator: tokens=N, kernel=expert-grouped-wmma"
echo " corex_fa2.py:NN → Using CoreX FA2 packed prefill: B=N Hq=4 Hkv=1 D=256 ..."
echo " corex_fa2.py:NN → Using CoreX paged decode: B=N Hq=4 Hkv=1 D=256 ..."