[JAX] SBHD reorder skip uses original shape instead of swapped tensor

[andrewwhitecdw]

This PR fixes the SBHD Striped skip logic in tests/jax/test_distributed_fused_attn.py so it reads the sequence length from the swapped tensor shape rather than the original shape.

Changes

• tests/jax/test_distributed_fused_attn.py: use tensor.shape[seq_dim] instead of shape[seq_dim] after the SBHD axis swap.

Details

When qkv_format == QKVFormat.SBHD, the test swaps axes 0 and 1 so seq_dim becomes 0. The old skip guard read shape[seq_dim], which is the original (unswapped) batch size for SBHD, causing incorrect skips. The fix reads tensor.shape[seq_dim], which reflects the swapped sequence length.

-            seq_lens = shape[seq_dim]
+            seq_lens = tensor.shape[seq_dim]

Tests

• Covered by the existing parametrized TestReorderCausalLoadBalancing.test cases that exercise SBHD inputs.

Contributor guidelines

• Signed-off-by: Andrew White andrewwhitecdw@users.noreply.github.com

[greptile-apps]

Greptile Summary

This PR corrects the JAX distributed fused-attention test’s SBHD Striped skip guard to derive sequence length from the axis-swapped tensor.

• Uses tensor.shape[seq_dim] after SBHD swaps the batch and sequence axes.
• Preserves the existing minimum-length check for Striped reordering.

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	Correctly reads the effective sequence dimension from the reordered tensor when deciding whether Striped cases are eligible.

_{Reviews (5): Last reviewed commit: "Merge branch 'main' into bugfix/test-dis..." | 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)

[andrewwhitecdw]

@KshitijLakhani Thanks for the review. I removed the extra regression test as requested — the existing parametrized SBHD Striped cases now exercise the corrected skip logic directly. The branch is a single signed-off commit rebased onto main.

[KshitijLakhani]

/te-ci jax L2

[KshitijLakhani]

I launched the L1 tests manually as the te_ci pipeline launch from the PR triggered the L0 and L2 tests only.
Behavior is as expected - no more unintended skips for the SBHD case and they all pass for all intended test levels

[KshitijLakhani]

LGTM !