Compile SDPA with bounded tensor-dependent batch sizes (#1257)

# 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:

```python
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)
```

```text
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.

Author: aw471

sdm/nn/attention.py

@@ -171,7 +171,7 @@ def forward(
         key = key.transpose(-3, -2)  # [B, Hkv, KV, C]
         value = value.transpose(-3, -2)  # [B, Hkv, KV, C]
 
-        # CUDA attention kernels accept at most 65,535 batch entries per call:
+        # CUDA attention kernels accept at most 65,535 batch entries per call.
         outs = [
             F.scaled_dot_product_attention(
                 query=query[i : i + 65_535],
@@ -185,7 +185,11 @@ def forward(
                 is_causal=is_causal,
                 scale=self.scale,
             ).transpose(-3, -2)  # [B, Q, Hq, C]
-            for i in range(0, query.size(0), 65_535)
+            for i in (
+                (0,)
+                if query.size(0) <= 65_535
+                else range(0, query.size(0), 65_535)
+            )
         ]
         out = outs[0] if len(outs) == 1 else torch.cat(outs)