mjolnir-kernel

Mjolnir — a paged decode-attention kernel for NVIDIA Jetson AGX Thor (sm_110, 20 SMs), written in CuTe DSL on Blackwell tcgen05 tensor cores, and packaged as a vLLM attention backend.

vllm serve <model> --attention-backend MJOLNIR      # the mjolnir-vllm image
vllm serve <model> --attention-backend CUSTOM       # stock vLLM 0.31.0 + this package

This package is the kernel only. Serving, calibration and benchmarking are the mjolnir-cli package; benchmarks (Mjolnir vs FlashInfer vs FlashAttention 2 across context lengths), the design, and one chapter per kernel feature live in the project repository: thor-arena/mjolnir — numbers there are rendered from the raw results, so they are never stale.

What it does

The backend subclasses a stock vLLM backend — the fallback — and inherits its KV-cache allocation, cache write, prefill and validation:

• on the image, FlashAttention 4 (FLASH_ATTN, resolved to FA4 on sm_110 by the image's patch — its prefill is ~40% faster than FlashInfer's on Thor). If FA4's one-time fp8-KV probe fails on the device, the backend logs it and uses FlashInfer instead;
• on a stock vLLM, FlashInfer (stock FA on sm_110 is FA2, which cannot read an fp8 KV cache);
• MJOLNIR_ATTN_FALLBACK picks another either way.

Steps that are decode only run on Mjolnir:

• every request has the same q_len ≤ 1 + num_speculative_tokens (so MTP / n-gram speculative verification runs on Mjolnir too);
• CUDA-graph padding requests (zero tokens) may trail the batch;
• the step is causal — or the non-causal block of a parallel drafter (DSpark, DFlash: every query token sees every key) — and has no cascade prefix.

Every other step (prefill, mixed, cascade) is built and run by the base backend, unchanged.

A decode step:

1. Plan — an occupancy-pinned split-KV plan (plan_split_kv): exactly one wave of 2 CTAs × 20 SMs. The plan lives in persistent device buffers padded to a fixed grid, so full CUDA graphs replay new plans without re-capture.
2. Quantize Q — dynamic e4m3 per (sequence, KV head) by default: fp8 Q, fp8 KV and fp8 P go straight into the tensor cores, no conversion warps.
3. Stream KV — TMA loads pages into a shared-memory ring (L2 evict-first: KV is read once), tcgen05 MMAs accumulate S and O in TMEM, softmax warpgroups run the online softmax.
4. Merge — a small second kernel combines the split-KV partials by LSE.

The KV cache is read through runtime strides: vLLM's (B,H,N,2D), FlashInfer NHD/HND, or FlashAttention (2,B,…), uint8 or fp8, with per-tensor or per-head scales. Q and O are zero-copy for q_len == 1, even as strided views of the fused QKV buffer.

Install

On the image (narandill/mjolnir-vllm:v0.31.0) it is pre-installed and the MJOLNIR backend name is patched into vLLM — nothing to do.

On a stock vLLM 0.31.0 (inside a container that has CUDA 13 and sm_110):

pip install "mjolnir-kernel[dsl]"          # [dsl] pulls nvidia-cutlass-dsl
VLLM_KV_CACHE_LAYOUT=HND vllm serve Qwen/Qwen3-8B \
    --kv-cache-dtype fp8 --attention-backend CUSTOM

The package registers a vllm.general_plugins entry point (mjolnir:register) that maps AttentionBackendEnum.CUSTOM — and MJOLNIR, when the vLLM build has that member — to mjolnir.backend.MjolnirAttentionBackend in every vLLM process. The two-line enum patch for vLLM is thor-mjolnir-attention-backend.patch.

Look for this line in the server log:

Mjolnir decode enabled: h_q=… h_k=… d=… page=… q_len<=… Q=fp8 ctas/SM=2 …

Requirements

• An fp8 (fp8 / fp8_e4m3) KV cache — or, on the image, a 6-bit fp6_e2m3 one with calibrated scales (see below); head-major pages (VLLM_KV_CACHE_LAYOUT=HND) are fastest.
• head_size % 128 == 0; block size 16, 32, 64 or 128.
• A bf16 or fp16 model.
• No sliding window, soft-cap, sinks or ALiBi; no DCP or PCP.

If a layer does not qualify, the builder logs the reason and that KV-cache group runs on the fallback backend.

Configuration

The pre-rename THOR_T5_* spellings are still honored. Each default is backed by a measurement — see the kernel chapters.

fp6 (e2m3) KV in vLLM

On the image, --kv-cache-dtype fp6_e2m3 stores K/V as packed 6-bit e2m3 codes (25% fewer bytes than fp8) with static per-(layer, KV head, channel) scales — per channel because K has outlier channels a head-wide scale would crush the rest under. Mjolnir writes that cache (mjolnir/kvfp6.py), decodes from it directly (K's channel scales folded into Q, V's into the output), and expands the pages a prefill / mixed step reads to e4m3 for the FA4 fallback (channel scales applied, one descale per head). The scales come from a calibration run on fp8 KV:

python -m mjolnir.calibrate --config <serve.yaml> --out kv_scales.json   # in the image
MJOLNIR_ATTN_KV_SCALES=kv_scales.json vllm serve ... --kv-cache-dtype fp6_e2m3 --attention-backend MJOLNIR

(mjolnir calibrate / mjolnir serve do both for you — docs/faq/fp6-kv.md.)

Beyond fp8: the harness

mjolnir/t5_decode.py is the kernel and its standalone harness (python -m mjolnir.t5_decode --help) with a bit-faithful emulator as the correctness reference. The harness can also read KV formats vLLM cannot write yet — fp6 (e2m3/e3m2) and fp4 (e2m1) read directly by tcgen05 kind::f8f6f4, and nvfp4 — and defaults to e2m3 on head-major pages for fp8 Q. How they compare in speed and accuracy: kvfmt and quant. Wiring a vLLM cache-write path for fp6 is an open contribution.

Tests

From a checkout, on a Thor:

mjolnir verify correctness      # kernel vs the emulator, the full case matrix
mjolnir verify vllm             # runner tests + greedy MJOLNIR-vs-FLASHINFER smoke runs

or directly in the container: python tests/kernel/test_runner.py, python tests/kernel/smoke_vllm.py [--eager] [--spec 2] (they find the package installed, via $MJOLNIR_SRC, or by walking up to src/mjolnir).

Mjolnir rounds P (and, by default, Q) to e4m3, so greedy outputs can flip on near-tie tokens and then drift apart — FlashInfer's own greedy output differs between CUDA-graph and eager mode. The smoke test therefore compares per-step logprobs, not text: the chosen tokens must get the same logprob from both backends on the shared prefix, and every divergence must be a near-tie (each token in the other's top-5, at most 1 nat apart). The matching prefix is printed as information.

Files

Apache-2.0, except mjolnir/t5_decode.py, which is derived from an NVIDIA CUTLASS CuTe DSL example and stays under BSD-3-Clause — see LICENSE and NOTICE in this package.