mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-04-06 12:37:04 +08:00
2c81fbbb3c
Instead of changing library names (not scalable), create a shared test_operations.py module that: - Provides a single "silly" library for all compilation tests - Registers a unified attention operation that can handle both standard and counting modes - Eliminates duplicate registration errors when running all tests together - Maintains backward compatibility with existing test behavior Addresses feedback to make the solution more scalable and maintainable. Co-authored-by: ProExpertProg <11367180+ProExpertProg@users.noreply.github.com>
110 lines
3.7 KiB
Python
110 lines
3.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Test the piecewise compilation with a simple model so that we
|
|
can exactly calculate the expected output and side effects.
|
|
"""
|
|
import pytest
|
|
import torch
|
|
from torch import nn
|
|
|
|
from vllm.compilation.counter import compilation_counter
|
|
from vllm.compilation.decorators import support_torch_compile
|
|
from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
|
|
VllmConfig, set_current_vllm_config)
|
|
from vllm.envs import VLLM_USE_V1
|
|
from vllm.forward_context import BatchDescriptor, set_forward_context
|
|
|
|
# Import shared test operations
|
|
from tests.compile.test_operations import (
|
|
get_global_counter, reset_global_counter, enable_counting_mode
|
|
)
|
|
|
|
# Enable counting mode for this test's specific behavior
|
|
enable_counting_mode()
|
|
|
|
|
|
@support_torch_compile
|
|
class SillyModel(nn.Module):
|
|
|
|
def __init__(self,
|
|
*,
|
|
vllm_config: VllmConfig,
|
|
prefix: str = '',
|
|
**kwargs) -> None:
|
|
super().__init__()
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Overall effect:
|
|
x += 1
|
|
x[0] += 2
|
|
global_counter += 2
|
|
"""
|
|
x = x + 1
|
|
x = x + 2
|
|
out = torch.empty_like(x)
|
|
torch.ops.silly.attention(x, x, x, out)
|
|
x = out
|
|
x = x - 2
|
|
x = x - 1
|
|
out = torch.empty_like(x)
|
|
torch.ops.silly.attention(x, x, x, out)
|
|
x = out
|
|
x = x + 1
|
|
return x
|
|
|
|
|
|
@pytest.mark.parametrize("use_inductor", [True, False])
|
|
def test_simple_piecewise_compile(use_inductor):
|
|
assert VLLM_USE_V1
|
|
|
|
vllm_config = VllmConfig(compilation_config=CompilationConfig(
|
|
level=CompilationLevel.PIECEWISE,
|
|
use_cudagraph=True,
|
|
use_inductor=use_inductor,
|
|
splitting_ops=["silly.attention"],
|
|
cudagraph_copy_inputs=True,
|
|
cudagraph_capture_sizes=[1, 2],
|
|
))
|
|
with set_current_vllm_config(vllm_config):
|
|
model = SillyModel(vllm_config=vllm_config, prefix='')
|
|
|
|
inputs = torch.randn(100).cuda()
|
|
|
|
with compilation_counter.expect(
|
|
num_graphs_seen=1, # one graph for the model
|
|
num_piecewise_graphs_seen=5, # 2 * num_layers + 1
|
|
num_piecewise_capturable_graphs_seen=3, # 1 + num_layers
|
|
num_backend_compilations=3, # num_piecewise_capturable_graphs_seen
|
|
num_cudagraph_captured=
|
|
6, # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
|
|
), set_forward_context(None,
|
|
vllm_config=vllm_config): # background context
|
|
# warm up with background context
|
|
model(inputs)
|
|
|
|
# capturing/replaying should under context of cudagraph dispatching
|
|
with set_forward_context(
|
|
None,
|
|
vllm_config=vllm_config,
|
|
cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
|
|
batch_descriptor=BatchDescriptor(num_tokens=2, )):
|
|
model(torch.randn(2).cuda())
|
|
with set_forward_context(
|
|
None,
|
|
vllm_config=vllm_config,
|
|
cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
|
|
batch_descriptor=BatchDescriptor(num_tokens=1, )):
|
|
model(torch.randn(1).cuda())
|
|
|
|
input = torch.zeros(2).cuda()
|
|
reset_global_counter()
|
|
with set_forward_context(
|
|
None,
|
|
vllm_config=vllm_config,
|
|
cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
|
|
batch_descriptor=BatchDescriptor(num_tokens=2, )):
|
|
output = model(input)
|
|
assert get_global_counter() == 2
|
|
assert torch.allclose(output.cpu(), torch.tensor([3., 1.])) |