File size: 3,016 Bytes
cc96dd8 ca49e5d cc96dd8 ca49e5d 4741dbb cc96dd8 22021b2 cc96dd8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | ---
license: apache-2.0
library_name: kernels
tags:
- kernels
- triton
- quantization
- moe
---
# dg_w4_kernels
Triton kernels from the [DiffusionGemma-26B-A4B W4A16](https://huggingface.co/GoedelMachines/diffusiongemma-26B-A4B-w4a16)
release. Pure Triton, no compiled code, so it runs anywhere Triton does.
```python
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:
```bash
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:
```python
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.
```python
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.
```python
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.
```python
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
```python
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.
|