diff --git a/spyre_inference/v1/attention/backends/spyre_attn.py b/spyre_inference/v1/attention/backends/spyre_attn.py index 081393c16..371a47b21 100644 --- a/spyre_inference/v1/attention/backends/spyre_attn.py +++ b/spyre_inference/v1/attention/backends/spyre_attn.py @@ -22,6 +22,7 @@ from typing import ClassVar, NamedTuple import torch +from torch._dynamo.utils import counters from vllm.config import CompilationMode, VllmConfig, get_current_vllm_config from vllm.config.cache import CacheDType from vllm.logger import init_logger @@ -436,6 +437,36 @@ def _batched_decode_kernel( _page_attn_compiled = torch.compile(_page_attn_kernel, dynamic=False) _batched_decode_compiled = torch.compile(_batched_decode_kernel, dynamic=False) +_warmup_complete = False + + +def mark_warmup_complete() -> None: + """Arm the late-compile warning, once warmup has claimed full variant coverage.""" + global _warmup_complete + _warmup_complete = True + + +def _call_kernel(label: str, fn, *args): + """Dispatch a kernel, warning if it compiles once warmup has claimed coverage. + + Dynamo's counter is process-wide but attributable across just this call: a + compiled region runs no eager ops, and torch-spyre compiles every eager aten op. + That assumes nothing else compiles concurrently on another thread, which holds for + a single-tenant serving process; if it ever stops holding, the cost is a spurious + warning, not a wrong result. + """ + if not _warmup_complete: + return fn(*args) + before = counters["stats"]["unique_graphs"] + result = fn(*args) + if counters["stats"]["unique_graphs"] != before: + logger.warning_once( + "%s compiled outside warmup, which costs a full Inductor compile mid-request. " + "Re-run with TORCH_LOGS=recompiles to see which guard failed.", + label, + ) + return result + @dataclass class SpyreAttentionMetadata(AttentionMetadata): @@ -1576,7 +1607,9 @@ def _run_batched_decode_dispatch( and output.storage_offset() == 0 and output.is_contiguous() ) - result = self._decode_fn( + result = _call_kernel( + "batched decode attention", + self._decode_fn, query_dev, attn_metadata.query_row_ids_dev if needs_gather else None, k_pages, @@ -1745,7 +1778,9 @@ def _online_softmax_attention( row_table = attn_metadata.query_row_tables[seq_idx] # Run attention on target device - result = self._attn_fn( + result = _call_kernel( + "page attention", + self._attn_fn, q_staging, row_table, k_pages, diff --git a/spyre_inference/v1/worker/spyre_model_runner.py b/spyre_inference/v1/worker/spyre_model_runner.py index bd74cdd2e..e2ab15023 100644 --- a/spyre_inference/v1/worker/spyre_model_runner.py +++ b/spyre_inference/v1/worker/spyre_model_runner.py @@ -87,6 +87,7 @@ SpyreAttentionImpl, SpyrePagedKVCache, allocate_staging_buffers, + mark_warmup_complete, ) from spyre_inference.v1.attention.spyre_attn_bucketer import SpyreAttnBucketer from spyre_inference.v1.pool import ( @@ -766,6 +767,8 @@ def _record_attention_graphs(self, token_counts: list[int]) -> None: total, time.time() - t0, ) + # Past the early returns: with recording off, first-use compiles are intended. + mark_warmup_complete() def _resolve_builder_attn_bucketer(self) -> SpyreAttnBucketer | None: """The attention bucketer the metadata builders dispatch against. diff --git a/tests/attention/test_spyre_attn_recorder.py b/tests/attention/test_spyre_attn_recorder.py index 170b25e10..8a0cdbd50 100644 --- a/tests/attention/test_spyre_attn_recorder.py +++ b/tests/attention/test_spyre_attn_recorder.py @@ -21,13 +21,16 @@ Spyre), which is enough to exercise dummy-arg construction and the guards. """ +import logging from unittest.mock import MagicMock import pytest import torch from torch._dynamo.utils import counters from vllm.config import CompilationMode, get_current_vllm_config +from vllm.logger import _print_warning_once +from spyre_inference.v1.attention.backends import spyre_attn from spyre_inference.v1.attention.backends.spyre_attn import ( SpyreAttentionImpl, SpyrePagedKVCache, @@ -236,3 +239,57 @@ def test_limit_is_restored_even_when_recording_raises(self, impl, kv_cache, monk with pytest.raises(RuntimeError): impl.record_graphs(torch.device("cpu"), make_bucketer(), kv_cache) assert torch._dynamo.config.accumulated_recompile_limit == before + + +def _toy_kernel(x, n): + """`n` is a plain int, so ``dynamic=False`` gives one graph per value, like num_blocks.""" + for _ in range(n): + x = x + 1 + return x + + +class TestLateCompileWarning: + """The runtime half of the acceptance criterion, for what the tests above cannot see: + a real config whose buckets miss something, or the batched decode kernel, which the + recorder never traces. ``backend="eager"`` suffices since the counter is Dynamo's. + """ + + @pytest.fixture(autouse=True) + def _isolated(self, monkeypatch): + torch._dynamo.reset() + # warning_once is lru_cached process-wide, so a prior emit would mask ours. + _print_warning_once.cache_clear() + monkeypatch.setattr(spyre_attn, "_warmup_complete", False) + yield + _print_warning_once.cache_clear() + + def test_quiet_before_warmup_is_marked(self, caplog): + fn = torch.compile(_toy_kernel, dynamic=False, backend="eager") + with caplog.at_level(logging.WARNING): + spyre_attn._call_kernel("page attention", fn, torch.ones(4), 1) + assert "outside warmup" not in caplog.text + + def test_warns_when_an_unrecorded_variant_compiles(self, caplog): + fn = torch.compile(_toy_kernel, dynamic=False, backend="eager") + spyre_attn._call_kernel("page attention", fn, torch.ones(4), 1) + spyre_attn.mark_warmup_complete() + + with caplog.at_level(logging.WARNING): + spyre_attn._call_kernel("page attention", fn, torch.ones(4), 2) + + assert "page attention compiled outside warmup" in caplog.text + + def test_quiet_when_the_variant_was_already_recorded(self, caplog): + fn = torch.compile(_toy_kernel, dynamic=False, backend="eager") + spyre_attn._call_kernel("page attention", fn, torch.ones(4), 1) + spyre_attn.mark_warmup_complete() + + with caplog.at_level(logging.WARNING): + spyre_attn._call_kernel("page attention", fn, torch.ones(4), 1) + + assert "outside warmup" not in caplog.text + + def test_mark_warmup_complete_arms_the_check(self): + assert spyre_attn._warmup_complete is False + spyre_attn.mark_warmup_complete() + assert spyre_attn._warmup_complete is True