forked from EngineX-MetaX/enginex-c_series-vllm
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import copy
|
|
import dataclasses
|
|
from contextlib import contextmanager
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class CompilationCounter:
|
|
num_models_seen: int = 0
|
|
num_graphs_seen: int = 0
|
|
# including the splitting ops
|
|
num_piecewise_graphs_seen: int = 0
|
|
# not including the splitting ops
|
|
num_piecewise_capturable_graphs_seen: int = 0
|
|
num_backend_compilations: int = 0
|
|
num_cudagraph_captured: int = 0
|
|
# InductorAdapter.compile calls
|
|
num_inductor_compiles: int = 0
|
|
# EagerAdapter.compile calls
|
|
num_eager_compiles: int = 0
|
|
|
|
def clone(self) -> "CompilationCounter":
|
|
return copy.deepcopy(self)
|
|
|
|
@contextmanager
|
|
def expect(self, **kwargs):
|
|
old = self.clone()
|
|
yield
|
|
for k, v in kwargs.items():
|
|
assert getattr(self, k) - getattr(old, k) == v, (
|
|
f"{k} not as expected, before it is {getattr(old, k)}"
|
|
f", after it is {getattr(self, k)}, "
|
|
f"expected diff is {v}")
|
|
|
|
|
|
compilation_counter = CompilationCounter()
|