From de0247a0d2a04557c7ea57b88435ef63e4ac3e56 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Thu, 30 Jul 2026 18:54:24 +0100 Subject: [PATCH 1/3] feat(moe): resident-EP BF16 decode seam Decode-only companion to the streamed pure-DP path: per-rank stacked shard built once from the copy-engine host tensors; pad -> all_gather -> global router -> fused_moe_bf16 with non-local slots weight-masked -> all_reduce(SUM) -> local slice. Empty ranks participate fully in the collectives (their experts serve other ranks). Sized purely from the worker-synced num_tokens_per_rank scalar. --- batchgen/moe/fused_moe_bf16_resident.py | 205 ++++++++++++++++++++++++ 1 file changed, 205 insertions(+) create mode 100644 batchgen/moe/fused_moe_bf16_resident.py diff --git a/batchgen/moe/fused_moe_bf16_resident.py b/batchgen/moe/fused_moe_bf16_resident.py new file mode 100644 index 00000000..1fda5e6c --- /dev/null +++ b/batchgen/moe/fused_moe_bf16_resident.py @@ -0,0 +1,205 @@ +# ---------------------------------------------------------------------------- # +# BatchGen # +# copyright (c) EfficientMoE team 2025 # +# # +# licensed under the apache license, version 2.0 (the "license"); # +# ---------------------------------------------------------------------------- # + +"""Resident-EP decode MoE seam: stacked BF16 shards + fused grouped GEMM + +cross-rank combine (Kimi-Linear M4 P0.3). + +DECODE-ONLY companion to the streamed pure-DP MoE path. Each rank materializes +the stacked BF16 weights of its EP shard (num_experts / world_size experts per +MoE layer) ONCE at configure_decoding — from the same host copy-engine tensors +the streamed path consumes (core_engine.get_tensor) — and then computes every +decode step entirely from HBM (no per-step H2D): + + pad local rows to the synced per-rank layout -> comm.all_gather -> + router on the gathered GLOBAL tokens -> fused_moe_bf16 over the LOCAL + shard with non-local expert slots weight-masked -> comm.all_reduce(SUM) + over the global layout -> extract the local slice. + +Cross-rank layout contract (mirrors KimiK25MoE._forward_decode): + - ``num_tokens_per_rank`` (ntp) is the per-step MAX decode rows over all + ranks, synced by the worker's ``_sync_decode_moe_rank_counts()`` collective + and delivered via ``PSM.set_num_tokens_per_rank()`` BEFORE the forward. + Every rank — including an EMPTY one — therefore knows the global buffer + shape ``(world_size * ntp, H)`` from the synced scalar alone; no extra + communication is needed to size the collectives. + - An empty rank (0 local decode rows) skips NOTHING except the final local + slice: it contributes zero-padded rows to the all_gather and MUST run the + router + kernel + all_reduce, because its resident experts serve the OTHER + ranks' tokens. Zero-padded rows are harmless — a bias-free expert MLP maps + x = 0 to exactly 0 (w1·0 = 0, silu(0)·0 = 0, w2·0 = 0), so padding adds + nothing to the reduced sum and the padded slice is discarded on extract. + - Decode-only safety: all ranks always step decode together (worker :9746 + no-skip invariant), so the collectives cannot deadlock. Never call this + from prefill — prefill ranks may run different module sequences. +""" + +import logging +import time + +import torch +from torch.distributed import ReduceOp + +from batchgen_kernels.triton.fused_moe_bf16 import fused_moe_bf16 + + +def build_layer_shard(get_tensor, layer_idx, expert_start, num_local, + hidden_size, intermediate_size, device, + dtype=torch.bfloat16): + """Materialize one MoE layer's local expert shard as stacked GPU tensors. + + Copies each local expert's BF16 weights host->GPU exactly once, straight + from the copy-engine host source (``core_engine.get_tensor``, keys + ``routed_expert_{layer}_{expert}`` -> {"w1.weight", "w2.weight", + "w3.weight"}; w1 = gate, w3 = up, w2 = down). + + HBM budget (per rank, under the server's --gpu-memory-frac ceiling; + num_local = 256/8 = 32, 26 MoE layers, H = 2304, I = 1024): + - resident EP shards: 26 x 32 x (2*1024*2304 + 2304*1024) el x 2 B + = 26 x 453.5 MiB ~= 11.8 GB + - KDA state pools: kda_state_slots(256) x 20 layers x + (n_h*d_h^2 x 4 B fp32 recurrent + + 3 x 4096 x (W-1) x 2 B conv) + ~= 5.5 GB (d_h = 64) - 11.1 GB (d_h = 128 worst + case); scales linearly with kda_state_slots + - resident non-expert (attn/KDA/shared/skeleton/lm_head): ~4 GB + -> ~21-27 GB static; leaves > 55 GB under the 86.4 GB ceiling for the + paged-KV pool, staging buffers, NCCL and activations — fits. + + Returns: + w13: (num_local, 2*I, H) stacked [gate; up] weights (gate first — + the packing fused_moe_bf16 expects). + w2: (num_local, H, I) stacked down-projection weights. + """ + I, H = intermediate_size, hidden_size + w13 = torch.empty((num_local, 2 * I, H), dtype=dtype, device=device) + w2 = torch.empty((num_local, H, I), dtype=dtype, device=device) + for i in range(num_local): + tensors = get_tensor(f"routed_expert_{layer_idx}_{expert_start + i}") + w13[i, :I].copy_(tensors["w1.weight"]) + w13[i, I:].copy_(tensors["w3.weight"]) + w2[i].copy_(tensors["w2.weight"]) + return w13, w2 + + +class ResidentEPMoELayer: + """Per-layer resident-EP decode MoE forward (thin: kernel + collectives). + + Holds the layer's stacked shard and the NCCL communicator; the router + (KimiMoEGate) stays owned by the model and is passed into ``forward`` so + it runs — unchanged — on the gathered global tokens (K2.5 pattern; gate + weights are DP-replicated, so all ranks compute identical routing). + """ + + # Per-step padded rows per rank; synced across ranks by the worker + # (dist.all_gather_into_tensor over local batch sizes) and written here + # through PSM.set_num_tokens_per_rank before every decode forward. + num_tokens_per_rank = None + + def __init__(self, layer_idx, w13, w2, comm, world_size, rank, + expert_start): + self.layer_idx = layer_idx + self.w13 = w13 + self.w2 = w2 + self.comm = comm + self.world_size = world_size + self.rank = rank + self.expert_start = expert_start + self.num_local_experts = w13.shape[0] + + @classmethod + def set_num_tokens_per_rank(cls, num_tokens_per_rank): + cls.num_tokens_per_rank = int(num_tokens_per_rank) + + def forward(self, x_local, gate): + """Resident-EP decode MoE over DP-sharded tokens. + + Args: + x_local: (num_local_tokens, H) this rank's decode rows; may be + (0, H) on an empty rank — the collectives still run. + gate: router module; called on the gathered (num_global, 1, H) + tokens, returns (topk_idx, topk_weight[, ...]). + + Returns: + (num_local_tokens, H) summed routed-expert output for the local + rows (all 256 experts, combined across ranks). + """ + ntp = ResidentEPMoELayer.num_tokens_per_rank + num_tokens, H = x_local.shape + assert ntp is not None and ntp > 0, ( + "resident-EP decode requires the worker rank-count sync " + "(PSM.set_num_tokens_per_rank) before the MoE forward" + ) + assert num_tokens <= ntp, ( + f"local decode rows {num_tokens} exceed synced ntp {ntp}" + ) + num_global = self.world_size * ntp + + # Fixed per-rank layout: zero-pad local rows to ntp, gather globally. + padded = x_local.new_zeros((ntp, H)) + if num_tokens > 0: + padded[:num_tokens].copy_(x_local) + all_tokens = x_local.new_empty((num_global, H)) + with self.comm.change_state(enable=True): + self.comm.all_gather(all_tokens, padded) + + # Router on the global tokens (identical on every rank). + gate_out = gate(all_tokens.view(num_global, 1, H)) + topk_idx, topk_weight = gate_out[0], gate_out[1] + + # Mask routing to the local shard: non-local slots keep weight 0 and + # are pointed at local expert 0 — they contribute exactly 0 through + # the weighted top-k reduction, and the extra rows cost no additional + # weight traffic in the BW-bound decode regime (same experts read). + local_ids = topk_idx - self.expert_start + local_mask = (local_ids >= 0) & (local_ids < self.num_local_experts) + local_ids = torch.where(local_mask, local_ids, + torch.zeros_like(local_ids)) + masked_weight = topk_weight * local_mask.to(topk_weight.dtype) + + partial = fused_moe_bf16(all_tokens, self.w13, self.w2, + masked_weight, local_ids) + + # Combine expert shards: SUM over ranks in the global layout (BF16 + # over NCCL, K2.5-class numerics; per-rank top-k accumulation above + # is fp32 inside moe_weighted_sum). + with self.comm.change_state(enable=True): + self.comm.all_reduce(partial, op=ReduceOp.SUM) + + start = self.rank * ntp + return partial[start:start + num_tokens] + + +def build_resident_ep_layers(model_layers, get_tensor, comm, world_size, rank, + expert_start, num_local, intermediate_size, + device): + """Materialize shards for every MoE layer and attach a ResidentEPMoELayer + to each ``block_sparse_moe`` as ``_resident_ep_moe`` (consumed by + ``moe_forward_serving``'s decode dispatch). Returns total bytes resident. + """ + start_t = time.perf_counter() + total_bytes = 0 + num_layers = 0 + for layer_idx, layer in enumerate(model_layers): + moe = getattr(layer, "block_sparse_moe", None) + if moe is None or moe.experts is None: + continue + w13, w2 = build_layer_shard( + get_tensor, layer_idx, expert_start, num_local, + moe.moe_hidden_size, intermediate_size, device, + ) + moe._resident_ep_moe = ResidentEPMoELayer( + layer_idx, w13, w2, comm, world_size, rank, expert_start, + ) + total_bytes += (w13.numel() * w13.element_size() + + w2.numel() * w2.element_size()) + num_layers += 1 + logging.info( + f"Rank {rank}: resident EP shards materialized — {num_layers} MoE " + f"layers x {num_local} experts, {total_bytes / (1024**3):.2f} GiB " + f"({time.perf_counter() - start_t:.1f}s, one-time H2D)" + ) + return total_bytes From 27fccf14fe13aa2880bf09d21a336a8707a08d8c Mon Sep 17 00:00:00 2001 From: tairanxu Date: Wed, 5 Aug 2026 15:52:19 +0100 Subject: [PATCH 2/3] feat(moe): extend marlin to MXFP4 (E2M1 + E8M0) for K3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Repack, dequant, SiTU epilogue, hard-fail contract. Nibble order and E8M0 semantics settled against compressed-tensors 0.17.1 as an INDEPENDENT oracle (ours was bent to match our own kernel, so our pair proved nothing), on a real K3 expert tensor pulled from the checkpoint: verdict frozen into mxfp4_oracle_vector.py with sha256 pins and a clamp-free reference dequant, so it never needs re-deriving. marlin_weight_prep.py: repack_mxfp4_to_marlin_gs32 — pure nibble rearrangement carrying E8M0 uint8 scales through (E2M1 magnitudes are non-uniform, so INT4 conversion would be a second lossy quantization and is refused), exact E8M0->bf16 expansion (<<7; edge bytes 0x00/0xFF raise), the w1||w3 storage-adjacent fused repack, an exact CPU inverse for the round-trip proof, and R1-R8 hard-fail contract checks. marlin_grouped_gemm.cu: dequant_e2m1 + WCodec/Act templating of the grouped kernels, a SiTU epilogue, and two new pybind entries (grouped_marlin_gemm_m16_mxfp4, _m16_s1_mxfp4_situ) with TORCH_CHECK hard-fail seams so the raw bindings cannot bypass the contract. HONESTY NOTE: the existing INT4 kernel bodies are TEMPLATED (dequant_u4b8 -> dequant_w4), behavior-preserving by construction but not byte-identical -- the staged GPU ladder carries INT4 M16 and fused-S1 regression arms (T6/T6b) that MUST pass before any merge to main. K2.5's python surface is append-only. marlin_grouped_moe.py: K3 wrappers with L1-L5 + activation-contract hard-fail checks per the 2026-08-04 POIS ruling -- no operational fallback; the reference path exists only behind an explicit debug opt-in whose server-side consumer is a named follow-up. --- batchgen/moe/marlin_grouped_moe.py | 279 ++++++++ batchgen/moe/marlin_weight_prep.py | 328 +++++++++- batchgen/moe/mxfp4_grouped_gemm.py | 52 +- batchgen/moe/mxfp4_oracle_vector.py | 147 +++++ .../src/moe/marlin_grouped_gemm.cu | 268 +++++++- tests/moe/__init__.py | 0 tests/moe/_loader.py | 35 + tests/moe/gpu_parity_mxfp4_marlin.py | 616 ++++++++++++++++++ tests/moe/test_mxfp4_marlin_repack.py | 380 +++++++++++ 9 files changed, 2077 insertions(+), 28 deletions(-) create mode 100644 batchgen/moe/mxfp4_oracle_vector.py create mode 100644 tests/moe/__init__.py create mode 100644 tests/moe/_loader.py create mode 100644 tests/moe/gpu_parity_mxfp4_marlin.py create mode 100644 tests/moe/test_mxfp4_marlin_repack.py diff --git a/batchgen/moe/marlin_grouped_moe.py b/batchgen/moe/marlin_grouped_moe.py index b55bd561..ecd217a3 100644 --- a/batchgen/moe/marlin_grouped_moe.py +++ b/batchgen/moe/marlin_grouped_moe.py @@ -8,6 +8,12 @@ Both use GROUP_BLOCKS=2 (gs=32, K2.5 native). All buffers and pointer arrays pre-computed at init time. Per-step forward: 2 kernel launches (GEMM + SiLU), zero Python loops or allocations. + +K2.5 wrappers above keep that zero-overhead contract. The K3 MXFP4 wrappers +below add host-side hard-fail contract checks per the 2026-08-04 ledger — +cheap (a handful of attribute reads per launch, no device work) but not zero; +if the model-side integration calls them per-step in eager mode, hoisting the +static checks (L1/L4) to plan build is a named integration follow-up. """ import logging @@ -242,3 +248,276 @@ def marlin_grouped_stage1_unified( intermediate_3d, up_buf, expert_counts, num_experts, mtp, compact_stride, N, ) + + +# ============================================================================ +# Kimi-K3 MXFP4 (E2M1 + E8M0) wrappers. HARD-FAIL policy: +# every contract violation RAISES; there is no warn-and-degrade. The unfused +# reference path survives only behind the explicit batchgen_debug opt-in +# (`k3_moe_reference`) as the parity oracle — it is not a fallback. NOTE the +# opt-in's consumer is model/server-side wiring (outside this kernel PR's +# allowlist) and is a NAMED FOLLOW-UP; until it lands there is no reference +# forward path at all — the flag name below is forward-declared, not live. +# ============================================================================ + +_MXFP4_KERNEL_ENTRIES = ( + "grouped_marlin_gemm_m16_mxfp4", + "grouped_marlin_gemm_m16_s1_mxfp4_situ", +) + +_warned_mxfp4 = False + + +def is_marlin_mxfp4_available() -> bool: + return _module is not None and all( + hasattr(_module, k) for k in _MXFP4_KERNEL_ENTRIES) + + +def _require_mxfp4_kernels(): + """L1: the marlin MXFP4 entries must exist — K3 refuses to run otherwise.""" + missing = [k for k in _MXFP4_KERNEL_ENTRIES if not hasattr(_module, k)] + if missing: + raise RuntimeError( + f"Marlin MXFP4 kernel entries missing from " + f"batchgen_kernels.moe._C_marlin_grouped_gemm: {missing}. " + f"K3 refuses to run (stale batchgen_kernels build — rebuild). " + f"The designated parity-debug opt-in is batchgen_debug." + f"k3_moe_reference; its model-side wiring is a named follow-up " + f"of the K3 MXFP4 work — if it is not wired yet there is NO alternative " + f"path and rebuilding is the only fix.") + + +def _check_ptr_array(name: str, t: torch.Tensor, length: int): + if t.dtype != torch.int64 or not t.is_cuda or t.numel() != length: + raise ValueError( + f"group metadata contract violated: {name} must be int64 CUDA " + f"[{length}], got {t.dtype} {t.device} numel={t.numel()}") + + +def _check_counts(name: str, t: torch.Tensor, length: int): + if t.dtype != torch.int32 or not t.is_cuda or t.numel() != length: + raise ValueError( + f"group metadata contract violated: {name} must be int32 CUDA " + f"[{length}], got {t.dtype} {t.device} numel={t.numel()}") + + +def _check_activation(name: str, x: torch.Tensor, last_dim: int, + require_contiguous: bool = True): + """Activation-tensor contract at the hard-fail seams. An fp16 (or fp32) + activation is byte-compatible with the kernel's bf16 reinterpret and + produces finite silent garbage — exactly the class the HARD-FAIL ledger + targets — so dtype/device/shape are checked, not assumed.""" + if (x.dtype != torch.bfloat16 or not x.is_cuda or x.shape[-1] != last_dim + or (require_contiguous and not x.is_contiguous())): + raise ValueError( + f"activation contract violated: {name} must be contiguous bf16 " + f"CUDA [..., {last_dim}], got {x.dtype} {x.device} " + f"{tuple(x.shape)} contiguous={x.is_contiguous()}") + + +def _check_m16_shapes(prob_n: int, prob_k: int): + """L4: kernel tiling constraints (prob_n%256 for n_tiles, prob_k%128 per + pipeline stage). K3 shapes 3072/3584/6144 all pass; anything else is a + wiring bug.""" + if prob_n % 256 != 0 or prob_k % 128 != 0: + raise ValueError( + f"marlin M16 kernel constraint violated: prob_n%256==0 and " + f"prob_k%128==0 required, got prob_n={prob_n}, prob_k={prob_k}") + + +def _check_m_tile_bound(max_m_tiles: int, mtp: int, total_rows: int): + """L5 (host-static, plan-build): CTAs beyond max_m_tiles*16 rows would be + silently dropped. The dispatcher guarantees counts[e] <= min(mtp, + total_rows), so this bound makes drops impossible.""" + admissible = min(int(mtp), int(total_rows)) + if max_m_tiles * 16 < admissible: + raise ValueError( + f"M-tile bound below admissible per-expert tokens: " + f"max_m_tiles={max_m_tiles} covers {max_m_tiles * 16} rows < " + f"min(mtp={mtp}, total_rows={total_rows})={admissible} — CTAs " + f"would silently drop rows") + + +def _check_marlin_mxfp4_tensors(name: str, qw: torch.Tensor, scale: torch.Tensor, + prob_n: int, prob_k: int): + """L2 (tensor-visible call sites only): marlin layout + bf16 scales at the + kernel boundary. E8M0 bytes must be expanded (exactly) at fill time via + marlin_weight_prep.mxfp4_scale_e8m0_to_bf16 — never value-cast.""" + if qw.dtype != torch.int32: + raise ValueError( + f"{name}: marlin_qw must be int32 marlin-packed, got {qw.dtype}") + if scale.dtype != torch.bfloat16: + raise ValueError( + f"{name}: scale dtype != bf16 at kernel boundary (got {scale.dtype}) " + f"— E8M0 uint8 bytes must be expanded exactly at fill " + f"(mxfp4_scale_e8m0_to_bf16), never value-cast, never fed raw") + if tuple(qw.reshape(-1, qw.shape[-1]).shape) != (prob_k // 16, prob_n * 2): + raise ValueError( + f"{name}: marlin_qw shape {tuple(qw.shape)} != " + f"[{prob_k // 16}, {prob_n * 2}] for prob_k={prob_k}, prob_n={prob_n}") + if tuple(scale.reshape(-1, scale.shape[-1]).shape) != (prob_k // 32, prob_n): + raise ValueError( + f"{name}: marlin scale shape {tuple(scale.shape)} != " + f"[{prob_k // 32}, {prob_n}] for prob_k={prob_k}, prob_n={prob_n}") + + +def marlin_grouped_stage1_fused_mxfp4_situ( + dispatched_x_3d: torch.Tensor, + intermediate_3d: torch.Tensor, + expert_counts: torch.Tensor, + expert_starts: torch.Tensor, + gate_B_ptrs: torch.Tensor, + gate_scales_ptrs: torch.Tensor, + up_B_ptrs: torch.Tensor, + up_scales_ptrs: torch.Tensor, + C_ptrs: torch.Tensor, + N: int, + K: int, + workspace: torch.Tensor, + max_m_tiles: int, + mtp: int, + num_experts: int, + total_rows: int, +) -> None: + """K3 fused S1: gate(w1)+up(w3) MXFP4 GEMM + SiTU in a single kernel. + + Same seam as marlin_grouped_stage1_fused (K2.5 production), with: + - E2M1 in-kernel weight decode (dequant_e2m1) instead of (q-8), + - SiTU epilogue (beta=4, linear_beta=25) instead of SiLU, + - hard-fail contract checks L1/L3/L4/L5 + activation contract (dtype/ + device/shape/contiguity of both activation tensors; no warn-and-degrade), + - `total_rows` REQUIRED so the L5 M-tile bound is enforceable here + (K2.5 computes it at plan build; the K3 seam must not trust it). + + Pointer arrays must point at tensors produced by + marlin_weight_prep.repack_mxfp4_to_marlin_gs32 with bf16 scales at the + kernel boundary (L2 is checked at the tensor-visible call sites; the + checkpoint stamp check L6 is model-side). + + Zero-token experts are handled by the kernel's per-CTA early-exit — the + caller must NOT filter empty experts (graph-static pointer arrays). + """ + global _warned_mxfp4 + _require_mxfp4_kernels() + _check_activation("dispatched_x_3d", dispatched_x_3d, K) + _check_activation("intermediate_3d", intermediate_3d, N) + E = int(num_experts) + _check_counts("expert_counts", expert_counts, E) + _check_counts("expert_starts", expert_starts, E) + for name, t in (("gate_B_ptrs", gate_B_ptrs), ("up_B_ptrs", up_B_ptrs), + ("gate_scales_ptrs", gate_scales_ptrs), + ("up_scales_ptrs", up_scales_ptrs), ("C_ptrs", C_ptrs)): + _check_ptr_array(name, t, E) + _check_m16_shapes(N, K) + _check_m_tile_bound(max_m_tiles, mtp, total_rows) + + if not _warned_mxfp4: + logging.info("[Marlin] Using fused M16 Marlin MXFP4 S1 (gate+up+SiTU, K3)") + _warned_mxfp4 = True + + mod = _load_module() + n_tiles = N // 256 + mod.grouped_marlin_gemm_m16_s1_mxfp4_situ( + dispatched_x_3d, + gate_B_ptrs, up_B_ptrs, C_ptrs, + gate_scales_ptrs, up_scales_ptrs, + expert_starts, expert_counts, + E, N, K, workspace, n_tiles, max_m_tiles, + ) + + +def marlin_grouped_m16_mxfp4( + A: torch.Tensor, + B_ptrs: torch.Tensor, + C_ptrs: torch.Tensor, + scales_ptrs: torch.Tensor, + expert_starts: torch.Tensor, + expert_counts: torch.Tensor, + num_experts: int, + prob_n: int, + prob_k: int, + workspace: torch.Tensor, + num_matrices: int, + n_tiles: int, + max_m_tiles: int, +) -> None: + """K3 M16 MXFP4 grouped GEMM (S3 down projection / standalone).""" + _require_mxfp4_kernels() + _check_activation("A", A, prob_k) + E = int(num_experts) + _check_counts("expert_counts", expert_counts, E) + _check_counts("expert_starts", expert_starts, E) + _check_ptr_array("B_ptrs", B_ptrs, num_matrices) + _check_ptr_array("C_ptrs", C_ptrs, num_matrices) + _check_ptr_array("scales_ptrs", scales_ptrs, num_matrices) + _check_m16_shapes(prob_n, prob_k) + if n_tiles != prob_n // 256: + raise ValueError(f"n_tiles={n_tiles} != prob_n//256={prob_n // 256}") + + mod = _load_module() + mod.grouped_marlin_gemm_m16_mxfp4( + A, B_ptrs, C_ptrs, scales_ptrs, + expert_starts, expert_counts, + E, prob_n, prob_k, workspace, num_matrices, n_tiles, max_m_tiles, + ) + + +def single_expert_marlin_mxfp4_decode( + x: torch.Tensor, + gate_qw: torch.Tensor, gate_scale: torch.Tensor, + up_qw: torch.Tensor, up_scale: torch.Tensor, + down_qw: torch.Tensor, down_scale: torch.Tensor, + N: int, K: int, +) -> torch.Tensor: + """K3 W4A16 MXFP4 decode for ONE expert (streamed/offloaded experts). + + Mirror of single_expert_marlin_decode with the _mxfp4 + SiTU kernel + entries and full tensor-level contract checks (this is the one seam where + the real tensors — not just pointers — are visible, so L2 is enforced). + + Args: + x: [t, K] BF16 gathered tokens routed to this expert. + gate_qw/up_qw: [K//16, N*2] int32 marlin MXFP4 (from + repack_mxfp4_to_marlin_gs32); gate = w1, up = w3 (gate-first — + swapped branches are silent, pinned by the GPU mutation test). + gate_scale/up_scale: [K//32, N] bf16 (exact E8M0 expansion). + down_qw: [N//16, K*2] int32; down_scale: [N//32, K] bf16. + N: moe_intermediate_size (K3: 3072); K: hidden_size (K3: 3584). + Returns: [t, K] BF16. + """ + _require_mxfp4_kernels() + _check_activation("x", x, K, require_contiguous=False) # .contiguous() below + _check_m16_shapes(N, K) # S1: prob_n=N, prob_k=K + _check_m16_shapes(K, N) # S3: prob_n=K, prob_k=N + _check_marlin_mxfp4_tensors("gate(w1)", gate_qw, gate_scale, N, K) + _check_marlin_mxfp4_tensors("up(w3)", up_qw, up_scale, N, K) + _check_marlin_mxfp4_tensors("down(w2)", down_qw, down_scale, K, N) + + mod = _load_module() + device = x.device + t = x.shape[0] + x = x.contiguous() + + def _p(tensor): + return torch.tensor([tensor.data_ptr()], dtype=torch.int64, device=device) + + gate_B, up_B, down_B = _p(gate_qw), _p(up_qw), _p(down_qw) + gate_sB, up_sB, down_sB = _p(gate_scale), _p(up_scale), _p(down_scale) + expert_starts = torch.zeros(1, dtype=torch.int32, device=device) + expert_counts = torch.tensor([t], dtype=torch.int32, device=device) + intermediate = torch.empty(t, N, dtype=torch.bfloat16, device=device) + expert_out = torch.empty(t, K, dtype=torch.bfloat16, device=device) + s1_C, s3_C = _p(intermediate), _p(expert_out) + s1_ws = torch.zeros(N // 256 + 17, dtype=torch.int32, device=device) + s3_ws = torch.zeros(K // 256 + 17, dtype=torch.int32, device=device) + max_m_tiles = (t + 15) // 16 + + # Stage 1: fused gate + up + SiTU -> intermediate [t, N] + mod.grouped_marlin_gemm_m16_s1_mxfp4_situ( + x, gate_B, up_B, s1_C, gate_sB, up_sB, + expert_starts, expert_counts, 1, N, K, s1_ws, N // 256, max_m_tiles) + # Stage 3: down -> expert_out [t, K] + mod.grouped_marlin_gemm_m16_mxfp4( + intermediate, down_B, s3_C, down_sB, expert_starts, expert_counts, + 1, K, N, s3_ws, 1, K // 256, max_m_tiles) + return expert_out diff --git a/batchgen/moe/marlin_weight_prep.py b/batchgen/moe/marlin_weight_prep.py index c99f041f..ce928e38 100644 --- a/batchgen/moe/marlin_weight_prep.py +++ b/batchgen/moe/marlin_weight_prep.py @@ -1,9 +1,18 @@ -"""Marlin weight preprocessing: convert K2.5 INT4 → Marlin packed format.""" +"""Marlin weight preprocessing: convert K2.5 INT4 / K3 MXFP4 → Marlin packed format.""" import logging import numpy as np import torch +from batchgen.moe.mxfp4_oracle_vector import ( + MXFP4_E2M1_LUT, + MXFP4_E8M0_BIAS, + MXFP4_FORBIDDEN_SCALE_BYTES, + MXFP4_GROUP_SIZE, + MXFP4_LOW_NIBBLE_IS_EVEN_K, + check_dequant_fn, +) + GPTQ_MARLIN_TILE = 16 INT4_GROUP_SIZE = 32 # K2.5 checkpoint group size @@ -224,3 +233,320 @@ def convert_int4_to_marlin( marlin_s = marlin_s.to(compute_dtype) return marlin_qw, marlin_s + + +# ============================================================================ +# K3 MXFP4 (E2M1 nibbles + E8M0 uint8 scales) → Marlin converter +# +# The Marlin tile permutation moves opaque 4-bit codes; it performs NO +# arithmetic on them, so E2M1 codes ride through the exact machinery the K2.5 +# INT4 production path uses. The ONLY format-specific pieces are: +# (a) the source nibble-extraction convention (low nibble = even K index — +# frozen against the compressed-tensors oracle, see mxfp4_oracle_vector), +# (b) E8M0 uint8 scale handling (index-permute only; optional EXACT bf16 +# materialization via bit shift — never a value cast), +# (c) the in-kernel decode (dequant_e2m1 + SiTU, marlin_grouped_gemm.cu). +# +# FORBIDDEN per the 2026-08-04 POIS decision ledger (HARD-FAIL policy): +# - any decode→requantize path (E2M1 magnitudes are non-uniform; INT4+scale +# re-quantization is a second lossy quantization). Only nibble +# rearrangement is allowed here. Contrast convert_int4_to_marlin above, +# which REQUANTIZES and must never be extended to MXFP4. +# - .to(torch.float16)/.to(float) on scale BYTES or scale VALUES +# (fp16 overflows at byte >= 143; byte-value casts are garbage). +# - silent clamping of E8M0 edge bytes 0x00/0xFF: the repack RAISES instead. +# ============================================================================ + +_mxfp4_convention_verified = False + + +def _unpack_mxfp4_nibbles(weight_packed: torch.Tensor, K: int, N: int) -> torch.Tensor: + """Unpack compressed-tensors mxfp4-pack-quantized [N, K//2] uint8 → [N, K] codes. + + Convention (frozen, oracle-verified): low nibble of byte j = K index 2j, + high nibble = K index 2j+1. Because int32 little-endian byte order composes + with intra-byte low-first, this is bit-identical to the K2.5 raw INT4 + int32 unpacking (nibble i of int32 word w = K index 8w+i) — so we view the + uint8 buffer as int32 and reuse the exact production unpack loop. + + Returns [N, K] int32 with values 0..15 (raw E2M1 codes, NOT decoded). + """ + if not MXFP4_LOW_NIBBLE_IS_EVEN_K: + # The int32-view shortcut below encodes low-first. If the frozen verdict + # ever changes, this function must be rewritten — refuse loudly. + raise ValueError( + "MXFP4_LOW_NIBBLE_IS_EVEN_K is no longer True; _unpack_mxfp4_nibbles " + "hard-codes the low-nibble-first convention and must be updated.") + packed_i32 = weight_packed.contiguous().view(N, K // 2).view(torch.int32) # [N, K//8] + unpacked = torch.empty(N, K // 8, 8, dtype=torch.int32, device=weight_packed.device) + for i in range(8): + unpacked[:, :, i] = (packed_i32 >> (i * 4)) & 0xF + return unpacked.view(N, K) + + +def _mxfp4_dequant_via_unpack(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """Dequant path built on THIS module's unpack helper (used by the R8 self-check).""" + N, half_k = packed.shape + K = half_k * 2 + codes = _unpack_mxfp4_nibbles(packed, K, N) + lut = torch.tensor(list(MXFP4_E2M1_LUT) + [-v for v in MXFP4_E2M1_LUT], + dtype=torch.float32) + vals = lut[codes.long()] + exps = (scales.to(torch.int32) - MXFP4_E8M0_BIAS).repeat_interleave( + MXFP4_GROUP_SIZE, dim=-1) + return torch.ldexp(vals, exps).to(torch.bfloat16) + + +def _verify_mxfp4_convention() -> None: + """R8: one-time self-check of the nibble convention against the frozen + real-checkpoint oracle vector. Raises if the module's unpack logic ever + diverges from the recorded compressed-tensors verdict.""" + global _mxfp4_convention_verified + if _mxfp4_convention_verified: + return + check_dequant_fn(_mxfp4_dequant_via_unpack) + _mxfp4_convention_verified = True + + +def mxfp4_scale_e8m0_to_bf16(scale_u8: torch.Tensor) -> torch.Tensor: + """EXACT E8M0 uint8 → bf16 conversion: 2^(e8 - 127) for e8 in [1, 254]. + + bf16 layout is 1s/8e/7m, so the bit pattern uint16(e8) << 7 IS the value + 2^(e8-127) exactly — no rounding anywhere. This is a bit shift, not a + value cast; power-of-two scales times E2M1 magnitudes (<=2 significant + bits) stay exact in bf16 through the kernel's scale_op. + + Edge bytes RAISE (hard-fail policy): 0x00 would need the bf16 subnormal + 0x0040 (not e8<<7) and 0xFF is NaN per OCP MX — neither occurs in K3 data. + """ + if scale_u8.dtype != torch.uint8: + raise ValueError( + f"mxfp4_scale_e8m0_to_bf16 expects uint8 E8M0 bytes, got {scale_u8.dtype}") + n_bad = int(((scale_u8 == MXFP4_FORBIDDEN_SCALE_BYTES[0]) | + (scale_u8 == MXFP4_FORBIDDEN_SCALE_BYTES[1])).sum()) + if n_bad: + raise ValueError( + f"E8M0 edge byte 0x00/0xFF present (count={n_bad}, observed byte range " + f"[{int(scale_u8.min())}, {int(scale_u8.max())}]): outside the validated " + f"exact window [1, 254]. Edge semantics per the compressed-tensors " + f"verdict (0x00 -> 2^-127, 0xFF -> inf/NaN) — refusing, no silent clamp.") + return (scale_u8.to(torch.int16) << 7).view(torch.bfloat16) + + +def _check_mxfp4_repack_contract( + weight_packed: torch.Tensor, weight_scale: torch.Tensor, K: int, N: int, +) -> None: + """Hard-fail contract checks R1–R7 (ValueError, never assert).""" + # R1 + if weight_packed.dtype != torch.uint8: + raise ValueError( + f"MXFP4 weight_packed must be uint8 [N, K//2], got {weight_packed.dtype} " + f"{tuple(weight_packed.shape)}") + # R2 + if weight_scale.dtype != torch.uint8: + raise ValueError( + f"MXFP4 weight_scale must be uint8 E8M0 [N, K//32], got {weight_scale.dtype} " + f"— bf16 scales mean an INT4 checkpoint; use repack_int4_to_marlin_gs32") + # R3 + if tuple(weight_packed.shape) != (N, K // 2): + raise ValueError( + f"MXFP4 packed dim != K//2: got {tuple(weight_packed.shape)}, " + f"expected ({N}, {K // 2}) for K={K}") + # R4 + if K % MXFP4_GROUP_SIZE != 0 or weight_scale.shape[-1] != K // MXFP4_GROUP_SIZE: + raise ValueError( + f"MXFP4 scale groups != K/32: got {tuple(weight_scale.shape)}, expected " + f"({N}, {K // MXFP4_GROUP_SIZE}) — group size 32 is the only supported " + f"MXFP4 group (kernel GROUP_BLOCKS=2)") + # R5 + if weight_scale.shape[0] != weight_packed.shape[0]: + raise ValueError( + f"N mismatch packed vs scale: {weight_packed.shape[0]} vs " + f"{weight_scale.shape[0]}") + # R6 (K%16 for marlin tiles — already subsumed by R4's K%32, kept as + # defense in depth for the perm math; N%64 because the weight/scale + # permutations act on 64-column blocks) + if K % 16 != 0 or N % 64 != 0: + raise ValueError( + f"Marlin tiling requires K%16==0 and N%64==0 (perm block = 64 N-cols), " + f"got K={K}, N={N}") + # R7 + n_bad = int(((weight_scale == MXFP4_FORBIDDEN_SCALE_BYTES[0]) | + (weight_scale == MXFP4_FORBIDDEN_SCALE_BYTES[1])).sum()) + if n_bad: + raise ValueError( + f"E8M0 edge byte 0x00/0xFF present in weight_scale (count={n_bad}, " + f"observed byte range [{int(weight_scale.min())}, {int(weight_scale.max())}]): " + f"outside the validated exact window; refusing per the hard-fail policy " + f"(no silent clamp).") + + +def repack_mxfp4_to_marlin_gs32( + weight_packed: torch.Tensor, + weight_scale: torch.Tensor, + K: int, N: int, + emit_scale: str = "e8m0", +) -> tuple: + """Repack K3 MXFP4 gs=32 → Marlin tile layout. PURE NIBBLE REARRANGEMENT. + + No value is decoded, quantized, or clamped anywhere in this function. + E2M1 codes move as opaque 4-bit fields through the identical permutation + machinery the K2.5 INT4 production path uses (get_weight_perm(4) + + _marlin_pack_weights); E8M0 scale BYTES are index-permuted only. + + Args: + weight_packed: [N, K//2] uint8 (compressed-tensors mxfp4-pack-quantized; + low nibble = even K index — frozen oracle verdict) + weight_scale: [N, K//32] uint8 E8M0 + K: in_features (K3: 3584 for w1/w3, 3072 for w2) + N: out_features (K3: 3072 for w1/w3, 3584 for w2) + emit_scale: + "e8m0" — carry E8M0 uint8 bytes through (index-permuted only). + For SHM/checkpoint storage; expand to bf16 at slot fill with + mxfp4_scale_e8m0_to_bf16 before the kernel sees them. + "bf16" — materialize the EXACT bf16 values 2^(e8-127) here + (bit shift, provably lossless; see mxfp4_scale_e8m0_to_bf16). + This is the dtype the marlin kernel consumes. + + Returns: + marlin_qw: [K//16, N*2] int32 — same nibbles, Marlin tile order + (byte count unchanged vs source) + marlin_s: [K//32, N] uint8 (emit_scale="e8m0") or bf16 ("bf16"), + Marlin scale-permuted + """ + _check_mxfp4_repack_contract(weight_packed, weight_scale, K, N) + _verify_mxfp4_convention() # R8 + + # Step 1: unpack nibbles [N, K] (raw codes 0-15, NOT decoded) + q_w_nk = _unpack_mxfp4_nibbles(weight_packed, K, N) + + # Step 2: transpose to [K, N] for Marlin layout + q_w = q_w_nk.t().contiguous() + + # Step 3: Marlin tile permutation + packing (value-agnostic) + perm = get_weight_perm(4) + marlin_qw = _marlin_pack_weights(q_w, K, N, perm) + + # Step 4: scales — transpose [N, K//32] → [K//32, N], index-permute. + # NOTE: no dtype conversion here (contrast the INT4 path's .to(float16), + # which would silently corrupt E8M0 bytes — see module banner). + s = weight_scale.t().contiguous() + marlin_s = _marlin_permute_scales(s, K, N, MXFP4_GROUP_SIZE) + if emit_scale == "bf16": + marlin_s = mxfp4_scale_e8m0_to_bf16(marlin_s) + elif emit_scale != "e8m0": + raise ValueError(f"emit_scale must be 'e8m0' or 'bf16', got {emit_scale!r}") + + return marlin_qw, marlin_s + + +def repack_mxfp4_w13_to_marlin_gs32( + w1_packed: torch.Tensor, w1_scale: torch.Tensor, + w3_packed: torch.Tensor, w3_scale: torch.Tensor, + K: int, N: int, + emit_scale: str = "e8m0", +) -> tuple: + """Fused w1‖w3 repack: two complete Marlin tensors, storage-adjacent. + + The fused _s1 kernel takes SEPARATE gate/up pointer arrays and derives its + B row stride from the per-branch prob_n — a column slice of one wide + [K, 2N] marlin tensor has the WRONG stride and can never be passed. The + checkpoint-coordinate concat [2N, K//2] is legal (packing is along K) but + is only the debug-reference layout. Production "fusion" = adjacency: + + qw[0] = gate (w1) marlin tensor, qw[1] = up (w3) marlin tensor + up_qw_ptr = gate_qw_ptr + K//16 * N*2 * 4 bytes (qw is contiguous) + + Gate-first order matches the HF reference (gate_up = cat([w1, w3])) and is + SILENT if swapped — the GPU SiTU parity mutation test is the guard. + + Args (per branch): same contracts as repack_mxfp4_to_marlin_gs32. + K = 3584, N = 3072 for K3 routed experts. + + Returns: + qw: [2, K//16, N*2] int32 contiguous (index 0 = gate/w1, 1 = up/w3) + s: [2, K//32, N] uint8 or bf16 contiguous (same order) + """ + if tuple(w1_packed.shape) != tuple(w3_packed.shape): + raise ValueError( + f"w1/w3 packed shape mismatch: {tuple(w1_packed.shape)} vs " + f"{tuple(w3_packed.shape)}") + if tuple(w1_scale.shape) != tuple(w3_scale.shape): + raise ValueError( + f"w1/w3 scale shape mismatch: {tuple(w1_scale.shape)} vs " + f"{tuple(w3_scale.shape)}") + gate_qw, gate_s = repack_mxfp4_to_marlin_gs32(w1_packed, w1_scale, K, N, emit_scale) + up_qw, up_s = repack_mxfp4_to_marlin_gs32(w3_packed, w3_scale, K, N, emit_scale) + qw = torch.stack([gate_qw, up_qw], dim=0).contiguous() + s = torch.stack([gate_s, up_s], dim=0).contiguous() + return qw, s + + +# --- exact inverse (round-trip proof; local copies of the inverse perms so +# this module stays importable without compiled batchgen_kernels, unlike +# marlin_transform.py which imports the CUDA extension at module level) --- + +def _inverse_weight_perm(num_bits: int = 4) -> torch.Tensor: + perm = get_weight_perm(num_bits) + inv_perm = torch.empty_like(perm) + inv_perm[perm] = torch.arange(len(perm)) + return inv_perm + + +def _inverse_scale_perm() -> list: + scale_perm, _ = _get_scale_perms() + inv = [0] * len(scale_perm) + for i, p in enumerate(scale_perm): + inv[p] = i + return inv + + +def marlin_mxfp4_to_raw_cpu( + marlin_qw: torch.Tensor, + marlin_s: torch.Tensor, + K: int, N: int, +) -> tuple: + """Exact inverse of repack_mxfp4_to_marlin_gs32 (pure rearrangement proof). + + Accepts marlin_s as uint8 (e8m0 passthrough) or bf16 (exact expansion; + inverted losslessly via bits >> 7). + + Returns: + raw_packed: [N, K//2] uint8 — must be BYTE-IDENTICAL to the source + raw_scale: [N, K//32] uint8 + """ + if tuple(marlin_qw.shape) != (K // GPTQ_MARLIN_TILE, N * 2): + raise ValueError( + f"Expected marlin_qw [{K // GPTQ_MARLIN_TILE}, {N * 2}], " + f"got {tuple(marlin_qw.shape)}") + + # Weights: unpack int32 nibbles → inverse tile perm → undo tile transpose + flat = marlin_qw.reshape(-1) + unpacked = torch.empty(flat.numel(), 8, dtype=torch.int32, device=flat.device) + for i in range(8): + unpacked[:, i] = (flat >> (i * 4)) & 0xF + q_marlin = unpacked.view(K // GPTQ_MARLIN_TILE, N * GPTQ_MARLIN_TILE) + inv_perm = _inverse_weight_perm(4).to(flat.device) + q_tiled = q_marlin.reshape(-1, inv_perm.numel())[:, inv_perm].reshape(q_marlin.shape) + q_tiled = q_tiled.reshape(K // GPTQ_MARLIN_TILE, N // GPTQ_MARLIN_TILE, + GPTQ_MARLIN_TILE, GPTQ_MARLIN_TILE) + q_raw_kn = q_tiled.permute(0, 2, 1, 3).reshape(K, N) + q_raw_nk = q_raw_kn.t().contiguous() # [N, K] + + # Pack back to uint8, low nibble = even K index + raw_packed = (q_raw_nk[:, 0::2] | (q_raw_nk[:, 1::2] << 4)).to(torch.uint8) + + # Scales + if marlin_s.dtype == torch.bfloat16: + s_kn = (marlin_s.view(torch.int16) >> 7).to(torch.uint8) + elif marlin_s.dtype == torch.uint8: + s_kn = marlin_s + else: + raise ValueError( + f"marlin_s must be uint8 (e8m0) or bf16 (exact expansion), " + f"got {marlin_s.dtype}") + inv_scale_perm = _inverse_scale_perm() + s_inv = s_kn.reshape(-1, len(inv_scale_perm))[:, inv_scale_perm] + raw_scale = s_inv.reshape(-1, N).t().contiguous() # [N, K//32] + + return raw_packed, raw_scale diff --git a/batchgen/moe/mxfp4_grouped_gemm.py b/batchgen/moe/mxfp4_grouped_gemm.py index b19a63a3..ca82990c 100644 --- a/batchgen/moe/mxfp4_grouped_gemm.py +++ b/batchgen/moe/mxfp4_grouped_gemm.py @@ -1,7 +1,29 @@ -"""Fused MXFP4 dequantization and grouped GEMM for MoE layers. - -This module implements fused dequantization of MXFP4 weights during matrix -multiplication, avoiding the memory overhead of materializing full BF16 weights. +"""MXFP4 utilities + DEAD Triton grouped-GEMM kernels (gpt-oss lineage). + +============================================================================ +!! WARNING — THE GROUPED TRITON KERNELS IN THIS MODULE ARE DEAD CODE !! + +The grouped MXFP4 GEMM surface here (fused_mxfp4_grouped_gemm, +grouped_mxfp4_gemm_3d[_tunable], grouped_mxfp4_moe_forward[_3d,_cuda_routing]) +has ZERO production callers, was NEVER gated by a numerics test, and carries +12 confirmed defects — including an unconditional scale-group-aliasing bug +that produces finite, plausible, WRONG numbers on the first K3-shaped call. +See batchgen_design/model_support/kimi_k3/KERNEL_WORKUNIT.md. + +Per the 2026-08-04 POIS decision: Kimi-K3 MXFP4 MoE numerics run +on the PRODUCTION-PROVEN Marlin machinery instead — + - repack: batchgen/moe/marlin_weight_prep.py::repack_mxfp4_to_marlin_gs32 + - kernels: batchgen_kernels/src/moe/marlin_grouped_gemm.cu + (grouped_marlin_gemm_m16_mxfp4, ..._m16_s1_mxfp4_situ) + - wrappers: batchgen/moe/marlin_grouped_moe.py (hard-fail contracts) +Do NOT wire K3 (or any new model) through the grouped kernels below; they +raise NotImplementedError to keep this path loudly closed (this repo has been +bitten by dead-but-inviting kernels before). + +Still LIVE (gpt-oss-120b + int4_grouped_gemm) and NOT covered by the warning: + moe_token_dispatch, reshape_to_3d_expert_layout, gather_from_3d_expert_layout, + setup_expert_weight_pointers, fused_mxfp4_single_gemm, mxfp4_linear. +============================================================================ MXFP4 Format: - 32 FP4 values packed in 16 bytes (2 values per uint8) @@ -34,6 +56,22 @@ MXFP4_PACKED_BLOCK_SIZE = 16 # Bytes per scale (32 values / 2 per byte) +def _refuse_dead_triton_grouped(entry: str): + """Hard tombstone for the never-validated Triton grouped MXFP4 kernels. + + See the module docstring: 12 confirmed defects, zero callers, no gate ever + existed. The supported MXFP4 grouped path is the Marlin machinery + (batchgen/moe/marlin_grouped_moe.py + marlin_grouped_gemm.cu). + """ + raise NotImplementedError( + f"{entry} is dead, unvalidated code (12 confirmed defects incl. " + f"unconditional scale-group aliasing — see " + f"batchgen_design/model_support/kimi_k3/KERNEL_WORKUNIT.md). " + f"Use the Marlin MXFP4 path: batchgen.moe.marlin_grouped_moe " + f"(grouped_marlin_gemm_m16_s1_mxfp4_situ / grouped_marlin_gemm_m16_mxfp4) " + f"with batchgen.moe.marlin_weight_prep.repack_mxfp4_to_marlin_gs32.") + + @triton.jit def _fp4_lookup(idx): """Lookup FP4 value from 4-bit index (LEGACY - slow, 16 tl.where calls). @@ -390,6 +428,7 @@ def fused_mxfp4_grouped_gemm( Returns: Output tensor [M, N] in BF16 """ + _refuse_dead_triton_grouped("fused_mxfp4_grouped_gemm") assert lhs.dtype == torch.bfloat16, "lhs must be BF16" assert all(r.dtype == torch.uint8 for r in rhs_packed_list), "packed weights must be uint8" assert all(s.dtype == torch.uint8 for s in rhs_scales_list), "scales must be uint8" @@ -758,6 +797,7 @@ def grouped_mxfp4_gemm_3d( Returns: output_3d: [E, M_max, N] in BF16 """ + _refuse_dead_triton_grouped("grouped_mxfp4_gemm_3d") num_experts = hidden_3d.shape[0] M_max = hidden_3d.shape[1] K = hidden_3d.shape[2] @@ -827,6 +867,7 @@ def grouped_mxfp4_gemm_3d_tunable( Returns: output_3d: [E, M_max, N] in BF16 """ + _refuse_dead_triton_grouped("grouped_mxfp4_gemm_3d_tunable") num_experts = hidden_3d.shape[0] M_max = hidden_3d.shape[1] K = hidden_3d.shape[2] @@ -912,6 +953,7 @@ def grouped_mxfp4_moe_forward_3d( Returns: Output [batch*seq, hidden] in BF16 """ + _refuse_dead_triton_grouped("grouped_mxfp4_moe_forward_3d") num_tokens, hidden_size = hidden_states.shape device = hidden_states.device @@ -1024,6 +1066,7 @@ def grouped_mxfp4_moe_forward_cuda_routing( num_local_experts: Number of local experts Other args: Same as grouped_mxfp4_moe_forward_3d """ + _refuse_dead_triton_grouped("grouped_mxfp4_moe_forward_cuda_routing") from batchgen.moe.routing import dispatch_count_gather_cuda, reduce_weighted_scatter_cuda num_tokens, hidden_size = hidden_states.shape @@ -1144,6 +1187,7 @@ def grouped_mxfp4_moe_forward( num_tokens, hidden = hidden_states.shape num_experts = len(gate_weights) num_experts_per_tok = topk_indices.shape[1] + _refuse_dead_triton_grouped("grouped_mxfp4_moe_forward") device = hidden_states.device intermediate_size = gate_weights[0].shape[0] # N dimension diff --git a/batchgen/moe/mxfp4_oracle_vector.py b/batchgen/moe/mxfp4_oracle_vector.py new file mode 100644 index 00000000..b1720c0d --- /dev/null +++ b/batchgen/moe/mxfp4_oracle_vector.py @@ -0,0 +1,147 @@ +"""FROZEN MXFP4 oracle verdict for Kimi-K3 (gate settled 2026-08-04). + +Single source of truth for the MXFP4 packing convention. Production code +(marlin_weight_prep.repack_mxfp4_to_marlin_gs32) and the test suite both import +from here. Do NOT re-derive these facts; they were settled against an +independent oracle on a real K3 checkpoint tensor. + +ORACLE: compressed-tensors 0.17.1 (pure-torch CPU reference for the +"mxfp4-pack-quantized" format K3 is stored in). Source citations (paths within +the installed package): + - Nibble order: compressors/nvfp4/helpers.py:72 (pack) + packed = indices[:, 0] | (indices[:, 1] << 4) + and compressors/nvfp4/helpers.py:96-100 (unpack): low = byte & 0x0F is the + FIRST (even-K) element, high = byte >> 4 is the SECOND (odd-K) element. + - FP4 code: sign-magnitude. bit3 (0x08) = sign, bits0-2 index the E2M1 LUT + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] (helpers.py:29-31, 103-108). + - E8M0 scale: compressors/mx_utils.py:43-44 + scale_float = 2.0 ** (uint8 - 127) # NO clamp, NO special cases + => byte 0x00 -> 2^-127 (subnormal, valid); byte 0xFF -> 2^128 -> +inf in + bf16 (OCP MX spec says 0xFF is NaN; compressed_tensors yields inf). + - Group mapping: quantization/lifecycle/forward_helpers.py:149-151 — + scale[r, j] covers contiguous columns [32*j, 32*j+32) of the unpacked K dim. + - Format registration: compressors/mxfp4/base.py:25-26 (MXFP4PackedCompressor). + +REAL-TENSOR VERDICT (language_model.model.layers.4.block_sparse_moe.experts.0.w1, +model-00005-of-000096.safetensors, packed U8[3072,1792] at abs offset 1268562960, +scale U8[3072,112] at abs offset 1274067984): + - BatchGen batchgen/quantization/mxfp4.py::mxfp4_dequantize_reference is + BIT-EXACT vs MXFP4PackedCompressor.decompress on all 11,010,048 elements + (bf16 and fp32). Nibble convention low-first is CONFIRMED independently. + - Swapped-nibble mutation mismatches 91.72% of elements -> test has teeth. + - Scale-byte range observed: w1/w3 in [112, 122] (2^-15..2^-5), w2 in + [119, 122]. ZERO occurrences of 0x00 or 0xFF in any of the three tensors. + - The exponent clamp to [-126, 127] in quantization/mxfp4.py diverges from + the oracle ONLY for bytes 0x00 (2x too large: 2^-126 vs 2^-127) and 0xFF + (2^127 vs +inf). Neither byte occurs in K3 data; per the HARD-FAIL policy + the marlin-MXFP4 load contract asserts scale bytes not in {0x00, 0xFF} + instead of silently clamping (see repack_mxfp4_to_marlin_gs32). + +Raw-source checksums (bytes sliced straight from the shard, md5 verified +remote==local): + sha256(w1 packed bytes) = cf822517403f5ccb418150e10b303568f617f3099bea6dcc9af3a7b6a48e3501 + sha256(w1 scale bytes) = b1cd3499f23097edbbc3f4dc3304c573dbdc49602404b90bdfb1b3dbb0b4ea92 +""" +import torch + +# --- packing convention constants (from compressed_tensors, independently verified) --- +MXFP4_LOW_NIBBLE_IS_EVEN_K = True # byte & 0x0F -> element 2i; byte >> 4 -> 2i+1 +MXFP4_SIGN_BIT_MASK = 0x08 # bit 3 of each nibble +MXFP4_E2M1_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +MXFP4_E8M0_BIAS = 127 # scale_float = 2 ** (uint8 - 127) +MXFP4_GROUP_SIZE = 32 # scale[r, j] covers K columns [32j, 32j+32) +MXFP4_FORBIDDEN_SCALE_BYTES = (0x00, 0xFF) # never occur in K3; loader must HARD-FAIL +# Observed on layer4/expert0 (w1, w2, w3): all scale bytes within [112, 122]. + +# Identifier stamped into converted-checkpoint metadata by the conversion seam +# (R11 in the task-#34 design). The kernel-side contract checks match on this. +MXFP4_MARLIN_FORMAT_ID = "mxfp4_marlin_gs32_v1" + +# --- frozen real-data test vector: layer 4, expert 0, w1, row 0, group 0 --- +# 16 packed bytes + 1 E8M0 scale byte -> 32 bf16 outputs (exact bit patterns). +VEC_PACKED = bytes([0xAB, 0x4C, 0x18, 0x65, 0x2C, 0x91, 0x04, 0x39, + 0x48, 0x94, 0x58, 0x33, 0x94, 0xB5, 0x5B, 0x30]) +VEC_SCALE_BYTE = 0x79 # 121 -> 2^-6 = 0.015625 +VEC_EXPECTED_BF16_BITS = ( + 0xBCC0, 0xBC80, 0xBD00, 0x3D00, 0x8000, 0x3C00, 0x3D40, 0x3D80, + 0xBD00, 0x3C80, 0x3C00, 0xBC00, 0x3D00, 0x0000, 0xBC00, 0x3CC0, + 0x8000, 0x3D00, 0x3D00, 0xBC00, 0x8000, 0x3D40, 0x3CC0, 0x3CC0, + 0x3D00, 0xBC00, 0x3D40, 0xBCC0, 0xBCC0, 0x3D40, 0x0000, 0x3CC0, +) +# Full-tensor pin: sha256 of the little-endian bf16 buffer of the dequantized +# [3072, 3584] w1 weight (identical from ct's bf16 pipeline and an fp32-exact path). +VEC_FULL_W1_DEQUANT_BF16_SHA256 = \ + "eaeed3d0fd8378496f60174c74737f6c29dd97ec2b854da64b0966ead7f2090f" +VEC_SOURCE = ("Kimi-K3 model-00005-of-000096.safetensors " + "language_model.model.layers.4.block_sparse_moe.experts.0.w1 " + "row 0, K-group 0; oracle compressed-tensors 0.17.1") + + +def mxfp4_dequantize_oracle( + packed: torch.Tensor, + scales: torch.Tensor, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Oracle-faithful MXFP4 dequant (pure torch, CPU-safe, NO exponent clamp). + + Bit-exact reimplementation of compressed_tensors 0.17.1 semantics: + low nibble first, sign-magnitude E2M1 LUT, scale = 2^(uint8 - 127) via + torch.ldexp with no clamp (0x00 -> 2^-127, 0xFF -> inf in bf16). + + This is the parity reference for all marlin-MXFP4 tests. It intentionally + does NOT import batchgen.quantization.mxfp4 (which carries a clamp and a + triton import). + + Args: + packed: [..., K//2] uint8 + scales: [..., K//32] uint8 (E8M0) + Returns: [..., K] in `dtype`. + """ + if packed.dtype != torch.uint8 or scales.dtype != torch.uint8: + raise ValueError( + f"mxfp4_dequantize_oracle expects uint8 packed/scales, got " + f"{packed.dtype}/{scales.dtype}") + lut = torch.tensor( + list(MXFP4_E2M1_LUT) + [-v for v in MXFP4_E2M1_LUT], + dtype=torch.float32, device=packed.device) + idx_lo = (packed & 0x0F).to(torch.long) + idx_hi = (packed >> 4).to(torch.long) + out_shape = packed.shape[:-1] + (packed.shape[-1] * 2,) + vals = torch.empty(out_shape, dtype=torch.float32, device=packed.device) + vals[..., 0::2] = lut[idx_lo] # low nibble -> even K index + vals[..., 1::2] = lut[idx_hi] # high nibble -> odd K index + exponents = scales.to(torch.int32) - MXFP4_E8M0_BIAS # NO clamp (oracle semantics) + exponents = exponents.repeat_interleave(MXFP4_GROUP_SIZE, dim=-1) + return torch.ldexp(vals, exponents).to(dtype) + + +def vec_expected_bf16() -> torch.Tensor: + """The 32 expected dequantized values as bf16 (exact bits).""" + signed = [b - 0x10000 if b >= 0x8000 else b for b in VEC_EXPECTED_BF16_BITS] + return torch.tensor(signed, dtype=torch.int16).view(torch.bfloat16) + + +def vec_packed_tensor() -> torch.Tensor: + return torch.tensor(list(VEC_PACKED), dtype=torch.uint8).unsqueeze(0) # [1, 16] + + +def vec_scale_tensor() -> torch.Tensor: + return torch.tensor([[VEC_SCALE_BYTE]], dtype=torch.uint8) # [1, 1] + + +def check_dequant_fn(dequant_fn) -> None: + """Assert dequant_fn(packed[1,16], scale[1,1]) -> bf16 [1,32] bit-exact. + + Any wrong nibble order, wrong LUT, wrong sign bit, wrong bias, or wrong + group mapping fails this check (verified by mutation on the real tensor). + """ + out = dequant_fn(vec_packed_tensor(), vec_scale_tensor()).reshape(-1) + if out.dtype != torch.bfloat16: + raise AssertionError(f"expected bf16, got {out.dtype}") + exp = vec_expected_bf16() + same = (out.view(torch.int16) == exp.view(torch.int16)) + if not bool(same.all()): + raise AssertionError( + f"MXFP4 dequant mismatch vs frozen compressed-tensors oracle vector " + f"({VEC_SOURCE}): got {out.float().tolist()} " + f"expected {exp.float().tolist()}") diff --git a/batchgen_kernels/src/moe/marlin_grouped_gemm.cu b/batchgen_kernels/src/moe/marlin_grouped_gemm.cu index 4b202673..1a7fd262 100644 --- a/batchgen_kernels/src/moe/marlin_grouped_gemm.cu +++ b/batchgen_kernels/src/moe/marlin_grouped_gemm.cu @@ -118,6 +118,97 @@ __device__ inline void dequant_u4b8(int q, scalar_t2* frag_b) { #endif } +// ---------------------------------------------------------------------------- +// MXFP4 (E2M1) dequant — Kimi-K3. +// +// E2M1 nibble is SIGN-MAGNITUDE: bit3 = sign, bits2:0 = eem index into +// {0, 0.5, 1, 1.5, 2, 3, 4, 6}. The additive magic-number trick used by +// dequant_u4b8 (0x4300 bias / sub 0x4308) is intrinsically uint4-zero-point-8 +// and CANNOT be parameterized to produce sign-magnitude E2M1 — this is a full +// replacement of the decode, with the identical register contract +// (int q in; frag_b[0] = lanes from nibbles at bits[3:0]/[19:16], +// frag_b[1] = lanes from bits[7:4]/[23:20]; caller does q >> 8 for the rest). +// +// Recipe (branch-free): plant eem at bf16 bits [8:6] (low 2 exponent bits + +// mantissa MSB), sign at bit 15, then rebias with one mul by 2^126: +// eem=0 -> bits 0x0000 -> 0.0 (x 2^126 = 0) +// eem=1 -> bits 0x0040 = 2^-127 (subnormal) (x 2^126 = 0.5) +// eem=2e+m,e>0 -> (1+m/2) * 2^(e-127) (x 2^126 = {1,1.5,2,3,4,6}) +// Exact for all 8 magnitudes. NOTE the eem=1 path transits a bf16 SUBNORMAL +// input to mul.rn.bf16x2 (no .ftz variant exists for bf16x2 on sm_90) — the +// GPU parity suite is deliberately +-0.5-heavy to pin this; the fallback if a +// device flushes it is a 2x PRMT byte-LUT (hi {00,3F,3F,3F,40,40,40,40}, +// lo {00,00,80,C0,00,40,80,C0}) + sign OR. +// +// Scales are handled OUTSIDE this function: E8M0 uint8 bytes are expanded to +// exact bf16 powers of two at repack/fill time (marlin_weight_prep. +// mxfp4_scale_e8m0_to_bf16), so scale_op and the whole scale SMEM pipeline are +// reused byte-for-byte. Residual: scale byte 0x01 (2^-126) times a +-0.5 code +// makes the scale_op PRODUCT itself a bf16 subnormal — unreachable for K3 +// (observed scale floor is 112, and repack hard-fails only 0x00/0xFF), but +// revisit this if the legal scale window is ever widened toward the bottom. +// ---------------------------------------------------------------------------- +__device__ inline void dequant_e2m1(int q, scalar_t2* frag_b) { +#if defined(USE_BF16_COMPUTE) + static constexpr uint32_t EEM = 0x00070007; // magnitude bits per lane + static constexpr uint32_t SGN = 0x00080008; // sign bit per lane + static constexpr uint32_t REBIAS = 0x7E807E80; // bf16x2 {2^126, 2^126} + uint32_t uq = (uint32_t)q; + uint32_t lo = ((uq & EEM) << 6) | ((uq & SGN) << 12); + uint32_t uq_hi = uq >> 4; + uint32_t hi = ((uq_hi & EEM) << 6) | ((uq_hi & SGN) << 12); + uint32_t res_lo, res_hi; + asm("mul.rn.bf16x2 %0, %1, %2;\n" : "=r"(res_lo) : "r"(lo), "r"(REBIAS)); + asm("mul.rn.bf16x2 %0, %1, %2;\n" : "=r"(res_hi) : "r"(hi), "r"(REBIAS)); + reinterpret_cast(frag_b)[0] = res_lo; + reinterpret_cast(frag_b)[1] = res_hi; +#else + // Both build registrations (setup.py + _jit_registry.py) define + // USE_BF16_COMPUTE; an FP16 build of the MXFP4 path was never validated and + // must fail loudly rather than ship an untested decode. +#error "Marlin MXFP4 (E2M1) dequant requires USE_BF16_COMPUTE" +#endif +} + +// ---------------------------------------------------------------------------- +// Compile-time functors: weight codec + fused-epilogue activation. +// The U4B8/SILU instantiations must stay semantically identical to the +// pre-template K2.5 production kernels (if-constexpr resolves at compile time; +// verify with -Xptxas -v / SASS diff on the GPU stage). +// ---------------------------------------------------------------------------- +enum class WCodec { U4B8, E2M1 }; +enum class Act { SILU, SITU }; + +template +__device__ inline void dequant_w4(int q, scalar_t2* frag_b) { + if constexpr (CODEC == WCodec::E2M1) { + dequant_e2m1(q, frag_b); + } else { + dequant_u4b8(q, frag_b); + } +} + +// Fused S1 epilogue: combine gate-branch value g (pass 1 / w1) with +// linear-branch value u (pass 2 / w3). fp32 scalar, per output element. +template +__device__ inline float act_gate_mul(float g, float u) { + if constexpr (ACT == Act::SITU) { + // Kimi-K3 SiTU (modeling_kimi_linear.py:75-82; config beta=4.0, + // linear_beta=25.0). fp32 interior, tanh-soft-clamped both branches: + // situ_a = 4 * tanh(g/4) * sigmoid(g) in (-0.2698, 4) + // u_c = 25 * tanh(u/25) in (-25, 25) + // NOTE branch order is SILENT if swapped — pinned by the GPU mutation + // test. Under --use_fast_math tanhf/__expf lower to SFU approximations + // (~2^-11 rel err), inside the 1.6e-2 parity gate; de-fast-math this + // epilogue only if parity fails. + float situ_a = 4.0f * tanhf(0.25f * g) * (1.0f / (1.0f + __expf(-g))); + float u_c = 25.0f * tanhf(0.04f * u); + return situ_a * u_c; + } else { + return g / (1.0f + __expf(-g)) * u; + } +} + // mma_trans: swaps A and B operand positions for m_block_size_8 // B fragments go into the 4-register A operand slot (interleaved b0, b1) // A fragment goes into the 2-register B operand slot (ldsm<2>) @@ -542,9 +633,13 @@ static constexpr int m16_sh_a_size = STAGES * m16_a_sh_stage; // 1024 static constexpr int sh_gate_size = sh_red_size; // 528 int4 — stores gate BF16 result // ============================================================================ -// Fused S1 kernel: gate+up+SiLU in single kernel, no temp buffer +// Fused S1 kernel: gate+up+activation in single kernel, no temp buffer. +// Templated on weight codec (U4B8 = K2.5 INT4, E2M1 = K3 MXFP4) and epilogue +// activation (SILU = K2.5, SITU = K3). is the production K2.5 +// instantiation and must stay bit-identical to the pre-template kernel. // ============================================================================ +template __global__ void MarlinGrouped_M16_S1( const int4* __restrict__ A, const int4* const* __restrict__ gate_B_ptrs, // [E] gate weight ptrs @@ -696,8 +791,8 @@ __global__ void MarlinGrouped_M16_S1( FragB frag_b0, frag_b1; int b_quant_0 = frag_b_quant[k2][0][j]; int b_quant_1 = b_quant_0 >> 8; - dequant_u4b8(b_quant_0, reinterpret_cast(&frag_b0)); - dequant_u4b8(b_quant_1, reinterpret_cast(&frag_b1)); + dequant_w4(b_quant_0, reinterpret_cast(&frag_b0)); + dequant_w4(b_quant_1, reinterpret_cast(&frag_b1)); scale_op(frag_b0, frag_s[k2][j], 0); scale_op(frag_b1, frag_s[k2][j], 1); #pragma unroll @@ -890,7 +985,8 @@ __global__ void MarlinGrouped_M16_S1( } // ================================================================ - // FUSED WRITE-BACK: SiLU(gate) * up → output C + // FUSED WRITE-BACK: act(gate, up) → output C + // (SILU: SiLU(gate) * up — K2.5; SITU: Kimi-K3, see act_gate_mul) // ================================================================ { int c_gl_stride = prob_n / 8; @@ -918,7 +1014,7 @@ __global__ void MarlinGrouped_M16_S1( for (int k = 0; k < 8; k++) { float g = num2float(g_ptr[k]); float u = num2float(u_ptr[k]); - r_ptr[k] = float2num(g / (1.0f + __expf(-g)) * u); + r_ptr[k] = float2num(act_gate_mul(g, u)); } C[c_gl_wr] = result; c_gl_wr += c_gl_wr_delta; @@ -930,8 +1026,10 @@ __global__ void MarlinGrouped_M16_S1( // ============================================================================ // M16 kernel (for S2 and standalone use) +// Templated on weight codec; is the production K2.5 instantiation. // ============================================================================ +template __global__ void MarlinGrouped_M16( const int4* __restrict__ A, const int4* const* __restrict__ B_ptrs, @@ -1105,8 +1203,8 @@ __global__ void MarlinGrouped_M16( FragB frag_b0, frag_b1; int b_quant_0 = frag_b_quant[k2][0][j]; int b_quant_1 = b_quant_0 >> 8; - dequant_u4b8(b_quant_0, reinterpret_cast(&frag_b0)); - dequant_u4b8(b_quant_1, reinterpret_cast(&frag_b1)); + dequant_w4(b_quant_0, reinterpret_cast(&frag_b0)); + dequant_w4(b_quant_1, reinterpret_cast(&frag_b1)); scale_op(frag_b0, frag_s[k2][j], 0); scale_op(frag_b1, frag_s[k2][j], 1); #pragma unroll @@ -1367,24 +1465,25 @@ void silu_mul_scatter( num_experts, compact_stride, output_stride, N); } -void grouped_marlin_gemm_m16_s1( - torch::Tensor A, - torch::Tensor gate_B_ptrs, torch::Tensor up_B_ptrs, - torch::Tensor C_ptrs, - torch::Tensor gate_scales_ptrs, torch::Tensor up_scales_ptrs, - torch::Tensor expert_starts, torch::Tensor expert_counts, +template +static void launch_m16_s1( + torch::Tensor& A, + torch::Tensor& gate_B_ptrs, torch::Tensor& up_B_ptrs, + torch::Tensor& C_ptrs, + torch::Tensor& gate_scales_ptrs, torch::Tensor& up_scales_ptrs, + torch::Tensor& expert_starts, torch::Tensor& expert_counts, int num_experts, int prob_n, int prob_k, - torch::Tensor workspace, int n_tiles, int max_m_tiles) + int n_tiles, int max_m_tiles) { auto stream = at::cuda::getCurrentCUDAStream(); // SMEM: M16 base (90112) + gate result (528 * 16 = 8448) = 98560 bytes constexpr int smem_bytes = 98560; - cudaFuncSetAttribute((void*)MarlinGrouped_M16_S1, + cudaFuncSetAttribute((void*)MarlinGrouped_M16_S1, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); int total_ctas = n_tiles * max_m_tiles * num_experts; - MarlinGrouped_M16_S1<<>>( + MarlinGrouped_M16_S1<<>>( reinterpret_cast(A.data_ptr()), reinterpret_cast(gate_B_ptrs.data_ptr()), reinterpret_cast(up_B_ptrs.data_ptr()), @@ -1397,23 +1496,107 @@ void grouped_marlin_gemm_m16_s1( n_tiles, max_m_tiles); } -void grouped_marlin_gemm_m16( - torch::Tensor A, torch::Tensor B_ptrs, torch::Tensor C_ptrs, - torch::Tensor scales_ptrs, +void grouped_marlin_gemm_m16_s1( + torch::Tensor A, + torch::Tensor gate_B_ptrs, torch::Tensor up_B_ptrs, + torch::Tensor C_ptrs, + torch::Tensor gate_scales_ptrs, torch::Tensor up_scales_ptrs, torch::Tensor expert_starts, torch::Tensor expert_counts, int num_experts, int prob_n, int prob_k, - torch::Tensor workspace, int num_matrices, int n_tiles, - int max_m_tiles) + torch::Tensor workspace, int n_tiles, int max_m_tiles) +{ + launch_m16_s1( + A, gate_B_ptrs, up_B_ptrs, C_ptrs, gate_scales_ptrs, up_scales_ptrs, + expert_starts, expert_counts, num_experts, prob_n, prob_k, + n_tiles, max_m_tiles); +} + +// ---------------------------------------------------------------------------- +// K3 MXFP4 entry-point hard-fail checks. The established integration pattern +// (K2.5: kimi_k25/model.py calls the pybind symbols directly) bypasses the +// python wrappers entirely, so the C++ entries must hard-fail on malformed +// metadata themselves: distinct symbols stop codec mixups, these TORCH_CHECKs +// stop everything else that is visible host-side. INT4 entries are left +// byte-identical to production on purpose. +// ---------------------------------------------------------------------------- +static void check_mxfp4_ptr_array(const torch::Tensor& t, const char* name, + int64_t len) { + TORCH_CHECK(t.scalar_type() == at::kLong && t.is_cuda() && t.numel() == len, + "mxfp4 launch: ", name, " must be int64 CUDA [", len, "], got ", + t.scalar_type(), " numel=", t.numel()); +} + +static void check_mxfp4_counts(const torch::Tensor& t, const char* name, + int64_t len) { + TORCH_CHECK(t.scalar_type() == at::kInt && t.is_cuda() && t.numel() == len, + "mxfp4 launch: ", name, " must be int32 CUDA [", len, "], got ", + t.scalar_type(), " numel=", t.numel()); +} + +static void check_mxfp4_common(const torch::Tensor& A, + const torch::Tensor& expert_starts, + const torch::Tensor& expert_counts, + int num_experts, int prob_n, int prob_k, + int n_tiles, int max_m_tiles) { + TORCH_CHECK(A.scalar_type() == at::kBFloat16 && A.is_cuda() + && A.is_contiguous() && A.size(-1) == prob_k, + "mxfp4 launch: A must be contiguous bf16 CUDA [*, ", prob_k, + "], got ", A.scalar_type(), " last dim ", A.size(-1)); + TORCH_CHECK(prob_n % 256 == 0 && prob_k % 128 == 0, + "mxfp4 launch: prob_n%256==0 and prob_k%128==0 required, got prob_n=", + prob_n, " prob_k=", prob_k); + TORCH_CHECK(n_tiles == prob_n / 256, + "mxfp4 launch: n_tiles=", n_tiles, " != prob_n/256=", prob_n / 256); + TORCH_CHECK(num_experts >= 1 && max_m_tiles >= 1, + "mxfp4 launch: num_experts=", num_experts, " max_m_tiles=", + max_m_tiles, " must both be >= 1"); + check_mxfp4_counts(expert_starts, "expert_starts", num_experts); + check_mxfp4_counts(expert_counts, "expert_counts", num_experts); +} + +// K3 MXFP4 fused S1: E2M1 weight decode + SiTU epilogue. Separate entry point +// on purpose — pointer arrays are opaque to the kernel, so distinct pybind +// symbols are the enforceable seam preventing INT4 entries from silently +// consuming E2M1 codes ((q-8)*s on E2M1 codes is finite plausible garbage). +void grouped_marlin_gemm_m16_s1_mxfp4_situ( + torch::Tensor A, + torch::Tensor gate_B_ptrs, torch::Tensor up_B_ptrs, + torch::Tensor C_ptrs, + torch::Tensor gate_scales_ptrs, torch::Tensor up_scales_ptrs, + torch::Tensor expert_starts, torch::Tensor expert_counts, + int num_experts, int prob_n, int prob_k, + torch::Tensor workspace, int n_tiles, int max_m_tiles) +{ + check_mxfp4_common(A, expert_starts, expert_counts, + num_experts, prob_n, prob_k, n_tiles, max_m_tiles); + check_mxfp4_ptr_array(gate_B_ptrs, "gate_B_ptrs", num_experts); + check_mxfp4_ptr_array(up_B_ptrs, "up_B_ptrs", num_experts); + check_mxfp4_ptr_array(C_ptrs, "C_ptrs", num_experts); + check_mxfp4_ptr_array(gate_scales_ptrs, "gate_scales_ptrs", num_experts); + check_mxfp4_ptr_array(up_scales_ptrs, "up_scales_ptrs", num_experts); + launch_m16_s1( + A, gate_B_ptrs, up_B_ptrs, C_ptrs, gate_scales_ptrs, up_scales_ptrs, + expert_starts, expert_counts, num_experts, prob_n, prob_k, + n_tiles, max_m_tiles); +} + +template +static void launch_m16( + torch::Tensor& A, torch::Tensor& B_ptrs, torch::Tensor& C_ptrs, + torch::Tensor& scales_ptrs, + torch::Tensor& expert_starts, torch::Tensor& expert_counts, + int num_experts, int prob_n, int prob_k, + int num_matrices, int n_tiles, int max_m_tiles) { auto stream = at::cuda::getCurrentCUDAStream(); // SMEM: max(528, 4096) + 512 + 1024 = 5632 int4 = 90112 bytes constexpr int smem_bytes = 90112; - cudaFuncSetAttribute((void*)MarlinGrouped_M16, + cudaFuncSetAttribute((void*)MarlinGrouped_M16, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); int total_ctas = n_tiles * max_m_tiles * num_matrices; - MarlinGrouped_M16<<>>( + MarlinGrouped_M16<<>>( reinterpret_cast(A.data_ptr()), reinterpret_cast(B_ptrs.data_ptr()), reinterpret_cast(C_ptrs.data_ptr()), @@ -1424,6 +1607,41 @@ void grouped_marlin_gemm_m16( n_tiles, max_m_tiles); } +void grouped_marlin_gemm_m16( + torch::Tensor A, torch::Tensor B_ptrs, torch::Tensor C_ptrs, + torch::Tensor scales_ptrs, + torch::Tensor expert_starts, torch::Tensor expert_counts, + int num_experts, int prob_n, int prob_k, + torch::Tensor workspace, int num_matrices, int n_tiles, + int max_m_tiles) +{ + launch_m16( + A, B_ptrs, C_ptrs, scales_ptrs, expert_starts, expert_counts, + num_experts, prob_n, prob_k, num_matrices, n_tiles, max_m_tiles); +} + +// K3 MXFP4 M16 (S3 down projection / standalone): E2M1 weight decode. +void grouped_marlin_gemm_m16_mxfp4( + torch::Tensor A, torch::Tensor B_ptrs, torch::Tensor C_ptrs, + torch::Tensor scales_ptrs, + torch::Tensor expert_starts, torch::Tensor expert_counts, + int num_experts, int prob_n, int prob_k, + torch::Tensor workspace, int num_matrices, int n_tiles, + int max_m_tiles) +{ + check_mxfp4_common(A, expert_starts, expert_counts, + num_experts, prob_n, prob_k, n_tiles, max_m_tiles); + TORCH_CHECK(num_matrices >= num_experts, + "mxfp4 launch: num_matrices=", num_matrices, " < num_experts=", + num_experts); + check_mxfp4_ptr_array(B_ptrs, "B_ptrs", num_matrices); + check_mxfp4_ptr_array(C_ptrs, "C_ptrs", num_matrices); + check_mxfp4_ptr_array(scales_ptrs, "scales_ptrs", num_matrices); + launch_m16( + A, B_ptrs, C_ptrs, scales_ptrs, expert_starts, expert_counts, + num_experts, prob_n, prob_k, num_matrices, n_tiles, max_m_tiles); +} + void silu_mul_dual_stride( torch::Tensor gate_inplace, torch::Tensor up, torch::Tensor expert_counts, @@ -1446,6 +1664,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("grouped_marlin_gemm", &grouped_marlin_gemm, "Marlin M8 grouped GEMM"); m.def("grouped_marlin_gemm_m16", &grouped_marlin_gemm_m16, "Marlin M16 grouped GEMM with CTA M-tiling"); m.def("grouped_marlin_gemm_m16_s1", &grouped_marlin_gemm_m16_s1, "Marlin M16 fused S1 (gate+up+SiLU)"); + m.def("grouped_marlin_gemm_m16_mxfp4", &grouped_marlin_gemm_m16_mxfp4, + "Marlin M16 grouped GEMM, MXFP4 (E2M1) weights — Kimi-K3"); + m.def("grouped_marlin_gemm_m16_s1_mxfp4_situ", &grouped_marlin_gemm_m16_s1_mxfp4_situ, + "Marlin M16 fused S1, MXFP4 (E2M1) weights + SiTU epilogue — Kimi-K3"); m.def("silu_mul", &silu_mul, "Element-wise SiLU(gate) * up"); m.def("silu_mul_scatter", &silu_mul_scatter, "SiLU with expert_counts scatter"); m.def("silu_mul_dual_stride", &silu_mul_dual_stride, "SiLU with dual-stride layout"); diff --git a/tests/moe/__init__.py b/tests/moe/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/moe/_loader.py b/tests/moe/_loader.py new file mode 100644 index 00000000..b96b988f --- /dev/null +++ b/tests/moe/_loader.py @@ -0,0 +1,35 @@ +"""Import batchgen.moe MXFP4/marlin modules for CPU tests. + +On GPU machines with the full package installed, a plain import works. On +CPU-only dev machines, executing batchgen/__init__.py fails (it imports the +client and checks compiled batchgen_kernels), so we fall back to lightweight +namespace-package stubs pointing at the real source directories — the target +modules themselves (mxfp4_oracle_vector, marlin_weight_prep) import only +torch/numpy/logging and are fully CPU-safe. +""" + +import importlib +import sys +import types +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def load_moe_modules(): + """Return (mxfp4_oracle_vector, marlin_weight_prep) modules.""" + try: + oracle = importlib.import_module("batchgen.moe.mxfp4_oracle_vector") + mwp = importlib.import_module("batchgen.moe.marlin_weight_prep") + return oracle, mwp + except Exception: + pass + + for pkg, rel in (("batchgen", "batchgen"), ("batchgen.moe", "batchgen/moe")): + if pkg not in sys.modules: + mod = types.ModuleType(pkg) + mod.__path__ = [str(REPO_ROOT / rel)] + sys.modules[pkg] = mod + oracle = importlib.import_module("batchgen.moe.mxfp4_oracle_vector") + mwp = importlib.import_module("batchgen.moe.marlin_weight_prep") + return oracle, mwp diff --git a/tests/moe/gpu_parity_mxfp4_marlin.py b/tests/moe/gpu_parity_mxfp4_marlin.py new file mode 100644 index 00000000..4dc9ee4f --- /dev/null +++ b/tests/moe/gpu_parity_mxfp4_marlin.py @@ -0,0 +1,616 @@ +#!/usr/bin/env python3 +"""STAGED GPU parity ladder for the K3 Marlin-MXFP4 kernels. + +NOT run by the CPU workflow. Staged for a dedicated CUDA GPU (keep it off +the GPU the model workstream is using). No benchmarks here — parity only. + +How to run (on the GPU machine): + cd # sync via git pull — never scp into the repo + + # Build + register-tier gate (runs `setup.py build_ext --inplace` with + # `-Xptxas -v` captured, greps per-kernel register counts, and FAILS if an + # E2M1 instantiation jumps more than 16 regs past its U4B8 counterpart — + # the INT4 M16 tier is ~130 regs / 12.5%% occupancy): + CUDA_VISIBLE_DEVICES=1 python tests/moe/gpu_parity_mxfp4_marlin.py --build + + CUDA_VISIBLE_DEVICES=1 python tests/moe/gpu_parity_mxfp4_marlin.py --smoke + CUDA_VISIBLE_DEVICES=1 python tests/moe/gpu_parity_mxfp4_marlin.py + +Ladder (tolerance gate per KERNEL_WORKUNIT.md: tol = 1e-5 + 1.6e-2*|ref|, +PASS iff finite AND fail_frac < 1e-4, plus max relative error < 1.6e-2 on the +well-conditioned subset |ref| > 0.1*rms(ref). fail_frac == 0.0 is NOT +asserted — fp32 summation-order noise on cancelled outputs): + + T0 SMOKE single expert, t=16, K3 shapes, finite output. + T1 DECODE bit-exact in-kernel E2M1+E8M0 decode via one-hot activations + (no accumulation => exactness is legitimate). +-0.5-heavy + codes specifically pin the bf16-subnormal rebias path of + dequant_e2m1 (eem=1 -> 0x0040 * 2^126). Zeros are + CANONICALIZED before the bit compare: nibble 0x8 decodes to + -0.0 in the reference, but the kernel's fp32 accumulator + starts at +0.0 and (+0.0) + (1.0 * -0.0) = +0.0 under IEEE + RN, so a CORRECT kernel emits +0.0 there. Bits stay strict + for every nonzero output. + T2 M16 grouped_marlin_gemm_m16_mxfp4 vs oracle-dequant + matmul at + both K3 shapes, M sweep {1,16,63,64,65,512,4096}. + T3 S1+SiTU grouped_marlin_gemm_m16_s1_mxfp4_situ vs eager fp32 SiTU + reference (modeling_kimi_linear.py:75-82 semantics), M sweep. + T4 GROUPED E=32 dense grid with >half zero-token experts (caller does + NOT filter empties — C-EMPTYGRP). Uses the PRODUCTION fused + layout: per-expert stacked [2, K//16, N*2] blobs from + repack_mxfp4_w13_to_marlin_gs32, with up pointers derived by + byte arithmetic off the gate pointers (storage adjacency). + The all-empty launch must write NOTHING (asserted via a NaN + canary in the output buffer). + T5 MUTATIONS each deliberately-broken arm must FAIL the gate: + m1 INT4 (u4b8+SiLU) entry consuming E2M1 tensors + m2 gate/up pointer swap (silent at kernel level otherwise) + m3 off-by-one scale group (rolled marlin scales) + m4 SiLU-instead-of-SiTU reference vs the SiTU kernel + Reports the catch count — must be 4/4. + T6 REGRESSION K2.5 INT4 M16 through the templated instantiation + still passes its own parity (templating did not disturb the + production kernels). + T6b REGRESSION K2.5 fused S1 — the decode DEFAULT and the + kernel that got the heavier template surgery (epilogue + swapped to act_gate_mul) — vs SiLU(gate)*up reference. + T7 NEGATIVES every wrapper hard-fail check (L2/L3/L4/L5 + activation + contract) and the raw-pybind TORCH_CHECK seam must RAISE. + Reports the catch count — must be N/N. +""" + +import argparse +import os +import re +import subprocess +import sys +from pathlib import Path + +import torch + +from batchgen.moe import marlin_grouped_moe as mgm +from batchgen.moe import marlin_weight_prep as mwp +from batchgen.moe import mxfp4_oracle_vector as oracle + +DEV = "cuda" +K3_K, K3_N = 3584, 3072 # w1/w3 branch: prob_k, prob_n +M_SWEEP = [1, 16, 63, 64, 65, 512, 4096] +SCALE_LO, SCALE_HI = 112, 122 # observed K3 range (frozen verdict) + +_results = [] + + +def report(name, ok, detail=""): + _results.append((name, ok)) + print(f"[{'PASS' if ok else 'FAIL'}] {name} {detail}") + + +def gate(out, ref, name): + """Project numerical gate. Returns True iff the gate PASSES.""" + out = out.float() + ref = ref.float() + finite = bool(torch.isfinite(out).all()) + err = (out - ref).abs() + tol = 1e-5 + 1.6e-2 * ref.abs() + fail_frac = float((err > tol).float().mean()) + rms = float(ref.pow(2).mean().sqrt()) + mask = ref.abs() > 0.1 * rms + max_rel = float((err[mask] / ref.abs()[mask]).max()) if mask.any() else 0.0 + passed = finite and fail_frac < 1e-4 and max_rel < 1.6e-2 + tag = "gate-PASS" if passed else "gate-FAIL" + print(f" {name}: {tag} fail_frac={fail_frac:.2e} max_rel={max_rel:.2e} " + f"finite={finite}") + return passed + + +def rand_expert(K, N, seed, half_heavy=False, scale_lo=SCALE_LO, scale_hi=SCALE_HI): + """Random MXFP4 expert weight: packed [N, K//2] u8 + scale [N, K//32] u8.""" + g = torch.Generator().manual_seed(seed) + if half_heavy: + # mostly +-0.5 codes (nibbles 0x1/0x9) to hammer the subnormal path + nib = torch.where(torch.rand(N, K, generator=g) < 0.8, + torch.where(torch.rand(N, K, generator=g) < 0.5, + torch.tensor(0x1), torch.tensor(0x9)), + torch.randint(0, 16, (N, K), generator=g)) + packed = (nib[:, 0::2] | (nib[:, 1::2] << 4)).to(torch.uint8) + else: + packed = torch.randint(0, 256, (N, K // 2), generator=g, + dtype=torch.int16).to(torch.uint8) + scale = torch.randint(scale_lo, scale_hi + 1, (N, K // 32), generator=g, + dtype=torch.int16).to(torch.uint8) + return packed, scale + + +def marlinize(packed, scale, K, N): + """CPU repack (proven by the CPU suite) -> CUDA marlin tensors, bf16 scales.""" + qw, s = mwp.repack_mxfp4_to_marlin_gs32(packed, scale, K, N, emit_scale="bf16") + return qw.to(DEV), s.to(DEV) + + +def situ_ref_fp32(g, u): + """Eager SiTU (modeling_kimi_linear.py:75-82; beta=4, linear_beta=25).""" + g = g.float() + u = u.float() + a = 4.0 * torch.tanh(g / 4.0) * torch.sigmoid(g) + return a * (25.0 * torch.tanh(u / 25.0)) + + +def silu_ref_fp32(g, u): + g = g.float() + u = u.float() + return g * torch.sigmoid(g) * u + + +def dense_expert_bf16(packed, scale): + """Oracle-dequant reference weight [N, K] bf16 (exact).""" + return oracle.mxfp4_dequantize_oracle(packed, scale).to(DEV) + + +# --------------------------------------------------------------------------- + +def t0_smoke(): + torch.manual_seed(0) + w = {} + for name, (K, N) in (("w1", (K3_K, K3_N)), ("w3", (K3_K, K3_N)), + ("w2", (K3_N, K3_K))): + p, s = rand_expert(K, N, seed=hash(name) % 2**31) + w[name] = (marlinize(p, s, K, N), (p, s)) + x = torch.randn(16, K3_K, dtype=torch.bfloat16, device=DEV) + out = mgm.single_expert_marlin_mxfp4_decode( + x, + *w["w1"][0], *w["w3"][0], *w["w2"][0], + N=K3_N, K=K3_K) + report("T0 smoke", bool(torch.isfinite(out).all()), + f"out {tuple(out.shape)} finite") + + +def t1_decode_bitexact(): + """One-hot activations => out row m = dequant(W)[:, k_m] with a single + nonzero product per accumulator: bit-exact vs the oracle is legitimate.""" + K, N = 512, 512 + ok = True + for tag, heavy, lo, hi in (("uniform", False, 100, 140), + ("half-heavy", True, SCALE_LO, SCALE_HI)): + p, s = rand_expert(K, N, seed=101, half_heavy=heavy, + scale_lo=lo, scale_hi=hi) + qw, ms = marlinize(p, s, K, N) + w_ref = dense_expert_bf16(p, s) # [N, K] + k_idx = torch.randperm(K)[:64] + A = torch.zeros(64, K, dtype=torch.bfloat16, device=DEV) + A[torch.arange(64), k_idx] = 1.0 + + mod = mgm._load_module() + C = torch.empty(64, N, dtype=torch.bfloat16, device=DEV) + starts = torch.zeros(1, dtype=torch.int32, device=DEV) + counts = torch.tensor([64], dtype=torch.int32, device=DEV) + ws = torch.zeros(N // 256 + 17, dtype=torch.int32, device=DEV) + + def ptr(t): + return torch.tensor([t.data_ptr()], dtype=torch.int64, device=DEV) + + mod.grouped_marlin_gemm_m16_mxfp4( + A, ptr(qw), ptr(C), ptr(ms), starts, counts, + 1, N, K, ws, 1, N // 256, 4) + torch.cuda.synchronize() + + ref = w_ref[:, k_idx].t().contiguous() # [64, N] bf16 exact + # Canonicalize +-0.0 before the bit compare (see ladder docstring): + # nibble 0x8 puts -0.0 (0x8000) in ref, the kernel's +0-seeded fp32 + # accumulator legitimately yields +0.0 (0x0000). Nonzero bits strict. + Cc, refc = C.cpu(), ref.cpu() + Cc = torch.where(Cc == 0, Cc.abs(), Cc) + refc = torch.where(refc == 0, refc.abs(), refc) + exact = torch.equal(Cc.view(torch.int16), refc.view(torch.int16)) + ok = ok and exact + print(f" T1[{tag}]: bit-exact={exact}") + report("T1 in-kernel E2M1 decode bit-exact (incl. +-0.5 subnormal path)", ok) + + +def t2_m16_parity(): + ok = True + for shape_tag, (K, N) in (("w13-branch", (K3_K, K3_N)), ("w2", (K3_N, K3_K))): + p, s = rand_expert(K, N, seed=202) + qw, ms = marlinize(p, s, K, N) + w_ref = dense_expert_bf16(p, s) + mod = mgm._load_module() + for M in M_SWEEP: + A = torch.randn(M, K, dtype=torch.bfloat16, device=DEV) + C = torch.empty(M, N, dtype=torch.bfloat16, device=DEV) + starts = torch.zeros(1, dtype=torch.int32, device=DEV) + counts = torch.tensor([M], dtype=torch.int32, device=DEV) + ws = torch.zeros(N // 256 + 17, dtype=torch.int32, device=DEV) + ptr = lambda t: torch.tensor([t.data_ptr()], dtype=torch.int64, device=DEV) + mod.grouped_marlin_gemm_m16_mxfp4( + A, ptr(qw), ptr(C), ptr(ms), starts, counts, + 1, N, K, ws, 1, N // 256, (M + 15) // 16) + torch.cuda.synchronize() + ref = A.float() @ w_ref.float().t() + ok = gate(C, ref, f"T2[{shape_tag}] M={M}") and ok + report("T2 M16 MXFP4 GEMM parity (K3 shapes, M sweep)", ok) + + +def _run_s1(x, gate_qw, gate_s, up_qw, up_s, N, K, entry): + mod = mgm._load_module() + t = x.shape[0] + inter = torch.empty(t, N, dtype=torch.bfloat16, device=DEV) + starts = torch.zeros(1, dtype=torch.int32, device=DEV) + counts = torch.tensor([t], dtype=torch.int32, device=DEV) + ws = torch.zeros(N // 256 + 17, dtype=torch.int32, device=DEV) + ptr = lambda tt: torch.tensor([tt.data_ptr()], dtype=torch.int64, device=DEV) + getattr(mod, entry)( + x, ptr(gate_qw), ptr(up_qw), ptr(inter), ptr(gate_s), ptr(up_s), + starts, counts, 1, N, K, ws, N // 256, (t + 15) // 16) + torch.cuda.synchronize() + return inter + + +def t3_s1_situ_parity(): + K, N = K3_K, K3_N + p1, s1 = rand_expert(K, N, seed=301) + p3, s3 = rand_expert(K, N, seed=302, half_heavy=True) + g_qw, g_s = marlinize(p1, s1, K, N) + u_qw, u_s = marlinize(p3, s3, K, N) + w1_ref = dense_expert_bf16(p1, s1) + w3_ref = dense_expert_bf16(p3, s3) + ok = True + for M in M_SWEEP: + x = torch.randn(M, K, dtype=torch.bfloat16, device=DEV) + out = _run_s1(x, g_qw, g_s, u_qw, u_s, N, K, + "grouped_marlin_gemm_m16_s1_mxfp4_situ") + # eager reference: fp32 GEMMs -> bf16 (kernel stores bf16 pass results + # in SMEM before the epilogue) -> fp32 SiTU + g_ref = (x.float() @ w1_ref.float().t()).to(torch.bfloat16) + u_ref = (x.float() @ w3_ref.float().t()).to(torch.bfloat16) + ref = situ_ref_fp32(g_ref, u_ref) + ok = gate(out, ref, f"T3 M={M}") and ok + report("T3 fused S1 SiTU parity (M sweep)", ok) + return (g_qw, g_s, u_qw, u_s, w1_ref, w3_ref) + + +def t4_grouped_zero_token(): + K, N = K3_K, K3_N + E, mtp = 32, 64 + torch.manual_seed(404) + # PRODUCTION fused layout: per-expert stacked [2, K//16, N*2] blob from + # repack_mxfp4_w13_to_marlin_gs32; up pointers derived by BYTE ARITHMETIC + # off the gate pointers, exactly as the model side must (storage + # adjacency — a wide [K, 2N] slice would carry the wrong b_gl_stride). + weights = [] # keep python refs alive for the pointer arrays + for e in range(E): + p1, s1 = rand_expert(K, N, seed=1000 + e) + p3, s3 = rand_expert(K, N, seed=2000 + e) + qw, sc = mwp.repack_mxfp4_w13_to_marlin_gs32(p1, s1, p3, s3, K, N, + emit_scale="bf16") + weights.append((qw.to(DEV), sc.to(DEV), + dense_expert_bf16(p1, s1), dense_expert_bf16(p3, s3))) + qw_branch_bytes = weights[0][0][0].numel() * weights[0][0].element_size() + s_branch_bytes = weights[0][1][0].numel() * weights[0][1].element_size() + for qw, sc, _, _ in weights: # pin the adjacency arithmetic per expert + assert qw.data_ptr() + qw_branch_bytes == qw[1].data_ptr() + assert sc.data_ptr() + s_branch_bytes == sc[1].data_ptr() + counts_host = torch.zeros(E, dtype=torch.int32) + for e in range(0, E, 3): # >half of the experts stay empty + counts_host[e] = int(torch.randint(1, mtp + 1, (1,))) + A = torch.randn(E * mtp, K, dtype=torch.bfloat16, device=DEV) + inter = torch.full((E * mtp, N), float("nan"), dtype=torch.bfloat16, device=DEV) + starts = (torch.arange(E, dtype=torch.int32) * mtp).to(DEV) + counts = counts_host.to(DEV) + gate_B = torch.tensor([w[0].data_ptr() for w in weights], dtype=torch.int64, device=DEV) + gate_S = torch.tensor([w[1].data_ptr() for w in weights], dtype=torch.int64, device=DEV) + up_B = gate_B + qw_branch_bytes # storage-adjacency derivation + up_S = gate_S + s_branch_bytes + C_ptrs = torch.tensor([inter.data_ptr() + e * mtp * N * 2 for e in range(E)], + dtype=torch.int64, device=DEV) + ws = torch.zeros(N // 256 + 17, dtype=torch.int32, device=DEV) + + mgm.marlin_grouped_stage1_fused_mxfp4_situ( + A, inter, counts, starts, gate_B, gate_S, up_B, up_S, C_ptrs, + N, K, ws, max_m_tiles=(mtp + 15) // 16, mtp=mtp, num_experts=E, + total_rows=int(counts_host.sum())) + torch.cuda.synchronize() + + ok = True + for e in range(E): + c = int(counts_host[e]) + if c == 0: + continue + x = A[e * mtp:e * mtp + c] + g_ref = (x.float() @ weights[e][2].float().t()).to(torch.bfloat16) + u_ref = (x.float() @ weights[e][3].float().t()).to(torch.bfloat16) + ref = situ_ref_fp32(g_ref, u_ref) + ok = gate(inter[e * mtp:e * mtp + c], ref, f"T4 expert {e} c={c}") and ok + + # all-empty launch must be safe AND write nothing (C-EMPTYGRP): + # NaN canary — every output byte must survive the launch untouched. + inter.fill_(float("nan")) + zero = torch.zeros(E, dtype=torch.int32, device=DEV) + mgm.marlin_grouped_stage1_fused_mxfp4_situ( + A, inter, zero, starts, gate_B, gate_S, up_B, up_S, C_ptrs, + N, K, ws, max_m_tiles=(mtp + 15) // 16, mtp=mtp, num_experts=E, + total_rows=0) + torch.cuda.synchronize() + all_nan = bool(torch.isnan(inter).all()) + print(f" T4 all-empty launch wrote nothing: {all_nan}") + report("T4 grouped fused S1 with zero-token experts (dense grid, no filtering)", + ok and all_nan) + + +def t5_mutations(t3_tensors): + g_qw, g_s, u_qw, u_s, w1_ref, w3_ref = t3_tensors + K, N = K3_K, K3_N + M = 64 + torch.manual_seed(505) + x = torch.randn(M, K, dtype=torch.bfloat16, device=DEV) + g_ref = (x.float() @ w1_ref.float().t()).to(torch.bfloat16) + u_ref = (x.float() @ w3_ref.float().t()).to(torch.bfloat16) + ref = situ_ref_fp32(g_ref, u_ref) + + caught = 0 + # m1: INT4 entry consuming E2M1 codes — finite plausible garbage, must FAIL + out = _run_s1(x, g_qw, g_s, u_qw, u_s, N, K, "grouped_marlin_gemm_m16_s1") + if not gate(out, ref, "T5.m1 int4-entry-on-e2m1"): + caught += 1 + # m2: gate/up swap — silent at kernel level, must FAIL vs the ordered ref + out = _run_s1(x, u_qw, u_s, g_qw, g_s, N, K, + "grouped_marlin_gemm_m16_s1_mxfp4_situ") + if not gate(out, ref, "T5.m2 gate/up-swap"): + caught += 1 + # m3: off-by-one scale group + g_s_mut = torch.roll(g_s, 1, dims=0).contiguous() + out = _run_s1(x, g_qw, g_s_mut, u_qw, u_s, N, K, + "grouped_marlin_gemm_m16_s1_mxfp4_situ") + if not gate(out, ref, "T5.m3 off-by-one-scale"): + caught += 1 + # m4: SiLU-instead-of-SiTU reference vs the SiTU kernel + out = _run_s1(x, g_qw, g_s, u_qw, u_s, N, K, + "grouped_marlin_gemm_m16_s1_mxfp4_situ") + silu_ref = silu_ref_fp32(g_ref, u_ref) + if not gate(out, silu_ref, "T5.m4 silu-vs-situ"): + caught += 1 + + report("T5 mutation arms caught", caught == 4, f"catch count {caught}/4") + + +def _int4_expert(K, N, seed): + """Random K2.5-style INT4 expert -> (marlin qw, marlin scales, dense ref).""" + g = torch.Generator().manual_seed(seed) + q = torch.randint(0, 16, (N, K), generator=g, dtype=torch.int32) + raw = torch.zeros(N, K // 8, dtype=torch.int32) + for i in range(8): + raw |= (q[:, i::8] & 0xF) << (i * 4) + scales = (torch.rand(N, K // 32, generator=g) * 0.02 + 0.001).to(torch.bfloat16) + qw, ms = mwp.repack_int4_to_marlin_gs32(raw, scales, K, N) + w_ref = ((q - 8).float().view(N, K // 32, 32) + * scales.float().unsqueeze(-1)).view(N, K).to(torch.bfloat16).to(DEV) + return qw.to(DEV), ms.to(DEV), w_ref + + +def t6_int4_regression(): + """K2.5 INT4 M16 through the now-templated kernel must still pass.""" + K, N = 1024, 1024 + qw, ms, w_ref = _int4_expert(K, N, seed=606) + + mod = mgm._load_module() + ok = True + for M in (1, 64, 512): + A = torch.randn(M, K, dtype=torch.bfloat16, device=DEV) + C = torch.empty(M, N, dtype=torch.bfloat16, device=DEV) + starts = torch.zeros(1, dtype=torch.int32, device=DEV) + counts = torch.tensor([M], dtype=torch.int32, device=DEV) + ws = torch.zeros(N // 256 + 17, dtype=torch.int32, device=DEV) + ptr = lambda t: torch.tensor([t.data_ptr()], dtype=torch.int64, device=DEV) + mod.grouped_marlin_gemm_m16( + A, ptr(qw), ptr(C), ptr(ms), starts, counts, + 1, N, K, ws, 1, N // 256, (M + 15) // 16) + torch.cuda.synchronize() + ref = A.float() @ w_ref.float().t() + ok = gate(C, ref, f"T6 M={M}") and ok + report("T6 INT4 M16 regression (templating did not disturb production)", ok) + + +def t6b_int4_s1_regression(): + """K2.5 fused S1 — the decode DEFAULT — must still pass its + own parity. This kernel got the heavier template surgery (epilogue + expression swapped to act_gate_mul); T5.m1 alone cannot certify it + (a must-diverge arm passes even if SiLU output is subtly wrong).""" + K, N = 1024, 1024 + g_qw, g_s, g_w = _int4_expert(K, N, seed=616) + u_qw, u_s, u_w = _int4_expert(K, N, seed=617) + ok = True + for M in (1, 64, 512): + x = torch.randn(M, K, dtype=torch.bfloat16, device=DEV) + out = _run_s1(x, g_qw, g_s, u_qw, u_s, N, K, "grouped_marlin_gemm_m16_s1") + g_ref = (x.float() @ g_w.float().t()).to(torch.bfloat16) + u_ref = (x.float() @ u_w.float().t()).to(torch.bfloat16) + ok = gate(out, silu_ref_fp32(g_ref, u_ref), f"T6b M={M}") and ok + report("T6b INT4 fused S1 regression (K2.5 decode default)", ok) + + +def t7_wrapper_hardfail_negatives(): + """Every hard-fail seam must RAISE: python wrapper checks (ValueError) and + the raw-pybind TORCH_CHECK seam in the C++ entries (RuntimeError). Uses + zero counts + real allocations so a REGRESSED (non-raising) arm degrades + to an empty launch and a clean FAIL, never a wild pointer deref.""" + K, N = K3_K, K3_N + E, mtp = 4, 16 + t = 8 + x = torch.randn(t, K, dtype=torch.bfloat16, device=DEV) + p, s = rand_expert(K, N, seed=701) + qw, ms = marlinize(p, s, K, N) + qw_e8, ms_e8 = mwp.repack_mxfp4_to_marlin_gs32(p, s, K, N) # e8m0 uint8 + qw_e8, ms_e8 = qw_e8.to(DEV), ms_e8.to(DEV) + pd, sd = rand_expert(N, K, seed=702) + dqw, dms = marlinize(pd, sd, N, K) + + A = torch.randn(E * mtp, K, dtype=torch.bfloat16, device=DEV) + inter = torch.empty(E * mtp, N, dtype=torch.bfloat16, device=DEV) + counts = torch.zeros(E, dtype=torch.int32, device=DEV) # empty: safe + starts = (torch.arange(E, dtype=torch.int32) * mtp).to(DEV) + bp = torch.tensor([qw.data_ptr()] * E, dtype=torch.int64, device=DEV) + sp = torch.tensor([ms.data_ptr()] * E, dtype=torch.int64, device=DEV) + cp = torch.tensor([inter.data_ptr() + e * mtp * N * 2 for e in range(E)], + dtype=torch.int64, device=DEV) + ws = torch.zeros(N // 256 + 17, dtype=torch.int32, device=DEV) + mod = mgm._load_module() + + def fused(**kw): + args = dict(dispatched_x_3d=A, intermediate_3d=inter, + expert_counts=counts, expert_starts=starts, + gate_B_ptrs=bp, gate_scales_ptrs=sp, + up_B_ptrs=bp, up_scales_ptrs=sp, C_ptrs=cp, + N=N, K=K, workspace=ws, max_m_tiles=1, mtp=mtp, + num_experts=E, total_rows=0) + args.update(kw) + mgm.marlin_grouped_stage1_fused_mxfp4_situ(**args) + + arms = [ + ("L2 raw-e8m0-scale-at-kernel", lambda: mgm.single_expert_marlin_mxfp4_decode( + x, qw_e8, ms_e8, qw, ms, dqw, dms, N=K3_N, K=K3_K)), + ("L2 wrong-marlin-shape", lambda: mgm.single_expert_marlin_mxfp4_decode( + x, dqw, dms, qw, ms, dqw, dms, N=K3_N, K=K3_K)), + ("L3 ptr-array-int32", lambda: fused(gate_B_ptrs=bp.to(torch.int32))), + ("L3 counts-wrong-length", lambda: fused(expert_counts=counts[:-1])), + # intermediate resized to match so the ACT check passes and L4 itself + # (prob_n % 256) is the check that fires + ("L4 N-not-256-multiple", lambda: fused( + N=N - 64, intermediate_3d=torch.empty( + E * mtp, N - 64, dtype=torch.bfloat16, device=DEV))), + ("L5 m-tile-bound-too-small", lambda: fused(mtp=64, total_rows=64, + max_m_tiles=1)), + ("ACT x-fp16", lambda: fused(dispatched_x_3d=A.half())), + ("ACT single-expert-x-fp16", lambda: mgm.single_expert_marlin_mxfp4_decode( + x.half(), qw, ms, qw, ms, dqw, dms, N=K3_N, K=K3_K)), + # raw-pybind seam: TORCH_CHECKs inside the C++ entries must hard-fail + # even when the python wrappers are bypassed (K2.5 integration pattern) + ("RAW n_tiles-mismatch", lambda: mod.grouped_marlin_gemm_m16_mxfp4( + A, bp, cp, sp, starts, counts, E, N, K, ws, E, N // 256 - 1, 1)), + ("RAW ptr-array-int32", lambda: mod.grouped_marlin_gemm_m16_s1_mxfp4_situ( + A, bp.to(torch.int32), bp, cp, sp, sp, starts, counts, + E, N, K, ws, N // 256, 1)), + ] + caught = 0 + for name, fn in arms: + try: + fn() + torch.cuda.synchronize() + print(f" T7[{name}]: NOT caught") + except (ValueError, RuntimeError) as e: + caught += 1 + print(f" T7[{name}]: caught ({str(e).splitlines()[0][:64]})") + report("T7 hard-fail negative arms", caught == len(arms), + f"catch count {caught}/{len(arms)}") + + +def do_build(): + """Rebuild the marlin TU with `-Xptxas -v` captured; print per-kernel + register counts and FAIL on an E2M1-vs-U4B8 register-tier jump (>16).""" + repo = Path(__file__).resolve().parents[2] + kdir = repo / "batchgen_kernels" + src = kdir / "src" / "moe" / "marlin_grouped_gemm.cu" + src.touch() # force recompilation of just this TU (incremental build) + env = os.environ.copy() + env["NVCC_APPEND_FLAGS"] = (env.get("NVCC_APPEND_FLAGS", "") + + " -Xptxas -v").strip() + proc = subprocess.run( + [sys.executable, "setup.py", "build_ext", "--inplace"], + cwd=kdir, env=env, capture_output=True, text=True) + log = proc.stdout + proc.stderr + log_path = kdir / "build_ptxas_marlin.log" + log_path.write_text(log) + print(f"build log: {log_path}") + if proc.returncode != 0: + print(log[-4000:]) + print("FATAL: build failed") + sys.exit(2) + + # pair each 'Compiling entry function' with the next 'Used N registers' + entries, cur = [], None + for line in log.splitlines(): + m = re.search(r"Compiling entry function '([^']+)'", line) + if m: + cur = m.group(1) + continue + m = re.search(r"Used (\d+) registers", line) + if m and cur is not None: + entries.append((cur, int(m.group(1)))) + cur = None + + def demangle(n): + try: + out = subprocess.run(["c++filt", n], capture_output=True, + text=True).stdout.strip() + return out or n + except OSError: + return n + + marlin = [(demangle(n), n, r) for n, r in entries if "Marlin" in n] + if not marlin: + print("FATAL: no Marlin ptxas lines captured — object was not " + "recompiled? Delete the build dir / .so and rerun --build.") + sys.exit(2) + regs = {} + for dem, mangled, r in marlin: + print(f" {r:4d} regs {dem}") + for fam in ("MarlinGrouped_M16_S1", "MarlinGrouped_M16"): + if fam + "I" in mangled or fam + "<" in dem: + codec = ("E2M1" if ("L6WCodec1E" in mangled or "WCodec)1" in dem) + else "U4B8" if ("L6WCodec0E" in mangled or "WCodec)0" in dem) + else "?") + regs[(fam, codec)] = r + break + ok = True + for fam in ("MarlinGrouped_M16_S1", "MarlinGrouped_M16"): + u, e = regs.get((fam, "U4B8")), regs.get((fam, "E2M1")) + if u is not None and e is not None: + tier_ok = e <= u + 16 + print(f" {fam}: U4B8={u} E2M1={e} regs " + f"({'same tier' if tier_ok else 'TIER JUMP — investigate'})") + ok = ok and tier_ok + else: + print(f" {fam}: could not classify both codecs (U4B8={u}, " + f"E2M1={e}) — inspect {log_path} manually") + ok = False + if not ok: + sys.exit(2) + print("build + register-tier check OK") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--smoke", action="store_true", help="run T0 only") + ap.add_argument("--build", action="store_true", + help="rebuild with ptxas -v and gate the register tier") + args = ap.parse_args() + + if args.build: + do_build() + return + + if not mgm.is_marlin_mxfp4_available(): + print("FATAL: marlin MXFP4 kernel entries missing — rebuild " + "batchgen_kernels (this is the L1 hard-fail).") + sys.exit(2) + + torch.cuda.init() + print(f"device: {torch.cuda.get_device_name()}") + + t0_smoke() + if not args.smoke: + t1_decode_bitexact() + t2_m16_parity() + t3_tensors = t3_s1_situ_parity() + t4_grouped_zero_token() + t5_mutations(t3_tensors) + t6_int4_regression() + t6b_int4_s1_regression() + t7_wrapper_hardfail_negatives() + + failed = [n for n, ok in _results if not ok] + print(f"\n{len(_results) - len(failed)}/{len(_results)} ladder stages passed") + if failed: + print("FAILED:", failed) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/moe/test_mxfp4_marlin_repack.py b/tests/moe/test_mxfp4_marlin_repack.py new file mode 100644 index 00000000..fdcb1974 --- /dev/null +++ b/tests/moe/test_mxfp4_marlin_repack.py @@ -0,0 +1,380 @@ +"""CPU suite for the K3 MXFP4 → Marlin repack. + +Proves, on CPU with no compiled kernels: + 1. The frozen compressed-tensors oracle vector pins the dequant convention + (and catches the known mutant conventions). + 2. repack_mxfp4_to_marlin_gs32 is PURE REARRANGEMENT: marlin → inverse is + BYTE-IDENTICAL to the source (weights and E8M0 scale bytes), at the real + K3 shapes including the w1/w3 branch shape. + 3. E8M0 → bf16 scale expansion is exact over the whole legal window and + hard-fails on the forbidden edge bytes. + 4. Every contract check (R1–R7) raises. + 5. Deliberately broken variants (wrong nibble order, transposed/wrong perm, + off-by-one scale group) are CAUGHT by these tests. + 6. The w1‖w3 commutation theorem: marlin(w1 ‖_N w3) == hcat of per-branch + marlin tensors; the fused repack is gate-first storage adjacency. + +What CPU cannot prove (GPU-staged, see gpu_parity_mxfp4_marlin.py): the +in-kernel E2M1 decode (incl. the bf16-subnormal 0.5 path), the SiTU epilogue, +and GEMM parity under the project tolerance gate. +""" + +import hashlib +import os +from pathlib import Path + +import pytest +import torch + +from tests.moe._loader import load_moe_modules + +oracle, mwp = load_moe_modules() + +# (K, N) per projection; K3 real shapes: w1/w3 branch and w2 +K3_W13 = (3584, 3072) +K3_W2 = (3072, 3584) +SMALL = (128, 256) + + +def _rand_mxfp4(K, N, seed, scale_lo=112, scale_hi=134): + g = torch.Generator().manual_seed(seed) + packed = torch.randint(0, 256, (N, K // 2), generator=g, dtype=torch.int16).to(torch.uint8) + scale = torch.randint(scale_lo, scale_hi + 1, (N, K // 32), generator=g, + dtype=torch.int16).to(torch.uint8) + return packed, scale + + +# --------------------------------------------------------------------------- +# 1. Frozen oracle vector +# --------------------------------------------------------------------------- + +def test_frozen_vector_pins_oracle_dequant(): + oracle.check_dequant_fn(oracle.mxfp4_dequantize_oracle) + + +def test_format_id_pinned(): + """The R9-R11 follow-up (ckpt_converter stamps this id; model-side L6 + matches on it) is a cross-PR contract — pin the string so a silent edit + breaks loudly instead of desynchronizing converter and loader.""" + assert oracle.MXFP4_MARLIN_FORMAT_ID == "mxfp4_marlin_gs32_v1" + + +def test_frozen_vector_pins_repack_unpack_path(): + # R8's exact self-check: the dequant built on the repack module's own + # nibble unpack must match the frozen real-checkpoint vector. + oracle.check_dequant_fn(mwp._mxfp4_dequant_via_unpack) + + +def test_frozen_vector_catches_mutant_dequants(): + lut = torch.tensor(list(oracle.MXFP4_E2M1_LUT) + + [-v for v in oracle.MXFP4_E2M1_LUT], dtype=torch.float32) + + def _dequant(packed, scales, *, swap=False, bias=oracle.MXFP4_E8M0_BIAS, + int4_style=False): + lo = (packed & 0x0F).to(torch.long) + hi = (packed >> 4).to(torch.long) + if swap: + lo, hi = hi, lo + out = torch.empty(packed.shape[0], packed.shape[1] * 2, dtype=torch.float32) + if int4_style: + out[:, 0::2] = lo.float() - 8.0 + out[:, 1::2] = hi.float() - 8.0 + else: + out[:, 0::2] = lut[lo] + out[:, 1::2] = lut[hi] + exp = (scales.to(torch.int32) - bias).repeat_interleave(32, dim=-1) + return torch.ldexp(out, exp).to(torch.bfloat16) + + mutants = { + "swapped nibbles": lambda p, s: _dequant(p, s, swap=True), + "bias off-by-one (126)": lambda p, s: _dequant(p, s, bias=126), + "INT4 (q-8)*scale instead of E2M1 LUT": lambda p, s: _dequant(p, s, int4_style=True), + } + caught = 0 + for name, fn in mutants.items(): + try: + oracle.check_dequant_fn(fn) + except AssertionError: + caught += 1 + else: + pytest.fail(f"mutant dequant NOT caught by frozen vector: {name}") + assert caught == 3 + + +# --------------------------------------------------------------------------- +# 2. Round-trip byte identity (pure-rearrangement proof) +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("K,N", [K3_W13, K3_W2, SMALL]) +@pytest.mark.parametrize("emit_scale", ["e8m0", "bf16"]) +def test_repack_roundtrip_byte_identical(K, N, emit_scale): + packed, scale = _rand_mxfp4(K, N, seed=K + N) + qw, s = mwp.repack_mxfp4_to_marlin_gs32(packed, scale, K, N, emit_scale=emit_scale) + assert qw.dtype == torch.int32 and tuple(qw.shape) == (K // 16, N * 2) + expected_s_dtype = torch.uint8 if emit_scale == "e8m0" else torch.bfloat16 + assert s.dtype == expected_s_dtype and tuple(s.shape) == (K // 32, N) + # byte count of the nibble payload is unchanged + assert qw.numel() * 4 == packed.numel() + + rp, rs = mwp.marlin_mxfp4_to_raw_cpu(qw, s, K, N) + assert torch.equal(rp, packed), "marlin round-trip weights not BYTE-IDENTICAL" + assert torch.equal(rs, scale), "marlin round-trip E8M0 scale bytes not BYTE-IDENTICAL" + + +@pytest.mark.parametrize("K,N", [SMALL, K3_W13]) +def test_dequant_of_repacked_matches_oracle(K, N): + packed, scale = _rand_mxfp4(K, N, seed=7) + qw, s = mwp.repack_mxfp4_to_marlin_gs32(packed, scale, K, N, emit_scale="e8m0") + rp, rs = mwp.marlin_mxfp4_to_raw_cpu(qw, s, K, N) + a = oracle.mxfp4_dequantize_oracle(rp, rs) + b = oracle.mxfp4_dequantize_oracle(packed, scale) + assert torch.equal(a.view(torch.int16), b.view(torch.int16)), \ + "dequant of repacked-then-unpacked diverges from oracle dequant of source" + + +# --------------------------------------------------------------------------- +# 3. Exact E8M0 -> bf16 expansion +# --------------------------------------------------------------------------- + +def test_scale_e8m0_to_bf16_exact_full_window(): + e8 = torch.arange(1, 255, dtype=torch.int16).to(torch.uint8) + got = mwp.mxfp4_scale_e8m0_to_bf16(e8).float() + ref = torch.ldexp(torch.ones(254), e8.to(torch.int32) - 127) + assert torch.equal(got.view(torch.int32), ref.view(torch.int32)), \ + "E8M0->bf16 expansion is not exact over [1, 254]" + + +@pytest.mark.parametrize("edge", [0x00, 0xFF]) +def test_scale_e8m0_edge_bytes_raise(edge): + s = torch.full((4,), edge, dtype=torch.uint8) + with pytest.raises(ValueError, match="edge byte"): + mwp.mxfp4_scale_e8m0_to_bf16(s) + + +def test_bf16_scale_roundtrip_lossless(): + _, scale = _rand_mxfp4(*SMALL, seed=11, scale_lo=1, scale_hi=254) + K, N = SMALL + packed, _ = _rand_mxfp4(K, N, seed=12) + qw, s_bf16 = mwp.repack_mxfp4_to_marlin_gs32(packed, scale, K, N, emit_scale="bf16") + _, rs = mwp.marlin_mxfp4_to_raw_cpu(qw, s_bf16, K, N) + assert torch.equal(rs, scale) + + +# --------------------------------------------------------------------------- +# 4. Contract checks R1-R7 (hard-fail negatives) +# --------------------------------------------------------------------------- + +def test_contract_checks_raise(): + K, N = SMALL + packed, scale = _rand_mxfp4(K, N, seed=3) + + # R1: wrong packed dtype + with pytest.raises(ValueError, match="must be uint8"): + mwp.repack_mxfp4_to_marlin_gs32(packed.to(torch.int32), scale, K, N) + # R2: bf16 scales = INT4 checkpoint + with pytest.raises(ValueError, match="INT4 checkpoint"): + mwp.repack_mxfp4_to_marlin_gs32(packed, scale.to(torch.bfloat16), K, N) + # R3: packed dim != K//2 + with pytest.raises(ValueError, match="K//2"): + mwp.repack_mxfp4_to_marlin_gs32(packed[:, :-1], scale, K, N) + # R4: scale groups != K/32 + with pytest.raises(ValueError, match="K/32"): + mwp.repack_mxfp4_to_marlin_gs32(packed, scale[:, :-1], K, N) + # R5: N mismatch packed vs scale + with pytest.raises(ValueError, match="N mismatch"): + mwp.repack_mxfp4_to_marlin_gs32(packed, scale[:-1], K, N) + # R6: N % 64 != 0 + p32, s32 = _rand_mxfp4(K, 32, seed=4) + with pytest.raises(ValueError, match="N%64"): + mwp.repack_mxfp4_to_marlin_gs32(p32, s32, K, 32) + # K-divisibility leg: K % 32 != 0 is caught by R4 (which subsumes R6's + # K%16 clause — that clause is defense in depth and unreachable while R4 + # runs first). K=48: packed [N, 24] passes R3, then R4 raises. + p48 = torch.randint(0, 256, (N, 24), dtype=torch.int16).to(torch.uint8) + s48 = torch.full((N, 1), 120, dtype=torch.uint8) + with pytest.raises(ValueError, match="K/32"): + mwp.repack_mxfp4_to_marlin_gs32(p48, s48, 48, N) + # R7: forbidden E8M0 edge bytes + for edge in (0x00, 0xFF): + bad = scale.clone() + bad[0, 0] = edge + with pytest.raises(ValueError, match="edge byte"): + mwp.repack_mxfp4_to_marlin_gs32(packed, bad, K, N) + # emit_scale validation + with pytest.raises(ValueError, match="emit_scale"): + mwp.repack_mxfp4_to_marlin_gs32(packed, scale, K, N, emit_scale="fp16") + + +# --------------------------------------------------------------------------- +# 5. Mutation arms — each deliberately broken variant must be CAUGHT +# --------------------------------------------------------------------------- + +def test_mutation_swapped_nibble_order_caught(): + """A hi-nibble-first unpack must be caught (a) by the frozen vector and + (b) by byte-identity against the source.""" + K, N = SMALL + packed, scale = _rand_mxfp4(K, N, seed=21) + + swapped = ((packed & 0x0F) << 4) | (packed >> 4) # mutant source unpack order + + # (a) dequant-level: swapping changes values on real-vector data + def mutant_dequant(p, s): + return oracle.mxfp4_dequantize_oracle(((p & 0x0F) << 4) | (p >> 4), s) + with pytest.raises(AssertionError): + oracle.check_dequant_fn(mutant_dequant) + + # (b) byte-level: repacking the swapped bytes cannot round-trip to the + # original source + qw, s = mwp.repack_mxfp4_to_marlin_gs32(swapped, scale, K, N) + rp, _ = mwp.marlin_mxfp4_to_raw_cpu(qw, s, K, N) + assert not torch.equal(rp, packed), \ + "swapped-nibble mutation NOT caught by byte-identity" + + +def test_mutation_wrong_perm_caught(): + """Packing with a wrong (inverse-instead-of-forward) tile permutation must + break the byte-identity round trip.""" + K, N = SMALL + packed, scale = _rand_mxfp4(K, N, seed=22) + q_w_nk = mwp._unpack_mxfp4_nibbles(packed, K, N) + q_w = q_w_nk.t().contiguous() + + wrong_perm = mwp._inverse_weight_perm(4) # mutant: inverse used as forward + qw_mut = mwp._marlin_pack_weights(q_w, K, N, wrong_perm) + + rp, _ = mwp.marlin_mxfp4_to_raw_cpu( + qw_mut, mwp._marlin_permute_scales(scale.t().contiguous(), K, N, 32), K, N) + assert not torch.equal(rp, packed), "wrong-perm mutation NOT caught" + + +def test_mutation_transposed_tiles_caught(): + """Skipping the [K,N] transpose (packing N-major) must break round-trip.""" + K, N = (256, 256) # square so shapes still line up — hardest case + packed, scale = _rand_mxfp4(K, N, seed=23) + q_w_nk = mwp._unpack_mxfp4_nibbles(packed, K, N) + + qw_mut = mwp._marlin_pack_weights(q_w_nk.contiguous(), K, N, + mwp.get_weight_perm(4)) # mutant: no .t() + rp, _ = mwp.marlin_mxfp4_to_raw_cpu( + qw_mut, mwp._marlin_permute_scales(scale.t().contiguous(), K, N, 32), K, N) + assert not torch.equal(rp, packed), "transposed-tiles mutation NOT caught" + + +def test_mutation_off_by_one_scale_group_caught(): + K, N = SMALL + packed, scale = _rand_mxfp4(K, N, seed=24, scale_lo=100, scale_hi=140) + qw, s = mwp.repack_mxfp4_to_marlin_gs32(packed, scale, K, N) + s_mut = torch.roll(s, shifts=1, dims=0) # off-by-one K-group + rp, rs = mwp.marlin_mxfp4_to_raw_cpu(qw, s_mut, K, N) + a = oracle.mxfp4_dequantize_oracle(rp, rs) + b = oracle.mxfp4_dequantize_oracle(packed, scale) + assert not torch.equal(a.view(torch.int16), b.view(torch.int16)), \ + "off-by-one scale-group mutation NOT caught" + + +# --------------------------------------------------------------------------- +# 6. w1 ‖ w3 fusion +# --------------------------------------------------------------------------- + +def test_w13_commutation_theorem(): + """marlin(w1 ‖_N w3) == hcat(marlin(w1), marlin(w3)) — the permutation + acts within 64-N-column blocks, and 3072 % 64 == 0, so no permuted unit + crosses the branch boundary. Verified here at a small N multiple of 64.""" + K, N = (128, 128) + w1_p, w1_s = _rand_mxfp4(K, N, seed=31) + w3_p, w3_s = _rand_mxfp4(K, N, seed=32) + + cat_p = torch.cat([w1_p, w3_p], dim=0) # [2N, K//2] source-coordinate concat + cat_s = torch.cat([w1_s, w3_s], dim=0) + qw_cat, s_cat = mwp.repack_mxfp4_to_marlin_gs32(cat_p, cat_s, K, 2 * N) + + qw1, s1 = mwp.repack_mxfp4_to_marlin_gs32(w1_p, w1_s, K, N) + qw3, s3 = mwp.repack_mxfp4_to_marlin_gs32(w3_p, w3_s, K, N) + + assert torch.equal(qw_cat, torch.cat([qw1, qw3], dim=1)), \ + "w13 commutation theorem violated (weights)" + assert torch.equal(s_cat, torch.cat([s1, s3], dim=1)), \ + "w13 commutation theorem violated (scales)" + + +def test_w13_fused_repack_gate_first_adjacency(): + K, N = SMALL + w1_p, w1_s = _rand_mxfp4(K, N, seed=41) + w3_p, w3_s = _rand_mxfp4(K, N, seed=42) + + qw, s = mwp.repack_mxfp4_w13_to_marlin_gs32(w1_p, w1_s, w3_p, w3_s, K, N) + qw1, s1 = mwp.repack_mxfp4_to_marlin_gs32(w1_p, w1_s, K, N) + qw3, s3 = mwp.repack_mxfp4_to_marlin_gs32(w3_p, w3_s, K, N) + + assert tuple(qw.shape) == (2, K // 16, N * 2) and qw.is_contiguous() + assert torch.equal(qw[0], qw1) and torch.equal(qw[1], qw3), \ + "fused repack is not gate(w1)-first / up(w3)-second" + assert torch.equal(s[0], s1) and torch.equal(s[1], s3) + + # storage adjacency: up tensor starts exactly one branch after gate + assert qw[1].data_ptr() == qw.data_ptr() + qw[0].numel() * qw.element_size() + + # gate/up swap is a DIFFERENT blob (silent at kernel level — pinned here + # and by the GPU SiTU mutation test) + qw_sw, _ = mwp.repack_mxfp4_w13_to_marlin_gs32(w3_p, w3_s, w1_p, w1_s, K, N) + assert not torch.equal(qw_sw, qw) + + # branch shape mismatch raises + with pytest.raises(ValueError, match="mismatch"): + mwp.repack_mxfp4_w13_to_marlin_gs32(w1_p, w1_s, w3_p[:-1], w3_s[:-1], K, N) + + +# --------------------------------------------------------------------------- +# 7. Optional heavier oracles +# --------------------------------------------------------------------------- + +def test_oracle_matches_compressed_tensors_package(): + """Independent cross-check against the installed compressed-tensors + package (the checkpoint's own format implementation). Skipped when the + package is not installed; the frozen vector pins the convention anyway.""" + pytest.importorskip("compressed_tensors") + try: + from compressed_tensors.compressors.nvfp4.helpers import unpack_fp4_from_uint8 + from compressed_tensors.compressors.mx_utils import decompress_mx_scale + except ImportError: + pytest.skip("compressed-tensors version lacks nvfp4/mx helpers") + + K, N = (256, 128) + packed, scale = _rand_mxfp4(K, N, seed=51, scale_lo=100, scale_hi=140) + vals = unpack_fp4_from_uint8(packed, N, K, dtype=torch.float32) + ref = (vals.unflatten(-1, (K // 32, 32)) + * decompress_mx_scale(scale).to(torch.float32).unsqueeze(-1) + ).flatten(-2) + ours = oracle.mxfp4_dequantize_oracle(packed, scale, dtype=torch.float32) + assert torch.equal(ours.view(torch.int32), ref.view(torch.int32)), \ + "mxfp4_dequantize_oracle diverges from compressed-tensors" + + +REAL_DIR = os.environ.get("MXFP4_W1_REAL_DIR", "") + + +@pytest.mark.skipif( + not (REAL_DIR and Path(REAL_DIR, "w1_packed.bin").exists()), + reason="set MXFP4_W1_REAL_DIR to a dir with w1_packed.bin/w1_scale.bin " + "(real K3 layer-4 expert-0 w1 bytes) to run the real-tensor pin") +def test_real_w1_tensor_pins_and_roundtrip(): + N, K = 3072, 3584 + packed = torch.frombuffer( + bytearray(Path(REAL_DIR, "w1_packed.bin").read_bytes()), + dtype=torch.uint8).view(N, K // 2) + scale = torch.frombuffer( + bytearray(Path(REAL_DIR, "w1_scale.bin").read_bytes()), + dtype=torch.uint8).view(N, K // 32) + + assert hashlib.sha256(packed.numpy().tobytes()).hexdigest() == \ + "cf822517403f5ccb418150e10b303568f617f3099bea6dcc9af3a7b6a48e3501" + assert hashlib.sha256(scale.numpy().tobytes()).hexdigest() == \ + "b1cd3499f23097edbbc3f4dc3304c573dbdc49602404b90bdfb1b3dbb0b4ea92" + + deq = oracle.mxfp4_dequantize_oracle(packed, scale) + assert hashlib.sha256(deq.view(torch.int16).numpy().tobytes()).hexdigest() == \ + oracle.VEC_FULL_W1_DEQUANT_BF16_SHA256 + + qw, s = mwp.repack_mxfp4_to_marlin_gs32(packed, scale, K, N) + rp, rs = mwp.marlin_mxfp4_to_raw_cpu(qw, s, K, N) + assert torch.equal(rp, packed) and torch.equal(rs, scale) From 765971a2d89396b83de4f228b6592741fb21285a Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 22:52:19 +0800 Subject: [PATCH 3/3] fix(moe): name the loaded .so in the marlin MXFP4 L1 hard-fail MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A stale installed copy in site-packages shadows the repo's in-tree build whenever the repo root is not on sys.path, and the old message said "rebuild" — the wrong fix for that failure. The L1 error now names the extension file actually loaded so an import-path problem is diagnosed as one. Split out of the Kimi-K3 M2 model commit (kernel-scope file). --- batchgen/moe/marlin_grouped_moe.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/batchgen/moe/marlin_grouped_moe.py b/batchgen/moe/marlin_grouped_moe.py index ecd217a3..5f22d595 100644 --- a/batchgen/moe/marlin_grouped_moe.py +++ b/batchgen/moe/marlin_grouped_moe.py @@ -277,10 +277,19 @@ def _require_mxfp4_kernels(): """L1: the marlin MXFP4 entries must exist — K3 refuses to run otherwise.""" missing = [k for k in _MXFP4_KERNEL_ENTRIES if not hasattr(_module, k)] if missing: + # Name the .so actually loaded. A stale copy installed in + # site-packages shadows the repo's in-tree build whenever the repo + # root is not on sys.path (e.g. `python path/to/script.py`, whose + # sys.path[0] is the SCRIPT's directory) — in which case the fix is + # the import path, not a rebuild. + loaded = getattr(_module, "__file__", "") raise RuntimeError( f"Marlin MXFP4 kernel entries missing from " f"batchgen_kernels.moe._C_marlin_grouped_gemm: {missing}. " - f"K3 refuses to run (stale batchgen_kernels build — rebuild). " + f"The extension actually loaded is {loaded} — if that path is " + f"not inside this repo, a stale installed copy is shadowing the " + f"in-tree build; fix sys.path/PYTHONPATH rather than rebuilding. " + f"K3 refuses to run. " f"The designated parity-debug opt-in is batchgen_debug." f"k3_moe_reference; its model-side wiring is a named follow-up " f"of the K3 MXFP4 work — if it is not wired yet there is NO alternative "