fix(dynamicemb): retain CUB scratch tensors through selection

[lidongdongbuaa]

The CUB selection helpers take a raw pointer from a temporary at::empty(...) tensor. The temporary is destroyed at the end of the assignment statement, before the following DeviceSelect::Flagged call uses the pointer.

Keep the owning tensor in scope through the CUB call in select_async, select_index_async, and flagged_compact_impl. Allocation and launch use the existing current-stream path; no synchronization or API change is introduced.

Validation on upstream commit 1ed1ab716cdbae8d103b1007d8beb40c86c4a18a, with the identical index_calculation.h present in the current base (97062d9):

• H100 80GB, Python 3.11, PyTorch 2.11.0, CUDA 13.2.
• Before: the existing test_unique_op.py -k flagged_compact tests under Compute Sanitizer reported an invalid global write in CUB ScanTileState::InitializeStatus (18 reported errors).
• After rebuilding with this patch: all 10 flagged-compaction cases pass with zero sanitizer errors; all 19 tests in test_unique_op.py pass.
• Environment: PYTORCH_NO_CUDA_MEMORY_CACHING=1 CUDA_LAUNCH_BLOCKING=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1.
• Changed-file pre-commit checks pass.

Reproduce the sanitizer check:

PYTORCH_NO_CUDA_MEMORY_CACHING=1 CUDA_LAUNCH_BLOCKING=1 \
  PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \
  compute-sanitizer --tool memcheck --error-exitcode 86 \
  python -m pytest corelib/dynamicemb/test/test_unique_op.py \
  -q -k flagged_compact --maxfail=1

This fixes the independently reproduced scratch-storage lifetime defect; it is not a claim that every illegal-memory-access failure in an embedding workload has this cause. The entire current-main library and distributed training suite were not rebuilt or rerun.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge and correctly addresses the temporary scratch-storage lifetime defect without changing APIs or stream behavior.

Summary

This PR fixes a CUDA scratch-storage lifetime defect in DynamicEmb’s CUB selection helpers.

• Retains each at::Tensor that owns temporary CUB storage through the corresponding DeviceSelect::Flagged launch.
• Applies the same lifetime correction to value selection, index selection, and flagged compaction.
• Preserves the existing device, stream, synchronization, and API behavior.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Query CUB scratch size] --> B[Allocate scratch tensor]
    B --> C[Obtain scratch pointer]
    C --> D[Enqueue DeviceSelect::Flagged]
    D --> E[Scratch tensor leaves scope]
    E --> F[Allocator defers reuse until stream work completes]

Loading

_{Reviews (1) · Last reviewed commit: "fix(dynamicemb): retain CUB scratch tens..."}