ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: standardize packed 3d inference

Browse Source 
- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch

Benchmark: NVIDIA L20, BF16, 1B model, paged KV cache, CUDA Graph, prompt 512, generation 128 (median of 3 alternating runs)
- batch 1: 234.5 -> 242.6 tok/s (1.034x, +3.4%)
- batch 8: 1243.1 -> 1286.6 tok/s (1.035x, +3.5%)
This commit is contained in:
ViperEkura
2026-08-19 00:36:53 +08:00
parent f7f14d0e5f
commit 3d3ea47d37
9 changed files with 151 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/backend/attention.py
+2 -26
View File
@@ -253,28 +253,6 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
)
def _write_and_gather_kv(
kv_cache: "KVCache",
k: Tensor,
v: Tensor,
layer_id: int,
q: Tensor,
attn_mask: Optional[Tensor],
) -> tuple[Tensor, Tensor]:
kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
max_len = kv_cache.max_len
indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
if q.size(1) == 1 and attn_mask is not None and attn_mask.dim() == 4:
pos_mask = attn_mask[:, 0, 0]
else:
pos_mask = (
torch.arange(max_len, device=q.device)[None, :] < kv_cache.seq_lens[:, None]
)
indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
return kv_cache.k_buffer[layer_id, indices], kv_cache.v_buffer[layer_id, indices]
def attention(
q: Tensor,
k: Tensor,
@@ -565,10 +543,6 @@ class CudaBackend(AttentionBackend):
if kv_cache is None:
raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
loc = kv_cache.out_cache_loc
kv_cache.k_buffer[layer_id, loc] = k
kv_cache.v_buffer[layer_id, loc] = v
kv_indptr = kv_cache.kv_indptr
out = attn_paged_decode(
@@ -578,6 +552,8 @@ class CudaBackend(AttentionBackend):
kv_cache.req_to_token,
kv_cache.req_pool_indices,
kv_indptr,
new_k=k,
new_v=v,
is_causal=True,
o_part_buf=kv_cache.decode_o_part,
ml_part_buf=kv_cache.decode_ml_part,
astrai/extension/ops/attention.py
+6
View File
@@ -97,6 +97,8 @@ def attn_paged_decode(
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
kv_indptr: torch.Tensor,
new_k: Optional[torch.Tensor] = None,
new_v: Optional[torch.Tensor] = None,
mask: Optional[torch.Tensor] = None,
is_causal: bool = False,
o_part_buf: Optional[torch.Tensor] = None,
@@ -116,6 +118,8 @@ def attn_paged_decode(
req_to_token: [num_reqs, max_context_len] (int32) — token -> slot
req_pool_indices: [batch] (int32) — rows into req_to_token
kv_indptr: [batch+1] (int32) — prefix sum of per-request seq_lens
new_k: current-token K to append, [batch, n_kv_heads, head_dim]
new_v: current-token V to append, same shape as new_k
mask: 2D [batch, max_context_len] (bool, True=keep) or None
is_causal: apply causal mask
o_part_buf: pre-allocated split-KV o partial buffer (workflow bypass)
@@ -134,6 +138,8 @@ def attn_paged_decode(
req_to_token,
req_pool_indices,
kv_indptr,
new_k=new_k,
new_v=new_v,
mask=mask,
causal_offset=causal_offset,
o_part_buf=o_part_buf,
csrc/kernels/attn_common.h
+5
View File
@@ -36,6 +36,8 @@ struct AttentionParams {
const T* __restrict__ q_ptr;
const T* __restrict__ k_ptr;
const T* __restrict__ v_ptr;
const T* __restrict__ new_k_ptr;
const T* __restrict__ new_v_ptr;
T* __restrict__ o_ptr;
const bool* __restrict__ mask;
@@ -50,6 +52,9 @@ struct AttentionParams {
int kv_l_stride;
int kv_d_stride;
int new_kv_b_stride;
int new_kv_h_stride;
int mask_b_stride;
int mask_h_stride;
int mask_l_stride;
csrc/kernels/attn_decode_split_kv.cuh
+2 -2
View File
@@ -57,8 +57,8 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
int s = i / p.head_dim;
int d_dim = i % p.head_dim;
int kc = chunk_start + s;
int token = KV::resolve_token(p, kctx, kc, true);
KVAddr a = KV::kv_addr_from_token(p, kctx, token, d_dim);
KVAddr a = KV::template decode_addr<1>(
p, kctx, batch, kv_head, kc, d_dim, true, true);
k_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.k) : (bf16)0.f;
v_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.v) : (bf16)0.f;
}
csrc/kernels/attn_decode_split_kv_mma.cuh
+4 -2
View File
@@ -73,8 +73,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
int kc = kv0 + r;
bool valid = kc < seq_len;
int token = KV::resolve_token(p, kctx, kc, valid);
KVAddr a = KV::kv_addr_from_token(p, kctx, token, d);
// All GQA passes consume new K/V directly. Only the first pass
// persists it, so no cross-block synchronization is required.
KVAddr a = KV::template decode_addr<Traits::VEC>(
p, kctx, batch, kv_head, kc, d, valid, pass == 0);
int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
cp_async_16_pred(&dK[off], a.k, a.valid);
cp_async_16_pred(&dV[off], a.v, a.valid);
csrc/kernels/attn_entry_utils.cuh
+33
View File
@@ -130,6 +130,8 @@ inline void attn_pack_params(
p.q_ptr = (const T*)q.data_ptr();
p.k_ptr = (const T*)k.data_ptr();
p.v_ptr = (const T*)v.data_ptr();
p.new_k_ptr = nullptr;
p.new_v_ptr = nullptr;
p.o_ptr = nullptr;
p.o_part = nullptr;
p.ml_part = nullptr;
@@ -148,6 +150,8 @@ inline void attn_pack_paged_decode_params(
torch::Tensor req_to_token,
torch::Tensor req_pool_indices,
torch::Tensor kv_indptr,
const c10::optional<torch::Tensor>& new_k,
const c10::optional<torch::Tensor>& new_v,
c10::optional<torch::Tensor> mask,
int64_t causal_offset,
double scale,
@@ -191,6 +195,33 @@ inline void attn_pack_paged_decode_params(
p.qo_indptr = nullptr;
p.max_context_len = (int)req_to_token.size(1);
TORCH_CHECK(new_k.has_value() == new_v.has_value(),
"new_k and new_v must be provided together");
if (new_k.has_value()) {
auto nk = new_k.value();
auto nv = new_v.value();
TORCH_CHECK(nk.is_cuda() && nv.is_cuda(), "new K/V must be CUDA tensors");
TORCH_CHECK(nk.dtype() == torch::kBFloat16 && nv.dtype() == torch::kBFloat16,
"new K/V must be bf16");
TORCH_CHECK(nk.dim() == 3 && nv.dim() == 3,
"new K/V must be 3D [batch, kv_head, head_dim]");
TORCH_CHECK(nk.sizes() == nv.sizes(), "new K and V must have identical shapes");
TORCH_CHECK(nk.strides() == nv.strides(),
"new K and V must have identical strides");
TORCH_CHECK(nk.size(0) == p.batch && nk.size(1) == p.kv_head
&& nk.size(2) == p.head_dim, "new K/V shape mismatch");
TORCH_CHECK(nk.stride(2) == 1 && nv.stride(2) == 1,
"new K/V head_dim must be contiguous");
p.new_k_ptr = (const T*)nk.data_ptr();
p.new_v_ptr = (const T*)nv.data_ptr();
p.new_kv_b_stride = (int)nk.stride(0);
p.new_kv_h_stride = (int)nk.stride(1);
} else {
p.new_k_ptr = nullptr;
p.new_v_ptr = nullptr;
p.new_kv_b_stride = p.new_kv_h_stride = 0;
}
p.causal_offset = (int)causal_offset;
p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
@@ -279,6 +310,8 @@ inline void attn_pack_paged_prefill_params(
p.k_ptr = (const T*)k_cache.data_ptr();
p.v_ptr = (const T*)v_cache.data_ptr();
p.new_k_ptr = nullptr;
p.new_v_ptr = nullptr;
p.q_ptr = (const T*)q.data_ptr();
p.req_to_token = req_to_token.data_ptr<int>();
p.req_pool_indices = req_pool_indices.data_ptr<int>();
csrc/kernels/attn_layout_policies.cuh
+43
View File
@@ -160,6 +160,14 @@ struct ContigKV {
+ (int64_t)d * p.kv_d_stride;
return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid};
}
template <int VEC>
DEVICE_FORCEINLINE KVAddr decode_addr(
const AttentionParams<bf16>& p, const KVContext& c,
int, int, int kc, int d, bool valid, bool) {
int token = resolve_token(p, c, kc, valid);
return kv_addr_from_token(p, c, token, d);
}
};
// ---- Paged (SGLang-style flat pool) K/V ----
@@ -209,4 +217,39 @@ struct PagedKV {
const int64_t gmem_off = (int64_t)safe_slot * c.pool_stride + c.head_off + d;
return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid};
}
DEVICE_FORCEINLINE KVAddr new_kv_addr(
const AttentionParams<bf16>& p, int batch, int kv_head, int d) {
const int64_t off = (int64_t)batch * p.new_kv_b_stride
+ (int64_t)kv_head * p.new_kv_h_stride + d;
return {&p.new_k_ptr[off], &p.new_v_ptr[off], true};
}
DEVICE_FORCEINLINE void store_new_kv(
const AttentionParams<bf16>& p, const KVContext& c,
int seq_len, int d, const KVAddr& src) {
int slot = resolve_token(p, c, seq_len - 1, true);
const int64_t off = (int64_t)slot * c.pool_stride + c.head_off + d;
const_cast<bf16*>(p.k_ptr)[off] = *reinterpret_cast<const bf16*>(src.k);
const_cast<bf16*>(p.v_ptr)[off] = *reinterpret_cast<const bf16*>(src.v);
}
template <int VEC>
DEVICE_FORCEINLINE KVAddr decode_addr(
const AttentionParams<bf16>& p, const KVContext& c,
int batch, int kv_head, int kc, int d, bool valid, bool persist) {
if (p.new_k_ptr && valid && kc == kv_len(p, batch) - 1) {
KVAddr src = new_kv_addr(p, batch, kv_head, d);
if (persist) {
#pragma unroll
for (int j = 0; j < VEC; j++) {
KVAddr value = new_kv_addr(p, batch, kv_head, d + j);
store_new_kv(p, c, kc + 1, d + j, value);
}
}
return src;
}
int token = resolve_token(p, c, kc, valid);
return kv_addr_from_token(p, c, token, d);
}
};
csrc/kernels/attn_paged_decode.cu
+5
View File
@@ -8,6 +8,8 @@ torch::Tensor attn_paged_decode(
torch::Tensor req_to_token,
torch::Tensor req_pool_indices,
torch::Tensor kv_indptr,
c10::optional<torch::Tensor> new_k,
c10::optional<torch::Tensor> new_v,
c10::optional<torch::Tensor> mask,
int64_t causal_offset,
double scale,
@@ -21,6 +23,7 @@ torch::Tensor attn_paged_decode(
AttentionParams<bf16> p;
attn_pack_paged_decode_params(q, k_cache, v_cache,
req_to_token, req_pool_indices, kv_indptr,
new_k, new_v,
mask, causal_offset, scale, p);
torch::Tensor O;
@@ -71,6 +74,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::arg("req_to_token"),
py::arg("req_pool_indices"),
py::arg("kv_indptr"),
py::arg("new_k") = py::none(),
py::arg("new_v") = py::none(),
py::arg("mask") = py::none(),
py::arg("causal_offset") = -1,
py::arg("scale") = 0.0,
tests/extension/test_backend_equivalence.py
+51
View File
@@ -7,6 +7,7 @@ seq_lens with padding mask), and end-to-end scheduler.run_batch.
import torch
from astrai.extension import ATTN_BACKEND, attn_backend
from astrai.extension.ops.attention import attn_paged_decode
from astrai.inference.cache import PagePool, TaskCacheManager
from astrai.inference.runtime.graph import CudaGraphContext
from astrai.inference.scheduler import InferenceScheduler
@@ -160,6 +161,56 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
@skip_no_kernel
def test_paged_decode_appends_new_kv_in_kernel():
"""Fused decode writes current-token K/V to each request's paged slot."""
pool = PagePool(
n_layers=1,
n_kv_heads=1,
head_dim=D,
max_batch_size=2,
max_seq_len=64,
device="cuda",
dtype=torch.bfloat16,
page_size=8,
n_tokens=128,
)
task_cache = _mk_task_cache(pool)
ws = _ws(pool)
task_cache.task_alloc("t1", list(range(8)))
task_cache.task_alloc("t2", list(range(6)))
task_cache.task_extend("t1", 8)
task_cache.task_extend("t2", 6)
kv_cache = task_cache.bind(["t1", "t2"], ws)
q = torch.randn(2, 2, D, device="cuda", dtype=torch.bfloat16)
new_k = torch.randn(2, 1, D, device="cuda", dtype=torch.bfloat16)
new_v = torch.randn(2, 1, D, device="cuda", dtype=torch.bfloat16)
out = attn_paged_decode(
q,
kv_cache.k_buffer[0],
kv_cache.v_buffer[0],
kv_cache.req_to_token,
kv_cache.req_pool_indices,
kv_cache.kv_indptr,
new_k=new_k,
new_v=new_v,
is_causal=True,
o_part_buf=kv_cache.decode_o_part,
ml_part_buf=kv_cache.decode_ml_part,
out_buf=kv_cache.decode_out,
)
torch.cuda.synchronize()
torch.testing.assert_close(
kv_cache.k_buffer[0, kv_cache.out_cache_loc], new_k, rtol=0, atol=0
)
torch.testing.assert_close(
kv_cache.v_buffer[0, kv_cache.out_cache_loc], new_v, rtol=0, atol=0
)
assert torch.isfinite(out).all()
@skip_no_kernel
def test_decode_cuda_graph_replay_is_exact(cuda_model):
"""INT32 cache indices must remain graph-capturable and replay exactly."""