@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b88916617d

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

Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

stefankoncarevic committed Jul 30, 2026 at 06:31 UTC b88916617d3d2249bff0dae5cecb6b727c980a20

2 files changed +525 -40

tests/kernels/attention/test_prefix_prefill.py

+336

@@ -323,6 +323,342 @@ def test_contexted_kv_attention(

323 torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)

324

325

326 +@pytest.mark.parametrize("num_heads", NUM_HEADS)

327 +@pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)

328 +@pytest.mark.parametrize("head_size", HEAD_SIZES)

329 +@pytest.mark.parametrize("dtype", DTYPES)

330 +@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)

331 +@pytest.mark.parametrize("device", CUDA_DEVICES)

332 +@torch.inference_mode()

333 +def test_contexted_kv_attention_cached_kv(

334 + num_heads: int,

335 + num_queries_per_kv: int,

336 + head_size: int,

337 + dtype: torch.dtype,

338 + kv_cache_dtype: str,

339 + device: str,

340 + block_size: int = 32,

341 +) -> None:

342 + # Exercises the KV_FROM_CACHE path of context_attention_fwd: the current

343 + # chunk K/V are not passed as dense tensors (k=v=None); they are read back

344 + # from the paged KV cache, as done by layers that re-attend an already

345 + # cached sequence with the query only (e.g. IQuest LoopCoder's

346 + # attn(q, None, None)). The whole sequence therefore lives in the cache and

347 + # the result must still match a dense causal SDPA reference.

348 + if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):

349 + pytest.skip(

350 + "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"

351 + )

352 +

353 + set_random_seed(0)

354 + torch.set_default_device(device)

355 + torch.accelerator.set_device_index(device)

356 +

357 + MAX_SEQ_LEN = 1024

358 + MAX_CTX_LEN = 1024

359 + BS = 10

360 + cache_size = 640

361 + max_block_per_request = 64

362 + query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]

363 + # ensure one sequence in batch is a decode

364 + query_lens[-1] = 1

365 + ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]

366 + seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]

367 + num_kv_heads = num_heads // num_queries_per_kv

368 +

369 + num_tokens = sum(query_lens)

370 + query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)

371 + query.uniform_(-1e-3, 1e-3)

372 + output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)

373 +

374 + kv = torch.empty(sum(seq_lens), 2, num_kv_heads, head_size, dtype=dtype)

375 + kv.uniform_(-1e-3, 1e-3)

376 + key, value = kv.unbind(dim=1)

377 +

378 + if kv_cache_dtype == "auto":

379 + cache_dtype = dtype

380 + else:

381 + cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype]

382 + k_cache = torch.zeros(

383 + cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype

384 + )

385 + v_cache = torch.zeros(

386 + cache_size, block_size, num_kv_heads, head_size, dtype=cache_dtype

387 + )

388 + values = torch.arange(0, cache_size, dtype=torch.int32)

389 + values = values[torch.randperm(cache_size)]

390 + block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)

391 + b_seq_len = torch.tensor(seq_lens, dtype=torch.int32)

392 + b_start_loc = torch.cumsum(torch.tensor([0] + query_lens), dim=0).to(torch.int32)

393 + max_input_len = MAX_SEQ_LEN

394 + b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1]), dim=0).to(

395 + torch.int32

396 + )

397 + # Unlike the dense test, write the WHOLE sequence (context + current chunk)

398 + # into the paged cache, since the current chunk is read back from the cache.

399 + for i in range(BS):

400 + cur = 0

401 + block_id = 0

402 + while cur < seq_lens[i]:

403 + start_loc = b_seq_start_loc[i] + cur

404 + if cur + block_size > seq_lens[i]:

405 + end_loc = b_seq_start_loc[i] + seq_lens[i]

406 + else:

407 + end_loc = start_loc + block_size

408 + start_slot = block_table[i, block_id] * block_size

409 + end_slot = start_slot + end_loc - start_loc

410 + k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(

411 + key[start_loc:end_loc]

412 + )

413 + v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(

414 + value[start_loc:end_loc]

415 + )

416 + cur += block_size

417 + block_id += 1

418 + # transpose to the paged cache layouts the kernel expects:

419 + # K_cache[num_blocks, num_kv_heads, head_size/8, block_size, 8]

420 + # V_cache[num_blocks, num_kv_heads, head_size, block_size]

421 + k_cache = (

422 + k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)

423 + .permute(0, 2, 3, 1, 4)

424 + .contiguous()

425 + )

