Fix SM90 MQA logits (#229)

2025-11-19 10:50:36 +08:00
parent 2f9d87877e
commit ec5e9ed0b8
2 changed files with 2 additions and 2 deletions
@@ -87,7 +87,7 @@ static void smxx_fp8_mqa_logits(const torch::Tensor& q,
    constexpr int block_qh = 128;
    constexpr int block_kv = 256;
    constexpr int num_specialized_threads = 128;
-    constexpr int num_math_threads = 256;
+    const int num_math_threads = (device_runtime->get_arch_major() == 10 ? 256 : 512);
    constexpr int num_q_stages = 3, num_kv_stages = 3;
    const int block_q = block_qh / num_heads;
    DG_HOST_ASSERT(block_qh % num_heads == 0);