need to handle braodcasted grads

Author: drisspg

benchmarks/bench_dynamic_linear_fused_cast.py

@@ -116,7 +116,8 @@ def main(
         ref_forw_backward = lambda: linear_ref(input_tensor).sum().backward()
 
         def float8_forw_backward():
-            linear_float8(input_tensor).sum().backward()
+            out = linear_float8(input_tensor)
+            out.sum().backward()
 
         def n_times(n, fn, *args, **kwargs):
             def wrapper(*args, **kwargs):

float8_experimental/float8_tensor.py

@@ -48,6 +48,7 @@ def to_fp8_no_autograd(
         from driss_torch import saturated_cast
 
         bits_fp8 = saturated_cast(x, x_scale, float8_dtype)
+
     else:
         x_scaled = x * x_scale
         bits_fp8 = to_fp8_saturated(x_scaled, float8_dtype)

float8_experimental/fused_kernels/fused_casting_kernels.py

@@ -28,6 +28,7 @@ def maximum(a, b):
         mask |= a != a
     return tl.where(mask, a, b)
 
+
 @triton.jit
 def abs_max_kernel(
     x_ptr,