已合并
test: add math SDP fp16 bf16 reduction API validation #41428
zhaoziyi-2026创建于 7月13日
test: add math SDP fp16 bf16 reduction API validation #41428
已合并
共 1 个文件变更+51-0
| @@ -1,3 +1,13 @@ | |||
| 1 | +""" | ||
| 2 | +Add validation cases for torch._C._set_math_sdp_allow_fp16_bf16_reduction API on NPU: | ||
| 3 | + | ||
| 4 | + 1. PyTorch community lacks sufficient and direct API validation for this API, so this case is added. | ||
| 5 | + 2. This file validates setter/getter state switching, backend wrapper linkage, and invalid input handling. | ||
| 6 | + | ||
| 7 | +Backend API tests for NPU SDP switches and the private math SDP | ||
| 8 | +fp16/bf16 reduction flag shared by torch._C and backend wrappers. | ||
| 9 | +""" | ||
| 10 | + | ||
| 1 | import torch | 11 | import torch |
| 2 | import torch.nn as nn | 12 | import torch.nn as nn |
| 3 | 13 | ||
| @@ -42,6 +52,47 @@ class TorchBackendsApiTestCase(TestCase): | |||
| 42 | self.assertEqual(mem_mem_efficient_res, False) | 52 | self.assertEqual(mem_mem_efficient_res, False) |
| 43 | self.assertEqual(math_res, False) | 53 | self.assertEqual(math_res, False) |
| 44 | 54 | ||
| 55 | + def test_math_sdp_allow_fp16_bf16_reduction_setter(self): | ||
| 56 | + # The private setter should update the matching private getter. | ||
| 57 | + setter = torch._C._set_math_sdp_allow_fp16_bf16_reduction | ||
| 58 | + getter = torch._C._get_math_sdp_allow_fp16_bf16_reduction | ||
| 59 | + original = getter() | ||
| 60 | + self.addCleanup(setter, original) | ||
| 61 | + | ||
| 62 | + setter(True) | ||
| 63 | + self.assertEqual(getter(), True) | ||
| 64 | + setter(False) | ||
| 65 | + self.assertEqual(getter(), False) | ||
| 66 | + | ||
| 67 | + def test_math_sdp_allow_fp16_bf16_reduction_backend_wrapper(self): | ||
| 68 | + # The public backend wrapper should control the same underlying flag. | ||
| 69 | + setter = torch._C._set_math_sdp_allow_fp16_bf16_reduction | ||
| 70 | + getter = torch._C._get_math_sdp_allow_fp16_bf16_reduction | ||
| 71 | + wrapper = torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp | ||
| 72 | + original = getter() | ||
| 73 | + self.addCleanup(setter, original) | ||
| 74 | + | ||
| 75 | + wrapper(True) | ||
| 76 | + self.assertEqual(getter(), True) | ||
| 77 | + wrapper(False) | ||
| 78 | + self.assertEqual(getter(), False) | ||
| 79 | + | ||
| 80 | + def test_math_sdp_allow_fp16_bf16_reduction_invalid_value(self): | ||
| 81 | + # Invalid inputs should raise and preserve the current math SDP flag. | ||
| 82 | + setter = torch._C._set_math_sdp_allow_fp16_bf16_reduction | ||
| 83 | + getter = torch._C._get_math_sdp_allow_fp16_bf16_reduction | ||
| 84 | + original = getter() | ||
| 85 | + self.addCleanup(setter, original) | ||
| 86 | + | ||
| 87 | + setter(True) | ||
| 88 | + with self.assertRaises((RuntimeError, TypeError)): | ||
| 89 | + setter(1) | ||
| 90 | + self.assertEqual(getter(), True) | ||
| 91 | + | ||
| 92 | + with self.assertRaises((RuntimeError, TypeError)): | ||
| 93 | + setter(None) | ||
| 94 | + self.assertEqual(getter(), True) | ||
| 95 | + | ||
| 45 | def test_aclnn_allow_hf32(self): | 96 | def test_aclnn_allow_hf32(self): |
| 46 | res = torch.npu.aclnn.allow_hf32 | 97 | res = torch.npu.aclnn.allow_hf32 |
| 47 | self.assertEqual(res, True) | 98 | self.assertEqual(res, True) |