Use larger MMA shape (#227)

2025-11-14 11:38:15 +08:00
parent c9f8b34dcd
commit 2f9d87877e
2 changed files with 46 additions and 56 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 = 512;
+    constexpr int num_math_threads = 256;
    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);