Attention: per-layer local_context_parallel_strategy and halo KV exchange for LOCAL_SLIDING by csgoogle · Pull Request #5608 · AI-Hypercomputer/maxtext

csgoogle · 2026-10-08T08:40:19Z

Summary

Adds local_context_parallel_strategy ("", "all_gather", "halo") so hybrid models with both GLOBAL and LOCAL_SLIDING attention layers (such as Gemma4-26B) can override context_parallel_strategy on LOCAL_SLIDING layers, and implements a ppermute-based halo KV exchange in AttentionOp.tpu_halo_flash_attention for LOCAL_SLIDING context parallelism.

Per-layer CP strategy (local_context_parallel_strategy): AttentionOp.__init__ resolves self.context_parallel_strategy using local_context_parallel_strategy when self.attention_type == AttentionType.LOCAL_SLIDING, allowing GLOBAL layers to use ulysses, ring, or all_gather while LOCAL_SLIDING layers use halo or all_gather.
Separated Halo Flash Attention (AttentionOp.tpu_halo_flash_attention): Instead of all-gathering and unpermuting the full [B, S, H_kv, D] K/V sequence across the context axis on LOCAL_SLIDING layers, each CP rank exchanges only the block-aligned sliding-window tail (halo_pad_width) of the preceding chunk via lax.ppermute. Supports both contiguous (context_parallel_load_balance=False) and DUAL_CHUNK_SWAP (context_parallel_load_balance=True) layouts without splitting the batch dimension.

Related PRs

Gemma4 26B MoE on TPU v6e: MaxText layer/model optimizations (one commit per change) #5572 (MaxText layer/model optimizations)
MoE: fused TensorCore Pallas permute/combine kernels and split-expert wi layout for TPU v6e #5573 (Fused TensorCore Pallas permute/combine kernels)
Megablox: transposed-RHS gmm_v2 kernel for the backward dlhs #5574 (Megablox transposed-RHS gmm_v2 backward dlhs kernel)

gemini-code-assist

Code Review

This pull request introduces support for a per-layer context parallel (CP) strategy and a new 'halo' sliding-window CP attention mechanism in MaxText. It adds configuration options, validation logic, and the core implementation of the halo strategy using ppermute communication. Comprehensive unit tests are also added. The review comments identify two important improvements: guarding against a potential TypeError when validating sliding_window_size if it is None, and ensuring XLA_FLAGS is set before importing JAX in the test file to guarantee proper device initialization.

gemini-code-assist · 2026-10-08T08:42:03Z

+    if (
+        context_parallel_strategy == "halo" or local_context_parallel_strategy == "halo"
+    ) and self.sliding_window_size <= 0:
+      raise ValueError("Halo context parallelism requires sliding_window_size > 0.")


If self.sliding_window_size is None, the comparison self.sliding_window_size <= 0 will raise a TypeError. Guarding against None values ensures robust configuration validation.

Suggested change

if (

context_parallel_strategy == "halo" or local_context_parallel_strategy == "halo"

) and self.sliding_window_size <= 0:

raise ValueError("Halo context parallelism requires sliding_window_size > 0.")

if (

context_parallel_strategy == "halo" or local_context_parallel_strategy == "halo"

) and (self.sliding_window_size is None or self.sliding_window_size <= 0):

raise ValueError("Halo context parallelism requires sliding_window_size > 0.")

gemini-code-assist · 2026-10-08T08:42:04Z

+import os
+import types
+from absl.testing import absltest
+from absl.testing import parameterized
+from flax import linen as nn
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+from jax.sharding import NamedSharding
+from jax.sharding import PartitionSpec as P
+from maxtext.common.common_types import AttentionType
+from maxtext.common.common_types import MODEL_MODE_TRAIN
+from maxtext.configs import pyconfig
+from maxtext.layers import attention_op as attention_op_lib
+from maxtext.utils import max_utils
+import numpy as np
+
+os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=8")


Setting XLA_FLAGS via os.environ after importing jax can be unreliable because JAX may initialize its backend and devices upon import. To guarantee that the CPU device count is correctly set to 8, os.environ should be configured before import jax.

Suggested change

import os

import types

from absl.testing import absltest

from absl.testing import parameterized

from flax import linen as nn

from flax import nnx

import jax

import jax.numpy as jnp

from jax.sharding import Mesh

from jax.sharding import NamedSharding

from jax.sharding import PartitionSpec as P

from maxtext.common.common_types import AttentionType

from maxtext.common.common_types import MODEL_MODE_TRAIN

from maxtext.configs import pyconfig

from maxtext.layers import attention_op as attention_op_lib

from maxtext.utils import max_utils

import numpy as np

os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=8")

import os

os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=8")

import types

from absl.testing import absltest

from absl.testing import parameterized

from flax import linen as nn

from flax import nnx

import jax

import jax.numpy as jnp

from jax.sharding import Mesh

from jax.sharding import NamedSharding

from jax.sharding import PartitionSpec as P

from maxtext.common.common_types import AttentionType

from maxtext.common.common_types import MODEL_MODE_TRAIN

from maxtext.configs import pyconfig

