FlashKernel — Custom CUDA Kernels for Transformer Inference

Problem and scope

Transformer inference combines operations with different compute and memory costs. FlashKernel exposes those operations as small implementations that can be inspected, compared with reference code and profiled separately. It is an experimental kernel library; GPU performance and general deployment suitability remain to be validated.

Engineering questions

When can tiling avoid materializing a full attention matrix?
Which intermediate reads and writes can activation fusion remove?
What semantics must an optimized operation preserve when inserted into a real model?

Implementation inventory

Family	Mechanism	Source
Attention	Tiled computation and online softmax	CUDA implementations
Linear + GeLU	Apply the activation before writing the projection output	Triton implementations
RoPE	On-the-fly trigonometry or precomputed tables	Standalone positional-encoding kernels
Paged KV-cache	Append and gather through slot and block mappings	Standalone cache operations
Reduction	Warp-level reduction and intermediate partial sums	Reference comparisons in the test suite

An implementation in the source tree is distinct from a successfully validated GPU run. Inspect the kernel tests.

Technical decisions

Keep the kernel interface narrow and compare it against a reference before integrating it. The most consequential integration constraint is attention shape: cached decoding has a shorter query sequence than the accumulated keys and values. A kernel accepting equal sequence lengths cannot replace that path unchanged.

GPT-2 also uses learned absolute positional embeddings. RoPE and paged KV-cache experiments are separate from its attention/MLP patch; combining all library components in one diagram would misrepresent the implemented model.

Online softmax

Code

m_new = max(m_old, max(scores_tile))
alpha = exp(m_old - m_new)
p = exp(scores_tile - m_new)
l_new = alpha * l_old + sum(p)
acc_new = alpha * acc_old + p @ values_tile
output = acc_new / l_new

Algorithm sketch for an unnormalized output accumulator. A working kernel must additionally handle causal masking, bounds, synchronization and rows without valid keys.

Performance evidence

The repository contains a historical metrics summary and roofline plots. The reviewed revision does not include the underlying Nsight reports or per-run timing records needed to verify the headline throughput values. Those values are not presented here as established performance.

Correctness first: compare outputs with the PyTorch reference across causal masks, sequence lengths, head dimensions and precision. Check model logits and generated tokens before timing an integration.
Comparable timing: record the code revision, GPU, driver, CUDA/PyTorch/Triton versions, tensor shapes, warm-up, repetitions and synchronization. Report variation and which backend actually executed.
Traceable profiling: retain raw Nsight reports and their exports alongside derived plots. A plot generated from an existing JSON file is not a new measurement.
Separate conclusions: report individual kernel latency, prefill and cached decoding independently. A faster isolated operation does not establish a faster model.

GPT-2 integration boundaries

The integration module patches attention and the MLP projection/activation. Model correctness depends on preserving cache handling, masks, dropout, activation choice and the installed Transformers API. No end-to-end speedup is claimed here.

CPU regression tests can validate dispatch and reference behavior. They do not validate CUDA arithmetic, Triton compilation or NVIDIA T4 throughput. GPU parity checks and raw timing artifacts remain necessary.

Integration structure

flowchart TB Input["Tokens + learned absolute positions"] --> GPT["GPT-2 attention and MLP"] GPT --> Gate{"Supported inference inputs?"} Gate -->|Yes| Kernel["Custom attention or fused linear + GeLU"] Gate -->|No| Original["Original PyTorch implementation"] Kernel --> Check["Output correctness before timing"] Original --> Check Standalone["RoPE and paged KV-cache: separate kernel experiments"]

Validation update · 26 September 2026

Guarded GPT-2 integration. The integration repair preserves the reference path for caches, masks, training and unsupported inputs, and records which path actually executes. Offline checks on Python 3.11.14, PyTorch 2.8.0 and Transformers 4.44.2 passed 58 tests (14 skipped; 350 CUDA-marked cases excluded). These are CPU integration checks, not validation of native GPU arithmetic or speed.

Inspect the tested revision and validation details. This repair is available on a review branch; the repository’s main branch has not yet been updated.

Inspect and validate

Start with the pinned source and inspect the test and benchmark entry points before running them: