biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

router: catch CuTeDSL warmup failures fast, don't let MLIR errors slow down init

Browse Source
This commit is contained in:
biondizzle
2026-06-01 00:00:07 +00:00
parent 56dff8d185
commit be476b2ce2
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
dsv4/ops/router.py
View File

@@ -36,11 +36,15 @@ def warmup_router_compilation(router) -> None:
""" """
if router.mode == "dense": if router.mode == "dense":
# Dummy forward at small N triggers decode-path compile. # Dummy forward at small N triggers decode-path compile.
# CuTeDSL fused kernel is WIP — falls through to prefill path.
dummy = torch.zeros( dummy = torch.zeros(
1, router.hidden_size, 1, router.hidden_size,
dtype=torch.bfloat16, device=router.device, dtype=torch.bfloat16, device=router.device,
) )
router._run_dense_impl(dummy) try:
router._run_dense_impl(dummy)
except Exception:
pass # CuTeDSL kernel not yet working; prefill path is fine
else: else:
dummy = torch.zeros(1, dtype=torch.int32, device=router.device) dummy = torch.zeros(1, dtype=torch.int32, device=router.device)
router._run_hash_impl(dummy) router._run_hash_impl(dummy)