from maxtext.layers import attention_op as attention_op_lib

from maxtext.utils import max_utils

import numpy as np

csgoogle · 2026-10-08T10:05:23Z

 ### Determine if we want to use load balance for context parallelism
 context_parallel_load_balance: true
-context_parallel_strategy: "all_gather" # "all_gather", "ring", "ulysses", or "usp"
+context_parallel_strategy: "all_gather" # "all_gather", "ring", "ulysses", "usp", or "halo"


We can remove halo here

csgoogle · 2026-10-08T12:35:34Z

 import numpy as np
-from tokamax._src.ops.attention import base as tokamax_attention_base
-from tokamax._src.ops.attention import pallas_triton as tokamax_pallas_triton
-from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_kernel as tokamax_splash_kernel
-from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_mask as tokamax_splash_mask
+try:
+  from tokamax._src.ops.attention import base as tokamax_attention_base
+  from tokamax._src.ops.attention import pallas_triton as tokamax_pallas_triton
+  from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_kernel as tokamax_splash_kernel


csgoogle · 2026-10-08T12:36:43Z

-      raise ValueError("context_parallel_strategy must be one of 'all_gather', 'ring', 'ulysses', or 'usp'.")
+    if context_parallel_strategy not in ("all_gather", "ring", "ulysses", "usp", "halo"):
+      raise ValueError("context_parallel_strategy must be one of 'all_gather', 'ring', 'ulysses', 'usp', or 'halo'.")
    self.context_parallel_strategy = context_parallel_strategy


make it simple, like I mentioned, keep the local only to halo..

csgoogle · 2026-10-08T17:18:09Z

      single_head_mask = mask  # tokamax now just uses a single mask and assumes broadcast to all heads
      if self.config.use_max_logit_estimate > 0:
        sa_config = dataclasses.replace(sa_config, max_logit_const=self.config.use_max_logit_estimate)
+      if use_halo and jax.default_backend() == "cpu":


csgoogle · 2026-10-08T17:19:15Z

        return splash_kernel

-      segment_axis_names_splash_kernel = self._logical_to_mesh_axes((Q_LENGTH,))
+      segment_axis_names_splash_kernel = (


better seperate the halo attention, it will cause less changes and will be more clear and readable

csgoogle · 2026-10-08T17:20:11Z

  )
  use_dq_carry = config.dq_reduction_steps == 3
-  dq_accum = jnp.zeros((3, *q.shape), dtype=jnp.float32) if use_dq_carry else jnp.zeros(q.shape, dtype=jnp.float32)
+  dq_accum = jnp.zeros((3, *q.shape), dtype=q.dtype) if use_dq_carry else jnp.zeros(q.shape, dtype=jnp.float32)


should be fp32 always, revert

csgoogle · 2026-10-08T17:21:59Z

      mask_value=mask_value,
      is_mqa=is_mqa,
-      config=config,
+      config=dataclasses.replace(config, residual_checkpoint_name=None),


why is this needed, if needed, do a proper fix.

…ange for LOCAL_SLIDING Adds local_context_parallel_strategy ('', 'all_gather', 'halo') so hybrid models like Gemma4 can override context_parallel_strategy on LOCAL_SLIDING layers (e.g., global layers using ring/ulysses/all_gather while local layers use halo). Separates the halo KV exchange path into AttentionOp.tpu_halo_flash_attention.

gemini-code-assist Bot reviewed Oct 8, 2026

View reviewed changes

csgoogle commented Oct 8, 2026

View reviewed changes

csgoogle force-pushed the g4-v6e128-attention branch 3 times, most recently from d34004f to 89f93af Compare October 8, 2026 17:46

csgoogle changed the title ~~Attention: per-layer Halo context parallelism for LOCAL_SLIDING layers & Tokamax Ring Attention fixes~~ Attention: per-layer local_context_parallel_strategy and halo KV exchange for LOCAL_SLIDING Oct 8, 2026

csgoogle force-pushed the g4-v6e128-attention branch from 89f93af to 825cbdd Compare October 8, 2026 18:03

Navigation Menu

Repository navigation

Attention: per-layer local_context_parallel_strategy and halo KV exchange for LOCAL_SLIDING - #5608

Attention: per-layer local_context_parallel_strategy and halo KV exchange for LOCAL_SLIDING#5608
csgoogle wants to merge 1 commit into
AI-Hypercomputer:mainfrom
csgoogle:g4-v6e128-attention

csgoogle commented Oct 8, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Oct 8, 2026

Uh oh!

gemini-code-assist Bot Oct 8, 2026

Uh oh!

csgoogle Oct 8, 2026

Uh oh!

csgoogle Oct 8, 2026

Uh oh!

csgoogle Oct 8, 2026

Uh oh!

csgoogle Oct 8, 2026

Uh oh!

csgoogle Oct 8, 2026

Uh oh!

csgoogle Oct 8, 2026

Uh oh!

csgoogle Oct 8, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

csgoogle commented Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Related PRs

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

csgoogle Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

csgoogle Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

csgoogle Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

csgoogle Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

csgoogle Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

csgoogle Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

csgoogle Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

csgoogle commented Oct 8, 2026 •

edited

Loading