cache offsets / lengths on updates

[murphymatt]

Summary

Currently, in the UVM KJT pools, on every lookup, we recompute offset values for all lengths of the lookup ids. This can be extremely expensive if the pool size is very large and if the id lookup size is large. Furthermore, given that these offsets should only change upon updates to the pool, it should be unnecessary to recompute these values on every read path. Instead, we can compute these offsets on the pool's write path, and reference the offsets at the correct lookup indices, to significantly cut down on lookup latency.

Testing

torchx run -s local_cwd utils.python --script torchrec/modules/tests/test_kjt_pool_lookup.py

torchx 2025-06-24 16:02:22 INFO     Tracker configurations: {}
torchx 2025-06-24 16:02:22 INFO     Log directory not set in scheduler cfg. Creating a temporary log dir that will be deleted on exit. To preserve log directory set the `log_dir` cfg option
torchx 2025-06-24 16:02:22 INFO     Log directory is: /tmp/torchx_9keqq27r
local_cwd://torchx/torchx_utils_python-q6bc5grqt7hv3c
torchx 2025-06-24 16:02:22 INFO     Waiting for the app to finish...
python/0 torch.cuda.is_available()=True
python/0 ..
python/0 ----------------------------------------------------------------------
python/0 Ran 2 tests in 0.458s
python/0
python/0 OK
torchx 2025-06-24 16:02:31 INFO     Job finished: SUCCEEDED

[murphymatt]

Opening for review

[TroyGarden]

@murphymatt thanks for the PR. could you please address the test failures

=========================== short test summary info ============================
FAILED torchrec/modules/tests/test_keyed_jagged_tensor_pool.py::KeyedJaggedTensorPoolTest::test_conflict - AttributeError: 'TensorJaggedIndexSelectLookup' object has no attribute '_tbe_state'
FAILED torchrec/modules/tests/test_keyed_jagged_tensor_pool.py::KeyedJaggedTensorPoolTest::test_empty_lookup - AttributeError: 'TensorJaggedIndexSelectLookup' object has no attribute '_tbe_state'
FAILED torchrec/modules/tests/test_keyed_jagged_tensor_pool.py::KeyedJaggedTensorPoolTest::test_input_permute - AttributeError: 'TensorJaggedIndexSelectLookup' object has no attribute '_tbe_state'
FAILED torchrec/modules/tests/test_keyed_jagged_tensor_pool.py::KeyedJaggedTensorPoolTest::test_update_lookup - AttributeError: 'TensorJaggedIndexSelectLookup' object has no attribute '_tbe_state'