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 parallelised

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 can compose its local chunks into one (A, B) pair with a local scan, all-gather the D pairs, take the exclusive prefix to get its incoming state, and replay locally. A and B are 128x128 for every published Qwen 3.5 size, so the gathered volume is small.

Status of my implementation

I have this working. I'd like guidance before proposing it upstream.

• Exact to 1.1e-08 against the stock sequential scan.
• A 1,048,576-token context trains end to end on 4 v5p chips, and the same configuration trains under Pathways with enable_single_controller.
• The local composition needs jax.checkpoint on the scan body. Without it the autodiff residuals reach 103 GB at sequence 262,144.

Two things I'm unsure about. That's why this is an issue and not a pull request.

1. It adds a whole module instead of changing a few lines. Where should it live, and how should the context-parallel axis reach jax_chunk_gated_delta_rule?
2. I only tested the shapes I needed. No test covers this path today, so a pull request has to bring its own. Let me know what they should cover first.

I can open a pull request. Or if someone else owns this area, I'll hand over the derivation and the parity test.

cc @mmcsa