[Core][Diffusion] Allocate request-scoped KV contexts for SenseNova-U1.5 B1

.buildkite/cuda/test-ready.yml

@@ -43,7 +43,7 @@ steps:
         commands:
           - uv pip install --system --no-deps 'fastvideo-kernel==0.3.4'
           - python3 -c 'from fastvideo_kernel.block_sparse_attn import block_sparse_attn_triton'
-          - timeout 15m pytest -s -v tests/diffusion/attention/ tests/diffusion/utils/test_chunked_transfer.py -m 'core_model and cuda and L4 and cards_1' --run-level "core_model"
+          - timeout 15m pytest -s -v tests/diffusion/attention/ tests/diffusion/diffusion_kv/ tests/diffusion/utils/test_chunked_transfer.py -m 'core_model and cuda and L4 and cards_1' --run-level "core_model"
         mirror_hardwares: l4_1
 
       - label: "Diffusion · Offloader Test"

docs/design/module/diffusion/continuous_batching.md

@@ -9,10 +9,12 @@ primary_code_paths:
   - vllm_omni/diffusion/sched/**
 related_code_paths:
   - vllm_omni/diffusion/executor/**
+  - vllm_omni/diffusion/diffusion_kv/**
 depends_on:
   - diffusion_runtime.md
 validation_paths:
   - tests/diffusion/batching/**
+  - tests/diffusion/diffusion_kv/**
 upstream_refs: []
 last_reviewed: 2026-07-16
 ---

tests/diffusion/diffusion_kv/test_block_tables.py

@@ -184,14 +184,32 @@ def copy_to_uva(self) -> None:
 class _FakeBlockTables:
     instances: list[_FakeBlockTables] = []
 
-    def __init__(self, **kwargs: object) -> None:
-        self.kwargs = kwargs
-        self.block_sizes = list(kwargs["block_sizes"])
-        self.max_num_reqs = int(kwargs["max_num_reqs"])
-        self.max_num_batched_tokens = int(kwargs["max_num_batched_tokens"])
-        self.cp_size = int(kwargs["cp_size"])
-        self.cp_rank = int(kwargs["cp_rank"])
-        self.cp_interleave = int(kwargs["cp_interleave"])
+    def __init__(
+        self,
+        *,
+        block_sizes: list[int],
+        max_num_reqs: int,
+        max_num_batched_tokens: int,
+        cp_size: int,
+        cp_rank: int,
+        cp_interleave: int,
+        **kwargs: object,
+    ) -> None:
+        self.kwargs = dict(
+            kwargs,
+            block_sizes=block_sizes,
+            max_num_reqs=max_num_reqs,
+            max_num_batched_tokens=max_num_batched_tokens,
+            cp_size=cp_size,
+            cp_rank=cp_rank,
+            cp_interleave=cp_interleave,
+        )
+        self.block_sizes = list(block_sizes)
+        self.max_num_reqs = max_num_reqs
+        self.max_num_batched_tokens = max_num_batched_tokens
+        self.cp_size = cp_size
+        self.cp_rank = cp_rank
+        self.cp_interleave = cp_interleave
         self.blocks_per_kv_block = [1] * len(self.block_sizes)
         self.block_tables = [_FakeStagedTable() for _ in self.block_sizes]
         self.num_blocks = _FakeNumBlocks(len(self.block_sizes), self.max_num_reqs)
@@ -345,8 +363,12 @@ def _runner(
         device="cuda:1",
     )
     backend._kv_cache_layer_adapters = {
-        "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None),
-        "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None),
+        "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None, cache_role="primary"),
+        "layer-1": SimpleNamespace(
+            spec=SimpleNamespace(non_causal=False),
+            kv_cache=None,
+            cache_role="cross_attention",
+        ),
     }
 
     config = SimpleNamespace(
@@ -420,8 +442,16 @@ def test_initialize_failure_does_not_publish_partial_cache_state(
         device="cuda:1",
     )
     backend._kv_cache_layer_adapters = {
-        "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-0"),
-        "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-1"),
+        "layer-0": SimpleNamespace(
+            spec=SimpleNamespace(non_causal=False),
+            kv_cache="placeholder-0",
+            cache_role="primary",
+        ),
+        "layer-1": SimpleNamespace(
+            spec=SimpleNamespace(non_causal=False),
+            kv_cache="placeholder-1",
+            cache_role="cross_attention",
+        ),
     }
     config = SimpleNamespace(
         num_blocks=32,
@@ -460,7 +490,11 @@ def test_initialize_rejects_config_for_different_registered_layers() -> None:
         device="cuda:1",
     )
     backend._kv_cache_layer_adapters = {
-        "expected-layer": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None)
+        "expected-layer": SimpleNamespace(
+            spec=SimpleNamespace(non_causal=False),
+            kv_cache=None,
+            cache_role="primary",
+        )
     }
     config = SimpleNamespace(
         kv_cache_groups=[SimpleNamespace(layer_names=["other-layer"])],
@@ -494,6 +528,19 @@ def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest
         5,
         ((6, 7), (8,)),
     )
+    assert sequence_binding.cache_role == "primary"
+    assert context_binding.cache_role == "cross_attention"
+
+
+def test_context_install_rejects_unregistered_cache_role(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    unknown = replace(_metadata().contexts[0], cache_role="unknown")
+
+    with pytest.raises(ValueError, match="unregistered cache role 'unknown'"):
+        runner.install_diffusion_kv_metadata(_metadata(contexts=(unknown,)))
+
+    assert block_tables.append_calls == []
+    assert runner._diffusion_kv_identity_to_row == {}
 
 
 def test_block_table_mutations_invalidate_prepared_attention_batches(

tests/diffusion/diffusion_kv/test_initialization.py

@@ -59,8 +59,8 @@ def _od_config(**overrides):
         ),
         gpu_memory_utilization=0.9,
         kv_cache_memory_bytes=None,
+        diffusion_kv_max_rows_per_request=4,
         max_num_seqs=1,
-        diffusion_kv_max_rows_per_request=1,
         max_num_batched_tokens=64,
         num_gpus=1,
     )

tests/diffusion/diffusion_kv/test_manager.py

@@ -57,17 +57,57 @@ def _manager(
     num_blocks: int,
     *,
     max_model_len: int = 64,
+    max_rows_per_request: int = 4,
+    max_num_rows: int = 64,
     enable_prefix_caching: bool = False,
 ) -> DiffusionKVCacheManager:
     return DiffusionKVCacheManager(
         _config(num_blocks),
         max_model_len=max_model_len,
         scheduler_block_size=BLOCK_SIZE,
         hash_block_size=BLOCK_SIZE,
+        max_rows_per_request=max_rows_per_request,
+        max_num_rows=max_num_rows,
         enable_prefix_caching=enable_prefix_caching,
     )
 
 
+@pytest.mark.parametrize("max_num_rows", [0, -1, True, 1.0, None])
+def test_rejects_invalid_worker_row_capacity(max_num_rows) -> None:
+    with pytest.raises(ValueError, match="max_num_rows must be a positive integer"):
+        _manager(32, max_num_rows=max_num_rows)
+
+
+def test_live_reservations_are_not_bounded_by_the_worker_row_pool() -> None:
+    # Worker rows hold only scheduled requests; a prefetched reservation beside
+    # a running one may exceed the pool and is admitted while blocks fit.
+    manager = _manager(32, max_num_rows=4)
+    context = DiffusionKVContext(context_id="ar", cache_role="ar_decode", num_tokens=8)
+
+    def requests(public_id):
+        return tuple(_request(public_id, index, kv_contexts=(context,)) for index in range(2))
+
+    first = manager.reserve_request("first", requests("first"))
+    assert first is not None
+    assert len(first.sequences) == 2
+    assert len(first.contexts) == 1
+    assert manager.reserve_request("second", requests("second")) is not None
+    manager.free_request("first")
+    manager.free_request("second")
+    assert not manager._internal_request_ids
+
+
+def test_request_exceeding_empty_worker_rows_is_a_permanent_admission_error(monkeypatch) -> None:
+    manager = _manager(32, max_num_rows=2)
+    contexts = tuple(
+        DiffusionKVContext(context_id=str(index), cache_role="ar_decode", num_tokens=8) for index in range(2)
+    )
+    monkeypatch.setattr(manager.native_manager, "allocate_slots", lambda *args, **kwargs: pytest.fail("allocated"))
+
+    with pytest.raises(DiffusionKVAdmissionError, match="requires 3 rows; Worker capacity is 2"):
+        manager.reserve_request("too-large", (_request("too-large", 0, kv_contexts=contexts),))
+
+
 def test_successful_request_publishes_prefix_for_a_warm_hit() -> None:
     manager = _manager(8, enable_prefix_caching=True)
     cold = _request("cold", 0, cache_token_ids=range(4))
@@ -299,11 +339,196 @@ def fail_second_allocation(*args, **kwargs):
     assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
 
 
-def test_rejects_independent_context_until_role_routing_is_implemented() -> None:
+def test_reserves_request_scoped_context_and_releases_it_with_the_request() -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
+    requests = (_request("public", 0, kv_contexts=(context,)),)
+
+    metadata = manager.reserve_request("public", requests)
+
+    assert metadata is not None
+    assert metadata.sequences[0].context_ids == ("text",)
+    assert len(metadata.contexts) == 1
+    assert metadata.contexts[0].context_id == "text"
+    assert metadata.contexts[0].cache_role == "cross.text"
+    assert metadata.contexts[0].num_tokens == 8
+    assert len(metadata.contexts[0].block_ids[0]) == 2
+    assert manager._requests["public"] == requests
+    assert [request.request_id for request in manager._context_requests["public"]] == [
+        "public/diffusion-kv/context/text"
+    ]
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before - 4
+
+    manager.free_request("public")
+    assert "public" not in manager._context_requests
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_context_with_block_hashes_never_reads_the_prefix_cache(monkeypatch) -> None:
+    manager = _manager(8, enable_prefix_caching=True)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    context = DiffusionKVContext(context_id="ar", cache_role="ar_decode", num_tokens=8, block_hashes=(b"h0", b"h1"))
+    looked_up = []
+    real_lookup = manager.native_manager.get_computed_blocks
+
+    def recording_lookup(request):
+        looked_up.append(request.request_id)
+        return real_lookup(request)
+
+    monkeypatch.setattr(manager.native_manager, "get_computed_blocks", recording_lookup)
+
+    metadata = manager.reserve_request("public", (_request("public", 0, kv_contexts=(context,)),))
+
+    assert metadata is not None
+    assert all(shim.skip_reading_prefix_cache for shim in manager._context_requests["public"])
+    assert "public/diffusion-kv/context/ar" not in looked_up
+    manager.free_request("public")
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_shared_context_is_allocated_once_for_multiple_sequences() -> None:
     manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
     context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
 
-    with pytest.raises(DiffusionKVAdmissionError, match="DiffusionKVContext"):
+    metadata = manager.reserve_request(
+        "public",
+        (
+            _request("public", 0, kv_contexts=(context,)),
+            _request("public", 1, kv_contexts=(context,)),
+        ),
+    )
+
+    assert metadata is not None
+    assert [sequence.context_ids for sequence in metadata.sequences] == [("text",), ("text",)]
+    assert [context_metadata.context_id for context_metadata in metadata.contexts] == ["text"]
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before - 6
+
+
+def test_conflicting_shared_context_definitions_are_rejected_before_allocation() -> None:
+    manager = _manager(8)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    text = DiffusionKVContext(context_id="shared", cache_role="cross.text", num_tokens=8)
+    image = DiffusionKVContext(context_id="shared", cache_role="cross.image", num_tokens=8)
+
+    with pytest.raises(ValueError, match="conflicting request-scoped contexts"):
+        manager.reserve_request(
+            "public",
+            (
+                _request("public", 0, kv_contexts=(text,)),
+                _request("public", 1, kv_contexts=(image,)),
+            ),
+        )
+
+    assert manager.has_request("public") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_context_allocation_failure_rolls_back_sequences_and_earlier_contexts(monkeypatch) -> None:
+    manager = _manager(12)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    native_allocate = manager.native_manager.allocate_slots
+    num_calls = 0
+
+    def fail_third_allocation(*args, **kwargs):
+        nonlocal num_calls
+        num_calls += 1
+        if num_calls == 3:
+            raise RuntimeError("injected context allocation failure")
+        return native_allocate(*args, **kwargs)
+
+    monkeypatch.setattr(manager.native_manager, "allocate_slots", fail_third_allocation)
+    contexts = (
+        DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8),
+        DiffusionKVContext(context_id="image", cache_role="cross.image", num_tokens=8),
+    )
+
+    with pytest.raises(RuntimeError, match="injected context allocation failure"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0, kv_contexts=contexts),),
+        )
+
+    assert manager.has_request("public") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+@pytest.mark.parametrize("failure_index", [1, 2, 3, 4])
+def test_allocation_exception_after_native_mutation_rolls_back_and_can_retry(monkeypatch, failure_index) -> None:
+    manager = _manager(24)
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    native_allocate = manager.native_manager.allocate_slots
+    attempted_ids: list[str] = []
+    contexts = (
+        DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8),
+        DiffusionKVContext(context_id="image", cache_role="cross.image", num_tokens=8),
+    )
+    requests = tuple(_request("public", i, kv_contexts=contexts) for i in range(2))
+
+    def fail_after_allocation(request, *args, **kwargs):
+        blocks = native_allocate(request, *args, **kwargs)
+        assert blocks is not None
+        attempted_ids.append(request.request_id)
+        if len(attempted_ids) == failure_index:
+            raise RuntimeError("injected failure after native allocation")
+        return blocks
+
+    monkeypatch.setattr(manager.native_manager, "allocate_slots", fail_after_allocation)
+    with pytest.raises(RuntimeError, match="injected failure after native allocation"):
+        manager.reserve_request("public", requests)
+
+    assert manager.has_request("public") is False
+    assert manager._internal_request_ids == set()
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+    for request_id in attempted_ids:
+        assert all(not group for group in manager.native_manager.get_block_ids(request_id))
+
+    monkeypatch.setattr(manager.native_manager, "allocate_slots", native_allocate)
+    metadata = manager.reserve_request("public", requests)
+    assert metadata is not None
+    assert len(metadata.sequences) == 2
+    assert len(metadata.contexts) == 2
+    manager.free_request("public")
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_context_capacity_pressure_rolls_back_partial_request() -> None:
+    manager = _manager(5)
+    assert manager.reserve_request("running", (_request("running", 0),)) is not None
+    free_before = manager.native_manager.block_pool.get_num_free_blocks()
+    context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
+
+    allocation = manager.reserve_request(
+        "waiting",
+        (_request("waiting", 0, kv_contexts=(context,)),),
+    )
+
+    assert allocation is None
+    assert manager.has_request("waiting") is False
+    assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_rejects_request_whose_contexts_can_never_fit_the_empty_pool() -> None:
+    # Three free blocks hold the two-block sequence but not the two-block
+    # context beside it. Without the context term the bound would admit the
+    # request, allocation would fail on every tick, and it would wait forever.
+    manager = _manager(4)
+    context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
+
+    with pytest.raises(DiffusionKVAdmissionError, match="cannot fit even when the block pool is empty"):
+        manager.reserve_request(
+            "public",
+            (_request("public", 0, kv_contexts=(context,)),),
+        )
+    assert manager.has_request("public") is False
+
+
+def test_rejects_context_longer_than_admission_bound() -> None:
+    manager = _manager(8, max_model_len=8)
+    context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=9)
+
+    with pytest.raises(DiffusionKVAdmissionError, match="context 'text' exceeds max_model_len"):
         manager.reserve_request(
             "public",
             (_request("public", 0, kv_contexts=(context,)),),