426 + v_cache = (

427 + v_cache.view(-1, block_size, num_kv_heads, head_size)

428 + .permute(0, 2, 3, 1)

429 + .contiguous()

430 + )

431 + k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)

432 +

433 + # Cached-K/V path: current-chunk k and v are None.

434 + context_attention_fwd(

435 + query,

436 + None,

437 + None,

438 + output,

439 + kv_cache_dtype,

440 + k_cache,

441 + v_cache,

442 + block_table,

443 + b_start_loc,

444 + b_seq_len,

445 + MAX_CTX_LEN,

446 + max_input_len,

447 + k_scale,

448 + v_scale,

449 + sliding_window=0,

450 + )

451 + torch.accelerator.synchronize()

452 +

453 + scale = float(1.0 / (head_size**0.5))

454 +

455 + query_sdpa = query.view(num_tokens, num_kv_heads, num_queries_per_kv, head_size)

456 + query_sdpa = query_sdpa.permute(1, 2, 0, 3).reshape(

457 + 1, num_heads, num_tokens, head_size

458 + )

459 + key_sdpa = key[:, :, None, :].expand(

460 + key.shape[0], num_kv_heads, num_queries_per_kv, key.shape[-1]

461 + )

462 + key_sdpa = key_sdpa.permute(1, 2, 0, 3).reshape(

463 + 1, num_heads, sum(seq_lens), head_size

464 + )

465 + value_sdpa = value[:, :, None, :].expand(

466 + value.shape[0], num_kv_heads, num_queries_per_kv, value.shape[-1]

467 + )

468 + value_sdpa = value_sdpa.permute(1, 2, 0, 3).reshape(

469 + 1, num_heads, sum(seq_lens), head_size

470 + )

471 +

472 + attn_mask = create_causal_attention_mask_for_sdpa(

473 + query_lens, seq_lens, 0, device=device, dtype=dtype

474 + )

475 + output_ref = F.scaled_dot_product_attention(

476 + query_sdpa,

477 + key_sdpa,

478 + value_sdpa,

479 + attn_mask=attn_mask,

480 + dropout_p=0.0,

481 + scale=scale,

482 + )

483 + output_ref = output_ref.view(num_heads, num_tokens, head_size)

484 + output_ref = output_ref.permute(1, 0, 2).contiguous()

485 + atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-4

486 + torch.testing.assert_close(output, output_ref, atol=atol, rtol=0)

487 +

488 +

489 +@pytest.mark.parametrize("device", CUDA_DEVICES)

490 +@torch.inference_mode()

491 +def test_contexted_kv_attention_cached_kv_block_table_boundary(device: str) -> None:

492 + # Boundary guard for the KV_FROM_CACHE block-table load. With an

493 + # exact-sized block table (row length == number of blocks for the

494 + # sequence) and a query that ends the sequence, the last K/V tile has

495 + # padded lanes whose absolute positions step past the sequence end. Those

496 + # lanes must not read a block-table entry past this batch's row.

497 + #

498 + # seq_len is a whole number of blocks, so bn_logical for a padded lane

499 + # lands exactly on num_blocks (one past the last valid row entry), and

500 + # query_len is not tile-aligned so the overshoot is actually exercised.

501 + # The result must still match a dense causal SDPA reference.

502 + set_random_seed(0)

503 + torch.set_default_device(device)

504 + torch.accelerator.set_device_index(device)

505 +

506 + dtype = torch.float16

507 + kv_cache_dtype = "auto"

508 + num_heads = 4

509 + num_queries_per_kv = 1

510 + num_kv_heads = num_heads // num_queries_per_kv

511 + head_size = 32

512 + block_size = 32

513 + num_blocks = 4

514 + seq_len = num_blocks * block_size # 128 == exact multiple of block_size

515 + query_len = 99 # < seq_len and not a multiple of the kernel tile

516 +

517 + query = torch.empty(query_len, num_heads, head_size, dtype=dtype)

518 + query.uniform_(-1e-3, 1e-3)

519 + output = torch.empty(query_len, num_heads, head_size, dtype=dtype)

520 +

521 + kv = torch.empty(seq_len, 2, num_kv_heads, head_size, dtype=dtype)

522 + kv.uniform_(-1e-3, 1e-3)

523 + key, value = kv.unbind(dim=1)

524 +

525 + k_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype)

526 + v_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype)

527 + # Exact-sized block table: row length == num_blocks (identity mapping).

528 + block_table = torch.arange(num_blocks, dtype=torch.int32).view(1, num_blocks)

