claude
|
a50adefdfc
|
feat: xllm MoE CUDA kernels — fused_topk + compute_index + combine
3 MoE kernel files adapted for corex:
moe_fused_topk.cu: LOG(FATAL)→TORCH_CHECK, +torch/extension.h
moe_compute_index.cu: CHECK_LE→TORCH_CHECK, uses cub::BlockScan (corex CUB)
moe_combine.cu: fixed duplicate include, +torch/extension.h
New pybind binding: xllm_moe_bind.cpp
→ moe_fused_topk(gating, topk, renormalize, bias, scoring_func)
→ moe_compute_index(expert_id, num_experts)
→ moe_combine_result(gemm2, weights, N, topk)
AST verification added for all 3 functions
|
2026-08-14 11:23:49 +00:00 |
|
claude
|
7fc0c1defa
|
test: AST call chain verification for all 4 xllm kernel .so
verify_ast_chain.py tests every exported function:
xllm_norm.so: rms_norm (vs PyTorch), fused_add_rms_norm (vs PyTorch)
xllm_activation.so: silu_and_mul (vs F.silu), gelu_and_mul (vs F.gelu)
xllm_rope.so: rotary_embedding (verify rotation applied)
xllm_cache.so: reshape_paged_cache (verify slot write)
vs_ixformer: compare our .so output vs base image ixformer
All 4 .so compiled successfully:
xllm_norm.so 1.4MB [rms_norm, fused_add_rms_norm]
xllm_activation.so 1.3MB [silu_and_mul, gelu_and_mul, act_and_mul]
xllm_rope.so 1.2MB [rotary_embedding]
xllm_cache.so 1.3MB [block_copy, reshape_paged_cache]
|
2026-08-14 11:14:58 +00:00 |
|