Fix torch.compile: use custom autograd Function instead of @torch.compiler.disable

torch.compile fullgraph mode can't handle @torch.compiler.disable (skips
the function and refuses to compile). Custom autograd Functions are treated
as opaque ops by torch.compile — they execute eagerly without the compiler
trying to trace into CuTeDSL internals (JIT, Path.cwd, etc).

cutedsl/nvfp4_linear.py

@@ -21,6 +21,19 @@ from cutedsl.kernel.moe.torch_scaled_grouped_mm import (
 )
 
 
+class _Nvfp4LinearApply(torch.autograd.Function):
+    """Custom autograd function to make CuTeDSL runner opaque to torch.compile.
+    
+    torch.compile (fullgraph mode) can't trace through CuTeDSL internals
+    (JIT compilation, Path.cwd, etc.). By routing through a custom autograd
+    function, torch.compile treats it as an opaque op and doesn't try to
+    inline it.
+    """
+    @staticmethod
+    def forward(ctx, runner, hidden_states):
+        return runner._run_impl(hidden_states)
+
+
 class CuTeDSLNvfp4Linear:
     """Single NVFP4 GEMM using CuTeDSL (num_groups=1).
 
@@ -110,9 +123,12 @@ 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."""
+        return _Nvfp4LinearApply.apply(self, hidden_states)
+
+    def _run_impl(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Actual implementation — called via custom autograd to be torch.compile-safe."""
         self._ensure_initialized()
 
         num_tokens = hidden_states.shape[0]

cutedsl/shared_expert_pipeline.py

@@ -22,6 +22,13 @@ import torch
 
 from cutedsl.bridge import (
     quantize_activation_nvfp4,
+
+
+class _SharedExpertApply(torch.autograd.Function):
+    """Custom autograd function to make CuTeDSL runner opaque to torch.compile."""
+    @staticmethod
+    def forward(ctx, runner, hidden_states):
+        return runner._run_impl(hidden_states)
     quantize_to_nvfp4,
     make_b_k_major,
     assemble_scales_3d_side,
@@ -274,9 +281,12 @@ 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."""
+        return _SharedExpertApply.apply(self, hidden_states)
+
+    def _run_impl(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Actual implementation — called via custom autograd to be torch.compile-safe."""
         self._ensure_initialized()
 
         l1_out = self._run_l1(hidden_states)

vllm/nvfp4_cutedsl.py

@@ -18,6 +18,13 @@ import torch
 from cutedsl.bridge import (
     quantize_activation_nvfp4,
     quantize_weight_to_nvfp4,
+
+
+class _MoEApply(torch.autograd.Function):
+    """Custom autograd function to make CuTeDSL MoE runner opaque to torch.compile."""
+    @staticmethod
+    def forward(ctx, runner, hidden_states, topk_weights, topk_ids, expert_indices):
+        return runner._run_impl(hidden_states, topk_weights, topk_ids, expert_indices)
     quantize_to_nvfp4,
     make_b_k_major,
     assemble_scales_3d_side,
@@ -374,8 +381,11 @@ class CuTeDSLMoERunner:
         
 
 
-    @torch.compiler.disable
     def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None):
+        """Forward: route tokens to experts, GEMM, combine."""
+        return _MoEApply.apply(self, hidden_states, topk_weights, topk_ids, expert_indices)
+
+    def _run_impl(self, hidden_states, topk_weights, topk_ids, expert_indices=None):
         """Run the NVFP4 MoE forward pass.
         
         Handles global→local expert ID remapping for expert parallelism.