529 + b_seq_len = torch.tensor([seq_len], dtype=torch.int32)

530 + b_start_loc = torch.tensor([0, query_len], dtype=torch.int32)

531 +

532 + # Write the whole sequence (context + current chunk) into the paged cache.

533 + for cur in range(0, seq_len, block_size):

534 + block_id = cur // block_size

535 + end = min(cur + block_size, seq_len)

536 + start_slot = block_table[0, block_id] * block_size

537 + end_slot = start_slot + (end - cur)

538 + k_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(

539 + key[cur:end]

540 + )

541 + v_cache.view(-1, num_kv_heads, head_size)[start_slot:end_slot].copy_(

542 + value[cur:end]

543 + )

544 +

545 + k_cache = (

546 + k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)

547 + .permute(0, 2, 3, 1, 4)

548 + .contiguous()

549 + )

550 + v_cache = (

551 + v_cache.view(-1, block_size, num_kv_heads, head_size)

552 + .permute(0, 2, 3, 1)

553 + .contiguous()

554 + )

555 + k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)

556 +

557 + # Cached-K/V path: current-chunk k and v are None.

558 + context_attention_fwd(

559 + query,

560 + None,

561 + None,

562 + output,

563 + kv_cache_dtype,

564 + k_cache,

565 + v_cache,

566 + block_table,

567 + b_start_loc,

568 + b_seq_len,

569 + seq_len,

570 + query_len,

571 + k_scale,

572 + v_scale,

573 + sliding_window=0,

574 + )

575 + torch.accelerator.synchronize()

576 +

577 + scale = float(1.0 / (head_size**0.5))

578 + query_sdpa = query.view(query_len, num_kv_heads, num_queries_per_kv, head_size)

579 + query_sdpa = query_sdpa.permute(1, 2, 0, 3).reshape(

580 + 1, num_heads, query_len, head_size

581 + )

582 + key_sdpa = key[:, :, None, :].expand(

583 + seq_len, num_kv_heads, num_queries_per_kv, head_size

584 + )

585 + key_sdpa = key_sdpa.permute(1, 2, 0, 3).reshape(1, num_heads, seq_len, head_size)

586 + value_sdpa = value[:, :, None, :].expand(

587 + seq_len, num_kv_heads, num_queries_per_kv, head_size

588 + )

589 + value_sdpa = value_sdpa.permute(1, 2, 0, 3).reshape(

590 + 1, num_heads, seq_len, head_size

591 + )

592 +

593 + attn_mask = create_causal_attention_mask_for_sdpa(

594 + [query_len], [seq_len], 0, device=device, dtype=dtype

595 + )

596 + output_ref = F.scaled_dot_product_attention(

597 + query_sdpa,

598 + key_sdpa,

599 + value_sdpa,

600 + attn_mask=attn_mask,

601 + dropout_p=0.0,

602 + scale=scale,

603 + )

604 + output_ref = output_ref.view(num_heads, query_len, head_size)

605 + output_ref = output_ref.permute(1, 0, 2).contiguous()

606 + torch.testing.assert_close(output, output_ref, atol=1e-4, rtol=0)

607 +

608 +

609 +@pytest.mark.parametrize("device", CUDA_DEVICES)

610 +@torch.inference_mode()

611 +def test_contexted_kv_attention_cached_kv_alibi_unsupported(device: str) -> None:

612 + # The cached-K/V (k=None) path is not supported together with ALiBi; the

613 + # entry point must reject it up-front with a clear NotImplementedError

614 + # rather than launching the kernel with an unsupported combination.

615 + set_random_seed(0)

616 + torch.set_default_device(device)

617 + torch.accelerator.set_device_index(device)

618 +

619 + num_heads = 4

620 + num_kv_heads = 4

621 + head_size = 16

622 + x = 8

623 + block_size = 16

624 + num_blocks = 4

625 + query_len = 8

626 +

627 + query = torch.empty(query_len, num_heads, head_size, dtype=torch.float16)

628 + query.uniform_(-1e-3, 1e-3)

629 + output = torch.empty_like(query)

630 + k_cache = torch.zeros(

631 + num_blocks, num_kv_heads, head_size // x, block_size, x, dtype=torch.float16

632 + )

633 + v_cache = torch.zeros(

634 + num_blocks, num_kv_heads, head_size, block_size, dtype=torch.float16

635 + )

636 + block_table = torch.arange(num_blocks, dtype=torch.int32).view(1, num_blocks)

