# Copyright 2026 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import functools from contextlib import contextmanager import torch import triton import triton.language as tl from torch.library import triton_op, wrap_triton from ._ops import add_op_namespace_prefix, ops # ── Format constants ────────────────────────────────────────────────────────── # FP8 (E4M3) is the main format for weights and activations; FP8_DTYPE = torch.float8_e4m3fn # FP4 (E2M1) packs two 4-bit nibbles per byte. MX formats (MXFP4 weights, MXFP8 # E4M3 weights/activations) share one UE8M0 scale per 32-element K-group — the OCP # MX block size, consumed by ``tl.dot_scaled``. Format constants, not tunables. NIBBLES_PER_BYTE = 2 MX_SCALE_GROUP_K = 32 # ── Host-side helpers ───────────────────────────────────────────────────────── @contextmanager def device_context(device: torch.device): """Context manager that sets the active device for any backend (cuda, xpu, etc.).""" backend = getattr(torch, device.type, None) if backend is not None and hasattr(backend, "device"): with backend.device(device): yield else: yield @functools.lru_cache(maxsize=8) def sm_count(device_index: int) -> int: """Processor (CUDA SM / XPU Xe-core) count for a device (a constant) — cached so callers don't re-query it each launch. Read from Triton's active driver so it works on any backend, with a CUDA fallback.""" try: return triton.runtime.driver.active.utils.get_device_properties(device_index)[ "multiprocessor_count" ] except Exception: active_device = get_active_device_type() if active_device == "cuda": return torch.cuda.get_device_properties(device_index).multi_processor_count elif active_device == "xpu": return torch.xpu.get_device_properties(device_index).multi_processor_count else: raise RuntimeError( f"Unsupported device type {active_device} for sm_count; only cuda/xpu are supported." ) # H-tile each topk_reduce program reduces (one (token, tile) program per BLOCK_H span). The # reduce is bandwidth-bound and saturates at BLOCK_H>=512: a sweep over {128..2048} x prefill # shapes put 512/1024/2048 within ~7% (mostly noise) while 128/256 lagged — not worth a tuning # axis for this small slice of runtime. 512 = smallest that saturates (best at low token counts, # where bigger tiles waste masked lanes). Power of 2 (tl.arange). TOPK_REDUCE_BLOCK_H = 512 @triton.jit def topk_reduce_kernel( ProjOut, Out, ExpertIds, H, stride_pm, stride_ph, stride_om, stride_oh, stride_ei_m, stride_ei_k, NUM_TOP_K: tl.constexpr, NUM_EXPERTS: tl.constexpr, BLOCK_H: tl.constexpr, ): """Sum the ``NUM_TOP_K`` flat ``(token, slot)`` rows of ``ProjOut`` into ``Out[token]`` — a tight replacement for ``proj_out.view(T, K, H).sum(1)`` (~2.8x faster than torch's generic reduce; fp32 accumulate, so bit-identical). Slots whose expert is non-local (EP sentinel id ``>= NUM_EXPERTS``) are skipped — the experts kernels never write those rows, so they must contribute 0. Caller allocs Out and launches it.""" t = tl.program_id(0) offs_h = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) mask = offs_h < H acc = tl.zeros((BLOCK_H,), tl.float32) for k in tl.static_range(NUM_TOP_K): local = tl.load(ExpertIds + t * stride_ei_m + k * stride_ei_k) < NUM_EXPERTS acc += tl.load( ProjOut + (t * NUM_TOP_K + k) * stride_pm + offs_h * stride_ph, mask=mask & local, other=0.0, ).to(tl.float32) tl.store( Out + t * stride_om + offs_h * stride_oh, acc.to(Out.dtype.element_ty), mask=mask, ) def ue8m0_as_uint8(scale: torch.Tensor) -> torch.Tensor: """View UE8M0 (``float8_e8m0fnu``) weight scales as ``uint8`` for the Triton binder, which doesn't recognize the dtype; kernels decode ``2^(exp-127)`` inline. fp32 (non-UE8M0) scales pass through unchanged.""" return scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale def e2m1_as_uint8(weight: torch.Tensor) -> torch.Tensor: """View an ``int8``-stored MXFP4 (packed E2M1) weight as ``uint8`` — a zero-cost reinterpret. ``tl.dot_scaled`` requires the packed rhs as ``uint8``, so do it once here instead of casting in-kernel at every load. E4M3 (MXFP8) weights pass through unchanged.""" return weight.view(torch.uint8) if weight.dtype == torch.int8 else weight # UE8M0 group-32 scales arrive either as ``float8_e8m0fnu`` or as raw ``uint8`` — the same 8 # exponent bits, and a common on-disk encoding (e.g. group-32 "mxfp8" checkpoints store the # scale tensor as uint8). Both are valid MX scales: ``ue8m0_as_uint8`` reinterprets to uint8 # and the kernels decode ``2^(exp-127)`` inline, so the detectors accept either dtype. UE8M0_SCALE_DTYPES = (torch.float8_e8m0fnu, torch.uint8) def is_mxfp8(weight: torch.Tensor, scale: torch.Tensor) -> bool: """MXFP8 weight/scale pair: E4M3 weights with UE8M0 group-32 scales — last dim ``scale.shape[-1] == weight.shape[-1] // MX_SCALE_GROUP_K``, matching leading dims. Works for 2D ``(N, K)`` and 3D ``(E, N, K)`` weights. The group-32 layout is what separates MXFP8 from 128-block FP8 (which may also carry UE8M0 scales).""" return ( weight.dtype == torch.float8_e4m3fn and scale.dtype in UE8M0_SCALE_DTYPES and scale.ndim == weight.ndim and scale.shape[:-1] == weight.shape[:-1] and scale.shape[-1] == weight.shape[-1] // MX_SCALE_GROUP_K ) def is_mxfp4(weight: torch.Tensor, scale: torch.Tensor) -> bool: """MXFP4 weight/scale pair: packed E2M1 weights (``int8``, two codes/byte) with UE8M0 group-32 scales — ``scale.shape[-1] == weight.shape[-1] * NIBBLES_PER_BYTE // MX_SCALE_GROUP_K`` (unpacked K = ``2 * K_half``), matching leading dims. 2D or 3D.""" return ( weight.dtype == torch.int8 and scale.dtype in UE8M0_SCALE_DTYPES and scale.ndim == weight.ndim and scale.shape[:-1] == weight.shape[:-1] and scale.shape[-1] == (weight.shape[-1] * NIBBLES_PER_BYTE) // MX_SCALE_GROUP_K ) def is_mxfp(weight: torch.Tensor, scale: torch.Tensor) -> bool: """Any MX weight/scale pair — MXFP8 (``float8_e4m3fn``, one value/byte) or MXFP4 (``int8``, two E2M1 codes/byte), both with UE8M0 group-32 scales. ``values_per_byte`` folds the two cases: the scale's last dim covers the unpacked K (``weight.shape[-1] * values_per_byte``) in groups of ``MX_SCALE_GROUP_K``. The dispatchers route on this; the op picks the format from ``weight.dtype``.""" if weight.dtype == torch.float8_e4m3fn: values_per_byte = 1 elif weight.dtype == torch.int8: values_per_byte = NIBBLES_PER_BYTE else: return False return ( scale.dtype in UE8M0_SCALE_DTYPES and scale.ndim == weight.ndim and scale.shape[:-1] == weight.shape[:-1] and scale.shape[-1] == (weight.shape[-1] * values_per_byte) // MX_SCALE_GROUP_K ) def is_tensor_wide(block_size, weight: torch.Tensor) -> bool: """True when ``block_size`` selects per-tensor (tensor-dynamic) scaling: ``None`` or equal to the weight's full ``(N, K)`` — one scale block spanning the whole matrix. Handles 2D ``(N, K)`` and 3D ``(E, N, K)`` weights via the last two dims.""" return block_size is None or ( block_size[0] == weight.shape[-2] and block_size[1] == weight.shape[-1] ) # The per-token batched/fused kernels are decode-shaped — one routed row per program, so # BLOCK_SIZE_M is 1 (a larger tile would just recompute the same row). The 2D matmul sizes # its M tile to the workload instead, via ``adaptive_block_size_m``. DECODE_BLOCK_SIZE_M = 1 def adaptive_block_size_m(target_m: int) -> int: """Smallest power-of-2 >= ``target_m``, floored at 16 and capped at 128. Used by all matmul wrappers (single / batched / grouped) to size the M tile to the workload — small per-expert M wants smaller tiles, large M caps out at 128 to keep register pressure bounded. Pass ``M`` for single matmul, or ``(S + E - 1) // E`` (avg tokens per expert) for batched/grouped. """ return min(max(triton.next_power_of_2(target_m), 16), 128) @functools.cache def get_active_device_type() -> str: """Active torch device type for the current Triton backend (``"cuda"``, ``"xpu"``, ...). Falls back to ``"cuda"`` when no driver is loaded — Triton raises ``RuntimeError: 0 active drivers ([])`` on driverless build boxes, and the autotune-config builder is evaluated at module-import time under the ``@triton.autotune`` decorator (no kernel launches there, so the default is only used to shape the config list). """ try: return triton.runtime.driver.active.get_active_torch_device().type except RuntimeError: return "cuda" def get_accelerator_autotuning_configs( *, tune_block_m: bool = False, tune_block_nk: bool = False, ): """Autotune search grid for the current accelerator. ``num_warps``, ``num_stages`` and ``blocks`` (the ``(BLOCK_SIZE_N, BLOCK_SIZE_K)`` tile shapes) are fixed up front from ``(is_xpu, tune_block_nk)``, then crossed into the config list. ``tune_block_nk=True`` sweeps the tile: used by kernels that have no caller ``block_size`` to fix it — the MX ``tl.dot_scaled`` paths AND the tensor-dynamic FP8 paths. ``tune_block_nk=False`` emits a single empty meta-dict (block-scaled kernels take the tile from the caller's ``block_size``). ``tune_block_m=True`` also crosses in BLOCK_SIZE_M ∈ ``(16, 32, 64, 128)`` (fused grouped, which lets the tuner pick M instead of computing it via ``adaptive_block_size_m``). The CUDA tile set is a data-driven prune of a B200 sweep across single (BM=128), grouped MoE (BM=16/64) and decode (BM=1): winners only ever used these 4 tiles, num_warps in {4,8,16}, num_stages in {2,3} (warps=2, stages=4 and 256x256 never won) — 108 → 24. (Tuned on dot_scaled; tensor-dynamic tl.dot reuses it.) """ is_xpu = get_active_device_type() == "xpu" num_warps = [8, 16] if is_xpu else [2, 4, 8, 16] num_stages = [2, 3, 4] if tune_block_nk: tiles = ( [(128, 128)] if is_xpu else [(128, 128), (256, 128), (128, 64), (64, 256)] ) blocks = [{"BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk} for bn, bk in tiles] else: blocks = [{}] if tune_block_m: blocks = [{**b, "BLOCK_SIZE_M": bm} for b in blocks for bm in (16, 32, 64, 128)] return [ triton.Config(b, num_warps=w, num_stages=s) for b in blocks for w in num_warps for s in num_stages ] def get_mxfp_autotuning_configs( pre_hook=None, compute_modes=("dot_scaled", "dot", "scalar"), memory_modes=(None,), tune_block_m=False, ): """Autotune grid for the MXFP8 MoE/matmul kernels. ``COMPUTE_MODE`` (a constexpr string) picks the multiply-accumulate, and the Bayesian tuner selects it per workload (token count is in the key): - ``"dot_scaled"``: native group-32 scaled MMA (wide K, ``BLOCK_SIZE_K`` ∈ {128, 256}) — wins once the grid saturates (prefill, ~S≥32). - ``"dot"``: fp8 ``tl.dot`` + per-group software rescale, one scale group per K-step (``BLOCK_SIZE_K == 32``). - ``"scalar"``: scalar CUDA-core FMA reduction, no tensor core — same per-group structure as ``"dot"`` (hence bit-exact with it) but avoids the scaled-MMA's M→16 pad, so it wins for memory-bound **decode** (M=1; ≈2× over ``dot_scaled`` at S≥32 in a sweep). Only emit for kernels that implement the branch. ``num_warps=16`` and ``BLOCK_SIZE_K=64`` are omitted (dead at M=1 per a sweep). ``memory_modes`` (fused_grouped only) is the ``MEMORY_MODE`` axis for the weight load: pass ``("descriptor", "pointer")`` — ``"descriptor"`` resolves here to the device's tensor-descriptor flavor (host-built / NVIDIA-TMA on CUDA, device-built in-kernel tensormap on XPU), ``"pointer"`` is explicit pointers. ``(None,)`` (default) emits no axis. A given pre_hook is attached to every config and self-guards on ``MEMORY_MODE``.""" # "descriptor" -> the device's descriptor flavor: host-built (NVIDIA TMA) on CUDA, device-built # (in-kernel tensormap) on XPU. "pointer" / "None" pass through. descriptor = ( "device_descriptor" if get_active_device_type() == "xpu" else "host_descriptor" ) memory_modes = tuple(descriptor if m == "descriptor" else m for m in memory_modes) if "device_descriptor" in memory_modes: # device-side descriptors build in-kernel tensormaps that need a scratch allocator (the # XPU path; CUDA keeps host-built descriptors and never reaches this). triton.set_allocator( lambda size, alignment, stream: torch.empty( size, dtype=torch.int8, device="xpu" ) ) bms = (16, 32, 64, 128) if tune_block_m else (None,) return [ triton.Config( { "COMPUTE_MODE": mode, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, **({"MEMORY_MODE": mm} if mm is not None else {}), **({"BLOCK_SIZE_M": bm} if bm is not None else {}), }, num_warps=w, num_stages=s, pre_hook=pre_hook, ) for mm in memory_modes for mode in compute_modes for bn in [32, 64, 128, 256] # only "dot" pins BK to the group (one scale group per fp8 tl.dot step); "dot_scaled" # and "scalar" handle multiple groups per K-chunk, so they use a wide BK. for bk in ([32] if mode == "dot" else [128, 256]) for s in [2, 3, 4, 5, 6] for w in [2, 4, 8] for bm in bms ] @functools.lru_cache(maxsize=None) def sm_shared_memory_limit(device_index: int) -> int: """Max dynamic shared memory per block (bytes) for a CUDA device — the cap Triton reports as the 'Hardware limit' on an ``out of resource: shared memory`` failure (~232 KB on B200, ~227 KB on H100, much less on older/consumer parts). Queried from the driver so the prune adapts to the hardware instead of hardcoding one GPU.""" try: return triton.runtime.driver.active.utils.get_device_properties(device_index)[ "max_shared_mem" ] except Exception: if get_active_device_type() == "cuda": return torch.cuda.get_device_properties( device_index ).shared_memory_per_block_optin elif get_active_device_type() == "xpu": return torch.xpu.get_device_properties( device_index ).shared_memory_per_block_optin else: raise RuntimeError( f"Unsupported device type {get_active_device_type()} for sm_shared_memory_limit; only cuda/xpu are supported." ) def smem_config_pruner( act_bytes: int, n_weight_tiles: int, weight_bytes: int = 1, reduction_dim: str | None = None, ): """Build an ``early_config_prune`` that drops configs whose pipelined operand shared memory would overflow the SM — the source of ``out of resource: shared memory`` autotune failures (and the wasted compiles they cause). Per-stage estimate (bytes) = ``BK · (act_bytes·BM + weight_bytes·n_weight_tiles·BN)``: one activation tile ``[BM, BK]`` plus ``n_weight_tiles`` weight tiles ``[BN, BK]`` (gate_up fuses 2; down has 1), times ``num_stages``. ``BM`` is the routing-derived tile, read from the launch args. MXFP4 packs 2 weights/byte, so ``weight_bytes=1`` (MXFP8) is a safe upper bound. The limit is read from the active device. Never returns empty — keeps the smallest-footprint config as a fallback. ``reduction_dim`` (the K-loop bound arg name, e.g. ``"INTERMEDIATE_DIM"``) additionally drops ``dot_scaled`` configs with ``BLOCK_SIZE_K >= reduction_dim`` — a single-trip K-loop, which on sm_100 trips a Triton ``optimize-accumulator-init`` bug (uninitialized TMEM alloc must be mutable) and never compiles. The autotuner discards them anyway; pruning up front just skips the doomed compile and its MLIR error spam. ``None`` (default) disables the check.""" def prune(configs, named_args, **kwargs): # The estimate needs all three tile dims. Each is either an autotuned config meta # (MX) or a launch arg (block-dynamic), so look in both; raise clearly if missing. all_args = {**named_args, **kwargs} dev = triton.runtime.driver.active.get_active_torch_device() limit = sm_shared_memory_limit(dev.index if dev.index is not None else 0) def dim(c, name): v = c.kwargs.get(name, all_args.get(name)) if v is None: raise ValueError( f"smem_config_pruner needs {name} (autotune config meta or launch arg) " "to estimate shared memory; none found." ) return v def smem(c): BM, BN, BK = ( dim(c, n) for n in ("BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K") ) return ( c.num_stages * BK * (act_bytes * BM + weight_bytes * n_weight_tiles * BN) ) def single_trip_dot_scaled(c): return ( reduction_dim is not None and c.kwargs.get("COMPUTE_MODE") == "dot_scaled" and c.kwargs.get("BLOCK_SIZE_K", 0) >= all_args[reduction_dim] ) kept = [ c for c in configs if smem(c) <= limit and not single_trip_dot_scaled(c) ] return kept or [min(configs, key=smem)] return prune # ── Triton-side helpers (inlined by ``@triton.jit`` callers) ────────────────── @triton.jit def fp8_act_quant_inline(a_raw): """Inline FP8 (E4M3) activation quant for the W8A8 block-scale path. Per-row amax → fp32 scale ``amax/448`` (floored at 1e-12 against zero rows) → cast values to FP8. Returns ``(a_fp8, a_s)`` with shapes ``(M, K)`` and ``(M,)``. """ a_s = tl.max(tl.abs(a_raw), axis=1) / 448.0 a_fp8 = (a_raw / tl.maximum(a_s[:, None], 1e-12)).to(tl.float8e4nv) return a_fp8, a_s @triton.jit def mxfp_act_quant_inline( a_raw, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, SCALE_GROUP_K: tl.constexpr, ): """Inline E4M3 activation quant for the MX paths (W4A8 MXFP4 / W8A8 MXFP8). Per-row, per-K-group amax → UE8M0 scale (ceil to next power-of-2 via the mantissa-nonzero bump trick) → cast values to FP8. Returns ``(a_fp8, a_scale_u8)`` with shapes ``(M, K)`` and ``(M, K // SCALE_GROUP_K)``. """ a_groups = tl.reshape( a_raw, (BLOCK_SIZE_M, BLOCK_SIZE_K // SCALE_GROUP_K, SCALE_GROUP_K) ) a_s_fp32 = tl.max(tl.abs(a_groups), axis=2) / 448.0 bits = a_s_fp32.to(tl.int32, bitcast=True) # ceil_to_ue8m0: bump exponent by 1 when mantissa is non-zero. exp_ceil = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) exp_ceil = tl.minimum(tl.maximum(exp_ceil, 1), 254) a_scale_u8 = exp_ceil.to(tl.uint8) a_s_pow2 = (exp_ceil << 23).to(tl.float32, bitcast=True) a_fp8 = tl.reshape( a_groups / tl.maximum(a_s_pow2[:, :, None], 1e-12), (BLOCK_SIZE_M, BLOCK_SIZE_K), ).to(tl.float8e4nv) return a_fp8, a_scale_u8 @triton.jit def decode_ue8m0_scale(scale): """Decode a UE8M0 weight scale to fp32: when it was loaded as ``uint8`` exponent bits, ``value = 2^(exp - 127)``, built directly as the fp32 bit pattern. fp32 scales (block-dynamic with float scales) pass through. The dtype branch is a compile-time constant, so only the taken path is emitted (single return — Triton requires all ``return`` statements to share a type).""" if scale.dtype == tl.uint8: scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True) return scale @triton.jit def expand_ue8m0_scale( scale_u8, ROWS: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, SCALE_GROUP_K: tl.constexpr, ): """Decode a UE8M0 group-scale tile ``[ROWS, BLOCK_SIZE_K // SCALE_GROUP_K]`` and broadcast each group scale across its ``SCALE_GROUP_K`` elements → ``[ROWS, BLOCK_SIZE_K]`` fp32 (one scale per K element). Used by the scalar (CUDA-core FMA) MX paths to dequant the whole K chunk.""" s = decode_ue8m0_scale(scale_u8) return tl.reshape( tl.broadcast_to( s[:, :, None], (ROWS, BLOCK_SIZE_K // SCALE_GROUP_K, SCALE_GROUP_K) ), (ROWS, BLOCK_SIZE_K), ) @triton.jit def mx_dot_scaled(acc, a, a_scale, w, w_scale, VALUES_PER_BYTE: tl.constexpr): """MX 'dot_scaled' path: scaled MMA folding the UE8M0 group scales into the tensor core — ``a`` (E4M3) @ ``w`` (packed E2M1 if MXFP4, else E4M3). The rhs format is picked from ``VALUES_PER_BYTE``. Caller pre-shapes ``w``/``w_scale`` (e.g. ``tl.trans(gu)``).""" rhs_format: tl.constexpr = "e2m1" if VALUES_PER_BYTE == 2 else "e4m3" return tl.dot_scaled(a, a_scale, "e4m3", w, w_scale, rhs_format, acc) @triton.jit def mx_dot_rescale(acc, a, w, a_scale, w_scale, VALUES_PER_BYTE: tl.constexpr): """MX 'dot' path (BK == group): unpack MXFP4 weights to E4M3, fp8 ``tl.dot`` + per-group software rescale (decoding both UE8M0 scales internally), accumulating into ``acc`` (returned updated). Single weight — for the gate∪up pair (shared activation scale) use ``mx_dot_rescale_gate_up``.""" wq = mxfp4_e2m1_to_e4m3(w) if VALUES_PER_BYTE == 2 else w return acc + tl.dot(a, wq) * decode_ue8m0_scale(a_scale) * tl.trans( decode_ue8m0_scale(w_scale) ) @triton.jit def mx_dot_rescale_gate_up( acc_gate, acc_up, a, w_gate, w_up, a_scale, gate_scale, up_scale ): """Gate∪up 'dot' path: decode the SHARED activation scale once, rescale both projections and accumulate into ``acc_gate``/``acc_up`` (returned updated).""" a_s = decode_ue8m0_scale(a_scale) acc_gate += tl.dot(a, w_gate) * a_s * tl.trans(decode_ue8m0_scale(gate_scale)) acc_up += tl.dot(a, w_up) * a_s * tl.trans(decode_ue8m0_scale(up_scale)) return acc_gate, acc_up @triton.jit def mx_scalar_reduce( acc, a, a_scale, w, w_scale, BLOCK_SIZE_M: tl.constexpr, ROWS_W: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, SCALE_GROUP_K: tl.constexpr, VALUES_PER_BYTE: tl.constexpr, ): """MX 'scalar' path: CUDA-core FMA GEMV, unpacking MXFP4 weights to E4M3 then dequantizing activation + weight per-element by their expanded group scales, reducing and accumulating into ``acc`` (returned updated). No tensor core (so no M→16 MMA pad) — wins for the memory-bound decode GEMV (M=1). Single weight — for the gate∪up pair use ``mx_scalar_reduce_gate_up``.""" wq = mxfp4_e2m1_to_e4m3(w) if VALUES_PER_BYTE == 2 else w a_deq = tl.trans( a.to(tl.float32) * expand_ue8m0_scale(a_scale, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K) ) w_exp = tl.trans(expand_ue8m0_scale(w_scale, ROWS_W, BLOCK_SIZE_K, SCALE_GROUP_K)) return acc + tl.sum(a_deq * (wq.to(tl.float32) * w_exp), axis=0)[None, :] @triton.jit def mx_scalar_reduce_gate_up( acc_gate, acc_up, a, w_gate, w_up, a_scale, gate_scale, up_scale, BLOCK_SIZE_M: tl.constexpr, ROWS_W: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, SCALE_GROUP_K: tl.constexpr, ): """Gate∪up 'scalar' path: dequant the SHARED activation once, reduce both projections and accumulate into ``acc_gate``/``acc_up`` (returned updated).""" a_deq = tl.trans( a.to(tl.float32) * expand_ue8m0_scale(a_scale, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K) ) acc_gate += tl.sum( a_deq * ( w_gate.to(tl.float32) * tl.trans( expand_ue8m0_scale(gate_scale, ROWS_W, BLOCK_SIZE_K, SCALE_GROUP_K) ) ), axis=0, )[None, :] acc_up += tl.sum( a_deq * ( w_up.to(tl.float32) * tl.trans( expand_ue8m0_scale(up_scale, ROWS_W, BLOCK_SIZE_K, SCALE_GROUP_K) ) ), axis=0, )[None, :] return acc_gate, acc_up @triton.jit def mx_compute_gate_up( acc_gate, acc_up, a, a_scale, b_gate, b_up, gate_scale, up_scale, COMPUTE_MODE: tl.constexpr, VALUES_PER_BYTE: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, SCALE_GROUP_K: tl.constexpr, ): """Gate∪up MMA step dispatched on ``COMPUTE_MODE``: scaled-MMA on the raw weights (``b_gate``/``b_up``), or fp8 ``tl.dot`` + per-group rescale / scalar reduce on the E4M3-decoded weights. Returns the updated ``(acc_gate, acc_up)``.""" w_gate = mxfp4_e2m1_to_e4m3(b_gate) if VALUES_PER_BYTE == 2 else b_gate w_up = mxfp4_e2m1_to_e4m3(b_up) if VALUES_PER_BYTE == 2 else b_up if COMPUTE_MODE == "dot_scaled": acc_gate = mx_dot_scaled( acc_gate, a, a_scale, b_gate, gate_scale, VALUES_PER_BYTE ) acc_up = mx_dot_scaled(acc_up, a, a_scale, b_up, up_scale, VALUES_PER_BYTE) elif COMPUTE_MODE == "dot": acc_gate, acc_up = mx_dot_rescale_gate_up( acc_gate, acc_up, a, w_gate, w_up, a_scale, gate_scale, up_scale ) else: # scalar acc_gate, acc_up = mx_scalar_reduce_gate_up( acc_gate, acc_up, a, w_gate, w_up, a_scale, gate_scale, up_scale, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, ) return acc_gate, acc_up @triton.jit def mx_compute( acc, a, a_scale, w, w_scale, COMPUTE_MODE: tl.constexpr, VALUES_PER_BYTE: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, SCALE_GROUP_K: tl.constexpr, ): """Single-projection MMA step dispatched on ``COMPUTE_MODE``: scaled-MMA on the raw weight (``w``), or fp8 ``tl.dot`` + per-group rescale / scalar reduce on the E4M3-decoded weight. Returns the updated ``acc``.""" if COMPUTE_MODE == "dot_scaled": acc = mx_dot_scaled(acc, a, a_scale, w, w_scale, VALUES_PER_BYTE) elif COMPUTE_MODE == "dot": acc = mx_dot_rescale(acc, a, w, a_scale, w_scale, VALUES_PER_BYTE) else: # scalar acc = mx_scalar_reduce( acc, a, a_scale, w, w_scale, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, VALUES_PER_BYTE, ) return acc @triton.jit def glu( gate, up, ACT_FN: tl.constexpr = "silu", SWIGLU_ALPHA: tl.constexpr = None, SWIGLU_LIMIT: tl.constexpr = None, OUT_DTYPE: tl.constexpr = tl.float32, SIMULATE_UNFUSED: tl.constexpr = False, ): """Gated linear unit on the gate/up matmul accumulators. ``SWIGLU_LIMIT`` clamps gate above and up to ``[-LIMIT, LIMIT]``; ``SWIGLU_ALPHA`` gives the clamped/scaled SwiGLU ``(up + 1) * gate * sigmoid(ALPHA * gate)`` (GPT-OSS / MiniMax), else ``ACT_FN(gate) * up`` (``ACT_FN`` in {silu, gelu, relu}, gelu exact via erf). ``SIMULATE_UNFUSED`` rounds each materialized value through ``OUT_DTYPE`` to match the unfused (separate-kernel) path, where every intermediate lands in that dtype.""" g = gate u = up if SIMULATE_UNFUSED: g = g.to(OUT_DTYPE).to(tl.float32) u = u.to(OUT_DTYPE).to(tl.float32) if SWIGLU_LIMIT is not None: g = tl.minimum(g, SWIGLU_LIMIT) u = tl.minimum(tl.maximum(u, -SWIGLU_LIMIT), SWIGLU_LIMIT) if SWIGLU_ALPHA is not None: gate_scaled = g * SWIGLU_ALPHA if SIMULATE_UNFUSED: gate_scaled = gate_scaled.to(OUT_DTYPE).to(tl.float32) sig = tl.sigmoid(gate_scaled) if SIMULATE_UNFUSED: sig = sig.to(OUT_DTYPE).to(tl.float32) act = g * sig u = u + 1.0 elif ACT_FN == "silu": act = g * tl.sigmoid(g) elif ACT_FN == "gelu": act = 0.5 * g * (1.0 + tl.erf(g * 0.7071067811865476)) elif ACT_FN == "relu": act = tl.maximum(g, 0.0) else: tl.static_assert( False, "unsupported ACT_FN; expected 'silu', 'gelu', or 'relu'" ) if SIMULATE_UNFUSED: act = act.to(OUT_DTYPE).to(tl.float32) u = u.to(OUT_DTYPE).to(tl.float32) gated = act * u if SIMULATE_UNFUSED: gated = gated.to(OUT_DTYPE).to(tl.float32) return gated @triton.jit def _e2m1_code_to_f32(code): """One E2M1 4-bit code -> fp32. Layout ``[sign | exp(2) | mant(1)]``; the 8 magnitudes are ``{0, .5, 1, 1.5, 2, 3, 4, 6}`` (exp==0 is the 0/0.5 subnormal).""" code = code.to(tl.int32) s = (code >> 3) & 1 e = (code >> 1) & 3 m = (code & 1).to(tl.float32) pow2 = tl.exp2((e - 1).to(tl.float32)) # e in 0..3 -> 0.5, 1, 2, 4 mag = tl.where(e == 0, 0.5 * m, (1.0 + 0.5 * m) * pow2) return (1.0 - 2.0 * s.to(tl.float32)) * mag @triton.jit def mxfp4_e2m1_to_e4m3(b_packed): """Unpack packed MXFP4 (E2M1, two nibbles/byte along K) to E4M3, doubling the K (row) dim: ``(R, C) uint8 -> (2R, C) E4M3``. E2M1's 8 magnitudes are all exact in E4M3, so this is lossless — it lets the FP8 ``tl.dot`` path stand in for ``tl.dot_scaled`` at decode (avoiding its M->128 pad). K order is the low nibble first: ``[byte0_lo, byte0_hi, byte1_lo, ...]``.""" lo = _e2m1_code_to_f32(b_packed & 0xF) hi = _e2m1_code_to_f32(b_packed >> 4) # interleave along the K (row) dim via trans -> interleave-last-dim -> trans back unpacked = tl.trans(tl.interleave(tl.trans(lo), tl.trans(hi))) return unpacked.to(tl.float8e4nv) # ── fp8_act_quant kernel (used by tensor-mode FP8 wrappers) ─────────────────── @triton.jit def _fp8_act_quant_kernel( x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr, PADDED_BLOCK: tl.constexpr ): # ``tl.arange`` needs a power-of-2 length, so iterate over PADDED_BLOCK (the next # power of 2) and mask the tail — lets block_size be non-power-of-2 (e.g. a full # row K=14336 in tensor-mode). Masked lanes load 0, which can't affect ``amax``. pid = tl.program_id(axis=0) cols = tl.arange(0, PADDED_BLOCK) mask = cols < BLOCK_SIZE offs = pid * BLOCK_SIZE + cols x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) s = tl.max(tl.abs(x)) / 448.0 # float8_e4m3fn max y = (x / tl.maximum(s, 1e-12)).to(y_ptr.dtype.element_ty) tl.store(y_ptr + offs, y, mask=mask) tl.store(s_ptr + pid, s) @triton_op(add_op_namespace_prefix("fp8_act_quant"), mutates_args=()) def _fp8_act_quant( x: torch.Tensor, block_size: int = 128 ) -> tuple[torch.Tensor, torch.Tensor]: assert x.is_contiguous() assert x.shape[-1] % block_size == 0 y = torch.empty_like(x, dtype=FP8_DTYPE) grid = (triton.cdiv(x.numel(), block_size),) s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32) with device_context(x.device): wrap_triton(_fp8_act_quant_kernel)[grid]( x, y, s, BLOCK_SIZE=block_size, PADDED_BLOCK=triton.next_power_of_2(block_size), ) return y, s def fp8_act_quant( x: torch.Tensor, block_size: int = 128 ) -> tuple[torch.Tensor, torch.Tensor]: """Quantize activations to FP8 with per-block dynamic scaling. Splits the last dimension of ``x`` into blocks of ``block_size`` elements, computes ``scale = max(|x_block|) / 448`` per block, and quantizes to ``float8_e4m3fn``. Args: x: Input tensor in bf16/fp16/fp32. Last dimension must be divisible by ``block_size`` and the tensor must be contiguous. block_size: Number of elements per quantization block (default: 128). Returns: A tuple ``(quantized, scales)`` where ``quantized`` has dtype ``float8_e4m3fn`` with the same shape as ``x``, and ``scales`` has shape ``(*x.shape[:-1], x.shape[-1] // block_size)`` in float32. """ return ops.fp8_act_quant(x, block_size)