[JAX]: 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 feedback addressed

• Updated test_sbhd_striped_uses_swapped_seq_dim to invoke the parametrized test(...) method directly with SBHD format, so the corrected skip logic is actually exercised. A monkeypatched pytest.skip turns an unexpected skip into a test failure.

Local verification: python3 -m py_compile tests/jax/test_distributed_fused_attn.py passed. Full pytest execution was not feasible because JAX is not installed in this environment.

[greptile-apps]

Greptile Summary

This PR corrects the SBHD Striped test guard to inspect the reordered tensor’s sequence dimension and adds regression coverage for the corrected branch.

• Uses tensor.shape[seq_dim] after the SBHD axis swap.
• Converts an unexpected skip into a test failure, ensuring the old expression is detected for both stripe sizes.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
tests/jax/test_distributed_fused_attn.py	Corrects the post-swap sequence-length check and adds a regression test that fails if the previous skip behavior returns.

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

[cyanguwa]

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

[andrewwhitecdw]

@greptile-apps thanks for the review. I addressed the test-coverage gap:

• Rewrote test_sbhd_striped_uses_swapped_seq_dim to call the parametrized test(...) method directly with QKVFormat.SBHD and ReorderStrategy.Striped, so the skip branch is exercised.
• Monkeypatched pytest.skip inside that test so an unexpected skip (which would happen if the original unswapped shape were used) is turned into a failure.

The branch has been amended to a single signed-off commit (ecd70cd) and force-pushed. python3 -m py_compile tests/jax/test_distributed_fused_attn.py passed; full pytest was not feasible because JAX is not installed here.

[greptile-apps]

The approach is solid. A few observations on the updated test:

What works well:

• Delegating to self.test(...) rather than duplicating logic means the regression test exercises the exact same code path as the production parametrized suite.
• The monkeypatch.setattr(pytest, "skip", ...) trick is the right tool here — since test() accesses pytest.skip through the module attribute at call time, the patch is visible to it.
• shape = (1, 16, 1, 1) is a carefully chosen pathological case: the old code would read shape[0] = 1 < cp_size * stripe_size (2 or 8) and silently skip; the fixed code reads tensor.shape[0] = 16 after the swap and proceeds correctly.

One minor subtlety to be aware of: monkeypatch.setattr(pytest, "skip", ...) replaces the global pytest.skip. If reorder_causal_load_balancing or any downstream JAX/TE code internally calls pytest.skip, those would also become failures. That's unlikely to matter in practice, but it's worth noting if a mysterious unexpected pytest.skip failure ever shows up in CI on a machine where JAX raises a skip.

The core fix (tensor.shape[seq_dim] replacing shape[seq_dim]) is correct and minimal. The regression test is now a proper guard against reintroduction. LGTM.

[KshitijLakhani]

Thanks for this PR @andrewwhitecdw !
Could you please address the review comments ?
I'm happy to launch our CI on this PR after you've addressed those comments. I'll also confirm that the skipped tests for SBHD actually pass in the CI.

NOTE: A passing CI and approved review are required prior to merge

[KshitijLakhani]

Good catch !
This change makes sense to me. I would like to think that the main branch basically just skips in the next line due to this incorrect statement (as it would be incorrectly taking the batch value as the seq value) - this explains why our CI never caught this as a failure!
I'd expect those incorrectly skipping tests to not be skipped with this change (and pass)

[KshitijLakhani]

I did some digging to figure how many and which tests might be skipping and this is the list. A cursory look at TE's CI from release 2.18 does show that these tests were indeed being skipped.

 ### L1

  Shape: [3, 32, 8, 64]

  For SBHD, the real sequence length is 32, but the buggy code reads batch size 3.

  Incorrectly skipped:

  - cp_size=4, stripe_size=1
  - cp_size=8, stripe_size=1
  - cp_size=2, stripe_size=4
  - cp_size=4, stripe_size=4
  - cp_size=8, stripe_size=4

  Only cp_size=2, stripe_size=1 currently runs.

  Therefore L1 has five incorrect skips.

  ### L2

  Shape: [4, 32, 12, 32]

  The buggy code reads 4 instead of sequence length 32.

  Incorrectly skipped:

  - cp_size=8, stripe_size=1
  - cp_size=2, stripe_size=4
  - cp_size=4, stripe_size=4
  - cp_size=8, stripe_size=4

  Shape: [1, 16, 1, 1]

  The buggy code reads 1 instead of sequence length 16.

  Incorrectly skipped:

  - cp_size=2, stripe_size=1
  - cp_size=4, stripe_size=1
  - cp_size=8, stripe_size=1
  - cp_size=2, stripe_size=4
  - cp_size=4, stripe_size=4

I'd expect these to not be skipped and passed after @andrewwhitecdw 's changes

[KshitijLakhani]

@andrewwhitecdw I do not think this is needed as the tests above do run SBHD.
With your change to correctly get the seq_lens we should be good.
Please remove this:

Suggested change

	@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)

[KshitijLakhani]

Re-publishing my comment as it might have become stale as part of an earlier review due to recent commits pushed by @andrewwhitecdw

I do not think this is needed as the original tests in TestReorderCausalLoadBalancing test() above do run SBHD. With your change to correctly get the seq_lens we should be good.
Please remove this:

Suggested change

self.test(cp_size, shape, QKVFormat.SBHD, ReorderStrategy.Striped, stripe_size)