[Bug] Fix V4-Pro NaN on Blackwell by converting fp8_einsum input scale to ue8m0 (#25733)

2026-05-18 23:48:34 -07:00
parent 7c3f614e23
commit 79ea30d1f1
1 changed files with 1 additions and 0 deletions
@@ -623,6 +623,7 @@ class MQALayer(nn.Module):
                o.reshape(T * G, D).contiguous(),
                group_size=128,
            )
+            o_s = deep_gemm.ceil_to_ue8m0(o_s)
            output = torch.empty(T, G, R, device=o.device, dtype=torch.bfloat16)
            deep_gemm.fp8_einsum(
                "bhr,hdr->bhd",