[aw471]

We want to make things pytorch native if there is a meaningful lift.

Summary

Use native PyTorch string-category lookup on CUDA in eager and compiled execution. Keep the existing Arrow path on CPU.

Description of the change

Compact byte keys find candidates; exact comparisons preserve missing values and rightmost duplicates. The CUDA path supports unknown length metadata and views with unused backing storage, without an input string-length limit or registered join. Depends on #1274.

Validation

Focused rebase checks on one RTX PRO 6000 Blackwell GPU. Timings include key packing and matching; these are isolated category-lookup measurements, not end-to-end model/fit/predict latency. Five warmups and twenty balanced, synchronized-wall samples per arm.

Lookup profile	cuDF ms	Native eager ms	Native compiled ms	Eager speedup	Compiled speedup
3 column pairs: 35/6/4 categories, 128-byte strings	5.249	4.564	0.663	1.15x	7.91x
35 / 35 categories, visible 128-byte strings, unused 50 KB tail	5.819	1.689	0.370	3.45x	15.73x

Peak logical working GPU memory, above the same resident inputs:

Lookup profile	cuDF KiB	Native eager KiB	Native compiled KiB	Eager change	Compiled change
3 column pairs: 35/6/4 categories, 128-byte strings	35.69	17.98	17.74	-49.6%	-50.3%
35 / 35 categories, visible 128-byte strings, unused 50 KB tail	29.97	17.98	23.52	-40.0%	-21.5%

Native eager produced zero Dynamo graphs. Compiled lookup used Inductor with fullgraph=True, dynamic=False: one graph and zero breaks in each profile. The first compiled call for the three-column case took 128.7 seconds, excluded from warm timings. Input conversion and query-row code gathering are also excluded.

Additional details

PyTorch 2.14.0+cu130, cuDF 26.10.02, RMM 26.10.00. Memory measures simultaneous live allocation using one shared RMM tracker, including retained buffers and outputs. Default-allocator latency is measured separately. No special allocator or compiler setting is required by the implementation. These steady-state memory figures exclude compilation-time allocations, host compiler RAM and allocator-reserved memory. Default allocation granularity can change allocated-byte peaks, so logical memory is not a claim of universally lower default-allocator memory.

The tables are from the two fresh GPU reruns with #1274’s metadata changes present. #1275 remains stacked on #1274. Subsequent parent updates changed documentation/equivalent metadata expressions and removed an unused packing helper; the measured native lookup implementation is unchanged. Restored stacked CPU checks: 33 passed, 25 skipped; all pre-commit hooks passed.

Existing stack testing found a dynamic=True strided-StringTensor tracing assertion even with lambda x: x._data.sum() before lookup logic. Large hash-collision groups can require extra exact comparisons; the measurements are not a universal performance guarantee.

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

[codecov]

Codecov Report

❌ Patch coverage is 90.90909% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.01%. Comparing base (73602df) to head (d86e15f).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
sdm/processing/categorical/_string_lookup.py	91.47%	11 Missing ⚠️
sdm/processing/categorical/align.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1275      +/-   ##
==========================================
+ Coverage   82.90%   83.01%   +0.11%     
==========================================
  Files         177      178       +1     
  Lines       11455    11577     +122     
==========================================
+ Hits         9497     9611     +114     
- Misses       1958     1966       +8     

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