[aw471]

Summary

On CUDA, chunked processors fail under torch.compile with ConstraintViolationError when their row count is marked dynamic. split_size now returns the item count directly when one chunk provably holds all items.

Description of the change

On a CUDA device with gradients off, split_size balances the chunks as ceil(n / ceil(n / capacity)). For a symbolic item count n, the result is that whole expression, not n. _split_tensor then cannot prove that the tensor fits one chunk, calls Tensor.split, and PyTorch adds a guard that the expression gives one chunk. The guard is on the marked size, which contradicts the dynamic mark:

create_symbol s54 = 5 for L['inp'].size()[0] [2, 16]
eval Eq(((s54 + Max(1, -(((-s54)//Max(1, -(((-s54)//256)))))) - 1)//Max(1, -(((-s54)//Max(1, -(((-s54)//256))))))), 1) [guard added] return tensor.split(size, dim=dim)  # sdm/_memory.py:16 in _split_tensor
ConstraintViolationError: Constraints violated (L['inp'].size()[0])!

split_size now returns num_items before the arithmetic when num_items <= capacity is statically known. The balanced size equals num_items in that case anyway, so eager results are unchanged. When the count can exceed the capacity, the arithmetic runs as before.

In Kumo Tabular, ClipSigma and RankGaussian call split_size with the number of query rows during predict. On CPU split_size returns at its first line, so only CUDA is affected.

Reproduction (PyTorch 2.14.1; no GPU needed, split_size only reads the device type):

import torch
import sdm._memory
from sdm._memory import _split_tensor, split_size

sdm._memory.chunk_memory_limit = lambda device: 1024  # 256 items of 4 bytes


def chunks(inp):
    size = split_size(inp.size(0), 4, torch.device("cuda"))
    return torch.cat([chunk * 2 for chunk in _split_tensor(inp, size, dim=0)])


compiled = torch.compile(chunks, fullgraph=True, dynamic=False)
with torch.inference_mode():
    for rows in (5, 9, 3, 16):
        inp = torch.randn(rows, 3)
        torch._dynamo.mark_dynamic(inp, 0, min=2, max=16)
        compiled(inp)  # ConstraintViolationError before this change

Validation

Ran the script above with TORCH_LOGS="recompiles,dynamic,graph_breaks" and the eager and inductor backends, before and after this change. Both backends give the same result (BEGIN/END lines are printed by the script):

Before: every call ends with ConstraintViolationError.

After:

BEGIN call=0 shape=(5, 3)
END call=0 graphs=1 matches_eager=True
BEGIN call=1 shape=(9, 3)
END call=1 graphs=1 matches_eager=True
BEGIN call=2 shape=(3, 3)
END call=2 graphs=1 matches_eager=True
BEGIN call=3 shape=(16, 3)
END call=3 graphs=1 matches_eager=True

On a Tesla T4, compiled Kumo Tabular predict with the query rows marked fails with this error for all 6 inputs tried without the change. With it, the error is gone.

FULL_TEST=1 pytest test/processing: 451 passed.

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

Important

Review skipped

Review was skipped as selected files did not have any reviewable changes.

Run configuration

• Source mode: Auto
• Repository: NVIDIA/structured-data-models
• Branch: compile-split-size
• File: .coderabbit.yaml
• Commit: 9cdea2b
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: b846a3a8-d902-4ba8-bbea-f18cb91c17cb

Commits

Reviewing files that changed from the base of the PR and between 0da9c57 and 9cdea2b.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• Trigger review

No actionable comments were generated in the recent review.

Summary

Summary by CodeRabbit

• Bug Fixes
  • Improved chunk sizing for inputs that fit within the available capacity, including empty inputs, to ensure a valid chunk count is returned.

split_size now returns at least one when num_items <= capacity is statically known true. Otherwise, it uses the existing balanced-chunk calculation.

Changes: Dynamic split sizing

Layer / File(s)	Summary
Split sizing shortcut sdm/_memory.py	When num_items <= capacity is statically known true, split_size returns max(num_items, 1) before the balanced-chunk calculation.

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

Merge Risk: Minimal · previous review 0da9c

This change avoids a torch.compile constraint error when a row count is marked dynamic, and eager results stay the same. No merge-blocking risk was found.

Architecture risk: Low · up to 0da9c

The change affects 1 system.

Changed systems: sdm

Architecture concerns
No architecture-level concerns identified.

Review details

Systems and components

• observed — sdm (service) was modified; 1 changed file maps to changed impact.

Before / after behavior

• observed — Modified behavior in sdm/_memory.py: split_size now returns at least 1 directly when num_items <= capacity is statically known true; otherwise, execution continues to the existing balanced-chunk calculation.

Pre-merge checks: 5

Check name	Status	Explanation
Title check		The title clearly identifies the main change: skipping balanced split sizing when one chunk can hold all items, including the dynamic-shape context.
Description check		The description directly explains the CUDA dynamic-shape failure, the split_size change, the affected callers, reproduction steps, and validation results.
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 with cyclomatic complexity of at least 2. Documentat…
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.

Full details: Docstring Coverage

Explanation

Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 1 functions across 1 files.

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

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/structured-data-models
• Branch: compile-split-size
• File: .coderabbit.yaml
• Commit: 0da9c57
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: e7788016-1598-48f9-b6e3-28d9282941aa

Commits

Reviewing files that changed from the base of the PR and between a00920e and 0da9c57.

Files selected for processing: 1

• sdm/_memory.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.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.25%. Comparing base (df0da4c) to head (9cdea2b).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
sdm/_memory.py	0.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1359      +/-   ##
==========================================
- Coverage   82.26%   82.25%   -0.02%     
==========================================
  Files         186      186              
  Lines       11966    11968       +2     
==========================================
  Hits         9844     9844              
- Misses       2122     2124       +2     

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