[Perf] Support internal prefill checkpoints for Mamba prefix caching, 9%~25% TTFT improvement (#52789)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Ziming Huang <zelda.huanghuang@gmail.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Wentao Ye
co-authored by
Ziming Huang
OpenAI Codex
Ziming Huang
OpenAI Codex
parent
020829f66b
commit
9eb9d9d395
@@ -13,7 +13,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
flashkda
|
||||
GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git
|
||||
GIT_TAG 053de1b716ef3255873e02d2d28f4adf09951978
|
||||
GIT_TAG ee0be888cd0e972f9409bf53756f8c38c6652173
|
||||
GIT_PROGRESS TRUE
|
||||
GIT_SUBMODULES cutlass
|
||||
)
|
||||
|
||||
@@ -10,7 +10,8 @@ STABLE_TORCH_LIBRARY(_flashkda_C, m) {
|
||||
"Tensor(a!) out, Tensor(c!) workspace, Tensor A_log, Tensor dt_bias, "
|
||||
"float lower_bound, "
|
||||
"Tensor? initial_state=None, Tensor(b!)? final_state=None, "
|
||||
"Tensor? cu_seqlens=None) -> ()");
|
||||
"Tensor? cu_seqlens=None, Tensor(d!)? checkpoint_state=None, "
|
||||
"Tensor? checkpoint_offsets=None) -> ()");
|
||||
}
|
||||
|
||||
STABLE_TORCH_LIBRARY_IMPL(_flashkda_C, CompositeExplicitAutograd, m) {
|
||||
|
||||
@@ -21,6 +21,8 @@ from vllm.models.kimi_k3.amd.ops.third_party.kda import (
|
||||
fused_recurrent_kda_packed_decode as fused_recurrent_kda_packed_decode_amd,
|
||||
)
|
||||
from vllm.models.kimi_k3.nvidia.kda import (
|
||||
_flashkda_prefill,
|
||||
_store_cache_checkpoints_kernel,
|
||||
is_flashkda_supported,
|
||||
is_fused_kda_decode_supported,
|
||||
)
|
||||
@@ -40,6 +42,7 @@ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
|
||||
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
|
||||
|
||||
DEVICE = current_platform.device_type
|
||||
|
||||
@@ -1088,6 +1091,17 @@ def test_flashkda_correctness():
|
||||
expected_states.append(final_state)
|
||||
expected_out = torch.cat(expected_outputs, dim=1)
|
||||
expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous()
|
||||
_, expected_checkpoint = naive_recurrent_kda(
|
||||
q_norm[:, :16],
|
||||
k_norm[:, :16],
|
||||
v[:, :16],
|
||||
gate[:, :16],
|
||||
beta[:, :16],
|
||||
initial_state=initial_state[0:1].transpose(-1, -2),
|
||||
output_final_state=True,
|
||||
)
|
||||
assert expected_checkpoint is not None
|
||||
expected_checkpoint = expected_checkpoint.transpose(-1, -2).contiguous()
|
||||
|
||||
actual_out = torch.empty_like(v)
|
||||
actual_state = torch.empty_like(initial_state)
|
||||
@@ -1115,3 +1129,70 @@ def test_flashkda_correctness():
|
||||
|
||||
assert_close("o", expected_out, actual_out, 0.01)
|
||||
assert_close("ht", expected_state, actual_state, 0.01)
|
||||
|
||||
checkpoint_out = torch.empty_like(v)
|
||||
checkpoint_final_state = torch.empty_like(initial_state)
|
||||
checkpoint_state = torch.empty_like(initial_state)
|
||||
checkpoint_offsets = torch.tensor([16, 31], dtype=torch.int32, device=DEVICE)
|
||||
_flashkda_prefill(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=raw_g,
|
||||
beta=beta_logits,
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
lower_bound=lower_bound,
|
||||
initial_state=initial_state,
|
||||
cu_seqlens=cu_seqlens,
|
||||
out=checkpoint_out,
|
||||
final_state=checkpoint_final_state,
|
||||
workspace=workspace,
|
||||
checkpoint_state=checkpoint_state,
|
||||
checkpoint_offsets=checkpoint_offsets,
|
||||
)
|
||||
|
||||
assert_close("checkpoint_o", expected_out, checkpoint_out, 0.01)
|
||||
assert_close("checkpoint_ht", expected_state, checkpoint_final_state, 0.01)
|
||||
assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], 0.01)
|
||||
|
||||
conv_state = torch.zeros(2, H * D, 3, dtype=q.dtype, device=DEVICE)
|
||||
recurrent_storage = torch.zeros(2, H * D * D + 8, device=DEVICE)
|
||||
recurrent_state = recurrent_storage[:, : H * D * D].view(2, H, D, D)
|
||||
conv_input = q[0].flatten(1)
|
||||
checkpoint_state_indices = torch.tensor(
|
||||
[1, NULL_BLOCK_ID], dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
state_len = conv_state.shape[-1]
|
||||
width = H * D
|
||||
recurrent_row_size = checkpoint_state[0].numel()
|
||||
block_size = 256
|
||||
_store_cache_checkpoints_kernel[
|
||||
(
|
||||
checkpoint_state_indices.numel(),
|
||||
(max(width * state_len, recurrent_row_size) + block_size - 1) // block_size,
|
||||
)
|
||||
](
|
||||
conv_input,
|
||||
conv_state,
|
||||
checkpoint_state,
|
||||
recurrent_state,
|
||||
cu_seqlens,
|
||||
checkpoint_offsets,
|
||||
checkpoint_state_indices,
|
||||
conv_input.stride(0),
|
||||
conv_input.stride(1),
|
||||
conv_state.stride(0),
|
||||
conv_state.stride(1),
|
||||
conv_state.stride(2),
|
||||
checkpoint_state.stride(0),
|
||||
recurrent_state.stride(0),
|
||||
checkpoint_offsets.stride(0),
|
||||
state_len,
|
||||
width,
|
||||
recurrent_row_size,
|
||||
NULL_BLOCK_ID,
|
||||
block_size,
|
||||
)
|
||||
torch.testing.assert_close(conv_state[1], q[0, 13:16].flatten(1).transpose(0, 1))
|
||||
torch.testing.assert_close(recurrent_state[1], checkpoint_state[0])
|
||||
|
||||
@@ -87,6 +87,7 @@ def _make_builder(
|
||||
device: torch.device = DEVICE,
|
||||
mamba_cache_mode: str = "none",
|
||||
use_recoverssm: bool = False,
|
||||
num_prefill_checkpoint_blocks: int = 0,
|
||||
) -> AttentionMetadataBuilder:
|
||||
vllm_config = create_vllm_config(
|
||||
model_name="Qwen/Qwen3.5-0.8B",
|
||||
@@ -110,6 +111,7 @@ def _make_builder(
|
||||
dtypes=(torch.float16,),
|
||||
mamba_cache_mode=mamba_cache_mode,
|
||||
num_speculative_blocks=(0 if use_recoverssm else num_speculative_tokens),
|
||||
num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
|
||||
),
|
||||
layer_names=["layer.0"],
|
||||
vllm_config=vllm_config,
|
||||
@@ -121,6 +123,54 @@ def _make_builder(
|
||||
return builder
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
||||
def test_internal_checkpoint_metadata_targets_last_aligned_boundary():
|
||||
device = torch.device("cuda")
|
||||
batch = BatchSpec(seq_lens=[50, 32], query_lens=[50, 16])
|
||||
common_attn_metadata = create_common_attn_metadata(
|
||||
batch, BLOCK_SIZE, device, arange_block_indices=True
|
||||
).replace(is_prefilling=torch.tensor([True, True]))
|
||||
actual = _make_builder(
|
||||
KimiK3KDAMetadataBuilder,
|
||||
num_speculative_tokens=0,
|
||||
full_cuda_graph=False,
|
||||
mamba_cache_mode="align",
|
||||
num_prefill_checkpoint_blocks=1,
|
||||
device=device,
|
||||
).build(0, common_attn_metadata)
|
||||
|
||||
assert actual.checkpoint is not None
|
||||
torch.testing.assert_close(
|
||||
actual.checkpoint.state_indices,
|
||||
torch.tensor([2, NULL_BLOCK_ID], dtype=torch.int32, device=device),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
actual.checkpoint.checkpoint_offsets,
|
||||
torch.tensor([48, 0], dtype=torch.int32, device=device),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
||||
def test_internal_checkpoint_metadata_skips_unaligned_offset():
|
||||
device = torch.device("cuda")
|
||||
# The checkpoint block boundary is 48, but this query starts at token 1,
|
||||
# making the real checkpoint offset 47, which is not FlashKDA-aligned.
|
||||
batch = BatchSpec(seq_lens=[50], query_lens=[49])
|
||||
common_attn_metadata = create_common_attn_metadata(
|
||||
batch, BLOCK_SIZE, device, arange_block_indices=True
|
||||
).replace(is_prefilling=torch.tensor([True]))
|
||||
actual = _make_builder(
|
||||
KimiK3KDAMetadataBuilder,
|
||||
num_speculative_tokens=0,
|
||||
full_cuda_graph=False,
|
||||
mamba_cache_mode="align",
|
||||
num_prefill_checkpoint_blocks=1,
|
||||
device=device,
|
||||
).build(0, common_attn_metadata)
|
||||
|
||||
assert actual.checkpoint is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
(
|
||||
"batch",
|
||||
|
||||
@@ -77,6 +77,7 @@ def make_full_mamba_manager(
|
||||
mamba_block_size: int = 4,
|
||||
num_blocks: int = 32,
|
||||
use_eagle: bool = False,
|
||||
num_prefill_checkpoint_blocks: int = 0,
|
||||
):
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
@@ -98,6 +99,7 @@ def make_full_mamba_manager(
|
||||
shapes=(1, 1),
|
||||
dtypes=(torch.float32,),
|
||||
mamba_cache_mode="align",
|
||||
num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
|
||||
),
|
||||
),
|
||||
],
|
||||
@@ -135,6 +137,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
|
||||
dcp_world_size=dcp_world_size,
|
||||
scheduler_block_size=scheduler_block_size,
|
||||
mamba_partial_cache_hit=True,
|
||||
mamba_has_prefill_checkpoint_blocks=False,
|
||||
)
|
||||
split = Scheduler._mamba_block_aligned_split
|
||||
|
||||
@@ -180,6 +183,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget():
|
||||
use_eagle=False,
|
||||
hash_block_size=32,
|
||||
mamba_partial_cache_hit=False,
|
||||
mamba_has_prefill_checkpoint_blocks=False,
|
||||
)
|
||||
req = make_request("0", [0] * prompt_length, 32, sha256)
|
||||
split = Scheduler._mamba_block_aligned_split
|
||||
@@ -218,6 +222,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
|
||||
use_eagle=False,
|
||||
hash_block_size=32,
|
||||
mamba_partial_cache_hit=False,
|
||||
mamba_has_prefill_checkpoint_blocks=False,
|
||||
)
|
||||
req = make_request("0", [0] * prompt_length, 32, sha256)
|
||||
split = Scheduler._mamba_block_aligned_split
|
||||
@@ -462,6 +467,85 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
|
||||
assert moved[0].block_hash_num_tokens == 6
|
||||
|
||||
|
||||
def test_partial_hit_then_internal_checkpoint_uses_distinct_mamba_blocks():
|
||||
hash_block_size = 2
|
||||
mamba_block_size = 4
|
||||
manager = make_full_mamba_manager(
|
||||
dcp_world_size=1,
|
||||
hash_block_size=hash_block_size,
|
||||
full_block_size=hash_block_size,
|
||||
mamba_block_size=mamba_block_size,
|
||||
num_prefill_checkpoint_blocks=1,
|
||||
)
|
||||
|
||||
owner = make_request("owner", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
|
||||
computed_blocks, num_computed, _ = manager.get_computed_blocks(owner)
|
||||
assert manager.allocate_slots(owner, 6, num_computed, computed_blocks) is not None
|
||||
manager.free(owner)
|
||||
manager.new_step_starts()
|
||||
|
||||
partial_hash = owner.block_hashes[2]
|
||||
partial_block = manager.block_pool.get_cached_block(partial_hash, [1])
|
||||
assert partial_block is not None
|
||||
partial_block_id = partial_block[0].block_id
|
||||
|
||||
replay = make_request(
|
||||
"replay",
|
||||
[0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6],
|
||||
hash_block_size,
|
||||
sha256,
|
||||
)
|
||||
computed_blocks, num_computed, _ = manager.get_computed_blocks(replay)
|
||||
assert num_computed == 6
|
||||
|
||||
realigned_blocks = manager.allocate_slots(replay, 2, num_computed, computed_blocks)
|
||||
assert realigned_blocks is not None
|
||||
mamba_cow_block_id = realigned_blocks.get_block_ids()[1][0]
|
||||
copies, retained = manager.take_kv_cache_block_copies()
|
||||
assert KVCacheBlockCopy(partial_block_id, mamba_cow_block_id) in copies
|
||||
manager.block_pool.free_blocks(retained)
|
||||
|
||||
replay.num_computed_tokens = 8
|
||||
manager.new_step_starts()
|
||||
final_blocks = manager.allocate_slots(replay, 6)
|
||||
assert final_blocks is not None
|
||||
|
||||
checkpoint_block_id, running_block_id = final_blocks.get_block_ids()[1]
|
||||
assert len({mamba_cow_block_id, checkpoint_block_id, running_block_id}) == 3
|
||||
mamba_blocks = manager.get_blocks(replay.request_id).blocks[1]
|
||||
assert mamba_blocks[2].block_id == checkpoint_block_id
|
||||
assert mamba_blocks[3].block_id == running_block_id
|
||||
|
||||
|
||||
def test_internal_checkpoint_requires_block_aligned_start():
|
||||
hash_block_size = 2
|
||||
mamba_block_size = 16
|
||||
manager = make_full_mamba_manager(
|
||||
dcp_world_size=1,
|
||||
hash_block_size=hash_block_size,
|
||||
full_block_size=hash_block_size,
|
||||
mamba_block_size=mamba_block_size,
|
||||
num_prefill_checkpoint_blocks=1,
|
||||
)
|
||||
request = make_request("producer", list(range(50)), hash_block_size, sha256)
|
||||
|
||||
# Compute token 0 first, so the next query starts at token 1, which is not
|
||||
# aligned to the Mamba block size.
|
||||
assert manager.allocate_slots(request, 1) is not None
|
||||
request.num_computed_tokens = 1
|
||||
manager.new_step_starts()
|
||||
|
||||
new_blocks = manager.allocate_slots(request, 49)
|
||||
|
||||
assert new_blocks is not None
|
||||
mamba_blocks = manager.get_blocks(request.request_id).blocks[1]
|
||||
checkpoint_block_idx = 48 // mamba_block_size - 1
|
||||
assert mamba_blocks[checkpoint_block_idx].is_null
|
||||
assert not mamba_blocks[-1].is_null
|
||||
checkpoint_hash = request.block_hashes[48 // hash_block_size - 1]
|
||||
assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None
|
||||
|
||||
|
||||
def test_external_mamba_hit_same_block_uses_running_cow_on_continue():
|
||||
"""An external mid-block hit must become a running request even when its
|
||||
first continuation does not need another Mamba block."""
|
||||
@@ -560,6 +644,7 @@ def test_take_partial_tail_offloads_returns_cow_target():
|
||||
shapes=(1, 1),
|
||||
dtypes=(torch.float32,),
|
||||
mamba_cache_mode="align",
|
||||
num_prefill_checkpoint_blocks=1,
|
||||
),
|
||||
),
|
||||
],
|
||||
|
||||
@@ -37,7 +37,9 @@ PROMPT_LEN = 2002
|
||||
MAMBA_GROUP_ID = 1
|
||||
|
||||
|
||||
def _make_hybrid_kv_cache_manager() -> KVCacheManager:
|
||||
def _make_hybrid_kv_cache_manager(
|
||||
num_prefill_checkpoint_blocks: int = 0,
|
||||
) -> KVCacheManager:
|
||||
config = KVCacheConfig(
|
||||
num_blocks=10000,
|
||||
kv_cache_tensors=[],
|
||||
@@ -59,6 +61,7 @@ def _make_hybrid_kv_cache_manager() -> KVCacheManager:
|
||||
dtypes=(torch.float32,),
|
||||
mamba_cache_mode="align",
|
||||
num_speculative_blocks=NUM_SPEC,
|
||||
num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
|
||||
),
|
||||
),
|
||||
],
|
||||
@@ -78,6 +81,7 @@ def _split(
|
||||
num_new_tokens: int,
|
||||
use_eagle: bool = True,
|
||||
partial_hit: bool = False,
|
||||
num_prefill_checkpoint_blocks: int = 0,
|
||||
) -> int:
|
||||
"""Call the real `Scheduler._mamba_block_aligned_split` on a stub self."""
|
||||
stub = SimpleNamespace(
|
||||
@@ -88,10 +92,42 @@ def _split(
|
||||
# `prefix_match_unit` finer than the block size (#46384).
|
||||
mamba_partial_cache_hit=partial_hit,
|
||||
hash_block_size=ATTN_BLOCK_SIZE,
|
||||
mamba_has_prefill_checkpoint_blocks=(
|
||||
num_prefill_checkpoint_blocks > 0 and not use_eagle
|
||||
),
|
||||
)
|
||||
return Scheduler._mamba_block_aligned_split(stub, request, num_new_tokens)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("prompt_len", "num_new_tokens", "use_eagle", "expected"),
|
||||
[
|
||||
(2002, 2002, False, 2002),
|
||||
(3602, 2000, False, 2000),
|
||||
(3602, 3602, True, MAMBA_BLOCK_SIZE),
|
||||
],
|
||||
)
|
||||
def test_internal_checkpoint_split(
|
||||
prompt_len: int, num_new_tokens: int, use_eagle: bool, expected: int
|
||||
) -> None:
|
||||
(request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
|
||||
assert (
|
||||
_split(
|
||||
request,
|
||||
num_new_tokens,
|
||||
use_eagle=use_eagle,
|
||||
num_prefill_checkpoint_blocks=1,
|
||||
)
|
||||
== expected
|
||||
)
|
||||
if not use_eagle:
|
||||
manager = _make_hybrid_kv_cache_manager(num_prefill_checkpoint_blocks=1)
|
||||
assert manager.allocate_slots(request, expected) is not None
|
||||
mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID]
|
||||
blocks = mamba_manager.req_to_blocks[request.request_id]
|
||||
assert all(not block.is_null for block in blocks) # checkpoint + running state
|
||||
|
||||
|
||||
def _run_chunked_prefill(
|
||||
manager: KVCacheManager, request: Request, budgets: list[int]
|
||||
) -> dict[int, int]:
|
||||
@@ -211,8 +247,15 @@ def test_poisoning_is_block_size_independent(
|
||||
|
||||
@pytest.mark.parametrize("partial_hit", [False, True])
|
||||
@pytest.mark.parametrize("resume_at", [331, 1599, 1601, 2531, 3011])
|
||||
@pytest.mark.parametrize(
|
||||
("num_prefill_checkpoint_blocks", "use_eagle"),
|
||||
[(0, True), (1, False)],
|
||||
)
|
||||
def test_unaligned_resume_never_runs_past_its_block(
|
||||
partial_hit: bool, resume_at: int
|
||||
partial_hit: bool,
|
||||
resume_at: int,
|
||||
num_prefill_checkpoint_blocks: int,
|
||||
use_eagle: bool,
|
||||
) -> None:
|
||||
"""A prefill resuming mid-block must re-align before crossing a boundary.
|
||||
|
||||
@@ -227,7 +270,13 @@ def test_unaligned_resume_never_runs_past_its_block(
|
||||
pos, ends = resume_at, []
|
||||
while pos < prompt_len:
|
||||
request.num_computed_tokens = pos
|
||||
num_new = _split(request, prompt_len - pos, partial_hit=partial_hit)
|
||||
num_new = _split(
|
||||
request,
|
||||
prompt_len - pos,
|
||||
use_eagle=use_eagle,
|
||||
partial_hit=partial_hit,
|
||||
num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
|
||||
)
|
||||
assert num_new > 0, f"no progress at {pos}"
|
||||
if pos % MAMBA_BLOCK_SIZE != 0:
|
||||
block_end = (pos // MAMBA_BLOCK_SIZE + 1) * MAMBA_BLOCK_SIZE
|
||||
|
||||
@@ -1089,6 +1089,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
|
||||
use_eagle=False,
|
||||
hash_block_size=block_size,
|
||||
mamba_partial_cache_hit=False,
|
||||
mamba_has_prefill_checkpoint_blocks=False,
|
||||
)
|
||||
req_2.shared_prefix_boundary = shared_prefix_boundary
|
||||
num_new_tokens_adjusted = Scheduler._mamba_block_aligned_split(
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import replace
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
@@ -45,7 +46,10 @@ from vllm.models.kimi_k3.nvidia.kda_metadata import (
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
|
||||
from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
|
||||
from vllm.triton_utils import tl, triton
|
||||
from vllm.v1.attention.backend import AttentionBackend
|
||||
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
|
||||
from vllm.v1.kv_cache_interface import MambaSpec
|
||||
from vllm.v1.worker.workspace import current_workspace_manager
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -199,6 +203,8 @@ def _flashkda_prefill(
|
||||
out: torch.Tensor,
|
||||
final_state: torch.Tensor,
|
||||
workspace: torch.Tensor,
|
||||
checkpoint_state: torch.Tensor | None = None,
|
||||
checkpoint_offsets: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
import vllm._flashkda_C # noqa: F401
|
||||
|
||||
@@ -221,10 +227,75 @@ def _flashkda_prefill(
|
||||
initial_state.contiguous(),
|
||||
final_state,
|
||||
cu_seqlens.contiguous(),
|
||||
checkpoint_state,
|
||||
checkpoint_offsets.contiguous() if checkpoint_offsets is not None else None,
|
||||
)
|
||||
return out, final_state
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _store_cache_checkpoints_kernel(
|
||||
x_ptr,
|
||||
conv_state_ptr,
|
||||
recurrent_checkpoint_ptr,
|
||||
recurrent_state_ptr,
|
||||
query_start_loc_ptr,
|
||||
checkpoint_offsets_ptr,
|
||||
checkpoint_state_indices_ptr,
|
||||
x_stride_0: tl.constexpr,
|
||||
x_stride_1: tl.constexpr,
|
||||
state_stride_0: tl.constexpr,
|
||||
state_stride_1: tl.constexpr,
|
||||
state_stride_2: tl.constexpr,
|
||||
checkpoint_stride_0: tl.constexpr,
|
||||
recurrent_state_stride_0: tl.constexpr,
|
||||
checkpoint_offset_stride: tl.constexpr,
|
||||
STATE_LEN: tl.constexpr,
|
||||
WIDTH: tl.constexpr,
|
||||
RECURRENT_ROW_SIZE: tl.constexpr,
|
||||
NULL_STATE_IDX: tl.constexpr,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
):
|
||||
# store checkpoints to cache
|
||||
seq_idx = tl.program_id(0)
|
||||
cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx)
|
||||
checkpoint_offset = tl.load(
|
||||
checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride
|
||||
)
|
||||
valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0)
|
||||
valid_conv = (
|
||||
(cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN)
|
||||
)
|
||||
width_idx = cols // STATE_LEN
|
||||
history_idx = cols % STATE_LEN
|
||||
checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset
|
||||
token_idx = checkpoint_end - STATE_LEN + history_idx
|
||||
values = tl.load(
|
||||
x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1,
|
||||
mask=valid_conv,
|
||||
)
|
||||
tl.store(
|
||||
conv_state_ptr
|
||||
+ state_idx * state_stride_0
|
||||
+ width_idx * state_stride_1
|
||||
+ history_idx * state_stride_2,
|
||||
values,
|
||||
mask=valid_conv,
|
||||
)
|
||||
|
||||
valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint
|
||||
recurrent = tl.load(
|
||||
recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols,
|
||||
mask=valid_recurrent,
|
||||
)
|
||||
tl.store(
|
||||
recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols,
|
||||
recurrent,
|
||||
mask=valid_recurrent,
|
||||
)
|
||||
|
||||
|
||||
def resolve_kda_prefill_backend(
|
||||
backend: str,
|
||||
head_dim: int,
|
||||
@@ -471,6 +542,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
|
||||
self._flashkda_buffer_specs = (
|
||||
((1, T, H, D), self.model_config.dtype),
|
||||
((N, H, D, D), self.get_state_dtype()[1]),
|
||||
((N, H, D, D), self.get_state_dtype()[1]),
|
||||
((workspace_size,), torch.uint8),
|
||||
)
|
||||
|
||||
@@ -518,6 +590,14 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
|
||||
raise ValueError(f"Duplicate layer name: {prefix}")
|
||||
compilation_config.static_forward_context[prefix] = self
|
||||
|
||||
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:
|
||||
spec = super().get_kv_cache_spec(vllm_config)
|
||||
assert isinstance(spec, MambaSpec)
|
||||
return replace(
|
||||
spec,
|
||||
num_prefill_checkpoint_blocks=int(self.kda_prefill_backend == "flashkda"),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
@@ -590,6 +670,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
|
||||
spec_query_start_loc = m.spec_query_start_loc
|
||||
num_accepted_tokens = m.num_accepted_tokens
|
||||
num_actual_tokens = m.num_actual_tokens
|
||||
checkpoint = m.checkpoint
|
||||
has_spec_decode = m.num_spec_decodes > 0
|
||||
mixed_qkv = mixed_qkv[:num_actual_tokens]
|
||||
g1 = g1[:, :num_actual_tokens]
|
||||
@@ -780,7 +861,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
|
||||
if self.kda_prefill_backend == "flashkda":
|
||||
assert self.gate_lower_bound is not None
|
||||
assert self._flashkda_buffer_specs is not None
|
||||
workspace_out, final_state, workspace = (
|
||||
workspace_out, final_state, checkpoint_state, workspace = (
|
||||
current_workspace_manager().get_simultaneous(
|
||||
*self._flashkda_buffer_specs
|
||||
)
|
||||
@@ -788,6 +869,67 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
|
||||
flashkda_out = (
|
||||
workspace_out if has_spec_decode else core_attn_out
|
||||
)[:, : q_ns.shape[1]]
|
||||
if checkpoint is not None:
|
||||
assert non_spec_query_start_loc is not None
|
||||
num_sequences = initial_state.shape[0]
|
||||
assert checkpoint.checkpoint_offsets.shape == (num_sequences,)
|
||||
final_state = final_state[:num_sequences]
|
||||
checkpoint_state = checkpoint_state[:num_sequences]
|
||||
checkpoint_offsets = checkpoint.checkpoint_offsets
|
||||
_flashkda_prefill(
|
||||
q=q_ns,
|
||||
k=k_ns,
|
||||
v=v_ns,
|
||||
g=g1_ns,
|
||||
beta=beta_ns,
|
||||
A_log=self.A_log,
|
||||
dt_bias=self.dt_bias,
|
||||
lower_bound=self.gate_lower_bound,
|
||||
initial_state=initial_state,
|
||||
cu_seqlens=non_spec_query_start_loc,
|
||||
out=flashkda_out,
|
||||
final_state=final_state,
|
||||
workspace=workspace,
|
||||
checkpoint_state=checkpoint_state,
|
||||
checkpoint_offsets=checkpoint_offsets,
|
||||
)
|
||||
core_attn_out_non_spec = flashkda_out
|
||||
last_recurrent_state = final_state
|
||||
state_len = conv_state.shape[-1]
|
||||
width = mixed_qkv_ns.shape[-1]
|
||||
recurrent_row_size = checkpoint_state[0].numel()
|
||||
block_size = 256
|
||||
_store_cache_checkpoints_kernel[
|
||||
(
|
||||
checkpoint_offsets.numel(),
|
||||
triton.cdiv(
|
||||
max(width * state_len, recurrent_row_size),
|
||||
block_size,
|
||||
),
|
||||
)
|
||||
](
|
||||
mixed_qkv_ns,
|
||||
conv_state,
|
||||
checkpoint_state,
|
||||
recurrent_state,
|
||||
non_spec_query_start_loc,
|
||||
checkpoint_offsets,
|
||||
checkpoint.state_indices,
|
||||
mixed_qkv_ns.stride(0),
|
||||
mixed_qkv_ns.stride(1),
|
||||
conv_state.stride(0),
|
||||
conv_state.stride(1),
|
||||
conv_state.stride(2),
|
||||
checkpoint_state.stride(0),
|
||||
recurrent_state.stride(0),
|
||||
checkpoint_offsets.stride(0),
|
||||
state_len,
|
||||
width,
|
||||
recurrent_row_size,
|
||||
NULL_BLOCK_ID,
|
||||
block_size,
|
||||
)
|
||||
else:
|
||||
(
|
||||
core_attn_out_non_spec,
|
||||
last_recurrent_state,
|
||||
|
||||
@@ -44,6 +44,9 @@ if TYPE_CHECKING:
|
||||
)
|
||||
|
||||
|
||||
FLASHKDA_CHUNK_SIZE = 16
|
||||
|
||||
|
||||
@cache
|
||||
def _metadata_launch_pdl() -> bool:
|
||||
return current_platform.is_arch_support_pdl()
|
||||
@@ -257,12 +260,19 @@ class KDARecoverSSMCommitMetadata:
|
||||
align: KDARecoverSSMAlignMetadata | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class KDACheckpointMetadata:
|
||||
checkpoint_offsets: torch.Tensor
|
||||
state_indices: torch.Tensor
|
||||
|
||||
|
||||
@dataclass
|
||||
class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
|
||||
recoverssm_commit: KDARecoverSSMCommitMetadata | None = None
|
||||
recoverssm_context: "KDARecoverSSMCommitContext | None" = field(
|
||||
default=None, repr=False, compare=False
|
||||
)
|
||||
checkpoint: KDACheckpointMetadata | None = None
|
||||
|
||||
def commit_recoverssm_state(
|
||||
self, num_accepted_tokens: torch.Tensor
|
||||
@@ -557,6 +567,62 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
|
||||
else:
|
||||
has_initial_state = None
|
||||
|
||||
checkpoint = None
|
||||
if (
|
||||
num_prefills > 0
|
||||
and self.kv_cache_spec.num_prefill_checkpoint_blocks > 0
|
||||
and self.vllm_config.cache_config.mamba_cache_mode == "align"
|
||||
):
|
||||
# prepare checkpoint metadata
|
||||
assert m.seq_lens_cpu_upper_bound is not None
|
||||
request_rows = list(range(m.num_reqs))
|
||||
if spec_sequence_masks_cpu is not None:
|
||||
# get non spec request rows
|
||||
request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist()
|
||||
all_query_lens = query_start_loc_cpu.diff().tolist()
|
||||
query_lens = [all_query_lens[row] for row in request_rows]
|
||||
seq_lens = m.seq_lens_cpu_upper_bound.tolist()
|
||||
block_size = self.kv_cache_spec.block_size
|
||||
checkpoint_splits = []
|
||||
checkpoint_cols = []
|
||||
for row, query_len in zip(request_rows, query_lens):
|
||||
seq_len = seq_lens[row]
|
||||
offset = seq_len // block_size * block_size - (seq_len - query_len)
|
||||
# offset should be less than query_len
|
||||
valid = (
|
||||
seq_len % block_size != 0
|
||||
and 0 < offset < query_len
|
||||
and offset % FLASHKDA_CHUNK_SIZE == 0
|
||||
)
|
||||
offset = offset if valid else 0
|
||||
first_len = offset or query_len
|
||||
checkpoint_splits.append((first_len, query_len - first_len))
|
||||
checkpoint_cols.append(seq_len // block_size - 1 if valid else -1)
|
||||
if any(tail for _, tail in checkpoint_splits):
|
||||
checkpoint_offsets_tensor = async_tensor_h2d(
|
||||
[first if tail else 0 for first, tail in checkpoint_splits],
|
||||
dtype=torch.int32,
|
||||
device=query_start_loc.device,
|
||||
)
|
||||
request_rows_tensor = async_tensor_h2d(
|
||||
request_rows, dtype=torch.int64, device=query_start_loc.device
|
||||
)
|
||||
checkpoint_cols_tensor = async_tensor_h2d(
|
||||
checkpoint_cols, dtype=torch.int64, device=query_start_loc.device
|
||||
)
|
||||
checkpoint_state_indices = m.block_table_tensor[
|
||||
request_rows_tensor, checkpoint_cols_tensor
|
||||
]
|
||||
checkpoint_state_indices = torch.where(
|
||||
checkpoint_cols_tensor >= 0,
|
||||
checkpoint_state_indices,
|
||||
NULL_BLOCK_ID,
|
||||
)
|
||||
checkpoint = KDACheckpointMetadata(
|
||||
checkpoint_offsets_tensor,
|
||||
checkpoint_state_indices,
|
||||
)
|
||||
|
||||
# Prepare per-request tensors for cudagraph replay. num_actual_tokens
|
||||
# may be token-padded, while state/query/acceptance metadata is indexed
|
||||
# by request.
|
||||
@@ -647,6 +713,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
|
||||
nums_dict=nums_dict,
|
||||
batch_ptr=batch_ptr,
|
||||
token_chunk_offset_ptr=token_chunk_offset_ptr,
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -53,7 +53,7 @@ from vllm.v1.core.sched.request_queue import (
|
||||
)
|
||||
from vllm.v1.core.sched.utils import check_stop, remove_all
|
||||
from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
|
||||
from vllm.v1.kv_cache_interface import KVCacheConfig
|
||||
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
|
||||
from vllm.v1.metrics.perf import ModelMetrics, PerfStats
|
||||
from vllm.v1.metrics.stats import (
|
||||
PrefixCacheStats,
|
||||
@@ -324,6 +324,16 @@ class Scheduler(SchedulerInterface):
|
||||
self.need_mamba_block_aligned_split = (
|
||||
self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align"
|
||||
)
|
||||
self.mamba_has_prefill_checkpoint_blocks = (
|
||||
self.has_mamba_layers
|
||||
# TODO: support spec decoding
|
||||
and not self.use_eagle
|
||||
and all(
|
||||
not isinstance(group.kv_cache_spec, MambaSpec)
|
||||
or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0
|
||||
for group in kv_cache_config.kv_cache_groups
|
||||
)
|
||||
)
|
||||
# A finer prefix_match_unit is configured: a mamba partial tail entry
|
||||
# can only be registered by a step ending exactly at the prompt's last
|
||||
# hash boundary, so the split adds that stop.
|
||||
@@ -406,12 +416,17 @@ class Scheduler(SchedulerInterface):
|
||||
last_cache_position = max(last_cache_position - block_size, 0)
|
||||
|
||||
end = start + num_new_tokens
|
||||
use_internal_checkpoint = (
|
||||
self.mamba_has_prefill_checkpoint_blocks and start % block_size == 0
|
||||
)
|
||||
if use_internal_checkpoint:
|
||||
last_cache_position = 0
|
||||
# Invariant: slot p holds the state after exactly (p + 1) * block_size
|
||||
# tokens. State is written at chunk ends, so chunk ends must be block
|
||||
# aligned. Exempt: the prompt's last chunk, whose slot decode advances
|
||||
# to the boundary. A block too wide for one chunk advances sub-block
|
||||
# and re-aligns at the next boundary.
|
||||
if end < prefill_end:
|
||||
if end < prefill_end and not use_internal_checkpoint:
|
||||
max_prefill_tokens = self.max_num_scheduled_tokens
|
||||
long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold
|
||||
if long_prefill_threshold > 0:
|
||||
|
||||
@@ -1285,6 +1285,9 @@ class MambaManager(SingleTypeKVCacheManager):
|
||||
self.last_state_block_idx: dict[str, int] = {}
|
||||
# The set of the requests that have been allocated blocks
|
||||
self._allocated_block_reqs: set[str] = set()
|
||||
# Number of internal checkpoint blocks required by each request's
|
||||
# current allocation.
|
||||
self._num_checkpoint_blocks: dict[str, int] = {}
|
||||
# Requests that registered their own last-prompt-boundary partial
|
||||
# tail (producers). On the next step's CoW the boundary state moves
|
||||
# into a private cow_block; we record that block for connector
|
||||
@@ -1464,6 +1467,23 @@ class MambaManager(SingleTypeKVCacheManager):
|
||||
"""
|
||||
return 0
|
||||
|
||||
def _needs_internal_checkpoint(
|
||||
self,
|
||||
request_id: str,
|
||||
num_tokens: int,
|
||||
num_computed_tokens: int,
|
||||
) -> bool:
|
||||
assert isinstance(self.kv_cache_spec, MambaSpec)
|
||||
checkpoint_idx = cdiv(num_tokens, self.block_size) - 2
|
||||
blocks = self.req_to_blocks[request_id]
|
||||
return (
|
||||
self.kv_cache_spec.num_prefill_checkpoint_blocks > 0
|
||||
and num_tokens % self.block_size != 0
|
||||
and num_computed_tokens % self.block_size == 0
|
||||
and checkpoint_idx >= 0
|
||||
and (checkpoint_idx >= len(blocks) or blocks[checkpoint_idx].is_null)
|
||||
)
|
||||
|
||||
def get_num_blocks_to_allocate(
|
||||
self,
|
||||
request_id: str,
|
||||
@@ -1526,18 +1546,17 @@ class MambaManager(SingleTypeKVCacheManager):
|
||||
)
|
||||
if has_partial_hit:
|
||||
num_new_blocks = max(num_new_blocks, 0) + 1
|
||||
if num_new_blocks > 0:
|
||||
if request_id in self._allocated_block_reqs:
|
||||
# Old request. Needs at most 1 more blocks as we can reuse the
|
||||
# speculative blocks in previous step.
|
||||
num_new_blocks = 1 + int(has_partial_hit)
|
||||
else:
|
||||
# First prefill. Allocate 1 block for running state, the
|
||||
# speculative blocks, and one extra block if a partial cache
|
||||
# hit must be copy-on-written before the new tokens run.
|
||||
num_new_blocks = (
|
||||
1 + self.num_speculative_blocks + int(has_partial_hit)
|
||||
checkpoint_block = int(
|
||||
self._needs_internal_checkpoint(
|
||||
request_id, num_tokens, total_computed_tokens
|
||||
)
|
||||
)
|
||||
if not apply_admission_cap:
|
||||
self._num_checkpoint_blocks[request_id] = checkpoint_block
|
||||
if num_new_blocks > 0:
|
||||
num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block
|
||||
if request_id not in self._allocated_block_reqs:
|
||||
num_new_blocks += self.num_speculative_blocks
|
||||
|
||||
num_evictable_computed_blocks = self._get_num_evictable_blocks(
|
||||
new_computed_blocks
|
||||
@@ -1569,6 +1588,7 @@ class MambaManager(SingleTypeKVCacheManager):
|
||||
num_required_blocks = (
|
||||
cdiv(num_tokens, self.block_size) + self.num_speculative_blocks
|
||||
)
|
||||
checkpoint_block = self._num_checkpoint_blocks.get(request_id, 0)
|
||||
partial_hit = self._partial_hit_reqs.get(request_id)
|
||||
has_partial_hit = partial_hit is not None
|
||||
# `num_required_blocks` might be less than `len(req_blocks)` if blocks are
|
||||
@@ -1596,11 +1616,11 @@ class MambaManager(SingleTypeKVCacheManager):
|
||||
)
|
||||
# null blocks
|
||||
if prev_block_len < num_skipped_blocks:
|
||||
# minus the internal checkpoint block
|
||||
# so we don't set null for that block
|
||||
null_end = num_skipped_blocks - checkpoint_block
|
||||
req_blocks.extend(
|
||||
[
|
||||
self._null_block
|
||||
for _ in range(prev_block_len, num_skipped_blocks)
|
||||
]
|
||||
[self._null_block for _ in range(prev_block_len, null_end)]
|
||||
)
|
||||
|
||||
if blocks_allocated:
|
||||
@@ -1616,12 +1636,10 @@ class MambaManager(SingleTypeKVCacheManager):
|
||||
num_new_blocks = num_required_blocks - len(req_blocks)
|
||||
if has_partial_hit:
|
||||
num_new_blocks = max(num_new_blocks, 0) + 1
|
||||
if blocks_allocated:
|
||||
assert num_new_blocks <= 1 + int(has_partial_hit)
|
||||
else:
|
||||
assert num_new_blocks <= self.num_speculative_blocks + 1 + int(
|
||||
has_partial_hit
|
||||
)
|
||||
max_new_blocks = 1 + int(has_partial_hit) + checkpoint_block
|
||||
if not blocks_allocated:
|
||||
max_new_blocks += self.num_speculative_blocks
|
||||
assert num_new_blocks <= max_new_blocks
|
||||
new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
|
||||
returned_blocks = req_blocks[prev_block_len:]
|
||||
if partial_hit is not None:
|
||||
@@ -1670,6 +1688,7 @@ class MambaManager(SingleTypeKVCacheManager):
|
||||
if self.mamba_cache_mode == "align":
|
||||
self._allocated_block_reqs.discard(request_id)
|
||||
self.last_state_block_idx.pop(request_id, None)
|
||||
self._num_checkpoint_blocks.pop(request_id, None)
|
||||
self._producer_partial_tail_reqs.pop(request_id, None)
|
||||
# A hand-off whose request died in this same scheduling pass must
|
||||
# not reach the connector: its unpin hook (free) has already run.
|
||||
|
||||
@@ -815,6 +815,7 @@ class MambaSpec(KVCacheSpec):
|
||||
mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2
|
||||
mamba_cache_mode: str = "none"
|
||||
num_speculative_blocks: int = 0
|
||||
num_prefill_checkpoint_blocks: int = 0
|
||||
num_heads: int = 1
|
||||
tokens_per_state: int = -1
|
||||
|
||||
@@ -843,7 +844,9 @@ class MambaSpec(KVCacheSpec):
|
||||
cdiv(max_model_len, self.block_size) + self.num_speculative_blocks
|
||||
) * self.page_size_bytes
|
||||
elif vllm_config.cache_config.mamba_cache_mode == "align":
|
||||
return self.page_size_bytes * (2 + self.num_speculative_blocks)
|
||||
return self.page_size_bytes * (
|
||||
2 + self.num_speculative_blocks + self.num_prefill_checkpoint_blocks
|
||||
)
|
||||
else:
|
||||
return self.page_size_bytes * (1 + self.num_speculative_blocks)
|
||||
|
||||
@@ -865,6 +868,7 @@ class MambaSpec(KVCacheSpec):
|
||||
return all(
|
||||
isinstance(spec, MambaSpec)
|
||||
and spec.num_speculative_blocks == self.num_speculative_blocks
|
||||
and spec.num_prefill_checkpoint_blocks == self.num_prefill_checkpoint_blocks
|
||||
for spec in kv_cache_specs.values()
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user