Fix layernorm_fp8 scale-storage bug and bandwidth miscalculation

benchmarks/bench_layernorm_fp8.py

@@ -12,7 +12,7 @@
         line_names=['PyTorch BF16', 'Triton Fused FP8'], 
         styles=[('green', '-'), ('blue', '-')], 
         ylabel='GB/s', 
-        plot_name='layernorm-fp8-vs-bf16',
+        plot_name='layernorm-fp8',
         args={'M': 2048}, 
     )
 )
@@ -29,9 +29,10 @@ def benchmark(M, N, provider):
         ms, min_ms, max_ms = triton.testing.do_bench(lambda: layernorm_fp8(x, w, b), quantiles=quantiles)
 
     # Calculate Bandwidth: (Reads + Writes) / time
-    # BF16: M*N*2 (read X) + M*N*2 (read W/B) + M*N*2 (write Y) 
+    # BF16: M*N*2 (read X) + M*N*2 (read W/B) + M*N*2 (write Y)
     # FP8:  M*N*2 (read X) + M*N*2 (read W/B) + M*N*1 (write Y) -> 25% less traffic!
-    gbps = lambda ms: (M * N * 6) / ms / 1e6
+    bytes_per_elem = 6 if provider == 'torch_bf16' else 5
+    gbps = lambda ms: (M * N * bytes_per_elem) / ms / 1e6
     return gbps(ms), gbps(max_ms), gbps(min_ms)
 
 if __name__ == "__main__":

benchmarks/bench_relu.py

@@ -27,7 +27,7 @@ def simple_bench(fn, iterations=10):
         line_names=['Triton (Interpreted)', 'PyTorch (CPU)'], 
         styles=[('blue', '-'), ('green', '-')], 
         ylabel='Execution Time (ms)', # Switched to ms because GB/s is misleading on CPU
-        plot_name='relu-performance-local',
+        plot_name='relu-performance',
         args={}, 
     )
 )

kernels/layer_norm_fp8.py

@@ -34,10 +34,9 @@ def layernorm_quant_kernel(
 
     y_fp8 = (y_fp32 / scale).to(tl.float8e4m3fn)
 
-    # 3. Store the FP8 data and the scale factor
+    # 3. Store the FP8 data and this row's scale factor
     tl.store(Y_ptr + cols, y_fp8, mask=mask)
-    if tl.program_id(0) == 0: # Store the scale once for the whole tensor or per row
-        tl.store(Scale_ptr + row_idx, scale)
+    tl.store(Scale_ptr + row_idx, scale)
 
 def layernorm_fp8(x, w, b):
     M, N = x.shape

requirements.txt

@@ -38,6 +38,5 @@ six==1.17.0
 sympy==1.14.0
 torch==2.11.0
 triton==3.6.0
--e git+https://github.com/pauliano22/triton-gpu-kernels.git@46ece4d30cc46844e9de73ca933ed94cfbb15395#egg=triton_kernels
 typing_extensions==4.15.0
 tzdata==2025.3