Add linear attention API - #268
Open
jhjpark wants to merge 4 commits into
Open
Conversation
Anerudhan
added
the
cat-feature
Collaborator
|
Can you add documentation? PR description .. Also, seems like we are building the op out of pure Pytorch ops.. Is that the expectation? |
jhjpark
force-pushed
the
jopark/linear_attention
branch
from
June 29, 2026 22:34
7e8217a to
4d9bf15
Compare
Contributor
📝 WalkthroughWalkthroughAdds a new ChangesGated DeltaNet Linear Attention Op
Estimated code review effort: 4 (Complex) | ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
jhjpark
force-pushed
the
jopark/linear_attention
branch
from
June 29, 2026 22:37
4d9bf15 to
f595eda
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
python/cudnn/experimental/ops/linear_attention/gdn.py (1)
152-169: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the real runtime contract on
gated_delta_net.This docstring is the public API surface, but it currently omits the operational constraints that users will hit first: CUDA +
cuda.tileare required, and the only covered path in this PR is the CUDA/bf16-style path reflected by the tests. Please spell that out here instead of making users discover it at runtime. As per path instructions, focus on documentation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/experimental/ops/linear_attention/gdn.py` around lines 152 - 169, Update the `gated_delta_net` docstring to explicitly state the runtime contract: it requires CUDA and `cuda.tile`, and the supported path in this implementation is the CUDA/bf16-style execution covered by the existing tests. Keep the existing argument/return docs, but add a short “runtime requirements” note near the function description so users see the constraint before calling `gated_delta_net`.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/experimental/ops/__init__.py`:
- Around line 3-8: Add FE API coverage for the new public `gated_delta_net`
export by adding a matching case under the `test/python/fe_api` suite. Update
the FE API tests around the linear attention/GDN path to import and exercise
`gated_delta_net` similarly to how `test_linear_attention_gdn.py` validates the
operator, so the new symbol in `__all__` is covered end-to-end.
In `@python/cudnn/experimental/ops/linear_attention/_common.py`:
- Around line 109-117: The pad_to_multiple helper currently uses pad_spec math
that assumes a non-negative dim, so negative axes like dim=-1 build the wrong
padding tuple. Normalize dim to a canonical non-negative index at the start of
pad_to_multiple (or explicitly reject negative dims) before computing n and
pad_spec, and keep the existing torch.nn.functional.pad behavior unchanged.
In `@python/cudnn/experimental/ops/linear_attention/gdn.py`:
- Around line 51-55: The `chunk_size` argument in `gated_delta_net` is being
ignored, so both `_gdn_fwd` and `_gdn_fwd_fake` still behave as if the chunk
size were fixed at 64. Update the `gdn.py` implementation so `chunk_size` is
threaded through the real forward path and used when constructing the fake/meta
tensors in `_gdn_fwd_fake`, instead of hard-coding `_BT = 64`. Make sure the
helper functions and any related shape logic use the passed `chunk_size`
consistently so `gated_delta_net(..., chunk_size=32)` actually respects the
caller’s value.
In `@test/python/test_linear_attention_gdn.py`:
- Around line 120-121: The chained requires_grad_ calls in the linear attention
test are triggering Ruff E702, so split them into separate standalone
statements. Update the tensor setup in test_linear_attention_gdn so q, k, v, g,
and beta each call requires_grad_ on their own line, preserving the same
behavior while satisfying lint.
---
Nitpick comments:
In `@python/cudnn/experimental/ops/linear_attention/gdn.py`:
- Around line 152-169: Update the `gated_delta_net` docstring to explicitly
state the runtime contract: it requires CUDA and `cuda.tile`, and the supported
path in this implementation is the CUDA/bf16-style execution covered by the
existing tests. Keep the existing argument/return docs, but add a short “runtime
requirements” note near the function description so users see the constraint
before calling `gated_delta_net`.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 60603b18-6e9f-4e03-9ea4-900fe458c3cd
📒 Files selected for processing (6)
python/cudnn/experimental/ops/__init__.pypython/cudnn/experimental/ops/linear_attention/__init__.pypython/cudnn/experimental/ops/linear_attention/_common.pypython/cudnn/experimental/ops/linear_attention/_gdn_chunk_cutile.pypython/cudnn/experimental/ops/linear_attention/gdn.pytest/python/test_linear_attention_gdn.py
Comment on lines
+3
to
+8
| from .linear_attention import gated_delta_net | ||
|
|
||
| __all__ = [ | ||
| "scaled_dot_product_attention", | ||
| "moe_grouped_matmul", | ||
| "gated_delta_net", |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Looking for FE API coverage for gated_delta_net..."
if [ -d test/python/fe_api ]; then
rg -n 'gated_delta_net|linear_attention' test/python/fe_api || true
else
echo "test/python/fe_api is not present in this checkout"
fi
echo
echo "All matching tests under test/python:"
rg -n 'gated_delta_net|linear_attention' test/python || trueRepository: NVIDIA/cudnn-frontend
Length of output: 1643
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Top-level test tree:"
find test/python -maxdepth 2 -type d | sort
echo
echo "Files mentioning cudnn.experimental.ops.gated_delta_net:"
rg -n 'cudnn\.experimental\.ops\.gated_delta_net|from cudnn\.experimental\.ops import gated_delta_net|gated_delta_net\(' test/python || true
echo
echo "Check whether test/python/fe_api exists and what is inside:"
if [ -d test/python/fe_api ]; then
find test/python/fe_api -maxdepth 2 -type f | sort
else
echo "MISSING: test/python/fe_api"
fiRepository: NVIDIA/cudnn-frontend
Length of output: 3765
Add FE API coverage for gated_delta_net test/python/test_linear_attention_gdn.py covers the operator, but there’s no matching case under test/python/fe_api for this new public export.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/experimental/ops/__init__.py` around lines 3 - 8, Add FE API
coverage for the new public `gated_delta_net` export by adding a matching case
under the `test/python/fe_api` suite. Update the FE API tests around the linear
attention/GDN path to import and exercise `gated_delta_net` similarly to how
`test_linear_attention_gdn.py` validates the operator, so the new symbol in
`__all__` is covered end-to-end.
Source: Path instructions
Comment on lines
+109
to
+117
| def pad_to_multiple(t: torch.Tensor, multiple: int, dim: int, value: float = 0.0) -> Tuple[torch.Tensor, int]: | ||
| """Pad ``t`` along ``dim`` to the next multiple of ``multiple``.""" | ||
| n = t.shape[dim] | ||
| pad = (-n) % multiple | ||
| if pad == 0: | ||
| return t, 0 | ||
| # torch.nn.functional.pad operates on the last dims; convert to that form. | ||
| pad_spec = [0, 0] * (t.ndim - 1 - dim) + [0, pad] | ||
| return torch.nn.functional.pad(t, pad_spec, value=value), pad |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Normalize negative dim values before building pad_spec.
t.shape[dim] already accepts negative axes, but the pad_spec math does not. For example, dim=-1 produces the wrong pad tuple length/placement. Normalize first (dim %= t.ndim) or reject negative dims explicitly.
Suggested fix
def pad_to_multiple(t: torch.Tensor, multiple: int, dim: int, value: float = 0.0) -> Tuple[torch.Tensor, int]:
"""Pad ``t`` along ``dim`` to the next multiple of ``multiple``."""
+ dim %= t.ndim
n = t.shape[dim]
pad = (-n) % multiple
if pad == 0:
return t, 0📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def pad_to_multiple(t: torch.Tensor, multiple: int, dim: int, value: float = 0.0) -> Tuple[torch.Tensor, int]: | |
| """Pad ``t`` along ``dim`` to the next multiple of ``multiple``.""" | |
| n = t.shape[dim] | |
| pad = (-n) % multiple | |
| if pad == 0: | |
| return t, 0 | |
| # torch.nn.functional.pad operates on the last dims; convert to that form. | |
| pad_spec = [0, 0] * (t.ndim - 1 - dim) + [0, pad] | |
| return torch.nn.functional.pad(t, pad_spec, value=value), pad | |
| def pad_to_multiple(t: torch.Tensor, multiple: int, dim: int, value: float = 0.0) -> Tuple[torch.Tensor, int]: | |
| """Pad ``t`` along ``dim`` to the next multiple of ``multiple``.""" | |
| dim %= t.ndim | |
| n = t.shape[dim] | |
| pad = (-n) % multiple | |
| if pad == 0: | |
| return t, 0 | |
| # torch.nn.functional.pad operates on the last dims; convert to that form. | |
| pad_spec = [0, 0] * (t.ndim - 1 - dim) + [0, pad] | |
| return torch.nn.functional.pad(t, pad_spec, value=value), pad |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/experimental/ops/linear_attention/_common.py` around lines 109 -
117, The pad_to_multiple helper currently uses pad_spec math that assumes a
non-negative dim, so negative axes like dim=-1 build the wrong padding tuple.
Normalize dim to a canonical non-negative index at the start of pad_to_multiple
(or explicitly reject negative dims) before computing n and pad_spec, and keep
the existing torch.nn.functional.pad behavior unchanged.
Comment on lines
+51
to
+55
| scale: float, | ||
| chunk_size: int, | ||
| initial_state: Optional[torch.Tensor] = None, | ||
| output_final_state: bool = False, | ||
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Don't silently ignore chunk_size.
The API accepts chunk_size, but both execution paths are still hard-coded to _BT = 64: _gdn_fwd() never passes chunk_size downstream, and _gdn_fwd_fake() bakes 64 into A's shape. gated_delta_net(..., chunk_size=32) will therefore behave like 64 while the docstring suggests configurability.
Suggested fix
def _gdn_fwd(
@@
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ if chunk_size != _BT:
+ raise ValueError(f"Unsupported chunk_size={chunk_size}; expected {_BT}")
"""GDN forward.
@@
`@_gdn_fwd.register_fake`
def _gdn_fwd_fake(q, k, v, g, beta, scale, chunk_size, initial_state=None, output_final_state=False):
+ if chunk_size != _BT:
+ raise ValueError(f"Unsupported chunk_size={chunk_size}; expected {_BT}")
B, T, H, K = q.shape
@@
- chunk_size: chunk length (cuTile kernel uses 64).
+ chunk_size: chunk length. The current cuTile kernel requires ``64``.Also applies to: 79-88, 150-180
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/experimental/ops/linear_attention/gdn.py` around lines 51 - 55,
The `chunk_size` argument in `gated_delta_net` is being ignored, so both
`_gdn_fwd` and `_gdn_fwd_fake` still behave as if the chunk size were fixed at
64. Update the `gdn.py` implementation so `chunk_size` is threaded through the
real forward path and used when constructing the fake/meta tensors in
`_gdn_fwd_fake`, instead of hard-coding `_BT = 64`. Make sure the helper
functions and any related shape logic use the passed `chunk_size` consistently
so `gated_delta_net(..., chunk_size=32)` actually respects the caller’s value.
Source: Linters/SAST tools
jhjpark
force-pushed
the
jopark/linear_attention
branch
from
July 2, 2026 20:16
f595eda to
598da12
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
test/python/gdn/test_gdn_bprop.py (1)
134-178: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding a zero-length-sequence backward test.
test_gdn_fprop.pyhastest_fprop_zero_length_sequenceverifying empty segments in a varlen batch don't perturb others and their final state stays zero, but there's no analogous backward-pass test here. Gradients flowing through a zero-length segment (e.g., ensuring no NaN/garbage gradient leaks into adjacent segments) are untested.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/python/gdn/test_gdn_bprop.py` around lines 134 - 178, Add a backward-path regression test for zero-length varlen segments in test_gdn_bprop.py, mirroring the existing zero-length forward coverage in test_fprop_zero_length_sequence. Use the same _run_bprop_case and cu_seqlens-style setup as test_bprop_varlen_ragged, but include an empty segment in the sequence boundaries and assert the backward pass produces valid gradients without leaking NaN/garbage into neighboring segments. Reference the varlen/backward helpers and the zero-length-sequence case so the new test stays aligned with the existing bprop coverage.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/experimental/ops/linear_attention/gdn.py`:
- Around line 156-165: Update the docstring in gdn.py so the tensor shapes for
g, beta, initial_state, and final_state match the implementation and tests:
document g and beta as value-head dimensioned [B, T, HV] rather than [B, T, H],
and initial_state/final_state as [B, HV, K, V] rather than [B, H, K, V]. Also
note in the gdn op docs that the current implementation requires H == HV so the
public API contract is clear.
---
Nitpick comments:
In `@test/python/gdn/test_gdn_bprop.py`:
- Around line 134-178: Add a backward-path regression test for zero-length
varlen segments in test_gdn_bprop.py, mirroring the existing zero-length forward
coverage in test_fprop_zero_length_sequence. Use the same _run_bprop_case and
cu_seqlens-style setup as test_bprop_varlen_ragged, but include an empty segment
in the sequence boundaries and assert the backward pass produces valid gradients
without leaking NaN/garbage into neighboring segments. Reference the
varlen/backward helpers and the zero-length-sequence case so the new test stays
aligned with the existing bprop coverage.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9c1a7ce0-d127-40bb-ab30-475f2e6cbc83
📒 Files selected for processing (12)
python/cudnn/experimental/ops/__init__.pypython/cudnn/experimental/ops/linear_attention/__init__.pypython/cudnn/experimental/ops/linear_attention/_common.pypython/cudnn/experimental/ops/linear_attention/_gdn_chunk_cutile.pypython/cudnn/experimental/ops/linear_attention/gdn.pytest/python/gdn/__init__.pytest/python/gdn/common.pytest/python/gdn/conftest.pytest/python/gdn/reference_gdn.pytest/python/gdn/test_gdn_bprop.pytest/python/gdn/test_gdn_fprop.pytest/python/test_linear_attention_gdn.py
✅ Files skipped from review due to trivial changes (1)
- python/cudnn/experimental/ops/linear_attention/init.py
🚧 Files skipped from review as they are similar to previous changes (3)
- python/cudnn/experimental/ops/init.py
- test/python/test_linear_attention_gdn.py
- python/cudnn/experimental/ops/linear_attention/_common.py
Comment on lines
+156
to
+165
| g: log-space scalar decay per token, ``[B, T, H]`` | ||
| (``alpha = exp(g) in (0, 1]``). | ||
| beta: per-token write strength, ``[B, T, H]``. | ||
| scale: attention scale applied to ``q``. Defaults to ``1 / sqrt(K)``. | ||
| initial_state: optional recurrent state ``[B, H, K, V]`` (otherwise zero). | ||
| output_final_state: if ``True``, also return the state after the last token. | ||
|
|
||
| Returns: | ||
| ``(o, final_state)`` where ``o`` is ``[B, T, H, V]``. ``final_state`` is | ||
| ``[B, H, K, V]`` when ``output_final_state=True``, otherwise empty. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Docstring head dimensions for g/beta/state don't match the code and tests.
The code derives HV = beta.shape[2] and the tests build g/beta with HV and initial_state as [N, HV, K, V], but the docstring documents g/beta as [B, T, H] and initial_state/final_state as [B, H, K, V]. Since these gates and the recurrent state are value-head dimensioned, please document them as HV (and note the H == HV requirement of the current op) so the public API contract is unambiguous.
As per path instructions for python/cudnn/**: "Focus on documentation."
Suggested doc update
- g: log-space scalar decay per token, ``[B, T, H]``
+ g: log-space scalar decay per token, ``[B, T, HV]``
(``alpha = exp(g) in (0, 1]``).
- beta: per-token write strength, ``[B, T, H]``.
+ beta: per-token write strength, ``[B, T, HV]``.
scale: attention scale applied to ``q``. Defaults to ``1 / sqrt(K)``.
- initial_state: optional recurrent state ``[B, H, K, V]`` (otherwise zero).
+ initial_state: optional recurrent state ``[B, HV, K, V]`` (otherwise zero).
output_final_state: if ``True``, also return the state after the last token.
Returns:
``(o, final_state)`` where ``o`` is ``[B, T, H, V]``. ``final_state`` is
- ``[B, H, K, V]`` when ``output_final_state=True``, otherwise empty.
+ ``[B, HV, K, V]`` when ``output_final_state=True``, otherwise empty.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| g: log-space scalar decay per token, ``[B, T, H]`` | |
| (``alpha = exp(g) in (0, 1]``). | |
| beta: per-token write strength, ``[B, T, H]``. | |
| scale: attention scale applied to ``q``. Defaults to ``1 / sqrt(K)``. | |
| initial_state: optional recurrent state ``[B, H, K, V]`` (otherwise zero). | |
| output_final_state: if ``True``, also return the state after the last token. | |
| Returns: | |
| ``(o, final_state)`` where ``o`` is ``[B, T, H, V]``. ``final_state`` is | |
| ``[B, H, K, V]`` when ``output_final_state=True``, otherwise empty. | |
| g: log-space scalar decay per token, ``[B, T, HV]`` | |
| (``alpha = exp(g) in (0, 1]``). | |
| beta: per-token write strength, ``[B, T, HV]``. | |
| scale: attention scale applied to ``q``. Defaults to ``1 / sqrt(K)``. | |
| initial_state: optional recurrent state ``[B, HV, K, V]`` (otherwise zero). | |
| output_final_state: if ``True``, also return the state after the last token. | |
| Returns: | |
| ``(o, final_state)`` where ``o`` is ``[B, T, H, V]``. ``final_state`` is | |
| ``[B, HV, K, V]`` when ``output_final_state=True``, otherwise empty. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/experimental/ops/linear_attention/gdn.py` around lines 156 -
165, Update the docstring in gdn.py so the tensor shapes for g, beta,
initial_state, and final_state match the implementation and tests: document g
and beta as value-head dimensioned [B, T, HV] rather than [B, T, H], and
initial_state/final_state as [B, HV, K, V] rather than [B, H, K, V]. Also note
in the gdn op docs that the current implementation requires H == HV so the
public API contract is clear.
Source: Path instructions
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
gated_delta_netlinear-attention operator to the experimental ops package public namespace.initial_stateand can return afinal_statewhen requested.torch.compileparity.