finegrained-kernels: decode losslessly compressed NVFP4 scales (CSF)

finegrained-kernels/tests/test_ops.py

@@ -45,8 +45,8 @@
     dq_grouped,
     dq_scale,
     make_static_activation_scale,
-    maybe_compile,
     quant_dequant_a,
+    prepare_weight_scales,
 )
 
 import finegrained_kernels  # type: ignore
@@ -79,6 +79,7 @@ class Problem:
     per_expert_globals: bool = False  # activation (and requant output) globals calibrated per expert
     prequant: bool = False  # pass As explicitly (must be bit-identical to raw A)
     static: bool = False  # calibrated (static) activation scale instead of a runtime one
+    csf_encoding: bool = False  # lossless CSF compression of NVFP4 scales
     swizzled: bool = False  # pass MX weight scales pre-swizzled (5D SWIZZLE_32_4_4 fast path)
     sentinel_fraction: float = 0.0
     noncontiguous: bool = False
@@ -106,6 +107,8 @@ def id(self):
             tag += "_prequant"
         if self.static:
             tag += "_static"
+        if self.csf_encoding:
+            tag += "_csf"
         if self.swizzled:
             tag += "_swizzled"
         if self.sentinel_fraction:
@@ -219,6 +222,10 @@ def scenarios() -> list[Problem]:
         # per-projection input_scale), which the act quant applies by the row's expert
         Problem(weights="nvfp4", per_expert_globals=True),
         Problem(weights="nvfp4", gate=True, activation_format="nvfp4", quantize_output=True, per_expert_globals=True),
+        Problem(weights="nvfp4", csf_encoding=True, S=8),
+        Problem(weights="nvfp4", csf_encoding=True, empty_expert=True, sentinel_fraction=0.25, compile=True),
+        Problem(weights="nvfp4", csf_encoding=True, activation_format="bf16"),
+        Problem(weights="nvfp4", csf_encoding=True, gate=True, activation_format="nvfp4", quantize_output=True, compile=True),
         # swizzled decode (S=8) — the bench's pre-swizzled batched decode arm per fp4 family
         Problem(weights="mxfp4", swizzled=True, S=8),
         Problem(weights="nvfp4", swizzled=True, S=8),
