[aw471]

Summary

Allow SDPA compilation when a tensor-dependent batch size is provably within the existing 65,535-entry chunk limit.

Description of the change

PyTorch 2.14 CPU, default Inductor backend:

import torch
from sdm.nn import SDPA

attention = SDPA(1).eval()

def run(query, ids):
    selected = ids.unique()
    torch._check(selected.numel() > 0)
    query = query.index_select(0, selected)
    return attention(query, query, query)

compiled = torch.compile(run, fullgraph=True)
with torch.no_grad():
    compiled(torch.randn(8, 1, 1, 1), torch.arange(8) % 3)

Could not extract specialized integer from data-dependent expression u0
Caused by: for i in range(0, query.size(0), 65_535)

The original Python loop requires a concrete batch size. Use the single iteration (0,) when at most one chunk is needed; retain the original loop for larger batches. Attention computation and chunk sizes are unchanged, with no custom operation.

This covers provably bounded batches. Strict compilation across an unknown 65,535-entry boundary remains unsupported.

Validation

• PyTorch 2.14 CPU Inductor: fullgraph=False and True, tensor-dependent batches changing 2 → 8 → 1, matching eager outputs.
• 21 parity cases covering FP32/FP64/BF16, batches 1/65,535/65,536, masks, causal attention, dropout/RNG and gradients.
• Existing SDPA tests: 8 passed, 7 CUDA-only skipped. Pre-commit passed. CUDA execution and performance remain unmeasured.

[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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.91%. Comparing base (ec0203b) to head (765c509).
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1257   +/-   ##
=======================================
  Coverage   82.91%   82.91%           
=======================================
  Files         173      173           
  Lines       11191    11191           
=======================================
  Hits         9279     9279           
  Misses       1912     1912           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Source mode: Auto
• Repository: NVIDIA/structured-data-models
• Branch: compile-attention-symbolic-chunk-loop
• File: .coderabbit.yaml
• Commit: 765c509
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: a13b4f36-46de-4340-9805-751404ab890a

📥 Commits

Reviewing files that changed from the base of the PR and between cfef9e0 and 765c509.

📒 Files selected for processing: 1

• sdm/nn/attention.py

💤 Files with no reviewable changes: 1

• sdm/nn/attention.py

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

📝 Summary

Summary by CodeRabbit

• Improvements
  • Attention processing handles batches of up to 65,535 entries as a single group and divides larger batches into groups that stay within the established size limit. This keeps the processing behavior consistent across batch sizes while avoiding unnecessary grouping for smaller batches.

SDPA.forward uses one chunk start for flattened batches of at most 65,535 entries. Larger batches use chunk starts spaced 65,535 entries apart.

Changes: SDPA batch chunking

Layer / File(s)	Summary
Select and use batch chunk starts sdm/nn/attention.py	SDPA.forward selects (0,) for batches of at most 65,535 entries and uses 65,535-entry steps for larger batches. The attention-call loop uses these chunk starts. The CUDA kernel limit comment now ends with a period.

Priority: ➖ Normal

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

Merge Risk: ⚪ Minimal · up to 765c5

No actionable merge-blocking issue is established for this change; normal validation can proceed.

Pre-merge checks: 5

Check name	Status	Explanation
Description check		The description directly explains the SDPA compilation change, the bounded 65,535-entry behavior, the implementation approach, and validation results.
Title check		The title clearly and concisely identifies the main change: enabling SDPA compilation for bounded tensor-dependent batch sizes.
Docstring Coverage		Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files.
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.

✨ Finishing Touches

📝 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.}