perf: batch packed torch-native attention into one padded sdpa call

- replace the per-request python loop in the torch-native 3-d inference path with a single F.scaled_dot_product_attention call over [B, max_q, max_kv] padded tensors
- fold the per-request causal offset (seq_len - q_len) and kv padding into one bool mask [B, 1, max_q, max_kv]; padded q rows gather row 0 and are dropped by the [q_valid] unpack, which restores the flat qo_indptr order
- decode drops from B sdpa launches plus ~3B host syncs (int() on qo_indptr/seq_lens per request) to one call with two syncs
- masked lanes contribute exact-zero weights, so prefill logits stay bitwise identical and test_prefill_with_kv_cache_matches_torch still passes its diff == 0.0 requirement

Benchmark: NVIDIA L20, CUDA 12.8, torch 2.11.0+cu128, bf16 decode microbench (GQA 32/8 heads, head_dim 128, batch 32, seq_lens 200-2000), 20 trials
- per-layer decode attention: 117.5 -> 9.5 ms (~12x)
- max abs diff vs the per-request loop: 0.002 (bf16 noise; the decode invariant tolerates 0.05)
- tests: extension + inference suites 350 passed

astrai/extension/backend/attention.py

@@ -459,6 +459,10 @@ class TorchNativeBackend(AttentionBackend):
     via ``req_to_token`` indirect indexing, then calls
     ``F.scaled_dot_product_attention``.
 
+    Packed inference (3-D q) pads the ragged batch to [B, max_q, max_kv]
+    and runs a single batched SDPA call with a combined causal+padding
+    mask, then unpacks back to the flat layout.
+
     For training (``kv_cache is None``), skips cache I/O entirely and
     runs SDPA directly on the projected q/k/v.
     """
@@ -510,31 +514,46 @@ class TorchNativeBackend(AttentionBackend):
             raise ValueError("packed attention requires KV cache metadata")
         kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
         kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
-        outputs = []
+
+        # Pad the ragged batch to [B, max_q, max_kv] so one batched SDPA call
+        # replaces B per-request calls. The bool mask folds the per-request
+        # causal offset (seq_len - q_len) and the kv padding together; padded
+        # q rows gather slot/row 0 and are dropped by the [q_valid] unpack,
+        # which restores the packed qo_indptr order.
+        qo_indptr = kv_cache.qo_indptr
+        q_lens = qo_indptr[1:] - qo_indptr[:-1]
+        seq_lens = kv_cache.seq_lens
+        max_q = int(q_lens.max())
+        max_kv = int(seq_lens.max())
+
+        pos = torch.arange(max_kv, device=q.device)
+        kv_valid = pos.unsqueeze(0) < seq_lens.unsqueeze(1)
+        token_index = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_kv]
+        token_index = token_index.masked_fill(~kv_valid, 0)
+        k_b = kv_cache.k_buffer[layer_id, token_index]
+        v_b = kv_cache.v_buffer[layer_id, token_index]
         n_rep = q.size(1) // k.size(1)
-        for i in range(kv_cache.req_pool_indices.numel()):
-            q_start = int(kv_cache.qo_indptr[i])
-            q_end = int(kv_cache.qo_indptr[i + 1])
-            indices = kv_cache.req_to_token[
-                kv_cache.req_pool_indices[i], : kv_cache.seq_lens[i]
-            ]
-            k_i = kv_cache.k_buffer[layer_id, indices]
-            v_i = kv_cache.v_buffer[layer_id, indices]
         if n_rep > 1:
-                k_i = repeat_kv(k_i, n_rep)
+            k_b = repeat_kv(k_b, n_rep)
-                v_i = repeat_kv(v_i, n_rep)
+            v_b = repeat_kv(v_b, n_rep)
-            q_len = q_end - q_start
+
-            kv_len = k_i.size(0)
+        q_pos = torch.arange(max_q, device=q.device)
-            q_pos = torch.arange(kv_len - q_len, kv_len, device=q.device)
+        q_valid = q_pos.unsqueeze(0) < q_lens.unsqueeze(1)
-            causal_mask = q_pos[:, None] >= torch.arange(kv_len, device=q.device)
+        q_index = (qo_indptr[:-1].unsqueeze(1) + q_pos.unsqueeze(0)).masked_fill(
-            out = F.scaled_dot_product_attention(
+            ~q_valid, 0
-                q[q_start:q_end].transpose(0, 1).unsqueeze(0),
-                k_i.transpose(0, 1).unsqueeze(0),
-                v_i.transpose(0, 1).unsqueeze(0),
-                attn_mask=causal_mask,
         )
-            outputs.append(out.squeeze(0).transpose(0, 1))
+        causal = (
-        return torch.cat(outputs)
+            (seq_lens - q_lens).unsqueeze(1).unsqueeze(2) + q_pos.view(1, max_q, 1)
+        ) >= pos.view(1, 1, max_kv)
+        mask = (causal & kv_valid.unsqueeze(1)).unsqueeze(1)
+
+        out = F.scaled_dot_product_attention(
+            q[q_index].permute(0, 2, 1, 3),
+            k_b.permute(0, 2, 1, 3),
+            v_b.permute(0, 2, 1, 3),
+            attn_mask=mask,
+        )
+        return out.permute(0, 2, 1, 3)[q_valid]
 
 
 @AttentionBackendFactory.register(ATTN_BACKEND.CUDA.value)