baladhurgesh97/triton-gdn-decode

gdn-packed-decode

Triton packed Gated DeltaNet decode kernel. One CTA per head loops over V-tiles (BV=32, num_warps=4) instead of launching NV=4 CTAs per head. Tuned on NVIDIA GB10 for AxionML/Qwen3.5-0.8B-NVFP4.

Weights are unchanged. This is a kernel repo, not a model checkpoint.

Load from the Hub

from kernels import get_kernel

k = get_kernel("<this-repo>", trust_remote_code=True)
# k.fused_recurrent_gated_delta_rule_packed_decode(...)

trust_remote_code=True is required unless the publisher is on the kernels trusted list. Replace <this-repo> with this repository id.

Patch into vLLM

hf download <this-repo> --local-dir ./gdn-packed-decode
python3 ./gdn-packed-decode/apply_kernel.py apply
# restart vLLM so EngineCore reimports, then:
#   vllm serve AxionML/Qwen3.5-0.8B-NVFP4

GSM8K-10 (GB10, vLLM 0.27.1, frozen indices, thinking off)

Kernel	Acc	Median decode tok/s
Upstream vLLM	5/10	205.45
This kernel (run 1)	5/10	219.90
This kernel (run 2)	5/10	219.89

+7.0% e2e. Correctness vs upstream FLA: 6/6 (atol=rtol=2e-2).

Functions

fused_recurrent_gated_delta_rule_packed_decode

Same signature as vLLM FLA packed decode. Returns (out, initial_state).