Fix torch.compile crash: remove threading.Lock from LUT cache path

The _NVFP4_STEP_LUT_LOCK caused 'Unsupported context manager' under
torch.compile/cudagraph. LUT is now pre-populated during warmup so
the fast path (cache hit) never hits a lock.

Also removed all init/warmup debug prints from CuTeDSL kernels.

cutedsl/bridge.py

@@ -9,7 +9,6 @@ the ScaledGroupedGemmKernel expects:
 - Expert offset computation
 """
 import math
-import threading
 import torch
 import cutlass
 import cutlass.cute as cute
@@ -34,21 +33,23 @@ _compiled_kernel_cache = {}
 
 # Cached LUT for E2M1 quantization (created once per device, cudagraph-safe)
 _NVFP4_STEP_LUT_CACHE = {}
-_NVFP4_STEP_LUT_LOCK = threading.Lock()
-
-
 def _get_step_to_idx_lut(device):
     """Get or create the E2M1 step-to-index LUT for the given device.
     
     Cached per device to avoid CPU->CUDA copies during cudagraph capture.
+    Must be pre-populated during warmup (before torch.compile/cudagraph capture)
+    so the lock is never entered on the compiled path.
     """
-    with _NVFP4_STEP_LUT_LOCK:
-        if device not in _NVFP4_STEP_LUT_CACHE:
-            _NVFP4_STEP_LUT_CACHE[device] = torch.as_tensor(
-                [0, 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 7, 7],
-                dtype=torch.int8, device=device,
-            )
+    # Fast path: already cached — no lock needed (torch.compile-safe)
+    if device in _NVFP4_STEP_LUT_CACHE:
         return _NVFP4_STEP_LUT_CACHE[device]
+    # Slow path: first call, create the LUT
+    lut = torch.as_tensor(
+        [0, 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 7, 7],
+        dtype=torch.int8, device=device,
+    )
+    _NVFP4_STEP_LUT_CACHE[device] = lut
+    return lut
 SF_VEC_SIZE = 16  # NVFP4 block size
 
 

cutedsl/nvfp4_linear.py

@@ -62,9 +62,6 @@ class CuTeDSLNvfp4Linear:
         self._gsa_buf = None
         self._buffers_allocated = False
 
-        print(f"  Nvfp4Linear init: in={in_features} out={out_features} "
-              f"max_tokens={max_num_tokens}", flush=True)
-
     def finalize_weights(self):
         """Process weights for CuTeDSL GEMM."""
         self._mat_b = make_b_k_major(torch.stack(self.fp4))  # (1, K_packed, N_packed)
@@ -111,7 +108,7 @@ class CuTeDSLNvfp4Linear:
         with torch.no_grad():
             _, _, gs = quantize_to_nvfp4(hidden_states_sample)
             self._activation_global_scale = gs
-        print(f"  Nvfp4Linear warmup: gs={gs:.10f}", flush=True)
+
 
     def run(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Forward: BF16 input → NVFP4 GEMM → BF16 output."""

cutedsl/shared_expert_pipeline.py

@@ -83,9 +83,6 @@ class CuTeDSLSharedExpertRunner:
         self._expert_offsets_buf = None
         self._buffers_allocated = False
 
-        print(f"  SharedExpert init: hidden={hidden_size} intermediate={intermediate_size} "
-              f"max_tokens={max_num_tokens}", flush=True)
-
     def set_swiglu_limit(self, limit: float):
         self.swiglu_limit = limit
 
@@ -198,8 +195,7 @@ class CuTeDSLSharedExpertRunner:
                 _, _, l2_gs = quantize_to_nvfp4(activated)
                 self._l2_activation_global_scale = l2_gs
 
-        print(f"  SharedExpert warmup: L1 gs={self._l1_activation_global_scale:.10f} "
-              f"L2 gs={self._l2_activation_global_scale:.10f}", flush=True)
+
 
     def _run_l1(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """L1 GEMM: activation × gate_up_weight → BF16."""

vllm/nvfp4_cutedsl.py

@@ -372,8 +372,7 @@ class CuTeDSLMoERunner:
         self._l1_activation_global_scale = l1_gs
         self._l2_activation_global_scale = l2_gs
         
-        print(f"  Warmup L1 gs: {l1_gs:.10f}")
-        print(f"  Warmup L2 gs: {l2_gs:.10f}")
+
 
     def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None):
         """Run the NVFP4 MoE forward pass.