diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index aec5e345d..78ecf7bbe 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -331,9 +331,11 @@ def __init__( self.seq_q_lens_present = bool(seq_q_lens_present) # cu_seq_len form (cuDNN 9.24+): the corresponding seq-lens execute # argument arrives as a (B+1,)-int32 PREFIX-SUM tensor instead of - # (B,) per-batch lengths. THD-only today: the ragged lowering derives - # both forms host-side from its inherent tolist round-trip; the dense - # kernels have no CU read mode yet (check_support rejects). + # (B,) per-batch lengths. THD derives both forms host-side from its + # inherent tolist round-trip; dense graphs bind the (B+1,) tensor + # directly and the f16 kernels read len = cu[b+1] - cu[b] on device + # (the sync-free dense hot path stays sync-free). The FP8/MXFP8 + # kernels are not plumbed (check_support rejects). self.cu_seq_q_lens = bool(cu_seq_q_lens) self.cu_seq_kv_lens = bool(cu_seq_kv_lens) self.has_sink = bool(has_sink) @@ -856,10 +858,24 @@ def check_support(self) -> bool: ) if self.thd: self.seq_kv_lens_present = True + # cu_seq_len form: THD consumes either form host-side; dense graphs + # use the f16 kernels' CU read mode (len = cu[b+1] - cu[b]). The + # FP8/MXFP8 execute paths are not plumbed for the (B+1,) form. self._not_implemented_error_if( - (self.cu_seq_q_lens or self.cu_seq_kv_lens) and not self.thd, - "cu_seq_len_* is THD-only (the dense kernels have no CU read mode yet)", + (self.cu_seq_q_lens or self.cu_seq_kv_lens) and self._fp8 and not self.thd, + "dense cu_seq_len_* is not plumbed for the FP8/MXFP8 kernels", ) + if not self.thd: + # The cu flags declare the FORM of the corresponding lengths + # argument; the *_present flags still declare its PRESENCE. + self._value_error_if( + self.cu_seq_kv_lens and not self.seq_kv_lens_present, + "cu_seq_kv_lens declares the form of the KV lengths and requires seq_kv_lens_present=True", + ) + self._value_error_if( + self.cu_seq_q_lens and not self.seq_q_lens_present, + "cu_seq_q_lens declares the form of the Q lengths and requires seq_q_lens_present=True", + ) # Dense padded-Q trim backstops (engines.lower_dsl_prefill never sets # these combinations; a direct caller could). self._value_error_if( @@ -941,6 +957,10 @@ def compile(self) -> None: has_sink=self.has_sink, seq_kv_lens_present=self.seq_kv_lens_present, seq_q_lens_present=self.seq_q_lens_present, + # Dense-only kernel CU read mode; the THD cu form is consumed + # host-side and compiles to the same THD specialization. + seq_kv_lens_cu=self.cu_seq_kv_lens and not self.thd, + seq_q_lens_cu=self.cu_seq_q_lens and not self.thd, sched_policy=sched_policy, thd_varlen=self.thd, fused_ldtm_stat=fused_ldtm_stat, @@ -1163,17 +1183,22 @@ def execute( else self._dummy("sinks", device, lambda: torch.zeros(self.h_q, dtype=torch.float32, device=device)) ) seq_kv_t = ( - self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") + (self._checked_cu_seq_lens(seq_kv_lens, "cu_seq_len_kv") if self.cu_seq_kv_lens else self._checked_seq_lens(seq_kv_lens, "seq_kv_lens")) if seq_kv_lens is not None else self._dummy("seq_kv", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device)) ) # Dense padded-Q trim: per-batch Q lengths are their OWN kernel # parameter (compiled in only when seq_q_lens_present — the kernel # signature is specialized on `None`, so the flag-off ABI is - # unchanged). The caller's (B,)-int32 device tensor is bound directly + # unchanged). The caller's device tensor — (B,)-int32 lengths or the + # (B+1,)-int32 cu prefix sums under cu_seq_q_lens — is bound directly # as a validated view — zero allocations/copies on the execute hot # path, stable pointer (CUDA-graph-capture friendly). - seq_q_t = self._checked_seq_lens(seq_q_lens, "seq_q_lens") if self.seq_q_lens_present else None + seq_q_t = ( + (self._checked_cu_seq_lens(seq_q_lens, "cu_seq_len_q") if self.cu_seq_q_lens else self._checked_seq_lens(seq_q_lens, "seq_q_lens")) + if self.seq_q_lens_present + else None + ) o_desc_dummy = self._dummy("o_desc", device, lambda: torch.zeros(1, dtype=torch.int64, device=device)) import cutlass @@ -1770,10 +1795,19 @@ def check_support(self) -> bool: if self.thd: self._value_error_if(self.seq_q_lens_present, "seq_q_lens_present is dense-only (THD carries per-sequence Q lengths via cu_seqlens)") self.seq_kv_lens_present = True - self._not_implemented_error_if( - (self.cu_seq_q_lens or self.cu_seq_kv_lens) and not self.thd, - "cu_seq_len_* is THD-only (the dense kernels have no CU read mode yet)", - ) + # cu_seq_len form: THD consumes either form host-side; dense graphs + # use the f16 kernel's CU read mode (len = cu[b+1] - cu[b]). The cu + # flags declare the FORM of the corresponding lengths argument; the + # *_present flags still declare its PRESENCE. + if not self.thd: + self._value_error_if( + self.cu_seq_kv_lens and not self.seq_kv_lens_present, + "cu_seq_kv_lens declares the form of the KV lengths and requires seq_kv_lens_present=True", + ) + self._value_error_if( + self.cu_seq_q_lens and not self.seq_q_lens_present, + "cu_seq_q_lens declares the form of the Q lengths and requires seq_q_lens_present=True", + ) self._value_error_if( self.sched_policy is not None and self.sched_policy != SCHED_NATURAL, f"SM120 DSL SDPA only supports sched_policy={SCHED_NATURAL}", @@ -1896,6 +1930,12 @@ 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") + # The fp8 kernel has no dense CU read mode (its template rejects + # the cu flags); THD cu is consumed host-side and stays served. + self._not_implemented_error_if( + (self.cu_seq_q_lens or self.cu_seq_kv_lens) and not self.thd, + "cu_seq_len_* is not plumbed for the SM120 fp8 kernel", + ) self._value_error_if( 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 " @@ -1998,6 +2038,10 @@ def compile(self) -> None: bottom_right=self.causal_bottom_right, seq_q_lens_present=self.seq_q_lens_present, seq_kv_lens_present=self.seq_kv_lens_present, + # Dense-only kernel CU read mode; the THD cu form is consumed + # host-side and compiles to the same THD specialization. + seq_q_lens_cu=self.cu_seq_q_lens and not self.thd, + seq_kv_lens_cu=self.cu_seq_kv_lens and not self.thd, has_sink=self.has_sink, thd_varlen=self.thd, q_tile=self.q_tile, @@ -2117,7 +2161,7 @@ def execute( lse = self._checked_lse_view(lse_tensor) if lse_tensor is not None else None sinks_t = self._checked_sinks_1d(sinks) if sinks is not None else None seq_q_lens = ( - self._checked_seq_lens(seq_q_lens, "seq_q_lens") + (self._checked_cu_seq_lens(seq_q_lens, "cu_seq_len_q") if self.cu_seq_q_lens else self._checked_seq_lens(seq_q_lens, "seq_q_lens")) if seq_q_lens is not None else self._dummy( "seq_q_lens", @@ -2126,7 +2170,7 @@ def execute( ) ) seq_kv_lens = ( - self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") + (self._checked_cu_seq_lens(seq_kv_lens, "cu_seq_len_kv") if self.cu_seq_kv_lens else self._checked_seq_lens(seq_kv_lens, "seq_kv_lens")) if seq_kv_lens is not None else self._dummy( "seq_kv_lens", diff --git a/python/cudnn/sdpa/fwd/config_sm100.py b/python/cudnn/sdpa/fwd/config_sm100.py index 1dedb6c5e..cb0e211d1 100644 --- a/python/cudnn/sdpa/fwd/config_sm100.py +++ b/python/cudnn/sdpa/fwd/config_sm100.py @@ -85,6 +85,13 @@ class TemplateParams: # convention). Dense-only — THD carries per-sequence Q lengths via # cu_seqlens instead. seq_q_lens_present: bool = False + # cu_seq_len form (cuDNN 9.24+): the corresponding seq-lens kernel + # parameter is the (B+1,)-int32 PREFIX-SUM tensor instead of (B,) + # per-batch lengths; the kernels read len = cu[b+1] - cu[b] on device (the + # sync-free dense hot path stays sync-free). Dense-only — THD already + # carries prefix sums in its packed [kv_lens | cu_q | cu_kv] metadata. + seq_kv_lens_cu: bool = False + seq_q_lens_cu: bool = False sched_policy: int = SCHED_NATURAL thd_varlen: bool = False # cc10.3+ fuses the S_acc row-max into the LDTM (tcgen05.ld.red.f32.max); cc10.0 @@ -120,6 +127,13 @@ def _validate_params(flavor: str, k: TemplateParams) -> None: raise ValueError(f"{flavor}: SEQ_Q_LENS_PRESENT is dense-only (THD carries per-sequence Q lengths via cu_seqlens)") if not k.seq_kv_lens_present: raise ValueError(f"{flavor}: SEQ_Q_LENS_PRESENT requires SEQ_KV_LENS_PRESENT (padding mask)") + if k.seq_kv_lens_cu or k.seq_q_lens_cu: + if k.thd_varlen: + raise ValueError(f"{flavor}: seq_*_lens_cu is dense-only (the THD metadata buffer already carries cu_seqlens)") + if k.seq_kv_lens_cu and not k.seq_kv_lens_present: + raise ValueError(f"{flavor}: SEQ_KV_LENS_CU declares the FORM of the KV lengths and requires SEQ_KV_LENS_PRESENT") + if k.seq_q_lens_cu and not k.seq_q_lens_present: + raise ValueError(f"{flavor}: SEQ_Q_LENS_CU declares the FORM of the Q lengths and requires SEQ_Q_LENS_PRESENT") if k.sched_policy not in (SCHED_NATURAL, SCHED_LPT): raise ValueError(f"{flavor}: only SCHED_NATURAL (0) / SCHED_LPT (1) are wired up; got {k.sched_policy}") @@ -288,6 +302,10 @@ class CfgD256: SEQ_KV_LENS_PRESENT: int = 0 SEQ_Q_LENS_PRESENT: int = 0 + # cu_seq_len form: the corresponding lens parameter is the (B+1,) + # prefix-sum tensor; per-batch length = cu[b+1] - cu[b]. Dense-only. + SEQ_KV_LENS_CU: int = 0 + SEQ_Q_LENS_CU: int = 0 THD_VARLEN: int = 0 @@ -335,6 +353,8 @@ def make_cfg_d256(params: TemplateParams) -> Tuple[CfgD256, TmaIters]: SCHEDULER_POLICY=params.sched_policy, SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), + SEQ_KV_LENS_CU=int(params.seq_kv_lens_cu), + SEQ_Q_LENS_CU=int(params.seq_q_lens_cu), THD_VARLEN=int(params.thd_varlen), ) _validate_cfg_d256(cfg) @@ -427,6 +447,10 @@ class CfgD512: SEQ_KV_LENS_PRESENT: int = 0 SEQ_Q_LENS_PRESENT: int = 0 + # cu_seq_len form: the corresponding lens parameter is the (B+1,) + # prefix-sum tensor; per-batch length = cu[b+1] - cu[b]. Dense-only. + SEQ_KV_LENS_CU: int = 0 + SEQ_Q_LENS_CU: int = 0 THD_VARLEN: int = 0 @@ -478,6 +502,8 @@ def make_cfg_d512(params: TemplateParams) -> Tuple[CfgD512, TmaIters]: SCHEDULER_POLICY=params.sched_policy, SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), + SEQ_KV_LENS_CU=int(params.seq_kv_lens_cu), + SEQ_Q_LENS_CU=int(params.seq_q_lens_cu), THD_VARLEN=int(params.thd_varlen), ) _validate_cfg_d512(cfg) @@ -573,6 +599,10 @@ class CfgD128: SEQ_KV_LENS_PRESENT: int = 0 SEQ_Q_LENS_PRESENT: int = 0 + # cu_seq_len form: the corresponding lens parameter is the (B+1,) + # prefix-sum tensor; per-batch length = cu[b+1] - cu[b]. Dense-only. + SEQ_KV_LENS_CU: int = 0 + SEQ_Q_LENS_CU: int = 0 THD_VARLEN: int = 0 @@ -638,6 +668,8 @@ def make_cfg_d128(params: TemplateParams) -> Tuple[CfgD128, TmaIters]: SCHEDULER_POLICY=params.sched_policy, SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), + SEQ_KV_LENS_CU=int(params.seq_kv_lens_cu), + SEQ_Q_LENS_CU=int(params.seq_q_lens_cu), THD_VARLEN=int(params.thd_varlen), ) _validate_cfg_d128(cfg) @@ -709,6 +741,8 @@ def make_cfg_d192(params: TemplateParams) -> Tuple[CfgD192, TmaIters]: CORRECTION_REGS=40 if _mask_flags_from(params) == MASK_NONE else 88, SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), + SEQ_KV_LENS_CU=int(params.seq_kv_lens_cu), + SEQ_Q_LENS_CU=int(params.seq_q_lens_cu), THD_VARLEN=int(params.thd_varlen), ) _validate_cfg_d192(cfg) diff --git a/python/cudnn/sdpa/fwd/config_sm120.py b/python/cudnn/sdpa/fwd/config_sm120.py index 1277ea5a6..7e9b90116 100644 --- a/python/cudnn/sdpa/fwd/config_sm120.py +++ b/python/cudnn/sdpa/fwd/config_sm120.py @@ -61,6 +61,12 @@ class TemplateParams: bottom_right: bool = False seq_q_lens_present: bool = False seq_kv_lens_present: bool = False + # cu_seq_len form (cuDNN 9.24+): the corresponding seq-lens kernel + # argument is the (B+1,)-int32 PREFIX-SUM tensor instead of (B,) + # per-batch lengths; the kernel reads len = cu[b+1] - cu[b] on device. + # Dense-only — THD already carries prefix sums in its metadata tensor. + seq_q_lens_cu: bool = False + seq_kv_lens_cu: bool = False has_sink: bool = False thd_varlen: bool = False q_tile: int = SEQ_Q_TILES[0] @@ -71,6 +77,7 @@ def validate_params( params: TemplateParams, allowed_dtypes: tuple[int, ...] = (DTYPE_BF16, DTYPE_FP16), allow_right_band: bool = True, + allow_cu: bool = True, ) -> None: """Validate the SM120 template specialization. @@ -78,7 +85,8 @@ def validate_params( capabilities or adapter support checks; this validation is a backstop for direct template use. ``allowed_dtypes`` defaults to the FP16/BF16 template's set; the FP8 template passes its own. ``allow_right_band=False`` also - rejects a widened right band, which the FP8 template does not plumb. + rejects a widened right band, which the FP8 template does not plumb; + ``allow_cu=False`` likewise rejects the dense cu_seq_len read mode. """ if params.dtype_qkv not in allowed_dtypes: @@ -100,3 +108,12 @@ def validate_params( raise ValueError("SM120 SDPA: thd_varlen requires seq_kv_lens_present (the THD metadata tensor)") if params.seq_q_lens_present: raise ValueError("SM120 SDPA: seq_q_lens_present is dense-only (THD carries per-sequence Q lengths via cu_seqlens)") + if params.seq_q_lens_cu or params.seq_kv_lens_cu: + if not allow_cu: + raise ValueError("SM120 SDPA: the dense cu_seq_len read mode is not plumbed for this template") + if params.thd_varlen: + raise ValueError("SM120 SDPA: seq_*_lens_cu is dense-only (the THD metadata tensor already carries cu_seqlens)") + if params.seq_q_lens_cu and not params.seq_q_lens_present: + raise ValueError("SM120 SDPA: seq_q_lens_cu declares the FORM of the Q lengths and requires seq_q_lens_present") + if params.seq_kv_lens_cu and not params.seq_kv_lens_present: + raise ValueError("SM120 SDPA: seq_kv_lens_cu declares the FORM of the KV lengths and requires seq_kv_lens_present") diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 45b8698b5..22fc73d48 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -165,9 +165,15 @@ class Capabilities: thd: bool = False # cu_seq_len_q / cu_seq_len_kv (B+1,) prefix sums (cuDNN 9.24+). Serving # rows consume the form on THD host-side (lens = adjacent differences of - # the inherent tolist); dense cu graphs stay declined until the kernels - # grow a CU read mode (len = cu[b+1] - cu[b]) — see mismatch(). + # the inherent tolist); dense cu graphs additionally need the kernels' + # CU read mode — see dense_cu_seq_len and mismatch(). cu_seq_len: bool = False + # Dense graphs carrying the cu form: the row's kernels read + # len = cu[b+1] - cu[b] straight from the bound (B+1,) tensor, so the + # sync-free dense hot path stays sync-free (no host round-trip, no + # conversion kernel). Only meaningful with cu_seq_len; rows without it + # (the FP8/MXFP8 flavors, SM80) keep declining dense cu graphs. + dense_cu_seq_len: bool = False # Dense padded + stats needs the per-batch seq_len_q LSE trim (padded # q-rows write LSE=-inf / O=0, cuDNN >= 9.14). Plumbed for the half # kernels via SEQ_Q_LENS_PRESENT; the FP8/MXFP8 kernels lack the epilogue @@ -335,13 +341,12 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti if facts.has_cu_seq_len: # cu_seq_len_* ((B+1,) prefix sums, cuDNN 9.24+). The THD lowering - # consumes either length form host-side; the dense kernels' CU read - # mode (len = cu[b+1] - cu[b]) is not plumbed yet, so dense cu graphs - # stay declined even on serving rows. + # consumes either length form host-side; dense graphs need the row's + # kernel CU read mode (len = cu[b+1] - cu[b], dense_cu_seq_len). if not capabilities.cu_seq_len: return "graph uses cu_seq_len_q / cu_seq_len_kv, which this engine does not support" - if not facts.thd: - return "cu_seq_len_* on dense graphs is not supported yet (kernel CU read mode not plumbed)" + if not facts.thd and not capabilities.dense_cu_seq_len: + return "cu_seq_len_* on dense graphs is not supported by this engine (no kernel CU read mode)" if (facts.seq_q_t is not None and facts.cu_seq_q_t is not None) or (facts.seq_kv_t is not None and facts.cu_seq_kv_t is not None): return "seq_len_* and cu_seq_len_* on the same side is ambiguous (backend precedence is not replicated here)" @@ -352,7 +357,7 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti return "graph uses bottom-right causal, which this kernel does not support" if facts.window_left is not None and not capabilities.bottom_right_with_swa: return "bottom-right causal combined with a sliding window is not supported" - if facts.padded and not facts.thd and facts.seq_q_t is not None and not capabilities.bottom_right_padded_seq_q: + if facts.padded and not facts.thd and (facts.seq_q_t is not None or facts.cu_seq_q_t is not None) and not capabilities.bottom_right_padded_seq_q: return ( "bottom-right causal with a dense padding mask carrying per-batch seq_len_q is not " "supported (kernel anchors the BR diagonal at the global S_q, not seq_len_q[b])" @@ -411,6 +416,7 @@ def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: lse_optional=True, thd=True, cu_seq_len=True, + dense_cu_seq_len=True, padded_stats=True, # The f16/bf16 lowering serves any dense B/H/S stride permutation # (padded strides included) with the head dim innermost; the @@ -649,6 +655,7 @@ def _sm120_spec() -> EngineSpec: # padding and no padded-path cost. skv_tile=0, cu_seq_len=True, + dense_cu_seq_len=True, layouts=frozenset({"bshd", "dense_flex"}), sched_policies=frozenset({SCHED_NATURAL}), tile_ms=frozenset({64, 128}), @@ -716,7 +723,7 @@ def lower_dsl_prefill( # buffer the FP8/MXFP8 kernels can't honor (dense padded-Q trim is not # plumbed there — known gap) is dropped here rather than erroring at # execute. - seq_q_lens_present = facts.padded and not facts.thd and facts.seq_q_t is not None and not (facts.is_mxfp8 or facts.is_fp8) + seq_q_lens_present = facts.padded and not facts.thd and (facts.seq_q_t is not None or facts.cu_seq_q_t is not None) and not (facts.is_mxfp8 or facts.is_fp8) api = _adapter(api_type)( sample_q=ga.tensor_desc_from_ir(facts.q_t, name="q"), sample_k=ga.tensor_desc_from_ir(facts.k_t, name="k"), @@ -736,8 +743,10 @@ def lower_dsl_prefill( # THD carries Q lengths via cu_seqlens; the FP8/MXFP8 kernels are not # plumbed (their specs also keep padded_stats=False). seq_q_lens_present=seq_q_lens_present, - # cu_seq_len form (THD-only; the probe declined dense cu graphs): the - # adapter's seq-lens execute arguments carry (B+1,) prefix sums. + # cu_seq_len form: the adapter's seq-lens execute arguments carry + # (B+1,) prefix sums — consumed host-side for THD, bound directly for + # dense (the kernels' CU read mode; the probe admits dense cu graphs + # only on dense_cu_seq_len rows). cu_seq_q_lens=facts.cu_seq_q_t is not None, cu_seq_kv_lens=facts.cu_seq_kv_t is not None, has_sink=facts.has_sink, diff --git a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py index ff47dc4cd..2693f27cf 100644 --- a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py @@ -389,7 +389,12 @@ def _bounds_for_tile_qtrim(q_super_idx, seqlen_q, seqlen_kv, cta_in_pair, seq_q_ cga_base_super = q_super_idx - cta_in_pair q_row_coord = cga_base_super * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) arr = cutlass.make_array_view(seq_q_lens_tensor) - q_len_b = cutlass.Int32(arr[batch_idx]) + if cutlass.const_expr(int(getattr(CFG, "SEQ_Q_LENS_CU", 0)) == 1): + # cu_seq_len form (cuDNN 9.24+): the parameter is the (B+1,) + # prefix-sum tensor; the length is the adjacent difference. + q_len_b = cutlass.Int32(arr[batch_idx + cutlass.Int32(1)]) - cutlass.Int32(arr[batch_idx]) + else: + q_len_b = cutlass.Int32(arr[batch_idx]) tile_dead = q_row_coord >= q_len_b dead_lo = cutlass.Int32(arith.select(tile_dead.ir_value(), b.left.ir_value(), b.unmasked_lo.ir_value())) dead_hi = cutlass.Int32(arith.select(tile_dead.ir_value(), b.left.ir_value(), b.unmasked_hi.ir_value())) @@ -401,6 +406,12 @@ def _bounds_for_tile_qtrim(q_super_idx, seqlen_q, seqlen_kv, cta_in_pair, seq_q_ def _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, scalar_seqlen_kv): if cutlass.const_expr(CFG.SEQ_KV_LENS_PRESENT == 1): arr = cutlass.make_array_view(seq_kv_lens_tensor) + if cutlass.const_expr(int(getattr(CFG, "SEQ_KV_LENS_CU", 0)) == 1): + # cu_seq_len form (cuDNN 9.24+, dense-only — THD reads its + # packed metadata's kv-lens region instead): the parameter is + # the (B+1,) prefix-sum tensor; the per-batch KV length is the + # adjacent difference, read sync-free on device. + return cutlass.Int32(arr[batch_idx + cutlass.Int32(1)]) - cutlass.Int32(arr[batch_idx]) return cutlass.Int32(arr[batch_idx]) return scalar_seqlen_kv diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py index fd70b3cf3..b8148975e 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -1819,7 +1819,12 @@ def _correction_warp_group( # purpose — a trimmed row is dead even with a sink. Per-batch # q lens come in via the dedicated seq_q_lens_tensor parameter. _sq_arr = cutlass.make_array_view(seq_q_lens_tensor) - _q_len_b = cutlass.Int32(_sq_arr[batch_idx]) + if cutlass.const_expr(CFG.SEQ_Q_LENS_CU == 1): + # cu_seq_len form: (B+1,) prefix sums; the length is the + # adjacent difference, read sync-free on device. + _q_len_b = cutlass.Int32(_sq_arr[batch_idx + cutlass.Int32(1)]) - cutlass.Int32(_sq_arr[batch_idx]) + else: + _q_len_b = cutlass.Int32(_sq_arr[batch_idx]) row_trim = q_row_global >= _q_len_b neg_inf_trim = cutlass.Float32(float("-inf")) lse_val = cutlass.Float32(arith.select(row_trim.ir_value(), neg_inf_trim.ir_value(), lse_val.ir_value())) @@ -2149,7 +2154,8 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): # seq_kv_lens always part of the ABI; read only when CFG.SEQ_KV_LENS_PRESENT == 1 # (compile-time fold). THD overloads it as the [seq_kv_lens(B)|cu_q(B+1)| # cu_k(B+1)] metadata buffer (length 3B+2). - _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b + # Dense cu_seq_len form: the seq_kv parameter is the (B+1,) prefix sums. + _skv_len = (3 * b + 2) if CFG.THD_VARLEN else ((b + 1) if CFG.SEQ_KV_LENS_CU else b) fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (_skv_len,), @@ -2163,7 +2169,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_seq_q_lens = ( cute.runtime.make_fake_compact_tensor( cutlass.Int32, - (b,), + ((b + 1) if CFG.SEQ_Q_LENS_CU else b,), # cu form: (B+1,) prefix sums stride_order=(0,), assumed_align=4, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py index 8f2face2b..3707ae9b6 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py @@ -1907,7 +1907,12 @@ def _correction_warp_group( # purpose — a trimmed row is dead even with a sink. Per-batch # q lens come in via the dedicated seq_q_lens_tensor parameter. _sq_arr = cutlass.make_array_view(seq_q_lens_tensor) - _q_len_b = cutlass.Int32(_sq_arr[batch_idx]) + if cutlass.const_expr(CFG.SEQ_Q_LENS_CU == 1): + # cu_seq_len form: (B+1,) prefix sums; the length is the + # adjacent difference, read sync-free on device. + _q_len_b = cutlass.Int32(_sq_arr[batch_idx + cutlass.Int32(1)]) - cutlass.Int32(_sq_arr[batch_idx]) + else: + _q_len_b = cutlass.Int32(_sq_arr[batch_idx]) row_trim = q_row_global >= _q_len_b neg_inf_trim = cutlass.Float32(float("-inf")) lse_val = cutlass.Float32(arith.select(row_trim.ir_value(), neg_inf_trim.ir_value(), lse_val.ir_value())) @@ -2236,7 +2241,8 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): # seq_kv_lens always part of the ABI; read only when CFG.SEQ_KV_LENS_PRESENT == 1 # (compile-time fold). THD overloads it as the [seq_kv_lens(B)|cu_q(B+1)| # cu_k(B+1)] metadata buffer (length 3B+2). - _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b + # Dense cu_seq_len form: the seq_kv parameter is the (B+1,) prefix sums. + _skv_len = (3 * b + 2) if CFG.THD_VARLEN else ((b + 1) if CFG.SEQ_KV_LENS_CU else b) fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (_skv_len,), @@ -2250,7 +2256,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_seq_q_lens = ( cute.runtime.make_fake_compact_tensor( cutlass.Int32, - (b,), + ((b + 1) if CFG.SEQ_Q_LENS_CU else b,), # cu form: (B+1,) prefix sums stride_order=(0,), assumed_align=4, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py index 36a4295e5..79f88e50f 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -1492,7 +1492,12 @@ def _correction_warp_group( # a trimmed row is dead even with a sink. Per-batch q lens come in # via the dedicated seq_q_lens_tensor parameter. _sq_arr = cutlass.make_array_view(seq_q_lens_tensor) - _q_len_b = cutlass.Int32(_sq_arr[batch_idx]) + if cutlass.const_expr(CFG.SEQ_Q_LENS_CU == 1): + # cu_seq_len form: (B+1,) prefix sums; the length is the + # adjacent difference, read sync-free on device. + _q_len_b = cutlass.Int32(_sq_arr[batch_idx + cutlass.Int32(1)]) - cutlass.Int32(_sq_arr[batch_idx]) + else: + _q_len_b = cutlass.Int32(_sq_arr[batch_idx]) row_trim = q_row_global >= _q_len_b neg_inf_trim = cutlass.Float32(float("-inf")) lse_val = cutlass.Float32(arith.select(row_trim.ir_value(), neg_inf_trim.ir_value(), lse_val.ir_value())) @@ -1782,7 +1787,8 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): stride_order=(0,), assumed_align=16, ) - _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b + # Dense cu_seq_len form: the seq_kv parameter is the (B+1,) prefix sums. + _skv_len = (3 * b + 2) if CFG.THD_VARLEN else ((b + 1) if CFG.SEQ_KV_LENS_CU else b) fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (_skv_len,), @@ -1796,7 +1802,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_seq_q_lens = ( cute.runtime.make_fake_compact_tensor( cutlass.Int32, - (b,), + ((b + 1) if CFG.SEQ_Q_LENS_CU else b,), # cu form: (B+1,) prefix sums stride_order=(0,), assumed_align=4, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py index 374eb9bf2..5bb675994 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -1106,7 +1106,12 @@ def _compute_warp_group( # so the trim recomputes the row index here — beta feeds the # O normalization in that loop.) _sq_arr = cutlass.make_array_view(seq_q_lens_tensor) - _q_len_b = cutlass.Int32(_sq_arr[batch_idx]) + if cutlass.const_expr(CFG.SEQ_Q_LENS_CU == 1): + # cu_seq_len form: (B+1,) prefix sums; the length is the + # adjacent difference, read sync-free on device. + _q_len_b = cutlass.Int32(_sq_arr[batch_idx + cutlass.Int32(1)]) - cutlass.Int32(_sq_arr[batch_idx]) + else: + _q_len_b = cutlass.Int32(_sq_arr[batch_idx]) _q_row_trim = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) + tid_in_wg row_trim = _q_row_trim >= _q_len_b neg_inf_trim = cutlass.Float32(float("-inf")) @@ -1984,7 +1989,8 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): stride_order=(0,), assumed_align=16, ) - _skv_len = (3 * b + 2) if CFG.THD_VARLEN else b + # Dense cu_seq_len form: the seq_kv parameter is the (B+1,) prefix sums. + _skv_len = (3 * b + 2) if CFG.THD_VARLEN else ((b + 1) if CFG.SEQ_KV_LENS_CU else b) fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (_skv_len,), @@ -1998,7 +2004,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_seq_q_lens = ( cute.runtime.make_fake_compact_tensor( cutlass.Int32, - (b,), + ((b + 1) if CFG.SEQ_Q_LENS_CU else b,), # cu form: (B+1,) prefix sums stride_order=(0,), assumed_align=4, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py index 72d9f3d34..744a29bb6 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py @@ -196,6 +196,8 @@ def __init__( window_size_right: int | None = None, seq_q_lens_present: bool = False, seq_kv_lens_present: bool = False, + seq_q_lens_cu: bool = False, + seq_kv_lens_cu: bool = False, has_sink: bool = False, thd_varlen: bool = False, thd_batch: int = 1, @@ -220,6 +222,11 @@ def __init__( this offset. :param seq_q_lens_present: Read per-batch query lengths at runtime. :param seq_kv_lens_present: Read per-batch key/value lengths at runtime. + :param seq_q_lens_cu: ``seq_q_lens`` is the ``(B+1,)`` cu_seq_len + prefix-sum form (cuDNN 9.24+); the kernel reads + ``len = cu[b+1] - cu[b]``. Dense-only; requires + ``seq_q_lens_present``. + :param seq_kv_lens_cu: Same for ``seq_kv_lens``. :param has_sink: Fold the per-Q-head sink logit from the ``sinks`` tensor into the softmax denominator; when ``False`` the ``sinks`` argument is an unused dummy. @@ -261,8 +268,16 @@ def __init__( # A translated diagonal (bottom-right anchoring or a right band) can # straddle one more KV tile than the tile-aligned top-left one. self.diag_shifted = bottom_right or self.right_slack > 0 + if (seq_q_lens_cu or seq_kv_lens_cu) and thd_varlen: + raise ValueError("seq_*_lens_cu is dense-only (the THD metadata tensor already carries cu_seqlens)") + if seq_q_lens_cu and not seq_q_lens_present: + raise ValueError("seq_q_lens_cu declares the FORM of the Q lengths and requires seq_q_lens_present") + if seq_kv_lens_cu and not seq_kv_lens_present: + raise ValueError("seq_kv_lens_cu declares the FORM of the KV lengths and requires seq_kv_lens_present") self.seq_q_lens_present = seq_q_lens_present self.seq_kv_lens_present = seq_kv_lens_present + self.seq_q_lens_cu = seq_q_lens_cu + self.seq_kv_lens_cu = seq_kv_lens_cu self.has_sink = has_sink self.thd_varlen = thd_varlen self.thd_batch = thd_batch @@ -847,15 +862,26 @@ def kernel( kv_row_base = cutlass.Int32(meta[2 * n_batch + 1 + batch_idx]) seqlen_k = cutlass.Int32(meta[2 * n_batch + 1 + batch_idx + 1]) - kv_row_base else: + # Dense per-batch lengths: (B,) lengths, or the (B+1,) cu_seq_len + # prefix sums (adjacent difference) under the *_cu specialization — + # read sync-free on device either way. if cutlass.const_expr(self.seq_q_lens_present): + if cutlass.const_expr(self.seq_q_lens_cu): + q_len_b = cutlass.Int32(seq_q_lens[batch_idx + 1]) - cutlass.Int32(seq_q_lens[batch_idx]) + else: + q_len_b = cutlass.Int32(seq_q_lens[batch_idx]) seqlen_q = cute.math.max( cutlass.Int32(0), - cute.math.min(seq_q_lens[batch_idx], cutlass.Int32(q.shape[1])), + cute.math.min(q_len_b, cutlass.Int32(q.shape[1])), ) if cutlass.const_expr(self.seq_kv_lens_present): + if cutlass.const_expr(self.seq_kv_lens_cu): + kv_len_b = cutlass.Int32(seq_kv_lens[batch_idx + 1]) - cutlass.Int32(seq_kv_lens[batch_idx]) + else: + kv_len_b = cutlass.Int32(seq_kv_lens[batch_idx]) seqlen_k = cute.math.max( cutlass.Int32(0), - cute.math.min(seq_kv_lens[batch_idx], cutlass.Int32(k.shape[1])), + cute.math.min(kv_len_b, cutlass.Int32(k.shape[1])), ) num_heads_q = q.shape[2] @@ -1361,6 +1387,11 @@ def __call__( raise ValueError("THD Q/K/V/O must be packed batch-1 views") if cutlass.const_expr(seq_kv_lens.shape != (3 * self.thd_batch + 2,)): raise ValueError("THD seq_kv_lens must be the (3*B+2,) metadata tensor") + else: + if cutlass.const_expr(self.seq_q_lens_cu and seq_q_lens.shape != (q.shape[0] + 1,)): + raise ValueError("cu-form seq_q_lens must be the (B+1,) prefix-sum tensor") + if cutlass.const_expr(self.seq_kv_lens_cu and seq_kv_lens.shape != (k.shape[0] + 1,)): + raise ValueError("cu-form seq_kv_lens must be the (B+1,) prefix-sum tensor") # Split D into I contiguous C-element chunks while preserving the # per-tensor TMA descriptor over the compact (B, S, H, D) storage. @@ -1476,6 +1507,8 @@ def compile( # noqa: A001 window_size_right=PARAMS.window_right, seq_q_lens_present=PARAMS.seq_q_lens_present, seq_kv_lens_present=PARAMS.seq_kv_lens_present, + seq_q_lens_cu=PARAMS.seq_q_lens_cu, + seq_kv_lens_cu=PARAMS.seq_kv_lens_cu, has_sink=PARAMS.has_sink, thd_varlen=PARAMS.thd_varlen, thd_batch=b, @@ -1521,15 +1554,16 @@ def _fake_bshd(shape, stride): if PARAMS.has_sink else None ) + # Dense cu form (seq_*_lens_cu): the argument is the (B+1,) prefix sums. fake_seq_q_lens = cute.runtime.make_fake_compact_tensor( cutlass.Int32, - (b,), + ((b + 1) if PARAMS.seq_q_lens_cu else b,), stride_order=(0,), assumed_align=4, ) fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor( cutlass.Int32, - (3 * b + 2,) if PARAMS.thd_varlen else (b,), # THD: [ seq_kv(B) | cu_q(B+1) | cu_k(B+1) ] + (3 * b + 2,) if PARAMS.thd_varlen else ((b + 1) if PARAMS.seq_kv_lens_cu else b,), # THD: [ seq_kv(B) | cu_q(B+1) | cu_k(B+1) ] stride_order=(0,), assumed_align=4, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py index 7d36e7ca5..47de6cbf7 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py @@ -71,7 +71,7 @@ # The FROST loader injects one immutable specialization before executing this # module. A direct import uses dense e4m3 defaults. PARAMS: TemplateParams = globals().get("FROST_TEMPLATE_PARAMS", TemplateParams(dtype_qkv=DTYPE_E4M3)) -validate_params(PARAMS, allowed_dtypes=(DTYPE_E4M3,), allow_right_band=False) +validate_params(PARAMS, allowed_dtypes=(DTYPE_E4M3,), allow_right_band=False, allow_cu=False) # e4m3 travels as raw bytes: TMA, ldmatrix, and the MMA consume bit patterns, # so Uint8 storage sidesteps Float8 element support in the DSL plumbing. diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py index 657898709..7934d720d 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -341,6 +341,102 @@ def test_dsl_sm100_padded(dtype, d): torch.testing.assert_close(o, o_ref, atol=5e-2, rtol=3e-2) +def _run_dense_cu_graph(q, k, v, *, scale, dtype, cu_kv, cu_q=None, check_stats=False): + """Dense padded graph carrying the cu_seq_len ((B+1,) prefix-sum, cuDNN + 9.24+) length form, executed through the matching FROST engine. The + kernels read len = cu[b+1] - cu[b] on device — no host round-trip, no + conversion kernel (the sync-free dense hot path stays sync-free).""" + import cudnn + + b, h_q, s_q, _ = q.shape + d_v = v.shape[-1] + o_gpu = torch.empty(b, s_q, h_q, d_v, device="cuda", dtype=dtype).transpose(1, 2) + io = cudnn.data_type.HALF if dtype == torch.float16 else cudnn.data_type.BFLOAT16 + g = cudnn.pygraph(io_data_type=io, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + tq, tk, tv = g.tensor_like(q), g.tensor_like(k), g.tensor_like(v) + t_cu_kv = g.tensor_like(cu_kv) + kw = dict(name="sdpa", q=tq, k=tk, v=tv, generate_stats=check_stats, attn_scale=scale, use_padding_mask=True, cu_seq_len_kv=t_cu_kv) + vp = {tq: q, tk: k, tv: v, t_cu_kv: cu_kv} + if cu_q is not None: + t_cu_q = g.tensor_like(cu_q) + kw["cu_seq_len_q"] = t_cu_q + vp[t_cu_q] = cu_q + o, stats = g.sdpa(**kw) + o.set_output(True).set_dim(o_gpu.shape).set_stride(o_gpu.stride()) + lse_gpu = None + if check_stats: + lse_gpu = torch.empty(b, h_q, s_q, 1, dtype=torch.float32, device="cuda") + stats.set_output(True).set_dim(lse_gpu.shape).set_stride(lse_gpu.stride()) + stats.set_data_type(cudnn.data_type.FLOAT) + vp[stats] = lse_gpu + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + _select_engine(g, engine_name(q.shape[-1], d_v=d_v)) + g.check_support() + g.build_plans() + vp[o] = o_gpu + g.execute(vp, torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8)) + torch.cuda.synchronize() + return o_gpu, lse_gpu + + +_CU_FLAVORS = [(512, 512), (256, 256), (192, 128), (128, 128)] +_CU_FLAVOR_IDS = ["dsv4_d512", "qwen_d256", "dsv3_d192_d128", "llama_d128"] + + +@pytest.mark.L0 +@pytest.mark.parametrize("d,d_v", _CU_FLAVORS, ids=_CU_FLAVOR_IDS) +@pytest.mark.parametrize("dtype", _DTYPES, ids=_DTYPE_IDS) +@torch_fork_set_rng(seed=0) +def test_dsl_sm100_padded_cu_seq_len(dtype, d, d_v): + """Dense padded mask carried in the cu_seq_len ((B+1,) prefix-sum) form: + the kernel's CU read mode recovers seq_len_kv[b] = cu[b+1] - cu[b]. The + full-length cu_seq_len_q companion also exercises the CU form of the + padded-Q trim read (a no-op at full lengths).""" + _require_dsl() + b, h, s = 2, 8, 256 + scale = 1.0 / math.sqrt(d) + q = _bhsd(b, h, s, d, dtype) + k = _bhsd(b, h, s, d, dtype) + v = _bhsd(b, h, s, d_v, dtype) + kv_lens = [180, 240] + cu_kv = torch.tensor([0, kv_lens[0], kv_lens[0] + kv_lens[1]], dtype=torch.int32, device="cuda").view(b + 1, 1, 1, 1) + cu_q = torch.tensor([0, s, 2 * s], dtype=torch.int32, device="cuda").view(b + 1, 1, 1, 1) + o, _ = _run_dense_cu_graph(q, k, v, scale=scale, dtype=dtype, cu_kv=cu_kv, cu_q=cu_q) + o_ref = _ref_sdpa_full(q, k, v, scale=scale, seq_kv_lens=torch.tensor(kv_lens, device="cuda")) + torch.testing.assert_close(o, o_ref, atol=5e-2, rtol=3e-2) + + +@pytest.mark.L0 +@pytest.mark.parametrize("d,d_v", _CU_FLAVORS, ids=_CU_FLAVOR_IDS) +@torch_fork_set_rng(seed=0) +def test_dsl_sm100_padded_cu_q_trim_stats(d, d_v): + """Dense padded-Q trim with BOTH lengths in the cu form + Stats: q rows >= + cu_q[b+1] - cu_q[b] must come back O := 0 / LSE := -inf (cuDNN >= 9.14), + live rows must match the reference — per f16 kernel flavor (each has its + own epilogue trim read).""" + _require_dsl() + dtype = torch.float16 + b, h, s = 2, 4, 256 + scale = 1.0 / math.sqrt(d) + q = _bhsd(b, h, s, d, dtype) + k = _bhsd(b, h, s, d, dtype) + v = _bhsd(b, h, s, d_v, dtype) + q_lens = [130, 256] + kv_lens = [180, 240] + cu_q = torch.tensor([0, q_lens[0], q_lens[0] + q_lens[1]], dtype=torch.int32, device="cuda").view(b + 1, 1, 1, 1) + cu_kv = torch.tensor([0, kv_lens[0], kv_lens[0] + kv_lens[1]], dtype=torch.int32, device="cuda").view(b + 1, 1, 1, 1) + o, lse = _run_dense_cu_graph(q, k, v, scale=scale, dtype=dtype, cu_kv=cu_kv, cu_q=cu_q, check_stats=True) + o_ref, lse_ref = _ref_sdpa_full(q, k, v, scale=scale, seq_kv_lens=torch.tensor(kv_lens, device="cuda"), return_stats=True) + lse = lse.view(b, h, s) + for bi, q_len in enumerate(q_lens): + torch.testing.assert_close(o[bi, :, :q_len], o_ref[bi, :, :q_len], atol=5e-2, rtol=3e-2) + torch.testing.assert_close(lse[bi, :, :q_len], lse_ref[bi, :, :q_len], atol=5e-2, rtol=3e-2) + assert (o[bi, :, q_len:] == 0).all(), f"batch {bi}: trimmed O rows must be zero" + assert (lse[bi, :, q_len:] == float("-inf")).all(), f"batch {bi}: trimmed LSE rows must be -inf" + + @pytest.mark.L0 @pytest.mark.parametrize("d", _FLAVORS, ids=_FLAVOR_IDS) @pytest.mark.parametrize("dtype", _DTYPES, ids=_DTYPE_IDS) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py index b6b38e0ca..4cbcbfbab 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -165,6 +165,7 @@ def _run_case( window_size_right: int | None = None, seq_q_lens: torch.Tensor | None = None, seq_kv_lens: torch.Tensor | None = None, + cu_lens: bool = False, scale: float | None = None, with_sink: bool = False, check_stats: bool = False, @@ -191,6 +192,7 @@ def _run_case( kv_tile=kv_tile, scale=scale, return_stats=check_stats, + cu_lens=cu_lens, **mask_kwargs, ) if check_stats: @@ -237,6 +239,7 @@ def _run_dsl_graph( window_size_right: int | None = None, seq_q_lens: torch.Tensor | None = None, seq_kv_lens: torch.Tensor | None = None, + cu_lens: bool = False, sinks: torch.Tensor | None = None, q_tile: int | None = None, kv_tile: int | None = None, @@ -289,14 +292,30 @@ def _run_dsl_graph( ) if seq_q_lens is not None or seq_kv_lens is not None: assert seq_q_lens is not None and seq_kv_lens is not None - seq_q = graph.tensor_like(seq_q_lens, name="seq_q") - seq_kv = graph.tensor_like(seq_kv_lens, name="seq_kv") - sdpa_kwargs.update( - use_padding_mask=True, - seq_len_q=seq_q, - seq_len_kv=seq_kv, - ) - variant_pack.update({seq_q: seq_q_lens, seq_kv: seq_kv_lens}) + if cu_lens: + # Dense cu_seq_len form (cuDNN 9.24+): bind the (B+1,) prefix + # sums; the kernel reads len = cu[b+1] - cu[b] on device. + def _cu(lens): + return torch.cat([torch.zeros(1, dtype=torch.int32, device=lens.device), lens.reshape(-1).cumsum(0).to(torch.int32)]).view(-1, 1, 1, 1) + + cuq_t, cukv_t = _cu(seq_q_lens), _cu(seq_kv_lens) + seq_q = graph.tensor_like(cuq_t, name="cu_seq_q") + seq_kv = graph.tensor_like(cukv_t, name="cu_seq_kv") + sdpa_kwargs.update( + use_padding_mask=True, + cu_seq_len_q=seq_q, + cu_seq_len_kv=seq_kv, + ) + variant_pack.update({seq_q: cuq_t, seq_kv: cukv_t}) + else: + seq_q = graph.tensor_like(seq_q_lens, name="seq_q") + seq_kv = graph.tensor_like(seq_kv_lens, name="seq_kv") + sdpa_kwargs.update( + use_padding_mask=True, + seq_len_q=seq_q, + seq_len_kv=seq_kv, + ) + variant_pack.update({seq_q: seq_q_lens, seq_kv: seq_kv_lens}) if sinks is not None: sink_t = graph.tensor_like(sinks, name="sink") sdpa_kwargs["sink_token"] = sink_t @@ -715,6 +734,48 @@ def test_dsl_sm120_stats_padded_trim(): ) +@pytest.mark.L0 +@torch_fork_set_rng(seed=14) +def test_dsl_sm120_padded_cu_seq_len(): + """Dense padded mask carried in the cu_seq_len ((B+1,) prefix-sum) form: + the kernel's CU read mode recovers len = cu[b+1] - cu[b] on device — no + host round-trip, no conversion kernel.""" + + seq_q_lens = torch.tensor([256, 256], dtype=torch.int32, device="cuda") + seq_kv_lens = torch.tensor([180, 240], dtype=torch.int32, device="cuda") + _run_case( + batch=2, + h_q=4, + h_kv=2, + s_q=256, + s_kv=256, + head_dim=128, + seq_q_lens=seq_q_lens, + seq_kv_lens=seq_kv_lens, + cu_lens=True, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=14) +def test_dsl_sm120_stats_padded_trim_cu_seq_len(): + """The padded-Q trim (LSE = -inf past seq_len_q[b]) with BOTH lengths in + the cu form — the dense Q-length CU read feeds the same trim path.""" + + seq_q_lens = torch.tensor([96, 128], dtype=torch.int32, device="cuda") + seq_kv_lens = torch.tensor([0, 64], dtype=torch.int32, device="cuda") + _run_case( + batch=2, + s_q=128, + s_kv=128, + head_dim=64, + seq_q_lens=seq_q_lens, + seq_kv_lens=seq_kv_lens, + cu_lens=True, + check_stats=True, + ) + + @pytest.mark.L0 @torch_fork_set_rng(seed=18) def test_dsl_sm120_sink(): diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index ac9b509ea..d14c7c235 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -816,15 +816,73 @@ def test_probe_rejects_thd_cu_plus_seq_len(): assert not _eligible(_mk_thd_cu_graph(extra_seq_len=True)) +def _mk_dense_cu_graph(*, d=128, cu_q=True, cu_kv=True, seq_q=False, seq_kv=False, bottom_right=False): + """Dense padded graph carrying the cu_seq_len (B+1,) prefix-sum form on + the selected sides (per-batch (B,) seq_len_* on the others).""" + g = _mk_graph() + q, k, v, dims, strides = _mk_qkv(g, d=d) + kw = {} + if cu_q: + kw["cu_seq_len_q"] = g.tensor(dim=(B + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="cu_q") + if cu_kv: + kw["cu_seq_len_kv"] = g.tensor(dim=(B + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="cu_kv") + if seq_q: + kw["seq_len_q"] = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="seq_q") + if seq_kv: + kw["seq_len_kv"] = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="seq_kv") + if bottom_right: + kw["use_causal_mask_bottom_right"] = True + o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True, use_padding_mask=True, **kw) + _finish_output(o, dims, strides) + return g + + +def test_probe_accepts_dense_cu_seq_len(): + """Dense padded graphs carrying the cu_seq_len (B+1,) prefix-sum form are + served by the SM100 f16 rows via the kernels' CU read mode + (len = cu[b+1] - cu[b], read sync-free on device); the FP8/MXFP8 rows have + no CU read mode and keep declining.""" + g = _mk_dense_cu_graph() + names = _eligible(g) + assert engines.engine_name(128) in names + assert engines.engine_name(128, mxfp8=True) not in names + assert engines.engine_name(128, fp8=True) not in names + facts = _facts(g) + assert facts.has_cu_seq_len and facts.padded and not facts.thd + + +def test_probe_accepts_dense_cu_kv_only(): + """cu form on the KV side only (per-batch (B,) lengths on Q) — the forms + are declared per side and may be mixed across sides.""" + assert engines.engine_name(128) in _eligible(_mk_dense_cu_graph(cu_q=False, seq_q=True)) + + +def test_probe_rejects_dense_cu_bottom_right_q_lens(): + """Dense BR + per-batch Q lengths hits the same kernel gap regardless of + the length FORM: the SM100 kernels anchor the BR diagonal at the global + S_q, so cu-form Q lengths must be declined exactly like seq_len_q.""" + assert not _eligible(_mk_dense_cu_graph(bottom_right=True)) + + +def test_sm120_probe_accepts_dense_cu_seq_len(monkeypatch): + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + g = _mk_dense_cu_graph() + names = _eligible(g) + assert engines.engine_name(arch="sm120") in names + assert "sdpa_fwd_prefill_sm120_fp8" not in names + # SM120's kernel anchors the BR diagonal per batch, so dense cu-form Q + # lengths stay served with bottom-right (bottom_right_padded_seq_q). + assert engines.engine_name(arch="sm120") in _eligible(_mk_dense_cu_graph(bottom_right=True)) + + @pytest.mark.parametrize("side", ["cu_seq_len_q", "cu_seq_len_kv"]) -def test_cu_seq_len_is_declined(side): - """cu_seq_len_* (cuDNN 9.24+) are prefix sums — a different contract from - seq_len_* and from ragged_offset, and these kernels implement neither. - Reading such a graph as plain padded silently produced wrong output: 14.9% - of O on test_sdpa_mixed_seq_len_forms_L0[cu_q_brcm]. +def test_cu_plus_seq_len_same_side_is_declined(side): + """Both length forms on ONE side is ambiguous (the backend has its own + precedence, which the python engines do not replicate) — declined even on + rows that serve the dense cu form. A FACT, not a verdict: ``invalid`` means malformed-for-everyone, so putting - this there would also bar the engine that eventually implements it.""" + this there would also bar an engine that replicates the precedence.""" g = _mk_graph() q, k, v, dims, strides = _mk_qkv(g, d=128) seq_kv = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="seq_kv") @@ -846,7 +904,7 @@ def test_cu_seq_len_is_declined(side): facts = ga.analyze(g) assert facts.invalid is None, facts.invalid assert facts.has_cu_seq_len - assert not _eligible(g), "no engine may claim a graph carrying cu_seq_len" + assert not _eligible(g), "no engine may claim a graph with both length forms on one side" def test_sm120_probe_accepts_padding_mask_with_seq_lens(monkeypatch):