Skip to content

add out parameter for dsa api #470

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Anerudhan merged 2 commits into from
Aug 3, 2026
Merged

add out parameter for dsa api #470

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9852848
add out parameter for dsa api
terminator123 Aug 3, 2026
53091b2
add out paramter for dsa indexer forward
terminator123 Aug 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,8 @@ def indexer_fwd(
w: torch.Tensor,
ratio: int = 4,
qhead_per_kv_head: Optional[int] = None,
out: Optional[torch.Tensor] = None,
sm_scale: float = 1.0,
out: Optional[torch.Tensor] = None,
cu_seqlens_q: Optional[torch.Tensor] = None,
cu_seqlens_k: Optional[torch.Tensor] = None,
max_seqlen_q: Optional[int] = None,
Expand Down
3 changes: 3 additions & 0 deletions python/cudnn/deepseek_sparse_attention/indexer_forward/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,7 @@ def indexer_forward_wrapper(
q_stage: int = 2,
kv_stage: int = 4,
sm_scale: float = 1.0,
out: Optional[torch.Tensor] = None,

@coderabbitai coderabbitai Bot Aug 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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 -250

Repository: 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])
PY

Repository: 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 -300

Repository: 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.py

Repository: 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.

stream: Optional[cuda.CUstream] = None,
cu_seqlens_q: Optional[torch.Tensor] = None,
cu_seqlens_k: Optional[torch.Tensor] = None,
Expand Down Expand Up @@ -267,6 +268,7 @@ def indexer_forward_wrapper(
ratio=ratio,
qhead_per_kv_head=qhead_per_kv_head,
sm_scale=sm_scale,
out=out,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
max_seqlen_q=max_seqlen_q,
Expand All @@ -286,6 +288,7 @@ def indexer_forward_wrapper(
w,
ratio=ratio,
qhead_per_kv_head=qhead_per_kv_head,
out=out,
m_block_size=m_block_size,
n_block_size=n_block_size,
num_threads=384,
Expand Down