diff --git a/python/cudnn/engines/heuristics.py b/python/cudnn/engines/heuristics.py index d649af9b2..53ee8b188 100644 --- a/python/cudnn/engines/heuristics.py +++ b/python/cudnn/engines/heuristics.py @@ -5,35 +5,39 @@ ``create_execution_plans([heur_mode.A, ...])`` collects the inputs — the parsed facts, the engine ids on offer, and the backend's own entries tagged with the -mode that produced each — and hands all of it to the graph's family. What the -family returns IS ``graph.plans``, position for position. +mode that produced each — and assembles ``graph.plans``. -One function, everything in view. Ranking is a comparison, so whoever ranks has -to see both sides: an engine cannot see its siblings, and a family that only -saw its own engines could not decide whether the backend belongs in front of -them. That is why there is no per-engine ``propose_plans`` and no second -merge step after this one — splitting the decision is what forced the previous -design to concatenate and call it ranking. +Two layers, deliberately separated: + +- The FAMILY hook is :func:`recommend`-shaped: ``(kind, facts, offered) -> + [PlanConfig]`` — pure, backend-blind, import-light. It answers one question: + which of MY engines serve these facts, with which complete knob assignments, + best first. It never sees the backend, modes, or another family. + +- PLACEMENT lives HERE (:func:`_assemble`), once for every family: python + proposals lead the backend's entries inside each mode block. That is a + standing assumption, not a measurement — an OSS engine that loses to the + backend gets fixed or pulled, not demoted; and an autotune (build ALL) pass + measures every entry regardless of order, so the order only decides the + default winner. The delegating entry, dedup, and the mode strip are all + placement bookkeeping and stay out of the families. An engine answers two questions only: can I serve this graph (``check_support``) and compile me this config (``build_plan``). -A family that declares no ``heuristics`` hook falls back to one default plan -per accepting engine, ahead of the backend's entries. Every python engine -belongs to a family — the manifest is the only way one exists — so a graph has -a family's opinion or only the backend's. +A family that declares no hook falls back to one default plan per accepting +engine, ahead of the backend's entries. Every python engine belongs to a +family — the manifest is the only way one exists — so a graph has a family's +proposals or only the backend's. """ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, List, Optional +from typing import Any, Callable, List, Optional from .base import BaseEngine, PlanConfig, decline_types -if TYPE_CHECKING: - from .._pygraph import pygraph - _LOG = logging.getLogger("cudnn.engines.heuristics") @@ -44,7 +48,7 @@ def default_modes() -> List[Any]: return [cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK] -def accepts(engine: BaseEngine, graph: "pygraph") -> bool: +def accepts(engine: BaseEngine, graph) -> bool: """Whether ``engine`` will serve ``graph``, declines being routing not error.""" try: engine.check_support(graph) @@ -54,12 +58,67 @@ def accepts(engine: BaseEngine, graph: "pygraph") -> bool: return True -def _unranked(graph: "pygraph", engines: List[BaseEngine], backend_plans: List[PlanConfig]) -> List[PlanConfig]: - """The ranking for a family with no heuristics hook: accepting engines, then the backend.""" - return [PlanConfig(e.engine_id, None) for e in engines if accepts(e, graph)] + list(backend_plans) +def _unranked(graph, engines: List[BaseEngine], backend_plans: List[PlanConfig]) -> List[PlanConfig]: + """The list for a family with no hook: accepting engines, then the backend.""" + return [PlanConfig(e.engine_id, None) for e in engines if accepts(e, graph)] + [_strip(c) for c in backend_plans] + + +def _strip(cfg: PlanConfig) -> PlanConfig: + """A final-list entry: (engine_id, knobs[, cpp_index]) — the mode tag is + assembly bookkeeping and never reaches ``graph.plans``.""" + if cfg.mode is None and cfg.cpp_index is None: + return cfg + return PlanConfig(cfg.engine_id, cfg.knobs, cpp_index=cfg.cpp_index) + +def _assemble(modes: List[Any], recommend: Callable[[str], List[PlanConfig]], backend_plans: List[PlanConfig]) -> List[PlanConfig]: + """The final ranked list: mode block by mode block in the caller's order, + python proposals leading the backend's entries inside each block. + + ``recommend(kind)`` is the family's hook already bound to (facts, offered): + ``kind`` is ``"A"`` (candidates worth timing, best first — also the answer + to B, which asks for a wider search the families have none to give) or + ``"FALLBACK"`` (the config expected to build where A's choice may not). + + An untagged backend entry is the delegating one: OSS candidates C++ holds + but never exposes as plans, so it cannot be enumerated. It belongs to no + mode, and it is NOT a pure OSS entry — Graph::build_plans tries the OSS + engine and, if that one declines, falls through to the native + engine_configs already enqueued. So it leads the BACKEND's entries but not + ours: ahead of our OPENSOURCE block it would answer an OSS-coverage + question with a native kernel. + + Asking for ``[A, FALLBACK]`` puts every tuned candidate — both sides' — + ahead of every fallback. A plan repeated across blocks keeps its first + position. Identity is (engine, knobs): cpp_index is only WHERE one backend + query put a plan, so keying on it would let one config both modes return + through as two entries — and an autotuner would build and time it twice. + """ + import cudnn -def rank(graph: "pygraph", engines: List[BaseEngine], backend_plans: List[PlanConfig], modes: Optional[List[Any]] = None) -> List[PlanConfig]: + delegating = [c for c in backend_plans if c.mode is None] + out: List[PlanConfig] = [] + for mode in modes: + if mode == cudnn.heur_mode.OPENSOURCE: + out += recommend("A") + delegating + elif mode in (cudnn.heur_mode.A, cudnn.heur_mode.B): + out += recommend("A") + delegating + [c for c in backend_plans if c.mode == mode] + elif mode == cudnn.heur_mode.FALLBACK: + out += recommend("FALLBACK") + delegating + [c for c in backend_plans if c.mode == mode] + # A delegate with no mode asked for it (the backend has engines but exposed + # no plans) would otherwise be dropped. + out += delegating + + seen, ranked = set(), [] + for cfg in out: + key = (cfg.engine_id, repr(cfg.knobs)) + if key not in seen: + seen.add(key) + ranked.append(_strip(cfg)) + return ranked + + +def rank(graph, engines: List[BaseEngine], backend_plans: List[PlanConfig], modes: Optional[List[Any]] = None) -> List[PlanConfig]: """The ranked plan list for ``graph`` — what ``create_execution_plans`` stores. ``engines`` are this graph's python candidates and ``backend_plans`` the @@ -77,11 +136,11 @@ def rank(graph: "pygraph", engines: List[BaseEngine], backend_plans: List[PlanCo facts = graph._facts_for(analyzer) if analyzer is not None else None if facts is None: # The family claims the graph by node type but its analyzer cannot - # express it. Nothing to rank its engines on; the backend serves it. - return list(backend_plans) + # express it. Nothing to recommend on; the backend serves it. + return [_strip(c) for c in backend_plans] offered = {e.name: e.engine_id for e in engines} - plans = list(recommend(modes, facts, offered, list(backend_plans))) + plans = _assemble(modes, lambda kind: list(recommend(kind, facts, offered)), backend_plans) own = set(offered.values()) for cfg in plans: from .engine_ids import is_python_engine diff --git a/python/cudnn/engines/manifest.py b/python/cudnn/engines/manifest.py index 153e4ab2b..be040eb90 100644 --- a/python/cudnn/engines/manifest.py +++ b/python/cudnn/engines/manifest.py @@ -253,7 +253,12 @@ def _resolve(family: EngineFamily, ref: Optional[Tuple[str, str]], what: str): def resolve_heuristics(family: EngineFamily): - """The family's plan-ranking callable, or None when it declares none.""" + """The family's proposal callable, or None when it declares none. + + The contract is ``recommend(kind, facts, offered) -> [PlanConfig]`` — pure + and backend-blind; placement against the backend's entries happens once + for every family in ``engines/heuristics._assemble``. + """ return _resolve(family, family.heuristics, "heuristics") diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 044bd44cd..660b49b57 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -297,6 +297,8 @@ def __init__( tile_m: Optional[int] = None, tile_n: Optional[int] = None, cga: Optional[int] = None, + split_kv: Optional[int] = None, + softmax_precision: Optional[int] = None, ) -> None: """Capture the common SDPA operation and tuning contract. @@ -357,11 +359,20 @@ def __init__( self._pertensor = bool(pertensor_fp8) self._device_cc = None # (major, minor); set in check_support # Tuning-knob choice, already validated against the engine's - # Capabilities domain by the probe (engines.mismatch). - self.sched_policy = SCHED_NATURAL if sched_policy is None else int(sched_policy) + # Capabilities domain by the probe (engines.mismatch). None means the + # caller stated NO preference: the graph path always arrives with an + # explicit value (the heuristic emits complete assignments), so a None + # here is the standalone-wrapper tier, where compile() derives the + # policy itself. An explicit value — including NATURAL — is honored + # verbatim, never re-derived. + self.sched_policy = None if sched_policy is None else int(sched_policy) self.tile_m = None if tile_m is None else int(tile_m) self.tile_n = None if tile_n is None else int(tile_n) self.cga = None if cga is None else int(cga) + self.split_kv = 1 if split_kv is None else int(split_kv) + # Framework axis: no forward kernel serves a softmax-precision choice + # yet, so anything non-None is rejected in check_support. + self.softmax_precision = softmax_precision self.batch_size: Optional[int] = None self.s_q_max: Optional[int] = None @@ -642,6 +653,41 @@ def _check_seq_lens_contract(self, seq_q_lens, seq_kv_lens) -> None: def scratch_workspace_bytes(self) -> int: """Return the per-execution scratch requirement for this implementation.""" + # -- KV-split shared helpers (SM100 + SM120 dense split paths) ----------- + + def _o_itemsize(self) -> int: + return 2 # f16 / bf16; the split path is half-precision-O only + + def _combine_dtype_tag(self) -> str: + # The combine reduces INTO the O dtype: the graph's dtype_o on the + # quantized rows (half-gated by check_support), Q's dtype elsewhere. + o_dtype = self.dtype_o if (self._fp8 and self.dtype_o is not None) else self.dtype + return "bf16" if o_dtype == torch.bfloat16 else "f16" + + def _split_partials(self, workspace, o_like, device, current_stream=None): + """The split-major (O, LSE) partial buffers, carved from the caller's + workspace when there is one and torch-allocated otherwise (standalone + use, matching what the rest of this adapter does). + + The allocation happens ON the launch stream: the caching allocator tags + a block with the stream it was allocated on, and the kernels that write + and read these buffers run on ``current_stream``. Allocating on torch's + current stream instead would leave a later free/reuse unordered against + those launches.""" + rows = self.split_kv * self.batch_size + o_shape = (rows, self.s_q_max, self.h_q, self.head_dim_v) + lse_shape = (rows, self.h_q, self.s_q_max) + if workspace is None: + with _torch_stream_context(current_stream, device): + return ( + torch.empty(o_shape, dtype=o_like.dtype, device=device), + torch.empty(lse_shape, dtype=torch.float32, device=device), + ) + carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), f"{type(self).__name__} (KV split)") + o_part = carver.take(rows * self.s_q_max * self.h_q * self.head_dim_v, o_like.dtype).view(o_shape) + lse_part = carver.take(rows * self.h_q * self.s_q_max, torch.float32).view(lse_shape) + return o_part, lse_part + @abstractmethod def execute( self, @@ -830,8 +876,8 @@ def check_support(self) -> bool: ) self.flavor = _pick_flavor(d_qk, d_v) self._value_error_if( - self.sched_policy != SCHED_NATURAL, - f"SM100 DSL SDPA only supports sched_policy={SCHED_NATURAL}", + self.sched_policy is not None and self.sched_policy not in (SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2), + f"SM100 DSL SDPA sched_policy must be NATURAL/LPT/LPT_L2 (or None to derive); got {self.sched_policy}", ) for requested, supported, name in ( (self.tile_m, 128, "tile_m"), @@ -842,6 +888,25 @@ def check_support(self) -> bool: requested is not None and requested != supported, f"SM100 DSL SDPA only supports {name}={supported}", ) + self._value_error_if( + self.softmax_precision is not None, + "SM100 DSL SDPA has no softmax-precision arm yet (softmax_precision must be unset)", + ) + if self.split_kv > 1: + # Split-KV: partials weighted by the per-split LSE, recombined by + # split_combine_sm100 (which also owns the FP8 amax of the + # recombined O). Structural limits mirror mismatch()'s + # facts x knobs gate so the standalone API declines identically. + self._not_implemented_error_if( + self._fp8 and self.dtype_o not in (torch.float16, torch.bfloat16), + "split_kv > 1 on a quantized graph requires a bf16/fp16 O (the combine reduces half-precision partials)", + ) + self._not_implemented_error_if(self.thd, "split_kv > 1 is dense-only (THD packs its own flat grid)") + self._value_error_if(self.has_sink, "split_kv > 1 with an attention sink is not supported") + self._value_error_if( + self.seq_kv_lens_present or self.seq_q_lens_present, + "split_kv > 1 serves unpadded dense graphs only", + ) swa_left = self.window_size_left self._value_error_if( @@ -932,18 +997,24 @@ def compile(self) -> None: # picks the fused path with no user action. mxfp8 = self._fp8 and not self._pertensor fused_ldtm_stat = mxfp8 and (self._device_cc == (10, 3)) + # None = the standalone-wrapper tier stated no preference: derive the + # causal-balancing policy here. The graph path never hits this branch — + # the heuristic emits an explicit policy (the same primary this + # derivation picks) and it is honored verbatim, NATURAL included. sched_policy = self.sched_policy - if sched_policy == SCHED_NATURAL and self.window_right is not None: - # Causal: balance the triangular load; pick the LPT variant by working set. - _, _, s_kv_sched, _ = self.k_desc.shape - _, _, _, d_qk_sched = self.q_desc.shape - _, _, _, d_v_sched = self.v_desc.shape - sched_policy = _causal_sched_policy( - s_kv=s_kv_sched, - d_qk=d_qk_sched, - d_v=d_v_sched, - elem_bytes=1 if self._fp8 else 2, - ) + if sched_policy is None: + sched_policy = SCHED_NATURAL + if self.window_right is not None: + # Causal: balance the triangular load; pick the LPT variant by working set. + _, _, s_kv_sched, _ = self.k_desc.shape + _, _, _, d_qk_sched = self.q_desc.shape + _, _, _, d_v_sched = self.v_desc.shape + sched_policy = _causal_sched_policy( + s_kv=s_kv_sched, + d_qk=d_qk_sched, + d_v=d_v_sched, + elem_bytes=1 if self._fp8 else 2, + ) lpt_head_group = 1 if self._fp8 and self.flavor == (192, 128) and not self.thd and (self.batch_size * self.h_q) % 8 == 0: lpt_head_group = 8 @@ -977,6 +1048,7 @@ def compile(self) -> None: lpt_q_tiles=lpt_q_tiles, thd_varlen=self.thd, fused_ldtm_stat=fused_ldtm_stat, + split_kv=self.split_kv, ) self._k_mod = _load_sm100_kernel_module(self.flavor, params, fp8=self._fp8, pertensor=self._pertensor, rubin=(self._device_cc == (10, 7))) if self.thd: @@ -992,14 +1064,17 @@ def compile(self) -> None: # FP8/MXFP8 kernels use exact native shapes (gated in check_support); # their compile() has no envelope head-dim parameters. has_lse=False # (no Stats output) compiles the LSE store out — no dummy buffer at - # any level (the amax_o atomicMax write is independent). + # any level (the amax_o atomicMax write is independent). A split + # REQUIRES the in-kernel LSE: the per-split LSE is the combine + # weight (the kernel skips its own amax write under a split; the + # combine reports the amax of the RECOMBINED O instead). self._compiled_kernel = self._k_mod.compile( b=self.batch_size, qh=self.h_q, kh=self.h_kv, sq=self.s_q_max, skv=self.s_k_max, - has_lse=self.lse_desc is not None, + has_lse=(self.lse_desc is not None) or self.split_kv > 1, ) else: # ENVELOPE: hand the f16/bf16 kernel the ACTUAL head dims so its @@ -1007,6 +1082,8 @@ def compile(self) -> None: # zero-fill, O stores past d_v clip); the tile box stays the # flavor's compile-time D. has_lse=False (no Stats output) # compiles the LSE store out — no dummy buffer at any level. + # Split-KV REQUIRES the in-kernel LSE regardless of a Stats + # output: the per-split LSE is the combine weight. self._compiled_kernel = self._k_mod.compile( b=self.batch_size, qh=self.h_q, @@ -1015,7 +1092,26 @@ def compile(self) -> None: skv=self.s_k_max, d_qk=self.head_dim_qk, d_v=self.head_dim_v, + has_lse=(self.lse_desc is not None) or self.split_kv > 1, + ) + self._combine_kernel = None + if self.split_kv > 1: + # The recombine pass compiles at PLAN time like everything else; + # execute() only rebinds the partial slabs it carves. On the FP8 + # families the combine also owns the amax of the recombined O + # (a max over per-split partials would over-report — each split's + # O is normalized by its own running sum). + from cudnn.sdpa.fwd.kernels import split_combine_sm100 as _split_combine + + self._combine_kernel = _split_combine.compile( + b=self.batch_size, + h=self.h_q, + sq=self.s_q_max, + d_v=self.head_dim_v, + splits=self.split_kv, + dtype_o=self._combine_dtype_tag(), has_lse=self.lse_desc is not None, + has_amax=self._fp8, ) self._logger.debug("compile completed") @@ -1047,8 +1143,17 @@ def scratch_workspace_bytes(self) -> int: # clamped K/V runtime descriptors (see the kernels' THD closures). o_desc_slots = b + (3 if self._fp8 else 1) return ws_align((3 * b + 2) * 4) + ws_align(o_desc_slots * 16 * 8) + (0 if self.has_sink else ws_align(qh * 4)) - if self._fp8: + if self._fp8 and self.split_kv == 1: return 0 # dense FP8/MXFP8: no per-execute scratch (dummies are cached one-time) + if self.split_kv > 1: + # Split-major partial slabs the main kernel writes and the combine + # pass reduces: O_s [splits*B, S_q, H, d_v] in the O dtype (half — + # the split path requires a bf16/fp16 O even on the FP8 families) + # and lse_s [splits*B, H, S_q] fp32. Carved from the caller's + # workspace — zero per-execute allocations (Hard Rule 1). + o_bytes = self.split_kv * b * self.s_q_max * qh * self.head_dim_v * self._o_itemsize() + lse_bytes = self.split_kv * b * qh * self.s_q_max * 4 + return ws_align(o_bytes) + ws_align(lse_bytes) # Dense padded-Q lens bind directly as their own kernel parameter # (no combine buffer since the seq_len_q-as-parameter change) — no scratch. return 0 @@ -1211,21 +1316,58 @@ def execute( import cutlass - self._compiled_kernel( - Q, - K, - V, - O_scratch if o_needs_copy_back else O_view, - lse_tensor.reshape(self.batch_size, self.h_q, self.s_q_max) if lse_tensor is not None else None, - sinks_t, - seq_kv_t, - o_desc_dummy, - (self.batch_size, self.h_q, self.h_kv, self.s_q_max, self.s_k_max, 0), - cutlass.Float32(scale_softmax_log2), - cutlass.Int32(0), - seq_q_t, - stream=current_stream, - ) + o_arg = O_scratch if o_needs_copy_back else O_view + lse_arg = lse_tensor.reshape(self.batch_size, self.h_q, self.s_q_max) if lse_tensor is not None else None + if self.split_kv > 1: + # Redirect the mainloop into split-major partial slabs, then reduce + # into the caller's O/LSE with the plan-time-compiled combine pass. + # Slabs are carved from the caller's workspace (Hard Rule 1); + # standalone use (no workspace) torch-allocates, like the THD path. + # No zero-fill: the split kernel writes EVERY (split, batch) slot, + # emitting O := 0 / lse := -inf for empty split ranges itself. + s, b, h, sq, dv = self.split_kv, self.batch_size, self.h_q, self.s_q_max, self.head_dim_v + o_partial, lse_partial = self._split_partials(workspace, o_arg, device, current_stream) + self._compiled_kernel( + Q, + K, + V, + o_partial, + lse_partial, + sinks_t, + seq_kv_t, + o_desc_dummy, + (b, h, self.h_kv, sq, self.s_k_max, 0), + cutlass.Float32(scale_softmax_log2), + cutlass.Int32(0), + seq_q_t, + stream=current_stream, + ) + self._combine_kernel( + o_partial, + lse_partial, + o_arg, + lse_arg, + None, + (b, h, sq, dv), + cutlass.Int32(s), + stream=current_stream, + ) + else: + self._compiled_kernel( + Q, + K, + V, + o_arg, + lse_arg, + sinks_t, + seq_kv_t, + o_desc_dummy, + (self.batch_size, self.h_q, self.h_kv, self.s_q_max, self.s_k_max, 0), + cutlass.Float32(scale_softmax_log2), + cutlass.Int32(0), + seq_q_t, + stream=current_stream, + ) if o_needs_copy_back: O_view.copy_(O_scratch) self._logger.debug("execute completed") @@ -1671,15 +1813,22 @@ def _execute_mxfp8( amax_o_buf.zero_() o_desc_dummy = self._dummy("o_desc", device, lambda: torch.zeros(1, dtype=torch.int64, device=device)) + # Split-KV: the mainloop writes split-major partials (skipping its own + # amax — each split's O is normalized by its own running sum, so a max + # over partials over-reports); the combine reduces them into the + # caller's O/LSE and owns the recombined amax. + O_dst, lse_dst = O, lse + if self.split_kv > 1: + O_dst, lse_dst = self._split_partials(workspace, O, device, current_stream) self._compiled_kernel( Q, K, V, - O, + O_dst, sf_q_v, sf_k_v, sf_v_v, - lse, + lse_dst, amax_o_buf, sinks_t, seq_kv_t, @@ -1689,6 +1838,17 @@ def _execute_mxfp8( cutlass.Int32(0), stream=current_stream, ) + if self.split_kv > 1: + self._combine_kernel( + O_dst, + lse_dst, + O, + lse, + amax_o_buf, + (b, h_q, sq, self.head_dim_v), + cutlass.Int32(self.split_kv), + stream=current_stream, + ) if o_needs_copy_back: O_view.copy_(O) self._logger.debug("execute (MXFP8) completed") @@ -1822,12 +1982,18 @@ def _execute_fp8( amax_o_buf.zero_() o_desc_dummy = self._dummy("o_desc", device, lambda: torch.zeros(1, dtype=torch.int64, device=device)) + # Split-KV: mainloop into split-major partials (the kernel skips its + # in-kernel amax under a split), combine into the caller's O/LSE with + # the recombined amax. + O_dst, lse_dst = O, lse + if self.split_kv > 1: + O_dst, lse_dst = self._split_partials(workspace, O, device, current_stream) self._compiled_kernel( Q, K, V, - O, - lse, + O_dst, + lse_dst, sinks_t, seq_kv_t, o_desc_dummy, @@ -1842,6 +2008,17 @@ def _execute_fp8( amax_o_buf, stream=current_stream, ) + if self.split_kv > 1: + self._combine_kernel( + O_dst, + lse_dst, + O, + lse, + amax_o_buf, + (b, h_q, sq, self.head_dim_v), + cutlass.Int32(self.split_kv), + stream=current_stream, + ) if o_needs_copy_back: O_view.copy_(O) if amax_o is not None: @@ -2275,6 +2452,28 @@ def _smem_bytes(kv_tile: int) -> int: if self.scale_softmax is None or self.scale_softmax == 0.0: self.scale_softmax = 1.0 / math.sqrt(d_q) + self._value_error_if( + self.sched_policy is not None and self.sched_policy not in (SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2), + f"SM120 DSL SDPA sched_policy must be NATURAL/LPT/LPT_L2 (or None to derive); got {self.sched_policy}", + ) + if self.split_kv > 1: + # The SM120 kernel's inline split chunking + the shared (arch- + # agnostic, one block per row) split_combine pass. The config + # backstop additionally bars a split under the LPT remaps — + # validated at compile via make_cfg, and the heuristic's split + # sets ride SCHED_NATURAL. + self._not_implemented_error_if(self._fp8, "SM120 split_kv > 1 is f16/bf16-only (the fp8 kernel has no split path)") + self._not_implemented_error_if(self.thd, "split_kv > 1 is dense-only (THD packs its own flat grid)") + self._value_error_if(self.has_sink, "split_kv > 1 with an attention sink is not supported") + self._value_error_if( + self.seq_kv_lens_present or self.seq_q_lens_present, + "split_kv > 1 serves unpadded dense graphs only", + ) + self._value_error_if( + self.softmax_precision is not None, + "SM120 DSL SDPA has no softmax-precision arm yet (softmax_precision must be unset)", + ) + self.batch_size = int(b) self.s_q_max = int(s_q) self.s_k_max = int(s_kv) @@ -2295,18 +2494,23 @@ def compile(self) -> None: if self._compiled_kernel is not None: return - # Causal: balance the triangular load; pick the LPT variant by working set. + # None = the standalone-wrapper tier stated no preference: derive the + # causal-balancing policy here. The graph path arrives with an explicit + # policy from the heuristic and it is honored verbatim, NATURAL included. sched_policy = self.sched_policy - if sched_policy == SCHED_NATURAL and self.window_right is not None: - _, _, s_kv_sched, _ = self.k_desc.shape - _, _, _, d_qk_sched = self.q_desc.shape - _, _, _, d_v_sched = self.v_desc.shape - sched_policy = _causal_sched_policy( - s_kv=s_kv_sched, - d_qk=d_qk_sched, - d_v=d_v_sched, - elem_bytes=1 if self._fp8 else 2, - ) + if sched_policy is None: + sched_policy = SCHED_NATURAL + if self.window_right is not None: + # Causal: balance the triangular load; pick the LPT variant by working set. + _, _, s_kv_sched, _ = self.k_desc.shape + _, _, _, d_qk_sched = self.q_desc.shape + _, _, _, d_v_sched = self.v_desc.shape + sched_policy = _causal_sched_policy( + s_kv=s_kv_sched, + d_qk=d_qk_sched, + d_v=d_v_sched, + elem_bytes=1 if self._fp8 else 2, + ) params = Sm120TemplateParams( dtype_qkv=_SM120_DTYPE_QKV_CODE[self.dtype], dtype_o=_SM120_DTYPE_QKV_CODE[self.o_desc.dtype], @@ -2320,6 +2524,7 @@ def compile(self) -> None: thd_varlen=self.thd, q_tile=self.q_tile, kv_tile=self.kv_tile, + split_kv=self.split_kv, ) self._k_mod = _load_sm120_kernel_module(params, fp8=self._fp8) if self.thd: @@ -2342,9 +2547,27 @@ def compile(self) -> None: d_qk=self.head_dim_qk, d_v=self.head_dim_v, # No sample_lse -> the LSE store is compiled out; execute() then - # binds no LSE buffer at all (no dummy, no allocation). - has_lse=self.lse_desc is not None, - ) + # binds no LSE buffer at all (no dummy, no allocation). A split + # REQUIRES it: the per-split LSE is the combine weight. + has_lse=(self.lse_desc is not None) or self.split_kv > 1, + ) + self._combine_kernel = None + if self.split_kv > 1: + # The recombine pass compiles at PLAN time; execute() only rebinds + # the partial slabs it carves. The combine kernel is arch-agnostic + # (one block per (q_row, head, batch), no cluster/TMEM features). + from cudnn.sdpa.fwd.kernels import split_combine_sm100 as _split_combine + + self._combine_kernel = _split_combine.compile( + b=self.batch_size, + h=self.h_q, + sq=self.s_q_max, + d_v=self.head_dim_v, + splits=self.split_kv, + dtype_o=self._combine_dtype_tag(), + has_lse=self.lse_desc is not None, + has_amax=False, + ) self._logger.debug("compile completed") def execute( @@ -2464,12 +2687,18 @@ def execute( v = self._to_bshd(v_tensor) o_view, o_needs_copy_back, o_scratch = self._to_bshd_writable(o_tensor) o = o_scratch if o_needs_copy_back else o_view + # Split-KV: the kernel's inline chunking writes split-major partials + # (every (split, batch) slot — empty ranges as O := 0 / lse := -inf); + # the shared combine reduces them into the caller's O/LSE. + o_dst, lse_dst = o, lse + if self.split_kv > 1: + o_dst, lse_dst = self._split_partials(workspace, o, q_tensor.device, current_stream) self._compiled_kernel( q, k, v, - o, - lse, + o_dst, + lse_dst, sinks_t, seq_q_lens, seq_kv_lens, @@ -2480,6 +2709,17 @@ def execute( None, current_stream, ) + if self.split_kv > 1: + self._combine_kernel( + o_dst, + lse_dst, + o, + lse, + None, + (self.batch_size, self.h_q, self.s_q_max, self.head_dim_v), + cutlass.Int32(self.split_kv), + stream=current_stream, + ) if o_needs_copy_back: o_view.copy_(o_scratch) @@ -2882,6 +3122,13 @@ def scratch_workspace_bytes(self) -> int: # guarded GMEM stores, so THD needs no per-sequence tensor maps. b = self.batch_size return ws_align((3 * b + 2) * 4) + if self.split_kv > 1: + # Split-major partial slabs (see the SM100 sibling): O_s in the O + # dtype (half) + lse_s fp32, carved from the caller's workspace. + b, qh = self.batch_size, self.h_q + o_bytes = self.split_kv * b * self.s_q_max * qh * self.head_dim_v * self._o_itemsize() + lse_bytes = self.split_kv * b * qh * self.s_q_max * 4 + return ws_align(o_bytes) + ws_align(lse_bytes) return 0 @@ -3189,7 +3436,8 @@ class SdpaFwdDslSm80(SdpaFwdDsl): 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. + # and carries the heuristic's explicit sched_policy knob instead + # (None = derive via "auto", explicit ints map to their tokens). self._scheduler_token = scheduler self._bias_present = bool(bias_present) self._bias_fp32 = bool(bias_fp32) @@ -3285,6 +3533,14 @@ def check_support(self) -> bool: self._rope_max_s and (self.seq_kv_lens_present or self.seq_q_lens_present), "SM80 SDPA: RoPE fusion is dense-unpadded-only", ) + self._not_implemented_error_if( + self.split_kv > 1, + "SM80 SDPA has no split-KV path (no partial slabs, no combine kernel)", + ) + self._value_error_if( + self.softmax_precision is not None, + "SM80 SDPA has no softmax-precision arm yet (softmax_precision must be unset)", + ) self._value_error_if(not torch.cuda.is_available(), "CUDA must be available for SM80 SDPA") device = self.q_desc.device @@ -3327,11 +3583,17 @@ def check_support(self) -> bool: token = self._scheduler_token if token is None: - token = {SCHED_NATURAL: "auto", SCHED_LPT: "lpt", SCHED_LPT_L2: "lpt_l2"}.get(self.sched_policy) - self._value_error_if( - token is None, - f"SM80 SDPA: unsupported sched_policy {self.sched_policy}", - ) + # None = no preference anywhere -> "auto" (the adapter derives, a + # standalone-wrapper convenience). An explicit knob is honored + # verbatim — NATURAL included — never re-derived. + if self.sched_policy is None: + token = "auto" + else: + token = {SCHED_NATURAL: "default", SCHED_LPT: "lpt", SCHED_LPT_L2: "lpt_l2"}.get(self.sched_policy) + self._value_error_if( + token is None, + f"SM80 SDPA: unsupported sched_policy {self.sched_policy}", + ) else: _VALID = ("auto", "natural", "default", "lpt", "lpt_l2") self._value_error_if(token not in _VALID, f"scheduler must be one of {_VALID}; got {token!r}") diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 7f0230b2d..373a885e9 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -33,7 +33,7 @@ import cudnn -from cudnn.frost.tile_dsl.constants import SCHED_NATURAL +from cudnn.frost.tile_dsl.constants import SCHED_LPT, SCHED_LPT_L2, SCHED_NATURAL from cudnn.frost.buffers import CUTEDSL_MIN_VERSION, cutedsl_state, cutedsl_too_old from cudnn.sdpa import graph_analyzer as ga @@ -87,6 +87,14 @@ class SdpaFwdKnobs: tile_m: Optional[int] = None # Q sequence tile width tile_n: Optional[int] = None # KV sequence tile width cga: Optional[int] = None # cluster size (CTAs cooperating per tile) + # KV-split count: each Q tile's KV range cut into this many chunks, each + # run by its own CTA, recombined by the split_combine pass. 1 = off. + split_kv: Optional[int] = None + # Softmax accumulation precision as a cudnn.data_type value. Framework + # axis: no forward kernel declares a domain yet, so any explicit request + # declines; the first kernel that grows an f16-softmax arm lights it up + # by declaring {FLOAT, HALF} on its row. + softmax_precision: Optional[int] = None @dataclass(frozen=True) @@ -221,6 +229,12 @@ class Capabilities: tile_ms: frozenset[int] = frozenset() tile_ns: frozenset[int] = frozenset() cgas: frozenset[int] = frozenset() + # Split-KV domain. {1} = the axis exists but only "off" is served; rows + # whose kernels wire the split path AND whose adapter launches the combine + # widen this (the SM100 f16 rows today). + split_kvs: frozenset[int] = frozenset({1}) + # Softmax-precision domain (cudnn.data_type values). Empty = unserved. + softmax_precisions: frozenset[int] = frozenset() def _band_covers_kv_tail(facts: "ga.SdpaGraphFacts") -> bool: @@ -260,9 +274,29 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti (knobs.tile_m, capabilities.tile_ms, "tile_m"), (knobs.tile_n, capabilities.tile_ns, "tile_n"), (knobs.cga, capabilities.cgas, "cga"), + (knobs.split_kv, capabilities.split_kvs, "split_kv"), + (knobs.softmax_precision, capabilities.softmax_precisions, "softmax_precision"), ): if value is not None and value not in domain: return f"requested {label}={value} is outside this engine's domain {sorted(domain)}" + if knobs.split_kv is not None and knobs.split_kv > 1: + # Facts x knobs: the split path is structurally dense-only (the + # per-split LSE is the combine weight; the THD/sink/padded paths + # do not produce per-split partials). Declined HERE so a split + # request never reaches a kernel that cannot honor it. + if facts.thd or facts.has_sink or facts.padded or facts.seq_q_trim: + return "split_kv > 1 serves dense, unpadded, sink-free graphs only" + if capabilities.skv_tail_via_padding and facts.s_kv % (capabilities.skv_tile or 128) != 0 and not _band_covers_kv_tail(facts): + # The lowering would serve this ragged S_kv through the padded + # kernel path (synthesized per-batch KV lengths) — the same + # path the split cannot ride. Mirror lower_dsl_prefill's + # synth_kv_padding predicate so the plan is never listed. + return "split_kv > 1 cannot ride the synthesized KV-tail padding this S_kv needs" + if (facts.is_fp8 or facts.is_mxfp8) and facts.dtype_o not in (cudnn.data_type.HALF, cudnn.data_type.BFLOAT16): + # The combine reduces partials in half precision; reducing + # QUANTIZED partials would lose what the split must be + # numerically neutral about. + return "split_kv > 1 on a quantized graph requires a bf16/fp16 O" cc = facts.device_cc sm = None if cc is None else cc[0] * 10 + cc[1] if sm is None or not (capabilities.sm_lo <= sm <= capabilities.sm_hi): @@ -443,10 +477,14 @@ def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: # FP8/MXFP8 rows stay on the strict BSHD gate until their padded / # scale-factor paths are validated against relaxed layouts. layouts=frozenset({"bshd", "dense_flex"}), - sched_policies=frozenset({SCHED_NATURAL}), + sched_policies=frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}), tile_ms=frozenset({128}), tile_ns=frozenset({128}), cgas=frozenset({2}), + # All four f16 flavor kernels wire SplitHelpers, and the adapter + # carves the partial slabs + launches split_combine_sm100 when + # split_kv > 1 (dense f16 only; see mismatch's facts x knobs gate). + split_kvs=frozenset({1, 2, 4}), ), lower=partial(lower_dsl_prefill, api_type=_SM100), ) @@ -487,10 +525,13 @@ def _sm100_mxfp8_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: lse_optional=True, thd=thd, cu_seq_len=thd, - sched_policies=frozenset({SCHED_NATURAL}), + sched_policies=frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}), tile_ms=frozenset({128}), tile_ns=frozenset({128}), cgas=frozenset({2}), + # Only the d128 mxfp8 kernel wires SplitHelpers; the split path + # also needs a half-precision O (mismatch's facts x knobs gate). + split_kvs=frozenset({1, 2, 4}) if d == 128 else frozenset({1}), ), lower=partial(lower_dsl_prefill, api_type=_SM100), ) @@ -554,10 +595,15 @@ def _sm100_fp8_spec( # race was fixed with the mb_stats_read barrier (verified on the # gated 132/192/200-cluster repros, 3x each). skv_tail_via_padding=True, - sched_policies=frozenset({SCHED_NATURAL}), + sched_policies=frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}), tile_ms=frozenset({128}), tile_ns=frozenset({128}), cgas=frozenset({2}), + # Only the d128 fp8 kernel wires SplitHelpers (the d192x128 file + # forks its own scheduler and has no split path). Split partials + # reduce in half precision, so mismatch()'s facts x knobs gate + # additionally requires a bf16/fp16 O on the quantized rows. + split_kvs=frozenset({1, 2, 4}) if d == 128 else frozenset({1}), ), lower=partial(lower_dsl_prefill, api_type=_SM100), ) @@ -599,7 +645,10 @@ def _sm80_spec() -> EngineSpec: lse_optional=True, layouts=frozenset({"bshd", "dense_flex"}), skv_tile=0, # the kernels' is_even_k path serves ragged S_kv - sched_policies=frozenset(), + # The static-grid remap serves all three policies (the template's + # sched_policy field); the adapter maps the explicit int to its + # kernel token and derives only when the knob is None. + sched_policies=frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}), ), lower=partial(lower_dsl_prefill, api_type=_SM80), ) @@ -637,7 +686,12 @@ def _sm120_spec() -> EngineSpec: skv_tile=0, cu_seq_len=True, layouts=frozenset({"bshd", "dense_flex"}), - sched_policies=frozenset({SCHED_NATURAL}), + sched_policies=frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}), + # The kernel's inline chunking + the shared split_combine pass + # (the combine is one block per row — arch-agnostic). The config + # backstop bars a split under the LPT remaps, so the heuristic's + # split sets ride SCHED_NATURAL. + split_kvs=frozenset({1, 2, 4}), tile_ms=frozenset({64, 128}), tile_ns=frozenset({64, 128}), cgas=frozenset({1}), @@ -735,6 +789,8 @@ 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, + split_kv=knobs.split_kv if knobs is not None else None, + softmax_precision=knobs.softmax_precision 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 @@ -1000,7 +1056,7 @@ def _sm120_fp8_spec() -> EngineSpec: # by the shared adapter (one gather copy in, one scatter copy # back for O; zero-copy when already BSHD-physical). layouts=frozenset({"bshd", "dense_flex"}), - sched_policies=frozenset({SCHED_NATURAL}), + sched_policies=frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}), tile_ms=frozenset({64, 128}), tile_ns=frozenset({64, 128}), cgas=frozenset({1}), diff --git a/python/cudnn/sdpa/fwd/heuristics.py b/python/cudnn/sdpa/fwd/heuristics.py index dd3acc6dd..fefd54c31 100644 --- a/python/cudnn/sdpa/fwd/heuristics.py +++ b/python/cudnn/sdpa/fwd/heuristics.py @@ -1,50 +1,159 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""How the SDPA-forward family ranks plans for a graph. - -``engines/heuristics.rank`` hands over the parsed facts, this family's engine -ids, and the backend's entries tagged by mode; what :func:`recommend` returns is -``graph.plans``, position for position. The comparison is here because a cell -cannot see its siblings and neither side of the FROST/backend split can place -the other. - -Per mode: - -- **A** — candidates worth running, best guess first, runners-up behind it for - a caller that autotunes. -- **FALLBACK** — the config expected to build where mode A's choice may not. - Nothing here is chosen for speed. -- **OPENSOURCE** — mode A without the backend's recommendation, since these - cells ARE the open-source implementation. -- **B** — answered as A: it asks for a wider search this family has none to - give. - -To add a rule for a cell: write the function (:func:`_sm120_tiles` is the worked -example), list the cell in ``_TILE_RULE_CELLS``, put the measurement in the -commit. A cell absent from that set runs its row's sole point per axis, which is +"""The SDPA-forward family's proposals: which cells, with which knob sets. + +:func:`recommend` is the family's ENTIRE heuristic surface — the PURE core: +``(kind, facts, offered) -> [PlanConfig]``. Backend-blind, graph-blind, +import-light. For every offered cell whose capability row admits the facts, +the cell's rule emits an ORDERED list of COMPLETE knob assignments (every +axis the row declares a domain for carries a concrete value; ``None`` only on +undeclared axes), each re-validated through ``mismatch(caps, facts, knobs)`` +— a set is honored or never listed. The same engine appears once per +surviving set. Standalone callers (wrappers, autotuners) invoke this directly +with a hand-built :class:`~cudnn.sdpa.graph_analyzer.SdpaGraphFacts`; nothing +here touches the backend, a graph object, or heuristic modes. + +``kind`` is ``"A"`` (candidates worth timing, best guess first, runners-up +behind for a caller that autotunes) or ``"FALLBACK"`` (the config expected to +build where A's choice may not — nothing chosen for speed). + +Everything else about the final list — mode blocks, the backend's entries, +the delegating entry, dedup, the mode strip — is PLACEMENT, and placement is +not a family opinion: it lives once in ``engines/heuristics._assemble``, +under the standing assumption that these proposals lead the backend's entries +(an OSS engine measured behind the backend gets fixed or pulled, not +demoted). + +Cross-ENGINE order within a proposal batch is ``ENGINE_SPECS`` declaration +order. Today that is unambiguous in practice — co-eligible cells are the +envelope-overlap family, which all lower to the same kernel — and the seam +for a real ranking, when one is measured, is a score stage here in +:func:`recommend`, not a new layer. + +To add a rule for a cell: write a generator (:func:`_sm120_tiles` is the +worked example), register the cell in ``_TILE_RULE_CELLS`` (or grow a new +axis via the five-part checklist in ``engines.py``), put the measurement in +the commit. A cell with no rule runs its row's sole point per axis, which is the honest answer while nobody has timed it. """ from __future__ import annotations -from typing import Any, Dict, List, Tuple +from dataclasses import replace +from typing import Dict, Iterator, List, Optional, Tuple import cudnn from cudnn.engines.base import PlanConfig +from cudnn.frost.tile_dsl.constants import SCHED_LPT, SCHED_LPT_L2, SCHED_NATURAL from cudnn.sdpa.fwd.config_sm120 import FP8_HEAD_TILE_GRANULE, HEAD_TILE_GRANULE, SMEM_CAPACITY_BYTES, smem_bytes -from cudnn.sdpa.fwd.engines import ENGINE_SPECS, Capabilities, SdpaFwdKnobs, mismatch - -# Cells timed against the backend's kernel and found SLOWER, so the backend's -# mode-A entries lead them. Empty: absent means faster OR never timed, and both -# keep the order this dispatch has always had. Moving a cell in needs a -# measurement. -_MEASURED_BEHIND: frozenset = frozenset() +from cudnn.sdpa.fwd.engines import ENGINE_SPECS, Capabilities, EngineSpec, SdpaFwdKnobs, _band_covers_kv_tail, mismatch # Cells whose (tile_m, tile_n) choice _sm120_tiles makes. _TILE_RULE_CELLS = frozenset({"sdpa_fwd_prefill_sm120", "sdpa_fwd_prefill_sm120_fp8"}) +# Cap on complete knob sets emitted per engine per kind. The combiner grows +# Σ|axis runners|, never the cartesian product; this bound keeps the plan list +# legible and an autotune-ALL pass affordable even as axes accumulate. +_MAX_SETS_PER_ENGINE = 6 + +# The causal-balancing budget for the CLC/static LPT_L2 policy on SM100/SM120: +# LPT_L2's block-cyclic head grouping only pays when ONE head's K+V working set +# can actually stay L2-resident. +_SM100_L2_BUDGET_BYTES = 50 * 1024 * 1024 + +# The SM80 kernels' L2 grouping budget is a per-flavor MiB table fed to the +# template (sched_l2_mib); the adapter owns that table. For POINT ORDERING all +# that matters here is that SM80's measured primary for causal is LPT_L2. + + +# --------------------------------------------------------------------------- +# axis generators — each returns an ORDERED candidate list, best first +# --------------------------------------------------------------------------- + + +def _sole(values): + """The only value on an axis, or None where the row declares no domain.""" + return next(iter(values)) if len(values) == 1 else None + + +def _ceil_div(a: int, b: int) -> int: + return -(-a // b) + + +# --- KV split (see choose_split_kv) ---------------------------------------- +# Largest split considered; past this the reduction outgrows the parallelism. +_SPLIT_KV_MAX = 16 +# A split thinner than this is prologue/epilogue dominated. +_SPLIT_KV_MIN_TILES = 2 +# What a CTA-tile costs beyond its KV loop (Q load, prologue, epilogue), in +# units of one KV tile. Empirical: re-measure if the per-tile fixed cost moves. +_SPLIT_KV_CTA_COST = 21.0 + + +def choose_split_kv( + *, + q_tiles: int, + heads_q: int, + batch: int, + kv_tiles: int, + sm_count: int, + ctas_per_tile: int = 1, + max_split: int = _SPLIT_KV_MAX, +) -> int: + """How many KV chunks to cut each Q tile into; 1 = do not split. + + A prefill launch is ``q_tiles * heads_q * batch`` independent tiles, each + walking the whole KV loop. When that product is below the SM count the chip + idles however long the loop is; splitting multiplies the tile count by ``s`` + and divides each tile's KV work by it, then pays one reduction over the + partials. + + A CTA holds its tile for the whole loop, so a launch costs whole WAVES. + Minimise, over powers of two: + + waves(s) = ceil(base_ctas * s / sm_count) + cost(s) = waves(s) * (ceil(kv_tiles / s) + CTA_COST) + + CTA_COST is what a tile re-pays whatever its loop length, so it sits inside + the wave term -- once per CTA-tile, not once per split. + + What falls out: an under-full launch splits until the wave is full; an + over-full one with a partial-wave tail splits FINER to smooth it, even past + the SM count; an exactly balanced one (base_ctas = k * sm_count) has no tail + and never splits. + + Powers of two only, because ``split_kv`` is a TemplateParams field and so a + kernel-module cache key -- an unrestricted choice mints a compiled + specialization per shape. + + Returns 1 when there is nothing to split or nothing beats not splitting. + Bounded by ``max_split``, by ``kv_tiles`` (more splits than tiles would + leave some provably empty) and by ``_SPLIT_KV_MIN_TILES``. + """ + if min(q_tiles, heads_q, batch, kv_tiles, sm_count, ctas_per_tile) <= 0: + return 1 + base_ctas = q_tiles * heads_q * batch * ctas_per_tile + if kv_tiles <= 1: + return 1 + + best_split = 1 + best_cost = float(_ceil_div(base_ctas, sm_count) * (kv_tiles + _SPLIT_KV_CTA_COST)) + split = 2 + while split <= min(max_split, kv_tiles): + # Every split must stay thick enough to amortise its own prologue and + # epilogue. The chunking hands the remainder to the leading splits, so + # the THINNEST gets floor(kv_tiles / split) -- that is what must clear. + if kv_tiles // split < _SPLIT_KV_MIN_TILES: + break + waves = _ceil_div(base_ctas * split, sm_count) + cost = waves * (_ceil_div(kv_tiles, split) + _SPLIT_KV_CTA_COST) + if cost < best_cost: + best_split, best_cost = split, cost + split <<= 1 + return best_split + def _sm120_tiles(caps: Capabilities, facts) -> Tuple[int, int]: """(tile_m, tile_n) for the SM120 SDPA-forward prefill cell. @@ -89,19 +198,189 @@ def _sm120_tiles(caps: Capabilities, facts) -> Tuple[int, int]: return tile_m, (fits[0] if fits else min(caps.tile_ns)) -def _sole(values): - """The only value on an axis, or None where the row declares no domain.""" - return next(iter(values)) if len(values) == 1 else None +def _tile_points(spec: EngineSpec, facts) -> List[Tuple[Optional[int], Optional[int]]]: + """Ordered (tile_m, tile_n) candidates: the rule's best guess first, then + the rest of the SMEM-fitting domain for a caller that autotunes. Configs + the kernel cannot fit are not runners-up — they would sit in the list only + to decline at build.""" + caps = spec.capabilities + if spec.name not in _TILE_RULE_CELLS: + # No rule measured for this cell: its capability row has one point per + # axis, so there is nothing to choose between anyway. + return [(_sole(caps.tile_ms), _sole(caps.tile_ns))] + best = _sm120_tiles(caps, facts) + qkv_itemsize, o_itemsize = (1, 2) if facts.is_fp8 else (2, 2) + # Same envelope rounding as _sm120_tiles / the adapter's SMEM check. + granule = FP8_HEAD_TILE_GRANULE if facts.is_fp8 else HEAD_TILE_GRANULE + d_qp = -(-facts.d_qk // granule) * granule + d_vp = -(-facts.d_v // granule) * granule + domain = [(m, n) for m in caps.tile_ms for n in caps.tile_ns if smem_bytes(d_qp, d_vp, m, n, qkv_itemsize, o_itemsize) <= SMEM_CAPACITY_BYTES] + return sorted(domain or [best], key=lambda mn: (mn != best, mn[1] != best[1], -mn[0])) + +def _sched_points(caps: Capabilities, facts) -> List[Optional[int]]: + """Ordered scheduler-policy candidates. -def _knobs(caps: Capabilities, tile_m, tile_n) -> SdpaFwdKnobs: - """A knob request for one point. A field is None ONLY where the capability - row declares no domain for that axis — never "engine, pick for me", which - is how the same choice ended up being made here and again in the adapter.""" - return SdpaFwdKnobs(sched_policy=_sole(caps.sched_policies), tile_m=tile_m, tile_n=tile_n, cga=_sole(caps.cgas)) + The PRIMARY reproduces what each adapter's internal derivation historically + chose for the graph path, so promoting the decision into the ranked list + changes nothing for a caller that builds the first plan; the remaining + domain follows for autotune. This is the one causal LPT/LPT_L2 oracle on + the graph path — the adapters keep a None-input derivation only for + standalone wrapper users who bypass ranking. + """ + domain = caps.sched_policies + if len(domain) <= 1: + return [_sole(domain)] + causal_ish = facts.causal or facts.right_band_widening + if caps.sm_hi == 80: + # SM80's measured choices (see the adapter's flavor table): causal + # always groups for L2; pure SWA prefers plain LPT only in the band + # where the window walk is long enough to imbalance rows. + if causal_ish: + primary = SCHED_LPT_L2 + elif facts.window_left is not None and facts.window_left >= 0: + primary = SCHED_LPT if 1024 <= facts.s_kv <= 16384 else SCHED_NATURAL + else: + primary = SCHED_NATURAL + elif causal_ish: + # SM100/SM120: balance the triangular load; pick the LPT variant by + # whether one head's K+V working set fits the L2 budget. + elem = 1 if (facts.is_fp8 or facts.is_mxfp8) else 2 + one_head_bytes = int(facts.s_kv) * (int(facts.d_qk) + int(facts.d_v)) * elem + primary = SCHED_LPT_L2 if one_head_bytes <= _SM100_L2_BUDGET_BYTES else SCHED_LPT + else: + primary = SCHED_NATURAL + order = {SCHED_LPT_L2: (SCHED_LPT, SCHED_NATURAL), SCHED_LPT: (SCHED_LPT_L2, SCHED_NATURAL), SCHED_NATURAL: (SCHED_LPT, SCHED_LPT_L2)} + runners = [p for p in order[primary] if p in domain] + # A mask-free graph gains nothing from either LPT remap — the grid is + # already balanced — so don't spend autotune slots on them. + if not causal_ish and facts.window_left is None: + runners = [] + return [primary if primary in domain else _sole(domain) or SCHED_NATURAL] + runners + + +def _split_points(caps: Capabilities, facts, tile_m: Optional[int], tile_n: Optional[int], cga: Optional[int]) -> List[Optional[int]]: + """Ordered split-KV candidates for the chosen tile geometry. + + The value comes from :func:`choose_split_kv`'s wave-cost model, fed the + facts-level launch geometry (``tile_m*cga`` rows per tile — the recommend + tier's approximation of the kernel Cfg's exact ``TILES_Q*TILE_M*CTA_MMA``). + The generator respects the split path's structural limits (dense-only, no + sink — mismatch() enforces the same, so an emitted >1 never reaches a + kernel that cannot honor it). + + The split point is deliberately a RUNNER-UP behind no-split until sweeps + justify flipping the default: first-build behavior stays exactly what this + dispatch has always done, and autotune / select_plan reach the split plan + today. + """ + domain = caps.split_kvs + if len(domain) <= 1: + return [_sole(domain)] + no_split = 1 if 1 in domain else min(domain) + if facts.thd or facts.has_sink or facts.padded or facts.seq_q_trim: + return [no_split] + if caps.skv_tail_via_padding and facts.s_kv % (caps.skv_tile or 128) != 0 and not _band_covers_kv_tail(facts): + # This S_kv would be served through the synthesized KV-tail padding, + # which the split cannot ride (mismatch declines the same combination). + return [no_split] + if (facts.is_fp8 or facts.is_mxfp8) and facts.dtype_o not in (cudnn.data_type.HALF, cudnn.data_type.BFLOAT16): + # The combine reduces partials in half precision; reducing QUANTIZED + # partials would lose what the split is meant to be neutral about. + return [no_split] + sm_count = facts.device_sm_count or 0 + if sm_count <= 0: + return [no_split] + rows_per_tile = (tile_m or 128) * (cga or 1) + split = choose_split_kv( + q_tiles=_ceil_div(facts.s_q, rows_per_tile), + heads_q=facts.h_q, + batch=facts.b, + kv_tiles=_ceil_div(facts.s_kv, tile_n or 128), + sm_count=sm_count, + ctas_per_tile=cga or 1, + max_split=max(domain), + ) + # Snap the model's power-of-two answer down into the declared domain. + usable = [s for s in sorted(domain) if 1 < s <= split] + if not usable: + return [no_split] + return [no_split, usable[-1]] + + +def _softmax_points(caps: Capabilities) -> List[Optional[int]]: + """Softmax-precision candidates — sole-point until a kernel serves more.""" + return [_sole(caps.softmax_precisions)] + + +def _cga_points(caps: Capabilities) -> List[Optional[int]]: + """CGA candidates — every row today declares a single honest point.""" + return [_sole(caps.cgas)] + + +# --------------------------------------------------------------------------- +# the combiner — complete assignments, Σ growth, never the cartesian product +# --------------------------------------------------------------------------- + + +def _knob_sets(spec: EngineSpec, facts) -> List[SdpaFwdKnobs]: + """The cell's ordered COMPLETE knob assignments. + + The baseline takes the best value on every axis; runners-up deviate on ONE + axis at a time in impact order (tiles, sched, split) with the other axes + held at their best, capped at ``_MAX_SETS_PER_ENGINE``. Axis interactions + the kernels cannot serve are the generators'/mismatch's job — nothing here + multiplies domains together. + """ + caps = spec.capabilities + tiles = _tile_points(spec, facts) + scheds = _sched_points(caps, facts) + cga = _cga_points(caps)[0] + splits = _split_points(caps, facts, tiles[0][0], tiles[0][1], cga) + base = SdpaFwdKnobs( + sched_policy=scheds[0], + tile_m=tiles[0][0], + tile_n=tiles[0][1], + cga=cga, + split_kv=splits[0], + softmax_precision=_softmax_points(caps)[0], + ) + out = [base] + for tile_m, tile_n in tiles[1:]: + out.append(replace(base, tile_m=tile_m, tile_n=tile_n)) + for policy in scheds[1:]: + out.append(replace(base, sched_policy=policy)) + for split in splits[1:]: + # Split sets ride the plain scheduler: the SM120 config bars a split + # under the LPT remaps, and in the underfilled regime a split targets + # the LPT balancing is moot — the split itself levels the grid. + out.append(replace(base, split_kv=split, sched_policy=SCHED_NATURAL if SCHED_NATURAL in caps.sched_policies else base.sched_policy)) + seen, unique = set(), [] + for knobs in out: + if knobs not in seen: + seen.add(knobs) + unique.append(knobs) + return unique[:_MAX_SETS_PER_ENGINE] + + +def _fallback_knobs(caps: Capabilities) -> SdpaFwdKnobs: + """The config expected to build where the tuned choice may not. + + Today this is the smallest tile the row admits with the plain scheduler + and no split — the config that asks least of the device, which is the one + thing a fallback must be. + """ + return SdpaFwdKnobs( + sched_policy=SCHED_NATURAL if SCHED_NATURAL in caps.sched_policies else _sole(caps.sched_policies), + tile_m=min(caps.tile_ms, default=None), + tile_n=min(caps.tile_ns, default=None), + cga=_sole(caps.cgas), + split_kv=1 if 1 in caps.split_kvs else _sole(caps.split_kvs), + softmax_precision=_sole(caps.softmax_precisions), + ) -def _eligible(facts, offered: Dict[str, int]): +def _eligible(facts, offered: Dict[str, int]) -> Iterator[Tuple[int, EngineSpec]]: """(engine_id, spec) for each offered cell whose capability row admits ``facts``.""" for spec in ENGINE_SPECS: engine_id = offered.get(spec.name) @@ -109,101 +388,32 @@ def _eligible(facts, offered: Dict[str, int]): yield engine_id, spec -def _admissible(caps: Capabilities, facts, knobs: SdpaFwdKnobs) -> bool: - return mismatch(caps, facts, knobs) is None - - -def _mode_a(facts, offered: Dict[str, int], mode) -> List[PlanConfig]: - """Candidates worth timing, best guess first.""" - out = [] - for engine_id, spec in _eligible(facts, offered): - caps = spec.capabilities - if spec.name not in _TILE_RULE_CELLS: - # No rule measured for this cell: its capability row has one point - # per axis, so there is nothing to choose between anyway. - knobs = _knobs(caps, _sole(caps.tile_ms), _sole(caps.tile_ns)) - if _admissible(caps, facts, knobs): - out.append(PlanConfig(engine_id, knobs, mode=mode)) - continue - best = _sm120_tiles(caps, facts) - # The guess first, then the rest of the domain as autotune candidates: - # the rule's regret is small but not zero, so the runners-up are worth - # offering to a caller who measures. Configs the kernel cannot fit are - # not runners-up -- they would sit in the list only to decline at build. - qkv_itemsize, o_itemsize = (1, 2) if facts.is_fp8 else (2, 2) - # Same envelope rounding as _sm120_tiles / the adapter's SMEM check. - granule = FP8_HEAD_TILE_GRANULE if facts.is_fp8 else HEAD_TILE_GRANULE - d_qp = -(-facts.d_qk // granule) * granule - d_vp = -(-facts.d_v // granule) * granule - domain = [(m, n) for m in caps.tile_ms for n in caps.tile_ns if smem_bytes(d_qp, d_vp, m, n, qkv_itemsize, o_itemsize) <= SMEM_CAPACITY_BYTES] - ordered = sorted(domain or [best], key=lambda mn: (mn != best, mn[1] != best[1], -mn[0])) - for tile_m, tile_n in ordered: - knobs = _knobs(caps, tile_m, tile_n) - if _admissible(caps, facts, knobs): - out.append(PlanConfig(engine_id, knobs, mode=mode)) - return out +# --------------------------------------------------------------------------- +# recommend — the pure, backend-blind core (also the standalone entry point) +# --------------------------------------------------------------------------- -def _mode_fallback(facts, offered: Dict[str, int]) -> List[PlanConfig]: - """Configs expected to build where mode A's choice may not. +def recommend(kind: str, facts, offered: Dict[str, int]) -> List[PlanConfig]: + """Ordered candidate plans for ``facts`` — no backend, no graph, no modes. - TODO: today this is the smallest tile the row admits — the config that asks - least of the device, which is the one thing a fallback must be. Once a cell - has features its largest tiles cannot serve, this becomes the handful of - configs that between them cover the whole plane, chosen from measurements. + ``kind`` is ``"A"`` (candidates worth timing, best guess first) or + ``"FALLBACK"`` (least-demanding configs). Every returned entry carries a + complete knob assignment validated through ``mismatch(caps, facts, knobs)`` + — honored-or-never-listed — and NO mode. Standalone callers (wrappers, + autotuners) use this directly: build a ``SdpaGraphFacts``, pass the + family's ``offered_ids()``, run or time the sets in order. """ - out = [] + out: List[PlanConfig] = [] for engine_id, spec in _eligible(facts, offered): caps = spec.capabilities - knobs = _knobs(caps, min(caps.tile_ms, default=None), min(caps.tile_ns, default=None)) - if _admissible(caps, facts, knobs): - out.append(PlanConfig(engine_id, knobs, mode=cudnn.heur_mode.FALLBACK)) + sets = _knob_sets(spec, facts) if kind == "A" else [_fallback_knobs(caps)] + for knobs in sets: + if mismatch(caps, facts, knobs) is None: + out.append(PlanConfig(engine_id, knobs)) return out -def _leads(offered: Dict[str, int], plans: List[PlanConfig]) -> bool: - """Whether this family's mode-A plans outrank the backend's. See _MEASURED_BEHIND.""" - behind = {offered[name] for name in _MEASURED_BEHIND if name in offered} - return bool(plans) and not all(cfg.engine_id in behind for cfg in plans) - - -def recommend(modes: List[Any], facts, offered: Dict[str, int], backend_plans: List[PlanConfig]) -> List[PlanConfig]: - """The ranked plan list for this graph, mode by mode in the caller's order. - - Each mode contributes a block and the blocks concatenate, so asking for - ``[A, FALLBACK]`` puts every tuned candidate — both sides' — ahead of every - fallback. A plan repeated across modes keeps its first position: building - the same config twice only costs the caller a JIT compile. - """ - # An untagged backend entry is the delegating one: OSS candidates C++ holds - # but never exposes as plans, so it cannot be enumerated. It belongs to no - # mode, and it is NOT a pure OSS entry -- Graph::build_plans tries the OSS - # engine and, if that one declines, falls through to the native - # engine_configs already enqueued. So it leads the BACKEND's entries but not - # ours: ahead of our OPENSOURCE block it would answer an OSS-coverage - # question with a native kernel. - delegating = [c for c in backend_plans if c.mode is None] - out: List[PlanConfig] = [] - for mode in modes: - if mode == cudnn.heur_mode.OPENSOURCE: - out += _mode_a(facts, offered, cudnn.heur_mode.A) + delegating - elif mode in (cudnn.heur_mode.A, cudnn.heur_mode.B): - ours = _mode_a(facts, offered, mode) - theirs = delegating + [c for c in backend_plans if c.mode == mode] - out += (ours + theirs) if _leads(offered, ours) else (theirs + ours) - elif mode == cudnn.heur_mode.FALLBACK: - out += _mode_fallback(facts, offered) + delegating + [c for c in backend_plans if c.mode == mode] - # A delegate with no mode asked for it (the backend has engines but exposed - # no plans) would otherwise be dropped. - out += delegating - - # Identity is (engine, knobs). cpp_index is only WHERE one backend query put - # a plan, so keying on it would let [A, A], or one config both modes return, - # through as two entries -- and an autotuner would build and time it twice. - seen, ranked = set(), [] - for cfg in out: - key = (cfg.engine_id, repr(cfg.knobs)) - if key not in seen: - seen.add(key) - ranked.append(cfg) - return ranked +# Placement — mode blocks, the backend's entries, the delegating entry, dedup, +# the mode strip — is NOT this family's business: it happens once for every +# family in ``engines/heuristics._assemble``, with these proposals leading the +# backend's entries inside each block by standing assumption. diff --git a/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py b/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py new file mode 100644 index 000000000..1e3940d6d --- /dev/null +++ b/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The SDPA-forward heuristic's recommend contract. + +Unit tier (no GPU): recommend() emits ordered COMPLETE knob assignments — the +same engine repeated with different sets, every set admissible, no cartesian +blowup, cross-axis constraints never emitted, mode never on an entry. + +Executable tier (SM100): the ranked list carries knob-suffixed duplicates of +one cell; the split_kv entry, pinned by name, builds, carves its partial slabs +from the caller workspace, recombines correctly (O and Stats), and its +(engine_id, knobs) tuple replays on a fresh graph. +""" + +import math + +import pytest +import torch + +import cudnn +from cudnn.engines.base import PlanConfig +from cudnn.sdpa.fwd import engines +from cudnn.engines.heuristics import _assemble +from cudnn.sdpa.fwd.heuristics import _MAX_SETS_PER_ENGINE, recommend +from cudnn.sdpa.graph_analyzer import SdpaGraphFacts + +_D128 = "sdpa_fwd_prefill_sm100_d128" +_OFFERED = {_D128: 20500, "sdpa_fwd_prefill_sm100_d256": 20501} + + +def _facts(**over): + base = dict( + b=1, + h_q=4, + h_kv=4, + s_q=128, + s_kv=8192, + d_qk=128, + d_v=128, + dtype=cudnn.data_type.HALF, + causal=True, + device_cc=(10, 0), + device_sm_count=148, + ) + base.update(over) + return SdpaGraphFacts(**base) + + +@pytest.mark.L0 +def test_recommend_emits_multiple_complete_sets_per_engine(): + plans = recommend("A", _facts(), _OFFERED) + d128 = [p for p in plans if p.engine_id == 20500] + assert len(d128) >= 3, "expected sched + split runners behind the primary" + for p in d128: + k = p.knobs + # Complete assignment: every axis the row declares carries a value. + assert None not in (k.sched_policy, k.tile_m, k.tile_n, k.cga, k.split_kv) + assert p.mode is None and p.cpp_index is None + assert len({p.knobs for p in d128}) == len(d128), "duplicate knob sets emitted" + assert len(d128) <= _MAX_SETS_PER_ENGINE + + +@pytest.mark.L0 +def test_recommend_primary_reproduces_the_derived_scheduler(): + # Behavior preservation: the first set carries exactly what the adapter's + # internal derivation historically chose (causal + small working set -> + # LPT_L2; mask-free -> NATURAL with no sched runners). + causal = recommend("A", _facts(), _OFFERED) + assert causal[0].knobs.sched_policy == 2 # SCHED_LPT_L2 + dense = recommend("A", _facts(causal=False), _OFFERED) + dense_d128 = [p for p in dense if p.engine_id == 20500] + assert dense_d128[0].knobs.sched_policy == 0 # SCHED_NATURAL + assert all(p.knobs.sched_policy == 0 for p in dense_d128), "mask-free graphs gain nothing from LPT runners" + + +@pytest.mark.L0 +def test_recommend_split_is_a_runner_up_and_respects_structure(): + plans = [p for p in recommend("A", _facts(), _OFFERED) if p.engine_id == 20500] + assert plans[0].knobs.split_kv == 1, "no-split stays the default winner until sweeps flip it" + assert any(p.knobs.split_kv > 1 for p in plans), "underfilled decode-ish grid must offer a split runner" + for bad in (dict(has_sink=True), dict(thd=True, padded=True), dict(padded=True), dict(s_q=8192)): + got = [p for p in recommend("A", _facts(**bad), _OFFERED) if p.engine_id == 20500] + assert all(p.knobs.split_kv == 1 for p in got), f"split emitted under {bad}" + + +@pytest.mark.L0 +def test_recommend_every_set_is_admissible(): + facts = _facts() + for p in recommend("A", facts, _OFFERED): + spec = next(s for s in engines.ENGINE_SPECS if _OFFERED.get(s.name) == p.engine_id) + assert engines.mismatch(spec.capabilities, facts, p.knobs) is None + + +@pytest.mark.L0 +def test_assemble_strips_mode_dedups_and_our_proposals_lead(): + """Placement is the SHARED layer's job (engines/heuristics._assemble): + proposals lead the backend's entries inside each mode block by standing + assumption, the delegating entry never leads an OPENSOURCE block, one + config repeated across blocks keeps its first position, and no final + entry carries a mode.""" + ours = [PlanConfig(20500, "set-a"), PlanConfig(20500, "set-b")] + backend = [ + PlanConfig(-1, None), # delegating (mode None) + PlanConfig(7, {"k": 1}, cpp_index=0, mode=cudnn.heur_mode.A), + PlanConfig(7, {"k": 1}, cpp_index=1, mode=cudnn.heur_mode.FALLBACK), # same config, later block + ] + final = _assemble([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK], lambda kind: ours if kind == "A" else [], backend) + assert all(p.mode is None for p in final), "mode must never reach final entries" + assert [p.engine_id for p in final[:2]] == [20500, 20500], "our proposals lead the backend inside the block" + assert sum(1 for p in final if p.engine_id == 7) == 1, "one backend config repeated across modes must dedup" + assert next(p for p in final if p.engine_id == 7).cpp_index == 0, "first position wins" + assert sum(1 for p in final if p.engine_id == -1) == 1 + # OPENSOURCE: ours + delegating, and never the backend's own entries. + oss = _assemble([cudnn.heur_mode.OPENSOURCE], lambda kind: ours, backend) + assert [p.engine_id for p in oss] == [20500, 20500, -1] + + +@pytest.mark.L0 +def test_fallback_kind_is_least_demanding(): + for p in recommend("FALLBACK", _facts(), _OFFERED): + assert p.knobs.split_kv == 1 + assert p.knobs.sched_policy == 0 # SCHED_NATURAL + + +# --------------------------------------------------------------------------- +# Executable tier — SM100 graph path +# --------------------------------------------------------------------------- + + +def _is_sm100() -> bool: + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability(0) + return major == 10 + + +def _dsl_available() -> bool: + try: + import cutlass.experimental # noqa: F401 + except ImportError: + return False + return True + + +def _build_decodeish_graph(): + B, H, SQ, SKV, D = 1, 4, 128, 8192, 128 + g = cudnn.pygraph( + io_data_type=cudnn.data_type.HALF, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + q = g.tensor(dim=(B, H, SQ, D), stride=(SQ * H * D, D, H * D, 1), data_type=cudnn.data_type.HALF, name="q") + k = g.tensor(dim=(B, H, SKV, D), stride=(SKV * H * D, D, H * D, 1), data_type=cudnn.data_type.HALF, name="k") + v = g.tensor(dim=(B, H, SKV, D), stride=(SKV * H * D, D, H * D, 1), data_type=cudnn.data_type.HALF, name="v") + o, st = g.sdpa(name="sdpa", q=q, k=k, v=v, attn_scale=1.0 / math.sqrt(D), is_inference=False, use_causal_mask=True) + o.set_output(True).set_dim((B, H, SQ, D)).set_stride((SQ * H * D, D, H * D, 1)) + st.set_output(True).set_data_type(cudnn.data_type.FLOAT) + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + return g, (q, k, v, o, st), (B, H, SQ, SKV, D) + + +@pytest.mark.L1 +@pytest.mark.skipif(not (_is_sm100() and _dsl_available()), reason="needs an SM100 device and nvidia-cutlass-dsl") +def test_split_kv_plan_pinned_by_name_matches_reference(): + """Issue F-2 regression: the split plan is graph-reachable, carves its + slabs from the caller workspace, and recombines exactly.""" + g, (q, k, v, o, st), (B, H, SQ, SKV, D) = _build_decodeish_graph() + # The split value depends on this device's SM count — ask the chooser + # rather than hard-coding one that only holds at one part's geometry. + from cudnn._device import device_info + from cudnn.sdpa.fwd.heuristics import choose_split_kv + + want = choose_split_kv( + q_tiles=-(-SQ // 256), + heads_q=H, + batch=B, + kv_tiles=-(-SKV // 128), + sm_count=device_info(torch.cuda.current_device()).sm_count, + ctas_per_tile=2, + max_split=4, + ) + if want == 1: + pytest.skip("this part is small enough that the shape already fills it") + names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + d128 = [n for n in names if "sm100_d128" in n] + assert len(d128) >= 3, f"expected knob-suffixed duplicates of the d128 cell: {d128}" + split_idx = next(i for i, n in enumerate(names) if "sm100_d128" in n and f"split_kv={want}" in n) + g.select_plan(split_idx) + g.check_support() + g.build_plans() + assert g.get_workspace_size() > 0, "the split plan must report its partial-slab workspace" + + torch.manual_seed(0) + q_gpu = torch.randn(B, SQ, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) + k_gpu = torch.randn(B, SKV, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) + v_gpu = torch.randn(B, SKV, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) + o_gpu = torch.empty(B, SQ, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) + st_gpu = torch.empty(B, H, SQ, 1, device="cuda", dtype=torch.float32) + ws = torch.empty(g.get_workspace_size(), device="cuda", dtype=torch.uint8) + g.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu, st: st_gpu}, ws) + torch.cuda.synchronize() + + s = torch.einsum("bhqd,bhkd->bhqk", q_gpu.float(), k_gpu.float()) / math.sqrt(D) + i = torch.arange(SQ, device="cuda").view(SQ, 1) + j = torch.arange(SKV, device="cuda").view(1, SKV) + s = s.masked_fill(j > i, float("-inf")) + torch.testing.assert_close(o_gpu, torch.einsum("bhqk,bhkd->bhqd", torch.softmax(s, dim=-1), v_gpu.float()).half(), atol=5e-2, rtol=3e-2) + torch.testing.assert_close(st_gpu.squeeze(-1), torch.logsumexp(s, dim=-1), atol=2e-3, rtol=2e-3) + + # Autotune replay: the split entry round-trips through (engine_id, knobs). + eng_id, knobs = g.get_engine_and_knobs_at_index(split_idx) + assert knobs.split_kv == want + g2, handles2, _ = _build_decodeish_graph() + cfg = g2.create_execution_plan(eng_id, knobs) + assert cfg is not None + + +@pytest.mark.L1 +@pytest.mark.skipif(not (_is_sm100() and _dsl_available()), reason="needs an SM100 device and nvidia-cutlass-dsl") +def test_runner_up_sched_plan_builds_and_matches_the_winner(): + """select_plan on a runner-up knob set compiles the adapter with exactly + that set and executes correctly — honored, not silently degraded.""" + g, (q, k, v, o, st), (B, H, SQ, SKV, D) = _build_decodeish_graph() + names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + nat_idx = next(i for i, n in enumerate(names) if "sm100_d128" in n and "sched_policy=0" in n and "split_kv=1" in n) + g.select_plan(nat_idx) + g.check_support() + g.build_plans() + eng_id, knobs = g.get_engine_and_knobs_at_index(nat_idx) + assert knobs.sched_policy == 0 and knobs.split_kv == 1 + + torch.manual_seed(0) + q_gpu = torch.randn(B, SQ, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) + k_gpu = torch.randn(B, SKV, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) + v_gpu = torch.randn(B, SKV, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) + o_gpu = torch.empty(B, SQ, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) + st_gpu = torch.empty(B, H, SQ, 1, device="cuda", dtype=torch.float32) + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + g.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu, st: st_gpu}, ws) + torch.cuda.synchronize() + s = torch.einsum("bhqd,bhkd->bhqk", q_gpu.float(), k_gpu.float()) / math.sqrt(D) + i = torch.arange(SQ, device="cuda").view(SQ, 1) + j = torch.arange(SKV, device="cuda").view(1, SKV) + s = s.masked_fill(j > i, float("-inf")) + torch.testing.assert_close(o_gpu, torch.einsum("bhqk,bhkd->bhqd", torch.softmax(s, dim=-1), v_gpu.float()).half(), atol=5e-2, rtol=3e-2) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py index 22619c043..4934d222e 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py @@ -821,3 +821,178 @@ def test_split_kv_padded_q_trim(flavor, splits): assert got[dead.expand_as(got)].abs().max().item() == 0.0, "trimmed Q rows are not zero" live = ~dead assert (got - ref)[live.expand_as(got)].abs().max().item() <= 2e-2 + + +# --- the adapter honors the heuristic's split knob --------------------------- + + +def _expected_split(b, h_q, s_q, s_kv, *, rows_per_tile=512, ctas_per_tile=2, kv_tile=128): + """What the chooser asks for on THIS device, so the tests assert the knob + route delivers it rather than hard-coding a split that only holds at one + SM count.""" + from cudnn._device import device_info + from cudnn.sdpa.fwd.heuristics import choose_split_kv + + return choose_split_kv( + q_tiles=-(-s_q // rows_per_tile), + heads_q=h_q, + batch=b, + kv_tiles=-(-s_kv // kv_tile), + sm_count=device_info(torch.cuda.current_device()).sm_count, + ctas_per_tile=ctas_per_tile, + ) + + +def _api_case(b, h_q, h_kv, s_q, s_kv, *, with_lse=False, workspace=True): + """Drive SdpaFwdDslSm100 the way the graph path does — the chooser's value + arrives as the explicit ``split_kv`` constructor knob, exactly as + ``lower_dsl_prefill`` forwards a plan's knobs; return (split, O, ref).""" + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)): + pytest.skip("half-precision SM100 prefill requires cc10.0 / cc10.3") + d = 128 + dev = "cuda" + torch.manual_seed(0) + q = torch.randn(b, h_q, s_q, d, device=dev, dtype=torch.float16) # BHSD samples + k = torch.randn(b, h_kv, s_kv, d, device=dev, dtype=torch.float16) + v = torch.randn(b, h_kv, s_kv, d, device=dev, dtype=torch.float16) + o = torch.zeros_like(q) + lse = torch.zeros(b, h_q, s_q, device=dev, dtype=torch.float32) if with_lse else None + + split_knob = _expected_split(b, h_q, s_q, s_kv) + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, split_kv=split_knob) + assert api.check_support() + split = api.split_kv + assert split == split_knob, "the knob is honored verbatim, never re-derived" + ws_bytes = api.scratch_workspace_bytes() + api.compile() + ws = torch.empty(ws_bytes, dtype=torch.uint8, device=dev) if (workspace and ws_bytes) else None + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse, workspace=ws) + torch.cuda.synchronize() + + qb, kb, vb = q.float(), k.float(), v.float() + if h_q != h_kv: + kb = kb.repeat_interleave(h_q // h_kv, dim=1) + vb = vb.repeat_interleave(h_q // h_kv, dim=1) + p = torch.softmax(torch.matmul(qb, kb.transpose(-1, -2)) / math.sqrt(d), dim=-1) + return split, o.float(), torch.matmul(p, vb), ws_bytes + + +@pytest.mark.L0 +def test_api_splits_a_decode_shape_and_is_correct(): + """8 heads over a long KV run cannot fill the part: the chooser wants a + split, the knob delivers it, the adapter sizes its own workspace and still + matches fp32. S_kv is the smallest that still splits -- the fp32 reference + is O(h * s_q * s_kv) and this is L0.""" + split, got, ref, ws_bytes = _api_case(1, 8, 1, 512, 16384) + assert split == _expected_split(1, 8, 512, 16384) > 1 + assert ws_bytes > 0, "a split needs the split-major O/LSE partials in workspace" + assert (got - ref).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +def test_api_does_not_split_a_full_chip(): + """A launch that fills the machine is left alone. The expectation comes from + the chooser on THIS device: whether 2048x16 fills it depends on the SM count, + so hard-coding split==1 would fail on a smaller SM100 part for a device + reason rather than a policy one.""" + want = _expected_split(1, 16, 2048, 8192) + split, got, ref, ws_bytes = _api_case(1, 16, 16, 2048, 8192) + assert split == want + assert (ws_bytes > 0) == (split > 1), "workspace is needed exactly when we split" + assert (got - ref).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +@pytest.mark.parametrize("workspace", [True, False], ids=["carved", "standalone"]) +def test_api_split_with_and_without_workspace(workspace): + """With a workspace the partials are carved from it; without one they are + torch-allocated (standalone use). Same answer either way.""" + split, got, ref, _ = _api_case(1, 8, 1, 512, 16384, workspace=workspace) + assert split > 1 + assert (got - ref).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +def test_api_split_writes_the_recombined_lse(): + """A Stats output under a split comes from the combine, not from any one + chunk: the per-chunk LSE is compiled in even when the caller wants none.""" + split, got, ref, _ = _api_case(1, 8, 1, 512, 16384, with_lse=True) + assert split > 1 + assert (got - ref).abs().max().item() <= 2e-2 + + +# --- the adapter splits the FP8 family too ---------------------------------- + + +def _api_fp8_case(h_q, h_kv, s_q, s_kv, *, mx): + """FP8 / MXFP8 through the adapter; returns (split, O, O_unsplit, amax).""" + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + _cc = torch.cuda.get_device_capability() + if _cc not in ((10, 0), (10, 3)) and not (_cc == (10, 7) and not mx): + pytest.skip("MXFP8 requires cc10.0 / cc10.3; per-tensor FP8 also runs on cc10.7") + b, d, dev = 1, 128, "cuda" + + def build(force_one): + torch.manual_seed(0) + if mx: + from sdpa.mxfp8_quant import quantize_to_mxfp8 + + def qz(shape): + r = torch.randn(*shape, device=dev) * 0.5 + a, _adq, aswz, *_ = quantize_to_mxfp8(r, shape[0], shape[1], shape[2], shape[3]) + return a.reshape(*shape), aswz + + q, sf_q = qz((b, h_q, s_q, d)) + k, sf_k = qz((b, h_kv, s_kv, d)) + v, sf_v = qz((b, h_kv, s_kv, d)) + extra = dict(sf_q=sf_q, sf_k=sf_k, sf_v=sf_v) + kw = {} + else: + mk = lambda *sh: (torch.randn(*sh, device=dev) * 0.5).clamp(-448, 448).to(torch.float8_e4m3fn) + q, k, v = mk(b, h_q, s_q, d), mk(b, h_kv, s_kv, d), mk(b, h_kv, s_kv, d) + one = lambda: torch.ones(1, dtype=torch.float32, device=dev) + extra = dict(descale_q=one(), descale_k=one(), descale_v=one(), scale_o=one()) + kw = dict(pertensor_fp8=True) + + o = torch.zeros(b, h_q, s_q, d, device=dev, dtype=torch.float16) # HALF out + amax = torch.zeros(1, dtype=torch.float32, device=dev) + split_knob = 1 if force_one else _expected_split(b, h_q, s_q, s_kv) + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, dtype_o=torch.float16, split_kv=split_knob, **kw) + assert api.check_support() + split = api.split_kv + ws_bytes = api.scratch_workspace_bytes() + api.compile() + ws = torch.empty(ws_bytes, dtype=torch.uint8, device=dev) if ws_bytes else None + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, amax_o=amax, workspace=ws, **extra) + torch.cuda.synchronize() + return split, o.float().clone(), amax.item() + + _, o_one, _ = build(True) + split, o_split, amax = build(False) + return split, o_split, o_one, amax + + +@pytest.mark.L0 +@pytest.mark.parametrize("mx", [False, True], ids=["fp8", "mxfp8"]) +def test_api_splits_the_fp8_family(mx): + """A decode-shaped FP8 graph with a half output splits, and the split is + numerically neutral against the same graph forced to split_kv=1.""" + split, got, unsplit, _ = _api_fp8_case(8, 1, 512, 16384, mx=mx) + assert split > 1 + assert (got - unsplit).abs().max().item() <= 5e-2 + + +@pytest.mark.L0 +@pytest.mark.parametrize("mx", [False, True], ids=["fp8", "mxfp8"]) +def test_api_fp8_amax_describes_the_recombined_output(mx): + """The per-split epilogues stand down under a split; amax_o has to come + from the combine, over the RECOMBINED O. Maxing the partials instead + over-reports, so compare against the output the caller receives.""" + split, got, _unsplit, amax = _api_fp8_case(8, 1, 512, 16384, mx=mx) + assert split > 1 + true_amax = got.abs().max().item() + assert amax >= true_amax * 0.99, f"amax {amax} under-reports |O| {true_amax}" + assert amax <= true_amax * 1.01, f"amax {amax} over-reports |O| {true_amax} — taken over partials?" diff --git a/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm120.py new file mode 100644 index 000000000..8af3b6bcf --- /dev/null +++ b/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm120.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""KV split on SM120, driven through the adapter rather than the template. + +KNOB ROUTE: the adapter never decides the split itself — the heuristic's +chooser picks a value and the test passes it explicitly as the ``split_kv`` +constructor knob, exactly as ``lower_dsl_prefill`` forwards a plan's knobs. +The chooser and the two-launch execute are shared with SM100; what differs here +is the geometry (one CTA per tile, no cluster) and that the config bars a split +under the flattened LPT schedulers. +""" + +import math + +import pytest +import torch + +from frost_test_utils import requires_dsl + +pytestmark = [requires_dsl, pytest.mark.L0] + + +def _expected_split(api): + """What the chooser asks for on THIS device, from the adapter's OWN tile + geometry. SM120 runs one CTA per tile (no cluster), and the adapter may pick + either q_tile, so reading them off `api` keeps the expectation tied to the + launch that actually happens rather than to one part's SM count.""" + from cudnn._device import device_info + from cudnn.sdpa.fwd.heuristics import choose_split_kv + + return choose_split_kv( + q_tiles=-(-api.s_q_max // api.q_tile), + heads_q=api.h_q, + batch=api.batch_size, + kv_tiles=-(-api.s_k_max // api.kv_tile), + sm_count=device_info(torch.cuda.current_device()).sm_count, + ctas_per_tile=1, + ) + + +def _sm120_case(h_q, h_kv, s_q, s_kv, *, with_lse=False, workspace=True, causal=False): + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm120 + + if torch.cuda.get_device_capability()[0] != 12: + pytest.skip("SM120 part required") + b, d, dev = 1, 128, "cuda" + torch.manual_seed(0) + q = torch.randn(b, h_q, s_q, d, device=dev, dtype=torch.float16) + k = torch.randn(b, h_kv, s_kv, d, device=dev, dtype=torch.float16) + v = torch.randn(b, h_kv, s_kv, d, device=dev, dtype=torch.float16) + o = torch.zeros_like(q) + lse = torch.zeros(b, h_q, s_q, device=dev, dtype=torch.float32) if with_lse else None + + kw = dict(is_causal=True) if causal else {} + # Probe pass: the chooser reads the adapter's own tile geometry. + probe = SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, **kw) + assert probe.check_support() + expected = _expected_split(probe) + # Knob route: the chosen split arrives as an explicit constructor knob + # (split sets ride SCHED_NATURAL — the config bars a split under the LPT + # remaps a causal graph would otherwise derive). + if expected > 1: + kw.update(split_kv=expected, sched_policy=0) + api = SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, **kw) + assert api.check_support() + split = api.split_kv + ws_bytes = api.scratch_workspace_bytes() + api.compile() + ws = torch.empty(ws_bytes, dtype=torch.uint8, device=dev) if (workspace and ws_bytes) else None + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse, workspace=ws) + torch.cuda.synchronize() + + qb, kb, vb = q.float(), k.float(), v.float() + if h_q != h_kv: + kb = kb.repeat_interleave(h_q // h_kv, dim=1) + vb = vb.repeat_interleave(h_q // h_kv, dim=1) + scores = torch.matmul(qb, kb.transpose(-1, -2)) / math.sqrt(d) + if causal: + i = torch.arange(s_q, device=scores.device).view(s_q, 1) + j = torch.arange(s_kv, device=scores.device).view(1, s_kv) + scores = scores.masked_fill(j > i, float("-inf")) + p = torch.softmax(scores, dim=-1) + return split, o.float(), torch.matmul(p, vb), ws_bytes, expected + + +def test_sm120_splits_a_decode_shape(): + """A decode shape the chooser wants split -- and the adapter must honor it.""" + split, got, ref, ws_bytes, expected = _sm120_case(8, 1, 128, 32768) + if expected == 1: + pytest.skip("this part is small enough that the shape already fills it") + assert split == expected > 1 + assert ws_bytes > 0 + assert (got - ref).abs().max().item() <= 2e-2 + + +def test_sm120_does_not_split_a_full_part(): + """A launch that fills the part is left alone. Whether 1024x64 fills it + depends on the SM count, so the expectation comes from the chooser rather + than a fixed 1 that only holds on one device.""" + split, got, ref, ws_bytes, expected = _sm120_case(64, 8, 1024, 16384) + assert split == expected + assert (ws_bytes > 0) == (split > 1), "workspace is needed exactly when we split" + assert (got - ref).abs().max().item() <= 2e-2 + + +@pytest.mark.parametrize("workspace", [True, False], ids=["carved", "standalone"]) +def test_sm120_split_with_and_without_workspace(workspace): + split, got, ref, _, expected = _sm120_case(8, 1, 128, 32768, workspace=workspace) + assert split == expected + assert (got - ref).abs().max().item() <= 2e-2 + + +def test_sm120_split_writes_the_recombined_lse(): + split, got, ref, _, expected = _sm120_case(8, 1, 128, 32768, with_lse=True) + assert split == expected + assert (got - ref).abs().max().item() <= 2e-2 + + +def test_sm120_causal_split_requires_the_natural_scheduler(): + """The config bars a split under the LPT remaps a causal graph derives. + Knob-route contract, both directions: a split WITHOUT an explicit + scheduler lets the adapter derive LPT and must fail loudly at compile + (honored-or-error, never silently degraded), while split + explicit + NATURAL — what the heuristic actually emits — compiles and matches.""" + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm120 + + if torch.cuda.get_device_capability()[0] != 12: + pytest.skip("SM120 part required") + b, d, dev = 1, 128, "cuda" + torch.manual_seed(0) + q = torch.randn(b, 8, 512, d, device=dev, dtype=torch.float16) + k = torch.randn(b, 1, 8192, d, device=dev, dtype=torch.float16) + v = torch.randn(b, 1, 8192, d, device=dev, dtype=torch.float16) + o = torch.zeros_like(q) + + bad = SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, is_causal=True, split_kv=2) + assert bad.check_support() + with pytest.raises(ValueError, match="split_kv"): + bad.compile() # derived LPT + split: the config backstop rejects + + split, got, ref, _, expected = _sm120_case(8, 1, 512, 8192, causal=True) + if expected == 1: + pytest.skip("this part is small enough that the causal shape already fills it") + assert split == expected > 1, "the causal arm must actually exercise the split" + assert (got - ref).abs().max().item() <= 2e-2 diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index 53bd23116..f0b2c56e1 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -562,9 +562,19 @@ def test_knob_request_within_domain_keeps_engine_eligible(): def test_knob_request_outside_domain_rejects_engine(): - # No engine advertises LPT scheduling yet: honored or ineligible, never degraded. + # A value no row's domain contains: honored or ineligible, never degraded. g = _mk_eligible_graph() - assert not _eligible(g, engines.SdpaFwdKnobs(sched_policy=1)) + assert not _eligible(g, engines.SdpaFwdKnobs(sched_policy=99)) + # softmax_precision is a framework axis with no serving kernel yet: any + # explicit request declines everywhere. + assert not _eligible(g, engines.SdpaFwdKnobs(softmax_precision=1)) + + +def test_knob_request_lpt_sched_is_in_domain(): + # The SM100 rows advertise all three scheduler policies (the static/CLC + # remap serves them); an explicit LPT request stays eligible. + g = _mk_eligible_graph() + assert engines.engine_name(512) in _eligible(g, engines.SdpaFwdKnobs(sched_policy=1)) def test_knob_request_unsupported_tile_rejects_engine(): @@ -974,7 +984,15 @@ def test_sm120_knob_domains(monkeypatch): g = _mk_sm120_graph(use_causal_mask=True) assert _SM120 in _eligible(g, engines.SdpaFwdKnobs(tile_m=64, tile_n=64, cga=1)) assert not _eligible(g, engines.SdpaFwdKnobs(cga=2)) - assert not _eligible(g, engines.SdpaFwdKnobs(sched_policy=1)) + # All three scheduler policies are in the SM120 domain (static-grid + # remap); a value outside the vocabulary still declines. + assert _SM120 in _eligible(g, engines.SdpaFwdKnobs(sched_policy=1)) + assert not _eligible(g, engines.SdpaFwdKnobs(sched_policy=99)) + # split_kv: the SM120 row serves {1, 2, 4} (inline chunking + the shared + # combine); a value outside the domain still declines. + assert _SM120 in _eligible(g, engines.SdpaFwdKnobs(split_kv=1)) + assert _SM120 in _eligible(g, engines.SdpaFwdKnobs(split_kv=4)) + assert not _eligible(g, engines.SdpaFwdKnobs(split_kv=8)) # --------------------------------------------------------------------------- diff --git a/test/python/sdpa/frost/test_split_kv_heuristic.py b/test/python/sdpa/frost/test_split_kv_heuristic.py new file mode 100644 index 000000000..84cebf106 --- /dev/null +++ b/test/python/sdpa/frost/test_split_kv_heuristic.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the KV-split chooser. + +Pure arithmetic over (tile count, KV length, SM count) — no GPU, no kernel +build, so these run anywhere and pin the POLICY rather than any one device. +The B200 numbers below are the shapes the split was built for, taken from the +d128 geometry: a cga2 cluster covers TILES_Q * TILE_M * CTA_MMA = 512 Q rows on +2 CTAs, and TILE_N = 128 sets the KV tile. +""" + +import pytest + +from cudnn.sdpa.fwd.engines import Capabilities, SdpaFwdKnobs, mismatch +from cudnn.sdpa.fwd.heuristics import _SPLIT_KV_MAX, _SPLIT_KV_MIN_TILES, choose_split_kv + +# Pure arithmetic — no device, no kernel build — so every case is L0. +pytestmark = pytest.mark.L0 + +B200_SMS = 148 + + +def _d128_cga2(s_q, s_kv, heads, batch, sm_count=B200_SMS, **kw): + """choose_split_kv for the d128 cga2 geometry.""" + return choose_split_kv( + q_tiles=-(-s_q // 512), + heads_q=heads, + batch=batch, + kv_tiles=-(-s_kv // 128), + sm_count=sm_count, + ctas_per_tile=2, + **kw, + ) + + +# --- the case the feature exists for -------------------------------------- + + +def _single_wave(base_ctas, split, sm_count=B200_SMS): + """Does the split leave the launch inside one wave? True for an UNDER-full + launch, where the idle SMs are free to fill; not a global rule (see + test_may_exceed_the_sm_count_to_smooth_a_tail).""" + return base_ctas * split <= sm_count + + +def test_decode_shape_splits(): + """S_q=128, S_kv=32K, H=16, B=1 — 32 CTAs on a 148-SM part.""" + split = _d128_cga2(128, 32768, 16, 1) + assert split > 1, "a launch filling 32 of 148 SMs must split" + assert _single_wave(32, split), "the chosen split must stay inside one wave" + assert not _single_wave(32, split * 2), "and be the LARGEST power of two that does" + + +def test_more_heads_need_less_split(): + """Heads are already parallelism: the more there are, the less splitting.""" + few = _d128_cga2(128, 32768, 4, 1) + many = _d128_cga2(128, 32768, 64, 1) + assert few >= many + + +def test_bigger_batch_needs_less_split(): + assert _d128_cga2(128, 32768, 8, 1) >= _d128_cga2(128, 32768, 8, 8) + + +def test_smaller_chip_needs_less_split(): + """The same shape on a smaller part needs fewer splits to fill it.""" + big = _d128_cga2(128, 32768, 16, 1, sm_count=148) + small = _d128_cga2(128, 32768, 16, 1, sm_count=48) + assert small <= big + + +# --- cases that must NOT split -------------------------------------------- + + +def test_chip_already_full_does_not_split(): + """Long prefill: 4096 Q rows x 16 heads is 8 * 16 * 2 = 256 CTAs > 148.""" + assert _d128_cga2(4096, 2048, 16, 1) == 1 + + +def test_nearly_full_chip_does_not_split(): + """128 CTAs on 148 SMs — 86% — must not split: doubling tips it into a + second partial wave, so the makespan is flat and only the reduction is + added. Measured, not assumed.""" + base_ctas = 4 * 16 * 2 # s_q=2048 -> 4 q tiles, 16 heads, cga2 + assert base_ctas == 128 + assert _d128_cga2(2048, 65536, 16, 1) == 1 + + +def test_exactly_full_does_not_split(): + """base_ctas == sm_count is 'filled' — no reduction for zero gain.""" + assert choose_split_kv(q_tiles=1, heads_q=B200_SMS, batch=1, kv_tiles=256, sm_count=B200_SMS) == 1 + + +def test_single_kv_tile_cannot_split(): + assert _d128_cga2(128, 128, 1, 1) == 1 + + +def test_short_kv_does_not_over_split(): + """Splits below _SPLIT_KV_MIN_TILES KV tiles are prologue-dominated.""" + kv_tiles = 4 + split = choose_split_kv(q_tiles=1, heads_q=1, batch=1, kv_tiles=kv_tiles, sm_count=B200_SMS) + assert split <= kv_tiles // _SPLIT_KV_MIN_TILES + + +def test_degenerate_inputs_do_not_split(): + for kw in ( + {"q_tiles": 0}, + {"heads_q": 0}, + {"batch": 0}, + {"kv_tiles": 0}, + {"sm_count": 0}, + {"sm_count": -1}, + ): + args = {"q_tiles": 1, "heads_q": 1, "batch": 1, "kv_tiles": 256, "sm_count": B200_SMS} + args.update(kw) + assert choose_split_kv(**args) == 1, f"{kw} must fall back to no split" + + +# --- invariants over a sweep ---------------------------------------------- + + +@pytest.mark.parametrize("s_kv", [1024, 4096, 16384, 32768, 131072]) +@pytest.mark.parametrize("heads", [1, 2, 8, 16, 64]) +def test_invariants(s_kv, heads): + kv_tiles = -(-s_kv // 128) + split = _d128_cga2(128, s_kv, heads, 1) + assert 1 <= split <= _SPLIT_KV_MAX + assert split <= kv_tiles, "more splits than KV tiles would leave empty splits" + if split > 1: + assert -(-kv_tiles // split) >= _SPLIT_KV_MIN_TILES + + +def test_longer_kv_never_needs_fewer_splits(): + """Monotone in KV length: a longer loop is never served by less splitting.""" + prev = 0 + for s_kv in (256, 512, 1024, 2048, 4096, 8192, 16384, 32768): + split = _d128_cga2(128, s_kv, 16, 1) + assert split >= prev, f"S_kv={s_kv} chose {split} after {prev}" + prev = split + + +@pytest.mark.parametrize("heads", [1, 2, 3, 5, 8, 11, 16, 32, 64]) +@pytest.mark.parametrize("s_kv", [4096, 65536, 131072]) +def test_choice_is_always_a_power_of_two(heads, s_kv): + """split_kv is a compile-cache key, so the set is bounded to {1,2,4,8,16}.""" + split = _d128_cga2(512, s_kv, heads, 1) + assert split & (split - 1) == 0, f"{split} is not a power of two" + assert split <= _SPLIT_KV_MAX + + +def test_may_exceed_the_sm_count_to_smooth_a_tail(): + """Over-subscribing the SMs is allowed, and sometimes required: 160 CTAs on + 148 SMs already wastes most of a second wave, and splitting finer shrinks + that tail rather than adding a wave.""" + split = choose_split_kv(q_tiles=1, heads_q=80, batch=1, kv_tiles=512, sm_count=148, ctas_per_tile=2) + assert split > 1 + assert 160 * split > 148, "this shape is exactly the case that wants over-subscription" + + +def test_exactly_balanced_launch_never_splits(): + """base_ctas = k * sm_count has no tail to recover, so every split is pure + overhead. Provable rather than fitted.""" + for sm_count in (16, 48, 108, 148, 256): + for k in (1, 2, 3, 4): + for kv_tiles in (16, 64, 256, 512, 1024): + split = choose_split_kv(q_tiles=1, heads_q=k * sm_count, batch=1, kv_tiles=kv_tiles, sm_count=sm_count, ctas_per_tile=1) + assert split == 1, f"base={k * sm_count} == {k}x{sm_count} SMs: nothing to smooth, got {split}" + + +# (base_ctas, chosen split) pinned against a per-split sweep on B300 (148 SMs, +# d128, 512 KV tiles). A change that moves any of these is a policy change and +# needs its own measurement. +_B300_FIT = [(8, 16), (16, 8), (32, 4), (64, 2), (88, 8), (100, 4), (120, 1), (128, 1), (150, 4), (160, 4), (200, 2), (296, 1)] + + +@pytest.mark.parametrize("base_ctas,expected", _B300_FIT, ids=[f"{b}ctas" for b, _ in _B300_FIT]) +def test_reproduces_the_b300_fit(base_ctas, expected): + got = choose_split_kv(q_tiles=1, heads_q=base_ctas // 2, batch=1, kv_tiles=512, sm_count=148, ctas_per_tile=2) + assert got == expected + + +def test_response_is_not_monotone_in_occupancy(): + """At 88 and 100 CTAs split 2 loses while 4 and 8 win, because 2 lands just + over a wave boundary and 4 does not. The chooser must search, not + interpolate.""" + assert choose_split_kv(q_tiles=1, heads_q=44, batch=1, kv_tiles=512, sm_count=148, ctas_per_tile=2) != 2 + assert choose_split_kv(q_tiles=1, heads_q=50, batch=1, kv_tiles=512, sm_count=148, ctas_per_tile=2) != 2 + + +def test_max_split_is_respected(): + assert _d128_cga2(128, 1 << 20, 1, 1, max_split=4) <= 4 + + +# --- the knob plumbing ------------------------------------------------------ + + +@pytest.mark.parametrize("requested", [2, 4, 8, 16]) +def test_split_request_outside_the_domain_makes_the_engine_ineligible(requested): + """The default row serves only split_kv=1 ("off"): a split request on a row + whose lowering has no split path is honored-or-ineligible, never silently + dropped.""" + caps = Capabilities(sm_lo=100, sm_hi=100, phase="prefill", d_qk=frozenset({128}), d_v=frozenset({128})) + assert caps.split_kvs == frozenset({1}) + why = mismatch(caps, _facts(), SdpaFwdKnobs(split_kv=requested)) + assert why is not None and "split_kv" in why + + +def test_no_split_and_explicit_one_leave_the_engine_eligible(): + """No preference passes; so does an EXPLICIT split_kv=1 — "do not split" + is a real point on the axis, not an ignored request.""" + caps = Capabilities(sm_lo=100, sm_hi=100, phase="prefill", d_qk=frozenset({128}), d_v=frozenset({128})) + for knobs in (SdpaFwdKnobs(split_kv=None), SdpaFwdKnobs(split_kv=1)): + why = mismatch(caps, _facts(), knobs) or "" + assert "split_kv" not in why + + +def _facts(): + from cudnn.sdpa import graph_analyzer as ga + + return ga.SdpaGraphFacts() + + +def test_split_declines_when_the_kv_tail_needs_synthesized_padding(): + """A ragged S_kv on a skv_tail_via_padding row is served through the + padded kernel path (synthesized per-batch KV lengths) — the one path the + split cannot ride. The gate must mirror lower_dsl_prefill's predicate so + the plan is never listed, not declined at build.""" + from cudnn.sdpa import graph_analyzer as ga + + caps = Capabilities( + sm_lo=100, + sm_hi=100, + phase="prefill", + d_qk=frozenset({128}), + d_v=frozenset({128}), + skv_tail_via_padding=True, + split_kvs=frozenset({1, 2, 4}), + ) + ragged = ga.SdpaGraphFacts(s_q=128, s_kv=1000) # 1000 % 128 != 0, mask-free + why = mismatch(caps, ragged, SdpaFwdKnobs(split_kv=2)) + assert why is not None and "split_kv" in why + # A causal band that provably masks the tail needs no synthesized padding, + # so the same request passes the split gate (later gates may still apply). + covered = ga.SdpaGraphFacts(s_q=128, s_kv=1000, causal=True) + why = mismatch(caps, covered, SdpaFwdKnobs(split_kv=2)) or "" + assert "split_kv" not in why + + +def test_split_domains_match_the_wired_lowerings(): + """Guards the pairing: a row advertises split_kvs > {1} exactly when its + adapter forwards the knob into TemplateParams and launches the combine. + Widening one without the plumbing reintroduces the silently-dropped knob.""" + from cudnn.sdpa.fwd.engines import ENGINE_SPECS + + advertising = {sp.name for sp in ENGINE_SPECS if sp.capabilities.split_kvs != frozenset({1})} + assert advertising == { + "sdpa_fwd_prefill_sm100_d128", + "sdpa_fwd_prefill_sm100_d256", + "sdpa_fwd_prefill_sm100_d512", + "sdpa_fwd_prefill_sm100_d192_d128", + "sdpa_fwd_prefill_sm100_d128_mxfp8", + "sdpa_fwd_prefill_sm100_d128_fp8", + "sdpa_fwd_prefill_sm120", + }, f"split domains drifted from the wired lowerings: {sorted(advertising)}"