637 + b_seq_len = torch.tensor([query_len], dtype=torch.int32)

638 + b_start_loc = torch.tensor([0, query_len], dtype=torch.int32)

639 + k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)

640 + alibi_slopes = torch.ones(num_heads, dtype=torch.float32, device=device)

641 +

642 + with pytest.raises(NotImplementedError):

643 + context_attention_fwd(

644 + query,

645 + None,

646 + None,

647 + output,

648 + "auto",

649 + k_cache,

650 + v_cache,

651 + block_table,

652 + b_start_loc,

653 + b_seq_len,

654 + query_len,

655 + query_len,

656 + k_scale,

657 + v_scale,

658 + alibi_slopes=alibi_slopes,

659 + )

660 +

661 +

662 @pytest.mark.parametrize("num_heads", NUM_HEADS)

663 @pytest.mark.parametrize("num_queries_per_kv", NUM_QUERIES_PER_KV)

664 @pytest.mark.parametrize("head_size", HEAD_SIZES)

vllm/v1/attention/ops/prefix_prefill.py

+189 -40

@@ -35,6 +35,74 @@ float8_info = torch.finfo(current_platform.fp8_dtype())

35 # ],

36 # key=["BLOCK_SIZE", "MAX_Q_LEN", "MAX_CTX_LEN"]

37 # )

38 +@triton.jit

39 +def _paged_kv_cache_offsets(

40 + B_Loc,

41 + cur_batch,

42 + token_indices,

43 + token_valid,

44 + offs_d,

45 + cur_kv_head,

46 + x,

47 + stride_b_loc_b,

48 + stride_b_loc_s,

49 + stride_k_cache_bs,

50 + stride_k_cache_h,

51 + stride_k_cache_d,

52 + stride_k_cache_bl,

53 + stride_k_cache_x,

54 + stride_v_cache_bs,

55 + stride_v_cache_h,

56 + stride_v_cache_d,

57 + stride_v_cache_bl,

58 + PHYSICAL_BLOCK_SIZE: tl.constexpr,

59 + MASK_BLOCK_TABLE: tl.constexpr = False,

60 +):

61 + """Compute paged K/V cache element offsets for a tile of token positions.

62 +

63 + `token_indices` holds the absolute sequence positions of the tile. The

64 + logical block for each token is looked up in `B_Loc` (handling cross-block

65 + tiles), then the physical element offsets are built with the same layout the

66 + kernel uses everywhere:

67 + K cache: [num_blocks, num_kv_heads, head_size // x, block_size, x]

68 + V cache: [num_blocks, num_kv_heads, head_size, block_size]

69 + Returns `(off_k, off_v)` with shapes [D, N] and [N, D] respectively.

70 +

71 + When `MASK_BLOCK_TABLE` is set, the block-table load is masked with

72 + `token_valid` so padded lanes (a tile that extends past the end of the

73 + sequence) cannot read a block-table entry past this batch's row. Invalid

74 + lanes resolve to block 0 (always in-bounds) and are dropped by the caller's

75 + later K/V load mask. When it is not set, `token_valid` is unused and the

76 + load is emitted exactly as before.

77 + """

78 + bn_logical = token_indices // PHYSICAL_BLOCK_SIZE

79 + if MASK_BLOCK_TABLE:

80 + bn = tl.load(

81 + B_Loc + cur_batch * stride_b_loc_b + bn_logical * stride_b_loc_s,

82 + mask=token_valid,

83 + other=0,

84 + ).to(tl.int64)

85 + else:

86 + bn = tl.load(

87 + B_Loc + cur_batch * stride_b_loc_b + bn_logical * stride_b_loc_s

88 + ).to(tl.int64)

89 + internal = token_indices % PHYSICAL_BLOCK_SIZE

90 + off_k = (

91 + bn[None, :] * stride_k_cache_bs

92 + + cur_kv_head * stride_k_cache_h

93 + + (offs_d[:, None] // x) * stride_k_cache_d

94 + + internal[None, :] * stride_k_cache_bl

95 + + (offs_d[:, None] % x) * stride_k_cache_x

96 + )

97 + off_v = (

98 + bn[:, None] * stride_v_cache_bs

99 + + cur_kv_head * stride_v_cache_h

100 + + offs_d[None, :] * stride_v_cache_d

101 + + internal[:, None] * stride_v_cache_bl

102 + )

103 + return off_k, off_v

104 +

105 +

106 @triton.jit

107 def _fwd_kernel(

108 Q,

@@ -94,6 +162,11 @@ def _fwd_kernel(

162 MAX_CTX_LEN: tl.constexpr = 0,

163 FP8_MIN: tl.constexpr = float8_info.min,

164 FP8_MAX: tl.constexpr = float8_info.max,

165 + # When True, the current-chunk K/V are not provided as dense tensors

166 + # (`K`/`V` are unused placeholders); instead they are read back from the

167 + # paged KV cache. This supports layers that re-attend an already-cached

168 + # sequence with query only (e.g. IQuest LoopCoder's `attn(q, None, None)`).

169 + KV_FROM_CACHE: tl.constexpr = False,

170 ):

171 cur_batch = tl.program_id(0)

172 cur_head = tl.program_id(1)

@@ -162,33 +235,34 @@ def _fwd_kernel(

235 # replace one physical block every 17 32-Tile blocks

236 # Calculate the logical block index of each of the 32 tokens

237 # in the current Tile (handling cross-block cases).

238 + # Physical-block K/V cache offsets for the context tokens in this tile

239 + # (handles cross-block tiles via B_Loc). Same addressing is reused for

240 + # the current-chunk cache read below.

241 token_indices = start_n + offs_bs_n

166 - bn_logical_indices = token_indices // PHYSICAL_BLOCK_SIZE

167 -

168 - # 2. Vectorized loading of physical block IDs from B_Loc

169 - bn = tl.load(

170 - B_Loc + cur_batch * stride_b_loc_b + bn_logical_indices * stride_b_loc_s

171 - ).to(tl.int64)

172 -

173 - # 3. Calculate the exact offset of

174 - # each token within its physical block.

175 - internal_offsets = token_indices % PHYSICAL_BLOCK_SIZE

176 -

177 - # Addressing of K (5D)

178 - off_k = (

179 - bn[None, :] * stride_k_cache_bs

180 - + cur_kv_head * stride_k_cache_h

181 - + (offs_d[:, None] // x) * stride_k_cache_d

182 - + internal_offsets[None, :] * stride_k_cache_bl

183 - + (offs_d[:, None] % x) * stride_k_cache_x

184 - )

185 -

186 - # Addressing of V (4D)

187 - off_v = (

188 - bn[:, None] * stride_v_cache_bs

189 - + cur_kv_head * stride_v_cache_h

190 - + offs_d[None, :] * stride_v_cache_d

191 - + internal_offsets[:, None] * stride_v_cache_bl

242 + # Context tokens are always followed by the current chunk, so the tile

243 + # never steps past this batch's block-table row; the block-table load

244 + # stays unmasked (MASK_BLOCK_TABLE=False) and offs_bs_n is an unused

245 + # placeholder for token_valid.

246 + off_k, off_v = _paged_kv_cache_offsets(

247 + B_Loc,

248 + cur_batch,

249 + token_indices,

250 + offs_bs_n,

251 + offs_d,

252 + cur_kv_head,

253 + x,

254 + stride_b_loc_b,

255 + stride_b_loc_s,

256 + stride_k_cache_bs,

257 + stride_k_cache_h,

258 + stride_k_cache_d,

259 + stride_k_cache_bl,

260 + stride_k_cache_x,

261 + stride_v_cache_bs,

262 + stride_v_cache_h,

263 + stride_v_cache_d,

264 + stride_v_cache_bl,

265 + PHYSICAL_BLOCK_SIZE,

266 )

267

268 if (

@@ -300,12 +374,54 @@ def _fwd_kernel(

374 ):

375 start_n = tl.multiple_of(start_n, BLOCK_N)

376 # -- compute qk ----

303 - k = tl.load(

304 - k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,

305 - mask=dim_mask[:, None]

306 - & ((start_n + offs_n[None, :]) < cur_batch_query_len),

307 - other=0.0,

308 - )

377 + if KV_FROM_CACHE:

378 + # Current-chunk K/V are not passed densely; read them from the paged

379 + # cache at their absolute positions [ctx_len, ctx_len + query_len),

380 + # reusing the same addressing as the context loop above.

381 + cache_token_idx = cur_batch_ctx_len + start_n + offs_n

382 + # Padded lanes of the last tile can index past the sequence end;

383 + # mask the block-table load so they don't read past this batch's

384 + # row (they are dropped by the K/V load mask below anyway).

385 + cache_token_valid = (start_n + offs_n) < cur_batch_query_len

386 + off_k_cur, off_v_cur = _paged_kv_cache_offsets(

387 + B_Loc,

388 + cur_batch,

389 + cache_token_idx,

390 + cache_token_valid,

391 + offs_d,

392 + cur_kv_head,

393 + x,

394 + stride_b_loc_b,

395 + stride_b_loc_s,

396 + stride_k_cache_bs,

397 + stride_k_cache_h,

398 + stride_k_cache_d,

399 + stride_k_cache_bl,

400 + stride_k_cache_x,

401 + stride_v_cache_bs,

402 + stride_v_cache_h,

403 + stride_v_cache_d,

404 + stride_v_cache_bl,

405 + PHYSICAL_BLOCK_SIZE,

406 + MASK_BLOCK_TABLE=True,

407 + )

408 + k_cur_load = tl.load(

409 + K_cache + off_k_cur,

410 + mask=dim_mask[:, None]

411 + & ((start_n + offs_n[None, :]) < cur_batch_query_len),

412 + other=0.0,

413 + )

414 + if k_cur_load.dtype.is_fp8():

415 + k = (k_cur_load.to(tl.float32) * tl.load(k_scale)).to(q.dtype)

416 + else:

417 + k = k_cur_load

418 + else:

419 + k = tl.load(

420 + k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,

421 + mask=dim_mask[:, None]

422 + & ((start_n + offs_n[None, :]) < cur_batch_query_len),

423 + other=0.0,

424 + )

425

426 qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)

427 qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)

@@ -333,12 +449,25 @@ def _fwd_kernel(

449 acc = acc * alpha[:, None]

450

451 # update acc

336 - v = tl.load(

337 - v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,

338 - mask=dim_mask[None, :]

339 - & ((start_n + offs_n[:, None]) < cur_batch_query_len),

340 - other=0.0,

341 - )

452 + if KV_FROM_CACHE:

453 + # off_v_cur was computed together with off_k_cur above.

454 + v_cur_load = tl.load(

455 + V_cache + off_v_cur,

456 + mask=dim_mask[None, :]

457 + & ((start_n + offs_n[:, None]) < cur_batch_query_len),

458 + other=0.0,

459 + )

460 + if v_cur_load.dtype.is_fp8():

461 + v = (v_cur_load.to(tl.float32) * tl.load(v_scale)).to(q.dtype)

462 + else:

463 + v = v_cur_load

464 + else:

465 + v = tl.load(

466 + v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,

467 + mask=dim_mask[None, :]

468 + & ((start_n + offs_n[:, None]) < cur_batch_query_len),

469 + other=0.0,

470 + )

471 p = p.to(v.dtype)

472

473 acc = tl.dot(p, v, acc=acc, input_precision=IN_PRECISION)

@@ -703,8 +832,27 @@ def context_attention_fwd(

832 FP8 KV Cache prefill kernel"

833 )

834

835 + assert (k is None) == (v is None), (

836 + "k and v must both be None (cached-K/V path) or both be tensors"

837 + )

838 + kv_from_cache = k is None

839 + if kv_from_cache:

840 + if alibi_slopes is not None:

841 + raise NotImplementedError(

842 + "context_attention_fwd with cached K/V (key=None) is not "

843 + "supported together with ALiBi slopes."

844 + )

845 + # k_cache: [num_blocks, num_kv_heads, head_size // x, block_size, x]

846 + num_kv_heads = k_cache.shape[1]

847 + Lq = Lk = Lv = q.shape[-1]

848 + # q is a never-dereferenced placeholder for the dense K/V pointer args.

849 + k = q

850 + v = q

851 + else:

852 + num_kv_heads = k.shape[1]

853 + Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]

854 +

855 # shape constraints

707 - Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]

856 assert Lq == Lk and Lk == Lv

857 # round up Lk to a power of 2 - this is required for Triton block size

858 Lk_padded = triton.next_power_of_2(Lk)

@@ -712,7 +860,7 @@ def context_attention_fwd(

860 if sm_scale is None:

861 sm_scale = 1.0 / (Lq**0.5)

862 batch, head = b_seq_len.shape[0], q.shape[1]

715 - num_queries_per_kv = q.shape[1] // k.shape[1]

863 + num_queries_per_kv = q.shape[1] // num_kv_heads

864

865 assert batch + 1 == len(b_start_loc)

866

@@ -873,6 +1021,7 @@ def context_attention_fwd(

1021 num_stages=1,

1022 USE_SINKS=sinks is not None,

1023 CAUSAL=causal,

1024 + KV_FROM_CACHE=kv_from_cache,

1025 **extra_kargs,

1026 )

1027 return