Skip the attention kernel for single-key SDPA

[guerrantif]

With a single key, softmax assigns it weight one, so attention returns its value for every query. SDPA now returns the value directly in that case, repeated across query heads for GQA/MQA and broadcast over the batch, instead of calling F.scaled_dot_product_attention. The shortcut applies only when neither attn_mask nor seqused_key_value is given, so a mask that blocks the key still takes the regular path.

This case comes up in variate attention over a single variate, e.g. TimesFM3 univariate forecasting, which runs B × N attention problems of length one.

TimesFM3, univariate target without covariates, fp32 on an L4:

series × context × horizon	before	after
64 × 1024 × 128	363 ms	214 ms
256 × 4096 × 512	6.10 s	3.57 s

Forecasts match to fp32 rounding (max. relative difference 1.2e-5), and peak memory is unchanged.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/structured-data-models/.coderabbit.yaml
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 318c9aae-0557-4f99-be12-d72581a52c79

📥 Commits

Reviewing files that changed from the base of the PR and between 6b69c4f and 98b0b29.

📒 Files selected for processing: 2

• sdm/nn/attention.py
• test/nn/test_attention.py

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

📝 Summary

Summary by CodeRabbit

• Bug Fixes
  • Improved attention results for inputs with a single key/value pair, including broadcast batch dimensions.
  • Preserved expected gradients for query, key, and value in this case.

SDPA.forward now returns expanded values directly for eligible single-key inputs. The test compares output values and gradients for query, key, and value against the reference implementation.

Changes: Single-key SDPA

Layer / File(s)	Summary
Single-key path and validation sdm/nn/attention.py, test/nn/test_attention.py	When the key sequence has length one and no attention mask or seqused_key_value is supplied, SDPA.forward expands and returns the value. With gradients enabled, it preserves zero-valued gradient connections to query and key. The test checks outputs and gradients against the reference across MHA, MQA, and GQA.

Priority: ➖ Normal

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

Merge Risk: ⚪ Minimal · up to 98b0b

No actionable issue remains with the single-key optimization. Normal validation is appropriate before merging.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

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

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly and concisely states the main change: bypassing the attention kernel for single-key SDPA.
Description check		The description accurately explains the single-key optimization, its conditions, tensor behavior, use case, and reported performance results.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches: 💡 1

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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