From 8e1801895fbdbf53de318879790c6c61b9266027 Mon Sep 17 00:00:00 2001 From: yanzhuoc Date: Tue, 18 Aug 2026 18:04:30 -0700 Subject: [PATCH] sdpa fwd: KV split for the SM100/SM120 prefill kernels Add an optional split over the KV sequence: with split_kv > 1 each Q tile's KV range is cut into contiguous chunks, every chunk runs as its own persistent tile, and the per-chunk (O, LSE) partials are reduced by kernels/split_combine_sm100.py. At split_kv == 1 the added closures fold away and the traced code is unchanged. Flavors: sm100 d128 (f16/bf16, fp8, mxfp8), d192/128, d256 and d512, plus sm120 f16/bf16. The knob is gated per flavor so a flavor that does not thread it cannot silently accept it, and the config backstop rejects THD, attention sink and the flattened scheduler grids. Also adds an optional cga1 cluster width (cta_mma=1) for d128 and d192/128. cga1 has no collective MMA to halve per-CTA K/V, so d128 recovers the extra SMEM through the existing Q/O alias and the fp8 family scales the KV stage depth with the cluster width. For the FP8 family the combine also reports amax_o over the recombined O: a per-split epilogue sees only its own partial, and O is a convex combination of those, so a max over partials over-reports the output amax. Measured 3-5x end to end at S_q=128 over a 32K KV run on B200, with numerics unchanged against both the unsplit kernel and an fp32 reference. Tests: test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py -- 56 cases over even and uneven splits (including empty ones), dense / causal / SWA / bottom-right / padded masks, GQA, bf16, fp8 and mxfp8 at both cluster widths, the recombined LSE, and amax_o. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/sdpa/fwd/config_sm100.py | 161 +++- python/cudnn/sdpa/fwd/config_sm120.py | 21 + .../cudnn/sdpa/fwd/kernels/_common_sm100.py | 198 +++++ .../fwd/kernels/prefill_d128_f16_sm100.py | 167 +++- .../fwd/kernels/prefill_d128_fp8_sm100.py | 149 +++- .../fwd/kernels/prefill_d128_mxfp8_sm100.py | 151 +++- .../kernels/prefill_d192_d128_f16_sm100.py | 138 ++- .../fwd/kernels/prefill_d256_f16_sm100.py | 129 ++- .../fwd/kernels/prefill_d512_f16_sm100.py | 128 ++- .../sdpa/fwd/kernels/prefill_f16_sm120.py | 69 +- .../sdpa/fwd/kernels/split_combine_sm100.py | 199 +++++ .../frost/test_sdpa_fwd_split_kv_sm100.py | 814 ++++++++++++++++++ 12 files changed, 2079 insertions(+), 245 deletions(-) create mode 100644 python/cudnn/sdpa/fwd/kernels/split_combine_sm100.py create mode 100644 test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py diff --git a/python/cudnn/sdpa/fwd/config_sm100.py b/python/cudnn/sdpa/fwd/config_sm100.py index f69331962..a01a4fa28 100644 --- a/python/cudnn/sdpa/fwd/config_sm100.py +++ b/python/cudnn/sdpa/fwd/config_sm100.py @@ -88,12 +88,39 @@ class TemplateParams: seq_q_lens_present: bool = False sched_policy: int = SCHED_NATURAL thd_varlen: bool = False + # KV split: each Q tile's KV loop range is cut into ``split_kv`` contiguous + # chunks, each run as its own persistent tile writing a partial (O, LSE) + # that kernels/split_combine_sm100.py reduces. 1 = off (byte-identical + # codegen to the single-pass kernel). + split_kv: int = 1 + # MMA cluster width: 2 = cga2 collective tcgen05.mma.cta_group::2 (a CTA + # pair share one MMA, each holding half of every K/V tile); 1 = cga1, one + # independent CTA per tile. + # + # cga1 has no collective MMA to halve per-CTA K/V, so at a fixed STAGES_KV + # it doubles that footprint. d128 buys the 64 KiB back by aliasing Q and O + # into one slab (make_cfg_d128 turns QO_ALIAS on for cga1; see + # _validate_cfg_d128's SMEM check); the fp8 family instead scales the stage + # count with the width so stages x per-CTA-buffer stays constant. + cta_mma: int = 2 # cc10.3+ fuses the S_acc row-max into the LDTM (tcgen05.ld.red.f32.max); cc10.0 # lacks it and uses the manual load + software reduction. Auto-set from the device # capability at compile time (MXFP8 only; the f16/fp8 kernels do not read it). fused_ldtm_stat: bool = False +# split_kv / cta_mma live on the TemplateParams shared by every SM100 flavor, but +# a flavor only honours them once its make_cfg_* threads them into a Cfg AND its +# kernel reads them. Accepting them elsewhere would silently ignore them — and +# for split_kv that is not merely surprising but WRONG: the caller sizes an +# (S*B)-batch partial workspace and runs the combine, while the kernel keeps +# writing only slots [0, B). The untouched slots keep lse_partial = 0 rather +# than -inf, so they carry weight exp(0 - M) != 0 through the log-sum-exp and +# corrupt the result instead of dropping out. Grow these sets as flavors land. +_SPLIT_KV_FLAVORS = frozenset({"d128", "d192", "d256", "d512"}) +_CTA_MMA_FLAVORS = frozenset({"d128", "d192"}) + + def _validate_params(flavor: str, k: TemplateParams) -> None: if k.dtype_qkv not in (DTYPE_E4M3, DTYPE_E5M2, DTYPE_BF16, DTYPE_FP16): raise ValueError(f"{flavor}: DTYPE_QKV must be E4M3/E5M2/BF16/FP16 (0..3); got {k.dtype_qkv}") @@ -121,6 +148,31 @@ def _validate_params(flavor: str, k: TemplateParams) -> None: raise ValueError(f"{flavor}: SEQ_Q_LENS_PRESENT requires SEQ_KV_LENS_PRESENT (padding mask)") if k.sched_policy not in (SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2): raise ValueError(f"{flavor}: only SCHED_NATURAL (0) / SCHED_LPT (1) / SCHED_LPT_L2 (2) are wired up; got {k.sched_policy}") + if k.cta_mma not in (1, 2): + raise ValueError(f"{flavor}: cta_mma must be 1 (cga1) or 2 (cga2); got {k.cta_mma}") + # split_kv / cta_mma live on the TemplateParams shared by every SM100 flavor, + # but only make_cfg_d128 threads them into a Cfg and only the d128 kernel + # reads them. Accepting them elsewhere would silently ignore them — and for + # split_kv that is not merely surprising but WRONG: the caller sizes an + # (S*B)-batch partial workspace and runs the combine, while the kernel keeps + # writing only slots [0, B). The untouched slots keep lse_partial = 0 rather + # than -inf, so they carry weight exp(0 - M) != 0 through the log-sum-exp and + # corrupt the result instead of dropping out. Reject at the door. + if k.split_kv != 1 and flavor not in _SPLIT_KV_FLAVORS: + raise ValueError(f"{flavor}: split_kv is not implemented on this flavor (got {k.split_kv}); supported: {sorted(_SPLIT_KV_FLAVORS)}") + if k.cta_mma != 2 and flavor not in _CTA_MMA_FLAVORS: + raise ValueError(f"{flavor}: cta_mma is not selectable on this flavor (got {k.cta_mma}); supported: {sorted(_CTA_MMA_FLAVORS)}") + if k.split_kv < 1: + raise ValueError(f"{flavor}: split_kv must be >= 1 (1 = KV-split off); got {k.split_kv}") + if k.split_kv > 1: + # Each of these would need extra machinery in the combine pass, so the + # backstop rejects them rather than silently producing a wrong answer. + if k.thd_varlen: + raise ValueError(f"{flavor}: split_kv > 1 is dense-only (THD packs its own flat grid)") + if k.has_sink: + # The sink logit is folded into the softmax denominator in the + # per-tile epilogue, so every split would add its own copy of it. + raise ValueError(f"{flavor}: split_kv > 1 with attention sink is not supported (the sink would be counted once per split)") def _mask_flags_from(params: TemplateParams) -> int: @@ -290,6 +342,9 @@ class CfgD256: THD_VARLEN: int = 0 + # KV split; 1 = off. See TemplateParams.split_kv. + SPLIT_KV: int = 1 + def _validate_cfg_d256(cfg: CfgD256) -> None: """Consistency checks on the (mostly hardcoded) d256 geometry.""" @@ -335,6 +390,7 @@ def make_cfg_d256(params: TemplateParams) -> Tuple[CfgD256, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), + SPLIT_KV=int(params.split_kv), ) _validate_cfg_d256(cfg) return cfg, _tma_iters(cfg) @@ -429,6 +485,9 @@ class CfgD512: THD_VARLEN: int = 0 + # KV split; 1 = off. See TemplateParams.split_kv. + SPLIT_KV: int = 1 + def _validate_cfg_d512(cfg: CfgD512) -> None: """Consistency checks on the (mostly hardcoded) d512 geometry.""" @@ -478,6 +537,7 @@ def make_cfg_d512(params: TemplateParams) -> Tuple[CfgD512, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), + SPLIT_KV=int(params.split_kv), ) _validate_cfg_d512(cfg) return cfg, _tma_iters(cfg) @@ -575,6 +635,34 @@ class CfgD128: THD_VARLEN: int = 0 + # KV split; 1 = off. See TemplateParams.split_kv. + SPLIT_KV: int = 1 + + +# Blackwell SM100 per-CTA dynamic SMEM cap (228 KiB physical, 227 KiB usable). +_SM100_MAX_DYN_SMEM = 227 * 1024 + + +def _d128_smem_bytes(cfg) -> int: + """Data-buffer SMEM for the d128 pipeline (barriers/TMEM ptr are noise). + + Q and O are TILES_Q slabs each; under QO_ALIAS they share one slab sized to + the larger. K/V are STAGES_KV buffers each, and their PER-CTA size is + divided by CTA_MMA because the cga2 collective MMA lets a CTA pair hold half + of every tile. That divisor is exactly what cga1 gives up, which is why + cga1 needs the Q/O alias to break even: + + cga2, no alias : 64(Q) + 64(O) + 32(K) + 32(V) = 192 KiB + cga1, no alias : 64(Q) + 64(O) + 64(K) + 64(V) = 256 KiB (over cap) + cga1, alias : 64(Q u O) + 64(K) + 64(V) = 192 KiB + """ + q_slab = cfg.TILE_M * cfg.TILE_K * cfg.BPE + o_slab = cfg.TILE_M * cfg.TILE_O * cfg.BPE_O + qo = cfg.TILES_Q * (max(q_slab, o_slab) if cfg.QO_ALIAS else q_slab + o_slab) + k = cfg.STAGES_KV * (cfg.TILE_N * cfg.TILE_K * cfg.BPE // cfg.CTA_MMA) + v = cfg.STAGES_KV * (cfg.TILE_O * cfg.TILE_N * cfg.BPE // cfg.CTA_MMA) + return qo + k + v + def _validate_cfg_d128(cfg: CfgD128) -> None: """Consistency checks on the (mostly hardcoded) d128 (llama) geometry.""" @@ -583,14 +671,27 @@ def _validate_cfg_d128(cfg: CfgD128) -> None: (cfg.MMA_REGS == cfg.TMALDG_REGS == cfg.TMASTG_REGS == cfg.SCHEDULER_REGS, "d128: MMA/TMALDG/TMASTG/SCHEDULER regs must match"), (cfg.MMA_REGS + cfg.CORRECTION_REGS + cfg.SOFTMAX_WARPGROUPS * cfg.SOFTMAX_REGS <= 512, "d128: register budget over 512"), (cfg.MMA_REGS % 8 == 0 and cfg.CORRECTION_REGS % 8 == 0 and cfg.SOFTMAX_REGS % 8 == 0, "d128: per-role regs must be multiples of 8"), - (cfg.CGA_M == 2 and cfg.CTA_MMA == 2, "d128 SM100 is cga2-only (CGA_M == CTA_MMA == 2)"), + (cfg.CGA_M == cfg.CTA_MMA and cfg.CTA_MMA in (1, 2), "d128 SM100: CGA_M must equal CTA_MMA, and CTA_MMA must be 1 (cga1) or 2 (cga2)"), + ( + cfg.QO_ALIAS == 1 if cfg.CTA_MMA == 1 else True, + "d128 cga1: QO_ALIAS is mandatory — cga1 doubles per-CTA K/V (no collective MMA to halve it), " + "so Q and O must share one slab to stay inside the SMEM cap", + ), + ( + _d128_smem_bytes(cfg) <= _SM100_MAX_DYN_SMEM, + f"d128: SMEM {_d128_smem_bytes(cfg) // 1024} KiB over the SM100 {_SM100_MAX_DYN_SMEM // 1024} KiB per-CTA cap", + ), (cfg.TILE_K == 128 and cfg.TILE_O == 128, "d128: d_qk = d_v = 128"), (cfg.TILES_Q == 2, "d128 (llama): TILES_Q must be 2"), (cfg.SOFTMAX_WARPGROUPS == 2, "d128 (llama): SOFTMAX_WARPGROUPS must be 2"), (cfg.CORRECTION_WARPS == 4, "d128 (llama): CORRECTION_WARPS must be 4"), (cfg.TOTAL_WARPS == 16 and cfg.THREADS_PER_CTA == 512, "d128 (llama): 16 warps / 512 threads"), (cfg.READ_TILE_ARRIVERS == 15, f"d128 llama: expected READ_TILE_ARRIVERS=15, got {cfg.READ_TILE_ARRIVERS}"), - (cfg.STAGES_KV == (4 if _fp8 else 2), "d128 SM100: STAGES_KV must be 4 (fp8/mxfp8) / 2 (f16, 192 KiB SMEM budget)"), + ( + cfg.STAGES_KV == ((2 if cfg.CTA_MMA == 1 else 4) if _fp8 else 2), + "d128 SM100: STAGES_KV must be 2 (f16/bf16) or, for fp8/mxfp8, 4 at cga2 and 2 at cga1 — " + "the stage depth scales with the cluster width so stages x per-CTA-buffer stays constant", + ), ( cfg.TILE_K_HW_BMM1 == (32 if _fp8 else 16) and cfg.TILE_K_HW_BMM2 == (32 if _fp8 else 16), "d128: TILE_K_HW must be 32 (fp8/mxfp8 K=32 QMMA) / 16 (f16, 1-chunk on SM10x)", @@ -621,14 +722,27 @@ def make_cfg_d128(params: TemplateParams) -> Tuple[CfgD128, TmaIters]: DTYPE_O=dtype_o, BPE=b, BPE_O=b_o, + CGA_M=params.cta_mma, + CTA_MMA=params.cta_mma, + # cga1 has no collective MMA to halve per-CTA K/V, so Q and O must share + # one slab to stay under the SMEM cap (_validate_cfg_d128 enforces it). + QO_ALIAS=1 if params.cta_mma == 1 else 0, Q_SWZ_BYTES=q_swz_bytes(128, b), K_SWZ_BYTES=q_swz_bytes(128, b), - V_SWZ_BYTES=v_swz_bytes(128, 2, b), + V_SWZ_BYTES=v_swz_bytes(128, params.cta_mma, b), O_SWZ_BYTES=o_swz_bytes(128, b_o), RESCALE_THRESHOLD=rescale_threshold(params.dtype_qkv), TILE_K_HW_BMM1=tile_k_hw_fp8, TILE_K_HW_BMM2=tile_k_hw_fp8, - STAGES_KV=4 if fp8 else 2, + # KV stage depth scales with the cluster width, as in cuDNN's own + # kernels (stages_kv = N * CTA_MMA): cga1 has no collective MMA to halve + # per-CTA K/V, so the stage count halves instead to keep the product -- + # and hence the SMEM -- constant. Only the fp8 family needs this here: + # f16/bf16 already fit at cga1 by aliasing Q and O, and their verified + # cga1 configuration keeps STAGES_KV=2. mxfp8 additionally stages E8M0 + # scale factors that the SMEM model cannot see, and at STAGES_KV=4 that + # pushed a cga1 CTA to 237024 B against the 232448 B cap. + STAGES_KV=(2 if params.cta_mma == 1 else 4) if fp8 else 2, MASK_FLAGS=_mask_flags_from(params), WINDOW_LEFT=params.window_left or 0, WINDOW_RIGHT=params.window_right or 0, @@ -638,6 +752,7 @@ def make_cfg_d128(params: TemplateParams) -> Tuple[CfgD128, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), + SPLIT_KV=int(params.split_kv), ) _validate_cfg_d128(cfg) return cfg, _tma_iters(cfg) @@ -657,6 +772,24 @@ class CfgD192(CfgD128): CORRECTION_REGS: int = 88 +def _d192_smem_bytes(cfg) -> int: + """Data-buffer SMEM for the d192 pipeline (Q/O always aliased). + + Same shape as _d128_smem_bytes, but d_qk = 192 makes the Q and K slabs 1.5x + the d128 ones, which is why this flavor needs a shallower KV pipeline: + + cga2, STAGES_KV=2 : 96(Q u O) + 48(K) + 32(V) = 176 KiB + cga1, STAGES_KV=2 : 96 + 96 + 64 = 256 KiB (over cap) + cga1, STAGES_KV=1 : 96 + 48 + 32 = 176 KiB + """ + q_slab = cfg.TILE_M * cfg.TILE_K * cfg.BPE + o_slab = cfg.TILE_M * cfg.TILE_O * cfg.BPE_O + qo = cfg.TILES_Q * (max(q_slab, o_slab) if cfg.QO_ALIAS else q_slab + o_slab) + k = cfg.STAGES_KV * (cfg.TILE_N * cfg.TILE_K * cfg.BPE // cfg.CTA_MMA) + v = cfg.STAGES_KV * (cfg.TILE_O * cfg.TILE_N * cfg.BPE // cfg.CTA_MMA) + return qo + k + v + + def _validate_cfg_d192(cfg: CfgD192) -> None: """Consistency checks on the native DSv3 d192/d128 geometry.""" checks = ( @@ -665,7 +798,16 @@ def _validate_cfg_d192(cfg: CfgD192) -> None: (cfg.MMA_REGS == cfg.TMALDG_REGS == cfg.TMASTG_REGS == cfg.SCHEDULER_REGS, "d192: MMA/TMALDG/TMASTG/SCHEDULER regs must match"), (cfg.MMA_REGS + cfg.CORRECTION_REGS + cfg.SOFTMAX_WARPGROUPS * cfg.SOFTMAX_REGS <= 512, "d192: register budget over 512"), (cfg.MMA_REGS % 8 == 0 and cfg.CORRECTION_REGS % 8 == 0 and cfg.SOFTMAX_REGS % 8 == 0, "d192: per-role regs must be multiples of 8"), - (cfg.CGA_M == 2 and cfg.CTA_MMA == 2, "d192 SM100 is cga2-only (CGA_M == CTA_MMA == 2)"), + (cfg.CGA_M == cfg.CTA_MMA and cfg.CTA_MMA in (1, 2), "d192 SM100: CGA_M must equal CTA_MMA, and CTA_MMA must be 1 (cga1) or 2 (cga2)"), + ( + cfg.STAGES_KV == (1 if cfg.CTA_MMA == 1 else 2), + "d192: STAGES_KV must scale with the cluster width (2 at cga2, 1 at cga1) — cga1 doubles per-CTA K/V, " + "and cuDNN's own d192 kernel uses stages_kv = 1 * CTA_MMA for exactly this reason", + ), + ( + _d192_smem_bytes(cfg) <= _SM100_MAX_DYN_SMEM, + f"d192: SMEM {_d192_smem_bytes(cfg) // 1024} KiB over the SM100 {_SM100_MAX_DYN_SMEM // 1024} KiB per-CTA cap", + ), (cfg.TILE_K == 192 and cfg.TILE_O == 128, "d192: expected D_QK tile 192 and D_V tile 128"), (cfg.QO_ALIAS == 1, "d192: Q/O SMEM alias is required to stay within SM100 SMEM budget"), (cfg.TILES_Q == 2, "d192: TILES_Q must be 2"), @@ -673,7 +815,6 @@ def _validate_cfg_d192(cfg: CfgD192) -> None: (cfg.CORRECTION_WARPS == 4, "d192: CORRECTION_WARPS must be 4"), (cfg.TOTAL_WARPS == 16 and cfg.THREADS_PER_CTA == 512, "d192: 16 warps / 512 threads"), (cfg.READ_TILE_ARRIVERS == 15, f"d192: expected READ_TILE_ARRIVERS=15, got {cfg.READ_TILE_ARRIVERS}"), - (cfg.STAGES_KV == 2, "d192 SM100: STAGES_KV must be 2 for BF16/FP16"), (cfg.TILE_K_HW_BMM1 == 16 and cfg.TILE_K_HW_BMM2 == 16, "d192: TILE_K_HW must be 16 for BF16/FP16 on SM10x"), (cfg.Q_SWZ_BYTES == 128 and cfg.K_SWZ_BYTES == 128, "d192: Q/K swizzle must be 128B"), (cfg.V_SWZ_BYTES == 128 and cfg.O_SWZ_BYTES == 128, "d192: V/O swizzle must be 128B"), @@ -691,9 +832,15 @@ def make_cfg_d192(params: TemplateParams) -> Tuple[CfgD192, TmaIters]: DTYPE_O=params.dtype_qkv, BPE=b, BPE_O=b, + SPLIT_KV=int(params.split_kv), Q_SWZ_BYTES=q_swz_bytes(192, b), K_SWZ_BYTES=q_swz_bytes(192, b), - V_SWZ_BYTES=v_swz_bytes(128, 2, b), + CGA_M=params.cta_mma, + CTA_MMA=params.cta_mma, + # cuDNN's d192 kernel uses stages_kv = 1 * CTA_MMA; cga1 doubles per-CTA + # K/V, so the stage count has to halve to keep the CTA inside the cap. + STAGES_KV=1 if params.cta_mma == 1 else 2, + V_SWZ_BYTES=v_swz_bytes(128, params.cta_mma, b), O_SWZ_BYTES=o_swz_bytes(128, b), RESCALE_THRESHOLD=rescale_threshold(params.dtype_qkv), TILE_K_HW_BMM1=tile_k_hw(params.dtype_qkv), diff --git a/python/cudnn/sdpa/fwd/config_sm120.py b/python/cudnn/sdpa/fwd/config_sm120.py index 031ecb0ec..9726e04ec 100644 --- a/python/cudnn/sdpa/fwd/config_sm120.py +++ b/python/cudnn/sdpa/fwd/config_sm120.py @@ -71,6 +71,11 @@ class TemplateParams: sched_policy: int = SCHED_NATURAL q_tile: int = SEQ_Q_TILES[0] kv_tile: int = SEQ_KV_TILES[0] + # KV split: each Q tile's KV-tile range [min_kv_tile, num_kv_tiles) is cut + # into ``split_kv`` contiguous chunks, each run by its own CTA writing a + # partial (O, LSE) that kernels/split_combine_sm100.py reduces. 1 = off. + # Opt-in only -- the graph front door never selects it. + split_kv: int = 1 def validate_params( @@ -94,6 +99,22 @@ def validate_params( raise ValueError(f"SM120 SDPA: dtype_qkv must be one of {allowed_dtypes}; got {params.dtype_qkv}") if params.dtype_o not in allowed_o_dtypes: raise ValueError(f"SM120 SDPA: dtype_o must be one of {allowed_o_dtypes}; got {params.dtype_o}") + if params.split_kv < 1: + raise ValueError(f"SM120 SDPA: split_kv must be >= 1 (1 = KV-split off); got {params.split_kv}") + if params.split_kv > 1: + # Each of these would need extra machinery in the combine pass. + if params.thd_varlen: + raise ValueError("SM120 SDPA: split_kv > 1 is dense-only (THD packs its own flat grid)") + if params.has_sink: + # The sink logit is folded into the softmax denominator per tile, so + # every split would add its own copy of it. + raise ValueError("SM120 SDPA: split_kv > 1 with an attention sink is not supported") + if params.sched_policy != SCHED_NATURAL: + # The split index rides the NATURAL grid's batch axis (y = batch + + # split*B); LPT / LPT_L2 flatten the grid to 1-D and derive the batch + # from the linear tile id, so there is no axis left to fold it into. + raise ValueError(f"SM120 SDPA: split_kv > 1 currently requires SCHED_NATURAL; got sched_policy={params.sched_policy}") + if params.window_right is not None and params.window_right < 0: raise ValueError(f"SM120 SDPA: window_right must be None (unbounded) or >= 0 (0 = plain causal); got {params.window_right}") if not allow_right_band and params.window_right not in (None, 0): diff --git a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py index 336fca559..9023cbf53 100644 --- a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py @@ -302,6 +302,204 @@ def compute_kv_loop_bounds( ) +class SplitHelpers(NamedTuple): + """Split-aware decode / bounds closures, plus the two flags kernels fold on.""" + + SPLIT_KV: int + # True when a tile's KV range can come out empty (right <= left). See + # make_split_helpers for why KV split makes this reachable without a mask. + MAY_BE_EMPTY: bool + split_chunk: object + decode_initial_split: object + decode_payload_split: object + bounds_for_tile_split: object + nomask_range_split: object + partial_batch: object + + +def make_split_helpers(CFG, *, bounds_for_tile, dispatch_decode_initial, dispatch_decode_payload) -> SplitHelpers: + """Split-aware decode / bounds closures shared by the SM100 prefill flavors. + + ``bounds_for_tile`` is the caller's own bounds closure, taking + ``(q_super_idx, seqlen_q, seqlen_kv, cta_in_pair, seq_q_lens_tensor, + batch_idx)`` — flavors differ in whether they apply the dead-Q-tile trim, so + the split narrowing composes on top of whatever they already do. + + At SPLIT_KV == 1 every closure below folds away and the traced code is the + classic single-pass kernel. + """ + SPLIT_KV = int(getattr(CFG, "SPLIT_KV", 1)) + + # The split index rides the BATCH axis (grid.z), not grid.x: the decode + # already recovers the batch coordinate on BOTH the blockIdx and the + # scheduler-handout paths -- it is the high half of the packed head|batch + # word -- so a composite z = batch + split*B travels with it for free, with + # no in-place mutation of the shared tile id and no dependence on the grid's + # x extent (which is q_clusters * CGA_M, NOT the n_q_supers the kernel is + # handed, on any flavor where CGA_M != CTA_MMA -- d512). + + # Can a tile's KV range come out EMPTY (right <= left)? + # + # Before KV split the answer was "only under a mask" — a SWA/causal/padded + # tile can fall entirely outside the band — so the empty-tile handshake + # (mb_empty_mainloop: correction arrives, MMA waits, TMA-LDG skips its + # loads) was gated on MASK_FLAGS != 0 and folded away at MASK_NONE, where + # [0, S_kv/TILE_N) is never empty. + # + # KV split breaks that WITHOUT a mask: a split past the end of a short range + # legitimately gets zero tiles. Correction detects empties with a RUNTIME + # test, so if the gate const-folds to False the warp groups disagree — + # correction jumps to its epilogue while MMA waits on mb_q_full and TMA-LDG + # issues loads nobody consumes — and the kernel deadlocks. + MAY_BE_EMPTY = (CFG.MASK_FLAGS != 0) or (SPLIT_KV > 1) + + # Which grid does this flavor launch? SCHED_NATURAL uses a 3-D + # (q_super, head, batch) grid; the LPT policy flattens everything into x. + # NOTE this is the flavor's EFFECTIVE policy, not the requested one: + # make_cfg_d192 hardcodes SCHEDULER_POLICY=1 regardless of params, so a + # params-level check would miss it (and did -- d192 silently launched the + # unsplit grid because the split multiplier was only on the NATURAL branch). + IS_LPT = CFG.SCHEDULER_POLICY != SCHED_NATURAL + + @cute.jit + def _lpt_split_of(raw, n_q_supers, n_qh, n_batch): + """(within-split raw x, split) for the flattened LPT grid. + + The LPT tile space is q_tiles * n_qh * n_batch clusters; KV split + appends SPLIT_KV copies of it, split-major, so the split is the high + digit of the cluster index. The CGA lane (raw % CGA_M) is preserved so + the caller's decode still sees a well-formed x coordinate. + """ + cga = cutlass.Int32(CFG.CGA_M) + linear = raw // cga + q_tiles = n_q_supers // cutlass.Int32(CFG.CTA_MMA) + per_split = q_tiles * n_qh * n_batch + split = linear // per_split + rest = (linear % per_split) * cga + (raw % cga) + return rest, split + + @cute.jit + def _split_chunk(left, right, split_idx): + """Cut ``[left, right)`` into SPLIT_KV near-equal chunks. + + The FIRST ``rem`` splits get one extra tile, so chunk sizes differ by at + most 1 however the mask has already narrowed the range — a balanced cut + matters because the slowest split sets the critical path. A split past + the end collapses to ``lo == hi``; the existing empty-mainloop path then + writes O := 0 / LSE := -inf, exactly the identity of the combine's + log-sum-exp, so no special case is needed downstream. + """ + n_tiles = right - left + per = n_tiles // cutlass.Int32(SPLIT_KV) + rem = n_tiles % cutlass.Int32(SPLIT_KV) + lo = left + split_idx * per + cute.math.min(split_idx, rem) + extra = cutlass.Int32( + arith.select( + (split_idx < rem).ir_value(), + cutlass.Int32(1).ir_value(), + cutlass.Int32(0).ir_value(), + ) + ) + return lo, lo + per + extra + + @cute.jit + def _decode_initial_split(bidx, bidy, bidz, cta_in_pair, n_q_supers, n_qh, n_batch, seq_kv_lens_t, qh_per_kh=None, seqlen_kv=None): + """decode_initial + this tile's split index. + + NATURAL: the split rides the BATCH axis (see the note above on why not + grid.x). The host launches z = B * SPLIT_KV, so the split falls out of + the DECODED batch coordinate as ``b // n_batch``, leaving the real batch + as ``b % n_batch``. + + LPT / LPT_L2: the grid is flat, so there is no batch axis to ride and the + split is folded into the linear tile id instead; ``_lpt_split_of`` peels + it back off before the flavor's dispatcher sees the id. + + ``qh_per_kh`` / ``seqlen_kv`` are the LPT_L2 cost-model inputs; they are + opaque here and forwarded to the flavor's dispatcher unchanged. + """ + if cutlass.const_expr(SPLIT_KV > 1 and IS_LPT): + raw, split = _lpt_split_of(bidx, n_q_supers, n_qh, n_batch) + q, h, b = dispatch_decode_initial(raw, bidy, bidz, cta_in_pair, n_q_supers, n_qh, n_batch, seq_kv_lens_t, qh_per_kh, seqlen_kv) + return q, h, b, split + q, h, b = dispatch_decode_initial(bidx, bidy, bidz, cta_in_pair, n_q_supers, n_qh, n_batch, seq_kv_lens_t, qh_per_kh, seqlen_kv) + if cutlass.const_expr(SPLIT_KV == 1): + return q, h, b, cutlass.Int32(0) + return q, h, b % n_batch, b // n_batch + + @cute.jit + def _decode_payload_split(t0, t1, cta_in_pair, n_q_supers, n_qh, n_batch, seq_kv_lens_t, qh_per_kh=None, seqlen_kv=None): + """decode_payload + split index; ``t0`` is the try_cancel cluster-base id.""" + if cutlass.const_expr(SPLIT_KV > 1 and IS_LPT): + raw, split = _lpt_split_of(t0, n_q_supers, n_qh, n_batch) + q, h, b = dispatch_decode_payload(raw, t1, cta_in_pair, n_q_supers, n_qh, n_batch, seq_kv_lens_t, qh_per_kh, seqlen_kv) + return q, h, b, split + q, h, b = dispatch_decode_payload(t0, t1, cta_in_pair, n_q_supers, n_qh, n_batch, seq_kv_lens_t, qh_per_kh, seqlen_kv) + if cutlass.const_expr(SPLIT_KV == 1): + return q, h, b, cutlass.Int32(0) + return q, h, b % n_batch, b // n_batch + + @cute.jit + def _bounds_for_tile_split(q_super_idx, seqlen_q, seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx): + """The flavor's bounds, narrowed to this split's slice of the KV range. + + Splitting the ALREADY-masked ``[left, right)`` rather than the raw KV + extent is what keeps causal / SWA correct AND balanced: each split gets + an equal share of the tile's real work, not of the sequence. The + unmasked band is clamped into the slice, which preserves the + ``left <= unmasked_lo <= unmasked_hi <= right`` invariant the mainloop + relies on, because clamping is monotone. + """ + b = bounds_for_tile(q_super_idx, seqlen_q, seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + if cutlass.const_expr(SPLIT_KV == 1): + return b + lo, hi = _split_chunk(b.left, b.right, split_idx) + return KvLoopBounds( + left=lo, + unmasked_lo=cute.math.min(cute.math.max(b.unmasked_lo, lo), hi), + unmasked_hi=cute.math.min(cute.math.max(b.unmasked_hi, lo), hi), + right=hi, + ) + + @cute.jit + def _nomask_range_split(seqlen_kv, split_idx): + """MASK_NONE fast-path KV range, split-aware. + + MMA / TMA-LDG take this path while softmax / correction go through + bounds_for_tile_split; every warp group must land on the SAME chunk + boundaries or their mbarrier handshakes desync. The split branch + therefore divides with the div-up compute_kv_loop_bounds uses, not the + floor of the historical fast path. + """ + if cutlass.const_expr(SPLIT_KV == 1): + return cutlass.Int32(0), seqlen_kv // cutlass.Int32(CFG.TILE_N) + n_tiles = (seqlen_kv + cutlass.Int32(CFG.TILE_N - 1)) // cutlass.Int32(CFG.TILE_N) + return _split_chunk(cutlass.Int32(0), n_tiles, split_idx) + + @cute.jit + def _partial_batch(batch_idx, split_idx, n_batch): + """Batch coord of this split's partial O / LSE slot (split-major). + + Stacking the partials on the BATCH axis (extent B*SPLIT_KV) means the O + TMA descriptor is untouched — only the coord shifts. Folds to batch_idx + at SPLIT_KV == 1. + """ + if cutlass.const_expr(SPLIT_KV == 1): + return batch_idx + return batch_idx + split_idx * n_batch + + return SplitHelpers( + SPLIT_KV=SPLIT_KV, + MAY_BE_EMPTY=MAY_BE_EMPTY, + split_chunk=_split_chunk, + decode_initial_split=_decode_initial_split, + decode_payload_split=_decode_payload_split, + bounds_for_tile_split=_bounds_for_tile_split, + nomask_range_split=_nomask_range_split, + partial_batch=_partial_batch, + ) + + class SdpaHelpers(NamedTuple): decode_initial: object decode_payload: object diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py index 0d313815a..9cfdd3321 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -9,9 +9,13 @@ threads, persistent try_cancel scheduler, no Q∪O alias. SM100 resource layout: - 1. **cga2-only, STAGES_KV=2**: per-CTA K/V halved by the collective MMA → - SMEM 192 KiB (sQ 64 + sK 32 + sV 32 + sO 64), under the Blackwell - ~228 KiB cap. + 1. **STAGES_KV=2, both cluster widths at 192 KiB**: at cga2 (CTA_MMA=2, + the default) the collective MMA halves per-CTA K/V → sQ 64 + sK 32 + + sV 32 + sO 64. At cga1 (``cta_mma=1``) there is no collective MMA, so + K/V double to 64 + 64 and the 64 KiB is bought back by aliasing Q and O + into one slab (``QO_ALIAS``, mandatory at cga1 and enforced by + ``_validate_cfg_d128``) → 64(Q∪O) + 64 + 64. Both land under the + Blackwell ~228 KiB cap; ``_d128_smem_bytes`` is the checked model. 2. **TMEM stats** (512-col Blackwell cap): S_acc 0/128 + O 256..511 fill all 512 cols, so stats ride the FREE HEAD of each sub-tile's S_acc slot — sub-tile 0 → col 0, sub-tile 1 → col 128 (P only aliases the @@ -42,6 +46,30 @@ FP8 (E4M3 / E5M2, incl. output-dtype override) lives in the sibling ``prefill_d128_fp8_sm100.py`` (shares this flavor's config). + +KV split (``CFG.SPLIT_KV > 1``): each chunk runs the UNCHANGED mainloop and +epilogue, writing its partial O and LSE into a split-major workspace at batch +coord ``b + s*B`` -- so the O TMA descriptor is untouched, only the batch coord +shifts. ``split_combine_sm100.py`` reduces over the split axis. + +KV split composes with the cluster width, and the pair is what closes the gap +to cuDNN's own SM100 prefill split-K. cga1 halves both the wasted MMA work at +small S_q (a tile covers 256 Q rows, not 512) and the CTAs per tile, so twice +as many splits fit in one wave. Measured on B200 at B=1, H=16, S_q=128, +S_kv=32K, d=128, end-to-end incl. combine, max|O-ref| 2e-5 throughout: + + cga2 SPLIT_KV=1 414 us (32 CTAs — the classic kernel) + cga2 SPLIT_KV=4 117 us (128 CTAs) 3.6x + cga1 SPLIT_KV=8 73 us (128 CTAs) 5.7x <- best + (cuDNN 9.26/9.30's best plan for the same shape: 69 us, cga1 + 8 splits) + +On square shapes (S_q == S_kv, 1K/4K/8K) cga1 and cga2 measure within noise — +same CTA count, same work per CTA — so cga1 is a small-S_q lever, not a +regression risk elsewhere. +Gated (config_sm100._validate_params) to SCHED_NATURAL, dense (non-THD), no +sink; requires ``has_lse=True`` since the per-split LSE drives the combine. +At SPLIT_KV == 1 every split helper const-folds away and the traced code is the +classic single-pass kernel. """ from functools import lru_cache @@ -135,6 +163,7 @@ from cudnn.sdpa.fwd.kernels._common_sm100 import ( Bars, + make_split_helpers, KvLoopBounds, make_classic_bars, compute_kv_loop_bounds, @@ -200,6 +229,28 @@ _thd_tma_offsets = _sdpa_h.thd_tma_offsets +# === KV split === +# +# The mechanics live in _common_sm100.make_split_helpers (shared with the other +# SM100 prefill flavors); this is the small-S_q lever for this kernel: a cga2 +# cluster covers TILES_Q*TILE_M*CTA_MMA = 512 Q rows, so at S_q = 128 the whole +# problem is ceil(S_q/512) * H * B clusters (32 CTAs at H=16, B=1) on a 148-SM +# part, each walking the entire 32K-token KV loop alone. +_split_h = make_split_helpers( + CFG, + bounds_for_tile=_bounds_for_tile, + dispatch_decode_initial=_dispatch_decode_initial, + dispatch_decode_payload=_dispatch_decode_payload, +) +SPLIT_KV = _split_h.SPLIT_KV +MAY_BE_EMPTY = _split_h.MAY_BE_EMPTY +_decode_initial_split = _split_h.decode_initial_split +_decode_payload_split = _split_h.decode_payload_split +_bounds_for_tile_split = _split_h.bounds_for_tile_split +_nomask_range_split = _split_h.nomask_range_split +_partial_batch = _split_h.partial_batch + + # P (BMM2 operand) aliases the TAIL of each 128-col S_acc slot since BMM1 # finishes (and softmax loads S into registers) before P is written. # fp16/bf16 pack 2 probs per FP32 cell → P width = TILE_N/2 ≤ 64 cols → P @@ -617,7 +668,7 @@ def _tmaldg_warp_group( tma_k = GmemTileTma(tma_k_desc) tma_v = GmemTileTma(tma_v_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -634,13 +685,15 @@ def _tmaldg_warp_group( q_row_base = cute.arch.make_warp_uniform(q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M)) q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) + elif cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -654,7 +707,7 @@ def _tmaldg_warp_group( while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): pass else: # Prologue interleave Q[0] -> K[first] -> Q[1] -> V[first] -> mainloop — @@ -760,7 +813,7 @@ def _tmaldg_warp_group( nxt_q = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load()) nxt_hb = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load()) nxt_v = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load()) - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -777,10 +830,12 @@ def _tmaldg_warp_group( q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) - if cutlass.const_expr(CFG.MASK_FLAGS != 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV > 1): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -822,7 +877,7 @@ def _tmastg_warp_group( tma_o = GmemTileTma(tma_o_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -841,6 +896,10 @@ def _tmastg_warp_group( read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) q_row_base = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) + # KV split: partials are stacked split-major on the BATCH axis of the + # workspace (extent B*SPLIT_KV), so the store needs no new descriptor — + # only a shifted batch coord. Folds to batch_idx at SPLIT_KV == 1. + o_batch = _partial_batch(batch_idx, split_idx, n_batch) for qs in cutlass.range_constexpr(CFG.TILES_Q): bars.mb_o_full[qs].wait(o_full_phase) @@ -861,7 +920,7 @@ def _tmastg_warp_group( else: tma_store_tile( sO[qs], - tma_o(cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), batch_idx), + tma_o(cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), o_batch), ) tma_store_commit() @@ -879,7 +938,7 @@ def _tmastg_warp_group( nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1022,11 +1081,13 @@ def _mma_warp_group( desc_Q0 = sQ[0].desc() desc_Q1 = sQ[1].desc() - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) else: - q_super_idx, _hd, batch_idx = _dispatch_decode_initial( + # Under KV split the MASK_NONE path still has to decode, because the + # split index (and hence this tile's KV slice) lives in the tile id. + q_super_idx, _hd, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1038,11 +1099,14 @@ def _mma_warp_group( qh_per_kh, seqlen_kv, ) - eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) - kv_left = bounds_init.left - kv_right = bounds_init.right + if cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + else: + eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) + kv_left = bounds_init.left + kv_right = bounds_init.right q_full_phase = cutlass.Int32(0) kv_state = PipelineState.start(phase=0) @@ -1065,7 +1129,7 @@ def _mma_warp_group( while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): # Empty mainloop: fire both bmm2_done so softmax/correction phase trackers stay in lockstep. bars.mb_empty_mainloop.wait(empty_mainloop_phase) empty_mainloop_phase = empty_mainloop_phase ^ cutlass.Int32(1) @@ -1224,14 +1288,14 @@ def _mma_warp_group( nvvm.bar_warp_sync(cute.arch.FULL_MASK) wait(sched.mb_scheduler.subview(sched_state.idx), sched_state.phase) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() is_valid_tile = nxt_v & cutlass.Int32(1) else: nxt_q = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load()) nxt_hb = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load()) nxt_v = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load()) - q_super_idx, _hd, batch_idx = _dispatch_decode_payload( + q_super_idx, _hd, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1243,11 +1307,14 @@ def _mma_warp_group( seqlen_kv, ) is_valid_tile = nxt_v & cutlass.Int32(1) - eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) - kv_left = bounds_next.left - kv_right = bounds_next.right + if cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + else: + eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) + kv_left = bounds_next.left + kv_right = bounds_next.right sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) bars.mb_tmem_dealloc.wait(cutlass.Int32(0)) @@ -1509,7 +1576,7 @@ def _softmax_warp_group( cutlass.Float32, ) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1527,7 +1594,7 @@ def _softmax_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) softmax_wg_base_const = CFG.SOFTMAX_WG0_BASE if sub_tile_id == 0 else CFG.SOFTMAX_WG1_BASE tid_in_wg = cute.arch.thread_idx()[0] - cutlass.Int32(softmax_wg_base_const * 32) @@ -1644,7 +1711,7 @@ def _softmax_warp_group( nxt_q = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load()) nxt_hb = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load()) nxt_v = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load()) - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1659,7 +1726,7 @@ def _softmax_warp_group( sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) @cute.jit @@ -1713,7 +1780,7 @@ def _correction_warp_group( bmm2_done_phase = cutlass.Int32(0) o_empty_phase = cutlass.Int32(1) # bootstrap - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1731,7 +1798,7 @@ def _correction_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) @@ -1884,7 +1951,12 @@ def _correction_warp_group( else: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) - lse_row = lse_arr[batch_idx, head_idx, :] + # KV split: this chunk's LSE goes to its own split-major + # slot (batch extent B*SPLIT_KV), matching where TMA-STG put + # the chunk's O. The pair (O_s, lse_s) is everything the + # combine needs. Folds to batch_idx at SPLIT_KV == 1. + lse_batch = _partial_batch(batch_idx, split_idx, n_batch) + lse_row = lse_arr[lse_batch, head_idx, :] lse_row[q_row_global] = lse_val sO_sub_base = sO[qs].base @@ -1929,7 +2001,7 @@ def _correction_warp_group( nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1944,7 +2016,7 @@ def _correction_warp_group( sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) # End-of-warp tmem_dealloc: under cga2 each corr lane ALSO DSMEM-arrives # on the peer so the peer's local mbar accumulates the full CGA-total count. @@ -2072,7 +2144,13 @@ def _tma_swz(byte_w: int): grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1)) else: # Grid Python-folds on Cfg constant (avoids DSL if staging). - grid_shape = (grid_q_supers, QH, B) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B, 1, 1) + # KV split rides the BATCH axis: z = batch + split*B. The decode + # already recovers the batch coord on both the blockIdx and the + # scheduler-handout paths, so the split travels with it for free. SPLIT_KV > 1 is + # gated to SCHED_NATURAL by the config validator. + grid_shape = ( + (grid_q_supers, QH, B * SPLIT_KV) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B * SPLIT_KV, 1, 1) + ) _kernel( tma_q_desc, tma_k_desc, @@ -2147,6 +2225,10 @@ def compile( # noqa: A001 raise ValueError(f"d128 envelope: need 0 < d_qk <= {CFG.TILE_K} and 0 < d_v <= {CFG.TILE_O}; got ({d_qk}, {d_v})") if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d128 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") + if SPLIT_KV > 1 and not has_lse: + # Each split's LSE is not optional under KV split — it IS the weight the + # combine reduces with. Without it the partials cannot be recombined. + raise ValueError("d128: split_kv > 1 requires has_lse=True (the per-split LSE drives the combine)") _fake_batch = 1 if CFG.THD_VARLEN else b if CFG.THD_VARLEN: # Dynamic packed token totals: one symbol per ragged group (Q/O and @@ -2154,6 +2236,11 @@ def compile( # noqa: A001 # compiled artifact instead of minting a new one (issue #552). sq = cute.sym_int(divisibility=1) skv = cute.sym_int(divisibility=1) + # KV split: O and LSE are the PARTIAL workspaces, stacked split-major on the + # batch axis (B*SPLIT_KV). Q/K/V keep the real batch — only the outputs + # grow. At SPLIT_KV == 1 these are the plain output tensors. + _o_batch = _fake_batch * SPLIT_KV + _lse_batch = b * SPLIT_KV def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: @@ -2172,7 +2259,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) - fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=STORAGE_DTYPE) + fake_o = _fake_bshd((_o_batch, sq, qh, d_v), o_stride, dtype=STORAGE_DTYPE) if not has_lse: # No Stats output: the LSE argument is None-specialized and the store # is compiled out entirely — no dummy buffer exists at any level. @@ -2212,7 +2299,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): raise ValueError("lse_head_major / lse_head_stride are THD-only (dense LSE is compact (B, H, Sq))") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (b, qh, sq), + (_lse_batch, qh, sq), stride_order=(2, 1, 0), assumed_align=16, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py index 84d2e5c12..fd9118a39 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py @@ -151,6 +151,7 @@ from cudnn.sdpa.fwd.kernels._common_sm100 import ( + make_split_helpers, Bars, KvLoopBounds, make_classic_bars, @@ -195,6 +196,41 @@ _thd_tma_offsets = _sdpa_h.thd_tma_offsets +# === KV split === +# +# Mechanics live in _common_sm100.make_split_helpers, shared with the other +# SM100 prefill flavors: each Q tile's KV loop range is cut into SPLIT_KV +# contiguous chunks, each run as its own persistent tile, and each writing a +# normalized partial O + its own LSE into a split-major workspace that +# split_combine_sm100 folds with the exact log-sum-exp identity. At +# SPLIT_KV == 1 every closure folds away and this is the classic kernel. + + +@cute.jit +def _bounds_for_tile_uniform(q_super_idx, seqlen_q, seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx): + """Uniform 6-arg bounds signature for make_split_helpers. + + This flavor has no dead-Q-tile trim, so the trailing two args are + accepted and ignored — callers pass None for them. + """ + return _bounds_for_tile(q_super_idx, seqlen_q, seqlen_kv, cta_in_pair) + + +_split_h = make_split_helpers( + CFG, + bounds_for_tile=_bounds_for_tile_uniform, + dispatch_decode_initial=_dispatch_decode_initial, + dispatch_decode_payload=_dispatch_decode_payload, +) +SPLIT_KV = _split_h.SPLIT_KV +MAY_BE_EMPTY = _split_h.MAY_BE_EMPTY +_decode_initial_split = _split_h.decode_initial_split +_decode_payload_split = _split_h.decode_payload_split +_bounds_for_tile_split = _split_h.bounds_for_tile_split +_nomask_range_split = _split_h.nomask_range_split +_partial_batch = _split_h.partial_batch + + @dataclass(frozen=True) class KernelTmemLayout: """Column offsets for the classic 2-sub-tile SDPA pipeline (FP8). @@ -586,7 +622,7 @@ def _tmaldg_warp_group( tma_k = GmemTileTma(tma_k_desc) tma_v = GmemTileTma(tma_v_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -603,13 +639,15 @@ def _tmaldg_warp_group( q_row_base = cute.arch.make_warp_uniform(q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M)) q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) + elif cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -623,7 +661,7 @@ def _tmaldg_warp_group( while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): pass else: # Prologue interleave: Q[0] → K[first] → Q[1] → V[first] → mainloop. @@ -722,7 +760,7 @@ def _tmaldg_warp_group( nxt_q = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load()) nxt_hb = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load()) nxt_v = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load()) - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -739,10 +777,12 @@ def _tmaldg_warp_group( q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) - if cutlass.const_expr(CFG.MASK_FLAGS != 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV > 1): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -780,7 +820,7 @@ def _tmastg_warp_group( tma_o = GmemTileTma(tma_o_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -799,6 +839,10 @@ def _tmastg_warp_group( read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) q_row_base = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) + # KV split: partials are stacked split-major on the workspace BATCH axis + # (extent B*SPLIT_KV), so the store needs no new descriptor — only a + # shifted batch coord. Folds to batch_idx at SPLIT_KV == 1. + o_batch = _partial_batch(batch_idx, split_idx, n_batch) for qs in cutlass.range_constexpr(CFG.TILES_Q): bars.mb_o_full[qs].wait(o_full_phase) @@ -806,7 +850,7 @@ def _tmastg_warp_group( # O TMA params follow O's swizzle, not V's (V and O swizzles may differ). tma_store_tile( sO[qs], - tma_o(cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), batch_idx), + tma_o(cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), o_batch), ) tma_store_commit() @@ -820,7 +864,7 @@ def _tmastg_warp_group( nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -955,11 +999,11 @@ def _mma_warp_group( desc_Q0 = sQ[0].desc() desc_Q1 = sQ[1].desc() - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) else: - q_super_idx, _hd, batch_idx = _dispatch_decode_initial( + q_super_idx, _hd, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -971,11 +1015,14 @@ def _mma_warp_group( qh_per_kh, seqlen_kv, ) - eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) - kv_left = bounds_init.left - kv_right = bounds_init.right + if cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + else: + eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) + kv_left = bounds_init.left + kv_right = bounds_init.right q_full_phase = cutlass.Int32(0) kv_state = PipelineState.start(phase=0) @@ -999,7 +1046,7 @@ def _mma_warp_group( while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): # Empty-kv tile: fire bmm2_done so softmax/corr phases stay in lockstep. bars.mb_empty_mainloop.wait(empty_mainloop_phase) empty_mainloop_phase = empty_mainloop_phase ^ cutlass.Int32(1) @@ -1150,14 +1197,14 @@ def _mma_warp_group( nvvm.bar_warp_sync(cute.arch.FULL_MASK) wait(sched.mb_scheduler.subview(sched_state.idx), sched_state.phase) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() is_valid_tile = nxt_v & cutlass.Int32(1) else: nxt_q = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load()) nxt_hb = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load()) nxt_v = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load()) - q_super_idx, _hd, batch_idx = _dispatch_decode_payload( + q_super_idx, _hd, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1169,11 +1216,14 @@ def _mma_warp_group( seqlen_kv, ) is_valid_tile = nxt_v & cutlass.Int32(1) - eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) - kv_left = bounds_next.left - kv_right = bounds_next.right + if cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + else: + eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) + kv_left = bounds_next.left + kv_right = bounds_next.right sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) bars.mb_tmem_dealloc.wait(cutlass.Int32(0)) @@ -1404,7 +1454,7 @@ def _softmax_warp_group( cutlass.Float32, ) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1422,7 +1472,7 @@ def _softmax_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) softmax_wg_base_const = CFG.SOFTMAX_WG0_BASE if sub_tile_id == 0 else CFG.SOFTMAX_WG1_BASE tid_in_wg = cute.arch.thread_idx()[0] - cutlass.Int32(softmax_wg_base_const * 32) @@ -1539,7 +1589,7 @@ def _softmax_warp_group( nxt_q = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load()) nxt_hb = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load()) nxt_v = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load()) - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1554,7 +1604,7 @@ def _softmax_warp_group( sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) @cute.jit @@ -1608,7 +1658,7 @@ def _correction_warp_group( bmm2_done_phase = cutlass.Int32(0) o_empty_phase = cutlass.Int32(1) # bootstrap pre-armed at phase 1 so first wait passes - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1626,7 +1676,7 @@ def _correction_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) @@ -1748,7 +1798,11 @@ def _correction_warp_group( if _row_valid: if cutlass.const_expr(lse_tensor is not None): lse_arr = cutlass.make_array_view(lse_tensor) - lse_row = lse_arr[batch_idx, head_idx, :] + # This chunk's LSE goes to its own split-major slot, matching where + # TMA-STG put the chunk's O. The pair (O_s, lse_s) is everything + # the combine needs. + lse_batch = _partial_batch(batch_idx, split_idx, n_batch) + lse_row = lse_arr[lse_batch, head_idx, :] lse_row[q_row_global] = lse_val # amax_o = max over valid rows of |o_scaled| (the fp32 pre-cast output). Divided @@ -1837,8 +1891,14 @@ def _correction_warp_group( bars.mb_o_empty[qs].wait(o_empty_phase) smem_ptr.store_swizzled(o_half, alignment=64, swizzle=_O_SMEM_SWIZZLE) - if _row_valid: - nvvm.atomicrmw(nvvm.AtomicOp.MAX, _amax_o_ptr, _amax_o_local.bitcast(cutlass.Int32)) + # Under KV split this epilogue sees only its OWN partial, and the + # recombined O is a convex combination of the partials -- so a max + # over partials over-reports the output amax (~2.9x at 8 splits). + # split_combine_sm100 computes it over the recombined O instead; + # this write has to stay out of the way, since atomicMax only grows. + if cutlass.const_expr(SPLIT_KV == 1): + if _row_valid: + nvvm.atomicrmw(nvvm.AtomicOp.MAX, _amax_o_ptr, _amax_o_local.bitcast(cutlass.Int32)) # fence_proxy needed before TMA reads SMEM written by stores above. nvvm.fence_proxy("async.shared", space="cta") @@ -1855,7 +1915,7 @@ def _correction_warp_group( nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1870,7 +1930,7 @@ def _correction_warp_group( sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) # tmem_dealloc fan-out: fire one arrive per lane; cga2 also DSMEM-arrives # on the peer so peer's local mbar accumulates the full CGA count. @@ -1956,7 +2016,10 @@ def _tma_swz(byte_w: int): q_clusters = (SQ + rows_per_cluster - 1) // rows_per_cluster grid_q_supers = q_clusters * CFG.CTA_MMA q_supers = grid_q_supers - grid_shape = (grid_q_supers, QH, B) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B, 1, 1) + # KV split rides the BATCH axis: z = batch + split*B. The decode + # already recovers the batch coord on both the blockIdx and the + # scheduler-handout paths, so the split travels with it for free. + grid_shape = (grid_q_supers, QH, B * SPLIT_KV) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B * SPLIT_KV, 1, 1) _kernel( tma_q_desc, tma_k_desc, @@ -1993,6 +2056,14 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, ``has_lse=False`` compiles the LSE store out (the kernel specializes on a ``None`` LSE argument) — callers without a Stats output pass no LSE buffer at all; the amax_o atomicMax write is independent and unchanged.""" + if SPLIT_KV > 1 and not has_lse: + # Each split's LSE is not optional under KV split — it IS the weight + # the combine reduces with. Without it the partials cannot be recombined. + raise ValueError("split_kv > 1 requires has_lse=True (the per-split LSE drives the combine)") + # KV split: O and LSE are the PARTIAL workspaces, stacked split-major on + # the batch axis (B*SPLIT_KV). Q/K/V keep the real batch. + _o_batch = b * SPLIT_KV + _lse_batch = b * SPLIT_KV fake_q = cute.runtime.make_fake_compact_tensor( STORAGE_DTYPE, (b, sq, qh, CFG.TILE_K), @@ -2013,7 +2084,7 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, ) fake_o = cute.runtime.make_fake_compact_tensor( OUT_STORAGE_DTYPE, - (b, sq, qh, CFG.TILE_O), + (_o_batch, sq, qh, CFG.TILE_O), stride_order=(3, 2, 1, 0), assumed_align=16, ) @@ -2024,7 +2095,7 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, else: fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (b, qh, sq), + (_lse_batch, qh, sq), stride_order=(2, 1, 0), assumed_align=16, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py index e629faa0c..c6870fed0 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py @@ -196,6 +196,7 @@ def _round_up(a: int, b: int) -> int: from cudnn.sdpa.fwd.kernels._common_sm100 import ( + make_split_helpers, Bars, KvLoopBounds, make_classic_bars, @@ -238,6 +239,41 @@ def _round_up(a: int, b: int) -> int: _dispatch_decode_initial = _sdpa_h.dispatch_decode_initial _dispatch_decode_payload = _sdpa_h.dispatch_decode_payload _thd_tma_offsets = _sdpa_h.thd_tma_offsets + + +# === KV split === +# +# Mechanics live in _common_sm100.make_split_helpers, shared with the other +# SM100 prefill flavors: each Q tile's KV loop range is cut into SPLIT_KV +# contiguous chunks, each run as its own persistent tile, and each writing a +# normalized partial O + its own LSE into a split-major workspace that +# split_combine_sm100 folds with the exact log-sum-exp identity. At +# SPLIT_KV == 1 every closure folds away and this is the classic kernel. + + +@cute.jit +def _bounds_for_tile_uniform(q_super_idx, seqlen_q, seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx): + """Uniform 6-arg bounds signature for make_split_helpers. + + This flavor has no dead-Q-tile trim, so the trailing two args are + accepted and ignored — callers pass None for them. + """ + return _bounds_for_tile(q_super_idx, seqlen_q, seqlen_kv, cta_in_pair) + + +_split_h = make_split_helpers( + CFG, + bounds_for_tile=_bounds_for_tile_uniform, + dispatch_decode_initial=_dispatch_decode_initial, + dispatch_decode_payload=_dispatch_decode_payload, +) +SPLIT_KV = _split_h.SPLIT_KV +MAY_BE_EMPTY = _split_h.MAY_BE_EMPTY +_decode_initial_split = _split_h.decode_initial_split +_decode_payload_split = _split_h.decode_payload_split +_bounds_for_tile_split = _split_h.bounds_for_tile_split +_nomask_range_split = _split_h.nomask_range_split +_partial_batch = _split_h.partial_batch _thd_sf_tile_bases = _sdpa_h.thd_sf_tile_bases @@ -773,7 +809,7 @@ def _tmaldg_warp_group( k_sf_peer_row = cta_in_pair * cutlass.Int32(K_SF_ROWS_PER_PEER) v_sf_peer_row = cta_in_pair * cutlass.Int32(V_SF_ROWS_PER_PEER) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -793,13 +829,15 @@ def _tmaldg_warp_group( q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) cu_sf_q_base, cu_sf_k_base = _thd_sf_tile_bases(seq_kv_lens_tensor, batch_idx, n_batch) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) + elif cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -818,7 +856,7 @@ def _tmaldg_warp_group( read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) # Empty-kv tile: MMA's empty-mainloop branch handles the matching mbar phases. - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): pass else: q_sf_tile_base = q_row_base // cutlass.Int32(CFG.TILE_M) @@ -976,7 +1014,7 @@ def _tmaldg_warp_group( nxt_q = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load()) nxt_hb = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load()) nxt_v = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load()) - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -993,10 +1031,12 @@ def _tmaldg_warp_group( cu_sf_q_base, cu_sf_k_base = _thd_sf_tile_bases(seq_kv_lens_tensor, batch_idx, n_batch) is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) - if cutlass.const_expr(CFG.MASK_FLAGS != 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV > 1): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -1031,7 +1071,7 @@ def _tmastg_warp_group( tma_o = GmemTileTma(tma_o_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1050,6 +1090,10 @@ def _tmastg_warp_group( read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) q_row_base = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) + # KV split: partials are stacked split-major on the workspace BATCH axis + # (extent B*SPLIT_KV), so the store needs no new descriptor — only a + # shifted batch coord. Folds to batch_idx at SPLIT_KV == 1. + o_batch = _partial_batch(batch_idx, split_idx, n_batch) for qs in cutlass.range_constexpr(CFG.TILES_Q): bars.mb_o_full[qs].wait(o_full_phase) @@ -1057,7 +1101,7 @@ def _tmastg_warp_group( # O TMA params follow O's swizzle, NOT V's — required when V_SWZ_B != O_SWZ_B. tma_store_tile( sO[qs], - tma_o(cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), batch_idx), + tma_o(cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), o_batch), ) tma_store_commit() @@ -1071,7 +1115,7 @@ def _tmastg_warp_group( nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1255,11 +1299,11 @@ def _utccp_bmm2_sf(tmem_sf_p, tmem_sf_v, smem_desc_p, smem_desc_v): nvvm.tcgen05_cp(nvvm.Tcgen05CpShape.SHAPE_32X128B, tmem_sf_p, smem_desc_p, group=CTA_GROUP_KIND, multicast=nvvm.Tcgen05CpMulticast.WARPX4) nvvm.tcgen05_cp(nvvm.Tcgen05CpShape.SHAPE_32X128B, tmem_sf_v, smem_desc_v, group=CTA_GROUP_KIND, multicast=nvvm.Tcgen05CpMulticast.WARPX4) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) else: - q_super_idx, _hd, batch_idx = _dispatch_decode_initial( + q_super_idx, _hd, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1271,11 +1315,14 @@ def _utccp_bmm2_sf(tmem_sf_p, tmem_sf_v, smem_desc_p, smem_desc_v): qh_per_kh, seqlen_kv, ) - eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) - kv_left = bounds_init.left - kv_right = bounds_init.right + if cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + else: + eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) + kv_left = bounds_init.left + kv_right = bounds_init.right q_full_phase = cutlass.Int32(0) kv_state = PipelineState.start(phase=0) @@ -1304,7 +1351,7 @@ def _utccp_bmm2_sf(tmem_sf_p, tmem_sf_v, smem_desc_p, smem_desc_v): while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): # Empty mainloop — keep softmax/correction phase trackers in lockstep with non-empty path. bars.mb_empty_mainloop.wait(empty_mainloop_phase) empty_mainloop_phase = empty_mainloop_phase ^ cutlass.Int32(1) @@ -1552,14 +1599,14 @@ def _utccp_bmm2_sf(tmem_sf_p, tmem_sf_v, smem_desc_p, smem_desc_v): nvvm.bar_warp_sync(cute.arch.FULL_MASK) wait(sched.mb_scheduler.subview(sched_state.idx), sched_state.phase) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() is_valid_tile = nxt_v & cutlass.Int32(1) else: nxt_q = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load()) nxt_hb = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load()) nxt_v = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load()) - q_super_idx, _hd, batch_idx = _dispatch_decode_payload( + q_super_idx, _hd, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1571,11 +1618,14 @@ def _utccp_bmm2_sf(tmem_sf_p, tmem_sf_v, smem_desc_p, smem_desc_v): seqlen_kv, ) is_valid_tile = nxt_v & cutlass.Int32(1) - eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) - kv_left = bounds_next.left - kv_right = bounds_next.right + if cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + else: + eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) + kv_left = bounds_next.left + kv_right = bounds_next.right sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) bars.mb_tmem_dealloc.wait(cutlass.Int32(0)) @@ -1838,7 +1888,7 @@ def _softmax_warp_group( cutlass.Float32, ) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1856,7 +1906,7 @@ def _softmax_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) softmax_wg_base_const = CFG.SOFTMAX_WG0_BASE if sub_tile_id == 0 else CFG.SOFTMAX_WG1_BASE tid_in_wg = cute.arch.thread_idx()[0] - cutlass.Int32(softmax_wg_base_const * 32) @@ -1972,7 +2022,7 @@ def _softmax_warp_group( nxt_q = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load()) nxt_hb = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load()) nxt_v = cute.arch.make_warp_uniform((sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load()) - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1987,7 +2037,7 @@ def _softmax_warp_group( sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) @cute.jit @@ -2035,7 +2085,7 @@ def _correction_warp_group( bmm2_done_phase = cutlass.Int32(0) o_empty_phase = cutlass.Int32(1) # bootstrap pre-armed at phase 1 so first wait passes immediately - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -2053,7 +2103,7 @@ def _correction_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) @@ -2166,7 +2216,11 @@ def _correction_warp_group( if cutlass.const_expr(lse_tensor is not None): if _row_valid: lse_arr = cutlass.make_array_view(lse_tensor) - lse_row = lse_arr[batch_idx, head_idx, :] + # This chunk's LSE goes to its own split-major slot, matching where + # TMA-STG put the chunk's O. The pair (O_s, lse_s) is everything + # the combine needs. + lse_batch = _partial_batch(batch_idx, split_idx, n_batch) + lse_row = lse_arr[lse_batch, head_idx, :] lse_row[q_row_global] = lse_val sO_sub_base = sO[qs].base @@ -2200,8 +2254,15 @@ def _correction_warp_group( smem_ptr.store_swizzled(o_out, alignment=64, swizzle=_O_SMEM_SWIZZLE) # One atomic per valid row (invalid/OOB rows must not poison the global amax). - if _row_valid: - nvvm.atomicrmw(nvvm.AtomicOp.MAX, _amax_o_ptr, _amax_o_local.bitcast(cutlass.Int32)) + # + # Under KV split this epilogue sees only its OWN partial, and the + # recombined O is a convex combination of the partials -- so a max + # over partials over-reports the output amax. split_combine_sm100 + # computes it over the recombined O instead; this write has to stay + # out of the way, since atomicMax only grows. + if cutlass.const_expr(SPLIT_KV == 1): + if _row_valid: + nvvm.atomicrmw(nvvm.AtomicOp.MAX, _amax_o_ptr, _amax_o_local.bitcast(cutlass.Int32)) # fence_proxy needed before TMA reads SMEM written by tcgen05_st. nvvm.fence_proxy("async.shared", space="cta") @@ -2217,7 +2278,7 @@ def _correction_warp_group( nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -2232,7 +2293,7 @@ def _correction_warp_group( sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, None, None, split_idx) # mb_tmem_dealloc fan-out — all-lanes arrive + DSMEM-arrive on cross-pair peer under cga2. if cutlass.const_expr(CFG.CTA_MMA == 2): @@ -2353,7 +2414,10 @@ def _build_sf_desc(sf_tensor, num_tiles, sf_smem_size, num_rows_box, num_heads): q_clusters = (SQ + rows_per_cluster - 1) // rows_per_cluster grid_q_supers = q_clusters * CFG.CTA_MMA q_supers = grid_q_supers - grid_shape = (grid_q_supers, QH, B) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B, 1, 1) + # KV split rides the BATCH axis: z = batch + split*B. The decode + # already recovers the batch coord on both the blockIdx and the + # scheduler-handout paths, so the split travels with it for free. + grid_shape = (grid_q_supers, QH, B * SPLIT_KV) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B * SPLIT_KV, 1, 1) _kernel( tma_q_desc, tma_k_desc, @@ -2395,6 +2459,15 @@ def compile( # noqa: A001 ``has_lse=False`` compiles the LSE store out (the kernel specializes on a ``None`` LSE argument) — callers without a Stats output pass no LSE buffer at all; the amax_o atomicMax write is independent and unchanged.""" + if SPLIT_KV > 1 and not has_lse: + # Each split's LSE is not optional under KV split — it IS the weight + # the combine reduces with. Without it the partials cannot be recombined. + raise ValueError("split_kv > 1 requires has_lse=True (the per-split LSE drives the combine)") + # KV split: O and LSE are the PARTIAL workspaces, stacked split-major on + # the batch axis (B*SPLIT_KV). Q/K/V keep the real batch. + _o_batch = b * SPLIT_KV + _lse_batch = b * SPLIT_KV + # Q SF tiles TILE_M-row wide → num_tiles = SQ/TILE_M; K/V SF TILE_N-row wide → num_tiles = SKV/TILE_N. sq_tiles = (sq + CFG.TILE_M - 1) // CFG.TILE_M skv_tiles = (skv + CFG.TILE_N - 1) // CFG.TILE_N @@ -2421,7 +2494,7 @@ def compile( # noqa: A001 ) fake_o = cute.runtime.make_fake_compact_tensor( OUT_STORAGE_DTYPE, - (b, sq, qh, CFG.TILE_O), + (_o_batch, sq, qh, CFG.TILE_O), stride_order=(3, 2, 1, 0), assumed_align=16, ) @@ -2452,7 +2525,7 @@ def compile( # noqa: A001 else: fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (b, qh, sq), + (_lse_batch, qh, sq), stride_order=(2, 1, 0), assumed_align=16, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py index 6eac98386..d0b4fb547 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py @@ -137,6 +137,7 @@ def _wait_mbarrier(mb, phase): from cudnn.sdpa.fwd.kernels._common_sm100 import ( + make_split_helpers, make_classic_bars, row_max_for_exp2, make_sdpa_helpers, @@ -211,6 +212,38 @@ def _bounds_for_tile( _thd_tma_offsets = _sdpa_h.thd_tma_offsets +# === KV split === +# +# Mechanics live in _common_sm100.make_split_helpers, shared with the other +# SM100 prefill flavors: each Q tile's KV loop range is cut into SPLIT_KV +# contiguous chunks, each run as its own persistent tile, and each writing a +# normalized partial O + its own LSE into a split-major workspace that +# split_combine_sm100 folds with the exact log-sum-exp identity. At +# SPLIT_KV == 1 every closure folds away and this is the classic kernel. +_split_h = make_split_helpers( + CFG, + bounds_for_tile=_bounds_for_tile, + dispatch_decode_initial=_dispatch_decode_initial, + dispatch_decode_payload=_dispatch_decode_payload, +) +SPLIT_KV = _split_h.SPLIT_KV +MAY_BE_EMPTY = _split_h.MAY_BE_EMPTY +_decode_initial_split = _split_h.decode_initial_split +_decode_payload_split = _split_h.decode_payload_split +_bounds_for_tile_split = _split_h.bounds_for_tile_split +_nomask_range_split = _split_h.nomask_range_split +_partial_batch = _split_h.partial_batch + +# This flavor carries its OWN empty-KV predicate (defined above), narrower than +# MAY_BE_EMPTY: it counts only PADDED / SWA / bottom-right, the masks that can +# empty a tile. KV split makes an empty range reachable without any of them — a +# split past the end of a short range gets zero tiles — so the nine sites gated +# on CAN_HAVE_EMPTY_KV (the empty-tile handshake, the QO_ALIAS drain, and the +# epilogue's O-store guards) have to compile in under split as well. Leaving it +# const-folded to False is what desynchronises the warp groups. +CAN_HAVE_EMPTY_KV = CAN_HAVE_EMPTY_KV or SPLIT_KV > 1 + + @cute.jit def _apply_top_left_causal_mask_chunk(reg_S, q_abs, kv_col_base, N: int = 64): neg_inf = cutlass.Float32(float("-inf")) @@ -671,7 +704,7 @@ def _tmaldg_warp_group( tma_k = GmemTileTma(tma_k_desc) tma_v = GmemTileTma(tma_v_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -688,13 +721,15 @@ def _tmaldg_warp_group( q_row_base = cute.arch.make_warp_uniform(q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M)) q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) + elif cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -708,7 +743,7 @@ def _tmaldg_warp_group( while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): if cutlass.const_expr(CAN_HAVE_EMPTY_KV and IS_QO_ALIAS): # TMA-STG advances the Q/O alias gate for every tile, including # empty ones. Consume that transaction even though no Q reload @@ -830,7 +865,7 @@ def _tmaldg_warp_group( nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -847,10 +882,12 @@ def _tmaldg_warp_group( q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) - if cutlass.const_expr(CFG.MASK_FLAGS != 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV > 1): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -895,7 +932,7 @@ def _tmastg_warp_group( tma_o = GmemTileTma(tma_o_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -914,6 +951,10 @@ def _tmastg_warp_group( read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) q_row_base = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) + # KV split: partials are stacked split-major on the workspace BATCH axis + # (extent B*SPLIT_KV), so the store needs no new descriptor — only a + # shifted batch coord. Folds to batch_idx at SPLIT_KV == 1. + o_batch = _partial_batch(batch_idx, split_idx, n_batch) for qs in cutlass.range_constexpr(CFG.TILES_Q): _wait_mbarrier(bars.mb_o_full[qs], o_full_phase) @@ -934,7 +975,7 @@ def _tmastg_warp_group( else: tma_store_tile( sO[qs], - tma_o(cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), batch_idx), + tma_o(cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), o_batch), ) tma_store_commit() @@ -963,7 +1004,7 @@ def _tmastg_warp_group( nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1106,11 +1147,11 @@ def _mma_warp_group( desc_Q0 = sQ[0].desc() desc_Q1 = sQ[1].desc() - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) else: - q_super_idx, _hd, batch_idx = _dispatch_decode_initial( + q_super_idx, _hd, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1122,11 +1163,14 @@ def _mma_warp_group( qh_per_kh, seqlen_kv, ) - eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) - kv_left = bounds_init.left - kv_right = bounds_init.right + if cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + else: + eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) + kv_left = bounds_init.left + kv_right = bounds_init.right q_full_phase = cutlass.Int32(0) kv_state = PipelineState.start(phase=0) @@ -1149,7 +1193,7 @@ def _mma_warp_group( while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): if cutlass.const_expr(not CAN_HAVE_EMPTY_KV): _wait_mbarrier(bars.mb_empty_mainloop, empty_mainloop_phase) empty_mainloop_phase = empty_mainloop_phase ^ cutlass.Int32(1) @@ -1307,14 +1351,14 @@ def _mma_warp_group( nvvm.bar_warp_sync(cute.arch.FULL_MASK) _wait_ptr(sched.mb_scheduler.subview(sched_state.idx), sched_state.phase) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() is_valid_tile = nxt_v & cutlass.Int32(1) else: nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, _hd, batch_idx = _dispatch_decode_payload( + q_super_idx, _hd, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1326,11 +1370,14 @@ def _mma_warp_group( seqlen_kv, ) is_valid_tile = nxt_v & cutlass.Int32(1) - eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) - kv_left = bounds_next.left - kv_right = bounds_next.right + if cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + else: + eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) + kv_left = bounds_next.left + kv_right = bounds_next.right sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) @@ -1617,7 +1664,7 @@ def _softmax_warp_group( cutlass.Float32, ) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1635,7 +1682,7 @@ def _softmax_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) softmax_wg_base_const = CFG.SOFTMAX_WG0_BASE if sub_tile_id == 0 else CFG.SOFTMAX_WG1_BASE tid_in_wg = cute.arch.thread_idx()[0] - cutlass.Int32(softmax_wg_base_const * 32) @@ -1755,7 +1802,7 @@ def _softmax_warp_group( nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1770,7 +1817,7 @@ def _softmax_warp_group( sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) @cute.jit @@ -1824,7 +1871,7 @@ def _correction_warp_group( bmm2_done_phase = cutlass.Int32(0) o_empty_phase = cutlass.Int32(1) # bootstrap - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1842,7 +1889,7 @@ def _correction_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) @@ -1997,7 +2044,11 @@ def _correction_warp_group( else: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) - lse_row = lse_arr[batch_idx, head_idx, :] + # This chunk's LSE goes to its own split-major slot, matching + # where TMA-STG put the chunk's O. The pair (O_s, lse_s) is + # everything the combine needs. + lse_batch = _partial_batch(batch_idx, split_idx, n_batch) + lse_row = lse_arr[lse_batch, head_idx, :] lse_row[q_row_global] = lse_val sO_sub_base = sO[qs].base @@ -2050,7 +2101,7 @@ def _correction_warp_group( nxt_q = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0))).load() nxt_hb = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1))).load() nxt_v = (sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2))).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -2065,7 +2116,7 @@ def _correction_warp_group( sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) # End-of-warp tmem_dealloc: under cga2 each corr lane ALSO DSMEM-arrives # on the peer so the peer's local mbar accumulates the full CGA-total count. @@ -2191,7 +2242,12 @@ def _tma_swz(byte_w: int): grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1)) else: # Grid Python-folds on Cfg constant (avoids DSL if staging). - grid_shape = (grid_q_supers, QH, B) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B, 1, 1) + # KV split rides the BATCH axis: z = batch + split*B. The decode + # already recovers the batch coord on both the blockIdx and the + # scheduler-handout paths, so the split travels with it for free. + grid_shape = ( + (grid_q_supers, QH, B * SPLIT_KV) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B * SPLIT_KV, 1, 1) + ) _kernel( tma_q_desc, tma_k_desc, @@ -2258,6 +2314,10 @@ def compile( # noqa: A001 raise ValueError(f"d192 envelope: need 0 < d_qk <= {CFG.TILE_K} and 0 < d_v <= {CFG.TILE_O}; got ({d_qk}, {d_v})") if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d192 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") + if SPLIT_KV > 1 and not has_lse: + # Each split's LSE is not optional under KV split — it IS the weight + # the combine reduces with. Without it the partials cannot be recombined. + raise ValueError("split_kv > 1 requires has_lse=True (the per-split LSE drives the combine)") _fake_batch = 1 if CFG.THD_VARLEN else b if CFG.THD_VARLEN: # Dynamic packed token totals: one symbol per ragged group (Q/O and @@ -2265,6 +2325,10 @@ def compile( # noqa: A001 # compiled artifact instead of minting a new one (issue #552). sq = cute.sym_int(divisibility=1) skv = cute.sym_int(divisibility=1) + # KV split: O and LSE are the PARTIAL workspaces, stacked split-major on + # the batch axis (B*SPLIT_KV). Q/K/V keep the real batch. + _o_batch = _fake_batch * SPLIT_KV + _lse_batch = b * SPLIT_KV def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: @@ -2283,7 +2347,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) - fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=STORAGE_DTYPE) + fake_o = _fake_bshd((_o_batch, sq, qh, d_v), o_stride, dtype=STORAGE_DTYPE) if not has_lse: # No Stats output: the LSE argument is None-specialized and the store # is compiled out entirely — no dummy buffer exists at any level. @@ -2323,7 +2387,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): raise ValueError("lse_head_major / lse_head_stride are THD-only (dense LSE is compact (B, H, Sq))") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (b, qh, sq), + (_lse_batch, qh, sq), stride_order=(2, 1, 0), assumed_align=16, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py index 936990e32..a3ddef5a3 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -83,6 +83,7 @@ from cudnn.sdpa.fwd.kernels._common_sm100 import ( + make_split_helpers, D256Bars as Bars, KvLoopBounds, make_d256_bars, @@ -131,6 +132,29 @@ _thd_tma_offsets = _sdpa_h.thd_tma_offsets +# === KV split === +# +# Mechanics live in _common_sm100.make_split_helpers, shared with the other +# SM100 prefill flavors: each Q tile's KV loop range is cut into SPLIT_KV +# contiguous chunks, each run as its own persistent tile, and each writing a +# normalized partial O + its own LSE into a split-major workspace that +# split_combine_sm100 folds with the exact log-sum-exp identity. At +# SPLIT_KV == 1 every closure folds away and this is the classic kernel. +_split_h = make_split_helpers( + CFG, + bounds_for_tile=_bounds_for_tile, + dispatch_decode_initial=_dispatch_decode_initial, + dispatch_decode_payload=_dispatch_decode_payload, +) +SPLIT_KV = _split_h.SPLIT_KV +MAY_BE_EMPTY = _split_h.MAY_BE_EMPTY +_decode_initial_split = _split_h.decode_initial_split +_decode_payload_split = _split_h.decode_payload_split +_bounds_for_tile_split = _split_h.bounds_for_tile_split +_nomask_range_split = _split_h.nomask_range_split +_partial_batch = _split_h.partial_batch + + @dataclass(frozen=True) class KernelTmemLayout: TOTAL_COLS: int = 512 @@ -480,7 +504,7 @@ def _tmaldg_warp_group( tma_k = GmemTileTma(tma_k_desc) tma_v = GmemTileTma(tma_v_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -496,13 +520,15 @@ def _tmaldg_warp_group( q_row_base = cute.arch.make_warp_uniform(q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M)) q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) + elif cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -518,7 +544,7 @@ def _tmaldg_warp_group( bars.mb_q_o_alias.wait(q_o_alias_phase) q_o_alias_phase = q_o_alias_phase ^ cutlass.Int32(1) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): pass else: if cutlass.const_expr(CFG.CTA_MMA == 2): @@ -572,7 +598,7 @@ def _tmaldg_warp_group( nxt_q = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0)).load()) nxt_hb = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1)).load()) nxt_v = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load()) - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -588,10 +614,12 @@ def _tmaldg_warp_group( q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) - if cutlass.const_expr(CFG.MASK_FLAGS != 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV > 1): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -624,7 +652,7 @@ def _tmastg_warp_group( tma_o = GmemTileTma(tma_o_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -649,6 +677,10 @@ def _tmastg_warp_group( bars.mb_o_full[chunk].wait(o_full_phase) q_row_coord = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) + # KV split: partials are stacked split-major on the workspace BATCH + # axis (extent B*SPLIT_KV), so the store needs no new descriptor — + # only a shifted batch coord. Folds to batch_idx at SPLIT_KV == 1. + o_batch = _partial_batch(batch_idx, split_idx, n_batch) if cutlass.const_expr(CFG.THD_VARLEN): # DEAD unit (batch == n_batch, envelope grid — issue #552): no O @@ -661,7 +693,7 @@ def _tmastg_warp_group( else: tma_store_tile( sO[0], - tma_o(cutlass.Int32(0), head_idx, q_row_coord, batch_idx), + tma_o(cutlass.Int32(0), head_idx, q_row_coord, o_batch), ) tma_store_commit() tma_store_wait(0) @@ -676,7 +708,7 @@ def _tmastg_warp_group( nxt_q = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0)).load() nxt_hb = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1)).load() nxt_v = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -771,11 +803,11 @@ def _mma_warp_group( desc_Q = sQ[0].desc() - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) else: - q_super_idx, _hd, batch_idx = _dispatch_decode_initial( + q_super_idx, _hd, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -787,11 +819,14 @@ def _mma_warp_group( qh_per_kh, seqlen_kv, ) - eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) - kv_left = bounds_init.left - kv_right = bounds_init.right + if cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + else: + eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) + kv_left = bounds_init.left + kv_right = bounds_init.right q_full_phase = cutlass.Int32(0) kv_state_K = PipelineState.start(phase=0) @@ -805,7 +840,7 @@ def _mma_warp_group( while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): bars.mb_empty_mainloop.wait(empty_mainloop_phase) empty_mainloop_phase = empty_mainloop_phase ^ cutlass.Int32(1) elect_p = nvvm.elect_sync() @@ -913,14 +948,14 @@ def _mma_warp_group( nvvm.bar_warp_sync(cute.arch.FULL_MASK) wait(sched.mb_scheduler.subview(sched_state.idx), sched_state.phase) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): nxt_v = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load() is_valid_tile = nxt_v & cutlass.Int32(1) else: nxt_q = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0)).load()) nxt_hb = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1)).load()) nxt_v = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load()) - q_super_idx, _hd, batch_idx = _dispatch_decode_payload( + q_super_idx, _hd, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -932,11 +967,14 @@ def _mma_warp_group( seqlen_kv, ) is_valid_tile = nxt_v & cutlass.Int32(1) - eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) - eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) - kv_left = bounds_next.left - kv_right = bounds_next.right + if cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + else: + eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) + eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) + kv_left = bounds_next.left + kv_right = bounds_next.right sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) bars.mb_tmem_dealloc.wait(cutlass.Int32(0)) @@ -970,7 +1008,7 @@ def _softmax_warp_group( NEG_INF = cutlass.Float32(float("-inf")) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -988,7 +1026,7 @@ def _softmax_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) tid_in_wg = cute.arch.thread_idx()[0] - cutlass.Int32(CFG.SOFTMAX_WG0_BASE * 32) @@ -1338,7 +1376,7 @@ def _softmax_warp_group( nxt_q = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0)).load()) nxt_hb = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1)).load()) nxt_v = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load()) - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1353,7 +1391,7 @@ def _softmax_warp_group( sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) @cute.jit @@ -1386,7 +1424,7 @@ def _correction_warp_group( stat_mbar_state = cutlass.Int32(0) epilogue_state = cutlass.Int32(1) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1404,7 +1442,7 @@ def _correction_warp_group( eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) O_CHUNK = 16 N_CHUNKS_O = CFG.TILE_O // O_CHUNK @@ -1552,7 +1590,11 @@ def _correction_warp_group( else: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) - lse_row = lse_arr[batch_idx, head_idx, :] + # This chunk's LSE goes to its own split-major slot, matching where + # TMA-STG put the chunk's O. The pair (O_s, lse_s) is everything + # the combine needs. + lse_batch = _partial_batch(batch_idx, split_idx, n_batch) + lse_row = lse_arr[lse_batch, head_idx, :] lse_row[q_row_global] = lse_val parity_last_rt = cutlass.Int32(0) @@ -1602,7 +1644,7 @@ def _correction_warp_group( nxt_q = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0)).load() nxt_hb = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1)).load() nxt_v = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1617,7 +1659,7 @@ def _correction_warp_group( sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) if cutlass.const_expr(CFG.CTA_MMA == 2): peer_cta = cta_id_x ^ cutlass.Int32(1) @@ -1720,7 +1762,12 @@ def _tma_swz(byte_w: int): ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1)) else: - grid_shape = (grid_q_supers, QH, B) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B, 1, 1) + # KV split rides the BATCH axis: z = batch + split*B. The decode + # already recovers the batch coord on both the blockIdx and the + # scheduler-handout paths, so the split travels with it for free. + grid_shape = ( + (grid_q_supers, QH, B * SPLIT_KV) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B * SPLIT_KV, 1, 1) + ) _kernel( tma_q_desc, tma_k_desc, @@ -1779,6 +1826,10 @@ def compile( # noqa: A001 raise ValueError(f"d256 envelope: need 0 < d_qk <= {CFG.TILE_K} and 0 < d_v <= {CFG.TILE_O}; got ({d_qk}, {d_v})") if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d256 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") + if SPLIT_KV > 1 and not has_lse: + # Each split's LSE is not optional under KV split — it IS the weight + # the combine reduces with. Without it the partials cannot be recombined. + raise ValueError("split_kv > 1 requires has_lse=True (the per-split LSE drives the combine)") _fake_batch = 1 if CFG.THD_VARLEN else b if CFG.THD_VARLEN: # Dynamic packed token totals: one symbol per ragged group (Q/O and @@ -1786,6 +1837,10 @@ def compile( # noqa: A001 # compiled artifact instead of minting a new one (issue #552). sq = cute.sym_int(divisibility=1) skv = cute.sym_int(divisibility=1) + # KV split: O and LSE are the PARTIAL workspaces, stacked split-major on + # the batch axis (B*SPLIT_KV). Q/K/V keep the real batch. + _o_batch = _fake_batch * SPLIT_KV + _lse_batch = b * SPLIT_KV def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: @@ -1804,7 +1859,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) - fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE, bpe=CFG.BPE_O) + fake_o = _fake_bshd((_o_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE, bpe=CFG.BPE_O) if not has_lse: # No Stats output: the LSE argument is None-specialized and the store # is compiled out entirely — no dummy buffer exists at any level. @@ -1844,7 +1899,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): raise ValueError("lse_head_major / lse_head_stride are THD-only (dense LSE is compact (B, H, Sq))") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (b, qh, sq), + (_lse_batch, qh, sq), stride_order=(2, 1, 0), assumed_align=16, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py index 3eea02ac0..78c914d7d 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -91,6 +91,7 @@ def _require(cond, msg): ) from cudnn.sdpa.fwd.kernels._common_sm100 import ( + make_split_helpers, KvLoopBounds, compute_kv_loop_bounds, row_max_for_exp2, @@ -322,6 +323,29 @@ class KernelTmemLayout: _thd_tma_offsets = _sdpa_h.thd_tma_offsets +# === KV split === +# +# Mechanics live in _common_sm100.make_split_helpers, shared with the other +# SM100 prefill flavors: each Q tile's KV loop range is cut into SPLIT_KV +# contiguous chunks, each run as its own persistent tile, and each writing a +# normalized partial O + its own LSE into a split-major workspace that +# split_combine_sm100 folds with the exact log-sum-exp identity. At +# SPLIT_KV == 1 every closure folds away and this is the classic kernel. +_split_h = make_split_helpers( + CFG, + bounds_for_tile=_bounds_for_tile, + dispatch_decode_initial=_dispatch_decode_initial, + dispatch_decode_payload=_dispatch_decode_payload, +) +SPLIT_KV = _split_h.SPLIT_KV +MAY_BE_EMPTY = _split_h.MAY_BE_EMPTY +_decode_initial_split = _split_h.decode_initial_split +_decode_payload_split = _split_h.decode_payload_split +_bounds_for_tile_split = _split_h.bounds_for_tile_split +_nomask_range_split = _split_h.nomask_range_split +_partial_batch = _split_h.partial_batch + + @cute.kernel def _kernel( tma_q_desc: cutlass.GridConstant[tmap.TensorMap], @@ -807,7 +831,7 @@ def _compute_warp_group( wid_in_wg = tid_in_wg // cutlass.Int32(32) is_lead_warp = wid_in_wg == cutlass.Int32(0) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -822,16 +846,22 @@ def _compute_warp_group( is_valid_tile = cutlass.Int32(1) sched_state = PipelineState.start() - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_unmasked_lo = cutlass.Int32(0) kv_unmasked_hi = seqlen_kv // cutlass.Int32(CFG.TILE_N) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) eff_seqlen_kv = seqlen_kv + elif cutlass.const_expr(CFG.MASK_FLAGS == 0): + # Unmasked, so this split's whole slice is the unmasked band. + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + kv_unmasked_lo = kv_left + kv_unmasked_hi = kv_right + eff_seqlen_kv = seqlen_kv else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_unmasked_lo = bounds_init.unmasked_lo kv_unmasked_hi = bounds_init.unmasked_hi @@ -875,7 +905,7 @@ def _compute_warp_group( q_abs = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) + tid_in_wg - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): pass else: if cutlass.const_expr(CFG.MASK_FLAGS == 0): @@ -995,7 +1025,7 @@ def _compute_warp_group( else: tmem_O_base = tmem_base_addr + cutlass.Int32(LAYOUT.O_OFF) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): bars.mb_empty_mainloop.arrive_on_peer(leader_cta_id, pred=is_lead_warp & nvvm.elect_sync()) else: cur_parity_0 = alpha_full_state.idx @@ -1178,14 +1208,15 @@ def _compute_warp_group( else: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) - lse_row = lse_arr[batch_idx, head_idx, :] + lse_batch = _partial_batch(batch_idx, split_idx, n_batch) + lse_row = lse_arr[lse_batch, head_idx, :] lse_row[q_row_global] = lse wait(sched.mb_scheduler.subview(sched_state.idx), sched_state.phase) nxt_q = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0)).load()) nxt_hb = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1)).load()) nxt_v = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load()) - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1198,10 +1229,12 @@ def _compute_warp_group( ) is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) - if cutlass.const_expr(CFG.MASK_FLAGS != 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV > 1): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_unmasked_lo = bounds_next.unmasked_lo kv_unmasked_hi = bounds_next.unmasked_hi @@ -1293,7 +1326,7 @@ def _mma_warp_group( kind=MMA_KIND, ) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1308,13 +1341,15 @@ def _mma_warp_group( is_valid_tile = cutlass.Int32(1) sched_state = PipelineState.start() - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) + elif cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -1351,7 +1386,7 @@ def _mma_warp_group( read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) if is_sg0: - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): pass else: bars.mb_tma_q_full.wait(q_full_state.phase) @@ -1386,7 +1421,7 @@ def _mma_warp_group( bars.mb_tma_q_empty.arrive(cta_group=CFG.CTA_MMA, mcast_mask=mcast_mask, pred=nvvm.elect_sync()) else: - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): bars.mb_empty_mainloop.wait(empty_mainloop_state.phase) empty_mainloop_state = advance(empty_mainloop_state, 1) bars.mb_bmm2_done[bmm2_done_prod_state.idx].arrive(cta_group=CFG.CTA_MMA, mcast_mask=mcast_mask, pred=nvvm.elect_sync()) @@ -1447,7 +1482,7 @@ def _mma_warp_group( nxt_q = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0)).load() nxt_hb = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1)).load() nxt_v = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1460,10 +1495,12 @@ def _mma_warp_group( ) is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) - if cutlass.const_expr(CFG.MASK_FLAGS != 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV > 1): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -1503,7 +1540,7 @@ def _mma_warp_non_leader( _UTCCP_N_CALLS = LAYOUT.Q_TMEM_COLS // _UTCCP_TMEM_COLS_PER_CALL if is_sg1: - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1518,13 +1555,15 @@ def _mma_warp_non_leader( is_valid_tile = cutlass.Int32(1) sched_state = PipelineState.start() - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) + elif cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -1533,7 +1572,7 @@ def _mma_warp_non_leader( while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): pass else: for _kv in cutlass.range(kv_left, kv_right, 1, unroll=1): @@ -1550,7 +1589,7 @@ def _mma_warp_non_leader( nxt_q = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0)).load() nxt_hb = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1)).load() nxt_v = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1563,10 +1602,12 @@ def _mma_warp_non_leader( ) is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) - if cutlass.const_expr(CFG.MASK_FLAGS != 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV > 1): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right bars.mb_tmem_dealloc.wait(cutlass.Int32(0)) @@ -1606,7 +1647,7 @@ def _tmaldg_warp_group( tma_k = GmemTileTma(tma_k_desc) tma_v = GmemTileTma(tma_v_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1622,13 +1663,15 @@ def _tmaldg_warp_group( q_row_base = cute.arch.make_warp_uniform(q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M)) q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) - if cutlass.const_expr(CFG.MASK_FLAGS == 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV == 1): kv_left = cutlass.Int32(0) kv_right = seqlen_kv // cutlass.Int32(CFG.TILE_N) + elif cutlass.const_expr(CFG.MASK_FLAGS == 0): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) else: eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_init = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_init = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_init.left kv_right = bounds_init.right @@ -1641,7 +1684,7 @@ def _tmaldg_warp_group( while is_valid_tile > cutlass.Int32(0): read_tile_id_arrive(sched.mb_read_tile_id.subview(sched_state.idx), CGA_SIZE) - if cutlass.const_expr(CFG.MASK_FLAGS != 0) and (kv_right <= kv_left): + if cutlass.const_expr(MAY_BE_EMPTY) and (kv_right <= kv_left): # Empty KV loop (SWA window past the padded KV tail, or a dead-Q-tile # collapse): no loads — but the O∪V SMEM alias gate must stay in # phase. The sg1 compute epilogue stores (trimmed) O and TMA-STG @@ -1705,7 +1748,7 @@ def _tmaldg_warp_group( nxt_q = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0)).load()) nxt_hb = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1)).load()) nxt_v = cute.arch.make_warp_uniform(sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load()) - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1721,10 +1764,12 @@ def _tmaldg_warp_group( q_seq_off, kv_seq_off, tma_batch = _thd_tma_offsets(seq_kv_lens_tensor, batch_idx, n_batch) is_valid_tile = nxt_v & cutlass.Int32(1) sched_state = advance(sched_state, CFG.SCHEDULER_STAGES) - if cutlass.const_expr(CFG.MASK_FLAGS != 0): + if cutlass.const_expr(CFG.MASK_FLAGS == 0 and SPLIT_KV > 1): + kv_left, kv_right = _nomask_range_split(seqlen_kv, split_idx) + elif cutlass.const_expr(CFG.MASK_FLAGS != 0): eff_seqlen_kv = _resolve_seqlen_kv(seq_kv_lens_tensor, batch_idx, seqlen_kv) eff_seqlen_q = _resolve_seqlen_q(seq_kv_lens_tensor, batch_idx, seqlen_q, n_batch) - bounds_next = _bounds_for_tile(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx) + bounds_next = _bounds_for_tile_split(q_super_idx, eff_seqlen_q, eff_seqlen_kv, cta_in_pair, seq_q_lens_tensor, batch_idx, split_idx) kv_left = bounds_next.left kv_right = bounds_next.right @@ -1761,7 +1806,7 @@ def _tmastg_warp_group( tma_o = GmemTileTma(tma_o_desc) - q_super_idx, head_idx, batch_idx = _dispatch_decode_initial( + q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, sched.bidy_init, sched.bidz_init, @@ -1784,6 +1829,8 @@ def _tmastg_warp_group( bars.mb_tma_o_full[chunk].wait(o_full_state.phase) q_row_coord = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) + # KV split: partials stack split-major on the workspace BATCH axis. + o_batch = _partial_batch(batch_idx, split_idx, n_batch) if cutlass.const_expr(CFG.THD_VARLEN): # DEAD unit (batch == n_batch, envelope grid — issue #552): no O # rows exist and descriptor slot n_batch is never built, so skip @@ -1795,7 +1842,7 @@ def _tmastg_warp_group( else: tma_store_tile( sO[0], - tma_o(cutlass.Int32(0), head_idx, q_row_coord, batch_idx), + tma_o(cutlass.Int32(0), head_idx, q_row_coord, o_batch), ) tma_store_commit() tma_store_wait(0) @@ -1809,7 +1856,7 @@ def _tmastg_warp_group( nxt_q = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(0)).load() nxt_hb = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(1)).load() nxt_v = sched.tile_id_smem.subview(sched_state.idx * cutlass.Int32(8) + cutlass.Int32(2)).load() - q_super_idx, head_idx, batch_idx = _dispatch_decode_payload( + q_super_idx, head_idx, batch_idx, split_idx = _decode_payload_split( nxt_q, nxt_hb, cta_in_pair, @@ -1919,7 +1966,10 @@ def _tma_swz(byte_w: int): ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1)) else: - grid_shape = (grid_q_supers, QH, B) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B, 1, 1) + # KV split rides the BATCH axis: z = batch + split*B. + grid_shape = ( + (grid_q_supers, QH, B * SPLIT_KV) if cutlass.const_expr(CFG.SCHEDULER_POLICY == SCHED_NATURAL) else (grid_q_supers * QH * B * SPLIT_KV, 1, 1) + ) _kernel( tma_q_desc, tma_k_desc, @@ -1980,6 +2030,8 @@ def compile( # noqa: A001 raise ValueError(f"d512 envelope: need 0 < d_qk <= {CFG.TILE_K} and 0 < d_v <= {CFG.TILE_O}; got ({d_qk}, {d_v})") if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d512 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") + if SPLIT_KV > 1 and not has_lse: + raise ValueError("split_kv > 1 requires has_lse=True (the per-split LSE drives the combine)") _fake_batch = 1 if CFG.THD_VARLEN else b if CFG.THD_VARLEN: # Dynamic packed token totals: one symbol per ragged group (Q/O and @@ -1987,6 +2039,8 @@ def compile( # noqa: A001 # compiled artifact instead of minting a new one (issue #552). sq = cute.sym_int(divisibility=1) skv = cute.sym_int(divisibility=1) + _o_batch = _fake_batch * SPLIT_KV + _lse_batch = b * SPLIT_KV def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: @@ -2005,7 +2059,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) - fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE, bpe=CFG.BPE_O) + fake_o = _fake_bshd((_o_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE, bpe=CFG.BPE_O) if not has_lse: # No Stats output: the LSE argument is None-specialized and the store # is compiled out entirely — no dummy buffer exists at any level. @@ -2045,7 +2099,7 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): raise ValueError("lse_head_major / lse_head_stride are THD-only (dense LSE is compact (B, H, Sq))") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (b, qh, sq), + (_lse_batch, qh, sq), stride_order=(2, 1, 0), assumed_align=16, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py index 545648226..4a4016bf0 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py @@ -210,6 +210,7 @@ def __init__( seq_kv_lens_present: bool = False, has_sink: bool = False, thd_varlen: bool = False, + split_kv: int = 1, thd_batch: int = 1, thd_lse_head_major: bool = False, head_tile_qk: int = 128, @@ -282,6 +283,9 @@ def __init__( self.head_tile_v = head_tile_v self.q_tile = q_tile self.kv_tile = kv_tile + # KV split: SPLIT_KV CTAs per (q_tile, batch, head), each covering a + # contiguous slice of that tile's KV-tile range. 1 = off (folds away). + self.split_kv = split_kv # Warp roles if self.q_tile == 128: @@ -831,6 +835,18 @@ def kernel( """ tidx, _, _ = cute.arch.thread_idx() q_tile_idx, batch_idx, head_idx = cute.arch.block_idx() + # KV split rides the BATCH axis: grid.y = batch + split*B. Q/K/V and the + # per-batch seqlens must use the REAL batch; only the O/LSE partial slot + # uses the composite. Folds away entirely at split_kv == 1. NATURAL-only + # (config_sm120 enforces it): LPT / LPT_L2 flatten the grid to 1-D and + # derive the batch from the linear tile id, leaving no axis to ride. + split_idx = cutlass.Int32(0) + o_batch_idx = batch_idx + if cutlass.const_expr(self.split_kv > 1): + n_batch_real = cutlass.Int32(q.shape[0]) + split_idx = batch_idx // n_batch_real + batch_idx = batch_idx % n_batch_real + o_batch_idx = split_idx * n_batch_real + batch_idx if cutlass.const_expr(self.sched_policy != SCHED_NATURAL): _n_qh = cutlass.Int32(q.shape[2]) _n_batch = cutlass.Int32(self.thd_batch if cutlass.const_expr(self.thd_varlen) else q.shape[0]) @@ -850,6 +866,7 @@ def kernel( ) else: q_tile_idx, head_idx, batch_idx = lpt_tile_coords(q_tile_idx, _n_qh, _n_batch, _q_tiles) + o_batch_idx = batch_idx elif cutlass.const_expr(self.is_causal): # Diagonal-bounded work grows with the Q tile. Launch long tiles # first to avoid leaving a few expensive CTAs in the final @@ -908,7 +925,7 @@ def kernel( o_head_off = q_row_base * o_seq_stride + head_idx * o_head_stride else: q_head_off = batch_idx * q_batch_stride + head_idx * q_head_stride - o_head_off = batch_idx * o_batch_stride + head_idx * o_head_stride + o_head_off = o_batch_idx * o_batch_stride + head_idx * o_head_stride kv_head_idx = head_idx // (num_heads_q // num_heads_kv) num_kv_tiles = ceil_div(seqlen_k, self.kv_tile) @@ -934,6 +951,20 @@ def kernel( first_q_position += seqlen_k - seqlen_q first_valid_col = cute.math.max(cutlass.Int32(0), first_q_position - self.window_size_left) min_kv_tile = first_valid_col // self.kv_tile + if cutlass.const_expr(self.split_kv > 1): + # Cut the ALREADY-masked [min_kv_tile, num_kv_tiles) into SPLIT_KV + # near-equal chunks; the first `rem` splits take one extra tile so + # the slowest split (which sets the critical path) is minimal. A + # split past the end collapses to lo == hi, which drives has_kv_work + # false below -- the epilogue then produces row_sum = 0, i.e. + # O := 0 / LSE := -inf, the identity of the combine's log-sum-exp. + _span = num_kv_tiles - min_kv_tile + _per = _span // cutlass.Int32(self.split_kv) + _rem = _span % cutlass.Int32(self.split_kv) + _lo = min_kv_tile + split_idx * _per + cute.math.min(split_idx, _rem) + _extra = cutlass.Int32(1) if split_idx < _rem else cutlass.Int32(0) + min_kv_tile = _lo + num_kv_tiles = _lo + _per + _extra has_kv_work = num_kv_tiles > 0 and (num_kv_tiles - 1) >= min_kv_tile # Shared-memory layout: @@ -1161,7 +1192,13 @@ def kernel( ) kv_tile_idx -= 1 else: - if kv_tile_idx >= 0: + # Guard against the split's START, not 0. Without KV split + # min_kv_tile is 0 on this branch so the two agree, but a split + # begins partway into the KV range -- and the load warp bounds + # its loop by min_kv_tile, so testing >= 0 here makes the compute + # warps run one tile the loader never fetches and the CTA + # deadlocks on the TMA barrier. + if kv_tile_idx >= min_kv_tile: self.compute_one_kv_tile( basic_params, mma_params, @@ -1236,7 +1273,7 @@ def kernel( if lse_q_idx >= seqlen_q: lse_out = -cutlass.Float32.inf if lse_q_idx < q.shape[1]: - lse_row = lse_arr[batch_idx, head_idx, :] + lse_row = lse_arr[o_batch_idx, head_idx, :] lse_row[lse_q_idx] = lse_out prims.barrier_cta_sync(self.bar_compute_sync, thread_count=self.threads_compute) @@ -1372,10 +1409,13 @@ def __call__( def _static_neq(a, b): return isinstance(a, int) and isinstance(b, int) and a != b + # Under KV split, O is the split-major PARTIAL workspace: its batch mode + # is B*SPLIT_KV while Q/K/V keep the real batch, so O's batch is checked + # against that multiple rather than against Q's. if cutlass.const_expr( _static_neq(q.shape[0], k.shape[0]) or any(_static_neq(a, b) for a, b in zip(k.shape[:3], v.shape[:3])) - or _static_neq(q.shape[0], o.shape[0]) + or _static_neq(q.shape[0] * self.split_kv, o.shape[0]) or _static_neq(q.shape[1], o.shape[1]) or _static_neq(q.shape[2], o.shape[2]) or q.shape[2] % k.shape[2] != 0 @@ -1405,8 +1445,10 @@ def _static_neq(a, b): if cutlass.const_expr(lse.stride != (q.shape[2], 1)): raise ValueError("THD LSE must be compact token-major") else: - if cutlass.const_expr(lse.shape != (q.shape[0], q.shape[2], q.shape[1])): - raise ValueError("LSE must have shape (B, H, Sq)") + # Under KV split the LSE is the split-major partial workspace, + # batch mode B*SPLIT_KV (same as O). + if cutlass.const_expr(lse.shape != (q.shape[0] * self.split_kv, q.shape[2], q.shape[1])): + raise ValueError("LSE must have shape (B * split_kv, H, Sq)") if cutlass.const_expr(lse.stride != (q.shape[2] * q.shape[1], q.shape[1], 1)): raise ValueError("LSE must be compact row-major") if cutlass.const_expr(self.has_sink != (sinks is not None)): @@ -1485,7 +1527,9 @@ def kv_tma_desc(t, head_dim, head_tile, swizzle, swizzle_chunks, swizzle_chunk_e if cutlass.const_expr(self.sched_policy != SCHED_NATURAL): grid = (n_q_tiles * n_batch * n_head, 1, 1) else: - grid = (n_q_tiles, n_batch, n_head) + # KV split rides the batch axis: y = batch + split*B (config_sm120 + # allows split_kv > 1 only under NATURAL, whose 3-D grid has one). + grid = (n_q_tiles, n_batch * self.split_kv, n_head) self.kernel( q, k, @@ -1566,7 +1610,10 @@ def compile( # noqa: A001 head_tile_v=round_up_head_tile(d_v), q_tile=PARAMS.q_tile, kv_tile=PARAMS.kv_tile, + split_kv=PARAMS.split_kv, ) + if PARAMS.split_kv > 1 and not has_lse: + raise ValueError("SM120 SDPA: split_kv > 1 requires an LSE output (the per-split LSE drives the combine)") fake_batch = 1 if PARAMS.thd_varlen else b if PARAMS.thd_varlen: # Dynamic packed token totals: one symbol per ragged group (Q/O and @@ -1574,6 +1621,10 @@ def compile( # noqa: A001 # compiled artifact instead of minting a new one (issue #552). sq = cute.sym_int(divisibility=1) skv = cute.sym_int(divisibility=1) + # KV split: O and LSE are the PARTIAL workspaces, stacked split-major on the + # batch axis (B*SPLIT_KV). Q/K/V keep the real batch. + o_fake_batch = fake_batch * PARAMS.split_kv + lse_fake_batch = fake_batch * PARAMS.split_kv def _fake_bshd(shape, stride): if stride is None: @@ -1587,11 +1638,11 @@ def _fake_bshd(shape, stride): fake_q = _fake_bshd((fake_batch, sq, qh, d_qk), q_stride) fake_k = _fake_bshd((fake_batch, skv, kh, d_qk), k_stride) fake_v = _fake_bshd((fake_batch, skv, kh, d_v), v_stride) - fake_o = _fake_bshd((fake_batch, sq, qh, d_v), o_stride) + fake_o = _fake_bshd((o_fake_batch, sq, qh, d_v), o_stride) if PARAMS.thd_varlen: fake_lse_shape = (qh, lse_head_stride) if lse_head_major else (sq, qh) else: - fake_lse_shape = (fake_batch, qh, sq) + fake_lse_shape = (lse_fake_batch, qh, sq) fake_lse = ( cute.runtime.make_fake_compact_tensor( cutlass.Float32, diff --git a/python/cudnn/sdpa/fwd/kernels/split_combine_sm100.py b/python/cudnn/sdpa/fwd/kernels/split_combine_sm100.py new file mode 100644 index 000000000..48aa41800 --- /dev/null +++ b/python/cudnn/sdpa/fwd/kernels/split_combine_sm100.py @@ -0,0 +1,199 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""KV-split combine (reduction) pass for the SM100 DSL prefill SDPA kernels. + +Each split writes ``O_s`` (normalized by its own running sum) and ``lse_s`` into +a split-major workspace at batch coord ``b + s*B``; this pass reduces over ``s``. + +A split whose range came out empty ends with total_sum == 0, which the epilogue +turns into ``O := 0 / lse := -inf`` — the identity here, so empty splits need no +special case. + +One block per (q_row, head, batch); the block's threads stride over d_v, and +each thread walks the split axis in registers. +""" + +from typing import Callable, Optional, Tuple + +from functools import lru_cache + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import arith +from cutlass.base_dsl.typing import Pointer +from cutlass.experimental import primitives as nvvm +import cuda.bindings.driver as _cuda_driver # noqa: F401 (cute.compile pulls cuda) + +# One block per (q_row, head, batch); 128 lanes stride over d_v. d_v <= 128 for +# every flavor that uses this pass today, so the stride loop runs once. +THREADS = 128 + +NEG_INF = float("-inf") + + +@cute.kernel +def _combine_kernel( + o_partial: cute.Tensor, # [S*B, S_q, H, D] — split-major partial O + lse_partial: cute.Tensor, # [S*B, H, S_q] — split-major partial LSE + o_out: cute.Tensor, # [B, S_q, H, D] + lse_out: Optional[cute.Tensor], # [B, H, S_q] or None (None-specialized) + # 1-element fp32 amax of the RECOMBINED O. The per-split epilogues cannot + # compute this: each sees only its own partial, and O is a convex + # combination of those, so a max over partials over-reports (~2.9x at 8 + # splits). The FP8 kernels therefore skip their in-kernel amax write when + # SPLIT_KV > 1 and leave it to this pass. None-specialized off otherwise. + amax_o: Optional[cute.Tensor], + n_batch: cutlass.Int32, + n_splits: cutlass.Int32, + d_v: cutlass.Int32, +) -> None: + tidx, _, _ = cute.arch.thread_idx() + q_row = cute.arch.block_idx()[0] + head = cute.arch.block_idx()[1] + batch = cute.arch.block_idx()[2] + + op = cutlass.make_array_view(o_partial) + lp = cutlass.make_array_view(lse_partial) + oo = cutlass.make_array_view(o_out) + + # --- pass 1: M = max_s lse_s, then den = sum_s exp(lse_s - M) --- + # Every lane redundantly walks the (very short) split axis; the values are + # block-uniform and hit L1, which is cheaper than staging them through SMEM. + m = cutlass.Float32(NEG_INF) + for s in cutlass.range(0, n_splits, 1, unroll=1): + lse_row = lp[batch + s * n_batch, head, :] + m = cute.math.max(m, cutlass.Float32(lse_row[q_row])) + + # All splits dead (every row fully masked): emit O := 0 / lse := -inf rather + # than exp(-inf - -inf) == NaN. m_safe only feeds the exponentials. + all_dead = m == cutlass.Float32(NEG_INF) + m_safe = cutlass.Float32(arith.select(all_dead.ir_value(), cutlass.Float32(0.0).ir_value(), m.ir_value())) + + # Same reasoning as pass 2: skip dead splits rather than trusting a fastmath + # exp(-inf) to be exactly 0. + den = cutlass.Float32(0.0) + for s in cutlass.range(0, n_splits, 1, unroll=1): + lse_row = lp[batch + s * n_batch, head, :] + lse_s = cutlass.Float32(lse_row[q_row]) + if lse_s > cutlass.Float32(NEG_INF): + den = den + cute.math.exp(lse_s - m_safe, fastmath=True) + + inv_den = cutlass.Float32(1.0) / cute.math.max(den, cutlass.Float32(1e-30)) + inv_den = cutlass.Float32(arith.select(all_dead.ir_value(), cutlass.Float32(0.0).ir_value(), inv_den.ir_value())) + + # --- pass 2: O = sum_s w_s O_s / den, accumulated in fp32 --- + # + # A dead split (empty KV range) carries lse_s = -inf, so its weight is + # exp(-inf) == 0 and it should contribute nothing. Relying on the ARITHMETIC + # to erase it is not safe: 0 * x is NaN for a non-finite x, and under + # fastmath the weight itself is only approximately zero. Skip such splits + # outright -- they are the identity element of this reduction by + # construction, so branching is exact where multiplying is not. (Observed: + # d512 with 5 KV tiles over 8 splits produced NaN in the recombined O + # without this guard, even though every partial slot held a clean + # -inf / 0.) + neg_inf = cutlass.Float32(NEG_INF) + amax_local = cutlass.Float32(0.0) + for d0 in cutlass.range(tidx, d_v, THREADS, unroll=1): + acc = cutlass.Float32(0.0) + for s in cutlass.range(0, n_splits, 1, unroll=1): + lse_row = lp[batch + s * n_batch, head, :] + lse_s = cutlass.Float32(lse_row[q_row]) + if lse_s > neg_inf: + w = cute.math.exp(lse_s - m_safe, fastmath=True) + o_row = op[batch + s * n_batch, q_row, head, :] + acc = acc + w * cutlass.Float32(o_row[d0]) + out_row = oo[batch, q_row, head, :] + o_val = acc * inv_den + out_row[d0] = o_val.to(o_out.element_type) + if cutlass.const_expr(amax_o is not None): + # amax over the fp32 PRE-CAST value, matching what the single-pass + # epilogue reports. + amax_local = cute.math.max(amax_local, cute.math.max(o_val, -o_val)) + + # One atomic per lane. The value is non-negative, so its fp32 bit pattern + # orders the same as the float and an integer atomicMax is exact -- the same + # trick the kernels' own epilogues use. + if cutlass.const_expr(amax_o is not None): + _amax_ptr = Pointer(amax_o.iterator.raw_ptr(), dtype=cutlass.Int32) + nvvm.atomicrmw(nvvm.AtomicOp.MAX, _amax_ptr, amax_local.bitcast(cutlass.Int32)) + + # --- the recombined LSE (only when the caller asked for Stats) --- + if cutlass.const_expr(lse_out is not None): + if tidx == cutlass.Int32(0): + lo = cutlass.make_array_view(lse_out) + lse_val = m_safe + cute.math.log(cute.math.max(den, cutlass.Float32(1e-30)), fastmath=True) + lse_val = cutlass.Float32(arith.select(all_dead.ir_value(), cutlass.Float32(NEG_INF).ir_value(), lse_val.ir_value())) + lse_row_out = lo[batch, head, :] + lse_row_out[q_row] = lse_val + + +@cute.jit +def _host( + o_partial: cute.Tensor, + lse_partial: cute.Tensor, + o_out: cute.Tensor, + lse_out: Optional[cute.Tensor], + amax_o: Optional[cute.Tensor], + problem_size: Tuple[int, int, int, int], + n_splits: cutlass.Int32, + stream: _cuda_driver.CUstream = None, +) -> None: + B, H, SQ, D = problem_size + _combine_kernel( + o_partial, + lse_partial, + o_out, + lse_out, + amax_o, + cutlass.Int32(B), + n_splits, + cutlass.Int32(D), + ).launch( + grid=(SQ, H, B), + block=[THREADS, 1, 1], + stream=stream, + ) + + +@lru_cache(maxsize=None) +def compile( # noqa: A001 + b: int, + h: int, + sq: int, + d_v: int, + splits: int, + dtype_o: str = "f16", + has_lse: bool = False, + has_amax: bool = False, +) -> Callable: + """Compile the combine pass for one concrete (B, H, S_q, d_v, splits) shape. + + ``splits`` is baked into the workspace EXTENTS (batch dim ``splits*b``) but + passed to the kernel as a runtime count, so the split axis is a dynamic loop + — the pass is bandwidth-bound, so unrolling it buys nothing. ``has_lse`` + controls whether the recombined LSE is written at all; with ``False`` the + store is None-specialized out of the traced code. ``has_amax`` does the + same for the FP8-family amax of the recombined O. + """ + elem = {"f16": cutlass.Float16, "bf16": cutlass.BFloat16}[dtype_o] + + fake_o_partial = cute.runtime.make_fake_compact_tensor(elem, (splits * b, sq, h, d_v), stride_order=(3, 2, 1, 0), assumed_align=16) + fake_lse_partial = cute.runtime.make_fake_compact_tensor(cutlass.Float32, (splits * b, h, sq), stride_order=(2, 1, 0), assumed_align=16) + fake_o_out = cute.runtime.make_fake_compact_tensor(elem, (b, sq, h, d_v), stride_order=(3, 2, 1, 0), assumed_align=16) + fake_lse_out = cute.runtime.make_fake_compact_tensor(cutlass.Float32, (b, h, sq), stride_order=(2, 1, 0), assumed_align=16) if has_lse else None + fake_amax_o = cute.runtime.make_fake_compact_tensor(cutlass.Float32, (1,), stride_order=(0,), assumed_align=4) if has_amax else None + + return cute.compile( + _host, + fake_o_partial, + fake_lse_partial, + fake_o_out, + fake_lse_out, + fake_amax_o, + (b, h, sq, d_v), + cutlass.Int32(0), + stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False), + options="--enable-tvm-ffi", + ) 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 new file mode 100644 index 000000000..62fad4b62 --- /dev/null +++ b/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py @@ -0,0 +1,814 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""KV-split tests for the SM100 d128 f16/bf16 prefill kernel. + +The split must change NOTHING numerically: the recombined output has to match +both a torch fp32 reference and the unsplit kernel to fp16 rounding. + +Coverage: split counts that divide the KV tile count evenly and ones that do not +(including more splits than tiles, which leaves some splits empty — the case +that exercises the empty-mainloop / LSE = -inf identity path), dense and causal +masks, MHA and GQA. +""" + +import math +import os + +import pytest +import torch + +from frost_test_utils import requires_blackwell, requires_dsl + +pytestmark = [requires_blackwell, requires_dsl] + +D = 128 +TILE_N = 128 +# One cga2 cluster covers TILES_Q * TILE_M * CTA_MMA Q rows. +ROWS_PER_CLUSTER = 2 * 128 * 2 + + +def _ref_sdpa(q, k, v, scale, is_causal, kh): + """fp32 reference over BSHD inputs; returns BSHD.""" + if kh != q.shape[2]: + rep = q.shape[2] // kh + k = k.repeat_interleave(rep, dim=2) + v = v.repeat_interleave(rep, dim=2) + qb, kb, vb = (t.float().permute(0, 2, 1, 3) for t in (q, k, v)) + s = torch.matmul(qb, kb.transpose(-1, -2)) * scale + if is_causal: + s_q, s_kv = qb.shape[2], kb.shape[2] + i = torch.arange(s_q, device=q.device).view(s_q, 1) + j = torch.arange(s_kv, device=q.device).view(1, s_kv) + s = s.masked_fill(j > i, float("-inf")) # top-left causal + p = torch.softmax(s, dim=-1) + return torch.matmul(p, vb).permute(0, 2, 1, 3) + + +def _kernel_module(splits, dtype_qkv, causal, cta_mma=2): + from cudnn.frost.template_loader import load_template + from cudnn.sdpa.fwd import api_dsl + from cudnn.sdpa.fwd.config_sm100 import TemplateParams + + path = os.path.join(os.path.dirname(os.path.abspath(api_dsl.__file__)), "kernels", "prefill_d128_f16_sm100.py") + kw = {"dtype_qkv": dtype_qkv, "split_kv": splits, "cta_mma": cta_mma} + if causal: + kw["window_right"] = 0 + return load_template(path, TemplateParams(**kw), tag=f"splitkv{splits}_d{dtype_qkv}_{'caus' if causal else 'dense'}_cga{cta_mma}") + + +def _run(splits, B, H, KH, SQ, SKV, dtype, causal, cta_mma=2): + """Launch the split kernel + DSL combine; returns the recombined O (BSHD fp32).""" + import cutlass + import cuda.bindings.driver as cuda_driver + + from cudnn.sdpa.fwd.kernels import split_combine_sm100 as comb + + dev = "cuda" + scale = 1.0 / math.sqrt(D) + torch.manual_seed(0) + q = torch.randn(B, SQ, H, D, device=dev, dtype=dtype) + k = torch.randn(B, SKV, KH, D, device=dev, dtype=dtype) + v = torch.randn(B, SKV, KH, D, device=dev, dtype=dtype) + + mod = _kernel_module(splits, 3 if dtype == torch.float16 else 2, causal, cta_mma=cta_mma) + fn = mod.compile(b=B, qh=H, kh=KH, sq=SQ, skv=SKV, d_qk=D, d_v=D, has_lse=True) + + o_p = torch.zeros(splits * B, SQ, H, D, device=dev, dtype=dtype) + lse_p = torch.zeros(splits * B, H, SQ, device=dev, dtype=torch.float32) + stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream) + fn( + q, + k, + v, + o_p, + lse_p, + torch.zeros(H, dtype=torch.float32, device=dev), # sinks (ABI slot) + torch.zeros(B, dtype=torch.int32, device=dev), # seq_kv_lens (unused) + torch.zeros(1, dtype=torch.int64, device=dev), # o_desc (THD only) + (B, H, KH, SQ, SKV, 0), + cutlass.Float32(scale * math.log2(math.e)), + cutlass.Int32(0), + None, + stream=stream, + ) + if splits == 1: + torch.cuda.synchronize() + return o_p.float(), (q, k, v, scale) + + o_out = torch.zeros(B, SQ, H, D, device=dev, dtype=dtype) + lse_out = torch.zeros(B, H, SQ, device=dev, dtype=torch.float32) + cfn = comb.compile(b=B, h=H, sq=SQ, d_v=D, splits=splits, dtype_o="f16" if dtype == torch.float16 else "bf16", has_lse=True) + cfn(o_p, lse_p, o_out, lse_out, None, (B, H, SQ, D), cutlass.Int32(splits), stream=stream) + torch.cuda.synchronize() + assert not torch.isnan(o_out).any(), "NaN in combined O" + return o_out.float(), (q, k, v, scale) + + +@pytest.mark.L0 +@pytest.mark.parametrize("splits", [2, 4, 8], ids=lambda s: f"split{s}") +def test_split_kv_matches_reference_dense(splits): + """The target shape class: tiny S_q against a long KV run.""" + B, H, SQ, SKV = 1, 4, 128, 2048 + got, (q, k, v, scale) = _run(splits, B, H, H, SQ, SKV, torch.float16, causal=False) + ref = _ref_sdpa(q, k, v, scale, is_causal=False, kh=H) + assert (got - ref).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +@pytest.mark.parametrize("splits", [3, 4, 8], ids=lambda s: f"split{s}") +def test_split_kv_uneven_and_empty_splits(splits): + """5 KV tiles over 3/4/8 splits: unequal chunks, and at 8 three EMPTY splits. + + An empty split ends with total_sum == 0, so its epilogue writes O := 0 / + LSE := -inf — the identity of the combine's log-sum-exp. Getting this wrong + shows up as NaN or as a wrong normalization, not as a small drift. + """ + B, H, SQ, SKV = 1, 4, 128, 5 * TILE_N + assert SKV // TILE_N == 5 + got, (q, k, v, scale) = _run(splits, B, H, H, SQ, SKV, torch.float16, causal=False) + ref = _ref_sdpa(q, k, v, scale, is_causal=False, kh=H) + assert (got - ref).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +@pytest.mark.parametrize("splits", [2, 4], ids=lambda s: f"split{s}") +def test_split_kv_causal(splits): + """Causal: the split must cut the MASKED range, so each split gets real work.""" + B, H, SQ, SKV = 1, 4, 1024, 1024 + got, (q, k, v, scale) = _run(splits, B, H, H, SQ, SKV, torch.float16, causal=True) + ref = _ref_sdpa(q, k, v, scale, is_causal=True, kh=H) + assert (got - ref).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +@pytest.mark.parametrize("splits", [4], ids=lambda s: f"split{s}") +def test_split_kv_gqa(splits): + B, H, KH, SQ, SKV = 2, 8, 2, 128, 2048 + got, (q, k, v, scale) = _run(splits, B, H, KH, SQ, SKV, torch.float16, causal=False) + ref = _ref_sdpa(q, k, v, scale, is_causal=False, kh=KH) + assert (got - ref).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +def test_split_kv_matches_unsplit(): + """Split vs unsplit on identical inputs — the split is occupancy-only.""" + B, H, SQ, SKV = 1, 4, 128, 2048 + base, _ = _run(1, B, H, H, SQ, SKV, torch.float16, causal=False) + for splits in (2, 4, 8): + got, _ = _run(splits, B, H, H, SQ, SKV, torch.float16, causal=False) + # Both round to fp16; the split reassociates the sum, so allow one ulp + # of drift at fp16 magnitudes rather than demanding bit-equality. + assert (got - base).abs().max().item() <= 2e-2, f"split{splits} diverged from unsplit" + + +@pytest.mark.L0 +def test_split_kv_bf16(): + B, H, SQ, SKV = 1, 4, 128, 2048 + got, (q, k, v, scale) = _run(4, B, H, H, SQ, SKV, torch.bfloat16, causal=False) + ref = _ref_sdpa(q, k, v, scale, is_causal=False, kh=H) + assert (got - ref).abs().max().item() <= 1e-1 # bf16 has ~8 mantissa bits + + +@pytest.mark.L0 +def test_split_kv_rejects_unsupported_combos(): + """The config validator is the backstop for what the combine cannot express.""" + from cudnn.frost.tile_dsl.scheduler import SCHED_LPT + from cudnn.sdpa.fwd.config_sm100 import TemplateParams, make_cfg_d128 + + with pytest.raises(ValueError, match="split_kv"): + make_cfg_d128(TemplateParams(split_kv=0)) + # Both scheduler policies are supported: NATURAL carries the split on the + # batch axis, LPT on the flattened x axis (make_split_helpers._lpt_split_of). + assert make_cfg_d128(TemplateParams(split_kv=4, sched_policy=SCHED_LPT))[0].SPLIT_KV == 4 + with pytest.raises(ValueError, match="sink"): + make_cfg_d128(TemplateParams(split_kv=4, has_sink=True)) + with pytest.raises(ValueError, match="dense-only"): + make_cfg_d128(TemplateParams(split_kv=4, thd_varlen=True, seq_kv_lens_present=True)) + + +@pytest.mark.L0 +def test_split_kv_requires_lse(): + """has_lse=False + split is rejected: the per-split LSE IS the combine weight.""" + mod = _kernel_module(4, 3, causal=False) + with pytest.raises(ValueError, match="has_lse"): + mod.compile(b=1, qh=4, kh=4, sq=128, skv=2048, d_qk=D, d_v=D, has_lse=False) + + +# --- cga1 (CTA_MMA=1) --------------------------------------------------- +# +# cga1 drops the collective 2-CTA MMA: one independent CTA per tile, covering +# TILES_Q*TILE_M = 256 Q rows instead of 512. It halves both the wasted MMA +# work at small S_q and the CTAs per tile, so with KV split twice as many splits +# fit in one wave. It is SMEM-neutral only because make_cfg_d128 turns QO_ALIAS +# on for cga1 (no collective MMA to halve per-CTA K/V), which is the part most +# likely to break -- hence the direct cga1-vs-cga2 output comparison below. + + +@pytest.mark.L0 +@pytest.mark.parametrize("splits", [1, 4, 8], ids=lambda s: f"split{s}") +def test_cga1_matches_reference(splits): + B, H, SQ, SKV = 1, 4, 128, 2048 + got, (q, k, v, scale) = _run(splits, B, H, H, SQ, SKV, torch.float16, causal=False, cta_mma=1) + ref = _ref_sdpa(q, k, v, scale, is_causal=False, kh=H) + assert (got - ref).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +def test_cga1_matches_cga2(): + """Same inputs through both cluster widths — cga1 is an occupancy change only.""" + B, H, SQ, SKV = 1, 4, 128, 2048 + for splits in (1, 4): + a, _ = _run(splits, B, H, H, SQ, SKV, torch.float16, causal=False, cta_mma=2) + b, _ = _run(splits, B, H, H, SQ, SKV, torch.float16, causal=False, cta_mma=1) + assert (a - b).abs().max().item() <= 2e-2, f"cga1 diverged from cga2 at split{splits}" + + +@pytest.mark.L0 +def test_cga1_causal_and_gqa(): + B, H, KH, SQ, SKV = 1, 8, 2, 1024, 1024 + got, (q, k, v, scale) = _run(4, B, H, KH, SQ, SKV, torch.float16, causal=True, cta_mma=1) + ref = _ref_sdpa(q, k, v, scale, is_causal=True, kh=KH) + assert (got - ref).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +def test_cga1_requires_qo_alias_and_smem_fits(): + """cga1 must enable QO_ALIAS, and both widths must stay inside the SMEM cap.""" + from dataclasses import replace + + from cudnn.sdpa.fwd.config_sm100 import TemplateParams, make_cfg_d128, _validate_cfg_d128, _d128_smem_bytes, _SM100_MAX_DYN_SMEM + + cfg1, _ = make_cfg_d128(TemplateParams(cta_mma=1)) + cfg2, _ = make_cfg_d128(TemplateParams(cta_mma=2)) + assert cfg1.QO_ALIAS == 1 and cfg2.QO_ALIAS == 0 + assert cfg1.CGA_M == cfg1.CTA_MMA == 1 and cfg2.CGA_M == cfg2.CTA_MMA == 2 + # Both fit, and cga1 only fits *because* of the alias. + assert _d128_smem_bytes(cfg1) <= _SM100_MAX_DYN_SMEM + assert _d128_smem_bytes(cfg2) <= _SM100_MAX_DYN_SMEM + assert _d128_smem_bytes(replace(cfg1, QO_ALIAS=0)) > _SM100_MAX_DYN_SMEM + with pytest.raises(ValueError, match="QO_ALIAS is mandatory"): + _validate_cfg_d128(replace(cfg1, QO_ALIAS=0)) + with pytest.raises(ValueError, match="cta_mma"): + make_cfg_d128(TemplateParams(cta_mma=3)) + + +@pytest.mark.L0 +def test_split_kv_and_cta_mma_flavor_gating(): + """Flavors must REJECT knobs they do not honour, not ignore them. + + ``split_kv`` / ``cta_mma`` sit on the TemplateParams shared by every SM100 + flavor, but a flavor only honours them once its make_cfg_* threads them into + a Cfg AND its kernel reads them. Silently ignoring split_kv is a + wrong-answer bug, not a no-op: the caller sizes an (S*B)-batch partial + workspace and runs the combine while the kernel writes only slots [0, B), + leaving the rest at lse_partial = 0 instead of -inf — weight exp(0 - M) != 0, + so they corrupt the reduction rather than dropping out. + """ + from cudnn.sdpa.fwd.config_sm100 import ( + _CTA_MMA_FLAVORS, + _SPLIT_KV_FLAVORS, + TemplateParams, + make_cfg_d128, + make_cfg_d192, + make_cfg_d256, + make_cfg_d512, + ) + + mk = {"d128": make_cfg_d128, "d192": make_cfg_d192, "d256": make_cfg_d256, "d512": make_cfg_d512} + for name, f in mk.items(): + f(TemplateParams()) # defaults must always build + if name in _SPLIT_KV_FLAVORS: + assert f(TemplateParams(split_kv=4))[0].SPLIT_KV == 4 + else: + with pytest.raises(ValueError, match="split_kv is not implemented"): + f(TemplateParams(split_kv=4)) + if name in _CTA_MMA_FLAVORS: + assert f(TemplateParams(cta_mma=1))[0].CTA_MMA == 1 + else: + with pytest.raises(ValueError, match="cta_mma is not selectable"): + f(TemplateParams(cta_mma=1)) + + +# --- empty-split coverage across every f16 flavor ------------------------- +# +# More splits than KV tiles leaves some splits with an EMPTY range. That path is +# what deadlocked d128 during bring-up: the empty-tile handshake +# (mb_empty_mainloop) was compile-time gated on MASK_FLAGS != 0, so at MASK_NONE +# the producers ran a full prologue while correction took the empty path and the +# kernel hung. d192 additionally has its own CAN_HAVE_EMPTY_KV predicate gating +# nine more sites. Every flavor therefore needs this exercised, not just d128 -- +# a regression here shows up as a HANG, not a wrong number. + +_F16_FLAVORS = { + "d128": ("prefill_d128_f16_sm100.py", 128, 128), + "d192": ("prefill_d192_d128_f16_sm100.py", 192, 128), + "d256": ("prefill_d256_f16_sm100.py", 256, 256), + "d512": ("prefill_d512_f16_sm100.py", 512, 512), +} + + +@pytest.mark.L0 +@pytest.mark.parametrize("flavor", sorted(_F16_FLAVORS)) +def test_empty_splits_every_flavor(flavor): + """5 KV tiles over 8 splits -> 3 empty splits, on every f16 flavor.""" + import math as _math + import os as _os + + import cutlass + import cuda.bindings.driver as cuda_driver + + from cudnn.frost.template_loader import load_template + from cudnn.sdpa.fwd import api_dsl + from cudnn.sdpa.fwd.config_sm100 import TemplateParams + from cudnn.sdpa.fwd.kernels import split_combine_sm100 as comb + + kmod, d_qk, d_v = _F16_FLAVORS[flavor] + B, H, SQ, SKV, S = 1, 4, 128, 5 * TILE_N, 8 + assert SKV // TILE_N < S, "this test must leave some splits empty" + dev = "cuda" + scale = 1.0 / _math.sqrt(d_qk) + torch.manual_seed(0) + + q = torch.randn(B, SQ, H, d_qk, device=dev, dtype=torch.float16) + k = torch.randn(B, SKV, H, d_qk, device=dev, dtype=torch.float16) + v = torch.randn(B, SKV, H, d_v, device=dev, dtype=torch.float16) + + path = _os.path.join(_os.path.dirname(_os.path.abspath(api_dsl.__file__)), "kernels", kmod) + mod = load_template(path, TemplateParams(dtype_qkv=3, split_kv=S), tag=f"empty_{flavor}_{S}") + fn = mod.compile(b=B, qh=H, kh=H, sq=SQ, skv=SKV, d_qk=d_qk, d_v=d_v, has_lse=True) + + o_p = torch.zeros(S * B, SQ, H, d_v, device=dev, dtype=torch.float16) + lse_p = torch.zeros(S * B, H, SQ, device=dev, dtype=torch.float32) + stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream) + fn( + q, + k, + v, + o_p, + lse_p, + torch.zeros(H, dtype=torch.float32, device=dev), + torch.zeros(B, dtype=torch.int32, device=dev), + torch.zeros(1, dtype=torch.int64, device=dev), + (B, H, H, SQ, SKV, 0), + cutlass.Float32(scale * _math.log2(_math.e)), + cutlass.Int32(0), + None, + stream=stream, + ) + o_out = torch.zeros(B, SQ, H, d_v, device=dev, dtype=torch.float16) + cfn = comb.compile(b=B, h=H, sq=SQ, d_v=d_v, splits=S, dtype_o="f16", has_lse=False) + cfn(o_p, lse_p, o_out, None, None, (B, H, SQ, d_v), cutlass.Int32(S), stream=stream) + torch.cuda.synchronize() + + assert not torch.isnan(o_out).any(), f"{flavor}: NaN from an empty split" + ref = _ref_sdpa(q, k, v, scale, is_causal=False, kh=H) + assert (o_out.float() - ref).abs().max().item() <= 2e-2 + + +# --- cga1 dtype coverage -------------------------------------------------- +# +# The d128 flavor's config gate admits cta_mma=1 for every dtype, because +# f16/bf16, fp8 and mxfp8 all share make_cfg_d128. That is correct for +# f16/bf16/fp8 but NOT for mxfp8: its per-32-block E8M0 scale factors are staged +# in SMEM by the mxfp8 kernel alone, which config_sm100._d128_smem_bytes cannot +# see, and they push a cga1 CTA to 237024 B against the 232448 B cap. Without +# the kernel-local guard the launch is rejected by the driver at runtime. + + +@pytest.mark.L0 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) +def test_cga1_half_dtypes(dtype): + B, H, SQ, SKV = 1, 4, 128, 2048 + got, (q, k, v, scale) = _run(4, B, H, H, SQ, SKV, dtype, causal=False, cta_mma=1) + ref = _ref_sdpa(q, k, v, scale, is_causal=False, kh=H) + tol = 2e-2 if dtype == torch.float16 else 1e-1 + assert (got - ref).abs().max().item() <= tol + + +@pytest.mark.L0 +def test_cga1_stage_depth_scales_with_cluster_width(): + """cga1 halves the KV stage depth for the fp8 family, as cuDNN's kernels do. + + cuDNN scales the stage count with the cluster width (stages_kv = N * CTA_MMA) + so that stages x per-CTA-buffer -- and hence SMEM -- stays constant. FROST + needs the same for fp8/mxfp8: at STAGES_KV=4 a cga1 CTA asked for 237024 B + against the 232448 B sm_100a cap, because mxfp8 also stages E8M0 scale + factors that _d128_smem_bytes cannot see. f16/bf16 already fit at cga1 via + the Q/O alias and keep STAGES_KV=2. + """ + from cudnn.frost.tile_dsl.constants import DTYPE_E4M3, DTYPE_FP16 + from cudnn.sdpa.fwd.config_sm100 import TemplateParams, make_cfg_d128 + + for cta_mma, want_fp8_stages in ((2, 4), (1, 2)): + cfg_fp8, _ = make_cfg_d128(TemplateParams(dtype_qkv=DTYPE_E4M3, dtype_o=DTYPE_FP16, cta_mma=cta_mma)) + assert cfg_fp8.STAGES_KV == want_fp8_stages, (cta_mma, cfg_fp8.STAGES_KV) + cfg_f16, _ = make_cfg_d128(TemplateParams(dtype_qkv=DTYPE_FP16, cta_mma=cta_mma)) + assert cfg_f16.STAGES_KV == 2, (cta_mma, cfg_f16.STAGES_KV) + # cga1 must also turn the Q/O alias on, which is what pays for the doubled + # per-CTA K/V in the first place. + assert make_cfg_d128(TemplateParams(dtype_qkv=DTYPE_FP16, cta_mma=1))[0].QO_ALIAS == 1 + + +# --- fp8 / mxfp8 split coverage ------------------------------------------ +# +# These two flavors share make_cfg_d128 with f16/bf16 but have their own ABIs +# (amax outputs for fp8; per-32-block E8M0 scale factors for mxfp8), so they are +# not reachable through _run above. Both cluster widths are exercised: mxfp8 at +# cga1 in particular only fits once STAGES_KV halves with the cluster width. + + +def _fp8_family_split(kfile, dtype_qkv, splits, cta_mma, mx): + import math as _math + import os as _os + + import cutlass + import cuda.bindings.driver as cuda_driver + + from cudnn.frost.template_loader import load_template + from cudnn.frost.tile_dsl.constants import DTYPE_FP16 + from cudnn.sdpa.fwd import api_dsl + from cudnn.sdpa.fwd.config_sm100 import TemplateParams + from cudnn.sdpa.fwd.kernels import split_combine_sm100 as comb + + B, H, SQ, SKV, D = 1, 4, 128, 2048, 128 + dev = "cuda" + scale = 1.0 / _math.sqrt(D) + torch.manual_seed(0) + kdir = _os.path.join(_os.path.dirname(_os.path.abspath(api_dsl.__file__)), "kernels") + params = TemplateParams(dtype_qkv=dtype_qkv, dtype_o=DTYPE_FP16, split_kv=splits, cta_mma=cta_mma) + mod = load_template(_os.path.join(kdir, kfile), params, tag=f"t_{kfile[:16]}_{splits}_{cta_mma}") + fn = mod.compile(b=B, qh=H, kh=H, sq=SQ, skv=SKV, has_lse=True) + stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream) + o_p = torch.zeros(splits * B, SQ, H, D, device=dev, dtype=torch.float16) + lse_p = torch.zeros(splits * B, H, SQ, device=dev, dtype=torch.float32) + amax_o = torch.zeros(1, dtype=torch.float32, device=dev) + zH = torch.zeros(H, dtype=torch.float32, device=dev) + zB = torch.zeros(B, dtype=torch.int32, device=dev) + ps = (B, H, H, SQ, SKV, 0) + log2e = cutlass.Float32(scale * _math.log2(_math.e)) + + if not mx: + mk = lambda *sh: (torch.randn(*sh, device=dev) * 0.5).clamp(-448, 448).to(torch.float8_e4m3fn) + q, k, v = mk(B, SQ, H, D), mk(B, SKV, H, D), mk(B, SKV, H, D) + # The FP8 entry takes four 1-element fp32 DEVICE scale tensors + # (descale_q/k/v, scale_o) — the scales fold in-kernel — and no Amax_S. + one = lambda: torch.ones(1, dtype=torch.float32, device=dev) + fn(q, k, v, o_p, lse_p, zH, zB, ps, log2e, cutlass.Float32(1.0), one(), one(), one(), one(), amax_o, stream=stream) + qf, kf, vf = (t.float().permute(0, 2, 1, 3) for t in (q, k, v)) + else: + from sdpa.mxfp8_quant import quantize_to_mxfp8 + + qr = torch.randn(B, H, SQ, D, device=dev) * 0.5 + kr = torch.randn(B, H, SKV, D, device=dev) * 0.5 + vr = torch.randn(B, H, SKV, D, device=dev) * 0.5 + a, adq, aswz, b_, bdq, bswz = quantize_to_mxfp8(qr, B, H, SQ, D) + q8, sfq, qf = a, aswz, adq.reshape(B, H, SQ, D).float() + a, adq, aswz, b_, bdq, bswz = quantize_to_mxfp8(kr, B, H, SKV, D) + k8, sfk, kf = a, aswz, adq.reshape(B, H, SKV, D).float() + a, adq, aswz, b_, bdq, bswz = quantize_to_mxfp8(vr, B, H, SKV, D) + v8, sfv, vf = b_, bswz, bdq.reshape(B, H, SKV, D).float() + sfq, sfk, sfv = (t.reshape(B, H, -1, 512).view(torch.int8).contiguous() for t in (sfq, sfk, sfv)) + q8 = q8.reshape(B, H, SQ, D).permute(0, 2, 1, 3).contiguous() + k8 = k8.reshape(B, H, SKV, D).permute(0, 2, 1, 3).contiguous() + v8 = v8.reshape(B, H, SKV, D).permute(0, 2, 1, 3).contiguous() + fn(q8, k8, v8, o_p, sfq, sfk, sfv, lse_p, amax_o, zH, zB, ps, log2e, stream=stream) + + ref = torch.matmul(torch.softmax(torch.matmul(qf, kf.transpose(-1, -2)) * scale, -1), vf).permute(0, 2, 1, 3) + if splits == 1: + torch.cuda.synchronize() + return o_p.float(), ref, amax_o + o_out = torch.zeros(B, SQ, H, D, device=dev, dtype=torch.float16) + # has_amax: at splits > 1 the per-split epilogues skip their amax write, so + # the combine is what reports it -- over the RECOMBINED O. + cfn = comb.compile(b=B, h=H, sq=SQ, d_v=D, splits=splits, dtype_o="f16", has_lse=False, has_amax=True) + cfn(o_p, lse_p, o_out, None, amax_o, (B, H, SQ, D), cutlass.Int32(splits), stream=stream) + torch.cuda.synchronize() + assert not torch.isnan(o_out).any(), "NaN in combined fp8-family O" + return o_out.float(), ref, amax_o + + +def _assert_amax_is_of_the_output(amax_o, got): + """amax_o must describe the OUTPUT the caller receives, at any split count. + + The per-split epilogues each see only their own partial, and the recombined + O is a convex combination of those, so |O| <= max_s |O_s|: a max taken over + partials silently over-reports (measured 1.5x at 2 splits, 2.9x at 8) and + would hand the caller a far-too-loose quantization scale. Comparing against + the recombined tensor is what catches that; comparing against the partials + would pass either way. + """ + reported = amax_o.item() + true_amax = got.abs().max().item() + # The kernel takes its amax on the fp32 value before the half store, so the + # readback can differ by one rounding step in that direction only. + assert reported >= true_amax * 0.99, f"amax_o {reported} under-reports |O| {true_amax}" + assert reported <= true_amax * 1.01, f"amax_o {reported} over-reports |O| {true_amax} — computed over partials?" + + +@pytest.mark.L0 +@pytest.mark.parametrize("cta_mma", [2, 1], ids=["cga2", "cga1"]) +@pytest.mark.parametrize("splits", [1, 4], ids=lambda s: f"split{s}") +def test_split_kv_fp8(splits, cta_mma): + from cudnn.frost.tile_dsl.constants import DTYPE_E4M3 + + got, ref, amax_o = _fp8_family_split("prefill_d128_fp8_sm100.py", DTYPE_E4M3, splits, cta_mma, mx=False) + assert (got - ref).abs().max().item() <= 5e-2 + _assert_amax_is_of_the_output(amax_o, got) + + +@pytest.mark.L0 +@pytest.mark.parametrize("cta_mma", [2, 1], ids=["cga2", "cga1"]) +@pytest.mark.parametrize("splits", [1, 4], ids=lambda s: f"split{s}") +def test_split_kv_mxfp8(splits, cta_mma): + from cudnn.frost.tile_dsl.constants import DTYPE_E4M3 + + got, ref, amax_o = _fp8_family_split("prefill_d128_mxfp8_sm100.py", DTYPE_E4M3, splits, cta_mma, mx=True) + assert (got - ref).abs().max().item() <= 1.5e-1 + _assert_amax_is_of_the_output(amax_o, got) + + +# --- gaps found by auditing the suite against what the config permits ----- + + +@pytest.mark.L0 +@pytest.mark.parametrize("splits", [1, 4], ids=lambda s: f"split{s}") +def test_combine_lse_matches_reference(splits): + """The RECOMBINED LSE is an output too, and nothing else here checks it. + + split_combine_sm100 computes lse = M + log(sum_s exp(lse_s - M)); every other + test only compares O, so a wrong LSE would pass all of them. + """ + import cutlass + import cuda.bindings.driver as cuda_driver + + from cudnn.sdpa.fwd.kernels import split_combine_sm100 as comb + + B, H, SQ, SKV = 2, 4, 128, 2048 + dev = "cuda" + scale = 1.0 / math.sqrt(D) + torch.manual_seed(0) + q = torch.randn(B, SQ, H, D, device=dev, dtype=torch.float16) + k = torch.randn(B, SKV, H, D, device=dev, dtype=torch.float16) + v = torch.randn(B, SKV, H, D, device=dev, dtype=torch.float16) + + mod = _kernel_module(splits, 3, causal=False) + fn = mod.compile(b=B, qh=H, kh=H, sq=SQ, skv=SKV, d_qk=D, d_v=D, has_lse=True) + o_p = torch.zeros(splits * B, SQ, H, D, device=dev, dtype=torch.float16) + lse_p = torch.zeros(splits * B, H, SQ, device=dev, dtype=torch.float32) + stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream) + fn( + q, + k, + v, + o_p, + lse_p, + torch.zeros(H, dtype=torch.float32, device=dev), + torch.zeros(B, dtype=torch.int32, device=dev), + torch.zeros(1, dtype=torch.int64, device=dev), + (B, H, H, SQ, SKV, 0), + cutlass.Float32(scale * math.log2(math.e)), + cutlass.Int32(0), + None, + stream=stream, + ) + + # Reference LSE = logsumexp of the scaled scores, in natural log. + qb, kb = (t.float().permute(0, 2, 1, 3) for t in (q, k)) + ref_lse = torch.logsumexp(torch.matmul(qb, kb.transpose(-1, -2)) * scale, dim=-1) # [B,H,SQ] + + if splits == 1: + torch.cuda.synchronize() + got_lse = lse_p.view(B, H, SQ) + else: + o_out = torch.zeros(B, SQ, H, D, device=dev, dtype=torch.float16) + lse_out = torch.zeros(B, H, SQ, device=dev, dtype=torch.float32) + cfn = comb.compile(b=B, h=H, sq=SQ, d_v=D, splits=splits, dtype_o="f16", has_lse=True) + cfn(o_p, lse_p, o_out, lse_out, None, (B, H, SQ, D), cutlass.Int32(splits), stream=stream) + torch.cuda.synchronize() + got_lse = lse_out + assert not torch.isnan(got_lse).any(), "NaN in recombined LSE" + assert (got_lse - ref_lse).abs().max().item() <= 5e-3 + + +@pytest.mark.L0 +@pytest.mark.parametrize("flavor", sorted(_F16_FLAVORS)) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) +def test_even_splits_every_flavor_batched(flavor, dtype): + """Even splits, B > 1, both half dtypes, on every flavor. + + B > 1 matters specifically: the split-major slot is batch + split*B and the + grid.z decode is b % n_batch / b // n_batch, all of which are trivially + correct at B == 1 and so were barely exercised. + """ + import os as _os + + import cutlass + import cuda.bindings.driver as cuda_driver + + from cudnn.frost.template_loader import load_template + from cudnn.sdpa.fwd import api_dsl + from cudnn.sdpa.fwd.config_sm100 import TemplateParams + from cudnn.sdpa.fwd.kernels import split_combine_sm100 as comb + + kmod, d_qk, d_v = _F16_FLAVORS[flavor] + B, H, SQ, SKV, S = 3, 4, 128, 2048, 4 + dev = "cuda" + scale = 1.0 / math.sqrt(d_qk) + torch.manual_seed(0) + q = torch.randn(B, SQ, H, d_qk, device=dev, dtype=dtype) + k = torch.randn(B, SKV, H, d_qk, device=dev, dtype=dtype) + v = torch.randn(B, SKV, H, d_v, device=dev, dtype=dtype) + + path = _os.path.join(_os.path.dirname(_os.path.abspath(api_dsl.__file__)), "kernels", kmod) + params = TemplateParams(dtype_qkv=3 if dtype == torch.float16 else 2, split_kv=S) + mod = load_template(path, params, tag=f"even_{flavor}_{dtype}_{S}") + fn = mod.compile(b=B, qh=H, kh=H, sq=SQ, skv=SKV, d_qk=d_qk, d_v=d_v, has_lse=True) + o_p = torch.zeros(S * B, SQ, H, d_v, device=dev, dtype=dtype) + lse_p = torch.zeros(S * B, H, SQ, device=dev, dtype=torch.float32) + stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream) + fn( + q, + k, + v, + o_p, + lse_p, + torch.zeros(H, dtype=torch.float32, device=dev), + torch.zeros(B, dtype=torch.int32, device=dev), + torch.zeros(1, dtype=torch.int64, device=dev), + (B, H, H, SQ, SKV, 0), + cutlass.Float32(scale * math.log2(math.e)), + cutlass.Int32(0), + None, + stream=stream, + ) + o_out = torch.zeros(B, SQ, H, d_v, device=dev, dtype=dtype) + cfn = comb.compile(b=B, h=H, sq=SQ, d_v=d_v, splits=S, dtype_o="f16" if dtype == torch.float16 else "bf16", has_lse=False) + cfn(o_p, lse_p, o_out, None, None, (B, H, SQ, d_v), cutlass.Int32(S), stream=stream) + torch.cuda.synchronize() + assert not torch.isnan(o_out).any(), f"{flavor}: NaN" + ref = _ref_sdpa(q, k, v, scale, is_causal=False, kh=H) + assert (o_out.float() - ref).abs().max().item() <= (2e-2 if dtype == torch.float16 else 1e-1) + + +# --- masks the config permits with split, previously untested ------------- +# +# split_kv only rejects sinks and THD, so SWA / bottom-right / padded are all +# reachable. They are also where _bounds_for_tile_split does its real work: +# it slices the ALREADY-masked [left, right) and clamps unmasked_lo/hi into the +# slice, and none of that was exercised. The reference is the one the sibling +# suite uses to model this kernel's exact mask semantics. + + +def _run_masked(kfile, d_qk, d_v, splits, *, B, H, KH, SQ, SKV, tp_kwargs, seq_kv_lens=None, seq_q_lens=None, dtype=torch.float16, cta_mma=2): + """Launch split kernel + combine under an arbitrary mask; returns (got, q, k, v, scale).""" + import os as _os + + import cutlass + import cuda.bindings.driver as cuda_driver + + from cudnn.frost.template_loader import load_template + from cudnn.sdpa.fwd import api_dsl + from cudnn.sdpa.fwd.config_sm100 import TemplateParams + from cudnn.sdpa.fwd.kernels import split_combine_sm100 as comb + + dev = "cuda" + scale = 1.0 / math.sqrt(d_qk) + torch.manual_seed(0) + q = torch.randn(B, SQ, H, d_qk, device=dev, dtype=dtype) + k = torch.randn(B, SKV, KH, d_qk, device=dev, dtype=dtype) + v = torch.randn(B, SKV, KH, d_v, device=dev, dtype=dtype) + + path = _os.path.join(_os.path.dirname(_os.path.abspath(api_dsl.__file__)), "kernels", kfile) + params = TemplateParams(dtype_qkv=3 if dtype == torch.float16 else 2, split_kv=splits, cta_mma=cta_mma, **tp_kwargs) + mod = load_template(path, params, tag=f"mask_{kfile[:16]}_{splits}_{cta_mma}_{sorted(tp_kwargs.items())}") + fn = mod.compile(b=B, qh=H, kh=KH, sq=SQ, skv=SKV, d_qk=d_qk, d_v=d_v, has_lse=True) + + o_p = torch.zeros(splits * B, SQ, H, d_v, device=dev, dtype=dtype) + lse_p = torch.zeros(splits * B, H, SQ, device=dev, dtype=torch.float32) + skv_t = seq_kv_lens if seq_kv_lens is not None else torch.zeros(B, dtype=torch.int32, device=dev) + stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream) + fn( + q, + k, + v, + o_p, + lse_p, + torch.zeros(H, dtype=torch.float32, device=dev), + skv_t, + torch.zeros(1, dtype=torch.int64, device=dev), + (B, H, KH, SQ, SKV, 0), + cutlass.Float32(scale * math.log2(math.e)), + cutlass.Int32(0), + seq_q_lens, + stream=stream, + ) + if splits == 1: + torch.cuda.synchronize() + return o_p.float(), q, k, v, scale + o_out = torch.zeros(B, SQ, H, d_v, device=dev, dtype=dtype) + cfn = comb.compile(b=B, h=H, sq=SQ, d_v=d_v, splits=splits, dtype_o="f16" if dtype == torch.float16 else "bf16", has_lse=False) + cfn(o_p, lse_p, o_out, None, None, (B, H, SQ, d_v), cutlass.Int32(splits), stream=stream) + torch.cuda.synchronize() + assert not torch.isnan(o_out).any(), "NaN under mask+split" + return o_out.float(), q, k, v, scale + + +def _bhsd(t): + return t.permute(0, 2, 1, 3) + + +@pytest.mark.L0 +@pytest.mark.parametrize("splits", [1, 4], ids=lambda s: f"split{s}") +def test_split_kv_swa_causal(splits): + """Causal + sliding window: both mask bits, so unmasked_lo AND hi are live.""" + from test_sdpa_fwd_dsl_sm100 import _ref_sdpa_full + + W = 256 + B, H, SQ, SKV = 1, 4, 1024, 1024 + got, q, k, v, scale = _run_masked( + "prefill_d128_f16_sm100.py", 128, 128, splits, B=B, H=H, KH=H, SQ=SQ, SKV=SKV, tp_kwargs=dict(window_right=0, window_left=W) + ) + ref = _ref_sdpa_full(_bhsd(q), _bhsd(k), _bhsd(v), scale=scale, is_causal=True, swa_window=W) + assert (got - ref.float().permute(0, 2, 1, 3)).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +@pytest.mark.parametrize("splits", [1, 4], ids=lambda s: f"split{s}") +def test_split_kv_bottom_right_causal(splits): + """Bottom-right anchored causal — the diagonal sits at (S_q, S_kv).""" + from test_sdpa_fwd_dsl_sm100 import _ref_sdpa_full + + B, H, SQ, SKV = 1, 4, 128, 2048 + got, q, k, v, scale = _run_masked( + "prefill_d128_f16_sm100.py", 128, 128, splits, B=B, H=H, KH=H, SQ=SQ, SKV=SKV, tp_kwargs=dict(window_right=0, bottom_right=True) + ) + ref = _ref_sdpa_full(_bhsd(q), _bhsd(k), _bhsd(v), scale=scale, is_causal=True, bottom_right=True) + assert (got - ref.float().permute(0, 2, 1, 3)).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +@pytest.mark.parametrize("splits", [1, 4], ids=lambda s: f"split{s}") +def test_split_kv_padded_kv(splits): + """Per-batch KV padding: the split must slice each batch's OWN live range.""" + from test_sdpa_fwd_dsl_sm100 import _ref_sdpa_full + + B, H, SQ, SKV = 2, 4, 128, 2048 + lens = torch.tensor([2048, 1531], dtype=torch.int32, device="cuda") # 2nd ends mid-tile + got, q, k, v, scale = _run_masked( + "prefill_d128_f16_sm100.py", 128, 128, splits, B=B, H=H, KH=H, SQ=SQ, SKV=SKV, tp_kwargs=dict(seq_kv_lens_present=True), seq_kv_lens=lens + ) + ref = _ref_sdpa_full(_bhsd(q), _bhsd(k), _bhsd(v), scale=scale, seq_kv_lens=lens) + assert (got - ref.float().permute(0, 2, 1, 3)).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +@pytest.mark.parametrize("flavor", ["d192", "d256", "d512"]) +def test_split_kv_causal_other_flavors(flavor): + """Causal + split beyond d128 — _split_chunk slices the masked range here too.""" + from test_sdpa_fwd_dsl_sm100 import _ref_sdpa_full + + kmod, d_qk, d_v = _F16_FLAVORS[flavor] + B, H, SQ, SKV = 1, 4, 1024, 1024 + got, q, k, v, scale = _run_masked(kmod, d_qk, d_v, 4, B=B, H=H, KH=H, SQ=SQ, SKV=SKV, tp_kwargs=dict(window_right=0)) + ref = _ref_sdpa_full(_bhsd(q), _bhsd(k), _bhsd(v), scale=scale, is_causal=True) + assert (got - ref.float().permute(0, 2, 1, 3)).abs().max().item() <= 2e-2 + + +@pytest.mark.L0 +@pytest.mark.parametrize("splits", [1, 4], ids=lambda s: f"split{s}") +def test_split_kv_padded_q_trim(splits): + """Dense padded-Q trim (seq_q_lens) + split. + + Q rows at or past the batch's actual length must come back O := 0 (cuDNN + >= 9.14 convention). Under split that has to hold for EVERY split's partial, + or the combine mixes live and dead rows -- the trim is applied per split in + the epilogue, after which the dead row's lse = -inf makes it drop out. + """ + from test_sdpa_fwd_dsl_sm100 import _ref_sdpa_full + + B, H, SQ, SKV = 2, 4, 256, 2048 + kv_lens = torch.tensor([2048, 2048], dtype=torch.int32, device="cuda") + q_lens = torch.tensor([256, 137], dtype=torch.int32, device="cuda") # 2nd trims mid-tile + got, q, k, v, scale = _run_masked( + "prefill_d128_f16_sm100.py", + 128, + 128, + splits, + B=B, + H=H, + KH=H, + SQ=SQ, + SKV=SKV, + tp_kwargs=dict(seq_kv_lens_present=True, seq_q_lens_present=True), + seq_kv_lens=kv_lens, + seq_q_lens=q_lens, + ) + ref = _ref_sdpa_full(_bhsd(q), _bhsd(k), _bhsd(v), scale=scale, seq_kv_lens=kv_lens) + ref = ref.float().permute(0, 2, 1, 3) # -> BSHD + # Rows past the per-batch Q length must be exactly zero. + rows = torch.arange(SQ, device=got.device).view(1, SQ, 1, 1) + dead = rows >= q_lens.view(B, 1, 1, 1) + 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