CUDA Kernel Engineering

This is the chapter you read when you're about to write a kernel. It covers the conventions every Metrale Engine kernel follows, the tools that matter on GB10 SM121, and the workflow that takes an idea from a profile to a merged PR.

The kernel inventory

A GB10 target compiles the shared baseline in kernels/gb10/common/ plus its own leaf's overrides. Canonical roles:

RoleFile exampleWhat it does
Prefill attentionattn_prefill_v47.cuFlash Attention v2, cp.async pipelining, mma.sync.aligned.m16n8k16 tensor cores, 2 CTAs/SM
Decode attentionpaged_decode_attn_turbo3_128.cuOnline softmax, split-K, adaptive split count
Prefill attention (FP8 KV)attn_prefill_fp8kv.cuSame but with FP8 KV read path
KV appendreshape_and_cache.cu, fused_k_norm_rope_cache.cuPer-token K/V write into paged cache
MoE prefillmoe_prefill.cuFused dequant + grouped GEMM, 256 experts, topk=10
MoE decode — shared expertmoe_shared_expert_fused_fp8.cuShared-expert path with fused FP8 GEMM
MoE decode — expertmoe_expert_relu2_down_shared.cuToken-level MoE for decode
Dense GEMMdense_gemm_bf16.cu, w8a16_gemv.cuNon-MoE FFN GEMMs
SSM — preprocessssm_preprocess.cuFused QKVZ deinterleaving + GDN gate (softplus + sigmoid)
SSM — Gated Delta Rulegated_delta_rule*.cuGated delta-net SSM (prefill + decode variants)
SSM — causal conv1dcausal_conv1d.cuMamba's 1D convolution
Primitive — RMSNormrms_norm.cuSingle-block tree reduction
Primitive — SiLU×Mulmoe_silu_mul.cuSwiGLU for the MoE path
Primitive — RoPErope.cuGQA-aware rotary embedding
Primitive — argmax BF16argmax_bf16.cuSingle-block tree reduction, 4-byte result
E2M1 conversione2m1_branchless.cuSoftware FP32 → E2M1 conversion for SM121
MoE gatingmoe_gate_topk.cuExpert selection pre-dispatch
WHTwht_bf16.cuWalsh-Hadamard for TurboQuant KV
Element-wisebf16_add.cu, fp8_scale_transpose.cuSmall utilities

A model whose kernel differs from the baseline keeps its own copy under the same file name in its (model, quant) leaf, declared in that leaf's KERNEL.toml [shadow] table with a reason.

SM121 hardware budget (quick reference)

Grace-Blackwell GB10 / SM121 numbers you will care about when writing kernels:

QuantityValue
Global memory119.7 GB LPDDR5X (unified)
Peak memory BW273 GB/s
SMs32 (typical, configurable)
Warps per SM48 max concurrent
Registers per SM65,536
Shared memory per SM100 KB effective
L2 cache~ 32 MB (large enough that benchmark reports up to 599 GB/s achieved BW at small sizes — the "L2 cache effect" in the kernel tables)
Tensor core throughput (BF16)high, SM121-specific
cp.asyncsupported
Native FP4 MMAnot available on SM121 — see below

The native FP4 MMA caveat is load-bearing: SM120/SM121 does not expose the cvt.rn.satfinite.e2m1x2.f32 instruction or native FP4 tensor-core paths. Every NVFP4 kernel on GB10 uses software E2M1 conversion (the "branchless" kernel) and dequantises-to-BF16 for the MMA. This is not a performance bug — it is the silicon. The community benchmarks that cite "native FP4 throughput" on newer Blackwell parts do not transfer. The NVFP4 deep dive walks the workaround in detail.

Conventions every Metrale Engine kernel follows

  • SPDX header line 1. // SPDX-License-Identifier: MIT OR Apache-2.0. Enforced by the license-headers job in CI.
  • extern "C" __global__ entry points with a stable name. The name is what GpuBackend::kernel(module, func) looks up.
  • All pointer args are typed at the right level — const __nv_bfloat16*, const int8_t*, not const void*. The BF16 + E2M1 types come from <cuda_bf16.h> and <cuda_fp8.h>; module-local aliases are fine but don't hide the precision.
  • Grid/block dimensions are passed from the Rust side. Never compute block dims from runtime GPU properties inside the kernel — the launch site owns them.
  • One kernel per file where possible. Fused variants belong in their own files (e.g. moe_shared_expert_fused_fp8.cu vs moe_shared_expert_fused.cu) so the KERNEL.toml [modules] override maps cleanly.
  • No <iostream>, no printf in hot paths. Use #if 0 stubs during development, strip before merging. nvcc warns on printf inside __device__ code when it bloats the PTX.
  • Shared-memory layouts are always annotated. A comment near each __shared__ alias documents the row/column order and any padding added to avoid bank conflicts.

The profiling workflow

  1. Start with nsys profile against the live server running a benchmark. The first question is always which kernel is the bottleneck — do not tune in the abstract. 
    nsys profile --trace=cuda,cudnn,cublas,osrt -o metrale.qdrep \
  /path/to/met serve <model>
