kernels/pack_experts.py

"""pack_experts.py — swap DiffusionGemmaTextExperts BF16 weights for packed W4A16 + a Triton grouped
forward. Mirrors the stock forward EXACTLY (routing/gather/act/index_add unchanged); only the two
nn.functional.linear calls become w4a16_linear (packed 4-bit, dequant inline). The evaluated model is
byte-identical to what the kernel computes (kill-test proves it).

gate_up_proj: per-expert [2*704, 2816], K=2816 -> group 128. down_proj: [2816, 704], K=704 -> group 64.
"""
import sys, pathlib, torch, torch.nn as nn, torch.nn.functional as F
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
from w4a16 import quantize_w4, w4a16_linear

GU_GS, DN_GS = 128, 64

# MoE ROUTING PROBE (MOE_STATS=1): does the expert UNION shrink if we drop top-k? Dynamic-k only cuts
# MoE bandwidth if top-4 union << top-8 union (else coupon-collector floors the union near num_experts).
import os as _os, collections as _coll
_MOE_ACC = _coll.defaultdict(list)


def dump_moe_stats():
    if not _MOE_ACC.get("u8"):
        return
    m = lambda x: sum(x) / len(x)
    E = _MOE_ACC["E"][0]
    print(f"\n=== MoE ROUTING PROBE (n={len(_MOE_ACC['u8'])} MoE-layer-forwards, {E} experts) ===", flush=True)
    print(f"  expert UNION loaded per forward:  top-8={m(_MOE_ACC['u8']):.0f}  top-6={m(_MOE_ACC['u6']):.0f}  "
          f"top-4={m(_MOE_ACC['u4']):.0f}  top-2={m(_MOE_ACC['u2']):.0f}   (of {E})", flush=True)
    print(f"  MoE bandwidth vs top-8:  top-6={100*m(_MOE_ACC['u6'])/m(_MOE_ACC['u8']):.0f}%  "
          f"top-4={100*m(_MOE_ACC['u4'])/m(_MOE_ACC['u8']):.0f}%  top-2={100*m(_MOE_ACC['u2'])/m(_MOE_ACC['u8']):.0f}%", flush=True)
    print(f"  experts holding 80% of routing mass: {m(_MOE_ACC['conc']):.0f}/{E}", flush=True)
    print(f"  VERDICT: dynamic-k moves tok/s IFF top-4 union << top-8 (concentrated). If ~equal -> coupon-collector, no BW win.", flush=True)


@torch.no_grad()
def _pack_3d(W3d, gs):
    """W3d [E, N, K] -> lists of per-expert (qpacked, scale, zero) stacked to [E, ...]."""
    qs, ss, zs = [], [], []
    for e in range(W3d.shape[0]):
        q, s, z = quantize_w4(W3d[e].float(), gs)
        qs.append(q); ss.append(s); zs.append(z)
    return torch.stack(qs), torch.stack(ss), torch.stack(zs)


def packed_experts_forward(self, hidden_states, top_k_index, top_k_weights):
    """Identical to DiffusionGemmaTextExperts.forward but the two linears read packed W4."""
    final_hidden_states = torch.zeros_like(hidden_states)
    with torch.no_grad():
        expert_mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0)
        expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
    for expert_idx in expert_hit:
        e = expert_idx[0]
        if e == self.num_experts:
            continue
        top_k_pos, token_idx = torch.where(expert_mask[e])
        current_state = hidden_states[token_idx]
        gate_up = w4a16_linear(current_state, self.gu_q[e], self.gu_s[e], self.gu_z[e], BK=GU_GS)
        gate, up = gate_up.chunk(2, dim=-1)
        h = self.act_fn(gate) * up
        h = w4a16_linear(h, self.dn_q[e], self.dn_s[e], self.dn_z[e], BK=DN_GS)
        h = h * top_k_weights[token_idx, top_k_pos, None]
        final_hidden_states.index_add_(0, token_idx, h.to(final_hidden_states.dtype))
    return final_hidden_states


def fused_experts_forward(self, hidden_states, top_k_index, top_k_weights):
    """Single fused grouped-GEMM (one launch for all experts) — replaces the launch-bound loop.
    MOE_V2=1 selects the deeper-fused kernels (act into GEMM1, scatter into GEMM2)."""
    import os
    if os.environ.get("MOE_STATS"):
        with torch.no_grad():                                    # measure expert union at each top-k
            order = top_k_weights.argsort(dim=-1, descending=True)   # ensure [:, :k] = true top-k
            tis = torch.gather(top_k_index, -1, order)
            tk = top_k_index.shape[1]
            for k, key in ((8, "u8"), (6, "u6"), (4, "u4"), (2, "u2")):
                _MOE_ACC[key].append(int(tis[:, :min(k, tk)].unique().numel()))
            hist = torch.bincount(top_k_index.reshape(-1), minlength=self.num_experts).float()
            cum = torch.cumsum(torch.sort(hist, descending=True).values, 0) / hist.sum().clamp(min=1)
            _MOE_ACC["conc"].append(int((cum < 0.8).sum()) + 1)
            _MOE_ACC["E"].append(self.num_experts)
    if os.environ.get("MOE_V2") == "1":
        from fused_moe_w4_v2 import fused_moe_w4_v2
        return fused_moe_w4_v2(hidden_states, self, top_k_index, top_k_weights)
    from fused_moe_w4 import fused_moe_w4
    return fused_moe_w4(hidden_states, self, top_k_index, top_k_weights)


