chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,514 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user