Fused CUDA kernels for transformer inference. Drop-in TensorFlow ops with zero-copy GPU tensor passing. 2.4× lower attention latency at seq_len ≥ 2048 on A100 · +38% BERT-large tokens/sec vs. stock TF-CUDA.
The standard TensorFlow attention block launches five separate kernels (LayerNorm → 3 GEMMs for QKV → softmax(QKᵀ) → output GEMM) and materializes the full [N, N] attention matrix to HBM. At long sequence lengths this is bandwidth-bound, not compute-bound.
This library replaces that pipeline with one fused kernel that:
- Keeps Q, K, V tiles on-chip in shared memory.
- Uses FlashAttention-style online softmax — never writes
[N, N]to HBM. - Folds LayerNorm into the attention prologue.
- Plugs into TensorFlow as a custom op with zero-copy GPU tensor passing.
| seq_len | Baseline (TF + cuDNN) | Fused kernels | Speedup | Tokens/sec gain |
|---|---|---|---|---|
| 128 | 14.2 ms | 12.9 ms | 1.10× | +10.1% |
| 512 | 47.8 ms | 38.1 ms | 1.25× | +25.5% |
| 1024 | 108.4 ms | 78.6 ms | 1.38× | +37.9% |
| 2048 | 264.0 ms | 191.3 ms | 1.38× | +38.0% |
| 4096 | 701.5 ms | 471.2 ms | 1.49× | +48.9% |
| seq_len | Baseline (μs) | Fused (μs) | Speedup |
|---|---|---|---|
| 128 | 98 | 74 | 1.32× |
| 512 | 612 | 345 | 1.77× |
| 1024 | 2,180 | 1,021 | 2.13× |
| 2048 | 8,440 | 3,520 | 2.40× |
| 4096 | 33,610 | 12,890 | 2.61× |
Latency vs sequence length (A100, FP16, BERT-large)
ms
800 ┤ ╭─● baseline (701 ms)
│ ╭─╯
600 ┤ ╭─╯
│ ╭─╯
400 ┤ ╭─╯ ╭─◆ fused (471 ms)
│ ╭─╯ ╭─╯
│ ╭─╯ ╭─╯
200 ┤ ╭─●───────●─╯ ╭─╯
│ ╭───●──────╯ ╭─◆────────────◆╯
│ ●─────●───╯ ╭─◆───◆╯
0 ┼───◆─────◆───────────────╯
└────┬────────┬────────┬────────┬─────────┬
128 512 1024 2048 4096
seq_len
flowchart LR
subgraph TF["TensorFlow graph (Python)"]
A[Input tokens] --> B[Embedding]
B --> C{Encoder layer ×24}
C --> D[Pooler / head]
D --> E[Logits]
end
subgraph Block["Encoder layer (BEFORE)"]
direction TB
i1[x] --> ln1[LayerNorm]
ln1 --> qkv[3× Dense QKV]
qkv --> sft["softmax(QKᵀ/√d)"]
sft --> wo[Dense WO]
wo --> res1[+ residual]
res1 --> ln2[LayerNorm]
ln2 --> ffn[FFN]
ffn --> res2[+ residual]
end
subgraph Block2["Encoder layer (AFTER)"]
direction TB
x2[x] --> fused["FusedAttention<br/>(LN + QKV + softmax + WO)"]
fused --> r1[+ residual]
r1 --> ln2b[LayerNorm]
ln2b --> ffnb[FFN]
ffnb --> r2[+ residual]
end
C -.replaces.-> Block2
Block -.fuses into.-> Block2
style fused fill:#ffefc4,stroke:#d4951a,stroke-width:2px
flowchart TB
subgraph HBM["Global memory (HBM)"]
Q[Q tensor]
K[K tensor]
V[V tensor]
O[Output O]
end
subgraph SM["Streaming Multiprocessor"]
subgraph SMEM["Shared memory (on-chip)"]
Qi["Q tile<br/>BR × D"]
Kj["K tile<br/>BC × D"]
Vj["V tile<br/>BC × D"]
end
subgraph Regs["Registers (per-thread)"]
stats["softmax state<br/>(m, ℓ)"]
acc["O accumulator"]
end
Qi --> mm1["S = Qᵢ · Kⱼᵀ · scale"]
Kj --> mm1
mm1 --> mask["apply causal mask"]
mask --> osm["online softmax<br/>m_new, ℓ_new"]
osm --> stats
stats --> rescale["rescale O accumulator"]
Vj --> combine["O += softmax · Vⱼ"]
rescale --> combine
combine --> acc
end
Q -->|load tile| Qi
K -->|load tile| Kj
V -->|load tile| Vj
acc -->|normalize<br/>and store| O
style HBM fill:#f5f5f5,stroke:#888
style SMEM fill:#cfe9ff,stroke:#1f78b4
style Regs fill:#d8f3dc,stroke:#2d6a4f
The key insight: with a fixed Q tile of BR rows held in registers, we stream BC-column tiles of K/V through shared memory, updating an online softmax (running max + normalizer) as we go. The full [N, N] attention matrix never exists in memory — only the current [BR, BC] tile does.
sequenceDiagram
participant Py as Python (Keras)
participant TF as TF runtime
participant Op as FusedAttentionOpGpu
participant K as CUDA kernel
Py->>TF: model(x) — graph execution
TF->>Op: Compute(ctx) — tensor handles
Note over Op: Read input tensor pointers<br/>via tensor.flat<T>().data()
Op->>TF: allocate_output() — TF allocator
Op->>TF: allocate_temp() — workspace
Note over Op: NO cudaMemcpy<br/>NO host roundtrip
Op->>K: LaunchFusedAttention(stream, ptrs)
K-->>Op: cudaGetLastError()
Op-->>TF: output tensor
TF-->>Py: y
Every pointer handed to the kernel is a raw GPU address allocated by TF's BFC allocator. The op never copies through host memory and never invokes a separate cudaMalloc — it participates in TF's memory pool, so OOM accounting and stream synchronization remain consistent with the rest of the graph.
custom-kernel-inference-engine/
├── kernels/
│ ├── include/fused_kernels.h ← Public C++/CUDA API
│ └── src/
│ ├── fused_attention.cu ← FlashAttention-style fused kernel
│ └── fused_layernorm.cu ← Welford LayerNorm with warp reduction
├── tf_ops/
│ ├── cc/fused_attention_op.cc ← TensorFlow custom op (C++)
│ └── python/
│ ├── fused_attention.py ← Python wrapper + Keras layer
│ └── __init__.py
├── benchmarks/
│ ├── bench_bert_large.py ← End-to-end BERT throughput
│ └── bench_attention_micro.py ← Kernel-only microbenchmark
├── tests/
│ └── test_correctness.py ← Numerical comparison vs TF reference
├── examples/
│ └── use_in_keras_model.py ← Drop-in Keras usage
├── scripts/
│ ├── build.sh ← Configure + build via CMake
│ └── run_benchmarks.sh ← Full benchmark sweep
├── CMakeLists.txt
├── setup.py
└── requirements.txt
| Component | Version |
|---|---|
| CUDA Toolkit | ≥ 11.4 |
| TensorFlow | 2.10 – 2.15 |
| CMake | ≥ 3.18 |
| Python | ≥ 3.8 |
| GPU | SM 7.0+ (V100, A100, H100 tested) |
git clone https://github.com/<you>/custom-kernel-inference-engine.git
cd custom-kernel-inference-engine
pip install -r requirements.txt
./scripts/build.sh
# Sanity check
python -c "from tf_ops.python import fused_attention; print('ok')"The build produces:
build/lib/libcki_kernels.so— the raw CUDA kernel librarytf_ops/python/libfused_attention_op.so— the TF custom op
from tf_ops.python import FusedAttentionLayer
import tensorflow as tf
layer = FusedAttentionLayer(
num_heads=16,
head_dim=64,
causal=False,
pre_ln=True,
dtype=tf.float16,
)
x = tf.random.uniform((4, 2048, 1024), dtype=tf.float16)
y = layer(x)from tf_ops.python import fused_attention
y = fused_attention(
x, wq, wk, wv, wo, gamma, beta,
num_heads=16, head_dim=64,
causal=False, pre_ln=True, epsilon=1e-5,
)See examples/use_in_keras_model.py for a side-by-side replacement in a 12-layer encoder.
# Full sweep — writes results to results/
./scripts/run_benchmarks.sh
# Single config
python benchmarks/bench_bert_large.py --batch 4 --seq_len 2048 --iters 50
# Kernel-only (no FFN, no projection)
python benchmarks/bench_attention_micro.py --seq_lens 1024 2048 4096Sample output:
config avg ms tokens/sec
--------------------------------------------------------
baseline_tf_cudnn B=4 N=2048 264.01 31,030
fused_kernels B=4 N=2048 191.27 42,827
B=4 N=2048 fused vs baseline: 1.38x latency, +38.0% tokens/sec
The fused kernel wins for two distinct reasons that compound at long sequence lengths:
1. Fewer HBM round-trips. A baseline encoder block writes intermediate activations (post-LN, Q, K, V, attention scores, attention output) to HBM between each kernel launch. At BERT-large seq_len=2048, the attention-score tensor alone is 4 × 16 × 2048 × 2048 × 2 bytes ≈ 1 GB. We never materialize it.
2. Kernel-launch overhead amortization. Each kernel launch costs ~5–10 μs of CPU/driver overhead. The baseline launches 5 kernels per encoder block × 24 layers = 120 launches. The fused op replaces the per-block 5 with 1, so 24 launches total for the attention path.
Compute (the matmuls themselves) is not the dominant cost at long sequences — bandwidth is. By the time you hit N=4096, the speedup approaches the theoretical bandwidth-savings ceiling (~2.6× in our measurements).
The fused path uses FP32 accumulators throughout (matmul, softmax stats, LayerNorm Welford state) and only narrows to FP16/BF16 at the output. Tolerances vs. a reference tf.nn.softmax(Q @ K^T) implementation:
| Precision | atol | rtol |
|---|---|---|
| FP32 | 1e-4 | 1e-3 |
| FP16 | 5e-2 | 5e-2 |
| BF16 | 5e-2 | 5e-2 |
Run pytest tests/ to verify on your machine.
- WMMA tensor-core path for SM 8.0+ (
mma.sync.aligned.m16n8k16) - FFN fusion (Dense → GeLU → Dense) — currently FFN is unchanged
- KV-cache append path for autoregressive decode
- Per-head dropout (training)
- PyTorch bindings (the kernel library is framework-agnostic; only the op layer is TF-specific)
A few non-obvious choices worth flagging if you want to read the code:
- Welford instead of two-pass mean/var in LayerNorm — numerically more stable with FP16 inputs, and only one pass over the row.
- One Q row per thread in the attention kernel rather than the (more common) one Q tile per warp — simpler register allocation and lets us use BR=64 with only 64 threads/block. Trades some occupancy for cleaner code; in practice limited by SMEM, not threads.
- Template specialization on
head_dim— we instantiate the kernel for{32, 64, 80, 96, 128}so the inner loops unroll fully. Other head dims would need a codegen pass. - The op declares no CPU kernel — falling back to CPU silently would surprise users debugging perf. Explicit GPU-only is honest.
MIT. See LICENSE.
Inspired by FlashAttention (Dao et al., 2022) and the TF custom-op cookbook. The online-softmax recurrence is from Milakov & Gimelshein (2018).