[Gemma4]: Fix FP8 Triton scale layout (#25286)

2026-05-19 14:00:23 -07:00
parent 8322fe09a7
commit fab097d66d
2 changed files with 24 additions and 2 deletions
@@ -58,12 +58,14 @@ class TestScaledMM(CustomTestCase):
        """Test core functionality with reduced precision requirements"""
        test_configs = [
            (32, 32, 32, torch.int8, torch.float16, False),
+            (17, 64, 96, torch.int8, torch.float16, False),
            (64, 64, 64, torch.int8, torch.float16, True),
        ]

        try:
            torch.tensor([1.0], dtype=torch.float8_e4m3fn, device=self._device)
            test_configs.append((32, 32, 32, torch.float8_e4m3fn, torch.float16, False))
+            test_configs.append((17, 64, 96, torch.float8_e4m3fn, torch.float16, False))
        except:
            print("FP8 not supported, skipping")

@@ -98,6 +100,14 @@ class TestScaledMM(CustomTestCase):

                torch.testing.assert_close(triton_out, ref_out, rtol=rtol, atol=atol)

+                scale_b_row = scale_b.t().contiguous()
+                triton_out_row_scale = triton_scaled_mm(
+                    input, weight, scale_a, scale_b_row, out_dtype, bias
+                )
+                torch.testing.assert_close(
+                    triton_out_row_scale, ref_out, rtol=rtol, atol=atol
+                )
+

 if __name__ == "__main__":
    unittest.main(verbosity=2)