write-fp8-kernel

Skill: Write an FP8 Kernel

Purpose

Guide the agent through designing and implementing FP8 compute kernels for inference and training on NVIDIA Hopper (sm_90) and Ada Lovelace (sm_89) hardware, covering FP8 format selection, scaling strategy, tensor core usage via WGMMA or cuBLAS, and dequantization epilogue design.

Use this when

• Writing a high-throughput GEMM, attention, or linear layer kernel targeting H100 or H200 (sm_90a) hardware where FP8 tensor cores provide 2x the throughput of FP16.
• Implementing FP8 training (forward pass in E4M3, gradient computation in E5M2) following the transformer engine or similar recipe.
• Evaluating whether FP8 achieves sufficient accuracy for a given model and activation distribution.
• Fusing quantization, GEMM, and dequantization into a single pass to avoid expensive fp32 ↔ fp8 conversion round-trips in memory.

Do not use this when

• The target hardware is pre-Ada (sm_80 or earlier) — FP8 tensor core instructions do not exist. fp16 or bf16 is the correct choice.
• The activation range is highly dynamic or the per-tensor scale would need to be recomputed at sub-batch granularity with prohibitive overhead.
• Numerical accuracy has not been validated for the target model. FP8 E4M3 has a much narrower dynamic range than fp16; overflows and underflows silently corrupt outputs without per-tensor or per-channel scaling.
• The simpler option (FP16 GEMM via cuBLAS or CUTLASS with a Flash Attention kernel) already meets the performance target.

Inputs the agent should gather first

• Hardware target: H100/H200 (sm_90a), RTX 4090/Ada (sm_89), or other. WGMMA instructions are H100-specific (sm_90a). Ada exposes FP8 via cuBLAS/CUTLASS but without WGMMA.
• FP8 format: E4M3 (4-bit exponent, 3-bit mantissa, max representable value ≈ 448) for forward pass weights and activations; E5M2 (5-bit exponent, 2-bit mantissa, max ≈ 57344) for gradients. Clarify which is needed.
• Scaling strategy: per-tensor (one scale for the whole tensor), per-channel/per-row (one scale per row or column), or delayed scaling (track amax history across steps and update scales at a fixed interval). Transformer Engine uses delayed scaling.
• Accumulation dtype: fp32 is mandatory for FP8 tensor core accumulation. DO NOT use fp16 accumulators with FP8 inputs — precision loss is severe.
• Output dtype: fp16, bf16, or fp32 after dequantization. This determines the epilogue.
• Operation type: GEMM only, GEMM + bias, attention (QK^T scaling), or custom fused operation.
• Whether cuBLAS FP8 API or CUTLASS FP8 templates are acceptable: these should be the default choice unless there is a specific reason for a custom kernel.