1489 lines
53 KiB
Python
1489 lines
53 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from contextlib import nullcontext
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from torch import nn
|
|
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
|
|
|
from tests.helpers.kv_layout import layout_for_backend
|
|
from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
|
|
from vllm_omni.diffusion.attention.layer import Attention
|
|
from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module
|
|
from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
|
|
DiffusionPagedAttentionAdapter,
|
|
DiffusionPagedAttentionMetadata,
|
|
DiffusionPagedAttentionRow,
|
|
DiffusionPagedAttentionRowBinding,
|
|
DiffusionPagedAttentionRuntime,
|
|
)
|
|
from vllm_omni.diffusion.forward_context import (
|
|
get_forward_context,
|
|
override_paged_kv_adapter,
|
|
set_forward_context,
|
|
set_forward_context_denoise_step_idx,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
class _FakeBlockTables:
|
|
def __init__(self) -> None:
|
|
self.max_num_reqs = 4
|
|
self.max_num_batched_tokens = 16
|
|
self.cp_size = 1
|
|
self.cp_rank = 0
|
|
self.cp_interleave = 1
|
|
self.block_sizes = [4]
|
|
self.blocks_per_kv_block = [1]
|
|
self.num_blocks = SimpleNamespace(np=np.full((1, self.max_num_reqs), 4, dtype=np.int32))
|
|
self.gather_calls: list[tuple[torch.Tensor, int]] = []
|
|
self.slot_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]] = []
|
|
|
|
def gather_block_tables(
|
|
self,
|
|
idx_mapping: torch.Tensor,
|
|
num_reqs_padded: int,
|
|
) -> tuple[torch.Tensor, ...]:
|
|
self.gather_calls.append((idx_mapping.clone(), num_reqs_padded))
|
|
return (torch.tensor([[3, 4], [7, 8]], dtype=torch.int32),)
|
|
|
|
def compute_slot_mappings(
|
|
self,
|
|
idx_mapping: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
num_tokens_padded: int,
|
|
) -> torch.Tensor:
|
|
self.slot_calls.append(
|
|
(
|
|
idx_mapping.clone(),
|
|
query_start_loc.clone(),
|
|
positions.clone(),
|
|
num_tokens_padded,
|
|
)
|
|
)
|
|
return positions[:num_tokens_padded].to(torch.int64).unsqueeze(0)
|
|
|
|
|
|
class _FakeSpec:
|
|
non_causal: bool
|
|
block_size = 4
|
|
|
|
def __init__(self, *, non_causal: bool) -> None:
|
|
self.non_causal = non_causal
|
|
|
|
def max_num_blocks_per_req(self, _vllm_config, max_len: int) -> int:
|
|
return math.ceil(max_len / self.block_size)
|
|
|
|
|
|
class _FakeAttentionGroup:
|
|
def __init__(self, reorder_batch_threshold: int | None) -> None:
|
|
self.builder = SimpleNamespace(reorder_batch_threshold=reorder_batch_threshold)
|
|
|
|
def get_metadata_builder(self, _index: int):
|
|
return self.builder
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _FakeNativeMetadata:
|
|
build_id: int
|
|
causal: bool
|
|
seq_lens: torch.Tensor
|
|
query_start_loc_cpu: torch.Tensor
|
|
positions: torch.Tensor
|
|
slot_mappings: torch.Tensor
|
|
max_num_splits: int = 0
|
|
|
|
|
|
class _FakeLayer:
|
|
def __init__(self, *, non_causal: bool, num_heads: int = 2, num_kv_heads: int = 2) -> None:
|
|
self.num_heads = num_heads
|
|
self.num_kv_heads = num_kv_heads
|
|
self.head_size = 4
|
|
self.head_size_v = 4
|
|
self.spec = _FakeSpec(non_causal=non_causal)
|
|
self.updates: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
|
|
self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, _FakeNativeMetadata]] = []
|
|
self.native_events: list[str] = []
|
|
self.layer_name = "layer-0"
|
|
self.kv_cache = object()
|
|
self.attn_backend = SimpleNamespace(forward_includes_kv_cache_update=False)
|
|
self.impl = _FakeNativeImpl(self)
|
|
|
|
|
|
class _FakeNativeImpl:
|
|
vllm_flash_attn_version: int
|
|
|
|
def __init__(self, layer: _FakeLayer) -> None:
|
|
self.layer = layer
|
|
|
|
def do_kv_cache_update(
|
|
self,
|
|
_layer,
|
|
key: torch.Tensor | None,
|
|
value: torch.Tensor | None,
|
|
_kv_cache,
|
|
slot_mapping: torch.Tensor,
|
|
) -> None:
|
|
self.layer.native_events.append("update")
|
|
self.layer.updates.append((key, value, slot_mapping))
|
|
|
|
def forward(
|
|
self,
|
|
_layer,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
_kv_cache,
|
|
metadata: _FakeNativeMetadata,
|
|
output: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
self.layer.native_events.append("forward")
|
|
self.layer.calls.append((query, key, value, metadata))
|
|
return output.copy_(query.reshape_as(output))
|
|
|
|
|
|
def _make_adapter(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
*,
|
|
non_causal: bool = True,
|
|
reorder_batch_threshold: int | None = None,
|
|
num_heads: int = 2,
|
|
num_kv_heads: int = 2,
|
|
capacity: int = 16,
|
|
) -> tuple[DiffusionPagedAttentionAdapter, _FakeBlockTables, _FakeLayer, list[tuple]]:
|
|
events: list[tuple] = []
|
|
monkeypatch.setattr(
|
|
adapter_module.current_omni_platform,
|
|
"requires_diffusion_paged_kv_prewrite",
|
|
lambda: False,
|
|
)
|
|
block_tables = _FakeBlockTables()
|
|
block_tables.max_num_batched_tokens = capacity
|
|
block_tables.num_blocks.np.fill(math.ceil(capacity / block_tables.block_sizes[0]))
|
|
layer = _FakeLayer(non_causal=non_causal, num_heads=num_heads, num_kv_heads=num_kv_heads)
|
|
kv_cache_config = SimpleNamespace(
|
|
kv_cache_groups=[SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=layer.spec)],
|
|
)
|
|
row_map = {
|
|
("req-0", 0, None): DiffusionPagedAttentionRowBinding(2, capacity),
|
|
("req-1", 1, None): DiffusionPagedAttentionRowBinding(3, capacity),
|
|
("req-0", None, "text"): DiffusionPagedAttentionRowBinding(1, capacity),
|
|
}
|
|
|
|
monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
|
|
|
|
def build_metadata(**kwargs):
|
|
metadata = _FakeNativeMetadata(
|
|
build_id=len(events),
|
|
causal=kwargs["causal"],
|
|
seq_lens=kwargs["seq_lens"].clone(),
|
|
query_start_loc_cpu=kwargs["query_start_loc_cpu"].clone(),
|
|
positions=kwargs["positions"].clone(),
|
|
slot_mappings=kwargs["slot_mappings"].clone(),
|
|
)
|
|
events.append(("build", kwargs, metadata))
|
|
return {"layer-0": metadata}
|
|
|
|
monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata)
|
|
monkeypatch.setattr(
|
|
adapter_module,
|
|
"build_slot_mappings_by_layer",
|
|
lambda slot_mappings, _config: {"layer-0": slot_mappings[0]},
|
|
)
|
|
|
|
config = SimpleNamespace(
|
|
name="vllm-config",
|
|
model_config=SimpleNamespace(dtype=torch.float32, rswa_window=None),
|
|
)
|
|
adapter = DiffusionPagedAttentionAdapter(
|
|
vllm_config=config,
|
|
device=torch.device("cpu"),
|
|
kv_cache_config=kv_cache_config,
|
|
block_tables=block_tables,
|
|
attn_groups=[[_FakeAttentionGroup(reorder_batch_threshold)]],
|
|
layers={"layer-0": layer},
|
|
resolve_row=lambda request_id, sequence_id, context_id: row_map[(request_id, sequence_id, context_id)],
|
|
)
|
|
return adapter, block_tables, layer, events
|
|
|
|
|
|
def _run_omni_paged_backend(context):
|
|
backend = FlashAttentionImpl(
|
|
num_heads=context.layer.num_heads,
|
|
head_size=context.layer.head_size,
|
|
softmax_scale=0.5,
|
|
num_kv_heads=context.layer.num_kv_heads,
|
|
)
|
|
return backend.forward_paged(context)
|
|
|
|
|
|
def test_prepare_batch_reuses_native_block_table_and_metadata_builders(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, block_tables, _, events = _make_adapter(monkeypatch)
|
|
rows = (
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=3,
|
|
seq_len=8,
|
|
kv_start_pos=4,
|
|
),
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-1",
|
|
sequence_id=1,
|
|
query_len=2,
|
|
seq_len=2,
|
|
),
|
|
)
|
|
|
|
batch = adapter.prepare_batch(rows)
|
|
|
|
assert batch.num_tokens == 5
|
|
assert batch.row_indices.tolist() == [2, 3]
|
|
assert batch.query_start_loc.tolist() == [0, 3, 5]
|
|
assert batch.seq_lens.tolist() == [8, 2]
|
|
assert batch.positions.tolist() == [4, 5, 6, 0, 1]
|
|
assert block_tables.gather_calls[0][1] == 2
|
|
assert block_tables.gather_calls[0][0].tolist() == [2, 3]
|
|
assert block_tables.slot_calls[0][3] == 5
|
|
build_kwargs = events[0][1]
|
|
assert build_kwargs["causal"] is False
|
|
assert build_kwargs["max_query_len"] == 3
|
|
assert build_kwargs["max_seq_len"] == 8
|
|
assert build_kwargs["positions"] is batch.positions
|
|
assert batch.attn_metadata["layer-0"] is events[0][2]
|
|
assert events[0][2].build_id == 0
|
|
assert batch.slot_mappings_by_layer["layer-0"].tolist() == [4, 5, 6, 0, 1]
|
|
|
|
|
|
def test_omni_paged_backend_consumes_context_and_restores_diffusion_shape(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, _, layer, events = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=5,
|
|
seq_len=5,
|
|
)
|
|
]
|
|
)
|
|
query = torch.randn(1, 5, 2, 4)
|
|
key = torch.randn_like(query)
|
|
value = torch.randn_like(query)
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context("layer-0", query, key, value)
|
|
output = _run_omni_paged_backend(context)
|
|
|
|
assert output.shape == query.shape
|
|
assert torch.equal(output, query)
|
|
assert layer.calls[0][0].shape == (5, 2, 4)
|
|
assert layer.calls[0][1] is None
|
|
assert layer.calls[0][2] is None
|
|
assert layer.native_events == ["update", "forward"]
|
|
assert layer.calls[0][3] is events[0][2]
|
|
|
|
|
|
def test_paged_adapter_accepts_native_gqa_shapes(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, layer, _ = _make_adapter(monkeypatch, num_heads=32, num_kv_heads=8)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)]
|
|
)
|
|
query = torch.randn(1, 3, 32, 4)
|
|
key = torch.randn(1, 3, 8, 4)
|
|
value = torch.randn_like(key)
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context("layer-0", query, key, value)
|
|
|
|
assert layer.num_heads == 32
|
|
assert layer.num_kv_heads == 8
|
|
assert context.query.shape == (3, 32, 4)
|
|
assert context.key_write.shape == (3, 8, 4)
|
|
assert context.value_write.shape == (3, 8, 4)
|
|
|
|
|
|
def test_paged_adapter_preserves_strided_kv_for_platform_writer(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
|
|
)
|
|
packed_qkv = torch.randn(1, 5, 3 * 2 * 4)
|
|
query, key, value = packed_qkv.split(2 * 4, dim=-1)
|
|
query = query.reshape(1, 5, 2, 4)
|
|
key = key.reshape(1, 5, 2, 4)
|
|
value = value.reshape(1, 5, 2, 4)
|
|
assert not key.is_contiguous()
|
|
assert not value.is_contiguous()
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context("layer-0", query, key, value)
|
|
|
|
assert context.query.is_contiguous()
|
|
assert not context.key_write.is_contiguous()
|
|
assert not context.value_write.is_contiguous()
|
|
assert context.key_write.data_ptr() == key.data_ptr()
|
|
assert context.value_write.data_ptr() == value.data_ptr()
|
|
|
|
|
|
def test_paged_adapter_packs_noncontiguous_kv_head_layout(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
|
|
)
|
|
query = torch.randn(1, 5, 2, 4)
|
|
key = torch.randn(1, 5, 4, 2).transpose(-1, -2)
|
|
value = torch.randn(1, 5, 4, 2).transpose(-1, -2)
|
|
assert key.shape == (1, 5, 2, 4)
|
|
assert key.stride(-1) != 1
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context("layer-0", query, key, value)
|
|
|
|
assert context.key_write.is_contiguous()
|
|
assert context.value_write.is_contiguous()
|
|
|
|
|
|
def test_omni_paged_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
|
layer.attn_backend.forward_includes_kv_cache_update = True
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)]
|
|
)
|
|
qkv = torch.randn(1, 3, 2, 4)
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
|
output = _run_omni_paged_backend(context)
|
|
|
|
assert torch.equal(output, qkv)
|
|
assert layer.native_events == ["forward"]
|
|
assert layer.updates == []
|
|
|
|
|
|
def test_omni_piecewise_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, layer, events = _make_adapter(monkeypatch)
|
|
layer.attn_backend.forward_includes_kv_cache_update = True
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
|
|
)
|
|
qkv = torch.randn(1, 4, 2, 4)
|
|
metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context(
|
|
"layer-0",
|
|
qkv,
|
|
qkv,
|
|
qkv,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
output = _run_omni_paged_backend(context)
|
|
|
|
assert torch.equal(output, qkv)
|
|
assert layer.native_events == ["forward", "forward", "forward"]
|
|
assert layer.updates == []
|
|
assert [call[0].shape[0] for call in layer.calls] == [1, 2, 1]
|
|
assert [call[1].shape[0] for call in layer.calls] == [1, 2, 1]
|
|
assert [call[2].shape[0] for call in layer.calls] == [1, 2, 1]
|
|
assert all(call[3] is event[2] for call, event in zip(layer.calls, events[1:], strict=True))
|
|
segment_metadata = [call[3] for call in layer.calls]
|
|
assert [metadata.causal for metadata in segment_metadata] == [True, False, True]
|
|
assert [metadata.seq_lens.tolist() for metadata in segment_metadata] == [[1], [3], [4]]
|
|
assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [[0, 1], [0, 2], [0, 1]]
|
|
assert [metadata.positions.tolist() for metadata in segment_metadata] == [[0], [1, 2], [3]]
|
|
assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [[[0]], [[1, 2]], [[3]]]
|
|
for event in events[1:]:
|
|
build_kwargs = event[1]
|
|
assert build_kwargs["positions"].untyped_storage().data_ptr() == batch.positions.untyped_storage().data_ptr()
|
|
assert (
|
|
build_kwargs["slot_mappings"].untyped_storage().data_ptr()
|
|
== batch.slot_mappings.untyped_storage().data_ptr()
|
|
)
|
|
|
|
|
|
def test_platform_prewrite_updates_once_then_piecewise_reads_cache(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
|
layer.attn_backend.forward_includes_kv_cache_update = True
|
|
monkeypatch.setattr(
|
|
adapter_module.current_omni_platform,
|
|
"requires_diffusion_paged_kv_prewrite",
|
|
lambda: True,
|
|
)
|
|
adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding(
|
|
row_index=2,
|
|
max_seq_len=16,
|
|
block_ids=((3, 7, 9, 11),),
|
|
)
|
|
row = DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)
|
|
batch = adapter.prepare_batch([row])
|
|
qkv = torch.randn(1, 4, 2, 4)
|
|
metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata)
|
|
output = _run_omni_paged_backend(context)
|
|
|
|
assert torch.equal(output, qkv)
|
|
assert layer.native_events == ["update", "forward", "forward", "forward"]
|
|
assert len(layer.updates) == 1
|
|
assert all(call[1] is None and call[2] is None for call in layer.calls)
|
|
|
|
|
|
def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, layer, events = _make_adapter(monkeypatch)
|
|
layer.impl.vllm_flash_attn_version = 2
|
|
batch = adapter.prepare_batch(
|
|
[
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=6,
|
|
seq_len=9,
|
|
kv_start_pos=3,
|
|
),
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-1",
|
|
sequence_id=1,
|
|
query_len=6,
|
|
seq_len=9,
|
|
kv_start_pos=3,
|
|
),
|
|
]
|
|
)
|
|
query = torch.arange(2 * 6 * 2 * 4, dtype=torch.float32).reshape(2, 6, 2, 4)
|
|
key = torch.arange(2 * 6 * 2 * 4, dtype=torch.float32).reshape(2, 6, 2, 4) + 100
|
|
value = key + 100
|
|
attn_mask = torch.ones(9, 9, dtype=torch.bool).tril().repeat(2, 1, 1)
|
|
attn_mask[:, 5:8, 5:8] = True
|
|
metadata = SimpleNamespace(
|
|
attn_mask=attn_mask.unsqueeze(1),
|
|
full_attn_spans=[[(5, 8)], [(5, 8)]],
|
|
query_ranges=None,
|
|
extra={},
|
|
)
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context(
|
|
"layer-0",
|
|
query,
|
|
key,
|
|
value,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
# Every model layer sees the same spans, so the native segment metadata
|
|
# is built once per active batch and then reused.
|
|
adapter.prepare_layer_context(
|
|
"layer-0",
|
|
query,
|
|
key,
|
|
value,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
output = _run_omni_paged_backend(context)
|
|
|
|
assert torch.equal(output, query)
|
|
assert layer.native_events == ["update", "forward", "forward", "forward"]
|
|
assert len(layer.updates) == 1
|
|
assert layer.updates[0][0].shape[0] == 12
|
|
assert torch.equal(layer.updates[0][0], key.reshape(12, 2, 4))
|
|
assert torch.equal(layer.updates[0][1], value.reshape(12, 2, 4))
|
|
assert [call[0].shape[0] for call in layer.calls] == [4, 6, 2]
|
|
assert all(call[1] is None and call[2] is None for call in layer.calls)
|
|
assert all(call[3] is event[2] for call, event in zip(layer.calls, events[1:], strict=True))
|
|
assert [metadata.build_id for metadata in context.piecewise_native_metadata] == [1, 2, 3]
|
|
assert [metadata.max_num_splits for metadata in context.piecewise_native_metadata] == [1, 1, 1]
|
|
|
|
# The first metadata build is the normal whole-query path. Piecewise calls
|
|
# then use causal [3, 5), full [5, 8), and causal [8, 9) segments.
|
|
assert len(events) == 4
|
|
segment_builds = [event[1] for event in events[1:]]
|
|
assert [build["causal"] for build in segment_builds] == [True, False, True]
|
|
assert [build["seq_lens"].tolist() for build in segment_builds] == [
|
|
[5, 5],
|
|
[8, 8],
|
|
[9, 9],
|
|
]
|
|
assert [build["query_start_loc_cpu"].tolist() for build in segment_builds] == [
|
|
[0, 2, 4],
|
|
[0, 3, 6],
|
|
[0, 1, 2],
|
|
]
|
|
segment_metadata = [call[3] for call in layer.calls]
|
|
assert [metadata.causal for metadata in segment_metadata] == [True, False, True]
|
|
assert [metadata.seq_lens.tolist() for metadata in segment_metadata] == [[5, 5], [8, 8], [9, 9]]
|
|
assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [
|
|
[0, 2, 4],
|
|
[0, 3, 6],
|
|
[0, 1, 2],
|
|
]
|
|
assert [metadata.positions.tolist() for metadata in segment_metadata] == [
|
|
[3, 4, 3, 4],
|
|
[5, 6, 7, 5, 6, 7],
|
|
[8, 8],
|
|
]
|
|
assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [
|
|
[[3, 4, 3, 4]],
|
|
[[5, 6, 7, 5, 6, 7]],
|
|
[[8, 8]],
|
|
]
|
|
|
|
|
|
def test_homogeneous_piecewise_packs_strided_query_only_per_segment(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=6,
|
|
seq_len=9,
|
|
kv_start_pos=3,
|
|
),
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-1",
|
|
sequence_id=1,
|
|
query_len=6,
|
|
seq_len=9,
|
|
kv_start_pos=3,
|
|
),
|
|
]
|
|
)
|
|
packed_qkv = torch.randn(2, 6, 3, 2, 4)
|
|
query, key, value = packed_qkv.unbind(dim=2)
|
|
assert not query.is_contiguous()
|
|
metadata = SimpleNamespace(
|
|
attn_mask=None,
|
|
full_attn_spans=[[(5, 8)], [(5, 8)]],
|
|
query_ranges=None,
|
|
extra={},
|
|
)
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context(
|
|
"layer-0",
|
|
query,
|
|
key,
|
|
value,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
|
|
assert not context.query.is_contiguous()
|
|
assert context.query.data_ptr() == query.data_ptr()
|
|
|
|
|
|
def test_piecewise_metadata_snapshots_reused_native_scheduler_buffer(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
shared_scheduler_metadata = torch.zeros(1, dtype=torch.int32)
|
|
build_id = 0
|
|
|
|
def build_metadata(**kwargs):
|
|
nonlocal build_id
|
|
shared_scheduler_metadata.fill_(build_id)
|
|
metadata = SimpleNamespace(
|
|
build_id=build_id,
|
|
causal=kwargs["causal"],
|
|
seq_lens=kwargs["seq_lens"].clone(),
|
|
query_start_loc_cpu=kwargs["query_start_loc_cpu"].clone(),
|
|
positions=kwargs["positions"].clone(),
|
|
slot_mappings=kwargs["slot_mappings"].clone(),
|
|
scheduler_metadata=shared_scheduler_metadata,
|
|
prefix_scheduler_metadata=None,
|
|
)
|
|
build_id += 1
|
|
return {"layer-0": metadata}
|
|
|
|
monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
|
|
)
|
|
qkv = torch.randn(1, 4, 2, 4)
|
|
metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context(
|
|
"layer-0",
|
|
qkv,
|
|
qkv,
|
|
qkv,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
|
|
segment_schedules = [item.scheduler_metadata for item in context.piecewise_native_metadata]
|
|
assert [item.item() for item in segment_schedules] == [1, 2, 3]
|
|
assert len({item.data_ptr() for item in segment_schedules}) == 3
|
|
|
|
|
|
def test_omni_paged_backend_treats_empty_full_spans_as_causal(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, layer, events = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
|
|
)
|
|
qkv = torch.randn(1, 4, 2, 4)
|
|
metadata = SimpleNamespace(
|
|
attn_mask=torch.ones(1, 1, 4, 4, dtype=torch.bool).tril(),
|
|
full_attn_spans=[[]],
|
|
query_ranges=None,
|
|
extra={},
|
|
)
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context(
|
|
"layer-0",
|
|
qkv,
|
|
qkv,
|
|
qkv,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
output = _run_omni_paged_backend(context)
|
|
|
|
assert torch.equal(output, qkv)
|
|
assert layer.native_events == ["update", "forward"]
|
|
assert len(events) == 2
|
|
assert events[1][1]["causal"] is True
|
|
assert events[1][1]["seq_lens"].tolist() == [4]
|
|
|
|
|
|
def test_omni_paged_backend_runs_unequal_hunyuan_prefixes_without_padding(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, _, layer, events = _make_adapter(monkeypatch, capacity=64)
|
|
batch = adapter.prepare_batch(
|
|
[
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=17,
|
|
seq_len=50,
|
|
kv_start_pos=33,
|
|
),
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-1",
|
|
sequence_id=1,
|
|
query_len=17,
|
|
seq_len=54,
|
|
kv_start_pos=37,
|
|
),
|
|
]
|
|
)
|
|
query = torch.arange(2 * 17 * 2 * 4, dtype=torch.float32).reshape(2, 17, 2, 4)
|
|
key = query + 100
|
|
value = query + 200
|
|
metadata = SimpleNamespace(
|
|
attn_mask=torch.ones(2, 1, 17, 54, dtype=torch.bool),
|
|
full_attn_spans=[[(33, 50)], [(37, 54)]],
|
|
query_ranges=None,
|
|
extra={},
|
|
)
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context(
|
|
"layer-0",
|
|
query,
|
|
key,
|
|
value,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
output = _run_omni_paged_backend(context)
|
|
|
|
assert torch.equal(output, query)
|
|
assert batch.query_start_loc.tolist() == [0, 17, 34]
|
|
assert batch.seq_lens.tolist() == [50, 54]
|
|
assert batch.positions.tolist() == [*range(33, 50), *range(37, 54)]
|
|
assert torch.equal(context.key_write, key.reshape(34, 2, 4))
|
|
assert torch.equal(context.value_write, value.reshape(34, 2, 4))
|
|
assert layer.native_events == ["update", "forward"]
|
|
assert layer.updates[0][2].tolist() == batch.positions.tolist()
|
|
assert layer.calls[0][3].seq_lens.tolist() == [50, 54]
|
|
assert layer.calls[0][3].query_start_loc_cpu.tolist() == [0, 17, 34]
|
|
assert len(events) == 2
|
|
|
|
|
|
def test_forward_rejects_unaligned_heterogeneous_piecewise_spans_before_cache_write(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[
|
|
DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4),
|
|
DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=4, seq_len=4),
|
|
]
|
|
)
|
|
qkv = torch.randn(2, 4, 2, 4)
|
|
metadata = SimpleNamespace(
|
|
attn_mask=torch.ones(2, 1, 4, 4, dtype=torch.bool),
|
|
full_attn_spans=[[(1, 3)], [(2, 4)]],
|
|
query_ranges=None,
|
|
extra={},
|
|
)
|
|
|
|
with adapter.activate(batch), pytest.raises(ValueError, match="produce aligned segments"):
|
|
adapter.prepare_layer_context(
|
|
"layer-0",
|
|
qkv,
|
|
qkv,
|
|
qkv,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
|
|
assert layer.native_events == []
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"spans",
|
|
[
|
|
[(1, 3), (2, 4)],
|
|
[(1, 5)],
|
|
],
|
|
)
|
|
def test_forward_rejects_invalid_piecewise_spans_before_cache_write(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
spans: list[tuple[int, int]],
|
|
) -> None:
|
|
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
|
|
)
|
|
qkv = torch.randn(1, 4, 2, 4)
|
|
metadata = SimpleNamespace(full_attn_spans=[spans], extra={})
|
|
|
|
with adapter.activate(batch), pytest.raises(ValueError, match="sorted, non-overlapping, and within the sequence"):
|
|
adapter.prepare_layer_context(
|
|
"layer-0",
|
|
qkv,
|
|
qkv,
|
|
qkv,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
|
|
assert layer.native_events == []
|
|
|
|
|
|
def test_forward_rejects_non_4d_mask_with_piecewise_spans_before_cache_write(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
|
|
)
|
|
qkv = torch.randn(1, 4, 2, 4)
|
|
metadata = SimpleNamespace(
|
|
attn_mask=torch.ones(1, 4, 4, dtype=torch.bool),
|
|
full_attn_spans=[[(1, 3)]],
|
|
extra={},
|
|
)
|
|
|
|
with adapter.activate(batch), pytest.raises(NotImplementedError, match="requires a 4D tensor"):
|
|
adapter.prepare_layer_context(
|
|
"layer-0",
|
|
qkv,
|
|
qkv,
|
|
qkv,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
|
|
assert layer.native_events == []
|
|
|
|
|
|
def test_forward_rejects_piecewise_query_ranges_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
|
|
)
|
|
qkv = torch.randn(1, 4, 2, 4)
|
|
metadata = SimpleNamespace(
|
|
attn_mask=torch.ones(1, 1, 4, 4, dtype=torch.bool),
|
|
full_attn_spans=[[(1, 3)]],
|
|
query_ranges=(object(),),
|
|
extra={},
|
|
)
|
|
|
|
with adapter.activate(batch), pytest.raises(NotImplementedError, match="query_ranges"):
|
|
adapter.prepare_layer_context(
|
|
"layer-0",
|
|
qkv,
|
|
qkv,
|
|
qkv,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
|
|
assert layer.native_events == []
|
|
|
|
|
|
def test_forward_rejects_piecewise_non_suffix_query_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=2,
|
|
seq_len=6,
|
|
kv_start_pos=2,
|
|
)
|
|
]
|
|
)
|
|
qkv = torch.randn(1, 2, 2, 4)
|
|
metadata = SimpleNamespace(full_attn_spans=[[(2, 4)]], extra={})
|
|
|
|
with adapter.activate(batch), pytest.raises(ValueError, match="requires each query/write span to end"):
|
|
adapter.prepare_layer_context(
|
|
"layer-0",
|
|
qkv,
|
|
qkv,
|
|
qkv,
|
|
omni_attn_metadata=metadata,
|
|
)
|
|
|
|
assert layer.native_events == []
|
|
|
|
|
|
def test_forward_accepts_non_aligned_current_write(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=2,
|
|
seq_len=5,
|
|
kv_start_pos=3,
|
|
)
|
|
]
|
|
)
|
|
query = torch.randn(1, 2, 2, 4)
|
|
key = torch.arange(1 * 2 * 2 * 4, dtype=torch.float32).reshape(1, 2, 2, 4)
|
|
value = key + 100
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context("layer-0", query, key, value)
|
|
|
|
assert context.query.shape == (2, 2, 4)
|
|
assert torch.equal(context.key_write, key.reshape(2, 2, 4))
|
|
assert torch.equal(context.value_write, value.reshape(2, 2, 4))
|
|
assert context.slot_mapping.tolist() == [3, 4]
|
|
assert layer.native_events == []
|
|
|
|
|
|
def test_forward_rejects_untranslated_omni_metadata_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
|
|
)
|
|
qkv = torch.randn(1, 2, 4)
|
|
metadata = SimpleNamespace(attn_mask=torch.ones(1, 1), extra={})
|
|
|
|
with adapter.activate(batch), pytest.raises(NotImplementedError, match="cannot translate"):
|
|
adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata)
|
|
|
|
assert layer.updates == []
|
|
|
|
|
|
def test_forward_accepts_flattened_hidden_dimension(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)]
|
|
)
|
|
query = torch.randn(3, 8)
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context("layer-0", query, query, query)
|
|
|
|
assert context.query.shape == (3, 2, 4)
|
|
assert context.restore_output(context.query).shape == query.shape
|
|
|
|
|
|
def test_forward_accepts_one_flattened_token_for_single_head(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch, num_heads=1, num_kv_heads=1)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
|
|
)
|
|
qkv = torch.randn(1, 4)
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
|
output = _run_omni_paged_backend(context)
|
|
|
|
assert context.query.shape == (1, 1, 4)
|
|
assert context.query_token_shape == (1,)
|
|
assert not context.query_has_head_dims
|
|
assert output.shape == qkv.shape
|
|
|
|
|
|
def test_forward_requires_active_prepared_batch(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
qkv = torch.randn(1, 2, 4)
|
|
|
|
with pytest.raises(RuntimeError, match="adapter.activate"):
|
|
adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
|
|
|
|
|
def test_new_preparation_invalidates_native_buffer_views_in_older_batch(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
old_batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
|
|
)
|
|
adapter.prepare_batch([DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1)])
|
|
|
|
with pytest.raises(ValueError, match="stale"):
|
|
with adapter.activate(old_batch):
|
|
pass
|
|
|
|
|
|
def test_block_table_change_invalidates_prepared_native_buffer_views(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
|
|
)
|
|
|
|
adapter.invalidate_prepared_batches()
|
|
|
|
with pytest.raises(ValueError, match="stale"):
|
|
with adapter.activate(batch):
|
|
pass
|
|
|
|
|
|
def test_block_table_change_is_rejected_during_active_forward(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
|
|
)
|
|
|
|
with adapter.activate(batch), pytest.raises(RuntimeError, match="during an active forward"):
|
|
adapter.invalidate_prepared_batches()
|
|
|
|
|
|
def test_prepare_batch_rejects_active_native_buffers(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
|
|
)
|
|
|
|
with adapter.activate(batch), pytest.raises(RuntimeError, match="during an active forward"):
|
|
adapter.prepare_batch([DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1)])
|
|
|
|
|
|
def test_forward_rejects_tensor_count_different_from_prepared_batch(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)]
|
|
)
|
|
qkv = torch.randn(2, 2, 4)
|
|
|
|
with adapter.activate(batch), pytest.raises(ValueError, match="token count"):
|
|
adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
|
|
|
|
|
def test_forward_rejects_dtype_different_from_model_activation_dtype(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
|
|
)
|
|
qkv = torch.randn(1, 2, 4, dtype=torch.float64)
|
|
|
|
with adapter.activate(batch), pytest.raises(ValueError, match="model activation dtype"):
|
|
adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
|
|
|
|
|
def test_forward_rejects_invalid_key_token_count(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
|
|
)
|
|
query = torch.randn(4, 2, 4)
|
|
key = torch.randn(3, 2, 4)
|
|
|
|
with adapter.activate(batch), pytest.raises(ValueError, match="key token count"):
|
|
adapter.prepare_layer_context("layer-0", query, key, query)
|
|
|
|
|
|
def test_forward_rejects_batched_layout_that_does_not_match_rows(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[
|
|
DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1),
|
|
DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=3, seq_len=3),
|
|
]
|
|
)
|
|
qkv = torch.randn(2, 2, 2, 4)
|
|
|
|
with adapter.activate(batch), pytest.raises(ValueError, match="must match prepared rows"):
|
|
adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
|
|
|
|
|
def test_forward_packs_and_restores_right_padded_heterogeneous_rows(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[
|
|
DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1),
|
|
DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=3, seq_len=3),
|
|
]
|
|
)
|
|
qkv = torch.arange(2 * 3 * 2 * 4, dtype=torch.float32).reshape(2, 3, 2, 4)
|
|
|
|
with adapter.activate(batch):
|
|
context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
|
|
|
assert context.output_scatter_indices.tolist() == [0, 3, 4, 5]
|
|
torch.testing.assert_close(context.query, qkv.reshape(-1, 2, 4)[[0, 3, 4, 5]])
|
|
restored = context.restore_output(context.query)
|
|
assert restored.shape == qkv.shape
|
|
torch.testing.assert_close(restored[0, 0], qkv[0, 0])
|
|
assert torch.count_nonzero(restored[0, 1:]) == 0
|
|
torch.testing.assert_close(restored[1], qkv[1])
|
|
|
|
|
|
def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
batch = adapter.prepare_batch(
|
|
[
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=1,
|
|
seq_len=3,
|
|
kv_start_pos=2,
|
|
),
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-1",
|
|
sequence_id=1,
|
|
query_len=1,
|
|
seq_len=5,
|
|
kv_start_pos=4,
|
|
),
|
|
]
|
|
)
|
|
query = torch.randn(2, 1, 2, 4)
|
|
full_kv = torch.randn(2, 4, 2, 4)
|
|
|
|
with adapter.activate(batch), pytest.raises(ValueError, match="batched key layout must match prepared rows"):
|
|
adapter.prepare_layer_context("layer-0", query, full_kv, full_kv)
|
|
|
|
|
|
def test_runner_runtime_switches_from_prefill_to_reusable_denoise_metadata(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch)
|
|
metadata = DiffusionPagedAttentionMetadata(
|
|
prefill_rows=(DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=6, seq_len=6),),
|
|
denoise_rows=(
|
|
DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=2,
|
|
seq_len=5,
|
|
kv_start_pos=3,
|
|
),
|
|
),
|
|
)
|
|
runtime = DiffusionPagedAttentionRuntime(adapter, metadata)
|
|
|
|
with set_forward_context(paged_kv_runtime=runtime), runtime.activate():
|
|
set_forward_context_denoise_step_idx(0)
|
|
assert Attention._active_paged_kv_adapter() is adapter
|
|
assert adapter._active_batch.rows == metadata.prefill_rows
|
|
|
|
set_forward_context_denoise_step_idx(1)
|
|
assert Attention._active_paged_kv_adapter() is adapter
|
|
denoise_batch = adapter._active_batch
|
|
assert denoise_batch.rows == metadata.denoise_rows
|
|
|
|
set_forward_context_denoise_step_idx(2)
|
|
assert Attention._active_paged_kv_adapter() is adapter
|
|
assert adapter._active_batch is denoise_batch
|
|
|
|
assert adapter._active_batch is None
|
|
|
|
|
|
def test_causal_batch_rejects_non_suffix_write(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, _, _, _ = _make_adapter(monkeypatch, non_causal=False)
|
|
row = DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=2,
|
|
seq_len=8,
|
|
kv_start_pos=3,
|
|
)
|
|
|
|
batch = adapter.prepare_batch([row])
|
|
qkv = torch.randn(2, 2, 4)
|
|
|
|
with adapter.activate(batch), pytest.raises(ValueError, match="requires each query/write span to end"):
|
|
adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
|
|
|
|
|
def test_prepare_batch_rejects_native_decode_after_prefill_order(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
adapter, block_tables, _, _ = _make_adapter(
|
|
monkeypatch,
|
|
non_causal=False,
|
|
reorder_batch_threshold=1,
|
|
)
|
|
rows = [
|
|
DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=2, seq_len=2),
|
|
DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1),
|
|
]
|
|
|
|
with pytest.raises(ValueError, match="decode/short-query rows before"):
|
|
adapter.prepare_batch(rows)
|
|
|
|
assert block_tables.gather_calls == []
|
|
assert block_tables.slot_calls == []
|
|
|
|
|
|
def test_prepare_batch_rejects_seq_len_beyond_installed_row_capacity(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, block_tables, _, _ = _make_adapter(monkeypatch)
|
|
block_tables.num_blocks.np[0, 2] = 1
|
|
row = DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=5,
|
|
seq_len=5,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="installed Worker row contains only 1 blocks"):
|
|
adapter.prepare_batch([row])
|
|
|
|
assert block_tables.gather_calls == []
|
|
assert block_tables.slot_calls == []
|
|
|
|
|
|
def test_prepare_batch_rejects_seq_len_beyond_installed_logical_length(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, block_tables, _, _ = _make_adapter(monkeypatch)
|
|
adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding(2, 4)
|
|
row = DiffusionPagedAttentionRow(
|
|
request_id="req-0",
|
|
sequence_id=0,
|
|
query_len=5,
|
|
seq_len=5,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="installed allocation has logical length 4"):
|
|
adapter.prepare_batch([row])
|
|
|
|
assert block_tables.gather_calls == []
|
|
assert block_tables.slot_calls == []
|
|
|
|
|
|
def test_prepare_batch_passes_dcp_local_seq_lens_to_native_builder(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
adapter, block_tables, _, events = _make_adapter(monkeypatch)
|
|
block_tables.cp_size = 2
|
|
block_tables.cp_rank = 1
|
|
block_tables.cp_interleave = 2
|
|
block_tables.num_blocks.np[0, 2] = 1
|
|
|
|
batch = adapter.prepare_batch(
|
|
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
|
|
)
|
|
|
|
assert batch.row_indices.tolist() == [2]
|
|
assert events[0][1]["dcp_local_seq_lens"].tolist() == [2]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("kwargs", "message"),
|
|
[
|
|
({"sequence_id": 0, "context_id": "text"}, "exactly one"),
|
|
({}, "exactly one"),
|
|
({"sequence_id": 0, "query_len": 0}, "query_len"),
|
|
({"sequence_id": 0, "kv_start_pos": 4, "query_len": 2, "seq_len": 5}, "exceeds seq_len"),
|
|
],
|
|
)
|
|
def test_row_contract_rejects_invalid_identity_or_span(kwargs: dict, message: str) -> None:
|
|
values = {
|
|
"request_id": "req-0",
|
|
"query_len": 1,
|
|
"seq_len": 1,
|
|
**kwargs,
|
|
}
|
|
with pytest.raises(ValueError, match=message):
|
|
DiffusionPagedAttentionRow(**values)
|
|
|
|
|
|
@pytest.mark.parametrize("pcp_size", [1, 2])
|
|
@pytest.mark.parametrize("selection_fails", [False, True])
|
|
def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
pcp_size: int,
|
|
selection_fails: bool,
|
|
) -> None:
|
|
from vllm.v1.attention import selector
|
|
|
|
selected_backends = []
|
|
specialized_backends = []
|
|
impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock()))
|
|
native_backend = SimpleNamespace(
|
|
get_name=lambda: "ASCEND",
|
|
indexes_kv_by_block_stride=lambda: True,
|
|
get_impl_cls=lambda: impl_cls,
|
|
forward_includes_kv_cache_update=True,
|
|
)
|
|
|
|
original_backend_per_kind = {"full": object()}
|
|
config = SimpleNamespace(
|
|
attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind, use_non_causal=True),
|
|
parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size, decode_context_parallel_size=1),
|
|
model_config=SimpleNamespace(dtype=torch.float16, rswa_window=None),
|
|
cache_config=SimpleNamespace(cache_dtype="auto", user_specified_block_size=False),
|
|
kv_transfer_config=None,
|
|
speculative_config=None,
|
|
)
|
|
|
|
def select_backend(*, backend, attn_selector_config, num_heads, _run_kv_cache_dtype):
|
|
# Exercise the real upstream selector, replacing only platform resolution.
|
|
# Ulysses has already gathered tokens and sharded heads before the
|
|
# native kernel; the MoE PCP mapping must not request PCP attention.
|
|
assert not attn_selector_config.use_pcp
|
|
assert not attn_selector_config.use_dcp
|
|
assert attn_selector_config.use_non_causal
|
|
assert num_heads == 4
|
|
assert _run_kv_cache_dtype == config.cache_config.cache_dtype
|
|
selected_backends.append((backend, config.attention_config.backend_per_kind))
|
|
if selection_fails:
|
|
raise ValueError("test backend unavailable")
|
|
return native_backend
|
|
|
|
monkeypatch.setattr("vllm.config.get_current_vllm_config", lambda: config)
|
|
monkeypatch.setattr(selector, "_cached_get_attn_backend", select_backend)
|
|
monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
|
|
|
|
def specialize_backend(backend, *, ulysses_degree):
|
|
specialized_backends.append((backend, ulysses_degree))
|
|
return backend
|
|
|
|
monkeypatch.setattr(
|
|
adapter_module.current_omni_platform,
|
|
"get_diffusion_paged_kv_attn_backend",
|
|
specialize_backend,
|
|
)
|
|
layer = SimpleNamespace(num_heads=8, softmax_scale=0.125)
|
|
spec = FullAttentionSpec(
|
|
block_size=16,
|
|
num_kv_heads=4,
|
|
head_size=8,
|
|
dtype=torch.float16,
|
|
non_causal=True,
|
|
)
|
|
|
|
kwargs = dict(
|
|
layer_name="layer-0",
|
|
layer=layer,
|
|
spec=spec,
|
|
vllm_config=config,
|
|
device=torch.device("cpu"),
|
|
ulysses_degree=2,
|
|
)
|
|
if selection_fails:
|
|
with pytest.raises(ValueError, match="test backend unavailable"):
|
|
adapter_module.DiffusionPagedAttentionLayerAdapter(**kwargs)
|
|
else:
|
|
native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter(**kwargs)
|
|
|
|
assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {})]
|
|
assert config.attention_config.backend is None
|
|
assert config.attention_config.backend_per_kind is original_backend_per_kind
|
|
assert config.parallel_config.prefill_context_parallel_size == pcp_size
|
|
if selection_fails:
|
|
assert not specialized_backends
|
|
impl_cls.assert_not_called()
|
|
return
|
|
assert specialized_backends == [(native_backend, 2)]
|
|
assert native_layer.num_heads == 4
|
|
assert native_layer.num_kv_heads == 2
|
|
assert native_layer.spec.num_kv_heads == 2
|
|
# 0.29: the flag no longer rides on the spec, it selects the layout.
|
|
assert layout_for_backend(native_layer.attn_backend).is_block_outermost is True
|
|
assert native_layer._q_scale_float == 1.0
|
|
assert native_layer._k_scale_float == 1.0
|
|
assert native_layer._v_scale_float == 1.0
|
|
impl_kwargs = impl_cls.call_args.kwargs
|
|
assert impl_kwargs["num_heads"] == 4
|
|
assert impl_kwargs["num_kv_heads"] == 2
|
|
|
|
|
|
def test_omni_attention_wraps_paged_kernel_with_sp_hooks() -> None:
|
|
events: list[str] = []
|
|
|
|
class Strategy:
|
|
name = "ulysses"
|
|
|
|
def pre_attention(self, query, key, value, metadata):
|
|
events.append("pre")
|
|
return query + 1, key + 2, value + 3, metadata, object()
|
|
|
|
def post_attention(self, output, _ctx):
|
|
events.append("post")
|
|
return output + 4
|
|
|
|
class Adapter:
|
|
def prepare_layer_context(self, layer_name, query, key, value, *, omni_attn_metadata):
|
|
events.append("prepare")
|
|
assert layer_name == "layer-0"
|
|
assert omni_attn_metadata is None
|
|
assert torch.equal(key, original + 2)
|
|
assert torch.equal(value, original + 3)
|
|
return SimpleNamespace(query=query)
|
|
|
|
class Backend:
|
|
def forward_paged(self, context):
|
|
events.append("kernel")
|
|
return context.query * 2
|
|
|
|
layer = Attention.__new__(Attention)
|
|
nn.Module.__init__(layer)
|
|
layer.prefix = "layer-0"
|
|
layer.paged_kv_cache_role = "primary"
|
|
layer.attn_backend = SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN")
|
|
layer.attention = Backend()
|
|
layer.use_ring = False
|
|
layer._no_parallel_strategy = object()
|
|
layer._get_active_parallel_strategy = lambda: Strategy()
|
|
layer._with_kv_cache_dtype = lambda metadata: metadata
|
|
original = torch.zeros(1, 2, 2, 4)
|
|
|
|
with set_forward_context(), override_paged_kv_adapter(Adapter()):
|
|
assert layer.is_paged_kv_active()
|
|
output = layer._forward_impl(original, original, original)
|
|
|
|
assert events == ["pre", "prepare", "kernel", "post"]
|
|
assert torch.equal(output, torch.full_like(original, 6))
|
|
|
|
|
|
def test_omni_attention_strips_paged_ulysses_padding_around_kernel() -> None:
|
|
events: list[tuple[str, tuple[int, ...]]] = []
|
|
|
|
class Strategy:
|
|
name = "ulysses"
|
|
|
|
def pre_attention(self, query, key, value, metadata):
|
|
return query, key, value, metadata, SimpleNamespace(joint_len=0, joint_strategy="front")
|
|
|
|
def post_attention(self, output, _ctx):
|
|
events.append(("post", tuple(output.shape)))
|
|
return output
|
|
|
|
class Adapter:
|
|
def prepare_layer_context(self, _layer_name, query, key, value, *, omni_attn_metadata):
|
|
assert omni_attn_metadata is None
|
|
events.append(("kernel_input", tuple(query.shape)))
|
|
assert query[:, -1].tolist() == [[[3.0]]]
|
|
assert torch.equal(query, key)
|
|
assert torch.equal(query, value)
|
|
return SimpleNamespace(query=query)
|
|
|
|
class Backend:
|
|
def forward_paged(self, context):
|
|
return context.query
|
|
|
|
layer = Attention.__new__(Attention)
|
|
nn.Module.__init__(layer)
|
|
layer.prefix = "layer-0"
|
|
layer.paged_kv_cache_role = "primary"
|
|
layer.attn_backend = SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN")
|
|
layer.attention = Backend()
|
|
layer.use_ring = False
|
|
layer._no_parallel_strategy = object()
|
|
layer._get_active_parallel_strategy = lambda: Strategy()
|
|
tensor = torch.arange(5, dtype=torch.float32).reshape(1, 5, 1, 1)
|
|
|
|
with set_forward_context(), override_paged_kv_adapter(Adapter()):
|
|
get_forward_context().sp_padding_size = 1
|
|
output = layer._forward_impl(tensor, tensor, tensor)
|
|
|
|
assert events == [("kernel_input", (1, 4, 1, 1)), ("post", (1, 5, 1, 1))]
|
|
assert output.flatten().tolist() == [0.0, 1.0, 2.0, 3.0, 0.0]
|
|
|
|
|
|
def test_omni_attention_rejects_paged_request_for_backend_without_paged_support() -> None:
|
|
class Adapter:
|
|
def prepare_layer_context(self, *args, **kwargs):
|
|
raise AssertionError("an unsupported backend must fail before adapter preparation")
|
|
|
|
layer = Attention.__new__(Attention)
|
|
nn.Module.__init__(layer)
|
|
layer.prefix = "layer-0"
|
|
layer.paged_kv_cache_role = "primary"
|
|
layer.attn_backend = SimpleNamespace(supports_paged_kv=False, get_name=lambda: "SDPA")
|
|
layer.use_ring = False
|
|
layer._no_parallel_strategy = object()
|
|
layer._get_active_parallel_strategy = lambda: layer._no_parallel_strategy
|
|
|
|
with (
|
|
set_forward_context(),
|
|
override_paged_kv_adapter(Adapter()),
|
|
pytest.raises(NotImplementedError, match="backend with paged support"),
|
|
):
|
|
layer._forward_impl(torch.zeros(1, 2, 2, 4), torch.zeros(1, 2, 2, 4), torch.zeros(1, 2, 2, 4))
|
|
|
|
|
|
def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None:
|
|
events: list[str] = []
|
|
|
|
class Strategy:
|
|
name = "ulysses"
|
|
|
|
def pre_attention(self, query, key, value, metadata):
|
|
events.append("pre")
|
|
return query, key, value, metadata, object()
|
|
|
|
def post_attention(self, output, _ctx):
|
|
events.append("post")
|
|
return output
|
|
|
|
layer = Attention.__new__(Attention)
|
|
nn.Module.__init__(layer)
|
|
layer.paged_kv_cache_role = "primary"
|
|
layer.use_ring = False
|
|
layer._no_parallel_strategy = object()
|
|
layer._get_active_parallel_strategy = lambda: Strategy()
|
|
layer._with_kv_cache_dtype = lambda metadata: metadata
|
|
|
|
def run_local_attention(query, _key, _value, _metadata):
|
|
events.append("dense")
|
|
return query
|
|
|
|
layer._run_local_attention = run_local_attention
|
|
qkv = torch.zeros(1, 2, 2, 4)
|
|
|
|
assert not layer.is_paged_kv_active()
|
|
output = layer._forward_impl(qkv, qkv, qkv)
|
|
|
|
assert events == ["pre", "dense", "post"]
|
|
assert output is qkv
|
|
|
|
|
|
def test_omni_attention_rejects_scheduler_paged_forward_without_active_adapter() -> None:
|
|
layer = Attention.__new__(Attention)
|
|
nn.Module.__init__(layer)
|
|
layer._scheduler_paged_kv = True
|
|
layer.paged_kv_cache_role = "primary"
|
|
layer._get_active_parallel_strategy = lambda: object()
|
|
qkv = torch.zeros(1, 2, 2, 4)
|
|
|
|
with set_forward_context(), pytest.raises(RuntimeError, match="without an active Worker adapter"):
|
|
layer._forward_impl(qkv, qkv, qkv)
|