refactor(attn): rename the bucketed decode kernel to batched decode - #791
Conversation
|
👋 Hi! Thank you for contributing. We also recommend installing prek and configuring it to check your code before every local commit. |
tdoublep
left a comment
tdoublep
left a comment
There was a problem hiding this comment.
This is Claude, reviewing on behalf of @tdoublep.
The rename is clean. I checked it is behaviour-free by mapping every case of bucketed and batched to the same token on both sides and diffing: nothing else changed. Tests still collect.
One thing to fix: the old name is still in six comments and docstrings. The grep in the test plan was case-sensitive, so it missed the ones that start a sentence, and grep -i finds them all. They sit outside the diff so I could not comment inline:
spyre_attn.py:437- the renamed function's own docstringspyre_attn.py:608and:1065- "Bucketed-decode precomputes"test_spyre_attn.py:1711,:1767,:1899- test docstrings
On the deprecation question in the description: no shim needed. SPYRE_BUCKETED_DECODE only landed a week ago in #701, there is no release with it, it is off by default, and nothing outside the code sets it.
Not a formal approval, leaving that to Thomas.
| # Batches below this fall back to the per-seq loop: the bucketed matmul's | ||
| # Batches below this fall back to the per-seq loop: the batched matmul's | ||
| # padded-row overhead exceeds the per-seq cost at small N. | ||
| _MIN_SEQS_BUCKET = 4 |
There was a problem hiding this comment.
Optional: _MIN_SEQS_BUCKET and the bucket_num_seqs / bucket_num_blocks metadata fields still say "bucket" for this kernel's dimensions, sitting right next to the other kernel's bucketing. That is the ambiguity the rename is removing, so it may be worth finishing here.
There was a problem hiding this comment.
Done!
bucket_num_seqs -> padded_num_seqs, bucket_num_blocks -> padded_batch_blocks, _MIN_SEQS_BUCKET -> _MIN_BATCHED_SEQS.
I agree that this is the ambiguity the rename exists to remove, and the deciding evidence was the sibling kernel: the per-sequence path is also bucketed and it names its ladder-rounded dimensions padded_num_blocks / padded_query_len. So "padded" is already the convention for a value rounded onto a ladder, and bucket_num_* was the outlier. The batched kernel even calls them b_seqs / b_blocks locally.
The asymmetry in the new names is deliberate: padded_num_blocks is already taken by the per-seq list[int] field, so the blocks field needs the batch_ qualifier to distinguish it while the seqs field doesn't. A symmetric batch_padded_* on both would have put the emphasis on the batching, when the reason these fields exist is the rounding.
_MIN_BATCHED_SEQS because the threshold is about whether a batch is large enough for this kernel, not a position on a ladder. The comment above it already said "Batches below this fall back to the per-seq loop." The value is compared against a bucket but isn't itself one.
Also fixed the six stale docstrings and comments from your other note. Verified: no old names remain, ty check passes (it resolves the dataclass field access, so a missed reference would fail), and the batched-decode device tests pass.
6563ca9 to
1de2c70
Compare
Signed-off-by: Jan van Lunteren <jvl@zurich.ibm.com>
1de2c70 to
d8dd7ef
Compare
|
@jvlunteren how should we proceed for the PRs merging? this one (#791) then you rebase #792 or you directly integrate your changes to #792 and we merge that one? |
|
Let's get this one in |
❌ merge-queue-integration: failurePlan (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 · 37b5e65c9f23"]
end
subgraph Lamd64_1["amd64 L1 · 1 parallel"]
n_amd64_hf_adapters_hf_adapters_dev["hf-adapters/hf-adapters-dev 🟢<br/>image · 41b5f97788b1"]
end
subgraph Lamd64_2["amd64 L2 · 1 parallel"]
n_amd64_spyre_inference_spyre_inference_dev["spyre-inference/spyre-inference-dev 🔴<br/>image · 2ec2df940a24"]
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 sOk;
class n_amd64_hf_adapters_hf_adapters_dev sOk;
class n_amd64_spyre_inference_spyre_inference_dev sOk;
✅ orch trigger-pr-validation — green · arches amd64 · fp amd64=0b7aa211
GHA test runs:
|
|
GHA test runs:
|
|
GHA test runs:
|
|
GHA test runs:
|
Takes main's encoder attention wholesale (torch-spyre#769 scatter-pack + per-step cache, torch-spyre#799 Dynamo-managed kernels) and keeps only the pooler-side compile fixes from this branch. The per-sequence encoder attention loop is deferred: torch-spyre#799 removed _maybe_compile and both kernel dicts it was built on, so it needs rewriting against the new module-level-kernel pattern before it can be measured against torch-spyre#769's scatter-pack. It stays available in this branch's history via the merge's first parent. Resolutions: - spyre_encoder_attn.py, its tests, envs.py, spyre_model_runner.py: took main. SPYRE_BUCKETED_ENCODE goes with the loop (and torch-spyre#791 renamed its decoder counterpart to SPYRE_BATCHED_DECODE, so the name was stale). - spyre_pooler.py: kept both new classes, main's SpyreMeanPool and this branch's SpyreDispatchPooler/SpyreTokenPooler. The conflicts were adjacency only. - SpyreMeanPool now crops on the host, after the D2H. Its crop only fires once SpyreDispatchPooler stops upstream's real-length slice, and on device it needed a real-length index_select -- one torch.compile specialization per distinct prompt length, the defect the dispatcher exists to remove. A host slice costs nothing; the D2H is a fixed, bucket-sized transfer either way. Signed-off-by: gkumbhat <Gaurav.Kumbhat@ibm.com>
3 participants
Description
Renames the multi-sequence decode kernel from "bucketed decode" to "batched decode", per review feedback on #772.
Both attention kernels are bucketed now, so "bucketed" no longer distinguishes them. What sets this one apart is the extra batch dimension: it attends for several sequences in one matmul rather than looping per sequence.
Pure rename, no behaviour change. Identifiers, the env var, test names, and the comments that referred to this kernel by the old name. Prose about bucket rounding is untouched where it means rounding rather than this kernel.
SPYRE_BUCKETED_DECODEbecomesSPYRE_BATCHED_DECODE. That is user-visible: anyone setting the old name would silently get the default instead. The flag is off by default so the application radius is small, but if a deprecation path is wanted the old name could be accepted with a warning.Related Issues
#771
Test Plan
pytest tests/attention/test_spyre_attn.py -m 'not upstream' -k 'not device_spyre': all passedChecklist
bash format.sh)Signed-off-by:line (DCO compliance)