[aw471]

Summary

Preserve asynchronous KV-cache prefetching in compiled prediction without allowing copies to overwrite buffers still used by computation. Eager prediction keeps its existing transfer behavior.

Description of the change

On PyTorch 2.14.1+cu130, compiled prediction could silently return incorrect warm outputs: maximum probability error 0.39 in KumoTabular and error 0.688 in the ordinary-tensor reproducer below. There was no compiler exception or graph break; generated code reused storage across the copy and compute streams.

Put CPU-to-GPU tensor copies in a nested compiled helper, whose output allocations are separate from the parent graph's buffer reuse. The next estimator's cache still transfers while the current estimator computes, using the existing streams, events, and lifetime tracking. Already-resident tensors keep their identity. Prepare the transfer stream during eager fit so the first prediction does not mutate stream state and trigger another compilation.

This replaces the previous revision's workaround, which put compiled copies and computation on the same stream. No custom operations, compiler settings, persistent staging buffers, or PyTorch changes are needed. The feature check preserves eager imports on older supported PyTorch versions; compilation was validated on 2.14.

PyTorch references: nested compiled regions explain the helper API; CUDA streams and memory lifetime explain overlap and synchronization. The allocation-boundary workaround is validated by this PR's reproducer and GPU runs; it is not a documented guarantee that this API fixes every cross-stream buffer-reuse problem.

Bug investigation:

• Eager execution passes, while repeated compiled calls corrupt outputs. Disabling only Inductor's allow_buffer_reuse makes the ordinary-tensor reproducer pass; disabling only inplace_buffers does not.
• In the instrumented PyTorch 2.14.0 reproducer, a CPU-to-GPU DeviceCopy runs on transfer stream 1, but get_buf_stream() reports stream 0 after following mutation_renames to a later alias. This gives the copy destination the same reuse key as a compute buffer, admitting unsafe reuse across the two streams. Instrumentation logs the original planner's decisions without changing them. The full-model failure also reproduces on 2.14.1; this precise internal diagnostic was performed on 2.14.0.
• The nested copy helper avoids that parent-graph reuse decision by allocating and returning its copy outputs separately. Correctness and actual copy/compute overlap were verified at 20k and 40k context with default compiler options.

Related upstream context: PyTorch #178549 — Prevent cross-stream memory planning buffer reuse describes the same category of race: one stream overwrites reused storage while another still reads it. It is not a confirmed duplicate or a confirmed fix for our case. The instrumented runtime already has stream-aware reuse keys; the diagnostic above shows incorrect stream attribution despite that safeguard. No exact matching upstream issue was found.

Validation

Same physical RTX PRO 6000 Blackwell; PyTorch 2.14.1+cu130; Diabetes130US; default pretrained KumoTabular LARGE classifier and recipe, 8 estimators, 512 queries, FP16 autocast, eager fit. Each case uses a fresh process, five warmups and twenty synchronized public predict measurements. Native modes include #1275; the copy fix itself is independent of #1275.

Warm prediction, including preprocessing. Memory is peak PyTorch allocated/reserved GPU memory during prediction, in decimal GB; reserved memory includes allocator blocks retained from fit and warmup.

Context rows	Implementation	Median latency (ms)	Allocated (GB)	Reserved (GB)
20,000	Main eager, cuDF	206.46	1.871	11.673
20,000	Native eager	196.89	1.871	11.673
20,000	Native compiled, previous same-stream workaround	124.19	1.028	10.687
20,000	Native compiled, this fix	78.93	1.656	11.778
40,000	Main eager, cuDF	200.87	2.636	14.806
40,000	Native eager	197.87	2.636	14.806
40,000	Native compiled, previous same-stream workaround	189.69	1.036	13.067
40,000	Native compiled, this fix	112.81	2.168	15.016

Relative to the previous workaround, latency falls 36.4% / 40.5%, while reserved memory rises 1.09 / 1.95 GB. Relative to main eager, reserved memory rises 0.9% / 1.4%; allocated memory is lower. Observed whole-device NVML peaks rise 0.8% / 1.3% versus eager. Preserving overlap therefore has a memory cost compared with running copies and compute on one stream.

• torch.compile(..., fullgraph=True) with default options and buffer reuse enabled: one graph, no graph breaks, guard failures, or measured recompiles. First predictions, warmups, and every timed prediction match both eager references; maximum probability error is 0.000322. Input and sampled fitted-cache hashes remain unchanged.
• A separate three-prediction GPU trace at 40k confirms that 1,390 of 1,488 CPU-to-GPU copies overlap computation. Repeated ordinary-tensor and nested-cache stress probes also pass, including delayed compute/copy streams and varied tensor layouts.
• First compilation took about 4.7 minutes and is excluded above. Startup conditions were not controlled for a fair cold-compilation comparison. CPU import/prefetch smoke checks pass; no CPU speedup is claimed. No new tests are added.

