# SPDX-License-Identifier: CC-BY-NC-4.0 # Copyright (c) 2026 Jakub Ćwirlej """EUGEN sparse SwiGLU FFN: grouped top-k neuron selection + Triton kernels (forward and backward). HOW IT WORKS FFN neurons are split into groups of `group_size` consecutive neurons. For every token, `grouped_topk_select` scores each group by its summed |SiLU(gate)| energy (optionally on a fixed subsample of neurons per group, see `make_subsample_index`) and keeps the top `num_active_groups`. `sparse_swiglu_ffn` then computes gate, up AND down only for the selected groups — the full [tokens, d_ff] activation is never materialised, neither in the forward nor in the backward pass. `update_load_balance_bias` keeps group usage close to uniform (auxiliary-loss-free balancing). Inside `sparse_swiglu_ffn` (one torch.library.triton_op, traced by torch.compile, no graph breaks): * token->group pairs are bucketed per group (CSR) by a deterministic Triton counting sort, split into token chunks so that the atomic accumulators of one chunk stay in L2; * forward: gate and up GEMMs in one pass over X (SiLU stays in registers), then the down GEMM accumulated into y with relaxed fp32 atomics; * backward: dP with fused epilogues (SiLU' computed in place), dX for both input projections in one pass, weight gradients per group (deterministic when there is a single token chunk). Without CUDA/Triton, or for an unsupported shape, the same math runs as a dense masked PyTorch path (slow; a warning is printed once on CUDA tensors). PRACTICAL GUIDANCE 1. Density (active fraction = num_active_groups * group_size / d_ff) - Recommended: 5-10% active. Below ~5% the layer becomes too sparse; above ~15% the sparse path stops paying off against a dense FFN of the same width. Our training configuration is 6.25%: d_ff 4096, group_size 32, 8 active groups, gate subsample of 8 rows per group. - Measured at 6.25%: a single FFN block ran 1.26x (d_ff 4096) and 1.34x (d_ff 8192) faster than its dense equivalent; a full model (d_model 512, d_ff 4096) trained 1.23x faster than the dense model of the same FFN capacity, before further kernel optimisations (+36% throughput). - The gain grows with FFN width: prefer a wide FFN at low density over a narrow FFN at high density. d_ff does not have to be a power of two (e.g. 1536, 3072, 3200) — only a multiple of group_size. - Keep enough tokens per group in a micro-batch: num_tokens * num_active_groups / num_groups should be well above the token tile (64), ideally in the hundreds; small micro-batches leave tiles half empty. 2. Shapes and tiles (Triton, bf16 tensor cores) - group_size must be a power of two >= 16 (it is a tl.arange extent and a tl.dot dimension). 32 is our training value; up to 64 the fused kernels are used, above that an unfused pair of kernels (same result, more memory traffic) because of the shared-memory limit. - d_ff must be a multiple of group_size (num_groups = d_ff // group_size, any number >= 2). - d_model must be a multiple of 16; the d_model tiles (64 / 128 / 128 by default) shrink to the largest power of two that divides d_model, so multiples of 128 run at full tile width. - Use bf16: tl.dot runs on tensor cores with fp32 accumulation. The kernels inherit the input dtype, so under autocast `sparse_swiglu_ffn` casts x and the weights to the autocast dtype itself. - Gradients: x, gate, up and down weights; only the selected groups receive weight gradients. The selection itself is a hard top-k (not differentiable). 3. Correctness check - reference_sparse_swiglu_ffn is a dense masked reference: compare outputs and gradients against it on small shapes before training (test_eugen_sparse.py does this; on a machine without a GPU run it with TRITON_INTERPRET=1). Tuning knobs (environment, read at import; defaults = the configuration we trained with): EUGEN_SPARSE_CH (4) token chunks, _NT (8) token programs per (group, chunk), _BT (64) token tile, _BD (64) / _BDY (128) / _BDW (128) d_model tiles, _NS (1), _DW_DET (0; 1 = always deterministic weight gradients), _WARPS (8), _WARPS_W (8), _FUSION_G_MAX (64), _CSR_CTA (256), _CSR_CHUNK (0 = auto). """ import os import warnings from typing import Tuple import torch import torch.nn.functional as F try: import triton import triton.language as tl _HAS_TRITON = True except Exception: # pragma: no cover _HAS_TRITON = False try: from torch.library import triton_op, wrap_triton _HAS_TRITON_OP = True except Exception: # pragma: no cover _HAS_TRITON_OP = False __all__ = [ "make_subsample_index", "grouped_topk_select", "update_load_balance_bias", "sparse_swiglu_ffn", "reference_sparse_swiglu_ffn", "kernel_supported", "plan", ] def _env(name, default): return int(os.environ.get('EUGEN_SPARSE_' + name, str(default)) or default) _CH = max(1, _env('CH', 4)) _NT = _env('NT', 8) _BT = _env('BT', 64) _BD = _env('BD', 64) _BDY = _env('BDY', 128) _BDW = _env('BDW', 128) _NS = max(1, _env('NS', 1)) _DW_DET = os.environ.get('EUGEN_SPARSE_DW_DET', '0') == '1' _WARPS = _env('WARPS', 8) _WARPS_W = _env('WARPS_W', 8) _FUSION_G_MAX = _env('FUSION_G_MAX', 64) _CSR_CTA = _env('CSR_CTA', 256) _CSR_CHUNK = _env('CSR_CHUNK', 0) _SUB = 64 # counting-sort sub-tile (rank O(SUB^2) in registers) _CHUNK = 8192 # reference chunk; defines the token chunks def _fusion_ok(G): return G <= _FUSION_G_MAX def _pick_block(dim, preferred): """Largest power of two <= preferred that divides dim (0 if none >= 16).""" block = 1 << (min(preferred, dim).bit_length() - 1) while block >= 16 and dim % block != 0: block //= 2 return block if block >= 16 else 0 # ══════════════════════════════════════════════════════════════════════════════ # Selection, load balancing and the dense masked reference (PyTorch) # ══════════════════════════════════════════════════════════════════════════════ def make_subsample_index(d_ff, group_size, samples_per_group, seed=0): """Fixed gate-row subsample for grouped_topk_select: the same number of random rows inside every group. Draw ONCE per layer and keep it with the model (register_buffer(..., persistent=True)), so the network can make these rows representative and a resumed run uses the same rows. Returns (subsample_idx int64 [num_groups * samples_per_group], samples_per_group).""" num_groups = d_ff // group_size spg = max(1, min(group_size, samples_per_group)) gen = torch.Generator().manual_seed(seed) loc = torch.stack([torch.randperm(group_size, generator=gen)[:spg] for _ in range(num_groups)]) return (loc + torch.arange(num_groups).unsqueeze(1) * group_size).reshape(-1), spg @torch.no_grad() def grouped_topk_select(x, gate_weight, group_size, num_active_groups, subsample_idx=None, samples_per_group=None, load_balance_bias=None): """Indices [..., num_active_groups] of the groups with the largest |SiLU(gate)| energy. With subsample_idx the energy is estimated from those gate rows only and scaled to the full group (unbiased); load_balance_bias [num_groups] multiplies the energy by exp(bias).""" d_model = x.shape[-1] num_groups = gate_weight.shape[0] // group_size x2 = x.reshape(-1, d_model) if subsample_idx is not None: assert samples_per_group is not None, \ "samples_per_group is required together with subsample_idx" gate_act = F.silu(F.linear(x2, gate_weight[subsample_idx])) energy = gate_act.abs().view(-1, num_groups, samples_per_group).sum(-1) energy = energy * (group_size / samples_per_group) else: gate_act = F.silu(F.linear(x2, gate_weight)) energy = gate_act.abs().view(-1, num_groups, group_size).sum(-1) if load_balance_bias is not None: energy = energy * torch.exp(load_balance_bias.to(energy.dtype)) idx = energy.topk(num_active_groups, dim=-1).indices return idx.view(*x.shape[:-1], num_active_groups) @torch.no_grad() def update_load_balance_bias(bias, selection_freq_ema, current_freq, target_freq, gamma, bias_max, momentum): """Call once per optimizer step. current_freq [num_groups] = fraction of tokens that selected each group in this step; target_freq = num_active_groups / num_groups. A group whose smoothed frequency leaves the dead zone [0.5, 2] x target gets its bias moved by -gamma * sign(freq - target). Values we trained with: gamma 1e-3, bias_max 0.5, momentum 0.99.""" selection_freq_ema.mul_(momentum).add_(current_freq, alpha=1.0 - momentum) outside = (selection_freq_ema < 0.5 * target_freq) | (selection_freq_ema > 2.0 * target_freq) step = -gamma * torch.sign(selection_freq_ema - target_freq) bias.add_(torch.where(outside, step, torch.zeros_like(bias))) bias.clamp_(-bias_max, bias_max) def reference_sparse_swiglu_ffn(x, gate_weight, up_weight, down_weight, group_idx, group_size): """Dense masked reference: same math as the kernel, plain PyTorch, differentiable.""" d_model = x.shape[-1] x2 = x.reshape(-1, d_model) d_ff = gate_weight.shape[0] num_groups = d_ff // group_size idx2 = group_idx.reshape(-1, group_idx.shape[-1]).to(torch.int64) num_tokens = x2.shape[0] group_mask = torch.zeros(num_tokens, num_groups, dtype=torch.bool, device=x2.device) group_mask.scatter_(1, idx2, True) full_mask = group_mask.unsqueeze(-1).expand(num_tokens, num_groups, group_size) full_mask = full_mask.reshape(num_tokens, d_ff) gate = F.silu(F.linear(x2, gate_weight)) up = F.linear(x2, up_weight) hidden = gate * up * full_mask.to(gate.dtype) y = F.linear(hidden, down_weight) return y.view(*x.shape[:-1], d_model) def _plan_chunks(P): """(ncta, ch, r, chunk). ch = number of token chunks (affects numerics: ch == 1 and NS == 1 give deterministic weight gradients), derived ONLY from the reference chunk; ncta = counting-sort CTAs (pure performance, no effect on the result), raised towards EUGEN_SPARSE_CSR_CTA by halving the chunk while keeping ncta divisible by ch.""" ncta_ref = triton.cdiv(P, _CHUNK) if _HAS_TRITON else max(1, -(-P // _CHUNK)) ch = min(_CH, ncta_ref) while ch > 1 and ncta_ref % ch != 0: ch -= 1 if not _HAS_TRITON: return ncta_ref, ch, ncta_ref // ch, _CHUNK if _CSR_CHUNK > 0: n = triton.cdiv(P, _CSR_CHUNK) if n % ch == 0: return n, ch, n // ch, _CSR_CHUNK return ncta_ref, ch, ncta_ref // ch, _CHUNK chunk, ncta = _CHUNK, ncta_ref c = _CHUNK while c > _SUB and ncta < _CSR_CTA: c //= 2 n = triton.cdiv(P, c) if n % ch == 0: chunk, ncta = c, n assert ncta % ch == 0, (ncta, ch) return ncta, ch, ncta // ch, chunk # ══════════════════════════════════════════════════════════════════════════════ # Triton kernels. A "pair" p = tok * kg + slot indexes group_idx [M, kg]; the CSR sorts pairs by # group. Segment = (token chunk cid, group gid); off2 [ch+1, ng] holds its bounds. # ══════════════════════════════════════════════════════════════════════════════ if _HAS_TRITON: @triton.jit def _csr_hist(GIDX, HIST, P, NG: tl.constexpr, NG_POW: tl.constexpr, SUB: tl.constexpr, CHUNK: tl.constexpr): """Per-CTA histogram HIST[cta, g]. Register block NG_POW = next_power_of_2(NG); masked elements use the sentinel -1, which never matches a bin.""" pid = tl.program_id(0) bins = tl.arange(0, NG_POW) m_b = bins < NG run = tl.zeros([NG_POW], dtype=tl.int32) start = pid * CHUNK end = tl.minimum(start + CHUNK, P) for s0 in range(start, end, SUB): offs = s0 + tl.arange(0, SUB) m = offs < end gv = tl.load(GIDX + offs, mask=m, other=0) g = tl.where(m, gv, -1) oh = (g[:, None] == bins[None, :]) run += tl.sum(oh.to(tl.int32), axis=0) tl.store(HIST + pid * NG + bins, run, mask=m_b) @triton.jit def _csr_scan(HIST, COLPRE, OFF, NCTA, NG: tl.constexpr, NG_POW: tl.constexpr): """Single CTA: per-CTA column prefixes COLPRE (NCTA+1 rows; the last row holds the per-group totals) and the exclusive scan of group sizes -> OFF[ng+1].""" bins = tl.arange(0, NG_POW) m_b = bins < NG run = tl.zeros([NG_POW], dtype=tl.int32) for c in range(0, NCTA): h = tl.load(HIST + c * NG + bins, mask=m_b, other=0) tl.store(COLPRE + c * NG + bins, run, mask=m_b) run += h tl.store(COLPRE + NCTA * NG + bins, run, mask=m_b) inc = tl.cumsum(run, axis=0) tl.store(OFF + bins, inc - run, mask=m_b) tl.store(OFF + NG, tl.sum(run, axis=0)) @triton.jit def _csr_scatter(GIDX, COLPRE, OFF, PAIRS, P, NG: tl.constexpr, NG_POW: tl.constexpr, SUB: tl.constexpr, CHUNK: tl.constexpr): """Scatter with a deterministic rank: position = off[g] + CTA prefix + prefix of earlier sub-tiles + rank inside the sub-tile (O(SUB^2)). Equals a stable argsort.""" pid = tl.program_id(0) bins = tl.arange(0, NG_POW) m_b = bins < NG a = tl.arange(0, SUB) base = (tl.load(OFF + bins, mask=m_b, other=0) + tl.load(COLPRE + pid * NG + bins, mask=m_b, other=0)) run = tl.zeros([NG_POW], dtype=tl.int32) start = pid * CHUNK end = tl.minimum(start + CHUNK, P) for s0 in range(start, end, SUB): offs = s0 + a m = offs < end gv = tl.load(GIDX + offs, mask=m, other=0) g = tl.where(m, gv, -1) oh = (g[:, None] == bins[None, :]) eq = (g[:, None] == g[None, :]) & (a[None, :] < a[:, None]) lrank = tl.sum(eq.to(tl.int32), axis=1) pref = tl.sum(tl.where(oh, (base + run)[None, :], 0), axis=1) tl.store(PAIRS + pref + lrank, offs.to(tl.int32), mask=m) run += tl.sum(oh.to(tl.int32), axis=0) @triton.jit def _gate_up(X, WG, WU, GACC, UP, PSEL, PAIRS, OFF2, D: tl.constexpr, FDIM: tl.constexpr, G: tl.constexpr, KG: tl.constexpr, NG: tl.constexpr, NT: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr): """Gate and up projections for the selected groups in ONE pass over the CSR (one gather of X): GACC = X @ WG[g]^T (rounded to the input dtype BEFORE silu, exactly as F.silu(F.linear(...)) does under autocast), UP = X @ WU[g]^T, PSEL = silu(GACC) * UP. silu(GACC) lives only in registers; the backward recomputes it from GACC.""" gc = tl.program_id(0) gid = gc % NG cid = gc // NG pid_t = tl.program_id(1) lo = tl.load(OFF2 + cid * NG + gid) hi = tl.load(OFF2 + (cid + 1) * NG + gid) offs_g = tl.arange(0, G) start = lo + pid_t * BT while start < hi: offs_t = start + tl.arange(0, BT) m_t = offs_t < hi par = tl.load(PAIRS + offs_t, mask=m_t, other=0) tq = par // KG slot = par - tq * KG tok = tq.to(tl.int64) acc_g = tl.zeros((BT, G), dtype=tl.float32) acc_u = tl.zeros((BT, G), dtype=tl.float32) for d0 in range(0, D, BD): offs_d = d0 + tl.arange(0, BD) xb = tl.load(X + tok[:, None] * D + offs_d[None, :], mask=m_t[:, None], other=0.0) wgb = tl.load(WG + (gid * G + offs_g)[None, :].to(tl.int64) * D + offs_d[:, None]) wub = tl.load(WU + (gid * G + offs_g)[None, :].to(tl.int64) * D + offs_d[:, None]) acc_g += tl.dot(xb, wgb) acc_u += tl.dot(xb, wub) base = tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :] gb = acc_g.to(GACC.dtype.element_ty) tl.store(GACC + base, gb, mask=m_t[:, None]) gf = gb.to(tl.float32) sb = gf * tl.sigmoid(gf) sbc = sb.to(PSEL.dtype.element_ty) ub = acc_u.to(UP.dtype.element_ty) tl.store(UP + base, ub, mask=m_t[:, None]) tl.store(PSEL + base, sbc * ub, mask=m_t[:, None]) start += NT * BT @triton.jit def _gate(X, WG, GACC, SIL, PAIRS, OFF2, D: tl.constexpr, G: tl.constexpr, KG: tl.constexpr, NG: tl.constexpr, NT: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr): """Gate projection for the selected groups only (used when G is above the fusion threshold): GACC = X @ WG[g]^T rounded to the input dtype, SIL = silu(GACC).""" gc = tl.program_id(0) gid = gc % NG cid = gc // NG pid_t = tl.program_id(1) lo = tl.load(OFF2 + cid * NG + gid) hi = tl.load(OFF2 + (cid + 1) * NG + gid) offs_g = tl.arange(0, G) start = lo + pid_t * BT while start < hi: offs_t = start + tl.arange(0, BT) m_t = offs_t < hi par = tl.load(PAIRS + offs_t, mask=m_t, other=0) tq = par // KG slot = par - tq * KG tok = tq.to(tl.int64) acc = tl.zeros((BT, G), dtype=tl.float32) for d0 in range(0, D, BD): offs_d = d0 + tl.arange(0, BD) xb = tl.load(X + tok[:, None] * D + offs_d[None, :], mask=m_t[:, None], other=0.0) wb = tl.load(WG + (gid * G + offs_g)[None, :].to(tl.int64) * D + offs_d[:, None]) acc += tl.dot(xb, wb) base = tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :] gb = acc.to(GACC.dtype.element_ty) tl.store(GACC + base, gb, mask=m_t[:, None]) gf = gb.to(tl.float32) sb = gf * tl.sigmoid(gf) tl.store(SIL + base, sb.to(SIL.dtype.element_ty), mask=m_t[:, None]) start += NT * BT @triton.jit def _up(X, WU, SIL, UP, PSEL, PAIRS, OFF2, D: tl.constexpr, FDIM: tl.constexpr, G: tl.constexpr, KG: tl.constexpr, NG: tl.constexpr, NT: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr): """UP = X @ WU[g]^T for the selected groups, PSEL = SIL * UP (used when G is above the fusion threshold).""" gc = tl.program_id(0) gid = gc % NG cid = gc // NG pid_t = tl.program_id(1) lo = tl.load(OFF2 + cid * NG + gid) hi = tl.load(OFF2 + (cid + 1) * NG + gid) offs_g = tl.arange(0, G) start = lo + pid_t * BT while start < hi: offs_t = start + tl.arange(0, BT) m_t = offs_t < hi par = tl.load(PAIRS + offs_t, mask=m_t, other=0) tq = par // KG slot = par - tq * KG tok = tq.to(tl.int64) acc = tl.zeros((BT, G), dtype=tl.float32) for d0 in range(0, D, BD): offs_d = d0 + tl.arange(0, BD) xb = tl.load(X + tok[:, None] * D + offs_d[None, :], mask=m_t[:, None], other=0.0) wb = tl.load(WU + (gid * G + offs_g)[None, :].to(tl.int64) * D + offs_d[:, None]) acc += tl.dot(xb, wb) ub = acc.to(UP.dtype.element_ty) sb = tl.load(SIL + tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :], mask=m_t[:, None], other=0.0) base = tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :] tl.store(UP + base, ub, mask=m_t[:, None]) tl.store(PSEL + base, sb * ub, mask=m_t[:, None]) start += NT * BT @triton.jit def _down(PSEL, WD, Y, PAIRS, OFF2, D: tl.constexpr, FDIM: tl.constexpr, G: tl.constexpr, KG: tl.constexpr, NG: tl.constexpr, NT: tl.constexpr, BT: tl.constexpr, BDY: tl.constexpr, WDT: tl.constexpr = 0): """Y[tok, :] += PSEL[tok, slot, :] @ WD[:, g]^T, fp32 atomic_add with relaxed semantics (Y is read only after the kernel finishes). D tile = BDY.""" gc = tl.program_id(0) gid = gc % NG cid = gc // NG pid_t = tl.program_id(1) pid_d = tl.program_id(2) lo = tl.load(OFF2 + cid * NG + gid) hi = tl.load(OFF2 + (cid + 1) * NG + gid) offs_g = tl.arange(0, G) offs_d = pid_d * BDY + tl.arange(0, BDY) start = lo + pid_t * BT while start < hi: offs_t = start + tl.arange(0, BT) m_t = offs_t < hi par = tl.load(PAIRS + offs_t, mask=m_t, other=0) tq = par // KG slot = par - tq * KG tok = tq.to(tl.int64) pb = tl.load(PSEL + tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :], mask=m_t[:, None], other=0.0) if WDT: wb = tl.load(WD + (gid * G + offs_g)[:, None].to(tl.int64) * D + offs_d[None, :]) else: wb = tl.load(WD + offs_d[None, :].to(tl.int64) * FDIM + (gid * G + offs_g)[:, None]) acc = tl.dot(pb, wb) tl.atomic_add(Y + tok[:, None] * D + offs_d[None, :], acc, mask=m_t[:, None], sem='relaxed') start += NT * BT @triton.jit def _dp(DY, WD, GACC, UPSEL, DGATE, DUP, PAIRS, OFF2, D: tl.constexpr, FDIM: tl.constexpr, G: tl.constexpr, KG: tl.constexpr, NG: tl.constexpr, NT: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr, WDT: tl.constexpr = 0): """dP = dY @ WD[:, g] with fused epilogues: dgate = dP * up * silu'(gacc) and dup = dP * silu(gacc); dP never touches memory, silu is recomputed from GACC.""" gc = tl.program_id(0) gid = gc % NG cid = gc // NG pid_t = tl.program_id(1) lo = tl.load(OFF2 + cid * NG + gid) hi = tl.load(OFF2 + (cid + 1) * NG + gid) offs_g = tl.arange(0, G) start = lo + pid_t * BT while start < hi: offs_t = start + tl.arange(0, BT) m_t = offs_t < hi par = tl.load(PAIRS + offs_t, mask=m_t, other=0) tq = par // KG slot = par - tq * KG tok = tq.to(tl.int64) acc = tl.zeros((BT, G), dtype=tl.float32) for d0 in range(0, D, BD): offs_d = d0 + tl.arange(0, BD) dyb = tl.load(DY + tok[:, None] * D + offs_d[None, :], mask=m_t[:, None], other=0.0) if WDT: wb = tl.load(WD + (gid * G + offs_g)[None, :].to(tl.int64) * D + offs_d[:, None]) else: wb = tl.load(WD + offs_d[:, None].to(tl.int64) * FDIM + gid * G + offs_g[None, :]) acc += tl.dot(dyb, wb) base = tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :] gb = tl.load(GACC + base, mask=m_t[:, None], other=0.0).to(tl.float32) sg = tl.sigmoid(gb) sb = gb * sg ub = tl.load(UPSEL + base, mask=m_t[:, None], other=0.0).to(tl.float32) tl.store(DGATE + base, (acc * ub * sg * (1.0 + gb * (1.0 - sg))).to(DGATE.dtype.element_ty), mask=m_t[:, None]) tl.store(DUP + base, (acc * sb).to(DUP.dtype.element_ty), mask=m_t[:, None]) start += NT * BT @triton.jit def _dx2(DUP, WU, DGATE, WG, DX, PAIRS, OFF2, D: tl.constexpr, FDIM: tl.constexpr, G: tl.constexpr, KG: tl.constexpr, NG: tl.constexpr, NT: tl.constexpr, BT: tl.constexpr, BDY: tl.constexpr): """dX[tok, :] += dup @ WU[g] + dgate @ WG[g] in one pass (both contributions summed in registers before a single relaxed fp32 atomic_add).""" gc = tl.program_id(0) gid = gc % NG cid = gc // NG pid_t = tl.program_id(1) pid_d = tl.program_id(2) lo = tl.load(OFF2 + cid * NG + gid) hi = tl.load(OFF2 + (cid + 1) * NG + gid) offs_g = tl.arange(0, G) offs_d = pid_d * BDY + tl.arange(0, BDY) start = lo + pid_t * BT while start < hi: offs_t = start + tl.arange(0, BT) m_t = offs_t < hi par = tl.load(PAIRS + offs_t, mask=m_t, other=0) tq = par // KG slot = par - tq * KG tok = tq.to(tl.int64) base = tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :] dub = tl.load(DUP + base, mask=m_t[:, None], other=0.0) dgb = tl.load(DGATE + base, mask=m_t[:, None], other=0.0) wrow = (gid * G + offs_g)[:, None].to(tl.int64) * D + offs_d[None, :] wu = tl.load(WU + wrow) wg = tl.load(WG + wrow) acc = tl.dot(dub, wu) + tl.dot(dgb, wg) tl.atomic_add(DX + tok[:, None] * D + offs_d[None, :], acc, mask=m_t[:, None], sem='relaxed') start += NT * BT @triton.jit def _dx(DUP, WU, DX, PAIRS, OFF2, D: tl.constexpr, FDIM: tl.constexpr, G: tl.constexpr, KG: tl.constexpr, NG: tl.constexpr, NT: tl.constexpr, BT: tl.constexpr, BDY: tl.constexpr): """dX[tok, :] += dsel @ W[g] (relaxed fp32 atomic_add); used twice when G is above the fusion threshold.""" gc = tl.program_id(0) gid = gc % NG cid = gc // NG pid_t = tl.program_id(1) pid_d = tl.program_id(2) lo = tl.load(OFF2 + cid * NG + gid) hi = tl.load(OFF2 + (cid + 1) * NG + gid) offs_g = tl.arange(0, G) offs_d = pid_d * BDY + tl.arange(0, BDY) start = lo + pid_t * BT while start < hi: offs_t = start + tl.arange(0, BT) m_t = offs_t < hi par = tl.load(PAIRS + offs_t, mask=m_t, other=0) tq = par // KG slot = par - tq * KG tok = tq.to(tl.int64) dub = tl.load(DUP + tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :], mask=m_t[:, None], other=0.0) wb = tl.load(WU + (gid * G + offs_g)[:, None].to(tl.int64) * D + offs_d[None, :]) acc = tl.dot(dub, wb) tl.atomic_add(DX + tok[:, None] * D + offs_d[None, :], acc, mask=m_t[:, None], sem='relaxed') start += NT * BT @triton.jit def _dwd(DY, GACC, UPSEL, DWD, PAIRS, OFF2, D: tl.constexpr, FDIM: tl.constexpr, G: tl.constexpr, KG: tl.constexpr, NG: tl.constexpr, NS: tl.constexpr, BT: tl.constexpr, BDW: tl.constexpr, ATOM: tl.constexpr): """dWD[d tile, g] = sum over the segment of dY^T @ (silu(GACC) * UP). ATOM=False (one token chunk, NS=1): direct deterministic store; otherwise fp32 atomic_add.""" gc = tl.program_id(0) gid = gc % NG cid = gc // NG pid_d = tl.program_id(1) pid_s = tl.program_id(2) lo = tl.load(OFF2 + cid * NG + gid) hi = tl.load(OFF2 + (cid + 1) * NG + gid) n = hi - lo seg = (n + NS - 1) // NS s0 = lo + pid_s * seg s1 = tl.minimum(hi, s0 + seg) offs_g = tl.arange(0, G) offs_d = pid_d * BDW + tl.arange(0, BDW) acc = tl.zeros((BDW, G), dtype=tl.float32) for st in range(s0, s1, BT): offs_t = st + tl.arange(0, BT) m_t = offs_t < s1 par = tl.load(PAIRS + offs_t, mask=m_t, other=0) tq = par // KG slot = par - tq * KG tok = tq.to(tl.int64) dyb = tl.load(DY + tok[:, None] * D + offs_d[None, :], mask=m_t[:, None], other=0.0) _b = tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :] _g = tl.load(GACC + _b, mask=m_t[:, None], other=0.0).to(tl.float32) sb = (_g * tl.sigmoid(_g)).to(dyb.dtype) ub = tl.load(UPSEL + _b, mask=m_t[:, None], other=0.0) acc += tl.dot(tl.trans(dyb), sb * ub) ptr = DWD + offs_d[:, None].to(tl.int64) * FDIM + gid * G + offs_g[None, :] if ATOM: tl.atomic_add(ptr, acc, sem='relaxed') else: tl.store(ptr, acc.to(DWD.dtype.element_ty)) @triton.jit def _dwu2(X, DUP, DGATE, DWU, DWG, PAIRS, OFF2, D: tl.constexpr, FDIM: tl.constexpr, G: tl.constexpr, KG: tl.constexpr, NG: tl.constexpr, NS: tl.constexpr, BT: tl.constexpr, BDW: tl.constexpr, ATOM: tl.constexpr): """dWU[g, d tile] and dWG[g, d tile] in one pass over the segment (one gather of X).""" gc = tl.program_id(0) gid = gc % NG cid = gc // NG pid_d = tl.program_id(1) pid_s = tl.program_id(2) lo = tl.load(OFF2 + cid * NG + gid) hi = tl.load(OFF2 + (cid + 1) * NG + gid) n = hi - lo seg = (n + NS - 1) // NS s0 = lo + pid_s * seg s1 = tl.minimum(hi, s0 + seg) offs_g = tl.arange(0, G) offs_d = pid_d * BDW + tl.arange(0, BDW) acc_u = tl.zeros((G, BDW), dtype=tl.float32) acc_g = tl.zeros((G, BDW), dtype=tl.float32) for st in range(s0, s1, BT): offs_t = st + tl.arange(0, BT) m_t = offs_t < s1 par = tl.load(PAIRS + offs_t, mask=m_t, other=0) tq = par // KG slot = par - tq * KG tok = tq.to(tl.int64) base = tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :] dub = tl.load(DUP + base, mask=m_t[:, None], other=0.0) dgb = tl.load(DGATE + base, mask=m_t[:, None], other=0.0) xb = tl.load(X + tok[:, None] * D + offs_d[None, :], mask=m_t[:, None], other=0.0) acc_u += tl.dot(tl.trans(dub), xb) acc_g += tl.dot(tl.trans(dgb), xb) row = (gid * G + offs_g)[:, None].to(tl.int64) * D + offs_d[None, :] if ATOM: tl.atomic_add(DWU + row, acc_u, sem='relaxed') tl.atomic_add(DWG + row, acc_g, sem='relaxed') else: tl.store(DWU + row, acc_u.to(DWU.dtype.element_ty)) tl.store(DWG + row, acc_g.to(DWG.dtype.element_ty)) @triton.jit def _dwu(X, DUP, DWU, PAIRS, OFF2, D: tl.constexpr, FDIM: tl.constexpr, G: tl.constexpr, KG: tl.constexpr, NG: tl.constexpr, NS: tl.constexpr, BT: tl.constexpr, BDW: tl.constexpr, ATOM: tl.constexpr): """dW[g, d tile] = sum over the segment of dsel^T @ X; used twice when G is above the fusion threshold.""" gc = tl.program_id(0) gid = gc % NG cid = gc // NG pid_d = tl.program_id(1) pid_s = tl.program_id(2) lo = tl.load(OFF2 + cid * NG + gid) hi = tl.load(OFF2 + (cid + 1) * NG + gid) n = hi - lo seg = (n + NS - 1) // NS s0 = lo + pid_s * seg s1 = tl.minimum(hi, s0 + seg) offs_g = tl.arange(0, G) offs_d = pid_d * BDW + tl.arange(0, BDW) acc = tl.zeros((G, BDW), dtype=tl.float32) for st in range(s0, s1, BT): offs_t = st + tl.arange(0, BT) m_t = offs_t < s1 par = tl.load(PAIRS + offs_t, mask=m_t, other=0) tq = par // KG slot = par - tq * KG tok = tq.to(tl.int64) dub = tl.load(DUP + tok[:, None] * (KG * G) + slot[:, None] * G + offs_g[None, :], mask=m_t[:, None], other=0.0) xb = tl.load(X + tok[:, None] * D + offs_d[None, :], mask=m_t[:, None], other=0.0) acc += tl.dot(tl.trans(dub), xb) ptr = DWU + (gid * G + offs_g)[:, None].to(tl.int64) * D + offs_d[None, :] if ATOM: tl.atomic_add(ptr, acc, sem='relaxed') else: tl.store(ptr, acc.to(DWU.dtype.element_ty)) # ══════════════════════════════════════════════════════════════════════════════ # triton_op: the op body is traced by dynamo/AOT (allocations, casts and zeros become graph # nodes); kernels enter the graph through wrap_triton. # ══════════════════════════════════════════════════════════════════════════════ if _HAS_TRITON and _HAS_TRITON_OP: # TRITON_INTERPRET=1 (CPU debugging): interpreted kernels are not triton.jit objects and # wrap_triton rejects them, so they are launched directly (eager only). _wrap = (lambda fn: fn) if os.environ.get('TRITON_INTERPRET', '0') == '1' else wrap_triton def _csr(idx, ng, kg): """(pairs [P] int32, off2 [ch+1, ng] int32, ch). off2[j, g] = start of the pairs of group g coming from token chunk j; off2[ch, g] = end of the segment of group g.""" M = idx.shape[0] P = M * kg g = idx.reshape(-1) ncta, ch, r, chunk = _plan_chunks(P) ngp = triton.next_power_of_2(ng) hist = torch.empty(ncta, ng, dtype=torch.int32, device=idx.device) colpre = torch.empty(ncta + 1, ng, dtype=torch.int32, device=idx.device) off = torch.empty(ng + 1, dtype=torch.int32, device=idx.device) pairs = torch.empty(P, dtype=torch.int32, device=idx.device) _wrap(_csr_hist)[(ncta,)](g, hist, P, NG=ng, NG_POW=ngp, SUB=_SUB, CHUNK=chunk) _wrap(_csr_scan)[(1,)](hist, colpre, off, ncta, NG=ng, NG_POW=ngp) _wrap(_csr_scatter)[(ncta,)](g, colpre, off, pairs, P, NG=ng, NG_POW=ngp, SUB=_SUB, CHUNK=chunk) rows = torch.arange(0, ch + 1, device=idx.device, dtype=torch.int64) * r off2 = off[:ng].unsqueeze(0) + colpre.index_select(0, rows) return pairs, off2, ch def _fwd_core(xc, gwc, uwc, dwc, idx, G): M, D = xc.shape bd, bdy, bdw = _pick_block(D, _BD), _pick_block(D, _BDY), _pick_block(D, _BDW) Fdim = uwc.shape[0] ng = Fdim // G kg = idx.shape[1] pairs, off2, ch = _csr(idx, ng, kg) gacc = torch.empty(M, kg * G, device=xc.device, dtype=xc.dtype) up_sel = torch.empty(M, kg * G, device=xc.device, dtype=xc.dtype) p_sel = torch.empty(M, kg * G, device=xc.device, dtype=xc.dtype) if _fusion_ok(G): _wrap(_gate_up)[(ng * ch, _NT)]( xc, gwc, uwc, gacc, up_sel, p_sel, pairs, off2, D=D, FDIM=Fdim, G=G, KG=kg, NG=ng, NT=_NT, BT=_BT, BD=bd, num_warps=_WARPS) else: # large G: unfused pair (shared memory limit) sil_tmp = torch.empty(M, kg * G, device=xc.device, dtype=xc.dtype) _wrap(_gate)[(ng * ch, _NT)]( xc, gwc, gacc, sil_tmp, pairs, off2, D=D, G=G, KG=kg, NG=ng, NT=_NT, BT=_BT, BD=bd, num_warps=_WARPS) _wrap(_up)[(ng * ch, _NT)]( xc, uwc, sil_tmp, up_sel, p_sel, pairs, off2, D=D, FDIM=Fdim, G=G, KG=kg, NG=ng, NT=_NT, BT=_BT, BD=bd, num_warps=_WARPS) y32 = torch.zeros(M, D, device=xc.device, dtype=torch.float32) _wrap(_down)[(ng * ch, _NT, D // bdy)]( p_sel, dwc, y32, pairs, off2, D=D, FDIM=Fdim, G=G, KG=kg, NG=ng, NT=_NT, BT=_BT, BDY=bdy, num_warps=_WARPS) return y32.to(xc.dtype), gacc, up_sel, pairs, off2 @triton_op("eugen_sparse::sparse_swiglu_fwd", mutates_args={}) def sparse_swiglu_fwd(x: torch.Tensor, gate_w: torch.Tensor, up_w: torch.Tensor, down_w: torch.Tensor, idx: torch.Tensor, G: int ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """x [M, D]; gate_w, up_w [F, D]; down_w [D, F]; idx [M, kg]. Returns (y [M, D], gacc [M, kg*G], up_sel [M, kg*G], pairs, off2) — compact, no [M, F].""" return _fwd_core(x.contiguous(), gate_w.contiguous(), up_w.contiguous(), down_w.contiguous(), idx.to(torch.int32).contiguous(), G) def _bwd(dy, x, gate_w, up_w, down_w, gacc, up_sel, pairs, off2, G): """Returns (dX, dW_gate, dW_up, dW_down).""" dyc = dy.contiguous(); xc = x.contiguous() gwc = gate_w.contiguous(); uwc = up_w.contiguous(); dwc = down_w.contiguous() gaccc = gacc.contiguous() M, D = xc.shape bd, bdy, bdw = _pick_block(D, _BD), _pick_block(D, _BDY), _pick_block(D, _BDW) Fdim = uwc.shape[0] K = up_sel.shape[1] kg = K // G ng = Fdim // G ch = off2.shape[0] - 1 dgate = torch.empty(M, K, device=dyc.device, dtype=xc.dtype) dup = torch.empty(M, K, device=dyc.device, dtype=xc.dtype) _wrap(_dp)[(ng * ch, _NT)]( dyc, dwc, gaccc, up_sel, dgate, dup, pairs, off2, D=D, FDIM=Fdim, G=G, KG=kg, NG=ng, NT=_NT, BT=_BT, BD=bd, num_warps=_WARPS) dX32 = torch.zeros(M, D, device=dyc.device, dtype=torch.float32) if _fusion_ok(G): _wrap(_dx2)[(ng * ch, _NT, D // bdy)]( dup, uwc, dgate, gwc, dX32, pairs, off2, D=D, FDIM=Fdim, G=G, KG=kg, NG=ng, NT=_NT, BT=_BT, BDY=bdy, num_warps=_WARPS) else: for dsel, w in ((dup, uwc), (dgate, gwc)): _wrap(_dx)[(ng * ch, _NT, D // bdy)]( dsel, w, dX32, pairs, off2, D=D, FDIM=Fdim, G=G, KG=kg, NG=ng, NT=_NT, BT=_BT, BDY=bdy, num_warps=_WARPS) dX = dX32.to(dy.dtype) if _DW_DET or (ch == 1 and _NS == 1): # whole group segment in one CTA: direct, deterministic store off2w = off2 if ch == 1 else torch.stack([off2[0], off2[-1]]) chw, atom = 1, False dWd = torch.empty(D, Fdim, device=dyc.device, dtype=down_w.dtype) dWu = torch.empty(Fdim, D, device=dyc.device, dtype=up_w.dtype) dWg = torch.empty(Fdim, D, device=dyc.device, dtype=gate_w.dtype) else: off2w = off2; chw, atom = ch, True dWd = torch.zeros(D, Fdim, device=dyc.device, dtype=torch.float32) dWu = torch.zeros(Fdim, D, device=dyc.device, dtype=torch.float32) dWg = torch.zeros(Fdim, D, device=dyc.device, dtype=torch.float32) _wrap(_dwd)[(ng * chw, D // bdw, _NS)]( dyc, gaccc, up_sel, dWd, pairs, off2w, D=D, FDIM=Fdim, G=G, KG=kg, NG=ng, NS=_NS, BT=_BT, BDW=bdw, ATOM=atom, num_warps=_WARPS_W) if _fusion_ok(G): _wrap(_dwu2)[(ng * chw, D // bdw, _NS)]( xc, dup, dgate, dWu, dWg, pairs, off2w, D=D, FDIM=Fdim, G=G, KG=kg, NG=ng, NS=_NS, BT=_BT, BDW=bdw, ATOM=atom, num_warps=_WARPS_W) else: for dsel, dw in ((dup, dWu), (dgate, dWg)): _wrap(_dwu)[(ng * chw, D // bdw, _NS)]( xc, dsel, dw, pairs, off2w, D=D, FDIM=Fdim, G=G, KG=kg, NG=ng, NS=_NS, BT=_BT, BDW=bdw, ATOM=atom, num_warps=_WARPS_W) if atom: dWd = dWd.to(down_w.dtype) dWu = dWu.to(up_w.dtype) dWg = dWg.to(gate_w.dtype) return dX, dWg, dWu, dWd def _setup_ctx(ctx, inputs, output): x, gate_w, up_w, down_w, idx, G = inputs y, gacc, up_sel, pairs, off2 = output ctx.save_for_backward(x, gate_w, up_w, down_w, gacc, up_sel, pairs, off2) ctx.G = G def _backward(ctx, dy, d_gacc, d_up_sel, d_pairs, d_off2): x, gate_w, up_w, down_w, gacc, up_sel, pairs, off2 = ctx.saved_tensors dX, dWg, dWu, dWd = _bwd(dy, x, gate_w, up_w, down_w, gacc, up_sel, pairs, off2, ctx.G) return dX, dWg, dWu, dWd, None, None torch.library.register_autograd("eugen_sparse::sparse_swiglu_fwd", _backward, setup_context=_setup_ctx) # ══════════════════════════════════════════════════════════════════════════════ # Public entry point # ══════════════════════════════════════════════════════════════════════════════ def kernel_supported(is_cuda, d_model, d_ff, group_size, num_active_groups=None): """True if the Triton path will be used: CUDA, Triton with torch.library.triton_op, d_model a multiple of 16, group_size a power of two >= 16, d_ff a multiple of group_size, >= 2 groups.""" G = group_size ok = (_HAS_TRITON and _HAS_TRITON_OP and is_cuda and _pick_block(d_model, _BD) >= 16 and G >= 16 and (G & (G - 1)) == 0 and d_ff % G == 0 and d_ff // G >= 2) return ok and (num_active_groups is None or 1 <= num_active_groups <= d_ff // G) def plan(num_tokens, d_model, d_ff, group_size, num_active_groups): """Diagnostics (no GPU work): will the kernel run and how the work is split.""" P = num_tokens * num_active_groups ncta, ch, r, chunk = _plan_chunks(P) return dict(kernel=kernel_supported(True, d_model, d_ff, group_size, num_active_groups), num_groups=d_ff // group_size, pairs=P, token_chunks=ch, csr_ctas=ncta, csr_chunk=chunk, tiles=dict(BD=_pick_block(d_model, _BD), BDY=_pick_block(d_model, _BDY), BDW=_pick_block(d_model, _BDW)), fused=_fusion_ok(group_size), deterministic_weight_grads=bool(_DW_DET or (ch == 1 and _NS == 1))) _WARNED = [False] def _autocast_dtype(device_type): try: if torch.is_autocast_enabled(device_type): return torch.get_autocast_dtype(device_type) except TypeError: # older PyTorch API if device_type == 'cuda' and torch.is_autocast_enabled(): return torch.get_autocast_gpu_dtype() return None def sparse_swiglu_ffn(x, gate_weight, up_weight, down_weight, group_idx, group_size): """x [..., d_model]; gate_weight, up_weight [d_ff, d_model]; down_weight [d_model, d_ff] (nn.Linear layout); group_idx [..., num_active_groups] from grouped_topk_select. Returns y [..., d_model] = sum over the selected groups of (SiLU(x W_g^T) * x W_u^T) W_d^T.""" d_model = x.shape[-1] x2 = x.reshape(-1, d_model) i2 = group_idx.reshape(-1, group_idx.shape[-1]) if kernel_supported(x.is_cuda, d_model, up_weight.shape[0], group_size, i2.shape[-1]): dt = _autocast_dtype(x.device.type) args = (x2, gate_weight, up_weight, down_weight) if dt is not None: # kernels inherit the input dtype args = tuple(t.to(dt) for t in args) y = sparse_swiglu_fwd(*args, i2, group_size)[0] else: if x.is_cuda and not _WARNED[0]: _WARNED[0] = True warnings.warn('eugen_sparse: Triton kernel not used for this shape/environment ' '(see kernel_supported); running the slow dense masked path') y = reference_sparse_swiglu_ffn(x2, gate_weight, up_weight, down_weight, i2, group_size) return y.view(*x.shape[:-1], d_model)