fix(xtoken): return cached tuple in get_topk_projection; align HSDP test

get_topk_projection's cache-miss path returned a freshly-built
(indices, likelihoods) tuple instead of the one stored in
_TOPK_PROJECTION_CACHE, so the first (miss) call and subsequent (hit)
calls handed back different tuple objects. This broke
test_repeat_call_hits_cache's identity assertion (a is b). Store and
return the same tuple.

Also update test_dtensor_hsdp_dispatches_distinct_batches: Policy.train
now threads check_dim_skip_keys through common_kwargs (added in this PR
for cross-tokenizer skip-keys plumbing), so the expected
run_all_workers_sharded_data call must include 'check_dim_skip_keys': None.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Adithya Hanasoge <adithyakvh@gmail.com>

Author: adithyakvh

nemo_rl/algorithms/x_token/loss_utils.py

@@ -425,8 +425,9 @@ def get_topk_projection(
         )
     indices = data["indices"].long().to(device)
     likelihoods = data["likelihoods"].float().to(device)
-    _TOPK_PROJECTION_CACHE[key] = (indices, likelihoods)
-    return indices, likelihoods
+    result = (indices, likelihoods)
+    _TOPK_PROJECTION_CACHE[key] = result
+    return result
 
 
 # Process-local cache. Keyed by every input that affects the partition:

tests/unit/models/policy/test_policy_validation.py

@@ -327,6 +327,7 @@ def test_dtensor_hsdp_dispatches_distinct_batches(
             "eval_mode": False,
             "gbs": 32,
             "mbs": 4,
+            "check_dim_skip_keys": None,
         },
     )