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")

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

Tuned on GB10 (sm_121). An sm_90 config is included for H100 and H200. Other architectures fall back to a safe default and will be slower. Requires triton and torch>=2.5.

Apache-2.0.