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:
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user