mixle.experimental.context_parallel_spine module

E8: context parallelism for SlidingWindowSpine.

See notes/designs/E8.md for the full design (why torch-native CP doesn’t suffice – SlidingWindowSpine.step computes attention by hand, there is no single scaled_dot_product_attention call for the native CP context manager to intercept – and why F1’s own CP (tensor_pipeline_context_parallel.py) doesn’t drop in as-is – it shards a whole-block STATELESS forward, not a carried-state streaming step).

This module reuses F1’s proven PATTERN (chunk the KV axis, all-gather once per attention op, reconstruct with an explicit offset mask, verify by exact match against a dense reference) rather than its entry points, purpose-built for ContextMechanism.step’s carried-state shape: each call shards the CURRENT step’s KV axis (cache ++ chunk, not a whole logical stream) across cp_size simulated ranks.

Same scope cut as F1’s CP, for the same reason (this environment has <=1 real GPU): eager chunked-all-gather CP, exact and testable at small scale, NOT incremental ring attention (no rank-to-rank streaming overlap). See the design note’s Risks section.

RoPE is applied PER SHARD, using each shard’s own absolute key_positions_shard, BEFORE the gather (not once on the full gathered K after). This is deliberately closer to what a real distributed rank would do (a rank never needs another rank’s raw K to compute its own RoPE) at the cost of being only mathematically – not bit-exactly – identical to the single-device code path (float non-associativity of RoPE-per-slice-then-concat vs. RoPE-on-the-full-array). See notes/designs/E8.md’s Risks section for the fallback (gather-raw-K-then-RoPE-once) if this tolerance had not held in practice; the test suite (mixle/tests/context_parallel_spine_test.py) measures this directly and it DOES hold, at low-1e-6 absolute divergence even under real multi-chunk SGD training, comfortably inside the documented rtol=1e-4 tolerance (F1’s own precedent).

Cache convention: SlidingWindowSpine.step caches raw K/V together with explicit absolute positions. RoPE is applied once, after the current raw keys and raw cached keys are assembled for attention. Both dense and context-parallel paths use this convention, preventing older keys from being rotated again on every streamed chunk.

Unlike F1’s TP (n_head % tp_size == 0) and CP (block % cp_size == 0), which both fail fast on uneven division, this module’s KV axis length (cache_len + t) varies chunk to chunk and is not required to divide evenly by cp_size – torch.chunk already degrades gracefully (size-balanced, possibly-uneven chunks, or fewer than cp_size chunks if the axis is shorter than cp_size) rather than erroring, and this module relies on exactly that behavior instead of adopting F1’s stricter validate-and-reject posture. See the design note’s Risks section for why this is a deliberate difference, not an oversight.