diff --git a/finegrained-kernels/bench/bench_moe_csf_latency.png b/finegrained-kernels/bench/bench_moe_csf_latency.png new file mode 100644 index 000000000..512c8294e Binary files /dev/null and b/finegrained-kernels/bench/bench_moe_csf_latency.png differ diff --git a/finegrained-kernels/bench/bench_moe_csf_memory.png b/finegrained-kernels/bench/bench_moe_csf_memory.png new file mode 100644 index 000000000..80d928c1e Binary files /dev/null and b/finegrained-kernels/bench/bench_moe_csf_memory.png differ diff --git a/finegrained-kernels/tests/test_ops.py b/finegrained-kernels/tests/test_ops.py index b75d3582e..5d691b89a 100644 --- a/finegrained-kernels/tests/test_ops.py +++ b/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: diff --git a/finegrained-kernels/tests/utils.py b/finegrained-kernels/tests/utils.py index 5a862bd03..ac4862fb3 100644 --- a/finegrained-kernels/tests/utils.py +++ b/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) diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/__init__.py b/finegrained-kernels/torch-ext/finegrained_kernels/__init__.py index f5bdfceb5..b4fd35f69 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/__init__.py +++ b/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", diff --git a/finegrained-kernels/torch-ext/finegrained_kernels/loading/scales.py b/finegrained-kernels/torch-ext/finegrained_kernels/loading/scales.py index 237f45f5d..49ce7c783 100644 --- a/finegrained-kernels/torch-ext/finegrained_kernels/loading/scales.py +++ b/finegrained-kernels/torch-ext/finegrained_kernels/loading/scales.py @@ -12,9 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections.abc import Sequence + +import torch import triton import triton.language as tl +from ..compat import ( + add_op_namespace_prefix, compile_time_only_triton_op, compile_time_only_triton_wrap, + device_context, sm_count, +) from ..quant import fp8_act_quant_inline, mx_act_quant_inline @@ -323,3 +330,238 @@ def mx_2d_scale_ptrs( as_mask = offs_am_lin < M bs_mask = offs_bn_lin < (2 * N if GATE else N) return as_ptrs, bs_ptrs, as_mask, bs_mask + + +# CSF byte-window scales: stream preparation and staged decoding before ordinary GEMMs. + + +def prepare_nvfp4_csf_scales( + fixed: torch.Tensor, exceptions: torch.Tensor | Sequence[torch.Tensor] +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepare one matrix or an expert stack without expanding or moving its scales. + + `fixed` is uint8 `[rows/16, 16*(1+cols/2)]`, or the same shape with an + expert dimension. Exceptions are uint32 `[count]` for one matrix, or one + tensor per expert; low 24 bits locate a scale and high 8 bits replace it. + Returns (fixed, concatenated exceptions, expert offsets) on the input device. + Decode on the accelerator at load time for dense residency, or into shared + scratch immediately before GEMM. + """ + if fixed.dtype != torch.uint8 or fixed.ndim not in (2, 3): + raise ValueError("CSF fixed stream must be a 2D or 3D uint8 tensor") + fixed = fixed.unsqueeze(0) if fixed.ndim == 2 else fixed + experts, slabs, slab_bytes = fixed.shape + if experts == 0 or slabs == 0 or slab_bytes <= 16 or slab_bytes % 16: + raise ValueError("Invalid CSF slab shape") + rows, cols = slabs * 16, (slab_bytes // 16 - 1) * 2 + if rows * cols > 1 << 24: + raise ValueError( + "CSF exception positions support at most 2**24 scales per matrix" + ) + if bool((fixed[:, :, :16] > 240).any()): + raise ValueError("CSF row bases must be in 0..240") + streams = [exceptions] if isinstance(exceptions, torch.Tensor) else list(exceptions) + if len(streams) != experts: + raise ValueError("CSF needs one exception stream per expert") + offsets, words = [0], [] + for stream in streams: + if ( + stream.dtype != torch.uint32 + or stream.ndim != 1 + or stream.device != fixed.device + ): + raise ValueError( + "CSF exception streams must be 1D uint32 tensors on the fixed stream's device" + ) + values = stream.to(torch.int64) + positions = values & 0xFFFFFF + if positions.numel() and ( + int(positions[-1]) >= rows * cols + or bool((positions[1:] <= positions[:-1]).any()) + ): + raise ValueError( + "CSF exception positions must be sorted, unique and in range" + ) + offsets.append(offsets[-1] + stream.numel()) + words.append(stream) + if offsets[-1] >= 1 << 32: + raise ValueError("CSF exception directory exceeds uint32 offsets") + return ( + fixed.contiguous(), + torch.cat(words), + torch.tensor(offsets, dtype=torch.uint32, device=fixed.device), + ) + + +@triton.jit +def _decode_fixed_kernel( + Fixed, + Out, + ExpertStart, + E: tl.constexpr, + ROWS: tl.constexpr, + COLS: tl.constexpr, + BLOCK: tl.constexpr, + ALIGNED_OUTPUT: tl.constexpr, +): + pairs: tl.constexpr = ROWS * (COLS // 2) + tiles: tl.constexpr = tl.cdiv(pairs, BLOCK) + offsets = tl.arange(0, BLOCK) + # Contiguous pairs avoid padding each row to a power-of-two tile width. + # Keep indices local to an expert; only the expert base needs 64-bit addressing. + for tile in range(tl.program_id(0), E * tiles, tl.num_programs(0)): + expert = tile // tiles + active = True + if ExpertStart is not None: + active = tl.load(ExpertStart + expert + 1) > tl.load(ExpertStart + expert) + if active: + pair = (tile % tiles) * BLOCK + offsets + row = pair // (COLS // 2) + slab = expert.to(tl.int64) * (ROWS // 16) + row // 16 + slab_bytes: tl.constexpr = 16 * (1 + COLS // 2) + base = tl.load(Fixed + slab * slab_bytes + row % 16, pair < pairs, other=0).to( + tl.uint16 + ) + packed = tl.load( + Fixed + slab * slab_bytes + 16 + (row % 16) * (COLS // 2) + pair % (COLS // 2), + pair < pairs, + other=0, + ).to(tl.uint16) + low = base + (packed & 15) + high = base + (packed >> 4) + # A contiguous sliced workspace can start at an odd byte. + if ALIGNED_OUTPUT: + output = Out.to(tl.pointer_type(tl.uint16)) + tl.store( + output + expert.to(tl.int64) * pairs + pair, + low | (high << 8), + pair < pairs, + ) + else: + values = tl.interleave(low, high).to(tl.uint8).to(tl.float8e4nv, bitcast=True) + output_offsets = (tile % tiles) * BLOCK * 2 + tl.arange(0, BLOCK * 2) + tl.store( + Out + expert.to(tl.int64) * pairs * 2 + output_offsets, + values, + output_offsets < pairs * 2, + ) + + +@triton.jit +def _apply_exceptions_kernel( + Exceptions, + ExpertOffsets, + Out, + ExpertStart, + PLANE: tl.constexpr, + BLOCK: tl.constexpr, +): + expert = tl.program_id(0) + if ExpertStart is not None: + if tl.load(ExpertStart + expert + 1) == tl.load(ExpertStart + expert): + return + start = tl.load(ExpertOffsets + expert).to(tl.int64) + end = tl.load(ExpertOffsets + expert + 1).to(tl.int64) + for base in range(start, end, BLOCK): + index = base + tl.arange(0, BLOCK) + word = tl.load(Exceptions + index, index < end, other=0) + position = word & 0xFFFFFF + value = (word >> 24).to(tl.uint8).to(tl.float8e4nv, bitcast=True) + tl.store(Out + expert.to(tl.int64) * PLANE + position, value, index < end) + + +@compile_time_only_triton_op( + add_op_namespace_prefix("decompress_nvfp4_csf_scales_into"), + mutates_args={"out"}, + opaque=True, +) +def _decompress_into( + fixed: torch.Tensor, + exceptions: torch.Tensor, + expert_offsets: torch.Tensor, + out: torch.Tensor, + expert_start: torch.Tensor | None, +) -> None: + experts, rows, cols = out.shape + with device_context(fixed.device): + programs = min( + experts * triton.cdiv(rows * cols // 2, 4096), + 16 * sm_count(fixed.device.index), + ) + compile_time_only_triton_wrap(_decode_fixed_kernel)[(programs,)]( + fixed, + out, + expert_start, + experts, + rows, + cols, + 4096, + ALIGNED_OUTPUT=out.storage_offset() % 2 == 0, + ) + if exceptions.numel(): + compile_time_only_triton_wrap(_apply_exceptions_kernel)[(experts,)]( + exceptions, expert_offsets, out, expert_start, rows * cols, 256 + ) + + +def decompress_nvfp4_csf_scales( + fixed: torch.Tensor, + exceptions: torch.Tensor, + expert_offsets: torch.Tensor, + *, + out: torch.Tensor | None = None, + expert_start: torch.Tensor | None = None, +) -> torch.Tensor: + """Decode prepared CSF streams to dense E4M3 scales for an ordinary NVFP4 GEMM. + + With no `out`, allocate dense scales on the accelerator at checkpoint loading. + Pass reusable contiguous [experts, rows, cols] workspace to keep the streams + compressed between layers. Consume the workspace before overwriting it; it + must not be retained for backward or shared by concurrent consumers. + A shared dense workspace saves memory across layers, not within one layer. + Pass the grouped GEMM's cumulative `expert_start` to decode only experts with + routed tokens. Inactive workspace planes stay untouched and must not be read. + """ + if fixed.dtype != torch.uint8 or fixed.ndim != 3 or not fixed.is_contiguous(): + raise ValueError("CSF fixed stream must be contiguous 3D uint8") + if fixed.device.type == "cpu": + raise ValueError("CSF decompression requires an accelerator") + experts, slabs, slab_bytes = fixed.shape + if experts == 0 or slabs == 0 or slab_bytes <= 16 or slab_bytes % 16: + raise ValueError("Invalid CSF slab shape") + if ( + exceptions.dtype != torch.uint32 + or exceptions.ndim != 1 + or not exceptions.is_contiguous() + or expert_offsets.dtype != torch.uint32 + or expert_offsets.shape != (experts + 1,) + or not expert_offsets.is_contiguous() + ): + raise ValueError( + "CSF exceptions and expert offsets must be contiguous uint32 vectors of prepared shape" + ) + rows, cols = slabs * 16, (slab_bytes // 16 - 1) * 2 + if rows * cols > 1 << 24: + raise ValueError( + "CSF exception positions support at most 2**24 scales per matrix" + ) + shape = (experts, rows, cols) + if out is None: + out = torch.empty(shape, dtype=torch.float8_e4m3fn, device=fixed.device) + if ( + out.shape != shape + or out.dtype != torch.float8_e4m3fn + or not out.is_contiguous() + ): + raise ValueError(f"CSF workspace must be contiguous E4M3 with shape {shape}") + if any(t.device != out.device for t in (fixed, exceptions, expert_offsets)): + raise ValueError("CSF buffers and workspace must share a device") + if expert_start is not None and ( + expert_start.shape != (experts + 1,) + or expert_start.dtype not in (torch.int32, torch.int64) + or not expert_start.is_contiguous() + or expert_start.device != out.device + ): + raise ValueError("CSF expert_start must be a contiguous device int32/int64 vector of length experts + 1") + _decompress_into(fixed, exceptions, expert_offsets, out, expert_start) + return out