chore: import upstream snapshot with attribution

2026-07-13 12:55:37 +08:00
commit 7ce4c8e27e
5900 changed files with 1668062 additions and 0 deletions
@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.v1.worker.gpu.block_table import BlockTables
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="requires CUDA",
+)
+
+
+def test_block_tables_apply_staged_writes_fuses_kv_groups(monkeypatch):
+    device = torch.device("cuda")
+    block_tables = BlockTables(
+        block_sizes=[16, 32, 8],
+        max_num_reqs=4,
+        max_num_batched_tokens=64,
+        max_num_blocks_per_group=[8, 8, 8],
+        device=device,
+        kernel_block_sizes=[16, 16, 8],
+    )
+
+    def fail_if_apply_write_called():
+        pytest.fail("multi-group writes should use the fused apply kernel")
+
+    for block_table in block_tables.block_tables:
+        monkeypatch.setattr(block_table, "apply_write", fail_if_apply_write_called)
+
+    block_tables.append_block_ids(
+        req_index=0,
+        new_block_ids=([1, 2], [10, 11], []),
+        overwrite=True,
+    )
+    block_tables.append_block_ids(
+        req_index=1,
+        new_block_ids=([3], [12], [5, 6]),
+        overwrite=True,
+    )
+    block_tables.apply_staged_writes()
+    torch.accelerator.synchronize()
+
+    assert torch.equal(
+        block_tables.block_tables[0].gpu[0, :2],
+        torch.tensor([1, 2], dtype=torch.int32, device=device),
+    )
+    # Group 1 has blocks_per_kv_block == 2, so each KV block expands to two
+    # kernel block IDs.
+    assert torch.equal(
+        block_tables.block_tables[1].gpu[0, :4],
+        torch.tensor([20, 21, 22, 23], dtype=torch.int32, device=device),
+    )
+    assert torch.equal(
+        block_tables.block_tables[0].gpu[1, :1],
+        torch.tensor([3], dtype=torch.int32, device=device),
+    )
+    assert torch.equal(
+        block_tables.block_tables[1].gpu[1, :2],
+        torch.tensor([24, 25], dtype=torch.int32, device=device),
+    )
+    assert torch.equal(
+        block_tables.block_tables[2].gpu[1, :2],
+        torch.tensor([5, 6], dtype=torch.int32, device=device),
+    )
+    assert block_tables.num_blocks.np[0, 0] == 2
+    assert block_tables.num_blocks.np[1, 0] == 4
+    assert block_tables.num_blocks.np[2, 0] == 0
+    assert block_tables.num_blocks.np[0, 1] == 1
+    assert block_tables.num_blocks.np[1, 1] == 2
+    assert block_tables.num_blocks.np[2, 1] == 2
+    assert torch.equal(
+        block_tables.num_blocks.gpu[:, :2],
+        torch.tensor([[2, 1], [4, 2], [0, 2]], dtype=torch.int32, device=device),
+    )
+
+    for block_table in block_tables.block_tables:
+        assert not block_table._staged_write_indices
+        assert not block_table._staged_write_starts
+        assert not block_table._staged_write_contents
+        assert not block_table._staged_write_cu_lens
+
+    block_tables.append_block_ids(
+        req_index=0,
+        new_block_ids=([7], [13], [8]),
+        overwrite=False,
+    )
+    block_tables.apply_staged_writes()
+    torch.accelerator.synchronize()
+
+    assert torch.equal(
+        block_tables.block_tables[0].gpu[0, :3],
+        torch.tensor([1, 2, 7], dtype=torch.int32, device=device),
+    )
+    assert torch.equal(
+        block_tables.block_tables[1].gpu[0, :6],
+        torch.tensor([20, 21, 22, 23, 26, 27], dtype=torch.int32, device=device),
+    )
+    assert torch.equal(
+        block_tables.block_tables[2].gpu[0, :1],
+        torch.tensor([8], dtype=torch.int32, device=device),
+    )
+    assert block_tables.num_blocks.np[0, 0] == 3
+    assert block_tables.num_blocks.np[1, 0] == 6
+    assert block_tables.num_blocks.np[2, 0] == 1
+
+
+def test_block_tables_apply_staged_writes_single_group():
+    device = torch.device("cuda")
+    block_tables = BlockTables(
+        block_sizes=[16],
+        max_num_reqs=2,
+        max_num_batched_tokens=16,
+        max_num_blocks_per_group=[4],
+        device=device,
+        kernel_block_sizes=[16],
+    )
+
+    block_tables.append_block_ids(
+        req_index=0,
+        new_block_ids=([1, 2],),
+        overwrite=True,
+    )
+    block_tables.apply_staged_writes()
+    torch.accelerator.synchronize()
+
+    assert torch.equal(
+        block_tables.block_tables[0].gpu[0, :2],
+        torch.tensor([1, 2], dtype=torch.int32, device=device),
+    )