dflash-ops

CUDA ops for DFlash block-diffusion drafting (arXiv:2602.06036), the speculative-decoding scheme Meta ships with Muse Glimmer 30B: a 5-layer drafter denoises a 16-token noise block in one forward pass against a sliding-window KV cache injected with fused target-model features, and the target verifies the proposed block in parallel. The ops cover everything in the drafter's block step except the GEMMs, which stay in cuBLAS: feature injection into the ring (RMS norm, NEOX rope, scatter, one launch per layer), the per-head norm+rope of the noise block, the non-causal sliding-window GQA attention of the block against ring plus itself, and the first-mismatch acceptance scan. A DFlashDrafter class assembles the full step and captures it into a CUDA graph.

The package also ships an eager fp32 reference (DFlashRef) transcribed from llama.cpp's llama_model_dflash graph, and a GGUF loader for the released drafter weights. The reference defines the semantics; the ops are certified against it.

Usage

import torch
from kernels import get_kernel

dfo = get_kernel("phanerozoic/dflash-ops", version=1, trust_remote_code=True)

cfg, w = dfo.load_drafter("dflash-kquant.gguf")   # 58 tensors, metadata-derived config
w["tok_embd"] = target_tok_embd                   # the drafter borrows both from
w["lm_head"] = target_lm_head                     # the target model

m = dfo.DFlashDrafter(cfg, w, dtype="bf16", ring_dtype="bf16")
m.inject(m.encode(target_features), positions)    # features: [T, 5 * n_embd]
step = m.capture_step()                           # CUDA-graphed block step
logits = step(id_last, n_past)                    # [16, n_vocab]
drafts = logits[1:].argmax(-1)                    # greedy block, positions 1..15
n_ok = m.accept(drafts, target_tokens)            # device-side, no host round trip

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Method

Injection follows the llama.cpp graph exactly: fused features go through wk/wv with no attn_norm in front, K is per-head RMS-normalized and roped, V is stored raw. The noise block attends non-causally, so later mask positions inform earlier ones and the denoising is joint; the sliding window masks only the deep past (0 <= qpos - kpos < window), and drafts are read greedily at block positions 1..15 with position 0 unused. Rope angles come from a host-computed fp64 inverse-frequency table passed to the kernels, so no device pow is involved and the table is identical on every device. Arithmetic is fp32 throughout regardless of storage dtype.

capture_step allocates nothing inside the capture region: every intermediate is a held static buffer written with out= and in-place ops, so the graph owns no memory the allocator could later hand to eager work running between replays, which is the normal traffic of speculative decoding (the target verifies between drafter blocks). The test suite replays the captured step interleaved with eager forwards and requires bitwise equality throughout.

Measured

RTX 6000 Ada (sm89), real drafter dimensions (5 layers, n_embd 6656, heads 32/8, head_dim 128, n_ff 19968, window 2048, block 16), ring filled to 2048, vocabulary reduced to 8192 for the benchmark (tok_embd and lm_head belong to the target model and touch none of the ops). Instrument: metakernel.

Per block step:

Phase decomposition at fp32, per step (5 layers):

The step is weight-bandwidth bound: measured GEMM streaming runs at 720-736 GB/s of the card's 771 GB/s read bandwidth, so the fused work is 18% of the fp32 step, 30% at bf16, and grows as the weights shrink; with a ~4.6-bit quantized drafter the GEMM floor is ~3.2 ms and the same fused savings project to a ~1.9x step.

Acceptance, measured with the released quantized drafter under llama.cpp against Muse Glimmer Q4_K_XL on the same device: 0.25-0.70 per block with a mean accepted run of ~2.8 from a 15-token draft; end-to-end decode reaches 1.44x over no drafter at a draft length of 4, which outperforms lengths 8 and 15 (54.9 vs 54.7 vs 51.8 tok/s) because acceptance saturates well below the block size.

Correctness

The eager reference is compared against structural invariants (ring wraparound, window masking proven by perturbing out-of-window slots, non-causal joint denoising, determinism), and the CUDA ops against the reference: logits to tight tolerance, greedy draft tokens exactly, bf16 paths against a reference rounded the same way, and the graphed step bitwise against the eager step before and after interleaved eager traffic. The suite runs on sm86 and sm89.

Requirements and limits

• NVIDIA GPU, compute capability 8.0+; torch with CUDA.
• head_dim must be a multiple of 32 (it is the launch width of the norm and rope kernels); the shipped drafter's is 128.
• The attention kernel stages one query head plus all scores in shared memory: (head_dim + window + block) * 4 bytes, so window <= ~12000 at the 48 KB static limit. The shipped drafter's window is 2048.
• At most window tokens per dflash_inject call (slots are pos % window; callers chunk, later chunks overwrite earlier ones).
• A 2048-slot ring costs 41.9 MB across 5 layers at fp32, 21.0 MB at bf16.
• tok_embd and lm_head are not in the drafter GGUF; they come from the target model.
• fp8 storage and a fused whole-block megakernel are not implemented.

References

Meta Superintelligence Lab, "DFlash: Block-Diffusion Drafting" (arXiv:2602.06036); Muse Glimmer 30B model card and released drafter (dflash-kquant.gguf); llama.cpp llama_model_dflash and common_speculative_impl_draft_dflash (the reference semantics).

License

Apache-2.0.