402 lines
14 KiB
Python
402 lines
14 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
# Attempt to import the optional module
|
|
try:
|
|
from dynamo.llm import BlockManager
|
|
except ImportError:
|
|
pytest.importorskip(
|
|
"optional_module", reason="block-manager feature is not enabled"
|
|
)
|
|
|
|
pytestmark = pytest.mark.pre_merge
|
|
|
|
|
|
WORKER_ID = 0
|
|
NUM_LAYER = 5
|
|
OUTER_DIM = 2
|
|
PAGE_SIZE = 4
|
|
INNER_DIM = 13
|
|
DTYPE, TORCH_DTYPE = "FP32", torch.float32
|
|
HOST_NUM_BLOCKS = 16
|
|
DEVICE_NUM_BLOCKS = 16
|
|
DEVICE_ID = 0
|
|
|
|
|
|
def new_block_manager():
|
|
return BlockManager(
|
|
WORKER_ID,
|
|
NUM_LAYER,
|
|
OUTER_DIM,
|
|
PAGE_SIZE,
|
|
INNER_DIM,
|
|
DTYPE,
|
|
HOST_NUM_BLOCKS,
|
|
DEVICE_NUM_BLOCKS,
|
|
DEVICE_ID,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def block_manager():
|
|
return new_block_manager()
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
|
|
async def test_block_manager_initialization():
|
|
# Python should drop the BlockManager instance as soon as it goes out of scope, but
|
|
# it may not be garbage collected immediately, depending on the garbage collector.
|
|
BlockManager(WORKER_ID, NUM_LAYER, OUTER_DIM, PAGE_SIZE, INNER_DIM)
|
|
BlockManager(WORKER_ID, NUM_LAYER, OUTER_DIM, PAGE_SIZE, INNER_DIM, DTYPE)
|
|
BlockManager(
|
|
WORKER_ID, NUM_LAYER, OUTER_DIM, PAGE_SIZE, INNER_DIM, DTYPE, HOST_NUM_BLOCKS
|
|
)
|
|
BlockManager(
|
|
WORKER_ID,
|
|
NUM_LAYER,
|
|
OUTER_DIM,
|
|
PAGE_SIZE,
|
|
INNER_DIM,
|
|
DTYPE,
|
|
device_num_blocks=DEVICE_NUM_BLOCKS,
|
|
)
|
|
BlockManager(
|
|
WORKER_ID,
|
|
NUM_LAYER,
|
|
OUTER_DIM,
|
|
PAGE_SIZE,
|
|
INNER_DIM,
|
|
DTYPE,
|
|
HOST_NUM_BLOCKS,
|
|
DEVICE_NUM_BLOCKS,
|
|
)
|
|
BlockManager(
|
|
WORKER_ID,
|
|
NUM_LAYER,
|
|
OUTER_DIM,
|
|
PAGE_SIZE,
|
|
INNER_DIM,
|
|
DTYPE,
|
|
device_num_blocks=DEVICE_NUM_BLOCKS,
|
|
device_id=DEVICE_ID,
|
|
)
|
|
BlockManager(
|
|
WORKER_ID,
|
|
NUM_LAYER,
|
|
OUTER_DIM,
|
|
PAGE_SIZE,
|
|
INNER_DIM,
|
|
DTYPE,
|
|
HOST_NUM_BLOCKS,
|
|
DEVICE_NUM_BLOCKS,
|
|
DEVICE_ID,
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
|
|
async def test_cpu_block_access(block_manager: BlockManager):
|
|
block_count = 2
|
|
block_list = block_manager.allocate_host_blocks_blocking(block_count)
|
|
blocks = block_list.to_list()
|
|
assert len(blocks) == block_count
|
|
tensors = [torch.from_dlpack(b) for b in blocks]
|
|
for tensor in tensors:
|
|
assert tensor.get_device() == -1 # CPU
|
|
assert tensor.shape == (1, NUM_LAYER, OUTER_DIM, PAGE_SIZE, INNER_DIM)
|
|
assert tensor.dtype == TORCH_DTYPE
|
|
# print(tensors)
|
|
for tensor in tensors:
|
|
tensor[0][0][0][0][0] = 1.0
|
|
tensor[0][NUM_LAYER - 1][OUTER_DIM - 1][PAGE_SIZE - 1][INNER_DIM - 1] = 1.0
|
|
# print(tensors)
|
|
blocks_ = block_list.to_list()
|
|
assert blocks is not blocks_
|
|
assert len(blocks) == len(blocks_)
|
|
tensors_ = [torch.from_dlpack(b) for b in blocks_]
|
|
for tensor, tensor_ in zip(tensors, tensors_):
|
|
assert tensor is not tensor_
|
|
assert tensor.shape == tensor_.shape
|
|
assert tensor.dtype == tensor_.dtype
|
|
assert torch.allclose(tensor, tensor_)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
|
|
async def test_gpu_block_access(block_manager: BlockManager):
|
|
block_count = 6
|
|
block_list = block_manager.allocate_device_blocks_blocking(block_count)
|
|
blocks = block_list.to_list()
|
|
assert len(blocks) == block_count
|
|
tensors = [torch.from_dlpack(b) for b in blocks]
|
|
for tensor in tensors:
|
|
assert tensor.get_device() == DEVICE_ID # GPU
|
|
assert tensor.shape == (1, NUM_LAYER, OUTER_DIM, PAGE_SIZE, INNER_DIM)
|
|
assert tensor.dtype == TORCH_DTYPE
|
|
# print(tensors)
|
|
for tensor in tensors:
|
|
tensor[0][0][0][0][0] = 1.0
|
|
tensor[0][NUM_LAYER - 1][OUTER_DIM - 1][PAGE_SIZE - 1][INNER_DIM - 1] = 1.0
|
|
# print(tensors)
|
|
blocks_ = block_list.to_list()
|
|
assert blocks is not blocks_
|
|
assert len(blocks) == len(blocks_)
|
|
tensors_ = [torch.from_dlpack(b) for b in blocks_]
|
|
for tensor, tensor_ in zip(tensors, tensors_):
|
|
assert tensor is not tensor_
|
|
assert tensor.shape == tensor_.shape
|
|
assert tensor.dtype == tensor_.dtype
|
|
assert torch.allclose(tensor, tensor_)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
|
|
async def test_block_list_iteration(block_manager: BlockManager):
|
|
block_count = 4
|
|
block_list = await block_manager.allocate_host_blocks(block_count)
|
|
# Test __len__()
|
|
assert len(block_list) == block_count
|
|
# Test __getitem__()
|
|
for i in range(block_count):
|
|
block = block_list[i]
|
|
tensor = torch.from_dlpack(block)
|
|
tensor[0][0][0][0][0] = 1.0 + i
|
|
# Test __iter__() and __next__()
|
|
idx = 1.0
|
|
for block in block_list:
|
|
tensor = torch.from_dlpack(block)
|
|
assert tensor[0][0][0][0][0] == idx
|
|
tensor[0][0][0][0][0] += 0.5
|
|
idx += 1.0
|
|
assert idx == 1.0 + block_count
|
|
# Test __iter__() should reset current index
|
|
idx = 1.0
|
|
for block in block_list:
|
|
tensor = torch.from_dlpack(block)
|
|
assert tensor[0][0][0][0][0] == idx + 0.5
|
|
idx += 1.0
|
|
assert idx == 1.0 + block_count
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
|
|
async def test_block_copy_g1_g2(block_manager: BlockManager):
|
|
# Allocate device (G1) and host (G2) block
|
|
host_block_list = await block_manager.allocate_host_blocks(1)
|
|
device_block_list = await block_manager.allocate_device_blocks(1)
|
|
# Populate host block with unique values
|
|
host_tensor = torch.from_dlpack(host_block_list[0])
|
|
for i in range(NUM_LAYER):
|
|
for j in range(OUTER_DIM):
|
|
for k in range(PAGE_SIZE):
|
|
for w in range(INNER_DIM):
|
|
host_tensor[0][i][j][k][w] = (
|
|
i * OUTER_DIM * PAGE_SIZE * INNER_DIM
|
|
+ j * PAGE_SIZE * INNER_DIM
|
|
+ k * INNER_DIM
|
|
+ w
|
|
)
|
|
# Copy host block to device block after permuting
|
|
permute_dims = (0, 2, 4, 3, 1)
|
|
device_tensor_ = torch.from_dlpack(device_block_list[0]).permute(*permute_dims)
|
|
device_tensor_.copy_(host_tensor.permute(*permute_dims))
|
|
# Assert device block is contiguous and updated in block manager
|
|
device_tensor = torch.from_dlpack(device_block_list[0])
|
|
for i in range(NUM_LAYER):
|
|
for j in range(OUTER_DIM):
|
|
for k in range(PAGE_SIZE):
|
|
for w in range(INNER_DIM):
|
|
assert (
|
|
device_tensor[0][i][j][k][w]
|
|
== i * OUTER_DIM * PAGE_SIZE * INNER_DIM
|
|
+ j * PAGE_SIZE * INNER_DIM
|
|
+ k * INNER_DIM
|
|
+ w
|
|
)
|
|
# Set host block to zero and assert updated in block manager
|
|
host_tensor_ = torch.from_dlpack(host_block_list[0]).permute(*permute_dims)
|
|
host_tensor_.zero_()
|
|
assert torch.all(host_tensor == 0)
|
|
# Copy device block back to host block
|
|
host_tensor_.copy_(device_tensor_)
|
|
# Assert host block is updated in block manager
|
|
for i in range(NUM_LAYER):
|
|
for j in range(OUTER_DIM):
|
|
for k in range(PAGE_SIZE):
|
|
for w in range(INNER_DIM):
|
|
assert (
|
|
host_tensor[0][i][j][k][w]
|
|
== i * OUTER_DIM * PAGE_SIZE * INNER_DIM
|
|
+ j * PAGE_SIZE * INNER_DIM
|
|
+ k * INNER_DIM
|
|
+ w
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
|
|
async def test_cpu_layer_access(block_manager: BlockManager):
|
|
block_list = block_manager.allocate_host_blocks_blocking(1)
|
|
block = block_list[0]
|
|
layers = block.to_list()
|
|
assert len(layers) == NUM_LAYER
|
|
tensors = [torch.from_dlpack(bl) for bl in layers]
|
|
for tensor in tensors:
|
|
assert tensor.get_device() == -1 # CPU
|
|
assert tensor.shape == (1, 1, OUTER_DIM, PAGE_SIZE, INNER_DIM)
|
|
assert tensor.dtype == TORCH_DTYPE
|
|
# print(tensors)
|
|
for tensor in tensors:
|
|
tensor[0][0][0][0][0] = 1.0
|
|
tensor[0][0][OUTER_DIM - 1][PAGE_SIZE - 1][INNER_DIM - 1] = 1.0
|
|
# print(tensors)
|
|
layers_ = block.to_list()
|
|
assert layers is not layers_
|
|
assert len(layers) == len(layers_)
|
|
tensors_ = [torch.from_dlpack(bl) for bl in layers_]
|
|
for tensor, tensor_ in zip(tensors, tensors_):
|
|
assert tensor is not tensor_
|
|
assert tensor.shape == tensor_.shape
|
|
assert tensor.dtype == tensor_.dtype
|
|
assert torch.allclose(tensor, tensor_)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
|
|
async def test_gpu_layer_access(block_manager: BlockManager):
|
|
block_list = block_manager.allocate_device_blocks_blocking(1)
|
|
block = block_list[0]
|
|
layers = block.to_list()
|
|
assert len(layers) == NUM_LAYER
|
|
tensors = [torch.from_dlpack(bl) for bl in layers]
|
|
for tensor in tensors:
|
|
assert tensor.get_device() == DEVICE_ID # GPU
|
|
assert tensor.shape == (1, 1, OUTER_DIM, PAGE_SIZE, INNER_DIM)
|
|
assert tensor.dtype == TORCH_DTYPE
|
|
# print(tensors)
|
|
for tensor in tensors:
|
|
tensor[0][0][0][0][0] = 1.0
|
|
tensor[0][0][OUTER_DIM - 1][PAGE_SIZE - 1][INNER_DIM - 1] = 1.0
|
|
# print(tensors)
|
|
layers_ = block.to_list()
|
|
assert layers is not layers_
|
|
assert len(layers) == len(layers_)
|
|
tensors_ = [torch.from_dlpack(bl) for bl in layers_]
|
|
for tensor, tensor_ in zip(tensors, tensors_):
|
|
assert tensor is not tensor_
|
|
assert tensor.shape == tensor_.shape
|
|
assert tensor.dtype == tensor_.dtype
|
|
assert torch.allclose(tensor, tensor_)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
|
|
async def test_block_iteration(block_manager: BlockManager):
|
|
block = (await block_manager.allocate_host_blocks(1))[0]
|
|
# Test __len__()
|
|
assert len(block) == NUM_LAYER
|
|
# Test __getitem__()
|
|
for i in range(NUM_LAYER):
|
|
layer = block[i]
|
|
tensor = torch.from_dlpack(layer)
|
|
tensor[0][0][0][0][0] = 1.0 + i
|
|
# Test __iter__() and __next__()
|
|
idx = 1.0
|
|
for layer in block:
|
|
tensor = torch.from_dlpack(layer)
|
|
assert tensor[0][0][0][0][0] == idx
|
|
tensor[0][0][0][0][0] += 0.5
|
|
idx += 1.0
|
|
assert idx == 1.0 + NUM_LAYER
|
|
# Test __iter__() should reset current index
|
|
idx = 1.0
|
|
for layer in block:
|
|
tensor = torch.from_dlpack(layer)
|
|
assert tensor[0][0][0][0][0] == idx + 0.5
|
|
idx += 1.0
|
|
assert idx == 1.0 + NUM_LAYER
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
|
|
async def test_block_layer_copy_g1_g2(block_manager: BlockManager):
|
|
# Allocate device (G1) and host (G2) block
|
|
host_block = (await block_manager.allocate_host_blocks(1))[0]
|
|
device_block = (await block_manager.allocate_device_blocks(1))[0]
|
|
# Populate host block at layer level with unique values
|
|
host_layer_tensors = [torch.from_dlpack(bl) for bl in host_block]
|
|
for i in range(NUM_LAYER):
|
|
host_layer_tensor = host_layer_tensors[i]
|
|
for j in range(OUTER_DIM):
|
|
for k in range(PAGE_SIZE):
|
|
for w in range(INNER_DIM):
|
|
host_layer_tensor[0][0][j][k][w] = (
|
|
i * OUTER_DIM * PAGE_SIZE * INNER_DIM
|
|
+ j * PAGE_SIZE * INNER_DIM
|
|
+ k * INNER_DIM
|
|
+ w
|
|
)
|
|
# Copy host block to device block after permuting
|
|
permute_dims = (0, 2, 4, 3, 1)
|
|
host_block_tensor_ = torch.from_dlpack(host_block).permute(*permute_dims)
|
|
device_block_tensor_ = torch.from_dlpack(device_block).permute(*permute_dims)
|
|
device_block_tensor_.copy_(host_block_tensor_)
|
|
# Assert device block is contiguous and updated in block manager at layer level
|
|
device_layer_tensors = [torch.from_dlpack(bl) for bl in device_block]
|
|
for i in range(NUM_LAYER):
|
|
device_layer_tensor = device_layer_tensors[i]
|
|
for j in range(OUTER_DIM):
|
|
for k in range(PAGE_SIZE):
|
|
for w in range(INNER_DIM):
|
|
assert (
|
|
device_layer_tensor[0][0][j][k][w]
|
|
== i * OUTER_DIM * PAGE_SIZE * INNER_DIM
|
|
+ j * PAGE_SIZE * INNER_DIM
|
|
+ k * INNER_DIM
|
|
+ w
|
|
)
|
|
# Set host block to zero and assert updated in block manager
|
|
host_block_tensor = torch.from_dlpack(host_block)
|
|
host_block_tensor.zero_()
|
|
assert torch.all(host_block_tensor_ == 0)
|
|
# Copy device block back to host block
|
|
host_block_tensor_.copy_(device_block_tensor_)
|
|
# Assert host block is updated in block manager
|
|
for i in range(NUM_LAYER):
|
|
for j in range(OUTER_DIM):
|
|
for k in range(PAGE_SIZE):
|
|
for w in range(INNER_DIM):
|
|
assert (
|
|
host_block_tensor[0][i][j][k][w]
|
|
== i * OUTER_DIM * PAGE_SIZE * INNER_DIM
|
|
+ j * PAGE_SIZE * INNER_DIM
|
|
+ k * INNER_DIM
|
|
+ w
|
|
)
|
|
|
|
|
|
async def main():
|
|
await test_block_manager_initialization()
|
|
await test_cpu_block_access(new_block_manager())
|
|
await test_gpu_block_access(new_block_manager())
|
|
await test_block_list_iteration(new_block_manager())
|
|
await test_block_copy_g1_g2(new_block_manager())
|
|
await test_cpu_layer_access(new_block_manager())
|
|
await test_gpu_layer_access(new_block_manager())
|
|
await test_block_iteration(new_block_manager())
|
|
await test_block_layer_copy_g1_g2(new_block_manager())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|