GatedDeltaNet does not support ici_context_parallelism

[WandLZhang]

Setting ici_context_parallelism > 1 doesn't shard the sequence inside Qwen3_5GatedDeltaNet. Two separate things prevent it.

1. Three pspecs pin the sequence axis to replicated. In models/qwen3.py, qkv_pspec, g_beta_pspec and qkvz_pspec are built as (KV_BATCH, None, KV_HEAD, None). The None in the sequence position tells XLA to gather the full sequence onto every device, so a with_sharding_constraint undoes the sharding the mesh asks for.

2. The inter-chunk recurrence is a sequential scan. jax_chunk_gated_delta_rule runs lax.scan over chunks. A sequential scan can't be split across devices, because each chunk needs the state the previous chunk produced.

The recurrence is affine, so it can be parallelized

The GatedDeltaNet inter-chunk step is

h_new = exp(g_last) * h + k_g^T (u - w h)
      = (exp(g_last) I - k_g^T w) h + k_g^T u
      = A h + B

A and B don't depend on h. Affine maps compose associatively:

(A2, B2) o (A1, B1) = (A2 A1, A2 B1 + B2)

So each device composes its local chunks into one (A, B) pair with a local scan, the pairs compose across devices with a log-depth prefix scan over ppermute, and each device replays locally from its incoming state. A and B are 128x128 for every published Qwen 3.5 size, so the cross-device payload is small.

Status

Filed as #4968.

• 1,048,576 tokens train end to end on 256 v5p at ctx=256, loss 12.923 to 10.583, 485.5 tokens/s/chip. The same configuration trains under Pathways with enable_single_controller.
• tests/unit/gdn_cp_test.py checks the composition against a sequential single-device reference on a forced 8-device CPU mesh, with a negative control. The whole path matched the stock sequential scan to 1.1e-08 in the earlier all-gather form.
• The local composition needs jax.checkpoint on the scan body. Without it the autodiff residuals reach 103 GB at sequence 262,144.
• An all-gather across devices is O(D) per layer, about 38 GB across 48 GatedDeltaNet layers at ctx=256. The prefix scan is log2(D) exchanges holding one pair.

The module is models/gdn_cp.py. The axis reaches jax_chunk_gated_delta_rule as a cp_axis keyword built from whichever context knob is set.

Running any of this needs per_device_batch_size = 1/ici_context_parallelism. See #4933.

cc @mmcsa