From b88916617d3d2249bff0dae5cecb6b727c980a20 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Thu, 30 Jul 2026 06:31:51 +0200 Subject: [PATCH] [ROCm] [CI] Support cached K/V (key/value=None) in Triton prefix-prefill (#48257) Signed-off-by: Stefan Koncarevic Co-authored-by: Andreas Karatzas --- .../kernels/attention/test_prefix_prefill.py | 336 ++++++++++++++++++ vllm/v1/attention/ops/prefix_prefill.py | 229 +++++++++--- 2 files changed, 525 insertions(+), 40 deletions(-) diff --git a/tests/kernels/attention/test_prefix_prefill.py b/tests/kernels/attention/test_prefix_prefill.py index f1c591fb67..e63d2c0adf 100644 --- a/tests/kernels/attention/test_prefix_prefill.py +++ b/tests/kernels/attention/test_prefix_prefill.py @@ -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) diff --git a/vllm/v1/attention/ops/prefix_prefill.py b/vllm/v1/attention/ops/prefix_prefill.py index 8488c72aea..6bbe6775e4 100644 --- a/vllm/v1/attention/ops/prefix_prefill.py +++ b/vllm/v1/attention/ops/prefix_prefill.py @@ -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