whisper_blaze — Hopper-Native CUDA Kernels for Whisper

The Challenge

An H100 that Whisper barely touches

Whisper large-v3 is a 1.5-billion-parameter speech model, and the standard way to run it — Hugging Face Transformers on stock PyTorch — leaves most of an H100 idle. The gap isn't the model. It's the software between the model and the silicon.

The stock stack

Mel spectrograms computed on the CPU with librosa-style code — the GPU waits on the preprocessor before inference even starts
Generic PyTorch kernels written for every GPU since Pascal — no Hopper-specific instructions
Attention materializes full score matrices in HBM — memory traffic grows with the square of sequence length
Everything runs in FP16 while the H100's FP8 tensor-core path — double the math throughput — sits unused

What Hopper silicon actually offers

WGMMA — asynchronous matrix-multiply issued by a whole 128-thread warpgroup at once
TMA — a dedicated copy engine that moves whole tensor tiles into shared memory while compute proceeds
FP8 — E4M3 and E5M2 formats at twice FP16 tensor throughput and half the memory traffic per element
mbarrier sync, thread-block clusters, 228 KB of shared memory per SM

Why it's hard: none of these fast paths are reachable from ordinary CUDA C++. WGMMA and TMA are PTX instructions (wgmma.mma_async, cp.async.bulk.tensor) with an asynchronous producer/consumer programming model that is nothing like thread-per-element CUDA. And FP8 gives you 3–4 bits of mantissa — where you quantize matters as much as how.

The Approach

Rebuild the hot path, keep the model

No fork of Whisper, no retraining — the original large-v3 weights load straight from the Hugging Face Hub. Every performance-critical operation underneath them is replaced with a hand-written Hopper kernel.

Kernels, not a fork bindings.cpp

The kernels register as PyTorch custom ops, so WhisperBlaze.from_pretrained() looks and behaves like the Transformers API it replaces. Adoption cost: one import.

PTX where it counts wgmma_utils.cuh

A thin header library wraps the raw PTX — WGMMA in every tile shape, TMA bulk copies, mbarrier arrive/wait, cluster launch — so the five kernels read like CUDA, not assembly.

Precision is a per-layer decision precision.py

Instead of quantizing the whole model, a PrecisionConfig assigns weight and activation precision to each of six layer sites independently — with presets full_fp16, mixed_fp8 and aggressive_fp8.

Kill the CPU preprocessor mel_spectrogram.cu

The mel spectrogram — Whisper's entire audio front-end — moves onto the GPU via cuFFT, with overlapping-chunk support for long audio. The H100 never waits on librosa again.

Inference pipeline under the mixed_fp8 preset — every colored block is a custom kernel; precision assigned per layer site.

Layer site	Weights	Activations	Why
encoder_ffn	fp8 e4m3	fp8 e4m3	Biggest matmuls — safest FP8 win
encoder_attn_qkv	fp8 e4m3	fp8 e4m3	Large projection, tolerant inputs
encoder_attn_out	fp16	fp16	Post-softmax activations vary widely
decoder_self_attn	fp16	fp16	Causal attention is precision-sensitive
decoder_cross_attn	fp16	fp16	Cross-attn scores span a large range
decoder_ffn	fp8 e4m3	fp8 e4m3	Big matmuls, well-behaved values

The mixed_fp8() preset, verbatim from precision.py. The aggressive_fp8() preset pushes FP8 everywhere and switches cross-attention activations to E5M2 for its wider dynamic range.

The Solution

Five kernels and a clean Python surface

Roughly 2,500 lines of CUDA C++ and 800 lines of Python: five kernels, three PTX header libraries, and an API that hides all of it behind from_pretrained().

gemm_hopper.cu

The workhorse GEMM: TMA streams tiles into shared memory while warpgroups issue asynchronous WGMMA — in FP16 and both FP8 formats.

wgmma.mma_asynccp.async.bulk.tensormbarrier