Minimal reproducer for the incorrect-output bug

import torch
from sdm.models.base.model import Model

class Probe(Model):
    @classmethod
    def default_recipe(cls):
        raise NotImplementedError

    def forward(self, x):
        values = []
        with self._prefetch(self.items, x.device) as prefetched:
            for item in prefetched:
                y = (item @ x).sin().square()
                values.append(y.mean())
        return torch.stack(values)

probe = Probe().eval()
probe.items = [torch.full((1024, 256), float(i + 1), pin_memory=True) for i in range(8)]
x = torch.eye(256, device="cuda")
expected = torch.arange(1, 9, device=x.device, dtype=x.dtype).sin().square()
compiled = torch.compile(probe, fullgraph=True)
with torch.inference_mode():
    for _ in range(5):
        torch.cuda._sleep(100_000_000)  # Expose asynchronous buffer reuse.
        actual = compiled(x)
        torch.cuda.synchronize()
        torch.testing.assert_close(actual, expected)

[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 31.57895% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.33%. Comparing base (35272bf) to head (d4969b4).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
sdm/models/base/model.py	29.41%	12 Missing ⚠️
sdm/models/base/icl.py	50.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1330      +/-   ##
==========================================
- Coverage   82.38%   82.33%   -0.06%     
==========================================
  Files         185      185              
  Lines       11779    11794      +15     
==========================================
+ Hits         9704     9710       +6     
- Misses       2075     2084       +9     

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

Summary

Summary by CodeRabbit

• Bug Fixes
  • Improved data prefetching during PyTorch compilation and when transferring data between CUDA and non-CUDA devices, helping keep data movement consistent across execution modes.
  • CUDA-based model fitting now sets up a transfer stream when needed, supporting data transfers during fitting on CUDA devices.

CUDA prefetching uses compilation-aware transfers when a transfer stream is present. ICLModel.fit creates a transfer stream for a CUDA input device when no stream is already registered.

Changes: CUDA prefetch transfer behavior

Layer / File(s)	Summary
Compilation-aware prefetch transfers sdm/models/base/model.py, sdm/models/base/icl.py	_transfer handles DeviceMixin containers recursively during compilation. It leaves tensors on the target device unchanged and copies other tensors through _copy_to_device. The prefetcher uses _transfer for initial and subsequent items when a transfer stream is present. ICLModel.fit creates a CUDA stream for x.device when x is CUDA and no stream is registered for that device.

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

Merge Risk: Moderate · up to d4969

Compiled prediction can fail during prefetch for workloads with sufficiently heterogeneous cache shapes; raise the reuse limit before merging.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title "Fix cache copy ordering in compiled prediction" names the main change: CPU-to-GPU cache copies are moved into a nested compiled helper so buffer reuse no longer overwrites in-use buffers du…
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 description is specific and related to the changeset. It explains the bug (incorrect warm outputs in compiled prediction on PyTorch 2.14.1+cu130 caused by cross-stream buffer reuse), the fix (CPU-…

Full details: Title check

Explanation

The title "Fix cache copy ordering in compiled prediction" names the main change: CPU-to-GPU cache copies are moved into a nested compiled helper so buffer reuse no longer overwrites in-use buffers during compiled prediction. The title is specific, short, and matches the objectives, which describe the same fix.

Full details: Description check

Explanation

The description is specific and related to the changeset. It explains the bug (incorrect warm outputs in compiled prediction on PyTorch 2.14.1+cu130 caused by cross-stream buffer reuse), the fix (CPU-to-GPU copies moved into a nested compiled helper, transfer stream prepared during eager fit), validation results, and includes a reproducer. This matches the summary of changes in model.py and icl.py.

Full details: Docstring Coverage

Explanation

Docstring coverage is 50.00% which is insufficient. 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 4 functions across 2 files.

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

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

Warning

Some tools did not complete. Review the errors below.

ast-grep (0.51.1)

ast-grep could not parse rule config: /ast-grep-rules/ast-grep-essentials/rules/kotlin/security/sql-injection-statement-execute-concat-kotlin.yml

Tip

🎯 Customize your risk assessments

Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.

✅ Automatic approvals, with your rules

Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.

CodeRabbit Verify

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

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/structured-data-models
• Branch: compile/cache-copy-ordering
• File: .coderabbit.yaml
• Commit: d4969b4
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: d825ec98-230c-40f6-86cb-148675e01b30

Commits

Reviewing files that changed from the base of the PR and between f09b399 and d4969b4.

Files selected for processing: 1

• sdm/models/base/model.py

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

[aw471]

@CodeRabbit full review

[coderabbitai]

Action performed

Full review finished.

[aw471]

/ok to test 54d5bab

[coderabbitai]

• Fix comments: Fix CodeRabbit comments on this PR

[aw471]

/ok to test d4969b4