fix(fmha): support >4GB KV cache in batch prefill via runtime dispatch

csrc/cpp_itfs/mha_fwd_batch_prefill.cu

@@ -47,7 +47,13 @@ float mha_batch_prefill(mha_batch_prefill_args args,
     int head_size_q  = args.hdim_q;
     int head_size_v  = args.hdim_v;
     bool has_dropout = args.p_drop > 0.f;
-    auto traits      = get_mha_batch_prefill_traits(head_size_q,
+
+    // The kUseGlobalLoad decision (>2GB KV cache → use `global_load_lds_*`
+    // instead of SRD `buffer_load_*`) is made per-arm inside the auto-generated
+    // dispatcher in fmha_batch_prefill_api.cpp, where each arm knows its own
+    // compile-time bn0 and dtype element size. The wrapper just forwards args;
+    // no runtime trait field for it.
+    auto traits = get_mha_batch_prefill_traits(head_size_q,
                                                head_size_v,
                                                q_dtype_str,
                                                is_group_mode,

csrc/py_itfs_ck/mha_batch_prefill_kernels.cu

@@ -817,7 +817,13 @@ mha_batch_prefill(at::Tensor& q,       // [total_q, hq, d]
                                            has_lse,
                                            qscale_type,
                                            false);
-        TORCH_CHECK(t >= 0, "invalid argument for batch_prefill");
+        TORCH_CHECK(t >= 0,
+                    "invalid argument for batch_prefill: no matching kernel found. "
+                    "page_size=", args.page_block_size,
+                    ", num_pages=", args.num_total_pages,
+                    ", dtype=", dtype_str,
+                    ". If KV cache exceeds 2GB (INT32_MAX byte offset) with page_size < kN0, "
+                    "CDNA3+ GPU (MI300/MI350) is required.");
     }
     else
     {