From 84990135cf93e50a73b55878b18e67130e30b70b Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 15:07:04 -0700 Subject: [PATCH 1/4] test(sdpa): fuzz per-tensor ragged token-stride gaps in test_mhas_v2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ragged Q/K/V/O were always generated with canonical packed strides (token stride == h*d): the gaps stride machinery was dense-only, and PR #462 varies only the STATS layout. So no sweep could ever produce a tensor like a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d) — the layout torch.nn.attention.varlen users get by slicing a fused KV projection. New with_ragged_token_gap knob in both ragged random sweeps: each of Q/K/V/O independently draws token stride h*d + gap, gap in {0, 8, 64, roundup8(h*d)} seeded from rng_geom_seed (deterministic through serialize/deserialize repro; explicit strides still win). Gaps stay multiples of 8 elements so ragged base addresses keep the packed layout's alignment class (the graph API requires 16-byte-aligned pointers). Harness: ragged buffers (incl. gradients) are allocated with the configured strides, and ragged offsets scale by each tensor's actual stride[2] — the same generalization #462 made for stats offsets. KNOWN FAILURES this exposes (intentionally not masked): the FROST THD forward engines (sdpa_fwd_prefill_sm120, sdpa_fwd_prefill_sm100_d128) claim non-packed-stride THD graphs and silently mis-address them (100% of O wrong; the stride ORDER is still BSHD, so the order-only layout gate passes). Backend engines serve every gapped combination correctly, fwd and bwd. The failing configs under CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 are the repro set for fixing the THD lowerings to honor declared strides. Co-Authored-By: Claude Fable 5 --- test/python/sdpa/fp16.py | 33 ++++++++++++------- test/python/sdpa/random_config.py | 53 ++++++++++++++++++++++++++----- test/python/test_mhas_v2.py | 2 ++ 3 files changed, 68 insertions(+), 20 deletions(-) diff --git a/test/python/sdpa/fp16.py b/test/python/sdpa/fp16.py index c1d92d207..b104f94e5 100644 --- a/test/python/sdpa/fp16.py +++ b/test/python/sdpa/fp16.py @@ -140,10 +140,17 @@ def allocate_tensors(cfg, rng_data_gen, perf=False): si = not perf if cfg.is_ragged: - allocs[TensorUid.q] = alloc_tensor((max_t_q, cfg.h_q, cfg.d_qk), cfg.data_type, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si) - allocs[TensorUid.k] = alloc_tensor((max_t_kv, cfg.h_k, cfg.d_qk), cfg.data_type, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si) - allocs[TensorUid.v] = alloc_tensor((max_t_kv, cfg.h_v, cfg.d_v), cfg.data_type, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si) - allocs[TensorUid.o] = alloc_tensor((max_t_q, cfg.h_q, cfg.d_v), cfg.data_type) + # 3-D (token, head, elem) strides come from the 4-D configs, so + # per-tensor token-stride gaps (with_ragged_token_gap) reach the + # actual buffers. + q_strides = (cfg.stride_q[2], cfg.stride_q[1], cfg.stride_q[3]) + k_strides = (cfg.stride_k[2], cfg.stride_k[1], cfg.stride_k[3]) + v_strides = (cfg.stride_v[2], cfg.stride_v[1], cfg.stride_v[3]) + o_strides = (cfg.stride_o[2], cfg.stride_o[1], cfg.stride_o[3]) + allocs[TensorUid.q] = alloc_tensor((max_t_q, cfg.h_q, cfg.d_qk), cfg.data_type, strides=q_strides, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si) + allocs[TensorUid.k] = alloc_tensor((max_t_kv, cfg.h_k, cfg.d_qk), cfg.data_type, strides=k_strides, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si) + allocs[TensorUid.v] = alloc_tensor((max_t_kv, cfg.h_v, cfg.d_v), cfg.data_type, strides=v_strides, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si) + allocs[TensorUid.o] = alloc_tensor((max_t_q, cfg.h_q, cfg.d_v), cfg.data_type, strides=o_strides) # cfg.stride_stats is 4-D (b, h, s, 1); its [1] and [2] entries are the head and token # strides of the packed buffer, which is exactly the (h, s) part of the 3-D alloc below. stats_strides = (cfg.stride_stats[2], cfg.stride_stats[1], 1) @@ -151,10 +158,10 @@ def allocate_tensors(cfg, rng_data_gen, perf=False): allocs[TensorUid.score_max] = alloc_tensor((max_t_q, cfg.h_q, 1), torch.float32, strides=stats_strides) if cfg.with_score_max else (None, None, None) allocs[TensorUid.score_sum_exp] = alloc_tensor((max_t_q, cfg.h_q, 1), torch.float32, strides=stats_strides) if cfg.with_score_sum_exp else (None, None, None) if cfg.is_train: - allocs[TensorUid.dQ] = alloc_tensor((max_t_q, cfg.h_q, cfg.d_qk), cfg.data_type) - allocs[TensorUid.dK] = alloc_tensor((max_t_kv, cfg.h_k, cfg.d_qk), cfg.data_type) - allocs[TensorUid.dV] = alloc_tensor((max_t_kv, cfg.h_v, cfg.d_v), cfg.data_type) - allocs[TensorUid.dO] = alloc_tensor((max_t_q, cfg.h_q, cfg.d_v), cfg.data_type, rng=rng_data_gen, mean=0.0, std=0.1, sparse_int=si) + allocs[TensorUid.dQ] = alloc_tensor((max_t_q, cfg.h_q, cfg.d_qk), cfg.data_type, strides=q_strides) + allocs[TensorUid.dK] = alloc_tensor((max_t_kv, cfg.h_k, cfg.d_qk), cfg.data_type, strides=k_strides) + allocs[TensorUid.dV] = alloc_tensor((max_t_kv, cfg.h_v, cfg.d_v), cfg.data_type, strides=v_strides) + allocs[TensorUid.dO] = alloc_tensor((max_t_q, cfg.h_q, cfg.d_v), cfg.data_type, strides=o_strides, rng=rng_data_gen, mean=0.0, std=0.1, sparse_int=si) else: allocs[TensorUid.q] = alloc_tensor(cfg.shape_q, cfg.data_type, strides=cfg.stride_q, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si) allocs[TensorUid.k] = alloc_tensor(cfg.shape_k, cfg.data_type, strides=cfg.stride_k, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si) @@ -194,10 +201,12 @@ def allocate_tensors(cfg, rng_data_gen, perf=False): k_off_mult = cfg.d_qk if cfg.with_ragged_offset_multiplier else 1 v_off_mult = cfg.d_v if cfg.with_ragged_offset_multiplier else 1 o_off_mult = cfg.d_v if cfg.with_ragged_offset_multiplier else 1 - allocs[TensorUid.q_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.h_q * cfg.d_qk // q_off_mult).to(torch.int64), None, None) - allocs[TensorUid.k_ragged_offset] = ((prefix_sum(seq_len_kv_gpu) * cfg.h_k * cfg.d_qk // k_off_mult).to(torch.int64), None, None) - allocs[TensorUid.v_ragged_offset] = ((prefix_sum(seq_len_kv_gpu) * cfg.h_v * cfg.d_v // v_off_mult).to(torch.int64), None, None) - allocs[TensorUid.o_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.h_q * cfg.d_v // o_off_mult).to(torch.int64), None, None) + # Offsets scale by each tensor's ACTUAL token stride (stride[2]), not an + # assumed-packed h*d — K/V may carry a token-stride gap (kv-interleaved). + allocs[TensorUid.q_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.stride_q[2] // q_off_mult).to(torch.int64), None, None) + allocs[TensorUid.k_ragged_offset] = ((prefix_sum(seq_len_kv_gpu) * cfg.stride_k[2] // k_off_mult).to(torch.int64), None, None) + allocs[TensorUid.v_ragged_offset] = ((prefix_sum(seq_len_kv_gpu) * cfg.stride_v[2] // v_off_mult).to(torch.int64), None, None) + allocs[TensorUid.o_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.stride_o[2] // o_off_mult).to(torch.int64), None, None) # Stats offsets are in elements and scale by its token stride: h_q for token-major stats, # 1 for head-major. allocs[TensorUid.stats_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.stride_stats[2]).to(torch.int64), None, None) diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index 082bbfe6f..e0788e8ad 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -63,12 +63,20 @@ def compute_default_BHSD_strides(shape): return tuple(strides) -def compute_packed_strides(shape): - """Compute packed (ragged) BSHD strides for BHSD shape: (s*h*d, d, h*d, 1).""" +def compute_packed_strides(shape, token_gap=0): + """Compute packed (ragged) BSHD strides for BHSD shape: (s*h*d, d, h*d, 1). + + ``token_gap`` widens the token stride to ``h*d + token_gap`` elements — + the layout of a tensor VIEW into a larger per-token record. With + ``token_gap == h*d`` this is exactly a K or V view of a kv-interleaved + ``[T, 2, H, D]`` buffer (token stride ``2*h*d``), the layout + ``torch.nn.attention.varlen`` users produce by slicing a fused KV + projection.""" if shape is None: return None b, h, s, d = shape - return (s * h * d, d, h * d, 1) + token_stride = h * d + token_gap + return (s * token_stride, d, token_stride, 1) @dataclass @@ -114,6 +122,15 @@ class ExecConfig: with_unfuse_fma: bool = False with_rope: bool = False with_ragged_offset_multiplier: bool = False + # Each ragged tensor (Q/K/V/O and gradients) independently draws a token + # stride h*d + gap, gap in {0, 8, 64, roundup8(h*d)} seeded from + # rng_geom_seed — so mixed combinations occur (e.g. gapped Q/K with packed + # V/O). gap≈h*d is a view of an interleaved [T, 2, H, D] buffer, the + # layout torch.nn.attention.varlen users produce by slicing a fused KV + # projection. Gaps are multiples of 8 elements so ragged base addresses + # keep the alignment class of the packed layout (the graph API requires + # 16-byte-aligned pointers; an odd gap would make every offset illegal). + with_ragged_token_gap: bool = False rescale_threshold: float = None diag_align: cudnn.diagonal_alignment = None @@ -183,16 +200,36 @@ def fill_derived_fields(self): if self.shape_stats is None and all(x is not None for x in [self.batches, self.h_q, self.s_q]): self.shape_stats = (self.batches, self.h_q, self.s_q, 1) - # Compute strides if not provided (packed for ragged, default BHSD otherwise) + # Compute strides if not provided (packed for ragged, default BHSD otherwise). + # with_ragged_token_gap: per-tensor token-stride gaps, re-derived + # deterministically from rng_geom_seed (so serialize/deserialize repro + # reproduces the same strides). + if self.is_ragged and self.with_ragged_token_gap: + _gap_rng = random.Random((self.rng_geom_seed or 0) ^ 0xA80517) + + def _gapped(shape): + if shape is None: + return None + h, d = shape[1], shape[3] + hd8 = ((h * d + 7) // 8) * 8 # interleaved-buffer gap, alignment-preserving + return compute_packed_strides(shape, _gap_rng.choice([0, 8, 64, hd8])) + + gap_fn = _gapped + elif self.is_ragged: + gap_fn = compute_packed_strides + else: + gap_fn = compute_default_BHSD_strides stride_fn = compute_packed_strides if self.is_ragged else compute_default_BHSD_strides if self.stride_q is None and self.shape_q is not None: - self.stride_q = stride_fn(self.shape_q) + self.stride_q = gap_fn(self.shape_q) if self.stride_k is None and self.shape_k is not None: - self.stride_k = stride_fn(self.shape_k) + self.stride_k = gap_fn(self.shape_k) if self.stride_v is None and self.shape_v is not None: - self.stride_v = stride_fn(self.shape_v) + self.stride_v = gap_fn(self.shape_v) if self.stride_o is None and self.shape_o is not None: - self.stride_o = stride_fn(self.shape_o) + self.stride_o = gap_fn(self.shape_o) + # stats keeps the packed default — its layout is fuzzed separately + # via ragged_stats_layout. if self.stride_stats is None and self.shape_stats is not None: self.stride_stats = stride_fn(self.shape_stats) diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 61f22b384..a550ce2c3 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -357,6 +357,7 @@ def test_sdpa_random_fwd_ragged_L0(env_info, test_no, request, cudnn_handle): is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "padded" : 0, "full" : 0}), with_sink_token=RandomChoice({True : 1, False : 3}), ragged_stats_layout=RandomChoice({"token_major" : 1, "head_major" : 1}), + with_ragged_token_gap=RandomChoice({True : 1, False : 3}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) @@ -458,6 +459,7 @@ def test_sdpa_random_bwd_ragged_L0(env_info, test_no, request, cudnn_handle): is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "padded" : 0, "full" : 0}), is_deterministic=RandomChoice({True : 3, False : 1}), ragged_stats_layout=RandomChoice({"token_major" : 1, "head_major" : 1}), + with_ragged_token_gap=RandomChoice({True : 1, False : 3}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) From b7d8c10fa10abafc3dbde52789fb5e8c36875319 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 22:45:12 -0700 Subject: [PATCH 2/4] test(sdpa): simplify ragged token gaps to whole tokens; always-on in the fp16 ragged sweeps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review: drop the True/False dice — the fp16 ragged sweeps enable the knob unconditionally, and the packed case comes from the gap draw itself. Each of Q/K/V/O independently draws gap = n*h*d tokens, n in 0..3 (n=0 packed, n=1 a kv-interleaved [T, 2, H, D] view, n=2 a [T, 3, H, D] QKV-interleave). Whole-token gaps keep every ragged base address in the packed layout's alignment class by construction, replacing the previous multiples-of-8-elements rule. The config field stays default-False: fixed configs and the fp8/mxfp8 harnesses still assume packed allocations, and the cu_ragged form derives offsets internally. Verified: backend engines pass 5/5 seeded repros, a bf16 training config, and a 128-test slice of the fwd ragged L0 sweep; FROST THD fwd engines keep failing gapped draws (the intended standing repro set). Co-Authored-By: Claude Fable 5 --- test/python/sdpa/random_config.py | 21 ++++++++++++--------- test/python/test_mhas_v2.py | 4 ++-- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index e0788e8ad..81bb12e42 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -123,13 +123,17 @@ class ExecConfig: with_rope: bool = False with_ragged_offset_multiplier: bool = False # Each ragged tensor (Q/K/V/O and gradients) independently draws a token - # stride h*d + gap, gap in {0, 8, 64, roundup8(h*d)} seeded from - # rng_geom_seed — so mixed combinations occur (e.g. gapped Q/K with packed - # V/O). gap≈h*d is a view of an interleaved [T, 2, H, D] buffer, the - # layout torch.nn.attention.varlen users produce by slicing a fused KV - # projection. Gaps are multiples of 8 elements so ragged base addresses - # keep the alignment class of the packed layout (the graph API requires - # 16-byte-aligned pointers; an odd gap would make every offset illegal). + # stride of 1-4 whole tokens (gap = n*h*d, n in 0..3, seeded from + # rng_geom_seed): n=0 is the plain packed case, n=1 is exactly a view of + # an interleaved [T, 2, H, D] buffer (the layout + # torch.nn.attention.varlen users produce by slicing a fused KV + # projection), n=2 a [T, 3, H, D] QKV-interleave, and so on. Whole-token + # gaps keep every ragged base address in the packed layout's alignment + # class by construction (sub-token gaps can violate the graph API's + # 16-byte pointer-alignment contract — an odd-element gap is illegal for + # every engine). Default False: fixed configs and the fp8/mxfp8 + # harnesses allocate assuming packed strides; the fp16 ragged sweeps + # enable this unconditionally. with_ragged_token_gap: bool = False rescale_threshold: float = None @@ -211,8 +215,7 @@ def _gapped(shape): if shape is None: return None h, d = shape[1], shape[3] - hd8 = ((h * d + 7) // 8) * 8 # interleaved-buffer gap, alignment-preserving - return compute_packed_strides(shape, _gap_rng.choice([0, 8, 64, hd8])) + return compute_packed_strides(shape, _gap_rng.randint(0, 3) * h * d) gap_fn = _gapped elif self.is_ragged: diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index a550ce2c3..7d4414e22 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -357,7 +357,7 @@ def test_sdpa_random_fwd_ragged_L0(env_info, test_no, request, cudnn_handle): is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "padded" : 0, "full" : 0}), with_sink_token=RandomChoice({True : 1, False : 3}), ragged_stats_layout=RandomChoice({"token_major" : 1, "head_major" : 1}), - with_ragged_token_gap=RandomChoice({True : 1, False : 3}), + with_ragged_token_gap=RandomChoice({True : 1}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) @@ -459,7 +459,7 @@ def test_sdpa_random_bwd_ragged_L0(env_info, test_no, request, cudnn_handle): is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "padded" : 0, "full" : 0}), is_deterministic=RandomChoice({True : 3, False : 1}), ragged_stats_layout=RandomChoice({"token_major" : 1, "head_major" : 1}), - with_ragged_token_gap=RandomChoice({True : 1, False : 3}), + with_ragged_token_gap=RandomChoice({True : 1}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) From 17be55680bb511ce64924e0e4f4617dc2e929fc6 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Mon, 10 Aug 2026 23:49:11 -0700 Subject: [PATCH 3/4] test(sdpa): draw all four ragged token gaps up front; drop unused binding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gap RNG was consumed lazily, one draw per stride left unspecified — so explicitly pinning e.g. stride_q shifted the gaps K/V/O derive from the same rng_geom_seed. Draw all four values up front in fixed Q/K/V/O order and apply each only where the stride is missing: per-tensor layouts are now a function of the seed alone. The all-defaults path (the sweeps) consumes draws in the same order as before, so existing seeded repros reproduce identical strides. Adds a GPU-free regression test pinning stride_q and asserting the K/V/O strides match the all-defaults derivation (seed chosen so the old lazy behavior visibly shifts two of the three gaps). Also renames the unused batch binding in compute_packed_strides (RUF059). Addresses CodeRabbit review feedback on #516. Co-Authored-By: Claude Fable 5 --- test/python/sdpa/random_config.py | 38 +++++++++++++++++++------------ test/python/test_mhas_v2.py | 23 +++++++++++++++++++ 2 files changed, 46 insertions(+), 15 deletions(-) diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index 81bb12e42..a9829c2ba 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -74,7 +74,7 @@ def compute_packed_strides(shape, token_gap=0): projection.""" if shape is None: return None - b, h, s, d = shape + _, h, s, d = shape token_stride = h * d + token_gap return (s * token_stride, d, token_stride, 1) @@ -210,27 +210,35 @@ def fill_derived_fields(self): # reproduces the same strides). if self.is_ragged and self.with_ragged_token_gap: _gap_rng = random.Random((self.rng_geom_seed or 0) ^ 0xA80517) - - def _gapped(shape): - if shape is None: - return None - h, d = shape[1], shape[3] - return compute_packed_strides(shape, _gap_rng.randint(0, 3) * h * d) - - gap_fn = _gapped + # Draw ALL FOUR gaps up front, in fixed Q/K/V/O order: an + # explicitly provided stride must not shift the gaps the + # remaining tensors get (same rng_geom_seed -> same per-tensor + # layouts regardless of which strides were overridden). + _gaps = {name: _gap_rng.randint(0, 3) for name in ("q", "k", "v", "o")} + + def _make_gap_fn(gap_tokens): + def _gapped(shape): + if shape is None: + return None + h, d = shape[1], shape[3] + return compute_packed_strides(shape, gap_tokens * h * d) + + return _gapped + + gap_q, gap_k, gap_v, gap_o = (_make_gap_fn(_gaps[n]) for n in ("q", "k", "v", "o")) elif self.is_ragged: - gap_fn = compute_packed_strides + gap_q = gap_k = gap_v = gap_o = compute_packed_strides else: - gap_fn = compute_default_BHSD_strides + gap_q = gap_k = gap_v = gap_o = compute_default_BHSD_strides stride_fn = compute_packed_strides if self.is_ragged else compute_default_BHSD_strides if self.stride_q is None and self.shape_q is not None: - self.stride_q = gap_fn(self.shape_q) + self.stride_q = gap_q(self.shape_q) if self.stride_k is None and self.shape_k is not None: - self.stride_k = gap_fn(self.shape_k) + self.stride_k = gap_k(self.shape_k) if self.stride_v is None and self.shape_v is not None: - self.stride_v = gap_fn(self.shape_v) + self.stride_v = gap_v(self.shape_v) if self.stride_o is None and self.shape_o is not None: - self.stride_o = gap_fn(self.shape_o) + self.stride_o = gap_o(self.shape_o) # stats keeps the packed default — its layout is fuzzed separately # via ragged_stats_layout. if self.stride_stats is None and self.shape_stats is not None: diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 7d4414e22..9c45e25fc 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -366,6 +366,29 @@ def test_sdpa_random_fwd_ragged_L0(env_info, test_no, request, cudnn_handle): exec_sdpa(test.cfg, request, cudnn_handle) +@pytest.mark.L0 +def test_ragged_token_gap_stable_under_stride_overrides(): + """Regression: the seeded per-tensor token gaps must not depend on which + strides were explicitly provided — pinning stride_q must leave the gaps + K/V/O derive from the same rng_geom_seed unchanged (the gap RNG draws all + four values up front, not lazily per missing stride).""" + from sdpa.random_config import ExecConfig + + base = dict( + batches=2, h_q=8, h_k=8, h_v=8, s_q=64, s_kv=64, d_qk=128, d_v=128, + is_ragged=True, with_ragged_token_gap=True, rng_geom_seed=7, + ) + plain = ExecConfig(**base) + plain.fill_derived_fields() + + pinned_q = (64 * 8 * 128, 128, 8 * 128, 1) # explicit packed Q, no gap + pinned = ExecConfig(**base, stride_q=pinned_q) + pinned.fill_derived_fields() + + assert pinned.stride_q == pinned_q + assert (pinned.stride_k, pinned.stride_v, pinned.stride_o) == (plain.stride_k, plain.stride_v, plain.stride_o) + + @pytest.mark.parametrize("test_no", generate_test_seeds(num_tests=128, rng_seed=888), ids=lambda p: f"test{p[0]}") @pytest.mark.L1 def test_sdpa_random_fwd_ragged_unified_L1(env_info, test_no, request, cudnn_handle): From 746ef17e8629e48fdb1da1f66a5a75819149f359 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Tue, 11 Aug 2026 00:05:55 -0700 Subject: [PATCH 4/4] test(sdpa): ragged token gaps on by default; auto-packed where not yet expressible MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop the per-sweep with_ragged_token_gap opt-in: the field now defaults True, so EVERY ragged config fuzzes its per-tensor token strides. fill_derived_fields falls back to packed automatically where a gap is not yet expressible or handled — cu / offset-multiplier forms bind offsets as cu (x multiplier) and cannot declare a token gap (#538), and the fp8/mxfp8 harnesses (1-byte data types) allocate assuming packed (#537) — so mixed sweeps (ragged + cu_ragged in one RandomChoice) gap exactly the draws that support it. Explicit strides are never touched (the gap only fills strides left None), so recorded repro dicts reproduce exactly. The regression test also locks in the new semantics: default-on for plain ragged, packed for the three fallback forms. Co-Authored-By: Claude Fable 5 --- test/python/sdpa/random_config.py | 30 ++++++++++++++++++++++-------- test/python/test_mhas_v2.py | 22 +++++++++++++++++----- 2 files changed, 39 insertions(+), 13 deletions(-) diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index a9829c2ba..6a1aca02e 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -131,10 +131,12 @@ class ExecConfig: # gaps keep every ragged base address in the packed layout's alignment # class by construction (sub-token gaps can violate the graph API's # 16-byte pointer-alignment contract — an odd-element gap is illegal for - # every engine). Default False: fixed configs and the fp8/mxfp8 - # harnesses allocate assuming packed strides; the fp16 ragged sweeps - # enable this unconditionally. - with_ragged_token_gap: bool = False + # every engine). Default True: every ragged config fuzzes its layouts. + # fill_derived_fields auto-falls-back to packed where a gap is not yet + # expressible or handled: cu / offset-multiplier forms (#538) and the + # fp8/mxfp8 harnesses (#537). Configs with explicit strides are + # unaffected (the gap only fills strides left None). + with_ragged_token_gap: bool = True rescale_threshold: float = None diag_align: cudnn.diagonal_alignment = None @@ -205,10 +207,22 @@ def fill_derived_fields(self): self.shape_stats = (self.batches, self.h_q, self.s_q, 1) # Compute strides if not provided (packed for ragged, default BHSD otherwise). - # with_ragged_token_gap: per-tensor token-stride gaps, re-derived - # deterministically from rng_geom_seed (so serialize/deserialize repro - # reproduces the same strides). - if self.is_ragged and self.with_ragged_token_gap: + # with_ragged_token_gap (default True): per-tensor token-stride gaps, + # re-derived deterministically from rng_geom_seed (so + # serialize/deserialize repro reproduces the same strides). Auto-packed + # where a gap is not yet expressible or handled: + # - cu / offset-multiplier forms bind offsets as cu (x multiplier) + # and cannot declare a token gap (#538); + # - the fp8/mxfp8 harnesses (1-byte data_type) allocate assuming + # packed strides (#537). + _gap_applicable = ( + self.is_ragged + and self.with_ragged_token_gap + and not self.is_cu_seq_len + and not self.with_ragged_offset_multiplier + and not (self.data_type is not None and self.data_type.itemsize == 1) + ) + if _gap_applicable: _gap_rng = random.Random((self.rng_geom_seed or 0) ^ 0xA80517) # Draw ALL FOUR gaps up front, in fixed Q/K/V/O order: an # explicitly provided stride must not shift the gaps the diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 9c45e25fc..129313552 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -357,7 +357,6 @@ def test_sdpa_random_fwd_ragged_L0(env_info, test_no, request, cudnn_handle): is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "padded" : 0, "full" : 0}), with_sink_token=RandomChoice({True : 1, False : 3}), ragged_stats_layout=RandomChoice({"token_major" : 1, "head_major" : 1}), - with_ragged_token_gap=RandomChoice({True : 1}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) @@ -371,12 +370,15 @@ def test_ragged_token_gap_stable_under_stride_overrides(): """Regression: the seeded per-tensor token gaps must not depend on which strides were explicitly provided — pinning stride_q must leave the gaps K/V/O derive from the same rng_geom_seed unchanged (the gap RNG draws all - four values up front, not lazily per missing stride).""" - from sdpa.random_config import ExecConfig + four values up front, not lazily per missing stride). Also locks in the + default-on semantics: gaps apply to plain ragged configs by default, and + auto-fall-back to packed for the forms that cannot express or handle + them yet (cu / offset-multiplier: #538; fp8 harness: #537).""" + from sdpa.random_config import ExecConfig, compute_packed_strides base = dict( batches=2, h_q=8, h_k=8, h_v=8, s_q=64, s_kv=64, d_qk=128, d_v=128, - is_ragged=True, with_ragged_token_gap=True, rng_geom_seed=7, + is_ragged=True, rng_geom_seed=7, ) plain = ExecConfig(**base) plain.fill_derived_fields() @@ -388,6 +390,17 @@ def test_ragged_token_gap_stable_under_stride_overrides(): assert pinned.stride_q == pinned_q assert (pinned.stride_k, pinned.stride_v, pinned.stride_o) == (plain.stride_k, plain.stride_v, plain.stride_o) + # Default-on: seed 7 draws at least one non-packed layout for plain ragged. + packed = {n: compute_packed_strides(getattr(plain, f"shape_{n}")) for n in ("q", "k", "v", "o")} + assert any(getattr(plain, f"stride_{n}") != packed[n] for n in ("q", "k", "v", "o")) + + # Auto-packed fallbacks: cu / multiplier offset forms (#538) and 1-byte + # (fp8) data types (#537) derive packed strides regardless of the default. + for override in (dict(is_cu_seq_len=True), dict(with_ragged_offset_multiplier=True), dict(data_type=torch.float8_e4m3fn)): + cfg = ExecConfig(**base, **override) + cfg.fill_derived_fields() + assert all(getattr(cfg, f"stride_{n}") == packed[n] for n in ("q", "k", "v", "o")), override + @pytest.mark.parametrize("test_no", generate_test_seeds(num_tests=128, rng_seed=888), ids=lambda p: f"test{p[0]}") @pytest.mark.L1 @@ -482,7 +495,6 @@ def test_sdpa_random_bwd_ragged_L0(env_info, test_no, request, cudnn_handle): is_ragged_or_padded_or_full=RandomChoice({"ragged" : 1, "padded" : 0, "full" : 0}), is_deterministic=RandomChoice({True : 3, False : 1}), ragged_stats_layout=RandomChoice({"token_major" : 1, "head_major" : 1}), - with_ragged_token_gap=RandomChoice({True : 1}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed)