diff --git a/astrai/extension/backend/attention.py b/astrai/extension/backend/attention.py index 7f2f327..782e633 100644 --- a/astrai/extension/backend/attention.py +++ b/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) - v_i = repeat_kv(v_i, n_rep) - q_len = q_end - q_start - kv_len = k_i.size(0) - q_pos = torch.arange(kv_len - q_len, kv_len, device=q.device) - causal_mask = q_pos[:, None] >= torch.arange(kv_len, device=q.device) - out = F.scaled_dot_product_attention( - 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)) - return torch.cat(outputs) + if n_rep > 1: + k_b = repeat_kv(k_b, n_rep) + v_b = repeat_kv(v_b, n_rep) + + q_pos = torch.arange(max_q, device=q.device) + q_valid = q_pos.unsqueeze(0) < q_lens.unsqueeze(1) + q_index = (qo_indptr[:-1].unsqueeze(1) + q_pos.unsqueeze(0)).masked_fill( + ~q_valid, 0 + ) + causal = ( + (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)