feat(attn): warn when an attention kernel compiles after warmup

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,

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.

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