feat(attn): warn when an attention kernel compiles after warmup - #800
Conversation
|
👋 Hi! Thank you for contributing. We also recommend installing prek and configuring it to check your code before every local commit. |
178ca41 to
cc8d56e
Compare
The recorder claims warmup covers every variant a run can reach, but nothing reported a miss at runtime: the log line that fired on a lazy compile is gone, and it only ever covered the page kernel anyway. The batched decode kernel is not recorded at all, so under SPYRE_BATCHED_DECODE it has no coverage guarantee. Both dispatch sites now read Dynamo's graph counter across the call and warn once if it moved. Armed at the end of attention recording, so the paths that skip recording keep compiling on first use without warning. Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
cc8d56e to
241eaa8
Compare
jvlunteren
left a comment
jvlunteren
left a comment
There was a problem hiding this comment.
The gap is genuine. #799 removed the log line that fired on lazy compiles, and, as the body says, it only ever covered the page kernel. The batched decode kernel is never recorded, so under SPYRE_BATCHED_DECODE there is no coverage guarantee at all.
The mechanism checks out. I verified counters["stats"]["unique_graphs"] increments once per graph compilation at output_graph.py:2933, inside compile_and_call_fx_graph, so a delta across the call detects a genuinely new graph, not a cache hit. That's the right counter, and it is the distinction #792 got wrong (a dict hit is not a recompile guarantee).
The attributability argument is the important claim, and it is stated in the docstring:
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.
I would accept that, with one reservation: it holds only while nothing else on another thread compiles concurrently. In a single-tenant serving process that is fine, and a false positive here costs one spurious warning, not a failure. Worth noting rather than blocking.
The counter delta is attributable to a single call only while nothing else compiles on another thread. Holds for a single-tenant serving process, and a violation costs a spurious warning rather than a wrong result. Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
|
⚙️ Auto-merge enabled — dispatched Automated: enabling auto-merge requests the gate, so the PR is not left waiting on a check nobody asked for. Re-enabling auto-merge on this same commit will not dispatch again. |
🔄 spyre-test: runningTriggered by: /spyre-test comment Plan (build waves + dependencies, per arch)amd64 flowchart LR
subgraph Lamd64_0["amd64 L0 · 1 parallel"]
n_amd64_torch_spyre_torch_spyre_dev["torch-spyre/torch-spyre-dev 🟢 ↺<br/>image · 72411d836e22"]
end
subgraph Lamd64_1["amd64 L1 · 1 parallel"]
n_amd64_hf_adapters_hf_adapters_dev["hf-adapters/hf-adapters-dev 🟢<br/>image · 5247b8e41856"]
end
subgraph Lamd64_2["amd64 L2 · 1 parallel"]
n_amd64_spyre_inference_spyre_inference_dev["spyre-inference/spyre-inference-dev 🟢<br/>image · 8f1976c5cef2"]
end
n_amd64_torch_spyre_torch_spyre_dev --> n_amd64_hf_adapters_hf_adapters_dev
n_amd64_hf_adapters_hf_adapters_dev --> n_amd64_spyre_inference_spyre_inference_dev
classDef sPending fill:#eceff1,stroke:#90a4ae,color:#37474f
classDef sBuilding fill:#fff8e1,stroke:#f9a825,color:#5d4037,stroke-width:2px
classDef sOk fill:#e8f5e9,stroke:#43a047,color:#1b5e20
classDef sReused fill:#e3f2fd,stroke:#1e88e5,color:#0d47a1
classDef sFailed fill:#ffebee,stroke:#e53935,color:#b71c1c,stroke-width:2px
classDef sDropped fill:#f5f5f5,stroke:#bdbdbd,color:#9e9e9e
class n_amd64_torch_spyre_torch_spyre_dev sReused;
class n_amd64_hf_adapters_hf_adapters_dev sOk;
class n_amd64_spyre_inference_spyre_inference_dev sOk;
🔨 orch trigger-pr-validation · arches amd64 · fp amd64=dc820e01
|
🔄 merge-queue-integration: runningPlan (build waves + dependencies, per arch)amd64 flowchart LR
subgraph Lamd64_0["amd64 L0 · 1 parallel"]
n_amd64_hf_adapters_hf_adapters_dev["hf-adapters/hf-adapters-dev 🔴<br/>image · bc0ac51daa38"]
end
subgraph Lamd64_1["amd64 L1 · 1 parallel"]
n_amd64_spyre_inference_spyre_inference_dev["spyre-inference/spyre-inference-dev 🔴<br/>image · 29d17914b03c"]
end
n_amd64_hf_adapters_hf_adapters_dev --> n_amd64_spyre_inference_spyre_inference_dev
classDef sPending fill:#eceff1,stroke:#90a4ae,color:#37474f
classDef sBuilding fill:#fff8e1,stroke:#f9a825,color:#5d4037,stroke-width:2px
classDef sOk fill:#e8f5e9,stroke:#43a047,color:#1b5e20
classDef sReused fill:#e3f2fd,stroke:#1e88e5,color:#0d47a1
classDef sFailed fill:#ffebee,stroke:#e53935,color:#b71c1c,stroke-width:2px
classDef sDropped fill:#f5f5f5,stroke:#bdbdbd,color:#9e9e9e
class n_amd64_hf_adapters_hf_adapters_dev sOk;
class n_amd64_spyre_inference_spyre_inference_dev sOk;
🔨 orch trigger-pr-validation · arches amd64 · fp amd64=7763559f
|
|
GHA test runs:
|
|
GHA test runs:
|
|
GHA test runs:
|
|
GHA test runs:
|
|
GHA test runs:
|
_record_one dispatched to self._attn_fn unconditionally, so with SPYRE_LX_KV_LAYOUT=1 warmup fed the (page, kv_head)-folded cache into the slot-major kernel and every variant died in k_page.squeeze(0).permute(1, 0, 2): Recorded 0/12 for all 40 layers. Only forward() had learned about _lx_attn_fn. The variants then compiled in the serving path, costing a 325s warmup iteration against ~17s once recorded. The recorder now builds the folded kernel's own arguments -- per-block [num_kv_heads, 1] index tables (one real tensor per block, torch-spyre#3770), head index tables, block_size, per-group ALiBi tiles, fused-store row tables -- and calls it inside _capped_attn_cores, so the recorded graph is the one dispatch reuses rather than an uncapped near-miss. Also routes the LX dispatch site through torch-spyre#800's _call_kernel, which that PR predates, so the folded path is covered by the late-compile warning too. Measured on granite-3.3-8b, 2048 in / 64 out / bs 1, torch-spyre built from #4153: 14.333s on main vs 13.829s with SPYRE_LX_KV_LAYOUT=1, i.e. -3.5% latency, distributions non-overlapping, both legs reporting no late compiles. Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
2 participants
Description
Warmup records every attention variant the bucketer enumerates, but nothing tells you at
runtime when that coverage falls short. #799 removed the log line that fired on a lazy
compile, and it only ever covered the page kernel in any case — the batched decode kernel
is never recorded, so under
SPYRE_BATCHED_DECODEit has no coverage guarantee at all.Both serving dispatch sites now read Dynamo's graph counter across the call and warn once
if it moved. It is armed at the end of attention recording, so the paths that deliberately
skip recording (
SPYRE_ATTN_RECORD=0,CompilationMode.NONE) keep compiling on first usewithout complaint.
It should stay silent on the default path, where the recorder covers every variant
dispatch can reach — if it ever fires there, that is the bug it exists to report. Today a
missed variant is invisible except as an unexplained latency spike.
Test Plan
pytest tests/attention/test_spyre_attn_recorder.py tests/runtime/test_spyre_attn_bucketer.py -m "not upstream"— 63 passed, 5 skipped.tests/attention/test_spyre_attn.py, subset covering both dispatch sites, on a card — 48 passed, 11 skipped.bash format.sh— all hooks pass exceptty, which reports the same diagnostics on this branch as on its base.Checklist
bash format.sh)Signed-off-by:line (DCO compliance)