perf: optimize small-head causal prefill

- map D=32 and D=64 causal prefill to BC=64 tiles

- add small-head correctness and benchmark coverage

csrc/kernels/attn_dispatchers.cuh

@@ -59,14 +59,34 @@ inline int compute_num_splits(int base_blocks, int tiles_total,
 // ======================================================================
 
 #ifndef ASTRAI_NO_MMA
+template <int BC_>
+struct PrefillKernelConfig {
+    static constexpr int BC = BC_;
+    static constexpr int WARPS = 4;
+    static constexpr int STAGES = 2;
+};
+
+// Compile-time configuration map shared by contiguous and paged prefill.
+// Unsupported head dimensions intentionally have no mapping.
+template <int HEAD_DIM, bool IsCausal>
+struct PrefillConfigMap;
+
+template <> struct PrefillConfigMap<32, false>  : PrefillKernelConfig<32> {};
+template <> struct PrefillConfigMap<32, true>   : PrefillKernelConfig<64> {};
+template <> struct PrefillConfigMap<64, false>  : PrefillKernelConfig<32> {};
+template <> struct PrefillConfigMap<64, true>   : PrefillKernelConfig<64> {};
+template <> struct PrefillConfigMap<128, false> : PrefillKernelConfig<32> {};
+template <> struct PrefillConfigMap<128, true>  : PrefillKernelConfig<32> {};
+template <> struct PrefillConfigMap<256, false> : PrefillKernelConfig<16> {};
+template <> struct PrefillConfigMap<256, true>  : PrefillKernelConfig<16> {};
+
 template <typename KV>
 struct PrefillLauncherMMA {
     template <int HEAD_DIM, bool IsCausal, bool HasMask>
     static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
-        constexpr int WARPS = 4;
+        using Config = PrefillConfigMap<HEAD_DIM, IsCausal>;
-        constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
+        using Traits = KernelTraits<HEAD_DIM, Config::BC, Config::WARPS, Config::STAGES>;
-        using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
+        constexpr int ROWS = Traits::BR * Config::WARPS;
-        constexpr int ROWS = Traits::BR * WARPS;
         dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head,
                   KV::kPaged ? 1 : p.batch);
         dim3 block(Traits::NUM_THREADS);

csrc/tests/attn_test.cu

@@ -118,7 +118,8 @@ static void bench_decode() {
     printf("\n===== DECODE BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
     print_bench_header();
 
-    for (int ci = 0; ci < 6; ci++) {
+    int n = sizeof(cfgs) / sizeof(cfgs[0]);
+    for (int ci = 0; ci < n; ci++) {
         int B = cfgs[ci][0], Hq = cfgs[ci][1], Hk = cfgs[ci][2];
         int sl = cfgs[ci][3], D = cfgs[ci][4];
         size_t nQ = (size_t)B * Hq * D;
@@ -229,6 +230,12 @@ static int run_prefill_test(int B, int Hq, int Hk, int ql, int kl, int D, int ca
 
 static void bench_prefill() {
     const int cfgs[][7] = {
+        {1,32,4,1024,1024,32,0},
+        {1,32,4,1024,1024,32,1},
+        {1,32,4,4096,4096,32,1},
+        {1,32,4,1024,1024,64,0},
+        {1,32,4,1024,1024,64,1},
+        {1,32,4,4096,4096,64,1},
         {1,32,4,512,512,128,0},
         {1,32,4,1024,1024,128,0},
         {1,32,4,2048,2048,128,0},
@@ -324,7 +331,9 @@ int main() {
     {
         const int configs[][7] = {
             {1,2,1,64,128,32,0},     // scalar fallback D=32
+            {1,4,2,256,256,32,1},     // causal D=32 dispatch
             {1,2,1,64,128,64,0},     // tiny: B,Hq,Hk,q,kv,D,causal
+            {1,4,2,256,256,64,1},     // causal D=64 dispatch
             {1,32,4,512,512,128,0},  // standard
             {1,32,4,128,256,128,0},  // medium
             {1,4,2,256,256,128,1},   // causal