Fix D192 BF16 short-sequence performance regression

[adshen]

Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes.
• I reviewed the Hard Rules in the AGENTS.md for each directory this PR touches and the changes comply.
• I added the required category, area, and origin labels.

Affected area

FE OSS kernels or CuTeDSL; benchmarks or performance.

Summary

Fix the SM100 D192/D128 BF16 short-sequence performance regression by selecting the measured D192 scheduler and matching source schedule for square top-left causal sequences from 512 through 3584 tokens.

Why

The D192 half kernel previously used plain LPT. After scheduler selection moved to the generic heuristic, short BF16 sequences selected LPT-L2. The public-source exp2/descriptor/election schedule used by the long-sequence path also costs additional time when paired with plain LPT at these short lengths. Earlier validation focused on 8K and did not expose the short-sequence crossover.

This change:

• selects LPT as the D192-specific primary only in the measured short BF16 region;
• selects the matching existing source schedule through a private codegen flag;
• preserves explicit user knobs, the LPT-L2 and NATURAL runner configurations, and NATURAL fallback behavior;
• leaves FP16, FP8, MXFP8, no-mask, bottom-right, GQA, sink, SWA, THD/runtime-length, CTA1, and split-KV paths unchanged.

The remaining 2K gap to a historical build is about 2.4-3.0%. That build used internal compiler controls which cannot be used in public source; this PR recovers the public-source regression without restoring them. Later public-source tuning has already recovered essentially all of the corresponding 8K gap.

Paired measurements used a local SM100 GPU at fixed 847 MHz SM / 4000 MHz memory clocks, with 10 warmups and 40 measured iterations per side. Lower is better.

BF16 top-left, Sq=Skv=2048	Upstream	PR	Change
Kimi K2.6, B2 H64	0.341 ms	0.323 ms	-5.28%
Kimi K3, B2 H96	0.485 ms	0.474 ms	-2.27%
DeepSeek V3, B2 H128	0.660 ms	0.620 ms	-6.06%

The non-target regression sweep covered 87 BF16/per-tensor-FP8/MXFP8 model cells across 2K/4K/8K/16K/32K and top-left/no-mask. Median change was 0.000%; the largest slowdown was 0.321% (1 microsecond). A supplementary 30-cell FP16 sweep also had a 0.000% median; its only initial outlier disappeared in four crossed-order repeats.

Related issues

None.

API and compatibility impact

No API or engine-capability change. Explicit scheduler/CGA choices remain authoritative.

Testing

cd test/python
python3 -m pytest -q sdpa/frost/test_sdpa_fwd_heuristics.py
python3 -m pytest -q -m "L0 or L1" \
  sdpa/frost/test_sdpa_fwd_dsl_sm100.py::test_dsl_sm100_d192_d128 \
  sdpa/frost/test_sdpa_fwd_dsl_sm100.py::test_dsl_sm100_d192_d128_short_bf16_lpt

• Heuristic suite: 24 passed.
• D192 FP16/BF16 dense/causal plus the exact 2K BF16 O/LSE reference case: 5 passed.
• Representative 2K path: Compute Sanitizer synccheck and memcheck both reported 0 errors.
• Pre-commit hooks passed on all changed files.

Summary by CodeRabbit

• Performance Improvements

  • Improved scheduling for short, dense BF16 attention workloads on supported hardware.
  • Added specialized handling for D192/D128 causal attention configurations.

• Bug Fixes

  • Corrected execution-path selection for qualifying BF16 causal workloads while preserving existing behavior for other configurations.

• Tests

  • Added regression coverage comparing outputs and statistics with FP32 references.
  • Expanded scheduler-selection tests across sequence lengths and configuration variations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: be3e0cb6-76dd-456e-8abd-8ec70d0e7cff

📥 Commits

Reviewing files that changed from the base of the PR and between 2cc71fb and f94abb9.

📒 Files selected for processing (5)

• python/cudnn/sdpa/fwd/config_sm100.py
• python/cudnn/sdpa/fwd/heuristics.py
• python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py
• test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
• test/python/sdpa/frost/test_sdpa_fwd_heuristics.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

The change adds short-sequence BF16 D192 LPT detection, scheduler selection, template propagation, and kernel dispatch behavior. New tests validate scheduler decisions, derived parameters, output values, and LSE statistics.

Changes

D192 BF16 LPT specialization

Layer / File(s)	Summary
Short-region contract python/cudnn/sdpa/fwd/config_sm100.py	Defines TemplateParams.d192_short_bf16_lpt. Derives the flag for dense square BF16 D192 configurations that meet the causal, topology, and sequence-length constraints.
Scheduler selection and validation python/cudnn/sdpa/fwd/heuristics.py, test/python/sdpa/frost/test_sdpa_fwd_heuristics.py	Selects SCHED_LPT for qualifying short BF16 D192 cases. Preserves other scheduler behavior and validates short, long, split, recommendation, and fallback cases.
Kernel specialization dispatch python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py, test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py	Adds SCHED_LPT dispatch and excludes the short BF16 specialization from the dense causal mixed-exp2 path. Validates output and LSE statistics against the FP32 reference.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to f94ab

