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.