@torch.no_grad()
def pack_experts_module(exp, free_bf16=True, fused=True):
    """Pack one DiffusionGemmaTextExperts in place: add packed buffers, bind packed forward."""
    gu_q, gu_s, gu_z = _pack_3d(exp.gate_up_proj.data, GU_GS)
    I = exp.down_proj.shape[-1]                              # 704 (contraction dim of down)
    # down stays UNPADDED GS=64/BK=64: the 704->768 pad (BK=128) was 1.27x on the isolated kernel but its
    # GS 64->128 coarsening cut commits/fwd 44.4->41.9 -> NET -10 tok/s end-to-end (pad_ab.log). D2F coupling.
    dn_q, dn_s, dn_z = _pack_3d(exp.down_proj.data, DN_GS)
    dev = exp.gate_up_proj.device
    for nm, t in [("gu_q", gu_q), ("gu_s", gu_s), ("gu_z", gu_z), ("dn_q", dn_q), ("dn_s", dn_s), ("dn_z", dn_z)]:
        exp.register_buffer(nm, t.to(dev), persistent=False)
    if free_bf16:
        exp.gate_up_proj = None
        exp.down_proj = None
    if _os.environ.get("MOE_NS") == "1" and _os.environ.get("MOE_V2") == "1":
        from fused_moe_w4_v2 import enable_ns_down               # preshuffled-T down pack (v2-only)
        enable_ns_down(exp)
    fwd = fused_experts_forward if fused else packed_experts_forward
    exp.forward = fwd.__get__(exp, exp.__class__)
    return exp


@torch.no_grad()
def pack_all_experts(model, free_bf16=True, verbose=True):
    from transformers.models.diffusion_gemma.modeling_diffusion_gemma import DiffusionGemmaTextExperts
    n, gu_bits, dn_bits = 0, 0, 0
    for name, mod in model.named_modules():
        if isinstance(mod, DiffusionGemmaTextExperts):
            gu_bits += mod.gate_up_proj.numel(); dn_bits += mod.down_proj.numel()
            pack_experts_module(mod, free_bf16)
            n += 1
    torch.cuda.empty_cache()
    if verbose:
        print(f"[pack] {n} expert modules -> W4 (gate_up gs{GU_GS}, down gs{DN_GS}); "
              f"experts {2*(gu_bits+dn_bits)/2**30:.1f}GiB bf16 -> ~{(gu_bits+dn_bits)/2/2**30:.1f}GiB W4", flush=True)
    return n


def _killtest():
    from transformers import AutoConfig
    from transformers.models.diffusion_gemma.modeling_diffusion_gemma import DiffusionGemmaTextExperts
    cfg = AutoConfig.from_pretrained("google/diffusiongemma-26B-A4B-it").get_text_config()
    cfg.num_experts = 128
    torch.manual_seed(0)
    exp = DiffusionGemmaTextExperts(cfg).to("cuda", torch.float16)
    with torch.no_grad():
        exp.gate_up_proj.normal_(0, 0.02); exp.down_proj.normal_(0, 0.02)
    T, topk = 256, cfg.top_k_experts if hasattr(cfg, "top_k_experts") else 8
    H = cfg.hidden_size
    from w4a16 import quantize_w4, dequant_w4
    hs = torch.randn(T, H, device="cuda", dtype=torch.float16) * 0.1
    idx = torch.randint(0, 128, (T, topk), device="cuda")
    wts = torch.softmax(torch.randn(T, topk, device="cuda", dtype=torch.float16), -1)

    bf16_ref = exp.forward(hs, idx, wts).float()                          # stock BF16 (full precision)
    # dequant reference: stock forward but with W4-ROUNDED weights -> isolates kernel+plumbing from quant error
    gu_dq = torch.stack([dequant_w4(*quantize_w4(exp.gate_up_proj[e].float(), GU_GS), GU_GS) for e in range(128)])
    dn_dq = torch.stack([dequant_w4(*quantize_w4(exp.down_proj[e].float(), DN_GS), DN_GS) for e in range(128)])
    with torch.no_grad():
        exp.gate_up_proj.copy_(gu_dq); exp.down_proj.copy_(dn_dq)
    dq_ref = exp.forward(hs, idx, wts).float()                           # W4 values via torch linear

    pack_experts_module(exp, free_bf16=True)
    out = exp.forward(hs, idx, wts).float()                             # W4 values via Triton kernel

    scale = dq_ref.abs().max().item() + 1e-9
    rel_max = (out - dq_ref).abs().max().item() / scale
    rel_mean = (out - dq_ref).abs().mean().item() / (dq_ref.abs().mean().item() + 1e-9)
    rel_quant = (dq_ref - bf16_ref).abs().max().item() / (bf16_ref.abs().max().item() + 1e-9)
    ok = rel_mean < 5e-3 and rel_max < 3e-2          # 2-GEMM fp16 MoE chain; eval is the real judge
    print(f"  kernel+plumbing (packed vs dequant-ref): rel_mean={rel_mean:.2e} rel_max={rel_max:.2e} "
          f"{'OK — grouped Triton forward CORRECT (fp16-level)' if ok else '**FAIL**'}")
    print(f"  (info) W4-vs-BF16 quant error on RANDOM weights: rel={rel_quant:.2e}  "
          f"(real-weight quality is judged by GSM8K/HumanEval, not this)")


if __name__ == "__main__":
    _killtest()