add out parameter for dsa api - #470
Merged
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughThe indexer forward wrapper now accepts an optional ChangesIndexer forward output
Estimated code review effort: 1 (Trivial) | ~5 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![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: 3
🤖 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/deepseek_sparse_attention/indexer_forward/api.py`:
- Line 227: Update the docstring for the wrapper containing the out parameter to
state that output allocation occurs only when out is None; when provided, out
must have the exact expected shape, torch.float32 dtype, CUDA storage, and
contiguous layout.
- Line 227: Update the affected function signature in the indexer forward API so
the new optional out parameter is appended after q_causal_offsets, rather than
inserted before m_block_size. Preserve the existing positional order of all
tuning parameters and keep out optional.
- Line 270: Update the SM100 dispatch wrapper call to indexer_fwd to pass the
existing out buffer as out=out, ensuring the caller-provided output receives the
computed result.
🪄 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: e26a8c47-3ca9-4fbb-877a-42c6cbdaf455
📒 Files selected for processing (1)
python/cudnn/deepseek_sparse_attention/indexer_forward/api.py
| w: torch.Tensor, | ||
| ratio: int = 4, | ||
| qhead_per_kv_head: Optional[int] = None, | ||
| out: Optional[torch.Tensor] = None, |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document the out contract.
The docstring says that the wrapper allocates the output buffer, but out bypasses allocation. Document that allocation occurs only when out is None. Also document the SM90 requirements: exact output shape, torch.float32, CUDA storage, and contiguous layout.
🤖 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/deepseek_sparse_attention/indexer_forward/api.py` at line 227,
Update the docstring for the wrapper containing the out parameter to state that
output allocation occurs only when out is None; when provided, out must have the
exact expected shape, torch.float32 dtype, CUDA storage, and contiguous layout.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Preserve the existing positional parameter layout.
out is inserted before m_block_size. Existing positional callers now bind tuning parameters to the wrong arguments. Append out after q_causal_offsets so existing calls remain compatible.
Proposed signature change
q_causal_offsets: Optional[torch.Tensor] = None,
+ out: Optional[torch.Tensor] = None,🤖 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/deepseek_sparse_attention/indexer_forward/api.py` at line 227,
Update the affected function signature in the indexer forward API so the new
optional out parameter is appended after q_causal_offsets, rather than inserted
before m_block_size. Preserve the existing positional order of all tuning
parameters and keep out optional.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/deepseek_sparse_attention/indexer_forward/api.py`:
- Line 232: Move the out parameter in indexer_forward/api.py:232 and
indexer_forward/_interface_sm90.py:51 out of the middle of each callable
signature by making it keyword-only or appending it after q_causal_offsets in
api.py and after current_stream in _interface_sm90.py; preserve all existing
positional argument bindings.
🪄 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: 6ae61b93-2978-4eab-b939-e1952e0539cf
📒 Files selected for processing (2)
python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.pypython/cudnn/deepseek_sparse_attention/indexer_forward/api.py
| q_stage: int = 2, | ||
| kv_stage: int = 4, | ||
| sm_scale: float = 1.0, | ||
| out: Optional[torch.Tensor] = None, |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
for f in \
python/cudnn/deepseek_sparse_attention/indexer_forward/api.py \
python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py
do
echo "=== $f ==="
sed -n '205,255p' "$f" 2>/dev/null || true
sed -n '25,75p' "$f" 2>/dev/null || true
done
echo "=== definitions and call sites ==="
rg -n --glob '*.py' \
'def (indexer_forward|.*indexer.*forward)|IndexerForward|indexer_forward\(' \
python test 2>/dev/null | head -250Repository: NVIDIA/cudnn-frontend
Length of output: 8739
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
import ast
from pathlib import Path
files = [
Path("python/cudnn/deepseek_sparse_attention/indexer_forward/api.py"),
Path("python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py"),
]
for path in files:
print(f"=== {path} ===")
tree = ast.parse(path.read_text())
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
if node.name in {"__call__", "execute", "indexer_forward"} or "forward" in node.name:
args = node.args
positional = args.posonlyargs + args.args
defaults = [None] * (len(positional) - len(args.defaults)) + list(args.defaults)
print(node.name, node.lineno)
print(" positional:", [(a.arg, ast.unparse(d) if d else None) for a, d in zip(positional, defaults)])
print(" kwonly:", [a.arg for a in args.kwonlyargs])
PYRepository: NVIDIA/cudnn-frontend
Length of output: 870
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
import ast
from pathlib import Path
targets = [
Path("python/cudnn/deepseek_sparse_attention/indexer_forward/api.py"),
Path("python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py"),
]
for path in targets:
print(f"=== {path} ===")
text = path.read_text()
tree = ast.parse(text)
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
if node.lineno <= 260:
args = node.args
positional = args.posonlyargs + args.args
defaults = [None] * (len(positional) - len(args.defaults)) + list(args.defaults)
print(
node.name,
node.lineno,
"positional=",
[(a.arg, ast.unparse(d) if d else None) for a, d in zip(positional, defaults)],
"kwonly=",
[a.arg for a in args.kwonlyargs],
)
PY
echo "=== all references ==="
rg -n --glob '*.py' 'indexer_forward|IndexerForward|_interface_sm90' . | head -300Repository: NVIDIA/cudnn-frontend
Length of output: 7192
🏁 Script executed:
#!/bin/bash
set -eu
echo "=== wrapper dispatch and output handling ==="
sed -n '245,345p' python/cudnn/deepseek_sparse_attention/indexer_forward/api.py
echo "=== SM90 interface implementation ==="
sed -n '40,175p' python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py
echo "=== test call shapes ==="
sed -n '50,180p' test/python/fe_api/dsa/test_DSA_indexer_forward.pyRepository: NVIDIA/cudnn-frontend
Length of output: 12496
Do not insert out into the middle of either callable signature.
Make out keyword-only, or append it after q_causal_offsets in api.py and after current_stream in _interface_sm90.py. This preserves existing positional bindings.
📍 Affects 2 files
python/cudnn/deepseek_sparse_attention/indexer_forward/api.py#L232-L232(this comment)python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py#L51-L51
🤖 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/deepseek_sparse_attention/indexer_forward/api.py` at line 232,
Move the out parameter in indexer_forward/api.py:232 and
indexer_forward/_interface_sm90.py:51 out of the middle of each callable
signature by making it keyword-only or appending it after q_causal_offsets in
api.py and after current_stream in _interface_sm90.py; preserve all existing
positional argument bindings.
Anerudhan
self-requested a review
August 3, 2026 20:47
Anerudhan
approved these changes
Aug 3, 2026
Anerudhan
added
orig-external
Merged
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.
Before submitting
pre-commit runand committed any formatting changes.Affected area
Summary
Add
outparameter toindexer_forward_wrapperto enable tensor reuse across multiple calls.This PR allows users to pass a pre-allocated output tensor for the cuDNN DSA indexer forward pass, avoiding repeated internal tensor creation when
indexer_forward_wrapperis called iteratively.Key changes:
output_tensorparameter toindexer_forward_wrapperNone)Why
In Megatron-LM,
_indexer_topk_from_score_chunksmay callindexer_forward_wrapperin a loop during training. Without output tensor reuse, the cuDNN DSA indexer creates new output tensors on every call, leading to:This approach:
Summary by CodeRabbit