[aw471]

Summary

Keep fitted class labels on GPU while offloading the model cache. This avoids the string-label transfer that prevents full-graph KumoTabular prediction.

Description of the change

Temporarily remove class metadata during the existing pinned-CPU offload and restore it in the existing finally block. The large cached tensors still offload; labels remain on their original GPU. On the tested two-class dataset, this retains 2 KiB of additional GPU allocation.

Main fails during cache prefetch with torch._dynamo.exc.UserError: Could not extract specialized integer from data-dependent expression u0 (unhinted: u0). (Size-like symbols: none) on PyTorch 2.14.1.

Validation

Class identity, nested tensor copying, and restoration after copy failure were checked manually. The GPU reproduction confirms that cache prefetch progresses past the previous string-copy failure. Other independent fixes are needed to compile the complete predict call.

Combined with the other three predict fixes: pretrained KumoTabular large, default recipe, 8 estimators, 10k context / 512 queries, RTX PRO 6000, PyTorch 2.14.1+cu130 and FP16 autocast. Main and #1275 each passed three actual-Inductor fullgraph=True predictions with one graph, zero breaks/recompiles, and max probability error 1.47e-4 versus eager; predicted classes matched. This is correctness validation, not a performance benchmark.

Minimal reproducer

import pyarrow as pa
import torch
from sdm import TableTensor
from sdm.models import KumoTabular

x = TableTensor(numerical=torch.randn(128, 2, device="cuda"))
query = TableTensor(numerical=torch.randn(16, 2, device="cuda"))
y = TableTensor.from_arrow(
    pa.table({"label": ["No", "Yes"] * 64}),
    stypes={"label": "categorical"},
    device="cuda",
)
model = KumoTabular(task="classification", size="small", pretrained=False, device="cuda")
with torch.amp.autocast("cuda", dtype=torch.float16):
    model.fit(x, y, num_estimators=8)
with torch.inference_mode(), torch.amp.autocast("cuda", dtype=torch.float16):
    torch.compile(model.predict, fullgraph=True)(query)

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

[aw471]

The author requested and pre-approved a full review of this PR.

@coderabbitai full review

[coderabbitai]

Action performed

Full review finished.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.38%. Comparing base (e32a878) to head (565fccb).
✅ All tests successful. No failed tests found.

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

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1329      +/-   ##
==========================================
- Coverage   82.39%   82.38%   -0.02%     
==========================================
  Files         185      185              
  Lines       11777    11779       +2     
==========================================
  Hits         9704     9704              
- Misses       2073     2075       +2     

☔ 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/cache-class-labels-on-device
• File: .coderabbit.yaml
• Commit: 565fccb
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: 9330b8a9-db37-4b2a-930f-5771f479d5df

Commits

Reviewing files that changed from the base of the PR and between e32a878 and 565fccb.

Files selected for processing: 1

• sdm/models/base/icl.py

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

Summary

Summary by CodeRabbit

• Bug Fixes
  • Improved CUDA model fitting when using multiple contexts. Class labels are preserved on the GPU while other cached data is copied to CPU memory, and fitting now synchronizes the CUDA stream after the transfer.

During CUDA fitting with multiple contexts, fit removes the cache’s classes entry before transferring the other cache tensors to pinned CPU memory. A finally block restores classes, including if the transfer raises.

Changes: CUDA cache handling

Layer / File(s)	Summary
Preserve classes during cache copy sdm/models/base/icl.py	fit removes classes before transferring cache tensors to pinned CPU memory and restores it in a finally block.

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

Merge Risk: Minimal · up to 565fc

No actionable merge-blocking risk remains in this change.

Pre-merge checks: 5

Check name	Status	Explanation
Title check		The title clearly and concisely describes the main change: keeping fitted class labels on the GPU for compiled prediction.
Description check		The description is directly related to the changeset. It explains the GPU label retention, cache offload behavior, validation, and compilation issue addressed.
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

CodeRabbit Verify

• Run CodeRabbit Verify to test the changes in this pull request.