susun-123/kern-kernels
kern-kernels Reproducible attention kernel recipes, ABI manifests, checksums and measured results. First profile: GB300 / Qwen3.8-27B / BF16 / Q24-KV4-D256 / page 64. Uses unmodified TRTLLM-GEN full attention, not MLA. The model's GDN layers are unchanged. Model weights and the base export's other kernels are not included. NVIDIA binaries are downloaded directly from pinned upstream URLs and verified by SHA256; this repository does not mirror them. The small Apache-2.0 vLLM KV… See the full description on the dataset page: https://huggingface.co/datasets/susun-123/kern-kernels.
kern-kernels
Reproducible attention kernel recipes, ABI manifests, checksums and measured results. First profile: GB300 / Qwen3.8-27B / BF16 / Q24-KV4-D256 / page 64. Uses unmodified TRTLLM-GEN full attention, not MLA. The model's GDN layers are unchanged. Model weights and the base export's other kernels are not included.
NVIDIA binaries are downloaded directly from pinned upstream URLs and verified by SHA256; this repository does not mirror them. The small Apache-2.0 vLLM KV append specialization is included, with its build recipe and license.
Compared with original kern Triton
Direct A/B of the original pinned Triton cubins against the integrated TRTLLM-GEN kernels, using identical logical Q/K/V and the same Rust CUDA Graph timing harness. Average KV length 200,000; one query token per sequence.
All seven A/B cases passed the numerical gate. Original Triton retains page 784 and its existing dispatch; TRTLLM-GEN uses page 64 and fixed 38 splits. Packing is excluded. Both use a two-layer interleaved physical KV fixture, not the full model's 16-layer stride. The original batch-one decode has a 16-way split/reduce path; its batched decode uses the unsplit unified kernel. That low parallelism explains the particularly large improvement at batch 2.
- Batch-one prefill, 2048 query / 2048 total KV tokens: 0.385 → 0.063 ms (6.1x).
- Batch-one prefill, 1536 query / 32768 total KV tokens: 8.255 → 0.633 ms (13.0x).
These are single attention-op measurements. The earlier short-prompt whole-model comparison was approximately 11.67 → 11.54 ms/step, about 1.1% faster. It does not measure long-context serving. See the direct comparison recipe.
FlashInfer cross-check
Single-layer decode, average KV length 200,000, randomized disjoint pages, CUDA Graph. For batch > 1, per-sequence lengths span 150,000–250,000.
The two measurements use different physical KV strides and graph timing harnesses. They establish comparable performance, not a controlled speedup. These are attention-op times, not complete model serving latency. A 200k × 32 single-layer test does not imply that all model weights and KV layers fit on one GPU.
All 14 integration cases pass, including page boundaries, nonzero layer offsets, ragged batches, chunked prefill, 8-token verify and repeated graph replay. Gates: finite output, relative L2 < 0.02, max absolute error < 0.02 against unmodified FlashInfer. See validation/attention.jsonl.
Five ordinary-generation prompts were checked for 128 tokens against a fresh same-configuration kern baseline. Matching prefixes were 66, 128, 128, 116, 128. The two first flips had a reference margin of 0.125 and a maximum logit change of 0.125 (one row-scale BF16 ULP). All 182 examined matched-prefix logits rows stayed within four row-scale ULP. This is numerical-tolerance validation, not a claim of bit-identical whole-model generation. The DFlash2 smoke test also passed: one prompt, 64 tokens matching the stored vLLM reference, 8-row CUDA Graph decode/verify. See validation/ for the evidence.
The upstream comparison covers FlashAttention-4, TRTLLM-GEN and CuTe DSL: 92 successful timings, 16 unsupported CuTe D256 prefill cases. FA4 D256 decode uses ragged KV with repacking excluded from timing. Read the benchmark methodology before comparing.
How to wire ops in JSON
The key is lowercase `ops`. Separate the pinned binary (modules), the kernel ABI (ops.<name>.impl.launches), and each invocation (programs.<program>.calls). The complete snippets below contain every packed field, without omitted bytes or placeholder hashes.
1. Add the module and op definition
For the shipped Qwen3.8-27B state layout, merge these top-level dictionaries into your manifest:
- decode-ops.json: one-token-per-sequence decode, fixed 38-way split-KV, including counter/statistics/partial-output scratch.
- prefill-ops.json: causal prefill and multi-token verification, including a pre-existing KV prefix.
These two fragments require the model profile: tokens.max = 2048, seqs.max = 128, maximum context 262144, page size 64, and 16 interleaved KV layers. They use in state for K/V. They are merge fragments, not complete manifests. For isolated buffer-based tests, use the complete smaller examples in step 4; their bounds and KV strides are intentionally different.
<details> <summary>Complete decode modules + ops JSON</summary>
{
"modules": {
"fmhaSm100fKernel_QkvBfloat16OBfloat16H256PagedKvCausalP64MultiCtasKvVarSeqQ8Kv128StaticSwapsAbForGen": {
"source": "fmhaSm100fKernel_QkvBfloat16OBfloat16H256PagedKvCausalP64MultiCtasKvVarSeqQ8Kv128StaticSwapsAbForGen.cubin",
"sha256": "a5c3300d05153675d2c6c432c50be888f3e0a8bb6c3ffbdaed48fbb25e52f86e"
}
},
"ops": {
"attn_decode": {
"params": ["out buffer<bf16>", "in buffer<bf16>", "in state", "in state", "in buffer<i32>", "in buffer<i32>", "in buffer<i32>", "i32", "i32"],
"impl": {
"scratch": {
"counter": {
"dtype": "i32",
"shape": [3072]
},
"stats": {
"dtype": "f32",
"shape": [128, 4, 38, 8, 2]
},
"partial": {
"dtype": "f32",
"shape": [128, 4, 38, 8, 256]
}
},
"launches": [
{
"module": "fmhaSm100fKernel_QkvBfloat16OBfloat16H256PagedKvCausalP64MultiCtasKvVarSeqQ8Kv128StaticSwapsAbForGen",
"entry": "fmhaSm100fKernel_QkvBfloat16OBfloat16H256PagedKvCausalP64MultiCtasKvVarSeqQ8Kv128StaticSwapsAbForGen",
"params": ["bytes<1280>"],
"block": [512, 1, 1],
"grid": [38, 4, "seqs"],
"shared_mem": 144992,
"args": [
{
"pack": {
"size": 1280,
"fields": [
{"at": 0, "tensormap": {"param": 1, "dtype": "bf16", "dims": [256, 6, 4, 2048], "strides": [512, 3072, 12288], "box": [64, 6, 1, 1], "swizzle": 128, "l2_promotion": 128}},
{"at": 128, "tensormap": {"param": 2, "dtype": "bf16", "dims": [256, 64, 4, 0], "strides": [4096, 1024, 4194304], "box": [64, 64, 1, 1], "swizzle": 128, "l2_promotion": 128}},
{"at": 384, "tensormap": {"param": 3, "dtype": "bf16", "dims": [256, 64, 4, 0], "strides": [4096, 1024, 4194304], "box": [64, 64, 1, 1], "swizzle": 128, "l2_promotion": 128}},
{"at": 512, "tensormap": {"param": 0, "dtype": "bf16", "dims": [256, 2048, 4, 6, 1], "strides": [12288, 512, 2048, 16], "box": [64, 8, 1, 1, 1], "swizzle": 128, "l2_promotion": 128}},
{"at": 912, "param": 0},
{"at": 1000, "param": 4},
{"at": 1080, "param": 5},
{"at": 1116, "param": 7},
{"at": 1244, "param": 8},
{"at": 1112, "i32": 2147483647},
{"at": 1144, "i32": 1},
{"at": 1148, "i32": 262144},
{"at": 1152, "i32": 1},
{"at": 1156, "i32": 38},
{"at": 1160, "i32": 4096},
{"at": 1164, "i32": 4},
{"at": 1168, "i32": 24},
{"at": 1172, "i32": 6},
{"at": 1176, "i32": 6},
{"at": 1180, "i32": -1431655765},
{"at": 1188, "i32": 2},
{"at": 1192, "i32": 6144},
{"at": 1200, "i32": 1048576},
{"at": 1204, "i32": 1},
{"at": 1208, "i32": 6},
{"at": 1212, "i32": 1},
{"at": 1216, "i32": 1065353216},
{"at": 1220, "i32": 1035512379},
{"at": 1224, "i32": 1065353216},
{"at": 1228, "i32": -1082130432},
{"at": 1252, "i32": 256},
{"at": 984, "scratch": "counter"},
{"at": 1008, "scratch": "partial"},
{"at": 1016, "scratch": "stats"}
]
}
}
]
}
]
}
}
}
}</details>
<details> <summary>Complete prefill modules + ops JSON</summary>
{
"modules": {
"fmhaSm103aKernel_QkvBfloat16OBfloat16H256PagedKvCausalP64VarSeqQ128Kv128PersistentContext": {
"source": "fmhaSm103aKernel_QkvBfloat16OBfloat16H256PagedKvCausalP64VarSeqQ128Kv128PersistentContext.cubin",
"sha256": "270e5246cd40181234f075ae4e05e0e53fbe662e46640807b44daf272cd5aed7"
}
},
"ops": {
"attn_prefill": {
"params": ["out buffer<bf16>", "in buffer<bf16>", "in state", "in state", "in buffer<i32>", "in buffer<i32>", "in buffer<i32>", "i32", "i32"],
"impl": {
"scratch": {},
"launches": [
{
"module": "fmhaSm103aKernel_QkvBfloat16OBfloat16H256PagedKvCausalP64VarSeqQ128Kv128PersistentContext",
"entry": "fmhaSm103aKernel_QkvBfloat16OBfloat16H256PagedKvCausalP64VarSeqQ128Kv128PersistentContext",
"params": ["bytes<1280>"],
"block": [384, 1, 1],
"grid": [
{
"ceil_div": ["tokens", 128]
},
24,
"seqs"
],
"shared_mem": 197872,
"args": [
{
"pack": {
"size": 1280,
"fields": [
{"at": 0, "tensormap": {"param": 1, "dtype": "bf16", "dims": [256, 1, 24, 2048], "strides": [512, 512, 12288], "box": [64, 1, 1, 128], "swizzle": 128, "l2_promotion": 128}},
{"at": 128, "tensormap": {"param": 2, "dtype": "bf16", "dims": [256, 64, 4, 0], "strides": [4096, 1024, 4194304], "box": [64, 64, 1, 1], "swizzle": 128, "l2_promotion": 128}},
{"at": 384, "tensormap": {"param": 3, "dtype": "bf16", "dims": [256, 64, 4, 0], "strides": [4096, 1024, 4194304], "box": [64, 64, 1, 1], "swizzle": 128, "l2_promotion": 128}},
{"at": 512, "tensormap": {"param": 0, "dtype": "bf16", "dims": [256, 2048, 4, 6, 1], "strides": [12288, 512, 2048, 16], "box": [64, 128, 1, 1, 1], "swizzle": 128, "l2_promotion": 128}},
{"at": 912, "param": 0},
{"at": 1000, "param": 4},
{"at": 1080, "param": 5},
{"at": 1116, "param": 7},
{"at": 1244, "param": 8},
{"at": 936, "param": 6},
{"at": 1144, "param": 8},
{"at": 1152, "expr": {"ceil_div": ["tokens", 128]}},
{"at": 1112, "i32": 2147483647},
{"at": 1148, "i32": 262144},
{"at": 1156, "i32": 1},
{"at": 1160, "i32": 4096},
{"at": 1164, "i32": 4},
{"at": 1168, "i32": 24},
{"at": 1172, "i32": 6},
{"at": 1176, "i32": 6},
{"at": 1180, "i32": -1431655765},
{"at": 1188, "i32": 2},
{"at": 1192, "i32": 6144},
{"at": 1200, "i32": 1048576},
{"at": 1204, "i32": 128},
{"at": 1208, "i32": 6},
{"at": 1212, "i32": 1},
{"at": 1216, "i32": 1065353216},
{"at": 1220, "i32": 1035512379},
{"at": 1224, "i32": 1065353216},
{"at": 1228, "i32": -1082130432},
{"at": 1252, "i32": 256}
]
}
}
]
}
]
}
}
}
}</details>
2. Bind the nine op parameters in a call
Example: decode at the second full-attention KV slot (zero-based slot 1):
{
"label": "full_attn_slot_1.decode",
"op": "attn_decode",
"args": [
{"buf": "attn_out"},
{"buf": "q_n"},
{"state": "kv", "offset": 262144},
{"state": "kv", "offset": 262656},
{"buf": "block_table"},
{"buf": "seq_lens"},
{"buf": "cu_seqlens_q"},
{"var": "seqs"},
{"var": "tokens"}
]
}Put this object in your decode program's calls array. All buffer/state names are references to declarations in that same manifest. For prefill or verify, use the same binding order and change "op" to "attn_prefill".
All offsets and TMA strides in these manifests are bytes. One layer's page occupies 64 * 4 * 2 * 256 * 2 = 262144 bytes. For full-attention slot L, K offset is L * 262144, and V offset is K + 512. The physical stride between pages is 16 * 262144 = 4194304 bytes. L is the KV slot among the 16 full attention layers, not the absolute transformer layer number.
3. Set live query/KV metadata
For two sequences with total KV lengths [385, 641]:
The KV pages must already contain the current query tokens' K/V (after KV append). Causal masking is bottom-right aligned: for query row j within a sequence, visible KV ends at kv_len - q_len + j. Verify computes attention for the candidates; acceptance/rejection is handled by the surrounding model program. tokens is not the context length.
The op receives one 1280-byte packed launch argument. A param inside that pack indexes the nine op parameters above. at is the byte offset in the pack. tensormap fields let kern create the TMA descriptors from the actual bound allocations; the zero outer KV dimension means span the bound storage. scratch names refer to this op's impl.scratch, allocated by kern rather than passed by the caller. Split counters start zero and are reused across replays. Do not share one scratch allocation across concurrent launches on independent streams.
4. Complete, small manifests you can run
- decode.json: batch 2, one query token each.
- prefill.json: batch 2, 129 query tokens each.
- verify.json: batch 2, 8 query tokens each.
These are complete schema-v4 manifests, with buffers, variables, modules, ops and calls. They use a two-layer KV buffer, not the model's opaque 16-layer state: [22 pages, 2 layers, 64 slots, 4 heads, 2 (K/V), 256]. Calls select layer 1. Each sequence has 11 page-table entries; maximum context is 704. Buffer-backed K/V parameters are in buffer<bf16> and the call uses {"buf": "kv", "offset": ...}. The shape is different from step 1, so use each complete file intact.
After the installation steps below, verify them without a GPU:
cargo run -q -p kern-manifest --example verify -- \
"$DATASET/examples/decode.json" \
"$DATASET/examples/prefill.json" \
"$DATASET/examples/verify.json"These are op-only fixtures for program_io, so the verifier prints verified followed by missing serving-protocol fills. That is expected; they are not standalone model-serving manifests for kern run.
For a ready-to-run correctness demo, this script creates the inputs and checks each exact JSON file against FlashInfer after repeated graph replay. All three pass on GB300; recorded results are included:
CUDA_VISIBLE_DEVICES=0 python "$DATASET/examples/run_examples.py" \
--runner target/release/examples/program_io --cubins kernels-trtllm \
--output results/example-validation.jsonWith raw contiguous BF16 Q/KV and int32 metadata prepared for the chosen file:
target/release/examples/program_io --gpu 0 \
--manifest "$DATASET/examples/decode.json" --cubins kernels-trtllm \
--env seqs=2 --env tokens=2 --graph --iters 9 \
--in q=q.bin --in kv=kv.bin --in table=table.bin \
--in seq=seq.bin --in cuq=cuq.bin --out out=out.binUse tokens=258 for prefill or tokens=16 for verify, with the corresponding JSON and query inputs. Read only the first tokens * 24 * 256 BF16 output values; the declared Q/O capacity may exceed the live row count.
5. Generate another supported capacity
Use the recipe to regenerate the packed constants, descriptors and scratch bounds together, instead of changing only the grid:
from trtllm_attention import op
# Run with PYTHONPATH pointing at the kern checkout's tools/ directory.
decode = op("decode", max_rows=2048, max_seqs=32,
max_context=262144, layers=16, splits=38)
prefill = op("prefill", max_rows=2048, max_seqs=32,
max_context=262144, layers=16)
# Add these to your generator's manifest["ops"], wire calls and buffers,
# then call kern_manifest.normalize(manifest) before writing JSON.The Python generator uses inline cubin/sha256 fields; normalize hoists them into the final modules table. Fetch pinned binaries with python tools/trtllm_attention.py --download kernels-trtllm. Only BF16 Q24/KV4/D256, page 64, and the documented SM103 profile are covered.
Download and install
from huggingface_hub import snapshot_download
path = snapshot_download("susun-123/kern-kernels", repo_type="dataset")
print(path)Pin revision to a dataset commit for reproducibility. SHA256SUMS covers the package files. Set DATASET to the downloaded directory, then:
git clone https://github.com/pegainfer-project/kern.git
cd kern
git checkout 5edac5990e0ebde208f2f6364dae488fcf3c7bad
git apply "$DATASET/patches/kern-integration.patch"
cp -R "$DATASET/tools/." tools/
# Standard kern CUDA/Rust build environment; see its README.
cargo build --release -p kern-run --bin kern --example program_io
mkdir -p kernels-trtllm
python tools/trtllm_attention.py --download kernels-trtllm
cp "$DATASET/artifacts/"* kernels-trtllm/
# Select an idle GPU (zero-based index). Requires PyTorch + FlashInfer and ninja.
CUDA_VISIBLE_DEVICES=0 python tools/test_trtllm_attention.py \
--cubins kernels-trtllm --work-dir results/trtllm-validation --quickDrop --quick for the 200k batch matrix; it needs a large-memory GPU and about 61 GiB of temporary fixture storage. The tests require the versions recorded in package.json; the runtime's NCCL/CUDA prerequisites still apply.
For a configured Qwen3.8 kern export, copy these new cubins into its existing kernel directory, then run with --manifest "$DATASET/manifests/qwen3.8-27b.json". The DFlash2 manifest is provided alongside it. Start with a fresh KV cache: the target page size changes from 784 to 64. See the complete recipe for conversion and rebuilding the KV append artifact. --splits controls a fixed decode profile; this package does not implement FlashInfer's adaptive dispatch policy.
Scope and licenses
Validated on SM103 GB300 only, with the pinned BF16 attention geometry. The prefill cubin is SM103-specific. No claim is made for other Blackwell variants, FP8, different head dimensions or MLA.
See NOTICE and LICENSE. Newly authored recipes, benchmark scripts and validation data use Apache-2.0. The vLLM cache kernel retains its Apache-2.0 attribution. Pre-existing kern code retains its original terms. NVIDIA artifacts remain subject to upstream terms; this dataset grants no rights to those external artifacts.