iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm] [CI] Support cached K/V (key/value=None) in Triton prefix-prefill (#48257)

Browse Source 
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
stefankoncarevic
2026-07-29 23:31:51 -05:00
committed by GitHub
co-authored by Andreas Karatzas
parent 1ad5182ba9
commit b88916617d
2 changed files with 525 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_prefix_prefill.py
+336
View File
@@ -323,6 +323,342 @@ def test_contexted_kv_attention(
torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_contexted_kv_attention_cached_kv(
num_heads: int,
num_queries_per_kv: int,
head_size: int,
dtype: torch.dtype,
kv_cache_dtype: str,
device: str,
block_size: int = 32,
) -> None:
# Exercises the KV_FROM_CACHE path of context_attention_fwd: the current
# chunk K/V are not passed as dense tensors (k=v=None); they are read back
# from the paged KV cache, as done by layers that re-attend an already
# cached sequence with the query only (e.g. IQuest LoopCoder's
# attn(q, None, None)). The whole sequence therefore lives in the cache and
# the result must still match a dense causal SDPA reference.
if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):
pytest.skip(
"Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
)
set_random_seed(0)
torch.set_default_device(device)
torch.accelerator.set_device_index(device)
MAX_SEQ_LEN = 1024
MAX_CTX_LEN = 1024
BS = 10
cache_size = 640
max_block_per_request = 64
query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]
# ensure one sequence in batch is a decode
query_lens[-1] = 1
ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]
seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
num_kv_heads = num_heads // num_queries_per_kv
num_tokens = sum(query_lens)
query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
query.uniform_(-1e-3, 1e-3)
output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
kv = torch.empty(sum(seq_lens), 2, num_kv_heads, head_size, dtype=dtype)
kv.uniform_(-1e-3, 1e-3)
key, value = kv.unbind(dim=1)
if kv_cache_dtype == "auto":
cache_dtype = dtype
else:
cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]
k_cache = torch.zeros(
cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
)
v_cache = torch.zeros(
cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype
)
values = torch.arange(0, cache_size, dtype=torch.int32)
values = values[torch.randperm(cache_size)]
block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)
b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)
max_input_len = MAX_SEQ_LEN
b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(
torch.int32
)
# Unlike the dense test, write the WHOLE sequence (context + current chunk)
# into the paged cache, since the current chunk is read back from the cache.
for i in range(BS):
cur = 0
block_id = 0
while cur < seq_lens[i]:
start_loc = b_seq_start_loc[i] + cur
if cur + block_size > seq_lens[i]:
end_loc = b_seq_start_loc[i] + seq_lens[i]
else:
end_loc = start_loc + block_size
start_slot = block_table[i, block_id] * block_size
end_slot = start_slot + end_loc - start_loc
k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
key[start_loc:end_loc]
)
v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
value[start_loc:end_loc]
)
cur += block_size
block_id += 1
# transpose to the paged cache layouts the kernel expects:
# K_cache[num_blocks, num_kv_heads, head_size/8, block_size, 8]
# V_cache[num_blocks, num_kv_heads, head_size, block_size]
k_cache = (
k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)
.permute(0, 2, 3, 1, 4)
.contiguous()
)
v_cache = (
v_cache.view(-1, block_size, num_kv_heads, head_size)
.permute(0, 2, 3, 1)
.contiguous()
)
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
# Cached-K/V path: current-chunk k and v are None.
context_attention_fwd(
query,
None,
None,
output,
kv_cache_dtype,
k_cache,
v_cache,
block_table,
b_start_loc,
b_seq_len,
MAX_CTX_LEN,
max_input_len,
k_scale,
v_scale,
sliding_window=0,
)
torch.accelerator.synchronize()
scale = float(1.0 / (head_size**0.5))
query_sdpa = query.view(num_tokens, num_kv_heads, num_queries_per_kv, head_size)
query_sdpa = query_sdpa.permute(1, 2, 0, 3).reshape(
1, num_heads, num_tokens, head_size
)
key_sdpa = key[:, :, None, :].expand(
key.shape[0], num_kv_heads, num_queries_per_kv, key.shape[-1]
)
key_sdpa = key_sdpa.permute(1, 2, 0, 3).reshape(
1, num_heads, sum(seq_lens), head_size
)
value_sdpa = value[:, :, None, :].expand(
value.shape[0], num_kv_heads, num_queries_per_kv, value.shape[-1]
)
value_sdpa = value_sdpa.permute(1, 2, 0, 3).reshape(
1, num_heads, sum(seq_lens), head_size
)
attn_mask = create_causal_attention_mask_for_sdpa(
query_lens, seq_lens, 0, device=device, dtype=dtype
)
output_ref = F.scaled_dot_product_attention(
query_sdpa,
key_sdpa,
value_sdpa,
attn_mask=attn_mask,
dropout_p=0.0,
scale=scale,
)
output_ref = output_ref.view(num_heads, num_tokens, head_size)
output_ref = output_ref.permute(1, 0, 2).contiguous()
atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-4
torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> None:
# Boundary guard for the KV_FROM_CACHE block-table load. With an
# exact-sized block table (row length == number of blocks for the
# sequence) and a query that ends the sequence, the last K/V tile has
# padded lanes whose absolute positions step past the sequence end. Those
# lanes must not read a block-table entry past this batch's row.
#
# seq_len is a whole number of blocks, so bn_logical for a padded lane
# lands exactly on num_blocks (one past the last valid row entry), and
# query_len is not tile-aligned so the overshoot is actually exercised.
# The result must still match a dense causal SDPA reference.
set_random_seed(0)
torch.set_default_device(device)
torch.accelerator.set_device_index(device)
dtype = torch.float16
kv_cache_dtype = "auto"
num_heads = 4
num_queries_per_kv = 1
num_kv_heads = num_heads // num_queries_per_kv
head_size = 32
block_size = 32
num_blocks = 4
seq_len = num_blocks * block_size # 128 == exact multiple of block_size
query_len = 99 # < seq_len and not a multiple of the kernel tile
query = torch.empty(query_len, num_heads, head_size, dtype=dtype)
query.uniform_(-1e-3, 1e-3)
output = torch.empty(query_len, num_heads, head_size, dtype=dtype)
kv = torch.empty(seq_len, 2, num_kv_heads, head_size, dtype=dtype)
kv.uniform_(-1e-3, 1e-3)
key, value = kv.unbind(dim=1)
k_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype)
v_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype)
# Exact-sized block table: row length == num_blocks (identity mapping).
block_table = torch.arange(num_blocks, dtype=torch.int32).view(1, num_blocks)
b_seq_len = torch.tensor([seq_len], dtype=torch.int32)
b_start_loc = torch.tensor([0, query_len], dtype=torch.int32)
# Write the whole sequence (context + current chunk) into the paged cache.
for cur in range(0, seq_len, block_size):
block_id = cur // block_size
end = min(cur + block_size, seq_len)
start_slot = block_table[0, block_id] * block_size
end_slot = start_slot + (end - cur)
k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
key[cur:end]
)
v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(
value[cur:end]
)
k_cache = (
k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)
.permute(0, 2, 3, 1, 4)
.contiguous()
)
v_cache = (
v_cache.view(-1, block_size, num_kv_heads, head_size)
.permute(0, 2, 3, 1)
.contiguous()
)
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
# Cached-K/V path: current-chunk k and v are None.
context_attention_fwd(
query,
None,
None,
output,
kv_cache_dtype,
k_cache,
v_cache,
block_table,
b_start_loc,
b_seq_len,
seq_len,
query_len,
k_scale,
v_scale,
sliding_window=0,
)
torch.accelerator.synchronize()
scale = float(1.0 / (head_size**0.5))
query_sdpa = query.view(query_len, num_kv_heads, num_queries_per_kv, head_size)
query_sdpa = query_sdpa.permute(1, 2, 0, 3).reshape(
1, num_heads, query_len, head_size
)
key_sdpa = key[:, :, None, :].expand(
seq_len, num_kv_heads, num_queries_per_kv, head_size
)
key_sdpa = key_sdpa.permute(1, 2, 0, 3).reshape(1, num_heads, seq_len, head_size)
value_sdpa = value[:, :, None, :].expand(
seq_len, num_kv_heads, num_queries_per_kv, head_size
)
value_sdpa = value_sdpa.permute(1, 2, 0, 3).reshape(
1, num_heads, seq_len, head_size
)
attn_mask = create_causal_attention_mask_for_sdpa(
[query_len], [seq_len], 0, device=device, dtype=dtype
)
output_ref = F.scaled_dot_product_attention(
query_sdpa,
key_sdpa,
value_sdpa,
attn_mask=attn_mask,
dropout_p=0.0,
scale=scale,
)
output_ref = output_ref.view(num_heads, query_len, head_size)
output_ref = output_ref.permute(1, 0, 2).contiguous()
torch.testing.assert_close(output, output_ref, atol=1e-4, rtol=0)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_contexted_kv_attention_cached_kv_alibi_unsupported(device: str) -> None:
# The cached-K/V (k=None) path is not supported together with ALiBi; the
# entry point must reject it up-front with a clear NotImplementedError
# rather than launching the kernel with an unsupported combination.
set_random_seed(0)
torch.set_default_device(device)
torch.accelerator.set_device_index(device)
num_heads = 4
num_kv_heads = 4
head_size = 16
x = 8
block_size = 16
num_blocks = 4
query_len = 8
query = torch.empty(query_len, num_heads, head_size, dtype=torch.float16)
query.uniform_(-1e-3, 1e-3)
output = torch.empty_like(query)
k_cache = torch.zeros(
num_blocks, num_kv_heads, head_size // x, block_size, x, dtype=torch.float16
)
v_cache = torch.zeros(
num_blocks, num_kv_heads, head_size, block_size, dtype=torch.float16
)
block_table = torch.arange(num_blocks, dtype=torch.int32).view(1, num_blocks)
b_seq_len = torch.tensor([query_len], dtype=torch.int32)
b_start_loc = torch.tensor([0, query_len], dtype=torch.int32)
k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
alibi_slopes = torch.ones(num_heads, dtype=torch.float32, device=device)
with pytest.raises(NotImplementedError):
context_attention_fwd(
query,
None,
None,
output,
"auto",
k_cache,
v_cache,
block_table,
b_start_loc,
b_seq_len,
query_len,
query_len,
k_scale,
v_scale,
alibi_slopes=alibi_slopes,
)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
vllm/v1/attention/ops/prefix_prefill.py
+189 -40
View File
@@ -35,6 +35,74 @@ float8_info = torch.finfo(current_platform.fp8_dtype())
# ],
# key=["BLOCK_SIZE", "MAX_Q_LEN", "MAX_CTX_LEN"]
# )
@triton.jit
def _paged_kv_cache_offsets(
B_Loc,
cur_batch,
token_indices,
token_valid,
offs_d,
cur_kv_head,
x,
stride_b_loc_b,
stride_b_loc_s,
stride_k_cache_bs,
stride_k_cache_h,
stride_k_cache_d,
stride_k_cache_bl,
stride_k_cache_x,
stride_v_cache_bs,
stride_v_cache_h,
stride_v_cache_d,
stride_v_cache_bl,
PHYSICAL_BLOCK_SIZE: tl.constexpr,
MASK_BLOCK_TABLE: tl.constexpr = False,
):
"""Compute paged K/V cache element offsets for a tile of token positions.
`token_indices` holds the absolute sequence positions of the tile. The
logical block for each token is looked up in `B_Loc` (handling cross-block
tiles), then the physical element offsets are built with the same layout the
kernel uses everywhere:
K cache: [num_blocks, num_kv_heads, head_size // x, block_size, x]
V cache: [num_blocks, num_kv_heads, head_size, block_size]
Returns `(off_k, off_v)` with shapes [D, N] and [N, D] respectively.
When `MASK_BLOCK_TABLE` is set, the block-table load is masked with
`token_valid` so padded lanes (a tile that extends past the end of the
sequence) cannot read a block-table entry past this batch's row. Invalid
lanes resolve to block 0 (always in-bounds) and are dropped by the caller's
later K/V load mask. When it is not set, `token_valid` is unused and the
load is emitted exactly as before.
"""
bn_logical = token_indices // PHYSICAL_BLOCK_SIZE
if MASK_BLOCK_TABLE:
bn = tl.load(
B_Loc + cur_batch * stride_b_loc_b + bn_logical * stride_b_loc_s,
mask=token_valid,
other=0,
).to(tl.int64)
else:
bn = tl.load(
B_Loc + cur_batch * stride_b_loc_b + bn_logical * stride_b_loc_s
).to(tl.int64)
internal = token_indices % PHYSICAL_BLOCK_SIZE
off_k = (
bn[None, :] * stride_k_cache_bs
+ cur_kv_head * stride_k_cache_h
+ (offs_d[:, None] // x) * stride_k_cache_d
+ internal[None, :] * stride_k_cache_bl
+ (offs_d[:, None] % x) * stride_k_cache_x
)
off_v = (
bn[:, None] * stride_v_cache_bs
+ cur_kv_head * stride_v_cache_h
+ offs_d[None, :] * stride_v_cache_d
+ internal[:, None] * stride_v_cache_bl
)
return off_k, off_v
@triton.jit
def _fwd_kernel(
Q,
@@ -94,6 +162,11 @@ def _fwd_kernel(
MAX_CTX_LEN: tl.constexpr = 0,
FP8_MIN: tl.constexpr = float8_info.min,
FP8_MAX: tl.constexpr = float8_info.max,
# When True, the current-chunk K/V are not provided as dense tensors
# (`K`/`V` are unused placeholders); instead they are read back from the
# paged KV cache. This supports layers that re-attend an already-cached
# sequence with query only (e.g. IQuest LoopCoder's `attn(q, None, None)`).
KV_FROM_CACHE: tl.constexpr = False,
):
cur_batch = tl.program_id(0)
cur_head = tl.program_id(1)
@@ -162,33 +235,34 @@ def _fwd_kernel(
# replace one physical block every 17 32-Tile blocks
# Calculate the logical block index of each of the 32 tokens
# in the current Tile (handling cross-block cases).
# Physical-block K/V cache offsets for the context tokens in this tile
# (handles cross-block tiles via B_Loc). Same addressing is reused for
# the current-chunk cache read below.
token_indices = start_n + offs_bs_n
bn_logical_indices = token_indices // PHYSICAL_BLOCK_SIZE
# 2. Vectorized loading of physical block IDs from B_Loc
bn = tl.load(
B_Loc + cur_batch * stride_b_loc_b + bn_logical_indices * stride_b_loc_s
).to(tl.int64)
# 3. Calculate the exact offset of
# each token within its physical block.
internal_offsets = token_indices % PHYSICAL_BLOCK_SIZE
# Addressing of K (5D)
off_k = (
bn[None, :] * stride_k_cache_bs
+ cur_kv_head * stride_k_cache_h
+ (offs_d[:, None] // x) * stride_k_cache_d
+ internal_offsets[None, :] * stride_k_cache_bl
+ (offs_d[:, None] % x) * stride_k_cache_x
)
# Addressing of V (4D)
off_v = (
bn[:, None] * stride_v_cache_bs
+ cur_kv_head * stride_v_cache_h
+ offs_d[None, :] * stride_v_cache_d
+ internal_offsets[:, None] * stride_v_cache_bl
# Context tokens are always followed by the current chunk, so the tile
# never steps past this batch's block-table row; the block-table load
# stays unmasked (MASK_BLOCK_TABLE=False) and offs_bs_n is an unused
# placeholder for token_valid.
off_k, off_v = _paged_kv_cache_offsets(
B_Loc,
cur_batch,
token_indices,
offs_bs_n,
offs_d,
cur_kv_head,
x,
stride_b_loc_b,
stride_b_loc_s,
stride_k_cache_bs,
stride_k_cache_h,
stride_k_cache_d,
stride_k_cache_bl,
stride_k_cache_x,
stride_v_cache_bs,
stride_v_cache_h,
stride_v_cache_d,
stride_v_cache_bl,
PHYSICAL_BLOCK_SIZE,
)
if (
@@ -300,12 +374,54 @@ def _fwd_kernel(
):
start_n = tl.multiple_of(start_n, BLOCK_N)
# -- compute qk ----
k = tl.load(
k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
mask=dim_mask[:, None]
& ((start_n + offs_n[None, :]) < cur_batch_query_len),
other=0.0,
)
if KV_FROM_CACHE:
# Current-chunk K/V are not passed densely; read them from the paged
# cache at their absolute positions [ctx_len, ctx_len + query_len),
# reusing the same addressing as the context loop above.
cache_token_idx = cur_batch_ctx_len + start_n + offs_n
# Padded lanes of the last tile can index past the sequence end;
# mask the block-table load so they don't read past this batch's
# row (they are dropped by the K/V load mask below anyway).
cache_token_valid = (start_n + offs_n) < cur_batch_query_len
off_k_cur, off_v_cur = _paged_kv_cache_offsets(
B_Loc,
cur_batch,
cache_token_idx,
cache_token_valid,
offs_d,
cur_kv_head,
x,
stride_b_loc_b,
stride_b_loc_s,
stride_k_cache_bs,
stride_k_cache_h,
stride_k_cache_d,
stride_k_cache_bl,
stride_k_cache_x,
stride_v_cache_bs,
stride_v_cache_h,
stride_v_cache_d,
stride_v_cache_bl,
PHYSICAL_BLOCK_SIZE,
MASK_BLOCK_TABLE=True,
)
k_cur_load = tl.load(
K_cache + off_k_cur,
mask=dim_mask[:, None]
& ((start_n + offs_n[None, :]) < cur_batch_query_len),
other=0.0,
)
if k_cur_load.dtype.is_fp8():
k = (k_cur_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)
else:
k = k_cur_load
else:
k = tl.load(
k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
mask=dim_mask[:, None]
& ((start_n + offs_n[None, :]) < cur_batch_query_len),
other=0.0,
)
qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)
@@ -333,12 +449,25 @@ def _fwd_kernel(
acc = acc * alpha[:, None]
# update acc
v = tl.load(
v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
mask=dim_mask[None, :]
& ((start_n + offs_n[:, None]) < cur_batch_query_len),
other=0.0,
)
if KV_FROM_CACHE:
# off_v_cur was computed together with off_k_cur above.
v_cur_load = tl.load(
V_cache + off_v_cur,
mask=dim_mask[None, :]
& ((start_n + offs_n[:, None]) < cur_batch_query_len),
other=0.0,
)
if v_cur_load.dtype.is_fp8():
v = (v_cur_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)
else:
v = v_cur_load
else:
v = tl.load(
v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
mask=dim_mask[None, :]
& ((start_n + offs_n[:, None]) < cur_batch_query_len),
other=0.0,
)
p = p.to(v.dtype)
acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION)
@@ -703,8 +832,27 @@ def context_attention_fwd(
FP8 KV Cache prefill kernel"
)
assert (k is None) == (v is None), (
"k and v must both be None (cached-K/V path) or both be tensors"
)
kv_from_cache = k is None
if kv_from_cache:
if alibi_slopes is not None:
raise NotImplementedError(
"context_attention_fwd with cached K/V (key=None) is not "
"supported together with ALiBi slopes."
)
# k_cache: [num_blocks, num_kv_heads, head_size // x, block_size, x]
num_kv_heads = k_cache.shape[1]
Lq = Lk = Lv = q.shape[-1]
# q is a never-dereferenced placeholder for the dense K/V pointer args.
k = q
v = q
else:
num_kv_heads = k.shape[1]
Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
# shape constraints
Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
assert Lq == Lk and Lk == Lv
# round up Lk to a power of 2 - this is required for Triton block size
Lk_padded = triton.next_power_of_2(Lk)
@@ -712,7 +860,7 @@ def context_attention_fwd(
if sm_scale is None:
sm_scale = 1.0 / (Lq**0.5)
batch, head = b_seq_len.shape[0], q.shape[1]
num_queries_per_kv = q.shape[1] // k.shape[1]
num_queries_per_kv = q.shape[1] // num_kv_heads
assert batch + 1 == len(b_start_loc)
@@ -873,6 +1021,7 @@ def context_attention_fwd(
num_stages=1,
USE_SINKS=sinks is not None,
CAUSAL=causal,
KV_FROM_CACHE=kv_from_cache,
**extra_kargs,
)
return