diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py index a061430a..2f607070 100644 --- a/cutedsl/bridge.py +++ b/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, :]