diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 347fe2faa..3c9e26f4d 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -33,6 +33,7 @@ SEQ_KV_TILES as _SM120_KV_TILES, SEQ_Q_TILES as _SM120_Q_TILES, SUPPORTED_HEAD_TILE_MAX as _SM120_HEAD_TILE_MAX, + SUPPORTED_HEAD_TILES_FP8 as _SM120_FP8_HEAD_TILES, TemplateParams as Sm120TemplateParams, smem_bytes as _sm120_smem_bytes, ) @@ -61,11 +62,14 @@ def _require_reciprocal_s_scales(descale_s: float, scale_s: float) -> None: This row cannot apply Scale_S, and would gain nothing if it could: - - No headroom. The lazy-rescale skip (RESCALE_THRESHOLD=8) refreshes the - running max only when a tile exceeds it by 2^8, so P is bounded by 256, - not 1. e4m3 tops out at 448, so any scale_s > 448/256 = 1.75 can saturate - a lazily-skipped tile. Measured on B2xH8xS256 e4m3: max|O-ref| is flat - from scale_s 1 to 64 and degrades at 448 (swa .0239 -> .0807). + - No headroom. The lazy-rescale skip refreshes the running max only when + a tile exceeds it by RESCALE_THRESHOLD -- 4.0 for the fp8 dtypes + (config_sm100.rescale_threshold; 8.0 is the dataclass default the fp8 + path overrides) -- so P is bounded by 2^4 = 16, not 1. e4m3 tops out at + 448, so any scale_s > 448/16 = 28 can saturate a lazily-skipped tile. + Measured on B2xH8xS256 e4m3: max|O-ref| is flat from scale_s 1 to 64 + (the analytical bound is conservative) and degrades at 448 + (swa .0239 -> .0807). - Nothing to gain. e4m3 is floating point, so relative precision does not move with scale, and subtracting the row max already places P per ROW — strictly better than a per-tensor scale. Hence the flat error above. @@ -1839,8 +1843,9 @@ def check_support(self) -> bool: self._value_error_if(self.has_sink, "SM120 fp8 does not support attention sinks (Amax_S semantics)") self._value_error_if(self.seq_q_lens_present and not self.thd, "SM120 fp8 does not support per-batch seq_len_q") self._value_error_if( - (d_q, d_v) != (128, 128), - f"SM120 fp8 requires D_QK=D_V=128 (no zero-padding envelope on the 8-bit fragment path); got ({d_q}, {d_v})", + any(d not in _SM120_FP8_HEAD_TILES for d in (d_q, d_v)), + f"SM120 fp8 requires D_QK and D_V to be multiples of 32 within 32..256 (k32 contraction and 1-byte " + f"TMA swizzle span; no zero-padding envelope on the 8-bit fragment path); got ({d_q}, {d_v})", ) self._value_error_if( @@ -1896,8 +1901,6 @@ def check_support(self) -> bool: def _smem_bytes(kv_tile: int) -> int: # FP8 stages a byte per KV element but still writes O in half. - # (The FP8 row requires exact d128, so its padded dims are the - # actual ones; the envelope padding is the f16 cell's.) return _sm120_smem_bytes(d_qp, d_vp, self.q_tile, kv_tile, self.dtype.itemsize, 2 if self._fp8 else self.dtype.itemsize) if self.tile_n is None: diff --git a/python/cudnn/sdpa/fwd/config_sm120.py b/python/cudnn/sdpa/fwd/config_sm120.py index bd4b984d3..1277ea5a6 100644 --- a/python/cudnn/sdpa/fwd/config_sm120.py +++ b/python/cudnn/sdpa/fwd/config_sm120.py @@ -16,6 +16,8 @@ SUPPORTED_HEAD_TILE_MIN = 16 SUPPORTED_HEAD_TILE_MAX = 256 SUPPORTED_HEAD_TILES = tuple(range(SUPPORTED_HEAD_TILE_MIN, SUPPORTED_HEAD_TILE_MAX + 1, HEAD_TILE_GRANULE)) +FP8_HEAD_TILE_GRANULE = 32 +SUPPORTED_HEAD_TILES_FP8 = tuple(range(FP8_HEAD_TILE_GRANULE, SUPPORTED_HEAD_TILE_MAX + 1, FP8_HEAD_TILE_GRANULE)) # SMEM the SM120 parts expose to a kernel. The adapter asks cutlass for the # authoritative number at build time; this constant lets the ranking answer diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index b3e8930fd..014538e9a 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -267,11 +267,11 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti return f"serves D_QK in {sorted(capabilities.d_qk)}/D_V in {sorted(capabilities.d_v)}; graph has D_QK={facts.d_qk}/D_V={facts.d_v}" if facts.dtype not in capabilities.dtypes: return f"dtype {facts.dtype} not in {sorted(str(d) for d in capabilities.dtypes)}" - if (capabilities.is_fp8 or capabilities.is_mxfp8) and facts.dtype_o not in capabilities.out_dtypes: - return f"O dtype {facts.dtype_o} not in {sorted(str(d) for d in capabilities.out_dtypes)}" if (facts.is_mxfp8, facts.is_fp8) != (capabilities.is_mxfp8, capabilities.is_fp8): quant = "block-scale MXFP8 (sdpa_mxfp8)" if capabilities.is_mxfp8 else "per-tensor FP8 (sdpa_fp8)" if capabilities.is_fp8 else "half (sdpa)" return f"this engine serves only {quant} graphs" + if (capabilities.is_fp8 or capabilities.is_mxfp8) and facts.dtype_o not in capabilities.out_dtypes: + return f"O dtype {facts.dtype_o} not in {sorted(str(d) for d in capabilities.out_dtypes)}" if not facts.uniform_dtype: return "K/V dtypes must match Q" if (facts.is_mxfp8 or facts.is_fp8) else "K/V/O dtypes must match Q" if facts.thd: @@ -718,10 +718,9 @@ def _execute(variant_pack, workspace=None, stream=None): # construction, and the packed layout gives each sequence its own # extent, so nothing is written past a valid length. if (facts.is_mxfp8 or facts.is_fp8) and not facts.thd and seq_q_buf is not None and not seq_q_lens_present: - if int(seq_q_buf.min().item()) < int(facts.s_q): - raise NotImplementedError( - f"per-tensor FP8/MXFP8: per-batch seq_len_q shorter than S_q={facts.s_q} is not plumbed; got min {int(seq_q_buf.min().item())}" - ) + min_seq_q = int(seq_q_buf.min().item()) + if min_seq_q < int(facts.s_q): + raise NotImplementedError(f"per-tensor FP8/MXFP8: per-batch seq_len_q shorter than S_q={facts.s_q} is not plumbed; got min {min_seq_q}") execute_kwargs = dict( q_tensor=q_buf, k_tensor=k_buf, @@ -811,11 +810,13 @@ def _sm120_fp8_spec() -> EngineSpec: m16n8k32 e4m3; ``descale_q*descale_k`` folds into the softmax scale and ``descale_v*scale_o`` into an epilogue scalar, so the kernel adds only the Amax_S/Amax_O atomics over the f16 sibling. E4M3 only (no E5M2 tag in the - kernel yet), FP16 O only, exact d128 (no zero-padding envelope on the + kernel yet), FP16 O only, head dims any multiple of 32 up to 256 with the + QK^T and P@V sides independent (exact — no zero-padding envelope on the 8-bit fragment path), and no sink (Amax_S semantics with a sink column are undefined here). THD (ragged) is served with token-major Stats; head-major ragged Stats stays f16-only (the fp8 kernel carries no such specialization). """ + from cudnn.sdpa.fwd.config_sm120 import SUPPORTED_HEAD_TILES_FP8 return EngineSpec( name="sdpa_fwd_prefill_sm120_fp8", @@ -823,8 +824,8 @@ def _sm120_fp8_spec() -> EngineSpec: sm_lo=_BLACKWELL_GEFORCE[0], sm_hi=_BLACKWELL_GEFORCE[1], phase="prefill", - d_qk=frozenset({128}), - d_v=frozenset({128}), + d_qk=frozenset(SUPPORTED_HEAD_TILES_FP8), + d_v=frozenset(SUPPORTED_HEAD_TILES_FP8), dtypes=frozenset({cudnn.data_type.FP8_E4M3}), out_dtypes=frozenset({cudnn.data_type.HALF}), is_fp8=True, diff --git a/python/cudnn/sdpa/fwd/heuristics.py b/python/cudnn/sdpa/fwd/heuristics.py index 39d91d0a7..b23f58f44 100644 --- a/python/cudnn/sdpa/fwd/heuristics.py +++ b/python/cudnn/sdpa/fwd/heuristics.py @@ -78,8 +78,8 @@ def _sm120_tiles(caps: Capabilities, facts) -> Tuple[int, int]: tile_m = 64 if fine else 128 # FP8 stages a byte per KV element but still writes O in half, so the two # SMEM terms size differently -- see config_sm120.smem_bytes. - qkv_item, o_item = (1, 2) if facts.is_fp8 else (2, 2) - fits = [n for n in sorted(caps.tile_ns, reverse=True) if smem_bytes(facts.d_qk, facts.d_v, tile_m, n, qkv_item, o_item) <= SMEM_CAPACITY_BYTES] + qkv_itemsize, o_itemsize = (1, 2) if facts.is_fp8 else (2, 2) + fits = [n for n in sorted(caps.tile_ns, reverse=True) if smem_bytes(facts.d_qk, facts.d_v, tile_m, n, qkv_itemsize, o_itemsize) <= SMEM_CAPACITY_BYTES] return tile_m, (fits[0] if fits else min(caps.tile_ns)) @@ -124,7 +124,10 @@ def _mode_a(facts, offered: Dict[str, int], mode) -> List[PlanConfig]: # the rule's regret is small but not zero, so the runners-up are worth # offering to a caller who measures. Configs the kernel cannot fit are # not runners-up -- they would sit in the list only to decline at build. - domain = [(m, n) for m in caps.tile_ms for n in caps.tile_ns if smem_bytes(facts.d_qk, facts.d_v, m, n) <= SMEM_CAPACITY_BYTES] + qkv_itemsize, o_itemsize = (1, 2) if facts.is_fp8 else (2, 2) + domain = [ + (m, n) for m in caps.tile_ms for n in caps.tile_ns if smem_bytes(facts.d_qk, facts.d_v, m, n, qkv_itemsize, o_itemsize) <= SMEM_CAPACITY_BYTES + ] ordered = sorted(domain or [best], key=lambda mn: (mn != best, mn[1] != best[1], -mn[0])) for tile_m, tile_n in ordered: knobs = _knobs(caps, tile_m, tile_n) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py index 6915c357a..7d36e7ca5 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py @@ -36,7 +36,7 @@ Constraints: * Input dtype: e4m3 only (as Uint8 storage); output dtype Float16 -* Head dimensions must be multiples of 16 between 16 and 256, inclusive +* Head dimensions must be multiples of 32 between 32 and 256, inclusive * Q heads must be divisible by the number of K/V heads * Q/K/V/O use compact BSHD storage (THD packs them to ``(1, T, H, D)``) * Supported CTA Q/KV tiles are 128 or 64 @@ -217,7 +217,7 @@ def __init__( :param thd_batch: THD only: the real sequence count B. :param thd_max_sq: THD only: the longest sequence's Q length. :param head_tile_qk: Q/K head dimension (the QK^T contraction width). - Must be a multiple of 16 between 16 and 256, inclusive. + Must be a multiple of 32 between 32 and 256, inclusive. :param head_tile_v: V/O head dimension (the P@V output width). Same constraint as ``head_tile_qk``. :param q_tile: Query sequence tile size. @@ -494,7 +494,7 @@ def online_softmax( row_max = softmax_params.row_max row_sum = softmax_params.row_sum softmax_scale_log2 = softmax_params.softmax_scale_log2 - scale_s = softmax_params.scale_s + log2_scale_s = softmax_params.log2_scale_s p_regs = cutlass.Array(cutlass.Uint16, self.qk_k_frags * 2) # Each lane owns four S registers split across two Q rows after Q@K^T. @@ -555,22 +555,42 @@ def online_softmax( else: row_max_prev = row_max[row_half] new_max = cute.arch.fmax(row_max_prev, cur_max) - need_correct = True - if cutlass.const_expr(in_mask_steps): - need_correct = new_max > -cutlass.Float32.inf - if need_correct: - old_scale = cute.math.exp2( - (row_max_prev - new_max) * softmax_scale_log2, - fastmath=True, - ) + # Keep this as inline PTX so old_scale lowers to one predicated + # EX2 with 1.0 as the default value. The equivalent Python DSL + # branch currently materializes extra MOV instructions in this + # hot softmax loop and increases issue pressure on SM120. + old_scale = cute.arch.inline_ptx( + ( + "{\n" + " .reg .pred p;\n" + " .reg .f32 delta;\n" + " sub.rn.f32 delta, $1, $2;\n" # delta = row_max_prev - new_max + " mul.rn.f32 delta, delta, $3;\n" # delta = delta * softmax_scale_log2 + " setp.gt.f32 p, $2, $1;\n" # p = new_max > row_max_prev + " mov.f32 $0, 0f3f800000;\n" # res = 1.0 + " @p ex2.approx.ftz.f32 $0, delta;\n" # if p: res = exp2(delta) + "}" + ), + write_only_types=[cutlass.Float32], + read_only_args=[row_max_prev, new_max, softmax_scale_log2], + ) row_max[row_half] = new_max + if cutlass.const_expr(not is_first_kv_tile): + for d_frag in cutlass.range_constexpr(self.pv_d_frags): + o_off = d_frag * 4 + row_half * 2 + if new_max > row_max_prev: + o_regs[o_off + 0], o_regs[o_off + 1] = fmul2( + (o_regs[o_off + 0], o_regs[o_off + 1]), + (old_scale, old_scale), + ) + # Compute P, accumulate the per-lane partial sum, and stage P. exp_max = new_max if cutlass.const_expr(in_mask_steps): if exp_max == -cutlass.Float32.inf: exp_max = cutlass.Float32(0.0) - neg_exp_max_scaled = -(exp_max * softmax_scale_log2) + neg_exp_max_scaled = log2_scale_s - exp_max * softmax_scale_log2 tile_sum = cutlass.Float32(0.0) for k_frag in cutlass.range_constexpr(self.qk_k_frags): s_off = k_frag * 4 @@ -585,29 +605,21 @@ def online_softmax( p1 = cute.math.exp2(in1, fastmath=True) tile_sum = tile_sum + (p0 + p1) # P stays in registers at the C-fragment coordinates; mma_pv - # redistributes it to the k32 A layout with shfl. row_sum above - # uses the UNSCALED fp32 P -- the denominator, and so amax_s, - # must not see Scale_S. Scale_S applies only to the value that - # gets cast, and descale_s is folded into o_scale_fused. - ps0, ps1 = fmul2((p0, p1), (scale_s, scale_s)) - p_regs[k_frag * 2 + row_half] = ptx_cvt_e4m3x2(ps1, ps0) + # redistributes it to the k32 A layout with shfl. p0/p1 carry + # Scale_S via the exp2 bias, so the cast input is the scaled + # value the contract asks for; descale_s stays folded into + # o_scale_fused. + p_regs[k_frag * 2 + row_half] = ptx_cvt_e4m3x2(p1, p0) # Reduce tile_sum across the four lanes that own one Q row. tile_sum = nvvm_threadquad_reduction_sum(tile_sum) - # Correct row_sum and rescale O when row_max changes. + # Correct row_sum (old_scale is exactly 1.0 when the max held). if cutlass.const_expr(is_first_kv_tile): row_sum[row_half] = tile_sum else: row_sum[row_half] = row_sum[row_half] * old_scale + tile_sum - for d_frag in cutlass.range_constexpr(self.pv_d_frags): - o_off = d_frag * 4 + row_half * 2 - o_regs[o_off + 0], o_regs[o_off + 1] = fmul2( - (o_regs[o_off + 0], o_regs[o_off + 1]), - (old_scale, old_scale), - ) - return p_regs @cute.jit @@ -1046,11 +1058,13 @@ def kernel( sV=sV, o_regs=o_regs, ) + log2_scale_s = cute.math.log2(scale_s, fastmath=True) + inv_scale_s = cutlass.Float32(1.0) / scale_s softmax_params = SimpleNamespace( row_max=row_max, row_sum=row_sum, softmax_scale_log2=softmax_scale_log2, - scale_s=scale_s, + log2_scale_s=log2_scale_s, ) # Load Q into registers. @@ -1063,6 +1077,8 @@ def kernel( if cutlass.const_expr(self.bottom_right): # The shifted diagonal can straddle one additional KV tile. mask_steps = ceil_div(self.q_tile + self.kv_tile - 1, self.kv_tile) + elif cutlass.const_expr(not self.seq_kv_lens_present and not self.thd_varlen and k.shape[1] % self.kv_tile == 0): + mask_steps = 0 left_mask_steps = 1 if cutlass.const_expr(self.window_size_left is not None): left_mask_steps = ceil_div(self.q_tile + self.kv_tile - 1, self.kv_tile) @@ -1131,38 +1147,24 @@ def kernel( # the four lanes that share a Q row, so every lane finalizes the two # rows it owns without further exchange. row_max holds the raw (unscaled) # score max; the scale is applied in log2 domain and converted with ln(2). - # With has_sink, the per-head sink logit joins the softmax denominator - # as a virtual column with no V row: it rescales O, enters the LSE, - # and gives a row with no visible key a finite LSE (the sink alone). LN2 = cutlass.Float32(0.6931471805599453) row_sum_inv = cutlass.Array(cutlass.Float32, 2, alignment=8) row_lse = cutlass.Array(cutlass.Float32, 2, alignment=8) for row_half in cutlass.range_constexpr(2): + row_sum[row_half] = row_sum[row_half] * inv_scale_s row_max_nat = row_max[row_half] * softmax_scale_log2 * LN2 - if cutlass.const_expr(self.has_sink): - sinks_arr = cutlass.make_array_view(sinks) - sink_logit = cutlass.Float32(sinks_arr[head_idx]) - new_max = cute.arch.fmax(row_max_nat, sink_logit) - # alpha re-normalizes the loop's accumulator and sum from - # row_max_nat to the sink-extended max; it is 0 for a row - # with no visible key, so O := 0 falls out. - alpha = cute.math.exp(row_max_nat - new_max, fastmath=True) - new_sum = row_sum[row_half] * alpha + cute.math.exp(sink_logit - new_max, fastmath=True) - row_sum_inv[row_half] = alpha / new_sum - row_lse[row_half] = new_max + cute.math.log(new_sum, fastmath=True) - else: - inv = cutlass.Float32(0.0) - if row_sum[row_half] > 0.0: - inv = cute.math.rcp(row_sum[row_half], approx=True, ftz=True) - row_sum_inv[row_half] = inv - lse_val = row_max_nat + cute.math.log( - cute.math.max(row_sum[row_half], cutlass.Float32(1e-30)), - fastmath=True, - ) - # Rows with no visible key write -inf / O := 0. - if row_sum[row_half] <= 0.0: - lse_val = -cutlass.Float32.inf - row_lse[row_half] = lse_val + inv = cutlass.Float32(0.0) + if row_sum[row_half] > 0.0: + inv = cute.math.rcp(row_sum[row_half], approx=True, ftz=True) + row_sum_inv[row_half] = inv + lse_val = row_max_nat + cute.math.log( + cute.math.max(row_sum[row_half], cutlass.Float32(1e-30)), + fastmath=True, + ) + # Rows with no visible key write -inf / O := 0. + if row_sum[row_half] <= 0.0: + lse_val = -cutlass.Float32.inf + row_lse[row_half] = lse_val # Amax_S = max over valid rows of the raw 1/row_sum (cuDNN # convention: the softmax-probability amax proxy), captured BEFORE @@ -1233,14 +1235,17 @@ def kernel( for row_half in cutlass.range_constexpr(2): amax_q_idx = q_seq_idx + q_warp_row0 + (lane // 4) + row_half * 8 row_valid[row_half] = cutlass.Float32(1.0) if amax_q_idx < seqlen_q else cutlass.Float32(0.0) - lane_amax_o = cutlass.Float32(0.0) + lane_amax_half = cutlass.Array(cutlass.Float32, 4, alignment=16) + for i in cutlass.range_constexpr(4): + lane_amax_half[i] = 0.0 for d_frag_pair in cutlass.range_constexpr(self.pv_d_frags // 2): o_off = (d_frag_pair * 2) * 4 o_scaled = fmul2(o_regs[o_off:8], row_sum_inv_vec) for i in cutlass.range_constexpr(8): # row_sum_inv_vec order: halves alternate 0,0,1,1,0,0,1,1. half = (i // 2) % 2 - lane_amax_o = cute.arch.fmax(lane_amax_o, cute.math.abs(o_scaled[i]) * row_valid[half]) + acc = half * 2 + (i % 2) + lane_amax_half[acc] = cute.arch.fmax(lane_amax_half[acc], cute.math.abs(o_scaled[i])) o_packed = o_scaled.to(self.out_dtype).bitcast(cutlass.Int32) sO_ptr = sO.data_ptr() + (compute_warp_idx * (self.pv_d_frags // 2) + d_frag_pair) * o_block_bytes + lane * 16 prims.stmatrix( @@ -1248,6 +1253,10 @@ def kernel( o_packed, prims.MMALayout.ROW, ) + lane_amax_o = cute.arch.fmax( + cute.arch.fmax(lane_amax_half[0], lane_amax_half[1]) * row_valid[0], + cute.arch.fmax(lane_amax_half[2], lane_amax_half[3]) * row_valid[1], + ) amax_o_arr = cutlass.make_array_view(amax_o) prims.atomicrmw( prims.AtomicOp.MAX, @@ -1328,7 +1337,7 @@ def __call__( :param amax_s: 1-element Int32 amax buffer (pre-zeroed; dummy ok). :param amax_o: 1-element Int32 amax buffer (pre-zeroed; dummy ok). :param softmax_scale_log2: ``softmax_scale * descale_q * descale_k * log2(e)``. - :param o_scale_fused: ``descale_v * scale_o``. + :param o_scale_fused: ``descale_s * descale_v * scale_o``. :param stream: CUDA stream used for the launch. """ head_dim_qk = q.shape[3] @@ -1365,10 +1374,8 @@ def __call__( raise ValueError("LSE must have shape (B, H, Sq)") if cutlass.const_expr(lse.stride != (q.shape[2] * q.shape[1], q.shape[1], 1)): raise ValueError("LSE must be compact row-major") - if cutlass.const_expr(self.has_sink != (sinks is not None)): - raise ValueError("sinks must be provided exactly when the kernel is configured with has_sink") - if cutlass.const_expr(sinks is not None and sinks.shape != (q.shape[2],)): - raise ValueError("sinks must have shape (H,)") + if cutlass.const_expr(sinks is not None): + raise ValueError("sinks must be None (the fp8 cell rejects has_sink)") if cutlass.const_expr(self.thd_varlen): if cutlass.const_expr(q.shape[0] != 1): raise ValueError("THD Q/K/V/O must be packed batch-1 views") diff --git a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py index 28b999e06..2345e200a 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py @@ -8,13 +8,14 @@ reference. ``Amax_S`` and ``Amax_O`` are both produced in-kernel (bitcast-int32 atomicMax over the pre-cast fp32 values); both are checked. -SM120 envelope (see engines._sm120_fp8_spec): E4M3 in / FP16 out only, exact -d128, causal / bottom-right / SWA / KV-padding masks, THD (ragged) with -token-major Stats; no sink (Amax_S semantics), no head-major ragged Stats. -Everything outside that envelope is a capability-row decline, checked by the -*_not_offered tests below. - -Requires: SM120/SM121 (consumer Blackwell), cutlass-dsl. Skips otherwise. +SM120 envelope (see engines._sm120_fp8_spec): E4M3 in / FP16 out only, head +dims any multiple of 32 up to 256 with the QK^T and P@V sides independent +(what graphs can reach is further gated by the C++ sdpa_fp8 node: +d_qk <= 128 x d_v <= 128 plus the (192, 128) MLA pair), causal / +bottom-right / SWA / KV-padding masks, THD (ragged) with token-major Stats; +no sink (Amax_S semantics), no head-major ragged Stats. + +Requires: SM120/SM121 (Blackwell GeForce), cutlass-dsl. Skips otherwise. """ import math @@ -25,7 +26,7 @@ from test_utils import torch_fork_set_rng from cudnn.sdpa.fwd.engines import engine_name -from frost_test_utils import requires_blackwell_geforce, requires_dsl, select_engine as _select_engine +from frost_test_utils import requires_blackwell_geforce, requires_dsl, select_engine as _select_engine, offers_engine pytestmark = [requires_blackwell_geforce, requires_dsl] @@ -62,14 +63,14 @@ def _ref(qd, kd, vd, *, scale, is_causal=False, bottom_right=False, swa_window=N return torch.matmul(probs, v_e), probs.max().item(), torch.logsumexp(scores, dim=-1) -def _run(B, H_q, H_kv, S_q, S_kv, *, scale, sdpa_kwargs, seq_lens_kv=None, tiles=None, s_descale_gain=1.0): +def _run(B, H_q, H_kv, S_q, S_kv, *, scale, sdpa_kwargs, seq_lens_kv=None, tiles=None, s_descale_gain=1.0, D=128, D_v=None): import cudnn dev = "cuda" - D = 128 + D_v = D if D_v is None else D_v Qf = torch.randn(B, H_q, S_q, D, device=dev) * 0.5 Kf = torch.randn(B, H_kv, S_kv, D, device=dev) * 0.5 - Vf = torch.randn(B, H_kv, S_kv, D, device=dev) * 0.5 + Vf = torch.randn(B, H_kv, S_kv, D_v, device=dev) * 0.5 Q8, dq = _quant(Qf) K8, dk = _quant(Kf) V8, dv = _quant(Vf) @@ -78,7 +79,7 @@ def bshd(x8): return x8.permute(0, 2, 1, 3).contiguous().transpose(1, 2) Qb, Kb, Vb = bshd(Q8), bshd(K8), bshd(V8) - Ob = torch.empty(B, S_q, H_q, D, device=dev, dtype=torch.float16).transpose(1, 2) + Ob = torch.empty(B, S_q, H_q, D_v, device=dev, dtype=torch.float16).transpose(1, 2) lse = torch.empty(B, H_q, S_q, 1, device=dev, dtype=torch.float32) amax_s = torch.zeros(1, 1, 1, 1, device=dev, dtype=torch.float32) amax_o = torch.zeros(1, 1, 1, 1, device=dev, dtype=torch.float32) @@ -214,6 +215,31 @@ def test_fp8_sm120_padding(causal): _check(*res) +@pytest.mark.L0 +@pytest.mark.parametrize("mask", ["none", "causal"]) +@pytest.mark.parametrize("s", [96, 97, 110], ids=lambda s: f"s{s}") +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_single_kv_tile(mask, s): + """Single-KV-tile shapes (s_kv <= kv_tile), repeated.""" + scale = 1.0 / math.sqrt(128) + seq_lens_kv = [s] if mask == "none" else None + for _ in range(3): + res = _run(1, 4, 4, s, s, scale=scale, sdpa_kwargs=_MASKS[mask], seq_lens_kv=seq_lens_kv) + _check(*res) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_single_kv_tile_padded(): + """Padded lengths inside a single KV tile (the other collapse-exposed + population: min(seq_kv_lens[b], shape) also bounds the trip count, so the + same collapsed-pipeline layout is emitted here).""" + scale = 1.0 / math.sqrt(128) + for _ in range(3): + res = _run(2, 4, 4, 97, 97, scale=scale, sdpa_kwargs=dict(use_causal_mask=True), seq_lens_kv=[97, 60]) + _check(*res) + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) def test_fp8_sm120_e5m2_not_offered(): @@ -255,7 +281,7 @@ def test_fp8_sm120_e5m2_not_offered(): assert engine_name(arch="sm120", fp8=True) not in names, f"E5M2 graph must not offer the sm120 fp8 engine; plans={names}" -def _fp8_graph_offers_sm120(io_dtype, o_dtype, D=128, sink=False): +def _fp8_graph_offers_sm120(io_dtype, o_dtype, D=128, D_v=None, sink=False): """Build one sdpa_fp8 graph and report whether the sm120 fp8 cell claims it. A capability rejection is the point, so nothing is executed; a graph that @@ -265,10 +291,12 @@ def _fp8_graph_offers_sm120(io_dtype, o_dtype, D=128, sink=False): dev = "cuda" B, H, S = 1, 4, 256 + D_v = D if D_v is None else D_v torch_in = torch.float8_e5m2 if io_dtype == cudnn.data_type.FP8_E5M2 else torch.float8_e4m3fn X = torch.randn(B, S, H, D, device=dev).to(torch_in).transpose(1, 2) + Xv = torch.randn(B, S, H, D_v, device=dev).to(torch_in).transpose(1, 2) g = cudnn.pygraph(io_data_type=io_dtype, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) - q, k, v = g.tensor_like(X), g.tensor_like(X), g.tensor_like(X) + q, k, v = g.tensor_like(X), g.tensor_like(X), g.tensor_like(Xv) scalars = [g.tensor(dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT) for _ in range(6)] kw = dict( q=q, @@ -286,7 +314,7 @@ def _fp8_graph_offers_sm120(io_dtype, o_dtype, D=128, sink=False): if sink: kw["sink_token"] = g.tensor(dim=[1, H, 1, 1], stride=[H, 1, 1, 1], data_type=cudnn.data_type.FLOAT) o, stats, amx_s, amx_o = g.sdpa_fp8(**kw) - o.set_output(True).set_dim([B, H, S, D]).set_stride([S * H * D, D, H * D, 1]).set_data_type(o_dtype) + o.set_output(True).set_dim([B, H, S, D_v]).set_stride([S * H * D_v, D_v, H * D_v, 1]).set_data_type(o_dtype) stats.set_output(True).set_dim([B, H, S, 1]).set_stride([H * S, S, 1, 1]).set_data_type(cudnn.data_type.FLOAT) for t in (amx_s, amx_o): t.set_output(True).set_dim([1, 1, 1, 1]).set_stride([1, 1, 1, 1]).set_data_type(cudnn.data_type.FLOAT) @@ -298,7 +326,7 @@ def _fp8_graph_offers_sm120(io_dtype, o_dtype, D=128, sink=False): # The op itself may refuse the shape before any engine is consulted; # for "this cell must not claim it" that is the same answer. return False - return engine_name(arch="sm120", fp8=True) in [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + return offers_engine(g, engine_name(arch="sm120", fp8=True)) @pytest.mark.L0 @@ -321,15 +349,197 @@ def test_fp8_sm120_fp8_output_not_offered(): @pytest.mark.L0 -@pytest.mark.parametrize("D", [64, 256]) +@pytest.mark.parametrize("D", [16, 48, 144, 288]) @torch_fork_set_rng(seed=0) -def test_fp8_sm120_non_128_head_dim_not_offered(D): - """The 8-bit fragment path has no zero-padding envelope, so d is exact.""" +def test_fp8_sm120_off_granule_head_dim_not_offered(D): + """No zero-padding envelope on the 8-bit fragment path: dims are exact + multiples of 32 (k32 contraction; 1-byte TMA swizzle span). 16-odd + multiples of 16 and out-of-range dims decline.""" import cudnn assert not _fp8_graph_offers_sm120(cudnn.data_type.FP8_E4M3, cudnn.data_type.HALF, D=D) +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_head_dim_domain_offered(): + """Every graph the C++ front door admits is served: the d_qk <= 128 x + d_v <= 128 cross (multiples of 32) plus the (192, 128) MLA pair.""" + import cudnn + + for D in range(32, 129, 32): + for D_v in range(32, 129, 32): + assert _fp8_graph_offers_sm120(cudnn.data_type.FP8_E4M3, cudnn.data_type.HALF, D=D, D_v=D_v), f"({D}, {D_v}) not offered" + assert _fp8_graph_offers_sm120(cudnn.data_type.FP8_E4M3, cudnn.data_type.HALF, D=192, D_v=128), "(192, 128) MLA not offered" + + +@pytest.mark.L0 +@pytest.mark.parametrize("D", [32, 64, 96], ids=lambda d: f"d{d}") +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_head_dims(D): + """Correctness across the widened head-dim domain (exact, no padding). + + The graph front door (the C++ sdpa_fp8 node) admits d_qk <= 128 x + d_v <= 128 plus the (192, 128) MLA pair, so >128 uniform dims cannot + reach any engine; the kernel itself serves multiples of 32 up to 256.""" + scale = 1.0 / math.sqrt(D) + res = _run(2, 4, 4, 256, 256, scale=scale, sdpa_kwargs=dict(use_causal_mask=True), D=D) + _check(*res) + + +@pytest.mark.L1 +@pytest.mark.parametrize("D", [32, 96], ids=lambda d: f"d{d}") +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_head_dims_no_mask(D): + scale = 1.0 / math.sqrt(D) + res = _run(1, 4, 4, 256, 256, scale=scale, sdpa_kwargs={}, D=D) + _check(*res) + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "D,D_v,mask", + [ + (192, 128, "causal"), + (192, 128, "causal_br"), + (192, 128, "swa"), + (192, 128, "padded"), + (64, 128, "none"), + (128, 32, "causal"), + ], +) +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_mixed_head_dims(D, D_v, mask): + """The QK^T and P@V sides are independent in the kernel; the 192/128 + shape is the MLA population the f16 cell already serves (the front + door's only >128 carve-out), so it is crossed with every mask family.""" + scale = 1.0 / math.sqrt(D) + seq_lens_kv = [256, 192] if mask == "padded" else None + kw = {} if mask == "padded" else _MASKS[mask] + res = _run(2, 4, 4, 256, 256, scale=scale, sdpa_kwargs=kw, seq_lens_kv=seq_lens_kv, D=D, D_v=D_v) + _check(*res) + + +@pytest.mark.L0 +@pytest.mark.parametrize("mask", ["causal_br", "swa", "padded"]) +@pytest.mark.parametrize("D", [32, 96], ids=lambda d: f"d{d}") +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_head_dim_mask_cross(D, mask): + """Non-d128 head dims crossed with the mask families (the base head-dims + test covers plain causal). Masking and the head-dim envelope are + independent axes in the kernel; a regression coupling them shows here.""" + scale = 1.0 / math.sqrt(D) + seq_lens_kv = [256, 192] if mask == "padded" else None + kw = {} if mask == "padded" else _MASKS[mask] + res = _run(2, 4, 4, 256, 256, scale=scale, sdpa_kwargs=kw, seq_lens_kv=seq_lens_kv, D=D) + _check(*res) + + +def _run_template_tail(D, D_v, *, mask, S=256): + """Compile and launch the fp8 kernel template directly (production loader + and adapter ABI) for head dims the graph front door cannot reach. + + The engine row declares the kernel's full domain — multiples of 32 up to + 256, QK^T/P@V sides independent — but the C++ sdpa_fp8 node admits only + d_qk <= 128 x d_v <= 128 plus (192, 128) today, so the >128 tail is + protected here at the template level. + + scale_s is 1.0 (P is cast to e4m3 unscaled and the reference does not + model that cast), so the error floor is the bare P-quantization step + (~0.055 measured on SM120); 0.15 separates it cleanly from real + corruption (a dropped 32-column group or swizzle fault lands > 1). + """ + import os + + import cutlass + import cuda.bindings.driver as cuda_driver + + from cudnn.frost.template_loader import load_template + from cudnn.frost.tile_dsl.constants import DTYPE_E4M3 + from cudnn.sdpa.fwd import api_dsl + from cudnn.sdpa.fwd.config_sm120 import TemplateParams + + B, H = (2, 2) if mask == "padded" else (1, 2) + kw = {"dtype_qkv": DTYPE_E4M3} + swa_window = None + seq_kv_lens = None + if mask in ("causal", "swa"): + kw["window_right"] = 0 + if mask == "swa": + swa_window = 64 # kernel window_left=W keeps kv in [q-W, q]; same W as _ref's swa_window + kw["window_left"] = swa_window + if mask == "padded": + kw["seq_kv_lens_present"] = True + seq_kv_lens = [S, S - 73] # batch 1 ends inside a KV tile at an odd offset + path = os.path.join(os.path.dirname(os.path.abspath(api_dsl.__file__)), "kernels", "prefill_fp8_sm120.py") + module = load_template(path, TemplateParams(**kw), tag=f"fp8_tail_d{D}_d{D_v}_{mask}") + fn = module.compile(compute_capability=torch.cuda.get_device_capability(), b=B, qh=H, kh=H, sq=S, skv=S, d_qk=D, d_v=D_v, has_lse=False) + + dev = "cuda" + + def mk(*shape): + return (torch.randn(*shape, device=dev) * 0.5).clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn) + + q8, k8, v8 = mk(B, S, H, D), mk(B, S, H, D), mk(B, S, H, D_v) # compact BSHD, the kernel contract + o = torch.zeros(B, S, H, D_v, device=dev, dtype=torch.float16) + seq_q = torch.full((B,), S, dtype=torch.int32, device=dev) + seq_kv = torch.tensor(seq_kv_lens, dtype=torch.int32, device=dev) if seq_kv_lens else seq_q + amax_s = torch.zeros(1, dtype=torch.float32, device=dev) + amax_o = torch.zeros(1, dtype=torch.float32, device=dev) + scale = 1.0 / math.sqrt(D) + fn( + q8.view(torch.uint8), # the adapter's ABI: e4m3 as uint8 storage + k8.view(torch.uint8), + v8.view(torch.uint8), + o, + None, # lse (has_lse=False) + None, # sinks (unsupported; ABI slot) + seq_q, + seq_kv, + amax_s.view(torch.int32), # bitcast-int32 atomicMax storage + amax_o.view(torch.int32), + cutlass.Float32(scale * math.log2(math.e)), # softmax_scale_log2 (descale_q = descale_k = 1) + cutlass.Float32(1.0), # o_scale_fused = descale_s * descale_v * scale_o, all 1.0 here + cutlass.Float32(1.0), # scale_s + cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream), + ) + torch.cuda.synchronize() + + def bhsd(t): + return t.float().permute(0, 2, 1, 3) + + o_ref, _, _ = _ref(bhsd(q8), bhsd(k8), bhsd(v8), scale=scale, is_causal=mask in ("causal", "swa"), swa_window=swa_window, seq_lens_kv=seq_kv_lens) + of = bhsd(o) + assert not torch.isnan(of).any(), "NaN in O" + diff = (of - o_ref).abs().max().item() + assert diff <= 0.15, f"max|O-ref|={diff:.4f} > 0.15 (P-quant floor ~0.055)" + + +@pytest.mark.L1 +@pytest.mark.parametrize( + "D,D_v,mask", + [ + (160, 160, "causal"), + (160, 160, "swa"), + (160, 160, "none"), + (256, 256, "causal"), + (256, 256, "none"), + (256, 256, "swa"), + (256, 256, "padded"), + (256, 128, "causal"), + (128, 256, "causal"), + (224, 160, "causal"), + ], +) +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_head_dim_tail_direct(D, D_v, mask): + """Correctness of the front-door-unreachable >128 head-dim tail (see + _run_template_tail): uncommon uniform dims (160), the 256 domain max + crossed with every mask family, 256 on each side alone (the QK^T and + P@V swizzle/fragment paths size independently), and a mixed >128 pair.""" + _run_template_tail(D, D_v, mask=mask) + + @pytest.mark.L0 @pytest.mark.parametrize("tiles", [(64, 64), (64, 128), (128, 64), (128, 128)]) @torch_fork_set_rng(seed=0)