Start with an inspectable operation
FlashKernel explores transformer operations through CUDA C++ and Triton implementations. The useful exercise is to connect data movement, numerical behavior and integration constraints. A kernel is only an optimization when it preserves the required semantics and improves a measured workload.
The source tree separates kernels, reference implementations and GPT-2 integration. The performance summaries previously shown here require stronger measurement provenance; this revision distinguishes the implementation from the evidence.
The hardware ceiling must match the work
NVIDIA’s T4 datasheet lists GDDR6 memory, not HBM2, and a mixed-precision Tensor Core peak of 65 TFLOPS. A scalar CUDA implementation cannot use that Tensor Core ceiling as evidence of its arithmetic utilization.
For a roofline plot, match precision, instruction path and traffic accounting to the measured kernel. Document the chosen bandwidth ceiling and raw counters. Arithmetic intensity and throughput must use the same operation count; a percent-of-peak label cannot substitute for those measurements.
Online softmax avoids a full score matrix
For each query tile, process key/value tiles while retaining a running maximum, normalization sum and output accumulator. Rescale earlier partial sums when a new maximum appears. This limits intermediate storage without changing the attention objective.
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_newThis sketch omits thread scheduling and masking. Tests must cover partial tiles, fully masked rows and unequal query/key lengths; threads participating in shared-memory loads must not exit before their data is available to the rest of the block.
Fuse a specific intermediate
The operation is GeLU(XWᵀ + b): a matrix projection followed by activation. Fusion can avoid writing and rereading the unactivated intermediate, but the benefit depends on tensor shape, data type, register pressure and the baseline library kernel.
The GeLU approximation is part of the model’s semantics. A tanh approximation and the exact error-function formulation are not interchangeable by assumption. Compare numerical errors and downstream logits before comparing latency. Fixed memory-saving or throughput claims need a named shape and actual run.
RoPE is a separate positional-encoding experiment
The library contains on-the-fly sine/cosine and precomputed-table variants. Their trade-off is arithmetic versus table traffic and reuse; a crossover point requires measurement on the target GPU.
GPT-2 uses learned absolute positions, as described in the Transformers documentation. Adding RoPE would change the model rather than simply accelerate its existing positional representation. These kernels are therefore not described as part of the GPT-2 patch.
Cache operations are not the whole decoding path
Paged storage maps logical token positions to physical blocks. Append and gather kernels test that mapping, but an end-to-end decoder must also preserve attention masks, past lengths and numerical behavior. The reviewed GPT-2 integration does not replace its cache with the standalone paged-cache implementation.
The current measurement record
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.
Preserve model semantics before dispatch
A safe integration needs an explicit supported fast path and an unchanged reference path for unsupported cases. Cached decoding, padding masks, head masks, requested attention weights, training/autograd and API differences all affect whether custom dispatch is valid.
A successful reference fallback establishes continued functionality, not accelerated execution. Benchmark output must identify actual custom calls and reject a “speedup” run in which no custom kernel executed. Generation comparisons also need identical output tokens and the actual number generated, including early stopping.
Keep the integration boundary visible
What would make this a complete performance case study
The next result should pair a tested source revision with GPU parity checks, raw timing distributions, profiler reports and explicit supported shapes. Separate prefill from decode and isolate the changed operation. Until those artifacts exist, the inspectable kernel and integration code is the contribution presented here; an end-to-end speedup remains unestablished.