forked from Karylab-cklius/vllm
360 lines
12 KiB
Python
360 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
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
|
|
yield
|
|
jit_monitor._active = False
|
|
jit_monitor._mode = "warn"
|
|
jit_monitor._verbose = False
|
|
jit_monitor._cutedsl_hook_installed = False
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# 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"
|
|
|
|
|
|
@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),
|
|
},
|
|
),
|
|
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)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# 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()
|