This change selects an LPT specialization for qualifying short BF16 D192 causal attention workloads while preserving stated fallback behavior. Functional and scheduler regression coverage is included, with no concrete current-head merge-blocking risk identified.

Suggested reviewers: vedaanta

Sequence Diagram(s)

sequenceDiagram
  participant Heuristics
  participant ConfigSM100
  participant D192Kernel
  Heuristics->>ConfigSM100: evaluate d192_short_bf16_lpt_region
  ConfigSM100-->>Heuristics: return short BF16 region result
  Heuristics->>D192Kernel: select SCHED_LPT
  D192Kernel->>D192Kernel: use non-mixed exp2 and issue-election path

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 12 functions across 5 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: fixing the D192 BF16 short-sequence performance regression.
Description check	✅ Passed	The description follows the repository template, identifies the affected areas, explains the problem and solution, documents compatibility impact, and lists testing commands with results. All required…
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[adshen]

@cudnn-ci-bot run frost

[cudnn-ci-bot]

🏁 Pipeline finished

SHA: f94abb9
Targets: frost
Branch: cudnn-gh/pr-932-f94abb9
Pipeline: 66608912
Last updated: 2026-09-07 11:02 UTC

[vedaanta]

@cudnn-ci-bot run frost

[cudnn-ci-bot]

🏁 Pipeline finished

SHA: f94abb9
Targets: frost
Branch: cudnn-gh/pr-932-f94abb9
Pipeline: 66864570
Last updated: 2026-09-09 00:27 UTC

[YangXu1990uiuc]

Codex bot review · model: gpt-6-astra

Reviewed f94abb9ac9dc against base 64f6b391144c.

The scoped scheduler/codegen change makes sense, and I found no actionable correctness issue. Explicit scheduler choices and the NATURAL fallback remain intact. Two non-blocking suggestions: record the measured GPU model/SM count alongside the clock settings, since the 2K gain did not transfer to the smaller SM100 device below; and put one short-path numerical case in L0, since the new 2K O/LSE regression is L1. Rebase the single import conflict in the renamed kernels/sm100/prefill_d192_d128_f16.py by retaining both read_clc_payload and SCHED_LPT. I verified that resolution only in a disposable integration checkout; the PR branch is unchanged.

Validation: Built matching bindings for the reviewed head and its exact merge base. With CuTe DSL 4.8 on SM100, the heuristic suite passed 26 tests and the focused D192 suite passed 11 tests under explicit L0/L1 selection. Eight additional full O/LSE reference cases passed at square lengths 511, 512, 513, 1025, 3583, 3584, 3585 and 4096. The actual loaded module enabled the short schedule only inside [512,3584]. NATURAL, LPT and LPT-L2 were each explicitly selected and verified at length 513, with numerical results checked. Paired merge-base/head measurements used ABBA order, actual template/scheduler checks, reference rows across all heads, 16 calls per captured graph, seven timing batches, and changed-V/poisoned-output replay checks. On the 68-SM SM100 device, B2/H8/S512 BF16 causal improved from 22.67 to 21.16 us (-6.69%). B2/H64/S2048 was 516.15 -> 520.85 us (+0.91%); B2/H128/S2048 was 1015.10 -> 1017.77 us (+0.26%). These small 2K slowdowns repeated in both arms; the claimed 2K speedup was not reproduced on this device. Non-target controls were essentially unchanged: BF16 causal S4096 +0.01%, BF16 non-causal S2048 +0.02%, FP16 causal S2048 +0.33% (all B2/H8). Warmed graph-call CPU time stayed about 28–34 us; observed differences were at most about 1.1 us. No clocks were changed. A temporary merge onto develop 64f6b39, resolving only the import conflict and building separate bindings for that integrated source, passed the three directly affected tests: scoped scheduler selection, NATURAL fallback and the new 2K O/LSE regression. The template-cache identity includes all frozen template parameters, so the new private flag is part of that identity. Current-head Style passed.

Limitations: The performance check is a component measurement on a 68-SM device with CuTe DSL 4.8; it does not independently reproduce the reported fixed-clock model grids or the broad 87-cell/30-cell sweeps. The reviewed head still needs the described rebase. The three temporary integration tests do not establish that the complete integrated CI suite passes. Existing current-head CI requests were reused for inspection rather than duplicated; SM103/SM107 execution was not independently rerun.

Approved under the trial's no-P0 threshold; any P1/P2 findings above remain for the owner to address. Merge timing stays with the owner.