Fix torch.compile crash: @torch.compiler.disable on all CuTeDSL run()

CuTeDSL internals (Path.cwd, threading, JIT) are incompatible with
torch.dynamo tracing. Marking run() as compiler-disabled makes the
runners opaque to torch.compile — they execute eagerly while the
rest of the model gets compiled.
This commit is contained in:
2026-05-18 21:07:35 +00:00
parent a94011ec92
commit 85e1cd3b69
3 changed files with 3 additions and 0 deletions

View File

@@ -110,6 +110,7 @@ class CuTeDSLNvfp4Linear:
self._activation_global_scale = gs
@torch.compiler.disable
def run(self, hidden_states: torch.Tensor) -> torch.Tensor:
"""Forward: BF16 input → NVFP4 GEMM → BF16 output."""
self._ensure_initialized()

View File

@@ -274,6 +274,7 @@ class CuTeDSLSharedExpertRunner:
return out[:num_tokens]
@torch.compiler.disable
def run(self, hidden_states: torch.Tensor) -> torch.Tensor:
"""Full shared expert forward: L1 → SiLU → L2 → output."""
self._ensure_initialized()

View File

@@ -374,6 +374,7 @@ class CuTeDSLMoERunner:
@torch.compiler.disable
def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None):
"""Run the NVFP4 MoE forward pass.