[aw471]

Compiled TransformerBlock can fail when category-dependent preprocessing produces a batch size known only at runtime. This was isolated while testing learned-class KumoTabular execution; the example replaces attention to expose the buffer check alone:

import torch
from sdm.nn import TransformerBlock

class SimpleAttention(torch.nn.Module):
    def forward(self, query, **kwargs):
        return query.sin()

block = TransformerBlock(8, 2, mlp=None).eval()
block.attn = SimpleAttention()

def run(query, ids):
    selected = ids.unique()
    torch._check(selected.numel() > 0)
    query = query.index_select(0, selected)
    out = query.new_empty((3, query.size(0), 8)).transpose(0, 1)
    return block(query, out=out)

with torch.no_grad():
    torch.compile(run, backend="inductor", fullgraph=True)(
        torch.randn(8, 3, 8), torch.tensor([0, 1] * 4)
    )

Could not guard on data-dependent expression Eq(24*u0, 0) | (Eq(u0, 1) & Eq(8*Max(1, u0), 8))

The traceback points to this condition in TransformerBlock._forward:

and not out.is_contiguous()

A transposed buffer becomes contiguous when that batch dimension is one, so tracing cannot decide the condition from its shape alone. Use the existing functional residual path for every compiled output buffer, then copy into the supplied buffer. Dtype conversion and buffer identity are preserved; eager execution is unchanged.

CPU validation, PyTorch 2.14	Result
22 layout/dtype cases, both graph policies	Pass; output buffer identity preserved
Contiguous buffers with real attention, three dtype/autocast configurations	Bitwise equal to previous compiled path
Existing attention tests	9 passed, 2 CUDA skipped

L4, PyTorch 2.14: six real-Transformer contiguous-buffer FP32/BF16 cases had bitwise-equal outputs, identical generated GPU code using a shared warmed compiler cache, and unchanged peak allocated memory. Incremental peaks for medium/large cases were 12/24 MiB in FP32 and 6/12 MiB in BF16, before and after.

Steady normal-call timings: 100 warmups, then 31 paired AB/BA blocks of 100 calls. Shapes are (batch, query length, key length, channels); percentage changes are paired medians.

BF16 shape	Before (ms)	After (ms)	Paired change (95% interval)
(1, 32, 32, 128)	0.284360	0.285580	+0.423% (+0.107% to +0.585%)
(4, 512, 2048, 256)	0.279823	0.279658	−0.215% (−0.864% to +0.267%)

The small case adds about 1.22 µs; the large-case interval includes no change. These are block measurements, not end-to-end KumoTabular or cold-compilation results. This fixes the buffer check, not other dynamic attention/chunking limitations.

[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 (7c356c7) to head (5d522ec).
⚠️ Report is 3 commits behind head on main.
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1259      +/-   ##
==========================================
- Coverage   82.93%   82.91%   -0.02%     
==========================================
  Files         173      173              
  Lines       11172    11191      +19     
==========================================
+ Hits         9265     9279      +14     
- Misses       1907     1912       +5     

☔ 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-transformer-symbolic-output
• File: .coderabbit.yaml
• Commit: 5d522ec
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 4bc4becc-19b5-4f5d-bc05-bfd784a1fe5b

📥 Commits

Reviewing files that changed from the base of the PR and between 84ffbc3 and 5d522ec.

📒 Files selected for processing: 1

• sdm/nn/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 compiled transformer output handling. Results are now copied into the provided output tensor consistently, including when it is contiguous, while preserving its data type.

When torch.compiler.is_compiling() and out is provided, TransformerBlock._forward now uses the temporary-tensor and copy path for both contiguous and non-contiguous output tensors.

Changes: Compiled output handling

Layer / File(s)	Summary
Compiled output handling sdm/nn/attention.py	The compiling path now handles every supplied out tensor. It computes the result in a temporary tensor, casts it to out.dtype, and copies it to out.

Priority: ➖ Normal

Estimated code review effort: 1 (Trivial) | ~3 minutes

Merge Risk: ⚪ Minimal · up to 5d522

No concrete regression is established for supported compiled output buffers; the change is mergeable with normal checks.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 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 describes the main change: preventing symbolic contiguity checks on compiled output buffers.
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.
Description check		The pull request description directly explains the compiled TransformerBlock failure, the symbolic contiguity check, the proposed fix, and validation results. It matches the changeset.

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