feat: scalar paged prefill fallback and decode causal fix

- Add scalar paged prefill kernel mirroring split-Q MMA indexing for sm<80 - Wire scalar path into dispatch_paged_prefill under ASTRAI_NO_MMA - Fix paged decode scalar causal mask dropping all kv>0 for decode
2026-08-01 16:52:01 +08:00
parent fda82ee232
commit 925cbedc93
3 changed files with 142 additions and 4 deletions
@@ -75,10 +75,9 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
                if (!p.mask[mask_base + kv_idx])
                    masked = true;
            }
-            if constexpr (IsCausal) {
-                if (kv_idx > p.causal_offset)
-                    masked = true;
-            }
+            // Decode: the query is the last token, so its valid range [0,
+            // seq_len) IS the causal range.  IsCausal is accepted for dispatch
+            // uniformity but must not apply causal_offset masking here.
            if (masked)
                partial = -FLT_MAX;