252 lines
8.9 KiB
Python
252 lines
8.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
import torch
|
||
from torch import nn
|
||
from torch.library import Library
|
||
|
||
from vllm.compilation.counter import compilation_counter
|
||
from vllm.compilation.decorators import (ignore_torch_compile,
|
||
support_torch_compile)
|
||
from vllm.config import (CacheConfig, CompilationConfig, CompilationLevel,
|
||
CUDAGraphMode, VllmConfig, set_current_vllm_config)
|
||
from vllm.forward_context import BatchDescriptor, set_forward_context
|
||
from vllm.utils import direct_register_custom_op
|
||
|
||
# create a library to hold the custom op
|
||
silly_lib = Library("silly", "FRAGMENT") # noqa
|
||
|
||
BATCH_SIZE = 32
|
||
MLP_SIZE = 128
|
||
|
||
|
||
def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
|
||
out: torch.Tensor) -> None:
|
||
out.copy_(q)
|
||
out += k
|
||
out += v
|
||
|
||
|
||
def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
|
||
out: torch.Tensor) -> None:
|
||
return
|
||
|
||
|
||
direct_register_custom_op(
|
||
op_name="attention",
|
||
op_func=silly_attention,
|
||
mutates_args=["out"],
|
||
fake_impl=silly_attention_fake,
|
||
target_lib=silly_lib,
|
||
)
|
||
|
||
|
||
@torch.inference_mode
|
||
def run_model(vllm_config: VllmConfig, model: nn.Module,
|
||
cudagraph_runtime_mode: CUDAGraphMode):
|
||
with set_forward_context({}, vllm_config=vllm_config):
|
||
# warmup for the model with cudagraph_mode NONE
|
||
model(torch.randn(BATCH_SIZE, MLP_SIZE).cuda())
|
||
|
||
# simulate cudagraphs capturing
|
||
with set_forward_context({},
|
||
vllm_config=vllm_config,
|
||
cudagraph_runtime_mode=cudagraph_runtime_mode,
|
||
batch_descriptor=BatchDescriptor(
|
||
num_tokens=2, )):
|
||
model(torch.randn(2, MLP_SIZE).cuda())
|
||
with set_forward_context({},
|
||
vllm_config=vllm_config,
|
||
cudagraph_runtime_mode=cudagraph_runtime_mode,
|
||
batch_descriptor=BatchDescriptor(
|
||
num_tokens=1, )):
|
||
model(torch.randn(1, MLP_SIZE).cuda())
|
||
|
||
# simulate cudagraphs replay
|
||
with set_forward_context({},
|
||
vllm_config=vllm_config,
|
||
cudagraph_runtime_mode=cudagraph_runtime_mode,
|
||
batch_descriptor=BatchDescriptor(
|
||
num_tokens=2, )):
|
||
output = model(torch.randn(2, MLP_SIZE).cuda())
|
||
|
||
output = output.cpu()
|
||
return output.cpu()
|
||
|
||
|
||
def test_ignore_torch_compile_decorator():
|
||
# piecewise
|
||
vllm_config = VllmConfig(compilation_config=CompilationConfig(
|
||
level=CompilationLevel.PIECEWISE,
|
||
use_cudagraph=True,
|
||
splitting_ops=["silly.attention"],
|
||
cudagraph_capture_sizes=[1, 2],
|
||
))
|
||
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
|
||
|
||
@support_torch_compile
|
||
class A(nn.Module):
|
||
|
||
def __init__(self,
|
||
*,
|
||
vllm_config: VllmConfig,
|
||
prefix: str = '',
|
||
**kwargs) -> None:
|
||
super().__init__()
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
x = x + x
|
||
attn_output = torch.empty_like(x)
|
||
torch.ops.silly.attention(x, x, x, attn_output)
|
||
x = attn_output
|
||
x = x * 3
|
||
return x
|
||
|
||
@ignore_torch_compile
|
||
class B(A):
|
||
...
|
||
|
||
@support_torch_compile
|
||
class C(B):
|
||
...
|
||
|
||
with set_current_vllm_config(vllm_config):
|
||
mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
|
||
|
||
# A has support_torch_compile
|
||
with compilation_counter.expect(
|
||
num_graphs_seen=1,
|
||
num_piecewise_graphs_seen=3,
|
||
num_piecewise_capturable_graphs_seen=2,
|
||
num_backend_compilations=2,
|
||
num_cudagraph_captured=4,
|
||
# num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
|
||
):
|
||
run_model(vllm_config, mod_A, cudagraph_runtime_mode)
|
||
|
||
with set_current_vllm_config(vllm_config):
|
||
mod_B = B(vllm_config=vllm_config, prefix='').eval().cuda()
|
||
|
||
# B's ignore_torch_compile should override A's support_torch_compile
|
||
with compilation_counter.expect(
|
||
num_graphs_seen=0,
|
||
num_piecewise_graphs_seen=0,
|
||
num_piecewise_capturable_graphs_seen=0,
|
||
num_backend_compilations=0,
|
||
num_cudagraph_captured=0,
|
||
):
|
||
run_model(vllm_config, mod_B, cudagraph_runtime_mode)
|
||
|
||
with set_current_vllm_config(vllm_config):
|
||
mod_C = C(vllm_config=vllm_config, prefix='').eval().cuda()
|
||
|
||
# C's support_torch_compile should override B's ignore_torch_compile
|
||
with compilation_counter.expect(
|
||
num_graphs_seen=1,
|
||
num_piecewise_graphs_seen=3,
|
||
num_piecewise_capturable_graphs_seen=2,
|
||
num_backend_compilations=2,
|
||
num_cudagraph_captured=4,
|
||
# num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
|
||
):
|
||
run_model(vllm_config, mod_C, cudagraph_runtime_mode)
|
||
|
||
|
||
# Only enable torch.compile if
|
||
# vllm_config.cache_config.kv_sharing_fast_prefill=True
|
||
@support_torch_compile(enable_if=lambda vllm_config: vllm_config.cache_config.
|
||
kv_sharing_fast_prefill)
|
||
class B(nn.Module):
|
||
|
||
def __init__(self,
|
||
*,
|
||
vllm_config: VllmConfig,
|
||
prefix: str = '',
|
||
**kwargs) -> None:
|
||
super().__init__()
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
x = x + x
|
||
attn_output = torch.empty_like(x)
|
||
torch.ops.silly.attention(x, x, x, attn_output)
|
||
x = attn_output
|
||
x = x + x
|
||
return x
|
||
|
||
|
||
# Only enable torch.compile if
|
||
# vllm_config.cache_config.kv_sharing_fast_prefill=False
|
||
@support_torch_compile(enable_if=lambda vllm_config: not vllm_config.
|
||
cache_config.kv_sharing_fast_prefill)
|
||
class A(nn.Module):
|
||
|
||
def __init__(self,
|
||
*,
|
||
vllm_config: VllmConfig,
|
||
prefix: str = '',
|
||
**kwargs) -> None:
|
||
super().__init__()
|
||
self.mod1 = B(vllm_config=vllm_config, prefix=prefix, **kwargs)
|
||
self.mod2 = B(vllm_config=vllm_config, prefix=prefix, **kwargs)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
x = self.mod1(x)
|
||
attn_output = torch.empty_like(x)
|
||
torch.ops.silly.attention(x, x, x, attn_output)
|
||
x = attn_output
|
||
x = self.mod2(x)
|
||
return x
|
||
|
||
|
||
def test_conditional_compile_enable_if():
|
||
vllm_config = VllmConfig(cache_config=CacheConfig(
|
||
kv_sharing_fast_prefill=True, ),
|
||
compilation_config=CompilationConfig(
|
||
level=CompilationLevel.PIECEWISE,
|
||
use_cudagraph=True,
|
||
splitting_ops=["silly.attention"],
|
||
cudagraph_capture_sizes=[1, 2],
|
||
))
|
||
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
|
||
|
||
with set_current_vllm_config(vllm_config):
|
||
mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
|
||
|
||
# A has support_torch_compile but enable_if fn returns False
|
||
# enalbe_if will be True for B, so we expect mod1 and mod2
|
||
# to be compiled
|
||
with compilation_counter.expect(
|
||
num_graphs_seen=2,
|
||
num_piecewise_graphs_seen=6,
|
||
# 3 piecewise graphs per instance of B()
|
||
num_piecewise_capturable_graphs_seen=4,
|
||
num_backend_compilations=4,
|
||
num_cudagraph_captured=8,
|
||
# num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
|
||
):
|
||
run_model(vllm_config, mod_A, cudagraph_runtime_mode)
|
||
|
||
# Set kv_sharing_fast_prefill=False
|
||
# which will cause A to be compiled and B to not be compiled
|
||
vllm_config = VllmConfig(cache_config=CacheConfig(
|
||
kv_sharing_fast_prefill=False, ),
|
||
compilation_config=CompilationConfig(
|
||
level=CompilationLevel.PIECEWISE,
|
||
use_cudagraph=True,
|
||
splitting_ops=["silly.attention"],
|
||
cudagraph_capture_sizes=[1, 2],
|
||
))
|
||
|
||
with set_current_vllm_config(vllm_config):
|
||
mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
|
||
|
||
with compilation_counter.expect(
|
||
num_graphs_seen=1,
|
||
num_piecewise_graphs_seen=7,
|
||
# 3 attn ops and 4 non-attn ops
|
||
num_piecewise_capturable_graphs_seen=4,
|
||
num_backend_compilations=4,
|
||
num_cudagraph_captured=8,
|
||
# num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
|
||
):
|
||
run_model(vllm_config, mod_A, cudagraph_runtime_mode)
|