flash_attn_hopper.cu

Flash Attention 3 for all three attention patterns Whisper needs — encoder self, decoder causal self, and decoder cross-attention — without ever writing score matrices to HBM.

online softmaxsmem tilingcausal mask

mel_spectrogram.cu

Whisper's whole audio front-end on the GPU: cuFFT-based STFT, 128-bin mel filterbank, fp16 output, overlapping-chunk mode for long recordings.

cuFFTfp16 outchunked audio

layernorm_fused.cu

Residual add and LayerNorm (or RMSNorm) fused into a single pass — one kernel launch and one trip through memory instead of three.

fused residualrmsnormwelford

fp8_quantize.cu

FP16 ↔ FP8 casts with per-tensor scaling for both E4M3 (more precision) and E5M2 (more range) — the numerical plumbing every FP8 layer depends on.

e4m3 · e5m2per-tensor scaleamax

+ 3 header libraries

wgmma_utils.cuh, hopper_utils.cuh and compute_traits.cuh wrap the PTX instruction set and type traits so kernels stay readable.

ptx wrapperstype traitscluster helpers

Using it three lines to FP8 Whisper

from whisper_blaze import WhisperBlaze
from whisper_blaze.precision import mixed_fp8

model = WhisperBlaze.from_pretrained(
    "openai/whisper-large-v3", precision=mixed_fp8(),
)
result = model.transcribe(audio, language="en")

# v0.1.7: batched multi-audio inference
results = model.transcribe_batch([a1, a2, a3])

Or go straight to the kernels direct kernel API

import whisper_blaze_kernels as k

fp8, scale = k.quantise_e4m3(x)          # fp8 + per-tensor scale
out = k.layernorm_fused(h, res, g, b, eps)
out = k.decoder_self_attn(Q, K, V)       # causal FA3
mel = k.mel_spectrogram(audio)           # [1, 128, T] fp16

Impact

What changed

Shipped as a real package

pip install whisper-blaze — live on PyPI at v0.1.7. No framework fork, no patched Transformers: the original weights, faster underneath.

The CPU left the pipeline

Audio preprocessing — previously a serial CPU stage in front of every inference — now runs on the same GPU as the model, so transcription starts when audio arrives.

FP8 where it's safe, FP16 where it matters

On FFN and QKV matmuls, FP8 halves memory traffic per element and engages the H100's double-rate tensor-core path — while attention keeps full FP16 precision.

Attention without the quadratic write

Flash Attention 3 keeps score matrices in shared memory, removing the O(T²) round-trips to HBM for all three of Whisper's attention patterns.

Benchmark harnesses (bench_gemm.py vs torch.matmul, bench_mel.py vs CPU librosa) ship with the library; measured throughput figures will be published here once formal runs are recorded.

Learnings

What writing Hopper kernels taught me

Precision is a per-layer property, not a model setting. Cross-attention activations span a dynamic range that E4M3 clips, and post-softmax encoder outputs are just as unforgiving — the working design quantizes six layer sites independently instead of flipping one global switch.
Hopper is a different programming model, not a faster GPU. WGMMA and TMA turn a kernel into an asynchronous producer/consumer pipeline — copy engines filling shared memory while warpgroups compute — and thread-per-element intuition from older CUDA actively misleads.
The preprocessor can be the bottleneck. Before any kernel tuning mattered, the CPU mel spectrogram in front of the model was the serial stage gating everything. Moving the front-end onto the GPU came first.
Shipping a CUDA extension is a project of its own. PyTorch's CUDA version must match the system toolkit, sm_90a gates the whole build behind CUDA 12.2+, and pip needs --no-build-isolation to see the user's torch — packaging consumed real engineering, not an afternoon.
FP8 lives or dies on scaling. E4M3 saturates at 448, so per-tensor scale factors computed from activation maxima are not an optimization — they are the difference between a working model and garbage tokens.