Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andrewwhitecdw wants to merge 1 commit into
base: main
Choose a base branch
from
+15 −1
Open

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

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion tests/jax/test_distributed_fused_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -835,7 +835,7 @@ def test(self, cp_size, shape, qkv_format, reorder_strategy, stripe_size):
seq_dim = 0

if reorder_strategy == ReorderStrategy.Striped:
seq_lens = shape[seq_dim]
seq_lens = tensor.shape[seq_dim]

@KshitijLakhani KshitijLakhani Aug 17, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 KshitijLakhani Aug 17, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

if seq_lens < (cp_size * stripe_size):
pytest.skip(f"{seq_lens=} must be larger than {cp_size*stripe_size=}")

Expand All @@ -848,3 +848,17 @@ def test(self, cp_size, shape, qkv_format, reorder_strategy, 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, monkeypatch):
"""Regression test: SBHD Striped skip must use the swapped sequence dim."""
cp_size = 2
shape = (1, 16, 1, 1) # original [batch, seq, heads, dim]

# If the skip logic reads the original unswapped batch dim (1), it would skip
# because 1 < cp_size * stripe_size. With the fix it reads the swapped seq
# dim (16), so the parametrized test should run to completion.
monkeypatch.setattr(
pytest, "skip", lambda reason: pytest.fail(f"unexpected pytest.skip: {reason}")
)
self.test(cp_size, shape, QKVFormat.SBHD, ReorderStrategy.Striped, stripe_size)

@KshitijLakhani KshitijLakhani Aug 17, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)