Lazy load TileLang MHC kernels

[mmangkad]

Summary

This keeps TileLang out of the DeepSeek V4 import path until we actually need it. Model-registry discovery imports a lot of model files, so it should not also load TileLang's native CUDA stubs just because one unrelated module was discovered.

deepseek_v4_rope.py only uses Triton kernels, so it no longer imports TileLang at module load. mhc.py uses a small lazy proxy, which lets the existing @tilelang.jit(...) definitions stay in place while delaying the real TileLang import until the first MHC TileLang kernel call. The split-k MHC path loads TileLang before using T.dynamic(...) / T.Tensor[...], and first-load / first-compile paths are guarded with locks.

---

CI States

Latest PR Test (Base): ✅ Run #29010339282
Latest PR Test (Extra): ❌ Run #29010339047

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[mmangkad]

/rerun-test test/registered/kernels/test_mhc_kernels.py test/registered/models_e2e/test_deepseek_v4*

[github-actions]

Results for /rerun-test test/registered/kernels/test_mhc_kernels.py test/registered/models_e2e/test_deepseek_v4*:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/test_mhc_kernels.py

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py
cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp4_megamoe_b200.py

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp4_h200.py
cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp8_h200.py

[hirakunaramuka2]

@merrymercy @Ying1123 @Fridge003 @ispobock @HaiShaw @ch-wan @BBuf @Edwardf0t1 — could you please take a look at this PR and merge it? It fixes the FlashInfer allreduce fusion regression introduced by the flashinfer 0.6.14 upgrade (#30875).

All CUDA CI checks pass (base-a through base-c, 1-gpu/2-gpu/4-gpu/8-gpu on H100/H200/B200/GB300 all green). The 15 failures are all on non-CUDA platforms (NPU/XPU/AMD) which are unrelated. GitHub shows the PR as mergeable.

Additionally, the MHC kernels and DeepSeek V4 e2e tests were explicitly re-run and pass:

• test/registered/kernels/test_mhc_kernels.py — ✅ (1-gpu-h100)
• test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py — ✅ (4-gpu-b200)
• test/registered/models_e2e/test_deepseek_v4_flash_fp4_megamoe_b200.py — ✅ (4-gpu-b200)
• test/registered/models_e2e/test_deepseek_v4_flash_fp4_h200.py — ✅ (8-gpu-h200)
• test/registered/models_e2e/test_deepseek_v4_flash_fp8_h200.py — ✅ (8-gpu-h200)

I also verified on lambda-h100-01 (single-node 8×H100 SM90) with TP=8 and --flashinfer-allreduce-fusion-backend trtllm that after applying this fix, workspace initialization succeeds on all ranks:

[2026-07-11 12:54:48 TP0] FlashInfer AllReduce Fusion enabled and workspace initialized: backend=trtllm, rank=0, world_size=8, max_token_num=2048, hidden_dim=6144
[2026-07-11 12:54:48 TP1] FlashInfer AllReduce Fusion enabled and workspace initialized: backend=trtllm, rank=1, world_size=8, max_token_num=2048, hidden_dim=6144
[2026-07-11 12:54:48 TP2] FlashInfer AllReduce Fusion enabled and workspace initialized: backend=trtllm, rank=2, world_size=8, max_token_num=2048, hidden_dim=6144
[2026-07-11 12:54:48 TP3] FlashInfer AllReduce Fusion enabled and workspace initialized: backend=trtllm, rank=3, world_size=8, max_token_num=2048, hidden_dim=6144
[2026-07-11 12:54:48 TP4] FlashInfer AllReduce Fusion enabled and workspace initialized: backend=trtllm, rank=4, world_size=8, max_token_num=2048, hidden_dim=6144
[2026-07-11 12:54:48 TP5] FlashInfer AllReduce Fusion enabled and workspace initialized: backend=trtllm, rank=5, world_size=8, max_token_num=2048, hidden_dim=6144
[2026-07-11 12:54:48 TP6] FlashInfer AllReduce Fusion enabled and workspace initialized: backend=trtllm, rank=6, world_size=8, max_token_num=2048, hidden_dim=6144
[2026-07-11 12:54:48 TP7] FlashInfer AllReduce Fusion enabled and workspace initialized: backend=trtllm, rank=7, world_size=8, max_token_num=2048, hidden_dim=6144

Note: there's a duplicate PR #30870 targeting the same issue, but this PR (#30580) is a better fix — it makes mhc.py itself lazy (defense in depth) rather than only deferring the import in deepseek_v4.py. See the detailed analysis in #30875 (comment).

[hirakunaramuka2]

Thank you @mmangkad for the fix and @mickqian for merging it! This resolves the FlashInfer allreduce fusion regression from the flashinfer 0.6.14 upgrade (#30875).

Confirmed merged on main as commit 14bef7cd. Verified the lazy-loading proxy is in place:

• deepseek_v4_rope.py no longer imports tilelang at module load
• mhc.py uses _LazyTilelang proxy, deferring real tilelang import until MHC kernel execution

Closing the related issue as fixed.