fix mha_fwd_kvcache_mla_nope_pe接口架构判断

5e577dee · zhanghj2 · 79c06a56 · 5e577dee
Commit 5e577dee authored Mar 09, 2026 by zhanghj2
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

csrc/extension/flash_api.h csrc/extension/flash_api.h +2 -2

No files found.
--- a/csrc/extension/flash_api.h
+++ b/csrc/extension/flash_api.h
@@ -337,8 +337,8 @@ mha_fwd_kvcache_mla_nope_pe(
    // bool is_sm90 = dprops->major == 9 && dprops->minor == 0;
    // TORCH_CHECK(is_sm90);
    Arch arch = Arch();
-    if (!arch.is_gfx93x() || !arch.is_gfx928()) {
-        TORCH_CHECK(false, "Dense decode MLA is only supported on gfx936 or gfx938 architecture");
+    if (!(arch.is_gfx93x() || arch.is_gfx928())) {
+        TORCH_CHECK(false, "Dense decode MLA is only supported on gfx928, gfx936 or gfx938 architecture");
    }
    at::Tensor vcache = vcache_.has_value() ? vcache_.value() : kcache;