From 85e1cd3b6948f5cbfc7d9c2635b1cf51810776f5 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Mon, 18 May 2026 21:07:35 +0000
Subject: [PATCH] Fix torch.compile crash: @torch.compiler.disable on all
 CuTeDSL run()
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 cutedsl/nvfp4_linear.py           | 1 +
 cutedsl/shared_expert_pipeline.py | 1 +
 vllm/nvfp4_cutedsl.py             | 1 +
 3 files changed, 3 insertions(+)

diff --git a/cutedsl/nvfp4_linear.py b/cutedsl/nvfp4_linear.py
index 7d36bf8f..2119176f 100644
--- a/cutedsl/nvfp4_linear.py
+++ b/cutedsl/nvfp4_linear.py
@@ -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()
diff --git a/cutedsl/shared_expert_pipeline.py b/cutedsl/shared_expert_pipeline.py
index 99913f1f..224ca1c7 100644
--- a/cutedsl/shared_expert_pipeline.py
+++ b/cutedsl/shared_expert_pipeline.py
@@ -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()
diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py
index 979efa17..f1ce8b67 100644
--- a/vllm/nvfp4_cutedsl.py
+++ b/vllm/nvfp4_cutedsl.py
@@ -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.