# in another terminal: drive bench load
nsys stats --report cuda_gpu_kern_sum metrale.qdrep | head -20
  2. For the top 2–3 kernels, drill into ncu (Nsight Compute). The metrics that matter on GB10:
    • smsp__cycles_active.avg.pct_of_peak_sustained_elapsed — SM utilisation.
    • l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum — shared-memory bank conflicts.
    • dram__bytes.sum vs dram__bytes_read.sum.peak_sustained — how close to the 273 GB/s ceiling.
    • sm__warps_active.avg.pct_of_peak_sustained_elapsed — warp occupancy.
  3. Know which side of the roofline you're on. GB10's compute-vs-BW roofline is steeper than a desktop GPU's (273 GB/s vs several thousand tensor-core TFLOPs). Most decode kernels are memory-bound; most prefill kernels are tensor-core-bound once cp.async is hiding the K/V load.

The three performance levers (in order of payoff)

  1. Tiling + cp.async. The biggest single win on SM121 is pipelining the next tile's global-memory load behind the current tile's compute. __pipeline_commit()/__pipeline_wait_prior(N) with two or three in-flight stages typically doubles attention throughput vs a naive loop. Prefill v47 uses 2 stages and two CTAs per SM.
  2. Shared-memory layout. Bank conflicts destroy kernels silently. The usual fix is an xor-swizzle or an extra padding column. The ncu metric above tells you how far you are from zero conflicts.
  3. Register budget. Hit the 255-register-per-thread cliff and you spill to local memory, which on GB10 means LPDDR. __launch_bounds__(256, 2) (max 256 threads, 2 blocks per SM) is the decoration that constrains nvcc's register allocator. Use it. Check with --ptxas-options=-v.

Things that matter less on GB10 than on an H100:

  • Shared-memory capacity (100 KB is generous for these shapes).
  • Warp specialisation. SM121's scheduler is good enough that explicit producer/consumer warp roles rarely pay back on kernels of the shapes Metrale Engine runs.
  • Distributed shared memory. No NVLS / multi-CTA clusters to lean on — NVLS_ENABLE=0 is forced in the NCCL env.

CUDA graphs

Every supported batch size gets a captured graph at startup. The engine replays graphs with cuGraphLaunch, eliminating per-launch cuLaunchKernel overhead (~microseconds each, compounded over ~300 kernels per forward pass). Graph stability requires:

  • Buffer addresses do not change. BufferArena pre-allocates at startup; the engine reuses the same DevicePtrs for every step. No alloc/free in the hot loop.
  • Kernel launch parameters are data-dependent in a bounded way. Per-layer kernel launches pass a handful of scalar args that vary per step (current token count, current block index). Those are captured as CUgraphNodeParams.

Turning graphs off (--profile) disables capture — useful when you are profiling under nsys because graphs collapse every kernel into a single graph-launch event and defeat per-kernel timing.

Writing a new kernel — the minimum

  1. Pick the directory. A kernel every GB10 target can use goes in kernels/gb10/common/; one only a single target needs goes in its leaf, kernels/gb10/<model>/<quant>/ (a leaf file with a common/ file's name is a shadow and must be declared in the leaf's KERNEL.toml [shadow]).
  2. Drop your_kernel.cu with the SPDX header and a extern "C" __global__ entry.
  3. Decide the module name. Default is the file stem (your_kernel). Override in KERNEL.toml if you want a short name.
  4. Call it from the layer. In crates/model-layers/src/layers/<your_layer>.rs, gpu.kernel("your_kernel_module", "your_kernel_function") returns a KernelHandle. Store it in the layer struct at load time, not per-step.
  5. Wire the launch with the typed launch helpers the layer's ops use (crates/model-layers/src/layers/ops/, crates/gpu-runtime/src/gpu_args.rs): grid, block, shared memory, the argument list and the stream, then GpuBackend::launch.
  6. Benchmark. Add a shape to a met benchmark benchmark (crates/bench) or a micro-benchmark: a *_microtest example in crates/model-arch/examples/, or a Criterion bench in the relevant crate. A kernel without a benchmark is not allowed to claim "faster".
  7. Verify correctness against a PyTorch reference on a fixture tensor. Numerical diff tolerance: for BF16 outputs, abs-tol 1e-3 / rel-tol 1e-2 is a typical starting point.

Anti-patterns

  • Don't branch on runtime flags inside the hot loop. If a kernel needs two variants (NVFP4 vs BF16 KV), make them two kernels.
  • Don't try to be generic. The whole point is specialisation. A kernel that works for three batch sizes is usually slower than three kernels that work for one each.
  • Don't call cudaDeviceSynchronize anywhere inside a kernel launch path. GpuBackend::synchronize(stream) exists for explicit syncs. Random device-sync calls break CUDA graph capture and defeat pipelining.
  • Don't mutate __constant__ memory at runtime. Upload it once at load time; use it forever.