Skip to content

test(sdpa): fuzz per-tensor ragged token-stride gaps in test_mhas_v2 #516

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

Merged
vedaanta merged 4 commits into from
Aug 11, 2026
Merged

test(sdpa): fuzz per-tensor ragged token-stride gaps in test_mhas_v2 #516

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8499013
test(sdpa): fuzz per-tensor ragged token-stride gaps in test_mhas_v2
vedaanta Aug 7, 2026
b7d8c10
test(sdpa): simplify ragged token gaps to whole tokens; always-on in …
vedaanta Aug 8, 2026
17be556
test(sdpa): draw all four ragged token gaps up front; drop unused bin…
vedaanta Aug 11, 2026
746ef17
test(sdpa): ragged token gaps on by default; auto-packed where not ye…
vedaanta Aug 11, 2026
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
33 changes: 21 additions & 12 deletions test/python/sdpa/fp16.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,21 +140,28 @@ 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)
allocs[TensorUid.stats] = alloc_tensor((max_t_q, cfg.h_q, 1), torch.float32, strides=stats_strides) if cfg.is_train else (None, None, None)
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)
Expand Down Expand Up @@ -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)
Expand Down
80 changes: 71 additions & 9 deletions test/python/sdpa/random_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
_, h, s, d = shape
token_stride = h * d + token_gap
return (s * token_stride, d, token_stride, 1)


@dataclass
Expand Down Expand Up @@ -114,6 +122,21 @@ 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 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 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
Comment on lines +125 to +139

@coderabbitai coderabbitai Bot Aug 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Do not default-enable a layout that FROST executes incorrectly.

The PR objective states that FROST THD forward engines address non-packed strides incorrectly and return wrong O values. Line 139 enables these layouts for all applicable ragged configurations, but no engine-side decline prevents FROST selection.

Add an eligibility decline for affected FROST engines before this default is enabled. Otherwise, keep with_ragged_token_gap disabled by default until the engine fix is available.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@test/python/sdpa/random_config.py` around lines 125 - 139, Prevent affected
FROST forward engines from being selected for configurations with non-packed
ragged strides before enabling with_ragged_token_gap by default. Add the
eligibility decline in the engine-selection or validation path used by
fill_derived_fields, and preserve packed-layout eligibility; if no such guard
exists, change with_ragged_token_gap back to False until the FROST fix is
available.

rescale_threshold: float = None

diag_align: cudnn.diagonal_alignment = None
Expand Down Expand Up @@ -183,16 +206,55 @@ 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 (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
# 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_q = gap_k = gap_v = gap_o = compute_packed_strides
else:
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 = stride_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 = stride_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 = stride_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 = stride_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:
self.stride_stats = stride_fn(self.shape_stats)

Expand Down
37 changes: 37 additions & 0 deletions test/python/test_mhas_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -365,6 +365,43 @@ 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). 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, 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)

# 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
def test_sdpa_random_fwd_ragged_unified_L1(env_info, test_no, request, cudnn_handle):
Expand Down