Fix FP8 LayerNorm scale bug, benchmark byte-count formula, and missing flash_attn wrapper

[pauliano22]

Summary

kernels/layer_norm_fp8.py — the kernel only stored the per-row dequant scale for row 0 (if tl.program_id(0) == 0), leaving every other row's entry in the scales tensor uninitialized. Dequantizing any row but the first (y_fp8 * scale) produced garbage. Now always stores the scale.

benchmarks/bench_layernorm_fp8.py — the GB/s formula always used the BF16 byte-traffic multiplier (M*N*6) even for the FP8 line, contradicting the file's own comment (which says FP8 should be M*N*5) and inflating the reported FP8 throughput by ~20% — this is the number quoted in the README (3,531.8 GB/s). Now selects the multiplier per provider.

kernels/flash_attn.py — benchmarks/bench_flash_attn.py imports flash_attn from this module, but only the low-level flash_attn_kernel existed; the benchmark has never been able to run (ImportError), which is presumably why the README only reports a PyTorch baseline for FlashAttention and no Triton number. Added the host-side flash_attn(q, k, v) wrapper, mirroring the pattern used in relu.py/layer_norm.py. While wiring it up, found two more issues that would've made the kernel fail to even compile:

• d_head was used as a tl.zeros/tl.arange shape argument without being marked tl.constexpr, which Triton requires for compile-time shapes.
• the K/V loop bound was misleadingly named n_heads when it's actually the sequence length (grid = one program per (batch, head) pair); renamed to seq_len.

requirements.txt — dropped a stale self-referential editable install of this repo pinned to an old commit hash (leftover from a pip freeze after pip install -e .).

Note: the flash_attn "Lite" kernel has no bounds masking on the query block, so SEQ_LEN must be a power of two — documented and asserted in the wrapper.

Test plan

No CUDA GPU was available to me, so I verified all three logic fixes via Triton's CPU interpreter mode (TRITON_INTERPRET=1) rather than skipping verification:

• layer_norm_fp8: reproduced the original bug (7/8 rows left uninitialized) and confirmed the fix gives every row a correct, distinct scale matching an fp32 reference layernorm.
• flash_attn: output matches torch.nn.functional.scaled_dot_product_attention(scale=1.0) (the kernel has no 1/sqrt(d_head) scaling, consistent with the benchmark's own unscaled naive_attn reference) within float16 tolerance, tested with both square dims and BATCH != N_HEADS != SEQ_LEN != D_HEAD, plus a per-block isolation check confirming the grid/stride indexing doesn't cross-contaminate between heads.
• bench_layernorm_fp8.py byte-formula fix is pure arithmetic, checked independently.
• Would appreciate a real H100 run to confirm the actual GB/s/TFLOPS numbers and update the README figures if they change materially.

---

Generated by Claude Code