diff --git a/README.md b/README.md index 830a94b..116d3d9 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,7 @@ CUDA kernels for Prime Intellect training stacks, shipped as one wheel, `prime-k │ ├── mxfp8.py │ └── csrc/ # the C++/CUDA sources compiled into prime_kernels.flash_moe._C ├── indexed_attention/ # Python-only TileLang indexed GQA forward + backward + ├── dsa_indexer_topk/ # Python-only Triton DeepSeek-V4.1 Lightning Indexer + top-k └── rmsnorm/ ├── __init__.py ├── csrc/ # the torch binding @@ -64,6 +65,54 @@ for context parallelism without gathering queries. It supports SM80, SM90, SM100, and SM103 (B300), and requires TileLang (validated with 0.1.12). Install TileLang separately; the registry reports it missing when unavailable. +`dsa_indexer_topk` is DeepSeek-V4.1's Lightning Indexer forward fused with its top-k, including +the two-level candidate-block filter, as a drop-in for prime-rl's `dsv41_index_topk` (same +arguments, same picks up to ties). It never materializes the `(queries, entries)` score matrix: +a dense FP8 GEMM keeps only per-group score maxima, the best groups are selected per query, and +only their entries are rescored and ranked; candidate consumer layers score only their candidate +entries. Forward only (the indexer is frozen). Triton, validated on SM90 (H200). + +`moe_experts` is the expert MLP of a rank's local experts for training on Hopper (SM90): +`down(silu(min(gate, l)) * clamp(up, -l, l))` over tokens already grouped by expert, the +DeepSeek-V4 clamped SwiGLU, bf16, differentiable in the tokens and all three weights. It is +the Hopper counterpart of the expert compute inside cuDNN's MegaMoE (which also fuses the +cross-GPU dispatch and is Blackwell only); dispatch and routing scores stay with the caller. +Its grouped GEMMs are persistent, warp-specialized Gluon kernels whose epilogue warpgroup +applies the SwiGLU (forward) or its derivative (backward, also recomputing the activation) +while the tensor cores run the next tile, so only `x`, `gate` and `up` are saved for backward. +The weight gradients use `torch._grouped_mm`. Group sizes must be multiples of +`moe_experts.TOKEN_GROUP_ALIGNMENT`; `tests/moe_experts/bench_moe_experts.py` compares it with +`torch._grouped_mm` plus an unfused activation. + +`moe_experts(..., fp8=True)` runs the same MLP in blockwise FP8 (DeepSeek-V3 recipe: e4m3, 1 x 128 +scales for activations and gradients along each GEMM's K, 128 x 128 for weights, fp32 +accumulation) for the forward, the data gradients and the weight gradients. The GEMMs are +DeepGEMM's grouped FP8 GEMMs (DeepGEMM must be installed); everything around them is fused into a +few Triton/Gluon passes: one read of `x` (and of the output gradient) writes both its row- and +column-quantized copies, the SwiGLU passes quantize their outputs, and DeepGEMM writes the outputs +in place at the tokens' rows. Hidden and intermediate sizes must be multiples of 128. + +`dsa_sparse_attn_bwd` is the backward of DeepSeek-V4.1's sparse attention on Hopper (SM90): +every query reads its own list of rows of one shared K = V latent buffer (`-1` = empty slot), with +a per-head attention sink, 64 heads of 512. It is a drop-in for prime-rl's +`dsv41_sparse_attn_backward` (same inputs, including FlashMLA's sink-free LSE; returns dq, dkv, +dsinks) as the custom op `prime_kernels::dsa_sparse_attn_bwd` (with a fake impl), and +`sparse_attn_backward_flat` takes the unbatched layout of cuDNN's `flash_attn_bwd_sm90`. A Triton +pass computes delta = rowsum(dO * O), the sink-aware LSE and dsinks and compacts each query's +non-empty slots to the front; the main kernel (CuTe DSL) is a persistent, two-warpgroup kernel +with one query per iteration: S / dP / softmax / dQ in registers over tiles of 64 gathered latent +rows, and dKV added with fp32 vector atomics into a buffer that a last Triton pass casts to bf16. +The atomics (2 KB per query-slot) bound it; at V4.1 shapes it is ~1.2x faster than cuDNN's SM90 +backward (`tests/dsa_sparse_attn_bwd/bench_dsa_sparse_attn_bwd.py`). + +`mhc_projection` is DeepSeek-V4.1's manifold-constrained hyper-connection (mHC) projection for +SM90, forward and backward, in Triton: one pass over the `(tokens, hc_mult, hidden)` streams +computes the RMS statistic, the projection to the `(2 + hc_mult) * hc_mult` gate logits and the +collapse of the streams by the previous sublayer's `pre` gate; a second kernel applies the gate +activations (sigmoid `pre`/`post`, Sinkhorn `comb`). `hyper_connection` chains both, with +`torch.library` custom ops (and fake impls) so it can sit inside `torch.compile`d blocks. +`tests/mhc_projection/bench_mhc_projection.py` compares it against prime-rl's path. + ## Installing prime-rl's `uv sync --extra kernels` installs the prebuilt wheels attached to a prime-kernels diff --git a/prime_kernels/dsa_indexer_topk/__init__.py b/prime_kernels/dsa_indexer_topk/__init__.py new file mode 100644 index 0000000..d105304 --- /dev/null +++ b/prime_kernels/dsa_indexer_topk/__init__.py @@ -0,0 +1,3 @@ +from prime_kernels.dsa_indexer_topk.indexer import dsv41_index_topk, unsupported_shape_reason + +__all__ = ["dsv41_index_topk", "unsupported_shape_reason"] diff --git a/prime_kernels/dsa_indexer_topk/indexer.py b/prime_kernels/dsa_indexer_topk/indexer.py new file mode 100644 index 0000000..55c46f3 --- /dev/null +++ b/prime_kernels/dsa_indexer_topk/indexer.py @@ -0,0 +1,269 @@ +import torch +import triton + +from prime_kernels.dsa_indexer_topk.kernels import ( + SUPER, + _gather_topk_kernel, + _group_max_kernel, + _group_topk_kernel, +) +from prime_kernels.dsa_indexer_topk.quant import per_token_quant_fp8 + +GROUP_MAX_BYTES = 1 << 29 +GROUP_MAX_BLOCK_M = 64 +GROUP_MAX_BLOCK_N = 128 +GATHER_TILE = 128 +SELECT_CHUNK = 1024 +SELECT_CAP = 1024 +# The persistent kernels are latency bound: capping registers fits 4 programs per SM. +MAX_REGISTERS = 128 + +_RESIDENT_PROGRAMS: dict[tuple, int] = {} + + +def unsupported_shape_reason(num_heads: int, head_dim: int, block_size: int) -> str | None: + """Why `dsv41_index_topk` cannot run these shapes, or None.""" + if num_heads < 16 or num_heads & (num_heads - 1): + return f"the number of index heads must be a power of two >= 16, got {num_heads}" + if head_dim < 32 or head_dim & (head_dim - 1): + return f"the index head dim must be a power of two >= 32, got {head_dim}" + if block_size < 1 or block_size > 16 or block_size & (block_size - 1): + return f"the candidate block size must be a power of two <= 16, got {block_size}" + return None + + +def _plain_group(max_entries_per_doc: int, topk: int) -> int: + """Group size when no candidate blocks are emitted, where it is free: larger groups shrink the + per-query group selection, smaller ones the entries rescored (`topk` groups' worth).""" + return min(8, max(2, triton.next_power_of_2(triton.cdiv(max_entries_per_doc, 8 * topk)))) + + +def _launch_persistent(kernel, n_rows: int, device: torch.device, make_args, **meta) -> None: + """Launch a persistent per-row kernel with as many programs as fit on the GPU at once. + + `make_args(n_programs)` returns the positional arguments, scratch sized for `n_programs`. + """ + key = (kernel.fn.__name__, device.index, tuple(sorted(meta.items()))) + if key not in _RESIDENT_PROGRAMS: + compiled = kernel.warmup(*make_args(1), grid=(1,), **meta) + compiled._init_handles() + props = torch.cuda.get_device_properties(device) + regs_per_warp = -(-compiled.n_regs * 32 // 256) * 256 + by_regs = props.regs_per_multiprocessor // regs_per_warp // meta["num_warps"] + by_smem = props.shared_memory_per_multiprocessor // (compiled.metadata.shared + 1024) + by_threads = props.max_threads_per_multi_processor // (32 * meta["num_warps"]) + _RESIDENT_PROGRAMS[key] = props.multi_processor_count * max(1, min(by_regs, by_smem, by_threads)) + n_programs = min(n_rows, _RESIDENT_PROGRAMS[key]) + kernel[(n_programs,)](*make_args(n_programs), **meta) + + +def _group_max(q_fp8, k_fp8, k_scale, w, ks, ke, n_groups: int, group: int) -> tuple[torch.Tensor, torch.Tensor]: + """`(S_q, n_groups)` group maxima and `(S_q, n_groups / SUPER)` super-group maxima.""" + S_q, H, D = q_fp8.shape + gm = torch.empty(S_q, n_groups, dtype=torch.float32, device=q_fp8.device) + sgm = torch.empty(S_q, triton.cdiv(n_groups, SUPER.value), dtype=torch.float32, device=q_fp8.device) + grid = (triton.cdiv(S_q, GROUP_MAX_BLOCK_M), triton.cdiv(n_groups * group, GROUP_MAX_BLOCK_N)) + _group_max_kernel[grid]( + q_fp8, + k_fp8, + k_scale, + w, + ks, + ke, + gm, + sgm, + S_q, + k_fp8.shape[0], + n_groups, + sgm.shape[1], + q_fp8.stride(0), + q_fp8.stride(1), + k_fp8.stride(0), + w.stride(0), + gm.stride(0), + sgm.stride(0), + H=H, + D=D, + GROUP=group, + BLOCK_M=GROUP_MAX_BLOCK_M, + BLOCK_N=GROUP_MAX_BLOCK_N, + num_warps=4, + num_stages=3, + ) + return gm, sgm + + +def _group_topk(gm, sgm, ks, ke, out, tau, k: int, group: int, pin_newest: bool) -> None: + """Best `k` groups per query into `out`; with `tau`, also the rescoring threshold.""" + + def make_args(n_programs): + return ( + gm, + sgm, + ks, + ke, + torch.empty(n_programs, gm.shape[1], dtype=torch.float32, device=gm.device), + torch.empty(n_programs, gm.shape[1], dtype=torch.int32, device=gm.device), + out, + out if tau is None else tau, + gm.shape[0], + gm.shape[1], + sgm.shape[1], + gm.stride(0), + sgm.stride(0), + out.stride(0), + ) + + _launch_persistent( + _group_topk_kernel, + gm.shape[0], + gm.device, + make_args, + K=k, + GROUP=group, + PIN_NEWEST=pin_newest, + WRITE_TAU=tau is not None, + CHUNK=SELECT_CHUNK, + CAP=max(SELECT_CAP, triton.next_power_of_2(k)), + num_warps=4, + maxnreg=MAX_REGISTERS, + ) + + +def _gather_topk(q_fp8, k_fp8, k_scale, w, ks, ke, blocks, tau, out, topk: int, group: int) -> None: + """Best `topk` entries per query among its `blocks`; `tau` from `_group_topk` when they are its picks.""" + S_q, H, D = q_fp8.shape + n_cols = blocks.shape[1] * group + + def make_args(n_programs): + return ( + q_fp8, + k_fp8, + k_scale, + w, + ks, + ke, + blocks, + blocks if tau is None else tau, + torch.empty(n_programs, 2 * n_cols + blocks.shape[1], dtype=torch.float32, device=q_fp8.device), + torch.empty(n_programs, n_cols, dtype=torch.int32, device=q_fp8.device), + out, + S_q, + k_fp8.shape[0], + blocks.shape[1], + q_fp8.stride(0), + q_fp8.stride(1), + k_fp8.stride(0), + w.stride(0), + blocks.stride(0), + out.stride(0), + ) + + _launch_persistent( + _gather_topk_kernel, + S_q, + q_fp8.device, + make_args, + H=H, + D=D, + GROUP=group, + TOPK=topk, + TILE=GATHER_TILE, + CHUNK=SELECT_CHUNK, + CAP=max(SELECT_CAP, triton.next_power_of_2(topk)), + KNOWN_TAU=tau is not None, + SCORE_STAGES=2, + maxnreg=MAX_REGISTERS, + num_warps=4, + ) + + +@torch.library.custom_op("prime_kernels::dsv41_index_topk", mutates_args=()) +def dsv41_index_topk( + q: torch.Tensor, + k: torch.Tensor, + w: torch.Tensor, + ks: torch.Tensor, + ke: torch.Tensor, + topk: int, + candidates: torch.Tensor | None, + emit_candidates: bool, + max_entries_per_doc: int, + block_size: int, + topk_blocks: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """DeepSeek-V4.1 Lightning Indexer top-k without a `(queries, entries)` score matrix. + + Drop-in for prime-rl's `dsv41_index_topk` (same arguments, same results up to ties): + `score[t, e] = sum_h relu(q[t, h] . k[e]) * w[t, h]` over e in `[ks[t], ke[t])`, computed on + per-token FP8 (UE8M0-scaled) q and k with fp32 accumulation. + + Args: + q: `(S_q, H, D)` bf16 index queries. k: `(S_k, D)` bf16 index keys, shared across heads. + w: `(S_q, H)` per-head weights. + ks, ke: `(S_q,)` int32, each query's readable entries `[ks, ke)`; `ks` is its document's first + entry and `ke - ks <= max_entries_per_doc`. + candidates: `(S_q, topk_blocks)` int32 doc-local block indices (-1 unused) restricting the + ranked entries, or None. + emit_candidates: also return each query's `topk_blocks` best doc-local blocks of `block_size` + entries, ranked by their best entry, with the block of the query's newest entry pinned. + + Returns `(S_q, topk)` int64 entry indices (unordered, -1 where a query has fewer readable + entries) and the `(S_q, topk_blocks)` int32 candidate blocks (an empty tensor unless + `emit_candidates`). + + Without `candidates`, a dense GEMM keeps only the max of every doc-local group of entries + (`block_size` entries when emitting candidates, else `_plain_group`; `(S_q, S_k / group)` fp32, + processed in query chunks of about `GROUP_MAX_BYTES`), the top `topk` groups are selected from + them, and only those groups' entries are rescored and ranked. With `candidates`, only the + candidate entries are scored. + """ + S_q, H, D = q.shape + reason = unsupported_shape_reason(H, D, block_size) + if reason is not None: + raise ValueError(reason) + device = q.device + out = torch.empty((S_q, topk), dtype=torch.int64, device=device) + cand_out = torch.empty((S_q, topk_blocks) if emit_candidates else (0, 0), dtype=torch.int32, device=device) + if S_q == 0 or topk == 0: + return out, cand_out + + q_fp8, q_scale = per_token_quant_fp8(q.contiguous()) + if k.shape[0] == 0: + # Nothing is readable; a zero row keeps every (masked) key load in bounds. + k = k.new_zeros(1, D) + k_fp8, k_scale = per_token_quant_fp8(k.contiguous()) + w = (w.float() * q_scale).contiguous() + ks = ks.to(torch.int32).contiguous() + ke = ke.to(torch.int32).contiguous() + torch._assert_async((ke - ks <= max_entries_per_doc).all(), "dsv41_index_topk: ke - ks > max_entries_per_doc") + + group = block_size if emit_candidates else _plain_group(max_entries_per_doc, topk) + n_groups = triton.cdiv(max_entries_per_doc, group) + best_groups = torch.empty((S_q, topk), dtype=torch.int32, device=device) if candidates is None else None + tau = torch.empty(S_q, dtype=torch.float32, device=device) if candidates is None else None + if best_groups is not None or emit_candidates: + rows_per_chunk = max(1, GROUP_MAX_BYTES // (4 * max(n_groups, 1))) + for r0 in range(0, S_q, rows_per_chunk): + r1 = min(S_q, r0 + rows_per_chunk) + gm, sgm = _group_max(q_fp8[r0:r1], k_fp8, k_scale, w[r0:r1], ks[r0:r1], ke[r0:r1], n_groups, group) + if best_groups is not None: + _group_topk( + gm, sgm, ks[r0:r1], ke[r0:r1], best_groups[r0:r1], tau[r0:r1], topk, group, pin_newest=False + ) + if emit_candidates: + _group_topk( + gm, sgm, ks[r0:r1], ke[r0:r1], cand_out[r0:r1], None, topk_blocks, block_size, pin_newest=True + ) + del gm, sgm # free this chunk's maxima before the next one is allocated + blocks = best_groups if candidates is None else candidates.to(torch.int32).contiguous() + _gather_topk(q_fp8, k_fp8, k_scale, w, ks, ke, blocks, tau, out, topk, group if candidates is None else block_size) + return out, cand_out + + +@dsv41_index_topk.register_fake +def _dsv41_index_topk_fake( + q, k, w, ks, ke, topk, candidates, emit_candidates, max_entries_per_doc, block_size, topk_blocks +): + cand_shape = (q.shape[0], topk_blocks) if emit_candidates else (0, 0) + return q.new_empty((q.shape[0], topk), dtype=torch.int64), q.new_empty(cand_shape, dtype=torch.int32) diff --git a/prime_kernels/dsa_indexer_topk/kernels.py b/prime_kernels/dsa_indexer_topk/kernels.py new file mode 100644 index 0000000..68f14b2 --- /dev/null +++ b/prime_kernels/dsa_indexer_topk/kernels.py @@ -0,0 +1,486 @@ +"""Triton kernels of the fused Lightning Indexer top-k. + +Scores are `score[t, e] = k_scale[e] * sum_h relu(q_fp8[t, h] . k_fp8[e]) * w[t, h]`, where `w` already +carries the per-(token, head) query scale. Three kernels share that formula: + +- `_group_max_kernel`: the dense query-block x entry-tile GEMM. It keeps no scores, only the max of + every doc-local group of `GROUP` entries (`GROUP` = the candidate block size) and of every + super-group of `SUPER` consecutive groups. +- `_group_topk_kernel`: per query, selects the best groups by their max. +- `_gather_topk_kernel`: per query, scores only the entries of a given list of groups and selects + the best `TOPK` of them. + +Every selection is hierarchical, which is exact for the following reason. Split a row into blocks; +if `tau` is the `k`-th largest block max, at least `k` values are `>= tau`, so every top-`k` value +is `>= tau`. Selecting the `k`-th largest block max, keeping only the values `>= tau` (a few more +than `k` for typical rows, at most `k` blocks' worth) and selecting among those gives the top `k`. +The same argument makes V4.1's candidate filter cheap: the dense pass only needs block maxima, and +the final ranking only rescores the entries of the top `topk` groups. + +Selection itself is a 4-pass 8-bit radix select over an order-preserving uint32 image of the fp32 +values, followed by a pass that writes the strictly larger values and the first ties in order. +""" + +import triton +import triton.language as tl + +SUPER = tl.constexpr(4) + + +@triton.jit +def _ordered_key(v): + """fp32 -> uint32 with the same order (larger float, larger key).""" + bits = v.to(tl.uint32, bitcast=True) + return tl.where((bits >> 31) != 0, ~bits, bits | 0x80000000) + + +@triton.jit +def _load_keys(row_ptr, cols, n_cols): + """Keys of `row_ptr[cols]` and which are valid (finite).""" + v = tl.load(row_ptr + cols, mask=cols < n_cols, other=float("-inf")) + return _ordered_key(v), v > float("-inf") + + +@triton.jit +def _digit_choice(hist, prefix, prefix_mask, remaining, SHIFT: tl.constexpr): + """Given the histogram of the 8-bit digit at `SHIFT` among the keys matching `prefix`, fix that + digit of the `remaining`-th largest key.""" + bins = tl.arange(0, 256) + at_or_above = tl.cumsum(hist, 0, reverse=True) + digit = tl.max(tl.where(at_or_above >= remaining, bins, -1)) + above = tl.sum(tl.where(bins > digit, hist, 0)) + prefix = prefix | (digit.to(tl.uint32) << SHIFT) + prefix_mask = prefix_mask | (tl.full([], 0xFF, tl.uint32) << SHIFT) + return prefix, prefix_mask, remaining - above + + +@triton.jit +def _radix_pass_regs(key, valid, prefix, prefix_mask, remaining, SHIFT: tl.constexpr): + match = valid & ((key & prefix_mask) == prefix) + hist = tl.histogram(((key >> SHIFT) & 0xFF).to(tl.int32), 256, mask=match) + return _digit_choice(hist, prefix, prefix_mask, remaining, SHIFT) + + +@triton.jit +def _radix_pass_mem(row_ptr, n_cols, prefix, prefix_mask, remaining, SHIFT: tl.constexpr, CHUNK: tl.constexpr): + hist = tl.zeros([256], dtype=tl.int32) + for c0 in range(0, n_cols, CHUNK): + key, valid = _load_keys(row_ptr, c0 + tl.arange(0, CHUNK), n_cols) + match = valid & ((key & prefix_mask) == prefix) + hist += tl.histogram(((key >> SHIFT) & 0xFF).to(tl.int32), 256, mask=match) + return _digit_choice(hist, prefix, prefix_mask, remaining, SHIFT) + + +@triton.jit +def _kth_key_regs(key, valid, k): + """Key of the `k`-th largest valid key and how many of the top `k` equal it (needs `k` valid).""" + prefix = tl.zeros([], dtype=tl.uint32) + prefix_mask = tl.zeros([], dtype=tl.uint32) + remaining = tl.zeros([], dtype=tl.int32) + k + prefix, prefix_mask, remaining = _radix_pass_regs(key, valid, prefix, prefix_mask, remaining, 24) + prefix, prefix_mask, remaining = _radix_pass_regs(key, valid, prefix, prefix_mask, remaining, 16) + prefix, prefix_mask, remaining = _radix_pass_regs(key, valid, prefix, prefix_mask, remaining, 8) + prefix, prefix_mask, remaining = _radix_pass_regs(key, valid, prefix, prefix_mask, remaining, 0) + return prefix, remaining + + +@triton.jit +def _kth_key_mem(row_ptr, n_cols, k, CHUNK: tl.constexpr): + """`_kth_key_regs` over `row_ptr[:n_cols]`, streamed in chunks.""" + prefix = tl.zeros([], dtype=tl.uint32) + prefix_mask = tl.zeros([], dtype=tl.uint32) + remaining = tl.zeros([], dtype=tl.int32) + k + prefix, prefix_mask, remaining = _radix_pass_mem(row_ptr, n_cols, prefix, prefix_mask, remaining, 24, CHUNK) + prefix, prefix_mask, remaining = _radix_pass_mem(row_ptr, n_cols, prefix, prefix_mask, remaining, 16, CHUNK) + prefix, prefix_mask, remaining = _radix_pass_mem(row_ptr, n_cols, prefix, prefix_mask, remaining, 8, CHUNK) + prefix, prefix_mask, remaining = _radix_pass_mem(row_ptr, n_cols, prefix, prefix_mask, remaining, 0, CHUNK) + return prefix, remaining + + +@triton.jit +def _block_threshold(bmax_ptr, n_blocks, K: tl.constexpr, CHUNK: tl.constexpr, CAP: tl.constexpr): + """Key of the `K`-th largest valid block max, or 0 (below every valid key) if fewer are valid.""" + threshold = tl.zeros([], dtype=tl.uint32) + if n_blocks <= CAP: + key, valid = _load_keys(bmax_ptr, tl.arange(0, CAP), n_blocks) + if tl.sum(valid.to(tl.int32)) >= K: + threshold, _ = _kth_key_regs(key, valid, K) + else: + n_valid = tl.zeros([], dtype=tl.int32) + for c0 in range(0, n_blocks, CHUNK): + _, valid = _load_keys(bmax_ptr, c0 + tl.arange(0, CHUNK), n_blocks) + n_valid += tl.sum(valid.to(tl.int32)) + if n_valid >= K: + threshold, _ = _kth_key_mem(bmax_ptr, n_blocks, K, CHUNK) + return threshold + + +@triton.jit +def _store_picks(out_ptr, pos, idx, take, blk_ptr, ks, GROUP: tl.constexpr, GATHERED: tl.constexpr): + """Write column `idx` (as an entry index when `GATHERED`) to `out_ptr[pos]` where `take`.""" + if GATHERED: + blk = tl.load(blk_ptr + idx // GROUP, mask=take, other=0) + value = ks.to(tl.int64) + blk.to(tl.int64) * GROUP + (idx % GROUP) + else: + value = idx + tl.store(out_ptr + pos, value.to(out_ptr.dtype.element_ty), mask=take) + + +@triton.jit +def _select_compacted( + cval_ptr, + cidx_ptr, + n_kept, + out_ptr, + blk_ptr, + ks, + K: tl.constexpr, + GROUP: tl.constexpr, + CHUNK: tl.constexpr, + CAP: tl.constexpr, + GATHERED: tl.constexpr, +): + """Write the columns `cidx_ptr[c]` of the `K` largest of `cval_ptr[:n_kept]` (all valid), -1 after. + + Ties at the `K`-th value are taken in compaction order. Up to `CAP` values are selected in + registers, more are streamed. Returns the key of the `K`-th largest value (0 if fewer). + """ + tl.static_assert(CAP >= K) + n_take = tl.minimum(n_kept, K) + if n_kept <= CAP: + offs = tl.arange(0, CAP) + valid = offs < n_kept + key = _ordered_key(tl.load(cval_ptr + offs, mask=valid, other=float("-inf"))) + if n_kept > K: + threshold, remaining = _kth_key_regs(key, valid, K) + bound = threshold + else: + threshold = tl.zeros([], dtype=tl.uint32) + remaining = tl.zeros([], dtype=tl.int32) + bound = tl.where(n_kept == K, tl.min(tl.where(valid, key, 0xFFFFFFFF)), 0).to(tl.uint32) + greater = (valid & (key > threshold)).to(tl.int32) + equal = (valid & (key == threshold)).to(tl.int32) + equal_rank = tl.cumsum(equal, 0) - 1 + take = (greater != 0) | ((equal != 0) & (equal_rank < remaining)) + pos = tl.where(greater != 0, tl.cumsum(greater, 0) - 1, n_take - remaining + equal_rank) + idx = tl.load(cidx_ptr + offs, mask=take, other=0) + _store_picks(out_ptr, pos, idx, take, blk_ptr, ks, GROUP, GATHERED) + else: + threshold, remaining = _kth_key_mem(cval_ptr, n_kept, K, CHUNK) + bound = threshold + seen_greater = tl.zeros([], dtype=tl.int32) + seen_equal = tl.zeros([], dtype=tl.int32) + for c0 in range(0, n_kept, CHUNK): + cols = c0 + tl.arange(0, CHUNK) + key, valid = _load_keys(cval_ptr, cols, n_kept) + greater = (valid & (key > threshold)).to(tl.int32) + equal = (valid & (key == threshold)).to(tl.int32) + equal_rank = seen_equal + tl.cumsum(equal, 0) - 1 + take = (greater != 0) | ((equal != 0) & (equal_rank < remaining)) + pos = tl.where(greater != 0, seen_greater + tl.cumsum(greater, 0) - 1, K - remaining + equal_rank) + idx = tl.load(cidx_ptr + cols, mask=take, other=0) + _store_picks(out_ptr, pos, idx, take, blk_ptr, ks, GROUP, GATHERED) + seen_greater += tl.sum(greater) + seen_equal += tl.sum(equal) + for p0 in range(0, K, CHUNK): + pos = p0 + tl.arange(0, CHUNK) + tl.store(out_ptr + pos, tl.full([CHUNK], -1, out_ptr.dtype.element_ty), mask=(pos >= n_take) & (pos < K)) + return bound + + +@triton.jit +def _key_to_float(key): + """Inverse of `_ordered_key`, with key 0 (below every valid key) mapped to -inf.""" + bits = tl.where((key >> 31) != 0, key ^ 0x80000000, ~key) + return tl.where(key == 0, float("-inf"), bits.to(tl.float32, bitcast=True)) + + +@triton.jit +def _compact(row_ptr, n_cols, pin, tau, cval_ptr, cidx_ptr, CHUNK: tl.constexpr): + """Copy the valid values `>= tau` of `row_ptr[:n_cols]` and their columns to `cval_ptr`/`cidx_ptr`; + column `pin` (if in range) is always kept, as +inf. Returns the counts kept and valid and the + largest unpinned value.""" + n_kept = tl.zeros([], dtype=tl.int32) + n_valid = tl.zeros([], dtype=tl.int32) + row_max = tl.full([], float("-inf"), tl.float32) + for c0 in range(0, n_cols, CHUNK): + cols = c0 + tl.arange(0, CHUNK) + v = tl.load(row_ptr + cols, mask=cols < n_cols, other=float("-inf")) + is_pin = (cols == pin) & (cols < n_cols) + valid = (v > float("-inf")) | is_pin + keep = (valid & (v >= tau)) | is_pin + rank = n_kept + tl.cumsum(keep.to(tl.int32), 0) - 1 + tl.store(cval_ptr + rank, tl.where(is_pin, float("inf"), v), mask=keep) + tl.store(cidx_ptr + rank, cols, mask=keep) + n_kept += tl.sum(keep.to(tl.int32)) + n_valid += tl.sum(valid.to(tl.int32)) + row_max = tl.maximum(row_max, tl.max(tl.where(is_pin, float("-inf"), v))) + return n_kept, n_valid, row_max + + +@triton.jit +def _topk_row( + row_ptr, + n_cols, + pin, + tau, + cval_ptr, + cidx_ptr, + out_ptr, + blk_ptr, + ks, + K: tl.constexpr, + GROUP: tl.constexpr, + CHUNK: tl.constexpr, + CAP: tl.constexpr, + GATHERED: tl.constexpr, +): + """Top-`K` valid columns of `row_ptr[:n_cols]`, column `pin` (if in range) counting as the largest. + + `tau` is a score no top-`K` value is expected below: only values `>= tau` are compacted into + `cval_ptr`/`cidx_ptr` (capacity `n_cols`) and selected from. Should fewer than `K` (or all the + valid ones) reach it, every valid value is. Returns the key of the `K`-th largest value (0 if + fewer) and the largest unpinned value. + """ + n_kept, n_valid, row_max = _compact(row_ptr, n_cols, pin, tau, cval_ptr, cidx_ptr, CHUNK) + if n_kept < tl.minimum(n_valid, K): + n_kept, n_valid, row_max = _compact(row_ptr, n_cols, pin, float("-inf"), cval_ptr, cidx_ptr, CHUNK) + tl.debug_barrier() + bound = _select_compacted(cval_ptr, cidx_ptr, n_kept, out_ptr, blk_ptr, ks, K, GROUP, CHUNK, CAP, GATHERED) + return bound, row_max + + +@triton.jit +def _rescore_threshold(bound, row_max): + """Score a rescoring pass may keep entries above: the `K`-th best group max `bound` (key, 0 if + there is none), lowered by a margin for rounding differences between the two score kernels.""" + tau = _key_to_float(bound) + return tl.where(bound == 0, float("-inf"), tau - 1e-3 * tl.maximum(tl.abs(tau), tl.abs(row_max))) + + +@triton.jit +def _score_tile(q, w, ks, ke, blk_row, cols, n_cols, K_BYTES, K_SCALE, stride_ks, offs_d, GROUP: tl.constexpr): + """Scores of the entries `cols` of a gathered row, -inf where invalid.""" + blk = tl.load(blk_row + cols // GROUP, mask=cols < n_cols, other=-1) + entry = ks + blk * GROUP + cols % GROUP + valid = (blk >= 0) & (entry < ke) + entry_safe = tl.where(valid, entry, 0).to(tl.int64) + # Masked rather than clamped: unused picks would all read one key row. + k_tile = tl.load(K_BYTES + entry_safe[:, None] * stride_ks + offs_d[None, :], mask=valid[:, None], other=0) + s = tl.dot(k_tile.to(tl.float8e4nv, bitcast=True), tl.trans(q)) + score = tl.sum(tl.maximum(s, 0.0) * w[None, :], axis=1) + score *= tl.load(K_SCALE + entry_safe, mask=valid, other=0.0) + return tl.where(valid, score, float("-inf")) + + +@triton.jit +def _group_max_kernel( + Q, + K, + K_SCALE, + W, + KS, + KE, + GM, + SGM, + S_Q, + S_K, + N_GROUPS, + N_SUPER, + stride_qs, + stride_qh, + stride_ks, + stride_ws, + stride_gm, + stride_sgm, + H: tl.constexpr, + D: tl.constexpr, + GROUP: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """`GM[t, g]` = max score over doc-local entries `[g * GROUP, (g + 1) * GROUP)` of query `t`, and + `SGM[t, s]` = max of `GM[t, s * SUPER : (s + 1) * SUPER]`. + + Program `(m, n)` covers queries `[m * BLOCK_M, ...)` and doc-local entries `[n * BLOCK_N, ...)`. + Queries of different documents read different entries, so the tile is computed once per + document among the block's queries (one pass unless the block straddles a document start). + """ + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + mask_m = offs_m < S_Q + offs_m_safe = tl.minimum(offs_m, S_Q - 1).to(tl.int64) + offs_d = tl.arange(0, D) + G_PER_TILE: tl.constexpr = BLOCK_N // GROUP + S_PER_TILE: tl.constexpr = G_PER_TILE // SUPER + + ks = tl.load(KS + offs_m, mask=mask_m, other=0) + n_valid = tl.load(KE + offs_m, mask=mask_m, other=0) - ks + local_lo = pid_n * BLOCK_N + offs_g = pid_n * G_PER_TILE + tl.arange(0, G_PER_TILE) + offs_s = pid_n * S_PER_TILE + tl.arange(0, S_PER_TILE) + gm_ptrs = GM + offs_m[:, None].to(tl.int64) * stride_gm + offs_g[None, :] + sgm_ptrs = SGM + offs_m[:, None].to(tl.int64) * stride_sgm + offs_s[None, :] + gm_mask = mask_m[:, None] & (offs_g[None, :] < N_GROUPS) + sgm_mask = mask_m[:, None] & (offs_s[None, :] < N_SUPER) + + needed = mask_m & (local_lo < n_valid) + tl.store(gm_ptrs, tl.full([BLOCK_M, G_PER_TILE], float("-inf"), tl.float32), mask=gm_mask & ~needed[:, None]) + tl.store(sgm_ptrs, tl.full([BLOCK_M, S_PER_TILE], float("-inf"), tl.float32), mask=sgm_mask & ~needed[:, None]) + done = ~needed + local_cols = local_lo + tl.arange(0, BLOCK_N) + while tl.min(done.to(tl.int32)) == 0: + doc_start = tl.min(tl.where(done, 2147483647, ks)) + active = ~done & (ks == doc_start) + offs_n = tl.minimum(doc_start + local_cols, S_K - 1).to(tl.int64) + k_tile = tl.load(K + offs_n[:, None] * stride_ks + offs_d[None, :]) + k_scale = tl.load(K_SCALE + offs_n) + acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + for h in range(H): + q_h = tl.load(Q + offs_m_safe[:, None] * stride_qs + h * stride_qh + offs_d[None, :]) + w_h = tl.load(W + offs_m_safe * stride_ws + h) + s = tl.dot(q_h, tl.trans(k_tile)) + acc += tl.maximum(s, 0.0) * w_h[:, None] + acc = acc * k_scale[None, :] + acc = tl.where(local_cols[None, :] < n_valid[:, None], acc, float("-inf")) + group_max = tl.max(tl.reshape(acc, [BLOCK_M, G_PER_TILE, GROUP]), axis=2) + super_max = tl.max(tl.reshape(group_max, [BLOCK_M, S_PER_TILE, SUPER]), axis=2) + tl.store(gm_ptrs, group_max, mask=gm_mask & active[:, None]) + tl.store(sgm_ptrs, super_max, mask=sgm_mask & active[:, None]) + done = done | active + + +@triton.jit +def _group_topk_kernel( + GM, + SGM, + KS, + KE, + SCRATCH_V, + SCRATCH_I, + OUT, + TAU, + S_Q, + N_GROUPS, + N_SUPER, + stride_gm, + stride_sgm, + stride_out, + K: tl.constexpr, + GROUP: tl.constexpr, + PIN_NEWEST: tl.constexpr, + WRITE_TAU: tl.constexpr, + CHUNK: tl.constexpr, + CAP: tl.constexpr, +): + """`OUT[t]` = doc-local indices of query `t`'s `K` best groups, -1 for unused picks. + + With `PIN_NEWEST` the group holding the query's newest entry always counts as the best. With + `WRITE_TAU`, `TAU[t]` is the score below which no entry of those groups can be in the top `K`. + Persistent, with one `N_GROUPS`-wide compaction row of scratch per program. + """ + pid = tl.program_id(0) + cval = SCRATCH_V + pid.to(tl.int64) * N_GROUPS + cidx = SCRATCH_I + pid.to(tl.int64) * N_GROUPS + for row in range(pid, S_Q, tl.num_programs(0)): + row64 = row.to(tl.int64) + row_ptr = GM + row64 * stride_gm + pin = -1 + if PIN_NEWEST: + pin = tl.maximum(tl.load(KE + row64) - tl.load(KS + row64) - 1, 0) // GROUP + tau = _key_to_float(_block_threshold(SGM + row64 * stride_sgm, N_SUPER, K, CHUNK, CAP)) + bound, row_max = _topk_row( + row_ptr, + N_GROUPS, + pin, + tau, + cval, + cidx, + OUT + row64 * stride_out, + GM, + 0, + K, + GROUP, + CHUNK, + CAP, + False, + ) + if WRITE_TAU: + tl.store(TAU + row64, _rescore_threshold(bound, row_max)) + tl.debug_barrier() + + +@triton.jit +def _gather_topk_kernel( + Q, + K, + K_SCALE, + W, + KS, + KE, + BLOCKS, + TAU, + SCRATCH_V, + SCRATCH_I, + OUT, + S_Q, + S_K, + N_BLOCKS, + stride_qs, + stride_qh, + stride_ks, + stride_ws, + stride_blk, + stride_out, + H: tl.constexpr, + D: tl.constexpr, + GROUP: tl.constexpr, + TOPK: tl.constexpr, + TILE: tl.constexpr, + CHUNK: tl.constexpr, + CAP: tl.constexpr, + KNOWN_TAU: tl.constexpr, + SCORE_STAGES: tl.constexpr = 1, +): + """`OUT[t]` = the `TOPK` best entries among the doc-local groups `BLOCKS[t]` (-1 unused). + + Persistent: program `p` handles queries `p, p + grid, ...` and stages one query's scores in its + own scratch rows, which stay in L2. With `KNOWN_TAU` (`BLOCKS` are `_group_topk_kernel`'s picks) + `TAU[t]` bounds the top scores from below; otherwise the bound is the `TOPK`-th best group max. + """ + pid = tl.program_id(0) + offs_h = tl.arange(0, H) + offs_d = tl.arange(0, D) + offs_t = tl.arange(0, TILE) + G_PER_TILE: tl.constexpr = TILE // GROUP + offs_tg = tl.arange(0, G_PER_TILE) + n_cols = N_BLOCKS * GROUP + K_BYTES = K.to(tl.pointer_type(tl.uint8)) + scores = SCRATCH_V + pid.to(tl.int64) * (2 * n_cols + N_BLOCKS) + cval = scores + n_cols + bmax = cval + n_cols + cidx = SCRATCH_I + pid.to(tl.int64) * n_cols + for row in range(pid, S_Q, tl.num_programs(0)): + row64 = row.to(tl.int64) + q = tl.load(Q + row64 * stride_qs + offs_h[:, None] * stride_qh + offs_d[None, :]) + w = tl.load(W + row64 * stride_ws + offs_h) + ks = tl.load(KS + row64) + ke = tl.load(KE + row64) + blk_row = BLOCKS + row64 * stride_blk + for c0 in tl.range(0, n_cols, TILE, num_stages=SCORE_STAGES): + cols = c0 + offs_t + score = _score_tile(q, w, ks, ke, blk_row, cols, n_cols, K_BYTES, K_SCALE, stride_ks, offs_d, GROUP) + tl.store(scores + cols, score, mask=cols < n_cols) + if not KNOWN_TAU: + g = c0 // GROUP + offs_tg + tl.store(bmax + g, tl.max(tl.reshape(score, [G_PER_TILE, GROUP]), axis=1), mask=g < N_BLOCKS) + tl.debug_barrier() + if KNOWN_TAU: + tau = tl.load(TAU + row64) + else: + tau = _key_to_float(_block_threshold(bmax, N_BLOCKS, TOPK, CHUNK, CAP)) + _topk_row( + scores, n_cols, -1, tau, cval, cidx, OUT + row64 * stride_out, blk_row, ks, TOPK, GROUP, CHUNK, CAP, True + ) + tl.debug_barrier() diff --git a/prime_kernels/dsa_indexer_topk/quant.py b/prime_kernels/dsa_indexer_topk/quant.py new file mode 100644 index 0000000..0d2c667 --- /dev/null +++ b/prime_kernels/dsa_indexer_topk/quant.py @@ -0,0 +1,39 @@ +"""Per-token FP8 (e4m3) quantization with power-of-two (UE8M0) scales. + +Vendored from vLLM (Apache 2.0) by way of prime-rl's `fp8_indexer`, so the indexer here scores +with exactly the quantized operands prime-rl's Triton indexer uses. +""" + +import torch +import triton +import triton.language as tl + +FP8_MAX = tl.constexpr(448.0) +FP8_EPS = tl.constexpr(1e-10) + + +@triton.jit +def _per_token_quant_fp8_kernel(x_ptr, x_q_ptr, x_s_ptr, row_stride, N: tl.constexpr, BLOCK: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + cols = tl.arange(0, BLOCK) + mask = cols < N + x = tl.load(x_ptr + row * row_stride + cols, mask=mask, other=0.0).to(tl.float32) + absmax = tl.maximum(tl.max(tl.abs(x)), FP8_EPS) + scale = tl.math.exp2(tl.ceil(tl.log2(absmax / FP8_MAX))) + x_q = tl.clamp(x / scale, -FP8_MAX, FP8_MAX).to(x_q_ptr.dtype.element_ty) + tl.store(x_q_ptr + row * N + cols, x_q, mask=mask) + tl.store(x_s_ptr + row, scale) + + +def per_token_quant_fp8(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """`(..., N)` -> contiguous `(..., N)` float8_e4m3fn and `(...)` fp32 scales, one per row.""" + assert x.stride(-1) == 1 + rows = x.reshape(-1, x.shape[-1]) + x_q = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn) + x_s = torch.empty(x.shape[:-1], device=x.device, dtype=torch.float32) + if rows.shape[0]: + block = triton.next_power_of_2(x.shape[-1]) + _per_token_quant_fp8_kernel[(rows.shape[0],)]( + rows, x_q, x_s, rows.stride(0), N=x.shape[-1], BLOCK=block, num_warps=min(max(block // 256, 1), 8) + ) + return x_q, x_s diff --git a/prime_kernels/dsa_sparse_attn_bwd/__init__.py b/prime_kernels/dsa_sparse_attn_bwd/__init__.py new file mode 100644 index 0000000..a98fa3f --- /dev/null +++ b/prime_kernels/dsa_sparse_attn_bwd/__init__.py @@ -0,0 +1,7 @@ +from prime_kernels.dsa_sparse_attn_bwd.interface import ( + dsa_sparse_attn_backward, + sparse_attn_backward_flat, + unsupported_shape_reason, +) + +__all__ = ["dsa_sparse_attn_backward", "sparse_attn_backward_flat", "unsupported_shape_reason"] diff --git a/prime_kernels/dsa_sparse_attn_bwd/_cute_utils.py b/prime_kernels/dsa_sparse_attn_bwd/_cute_utils.py new file mode 100644 index 0000000..2c655d9 --- /dev/null +++ b/prime_kernels/dsa_sparse_attn_bwd/_cute_utils.py @@ -0,0 +1,148 @@ +"""Small CuTe DSL helpers for the SM90 sparse attention backward.""" + +import cutlass +import cutlass.cute as cute +import cutlass.utils.hopper_helpers as sm90_utils +from cutlass import Float32 +from cutlass._mlir.dialects import arith, llvm +from cutlass.cute.nvgpu import warpgroup +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass.utils import LayoutEnum + + +def make_smem_layout(dtype, shape, stage=None): + """Row-major (K-major) smem tile with the widest swizzle its row allows, optionally staged.""" + atom = warpgroup.make_smem_layout_atom( + sm90_utils.get_smem_layout_atom(LayoutEnum.ROW_MAJOR, dtype, shape[1]), dtype + ) + if stage is None: + return cute.tile_to_shape(atom, shape, order=(0, 1)) + return cute.tile_to_shape(atom, (*shape, stage), order=(0, 1, 2)) + + +def make_tiled_mma(a_major, b_major, tiler_n, a_in_regs=False): + mode = {"K": cute.nvgpu.OperandMajorMode.K, "MN": cute.nvgpu.OperandMajorMode.MN} + return sm90_utils.make_trivial_tiled_mma( + cutlass.BFloat16, + cutlass.BFloat16, + mode[a_major], + mode[b_major], + Float32, + atom_layout_mnk=(1, 1, 1), + tiler_mn=(64, tiler_n), + a_source=warpgroup.OperandSource.RMEM if a_in_regs else warpgroup.OperandSource.SMEM, + ) + + +def transpose_view(a: cute.Tensor) -> cute.Tensor: + """Swap the first two modes of a smem tensor.""" + shape = (a.shape[1], a.shape[0], *a.shape[2:]) + order = (1, 0, *range(2, cute.rank(a))) + return cute.composition(a, cute.make_ordered_layout(shape, order=order)) + + +@cute.jit +def gemm( + tiled_mma: cute.TiledMma, acc: cute.Tensor, tA: cute.Tensor, tB: cute.Tensor, zero_init: cutlass.Constexpr[bool] +): + """Issue one warpgroup MMA over every k block of `tA`/`tB` and commit it (does not wait).""" + warpgroup.fence() + atom = cute.make_mma_atom(tiled_mma.op) + atom.set(warpgroup.Field.ACCUMULATE, not zero_init) + for k in cutlass.range_constexpr(cute.size(tA.shape[2])): + cute.gemm(atom, acc, tA[None, None, k], tB[None, None, k], acc) + atom.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.commit_group() + + +def acc_mn_view(acc: cute.Tensor) -> cute.Tensor: + """View an SM90 accumulator ((2, 2, V), M, N) as ((2, M), (2, V, N)): rows x columns.""" + layout = cute.make_layout(acc.layout.shape) + shape = ((layout.shape[0][1], layout.shape[1]), (layout.shape[0][0], layout.shape[0][2], layout.shape[2])) + stride = ((layout.stride[0][1], layout.stride[1]), (layout.stride[0][0], layout.stride[0][2], layout.stride[2])) + return cute.make_tensor(acc.iterator, cute.composition(acc.layout, cute.make_layout(shape, stride=stride))) + + +def acc_as_operand_a(acc_layout: cute.Layout) -> cute.Layout: + """Reinterpret a 16-bit copy of an accumulator ((2, 2, N/8), M, N') as the A fragment of the next MMA.""" + div = cute.logical_divide(acc_layout, ((None, None, 2), None, None)) + return cute.make_layout( + ((div.shape[0][0], div.shape[0][1], div.shape[0][2][0]), div.shape[1], (div.shape[0][2][1], div.shape[2])), + stride=( + (div.stride[0][0], div.stride[0][1], div.stride[0][2][0]), + div.stride[1], + (div.stride[0][2][1], div.stride[2]), + ), + ) + + +@dsl_user_op +def cvt_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> cutlass.Int32: + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@cute.jit +def to_bf16(src: cute.Tensor) -> cute.Tensor: + dst = cute.make_rmem_tensor(src.shape, cutlass.BFloat16) + dst_i32 = cute.recast_tensor(dst, cutlass.Int32) + for i in cutlass.range_constexpr(cute.size(dst_i32)): + dst_i32[i] = cvt_bf16x2(src[2 * i], src[2 * i + 1]) + return dst + + +@dsl_user_op +def select_f32(pred, a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + return Float32( + arith.select( + cutlass.Boolean(pred).ir_value(loc=loc, ip=ip), + Float32(a).ir_value(loc=loc, ip=ip), + Float32(b).ir_value(loc=loc, ip=ip), + ) + ) + + +@cute.jit +def gemm2( + tiled_mma: cute.TiledMma, acc: cute.Tensor, tA0: cute.Tensor, tB0: cute.Tensor, tA1: cute.Tensor, tB1: cute.Tensor +): + """acc = A0 B0 + A1 B1 as one committed warpgroup MMA group (does not wait).""" + warpgroup.fence() + atom = cute.make_mma_atom(tiled_mma.op) + atom.set(warpgroup.Field.ACCUMULATE, False) + for k in cutlass.range_constexpr(cute.size(tA0.shape[2])): + cute.gemm(atom, acc, tA0[None, None, k], tB0[None, None, k], acc) + atom.set(warpgroup.Field.ACCUMULATE, True) + for k in cutlass.range_constexpr(cute.size(tA1.shape[2])): + cute.gemm(atom, acc, tA1[None, None, k], tB1[None, None, k], acc) + warpgroup.commit_group() + + +@dsl_user_op +def red_add_v4_if(pred, ptr_i64, a: Float32, b: Float32, c: Float32, d: Float32, *, loc=None, ip=None) -> None: + """If `pred`: fire-and-forget float4 atomic add to a 16-byte aligned global address (no branch).""" + llvm.inline_asm( + None, + [ + cutlass.Int32(pred).ir_value(loc=loc, ip=ip), + cutlass.Int64(ptr_i64).ir_value(loc=loc, ip=ip), + Float32(a).ir_value(loc=loc, ip=ip), + Float32(b).ir_value(loc=loc, ip=ip), + Float32(c).ir_value(loc=loc, ip=ip), + Float32(d).ir_value(loc=loc, ip=ip), + ], + "{\n.reg .pred p;\nsetp.ne.b32 p, $0, 0;\n@p red.relaxed.gpu.global.add.v4.f32 [$1], {$2, $3, $4, $5};\n}", + "r,l,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) diff --git a/prime_kernels/dsa_sparse_attn_bwd/interface.py b/prime_kernels/dsa_sparse_attn_bwd/interface.py new file mode 100644 index 0000000..210dcb0 --- /dev/null +++ b/prime_kernels/dsa_sparse_attn_bwd/interface.py @@ -0,0 +1,231 @@ +"""Host side of the SM90 sparse attention backward: Triton pre/post-processing around the CuTe DSL +main kernel, exposed as a `torch.library` custom op with the contract of prime-rl's +`dsv41_sparse_attn_backward` (and cuDNN's `flash_attn_bwd_sm90`).""" + +import torch +import triton +import triton.language as tl + +HEADS = 64 +DIM = 512 +TILE_N = 64 + +_compiled: dict = {} + + +@triton.jit +def _preprocess_kernel( + out_ptr, + out_tail_ptr, + dout_ptr, + lse_ptr, + sink_ptr, + idx_ptr, + delta_ptr, + l2_ptr, + dsink_ptr, + idx_out_ptr, + count_ptr, + H: tl.constexpr, + D: tl.constexpr, + BLOCK_D: tl.constexpr, + K: tl.constexpr, + BLOCK_K: tl.constexpr, + TAIL: tl.constexpr, +): + """Per query: delta = rowsum(dO * O), the sink-aware LSE in log2 units, each head's sink gradient + -p_sink * delta, and the slot list with its non-empty slots moved to the front (then -1s) plus + their count. With `TAIL` > 0, each head's last `TAIL` channels of O are read from `out_tail`.""" + t = tl.program_id(0).to(tl.int64) + heads = tl.arange(0, H) + acc = tl.zeros([H], dtype=tl.float32) + for d0 in tl.static_range(0, D, BLOCK_D): + cols = d0 + tl.arange(0, BLOCK_D)[None, :] + offs = t * H * D + heads[:, None] * D + cols + if TAIL > 0 and d0 + BLOCK_D > D - TAIL: + in_tail = cols >= D - TAIL + tail_offs = t * H * TAIL + heads[:, None] * TAIL + cols - (D - TAIL) + tail = tl.load(out_tail_ptr + tail_offs, mask=in_tail) + o = tl.where(in_tail, tail, tl.load(out_ptr + offs, mask=~in_tail)).to(tl.float32) + else: + o = tl.load(out_ptr + offs).to(tl.float32) + g = tl.load(dout_ptr + offs).to(tl.float32) + acc += tl.sum(o * g, axis=1) + lse = tl.load(lse_ptr + t * H + heads) + sink = tl.load(sink_ptr + heads) + hi = tl.maximum(lse, sink) + lse_sink = hi + tl.log(tl.exp(lse - hi) + tl.exp(sink - hi)) + tl.store(delta_ptr + t * H + heads, acc) + tl.store(l2_ptr + t * H + heads, lse_sink * 1.4426950408889634) + tl.store(dsink_ptr + t * H + heads, -tl.exp(sink - lse_sink) * acc) + + slots = tl.arange(0, BLOCK_K) + in_row = slots < K + idx = tl.load(idx_ptr + t * K + slots, mask=in_row, other=-1) + valid = (idx >= 0).to(tl.int32) + empty = (in_row & (idx < 0)).to(tl.int32) + count = tl.sum(valid, axis=0) + pos = tl.where(valid == 1, tl.cumsum(valid, axis=0) - valid, count + tl.cumsum(empty, axis=0) - empty) + tl.store(idx_out_ptr + t * K + pos, tl.where(valid == 1, idx, -1), mask=in_row) + tl.store(count_ptr + t, count) + + +@triton.jit +def _postprocess_kernel(acc_ptr, dkv_ptr, n_rows, D: tl.constexpr, BLOCK_ROWS: tl.constexpr): + """Undo the scatter's column permutation (see `SparseAttnBwdSm90.scatter`) and cast to bf16.""" + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + cols = tl.arange(0, D) + w = cols % 16 + fake = (cols - w) + 4 * ((w % 8) // 2) + (w % 2) + 2 * (w // 8) + mask = rows[:, None] < n_rows + rows64 = rows.to(tl.int64)[:, None] + vals = tl.load(acc_ptr + rows64 * D + fake[None, :], mask=mask) + tl.store(dkv_ptr + rows64 * D + cols[None, :], vals.to(tl.bfloat16), mask=mask) + + +def _main_kernel(num_slots: int, device: torch.device): + key = (num_slots, device.index) + if key not in _compiled: + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + + from prime_kernels.dsa_sparse_attn_bwd.kernel import SparseAttnBwdSm90 + + def fake(shape, dtype): + t = torch.empty(shape, dtype=dtype, device=device) + return from_dlpack(t, assumed_align=16, enable_tvm_ffi=True).mark_layout_dynamic(leading_dim=t.ndim - 1) + + bf16, f32, i32 = torch.bfloat16, torch.float32, torch.int32 + args = ( + fake((2, HEADS, DIM), bf16), # q + fake((3, DIM), bf16), # kv + fake((2, HEADS, DIM), bf16), # grad_out + fake((2, HEADS), f32), # log2 LSE with sink + fake((2, HEADS), f32), # delta + fake((2, num_slots), i32), # compacted slots + fake((2,), i32), # slot counts + fake((2, HEADS, DIM), bf16), # dq + fake((3, DIM), f32), # dkv accumulator + ) + stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + _compiled[key] = cute.compile( + SparseAttnBwdSm90(num_slots), + *args, + cutlass.Float32(1.0), + cutlass.Int32(1), + stream, + options="--enable-tvm-ffi", + ) + return _compiled[key] + + +def unsupported_shape_reason(num_heads: int, head_dim: int) -> str | None: + if num_heads != HEADS or head_dim != DIM: + return f"dsa_sparse_attn_bwd supports {HEADS} heads of a {DIM}-wide latent, got {num_heads} x {head_dim}" + return None + + +def sparse_attn_backward_flat( + q: torch.Tensor, # (t, h, d) bf16 + kv: torch.Tensor, # (n, d) bf16, shared K = V latent rows + out: torch.Tensor, # (t, h, d) bf16, forward output (sink applied) + grad_out: torch.Tensor, # (t, h, d) bf16 + lse: torch.Tensor, # (t, h) f32, sink-free natural-log LSE (as FlashMLA returns it) + indices: torch.Tensor, # (t, k) int32 rows of kv, -1 = empty slot + sinks: torch.Tensor, # (h,) f32 + sm_scale: float, + out_tail: torch.Tensor | None = None, # (t, h, tail) bf16, each head's last channels of `out` +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """dq (t, h, d) bf16, dkv (n, d) bf16, dsinks (h,) f32 of sink-aware sparse attention. + + With `out_tail`, the forward output is `out` with each head's last `tail` channels replaced by + `out_tail`'s (only the backward's `delta` reads it), so a caller that rewrote those channels in + place keeps just them.""" + import cuda.bindings.driver as cuda + + t, h, d = q.shape + n = kv.shape[0] + reason = unsupported_shape_reason(h, d) + if reason is not None: + raise ValueError(reason) + if indices.shape[1] % TILE_N: + indices = torch.nn.functional.pad(indices, (0, TILE_N - indices.shape[1] % TILE_N), value=-1) + q, kv, out, grad_out, indices = (x.contiguous() for x in (q, kv, out, grad_out, indices)) + lse = lse.float().contiguous() + sinks = sinks.float().contiguous() + k = indices.shape[1] + + delta = torch.empty(t, h, dtype=torch.float32, device=q.device) + lse2 = torch.empty_like(delta) + dsink_part = torch.empty_like(delta) + slots = torch.empty_like(indices) + count = torch.empty(t, dtype=torch.int32, device=q.device) + dq = torch.empty_like(q) + dkv_acc = torch.zeros(n, d, dtype=torch.float32, device=q.device) + dkv = torch.empty(n, d, dtype=kv.dtype, device=q.device) + if t == 0: + return dq, dkv.zero_(), torch.zeros_like(sinks) + + tail = 0 if out_tail is None else out_tail.shape[-1] + _preprocess_kernel[(t,)]( + out, + out if out_tail is None else out_tail.contiguous(), + grad_out, + lse, + sinks, + indices, + delta, + lse2, + dsink_part, + slots, + count, + H=h, + D=d, + BLOCK_D=128, + K=k, + BLOCK_K=triton.next_power_of_2(k), + TAIL=tail, + ) + num_ctas = min(t, torch.cuda.get_device_properties(q.device).multi_processor_count) + stream = cuda.CUstream(torch.cuda.current_stream(q.device).cuda_stream) + _main_kernel(k, q.device)(q, kv, grad_out, lse2, delta, slots, count, dq, dkv_acc, sm_scale, num_ctas, stream) + _postprocess_kernel[(triton.cdiv(n, 16),)](dkv_acc, dkv, n, D=d, BLOCK_ROWS=16) + return dq, dkv, dsink_part.sum(0) + + +@torch.library.custom_op("prime_kernels::dsa_sparse_attn_bwd", mutates_args=()) +def dsa_sparse_attn_backward( + grad_out: torch.Tensor, + q: torch.Tensor, + kv: torch.Tensor, + out: torch.Tensor, + lse: torch.Tensor, + indices: torch.Tensor, + sinks: torch.Tensor, + sm_scale: float, + out_tail: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Drop-in for prime-rl's `dsv41_sparse_attn_backward`: `q`, `out`, `grad_out` (1, t, h, d), `kv` + (1, n, 1, d), `lse` (1, t, h), `indices` (1, t, 1, k) int32 (-1 = empty), `sinks` (h,) -> + dq like q, dkv like kv, dsinks like sinks. `out_tail` (1, t, h, tail), if given, holds each head's + last `tail` channels of the forward output in place of `out`'s (see `sparse_attn_backward_flat`).""" + _, t, h, d = q.shape + dq, dkv, dsinks = sparse_attn_backward_flat( + q.view(t, h, d), + kv.view(-1, d), + out.view(t, h, d), + grad_out.contiguous().view(t, h, d), + lse.view(t, h), + indices.reshape(t, -1), + sinks, + sm_scale, + None if out_tail is None else out_tail.view(t, h, -1), + ) + return dq.view_as(q), dkv.view_as(kv).to(kv.dtype), dsinks.to(sinks.dtype) + + +@dsa_sparse_attn_backward.register_fake +def _dsa_sparse_attn_backward_fake(grad_out, q, kv, out, lse, indices, sinks, sm_scale, out_tail=None): + return torch.empty_like(q), torch.empty_like(kv), torch.empty_like(sinks) diff --git a/prime_kernels/dsa_sparse_attn_bwd/kernel.py b/prime_kernels/dsa_sparse_attn_bwd/kernel.py new file mode 100644 index 0000000..e65e5db --- /dev/null +++ b/prime_kernels/dsa_sparse_attn_bwd/kernel.py @@ -0,0 +1,734 @@ +"""SM90 (Hopper) backward of DeepSeek-V4.1 sparse attention, in CuTe DSL. + +A persistent CTA per SM walks queries q = blockIdx.x, + gridDim.x, ...; for each query its 64 +heads form the M = 64 side of every MMA, and its non-empty slots (compacted to the front of the +list by the preprocess) are walked in tiles of 64 gathered latent rows (K = V, 512 wide): + + S = Q K^T, dP = dO K^T (64 heads x 64 slots, K = 512) + P = exp2(S * scale_log2 - L2), dS = P (dP - delta) * scale + dQ += dS K (64 heads x 512, K = 64 slots) + dKV[slots] += P^T dO + dS^T Q (64 slots x 512, K = 64 heads), scattered with + fp32 vector atomics into a column-permuted buffer + +Q and dO (64 x 512 each) stay resident; with them, one KV tile and P / dS the CTA uses ~215 KB of +shared memory, so the KV tile is single-buffered and split into two 256-column halves, each +refilled as soon as its last reader is done: + + WG0: S, dP (lo half then hi half) -> softmax -> P, dS to smem -> dQ[:, :256] += dS K_lo (dS from + registers) -> gather K_lo of the next tile -> the first `wg0_chunks` dKV chunks + WG1: P^T, dS^T to registers (ldmatrix.trans) -> dQ[:, 256:] += dS K_hi -> gather K_hi of the next + tile -> the remaining dKV chunks, their A operands read from registers + (64 columns each; rotating accumulators, so a chunk's atomics overlap the next MMAs) + +The next tile may belong to the next query: its slot list is prefetched into a second buffer +during the current query's first tile, and the halves of its Q / dO are loaded by TMA as soon as +the current query's last dKV chunk using them is done, so query switches do not drain the SM. +The dKV atomics (2 KB per query-slot, ~21 GB at V4.1 shapes) bound this kernel: on their own +they take ~6 ms at the ~3.5 TB/s an H200's L2 sustains for fp32 vector reductions. +""" + +import math + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr +from cutlass.cute.nvgpu import cpasync, warpgroup + +from prime_kernels.dsa_sparse_attn_bwd import _cute_utils as cu + +HEADS = 64 +DIM = 512 +HALF = 256 +TILE_N = 64 +CHUNK = 64 # dKV is computed and scattered 64 latent columns at a time +NUM_CHUNKS = DIM // CHUNK + +# named barriers (0 is __syncthreads) +BAR_PDS_FULL = 1 # WG0 -> WG1: P / dS of the current tile are in smem +BAR_PDS_FREE = 2 # WG1 -> WG0: WG1 is done reading P / dS +BAR_QDO_FREE = 3 # WG0 -> WG1: WG0's dKV chunks of a query's last tile are done (Q / dO reusable) + + +class SparseAttnBwdSm90: + def __init__(self, num_slots: int, wg0_chunks: int = 2, rs: int = 1): + assert num_slots % TILE_N == 0 + assert 0 <= wg0_chunks <= NUM_CHUNKS // 2 + self.wg0_chunks = wg0_chunks + # rs: WG1 reads P^T / dS^T for its dKV MMAs from registers, loaded once per tile, instead of from + # shared memory once per chunk (shared memory bandwidth is what the SS MMAs run at) + self.rs = rs + # dKV accumulators in rotation; 3 give a chunk's atomics two MMAs' time to read their registers, + # rs needs the 32 registers of the third for P^T / dS^T (and gains nothing from it) + self.num_acc = 2 if rs else 3 + self.num_slots = num_slots + self.num_threads = 256 + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, # (t, h, d) bf16 + mKV: cute.Tensor, # (n, d) bf16 + mdO: cute.Tensor, # (t, h, d) bf16 + mL2: cute.Tensor, # (t, h) f32, log2-domain sink-aware LSE + mDelta: cute.Tensor, # (t, h) f32 + mIdx: cute.Tensor, # (t, k) int32, non-empty slots first, then -1 + mCount: cute.Tensor, # (t,) int32, non-empty slots per query + mdQ: cute.Tensor, # (t, h, d) bf16 + mdKV: cute.Tensor, # (n, d) f32, column-permuted accumulator + scale: Float32, + num_ctas: Int32, + stream, + ): + def assume_aligned(t): + divby = 128 // t.element_type.width + strides = [s if isinstance(s, int) else cute.assume(s, divby=divby) for s in t.stride[:-1]] + return cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=(*strides, t.stride[-1]))) + + mQ, mKV, mdO, mdQ, mdKV, mIdx = [assume_aligned(t) for t in (mQ, mKV, mdO, mdQ, mdKV, mIdx)] + bf16 = cutlass.BFloat16 + sQ_layout = cu.make_smem_layout(bf16, (HEADS, HALF), 2) + sK_layout = cu.make_smem_layout(bf16, (TILE_N, HALF), 2) + sPdS_layout = cu.make_smem_layout(bf16, (HEADS, TILE_N)) + chunk_layout = cu.make_smem_layout(bf16, (HEADS, CHUNK)) + + # (h, d, t) views so a (64, 256) TMA box is one query's half + mQ_t, mdO_t = [cute.make_tensor(m.iterator, cute.select(m.layout, mode=[1, 2, 0])) for m in (mQ, mdO)] + half_layout = cute.select(sQ_layout, mode=[0, 1]) + tma_Q, tma_mQ = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), mQ_t, half_layout, (HEADS, HALF)) + tma_dO, tma_mdO = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), mdO_t, half_layout, (HEADS, HALF) + ) + self.half_bytes = HEADS * HALF * 2 + + mma_SdP = cu.make_tiled_mma("K", "K", TILE_N) + mma_dQ_rs = cu.make_tiled_mma("K", "MN", HALF, a_in_regs=True) + mma_dQ_ss = cu.make_tiled_mma("K", "MN", HALF) + mma_dKV = cu.make_tiled_mma("MN", "MN", CHUNK) + mma_dKV_rs = cu.make_tiled_mma("K", "MN", CHUNK, a_in_regs=True) + + @cute.struct + class SharedStorage: + # Q/dO halves (TMA), K halves (cp.async), slot list (cp.async) + mbar: cute.struct.MemRange[cutlass.Int64, 6] + sIdx: cute.struct.Align[cute.struct.MemRange[Int32, 2 * self.num_slots], 128] + sQ: cute.struct.Align[cute.struct.MemRange[bf16, cute.cosize(sQ_layout)], 1024] + sdO: cute.struct.Align[cute.struct.MemRange[bf16, cute.cosize(sQ_layout)], 1024] + sK: cute.struct.Align[cute.struct.MemRange[bf16, cute.cosize(sK_layout)], 1024] + sP: cute.struct.Align[cute.struct.MemRange[bf16, cute.cosize(sPdS_layout)], 1024] + sdS: cute.struct.Align[cute.struct.MemRange[bf16, cute.cosize(sPdS_layout)], 1024] + + self.shared_storage = SharedStorage + scale_log2 = scale * math.log2(math.e) + self.kernel( + tma_mQ, + tma_mdO, + tma_Q, + tma_dO, + mKV, + mL2, + mDelta, + mIdx, + mCount, + mdQ, + mdKV, + mma_SdP, + mma_dQ_rs, + mma_dQ_ss, + mma_dKV, + mma_dKV_rs, + sQ_layout, + sK_layout, + sPdS_layout, + chunk_layout, + scale, + scale_log2, + ).launch( + grid=[num_ctas, 1, 1], + block=[self.num_threads, 1, 1], + smem=SharedStorage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mdO: cute.Tensor, + tma_Q: cute.CopyAtom, + tma_dO: cute.CopyAtom, + mKV: cute.Tensor, + mL2: cute.Tensor, + mDelta: cute.Tensor, + mIdx: cute.Tensor, + mCount: cute.Tensor, + mdQ: cute.Tensor, + mdKV: cute.Tensor, + mma_SdP: cute.TiledMma, + mma_dQ_rs: cute.TiledMma, + mma_dQ_ss: cute.TiledMma, + mma_dKV: cute.TiledMma, + mma_dKV_rs: cute.TiledMma, + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + sPdS_layout: cute.ComposedLayout, + chunk_layout: cute.ComposedLayout, + scale: Float32, + scale_log2: Float32, + ): + tidx, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + grid, _, _ = cute.arch.grid_dim() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + wg = cute.arch.make_warp_uniform(tidx // 128) + wg_tidx = tidx % 128 + num_queries = mCount.shape[0] + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + mbar0 = storage.mbar.data_ptr() + mbar_QdO = (mbar0, mbar0 + 1) + mbar_K = (mbar0 + 2, mbar0 + 3) + mbar_idx = mbar0 + 4 + sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner) + sdO = storage.sdO.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner) + sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) + sP = storage.sP.get_tensor(sPdS_layout.outer, swizzle=sPdS_layout.inner) + sdS = storage.sdS.get_tensor(sPdS_layout.outer, swizzle=sPdS_layout.inner) + sIdx = storage.sIdx.get_tensor(cute.make_layout((self.num_slots, 2))) + + if warp_idx == 0: + cpasync.prefetch_descriptor(tma_Q) + cpasync.prefetch_descriptor(tma_dO) + cute.arch.mbarrier_init(mbar_QdO[0], 1) + cute.arch.mbarrier_init(mbar_QdO[1], 1) + cute.arch.mbarrier_init(mbar_K[0], 128) + cute.arch.mbarrier_init(mbar_K[1], 128) + cute.arch.mbarrier_init(mbar_idx, 128) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + + # first query with work: slot list (WG0), Q / dO (WG1), first KV tile (both halves) + q0 = self.next_query(mCount, bid - grid, grid, num_queries) + if q0 < num_queries: + if wg == 0: + self.load_slots(mIdx, sIdx, q0, 0, wg_tidx, mbar_idx) + elif warp_idx == 4: + for h in cutlass.range_constexpr(2): + self.load_QdO_half(mQ, mdO, sQ, sdO, tma_Q, tma_dO, q0, h, mbar_QdO[h]) + cute.arch.mbarrier_wait(mbar_idx, 0) + self.gather_half(mKV, sK, sIdx, 0, 0, wg, wg_tidx, mbar0 + 2 + wg) + + if wg == 0: + self.wg0_loop( + mKV, + mL2, + mDelta, + mIdx, + mCount, + mdQ, + mdKV, + sQ, + sdO, + sK, + sP, + sdS, + sIdx, + mbar_QdO, + mbar_K, + mbar_idx, + mma_SdP, + mma_dQ_rs, + mma_dKV, + chunk_layout, + scale, + scale_log2, + bid, + grid, + num_queries, + wg_tidx, + ) + else: + self.wg1_loop( + mQ, + mdO, + tma_Q, + tma_dO, + mKV, + mCount, + mdQ, + mdKV, + sQ, + sdO, + sK, + sP, + sdS, + sIdx, + mbar_QdO, + mbar_K, + mbar_idx, + mma_dQ_ss, + mma_dKV_rs if const_expr(self.rs) else mma_dKV, + chunk_layout, + bid, + grid, + num_queries, + wg_tidx, + warp_idx, + ) + + @cute.jit + def next_query(self, mCount: cute.Tensor, q: Int32, grid: Int32, num_queries) -> Int32: + """The CTA's next query after `q` that has at least one non-empty slot (num_queries if none).""" + qn = q + grid + while qn < num_queries and mCount[qn] == 0: + qn += grid + return cutlass.min(qn, num_queries) + + @cute.jit + def load_slots(self, mIdx: cute.Tensor, sIdx: cute.Tensor, q: Int32, buf: Int32, wg_tidx: Int32, mbar): + """cp.async query `q`'s slot list into buffer `buf`; arrives on `mbar` (128 arrivals).""" + cp_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), Int32, num_bits_per_copy=128 + ) + cp_thr = cute.make_tiled_copy_tv(cp_atom, cute.make_layout((1,)), cute.make_layout((4,))).get_slice(0) + num_vec = self.num_slots // 4 + g_vecs = cute.flat_divide(mIdx[q, None], (4,)) + s_vecs = cute.flat_divide(sIdx[None, buf], (4,)) + for i in cutlass.range_constexpr((num_vec + 127) // 128): + v = i * 128 + wg_tidx + if v < num_vec: + cute.copy(cp_atom, cp_thr.partition_S(g_vecs[None, v]), cp_thr.partition_D(s_vecs[None, v])) + cute.arch.cp_async_mbarrier_arrive_noinc(mbar) + + @cute.jit + def load_QdO_half(self, mQ, mdO, sQ, sdO, tma_Q, tma_dO, q: Int32, h: cutlass.Constexpr[int], mbar): + """One thread: TMA half `h` (256 columns) of query `q`'s Q and dO into smem.""" + gQ = cute.local_tile(mQ[None, None, q], (HEADS, HALF), (0, None)) + gdO = cute.local_tile(mdO[None, None, q], (HEADS, HALF), (0, None)) + tQs, tQg = cpasync.tma_partition( + tma_Q, 0, cute.make_layout(1), cute.group_modes(sQ, 0, 2), cute.group_modes(gQ, 0, 2) + ) + tdOs, tdOg = cpasync.tma_partition( + tma_dO, 0, cute.make_layout(1), cute.group_modes(sdO, 0, 2), cute.group_modes(gdO, 0, 2) + ) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar, 2 * self.half_bytes) + cute.copy(tma_Q, tQg[None, h], tQs[None, h], tma_bar_ptr=mbar) + cute.copy(tma_dO, tdOg[None, h], tdOs[None, h], tma_bar_ptr=mbar) + + @cute.jit + def gather_half( + self, + mKV: cute.Tensor, + sK: cute.Tensor, + sIdx: cute.Tensor, + tile: Int32, + buf: Int32, + half: Int32, + wg_tidx: Int32, + mbar, + ): + """cp.async the 256-column half `half` of the tile's 64 rows (empty slots read row 0, masked + later); completion arrives on `mbar` (one arrival per thread of the warpgroup).""" + cp_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), cutlass.BFloat16, num_bits_per_copy=128 + ) + cp_thr = cute.make_tiled_copy_tv(cp_atom, cute.make_layout((1,)), cute.make_layout((8,))).get_slice(0) + # 32 threads cover one row's 512-byte half: thread -> 16-byte chunk, rows strided by 4 + col = wg_tidx % 32 + for i in cutlass.range_constexpr(TILE_N // 4): + row = wg_tidx // 32 + 4 * i + kv_row = cutlass.max(sIdx[tile * TILE_N + row, buf], Int32(0)) + g_chunk = cute.flat_divide(mKV[kv_row, None], (8,))[None, half * 32 + col] + s_chunk = cute.flat_divide(sK[row, None, half], (8,))[None, col] + cute.copy(cp_atom, cp_thr.partition_S(g_chunk), cp_thr.partition_D(s_chunk)) + cute.arch.cp_async_mbarrier_arrive_noinc(mbar) + + @cute.jit + def gather_next( + self, + mKV, + sK, + sIdx, + tile, + num_tiles, + buf, + q_next, + num_queries, + nq, + half: cutlass.Constexpr[int], + wg_tidx, + mbar_K, + mbar_idx, + ): + """Gather this warpgroup's half of the tile after (query, tile): the next tile of this query, + or the first tile of the next query once its slot list has landed.""" + if tile + 1 < num_tiles: + self.gather_half(mKV, sK, sIdx, tile + 1, buf, half, wg_tidx, mbar_K) + elif q_next < num_queries: + cute.arch.mbarrier_wait(mbar_idx, (nq + 1) % 2) + self.gather_half(mKV, sK, sIdx, 0, 1 - buf, half, wg_tidx, mbar_K) + + @cute.jit + def wg0_loop( + self, + mKV, + mL2, + mDelta, + mIdx, + mCount, + mdQ, + mdKV, + sQ, + sdO, + sK, + sP, + sdS, + sIdx, + mbar_QdO, + mbar_K, + mbar_idx, + mma_SdP, + mma_dQ_rs, + mma_dKV, + chunk_layout, + scale, + scale_log2, + bid, + grid, + num_queries, + wg_tidx, + ): + dkv = self.dkv_setup(mma_dKV, sP, sdS, wg_tidx) + thr_SdP = mma_SdP.get_slice(wg_tidx) + tSrQ = thr_SdP.make_fragment_A(thr_SdP.partition_A(sQ)) + tSrdO = thr_SdP.make_fragment_A(thr_SdP.partition_A(sdO)) + tSrK = thr_SdP.make_fragment_B(thr_SdP.partition_B(sK)) + tScS = cu.acc_mn_view(thr_SdP.partition_C(cute.make_identity_tensor((HEADS, TILE_N)))) + + thr_dQ = mma_dQ_rs.get_slice(wg_tidx) + tdQrKt = thr_dQ.make_fragment_B(thr_dQ.partition_B(cu.transpose_view(sK))) + acc_dQ = cute.make_rmem_tensor(mma_dQ_rs.partition_shape_C((HEADS, HALF)), Float32) + + smem_store = cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=False, num_matrices=4), cutlass.BFloat16 + ) + thr_store = cute.make_tiled_copy_C(smem_store, mma_SdP).get_slice(wg_tidx) + tPsP = thr_store.partition_D(sP) + tPsdS = thr_store.partition_D(sdS) + row_l2 = cute.make_rmem_tensor((2,), Float32) + row_delta = cute.make_rmem_tensor((2,), Float32) + + it = Int32(0) # tiles done by this CTA (phase of the K barriers) + nq = Int32(0) # queries with work done by this CTA (phase of the Q/dO and slot barriers) + q = bid + while q < num_queries: + num_tiles = (mCount[q] + TILE_N - 1) // TILE_N + if num_tiles == 0: + acc_dQ.fill(0.0) + self.store_dQ(acc_dQ, thr_dQ, mdQ, q, 0) + q += grid + else: + buf = nq % 2 + q_next = self.next_query(mCount, q, grid, num_queries) + for r in cutlass.range_constexpr(2): + row_l2[r] = mL2[q, tScS[r, 0][0]] + row_delta[r] = mDelta[q, tScS[r, 0][0]] + acc_dQ.fill(0.0) + for tile in cutlass.range(num_tiles, unroll=1): + acc_S = cute.make_rmem_tensor(mma_SdP.partition_shape_C((HEADS, TILE_N)), Float32) + acc_dP = cute.make_rmem_tensor(mma_SdP.partition_shape_C((HEADS, TILE_N)), Float32) + for h in cutlass.range_constexpr(2): + cute.arch.mbarrier_wait(mbar_K[h], it % 2) + if tile == 0: + cute.arch.mbarrier_wait(mbar_QdO[h], nq % 2) + cute.arch.fence_view_async_shared() + cu.gemm( + mma_SdP, acc_S, tSrQ[None, None, None, h], tSrK[None, None, None, h], zero_init=(h == 0) + ) + cu.gemm( + mma_SdP, acc_dP, tSrdO[None, None, None, h], tSrK[None, None, None, h], zero_init=(h == 0) + ) + warpgroup.wait_group(0) + + S = cu.acc_mn_view(acc_S) + dP = cu.acc_mn_view(acc_dP) + for c in cutlass.range_constexpr(cute.size(S, mode=[1])): + valid = sIdx[tile * TILE_N + tScS[0, c][1], buf] >= 0 + for r in cutlass.range_constexpr(cute.size(S, mode=[0])): + p = cute.math.exp2(S[r, c] * scale_log2 - row_l2[r], fastmath=True) + p = cu.select_f32(valid, p, Float32(0.0)) + S[r, c] = p + dP[r, c] = p * (dP[r, c] - row_delta[r]) * scale + rP = cu.to_bf16(acc_S) + rdS = cu.to_bf16(acc_dP) + + if it > 0: + cute.arch.barrier(barrier_id=BAR_PDS_FREE, number_of_threads=256) + if tile == 0 and q_next < num_queries: + # WG1 is past the previous query: its slot buffer can take the next query's list + self.load_slots(mIdx, sIdx, q_next, 1 - buf, wg_tidx, mbar_idx) + cute.copy(thr_store, thr_store.retile(rP), tPsP) + cute.copy(thr_store, thr_store.retile(rdS), tPsdS) + cute.arch.fence_view_async_shared() + cute.arch.barrier_arrive(barrier_id=BAR_PDS_FULL, number_of_threads=256) + + tdQrdS = cute.make_tensor(rdS.iterator, cu.acc_as_operand_a(rdS.layout)) + cu.gemm(mma_dQ_rs, acc_dQ, tdQrdS, tdQrKt[None, None, None, 0], zero_init=False) + warpgroup.wait_group(0) + self.gather_next( + mKV, sK, sIdx, tile, num_tiles, buf, q_next, num_queries, nq, 0, wg_tidx, mbar_K[0], mbar_idx + ) + if const_expr(self.wg0_chunks > 0): + # the first chunks of this tile's dKV, while the next tile is being gathered + self.dkv_rows(dkv, sIdx, tile, buf, mdKV, wg_tidx) + for c in cutlass.range_constexpr(self.wg0_chunks + 1): + if const_expr(c < self.wg0_chunks): + self.dkv_chunk(dkv, mma_dKV, sQ, sdO, chunk_layout, c) + if const_expr(c > 0): + warpgroup.wait_group(1) + else: + warpgroup.wait_group(0) + if const_expr(c > 0): + self.scatter(dkv[3][(c - 1) % self.num_acc], dkv[5], dkv[6], c - 1, mdKV, wg_tidx % 32) + if tile + 1 == num_tiles and q_next < num_queries: + cute.arch.barrier_arrive(barrier_id=BAR_QDO_FREE, number_of_threads=256) + it += 1 + self.store_dQ(acc_dQ, thr_dQ, mdQ, q, 0) + self.zero_skipped_dQ(acc_dQ, thr_dQ, mdQ, q, q_next, grid, 0) + nq += 1 + q = q_next + if it > 0: + cute.arch.barrier(barrier_id=BAR_PDS_FREE, number_of_threads=256) + + @cute.jit + def dkv_setup(self, mma_dKV, sP, sdS, wg_tidx): + """Per warpgroup dKV state; with an RS `mma_dKV`, the P^T / dS^T fragments are registers.""" + thr_dKV = mma_dKV.get_slice(wg_tidx) + tKVrPt = thr_dKV.make_fragment_A(thr_dKV.partition_A(cu.transpose_view(sP))) + tKVrdSt = thr_dKV.make_fragment_A(thr_dKV.partition_A(cu.transpose_view(sdS))) + tKVcKV = cu.acc_mn_view(thr_dKV.partition_C(cute.make_identity_tensor((TILE_N, CHUNK)))) + acc_shape = mma_dKV.partition_shape_C((TILE_N, CHUNK)) + accs = tuple(cute.make_rmem_tensor(acc_shape, Float32) for _ in range(self.num_acc)) + rows = cute.make_rmem_tensor((2,), Int32) + addrs = cute.make_rmem_tensor((4,), cutlass.Int64) + preds = cute.make_rmem_tensor((4,), Int32) + return (thr_dKV, tKVrPt, tKVrdSt, accs, tKVcKV, addrs, preds, rows) + + @cute.jit + def dkv_rows(self, dkv, sIdx, tile, buf, mdKV, wg_tidx): + thr_dKV, tKVrPt, tKVrdSt, accs, tKVcKV, addrs, preds, rows = dkv + for r in cutlass.range_constexpr(2): + rows[r] = sIdx[tile * TILE_N + tKVcKV[r, 0][0], buf] + self.scatter_targets(rows, mdKV, wg_tidx % 32, addrs, preds) + + @cute.jit + def dkv_chunk(self, dkv, mma_dKV, sQ, sdO, chunk_layout, c: cutlass.Constexpr[int]): + """Issue (not wait) the MMA of dKV columns [64 c, 64 c + 64) into accumulator c % 2.""" + thr_dKV, tKVrPt, tKVrdSt, accs, tKVcKV, addrs, preds, rows = dkv + half, cc = c // (NUM_CHUNKS // 2), c % (NUM_CHUNKS // 2) + sdO_c = cute.make_tensor(sdO[None, None, half].iterator + cc * HEADS * CHUNK, chunk_layout.outer) + sQ_c = cute.make_tensor(sQ[None, None, half].iterator + cc * HEADS * CHUNK, chunk_layout.outer) + tKVrdOt = thr_dKV.make_fragment_B(thr_dKV.partition_B(cu.transpose_view(sdO_c))) + tKVrQt = thr_dKV.make_fragment_B(thr_dKV.partition_B(cu.transpose_view(sQ_c))) + cu.gemm2(mma_dKV, accs[c % self.num_acc], tKVrPt, tKVrdOt, tKVrdSt, tKVrQt) + + @cute.jit + def wg1_loop( + self, + mQ, + mdO, + tma_Q, + tma_dO, + mKV, + mCount, + mdQ, + mdKV, + sQ, + sdO, + sK, + sP, + sdS, + sIdx, + mbar_QdO, + mbar_K, + mbar_idx, + mma_dQ_ss, + mma_dKV, + chunk_layout, + bid, + grid, + num_queries, + wg_tidx, + warp_idx, + ): + thr_dQ = mma_dQ_ss.get_slice(wg_tidx) + tdQrKt = thr_dQ.make_fragment_B(thr_dQ.partition_B(cu.transpose_view(sK))) + tdQrdS = thr_dQ.make_fragment_A(thr_dQ.partition_A(sdS)) + acc_dQ = cute.make_rmem_tensor(mma_dQ_ss.partition_shape_C((HEADS, HALF)), Float32) + + dkv = self.dkv_setup(mma_dKV, sP, sdS, wg_tidx) + if const_expr(self.rs): + ld_atom = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), cutlass.BFloat16 + ) + thr_ld = cute.make_tiled_copy_A(ld_atom, mma_dKV).get_slice(wg_tidx) + tLsPt = thr_ld.partition_S(cu.transpose_view(sP)) + tLsdSt = thr_ld.partition_S(cu.transpose_view(sdS)) + tLrPt = thr_ld.retile(dkv[1]) + tLrdSt = thr_ld.retile(dkv[2]) + chunks = list(range(self.wg0_chunks, NUM_CHUNKS)) + + nq = Int32(0) + q = bid + while q < num_queries: + num_tiles = (mCount[q] + TILE_N - 1) // TILE_N + if num_tiles == 0: + acc_dQ.fill(0.0) + self.store_dQ(acc_dQ, thr_dQ, mdQ, q, 1) + q += grid + else: + buf = nq % 2 + q_next = self.next_query(mCount, q, grid, num_queries) + acc_dQ.fill(0.0) + for tile in cutlass.range(num_tiles, unroll=1): + last = tile + 1 == num_tiles + cute.arch.barrier(barrier_id=BAR_PDS_FULL, number_of_threads=256) + self.dkv_rows(dkv, sIdx, tile, buf, mdKV, wg_tidx) + if const_expr(self.rs): + cute.copy(thr_ld, tLsPt, tLrPt) + cute.copy(thr_ld, tLsdSt, tLrdSt) + cu.gemm(mma_dQ_ss, acc_dQ, tdQrdS, tdQrKt[None, None, None, 1], zero_init=False) + reload = last and q_next < num_queries + for i in cutlass.range_constexpr(len(chunks) + 1): + if const_expr(i < len(chunks)): + self.dkv_chunk(dkv, mma_dKV, sQ, sdO, chunk_layout, chunks[i]) + if const_expr(i == 0): + # dQ_hi (issued first) is done: the high half of the tile is free + warpgroup.wait_group(1) + if const_expr(self.rs): + # dQ_hi has read dS and chunk 0 has consumed the ldmatrix results: WG0 may refill + cute.arch.barrier_arrive(barrier_id=BAR_PDS_FREE, number_of_threads=256) + self.gather_next( + mKV, + sK, + sIdx, + tile, + num_tiles, + buf, + q_next, + num_queries, + nq, + 1, + wg_tidx, + mbar_K[1], + mbar_idx, + ) + else: + if const_expr(i < len(chunks)): + warpgroup.wait_group(1) + else: + warpgroup.wait_group(0) + if const_expr(not self.rs): + cute.arch.barrier_arrive(barrier_id=BAR_PDS_FREE, number_of_threads=256) + # once the last chunk reading a half of Q / dO is done, load the next query's + done = chunks[i - 1] if i > 0 else self.wg0_chunks - 1 + if const_expr(done == NUM_CHUNKS // 2 - 1 and (i > 0 or self.wg0_chunks == NUM_CHUNKS // 2)): + if reload: + if const_expr(self.wg0_chunks > 0): + cute.arch.barrier(barrier_id=BAR_QDO_FREE, number_of_threads=256) + if warp_idx == 4: + self.load_QdO_half(mQ, mdO, sQ, sdO, tma_Q, tma_dO, q_next, 0, mbar_QdO[0]) + if const_expr(i == len(chunks)): + if reload and warp_idx == 4: + self.load_QdO_half(mQ, mdO, sQ, sdO, tma_Q, tma_dO, q_next, 1, mbar_QdO[1]) + if const_expr(i > 0): + self.scatter( + dkv[3][chunks[i - 1] % self.num_acc], dkv[5], dkv[6], chunks[i - 1], mdKV, wg_tidx % 32 + ) + self.store_dQ(acc_dQ, thr_dQ, mdQ, q, 1) + self.zero_skipped_dQ(acc_dQ, thr_dQ, mdQ, q, q_next, grid, 1) + nq += 1 + q = q_next + + @cute.jit + def zero_skipped_dQ(self, acc_dQ, thr_mma, mdQ, q, q_next, grid, half: cutlass.Constexpr[int]): + """dQ = 0 for the queries between q and q_next that `next_query` skipped (no non-empty slot).""" + qq = q + grid + if qq < q_next: + acc_dQ.fill(0.0) + while qq < q_next: + self.store_dQ(acc_dQ, thr_mma, mdQ, qq, half) + qq += grid + + @cute.jit + def store_dQ(self, acc_dQ: cute.Tensor, thr_mma, mdQ: cute.Tensor, q: Int32, half: cutlass.Constexpr[int]): + """Write this warpgroup's (64 x 256) half of the query's dQ straight from registers.""" + rdQ = cu.to_bf16(acc_dQ) + gdQ = cute.local_tile(mdQ[q, None, None], (HEADS, HALF), (0, half)) + cute.autovec_copy(rdQ, thr_mma.partition_C(gdQ)) + + @cute.jit + def scatter_targets( + self, rows: cute.Tensor, mdKV: cute.Tensor, lane: Int32, addrs: cute.Tensor, preds: cute.Tensor + ): + """Per tile: the two latent rows (a, b) each of this thread's accumulator rows feeds after the + quad swap in `scatter`, as byte addresses (lane and half-group offsets folded in) and validity.""" + is_b = ((lane // 4) % 2) == 1 + base = mdKV.iterator.toint() + (lane % 4) * 16 + cutlass.select_(is_b, Int32(64), Int32(0)) + for r in cutlass.range_constexpr(2): + own_row = rows[r] + partner_row = cute.arch.shuffle_sync_bfly(own_row, offset=4) + row_a = cutlass.select_(is_b, partner_row, own_row) + row_b = cutlass.select_(is_b, own_row, partner_row) + addrs[2 * r] = base + cutlass.Int64(cutlass.max(row_a, Int32(0))) * (DIM * 4) + addrs[2 * r + 1] = base + cutlass.Int64(cutlass.max(row_b, Int32(0))) * (DIM * 4) + preds[2 * r] = Int32(row_a >= 0) + preds[2 * r + 1] = Int32(row_b >= 0) + + @cute.jit + def scatter( + self, + acc: cute.Tensor, + addrs: cute.Tensor, + preds: cute.Tensor, + chunk: cutlass.Constexpr[int], + mdKV: cute.Tensor, + lane: Int32, + ): + """Atomically add a (64 slots x 64 cols) accumulator chunk to the latent rows its slots name. + + A thread of quad i (lane // 4) holds rows i and i + 8 of its warp's 16, columns + {2 rank, 2 rank + 1} + 8 j (rank = lane % 4). Columns 16 g + {2 rank, 2 rank + 1, 8 + 2 rank, + 9 + 2 rank} are stored as the 4 contiguous floats at 16 g + 4 rank of the row's chunk (undone by + the postprocess). Quads 2p and 2p + 1 swap halves so that each atomic instruction covers 4 rows + x 128 contiguous bytes (8 rows x 64 bytes runs ~20% slower in L2). + """ + A = cu.acc_mn_view(acc) + is_b = ((lane // 4) % 2) == 1 + for r in cutlass.range_constexpr(2): + for gp in cutlass.range_constexpr(2): + g0, g1 = 2 * gp, 2 * gp + 1 + recv = cute.make_rmem_tensor((4,), Float32) + for e in cutlass.range_constexpr(4): + recv[e] = cute.arch.shuffle_sync_bfly( + cu.select_f32(is_b, A[r, 4 * g0 + e], A[r, 4 * g1 + e]), offset=4 + ) + offset = chunk * CHUNK * 4 + gp * 128 + cu.red_add_v4_if( + preds[2 * r], + addrs[2 * r] + offset, + cu.select_f32(is_b, recv[0], A[r, 4 * g0 + 0]), + cu.select_f32(is_b, recv[1], A[r, 4 * g0 + 1]), + cu.select_f32(is_b, recv[2], A[r, 4 * g0 + 2]), + cu.select_f32(is_b, recv[3], A[r, 4 * g0 + 3]), + ) + cu.red_add_v4_if( + preds[2 * r + 1], + addrs[2 * r + 1] + offset, + cu.select_f32(is_b, A[r, 4 * g1 + 0], recv[0]), + cu.select_f32(is_b, A[r, 4 * g1 + 1], recv[1]), + cu.select_f32(is_b, A[r, 4 * g1 + 2], recv[2]), + cu.select_f32(is_b, A[r, 4 * g1 + 3], recv[3]), + ) diff --git a/prime_kernels/fp8_dual_cast/__init__.py b/prime_kernels/fp8_dual_cast/__init__.py new file mode 100644 index 0000000..9222afd --- /dev/null +++ b/prime_kernels/fp8_dual_cast/__init__.py @@ -0,0 +1,3 @@ +from prime_kernels.fp8_dual_cast.interface import dual_cast_to_fp8, fp8_dual_cast + +__all__ = ["dual_cast_to_fp8", "fp8_dual_cast"] diff --git a/prime_kernels/fp8_dual_cast/interface.py b/prime_kernels/fp8_dual_cast/interface.py new file mode 100644 index 0000000..87456d8 --- /dev/null +++ b/prime_kernels/fp8_dual_cast/interface.py @@ -0,0 +1,100 @@ +"""Per-token and transposed per-token e4m3 casts of a bf16 matrix in one pass (optionally after DeepSeek-V4's +inverse RoPE), with the outputs of prime-rl's `per_token_cast_to_fp8_triton` and +`per_token_cast_to_fp8_tp_triton` (on the rotated matrix) bit for bit.""" + +import torch + +BLOCK = 128 +_compiled: dict = {} + + +def _kernel(rope: bool, head_dim: int, device: torch.device): + key = (rope, head_dim, device.index) + if key not in _compiled: + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + + from prime_kernels.fp8_dual_cast.kernel import DualCastSm90 + + def fake(shape, dtype): + t = torch.empty(shape, dtype=dtype, device=device) + return from_dlpack(t, assumed_align=16, enable_tvm_ffi=True).mark_layout_dynamic(leading_dim=t.ndim - 1) + + args = ( + fake((BLOCK, 2 * BLOCK), torch.int16), + fake((2, 64), torch.float32), + fake((BLOCK,), torch.int32), + fake((BLOCK, 2 * BLOCK), torch.uint8), + fake((2, BLOCK), torch.float32), + fake((2 * BLOCK, BLOCK), torch.uint8), + fake((1, 2 * BLOCK), torch.float32), + ) + stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + _compiled[key] = cute.compile( + DualCastSm90(rope, head_dim), *args, cutlass.Int32(1), stream, options="--enable-tvm-ffi" + ) + return _compiled[key] + + +def _casts(x, cos_sin, positions, head_dim): + import cuda.bindings.driver as cuda + + rows, cols = x.shape + prows = (rows + BLOCK - 1) // BLOCK * BLOCK + out = torch.empty(rows, cols, dtype=torch.float8_e4m3fn, device=x.device) + sf = torch.empty(cols // BLOCK, (rows + 3) // 4 * 4, dtype=torch.float32, device=x.device) + out_t = torch.empty(cols, prows, dtype=torch.float8_e4m3fn, device=x.device) + sf_t = torch.empty(prows // BLOCK, cols, dtype=torch.float32, device=x.device) + rope = cos_sin is not None + if rope: + assert cos_sin.shape[-1] == 64 and head_dim % BLOCK == 0 and cols % head_dim == 0 + cos_sin = cos_sin.float().contiguous() + positions = positions.reshape(rows).to(torch.int32) + else: + cos_sin = torch.zeros(1, 64, dtype=torch.float32, device=x.device) + positions = torch.zeros(max(rows, 1), dtype=torch.int32, device=x.device) + head_dim = BLOCK + if rows: + stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + _kernel(rope, head_dim, x.device)( + x.view(torch.int16), cos_sin, positions, out.view(torch.uint8), sf, out_t.view(torch.uint8), sf_t, rows, + stream, + ) + return out, sf, out_t, sf_t + + +@torch.library.custom_op("prime_kernels::fp8_dual_cast", mutates_args=()) +def fp8_dual_cast( + x: torch.Tensor, cos_sin: torch.Tensor | None, positions: torch.Tensor | None, head_dim: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """`dual_cast_to_fp8` as an op, returning the scale tensors whole: (x_fp8, sf (cols / 128, rows rounded up + to 4)), (x_t_fp8 (cols, padded rows), sf_t (padded rows / 128, cols)); the casts' scales are + `sf[:, :rows].T` and `sf_t.T`.""" + assert x.dim() == 2 and x.is_contiguous() and x.dtype == torch.bfloat16 and x.shape[1] % BLOCK == 0 + return _casts(x, cos_sin, positions, head_dim) + + +@fp8_dual_cast.register_fake +def _fp8_dual_cast_fake(x, cos_sin, positions, head_dim): + rows, cols = x.shape + prows = (rows + BLOCK - 1) // BLOCK * BLOCK + return ( + x.new_empty(rows, cols, dtype=torch.float8_e4m3fn), + x.new_empty(cols // BLOCK, (rows + 3) // 4 * 4, dtype=torch.float32), + x.new_empty(cols, prows, dtype=torch.float8_e4m3fn), + x.new_empty(prows // BLOCK, cols, dtype=torch.float32), + ) + + +def dual_cast_to_fp8( + x: torch.Tensor, # (rows, cols) bf16, contiguous + rope: tuple[torch.Tensor, torch.Tensor, int] | None = None, # ([cos | sin] f32 (positions, 64), positions, head_dim) +): + """((x_fp8, scales (rows, cols / 128) MN-major), (x_t_fp8 (cols, padded rows), scales (cols, padded rows / 128))) + of `x`, or of `dsv4_rope(x, inverse=True)` with `rope`.""" + assert x.dim() == 2 and x.is_contiguous() and x.dtype == torch.bfloat16 and x.shape[1] % BLOCK == 0 + rows = x.shape[0] + out, sf, out_t, sf_t = _casts(x, *(rope if rope is not None else (None, None, BLOCK))) + return (out, sf[:, :rows].T), (out_t, sf_t.T) diff --git a/prime_kernels/fp8_dual_cast/kernel.py b/prime_kernels/fp8_dual_cast/kernel.py new file mode 100644 index 0000000..c34cbfc --- /dev/null +++ b/prime_kernels/fp8_dual_cast/kernel.py @@ -0,0 +1,329 @@ +"""SM90 per-token (1 x 128) and transposed per-token (128 x 1) e4m3 casts of a bf16 matrix in one pass, +optionally after the inverse RoPE of each head's trailing 64 channels (the query gradient of DeepSeek-V4.1). + +One CTA (256 threads) per [128 tokens x 128 channels] tile: + 1. each half-warp loads one token row of the tile (16 B per thread), rotates it if the tile holds RoPE + channels, reduces the row's amax over its 16 lanes, writes the row's e4m3 values and scale, and parks + the (rotated) bf16 row in shared memory; + 2. each thread takes half of one channel column, the two halves combine their amax through shared + memory, and the column's e4m3 values go to a transposed shared tile; + 3. the transposed tile is written out, 128 contiguous bytes per channel. +The arithmetic is that of prime-rl's Triton kernels (`_triton_rope_inplace`, `_fp8_scale`, `_fp8_quantize`) +instruction for instruction, so the results are bitwise equal. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +from prime_kernels.dsa_sparse_attn_bwd._cute_utils import cvt_bf16x2 + +BLOCK = 128 +ROPE = 64 +T_PITCH = BLOCK + 4 # bytes per channel row of the transposed tile (conflict-free 4-byte stores) + + +@dsl_user_op +def ld_global_v4(ptr, *, loc=None, ip=None): + r = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + [cutlass.Int64(ptr).ir_value(loc=loc, ip=ip)], + "ld.global.nc.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return tuple(Int32(llvm.extractvalue(T.i32(), r, [i], loc=loc, ip=ip)) for i in range(4)) + + +def _asm(ret, asm, cons, args, side=False, loc=None, ip=None): + return llvm.inline_asm( + ret, [a.ir_value(loc=loc, ip=ip) for a in args], asm, cons, has_side_effects=side, is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def ld_global_v4_f32(ptr, *, loc=None, ip=None): + r = _asm(llvm.StructType.get_literal([T.f32()] * 4), "ld.global.nc.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,l", [cutlass.Int64(ptr)], loc=loc, ip=ip) + return tuple(Float32(llvm.extractvalue(T.f32(), r, [i], loc=loc, ip=ip)) for i in range(4)) + + +@dsl_user_op +def ld_shared_bf16x2_f32(addr, *, loc=None, ip=None): + """Two adjacent bf16 from shared memory, as fp32 (low address first).""" + r = _asm(llvm.StructType.get_literal([T.f32()] * 2), + "{\n.reg .b32 w;\nld.shared.b32 w, [$2];\nshl.b32 $0, w, 16;\nand.b32 $1, w, 0xffff0000;\n}", + "=f,=f,r", [Int32(addr)], side=True, loc=loc, ip=ip) + return tuple(Float32(llvm.extractvalue(T.f32(), r, [i], loc=loc, ip=ip)) for i in range(2)) + + +@dsl_user_op +def bf16_lo(v, *, loc=None, ip=None) -> Float32: + return Float32(_asm(T.f32(), "shl.b32 $0, $1, 16;", "=f,r", [Int32(v)], loc=loc, ip=ip)) + + +@dsl_user_op +def bf16_hi(v, *, loc=None, ip=None) -> Float32: + return Float32(_asm(T.f32(), "and.b32 $0, $1, 0xffff0000;", "=f,r", [Int32(v)], loc=loc, ip=ip)) + + +@dsl_user_op +def st_shared_v4_b32(addr, a, b, c, d, *, loc=None, ip=None) -> None: + _asm(None, "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", "r,r,r,r,r", [Int32(addr), Int32(a), Int32(b), Int32(c), + Int32(d)], side=True, loc=loc, ip=ip) + + +@dsl_user_op +def ld_shared_u16_f32(addr, *, loc=None, ip=None) -> Float32: + """A bf16 from shared memory, as fp32.""" + return Float32(_asm(T.f32(), "{\n.reg .b16 h;\nld.shared.b16 h, [$1];\ncvt.f32.bf16 $0, h;\n}", "=f,r", + [Int32(addr)], side=True, loc=loc, ip=ip)) + + +@dsl_user_op +def st_shared_b32(addr, v, *, loc=None, ip=None) -> None: + _asm(None, "st.shared.b32 [$0], $1;", "r,r", [Int32(addr), Int32(v)], side=True, loc=loc, ip=ip) + + +@dsl_user_op +def ld_shared_b32(addr, *, loc=None, ip=None) -> Int32: + return Int32(_asm(T.i32(), "ld.shared.b32 $0, [$1];", "=r,r", [Int32(addr)], side=True, loc=loc, ip=ip)) + + +@dsl_user_op +def st_shared_f32(addr, v, *, loc=None, ip=None) -> None: + _asm(None, "st.shared.f32 [$0], $1;", "r,f", [Int32(addr), Float32(v)], side=True, loc=loc, ip=ip) + + +@dsl_user_op +def ld_shared_f32(addr, *, loc=None, ip=None) -> Float32: + return Float32(_asm(T.f32(), "ld.shared.f32 $0, [$1];", "=f,r", [Int32(addr)], side=True, loc=loc, ip=ip)) + + +@dsl_user_op +def st_global_v2(ptr, a, b, *, loc=None, ip=None) -> None: + _asm(None, "st.global.v2.b32 [$0], {$1, $2};", "l,r,r", [cutlass.Int64(ptr), Int32(a), Int32(b)], side=True, + loc=loc, ip=ip) + + +@dsl_user_op +def st_global_b32(ptr, a, *, loc=None, ip=None) -> None: + _asm(None, "st.global.b32 [$0], $1;", "l,r", [cutlass.Int64(ptr), Int32(a)], side=True, loc=loc, ip=ip) + + +def _f32_asm(asm: str, *args, loc=None, ip=None) -> Float32: + constraints = "=f," + ",".join("f" for _ in args) + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip) for a in args], + asm, + constraints, + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def round_bf16(a: Float32, *, loc=None, ip=None) -> Float32: + """`a` rounded to bf16 (round to nearest even), as an fp32 value.""" + return _f32_asm("{\n.reg .b16 h;\ncvt.rn.bf16.f32 h, $1;\ncvt.f32.bf16 $0, h;\n}", a, loc=loc, ip=ip) + + +@dsl_user_op +def fma_rn(a: Float32, b: Float32, c: Float32, *, loc=None, ip=None) -> Float32: + return _f32_asm("fma.rn.f32 $0, $1, $2, $3;", a, b, c, loc=loc, ip=ip) + + +@dsl_user_op +def mul_rn(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + return _f32_asm("mul.rn.f32 $0, $1, $2;", a, b, loc=loc, ip=ip) + + +@dsl_user_op +def div_rn(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + return _f32_asm("div.rn.f32 $0, $1, $2;", a, b, loc=loc, ip=ip) + + +@dsl_user_op +def max_abs(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + """max(a, |b|).""" + return _f32_asm("{\n.reg .f32 t;\nabs.f32 t, $2;\nmax.f32 $0, $1, t;\n}", a, b, loc=loc, ip=ip) + + +@dsl_user_op +def quantize_e4m3x2(lo: Float32, hi: Float32, rcp: Float32, scale: Float32, *, loc=None, ip=None) -> cutlass.Int32: + """Two values / `scale` rounded to e4m3 (satfinite), as prime-rl's `_fp8_quantize` computes them + (q = x * rcp, one Markstein correction, q kept where it is +-0); `lo` in the low byte.""" + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [Float32(v).ir_value(loc=loc, ip=ip) for v in (lo, hi, rcp, scale)], + "{\n.reg .f32 q0, q1, r0, r1, y0, y1;\n.reg .pred z0, z1;\n.reg .b16 h;\n" + "mul.rn.f32 q0, $1, $3;\nmul.rn.f32 q1, $2, $3;\n" + "neg.f32 r0, q0;\nneg.f32 r1, q1;\n" + "fma.rn.f32 r0, r0, $4, $1;\nfma.rn.f32 r1, r1, $4, $2;\n" + "fma.rn.f32 y0, r0, $3, q0;\nfma.rn.f32 y1, r1, $3, q1;\n" + "setp.eq.f32 z0, q0, 0f00000000;\nsetp.eq.f32 z1, q1, 0f00000000;\n" + "selp.f32 y0, q0, y0, z0;\nselp.f32 y1, q1, y1, z1;\n" + "cvt.rn.satfinite.e4m3x2.f32 h, y1, y0;\ncvt.u32.u16 $0, h;\n}", + "=r,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def neg_triton(a: Float32, *, loc=None, ip=None) -> Float32: + """Triton's float negation, 0 - a (so -(+0) is +0).""" + return _f32_asm("sub.rn.f32 $0, 0f00000000, $1;", a, loc=loc, ip=ip) + + +class DualCastSm90: + def __init__(self, rope: bool, head_dim: int): + assert head_dim % BLOCK == 0 + self.rope = rope + self.head_dim = head_dim + + @cute.jit + def __call__( + self, + mX: cute.Tensor, # (rows, cols) bf16 (as int16) + mCS: cute.Tensor, # (positions, 64) f32 [cos | sin] + mPos: cute.Tensor, # (rows,) int32 + mOut: cute.Tensor, # (rows, cols) uint8 + mSF: cute.Tensor, # (cols / 128, >= rows) f32 + mOutT: cute.Tensor, # (cols, padded rows) uint8 + mSFT: cute.Tensor, # (padded rows / 128, cols) f32 + rows: Int32, + stream, + ): + cols = mX.shape[1] + prows = mOutT.shape[1] + + @cute.struct + class Smem: + tile: cute.struct.Align[cute.struct.MemRange[cutlass.Int16, BLOCK * BLOCK], 128] + red: cute.struct.Align[cute.struct.MemRange[Float32, 4 * BLOCK], 128] + + self.smem = Smem + self.kernel(mX, mCS, mPos, mOut, mSF, mOutT, mSFT, rows).launch( + grid=[cols // BLOCK, prows // BLOCK, 1], block=[256, 1, 1], smem=Smem.size_in_bytes(), stream=stream + ) + + @cute.kernel + def kernel(self, mX, mCS, mPos, mOut, mSF, mOutT, mSFT, rows: Int32): + tid, _, _ = cute.arch.thread_idx() + bk, bm, _ = cute.arch.block_idx() + cols = mX.shape[1] + prows = mOutT.shape[1] + storage = cutlass.utils.SmemAllocator().allocate(self.smem) + s_tile = storage.tile.data_ptr().toint() + s_t = s_tile # the transposed tile reuses the bf16 tile once every column is in registers + s_red = storage.red.data_ptr().toint() + x_base = mX.iterator.toint() + out_base = mOut.iterator.toint() + outt_base = mOutT.iterator.toint() + + c16 = tid % 16 + col0 = bk * BLOCK + 8 * c16 + rope_tile = (bk + 1) * BLOCK % self.head_dim == 0 + vals = cute.make_rmem_tensor((8,), Float32) + bits = cute.make_rmem_tensor((4,), Int32) + raw = cute.make_rmem_tensor((BLOCK // 16 * 4,), Int32) + # 1. rows: load (all eight rows of this thread first), rotate, per-token cast, park in shared memory + for i in cutlass.range_constexpr(BLOCK // 16): + grow = bm * BLOCK + tid // 16 + 16 * i + v = ld_global_v4(x_base + (cutlass.Int64(cutlass.min(grow, rows - 1)) * cols + col0) * 2) + for e in cutlass.range_constexpr(4): + raw[4 * i + e] = cutlass.select_(grow < rows, v[e], Int32(0)) + + for i in cutlass.range_constexpr(BLOCK // 16): + row = tid // 16 + 16 * i + grow = bm * BLOCK + row + for e in cutlass.range_constexpr(4): + bits[e] = raw[4 * i + e] + for e in cutlass.range_constexpr(4): + vals[2 * e] = bf16_lo(bits[e]) + vals[2 * e + 1] = bf16_hi(bits[e]) + if const_expr(self.rope): + if rope_tile and c16 >= (BLOCK - ROPE) // 8 and grow < rows: + cs_ptr = mCS.iterator.toint() + (cutlass.Int64(mPos[grow]) * ROPE + 4 * (c16 - (BLOCK - ROPE) // 8)) * 4 + coss = ld_global_v4_f32(cs_ptr) + sins = ld_global_v4_f32(cs_ptr + ROPE // 2 * 4) + for e in cutlass.range_constexpr(4): + cs = coss[e] + sn = neg_triton(sins[e]) + x1 = vals[2 * e] + x2 = vals[2 * e + 1] + y1 = round_bf16(fma_rn(x1, cs, neg_triton(mul_rn(x2, sn)))) + y2 = round_bf16(fma_rn(x1, sn, mul_rn(x2, cs))) + vals[2 * e] = y1 + vals[2 * e + 1] = y2 + bits[e] = cvt_bf16x2(y1, y2) + amax = Float32(0.0) + for e in cutlass.range_constexpr(8): + amax = max_abs(amax, vals[e]) + for k in cutlass.range_constexpr(4): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, offset=8 >> k)) + sc = div_rn(cute.arch.fmax(amax, Float32(1e-10)), Float32(448.0)) + rcp = div_rn(Float32(1.0), sc) + q01 = quantize_e4m3x2(vals[0], vals[1], rcp, sc) + q23 = quantize_e4m3x2(vals[2], vals[3], rcp, sc) + q45 = quantize_e4m3x2(vals[4], vals[5], rcp, sc) + q67 = quantize_e4m3x2(vals[6], vals[7], rcp, sc) + if grow < rows: + st_global_v2(out_base + cutlass.Int64(grow) * cols + col0, q01 | (q23 << 16), q45 | (q67 << 16)) + if c16 == 0: + mSF[bk, grow] = sc + st_shared_v4_b32(s_tile + (row * BLOCK + 8 * c16) * 2, bits[0], bits[1], bits[2], bits[3]) + cute.arch.sync_threads() + # 2. columns: thread = (column pair cp, token quarter tq); amax over the quarter, combined through shared + # memory, then the transposed cast of both columns into shared memory + cp = tid % (BLOCK // 2) + tq = tid // (BLOCK // 2) + lo = cute.make_rmem_tensor((BLOCK // 4,), Float32) + hi = cute.make_rmem_tensor((BLOCK // 4,), Float32) + amax_lo = Float32(0.0) + amax_hi = Float32(0.0) + for k in cutlass.range_constexpr(BLOCK // 4): + v = ld_shared_bf16x2_f32(s_tile + ((tq * (BLOCK // 4) + k) * BLOCK + 2 * cp) * 2) + lo[k] = v[0] + hi[k] = v[1] + amax_lo = max_abs(amax_lo, v[0]) + amax_hi = max_abs(amax_hi, v[1]) + st_shared_f32(s_red + (tq * BLOCK + 2 * cp) * 4, amax_lo) + st_shared_f32(s_red + (tq * BLOCK + 2 * cp + 1) * 4, amax_hi) + cute.arch.sync_threads() + for h in cutlass.range_constexpr(2): + c = 2 * cp + h + amax = ld_shared_f32(s_red + c * 4) + for j in cutlass.range_constexpr(1, 4): + amax = cute.arch.fmax(amax, ld_shared_f32(s_red + (j * BLOCK + c) * 4)) + sc = div_rn(cute.arch.fmax(amax, Float32(1e-10)), Float32(448.0)) + rcp = div_rn(Float32(1.0), sc) + if tq == 0: + mSFT[bm, bk * BLOCK + c] = sc + vv = lo if h == 0 else hi + for k4 in cutlass.range_constexpr(BLOCK // 16): + w = quantize_e4m3x2(vv[4 * k4], vv[4 * k4 + 1], rcp, sc) | ( + quantize_e4m3x2(vv[4 * k4 + 2], vv[4 * k4 + 3], rcp, sc) << 16 + ) + st_shared_b32(s_t + c * T_PITCH + tq * (BLOCK // 4) + 4 * k4, w) + cute.arch.sync_threads() + # 3. write the transposed tile: channel row r, 32 words of 4 tokens + lane = tid % 32 + warp = tid // 32 + for k in cutlass.range_constexpr(BLOCK // 8): + r = warp + 8 * k + w = ld_shared_b32(s_t + r * T_PITCH + 4 * lane) + st_global_b32(outt_base + cutlass.Int64(bk * BLOCK + r) * prows + bm * BLOCK + 4 * lane, w) diff --git a/prime_kernels/kernels.toml b/prime_kernels/kernels.toml index 5470c1a..aae3520 100644 --- a/prime_kernels/kernels.toml +++ b/prime_kernels/kernels.toml @@ -30,6 +30,36 @@ python-only = true requires = ["tilelang"] arch = ["8.0", "9.0", "10.0", "10.3"] +[dsa_indexer_topk] +description = "DeepSeek-V4.1 Lightning Indexer forward fused with top-k and two-level candidate blocks" +python-only = true +requires = ["triton"] +arch = ["9.0"] + +[moe_experts] +description = "Expert MLP of a rank's local experts (grouped GEMMs + fused clamped SwiGLU), forward and backward, on Hopper; bf16, or blockwise FP8 through DeepGEMM" +python-only = true +requires = ["triton.experimental.gluon"] +arch = ["9.0"] + +[mhc_projection] +description = "DeepSeek-V4.1 mHC projection (RMS + gate logits + stream collapse) and gate activations, forward and backward" +python-only = true +requires = ["triton"] +arch = ["9.0"] + +[dsa_sparse_attn_bwd] +description = "DeepSeek-V4.1 sparse attention backward (shared K = V latent, per-query slot lists, attention sink) on Hopper" +python-only = true +requires = ["cutlass", "triton"] +arch = ["9.0"] + +[fp8_dual_cast] +description = "Per-token and transposed per-token e4m3 casts of a bf16 matrix in one pass (optional DeepSeek-V4 inverse RoPE) on Hopper" +python-only = true +requires = ["cutlass"] +arch = ["9.0"] + # rmsnorm is not built yet: only its sources are committed. Uncomment the table below to # put it back into the build (and to make the registry report on it). # diff --git a/prime_kernels/mhc_projection/__init__.py b/prime_kernels/mhc_projection/__init__.py new file mode 100644 index 0000000..c7407fa --- /dev/null +++ b/prime_kernels/mhc_projection/__init__.py @@ -0,0 +1,48 @@ +import torch + +from prime_kernels.mhc_projection.gates import mhc_gates +from prime_kernels.mhc_projection.projection import mhc_projection, mhc_projection_streams + + +def hyper_connection( + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + pre_mix: torch.Tensor | None = None, + *, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """DeepSeek-V4.1 mHC gates for one sublayer, plus the collapse of its input streams. + + `x`: `(..., hc, d)` bf16 streams. `weight`: `((2 + hc) * hc, hc * d)`, `scale`: `(3,)`, + `base`: `((2 + hc) * hc,)`. `pre_mix`: `(..., hc)`, the gate collapsing these streams (V4.1 + computes it one sublayer earlier). Returns fp32 `pre`, `post`, `comb` and the collapsed + `(..., d)` sequence in `x.dtype` (None without `pre_mix`). + """ + mixes, collapsed = mhc_projection(x, weight, pre_mix, rms_eps) + pre, post, comb = mhc_gates(mixes, scale, base, x.shape[-2], sinkhorn_iters, hc_eps) + return pre, post, comb, collapsed + + +def hyper_connection_streams( + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + pre_mix: torch.Tensor | None = None, + *, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor]: + """`hyper_connection` plus the streams `x` as a view for their other readers (see + `mhc_projection_streams`): reading them through it saves a pass in backward.""" + mixes, collapsed, streams = mhc_projection_streams(x, weight, pre_mix, rms_eps) + pre, post, comb = mhc_gates(mixes, scale, base, x.shape[-2], sinkhorn_iters, hc_eps) + return pre, post, comb, collapsed, streams + + +__all__ = ["hyper_connection", "hyper_connection_streams", "mhc_gates", "mhc_projection", "mhc_projection_streams"] diff --git a/prime_kernels/mhc_projection/gates.py b/prime_kernels/mhc_projection/gates.py new file mode 100644 index 0000000..b35d4c6 --- /dev/null +++ b/prime_kernels/mhc_projection/gates.py @@ -0,0 +1,282 @@ +"""The mHC gate activations, forward and backward, in one Triton kernel each. + +From `mixes` `(T, (2 + hc) * hc)` fp32 (the RMS-normalized projection), `scale` `(3,)` and `base` +`((2 + hc) * hc,)`: + + pre = sigmoid(mixes[:, :hc] * scale[0] + base[:hc]) + eps + post = 2 * sigmoid(mixes[:, hc:2hc] * scale[1] + base[hc:2hc]) + comb = sinkhorn(mixes[:, 2hc:].view(hc, hc) * scale[2] + base[2hc:].view(hc, hc)) + +where `sinkhorn` is a row softmax plus `eps`, a column normalization, then `iters - 1` rounds of +row and column normalization, each dividing by `sum + eps`. The backward recomputes the Sinkhorn +iterates once into a scratch buffer and walks them in reverse. +""" + +import torch +import triton +import triton.language as tl + +BLOCK_T = 32 + + +@triton.jit +def _comb_logits(mixes_ptr, scale_ptr, base_ptr, rows, row_mask, HC: tl.constexpr, N: tl.constexpr): + r = tl.arange(0, HC) + offs = 2 * HC + r[:, None] * HC + r[None, :] + mask = row_mask[:, None, None] + mixes = tl.load(mixes_ptr + rows[:, None, None] * N + offs[None, :, :], mask=mask, other=0.0) + base = tl.load(base_ptr + offs) + return mixes * tl.load(scale_ptr + 2) + base[None, :, :] + + +@triton.jit +def _softmax_rows(z): + e = tl.exp(z - tl.max(z, axis=2)[:, :, None]) + return e / tl.sum(e, axis=2)[:, :, None] + + +@triton.jit +def _sinkhorn_prefix(z, num_ops, eps): + """The iterate after the first `num_ops` (>= 1) of the Sinkhorn steps. + + Step 0 is the row softmax plus eps; odd steps normalize columns, even steps rows. + """ + m = _softmax_rows(z) + eps + for step in range(1, num_ops): + if step % 2 == 1: + m = m / (tl.sum(m, axis=1)[:, None, :] + eps) + else: + m = m / (tl.sum(m, axis=2)[:, :, None] + eps) + return m + + +@triton.jit +def _gates_fwd_kernel( + mixes_ptr, + scale_ptr, + base_ptr, + pre_ptr, + post_ptr, + comb_ptr, + T, + eps, + HC: tl.constexpr, + N: tl.constexpr, + ITERS: tl.constexpr, + BLOCK_T: tl.constexpr, +): + rows = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + row_mask = rows < T + c = tl.arange(0, HC) + mask2 = row_mask[:, None] + + pre_z = tl.load(mixes_ptr + rows[:, None] * N + c[None, :], mask=mask2, other=0.0) + pre_z = pre_z * tl.load(scale_ptr) + tl.load(base_ptr + c)[None, :] + tl.store(pre_ptr + rows[:, None] * HC + c[None, :], tl.sigmoid(pre_z) + eps, mask=mask2) + + post_z = tl.load(mixes_ptr + rows[:, None] * N + HC + c[None, :], mask=mask2, other=0.0) + post_z = post_z * tl.load(scale_ptr + 1) + tl.load(base_ptr + HC + c)[None, :] + tl.store(post_ptr + rows[:, None] * HC + c[None, :], 2 * tl.sigmoid(post_z), mask=mask2) + + comb_z = _comb_logits(mixes_ptr, scale_ptr, base_ptr, rows, row_mask, HC, N) + comb = _sinkhorn_prefix(comb_z, 2 * ITERS, eps) + tl.store( + comb_ptr + rows[:, None, None] * (HC * HC) + c[None, :, None] * HC + c[None, None, :], + comb, + mask=row_mask[:, None, None], + ) + + +@triton.jit +def _gates_bwd_kernel( + mixes_ptr, + scale_ptr, + base_ptr, + grad_pre_ptr, + grad_post_ptr, + grad_comb_ptr, + grad_mixes_ptr, + grad_z_ptr, + scratch_ptr, + T, + eps, + HC: tl.constexpr, + N: tl.constexpr, + ITERS: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """Writes the gradients w.r.t. `mixes` and w.r.t. the pre-activation logits `z`. + + `scratch` holds the input of every Sinkhorn step after the first, `(2 * ITERS, T, HC * HC)`. + """ + rows = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + row_mask = rows < T + c = tl.arange(0, HC) + mask2 = row_mask[:, None] + offs2 = rows[:, None] * N + c[None, :] + + scale = tl.load(scale_ptr) + pre_z = tl.load(mixes_ptr + offs2, mask=mask2, other=0.0) * scale + tl.load(base_ptr + c)[None, :] + sig = tl.sigmoid(pre_z) + grad = tl.load(grad_pre_ptr + rows[:, None] * HC + c[None, :], mask=mask2, other=0.0) * sig * (1 - sig) + tl.store(grad_z_ptr + offs2, grad, mask=mask2) + tl.store(grad_mixes_ptr + offs2, grad * scale, mask=mask2) + + scale = tl.load(scale_ptr + 1) + post_z = tl.load(mixes_ptr + HC + offs2, mask=mask2, other=0.0) * scale + tl.load(base_ptr + HC + c)[None, :] + sig = tl.sigmoid(post_z) + grad = 2 * tl.load(grad_post_ptr + rows[:, None] * HC + c[None, :], mask=mask2, other=0.0) * sig * (1 - sig) + tl.store(grad_z_ptr + HC + offs2, grad, mask=mask2) + tl.store(grad_mixes_ptr + HC + offs2, grad * scale, mask=mask2) + + comb_z = _comb_logits(mixes_ptr, scale_ptr, base_ptr, rows, row_mask, HC, N) + mask3 = row_mask[:, None, None] + grad = tl.load( + grad_comb_ptr + rows[:, None, None] * (HC * HC) + c[None, :, None] * HC + c[None, None, :], + mask=mask3, + other=0.0, + ) + scratch = scratch_ptr + rows[:, None, None] * (HC * HC) + c[None, :, None] * HC + c[None, None, :] + m = _softmax_rows(comb_z) + eps + for step in range(1, 2 * ITERS): + tl.store(scratch + step * T * HC * HC, m, mask=mask3) + if step % 2 == 1: + m = m / (tl.sum(m, axis=1)[:, None, :] + eps) + else: + m = m / (tl.sum(m, axis=2)[:, :, None] + eps) + # Back through steps 2 * ITERS - 1 .. 1. For out = m / s with s = sum(m) + eps along an axis, + # d m = (d out - sum(d out * out)) / s along that axis. + for back in range(0, 2 * ITERS - 1): + step = 2 * ITERS - 1 - back + m = tl.load(scratch + step * T * HC * HC, mask=mask3, other=1.0) + if step % 2 == 1: + col_sum = tl.sum(m, axis=1)[:, None, :] + eps + grad = (grad - tl.sum(grad * (m / col_sum), axis=1)[:, None, :]) / col_sum + else: + row_sum = tl.sum(m, axis=2)[:, :, None] + eps + grad = (grad - tl.sum(grad * (m / row_sum), axis=2)[:, :, None]) / row_sum + # Step 0: softmax (+ eps, which has no gradient). + p = _softmax_rows(comb_z) + grad = p * (grad - tl.sum(grad * p, axis=2)[:, :, None]) + offs3 = rows[:, None, None] * N + 2 * HC + c[None, :, None] * HC + c[None, None, :] + tl.store(grad_z_ptr + offs3, grad, mask=mask3) + tl.store(grad_mixes_ptr + offs3, grad * tl.load(scale_ptr + 2), mask=mask3) + + +def _check(mixes: torch.Tensor, scale: torch.Tensor, base: torch.Tensor, hc: int) -> None: + n = (2 + hc) * hc + if mixes.dim() != 2 or mixes.shape[1] != n: + raise ValueError(f"expected mixes of shape (tokens, {n}), got {tuple(mixes.shape)}") + if hc & (hc - 1): + raise ValueError(f"hc must be a power of two, got {hc}") + if tuple(scale.shape) != (3,) or tuple(base.shape) != (n,): + raise ValueError(f"expected scale (3,) and base ({n},), got {tuple(scale.shape)}, {tuple(base.shape)}") + + +@torch.library.custom_op("prime_kernels::mhc_gates_forward", mutates_args=()) +def mhc_gates_forward( + mixes: torch.Tensor, scale: torch.Tensor, base: torch.Tensor, hc: int, iters: int, eps: float +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + _check(mixes, scale, base, hc) + tokens = mixes.shape[0] + pre = mixes.new_empty(tokens, hc, dtype=torch.float32) + post = mixes.new_empty(tokens, hc, dtype=torch.float32) + comb = mixes.new_empty(tokens, hc, hc, dtype=torch.float32) + if tokens > 0: + _gates_fwd_kernel[(triton.cdiv(tokens, BLOCK_T),)]( + mixes.float().contiguous(), + scale.float().contiguous(), + base.float().contiguous(), + pre, + post, + comb, + tokens, + eps, + HC=hc, + N=mixes.shape[1], + ITERS=iters, + BLOCK_T=BLOCK_T, + ) + return pre, post, comb + + +@mhc_gates_forward.register_fake +def _mhc_gates_forward_fake(mixes, scale, base, hc, iters, eps): + _check(mixes, scale, base, hc) + tokens = mixes.shape[0] + pre = mixes.new_empty(tokens, hc, dtype=torch.float32) + return pre, torch.empty_like(pre), mixes.new_empty(tokens, hc, hc, dtype=torch.float32) + + +@torch.library.custom_op("prime_kernels::mhc_gates_backward", mutates_args=()) +def mhc_gates_backward( + mixes: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + grad_pre: torch.Tensor, + grad_post: torch.Tensor, + grad_comb: torch.Tensor, + hc: int, + iters: int, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Returns `(grad_mixes, grad_z)`, the latter w.r.t. `mixes * scale + base`.""" + tokens = mixes.shape[0] + grad_mixes = torch.empty_like(mixes, dtype=torch.float32) + grad_z = torch.empty_like(mixes, dtype=torch.float32) + scratch = mixes.new_empty(2 * iters, tokens, hc * hc, dtype=torch.float32) + if tokens > 0: + _gates_bwd_kernel[(triton.cdiv(tokens, BLOCK_T),)]( + mixes.float().contiguous(), + scale.float().contiguous(), + base.float().contiguous(), + grad_pre.float().contiguous(), + grad_post.float().contiguous(), + grad_comb.float().contiguous(), + grad_mixes, + grad_z, + scratch, + tokens, + eps, + HC=hc, + N=mixes.shape[1], + ITERS=iters, + BLOCK_T=BLOCK_T, + ) + return grad_mixes, grad_z + + +@mhc_gates_backward.register_fake +def _mhc_gates_backward_fake(mixes, scale, base, grad_pre, grad_post, grad_comb, hc, iters, eps): + return torch.empty_like(mixes, dtype=torch.float32), torch.empty_like(mixes, dtype=torch.float32) + + +def _setup_context(ctx, inputs, output) -> None: + mixes, scale, base, hc, iters, eps = inputs + ctx.save_for_backward(mixes, scale, base) + ctx.hc, ctx.iters, ctx.eps = hc, iters, eps + + +def _backward(ctx, grad_pre, grad_post, grad_comb): + mixes, scale, base = ctx.saved_tensors + hc = ctx.hc + grad_mixes, grad_z = mhc_gates_backward( + mixes, scale, base, grad_pre, grad_post, grad_comb, hc, ctx.iters, ctx.eps + ) + grad_base = grad_z.sum(0) + grad_scale = torch.stack( + [part.sum() for part in (grad_z * mixes.float()).split([hc, hc, hc * hc], dim=1)] + ) + return grad_mixes.to(mixes.dtype), grad_scale.to(scale.dtype), grad_base.to(base.dtype), None, None, None + + +mhc_gates_forward.register_autograd(_backward, setup_context=_setup_context) + + +def mhc_gates( + mixes: torch.Tensor, scale: torch.Tensor, base: torch.Tensor, hc: int, iters: int, eps: float +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """`(pre, post, comb)` gates from `(..., (2 + hc) * hc)` mixes, all fp32.""" + lead = mixes.shape[:-1] + pre, post, comb = mhc_gates_forward(mixes.reshape(-1, mixes.shape[-1]), scale, base, hc, iters, eps) + return pre.view(*lead, hc), post.view(*lead, hc), comb.view(*lead, hc, hc) diff --git a/prime_kernels/mhc_projection/projection.py b/prime_kernels/mhc_projection/projection.py new file mode 100644 index 0000000..d2cc8df --- /dev/null +++ b/prime_kernels/mhc_projection/projection.py @@ -0,0 +1,450 @@ +"""Fused mHC projection: RMS over the flattened streams, projection to the gate logits, and the +stream collapse with the previous sublayer's `pre` gate, in one pass over the streams. + +For streams `x` of shape `(T, hc, d)` (flattened to `K = hc * d` per token), weight `w` of shape +`(n, K)` and an optional `pre_mix` of shape `(T, hc)`: + + rstd = rsqrt(mean(x ** 2) + eps) per token, fp32 + mixes = (x @ w.T) * rstd (T, n), fp32 + collapsed = sum_i pre_mix[:, i] * x[:, i, :] (T, d), x.dtype + +The projection runs on bf16 operands with fp32 accumulation (the weight is rounded to bf16 as +prime-rl's eager path does); everything else is fp32. The backward reads the streams once more +and writes `dx = (dmixes * rstd) @ w + coef * x + pre_mix * dcollapsed`, the weight gradient as +fixed-order split-token partials (deterministic), and the `pre_mix` gradient. Its GEMM operand +`dmixes * rstd` is rounded to bf16, again as prime-rl's eager path does. +""" + +import torch +import triton +import triton.language as tl +from triton.tools.tensor_descriptor import TensorDescriptor + +# Gate logits are padded to this width for the tensor-core tiles. +BLOCK_N = 32 + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_M": bm, "BLOCK_D": bd}, num_warps=nw, num_stages=ns) + for bm in (32, 64) + for bd in (64, 128, 256) + for nw in (4, 8) + for ns in (2, 3, 4) + if bm * bd <= 64 * 128 + ], + key=["HC", "D", "HAS_PRE"], +) +@triton.jit +def _projection_fwd_kernel( + x_ptr, + w_ptr, + pre_ptr, + mixes_ptr, + rstd_ptr, + collapsed_ptr, + T, + eps, + HC: tl.constexpr, + D: tl.constexpr, + N: tl.constexpr, + HAS_PRE: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, +): + K: tl.constexpr = HC * D + rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + row_mask = rows < T + rows64 = rows.to(tl.int64) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_D) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + sumsq = tl.zeros((BLOCK_M,), dtype=tl.float32) + for d0 in range(0, D, BLOCK_D): + d = d0 + offs_d + d_mask = d < D + tile_mask = row_mask[:, None] & d_mask[None, :] + collapsed = tl.zeros((BLOCK_M, BLOCK_D), dtype=tl.float32) + # Tensor-core accumulation loses precision over long K chains, so each d-block's partial + # product is promoted into `acc` with ordinary fp32 adds. + part = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for i in tl.static_range(HC): + k = i * D + d + x = tl.load(x_ptr + rows64[:, None] * K + k[None, :], mask=tile_mask, other=0.0) + # The weight is padded to BLOCK_N rows by the caller. + w = tl.load(w_ptr + offs_n[:, None] * K + k[None, :], mask=d_mask[None, :], other=0.0) + part = tl.dot(x, tl.trans(w), part) + xf = x.to(tl.float32) + sumsq += tl.sum(xf * xf, axis=1) + if HAS_PRE: + pre = tl.load(pre_ptr + rows * HC + i, mask=row_mask, other=0.0) + collapsed += pre[:, None] * xf + acc += part + if HAS_PRE: + tl.store( + collapsed_ptr + rows64[:, None] * D + d[None, :], + collapsed.to(collapsed_ptr.dtype.element_ty), + mask=tile_mask, + ) + + rstd = tl.rsqrt(sumsq / K + eps) + tl.store(rstd_ptr + rows, rstd, mask=row_mask) + tl.store( + mixes_ptr + rows[:, None] * N + offs_n[None, :], + acc * rstd[:, None], + mask=row_mask[:, None] & (offs_n[None, :] < N), + ) + + +@triton.jit +def _projection_bwd_prep_kernel( + grad_mixes_ptr, + mixes_ptr, + rstd_ptr, + grad_proj_ptr, + coef_ptr, + T, + K: tl.constexpr, + N: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_M: tl.constexpr, +): + """Per-token backward factors: `grad_proj = dmixes * rstd` (bf16, padded to BLOCK_N) and + `coef = -rstd^2 * sum(dmixes * mixes) / K`, the factor rstd's dependence on x contributes.""" + rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + row_mask = rows < T + offs_n = tl.arange(0, BLOCK_N) + mask = row_mask[:, None] & (offs_n[None, :] < N) + grad_mixes = tl.load(grad_mixes_ptr + rows[:, None] * N + offs_n[None, :], mask=mask, other=0.0) + mixes = tl.load(mixes_ptr + rows[:, None] * N + offs_n[None, :], mask=mask, other=0.0) + rstd = tl.load(rstd_ptr + rows, mask=row_mask, other=0.0) + tl.store( + grad_proj_ptr + rows[:, None] * BLOCK_N + offs_n[None, :], + (grad_mixes * rstd[:, None]).to(tl.bfloat16), + mask=row_mask[:, None], + ) + tl.store(coef_ptr + rows, -rstd * rstd * tl.sum(grad_mixes * mixes, axis=1) / K, mask=row_mask) + + +def _set_bwd_block_shapes(nargs): + block_m, block_d, block_n = nargs["BLOCK_M"], nargs["BLOCK_D"], nargs["BLOCK_N"] + nargs["x_desc"].block_shape = [block_m, 1, block_d] + nargs["grad_x_desc"].block_shape = [block_m, 1, block_d] + nargs["w_desc"].block_shape = [block_n, 1, block_d] + nargs["grad_proj_desc"].block_shape = [block_m, block_n] + nargs["grad_collapsed_desc"].block_shape = [block_m, block_d] + nargs["grad_x_addend_desc"].block_shape = [block_m, 1, block_d] + + +@triton.autotune( + configs=[ + triton.Config( + {"BLOCK_M": 64, "BLOCK_D": bd, "SPLITS": s}, + num_warps=nw, + num_stages=ns, + pre_hook=_set_bwd_block_shapes, + ) + for bd in (64, 128) + for s in (8, 16) + for nw in (4, 8) + for ns in (2, 3) + ], + key=["HC", "D", "HAS_PRE", "HAS_ADDEND"], +) +@triton.jit +def _projection_bwd_kernel( + x_desc, + w_desc, + pre_ptr, + grad_proj_desc, + coef_ptr, + grad_collapsed_desc, + grad_x_addend_desc, + grad_x_desc, + grad_w_partial_ptr, + grad_pre_partial_ptr, + T, + HC: tl.constexpr, + D: tl.constexpr, + HAS_PRE: tl.constexpr, + HAS_ADDEND: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, + SPLITS: tl.constexpr, +): + """Grid: (cdiv(D, BLOCK_D) * HC, SPLITS). One stream's d-block over a contiguous token range.""" + K: tl.constexpr = HC * D + pid = tl.program_id(0) + split = tl.program_id(1) + # The stream varies fastest so the HC programs sharing a `grad_collapsed` tile run together. + i = pid % HC + d_block = pid // HC + d0 = d_block * BLOCK_D + + w = w_desc.load([0, i, d0]).reshape(BLOCK_N, BLOCK_D) + grad_w = tl.zeros((BLOCK_D, BLOCK_N), dtype=tl.float32) + diagonal = tl.arange(0, BLOCK_M)[:, None] == tl.arange(0, BLOCK_M)[None, :] + + m_tiles = tl.cdiv(T, BLOCK_M) + per_split = tl.cdiv(m_tiles, SPLITS) + begin = split * per_split + end = tl.minimum(begin + per_split, m_tiles) + for m_tile in range(begin, end): + m0 = m_tile * BLOCK_M + rows = m0 + tl.arange(0, BLOCK_M) + row_mask = rows < T + grad_proj = grad_proj_desc.load([m0, 0]) + coef = tl.load(coef_ptr + rows, mask=row_mask, other=0.0) + x = x_desc.load([m0, i, d0]).reshape(BLOCK_M, BLOCK_D) + if HAS_PRE: + grad_collapsed = grad_collapsed_desc.load([m0, d0]) + pre = tl.load(pre_ptr + rows * HC + i, mask=row_mask, other=0.0) + xf = x.to(tl.float32) + grad_x = tl.dot(grad_proj, w) + coef[:, None] * xf + if HAS_PRE: + grad_x += pre[:, None] * grad_collapsed.to(tl.float32) + # The row-wise dot products of grad_collapsed and x are the diagonal of + # grad_collapsed @ x.T. The tensor cores compute it exactly (bf16 products, fp32 + # sums) and much faster than an elementwise product and cross-thread reduction. + gram = tl.dot(grad_collapsed, tl.trans(x)) + grad_pre = tl.sum(tl.where(diagonal, gram, 0.0), axis=1) + tl.store(grad_pre_partial_ptr + (d_block * HC + i) * T + rows, grad_pre, mask=row_mask) + grad_x = grad_x.to(grad_x_desc.dtype) + if HAS_ADDEND: + # Rounded first, then summed as autograd sums two bf16 gradients: bit for bit the same. + addend = grad_x_addend_desc.load([m0, i, d0]).reshape(BLOCK_M, BLOCK_D) + grad_x = (grad_x.to(tl.float32) + addend.to(tl.float32)).to(grad_x_desc.dtype) + grad_x_desc.store([m0, i, d0], grad_x.reshape(BLOCK_M, 1, BLOCK_D)) + grad_w += tl.dot(tl.trans(x), grad_proj) + + offs_d = d0 + tl.arange(0, BLOCK_D) + offs_n = tl.arange(0, BLOCK_N) + tl.store( + grad_w_partial_ptr + (split * BLOCK_N + offs_n[None, :]) * K + i * D + offs_d[:, None], + grad_w, + mask=(offs_d < D)[:, None], + ) + + +def _aligned(t: torch.Tensor) -> torch.Tensor: + """Contiguous with a 16-byte aligned base, as the TMA descriptors require.""" + t = t.contiguous() + return t if t.data_ptr() % 16 == 0 else t.clone() + + +def _padded_weight(weight: torch.Tensor) -> torch.Tensor: + padded = weight.new_zeros(BLOCK_N, weight.shape[1], dtype=torch.bfloat16) + padded[: weight.shape[0]] = weight + return padded + + +def _check(x: torch.Tensor, weight: torch.Tensor, pre_mix: torch.Tensor | None) -> None: + if x.dim() != 3: + raise ValueError(f"expected streams of shape (tokens, hc, d), got {tuple(x.shape)}") + if x.dtype != torch.bfloat16: + raise ValueError(f"expected bf16 streams, got {x.dtype}") + tokens, hc, d = x.shape + if weight.dim() != 2 or weight.shape[1] != hc * d: + raise ValueError(f"expected a weight of shape (n, {hc * d}), got {tuple(weight.shape)}") + if weight.shape[0] > BLOCK_N: + raise ValueError(f"at most {BLOCK_N} gate logits are supported, got {weight.shape[0]}") + if d % 8: + raise ValueError(f"d must be a multiple of 8 (16-byte rows for TMA), got {d}") + if pre_mix is not None and tuple(pre_mix.shape) != (tokens, hc): + raise ValueError(f"expected pre_mix of shape ({tokens}, {hc}), got {tuple(pre_mix.shape)}") + + +@torch.library.custom_op("prime_kernels::mhc_projection_forward", mutates_args=()) +def mhc_projection_forward( + x: torch.Tensor, weight: torch.Tensor, pre_mix: torch.Tensor | None, eps: float +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Returns `(mixes, rstd, collapsed)`; `collapsed` is empty when `pre_mix` is None.""" + _check(x, weight, pre_mix) + tokens, hc, d = x.shape + n = weight.shape[0] + x = x.contiguous() + w = _padded_weight(weight) + mixes = x.new_empty(tokens, n, dtype=torch.float32) + rstd = x.new_empty(tokens, dtype=torch.float32) + has_pre = pre_mix is not None + collapsed = x.new_empty(tokens, d) if has_pre else x.new_empty(0) + pre = pre_mix.float().contiguous() if has_pre else mixes + if tokens > 0: + grid = lambda meta: (triton.cdiv(tokens, meta["BLOCK_M"]),) + _projection_fwd_kernel[grid]( + x, w, pre, mixes, rstd, collapsed, tokens, eps, HC=hc, D=d, N=n, HAS_PRE=has_pre, BLOCK_N=BLOCK_N + ) + return mixes, rstd, collapsed + + +@mhc_projection_forward.register_fake +def _mhc_projection_forward_fake(x, weight, pre_mix, eps): + _check(x, weight, pre_mix) + tokens, hc, d = x.shape + mixes = x.new_empty(tokens, weight.shape[0], dtype=torch.float32) + rstd = x.new_empty(tokens, dtype=torch.float32) + collapsed = x.new_empty(tokens, d) if pre_mix is not None else x.new_empty(0) + return mixes, rstd, collapsed + + +@torch.library.custom_op("prime_kernels::mhc_projection_backward", mutates_args=()) +def mhc_projection_backward( + x: torch.Tensor, + weight: torch.Tensor, + pre_mix: torch.Tensor | None, + mixes: torch.Tensor, + rstd: torch.Tensor, + grad_mixes: torch.Tensor, + grad_collapsed: torch.Tensor | None, + grad_x_addend: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Returns `(grad_x, grad_weight, grad_pre_mix)`; `grad_pre_mix` is empty without `pre_mix`. + + `grad_weight` is fp32 regardless of the weight dtype. `grad_x_addend`, the gradient `x` receives + from its other readers, is added into `grad_x` (rounded as two separate bf16 gradients would be). + """ + tokens, hc, d = x.shape + n = weight.shape[0] + x = _aligned(x) + has_pre = pre_mix is not None and grad_collapsed is not None + grad_x = torch.empty_like(x) + if tokens == 0: + grad_weight = x.new_zeros(n, hc * d, dtype=torch.float32) + grad_pre_mix = x.new_zeros(0 if pre_mix is None else (0, hc), dtype=torch.float32) + return grad_x, grad_weight, grad_pre_mix + + grad_proj = x.new_empty(tokens, BLOCK_N, dtype=torch.bfloat16) + coef = x.new_empty(tokens, dtype=torch.float32) + _projection_bwd_prep_kernel[(triton.cdiv(tokens, 64),)]( + grad_mixes.float().contiguous(), mixes, rstd, grad_proj, coef, tokens, K=hc * d, N=n, BLOCK_N=BLOCK_N, BLOCK_M=64 + ) + pre = pre_mix.float().contiguous() if has_pre else coef + grad_collapsed = _aligned(grad_collapsed) if has_pre else x.view(-1, d) + has_addend = grad_x_addend is not None + grad_x_addend = _aligned(grad_x_addend.view(x.shape)) if has_addend else x + # Sized for the smallest BLOCK_D and the largest SPLITS the autotuner may pick; only the + # slots the chosen config writes are summed below. + max_d_blocks, max_splits = triton.cdiv(d, 64), 16 + grad_w_partial = x.new_empty(max_splits, BLOCK_N, hc * d, dtype=torch.float32) + grad_pre_partial = x.new_empty(max_d_blocks, hc, tokens, dtype=torch.float32) if has_pre else coef + # Block shapes are placeholders: the autotuner's pre-hook sets them per config. + _projection_bwd_kernel[lambda meta: (triton.cdiv(d, meta["BLOCK_D"]) * hc, meta["SPLITS"])]( + TensorDescriptor.from_tensor(x, [1, 1, 8]), + TensorDescriptor.from_tensor(_padded_weight(weight).view(BLOCK_N, hc, d), [1, 1, 8]), + pre, + TensorDescriptor.from_tensor(grad_proj, [1, 8]), + coef, + TensorDescriptor.from_tensor(grad_collapsed, [1, 8]), + TensorDescriptor.from_tensor(grad_x_addend, [1, 1, 8]), + TensorDescriptor.from_tensor(grad_x, [1, 1, 8]), + grad_w_partial, + grad_pre_partial, + tokens, + HC=hc, + D=d, + HAS_PRE=has_pre, + HAS_ADDEND=has_addend, + BLOCK_N=BLOCK_N, + ) + config = _projection_bwd_kernel.best_config.kwargs + # Fixed-order sums over the partials: deterministic. + grad_weight = grad_w_partial[: config["SPLITS"], :n].sum(0) + if has_pre: + grad_pre_mix = grad_pre_partial[: triton.cdiv(d, config["BLOCK_D"])].sum(0).t().contiguous() + elif pre_mix is not None: + grad_pre_mix = torch.zeros_like(pre_mix, dtype=torch.float32) + else: + grad_pre_mix = x.new_empty(0, dtype=torch.float32) + return grad_x, grad_weight, grad_pre_mix + + +@mhc_projection_backward.register_fake +def _mhc_projection_backward_fake(x, weight, pre_mix, mixes, rstd, grad_mixes, grad_collapsed, grad_x_addend=None): + grad_pre_mix = ( + x.new_empty(pre_mix.shape, dtype=torch.float32) if pre_mix is not None else x.new_empty(0, dtype=torch.float32) + ) + return torch.empty_like(x), x.new_empty(weight.shape, dtype=torch.float32), grad_pre_mix + + +def _setup_context(ctx, inputs, output) -> None: + x, weight, pre_mix, _ = inputs + mixes, rstd, _ = output + ctx.save_for_backward(x, weight, pre_mix, mixes, rstd) + ctx.mark_non_differentiable(rstd) + ctx.has_pre = pre_mix is not None + + +def _backward(ctx, grad_mixes, grad_rstd, grad_collapsed): + x, weight, pre_mix, mixes, rstd = ctx.saved_tensors + if grad_mixes is None: + grad_mixes = torch.zeros_like(mixes) + grad_x, grad_weight, grad_pre_mix = mhc_projection_backward( + x, weight, pre_mix, mixes, rstd, grad_mixes, grad_collapsed if ctx.has_pre else None + ) + grad_pre_mix = grad_pre_mix.to(pre_mix.dtype) if ctx.has_pre else None + return grad_x, grad_weight.to(weight.dtype), grad_pre_mix, None + + +mhc_projection_forward.register_autograd(_backward, setup_context=_setup_context) + + +def mhc_projection( + x: torch.Tensor, weight: torch.Tensor, pre_mix: torch.Tensor | None = None, eps: float = 1e-20 +) -> tuple[torch.Tensor, torch.Tensor | None]: + """RMS-normalized projection of the flattened streams, plus the optional `pre_mix` collapse. + + `x`: `(..., hc, d)` bf16 streams. `weight`: `(n, hc * d)`, `n <= 32`. `pre_mix`: `(..., hc)`. + Returns `mixes` `(..., n)` fp32 and `collapsed` `(..., d)` in `x.dtype` (None without `pre_mix`). + """ + lead = x.shape[:-2] + hc, d = x.shape[-2:] + flat_pre = pre_mix.reshape(-1, hc) if pre_mix is not None else None + mixes, _, collapsed = mhc_projection_forward(x.reshape(-1, hc, d), weight, flat_pre, eps) + mixes = mixes.view(*lead, weight.shape[0]) + return mixes, collapsed.view(*lead, d) if pre_mix is not None else None + + +class _ProjectionWithStreams(torch.autograd.Function): + """`mhc_projection_forward` that also hands `x` on, as a view, to its other readers. + + The gradient those readers send back reaches this backward as the view's gradient and is added + while `grad_x` is written, instead of by autograd in a separate pass over the streams. + """ + + @staticmethod + def forward(ctx, x, weight, pre_mix, eps): + mixes, rstd, collapsed = mhc_projection_forward(x, weight, pre_mix, eps) + ctx.save_for_backward(x, weight, pre_mix, mixes, rstd) + ctx.has_pre = pre_mix is not None + return mixes, collapsed, x.view_as(x) + + @staticmethod + def backward(ctx, grad_mixes, grad_collapsed, grad_streams): + x, weight, pre_mix, mixes, rstd = ctx.saved_tensors + if grad_mixes is None: + grad_mixes = torch.zeros_like(mixes) + grad_x, grad_weight, grad_pre_mix = mhc_projection_backward( + x, weight, pre_mix, mixes, rstd, grad_mixes, grad_collapsed if ctx.has_pre else None, grad_streams + ) + grad_pre_mix = grad_pre_mix.to(pre_mix.dtype) if ctx.has_pre else None + return grad_x, grad_weight.to(weight.dtype), grad_pre_mix, None + + +def mhc_projection_streams( + x: torch.Tensor, weight: torch.Tensor, pre_mix: torch.Tensor | None = None, eps: float = 1e-20 +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + """`mhc_projection` plus `x` itself (a view) for the streams' other readers. + + Reading the streams through that view rather than `x` folds their gradient into the projection's + backward, which saves autograd's separate sum of the two stream gradients (a full read of both + and a write). Values and gradients are bit for bit those of `mhc_projection`. + """ + lead = x.shape[:-2] + hc, d = x.shape[-2:] + flat_pre = pre_mix.reshape(-1, hc) if pre_mix is not None else None + mixes, collapsed, streams = _ProjectionWithStreams.apply(x.reshape(-1, hc, d), weight, flat_pre, eps) + mixes = mixes.view(*lead, weight.shape[0]) + return mixes, collapsed.view(*lead, d) if pre_mix is not None else None, streams.view(x.shape) diff --git a/prime_kernels/moe_experts/__init__.py b/prime_kernels/moe_experts/__init__.py new file mode 100644 index 0000000..6dd7305 --- /dev/null +++ b/prime_kernels/moe_experts/__init__.py @@ -0,0 +1,3 @@ +from prime_kernels.moe_experts.ops import TOKEN_GROUP_ALIGNMENT, moe_experts, unsupported_shape_reason + +__all__ = ["TOKEN_GROUP_ALIGNMENT", "moe_experts", "unsupported_shape_reason"] diff --git a/prime_kernels/moe_experts/fp8.py b/prime_kernels/moe_experts/fp8.py new file mode 100644 index 0000000..c982056 --- /dev/null +++ b/prime_kernels/moe_experts/fp8.py @@ -0,0 +1,389 @@ +"""FP8 expert MLP: DeepGEMM's blockwise FP8 grouped GEMMs with fused quantization passes. + +Numerics follow DeepSeek-V3: e4m3 values, fp32 scales, fp32 accumulation; activations and +gradients get one scale per 1 x 128 group along the GEMM's K dimension, weights one per +128 x 128 block. That holds for the forward, the data gradients and the weight gradients (whose +K dimension is the tokens, so their operands are quantized per 128 tokens of each channel). + +DeepGEMM's grouped GEMMs pick each 128-row block's expert from the block's first row, so the +experts' rows must start at multiples of 128. Two layouts make that hold without copies: + +- The *padded* layout: `TileTable` tile ``t`` (up to 128 rows of one expert) is padded rows + ``[128 t, 128 t + 128)``. The quantization of ``x`` / ``dout`` writes into it, so the gate/up + and dh GEMMs read it and write their outputs in it, which the SwiGLU passes read. The K-grouped + weight-gradient operands use the same padded rows, transposed. +- The *token-major* layout with a head GEMM, for the down and dx GEMMs, whose outputs are what + the caller sees: their operands are written at the tokens' own rows, so DeepGEMM writes + ``out`` / ``dx`` in place. Only an expert's *head rows*, those before the first multiple of 128 + in its group, sit in a block that starts in the previous expert and come out wrong; the SwiGLU + passes also copy them, at most 127 per expert, into one 128-row block per expert, a small + second GEMM computes those, and they are scattered over the wrong ones. + +Passes, forward: quantize ``x`` per row group and, for the weight gradient, per column group into +the K-grouped layout, from one read of ``x``; quantize the weights per block, also writing the +transposed copies the data gradients use; gate/up GEMM; clamped SwiGLU + quantization of ``h``; +down GEMM (+ head GEMM and scatter). Backward: quantize ``dout`` like ``x``; dh GEMM; SwiGLU +backward, writing FP8 ``dgate | dup`` per row group and per column group and FP8 ``h`` per column +group; dx GEMM (+ head GEMM and scatter); the two K-grouped weight-gradient GEMMs. The passes +over activation tiles are the Gluon kernels of `quant`. +""" + +import torch +import triton +import triton.language as tl + +from prime_kernels.moe_experts import quant +from prime_kernels.moe_experts.kernels import TileTable, zero_tail + +GROUP = 128 + + +def padded_rows(num_rows: int, num_experts: int) -> int: + """Rows of the padded layout: every tile `TileTable` can produce, ``GROUP`` rows each.""" + return (triton.cdiv(num_rows, GROUP) + num_experts) * GROUP + + +def unsupported_shape_reason(hidden_size: int, intermediate_size: int) -> str | None: + if hidden_size % GROUP or intermediate_size % GROUP: + return f"FP8 needs hidden ({hidden_size}) and intermediate ({intermediate_size}) sizes that are multiples of {GROUP}" + return None + + +@triton.jit +def _head_rows(expert_start, count): + """How many of an expert's first rows lie before the first multiple of 128 in its group.""" + return tl.minimum((128 - expert_start % 128) % 128, count) + + +@triton.jit +def _layouts_kernel( + padded_layout_ptr, + token_layout_ptr, + head_layout_ptr, + tile_expert_ptr, + tile_row_start_ptr, + tile_row_end_ptr, + expert_row_start_ptr, + counts_ptr, + num_tiles_ptr, + max_tiles, + num_experts, +): + """DeepGEMM's per-row expert indices (-1: no expert) of the padded and the token-major layouts + (from tile ``t``) and of the head operand (expert ``t``). Token-major rows after the last + group must be filled with -1 beforehand. A tile's padding rows get its expert, so DeepGEMM + computes them, from the zeros the quantization writes there: the padded GEMM outputs are then + zeros in those rows rather than never written, and the SwiGLU backward needs no masks.""" + t = tl.program_id(0) + rows = tl.arange(0, 128) + if t < max_tiles: + expert = tl.load(tile_expert_ptr + t) + row_start = tl.load(tile_row_start_ptr + t) + in_tile = t < tl.load(num_tiles_ptr) + tl.store(padded_layout_ptr + t * 128 + rows, tl.where(in_tile, expert, -1)) + valid = in_tile & (rows < tl.load(tile_row_end_ptr + t) - row_start) + tl.store(token_layout_ptr + row_start + rows, expert, mask=valid) + if t < num_experts: + head = _head_rows(tl.load(expert_row_start_ptr + t), tl.load(counts_ptr + t).to(tl.int32)) + tl.store(head_layout_ptr + t * 128 + rows, tl.where(rows < head, t, -1)) + + +@triton.jit +def _scatter_head_kernel(head_out_ptr, out_ptr, expert_row_start_ptr, counts_ptr, N, BLOCK_N: tl.constexpr): + """Copy expert ``e``'s head rows from the head GEMM's output over the token-major output.""" + e = tl.program_id(0) + start = tl.load(expert_row_start_ptr + e) + head = _head_rows(start, tl.load(counts_ptr + e).to(tl.int32)) + if head == 0: + return + rows = tl.arange(0, 128) + cols = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + mask = (rows < head)[:, None] & (cols < N)[None, :] + v = tl.load(head_out_ptr + (e * 128 + rows).to(tl.int64)[:, None] * N + cols[None, :], mask=mask) + tl.store(out_ptr + (start + rows).to(tl.int64)[:, None] * N + cols[None, :], v, mask=mask) + + +@triton.jit +def _transpose_fp8(q): + # Through 16 bits (exact both ways): Triton 3.7 scrambles a quarter of the elements of 8-bit + # transposes. + return tl.trans(q.to(tl.float16)).to(tl.float8e4nv) + + +@triton.jit +def _quantize_weight_kernel( + w_ptr, + w2_ptr, + q_ptr, + sf_ptr, + qt_ptr, + sf_t_ptr, + stride_e, + stride_e2, + R, + C, + SPLIT_BLOCKS, + TRANSPOSED: tl.constexpr, +): + """One 128 x 128 block of an ``[E, R, C]`` weight (row blocks from ``SPLIT_BLOCKS`` on read + from ``w2``): FP8 values and block scale, and with ``TRANSPOSED`` the ``[E, C, R]`` copy.""" + e = tl.program_id(0).to(tl.int64) + rb = tl.program_id(1) + cb = tl.program_id(2) + rows = tl.arange(0, 128) + cols = cb * 128 + tl.arange(0, 128) + if rb < SPLIT_BLOCKS: + src = w_ptr + e * stride_e + (rb * 128 + rows)[:, None] * C + cols[None, :] + else: + src = w2_ptr + e * stride_e2 + ((rb - SPLIT_BLOCKS) * 128 + rows)[:, None] * C + cols[None, :] + v = tl.load(src).to(tl.float32) + scale = tl.maximum(tl.max(tl.abs(v)), 1e-4) / 448.0 + q = (v * (1.0 / scale)).to(tl.float8e4nv) + num_rb = R // 128 + num_cb = C // 128 + tl.store(q_ptr + e * R * C + (rb * 128 + rows)[:, None] * C + cols[None, :], q) + tl.store(sf_ptr + (e * num_rb + rb) * num_cb + cb, scale) + if TRANSPOSED: + tl.store(qt_ptr + e * R * C + cols[:, None] * R + (rb * 128 + rows)[None, :], _transpose_fp8(q)) + tl.store(sf_t_ptr + (e * num_cb + cb) * num_rb + rb, scale) + + +@triton.jit +def _weight_grad_cast_kernel(d_ptr, a_ptr, b_ptr, SPLIT, REST, BLOCK: tl.constexpr): + """bf16 of an fp32 weight gradient, ``SPLIT + REST`` elements per expert: the first ``SPLIT`` + into ``a``, the rest into ``b``.""" + e = tl.program_id(0).to(tl.int64) + offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + total = SPLIT + REST + mask = offsets < total + v = tl.load(d_ptr + e * total + offsets, mask=mask).to(tl.bfloat16) + tl.store(a_ptr + e * SPLIT + offsets, v, mask=offsets < SPLIT) + tl.store(b_ptr + e * REST + offsets - SPLIT, v, mask=mask & (offsets >= SPLIT)) + + +def quantize_weight( + weight: torch.Tensor, weight2: torch.Tensor | None, transposed: bool +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """FP8 ``[E, R, C]`` blocks of ``weight`` (stacked on top of ``weight2`` along R when given) + and their scales, plus the ``[E, C, R]`` transposed copy when asked (empty otherwise).""" + num_experts, rows, cols = weight.shape + for w in (weight, weight2): + if w is not None and (w.stride(2) != 1 or w.stride(1) != cols): + raise ValueError(f"expert weights must have contiguous rows, got strides {w.stride()}") + R = rows + (weight2.shape[1] if weight2 is not None else 0) + q = weight.new_empty(num_experts, R, cols, dtype=torch.float8_e4m3fn) + sf = weight.new_empty(num_experts, R // GROUP, cols // GROUP, dtype=torch.float32) + if transposed: + qt = weight.new_empty(num_experts, cols, R, dtype=torch.float8_e4m3fn) + sf_t = weight.new_empty(num_experts, cols // GROUP, R // GROUP, dtype=torch.float32) + else: + qt = q.new_empty(0) + sf_t = sf.new_empty(0) + w2 = weight2 if weight2 is not None else weight + _quantize_weight_kernel[(num_experts, R // GROUP, cols // GROUP)]( + weight, + w2, + q, + sf, + qt if transposed else q, + sf_t if transposed else sf, + weight.stride(0), + w2.stride(0), + R, + cols, + rows // GROUP if weight2 is not None else R // GROUP, + TRANSPOSED=transposed, + num_warps=8, + ) + return q, sf, qt, sf_t + + +class Layout: + """One call's groups: the `TileTable` of 128-row tiles and DeepGEMM's per-row expert indices + of the padded layout, the token-major layout and the head operand (see the module doc).""" + + def __init__(self, num_tokens_per_expert: torch.Tensor, num_rows: int) -> None: + device = num_tokens_per_expert.device + num_experts = num_tokens_per_expert.numel() + self.table = TileTable(num_tokens_per_expert, num_rows, GROUP) + self.counts = num_tokens_per_expert + self.num_rows = num_rows + self.max_tiles = self.table.max_tiles + self.Mp = self.max_tiles * GROUP + self.M_head = num_experts * GROUP + self.padded_layout = torch.empty(self.Mp, dtype=torch.int32, device=device) + self.token_layout = torch.full((num_rows,), -1, dtype=torch.int32, device=device) + self.head_layout = torch.empty(self.M_head, dtype=torch.int32, device=device) + _layouts_kernel[(max(self.max_tiles, num_experts),)]( + self.padded_layout, + self.token_layout, + self.head_layout, + *self.tile_args(), + self.max_tiles, + num_experts, + ) + # K of each expert in the K-grouped GEMMs: its rows padded to whole groups. + self.ks = ((num_tokens_per_expert + GROUP - 1) // GROUP * GROUP).to(torch.int32) + + def tile_args(self): + t = self.table + return t.expert, t.row_start, t.row_end, t.expert_row_start, self.counts, t.num_tiles + + def check_columns(self, cols: int) -> None: + """The tile kernels address within an expert's K-grouped block with 32-bit offsets.""" + if self.Mp * cols >= 2**31: + raise ValueError(f"too many rows ({self.num_rows}) for {cols} columns in the FP8 path") + + +def quantize_activation(src: torch.Tensor, layout: Layout, columns: bool): + """``src`` [rows, C] per row group into the padded layout ([Mp, C] FP8, [C / 128, Mp] scales) + and, with ``columns``, per column group into the K-grouped layout (flat FP8 and + [Mp / 128, C] scales; empty without).""" + C = src.shape[1] + Mp = layout.Mp + layout.check_columns(C) + q = src.new_empty(Mp, C, dtype=torch.float8_e4m3fn) + q_sf = src.new_empty(C // GROUP, Mp, dtype=torch.float32) + if columns: + qt = src.new_empty(Mp * C, dtype=torch.float8_e4m3fn) + qt_sf = src.new_empty(Mp // GROUP, C, dtype=torch.float32) + else: + qt, qt_sf = q.new_empty(0), q_sf.new_empty(0, C) + quant.quantize_kernel[(layout.max_tiles, C // GROUP)]( + quant.descriptor(src), + q, + q_sf, + qt if columns else q, + qt_sf if columns else q_sf, + *layout.tile_args(), + Mp, + C, + COLUMNS=columns, + num_warps=quant.NUM_WARPS, + ) + return q, q_sf, qt, qt_sf + + +def _head_operand(layout: Layout, cols: int, like: torch.Tensor): + q = like.new_empty(layout.M_head, cols, dtype=torch.float8_e4m3fn) + sf = like.new_empty(cols // GROUP, layout.M_head, dtype=torch.float32) + return q, sf + + +def swiglu_quantize(gate_up: torch.Tensor, layout: Layout, limit: float): + """FP8 h in the token-major layout and its head operand, each with MN-major scales.""" + inter = gate_up.shape[1] // 2 + M = layout.num_rows + h = gate_up.new_empty(M, inter, dtype=torch.float8_e4m3fn) + h_sf = gate_up.new_empty(inter // GROUP, M, dtype=torch.float32) + head, head_sf = _head_operand(layout, inter, gate_up) + quant.swiglu_kernel[(layout.max_tiles, inter // GROUP)]( + quant.descriptor(gate_up), + h, + h_sf, + head, + head_sf, + *layout.tile_args(), + M, + layout.M_head, + inter, + limit, + num_warps=quant.NUM_WARPS, + ) + return h, h_sf, head, head_sf + + +def swiglu_backward_quantize(dh: torch.Tensor, gate_up: torch.Tensor, layout: Layout, limit: float): + """FP8 dgate | dup in the token-major layout + head operand (for dx) and in the K-grouped + layout, and FP8 h in the K-grouped layout.""" + Mp, inter = dh.shape + M = layout.num_rows + layout.check_columns(2 * inter) + dgu = dh.new_empty(M, 2 * inter, dtype=torch.float8_e4m3fn) + dgu_sf = dh.new_empty(2 * inter // GROUP, M, dtype=torch.float32) + head, head_sf = _head_operand(layout, 2 * inter, dh) + dgu_t = dh.new_empty(Mp * 2 * inter, dtype=torch.float8_e4m3fn) + dgu_t_sf = dh.new_empty(Mp // GROUP, 2 * inter, dtype=torch.float32) + h_t = dh.new_empty(Mp * inter, dtype=torch.float8_e4m3fn) + h_t_sf = dh.new_empty(Mp // GROUP, inter, dtype=torch.float32) + quant.swiglu_backward_kernel[(layout.max_tiles, inter // GROUP)]( + quant.descriptor(dh), + quant.descriptor(gate_up), + dgu, + dgu_sf, + head, + head_sf, + dgu_t, + dgu_t_sf, + h_t, + h_t_sf, + *layout.tile_args(), + M, + layout.M_head, + inter, + limit, + num_warps=quant.NUM_WARPS, + ) + return (dgu, dgu_sf, head, head_sf), (dgu_t, dgu_t_sf), (h_t, h_t_sf) + + +def _m_grouped(a, a_sf, b, b_sf, d, grouped_layout) -> None: + import deep_gemm + + deep_gemm.m_grouped_fp8_gemm_nt_contiguous((a, a_sf.T), (b, b_sf), d, grouped_layout) + + +def padded_gemm(a, a_sf, b, b_sf, layout: Layout, out_cols: int) -> torch.Tensor: + """``[Mp, out_cols]`` bf16 of ``a @ b[e].T`` over the padded layout; ``a``'s scales MN-major.""" + d = torch.empty(layout.Mp, out_cols, device=a.device, dtype=torch.bfloat16) + _m_grouped(a, a_sf, b, b_sf, d, layout.padded_layout) + return d + + +def token_gemm(a, a_sf, head, head_sf, b, b_sf, layout: Layout, out_cols: int) -> torch.Tensor: + """``[rows, out_cols]`` bf16 of ``a @ b[e].T`` over the token-major layout, the head rows + from the head GEMM; rows after the last group are zeros.""" + out = torch.empty(layout.num_rows, out_cols, device=a.device, dtype=torch.bfloat16) + _m_grouped(a, a_sf, b, b_sf, out, layout.token_layout) + head_out = torch.empty(layout.M_head, out_cols, device=a.device, dtype=torch.bfloat16) + _m_grouped(head, head_sf, b, b_sf, head_out, layout.head_layout) + block_n = 256 + _scatter_head_kernel[(layout.counts.numel(), triton.cdiv(out_cols, block_n))]( + head_out, out, layout.table.expert_row_start, layout.counts, out_cols, BLOCK_N=block_n, num_warps=4 + ) + zero_tail(out, layout.counts, torch.cuda.get_device_properties(out.device).multi_processor_count) + return out + + +def weight_grad_accumulate(a_t, a_t_sf, b_t, b_t_sf, layout: Layout, out: torch.Tensor) -> None: + """``out[e] += A_e^T B_e`` in fp32 (``out`` is ``[E, rows, cols]``), straight from DeepGEMM's + accumulator, so gradients accumulated over micro-batches need no cast or extra pass.""" + import deep_gemm + + if out.dtype != torch.float32 or not out.is_contiguous(): + raise ValueError(f"the accumulator must be contiguous fp32, got {out.dtype}") + ks = [layout.Mp] + [0] * (layout.counts.numel() - 1) + deep_gemm.k_grouped_fp8_gemm_nt_contiguous((a_t, a_t_sf.T), (b_t, b_t_sf.T), out, ks, layout.ks, out) + + +def weight_grad(a_t, a_t_sf, b_t, b_t_sf, layout: Layout, rows: int, cols: int, split: int | None = None): + """bf16 ``A_e^T B_e`` (``[E, rows, cols]``) over each expert's tokens, from operands in the + K-grouped layout; with ``split``, as its ``[E, split, cols]`` and ``[E, rows - split, cols]`` + halves (the second empty without).""" + import deep_gemm + + num_experts = layout.counts.numel() + # DeepGEMM accumulates into d (it requires that), and skips the experts without tokens. + d = torch.zeros(num_experts, rows, cols, device=a_t.device, dtype=torch.float32) + # The host K list only steers DeepGEMM's heuristics; the kernel walks the device ks. Passing + # the padded total keeps the group sizes on the device (no synchronization). + ks = [layout.Mp] + [0] * (num_experts - 1) + deep_gemm.k_grouped_fp8_gemm_nt_contiguous((a_t, a_t_sf.T), (b_t, b_t_sf.T), d, ks, layout.ks, d) + split = rows if split is None else split + a = d.new_empty(num_experts, split, cols, dtype=torch.bfloat16) + b = d.new_empty(num_experts, rows - split, cols, dtype=torch.bfloat16) + block = 4096 + _weight_grad_cast_kernel[(num_experts, triton.cdiv(rows * cols, block))]( + d, a, b, split * cols, (rows - split) * cols, BLOCK=block, num_warps=8 + ) + return a, b diff --git a/prime_kernels/moe_experts/gemm.py b/prime_kernels/moe_experts/gemm.py new file mode 100644 index 0000000..070607d --- /dev/null +++ b/prime_kernels/moe_experts/gemm.py @@ -0,0 +1,631 @@ +"""Token-major grouped GEMMs with fused SwiGLU epilogues, as Gluon kernels for Hopper. + +Every kernel is persistent and warp specialized into two partitions: + +- two consumer warpgroups (the default partition) run the K loop of one 128 x 256 output tile on + the tensor cores, issuing the TMA loads of the K steps ``STAGES - 1`` ahead (across tile + boundaries), and park the finished tile, as bf16, in a shared-memory epilogue buffer, +- an epilogue warpgroup applies the SwiGLU forward or backward to the parked tile and writes the + results while the consumers already run the next tile's K loop. + +The epilogue partition is what makes the fusion free: on Hopper the accumulator lives in the +consumers' registers, so without it the activation math and its extra global traffic would stall +the tensor cores. (A separate producer warp would need a fourth warpgroup, which caps every +thread at 128 registers, too few for a 64 x 256 wgmma accumulator.) + +Row tiles come from a `TileTable`: each covers rows of one expert only, so a tile's last rows may +belong to the next expert. Operand rows past an expert's end are loaded anyway (they only feed +output rows that are never written), and weights are addressed as one 2D ``[experts * rows, cols]`` +view, which is safe because a row past the end of one expert's weight either feeds an output +column that is never written or multiplies an activation column past ``K``, which TMA fills with +zeros. +""" + +import torch +from triton.experimental import gluon +from triton.experimental.gluon import language as gl +from triton.experimental.gluon.language.nvidia.hopper import ( + fence_async_shared, + mbarrier, + tma, + warpgroup_mma, + warpgroup_mma_wait, +) +from triton.experimental.gluon.nvidia.hopper import TensorDescriptor + +BLOCK_M = 128 + +MODE_PLAIN = 0 +MODE_FC1 = 1 +MODE_DSWIGLU = 2 + + +@gluon.jit +def _tile_coords(tile, num_m, num_n, GROUP_M: gl.constexpr): + group_size = GROUP_M * num_n + first_m = (tile // group_size) * GROUP_M + group_m = gl.minimum(num_m - first_m, GROUP_M) + m = first_m + (tile % group_size) % group_m + n = (tile % group_size) // group_m + return m, n + + +@gluon.jit +def _tile_operands(local, num_m, num_n, tile_expert_ptr, tile_row_start_ptr, W_ROWS, GROUP_M: gl.constexpr): + """First weight row, first token row and column tile of this program's ``local``-th tile.""" + m, n = _tile_coords(gl.program_id(0) + local * gl.num_programs(0), num_m, num_n, GROUP_M) + # Past the last tile (the loads only run ahead into it, nothing is issued) read a valid entry. + m = gl.maximum(gl.minimum(m, num_m - 1), 0) + return gl.load(tile_expert_ptr + m) * W_ROWS, gl.load(tile_row_start_ptr + m), n + + +@gluon.jit +def _issue_load( + counter, + step, + n, + w_row, + row_start, + k_tiles, + a_desc, + a2_desc, + b_desc, + b2_desc, + a_smem, + b_smem, + ready, + B2_ROW_OFFSET, + MODE: gl.constexpr, + B_N_MAJOR: gl.constexpr, + DUAL_K: gl.constexpr, + BLOCK_N: gl.constexpr, + BLOCK_K: gl.constexpr, + STAGES: gl.constexpr, +): + """Load the operands of K step ``step`` of a tile into stage ``counter % STAGES``.""" + s = counter % STAGES + bar = ready.index(s) + mbarrier.expect(bar, a_desc.block_type.nbytes + BLOCK_N * BLOCK_K * 2) + if DUAL_K: + # dx = dgate @ gate_proj[e] + dup @ up_proj[e]: one K loop over both products. + first = step < k_tiles + second = step >= k_tiles + k = (step % k_tiles) * BLOCK_K + tma.async_copy_global_to_shared(a_desc, [row_start, k], bar, a_smem.index(s), pred=first) + tma.async_copy_global_to_shared(a2_desc, [row_start, k], bar, a_smem.index(s), pred=second) + tma.async_copy_global_to_shared(b_desc, [w_row + k, n * BLOCK_N], bar, b_smem.index(s), pred=first) + tma.async_copy_global_to_shared( + b2_desc, [w_row + B2_ROW_OFFSET + k, n * BLOCK_N], bar, b_smem.index(s), pred=second + ) + else: + k = step * BLOCK_K + tma.async_copy_global_to_shared(a_desc, [row_start, k], bar, a_smem.index(s)) + if MODE == 1: + # fc1: the gate rows and the up rows of the same intermediate columns side by side, + # so a single 256-wide MMA computes both. + HALF: gl.constexpr = BLOCK_N // 2 + stage = b_smem.index(s) + tma.async_copy_global_to_shared(b_desc, [w_row + n * HALF, k], bar, stage.slice(0, HALF)) + tma.async_copy_global_to_shared( + b2_desc, [w_row + B2_ROW_OFFSET + n * HALF, k], bar, stage.slice(HALF, HALF) + ) + elif B_N_MAJOR: + tma.async_copy_global_to_shared(b_desc, [w_row + n * BLOCK_N, k], bar, b_smem.index(s)) + else: + tma.async_copy_global_to_shared(b_desc, [w_row + k, n * BLOCK_N], bar, b_smem.index(s)) + + +@gluon.jit +def _consumer( + a_desc, + a2_desc, + b_desc, + b2_desc, + a_smem, + b_smem, + epi_smem, + ready, + epi_full, + epi_free, + tile_expert_ptr, + tile_row_start_ptr, + num_tiles_ptr, + N, + K, + W_ROWS, + B2_ROW_OFFSET, + MODE: gl.constexpr, + B_N_MAJOR: gl.constexpr, + DUAL_K: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + BLOCK_K: gl.constexpr, + STAGES: gl.constexpr, + GROUP_M: gl.constexpr, +): + layout: gl.constexpr = gl.NVMMADistributedLayout( + version=[3, 0], warps_per_cta=[gl.num_warps(), 1], instr_shape=[16, BLOCK_N, 16] + ) + num_m = gl.load(num_tiles_ptr) + num_n = gl.cdiv(N, BLOCK_N) + k_tiles = gl.cdiv(K, BLOCK_K) + steps = k_tiles * (2 if DUAL_K else 1) + my_tiles = gl.cdiv(gl.maximum(num_m * num_n - gl.program_id(0), 0), gl.num_programs(0)) + total = my_tiles * steps + + # The loads run STAGES - 1 K steps ahead of the MMAs; (load_local, load_step) is the K step + # the next load is for, and (w_row, row_start, n) the operands of its tile. + load_step = 0 + load_local = 0 + w_row, row_start, n = _tile_operands(0, num_m, num_n, tile_expert_ptr, tile_row_start_ptr, W_ROWS, GROUP_M) + for i in gl.static_range(STAGES - 1): + if i < total: + _issue_load( + i, + load_step, + n, + w_row, + row_start, + k_tiles, + a_desc, + a2_desc, + b_desc, + b2_desc, + a_smem, + b_smem, + ready, + B2_ROW_OFFSET, + MODE, + B_N_MAJOR, + DUAL_K, + BLOCK_N, + BLOCK_K, + STAGES, + ) + load_step += 1 + if load_step == steps: + load_step = 0 + load_local += 1 + w_row, row_start, n = _tile_operands( + load_local, num_m, num_n, tile_expert_ptr, tile_row_start_ptr, W_ROWS, GROUP_M + ) + + counter = 0 + for local in range(my_tiles): + acc = gl.zeros([BLOCK_M, BLOCK_N], gl.float32, layout) + for step in range(steps): + s = counter % STAGES + mbarrier.wait(ready.index(s), (counter // STAGES) & 1) + if B_N_MAJOR: + b = b_smem.index(s).permute((1, 0)) + else: + b = b_smem.index(s) + acc = warpgroup_mma(a_smem.index(s), b, acc, is_async=True) + acc = warpgroup_mma_wait(num_outstanding=1, deps=(acc,)) + # The MMA of step counter - 1 has retired, so its stage takes step counter + STAGES - 1. + if counter + STAGES - 1 < total: + _issue_load( + counter + STAGES - 1, + load_step, + n, + w_row, + row_start, + k_tiles, + a_desc, + a2_desc, + b_desc, + b2_desc, + a_smem, + b_smem, + ready, + B2_ROW_OFFSET, + MODE, + B_N_MAJOR, + DUAL_K, + BLOCK_N, + BLOCK_K, + STAGES, + ) + load_step += 1 + if load_step == steps: + load_step = 0 + load_local += 1 + w_row, row_start, n = _tile_operands( + load_local, num_m, num_n, tile_expert_ptr, tile_row_start_ptr, W_ROWS, GROUP_M + ) + counter += 1 + acc = warpgroup_mma_wait(num_outstanding=0, deps=(acc,)) + + # Park the tile as bf16 (what the unfused GEMM would return) for the epilogue partition. + mbarrier.wait(epi_free, (local & 1) ^ 1) + epi_smem._reinterpret(gl.bfloat16, [BLOCK_M, BLOCK_N], epi_smem.layout).store(acc.to(gl.bfloat16)) + mbarrier.arrive(epi_full, count=1) + + +@gluon.jit +def _epilogue( + epi_smem, + epi_full, + epi_free, + out0_ptr, + out1_ptr, + out2_ptr, + in0_ptr, + in1_ptr, + tile_row_start_ptr, + tile_row_end_ptr, + num_tiles_ptr, + N, + OUT_STRIDE, + limit, + MODE: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + GROUP_M: gl.constexpr, +): + """Finish each parked tile; only rows before the tile's ``row_end`` are written, the rest + belongs to the next expert. Offsets inside a tile are 32-bit, only its base pointer is 64-bit. + The row-chunk loops stay rolled: unrolled, this partition's code evicts the consumers' K loop + from the instruction cache. + """ + CHUNK: gl.constexpr = epi_smem.shape[1] + HALF: gl.constexpr = BLOCK_N // 2 + layout: gl.constexpr = gl.BlockedLayout([1, 8], [1, 32], [gl.num_warps(), 1], [1, 0]) + half_layout: gl.constexpr = gl.BlockedLayout([1, 8], [2, 16], [gl.num_warps(), 1], [1, 0]) + num_m = gl.load(num_tiles_ptr) + num_n = gl.cdiv(N, BLOCK_N) + local = 0 + for tile in range(gl.program_id(0), num_m * num_n, gl.num_programs(0)): + m, n = _tile_coords(tile, num_m, num_n, GROUP_M) + row_start = gl.load(tile_row_start_ptr + m) + tile_rows = gl.load(tile_row_end_ptr + m) - row_start + mbarrier.wait(epi_full, local & 1) + if MODE == 0: + base = row_start.to(gl.int64) * N + cols = n * BLOCK_N + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, layout)) + for r in range(BLOCK_M // CHUNK): + rows = r * CHUNK + gl.arange(0, CHUNK, layout=gl.SliceLayout(1, layout)) + mask = (rows < tile_rows)[:, None] & (cols < N)[None, :] + value = epi_smem.index(r).load(layout) + gl.store(out0_ptr + base + rows[:, None] * N + cols[None, :], value, mask=mask) + elif MODE == 1: + # fc1: the tile holds gate (first half) and up (second half) of the intermediate + # columns n * HALF + [0, HALF). + intermediate = N // 2 + base = row_start.to(gl.int64) * intermediate + cols = n * HALF + gl.arange(0, HALF, layout=gl.SliceLayout(0, half_layout)) + for r in range(BLOCK_M // CHUNK): + rows = r * CHUNK + gl.arange(0, CHUNK, layout=gl.SliceLayout(1, half_layout)) + mask = (rows < tile_rows)[:, None] & (cols < intermediate)[None, :] + offsets = base + rows[:, None] * intermediate + cols[None, :] + chunk = epi_smem.index(r) + gate = chunk.slice(0, HALF, dim=1).load(half_layout) + up = chunk.slice(HALF, HALF, dim=1).load(half_layout) + gl.store(out0_ptr + offsets, gate, mask=mask) + gl.store(out1_ptr + offsets, up, mask=mask) + gate_c = gl.minimum(gate.to(gl.float32), limit) + up_c = gl.minimum(gl.maximum(up.to(gl.float32), -limit), limit) + h = gate_c / (1.0 + gl.exp(-gate_c)) * up_c + gl.store(out2_ptr + offsets, h.to(out2_ptr.dtype.element_ty), mask=mask) + else: + # dh -> dgate (out0), dup (out1) through the saved gate (in0) / up (in1); h (out2) is + # recomputed for the down projection's weight gradient. dgate and dup rows are + # OUT_STRIDE apart, so they can be the two halves of one [rows, 2 * N] tensor. + base = row_start.to(gl.int64) * N + out_base = row_start.to(gl.int64) * OUT_STRIDE + cols = n * BLOCK_N + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, layout)) + for r in range(BLOCK_M // CHUNK): + rows = r * CHUNK + gl.arange(0, CHUNK, layout=gl.SliceLayout(1, layout)) + mask = (rows < tile_rows)[:, None] & (cols < N)[None, :] + offsets = base + rows[:, None] * N + cols[None, :] + out_offsets = out_base + rows[:, None] * OUT_STRIDE + cols[None, :] + dh = epi_smem.index(r).load(layout).to(gl.float32) + gate = gl.load(in0_ptr + offsets, mask=mask, other=0.0).to(gl.float32) + up = gl.load(in1_ptr + offsets, mask=mask, other=0.0).to(gl.float32) + gate_c = gl.minimum(gate, limit) + up_c = gl.minimum(gl.maximum(up, -limit), limit) + sig = 1.0 / (1.0 + gl.exp(-gate_c)) + silu = gate_c * sig + dgate = gl.where(gate <= limit, dh * up_c * (sig * (1.0 + gate_c * (1.0 - sig))), 0.0) + dup = gl.where((up >= -limit) & (up <= limit), dh * silu, 0.0) + gl.store(out0_ptr + out_offsets, dgate.to(out0_ptr.dtype.element_ty), mask=mask) + gl.store(out1_ptr + out_offsets, dup.to(out1_ptr.dtype.element_ty), mask=mask) + gl.store(out2_ptr + offsets, (silu * up_c).to(out2_ptr.dtype.element_ty), mask=mask) + mbarrier.arrive(epi_free, count=1) + local += 1 + + +@gluon.jit +def _grouped_gemm_kernel( + a_desc, + a2_desc, + b_desc, + b2_desc, + out0_ptr, + out1_ptr, + out2_ptr, + in0_ptr, + in1_ptr, + tile_expert_ptr, + tile_row_start_ptr, + tile_row_end_ptr, + num_tiles_ptr, + N, + K, + W_ROWS, + B2_ROW_OFFSET, + OUT_STRIDE, + limit, + MODE: gl.constexpr, + B_N_MAJOR: gl.constexpr, + DUAL_K: gl.constexpr, + BLOCK_M: gl.constexpr, + BLOCK_N: gl.constexpr, + BLOCK_K: gl.constexpr, + STAGES: gl.constexpr, + GROUP_M: gl.constexpr, + EPILOGUE_CHUNK: gl.constexpr, + EPILOGUE_REGS: gl.constexpr, +): + a_smem = gl.allocate_shared_memory(a_desc.dtype, [STAGES] + a_desc.block_type.shape, a_desc.layout) + if B_N_MAJOR: + b_block: gl.constexpr = [BLOCK_N, BLOCK_K] + else: + b_block: gl.constexpr = [BLOCK_K, BLOCK_N] + b_smem = gl.allocate_shared_memory( + b_desc.dtype, [STAGES] + b_block, gl.NVMMASharedLayout.get_default_for(b_block, b_desc.dtype) + ) + # The parked tile is [row chunk, row, column] so the epilogue can walk row chunks with a rolled + # loop. The 2D swizzle repeats every 8 rows, so for chunks of a multiple of 8 rows this is the + # same memory layout as the plain [rows, columns] tile the consumers write. + gl.static_assert(EPILOGUE_CHUNK % 8 == 0) + epi_smem = gl.allocate_shared_memory( + gl.bfloat16, + [BLOCK_M // EPILOGUE_CHUNK, EPILOGUE_CHUNK, BLOCK_N], + gl.SwizzledSharedLayout(8, 1, 8, [1, 0]), + ) + ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + epi_bars = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout()) + for i in gl.static_range(STAGES): + mbarrier.init(ready.index(i), count=1) + mbarrier.init(epi_bars.index(0), count=1) + mbarrier.init(epi_bars.index(1), count=1) + fence_async_shared() + epi_full = epi_bars.index(0) + epi_free = epi_bars.index(1) + # Constexpr arguments must be spelled out in each tuple: building tuples by concatenation + # turns them into tensors. + gl.warp_specialize( + [ + ( + _consumer, + ( + a_desc, + a2_desc, + b_desc, + b2_desc, + a_smem, + b_smem, + epi_smem, + ready, + epi_full, + epi_free, + tile_expert_ptr, + tile_row_start_ptr, + num_tiles_ptr, + N, + K, + W_ROWS, + B2_ROW_OFFSET, + MODE, + B_N_MAJOR, + DUAL_K, + BLOCK_M, + BLOCK_N, + BLOCK_K, + STAGES, + GROUP_M, + ), + ), + ( + _epilogue, + ( + epi_smem, + epi_full, + epi_free, + out0_ptr, + out1_ptr, + out2_ptr, + in0_ptr, + in1_ptr, + tile_row_start_ptr, + tile_row_end_ptr, + num_tiles_ptr, + N, + OUT_STRIDE, + limit, + MODE, + BLOCK_M, + BLOCK_N, + GROUP_M, + ), + ), + ], + [4], + [EPILOGUE_REGS], + ) + + +def _desc(tensor: torch.Tensor, block: list[int]) -> TensorDescriptor: + return TensorDescriptor.from_tensor(tensor, block, gl.NVMMASharedLayout.get_default_for(block, gl.bfloat16)) + + +def weight_view(weight: torch.Tensor) -> tuple[torch.Tensor, int]: + """A 2D ``[rows, cols]`` view of an ``[E, rows_e, cols]`` weight whose experts are a fixed + number of rows apart (contiguous, or one half of a packed gate/up weight), and that stride.""" + num_experts, rows, cols = weight.shape + if weight.stride(2) != 1 or weight.stride(1) != cols or weight.stride(0) % cols: + raise ValueError(f"expert weights must be row-major with whole rows between experts, got {weight.stride()}") + expert_rows = weight.stride(0) // cols + view = weight.as_strided(((num_experts - 1) * expert_rows + rows, cols), (cols, 1)) + return view, expert_rows + + +DEFAULT_CONFIG = dict(block_n=256, block_k=64, stages=3, group_m=8, epilogue_chunk=8, epilogue_regs=128) + + +def launch( + mode: int, + a: torch.Tensor, + weight: torch.Tensor, + table, + *, + n_major: bool, + out_cols: int, + outputs: tuple[torch.Tensor, ...], + inputs: tuple[torch.Tensor, ...] = (), + a2: torch.Tensor | None = None, + weight2: torch.Tensor | None = None, + limit: float = 0.0, + num_sms: int, + block_n: int, + block_k: int, + stages: int, + group_m: int, + epilogue_chunk: int, + epilogue_regs: int, +) -> None: + """Run one grouped GEMM over ``table``'s row tiles; ``out_cols`` is the N the tiles cover.""" + assert table.block_m == BLOCK_M + K = a.shape[1] + w, w_rows = weight_view(weight) + if mode == MODE_FC1: + b_block = [block_n // 2, block_k] + else: + b_block = [block_n, block_k] if n_major else [block_k, block_n] + b_desc = _desc(w, b_block) + b2_desc = b_desc + b2_row_offset = 0 + if weight2 is not None: + w2, w2_rows = weight_view(weight2) + if w2_rows != w_rows: + raise ValueError("gate_proj and up_proj must have the same expert stride") + if w2.data_ptr() == w.data_ptr() + weight.shape[1] * w.stride(0) * w.element_size(): + # Packed gate/up halves: address up as rows of one view over the whole packed weight. + b2_row_offset = weight.shape[1] + b_desc = b2_desc = _desc(w.as_strided((weight.shape[0] * w_rows, w.shape[1]), w.stride()), b_block) + else: + b2_desc = _desc(w2, b_block) + a_desc = _desc(a, [BLOCK_M, block_k]) + a2_desc = _desc(a2, [BLOCK_M, block_k]) if a2 is not None else a_desc + outs = list(outputs) + [outputs[0]] * (3 - len(outputs)) + ins = list(inputs) + [outputs[0]] * (2 - len(inputs)) + grid = min(num_sms, table.max_tiles * -(-out_cols // block_n)) + _grouped_gemm_kernel[(grid,)]( + a_desc, + a2_desc, + b_desc, + b2_desc, + *outs, + *ins, + table.expert, + table.row_start, + table.row_end, + table.num_tiles, + out_cols, + K, + w_rows, + b2_row_offset, + outs[0].stride(0), + limit, + MODE=mode, + B_N_MAJOR=n_major, + DUAL_K=a2 is not None, + BLOCK_M=BLOCK_M, + BLOCK_N=block_n, + BLOCK_K=block_k, + STAGES=stages, + GROUP_M=group_m, + EPILOGUE_CHUNK=epilogue_chunk, + EPILOGUE_REGS=epilogue_regs, + num_warps=8, + ) + + +def fc1(x, gate_proj, up_proj, table, limit, *, num_sms, config=None): + """gate, up and h = clamped_swiglu(gate, up), each ``[rows, intermediate]``.""" + num_rows, intermediate = x.shape[0], gate_proj.shape[1] + gate, up, h = (x.new_empty(num_rows, intermediate) for _ in range(3)) + launch( + MODE_FC1, + x, + gate_proj, + table, + n_major=True, + out_cols=2 * intermediate, + outputs=(gate, up, h), + weight2=up_proj, + limit=limit, + num_sms=num_sms, + **(config or DEFAULT_CONFIG), + ) + return gate, up, h + + +def down(h, down_proj, table, *, num_sms, config=None): + """h @ down_proj[e].T per group, ``[rows, H]``.""" + out = h.new_empty(h.shape[0], down_proj.shape[1]) + launch( + MODE_PLAIN, + h, + down_proj, + table, + n_major=True, + out_cols=out.shape[1], + outputs=(out,), + num_sms=num_sms, + **(config or DEFAULT_CONFIG), + ) + return out + + +def dswiglu(dout, down_proj, gate, up, table, limit, *, num_sms, config=None): + """dgate and dup from dh = dout @ down_proj[e], as the two halves of one ``[rows, 2 * I]`` + tensor, and h recomputed from gate / up.""" + num_rows, intermediate = gate.shape + dgate_dup = gate.new_empty(num_rows, 2 * intermediate) + dgate, dup = dgate_dup[:, :intermediate], dgate_dup[:, intermediate:] + h = torch.empty_like(gate) + launch( + MODE_DSWIGLU, + dout, + down_proj, + table, + n_major=False, + out_cols=intermediate, + outputs=(dgate, dup, h), + inputs=(gate, up), + limit=limit, + num_sms=num_sms, + **(config or DEFAULT_CONFIG), + ) + return dgate_dup, h + + +def dx(dgate, dup, gate_proj, up_proj, table, *, num_sms, config=None): + """dgate @ gate_proj[e] + dup @ up_proj[e], ``[rows, H]``.""" + out = dgate.new_empty(dgate.shape[0], gate_proj.shape[2]) + launch( + MODE_PLAIN, + dgate, + gate_proj, + table, + n_major=False, + out_cols=out.shape[1], + outputs=(out,), + a2=dup, + weight2=up_proj, + num_sms=num_sms, + **(config or DEFAULT_CONFIG), + ) + return out diff --git a/prime_kernels/moe_experts/kernels.py b/prime_kernels/moe_experts/kernels.py new file mode 100644 index 0000000..b070040 --- /dev/null +++ b/prime_kernels/moe_experts/kernels.py @@ -0,0 +1,104 @@ +"""Small Triton helpers around the grouped GEMMs: the row-tile table and the tail fill.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _tile_table_kernel( + counts_ptr, + tile_expert_ptr, + tile_row_start_ptr, + tile_row_end_ptr, + num_tiles_ptr, + expert_row_start_ptr, + num_experts, + num_rows, + max_tiles, + BLOCK_M: tl.constexpr, + BLOCK_E: tl.constexpr, + BLOCK_T: tl.constexpr, +): + experts = tl.arange(0, BLOCK_E) + counts = tl.load(counts_ptr + experts, mask=experts < num_experts, other=0).to(tl.int32) + tiles = tl.cdiv(counts, BLOCK_M) + tile_end = tl.cumsum(tiles, axis=0) + row_end = tl.cumsum(counts, axis=0) + if tl.program_id(0) == 0: + tl.store(num_tiles_ptr, tl.sum(tiles)) + tl.store(expert_row_start_ptr + experts, row_end - counts, mask=experts < num_experts) + + t = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + # The expert of tile t is the number of experts whose tiles all come before t. + expert = tl.sum((tile_end[None, :] <= t[:, None]).to(tl.int32), axis=1) + owner = experts[None, :] == expert[:, None] + first_tile = tl.sum(tl.where(owner, tile_end - tiles, 0), axis=1) + expert_row_start = tl.sum(tl.where(owner, row_end - counts, 0), axis=1) + # Counts that overrun the activations are cut at their end, so no tile writes past it. + expert_row_end = tl.minimum(tl.sum(tl.where(owner, row_end, 0), axis=1), num_rows) + row_start = expert_row_start + (t - first_tile) * BLOCK_M + valid = t < max_tiles + tl.store(tile_expert_ptr + t, expert, mask=valid) + tl.store(tile_row_start_ptr + t, row_start, mask=valid) + tl.store(tile_row_end_ptr + t, expert_row_end, mask=valid) + + +class TileTable: + """Row tiles of ``block_m`` rows, each inside one expert's group, built on device. + + Expert ``e`` owns rows ``[start_e, start_e + num_tokens_per_expert[e])`` with the groups laid + out back to back from row 0. Tile ``t < num_tiles`` covers rows ``[row_start[t], row_start[t] + + block_m)`` of expert ``expert[t]``, of which only those before ``row_end[t]`` are its own. + """ + + def __init__(self, num_tokens_per_expert: torch.Tensor, num_rows: int, block_m: int) -> None: + num_experts = num_tokens_per_expert.numel() + device = num_tokens_per_expert.device + self.block_m = block_m + self.max_tiles = triton.cdiv(num_rows, block_m) + num_experts + self.expert = torch.empty(self.max_tiles, dtype=torch.int32, device=device) + self.row_start = torch.empty_like(self.expert) + self.row_end = torch.empty_like(self.expert) + self.num_tiles = torch.empty(1, dtype=torch.int32, device=device) + self.expert_row_start = torch.empty(num_experts, dtype=torch.int32, device=device) + self.num_tokens_per_expert = num_tokens_per_expert + block_t = 64 + _tile_table_kernel[(triton.cdiv(self.max_tiles, block_t),)]( + num_tokens_per_expert, + self.expert, + self.row_start, + self.row_end, + self.num_tiles, + self.expert_row_start, + num_experts, + num_rows, + self.max_tiles, + BLOCK_M=block_m, + BLOCK_E=triton.next_power_of_2(num_experts), + BLOCK_T=block_t, + ) + + +@triton.jit +def _zero_tail_kernel(out_ptr, counts_ptr, num_experts, numel, N, BLOCK_E: tl.constexpr, BLOCK: tl.constexpr): + experts = tl.arange(0, BLOCK_E) + used = tl.sum(tl.load(counts_ptr + experts, mask=experts < num_experts, other=0)).to(tl.int64) + for start in range(used * N + tl.program_id(0) * BLOCK, numel, tl.num_programs(0) * BLOCK): + offsets = start + tl.arange(0, BLOCK) + tl.store(out_ptr + offsets, tl.zeros((BLOCK,), dtype=out_ptr.dtype.element_ty), mask=offsets < numel) + + +def zero_tail(out: torch.Tensor, num_tokens_per_expert: torch.Tensor, num_programs: int) -> None: + """Zero the rows of ``out`` after the last expert's group, which no GEMM tile owns.""" + num_rows, N = out.shape + num_experts = num_tokens_per_expert.numel() + _zero_tail_kernel[(num_programs,)]( + out, + num_tokens_per_expert, + num_experts, + num_rows * N, + N, + BLOCK_E=triton.next_power_of_2(num_experts), + BLOCK=2048, + ) diff --git a/prime_kernels/moe_experts/ops.py b/prime_kernels/moe_experts/ops.py new file mode 100644 index 0000000..626bb23 --- /dev/null +++ b/prime_kernels/moe_experts/ops.py @@ -0,0 +1,508 @@ +"""Differentiable expert MLP of a rank's local experts: ``down(clamped_swiglu(gate(x), up(x)))``. + +Forward: one fused GEMM computes gate and up and applies the clamped SwiGLU, a second GEMM +the down projection. Only ``x``, ``gate`` and ``up`` are saved. Backward: one GEMM computes +``dh = dout @ down`` and turns it into ``dgate`` / ``dup`` (and recomputes ``h``), one GEMM +computes ``dx`` over both products. The weight gradients have nothing to fuse and stay with +``torch._grouped_mm`` (CUTLASS), which was faster than a Gluon version of them. + +With ``fp8=True`` the GEMMs are DeepGEMM's blockwise FP8 ones, with the quantization fused into +the passes around them (see `fp8`); the forward saves the padded gate/up GEMM output, ``x`` +quantized for the weight gradient and the transposed FP8 weights. +""" + +import importlib.util + +import torch + +from prime_kernels.moe_experts import fp8 as fp8_impl +from prime_kernels.moe_experts import gemm +from prime_kernels.moe_experts.kernels import TileTable, zero_tail + +# Every expert's group of rows must be a multiple of this many rows long (prime-rl's dispatcher +# pads the groups to its experts' token_group_alignment): torch._grouped_mm needs 16-byte aligned +# group starts along its K dimension for the weight gradients. +TOKEN_GROUP_ALIGNMENT = 8 + +_NUM_SMS: dict[int, int] = {} + + +def _num_sms(device: torch.device) -> int: + if device.index not in _NUM_SMS: + _NUM_SMS[device.index] = torch.cuda.get_device_properties(device).multi_processor_count + return _NUM_SMS[device.index] + + +def unsupported_shape_reason(hidden_size: int, intermediate_size: int, fp8: bool = False) -> str | None: + """Why these expert sizes cannot run, or None. Rows must be 16-byte aligned for TMA.""" + if fp8: + return fp8_impl.unsupported_shape_reason(hidden_size, intermediate_size) + if hidden_size % 8 or intermediate_size % 8: + return f"hidden ({hidden_size}) and intermediate ({intermediate_size}) sizes must be multiples of 8" + return None + + +def _gate_up(w1: torch.Tensor, w3: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]: + if w3 is not None: + return w1, w3 + intermediate = w1.shape[1] // 2 + return w1[:, :intermediate], w1[:, intermediate:] + + +def _check_groups(num_tokens_per_expert: torch.Tensor, num_rows: int) -> None: + # Asynchronous device asserts: the group sizes are device data, and reading them on the host + # would synchronize every MoE layer. + torch._assert_async((num_tokens_per_expert % TOKEN_GROUP_ALIGNMENT == 0).all()) + torch._assert_async(num_tokens_per_expert.sum() <= num_rows) + + +@torch.library.custom_op("prime_kernels::moe_experts_forward", mutates_args=()) +def _moe_experts_fwd( + x: torch.Tensor, + w1: torch.Tensor, + w3: torch.Tensor | None, + w2: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + swiglu_limit: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + gate_proj, up_proj = _gate_up(w1, w3) + _check_groups(num_tokens_per_expert, x.shape[0]) + if x.shape[0] == 0: + empty = x.new_empty(0, gate_proj.shape[1]) + return x.new_empty(x.shape), empty, empty.clone() + with torch.cuda.device(x.device): + num_sms = _num_sms(x.device) + table = TileTable(num_tokens_per_expert, x.shape[0], gemm.BLOCK_M) + gate, up, h = gemm.fc1(x, gate_proj, up_proj, table, swiglu_limit, num_sms=num_sms) + out = gemm.down(h, w2, table, num_sms=num_sms) + zero_tail(out, num_tokens_per_expert, num_sms) + return out, gate, up + + +@_moe_experts_fwd.register_fake +def _(x, w1, w3, w2, num_tokens_per_expert, swiglu_limit): + intermediate = w2.shape[2] + return x.new_empty(x.shape), x.new_empty(x.shape[0], intermediate), x.new_empty(x.shape[0], intermediate) + + +@torch.library.custom_op("prime_kernels::moe_experts_backward", mutates_args=()) +def _moe_experts_bwd( + dout: torch.Tensor, + x: torch.Tensor, + w1: torch.Tensor, + w3: torch.Tensor | None, + w2: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + gate: torch.Tensor, + up: torch.Tensor, + swiglu_limit: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """dx, dw1, dw3 (empty when w1 packs gate and up) and dw2.""" + gate_proj, up_proj = _gate_up(w1, w3) + dout = dout.contiguous() + intermediate = gate.shape[1] + if x.shape[0] == 0: + dw3 = x.new_empty(0) if w3 is None else x.new_zeros(w3.shape) + return x.new_empty(x.shape), x.new_zeros(w1.shape), dw3, x.new_zeros(w2.shape) + with torch.cuda.device(x.device): + num_sms = _num_sms(x.device) + table = TileTable(num_tokens_per_expert, x.shape[0], gemm.BLOCK_M) + dgate_dup, h = gemm.dswiglu(dout, w2, gate, up, table, swiglu_limit, num_sms=num_sms) + dgate, dup = dgate_dup[:, :intermediate], dgate_dup[:, intermediate:] + dx = gemm.dx(dgate, dup, gate_proj, up_proj, table, num_sms=num_sms) + zero_tail(dx, num_tokens_per_expert, num_sms) + offsets = torch.cumsum(num_tokens_per_expert, 0, dtype=torch.int32) + dw2 = torch._grouped_mm(dout.t(), h, offs=offsets) + if w3 is None: + # [E, 2 * I, H] straight from the two halves of dgate_dup: the packed weight's layout. + dw1 = torch._grouped_mm(dgate_dup.t(), x, offs=offsets) + dw3 = x.new_empty(0) + else: + dw1 = torch._grouped_mm(dgate.t(), x, offs=offsets) + dw3 = torch._grouped_mm(dup.t(), x, offs=offsets) + return dx, dw1, dw3, dw2 + + +@_moe_experts_bwd.register_fake +def _(dout, x, w1, w3, w2, num_tokens_per_expert, gate, up, swiglu_limit): + dw3 = x.new_empty(0) if w3 is None else x.new_empty(w3.shape) + return x.new_empty(x.shape), x.new_empty(w1.shape), dw3, x.new_empty(w2.shape) + + +def _setup_context(ctx, inputs, output): + x, w1, w3, w2, num_tokens_per_expert, swiglu_limit = inputs + _, gate, up = output + ctx.save_for_backward(x, w1, w3, w2, num_tokens_per_expert, gate, up) + ctx.swiglu_limit = swiglu_limit + # gate and up are only returned to be saved; without this autograd would fill zero gradients + # for them on every backward. + ctx.set_materialize_grads(False) + + +def _backward(ctx, dout, _dgate, _dup): + x, w1, w3, w2, num_tokens_per_expert, gate, up = ctx.saved_tensors + dx, dw1, dw3, dw2 = _moe_experts_bwd(dout, x, w1, w3, w2, num_tokens_per_expert, gate, up, ctx.swiglu_limit) + return dx, dw1, dw3 if w3 is not None else None, dw2, None, None + + +_moe_experts_fwd.register_autograd(_backward, setup_context=_setup_context) + + +@torch.library.custom_op("prime_kernels::moe_experts_fp8_forward", mutates_args=()) +def _moe_experts_fp8_fwd( + x: torch.Tensor, + w1: torch.Tensor, + w3: torch.Tensor | None, + w2: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + swiglu_limit: float, + save_for_backward: bool, +) -> list[torch.Tensor]: + """out, then what the backward reads (empty without ``save_for_backward``): the padded + gate/up GEMM output, ``x`` quantized per column group and the transposed FP8 weights.""" + _check_groups(num_tokens_per_expert, x.shape[0]) + gate_proj, up_proj = _gate_up(w1, w3) + intermediate = gate_proj.shape[1] + if x.shape[0] == 0: + return [x.new_empty(x.shape), *_fp8_saved(x, w2, save_for_backward)] + with torch.cuda.device(x.device): + layout = fp8_impl.Layout(num_tokens_per_expert, x.shape[0]) + x_q, x_sf, x_t, x_t_sf = fp8_impl.quantize_activation(x, layout, columns=save_for_backward) + w13_q, w13_sf, w13_t, w13_t_sf = fp8_impl.quantize_weight(gate_proj, up_proj, transposed=save_for_backward) + w2_q, w2_sf, w2_t, w2_t_sf = fp8_impl.quantize_weight(w2, None, transposed=save_for_backward) + gate_up = fp8_impl.padded_gemm(x_q, x_sf, w13_q, w13_sf, layout, 2 * intermediate) + del x_q, x_sf + h_q, h_sf, h_head, h_head_sf = fp8_impl.swiglu_quantize(gate_up, layout, swiglu_limit) + out = fp8_impl.token_gemm(h_q, h_sf, h_head, h_head_sf, w2_q, w2_sf, layout, x.shape[1]) + if not save_for_backward: + return [out, *_fp8_saved(x, w2, save_for_backward)] + return [out, gate_up, x_t, x_t_sf, w13_t, w13_t_sf, w2_t, w2_t_sf] + + +@torch.library.custom_op("prime_kernels::moe_experts_fp8_quantize_weights", mutates_args=()) +def _moe_experts_fp8_quantize_weights( + w1: torch.Tensor, w3: torch.Tensor | None, w2: torch.Tensor +) -> list[torch.Tensor]: + """The FP8 weights `moe_experts_fp8_forward_quantized` reads: the gate/up and down blocks with + their scales, then their transposed copies (for the backward) with theirs.""" + gate_proj, up_proj = _gate_up(w1, w3) + with torch.cuda.device(w2.device): + w13 = fp8_impl.quantize_weight(gate_proj, up_proj, transposed=True) + w2q = fp8_impl.quantize_weight(w2, None, transposed=True) + return [*w13[:2], *w2q[:2], *w13[2:], *w2q[2:]] + + +@_moe_experts_fp8_quantize_weights.register_fake +def _(w1, w3, w2): + gate_proj, up_proj = _gate_up(w1, w3) + E, I2, H = gate_proj.shape[0], gate_proj.shape[1] + (up_proj.shape[1] if up_proj is not None else 0), w2.shape[1] + I, G, f8 = w2.shape[2], fp8_impl.GROUP, torch.float8_e4m3fn + return [ + w2.new_empty(E, I2, H, dtype=f8), w2.new_empty(E, I2 // G, H // G, dtype=torch.float32), + w2.new_empty(E, H, I, dtype=f8), w2.new_empty(E, H // G, I // G, dtype=torch.float32), + w2.new_empty(E, H, I2, dtype=f8), w2.new_empty(E, H // G, I2 // G, dtype=torch.float32), + w2.new_empty(E, I, H, dtype=f8), w2.new_empty(E, I // G, H // G, dtype=torch.float32), + ] + + +@torch.library.custom_op("prime_kernels::moe_experts_fp8_forward_quantized", mutates_args=()) +def _moe_experts_fp8_fwd_quantized( + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + swiglu_limit: float, + w13_q: torch.Tensor, + w13_sf: torch.Tensor, + w2_q: torch.Tensor, + w2_sf: torch.Tensor, +) -> list[torch.Tensor]: + """`moe_experts_fp8_forward` with ``save_for_backward`` from weights quantized ahead (by + `moe_experts_fp8_quantize_weights`): out, the padded gate/up GEMM output and ``x`` quantized + per column group with its scales. The backward reads the quantized transposed weights too.""" + _check_groups(num_tokens_per_expert, x.shape[0]) + num_experts, hidden, intermediate = w2_q.shape + if x.shape[0] == 0: + Mp = fp8_impl.padded_rows(0, num_experts) + return [ + x.new_empty(x.shape), x.new_zeros(Mp, 2 * intermediate), + x.new_zeros(Mp * hidden, dtype=torch.float8_e4m3fn), x.new_zeros(Mp // fp8_impl.GROUP, hidden, dtype=torch.float32), + ] + with torch.cuda.device(x.device): + layout = fp8_impl.Layout(num_tokens_per_expert, x.shape[0]) + x_q, x_sf, x_t, x_t_sf = fp8_impl.quantize_activation(x, layout, columns=True) + gate_up = fp8_impl.padded_gemm(x_q, x_sf, w13_q, w13_sf, layout, 2 * intermediate) + del x_q, x_sf + h_q, h_sf, h_head, h_head_sf = fp8_impl.swiglu_quantize(gate_up, layout, swiglu_limit) + out = fp8_impl.token_gemm(h_q, h_sf, h_head, h_head_sf, w2_q, w2_sf, layout, x.shape[1]) + return [out, gate_up, x_t, x_t_sf] + + +@_moe_experts_fp8_fwd_quantized.register_fake +def _(x, num_tokens_per_expert, swiglu_limit, w13_q, w13_sf, w2_q, w2_sf): + num_experts, hidden, intermediate = w2_q.shape + Mp = fp8_impl.padded_rows(x.shape[0], num_experts) + return [ + x.new_empty(x.shape), x.new_empty(Mp, 2 * intermediate), + x.new_empty(Mp * hidden, dtype=torch.float8_e4m3fn), x.new_empty(Mp // fp8_impl.GROUP, hidden, dtype=torch.float32), + ] + + +def _fp8_saved(x, w2, save_for_backward): + """Zeros shaped like what the FP8 forward saves for backward (all empty without saving).""" + num_experts, hidden, intermediate = w2.shape + Mp = fp8_impl.padded_rows(x.shape[0], num_experts) if save_for_backward else 0 + G = fp8_impl.GROUP + E = num_experts if save_for_backward else 0 + return [ + x.new_zeros(Mp, 2 * intermediate), + x.new_zeros(Mp * hidden, dtype=torch.float8_e4m3fn), + x.new_zeros(Mp // G, hidden, dtype=torch.float32), + x.new_zeros(E, hidden, 2 * intermediate, dtype=torch.float8_e4m3fn), + x.new_zeros(E, hidden // G, 2 * intermediate // G, dtype=torch.float32), + x.new_zeros(E, intermediate, hidden, dtype=torch.float8_e4m3fn), + x.new_zeros(E, intermediate // G, hidden // G, dtype=torch.float32), + ] + + +@_moe_experts_fp8_fwd.register_fake +def _(x, w1, w3, w2, num_tokens_per_expert, swiglu_limit, save_for_backward): + return [x.new_empty(x.shape), *_fp8_saved(x, w2, save_for_backward)] + + +@torch.library.custom_op("prime_kernels::moe_experts_fp8_backward", mutates_args=()) +def _moe_experts_fp8_bwd( + dout: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + gate_up: torch.Tensor, + x_t: torch.Tensor, + x_t_sf: torch.Tensor, + w13_t: torch.Tensor, + w13_t_sf: torch.Tensor, + w2_t: torch.Tensor, + w2_t_sf: torch.Tensor, + swiglu_limit: float, + packed: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """dx, dw1, dw3 (empty when w1 packs gate and up) and dw2.""" + dout = dout.contiguous() + num_experts, hidden, two_intermediate = w13_t.shape + intermediate = two_intermediate // 2 + if dout.shape[0] == 0: + dw3 = dout.new_empty(0) if packed else dout.new_zeros(num_experts, intermediate, hidden) + dw1 = dout.new_zeros(num_experts, intermediate * (2 if packed else 1), hidden) + return dout.new_empty(dout.shape), dw1, dw3, dout.new_zeros(num_experts, hidden, intermediate) + with torch.cuda.device(dout.device): + layout = fp8_impl.Layout(num_tokens_per_expert, dout.shape[0]) + dout_q, dout_sf, dout_t, dout_t_sf = fp8_impl.quantize_activation(dout, layout, columns=True) + dh = fp8_impl.padded_gemm(dout_q, dout_sf, w2_t, w2_t_sf, layout, intermediate) + del dout_q, dout_sf + dgu, dgu_t, h_t = fp8_impl.swiglu_backward_quantize(dh, gate_up, layout, swiglu_limit) + del dh + dx = fp8_impl.token_gemm(*dgu, w13_t, w13_t_sf, layout, hidden) + del dgu + dw2, _ = fp8_impl.weight_grad(dout_t, dout_t_sf, *h_t, layout, hidden, intermediate) + dw1, dw3 = fp8_impl.weight_grad( + *dgu_t, x_t, x_t_sf, layout, two_intermediate, hidden, split=None if packed else intermediate + ) + if packed: + dw3 = dx.new_empty(0) + return dx, dw1, dw3, dw2 + + +@_moe_experts_fp8_bwd.register_fake +def _(dout, num_tokens_per_expert, gate_up, x_t, x_t_sf, w13_t, w13_t_sf, w2_t, w2_t_sf, swiglu_limit, packed): + num_experts, hidden, two_intermediate = w13_t.shape + intermediate = two_intermediate // 2 + dx = dout.new_empty(dout.shape) + dw2 = dout.new_empty(num_experts, hidden, intermediate) + if packed: + return dx, dout.new_empty(num_experts, two_intermediate, hidden), dout.new_empty(0), dw2 + dw = dout.new_empty(num_experts, intermediate, hidden) + return dx, dw, dw.clone(), dw2 + + +def _fp8_backward_data(dout, num_tokens_per_expert, gate_up, w13_t, w13_t_sf, w2_t, w2_t_sf, swiglu_limit): + """dx, then the weight gradients' FP8 operands in the K-grouped layout: ``dout``, ``h`` and + ``dgate | dup``, each with its scales.""" + hidden, intermediate = w13_t.shape[1], w13_t.shape[2] // 2 + layout = fp8_impl.Layout(num_tokens_per_expert, dout.shape[0]) + dout_q, dout_sf, dout_t, dout_t_sf = fp8_impl.quantize_activation(dout, layout, columns=True) + dh = fp8_impl.padded_gemm(dout_q, dout_sf, w2_t, w2_t_sf, layout, intermediate) + del dout_q, dout_sf + dgu, dgu_t, h_t = fp8_impl.swiglu_backward_quantize(dh, gate_up, layout, swiglu_limit) + del dh + dx = fp8_impl.token_gemm(*dgu, w13_t, w13_t_sf, layout, hidden) + return dx, layout, (dout_t, dout_t_sf, *h_t, *dgu_t) + + +def _fp8_weight_grad_accumulate(layout, dout_t, dout_t_sf, h_t, h_t_sf, dgu_t, dgu_t_sf, x_t, x_t_sf, dw13, dw2): + fp8_impl.weight_grad_accumulate(dout_t, dout_t_sf, h_t, h_t_sf, layout, dw2) + fp8_impl.weight_grad_accumulate(dgu_t, dgu_t_sf, x_t, x_t_sf, layout, dw13) + + +@torch.library.custom_op("prime_kernels::moe_experts_fp8_backward_accumulate", mutates_args=("dw13", "dw2")) +def _moe_experts_fp8_bwd_accumulate( + dout: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + gate_up: torch.Tensor, + x_t: torch.Tensor, + x_t_sf: torch.Tensor, + w13_t: torch.Tensor, + w13_t_sf: torch.Tensor, + w2_t: torch.Tensor, + w2_t_sf: torch.Tensor, + swiglu_limit: float, + dw13: torch.Tensor, + dw2: torch.Tensor, +) -> torch.Tensor: + """`moe_experts_fp8_backward` for packed gate/up weights that adds the weight gradients into + the fp32 accumulators ``dw13`` (``[E, 2 * intermediate, hidden]``) and ``dw2`` + (``[E, hidden, intermediate]``) and returns ``dx``.""" + dout = dout.contiguous() + if dout.shape[0] == 0: + return dout.new_empty(dout.shape) + with torch.cuda.device(dout.device): + dx, layout, operands = _fp8_backward_data( + dout, num_tokens_per_expert, gate_up, w13_t, w13_t_sf, w2_t, w2_t_sf, swiglu_limit + ) + _fp8_weight_grad_accumulate(layout, *operands, x_t, x_t_sf, dw13, dw2) + return dx + + +@_moe_experts_fp8_bwd_accumulate.register_fake +def _(dout, num_tokens_per_expert, gate_up, x_t, x_t_sf, w13_t, w13_t_sf, w2_t, w2_t_sf, swiglu_limit, dw13, dw2): + return dout.new_empty(dout.shape) + + +@torch.library.custom_op("prime_kernels::moe_experts_fp8_backward_data", mutates_args=()) +def _moe_experts_fp8_bwd_data( + dout: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + gate_up: torch.Tensor, + w13_t: torch.Tensor, + w13_t_sf: torch.Tensor, + w2_t: torch.Tensor, + w2_t_sf: torch.Tensor, + swiglu_limit: float, +) -> list[torch.Tensor]: + """The data-gradient half of `moe_experts_fp8_backward_accumulate`: ``dx``, then the operands + `moe_experts_fp8_weight_grad_accumulate` reads, so a caller can send ``dx`` on before the + weight-gradient GEMMs run.""" + dout = dout.contiguous() + if dout.shape[0] == 0: + return [dout.new_empty(dout.shape)] + with torch.cuda.device(dout.device): + dx, _, operands = _fp8_backward_data( + dout, num_tokens_per_expert, gate_up, w13_t, w13_t_sf, w2_t, w2_t_sf, swiglu_limit + ) + return [dx, *operands] + + +@_moe_experts_fp8_bwd_data.register_fake +def _(dout, num_tokens_per_expert, gate_up, w13_t, w13_t_sf, w2_t, w2_t_sf, swiglu_limit): + hidden, two_intermediate = w13_t.shape[1], w13_t.shape[2] + Mp, G, f8 = fp8_impl.padded_rows(dout.shape[0], w13_t.shape[0]), fp8_impl.GROUP, torch.float8_e4m3fn + return [ + dout.new_empty(dout.shape), + dout.new_empty(Mp * hidden, dtype=f8), dout.new_empty(Mp // G, hidden, dtype=torch.float32), + dout.new_empty(Mp * two_intermediate // 2, dtype=f8), dout.new_empty(Mp // G, two_intermediate // 2, dtype=torch.float32), + dout.new_empty(Mp * two_intermediate, dtype=f8), dout.new_empty(Mp // G, two_intermediate, dtype=torch.float32), + ] + + +@torch.library.custom_op("prime_kernels::moe_experts_fp8_weight_grad_accumulate", mutates_args=("dw13", "dw2")) +def _moe_experts_fp8_wgrad_accumulate( + num_tokens_per_expert: torch.Tensor, + dout_t: torch.Tensor, + dout_t_sf: torch.Tensor, + h_t: torch.Tensor, + h_t_sf: torch.Tensor, + dgu_t: torch.Tensor, + dgu_t_sf: torch.Tensor, + x_t: torch.Tensor, + x_t_sf: torch.Tensor, + num_rows: int, + dw13: torch.Tensor, + dw2: torch.Tensor, +) -> None: + """The weight-gradient half of `moe_experts_fp8_backward_accumulate`, from the operands + `moe_experts_fp8_backward_data` returned for the same ``num_rows`` rows.""" + if num_rows == 0: + return + with torch.cuda.device(dout_t.device): + layout = fp8_impl.Layout(num_tokens_per_expert, num_rows) + _fp8_weight_grad_accumulate(layout, dout_t, dout_t_sf, h_t, h_t_sf, dgu_t, dgu_t_sf, x_t, x_t_sf, dw13, dw2) + + +def _fp8_setup_context(ctx, inputs, output): + _, w1, w3, _, num_tokens_per_expert, swiglu_limit, save_for_backward = inputs + ctx.saved_for_backward = save_for_backward + ctx.save_for_backward(num_tokens_per_expert, *output[1:]) + ctx.swiglu_limit = swiglu_limit + ctx.packed = w3 is None + ctx.set_materialize_grads(False) + + +def _fp8_backward(ctx, grads): + if not ctx.saved_for_backward: + raise RuntimeError("the FP8 expert forward ran without saving for backward but is being differentiated") + num_tokens_per_expert, *saved = ctx.saved_tensors + dx, dw1, dw3, dw2 = _moe_experts_fp8_bwd(grads[0], num_tokens_per_expert, *saved, ctx.swiglu_limit, ctx.packed) + return dx, dw1, None if ctx.packed else dw3, dw2, None, None, None + + +_moe_experts_fp8_fwd.register_autograd(_fp8_backward, setup_context=_fp8_setup_context) + + +def moe_experts( + x: torch.Tensor, + gate_proj: torch.Tensor, + up_proj: torch.Tensor | None, + down_proj: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + swiglu_limit: float, + fp8: bool = False, +) -> torch.Tensor: + """Expert MLP over tokens grouped by expert, differentiable in ``x`` and the three weights. + + For expert ``e`` and its rows ``x_e``: ``down_proj[e] @ (silu(min(gate, l)) * clamp(up, -l, + l))`` with ``gate = gate_proj[e] @ x_e`` and ``up = up_proj[e] @ x_e``, both rounded to bf16 + first; routing scores are applied by the caller. + + Args: + x: ``[rows, hidden]`` bf16, contiguous. Expert ``e`` owns the ``num_tokens_per_expert[e]`` + rows after those of experts ``< e``; rows after the last group come out as zeros. + gate_proj, up_proj: ``[E, intermediate, hidden]`` bf16. Pass ``up_proj=None`` and the + packed ``[E, 2 * intermediate, hidden]`` weight (gate rows first) as ``gate_proj``. + down_proj: ``[E, hidden, intermediate]`` bf16, contiguous. + num_tokens_per_expert: ``[E]`` integer group sizes on the device, each a multiple of + ``TOKEN_GROUP_ALIGNMENT``, summing to at most ``rows`` (checked asynchronously). + swiglu_limit: the clamp ``l``. + fp8: run the GEMMs in blockwise FP8 (DeepSeek-V3 recipe, through DeepGEMM) for the + forward, the data gradients and the weight gradients; needs hidden and intermediate + sizes that are multiples of 128. + + Returns: + ``[rows, hidden]`` bf16. + """ + for name, tensor in (("x", x), ("gate_proj", gate_proj), ("up_proj", up_proj), ("down_proj", down_proj)): + if tensor is None: + continue + if tensor.dtype != torch.bfloat16 or not tensor.is_cuda: + raise ValueError(f"{name} must be a bfloat16 CUDA tensor, got {tensor.dtype} on {tensor.device}") + if not x.is_contiguous() or not down_proj.is_contiguous(): + raise ValueError("x and down_proj must be contiguous") + reason = unsupported_shape_reason(x.shape[1], down_proj.shape[2], fp8=fp8) + if reason is not None: + raise ValueError(reason) + if fp8: + if importlib.util.find_spec("deep_gemm") is None: + raise RuntimeError("moe_experts(fp8=True) requires DeepGEMM (deep_gemm), which is not installed") + save = torch.is_grad_enabled() and any( + t is not None and t.requires_grad for t in (x, gate_proj, up_proj, down_proj) + ) + out, *_ = _moe_experts_fp8_fwd( + x, gate_proj, up_proj, down_proj, num_tokens_per_expert, float(swiglu_limit), save + ) + return out + out, _, _ = _moe_experts_fwd(x, gate_proj, up_proj, down_proj, num_tokens_per_expert, float(swiglu_limit)) + return out diff --git a/prime_kernels/moe_experts/quant.py b/prime_kernels/moe_experts/quant.py new file mode 100644 index 0000000..f36807d --- /dev/null +++ b/prime_kernels/moe_experts/quant.py @@ -0,0 +1,272 @@ +"""Gluon kernels that quantize a 128 x 128 tile both per row and per column (transposed). + +The per-column outputs are what the K-grouped weight-gradient GEMMs read: each column of the +tile becomes a contiguous run of 128 bytes. Tiles are staged in shared memory by TMA so they can +be read back in a row-major register layout (vectorized along columns, for the per-row groups) +and in a column-major one in which each thread owns a run of rows of one column (for the +per-column groups and their transposed stores), without per-element global address math. +""" + +from triton.experimental import gluon +from triton.experimental.gluon import language as gl +from triton.experimental.gluon.language.nvidia.hopper import fence_async_shared, mbarrier, tma +from triton.experimental.gluon.nvidia.hopper import TensorDescriptor + +TILE = gl.constexpr(128) +NUM_WARPS = 8 +# Row pass: 16 rows at a time, 8 consecutive columns per thread. +ROW_CHUNK = gl.constexpr(16) +ROW_LAYOUT = gl.constexpr(gl.BlockedLayout([1, 8], [2, 16], [NUM_WARPS, 1], [1, 0])) +# Column pass: 32 columns at a time; each thread owns 16 consecutive rows of one column and two +# neighbouring threads 32, so every transposed store fills whole 32-byte sectors. +COL_CHUNK = gl.constexpr(32) +COL_LAYOUT = gl.constexpr(gl.BlockedLayout([16, 1], [2, 16], [4, 2], [0, 1])) +SMEM_LAYOUT = gl.constexpr(gl.NVMMASharedLayout(swizzle_byte_width=128, element_bitwidth=16, rank=2)) + + +def descriptor(tensor): + return TensorDescriptor.from_tensor(tensor, [128, 128], SMEM_LAYOUT.value) + + +@gluon.jit +def _tile(t, tile_expert_ptr, tile_row_start_ptr, tile_row_end_ptr, expert_row_start_ptr, counts_ptr): + """Tile ``t``'s expert and source rows, where its rows sit in its expert's K-grouped block, and + how many of them are head rows (see `fp8`).""" + expert = gl.load(tile_expert_ptr + t) + row_start = gl.load(tile_row_start_ptr + t) + num_valid = gl.load(tile_row_end_ptr + t) - row_start + expert_start = gl.load(expert_row_start_ptr + expert) + local_tile = (row_start - expert_start) // 128 + group_start = (t - local_tile).to(gl.int64) * 128 + group_rows = gl.cdiv(gl.load(counts_ptr + expert).to(gl.int32), 128) * 128 + head = gl.where(local_tile == 0, gl.minimum((128 - expert_start % 128) % 128, num_valid), 0) + return expert, row_start, num_valid, local_tile * 128, group_start, group_rows, head + + +@gluon.jit +def _store_rows(v, out_ptr, sf_ptr, first_row, rows, cols, out_cols, M, sf_row, mask=None): + """``v`` [rows, 128] quantized per row into rows ``first_row + rows`` of an ``[M, out_cols]`` + FP8 tensor with MN-major ``[K / 128, M]`` scales; ``sf_row`` is this K group. Offsets inside + the tile are 32-bit, only its base is 64-bit.""" + scale = gl.maximum(gl.max(gl.abs(v), axis=1), 1e-10) / 448.0 + q = gl.minimum(gl.maximum(v * (1.0 / scale)[:, None], -448.0), 448.0).to(gl.float8e4nv) + first_row = first_row.to(gl.int64) + out = out_ptr + first_row * out_cols + (rows[:, None] * out_cols + cols[None, :]) + sf = sf_ptr + (sf_row * M + first_row) + rows + if mask is None: + gl.store(out, q) + gl.store(sf, scale) + else: + gl.store(out, q, mask=mask[:, None]) + gl.store(sf, scale, mask=mask) + + +@gluon.jit +def _store_cols(v, out_ptr, sf_ptr, t, rows, cols, out_cols, local_row, group_start, group_rows): + """``v`` [128 rows, cols] quantized per column into the K-grouped layout: an expert's block is + ``[out_cols, group_rows]`` at element ``group_start * out_cols``, scales ``[Mp / 128, + out_cols]``. Rows past the expert's end must be zeros.""" + scale = gl.maximum(gl.max(gl.abs(v), axis=0), 1e-10) / 448.0 + q = gl.minimum(gl.maximum(v * (1.0 / scale)[None, :], -448.0), 448.0).to(gl.float8e4nv) + base = out_ptr + group_start * out_cols + local_row + gl.store(base + (cols[None, :] * group_rows + rows[:, None]), q) + gl.store(sf_ptr + t.to(gl.int64) * out_cols + cols, scale) + + +@gluon.jit +def _load_tiles(desc, row, col0, col1, buf0, buf1): + """TMA loads of the ``[TILE, TILE]`` tiles of ``desc`` at ``(row, col0)`` and ``(row, col1)``.""" + bar = gl.allocate_shared_memory(gl.int64, [1], mbarrier.MBarrierLayout()) + mbarrier.init(bar, count=1) + fence_async_shared() + mbarrier.expect(bar, 2 * TILE * TILE * 2) + tma.async_copy_global_to_shared(desc, [row, col0], bar, buf0) + tma.async_copy_global_to_shared(desc, [row, col1], bar, buf1) + mbarrier.wait(bar, 0) + + +@gluon.jit +def swiglu_kernel( + gate_up_desc, + h_ptr, + h_sf_ptr, + head_ptr, + head_sf_ptr, + tile_expert_ptr, + tile_row_start_ptr, + tile_row_end_ptr, + expert_row_start_ptr, + counts_ptr, + num_tiles_ptr, + M, + M_HEAD, + inter, + limit, +): + """h = clamped_swiglu(gate, up) from tile ``t`` of the padded gate/up GEMM output, quantized per + row into the token-major layout (and its head rows into the head operand).""" + t = gl.program_id(0) + c = gl.program_id(1) + if t >= gl.load(num_tiles_ptr): + return + expert, row_start, num_valid, _, _, _, head = _tile( + t, tile_expert_ptr, tile_row_start_ptr, tile_row_end_ptr, expert_row_start_ptr, counts_ptr + ) + gate_s = gl.allocate_shared_memory(gl.bfloat16, [TILE, TILE], SMEM_LAYOUT) + up_s = gl.allocate_shared_memory(gl.bfloat16, [TILE, TILE], SMEM_LAYOUT) + _load_tiles(gate_up_desc, t * TILE, c * TILE, inter + c * TILE, gate_s, up_s) + cols = c * TILE + gl.arange(0, TILE, layout=gl.SliceLayout(0, ROW_LAYOUT)) + for r in gl.static_range(0, TILE, ROW_CHUNK): + rows = r + gl.arange(0, ROW_CHUNK, layout=gl.SliceLayout(1, ROW_LAYOUT)) + gate = gate_s.slice(r, ROW_CHUNK, dim=0).load(ROW_LAYOUT).to(gl.float32) + up = up_s.slice(r, ROW_CHUNK, dim=0).load(ROW_LAYOUT).to(gl.float32) + gate_c = gl.minimum(gate, limit) + up_c = gl.minimum(gl.maximum(up, -limit), limit) + h = gate_c / (1.0 + gl.exp(-gate_c)) * up_c + _store_rows(h, h_ptr, h_sf_ptr, row_start, rows, cols, inter, M, c, rows < num_valid) + if r < head: + _store_rows(h, head_ptr, head_sf_ptr, expert * TILE, rows, cols, inter, M_HEAD, c, rows < head) + + +@gluon.jit +def _dswiglu(dh, gate, up, limit): + gate_c = gl.minimum(gate, limit) + up_c = gl.minimum(gl.maximum(up, -limit), limit) + sig = 1.0 / (1.0 + gl.exp(-gate_c)) + silu = gate_c * sig + dgate = gl.where(gate <= limit, dh * up_c * (sig * (1.0 + gate_c * (1.0 - sig))), 0.0) + dup = gl.where((up >= -limit) & (up <= limit), dh * silu, 0.0) + return dgate, dup, silu * up_c + + +@gluon.jit +def swiglu_backward_kernel( + dh_desc, + gate_up_desc, + dgu_ptr, + dgu_sf_ptr, + head_ptr, + head_sf_ptr, + dgu_t_ptr, + dgu_t_sf_ptr, + h_t_ptr, + h_t_sf_ptr, + tile_expert_ptr, + tile_row_start_ptr, + tile_row_end_ptr, + expert_row_start_ptr, + counts_ptr, + num_tiles_ptr, + M, + M_HEAD, + inter, + limit, +): + """dgate | dup from tile ``t`` of the padded dh and gate/up: per row into the token-major + layout (and the head operand) for dx, per column into the K-grouped layout for the gate/up + weight gradient; and h, per column, for the down weight gradient. The column pass recomputes + the SwiGLU backward from the staged inputs. The padding rows of the tile must be zeros (they + are: see `fp8.Layout`), which makes their dgate, dup and h zeros too.""" + t = gl.program_id(0) + c = gl.program_id(1) + if t >= gl.load(num_tiles_ptr): + return + expert, row_start, num_valid, local_row, group_start, group_rows, head = _tile( + t, tile_expert_ptr, tile_row_start_ptr, tile_row_end_ptr, expert_row_start_ptr, counts_ptr + ) + dh_s = gl.allocate_shared_memory(gl.bfloat16, [TILE, TILE], SMEM_LAYOUT) + gate_s = gl.allocate_shared_memory(gl.bfloat16, [TILE, TILE], SMEM_LAYOUT) + up_s = gl.allocate_shared_memory(gl.bfloat16, [TILE, TILE], SMEM_LAYOUT) + bar = gl.allocate_shared_memory(gl.int64, [1], mbarrier.MBarrierLayout()) + mbarrier.init(bar, count=1) + fence_async_shared() + mbarrier.expect(bar, 3 * TILE * TILE * 2) + tma.async_copy_global_to_shared(dh_desc, [t * TILE, c * TILE], bar, dh_s) + tma.async_copy_global_to_shared(gate_up_desc, [t * TILE, c * TILE], bar, gate_s) + tma.async_copy_global_to_shared(gate_up_desc, [t * TILE, inter + c * TILE], bar, up_s) + mbarrier.wait(bar, 0) + num_k = inter // TILE + + cols = c * TILE + gl.arange(0, TILE, layout=gl.SliceLayout(0, ROW_LAYOUT)) + for r in gl.static_range(0, TILE, ROW_CHUNK): + rows = r + gl.arange(0, ROW_CHUNK, layout=gl.SliceLayout(1, ROW_LAYOUT)) + valid = rows < num_valid + dh = dh_s.slice(r, ROW_CHUNK, dim=0).load(ROW_LAYOUT).to(gl.float32) + gate = gate_s.slice(r, ROW_CHUNK, dim=0).load(ROW_LAYOUT).to(gl.float32) + up = up_s.slice(r, ROW_CHUNK, dim=0).load(ROW_LAYOUT).to(gl.float32) + dgate, dup, _ = _dswiglu(dh, gate, up, limit) + _store_rows(dgate, dgu_ptr, dgu_sf_ptr, row_start, rows, cols, 2 * inter, M, c, valid) + _store_rows(dup, dgu_ptr, dgu_sf_ptr, row_start, rows, cols + inter, 2 * inter, M, c + num_k, valid) + if r < head: + in_head = rows < head + _store_rows(dgate, head_ptr, head_sf_ptr, expert * TILE, rows, cols, 2 * inter, M_HEAD, c, in_head) + _store_rows( + dup, head_ptr, head_sf_ptr, expert * TILE, rows, cols + inter, 2 * inter, M_HEAD, c + num_k, in_head + ) + + rows = gl.arange(0, TILE, layout=gl.SliceLayout(1, COL_LAYOUT)) + for cc in gl.static_range(0, TILE, COL_CHUNK): + chunk_cols = c * TILE + cc + gl.arange(0, COL_CHUNK, layout=gl.SliceLayout(0, COL_LAYOUT)) + dh = dh_s.slice(cc, COL_CHUNK, dim=1).load(COL_LAYOUT).to(gl.float32) + gate = gate_s.slice(cc, COL_CHUNK, dim=1).load(COL_LAYOUT).to(gl.float32) + up = up_s.slice(cc, COL_CHUNK, dim=1).load(COL_LAYOUT).to(gl.float32) + dgate, dup, h = _dswiglu(dh, gate, up, limit) + _store_cols(dgate, dgu_t_ptr, dgu_t_sf_ptr, t, rows, chunk_cols, 2 * inter, local_row, group_start, group_rows) + _store_cols( + dup, dgu_t_ptr, dgu_t_sf_ptr, t, rows, chunk_cols + inter, 2 * inter, local_row, group_start, group_rows + ) + _store_cols(h, h_t_ptr, h_t_sf_ptr, t, rows, chunk_cols, inter, local_row, group_start, group_rows) + + +@gluon.jit +def quantize_kernel( + src_desc, + q_ptr, + q_sf_ptr, + qt_ptr, + qt_sf_ptr, + tile_expert_ptr, + tile_row_start_ptr, + tile_row_end_ptr, + expert_row_start_ptr, + counts_ptr, + num_tiles_ptr, + Mp, + C, + COLUMNS: gl.constexpr, +): + """Tile ``t`` of the token-major ``src`` per row into the padded layout (its padding rows as + zeros) and, with ``COLUMNS``, per column into the K-grouped layout.""" + t = gl.program_id(0) + c = gl.program_id(1) + if t >= gl.load(num_tiles_ptr): + return + _, row_start, num_valid, local_row, group_start, group_rows, _ = _tile( + t, tile_expert_ptr, tile_row_start_ptr, tile_row_end_ptr, expert_row_start_ptr, counts_ptr + ) + buf = gl.allocate_shared_memory(gl.bfloat16, [TILE, TILE], SMEM_LAYOUT) + bar = gl.allocate_shared_memory(gl.int64, [1], mbarrier.MBarrierLayout()) + mbarrier.init(bar, count=1) + fence_async_shared() + mbarrier.expect(bar, TILE * TILE * 2) + tma.async_copy_global_to_shared(src_desc, [row_start, c * TILE], bar, buf) + mbarrier.wait(bar, 0) + if num_valid < TILE: + # The tile's last rows belong to the next expert: zero them once, in shared memory. + for r in gl.static_range(0, TILE, ROW_CHUNK): + rows = r + gl.arange(0, ROW_CHUNK, layout=gl.SliceLayout(1, ROW_LAYOUT)) + chunk = buf.slice(r, ROW_CHUNK, dim=0) + chunk.store(gl.where((rows < num_valid)[:, None], chunk.load(ROW_LAYOUT), 0.0)) + gl.barrier() + + cols = c * TILE + gl.arange(0, TILE, layout=gl.SliceLayout(0, ROW_LAYOUT)) + for r in gl.static_range(0, TILE, ROW_CHUNK): + rows = r + gl.arange(0, ROW_CHUNK, layout=gl.SliceLayout(1, ROW_LAYOUT)) + v = buf.slice(r, ROW_CHUNK, dim=0).load(ROW_LAYOUT).to(gl.float32) + _store_rows(v, q_ptr, q_sf_ptr, t * TILE, rows, cols, C, Mp, c) + if COLUMNS: + rows = gl.arange(0, TILE, layout=gl.SliceLayout(1, COL_LAYOUT)) + for cc in gl.static_range(0, TILE, COL_CHUNK): + chunk_cols = c * TILE + cc + gl.arange(0, COL_CHUNK, layout=gl.SliceLayout(0, COL_LAYOUT)) + v = buf.slice(cc, COL_CHUNK, dim=1).load(COL_LAYOUT).to(gl.float32) + _store_cols(v, qt_ptr, qt_sf_ptr, t, rows, chunk_cols, C, local_row, group_start, group_rows) diff --git a/tests/dsa_indexer_topk/bench_dsa_indexer_topk.py b/tests/dsa_indexer_topk/bench_dsa_indexer_topk.py new file mode 100644 index 0000000..d81d3e9 --- /dev/null +++ b/tests/dsa_indexer_topk/bench_dsa_indexer_topk.py @@ -0,0 +1,83 @@ +"""Benchmark `dsa_indexer_topk` against prime-rl's `dsv41_index_topk` at DeepSeek-V4.1 Flash shapes. + + python tests/dsa_indexer_topk/bench_dsa_indexer_topk.py [path/to/prime-rl/src] + +Shapes: 16k queries of one document, 32 index heads x 128, top-512, against 16k entries +(compress ratio 1) and 8k entries (ratio 2) on one GPU, and the last CP rank's 16k queries of a +131k-token document (CP=8) with the candidate filter active (2048 blocks of 8): plain, candidate +source (emit) and candidate consumer. With the prime-rl source path it also times prime-rl's op +and reports how much the two agree (mean |intersection| / |picks| per query). +""" + +import argparse +import sys + +import torch + +import prime_kernels + + +def bench(fn, iters=10): + fn() + torch.cuda.synchronize() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def overlap(a: torch.Tensor, b: torch.Tensor) -> float: + a, b = a.long(), b.long() + a_sorted = a.sort(-1).values + found = torch.searchsorted(a_sorted, b.contiguous()) + hit = (a_sorted.gather(1, found.clamp_max(a.shape[1] - 1)) == b) & (b >= 0) + return float(hit.sum() / (b >= 0).sum().clamp_min(1)) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("prime_rl_src", nargs="?", help="prime-rl's src/ directory, to compare against its op") + args = parser.parse_args() + ours = prime_kernels.load("dsa_indexer_topk").dsv41_index_topk + theirs = None + if args.prime_rl_src: + sys.path.insert(0, args.prime_rl_src) + from prime_rl.trainer.models.kernels.dsv41_indexer import dsv41_index_topk as theirs + + torch.manual_seed(0) + device = torch.device("cuda") + n_queries, heads, dim, topk, block, topk_blocks = 16384, 32, 128, 512, 8, 2048 + cases = [("ratio 1, 16k entries", 16384, 1, 0), ("ratio 2, 8k entries", 8192, 2, 0)] + cases.append(("CP rank 7/8, 131k entries", 131072, 1, 131072 - n_queries)) + for name, n_entries, ratio, query_offset in cases: + q = torch.randn(n_queries, heads, dim, device=device, dtype=torch.bfloat16) + k = torch.randn(n_entries, dim, device=device, dtype=torch.bfloat16) + w = torch.randn(n_queries, heads, device=device, dtype=torch.bfloat16) + pos = torch.arange(query_offset, query_offset + n_queries, device=device) + ks = torch.zeros(n_queries, dtype=torch.int32, device=device) + ke = ((pos + 1) // ratio).int() + modes = [("plain", None, False)] + if n_entries > topk_blocks * block: + _, cand = (theirs or ours)(q, k, w, ks, ke, topk, None, True, n_entries, block, topk_blocks) + modes += [("emit", None, True), ("consume", cand, False)] + for mode, cand_in, emit in modes: + + def run(op): + return op(q, k, w, ks, ke, topk, cand_in, emit, n_entries, block, topk_blocks) + + line = f"{name:28s} {mode:8s} fused {bench(lambda: run(ours)):8.2f} ms" + if theirs is not None: + line += f" | prime-rl {bench(lambda: run(theirs)):8.2f} ms" + (p_ours, c_ours), (p_theirs, c_theirs) = run(ours), run(theirs) + line += f" | top-k agreement {overlap(p_theirs, p_ours):.5f}" + if emit: + line += f", candidate agreement {overlap(c_theirs, c_ours):.5f}" + print(line, flush=True) + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/tests/dsa_indexer_topk/test_dsa_indexer_topk.py b/tests/dsa_indexer_topk/test_dsa_indexer_topk.py new file mode 100644 index 0000000..ddc5e32 --- /dev/null +++ b/tests/dsa_indexer_topk/test_dsa_indexer_topk.py @@ -0,0 +1,200 @@ +import pytest +import torch + +import prime_kernels + +REASON = prime_kernels.unavailable_reason("dsa_indexer_topk") +pytestmark = pytest.mark.skipif(REASON is not None, reason=REASON or "") + +FP8_MAX = 448.0 + + +def quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-row FP8 e4m3 with power-of-two scales, dequantized back to fp64.""" + x = x.double() + scale = torch.exp2(torch.ceil(torch.log2(x.abs().amax(-1, keepdim=True).clamp_min(1e-10) / FP8_MAX))) + x_q = (x / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + return x_q.double(), scale.squeeze(-1) + + +def reference_scores(q, k, w, ks, ke) -> torch.Tensor: + """`(S_q, S_k)` fp64 indexer scores of the FP8 operands, -inf outside `[ks, ke)`.""" + q_q, q_scale = quantize(q) + k_q, k_scale = quantize(k) + w = w.double() * q_scale + scores = torch.empty(q.shape[0], k.shape[0], dtype=torch.float64, device=q.device) + for r0 in range(0, q.shape[0], 256): + logits = torch.einsum("thd,ed->the", q_q[r0 : r0 + 256], k_q).relu() + scores[r0 : r0 + 256] = torch.einsum("the,th->te", logits, w[r0 : r0 + 256]) * k_scale + cols = torch.arange(k.shape[0], device=q.device) + valid = (cols >= ks[:, None]) & (cols < ke[:, None]) + return scores.masked_fill(~valid, float("-inf")) + + +def doc_local_blocks(scores, ks, width, block_size) -> torch.Tensor: + """`(S_q, n_blocks)` best score of every doc-local block, -inf where it holds no readable entry.""" + n_blocks = -(-width // block_size) + cols = ks[:, None].long() + torch.arange(n_blocks * block_size, device=scores.device) + local = scores.gather(1, cols.clamp_max(scores.shape[1] - 1)) + local = local.masked_fill(cols >= scores.shape[1], float("-inf")) + return local.unflatten(1, (n_blocks, block_size)).amax(-1) + + +def check_picks(picks: torch.Tensor, allowed: torch.Tensor, k: int) -> None: + """`picks` are `k` best columns of `allowed` (fp64 scores, -inf excluded) up to ties.""" + n_allowed = (allowed > float("-inf")).sum(-1) + n_taken = (picks >= 0).sum(-1) + torch.testing.assert_close(n_taken, n_allowed.clamp_max(k)) + # -1 padding is at the tail. + assert bool(((picks >= 0).long().diff(dim=-1) <= 0).all()) + safe = picks.clamp_min(0) + picked = allowed.gather(1, safe).masked_fill(picks < 0, float("inf")) + assert bool((picked > float("-inf")).all()), "a pick is outside the allowed entries" + sorted_picks = safe.masked_fill(picks < 0, -1).sort(-1).values + dup = (sorted_picks.diff(dim=-1) == 0) & (sorted_picks[:, 1:] >= 0) + assert not bool(dup.any()), "duplicate picks" + kth = allowed.topk(min(k, allowed.shape[1]), dim=-1).values[:, -1:] + # Hopper FP8 tensor cores accumulate with a reduced-precision mantissa (prime-rl's Triton + # indexer too), so near-ties can swap. + tol = 2e-3 * allowed.masked_fill(allowed == float("-inf"), 0).abs().amax(-1, keepdim=True) + 1e-12 + worst = picked.amin(-1, keepdim=True) + has = n_taken[:, None] > 0 + assert bool(((worst >= kth - tol) | ~has).all()), "a pick scores below the k-th best" + + +def packed_bounds(doc_lens, ratio, device, n_queries=None, query_offset=0): + """`ks, ke` in entry coordinates for queries of a packed row of documents (all queries by default).""" + ks, ke = [], [] + first_entry = 0 + for n in doc_lens: + pos = torch.arange(n) + ks.append(torch.full((n,), first_entry)) + ke.append(first_entry + (pos + 1) // ratio) + first_entry += n // ratio + ks, ke = torch.cat(ks), torch.cat(ke) + n_queries = ks.shape[0] - query_offset if n_queries is None else n_queries + sl = slice(query_offset, query_offset + n_queries) + return ks[sl].int().to(device), ke[sl].int().to(device), first_entry + + +def make_inputs(n_queries, n_entries, heads, dim, device, seed=0): + gen = torch.Generator(device="cpu").manual_seed(seed) + q = torch.randn(n_queries, heads, dim, generator=gen).to(device, torch.bfloat16) + k = torch.randn(n_entries, dim, generator=gen).to(device, torch.bfloat16) + w = torch.randn(n_queries, heads, generator=gen).to(device, torch.bfloat16) + return q, k, w + + +@pytest.fixture(scope="module") +def kernel(): + return prime_kernels.load("dsa_indexer_topk") + + +@pytest.mark.parametrize( + "doc_lens,ratio,heads,dim,topk,block_size,topk_blocks", + [ + ([700], 1, 16, 64, 64, 8, 32), + ([300, 5, 517, 1, 233], 1, 32, 128, 64, 8, 24), + ([1000, 37, 600], 2, 32, 128, 128, 8, 40), + ([129, 400, 3], 2, 16, 128, 16, 4, 8), + ], +) +def test_matches_reference(kernel, doc_lens, ratio, heads, dim, topk, block_size, topk_blocks): + device = torch.device("cuda") + ks, ke, n_entries = packed_bounds(doc_lens, ratio, device) + width = max(n // ratio for n in doc_lens) + q, k, w = make_inputs(ks.shape[0], n_entries, heads, dim, device) + scores = reference_scores(q, k, w, ks, ke) + + picks, cand = kernel.dsv41_index_topk(q, k, w, ks, ke, topk, None, True, width, block_size, topk_blocks) + check_picks(picks, scores, topk) + picks_plain, empty = kernel.dsv41_index_topk(q, k, w, ks, ke, topk, None, False, width, block_size, topk_blocks) + assert empty.numel() == 0 + check_picks(picks_plain, scores, topk) + + # Candidate blocks: the newest block is always kept, the rest are the best by block max. + blocks = doc_local_blocks(scores, ks, width, block_size) + newest = ((ke - ks - 1).clamp_min(0) // block_size).long() + assert bool((cand == newest[:, None]).any(-1).all()), "the newest block is not pinned" + pinned = blocks.scatter(1, newest[:, None], float("inf")) + check_picks(cand.long(), pinned, topk_blocks) + + # Consumer layer: rank only entries inside the given blocks. + picks_c, empty = kernel.dsv41_index_topk(q, k, w, ks, ke, topk, cand, False, width, block_size, topk_blocks) + assert empty.numel() == 0 + allowed = torch.full_like(scores, float("-inf")) + local = cand.long()[:, :, None] * block_size + torch.arange(block_size, device=device) + entries = (ks[:, None, None].long() + local).flatten(1) + ok = (cand.long()[:, :, None] >= 0).expand_as(local).flatten(1) & (entries < ke[:, None]) + rows = torch.arange(ks.shape[0], device=device)[:, None].expand_as(entries) + allowed[rows[ok], entries[ok]] = scores[rows[ok], entries[ok]] + check_picks(picks_c, allowed, topk) + + +def test_v41_shapes_with_binding_candidates(kernel): + """V4.1 Flash indexer shapes where the candidate filter binds (more than 2048 blocks of 8).""" + device = torch.device("cuda") + n_entries, n_queries = 20480, 512 + ks, ke, _ = packed_bounds([n_entries], 1, device, n_queries=n_queries, query_offset=n_entries - n_queries) + q, k, w = make_inputs(n_queries, n_entries, 32, 128, device, seed=1) + scores = reference_scores(q, k, w, ks, ke) + + picks, cand = kernel.dsv41_index_topk(q, k, w, ks, ke, 512, None, True, n_entries, 8, 2048) + check_picks(picks, scores, 512) + picks_plain, _ = kernel.dsv41_index_topk(q, k, w, ks, ke, 512, None, False, n_entries, 8, 2048) + check_picks(picks_plain, scores, 512) + blocks = doc_local_blocks(scores, ks, n_entries, 8) + newest = ((ke - ks - 1).clamp_min(0) // 8).long() + check_picks(cand.long(), blocks.scatter(1, newest[:, None], float("inf")), 2048) + + picks_c, _ = kernel.dsv41_index_topk(q, k, w, ks, ke, 512, cand, False, n_entries, 8, 2048) + allowed = torch.full_like(scores, float("-inf")) + entries = (cand.long()[:, :, None] * 8 + torch.arange(8, device=device)).flatten(1) + ok = (cand.long()[:, :, None] >= 0).expand(-1, -1, 8).flatten(1) & (entries < ke[:, None]) + rows = torch.arange(n_queries, device=device)[:, None].expand_as(entries) + allowed[rows[ok], entries[ok]] = scores[rows[ok], entries[ok]] + check_picks(picks_c, allowed, 512) + + +def test_rescoring_falls_back_when_the_threshold_overshoots(kernel): + """The rescoring pass keeps everything should its threshold leave fewer than `topk` entries.""" + from prime_kernels.dsa_indexer_topk import indexer + from prime_kernels.dsa_indexer_topk.quant import per_token_quant_fp8 + + device = torch.device("cuda") + ks, ke, n_entries = packed_bounds([900, 300], 1, device) + q, k, w = make_inputs(ks.shape[0], n_entries, 16, 64, device, seed=2) + scores = reference_scores(q, k, w, ks, ke) + group, topk = 4, 32 + n_groups = -(-900 // group) + blocks = torch.arange(n_groups, device=device, dtype=torch.int32).expand(ks.shape[0], -1).contiguous() + q_fp8, q_scale = per_token_quant_fp8(q) + k_fp8, k_scale = per_token_quant_fp8(k) + w = (w.float() * q_scale).contiguous() + picks = torch.empty(ks.shape[0], topk, dtype=torch.int64, device=device) + tau = torch.full((ks.shape[0],), float("inf"), device=device) + indexer._gather_topk(q_fp8, k_fp8, k_scale, w, ks, ke, blocks, tau, picks, topk, group) + check_picks(picks, scores, topk) + + +def test_no_readable_entries(kernel): + device = torch.device("cuda") + q, k, w = make_inputs(4, 8, 16, 64, device) + ks = torch.tensor([0, 0, 4, 4], dtype=torch.int32, device=device) + ke = torch.tensor([0, 1, 4, 6], dtype=torch.int32, device=device) + picks, cand = kernel.dsv41_index_topk(q, k, w, ks, ke, 4, None, True, 4, 2, 3) + assert picks[0].tolist() == [-1] * 4 and picks[2].tolist() == [-1] * 4 + assert picks[1].tolist() == [0, -1, -1, -1] + assert sorted(picks[3].tolist()) == [-1, -1, 4, 5] + # A query without entries keeps block 0, as the newest block is pinned. + assert cand[0].tolist() == [0, -1, -1] and cand[3].tolist() == [0, -1, -1] + + +def test_fake_tensor(kernel): + q = torch.empty(5, 32, 128, dtype=torch.bfloat16, device="meta") + k = torch.empty(9, 128, dtype=torch.bfloat16, device="meta") + w = torch.empty(5, 32, dtype=torch.bfloat16, device="meta") + ks = torch.empty(5, dtype=torch.int32, device="meta") + picks, cand = kernel.dsv41_index_topk(q, k, w, ks, ks, 512, None, True, 9, 8, 2048) + assert picks.shape == (5, 512) and picks.dtype == torch.int64 + assert cand.shape == (5, 2048) and cand.dtype == torch.int32 diff --git a/tests/dsa_sparse_attn_bwd/bench_dsa_sparse_attn_bwd.py b/tests/dsa_sparse_attn_bwd/bench_dsa_sparse_attn_bwd.py new file mode 100644 index 0000000..4b739e3 --- /dev/null +++ b/tests/dsa_sparse_attn_bwd/bench_dsa_sparse_attn_bwd.py @@ -0,0 +1,84 @@ +"""Benchmark `dsa_sparse_attn_bwd` against the backward prime-rl runs today (cuDNN frontend's SM90 DSA +backward, `flash_attn_bwd_sm90`) at DeepSeek-V4.1 Flash shapes on one GPU, with FlashMLA's sparse +prefill forward (prime-rl's forward) for reference. + + python tests/dsa_sparse_attn_bwd/bench_dsa_sparse_attn_bwd.py + +Shapes: 16384 queries x 64 heads x 512, kv = 16384 window positions + 8192 compressed entries, +640 slots per query (128-token causal window + top-512 random picks from the document's causal +compressed entries), for one 16k document and for 4 packed documents. TFLOP/s count every slot +(empty ones too), 4 * 64 * 512 * 640 flops per query forward and 10 * 64 * 512 * 640 backward +(S, dP, dQ, dK, dV), like the training trace's accounting. +""" + +import math +import sys +from pathlib import Path + +import torch + +import prime_kernels + +sys.path.insert(0, str(Path(__file__).parent)) +from test_dsa_sparse_attn_bwd import make_indices, rel_err # noqa: E402 + + +def bench(fn, iters=20): + fn() + torch.cuda.synchronize() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def main(): + from cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm90 import flash_attn_bwd_sm90 + from flash_mla import flash_mla_sparse_fwd + + ours = prime_kernels.load("dsa_sparse_attn_bwd").sparse_attn_backward_flat + torch.manual_seed(0) + h, d = 64, 512 + print( + f"{'case':10s} {'empty':>6s} | {'FlashMLA fwd':>17s} | {'cuDNN bwd':>17s} | {'ours bwd':>17s} | speedup | ours vs cuDNN (rel L2: dq, dkv, dsink)" + ) + for name, lens in [("1 doc", [16384]), ("4 docs", [6000, 3000, 4384, 3000])]: + idx, n = make_indices(lens) + t, k = idx.shape + q = torch.randn(t, h, d, device="cuda", dtype=torch.bfloat16) + kv = torch.randn(n, d, device="cuda", dtype=torch.bfloat16) + sinks = torch.randn(h, device="cuda") + grad_out = torch.randn(t, h, d, device="cuda", dtype=torch.bfloat16) + scale = 1 / math.sqrt(d) + + def fwd(): + return flash_mla_sparse_fwd(q, kv.view(n, 1, d), idx.view(t, 1, k), scale, d, attn_sink=sinks) + + out, _, lse = fwd() + + def theirs(): + return flash_attn_bwd_sm90( + q, kv, out, grad_out, lse, attn_sink=sinks, softmax_scale=scale, topk_idxs=idx, need_d_sink=True + ) + + def mine(): + return ours(q, kv, out, grad_out, lse, idx, sinks, scale) + + errs = ", ".join(f"{rel_err(a, b):.1e}" for a, b in zip(mine(), theirs())) + flops = 2 * t * k * h * d + t_fwd, t_theirs, t_mine = bench(fwd), bench(theirs), bench(mine) + empty = (idx < 0).float().mean().item() + print( + f"{name:10s} {empty:6.1%} | {t_fwd:6.2f} ms {2 * flops / t_fwd / 1e9:4.0f} TF/s | " + f"{t_theirs:6.2f} ms {5 * flops / t_theirs / 1e9:4.0f} TF/s | {t_mine:6.2f} ms {5 * flops / t_mine / 1e9:4.0f} TF/s | " + f"{t_theirs / t_mine:6.2f}x | {errs}", + flush=True, + ) + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/tests/dsa_sparse_attn_bwd/test_dsa_sparse_attn_bwd.py b/tests/dsa_sparse_attn_bwd/test_dsa_sparse_attn_bwd.py new file mode 100644 index 0000000..b778f56 --- /dev/null +++ b/tests/dsa_sparse_attn_bwd/test_dsa_sparse_attn_bwd.py @@ -0,0 +1,190 @@ +import math + +import pytest +import torch + +import prime_kernels + +REASON = prime_kernels.unavailable_reason("dsa_sparse_attn_bwd") +pytestmark = pytest.mark.skipif(REASON is not None, reason=REASON or "") + +HEADS, DIM = 64, 512 +# Relative L2 error vs the fp32 reference. The error is dominated by bf16 rounding of the inputs +# the kernel shares with cuDNN (out, grad_out) and of dq / dkv themselves; cuDNN's SM90 backward +# lands on the same numbers (~2.5e-3). +REF_TOL = 1e-2 +CUDNN_TOL = 2e-3 + + +@pytest.fixture(autouse=True) +def seed(): + torch.manual_seed(0) + + +@pytest.fixture(scope="module") +def dsa(): + return prime_kernels.load("dsa_sparse_attn_bwd") + + +def rel_err(actual: torch.Tensor, expected: torch.Tensor) -> float: + return ((actual.double() - expected.double()).norm() / expected.double().norm().clamp_min(1e-30)).item() + + +def make_indices(seq_lens, window=128, topk=512, ratio=2): + """DeepSeek-V4.1 slot lists for packed documents: a causal sliding window over the token stream + (rows [0, t)) followed by `topk` random picks from the query's document-local, causal compressed + entries (rows [t, t + entries)); -1 marks an empty slot. Returns (indices, number of kv rows).""" + device = "cuda" + seq_lens = torch.tensor(seq_lens, device=device) + t = int(seq_lens.sum()) + cu = torch.cat([seq_lens.new_zeros(1), seq_lens.cumsum(0)]) + tok = torch.arange(t, device=device) + doc = torch.searchsorted(cu[1:], tok, right=True) + start = cu[doc] + slots = torch.maximum(start, tok - window + 1)[:, None] + torch.arange(window, device=device) + win = torch.where(slots <= tok[:, None], slots, -1) + ent_cu = torch.cat([seq_lens.new_zeros(1), (seq_lens // ratio).cumsum(0)]) + avail = (tok - start + 1) // ratio + max_avail = max(int(avail.max()), 1) + scores = torch.rand(t, max_avail, device=device) + scores = scores.masked_fill(torch.arange(max_avail, device=device) >= avail[:, None], -1.0) + val, local = scores.topk(min(topk, max_avail), dim=-1) + picks = torch.where(val >= 0, local + ent_cu[doc][:, None] + t, -1) + picks = torch.nn.functional.pad(picks, (0, topk - picks.shape[1]), value=-1) + return torch.cat([win, picks], dim=-1).int().contiguous(), t + int(ent_cu[-1]) + + +def reference(q, kv, indices, sinks, scale, grad_out, chunk=256): + """fp32 sink-aware sparse attention forward (out, sink-free lse) and backward (dq, dkv, dsinks).""" + t, h, d = q.shape + out = torch.empty(t, h, d, dtype=torch.float32, device=q.device) + lse = torch.empty(t, h, dtype=torch.float32, device=q.device) + dq = torch.empty_like(out) + dkv = torch.zeros(kv.shape[0], d, dtype=torch.float32, device=q.device) + dsinks = torch.zeros(h, dtype=torch.float32, device=q.device) + kv32 = kv.float() + for s in range(0, t, chunk): + e = min(t, s + chunk) + qi, gi, ii = q[s:e].float(), grad_out[s:e].float(), indices[s:e].long() + valid = ii >= 0 + k = kv32[ii.clamp(min=0)] + logits = torch.einsum("thd,tkd->thk", qi, k) * scale + logits = logits.masked_fill(~valid[:, None, :], float("-inf")) + lse[s:e] = torch.logsumexp(logits, -1) + lse_sink = torch.logaddexp(lse[s:e], sinks[None]) + p = torch.exp(logits - lse_sink[..., None]) + o = torch.einsum("thk,tkd->thd", p, k) + out[s:e] = o + delta = (o * gi).sum(-1) + ds = p * (torch.einsum("thd,tkd->thk", gi, k) - delta[..., None]) + dq[s:e] = torch.einsum("thk,tkd->thd", ds, k) * scale + dk = torch.einsum("thk,thd->tkd", ds, qi) * scale + torch.einsum("thk,thd->tkd", p, gi) + dkv.index_add_(0, ii.clamp(min=0).flatten(), (dk * valid[..., None]).flatten(0, 1)) + dsinks -= (torch.exp(sinks[None] - lse_sink) * delta).sum(0) + return out, lse, dq, dkv, dsinks + + +def make_problem(indices, n_rows): + t = indices.shape[0] + q = torch.randn(t, HEADS, DIM, device="cuda", dtype=torch.bfloat16) + kv = torch.randn(n_rows, DIM, device="cuda", dtype=torch.bfloat16) + grad_out = torch.randn(t, HEADS, DIM, device="cuda", dtype=torch.bfloat16) + sinks = torch.randn(HEADS, device="cuda") + scale = 1 / math.sqrt(DIM) + out, lse, dq, dkv, dsinks = reference(q, kv, indices, sinks, scale, grad_out) + inputs = (q, kv, out.to(torch.bfloat16), grad_out, lse, indices, sinks, scale) + return inputs, (dq, dkv, dsinks) + + +def shuffled(indices, frac_empty=0.0): + """Permute each query's slots and blank a fraction of them (the result must not depend on slot order).""" + perm = torch.rand(indices.shape, device=indices.device).argsort(-1) + out = indices.gather(1, perm) + if frac_empty: + out = torch.where(torch.rand(out.shape, device=out.device) < frac_empty, -1, out) + return out.int().contiguous() + + +CASES = { + "2 docs, V4.1 slots": lambda: make_indices([300, 212]), + "ragged slot count (k = 56, padded)": lambda: make_indices([5, 70, 1], window=16, topk=40), + "shuffled slots, 30% empty": lambda: (lambda i, n: (shuffled(i, 0.3), n))(*make_indices([257, 400])), + "4k tokens, 3 docs": lambda: make_indices([2000, 1500, 596]), +} + + +@pytest.mark.parametrize("case", list(CASES)) +def test_matches_reference(dsa, case): + indices, n_rows = CASES[case]() + inputs, expected = make_problem(indices, n_rows) + actual = dsa.sparse_attn_backward_flat(*inputs) + for name, a, e in zip(("dq", "dkv", "dsinks"), actual, expected): + assert a.isfinite().all(), name + assert rel_err(a, e) < REF_TOL, f"{name}: {rel_err(a, e):.2e}" + + +def test_empty_queries_and_duplicate_slots(dsa): + indices, n_rows = make_indices([90, 60], window=32, topk=64) + indices[::7] = -1 # queries with nothing to attend to + indices[3, 5:9] = indices[3, 4] # the same row in several slots of a query + inputs, expected = make_problem(indices.contiguous(), n_rows) + actual = dsa.sparse_attn_backward_flat(*inputs) + assert (actual[0][::7] == 0).all() + for name, a, e in zip(("dq", "dkv", "dsinks"), actual, expected): + assert rel_err(a, e) < REF_TOL, f"{name}: {rel_err(a, e):.2e}" + + +@pytest.mark.parametrize("seq_lens", [[16384], [6000, 3000, 4384, 3000]], ids=["1 doc", "4 docs"]) +def test_v41_shapes(dsa, seq_lens): + """16k queries, 640 slots, 16k window positions + 8k compressed entries; also against cuDNN's SM90 + backward (prime-rl's current path) when it is installed.""" + indices, n_rows = make_indices(seq_lens) + inputs, expected = make_problem(indices, n_rows) + actual = dsa.sparse_attn_backward_flat(*inputs) + for name, a, e in zip(("dq", "dkv", "dsinks"), actual, expected): + assert rel_err(a, e) < REF_TOL, f"{name}: {rel_err(a, e):.2e}" + try: + from cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm90 import flash_attn_bwd_sm90 + except ImportError: + return + q, kv, out, grad_out, lse, idx, sinks, scale = inputs + theirs = flash_attn_bwd_sm90( + q, kv, out, grad_out, lse, attn_sink=sinks, softmax_scale=scale, topk_idxs=idx, need_d_sink=True + ) + for name, a, c in zip(("dq", "dkv", "dsinks"), actual, theirs): + assert rel_err(a, c) < CUDNN_TOL, f"{name} vs cuDNN: {rel_err(a, c):.2e}" + + +def test_custom_op(dsa): + """The torch op takes prime-rl's batched layout and has a fake impl for torch.compile.""" + indices, n_rows = make_indices([100, 28], window=16, topk=48) + (q, kv, out, grad_out, lse, idx, sinks, scale), _ = make_problem(indices, n_rows) + t = q.shape[0] + args = (grad_out[None], q[None], kv[None, :, None], out[None], lse[None], idx[None, :, None], sinks, scale) + dq, dkv, dsinks = dsa.dsa_sparse_attn_backward(*args) + flat = dsa.sparse_attn_backward_flat(q, kv, out, grad_out, lse, idx, sinks, scale) + assert dq.shape == (1, t, HEADS, DIM) and dkv.shape == (1, n_rows, 1, DIM) and dsinks.shape == (HEADS,) + torch.testing.assert_close(dq[0], flat[0], rtol=0, atol=0) + torch.testing.assert_close(dkv[0, :, 0], flat[1], rtol=1e-2, atol=1e-2) # atomics: summation order varies + + from torch._subclasses.fake_tensor import FakeTensorMode + + with FakeTensorMode() as mode: + fake = [mode.from_tensor(a) if isinstance(a, torch.Tensor) else a for a in args] + fdq, fdkv, fds = torch.ops.prime_kernels.dsa_sparse_attn_bwd(*fake) + assert fdq.shape == dq.shape and fdkv.shape == dkv.shape and fds.shape == dsinks.shape + + +def test_out_tail_replaces_the_last_channels(dsa): + """With `out_tail`, each head's last channels of the output come from it: a caller that rewrote + them in place (prime-rl rotates the output's RoPE channels) gets the same gradients bit for bit.""" + indices, n_rows = make_indices([300, 212], window=32, topk=64) + (q, kv, out, grad_out, lse, idx, sinks, scale), _ = make_problem(indices, n_rows) + tail = out[..., -64:].clone() + overwritten = out.clone() + overwritten[..., -64:] = torch.randn_like(tail) + expected = dsa.sparse_attn_backward_flat(q, kv, out, grad_out, lse, idx, sinks, scale) + actual = dsa.sparse_attn_backward_flat(q, kv, overwritten, grad_out, lse, idx, sinks, scale, tail) + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + torch.testing.assert_close(actual[2], expected[2], rtol=0, atol=0) + torch.testing.assert_close(actual[1], expected[1], rtol=1e-2, atol=1e-2) # atomics: summation order varies diff --git a/tests/mhc_projection/bench_mhc_projection.py b/tests/mhc_projection/bench_mhc_projection.py new file mode 100644 index 0000000..59e97ad --- /dev/null +++ b/tests/mhc_projection/bench_mhc_projection.py @@ -0,0 +1,153 @@ +"""Benchmark the fused mHC projection against prime-rl's DeepSeek-V4.1 path. + +Needs prime-rl importable (run with its venv). Times one sublayer's mHC gates plus the collapse +of its input streams, at V4.1 Flash shapes: + + today DeepseekV41HyperConnection.forward + collapse_streams, eager + today+compile the same under torch.compile (prime-rl compiles the decoder blocks) + fused prime_kernels.mhc_projection.hyper_connection, eager and under torch.compile + +Usage: python bench_mhc_projection.py [tokens] +""" + +import sys +from types import SimpleNamespace + +import torch +from triton.testing import do_bench + +import prime_kernels +from prime_rl.trainer.models.deepseek_v41.hyperconnections import DeepseekV41HyperConnection, collapse_streams + +HIDDEN, HC, ITERS, HC_EPS, RMS_EPS = 5120, 4, 20, 1e-6, 1e-20 + + +def today(module, x, pre_mix): + pre, post, comb = module(x) + return pre, post, comb, collapse_streams(x, pre_mix) + + +def fused(mhc, module, x, pre_mix): + return mhc.hyper_connection( + x, + module.fn, + module.scale, + module.base, + pre_mix, + rms_eps=RMS_EPS, + hc_eps=HC_EPS, + sinkhorn_iters=ITERS, + ) + + +def rel_err(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a.double() - b.double()).norm() / b.double().norm()).item() + + +def fp64_reference(module, x, pre_mix): + hc = HC + xd = x.double() + flat = xd.flatten(start_dim=2) + rstd = torch.rsqrt(flat.square().mean(-1, keepdim=True) + RMS_EPS) + mixes = (flat @ module.fn.to(torch.bfloat16).double().t()) * rstd + scale, base = module.scale.double(), module.base.double() + pre_m, post_m, comb_m = mixes.split([hc, hc, hc * hc], dim=-1) + pre_b, post_b, comb_b = base.split([hc, hc, hc * hc]) + pre = torch.sigmoid(pre_m * scale[0] + pre_b) + HC_EPS + post = 2 * torch.sigmoid(post_m * scale[1] + post_b) + comb = (comb_m * scale[2] + comb_b).unflatten(-1, (hc, hc)).softmax(-1) + HC_EPS + comb = comb / (comb.sum(-2, keepdim=True) + HC_EPS) + for _ in range(ITERS - 1): + comb = comb / (comb.sum(-1, keepdim=True) + HC_EPS) + comb = comb / (comb.sum(-2, keepdim=True) + HC_EPS) + collapsed = (pre_mix.double().unsqueeze(-1) * xd).sum(2) + return pre, post, comb, collapsed + + +def time_ms(fn, x, pre_mix, params, cotangents, backward): + """(wall ms from CUDA events, summed GPU kernel ms from the profiler) per step.""" + + def step(): + outputs = fn(x, pre_mix) + if backward: + torch.autograd.grad(outputs, (x, pre_mix, *params), cotangents) + + wall = do_bench(step, warmup=25, rep=200) + steps = 10 + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for _ in range(steps): + step() + torch.cuda.synchronize() + kernel_us = sum(event.device_time_total for event in prof.key_averages() if event.device_type.name == "CUDA") + return wall, kernel_us / steps / 1e3 + + +def main(): + tokens = int(sys.argv[1]) if len(sys.argv) > 1 else 16384 + mhc = prime_kernels.load("mhc_projection") + config = SimpleNamespace( + hc_mult=HC, hc_sinkhorn_iters=ITERS, hc_eps=HC_EPS, rms_norm_eps=RMS_EPS, hidden_size=HIDDEN + ) + module = DeepseekV41HyperConnection(config).cuda() + module.init_weights(0.02) + with torch.no_grad(): + module.base.normal_(0, 0.5) + module.scale.uniform_(0.5, 1.5) + params = tuple(module.parameters()) + + x = torch.randn(1, tokens, HC, HIDDEN, device="cuda", dtype=torch.bfloat16, requires_grad=True) + pre_mix = torch.rand(1, tokens, HC, device="cuda", requires_grad=True) + + paths = { + "today": lambda x, p: today(module, x, p), + "today+compile": torch.compile(lambda x, p: today(module, x, p)), + "fused": lambda x, p: fused(mhc, module, x, p), + "fused+compile": torch.compile(lambda x, p: fused(mhc, module, x, p)), + } + + # Accuracy against an fp64 evaluation of the same math (weight rounded to bf16 as both paths do). + inputs = (x, pre_mix, *params) + exact = fp64_reference(module, x, pre_mix) + # Cotangents in the dtypes the paths produce (fp32 gates, bf16 collapse), shared exactly. + cotangents = [torch.randn_like(t, dtype=dtype) for t, dtype in zip(exact, [torch.float32] * 3 + [x.dtype])] + exact_all = (*exact, *torch.autograd.grad(exact, inputs, [c.double() for c in cotangents])) + names = ("pre", "post", "comb", "collapsed", "d_x", "d_pre_mix", "d_fn", "d_base", "d_scale") + print("relative L2 error vs fp64:") + print(f" {'':10s} {'today':>9s} {'fused':>9s}") + errors = {} + for name in ("today", "fused"): + outputs = paths[name](x, pre_mix) + grads = torch.autograd.grad(outputs, inputs, cotangents) + errors[name] = [rel_err(a, b) for a, b in zip((*outputs, *grads), exact_all)] + for i, label in enumerate(names): + print(f" {label:10s} {errors['today'][i]:9.2e} {errors['fused'][i]:9.2e}") + + streams_gb = x.numel() * x.element_size() / 1e9 + print(f"\ntokens={tokens}, streams {streams_gb:.3f} GB (bf16)") + print("wall: CUDA events around one eager step (includes launch gaps); gpu: summed kernel time") + print(f"{'path':16s} {'fwd wall':>9s} {'fwd gpu':>8s} {'fwd+bwd wall':>13s} {'fwd+bwd gpu':>12s} (ms)") + for name, fn in paths.items(): + fwd = time_ms(fn, x, pre_mix, params, cotangents, backward=False) + both = time_ms(fn, x, pre_mix, params, cotangents, backward=True) + print(f"{name:16s} {fwd[0]:9.3f} {fwd[1]:8.3f} {both[0]:13.3f} {both[1]:12.3f}") + + # The two fused kernels alone, for their achieved bandwidth. + flat = x.detach().view(-1, HC, HIDDEN) + flat_pre = pre_mix.detach().view(-1, HC) + proj = mhc.projection + fwd = do_bench(lambda: proj.mhc_projection_forward(flat, module.fn, flat_pre, RMS_EPS), warmup=25, rep=200) + mixes, rstd, collapsed = proj.mhc_projection_forward(flat, module.fn, flat_pre, RMS_EPS) + grad_mixes, grad_collapsed = torch.randn_like(mixes), torch.randn_like(collapsed) + bwd = do_bench( + lambda: proj.mhc_projection_backward(flat, module.fn, flat_pre, mixes, rstd, grad_mixes, grad_collapsed), + warmup=25, + rep=200, + ) + fwd_bytes = streams_gb * 1.25 # read the streams, write the collapse + bwd_bytes = streams_gb * 2.25 # read the streams and d(collapse), write d(streams) + print(f"\nprojection fwd kernel {fwd:.3f} ms ({fwd_bytes / fwd:.2f} TB/s)") + print(f"projection bwd kernel {bwd:.3f} ms ({bwd_bytes / bwd:.2f} TB/s)") + + +if __name__ == "__main__": + main() diff --git a/tests/mhc_projection/test_mhc_projection.py b/tests/mhc_projection/test_mhc_projection.py new file mode 100644 index 0000000..e15f09c --- /dev/null +++ b/tests/mhc_projection/test_mhc_projection.py @@ -0,0 +1,173 @@ +import pytest +import torch + +import prime_kernels + +REASON = prime_kernels.unavailable_reason("mhc_projection") +pytestmark = pytest.mark.skipif(REASON is not None, reason=REASON or "") + +RMS_EPS = 1e-20 +HC_EPS = 1e-6 +ITERS = 20 + +# (tokens, hc, d): V4.1 Flash, then small shapes including ragged token and d tails. +SHAPES = [(4096, 4, 5120), (1, 4, 64), (37, 4, 320), (130, 4, 200), (77, 2, 192)] + + +@pytest.fixture(autouse=True) +def seed(): + torch.manual_seed(0) + + +@pytest.fixture(scope="module") +def mhc(): + return prime_kernels.load("mhc_projection") + + +def rel_err(actual: torch.Tensor, expected: torch.Tensor) -> float: + return ((actual.double() - expected.double()).norm() / expected.double().norm().clamp_min(1e-30)).item() + + +def reference_projection(x, weight, pre_mix, eps): + """Plain torch in fp64, so it is the exact value the fp32 kernel approximates. The weight is + rounded to bf16 first, which is the kernel's (and prime-rl's) contract.""" + xf = x.double() + flat = xf.flatten(-2) + rstd = torch.rsqrt(flat.square().mean(-1, keepdim=True) + eps) + mixes = (flat @ weight.to(torch.bfloat16).double().t()) * rstd + collapsed = (pre_mix.double().unsqueeze(-1) * xf).sum(-2) if pre_mix is not None else None + return mixes.float(), collapsed.float() if collapsed is not None else None + + +def reference_gates(mixes, scale, base, hc, iters, eps): + pre_m, post_m, comb_m = mixes.split([hc, hc, hc * hc], dim=-1) + pre_b, post_b, comb_b = base.split([hc, hc, hc * hc]) + pre = torch.sigmoid(pre_m * scale[0] + pre_b) + eps + post = 2 * torch.sigmoid(post_m * scale[1] + post_b) + comb = (comb_m * scale[2] + comb_b).unflatten(-1, (hc, hc)) + comb = comb.softmax(-1) + eps + comb = comb / (comb.sum(-2, keepdim=True) + eps) + for _ in range(iters - 1): + comb = comb / (comb.sum(-1, keepdim=True) + eps) + comb = comb / (comb.sum(-2, keepdim=True) + eps) + return pre, post, comb + + +def make_inputs(tokens, hc, d, with_pre=True): + n = (2 + hc) * hc + x = torch.randn(tokens, hc, d, device="cuda").to(torch.bfloat16).requires_grad_() + weight = (torch.randn(n, hc * d, device="cuda") * (hc * d) ** -0.5).requires_grad_() + pre_mix = (torch.rand(tokens, hc, device="cuda") + 0.1).requires_grad_() if with_pre else None + scale = (torch.rand(3, device="cuda") + 0.5).requires_grad_() + base = (torch.randn(n, device="cuda") * 0.5).requires_grad_() + return x, weight, pre_mix, scale, base + + +def grads(outputs, inputs, cotangents): + tensors = [t for t in inputs if t is not None] + found = torch.autograd.grad(outputs, tensors, cotangents) + return iter(found) + + +@pytest.mark.parametrize("with_pre", [True, False]) +@pytest.mark.parametrize("shape", SHAPES) +def test_projection_matches_reference(mhc, shape, with_pre): + x, weight, pre_mix, _, _ = make_inputs(*shape, with_pre=with_pre) + mixes, collapsed = mhc.mhc_projection(x, weight, pre_mix, RMS_EPS) + ref_mixes, ref_collapsed = reference_projection(x, weight, pre_mix, RMS_EPS) + + # bf16 x bf16 products are exact in fp32; only the accumulation order differs. + assert rel_err(mixes, ref_mixes) < 1e-5 + outputs, cotangents = [mixes], [torch.randn_like(mixes)] + ref_outputs = [ref_mixes] + if with_pre: + # Only the final bf16 rounding separates the two. + assert rel_err(collapsed.float(), ref_collapsed) < 3e-3 + outputs.append(collapsed) + cotangents.append(torch.randn_like(collapsed)) + ref_outputs.append(ref_collapsed) + + inputs = (x, weight, pre_mix) + got = grads(outputs, inputs, cotangents) + ref = grads(ref_outputs, inputs, cotangents) + # The kernel rounds d(projection) to bf16 for the tensor-core GEMMs (as prime-rl's eager + # path does); dx itself is stored in bf16. + assert rel_err(next(got), next(ref)) < 5e-3 # x + assert rel_err(next(got), next(ref)) < 5e-3 # weight + if with_pre: + assert rel_err(next(got), next(ref)) < 1e-5 # pre_mix + + +@pytest.mark.parametrize("tokens", [1, 333, 16384]) +def test_gates_match_reference(mhc, tokens): + hc = 4 + mixes = (torch.randn(tokens, (2 + hc) * hc, device="cuda") * 2).requires_grad_() + scale = (torch.rand(3, device="cuda") + 0.5).requires_grad_() + base = (torch.randn((2 + hc) * hc, device="cuda") * 0.5).requires_grad_() + got = mhc.mhc_gates(mixes, scale, base, hc, ITERS, HC_EPS) + ref = reference_gates(mixes, scale, base, hc, ITERS, HC_EPS) + for actual, expected in zip(got, ref): + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + + cotangents = [torch.randn_like(t) for t in ref] + got_grads = torch.autograd.grad(got, (mixes, scale, base), cotangents) + ref_grads = torch.autograd.grad(ref, (mixes, scale, base), cotangents) + for actual, expected in zip(got_grads, ref_grads): + assert rel_err(actual, expected) < 1e-4 + + +def test_hyper_connection_matches_reference_under_compile(mhc): + tokens, hc, d = 256, 4, 512 + x, weight, pre_mix, scale, base = make_inputs(tokens, hc, d) + x3 = x.detach().view(2, tokens // 2, hc, d).requires_grad_() + pre3 = pre_mix.detach().view(2, tokens // 2, hc).requires_grad_() + kwargs = dict(rms_eps=RMS_EPS, hc_eps=HC_EPS, sinkhorn_iters=ITERS) + + compiled = torch.compile(mhc.hyper_connection, fullgraph=True) + got = compiled(x3, weight, scale, base, pre3, **kwargs) + mixes, collapsed = reference_projection(x3, weight, pre3, RMS_EPS) + ref = (*reference_gates(mixes, scale, base, hc, ITERS, HC_EPS), collapsed) + assert [t.shape for t in got] == [t.shape for t in ref] + for actual, expected in zip(got, ref): + assert rel_err(actual.float(), expected) < 3e-3 # collapsed is bf16 + + inputs = (x3, weight, scale, base, pre3) + cotangents = [torch.randn_like(t) for t in ref] + got_grads = torch.autograd.grad(got, inputs, cotangents) + ref_grads = torch.autograd.grad(ref, inputs, cotangents) + for actual, expected in zip(got_grads, ref_grads): + assert rel_err(actual, expected) < 5e-3 + + +def test_backward_is_deterministic(mhc): + x, weight, pre_mix, _, _ = make_inputs(4096, 4, 5120) + + def run(): + mixes, collapsed = mhc.mhc_projection(x, weight, pre_mix, RMS_EPS) + torch.manual_seed(1) + cotangents = [torch.randn_like(mixes), torch.randn_like(collapsed)] + return torch.autograd.grad([mixes, collapsed], (x, weight, pre_mix), cotangents) + + for first, second in zip(run(), run()): + assert torch.equal(first, second) + + +def test_streams_view_folds_its_gradient_into_the_backward(mhc): + x, weight, pre_mix, _, _ = make_inputs(4096, 4, 5120) + other = torch.randn_like(x) # stands in for the streams' other reader, the post/comb update + + def run(streams_api): + if streams_api: + mixes, collapsed, streams = mhc.mhc_projection_streams(x, weight, pre_mix, RMS_EPS) + else: + (mixes, collapsed), streams = mhc.mhc_projection(x, weight, pre_mix, RMS_EPS), x + torch.manual_seed(1) + cotangents = [torch.randn_like(mixes), torch.randn_like(collapsed), torch.randn_like(x)] + outputs = [mixes, collapsed, streams * other] + return torch.autograd.grad(outputs, (x, weight, pre_mix), cotangents) + + (x_ref, weight_ref, pre_ref), (x_got, weight_got, pre_got) = run(False), run(True) + # The stream gradient is exact; the others are fp32 partial sums whose order follows the autotuned config. + assert torch.equal(x_got, x_ref) + assert rel_err(weight_got, weight_ref) < 1e-6 + assert rel_err(pre_got, pre_ref) < 1e-6 diff --git a/tests/moe_experts/bench_moe_experts.py b/tests/moe_experts/bench_moe_experts.py new file mode 100644 index 0000000..d887d95 --- /dev/null +++ b/tests/moe_experts/bench_moe_experts.py @@ -0,0 +1,313 @@ +"""Benchmark `moe_experts` against prime-rl's current expert compute at V4.1 Flash shapes. + +prime-rl today (`GroupedGemmExpertCompute` with `BF16ExpertCompute` and V4's `ClampedSwiglu`): +gate and up as two `torch._grouped_mm` calls (or one over the packed `gate_up_proj`), the clamped +SwiGLU in eager PyTorch (or under torch.compile when the decoder block is compiled), and the down +projection as a third `torch._grouped_mm`; autograd does the backward. Its FP8 path +(`DeepGemmFP8ExpertCompute`) is the same with prime-rl's `grouped_fp8_gemm` (DeepGEMM with +standalone quantization kernels) for each GEMM; it is benchmarked when prime-rl is importable +(put its `src` on PYTHONPATH). + +Per GPU, 16k tokens routed top-6 give 98304 expert rows: 4 local experts of ~24.5k rows each at +EP=96 and 48 local experts of ~2k rows each at EP=8. Group sizes get +-20% imbalance and are +padded to the kernel's TOKEN_GROUP_ALIGNMENT. + + python tests/moe_experts/bench_moe_experts.py [--breakdown] [--accuracy] +""" + +import argparse +import statistics + +import torch +import torch.nn.functional as F + +import prime_kernels + +HIDDEN = 5120 +INTERMEDIATE = 2304 +LIMIT = 10.0 + + +def clamped_swiglu(gate, up): + return F.silu(gate.clamp(max=LIMIT)) * up.clamp(min=-LIMIT, max=LIMIT) + + +compiled_swiglu = torch.compile(clamped_swiglu) + + +def grouped_mm(x, weight, offsets): + return torch._grouped_mm(x, weight, offs=offsets) + + +def today(x, gate_proj, up_proj, down_proj, offsets, activation, gemm=grouped_mm): + if up_proj is None: + gate, up = gemm(x, gate_proj.transpose(-2, -1), offsets).chunk(2, dim=-1) + else: + gate = gemm(x, gate_proj.transpose(-2, -1), offsets) + up = gemm(x, up_proj.transpose(-2, -1), offsets) + return gemm(activation(gate, up), down_proj.transpose(-2, -1), offsets) + + +def prime_rl_fp8_gemm(): + """prime-rl's `grouped_fp8_gemm`, or None when prime-rl is not importable.""" + try: + from prime_rl.trainer.models.layers.fp8_grouped_gemm import grouped_fp8_gemm + except ImportError: + return None + return grouped_fp8_gemm + + +def time_ms(fn, iters): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def bench(candidates, iters, repeats): + """Median time per candidate, the candidates interleaved so clock drift hits all alike.""" + for fn in candidates.values(): + fn() + fn() + torch.cuda.synchronize() + samples = {name: [] for name in candidates} + for _ in range(repeats): + for name, fn in candidates.items(): + samples[name].append(time_ms(fn, iters)) + return {name: statistics.median(times) for name, times in samples.items()} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--iters", type=int, default=5) + parser.add_argument("--repeats", type=int, default=7) + parser.add_argument("--per-op", action="store_true", help="also time each fused GEMM against its unfused ops") + parser.add_argument("--breakdown", action="store_true", help="also list the kernels of one forward + backward") + parser.add_argument("--accuracy", action="store_true", help="also report errors against an fp32 reference") + parser.add_argument("--no-bf16-today", action="store_true", help="skip prime-rl's bf16 variants") + args = parser.parse_args() + fp8_gemm = prime_rl_fp8_gemm() + if fp8_gemm is None: + print("prime-rl is not importable: skipping its FP8 path") + + moe = prime_kernels.load("moe_experts") + torch.manual_seed(0) + rows = 16384 * 6 + print( + f"{'load':8s} {'variant':26s} {'fwd ms':>8s} {'TFLOP/s':>8s} {'fwd+bwd ms':>11s} {'TFLOP/s':>8s} " + f"{'saved GiB':>10s}" + ) + for label, num_experts in (("EP=96", 4), ("EP=8", 48)): + mean = rows / num_experts + counts = torch.randint(int(mean * 0.8), int(mean * 1.2) + 1, (num_experts,)) + align = moe.TOKEN_GROUP_ALIGNMENT + counts = ((counts * rows / counts.sum()).long() + align - 1) // align * align + counts = counts.cuda() + num_rows = int(counts.sum()) + offsets = torch.cumsum(counts, 0, dtype=torch.int32) + flops = 2 * num_rows * HIDDEN * INTERMEDIATE * 3 + + x = torch.randn(num_rows, HIDDEN, device="cuda", dtype=torch.bfloat16, requires_grad=True) + scale = 7.0 / HIDDEN**0.5 + gate_up = (torch.randn(num_experts, 2 * INTERMEDIATE, HIDDEN, device="cuda") * scale).bfloat16() + gate_up.requires_grad_() + gate_proj = gate_up[:, :INTERMEDIATE].detach().clone().requires_grad_() + up_proj = gate_up[:, INTERMEDIATE:].detach().clone().requires_grad_() + down_proj = (torch.randn(num_experts, HIDDEN, INTERMEDIATE, device="cuda") / INTERMEDIATE**0.5).bfloat16() + down_proj.requires_grad_() + dout = torch.randn(num_rows, HIDDEN, device="cuda", dtype=torch.bfloat16) + + separate = [x, gate_proj, up_proj, down_proj] + packed = [x, gate_up, down_proj] + # name -> (forward, the tensors it is differentiated against) + variants = {} + if not args.no_bf16_today: + variants["today (eager act)"] = ( + lambda: today(x, gate_proj, up_proj, down_proj, offsets, clamped_swiglu), + separate, + ) + variants["today (compiled act)"] = ( + lambda: today(x, gate_proj, up_proj, down_proj, offsets, compiled_swiglu), + separate, + ) + variants["today packed gate_up"] = ( + lambda: today(x, gate_up, None, down_proj, offsets, compiled_swiglu), + packed, + ) + variants["moe_experts"] = (lambda: moe.moe_experts(x, gate_proj, up_proj, down_proj, counts, LIMIT), separate) + variants["moe_experts packed"] = (lambda: moe.moe_experts(x, gate_up, None, down_proj, counts, LIMIT), packed) + if fp8_gemm is not None: + variants["today fp8 (DeepGEMM)"] = ( + lambda: today(x, gate_proj, up_proj, down_proj, offsets, compiled_swiglu, fp8_gemm), + separate, + ) + variants["today fp8 packed"] = ( + lambda: today(x, gate_up, None, down_proj, offsets, compiled_swiglu, fp8_gemm), + packed, + ) + variants["moe_experts fp8"] = ( + lambda: moe.moe_experts(x, gate_proj, up_proj, down_proj, counts, LIMIT, fp8=True), + separate, + ) + variants["moe_experts fp8 packed"] = ( + lambda: moe.moe_experts(x, gate_up, None, down_proj, counts, LIMIT, fp8=True), + packed, + ) + + def forward(fn): + def run(): + with torch.no_grad(): + fn() + + return run + + def forward_backward(fn, inputs): + def run(): + # autograd.grad rather than backward(): no accumulation into .grad in the timing. + torch.autograd.grad(fn(), inputs, dout) + + return run + + saved = {} + for name, (fn, _) in variants.items(): + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + out = fn() + saved[name] = (torch.cuda.memory_allocated() - before - out.numel() * out.element_size()) / 2**30 + del out + + fwd = bench({name: forward(fn) for name, (fn, _) in variants.items()}, args.iters, args.repeats) + fwd_bwd = bench( + {name: forward_backward(fn, inputs) for name, (fn, inputs) in variants.items()}, args.iters, args.repeats + ) + for name in variants: + print( + f"{label:8s} {name:26s} {fwd[name]:8.3f} {flops / fwd[name] / 1e9:8.0f} " + f"{fwd_bwd[name]:11.3f} {3 * flops / fwd_bwd[name] / 1e9:8.0f} {saved[name]:10.2f}" + ) + if args.breakdown: + for name in ("moe_experts", "today fp8 (DeepGEMM)", "moe_experts fp8"): + if name in variants: + breakdown(label, name, forward_backward(*variants[name])) + if args.accuracy: + accuracy( + label, {name: variants[name] for name in variants if "fp8" in name}, inputs=separate, counts=counts + ) + if args.per_op: + per_op(moe, x, gate_proj, up_proj, down_proj, counts, offsets, dout, args) + + +def breakdown(label, name, fn): + """GPU time of each kernel of one forward + backward, the longest first.""" + from torch.profiler import ProfilerActivity, profile + + fn() + torch.cuda.synchronize() + iters = 3 + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + fn() + torch.cuda.synchronize() + events = [e for e in prof.key_averages() if e.device_type == torch.autograd.DeviceType.CUDA] + events.sort(key=lambda e: e.self_device_time_total, reverse=True) + total = sum(e.self_device_time_total for e in events) / iters / 1e3 + print(f"\n{label} {name}: kernels of one forward + backward ({total:.2f} ms of GPU time)") + for e in events: + ms = e.self_device_time_total / iters / 1e3 + if ms < 0.01: + continue + print(f" {ms:8.3f} ms x{e.count // iters:<3d} {e.key[:90]}") + + +def accuracy(label, variants, inputs, counts): + """Relative errors of each variant's output and gradients against an fp32 reference.""" + x, gate_proj, up_proj, down_proj = inputs + dout = torch.randn_like(x) + num_rows = int(counts.sum()) + + def reference(): + out = torch.zeros(x.shape, device=x.device, dtype=torch.float32) + start = 0 + for expert, count in enumerate(counts.tolist()): + rows = slice(start, start + count) + xe = x[rows].float() + h = clamped_swiglu(xe @ gate_proj[expert].float().T, xe @ up_proj[expert].float().T) + out[rows] = h @ down_proj[expert].float().T + start += count + return out + + def errors(out): + grads = torch.autograd.grad(out, inputs, dout.to(out.dtype)) + # Only the routed rows: prime-rl's FP8 path leaves the rest of dx uninitialized. + tensors = [out[:num_rows], grads[0][:num_rows], *grads[1:]] + return [t.float() for t in tensors] + + expected = errors(reference()) + for name, (fn, variant_inputs) in variants.items(): + if variant_inputs is not inputs: + continue + actual = errors(fn()) + rel = [((a - e).norm() / e.norm()).item() for a, e in zip(actual, expected)] + names = ("out", "dx", "dgate", "dup", "ddown") + print(f"{label:8s} {name:26s} rel err vs fp32: " + " ".join(f"{n} {r:.3e}" for n, r in zip(names, rel))) + + +def per_op(moe, x, gate_proj, up_proj, down_proj, counts, offsets, dout, args): + """Each fused GEMM of `moe_experts` next to the unfused ops it replaces (the weight gradients + are torch._grouped_mm in both).""" + from prime_kernels.moe_experts import gemm + from prime_kernels.moe_experts.kernels import TileTable + + x, gate_proj, up_proj, down_proj = (t.detach() for t in (x, gate_proj, up_proj, down_proj)) + num_sms = torch.cuda.get_device_properties(x.device).multi_processor_count + table = TileTable(counts, x.shape[0], gemm.BLOCK_M) + gate, up, h = gemm.fc1(x, gate_proj, up_proj, table, LIMIT, num_sms=num_sms) + dgate_dup, _ = gemm.dswiglu(dout, down_proj, gate, up, table, LIMIT, num_sms=num_sms) + dgate, dup = dgate_dup[:, :INTERMEDIATE], dgate_dup[:, INTERMEDIATE:] + flops = 2 * x.shape[0] * HIDDEN * INTERMEDIATE + + def mm(a, b): + return torch._grouped_mm(a, b, offs=offsets) + + def activation_backward(): + with torch.enable_grad(): + gate_ = gate.detach().requires_grad_() + up_ = up.detach().requires_grad_() + return torch.autograd.grad(compiled_swiglu(gate_, up_), (gate_, up_), h) + + ops = { + "fc1 + SwiGLU": ( + 2, + lambda: gemm.fc1(x, gate_proj, up_proj, table, LIMIT, num_sms=num_sms), + lambda: compiled_swiglu(mm(x, gate_proj.transpose(-2, -1)), mm(x, up_proj.transpose(-2, -1))), + ), + "down": ( + 1, + lambda: gemm.down(h, down_proj, table, num_sms=num_sms), + lambda: mm(h, down_proj.transpose(-2, -1)), + ), + "dh + SwiGLU backward": ( + 1, + lambda: gemm.dswiglu(dout, down_proj, gate, up, table, LIMIT, num_sms=num_sms), + lambda: (mm(dout, down_proj), activation_backward()), + ), + "dx": ( + 2, + lambda: gemm.dx(dgate, dup, gate_proj, up_proj, table, num_sms=num_sms), + lambda: mm(dgate, gate_proj) + mm(dup, up_proj), + ), + } + with torch.no_grad(): + for name, (gemms, fused, unfused) in ops.items(): + times = bench({"fused": fused, "unfused": unfused}, args.iters, args.repeats) + print( + f" {name:26s} moe_experts {times['fused']:7.3f} ms {gemms * flops / times['fused'] / 1e9:4.0f} " + f"TFLOP/s | today {times['unfused']:7.3f} ms {gemms * flops / times['unfused'] / 1e9:4.0f} TFLOP/s" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/moe_experts/test_moe_experts.py b/tests/moe_experts/test_moe_experts.py new file mode 100644 index 0000000..6d81822 --- /dev/null +++ b/tests/moe_experts/test_moe_experts.py @@ -0,0 +1,342 @@ +import pytest +import torch +import torch.nn.functional as F + +import prime_kernels + +LIMIT = 10.0 + + +@pytest.fixture(scope="module") +def moe_experts(): + reason = prime_kernels.unavailable_reason("moe_experts") + if reason is not None: + pytest.skip(reason) + return prime_kernels.load("moe_experts") + + +def make_problem(num_experts, tokens_per_expert, hidden, intermediate, *, tail_rows, packed, seed, alignment): + """Grouped tokens and weights scaled so that roughly 10% of the gates exceed the clamp.""" + generator = torch.Generator(device="cuda").manual_seed(seed) + counts = torch.randint( + int(tokens_per_expert * 0.75), + int(tokens_per_expert * 1.25) + 1, + (num_experts,), + device="cuda", + generator=generator, + ) + counts = (counts + alignment - 1) // alignment * alignment + if num_experts > 2: + counts[1] = 0 + rows = int(counts.sum()) + tail_rows + x = torch.randn(rows, hidden, device="cuda", dtype=torch.bfloat16, generator=generator) + x[rows - tail_rows :] = 0 + scale = 7.0 / hidden**0.5 + if packed: + gate_up = torch.randn(num_experts, 2 * intermediate, hidden, device="cuda", generator=generator) * scale + gate_proj = gate_up.bfloat16().requires_grad_() + up_proj = None + else: + gate_proj = torch.randn(num_experts, intermediate, hidden, device="cuda", generator=generator) * scale + up_proj = torch.randn(num_experts, intermediate, hidden, device="cuda", generator=generator) * scale + gate_proj = gate_proj.bfloat16().requires_grad_() + up_proj = up_proj.bfloat16().requires_grad_() + down_proj = torch.randn(num_experts, hidden, intermediate, device="cuda", generator=generator) / intermediate**0.5 + down_proj = down_proj.bfloat16().requires_grad_() + dout = torch.randn(rows, hidden, device="cuda", dtype=torch.bfloat16, generator=generator) + return x.requires_grad_(), gate_proj, up_proj, down_proj, counts, dout + + +def reference(x, gate_proj, up_proj, down_proj, counts, dtype): + """Per-expert loop in ``dtype``; rows after the last group give zeros.""" + if up_proj is None: + gate_proj, up_proj = gate_proj.chunk(2, dim=1) + out = torch.zeros(x.shape, device=x.device, dtype=dtype) + start = 0 + for expert, count in enumerate(counts.tolist()): + rows = slice(start, start + count) + xe = x[rows].to(dtype) + gate = xe @ gate_proj[expert].to(dtype).T + up = xe @ up_proj[expert].to(dtype).T + h = F.silu(gate.clamp(max=LIMIT)) * up.clamp(min=-LIMIT, max=LIMIT) + out[rows] = h @ down_proj[expert].to(dtype).T + start += count + return out + + +def grads(out, dout, tensors): + return torch.autograd.grad(out, tensors, dout.to(out.dtype)) + + +def rel_err(actual, expected): + return ((actual.float() - expected.float()).norm() / expected.float().norm().clamp_min(1e-30)).item() + + +CASES = { + "small": dict(num_experts=3, tokens_per_expert=300, hidden=256, intermediate=192, tail_rows=40), + "odd": dict(num_experts=5, tokens_per_expert=130, hidden=136, intermediate=264, tail_rows=0), + "v41_ep96": dict(num_experts=4, tokens_per_expert=2048, hidden=5120, intermediate=2304, tail_rows=64), + "v41_ep8": dict(num_experts=48, tokens_per_expert=256, hidden=5120, intermediate=2304, tail_rows=0), +} + + +@pytest.mark.parametrize("packed", [False, True], ids=["separate", "packed"]) +@pytest.mark.parametrize("case", list(CASES)) +def test_matches_fp32_reference(moe_experts, case, packed): + x, gate_proj, up_proj, down_proj, counts, dout = make_problem( + **CASES[case], packed=packed, seed=0, alignment=moe_experts.TOKEN_GROUP_ALIGNMENT + ) + weights = [w for w in (gate_proj, up_proj, down_proj) if w is not None] + tensors = [x, *weights] + + out = moe_experts.moe_experts(x, gate_proj, up_proj, down_proj, counts, LIMIT) + actual = [out, *grads(out, dout, tensors)] + expected_out = reference(x, gate_proj, up_proj, down_proj, counts, torch.float32) + expected = [expected_out, *grads(expected_out, dout, tensors)] + # The unfused bf16 computation (today's path), which rounds gate, up, h and dh to bf16 too. + bf16_out = reference(x, gate_proj, up_proj, down_proj, counts, torch.bfloat16) + bf16 = [bf16_out, *grads(bf16_out, dout, tensors)] + + names = ["out", "dx"] + [ + f"d{name}" for name, w in zip(("gate", "up", "down"), (gate_proj, up_proj, down_proj)) if w is not None + ] + for name, a, e, b in zip(names, actual, expected, bf16): + assert a.shape == e.shape and a.dtype == torch.bfloat16, name + assert torch.isfinite(a).all(), name + err, bf16_err = rel_err(a, e), rel_err(b, e) + # Gradients through the gate and up projections pass the clamps, whose masks flip for the + # values bf16 rounding moves across the limit: with ~10% of values clamped that alone is a + # ~4% relative difference from fp32, for the unfused bf16 computation as well. + tolerance = 1e-2 if name in ("out", "ddown") else 6e-2 + assert err < tolerance, f"{name}: rel err {err:.2e} vs fp32" + assert err < 1.05 * bf16_err + 1e-3, f"{name}: rel err {err:.2e}, unfused bf16 {bf16_err:.2e}" + tail = int(counts.sum()) + assert (out[tail:] == 0).all() and (actual[1][tail:] == 0).all() + + +def test_clamp_is_exercised(moe_experts): + x, gate_proj, up_proj, _, counts, _ = make_problem( + **CASES["v41_ep96"], packed=False, seed=0, alignment=moe_experts.TOKEN_GROUP_ALIGNMENT + ) + gate = x[: counts[0]].float() @ gate_proj[0].float().T + up = x[: counts[0]].float() @ up_proj[0].float().T + assert 0.02 < (gate > LIMIT).float().mean() < 0.5 + assert 0.02 < (up.abs() > LIMIT).float().mean() < 0.5 + + +def test_torch_compile(moe_experts): + x, gate_proj, up_proj, down_proj, counts, dout = make_problem( + **CASES["small"], packed=False, seed=1, alignment=moe_experts.TOKEN_GROUP_ALIGNMENT + ) + + def block(x, gate_proj, up_proj, down_proj, counts): + return moe_experts.moe_experts(x * 2, gate_proj, up_proj, down_proj, counts, LIMIT) + 1 + + compiled = torch.compile(block, fullgraph=True) + tensors = [x, gate_proj, up_proj, down_proj] + out = compiled(x, gate_proj, up_proj, down_proj, counts) + eager = block(x, gate_proj, up_proj, down_proj, counts) + torch.testing.assert_close(out, eager, rtol=0, atol=0) + for a, e in zip(grads(out, dout, tensors), grads(eager, dout, tensors)): + torch.testing.assert_close(a, e, rtol=0, atol=0) + + +def test_opcheck(moe_experts): + x, gate_proj, up_proj, down_proj, counts, _ = make_problem( + **CASES["small"], packed=False, seed=2, alignment=moe_experts.TOKEN_GROUP_ALIGNMENT + ) + torch.library.opcheck( + torch.ops.prime_kernels.moe_experts_forward.default, + (x, gate_proj, up_proj, down_proj, counts, LIMIT), + test_utils=("test_schema", "test_faketensor", "test_autograd_registration"), + ) + + +def test_unsupported_shape_reason(moe_experts): + assert moe_experts.unsupported_shape_reason(5120, 2304) is None + assert moe_experts.unsupported_shape_reason(5120, 2300) is not None + + +@pytest.mark.parametrize("rows", [0, 64]) +def test_no_routed_tokens(moe_experts, rows): + num_experts, hidden, intermediate = 3, 256, 192 + x = torch.randn(rows, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + gate_proj, up_proj = ( + torch.randn(num_experts, intermediate, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + for _ in range(2) + ) + down_proj = torch.randn(num_experts, hidden, intermediate, device="cuda", dtype=torch.bfloat16, requires_grad=True) + counts = torch.zeros(num_experts, dtype=torch.int64, device="cuda") + out = moe_experts.moe_experts(x, gate_proj, up_proj, down_proj, counts, LIMIT) + assert out.shape == x.shape and (out == 0).all() + dx, *dws = grads(out, torch.randn_like(out), [x, gate_proj, up_proj, down_proj]) + assert (dx == 0).all() and all((dw == 0).all() for dw in dws) + + +# FP8 (moe_experts(..., fp8=True)): blockwise e4m3, DeepSeek-V3 recipe through DeepGEMM. + + +@pytest.fixture(scope="module") +def moe_experts_fp8(moe_experts): + pytest.importorskip("deep_gemm") + return moe_experts + + +FP8_CASES = { + "small": dict(num_experts=3, tokens_per_expert=300, hidden=256, intermediate=384, tail_rows=40), + "odd": dict(num_experts=5, tokens_per_expert=90, hidden=384, intermediate=256, tail_rows=24), + "v41_ep96": CASES["v41_ep96"], + "v41_ep8": CASES["v41_ep8"], +} + + +def _fp8(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float8_e4m3fn).float() + + +def quant_groups(x: torch.Tensor) -> torch.Tensor: + """Quantize-dequantize per 1 x 128 group of the last dimension (zero-padded groups).""" + rows, cols = x.shape + groups = F.pad(x.float(), (0, -cols % 128)).view(rows, -1, 128) + scale = groups.abs().amax(-1, keepdim=True).clamp_min(1e-10) / 448 + return (_fp8(groups / scale) * scale).view(rows, -1)[:, :cols] + + +def quant_blocks(w: torch.Tensor) -> torch.Tensor: + """Quantize-dequantize per 128 x 128 block of the last two dimensions.""" + E, R, C = w.shape + blocks = w.float().view(E, R // 128, 128, C // 128, 128) + scale = blocks.abs().amax((2, 4), keepdim=True).clamp_min(1e-4) / 448 + return (_fp8(blocks / scale) * scale).view(E, R, C) + + +def fp8_emulation(x, gate_proj, up_proj, down_proj, counts, dout): + """out and the gradients of x and the weights, computed the way the FP8 path does: FP8 + operands (per-token groups along K for activations and gradients, per 128-token groups of + each channel for the weight-gradient operands, 128 x 128 blocks for weights), exact products + with fp32 accumulation, GEMM outputs rounded to bf16, SwiGLU math in fp32.""" + if up_proj is None: + gate_proj, up_proj = gate_proj.chunk(2, dim=1) + w13 = quant_blocks(torch.cat([gate_proj, up_proj], dim=1).detach()) + w2 = quant_blocks(down_proj.detach()) + intermediate = gate_proj.shape[1] + out, dx = torch.zeros_like(x, dtype=torch.float32), torch.zeros_like(x, dtype=torch.float32) + dw13, dw2 = torch.zeros_like(w13), torch.zeros_like(w2) + start = 0 + for e, count in enumerate(counts.tolist()): + rows = slice(start, start + count) + start += count + if count == 0: + continue + xe, de = x[rows].detach(), dout[rows] + gate_up = (quant_groups(xe) @ w13[e].T).bfloat16().float() + gate, up = gate_up[:, :intermediate], gate_up[:, intermediate:] + gate_c, up_c = gate.clamp(max=LIMIT), up.clamp(-LIMIT, LIMIT) + sig = torch.sigmoid(gate_c) + h = gate_c * sig * up_c + out[rows] = (quant_groups(h) @ w2[e].T).bfloat16().float() + dh = (quant_groups(de) @ w2[e]).bfloat16().float() + dgate = torch.where(gate <= LIMIT, dh * up_c * sig * (1 + gate_c * (1 - sig)), 0.0) + dup = torch.where(up.abs() <= LIMIT, dh * gate_c * sig, 0.0) + dgu = torch.cat([dgate, dup], dim=1) + dx[rows] = (quant_groups(dgu) @ w13[e]).bfloat16().float() + dw2[e] = quant_groups(de.T.float()) @ quant_groups(h.T).T + dw13[e] = quant_groups(dgu.T) @ quant_groups(xe.T.float()).T + grads = [dx, dw13[:, :intermediate], dw13[:, intermediate:], dw2] + return out, grads + + +@pytest.mark.parametrize("packed", [False, True], ids=["separate", "packed"]) +@pytest.mark.parametrize("case", list(FP8_CASES)) +def test_fp8_matches_reference(moe_experts_fp8, case, packed): + x, gate_proj, up_proj, down_proj, counts, dout = make_problem( + **FP8_CASES[case], packed=packed, seed=0, alignment=moe_experts_fp8.TOKEN_GROUP_ALIGNMENT + ) + weights = [w for w in (gate_proj, up_proj, down_proj) if w is not None] + tensors = [x, *weights] + out = moe_experts_fp8.moe_experts(x, gate_proj, up_proj, down_proj, counts, LIMIT, fp8=True) + actual = [out, *grads(out, dout, tensors)] + expected_out = reference(x, gate_proj, up_proj, down_proj, counts, torch.float32) + expected = [expected_out, *grads(expected_out, dout, tensors)] + emulated_out, emulated_grads = fp8_emulation(x, gate_proj, up_proj, down_proj, counts, dout) + dx, dgate, dup, ddown = emulated_grads + emulated = [emulated_out, dx, torch.cat([dgate, dup], 1) if packed else dgate, *([] if packed else [dup]), ddown] + + names = ["out", "dx"] + [ + f"d{name}" for name, w in zip(("gate", "up", "down"), (gate_proj, up_proj, down_proj)) if w is not None + ] + for name, a, e, m in zip(names, actual, expected, emulated): + assert a.shape == e.shape and a.dtype == torch.bfloat16, name + assert torch.isfinite(a).all(), name + err, emulated_err = rel_err(a, e), rel_err(m, e) + # FP8 rounding (~6% per GEMM operand pair here) moves more values across the clamps than + # bf16 does, so the gradients through gate and up are further from fp32 (see the bf16 test). + tolerance = 8e-2 if name in ("out", "ddown") else 1.6e-1 + assert err < tolerance, f"{name}: rel err {err:.2e} vs fp32" + assert err < 1.02 * emulated_err + 1e-3, f"{name}: rel err {err:.2e}, FP8 emulation {emulated_err:.2e}" + # Against the emulation of the same recipe only rounding-order differences remain (and the + # clamp decisions they flip). + assert rel_err(a, m) < 2e-2, f"{name}: rel err {rel_err(a, m):.2e} vs FP8 emulation" + tail = int(counts.sum()) + assert (out[tail:] == 0).all() and (actual[1][tail:] == 0).all() + + +def test_fp8_without_grad_matches(moe_experts_fp8): + x, gate_proj, up_proj, down_proj, counts, _ = make_problem( + **FP8_CASES["small"], packed=False, seed=3, alignment=moe_experts_fp8.TOKEN_GROUP_ALIGNMENT + ) + with_grad = moe_experts_fp8.moe_experts(x, gate_proj, up_proj, down_proj, counts, LIMIT, fp8=True) + with torch.no_grad(): + without = moe_experts_fp8.moe_experts(x, gate_proj, up_proj, down_proj, counts, LIMIT, fp8=True) + torch.testing.assert_close(with_grad, without, rtol=0, atol=0) + + +def test_fp8_torch_compile(moe_experts_fp8): + x, gate_proj, up_proj, down_proj, counts, dout = make_problem( + **FP8_CASES["small"], packed=False, seed=1, alignment=moe_experts_fp8.TOKEN_GROUP_ALIGNMENT + ) + + def block(x, gate_proj, up_proj, down_proj, counts): + return moe_experts_fp8.moe_experts(x * 2, gate_proj, up_proj, down_proj, counts, LIMIT, fp8=True) + 1 + + compiled = torch.compile(block, fullgraph=True) + tensors = [x, gate_proj, up_proj, down_proj] + out = compiled(x, gate_proj, up_proj, down_proj, counts) + eager = block(x, gate_proj, up_proj, down_proj, counts) + torch.testing.assert_close(out, eager, rtol=0, atol=0) + for a, e in zip(grads(out, dout, tensors), grads(eager, dout, tensors)): + torch.testing.assert_close(a, e, rtol=0, atol=0) + + +@pytest.mark.parametrize("save_for_backward", [True, False]) +def test_fp8_opcheck(moe_experts_fp8, save_for_backward): + x, gate_proj, up_proj, down_proj, counts, _ = make_problem( + **FP8_CASES["small"], packed=False, seed=2, alignment=moe_experts_fp8.TOKEN_GROUP_ALIGNMENT + ) + # No test_schema: it does arithmetic on the outputs, which torch lacks for FP8 tensors. + torch.library.opcheck( + torch.ops.prime_kernels.moe_experts_fp8_forward.default, + (x, gate_proj, up_proj, down_proj, counts, LIMIT, save_for_backward), + test_utils=("test_faketensor", "test_autograd_registration"), + ) + + +def test_fp8_unsupported_shape_reason(moe_experts_fp8): + assert moe_experts_fp8.unsupported_shape_reason(5120, 2304, fp8=True) is None + assert moe_experts_fp8.unsupported_shape_reason(5120, 2312, fp8=True) is not None + + +@pytest.mark.parametrize("rows", [0, 64]) +def test_fp8_no_routed_tokens(moe_experts_fp8, rows): + num_experts, hidden, intermediate = 3, 256, 128 + x = torch.randn(rows, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + gate_proj, up_proj = ( + torch.randn(num_experts, intermediate, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) + for _ in range(2) + ) + down_proj = torch.randn(num_experts, hidden, intermediate, device="cuda", dtype=torch.bfloat16, requires_grad=True) + counts = torch.zeros(num_experts, dtype=torch.int64, device="cuda") + out = moe_experts_fp8.moe_experts(x, gate_proj, up_proj, down_proj, counts, LIMIT, fp8=True) + assert out.shape == x.shape and (out == 0).all() + dx, *dws = grads(out, torch.randn_like(out), [x, gate_proj, up_proj, down_proj]) + assert (dx == 0).all() and all((dw == 0).all() for dw in dws)