From 12171872d169644e0f6721256729d1e1533ce3a1 Mon Sep 17 00:00:00 2001 From: andrewwhitecdw Date: Tue, 18 Aug 2026 15:52:48 -0500 Subject: [PATCH] fix: SBHD reorder skip uses original shape instead of swapped tensor Use tensor.shape[seq_dim] instead of shape[seq_dim] when deciding whether a Striped SBHD case is large enough. Signed-off-by: andrewwhitecdw --- tests/jax/test_distributed_fused_attn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index c1cc2b8282..b6e11b8bea 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -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] if seq_lens < (cp_size * stripe_size): pytest.skip(f"{seq_lens=} must be larger than {cp_size*stripe_size=}")