[Core] Add realtime AR-Diffusion tick sessions for LingBot World 2.0 (#5491)

Signed-off-by: zhangbin.327 <zhangbin.327@bytedance.com> Signed-off-by: Jack47 <scsvip@gmail.com> Co-authored-by: zhangbin.327 <zhangbin.327@bytedance.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-16 09:11:20 +08:00
1 parent 9798bc3a8d
commit 3875f8ded4
@@ -30,6 +30,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

 BLOCK = 16

+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+

 def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
    return ChunkWindowSpec(
@@ -95,7 +97,6 @@ def test_sliding_replace_snaps_to_chunk_boundary():
 def test_sink_chunks_protected():
    # sink = 1 chunk (16 tokens) is never skipped.
    skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
-    # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
    assert skip == 32


@@ -201,12 +202,149 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
        kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
        session_capacity=1,
        cross_attention_lengths={"text": L},
-        device=torch.device("cpu"),
    )
    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
-    expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
+    page_bytes = kv.spec.page_size_bytes * 2
+    scratch_bytes = kv.scratch_num_blocks * page_bytes
+    expected = (int(avail * 0.5) - cross_bytes - scratch_bytes) // page_bytes
    assert kv.num_blocks == expected
-    assert expected > 1 * (2 + 1) + 2  # above the local-slot floor, so the cross deduction is what's tested
+    assert kv.cross_attention_reserved_bytes == cross_bytes
+    assert (kv.num_blocks_total * page_bytes) + cross_bytes <= int(avail * 0.5)
+    assert expected > 1 * (2 + 1) + 2  # above the local-slot minimum, so the cross deduction is what's tested
+
+
+def _make_tiny_capacity_kv(
+    *,
+    requested_capacity: int,
+    available_bytes: int,
+    gpu_memory_fraction: float = 1.0,
+    model_owned_state_bytes_per_session: int = 0,
+) -> ARDiffusionKVCache:
+    return ARDiffusionKVCache(
+        ARDiffusionKVConfig(
+            enable=True,
+            chunk_size=1,
+            window_chunks=3,
+            sink_chunks=3,
+            gpu_memory_fraction=gpu_memory_fraction,
+        ),
+        num_layers=1,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        block_size=1,
+        max_model_len=128,
+        available_bytes=available_bytes,
+        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+        session_capacity=requested_capacity,
+        cross_attention_lengths={"text": 2},
+        frames_per_block=3,
+        model_owned_state_bytes_per_session=model_owned_state_bytes_per_session,
+        device=torch.device("cpu"),
+    )
+
+
+def test_capacity_two_retains_both_windows_and_allocates_next_block():
+    # LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
+    # 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
+    # reservations: 192 bytes exactly.
+    kv = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
+    assert kv.requested_session_capacity == 2
+    assert kv.session_capacity == 2
+    assert kv.managed_num_blocks == 17
+
+    adapters = [kv.begin_request(session_id) for session_id in ("first", "second")]
+    for adapter in adapters:
+        for _ in range(6):  # three sink + three window blocks
+            kv.allocate_token_slots(adapter, 1)
+            kv.commit_chunk(adapter)
+        assert len(kv.window_block_ids(adapter)) == 6
+
+    # Only one request is in flight, so its three-frame block is counted once.
+    kv.allocate_token_slots(adapters[0], 3)
+
+
+def test_requested_capacity_is_capped_and_cross_reservation_uses_effective_capacity():
+    kv = _make_tiny_capacity_kv(requested_capacity=64, available_bytes=192)
+
+    assert kv.requested_session_capacity == 64
+    assert kv.session_capacity == 2
+    assert kv.cross_attention_bytes_per_session == 16
+    assert kv.cross_attention_reserved_bytes == 32
+    assert kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes == kv.memory_budget_bytes
+
+
+def test_capacity_rejects_budget_that_cannot_fit_one_session():
+    # One session needs 11 managed + 3 scratch blocks and 16 cross bytes:
+    # 14 * 8 + 16 = 128 bytes.
+    with pytest.raises(ValueError, match="cannot fit one session.*available=127 bytes.*required=128 bytes"):
+        _make_tiny_capacity_kv(requested_capacity=1, available_bytes=127)
+
+
+def test_fraction_is_soft_floor_when_one_session_fits_actual_memory():
+    kv = _make_tiny_capacity_kv(
+        requested_capacity=2,
+        available_bytes=1280,
+        gpu_memory_fraction=0.05,
+    )
+
+    assert kv.configured_memory_budget_bytes == 64
+    assert kv.memory_budget_bytes == 128
+    assert kv.session_capacity == 1
+
+
+def test_model_owned_state_reduces_effective_session_capacity():
+    kv = _make_tiny_capacity_kv(
+        requested_capacity=2,
+        available_bytes=200,
+        model_owned_state_bytes_per_session=16,
+    )
+
+    assert kv.session_capacity == 1
+    assert kv.model_owned_state_reserved_bytes == 16
+    assert (
+        kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes + kv.model_owned_state_reserved_bytes
+        <= kv.memory_budget_bytes
+    )
+
+
+def _make_shipped_lingbot_geometry(*, gpu_memory_fraction: float) -> ARDiffusionKVCache:
+    return ARDiffusionKVCache(
+        ARDiffusionKVConfig(
+            enable=True,
+            chunk_size=1560,
+            window_chunks=9,
+            sink_chunks=9,
+            gpu_memory_fraction=gpu_memory_fraction,
+        ),
+        num_layers=40,
+        num_kv_heads=40,
+        head_size=128,
+        dtype=torch.bfloat16,
+        block_size=1560,
+        max_model_len=1 << 20,
+        available_bytes=100 * (1 << 30),
+        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+        session_capacity=2,
+        cross_attention_lengths={"text": 512},
+        frames_per_block=3,
+        model_owned_state_bytes_per_session=7_488_000,
+    )
+
+
+def test_shipped_lingbot_geometry_default_fraction_admits_one_session():
+    kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.1)
+
+    assert kv.configured_memory_budget_bytes == 10 * (1 << 30)
+    assert kv.memory_budget_bytes == 33_653_670_400
+    assert kv.session_capacity == 1
+
+
+def test_shipped_lingbot_geometry_tuned_fraction_admits_two_sessions():
+    kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.6)
+
+    assert kv.memory_budget_bytes == 60 * (1 << 30)
+    assert kv.session_capacity == 2


 def _make_kv(
@@ -221,15 +359,24 @@ def _make_kv(
        if local_branches == 1
        else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
    )
+    page_bytes = 2 * BLOCK * 4 * 64 * torch.float32.itemsize
+    declared_scratch_blocks = (max_scratch_tokens_per_branch + BLOCK - 1) // BLOCK
+    scratch_blocks = local_branches * (num_frame_per_block + declared_scratch_blocks)
+    managed_blocks = local_branches * (window_chunks + num_frame_per_block) + 2
    return ARDiffusionKVCache(
-        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
+        ARDiffusionKVConfig(
+            enable=True,
+            chunk_size=BLOCK,
+            window_chunks=window_chunks,
+            gpu_memory_fraction=1.0,
+        ),
        num_layers=1,
        num_kv_heads=4,
        head_size=64,
        dtype=torch.float32,
        block_size=BLOCK,
        max_model_len=4096,
-        available_bytes=1 << 16,  # tiny -> the floor binds
+        available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
        device=torch.device("cpu"),
        kv_branches=kv_branches,
        session_capacity=1,