GoedelMachines/dg-w4-kernels

dg_w4_kernels

Triton kernels from the DiffusionGemma-26B-A4B W4A16 release. Pure Triton, no compiled code, so it runs anywhere Triton does.

import kernels
k = kernels.get_kernel("GoedelMachines/dg-w4-kernels")

Verified with kernels==0.9.1 on an RTX 5090 (sm_120). Newer kernels releases (0.15+) query a repo_type="kernel" Hub endpoint that currently returns 404 for every kernel repo, including the official kernels-community ones, so pin an older release until that settles:

pip install "kernels==0.9.1"

On 0.15+ the call also needs version=1, trust_remote_code=True. Both a v1 branch and main are published here.

Or skip the loader entirely and import from a clone:

import sys; sys.path.insert(0, "<repo>/build/torch-universal")
import dg_w4_kernels as k

Sampler

The one worth borrowing. fused_entropy computes Categorical(logits).entropy() in a single streaming pass instead of the five-kernel logsumexp/sub/exp/mul/sum chain. On a [256, 262144] fp32 tensor that is roughly 1.9 GiB of traffic down to 268 MiB, which is the bandwidth floor.

h = k.fused_entropy(logits)                        # [..., V] -> [...] fp32 nats
samp, amax = k.gumbel_argmax_sample(logits, seed)  # Categorical sample AND argmax, one pass

gumbel_argmax_sample replaces softmax plus multinomial plus argmax. Gumbel-max samples the same Categorical distribution exactly, noise comes from Philox inline so no 268 MiB noise tensor is materialised, and the plain argmax falls out of the same reduction for free. The RNG stream differs from torch.multinomial, so it is a different draw from the same distribution.

W4A16

Asymmetric uint4, fp16 scale and zero-point per group, two nibbles per byte. The GEMM reads packed nibbles and dequantizes inline, so it streams real 4-bit weight traffic.

qw, scale, zero = k.quantize_w4(W, group_size=128)
y = k.w4a16_linear(x, qw, scale, zero, BK=128)     # == F.linear(x, dequant(W))

Grouped MoE

One launch for all experts instead of a per-expert Python loop, with the activation folded into the first GEMM's epilogue. Token to expert alignment is sync-free and CUDA-graph safe. The weighted combine uses a fixed-order reduction rather than index_add_, so results are reproducible.

out = k.fused_moe_w4_v2(hidden, expert_module, topk_ids, topk_weights)

The expert module needs packed buffers gu_q/gu_s/gu_z and dn_q/dn_s/dn_z. See the model repo for how they are produced.

RMSNorm

y = k.fused_rmsnorm(x, weight, eps)   # x / sqrt(mean(x^2) + eps) * weight, one kernel

Notes

Tile configs ship for GB10 (sm_121) and RTX 5090 (sm_120). Other architectures fall back to the GB10 config, which is safe rather than tuned. Requires triton and torch>=2.5.

Apache-2.0.