Give chunk_kda and recurrent_kda a string-enum impl selector - #318
Conversation
## Human Note
## Agent note
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel `naive_*` names, following the shape of gdn's existing mode/backend split and
FLA's ops-layer layout. `Impl` is a `str` enum ("fused" | "reference") shared from
`attn_gym.linear.impl`; both ops accept the enum or its string value, and unknown selectors list
the valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only `attn_gym.linear.kda.api`, which validates the shared
contract and imports the fused backends lazily inside the fused branch, so `chunk_kda` and
`recurrent_kda` are now eager exports and the PEP 562 lazy set shrinks to the standalone kernels.
`impl="reference"` runs the naive oracles in FP32 behind the identical packed contract (empty
padding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only: `fastmath` with the reference raises. The
`naive_*` functions leave the public interface and remain importable as test oracles from
`attn_gym.linear.kda.naive`.
## Test Plan
```bash
gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_impl_dispatch.py \
test/test_kda_recurrent.py test/test_kda.py test/test_kda_cute_forward.py \
test/test_kda_ragged_public_backward.py test/test_kda_training_example.py -q
# 86 passed on B200. New coverage: enum/string coercion and rejection, reference-vs-fused
# agreement for both ops (dense and packed with empty slots and capacity tails), reference
# differentiability plus fused gradient rejection, non-128 head dims on the reference, and
# fastmath rejection. The one failure, test_chunk_kda_single_sequence_metadata_matches_dense_path,
# reproduces on pristine origin/main (16c2563) on this machine and is unrelated.
```
stack-info: PR: #318, branch: drisspg/stack/80
abdb244 to
816bac6
Compare
meta-cla
Bot
added
the
CLA Signed
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 816bac6633
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| """ | ||
| match resolve_impl(impl): | ||
| case Impl.FUSED: | ||
| from attn_gym.linear.kda.fwd.cute.chunk_kda_fwd import chunk_kda as fused_chunk_kda |
There was a problem hiding this comment.
When a fresh process uses torch.compile(chunk_kda, fullgraph=True) before importing the fused backend, this branch imports chunk_kda_fwd and performs its torch.library registrations while Dynamo is already tracing the public function. That cold path can fail strict capture or depend on tracing-time side effects; the existing compile tests avoid it by importing the backend module first. Ensure registration occurs before graph capture while retaining lazy optional dependencies, and cover a cold public-operation compile.
AGENTS.md reference: AGENTS.md:L69-L73
Useful? React with 👍 / 👎.
| from attn_gym.linear.kda import ( | ||
| active_token_mask, | ||
| chunk_kda, | ||
| mask_inactive_token_gradients, | ||
| mask_inactive_tokens, | ||
| naive_chunk_kda, | ||
| naive_chunk_kda_from_cumulative, | ||
| naive_recurrent_kda, | ||
| recurrent_kda, | ||
| ) |
There was a problem hiding this comment.
Removing these eager exports leaves docs/linear.md lines 222–226 with three mkdocstrings directives for attn_gym.linear.naive_chunk_kda, naive_chunk_kda_from_cumulative, and naive_recurrent_kda, none of which now exists at that path. The generated linear API page therefore contains unresolved entries and still describes the superseded interface; update it to document the selector-based public functions or the new private oracle paths.
AGENTS.md reference: AGENTS.md:L57-L62
Useful? React with 👍 / 👎.
## Human Note
## Agent note
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel `naive_*` names, following the shape of gdn's existing mode/backend split and
FLA's ops-layer layout. `Impl` is a `str` enum ("fused" | "reference") shared from
`attn_gym.linear.impl`; both ops accept the enum or its string value, and unknown selectors list
the valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only `attn_gym.linear.kda.api`, which validates the shared
contract and imports the fused backends lazily inside the fused branch, so `chunk_kda` and
`recurrent_kda` are now eager exports and the PEP 562 lazy set shrinks to the standalone kernels.
`impl="reference"` runs the naive oracles in FP32 behind the identical packed contract (empty
padding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only: `fastmath` with the reference raises. The
`naive_*` functions leave the public interface and remain importable as test oracles from
`attn_gym.linear.kda.naive`.
## Test Plan
```bash
gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_impl_dispatch.py \
test/test_kda_recurrent.py test/test_kda.py test/test_kda_cute_forward.py \
test/test_kda_ragged_public_backward.py test/test_kda_training_example.py -q
# 86 passed on B200. New coverage: enum/string coercion and rejection, reference-vs-fused
# agreement for both ops (dense and packed with empty slots and capacity tails), reference
# differentiability plus fused gradient rejection, non-128 head dims on the reference, and
# fastmath rejection. The one failure, test_chunk_kda_single_sequence_metadata_matches_dense_path,
# reproduces on pristine origin/main (16c2563) on this machine and is unrelated.
```
stack-info: PR: #318, branch: drisspg/stack/80
816bac6 to
a9a5541
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a9a55411d4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| num_sequences = cu_seqlens.shape[0] - 1 | ||
| output = torch.zeros_like(v) | ||
| final_state = ( | ||
| torch.zeros(num_sequences, heads, key_dim, value_dim, device=q.device) |
There was a problem hiding this comment.
When the caller has changed PyTorch's default dtype (for example with torch.set_default_dtype(torch.float64)) and omits initial_state, this allocation produces an FP64 final state even though the public contract promises FP32 states. Besides returning the wrong dtype, that state cannot be passed into the next segmented call because validate_kda_inputs rejects FP64 inputs; specify dtype=torch.float32 or derive the allocation from the already-normalized FP32 query.
Useful? React with 👍 / 👎.
| if initial_state is None | ||
| else initial_state.clone() | ||
| ) | ||
| for sequence, (start, end) in enumerate(pairwise(cu_seqlens.cpu().tolist())): |
There was a problem hiding this comment.
When either public KDA operation is compiled with fullgraph=True, impl="reference", and packed cu_seqlens, converting the offsets through .cpu().tolist() extracts tensor data into Python and prevents Dynamo from capturing this implementation as one graph. The repository's custom-operator workflow requires strict compilation coverage for each materially different public implementation, so this packed dispatch needs compiler-visible tensor operations or an explicitly narrowed compile contract.
AGENTS.md reference: AGENTS.md:L69-L73
Useful? React with 👍 / 👎.
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel `naive_*` names, following the shape of gdn's existing mode/backend split and
FLA's ops-layer layout. `Impl` is a `str` enum ("fused" | "reference") shared from
`attn_gym.linear.impl`; both ops accept the enum or its string value, and unknown selectors list
the valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only `attn_gym.linear.kda.api`, which validates the shared
contract and imports the fused backends lazily inside the fused branch, so `chunk_kda` and
`recurrent_kda` are now eager exports and the PEP 562 lazy set shrinks to the standalone kernels.
`impl="reference"` runs the naive oracles in FP32 behind the identical packed contract (empty
padding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only: `fastmath` with the reference raises. The
`naive_*` functions leave the public interface and remain importable as test oracles from
`attn_gym.linear.kda.naive`.
```bash
gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_impl_dispatch.py \
test/test_kda_recurrent.py test/test_kda.py test/test_kda_cute_forward.py \
test/test_kda_ragged_public_backward.py test/test_kda_training_example.py -q
```
stack-info: PR: #318, branch: drisspg/stack/80
a9a5541 to
f8668b3
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f8668b30da
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if initial_state is None | ||
| else initial_state.clone() | ||
| ) | ||
| for sequence, (start, end) in enumerate(pairwise(cu_seqlens.cpu().tolist())): |
There was a problem hiding this comment.
When impl="reference" receives malformed offset values—such as a nonzero first offset, decreasing boundaries, or an endpoint beyond T—Python slicing silently clamps or overlaps spans and returns incorrect outputs/states. The fused path rejects these cases in chunk_scheduler.py via its begin >= 0, begin <= end, end <= tokens, and initial-zero assertions, so validate the host list against the same invariants before iterating to preserve the shared public contract.
Useful? React with 👍 / 👎.
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel `naive_*` names, following the shape of gdn's existing mode/backend split and
FLA's ops-layer layout. `Impl` is a `str` enum ("fused" | "reference") shared from
`attn_gym.linear.impl`; both ops accept the enum or its string value, and unknown selectors list
the valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only `attn_gym.linear.kda.api`, which validates the shared
contract and imports the fused backends lazily inside the fused branch, so `chunk_kda` and
`recurrent_kda` are now eager exports and the PEP 562 lazy set shrinks to the standalone kernels.
`impl="reference"` runs the naive oracles in FP32 behind the identical packed contract (empty
padding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only: `fastmath` with the reference raises. The
`naive_*` functions leave the public interface and remain importable as test oracles from
`attn_gym.linear.kda.naive`.
```bash
gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_impl_dispatch.py \
test/test_kda_recurrent.py test/test_kda.py test/test_kda_cute_forward.py \
test/test_kda_ragged_public_backward.py test/test_kda_training_example.py -q
```
stack-info: PR: #318, branch: drisspg/stack/80
f8668b3 to
5d12db5
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5d12db55d2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| applies only to the fused implementation. | ||
| """ | ||
| if resolve_impl(impl) is Impl.FUSED: | ||
| from attn_gym.linear.kda.fwd.cute.chunk_kda_fwd import chunk_kda as fused_chunk_kda |
There was a problem hiding this comment.
When a base installation without the linear extra calls chunk_kda(..., impl="fused"), this direct import exposes a low-level ModuleNotFoundError for cutlass instead of the actionable pip install attn-gym[linear] error previously provided by attn_gym.linear.kda.__getattr__. Wrap the lazy backend import so the documented optional-dependency behavior is retained.
Useful? React with 👍 / 👎.
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel `naive_*` names, following the shape of gdn's existing mode/backend split and
FLA's ops-layer layout. `Impl` is a `str` enum ("fused" | "reference") shared from
`attn_gym.linear.impl`; both ops accept the enum or its string value, and unknown selectors list
the valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only `attn_gym.linear.kda.api`, which validates the shared
contract and imports the fused backends lazily inside the fused branch, so `chunk_kda` and
`recurrent_kda` are now eager exports and the PEP 562 lazy set shrinks to the standalone kernels.
`impl="reference"` runs the naive oracles in FP32 behind the identical packed contract (empty
padding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only: `fastmath` with the reference raises. The
`naive_*` functions leave the public interface and remain importable as test oracles from
`attn_gym.linear.kda.naive`.
```bash
gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_impl_dispatch.py \
test/test_kda_recurrent.py test/test_kda.py test/test_kda_cute_forward.py \
test/test_kda_ragged_public_backward.py test/test_kda_training_example.py -q
```
stack-info: PR: #318, branch: drisspg/stack/80
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel `naive_*` names, following the shape of gdn's existing mode/backend split and
FLA's ops-layer layout. `Impl` is a `str` enum ("fused" | "reference") shared from
`attn_gym.linear.impl`; both ops accept the enum or its string value, and unknown selectors list
the valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only `attn_gym.linear.kda.api`, which now owns the shared
contract for *both* implementations: `validation.py` runs once at the API boundary, the fused
modules keep only their own constraints behind private `forward` entry points (no second
public-looking `chunk_kda` in `fwd/cute`), and the fused backends import lazily inside the fused
branch, so `chunk_kda` and `recurrent_kda` are eager exports and the PEP 562 lazy set shrinks to
the standalone kernels. `docs/linear.md` documents the new surface (the removed `naive_*` anchors
would have failed `mkdocs build --strict`) including the serving limitations that are deliberate
for now: dense state slots, out-of-place final state, unfused decode-step launches, and no
spec-decode rollback. `naive_recurrent_kda`'s own `cu_seqlens` branch is removed: it had no
remaining callers and narrower packed semantics than the public contract, whose reference
statement now lives solely in `api._packed_reference` (with the documented offset checks enforced
at its existing host read).
`impl="reference"` runs the naive oracles in FP32 behind the identical packed contract (empty
padding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only: `fastmath` with the reference raises. The
`naive_*` functions leave the public interface and remain importable as test oracles from
`attn_gym.linear.kda.naive`.
```bash
gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_impl_dispatch.py \
test/test_kda_recurrent.py test/test_kda.py test/test_kda_cute_forward.py \
test/test_kda_ragged_public_backward.py test/test_kda_training_example.py -q
```
stack-info: PR: #318, branch: drisspg/stack/80
154a084 to
8c573b3
Compare
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel `naive_*` names, following the shape of gdn's existing mode/backend split and
FLA's ops-layer layout. `Impl` is a `str` enum ("fused" | "reference") shared from
`attn_gym.linear.impl`; both ops accept the enum or its string value, and unknown selectors list
the valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only `attn_gym.linear.kda.api`, which now owns the shared
contract for *both* implementations: `validation.py` runs once at the API boundary, the fused
modules keep only their own constraints behind private `forward` entry points (no second
public-looking `chunk_kda` in `fwd/cute`), and the fused backends import lazily inside the fused
branch, so `chunk_kda` and `recurrent_kda` are eager exports and the PEP 562 lazy set shrinks to
the standalone kernels. `docs/linear.md` documents the new surface (the removed `naive_*` anchors
would have failed `mkdocs build --strict`) including the serving limitations that are deliberate
for now: dense state slots, out-of-place final state, unfused decode-step launches, and no
spec-decode rollback. `api._packed_reference` enforces the documented
offset contract at its existing host read.
`impl="reference"` runs the naive oracles in FP32 behind the identical packed contract (empty
padding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only: `fastmath` with the reference raises. The
`naive_*` functions leave the public interface and remain importable as test oracles from
`attn_gym.linear.kda.naive`.
```bash
gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_impl_dispatch.py \
test/test_kda_recurrent.py test/test_kda.py test/test_kda_cute_forward.py \
test/test_kda_ragged_public_backward.py test/test_kda_training_example.py -q
```
stack-info: PR: #318, branch: drisspg/stack/80
8c573b3 to
d481888
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d481888ed7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| q, k, v = (tensor.float() for tensor in (q, k, v)) | ||
| gate = gate.float() | ||
| beta = beta.float() |
There was a problem hiding this comment.
When impl="reference" is called inside a CUDA autocast region, converting the operands with .float() does not force subsequent autocast-eligible operations to remain in FP32; the einsum and matrix multiplications in both naive implementations can still execute in BF16/FP16. This violates the newly documented FP32-oracle contract and can make reference training or fused-backend comparisons substantially less accurate under the common AMP context, so run the dense reference operation with autocast disabled.
Useful? React with 👍 / 👎.
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel `naive_*` names, following the shape of gdn's existing mode/backend split and
FLA's ops-layer layout. `Impl` is a `str` enum ("fused" | "reference") shared from
`attn_gym.linear.impl`; both ops accept the enum or its string value, and unknown selectors list
the valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only `attn_gym.linear.kda.api`, which now owns the shared
contract for *both* implementations: `validation.py` runs once at the API boundary, the fused
modules keep only their own constraints behind private `forward` entry points (no second
public-looking `chunk_kda` in `fwd/cute`), and the fused backends import lazily inside the fused
branch, so `chunk_kda` and `recurrent_kda` are eager exports and the PEP 562 lazy set shrinks to
the standalone kernels. `docs/linear.md` documents the new surface (the removed `naive_*` anchors
would have failed `mkdocs build --strict`) including the serving limitations that are deliberate
for now: dense state slots, out-of-place final state, unfused decode-step launches, and no
spec-decode rollback. `api._packed_reference` enforces the documented
offset contract at its existing host read.
`impl="reference"` runs the naive oracles in FP32 behind the identical packed contract (empty
padding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only: `fastmath` with the reference raises. The
`naive_*` functions leave the public interface and remain importable as test oracles from
`attn_gym.linear.kda.naive`.
```bash
gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_impl_dispatch.py \
test/test_kda_recurrent.py test/test_kda.py test/test_kda_cute_forward.py \
test/test_kda_ragged_public_backward.py test/test_kda_training_example.py -q
```
stack-info: PR: #318, branch: drisspg/stack/80
d481888 to
6b6bcd1
Compare
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel `naive_*` names, following the shape of gdn's existing mode/backend split and
FLA's ops-layer layout. `Impl` is a `str` enum ("fused" | "reference") shared from
`attn_gym.linear.impl`; both ops accept the enum or its string value, and unknown selectors list
the valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only `attn_gym.linear.kda.api`, which now owns the shared
contract for *both* implementations: `validation.py` runs once at the API boundary, the fused
modules keep only their own constraints behind private `forward` entry points (no second
public-looking `chunk_kda` in `fwd/cute`), and the fused backends import lazily inside the fused
branch, so `chunk_kda` and `recurrent_kda` are eager exports and the PEP 562 lazy set shrinks to
the standalone kernels. `docs/linear.md` documents the new surface (the removed `naive_*` anchors
would have failed `mkdocs build --strict`) including the serving limitations that are deliberate
for now: dense state slots, out-of-place final state, unfused decode-step launches, and no
spec-decode rollback. `api._packed_reference` enforces the documented
offset contract at its existing host read.
`impl="reference"` runs the naive oracles in FP32 behind the identical packed contract (empty
padding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only: `fastmath` with the reference raises. The
`naive_*` functions leave the public interface and remain importable as test oracles from
`attn_gym.linear.kda.naive`.
```bash
gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_impl_dispatch.py \
test/test_kda_recurrent.py test/test_kda.py test/test_kda_cute_forward.py \
test/test_kda_ragged_public_backward.py test/test_kda_training_example.py -q
```
stack-info: PR: #318, branch: drisspg/stack/80
6b6bcd1 to
b51eb06
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0ad98988b8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| metadata: RaggedChunkMetadata | None, | ||
| *, | ||
| output_final_state: bool, | ||
| tune: bool, |
There was a problem hiding this comment.
When chunk_kda(..., tune=False) executes the fused forward path, this flag reaches the recurrence and output-composition stages but not chunk_kda_fwd_intra; that function calls recompute_w_u_fwd without tune, so its autotuner still benchmarks and selects a configuration using the default True. Consequently the advertised fixed-configuration behavior can still vary across machines or cold caches. Pass tune through chunk_kda_fwd_intra to its recompute call.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5a96c3650e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| tuple( | ||
| _canonicalize(kernel.compile_call(candidate, *runtime_args)) | ||
| for candidate in candidates | ||
| ), |
There was a problem hiding this comment.
When a process runs chunk_kda backward at multiple sequence lengths, ChunkKdaBwdWyDqkgTunable.compile_call produces the same values for every length because it includes only heads, dimensions, flags, and the candidate's grid waves. Hashing only those compile calls here therefore reuses the first length's persisted winner without benchmarking later workloads, even though the persistent-grid winner depends on the runtime token and chunk counts. This makes the default autotuning result shape-order-dependent and potentially substantially suboptimal; include performance-relevant runtime dimensions in the winner key rather than only compile-specialization arguments.
AGENTS.md reference: AGENTS.md:L74-L75
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6e1dc9593a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| winner_key = _winner_key(kernel, candidates, runtime_args) | ||
| selected = _load_winner(winner_key) |
There was a problem hiding this comment.
When callers use the documented benchmark= override, the winner key remains identical to one produced with the default benchmark or a different callback. A prior invocation can therefore cause this call to skip its requested timing policy entirely and reuse a configuration selected under different criteria, producing an incorrect tuning result. Include the benchmark policy in the cache identity, or bypass persistent winner reuse whenever a custom benchmark is supplied.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: aee6e958a6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ) | ||
| except TypeError: | ||
| fast_key = None | ||
| selected = _WINNERS_FAST.get(fast_key) if fast_key is not None else None |
There was a problem hiding this comment.
When a valid compile_call() tuple contains a pickleable but unhashable static value such as a list or dict, constructing fast_key succeeds because tuple construction does not hash its members, but _WINNERS_FAST.get(fast_key) then raises TypeError. This prevents run_tunable() from reaching the canonicalized persistent-key path even though jit_cache supports these static values; perform the lookup inside the try or canonicalize the fast key first.
AGENTS.md reference: AGENTS.md:L74-L75
Useful? React with 👍 / 👎.
| winner = pickle.loads(path.read_bytes()) | ||
| except (OSError, pickle.PickleError): |
There was a problem hiding this comment.
When a persisted winner file is empty or truncated, pickle.loads() raises EOFError, which is not a pickle.PickleError and therefore escapes this handler. A damaged best-effort cache entry consequently blocks every matching autotuned call until the file is manually removed instead of re-tuning as intended; catch EOFError alongside the existing cache-read failures.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d9f81e8eda
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| import torch | ||
|
|
||
| from attn_gym.linear.kda.fwd.triton.recurrent import forward as _fused_recurrent_forward |
There was a problem hiding this comment.
On CPU-only or macOS PyTorch installations where triton is unavailable, importing attn_gym.linear now fails here before impl="reference" can be selected. This makes the newly advertised hardware-independent reference implementation unusable in exactly those environments; import the Triton recurrent backend only inside the fused dispatch branch, and keep the reference-facing public import path free of fused-only dependencies.
Useful? React with 👍 / 👎.
| tuple( | ||
| _canonicalize(kernel.compile_call(candidate, *runtime_args)) | ||
| for candidate in candidates | ||
| ), |
There was a problem hiding this comment.
When the same process trains multiple sequence lengths, this key can reuse a winner selected for an unrelated workload because it contains only compile_call results. In particular, ChunkKdaBwdWyDqkgTunable.compile_call omits token/chunk capacity even though its grid_waves candidates control a runtime grid whose best choice depends on that capacity, so the first tuned shape permanently determines later shapes' schedule. Include performance-relevant runtime dimensions through a kernel-supplied tuning key rather than assuming compile-specialization arguments fully identify a tuning decision.
Useful? React with 👍 / 👎.
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel `naive_*` names, following the shape of gdn's existing mode/backend split and
FLA's ops-layer layout. `Impl` is a `str` enum ("fused" | "reference") shared from
`attn_gym.linear.impl`; both ops accept the enum or its string value, and unknown selectors list
the valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only `attn_gym.linear.kda.api`, which now owns the shared
contract for *both* implementations: `validation.py` runs once at the API boundary, the fused
modules keep only their own constraints behind private `forward` entry points (no second
public-looking `chunk_kda` in `fwd/cute`), and the fused backends import lazily inside the fused
branch, so `chunk_kda` and `recurrent_kda` are eager exports and the PEP 562 lazy set shrinks to
the standalone kernels. `docs/linear.md` documents the new surface (the removed `naive_*` anchors
would have failed `mkdocs build --strict`) including the serving limitations that are deliberate
for now: dense state slots, out-of-place final state, unfused decode-step launches, and no
spec-decode rollback. `api._packed_reference` enforces the documented
offset contract at its existing host read.
`impl="reference"` runs the naive oracles in FP32 behind the identical packed contract (empty
padding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only: `fastmath` with the reference raises. The
`naive_*` functions leave the public interface and remain importable as test oracles from
`attn_gym.linear.kda.naive`.
A perf-focused review round (measured on B200, warm launch-bound B=1/T=64/H=1/D=128, CPU enqueue
medians per the checked-in probe) then removed ~190us/step from fwd+bwd and ~72us from forward:
profiler ranges now compile to nullcontext unless a torch profiler is active, compile-target
detection is memoized per device, the warm winner path is O(1) (one default-config compile call
keys the memo; candidates materialize only on a miss), and `autotune=False` now genuinely reaches
every stage (the forward recompute and the ragged dAv fallback previously ignored it). Residual
warm `autotune=True` overhead (~70us/step) is Triton's own per-call autotuner dispatch;
`autotune=False` avoids it entirely.
## Test Plan
```bash
gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_impl_dispatch.py \
test/test_kda_recurrent.py test/test_kda.py test/test_kda_cute_forward.py \
test/test_kda_ragged_public_backward.py test/test_kda_training_example.py -q
```
stack-info: PR: #318, branch: drisspg/stack/80
1 participant
Give chunk_kda and recurrent_kda a string-enum impl selector
Human Note
Agent note
The two algorithmic forms stay separate functions because their gate contracts differ
(chunk-local cumulative vs per-token log2 decay), but naive-versus-fused is now an argument
instead of parallel
naive_*names, following the shape of gdn's existing mode/backend split andFLA's ops-layer layout.
Implis astrenum ("fused" | "reference") shared fromattn_gym.linear.impl; both ops accept the enum or its string value, and unknown selectors listthe valid set. There is deliberately no "auto" fallback: a training run should never silently land
on the eager reference.
The public wrappers move to torch-only
attn_gym.linear.kda.api, which validates the sharedcontract and imports the fused backends lazily inside the fused branch, so
chunk_kdaandrecurrent_kdaare now eager exports and the PEP 562 lazy set shrinks to the standalone kernels.impl="reference"runs the naive oracles in FP32 behind the identical packed contract (emptypadding slots pass state through, capacity tails stay out of contract, states come back FP32 with
one entry per logical sequence); it lifts the fused constraints — any hardware, any head dimension,
and a differentiable recurrent path where the fused scan is inference-only — at the cost of a host
read of the offsets. Fused-only knobs stay fused-only:
fastmathwith the reference raises. Thenaive_*functions leave the public interface and remain importable as test oracles fromattn_gym.linear.kda.naive.Test Plan