@@ -523,15 +530,15 @@ def _op(problem: Problem, op, A, expert_ids, B, Bs, Bs_global, As=None, As_globa
     bs = swizzle_mx_scales(Bs) if problem.swizzled and Bs is not None else Bs
     globals_kw = dict(a_global_scale=As_global, b_global_scale=Bs_global)
     if op == "matmul":
-        fn = maybe_compile(finegrained_kernels.matmul_2d, problem.compile)
+        fn = prepare_weight_scales(bs, problem.csf_encoding, finegrained_kernels.matmul_2d, problem.compile)
         return fn(A, B, As, bs, **globals_kw, **kw)
     if op == "batched":
-        fn = maybe_compile(finegrained_kernels.matmul_batched, problem.compile)
+        fn = prepare_weight_scales(bs, problem.csf_encoding, finegrained_kernels.matmul_batched, problem.compile)
         return fn(A, B, As, bs, expert_ids=expert_ids, **globals_kw, **kw)
     expert_start, gather_idx, scatter_idx = finegrained_kernels.scheduling.compute_grouped_scheduling(
         expert_ids, problem.E, 1
     )
-    fn = maybe_compile(finegrained_kernels.matmul_grouped, problem.compile)
+    fn = prepare_weight_scales(bs, problem.csf_encoding, finegrained_kernels.matmul_grouped, problem.compile)
     raw = fn(
         A,
         B,
@@ -626,7 +633,11 @@ def _check(problem: Problem, dq_out, ref_cmp, expert_ids, op):
 
 def _make_weights(problem: Problem, row, E):
     """The scenario's weights: ``(B, Bs, Bs_global)`` at the stack's row count (doubled under gate)."""
-    return row["make"](2 * problem.N if problem.gate else problem.N, problem.K, E)
+    weight, scales, global_scale = row["make"](2 * problem.N if problem.gate else problem.N, problem.K, E)
+    if problem.csf_encoding:
+        # Their E4M3 byte codes differ by more than 15, requiring an exception in each expert.
+        scales[..., 0, :2] = scales.new_tensor([0.125, 1.0])
+    return weight, scales, global_scale
 
 
 def _skip_moe_only(problem: Problem, op: str) -> None:

finegrained-kernels/tests/utils.py

@@ -23,6 +23,7 @@
     NVFP4_SCALE_GROUP_K,
     is_sm10x,
 )
+from finegrained_kernels.loading.scales import decompress_nvfp4_csf_scales, prepare_nvfp4_csf_scales
 
 
 # ── Device + capability ───────────────────────────────────────────────────────
@@ -476,3 +477,53 @@ def dequantize_weight(
     return out
 
 
+def encode_csf_scales(scales):
+    """Independent CSF fixture encoder; the original dense bytes remain the oracle."""
+    shape = scales.shape
+    values = scales.cpu().contiguous().view(torch.uint8)
+    values = values.reshape(-1, shape[-2], shape[-1])
+    experts, rows, cols = values.shape
+    bases = values.amin(dim=-1, keepdim=True).clamp_max(240)
+    delta = values.to(torch.int16) - bases.to(torch.int16)
+    nibbles = delta.clamp(0, 15).to(torch.uint8)
+    packed = nibbles[:, :, ::2] | (nibbles[:, :, 1::2] << 4)
+    fixed = torch.cat(
+        (
+            bases.reshape(experts, rows // 16, 16),
+            packed.reshape(experts, rows // 16, 16 * cols // 2),
+        ),
+        dim=-1,
+    )
+    streams = []
+    for expert in range(experts):
+        positions = (
+            ((delta[expert] < 0) | (delta[expert] > 15)).flatten().nonzero().flatten()
+        )
+        streams.append(
+            (positions | (values[expert].flatten()[positions].long() << 24)).to(
+                torch.uint32
+            )
+        )
+    return fixed, streams
+
+
+def prepare_weight_scales(scales, csf_encoding, fn, compiled=False):
+    """Wrap a matrix Problem's existing op; compiled cases capture decode + GEMM together."""
+    if not csf_encoding:
+        return maybe_compile(fn, compiled)
+    fixed, exceptions = encode_csf_scales(scales)
+    buffers = prepare_nvfp4_csf_scales(
+        fixed.to(scales.device), [stream.to(scales.device) for stream in exceptions]
+    )
+
+    # A contiguous workspace need not start on a two-byte boundary.
+    workspace = torch.empty(scales.numel() + 1, dtype=scales.dtype, device=scales.device)[1:]
+    workspace = workspace.reshape(-1, *scales.shape[-2:])
+
+    def run(A, B, As, Bs, **kwargs):
+        decoded = decompress_nvfp4_csf_scales(
+            *buffers, out=workspace, expert_start=kwargs.get("expert_start")
+        )
+        return fn(A, B, As, decoded.reshape(Bs.shape), **kwargs)
+
+    return maybe_compile(run, compiled)

finegrained-kernels/torch-ext/finegrained_kernels/__init__.py

@@ -21,6 +21,7 @@
 from .formats import get_supported_act_fns
 from .norm import get_supported_norms, rms_norm_rows
 from .swizzle import swizzle_mx_scales, unswizzle_mx_scales
+from .loading.scales import prepare_nvfp4_csf_scales, decompress_nvfp4_csf_scales
 from .quant import mxfp4_act_quant, mxfp8_act_quant, nvfp4_act_quant
 
 __all__ = [
@@ -42,6 +43,8 @@
     # quantizers a loader uses to quantize weights into the group formats
     "swizzle_mx_scales",
     "unswizzle_mx_scales",
+    "prepare_nvfp4_csf_scales",
+    "decompress_nvfp4_csf_scales",
     "mxfp8_act_quant",
     "mxfp4_act_quant",
     "nvfp4_act_quant",