Blame for eugen_sparse.py

# 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) against the dense model of the same FFN
     capacity: the first version of this kernel trained 1.23x faster; the current version reaches
     36% higher training throughput than the dense model.
   - 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)