fix: correct prefill mask index, unify GQA kernel interface

- Fix mask indexing: batch*q_len*kv_len -> batch*kv_len
- Add csrc/kernels/gqa_common.cuh with shared GQAParams struct
- Unify decode/prefill Python API: both accept (q,k,v,mask=None,...)
- Decode now supports optional mask, is_causal, causal_offset, scale
- Rename struct fields: B->batch, Hq->q_head, Hk->kv_head, D->head_dim
- Use py::arg() for correct None/defaults handling in pybind11
- Update pure C tests and build instructions (-arch=sm_89)

csrc/tests/gqa_decode_test.cu

@@ -1,5 +1,4 @@
-// Pure-C test for decode kernel
-// compile: nvcc -I csrc csrc/tests/gqa_decode_test.cu -o test && ./test
+// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_decode_test.cu -o test && ./test
 #include <cstdio>
 #include <cstdlib>
 #include <cmath>
@@ -85,6 +84,12 @@ int main() {
         cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
         cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
 
+        GQAParams p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
+        p.use_mask=1; p.is_causal=0; p.causal_offset=0;
+        p.scale=1.0f/sqrtf((float)D);
+        p.q=dQ; p.k=dK; p.v=dV; p.mask=dMask; p.o=dO;
+
         size_t smem=DC_CHUNK*D*sizeof(bf16);
         dim3 block(32, gs);
         dim3 grid(B*Hk);
@@ -92,8 +97,7 @@ int main() {
                grid.x, block.x, block.y, smem);
 
         double t0=now_ms();
-        gqa_decode_attn_kernel<<<grid,block,smem>>>(dQ,dK,dV,dMask,dO,
-            B,Hq,Hk,sl,D);
+        gqa_decode_attn_kernel<<<grid,block,smem>>>(p);
         cudaDeviceSynchronize();
         double kms=now_ms()-t0;
         cudaError_t err=cudaGetLastError();

csrc/tests/gqa_prefill_test.cu

@@ -1,4 +1,4 @@
-// Pure-C test: compile with nvcc -I csrc csrc/tests/gqa_prefill_test.cu -o test && ./test
+// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_prefill_test.cu -o test && ./test
 #include <cstdio>
 #include <cstdlib>
 #include <cmath>
@@ -81,6 +81,12 @@ int main() {
         for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
         cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
 
+        GQAParams p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
+        p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
+        p.scale=1.0f/sqrtf((float)D);
+        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
+
         dim3 grid((ql+Br-1)/Br, Hq, B);
         dim3 block(32, Br, 1);
         size_t smem=2*Bc*D*sizeof(bf16);
@@ -88,8 +94,7 @@ int main() {
                grid.x,grid.y,grid.z, block.x,block.y,block.z, smem);
 
         double t0=now_ms();
-        gqa_prefill_attn_kernel<<<grid,block,smem>>>(dQ,dK,dV,nullptr,dO,
-            B,Hq,Hk,ql,kl,D,causal,0,0);
+        gqa_prefill_attn_kernel<<<grid,block,smem>>>(p);
         cudaDeviceSynchronize();
         double kms=now_ms()-t0;
         cudaError_t err=cudaGetLastError();