Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
25 changes: 0 additions & 25 deletions test/python/test_mhas_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,8 +116,6 @@ def test_sdpa_random_fwd_L0(env_info, test_no, request, cudnn_handle):
with_sliding_mask=SlidingWindowMaskGenerator(causal=10, left_window_only=5, right_window_only=5, band_around_diag=10, no_mask=10),
diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}),
is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "full" : 1}),
with_score_max=RandomChoice({True : 1, False : 3}),
with_score_sum_exp=RandomChoice({True : 1, False : 3}),
with_sink_token=RandomChoice({True : 1, False : 3}),
) as randomization_ctx:
test.cfg = randomization_ctx(rng, data_seed, geom_seed)
Expand Down Expand Up @@ -148,18 +146,13 @@ def test_sdpa_random_fwd_unified_L1(env_info, test_no, request, cudnn_handle):
with_sliding_mask=SlidingWindowMaskGenerator(causal=10, left_window_only=5, right_window_only=5, band_around_diag=10, no_mask=10),
diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}),
is_bias=RandomChoice({True : 1, False : 3}),
is_alibi=RandomChoice({True : 1, False : 3}),
is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "cu_padded" : 1, "full" : 1}),
with_unfuse_fma=RandomChoice({True : 1, False : 1}), # Randomly enable unfuse_fma for SM100
with_score_max=RandomChoice({True : 1, False : 3}),
with_score_sum_exp=RandomChoice({True : 1, False : 3}),
with_sink_token=RandomChoice({True : 1, False : 3}),
is_dropout=RandomChoice({True : 1, False : 3}),
with_rope=RandomChoice({True : 1, False : 3}), # RoPE at end to preserve existing test distributions
) as randomization_ctx:
test.cfg = randomization_ctx(rng, data_seed, geom_seed)

test.cfg.dropout_prob = 0.1 if test.cfg.is_dropout else 0.0
test.cfg.implementation = getattr(cudnn.attention_implementation, request.config.getoption("--implementation") or "", cudnn.attention_implementation.UNIFIED)
test.showConfig(test_no, request)

Expand Down Expand Up @@ -235,8 +228,6 @@ def test_sdpa_random_sq1_L0(env_info, test_no, request, cudnn_handle):
with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10),
diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}),
is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 0, "full" : 1}),
with_score_max=RandomChoice({True : 1, False : 3}),
with_score_sum_exp=RandomChoice({True : 1, False : 3}),
# sink_token not supported with s_q==1
# dropout not supported with s_q==1
) as randomization_ctx:
Expand Down Expand Up @@ -268,8 +259,6 @@ def test_sdpa_random_sq1_unified_L1(env_info, test_no, request, cudnn_handle):
with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10),
diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 0}), # Modified from non-unified test
is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 0, "full" : 1}),
with_score_max=RandomChoice({True : 1, False : 3}),
with_score_sum_exp=RandomChoice({True : 1, False : 3}),
# sink_token not supported with s_q==1
# dropout not supported with s_q==1
) as randomization_ctx:
Expand Down Expand Up @@ -306,8 +295,6 @@ def test_sdpa_random_lean_attn_L0(env_info, test_no, request, cudnn_handle):
with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10),
diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}),
is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "full" : 1}),
with_score_max=RandomChoice({True : 1, False : 3}),
with_score_sum_exp=RandomChoice({True : 1, False : 3}),
# sink_token not supported with s_q==1
# dropout not supported with s_q==1
) as randomization_ctx:
Expand Down Expand Up @@ -339,8 +326,6 @@ def test_sdpa_random_lean_attn_unified_L1(env_info, test_no, request, cudnn_hand
with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10),
diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 0}), # Modified from non-unified test
is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "full" : 1}),
with_score_max=RandomChoice({True : 1, False : 3}),
with_score_sum_exp=RandomChoice({True : 1, False : 3}),
# sink_token not supported with s_q==1
# dropout not supported with s_q==1
) as randomization_ctx:
Expand Down Expand Up @@ -376,8 +361,6 @@ def test_sdpa_random_fwd_ragged_L0(env_info, test_no, request, cudnn_handle):
with_sliding_mask=SlidingWindowMaskGenerator(causal=10, left_window_only=5, right_window_only=5, band_around_diag=10, no_mask=10),
diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}),
is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "padded" : 0, "full" : 0}),
with_score_max=RandomChoice({True : 1, False : 3}),
with_score_sum_exp=RandomChoice({True : 1, False : 3}),
with_sink_token=RandomChoice({True : 1, False : 3}),
) as randomization_ctx:
test.cfg = randomization_ctx(rng, data_seed, geom_seed)
Expand Down Expand Up @@ -408,14 +391,10 @@ def test_sdpa_random_fwd_ragged_unified_L1(env_info, test_no, request, cudnn_han
with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10), # Modified from non-unified test
diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 0}), # Modified from non-unified test
is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "cu_ragged" : 1, "padded" : 0, "full" : 0}),
with_score_max=RandomChoice({True : 1, False : 3}),
with_score_sum_exp=RandomChoice({True : 1, False : 3}),
with_sink_token=RandomChoice({True : 1, False : 3}),
is_dropout=RandomChoice({True : 1, False : 3}),
) as randomization_ctx:
test.cfg = randomization_ctx(rng, data_seed, geom_seed)

