EX Engine
8e6adf20e6
refactor(EX): upstream-aligned kernels + FlashQLA GDN backend
Major changes based on upstream_ref analysis:
1. factor_moe_topk_softmax.cu v2.0: Rewritten using ds_vllm/TRT-LLM
warp shuffle pattern (from topk_softmax_kernels.cu). Key differences:
- Zero shared memory (all butterfly __shfl_xor_sync)
- VPT=2, THREADS_PER_ROW=32 (1 warp per token row)
- 4 warps per CTA (4 tokens per block)
- Iterative argmax with winner suppression for top-K
- NaN/Inf clamping to 0 (prevents duplicate expert IDs)
2. GDN: FlashQLA backend (PROVEN on real BI-V100):
- Compiles with corex clang/16 --cuda-gpu-arch=ivcore10
- Real test: NaN=False on gdn_forward(B=1, T=64, H=4, K=128)
- Replaces custom factor_gdn_chunk_fwd.cu (archived to .ref)
- patch_model.py now JIT-loads FlashQLA extension at runtime
3. build.sh: Correct corex flags from real compile log:
--cuda-gpu-arch=ivcore10 (NOT sm_70)
-D__ILUVATAR__ -D__ILUVATAR_WORKAROUND__ -D__ILUVATAR_DIAG__
-cl-single-precision-constant -mllvm --bonus-inst-threshold=0
Key insight from xllm/kernels/ilu/ixformer.h:
ixformer::infer::topk_softmax() EXISTS at C++ level but Python
ixformer.functions binding is missing. Our .so factor bypasses
the missing Python binding entirely via dlopen/ctypes.