srswti/axe-diffusion-ksl-xe

knivesysl-diffusion-xe

a block-diffusion language model served on a single gpu — one process, one model, two apis: a browser studio with live revisable drafts and an openai-compatible endpoint for routers and cockpits.

the weights are 4-bit quantized, and the fast kernels that run them are ours: we wrote and fused the kernels ourselves instead of relying on generic inference-stack implementations.

what the model is

this is a block-diffusion language model, not autoregressive. instead of emitting one token at a time, it denoises a 256-token canvas over up to 48 steps, committing tokens once their entropy drops below a bound. 26b total parameters with 4b active: a 128-expert moe routing top-8 per token, across 30 encoder and 30 decoder layers.

because generation is denoising rather than decoding, quantization and throughput are coupled: quantization noise feeds the entropy bound that decides how many tokens commit per step. the kernels we wrote target exactly that loop.

the kernels we built

the expert weights are asymmetric uint4 with an fp16 scale and zero-point per group. the gemm reads packed nibbles and dequantizes inline, streaming real 4-bit traffic — experts drop from 85 gib to ~11.5 gib, and since the expert path is bandwidth-bound that accounts for most of the win.

on top of that:

• one fused triton launch handles all 128 experts instead of a 128-launch python loop, with the activation folded into the first gemm's epilogue
• rmsnorm collapses from roughly six elementwise launches into one, 180 times per forward
• the weighted combine uses a fixed-order reduction instead of index_add_, because in block diffusion a 1-ulp difference from atomics flips an acceptance threshold and forks the entire trajectory
• the sampler computes entropy once per denoising step in a single streaming pass and replaces softmax + multinomial with one gumbel-max kernel — on the stock loop the same 268 mib logits tensor was processed twice per step, and fixing that is worth ~1.4x by itself

requirements

~18 gib vram for weights, ~20 gib peak (the fp32 logits tensor alone is 268 mib per denoising step). checkpoint format dg-w4a16-v1: asymmetric uint4, fp16 scale and zero-point per group, two nibbles per byte. gate_up_proj groups at 128 along the contraction dim, down_proj at 64 (its contraction dim is 704 and 128 does not divide it). the vision tower rides at bf16 and multimodal paths are untested.

the first forward jits the triton kernels — expect a one-time warmup on a freshly started server.