fix: SBHD reorder skip uses original shape instead of swapped tensor

[andrewwhitecdw]

This PR addresses the following issue in tests/jax/test_distributed_fused_attn.py: SBHD reorder skip uses original shape instead of swapped tensor.

Changes

• tests/jax/test_distributed_fused_attn.py: SBHD reorder skip uses original shape instead of swapped tensor.

Details

--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -1,8 +1,8 @@
-        if qkv_format == QKVFormat.SBHD:
-            tensor = tensor.swapaxes(0, 1)
-            seq_dim = 0
-
-        if reorder_strategy == ReorderStrategy.Striped:
-            seq_lens = shape[seq_dim]
-            if seq_lens < (cp_size * stripe_size):
-                pytest.skip(f"{seq_lens=} must be larger than {cp_size*stripe_size=}")
+        if qkv_format == QKVFormat.SBHD:
+            tensor = tensor.swapaxes(0, 1)
+            seq_dim = 0
+
+        if reorder_strategy == ReorderStrategy.Striped:
+            seq_lens = tensor.shape[seq_dim]
+            if seq_lens < (cp_size * stripe_size):
+                pytest.skip(f"{seq_lens=} must be larger than {cp_size*stripe_size=}")

Tests

• tests/jax/test_distributed_fused_attn.py

--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -424,6 +424,25 @@ class TestReorderCausalLoadBalancing:
         reordered = reorder(tensor, reorder_strategy, cp_size, seq_dim, stripe_size)
         inversed = inverse(reordered, reorder_strategy, cp_size, seq_dim, stripe_size)
 
         assert jnp.array_equal(inversed, ref)
+
+    @pytest.mark.parametrize("stripe_size", [1, 4])
+    def test_sbhd_striped_uses_swapped_seq_dim(self, stripe_size):
+        """Regression test: SBHD Striped skip must use the swapped sequence dim."""
+        cp_size = 2
+        shape = (1, 16, 1, 1)  # original [batch, seq, heads, dim]
+        tensor = random.normal(random.PRNGKey(42), shape, dtype=jnp.bfloat16)
+        tensor = tensor.swapaxes(0, 1)  # SBHD: [seq, batch, heads, dim]
+
+        # Old logic read original shape[0]=1 (batch) and skipped; seq_len after swap is 16.
+        reorder = jax.jit(reorder_causal_load_balancing, static_argnums=[1, 2, 3, 4])
+        inverse = jax.jit(inverse_reorder_causal_load_balancing, static_argnums=[1, 2, 3, 4])
+
+        reordered = reorder(tensor, ReorderStrategy.Striped, cp_size, 0, stripe_size)
+        inversed = inverse(reordered, ReorderStrategy.Striped, cp_size, 0, stripe_size)
+
+        assert jnp.array_equal(inversed, tensor)

[greptile-apps]

Greptile Summary

The PR corrects the SBHD Striped test’s skip calculation to inspect the sequence dimension of the swapped tensor.

• Uses tensor.shape[seq_dim] when deciding whether a Striped case is large enough.
• Adds an SBHD reorder/inverse round-trip test, although it does not exercise the corrected skip branch.

Confidence Score: 4/5

The PR appears safe to merge, with a non-blocking test-coverage gap because the new regression test does not exercise the skip logic being fixed.

The production behavior is unchanged and the corrected dimension matches the reordered tensor, but the added test only validates reorder/inverse round-trip behavior and therefore does not prevent the original skip regression from returning.

Files Needing Attention: tests/jax/test_distributed_fused_attn.py

Important Files Changed

Filename	Overview
tests/jax/test_distributed_fused_attn.py	The skip calculation is corrected, but the added regression test remains green if that correction is reverted because it bypasses the skip branch.

_{Reviews (1): Last reviewed commit: "fix: SBHD reorder skip uses original sha..." | Re-trigger Greptile}

[greptile-apps]

Regression test bypasses skip logic

This test calls the reorder and inverse functions directly instead of exercising the changed skip branch, so reverting tensor.shape[seq_dim] to the original expression leaves it green and allows the erroneous loss of SBHD Striped coverage to recur undetected.

Knowledge Base Used: Tests and QA

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[cyanguwa]

@KshitijLakhani, could you please help review this one? Thanks.