From 516775b01fc079137b76865eb1ef2fc927ccf1da Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Thu, 3 Sep 2026 04:03:09 +0000 Subject: [PATCH] fix(attention test): give the mask tests the axis rules the cp gate reads These three tests build a 2- or 4-device context mesh and assert that the mask is derived from the load-balanced `segment_positions`. That path is guarded by `_load_balanced_context_parallel()`, which sizes cp by resolving `activation_q_length` against the *ambient* flax rules - deliberately, so eval under `logical_axis_rules_for_eval` cannot claim an order the loader never applied. The tests never enter an `axis_rules` context, so the logical name resolves to `P(None)`, cp_size comes out as 1, load balancing is skipped and the op returns a plain index-order causal mask. Declare the `activation_q_length` rule and wrap the op in it, as the rest of the suite does. Also set `ulysses_context_sharding`, which the cp-size computation reads once the query axis actually resolves. Only a host with 2+ visible devices runs these (they skip otherwise), so CI has not been executing them. --- tests/unit/attention_test.py | 111 +++++++++++++++++++---------------- 1 file changed, 61 insertions(+), 50 deletions(-) diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 2e51d79725..2c1ef354d6 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -679,6 +679,8 @@ def test_dot_product_mask_uses_original_load_balanced_positions(self): causal_block_size=4, context_parallel_load_balance=True, context_sharding="context", + ulysses_context_sharding="context_usp_ulysses", + logical_axis_rules=[["activation_q_length", ["context"]]], shard_mode="auto", debug_sharding=False, eval_interval=-1, @@ -687,23 +689,24 @@ def test_dot_product_mask_uses_original_load_balanced_positions(self): if len(devices) < 2: self.skipTest("Need at least 2 devices to test chunk mask") mesh = Mesh(devices[:2], ["context"]) - op = AttentionOp( - config=config, - num_query_heads=1, - num_kv_heads=1, - max_target_length=sequence_length, - mesh=mesh, - attention_kernel="dot_product", - attention_type=AttentionType.BLOCK_DIFFUSION, - ) + with nn_partitioning.axis_rules(config.logical_axis_rules): + op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=sequence_length, + mesh=mesh, + attention_kernel="dot_product", + attention_type=AttentionType.BLOCK_DIFFUSION, + ) - mask = op.generate_attention_mask( - query, - key, - segment_ids, - MODEL_MODE_TRAIN, - segment_positions=positions, - ) + mask = op.generate_attention_mask( + query, + key, + segment_ids, + MODEL_MODE_TRAIN, + segment_positions=positions, + ) expected = np.asarray(positions[0])[:, None] // 4 >= np.asarray(positions[0])[None, :] // 4 np.testing.assert_array_equal(np.asarray(mask == 0.0)[0, 0, 0], expected) @@ -1305,6 +1308,7 @@ def test_dot_product_local_mask_uses_segment_positions(self): config = types.SimpleNamespace( context_parallel_load_balance=True, context_sharding="context", + ulysses_context_sharding="context_usp_ulysses", using_pipeline_parallelism=False, logical_axis_rules=[["segment_ids_batch", ["context"]]], shard_mode="auto", @@ -1321,24 +1325,27 @@ def test_dot_product_local_mask_uses_segment_positions(self): query = jnp.zeros((1, seq_len, 1, 128)) key = jnp.zeros((1, seq_len, 1, 128)) decoder_segment_ids = jnp.ones((1, seq_len), dtype=jnp.int32) - op = AttentionOp( - config=config, - num_query_heads=1, - num_kv_heads=1, - max_target_length=seq_len, - mesh=mesh, - attention_kernel="dot_product", - attention_type=AttentionType.LOCAL_SLIDING, - sliding_window_size=sliding_window_size, - ) + # Only the query-length rule: `segment_ids_batch` -> context would shard this batch of 1 + # across the 4-way context mesh. + with nn_partitioning.axis_rules([["activation_q_length", ["context"]]]): + op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=seq_len, + mesh=mesh, + attention_kernel="dot_product", + attention_type=AttentionType.LOCAL_SLIDING, + sliding_window_size=sliding_window_size, + ) - mask = op.generate_attention_mask( - query, - key, - decoder_segment_ids, - MODEL_MODE_TRAIN, - segment_positions=positions, - ) + mask = op.generate_attention_mask( + query, + key, + decoder_segment_ids, + MODEL_MODE_TRAIN, + segment_positions=positions, + ) expected_mask = np.zeros((seq_len, seq_len), dtype=np.bool_) for r, q_pos in enumerate(np.asarray(positions[0])): @@ -1352,6 +1359,7 @@ def test_dot_product_chunk_mask_uses_segment_positions(self): config = types.SimpleNamespace( context_parallel_load_balance=True, context_sharding="context", + ulysses_context_sharding="context_usp_ulysses", using_pipeline_parallelism=False, logical_axis_rules=[["segment_ids_batch", ["context"]]], shard_mode="auto", @@ -1368,24 +1376,27 @@ def test_dot_product_chunk_mask_uses_segment_positions(self): query = jnp.zeros((1, seq_len, 1, 128)) key = jnp.zeros((1, seq_len, 1, 128)) decoder_segment_ids = jnp.ones((1, seq_len), dtype=jnp.int32) - op = AttentionOp( - config=config, - num_query_heads=1, - num_kv_heads=1, - max_target_length=seq_len, - mesh=mesh, - attention_kernel="dot_product", - attention_type=AttentionType.CHUNK, - chunk_attn_window_size=chunk_size, - ) + # Only the query-length rule: `segment_ids_batch` -> context would shard this batch of 1 + # across the 4-way context mesh. + with nn_partitioning.axis_rules([["activation_q_length", ["context"]]]): + op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=seq_len, + mesh=mesh, + attention_kernel="dot_product", + attention_type=AttentionType.CHUNK, + chunk_attn_window_size=chunk_size, + ) - mask = op.generate_attention_mask( - query, - key, - decoder_segment_ids, - MODEL_MODE_TRAIN, - segment_positions=positions, - ) + mask = op.generate_attention_mask( + query, + key, + decoder_segment_ids, + MODEL_MODE_TRAIN, + segment_positions=positions, + ) expected_mask = np.zeros((seq_len, seq_len), dtype=np.bool_) for r, q_pos in enumerate(np.asarray(positions[0])):