benchmark: compute causal SDPA FLOP counts without allocating masks

[fallintoplace]

Why

The benchmark FLOP helper allocated torch.ones((q_seqlen, kv_seqlen)) for causal masks, which can allocate huge CPU masks (e.g. 32768 x 32768) before counting non-masked elements.

What changed

• Added count_causal_nonmasked_elems(...) to compute causal non-masked element counts using row-bound arithmetic.
• Replaced torch.tril(...).sum() usage in flops(...) with the new helper.
• Added unit tests covering:
  • square masks
  • rectangular masks
  • top-left causal and bottom-right causal
  • sliding-window variants
  • reference cross-check against a brute-force boolean-mask computation

Test coverage

Added test/python/test_sdpa_benchmark_single_sdpa.py with explicit assertions and a reference cross-check.

Summary by CodeRabbit

• New Features
  • Improved attention benchmark calculations to more accurately handle causal masking, including sliding-window behavior and different mask alignments.
• Tests
  • Added coverage for square and rectangular sequence lengths, plus sliding-window scenarios, to verify the benchmark’s attention counting matches expected results.

[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: 81c5fc1f-cd95-4eab-861d-9632009c4c05

📥 Commits

Reviewing files that changed from the base of the PR and between 52119ee and aa70d03.

📒 Files selected for processing (2)

• benchmark/sdpa_benchmark_training/benchmark_single_sdpa.py
• test/python/test_sdpa_benchmark_single_sdpa.py

---

📝 Walkthrough

Walkthrough

Adds a count_causal_nonmasked_elems helper function to compute non-masked (q, k) pairs for causal attention masks (top_left, bottom_right), with optional sliding-window support, replacing prior inline FLOPs counting logic. Adds a new test module validating this helper against a torch-based reference implementation.

Changes

Causal Element Counting

Layer / File(s)	Summary
Helper implementation and FLOPs integration benchmark/sdpa_benchmark_training/benchmark_single_sdpa.py	Adds count_causal_nonmasked_elems to compute non-masked causal (q, k) pairs for top_left/bottom_right masks with optional sliding window, and replaces prior tensor/closed-form counting in the flops(...) helper with a call to this function.
Reference implementation and unit tests test/python/test_sdpa_benchmark_single_sdpa.py	Adds _reference_causal_nonmasked_elems torch-based reference and multiple unit tests (square, rectangular, sliding-window, and a parametrized comparison test) validating the new helper.

Estimated code review effort: 2 (Simple) | ~12 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly captures the main change: computing causal SDPA FLOP counts without materializing masks.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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