diff --git a/cutedsl/nvfp4_linear.py b/cutedsl/nvfp4_linear.py index 2119176f..0e0d5ee0 100644 --- a/cutedsl/nvfp4_linear.py +++ b/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] diff --git a/cutedsl/shared_expert_pipeline.py b/cutedsl/shared_expert_pipeline.py index 224ca1c7..7d1429bd 100644 --- a/cutedsl/shared_expert_pipeline.py +++ b/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) diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index f1ce8b67..e22e60fa 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/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.