fmha: add sm_100 and sm_120 (#1890)

Aidyn-A · web-flow · commit a255aa6139d1 · 2025-06-26T08:37:38.000+09:00
diff --git a/apex/contrib/csrc/fmha/fmha_api.cpp b/apex/contrib/csrc/fmha/fmha_api.cpp
@@ -97,7 +97,9 @@ mha_fwd(const at::Tensor &qkv,         // total x num_heads x 3 x head_size, tot
     using namespace torch::indexing;
     auto dprops = at::cuda::getCurrentDeviceProperties();
     TORCH_CHECK((dprops->major == 8 && dprops->minor == 0) ||
-                (dprops->major == 9 && dprops->minor == 0));
+                (dprops->major == 9 && dprops->minor == 0) ||
+                (dprops->major == 10 && dprops->minor == 0) ||
+                (dprops->major == 12 && dprops->minor == 0));
     auto stream = at::cuda::getCurrentCUDAStream().stream();
     Launch_params<Fused_multihead_attention_fprop_params> launch_params(dprops, stream, is_training, is_nl);
 
@@ -193,7 +195,9 @@ mha_bwd(const at::Tensor &dout,  // total x num_heads, x head_size
     using namespace torch::indexing;
     auto dprops = at::cuda::getCurrentDeviceProperties();
     TORCH_CHECK((dprops->major == 8 && dprops->minor == 0) ||
-                (dprops->major == 9 && dprops->minor == 0));
+                (dprops->major == 9 && dprops->minor == 0) ||
+                (dprops->major == 10 && dprops->minor == 0) ||
+                (dprops->major == 12 && dprops->minor == 0));
     int seq_len = 512;
     auto launch = &run_fmha_dgrad_fp16_512_64_sm80;
     if( max_seq_len <= 128 ) {
diff --git a/apex/contrib/test/fmha/test_fmha.py b/apex/contrib/test/fmha/test_fmha.py
@@ -61,7 +61,7 @@ def py_mha(qkv, amask, b, s, h, d):
 
 
 @unittest.skipIf(SKIP_TEST, f"{SKIP_TEST}")
-@unittest.skipIf(not _get_device_properties() == (8, 0), "FMHA only supports sm80")
+@unittest.skipIf(_get_device_properties() not in [(8, 0), (9, 0), (10, 0), (12, 0)], "FMHA only supports sm80")
 class TestFMHA(unittest.TestCase):
 
     def run_test(self, s: int, b: int, zero_tensors: bool):
