test: old bridge + LUT quantization only (step 1 of cudagraph migration)

cutedsl/bridge.py

@@ -9,6 +9,7 @@ the ScaledGroupedGemmKernel expects:
 - Expert offset computation
 """
 import math
+import threading
 import torch
 import cutlass
 import cutlass.cute as cute
@@ -27,6 +28,24 @@ from cutedsl.kernel.moe.torch_scaled_grouped_mm import (
 # ── Constants ──────────────────────────────────────────────────────────
 
 E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
+
+# 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.
+    """
+    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,
+            )
+        return _NVFP4_STEP_LUT_CACHE[device]
 SF_VEC_SIZE = 16  # NVFP4 block size
 
 
@@ -71,11 +90,12 @@ def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE):
     block_sf_expanded = block_scale.float().unsqueeze(-1)
     x_scaled = x_reshaped / block_sf_expanded.clamp(min=1e-8)
 
-    magnitudes = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=x_bf16.device)
     signs = torch.sign(x_scaled)
-    abs_scaled = x_scaled.abs().unsqueeze(-1)
-    distances = (abs_scaled - magnitudes).abs()
-    indices = distances.argmin(dim=-1)
+    abs_scaled = x_scaled.abs().clamp(max=6.0)
+    
+    half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8)
+    step_to_idx = _get_step_to_idx_lut(x_bf16.device)
+    indices = step_to_idx[half_steps.long()]
 
     nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8)
     even = nibbles[..., ::2]
@@ -123,11 +143,12 @@ def quantize_weight_to_nvfp4(w_bf16, block_size=SF_VEC_SIZE):
     w_block_sf = w_sf.float().unsqueeze(1)
     w_scaled = w_reshaped / w_block_sf.clamp(min=1e-8)
 
-    magnitudes = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=w_bf16.device)
     signs = torch.sign(w_scaled)
-    abs_scaled = w_scaled.abs().unsqueeze(-1)
-    distances = (abs_scaled - magnitudes).abs()
-    indices = distances.argmin(dim=-1)
+    abs_scaled = w_scaled.abs().clamp(max=6.0)
+    
+    half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8)
+    step_to_idx = _get_step_to_idx_lut(w_bf16.device)
+    indices = step_to_idx[half_steps.long()]
     nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8)
 
     even = nibbles[:, ::2, :]