Shard the GatedDeltaNet sequence under ici_context_parallelism

[WandLZhang]

Fixes #4932.

Two pspecs pinned the GDN sequence axis to None, and the inter-chunk recurrence was a sequential lax.scan. Together they mean ici_context_parallelism can't shard the GDN sequence at all, so neither 256K nor 1M is reachable at any chip count.

The recurrence is affine in the state:

h' = (exp(g_last)·I − k_gᵀ·w)·h + k_gᵀ·u  =  A·h + B

so it evaluates as a two-pass associative scan. Each device folds its local chunks into one (A, B) pair with a lax.scan — not associative_scan, which would materialize the running composition per chunk, 17 GB per device at a million tokens — then an all-gather exchanges the pairs and an exclusive prefix gives each device its incoming state. A and B are 128x128 at every published Qwen 3.5 size, so the collective doesn't grow with the model.

The second pspec fix matters as much as the scan. qkvz_pspec had the sequence axis hardcoded replicated, which told XLA to gather the full sequence onto every device: at ctx=2, seq 32,768, six live bf16[2,32768,16,512] buffers of 1.07 GB each. Found by diffing XLA buffer assignment between ctx=1 and ctx=2.

Measured at 1,048,576 tokens on 64 v5p chips: 513 tokens/s/chip, 27.6% MFU, loss 6.736 to 4.335.

Interaction with #4348

The two can't both be live — the fused kernel returns above the state composition, so use_pallas is forced false whenever a context axis is set and sequence-sharded runs take the XLA scan. @bzantium verified that fallback is loss-preserving: 6e-6 relative on the output, 2e-5 on final state. Disjoint rather than stacked, so this goes against main and #4348 rebases on top.

Size

models/gdn_cp.py is new, 79 lines. models/qwen3.py is +42/-4. Applies to a pristine checkout of main with no dependency on #4348 or on my other open PRs.

Tests

ici_context_parallelism=4 at sequence 32,768 and ici_context_parallelism=8 at 262,144, both training with loss descending, on v5p and v6e. 1,048,576 tokens on 64 v5p chips at 513 tokens/s/chip. Memory is flat in ctx at a fixed local shard: 30.5 GB at ctx=2 against 30.7 GB at ctx=4 on a 2,048-token shard, so the sizing extrapolates. pyink clean.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces context-parallel evaluation for the GatedDeltaNet (GDN) inter-chunk recurrence to optimize memory usage during long-sequence training. It adds a new module gdn_cp.py to handle the two-pass affine composition of local chunks and cross-device state communication, and updates qwen3.py to support sequence sharding over context-parallel axes. A critical runtime issue was identified in gdn_cp.py where cp_axis is passed as a tuple of strings, but lax.axis_index only accepts a single string. A detailed code suggestion is provided to generalize incoming_state to support both single strings and tuples of strings, correctly handling multi-dimensional device grids.

[gemini-code-assist]

In qwen3.py, cp_axis is passed as a tuple of strings (e.g., ("context",) or ("context", "context_usp_ulysses")). However, lax.axis_index only accepts a single string axis name, and calling it with a tuple will raise a TypeError at runtime. Additionally, if there are multiple context parallel axes, lax.all_gather and lax.associative_scan need to be handled sequentially and flattened to correctly compute the 1D associative scan across the multi-dimensional device grid.

This suggestion generalizes incoming_state to support both a single string and a tuple of strings for cp_axis, gathers sequentially, flattens the gathered dimensions in the correct major-to-minor sequence order, and computes the correct flat device index.

def incoming_state(A_loc, B_loc, h_init, cp_axis):
  """State entering this device, plus the final state after all devices.

  Gathering D pairs of small matrices is the only cross-device traffic in the
  scheme. Must be called inside a shard_map over `cp_axis`.
  """
  if isinstance(cp_axis, str):
    cp_axes = (cp_axis,)
  else:
    cp_axes = cp_axis

  A_all = A_loc
  B_all = B_loc
  for axis in cp_axes:
    A_all = lax.all_gather(A_all, axis, axis=0, tiled=False)
    B_all = lax.all_gather(B_all, axis, axis=0, tiled=False)

  n_axes = len(cp_axes)
  if n_axes > 1:
    # Reverse the gathered dimensions to match the major-to-minor order of cp_axes
    perm = list(range(n_axes))[::-1] + list(range(n_axes, A_all.ndim))
    A_all = jnp.transpose(A_all, perm)
    B_all = jnp.transpose(B_all, perm)

  # Flatten the gathered dimensions into a single axis at 0
  orig_shape_A = A_all.shape
  orig_shape_B = B_all.shape
  num_devices = 1
  for axis in cp_axes:
    num_devices *= lax.psum(1, axis)

  A_all = A_all.reshape((num_devices,) + orig_shape_A[n_axes:])
  B_all = B_all.reshape((num_devices,) + orig_shape_B[n_axes:])

  A_cum, B_cum = lax.associative_scan(compose, (A_all, B_all), axis=0)

  idx = 0
  stride = 1
  for axis in reversed(cp_axes):
    idx += lax.axis_index(axis) * stride
    stride *= lax.psum(1, axis)

  prev = jnp.maximum(idx - 1, 0)
  carried = jnp.matmul(A_cum[prev], h_init, precision=_PREC) + B_cum[prev]
  h_in = jnp.where(idx == 0, h_init, carried)
  final_h = jnp.matmul(A_cum[-1], h_init, precision=_PREC) + B_cum[-1]
  return h_in, final_h