diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 1acf50d46..73d4ea9ed 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -242,10 +242,14 @@ class Capabilities: tile_ns: frozenset[int] = frozenset() cgas: frozenset[int] = frozenset() pack_gqas: frozenset[bool] = frozenset({False}) - # 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}) + # Does this row's lowering wire the KV-split path (kernel SplitHelpers + + # adapter carving the partial slabs + launching the combine)? A GATE, not a + # domain: WHICH splits are worth trying is a device-derived search space + # (heuristics.split_kv_candidates), not a per-row constant. Fail-closed — + # a row that accepts split_kv > 1 without the plumbing leaves untouched + # partial slots at lse_partial = 0, which corrupt the combine's + # log-sum-exp rather than raising. + split_kv_supported: bool = False # Shapes whose kernel flavors wire SplitHelpers. None = every flavor in # d_shapes does (f16/SM120). A set = split_kv > 1 is honored only when # the graph's dims are covered by a member (the quantized families wire @@ -297,7 +301,6 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti (knobs.tile_n, capabilities.tile_ns, "tile_n"), (knobs.cga, capabilities.cgas, "cga"), (knobs.pack_gqa, capabilities.pack_gqas, "pack_gqa"), - (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: @@ -305,7 +308,11 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti # members (softmax_precision), and the pybind enum defines no # ordering of its own. return f"requested {label}={value} is outside this engine's domain {sorted(domain, key=int)}" + if knobs.split_kv is not None and knobs.split_kv < 1: + return f"requested split_kv={knobs.split_kv} is not a split count (1 = off)" if knobs.split_kv is not None and knobs.split_kv > 1: + if not capabilities.split_kv_supported: + return "split_kv > 1 is not wired in this engine's lowering" # 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 @@ -526,7 +533,7 @@ def _sm100_spec() -> EngineSpec: # 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}), + split_kv_supported=True, pack_gqas=frozenset({False, True}), ), lower=partial(lower_dsl_prefill, api_type=_SM100), @@ -577,7 +584,7 @@ def _sm100_mxfp8_spec() -> EngineSpec: cgas=frozenset({2}), # The split path also needs a half-precision O (mismatch's # facts x knobs gate) and rides the d128 flavor (split_d_shapes). - split_kvs=frozenset({1, 2, 4}), + split_kv_supported=True, # PackGQA is currently not supported for the MXFP8 SDPA engine: # the F8_128x4 sf_q scale-factor atom bundles 128 rows of ONE # head, so a packed tile's interleaved (token, head) rows cannot @@ -609,7 +616,7 @@ def _sm100_fp8_spec(*, arch: str = "sm100") -> EngineSpec: - softmax_precisions: the f16x2 exponent arm lives only in the SM107 sibling kernel, so only that row admits HALF. FLOAT is the pipeline every flavor already runs. - - split_kvs / split_d_shapes: only the SM100 d128 kernel wires + - split_kv_supported / split_d_shapes: only the SM100 d128 kernel wires SplitHelpers; the SM107 sibling has no split path yet, and the d192x128 file forks its own scheduler and has none either. - sched_policies: the LPT/LPT_L2 remap is not yet ported to the SM107 @@ -678,7 +685,7 @@ def _sm100_fp8_spec(*, arch: str = "sm100") -> EngineSpec: # Split partials reduce in half precision, so mismatch()'s # facts x knobs gate additionally requires a bf16/fp16 O on the # quantized rows; split_d_shapes pins it to the d128 flavor. - split_kvs=frozenset({1}) if rubin_row else frozenset({1, 2, 4}), + split_kv_supported=not rubin_row, split_d_shapes=frozenset({(128, 128)}), pack_gqas=frozenset({False, True}), ), @@ -766,7 +773,7 @@ def _sm120_spec() -> EngineSpec: # (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}), + split_kv_supported=True, 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 ffab44ab8..7afd8b553 100644 --- a/python/cudnn/sdpa/fwd/heuristics.py +++ b/python/cudnn/sdpa/fwd/heuristics.py @@ -84,13 +84,54 @@ def _ceil_div(a: int, b: int) -> int: # --- 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 +# What ONE split's partials cost the combine pass, per wave of combine blocks, +# in units of one KV tile of main-kernel work. The combine's own occupancy +# (blocks/SM of split_combine_sm100) is ABSORBED into this coefficient: it is +# one fixed kernel, so blocks/SM is a constant, and folding it in keeps a +# cuOccupancy query -- which would need a compiled CUfunction -- off the +# planning path. Empirical: re-measure if split_combine_sm100 changes. +_SPLIT_KV_COMBINE_COST = 0.2 + + +def split_kv_candidates(*, sm_count: int, kv_tiles: int) -> List[int]: + """The splits worth scoring on this device, ascending, always starting at 1. + + THE single split-KV list -- what a row can BUILD is a separate boolean + (``Capabilities.split_kv_supported``), so this is free to be device-derived + rather than a hand-maintained per-row literal. + + Powers of two from 1 up to ``2**ceil(log2(sm_count))``: you never need more + CTA-tiles than the machine has SMs, so that is where the occupancy argument + for splitting runs out. Rounding UP rather than down offers the first + over-subscribing point and lets the cost model reject it on the wave term, + instead of the bound pre-judging it. Powers of two because ``split_kv`` is a + TemplateParams field and so a kernel-module cache key -- an unrestricted + choice mints a compiled specialization per shape. + + Also bounded by ``kv_tiles // _SPLIT_KV_MIN_TILES``. The chunking hands the + remainder to the LEADING splits, so the thinnest gets ``floor(kv_tiles/s)`` + tiles, and ``floor(kv_tiles/s) >= m`` is exactly ``s <= floor(kv_tiles/m)``. + On a short KV that bound binds first. + + No workspace bound here: the partial slabs grow with s, but so does the + combine term in :func:`choose_split_kv`, and it grows with the Q rows -- + which is what makes the slabs big in the first place. The model self-limits; + a caller needing a hard ceiling has ``deselect_workspace_greater_than``. + """ + if sm_count <= 0 or kv_tiles <= 0: + return [1] + hi = 1 << max(0, (sm_count - 1).bit_length()) # 2**ceil(log2(sm_count)) + hi = min(hi, max(1, kv_tiles // _SPLIT_KV_MIN_TILES)) + out, s = [], 1 + while s <= hi: + out.append(s) + s <<= 1 + return out def choose_split_kv( @@ -100,8 +141,9 @@ def choose_split_kv( batch: int, kv_tiles: int, sm_count: int, + combine_rows: int, ctas_per_tile: int = 1, - max_split: int = _SPLIT_KV_MAX, + candidates: Optional[List[int]] = None, ) -> int: """How many KV chunks to cut each Q tile into; 1 = do not split. @@ -111,48 +153,63 @@ def choose_split_kv( 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. + Splitting runs TWO kernels, so the model is two LATENCIES summed -- each one + (sequential rounds) x (what one round costs). Both terms must be latency: + mixing in an aggregate-work term would double-count the parallelism the wave + factor has already divided out. + + waves(s) = ceil(base_ctas * s / sm_count) # main grid + combine_waves = ceil(combine_rows / sm_count) # combine grid, NO s + cost(s) = waves(s) * (ceil(kv_tiles / s) + CTA_COST) + + combine_waves * (s * COMBINE_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. COMBINE_COST is + outside it: the combine is a separate launch whose grid is ``(S_q, H, B)`` + (split_combine_sm100), one block per output row and independent of ``s`` -- + only the per-block work grows with ``s``, since each block reduces ``s`` + partials. Hence ``combine_rows`` (= S_q * H_q * B) and not ``base_ctas``. + + Why the combine term matters: ``s`` reaches the first term ONLY through + ``waves(s)``, a step function. Between wave boundaries a larger split is + free there while the loop term keeps falling, so without a second term the + model always takes the largest split that fits the current wave. That is + harmless while the candidate list stops at 4 and a runaway once it does not. 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. + and never splits; and a long-S_q chunk splits less than a short one, because + its combine has more rows to reduce. 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``. + ``candidates`` defaults to :func:`split_kv_candidates` for the device. """ 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): + if candidates is None: + candidates = split_kv_candidates(sm_count=sm_count, kv_tiles=kv_tiles) + # The combine reads every partial of every output row, so its grid is sized + # by the rows; max(1, ...) because a decode-shaped launch has fewer rows + # than SMs and still pays one wave. + combine_waves = max(1, _ceil_div(max(0, combine_rows), sm_count)) + + best_split, best_cost = 1, None + for split in candidates: + if split < 1 or split > kv_tiles: + continue # 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 + if split > 1 and kv_tiles // split < _SPLIT_KV_MIN_TILES: + continue waves = _ceil_div(base_ctas * split, sm_count) - cost = waves * (_ceil_div(kv_tiles, split) + _SPLIT_KV_CTA_COST) - if cost < best_cost: + cost = waves * (_ceil_div(kv_tiles, split) + _SPLIT_KV_CTA_COST) + combine_waves * split * _SPLIT_KV_COMBINE_COST + if best_cost is None or cost < best_cost: best_split, best_cost = split, cost - split <<= 1 return best_split @@ -310,21 +367,23 @@ def _split_points(caps: Capabilities, facts, tile_m: Optional[int], tile_n: Opti — the packed grid is smaller, which is exactly when splitting pays. 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. + EXACT launch geometry via :func:`_pack_gqa_tile_q` — the Q rows one grid + tile covers, which on SM100 is the cluster's ``TILES_Q*TILE_M*CTA_MMA`` + (512 at d128/d192), not ``tile_m*cga`` (256). The distinction is the whole + model: fed 256 the chooser sees twice the tiles the launch actually has, + so it reads a half-empty machine as full and under-splits or declines to + split at all. 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). + + A split the model asks for LEADS, with no-split behind it as the runner-up + — so a plain ``build_plans()`` runs the split, and autotune / select_plan + can still reach the unsplit plan. Emitting it the other way round meant the + default build never used the split the model had just computed. """ - domain = caps.split_kvs - if len(domain) <= 1: - return [_sole(domain)] - no_split = 1 if 1 in domain else min(domain) + no_split = 1 + if not caps.split_kv_supported: + return [no_split] 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): @@ -338,21 +397,23 @@ def _split_points(caps: Capabilities, facts, tile_m: Optional[int], tile_n: Opti 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) + rows_per_tile = _pack_gqa_tile_q(caps, facts, tile_m) split = choose_split_kv( q_tiles=_ceil_div(facts.s_q * pack_g, rows_per_tile), heads_q=facts.h_q // pack_g, batch=facts.b, kv_tiles=_ceil_div(facts.s_kv, tile_n or 128), sm_count=sm_count, + # The combine's grid is (S_q, H, B) — the REAL head count, not the + # packed one: packing folds heads into Q rows for the main kernel, but + # the combine still reduces one block per (row, head, batch) of the + # graph's own output. + combine_rows=facts.s_q * facts.h_q * facts.b, 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: + if split <= 1: return [no_split] - return [no_split, usable[-1]] + return [split, no_split] def _softmax_points(caps: Capabilities) -> List[Optional[int]]: @@ -410,15 +471,25 @@ def _knob_sets(spec: EngineSpec, facts) -> List[SdpaFwdKnobs]: # The split model sees the launch geometry of the set it rides — the # packed grid when the baseline packs. splits = _split_points(caps, facts, base_tile[0], base_tile[1], cga, pack_g=(facts.h_q // facts.h_kv) if packed_first else 1) - base = SdpaFwdKnobs( - sched_policy=scheds[0], - tile_m=base_tile[0], - tile_n=base_tile[1], - cga=cga, - pack_gqa=True if packed_first else unpacked_pack, - split_kv=splits[0], - softmax_precision=_softmax_points(caps)[0], - ) + # A split set rides the plain scheduler: the SM120 config bars a split under + # the LPT remaps, and in the underfilled regime a split targets, LPT + # balancing is moot — the split itself levels the grid. The coupling is + # structural, so it binds whichever leg leads; it cannot live only on the + # runner-up loop or a leading split would inherit the derived LPT policy. + plain_sched = SCHED_NATURAL if SCHED_NATURAL in caps.sched_policies else scheds[0] + + def _leg(split: Optional[int]) -> SdpaFwdKnobs: + return SdpaFwdKnobs( + sched_policy=plain_sched if (split or 1) > 1 else scheds[0], + tile_m=base_tile[0], + tile_n=base_tile[1], + cga=cga, + pack_gqa=True if packed_first else unpacked_pack, + split_kv=split, + softmax_precision=_softmax_points(caps)[0], + ) + + base = _leg(splits[0]) out = [base] for tile_m, tile_n in tiles[1:]: # A packed baseline's tile runners keep the packing, so tiles the @@ -427,8 +498,11 @@ def _knob_sets(spec: EngineSpec, facts) -> List[SdpaFwdKnobs]: if base.pack_gqa is True and True not in _pack_gqa_points(caps, facts, tile_m or 128): continue out.append(replace(base, tile_m=tile_m, tile_n=tile_n)) + # Scheduler runners ride an UNSPLIT leg: a split set is pinned to the plain + # scheduler above, so an LPT runner is only a candidate without one. + sched_host = base if (base.split_kv or 1) == 1 else _leg(splits[-1]) for policy in scheds[1:]: - out.append(replace(base, sched_policy=policy)) + out.append(replace(sched_host, sched_policy=policy)) # The opposite pack_gqa leg, riding its own tile (packed: the largest # admitting tile; unpacked: the tile rule's best). if pack_tile is not None: @@ -437,10 +511,7 @@ def _knob_sets(spec: EngineSpec, facts) -> List[SdpaFwdKnobs]: else: out.append(replace(base, pack_gqa=True, tile_m=pack_tile[0], tile_n=pack_tile[1])) 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)) + out.append(_leg(split)) seen, unique = set(), [] for knobs in out: if knobs not in seen: @@ -462,7 +533,7 @@ def _fallback_knobs(caps: Capabilities) -> SdpaFwdKnobs: tile_n=min(caps.tile_ns, default=None), cga=_sole(caps.cgas), pack_gqa=False if False in caps.pack_gqas else _sole(caps.pack_gqas), - split_kv=1 if 1 in caps.split_kvs else _sole(caps.split_kvs), + split_kv=1, # the fallback never splits: least-demanding means one kernel, no partial workspace softmax_precision=_sole(caps.softmax_precisions), ) diff --git a/test/python/sdpa/frost/test_sdpa_fp8_sm107.py b/test/python/sdpa/frost/test_sdpa_fp8_sm107.py index 1121a93dc..1d3f87535 100644 --- a/test/python/sdpa/frost/test_sdpa_fp8_sm107.py +++ b/test/python/sdpa/frost/test_sdpa_fp8_sm107.py @@ -82,8 +82,8 @@ def test_per_tensor_fp8_rows_split_per_arch_line(): assert sm107.softmax_precisions == frozenset({_c.data_type.FLOAT, _c.data_type.HALF}) # The SM107 sibling has no split path and no LPT remap yet (issue #653). - assert sm107.split_kvs == frozenset({1}) - assert sm100.split_kvs == frozenset({1, 2, 4}) + assert sm107.split_kv_supported is False + assert sm100.split_kv_supported is True assert sm107.sched_policies == frozenset({SCHED_NATURAL}) assert sm100.sched_policies == frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py b/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py index 7f34147ed..76ee00f75 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py @@ -63,11 +63,13 @@ def test_recommend_emits_multiple_complete_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 + # Behavior preservation on the UNSPLIT leg: 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). A + # grid that fills the machine never splits, so it reads the derivation + # straight off the primary. + causal = recommend("A", _facts(s_q=8192), _OFFERED) + assert causal[0].knobs.split_kv == 1 and causal[0].knobs.sched_policy == 2 # SCHED_LPT_L2 dense = recommend("A", _facts(causal=False), _OFFERED) dense_f16 = [p for p in dense if p.engine_id == 20500] assert dense_f16[0].knobs.sched_policy == 0 # SCHED_NATURAL @@ -75,10 +77,29 @@ def test_recommend_primary_reproduces_the_derived_scheduler(): @pytest.mark.L0 -def test_recommend_split_is_a_runner_up_and_respects_structure(): +def test_split_and_scheduler_stay_coupled_whichever_leads(): + """A split set rides the plain scheduler — structural, so it must bind the + PRIMARY too, not just the runner-ups. config_sm120 raises outright on + split_kv > 1 under an LPT remap, so an LPT+split set is unbuildable there. + Regression: flipping the split to lead once let it inherit the derived + LPT_L2 policy on causal graphs.""" + for f in (_facts(), _facts(causal=False), _facts(s_q=8192), _facts(h_q=1, h_kv=1)): + for p in recommend("A", f, _OFFERED): + if (p.knobs.split_kv or 1) > 1: + assert p.knobs.sched_policy == 0, f"split set on a non-plain scheduler: {p.knobs}" + + +@pytest.mark.L0 +def test_recommend_split_leads_and_respects_structure(): + # A split the wave-cost model asks for is what a plain build_plans() runs, + # with no-split behind it for autotune / select_plan. Sweep justifying the + # lead (B300, ar_dit chunked prefill, bf16 B1xH9xD128, S_kv=62208, no mask): + # S_q=985 0.955 ms -> 0.556 ms (split 4, 1.72x) + # S_q=2048 0.960 ms -> 0.722 ms (split 2, 1.33x) + # S_q>=4096 unchanged (model declines to split a full grid) 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" + assert plans[0].knobs.split_kv > 1, "an underfilled grid runs the split the model chose" + assert any(p.knobs.split_kv == 1 for p in plans), "no-split must stay reachable as the runner-up" 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}" 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 be60a88bc..14f1cc71a 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 @@ -856,6 +856,8 @@ def _expected_split(b, h_q, s_q, s_kv, *, rows_per_tile=512, ctas_per_tile=2, kv kv_tiles=-(-s_kv // kv_tile), sm_count=device_info(torch.cuda.current_device()).sm_count, ctas_per_tile=ctas_per_tile, + # The combine's grid is (S_q, H, B) — see choose_split_kv. + combine_rows=b * h_q * s_q, ) diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index 58e0a2f5e..07fc5ca84 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -1037,11 +1037,15 @@ def test_sm120_knob_domains(monkeypatch): # 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. + # split_kv: the SM120 row WIRES the split path (inline chunking + the + # shared combine), which is a boolean gate — the kernel has no upper bound + # on the split count, so 8 is admissible too. WHICH splits get proposed is + # split_kv_candidates' device-derived ladder, not a per-row domain. A + # non-count 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)) + assert _SM120 in _eligible(g, engines.SdpaFwdKnobs(split_kv=8)) + assert not _eligible(g, engines.SdpaFwdKnobs(split_kv=0)) # --------------------------------------------------------------------------- diff --git a/test/python/sdpa/frost/test_split_kv_heuristic.py b/test/python/sdpa/frost/test_split_kv_heuristic.py index 5a5304878..cc0a2fd3d 100644 --- a/test/python/sdpa/frost/test_split_kv_heuristic.py +++ b/test/python/sdpa/frost/test_split_kv_heuristic.py @@ -13,7 +13,7 @@ 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 +from cudnn.sdpa.fwd.heuristics import _SPLIT_KV_MIN_TILES, choose_split_kv, split_kv_candidates # Pure arithmetic — no device, no kernel build — so every case is L0. pytestmark = pytest.mark.L0 @@ -22,7 +22,12 @@ def _d128_cga2(s_q, s_kv, heads, batch, sm_count=B200_SMS, **kw): - """choose_split_kv for the d128 cga2 geometry.""" + """choose_split_kv for the d128 cga2 geometry. + + ``combine_rows`` is the combine kernel's grid (S_q, H, B) — derived here so + every case exercises the real two-kernel model rather than a configuration + production never builds.""" + kw.setdefault("combine_rows", s_q * heads * batch) return choose_split_kv( q_tiles=-(-s_q // 512), heads_q=heads, @@ -34,6 +39,10 @@ def _d128_cga2(s_q, s_kv, heads, batch, sm_count=B200_SMS, **kw): ) +def _ladder(sm_count=B200_SMS, kv_tiles=1 << 20): + return split_kv_candidates(sm_count=sm_count, kv_tiles=kv_tiles) + + # --- the case the feature exists for -------------------------------------- @@ -89,7 +98,7 @@ def test_nearly_full_chip_does_not_split(): 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 + assert choose_split_kv(q_tiles=1, heads_q=B200_SMS, batch=1, kv_tiles=256, sm_count=B200_SMS, combine_rows=128 * B200_SMS) == 1 def test_single_kv_tile_cannot_split(): @@ -99,7 +108,7 @@ def test_single_kv_tile_cannot_split(): 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) + split = choose_split_kv(q_tiles=1, heads_q=1, batch=1, kv_tiles=kv_tiles, sm_count=B200_SMS, combine_rows=128) assert split <= kv_tiles // _SPLIT_KV_MIN_TILES @@ -112,7 +121,7 @@ def test_degenerate_inputs_do_not_split(): {"sm_count": 0}, {"sm_count": -1}, ): - args = {"q_tiles": 1, "heads_q": 1, "batch": 1, "kv_tiles": 256, "sm_count": B200_SMS} + args = {"q_tiles": 1, "heads_q": 1, "batch": 1, "kv_tiles": 256, "sm_count": B200_SMS, "combine_rows": 128} args.update(kw) assert choose_split_kv(**args) == 1, f"{kw} must fall back to no split" @@ -125,7 +134,7 @@ def test_degenerate_inputs_do_not_split(): 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 1 <= split <= max(_ladder(kv_tiles=kv_tiles)) assert split <= kv_tiles, "more splits than KV tiles would leave empty splits" if split > 1: assert -(-kv_tiles // split) >= _SPLIT_KV_MIN_TILES @@ -143,17 +152,17 @@ def test_longer_kv_never_needs_fewer_splits(): @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_kv is a compile-cache key, so the search is a power-of-two ladder.""" 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 + assert split in _ladder(kv_tiles=-(-s_kv // 128)) 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) + split = choose_split_kv(q_tiles=1, heads_q=80, batch=1, kv_tiles=512, sm_count=148, ctas_per_tile=2, combine_rows=512 * 80) assert split > 1 assert 160 * split > 148, "this shape is exactly the case that wants over-subscription" @@ -164,19 +173,33 @@ def test_exactly_balanced_launch_never_splits(): 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) + 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, combine_rows=128 * k * sm_count + ) assert split == 1, f"base={k * sm_count} == {k}x{sm_count} SMs: nothing to smooth, got {split}" +# The Q extent the fit below is pinned at: one full cga2 cluster tile +# (TILES_Q*TILE_M*CTA_MMA = 512 rows), so q_tiles == 1 and the combine grid is +# 512 * heads rows. +_FIT_S_Q = 512 + # (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)] +# d128, 512 KV tiles, S_q=512, bf16, mask-free), re-measured for the two-kernel +# cost model. A change that moves any of these is a policy change and needs its +# own measurement. +# +# The model matches the measured optimum on 8 of 12; mean regret 1.009, worst +# 1.035. The misses (88, 100, 150, 160) all sit within ~10% of a FULL machine, +# where the measured curve is nearly flat -- e.g. base=88 spans 2.264..2.512 ms +# across every split -- so the ranking there is worth little and the model +# prefers the cheap answer. Regret, not exact agreement, is the bar. +_B300_FIT = [(8, 16), (16, 8), (32, 4), (64, 2), (88, 1), (100, 1), (120, 1), (128, 1), (150, 2), (160, 2), (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) + got = choose_split_kv(q_tiles=1, heads_q=base_ctas // 2, batch=1, kv_tiles=512, sm_count=148, ctas_per_tile=2, combine_rows=_FIT_S_Q * (base_ctas // 2)) assert got == expected @@ -184,12 +207,15 @@ 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 + assert choose_split_kv(q_tiles=1, heads_q=44, batch=1, kv_tiles=512, sm_count=148, ctas_per_tile=2, combine_rows=_FIT_S_Q * 44) != 2 + assert choose_split_kv(q_tiles=1, heads_q=50, batch=1, kv_tiles=512, sm_count=148, ctas_per_tile=2, combine_rows=_FIT_S_Q * 50) != 2 -def test_max_split_is_respected(): - assert _d128_cga2(128, 1 << 20, 1, 1, max_split=4) <= 4 +def test_candidates_bound_the_choice(): + """The chooser never returns a split outside the list it was given.""" + for cand in ([1], [1, 2], [1, 2, 4], [1, 2, 4, 8, 16]): + got = _d128_cga2(128, 1 << 20, 1, 1, candidates=cand) + assert got in cand # --- the knob plumbing ------------------------------------------------------ @@ -201,7 +227,7 @@ def test_split_request_outside_the_domain_makes_the_engine_ineligible(requested) whose lowering has no split path is honored-or-ineligible, never silently dropped.""" caps = Capabilities(sm_lo=100, sm_hi=100, phase="prefill", d_shapes=frozenset({(128, 128)})) - assert caps.split_kvs == frozenset({1}) + assert caps.split_kv_supported is False why = mismatch(caps, _facts(), SdpaFwdKnobs(split_kv=requested)) assert why is not None and "split_kv" in why @@ -234,7 +260,7 @@ def test_split_declines_when_the_kv_tail_needs_synthesized_padding(): phase="prefill", d_shapes=frozenset({(128, 128)}), skv_tail_via_padding=True, - split_kvs=frozenset({1, 2, 4}), + split_kv_supported=True, ) ragged = ga.SdpaGraphFacts(s_q=128, s_kv=1000) # 1000 % 128 != 0, mask-free why = mismatch(caps, ragged, SdpaFwdKnobs(split_kv=2)) @@ -247,15 +273,106 @@ def test_split_declines_when_the_kv_tail_needs_synthesized_padding(): def test_split_domains_match_the_wired_lowerings(): - """Guards the pairing: a row advertises split_kvs > {1} exactly when its + """Guards the pairing: a row sets split_kv_supported exactly when its adapter forwards the knob into TemplateParams and launches the combine. - Widening one without the plumbing reintroduces the silently-dropped knob.""" + Setting 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})} + advertising = {sp.name for sp in ENGINE_SPECS if sp.capabilities.split_kv_supported} assert advertising == { "sdpa_fwd_prefill_sm100", "sdpa_fwd_prefill_sm100_mxfp8", "sdpa_fwd_prefill_sm100_fp8", "sdpa_fwd_prefill_sm120", }, f"split domains drifted from the wired lowerings: {sorted(advertising)}" + + +def test_split_points_feeds_the_exact_cluster_extent(): + """_split_points must measure the launch in CLUSTERS, not CTA tiles. + + The SM100 d128 cluster covers TILES_Q * TILE_M * CTA_MMA = 512 Q rows on + its 2 CTAs — the same extent every helper above assumes. Feeding the model + ``tile_m * cga`` (256) instead doubles the apparent tile count, so a + half-empty machine reads as full and the chooser under-splits: the ar_dit + chunked-prefill shape below measured 4 on the true geometry and 2 on the + approximation, worth 1.72x vs 1.33x on B300. + """ + import cudnn + from cudnn.sdpa import graph_analyzer as ga + from cudnn.sdpa.fwd.engines import ENGINE_SPECS + from cudnn.sdpa.fwd.heuristics import _split_points + + caps = next(sp for sp in ENGINE_SPECS if sp.name == "sdpa_fwd_prefill_sm100").capabilities + # ar_dit: B1 x H9 x D128 bf16, 985 new tokens against a 62208-token clip. + facts = ga.SdpaGraphFacts( + b=1, + h_q=9, + h_kv=9, + s_q=985, + s_kv=62208, + d_qk=128, + d_v=128, + dtype=cudnn.data_type.BFLOAT16, + dtype_o=cudnn.data_type.BFLOAT16, + device_sm_count=B200_SMS, + device_cc=(10, 0), + ) + points = _split_points(caps, facts, 128, 128, 2) + assert points[0] == 4, f"expected the 512-row cluster extent to choose 4, got {points}" + assert points[-1] == 1, "no-split must remain reachable behind the chosen split" + + +# --- the candidate ladder --------------------------------------------------- + + +@pytest.mark.parametrize("sm_count,top", [(148, 256), (132, 256), (108, 128), (84, 128), (16, 16), (1, 1)]) +def test_ladder_is_derived_from_the_device(sm_count, top): + """THE single split list, derived per device rather than declared per row: + powers of two up to 2**ceil(log2(sm_count)) — you never need more CTA-tiles + than the machine has SMs.""" + got = split_kv_candidates(sm_count=sm_count, kv_tiles=1 << 20) + assert got[0] == 1 and got[-1] == top + assert got == [1 << i for i in range(len(got))] + + +def test_ladder_is_bounded_by_the_thinnest_split(): + """kv_tiles // MIN_TILES is the largest split whose thinnest chunk still + clears the floor — the loop guard restated as a bound.""" + for kv_tiles in (4, 17, 64, 486, 512): + got = split_kv_candidates(sm_count=B200_SMS, kv_tiles=kv_tiles) + assert max(got) <= max(1, kv_tiles // _SPLIT_KV_MIN_TILES) + for s in got: + if s > 1: + assert kv_tiles // s >= _SPLIT_KV_MIN_TILES + + +def test_degenerate_device_yields_the_no_split_ladder(): + assert split_kv_candidates(sm_count=0, kv_tiles=512) == [1] + assert split_kv_candidates(sm_count=148, kv_tiles=0) == [1] + + +# --- the combine term ------------------------------------------------------- + + +@pytest.mark.parametrize("heads", [1, 2, 4, 8, 16]) +def test_longer_q_never_splits_more(heads): + """The combine's grid is (S_q, H, B), so more Q rows means more combine + waves to pay per split. At a FIXED base_ctas (every S_q here is one cga2 + cluster tile, so q_tiles == 1) a longer Q extent must never ask for a + LARGER split. Only the combine term can express this — the wave term does + not see S_q at all.""" + got = [ + choose_split_kv(q_tiles=1, heads_q=heads, batch=1, kv_tiles=512, sm_count=B200_SMS, ctas_per_tile=2, combine_rows=s_q * heads) + for s_q in (64, 128, 256, 512) + ] + assert all(a >= b for a, b in zip(got, got[1:])), f"S_q 64/128/256/512 -> {got}" + + +def test_decode_rows_barely_pay_for_the_combine(): + """A decode-shaped launch reduces one block per (row, head, batch) — far + fewer rows than SMs, so one combine wave. The reduction must not price it + out of the splitting it exists for.""" + decode = choose_split_kv(q_tiles=1, heads_q=8, batch=1, kv_tiles=512, sm_count=B200_SMS, ctas_per_tile=2, combine_rows=8) + prefill = choose_split_kv(q_tiles=1, heads_q=8, batch=1, kv_tiles=512, sm_count=B200_SMS, ctas_per_tile=2, combine_rows=8 * 4096) + assert decode > 1 + assert decode >= prefill