Share chunk-memory sizing

[ValterH]

Split from #996 to keep each review focused on one behavior or optimization.

Move the SDM_CHUNK_MEMORY_FRACTION budget used by TransformerBlock and the TabFM CellEmbedding into sdm._memory.chunk_memory_limit, and expose the automatic attention batch size limit as TransformerBlock.auto_batch_size_limit. Behavior is unchanged.

Shared base for numerical chunking and Kumo row embeddings. Split from #996; related to #994.

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

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

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: d677759f-52ed-42a6-ba5a-813b420f14c1

📥 Commits

Reviewing files that changed from the base of the PR and between 219ebe7 and 62280d8.

📒 Files selected for processing (2)

• sdm/models/tabfm/cell_embedding.py
• sdm/nn/attention.py

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

---

📝 Summary

Summary by CodeRabbit

• Improvements
  • Automatic batch sizing for cell embeddings and transformer attention now uses a consistent memory budget on CUDA devices, helping limit memory use during processing.
  • CUDA batch sizes account for memory available to the process and the size of the data being processed.
  • Transformer attention uses a fixed maximum batch-size limit on non-CUDA devices.

Walkthrough

Automatic chunk sizing now uses a shared CUDA memory-budget helper. Transformer attention delegates batch-size calculation to a method that also handles non-CUDA devices.

Changes

Automatic chunk sizing

Layer / File(s)	Summary
Shared memory budget and cell embedding sdm/_memory.py, sdm/models/tabfm/cell_embedding.py	chunk_memory_limit calculates a CUDA memory limit from device memory, the per-process memory fraction, and SDM_CHUNK_MEMORY_FRACTION. CellEmbedding uses this limit and retains its fixed-memory subtraction and batch-size calculation.
Transformer attention batch sizing sdm/nn/attention.py	TransformerBlock.forward delegates automatic sizing to auto_batch_size_limit. The method returns 65,535 for non-CUDA devices and calculates a clamped limit for CUDA devices using the shared memory budget and per-example memory estimate.

Priority: ⬇️ Low

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

Merge Risk: ⚪ Minimal · up to 62280

The shared chunk-sizing calculation appears to preserve existing behavior. No actionable merge-blocking risk remains after normal checks.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: centralizing and sharing chunk-memory sizing across attention and TabFM embedding code.
Description check	✅ Passed	The description directly explains the changes to chunk-memory budgeting and automatic attention batch-size sizing. It matches the pull request objectives and identifies related context.
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.

• 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

---

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