fix the round mode uncontrollable (#1535)

minmengdie · web-flow · commit 472513da2dde · 2025-12-05T14:53:24.000+08:00
diff --git a/csrc/py_itfs_cu/asm_mha_fwd.cu b/csrc/py_itfs_cu/asm_mha_fwd.cu
@@ -319,6 +319,7 @@ std::vector<at::Tensor> fmha_v3_fwd(at::Tensor &q, // [b, sq, hq, d]
                                  bias_type,
                                  has_lse,
                                  quant_scale_enum::no_scale,
+                                 true,
                                  how_v3_bf16_cvt);
         TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd");
     }
diff --git a/op_tests/test_mha.py b/op_tests/test_mha.py
@@ -106,7 +106,7 @@ def run_ck(
         deterministic,
         return_lse=return_lse,
         return_attn_probs=return_attn_probs,
-        how_v3_bf16_cvt=1,
+        how_v3_bf16_cvt=2,
         cu_seqlens_q=cu_seqlens_q,
         cu_seqlens_kv=cu_seqlens_kv,
         num_rotate_args=1,

Original file line number	Diff line number	Diff line change
`@@ -319,6 +319,7 @@ std::vector<at::Tensor> fmha_v3_fwd(at::Tensor &q, // [b, sq, hq, d]`
`319`	`319`	`bias_type,`
`320`	`320`	`has_lse,`
`321`	`321`	`quant_scale_enum::no_scale,`
	`322`	`+ true,`
`322`	`323`	`how_v3_bf16_cvt);`
`323`	`324`	`TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd");`
`324`	`325`	`}`