From 3b8b8aa06c501c3eb3fca35b4983aaf2119142cd Mon Sep 17 00:00:00 2001 From: randers Date: Wed, 23 Sep 2026 01:43:37 +0000 Subject: [PATCH 1/3] frost(sdpa): collapse the per-kernel MASK_FORM constant -- apply_mask_chunk IS the bit-word mask op #1192 (SM107) and #1197 (SM100) put the register-to-predicate mask behind a per-kernel `MASK_FORM = MASK_FORM_BITS` constant and an `apply_mask_chunk_form` dispatcher so each arch line could be A/B'd and rolled out separately. Both lines ship the bit-word form now, so the selector is dead weight: 27 module constants, 46 `form=MASK_FORM,` kwargs and a dispatcher that could only ever pick one arm. - `tile_dsl/mask.py`: `apply_mask_chunk` takes the bit-word body (band edges -> keep-word per 32 columns via saturating shifts -> one `arith.select` per cell, R2P + FSEL in SASS), same signature, same trace-time `MASK_BOUND_LIMIT` guard (messages now say `apply_mask_chunk:`). `apply_mask_chunk_bits`, `apply_mask_chunk_form`, `MASK_FORM_CELLS` / `MASK_FORM_BITS` / `MASK_FORMS` and the per-cell compare + select body are deleted; the header comment keeps the measured record of the form it replaced. `keep_below_word`, `keep_from_word`, `band_mask_words`, `apply_mask_words`, `MASK_WORD_COLS`, `MASK_BOUND_LIMIT` unchanged. - 26 kernels (12 sm107 prefill, 12 sm100 prefill, sm100 decode d128, bwd d512 stage 2): import `apply_mask_chunk`, drop the constant + its comment block, every site `apply_mask_chunk_form(..., form=MASK_FORM)` -> `apply_mask_chunk(...)` (site counts unchanged: 19 sm107, 25 sm100 prefill, 1 decode, 1 bwd). The five d512 `tcgen05_wait(LOAD)` comments and `config_sm100.canonicalize_d192_lowering` now say "bit-word form" / `apply_mask_chunk` instead of the retired names. - Pins rewritten for the new invariant: `test_sm107_every_mask_site_calls_apply_mask_chunk` / `test_sm100_every_mask_site_calls_apply_mask_chunk` (every masked site is a direct call, no `form=`, no module `MASK_FORM`, migration-table counts kept), `test_the_mask_op_has_one_form` (the vocabulary stays gone), the domain-guard test calls `apply_mask_chunk`. The sm_107a / sm_100a SASS pins (R2P > 0, ISETP / predicate-to-register / spill ceilings) are untouched -- they are form-independent. No behaviour change: the masked fp8 d128 causal+padded specialization trace-compiled on this box before and after is byte-identical -- sm_107a cubin md5 7f37f174467d399c8fcfaad019f28d62 (R2P 32 / ISETP 96), sm_100a cubin md5 e3424f1e441770396c6885a7185ba044 (R2P 32 / ISETP 94). Host tests on the worktree: test_tile_dsl_mask_bits.py + test_import_boundaries.py 59 passed; sm107 mask-site + SASS pins 20 passed; test_sm100_mask_form.py 18 passed. Co-Authored-By: Claude Fable 5.1 --- python/cudnn/frost/tile_dsl/mask.py | 195 ++++++------------ .../sdpa/bwd/kernels/sm100/bprop_d512_f16.py | 15 +- python/cudnn/sdpa/fwd/config_sm100.py | 4 +- .../sdpa/fwd/kernels/sm100/decode_d128_f16.py | 15 +- .../fwd/kernels/sm100/prefill_d128_f16.py | 10 +- .../fwd/kernels/sm100/prefill_d128_fp8.py | 10 +- .../fwd/kernels/sm100/prefill_d128_mxfp8.py | 13 +- .../kernels/sm100/prefill_d192_d128_f16.py | 22 +- .../kernels/sm100/prefill_d192_d128_fp8.py | 13 +- .../kernels/sm100/prefill_d192_d128_mxfp8.py | 16 +- .../fwd/kernels/sm100/prefill_d256_f16.py | 13 +- .../fwd/kernels/sm100/prefill_d256_fp8.py | 19 +- .../fwd/kernels/sm100/prefill_d256_mxfp8.py | 16 +- .../fwd/kernels/sm100/prefill_d512_f16.py | 17 +- .../fwd/kernels/sm100/prefill_d512_fp8.py | 17 +- .../fwd/kernels/sm100/prefill_d512_mxfp8.py | 21 +- .../fwd/kernels/sm107/prefill_d128_f16.py | 15 +- .../fwd/kernels/sm107/prefill_d128_fp8.py | 15 +- .../fwd/kernels/sm107/prefill_d128_mxfp8.py | 18 +- .../kernels/sm107/prefill_d192_d128_f16.py | 15 +- .../kernels/sm107/prefill_d192_d128_fp8.py | 15 +- .../kernels/sm107/prefill_d192_d128_mxfp8.py | 18 +- .../fwd/kernels/sm107/prefill_d256_f16.py | 18 +- .../fwd/kernels/sm107/prefill_d256_fp8.py | 18 +- .../fwd/kernels/sm107/prefill_d256_mxfp8.py | 24 +-- .../fwd/kernels/sm107/prefill_d512_f16.py | 19 +- .../fwd/kernels/sm107/prefill_d512_fp8.py | 19 +- .../fwd/kernels/sm107/prefill_d512_mxfp8.py | 19 +- test/python/sdpa/frost/frost_test_utils.py | 6 +- .../sdpa/frost/test_sdpa_fwd_dsl_sm107.py | 65 +++--- .../python/sdpa/frost/test_sm100_mask_form.py | 89 ++++---- .../sdpa/frost/test_tile_dsl_mask_bits.py | 65 +++--- 32 files changed, 245 insertions(+), 609 deletions(-) diff --git a/python/cudnn/frost/tile_dsl/mask.py b/python/cudnn/frost/tile_dsl/mask.py index 4ace932c9..d5f3d6fdf 100644 --- a/python/cudnn/frost/tile_dsl/mask.py +++ b/python/cudnn/frost/tile_dsl/mask.py @@ -14,117 +14,59 @@ _NEG_INF_BITS = -3.4028235e38 -def apply_mask_chunk( - reg_S, - q_abs, - kv_col_base, - seq_kv_len, - window_left: int, - mask_flags: int, - N: int = 64, - bottom_right: int = 0, - causal_diag=None, - mask_value: float = _NEG_INF_BITS, - window_right: int = 0, -): - # mask_value: what a masked score becomes. Default is the legacy finite - # sentinel; the f16 prefill kernels pass float("-inf") so a fully-masked - # row's max stays -inf under any scale and the canonical - # `max == -inf -> substitute 0` guard (row_max_for_exp2) applies. - if cutlass.const_expr(mask_flags == MASK_NONE): - return reg_S - - neg_inf = cutlass.Float32(mask_value) - # The whole band shifts with the diagonal: under BOTTOM_RIGHT the SWA - # lower limit is q + (S_kv - S_q) - W — the same causal_diag offset the - # upper (causal) limit uses below. Top-left keeps the plain q - W. - q_minus_w = None - if mask_flags & MASK_SWA: - swa_base = (q_abs + causal_diag) if bottom_right else q_abs - q_minus_w = swa_base - cutlass.Int32(window_left) - # window_right is the compile-time diagonal-band right bound (cuDNN - # diagonal_band_right_bound): kv columns up to q + window_right (plus the - # bottom-right diagonal offset) stay unmasked. 0 = plain causal. - if cutlass.const_expr((mask_flags & MASK_CAUSAL) and bottom_right): - q_caus_lim = q_abs + causal_diag - else: - q_caus_lim = q_abs - if cutlass.const_expr((mask_flags & MASK_CAUSAL) and window_right != 0): - q_caus_lim = q_caus_lim + cutlass.Int32(window_right) - - elems = [] - for i in range(N): - kv_abs = kv_col_base + cutlass.Int32(i) - masked = None - if cutlass.const_expr(mask_flags & MASK_PADDED): - term = kv_abs >= seq_kv_len - masked = term if masked is None else (masked | term) - if cutlass.const_expr(mask_flags & MASK_CAUSAL): - term = kv_abs > q_caus_lim - masked = term if masked is None else (masked | term) - if cutlass.const_expr(mask_flags & MASK_SWA): - term = kv_abs < q_minus_w - masked = term if masked is None else (masked | term) - val = cutlass.Float32( - arith.select( - masked.ir_value(), - neg_inf.ir_value(), - reg_S[i].ir_value(), - ) - ) - elems.append(val) - return cutlass.Vector.from_elements(tuple(elems), cutlass.Float32) - - # --------------------------------------------------------------------------- -# Per-cell mask, "bits" form: one keep-word per 32 columns, register-to-predicate +# Per-cell mask: one keep-word per 32 columns, register-to-predicate select # --------------------------------------------------------------------------- # -# `apply_mask_chunk` above compares EVERY cell against every active bound and -# ORs the terms: per cell one IADD (kv_col_base + i, never folded into the -# compare), one ISETP per term and one FSEL per term (the i1 OR lowers to nested -# selects), so a causal+SWA tile costs 5 instructions per cell and a padded -# causal+SWA one 7 -- 51-72 % of a masked softmax tile's instructions, serialized -# in front of the exp burst (sm_107a listings, 2026-09-22: a masked KV tile ran -# 805-1602 instructions per lane against 358-451 for a dense one). The information -# content of any of these masks is one or two band EDGES per row, so the form -# below spends its instructions there instead: +# `apply_mask_chunk` below masks a register chunk of S in the BIT-WORD form. The +# information content of a padded / causal / sliding-window mask is one or two +# band EDGES per row, so that is where the instructions go: # # 1. per 32-column word, build the KEEP mask from the band edges with two # saturating shifts (`shr.u32` for the upper edge, `shl.b32` for the lower; # PTX clamps a shift amount >= 32 to a zero result, which is exactly the # "edge is outside this word" case) -- ~7 integer ops per word, per lane; -# 2. per cell, test one bit of the word and feed the SAME `arith.select` -# `apply_mask_chunk` uses. ptxas turns 32 consecutive bit tests into -# `R2P` (register -> 7 predicates) + one `FSEL` per cell, so the per-cell -# cost is 1.4-1.6 instructions regardless of how many mask terms are on -# (sm_107a / sm_100a micro-kernel listings, 2026-09-22, against 3.0-7.6 for -# the per-cell form; the same idiom the block-sparse-attention SM90 backward -# uses, `predicate_bitmask_below`). +# 2. per cell, test one bit of the word and feed it to an `arith.select`. +# ptxas turns 32 consecutive bit tests into `R2P` (register -> 7 predicates) +# + one `FSEL` per cell, so the per-cell cost is 1.4-1.6 instructions +# regardless of how many mask terms are on (sm_107a / sm_100a micro-kernel +# listings, 2026-09-22; the same idiom the block-sparse-attention SM90 +# backward uses, `predicate_bitmask_below`). +# +# The form it replaced (shipped until PR #1192 / #1197 behind a per-kernel +# `MASK_FORM` constant, collapsed since) compared EVERY cell against every active +# bound and ORed the terms: per cell one IADD (kv_col_base + i, never folded into +# the compare), one ISETP per term and one FSEL per term (the i1 OR lowers to +# nested selects), so a causal+SWA tile cost 5 instructions per cell and a padded +# causal+SWA one 7 -- 51-72 % of a masked softmax tile's instructions, serialized +# in front of the exp burst (sm_107a listings, 2026-09-22: a masked KV tile ran +# 805-1602 instructions per lane against 358-451 for a dense one; one build had +# spilled its 128 live i1 values into GPR bits). Both forms mask the same set +# with the same sentinel through the same `arith.select`, so O / LSE were bitwise +# identical across the switch; `test_tile_dsl_mask_bits.py` keeps a host emulation +# of the per-cell formula as the reference this op is pinned against, and the +# sm_107a / sm_100a SASS pins hold the lowering (R2P > 0, ISETP within the +# measured ceiling). # -# Semantics contract (shared with `apply_mask_chunk`, bitwise): +# Semantics contract: # - a masked cell becomes `mask_value` (default: the finite fp32-min sentinel; # pass float("-inf") for the true -inf form); an unmasked cell is passed # through untouched (NaN payloads included); # - bounds are PER LANE (one lane = one q row): `lo`/`hi` are Int32 values # derived from that lane's q row, `kv_col_base` is the chunk's first # absolute kv column; a fully-masked row yields `mask_value` in every cell, -# so the caller's `row_max_for_exp2` guard applies exactly as today; +# so the caller's `row_max_for_exp2` guard applies exactly as before; # - a `tcgen05.ld 32x32b` chunk is contiguous along columns per lane (register # k = column k), so NO column remap is needed here -- unlike the WGMMA # accumulator layout the BSA kernel remaps with `sm90_col_to_predicate_idx`. -MASK_FORM_CELLS = "cells" # apply_mask_chunk: per-cell compare + select -MASK_FORM_BITS = "bits" # apply_mask_chunk_bits: keep-word + register-to-predicate select -MASK_FORMS = (MASK_FORM_CELLS, MASK_FORM_BITS) - MASK_WORD_COLS = 32 # columns per keep-word = bits per register # The band arithmetic (`lo - kv_col_base`, `hi - kv_col_base`, the per-word shift count) is Int32. # With every absolute row / column index below 2**28 (a TMA coordinate keeps S far under that) a # window bound below this limit cannot wrap it; a bound at or past it can, and a wrapped `lo` makes -# the bits form mask EVERYTHING where `apply_mask_chunk` masks nothing. Enforced at trace time on -# the Python-int bounds (no instruction); pinned by `test_bits_form_domain_guard`. +# the op mask EVERYTHING where the per-cell formula masks nothing. Enforced at trace time on the +# Python-int bounds (no instruction); pinned by `test_bits_form_domain_guard`. MASK_BOUND_LIMIT = 1 << 30 @@ -167,9 +109,9 @@ def band_mask_words(lo, hi, kv_col_base, n_cols: int): for a side no mask term defines -- that side folds out at trace time (one shift per word for a one-sided mask, two for a band), never an INT_MIN/INT_MAX sentinel, so ``lo - base`` cannot wrap as long as the caller keeps every bound within ``MASK_BOUND_LIMIT`` of the chunk - (:func:`apply_mask_chunk_bits` refuses a wider window). Returns a tuple of + (:func:`apply_mask_chunk` refuses a wider window). Returns a tuple of ``ceil(n_cols / 32)`` ``Uint32`` words. Trace-time helper (plain Python over traced - values), like :func:`apply_mask_chunk`.""" + values).""" if lo is None and hi is None: raise ValueError("band_mask_words: at least one of lo / hi must be given (a mask with no edge is no mask)") n_words = (n_cols + MASK_WORD_COLS - 1) // MASK_WORD_COLS @@ -193,9 +135,10 @@ def apply_mask_words(reg_S, words, mask_value: float = _NEG_INF_BITS, n_cols: in ``words`` is the tuple :func:`band_mask_words` returns (bit ``i`` of word ``s`` = column ``32 s + i`` is KEPT). The inner loop is a Python ``range`` over a compile-time bit index, which is what lets ptxas see 32 consecutive single-bit tests of one register and - emit ``R2P`` + one ``FSEL`` per cell. Same ``arith.select`` as :func:`apply_mask_chunk` - (only the predicate derivation differs), so the two forms are bitwise identical for the - same mask set, for either sentinel, NaN payloads included.""" + emit ``R2P`` + one ``FSEL`` per cell. One ``arith.select`` per cell, so a kept cell passes + through bit-exactly (NaN payloads included) and a masked one is exactly ``mask_value``, + for either sentinel -- the same select the per-cell compare form used, only the predicate + derivation differs.""" if n_cols is None: n_cols = MASK_WORD_COLS * len(words) if len(words) * MASK_WORD_COLS < n_cols: @@ -212,7 +155,7 @@ def apply_mask_words(reg_S, words, mask_value: float = _NEG_INF_BITS, n_cols: in return cutlass.Vector.from_elements(tuple(elems), cutlass.Float32) -def apply_mask_chunk_bits( +def apply_mask_chunk( reg_S, q_abs, kv_col_base, @@ -225,27 +168,36 @@ def apply_mask_chunk_bits( mask_value: float = _NEG_INF_BITS, window_right: int = 0, ): - """:func:`apply_mask_chunk` with the same signature and the same masked set, in the - "bits" form. The three terms map onto one band ``[lo, hi)`` per lane: - - - CAUSAL masks ``kv > q_caus_lim`` -> ``hi = q_caus_lim + 1`` with the same - ``q_caus_lim = q_abs (+ causal_diag under bottom_right) (+ window_right)``; - - PADDED masks ``kv >= seq_kv_len`` -> ``hi = min(hi, seq_kv_len)``; - - SWA masks ``kv < q_minus_w`` -> ``lo = q_minus_w`` (the same bottom-right anchor). + """Mask the ``N``-wide fp32 register chunk ``reg_S`` (one lane = one q row ``q_abs``, columns + ``kv_col_base .. kv_col_base + N``) for the compile-time ``mask_flags`` set, in the bit-word form. - A side no term defines is ``None`` and folds out. Fully-masked rows, the sentinel and - the caller's tile-level trimming are exactly as for :func:`apply_mask_chunk`. + The three terms map onto one band ``[lo, hi)`` per lane: - Domain: a compile-time ``window_left`` / ``window_right`` at or past ``MASK_BOUND_LIMIT`` - raises at trace time (see the constant); the per-cell form has no such limit because it - never subtracts the chunk base.""" + - CAUSAL masks ``kv > q_caus_lim`` -> ``hi = q_caus_lim + 1`` with + ``q_caus_lim = q_abs (+ causal_diag under bottom_right) (+ window_right)``. ``window_right`` + is the compile-time diagonal-band right bound (cuDNN diagonal_band_right_bound): kv columns + up to ``q + window_right`` (plus the bottom-right diagonal offset) stay unmasked; 0 = plain causal; + - PADDED masks ``kv >= seq_kv_len`` -> ``hi = min(hi, seq_kv_len)``; + - SWA masks ``kv < q_minus_w`` -> ``lo = q_minus_w = swa_base - window_left`` with the same + bottom-right anchor (``swa_base = q_abs + causal_diag`` under bottom_right, else ``q_abs``) -- + the whole band shifts with the diagonal. + + A side no term defines is ``None`` and folds out. ``mask_value`` is what a masked score becomes: + the default is the legacy finite sentinel; the f16 prefill kernels pass ``float("-inf")`` so a + fully-masked row's max stays -inf under any scale and the canonical ``max == -inf -> substitute 0`` + guard (``row_max_for_exp2``) applies. ``MASK_NONE`` returns ``reg_S`` itself. + + Domain: a compile-time ``window_left`` / ``window_right`` at or past ``MASK_BOUND_LIMIT`` raises + at trace time (see the constant): the band arithmetic is Int32 and a wrapped ``lo`` would mask the + whole chunk. Callers that invent a "no sequence reaches this" band (the dense per-tensor fp8 d192 + lowering) must stay below it -- `config_sm100.canonicalize_d192_lowering` uses ``MASK_BOUND_LIMIT - 1``.""" if cutlass.const_expr(mask_flags == MASK_NONE): return reg_S if (mask_flags & MASK_SWA) and isinstance(window_left, int) and window_left >= MASK_BOUND_LIMIT: - raise ValueError(f"apply_mask_chunk_bits: window_left must be < {MASK_BOUND_LIMIT} (got {window_left}); the Int32 band arithmetic would wrap") + raise ValueError(f"apply_mask_chunk: window_left must be < {MASK_BOUND_LIMIT} (got {window_left}); the Int32 band arithmetic would wrap") if (mask_flags & MASK_CAUSAL) and isinstance(window_right, int) and window_right >= MASK_BOUND_LIMIT: - raise ValueError(f"apply_mask_chunk_bits: window_right must be < {MASK_BOUND_LIMIT} (got {window_right}); the Int32 band arithmetic would wrap") + raise ValueError(f"apply_mask_chunk: window_right must be < {MASK_BOUND_LIMIT} (got {window_right}); the Int32 band arithmetic would wrap") lo = None hi = None @@ -264,35 +216,6 @@ def apply_mask_chunk_bits( return apply_mask_words(reg_S, words, mask_value=mask_value, n_cols=N) -def apply_mask_chunk_form( - reg_S, - q_abs, - kv_col_base, - seq_kv_len, - window_left: int, - mask_flags: int, - N: int = 64, - bottom_right: int = 0, - causal_diag=None, - mask_value: float = _NEG_INF_BITS, - window_right: int = 0, - *, - form: str, -): - """:func:`apply_mask_chunk` (``form=MASK_FORM_CELLS``) or :func:`apply_mask_chunk_bits` - (``form=MASK_FORM_BITS``) behind one signature. A kernel picks the form with ONE - module-level constant (``MASK_FORM``) that every masked call site passes -- the - ``DESC_VERSION`` discipline -- so an A/B of the two lowerings is a constant flip and a - test can count call sites against the constant. ``form`` is keyword-only and has no - default: a call site that forgets it fails at trace time instead of silently picking one.""" - kw = dict(N=N, bottom_right=bottom_right, causal_diag=causal_diag, mask_value=mask_value, window_right=window_right) - if form == MASK_FORM_BITS: - return apply_mask_chunk_bits(reg_S, q_abs, kv_col_base, seq_kv_len, window_left, mask_flags, **kw) - if form == MASK_FORM_CELLS: - return apply_mask_chunk(reg_S, q_abs, kv_col_base, seq_kv_len, window_left, mask_flags, **kw) - raise ValueError(f"apply_mask_chunk_form: form must be one of {MASK_FORMS}, got {form!r}") - - # --------------------------------------------------------------------------- # Tile-level mask bounds: which kv TILES a q tile has to visit at all. # --------------------------------------------------------------------------- diff --git a/python/cudnn/sdpa/bwd/kernels/sm100/bprop_d512_f16.py b/python/cudnn/sdpa/bwd/kernels/sm100/bprop_d512_f16.py index 98c4832e3..fb580dda2 100644 --- a/python/cudnn/sdpa/bwd/kernels/sm100/bprop_d512_f16.py +++ b/python/cudnn/sdpa/bwd/kernels/sm100/bprop_d512_f16.py @@ -60,7 +60,7 @@ cga_wait, ) from cudnn.frost.tile_dsl.handles import GmemTileTma, MmaDesc, SmemTile, tma_slice_runtime_desc -from cudnn.frost.tile_dsl.mask import MASK_CAUSAL, MASK_FORM_BITS, MASK_NONE, MASK_PADDED, apply_mask_chunk_form, compute_kv_loop_bounds +from cudnn.frost.tile_dsl.mask import MASK_CAUSAL, MASK_NONE, MASK_PADDED, apply_mask_chunk, compute_kv_loop_bounds from cudnn.frost.tile_dsl.mma import mma_ts from cudnn.frost.tile_dsl.scheduler import Sched, read_tile_id_arrive, scheduler_warp_loop, scheduler_warp_loop_persistent, read_clc_payload from cudnn.frost.tile_dsl.tma import ( @@ -86,16 +86,6 @@ xfer_bytes, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# the masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel (0.0 here: S is masked AFTER exp2), so -# S and every gradient downstream of it are bitwise identical; -# test_sm100_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # Injected by the loader before this body executes; a plain import gets the # all-defaults config (dense bf16), which keeps `python sm100/bprop_d512_f16.py` # usable as a standalone benchmark. @@ -869,7 +859,7 @@ def _compute_kv_iter( # split (sg1 receives the already-masked S). `q_row` is this # CTA's own row; only the tile bound is cluster-wide. if cutlass.const_expr(apply_mask): - s_post = apply_mask_chunk_form( + s_post = apply_mask_chunk( s_post, q_row, kv_loop * cutlass.Int32(CFG.TILE_N) + cutlass.Int32(chunk * S_D_BLOCK), @@ -881,7 +871,6 @@ def _compute_kv_iter( causal_diag=(seqlen_kv - seqlen_q) if CFG.BOTTOM_RIGHT else None, mask_value=0.0, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) # Rows past the real S_q: zero the whole row. Their LSE came from a # CLAMPED index and is meaningless, so the value must be discarded diff --git a/python/cudnn/sdpa/fwd/config_sm100.py b/python/cudnn/sdpa/fwd/config_sm100.py index 5fa662a31..8afa71551 100644 --- a/python/cudnn/sdpa/fwd/config_sm100.py +++ b/python/cudnn/sdpa/fwd/config_sm100.py @@ -1864,11 +1864,11 @@ def canonicalize_d192_lowering( # MASK_NONE x32 path, so the dense plan is lowered as MASK_CAUSAL with a # right band no sequence reaches. The band is a compile-time # `window_right` at the kernel's mask sites, so it must sit INSIDE the - # bits mask form's Int32 domain: `apply_mask_chunk_bits` raises at trace + # bit-word mask op's Int32 domain: `apply_mask_chunk` raises at trace # time from MASK_BOUND_LIMIT (1 << 30) on, and a trace-time raise is a # typed decline at engine.build_plan -- the former `1 << 30` dropped the # FROST fp8 row out of every dense per-tensor d192 graph once the kernels - # masked in that form. MASK_BOUND_LIMIT - 1 still exceeds any dense D192 + # masked in the bit-word form. MASK_BOUND_LIMIT - 1 still exceeds any dense D192 # sequence that fits in SM100 memory while leaving signed-int32 headroom # for q + R, and keeps the module key independent of S_kv. Imported here, # not at module level: tile_dsl.mask imports cutlass, which stays off the diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/decode_d128_f16.py b/python/cudnn/sdpa/fwd/kernels/sm100/decode_d128_f16.py index f55a4f4be..13e8a423a 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/decode_d128_f16.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/decode_d128_f16.py @@ -141,20 +141,10 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm100_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # Storage dtype + MMA kind dispatch — folded at trace time on CFG.DTYPE_QKV. if CFG.DTYPE_QKV == 2: STORAGE_DTYPE = cutlass.BFloat16 @@ -1250,7 +1240,7 @@ def _softmax_kv_body( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1262,7 +1252,6 @@ def _softmax_kv_body( causal_diag=causal_diag, mask_value=float("-inf"), window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_f16.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_f16.py index a9c5d7c75..db4fbce9b 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_f16.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_f16.py @@ -156,18 +156,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): every masked call site -# below passes `form=MASK_FORM`; both forms mask the same set with the same sentinel, so O / LSE are bitwise identical. -MASK_FORM: str = MASK_FORM_BITS - # Storage dtype + MMA kind dispatch — folded at trace time on CFG.DTYPE_QKV. if CFG.DTYPE_QKV == 2: STORAGE_DTYPE = cutlass.BFloat16 @@ -1698,7 +1693,7 @@ def _softmax_kv_body( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1710,7 +1705,6 @@ def _softmax_kv_body( causal_diag=causal_diag, mask_value=float("-inf"), window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_fp8.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_fp8.py index dc2ecadad..2a1569d9c 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_fp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_fp8.py @@ -124,18 +124,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): every masked call site -# below passes `form=MASK_FORM`; both forms mask the same set with the same sentinel, so O / LSE are bitwise identical. -MASK_FORM: str = MASK_FORM_BITS - # Storage dtype + MMA kind dispatch keyed off CFG.DTYPE_QKV. if CFG.DTYPE_QKV == 0: STORAGE_DTYPE = cutlass.Float8E4M3FN @@ -1496,7 +1491,7 @@ def _softmax_kv_body( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1507,7 +1502,6 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_mxfp8.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_mxfp8.py index 03a31c01d..d08b7b5b1 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_mxfp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d128_mxfp8.py @@ -151,18 +151,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, GmemTileLinear, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): every masked call site -# below passes `form=MASK_FORM`; both forms mask the same set with the same sentinel, so O / LSE are bitwise identical. -MASK_FORM: str = MASK_FORM_BITS - # MXFP8 storage dtype dispatch — keyed off CFG.DTYPE_QKV (0=E4M3, 1=E5M2). if CFG.DTYPE_QKV == 0: STORAGE_DTYPE = cutlass.Float8E4M3FN @@ -2018,7 +2013,7 @@ def _softmax_kv_body( # Bottom-right causal: runtime SKV-SQ diagonal offset (folds out when # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None - reg_S_a = apply_mask_chunk_form( + reg_S_a = apply_mask_chunk( reg_S_a, q_abs, kv_col_base_a, @@ -2029,9 +2024,8 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) - reg_S_b = apply_mask_chunk_form( + reg_S_b = apply_mask_chunk( reg_S_b, q_abs, kv_col_base_b, @@ -2042,7 +2036,6 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) max_a = row_max_reduction_64(reg_S_a) diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_f16.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_f16.py index 922267fb1..bfc51ebad 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_f16.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_f16.py @@ -122,18 +122,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): every masked call site -# below passes `form=MASK_FORM`; both forms mask the same set with the same sentinel, so O / LSE are bitwise identical. -MASK_FORM: str = MASK_FORM_BITS - from cudnn.block_sparse_attention.csrc.utils.kernel_utils import ex2_emulation_2 _PADDED_CAUSAL = CFG.MASK_FLAGS == (MASK_CAUSAL | MASK_PADDED) and CFG.WINDOW_RIGHT == 0 @@ -623,7 +618,7 @@ def _apply_padding_mask_if_needed(reg_S, kv_col_base, eff_seqlen_kv): """Apply the per-element padding predicate only to a partial KV chunk.""" result = reg_S if kv_col_base + cutlass.Int32(int(reg_S.shape[0])) > eff_seqlen_kv: - result = apply_mask_chunk_form( + result = apply_mask_chunk( reg_S, cutlass.Int32(0), kv_col_base, @@ -632,7 +627,6 @@ def _apply_padding_mask_if_needed(reg_S, kv_col_base, eff_seqlen_kv): MASK_PADDED, N=int(reg_S.shape[0]), mask_value=float("-inf"), - form=MASK_FORM, ) return result @@ -2077,7 +2071,7 @@ def _softmax_kv_body( # exactly the causal OR padded set, so the chunk sees a top-left causal mask # anchored at mask_q_abs (WINDOW_RIGHT == 0 in this arm by _PADDED_CAUSAL). chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], mask_q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -2087,14 +2081,13 @@ def _softmax_kv_body( N=CHUNK, mask_value=float("-inf"), window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] elif cutlass.const_expr(CFG.MASK_FLAGS == MASK_CAUSAL and CFG.BOTTOM_RIGHT == 0): # Top-left causal (+ the compile-time right band, cuDNN diagonal_band_right_bound). chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -2104,14 +2097,13 @@ def _softmax_kv_body( N=CHUNK, mask_value=float("-inf"), window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] elif cutlass.const_expr(CFG.MASK_FLAGS == MASK_CAUSAL and CFG.BOTTOM_RIGHT != 0): # Bottom-right causal: the diagonal sits causal_diag = S_kv - S_q columns right of top-left. chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -2123,14 +2115,13 @@ def _softmax_kv_body( causal_diag=causal_diag, mask_value=float("-inf"), window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] else: chunk_mask_flags = body_mask_flags & ~MASK_PADDED if CFG.MASK_FLAGS & MASK_SWA else body_mask_flags chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -2142,7 +2133,6 @@ def _softmax_kv_body( causal_diag=causal_diag, mask_value=float("-inf"), window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_fp8.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_fp8.py index cc99cb22d..4b4297363 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_fp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_fp8.py @@ -93,18 +93,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): every masked call site -# below passes `form=MASK_FORM`; both forms mask the same set with the same sentinel, so O / LSE are bitwise identical. -MASK_FORM: str = MASK_FORM_BITS - from cudnn.block_sparse_attention.csrc.utils.kernel_utils import ex2_emulation_2 _PADDED_CAUSAL = CFG.MASK_FLAGS == (MASK_CAUSAL | MASK_PADDED) and CFG.WINDOW_RIGHT == 0 @@ -196,7 +191,7 @@ def _apply_padding_mask_if_needed(reg_s, kv_col_base, eff_seqlen_kv, mask_value: """Apply the per-element padding predicate only to a partial KV chunk.""" result = reg_s if kv_col_base + cutlass.Int32(int(reg_s.shape[0])) > eff_seqlen_kv: - result = apply_mask_chunk_form( + result = apply_mask_chunk( reg_s, cutlass.Int32(0), kv_col_base, @@ -205,7 +200,6 @@ def _apply_padding_mask_if_needed(reg_s, kv_col_base, eff_seqlen_kv, mask_value: MASK_PADDED, N=int(reg_s.shape[0]), mask_value=mask_value, - form=MASK_FORM, ) return result @@ -1901,7 +1895,7 @@ def _softmax_kv_body( mask_q_abs = cute.math.min(q_abs, eff_seqlen_kv - cutlass.Int32(1)) chunk_mask_flags = MASK_CAUSAL chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], mask_q_abs - (kv_col_base + cutlass.Int32(c * CHUNK)), cutlass.Int32(0), @@ -1911,7 +1905,6 @@ def _softmax_kv_body( N=CHUNK, mask_value=mask_value, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_mxfp8.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_mxfp8.py index 9a9140182..125436a2a 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_mxfp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d192_d128_mxfp8.py @@ -123,18 +123,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): every masked call site -# below passes `form=MASK_FORM`; both forms mask the same set with the same sentinel, so O / LSE are bitwise identical. -MASK_FORM: str = MASK_FORM_BITS - _PADDED_CAUSAL = CFG.MASK_FLAGS == (MASK_CAUSAL | MASK_PADDED) and CFG.WINDOW_RIGHT == 0 _THD_HEAVY_ROWS_FIRST = CFG.THD_VARLEN and bool(CFG.MASK_FLAGS & MASK_CAUSAL) _DENSE_NATURAL_HEAVY_ROWS_FIRST = ( @@ -212,7 +207,7 @@ def _apply_padding_mask_if_needed(reg_s, kv_col_base, eff_seqlen_kv): """Apply the per-element padding predicate only to a partial KV chunk.""" result = reg_s if kv_col_base + cutlass.Int32(int(reg_s.shape[0])) > eff_seqlen_kv: - result = apply_mask_chunk_form( + result = apply_mask_chunk( reg_s, cutlass.Int32(0), kv_col_base, @@ -220,7 +215,6 @@ def _apply_padding_mask_if_needed(reg_s, kv_col_base, eff_seqlen_kv): 0, MASK_PADDED, N=int(reg_s.shape[0]), - form=MASK_FORM, ) return result @@ -2465,7 +2459,7 @@ def _softmax_kv_body( mask_causal_diag = None elif cutlass.const_expr(CFG.MASK_FLAGS & MASK_SWA): mask_flags = body_mask_flags & ~MASK_PADDED - reg_S_a = apply_mask_chunk_form( + reg_S_a = apply_mask_chunk( reg_S_a, mask_q_abs - kv_col_base_a, cutlass.Int32(0), @@ -2476,12 +2470,11 @@ def _softmax_kv_body( bottom_right=mask_bottom_right, causal_diag=mask_causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) if cutlass.const_expr(may_need_padding and (CFG.MASK_FLAGS & MASK_PADDED) and (CFG.MASK_FLAGS & MASK_SWA)): reg_S_a = _apply_padding_mask_if_needed(reg_S_a, kv_col_base_a, eff_seqlen_kv) reg_S_b = _apply_padding_mask_if_needed(reg_S_b, kv_col_base_b, eff_seqlen_kv) - reg_S_b = apply_mask_chunk_form( + reg_S_b = apply_mask_chunk( reg_S_b, mask_q_abs - kv_col_base_b, cutlass.Int32(0), @@ -2492,7 +2485,6 @@ def _softmax_kv_body( bottom_right=mask_bottom_right, causal_diag=mask_causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) max_a = row_max_reduction_64(reg_S_a) diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_f16.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_f16.py index 2ecbf7ed0..29a7fc0eb 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_f16.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_f16.py @@ -64,18 +64,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): every masked call site -# below passes `form=MASK_FORM`; both forms mask the same set with the same sentinel, so O / LSE are bitwise identical. -MASK_FORM: str = MASK_FORM_BITS - if CFG.DTYPE_QKV == 2: STORAGE_DTYPE = cutlass.BFloat16 P_STORAGE_DTYPE = cutlass.BFloat16 @@ -1349,7 +1344,7 @@ def _softmax_warp_group( ] causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1361,7 +1356,6 @@ def _softmax_warp_group( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(N_CHUNKS) ] @@ -1503,7 +1497,7 @@ def _softmax_warp_group( ] causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1515,7 +1509,6 @@ def _softmax_warp_group( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_fp8.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_fp8.py index c8e2f61d3..453708fc5 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_fp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_fp8.py @@ -73,17 +73,12 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_CAUSAL, MASK_PADDED, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): every masked call site -# below passes `form=MASK_FORM`; both forms mask the same set with the same sentinel, so O / LSE are bitwise identical. -MASK_FORM: str = MASK_FORM_BITS - from cudnn.block_sparse_attention.csrc.utils.kernel_utils import ex2_emulation_2 if CFG.DTYPE_QKV == 0: @@ -1818,7 +1813,7 @@ def _softmax_warp_group( for c in range(N_CHUNKS) ] chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], mask_q_abs - (kv_col_base + cutlass.Int32(c * CHUNK)), cutlass.Int32(0), @@ -1830,7 +1825,6 @@ def _softmax_warp_group( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(N_CHUNKS) ] @@ -1861,7 +1855,7 @@ def _softmax_warp_group( # store into cols 96-111 (mb_softmax_hi_loaded, #981). nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) mb_softmax_hi_loaded[parity_rt].arrive() - masked_chunk = apply_mask_chunk_form( + masked_chunk = apply_mask_chunk( raw_chunk, mask_q_abs - (kv_col_base + cutlass.Int32(CHUNK)), cutlass.Int32(0), @@ -1873,7 +1867,6 @@ def _softmax_warp_group( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) reg_S_half = RegTile(masked_chunk, size=CHUNK) @@ -1968,7 +1961,7 @@ def _softmax_warp_group( mask_bottom_right = CFG.BOTTOM_RIGHT causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], mask_q_abs - (kv_col_base + cutlass.Int32(c * CHUNK)), cutlass.Int32(0), @@ -1980,7 +1973,6 @@ def _softmax_warp_group( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(N_CHUNKS) ] @@ -2174,7 +2166,7 @@ def _softmax_warp_group( mask_bottom_right = CFG.BOTTOM_RIGHT causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], mask_q_abs - (kv_col_base + cutlass.Int32(c * CHUNK)), cutlass.Int32(0), @@ -2186,7 +2178,6 @@ def _softmax_warp_group( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_mxfp8.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_mxfp8.py index b524d967f..d6b1691e5 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_mxfp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d256_mxfp8.py @@ -71,18 +71,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_CAUSAL, MASK_PADDED, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): every masked call site -# below passes `form=MASK_FORM`; both forms mask the same set with the same sentinel, so O / LSE are bitwise identical. -MASK_FORM: str = MASK_FORM_BITS - from cudnn.block_sparse_attention.csrc.utils.kernel_utils import ex2_emulation_2 if CFG.DTYPE_QKV == 0: @@ -2185,7 +2180,7 @@ def _softmax_warp_group( mask_bottom_right = CFG.BOTTOM_RIGHT causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs - (kv_col_base + cutlass.Int32(c * CHUNK)), cutlass.Int32(0), @@ -2197,7 +2192,6 @@ def _softmax_warp_group( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(N_CHUNKS) ] @@ -2433,7 +2427,7 @@ def _softmax_warp_group( mask_q_abs = cute.math.min(q_abs, eff_seqlen_kv - cutlass.Int32(1)) mask_flags = MASK_CAUSAL chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], mask_q_abs - (kv_col_base + cutlass.Int32(c * CHUNK)), cutlass.Int32(0), @@ -2445,7 +2439,6 @@ def _softmax_warp_group( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(N_CHUNKS) ] @@ -2711,7 +2704,7 @@ def _fused_p1_step( if cutlass.const_expr(_PADDED_TOP_LEFT_CAUSAL): mask_q_abs = cute.math.min(mask_q_abs, mask_seq_kv - cutlass.Int32(1)) mask_flags = MASK_CAUSAL - raw_hi = apply_mask_chunk_form( + raw_hi = apply_mask_chunk( raw_hi, mask_q_abs, cutlass.Int32(0), @@ -2723,7 +2716,6 @@ def _fused_p1_step( causal_diag=None, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) reg_S_half = RegTile(raw_hi, size=64) diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_f16.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_f16.py index 1bccfa7a8..988c4089a 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_f16.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_f16.py @@ -86,23 +86,13 @@ def _require(cond, msg): from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm100_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - from cudnn.sdpa.fwd.kernels._common_blackwell import ( sdpa_operand_tensors, make_split_helpers, @@ -743,14 +733,14 @@ def _sg0_softmax_kv_iter( # Pin the loads AHEAD of this iteration's `mb_s_acc_empty` arrive. `tcgen05.ld` is asynchronous and # the arrive (below) has no data dependency on the loaded registers, so without this wait ptxas is free to # schedule the arrive between the two chunk loads -- and on the sm107 twin it did, once the mask code got - # shorter (the "bits" form): the parked MMA then overwrites the S parity slot under the still-pending + # shorter (the bit-word mask form): the parked MMA then overwrites the S parity slot under the still-pending # second read, which shows as a two-launch delta on O. The dense arm is ordered by its own wait(LOAD) # right after tmem_load_tile; the d128 / d192 / d256 kernels by their P `tcgen05.st` + `wait(STORE)` # data dependency. One instruction per masked KV tile. nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * SOFTMAX_CHUNK), @@ -762,7 +752,6 @@ def _sg0_softmax_kv_iter( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(SOFTMAX_N_CHUNKS_LOAD) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_fp8.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_fp8.py index 54b824fda..0d4ace0cb 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_fp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_fp8.py @@ -152,23 +152,13 @@ def _require(cond, msg): from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm100_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - from cudnn.sdpa.fwd.kernels._common_blackwell import ( make_split_helpers, KvLoopBounds, @@ -880,14 +870,14 @@ def _sg0_softmax_kv_iter( # Pin the loads AHEAD of this iteration's `mb_s_acc_empty` arrive. `tcgen05.ld` is asynchronous and # the arrive (below) has no data dependency on the loaded registers, so without this wait ptxas is free to # schedule the arrive between the two chunk loads -- and on the sm107 twin it did, once the mask code got - # shorter (the "bits" form): the parked MMA then overwrites the S parity slot under the still-pending + # shorter (the bit-word mask form): the parked MMA then overwrites the S parity slot under the still-pending # second read, which shows as a two-launch delta on O. The dense arm is ordered by its own wait(LOAD) # right after tmem_load_tile; the d128 / d192 / d256 kernels by their P `tcgen05.st` + `wait(STORE)` # data dependency. One instruction per masked KV tile. nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * SOFTMAX_CHUNK), @@ -899,7 +889,6 @@ def _sg0_softmax_kv_iter( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(SOFTMAX_N_CHUNKS_LOAD) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_mxfp8.py b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_mxfp8.py index c15a5f30b..767e4524c 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_mxfp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm100/prefill_d512_mxfp8.py @@ -73,8 +73,7 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_CAUSAL, MASK_PADDED, @@ -82,15 +81,6 @@ ) from cudnn.block_sparse_attention.csrc.utils.kernel_utils import ex2_emulation_2 -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm100_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - if CFG.DTYPE_QKV == 0: STORAGE_DTYPE = cutlass.Float8E4M3FN P_STORAGE_DTYPE = cutlass.Float8E4M3FN @@ -2345,7 +2335,7 @@ def _softmax_warp_group( chunks_S = [_mask_exact_swa_edge(raw_chunks[c], tid_in_wg, c * P_SUBCHUNK, left_edge=True) for c in range(CFG.TILE_N // P_SUBCHUNK)] else: chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs - (kv_col_base + cutlass.Int32(c * CHUNK)), cutlass.Int32(0), @@ -2357,7 +2347,6 @@ def _softmax_warp_group( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(N_CHUNKS) ] @@ -2619,7 +2608,7 @@ def _softmax_warp_group( chunks_S = [_mask_exact_swa_edge(raw_chunks[c], tid_in_wg, c * P_SUBCHUNK, left_edge=False) for c in range(CFG.TILE_N // P_SUBCHUNK)] else: chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], mask_q_abs - (kv_col_base + cutlass.Int32(c * CHUNK)), cutlass.Int32(0), @@ -2631,7 +2620,6 @@ def _softmax_warp_group( causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) for c in range(N_CHUNKS) ] @@ -3040,7 +3028,7 @@ def _fused_p1_step( if cutlass.const_expr(_PADDED_TOP_LEFT_CAUSAL): mask_q_abs = cute.math.min(mask_q_abs, mask_seq_kv - cutlass.Int32(1)) mask_flags = MASK_CAUSAL - raw_hi = apply_mask_chunk_form( + raw_hi = apply_mask_chunk( raw_hi, mask_q_abs, cutlass.Int32(0), @@ -3052,7 +3040,6 @@ def _fused_p1_step( causal_diag=None, window_right=CFG.WINDOW_RIGHT, mask_value=float("-inf"), - form=MASK_FORM, ) reg_S_half = RegTile(raw_hi, size=64) diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_f16.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_f16.py index a16f6be53..fbb740687 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_f16.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_f16.py @@ -147,23 +147,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # Storage dtype + MMA kind dispatch — folded at trace time on CFG.DTYPE_QKV. if CFG.DTYPE_QKV == 2: STORAGE_DTYPE = cutlass.BFloat16 @@ -1384,7 +1374,7 @@ def _softmax_kv_body( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1395,7 +1385,6 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_fp8.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_fp8.py index 2e30fbd32..b6a692d7e 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_fp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_fp8.py @@ -183,23 +183,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # Storage dtype + MMA kind dispatch keyed off CFG.DTYPE_QKV. if CFG.DTYPE_QKV == 0: STORAGE_DTYPE = cutlass.Float8E4M3FN @@ -1731,7 +1721,7 @@ def _softmax_kv_body( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1742,7 +1732,6 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_mxfp8.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_mxfp8.py index c99cc4bda..d70c083e8 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_mxfp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d128_mxfp8.py @@ -149,23 +149,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, GmemTileLinear, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # MXFP8 storage dtype dispatch — keyed off CFG.DTYPE_QKV (0=E4M3, 1=E5M2). if CFG.DTYPE_QKV == 0: STORAGE_DTYPE = cutlass.Float8E4M3FN @@ -1938,7 +1928,7 @@ def _softmax_kv_body( # Bottom-right causal: runtime SKV-SQ diagonal offset (folds out when # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None - reg_S_a = apply_mask_chunk_form( + reg_S_a = apply_mask_chunk( reg_S_a, q_abs, kv_col_base_a, @@ -1949,9 +1939,8 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) - reg_S_b = apply_mask_chunk_form( + reg_S_b = apply_mask_chunk( reg_S_b, q_abs, kv_col_base_b, @@ -1962,7 +1951,6 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) max_a = row_max_reduction_64(reg_S_a) diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_f16.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_f16.py index 1ff2ca692..9a27b2be1 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_f16.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_f16.py @@ -150,23 +150,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # Storage dtype + MMA kind dispatch — folded at trace time on CFG.DTYPE_QKV. if CFG.DTYPE_QKV == 2: STORAGE_DTYPE = cutlass.BFloat16 @@ -1387,7 +1377,7 @@ def _softmax_kv_body( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1398,7 +1388,6 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_fp8.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_fp8.py index e3c48bb92..7d658496d 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_fp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_fp8.py @@ -198,23 +198,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # Storage dtype + MMA kind dispatch keyed off CFG.DTYPE_QKV. if CFG.DTYPE_QKV == 0: STORAGE_DTYPE = cutlass.Float8E4M3FN @@ -1719,7 +1709,7 @@ def _softmax_kv_body( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1730,7 +1720,6 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_mxfp8.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_mxfp8.py index 59169c46c..89485451b 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_mxfp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d192_d128_mxfp8.py @@ -184,23 +184,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, GmemTileLinear, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # MXFP8 storage dtype dispatch — keyed off CFG.DTYPE_QKV (0=E4M3, 1=E5M2). if CFG.DTYPE_QKV == 0: STORAGE_DTYPE = cutlass.Float8E4M3FN @@ -1946,7 +1936,7 @@ def _softmax_kv_body( # Bottom-right causal: runtime SKV-SQ diagonal offset (folds out when # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None - reg_S_a = apply_mask_chunk_form( + reg_S_a = apply_mask_chunk( reg_S_a, q_abs, kv_col_base_a, @@ -1957,9 +1947,8 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) - reg_S_b = apply_mask_chunk_form( + reg_S_b = apply_mask_chunk( reg_S_b, q_abs, kv_col_base_b, @@ -1970,7 +1959,6 @@ def _softmax_kv_body( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) max_a = row_max_reduction_64(reg_S_a) diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_f16.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_f16.py index 437841d8c..6a31ce927 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_f16.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_f16.py @@ -149,23 +149,13 @@ def _needs_desc_v1(cfg) -> bool: from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - if CFG.DTYPE_QKV == 2: STORAGE_DTYPE = cutlass.BFloat16 P_STORAGE_DTYPE = cutlass.BFloat16 @@ -1476,7 +1466,7 @@ def _softmax_warp_group( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1487,7 +1477,6 @@ def _softmax_warp_group( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] @@ -1622,7 +1611,7 @@ def _softmax_warp_group( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1633,7 +1622,6 @@ def _softmax_warp_group( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_fp8.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_fp8.py index 4e04a2797..fe010789e 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_fp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_fp8.py @@ -179,23 +179,13 @@ def _needs_desc_v1(cfg) -> bool: from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - if CFG.DTYPE_QKV == 0: STORAGE_DTYPE = cutlass.Float8E4M3FN P_STORAGE_DTYPE = cutlass.Float8E4M3FN @@ -1501,7 +1491,7 @@ def _softmax_warp_group( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1512,7 +1502,6 @@ def _softmax_warp_group( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] @@ -1645,7 +1634,7 @@ def _softmax_warp_group( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * CHUNK), @@ -1656,7 +1645,6 @@ def _softmax_warp_group( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(N_CHUNKS) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_mxfp8.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_mxfp8.py index 56bb0f46e..195d9f3a8 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_mxfp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d256_mxfp8.py @@ -184,23 +184,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - if CFG.DTYPE_QKV == 0: STORAGE_DTYPE = cutlass.Float8E4M3FN P_STORAGE_DTYPE = cutlass.Float8E4M3FN @@ -1900,7 +1890,7 @@ def _softmax_warp_group( # Bottom-right causal: runtime SKV-SQ diagonal offset (folds out when # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None - reg_S_a = apply_mask_chunk_form( + reg_S_a = apply_mask_chunk( reg_S_a, q_abs, kv_col_base_a, @@ -1911,9 +1901,8 @@ def _softmax_warp_group( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) - reg_S_b = apply_mask_chunk_form( + reg_S_b = apply_mask_chunk( reg_S_b, q_abs, kv_col_base_b, @@ -1924,7 +1913,6 @@ def _softmax_warp_group( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) max_a = row_max_reduction_64(reg_S_a) max_b = row_max_reduction_64(reg_S_b) @@ -2039,7 +2027,7 @@ def _softmax_warp_group( # Bottom-right causal: runtime SKV-SQ diagonal offset (folds out when # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None - reg_S_a = apply_mask_chunk_form( + reg_S_a = apply_mask_chunk( reg_S_a, q_abs, kv_col_base_a, @@ -2050,9 +2038,8 @@ def _softmax_warp_group( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) - reg_S_b = apply_mask_chunk_form( + reg_S_b = apply_mask_chunk( reg_S_b, q_abs, kv_col_base_b, @@ -2063,7 +2050,6 @@ def _softmax_warp_group( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) max_a = row_max_reduction_64(reg_S_a) max_b = row_max_reduction_64(reg_S_b) diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_f16.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_f16.py index 2abdb1420..55ab589ef 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_f16.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_f16.py @@ -154,23 +154,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # Reuse sm107 SDPA pipeline-shared helpers (KvLoopBounds + closures factory). # Note: dsv4 forks Bars in this file (NOT imported from _sdpa_common) per the # per-pipeline fork pattern in the C++-to-DSL porting notes. @@ -976,8 +966,8 @@ def _sg0_softmax_kv_iter( ] # Pin the loads AHEAD of this iteration's `mb_s_acc_empty` arrive. `tcgen05.ld` is asynchronous and # the arrive (below) has no data dependency on the loaded registers, so without this wait ptxas is free to - # schedule the arrive between the two chunk loads -- and it did, once the mask code got shorter (the "bits" - # form): the parked MMA then overwrites the S parity slot under the still-pending second read, which shows + # schedule the arrive between the two chunk loads -- and it did, once the mask code got shorter (the bit-word + # mask form): the parked MMA then overwrites the S parity slot under the still-pending second read, which shows # as a two-launch delta on O. The dense arm is ordered by its fused inline-asm `tcgen05.ld.red`; the # d128 / d192 / d256 kernels by their P `tcgen05.st` + `wait(STORE)` data dependency. One instruction. nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) @@ -985,7 +975,7 @@ def _sg0_softmax_kv_iter( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * SOFTMAX_CHUNK), @@ -996,7 +986,6 @@ def _sg0_softmax_kv_iter( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(SOFTMAX_N_CHUNKS_LOAD) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_fp8.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_fp8.py index 54768f805..f07c126a2 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_fp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_fp8.py @@ -183,23 +183,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # Reuse sm107 SDPA pipeline-shared helpers (KvLoopBounds + closures factory). # Note: dsv4 forks Bars in this file (NOT imported from _sdpa_common) per the # per-pipeline fork pattern in the C++-to-DSL porting notes. @@ -1050,8 +1040,8 @@ def _sg0_softmax_kv_iter( ] # Pin the loads AHEAD of this iteration's `mb_s_acc_empty` arrive. `tcgen05.ld` is asynchronous and # the arrive (below) has no data dependency on the loaded registers, so without this wait ptxas is free to - # schedule the arrive between the two chunk loads -- and it did, once the mask code got shorter (the "bits" - # form): the parked MMA then overwrites the S parity slot under the still-pending second read, which shows + # schedule the arrive between the two chunk loads -- and it did, once the mask code got shorter (the bit-word + # mask form): the parked MMA then overwrites the S parity slot under the still-pending second read, which shows # as a two-launch delta on O. The dense arm is ordered by its fused inline-asm `tcgen05.ld.red`; the # d128 / d192 / d256 kernels by their P `tcgen05.st` + `wait(STORE)` data dependency. One instruction. nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) @@ -1059,7 +1049,7 @@ def _sg0_softmax_kv_iter( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - eff_seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * SOFTMAX_CHUNK), @@ -1070,7 +1060,6 @@ def _sg0_softmax_kv_iter( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(SOFTMAX_N_CHUNKS_LOAD) ] diff --git a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_mxfp8.py b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_mxfp8.py index bf42c6c41..0ce920cbd 100644 --- a/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_mxfp8.py +++ b/python/cudnn/sdpa/fwd/kernels/sm107/prefill_d512_mxfp8.py @@ -205,23 +205,13 @@ from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, GmemTileTma, tma_slice_runtime_desc from cudnn.frost.tile_dsl.tmem import tmem_alloc, tmem_dealloc from cudnn.frost.tile_dsl.mask import ( - apply_mask_chunk_form, - MASK_FORM_BITS, + apply_mask_chunk, MASK_NONE, MASK_PADDED, MASK_CAUSAL, MASK_SWA, ) -# Per-cell mask lowering, ONE constant per kernel (the DESC_VERSION discipline): -# every masked call site below passes `form=MASK_FORM`, so the two forms of the -# same mask -- "cells" (per-cell compare + select, 3-7 instructions per cell) and -# "bits" (one keep-word per 32 columns, register-to-predicate R2P + one FSEL per -# cell, 1.4-1.6 per cell) -- are an A/B by flipping this line. Both produce the -# same masked set with the same sentinel, so O / LSE are bitwise identical; -# test_sm107_every_mask_site_takes_the_module_mask_form counts the sites. -MASK_FORM: str = MASK_FORM_BITS - # Reuse sm107 SDPA pipeline-shared helpers (KvLoopBounds + closures factory). # Note: dsv4 forks Bars in this file (NOT imported from _sdpa_common) per the # per-pipeline fork pattern in the C++-to-DSL porting notes. @@ -1215,8 +1205,8 @@ def _sg0_softmax_kv_iter( ] # Pin the loads AHEAD of this iteration's `mb_s_acc_empty` arrive. `tcgen05.ld` is asynchronous and # the arrive (below) has no data dependency on the loaded registers, so without this wait ptxas is free to - # schedule the arrive between the two chunk loads -- and it did, once the mask code got shorter (the "bits" - # form): the parked MMA then overwrites the S parity slot under the still-pending second read, which shows + # schedule the arrive between the two chunk loads -- and it did, once the mask code got shorter (the bit-word + # mask form): the parked MMA then overwrites the S parity slot under the still-pending second read, which shows # as a two-launch delta on O. The dense arm is ordered by its fused inline-asm `tcgen05.ld.red`; the # d128 / d192 / d256 kernels by their P `tcgen05.st` + `wait(STORE)` data dependency. One instruction. nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) @@ -1224,7 +1214,7 @@ def _sg0_softmax_kv_iter( # CFG.BOTTOM_RIGHT is 0 — top-left masking is unchanged). causal_diag = eff_seqlen_kv - seqlen_q if cutlass.const_expr(CFG.BOTTOM_RIGHT) else None chunks_S = [ - apply_mask_chunk_form( + apply_mask_chunk( raw_chunks[c], q_abs, kv_col_base + cutlass.Int32(c * SOFTMAX_CHUNK), @@ -1235,7 +1225,6 @@ def _sg0_softmax_kv_iter( bottom_right=CFG.BOTTOM_RIGHT, causal_diag=causal_diag, window_right=CFG.WINDOW_RIGHT, - form=MASK_FORM, ) for c in range(SOFTMAX_N_CHUNKS_LOAD) ] diff --git a/test/python/sdpa/frost/frost_test_utils.py b/test/python/sdpa/frost/frost_test_utils.py index 7c44d0911..a95662043 100644 --- a/test/python/sdpa/frost/frost_test_utils.py +++ b/test/python/sdpa/frost/frost_test_utils.py @@ -279,9 +279,9 @@ def P(t, align=16): "BSSY": ("BSSY",), "SYNCS_ARRIVE": (" SYNCS.ARRIVE",), } -# The masked-softmax-arm pins (`tile_dsl/mask.py`, MASK_FORM): under the "bits" form every masked KV-tile body carries 4 R2P -# per 32-column keep-word and ~0.04 ISETP per cell; under "cells" it carries 0 R2P and one ISETP per cell per mask term, and a -# build that runs out of predicate registers spills them into GPRs through predicate-to-register moves. +# The masked-softmax-arm pins (`tile_dsl.mask.apply_mask_chunk`, the bit-word form): every masked KV-tile body carries 4 R2P +# per 32-column keep-word and ~0.04 ISETP per cell; the per-cell compare + select form it replaced carried 0 R2P and one ISETP +# per cell per mask term, and a build that ran out of predicate registers spilled them into GPRs through predicate-to-register moves. MASK_SASS_OPCODE_COUNTS = { **SASS_OPCODE_COUNTS, "R2P": (" R2P ",), diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm107.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm107.py index 0aef1ae90..47b06c674 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm107.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm107.py @@ -189,47 +189,34 @@ def test_sm107_ring_waits_take_the_module_spin_constant(flavor, kind, load_kw): assert code.count("spin=") == n_ring, f"{mod.__name__}: a spin= outside a .wait( call" -# The per-cell softmax mask has two lowerings behind one signature (`tile_dsl.mask.apply_mask_chunk_form`): -# "cells" = compare + select per cell per mask term (3-7 instructions per cell, 51-72 % of a masked softmax -# tile's instructions serialized ahead of the exp burst) and "bits" = one keep-word per 32 columns from two -# saturating shifts, then a register-to-predicate `R2P` + one `FSEL` per cell (1.4-1.6 per cell, independent -# of the number of active terms). Same masked set, same sentinel -> O / LSE bitwise identical; the forms differ -# ONLY in instruction count (sm_107a listings, 2026-09-22: masked body -208 (1 term) / -466 (2 terms) / -903 (the -# mxfp8 d512 SWA build, whose 128 live i1 values had spilled into GPR bits through predicate-to-register moves and LOP3) instructions per KV tile per lane). -# Every sm107 kernel picks the form with ONE module constant, `MASK_FORM`, the way it picks `DESC_VERSION`. -_MASK_FORM_EXPECTED = "bits" +# The per-cell softmax mask is ONE tile_dsl op, `tile_dsl.mask.apply_mask_chunk`: a keep-word per 32 columns from two +# saturating shifts, then a register-to-predicate `R2P` + one `FSEL` per cell (1.4-1.6 instructions per cell, independent +# of the number of active terms). It replaced a per-cell compare + select (3-7 instructions per cell, 51-72 % of a masked +# softmax tile's instructions serialized ahead of the exp burst; sm_107a listings, 2026-09-22: masked body -208 (1 term) +# / -466 (2 terms) / -903 (the mxfp8 d512 SWA build, whose 128 live i1 values had spilled into GPR bits through +# predicate-to-register moves and LOP3) instructions per KV tile per lane) -- first behind a per-kernel `MASK_FORM` +# constant (#1192 / #1197), then collapsed into the op itself. Same masked set, same sentinel -> O / LSE bitwise +# identical. What is left to pin: every masked site calls the op DIRECTLY, and no per-kernel form vocabulary comes back. @pytest.mark.parametrize("kind,load_kw", _DTYPE_FAMILIES, ids=[k for k, _ in _DTYPE_FAMILIES]) @pytest.mark.parametrize("flavor", _FLAVORS) -def test_sm107_mask_form_is_the_bits_form(flavor, kind, load_kw): - """Every sm107 prefill kernel masks in the register-to-predicate form. Asserted on the module's own - constant (what the call sites read at trace time), not on a substring a comment could supply. Flipping a - kernel back to "cells" is a legitimate A/B -- do it on a branch and re-measure, do not delete the check.""" - from cudnn.frost.tile_dsl.mask import MASK_FORMS - - mod = _load(flavor, rubin=True, **load_kw) - assert mod.MASK_FORM in MASK_FORMS, f"{mod.__name__}: MASK_FORM={mod.MASK_FORM!r} is not one of {MASK_FORMS}" - assert mod.MASK_FORM == _MASK_FORM_EXPECTED, f"{mod.__name__}: MASK_FORM={mod.MASK_FORM!r}, expected {_MASK_FORM_EXPECTED!r}" - - -@pytest.mark.parametrize("kind,load_kw", _DTYPE_FAMILIES, ids=[k for k, _ in _DTYPE_FAMILIES]) -@pytest.mark.parametrize("flavor", _FLAVORS) -def test_sm107_every_mask_site_takes_the_module_mask_form(flavor, kind, load_kw): - """MASK_FORM is only meaningful if EVERY masked call site passes it. A site that calls `apply_mask_chunk` - directly, or re-literals `form="cells"`, silently keeps the per-cell lowering on that one arm (the d256 and - mxfp8 kernels have 2-4 masked arms each), so count the sites against the constant.""" +def test_sm107_every_mask_site_calls_apply_mask_chunk(flavor, kind, load_kw): + """Every masked call site of every sm107 prefill kernel is a direct `apply_mask_chunk(` call -- no dispatcher, no + `form=` kwarg, no module `MASK_FORM` constant (the vocabulary the collapse removed). A reintroduced per-kernel + selector would let one arm (the d256 and mxfp8 kernels have 2-4 masked arms each) drift to a slower lowering with + bitwise-identical output, which no numerics test sees; the lowering itself is held by + test_sm107_masked_softmax_sass_is_register_to_predicate.""" import re mod = _load(flavor, rubin=True, **load_kw) + assert not hasattr(mod, "MASK_FORM"), f"{mod.__name__}: a MASK_FORM constant is back" with open(mod.__file__, encoding="utf-8") as fh: code = _code_lines(fh.read()) - n_sites = len(re.findall(r"\bapply_mask_chunk_form\(", code)) + n_sites = len(re.findall(r"\bapply_mask_chunk\(", code)) assert n_sites > 0, f"{mod.__name__}: no masked call site found" - n_wired = code.count("form=MASK_FORM,") - assert n_wired == n_sites, f"{mod.__name__}: {n_sites} apply_mask_chunk_form site(s) but {n_wired} pass form=MASK_FORM" - assert not re.search(r"\bapply_mask_chunk(_bits)?\(", code), f"{mod.__name__}: a direct apply_mask_chunk / apply_mask_chunk_bits call bypasses MASK_FORM" - assert not re.search(r"""form=["']""", code), f"{mod.__name__}: a re-literalled form= bypasses MASK_FORM" + for spelling in (r"\bapply_mask_chunk_form\b", r"\bapply_mask_chunk_bits\b", r"\bMASK_FORM", r"(? 165). Rows: (quant, d, dtype_o, mask spec, ISETP ceiling, -# predicate-to-register-move ceiling, spill ceiling); the -# ceilings are the measured "bits" counts (2026-09-22, sm_107a, production geometry) plus slack -- one masked arm falling back +# The masked softmax arm's SASS pin (sm_107a listings of both forms, 2026-09-22). In the bit-word form (`tile_dsl.mask.apply_mask_chunk`) +# every masked KV-tile body carries 4 R2P per 32-column word and ~0.04 ISETP per cell; the per-cell compare + select form it replaced +# carried 0 R2P and 1 ISETP per cell per mask term (611-764 ISETP whole-kernel on these two builds vs 107-108 now), and the mxfp8 d512 +# causal+SWA build ran out of predicate registers (152 predicate-to-register moves, REG 254 -> 165). Rows: (quant, d, dtype_o, mask +# spec, ISETP ceiling, predicate-to-register-move ceiling, spill ceiling); the +# ceilings are the measured bit-word counts (2026-09-22, sm_107a, production geometry) plus slack -- one masked arm falling back # to per-cell compares adds >= 128 ISETP, so a slack of 32 still catches a single arm. _SM107_MASK_SASS_ROWS = [ pytest.param("mxfp8", 512, _BF16_OUT, "causal_swa640", 108 + 32, 2 + 6, 0, id="mxfp8-d512-causal_swa640"), @@ -2768,7 +2755,7 @@ def _sm107a_sass_counts(dump, quant, d, dtype_o, mask, cands): @pytest.mark.parametrize("quant, d, dtype_o, mask, isetp_max, pred_spill_max, spill_max", _SM107_MASK_SASS_ROWS) def test_sm107_masked_softmax_sass_is_register_to_predicate(tmp_path, quant, d, dtype_o, mask, isetp_max, pred_spill_max, spill_max): - """The masked softmax arm masks through R2P + FSEL (the "bits" form), not one ISETP + FSEL per cell per term: R2P > 0, + """The masked softmax arm masks through R2P + FSEL (the bit-word form of apply_mask_chunk), not one ISETP + FSEL per cell per term: R2P > 0, whole-kernel ISETP within the measured ceiling, no predicate-register spill storm (predicate-to-register moves) and no new stack spills (the d128 causal builds carry 5 STL / 9 LDL per TILE on develop already -- that is the row's pre-existing count).""" if not _sm107a_known_to_the_dsl(): @@ -2780,7 +2767,7 @@ def test_sm107_masked_softmax_sass_is_register_to_predicate(tmp_path, quant, d, dump.mkdir() stats = _sm107a_sass_counts(dump, quant, d, dtype_o, mask, cands) print(f"\nsm107 {quant} d={d} {mask} sm_107a SASS: {stats}") - assert stats["R2P"] > 0, "no R2P: the masked arm is back to per-cell compare + select (MASK_FORM or apply_mask_chunk_bits regressed)" + assert stats["R2P"] > 0, "no R2P: the masked arm is back to per-cell compare + select (tile_dsl.mask.apply_mask_chunk regressed)" assert stats["ISETP"] <= isetp_max, f"{stats['ISETP']} ISETP > {isetp_max}: a masked arm is comparing per cell again" assert ( stats["PRED2GPR"] <= pred_spill_max diff --git a/test/python/sdpa/frost/test_sm100_mask_form.py b/test/python/sdpa/frost/test_sm100_mask_form.py index 72a50d7e4..5bcf7f883 100644 --- a/test/python/sdpa/frost/test_sm100_mask_form.py +++ b/test/python/sdpa/frost/test_sm100_mask_form.py @@ -1,20 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The sm100 masked softmax arms mask through ONE module constant, ``MASK_FORM`` (``tile_dsl/mask.py``). +"""The sm100 masked softmax arms mask through ONE tile_dsl op, ``tile_dsl.mask.apply_mask_chunk``. Every masked call site of the twelve sm100 prefill kernels, the d128 decode kernel and the d512 backward stage-2 -kernel goes through ``apply_mask_chunk_form(..., form=MASK_FORM)``: the per-cell "cells" form (one ``IADD`` + -``ISETP`` + ``FSEL`` per cell per mask term) or the "bits" form (one 32-column keep-word from two saturating shifts, -then a register-to-predicate ``R2P`` + one ``FSEL`` per cell, independent of the number of active terms). Same masked -set, same sentinel, same ``arith.select`` -> O / LSE (and, on the backward, S and every gradient downstream) are -bitwise identical; the forms differ ONLY in instruction count. ``MASK_FORM`` is spelled once per kernel next to the -mask import, the ``DESC_VERSION`` discipline, so an A/B of the two lowerings is a constant flip and no arm can drift. - -Host-only where possible: the constant and the site scan load each module through the adapter's own loaders (no GPU, -no compile); the SASS pins trace-compile two masked builds for sm_100a (``CUTE_DSL_ARCH`` needs no matching device) -and skip when no nvdisasm on ``$CUDA_PATH/bin`` or ``$PATH`` decodes the cubin. The Rubin twins live in -``test_sdpa_fwd_dsl_sm107.py``. +kernel calls ``apply_mask_chunk(...)`` directly: one 32-column keep-word from two saturating shifts, then a +register-to-predicate ``R2P`` + one ``FSEL`` per cell, independent of the number of active terms. It replaced a +per-cell compare + select (one ``IADD`` + ``ISETP`` + ``FSEL`` per cell per mask term) -- first behind a per-kernel +``MASK_FORM`` constant (#1197 on this arch line), then collapsed into the op itself once both arch lines shipped it. +Same masked set, same sentinel, same ``arith.select`` -> O / LSE (and, on the backward, S and every gradient +downstream) are bitwise identical; the forms differed ONLY in instruction count. What is left to pin: every masked +site calls the op DIRECTLY (no dispatcher, no ``form=``, no module constant a kernel could flip back), and the +lowering is the register-to-predicate one. + +Host-only where possible: the site scan loads each module through the adapter's own loaders (no GPU, no compile); the +SASS pins trace-compile two masked builds for sm_100a (``CUTE_DSL_ARCH`` needs no matching device) and skip when no +nvdisasm on ``$CUDA_PATH/bin`` or ``$PATH`` decodes the cubin. The Rubin twins live in ``test_sdpa_fwd_dsl_sm107.py``. """ import re @@ -32,7 +33,7 @@ _E4M3, _BF16 = 0, 2 # Every sm100 kernel with a masked softmax arm: (quantization kind, flavor) -> (the kernel file, the number of -# apply_mask_chunk_form sites it carries). "decode" is the d128 f16 flavor at TILE_CGA_M=1 (the decode tile, +# apply_mask_chunk sites it carries). "decode" is the d128 f16 flavor at TILE_CGA_M=1 (the decode tile, # `sm100/decode_d128_f16.py`); "bwd" is the d512 backward's stage-2 kernel. The site count is the migration table of # the port: a count that moves means an arm was added or dropped and the table (and this pin) must say so. _KERNELS = { @@ -53,12 +54,10 @@ } _KERNEL_IDS = [f"{k}-d{f[0]}" if f[0] == f[1] else f"{k}-d{f[0]}x{f[1]}" for k, f in _KERNELS] -# Both forms mask the same set with the same sentinel; "bits" is the form every sm100 kernel ships (sm_100a listings of both -# forms, 2026-09-22, cutlass-dsl 4.8.0: whole-kernel ISETP 598 -> 94 on the fp8 d128 causal + padded build and 590 -> 86 on the -# f16 d512 causal + SWA 640 build, R2P 0 -> 32 on both, the masked KV-tile body 1191 -> 723 / 1172 -> 719 instructions per lane -# while the dense bodies keep their 552 / 530). Flipping a kernel back to "cells" is a legitimate A/B -- do it on a branch -# and re-measure, do not delete the check. -_MASK_FORM_EXPECTED = "bits" +# The switch to the bit-word form, sm_100a listings of both forms (2026-09-22, cutlass-dsl 4.8.0): whole-kernel ISETP 598 -> 94 +# on the fp8 d128 causal + padded build and 590 -> 86 on the f16 d512 causal + SWA 640 build, R2P 0 -> 32 on both, the masked +# KV-tile body 1191 -> 723 / 1172 -> 719 instructions per lane while the dense bodies keep their 552 / 530. An A/B against +# the per-cell form is a branch that re-adds it to tile_dsl; the shipped tree has one op. def _load(kind, flavor): @@ -86,45 +85,31 @@ def _code_lines(src): @pytest.mark.parametrize("kind,flavor", list(_KERNELS), ids=_KERNEL_IDS) -def test_sm100_mask_form_is_the_bits_form(kind, flavor): - """Every sm100 kernel with a masked arm masks in the register-to-predicate form. Asserted on the module's own - constant (what the call sites read at trace time), not on a substring a comment could supply, and on the module the - adapter really loads for that kernel (the file name is checked, so a loader routing to a sibling cannot pass).""" - from cudnn.frost.tile_dsl.mask import MASK_FORMS - - mod = _load(kind, flavor) - file, _ = _KERNELS[(kind, flavor)] - assert mod.__file__.endswith(file), f"{mod.__name__}: loaded {mod.__file__}, expected {file}" - assert mod.MASK_FORM in MASK_FORMS, f"{mod.__name__}: MASK_FORM={mod.MASK_FORM!r} is not one of {MASK_FORMS}" - assert mod.MASK_FORM == _MASK_FORM_EXPECTED, f"{mod.__name__}: MASK_FORM={mod.MASK_FORM!r}, expected {_MASK_FORM_EXPECTED!r}" - - -@pytest.mark.parametrize("kind,flavor", list(_KERNELS), ids=_KERNEL_IDS) -def test_sm100_every_mask_site_takes_the_module_mask_form(kind, flavor): - """MASK_FORM is only meaningful if EVERY masked call site passes it. A site that calls `apply_mask_chunk` directly, - or re-literals `form="cells"`, silently keeps the per-cell lowering on that one arm (the d256 and mxfp8 kernels have - 2-4 masked arms each), so count the sites against the constant, and against the migration table above.""" +def test_sm100_every_mask_site_calls_apply_mask_chunk(kind, flavor): + """Every masked call site is a direct `apply_mask_chunk(` call, counted against the migration table above (a count + that moves means an arm was added or dropped and the table must say so), on the module the adapter really loads for + that kernel (the file name is checked, so a loader routing to a sibling cannot pass). No dispatcher, no `form=` + kwarg and no module `MASK_FORM` constant: a reintroduced per-kernel selector would let one arm (the d256 and mxfp8 + kernels have 2-4 masked arms each) drift to a slower lowering with bitwise-identical output, which no numerics test + sees; the lowering itself is held by test_sm100_masked_softmax_sass_is_register_to_predicate.""" mod = _load(kind, flavor) file, n_expected = _KERNELS[(kind, flavor)] assert mod.__file__.endswith(file), f"{mod.__name__}: loaded {mod.__file__}, expected {file}" + assert not hasattr(mod, "MASK_FORM"), f"{mod.__name__}: a MASK_FORM constant is back" with open(mod.__file__, encoding="utf-8") as fh: code = _code_lines(fh.read()) - n_sites = len(re.findall(r"\bapply_mask_chunk_form\(", code)) - assert n_sites == n_expected, f"{mod.__name__}: {n_sites} apply_mask_chunk_form site(s), the migration table says {n_expected} -- update both" - n_wired = code.count("form=MASK_FORM,") - assert n_wired == n_sites, f"{mod.__name__}: {n_sites} apply_mask_chunk_form site(s) but {n_wired} pass form=MASK_FORM" - assert not re.search(r"\bapply_mask_chunk(_bits)?\(", code), f"{mod.__name__}: a direct apply_mask_chunk / apply_mask_chunk_bits call bypasses MASK_FORM" - assert not re.search(r"""form=["']""", code), f"{mod.__name__}: a re-literalled form= bypasses MASK_FORM" - defs = re.findall(r"^MASK_FORM\b.*$", code, re.M) - assert defs == ["MASK_FORM: str = MASK_FORM_BITS"], f"{mod.__name__}: MASK_FORM must be defined exactly once from the tile_dsl constant, found {defs}" + n_sites = len(re.findall(r"\bapply_mask_chunk\(", code)) + assert n_sites == n_expected, f"{mod.__name__}: {n_sites} apply_mask_chunk site(s), the migration table says {n_expected} -- update both" + for spelling in (r"\bapply_mask_chunk_form\b", r"\bapply_mask_chunk_bits\b", r"\bMASK_FORM", r"(?= 128 ISETP, so _ISETP_SLACK = 32 still catches a single arm; STL / LDL are bounds @@ -198,7 +182,7 @@ def test_sm100_d192_dense_fp8_window_sentinel_is_inside_the_bits_domain(rubin): @pytest.mark.parametrize("kind, d, fields, compile_kw, pins", _SM100_MASK_SASS_ROWS) def test_sm100_masked_softmax_sass_is_register_to_predicate(tmp_path, kind, d, fields, compile_kw, pins): - """The masked softmax arm masks through R2P + FSEL (the "bits" form), not one ISETP + FSEL per cell per term: R2P > 0, + """The masked softmax arm masks through R2P + FSEL (the bit-word form of apply_mask_chunk), not one ISETP + FSEL per cell per term: R2P > 0, whole-kernel ISETP within the measured ceiling, no predicate-register spill storm (predicate-to-register moves) and no new stack spills (STL / LDL within SPILL_TOLERANCE of this toolchain's count). Compiled for sm_100a at the production geometry; skips when no nvdisasm decodes the cubin.""" @@ -206,9 +190,8 @@ def test_sm100_masked_softmax_sass_is_register_to_predicate(tmp_path, kind, d, f tmp_path, probe_src=_SM100_MASK_SASS_PROBE, arch="sm_100a", params={"kind": kind, "d": d, "compile": compile_kw, **fields}, tag=f"mask_{kind}_d{d}" ) assert probe.expect["MASK_FLAGS"] != 0, "the probe params select the DENSE specialization -- no masked arm is compiled in" - assert probe.expect["MASK_FORM_IS_BITS"] == 1, "the module's MASK_FORM is not the bits form" stats = probe.stats - assert stats["R2P"] > 0, "no R2P: the masked arm is back to per-cell compare + select (MASK_FORM or apply_mask_chunk_bits regressed)" + assert stats["R2P"] > 0, "no R2P: the masked arm is back to per-cell compare + select (tile_dsl.mask.apply_mask_chunk regressed)" assert ( stats["ISETP"] <= pins["ISETP"] + _ISETP_SLACK ), f"{stats['ISETP']} ISETP > {pins['ISETP']} + {_ISETP_SLACK}: a masked arm is comparing per cell again" diff --git a/test/python/sdpa/frost/test_tile_dsl_mask_bits.py b/test/python/sdpa/frost/test_tile_dsl_mask_bits.py index 4b470e058..7630354f8 100644 --- a/test/python/sdpa/frost/test_tile_dsl_mask_bits.py +++ b/test/python/sdpa/frost/test_tile_dsl_mask_bits.py @@ -1,25 +1,26 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The "bits" per-cell mask form (`tile_dsl.mask.apply_mask_chunk_bits`) masks EXACTLY the cells -`apply_mask_chunk` masks -- pinned on a host emulation of both formulas. - -`apply_mask_chunk` compares every cell against every active bound (padded: `kv >= seq_kv_len`; -causal: `kv > q_caus_lim`; SWA: `kv < q_minus_w`) and ORs the terms. `apply_mask_chunk_bits` -maps the same terms onto one band `[lo, hi)` per lane, builds a 32-column KEEP word per word -with two saturating shifts (PTX `shr.u32` / `shl.b32` clamp a shift amount >= 32 to a zero -result) and selects on one bit per cell. The two must agree on every cell for every flag -combination, every sentinel and every payload -- including the cases a random sweep rarely -hits: a band edge exactly on a word boundary, an edge left / right of the whole chunk (the -saturated shift), a fully-masked and a fully-unmasked row, NaN / +-inf / +-0 / the sentinel -itself as payload. Bitwise (`view(int32)`) so a NaN payload is compared as a bit pattern. +"""The bit-word per-cell mask op (`tile_dsl.mask.apply_mask_chunk`) masks EXACTLY the cells the +per-cell compare + select formula it replaced masks -- pinned on a host emulation of both formulas. + +The per-cell formula (what every kernel shipped before #1192 / #1197) compares every cell against +every active bound (padded: `kv >= seq_kv_len`; causal: `kv > q_caus_lim`; SWA: `kv < q_minus_w`) +and ORs the terms. `apply_mask_chunk` maps the same terms onto one band `[lo, hi)` per lane, +builds a 32-column KEEP word per word with two saturating shifts (PTX `shr.u32` / `shl.b32` clamp +a shift amount >= 32 to a zero result) and selects on one bit per cell. The two must agree on +every cell for every flag combination, every sentinel and every payload -- including the cases a +random sweep rarely hits: a band edge exactly on a word boundary, an edge left / right of the +whole chunk (the saturated shift), a fully-masked and a fully-unmasked row, NaN / +-inf / +-0 / +the sentinel itself as payload. Bitwise (`view(int32)`) so a NaN payload is compared as a bit +pattern. Device-independent: the emulation below mirrors `tile_dsl/mask.py` line by line in torch int64 arithmetic and wraps the intermediates the device computes in Int32 (`lo - kv_col_base`, the per-word shift count), so the reason for `MASK_BOUND_LIMIT` is visible here too. The GPU half -- -the real kernels, `MASK_FORM` cells vs bits -- was pinned by dumping O / LSE from develop and -from the branch on a Rubin GPU (82 kernel x mask x shape cases bitwise, 2026-09-22); the -lowering itself is pinned by `test_sm107_masked_softmax_sass_is_register_to_predicate`. +the real kernels, per-cell form vs bit-word form -- was pinned by dumping O / LSE from develop and +from the branch on a Rubin GPU (82 kernel x mask x shape cases bitwise, 2026-09-22); the lowering +itself is pinned by `test_sm107_masked_softmax_sass_is_register_to_predicate` and its sm100 twin. """ import itertools @@ -28,7 +29,8 @@ import torch from cudnn.frost.tile_dsl.constants import MASK_CAUSAL, MASK_PADDED, MASK_SWA -from cudnn.frost.tile_dsl.mask import _NEG_INF_BITS, MASK_BOUND_LIMIT, MASK_FORM_BITS, MASK_FORM_CELLS, MASK_FORMS, MASK_WORD_COLS, apply_mask_chunk_bits +from cudnn.frost.tile_dsl import mask as mask_mod +from cudnn.frost.tile_dsl.mask import _NEG_INF_BITS, MASK_BOUND_LIMIT, MASK_WORD_COLS, apply_mask_chunk pytestmark = [pytest.mark.L0] @@ -53,7 +55,7 @@ def _fill(S, value): def _band(q_abs, seq_kv_len, window_left, mask_flags, bottom_right, causal_diag, window_right): - """`apply_mask_chunk_bits`'s flag -> band mapping, on int64 tensors. Returns (lo, hi), either None.""" + """`apply_mask_chunk`'s flag -> band mapping, on int64 tensors. Returns (lo, hi), either None.""" lo = hi = None diag = causal_diag if bottom_right else torch.zeros_like(q_abs) if mask_flags & MASK_SWA: @@ -91,7 +93,7 @@ def _keep_words(lo, hi, kv_col_base, n_cols): def mask_bits(S, q_abs, kv_col_base, seq_kv_len, window_left, mask_flags, bottom_right, causal_diag, mask_value, window_right): - """`apply_mask_chunk_bits`, emulated: S [rows, n_cols] fp32, bounds [rows] int64.""" + """`apply_mask_chunk` (the bit-word op), emulated: S [rows, n_cols] fp32, bounds [rows] int64.""" n_cols = S.shape[1] lo, hi = _band(q_abs, seq_kv_len, window_left, mask_flags, bottom_right, causal_diag, window_right) words = _keep_words(lo, hi, kv_col_base, n_cols) @@ -106,7 +108,8 @@ def mask_bits(S, q_abs, kv_col_base, seq_kv_len, window_left, mask_flags, bottom def mask_cells(S, q_abs, kv_col_base, seq_kv_len, window_left, mask_flags, bottom_right, causal_diag, mask_value, window_right): - """`apply_mask_chunk`, emulated: the per-cell reference.""" + """The per-cell compare + select formula, emulated: the reference `apply_mask_chunk` is pinned against + (the form every kernel shipped before the bit-word op; it has no Int32 band arithmetic, hence no domain limit).""" n_cols = S.shape[1] kv = kv_col_base[:, None] + torch.arange(n_cols, dtype=torch.int64)[None, :] diag = causal_diag if bottom_right else torch.zeros_like(q_abs) @@ -212,9 +215,9 @@ def test_bits_form_masks_the_same_cells_at_word_edges(mask_flags, n_cols): def test_bits_form_domain_guard(): """The band arithmetic is Int32. A window bound at or past `MASK_BOUND_LIMIT` can wrap `lo - kv_col_base`, and the - bits form then masks EVERYTHING where the per-cell form masks nothing. Three pins: the two forms agree at the largest + bit-word op then masks EVERYTHING where the per-cell formula masks nothing. Three pins: the two agree at the largest in-domain window with every index below 2**28 (the documented domain); the Int32-faithful emulation shows the - divergence just past it (why the guard exists, not a claim about the device); and `apply_mask_chunk_bits` refuses such + divergence just past it (why the guard exists, not a claim about the device); and `apply_mask_chunk` refuses such a window at trace time, before it touches a register (Python ints, so the check costs no instruction).""" gen = torch.Generator().manual_seed(0xD0A1) rows, n_cols = 256, 128 @@ -234,16 +237,20 @@ def test_bits_form_domain_guard(): zero = torch.zeros(rows, dtype=torch.int64) base = torch.full((rows,), n_cols, dtype=torch.int64) args = (zero, base, zero, (1 << 31) - 1, MASK_SWA, 0, zero, _NEG_INF_BITS, 0) - _assert_same(mask_cells(S, *args), S, "per-cell form with the widest window leaves the chunk untouched") - _assert_same(mask_bits(S, *args), _fill(S, _NEG_INF_BITS), "bits form with the widest window masks the whole chunk") + _assert_same(mask_cells(S, *args), S, "per-cell formula with the widest window leaves the chunk untouched") + _assert_same(mask_bits(S, *args), _fill(S, _NEG_INF_BITS), "bit-word op with the widest window masks the whole chunk") with pytest.raises(ValueError, match="window_left must be <"): - apply_mask_chunk_bits(None, None, None, None, MASK_BOUND_LIMIT, MASK_SWA) + apply_mask_chunk(None, None, None, None, MASK_BOUND_LIMIT, MASK_SWA) with pytest.raises(ValueError, match="window_right must be <"): - apply_mask_chunk_bits(None, None, None, None, 0, MASK_CAUSAL, window_right=MASK_BOUND_LIMIT) + apply_mask_chunk(None, None, None, None, 0, MASK_CAUSAL, window_right=MASK_BOUND_LIMIT) -def test_mask_form_vocabulary(): - """The kernels' `MASK_FORM` constant takes one of exactly these two spellings; a third would silently fall - through `apply_mask_chunk_form` to its ValueError at trace time.""" - assert MASK_FORMS == (MASK_FORM_CELLS, MASK_FORM_BITS) == ("cells", "bits") +def test_the_mask_op_has_one_form(): + """`apply_mask_chunk` IS the bit-word op: the per-kernel `MASK_FORM` constant, the `apply_mask_chunk_form` + dispatcher and the `apply_mask_chunk_bits` twin of #1192 / #1197 were collapsed once both arch lines shipped the + form. A reintroduced selector would let a kernel drift back to the per-cell lowering with bitwise-identical + output -- a regression no numerics test can see -- so the vocabulary stays gone; the per-cell formula survives + only as the host emulation above.""" assert MASK_WORD_COLS == 32 + for name in ("apply_mask_chunk_form", "apply_mask_chunk_bits", "MASK_FORM", "MASK_FORMS", "MASK_FORM_BITS", "MASK_FORM_CELLS"): + assert not hasattr(mask_mod, name), f"tile_dsl.mask.{name} is back -- there is one mask op, apply_mask_chunk" From 3c756142bfd9e04941ae3660847c3a7296d0f986 Mon Sep 17 00:00:00 2001 From: randers Date: Wed, 23 Sep 2026 01:48:42 +0000 Subject: [PATCH 2/3] test(sdpa/frost): retire the d128-only-era Rubin skips in the per-tensor FP8 forward suite `_skip_on_rubin` / `_skip_d256_on_rubin` / `_skip_d512_on_rubin` all carried the reason "sm107 serves d128 only". It has not been true since the Rubin d192x128 / d256 / d512 per-tensor FP8 kernels landed (api_dsl._SM107_FP8_KERNEL_FILES, engines._sm100_fp8_spec(arch="sm107") claims all four flavors, THD on all four, the same envelope floors) -- so 33 decorators plus the runtime twin in `_check_fp8_strided_stats` were turning 97 cc 10.7 cases into skips, and a skip is a hole in the decline detector (frost-gotchas: a stale `_skip_on_rubin` hid the d192 dense fp8 window-sentinel decline on develop). - The three markers and the runtime skip are gone: the d192x128 masks / wide SWA / leading-zero-KV / THD / strided-stats tests, every d256 test and the whole d512 section now run on both arch lines. The d512 section header no longer says the flavor is sm100-only (nor that no d512 MXFP8 kernel exists). - TWO Rubin declines are still true and stay as markers with their real reason, each pinned host-side in test_sdpa_fwd_dsl_sm107.py (no GPU needed): * `_skip_pack_gqa_wide_on_rubin` -- the sm107 row packs GQA on d128 only (`pack_gqa_d_shapes = {(128, 128)}`), a DECLARED decline: the packed d192 tests and the d256 param of the packed d128/d256 tests (`_D128_D256_PACK_GQA`; the THD users keep the plain `_D128_D256`). Pin: `test_sm107_fp8_pack_gqa_is_d128_only` (packed d192/d256/d512 ineligible with a reason naming the knob, unpacked eligible, packed d128 eligible). * `_skip_strided_stats_d256_on_rubin` -- the sm107 fp8 d256 kernel has NOT ported strided Stats: `compile()` raises NotImplementedError("strided Stats not ported (contiguous [B, H, S] only)") while the row declares Stats on every flavor and mismatch()'s generic `dense_layout_ok` admits the layout -- a Capabilities GAP (also the d512 fp8 kernel and all four sm107 MXFP8 kernels; d128 / d192x128 fp8 ported it). Pin: `test_sm107_fp8_strided_stats_is_not_ported_beyond_d192` (the typed raise on the d256 / d512 templates under a strided `lse_stride`, no guard on d128 / d192x128, the per-kernel guard set for fp8 + mxfp8). Follow-up, not this PR: port it (the d256 f16 sibling's `lse_strides` is the model) or declare the layout per flavor; then invert the arm and drop the marker. - The two host-only tests that pinned the sm100 row / (10, 0) shape map alone now iterate both rows / both device lines (`test_fp8_large_flavors_serve_exact_shapes_only`, `test_fp8_d512_envelope_floor_declines_straddling_shapes`). Rubin proof (cc 10.7, 204-SM part, GPU 6, 2026-09-23; develop control first: the 97 ids all SKIPPED on develop @ 8efd9510): 94 of the 97 newly-unskipped ids PASS; `test_fp8_d256_strided_stats[e4m3|e5m2]` hit the strided-Stats gap above (now the second marker); `test_fp8_d192_d128_thd[e5m2-True-False]` reads max|O-ref| 0.0409 vs the shared e5m2 atol 0.0400 -- the same e5m2 dataset-edge class as the three 0.0404-0.0409 ids already classified on develop's Rubin gates. The tolerance is NOT widened and the id stays un-skipped; it is recorded for a decision (frost_dev/mask_sass/pr3/PR3_BODY.md). Co-Authored-By: Claude Fable 5.1 --- .../sdpa/frost/test_sdpa_fwd_dsl_sm107.py | 64 +++++++++++ .../sdpa/frost/test_sdpa_fwd_fp8_sm100.py | 104 +++++++++--------- 2 files changed, 114 insertions(+), 54 deletions(-) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm107.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm107.py index 47b06c674..645cc5ca8 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm107.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm107.py @@ -422,6 +422,70 @@ def test_sm107_f16_declines_split_kv_and_pack_gqa(): assert engines.mismatch(caps, _f16_facts(), engines.SdpaFwdKnobs(pack_gqa=True)) is not None +def test_sm107_fp8_pack_gqa_is_d128_only(): + """The Rubin per-tensor FP8 row packs GQA on the d128 flavor only (`pack_gqa_d_shapes = {(128, 128)}`: the + d192x128 / d256 / d512 siblings carry no PackGQA path) while it serves those flavors UNPACKED. This is the + typed decline behind `_skip_pack_gqa_wide_on_rubin` in test_sdpa_fwd_fp8_sm100.py -- the one Rubin marker that + survived retiring the d128-only-era skips -- so it is asserted on real facts here, where no GPU is needed: a + packed d192 / d256 / d512 graph is ineligible with a reason naming the knob, the same graph unpacked is eligible, + and packed d128 is eligible. When a wider kernel gains PackGQA: widen the row, INVERT that shape's packed + assertion and drop the marker (test/AGENTS.md: invert the counter assertion, do not delete it).""" + from cudnn.sdpa.fwd import engines + + caps = _caps("sdpa_fwd_prefill_sm107_fp8") + assert caps.pack_gqa_d_shapes == frozenset({(128, 128)}) + for d_qk, d_v in ((192, 128), (256, 256), (512, 512)): + facts = _f16_facts(**_fp8_ungated_kw(h_kv=2, d_qk=d_qk, d_v=d_v)) + assert engines.mismatch(caps, facts) is None, (d_qk, d_v, "the unpacked graph must be served") + why = engines.mismatch(caps, facts, engines.SdpaFwdKnobs(pack_gqa=True)) + assert why is not None and "pack_gqa" in why, (d_qk, d_v, why) + packed_d128 = _f16_facts(**_fp8_ungated_kw(h_kv=2, d_qk=128, d_v=128)) + assert engines.mismatch(caps, packed_d128, engines.SdpaFwdKnobs(pack_gqa=True)) is None + + +# The sm107 kernels that still raise at `compile()` on a DENSE strided Stats layout (`lse_stride` given, no padded +# rows): the guard reads "strided Stats not ported (contiguous [B, H, S] only)". The fp8 row's d128 / d192x128 +# kernels ported it (Rubin run 2026-09-23: test_fp8_strided_stats + test_fp8_strided_stats_other_flavors[d192_d128_*] +# PASS on cc 10.7); the fp8 d256 / d512 kernels and every MXFP8 kernel did not. +_SM107_STRIDED_STATS_NOT_PORTED = { + "fp8": {(256, 256), (512, 512)}, + "mxfp8": {(128, 128), (192, 128), (256, 256), (512, 512)}, +} + + +def test_sm107_fp8_strided_stats_is_not_ported_beyond_d192(): + """A Capabilities GAP, pinned so it is visible: the sm107 per-tensor FP8 row declares Stats on all four flavors and + `engines.mismatch` admits any dense-compatible Stats layout (`ga.dense_layout_ok`, no per-flavor field), but the + d256 and d512 sm107 fp8 kernels raise `NotImplementedError("strided Stats not ported ...")` from `compile()` on a + strided dense layout -- so on cc 10.7 the pinned engine dies at build_plans on `test_fp8_d256_strided_stats` + (measured 2026-09-23; that test carries `_skip_strided_stats_d256_on_rubin` with this reason). The raise is at + the top of `compile()`, before any JIT work, so this is host-only. The d128 / d192x128 kernels take the layout. + Follow-up (not this PR): port strided Stats to the d256 / d512 fp8 kernels (the d256 f16 sibling's `lse_strides` + is the model) OR declare the layout per flavor on the row; then INVERT the raise assertion for that shape and + drop the marker (test/AGENTS.md: invert the counter assertion, do not delete it).""" + caps = _caps("sdpa_fwd_prefill_sm107_fp8") + assert caps.stats is True and caps.d_shapes == frozenset({(128, 128), (192, 128), (256, 256), (512, 512)}) + assert not any(f.startswith("stats_layout") or f.startswith("strided_stats") for f in caps.__dataclass_fields__), "a per-flavor field exists now: use it" + fp8_kw = dict(_DTYPE_FAMILIES[1][1]) + for flavor in _FLAVORS: + mod = _load(flavor, rubin=True, **fp8_kw) + d_qk, d_v = flavor + sq = 128 + strided = (sq * 4 * 2, sq * 2, 2) # any declared (B, H, S) stride: the guard fires on `is not None` + with open(mod.__file__, encoding="utf-8") as fh: + has_guard = "strided Stats not ported" in _code_lines(fh.read()) + if flavor in _SM107_STRIDED_STATS_NOT_PORTED["fp8"]: + assert has_guard, f"{mod.__name__}: the guard is gone -- strided Stats ported? Invert this arm and drop the marker." + with pytest.raises(NotImplementedError, match="strided Stats not ported"): + mod.compile(b=2, qh=4, kh=2, sq=sq, skv=128, d_qk=d_qk, d_v=d_v, has_lse=True, lse_stride=strided) + else: + assert not has_guard, f"{mod.__name__}: a strided-Stats guard appeared on a flavor that had ported it" + for flavor in _FLAVORS: + mod = _load(flavor, rubin=True, **_DTYPE_FAMILIES[2][1]) + with open(mod.__file__, encoding="utf-8") as fh: + assert ("strided Stats not ported" in _code_lines(fh.read())) == (flavor in _SM107_STRIDED_STATS_NOT_PORTED["mxfp8"]), mod.__name__ + + def test_sm107_rows_carry_the_padded_stats_trim(): """Every Rubin template carries the per-batch seq_len_q trim (padded q rows write LSE=-inf, O=0), so the rows serve dense padded Stats and the diff --git a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py index b67743044..908a3ea1b 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py @@ -36,11 +36,34 @@ # ONE per-tensor FP8 engine per arch line (sm100 = pre-Rubin Blackwell # 100-106, sm107 = the Rubin line): pin the engine that serves the device -# under test. The d192xd128 kernel flavor exists on the sm100 engine only. +# under test. Both lines carry all four native flavors -- d128, d192xd128, +# d256 and d512 (api_dsl._SM107_FP8_KERNEL_FILES mirrors _SM100_FP8_KERNEL_FILES) +# -- so every flavor test below runs on both; the `_skip_*_on_rubin` markers +# from the d128-only era are retired. TWO Rubin declines are still true, each +# kept as a marker with its real reason and pinned host-side in +# test_sdpa_fwd_dsl_sm107.py: +# * the sm107 row packs GQA on d128 only (`pack_gqa_d_shapes = {(128, 128)}`, +# no PackGQA path in the wider siblings) -- a DECLARED decline +# (test_sm107_fp8_pack_gqa_is_d128_only); the packed d192 / d256 execute +# tests skip with that reason; +# * the sm107 fp8 d256 (and d512) kernels have NOT ported strided Stats +# (`compile()` raises "strided Stats not ported (contiguous [B, H, S] only)") +# while the row declares Stats on every flavor and the generic stride check +# admits any dense-compatible layout -- a Capabilities GAP, follow-up: port it +# or declare it per flavor (test_sm107_fp8_strided_stats_is_not_ported_beyond_d192). +# Measured on cc 10.7 (2026-09-23): the pinned engine raises the typed +# NotImplementedError at build_plans, so select_engine's strict pin FAILS. _D128_ARCH = "sm107" if _SM == 107 else "sm100" -_skip_on_rubin = pytest.mark.skipif(_SM == 107, reason="the d192xd128 per-tensor FP8 flavor has no Rubin kernel (sm107 serves d128 only)") -_skip_d256_on_rubin = pytest.mark.skipif(_SM == 107, reason="the d256xd256 per-tensor FP8 flavor has no Rubin kernel (sm107 serves d128 only)") -_D128_D256 = [pytest.param(128, id="d128"), pytest.param(256, id="d256", marks=_skip_d256_on_rubin)] +_skip_pack_gqa_wide_on_rubin = pytest.mark.skipif( + _SM == 107, reason="the Rubin per-tensor FP8 row packs GQA on d128 only (pack_gqa_d_shapes; pinned by test_sm107_fp8_pack_gqa_is_d128_only)" +) +_skip_strided_stats_d256_on_rubin = pytest.mark.skipif( + _SM == 107, + reason="sm107 fp8 d256 kernel: strided Stats not ported (contiguous [B, H, S] only); the row declares it -- Capabilities gap, follow-up " + "(pinned by test_sm107_fp8_strided_stats_is_not_ported_beyond_d192)", +) +_D128_D256 = [pytest.param(128, id="d128"), pytest.param(256, id="d256")] +_D128_D256_PACK_GQA = [pytest.param(128, id="d128"), pytest.param(256, id="d256", marks=_skip_pack_gqa_wide_on_rubin)] _FP8 = {"e4m3": torch.float8_e4m3fn, "e5m2": torch.float8_e5m2} _FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0} @@ -439,8 +462,6 @@ def _check(out, o_ref, out_dt, in_key, amax_o, amax_o_ref, atol=None): def _check_fp8_strided_stats(d_qk, d_v, in_key): - if torch.cuda.get_device_capability() == (10, 7) and d_qk != 128: - pytest.skip("SM107 per-tensor FP8 supports only d128") kwargs = dict( B=2, H_q=4, @@ -565,7 +586,6 @@ def test_fp8_d192_d128_output_dtypes(in_key, out_key): _check(out, o_ref, _OUT[out_key], in_key, a_o, a_o_ref) -@_skip_on_rubin @pytest.mark.L0 @pytest.mark.parametrize("mask", ["none", "causal_br", "swa"]) @torch_fork_set_rng(seed=0) @@ -588,7 +608,6 @@ def test_fp8_d192_d128_masks(mask): _check(out, o_ref, torch.float16, "e4m3", a_o, a_o_ref) -@_skip_on_rubin @pytest.mark.L0 @torch_fork_set_rng(seed=0) def test_fp8_d192_d128_wide_swa_boundary_dense(): @@ -610,7 +629,6 @@ def test_fp8_d192_d128_wide_swa_boundary_dense(): _check(out, o_ref, torch.float16, "e4m3", a_o, a_o_ref) -@_skip_on_rubin @pytest.mark.L0 @pytest.mark.parametrize( ("in_key", "out_key", "with_sink"), @@ -676,7 +694,6 @@ def test_fp8_head_dim_envelope(dims, mask): @pytest.mark.L0 -@_skip_on_rubin def test_fp8_large_flavors_serve_exact_shapes_only(): """The d192xd128 and D256 per-tensor FP8 flavors serve ONLY their exact shapes. With d_qk zero-padded into d192 (144/160/176) the output is wrong (test_mhas_v2 @@ -685,13 +702,15 @@ def test_fp8_large_flavors_serve_exact_shapes_only(): alone and fails in the battery). The engine row floors both (128 / 255), the adapter mirrors the floors, and both flavor selections honour them, so the inexact (128, 256) region is declined everywhere -- the classic backend verdict - applies -- while the d128 flavor's envelope and the d512 band are unchanged.""" + applies -- while the d128 flavor's envelope and the d512 band are unchanged. + The floor table is arch-independent (contract rule 8b'), so both rows are pinned.""" from cudnn.sdpa.fwd import engines as eng from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100, _SM100_FP8_ENVELOPE_FLOORS, _fp8_envelope_covers assert _SM100_FP8_ENVELOPE_FLOORS[(192, 128)] == 128 and _SM100_FP8_ENVELOPE_FLOORS[(256, 256)] == 255 - spec = next(sp for sp in eng.ENGINE_SPECS if sp.name == "sdpa_fwd_prefill_sm100_fp8") - assert dict(spec.capabilities.d_envelope_floors) == {(192, 128): 128, (256, 256): 255, (512, 512): 256} + for arch in ("sm100", "sm107"): + spec = next(sp for sp in eng.ENGINE_SPECS if sp.name == f"sdpa_fwd_prefill_{arch}_fp8") + assert dict(spec.capabilities.d_envelope_floors) == {(192, 128): 128, (256, 256): 255, (512, 512): 256}, arch shapes = frozenset({(128, 128), (192, 128), (256, 256), (512, 512)}) for d_qk, d_v in ((160, 96), (176, 128), (192, 96), (160, 160), (224, 208)): assert not _fp8_envelope_covers(d_qk, d_v, shapes), (d_qk, d_v) @@ -723,7 +742,6 @@ def _api(d_qk, d_v): @pytest.mark.L0 -@_skip_d256_on_rubin @pytest.mark.parametrize("out_key", ["fp16", "bf16", "e4m3", "e5m2"]) @pytest.mark.parametrize("in_key", _INS) @torch_fork_set_rng(seed=0) @@ -746,7 +764,6 @@ def test_fp8_d256_output_dtypes(in_key, out_key): @pytest.mark.L0 -@_skip_d256_on_rubin @pytest.mark.parametrize("in_key", _INS) @pytest.mark.parametrize("mask", ["none", "causal_br", "swa"]) @torch_fork_set_rng(seed=0) @@ -769,16 +786,15 @@ def test_fp8_d256_masks(in_key, mask): @pytest.mark.L1 -@_skip_d256_on_rubin +@_skip_strided_stats_d256_on_rubin @pytest.mark.parametrize("in_key", _INS) @torch_fork_set_rng(seed=59) def test_fp8_d256_strided_stats(in_key): - """D256 preserves the caller's non-contiguous Stats layout.""" + """D256 preserves the caller's non-contiguous Stats layout (sm100; the sm107 d256 kernel has not ported it -- see the marker).""" _check_fp8_strided_stats(256, 256, in_key) @pytest.mark.L1 -@_skip_d256_on_rubin @pytest.mark.parametrize("in_key", _INS) @torch_fork_set_rng(seed=0) def test_fp8_d256_bottom_right_rectangular(in_key): @@ -801,7 +817,6 @@ def test_fp8_d256_bottom_right_rectangular(in_key): @pytest.mark.L1 -@_skip_d256_on_rubin @pytest.mark.parametrize("case", ["right", "bottom_right", "right_swa"]) @torch_fork_set_rng(seed=0) def test_fp8_d256_right_band_combinations(case): @@ -831,7 +846,6 @@ def test_fp8_d256_right_band_combinations(case): @pytest.mark.L1 -@_skip_d256_on_rubin @pytest.mark.parametrize("in_key", _INS) @pytest.mark.parametrize("causal", [False, True]) @torch_fork_set_rng(seed=0) @@ -992,7 +1006,6 @@ def test_fp8_dense_q_trim_stats_sink(d, d_v): @pytest.mark.L1 -@_skip_d256_on_rubin @pytest.mark.parametrize("in_key", _INS) @torch_fork_set_rng(seed=0) def test_fp8_d256_gqa_sink(in_key): @@ -1018,7 +1031,6 @@ def test_fp8_d256_gqa_sink(in_key): @pytest.mark.L0 -@_skip_d256_on_rubin @pytest.mark.parametrize( ("in_key", "out_key", "with_sink"), [ @@ -1081,7 +1093,7 @@ def test_fp8_gqa(in_key): [(8, 4), (8, 2), (8, 1), (16, 1)], ids=["g2", "g4", "g8_mqa", "g16_mqa"], ) -@pytest.mark.parametrize("d", _D128_D256) +@pytest.mark.parametrize("d", _D128_D256_PACK_GQA) @torch_fork_set_rng(seed=0) def test_fp8_pack_gqa_ratios(d, h_q, h_kv): """Packed plans across GQA ratios, causal, tile-unaligned s_q, LSE checked.""" @@ -1124,7 +1136,7 @@ def test_fp8_pack_gqa_tiles(s_q): "mask", ["none_padded", "causal", "causal_br", "swa", "sink_causal"], ) -@pytest.mark.parametrize("d", _D128_D256) +@pytest.mark.parametrize("d", _D128_D256_PACK_GQA) @torch_fork_set_rng(seed=0) def test_fp8_pack_gqa_features(d, mask, out_dt): """Packed plans x the fp8 mask/sink envelope x both output dtypes.""" @@ -1165,7 +1177,7 @@ def test_fp8_pack_gqa_features(d, mask, out_dt): @pytest.mark.L1 -@pytest.mark.parametrize("d", _D128_D256) +@pytest.mark.parametrize("d", _D128_D256_PACK_GQA) @torch_fork_set_rng(seed=0) def test_fp8_pack_gqa_e5m2(d): """Packed e5m2 input path.""" @@ -1187,7 +1199,7 @@ def test_fp8_pack_gqa_e5m2(d): _check(out, o_ref, torch.float16, "e5m2", a_o, a_ref) -@_skip_on_rubin +@_skip_pack_gqa_wide_on_rubin @pytest.mark.L0 @pytest.mark.parametrize("h_q,h_kv", [(8, 4), (8, 2), (16, 2)], ids=["g2", "g4", "g8"]) @torch_fork_set_rng(seed=0) @@ -1213,7 +1225,7 @@ def test_fp8_pack_gqa_d192_d128_ratios(h_q, h_kv): torch.testing.assert_close(lse_v, lse_ref, atol=5e-2, rtol=3e-2) -@_skip_on_rubin +@_skip_pack_gqa_wide_on_rubin @pytest.mark.L0 @torch_fork_set_rng(seed=0) def test_fp8_pack_gqa_d192_d128_grouped_lpt(): @@ -1241,7 +1253,7 @@ def test_fp8_pack_gqa_d192_d128_grouped_lpt(): torch.testing.assert_close(lse_v, lse_ref, atol=5e-2, rtol=3e-2) -@_skip_on_rubin +@_skip_pack_gqa_wide_on_rubin @pytest.mark.L1 @torch_fork_set_rng(seed=0) def test_fp8_pack_gqa_d192_d128_e5m2_sink(): @@ -1534,7 +1546,6 @@ def test_fp8_thd(in_key, causal): _check(out, o_ref, torch.float16, in_key, a_o, a_o_ref) -@_skip_on_rubin @pytest.mark.L0 @pytest.mark.parametrize( ("in_key", "causal", "bottom_right"), @@ -1559,7 +1570,6 @@ def test_fp8_d192_d128_thd(in_key, causal, bottom_right): _check(out, o_ref, torch.float16, in_key, a_o, a_o_ref) -@_skip_on_rubin @pytest.mark.L0 @torch_fork_set_rng(seed=0) def test_fp8_d192_d128_wide_swa_boundary_thd(): @@ -1580,7 +1590,6 @@ def test_fp8_d192_d128_wide_swa_boundary_thd(): _check(out, o_ref, torch.float16, "e4m3", a_o, a_o_ref) -@_skip_on_rubin @pytest.mark.L0 @torch_fork_set_rng(seed=0) def test_fp8_d192_d128_thd_features(): @@ -1625,7 +1634,6 @@ def test_fp8_thd_multi_unit_per_cta(monkeypatch, d_qk, d_v): @pytest.mark.L0 -@_skip_d256_on_rubin @pytest.mark.parametrize("in_key", _INS) @pytest.mark.parametrize("causal", [False, True]) @torch_fork_set_rng(seed=0) @@ -1820,17 +1828,16 @@ def test_fp8_sm100_device_scales_execute_reads_no_device_memory(): # # Same cga4x1 role-split kernel family as the d512 f16 sibling; the FP8 fork # runs STAGES_KV = XFER_STAGES = 3 with a 128-column TMEM-resident Q and the -# Blackwell K=32 QMMA step. There is no d512 MXFP8 kernel and no Rubin d512 -# kernel, so this whole section is sm100-only. +# Blackwell K=32 QMMA step (the Rubin sibling sm107/prefill_d512_fp8.py takes +# the K=64 form). Both arch lines serve the flavor, so this section runs on +# both; the MXFP8 d512 kernels have their own suites. # =========================================================================== -_skip_d512_on_rubin = pytest.mark.skipif(_SM == 107, reason="the d512 per-tensor FP8 flavor has no Rubin kernel (sm107 serves d128 only)") _D512 = 512 _D512_SCALE = 1.0 / math.sqrt(_D512) @pytest.mark.L0 -@_skip_d512_on_rubin @pytest.mark.parametrize("in_key", _INS) @pytest.mark.parametrize("mask", list(_MASKS)) @torch_fork_set_rng(seed=0) @@ -1841,7 +1848,6 @@ def test_fp8_d512_masks(in_key, mask): @pytest.mark.L0 -@_skip_d512_on_rubin @pytest.mark.parametrize("out_key", ["fp16", "bf16", "e4m3", "e5m2"]) @pytest.mark.parametrize("in_key", _INS) @torch_fork_set_rng(seed=0) @@ -1856,7 +1862,6 @@ def test_fp8_d512_output_dtypes(in_key, out_key): @pytest.mark.L0 -@_skip_d512_on_rubin @torch_fork_set_rng(seed=0) def test_fp8_d512_sink(): """Attention sink on d512: one extra softmax column with V = 0.""" @@ -1866,7 +1871,6 @@ def test_fp8_d512_sink(): @pytest.mark.L0 -@_skip_d512_on_rubin @pytest.mark.parametrize("h_q,h_kv", [(8, 2), (8, 1)], ids=["g4", "mqa"]) @torch_fork_set_rng(seed=0) def test_fp8_d512_gqa(h_q, h_kv): @@ -1876,7 +1880,6 @@ def test_fp8_d512_gqa(h_q, h_kv): @pytest.mark.L0 -@_skip_d512_on_rubin @torch_fork_set_rng(seed=0) def test_fp8_d512_padded_kv(): """Per-batch KV padding (seq_len_kv) on d512: KV columns past the batch's @@ -1887,7 +1890,6 @@ def test_fp8_d512_padded_kv(): @pytest.mark.L0 -@_skip_d512_on_rubin @torch_fork_set_rng(seed=0) def test_fp8_d512_stats(): """generate_stats on d512: the dense LSE matches the natural-log reference.""" @@ -1898,7 +1900,6 @@ def test_fp8_d512_stats(): @pytest.mark.L0 -@_skip_d512_on_rubin @torch_fork_set_rng(seed=0) def test_fp8_d512_multi_tile(): """A shape that spans several Q tiles and several persistent waves, so the @@ -1914,7 +1915,6 @@ def test_fp8_d512_multi_tile(): @pytest.mark.L0 -@_skip_d512_on_rubin @pytest.mark.parametrize("in_key", _INS) @pytest.mark.parametrize("causal", [False, True]) @torch_fork_set_rng(seed=0) @@ -1926,7 +1926,6 @@ def test_fp8_d512_thd(in_key, causal): @pytest.mark.L0 -@_skip_d512_on_rubin @torch_fork_set_rng(seed=0) def test_fp8_d512_thd_cross_gqa(): """THD cross-attention on d512: unequal per-sequence Q and KV lengths with @@ -1936,7 +1935,6 @@ def test_fp8_d512_thd_cross_gqa(): @pytest.mark.L0 -@_skip_d512_on_rubin @torch_fork_set_rng(seed=0) def test_fp8_d512_thd_sink_stats(): """THD + sink + ragged Stats on d512 (packed token-major TH1 layout).""" @@ -1948,7 +1946,6 @@ def test_fp8_d512_thd_sink_stats(): @pytest.mark.L0 -@_skip_d512_on_rubin @torch_fork_set_rng(seed=0) def test_fp8_d512_thd_zero_len_kv(): """Zero-length KV and Q sequences on d512 (the d128 sibling's shapes). @@ -1970,7 +1967,6 @@ def test_fp8_d512_thd_zero_len_kv(): @pytest.mark.L0 -@_skip_d512_on_rubin @torch_fork_set_rng(seed=0) def test_fp8_d512_thd_cu_seq_len(): """THD on d512 via the cu_seq_len prefix-sum length form.""" @@ -1979,7 +1975,6 @@ def test_fp8_d512_thd_cu_seq_len(): @pytest.mark.L0 -@_skip_d512_on_rubin @pytest.mark.parametrize( ("d_qk", "d_v"), [(384, 448), (464, 368), (272, 272), (512, 496)], @@ -2002,21 +1997,22 @@ def test_fp8_d512_head_dim_envelope(d_qk, d_v, mask): @pytest.mark.L0 -@_skip_d512_on_rubin def test_fp8_d512_envelope_floor_declines_straddling_shapes(): """The D512 flavor serves the (256, 512] band on both head dims; shapes straddling its floor decline. Below it the D256 flavor is exact-shape only (its padded envelope is non-deterministic, see test_fp8_large_flavors_serve_exact_shapes_only), so a shape such as (160, 160) -- too wide for D192/128, not the D256 shape -- is declined too. + Both device lines carry the flavor, so both shape maps are pinned. """ from cudnn.sdpa.fwd.api_dsl import _fp8_envelope_covers, _sm100_fp8_shapes - shapes = _sm100_fp8_shapes(pertensor=True, device_cc=(10, 0)) - for d_qk, d_v in [(256, 256), (384, 448), (464, 368), (272, 272), (512, 512)]: - assert _fp8_envelope_covers(d_qk, d_v, shapes), (d_qk, d_v) - for d_qk, d_v in [(160, 160), (272, 256), (384, 128), (512, 256)]: - assert not _fp8_envelope_covers(d_qk, d_v, shapes), (d_qk, d_v) + for device_cc in ((10, 0), (10, 7)): + shapes = _sm100_fp8_shapes(pertensor=True, device_cc=device_cc) + for d_qk, d_v in [(256, 256), (384, 448), (464, 368), (272, 272), (512, 512)]: + assert _fp8_envelope_covers(d_qk, d_v, shapes), (device_cc, d_qk, d_v) + for d_qk, d_v in [(160, 160), (272, 256), (384, 128), (512, 256)]: + assert not _fp8_envelope_covers(d_qk, d_v, shapes), (device_cc, d_qk, d_v) @pytest.mark.L0 From 45e2c46211bc449c69f34fe6567ea459f6cadb3c Mon Sep 17 00:00:00 2001 From: randers Date: Wed, 23 Sep 2026 20:51:53 +0000 Subject: [PATCH 3/3] test(sdpa/frost): scope the five known e5m2 0.04-bound dataset edges of the per-tensor FP8 forward suite to their O comparison on cc 10.7 The five ids (test_fp8_d192_d128_output_dtypes[e5m2-bf16], test_fp8_pack_gqa_e5m2[d128], test_fp8_thd[True-e5m2], test_fp8_d192_d128_thd[e5m2-True-False], test_fp8_thd_sliding_window[False-e5m2-d128]) read max|O-ref| 0.0404-0.0423 against the shared 4e-2 e5m2 bound on the 204-SM Rubin dataset -- a data-dependent edge, four of them identical on develop, the fifth newly un-skipped here. Instead of a broad xfail, the ids carry a known_e5m2_edge_on_rubin marker and _check turns ONLY their O comparison into an imperative pytest.xfail when the deviation lands inside the measured window (atol, 0.0425]; Amax_O is asserted first, and a build error, a failed run, an Amax_O mismatch, a deviation past the window or any deviation on an unmarked id stay fatal. The bound itself is unchanged. Host pin test_known_e5m2_edge_is_scoped_to_the_o_compare exercises all five outcomes on synthetic tensors. Co-Authored-By: Claude Fable 5.1 --- test/python/pytest.ini | 1 + .../sdpa/frost/test_sdpa_fwd_fp8_sm100.py | 117 ++++++++++++++++-- 2 files changed, 108 insertions(+), 10 deletions(-) diff --git a/test/python/pytest.ini b/test/python/pytest.ini index fb1874c41..f19e35323 100644 --- a/test/python/pytest.ini +++ b/test/python/pytest.ini @@ -8,5 +8,6 @@ markers = L4: specifies L4 level (use -m L4) gpu_exclusive: tests that require exclusive GPU access (no concurrent kernels from other processes) moe_ep_multinode: torchrun-native multi-node MoE EP tests + known_e5m2_edge_on_rubin: sdpa/frost -- per-tensor FP8 ids whose max|O-ref| reads 0.0404-0.0409 on the 204-SM Rubin dataset; _check xfails ONLY that comparison inside the window addopts = -m L0 --tb=short --no-header diff --git a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py index 908a3ea1b..b69d2e61a 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py @@ -62,8 +62,39 @@ reason="sm107 fp8 d256 kernel: strided Stats not ported (contiguous [B, H, S] only); the row declares it -- Capabilities gap, follow-up " "(pinned by test_sm107_fp8_strided_stats_is_not_ported_beyond_d192)", ) +# Five e5m2 ids read `max|O-ref|` 0.0404-0.0423 against the shared 4e-2 e5m2 bound +# (`_half_atol`, tightened by #971) on the 204-SM Rubin part -- a DATASET edge +# (torch's Philox lays draws out by SM count, so the cc 10.7 tensors are not the +# CI's): four of them read identically to 17 digits on develop, the fifth +# (d192x128 THD e5m2 causal) was skipped there and is newly exposed by this file's +# retired Rubin skips. Tracked as a known issue; the bound is NOT widened, and the +# exception is SCOPED: `_check` turns ONLY the O comparison of a marked id into an +# imperative `pytest.xfail` when `max|O-ref|` lands inside the measured window +# `(atol, _E5M2_EDGE_MAX]` -- a build error, a failed `_run`/`_run_thd`, an `Amax_O` +# mismatch (asserted first) or a deviation past the window all stay FATAL, and a +# pass on another dataset is a plain pass (review on PR #1209). Attached per param +# combo (pytest.param) so the sibling e4m3 / dense / d256 combos are untouched. +_E5M2_EDGE_MAX = 4.25e-2 # upper edge of the measured window (0.0404, 0.0406, 0.0409, 0.0409, 0.0423); past it the id is a real failure +_known_e5m2_edge_on_rubin = pytest.mark.known_e5m2_edge_on_rubin +_CURRENT_ITEM_IS_KNOWN_EDGE = False + + +@pytest.fixture(autouse=True) +def _known_e5m2_edge_scope(request): + """Tell ``_check`` whether the running item is one of the marked ids on cc 10.7.""" + global _CURRENT_ITEM_IS_KNOWN_EDGE + _CURRENT_ITEM_IS_KNOWN_EDGE = _SM == 107 and request.node.get_closest_marker("known_e5m2_edge_on_rubin") is not None + yield + _CURRENT_ITEM_IS_KNOWN_EDGE = False + + _D128_D256 = [pytest.param(128, id="d128"), pytest.param(256, id="d256")] _D128_D256_PACK_GQA = [pytest.param(128, id="d128"), pytest.param(256, id="d256", marks=_skip_pack_gqa_wide_on_rubin)] +# test_fp8_pack_gqa_e5m2 only: its d128 combo is one of the five e5m2 edges. +_D128_D256_PACK_GQA_E5M2 = [ + pytest.param(128, id="d128", marks=_known_e5m2_edge_on_rubin), + pytest.param(256, id="d256", marks=_skip_pack_gqa_wide_on_rubin), +] _FP8 = {"e4m3": torch.float8_e4m3fn, "e5m2": torch.float8_e5m2} _FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0} @@ -444,10 +475,16 @@ def _check(out, o_ref, out_dt, in_key, amax_o, amax_o_ref, atol=None): if out_dt in (torch.float8_e4m3fn, torch.float8_e5m2): floor = (o_ref - o_ref.to(out_dt).float()).abs().max().item() atol = max(atol, 3.0 * floor) - assert diff <= atol, f"max|O-ref|={diff:.4f} > {atol:.4f}" # Amax_O is produced in-kernel (atomicMax over the pre-cast fp32 # values), so they match the exact fp32 reference for every output dtype, incl. FP8. + # Asserted BEFORE the O comparison so the known-edge xfail below can never hide it. assert abs(amax_o - amax_o_ref) <= 0.03, f"amax_o {amax_o:.4f} vs ref {amax_o_ref:.4f}" + if diff > atol and _CURRENT_ITEM_IS_KNOWN_EDGE and diff <= _E5M2_EDGE_MAX: + pytest.xfail( + f"known e5m2 dataset edge on the 204-SM Rubin part: max|O-ref|={diff:.4f} in ({atol:.4f}, {_E5M2_EDGE_MAX:.4f}] " + "-- the bound is not widened; build, run and Amax_O stayed fatal (PR #1209)" + ) + assert diff <= atol, f"max|O-ref|={diff:.4f} > {atol:.4f}" _INS = ["e4m3", "e5m2"] @@ -461,6 +498,32 @@ def _check(out, o_ref, out_dt, in_key, amax_o, amax_o_ref, atol=None): } +@pytest.mark.L0 +def test_known_e5m2_edge_is_scoped_to_the_o_compare(monkeypatch): + """The Rubin known-edge exception covers ONLY an O deviation inside (atol, _E5M2_EDGE_MAX] on a marked id. + + Everything else through ``_check`` stays fatal: a deviation past the window, an ``Amax_O`` mismatch + (asserted first), and any deviation on an unmarked id. Host-only: synthetic CPU tensors.""" + ref = torch.zeros(4, 8, dtype=torch.float32) + atol = _half_atol("e5m2", 8) # 4e-2 + + def run(diff, amax_off=0.0, marked=True): + monkeypatch.setitem(globals(), "_CURRENT_ITEM_IS_KNOWN_EDGE", marked) + out = ref.clone() + out[0, 0] = diff + _check(out, ref, torch.float16, "e5m2", 1.0 + amax_off, 1.0) + + run(atol - 1e-3) # inside the bound: a plain pass + with pytest.raises(pytest.xfail.Exception, match="known e5m2 dataset edge"): + run(0.0405) # the measured window -> xfail, not a pass + with pytest.raises(AssertionError, match=r"max\|O-ref\|"): + run(_E5M2_EDGE_MAX + 1e-3) # past the window -> fatal + with pytest.raises(AssertionError, match="amax_o"): + run(0.0405, amax_off=0.1) # Amax_O mismatch -> fatal even inside the window + with pytest.raises(AssertionError, match=r"max\|O-ref\|"): + run(0.0405, marked=False) # an unmarked id never gets the exception + + def _check_fp8_strided_stats(d_qk, d_v, in_key): kwargs = dict( B=2, @@ -563,9 +626,22 @@ def test_fp8_block_scaled_output_declines_wide_flavor(): _run(1, 2, 2, 256, 256, "e4m3", torch.float8_e4m3fn, scale=1.0 / 16, sdpa_kwargs={}, d_qk=256, d_v=256, block_scaled_o="nvfp4") +# in_key x out_key with the ids the former stacked parametrize produced (`-`); +# one combined list so ONLY e5m2-bf16 carries the Rubin e5m2-edge marker. +_D192_D128_OUTPUT_DTYPE_CASES = [ + pytest.param("e4m3", "fp16", id="e4m3-fp16"), + pytest.param("e4m3", "bf16", id="e4m3-bf16"), + pytest.param("e4m3", "e4m3", id="e4m3-e4m3"), + pytest.param("e4m3", "e5m2", id="e4m3-e5m2"), + pytest.param("e5m2", "fp16", id="e5m2-fp16"), + pytest.param("e5m2", "bf16", id="e5m2-bf16", marks=_known_e5m2_edge_on_rubin), + pytest.param("e5m2", "e4m3", id="e5m2-e4m3"), + pytest.param("e5m2", "e5m2", id="e5m2-e5m2"), +] + + @pytest.mark.L0 -@pytest.mark.parametrize("out_key", ["fp16", "bf16", "e4m3", "e5m2"]) -@pytest.mark.parametrize("in_key", ["e4m3", "e5m2"]) +@pytest.mark.parametrize(("in_key", "out_key"), _D192_D128_OUTPUT_DTYPE_CASES) @torch_fork_set_rng(seed=0) def test_fp8_d192_d128_output_dtypes(in_key, out_key): """Exact DSv3 shape: FP8 Q/K use d192 while V and O use d128.""" @@ -1177,7 +1253,7 @@ def test_fp8_pack_gqa_features(d, mask, out_dt): @pytest.mark.L1 -@pytest.mark.parametrize("d", _D128_D256_PACK_GQA) +@pytest.mark.parametrize("d", _D128_D256_PACK_GQA_E5M2) @torch_fork_set_rng(seed=0) def test_fp8_pack_gqa_e5m2(d): """Packed e5m2 input path.""" @@ -1535,9 +1611,18 @@ def _run(declare): assert torch.equal(declared[0], inferred[0]), "declaring the packed totals must not change O" +# causal x in_key with the ids the former stacked parametrize produced (`-`); +# ONLY True-e5m2 carries the Rubin e5m2-edge marker. +_THD_CASES = [ + pytest.param(False, "e4m3", id="False-e4m3"), + pytest.param(False, "e5m2", id="False-e5m2"), + pytest.param(True, "e4m3", id="True-e4m3"), + pytest.param(True, "e5m2", id="True-e5m2", marks=_known_e5m2_edge_on_rubin), +] + + @pytest.mark.L0 -@pytest.mark.parametrize("in_key", _INS) -@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize(("causal", "in_key"), _THD_CASES) @torch_fork_set_rng(seed=0) def test_fp8_thd(in_key, causal): """THD/varlen self-attention: two packed sequences of unequal, tile-ragged length.""" @@ -1549,7 +1634,7 @@ def test_fp8_thd(in_key, causal): @pytest.mark.L0 @pytest.mark.parametrize( ("in_key", "causal", "bottom_right"), - [("e4m3", False, False), ("e5m2", True, False), ("e4m3", False, True)], + [("e4m3", False, False), pytest.param("e5m2", True, False, marks=_known_e5m2_edge_on_rubin), ("e4m3", False, True)], ) @torch_fork_set_rng(seed=0) def test_fp8_d192_d128_thd(in_key, causal, bottom_right): @@ -1653,10 +1738,22 @@ def test_fp8_d256_thd(in_key, causal): _check(out, o_ref, torch.float16, in_key, a_o, a_o_ref) +# bottom_right x in_key x d with the ids the former stacked parametrize produced (`--`); +# ONLY False-e5m2-d128 carries the Rubin e5m2-edge marker. +_THD_SLIDING_WINDOW_CASES = [ + pytest.param(False, "e4m3", 128, id="False-e4m3-d128"), + pytest.param(False, "e4m3", 256, id="False-e4m3-d256"), + pytest.param(False, "e5m2", 128, id="False-e5m2-d128", marks=_known_e5m2_edge_on_rubin), + pytest.param(False, "e5m2", 256, id="False-e5m2-d256"), + pytest.param(True, "e4m3", 128, id="True-e4m3-d128"), + pytest.param(True, "e4m3", 256, id="True-e4m3-d256"), + pytest.param(True, "e5m2", 128, id="True-e5m2-d128"), + pytest.param(True, "e5m2", 256, id="True-e5m2-d256"), +] + + @pytest.mark.L0 -@pytest.mark.parametrize("d", _D128_D256) -@pytest.mark.parametrize("in_key", _INS) -@pytest.mark.parametrize("bottom_right", [False, True]) +@pytest.mark.parametrize(("bottom_right", "in_key", "d"), _THD_SLIDING_WINDOW_CASES) @torch_fork_set_rng(seed=0) def test_fp8_thd_sliding_window(d, in_key, bottom_right): """THD sliding window uses each sequence's local causal diagonal."""