diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 2f090ef14..304155a96 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -2812,7 +2812,6 @@ def sdpa_fwd_wrapper_dsl_sm120( # ``lower_dsl_prefill`` drives all three identically; converting the kernels # themselves to TemplateParams form is tracked as follow-up work. # ============================================================================= -import inspect as _sm80_inspect from cudnn.sdpa.fwd import config_sm80 as _sm80_config @@ -2845,22 +2844,6 @@ def sdpa_fwd_wrapper_dsl_sm120( # all others use the shared generic kernel. _SM80_D256_FLAVORS = ("qwen",) -_SM80_KERNEL_MOD: dict = {} - - -def _sm80_kernel_mod(flavor: str = ""): - """Lazily import + cache the SM80 kernel module for ``flavor``. qwen - (d=256) routes to ``prefill_d256_f16_sm80`` (symmetric K+V prefetch); the - rest use ``prefill_f16_sm80``.""" - key = "d256" if flavor in _SM80_D256_FLAVORS else "f16" - if key not in _SM80_KERNEL_MOD: - if key == "d256": - from .kernels import prefill_d256_f16_sm80 as _mod - else: - from .kernels import prefill_f16_sm80 as _mod - _SM80_KERNEL_MOD[key] = _mod - return _SM80_KERNEL_MOD[key] - def _sm80_pick_flavor(d_qk: int, d_v: int) -> str: """Smallest kernel flavor whose ``(D_QK, D_V)`` envelope covers @@ -2924,35 +2907,127 @@ def _sm80_resolve_scheduler( raise ValueError(f"SM80 SDPA: scheduler must be 'auto' / 'default' / 'natural' / 'lpt' / 'lpt_l2', got {scheduler!r}") +# --- SM80 template loading --------------------------------------------------- + +_LOG2E = math.log2(math.e) + +_SM80_KERNEL_FILES = { + "d256": "prefill_d256_f16_sm80.py", + "f16": "prefill_f16_sm80.py", +} + + +def _sm80_load_kernel_module(flavor: str, params): + """One uniquely-named module per (kernel file, TemplateParams) — the same + ``frost.template_loader`` mechanism the SM100/SM120 templates use. qwen + (d=256) routes to the symmetric-K+V-prefetch file; the rest share the + generic kernel.""" + filename = _SM80_KERNEL_FILES["d256" if flavor in _SM80_D256_FLAVORS else "f16"] + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels", filename) + return load_template(path, params, tag=f"sm80_{filename.rsplit('.', 1)[0]}") + + +def _sm80_sched_policy_int(token: str) -> int: + return {"default": SCHED_NATURAL, "lpt": SCHED_LPT, "lpt_l2": SCHED_LPT_L2}[token] + + +def _sm80_call( + fn, + *, + q, + k, + v, + o, + lse, + seq_kv, + seq_q, + sinks_log2, + bias, + cu_q, + cu_k, + rope_cs, + n_kv_tiles, + scale_log2, + sq, + skv, + d, + right_bound, + inv_scale, + thd_q_tiles, + n_batch_logical, + stream, +): + """Invoke one compiled SM80 artifact (the traced ``_sdpa_host`` ABI: + 12 tensors — LSE may be None-specialized — then 9 runtime scalars and the + launch stream).""" + import cutlass + from cutlass.cute.runtime import from_dlpack as _from_dlpack_raw + + def from_dlpack(t): + # The kernels compile with --enable-tvm-ffi, so host-side conversions + # must produce TVM-FFI tensors regardless of the env latch. + return _from_dlpack_raw(t, enable_tvm_ffi=True) + + fn( + from_dlpack(q), + from_dlpack(k), + from_dlpack(v), + from_dlpack(o), + from_dlpack(lse) if lse is not None else None, + from_dlpack(seq_kv), + from_dlpack(seq_q), + from_dlpack(sinks_log2), + from_dlpack(bias), + from_dlpack(cu_q), + from_dlpack(cu_k), + from_dlpack(rope_cs), + cutlass.Int32(n_kv_tiles), + cutlass.Float32(scale_log2), + cutlass.Int32(sq), + cutlass.Int32(skv), + cutlass.Int32(d), + cutlass.Int32(right_bound), + cutlass.Float32(inv_scale), + cutlass.Int32(thd_q_tiles), + cutlass.Int32(n_batch_logical), + stream, + ) + + class SdpaFwdDslSm80(SdpaFwdDsl): - """SM80 (A100) SDPA forward via the pre-TemplateParams CuTe-DSL kernels. - - Follows the SM100/SM120 adapter lifecycle (check_support → compile → - execute, caller buffers bound directly) on top of the self-caching SM80 - kernel modules. SM80-only features (bias, RoPE) arrive as extra optional - ``execute`` keywords, exactly as the ``SdpaFwdDsl.execute`` contract - permits. ALiBi, block_mask and the score-stat side outputs are - deliberately NOT served: the capability row declines such graphs and the - backend takes them. - - Known deviations, both pre-existing and shared with the old - ``SdpafwdSm80`` path rather than introduced here: - - - dense GQA/MQA is served by expanding K/V heads adapter-side - (``repeat_interleave``): the kernels' native dense-GQA path exists but - is not yet qualified by the upstream validation harness (see - ``graph_analyzer.expand_gqa_heads``). Dropping the expansion once that - path is qualified on A100 CI is the tracked cleanup. - - a head-dim between flavor points pads V host-side - (``d_pad_multiple=1`` in the capability row). + """SM80 (A100) SDPA forward via the FROST template kernels. + + Since the TemplateParams conversion this adapter has the same shape as + its SM100/SM120 siblings end to end: ``check_support`` resolves the + flavor/mask/scheduler into a plan-time :class:`config_sm80.TemplateParams`, + ``compile()`` loads the specialized template module and compiles the + per-shape artifact ONCE (THD packed token extents compile dynamic via + ``cute.sym_int`` — issue #604's key is gone), and ``execute()`` re-binds + caller buffers to the cached artifact (a compile-cache miss at execute is + a bug by contract). + + SM80-only compile axes that have no home in the shared constructor arrive + as extra keyword-only arguments (``bias_present`` / ``bias_fp32`` / + ``rope_max_s``); the engine lowering forwards them only when the graph + declares the operands. ALiBi, block_mask and the score-stat side outputs + are deliberately NOT served: the capability row declines such graphs and + the backend takes them. + + Known deviations, pre-existing and tracked rather than introduced here: + dense GQA expands K/V heads adapter-side until the kernels' native dense + GQA path is qualified (see ``graph_analyzer.expand_gqa_heads``); an + off-flavor head dim pads V (and O, via a scratch) host-side; sink logits + are rescaled to log2 units with one (H,)-element multiply per execute. """ - def __init__(self, *args, scheduler: Optional[str] = None, **kwargs) -> None: - # SM80-only scheduler-token override for standalone callers - # ("auto"/"default"/"natural"/"lpt"/"lpt_l2"). The graph path never - # sets it: the engine row's sched_policy domain is empty, so the - # base sched_policy stays NATURAL and maps to "auto" below. + def __init__(self, *args, scheduler: Optional[str] = None, bias_present: bool = False, bias_fp32: bool = False, rope_max_s: int = 0, **kwargs) -> None: + # SM80-only plan-time axes (see class docstring). ``scheduler`` is the + # token override for standalone callers; the graph path leaves it None + # and the sched_policy knob (NATURAL default) maps to "auto" below. self._scheduler_token = scheduler + self._bias_present = bool(bias_present) + self._bias_fp32 = bool(bias_fp32) + self._rope_max_s = int(rope_max_s) super().__init__(*args, **kwargs) def _initialize_implementation(self) -> None: @@ -2967,6 +3042,8 @@ def _initialize_implementation(self) -> None: self.mask_token: Optional[str] = None self.swa_window_runtime: int = 0 self.right_bound_runtime: int = 0 + self._k_mod = None + self._params = None # ------------------------------------------------------------------ def check_support(self) -> bool: @@ -3004,7 +3081,6 @@ def check_support(self) -> bool: f"H_q ({h_qo}) must be divisible by H_kv ({h_kv}) for GQA / MQA", ) - # ---- head-dim envelope ---------------------------------------- max_d_qk = max(fdqk for fdqk, _ in _SM80_FLAVOR_DIMS.values()) max_d_v = max(fdv for _, fdv in _SM80_FLAVOR_DIMS.values()) self._value_error_if( @@ -3014,7 +3090,6 @@ def check_support(self) -> bool: f"Larger heads are not yet ported.", ) - # ---- dtype: FP16 or BF16 (both ride one SM80 mma pipeline) ---- self.dtype = self._check_dtype(self.q_desc, [torch.float16, torch.bfloat16], name="Q") for desc in (self.k_desc, self.v_desc, self.o_desc): self._check_dtype( @@ -3032,17 +3107,19 @@ def check_support(self) -> bool: self._check_tensor_shape(self.lse_desc, (b, h_qo, s_qo), name="LSE") self._value_error_if(not self.lse_desc.is_contiguous(), "LSE must be contiguous on SM80") - # ---- modes this adapter does not serve ------------------------- self._not_implemented_error_if( self.thd or self.cu_seq_q_lens or self.cu_seq_kv_lens, - "SdpaFwdDslSm80 does not serve packed THD / cu_seq_len graphs; " "sdpa_fwd_wrapper_sm80's THD path launches them directly", + "SdpaFwdDslSm80 does not serve packed THD / cu_seq_len graphs; " "sdpa_fwd_wrapper_sm80's varlen path launches them directly", ) self._not_implemented_error_if( self.window_size_right is not None and not self.is_causal, "SM80 SDPA: window_size_right without is_causal=True has no diagonal to anchor to", ) + self._not_implemented_error_if( + self._rope_max_s and (self.seq_kv_lens_present or self.seq_q_lens_present), + "SM80 SDPA: RoPE fusion is dense-unpadded-only", + ) - # ---- arch ------------------------------------------------------- self._value_error_if(not torch.cuda.is_available(), "CUDA must be available for SM80 SDPA") device = self.q_desc.device major, minor = torch.cuda.get_device_capability(device) @@ -3052,34 +3129,22 @@ def check_support(self) -> bool: f"SdpaFwdDslSm80 requires SM80 (A100); found SM{major}{minor} on {device}", ) - # ---- flavor + tile knobs --------------------------------------- self.flavor = _sm80_pick_flavor(d_qk, d_v) self.flavor_d_qk, self.flavor_d_v = _SM80_FLAVOR_DIMS[self.flavor] tile_m_default, num_warps_default, tile_n_default = _SM80_FLAVOR_KNOBS[self.flavor] self.kernel_tile_m = tile_m_default if self.tile_m is None else int(self.tile_m) self.kernel_num_warps = num_warps_default self.kernel_tile_n = tile_n_default if self.tile_n is None else int(self.tile_n) - self._value_error_if( - self.kernel_tile_n not in (64, 128), - f"SM80 SDPA: tile_n must be 64 or 128; got {self.kernel_tile_n}", - ) - self._value_error_if( - self.kernel_tile_m % (self.kernel_num_warps * 16) != 0, - f"SM80 SDPA: tile_m ({self.kernel_tile_m}) must be a multiple of " f"num_warps*16 ({self.kernel_num_warps * 16})", - ) self._value_error_if( self.cga not in (None, 1), f"SM80 SDPA has no CGA clustering; cga must be 1 (or unset), got {self.cga}", ) - # Bottom-right alignment shifts the band to the corner; it must affect - # SOMETHING — a causal upper bound and/or a left sliding-window. self._value_error_if( self.causal_bottom_right and not (self.is_causal or (self.window_size_left is not None and self.window_size_left >= 0)), "SM80 SDPA: causal_bottom_right requires is_causal=True and/or a left sliding-window (window_size_left >= 0).", ) - # ---- mask token -------------------------------------------------- swa_left = -1 if self.window_size_left is None else int(self.window_size_left) swa_right = 0 if self.window_size_right is None else int(self.window_size_right) self.right_bound_runtime = 0 @@ -3094,10 +3159,6 @@ def check_support(self) -> bool: self.mask_token = "none" self.swa_window_runtime = 0 - # ---- scheduler --------------------------------------------------- - # Explicit token (standalone callers) wins; otherwise the sched_policy - # knob maps NATURAL → "auto" (the SM80 heuristic decides), LPT/LPT_L2 - # → their tokens. token = self._scheduler_token if token is None: token = {SCHED_NATURAL: "auto", SCHED_LPT: "lpt", SCHED_LPT_L2: "lpt_l2"}.get(self.sched_policy) @@ -3116,7 +3177,6 @@ def check_support(self) -> bool: skv=int(s_kv), ) - # ---- softmax scale ------------------------------------------- if self.scale_softmax is None or self.scale_softmax == 0.0: self.scale_softmax = 1.0 / math.sqrt(d_qk) @@ -3134,18 +3194,52 @@ def check_support(self) -> bool: # ------------------------------------------------------------------ def compile(self) -> None: - """Mark compiled — the kernel module owns the JIT cache. - - The SM80 kernels self-cache one artifact per (shape, feature) key - (``_compile_cached``'s lru_cache), so there is no template module to - load here; the first ``execute()`` triggers the JIT. On the graph path - every key component is plan-time data (dense B/H/S/D from the frozen - graph), so the cache key is Rule-4 clean; the THD wrapper path's - packed-total keys are the known issue #604 cleanup. + """Load the TemplateParams-specialized module and compile the artifact. + + Plan-time only (Hard Rule 4): every key component here is graph + declaration or capability data; execute()'s call into the module's + per-shape lru is a guaranteed hit. """ - self._logger.debug("Entering compile (no-op — kernel self-caches)") + self._logger.debug("Entering compile") self._ensure_support_checked() - self._compiled_kernel = True + from cudnn.sdpa.fwd import config_sm80 as _sm80_cfg + + self._params = _sm80_cfg.TemplateParams( + io_bf16=(self.dtype == torch.bfloat16), + d_qk=self.flavor_d_qk, + d_v=self.flavor_d_v, + tile_m=self.kernel_tile_m, + num_warps=self.kernel_num_warps, + tile_n=self.kernel_tile_n, + is_causal=self.mask_token in ("causal", "causal_swa"), + has_swa=self.mask_token in ("swa", "causal_swa"), + causal_bottom_right=self.causal_bottom_right, + has_seq_kv_lens=self.seq_kv_lens_present, + has_seq_q_lens=self.seq_q_lens_present, + has_sink=self.has_sink, + has_bias=self._bias_present, + bias_is_fp32=self._bias_fp32, + has_rope=self._rope_max_s > 0, + thd_varlen=False, + sched_policy=_sm80_sched_policy_int(self.sched_token), + sched_l2_mib=self.sched_l2_mib, + has_lse=self.lse_desc is not None, + ) + self._k_mod = _sm80_load_kernel_module(self.flavor, self._params) + self._compiled_kernel = self._k_mod.compile( + b=self.batch_size, + # Dense GQA is served by adapter-side K/V head expansion until the + # kernels' native dense-GQA path is qualified (class docstring), so + # the artifact is compiled against the EXPANDED head count — the + # shapes execute() actually binds. + h=self.h_q, + h_kv=self.h_q, + sq=self.s_q_max, + skv=self.s_k_max, + d=self.head_dim_qk, + swa_window=int(self.swa_window_runtime), + rope_max_s=self._rope_max_s, + ) self._logger.debug("compile completed") def scratch_workspace_bytes(self) -> int: @@ -3171,15 +3265,23 @@ def execute( self._logger.debug("Entering execute") if self._compiled_kernel is None: raise RuntimeError("SdpaFwdDslSm80 is not compiled") + p = self._params - # Graph Stats declarations arrive as (B, H, S, 1); the kernels write - # [B, H, SQ] — rebind the squeezed view (zero-cost, same storage). + # Init-time flags are compile-time specializations; execute must match + # them exactly, in both directions (Hard Rule 1). if lse_tensor is not None and lse_tensor.ndim == 4: lse_tensor = lse_tensor.squeeze(-1) + self._value_error_if(p.has_lse and lse_tensor is None, "compiled with a Stats output but execute() got no lse_tensor") + self._value_error_if(not p.has_lse and lse_tensor is not None, "lse_tensor provided but the plan compiled the LSE store out") + self._value_error_if(p.has_bias != (bias_tensor is not None), "bias presence must match the compiled specialization") + self._value_error_if((p.has_rope) != (rope_freqs is not None), "rope_freqs presence must match the compiled specialization") + self._value_error_if(p.has_sink != (sinks is not None), "sinks presence must match the compiled specialization") + self._value_error_if(p.has_seq_kv_lens != (seq_kv_lens is not None), "seq_kv_lens presence must match the compiled specialization") + self._value_error_if(p.has_seq_q_lens != (seq_q_lens is not None), "seq_q_lens presence must match the compiled specialization") scale_val = self.scale_softmax if (scale_softmax is None or scale_softmax == 0.0) else float(scale_softmax) - kernel = _sm80_kernel_mod(self.flavor) device = q_tensor.device + launch_stream = self._get_default_stream(current_stream) with _torch_stream_context(current_stream, device): # BHSD → BSHD views; a dense_flex layout that is not BSHD-physical @@ -3190,8 +3292,7 @@ def execute( V = self._to_bshd(v_tensor) if self.h_kv != self.h_q: # Dense GQA: expand K/V heads until the kernels' native dense - # GQA path is qualified (see class docstring). BSHD head dim - # is 2. + # GQA path is qualified (see class docstring). BSHD head dim is 2. reps = self.h_q // self.h_kv K = K.repeat_interleave(reps, dim=2) V = V.repeat_interleave(reps, dim=2) @@ -3200,72 +3301,127 @@ def execute( if pad_v: V = _sm80_pad_last_dim(V, self.flavor_d_v) - # Output binding: hand the kernel the caller's BSHD view directly. - # A non-BSHD-physical O (dense_flex) binds a contiguous scratch the - # kernel writes, copied back below (the SM100 dense path's - # grandfathered normalization); only the padded-V envelope case - # falls back to kernel-side allocation + slice/copy. + # Output binding: the compiled O ABI is (B, SQ, H, flavor_d_v). + # Direct-bind the caller's BSHD view when it matches; the padded-V + # envelope and dense_flex cases go through a scratch + copy-back + # (both pre-existing normalizations). o_view, o_needs_copyback, o_scratch = self._to_bshd_writable(o_tensor) if pad_v: - bind_o = None + o_kernel = torch.zeros(self.batch_size, self.s_q_max, self.h_q, self.flavor_d_v, dtype=q_tensor.dtype, device=device) elif o_needs_copyback: - bind_o = o_scratch + o_kernel = o_scratch else: - bind_o = o_view - bind_lse = lse_tensor if (lse_tensor is not None and lse_tensor.is_contiguous()) else None - - fwd_kwargs = dict( - scale=scale_val, - return_lse=True, - tile_m=self.kernel_tile_m, - num_warps=self.kernel_num_warps, - tile_n=self.kernel_tile_n, - d_qk=self.flavor_d_qk, - d_v=self.flavor_d_v, - mask=self.mask_token, - swa_window=int(self.swa_window_runtime), + o_kernel = o_view + # DEFENSIVE zero-fill, not load-bearing: the dense epilogue stores + # every in-bounds row unconditionally; kept so a bound buffer can + # never surface uninitialized memory if a future path skips rows. + # (The pad_v scratch above is allocated zeroed already.) + if (seq_q_lens is not None or seq_kv_lens is not None) and not pad_v: + o_kernel.zero_() + if lse_tensor is not None: + lse_tensor.zero_() + + # Dummies fill compiled-out ABI slots (Rule 1); the kernel never + # dereferences them. Base-class ``_dummy`` (key, device, factory). + seq_kv_b = ( + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") + if seq_kv_lens is not None + else self._dummy("seq_i32", device, lambda: torch.ones(1, dtype=torch.int32, device=device)) + ) + seq_q_b = ( + self._checked_seq_lens(seq_q_lens, "seq_q_lens") + if seq_q_lens is not None + else self._dummy("seq_i32", device, lambda: torch.ones(1, dtype=torch.int32, device=device)) + ) + if sinks is not None: + # log2-unit rescale: one (H,)-element multiply per execute + # (pre-existing SM80 contract; the kernels consume log2 units). + sinks_b = (self._checked_sinks_1d(sinks) * _LOG2E).contiguous() + else: + sinks_b = self._dummy("one_f32", device, lambda: torch.ones(1, dtype=torch.float32, device=device)) + if bias_tensor is not None: + self._value_error_if( + bias_tensor.dtype != (torch.float32 if p.bias_is_fp32 else q_tensor.dtype), + f"bias dtype must match the compiled specialization; got {bias_tensor.dtype}", + ) + self._value_error_if( + tuple(bias_tensor.shape[-3:]) != (self.h_q, self.s_q_max, self.s_k_max), + f"bias trailing dims must be (H, SQ, SKV) = ({self.h_q}, {self.s_q_max}, {self.s_k_max}); got {tuple(bias_tensor.shape)}", + ) + bias_b = bias_tensor[:1] if bias_tensor.shape[0] != 1 else bias_tensor + self._value_error_if(not bias_b.is_contiguous(), "bias must be contiguous") + else: + bias_dt = q_tensor.dtype + bias_b = self._dummy(f"one_{bias_dt}", device, lambda: torch.ones(1, dtype=bias_dt, device=device)) + if rope_freqs is not None: + # (cos, sin) table build — wrapper-only fusion (the engine row + # never admits RoPE); per-execute by contract, like the caller + # passing fresh angle tables. + d2 = self.flavor_d_qk // 2 + rf = rope_freqs.to(dtype=torch.float32, device=device).reshape(rope_freqs.shape[0], -1) + self._value_error_if(rf.shape[1] < d2, f"rope_freqs last dim ({rf.shape[1]}) must be >= d_qk//2 ({d2})") + self._value_error_if( + rf.shape[0] != self._rope_max_s, f"rope_freqs rows ({rf.shape[0]}) must equal the compiled rope_max_s ({self._rope_max_s})" + ) + angles = rf[:, :d2] + rope_b = torch.stack([angles.cos(), angles.sin()], dim=-1).contiguous() + else: + rope_b = self._dummy("one_f32", device, lambda: torch.ones(1, dtype=torch.float32, device=device)) + cu_dummy = self._dummy("seq_i32", device, lambda: torch.ones(1, dtype=torch.int32, device=device)) + + _sm80_call( + self._compiled_kernel, + q=Q, + k=K, + v=V, + o=o_kernel, + lse=lse_tensor if p.has_lse else None, + seq_kv=seq_kv_b, + seq_q=seq_q_b, + sinks_log2=sinks_b, + bias=bias_b, + cu_q=cu_dummy, + cu_k=cu_dummy, + rope_cs=rope_b, + n_kv_tiles=(self.s_k_max + p.tile_n - 1) // p.tile_n, + scale_log2=scale_val * _LOG2E, + sq=self.s_q_max, + skv=self.s_k_max, + d=self.head_dim_qk, right_bound=int(self.right_bound_runtime), - causal_bottom_right=self.causal_bottom_right, - seq_kv_lens=self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") if seq_kv_lens is not None else None, - seq_len_q=self._checked_seq_lens(seq_q_lens, "seq_q_lens") if seq_q_lens is not None else None, - bias=bias_tensor, - sinks=self._checked_sinks_1d(sinks) if sinks is not None else None, - sched=self.sched_token, - sched_l2_mib=self.sched_l2_mib, - rope_freqs=rope_freqs, - out_o=bind_o, - out_lse=bind_lse, + inv_scale=1.0 / float(scale_val), + thd_q_tiles=0, + n_batch_logical=1, + stream=launch_stream, ) - _accepted = _sm80_inspect.signature(kernel.forward).parameters - fwd_kwargs = {k: v for k, v in fwd_kwargs.items() if k in _accepted} - O_buf, LSE_buf = kernel.forward(Q, K, V, **fwd_kwargs) - - # Copy-back only on the fallback paths the binding above skipped. - if bind_o is None: - # pad_v: the kernel allocated a flavor-wide O — slice + copy. - o_view.copy_(O_buf[..., : self.head_dim_v]) + + if pad_v: + o_view.copy_(o_kernel[..., : self.head_dim_v]) elif o_needs_copyback: - # dense_flex O: the kernel wrote the contiguous scratch. o_view.copy_(o_scratch) - if lse_tensor is not None and bind_lse is None: - lse_tensor.copy_(LSE_buf) self._logger.debug("execute completed") -def _sm80_thd_forward(q, k, v, *, cu_q, cu_k, max_s_q, scale_softmax, is_causal, window_size, causal_bottom_right, bias_tensor, sinks): +def _sm80_thd_forward(q, k, v, *, cu_q, cu_k, max_s_q, scale_softmax, is_causal, window_size, causal_bottom_right, bias_tensor, sinks, current_stream=None): """THD / varlen forward: q/k/v are PACKED ``[1, T, H, D]`` (already BSHD — - no transpose), cu_q/cu_k are ``[B+1]`` cumulative seqlens. Routes straight - to the kernel's THD path (graph-safe over-provisioned grid), reusing the - flavor-pick + d-pad. Returns packed ``[1, T_q, H, D_v]`` O + LSE.""" + no transpose), cu_q/cu_k are ``[B+1]`` cumulative seqlens. Rides the same + TemplateParams-specialized module as the dense path; the packed token + extents compile DYNAMIC (``cute.sym_int``), so the compile key is + plan-time-only and a new token total re-binds the cached artifact (the + old per-total ``lru`` key — issue #604 — is gone). Returns packed + ``[1, T_q, H, D_v]`` O + packed ``[1, H, T_q]`` LSE.""" + from cudnn.sdpa.fwd import config_sm80 as _sm80_cfg + + if bias_tensor is not None: + raise NotImplementedError("SM80 SDPA THD does not support bias (varlen has no single [1,H,SQ,SKV] bias shape)") d_qk = q.shape[-1] d_v = v.shape[-1] h_q = q.shape[2] + h_kv = k.shape[2] + device = q.device flavor = _sm80_pick_flavor(d_qk, d_v) fdqk, fdv = _SM80_FLAVOR_DIMS[flavor] tile_m, num_warps, tile_n = _SM80_FLAVOR_KNOBS[flavor] - # Resolve the default scale from the USER's head dim before padding: the - # kernel would otherwise derive 1/sqrt(D) from the padded flavor width - # (e.g. 1/sqrt(128) for a d=96 llama-flavor call) — silently wrong. if scale_softmax is None or scale_softmax == 0.0: scale_softmax = 1.0 / math.sqrt(d_qk) if d_qk < fdqk: @@ -3274,42 +3430,86 @@ def _sm80_thd_forward(q, k, v, *, cu_q, cu_k, max_s_q, scale_softmax, is_causal, pad_v = d_v < fdv if pad_v: v = _sm80_pad_last_dim(v, fdv) - # mask token from cuDNN's (is_causal, window_size=(left,right)). wl, wr = window_size - if is_causal and wl >= 0: - mask_token, swa = "causal_swa", wl - elif is_causal: - mask_token, swa = "causal", 0 - elif wl >= 0: - mask_token, swa = "swa", wl - else: - mask_token, swa = "none", 0 right_bound = wr if (is_causal and wr is not None and wr > 0) else 0 - kernel = _sm80_kernel_mod(flavor) - fwd_kwargs = dict( - scale=scale_softmax, - return_lse=True, + + n_seqs = int(cu_q.numel()) - 1 + if n_seqs < 1: + raise ValueError("cu_seqlens_q must have >= 2 entries") + cu_q_t = cu_q.to(dtype=torch.int32, device=device).contiguous() + cu_k_t = cu_k.to(dtype=torch.int32, device=device).contiguous() + + params = _sm80_cfg.TemplateParams( + io_bf16=(q.dtype == torch.bfloat16), + d_qk=fdqk, + d_v=fdv, tile_m=tile_m, num_warps=num_warps, tile_n=tile_n, - d_qk=fdqk, - d_v=fdv, - mask=mask_token, - swa_window=int(swa), - right_bound=int(right_bound), + is_causal=bool(is_causal), + has_swa=wl is not None and wl >= 0, causal_bottom_right=bool(causal_bottom_right), - cu_seqlens_q=cu_q, - cu_seqlens_k=cu_k, - max_s_q=int(max_s_q), - bias=bias_tensor, - sinks=sinks, + has_sink=sinks is not None, + thd_varlen=True, + has_lse=True, + ) + mod = _sm80_load_kernel_module(flavor, params) + # Off-flavor d_qk was HOST-PADDED to fdqk above, so the compiled fakes and + # the runtime d must both be the padded width: the kernel derives its Q/K + # row strides from d_runtime (Q_ROW_STRIDE_E = H * d_runtime), and the + # zero columns are exact for the QK dot products. (Same contract as the + # pre-template forward(), which read d_runtime off the padded shape.) + fn = mod.compile( + b=1, + h=h_q, + h_kv=h_kv, + sq=0, + skv=0, + d=int(fdqk), + swa_window=int(max(0, wl)) if wl is not None and wl >= 0 else 0, + n_batch_logical=n_seqs, + ) + + t_q = q.shape[1] + o_buf = torch.zeros(1, t_q, h_q, fdv, dtype=q.dtype, device=device) + lse_buf = torch.zeros(1, h_q, t_q, dtype=torch.float32, device=device) + sinks_b = ( + (sinks.to(dtype=torch.float32, device=device).reshape(h_q) * _LOG2E).contiguous() + if sinks is not None + else torch.ones(1, dtype=torch.float32, device=device) + ) + dummy_i32 = torch.ones(1, dtype=torch.int32, device=device) + dummy_f32 = torch.ones(1, dtype=torch.float32, device=device) + dummy_io = torch.ones(1, dtype=q.dtype, device=device) + + _sm80_call( + fn, + q=q, + k=k, + v=v, + o=o_buf, + lse=lse_buf, + seq_kv=dummy_i32, + seq_q=dummy_i32, + sinks_log2=sinks_b, + bias=dummy_io, + cu_q=cu_q_t, + cu_k=cu_k_t, + rope_cs=dummy_f32, + n_kv_tiles=(int(k.shape[1]) + tile_n - 1) // tile_n, + scale_log2=float(scale_softmax) * _LOG2E, + sq=int(t_q), + skv=int(k.shape[1]), + d=int(fdqk), + right_bound=int(right_bound), + inv_scale=1.0 / float(scale_softmax), + thd_q_tiles=(int(max_s_q) + tile_m - 1) // tile_m, + n_batch_logical=n_seqs, + stream=current_stream if current_stream is not None else cuda.CUstream(torch.cuda.current_stream(device).cuda_stream), ) - acc = _sm80_inspect.signature(kernel.forward).parameters - fwd_kwargs = {kk: vv for kk, vv in fwd_kwargs.items() if kk in acc} - O_buf, LSE_buf = kernel.forward(q, k, v, **fwd_kwargs) if pad_v: - O_buf = O_buf[..., :d_v].contiguous() - return TupleDict(o_tensor=O_buf, lse_tensor=LSE_buf) + o_buf = o_buf[..., :d_v].contiguous() + return TupleDict(o_tensor=o_buf, lse_tensor=lse_buf) _sm80_wrapper_cache: dict = {} @@ -3339,8 +3539,8 @@ def sdpa_fwd_wrapper_sm80( Returns ``TupleDict(o_tensor=..., lse_tensor=...)`` matching the DSL wrappers' contract. Dense calls route through :class:`SdpaFwdDslSm80`; - packed THD calls (``cum_seqlen_*``) launch the kernel's varlen path - directly. ALiBi, block_mask and the score-stat side outputs are not + packed THD calls (``cum_seqlen_*``) ride the same template with dynamic + token extents. ALiBi, block_mask and the score-stat side outputs are not supported (use the graph API, which routes them to the cuDNN backend). """ if q_tensor.ndim != 4 or v_tensor.ndim != 4: @@ -3348,18 +3548,11 @@ def sdpa_fwd_wrapper_sm80( if scale_output not in (None, 1.0): raise NotImplementedError(f"SM80 SDPA: scale_output != 1.0 is not supported yet (got {scale_output})") - # THD / varlen: q/k/v are PACKED [1, T, H, D] (BSHD) + cu_seqlens. Handled - # by a dedicated path that skips the dense BHSD transpose + dense O alloc. if cum_seqlen_q_tensor is not None: if max_s_q is None: raise ValueError("THD path requires max_s_q (host int) for the grid") if causal_bottom_right and not (is_causal or window_size[0] >= 0): - # Same anchor rule check_support enforces on the dense path: a - # bare bottom-right alignment has nothing to align. raise ValueError("SM80 SDPA: causal_bottom_right requires is_causal=True and/or a left sliding-window (window_size_left >= 0).") - # Reject dense-only features up front: _sm80_thd_forward does not plumb - # them, and silently computing without a requested feature is worse - # than an error. for label, present in ( ("rope_freqs", rope_freqs is not None), ("seq_kv_lens", seq_kv_lens is not None), @@ -3382,24 +3575,22 @@ def sdpa_fwd_wrapper_sm80( causal_bottom_right=causal_bottom_right, bias_tensor=bias_tensor, sinks=sinks, + current_stream=current_stream, ) b, h_q, s_q, _ = q_tensor.shape d_v = v_tensor.shape[-1] - # O takes Q's leading shape but V's head dim — supports dsv3-style - # D_QK != D_V. Allocate as contiguous (B, S, H, D) then transpose to the - # (B, H, S, D) BSHD-physical view the adapter binds without copies. o_tensor = torch.empty( (b, s_q, h_q, d_v), dtype=q_tensor.dtype, device=q_tensor.device, ).transpose(1, 2) lse_tensor = _allocate_lse_tensor(q_tensor) + wl, wr = window_size if not is_causal and wr >= 0: - # A right bound has no diagonal to anchor to without is_causal — - # reject rather than silently pick a mask (matches the THD path). raise NotImplementedError("SM80 SDPA: window_size_right without is_causal=True has no effect; pass is_causal=True or a left window") + rope_max_s = int(rope_freqs.shape[0]) if rope_freqs is not None else 0 cache_key = ( q_tensor.shape, k_tensor.shape, @@ -3412,6 +3603,10 @@ def sdpa_fwd_wrapper_sm80( bool(causal_bottom_right), seq_kv_lens is not None, seq_len_q is not None, + sinks is not None, + bias_tensor is not None, + (bias_tensor.dtype if bias_tensor is not None else None), + rope_max_s, q_tensor.device, ) api = _sm80_wrapper_cache.get(cache_key) @@ -3429,7 +3624,11 @@ def sdpa_fwd_wrapper_sm80( scale_softmax=scale_softmax, seq_kv_lens_present=seq_kv_lens is not None, seq_q_lens_present=seq_len_q is not None, + has_sink=sinks is not None, scheduler=scheduler, + bias_present=bias_tensor is not None, + bias_fp32=(bias_tensor is not None and bias_tensor.dtype == torch.float32), + rope_max_s=rope_max_s, ) api.check_support() api.compile() diff --git a/python/cudnn/sdpa/fwd/config_sm80.py b/python/cudnn/sdpa/fwd/config_sm80.py index 5a6e315ab..902c68fc0 100644 --- a/python/cudnn/sdpa/fwd/config_sm80.py +++ b/python/cudnn/sdpa/fwd/config_sm80.py @@ -43,3 +43,98 @@ class Cfg: LLAMA_CFG = Cfg(D_QK=128, D_V=128, TILE_M=128, TILE_N=64, NUM_WARPS=8) DSV3_CFG = Cfg(D_QK=192, D_V=128, TILE_M=128, TILE_N=64, NUM_WARPS=8) QWEN_CFG = Cfg(D_QK=256, D_V=256, TILE_M=128, TILE_N=64, NUM_WARPS=8) + + +# --------------------------------------------------------------------------- +# TemplateParams — the compile-time identity of one SM80 kernel specialization. +# +# Mirrors config_sm100/config_sm120: a frozen, hashable record injected into +# the kernel template as the ``FROST_TEMPLATE_PARAMS`` module global by +# ``frost.template_loader.load_template``, so ``cutlass.const_expr`` folding +# specializes the traced code per parameter set. Everything here is PLAN-TIME +# data (graph declaration + capability row + knobs) — never a runtime tensor +# value (AGENTS.md Hard Rule 4). Shape axes (b/h/sq/skv, the actual head dim +# under the envelope, the SWA width) stay arguments of the template module's +# ``compile()`` and its per-shape lru cache; THD packed token totals compile +# DYNAMIC (``cute.sym_int``) there and are never part of any key. +# --------------------------------------------------------------------------- +from dataclasses import dataclass + +from cudnn.frost.tile_dsl.constants import SCHED_LPT, SCHED_LPT_L2, SCHED_NATURAL + +_SM80_SEQ_TILES_N = (64, 128) + + +@dataclass(frozen=True) +class TemplateParams: + """One SM80 prefill kernel specialization (the module-identity axes).""" + + # dtype: fp16 or bf16 I/O (one mma pipeline serves both). + io_bf16: bool = False + # Flavor envelope tile dims (the compile-time D box; the actual head dim + # is a compile() argument and may be smaller — loads zero-fill past it). + d_qk: int = 128 + d_v: int = 128 + # Tile geometry (swept-and-frozen per flavor; see the Cfg table above). + tile_m: int = 128 + num_warps: int = 8 + tile_n: int = 64 + # Mask family. right_bound stays a RUNTIME argument (it widens the causal + # band without changing the traced structure). + is_causal: bool = False + has_swa: bool = False + causal_bottom_right: bool = False + # Optional operands (compile-time ABI presence; a missing-but-required or + # provided-but-uncompiled operand raises at execute — Hard Rule 1). + has_seq_kv_lens: bool = False + has_seq_q_lens: bool = False + has_sink: bool = False + has_bias: bool = False + bias_is_fp32: bool = False + has_rope: bool = False + # Packed varlen (wrapper-only today; the engine row declares thd=False). + thd_varlen: bool = False + # Tile-scheduler policy, in the SHARED frost vocabulary + # (tile_dsl.constants.SCHED_*; the kernel's grid mapping interprets + # NATURAL as its plain 3-D grid). The L2 budget for SCHED_LPT_L2 is + # flavor-tuned plan-time data, so it rides here too. + sched_policy: int = SCHED_NATURAL + sched_l2_mib: int = 32 + # False compiles the LSE store out entirely (the template None-specializes + # the LSE argument) — a stats-less graph binds no LSE buffer at any level. + has_lse: bool = True + + +def validate_params(p: TemplateParams) -> None: + """Raising validator — a failure here means the capability row or the + adapter lied about what this template can serve (README Rule 2).""" + if p.tile_n not in _SM80_SEQ_TILES_N: + raise ValueError(f"sm80: tile_n must be one of {_SM80_SEQ_TILES_N}; got {p.tile_n}") + if p.num_warps not in (4, 8): + raise ValueError(f"sm80: num_warps must be 4 or 8; got {p.num_warps}") + if p.tile_m % (p.num_warps * 16) != 0: + raise ValueError(f"sm80: tile_m ({p.tile_m}) must be a multiple of num_warps*16 ({p.num_warps * 16})") + if p.d_qk % 16 != 0 or p.d_qk <= 0: + raise ValueError(f"sm80: template d_qk must be a positive multiple of 16 (m16n8k16 K); got {p.d_qk}") + if p.d_v % 16 != 0 or p.d_v <= 0: + raise ValueError(f"sm80: template d_v must be a positive multiple of 16 (cp.async + STG.128 epilogue); got {p.d_v}") + if (p.d_qk, p.d_v) not in {(cfg.D_QK, cfg.D_V) for cfg in (GPTOSS_CFG, LLAMA_CFG, DSV3_CFG, QWEN_CFG)}: + raise ValueError(f"sm80: (d_qk, d_v) = ({p.d_qk}, {p.d_v}) is not a swept flavor envelope") + if p.sched_policy not in (SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2): + raise ValueError(f"sm80: sched_policy must be a tile_dsl SCHED_* value; got {p.sched_policy}") + if p.sched_l2_mib <= 0: + raise ValueError(f"sm80: sched_l2_mib must be > 0; got {p.sched_l2_mib}") + if p.causal_bottom_right and not (p.is_causal or p.has_swa): + raise ValueError("sm80: causal_bottom_right requires is_causal and/or has_swa (nothing to align otherwise)") + if p.thd_varlen and (p.has_rope or p.has_seq_kv_lens or p.has_seq_q_lens or p.has_bias): + raise ValueError("sm80: THD carries lengths via cu_seqlens; rope / bias / dense seq-lens are dense-only") + + +def params_for_flavor(flavor: str, **overrides) -> TemplateParams: + """A TemplateParams seeded from one swept flavor's Cfg row.""" + cfg = {"gptoss": GPTOSS_CFG, "llama": LLAMA_CFG, "dsv3": DSV3_CFG, "qwen": QWEN_CFG}[flavor] + base = dict(d_qk=cfg.D_QK, d_v=cfg.D_V, tile_m=cfg.TILE_M, num_warps=cfg.NUM_WARPS, tile_n=cfg.TILE_N) + base.update(overrides) + p = TemplateParams(**base) + validate_params(p) + return p diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index d7ff1ff02..db78ec0be 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -349,6 +349,14 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti if facts.has_sink and capabilities.sink_dtypes is not None and facts.dtype not in capabilities.sink_dtypes: return f"sink token with dtype {facts.dtype} not in {sorted(str(d) for d in capabilities.sink_dtypes)}" + if facts.has_bias and capabilities.bias and facts.bias_t is not None: + # uniform_dtype covers K/V/O only; the serving adapters compile the + # bias load as fp32 or the io dtype, so anything else must decline + # HERE, not ValueError at execute. + bias_dt = facts.bias_t.get_data_type() + if bias_dt not in (cudnn.data_type.FLOAT, facts.dtype): + return f"bias dtype {bias_dt} must be fp32 or match the Q/K/V dtype ({facts.dtype})" + if facts.right_band_widening and facts.right_bound is not None and facts.right_bound < 0: return f"negative diagonal_band_right_bound ({facts.right_bound}) is not supported" @@ -728,6 +736,18 @@ def lower_dsl_prefill( tile_m=knobs.tile_m if knobs is not None else None, tile_n=knobs.tile_n if knobs is not None else None, cga=knobs.cga if knobs is not None else None, + # SM80-only PLAN-TIME axes (bias presence/dtype are compile-time + # specializations of that template): forwarded only to adapters whose + # constructor declares them — every other row's mismatch gated the + # operands off already. + **( + { + "bias_present": facts.bias_t is not None, + "bias_fp32": facts.bias_t is not None and facts.bias_t.get_data_type() == cudnn.data_type.FLOAT, + } + if "bias_present" in inspect.signature(_adapter(api_type).__init__).parameters + else {} + ), ) api.check_support() # raises ValueError / NotImplementedError if unsupported api.compile() diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py index caa4fc039..9cce73422 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm80.py @@ -3,6 +3,14 @@ """SM80 (Ampere / A100) SDPA prefill at d_qk = d_v = 256, FP16 (qwen). +This file is a TEMPLATE: ``frost.template_loader.load_template`` re-executes +it as a fresh module per ``config_sm80.TemplateParams`` (injected as the +``FROST_TEMPLATE_PARAMS`` module global), so the feature/config axes fold at +trace time; the remaining SHAPE axes compile through the module's own +``compile()`` (per-shape ``@lru_cache``; THD packed token totals are DYNAMIC +via ``cute.sym_int`` — plan-time-only keys, Hard Rule 4). The adapter +(``api_dsl.SdpaFwdDslSm80``) owns validation, operand binding and launch. + Sibling of ``prefill_f16_sm80.py`` — same online flash-attention recipe (rowwise max / sum, exp2 softmax with scale folded into the exponent, threadquad butterflies, RESCALE_THRESHOLD=8.0, mask pre-pass @@ -72,7 +80,7 @@ from functools import lru_cache from typing import Optional -import torch + import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute @@ -104,32 +112,31 @@ def from_dlpack(t, **kw): MASK_CAUSAL, MASK_SWA, ) - -# --------------------------------------------------------------------------- -# Compile-time shape constants. -# --------------------------------------------------------------------------- -# ``d_qk`` / ``d_v`` are now Constexpr params on the kernel (folded at trace -# time); the flavor (Llama d=128, GPT-OSS d=64, …) is picked by the driver -# by the adapter's per-flavor knob table -# (``cudnn.sdpa.fwd.api_dsl._SM80_FLAVOR_KNOBS``). Knob defaults -# below match the Llama flavor — overridden per-flavor in -# the adapter's knob table and threaded through ``forward()``. -DEFAULT_TILE_M = 128 -DEFAULT_NUM_WARPS = 8 -DEFAULT_TILE_N = 64 -DEFAULT_D_QK = 128 -DEFAULT_D_V = 128 +from cudnn.frost.tile_dsl.constants import ( # noqa: E402 + SCHED_LPT, + SCHED_NATURAL, +) +from cudnn.sdpa.fwd.config_sm80 import TemplateParams, validate_params # noqa: E402 ELEM_BYTES = 2 # fp16 (Phase 2/3 baseline) ELEMS_PER_LD = 8 # 8 fp16 per cp.async = 16 B (max throughput) -# Scheduler policy IDs. -SCHED_DEFAULT = 0 # Plain 3-D grid (q_tile, head, batch). No reorder. -SCHED_LPT = 1 # 1-D grid, reverse-row-order across (head, batch). -SCHED_LPT_L2 = 2 # 1-D grid, block-cyclic over L2-sized (batch, head) groups -# so each active block's K + V fits in L2. Reverse-row -# ordering WITHIN each block. +# Scheduler policy IDs — the shared frost vocabulary (tile_dsl.constants; +# identical 0/1/2 values the kernel always used). +SCHED_DEFAULT = SCHED_NATURAL # the kernel's plain 3-D grid (q_tile, head, batch) + + +# --------------------------------------------------------------------------- +# Template identity. +# --------------------------------------------------------------------------- +# Injected by ``frost.template_loader.load_template``; the module-level +# default keeps a direct import (tests, tooling) importable at the +# qwen-flavor (d=256) baseline. Tile geometry, head-dim envelope, dtype, mask +# family, operand presence, scheduler policy and LSE presence all live here +# and fold at trace time — ``compile()`` below carries only shape axes. +PARAMS: TemplateParams = globals().get("FROST_TEMPLATE_PARAMS", TemplateParams(d_qk=256, d_v=256, tile_m=128, num_warps=8, tile_n=64)) +validate_params(PARAMS) # --------------------------------------------------------------------------- @@ -141,7 +148,9 @@ def _sdpa_kernel( K: cute.Tensor, # [B, SKV, HQ, D_QK] fp16 (MHA: HK == HQ) V: cute.Tensor, # [B, SKV, HQ, D_V] fp16 O: cute.Tensor, # [B, SQ, HQ, D_V] fp16 (output) - LSE: cute.Tensor, # [B, HQ, SQ] fp32 (LSE in natural log) + LSE: Optional[cute.Tensor], # [B, HQ, SQ] fp32 (LSE in natural log); THD + # packed [1, HQ, T]. None ⇒ no Stats output — + # the whole LSE compute + store is compiled out. seq_kv_lens: cute.Tensor, # [B] int32 — per-batch KV length (padded mask) seq_len_q: cute.Tensor, # [B] int32 — per-batch effective Q length. Consulted # only when has_seq_len_q (BR diagonal under padding; @@ -363,7 +372,6 @@ def _sdpa_kernel( K_view = cutlass.make_array_view(K) V_view = cutlass.make_array_view(V) O_view = cutlass.make_array_view(O) - LSE_view = cutlass.make_array_view(LSE) # Shapes (B, S, H, D). GQA/MQA: H_q (Q heads) may exceed H_kv (K/V # heads) when H_q % H_kv == 0. ``heads_per_kv`` Q-heads share each @@ -443,11 +451,16 @@ def _sdpa_kernel( O_BASE = q_seq_abs64 * O_ROW_STRIDE_E.to(cutlass.Int64) + cutlass.Int64(head_idx) * d_v64 # LSE: dense [B, H_q, SQ]; THD packed [1, H_q, T] — element-contiguous along # the seq axis. THD uses the packed seq position (cu_q[b] + q_row_base). - if cutlass.const_expr(THD_VARLEN): - T_lse_i64 = cutlass.Int64(LSE.shape[2]) - LSE_BASE = cutlass.Int64(head_idx) * T_lse_i64 + q_seq_origin + cutlass.Int64(q_row_base) - else: - LSE_BASE = cutlass.Int64(batch_idx) * H_i64 * SQ_i64 + cutlass.Int64(head_idx) * SQ_i64 + cutlass.Int64(q_row_base) + # None-specialized: no Stats output ⇒ view/base/pointer are compiled out + # together with the epilogue store below. + if cutlass.const_expr(LSE is not None): + LSE_view = cutlass.make_array_view(LSE) + if cutlass.const_expr(THD_VARLEN): + T_lse_i64 = cutlass.Int64(LSE.shape[2]) + LSE_BASE = cutlass.Int64(head_idx) * T_lse_i64 + q_seq_origin + cutlass.Int64(q_row_base) + else: + LSE_BASE = cutlass.Int64(batch_idx) * H_i64 * SQ_i64 + cutlass.Int64(head_idx) * SQ_i64 + cutlass.Int64(q_row_base) + lse_gmem = LSE_view.data_ptr() + LSE_BASE # K/V offsets parameterised on kv_row_base (variable in mainloop). Uses # H_kv and kv_head_idx (Q-heads sharing a KV head land at the same K/V @@ -461,7 +474,6 @@ def _sdpa_kernel( q_gmem = Q_view.data_ptr() + Q_BASE o_gmem = O_view.data_ptr() + O_BASE - lse_gmem = LSE_view.data_ptr() + LSE_BASE k_gmem_batch_head = K_view.data_ptr() + K_BATCH_OFF + K_HEAD_OFF v_gmem_batch_head = V_view.data_ptr() + V_BATCH_OFF + V_HEAD_OFF # Per-iter element advance: ``+ TILE_N rows`` on a Float16-typed ptr. @@ -1373,50 +1385,53 @@ def _mask_term(col, q_abs, q_mw): # [B,SQ,H,D] tensor; THD must NOT write rows past this sequence (they # belong to the next packed sequence) → bound by the per-batch eff_sq. sq_store_bound = eff_sq if cutlass.const_expr(THD_VARLEN) else sq_runtime - for m_block in cutlass.range_constexpr(m_blocks): - block_warp_row = warp_m_base + m_block * 16 - block_row_top = block_warp_row + g_lane - block_row_bot = block_warp_row + g_lane + 8 - row_state_lo = m_block * 2 - row_state_hi = m_block * 2 + 1 - lse_top = LN2 * (row_max[row_state_lo] + cute.math.log2(row_sum[row_state_lo])) - lse_bot = LN2 * (row_max[row_state_hi] + cute.math.log2(row_sum[row_state_hi])) - # Dense PADDED (per-batch seq_len_q): padded query rows q >= eff_sq are - # within this batch's [B,SQ,..] slice but don't exist → lse = -inf - # (log-sum-exp of a fully-masked row; matches cuDNN >= 9.14, and what - # test_mhas_v2 expects for stats on padded rows). Applied BEFORE the - # store-path split below: the is_even_mn fast path stores every row, so - # trimming only in the predicated path left finite LSE on padded rows - # whenever SQ is tile-aligned. THD is bounded by sq_store_bound==eff_sq - # instead (must NOT write the next packed seq's rows), so this select is - # gated on has_seq_len_q only. Mirrors prefill_f16_sm80.py; the SM80 - # bprop reads this lse and masks P=0 for padded rows (inf->0 select, - # no NaN), so the -inf is safe downstream. - if cutlass.const_expr(has_seq_len_q): - _ninf = cutlass.Float32(float("-inf")) - trim_top = q_row_base_i32 + block_row_top - trim_bot = q_row_base_i32 + block_row_bot - lse_top = cutlass.Float32(arith.select((trim_top < eff_sq).ir_value(), lse_top.ir_value(), _ninf.ir_value())) - lse_bot = cutlass.Float32(arith.select((trim_bot < eff_sq).ir_value(), lse_bot.ir_value(), _ninf.ir_value())) - lse_top_ptr = lse_gmem + cutlass.Int64(block_row_top) - lse_bot_ptr = lse_gmem + cutlass.Int64(block_row_bot) - - # Row predication for OOB Q-rows when ~is_even_mn. All 4 lanes of a - # threadquad share the same row → predicate is uniform within the - # threadquad, so the if-branch traces identically across them. - def _stf(ptr, val): - cutlass.Array(ptr, (1,), dtype=cutlass.Float32)[0] = val - - if cutlass.const_expr(is_even_mn): - _stf(lse_top_ptr, lse_top) - _stf(lse_bot_ptr, lse_bot) - else: - top_abs = q_row_base_i32 + block_row_top - bot_abs = q_row_base_i32 + block_row_bot - if top_abs < sq_store_bound: + # LSE output is None-specialized: no Stats output -> the whole block + # (compute + stores) is compiled out. + if cutlass.const_expr(LSE is not None): + for m_block in cutlass.range_constexpr(m_blocks): + block_warp_row = warp_m_base + m_block * 16 + block_row_top = block_warp_row + g_lane + block_row_bot = block_warp_row + g_lane + 8 + row_state_lo = m_block * 2 + row_state_hi = m_block * 2 + 1 + lse_top = LN2 * (row_max[row_state_lo] + cute.math.log2(row_sum[row_state_lo])) + lse_bot = LN2 * (row_max[row_state_hi] + cute.math.log2(row_sum[row_state_hi])) + # Dense PADDED (per-batch seq_len_q): padded query rows q >= eff_sq are + # within this batch's [B,SQ,..] slice but don't exist → lse = -inf + # (log-sum-exp of a fully-masked row; matches cuDNN >= 9.14, and what + # test_mhas_v2 expects for stats on padded rows). Applied BEFORE the + # store-path split below: the is_even_mn fast path stores every row, so + # trimming only in the predicated path left finite LSE on padded rows + # whenever SQ is tile-aligned. THD is bounded by sq_store_bound==eff_sq + # instead (must NOT write the next packed seq's rows), so this select is + # gated on has_seq_len_q only. Mirrors prefill_f16_sm80.py; the SM80 + # bprop reads this lse and masks P=0 for padded rows (inf->0 select, + # no NaN), so the -inf is safe downstream. + if cutlass.const_expr(has_seq_len_q): + _ninf = cutlass.Float32(float("-inf")) + trim_top = q_row_base_i32 + block_row_top + trim_bot = q_row_base_i32 + block_row_bot + lse_top = cutlass.Float32(arith.select((trim_top < eff_sq).ir_value(), lse_top.ir_value(), _ninf.ir_value())) + lse_bot = cutlass.Float32(arith.select((trim_bot < eff_sq).ir_value(), lse_bot.ir_value(), _ninf.ir_value())) + lse_top_ptr = lse_gmem + cutlass.Int64(block_row_top) + lse_bot_ptr = lse_gmem + cutlass.Int64(block_row_bot) + + # Row predication for OOB Q-rows when ~is_even_mn. All 4 lanes of a + # threadquad share the same row → predicate is uniform within the + # threadquad, so the if-branch traces identically across them. + def _stf(ptr, val): + cutlass.Array(ptr, (1,), dtype=cutlass.Float32)[0] = val + + if cutlass.const_expr(is_even_mn): _stf(lse_top_ptr, lse_top) - if bot_abs < sq_store_bound: _stf(lse_bot_ptr, lse_bot) + else: + top_abs = q_row_base_i32 + block_row_top + bot_abs = q_row_base_i32 + block_row_bot + if top_abs < sq_store_bound: + _stf(lse_top_ptr, lse_top) + if bot_abs < sq_store_bound: + _stf(lse_bot_ptr, lse_bot) # ---- Final normalization + SMEM-staged STG.128 epilogue -------------- # Naïve scalar per-lane STG (the previous epilogue) emits 2-fp16 stores @@ -1551,7 +1566,7 @@ def _sdpa_host( K: cute.Tensor, V: cute.Tensor, O: cute.Tensor, - LSE: cute.Tensor, + LSE: Optional[cute.Tensor], seq_kv_lens: cute.Tensor, seq_len_q: cute.Tensor, sinks: cute.Tensor, @@ -1655,120 +1670,144 @@ def _sdpa_host( # --------------------------------------------------------------------------- -# Compile cache. +# Per-shape compile cache. # --------------------------------------------------------------------------- # ``cute.compile`` is expensive (trace + MLIR + NVVM + PTX → SASS, ~1-2 s on -# A100). Without a per-process cache, every ``forward()`` call re-traces -# even for identical (shape, dtype, flag) tuples — costing 1+ s/call on -# top of the actual kernel runtime. ``@lru_cache`` keyed on the cache- -# significant tuple gives us O(1) lookup; the *value* is the compiled fn -# handle which is reusable across calls. +# A100). The FEATURE / config axes are module identity — one loaded template +# module per ``TemplateParams`` via ``frost.template_loader`` — so this cache +# covers the remaining SHAPE axes only. Every key component is PLAN-TIME +# data (AGENTS.md Hard Rule 4): under ``PARAMS.thd_varlen`` the packed token +# totals compile DYNAMIC (``cute.sym_int``) and are never part of the key — +# callers pass ``sq = skv = 0`` there (a stray runtime total must not be +# passed: it would only mint a redundant cache entry for the same artifact). @lru_cache(maxsize=None) -def _compile_cached( - B: int, - H: int, - H_kv: int, - SQ: int, - SKV: int, - D: int, - tile_m: int, - num_warps: int, - tile_n: int, - d_qk: int, - d_v: int, - io_is_bf16: bool, - is_even_mn: bool, - is_even_k: bool, - mask_flags: int, - swa_window: int, - causal_bottom_right: bool, - has_seq_kv_lens: bool, - has_seq_len_q: bool, - has_sink: bool, - has_bias: bool, - bias_is_fp32: bool, - THD_VARLEN: bool, - n_batch_logical: int, - has_rope: bool, - rope_max_s: int, - sched_policy: int, - sched_l2_bytes: int, +def compile( # noqa: A001 — the template contract's entry point (matches the SM100 kernels) + b: int, + h: int, + h_kv: int, + sq: int, + skv: int, + d: int, + swa_window: int = 0, + rope_max_s: int = 0, + n_batch_logical: int = 0, ): - """Compile (or return cached) ``_sdpa_host`` for the given config. - - Uses ``cute.runtime.make_fake_compact_tensor`` for shape-stable trace - inputs so the cached binary is reusable across different torch tensor - instances with the same shape signature. + """Compile (or fetch) this template specialization for one shape. + + ``d`` is the ACTUAL Q/K head dim and may be < ``PARAMS.d_qk``: the + SMEM/reg tile stays ``PARAMS.d_qk`` wide and the missing columns + zero-fill via cp.async predication. V/O are always exactly + ``PARAMS.d_v`` wide. ``swa_window`` is the left-window width W (keep + k in [q-W, q]) — plan-time graph data, baked exactly as the old + ``forward()`` did. Dense evenness is derived here the same way the old + entry point derived it: ``is_even_mn = (sq % tile_m == 0) and + (skv % tile_n == 0)``; ``is_even_k = (d == PARAMS.d_qk)``. + + THD (``PARAMS.thd_varlen``): q/k/v/o are packed ``[1, T, H, D]`` and the + LSE is packed ``[1, H, T]``; the token extents compile DYNAMIC — one + ``cute.sym_int`` symbol shared by the Q/O/LSE group and one for K/V — so + one artifact re-binds any packed totals (issue #604). Pass ``b = 1``, + ``sq = skv = 0``; ``n_batch_logical`` (the logical sequence count) sizes + the ``cu_seqlens`` ABI and IS plan-time. THD always takes the + predicated-store path (``is_even_mn = False``) and the over-provisioned + SCHED_DEFAULT grid, driven by the runtime ``thd_q_tiles``/``thd_n_batch`` + launch arguments. + + ``PARAMS.has_lse = False`` compiles the LSE store out entirely (the LSE + argument is None-specialized) — no buffer and no dummy at any level. """ - # Q and K share the QK head dim (= D, possibly < d_qk when ~is_even_k); - # V and O follow d_v (DSv3: d_qk != d_v). - io_dtype = cutlass.BFloat16 if io_is_bf16 else cutlass.Float16 + p = PARAMS + if p.thd_varlen and p.has_bias: + raise ValueError("sm80: bias + THD is not supported (varlen has no single [1,H,SQ,SKV] bias shape)") + io_dtype = cutlass.BFloat16 if p.io_bf16 else cutlass.Float16 + mask_flags = (MASK_CAUSAL if p.is_causal else MASK_NONE) | (MASK_SWA if p.has_swa else 0) + sched_l2_bytes = p.sched_l2_mib * 1024 * 1024 + is_even_k = d == p.d_qk + if p.thd_varlen: + # One symbol per ragged group: Q/O (and the LSE's T axis) share t_q, + # K/V share t_kv — a new packed total re-binds the same artifact. + is_even_mn = False + t_q = cute.sym_int(divisibility=1) + t_kv = cute.sym_int(divisibility=1) + _b, _sq, _skv = 1, t_q, t_kv + else: + is_even_mn = (sq % p.tile_m == 0) and (skv % p.tile_n == 0) + _b, _sq, _skv = b, sq, skv + + # Q and K share the QK head dim (= d, possibly < PARAMS.d_qk when + # ~is_even_k); V and O follow PARAMS.d_v (DSv3: d_qk != d_v). fake_q = cute.runtime.make_fake_compact_tensor( io_dtype, - (B, SQ, H, D), + (_b, _sq, h, d), stride_order=(3, 2, 1, 0), assumed_align=16, ) fake_k = cute.runtime.make_fake_compact_tensor( io_dtype, - (B, SKV, H_kv, D), + (_b, _skv, h_kv, d), stride_order=(3, 2, 1, 0), assumed_align=16, ) fake_v = cute.runtime.make_fake_compact_tensor( io_dtype, - (B, SKV, H_kv, d_v), + (_b, _skv, h_kv, p.d_v), stride_order=(3, 2, 1, 0), assumed_align=16, ) fake_o = cute.runtime.make_fake_compact_tensor( io_dtype, - (B, SQ, H, d_v), + (_b, _sq, h, p.d_v), stride_order=(3, 2, 1, 0), assumed_align=16, ) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (B, H, SQ), - stride_order=(2, 1, 0), - assumed_align=16, - ) + # LSE: dense [B, H, SQ]; THD packed [1, H, T] (shares the Q/O token + # symbol, which is exactly what the kernel's LSE.shape[2] read needs). + if p.has_lse: + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (_b, h, _sq), + stride_order=(2, 1, 0), + assumed_align=16, + ) + else: + fake_lse = None + # Per-batch KV / Q lengths [B] int32 (or a 1-elem dummy when unused). fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor( cutlass.Int32, - (B if has_seq_kv_lens else 1,), + (b if p.has_seq_kv_lens else 1,), stride_order=(0,), assumed_align=4, ) fake_seq_len_q = cute.runtime.make_fake_compact_tensor( cutlass.Int32, - (B if has_seq_len_q else 1,), + (b if p.has_seq_q_lens else 1,), stride_order=(0,), assumed_align=4, ) # Per-Q-head sink logit [H] fp32 (log2 units) (or a 1-elem dummy when unused). fake_sinks = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (H if has_sink else 1,), + (h if p.has_sink else 1,), stride_order=(0,), assumed_align=4, ) # Additive bias [1, H, SQ, SKV] in io_dtype or fp32 (or a 1-elem dummy). - bias_io_dtype = cutlass.Float32 if bias_is_fp32 else io_dtype + bias_io_dtype = cutlass.Float32 if p.bias_is_fp32 else io_dtype fake_bias = cute.runtime.make_fake_compact_tensor( bias_io_dtype, - ((1, H, SQ, SKV) if has_bias else (1,)), - stride_order=((3, 2, 1, 0) if has_bias else (0,)), + ((1, h, sq, skv) if p.has_bias else (1,)), + stride_order=((3, 2, 1, 0) if p.has_bias else (0,)), assumed_align=16, ) # THD cumulative seqlens [B_logical + 1] int32 (or 1-elem dummies). - _cu_len = (n_batch_logical + 1) if THD_VARLEN else 1 + _cu_len = (n_batch_logical + 1) if p.thd_varlen else 1 fake_cu_q = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (_cu_len,), stride_order=(0,), assumed_align=4) fake_cu_k = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (_cu_len,), stride_order=(0,), assumed_align=4) # RoPE (cos, sin) table [max_s, d_qk//2, 2] fp32 (or a 1-elem dummy). fake_rope_cs = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - ((rope_max_s, d_qk // 2, 2) if has_rope else (1,)), - stride_order=((2, 1, 0) if has_rope else (0,)), + ((rope_max_s, p.d_qk // 2, 2) if p.has_rope else (1,)), + stride_order=((2, 1, 0) if p.has_rope else (0,)), assumed_align=16, ) fake_n_kv_tiles = cutlass.Int32(0) @@ -1797,25 +1836,25 @@ def _compile_cached( fake_cu_q, fake_cu_k, fake_rope_cs, - tile_m, - num_warps, - tile_n, - d_qk, - d_v, + p.tile_m, + p.num_warps, + p.tile_n, + p.d_qk, + p.d_v, io_dtype, is_even_mn, is_even_k, mask_flags, swa_window, - causal_bottom_right, - has_seq_kv_lens, - has_seq_len_q, - has_sink, - has_bias, - bias_is_fp32, - THD_VARLEN, - has_rope, - sched_policy, + p.causal_bottom_right, + p.has_seq_kv_lens, + p.has_seq_q_lens, + p.has_sink, + p.has_bias, + p.bias_is_fp32, + p.thd_varlen, + p.has_rope, + p.sched_policy, sched_l2_bytes, fake_n_kv_tiles, fake_scale, @@ -1829,387 +1868,3 @@ def _compile_cached( fake_stream, options="--enable-tvm-ffi", ) - - -# --------------------------------------------------------------------------- -# Python entry point. -# --------------------------------------------------------------------------- -def forward( - Q: torch.Tensor, # [B, SQ, H, D] fp16 - K: torch.Tensor, # [B, SKV, H, D] fp16 - V: torch.Tensor, # [B, SKV, H, D] fp16 - scale: Optional[float] = None, - return_lse: bool = False, - *, - tile_m: int = DEFAULT_TILE_M, - num_warps: int = DEFAULT_NUM_WARPS, - tile_n: int = DEFAULT_TILE_N, - d_qk: int = DEFAULT_D_QK, - d_v: int = DEFAULT_D_V, - mask: str = "none", # "none" / "causal" / "swa" - swa_window: int = 0, # window width when mask == "swa" - right_bound: int = 0, # extra causal right band (k <= q + right_bound); - # 0 = plain causal. Requires MASK_CAUSAL. - causal_bottom_right: bool = False, # bottom-right causal diagonal (k <= q+SKV-SQ) - seq_kv_lens: Optional[torch.Tensor] = None, # [B] int32 per-batch KV length (padded) - seq_len_q: Optional[torch.Tensor] = None, # [B] int32 per-batch effective Q length - # → bottom-right diagonal under padding - # (br_base = eff_skv - eff_sq); only used - # with causal_bottom_right. - sinks: Optional[torch.Tensor] = None, # [H] fp32 per-Q-head sink logits in - # SCALED-logit (natural) units — joins the - # softmax denominator only (V_sink = 0). - # forward multiplies by log2(e). None → - # no sink. - bias: Optional[torch.Tensor] = None, # additive attention bias, broadcast over - # batch: shape [1, H, SQ, SKV] (or - # [B, H, SQ, SKV] — only slice 0 is read). - # dtype defaults to QKV (fp16/bf16); fp32 - # also accepted. None → no bias. When - # bias is set and tile_m/num_warps are at - # their defaults, the bias-enabled path - # defaults to tile_m=64/num_warps=4 (lower - # SMEM → 2 CTAs/SM to hide the bias LDGs). - cu_seqlens_q: Optional[torch.Tensor] = None, # [B+1] int32 cumulative Q seqlens. - # When set → THD/varlen mode: Q/K/V/O are - # PACKED [1, T, H, D] and cu_seqlens_* - # define each sequence's span. Device or - # host int32; copied to device int32. - cu_seqlens_k: Optional[torch.Tensor] = None, # [B+1] int32 cumulative KV seqlens. - max_s_q: Optional[int] = None, # THD: max sequence Q length (for the - # over-provisioned grid). If None it is - # computed from cu_seqlens_q (one d2h sync — - # pass it explicitly for graph capture). - sched: str = "auto", # "auto" (default→none / lpt→mask) / "default" - # / "lpt" / "lpt_l2" - sched_l2_mib: int = 32, # L2 budget (MiB) for "lpt_l2". A100 = 40 MiB - # physical; ~32 MiB usable after texture / RO. - rope_freqs: Optional[torch.Tensor] = None, # RoPE angles, shape [max_s, 1, 1, d_qk] - # (cuDNN graph.rope freqs convention) or - # [max_s, d_qk]. Only the first d_qk//2 cols - # are used; cos/sin precomputed on the HOST - # (torch full-range) and the half-split - # rotate_half applied to Q AND K in-kernel. - # Dense-only (no THD); max_s must cover - # max(SQ, SKV). None → no RoPE. - out_o: Optional[torch.Tensor] = None, # [B, SQ, H, d_v] caller-provided output - # buffer (Q dtype, contiguous); None → - # allocate. Dense-only (THD allocates). - out_lse: Optional[torch.Tensor] = None, # [B, H, SQ] fp32 contiguous LSE buffer; - # None → allocate (used when return_lse). -): - """Run SM80 SDPA prefill for an MHA-shaped (B, S, H, D) Q/K/V triple. - - Returns a [B, SQ, H, D] fp16 output tensor. When ``return_lse`` is - true, returns ``(O, LSE)`` with LSE shape ``[B, H, SQ]`` fp32 in - natural-log (matches ``torch.logsumexp(scale*Q@K.T, dim=-1)``). - - Configurable kernel variants via ``(tile_m, num_warps)``: - * ``(64, 4)`` — default; 128 threads / CTA, M_BLOCKS = 1. Fits - 2 CTAs / SM implicitly (≤128 regs/thread, ≤48 KiB SMEM/CTA). - * ``(128, 4)`` — same 128 threads, but each warp owns 2 m16n8k16 - row-blocks (M_BLOCKS=2). Spills at this register footprint on - SM80 — kept only as a path for future GPT-OSS-style configs. - * ``(128, 8)`` — 256 threads / CTA, 1 m16n8k16 block per warp - (M_BLOCKS=1, no register-pressure regression). Same K/V load - cost amortized over 2× Q rows. Targets 1 CTA / SM. - """ - assert Q.dtype in (torch.float16, torch.bfloat16), f"Q dtype must be float16 or bfloat16 (got {Q.dtype})" - assert K.dtype == Q.dtype and V.dtype == Q.dtype, f"K/V dtype must match Q ({Q.dtype}); got K={K.dtype} V={V.dtype}" - io_is_bf16 = Q.dtype == torch.bfloat16 - assert Q.is_cuda and K.is_cuda and V.is_cuda - B, SQ, H, D = Q.shape - _, SKV, Hk, D_K = K.shape - _, _, _, D_V_actual = V.shape - # GQA/MQA: H_q (= H) must be a multiple of H_kv (= Hk). MHA is H == Hk. - assert H % Hk == 0, f"H_q ({H}) must be a multiple of H_kv ({Hk}) for GQA/MQA" - assert K.shape[2] == V.shape[2], f"K H ({K.shape[2]}) must equal V H ({V.shape[2]})" - # Q and K share the QK head dim; V has its own d_v (asymmetric on DSv3: - # d_qk=192, d_v=128). D may be < d_qk — the SMEM/reg tile is sized for - # d_qk and the missing cols are zero-padded via cp.async predication - # when ~is_even_k. - assert D == D_K, f"Q D ({D}) must equal K D ({D_K})" - assert D_V_actual == d_v, f"V D ({D_V_actual}) must equal compile-time d_v={d_v}" - assert D <= d_qk, f"D ({D}) must be <= compile-time d_qk={d_qk}" - assert d_qk % 16 == 0, f"d_qk ({d_qk}) must be a multiple of 16 (m16n8k16 K)" - assert d_v % 8 == 0, f"d_v ({d_v}) must be a multiple of 8 (SV n_frags)" - assert d_v % 16 == 0, f"d_v ({d_v}) must be a multiple of 16 (cp.async + STG.128 epilogue)" - assert D % 8 == 0, f"D ({D}) must be a multiple of 8 (cp.async chunk size)" - # Bias-enabled default tile: lower SMEM (64 Q rows / 128 threads) so 2 CTAs - # fit per SM, helping hide the per-iter bias LDGs. Only when the caller - # left tile_m / num_warps at their module defaults. - has_bias = bias is not None - if has_bias and tile_m == DEFAULT_TILE_M and num_warps == DEFAULT_NUM_WARPS: - tile_m, num_warps = 64, 4 - assert tile_m % (num_warps * 16) == 0, f"tile_m={tile_m} must be a multiple of num_warps*16={num_warps*16}" - assert tile_n in (64, 128), f"tile_n must be 64 or 128 (got {tile_n})" - assert mask in ("none", "causal", "swa", "causal_swa"), f"mask: 'none' | 'causal' | 'swa' | 'causal_swa' (got {mask!r})" - assert (not causal_bottom_right) or mask in ("causal", "causal_swa", "swa"), ( - "causal_bottom_right=True requires mask in causal/causal_swa/swa " "(BR shifts the causal upper and/or SWA lower bound to the corner)" - ) - # Bottom-right diagonal under per-batch padding uses br_base = eff_skv - - # eff_sq: seq_kv_lens supplies eff_skv and seq_len_q supplies eff_sq. When - # BR + seq_kv_lens is requested without seq_len_q, eff_sq falls back to the - # physical SQ — correct only when every sequence is full-length. Mirror the - # f16 kernel: pass seq_len_q for ragged/padded BR. - assert right_bound >= 0, f"right_bound must be >= 0 (got {right_bound})" - assert right_bound == 0 or mask in ("causal", "causal_swa"), "right_bound>0 (causal right band) requires a causal mask" - assert sched in ("auto", "default", "lpt", "lpt_l2"), f"sched: 'auto' | 'default' | 'lpt' | 'lpt_l2' (got {sched!r})" - assert sched_l2_mib > 0, f"sched_l2_mib must be > 0 (got {sched_l2_mib})" - - is_even_mn = (SQ % tile_m == 0) and (SKV % tile_n == 0) - is_even_k = D == d_qk - mask_flags = MASK_NONE - if mask == "causal": - mask_flags |= MASK_CAUSAL - elif mask == "swa": - assert swa_window >= 0, "mask='swa' requires swa_window >= 0 (0 = 1-token window, keep k>=q)" - mask_flags |= MASK_SWA - elif mask == "causal_swa": - # Causal sliding window [q-W, q]: both mask bits; body composes them. - assert swa_window >= 0, "mask='causal_swa' requires swa_window >= 0 (0 = diagonal-only)" - mask_flags |= MASK_CAUSAL | MASK_SWA - - if sched == "auto": - # MASK_NONE → SCHED_DEFAULT (preserves L2 reuse); causal / SWA → LPT. - sched_policy = SCHED_DEFAULT if mask_flags == MASK_NONE else SCHED_LPT - elif sched == "default": - sched_policy = SCHED_DEFAULT - elif sched == "lpt": - sched_policy = SCHED_LPT - else: # "lpt_l2" - sched_policy = SCHED_LPT_L2 - sched_l2_bytes = int(sched_l2_mib) * 1024 * 1024 - - # ---- THD/varlen setup ------------------------------------------------- - # Packed [1,T,H,D] Q/K/V/O + cu_seqlens. Forces the predicated-store path - # (is_even_mn=False) and a SCHED_DEFAULT over-provisioned 3-D grid - # (ceil(max_s_q/tile_m), H, B_logical); tiles past a sequence early-out. - THD_VARLEN = cu_seqlens_q is not None - if THD_VARLEN: - assert cu_seqlens_k is not None, "THD needs both cu_seqlens_q and cu_seqlens_k" - assert B == 1, f"THD: Q/K/V must be packed [1, T, H, D] (got batch dim {B})" - assert not has_bias, "bias + THD not supported (varlen has no single " "[1,H,SQ,SKV] bias shape)" - n_batch_logical = int(cu_seqlens_q.numel()) - 1 - assert n_batch_logical >= 1, "cu_seqlens_q must have >= 2 entries" - cu_q_t = cu_seqlens_q.to(dtype=torch.int32, device=Q.device).contiguous() - cu_k_t = cu_seqlens_k.to(dtype=torch.int32, device=Q.device).contiguous() - if max_s_q is None: - _d = cu_q_t[1:] - cu_q_t[:-1] - max_s_q = int(_d.max().item()) # one d2h sync (non-graph) - thd_q_tiles_v = (int(max_s_q) + tile_m - 1) // tile_m - is_even_mn = False - sched_policy = SCHED_DEFAULT - else: - n_batch_logical = 1 - cu_q_t = torch.ones(1, dtype=torch.int32, device=Q.device) - cu_k_t = torch.ones(1, dtype=torch.int32, device=Q.device) - thd_q_tiles_v = 0 - - if scale is None: - scale = 1.0 / (D**0.5) - import math - - scale_log2 = scale * math.log2(math.e) - - # Allocate output / LSE sized to the actual (possibly uneven) shapes — - # the kernel writes the full row x d_v range via STG (rows predicated on - # sq_store_bound only). Output dim follows V (= d_v), which can differ - # from Q's d_qk (DSv3). Caller-bound out_* buffers skip the allocation. - # The zero-fill below is DEFENSIVE, not load-bearing: the dense epilogue - # stores every in-bounds row unconditionally (zero-KV-iteration tiles - # route through the store with row_sum=0, trimmed rows are written - # explicitly with O=0 / LSE=-inf, and THD never receives bound outputs). - # Kept so a bound buffer can never surface uninitialized memory if a - # future feature path skips rows. - _needs_zero_init = (seq_len_q is not None) or (seq_kv_lens is not None) - if out_o is None: - out_o = torch.zeros(B, SQ, H, d_v, dtype=Q.dtype, device=Q.device) - else: - assert out_o.shape == (B, SQ, H, d_v) and out_o.dtype == Q.dtype and out_o.is_contiguous(), ( - f"out_o must be a contiguous [{B}, {SQ}, {H}, {d_v}] {Q.dtype} tensor; " f"got shape {tuple(out_o.shape)} dtype {out_o.dtype}" - ) - if _needs_zero_init: - out_o.zero_() - if out_lse is None: - LSE = torch.zeros(B, H, SQ, dtype=torch.float32, device=Q.device) - else: - assert out_lse.shape == (B, H, SQ) and out_lse.dtype == torch.float32 and out_lse.is_contiguous(), ( - f"out_lse must be a contiguous [{B}, {H}, {SQ}] fp32 tensor; " f"got shape {tuple(out_lse.shape)} dtype {out_lse.dtype}" - ) - LSE = out_lse - if _needs_zero_init: - LSE.zero_() - has_seq_kv_lens = seq_kv_lens is not None - if has_seq_kv_lens: - assert seq_kv_lens.shape == (B,), f"seq_kv_lens must be shape ({B},); got {tuple(seq_kv_lens.shape)}" - seq_kv_lens_t = seq_kv_lens.to(dtype=torch.int32, device=Q.device).contiguous() - else: - seq_kv_lens_t = torch.ones(1, dtype=torch.int32, device=Q.device) - - has_seq_len_q = seq_len_q is not None - if has_seq_len_q: - assert seq_len_q.shape == (B,), f"seq_len_q must be shape ({B},); got {tuple(seq_len_q.shape)}" - seq_len_q_t = seq_len_q.to(dtype=torch.int32, device=Q.device).contiguous() - else: - seq_len_q_t = torch.ones(1, dtype=torch.int32, device=Q.device) - - # Per-Q-head sink logits → log2 units (= sink * log2(e)) so they share the - # kernel's log2-of-scaled-logit max domain. - has_sink = sinks is not None - if has_sink: - assert sinks.shape == (H,), f"sinks must be shape ({H},); got {tuple(sinks.shape)}" - sinks_t = (sinks.to(dtype=torch.float32, device=Q.device) * math.log2(math.e)).contiguous() - else: - sinks_t = torch.ones(1, dtype=torch.float32, device=Q.device) - - # Additive bias: dtype defaults to QKV (fp16/bf16); fp32 also accepted. - # Broadcast over batch — keep only the [1, H, SQ, SKV] slice (the kernel - # reads slice 0). inv_scale = 1/scale is folded in-kernel. - inv_scale = 1.0 / float(scale) - if has_bias: - assert bias.dtype in (Q.dtype, torch.float32), f"bias dtype must be {Q.dtype} or float32 (got {bias.dtype})" - assert tuple(bias.shape[-3:]) == (H, SQ, SKV), f"bias trailing dims must be (H={H}, SQ={SQ}, SKV={SKV}); " f"got {tuple(bias.shape)}" - bias_is_fp32 = bias.dtype == torch.float32 - bias_t = bias[:1].contiguous() if bias.shape[0] != 1 else bias.contiguous() - else: - bias_is_fp32 = False - bias_t = torch.ones(1, dtype=Q.dtype, device=Q.device) - - # RoPE: precompute the (cos, sin) table on the host (torch full-range cos/sin - # → bit-matches the reference, no in-kernel MUFU range error). Packed - # [max_s, d2, 2] fp32; applied as the half-split rotate_half to Q AND K. - has_rope = rope_freqs is not None - if has_rope: - assert not THD_VARLEN, "RoPE is dense-only (no THD/varlen) on SM80" - assert D == d_qk, f"RoPE requires D == d_qk (got D={D}, d_qk={d_qk}); the partial " "rope_dim < head_dim (MLA) path is not implemented" - d2 = d_qk // 2 - rf = rope_freqs.to(dtype=torch.float32, device=Q.device) - rf = rf.reshape(rf.shape[0], -1) # [max_s, d_qk] (or wider) - rope_max_s = rf.shape[0] - assert rf.shape[1] >= d2, f"rope_freqs last dim ({rf.shape[1]}) must be >= d_qk//2 ({d2})" - assert rope_max_s >= max(SQ, SKV), f"rope_freqs max_s ({rope_max_s}) must cover max(SQ={SQ}, SKV={SKV})" - angles = rf[:, :d2] # [max_s, d2] - rope_cs_t = torch.stack([angles.cos(), angles.sin()], dim=-1).contiguous() - else: - rope_max_s = 1 - rope_cs_t = torch.ones(1, dtype=torch.float32, device=Q.device) - - cQ = from_dlpack(Q) - cK = from_dlpack(K) - cV = from_dlpack(V) - cO = from_dlpack(out_o) - cLSE = from_dlpack(LSE) - cSEQK = from_dlpack(seq_kv_lens_t) - cSEQQ = from_dlpack(seq_len_q_t) - cSINKS = from_dlpack(sinks_t) - cBIAS = from_dlpack(bias_t) - cCUQ = from_dlpack(cu_q_t) - cCUK = from_dlpack(cu_k_t) - cROPE = from_dlpack(rope_cs_t) - torch_stream = torch.cuda.current_stream() - stream = cuda.CUstream(torch_stream.cuda_stream) - - # Host computes ``round_up(SKV / TILE_N)`` and passes as a runtime - # Int32 — the kernel uses ``cutlass.range(..., unroll=1)`` over it so - # the body only traces ONCE regardless of how many KV tiles SKV - # contains. Slashes ``cute.compile`` time from ~60 s (constexpr - # unroll over 128 iters at SQ=8192) to a few seconds. - n_kv_tiles = cutlass.Int32((SKV + tile_n - 1) // tile_n) - sq_rt = cutlass.Int32(SQ) - skv_rt = cutlass.Int32(SKV) - d_rt = cutlass.Int32(D) - # ``_compile_cached`` returns the same compiled fn on cache hit (~µs); - # on cache miss it traces + lowers (~1 s). Cache key is the full - # Constexpr tuple — different (mask, shape) combos compile separately. - fn = _compile_cached( - B, - H, - Hk, - SQ, - SKV, - D, - tile_m, - num_warps, - tile_n, - d_qk, - d_v, - io_is_bf16, - is_even_mn, - is_even_k, - mask_flags, - swa_window, - bool(causal_bottom_right), - bool(has_seq_kv_lens), - bool(has_seq_len_q), - bool(has_sink), - bool(has_bias), - bool(bias_is_fp32), - bool(THD_VARLEN), - int(n_batch_logical), - bool(has_rope), - int(rope_max_s), - sched_policy, - sched_l2_bytes, - ) - fn( - cQ, - cK, - cV, - cO, - cLSE, - cSEQK, - cSEQQ, - cSINKS, - cBIAS, - cCUQ, - cCUK, - cROPE, - n_kv_tiles, - cutlass.Float32(scale_log2), - sq_rt, - skv_rt, - d_rt, - cutlass.Int32(right_bound), - cutlass.Float32(inv_scale), - cutlass.Int32(thd_q_tiles_v), - cutlass.Int32(n_batch_logical), - stream, - ) - if return_lse: - return (out_o, LSE) - return out_o - - -if __name__ == "__main__": - # Phase 3 validation: compare against torch SDPA reference. - torch.manual_seed(0) - import torch.nn.functional as F - - for dt, atol in [(torch.float16, 5e-3), (torch.bfloat16, 4e-2)]: - print(f"--- dtype={dt} ---") - for B, H, SQ, SKV in [ - (1, 1, 64, 64), - (1, 1, 64, 128), - (1, 4, 64, 256), - (2, 4, 128, 512), - ]: - Q = torch.randn(B, SQ, H, 128, dtype=dt, device="cuda") - K = torch.randn(B, SKV, H, 128, dtype=dt, device="cuda") - V = torch.randn(B, SKV, H, 128, dtype=dt, device="cuda") - out_o = forward(Q, K, V) - # Torch reference: scaled_dot_product_attention takes (B, H, S, D). - ref = F.scaled_dot_product_attention( - Q.transpose(1, 2), - K.transpose(1, 2), - V.transpose(1, 2), - ).transpose(1, 2) - diff = (out_o.float() - ref.float()).abs() - maxd = diff.max().item() - ok = maxd < atol - print( - f" B={B} H={H} SQ={SQ:4d} SKV={SKV:4d} " - f"max|out_o|={out_o.abs().max().item():.4f} " - f"max|out_o-ref|={maxd:.4f} {'PASS' if ok else 'FAIL'}" - ) - print("[sm80-sdpa] done.") diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py index 8c75afd9f..8275aa249 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm80.py @@ -3,6 +3,14 @@ """SM80 (Ampere / A100) SDPA prefill, FP16 in / FP16 out. +This file is a TEMPLATE: ``frost.template_loader.load_template`` re-executes +it as a fresh module per ``config_sm80.TemplateParams`` (injected as the +``FROST_TEMPLATE_PARAMS`` module global), so the feature/config axes fold at +trace time; the remaining SHAPE axes compile through the module's own +``compile()`` (per-shape ``@lru_cache``; THD packed token totals are DYNAMIC +via ``cute.sym_int`` — plan-time-only keys, Hard Rule 4). The adapter +(``api_dsl.SdpaFwdDslSm80``) owns validation, operand binding and launch. + Online flash-attention with rowwise max / sum tracked per warp lane (2 M-rows per lane), threadquad butterfly reductions across the 4 lanes that share a row, exp2-based softmax with the input ``softmax_scale * log2(e)`` @@ -88,7 +96,7 @@ from functools import lru_cache from typing import Optional -import torch + import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute @@ -120,32 +128,31 @@ def from_dlpack(t, **kw): MASK_CAUSAL, MASK_SWA, ) - -# --------------------------------------------------------------------------- -# Compile-time shape constants. -# --------------------------------------------------------------------------- -# ``d_qk`` / ``d_v`` are now Constexpr params on the kernel (folded at trace -# time); the flavor (Llama d=128, GPT-OSS d=64, …) is picked by the driver -# by the adapter's per-flavor knob table -# (``cudnn.sdpa.fwd.api_dsl._SM80_FLAVOR_KNOBS``). Knob defaults -# below match the Llama flavor — overridden per-flavor in -# the adapter's knob table and threaded through ``forward()``. -DEFAULT_TILE_M = 128 -DEFAULT_NUM_WARPS = 8 -DEFAULT_TILE_N = 64 -DEFAULT_D_QK = 128 -DEFAULT_D_V = 128 +from cudnn.frost.tile_dsl.constants import ( # noqa: E402 + SCHED_LPT, + SCHED_NATURAL, +) +from cudnn.sdpa.fwd.config_sm80 import TemplateParams, validate_params # noqa: E402 ELEM_BYTES = 2 # fp16 (Phase 2/3 baseline) ELEMS_PER_LD = 8 # 8 fp16 per cp.async = 16 B (max throughput) -# Scheduler policy IDs. -SCHED_DEFAULT = 0 # Plain 3-D grid (q_tile, head, batch). No reorder. -SCHED_LPT = 1 # 1-D grid, reverse-row-order across (head, batch). -SCHED_LPT_L2 = 2 # 1-D grid, block-cyclic over L2-sized (batch, head) groups -# so each active block's K + V fits in L2. Reverse-row -# ordering WITHIN each block. +# Scheduler policy IDs — the shared frost vocabulary (tile_dsl.constants; +# identical 0/1/2 values the kernel always used). +SCHED_DEFAULT = SCHED_NATURAL # the kernel's plain 3-D grid (q_tile, head, batch) + + +# --------------------------------------------------------------------------- +# Template identity. +# --------------------------------------------------------------------------- +# Injected by ``frost.template_loader.load_template``; the module-level +# default keeps a direct import (tests, tooling) importable at the +# llama-flavor baseline. Tile geometry, head-dim envelope, dtype, mask +# family, operand presence, scheduler policy and LSE presence all live here +# and fold at trace time — ``compile()`` below carries only shape axes. +PARAMS: TemplateParams = globals().get("FROST_TEMPLATE_PARAMS", TemplateParams()) +validate_params(PARAMS) # --------------------------------------------------------------------------- @@ -157,7 +164,9 @@ def _sdpa_kernel( K: cute.Tensor, # [B, SKV, HQ, D_QK] fp16 (MHA: HK == HQ) V: cute.Tensor, # [B, SKV, HQ, D_V] fp16 O: cute.Tensor, # [B, SQ, HQ, D_V] fp16 (output) - LSE: cute.Tensor, # [B, HQ, SQ] fp32 (LSE in natural log) + LSE: Optional[cute.Tensor], # [B, HQ, SQ] fp32 (LSE in natural log); THD + # packed [1, HQ, T]. None ⇒ no Stats output — + # the whole LSE compute + store is compiled out. seq_kv_lens: cute.Tensor, # [B] int32 — per-batch effective KV length (padded # mask). Consulted only when has_seq_kv_lens; a # 1-elem dummy is passed otherwise. @@ -390,7 +399,6 @@ def _sdpa_kernel( K_view = cutlass.make_array_view(K) V_view = cutlass.make_array_view(V) O_view = cutlass.make_array_view(O) - LSE_view = cutlass.make_array_view(LSE) # Shapes (B, S, H, D). GQA/MQA: H_q (Q heads) may exceed H_kv (K/V # heads) when H_q % H_kv == 0. ``heads_per_kv`` Q-heads share each @@ -471,11 +479,16 @@ def _sdpa_kernel( O_BASE = q_seq_abs64 * O_ROW_STRIDE_E.to(cutlass.Int64) + cutlass.Int64(head_idx) * d_v64 # LSE: dense [B, H_q, SQ]; THD packed [1, H_q, T] — element-contiguous along # the seq axis. THD uses the packed seq position (cu_q[b] + q_row_base). - if cutlass.const_expr(THD_VARLEN): - T_lse_i64 = cutlass.Int64(LSE.shape[2]) - LSE_BASE = cutlass.Int64(head_idx) * T_lse_i64 + q_seq_origin + cutlass.Int64(q_row_base) - else: - LSE_BASE = cutlass.Int64(batch_idx) * H_i64 * SQ_i64 + cutlass.Int64(head_idx) * SQ_i64 + cutlass.Int64(q_row_base) + # None-specialized: no Stats output ⇒ view/base/pointer are compiled out + # together with the epilogue store below. + if cutlass.const_expr(LSE is not None): + LSE_view = cutlass.make_array_view(LSE) + if cutlass.const_expr(THD_VARLEN): + T_lse_i64 = cutlass.Int64(LSE.shape[2]) + LSE_BASE = cutlass.Int64(head_idx) * T_lse_i64 + q_seq_origin + cutlass.Int64(q_row_base) + else: + LSE_BASE = cutlass.Int64(batch_idx) * H_i64 * SQ_i64 + cutlass.Int64(head_idx) * SQ_i64 + cutlass.Int64(q_row_base) + lse_gmem = LSE_view.data_ptr() + LSE_BASE # K/V offsets parameterised on kv_row_base (variable in mainloop). Uses # H_kv and kv_head_idx (Q-heads sharing a KV head land at the same K/V @@ -489,7 +502,6 @@ def _sdpa_kernel( q_gmem = Q_view.data_ptr() + Q_BASE o_gmem = O_view.data_ptr() + O_BASE - lse_gmem = LSE_view.data_ptr() + LSE_BASE k_gmem_batch_head = K_view.data_ptr() + K_BATCH_OFF + K_HEAD_OFF v_gmem_batch_head = V_view.data_ptr() + V_BATCH_OFF + V_HEAD_OFF # Per-iter element advance: ``+ TILE_N rows`` on a Float16-typed ptr. @@ -1372,50 +1384,53 @@ def _mask_term(col, q_abs, q_mw): # [B,SQ,H,D] tensor; THD must NOT write rows past this sequence (they # belong to the next packed sequence) → bound by the per-batch eff_sq. sq_store_bound = eff_sq if cutlass.const_expr(THD_VARLEN) else sq_runtime - for m_block in cutlass.range_constexpr(m_blocks): - block_warp_row = warp_m_base + m_block * 16 - block_row_top = block_warp_row + g_lane - block_row_bot = block_warp_row + g_lane + 8 - row_state_lo = m_block * 2 - row_state_hi = m_block * 2 + 1 - lse_top = LN2 * (row_max[row_state_lo] + cute.math.log2(row_sum[row_state_lo])) - lse_bot = LN2 * (row_max[row_state_hi] + cute.math.log2(row_sum[row_state_hi])) - # Dense PADDED (per-batch seq_len_q): padded query rows q >= eff_sq are - # within this batch's [B,SQ,..] slice but don't exist → lse = -inf - # (log-sum-exp of a fully-masked row; matches cuDNN >= 9.14, and what - # test_mhas_v2 expects for stats on padded rows). Applied BEFORE the - # store-path split below: the is_even_mn fast path stores every row, so - # trimming only in the predicated path left finite LSE on padded rows - # whenever SQ is tile-aligned. THD is bounded by sq_store_bound==eff_sq - # instead (must NOT write the next packed seq's rows), so this select is - # gated on has_seq_len_q only. The SM80 bprop reads this lse and masks - # P=0 for padded rows via a select (inf->0, no NaN), so the -inf is safe - # downstream. - if cutlass.const_expr(has_seq_len_q): - _ninf = cutlass.Float32(float("-inf")) - trim_top = q_row_base_i32 + block_row_top - trim_bot = q_row_base_i32 + block_row_bot - lse_top = cutlass.Float32(arith.select((trim_top < eff_sq).ir_value(), lse_top.ir_value(), _ninf.ir_value())) - lse_bot = cutlass.Float32(arith.select((trim_bot < eff_sq).ir_value(), lse_bot.ir_value(), _ninf.ir_value())) - lse_top_ptr = lse_gmem + cutlass.Int64(block_row_top) - lse_bot_ptr = lse_gmem + cutlass.Int64(block_row_bot) - - # Row predication for OOB Q-rows when ~is_even_mn. All 4 lanes of a - # threadquad share the same row → predicate is uniform within the - # threadquad, so the if-branch traces identically across them. - def _stf(ptr, val): - cutlass.Array(ptr, (1,), dtype=cutlass.Float32)[0] = val - - if cutlass.const_expr(is_even_mn): - _stf(lse_top_ptr, lse_top) - _stf(lse_bot_ptr, lse_bot) - else: - top_abs = q_row_base_i32 + block_row_top - bot_abs = q_row_base_i32 + block_row_bot - if top_abs < sq_store_bound: + # LSE output is None-specialized: no Stats output -> the whole block + # (compute + stores) is compiled out. + if cutlass.const_expr(LSE is not None): + for m_block in cutlass.range_constexpr(m_blocks): + block_warp_row = warp_m_base + m_block * 16 + block_row_top = block_warp_row + g_lane + block_row_bot = block_warp_row + g_lane + 8 + row_state_lo = m_block * 2 + row_state_hi = m_block * 2 + 1 + lse_top = LN2 * (row_max[row_state_lo] + cute.math.log2(row_sum[row_state_lo])) + lse_bot = LN2 * (row_max[row_state_hi] + cute.math.log2(row_sum[row_state_hi])) + # Dense PADDED (per-batch seq_len_q): padded query rows q >= eff_sq are + # within this batch's [B,SQ,..] slice but don't exist → lse = -inf + # (log-sum-exp of a fully-masked row; matches cuDNN >= 9.14, and what + # test_mhas_v2 expects for stats on padded rows). Applied BEFORE the + # store-path split below: the is_even_mn fast path stores every row, so + # trimming only in the predicated path left finite LSE on padded rows + # whenever SQ is tile-aligned. THD is bounded by sq_store_bound==eff_sq + # instead (must NOT write the next packed seq's rows), so this select is + # gated on has_seq_len_q only. The SM80 bprop reads this lse and masks + # P=0 for padded rows via a select (inf->0, no NaN), so the -inf is safe + # downstream. + if cutlass.const_expr(has_seq_len_q): + _ninf = cutlass.Float32(float("-inf")) + trim_top = q_row_base_i32 + block_row_top + trim_bot = q_row_base_i32 + block_row_bot + lse_top = cutlass.Float32(arith.select((trim_top < eff_sq).ir_value(), lse_top.ir_value(), _ninf.ir_value())) + lse_bot = cutlass.Float32(arith.select((trim_bot < eff_sq).ir_value(), lse_bot.ir_value(), _ninf.ir_value())) + lse_top_ptr = lse_gmem + cutlass.Int64(block_row_top) + lse_bot_ptr = lse_gmem + cutlass.Int64(block_row_bot) + + # Row predication for OOB Q-rows when ~is_even_mn. All 4 lanes of a + # threadquad share the same row → predicate is uniform within the + # threadquad, so the if-branch traces identically across them. + def _stf(ptr, val): + cutlass.Array(ptr, (1,), dtype=cutlass.Float32)[0] = val + + if cutlass.const_expr(is_even_mn): _stf(lse_top_ptr, lse_top) - if bot_abs < sq_store_bound: _stf(lse_bot_ptr, lse_bot) + else: + top_abs = q_row_base_i32 + block_row_top + bot_abs = q_row_base_i32 + block_row_bot + if top_abs < sq_store_bound: + _stf(lse_top_ptr, lse_top) + if bot_abs < sq_store_bound: + _stf(lse_bot_ptr, lse_bot) # ---- Final normalization + SMEM-staged STG.128 epilogue -------------- # Naïve scalar per-lane STG (the previous epilogue) emits 2-fp16 stores @@ -1550,7 +1565,7 @@ def _sdpa_host( K: cute.Tensor, V: cute.Tensor, O: cute.Tensor, - LSE: cute.Tensor, + LSE: Optional[cute.Tensor], seq_kv_lens: cute.Tensor, seq_len_q: cute.Tensor, sinks: cute.Tensor, @@ -1654,121 +1669,144 @@ def _sdpa_host( # --------------------------------------------------------------------------- -# Compile cache. +# Per-shape compile cache. # --------------------------------------------------------------------------- # ``cute.compile`` is expensive (trace + MLIR + NVVM + PTX → SASS, ~1-2 s on -# A100). Without a per-process cache, every ``forward()`` call re-traces -# even for identical (shape, dtype, flag) tuples — costing 1+ s/call on -# top of the actual kernel runtime. ``@lru_cache`` keyed on the cache- -# significant tuple gives us O(1) lookup; the *value* is the compiled fn -# handle which is reusable across calls. +# A100). The FEATURE / config axes are module identity — one loaded template +# module per ``TemplateParams`` via ``frost.template_loader`` — so this cache +# covers the remaining SHAPE axes only. Every key component is PLAN-TIME +# data (AGENTS.md Hard Rule 4): under ``PARAMS.thd_varlen`` the packed token +# totals compile DYNAMIC (``cute.sym_int``) and are never part of the key — +# callers pass ``sq = skv = 0`` there (a stray runtime total must not be +# passed: it would only mint a redundant cache entry for the same artifact). @lru_cache(maxsize=None) -def _compile_cached( - B: int, - H: int, - H_kv: int, - SQ: int, - SKV: int, - D: int, - tile_m: int, - num_warps: int, - tile_n: int, - d_qk: int, - d_v: int, - io_is_bf16: bool, - is_even_mn: bool, - is_even_k: bool, - mask_flags: int, - swa_window: int, - causal_bottom_right: bool, - has_seq_kv_lens: bool, - has_seq_len_q: bool, - has_sink: bool, - has_bias: bool, - bias_is_fp32: bool, - THD_VARLEN: bool, - n_batch_logical: int, - has_rope: bool, - rope_max_s: int, - sched_policy: int, - sched_l2_bytes: int, +def compile( # noqa: A001 — the template contract's entry point (matches the SM100 kernels) + b: int, + h: int, + h_kv: int, + sq: int, + skv: int, + d: int, + swa_window: int = 0, + rope_max_s: int = 0, + n_batch_logical: int = 0, ): - """Compile (or return cached) ``_sdpa_host`` for the given config. - - Uses ``cute.runtime.make_fake_compact_tensor`` for shape-stable trace - inputs so the cached binary is reusable across different torch tensor - instances with the same shape signature. + """Compile (or fetch) this template specialization for one shape. + + ``d`` is the ACTUAL Q/K head dim and may be < ``PARAMS.d_qk``: the + SMEM/reg tile stays ``PARAMS.d_qk`` wide and the missing columns + zero-fill via cp.async predication. V/O are always exactly + ``PARAMS.d_v`` wide. ``swa_window`` is the left-window width W (keep + k in [q-W, q]) — plan-time graph data, baked exactly as the old + ``forward()`` did. Dense evenness is derived here the same way the old + entry point derived it: ``is_even_mn = (sq % tile_m == 0) and + (skv % tile_n == 0)``; ``is_even_k = (d == PARAMS.d_qk)``. + + THD (``PARAMS.thd_varlen``): q/k/v/o are packed ``[1, T, H, D]`` and the + LSE is packed ``[1, H, T]``; the token extents compile DYNAMIC — one + ``cute.sym_int`` symbol shared by the Q/O/LSE group and one for K/V — so + one artifact re-binds any packed totals (issue #604). Pass ``b = 1``, + ``sq = skv = 0``; ``n_batch_logical`` (the logical sequence count) sizes + the ``cu_seqlens`` ABI and IS plan-time. THD always takes the + predicated-store path (``is_even_mn = False``) and the over-provisioned + SCHED_DEFAULT grid, driven by the runtime ``thd_q_tiles``/``thd_n_batch`` + launch arguments. + + ``PARAMS.has_lse = False`` compiles the LSE store out entirely (the LSE + argument is None-specialized) — no buffer and no dummy at any level. """ - # Q and K share the QK head dim (= D, possibly < d_qk when ~is_even_k); - # V and O follow d_v (DSv3: d_qk != d_v). - io_dtype = cutlass.BFloat16 if io_is_bf16 else cutlass.Float16 + p = PARAMS + if p.thd_varlen and p.has_bias: + raise ValueError("sm80: bias + THD is not supported (varlen has no single [1,H,SQ,SKV] bias shape)") + io_dtype = cutlass.BFloat16 if p.io_bf16 else cutlass.Float16 + mask_flags = (MASK_CAUSAL if p.is_causal else MASK_NONE) | (MASK_SWA if p.has_swa else 0) + sched_l2_bytes = p.sched_l2_mib * 1024 * 1024 + is_even_k = d == p.d_qk + if p.thd_varlen: + # One symbol per ragged group: Q/O (and the LSE's T axis) share t_q, + # K/V share t_kv — a new packed total re-binds the same artifact. + is_even_mn = False + t_q = cute.sym_int(divisibility=1) + t_kv = cute.sym_int(divisibility=1) + _b, _sq, _skv = 1, t_q, t_kv + else: + is_even_mn = (sq % p.tile_m == 0) and (skv % p.tile_n == 0) + _b, _sq, _skv = b, sq, skv + + # Q and K share the QK head dim (= d, possibly < PARAMS.d_qk when + # ~is_even_k); V and O follow PARAMS.d_v (DSv3: d_qk != d_v). fake_q = cute.runtime.make_fake_compact_tensor( io_dtype, - (B, SQ, H, D), + (_b, _sq, h, d), stride_order=(3, 2, 1, 0), assumed_align=16, ) fake_k = cute.runtime.make_fake_compact_tensor( io_dtype, - (B, SKV, H_kv, D), + (_b, _skv, h_kv, d), stride_order=(3, 2, 1, 0), assumed_align=16, ) fake_v = cute.runtime.make_fake_compact_tensor( io_dtype, - (B, SKV, H_kv, d_v), + (_b, _skv, h_kv, p.d_v), stride_order=(3, 2, 1, 0), assumed_align=16, ) fake_o = cute.runtime.make_fake_compact_tensor( io_dtype, - (B, SQ, H, d_v), + (_b, _sq, h, p.d_v), stride_order=(3, 2, 1, 0), assumed_align=16, ) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (B, H, SQ), - stride_order=(2, 1, 0), - assumed_align=16, - ) + # LSE: dense [B, H, SQ]; THD packed [1, H, T] (shares the Q/O token + # symbol, which is exactly what the kernel's LSE.shape[2] read needs). + if p.has_lse: + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (_b, h, _sq), + stride_order=(2, 1, 0), + assumed_align=16, + ) + else: + fake_lse = None # Per-batch KV / Q lengths [B] int32 (or a 1-elem dummy when unused). fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor( cutlass.Int32, - (B if has_seq_kv_lens else 1,), + (b if p.has_seq_kv_lens else 1,), stride_order=(0,), assumed_align=4, ) fake_seq_len_q = cute.runtime.make_fake_compact_tensor( cutlass.Int32, - (B if has_seq_len_q else 1,), + (b if p.has_seq_q_lens else 1,), stride_order=(0,), assumed_align=4, ) # Per-Q-head sink logit [H] fp32 (log2 units) (or a 1-elem dummy when unused). fake_sinks = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (H if has_sink else 1,), + (h if p.has_sink else 1,), stride_order=(0,), assumed_align=4, ) # Additive bias [1, H, SQ, SKV] in io_dtype or fp32 (or a 1-elem dummy). - bias_io_dtype = cutlass.Float32 if bias_is_fp32 else io_dtype + bias_io_dtype = cutlass.Float32 if p.bias_is_fp32 else io_dtype fake_bias = cute.runtime.make_fake_compact_tensor( bias_io_dtype, - ((1, H, SQ, SKV) if has_bias else (1,)), - stride_order=((3, 2, 1, 0) if has_bias else (0,)), + ((1, h, sq, skv) if p.has_bias else (1,)), + stride_order=((3, 2, 1, 0) if p.has_bias else (0,)), assumed_align=16, ) # THD cumulative seqlens [B_logical + 1] int32 (or 1-elem dummies). - _cu_len = (n_batch_logical + 1) if THD_VARLEN else 1 + _cu_len = (n_batch_logical + 1) if p.thd_varlen else 1 fake_cu_q = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (_cu_len,), stride_order=(0,), assumed_align=4) fake_cu_k = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (_cu_len,), stride_order=(0,), assumed_align=4) # RoPE (cos, sin) table [max_s, d_qk//2, 2] fp32 (or a 1-elem dummy). fake_rope_cs = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - ((rope_max_s, d_qk // 2, 2) if has_rope else (1,)), - stride_order=((2, 1, 0) if has_rope else (0,)), + ((rope_max_s, p.d_qk // 2, 2) if p.has_rope else (1,)), + stride_order=((2, 1, 0) if p.has_rope else (0,)), assumed_align=16, ) fake_n_kv_tiles = cutlass.Int32(0) @@ -1797,25 +1835,25 @@ def _compile_cached( fake_cu_q, fake_cu_k, fake_rope_cs, - tile_m, - num_warps, - tile_n, - d_qk, - d_v, + p.tile_m, + p.num_warps, + p.tile_n, + p.d_qk, + p.d_v, io_dtype, is_even_mn, is_even_k, mask_flags, swa_window, - causal_bottom_right, - has_seq_kv_lens, - has_seq_len_q, - has_sink, - has_bias, - bias_is_fp32, - THD_VARLEN, - has_rope, - sched_policy, + p.causal_bottom_right, + p.has_seq_kv_lens, + p.has_seq_q_lens, + p.has_sink, + p.has_bias, + p.bias_is_fp32, + p.thd_varlen, + p.has_rope, + p.sched_policy, sched_l2_bytes, fake_n_kv_tiles, fake_scale, @@ -1829,393 +1867,3 @@ def _compile_cached( fake_stream, options="--enable-tvm-ffi", ) - - -# --------------------------------------------------------------------------- -# Python entry point. -# --------------------------------------------------------------------------- -def forward( - Q: torch.Tensor, # [B, SQ, H, D] fp16 - K: torch.Tensor, # [B, SKV, H, D] fp16 - V: torch.Tensor, # [B, SKV, H, D] fp16 - scale: Optional[float] = None, - return_lse: bool = False, - *, - tile_m: int = DEFAULT_TILE_M, - num_warps: int = DEFAULT_NUM_WARPS, - tile_n: int = DEFAULT_TILE_N, - d_qk: int = DEFAULT_D_QK, - d_v: int = DEFAULT_D_V, - mask: str = "none", # "none" / "causal" / "swa" - swa_window: int = 0, # window width when mask == "swa" - right_bound: int = 0, # extra causal right band: admit up to this - # many future tokens (k <= q + right_bound); - # 0 = plain causal. Requires MASK_CAUSAL. - causal_bottom_right: bool = False, # align causal diagonal to bottom-right - # (k <= q + SKV-SQ); only with mask=="causal" - seq_kv_lens: Optional[torch.Tensor] = None, # [B] int32 per-batch effective KV - # length → padded mask (cols >= len -inf'd) - seq_len_q: Optional[torch.Tensor] = None, # [B] int32 per-batch effective Q length - # → bottom-right diagonal under padding - # (br_base = eff_skv - eff_sq); only used - # with causal_bottom_right. - sinks: Optional[torch.Tensor] = None, # [H] fp32 per-Q-head sink logits in - # SCALED-logit (natural) units — joins the - # softmax denominator only (V_sink = 0). - # forward multiplies by log2(e). None → - # no sink. - bias: Optional[torch.Tensor] = None, # additive attention bias, broadcast over - # batch: shape [1, H, SQ, SKV] (or - # [B, H, SQ, SKV] — only slice 0 is read). - # dtype defaults to QKV (fp16/bf16); fp32 - # also accepted. None → no bias. When - # bias is set and tile_m/num_warps are at - # their defaults, the bias-enabled path - # defaults to tile_m=64/num_warps=4 (lower - # SMEM → 2 CTAs/SM to hide the bias LDGs). - cu_seqlens_q: Optional[torch.Tensor] = None, # [B+1] int32 cumulative Q seqlens. - # When set → THD/varlen mode: Q/K/V/O are - # PACKED [1, T, H, D] and cu_seqlens_* - # define each sequence's span. Device or - # host int32; copied to device int32. - cu_seqlens_k: Optional[torch.Tensor] = None, # [B+1] int32 cumulative KV seqlens. - max_s_q: Optional[int] = None, # THD: max sequence Q length (for the - # over-provisioned grid). If None it is - # computed from cu_seqlens_q (one d2h sync — - # pass it explicitly for graph capture). - sched: str = "auto", # "auto" (default→none / lpt→mask) / "default" - # / "lpt" / "lpt_l2" - sched_l2_mib: int = 32, # L2 budget (MiB) for "lpt_l2". A100 = 40 MiB - # physical; ~32 MiB usable after texture / RO. - rope_freqs: Optional[torch.Tensor] = None, # RoPE angles, shape [max_s, 1, 1, d_qk] - # (cuDNN graph.rope freqs convention) or - # [max_s, d_qk]. Only the first d_qk//2 cols - # (the angles; second half is zeros) are used. - # cos/sin are precomputed on the HOST with - # torch (full-range, bit-matches the reference - # freqs.cos()/.sin()) and the half-split - # rotate_half is applied to Q AND K in-kernel. - # Dense-only (no THD); max_s must cover - # max(SQ, SKV). None → no RoPE. - out_o: Optional[torch.Tensor] = None, # [B, SQ, H, d_v] caller-provided output - # buffer (Q dtype, contiguous); None → - # allocate. Dense-only (THD allocates). - out_lse: Optional[torch.Tensor] = None, # [B, H, SQ] fp32 contiguous LSE buffer; - # None → allocate (used when return_lse). -): - """Run SM80 SDPA prefill for an MHA-shaped (B, S, H, D) Q/K/V triple. - - Returns a [B, SQ, H, D] fp16 output tensor. When ``return_lse`` is - true, returns ``(O, LSE)`` with LSE shape ``[B, H, SQ]`` fp32 in - natural-log (matches ``torch.logsumexp(scale*Q@K.T, dim=-1)``). - - Configurable kernel variants via ``(tile_m, num_warps)``: - * ``(64, 4)`` — default; 128 threads / CTA, M_BLOCKS = 1. Fits - 2 CTAs / SM implicitly (≤128 regs/thread, ≤48 KiB SMEM/CTA). - * ``(128, 4)`` — same 128 threads, but each warp owns 2 m16n8k16 - row-blocks (M_BLOCKS=2). Spills at this register footprint on - SM80 — kept only as a path for future GPT-OSS-style configs. - * ``(128, 8)`` — 256 threads / CTA, 1 m16n8k16 block per warp - (M_BLOCKS=1, no register-pressure regression). Same K/V load - cost amortized over 2× Q rows. Targets 1 CTA / SM. - """ - assert Q.dtype in (torch.float16, torch.bfloat16), f"Q dtype must be float16 or bfloat16 (got {Q.dtype})" - assert K.dtype == Q.dtype and V.dtype == Q.dtype, f"K/V dtype must match Q ({Q.dtype}); got K={K.dtype} V={V.dtype}" - io_is_bf16 = Q.dtype == torch.bfloat16 - assert Q.is_cuda and K.is_cuda and V.is_cuda - B, SQ, H, D = Q.shape - _, SKV, Hk, D_K = K.shape - _, _, _, D_V_actual = V.shape - # GQA/MQA: H_q (= H) must be a multiple of H_kv (= Hk). MHA is H == Hk. - assert H % Hk == 0, f"H_q ({H}) must be a multiple of H_kv ({Hk}) for GQA/MQA" - assert K.shape[2] == V.shape[2], f"K H ({K.shape[2]}) must equal V H ({V.shape[2]})" - # Q and K share the QK head dim; V has its own d_v (asymmetric on DSv3: - # d_qk=192, d_v=128). D may be < d_qk — the SMEM/reg tile is sized for - # d_qk and the missing cols are zero-padded via cp.async predication - # when ~is_even_k. - assert D == D_K, f"Q D ({D}) must equal K D ({D_K})" - assert D_V_actual == d_v, f"V D ({D_V_actual}) must equal compile-time d_v={d_v}" - assert D <= d_qk, f"D ({D}) must be <= compile-time d_qk={d_qk}" - assert d_qk % 16 == 0, f"d_qk ({d_qk}) must be a multiple of 16 (m16n8k16 K)" - assert d_v % 8 == 0, f"d_v ({d_v}) must be a multiple of 8 (SV n_frags)" - assert d_v % 16 == 0, f"d_v ({d_v}) must be a multiple of 16 (cp.async + STG.128 epilogue)" - assert D % 8 == 0, f"D ({D}) must be a multiple of 8 (cp.async chunk size)" - # Bias-enabled default tile: lower SMEM (64 Q rows / 128 threads) so 2 CTAs - # fit per SM, helping hide the per-iter bias LDGs. Only when the caller - # left tile_m / num_warps at their module defaults. - has_bias = bias is not None - if has_bias and tile_m == DEFAULT_TILE_M and num_warps == DEFAULT_NUM_WARPS: - tile_m, num_warps = 64, 4 - assert tile_m % (num_warps * 16) == 0, f"tile_m={tile_m} must be a multiple of num_warps*16={num_warps*16}" - assert tile_n in (64, 128), f"tile_n must be 64 or 128 (got {tile_n})" - assert mask in ("none", "causal", "swa", "causal_swa"), f"mask: 'none' | 'causal' | 'swa' | 'causal_swa' (got {mask!r})" - assert (not causal_bottom_right) or mask in ("causal", "causal_swa", "swa"), ( - "causal_bottom_right=True requires mask in causal/causal_swa/swa " "(BR shifts the causal upper and/or SWA lower bound to the corner)" - ) - assert right_bound >= 0, f"right_bound must be >= 0 (got {right_bound})" - assert right_bound == 0 or mask in ("causal", "causal_swa"), "right_bound>0 (causal right band) requires a causal mask" - assert sched in ("auto", "default", "lpt", "lpt_l2"), f"sched: 'auto' | 'default' | 'lpt' | 'lpt_l2' (got {sched!r})" - assert sched_l2_mib > 0, f"sched_l2_mib must be > 0 (got {sched_l2_mib})" - - is_even_mn = (SQ % tile_m == 0) and (SKV % tile_n == 0) - is_even_k = D == d_qk - mask_flags = MASK_NONE - if mask == "causal": - mask_flags |= MASK_CAUSAL - elif mask == "swa": - assert swa_window >= 0, "mask='swa' requires swa_window >= 0 (0 = 1-token window, keep k>=q)" - mask_flags |= MASK_SWA - elif mask == "causal_swa": - # Causal sliding window [q-W, q]: causal upper bound + SWA lower bound. - # The kernel body composes both bits (independent tile-prune trims + - # OR'd per-element terms); no body change needed. - assert swa_window >= 0, "mask='causal_swa' requires swa_window >= 0 (0 = diagonal-only)" - mask_flags |= MASK_CAUSAL | MASK_SWA - - if sched == "auto": - # MASK_NONE → SCHED_DEFAULT (preserves L2 reuse); causal / SWA → LPT. - sched_policy = SCHED_DEFAULT if mask_flags == MASK_NONE else SCHED_LPT - elif sched == "default": - sched_policy = SCHED_DEFAULT - elif sched == "lpt": - sched_policy = SCHED_LPT - else: # "lpt_l2" - sched_policy = SCHED_LPT_L2 - sched_l2_bytes = int(sched_l2_mib) * 1024 * 1024 - - # ---- THD/varlen setup ------------------------------------------------- - # Packed [1,T,H,D] Q/K/V/O + cu_seqlens. Forces the predicated-store path - # (is_even_mn=False) and a SCHED_DEFAULT over-provisioned 3-D grid - # (ceil(max_s_q/tile_m), H, B_logical); tiles past a sequence early-out. - THD_VARLEN = cu_seqlens_q is not None - if THD_VARLEN: - assert cu_seqlens_k is not None, "THD needs both cu_seqlens_q and cu_seqlens_k" - assert B == 1, f"THD: Q/K/V must be packed [1, T, H, D] (got batch dim {B})" - assert not has_bias, "bias + THD not supported (varlen has no single " "[1,H,SQ,SKV] bias shape)" - n_batch_logical = int(cu_seqlens_q.numel()) - 1 - assert n_batch_logical >= 1, "cu_seqlens_q must have >= 2 entries" - cu_q_t = cu_seqlens_q.to(dtype=torch.int32, device=Q.device).contiguous() - cu_k_t = cu_seqlens_k.to(dtype=torch.int32, device=Q.device).contiguous() - if max_s_q is None: - _d = cu_q_t[1:] - cu_q_t[:-1] - max_s_q = int(_d.max().item()) # one d2h sync (non-graph) - thd_q_tiles_v = (int(max_s_q) + tile_m - 1) // tile_m - is_even_mn = False - sched_policy = SCHED_DEFAULT - else: - n_batch_logical = 1 - cu_q_t = torch.ones(1, dtype=torch.int32, device=Q.device) - cu_k_t = torch.ones(1, dtype=torch.int32, device=Q.device) - thd_q_tiles_v = 0 - - if scale is None: - scale = 1.0 / (D**0.5) - import math - - scale_log2 = scale * math.log2(math.e) - - # Allocate output / LSE sized to the actual (possibly uneven) shapes — - # the kernel writes the full row x d_v range via STG (rows predicated on - # sq_store_bound only). Output dim follows V (= d_v), which can differ - # from Q's d_qk (DSv3). Caller-bound out_* buffers skip the allocation. - # The zero-fill below is DEFENSIVE, not load-bearing: the dense epilogue - # stores every in-bounds row unconditionally (zero-KV-iteration tiles - # route through the store with row_sum=0, trimmed rows are written - # explicitly with O=0 / LSE=-inf, and THD never receives bound outputs). - # Kept so a bound buffer can never surface uninitialized memory if a - # future feature path skips rows. - _needs_zero_init = (seq_len_q is not None) or (seq_kv_lens is not None) - if out_o is None: - out_o = torch.zeros(B, SQ, H, d_v, dtype=Q.dtype, device=Q.device) - else: - assert out_o.shape == (B, SQ, H, d_v) and out_o.dtype == Q.dtype and out_o.is_contiguous(), ( - f"out_o must be a contiguous [{B}, {SQ}, {H}, {d_v}] {Q.dtype} tensor; " f"got shape {tuple(out_o.shape)} dtype {out_o.dtype}" - ) - if _needs_zero_init: - out_o.zero_() - if out_lse is None: - LSE = torch.zeros(B, H, SQ, dtype=torch.float32, device=Q.device) - else: - assert out_lse.shape == (B, H, SQ) and out_lse.dtype == torch.float32 and out_lse.is_contiguous(), ( - f"out_lse must be a contiguous [{B}, {H}, {SQ}] fp32 tensor; " f"got shape {tuple(out_lse.shape)} dtype {out_lse.dtype}" - ) - LSE = out_lse - if _needs_zero_init: - LSE.zero_() - - # Per-batch padded mask: [B] int32 KV lengths (or a 1-elem dummy when - # unused — the kernel never reads it under has_seq_kv_lens=False). - has_seq_kv_lens = seq_kv_lens is not None - if has_seq_kv_lens: - assert seq_kv_lens.shape == (B,), f"seq_kv_lens must be shape ({B},); got {tuple(seq_kv_lens.shape)}" - seq_kv_lens_t = seq_kv_lens.to(dtype=torch.int32, device=Q.device).contiguous() - else: - seq_kv_lens_t = torch.ones(1, dtype=torch.int32, device=Q.device) - - has_seq_len_q = seq_len_q is not None - if has_seq_len_q: - assert seq_len_q.shape == (B,), f"seq_len_q must be shape ({B},); got {tuple(seq_len_q.shape)}" - seq_len_q_t = seq_len_q.to(dtype=torch.int32, device=Q.device).contiguous() - else: - seq_len_q_t = torch.ones(1, dtype=torch.int32, device=Q.device) - - # Per-Q-head sink logits → log2 units (= sink * log2(e)) so they share the - # kernel's log2-of-scaled-logit max domain. - has_sink = sinks is not None - if has_sink: - assert sinks.shape == (H,), f"sinks must be shape ({H},); got {tuple(sinks.shape)}" - sinks_t = (sinks.to(dtype=torch.float32, device=Q.device) * math.log2(math.e)).contiguous() - else: - sinks_t = torch.ones(1, dtype=torch.float32, device=Q.device) - - # Additive bias: dtype defaults to QKV (fp16/bf16); fp32 also accepted. - # Broadcast over batch — keep only the [1, H, SQ, SKV] slice (the kernel - # reads slice 0). inv_scale = 1/scale is folded in-kernel. - inv_scale = 1.0 / float(scale) - if has_bias: - assert bias.dtype in (Q.dtype, torch.float32), f"bias dtype must be {Q.dtype} or float32 (got {bias.dtype})" - assert tuple(bias.shape[-3:]) == (H, SQ, SKV), f"bias trailing dims must be (H={H}, SQ={SQ}, SKV={SKV}); " f"got {tuple(bias.shape)}" - bias_is_fp32 = bias.dtype == torch.float32 - bias_t = bias[:1].contiguous() if bias.shape[0] != 1 else bias.contiguous() - else: - bias_is_fp32 = False - bias_t = torch.ones(1, dtype=Q.dtype, device=Q.device) - - # RoPE: precompute the (cos, sin) table on the host with torch's full-range - # cos/sin (so it bit-matches the reference's freqs.cos()/.sin() — no in-kernel - # MUFU range-reduction error for large angles). freqs holds the ANGLES in - # [..., :d2]; the second half is zeros (unused). Packed as [max_s, d2, 2] fp32. - has_rope = rope_freqs is not None - if has_rope: - assert not THD_VARLEN, "RoPE is dense-only (no THD/varlen) on SM80" - assert D == d_qk, f"RoPE requires D == d_qk (got D={D}, d_qk={d_qk}); the partial " "rope_dim < head_dim (MLA) path is not implemented" - d2 = d_qk // 2 - rf = rope_freqs.to(dtype=torch.float32, device=Q.device) - rf = rf.reshape(rf.shape[0], -1) # [max_s, d_qk] (or wider) - rope_max_s = rf.shape[0] - assert rf.shape[1] >= d2, f"rope_freqs last dim ({rf.shape[1]}) must be >= d_qk//2 ({d2})" - assert rope_max_s >= max(SQ, SKV), f"rope_freqs max_s ({rope_max_s}) must cover max(SQ={SQ}, SKV={SKV})" - angles = rf[:, :d2] # [max_s, d2] - rope_cs_t = torch.stack([angles.cos(), angles.sin()], dim=-1).contiguous() - else: - rope_max_s = 1 - rope_cs_t = torch.ones(1, dtype=torch.float32, device=Q.device) - - cQ = from_dlpack(Q) - cK = from_dlpack(K) - cV = from_dlpack(V) - cO = from_dlpack(out_o) - cLSE = from_dlpack(LSE) - cSEQK = from_dlpack(seq_kv_lens_t) - cSEQQ = from_dlpack(seq_len_q_t) - cSINKS = from_dlpack(sinks_t) - cBIAS = from_dlpack(bias_t) - cCUQ = from_dlpack(cu_q_t) - cCUK = from_dlpack(cu_k_t) - cROPE = from_dlpack(rope_cs_t) - torch_stream = torch.cuda.current_stream() - stream = cuda.CUstream(torch_stream.cuda_stream) - - # Host computes ``round_up(SKV / TILE_N)`` and passes as a runtime - # Int32 — the kernel uses ``cutlass.range(..., unroll=1)`` over it so - # the body only traces ONCE regardless of how many KV tiles SKV - # contains. Slashes ``cute.compile`` time from ~60 s (constexpr - # unroll over 128 iters at SQ=8192) to a few seconds. - n_kv_tiles = cutlass.Int32((SKV + tile_n - 1) // tile_n) - sq_rt = cutlass.Int32(SQ) - skv_rt = cutlass.Int32(SKV) - d_rt = cutlass.Int32(D) - # ``_compile_cached`` returns the same compiled fn on cache hit (~µs); - # on cache miss it traces + lowers (~1 s). Cache key is the full - # Constexpr tuple — different (mask, shape) combos compile separately. - fn = _compile_cached( - B, - H, - Hk, - SQ, - SKV, - D, - tile_m, - num_warps, - tile_n, - d_qk, - d_v, - io_is_bf16, - is_even_mn, - is_even_k, - mask_flags, - swa_window, - bool(causal_bottom_right), - bool(has_seq_kv_lens), - bool(has_seq_len_q), - bool(has_sink), - bool(has_bias), - bool(bias_is_fp32), - bool(THD_VARLEN), - int(n_batch_logical), - bool(has_rope), - int(rope_max_s), - sched_policy, - sched_l2_bytes, - ) - fn( - cQ, - cK, - cV, - cO, - cLSE, - cSEQK, - cSEQQ, - cSINKS, - cBIAS, - cCUQ, - cCUK, - cROPE, - n_kv_tiles, - cutlass.Float32(scale_log2), - sq_rt, - skv_rt, - d_rt, - cutlass.Int32(right_bound), - cutlass.Float32(inv_scale), - cutlass.Int32(thd_q_tiles_v), - cutlass.Int32(n_batch_logical), - stream, - ) - if return_lse: - return (out_o, LSE) - return out_o - - -if __name__ == "__main__": - # Phase 3 validation: compare against torch SDPA reference. - torch.manual_seed(0) - import torch.nn.functional as F - - for dt, atol in [(torch.float16, 5e-3), (torch.bfloat16, 4e-2)]: - print(f"--- dtype={dt} ---") - for B, H, SQ, SKV in [ - (1, 1, 64, 64), - (1, 1, 64, 128), - (1, 4, 64, 256), - (2, 4, 128, 512), - ]: - Q = torch.randn(B, SQ, H, 128, dtype=dt, device="cuda") - K = torch.randn(B, SKV, H, 128, dtype=dt, device="cuda") - V = torch.randn(B, SKV, H, 128, dtype=dt, device="cuda") - out_o = forward(Q, K, V) - # Torch reference: scaled_dot_product_attention takes (B, H, S, D). - ref = F.scaled_dot_product_attention( - Q.transpose(1, 2), - K.transpose(1, 2), - V.transpose(1, 2), - ).transpose(1, 2) - diff = (out_o.float() - ref.float()).abs() - maxd = diff.max().item() - ok = maxd < atol - print( - f" B={B} H={H} SQ={SQ:4d} SKV={SKV:4d} " - f"max|out_o|={out_o.abs().max().item():.4f} " - f"max|out_o-ref|={maxd:.4f} {'PASS' if ok else 'FAIL'}" - ) - print("[sm80-sdpa] done.") diff --git a/test/python/fe_api/sdpa/test_sdpa_fwd_sm80.py b/test/python/fe_api/sdpa/test_sdpa_fwd_sm80.py index e9f259a42..5640ede0f 100644 --- a/test/python/fe_api/sdpa/test_sdpa_fwd_sm80.py +++ b/test/python/fe_api/sdpa/test_sdpa_fwd_sm80.py @@ -195,3 +195,97 @@ def test_sdpa_fwd_sm80_check_support_rejections(): api = SdpaFwdDslSm80(sample_q=q64, sample_k=q64, sample_v=q64, sample_o=q64, sample_lse=lse64) with pytest.raises(ValueError, match="dtype"): api.check_support() + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=0) +def test_sdpa_fwd_sm80_thd_off_flavor_head_dim(): + """Off-flavor THD head dim (d=96 rides the llama d=128 envelope): Q/K are + host-padded to the flavor width, and the kernel derives its Q/K row + strides from the runtime ``d`` — so the launch must pass the PADDED + width, or every row past the first reads the wrong address.""" + from cudnn.sdpa.fwd import sdpa_fwd_wrapper_sm80 + + H, D = 4, 96 + lens = [80, 33] + t = sum(lens) + cu = torch.tensor([0, lens[0], t], dtype=torch.int32, device="cuda") + q = torch.randn(1, t, H, D, dtype=torch.float16, device="cuda") + k = torch.randn(1, t, H, D, dtype=torch.float16, device="cuda") + v = torch.randn(1, t, H, D, dtype=torch.float16, device="cuda") + scale = 1.0 / math.sqrt(D) + out = sdpa_fwd_wrapper_sm80( + q, + k, + v, + is_causal=True, + scale_softmax=scale, + cum_seqlen_q_tensor=cu, + cum_seqlen_k_tensor=cu, + max_s_q=max(lens), + ) + o = out["o_tensor"] + for b0, b1 in zip(cu[:-1].tolist(), cu[1:].tolist()): + o_ref = _ref_sdpa( + q[0, b0:b1].permute(1, 0, 2)[None], + k[0, b0:b1].permute(1, 0, 2)[None], + v[0, b0:b1].permute(1, 0, 2)[None], + is_causal=True, + window_size=(-1, -1), + scale=scale, + ) + torch.testing.assert_close(o[0, b0:b1].permute(1, 0, 2)[None], o_ref, rtol=1e-2, atol=4e-3) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=0) +def test_sm80_thd_compile_key_plan_time_only(): + """Issue #604 regression: the packed THD token totals are RUNTIME values, + so two varlen calls with different totals must re-bind ONE compiled + artifact (the template module's per-shape lru sees a single miss) — + never mint a compile per step, which is the continuous-batching + pathology no correctness test catches.""" + from cudnn.frost import template_loader + from cudnn.sdpa.fwd import sdpa_fwd_wrapper_sm80 + + H, D = 4, 128 + + def varlen(lens): + import itertools + + t = int(sum(lens)) + cu = torch.tensor([0] + list(itertools.accumulate(lens)), dtype=torch.int32, device="cuda") + q = torch.randn(1, t, H, D, dtype=torch.float16, device="cuda") + k = torch.randn(1, t, H, D, dtype=torch.float16, device="cuda") + v = torch.randn(1, t, H, D, dtype=torch.float16, device="cuda") + return sdpa_fwd_wrapper_sm80( + q, + k, + v, + is_causal=True, + cum_seqlen_q_tensor=cu, + cum_seqlen_k_tensor=cu, + max_s_q=int(max(lens)), + ) + + def cache_totals(): + # The lru counters are session-global (earlier tests in a full run + # accumulate misses), so assert on DELTAS across our calls only. + mods = [m for (path, _params), m in template_loader._MODULES.items() if "sm80" in str(path)] + infos = [m.compile.cache_info() for m in mods if hasattr(m.compile, "cache_info")] + return sum(i.misses for i in infos), sum(i.hits for i in infos) + + varlen([96, 160]) # first call: one compile + n_modules_before = len(template_loader._MODULES) + misses_0, hits_0 = cache_totals() + varlen([128, 64, 320]) # different totals AND batch count... same artifact? + # Different logical batch counts legitimately re-specialize (the cu fake + # length is plan-time); different TOKEN TOTALS at the same batch count + # must not. + varlen([64, 192]) # same n_seqs as call 1, different totals + assert len(template_loader._MODULES) == n_modules_before, "a new template specialization was minted by runtime data" + misses_1, hits_1 = cache_totals() + # Call 2 (n_seqs=3) may legitimately re-specialize once; call 3 shares + # call 1's key (n_seqs=2, different token totals) and MUST cache-hit. + assert misses_1 - misses_0 <= 1, f"THD compile key leaked runtime data: {misses_1 - misses_0} new misses" + assert hits_1 - hits_0 >= 1, "expected a cache hit on the same-batch-count re-call"