diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py index cc05fbaa..6e9d28fc 100644 --- a/cutedsl/bridge.py +++ b/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 diff --git a/cutedsl/nvfp4_linear.py b/cutedsl/nvfp4_linear.py index 9370b5bc..7d36bf8f 100644 --- a/cutedsl/nvfp4_linear.py +++ b/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.""" diff --git a/cutedsl/shared_expert_pipeline.py b/cutedsl/shared_expert_pipeline.py index 7650fde5..99913f1f 100644 --- a/cutedsl/shared_expert_pipeline.py +++ b/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.""" diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 87a863e3..979efa17 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/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.