[FA] Fix tutorial blackwell perf

Differential Revision: D81343267

Pull Request resolved: #377

Author: manman-ren

tritonbench/kernels/blackwell_triton_fused_attention.py

@@ -402,7 +402,9 @@ def grid(META):
         ctx.grid = grid
         warp_specialize = baseVariant == "ws"
         if is_blackwell() and warp_specialize:
-            if HEAD_DIM_K == 128 and q.dtype == torch.float16:
+            if HEAD_DIM_K == 128 and (
+                q.dtype == torch.float16 or q.dtype == torch.bfloat16
+            ):
                 extra_kern_args["maxnreg"] = 168
             else:
                 extra_kern_args["maxnreg"] = 80