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..6a1aca02e 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) + _, h, s, d = shape + token_stride = h * d + token_gap + return (s * token_stride, d, token_stride, 1) @dataclass @@ -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 rescale_threshold: float = None diag_align: cudnn.diagonal_alignment = None @@ -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) diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 61f22b384..129313552 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -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):