test.cfg.dropout_prob = 0.1 if test.cfg.is_dropout else 0.0
test.cfg.implementation = getattr(cudnn.attention_implementation, request.config.getoption("--implementation") or "", cudnn.attention_implementation.UNIFIED)
test.showConfig(test_no, request)

Expand Down Expand Up @@ -450,8 +429,6 @@ def test_sdpa_random_fwd_ragged_offset_multiplier_unified_L1(env_info, test_no,
with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10),
diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 0}),
is_ragged_or_padded_or_full=RandomChoice({"ragged_mult" : 1, "cu_ragged_mult" : 1}),
with_score_max=RandomChoice({True : 1, False : 3}),
with_score_sum_exp=RandomChoice({True : 1, False : 3}),
with_sink_token=RandomChoice({True : 1, False : 3}),
) as randomization_ctx:
test.cfg = randomization_ctx(rng, data_seed, geom_seed)
Expand Down Expand Up @@ -522,8 +499,6 @@ def test_sdpa_fwd_paged_L0(env_info, test_no, request, cudnn_handle):
diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}),
is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "full" : 0}),
block_size=RandomBlockSize(min=1, max=1024, with_high_probability=[1,32,128]),
with_score_max=RandomChoice({True : 1, False : 3}),
with_score_sum_exp=RandomChoice({True : 1, False : 3}),
with_sink_token=RandomChoice({True : 1, False : 3}),
) as randomization_ctx:
test.cfg = randomization_ctx(rng, data_seed, geom_seed)
Expand Down
86 changes: 86 additions & 0 deletions test/python/test_sdpa_custom_features.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
"""
Dedicated coverage for rarely-used SDPA features that are no longer exercised
by the randomized configurations in test_mhas_v2.py: ALiBi masking, the
score_max/score_sum_exp softmax outputs, and dropout. A few deterministic
configs verify that basic graph lowering and numerics keep working.
"""

import cudnn
import pytest
import torch

from sdpa.fp16 import exec_sdpa
from sdpa.random_config import ExecConfig


def make_config(
*,
data_type=torch.bfloat16,
is_infer=True,
is_alibi=False,
is_dropout=False,
dropout_prob=0.0,
with_score_max=False,
with_score_sum_exp=False,
right_bound=None,
):
cfg = ExecConfig(
data_type=data_type,
rng_data_seed=1234,
rng_geom_seed=5678,
is_alibi=is_alibi,
is_infer=is_infer,
is_paged=False,
is_bias=False,
is_block_mask=False,
is_padding=False,
is_cu_seq_len=False,
is_ragged=False,
is_dropout=is_dropout,
is_determin=False,
Comment on lines +27 to +40

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.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== File outline ==\n'
ast-grep outline test/python/test_sdpa_unpopular_features.py --view expanded || true

printf '\n== Relevant symbols ==\n'
rg -n "class ExecConfig|def exec_sdpa|is_determin|determin" -S test python . || true

Repository: NVIDIA/cudnn-frontend

Length of output: 50377


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
p = Path("test/python/test_sdpa_unpopular_features.py")
print("lines", sum(1 for _ in p.open()))
for i, line in enumerate(p.open(), 1):
    if 1 <= i <= 220:
        if "ExecConfig" in line or "is_determin" in line or "exec_sdpa" in line or "L0" in line:
            print(f"{i}: {line.rstrip()}")
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 505


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Read the relevant source slices around the referenced lines and the ExecConfig/exec_sdpa definitions.
for f in test/python/test_sdpa_unpopular_features.py $(git ls-files | rg '(^|/)(exec_sdpa|ExecConfig).*\.py$' || true); do
  [ -f "$f" ] || continue
  echo "===== $f ====="
  nl -ba "$f" | sed -n '1,240p'
done

Repository: NVIDIA/cudnn-frontend

Length of output: 255


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
paths = [Path("test/python/test_sdpa_unpopular_features.py")]
# discover the definition files referenced by rg
import subprocess, shlex
out = subprocess.check_output(rg := ["bash", "-lc", "rg -n --glob '*.py' 'class ExecConfig|def exec_sdpa|is_determin' test ."], text=True)
print(out)
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 6341


🏁 Script executed:

#!/bin/bash
set -euo pipefail
rg -n "class ExecConfig|def exec_sdpa|is_determin" test . || true

Repository: NVIDIA/cudnn-frontend

Length of output: 12028


🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf 'Searching...\n'
rg -n "is_determin|ExecConfig|exec_sdpa" test/python test || true

Repository: NVIDIA/cudnn-frontend

Length of output: 12264


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path

def show(path, start, end):
    print(f"===== {path} {start}-{end} =====")
    lines = Path(path).read_text().splitlines()
    for i in range(start, min(end, len(lines)) + 1):
        print(f"{i:4d}: {lines[i-1]}")

show("test/python/test_sdpa_unpopular_features.py", 1, 120)
show("test/python/sdpa/fp16.py", 450, 620)
show("test/python/sdpa/random_config.py", 67, 110)
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 15272


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
for path, start, end in [
    ("test/python/test_sdpa_edge_cases.py", 1, 80),
    ("test/python/test_sdpa_unpopular_features.py", 1, 95),
]:
    print(f"===== {path} {start}-{end} =====")
    lines = Path(path).read_text().splitlines()
    for i in range(start, min(end, len(lines)) + 1):
        print(f"{i:4d}: {lines[i-1]}")
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 5717


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
path = Path("test/python/sdpa/fp16.py")
lines = path.read_text().splitlines()
for start, end in [(470, 590), (840, 885)]:
    print(f"===== {path} {start}-{end} =====")
    for i in range(start, min(end, len(lines)) + 1):
        print(f"{i:4d}: {lines[i-1]}")
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 9171


🏁 Script executed:

#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("test/python/test_sdpa_unpopular_features.py")
for i, line in enumerate(p.read_text().splitlines(), 1):
    if i <= 20 or 24 <= i <= 45:
        print(f"{i:4d}: {line}")
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 1490


🏁 Script executed:

#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for path in ["test/python/test_sdpa_edge_cases.py"]:
    print(f"===== {path} =====")
    for i, line in enumerate(Path(path).read_text().splitlines(), 1):
        if i <= 55:
            print(f"{i:4d}: {line}")
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 1824


🏁 Script executed:

#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("test/python/sdpa/fp16.py")
for i, line in enumerate(p.read_text().splitlines(), 1):
    if 475 <= i <= 590:
        print(f"{i:4d}: {line}")
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 7089


Make these cases deterministic cfg.is_determin drives both use_deterministic_algorithm and the rerun bitwise check in test/python/sdpa/fp16.py, so leaving it False means these backward/dropout cases never exercise the deterministic path described here. Set it to True.

🤖 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/test_sdpa_unpopular_features.py` around lines 27 - 40, Update the
ExecConfig construction in test_sdpa_unpopular_features.py so cfg.is_determin is
set to True, ensuring the backward/dropout cases exercise the deterministic
algorithm and bitwise rerun check.

with_score_max=with_score_max,
with_score_sum_exp=with_score_sum_exp,
batches=2,
d_qk=64,
d_v=64,
s_q=512,
s_kv=512,
h_q=4,
h_k=4,
h_v=4,
diag_align=cudnn.diagonal_alignment.TOP_LEFT,
left_bound=None,
right_bound=right_bound,
dropout_prob=dropout_prob,
implementation=cudnn.attention_implementation.AUTO,
)
cfg.fill_derived_fields()
return cfg


@pytest.mark.L0
@pytest.mark.parametrize("is_infer", [True, False], ids=["fwd", "bwd"])
def test_sdpa_alibi(is_infer, request, cudnn_handle):
# ALiBi requires a causal mask (right_bound=0, top-left alignment).
cfg = make_config(is_infer=is_infer, is_alibi=True, right_bound=0)
exec_sdpa(cfg, request, cudnn_handle)


@pytest.mark.L0
@pytest.mark.parametrize("is_infer", [True, False], ids=["fwd", "bwd"])
def test_sdpa_score_max_sum_exp(is_infer, request, cudnn_handle):
cfg = make_config(is_infer=is_infer, with_score_max=True, with_score_sum_exp=True)
exec_sdpa(cfg, request, cudnn_handle)


@pytest.mark.L0
@pytest.mark.parametrize("is_infer", [True, False], ids=["fwd", "bwd"])
def test_sdpa_dropout(is_infer, request, cudnn_handle):
cfg = make_config(
data_type=torch.float16,
is_infer=is_infer,
is_dropout=True,
dropout_prob=0.1,
right_bound=0,
)
exec_sdpa(cfg, request, cudnn_handle)
Comment on lines +61 to +86

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.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== test file ==\n'
cat -n test/python/test_sdpa_unpopular_features.py | sed -n '1,220p'

printf '\n== search for support checks ==\n'
rg -n "support check|is_supported|backend_version|device_capability|skipif|pytest\.mark\.skip|cudnn\.backend_version|torch\.cuda\.get_device_capability|validate_config|exec_sdpa" test/python -S

printf '\n== nearby AGENTS ==\n'
find test -name AGENTS.md -print
if [ -f test/AGENTS.md ]; then
  printf '\n== test/AGENTS.md ==\n'
  cat -n test/AGENTS.md
fi

Repository: NVIDIA/cudnn-frontend

Length of output: 29379


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== test/python/test_sdpa_unpopular_features.py =="
sed -n '1,220p' test/python/test_sdpa_unpopular_features.py | nl -ba

echo
echo "== matching helpers in test/python =="
rg -n "def (make_config|exec_sdpa|validate_config|is_supported|supports_|skip|backend_version|get_device_capability)" test/python -S

echo
echo "== AGENTS files under test =="
find test -name AGENTS.md -print

Repository: NVIDIA/cudnn-frontend

Length of output: 249


🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '== sdpa/fp16.py validate_config and exec_sdpa ==\n'
sed -n '1,180p' test/python/sdpa/fp16.py | cat -n

printf '\n== maybe related support helpers ==\n'
sed -n '180,260p' test/python/sdpa/fp16.py | cat -n
sed -n '840,940p' test/python/sdpa/fp16.py | cat -n

printf '\n== similar SDPA tests with explicit gating ==\n'
sed -n '150,220p' test/python/test_mhas_v2.py | cat -n
sed -n '680,860p' test/python/test_mhas_v2.py | cat -n
sed -n '1000,1045p' test/python/test_mhas_v2.py | cat -n

Repository: NVIDIA/cudnn-frontend

Length of output: 33998


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path
p = Path('test/python/sdpa/fp16.py')
text = p.read_text().splitlines()
for start, end, label in [
    (1, 180, 'validate_config and top-of-file'),
    (180, 260, 'mid helpers'),
    (840, 940, 'exec_sdpa'),
]:
    print(f'== {label} ==')
    for i in range(start-1, min(end, len(text))):
        print(f'{i+1:4d}: {text[i]}')
    print()
PY

python3 - <<'PY'
from pathlib import Path
p = Path('test/python/test_mhas_v2.py')
text = p.read_text().splitlines()
for start, end, label in [
    (150, 220, 'gated rope-related tests'),
    (680, 860, 'fp8 tests with capability checks'),
    (1000, 1045, 'repro driver'),
]:
    print(f'== {label} ==')
    for i in range(start-1, min(end, len(text))):
        print(f'{i+1:4d}: {text[i]}')
    print()
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 33440


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path

def show(path, start, end, label):
    text = Path(path).read_text().splitlines()
    print(f'== {label} ==')
    for i in range(start - 1, min(end, len(text))):
        print(f'{i+1:4d}: {text[i]}')
    print()

show('test/python/sdpa/fp16.py', 1, 180, 'validate_config and top-of-file')
show('test/python/sdpa/fp16.py', 840, 940, 'exec_sdpa')
show('test/python/test_mhas_v2.py', 150, 220, 'gated rope-related tests')
show('test/python/test_mhas_v2.py', 680, 860, 'fp8 tests with capability checks')
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 25788


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
from pathlib import Path

for path, ranges in {
    'test/python/sdpa/fp16.py': [(60, 140, 'validate_config'), (840, 900, 'exec_sdpa')],
    'test/python/test_mhas_v2.py': [(150, 210, 'existing skip patterns')],
}.items():
    text = Path(path).read_text().splitlines()
    for start, end, label in ranges:
        print(f'== {path} :: {label} ==')
        for i in range(start - 1, min(end, len(text))):
            print(f'{i+1:4d}: {text[i]}')
        print()
PY

Repository: NVIDIA/cudnn-frontend

Length of output: 10959


Gate these SDPA cases on backend and GPU capability. These tests always exercise BF16/FP16 paths, but exec_sdpa only applies a few cuDNN-version skips and doesn’t gate architecture or dtype support here. Add the same cudnn.backend_version() and torch.cuda.get_device_capability() checks used elsewhere before calling exec_sdpa, so unsupported combos skip instead of failing CI.

🤖 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/test_sdpa_unpopular_features.py` around lines 61 - 86, Add the
established cuDNN backend-version and GPU compute-capability guards to
test_sdpa_alibi, test_sdpa_score_max_sum_exp, and test_sdpa_dropout before
exec_sdpa. Reuse the same checks and skip behavior used elsewhere in the test
module, including dtype-specific support conditions, so unsupported BF16/FP16
combinations are skipped while supported cases retain their current coverage.

Source: Coding guidelines