wehub-resource-sync
515 lines
17 KiB
Python
515 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import inspect
|
|
import os
|
|
import sys
|
|
from contextlib import contextmanager
|
|
from types import ModuleType, SimpleNamespace
|
|
from typing import Any, cast
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from vllm.utils import jit_monitor
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_monitor():
|
|
"""Reset global monitor state between tests."""
|
|
jit_monitor._active = False
|
|
jit_monitor._mode = "warn"
|
|
jit_monitor._verbose = False
|
|
jit_monitor._cutedsl_hook_installed = False
|
|
jit_monitor._tilelang_hook_installed = False
|
|
jit_monitor._tilelang_jitimpl_compile_depth = 0
|
|
yield
|
|
jit_monitor._active = False
|
|
jit_monitor._mode = "warn"
|
|
jit_monitor._verbose = False
|
|
jit_monitor._cutedsl_hook_installed = False
|
|
jit_monitor._tilelang_hook_installed = False
|
|
jit_monitor._tilelang_jitimpl_compile_depth = 0
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Helpers — lightweight stand-ins for the modules ``activate()`` patches
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
def _make_fake_knobs(*, autotuning_print=False, jit_hook=None):
|
|
"""Build a minimal fake ``triton.knobs`` namespace."""
|
|
autotuning = SimpleNamespace(print=autotuning_print)
|
|
runtime = SimpleNamespace(jit_post_compile_hook=jit_hook)
|
|
return SimpleNamespace(autotuning=autotuning, runtime=runtime)
|
|
|
|
|
|
def _fake_cute_import_modules(compile_fn):
|
|
"""Fake Python's parent package + submodule for ``import cutlass.cute``."""
|
|
fake_cute = cast(Any, ModuleType("cutlass.cute"))
|
|
fake_cute.compile = compile_fn
|
|
fake_parent_package = cast(Any, ModuleType("cutlass"))
|
|
fake_parent_package.__path__ = []
|
|
fake_parent_package.cute = fake_cute
|
|
return {
|
|
"cutlass": fake_parent_package,
|
|
"cutlass.cute": fake_cute,
|
|
}
|
|
|
|
|
|
def _fake_cute_compile(*args, **kwargs):
|
|
return "compiled"
|
|
|
|
|
|
def _fake_tilelang_import_modules():
|
|
"""Fake Python's TileLang modules touched by ``jit_monitor.activate``."""
|
|
|
|
class FakeJITKernel:
|
|
def __init__(self, *args, **kwargs):
|
|
pass
|
|
|
|
class FakeJITImpl:
|
|
def __init__(self, func, signature):
|
|
self.func = func
|
|
self.signature = signature
|
|
self.mode = "lazy"
|
|
self._kernel_cache = {}
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
key, _ = self.func.parse_args(*args, **kwargs)
|
|
kernel = self._kernel_cache.get(key)
|
|
if kernel is None:
|
|
kernel = "compiled"
|
|
self._kernel_cache[key] = kernel
|
|
return kernel
|
|
|
|
fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel"))
|
|
fake_kernel.JITKernel = FakeJITKernel
|
|
|
|
fake_jit = cast(Any, ModuleType("tilelang.jit"))
|
|
fake_jit.JITImpl = FakeJITImpl
|
|
fake_jit.kernel = fake_kernel
|
|
|
|
fake_tilelang = cast(Any, ModuleType("tilelang"))
|
|
fake_tilelang.jit = fake_jit
|
|
|
|
return {
|
|
"tilelang": fake_tilelang,
|
|
"tilelang.jit": fake_jit,
|
|
"tilelang.jit.kernel": fake_kernel,
|
|
}
|
|
|
|
|
|
@contextmanager
|
|
def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile):
|
|
"""Patch the Triton and CuTeDSL imports touched by ``jit_monitor.activate``."""
|
|
fake_triton = cast(Any, ModuleType("triton"))
|
|
fake_triton.knobs = fake_knobs
|
|
with (
|
|
mock.patch.dict(
|
|
sys.modules,
|
|
{
|
|
"triton": fake_triton,
|
|
**_fake_cute_import_modules(cute_compile),
|
|
**_fake_tilelang_import_modules(),
|
|
},
|
|
),
|
|
mock.patch.object(jit_monitor, "HAS_TRITON", True),
|
|
):
|
|
yield
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Unit tests (no GPU required, triton is mocked)
|
|
# ------------------------------------------------------------------
|
|
|
|
|
|
class TestActivateBasic:
|
|
def test_sets_active(self):
|
|
assert not jit_monitor.is_active()
|
|
with _patch_jit_modules(_make_fake_knobs()):
|
|
jit_monitor.activate()
|
|
assert jit_monitor.is_active()
|
|
|
|
def test_idempotent(self):
|
|
fake = _make_fake_knobs()
|
|
with _patch_jit_modules(fake):
|
|
jit_monitor.activate()
|
|
first_hook = fake.runtime.jit_post_compile_hook
|
|
jit_monitor.activate()
|
|
assert fake.runtime.jit_post_compile_hook is first_hook
|
|
|
|
def test_logs_info_on_activation(self):
|
|
with (
|
|
mock.patch.object(jit_monitor.logger, "info") as m,
|
|
_patch_jit_modules(_make_fake_knobs()),
|
|
):
|
|
jit_monitor.activate()
|
|
m.assert_called_once()
|
|
assert "Kernel JIT monitor activated" in m.call_args[0][0]
|
|
|
|
def test_rejects_unknown_mode(self):
|
|
with pytest.raises(ValueError, match="Unsupported JIT monitor mode"):
|
|
jit_monitor.activate(mode="panic") # type: ignore[arg-type]
|
|
|
|
|
|
class TestAutotuningPrint:
|
|
def test_enables_autotuning_print(self):
|
|
fake = _make_fake_knobs(autotuning_print=False)
|
|
with _patch_jit_modules(fake):
|
|
jit_monitor.activate()
|
|
assert fake.autotuning.print is True
|
|
|
|
def test_respects_user_opt_out(self):
|
|
fake = _make_fake_knobs(autotuning_print=False)
|
|
with (
|
|
mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "0"}),
|
|
_patch_jit_modules(fake),
|
|
):
|
|
jit_monitor.activate()
|
|
assert fake.autotuning.print is False
|
|
|
|
def test_noop_when_user_already_enabled(self):
|
|
fake = _make_fake_knobs(autotuning_print=True)
|
|
with (
|
|
mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "1"}),
|
|
_patch_jit_modules(fake),
|
|
):
|
|
jit_monitor.activate()
|
|
assert fake.autotuning.print is True
|
|
|
|
|
|
class TestTritonJitHook:
|
|
def test_hook_registered(self):
|
|
fake = _make_fake_knobs()
|
|
assert fake.runtime.jit_post_compile_hook is None
|
|
with _patch_jit_modules(fake):
|
|
jit_monitor.activate()
|
|
assert fake.runtime.jit_post_compile_hook is not None
|
|
|
|
def test_hook_logs_warning(self):
|
|
fake = _make_fake_knobs()
|
|
with _patch_jit_modules(fake):
|
|
jit_monitor.activate()
|
|
|
|
hook = fake.runtime.jit_post_compile_hook
|
|
mock_fn = SimpleNamespace(name="test_kernel")
|
|
|
|
with (
|
|
mock.patch.object(jit_monitor.logger, "warning_once") as m,
|
|
mock.patch.object(jit_monitor.logger, "warning") as warning,
|
|
):
|
|
hook(
|
|
key="some_key",
|
|
repr="some_repr",
|
|
fn=mock_fn,
|
|
compile=lambda: None,
|
|
is_manual_warmup=False,
|
|
already_compiled=False,
|
|
)
|
|
|
|
m.assert_called_once()
|
|
warning.assert_not_called()
|
|
msg = m.call_args[0][0] % m.call_args[0][1:]
|
|
assert "Triton kernel JIT compilation during inference" in msg
|
|
assert "test_kernel" in msg
|
|
|
|
def test_hook_chains_existing_hook(self):
|
|
existing = mock.MagicMock(return_value="existing_result")
|
|
fake = _make_fake_knobs(jit_hook=existing)
|
|
with _patch_jit_modules(fake):
|
|
jit_monitor.activate()
|
|
|
|
hook = fake.runtime.jit_post_compile_hook
|
|
mock_fn = SimpleNamespace(name="chained_kernel")
|
|
kwargs = dict(
|
|
key="k",
|
|
repr="r",
|
|
fn=mock_fn,
|
|
compile=lambda: None,
|
|
is_manual_warmup=False,
|
|
already_compiled=False,
|
|
)
|
|
result = hook(**kwargs)
|
|
|
|
existing.assert_called_once()
|
|
assert result == "existing_result"
|
|
|
|
def test_hook_works_without_existing_hook(self):
|
|
fake = _make_fake_knobs(jit_hook=None)
|
|
with _patch_jit_modules(fake):
|
|
jit_monitor.activate()
|
|
|
|
hook = fake.runtime.jit_post_compile_hook
|
|
mock_fn = SimpleNamespace(name="solo_kernel")
|
|
result = hook(
|
|
key="k",
|
|
repr="r",
|
|
fn=mock_fn,
|
|
compile=lambda: None,
|
|
is_manual_warmup=False,
|
|
already_compiled=False,
|
|
)
|
|
assert result is None
|
|
|
|
def test_error_mode_raises(self):
|
|
fake = _make_fake_knobs()
|
|
with _patch_jit_modules(fake):
|
|
jit_monitor.activate(mode="error")
|
|
|
|
hook = fake.runtime.jit_post_compile_hook
|
|
mock_fn = SimpleNamespace(name="error_kernel")
|
|
with pytest.raises(RuntimeError, match="Triton kernel JIT compilation"):
|
|
hook(
|
|
key="k",
|
|
repr="r",
|
|
fn=mock_fn,
|
|
compile=lambda: None,
|
|
is_manual_warmup=False,
|
|
already_compiled=False,
|
|
)
|
|
|
|
|
|
class TestNoTritonFallback:
|
|
def test_activate_without_triton(self):
|
|
with mock.patch.object(jit_monitor, "HAS_TRITON", False):
|
|
jit_monitor.activate()
|
|
assert jit_monitor.is_active()
|
|
|
|
|
|
class TestCuTeDSLHook:
|
|
def test_compile_logs_warning(self):
|
|
def compile_fn(*args, **kwargs):
|
|
return "compiled"
|
|
|
|
with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn):
|
|
import cutlass.cute as cute
|
|
|
|
jit_monitor.activate()
|
|
with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
|
|
result = cute.compile(lambda: None, "arg", option=True)
|
|
|
|
assert result == "compiled"
|
|
warning_once.assert_called_once()
|
|
msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
|
|
assert "CuTeDSL JIT compilation during inference" in msg
|
|
|
|
def test_compile_logs_verbose_warning(self):
|
|
def compile_fn(*args, **kwargs):
|
|
return "compiled"
|
|
|
|
with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn):
|
|
import cutlass.cute as cute
|
|
|
|
jit_monitor.activate(verbose=True)
|
|
with mock.patch.object(jit_monitor.logger, "warning") as warning:
|
|
result = cute.compile(lambda: None, "arg", option=True)
|
|
|
|
assert result == "compiled"
|
|
warning.assert_called_once()
|
|
msg = warning.call_args[0][0] % warning.call_args[0][1:]
|
|
assert "CuTeDSL JIT compilation during inference" in msg
|
|
|
|
def test_error_mode_raises(self):
|
|
def compile_fn(*args, **kwargs):
|
|
return "compiled"
|
|
|
|
with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn):
|
|
import cutlass.cute as cute
|
|
|
|
jit_monitor.activate(mode="error")
|
|
with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"):
|
|
cute.compile(lambda: None, "arg", option=True)
|
|
|
|
|
|
class TestTileLangHook:
|
|
def test_jit_kernel_logs_warning(self):
|
|
with _patch_jit_modules(_make_fake_knobs()):
|
|
from tilelang.jit.kernel import JITKernel
|
|
|
|
func = SimpleNamespace(attrs={"global_symbol": "tl_kernel"})
|
|
jit_monitor.activate()
|
|
with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
|
|
JITKernel(func=func, out_idx=None, execution_backend="tvm_ffi")
|
|
|
|
warning_once.assert_called_once()
|
|
msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
|
|
assert "TileLang JIT compilation during inference" in msg
|
|
assert "tl_kernel" in msg
|
|
|
|
def test_jit_impl_logs_warning(self):
|
|
with _patch_jit_modules(_make_fake_knobs()):
|
|
from tilelang.jit import JITImpl
|
|
|
|
def tilelang_fn(
|
|
gemm_out_mul,
|
|
hidden_size: int,
|
|
n_splits: int = 1,
|
|
hc_mult: int = 4,
|
|
):
|
|
return None
|
|
|
|
class FakeFunc:
|
|
orig_func = tilelang_fn
|
|
|
|
def parse_args(self, *args, **kwargs):
|
|
return (
|
|
(
|
|
"tilelang_key",
|
|
kwargs["hidden_size"],
|
|
kwargs.get("n_splits", 1),
|
|
),
|
|
{},
|
|
)
|
|
|
|
def set_mode(self, mode):
|
|
self.mode = mode
|
|
|
|
tensor = SimpleNamespace(
|
|
shape=(2, 16, 24),
|
|
dtype="float32",
|
|
device="cuda:0",
|
|
)
|
|
impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn))
|
|
|
|
jit_monitor.activate()
|
|
with (
|
|
mock.patch.object(jit_monitor.logger, "warning_once") as warning_once,
|
|
mock.patch.object(jit_monitor.logger, "warning") as warning,
|
|
):
|
|
impl(tensor, hidden_size=7168, n_splits=2)
|
|
|
|
warning_once.assert_called_once()
|
|
warning.assert_not_called()
|
|
msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
|
|
assert "TileLang JIT compilation during inference" in msg
|
|
assert "tilelang_fn" in msg
|
|
|
|
def test_jit_impl_does_not_log_on_cache_hit(self):
|
|
with _patch_jit_modules(_make_fake_knobs()):
|
|
from tilelang.jit import JITImpl
|
|
|
|
def tilelang_fn(gemm_out_mul, n_splits: int = 1):
|
|
return None
|
|
|
|
class FakeFunc:
|
|
orig_func = tilelang_fn
|
|
|
|
def parse_args(self, *args, **kwargs):
|
|
return (("tilelang_key", kwargs.get("n_splits", 1)), {})
|
|
|
|
def set_mode(self, mode):
|
|
self.mode = mode
|
|
|
|
tensor = SimpleNamespace(shape=(2, 16, 24), dtype="float32")
|
|
impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn))
|
|
|
|
jit_monitor.activate()
|
|
with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
|
|
impl(tensor, n_splits=2)
|
|
impl(tensor, n_splits=2)
|
|
|
|
warning_once.assert_called_once()
|
|
|
|
def test_from_database_does_not_log(self):
|
|
with _patch_jit_modules(_make_fake_knobs()):
|
|
from tilelang.jit.kernel import JITKernel
|
|
|
|
func = SimpleNamespace(attrs={"global_symbol": "cached_tl_kernel"})
|
|
jit_monitor.activate()
|
|
with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
|
|
JITKernel(func=func, from_database=True)
|
|
|
|
warning_once.assert_not_called()
|
|
|
|
def test_error_mode_raises(self):
|
|
with _patch_jit_modules(_make_fake_knobs()):
|
|
from tilelang.jit.kernel import JITKernel
|
|
|
|
func = SimpleNamespace(attrs={"global_symbol": "error_tl_kernel"})
|
|
jit_monitor.activate(mode="error")
|
|
with pytest.raises(RuntimeError, match="TileLang JIT compilation"):
|
|
JITKernel(func=func)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Integration tests (real Triton + GPU)
|
|
# ------------------------------------------------------------------
|
|
|
|
try:
|
|
import torch
|
|
|
|
_HAS_CUDA = torch.cuda.is_available()
|
|
except ImportError:
|
|
_HAS_CUDA = False
|
|
|
|
try:
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
_HAS_TRITON = True
|
|
except ImportError:
|
|
_HAS_TRITON = False
|
|
|
|
_skip_no_gpu = pytest.mark.skipif(
|
|
not (_HAS_CUDA and _HAS_TRITON),
|
|
reason="Requires CUDA GPU and Triton",
|
|
)
|
|
|
|
|
|
if _HAS_TRITON:
|
|
|
|
@triton.jit
|
|
def _add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
|
|
pid = tl.program_id(0)
|
|
offs = pid * BLOCK + tl.arange(0, BLOCK)
|
|
mask = offs < n
|
|
x = tl.load(x_ptr + offs, mask=mask)
|
|
y = tl.load(y_ptr + offs, mask=mask)
|
|
tl.store(out_ptr + offs, x + y, mask=mask)
|
|
|
|
|
|
def _run_add_kernel(n: int, block: int = 256, offset: int = 0) -> None:
|
|
"""Launch ``_add_kernel`` with vectors of length *n*."""
|
|
x = torch.randn(n + offset, device="cuda")[offset:] # affect alignment
|
|
y = torch.randn(n, device="cuda")
|
|
out = torch.empty(n, device="cuda")
|
|
grid = ((n + block - 1) // block,)
|
|
_add_kernel[grid](x, y, out, n, BLOCK=block)
|
|
torch.accelerator.synchronize()
|
|
|
|
|
|
@_skip_no_gpu
|
|
class TestTritonJitHookIntegration:
|
|
"""End-to-end: real Triton kernel, real GPU, real hook."""
|
|
|
|
def test_no_warning_on_cached_shape(self):
|
|
_run_add_kernel(1024)
|
|
|
|
jit_monitor.activate()
|
|
with mock.patch.object(jit_monitor.logger, "warning_once") as w:
|
|
_run_add_kernel(1024)
|
|
w.assert_not_called()
|
|
|
|
def test_warning_on_new_constexpr(self):
|
|
_run_add_kernel(1024, block=256)
|
|
|
|
jit_monitor.activate()
|
|
with mock.patch.object(jit_monitor.logger, "warning_once") as w:
|
|
# Different BLOCK (a tl.constexpr) forces recompilation.
|
|
_run_add_kernel(1024, block=512)
|
|
w.assert_called()
|
|
msg = w.call_args[0][0] % w.call_args[0][1:]
|
|
assert "_add_kernel" in msg
|
|
|
|
def test_verbose_warning_on_each_new_pointer_alignment(self):
|
|
_run_add_kernel(1024)
|
|
|
|
jit_monitor.activate(verbose=True)
|
|
with (
|
|
mock.patch.object(jit_monitor.logger, "warning") as w,
|
|
mock.patch.object(jit_monitor.logger, "warning_once") as w_once,
|
|
):
|
|
_run_add_kernel(1024, offset=1)
|
|
assert w.called
|
|
w_once.assert_not_called()
|