Fix layernorm_fp8 scale-storage bug and bandwidth miscalculation#2
Open
pauliano22 wants to merge 1 commit into
Open
Fix layernorm_fp8 scale-storage bug and bandwidth miscalculation#2pauliano22 wants to merge 1 commit into
pauliano22 wants to merge 1 commit into
Conversation
- 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
layernorm_quant_kernelonly ever stored row 0's dequantization scale. The store was gated ontl.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 returnedscalestensor is left as whatevertorch.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.bench_layernorm_fp8.pyused 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.plot_nameinbench_layernorm_fp8.pyandbench_relu.pyto match the actual checked-inresults/*.csv/*.pngfilenames — 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.requirements.txt:-e git+https://github.com/pauliano22/triton-gpu-kernels.git@<old-sha>#egg=triton_kernels. This looks like a leftover from runningpip freezeagainst a localpip install -e ., and would pull a stale historical commit of the repo itself on a freshpip install -r requirements.txt.Test plan
python3 -m py_compileon all changed.pyfiles — syntax OKlayer_norm_fp8.pyfix by reading the kernel's execution model (one program per row viatl.program_id(0)) — the fix is a one-line removal of an incorrect conditionaltorch/tritoninstalled 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 checkingscaleshas no zero/garbage rows beyond row 0's expected value.Generated by Claude Code