Skip to content

Fix layernorm_fp8 scale-storage bug and bandwidth miscalculation#2

Open
pauliano22 wants to merge 1 commit into
mainfrom
fix/layernorm-fp8-scale-bug-and-bandwidth-calc
Open

Fix layernorm_fp8 scale-storage bug and bandwidth miscalculation#2
pauliano22 wants to merge 1 commit into
mainfrom
fix/layernorm-fp8-scale-bug-and-bandwidth-calc

Conversation

@pauliano22

Copy link
Copy Markdown
Owner

Summary

  • layernorm_quant_kernel only ever stored row 0's dequantization scale. The store was gated on tl.program_id(0) == 0, but each program instance handles exactly one row (row_idx = tl.program_id(0)). Every other row's entry in the returned scales tensor is left as whatever torch.empty((M,), ...) happened to allocate — uninitialized memory. Dequantizing any row other than 0 with that scale would silently produce garbage. Removed the guard so every program stores its own row's scale.
  • The FP8 vs. BF16 bandwidth comparison in bench_layernorm_fp8.py used the same byte multiplier (6) for both providers, even though the comment directly above it says FP8's output write is 1 byte instead of 2 (5 bytes/elem total vs. BF16's 6). This overstated the FP8 kernel's reported GB/s and speedup in the README by ~20%. Now uses 5 for the FP8 provider, 6 for BF16.
  • Fixed plot_name in bench_layernorm_fp8.py and bench_relu.py to match the actual checked-in results/*.csv/*.png filenames — as written, re-running either script would write to a different filename than the one committed, silently diverging from the results shown in the README.
  • Removed a stray self-referential line from requirements.txt: -e git+https://github.com/pauliano22/triton-gpu-kernels.git@<old-sha>#egg=triton_kernels. This looks like a leftover from running pip freeze against a local pip install -e ., and would pull a stale historical commit of the repo itself on a fresh pip install -r requirements.txt.

Test plan

  • python3 -m py_compile on all changed .py files — syntax OK
  • Verified the layer_norm_fp8.py fix by reading the kernel's execution model (one program per row via tl.program_id(0)) — the fix is a one-line removal of an incorrect conditional
  • Not able to run these kernels in this environment — no CUDA GPU, no torch/triton installed here. Would appreciate a sanity run on the H100 box before merging, e.g. layernorm_fp8(x, w, b) on a multi-row input and checking scales has no zero/garbage rows beyond row 0's expected value.

Generated by Claude Code

- layernorm_quant_kernel gated the scale store on
  `tl.program_id(0) == 0`, but each program instance handles exactly
  one row (row_idx = program_id(0)). This meant only row 0's entry in
  the returned `scales` tensor was ever written; every other row was
  left as uninitialized memory from torch.empty(), so dequantizing any
  row other than 0 would use garbage. Removed the guard so every
  program stores its own row's scale.

- bench_layernorm_fp8.py's bandwidth formula used a fixed 6-bytes-per-
  element multiplier for both providers, but the comment directly
  above it (and the actual I/O pattern) says FP8's output write is
  1 byte vs BF16's 2, i.e. 5 bytes/elem total for FP8 vs 6 for BF16.
  Using 6 for both overstated the FP8 kernel's reported bandwidth/
  speedup in the README by ~20%.

- Fixed bench_layernorm_fp8.py/bench_relu.py plot_name to match the
  actual checked-in results/ filenames, so re-running the scripts
  reproduces (rather than diverges from) the committed CSV/PNG.

- Removed a self-referential, pinned-to-an-old-commit editable install
  of this same repo from requirements.txt (an artifact of running
  `pip freeze` against a local `pip install -e .`), which would pull a
  stale historical commit of the repo itself on a fresh install.

Not able to execute these kernels in this environment (no CUDA GPU, no
torch/triton installed) - the layer_norm_fp8.py fix is a one-line
removal of an incorrect conditional and was verified by code reading
against the kernel's per-row execution model, but hasn't been run.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants