Signed-off-by: Shenglei Fu <sfu@confluent.io> Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Shenglei FuandClaude authored and GitHub committed
1 parent
bfaaa3c2c3
commit
f24a6165ff
4 files changed
+718
No files matched your search
@@ -1,3 +1,6 @@
|
||||
import math
|
||||
from typing import NamedTuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -317,6 +320,103 @@ def test_deferred_multimodal_cache_write_commits_on_completion():
|
||||
assert torch.equal(rows[8:16], expected)
|
||||
|
||||
|
||||
def test_discard_deferred_mm_outputs_drops_pending_chunks():
|
||||
"""discard_deferred_mm_outputs is the abort path: a request that
|
||||
leaves without a cache commit must not leak its staged chunks into
|
||||
the cache if commit_deferred_mm_outputs is later called for it."""
|
||||
cache = get_omni_pcache()
|
||||
mm_key = "codes.audio"
|
||||
block_table = torch.tensor(
|
||||
[
|
||||
[2, 3, 4, 5],
|
||||
[6, 7, 8, 9],
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
input_batch = MockInputBatch(
|
||||
num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
|
||||
block_table=block_table,
|
||||
)
|
||||
|
||||
cache.update_omni_tensor_prefix_cache(
|
||||
hidden_states=None,
|
||||
multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)},
|
||||
num_tokens_unpadded=4,
|
||||
slot_mapping=torch.arange(8, 12),
|
||||
skip_mm_cache_keys={mm_key},
|
||||
)
|
||||
|
||||
chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
|
||||
cache.stage_deferred_mm_outputs(
|
||||
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs={mm_key: chunk},
|
||||
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||
deferred_mm_cache_keys={mm_key},
|
||||
)
|
||||
assert "req1" in cache._deferred_mm_outputs
|
||||
|
||||
cache.discard_deferred_mm_outputs("req1")
|
||||
assert "req1" not in cache._deferred_mm_outputs
|
||||
|
||||
# A later commit call for the now-discarded req_id must be a no-op,
|
||||
# not a resurrection of the dropped chunk.
|
||||
cache.commit_deferred_mm_outputs({"req1"}, input_batch)
|
||||
rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
|
||||
assert torch.all(rows[8:12] == 0)
|
||||
|
||||
|
||||
def test_deferred_mm_outputs_request_id_reused_after_discard():
|
||||
"""A request ID reused after discard_deferred_mm_outputs must start
|
||||
clean: the new occupant's committed data must not be mixed with the
|
||||
previous (discarded) occupant's staged chunks."""
|
||||
cache = get_omni_pcache()
|
||||
mm_key = "codes.audio"
|
||||
block_table = torch.tensor(
|
||||
[
|
||||
[2, 3, 4, 5],
|
||||
[6, 7, 8, 9],
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
input_batch = MockInputBatch(
|
||||
num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
|
||||
block_table=block_table,
|
||||
)
|
||||
|
||||
cache.update_omni_tensor_prefix_cache(
|
||||
hidden_states=None,
|
||||
multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)},
|
||||
num_tokens_unpadded=4,
|
||||
slot_mapping=torch.arange(8, 12),
|
||||
skip_mm_cache_keys={mm_key},
|
||||
)
|
||||
|
||||
stale_chunk = torch.full((4, 2), -1, dtype=torch.long)
|
||||
cache.stage_deferred_mm_outputs(
|
||||
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs={mm_key: stale_chunk},
|
||||
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||
deferred_mm_cache_keys={mm_key},
|
||||
)
|
||||
cache.discard_deferred_mm_outputs("req1")
|
||||
|
||||
# "req1" is reused by a brand-new request with different data.
|
||||
fresh_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
|
||||
cache.stage_deferred_mm_outputs(
|
||||
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs={mm_key: fresh_chunk},
|
||||
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||
deferred_mm_cache_keys={mm_key},
|
||||
)
|
||||
cache.commit_deferred_mm_outputs({"req1"}, input_batch)
|
||||
|
||||
rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
|
||||
assert torch.equal(rows[8:12], fresh_chunk)
|
||||
|
||||
|
||||
def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit():
|
||||
cache = get_omni_pcache()
|
||||
mm_key = "codes.audio"
|
||||
@@ -599,3 +699,219 @@ def test_get_merged_multimodal_outputs(feat_dims, num_tokens_padded, mocker):
|
||||
|
||||
assert req2_merged_mm_outputs.shape == torch.Size([num_new_toks_req2, curr_feat_dim])
|
||||
assert torch.all(req2_merged_mm_outputs == req2_new_mm_outputs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"passthrough_shape",
|
||||
[(5,), (5, 3)],
|
||||
ids=["1d_token_metadata", "2d_feature_tensor"],
|
||||
)
|
||||
def test_get_merged_multimodal_outputs_slices_passthrough_tensor_per_request(passthrough_shape, mocker):
|
||||
"""A passthrough key that is itself a real tensor (not the string/list
|
||||
passthrough already covered above) must be sliced per request through
|
||||
to_payload_element's tensor branch, under the same mixed cache
|
||||
hit/miss + uneven scheduled-length batch as test_get_merged_multimodal_outputs.
|
||||
|
||||
Covers both a 1D token-aligned metadata tensor (e.g. ref_code_len,
|
||||
codec_streaming -- shape (total_tokens,)) and a 2D per-token feature
|
||||
tensor (shape (total_tokens, feat_dim)), since W2's batching could
|
||||
easily special-case one and silently mishandle the other.
|
||||
"""
|
||||
feat_dims = {"foo": 100, "bar": 100}
|
||||
cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN)
|
||||
|
||||
orig_num_tokens_unpadded = 8
|
||||
slot_offset = 8
|
||||
orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded)
|
||||
feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()}
|
||||
orig_mm_outputs = {
|
||||
key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
|
||||
}
|
||||
|
||||
cache.update_omni_tensor_prefix_cache(
|
||||
hidden_states=None,
|
||||
multimodal_outputs=orig_mm_outputs,
|
||||
num_tokens_unpadded=orig_num_tokens_unpadded,
|
||||
slot_mapping=orig_slot_mapping,
|
||||
num_tokens_padded=None,
|
||||
)
|
||||
|
||||
num_new_toks_req1 = 3
|
||||
num_new_toks_req2 = 2
|
||||
cache.add_prefix_cached_new_req_id("req1")
|
||||
|
||||
num_scheduled_tokens = {
|
||||
"req1": num_new_toks_req1,
|
||||
"req2": num_new_toks_req2,
|
||||
}
|
||||
|
||||
new_mm_outputs = {}
|
||||
for mm_key in cache.mm_cache_keys:
|
||||
new_mm_outputs[mm_key] = torch.rand(
|
||||
(num_new_toks_req1 + num_new_toks_req2, feature_dims[mm_key]),
|
||||
dtype=DTYPE,
|
||||
)
|
||||
total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2
|
||||
assert passthrough_shape[0] == total_scheduled_tokens
|
||||
passthrough_tensor = torch.arange(int(torch.prod(torch.tensor(passthrough_shape))), dtype=DTYPE).reshape(
|
||||
passthrough_shape
|
||||
)
|
||||
new_mm_outputs["passthrough_tensor"] = passthrough_tensor
|
||||
|
||||
input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0]))
|
||||
|
||||
mocker.patch(
|
||||
"vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
|
||||
new=fake_get_cached_block_ids,
|
||||
)
|
||||
merged_mm_outputs = cache.get_merged_multimodal_states(
|
||||
query_start_loc=[0, num_new_toks_req1],
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs=new_mm_outputs,
|
||||
num_scheduled_tokens=num_scheduled_tokens,
|
||||
)
|
||||
|
||||
assert "passthrough_tensor" not in cache.mm_cache_keys
|
||||
|
||||
passthrough_out = merged_mm_outputs["passthrough_tensor"]
|
||||
# req1 (cache hit) gets tokens [0:3), req2 (cache miss) gets tokens [3:5) --
|
||||
# a passthrough tensor is sliced by scheduled position, not merged with
|
||||
# any cached history (unlike the mm_cache_keys tensors above).
|
||||
assert torch.equal(passthrough_out["req1"], passthrough_tensor[0:num_new_toks_req1])
|
||||
assert torch.equal(passthrough_out["req2"], passthrough_tensor[num_new_toks_req1:total_scheduled_tokens])
|
||||
|
||||
|
||||
class MergeHarness(NamedTuple):
|
||||
"""A primed cache plus the arguments get_merged_multimodal_states needs."""
|
||||
|
||||
cache: OmniTensorPrefixCache
|
||||
input_batch: MockInputBatch
|
||||
num_scheduled_tokens: dict[str, int]
|
||||
query_start_loc: list[int]
|
||||
total_scheduled_tokens: int
|
||||
new_mm_outputs: dict[str, object]
|
||||
|
||||
|
||||
def build_mixed_hit_miss_harness(mocker, num_new_toks_req1: int = 3, num_new_toks_req2: int = 2) -> MergeHarness:
|
||||
"""Two requests over one cache key: req1 hits 8 cached tokens, req2 misses,
|
||||
and their scheduled lengths differ -- the same batch shape as
|
||||
test_get_merged_multimodal_outputs, which is what makes a per-request
|
||||
boundary error observable. Callers add their passthrough keys to
|
||||
``new_mm_outputs`` before calling get_merged_multimodal_states.
|
||||
"""
|
||||
cache = get_omni_pcache_with_mm_tensors({"foo": 100}, seq_len=DEFAULT_SEQ_LEN)
|
||||
|
||||
orig_num_tokens_unpadded = 8
|
||||
feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()}
|
||||
cache.update_omni_tensor_prefix_cache(
|
||||
hidden_states=None,
|
||||
multimodal_outputs={
|
||||
key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
|
||||
},
|
||||
num_tokens_unpadded=orig_num_tokens_unpadded,
|
||||
slot_mapping=torch.arange(8, 8 + orig_num_tokens_unpadded),
|
||||
num_tokens_padded=None,
|
||||
)
|
||||
cache.add_prefix_cached_new_req_id("req1")
|
||||
|
||||
total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2
|
||||
mocker.patch(
|
||||
"vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
|
||||
new=fake_get_cached_block_ids,
|
||||
)
|
||||
return MergeHarness(
|
||||
cache=cache,
|
||||
input_batch=MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])),
|
||||
num_scheduled_tokens={"req1": num_new_toks_req1, "req2": num_new_toks_req2},
|
||||
query_start_loc=[0, num_new_toks_req1],
|
||||
total_scheduled_tokens=total_scheduled_tokens,
|
||||
new_mm_outputs={
|
||||
key: torch.rand((total_scheduled_tokens, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_get_merged_multimodal_outputs_passes_per_request_tensor_list_through(mocker):
|
||||
"""The passthrough shape qwen3-tts actually emits: a nested ``codes``/``meta``
|
||||
dict whose ``codes.ref`` holds one tensor per request
|
||||
(qwen3_tts_talker.make_omni_output), carried through to_payload_element's
|
||||
pass_lists_through branch beside 1D int metadata that IS sliced per request.
|
||||
|
||||
Each request must receive the list *whole*. Per-request indexing happens
|
||||
later, in the runner's _unwrap_lists; a W2 that split the list here would
|
||||
make that second index select a different request's reference codes. The
|
||||
existing list coverage passes a list of strings, which cannot catch a
|
||||
tensor-specific regression, and no other test covers an int dtype.
|
||||
"""
|
||||
harness = build_mixed_hit_miss_harness(mocker)
|
||||
|
||||
ref_codes = [torch.tensor([11, 12], dtype=torch.long), torch.tensor([21], dtype=torch.long)]
|
||||
ref_code_len = torch.arange(harness.total_scheduled_tokens, dtype=torch.int32)
|
||||
harness.new_mm_outputs["codes"] = {"ref": ref_codes}
|
||||
harness.new_mm_outputs["meta"] = {"ref_code_len": ref_code_len}
|
||||
|
||||
merged = harness.cache.get_merged_multimodal_states(
|
||||
query_start_loc=harness.query_start_loc,
|
||||
input_batch=harness.input_batch,
|
||||
multimodal_outputs=harness.new_mm_outputs,
|
||||
num_scheduled_tokens=harness.num_scheduled_tokens,
|
||||
)
|
||||
|
||||
assert "codes" not in harness.cache.mm_cache_keys
|
||||
for req_id in ("req1", "req2"):
|
||||
req_refs = merged["codes"][req_id]["ref"]
|
||||
assert len(req_refs) == len(ref_codes)
|
||||
for got, want in zip(req_refs, ref_codes):
|
||||
assert torch.equal(got, want)
|
||||
assert got.dtype == torch.long
|
||||
|
||||
# Per-request clones: the branch exists so one request's downstream mutation
|
||||
# cannot reach another request's payload or the source list.
|
||||
merged["codes"]["req1"]["ref"][0] += 100
|
||||
assert torch.equal(merged["codes"]["req2"]["ref"][0], torch.tensor([11, 12], dtype=torch.long))
|
||||
assert torch.equal(ref_codes[0], torch.tensor([11, 12], dtype=torch.long))
|
||||
|
||||
# Token-aligned 1D metadata is still split per request, unlike the list.
|
||||
req1_len = harness.num_scheduled_tokens["req1"]
|
||||
assert torch.equal(merged["meta"]["req1"]["ref_code_len"], ref_code_len[:req1_len])
|
||||
assert torch.equal(merged["meta"]["req2"]["ref_code_len"], ref_code_len[req1_len:])
|
||||
assert merged["meta"]["req1"]["ref_code_len"].dtype == torch.int32
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"passthrough_shape",
|
||||
[(3, 4), (1, 4), (7,)],
|
||||
ids=["shorter_2d", "single_row_2d", "longer_1d"],
|
||||
)
|
||||
def test_get_merged_multimodal_outputs_broadcasts_non_token_aligned_passthrough(passthrough_shape, mocker):
|
||||
"""A passthrough tensor whose first dimension is not the scheduled-token
|
||||
count is request-invariant: to_payload_element hands every request the whole
|
||||
tensor, cloned, rather than slicing it.
|
||||
|
||||
W2 batching most plausibly breaks by assuming every passthrough tensor is
|
||||
token-aligned; that assumption would slice these and give each request a
|
||||
truncated view, or an empty one once start exceeds the first dimension.
|
||||
"""
|
||||
harness = build_mixed_hit_miss_harness(mocker)
|
||||
assert passthrough_shape[0] != harness.total_scheduled_tokens
|
||||
|
||||
passthrough = torch.arange(math.prod(passthrough_shape), dtype=DTYPE).reshape(passthrough_shape)
|
||||
expected = passthrough.clone()
|
||||
harness.new_mm_outputs["speaker_embedding"] = passthrough
|
||||
|
||||
merged = harness.cache.get_merged_multimodal_states(
|
||||
query_start_loc=harness.query_start_loc,
|
||||
input_batch=harness.input_batch,
|
||||
multimodal_outputs=harness.new_mm_outputs,
|
||||
num_scheduled_tokens=harness.num_scheduled_tokens,
|
||||
)
|
||||
|
||||
assert "speaker_embedding" not in harness.cache.mm_cache_keys
|
||||
out = merged["speaker_embedding"]
|
||||
for req_id in ("req1", "req2"):
|
||||
assert out[req_id].shape == passthrough_shape
|
||||
assert torch.equal(out[req_id], expected)
|
||||
|
||||
out["req1"] += 100
|
||||
assert torch.equal(out["req2"], expected)
|
||||
assert torch.equal(passthrough, expected)
|
||||
Reference in new issue
Block a user