[acisseJZhong]

Summary

CP linear-attention routing currently copies GPU document offsets and token permutations back to the CPU to construct routing plans. Reuse the microbatch's CPU positions during preprocessing to avoid this synchronization for contiguous and head-tail CP layouts.

• Pass CPU input metadata through training, validation, and decoder preprocessing. Build document boundaries and per-rank token fragments once per microbatch for KDA/GDN routing.
• Derive head-tail fragments with the same upstream partitioner on CPU. PTRR still reads its GPU-dependent permutation back once, instead of once per rank.
• Cache up to 16 CPU routing plans, keyed by document boundaries, all ranks' fragments, and the current rank. Each batch creates its own device routing tensors. Callers without CPU metadata retain the existing readback path.

Test plan

• 72 focused tests passed, including 17 new CPU cases covering exact routing equivalence, cache reuse/invalidation, independent device tensors, and avoidance of metadata readback. Four failures also reproduce on the unmodified base (0980e1cb1c9a) due to missing local PyTorch/SPMD APIs.
• Two-H100 BF16 numerical comparison passed for KDA/GDN with contiguous/head-tail CP. With seed 42 and deterministic algorithms enabled, outputs, loss, gradient norms, all gradients, and updated parameters are bitwise identical across three AdamW updates. This covers convolution halo exchange and distributed linear attention; it is a focused kernel/routing comparison, not full-model training. The new test is selected by the existing multi-GPU unit-test workflow.
• Non-Pyrefly pre-commit hooks passed. The pinned Pyrefly comparison reports no new diagnostics in changed files, but the full type check and trainer tests remain blocked by local dependency/API mismatches.

Local tests used the repository-pinned AttentionGym 0.0.17 in a temporary import path and postponed evaluation of one incompatible type annotation. No computation or collective implementation was replaced. End-to-end throughput has not been measured.