diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/__init__.py b/python/sglang/kernels/ops/attention/fa4_sm120/__init__.py new file mode 100644 index 000000000..738249815 --- /dev/null +++ b/python/sglang/kernels/ops/attention/fa4_sm120/__init__.py @@ -0,0 +1 @@ +"""SGLang-owned FlashAttention-4 kernels and launch policy for SM120.""" diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/dispatch.py b/python/sglang/kernels/ops/attention/fa4_sm120/dispatch.py new file mode 100644 index 000000000..49dada74d --- /dev/null +++ b/python/sglang/kernels/ops/attention/fa4_sm120/dispatch.py @@ -0,0 +1,58 @@ +# Copyright (c) 2026, SGLang Team. +"""Lightweight dispatch bridge for SGLang-owned SM120 FA4 kernels. + +The vendored FA4 interface dispatches through this module so the SM120 +implementation and its launch state remain outside ``flash_attn/cute``. +""" + +from functools import lru_cache + + +@lru_cache(maxsize=None) +def get_forward_host(arch: int): + """Return the optional architecture-owned forward host.""" + if arch // 10 == 12: + from sglang.kernels.ops.attention.fa4_sm120.runtime import ( + sm120_forward_host, + ) + + return sm120_forward_host + return None + + +def resolve_runtime_policy( + *, + device_capability: tuple[int, int], + deterministic: bool, +) -> tuple[int, int, bool]: + """Resolve generic and architecture-owned SplitKV launch policy.""" + arch = device_capability[0] * 10 + device_capability[1] + uses_arch_decode_policy = get_forward_host(arch) is not None + no_splitkv = device_capability < (9, 0) or uses_arch_decode_policy + num_splits = 1 if deterministic or no_splitkv else 0 + decode_num_splits = ( + 0 if uses_arch_decode_policy and not deterministic else num_splits + ) + return num_splits, decode_num_splits, uses_arch_decode_policy + + +@lru_cache(maxsize=None) +def get_forward_arch(device) -> int | None: + """Return the device arch when it has an architecture-owned forward host.""" + import torch + + major, minor = torch.cuda.get_device_capability(device) + arch = major * 10 + minor + return arch if get_forward_host(arch) is not None else None + + +def try_cached_paged_decode(*, arch: int, **kwargs): + """Try an architecture-owned paged-decode launch plan.""" + host = get_forward_host(arch) + return None if host is None else host.try_paged_decode(arch=arch, **kwargs) + + +def try_cached_varlen(*, arch: int, **kwargs): + """Try an architecture-owned varlen launch plan.""" + host = get_forward_host(arch) + return None if host is None else host.try_varlen(arch=arch, **kwargs) diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py b/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py new file mode 100644 index 000000000..4188b95c4 --- /dev/null +++ b/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd.py @@ -0,0 +1,4552 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# SM120 (Blackwell GeForce / DGX Spark) forward pass. + +import math +import operator +from functools import lru_cache, partial +from types import SimpleNamespace +from typing import Callable, Optional + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils_basic +from cutlass import Float32, Int32, const_expr +from cutlass.cute import FastDivmodDivisor +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.pipeline import ( + Agent, + CooperativeGroup, + PipelineAsync, + PipelineState, + PipelineTmaAsync, + pipeline_init_arrive, + pipeline_init_wait, +) +from quack import copy_utils, layout_utils + +from sglang.kernels.ops.attention.fa4_sm120.paged_kv import Sm120PagedKVManager +from sglang.kernels.ops.attention.fa4_sm120.policy import ( + LOW_HD_DECODE_SHAPES, + LOW_HD_DECODE_TILE_N, + low_hd_paged_decode_tile_m, + visible_decode_seqlen_k, +) +from sglang.kernels.ops.attention.fa4_sm120.scheduler import ( + Sm120UniformBatchScheduler, +) +from sglang.kernels.ops.attention.flash_attn.cute import pipeline as pipeline_custom +from sglang.kernels.ops.attention.flash_attn.cute import utils +from sglang.kernels.ops.attention.flash_attn.cute.block_info import BlockInfo +from sglang.kernels.ops.attention.flash_attn.cute.block_sparsity import ( + BlockSparseTensors, +) +from sglang.kernels.ops.attention.flash_attn.cute.cute_dsl_utils import ( + assume_tensor_aligned, +) +from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd import ( + FlashAttentionForwardBase, +) +from sglang.kernels.ops.attention.flash_attn.cute.mask import AttentionMask +from sglang.kernels.ops.attention.flash_attn.cute.named_barrier import NamedBarrierFwd +from sglang.kernels.ops.attention.flash_attn.cute.pack_gqa import ( + PackGQA, + pack_gqa_layout, +) +from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK +from sglang.kernels.ops.attention.flash_attn.cute.softmax import ( + Softmax, + apply_score_mod_inner, +) +from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import ( + SchedulingMode, + SingleTileScheduler, + SingleTileVarlenScheduler, + TileSchedulerArguments, + TileSchedulerProtocol, +) +from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData + + +class FlashAttentionForwardSm120(FlashAttentionForwardBase): + """SM120 warp-MMA forward kernel with TMA fused into the QK warps.""" + + # Experimental same-page paged-KV TMA path. Scratch benchmarks set this on + # the kernel instance; qualified dispatch keeps it disabled. + paged_tma = False + + def __init__( + self, + *args, + direct_uniform_batch: bool = False, + paged_kv: bool = False, + split_qk_n: bool = False, + split_kv_blocks_per_cta: int = 0, + has_bias: bool = False, + bias_block_size: int = 64, + rel_extent_padded: int = 128, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.direct_uniform_batch = direct_uniform_batch + self.paged_kv = paged_kv + self.split_qk_n = split_qk_n + self.split_kv_blocks_per_cta = split_kv_blocks_per_cta + self.has_bias = has_bias + self.bias_block_size = bias_block_size + self.rel_extent_padded = rel_extent_padded + if has_bias: + assert not self._uses_split_pv_warps() + assert self.tile_n == 128 + assert 0 < bias_block_size <= self.tile_m + assert bias_block_size % 8 == 0 + assert rel_extent_padded >= 128 + assert rel_extent_padded % 128 == 0 + self.bias_n_max = rel_extent_padded // self.tile_n if has_bias else 0 + + @cute.jit + def _get_n_block_min_max( + self, + block_info: BlockInfo, + seqlen: SeqlenInfoQK, + m_block: Int32, + split_idx: Int32, + num_splits: Int32, + ): + """Balance ragged decode requests without changing the workspace bound.""" + if const_expr(self.split_kv_blocks_per_cta <= 0): + return block_info.get_n_block_min_max( + seqlen, m_block, split_idx, num_splits + ) + n_block_min, n_block_max = block_info.get_n_block_min_max( + seqlen, + m_block, + absolute=True, + ) + dynamic_splits = cute.ceil_div( + cutlass.max(n_block_max - n_block_min, 1), + self.split_kv_blocks_per_cta, + ) + return block_info.get_n_block_min_max( + seqlen, + m_block, + split_idx, + dynamic_splits, + ) + + num_dma_threads = 32 + # Relative CTA critical-path costs keyed by the complete kernel config: + # (fixed work, work per N block, extra local-mask work). Sequence length, + # window, packed heads, and SM count remain analytic inputs to the generic + # LPT model below. + _lpt_cost_by_config = { + (256, 256, 32, 64): (213, 95, 95), + (256, 256, 48, 64): (221, 99, 74), + (256, 256, 64, 64): (212, 101, 67), + } + _lpt_tie_margin = 1 + _qualified_wave_tile_shapes = frozenset( + ((32, 32), (64, 64), (96, 96), (128, 128), (192, 128)) + ) + + @staticmethod + def _estimate_lpt_makespan( + tile_m: int, + tile_n: int, + cost: tuple[int, int, int], + *, + seqlen_q: int, + seqlen_k: int, + num_sms: int, + num_head_kv: int, + qhead_per_kvhead: int, + is_causal: bool, + is_local: bool, + window_size_left: int | None, + window_size_right: int | None, + ) -> int | None: + """Estimate the one- or two-wave LPT critical path in relative units.""" + if ( + seqlen_q <= 0 + or seqlen_k <= 0 + or num_sms <= 0 + or num_head_kv <= 0 + or qhead_per_kvhead <= 0 + ): + return None + + packed_q = seqlen_q * qhead_per_kvhead + num_m_blocks = (packed_q + tile_m - 1) // tile_m + num_ctas = num_m_blocks * num_head_kv + # The exact boundary expression below is intentionally limited to two + # physical waves. Returning None also keeps this estimator O(1). + if num_ctas > 2 * num_sms: + return None + + fixed_cost, n_block_cost, local_mask_cost = cost + fixed_cost += local_mask_cost if is_local else 0 + num_k_blocks = (seqlen_k + tile_n - 1) // tile_n + seqlen_delta = seqlen_k - seqlen_q + + def cta_cost(launch_idx: int) -> int: + # SingleTileVarlenScheduler repeats each reversed (LPT) M block + # across the packed KV heads before advancing to the next block. + m_block = num_m_blocks - 1 - launch_idx // num_head_kv + m_idx_min = m_block * tile_m // qhead_per_kvhead + m_idx_max = ( + (m_block + 1) * tile_m + qhead_per_kvhead - 1 + ) // qhead_per_kvhead + + if is_causal or (is_local and window_size_right is not None): + n_idx_right = m_idx_max + seqlen_delta + if not is_causal: + n_idx_right += window_size_right + n_block_max = min( + num_k_blocks, + max(0, (n_idx_right + tile_n - 1) // tile_n), + ) + else: + n_block_max = num_k_blocks + + n_block_min = 0 + if is_local and window_size_left is not None: + n_idx_left = m_idx_min + seqlen_delta - window_size_left + n_block_min = max(n_idx_left // tile_n, 0) + num_n_blocks = max(n_block_max - n_block_min, 0) + return fixed_cost + n_block_cost * num_n_blocks + + heaviest_cta = cta_cost(0) + if num_ctas <= num_sms: + return heaviest_cta + # With at most two waves, the first tail CTA is paired with the lightest + # CTA in the first hardware wave. This captures the discrete tail that + # an average-work occupancy model misses. + wave_boundary = cta_cost(num_sms - 1) + cta_cost(num_sms) + return max(heaviest_cta, wave_boundary) + + @staticmethod + def _fits_lpt_equivalent_wave( + tile_n: int, + resident_ctas_per_sm: int, + *, + seqlen_q: int, + seqlen_k: int, + num_sms: int, + num_head_kv: int, + qhead_per_kvhead: int, + is_causal: bool, + is_local: bool, + window_size_left: int | None, + window_size_right: int | None, + ) -> bool: + """Return whether structural CTA work fits one LPT-equivalent wave.""" + if resident_ctas_per_sm <= 0: + return False + packed_q = seqlen_q * qhead_per_kvhead + num_ctas = ((packed_q + 63) // 64) * num_head_kv + structural_cost = (0, 1, 0) + workload = FlashAttentionForwardSm120._estimate_lpt_makespan( + 64, + tile_n, + structural_cost, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + num_sms=num_sms * resident_ctas_per_sm, + num_head_kv=num_head_kv, + qhead_per_kvhead=qhead_per_kvhead, + is_causal=is_causal, + is_local=is_local, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + heaviest_cta = FlashAttentionForwardSm120._estimate_lpt_makespan( + 64, + tile_n, + structural_cost, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + num_sms=max(num_ctas, 1), + num_head_kv=num_head_kv, + qhead_per_kvhead=qhead_per_kvhead, + is_causal=is_causal, + is_local=is_local, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + return ( + workload is not None + and heaviest_cta is not None + and workload <= heaviest_cta + ) + + @staticmethod + def _select_qualified_tile_n( + head_dim: int, + head_dim_v: int, + *, + seqlen_q: int, + seqlen_k: int, + num_sms: int, + num_head_kv: int, + qhead_per_kvhead: int, + is_causal: bool, + is_local: bool, + window_size_left: int | None, + window_size_right: int | None, + ) -> int | None: + """Select N from LPT workload and SM120 residency without timing fits.""" + shape = (head_dim, head_dim_v) + if shape not in FlashAttentionForwardSm120._qualified_wave_tile_shapes: + return None + if ( + seqlen_q <= 0 + or seqlen_k <= 0 + or num_sms <= 0 + or num_head_kv <= 0 + or qhead_per_kvhead <= 0 + ): + return None + + has_steady_state_k_loop = seqlen_k >= 1024 + # A compact SM array reaches multi-wave scheduling much earlier. The + # narrower N64 loop wins there for several exact head shapes, while + # the 110- and 188-SM SM120 SKUs retain their calibrated wider tiles. + has_compact_sm_array = num_sms <= 64 + + def fits_lpt_wave(tile_n: int, resident_ctas_per_sm: int) -> bool: + return FlashAttentionForwardSm120._fits_lpt_equivalent_wave( + tile_n, + resident_ctas_per_sm, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + num_sms=num_sms, + num_head_kv=num_head_kv, + qhead_per_kvhead=qhead_per_kvhead, + is_causal=is_causal, + is_local=is_local, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + + if is_local: + if shape == (32, 32): + # N128 retains two resident CTAs/SM. + return 128 if fits_lpt_wave(128, 2) else 64 + if shape == (64, 64): + return 128 if fits_lpt_wave(128, 1) else 64 + return 64 + + if shape == (32, 32): + if has_compact_sm_array and seqlen_k >= 1536: + return 64 + # N128 is robust before the compact-SM steady state. N256's + # marginal single-CTA gain reverses under a nearly full wave. + return 128 + if shape == (64, 64): + if not has_steady_state_k_loop: + return 128 + return 192 if fits_lpt_wave(192, 1) else 64 + if shape == (96, 96): + if has_compact_sm_array and has_steady_state_k_loop: + return 64 + # N128's more efficient K loop wins on the larger SM120 arrays. + return 128 + if shape == (128, 128): + if has_compact_sm_array and 512 <= seqlen_k < 4096: + return 64 + return 128 + # HD192 N96 is at best tied on the smaller SKU and regresses the larger + # one, so keep the zero-spill N64 configuration. + return 64 + + @staticmethod + def _smem_usage_in_bytes( + head_dim, + head_dim_v, + tile_m, + tile_n, + num_stages, + Q_in_regs=False, + ) -> int: + """Return SMEM usage after padding head dimensions to kernel alignment.""" + head_dim = math.ceil(head_dim / 16) * 16 + head_dim_v = math.ceil(head_dim_v / 16) * 16 + element_size = 2 + smem_usage_Q = tile_m * head_dim * element_size + smem_usage_K = tile_n * head_dim * num_stages * element_size + smem_usage_V = tile_n * head_dim_v * num_stages * element_size + smem_usage_QV = ( + smem_usage_Q + smem_usage_V + if not Q_in_regs + else max(smem_usage_Q, smem_usage_V) + ) + smem_usage = smem_usage_QV + smem_usage_K + if (head_dim, head_dim_v, tile_m, tile_n) in ( + (256, 256, 16, 64), + (256, 256, 16, 80), + (256, 256, 32, 64), + (256, 256, 48, 64), + (256, 256, 64, 48), + (256, 256, 64, 64), + ): + # Q is copied to QK registers before the mainloop, so its allocation + # can hold both P stages and later the O epilogue tile. Four FP32 + # rows hold the per-stage rescale values and final scale/LSE. + smem_usage += 4 * tile_m * 4 + # Q/K/V/P/final full-empty mbarriers and conservative field alignment. + return smem_usage + 2048 + + @staticmethod + @lru_cache(maxsize=4096) + def get_fwd_tile_size( + head_dim: int, + head_dim_v: int, + total_q_rows: int | None = None, + num_sms: int | None = None, + num_batch: int | None = None, + seqlen_q: int | None = None, + seqlen_k: int | None = None, + num_head_kv: int | None = None, + qhead_per_kvhead: int | None = None, + is_causal: bool = False, + is_local: bool = False, + window_size_left: int | None = None, + window_size_right: int | None = None, + pack_gqa: bool = False, + paged_kv: bool = False, + ) -> tuple[int, int]: + """Select an SM120 tile that fits the architecture's 99 KB SMEM.""" + smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_120") + shape = (head_dim, head_dim_v) + qualified_shape = ( + shape in FlashAttentionForwardSm120._qualified_wave_tile_shapes + ) + has_compact_q_groups = pack_gqa or (paged_kv and qhead_per_kvhead == 1) + is_short_compact_q = ( + has_compact_q_groups + and seqlen_q is not None + and qhead_per_kvhead is not None + and seqlen_q * qhead_per_kvhead <= 16 + ) + visible_seqlen_k = ( + None + if seqlen_k is None + else visible_decode_seqlen_k( + seqlen_k, + is_local=is_local, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + ) + low_hd_m16_total_mblocks = None + if ( + num_batch is not None + and num_head_kv is not None + and seqlen_q is not None + and qhead_per_kvhead is not None + ): + packed_q_rows = seqlen_q * qhead_per_kvhead + low_hd_m16_total_mblocks = ( + num_batch * num_head_kv * ((packed_q_rows + 16 - 1) // 16) + ) + low_hd_tile_m = low_hd_paged_decode_tile_m( + head_dim=head_dim, + head_dim_v=head_dim_v, + paged_kv=paged_kv, + seqlen_q=seqlen_q, + visible_seqlen_k=visible_seqlen_k, + qhead_per_kvhead=qhead_per_kvhead, + num_sms=num_sms, + total_mblocks=low_hd_m16_total_mblocks, + ) + if low_hd_tile_m is not None: + fallback_candidates = ((low_hd_tile_m, LOW_HD_DECODE_TILE_N, 1),) + elif (head_dim, head_dim_v) == (256, 256): + if is_short_compact_q: + fallback_candidates = ( + (16, 64, 1), + (16, 80, 1), + (32, 64, 1), + (64, 64, 1), + ) + elif ( + total_q_rows is not None + and num_sms is not None + and total_q_rows > 76 * num_sms + and total_q_rows <= 92 * num_sms + ): + # Preserve the previous multi-batch/non-packed fallback where + # per-sequence LPT costs are unavailable on the host. + fallback_candidates = ( + (48, 64, 1), + (64, 64, 1), + (64, 48, 1), + (32, 64, 1), + ) + else: + fallback_candidates = ( + (64, 64, 1), + (64, 48, 1), + (48, 64, 1), + (32, 64, 1), + ) + elif qualified_shape: + # M64 is the zero-spill fallback across the qualified exact shapes. + # HD32 global benefits from N128 without giving up its second + # resident CTA; local masks retain the more parallel N64 tile. + safe_tile_n = 128 if shape == (32, 32) and not is_local else 64 + fallback_candidates = ((64, safe_tile_n, 1),) + ( + ((64, 64, 1),) if safe_tile_n != 64 else () + ) + else: + fallback_candidates = ( + ((128, 128, 1),) if max(head_dim, head_dim_v) <= 64 else () + ) + ((128, 64, 1), (64, 64, 1)) + + can_select_qualified_tile = ( + qualified_shape + and num_batch == 1 + and pack_gqa + and total_q_rows is not None + and num_sms is not None + and seqlen_q is not None + and seqlen_k is not None + and num_head_kv is not None + and qhead_per_kvhead is not None + and total_q_rows == seqlen_q * num_head_kv * qhead_per_kvhead + and (is_causal or is_local) + and not is_short_compact_q + ) + candidates = fallback_candidates + if can_select_qualified_tile: + preferred_tile_n = FlashAttentionForwardSm120._select_qualified_tile_n( + head_dim, + head_dim_v, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + num_sms=num_sms, + num_head_kv=num_head_kv, + qhead_per_kvhead=qhead_per_kvhead, + is_causal=is_causal, + is_local=is_local, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + if preferred_tile_n is not None: + preferred = (64, preferred_tile_n, 1) + candidates = (preferred,) + tuple( + candidate + for candidate in fallback_candidates + if candidate != preferred + ) + + # HD256 retains the calibrated LPT selection used by its dedicated + # bounded-SMEM pipeline. + model_candidates = tuple( + (config[2], config[3], 1) + for config in FlashAttentionForwardSm120._lpt_cost_by_config + if config[:2] == shape + and FlashAttentionForwardSm120._smem_usage_in_bytes(*config, 1) + <= smem_capacity + ) + can_model_lpt = ( + len(model_candidates) >= 2 + and num_batch == 1 + and pack_gqa + and total_q_rows is not None + and num_sms is not None + and seqlen_q is not None + and seqlen_k is not None + and num_head_kv is not None + and qhead_per_kvhead is not None + and total_q_rows == seqlen_q * num_head_kv * qhead_per_kvhead + and (is_causal or is_local) + and not is_short_compact_q + ) + if can_model_lpt: + scores = { + candidate: FlashAttentionForwardSm120._estimate_lpt_makespan( + candidate[0], + candidate[1], + FlashAttentionForwardSm120._lpt_cost_by_config[ + (head_dim, head_dim_v, candidate[0], candidate[1]) + ], + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + num_sms=num_sms, + num_head_kv=num_head_kv, + qhead_per_kvhead=qhead_per_kvhead, + is_causal=is_causal, + is_local=is_local, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + for candidate in model_candidates + } + valid_scores = { + candidate: score + for candidate, score in scores.items() + if score is not None + } + if valid_scores: + exact_best = min(valid_scores, key=valid_scores.get) + packed_q = seqlen_q * qhead_per_kvhead + + def num_ctas(candidate: tuple[int, int, int]) -> int: + return ((packed_q + candidate[0] - 1) // candidate[0]) * num_head_kv + + max_model_ctas = max( + num_ctas(candidate) for candidate in model_candidates + ) + if num_ctas(exact_best) == max_model_ctas: + # Keep a genuinely faster high-parallelism tile. Once that + # tile loses outright, a near tie favors fewer CTAs. + preferred = exact_best + else: + cutoff = ( + valid_scores[exact_best] + + FlashAttentionForwardSm120._lpt_tie_margin + ) + near = tuple( + candidate + for candidate, score in valid_scores.items() + if score <= cutoff + ) + preferred = min( + near, + key=lambda candidate: ( + num_ctas(candidate), + -(candidate[0] * candidate[1]), + ), + ) + candidates = (preferred,) + tuple( + candidate + for candidate in fallback_candidates + if candidate != preferred + ) + compact_mha_local_candidate = None + if ( + shape == (256, 256) + and num_sms is not None + and num_sms <= 64 + and is_local + and not paged_kv + and num_batch == 1 + and pack_gqa + and qhead_per_kvhead == 1 + and seqlen_q is not None + and seqlen_q >= 2048 + and window_size_left is not None + and window_size_left > 0 + and window_size_right == 0 + ): + # On the compact SM array, long-query MHA with a narrow left + # window benefits from smaller M tiles: they reduce masked local + # work while the long Q dimension amortizes the extra CTAs. Use + # Q/window ratios so the crossover scales with both dimensions. + if window_size_left <= 128 and seqlen_q >= 80 * window_size_left: + compact_mha_local_candidate = (32, 64, 1) + elif seqlen_q >= 40 * window_size_left: + compact_mha_local_candidate = (48, 64, 1) + if compact_mha_local_candidate is not None: + candidates = (compact_mha_local_candidate,) + tuple( + candidate + for candidate in candidates + if candidate != compact_mha_local_candidate + ) + if ( + shape == (256, 256) + and num_sms is not None + and num_sms <= 64 + and is_local + and seqlen_k is not None + and seqlen_k >= 512 + and qhead_per_kvhead is not None + and qhead_per_kvhead > 1 + ): + # The 48-SM class benefits from N48's lower masked-KV work once + # packed GQA local attention reaches its steady state. Keep MHA + # and larger SM arrays on the cross-SKU HD256 LPT calibration. + preferred = (64, 48, 1) + candidates = (preferred,) + tuple( + candidate for candidate in candidates if candidate != preferred + ) + for tile_m, tile_n, num_stages in candidates: + if ( + FlashAttentionForwardSm120._smem_usage_in_bytes( + head_dim, + head_dim_v, + tile_m, + tile_n, + num_stages, + ) + <= smem_capacity + ): + return tile_m, tile_n + raise ValueError( + f"(head_dim, head_dim_v)=({head_dim}, {head_dim_v}) exceeds " + f"SM120 shared-memory capacity ({smem_capacity} bytes)" + ) + + @staticmethod + def get_fwd_num_stages( + head_dim: int, head_dim_v: int, tile_m: int, tile_n: int + ) -> int: + """Return the public pipeline specialization depth.""" + return 1 + + @staticmethod + def get_fwd_num_threads( + head_dim: int, + head_dim_v: int, + tile_m: int, + tile_n: int, + paged_kv: bool = False, + ) -> int: + """Return the number of warp-MMA consumer threads for an SM120 tile.""" + if ( + paged_kv + and (head_dim, head_dim_v) in LOW_HD_DECODE_SHAPES + and tile_m in (16, 32) + and tile_n == LOW_HD_DECODE_TILE_N + ): + return tile_m * 2 + if paged_kv and head_dim == 256 and head_dim_v == 256: + if (tile_m, tile_n) == (16, 64): + # Decode is SMEM-limited to one CTA/SM. Four consumer warps + # expose QK/PV parallelism without reducing CTA residency. + return 128 + # The contiguous HD256 path assigns distinct QK and PV warp sets. + # Paged KV instead reserves one DMA warp for gather and has each + # consumer warp own the same 16 rows through both MMA phases. + return tile_m * 2 + config = (head_dim, head_dim_v, tile_m, tile_n) + if config in ((256, 256, 16, 64), (256, 256, 16, 80)): + return 64 + if config == (256, 256, 32, 64): + return 128 + if config == (256, 256, 48, 64): + return 192 + if config in ((256, 256, 64, 48), (256, 256, 64, 64)): + return 256 + if config == (256, 256, 96, 48): + return 192 + return 128 + + def _uses_split_pv_warps(self) -> bool: + """Whether dedicated QK/PV warp sets exchange P through SMEM.""" + config = ( + self.tile_hdim, + self.tile_hdimv, + self.tile_m, + self.tile_n, + self.num_threads, + ) + if self.paged_kv: + return config in ( + (256, 256, 16, 64, 64), + (256, 256, 16, 64, 128), + ) + return config in ( + (256, 256, 16, 64, 64), + (256, 256, 16, 80, 64), + (256, 256, 32, 64, 128), + (256, 256, 48, 64, 192), + (256, 256, 64, 48, 256), + (256, 256, 64, 64, 256), + ) + + def _q_in_regs_pipeline(self) -> bool: + """Whether Q remains resident while its SMEM allocation holds P.""" + if self.paged_kv: + return False + config = ( + self.tile_hdim, + self.tile_hdimv, + self.tile_m, + self.tile_n, + self.num_threads, + ) + return config in ( + (256, 256, 16, 64, 64), + (256, 256, 16, 80, 64), + (256, 256, 32, 64, 128), + (256, 256, 48, 64, 192), + (256, 256, 64, 48, 256), + (256, 256, 64, 64, 256), + ) + + def _num_k_stages(self) -> int: + return self.num_stages + + def _num_v_stages(self) -> int: + return 1 if self._uses_split_pv_warps() else self.num_stages + + def _num_p_stages(self) -> int: + return 2 + + def _num_softmax_stat_rows(self) -> int: + if self._uses_n_distributed_qk(): + num_qk_warps = self.num_qk_threads // cute.arch.WARP_SIZE + return 3 * num_qk_warps + 4 + return 4 + + def _uses_n_distributed_qk(self) -> bool: + return self.split_qk_n + + def _num_dma_threads(self) -> int: + return ( + self.num_dma_threads + if self.paged_kv or not self._uses_split_pv_warps() + else 0 + ) + + @staticmethod + def can_implement( + dtype, + head_dim, + head_dim_v, + tile_m, + tile_n, + num_stages, + num_threads, + is_causal, + Q_in_regs=False, + paged_kv=False, + ) -> bool: + """Check the constraints of the dedicated SM120 TMA kernel.""" + if dtype not in [cutlass.Float16, cutlass.BFloat16]: + return False + if head_dim % 8 != 0 or head_dim_v % 8 != 0: + return False + if tile_n % 16 != 0: + return False + if num_stages != FlashAttentionForwardSm120.get_fwd_num_stages( + head_dim, head_dim_v, tile_m, tile_n + ): + return False + if num_threads != FlashAttentionForwardSm120.get_fwd_num_threads( + head_dim, head_dim_v, tile_m, tile_n, paged_kv=paged_kv + ): + return False + if Q_in_regs: + return False + smem_usage = FlashAttentionForwardSm120._smem_usage_in_bytes( + head_dim, + head_dim_v, + tile_m, + tile_n, + num_stages, + Q_in_regs, + ) + if smem_usage > utils_basic.get_smem_capacity_in_bytes("sm_120"): + return False + if paged_kv and ( + head_dim, + head_dim_v, + tile_m, + tile_n, + num_threads, + ) in ( + (256, 256, 16, 64, 64), + (256, 256, 16, 64, 128), + ): + return True + if (head_dim, head_dim_v, tile_m, tile_n, num_threads) in ( + (256, 256, 16, 64, 64), + (256, 256, 16, 80, 64), + (256, 256, 32, 64, 128), + (256, 256, 48, 64, 192), + (256, 256, 64, 48, 256), + (256, 256, 64, 64, 256), + ): + return True + return (tile_m * 2) % num_threads == 0 + + def _get_smem_layout_atom(self): + sQ_layout_atom = self._make_smem_layout_atom( + self.dtype, self.tile_hdim, is_k_major=True + ) + sK_layout_atom = sQ_layout_atom + sV_layout_atom = self._make_smem_layout_atom( + self.dtype, self.tile_hdimv, is_k_major=True + ) + sO_layout_atom = sV_layout_atom + return sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, None + + def _setup_attributes(self): + super()._setup_attributes() + if const_expr(self._uses_split_pv_warps()): + sK_layout_atom = self._make_smem_layout_atom( + self.dtype, self.tile_hdim, is_k_major=True + ) + self.sK_layout = cute.tile_to_shape( + sK_layout_atom, + (self.tile_n, self.tile_hdim, self._num_k_stages()), + (0, 1, 2), + ) + sV_layout_atom = self._make_smem_layout_atom( + self.dtype, self.tile_hdimv, is_k_major=False + ) + self.sV_layout = cute.tile_to_shape( + sV_layout_atom, + (self.tile_hdimv, self.tile_n, self._num_v_stages()), + (1, 0, 2), + ) + self.sP_layout = None + if const_expr(self._uses_split_pv_warps()): + sP_layout_atom = self._make_smem_layout_atom( + self.dtype, self.tile_n, is_k_major=True + ) + self.sP_layout = cute.tile_to_shape( + sP_layout_atom, + (self.tile_m, self.tile_n, self._num_p_stages()), + (0, 1, 2), + ) + if const_expr(self.has_bias): + sBias_layout_atom = self._make_smem_layout_atom( + self.dtype, self.tile_n, is_k_major=True + ) + self.sBias_layout = cute.tile_to_shape( + sBias_layout_atom, + (self.bias_block_size, self.tile_n, self.num_stages), + (0, 1, 2), + ) + else: + self.sBias_layout = None + + @staticmethod + def _make_smem_layout_atom( + dtype: type[cutlass.Numeric], + major_dim: int, + *, + is_k_major: bool, + ) -> cute.ComposedLayout: + """Build a TMA-compatible SMEM layout for SM120 warp MMA.""" + major_mode_bits = const_expr(major_dim * dtype.width) + if const_expr(major_mode_bits % 1024 == 0): + contiguous_bits, swizzle_bits = 1024, 3 + elif const_expr(major_mode_bits % 512 == 0): + contiguous_bits, swizzle_bits = 512, 2 + elif const_expr(major_mode_bits % 256 == 0): + contiguous_bits, swizzle_bits = 256, 1 + else: + contiguous_bits, swizzle_bits = 128, 0 + contiguous_elems = const_expr(contiguous_bits // dtype.width) + layout = ( + cute.make_layout( + (8, contiguous_elems), + stride=(contiguous_elems, 1), + ) + if const_expr(is_k_major) + else cute.make_layout( + (contiguous_elems, 8), + stride=(1, contiguous_elems), + ) + ) + return cute.make_composed_layout( + cute.make_swizzle(swizzle_bits, 4, 3), + 0, + layout, + ) + + def _get_tiled_mma(self): + split_pv_warps = self._uses_split_pv_warps() + num_qk_warps = ( + self.num_threads // cute.arch.WARP_SIZE + if split_pv_warps and self._uses_n_distributed_qk() + else ( + self.tile_m // 16 + if split_pv_warps + else self.num_threads // cute.arch.WARP_SIZE + ) + ) + num_pv_warps_m = self.tile_m // 16 if split_pv_warps else self.num_threads // 32 + num_pv_warps_n = ( + self.num_threads // cute.arch.WARP_SIZE + if split_pv_warps and self.paged_kv + else 1 + ) + tiled_mma_qk = ( + cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)), + (1, num_qk_warps, 1), + permutation_mnk=(self.tile_m, self.tile_n, 16), + ) + if self.split_qk_n + else cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)), + (num_qk_warps, 1, 1), + permutation_mnk=(num_qk_warps * 16, 16, 16), + ) + ) + tiled_mma_pv = cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)), + (num_pv_warps_m, num_pv_warps_n, 1), + permutation_mnk=( + num_pv_warps_m * 16, + num_pv_warps_n * 16, + 16, + ), + ) + return tiled_mma_qk, tiled_mma_pv + + def _get_shared_storage_cls(self): + sQ_struct, sK_struct, sV_struct = [ + cute.struct.Align[ + cute.struct.MemRange[self.dtype, cute.cosize(layout)], 1024 + ] + for layout in (self.sQ_layout, self.sK_layout, self.sV_layout) + ] + mbar_Q_struct = cute.struct.MemRange[cutlass.Int64, 2] + mbar_K_struct = cute.struct.MemRange[cutlass.Int64, self._num_k_stages() * 2] + mbar_V_struct = cute.struct.MemRange[cutlass.Int64, self._num_v_stages() * 2] + mbar_P_struct = cute.struct.MemRange[ + cutlass.Int64, + 2 * self._num_p_stages() if self._uses_split_pv_warps() else 0, + ] + mbar_final_struct = cute.struct.MemRange[ + cutlass.Int64, 2 if self._uses_split_pv_warps() else 0 + ] + num_stats = ( + self._num_softmax_stat_rows() * self.tile_m + if self._uses_split_pv_warps() + else 0 + ) + softmax_stats_struct = cute.struct.MemRange[Float32, num_stats] + num_p_elements = ( + 0 + if self._q_in_regs_pipeline() or not self._uses_split_pv_warps() + else cute.cosize(self.sP_layout) + ) + sP_struct = cute.struct.Align[ + cute.struct.MemRange[ + self.dtype, + num_p_elements, + ], + 1024, + ] + mbar_Bias_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2] + sBias_struct = cute.struct.Align[ + cute.struct.MemRange[ + self.dtype, + cute.cosize(self.sBias_layout) if const_expr(self.has_bias) else 0, + ], + 1024, + ] + + @cute.struct + class SharedStorage: + mbar_Q: mbar_Q_struct + mbar_K: mbar_K_struct + mbar_V: mbar_V_struct + mbar_P: mbar_P_struct + mbar_final: mbar_final_struct + softmax_stats: softmax_stats_struct + sP: sP_struct + sV: sV_struct + sQ: sQ_struct + sK: sK_struct + + @cute.struct + class SharedStorageBias: + mbar_Q: mbar_Q_struct + mbar_K: mbar_K_struct + mbar_V: mbar_V_struct + mbar_P: mbar_P_struct + mbar_final: mbar_final_struct + mbar_Bias: mbar_Bias_struct + softmax_stats: softmax_stats_struct + sP: sP_struct + sV: sV_struct + sQ: sQ_struct + sK: sK_struct + sBias: sBias_struct + + return SharedStorageBias if const_expr(self.has_bias) else SharedStorage + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + softmax_scale: Float32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mPageTable: Optional[cute.Tensor] = None, + window_size_left: Int32 | int | None = None, + window_size_right: Int32 | int | None = None, + learnable_sink: Optional[cute.Tensor] = None, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + aux_data: AuxData = AuxData(), + mBias: Optional[cute.Tensor] = None, + launch_split_combine_early: Int32 = Int32(0), + stream: cuda.CUstream = None, + ): + assert blocksparse_tensors is None, "Block sparsity is not supported on SM120" + assert (mBias is not None) == self.has_bias + assert ( + mPageTable is None or self.paged_kv + ), "SM120 paged KV requires the dedicated DMA-warp specialization" + self._check_type( + *( + t.element_type if t is not None else None + for t in ( + mQ, + mK, + mV, + mO, + mLSE, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + ) + ) + ) + tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma() + self.num_qk_threads = tiled_mma_qk.size + self.num_mma_threads = tiled_mma_pv.size + self.num_producer_threads = self.num_threads + self.num_Q_load_threads = ( + self.num_qk_threads if self._uses_split_pv_warps() else self.num_threads + ) + self.num_epilogue_threads = ( + self.num_mma_threads if self._uses_split_pv_warps() else self.num_threads + ) + self.use_tma_O = False + self._setup_attributes() + SharedStorage = self._get_shared_storage_cls() + + mQ, mK, mV, mO = [assume_tensor_aligned(t) for t in (mQ, mK, mV, mO)] + if const_expr(mBias is not None): + assert mBias.element_type == self.dtype + mBias = assume_tensor_aligned(mBias) + Q_layout_transpose = ( + [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + ) + KV_layout_transpose = ( + [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1] + ) + mQ = cute.make_tensor( + mQ.iterator, cute.select(mQ.layout, mode=Q_layout_transpose) + ) + if const_expr(mBias is not None): + mBias = cute.make_tensor( + mBias.iterator, + cute.select(mBias.layout, mode=Q_layout_transpose), + ) + mK, mV = [ + cute.make_tensor( + t.iterator, cute.select(t.layout, mode=KV_layout_transpose) + ) + for t in (mK, mV) + ] + if const_expr(mPageTable is None): + V_layout_transpose = ( + [1, 0, 2, 3] if const_expr(mCuSeqlensK is None) else [1, 0, 2] + ) + mV = cute.make_tensor( + mV.iterator, cute.select(mV.layout, mode=V_layout_transpose) + ) + if const_expr(self.is_split_kv): + num_splits = mO.shape[0] + O_layout_transpose = ( + [2, 4, 3, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 3, 2, 0] + ) + LSE_layout_transpose = ( + [3, 2, 1, 0] if const_expr(mCuSeqlensQ is None) else [2, 1, 0] + ) + else: + num_splits = Int32(1) + O_layout_transpose = ( + [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + ) + LSE_layout_transpose = ( + [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0] + ) + mO = cute.make_tensor( + mO.iterator, cute.select(mO.layout, mode=O_layout_transpose) + ) + if const_expr(mLSE is not None): + mLSE = cute.make_tensor( + mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose) + ) + if const_expr(self.pack_gqa): + nheads_kv = mK.shape[2] + mQ = pack_gqa_layout(mQ, self.qhead_per_kvhead, nheads_kv, head_idx=2) + mO = pack_gqa_layout(mO, self.qhead_per_kvhead, nheads_kv, head_idx=2) + if const_expr(mLSE is not None): + mLSE = pack_gqa_layout( + mLSE, self.qhead_per_kvhead, nheads_kv, head_idx=1 + ) + if const_expr(mBias is not None): + mBias = pack_gqa_layout( + mBias, self.qhead_per_kvhead, nheads_kv, head_idx=2 + ) + + if const_expr(mPageTable is None or self.paged_tma): + tma_copy_op = cpasync.CopyBulkTensorTileG2SOp() + mV_tma = ( + mV + if const_expr(mPageTable is None) + else cute.make_tensor( + mV.iterator, + cute.select(mV.layout, mode=[1, 0, 2, 3]), + ) + ) + tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom( + tma_copy_op, + mK, + cute.select(self.sK_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdim), + 1, + ) + tma_atom_V, tma_tensor_V = cpasync.make_tiled_tma_atom( + tma_copy_op, + mV_tma, + cute.select(self.sV_layout, mode=[0, 1]), + (self.tile_hdimv, self.tile_n), + 1, + ) + self.tma_copy_bytes_K = cute.size_in_bytes( + mK.element_type, cute.select(self.sK_layout, mode=[0, 1]) + ) + self.tma_copy_bytes_V = cute.size_in_bytes( + mV_tma.element_type, + cute.select(self.sV_layout, mode=[0, 1]), + ) + else: + tma_atom_K = None + tma_atom_V = None + tma_tensor_K = mK + tma_tensor_V = mV + if const_expr(self.has_bias): + tma_atom_Bias, tma_tensor_Bias = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mBias, + cute.select(self.sBias_layout, mode=[0, 1]), + (self.bias_block_size, self.tile_n), + 1, + ) + self.tma_copy_bytes_Bias = cute.size_in_bytes( + mBias.element_type, + cute.select(self.sBias_layout, mode=[0, 1]), + ) + else: + tma_atom_Bias = None + tma_tensor_Bias = mBias + + is_varlen = const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None) + num_batch = ( + mCuSeqlensQ.shape[0] - 1 + if const_expr(mCuSeqlensQ is not None) + else cute.size(mQ.shape[3]) + ) + TileScheduler = ( + Sm120UniformBatchScheduler + if is_varlen and self.direct_uniform_batch + else SingleTileVarlenScheduler if is_varlen else SingleTileScheduler + ) + tile_sched_args = TileSchedulerArguments( + num_block=cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m), + num_head=cute.size(mQ.shape[2]), + num_batch=num_batch, + num_splits=num_splits, + seqlen_k=0, + headdim=mQ.shape[1], + headdim_v=mO.shape[1], + total_q=( + cute.size(mQ.shape[0]) + if const_expr(mCuSeqlensQ is not None) + else cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]) + ), + tile_shape_mn=(self.tile_m, self.tile_n), + lpt=(self.is_causal or self.is_local) and not self.direct_uniform_batch, + qhead_per_kvhead_packgqa=( + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1 + ), + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + is_persistent=False, + is_split_kv=self.is_split_kv, + ) + tile_sched_params = TileScheduler.to_underlying_arguments( + tile_sched_args, + scheduling_mode=SchedulingMode.STATIC, + ) + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + if const_expr(self.has_bias): + base_softmax_scale = softmax_scale + softmax_scale_log2, softmax_scale = utils.LOG2_E, None + else: + base_softmax_scale = None + softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2( + softmax_scale, self.score_mod + ) + window_size_left = ( + Int32(window_size_left) if window_size_left is not None else None + ) + window_size_right = ( + Int32(window_size_right) if window_size_right is not None else None + ) + fastdiv_mods = utils.compute_fastdiv_mods( + mQ, + mK, + self.qhead_per_kvhead, + self.pack_gqa, + aux_data.tensors, + mPageTable, + ) + + self.kernel( + mQ, + tma_tensor_K, + tma_tensor_V, + mO, + mLSE, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mPageTable, + tma_tensor_Bias, + tma_atom_K, + tma_atom_V, + tma_atom_Bias, + softmax_scale_log2, + softmax_scale, + base_softmax_scale, + window_size_left, + window_size_right, + learnable_sink, + self.sQ_layout, + self.sK_layout, + self.sV_layout, + self.sO_layout, + self.sP_layout, + self.sBias_layout, + self.gmem_tiled_copy_Q, + self.gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + SharedStorage, + tile_sched_params, + TileScheduler, + launch_split_combine_early, + aux_data, + fastdiv_mods, + ).launch( + grid=grid_dim, + block=[self.num_threads + self._num_dma_threads(), 1, 1], + smem=SharedStorage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + # The immediately preceding ShearingBias grid releases this + # dependent launch before all of its CTAs have retired. Bias + # readers synchronize before their first TMA issue below; the + # remaining kernel prologue can overlap the tail of the shear. + use_pdl=self.has_bias, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + mPageTable: Optional[cute.Tensor], + mBias: Optional[cute.Tensor], + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_V: Optional[cute.CopyAtom], + tma_atom_Bias: Optional[cute.CopyAtom], + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + base_softmax_scale: Optional[Float32], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + learnable_sink: Optional[cute.Tensor], + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + sV_layout: cute.ComposedLayout, + sO_layout: cute.ComposedLayout, + sP_layout: cute.ComposedLayout | None, + sBias_layout: cute.ComposedLayout | None, + gmem_tiled_copy_Q: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + SharedStorage: cutlass.Constexpr, + tile_sched_params, + TileScheduler: cutlass.Constexpr[Callable], + launch_split_combine_early: Int32, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + ): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + if const_expr(self.is_split_kv): + if launch_split_combine_early != 0 and warp_idx == 0: + cute.arch.griddepcontrol_launch_dependents() + + if const_expr(mPageTable is None or self.paged_tma): + if warp_idx == 0: + cpasync.prefetch_descriptor(tma_atom_K) + cpasync.prefetch_descriptor(tma_atom_V) + if const_expr(self.has_bias): + if warp_idx == 0: + cpasync.prefetch_descriptor(tma_atom_Bias) + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner) + sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) + sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner) + sO = storage.sQ.get_tensor(sO_layout.outer, swizzle=sO_layout.inner) + sBias = ( + storage.sBias.get_tensor(sBias_layout.outer, swizzle=sBias_layout.inner) + if const_expr(self.has_bias) + else None + ) + sP = None + sRowScale = None + sFinalScale = None + sLSE = None + if const_expr(sP_layout is not None): + if const_expr(self._q_in_regs_pipeline()): + sP = storage.sQ.get_tensor(sP_layout.outer, swizzle=sP_layout.inner) + else: + sP = storage.sP.get_tensor(sP_layout.outer, swizzle=sP_layout.inner) + sSoftmaxStats = storage.softmax_stats.get_tensor( + cute.make_layout( + (self._num_softmax_stat_rows(), self.tile_m), + stride=(self.tile_m, 1), + ) + ) + sRowScale = sSoftmaxStats + if const_expr(self._uses_n_distributed_qk()): + num_qk_warps = self.num_qk_threads // cute.arch.WARP_SIZE + sFinalScale = sSoftmaxStats[2 * num_qk_warps + 2, None] + sLSE = sSoftmaxStats[3 * num_qk_warps + 3, None] + else: + sFinalScale = sSoftmaxStats[2, None] + sLSE = sSoftmaxStats[3, None] + + tma_group = CooperativeGroup(Agent.Thread) + qk_group = CooperativeGroup( + Agent.Thread, self.num_qk_threads // cute.arch.WARP_SIZE + ) + pv_group = CooperativeGroup( + Agent.Thread, self.num_mma_threads // cute.arch.WARP_SIZE + ) + if const_expr(mPageTable is None or self.paged_tma): + pipeline_k = PipelineTmaAsync.create( + num_stages=self._num_k_stages(), + producer_group=tma_group, + consumer_group=qk_group, + tx_count=self.tma_copy_bytes_K, + barrier_storage=storage.mbar_K.data_ptr(), + defer_sync=True, + ) + pipeline_v = PipelineTmaAsync.create( + num_stages=self._num_v_stages(), + producer_group=tma_group, + consumer_group=pv_group, + tx_count=self.tma_copy_bytes_V, + barrier_storage=storage.mbar_V.data_ptr(), + defer_sync=True, + ) + else: + dma_group = CooperativeGroup(Agent.Thread, self.num_dma_threads) + k_consumer_group = CooperativeGroup( + Agent.Thread, + ( + self.num_qk_threads + if self._uses_split_pv_warps() + else self.num_threads + ), + ) + v_consumer_group = CooperativeGroup( + Agent.Thread, + ( + self.num_mma_threads + if self._uses_split_pv_warps() + else self.num_threads + ), + ) + pipeline_k = pipeline_custom.PipelineCpAsync.create( + num_stages=self._num_k_stages(), + producer_group=dma_group, + consumer_group=k_consumer_group, + barrier_storage=storage.mbar_K.data_ptr(), + defer_sync=True, + ) + pipeline_v = pipeline_custom.PipelineCpAsync.create( + num_stages=self._num_v_stages(), + producer_group=dma_group, + consumer_group=v_consumer_group, + barrier_storage=storage.mbar_V.data_ptr(), + defer_sync=True, + ) + pipeline_bias = ( + PipelineTmaAsync.create( + num_stages=self.num_stages, + producer_group=tma_group, + consumer_group=qk_group, + tx_count=self.tma_copy_bytes_Bias, + barrier_storage=storage.mbar_Bias.data_ptr(), + defer_sync=True, + ) + if const_expr(self.has_bias) + else None + ) + pipeline_p = None + pipeline_final = None + if const_expr(sP_layout is not None): + pipeline_p = PipelineAsync.create( + num_stages=self._num_p_stages(), + producer_group=CooperativeGroup(Agent.Thread, self.num_qk_threads), + consumer_group=CooperativeGroup(Agent.Thread, self.num_mma_threads), + barrier_storage=storage.mbar_P.data_ptr(), + defer_sync=True, + name="sm120_p", + ) + pipeline_final = PipelineAsync.create( + num_stages=1, + producer_group=CooperativeGroup(Agent.Thread, self.num_qk_threads), + consumer_group=CooperativeGroup(Agent.Thread, self.num_mma_threads), + barrier_storage=storage.mbar_final.data_ptr(), + defer_sync=True, + name="sm120_final", + ) + tile_scheduler = TileScheduler.create(tile_sched_params) + work_tile = tile_scheduler.initial_work_tile_info() + if work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoQK.create( + batch_idx=batch_idx, + seqlen_q_static=( + mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1] + ), + seqlen_k_static=( + mK.shape[0] + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1] + ), + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + tile_m=self.tile_m, + tile_n=self.tile_n, + ) + run_mainloop = True + if const_expr(self.is_split_kv): + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + self.is_split_kv, + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=( + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1 + ), + ) + n_block_min, n_block_max = self._get_n_block_min_max( + block_info, + seqlen, + m_block, + split_idx, + tile_scheduler.params.num_splits, + ) + is_empty_split = n_block_min >= n_block_max + if is_empty_split: + if const_expr(self._uses_split_pv_warps()): + if warp_idx >= self.num_qk_threads // cute.arch.WARP_SIZE: + self.epilogue_empty_split( + mO, + mLSE, + learnable_sink, + seqlen, + tiled_mma_pv, + tidx - self.num_qk_threads, + m_block, + head_idx, + batch_idx, + split_idx, + ) + else: + if 0 < warp_idx <= self.num_threads // cute.arch.WARP_SIZE: + self.epilogue_empty_split( + mO, + mLSE, + learnable_sink, + seqlen, + tiled_mma_pv, + tidx - self.num_dma_threads, + m_block, + head_idx, + batch_idx, + split_idx, + ) + run_mainloop = not is_empty_split + + if run_mainloop: + pipeline_init_arrive(cluster_shape_mn=(1, 1), is_relaxed=True) + pipeline_init_wait(cluster_shape_mn=(1, 1)) + + if const_expr(mPageTable is not None and self._uses_split_pv_warps()): + if warp_idx == 0: + self.load_paged_persistent( + mK, + mV, + mPageTable, + tma_atom_K, + tma_atom_V, + sK, + sV, + pipeline_k, + pipeline_v, + tile_scheduler, + mQ, + mCuSeqlensQ, + mSeqUsedQ, + mSeqUsedK, + window_size_left, + window_size_right, + tidx, + ) + elif const_expr(self._uses_n_distributed_qk()): + if warp_idx <= self.num_mma_threads // cute.arch.WARP_SIZE: + self.mma_persistent( + mQ, + mK, + mO, + mLSE, + sQ, + sK, + sV, + sO, + sP, + sRowScale, + sLSE, + learnable_sink, + pipeline_k, + pipeline_v, + gmem_tiled_copy_Q, + gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + tidx - self.num_dma_threads, + softmax_scale_log2, + softmax_scale, + base_softmax_scale, + tile_scheduler, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mPageTable, + window_size_left, + window_size_right, + True, + aux_data, + fastdiv_mods, + ) + elif warp_idx == 1: + self.mma_persistent( + mQ, + mK, + mO, + mLSE, + sQ, + sK, + sV, + sO, + sP, + sRowScale, + sLSE, + learnable_sink, + pipeline_k, + pipeline_v, + gmem_tiled_copy_Q, + gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + tidx - self.num_dma_threads, + softmax_scale_log2, + softmax_scale, + base_softmax_scale, + tile_scheduler, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mPageTable, + window_size_left, + window_size_right, + True, + aux_data, + fastdiv_mods, + ) + elif warp_idx <= self.num_mma_threads // cute.arch.WARP_SIZE: + self.mma_persistent( + mQ, + mK, + mO, + mLSE, + sQ, + sK, + sV, + sO, + sP, + sRowScale, + sLSE, + learnable_sink, + pipeline_k, + pipeline_v, + gmem_tiled_copy_Q, + gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + tidx - self.num_dma_threads, + softmax_scale_log2, + softmax_scale, + base_softmax_scale, + tile_scheduler, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mPageTable, + window_size_left, + window_size_right, + False, + aux_data, + fastdiv_mods, + ) + elif const_expr(self._uses_split_pv_warps()): + if warp_idx < self.num_qk_threads // cute.arch.WARP_SIZE: + self.mma_qk_pipeline_persistent( + mQ, + mK, + mV, + sQ, + sK, + sV, + tma_atom_K, + tma_atom_V, + sP, + sRowScale, + sFinalScale, + sLSE, + learnable_sink, + pipeline_k, + pipeline_v, + pipeline_p, + pipeline_final, + gmem_tiled_copy_Q, + tiled_mma_qk, + tidx, + softmax_scale_log2, + softmax_scale, + tile_scheduler, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + window_size_left, + window_size_right, + aux_data, + fastdiv_mods, + ) + else: + self.mma_pv_pipeline_persistent( + mQ, + mK, + mO, + mLSE, + sV, + sO, + sP, + sRowScale, + sFinalScale, + sLSE, + pipeline_v, + pipeline_p, + pipeline_final, + gmem_tiled_copy_O, + tiled_mma_pv, + tidx - self.num_qk_threads, + tile_scheduler, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + window_size_left, + window_size_right, + ) + elif warp_idx == 0: + if const_expr(mPageTable is None): + self.load_tma_persistent( + mK, + mV, + sK, + sV, + tma_atom_K, + tma_atom_V, + pipeline_k, + pipeline_v, + tile_scheduler, + mQ, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + window_size_left, + window_size_right, + mBias=mBias, + sBias=sBias, + tma_atom_Bias=tma_atom_Bias, + pipeline_bias=pipeline_bias, + ) + else: + self.load_paged_persistent( + mK, + mV, + mPageTable, + tma_atom_K, + tma_atom_V, + sK, + sV, + pipeline_k, + pipeline_v, + tile_scheduler, + mQ, + mCuSeqlensQ, + mSeqUsedQ, + mSeqUsedK, + window_size_left, + window_size_right, + tidx, + mBias=mBias, + sBias=sBias, + tma_atom_Bias=tma_atom_Bias, + pipeline_bias=pipeline_bias, + ) + elif warp_idx <= self.num_threads // cute.arch.WARP_SIZE: + self.mma_persistent( + mQ, + mK, + mO, + mLSE, + sQ, + sK, + sV, + sO, + sP, + sRowScale, + sLSE, + learnable_sink, + pipeline_k, + pipeline_v, + gmem_tiled_copy_Q, + gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + tidx - self.num_dma_threads, + softmax_scale_log2, + softmax_scale, + base_softmax_scale, + tile_scheduler, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mPageTable, + window_size_left, + window_size_right, + True, + aux_data, + fastdiv_mods, + sBias=sBias, + pipeline_bias=pipeline_bias, + ) + + @cute.jit + def epilogue_empty_split( + self, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + learnable_sink: Optional[cute.Tensor], + seqlen: SeqlenInfoQK, + tiled_mma_pv: cute.TiledMma, + tidx: Int32, + m_block: Int32, + head_idx: Int32, + batch_idx: Int32, + split_idx: Int32, + ): + """Write the reduction identity for a split with no visible K/V tile.""" + thr_mma_pv = tiled_mma_pv.get_slice(tidx) + cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) + tOcO_mn = layout_utils.reshape_acc_to_mn(thr_mma_pv.partition_C(cO)) + qhead_pack = self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1 + row_limit = seqlen.seqlen_q * qhead_pack + row_offset = ( + seqlen.offset_q * qhead_pack + if const_expr(seqlen.has_cu_seqlens_q) + else Int32(0) + ) + + if const_expr(mLSE is not None): + if tOcO_mn[0][1] == 0: + for r in cutlass.range(cute.size(tOcO_mn, mode=[0]), unroll_full=True): + row = m_block * self.tile_m + tOcO_mn[r, 0][0] + if row < row_limit: + lse = -Float32.inf + if const_expr(learnable_sink is not None): + if split_idx == 0: + q_head_idx = ( + row % self.qhead_per_kvhead + + head_idx * self.qhead_per_kvhead + if const_expr(self.pack_gqa) + else head_idx + ) + lse = Float32(learnable_sink[q_head_idx]) + if const_expr(seqlen.has_cu_seqlens_q): + mLSE[row_offset + row, head_idx, split_idx] = lse + else: + mLSE[row, head_idx, batch_idx, split_idx] = lse + + for r in cutlass.range(cute.size(tOcO_mn, mode=[0]), unroll_full=True): + row = m_block * self.tile_m + tOcO_mn[r, 0][0] + if row < row_limit: + for c in cutlass.range(cute.size(tOcO_mn, mode=[1]), unroll_full=True): + col = tOcO_mn[r, c][1] + if const_expr(not self.check_hdim_v_oob) or col < mO.shape[1]: + if const_expr(seqlen.has_cu_seqlens_q): + mO[row_offset + row, col, head_idx, split_idx] = ( + mO.element_type(0.0) + ) + else: + mO[row, col, head_idx, batch_idx, split_idx] = ( + mO.element_type(0.0) + ) + + @cute.jit + def _get_bias_load_info( + self, + block_info: BlockInfo, + seqlen: SeqlenInfoQK, + m_block: Int32, + n_block_min: Int32, + n_block_max: Int32, + ): + """Map the split-local right edge to the pre-sheared bias blocks.""" + _, n_block_max_abs = block_info.get_n_block_min_max( + seqlen, m_block, absolute=True + ) + bias_idx_offset = n_block_max_abs - n_block_max + bias_max_idx = self.bias_n_max - 1 - bias_idx_offset + num_bias_loads = min( + self.bias_n_max - bias_idx_offset, + n_block_max - n_block_min, + ) + return bias_max_idx, num_bias_loads + + @cute.jit + def load_tma_persistent( + self, + mK: cute.Tensor, + mV: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + tma_atom_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + tile_scheduler: TileSchedulerProtocol, + mQ: cute.Tensor, + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + mBias: Optional[cute.Tensor] = None, + sBias: Optional[cute.Tensor] = None, + tma_atom_Bias: Optional[cute.CopyAtom] = None, + pipeline_bias: Optional[PipelineAsync] = None, + ): + producer_state_k = PipelineState( + self._num_k_stages(), Int32(0), Int32(0), Int32(1) + ) + producer_state_v = PipelineState( + self._num_v_stages(), Int32(0), Int32(0), Int32(1) + ) + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + self.is_split_kv, + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=( + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1 + ), + ) + work_tile = tile_scheduler.initial_work_tile_info() + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoQK.create( + batch_idx=batch_idx, + seqlen_q_static=( + mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1] + ), + seqlen_k_static=mK.shape[0], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + tile_m=self.tile_m, + tile_n=self.tile_n, + ) + n_block_min, n_block_max = self._get_n_block_min_max( + block_info, + seqlen, + m_block, + split_idx, + tile_scheduler.params.num_splits, + ) + head_idx_kv = ( + head_idx if const_expr(self.pack_gqa) else head_idx // self.qhead_per_kvhead + ) + bias_max_idx, num_bias_loads = Int32(0), Int32(0) + if const_expr(self.has_bias): + bias_max_idx, num_bias_loads = self._get_bias_load_info( + block_info, + seqlen, + m_block, + n_block_min, + n_block_max, + ) + producer_state_k, producer_state_v = self.load_tma( + mK, + mV, + sK, + sV, + tma_atom_K, + tma_atom_V, + pipeline_k, + pipeline_v, + producer_state_k, + producer_state_v, + seqlen, + n_block_min, + n_block_max, + head_idx_kv, + batch_idx, + mBias, + sBias, + tma_atom_Bias, + pipeline_bias, + m_block, + head_idx, + bias_max_idx, + num_bias_loads, + ) + + pipeline_k.producer_tail(producer_state_k) + pipeline_v.producer_tail(producer_state_v) + + @cute.jit + def load_paged_persistent( + self, + mK: cute.Tensor, + mV: cute.Tensor, + mPageTable: cute.Tensor, + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_V: Optional[cute.CopyAtom], + sK: cute.Tensor, + sV: cute.Tensor, + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + tile_scheduler: TileSchedulerProtocol, + mQ: cute.Tensor, + mCuSeqlensQ: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + tidx: Int32, + mBias: Optional[cute.Tensor] = None, + sBias: Optional[cute.Tensor] = None, + tma_atom_Bias: Optional[cute.CopyAtom] = None, + pipeline_bias: Optional[PipelineAsync] = None, + ): + producer_state_k = PipelineState( + self._num_k_stages(), Int32(0), Int32(0), Int32(1) + ) + producer_state_v = PipelineState( + self._num_v_stages(), Int32(0), Int32(0), Int32(1) + ) + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + self.is_split_kv, + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=( + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1 + ), + ) + work_tile = tile_scheduler.initial_work_tile_info() + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoQK.create( + batch_idx=batch_idx, + seqlen_q_static=( + mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1] + ), + seqlen_k_static=mK.shape[0] * mPageTable.shape[1], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=None, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + tile_m=self.tile_m, + tile_n=self.tile_n, + ) + n_block_min, n_block_max = self._get_n_block_min_max( + block_info, + seqlen, + m_block, + split_idx, + tile_scheduler.params.num_splits, + ) + head_idx_kv = ( + head_idx if const_expr(self.pack_gqa) else head_idx // self.qhead_per_kvhead + ) + bias_max_idx, num_bias_loads = Int32(0), Int32(0) + if const_expr(self.has_bias): + bias_max_idx, num_bias_loads = self._get_bias_load_info( + block_info, + seqlen, + m_block, + n_block_min, + n_block_max, + ) + mBias_cur = seqlen.offset_batch_Q(mBias, batch_idx, dim=3)[ + None, None, head_idx + ] + gBias = cute.local_tile( + mBias_cur, + (self.bias_block_size, self.tile_n), + (None, None), + ) + tBsBias, tBgBias = cpasync.tma_partition( + tma_atom_Bias, + 0, + cute.make_layout(1), + cute.group_modes(sBias, 0, 2), + cute.group_modes(gBias, 0, 2), + ) + # Worktiles outside the materialized relative-bias band issue no + # bias TMA and therefore need not wait on the shear producer. + if num_bias_loads > 0: + cute.arch.griddepcontrol_wait() + if const_expr(self.paged_tma): + mK_cur = mK[None, None, head_idx_kv, None] + mV_cur = mV[None, None, head_idx_kv, None] + gK = cute.local_tile( + mK_cur, + (self.tile_n, self.tile_hdim), + (0, 0, None), + ) + gV = cute.local_tile( + mV_cur, + (self.tile_hdimv, self.tile_n), + (0, 0, None), + ) + copy_K, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_K, + 0, + cute.make_layout(1), + gK, + sK, + ) + copy_V, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_V, + 0, + cute.make_layout(1), + gV, + sV, + ) + num_n_blocks = cutlass.max(n_block_max - n_block_min, 1) + for n_tile in cutlass.range(num_n_blocks, unroll=1): + n_block = cutlass.max(n_block_max - 1 - n_tile, n_block_min) + page_idx = mPageTable[batch_idx, n_block] + + if const_expr(self.has_bias): + if n_tile < num_bias_loads: + pipeline_bias.producer_acquire(producer_state_k) + cute.copy( + tma_atom_Bias, + tBgBias[None, m_block, bias_max_idx - n_tile], + tBsBias[None, producer_state_k.index], + tma_bar_ptr=pipeline_bias.producer_get_barrier( + producer_state_k + ), + ) + pipeline_bias.producer_commit(producer_state_k) + pipeline_k.producer_acquire(producer_state_k) + copy_K( + src_idx=page_idx, + dst_idx=producer_state_k.index, + tma_bar_ptr=pipeline_k.producer_get_barrier(producer_state_k), + ) + pipeline_k.producer_commit(producer_state_k) + producer_state_k.advance() + + pipeline_v.producer_acquire(producer_state_v) + copy_V( + src_idx=page_idx, + dst_idx=producer_state_v.index, + tma_bar_ptr=pipeline_v.producer_get_barrier(producer_state_v), + ) + pipeline_v.producer_commit(producer_state_v) + producer_state_v.advance() + + pipeline_k.producer_tail(producer_state_k) + pipeline_v.producer_tail(producer_state_v) + return + + paged_kv_manager = Sm120PagedKVManager.create( + mPageTable, + mK, + mV, + FastDivmodDivisor(mK.shape[0]), + batch_idx, + head_idx_kv, + tidx, + seqlen.seqlen_k, + 0, + self.tile_n, + self.tile_hdim, + self.tile_hdimv, + self.num_dma_threads, + mK.element_type, + ) + num_n_blocks = cutlass.max(n_block_max - n_block_min, 1) + for n_tile in cutlass.range(num_n_blocks, unroll=1): + n_block = cutlass.max(n_block_max - 1 - n_tile, n_block_min) + paged_kv_manager.load_page_table(n_block) + + if const_expr(self.has_bias): + if n_tile < num_bias_loads: + pipeline_bias.producer_acquire(producer_state_k) + cute.copy( + tma_atom_Bias, + tBgBias[None, m_block, bias_max_idx - n_tile], + tBsBias[None, producer_state_k.index], + tma_bar_ptr=pipeline_bias.producer_get_barrier( + producer_state_k + ), + ) + pipeline_bias.producer_commit(producer_state_k) + pipeline_k.producer_acquire(producer_state_k) + paged_kv_manager.load_KV( + n_block, + sK[None, None, producer_state_k.index], + "K", + ) + cute.arch.cp_async_commit_group() + pipeline_k.producer_commit(producer_state_k) + producer_state_k.advance() + + pipeline_v.producer_acquire(producer_state_v) + sV_stage = layout_utils.transpose_view( + sV[None, None, producer_state_v.index] + ) + paged_kv_manager.load_KV(n_block, sV_stage, "V") + cute.arch.cp_async_commit_group() + pipeline_v.producer_commit(producer_state_v) + producer_state_v.advance() + + pipeline_k.producer_tail(producer_state_k) + pipeline_v.producer_tail(producer_state_v) + + @cute.jit + def mma_qk_pipeline_persistent( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + tma_atom_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + sP: cute.Tensor, + sRowScale: cute.Tensor, + sFinalScale: cute.Tensor, + sLSE: cute.Tensor, + learnable_sink: Optional[cute.Tensor], + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + pipeline_p: PipelineAsync, + pipeline_final: PipelineAsync, + gmem_tiled_copy_Q: cute.TiledCopy, + tiled_mma_qk: cute.TiledMma, + tidx: Int32, + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + tile_scheduler: TileSchedulerProtocol, + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + ): + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + self.is_split_kv, + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=( + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1 + ), + ) + work_tile = tile_scheduler.initial_work_tile_info() + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoQK.create( + batch_idx=batch_idx, + seqlen_q_static=( + mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1] + ), + seqlen_k_static=mK.shape[0], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + tile_m=self.tile_m, + tile_n=self.tile_n, + ) + n_block_min, n_block_max = self._get_n_block_min_max( + block_info, + seqlen, + m_block, + split_idx, + tile_scheduler.params.num_splits, + ) + head_idx_kv = ( + head_idx if const_expr(self.pack_gqa) else head_idx // self.qhead_per_kvhead + ) + self.mma_qk_pipeline( + mQ, + mK, + mV, + sQ, + sK, + sV, + tma_atom_K, + tma_atom_V, + sP, + sRowScale, + sFinalScale, + sLSE, + learnable_sink, + pipeline_k, + pipeline_v, + pipeline_p, + pipeline_final, + gmem_tiled_copy_Q, + tiled_mma_qk, + tidx, + softmax_scale_log2, + softmax_scale, + block_info, + seqlen, + n_block_min, + n_block_max, + m_block, + head_idx, + head_idx_kv, + batch_idx, + split_idx, + aux_data, + fastdiv_mods, + ) + + @cute.jit + def mma_pv_pipeline_persistent( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sV: cute.Tensor, + sO: cute.Tensor, + sP: cute.Tensor, + sRowScale: cute.Tensor, + sFinalScale: cute.Tensor, + sLSE: cute.Tensor, + pipeline_v: PipelineAsync, + pipeline_p: PipelineAsync, + pipeline_final: PipelineAsync, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma_pv: cute.TiledMma, + tidx: Int32, + tile_scheduler: TileSchedulerProtocol, + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + ): + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + self.is_split_kv, + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=( + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1 + ), + ) + work_tile = tile_scheduler.initial_work_tile_info() + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoQK.create( + batch_idx=batch_idx, + seqlen_q_static=( + mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1] + ), + seqlen_k_static=mK.shape[0], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + tile_m=self.tile_m, + tile_n=self.tile_n, + ) + n_block_min, n_block_max = self._get_n_block_min_max( + block_info, + seqlen, + m_block, + split_idx, + tile_scheduler.params.num_splits, + ) + self.mma_pv_pipeline( + mO, + mLSE, + sV, + sO, + sP, + sRowScale, + sFinalScale, + sLSE, + pipeline_v, + pipeline_p, + pipeline_final, + gmem_tiled_copy_O, + tiled_mma_pv, + tidx, + block_info, + seqlen, + n_block_min, + n_block_max, + m_block, + head_idx, + batch_idx, + split_idx, + ) + + @cute.jit + def _run_n_block_schedule( + self, + compute_one_n_block: Callable, + role_state, + block_info: BlockInfo, + seqlen: SeqlenInfoQK, + n_block_min: Int32, + n_block_max: Int32, + m_block: Int32, + mask_fn: Optional[Callable], + ): + if const_expr(mask_fn is not None): + mask_fn_seqlen = partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True) + mask_fn_no_seqlen = partial( + mask_fn, mask_mod=self.mask_mod, mask_seqlen=False + ) + else: + mask_fn_seqlen = None + mask_fn_no_seqlen = None + + n_block = cutlass.max(n_block_max - 1, 0) + role_state = compute_one_n_block( + n_block, + role_state, + mask_fn=mask_fn_seqlen, + is_first_n_block=True, + ) + n_block_upper = n_block + if const_expr(self.is_causal or self.is_local): + n_block_min_causal_local_mask = ( + block_info.get_n_block_min_causal_local_mask( + seqlen, m_block, n_block_min + ) + ) + for n_tile in cutlass.range( + n_block_max - 1 - n_block_min_causal_local_mask, unroll=1 + ): + n_block = n_block_max - 2 - n_tile + role_state = compute_one_n_block( + n_block, + role_state, + mask_fn=mask_fn_seqlen, + ) + n_block_upper = cutlass.min(n_block_upper, n_block_min_causal_local_mask) + n_block_min_before_local_mask = block_info.get_n_block_min_before_local_mask( + seqlen, m_block, n_block_min + ) + for n_tile in cutlass.range( + n_block_upper - n_block_min_before_local_mask, unroll=1 + ): + role_state = compute_one_n_block( + n_block_upper - n_tile - 1, + role_state, + mask_fn=mask_fn_no_seqlen, + ) + if const_expr(self.is_local and block_info.window_size_left is not None): + n_block_upper = cutlass.min(n_block_upper, n_block_min_before_local_mask) + for n_tile in cutlass.range(n_block_upper - n_block_min, unroll=1): + role_state = compute_one_n_block( + n_block_upper - n_tile - 1, + role_state, + mask_fn=mask_fn_no_seqlen, + ) + return role_state + + @cute.jit + def mma_qk_pipeline( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + tma_atom_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + sP: cute.Tensor, + sRowScale: cute.Tensor, + sFinalScale: cute.Tensor, + sLSE: cute.Tensor, + learnable_sink: Optional[cute.Tensor], + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + pipeline_p: PipelineAsync, + pipeline_final: PipelineAsync, + gmem_tiled_copy_Q: cute.TiledCopy, + tiled_mma_qk: cute.TiledMma, + tidx: Int32, + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + block_info: BlockInfo, + seqlen: SeqlenInfoQK, + n_block_min: Int32, + n_block_max: Int32, + m_block: Int32, + head_idx: Int32, + head_idx_kv: Int32, + batch_idx: Int32, + split_idx: Int32, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + ): + mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] + if const_expr(not self.pack_gqa): + gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0)) + if const_expr(not seqlen.has_cu_seqlens_k): + mK_cur = mK[None, None, head_idx_kv, batch_idx] + mV_cur = mV[None, None, head_idx_kv, batch_idx] + else: + mK_cur = cute.domain_offset( + (seqlen.offset_k, 0), mK[None, None, head_idx_kv] + ) + mV_cur = cute.domain_offset( + (0, seqlen.offset_k), mV[None, None, head_idx_kv] + ) + gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (None, 0)) + gV = cute.local_tile(mV_cur, (self.tile_hdimv, self.tile_n), (0, None)) + copy_K, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_K, 0, cute.make_layout(1), gK, sK + ) + copy_V, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_V, 0, cute.make_layout(1), gV, sV + ) + + thr_mma_qk = tiled_mma_qk.get_slice(tidx) + smem_copy_atom_qk = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), self.dtype + ) + smem_thr_copy_q = utils.make_tiled_copy_A( + smem_copy_atom_qk, tiled_mma_qk + ).get_slice(tidx) + smem_thr_copy_k = utils.make_tiled_copy_B( + smem_copy_atom_qk, tiled_mma_qk + ).get_slice(tidx) + tCrQ = None + if const_expr(self._q_in_regs_pipeline()): + tCrQ = thr_mma_qk.make_fragment_A(thr_mma_qk.partition_A(sQ)) + smem_store_atom_p = utils.get_smem_store_atom(120, self.dtype) + smem_thr_store_p = cute.make_tiled_copy_C( + smem_store_atom_p, tiled_mma_qk + ).get_slice(tidx) + tPsP_store = smem_thr_store_p.partition_D(sP) + + gmem_thr_copy_q = gmem_tiled_copy_Q.get_slice(tidx) + if const_expr(not self.pack_gqa): + self.load_Q( + gmem_thr_copy_q, + gQ, + sQ, + m_block, + seqlen=seqlen.seqlen_q, + headdim=mQ.shape[1], + ) + else: + PackGQA( + self.tile_m, + self.tile_hdim, + self.check_hdim_oob, + self.qhead_per_kvhead, + ).load_Q( + mQ_cur, + sQ, + gmem_tiled_copy_Q, + tidx, + m_block, + seqlen.seqlen_q, + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), + number_of_threads=self.num_Q_load_threads, + ) + if const_expr(self._q_in_regs_pipeline()): + tCsQ = smem_thr_copy_q.partition_S(sQ) + tCrQ_copy_view = smem_thr_copy_q.retile(tCrQ) + cute.copy(smem_thr_copy_q, tCsQ, tCrQ_copy_view) + # All QK warps must finish reading Q before its allocation becomes P. + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), + number_of_threads=self.num_Q_load_threads, + ) + + acc_shape_s = thr_mma_qk.partition_shape_C((self.tile_m, self.tile_n)) + softmax = Softmax.create( + softmax_scale_log2, + num_rows=acc_shape_s[0][0] * acc_shape_s[1], + softmax_scale=softmax_scale, + ) + softmax.reset() + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + tScS_mn = layout_utils.reshape_acc_to_mn(thr_mma_qk.partition_C(cS)) + mask = AttentionMask( + self.tile_m, + self.tile_n, + seqlen, + block_info.window_size_left, + block_info.window_size_right, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + enable_r2p_optimization=not self.split_qk_n, + ) + mask_fn = partial( + mask.apply_mask, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + thr_mma=thr_mma_qk, + mask_causal=self.is_causal, + mask_local=self.is_local, + aux_data=aux_data, + fastdiv_mods=( + fastdiv_mods if const_expr(self.mask_mod is not None) else None + ), + ) + compute_one_n_block = partial( + self.compute_one_n_block_qk_pipeline, + thr_mma_qk=thr_mma_qk, + sQ=sQ, + tCrQ=tCrQ, + sK=sK, + tPsP_store=tPsP_store, + smem_thr_copy_q=smem_thr_copy_q, + smem_thr_copy_k=smem_thr_copy_k, + smem_thr_store_p=smem_thr_store_p, + tScS_mn=tScS_mn, + copy_K=copy_K, + copy_V=copy_V, + tidx=tidx, + sRowScale=sRowScale, + softmax=softmax, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + pipeline_p=pipeline_p, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + role_state = ( + PipelineState(self._num_k_stages(), Int32(0), Int32(0), Int32(1)), + PipelineState(self._num_v_stages(), Int32(0), Int32(0), Int32(1)), + PipelineState(self._num_k_stages(), Int32(0), Int32(0), Int32(0)), + PipelineState(self._num_p_stages(), Int32(0), Int32(0), Int32(1)), + ) + producer_state_k, producer_state_v, k_state, p_state = ( + self._run_n_block_schedule( + compute_one_n_block, + role_state, + block_info, + seqlen, + n_block_min, + n_block_max, + m_block, + mask_fn, + ) + ) + + sink_val = None + if const_expr(learnable_sink is not None): + if const_expr(not self.pack_gqa): + sink_val = Float32(learnable_sink[head_idx]) + else: + sink_val = cute.make_rmem_tensor_like(softmax.row_max, Float32) + for r in cutlass.range(cute.size(sink_val), unroll_full=True): + row = m_block * self.tile_m + tScS_mn[r][0] + q_head_idx = ( + row % self.qhead_per_kvhead + head_idx * self.qhead_per_kvhead + ) + sink_val[r] = Float32(learnable_sink[q_head_idx]) + if const_expr(self.is_split_kv and learnable_sink is not None): + if const_expr(not self.pack_gqa): + sink_val = sink_val if split_idx == 0 else -Float32.inf + elif split_idx != 0: + sink_val.fill(-Float32.inf) + row_scale = softmax.finalize(sink_val=sink_val) + final_state = PipelineState(1, Int32(0), Int32(0), Int32(1)) + pipeline_final.producer_acquire(final_state) + if tScS_mn[0][1] == 0: + for r in cutlass.range(cute.size(row_scale), unroll_full=True): + row = tScS_mn[r][0] + sFinalScale[row] = row_scale[r] + sLSE[row] = softmax.row_sum[r] + cute.arch.fence_view_async_shared() + pipeline_final.producer_commit(final_state) + final_state.advance() + + pipeline_p.producer_tail(p_state) + pipeline_final.producer_tail(final_state) + if tidx < cute.arch.WARP_SIZE: + pipeline_k.producer_tail(producer_state_k) + pipeline_v.producer_tail(producer_state_v) + + @cute.jit + def compute_one_n_block_qk_pipeline( + self, + n_block: Int32, + role_state, + thr_mma_qk: cute.TiledMma, + sQ: cute.Tensor, + tCrQ: Optional[cute.Tensor], + sK: cute.Tensor, + tPsP_store: cute.Tensor, + smem_thr_copy_q: cute.TiledCopy, + smem_thr_copy_k: cute.TiledCopy, + smem_thr_store_p: cute.TiledCopy, + tScS_mn: cute.Tensor, + copy_K: Callable, + copy_V: Callable, + tidx: Int32, + sRowScale: cute.Tensor, + softmax: Softmax, + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + pipeline_p: PipelineAsync, + batch_idx: Int32, + head_idx: Int32, + m_block: Int32, + seqlen: SeqlenInfoQK, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + mask_fn: Optional[Callable] = None, + is_first_n_block: cutlass.Constexpr = False, + ): + producer_state_k, producer_state_v, k_state, p_state = role_state + if tidx < cute.arch.WARP_SIZE: + pipeline_k.producer_acquire(producer_state_k) + copy_K( + src_idx=n_block, + dst_idx=producer_state_k.index, + tma_bar_ptr=pipeline_k.producer_get_barrier(producer_state_k), + ) + pipeline_k.producer_commit(producer_state_k) + pipeline_v.producer_acquire(producer_state_v) + copy_V( + src_idx=n_block, + dst_idx=producer_state_v.index, + tma_bar_ptr=pipeline_v.producer_get_barrier(producer_state_v), + ) + pipeline_v.producer_commit(producer_state_v) + producer_state_k.advance() + producer_state_v.advance() + + pipeline_p.producer_acquire(p_state) + + k_wait_token = pipeline_k.consumer_try_wait(k_state) + pipeline_k.consumer_wait(k_state, k_wait_token) + + acc_shape_s = thr_mma_qk.partition_shape_C((self.tile_m, self.tile_n)) + acc_s = cute.make_rmem_tensor(acc_shape_s, Float32) + acc_s.fill(0.0) + if const_expr(self._q_in_regs_pipeline()): + self._gemm_qk_a_in_regs( + thr_mma_qk, + acc_s, + tCrQ, + sK[None, None, k_state.index], + smem_thr_copy_k, + ) + else: + self._gemm_qk_phase_local( + thr_mma_qk, + acc_s, + sQ, + sK[None, None, k_state.index], + smem_thr_copy_q, + smem_thr_copy_k, + ) + pipeline_k.consumer_release(k_state) + k_state.advance() + + if const_expr(self.score_mod is not None): + self.apply_score_mod( + thr_mma_qk, + batch_idx, + head_idx, + m_block, + acc_s, + n_block, + softmax_scale=softmax.softmax_scale, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + if const_expr(mask_fn is not None): + mask_fn(acc_s, n_block=n_block) + row_scale = softmax.online_softmax(acc_s, is_first=is_first_n_block) + rP = cute.make_fragment_like(acc_s, self.dtype) + rP.store(acc_s.load().to(self.dtype)) + tOrP_qk = layout_utils.reshape_acc_to_frgA(rP) + tPrP = smem_thr_store_p.retile(tOrP_qk) + cute.copy( + smem_thr_store_p, + tPrP, + tPsP_store[None, None, None, p_state.index], + ) + self._publish_row_scale(row_scale, tScS_mn, sRowScale[p_state.index, None]) + cute.arch.fence_view_async_shared() + pipeline_p.producer_commit(p_state) + p_state.advance() + return producer_state_k, producer_state_v, k_state, p_state + + @cute.jit + def mma_pv_pipeline( + self, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sV: cute.Tensor, + sO: cute.Tensor, + sP: cute.Tensor, + sRowScale: cute.Tensor, + sFinalScale: cute.Tensor, + sLSE: cute.Tensor, + pipeline_v: PipelineAsync, + pipeline_p: PipelineAsync, + pipeline_final: PipelineAsync, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma_pv: cute.TiledMma, + tidx: Int32, + block_info: BlockInfo, + seqlen: SeqlenInfoQK, + n_block_min: Int32, + n_block_max: Int32, + m_block: Int32, + head_idx: Int32, + batch_idx: Int32, + split_idx: Int32, + ): + thr_mma_pv = tiled_mma_pv.get_slice(tidx) + acc_shape_o = thr_mma_pv.partition_shape_C((self.tile_m, self.tile_hdimv)) + acc_o = cute.make_rmem_tensor(acc_shape_o, Float32) + acc_o.fill(0.0) + + smem_copy_atom_p = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), self.dtype + ) + smem_copy_atom_v = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), self.dtype + ) + smem_thr_copy_p = utils.make_tiled_copy_A( + smem_copy_atom_p, tiled_mma_pv + ).get_slice(tidx) + smem_thr_copy_v = utils.make_tiled_copy_B( + smem_copy_atom_v, tiled_mma_pv + ).get_slice(tidx) + tPsP = smem_thr_copy_p.partition_S(sP) + tOrP = thr_mma_pv.make_fragment_A(thr_mma_pv.partition_A(sP[None, None, 0])) + cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) + tOcO_mn = layout_utils.reshape_acc_to_mn(thr_mma_pv.partition_C(cO)) + compute_one_n_block = partial( + self.compute_one_n_block_pv_pipeline, + thr_mma_pv=thr_mma_pv, + acc_o=acc_o, + tOrP=tOrP, + tPsP=tPsP, + sV=sV, + sRowScale=sRowScale, + tOcO_mn=tOcO_mn, + smem_thr_copy_p=smem_thr_copy_p, + smem_thr_copy_v=smem_thr_copy_v, + pipeline_v=pipeline_v, + pipeline_p=pipeline_p, + ) + role_state = ( + PipelineState(self._num_v_stages(), Int32(0), Int32(0), Int32(0)), + PipelineState(self._num_p_stages(), Int32(0), Int32(0), Int32(0)), + ) + v_state, p_state = self._run_n_block_schedule( + compute_one_n_block, + role_state, + block_info, + seqlen, + n_block_min, + n_block_max, + m_block, + None, + ) + + final_state = PipelineState(1, Int32(0), Int32(0), Int32(0)) + pipeline_final.consumer_wait(final_state) + num_rows_pv = acc_o.shape[0][0] * acc_o.shape[1] + row_scale = cute.make_rmem_tensor(num_rows_pv, Float32) + lse = cute.make_rmem_tensor(num_rows_pv, Float32) + for r in cutlass.range(cute.size(row_scale), unroll_full=True): + row = tOcO_mn[r, 0][0] + row_scale[r] = sFinalScale[row] + lse[r] = sLSE[row] + pipeline_final.consumer_release(final_state) + final_state.advance() + self._rescale_O(acc_o, row_scale) + self.epilogue( + acc_o, + lse, + mO, + mLSE, + sO, + seqlen, + gmem_tiled_copy_O, + None, + tiled_mma_pv, + tidx, + m_block, + head_idx, + batch_idx, + split_idx, + ) + + @cute.jit + def compute_one_n_block_pv_pipeline( + self, + n_block: Int32, + role_state, + thr_mma_pv: cute.TiledMma, + acc_o: cute.Tensor, + tOrP: cute.Tensor, + tPsP: cute.Tensor, + sV: cute.Tensor, + sRowScale: cute.Tensor, + tOcO_mn: cute.Tensor, + smem_thr_copy_p: cute.TiledCopy, + smem_thr_copy_v: cute.TiledCopy, + pipeline_v: PipelineAsync, + pipeline_p: PipelineAsync, + mask_fn: Optional[Callable] = None, + is_first_n_block: cutlass.Constexpr = False, + ): + v_state, p_state = role_state + p_wait_token = pipeline_p.consumer_try_wait(p_state) + pipeline_p.consumer_wait(p_state, p_wait_token) + num_rows_pv = acc_o.shape[0][0] * acc_o.shape[1] + row_scale = cute.make_rmem_tensor(num_rows_pv, Float32) + for r in cutlass.range(cute.size(row_scale), unroll_full=True): + row_scale[r] = sRowScale[p_state.index, tOcO_mn[r, 0][0]] + self._rescale_O(acc_o, row_scale) + tOrP_copy_view = smem_thr_copy_p.retile(tOrP) + cute.copy( + smem_thr_copy_p, + tPsP[None, None, None, p_state.index], + tOrP_copy_view, + ) + pipeline_p.consumer_release(p_state) + p_state.advance() + + v_wait_token = pipeline_v.consumer_try_wait(v_state) + pipeline_v.consumer_wait(v_state, v_wait_token) + self._gemm_pv_phase_local( + thr_mma_pv, + acc_o, + tOrP, + sV[None, None, v_state.index], + smem_thr_copy_v, + ) + pipeline_v.consumer_release(v_state) + v_state.advance() + return v_state, p_state + + @cute.jit + def mma_persistent( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sO: cute.Tensor, + sP: Optional[cute.Tensor], + sRowScale: Optional[cute.Tensor], + sLSE: Optional[cute.Tensor], + learnable_sink: Optional[cute.Tensor], + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + gmem_tiled_copy_Q: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tidx: Int32, + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + base_softmax_scale: Optional[Float32], + tile_scheduler: TileSchedulerProtocol, + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + mPageTable: Optional[cute.Tensor], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + is_qk_owner: cutlass.Constexpr[bool], + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + sBias: Optional[cute.Tensor] = None, + pipeline_bias: Optional[PipelineAsync] = None, + ): + consumer_state = PipelineState(self.num_stages, Int32(0), Int32(0), Int32(0)) + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + self.is_split_kv, + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=( + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1 + ), + ) + work_tile = tile_scheduler.initial_work_tile_info() + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoQK.create( + batch_idx=batch_idx, + seqlen_q_static=( + mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1] + ), + seqlen_k_static=( + mK.shape[0] + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1] + ), + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + tile_m=self.tile_m, + tile_n=self.tile_n, + ) + n_block_min, n_block_max = self._get_n_block_min_max( + block_info, + seqlen, + m_block, + split_idx, + tile_scheduler.params.num_splits, + ) + mma_fn = partial( + self.mma, + mQ, + mO, + mLSE, + sQ, + sK, + sV, + sO, + sP, + sRowScale, + sLSE, + learnable_sink, + pipeline_k, + pipeline_v, + gmem_tiled_copy_Q, + gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + tidx, + softmax_scale_log2, + softmax_scale, + consumer_state, + block_info, + seqlen, + n_block_min, + n_block_max, + m_block, + head_idx, + batch_idx, + split_idx, + is_qk_owner, + aux_data, + fastdiv_mods, + ) + if const_expr(self.has_bias): + mma_fn( + base_softmax_scale=base_softmax_scale, + sBias=sBias, + pipeline_bias=pipeline_bias, + ) + else: + mma_fn() + + @cute.jit + def load_tma( + self, + mK: cute.Tensor, + mV: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + tma_atom_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + producer_state_k: PipelineState, + producer_state_v: PipelineState, + seqlen: SeqlenInfoQK, + n_block_min: Int32, + n_block_max: Int32, + head_idx: Int32, + batch_idx: Int32, + mBias: Optional[cute.Tensor] = None, + sBias: Optional[cute.Tensor] = None, + tma_atom_Bias: Optional[cute.CopyAtom] = None, + pipeline_bias: Optional[PipelineAsync] = None, + m_block: Int32 = Int32(0), + bias_head_idx: Int32 = Int32(0), + bias_max_idx: Int32 = Int32(0), + num_bias_loads: Int32 = Int32(0), + ): + if const_expr(not seqlen.has_cu_seqlens_k): + mK_cur = mK[None, None, head_idx, batch_idx] + mV_cur = mV[None, None, head_idx, batch_idx] + else: + mK_cur = cute.domain_offset((seqlen.offset_k, 0), mK[None, None, head_idx]) + mV_cur = cute.domain_offset((0, seqlen.offset_k), mV[None, None, head_idx]) + gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (None, 0)) + gV = cute.local_tile(mV_cur, (self.tile_hdimv, self.tile_n), (0, None)) + copy_K, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_K, 0, cute.make_layout(1), gK, sK + ) + copy_V, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_V, 0, cute.make_layout(1), gV, sV + ) + if const_expr(self.has_bias): + mBias_cur = seqlen.offset_batch_Q(mBias, batch_idx, dim=3)[ + None, None, bias_head_idx + ] + gBias = cute.local_tile( + mBias_cur, + (self.bias_block_size, self.tile_n), + (None, None), + ) + tBsBias, tBgBias = cpasync.tma_partition( + tma_atom_Bias, + 0, + cute.make_layout(1), + cute.group_modes(sBias, 0, 2), + cute.group_modes(gBias, 0, 2), + ) + if num_bias_loads > 0: + cute.arch.griddepcontrol_wait() + num_n_blocks = cutlass.max(n_block_max - n_block_min, 1) + for n_tile in cutlass.range(num_n_blocks, unroll=1): + n_block = cutlass.max(n_block_max - 1 - n_tile, n_block_min) + if const_expr(self.has_bias): + if n_tile < num_bias_loads: + pipeline_bias.producer_acquire(producer_state_k) + cute.copy( + tma_atom_Bias, + tBgBias[None, m_block, bias_max_idx - n_tile], + tBsBias[None, producer_state_k.index], + tma_bar_ptr=pipeline_bias.producer_get_barrier( + producer_state_k + ), + ) + pipeline_bias.producer_commit(producer_state_k) + pipeline_k.producer_acquire(producer_state_k) + copy_K( + src_idx=n_block, + dst_idx=producer_state_k.index, + tma_bar_ptr=pipeline_k.producer_get_barrier(producer_state_k), + ) + pipeline_k.producer_commit(producer_state_k) + producer_state_k.advance() + + pipeline_v.producer_acquire(producer_state_v) + copy_V( + src_idx=n_block, + dst_idx=producer_state_v.index, + tma_bar_ptr=pipeline_v.producer_get_barrier(producer_state_v), + ) + pipeline_v.producer_commit(producer_state_v) + producer_state_v.advance() + return producer_state_k, producer_state_v + + @cute.jit + def mma( + self, + mQ: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sO: cute.Tensor, + sP: Optional[cute.Tensor], + sRowScale: Optional[cute.Tensor], + sLSE: Optional[cute.Tensor], + learnable_sink: Optional[cute.Tensor], + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + gmem_tiled_copy_Q: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tidx: Int32, + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + consumer_state: PipelineState, + block_info: BlockInfo, + seqlen: SeqlenInfoQK, + n_block_min: Int32, + n_block_max: Int32, + m_block: Int32, + head_idx: Int32, + batch_idx: Int32, + split_idx: Int32, + is_qk_owner: cutlass.Constexpr[bool], + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + sBias: Optional[cute.Tensor] = None, + pipeline_bias: Optional[PipelineAsync] = None, + base_softmax_scale: Optional[Float32] = None, + ): + mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] + if const_expr(not self.pack_gqa): + gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0)) + num_bias_loads = Int32(0) + if const_expr(self.has_bias): + _, num_bias_loads = self._get_bias_load_info( + block_info, + seqlen, + m_block, + n_block_min, + n_block_max, + ) + + split_pv_warps = self._uses_split_pv_warps() + thr_mma_pv = tiled_mma_pv.get_slice(tidx) + acc_shape_O = thr_mma_pv.partition_shape_C((self.tile_m, self.tile_hdimv)) + acc_O = cute.make_rmem_tensor(acc_shape_O, Float32) + acc_O.fill(0.0) + + smem_copy_atom_QK = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), self.dtype + ) + smem_copy_atom_V = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), self.dtype + ) + smem_thr_copy_V = utils.make_tiled_copy_B( + smem_copy_atom_V, tiled_mma_pv + ).get_slice(tidx) + if const_expr(split_pv_warps): + tOrV = thr_mma_pv.make_fragment_B(thr_mma_pv.partition_B(sV[None, None, 0])) + tOsV = smem_thr_copy_V.partition_S(sV) + smem_thr_copy_P = utils.make_tiled_copy_A( + smem_copy_atom_QK, tiled_mma_pv + ).get_slice(tidx) + tPsP = smem_thr_copy_P.partition_S(sP) + tOrP = thr_mma_pv.make_fragment_A(thr_mma_pv.partition_A(sP[None, None, 0])) + if const_expr(not split_pv_warps or is_qk_owner): + thr_mma_qk = tiled_mma_qk.get_slice(tidx) + smem_thr_copy_Q = utils.make_tiled_copy_A( + smem_copy_atom_QK, tiled_mma_qk + ).get_slice(tidx) + smem_thr_copy_K = utils.make_tiled_copy_B( + smem_copy_atom_QK, tiled_mma_qk + ).get_slice(tidx) + if const_expr(split_pv_warps): + tSrQ = thr_mma_qk.make_fragment_A(thr_mma_qk.partition_A(sQ)) + tSrK = thr_mma_qk.make_fragment_B( + thr_mma_qk.partition_B(sK[None, None, 0]) + ) + tSsQ = smem_thr_copy_Q.partition_S(sQ) + tSsK = smem_thr_copy_K.partition_S(sK) + smem_store_atom_P = utils.get_smem_store_atom( + 120, + self.dtype, + ) + smem_thr_store_P = cute.make_tiled_copy_C( + smem_store_atom_P, tiled_mma_qk + ).get_slice(tidx) + tPsP_store = smem_thr_store_P.partition_D(sP) + + if const_expr(not split_pv_warps or is_qk_owner): + gmem_thr_copy_Q = gmem_tiled_copy_Q.get_slice(tidx) + if const_expr(not self.pack_gqa): + self.load_Q( + gmem_thr_copy_Q, + gQ, + sQ, + m_block, + seqlen=seqlen.seqlen_q, + headdim=mQ.shape[1], + ) + else: + PackGQA( + self.tile_m, + self.tile_hdim, + self.check_hdim_oob, + self.qhead_per_kvhead, + ).load_Q( + mQ_cur, + sQ, + gmem_tiled_copy_Q, + tidx, + m_block, + seqlen.seqlen_q, + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier( + barrier_id=1, + number_of_threads=self.num_Q_load_threads, + ) + + if const_expr(not split_pv_warps or is_qk_owner): + if const_expr(self._uses_n_distributed_qk()): + softmax = None + mma_params = SimpleNamespace( + thr_mma_qk=thr_mma_qk, + thr_mma_pv=thr_mma_pv, + tSrQ=tSrQ, + tSrK=tSrK, + tOrV=tOrV, + acc_O=acc_O, + tOrP=tOrP, + tidx=tidx, + ) + smem_copy_params = SimpleNamespace( + smem_thr_copy_Q=smem_thr_copy_Q, + smem_thr_copy_K=smem_thr_copy_K, + smem_thr_copy_V=smem_thr_copy_V, + tOsV=tOsV, + smem_thr_copy_P=smem_thr_copy_P, + tPsP=tPsP, + smem_thr_store_P=( + smem_thr_store_P if const_expr(self.split_qk_n) else None + ), + tPsP_store=(tPsP_store if const_expr(self.split_qk_n) else None), + sQ=sQ, + sK=sK, + sP=sP, + sRowScale=sRowScale, + sLSE=sLSE, + softmax_scale_log2=softmax_scale_log2, + softmax_scale=softmax_scale, + ) + if const_expr(self.split_qk_n): + mask = AttentionMask( + self.tile_m, + self.tile_n, + seqlen, + block_info.window_size_left, + block_info.window_size_right, + (self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1), + enable_r2p_optimization=False, + ) + mask_fn = partial( + mask.apply_mask, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + thr_mma=thr_mma_qk, + mask_causal=self.is_causal, + mask_local=self.is_local, + aux_data=aux_data, + fastdiv_mods=( + fastdiv_mods + if const_expr(self.mask_mod is not None) + else None + ), + ) + else: + mask_fn = None + else: + softmax = Softmax.create( + softmax_scale_log2, + num_rows=acc_O.shape[0][0] * acc_O.shape[1], + softmax_scale=softmax_scale, + ) + softmax.reset() + if const_expr(split_pv_warps): + mma_params = SimpleNamespace( + thr_mma_qk=thr_mma_qk, + thr_mma_pv=thr_mma_pv, + tSrQ=tSrQ, + tSrK=tSrK, + tOrV=tOrV, + acc_O=acc_O, + tOrP=tOrP, + ) + smem_copy_params = SimpleNamespace( + smem_thr_copy_Q=smem_thr_copy_Q, + smem_thr_copy_K=smem_thr_copy_K, + smem_thr_copy_V=smem_thr_copy_V, + tSsQ=tSsQ, + tSsK=tSsK, + tOsV=tOsV, + smem_thr_store_P=smem_thr_store_P, + tPsP_store=tPsP_store, + smem_thr_copy_P=smem_thr_copy_P, + tPsP=tPsP, + sRowScale=sRowScale, + sLSE=sLSE, + ) + else: + mma_params = SimpleNamespace( + thr_mma_qk=thr_mma_qk, + thr_mma_pv=thr_mma_pv, + acc_O=acc_O, + ) + smem_copy_params = SimpleNamespace( + smem_thr_copy_Q=smem_thr_copy_Q, + smem_thr_copy_K=smem_thr_copy_K, + smem_thr_copy_V=smem_thr_copy_V, + sQ=sQ, + sK=sK, + sV=sV, + ) + mask = AttentionMask( + self.tile_m, + self.tile_n, + seqlen, + block_info.window_size_left, + block_info.window_size_right, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + mask_fn = partial( + mask.apply_mask, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + thr_mma=thr_mma_qk, + mask_causal=self.is_causal, + mask_local=self.is_local, + aux_data=aux_data, + fastdiv_mods=( + fastdiv_mods if const_expr(self.mask_mod is not None) else None + ), + ) + else: + softmax = None + mma_params = SimpleNamespace( + thr_mma_pv=thr_mma_pv, + tOrV=tOrV, + acc_O=acc_O, + tOrP=tOrP, + ) + smem_copy_params = SimpleNamespace( + smem_thr_copy_V=smem_thr_copy_V, + tOsV=tOsV, + smem_thr_copy_P=smem_thr_copy_P, + tPsP=tPsP, + sRowScale=sRowScale, + sLSE=sLSE, + ) + mask_fn = None + if const_expr(split_pv_warps and not self._uses_n_distributed_qk()): + compute_one_n_block = ( + self.compute_one_n_block_split_pv_owner + if const_expr(is_qk_owner) + else self.compute_one_n_block_split_pv_helper + ) + elif const_expr(not self._uses_n_distributed_qk()): + compute_one_n_block = self.compute_one_n_block + if const_expr(not split_pv_warps or is_qk_owner): + mask_fn_seqlen = partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True) + mask_fn_no_seqlen = partial( + mask_fn, mask_mod=self.mask_mod, mask_seqlen=False + ) + else: + mask_fn_seqlen = None + mask_fn_no_seqlen = None + n_block = cutlass.max(n_block_max - 1, 0) + if const_expr(self._uses_n_distributed_qk()): + consumer_state = self.compute_one_n_block_split_pv_distributed_qk( + n_block, + consumer_state, + mma_params, + smem_copy_params, + None, + pipeline_k, + pipeline_v, + score_mod=self.score_mod, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + mask_fn=mask_fn_seqlen, + is_first_n_block=True, + is_last_n_block=n_block == n_block_min, + learnable_sink=learnable_sink, + split_idx=split_idx, + ) + for n_tile in cutlass.range(n_block - n_block_min, unroll=1): + consumer_state = self.compute_one_n_block_split_pv_distributed_qk( + n_block - n_tile - 1, + consumer_state, + mma_params, + smem_copy_params, + None, + pipeline_k, + pipeline_v, + score_mod=self.score_mod, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + mask_fn=mask_fn_no_seqlen, + is_last_n_block=n_block - n_tile - 1 == n_block_min, + learnable_sink=learnable_sink, + split_idx=split_idx, + ) + else: + consumer_state = compute_one_n_block( + n_block, + consumer_state, + mma_params, + smem_copy_params, + softmax, + pipeline_k, + pipeline_v, + score_mod=self.score_mod, + sBias=sBias, + pipeline_bias=pipeline_bias, + base_softmax_scale=base_softmax_scale, + apply_bias=n_block >= n_block_max - num_bias_loads, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + mask_fn=mask_fn_seqlen, + is_first_n_block=True, + ) + n_block_upper = n_block + if const_expr(self.is_causal or self.is_local): + n_block_min_causal_local_mask = ( + block_info.get_n_block_min_causal_local_mask( + seqlen, m_block, n_block_min + ) + ) + for n_tile in cutlass.range( + n_block_max - 1 - n_block_min_causal_local_mask, unroll=1 + ): + n_block = n_block_max - 2 - n_tile + consumer_state = compute_one_n_block( + n_block, + consumer_state, + mma_params, + smem_copy_params, + softmax, + pipeline_k, + pipeline_v, + score_mod=self.score_mod, + sBias=sBias, + pipeline_bias=pipeline_bias, + base_softmax_scale=base_softmax_scale, + apply_bias=n_block >= n_block_max - num_bias_loads, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + mask_fn=mask_fn_seqlen, + ) + n_block_upper = cutlass.min( + n_block_upper, n_block_min_causal_local_mask + ) + n_block_min_before_local_mask = ( + block_info.get_n_block_min_before_local_mask( + seqlen, m_block, n_block_min + ) + ) + for n_tile in cutlass.range( + n_block_upper - n_block_min_before_local_mask, unroll=1 + ): + consumer_state = compute_one_n_block( + n_block_upper - n_tile - 1, + consumer_state, + mma_params, + smem_copy_params, + softmax, + pipeline_k, + pipeline_v, + score_mod=self.score_mod, + sBias=sBias, + pipeline_bias=pipeline_bias, + base_softmax_scale=base_softmax_scale, + apply_bias=( + n_block_upper - n_tile - 1 >= n_block_max - num_bias_loads + ), + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + mask_fn=mask_fn_no_seqlen, + ) + if const_expr(self.is_local and block_info.window_size_left is not None): + n_block_upper = cutlass.min( + n_block_upper, n_block_min_before_local_mask + ) + for n_tile in cutlass.range(n_block_upper - n_block_min, unroll=1): + consumer_state = compute_one_n_block( + n_block_upper - n_tile - 1, + consumer_state, + mma_params, + smem_copy_params, + softmax, + pipeline_k, + pipeline_v, + score_mod=self.score_mod, + sBias=sBias, + pipeline_bias=pipeline_bias, + base_softmax_scale=base_softmax_scale, + apply_bias=( + n_block_upper - n_tile - 1 >= n_block_max - num_bias_loads + ), + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + mask_fn=mask_fn_no_seqlen, + ) + + if const_expr(self._uses_n_distributed_qk()): + cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) + tOcO_mn_finalize = layout_utils.reshape_acc_to_mn( + thr_mma_pv.partition_C(cO) + ) + num_rows_pv = acc_O.shape[0][0] * acc_O.shape[1] + row_scale_pv = cute.make_rmem_tensor(num_rows_pv, Float32) + lse = cute.make_rmem_tensor(num_rows_pv, Float32) + for r in cutlass.range(cute.size(row_scale_pv), unroll_full=True): + row = tOcO_mn_finalize[r, 0][0] + row_scale_pv[r] = sRowScale[0, row] + lse[r] = sLSE[row] + self._rescale_O(acc_O, row_scale_pv) + elif const_expr(split_pv_warps): + if const_expr(is_qk_owner): + sink_val = None + if const_expr(learnable_sink is not None): + if const_expr(not self.pack_gqa): + sink_val = Float32(learnable_sink[head_idx]) + else: + sink_val = cute.make_rmem_tensor_like(softmax.row_max, Float32) + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + tScS_mn_finalize = layout_utils.reshape_acc_to_mn( + thr_mma_qk.partition_C(cS) + ) + for r in cutlass.range(cute.size(sink_val), unroll_full=True): + row = m_block * self.tile_m + tScS_mn_finalize[r][0] + q_head_idx = ( + row % self.qhead_per_kvhead + + head_idx * self.qhead_per_kvhead + ) + sink_val[r] = Float32(learnable_sink[q_head_idx]) + if const_expr(self.is_split_kv and learnable_sink is not None): + if const_expr(not self.pack_gqa): + sink_val = sink_val if split_idx == 0 else -Float32.inf + elif split_idx != 0: + sink_val.fill(-Float32.inf) + row_scale_qk = softmax.finalize(sink_val=sink_val) + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + tScS_mn_finalize = layout_utils.reshape_acc_to_mn( + thr_mma_qk.partition_C(cS) + ) + if tScS_mn_finalize[0][1] == 0: + for r in cutlass.range(cute.size(row_scale_qk), unroll_full=True): + row = tScS_mn_finalize[r][0] + sRowScale[2, row] = row_scale_qk[r] + sLSE[row] = softmax.row_sum[r] + cute.arch.fence_view_async_shared() + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.PFull), + number_of_threads=self.num_mma_threads, + ) + cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) + tOcO_mn_finalize = layout_utils.reshape_acc_to_mn( + thr_mma_pv.partition_C(cO) + ) + num_rows_pv = acc_O.shape[0][0] * acc_O.shape[1] + row_scale_pv = cute.make_rmem_tensor(num_rows_pv, Float32) + lse = cute.make_rmem_tensor(num_rows_pv, Float32) + for r in cutlass.range(cute.size(row_scale_pv), unroll_full=True): + row = tOcO_mn_finalize[r, 0][0] + row_scale_pv[r] = sRowScale[2, row] + lse[r] = sLSE[row] + self._rescale_O(acc_O, row_scale_pv) + else: + sink_val = None + if const_expr(learnable_sink is not None): + if const_expr(not self.pack_gqa): + sink_val = Float32(learnable_sink[head_idx]) + else: + sink_val = cute.make_rmem_tensor_like(softmax.row_max, Float32) + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + tScS_mn_finalize = layout_utils.reshape_acc_to_mn( + thr_mma_qk.partition_C(cS) + ) + for r in cutlass.range(cute.size(sink_val), unroll_full=True): + row = m_block * self.tile_m + tScS_mn_finalize[r][0] + q_head_idx = ( + row % self.qhead_per_kvhead + + head_idx * self.qhead_per_kvhead + ) + sink_val[r] = Float32(learnable_sink[q_head_idx]) + if const_expr(self.is_split_kv and learnable_sink is not None): + if const_expr(not self.pack_gqa): + sink_val = sink_val if split_idx == 0 else -Float32.inf + elif split_idx != 0: + sink_val.fill(-Float32.inf) + row_scale = softmax.finalize(sink_val=sink_val) + softmax.rescale_O(acc_O, row_scale) + lse = softmax.row_sum + + self.epilogue( + acc_O, + lse, + mO, + mLSE, + sO, + seqlen, + gmem_tiled_copy_O, + None, + tiled_mma_pv, + tidx, + m_block, + head_idx, + batch_idx, + split_idx, + ) + return consumer_state + + @cute.jit + def compute_one_n_block_split_pv_distributed_qk( + self, + n_block: Int32, + consumer_state: PipelineState, + mma_params: SimpleNamespace, + smem_copy_params: SimpleNamespace, + softmax: None, + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + score_mod: Callable | None, + batch_idx: Int32, + head_idx: Int32, + m_block: Int32, + seqlen: SeqlenInfoQK, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + mask_fn: Optional[Callable] = None, + is_first_n_block: cutlass.Constexpr = False, + check_inf: cutlass.Constexpr = True, + is_last_n_block: cutlass.Boolean = False, + learnable_sink: Optional[cute.Tensor] = None, + split_idx: Int32 = 0, + ): + """Run N-distributed QK with a shared online-softmax reduction. + + Four warps own disjoint K/V-column slices. Each warp publishes local + row max/sum values; one lane per query row combines them with the + running state and publishes the scale used for the vector P store. + """ + p_stage = consumer_state.index + num_qk_warps = const_expr(self.num_qk_threads // cute.arch.WARP_SIZE) + local_sum_base = const_expr(num_qk_warps) + global_max_row = const_expr(2 * num_qk_warps) + global_sum_row = const_expr(global_max_row + 1) + old_o_scale_row = const_expr(global_max_row + 2) + warp_scale_base = const_expr(global_max_row + 3) + + acc_shape_S = mma_params.thr_mma_qk.partition_shape_C( + (self.tile_m, self.tile_n) + ) + acc_S = cute.make_rmem_tensor(acc_shape_S, Float32) + acc_S.fill(0.0) + k_wait_token = pipeline_k.consumer_try_wait(consumer_state) + pipeline_k.consumer_wait(consumer_state, k_wait_token) + + self._gemm_qk( + mma_params.thr_mma_qk, + acc_S, + mma_params.tSrQ, + mma_params.tSrK, + smem_copy_params.smem_thr_copy_Q.partition_S(smem_copy_params.sQ), + smem_copy_params.smem_thr_copy_K.partition_S(smem_copy_params.sK)[ + None, None, None, p_stage + ], + smem_copy_params.smem_thr_copy_Q, + smem_copy_params.smem_thr_copy_K, + ) + pipeline_k.consumer_release(consumer_state) + + if const_expr(score_mod is not None): + self.apply_score_mod( + mma_params.thr_mma_qk, + batch_idx, + head_idx, + m_block, + acc_S, + n_block, + softmax_scale=smem_copy_params.softmax_scale, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + if const_expr(mask_fn is not None): + mask_fn(acc_S, n_block=n_block) + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) + num_rows = cute.size(acc_S_mn.shape[0]) + row_max_local = cute.make_rmem_tensor(num_rows, Float32) + row_sum_local = cute.make_rmem_tensor(num_rows, Float32) + for r in cutlass.range(num_rows, unroll_full=True): + acc_S_row = acc_S_mn[r, None].load() + row_max = utils.fmax_reduce(acc_S_row) + row_max = cute.arch.warp_reduction_max(row_max, threads_in_group=4) + row_max_safe = 0.0 if row_max == -Float32.inf else row_max + acc_S_row_exp = cute.math.exp2( + (acc_S_row - row_max_safe) * smem_copy_params.softmax_scale_log2, + fastmath=True, + ) + row_sum = utils.fadd_reduce(acc_S_row_exp) + row_sum = utils.warp_reduce(row_sum, operator.add, width=4) + row_max_local[r] = row_max + row_sum_local[r] = row_sum + acc_S_mn[r, None].store(acc_S_row_exp) + + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + tScS_mn = layout_utils.reshape_acc_to_mn( + mma_params.thr_mma_qk.partition_C(cS), + ) + row_coord = const_expr(0) + col_coord = const_expr(1) + warp_idx = mma_params.tidx // cute.arch.WARP_SIZE + stat_writer_period = const_expr(8) + if tScS_mn[0, 0][col_coord] % stat_writer_period == 0: + for r in cutlass.range(num_rows, unroll_full=True): + row = tScS_mn[r, 0][row_coord] + smem_copy_params.sRowScale[warp_idx, row] = row_max_local[r] + smem_copy_params.sRowScale[local_sum_base + warp_idx, row] = ( + row_sum_local[r] + ) + cute.arch.fence_view_async_shared() + + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.PFull), + number_of_threads=self.num_mma_threads, + ) + if mma_params.tidx < self.tile_m: + row = mma_params.tidx + row_max = smem_copy_params.sRowScale[0, row] + for warp_idx_it in cutlass.range_constexpr(1, num_qk_warps): + row_max = utils.fmax( + row_max, + smem_copy_params.sRowScale[warp_idx_it, row], + ) + row_max_prev = ( + row_max + if const_expr(is_first_n_block) + else smem_copy_params.sRowScale[global_max_row, row] + ) + row_max_new = ( + row_max + if const_expr(is_first_n_block) + else utils.fmax(row_max_prev, row_max) + ) + row_max_new_safe = 0.0 if row_max_new == -Float32.inf else row_max_new + old_o_scale = ( + 1.0 + if const_expr(is_first_n_block) + else cute.math.exp2( + (row_max_prev - row_max_new_safe) + * smem_copy_params.softmax_scale_log2, + fastmath=True, + ) + ) + row_sum_new = ( + 0.0 + if const_expr(is_first_n_block) + else smem_copy_params.sRowScale[global_sum_row, row] * old_o_scale + ) + for warp_idx_it in cutlass.range_constexpr(num_qk_warps): + warp_scale = cute.math.exp2( + (smem_copy_params.sRowScale[warp_idx_it, row] - row_max_new_safe) + * smem_copy_params.softmax_scale_log2, + fastmath=True, + ) + smem_copy_params.sRowScale[warp_scale_base + warp_idx_it, row] = ( + warp_scale + ) + row_sum_new += ( + smem_copy_params.sRowScale[local_sum_base + warp_idx_it, row] + * warp_scale + ) + smem_copy_params.sRowScale[global_max_row, row] = row_max_new + smem_copy_params.sRowScale[global_sum_row, row] = row_sum_new + smem_copy_params.sRowScale[old_o_scale_row, row] = old_o_scale + if is_last_n_block: + row_max_final = row_max_new + row_sum_final = row_sum_new + if const_expr(learnable_sink is not None): + if split_idx == 0: + q_head_idx = ( + row % self.qhead_per_kvhead + + head_idx * self.qhead_per_kvhead + if const_expr(self.pack_gqa) + else head_idx + ) + sink_val = Float32(learnable_sink[q_head_idx]) + log2_e = math.log2(math.e) + if row_max_final == -Float32.inf: + row_max_final = sink_val * ( + log2_e / smem_copy_params.softmax_scale_log2 + ) + row_sum_final = 1.0 + else: + row_sum_final += cute.math.exp2( + sink_val * log2_e + - row_max_final * smem_copy_params.softmax_scale_log2, + fastmath=True, + ) + row_sum_is_zero_or_nan = ( + row_sum_final == 0.0 or row_sum_final != row_sum_final + ) + smem_copy_params.sRowScale[0, row] = cute.arch.rcp_approx( + row_sum_final if not row_sum_is_zero_or_nan else 1.0 + ) + smem_copy_params.sLSE[row] = ( + ( + row_max_final * smem_copy_params.softmax_scale_log2 + + cute.math.log2(row_sum_final, fastmath=True) + ) + * math.log(2.0) + if not row_sum_is_zero_or_nan + else -Float32.inf + ) + cute.arch.fence_view_async_shared() + + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.PEmpty), + number_of_threads=self.num_mma_threads, + ) + for r in cutlass.range(num_rows, unroll_full=True): + row = tScS_mn[r, 0][row_coord] + warp_scale = smem_copy_params.sRowScale[warp_scale_base + warp_idx, row] + acc_S_mn[r, None].store(acc_S_mn[r, None].load() * warp_scale) + rP = cute.make_fragment_like(acc_S, self.dtype) + rP.store(acc_S.load().to(self.dtype)) + tOrP_qk = layout_utils.reshape_acc_to_frgA(rP) + tPrP = smem_copy_params.smem_thr_store_P.retile(tOrP_qk) + cute.copy( + smem_copy_params.smem_thr_store_P, + tPrP, + smem_copy_params.tPsP_store[None, None, None, p_stage], + ) + cute.arch.fence_view_async_shared() + + return self._compute_one_n_block_split_pv_common( + consumer_state, + mma_params, + smem_copy_params, + pipeline_k, + pipeline_v, + False, + skip_p_empty=is_last_n_block, + ) + + @cute.jit + def compute_one_n_block_split_pv_owner( + self, + n_block: Int32, + consumer_state: PipelineState, + mma_params: SimpleNamespace, + smem_copy_params: SimpleNamespace, + softmax: Softmax, + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + score_mod: Callable | None, + sBias: Optional[cute.Tensor], + pipeline_bias: Optional[PipelineAsync], + base_softmax_scale: Optional[Float32], + apply_bias: bool, + batch_idx: Int32, + head_idx: Int32, + m_block: Int32, + seqlen: SeqlenInfoQK, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + mask_fn: Optional[Callable] = None, + is_first_n_block: cutlass.Constexpr = False, + check_inf: cutlass.Constexpr = True, + ): + """Compute QK/softmax in four owner warps, then join eight PV warps.""" + acc_shape_S = mma_params.thr_mma_qk.partition_shape_C( + (self.tile_m, self.tile_n) + ) + acc_S = cute.make_rmem_tensor(acc_shape_S, Float32) + acc_S.fill(0.0) + k_wait_token = pipeline_k.consumer_try_wait(consumer_state) + pipeline_k.consumer_wait(consumer_state, k_wait_token) + + self._gemm_qk( + mma_params.thr_mma_qk, + acc_S, + mma_params.tSrQ, + mma_params.tSrK, + smem_copy_params.tSsQ, + smem_copy_params.tSsK[None, None, None, consumer_state.index], + smem_copy_params.smem_thr_copy_Q, + smem_copy_params.smem_thr_copy_K, + ) + + if const_expr(score_mod is not None): + self.apply_score_mod( + mma_params.thr_mma_qk, + batch_idx, + head_idx, + m_block, + acc_S, + n_block, + softmax_scale=softmax.softmax_scale, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + if const_expr(mask_fn is not None): + mask_fn(acc_S, n_block=n_block) + row_scale = softmax.online_softmax( + acc_S, is_first=is_first_n_block, check_inf=check_inf + ) + rP = cute.make_fragment_like(acc_S, self.dtype) + rP.store(acc_S.load().to(self.dtype)) + tOrP_qk = layout_utils.reshape_acc_to_frgA(rP) + tPrP = smem_copy_params.smem_thr_store_P.retile(tOrP_qk) + cute.copy( + smem_copy_params.smem_thr_store_P, + tPrP, + smem_copy_params.tPsP_store[None, None, None, consumer_state.index], + ) + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + tScS_mn = layout_utils.reshape_acc_to_mn(mma_params.thr_mma_qk.partition_C(cS)) + self._publish_row_scale( + row_scale, + tScS_mn, + smem_copy_params.sRowScale[consumer_state.index, None], + ) + cute.arch.fence_view_async_shared() + + return self._compute_one_n_block_split_pv_common( + consumer_state, + mma_params, + smem_copy_params, + pipeline_k, + pipeline_v, + True, + ) + + @cute.jit + def compute_one_n_block_split_pv_helper( + self, + n_block: Int32, + consumer_state: PipelineState, + mma_params: SimpleNamespace, + smem_copy_params: SimpleNamespace, + softmax: None, + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + score_mod: Callable | None, + sBias: Optional[cute.Tensor], + pipeline_bias: Optional[PipelineAsync], + base_softmax_scale: Optional[Float32], + apply_bias: bool, + batch_idx: Int32, + head_idx: Int32, + m_block: Int32, + seqlen: SeqlenInfoQK, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + mask_fn: Optional[Callable] = None, + is_first_n_block: cutlass.Constexpr = False, + check_inf: cutlass.Constexpr = True, + ): + """Join the owner-published P tile using the four helper PV warps.""" + return self._compute_one_n_block_split_pv_common( + consumer_state, + mma_params, + smem_copy_params, + pipeline_k, + pipeline_v, + False, + ) + + @cute.jit + def _compute_one_n_block_split_pv_common( + self, + consumer_state: PipelineState, + mma_params: SimpleNamespace, + smem_copy_params: SimpleNamespace, + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + release_k: cutlass.Constexpr[bool], + skip_p_empty: cutlass.Boolean = False, + ): + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.PFull), + number_of_threads=self.num_mma_threads, + ) + cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) + tOcO_mn = layout_utils.reshape_acc_to_mn(mma_params.thr_mma_pv.partition_C(cO)) + num_rows_pv = mma_params.acc_O.shape[0][0] * mma_params.acc_O.shape[1] + row_scale_pv = cute.make_rmem_tensor(num_rows_pv, Float32) + for r in cutlass.range(cute.size(row_scale_pv), unroll_full=True): + row = tOcO_mn[r, 0][0] + row_scale_pv[r] = ( + smem_copy_params.sRowScale[ + 2 * (self.num_qk_threads // cute.arch.WARP_SIZE) + 2, + row, + ] + if const_expr(self._uses_n_distributed_qk()) + else smem_copy_params.sRowScale[consumer_state.index, row] + ) + self._rescale_O(mma_params.acc_O, row_scale_pv) + + v_wait_token = pipeline_v.consumer_try_wait(consumer_state) + pipeline_v.consumer_wait(consumer_state, v_wait_token) + tOrP_copy_view = smem_copy_params.smem_thr_copy_P.retile(mma_params.tOrP) + cute.copy( + smem_copy_params.smem_thr_copy_P, + smem_copy_params.tPsP[None, None, None, consumer_state.index], + tOrP_copy_view, + ) + self._gemm_pv( + mma_params.thr_mma_pv, + mma_params.acc_O, + mma_params.tOrP, + mma_params.tOrV, + smem_copy_params.tOsV[None, None, None, consumer_state.index], + smem_copy_params.smem_thr_copy_V, + ) + pipeline_v.consumer_release(consumer_state) + + if not skip_p_empty: + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.PEmpty), + number_of_threads=self.num_mma_threads, + ) + if const_expr(release_k): + pipeline_k.consumer_release(consumer_state) + consumer_state.advance() + return consumer_state + + @cute.jit + def _publish_row_scale( + self, + row_scale: cute.Tensor, + row_coords: cute.Tensor, + sRowScale: cute.Tensor, + ): + """Have one QK lane group publish each row without escaping state.""" + if row_coords[0][1] == 0: + for r in cutlass.range(cute.size(row_scale), unroll_full=True): + sRowScale[row_coords[r][0]] = row_scale[r] + + @cute.jit + def _rescale_O(self, acc_O: cute.Tensor, row_scale: cute.Tensor): + """Apply a shared-published row scale without carrying softmax state.""" + acc_O_mn = layout_utils.reshape_acc_to_mn(acc_O) + for r in cutlass.range(cute.size(row_scale), unroll_full=True): + acc_O_mn[r, None].store(acc_O_mn[r, None].load() * row_scale[r]) + + @cute.jit + def apply_sheared_bias( + self, + acc_S: cute.Tensor, + thr_mma_qk: cute.TiledMma, + sBias: cute.Tensor, + pipeline_bias: PipelineAsync, + consumer_state: PipelineState, + base_softmax_scale: Float32, + apply_bias: bool, + ): + """Scale QK and merge one asynchronously staged sheared-bias tile.""" + if apply_bias: + bias_wait_token = pipeline_bias.consumer_try_wait(consumer_state) + pipeline_bias.consumer_wait(consumer_state, bias_wait_token) + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + tScS = thr_mma_qk.partition_C(cS) + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + row = tScS[i][0] + col = tScS[i][1] + acc_S[i] = acc_S[i] * base_softmax_scale + if ( + const_expr(self.bias_block_size == self.tile_m) + or row < self.bias_block_size + ): + acc_S[i] = acc_S[i] + sBias[row, col, consumer_state.index].to( + self.qk_acc_dtype + ) + pipeline_bias.consumer_release(consumer_state) + else: + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + acc_S[i] = acc_S[i] * base_softmax_scale + + @cute.jit + def compute_one_n_block( + self, + n_block: Int32, + consumer_state: PipelineState, + mma_params: SimpleNamespace, + smem_copy_params: SimpleNamespace, + softmax: Softmax, + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + score_mod: Callable | None, + sBias: Optional[cute.Tensor], + pipeline_bias: Optional[PipelineAsync], + base_softmax_scale: Optional[Float32], + apply_bias: bool, + batch_idx: Int32, + head_idx: Int32, + m_block: Int32, + seqlen: SeqlenInfoQK, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + mask_fn: Optional[Callable] = None, + is_first_n_block: cutlass.Constexpr = False, + check_inf: cutlass.Constexpr = True, + ): + acc_shape_S = mma_params.thr_mma_qk.partition_shape_C( + (self.tile_m, self.tile_n) + ) + acc_S = cute.make_rmem_tensor(acc_shape_S, Float32) + acc_S.fill(0.0) + k_wait_token = pipeline_k.consumer_try_wait(consumer_state) + pipeline_k.consumer_wait(consumer_state, k_wait_token) + self._gemm_qk_phase_local( + mma_params.thr_mma_qk, + acc_S, + smem_copy_params.sQ, + smem_copy_params.sK[None, None, consumer_state.index], + smem_copy_params.smem_thr_copy_Q, + smem_copy_params.smem_thr_copy_K, + ) + pipeline_k.consumer_release(consumer_state) + + if const_expr(self.has_bias): + self.apply_sheared_bias( + acc_S, + mma_params.thr_mma_qk, + sBias, + pipeline_bias, + consumer_state, + base_softmax_scale, + apply_bias, + ) + if const_expr(score_mod is not None): + self.apply_score_mod( + mma_params.thr_mma_qk, + batch_idx, + head_idx, + m_block, + acc_S, + n_block, + softmax_scale=( + 1.0 if const_expr(self.has_bias) else softmax.softmax_scale + ), + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + if const_expr(mask_fn is not None): + mask_fn(acc_S, n_block=n_block) + row_scale = softmax.online_softmax( + acc_S, is_first=is_first_n_block, check_inf=check_inf + ) + softmax.rescale_O(mma_params.acc_O, row_scale) + rP = cute.make_fragment_like(acc_S, self.dtype) + rP.store(acc_S.load().to(self.dtype)) + tOrP = layout_utils.reshape_acc_to_frgA(rP) + + v_wait_token = pipeline_v.consumer_try_wait(consumer_state) + pipeline_v.consumer_wait(consumer_state, v_wait_token) + self._gemm_pv_phase_local( + mma_params.thr_mma_pv, + mma_params.acc_O, + tOrP, + smem_copy_params.sV[None, None, consumer_state.index], + smem_copy_params.smem_thr_copy_V, + ) + pipeline_v.consumer_release(consumer_state) + consumer_state.advance() + return consumer_state + + @cute.jit + def _gemm_qk_a_in_regs( + self, + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrQ: cute.Tensor, + sK: cute.Tensor, + smem_thr_copy_K: cute.TiledCopy, + ): + """Issue QK with the full Q fragment resident across N blocks.""" + tCrK = tiled_mma.make_fragment_B(tiled_mma.partition_B(sK)) + tCsK = smem_thr_copy_K.partition_S(sK) + tCrK_copy_view = smem_thr_copy_K.retile(tCrK) + cute.copy( + smem_thr_copy_K, + tCsK[None, None, 0], + tCrK_copy_view[None, None, 0], + ) + for k in cutlass.range_constexpr(cute.size(tCsK.shape[2])): + if k < cute.size(tCsK.shape[2]) - 1: + cute.copy( + smem_thr_copy_K, + tCsK[None, None, k + 1], + tCrK_copy_view[None, None, k + 1], + ) + cute.gemm( + tiled_mma, + acc, + tCrQ[None, None, k], + tCrK[None, None, k], + acc, + ) + + @cute.jit + def _gemm_qk_phase_local( + self, + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + sQ: cute.Tensor, + sK: cute.Tensor, + smem_thr_copy_Q: cute.TiledCopy, + smem_thr_copy_K: cute.TiledCopy, + ): + """Allocate Q/K fragments only for the QK phase.""" + tCrQ = tiled_mma.make_fragment_A(tiled_mma.partition_A(sQ)) + tCrK = tiled_mma.make_fragment_B(tiled_mma.partition_B(sK)) + self._gemm_qk( + tiled_mma, + acc, + tCrQ, + tCrK, + smem_thr_copy_Q.partition_S(sQ), + smem_thr_copy_K.partition_S(sK), + smem_thr_copy_Q, + smem_thr_copy_K, + ) + + @cute.jit + def _gemm_pv_phase_local( + self, + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrP: cute.Tensor, + sV: cute.Tensor, + smem_thr_copy_V: cute.TiledCopy, + ): + """Allocate the V fragment only for the PV phase.""" + tCrV = tiled_mma.make_fragment_B(tiled_mma.partition_B(sV)) + self._gemm_pv( + tiled_mma, + acc, + tCrP, + tCrV, + smem_thr_copy_V.partition_S(sV), + smem_thr_copy_V, + ) + + @cute.jit + def _gemm_qk( + self, + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrQ: cute.Tensor, + tCrK: cute.Tensor, + tCsQ: cute.Tensor, + tCsK: cute.Tensor, + smem_thr_copy_Q: cute.TiledCopy, + smem_thr_copy_K: cute.TiledCopy, + ): + """Issue the SM120 QK warp-MMA mainloop.""" + tCrQ_copy_view = smem_thr_copy_Q.retile(tCrQ) + tCrK_copy_view = smem_thr_copy_K.retile(tCrK) + cute.copy(smem_thr_copy_Q, tCsQ[None, None, 0], tCrQ_copy_view[None, None, 0]) + cute.copy(smem_thr_copy_K, tCsK[None, None, 0], tCrK_copy_view[None, None, 0]) + for k in cutlass.range_constexpr(cute.size(tCsQ.shape[2])): + if k < cute.size(tCsQ.shape[2]) - 1: + cute.copy( + smem_thr_copy_Q, + tCsQ[None, None, k + 1], + tCrQ_copy_view[None, None, k + 1], + ) + cute.copy( + smem_thr_copy_K, + tCsK[None, None, k + 1], + tCrK_copy_view[None, None, k + 1], + ) + cute.gemm( + tiled_mma, + acc, + tCrQ[None, None, k], + tCrK[None, None, k], + acc, + ) + + @cute.jit + def _gemm_pv( + self, + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrP: cute.Tensor, + tCrV: cute.Tensor, + tCsV: cute.Tensor, + smem_thr_copy_V: cute.TiledCopy, + ): + """Issue the SM120 PV warp-MMA mainloop.""" + tCrV_copy_view = smem_thr_copy_V.retile(tCrV) + for k in cutlass.range_constexpr(cute.size(tCrP.shape[2])): + cute.copy( + smem_thr_copy_V, + tCsV[None, None, k], + tCrV_copy_view[None, None, k], + ) + cute.gemm( + tiled_mma, + acc, + tCrP[None, None, k], + tCrV[None, None, k], + acc, + ) + + @cute.jit + def apply_score_mod( + self, + thr_mma_qk, + batch_idx, + head_idx, + m_block, + acc_S, + n_block, + softmax_scale, + seqlen, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + ): + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + cS = cute.domain_offset((m_block * self.tile_m, n_block * self.tile_n), cS) + tScS = thr_mma_qk.partition_C(cS) + apply_score_mod_inner( + acc_S, + tScS, + self.score_mod, + batch_idx, + head_idx, + softmax_scale, + self.score_vec_size, + self.qk_acc_dtype, + aux_data, + fastdiv_mods, + seqlen_info=seqlen, + constant_q_idx=None, + qhead_per_kvhead=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd_decode.py b/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd_decode.py new file mode 100644 index 000000000..e046f7b4f --- /dev/null +++ b/python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd_decode.py @@ -0,0 +1,711 @@ +# Copyright (c) 2026, SGLang Team. +"""End-to-end transposed SM120 paged-decode specialization. + +This path keeps the packed query axis on the N=8 dimension of warp MMA: + + scores.T = K @ Q.T # (64, 8) + output.T = V.T @ P.T # (256, 8) + +The dataflow is isolated from the general SM120 kernel because its page-TMA +transport, column-wise online softmax, and transposed epilogue form one compile +specialization. +""" + +import math +from typing import Optional + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr +from cutlass.cute.nvgpu import warp +from cutlass.pipeline import PipelineAsync, PipelineState +from quack import layout_utils + +from sglang.kernels.ops.attention.fa4_sm120.flash_fwd import ( + FlashAttentionForwardSm120, +) +from sglang.kernels.ops.attention.flash_attn.cute import utils +from sglang.kernels.ops.attention.flash_attn.cute.block_info import BlockInfo +from sglang.kernels.ops.attention.flash_attn.cute.named_barrier import NamedBarrierFwd +from sglang.kernels.ops.attention.flash_attn.cute.pack_gqa import PackGQA +from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK +from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData + + +class FlashAttentionForwardSm120DecodeTranspose(FlashAttentionForwardSm120): + """M64N8 QK and transposed PV for qualified packed single-token decode.""" + + # Paged TMA is part of this kernel's dataflow, not a runtime tuning knob. + # Keeping it on the distinct class identity prevents a gather-compiled + # specialization from being reused for the TMA tensor layout. + paged_tma = True + query_mma_n = 8 + query_in_regs = True + + def _uses_n_distributed_qk(self) -> bool: + # Reuse the base kernel's four-consumer-warp dispatch. All four warps + # participate in both transposed MMA phases. + return True + + def _uses_split_pv_warps(self) -> bool: + # Experimental mixed-HDV channel slices retain the same shared P + # handoff even though they are outside the base kernel's qualified + # (HDQ, HDV) configuration table. + return True + + def _setup_attributes(self): + super()._setup_attributes() + # The tiled MMA is deliberately warp-local. The base kernel normally + # derives its cooperative-group size from TiledMma.size, which would + # expose only one consumer warp here. Four physical consumer warps + # instead operate on disjoint K rows. + self.num_qk_threads = self.num_threads + self.num_mma_threads = self.num_threads + self.num_Q_load_threads = self.num_threads + self.num_epilogue_threads = self.num_threads + + def _get_tiled_mma(self): + mma_op = warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)) + tiled_mma_qk = cute.make_tiled_mma( + mma_op, + (1, 1, 1), + permutation_mnk=(16, self.query_mma_n, 16), + ) + # Each warp covers 64 value rows; the four disjoint SMEM views cover + # HDV=256 without a CTA-level M permutation. + tiled_mma_pv = cute.make_tiled_mma( + mma_op, + (1, 1, 1), + permutation_mnk=(64, self.query_mma_n, 16), + ) + return tiled_mma_qk, tiled_mma_pv + + @cute.jit + def _gemm_n8( + self, + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + tCsA: cute.Tensor, + tCsB: cute.Tensor, + smem_thr_copy_A: cute.TiledCopy, + smem_thr_copy_B: cute.TiledCopy, + B_in_regs: cutlass.Constexpr[bool] = False, + ): + """Issue an N=8 warp-MMA mainloop through the underlying MMA atom. + + CuTe DSL's tiled-MMA verifier rejects m16n8k16 when logical N is + exactly eight because it compares the raw A value mode against the C + value mode. The hardware atom itself has the correct native fragment + contract, so keep tiling for partitioning/copies and issue GEMM through + that atom. + """ + mma_atom = cute.make_mma_atom(tiled_mma.op) + tCrA_copy_view = smem_thr_copy_A.retile(tCrA) + tCrB_copy_view = smem_thr_copy_B.retile(tCrB) + cute.copy( + smem_thr_copy_A, + tCsA[None, None, 0], + tCrA_copy_view[None, None, 0], + ) + if const_expr(not B_in_regs): + cute.copy( + smem_thr_copy_B, + tCsB[None, None, 0], + tCrB_copy_view[None, None, 0], + ) + for k in cutlass.range_constexpr(cute.size(tCsA.shape[2])): + if k < cute.size(tCsA.shape[2]) - 1: + cute.copy( + smem_thr_copy_A, + tCsA[None, None, k + 1], + tCrA_copy_view[None, None, k + 1], + ) + if const_expr(not B_in_regs): + cute.copy( + smem_thr_copy_B, + tCsB[None, None, k + 1], + tCrB_copy_view[None, None, k + 1], + ) + cute.gemm( + mma_atom, + acc, + tCrA[None, None, k], + tCrB[None, None, k], + acc, + ) + + @cute.jit + def _rescale_transposed_o( + self, + acc_O: cute.Tensor, + tiled_mma_pv: cute.TiledMma, + tidx: Int32, + sRowScale: cute.Tensor, + scale_row: cutlass.Constexpr[int], + ): + lane_idx = tidx % cute.arch.WARP_SIZE + thr_mma_pv = tiled_mma_pv.get_slice(lane_idx) + acc_O_qd = layout_utils.reshape_acc_to_mn(acc_O, transpose=True) + cO = cute.make_identity_tensor((64, self.query_mma_n)) + tOcO_qd = layout_utils.reshape_acc_to_mn( + thr_mma_pv.partition_C(cO), transpose=True + ) + for r in cutlass.range(cute.size(acc_O_qd, mode=[0]), unroll_full=True): + query_row = tOcO_qd[r, 0][1] + acc_O_qd[r, None].store( + acc_O_qd[r, None].load() * sRowScale[scale_row, query_row] + ) + + @cute.jit + def _compute_one_n_block_transposed( + self, + n_block: Int32, + consumer_state: PipelineState, + acc_O: cute.Tensor, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sP: cute.Tensor, + sRowScale: cute.Tensor, + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + smem_thr_copy_K: cute.TiledCopy, + smem_thr_copy_Q: cute.TiledCopy, + smem_thr_copy_V: cute.TiledCopy, + smem_thr_copy_P: cute.TiledCopy, + tSrK: cute.Tensor, + tSrQ: cute.Tensor, + tOrV: cute.Tensor, + tOrP: cute.Tensor, + tKsK: cute.Tensor, + tQsQ: cute.Tensor, + tVsV: cute.Tensor, + tPsP: cute.Tensor, + tidx: Int32, + softmax_scale_log2: Float32, + seqlen: SeqlenInfoQK, + window_size_left: Optional[Int32], + is_first_n_block: cutlass.Constexpr[bool] = False, + ): + num_qk_warps = const_expr(4) + local_sum_base = const_expr(num_qk_warps) + global_max_row = const_expr(2 * num_qk_warps) + global_sum_row = const_expr(global_max_row + 1) + old_o_scale_row = const_expr(global_max_row + 2) + warp_scale_base = const_expr(global_max_row + 3) + p_stage = consumer_state.index + warp_idx = tidx // cute.arch.WARP_SIZE + lane_idx = tidx % cute.arch.WARP_SIZE + key_row_base = warp_idx * const_expr(16) + + k_wait_token = pipeline_k.consumer_try_wait(consumer_state) + pipeline_k.consumer_wait(consumer_state, k_wait_token) + + thr_mma_qk = tiled_mma_qk.get_slice(lane_idx) + acc_shape_S = thr_mma_qk.partition_shape_C((16, self.query_mma_n)) + acc_S = cute.make_rmem_tensor(acc_shape_S, Float32) + acc_S.fill(0.0) + self._gemm_n8( + tiled_mma_qk, + acc_S, + tSrK, + tSrQ, + tKsK[None, None, None, p_stage], + tQsQ, + smem_thr_copy_K, + smem_thr_copy_Q, + B_in_regs=self.query_in_regs, + ) + pipeline_k.consumer_release(consumer_state) + + acc_S_qk = layout_utils.reshape_acc_to_mn(acc_S, transpose=True) + cS = cute.make_identity_tensor((16, self.query_mma_n)) + tScS_qk = layout_utils.reshape_acc_to_mn( + thr_mma_qk.partition_C(cS), transpose=True + ) + num_query_rows = cute.size(acc_S_qk, mode=[0]) + row_max_local = cute.make_rmem_tensor(num_query_rows, Float32) + row_sum_local = cute.make_rmem_tensor(num_query_rows, Float32) + for r in cutlass.range(num_query_rows, unroll_full=True): + # The final paged tile can be only partially populated. The page + # gather leaves invalid SMEM rows untouched. Keep the test outside + # the unrolled fragment loop so complete 64-token tiles pay no + # per-element coordinate/comparison cost. + tile_start = n_block * self.tile_n + local_window_start = ( + cutlass.max( + seqlen.seqlen_k - 1 - window_size_left, + 0, + ) + if const_expr(self.is_local and window_size_left is not None) + else Int32(0) + ) + if tile_start + self.tile_n > seqlen.seqlen_k or ( + const_expr(self.is_local and window_size_left is not None) + and tile_start < local_window_start + ): + for c in cutlass.range(cute.size(acc_S_qk, mode=[1]), unroll_full=True): + key_row = tile_start + key_row_base + tScS_qk[r, c][0] + if key_row >= seqlen.seqlen_k or ( + const_expr(self.is_local and window_size_left is not None) + and key_row < local_window_start + ): + acc_S_qk[r, c] = -Float32.inf + + acc_S_row = acc_S_qk[r, None].load() + row_max = utils.fmax_reduce(acc_S_row) + # In an m16n8 accumulator, lanes with the same low two lane + # bits own the same pair of N/query columns. Reducing across M + # (keys) therefore uses the strided lane group + # {lane, lane^4, lane^8, lane^16}, not a contiguous width-4 + # group as in the ordinary row-wise QK layout. + for offset in cutlass.range_constexpr(2, 5): + row_max = utils.fmax( + row_max, + cute.arch.shuffle_sync_bfly(row_max, offset=1 << offset), + ) + row_max_safe = 0.0 if row_max == -Float32.inf else row_max + acc_S_row_exp = cute.math.exp2( + (acc_S_row - row_max_safe) * softmax_scale_log2, + fastmath=True, + ) + row_sum = utils.fadd_reduce(acc_S_row_exp) + for offset in cutlass.range_constexpr(2, 5): + row_sum += cute.arch.shuffle_sync_bfly(row_sum, offset=1 << offset) + row_max_local[r] = row_max + row_sum_local[r] = row_sum + acc_S_qk[r, None].store(acc_S_row_exp) + + keys_per_warp = const_expr(self.tile_n // num_qk_warps) + if tScS_qk[0, 0][0] % keys_per_warp == 0: + for r in cutlass.range(num_query_rows, unroll_full=True): + query_row = tScS_qk[r, 0][1] + sRowScale[warp_idx, query_row] = row_max_local[r] + sRowScale[local_sum_base + warp_idx, query_row] = row_sum_local[r] + cute.arch.fence_view_async_shared() + + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.PFull), + number_of_threads=self.num_mma_threads, + ) + if tidx < self.query_mma_n: + query_row = tidx + row_max = sRowScale[0, query_row] + for warp_idx_it in cutlass.range_constexpr(1, num_qk_warps): + row_max = utils.fmax(row_max, sRowScale[warp_idx_it, query_row]) + row_max_prev = ( + row_max + if const_expr(is_first_n_block) + else sRowScale[global_max_row, query_row] + ) + row_max_new = ( + row_max + if const_expr(is_first_n_block) + else utils.fmax(row_max_prev, row_max) + ) + row_max_new_safe = 0.0 if row_max_new == -Float32.inf else row_max_new + old_o_scale = ( + 1.0 + if const_expr(is_first_n_block) + else cute.math.exp2( + (row_max_prev - row_max_new_safe) * softmax_scale_log2, + fastmath=True, + ) + ) + row_sum_new = ( + 0.0 + if const_expr(is_first_n_block) + else sRowScale[global_sum_row, query_row] * old_o_scale + ) + for warp_idx_it in cutlass.range_constexpr(num_qk_warps): + warp_scale = cute.math.exp2( + (sRowScale[warp_idx_it, query_row] - row_max_new_safe) + * softmax_scale_log2, + fastmath=True, + ) + sRowScale[warp_scale_base + warp_idx_it, query_row] = warp_scale + row_sum_new += ( + sRowScale[local_sum_base + warp_idx_it, query_row] * warp_scale + ) + sRowScale[global_max_row, query_row] = row_max_new + sRowScale[global_sum_row, query_row] = row_sum_new + if const_expr(not is_first_n_block): + sRowScale[old_o_scale_row, query_row] = old_o_scale + cute.arch.fence_view_async_shared() + + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.PEmpty), + number_of_threads=self.num_mma_threads, + ) + for r in cutlass.range(num_query_rows, unroll_full=True): + query_row = tScS_qk[r, 0][1] + warp_scale = sRowScale[warp_scale_base + warp_idx, query_row] + acc_S_qk[r, None].store(acc_S_qk[r, None].load() * warp_scale) + for r in cutlass.range(num_query_rows, unroll_full=True): + query_row = tScS_qk[r, 0][1] + for c in cutlass.range(cute.size(acc_S_qk, mode=[1]), unroll_full=True): + key_row = key_row_base + tScS_qk[r, c][0] + sP[query_row, key_row, p_stage] = self.dtype(acc_S_qk[r, c]) + cute.arch.fence_view_async_shared() + + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.PFull), + number_of_threads=self.num_mma_threads, + ) + # acc_O is zero-initialized immediately before the first KV tile, so + # multiplying it by that tile's compile-time unit scale is pure + # overhead. Later tiles still rescale accumulated output before PV. + if const_expr(not is_first_n_block): + self._rescale_transposed_o( + acc_O, + tiled_mma_pv, + tidx, + sRowScale, + old_o_scale_row, + ) + + v_wait_token = pipeline_v.consumer_try_wait(consumer_state) + pipeline_v.consumer_wait(consumer_state, v_wait_token) + + self._gemm_n8( + tiled_mma_pv, + acc_O, + tOrV, + tOrP, + tVsV[None, None, None, p_stage], + tPsP[None, None, None, p_stage], + smem_thr_copy_V, + smem_thr_copy_P, + ) + pipeline_v.consumer_release(consumer_state) + + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.PEmpty), + number_of_threads=self.num_mma_threads, + ) + consumer_state.advance() + return consumer_state + + @cute.jit + def _store_transposed_output( + self, + acc_O: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sLSE: cute.Tensor, + tiled_mma_pv: cute.TiledMma, + tidx: Int32, + seqlen: SeqlenInfoQK, + m_block: Int32, + head_idx: Int32, + batch_idx: Int32, + split_idx: Int32, + ): + row_limit = seqlen.seqlen_q * self.qhead_per_kvhead + warp_idx = tidx // cute.arch.WARP_SIZE + lane_idx = tidx % cute.arch.WARP_SIZE + thr_mma_pv = tiled_mma_pv.get_slice(lane_idx) + acc_O_qd = layout_utils.reshape_acc_to_mn(acc_O, transpose=True) + cO = cute.make_identity_tensor((64, self.query_mma_n)) + tOcO_qd = layout_utils.reshape_acc_to_mn( + thr_mma_pv.partition_C(cO), transpose=True + ) + mO_cur = ( + seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx, split_idx] + if const_expr(self.is_split_kv) + else seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx] + ) + for r in cutlass.range(cute.size(acc_O_qd, mode=[0]), unroll_full=True): + query_row_local = tOcO_qd[r, 0][1] + query_row = m_block * self.tile_m + query_row_local + if query_row < row_limit: + for c in cutlass.range(cute.size(acc_O_qd, mode=[1]), unroll_full=True): + value_col = warp_idx * const_expr(64) + tOcO_qd[r, c][0] + mO_cur[query_row, value_col] = self.dtype(acc_O_qd[r, c]) + + if const_expr(mLSE is not None): + if tidx < self.query_mma_n: + query_row = m_block * self.tile_m + tidx + if query_row < row_limit: + mLSE_cur = ( + seqlen.offset_batch_Q(mLSE, batch_idx, dim=2)[ + None, head_idx, split_idx + ] + if const_expr(self.is_split_kv) + else seqlen.offset_batch_Q(mLSE, batch_idx, dim=2)[ + None, head_idx + ] + ) + mLSE_cur[query_row] = sLSE[tidx] + + @cute.jit + def mma( + self, + mQ: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sO: cute.Tensor, + sP: Optional[cute.Tensor], + sRowScale: Optional[cute.Tensor], + sLSE: Optional[cute.Tensor], + learnable_sink: Optional[cute.Tensor], + pipeline_k: PipelineAsync, + pipeline_v: PipelineAsync, + gmem_tiled_copy_Q: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tidx: Int32, + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + consumer_state: PipelineState, + block_info: BlockInfo, + seqlen: SeqlenInfoQK, + n_block_min: Int32, + n_block_max: Int32, + m_block: Int32, + head_idx: Int32, + batch_idx: Int32, + split_idx: Int32, + is_qk_owner: cutlass.Constexpr[bool], + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + ): + assert self.paged_kv + assert self.pack_gqa + assert self.is_causal or self.is_local + assert self.score_mod is None + assert self.mask_mod is None + assert self.tile_m == 16 and self.tile_n == 64 + assert self.tile_hdim == 256 + assert self.tile_hdimv in (64, 256) + assert self.qhead_per_kvhead <= self.query_mma_n + assert is_qk_owner + + mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] + warp_idx = tidx // cute.arch.WARP_SIZE + lane_idx = tidx % cute.arch.WARP_SIZE + thr_mma_qk = tiled_mma_qk.get_slice(lane_idx) + thr_mma_pv = tiled_mma_pv.get_slice(lane_idx) + acc_shape_O = thr_mma_pv.partition_shape_C((64, self.query_mma_n)) + acc_O = cute.make_rmem_tensor(acc_shape_O, Float32) + acc_O.fill(0.0) + + smem_copy_atom_k_major = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), + self.dtype, + ) + smem_copy_atom_v_major = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), + self.dtype, + ) + smem_thr_copy_K = utils.make_tiled_copy_A( + smem_copy_atom_k_major, tiled_mma_qk + ).get_slice(lane_idx) + smem_thr_copy_Q = utils.make_tiled_copy_B( + smem_copy_atom_k_major, tiled_mma_qk + ).get_slice(lane_idx) + smem_thr_copy_V = utils.make_tiled_copy_A( + smem_copy_atom_v_major, tiled_mma_pv + ).get_slice(lane_idx) + smem_thr_copy_P = utils.make_tiled_copy_B( + smem_copy_atom_k_major, tiled_mma_pv + ).get_slice(lane_idx) + + sK_warp = cute.local_tile( + sK, + (16, self.tile_hdim, self._num_k_stages()), + (warp_idx, 0, 0), + ) + sQ_query = cute.local_tile(sQ, (self.query_mma_n, self.tile_hdim), (0, 0)) + sP_query = cute.local_tile( + sP, + (self.query_mma_n, self.tile_n, self._num_p_stages()), + (0, 0, 0), + ) + tSrK = thr_mma_qk.make_fragment_A( + thr_mma_qk.partition_A(sK_warp[None, None, 0]) + ) + tSrQ = thr_mma_qk.make_fragment_B(thr_mma_qk.partition_B(sQ_query)) + sV_warp = cute.local_tile( + sV, + (64, self.tile_n, self._num_v_stages()), + (warp_idx, 0, 0), + ) + tOrV = thr_mma_pv.make_fragment_A( + thr_mma_pv.partition_A(sV_warp[None, None, 0]) + ) + tOrP = thr_mma_pv.make_fragment_B( + thr_mma_pv.partition_B(sP_query[None, None, 0]) + ) + tVsV = smem_thr_copy_V.partition_S(sV_warp) + tKsK = smem_thr_copy_K.partition_S(sK_warp) + tQsQ = smem_thr_copy_Q.partition_S(sQ_query) + tPsP = smem_thr_copy_P.partition_S(sP_query) + + PackGQA( + self.tile_m, + self.tile_hdim, + self.check_hdim_oob, + self.qhead_per_kvhead, + ).load_Q( + mQ_cur, + sQ, + gmem_tiled_copy_Q, + tidx, + m_block, + seqlen.seqlen_q, + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier( + barrier_id=1, + number_of_threads=self.num_Q_load_threads, + ) + + if const_expr(self.query_in_regs): + # Q.T is invariant across all KV tiles and its N=8 fragment is + # small. Load it while the first K TMA is in flight, then reuse it + # instead of reloading Q from SMEM per tile. + tSrQ_copy_view = smem_thr_copy_Q.retile(tSrQ) + for k in cutlass.range_constexpr(cute.size(tQsQ.shape[2])): + cute.copy( + smem_thr_copy_Q, + tQsQ[None, None, k], + tSrQ_copy_view[None, None, k], + ) + + n_block = cutlass.max(n_block_max - 1, n_block_min) + consumer_state = self._compute_one_n_block_transposed( + n_block, + consumer_state, + acc_O, + sQ, + sK, + sV, + sP, + sRowScale, + pipeline_k, + pipeline_v, + tiled_mma_qk, + tiled_mma_pv, + smem_thr_copy_K, + smem_thr_copy_Q, + smem_thr_copy_V, + smem_thr_copy_P, + tSrK, + tSrQ, + tOrV, + tOrP, + tKsK, + tQsQ, + tVsV, + tPsP, + tidx, + softmax_scale_log2, + seqlen, + block_info.window_size_left, + is_first_n_block=True, + ) + for n_tile in cutlass.range(n_block - n_block_min, unroll=1): + consumer_state = self._compute_one_n_block_transposed( + n_block - n_tile - 1, + consumer_state, + acc_O, + sQ, + sK, + sV, + sP, + sRowScale, + pipeline_k, + pipeline_v, + tiled_mma_qk, + tiled_mma_pv, + smem_thr_copy_K, + smem_thr_copy_Q, + smem_thr_copy_V, + smem_thr_copy_P, + tSrK, + tSrQ, + tOrV, + tOrP, + tKsK, + tQsQ, + tVsV, + tPsP, + tidx, + softmax_scale_log2, + seqlen, + block_info.window_size_left, + ) + + global_max_row = const_expr(8) + global_sum_row = const_expr(9) + final_scale_row = const_expr(0) + row_limit = seqlen.seqlen_q * self.qhead_per_kvhead + if tidx < self.query_mma_n: + query_row = tidx + row_max = sRowScale[global_max_row, query_row] + row_sum = sRowScale[global_sum_row, query_row] + if query_row < row_limit: + if const_expr(learnable_sink is not None): + if const_expr(not self.is_split_kv) or split_idx == 0: + q_head_idx = query_row + head_idx * self.qhead_per_kvhead + sink_val = Float32(learnable_sink[q_head_idx]) + log2_e = math.log2(math.e) + if row_max == -Float32.inf: + row_max = sink_val * (log2_e / softmax_scale_log2) + row_sum = 1.0 + else: + row_sum += cute.math.exp2( + sink_val * log2_e - row_max * softmax_scale_log2, + fastmath=True, + ) + row_sum_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum + sRowScale[final_scale_row, query_row] = cute.arch.rcp_approx( + row_sum if not row_sum_is_zero_or_nan else 1.0 + ) + sLSE[query_row] = ( + ( + row_max * softmax_scale_log2 + + cute.math.log2(row_sum, fastmath=True) + ) + * math.log(2.0) + if not row_sum_is_zero_or_nan + else -Float32.inf + ) + else: + sRowScale[final_scale_row, query_row] = 1.0 + sLSE[query_row] = -Float32.inf + cute.arch.fence_view_async_shared() + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.PFull), + number_of_threads=self.num_mma_threads, + ) + self._rescale_transposed_o( + acc_O, + tiled_mma_pv, + tidx, + sRowScale, + final_scale_row, + ) + self._store_transposed_output( + acc_O, + mO, + mLSE, + sLSE, + tiled_mma_pv, + tidx, + seqlen, + m_block, + head_idx, + batch_idx, + split_idx, + ) diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/paged_kv.py b/python/sglang/kernels/ops/attention/fa4_sm120/paged_kv.py new file mode 100644 index 000000000..921c8a4f1 --- /dev/null +++ b/python/sglang/kernels/ops/attention/fa4_sm120/paged_kv.py @@ -0,0 +1,216 @@ +import math +from dataclasses import dataclass +from typing import Type + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, const_expr +from cutlass.cute import FastDivmodDivisor +from cutlass.cute.nvgpu import cpasync +from quack.cute_dsl_utils import ParamsBase + +from sglang.kernels.ops.attention.flash_attn.cute import utils + + +@dataclass +class Sm120PagedKVManager(ParamsBase): + """SM120 paged-KV loader for the stage-sliced FA4 pipeline.""" + + mPageTable: cute.Tensor + mK_paged: cute.Tensor + mV_paged: cute.Tensor + thread_idx: Int32 + + page_size_divmod: FastDivmodDivisor + seqlen_k: Int32 + leftpad_k: Int32 + n_block_size: cutlass.Constexpr[Int32] + num_threads: cutlass.Constexpr[Int32] + head_dim_padded: cutlass.Constexpr[Int32] + head_dim_v_padded: cutlass.Constexpr[Int32] + + gmem_threads_per_row: cutlass.Constexpr[Int32] + page_entry_per_thread: cutlass.Constexpr[Int32] + async_copy_elems: cutlass.Constexpr[Int32] + + gmem_tiled_copy_KV: cute.TiledCopy + gmem_thr_copy_KV: cute.TiledCopy + tPrPage: cute.Tensor + tPrPageOffset: cute.Tensor + + @staticmethod + def create( + mPageTable: cute.Tensor, + mK_paged: cute.Tensor, + mV_paged: cute.Tensor, + page_size_divmod: FastDivmodDivisor, + bidb: Int32, + bidh: Int32, + thread_idx: Int32, + seqlen_k: Int32, + leftpad_k: Int32, + n_block_size: cutlass.Constexpr[Int32], + head_dim_padded: cutlass.Constexpr[Int32], + head_dim_v_padded: cutlass.Constexpr[Int32], + num_threads: cutlass.Constexpr[Int32], + dtype: Type[cutlass.Numeric], + ): + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // dtype.width + dtype_bytes = dtype.width // 8 + gmem_k_block_size = math.gcd( + head_dim_padded, + head_dim_v_padded, + 128 // dtype_bytes, + ) + assert gmem_k_block_size % async_copy_elems == 0 + gmem_threads_per_row = gmem_k_block_size // async_copy_elems + assert cute.arch.WARP_SIZE % gmem_threads_per_row == 0 + + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + dtype, + num_bits_per_copy=universal_copy_bits, + ) + thr_layout = cute.make_ordered_layout( + (num_threads // gmem_threads_per_row, gmem_threads_per_row), + order=(1, 0), + ) + val_layout = cute.make_layout((1, async_copy_elems)) + gmem_tiled_copy_KV = cute.make_tiled_copy_tv( + atom_async_copy, thr_layout, val_layout + ) + gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx) + + # SM120 decode tiles can have fewer rows than DMA threads. Keep one + # register entry per thread so those shapes do not create zero-sized + # register tensors. + page_entry_per_thread = max(1, (n_block_size + num_threads - 1) // num_threads) + tPrPage = cute.make_rmem_tensor((page_entry_per_thread,), Int32) + tPrPageOffset = cute.make_rmem_tensor((page_entry_per_thread,), Int32) + + return Sm120PagedKVManager( + mPageTable[bidb, None], + mK_paged[None, None, bidh, None], + mV_paged[None, None, bidh, None], + thread_idx, + page_size_divmod, + seqlen_k, + leftpad_k, + n_block_size, + num_threads, + head_dim_padded, + head_dim_v_padded, + gmem_threads_per_row, + page_entry_per_thread, + async_copy_elems, + gmem_tiled_copy_KV, + gmem_thr_copy_KV, + tPrPage, + tPrPageOffset, + ) + + @cute.jit + def _load_page_table_entry(self, i: Int32, n_block: Int32): + row = ( + i * self.num_threads + + (self.thread_idx % self.gmem_threads_per_row) + * (self.num_threads // self.gmem_threads_per_row) + + (self.thread_idx // self.gmem_threads_per_row) + ) + row_idx = n_block * self.n_block_size + row + page_idx, page_offset = divmod(row_idx + self.leftpad_k, self.page_size_divmod) + is_valid = ( + (i + 1) * self.num_threads <= self.n_block_size or row < self.n_block_size + ) and row_idx < self.seqlen_k + page = self.mPageTable[page_idx] if is_valid else 0 + self.tPrPage[i] = page + self.tPrPageOffset[i] = page_offset + + @cute.jit + def load_page_table(self, n_block: Int32): + # The entry count is a specialization constant for SM120. Expanding + # this small loop removes a measurable dynamic-loop cost in decode. + for i in cutlass.range_constexpr(self.page_entry_per_thread): + self._load_page_table_entry(i, n_block) + + @cute.jit + def compute_X_ptr(self, K_or_V: str): + tPrXPtr = cute.make_rmem_tensor((self.page_entry_per_thread,), cutlass.Int64) + mX = self.mK_paged if const_expr(K_or_V == "K") else self.mV_paged + for i in cutlass.range_constexpr(self.page_entry_per_thread): + page = self.tPrPage[i] + page_offset = self.tPrPageOffset[i] + # SGLang stores both paged K and paged V as + # (page_size, head_dim, num_pages). + tPrXPtr[i] = utils.elem_pointer(mX, (page_offset, 0, page)).toint() + return tPrXPtr + + @cute.jit + def _copy_row_async( + self, + tXsX: cute.Tensor, + tXcX: cute.Tensor, + mX_paged_cur_copy: cute.Tensor, + m: Int32, + should_load: cute.Tensor, + ): + for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])): + ki = tXcX[0, 0, k][1] // self.async_copy_elems + mX_paged_cur_copy_ki = mX_paged_cur_copy[None, ki] + tXsX_k = tXsX[None, m, k] + mX_paged_cur_copy_ki = cute.make_tensor( + mX_paged_cur_copy_ki.iterator, tXsX_k.layout + ) + cute.copy( + self.gmem_tiled_copy_KV, + mX_paged_cur_copy_ki, + tXsX_k, + pred=should_load, + ) + + @cute.jit + def load_KV(self, n_block: Int32, sX: cute.Tensor, K_or_V: str): + assert K_or_V in ("K", "V") + + tPrXPtr = self.compute_X_ptr(K_or_V) + + # The SM120 pipeline passes one stage at a time. V has already been + # transposed by the caller's shared-memory view. + sX_pi = cute.group_modes(sX, 0, 1) + head_dim = ( + self.head_dim_v_padded + if const_expr(K_or_V == "V") + else self.head_dim_padded + ) + cX = cute.make_identity_tensor((self.n_block_size, head_dim)) + tXsX = self.gmem_thr_copy_KV.partition_D(sX_pi) + tXcX = self.gmem_thr_copy_KV.partition_S(cX) + tXc0X = self.gmem_thr_copy_KV.get_slice(0).partition_S(cX) + + seqlenk_row_limit = ( + self.seqlen_k - n_block * self.n_block_size - tXcX[0][0] + if n_block >= 0 + else 0 + ) + for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])): + row_valid = tXc0X[0, m, 0][0] < seqlenk_row_limit + should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], cute.Boolean) + should_load.fill(row_valid) + + x_ptr_i64 = utils.shuffle_sync( + tPrXPtr[m // self.gmem_threads_per_row], + m % self.gmem_threads_per_row, + width=self.gmem_threads_per_row, + ) + x_gmem_ptr = cute.make_ptr( + self.mK_paged.element_type, + x_ptr_i64, + cute.AddressSpace.gmem, + assumed_align=16, + ) + mX_paged_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,))) + mX_paged_cur_copy = cute.tiled_divide( + mX_paged_cur, (self.async_copy_elems,) + ) + self._copy_row_async(tXsX, tXcX, mX_paged_cur_copy, m, should_load) diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/policy.py b/python/sglang/kernels/ops/attention/fa4_sm120/policy.py new file mode 100644 index 000000000..c607bce62 --- /dev/null +++ b/python/sglang/kernels/ops/attention/fa4_sm120/policy.py @@ -0,0 +1,90 @@ +# Copyright (c) 2026, SGLang Team. +"""Pure workload qualification shared by the SM120 FA4 host and kernel.""" + +from typing import Optional + +LOW_HD_DECODE_SHAPES = frozenset({(64, 64), (128, 128)}) +LOW_HD_DECODE_TILE_N = 64 +LOW_HD_DECODE_MIN_VISIBLE_K = 256 +LOW_HD_DECODE_SHORT_VISIBLE_K = 512 +LOW_HD_DECODE_MAX_SPLITS = 8 +LOW_HD_DECODE_SHORT_MIN_TILES_PER_CTA = 2 +LOW_HD_DECODE_LONG_MIN_TILES_PER_CTA = 8 + + +def visible_decode_seqlen_k( + max_seqlen_k: int, + *, + is_local: bool, + window_size_left: Optional[int], + window_size_right: Optional[int], +) -> int: + """Return the exact KV span visible to a single query position.""" + if not is_local: + return max(0, max_seqlen_k) + left = max_seqlen_k if window_size_left is None else max(0, window_size_left) + right = max_seqlen_k if window_size_right is None else max(0, window_size_right) + return max(0, min(max_seqlen_k, left + 1 + right)) + + +def low_hd_paged_decode_tile_m( + *, + head_dim: int, + head_dim_v: int, + paged_kv: bool, + seqlen_q: Optional[int], + visible_seqlen_k: Optional[int], + qhead_per_kvhead: Optional[int], + num_sms: Optional[int] = None, + total_mblocks: Optional[int] = None, +) -> Optional[int]: + """Return the qualified low-HD decode M tile, or ``None`` for fallback.""" + if ( + not paged_kv + or seqlen_q != 1 + or visible_seqlen_k is None + or visible_seqlen_k <= LOW_HD_DECODE_MIN_VISIBLE_K + or (head_dim, head_dim_v) not in LOW_HD_DECODE_SHAPES + ): + return None + qhead_ratio = 1 if qhead_per_kvhead is None else qhead_per_kvhead + if ( + (head_dim, head_dim_v) == (64, 64) + and qhead_ratio >= 8 + and visible_seqlen_k <= LOW_HD_DECODE_SHORT_VISIBLE_K + ): + if num_sms is not None and total_mblocks is not None: + num_n_blocks = ( + visible_seqlen_k + LOW_HD_DECODE_TILE_N - 1 + ) // LOW_HD_DECODE_TILE_N + max_short_splits = min( + LOW_HD_DECODE_MAX_SPLITS, + num_n_blocks // LOW_HD_DECODE_SHORT_MIN_TILES_PER_CTA, + ) + # Use the one-warp M16 CTA only when bounded SplitKV can fill an + # SM wave without reducing each partition below two KV tiles. + if total_mblocks * max_short_splits >= num_sms: + return 16 + return None + if qhead_ratio >= 8 and visible_seqlen_k <= LOW_HD_DECODE_SHORT_VISIBLE_K: + return 32 + return 16 + + +def is_low_hd_paged_decode_tile( + *, + head_dim: int, + head_dim_v: int, + paged_kv: bool, + seqlen_q: int, + tile_m: int, + tile_n: int, +) -> bool: + """Return whether a selected tile belongs to the qualified low-HD path.""" + return ( + paged_kv + and seqlen_q == 1 + and (head_dim, head_dim_v) in LOW_HD_DECODE_SHAPES + and tile_m in (16, 32) + and tile_n == LOW_HD_DECODE_TILE_N + ) diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/runtime.py b/python/sglang/kernels/ops/attention/fa4_sm120/runtime.py new file mode 100644 index 000000000..88a9fc875 --- /dev/null +++ b/python/sglang/kernels/ops/attention/fa4_sm120/runtime.py @@ -0,0 +1,1628 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +"""Host-side policy and launch state for the SM120 forward kernel. + +The generic FA4 interface owns argument normalization, compilation, and +architecture dispatch. This module owns the SM120-specific decisions that +must remain consistent across those phases: + +* tile, stage, and warp configuration; +* SplitKV sizing for paged decode; +* single-QK versus N-distributed-QK decode specialization; +* the direct uniform-batch decode specialization; +* cached TVM-FFI launch plans and their temporary workspaces. +""" + +import math +from collections import OrderedDict +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable, Optional + +import cutlass.utils as utils_basic +import torch +from cutlass import Int32 + +from sglang.kernels.ops.attention.fa4_sm120.flash_fwd import ( + FlashAttentionForwardSm120, +) +from sglang.kernels.ops.attention.fa4_sm120.policy import ( + LOW_HD_DECODE_LONG_MIN_TILES_PER_CTA, + LOW_HD_DECODE_MAX_SPLITS, + LOW_HD_DECODE_SHAPES, + is_low_hd_paged_decode_tile, + visible_decode_seqlen_k, +) +from sglang.kernels.ops.attention.flash_attn.cute import fa_logging +from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData + +_LAUNCH_PLAN_CAPACITY = 4096 +_WORKSPACE_CAPACITY = 64 +_WORKSPACE_VIEW_CAPACITY = 512 +_EMPTY_AUX_DATA = AuxData(None, None) +_DECODE_REFERENCE_CORE_GHZ = 2.4 +_DECODE_POWER_OF_TWO_SPLITS = (1, 2, 4, 8, 16, 32, 64, 128) + +# Joint fit of interleaved NCU duration curves from the 110-SM/384-bit and +# 188-SM/512-bit SM120 SKUs. The inputs below are physical quantities rather +# than product names or sequence-length thresholds. This calibration applies +# only to the qualified HD256 M16N64 paged-decode gather kernel. +_DECODE_LATENCY_FIXED_REF_US = 4.31046205978645 +_DECODE_LATENCY_PER_KV_TILE_REF_US = 1.3115656095100847 +_DECODE_MEMORY_FIXED_REF_US = 4.532105547934724 +_DECODE_MEMORY_SOL = 0.9415559971890943 +_DECODE_UNDERFILL_TAU_CTAS_PER_CHANNEL = 0.741265436766481 +_DECODE_COMBINE_FIXED_US = 3.654860520719416 +_DECODE_COMBINE_SPLIT_TO_8_US = 0.015181182195831644 +_DECODE_COMBINE_SPLIT_ABOVE_8_US = 0.12157753908563328 +_DECODE_COMBINE_OUTPUT_CTA_US = 0.019185022429169515 + + +@dataclass(frozen=True) +class Sm120ForwardConfig: + tile_m: int + tile_n: int + num_stages: int + num_threads: int + + @property + def compile_key(self) -> tuple: + """Return SM120 configuration state baked into generated code.""" + return ( + self.tile_m, + self.tile_n, + self.num_stages, + self.num_threads, + ) + + +@dataclass(frozen=True) +class Sm120ForwardPlan: + """Compile- and launch-time decisions for one normalized forward call.""" + + num_splits: int + split_num_n_blocks: int + direct_uniform_batch: bool + split_qk_n: bool + split_kv_blocks_per_cta: int + transpose_qk_pv: bool + launch_split_combine_early: bool + + @property + def compile_key(self) -> tuple: + """Return only specialization state baked into the generated kernel.""" + return ( + self.direct_uniform_batch, + self.split_qk_n, + self.split_kv_blocks_per_cta, + self.transpose_qk_pv, + ) + + +@dataclass(frozen=True) +class _VarlenLaunchPlan: + compiled_fn: Callable + compile_key: tuple + + +@dataclass(frozen=True) +class _PagedDecodeLaunchPlan: + compiled_fn: Callable + compile_key: tuple + actual_num_splits: int + compiled_combine: Optional[Callable] + partial_dtype: torch.dtype + + +@dataclass(frozen=True) +class _DecodeHardware: + num_sms: int + memory_channels: int + peak_memory_gbps: float + core_clock_ghz: float + is_integrated: bool + + +@lru_cache(maxsize=None) +def _get_decode_hardware(device: torch.device) -> _DecodeHardware: + properties = torch.cuda.get_device_properties(device) + memory_bus_width = properties.memory_bus_width + is_integrated = bool(getattr(properties, "is_integrated", False)) + # Discrete GPUs report the physical memory clock, while CUDA reports the + # effective LPDDR data rate on integrated devices. Convert kHz and bus bits + # to GB/s without doubling the already-effective integrated rate. + data_rate_multiplier = 1 if is_integrated else 2 + peak_memory_gbps = ( + properties.memory_clock_rate + * memory_bus_width + * data_rate_multiplier + / 8_000_000 + ) + return _DecodeHardware( + num_sms=properties.multi_processor_count, + memory_channels=max(1, memory_bus_width // 32), + peak_memory_gbps=peak_memory_gbps, + core_clock_ghz=properties.clock_rate / 1_000_000, + is_integrated=is_integrated, + ) + + +def _normalize_num_splits(num_splits: int, num_n_blocks: int) -> int: + if num_splits <= 1 or num_n_blocks <= 1: + return 1 + requested_splits = min(num_splits, num_n_blocks) + blocks_per_split = (num_n_blocks + requested_splits - 1) // requested_splits + return (num_n_blocks + blocks_per_split - 1) // blocks_per_split + + +def _predict_decode_split_us( + *, + num_splits: int, + num_n_blocks: int, + total_mblocks: int, + packed_q_rows: int, + tile_m: int, + tile_n: int, + head_dim: int, + head_dim_v: int, + element_size: int, + hardware: _DecodeHardware, +) -> float: + num_m_blocks = (packed_q_rows + tile_m - 1) // tile_m + batch_head_groups = total_mblocks // num_m_blocks + main_ctas = total_mblocks * num_splits + kv_tiles_per_cta = (num_n_blocks + num_splits - 1) // num_splits + clock_scale = _DECODE_REFERENCE_CORE_GHZ / hardware.core_clock_ghz + latency_main_us = ( + _DECODE_LATENCY_FIXED_REF_US + + _DECODE_LATENCY_PER_KV_TILE_REF_US * kv_tiles_per_cta + ) * clock_scale + ctas_per_channel = main_ctas / hardware.memory_channels + memory_fill = max( + 1e-6, + 1.0 - math.exp(-ctas_per_channel / _DECODE_UNDERFILL_TAU_CTAS_PER_CHANNEL), + ) + logical_kv_bytes = ( + batch_head_groups + * num_n_blocks + * tile_n + * (head_dim + head_dim_v) + * element_size + ) + theoretical_transfer_us = logical_kv_bytes / hardware.peak_memory_gbps / 1000.0 + memory_main_us = ( + _DECODE_MEMORY_FIXED_REF_US * clock_scale + + theoretical_transfer_us / _DECODE_MEMORY_SOL / memory_fill + ) + main_us = max(latency_main_us, memory_main_us) + if num_splits == 1: + return main_us + output_rows = batch_head_groups * packed_q_rows + combine_ctas = ((output_rows + 7) // 8) * ((head_dim_v + 127) // 128) + combine_us = ( + _DECODE_COMBINE_FIXED_US + + _DECODE_COMBINE_SPLIT_TO_8_US * min(num_splits, 8) + + _DECODE_COMBINE_SPLIT_ABOVE_8_US * max(num_splits - 8, 0) + + _DECODE_COMBINE_OUTPUT_CTA_US * combine_ctas + ) + return main_us + combine_us + + +def _select_decode_num_splits( + *, + head_dim: int, + head_dim_v: int, + element_size: int, + packed_q_rows: int, + tile_m: int, + tile_n: int, + total_mblocks: int, + num_n_blocks: int, + hardware: _DecodeHardware, +) -> int: + # At four or fewer N tiles, SplitKV's fixed combine launch costs more than + # the remaining serialized work. N=5 is the first measured crossover. + if num_n_blocks <= 4: + return 1 + + max_one_wave_splits = min( + 128, + num_n_blocks, + max(1, hardware.num_sms // total_mblocks), + ) + requested_candidates = ( + *_DECODE_POWER_OF_TWO_SPLITS, + max_one_wave_splits, + ) + candidates = { + _normalize_num_splits(requested, num_n_blocks) + for requested in requested_candidates + } + candidates = { + splits + for splits in candidates + if splits == 1 or total_mblocks * splits <= hardware.num_sms + } + if hardware.is_integrated: + if total_mblocks >= 3 * hardware.memory_channels: + # Three unsplit CTAs per LPDDR channel already saturate the LPDDR path. + return 1 + if num_n_blocks >= 64: + if total_mblocks >= 2 * hardware.memory_channels: + return 1 + target_main_ctas = (2 * hardware.num_sms + 2) // 3 + if total_mblocks == hardware.memory_channels: + # One base CTA per channel needs a second wave to cover UMA + # latency; retain the next balanced power-of-two split. + max_integrated_splits = 1 << ( + (hardware.num_sms // total_mblocks - 1).bit_length() + ) + candidates.add( + _normalize_num_splits(max_integrated_splits, num_n_blocks) + ) + elif total_mblocks < hardware.memory_channels: + max_integrated_splits = max( + 1, + (target_main_ctas + total_mblocks - 1) // total_mblocks, + ) + else: + max_integrated_splits = None + if max_integrated_splits is not None: + bounded_candidates = { + splits for splits in candidates if splits <= max_integrated_splits + } + if bounded_candidates: + candidates = bounded_candidates + # A 160-thread HD256 CTA has enough independent QK/PV work that roughly + # two thirds of an SM wave saturates both qualified SM120 SKUs. The fitted + # latency/bandwidth model is useful above that floor, but underestimates + # the cost of sparse grids at larger batch/head-group counts. + target_main_ctas = (2 * hardware.num_sms + 2) // 3 + fill_splits = min( + max_one_wave_splits, + max(1, (target_main_ctas + total_mblocks - 1) // total_mblocks), + ) + min_fill_splits = _normalize_num_splits(fill_splits, num_n_blocks) + filled_candidates = {splits for splits in candidates if splits >= min_fill_splits} + if filled_candidates: + candidates = filled_candidates + return min( + candidates, + key=lambda splits: _predict_decode_split_us( + num_splits=splits, + num_n_blocks=num_n_blocks, + total_mblocks=total_mblocks, + packed_q_rows=packed_q_rows, + tile_m=tile_m, + tile_n=tile_n, + head_dim=head_dim, + head_dim_v=head_dim_v, + element_size=element_size, + hardware=hardware, + ), + ) + + +def _select_low_hd_decode_num_splits( + *, + max_head_dim: int, + packed_q_rows: int, + total_mblocks: int, + num_n_blocks: int, + hardware: _DecodeHardware, +) -> int: + """Overschedule the short-M low-HD decode kernel by bounded waves. + + One hardware wave leaves each CTA with too long a serial K loop for this + one-warp kernel. For short KV, use the fewest splits that approximately + fill one device-specific SM wave; more waves only add combine work. For + longer KV, bound the launch by hardware waves and a minimum per-CTA KV + grain. Once eight-way splitting leaves a long grain, prefer balanced + power-of-two partitions over device-specific odd split counts. + """ + if num_n_blocks <= 4 or total_mblocks <= 0: + return 1 + num_sms = hardware.num_sms + if total_mblocks >= num_sms and (num_n_blocks <= 8 or num_n_blocks >= 64): + return 1 + if hardware.is_integrated and total_mblocks <= hardware.memory_channels // 2: + # Small batches need enough independent CTAs to cover UMA latency. + # Two thirds of the compact SM array is sufficient; larger grids only + # add combine work. + target_main_ctas = (2 * num_sms + 2) // 3 + fill_splits = min( + 32, + num_n_blocks, + max( + 1, + (target_main_ctas + total_mblocks - 1) // total_mblocks, + ), + ) + return _normalize_num_splits(fill_splits, num_n_blocks) + wave_split_budget = max(1, (3 * num_sms) // total_mblocks) + if num_n_blocks >= LOW_HD_DECODE_MAX_SPLITS * LOW_HD_DECODE_LONG_MIN_TILES_PER_CTA: + # Long KV can afford the next power-of-two split without making the + # per-CTA grain too small. Balanced partitions avoid partial waves and + # sustain memory bandwidth better than device-specific odd counts. + wave_split_budget = 1 << (wave_split_budget - 1).bit_length() + if num_n_blocks <= 8: + one_wave_splits = max(1, (num_sms + total_mblocks - 1) // total_mblocks) + if hardware.is_integrated: + one_wave_splits = 1 << (one_wave_splits - 1).bit_length() + kv_grain_splits = min(LOW_HD_DECODE_MAX_SPLITS, one_wave_splits) + else: + blocks_per_cta_floor = ( + 8 + if ( + max_head_dim >= 128 + and (hardware.is_integrated or 2 <= packed_q_rows <= 4) + ) + else 4 + ) + kv_grain_splits = min( + LOW_HD_DECODE_MAX_SPLITS, + num_n_blocks // blocks_per_cta_floor, + ) + return _normalize_num_splits( + min(wave_split_budget, kv_grain_splits), + num_n_blocks, + ) + + +def _tensor_signature(tensor: torch.Tensor) -> tuple: + return ( + type(tensor), + tensor.device, + tensor.dtype, + tensor.shape, + tensor.stride(), + tensor.requires_grad, + ) + + +def _resolve_causal_local_window( + causal: bool, + window_size_left: Optional[int], + window_size_right: Optional[int], +) -> tuple[bool, Optional[int], Optional[int]]: + if causal: + window_size_right = 0 + if ( + window_size_left is not None + and window_size_right is not None + and window_size_left + window_size_right < 0 + ): + window_size_left = None + window_size_right = None + if window_size_left is not None or window_size_right is not None: + if window_size_left is None and window_size_right == 0: + causal = True + window_size_right = None + else: + causal = False + return causal, window_size_left, window_size_right + + +class Sm120ForwardPolicy: + """Pure SM120 configuration, scheduling, and kernel-selection policy.""" + + @staticmethod + def supports_arch(arch: int) -> bool: + return arch // 10 == 12 + + @staticmethod + def use_graph_capture_split_combine_pdl( + *, + is_stream_capturing: bool, + is_split_kv: bool, + ) -> bool: + """Overlap SplitKV combine setup only when both grids are captured.""" + return is_stream_capturing and is_split_kv + + @staticmethod + @lru_cache(maxsize=1) + def implementation_token() -> tuple: + from sglang.kernels.ops.attention.fa4_sm120.flash_fwd_decode import ( + FlashAttentionForwardSm120DecodeTranspose, + ) + + return ( + FlashAttentionForwardSm120, + FlashAttentionForwardSm120.get_fwd_tile_size, + Sm120ForwardHost.resolve_plan, + Sm120ForwardHost.select_num_splits, + Sm120ForwardHost.select_paged_decode_split_kv_blocks_per_cta, + Sm120ForwardHost.use_paged_decode_transpose_qk_pv, + Sm120ForwardHost.make_kernel, + FlashAttentionForwardSm120DecodeTranspose, + FlashAttentionForwardSm120DecodeTranspose.paged_tma, + FlashAttentionForwardSm120DecodeTranspose.query_in_regs, + ) + + @classmethod + def resolve_plan( + cls, + *, + requested_num_splits: int, + generic_num_n_blocks: int, + head_dim: int, + head_dim_v: int, + batch_size: int, + num_head_kv: int, + paged_kv: bool, + page_size: Optional[int], + k: Optional[torch.Tensor], + v: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + pack_gqa: bool, + element_size: int, + packed_q_rows: int, + tile_m: int, + tile_n: int, + num_m_blocks: int, + total_mblocks: int, + num_sms: int, + total_q: int, + has_cu_seqlens_q: bool, + has_seqused_q: bool, + has_seqused_k: bool, + is_causal: bool, + is_local: bool, + window_size_left: Optional[int], + window_size_right: Optional[int], + has_score_or_mask_mod: bool, + is_stream_capturing: bool, + device: torch.device, + fake_mode: bool, + generic_heuristic: Callable[[int, int, int, int], int], + ) -> Sm120ForwardPlan: + """Resolve all SM120 dataflow decisions behind one host-side boundary.""" + has_compact_q_groups = pack_gqa or packed_q_rows == max_seqlen_q + split_num_n_blocks = cls.split_num_n_blocks( + generic_num_n_blocks=generic_num_n_blocks, + paged_kv=paged_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + is_local=is_local, + window_size_left=window_size_left, + window_size_right=window_size_right, + tile_n=tile_n, + ) + num_splits = cls.select_num_splits( + requested_num_splits=requested_num_splits, + head_dim=head_dim, + head_dim_v=head_dim_v, + paged_kv=paged_kv, + max_seqlen_q=max_seqlen_q, + has_compact_q_groups=has_compact_q_groups, + element_size=element_size, + packed_q_rows=packed_q_rows, + tile_m=tile_m, + tile_n=tile_n, + total_mblocks=total_mblocks, + num_sms=num_sms, + num_n_blocks=split_num_n_blocks, + device=device, + fake_mode=fake_mode, + generic_heuristic=generic_heuristic, + ) + is_split_kv = num_splits > 1 + direct_uniform_batch = cls.use_direct_uniform_batch( + batch_size=batch_size, + paged_kv=paged_kv, + has_cu_seqlens_q=has_cu_seqlens_q, + has_seqused_q=has_seqused_q, + total_q=total_q, + max_seqlen_q=max_seqlen_q, + num_m_blocks=num_m_blocks, + ) + split_qk_n = cls.use_paged_decode_split_qk_n( + head_dim=head_dim, + head_dim_v=head_dim_v, + paged_kv=paged_kv, + max_seqlen_q=max_seqlen_q, + has_compact_q_groups=has_compact_q_groups, + packed_q_rows=packed_q_rows, + tile_m=tile_m, + tile_n=tile_n, + num_n_blocks=generic_num_n_blocks, + is_causal=is_causal, + is_local=is_local, + has_score_or_mask_mod=has_score_or_mask_mod, + ) + dense_paged_kv = ( + page_size is not None + and k is not None + and k.stride(-1) == 1 + and v.stride(-1) == 1 + and k.stride(-2) == head_dim + and v.stride(-2) == head_dim_v + and k.stride(1) == num_head_kv * head_dim + and v.stride(1) == num_head_kv * head_dim_v + and k.stride(0) == page_size * num_head_kv * head_dim + and v.stride(0) == page_size * num_head_kv * head_dim_v + ) + transpose_qk_pv = cls.use_paged_decode_transpose_qk_pv( + head_dim=head_dim, + head_dim_v=head_dim_v, + batch_size=batch_size, + num_head_kv=num_head_kv, + paged_kv=paged_kv, + page_size=page_size, + dense_paged_kv=dense_paged_kv, + max_seqlen_q=max_seqlen_q, + pack_gqa=pack_gqa, + packed_q_rows=packed_q_rows, + tile_m=tile_m, + tile_n=tile_n, + num_splits=num_splits, + num_n_blocks=split_num_n_blocks, + is_causal=is_causal, + is_local=is_local, + has_score_or_mask_mod=has_score_or_mask_mod, + ) + if transpose_qk_pv: + split_qk_n = False + split_kv_blocks_per_cta = cls.select_paged_decode_split_kv_blocks_per_cta( + head_dim=head_dim, + head_dim_v=head_dim_v, + batch_size=batch_size, + paged_kv=paged_kv, + max_seqlen_q=max_seqlen_q, + has_compact_q_groups=has_compact_q_groups, + packed_q_rows=packed_q_rows, + is_split_kv=is_split_kv, + direct_uniform_batch=direct_uniform_batch, + has_seqused_k=has_seqused_k, + split_qk_n=split_qk_n, + num_splits=num_splits, + num_n_blocks=split_num_n_blocks, + ) + return Sm120ForwardPlan( + num_splits=num_splits, + split_num_n_blocks=split_num_n_blocks, + direct_uniform_batch=direct_uniform_batch, + split_qk_n=split_qk_n, + split_kv_blocks_per_cta=split_kv_blocks_per_cta, + transpose_qk_pv=transpose_qk_pv, + launch_split_combine_early=cls.use_graph_capture_split_combine_pdl( + is_stream_capturing=is_stream_capturing, + is_split_kv=is_split_kv, + ), + ) + + @staticmethod + def compile_arguments(plan: Sm120ForwardPlan) -> tuple: + """Extra CuTe arguments required by the selected SM120 kernel ABI.""" + return (Int32(0),) + + @staticmethod + def runtime_arguments(plan: Sm120ForwardPlan) -> tuple: + """Extra runtime arguments required by the selected SM120 kernel ABI.""" + return (int(plan.launch_split_combine_early),) + + @staticmethod + def select_config( + *, + head_dim: int, + head_dim_v: int, + tile_mn: Optional[tuple[int, int]], + has_bias: bool, + total_q_rows: int, + num_sms: Optional[int], + num_batch: int, + seqlen_q: Optional[int], + seqlen_k: Optional[int], + num_head_kv: int, + qhead_per_kvhead: int, + is_causal: bool, + is_local: bool, + window_size_left: Optional[int], + window_size_right: Optional[int], + pack_gqa: bool, + paged_kv: bool, + ) -> Sm120ForwardConfig: + if has_bias: + if max(head_dim, head_dim_v) > 128: + raise ValueError( + "SM120 relative bias currently supports head_dim and " + "head_dim_v up to 128" + ) + if tile_mn is not None and tile_mn != (64, 128): + raise ValueError("SM120 relative bias requires tile_mn=(64, 128)") + tile_m, tile_n = 64, 128 + elif tile_mn is None: + tile_m, tile_n = FlashAttentionForwardSm120.get_fwd_tile_size( + head_dim, + head_dim_v, + total_q_rows=total_q_rows, + num_sms=num_sms, + num_batch=num_batch, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + num_head_kv=num_head_kv, + qhead_per_kvhead=qhead_per_kvhead, + is_causal=is_causal, + is_local=is_local, + window_size_left=window_size_left, + window_size_right=window_size_right, + pack_gqa=pack_gqa, + paged_kv=paged_kv, + ) + else: + tile_m, tile_n = tile_mn + return Sm120ForwardConfig( + tile_m=tile_m, + tile_n=tile_n, + num_stages=FlashAttentionForwardSm120.get_fwd_num_stages( + head_dim, head_dim_v, tile_m, tile_n + ), + num_threads=FlashAttentionForwardSm120.get_fwd_num_threads( + head_dim, + head_dim_v, + tile_m, + tile_n, + paged_kv=paged_kv, + ), + ) + + @staticmethod + def select_num_splits( + *, + requested_num_splits: int, + head_dim: int, + head_dim_v: int, + paged_kv: bool, + max_seqlen_q: int, + has_compact_q_groups: bool, + element_size: int, + packed_q_rows: int, + tile_m: int, + tile_n: int, + total_mblocks: int, + num_sms: int, + num_n_blocks: int, + device: torch.device, + fake_mode: bool, + generic_heuristic: Callable[[int, int, int, int], int], + ) -> int: + num_splits = requested_num_splits + if num_splits < 1: + is_hd256_paged_decode = ( + head_dim == 256 + and head_dim_v == 256 + and paged_kv + and max_seqlen_q == 1 + and has_compact_q_groups + ) + is_low_hd_paged_decode = is_low_hd_paged_decode_tile( + head_dim=head_dim, + head_dim_v=head_dim_v, + paged_kv=paged_kv, + seqlen_q=max_seqlen_q, + tile_m=tile_m, + tile_n=tile_n, + ) + if is_low_hd_paged_decode and tile_m == 16: + num_splits = _select_low_hd_decode_num_splits( + max_head_dim=max(head_dim, head_dim_v), + packed_q_rows=packed_q_rows, + total_mblocks=total_mblocks, + num_n_blocks=num_n_blocks, + hardware=_get_decode_hardware(device), + ) + elif is_hd256_paged_decode: + if fake_mode: + max_splits = min(128, 1 << (num_sms.bit_length() - 1)) + num_splits = generic_heuristic( + total_mblocks, + num_sms, + num_n_blocks, + max_splits, + ) + else: + num_splits = _select_decode_num_splits( + head_dim=head_dim, + head_dim_v=head_dim_v, + element_size=element_size, + packed_q_rows=packed_q_rows, + tile_m=tile_m, + tile_n=tile_n, + total_mblocks=total_mblocks, + num_n_blocks=num_n_blocks, + hardware=_get_decode_hardware(device), + ) + else: + num_splits = generic_heuristic( + total_mblocks, + num_sms, + num_n_blocks, + 128, + ) + + if num_splits <= 1 or num_n_blocks == 0: + return 1 + + # BlockInfo assigns a uniform ceil-div chunk to every split. Normalize + # the count so no SM120 CTA owns an empty tail chunk. + return _normalize_num_splits(num_splits, num_n_blocks) + + @staticmethod + def split_num_n_blocks( + *, + generic_num_n_blocks: int, + paged_kv: bool, + max_seqlen_q: int, + max_seqlen_k: int, + is_local: bool, + window_size_left: Optional[int], + window_size_right: Optional[int], + tile_n: int, + ) -> int: + """Return the KV-block count used by the decode SplitKV policy. + + The generic local-attention bound includes a full query tile because + prefill tiles can span multiple query positions. Packed decode has one + query position regardless of how many GQA heads occupy the M tile, so + its exact visible span is just left + current + right. + """ + if not (paged_kv and max_seqlen_q == 1 and is_local): + return generic_num_n_blocks + visible_seqlen_k = visible_decode_seqlen_k( + max_seqlen_k, + is_local=True, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + return max(0, (visible_seqlen_k + tile_n - 1) // tile_n) + + @staticmethod + def use_direct_uniform_batch( + *, + batch_size: int, + paged_kv: bool, + has_cu_seqlens_q: bool, + has_seqused_q: bool, + total_q: int, + max_seqlen_q: int, + num_m_blocks: int, + ) -> bool: + """Use arithmetic scheduling for one uniform query row per request. + + This invariant is independent of whether GQA heads are folded into the + query-row mode. In particular, unpacked MHA should not pay the generic + varlen scheduler's prefix-sum walk merely to preserve its ordinary Q + loader. + """ + return ( + paged_kv + and has_cu_seqlens_q + and not has_seqused_q + and max_seqlen_q == 1 + and total_q == batch_size * max_seqlen_q + and num_m_blocks == 1 + ) + + @staticmethod + def use_paged_decode_split_qk_n( + *, + head_dim: int, + head_dim_v: int, + paged_kv: bool, + max_seqlen_q: int, + has_compact_q_groups: bool, + packed_q_rows: int, + tile_m: int, + tile_n: int, + num_n_blocks: int, + is_causal: bool, + is_local: bool, + has_score_or_mask_mod: bool, + ) -> bool: + """Select four-way N-distributed QK for qualified SM120 decode. + + N-distribution removes the single-QK-warp critical path, lowers the + register footprint, and is faster from the first K/V tile on both + qualified SM120 SKUs. Both structures use the same threads and + one-CTA-per-SM shared-memory residency. + """ + return ( + (head_dim, head_dim_v) == (256, 256) + and paged_kv + and max_seqlen_q == 1 + and has_compact_q_groups + and packed_q_rows <= 16 + and (tile_m, tile_n) == (16, 64) + and num_n_blocks >= 1 + and (is_causal or is_local) + and not has_score_or_mask_mod + ) + + @staticmethod + def select_paged_decode_split_kv_blocks_per_cta( + *, + head_dim: int, + head_dim_v: int, + batch_size: int, + paged_kv: bool, + max_seqlen_q: int, + has_compact_q_groups: bool, + packed_q_rows: int, + is_split_kv: bool, + direct_uniform_batch: bool, + has_seqused_k: bool, + split_qk_n: bool, + num_splits: int, + num_n_blocks: int, + ) -> int: + """Preserve the longest request's split grain for ragged decode.""" + if not ( + (head_dim, head_dim_v) == (256, 256) + and batch_size > 1 + and paged_kv + and max_seqlen_q == 1 + and has_compact_q_groups + and packed_q_rows <= 16 + and is_split_kv + and direct_uniform_batch + and has_seqused_k + and split_qk_n + and num_splits > 1 + and num_n_blocks > 0 + ): + return 0 + return (num_n_blocks + num_splits - 1) // num_splits + + @staticmethod + def use_paged_decode_transpose_qk_pv( + *, + head_dim: int, + head_dim_v: int, + batch_size: int, + num_head_kv: int, + paged_kv: bool, + page_size: Optional[int], + dense_paged_kv: bool, + max_seqlen_q: int, + pack_gqa: bool, + packed_q_rows: int, + tile_m: int, + tile_n: int, + num_splits: int, + num_n_blocks: int, + is_causal: bool, + is_local: bool, + has_score_or_mask_mod: bool, + ) -> bool: + """Select the qualified page-TMA M64N8 decode dataflow. + + Full QK+PV transpose removes M16 padding for two through eight packed + query rows. Interleaved NCU on the 110-SM/384-bit and 188-SM/512-bit + SM120 SKUs qualifies up to eight batch/head groups. One or two groups + amortize the transpose with two KV tiles per CTA; three through eight + groups require four. + """ + effective_splits = max(1, num_splits) + kv_tiles_per_cta = ( + (num_n_blocks + effective_splits - 1) // effective_splits + if num_n_blocks > 0 + else 0 + ) + batch_head_groups = batch_size * num_head_kv + min_kv_tiles_per_cta = 2 if batch_head_groups <= 2 else 4 + return ( + (head_dim, head_dim_v) == (256, 256) + and 1 <= batch_head_groups <= 8 + and paged_kv + and page_size == tile_n + and dense_paged_kv + and max_seqlen_q == 1 + and pack_gqa + and 2 <= packed_q_rows <= 8 + and (tile_m, tile_n) == (16, 64) + and kv_tiles_per_cta >= min_kv_tiles_per_cta + and (is_causal or is_local) + and not has_score_or_mask_mod + ) + + @staticmethod + def make_kernel( + *, + dtype, + head_dim: int, + head_dim_v: int, + qhead_per_kvhead: int, + is_causal: bool, + is_local: bool, + pack_gqa: bool, + config: Sm120ForwardConfig, + paged_kv: bool, + score_mod: Optional[Callable], + mask_mod: Optional[Callable], + has_aux_tensors: bool, + is_split_kv: bool, + has_bias: bool, + bias_block_size: int, + rel_extent_padded: int, + plan: Sm120ForwardPlan, + ) -> FlashAttentionForwardSm120: + if not FlashAttentionForwardSm120.can_implement( + dtype, + head_dim, + head_dim_v, + config.tile_m, + config.tile_n, + num_stages=config.num_stages, + num_threads=config.num_threads, + is_causal=is_causal, + Q_in_regs=False, + paged_kv=paged_kv, + ): + raise ValueError( + "The requested FlashAttention forward configuration exceeds " + "SM120 kernel constraints or shared-memory capacity" + ) + if has_bias: + bias_smem_bytes = ( + bias_block_size * config.tile_n * (dtype.width // 8) * config.num_stages + ) + total_smem_bytes = ( + FlashAttentionForwardSm120._smem_usage_in_bytes( + head_dim, + head_dim_v, + config.tile_m, + config.tile_n, + config.num_stages, + False, + ) + + bias_smem_bytes + ) + if total_smem_bytes > utils_basic.get_smem_capacity_in_bytes("sm_120"): + raise ValueError( + "The requested SM120 sheared-bias specialization exceeds " + "shared-memory capacity" + ) + Kernel = FlashAttentionForwardSm120 + if plan.transpose_qk_pv: + from sglang.kernels.ops.attention.fa4_sm120.flash_fwd_decode import ( + FlashAttentionForwardSm120DecodeTranspose, + ) + + Kernel = FlashAttentionForwardSm120DecodeTranspose + return Kernel( + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + is_causal=is_causal, + is_local=is_local, + pack_gqa=pack_gqa, + tile_m=config.tile_m, + tile_n=config.tile_n, + num_stages=config.num_stages, + num_threads=config.num_threads, + Q_in_regs=False, + score_mod=score_mod, + mask_mod=mask_mod, + has_aux_tensors=has_aux_tensors, + is_split_kv=is_split_kv, + has_bias=has_bias, + bias_block_size=bias_block_size, + rel_extent_padded=rel_extent_padded, + direct_uniform_batch=plan.direct_uniform_batch, + paged_kv=paged_kv, + split_qk_n=plan.split_qk_n, + split_kv_blocks_per_cta=plan.split_kv_blocks_per_cta, + ) + + +class Sm120ForwardHost(Sm120ForwardPolicy): + """Mutable direct-launch plans and workspaces for the SM120 policy.""" + + def __init__(self) -> None: + self._varlen_plans: OrderedDict[tuple, _VarlenLaunchPlan] = OrderedDict() + self._paged_plans: OrderedDict[tuple, _PagedDecodeLaunchPlan] = OrderedDict() + self._paged_plan_tiles: OrderedDict[tuple, tuple[int, int]] = OrderedDict() + self._paged_workspaces: OrderedDict[ + tuple[torch.device, int], tuple[torch.Tensor, torch.Tensor] + ] = OrderedDict() + self._paged_workspace_views: OrderedDict[ + tuple[torch.device, int, int, tuple[int, ...], int, torch.dtype], + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ] = OrderedDict() + + def _varlen_key( + self, + *, + arch: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool, + window_size_left: Optional[int], + window_size_right: Optional[int], + learnable_sink: Optional[torch.Tensor], + pack_gqa: bool, + ) -> tuple: + sink_signature = ( + None if learnable_sink is None else _tensor_signature(learnable_sink) + ) + return ( + "basic-varlen", + arch, + tuple(_tensor_signature(t) for t in (q, k, v, cu_seqlens_q, cu_seqlens_k)), + sink_signature, + max_seqlen_q, + max_seqlen_k, + causal, + window_size_left, + window_size_right, + pack_gqa, + self.implementation_token(), + fa_logging.get_fa_log_level(), + ) + + def try_varlen( + self, + *, + arch: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: Optional[int], + max_seqlen_k: Optional[int], + softmax_scale: Optional[float], + causal: bool, + window_size: tuple[Optional[int], Optional[int]], + learnable_sink: Optional[torch.Tensor], + pack_gqa: Optional[bool], + out: Optional[torch.Tensor], + ) -> Optional[tuple[torch.Tensor, None]]: + if ( + not self.supports_arch(arch) + or q.ndim != 3 + or k.ndim != 3 + or k.shape[1] == 0 + ): + return None + expected_out_shape = (*q.shape[:-1], v.shape[-1]) + if out is not None and ( + out.shape != expected_out_shape + or out.dtype != q.dtype + or out.device != q.device + or not out.is_contiguous() + ): + return None + actual_pack_gqa = q.shape[1] // k.shape[1] > 1 if pack_gqa is None else pack_gqa + actual_max_seqlen_q = q.shape[0] if max_seqlen_q is None else max_seqlen_q + actual_max_seqlen_k = k.shape[0] if max_seqlen_k is None else max_seqlen_k + causal, window_size_left, window_size_right = _resolve_causal_local_window( + causal, + window_size[0], + window_size[1], + ) + key = self._varlen_key( + arch=arch, + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=actual_max_seqlen_q, + max_seqlen_k=actual_max_seqlen_k, + causal=causal, + window_size_left=window_size_left, + window_size_right=window_size_right, + learnable_sink=learnable_sink, + pack_gqa=actual_pack_gqa, + ) + plan = self._varlen_plans.get(key) + if plan is None: + return None + self._varlen_plans.move_to_end(key) + if out is None: + out = torch.empty( + expected_out_shape, + dtype=q.dtype, + device=q.device, + ) + scale = 1.0 / math.sqrt(q.shape[-1]) if softmax_scale is None else softmax_scale + plan.compiled_fn( + q, + k, + v, + out, + None, + scale, + cu_seqlens_q, + cu_seqlens_k, + None, + None, + None, + window_size_left, + window_size_right, + learnable_sink, + None, + _EMPTY_AUX_DATA, + None, + 0, + ) + return out, None + + def register_varlen( + self, + *, + arch: int, + compiled_fn: Callable, + compile_key: tuple, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool, + window_size_left: Optional[int], + window_size_right: Optional[int], + learnable_sink: Optional[torch.Tensor], + pack_gqa: bool, + ) -> None: + key = self._varlen_key( + arch=arch, + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=causal, + window_size_left=window_size_left, + window_size_right=window_size_right, + learnable_sink=learnable_sink, + pack_gqa=pack_gqa, + ) + self._varlen_plans[key] = _VarlenLaunchPlan(compiled_fn, compile_key) + self._varlen_plans.move_to_end(key) + while len(self._varlen_plans) > _LAUNCH_PLAN_CAPACITY: + self._varlen_plans.popitem(last=False) + + def _paged_base_key( + self, + *, + arch: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + seqused_k: torch.Tensor, + page_table: torch.Tensor, + max_seqlen_q: int, + causal: bool, + window_size_left: Optional[int], + window_size_right: Optional[int], + learnable_sink: Optional[torch.Tensor], + pack_gqa: bool, + requested_num_splits: int, + ) -> tuple: + sink_signature = ( + None if learnable_sink is None else _tensor_signature(learnable_sink) + ) + return ( + "paged-forward", + arch, + tuple( + _tensor_signature(t) + for t in (q, k, v, cu_seqlens_q, seqused_k, page_table) + ), + sink_signature, + max_seqlen_q, + causal, + window_size_left, + window_size_right, + pack_gqa, + requested_num_splits, + self.implementation_token(), + fa_logging.get_fa_log_level(), + ) + + @staticmethod + def _paged_selection_key( + base_key: tuple, + tile_m: int, + tile_n: int, + max_seqlen_k: int, + window_size_left: Optional[int], + window_size_right: Optional[int], + ) -> tuple: + is_local = window_size_left is not None or window_size_right is not None + if is_local: + left = max_seqlen_k if window_size_left is None else window_size_left + right = max_seqlen_k if window_size_right is None else window_size_right + seqlen_k_loaded = max( + 0, + min(max_seqlen_k, left + 1 + right), + ) + else: + seqlen_k_loaded = max_seqlen_k + num_n_blocks = (seqlen_k_loaded + tile_n - 1) // tile_n + return (*base_key, (tile_m, tile_n, num_n_blocks)) + + def _paged_workspace( + self, + plan: _PagedDecodeLaunchPlan, + q: torch.Tensor, + v: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + stream_key = torch.cuda.current_stream(q.device).cuda_stream + key = (q.device, stream_key) + view_key = ( + q.device, + stream_key, + plan.actual_num_splits, + tuple(q.shape[:-1]), + v.shape[-1], + plan.partial_dtype, + ) + cached_view = self._paged_workspace_views.get(view_key) + if cached_view is not None: + self._paged_workspace_views.move_to_end(view_key) + return cached_view + out_numel = plan.actual_num_splits * math.prod(q.shape[:-1]) * v.shape[-1] + lse_numel = plan.actual_num_splits * q.shape[-2] * q.shape[0] + workspace = self._paged_workspaces.get(key) + if workspace is not None: + self._paged_workspaces.move_to_end(key) + if ( + workspace is None + or workspace[0].numel() < out_numel + or workspace[1].numel() < lse_numel + or workspace[0].dtype != plan.partial_dtype + ): + self._drop_workspace_views(key) + workspace = ( + torch.empty(out_numel, dtype=plan.partial_dtype, device=q.device), + torch.empty(lse_numel, dtype=torch.float32, device=q.device), + ) + self._cache_workspace(key, workspace) + out_partial = workspace[0][:out_numel].view( + plan.actual_num_splits, + *q.shape[:-1], + v.shape[-1], + ) + lse_partial = workspace[1][:lse_numel].view( + plan.actual_num_splits, + q.shape[-2], + q.shape[0], + ) + result = (out_partial, lse_partial, lse_partial.transpose(-1, -2)) + self._cache_workspace_view(view_key, result) + return result + + @staticmethod + def _supports_cached_paged_decode( + *, + head_dim: int, + head_dim_v: int, + effective_q_rows: int, + ) -> bool: + """Qualify shapes covered by the SM12x paged-decode launch cache.""" + head_dims = (head_dim, head_dim_v) + return ( + head_dims in LOW_HD_DECODE_SHAPES or head_dims == (256, 256) + ) and effective_q_rows <= 16 + + def try_paged_decode( + self, + *, + arch: int, + q: Optional[torch.Tensor], + k: Optional[torch.Tensor], + v: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor], + cu_seqlens_k: Optional[torch.Tensor], + seqused_q: Optional[torch.Tensor], + seqused_k: Optional[torch.Tensor], + page_table: Optional[torch.Tensor], + max_seqlen_q: Optional[int], + max_seqlen_k: Optional[int], + softmax_scale: Optional[float], + causal: bool, + window_size_left: Optional[int], + window_size_right: Optional[int], + learnable_sink: Optional[torch.Tensor], + requested_num_splits: int, + pack_gqa: Optional[bool], + out: Optional[torch.Tensor], + ) -> Optional[tuple[torch.Tensor, None]]: + if ( + not self.supports_arch(arch) + or q is None + or k is None + or q.ndim != 3 + or k.ndim != 4 + or cu_seqlens_q is None + or cu_seqlens_k is not None + or seqused_q is not None + or seqused_k is None + or page_table is None + or torch.cuda.is_current_stream_capturing() + or any(t.requires_grad for t in (q, k, v)) + or (learnable_sink is not None and learnable_sink.requires_grad) + ): + return None + expected_out_shape = (*q.shape[:-1], v.shape[-1]) + if out is not None and ( + out.shape != expected_out_shape + or out.dtype != q.dtype + or out.device != q.device + or not out.is_contiguous() + ): + return None + actual_pack_gqa = ( + q.shape[1] // k.shape[-2] > 1 if pack_gqa is None else pack_gqa + ) + actual_max_seqlen_q = q.shape[0] if max_seqlen_q is None else max_seqlen_q + actual_max_seqlen_k = ( + k.shape[0] * k.shape[1] if max_seqlen_k is None else max_seqlen_k + ) + qhead_per_kvhead = q.shape[1] // k.shape[-2] + effective_q_rows = actual_max_seqlen_q * ( + qhead_per_kvhead if actual_pack_gqa else 1 + ) + if not self._supports_cached_paged_decode( + head_dim=q.shape[-1], + head_dim_v=v.shape[-1], + effective_q_rows=effective_q_rows, + ): + return None + causal, window_size_left, window_size_right = _resolve_causal_local_window( + causal, + window_size_left, + window_size_right, + ) + base_key = self._paged_base_key( + arch=arch, + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + seqused_k=seqused_k, + page_table=page_table, + max_seqlen_q=actual_max_seqlen_q, + causal=causal, + window_size_left=window_size_left, + window_size_right=window_size_right, + learnable_sink=learnable_sink, + pack_gqa=actual_pack_gqa, + requested_num_splits=requested_num_splits, + ) + tile_mn = self._paged_plan_tiles.get(base_key) + if tile_mn is None: + return None + self._paged_plan_tiles.move_to_end(base_key) + selection_key = self._paged_selection_key( + base_key, + *tile_mn, + actual_max_seqlen_k, + window_size_left, + window_size_right, + ) + plan = self._paged_plans.get(selection_key) + if plan is None: + return None + self._paged_plans.move_to_end(selection_key) + if out is None: + out = torch.empty( + expected_out_shape, + dtype=q.dtype, + device=q.device, + ) + scale = 1.0 / math.sqrt(q.shape[-1]) if softmax_scale is None else softmax_scale + kernel_out = out + kernel_lse = None + lse_partial_transposed = None + if plan.actual_num_splits > 1: + if plan.compiled_combine is None: + return None + kernel_out, kernel_lse, lse_partial_transposed = self._paged_workspace( + plan, q, v + ) + plan.compiled_fn( + q, + k, + v, + kernel_out, + kernel_lse, + scale, + cu_seqlens_q, + None, + None, + seqused_k, + page_table, + window_size_left, + window_size_right, + learnable_sink, + None, + _EMPTY_AUX_DATA, + None, + 0, + ) + if plan.actual_num_splits > 1: + plan.compiled_combine( + kernel_out, + lse_partial_transposed, + out, + None, + cu_seqlens_q, + None, + None, + None, + None, + ) + return out, None + + def register_paged_decode( + self, + *, + arch: int, + compiled_fn: Callable, + compile_key: tuple, + compiled_combine: Optional[Callable], + actual_num_splits: int, + tile_m: int, + tile_n: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + seqused_k: torch.Tensor, + page_table: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool, + window_size_left: Optional[int], + window_size_right: Optional[int], + learnable_sink: Optional[torch.Tensor], + pack_gqa: bool, + requested_num_splits: int, + out_partial: Optional[torch.Tensor], + lse_partial: Optional[torch.Tensor], + ) -> None: + qhead_per_kvhead = q.shape[1] // k.shape[-2] + effective_q_rows = max_seqlen_q * (qhead_per_kvhead if pack_gqa else 1) + if not self._supports_cached_paged_decode( + head_dim=q.shape[-1], + head_dim_v=v.shape[-1], + effective_q_rows=effective_q_rows, + ): + return + base_key = self._paged_base_key( + arch=arch, + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + seqused_k=seqused_k, + page_table=page_table, + max_seqlen_q=max_seqlen_q, + causal=causal, + window_size_left=window_size_left, + window_size_right=window_size_right, + learnable_sink=learnable_sink, + pack_gqa=pack_gqa, + requested_num_splits=requested_num_splits, + ) + selection_key = self._paged_selection_key( + base_key, + tile_m, + tile_n, + max_seqlen_k, + window_size_left, + window_size_right, + ) + self._paged_plan_tiles[base_key] = (tile_m, tile_n) + self._paged_plan_tiles.move_to_end(base_key) + while len(self._paged_plan_tiles) > _LAUNCH_PLAN_CAPACITY: + stale_base_key, _ = self._paged_plan_tiles.popitem(last=False) + stale_selection_keys = [ + key for key in self._paged_plans if key[:-1] == stale_base_key + ] + for key in stale_selection_keys: + del self._paged_plans[key] + self._paged_plans[selection_key] = _PagedDecodeLaunchPlan( + compiled_fn=compiled_fn, + compile_key=compile_key, + actual_num_splits=actual_num_splits, + compiled_combine=compiled_combine, + partial_dtype=(out_partial.dtype if out_partial is not None else q.dtype), + ) + self._paged_plans.move_to_end(selection_key) + while len(self._paged_plans) > _LAUNCH_PLAN_CAPACITY: + self._paged_plans.popitem(last=False) + + if out_partial is None or lse_partial is None: + return + stream_key = torch.cuda.current_stream(q.device).cuda_stream + workspace_key = (q.device, stream_key) + current_workspace = self._paged_workspaces.get(workspace_key) + if current_workspace is not None: + self._paged_workspaces.move_to_end(workspace_key) + if ( + current_workspace is None + or current_workspace[0].numel() < out_partial.numel() + or current_workspace[1].numel() < lse_partial.numel() + or current_workspace[0].dtype != out_partial.dtype + ): + self._drop_workspace_views(workspace_key) + self._cache_workspace( + workspace_key, + ( + out_partial.view(-1), + lse_partial.view(-1), + ), + ) + view_key = ( + q.device, + stream_key, + actual_num_splits, + tuple(q.shape[:-1]), + v.shape[-1], + out_partial.dtype, + ) + workspace = self._paged_workspaces[workspace_key] + out_partial_view = workspace[0][: out_partial.numel()].view(out_partial.shape) + lse_partial_view = workspace[1][: lse_partial.numel()].view(lse_partial.shape) + self._cache_workspace_view( + view_key, + ( + out_partial_view, + lse_partial_view, + lse_partial_view.transpose(-1, -2), + ), + ) + + def _cache_workspace_view( + self, + key: tuple, + value: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ) -> None: + self._paged_workspace_views[key] = value + self._paged_workspace_views.move_to_end(key) + while len(self._paged_workspace_views) > _WORKSPACE_VIEW_CAPACITY: + self._paged_workspace_views.popitem(last=False) + + def _cache_workspace( + self, + key: tuple[torch.device, int], + value: tuple[torch.Tensor, torch.Tensor], + ) -> None: + self._paged_workspaces[key] = value + self._paged_workspaces.move_to_end(key) + while len(self._paged_workspaces) > _WORKSPACE_CAPACITY: + stale_key, _ = self._paged_workspaces.popitem(last=False) + self._drop_workspace_views(stale_key) + + def _drop_workspace_views(self, workspace_key: tuple[torch.device, int]) -> None: + stale_keys = [ + key for key in self._paged_workspace_views if key[:2] == workspace_key + ] + for key in stale_keys: + del self._paged_workspace_views[key] + + def clear_launch_plans(self) -> None: + self._varlen_plans.clear() + self._paged_plans.clear() + self._paged_plan_tiles.clear() + self._paged_workspaces.clear() + self._paged_workspace_views.clear() + + +sm120_forward_host = Sm120ForwardHost() diff --git a/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py b/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py new file mode 100644 index 000000000..9cc7cb365 --- /dev/null +++ b/python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py @@ -0,0 +1,150 @@ +# Copyright (c) 2026, SGLang Team. +"""Schedulers owned by the SGLang SM120 FA4 implementation.""" + +from dataclasses import dataclass +from typing import Tuple + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 +from quack.cute_dsl_utils import ParamsBase + +from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import ( + SchedulingMode, + TileSchedulerArguments, + WorkTileInfo, +) + + +class Sm120UniformBatchScheduler: + """Map uniform varlen batches without the generic prefix-sum walk. + + SM120 paged decode uses one equally sized query segment per request. Its + compile-time dispatch proves that invariant before selecting this scheduler, + so each CTA can recover ``(block, head, batch)`` arithmetically. + """ + + @dataclass + class Params(ParamsBase): + num_head: Int32 + num_batch: Int32 + total_q: Int32 + num_splits: Int32 + tile_shape_mn: cutlass.Constexpr[Tuple[int, int]] + is_split_kv: cutlass.Constexpr[bool] = False + + @staticmethod + @cute.jit + def create( + args: TileSchedulerArguments, *, loc=None, ip=None + ) -> "Sm120UniformBatchScheduler.Params": + assert args.cluster_shape_mn == ( + 1, + 1, + ), "SM120 uniform-batch scheduling requires a 1x1 cluster" + return Sm120UniformBatchScheduler.Params( + num_head=args.num_head, + num_batch=args.num_batch, + total_q=args.total_q, + num_splits=args.num_splits, + tile_shape_mn=args.tile_shape_mn, + is_split_kv=args.is_split_kv, + ) + + def __init__( + self, + params: Params, + tile_idx: Int32, + split_idx: Int32, + *, + loc=None, + ip=None, + ): + self.params = params + self._tile_idx = tile_idx + self._split_idx = split_idx + self._is_first_block = True + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + assert ( + scheduling_mode == SchedulingMode.STATIC + ), f"SM120 uniform-batch scheduler only supports STATIC, got {scheduling_mode!r}" + return Sm120UniformBatchScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + @cute.jit + def create( + params: Params, clc=None, *, loc=None, ip=None + ) -> "Sm120UniformBatchScheduler": + tile_idx, split_idx, _ = cute.arch.block_idx() + return Sm120UniformBatchScheduler(params, tile_idx, split_idx, loc=loc, ip=ip) + + @staticmethod + @cute.jit + def get_grid_shape( + params: Params, *, loc=None, ip=None + ) -> Tuple[Int32, Int32, Int32]: + rows_per_batch = params.total_q // params.num_batch + num_m_blocks = cute.ceil_div(rows_per_batch, params.tile_shape_mn[0]) + return ( + num_m_blocks * params.num_head * params.num_batch, + params.num_splits, + Int32(1), + ) + + @cute.jit + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + params = self.params + rows_per_batch = params.total_q // params.num_batch + num_m_blocks = cute.ceil_div(rows_per_batch, params.tile_shape_mn[0]) + mh_blocks_per_batch = num_m_blocks * params.num_head + batch_idx = self._tile_idx // mh_blocks_per_batch + mh_block = self._tile_idx - batch_idx * mh_blocks_per_batch + block = mh_block // params.num_head + head_idx = mh_block - block * params.num_head + split_idx = ( + self._split_idx if cutlass.const_expr(params.is_split_kv) else Int32(0) + ) + return WorkTileInfo( + (Int32(block), Int32(head_idx), Int32(batch_idx), split_idx), + self._is_first_block, + ) + + def initial_work_tile_info(self, *, loc=None, ip=None) -> WorkTileInfo: + return self.get_current_work(loc=loc, ip=ip) + + def prefetch_next_work(self, *, loc=None, ip=None): + pass + + def advance_to_next_work(self, *, loc=None, ip=None) -> WorkTileInfo: + self._is_first_block = False + return self.get_current_work(loc=loc, ip=ip) + + def producer_tail(self, *, loc=None, ip=None): + pass + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in (self.params, self._tile_idx, self._split_idx): + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + objects = [] + for obj, n_items in zip( + (self.params, self._tile_idx, self._split_idx), self._values_pos + ): + objects.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return Sm120UniformBatchScheduler(*objects, loc=self._loc) diff --git a/python/sglang/kernels/ops/attention/flash_attention.py b/python/sglang/kernels/ops/attention/flash_attention.py index cfd8ca4cb..5e9e39837 100644 --- a/python/sglang/kernels/ops/attention/flash_attention.py +++ b/python/sglang/kernels/ops/attention/flash_attention.py @@ -48,6 +48,7 @@ def flash_attn_with_kvcache( rel_bias_prep_cache=None, ver=3, out=None, + max_seqlen_k: Optional[int] = None, ): """ If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from @@ -194,6 +195,7 @@ def flash_attn_with_kvcache( page_table=page_table, cu_seqlens_q=cu_seqlens_q, max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, rotary_seqlens=rotary_seqlens, q_descale=q_descale, k_descale=k_descale, @@ -213,6 +215,7 @@ def flash_attn_with_kvcache( rel_bias=rel_bias, rel_bias_prep_cache=rel_bias_prep_cache, return_softmax_lse=return_softmax_lse, + out=out, ) else: raise RuntimeError(f"Unknown flash attention version {ver}") @@ -319,6 +322,7 @@ def flash_attn_varlen_func( rel_bias=rel_bias, rel_bias_prep_cache=rel_bias_prep_cache, return_softmax_lse=return_softmax_lse, + out=out, ) else: raise RuntimeError(f"Unknown flash attention version {ver}") diff --git a/python/sglang/kernels/ops/attention/flash_attention_v4_sm120.py b/python/sglang/kernels/ops/attention/flash_attention_v4_sm120.py new file mode 100644 index 000000000..460ad3ab3 --- /dev/null +++ b/python/sglang/kernels/ops/attention/flash_attention_v4_sm120.py @@ -0,0 +1,339 @@ +# Copyright (c) 2026, SGLang Team. +"""SGLang-facing FlashAttention-4 APIs specialized for SM12x.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Callable, Optional, Tuple, Union + +import torch + +from sglang.kernel_api_logging import debug_kernel_api +from sglang.kernels.ops.attention.flash_attention_v4 import ( + _flash_attn_import_error, + _flash_attn_varlen_func, + _maybe_contiguous, + _pad_mla_q_heads, + _unpad_mla_result, +) + +if os.environ.get("SGLANG_INKLING_FA4_USE_PIP") == "1": + # The pip escape hatch deliberately bypasses SGLang-owned SM12x kernels. + get_forward_arch = None + resolve_runtime_policy = None + try_cached_paged_decode = None +else: + from sglang.kernels.ops.attention.fa4_sm120.dispatch import ( + get_forward_arch, + resolve_runtime_policy, + try_cached_paged_decode, + ) + + +@dataclass(frozen=True) +class FlashAttentionV4SM120RuntimePolicy: + num_splits: int + decode_num_splits: int + decode_uses_static_max_seqlen_k: bool + + +def get_flash_attention_v4_sm120_runtime_policy( + *, + device_capability: tuple[int, int], + deterministic: bool, +) -> FlashAttentionV4SM120RuntimePolicy: + """Resolve the SM12x FA4 launch policy exposed to SGLang.""" + if resolve_runtime_policy is None: + num_splits = 1 if deterministic or device_capability < (9, 0) else 0 + return FlashAttentionV4SM120RuntimePolicy( + num_splits=num_splits, + decode_num_splits=num_splits, + decode_uses_static_max_seqlen_k=False, + ) + num_splits, decode_num_splits, decode_uses_static_max_seqlen_k = ( + resolve_runtime_policy( + device_capability=device_capability, + deterministic=deterministic, + ) + ) + return FlashAttentionV4SM120RuntimePolicy( + num_splits=num_splits, + decode_num_splits=decode_num_splits, + decode_uses_static_max_seqlen_k=decode_uses_static_max_seqlen_k, + ) + + +def _validate_out_contract(out: Optional[torch.Tensor]) -> None: + if out is None: + return + if out.requires_grad: + raise ValueError("out must not require gradients") + if out.stride(-1) != 1: + raise ValueError("out must have stride 1 in the last dimension") + + +@debug_kernel_api +def flash_attn_varlen_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + softcap: Optional[float] = None, + window_size: Tuple[Optional[int], Optional[int]] = (-1, -1), + learnable_sink: Optional[torch.Tensor] = None, + sinks: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + score_mod: Optional[Callable] = None, + aux_tensors: Optional[list] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + sfq: Optional[torch.Tensor] = None, + sfk: Optional[torch.Tensor] = None, + sfv: Optional[torch.Tensor] = None, + rel_bias: Optional[torch.Tensor] = None, + rel_bias_prep_cache: Optional[dict] = None, + return_softmax_lse: bool = False, + out: Optional[torch.Tensor] = None, + **_: object, +): + if _flash_attn_varlen_func is None: # pragma: no cover + raise ImportError( + "FlashAttention-4 CUTE is not available. Install flash-attn-4 with " + "its CUDA/CUTE dependencies, or run from a source tree where the " + "vendored FA4 package is importable." + ) from _flash_attn_import_error + + _validate_out_contract(out) + q, k, v, qv = [_maybe_contiguous(t) for t in (q, k, v, qv)] + if qv is None and q.shape[-1] == 256 and k.shape[-1] == 256 and v.shape[-1] == 256: + # The vendored hd256 kernel assumes dense Q/K/V strides. + q, k, v = [t.contiguous() for t in (q, k, v)] + q, qv, mla_head_padding = _pad_mla_q_heads(q, qv, v, pack_gqa) + if qv is not None and num_splits < 1: + # FA4 MLA does not implement split-KV; auto mode must use one split. + num_splits = 1 + cu_seqlens_q, cu_seqlens_k = [ + _maybe_contiguous(t) for t in (cu_seqlens_q, cu_seqlens_k) + ] + seqused_q, seqused_k = [_maybe_contiguous(t) for t in (seqused_q, seqused_k)] + page_table = _maybe_contiguous(page_table) + + if learnable_sink is None and sinks is not None: + learnable_sink = sinks + if window_size == (-1, -1): + window_size = (None, None) + + sf_kwargs = {} + if sfq is not None: + sf_kwargs["sfq"] = sfq + if sfk is not None: + sf_kwargs["sfk"] = sfk + if sfv is not None: + sf_kwargs["sfv"] = sfv + + descale_kwargs = {} + if q_descale is not None: + descale_kwargs["q_descale"] = q_descale + if k_descale is not None: + descale_kwargs["k_descale"] = k_descale + if v_descale is not None: + descale_kwargs["v_descale"] = v_descale + + rel_bias_kwargs = {} + if rel_bias is not None: + rel_bias_kwargs["rel_bias"] = rel_bias + if rel_bias_prep_cache is not None: + rel_bias_kwargs["rel_bias_prep_cache"] = rel_bias_prep_cache + + result = _flash_attn_varlen_func( + q=q, + k=k, + v=v, + qv=qv, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_q=seqused_q, + seqused_k=seqused_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + page_table=page_table, + softmax_scale=softmax_scale, + causal=causal, + softcap=softcap, + window_size=window_size, + learnable_sink=learnable_sink, + num_splits=num_splits, + pack_gqa=pack_gqa, + score_mod=score_mod, + aux_tensors=aux_tensors, + return_lse=return_softmax_lse, + out=out, + **sf_kwargs, + **descale_kwargs, + **rel_bias_kwargs, + ) + result = _unpad_mla_result(result, mla_head_padding) + + if return_softmax_lse: + return result + if isinstance(result, tuple): + return result[0] + return result + + +@debug_kernel_api +def flash_attn_with_kvcache( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k: Optional[torch.Tensor] = None, + v: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + cache_seqlens: Optional[Union[int, torch.Tensor]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size: Tuple[int, int] = (-1, -1), + attention_chunk: Optional[int] = None, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata=None, + num_splits: int = 0, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, + sinks: Optional[torch.Tensor] = None, + score_mod: Optional[Callable] = None, + aux_tensors: Optional[list] = None, + sfq: Optional[torch.Tensor] = None, + sfk: Optional[torch.Tensor] = None, + sfv: Optional[torch.Tensor] = None, + rel_bias: Optional[torch.Tensor] = None, + rel_bias_prep_cache: Optional[dict] = None, + return_softmax_lse: bool = False, + out: Optional[torch.Tensor] = None, + max_seqlen_k: Optional[int] = None, + **_: object, +): + _validate_out_contract(out) + if k is not None or v is not None: + raise NotImplementedError("FA4 does not support updating KV cache in-place.") + if rotary_cos is not None or rotary_sin is not None or rotary_seqlens is not None: + raise NotImplementedError("FA4 path does not support rotary embedding.") + if cache_batch_idx is not None or cache_leftpad is not None: + raise NotImplementedError( + "FA4 path does not support non-consecutive batch indices or left padding." + ) + if isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + + forward_arch = get_forward_arch(q.device) if get_forward_arch is not None else None + if ( + forward_arch is not None + and not return_softmax_lse + and softcap in (None, 0.0) + and all( + value is None + for value in ( + qv, + score_mod, + aux_tensors, + q_descale, + k_descale, + v_descale, + sfq, + sfk, + sfv, + rel_bias, + rel_bias_prep_cache, + ) + ) + ): + q, k_cache, v_cache = [_maybe_contiguous(t) for t in (q, k_cache, v_cache)] + cu_seqlens_q, cache_seqlens, page_table = [ + _maybe_contiguous(t) for t in (cu_seqlens_q, cache_seqlens, page_table) + ] + fast_result = try_cached_paged_decode( + arch=forward_arch, + q=q, + k=k_cache, + v=v_cache, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=None, + seqused_q=None, + seqused_k=cache_seqlens, + page_table=page_table, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + learnable_sink=sinks, + requested_num_splits=num_splits, + pack_gqa=pack_gqa, + out=out, + ) + if fast_result is not None: + return fast_result[0] + + result = flash_attn_varlen_func( + q=q, + k=k_cache, + v=v_cache, + qv=qv, + cu_seqlens_q=cu_seqlens_q, + seqused_k=cache_seqlens, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + page_table=page_table, + softmax_scale=softmax_scale, + causal=causal, + softcap=softcap if softcap != 0.0 else None, + window_size=window_size, + num_splits=num_splits, + pack_gqa=pack_gqa, + learnable_sink=sinks, + score_mod=score_mod, + aux_tensors=aux_tensors, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + sfq=sfq, + sfk=sfk, + sfv=sfv, + rel_bias=rel_bias, + rel_bias_prep_cache=rel_bias_prep_cache, + return_softmax_lse=return_softmax_lse if forward_arch is not None else True, + out=out, + ) + + if return_softmax_lse: + return result + if isinstance(result, tuple): + return result[0] + return result diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_sm120.py b/python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_sm120.py deleted file mode 100644 index 8c86e8944..000000000 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_sm120.py +++ /dev/null @@ -1,63 +0,0 @@ -# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. -# SM120 (Blackwell GeForce / DGX Spark) forward pass. -# -# SM120 uses the same SM80-era MMA instructions (mma.sync.aligned.m16n8k16) but has -# a smaller shared memory capacity (99 KB vs 163 KB on SM80). This module subclasses -# FlashAttentionForwardSm80 and overrides the SMEM capacity check accordingly. - -import cutlass -import cutlass.utils as utils_basic - -from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd import ( - FlashAttentionForwardSm80, -) - - -class FlashAttentionForwardSm120(FlashAttentionForwardSm80): - # Keep arch = 80 to use CpAsync code paths (no TMA for output). - # The compilation target is determined by the GPU at compile time, not this field. - arch = 80 - - @staticmethod - def can_implement( - dtype, - head_dim, - head_dim_v, - tile_m, - tile_n, - num_stages, - num_threads, - is_causal, - Q_in_regs=False, - ) -> bool: - """Check if the kernel can be implemented on SM120. - - Same logic as SM80 but uses SM120's shared memory capacity (99 KB). - """ - if dtype not in [cutlass.Float16, cutlass.BFloat16]: - return False - if head_dim % 8 != 0: - return False - if head_dim_v % 8 != 0: - return False - if tile_n % 16 != 0: - return False - if num_threads % 32 != 0: - return False - # Shared memory usage: Q tile + (K tile + V tile) - smem_usage_Q = tile_m * head_dim * 2 - smem_usage_K = tile_n * head_dim * num_stages * 2 - smem_usage_V = tile_n * head_dim_v * num_stages * 2 - smem_usage_QV = ( - (smem_usage_Q + smem_usage_V) - if not Q_in_regs - else max(smem_usage_Q, smem_usage_V) - ) - smem_usage = smem_usage_QV + smem_usage_K - # SM120 has 99 KB shared memory (vs 163 KB on SM80) - smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_120") - if smem_usage > smem_capacity: - return False - if (tile_m * 2) % num_threads != 0: - return False - return True diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/interface.py b/python/sglang/kernels/ops/attention/flash_attn/cute/interface.py index 29121bc0c..aefa4861f 100644 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/interface.py +++ b/python/sglang/kernels/ops/attention/flash_attn/cute/interface.py @@ -58,8 +58,10 @@ from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_sm100 import ( DescaleTensors, FlashAttentionForwardSm100, ) -from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_sm120 import ( - FlashAttentionForwardSm120, +from sglang.kernels.ops.attention.fa4_sm120.dispatch import ( + get_forward_host, + try_cached_paged_decode, + try_cached_varlen, ) from sglang.kernels.ops.attention.flash_attn.cute.shearing_bias import ShearingBias @@ -85,9 +87,8 @@ def _parse_arch_str(arch_str): def _get_device_arch(): """Cached device arch check. - Override with FLASH_ATTENTION_ARCH (e.g. 'sm_80' or '80') to select which - kernel path to use (SM80/SM90/SM100/SM120) independently of the compilation - target (CUTE_DSL_ARCH). + Override with FLASH_ATTENTION_ARCH (e.g. 'sm_80' or '80') to select the + kernel path independently of the compilation target (CUTE_DSL_ARCH). For CPU-only compilation (no GPU), set both: FLASH_ATTENTION_ARCH=sm_80 (kernel selection) @@ -100,6 +101,12 @@ def _get_device_arch(): return major * 10 + int(minor) +@lru_cache(maxsize=None) +def _get_device_num_sms(device: torch.device) -> int: + """Return the stable SM count without querying CUDA on every launch.""" + return torch.cuda.get_device_properties(device).multi_processor_count + + def _validate_head_dims( head_dim: int, head_dim_v: int, compute_capability: int, alignment: int ) -> None: @@ -342,6 +349,60 @@ def _flash_attn_fwd( aux_tensors: Some score_mods will want to read from global aux_tensors. This is how we thread them through to the inner kernel. aux_scalars: Runtime scalar captures used by score_mod or mask_mod. """ + fake_mode = is_fake_mode() + arch = _get_device_arch() if _arch is None else _arch + arch_forward_host = get_forward_host(arch) + requested_num_splits = num_splits + if ( + not fake_mode + and arch_forward_host is not None + and not return_lse + and lse is None + and softcap in (None, 0.0) + and all( + value is None + for value in ( + qv, + gather_kv_indices, + score_mod, + mask_mod, + block_sparse_tensors, + aux_tensors, + aux_scalars, + q_descale, + k_descale, + v_descale, + rel_bias, + sfq, + sfk, + sfv, + ) + ) + ): + fast_result = try_cached_paged_decode( + arch=arch, + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_q=seqused_q, + seqused_k=seqused_k, + page_table=page_table, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + window_size_left=window_size_left, + window_size_right=window_size_right, + learnable_sink=learnable_sink, + requested_num_splits=requested_num_splits, + pack_gqa=pack_gqa, + out=out, + ) + if fast_result is not None: + return fast_result + aux_scalars = tuple(aux_scalars) if aux_scalars else None q, k, v, qv = [maybe_contiguous(t) for t in (q, k, v, qv)] assert q is not None or qv is not None @@ -478,7 +539,7 @@ def _flash_attn_fwd( assert learnable_sink.shape == (num_head,) assert learnable_sink.dtype == torch.bfloat16, "learnable_sink must be bfloat16" - if not is_fake_mode(): + if not fake_mode: assert all( t is None or t.is_cuda for t in ( @@ -497,17 +558,13 @@ def _flash_attn_fwd( learnable_sink, ) ), "inputs must be on CUDA device" - arch = _get_device_arch() if _arch is None else _arch - assert arch // 10 in [ - 8, - 9, - 10, - 11, - 12, - ], "Unsupported compute capability. Supported: 8.x, 9.x, 10.x, 11.x, 12.x" + assert arch // 10 in [8, 9, 10, 11] or arch_forward_host is not None, ( + "Unsupported compute capability. Supported: 8.x, 9.x, 10.x, 11.x, " + "and architectures registered through the forward-host bridge" + ) assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv" alignment = 16 // v.element_size() - if arch // 10 not in [8, 12]: + if arch // 10 != 8 and arch_forward_host is None: _validate_head_dims(head_dim, head_dim_v, arch // 10, alignment) if softmax_scale is None: softmax_scale = ( @@ -565,6 +622,10 @@ def _flash_attn_fwd( device=device, ) else: + if out.requires_grad: + raise ValueError("out must not require gradients") + if out.stride(-1) != 1: + raise ValueError("out must have stride 1 in the last dimension") _validate_tensor( out, "out", @@ -619,21 +680,41 @@ def _flash_attn_fwd( current_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - # SM80/SM120: uses SM80 MMA, 128 threads (4 warps) - if arch // 10 in [8, 12]: + if arch // 10 == 8: num_threads = 128 + num_SMs = 132 if fake_mode else _get_device_num_sms(device) fwd_cfg = FwdConfig(128, 128, True, True) # default - if tile_mn is None: - if arch // 10 == 12: - # SM120 tile sizes tuned for 99 KB SMEM capacity: - # D<=64: 128x128 → 48 KB (good occupancy) - # D>64: 128x64 → 64 KB (128x128 would use 96 KB, hurting occupancy) - if head_dim <= 64: - fwd_cfg = FwdConfig(128, 128, True, True) - else: - fwd_cfg = FwdConfig(128, 64, True, True) - elif arch // 10 == 8: + arch_forward_config = None + if arch_forward_host is not None: + arch_forward_config = arch_forward_host.select_config( + head_dim=head_dim, + head_dim_v=head_dim_v, + tile_mn=tile_mn, + has_bias=rel_bias is not None, + total_q_rows=total_q * num_head, + num_sms=None if fake_mode else num_SMs, + num_batch=batch_size, + seqlen_q=(max_seqlen_q if max_seqlen_q is not None else seqlen_q), + seqlen_k=(max_seqlen_k if max_seqlen_k is not None else seqlen_k), + num_head_kv=num_head_kv, + qhead_per_kvhead=qhead_per_kvhead, + is_causal=causal, + is_local=local, + window_size_left=window_size_left, + window_size_right=window_size_right, + pack_gqa=pack_gqa, + paged_kv=page_table is not None, + ) + fwd_cfg = FwdConfig( + arch_forward_config.tile_m, + arch_forward_config.tile_n, + True, + True, + ) + num_threads = arch_forward_config.num_threads + elif tile_mn is None: + if arch // 10 == 8: fwd_cfg = FwdConfig(128, 64, True, True) # SM80, should tune elif arch // 10 == 9: sparse_q = get_sparse_q_block_size(block_sparse_tensors, seqlen_q) @@ -683,13 +764,58 @@ def _flash_attn_fwd( ) // m_block_size_effective total_mblocks = batch_size * num_head_kv * num_m_blocks num_n_blocks = (seqlen_k_loaded + tile_n - 1) // tile_n - num_SMs = ( - 132 - if is_fake_mode() - else torch.cuda.get_device_properties(device).multi_processor_count - ) - if num_splits < 1: - num_splits = num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, 128) + arch_forward_plan = None + if arch_forward_host is not None: + arch_forward_plan = arch_forward_host.resolve_plan( + requested_num_splits=num_splits, + generic_num_n_blocks=num_n_blocks, + head_dim=head_dim, + head_dim_v=head_dim_v, + batch_size=batch_size, + num_head_kv=num_head_kv, + paged_kv=page_table is not None, + page_size=page_size, + k=k, + v=v, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + pack_gqa=pack_gqa, + element_size=q.element_size(), + packed_q_rows=seqlen_q_packgqa, + tile_m=tile_m, + tile_n=tile_n, + num_m_blocks=num_m_blocks, + total_mblocks=total_mblocks, + num_sms=num_SMs, + total_q=total_q, + has_cu_seqlens_q=cu_seqlens_q is not None, + has_seqused_q=seqused_q is not None, + has_seqused_k=seqused_k is not None, + is_causal=causal, + is_local=local, + window_size_left=window_size_left, + window_size_right=window_size_right, + has_score_or_mask_mod=( + softcap is not None + or score_mod is not None + or mask_mod is not None + or rel_bias is not None + ), + is_stream_capturing=( + not fake_mode and torch.cuda.is_current_stream_capturing() + ), + device=device, + fake_mode=fake_mode, + generic_heuristic=num_splits_heuristic, + ) + num_splits = arch_forward_plan.num_splits + elif num_splits < 1: + num_splits = num_splits_heuristic( + total_mblocks, + num_SMs, + num_n_blocks, + 128, + ) # SplitKV uses float32 partial output, which doubles the O buffer size # in shared memory, causing OOM for diff-headdim (192, 128) @@ -859,8 +985,8 @@ def _flash_attn_fwd( disable_sparse_kv_bitmask = None p = row_max = None - # rel_bias -> sheared bias (Inkling relative attention). Produces `bias`, the column-aligned - # bias the SM100 kernel adds to pre-softmax scores via its dedicated TMA pipeline. + # Inkling relative attention. Shear the relative rows into the column-aligned + # tiles consumed by the architecture-specific attention mainloop. rel_extent = 0 rel_extent_padded = 0 bias = None @@ -868,16 +994,17 @@ def _flash_attn_fwd( cu_total_m_blocks_bias = None blocks_to_batch_idx = None if rel_bias is not None: - assert arch // 10 in [ - 9, - 10, - 11, - ], "rel_bias (sheared bias) is only supported on SM9x/10x" + assert arch // 10 in [9, 10, 11] or arch_forward_host is not None, ( + "rel_bias requires SM9x/10x/11x or an architecture-owned " + "forward implementation" + ) qhead_per_kvhead_packgqa = qhead_per_kvhead if pack_gqa else 1 rel_extent = rel_bias.shape[-1] rel_extent_padded = rel_extent + 256 assert rel_extent % 128 == 0 - assert tile_m == 128 and tile_n == 128 + assert tile_n == 128 + if arch_forward_host is None: + assert tile_m == 128 assert ( causal or window_size_left is None @@ -975,7 +1102,7 @@ def _flash_attn_fwd( current_stream, options="--enable-tvm-ffi", ) - if not is_fake_mode(): + if not fake_mode: _flash_attn_fwd.compile_cache_prepare_shear_bias[ compile_key_prepare ]( @@ -1003,6 +1130,7 @@ def _flash_attn_fwd( qhead_per_kvhead, rows_per_cta, group_tile_bias, + tile_m, max_m_blocks_leq_one, cu_total_m_blocks_bias is not None, blocks_to_batch_idx is not None, @@ -1039,6 +1167,7 @@ def _flash_attn_fwd( qhead_per_kvhead=qhead_per_kvhead, rows_per_cta=rows_per_cta, tile_m=group_tile_bias, + attention_tile_m=tile_m, max_m_blocks_leq_one=max_m_blocks_leq_one, use_pdl=use_pdl, ), @@ -1057,7 +1186,7 @@ def _flash_attn_fwd( current_stream, options="--enable-tvm-ffi", ) - if not is_fake_mode(): + if not fake_mode: _flash_attn_fwd.compile_cache_shear_bias[shear_compile_key]( rel_bias, bias, @@ -1111,6 +1240,8 @@ def _flash_attn_fwd( is_split_kv, pack_gqa, arch, + arch_forward_config.compile_key if arch_forward_config is not None else None, + arch_forward_plan.compile_key if arch_forward_plan is not None else None, page_size not in [None, tile_n], # paged KV non-TMA use_2cta_instrs, q_subtile_factor, @@ -1364,31 +1495,37 @@ def _flash_attn_fwd( ) ), ) - elif arch // 10 == 12: - # SM120 (Blackwell GeForce / DGX Spark): uses SM80 MMA with SM120 SMEM capacity - assert not use_block_sparsity, "Block sparsity not supported on SM 12.0" - assert page_table is None, "Paged KV not supported on SM 12.0 in this PR" - assert not is_split_kv, "SplitKV not supported on SM 12.0 in this PR" - fa_fwd = FlashAttentionForwardSm120( - dtype, - head_dim, - head_dim_v, - qhead_per_kvhead, + elif arch_forward_host is not None: + assert not use_block_sparsity, ( + "Block sparsity is not supported by the architecture-owned " + "forward implementation" + ) + assert arch_forward_config is not None + assert arch_forward_plan is not None + fa_fwd = arch_forward_host.make_kernel( + dtype=dtype, + head_dim=head_dim, + head_dim_v=head_dim_v, + qhead_per_kvhead=qhead_per_kvhead, is_causal=causal, is_local=local, pack_gqa=pack_gqa, - tile_m=tile_m, - tile_n=tile_n, - num_stages=1, - num_threads=num_threads, - Q_in_regs=False, + config=arch_forward_config, + paged_kv=page_table is not None, score_mod=score_mod, mask_mod=mask_mod, has_aux_tensors=aux_tensors is not None, + is_split_kv=is_split_kv, + has_bias=bias is not None, + bias_block_size=tile_bias, + rel_extent_padded=rel_extent_padded, + plan=arch_forward_plan, ) else: raise ValueError( - f"Unsupported compute capability: {arch}. Supported: 8.x, 9.x, 10.x, 11.x, 12.x" + f"Unsupported compute capability: {arch}. Supported: 8.x, 9.x, " + "10.x, 11.x, and architectures registered through the " + "forward-host bridge" ) # TODO: check @can_implement if qv is not None: @@ -1440,7 +1577,7 @@ def _flash_attn_fwd( AuxData(cute_aux_tensors, aux_scalars), ] ) - if arch // 10 in [9, 10, 11]: + if arch // 10 in [9, 10, 11] or arch_forward_host is not None: compile_args.append(bias_tensor) # mBias if arch // 10 in [10, 11]: if not use_dedicated_hd256_kernel: @@ -1453,12 +1590,17 @@ def _flash_attn_fwd( v_sf_vec_size, ] ) + if arch_forward_host is not None: + assert arch_forward_plan is not None + compile_args.extend( + arch_forward_host.compile_arguments(arch_forward_plan) + ) compile_args.append(current_stream) _flash_attn_fwd.compile_cache[compile_key] = cute.compile( *compile_args, options="--enable-tvm-ffi" ) - if not is_fake_mode(): + if not fake_mode: q_call, k_call, v_call, qv_call = [ t.detach() if t is not None else None for t in (q, k, v, qv) ] @@ -1543,7 +1685,7 @@ def _flash_attn_fwd( AuxData(aux_tensors, aux_scalars), ] ) - if arch // 10 in [9, 10, 11]: + if arch // 10 in [9, 10, 11] or arch_forward_host is not None: call_args.append(bias) # mBias if arch // 10 in [10, 11]: if not use_dedicated_hd256_kernel: @@ -1556,7 +1698,52 @@ def _flash_attn_fwd( sfv_call, # mSFV (None unless v_blockscaled) ] ) - _flash_attn_fwd.compile_cache[compile_key](*call_args) + if arch_forward_host is not None: + assert arch_forward_plan is not None + call_args.extend(arch_forward_host.runtime_arguments(arch_forward_plan)) + compiled_fwd = _flash_attn_fwd.compile_cache[compile_key] + if ( + arch_forward_host is not None + and cu_seqlens_q is not None + and cu_seqlens_k is not None + and seqused_q is None + and seqused_k is None + and not is_split_kv + and not requires_grad + and (learnable_sink is None or not learnable_sink.requires_grad) + and lse is None + and softcap is None + and score_mod is None + and mask_mod is None + and block_sparse_tensors is None + and aux_tensors is None + and aux_scalars is None + and q_descale is None + and k_descale is None + and v_descale is None + and rel_bias is None + and sfq is None + and sfk is None + and sfv is None + ): + arch_forward_host.register_varlen( + arch=arch, + compiled_fn=compiled_fwd, + compile_key=compile_key, + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=causal, + window_size_left=window_size_left, + window_size_right=window_size_right, + learnable_sink=learnable_sink, + pack_gqa=pack_gqa, + ) + compiled_fwd(*call_args) if is_split_kv: _flash_attn_fwd_combine( out_partial, @@ -1566,6 +1753,72 @@ def _flash_attn_fwd( cu_seqlens_q, seqused_q, ) + if ( + not fake_mode + and not torch.cuda.is_current_stream_capturing() + and arch_forward_host is not None + and q is not None + and k is not None + and qv is None + and cu_seqlens_q is not None + and cu_seqlens_k is None + and seqused_q is None + and seqused_k is not None + and page_table is not None + and not requires_grad + and (learnable_sink is None or not learnable_sink.requires_grad) + and lse is None + and softcap is None + and score_mod is None + and mask_mod is None + and block_sparse_tensors is None + and aux_tensors is None + and aux_scalars is None + and q_descale is None + and k_descale is None + and v_descale is None + and rel_bias is None + and sfq is None + and sfk is None + and sfv is None + ): + compiled_combine = None + if is_split_kv: + lse_partial_transposed = lse_partial.transpose(-1, -2) + combine_key = _fwd_combine_compile_key( + out_partial, + out, + None, + cu_seqlens_q, + None, + None, + ) + compiled_combine = _flash_attn_fwd_combine.compile_cache[combine_key] + arch_forward_host.register_paged_decode( + arch=arch, + compiled_fn=compiled_fwd, + compile_key=compile_key, + compiled_combine=compiled_combine, + actual_num_splits=num_splits, + tile_m=tile_m, + tile_n=tile_n, + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + seqused_k=seqused_k, + page_table=page_table, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=causal, + window_size_left=window_size_left, + window_size_right=window_size_right, + learnable_sink=learnable_sink, + pack_gqa=pack_gqa, + requested_num_splits=requested_num_splits, + out_partial=out_partial if is_split_kv else None, + lse_partial=lse_partial if is_split_kv else None, + ) return out, lse @@ -1689,6 +1942,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function): v_sf_vec_size: Optional[int] = None, rel_bias_prep_cache: Optional[dict] = None, return_lse: bool = False, + out: Optional[torch.Tensor] = None, ): aux_scalars = tuple(aux_scalars) if aux_scalars else None shared_kv = k is v @@ -1736,32 +1990,34 @@ class FlashAttnVarlenFunc(torch.autograd.Function): qk_sf_vec_size=qk_sf_vec_size, v_sf_vec_size=v_sf_vec_size, rel_bias_prep_cache=rel_bias_prep_cache, + out=out, ) - ctx.save_for_backward( - q, - k, - v, - out, - lse, - cu_seqlens_q, - cu_seqlens_k, - seqused_q, - seqused_k, - *(aux_tensors or ()), - ) - ctx.softmax_scale = softmax_scale - ctx.causal = causal - ctx.window_size = window_size - ctx.softcap = softcap - ctx.deterministic = deterministic - ctx.max_seqlen_q = max_seqlen_q - ctx.max_seqlen_k = max_seqlen_k - ctx.return_lse = return_lse - ctx.score_mod = score_mod - ctx.score_mod_bwd = score_mod_bwd - ctx.mask_mod = mask_mod - ctx.aux_scalars = aux_scalars - ctx.set_materialize_grads(False) + if ctx is not None: + ctx.save_for_backward( + q, + k, + v, + out, + lse, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + *(aux_tensors or ()), + ) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.return_lse = return_lse + ctx.score_mod = score_mod + ctx.score_mod_bwd = score_mod_bwd + ctx.mask_mod = mask_mod + ctx.aux_scalars = aux_scalars + ctx.set_materialize_grads(False) return out, lse @@ -1852,6 +2108,7 @@ def flash_attn_varlen_func( v_sf_vec_size: Optional[int] = None, rel_bias_prep_cache: Optional[dict] = None, return_lse: bool = False, + out: Optional[torch.Tensor] = None, ): """ Tensor arguments: @@ -1891,7 +2148,63 @@ def flash_attn_varlen_func( qk_sf_vec_size = 32 if v_sf_vec_size is None and sfv is not None and sfv.dtype == torch.float8_e8m0fnu: v_sf_vec_size = 32 - return FlashAttnVarlenFunc.apply( + if out is not None and out.requires_grad: + raise ValueError("out must not require gradients") + if out is not None and out.stride(-1) != 1: + raise ValueError("out must have stride 1 in the last dimension") + runtime_arch = _get_device_arch() + forward_host = None if is_fake_mode() else get_forward_host(runtime_arch) + if ( + forward_host is not None + and q is not None + and k is not None + and cu_seqlens_q is not None + and cu_seqlens_k is not None + and num_splits == 1 + and softcap in (None, 0.0) + and not return_lse + and all( + value is None + for value in ( + qv, + seqused_q, + seqused_k, + gather_kv_indices, + page_table, + score_mod, + mask_mod, + block_sparse_tensors, + aux_tensors, + aux_scalars, + q_descale, + k_descale, + v_descale, + rel_bias, + sfq, + sfk, + sfv, + ) + ) + ): + fast_result = try_cached_varlen( + arch=runtime_arch, + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + learnable_sink=learnable_sink, + pack_gqa=pack_gqa, + out=out, + ) + if fast_result is not None: + return fast_result + autograd_args = ( q, k, v, @@ -1930,7 +2243,34 @@ def flash_attn_varlen_func( v_sf_vec_size, rel_bias_prep_cache, return_lse, + out, ) + needs_autograd = False + if forward_host is not None or out is not None: + differentiable_tensors = ( + q, + k, + v, + qv, + learnable_sink, + q_descale, + k_descale, + v_descale, + rel_bias, + sfq, + sfk, + sfv, + *(aux_tensors or ()), + ) + needs_autograd = torch.is_grad_enabled() and any( + tensor is not None and tensor.requires_grad + for tensor in differentiable_tensors + ) + if needs_autograd and out is not None: + raise ValueError("out is only supported for forward-only inference") + if not needs_autograd and forward_host is not None: + return FlashAttnVarlenFunc.forward(None, *autograd_args) + return FlashAttnVarlenFunc.apply(*autograd_args) def _compile_fwd_combine( @@ -2039,6 +2379,35 @@ def _compile_fwd_combine( ) +def _fwd_combine_compile_key( + out_partial: torch.Tensor, + out: torch.Tensor, + lse: Optional[torch.Tensor], + cu_seqlens: Optional[torch.Tensor], + seqused: Optional[torch.Tensor], + varlen_batch_idx: Optional[torch.Tensor], +) -> tuple: + head_dim = out_partial.shape[-1] + num_splits = out_partial.shape[0] + k_block_size = 64 if head_dim <= 64 else 128 + tile_m = 8 if k_block_size % 128 == 0 else (16 if k_block_size % 64 == 0 else 32) + log_max_splits = max(math.ceil(math.log2(num_splits)), 4) + if tile_m == 8: + log_max_splits = max(log_max_splits, 5) + return ( + torch2cute_dtype_map[out.dtype], + torch2cute_dtype_map[out_partial.dtype], + head_dim, + tile_m, + k_block_size, + log_max_splits, + cu_seqlens is not None, + seqused is not None, + lse is not None, + varlen_batch_idx is not None, + ) + + def _flash_attn_fwd_combine( out_partial: torch.Tensor, lse_partial: torch.Tensor, @@ -2092,38 +2461,18 @@ def _flash_attn_fwd_combine( if not is_fake_mode(): assert t.is_cuda, f"{name} must be on CUDA device" assert t.is_contiguous(), f"{name} must be contiguous" - head_dim = out_partial.shape[-1] num_splits = out_partial.shape[0] assert num_splits <= 256 - # If hdim is 96 or 192, it's faster to round them to 128 or 256 respectively - # so that kBlockM is smaller and we have more parallelism. - k_block_size = 64 if head_dim <= 64 else 128 - # We want kBlockM to be as small as possible to maximize parallelism. - # E.g., if hdim is 64, we want kBlockM to be 16 so that we can use 256 threads, each reading 4 elements (floats). - tile_m = 8 if k_block_size % 128 == 0 else (16 if k_block_size % 64 == 0 else 32) - log_max_splits = max(math.ceil(math.log2(num_splits)), 4) - if tile_m == 8: - # If kBlockM == 8 then the minimum number of splits is 32. - # TODO: we can deal w this by using 128 threads instead - log_max_splits = max(log_max_splits, 5) - - # Create combine kernel configuration - dtype = torch2cute_dtype_map[out.dtype] - dtype_partial = torch2cute_dtype_map[out_partial.dtype] # Device architecture is invariant for the lifetime of this server/JIT # cache, so PDL does not belong in the compile key. use_pdl = is_arch_support_pdl() - compile_key = ( - dtype, - dtype_partial, - head_dim, - tile_m, - k_block_size, - log_max_splits, - cu_seqlens is not None, - seqused is not None, - lse is not None, - varlen_batch_idx is not None, + compile_key = _fwd_combine_compile_key( + out_partial, + out, + lse, + cu_seqlens, + seqused, + varlen_batch_idx, ) if compile_key not in _flash_attn_fwd_combine.compile_cache: _flash_attn_fwd_combine.compile_cache[compile_key] = _compile_fwd_combine( diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/mask.py b/python/sglang/kernels/ops/attention/flash_attn/cute/mask.py index 8b805d0a4..29fcef7b4 100644 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/mask.py +++ b/python/sglang/kernels/ops/attention/flash_attn/cute/mask.py @@ -168,6 +168,9 @@ class AttentionMask: 1 # only pass in if we're doing PackGQA ) swap_AB: cutlass.Constexpr[bool] = False + # R2P assumes the canonical row ownership of the QK accumulator. Kernels + # with a different accumulator mapping must use the direct predicate path. + enable_r2p_optimization: cutlass.Constexpr[bool] = True @property def seqlen_q(self) -> Int32: @@ -330,7 +333,7 @@ class AttentionMask: ) if const_expr(mask_causal): r2p = const_expr( - not self.swap_AB + not self.swap_AB and self.enable_r2p_optimization ) # R2P trick, see apply_mask_sm100 for r in cutlass.range( cute.size(tScS_mn.shape[0]), unroll_full=True @@ -375,7 +378,9 @@ class AttentionMask: if const_expr(self.window_size_left is not None) else None ) - r2p_local = const_expr(not self.swap_AB) + r2p_local = const_expr( + not self.swap_AB and self.enable_r2p_optimization + ) for r in cutlass.range( cute.size(tScS_mn.shape[0]), unroll_full=True ): diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/shearing_bias.py b/python/sglang/kernels/ops/attention/flash_attn/cute/shearing_bias.py index 1169fec3c..7a063d202 100644 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/shearing_bias.py +++ b/python/sglang/kernels/ops/attention/flash_attn/cute/shearing_bias.py @@ -35,6 +35,7 @@ class ShearingBias: qhead_per_kvhead: cutlass.Constexpr[int] = 1, rows_per_cta: int = 4, tile_m: int = 128, + attention_tile_m: int = 128, max_m_blocks_leq_one: bool = False, use_pdl: bool = False, clamp_subtiles: bool = True, @@ -67,6 +68,12 @@ class ShearingBias: # only used with block packed scheduling self.tile_m = tile_m + # The output columns are aligned to the rightmost N tile visible to the + # attention CTA. Keep this distinct from ``tile_m`` above: the shear + # scheduler may group rows in larger blocks than the attention kernel + # consumes per CTA. + assert attention_tile_m % self.rows_per_cta == 0 + self.attention_tile_m = attention_tile_m # Shrink the subtile grid dim to the rows a block can actually hold # (decode blocks hold qhead_per_kvhead*seqlen_q rows, not tile_m). self.clamp_subtiles = clamp_subtiles @@ -335,7 +342,7 @@ class ShearingBias: ) block_info = BlockInfo( - 128, + self.attention_tile_m, 128, self.is_causal, self.is_local, @@ -386,7 +393,7 @@ class ShearingBias: # Convention: inclusive min, exclusive max m_idx = m_block * self.rows_per_cta + warp_idx - attn_m_block = m_idx // 128 + attn_m_block = m_idx // self.attention_tile_m _, attn_n_block_max = block_info.get_n_block_min_max( seqlen_info, diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index ee2337f7e..b9dcd6746 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -253,6 +253,7 @@ class FlashAttentionBackend(AttentionBackend): # Select version self.fa_impl_ver = fa_impl_ver + device_capability = get_device_capability() if self.fa_impl_ver == 3: from sgl_kernel.flash_attn import ( flash_attn_varlen_func, @@ -261,11 +262,25 @@ class FlashAttentionBackend(AttentionBackend): ) self._get_scheduler_metadata = get_scheduler_metadata + self._get_fa_runtime_policy = None elif self.fa_impl_ver == 4: - from sglang.kernels.ops.attention.flash_attention_v4 import ( - flash_attn_varlen_func, - flash_attn_with_kvcache, - ) + if device_capability[0] == 12: + from sglang.kernels.ops.attention.flash_attention_v4_sm120 import ( + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_flash_attention_v4_sm120_runtime_policy, + ) + + self._get_fa_runtime_policy = ( + get_flash_attention_v4_sm120_runtime_policy + ) + else: + from sglang.kernels.ops.attention.flash_attention_v4 import ( + flash_attn_varlen_func, + flash_attn_with_kvcache, + ) + + self._get_fa_runtime_policy = None self._get_scheduler_metadata = None if model_runner.server_args.enable_deterministic_inference: @@ -295,15 +310,22 @@ class FlashAttentionBackend(AttentionBackend): ) self.has_softcap = _softcapping is not None and _softcapping > 0.0 - # If num_splits == 0, we use a heuristic to automatically determine the number of splits. - # We set nums splits to 1 if deterministic inference is enabled. - # See https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/ for more details. - fa4_no_splitkv = self.fa_impl_ver == 4 and get_device_capability() < (9, 0) - self.num_splits = ( - 1 - if model_runner.server_args.enable_deterministic_inference or fa4_no_splitkv - else 0 - ) + # num_splits == 0 delegates SplitKV sizing to the selected FA runtime. + deterministic = model_runner.server_args.enable_deterministic_inference + if self._get_fa_runtime_policy is None: + self.num_splits = 1 if deterministic else 0 + self.decode_num_splits = self.num_splits + self._decode_uses_static_max_seqlen_k = False + else: + runtime_policy = self._get_fa_runtime_policy( + device_capability=device_capability, + deterministic=deterministic, + ) + self.num_splits = runtime_policy.num_splits + self.decode_num_splits = runtime_policy.decode_num_splits + self._decode_uses_static_max_seqlen_k = ( + runtime_policy.decode_uses_static_max_seqlen_k + ) # Set (never getattr'd) so forward_extend can identity-check "is this the # full-CG prefill metadata?" to disable the pointer-keyed shear-bias # block-schedule cache (see forward_extend rel_bias handling). @@ -506,6 +528,12 @@ class FlashAttentionBackend(AttentionBackend): # Local attention and scheduler metadata require capture-time slice sizing. # Both depend on data already filled by replay above. metadata = self.decode_cuda_graph_metadata[bs] + if self._decode_uses_static_max_seqlen_k: + # FA4 bakes its N-tile grid and SplitKV specialization into + # the graph. Capture against the full replay bound, not the + # padded seq-len fill value (1), or a later long-context + # replay would leave K/V tiles uncovered. + metadata.max_seq_len_k = self.max_context_len self._maybe_update_local_attn_metadata_for_capture(metadata, bs) if self._sched_meta_buf is not None: sched = self._compute_scheduler_metadata( @@ -1860,6 +1888,8 @@ class FlashAttentionBackend(AttentionBackend): if layer.is_cross_attention: # Always use non-chunked logic for cross-attention + if self._decode_uses_static_max_seqlen_k: + kwargs["max_seqlen_k"] = metadata.encoder_max_seq_len_k o = flash_attn_with_kvcache( q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), k_cache=key_cache, @@ -1879,6 +1909,8 @@ class FlashAttentionBackend(AttentionBackend): ) elif use_local_attn: # Use chunked (local) attention batching for self-attention + if self._decode_uses_static_max_seqlen_k: + kwargs["max_seqlen_k"] = local_attn_metadata.local_max_seq_len o = flash_attn_with_kvcache( q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), k_cache=key_cache, @@ -1932,6 +1964,8 @@ class FlashAttentionBackend(AttentionBackend): and not pa_swa_active ): sched_meta = metadata.scheduler_metadata + if self._decode_uses_static_max_seqlen_k: + kwargs["max_seqlen_k"] = metadata.max_seq_len_k result = flash_attn_with_kvcache( q=q_reshaped, k_cache=key_cache, @@ -1945,7 +1979,13 @@ class FlashAttentionBackend(AttentionBackend): window_size=window_size, softcap=layer.logit_cap, return_softmax_lse=use_cascade_attn, - num_splits=self.num_splits, + num_splits=( + self.decode_num_splits + if not is_swa_layer + and not use_cascade_attn + and not pa_swa_active + else self.num_splits + ), out=_fa_out, ver=self.fa_impl_ver, scheduler_metadata=sched_meta, @@ -1953,6 +1993,10 @@ class FlashAttentionBackend(AttentionBackend): ) if use_cascade_attn: o, softmax_lse, *rest = result + if self._decode_uses_static_max_seqlen_k: + kwargs["max_seqlen_k"] = ( + self.forward_metadata_spec_decode_expand.max_seq_len_k + ) o_expand, softmax_lse_expand, *rest_expand = ( flash_attn_with_kvcache( q=q_reshaped, diff --git a/test/registered/kernels/ops/attention/test_flash_attention_4_sm120.py b/test/registered/kernels/ops/attention/test_flash_attention_4_sm120.py new file mode 100644 index 000000000..235856657 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_flash_attention_4_sm120.py @@ -0,0 +1,2160 @@ +"""Focused SM120 FlashAttention-4 regression tests.""" + +import math + +import cutlass.cute as cute +import pytest +import torch + +from sglang.kernels.ops.attention.fa4_sm120.policy import ( + low_hd_paged_decode_tile_m, + visible_decode_seqlen_k, +) +from sglang.kernels.ops.attention.fa4_sm120.runtime import ( + Sm120ForwardHost, + sm120_forward_host, +) +from sglang.kernels.ops.attention.flash_attention_v4_sm120 import ( + flash_attn_varlen_func, + flash_attn_with_kvcache, + get_flash_attention_v4_sm120_runtime_policy, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=240, + stage="base-b", + runner_config="1-gpu-small", +) + +if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 12): + pytest.skip( + "SM12x FlashAttention-4 test requires CUDA SM 12.x.", + allow_module_level=True, + ) + + +def test_sm120_runtime_policy_delegates_decode_shapes(): + policy = get_flash_attention_v4_sm120_runtime_policy( + device_capability=(12, 0), + deterministic=False, + ) + assert policy.num_splits == 1 + assert policy.decode_num_splits == 0 + assert policy.decode_uses_static_max_seqlen_k + + deterministic_policy = get_flash_attention_v4_sm120_runtime_policy( + device_capability=(12, 0), + deterministic=True, + ) + assert deterministic_policy.num_splits == 1 + assert deterministic_policy.decode_num_splits == 1 + assert deterministic_policy.decode_uses_static_max_seqlen_k + + future_policy = get_flash_attention_v4_sm120_runtime_policy( + device_capability=(13, 0), + deterministic=False, + ) + assert future_policy.num_splits == 0 + assert future_policy.decode_num_splits == 0 + assert not future_policy.decode_uses_static_max_seqlen_k + + +def test_sm120_preallocated_output_contract_is_validated_before_launch(): + q = torch.empty((1, 1, 1, 32), device="cuda", dtype=torch.bfloat16) + k = torch.empty((1, 1, 1, 32), device="cuda", dtype=torch.bfloat16) + v = torch.empty_like(k) + + with pytest.raises(ValueError, match="must not require gradients"): + flash_attn_with_kvcache( + q, + k, + v, + out=torch.empty_like(q, requires_grad=True), + ) + + strided_out = torch.empty( + (1, 1, 1, 64), + device="cuda", + dtype=torch.bfloat16, + )[..., ::2] + with pytest.raises(ValueError, match="must have stride 1"): + flash_attn_with_kvcache(q, k, v, out=strided_out) + + +@pytest.mark.parametrize( + ( + "head_dim", + "head_dim_v", + "visible_seqlen_k", + "qhead_per_kvhead", + "expected_tile_m", + ), + [ + pytest.param(64, 64, 256, 1, None, id="minimum-exclusive"), + pytest.param(64, 64, 512, 8, None, id="hd64-short-mqa-fallback"), + pytest.param(64, 64, 513, 8, 16, id="hd64-long-mqa"), + pytest.param(128, 128, 512, 8, 32, id="hd128-short-mqa"), + pytest.param(128, 128, 2048, 1, 16, id="hd128-mha"), + pytest.param(64, 128, 2048, 1, None, id="asymmetric-fallback"), + pytest.param(96, 96, 2048, 1, None, id="unqualified-fallback"), + ], +) +def test_sm120_low_hd_decode_tile_qualification( + head_dim, + head_dim_v, + visible_seqlen_k, + qhead_per_kvhead, + expected_tile_m, +): + assert ( + low_hd_paged_decode_tile_m( + head_dim=head_dim, + head_dim_v=head_dim_v, + paged_kv=True, + seqlen_q=1, + visible_seqlen_k=visible_seqlen_k, + qhead_per_kvhead=qhead_per_kvhead, + ) + == expected_tile_m + ) + + +def test_sm120_decode_visible_k_uses_exact_local_window(): + assert ( + visible_decode_seqlen_k( + 8192, + is_local=False, + window_size_left=256, + window_size_right=0, + ) + == 8192 + ) + assert ( + visible_decode_seqlen_k( + 8192, + is_local=True, + window_size_left=256, + window_size_right=0, + ) + == 257 + ) + assert ( + visible_decode_seqlen_k( + 128, + is_local=True, + window_size_left=256, + window_size_right=0, + ) + == 128 + ) + + +@cute.jit +def _coordinate_score_mod( + score, + batch_idx, + head_idx, + q_idx, + kv_idx, + seqlen_info, + aux_tensors, +): + """Exercise logical batch/head/token indices, including packed GQA.""" + return ( + score + + batch_idx * 0.00390625 + + head_idx * 0.0078125 + + q_idx * 0.015625 + - kv_idx * 0.0078125 + ) + + +@cute.jit +def _aux_score_mod( + score, + batch_idx, + head_idx, + q_idx, + kv_idx, + seqlen_info, + aux_tensors, +): + """Exercise runtime aux-tensor threading without shape-specific indexing.""" + return score + aux_tensors[0][0] + + +def _attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + batch_idx: int = 0, + causal: bool = False, + window_size: tuple[int | None, int | None] = (None, None), + softmax_scale: float | None = None, + softcap: float = 0.0, + sinks: torch.Tensor | None = None, + score_mod: str | None = None, + aux_score: float = 0.0, + rel_bias: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return FP32 attention output and head-major LSE for one sequence.""" + q_len, num_q_heads, head_dim = q.shape + k_len, num_kv_heads, _ = k.shape + assert num_q_heads % num_kv_heads == 0 + scale = head_dim**-0.5 if softmax_scale is None else softmax_scale + k = k.float().repeat_interleave(num_q_heads // num_kv_heads, dim=1) + v = v.float().repeat_interleave(num_q_heads // num_kv_heads, dim=1) + scores = torch.einsum("qhd,khd->hqk", q.float(), k) * scale + + q_idx = torch.arange(q_len, device=q.device) + kv_idx = torch.arange(k_len, device=q.device) + if score_mod == "coordinate": + scores += ( + batch_idx * 0.00390625 + + torch.arange(num_q_heads, device=q.device)[:, None, None] * 0.0078125 + + q_idx[None, :, None] * 0.015625 + - kv_idx[None, None, :] * 0.0078125 + ) + elif score_mod == "aux": + scores += aux_score + elif score_mod is not None: + raise ValueError(f"unknown score_mod reference: {score_mod}") + + relative_position = q_idx[:, None] + k_len - q_len - kv_idx[None, :] + if rel_bias is not None: + bias_index = relative_position.clamp(0, rel_bias.shape[-1] - 1) + bias = ( + rel_bias.float() + .permute(1, 0, 2) + .gather( + 2, + bias_index[None].expand(num_q_heads, -1, -1), + ) + ) + bias.masked_fill_( + ~((relative_position >= 0) & (relative_position < rel_bias.shape[-1]))[ + None + ], + 0.0, + ) + scores += bias + if softcap > 0: + scores = torch.tanh(scores / softcap) * softcap + + visible = torch.ones_like(relative_position, dtype=torch.bool) + if causal: + visible &= relative_position >= 0 + if window_size[0] is not None: + visible &= relative_position <= window_size[0] + if window_size[1] is not None: + visible &= relative_position >= -window_size[1] + scores.masked_fill_(~visible[None], -torch.inf) + + if sinks is None: + probabilities = torch.softmax(scores, dim=-1) + lse = torch.logsumexp(scores, dim=-1) + else: + row_max = torch.maximum(scores.amax(dim=-1), sinks.float()[:, None]) + weights = torch.exp(scores - row_max[:, :, None]) + denominator = weights.sum(dim=-1) + torch.exp(sinks.float()[:, None] - row_max) + probabilities = weights / denominator[:, :, None] + lse = torch.log(denominator) + row_max + output = torch.einsum("hqk,khd->qhd", probabilities, v) + return output, lse + + +def _assert_attention_close( + output: torch.Tensor, + reference: torch.Tensor, + lse: torch.Tensor | None = None, + lse_reference: torch.Tensor | None = None, +) -> None: + error = (output.float() - reference.float()).abs() + assert error.max().item() < 1e-2 + assert error.mean().item() < 5e-4 + if lse is not None: + assert lse_reference is not None + lse_error = (lse.float() - lse_reference.float()).abs() + assert lse_error.max().item() < 2e-4 + + +def _reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sinks: torch.Tensor, + window_left: int | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return FP32 causal attention-with-sink output and LSE.""" + seq = q.shape[0] + scale = q.shape[-1] ** -0.5 + scores = ( + torch.einsum( + "qhd,kd->hqk", + q.float(), + k[:, 0].float(), + ) + * scale + ) + q_idx = torch.arange(seq, device=q.device)[:, None] + kv_idx = torch.arange(seq, device=q.device)[None, :] + mask = kv_idx <= q_idx + if window_left is not None: + mask &= kv_idx >= q_idx - window_left + scores.masked_fill_(~mask[None], -torch.inf) + + row_max = torch.maximum(scores.amax(dim=-1), sinks.float()[:, None]) + weights = torch.exp(scores - row_max[:, :, None]) + denominator = weights.sum(dim=-1) + torch.exp(sinks.float()[:, None] - row_max) + output = ( + torch.einsum( + "hqk,kd->qhd", + weights, + v[:, 0].float(), + ) + / denominator.T[:, :, None] + ) + lse = torch.log(denominator) + row_max + return output, lse + + +def _relative_bias_reference( + q_parts: list[torch.Tensor], + k_parts: list[torch.Tensor], + v_parts: list[torch.Tensor], + bias_parts: list[torch.Tensor], + *, + causal: bool, + window_size: tuple[int | None, int | None], + softcap: float = 0.0, +) -> torch.Tensor: + outputs = [] + for q, k, v, rel_bias in zip(q_parts, k_parts, v_parts, bias_parts): + q_len, k_len = q.shape[0], k.shape[0] + k = k.float().repeat_interleave(q.shape[1] // k.shape[1], dim=1) + v = v.float().repeat_interleave(q.shape[1] // v.shape[1], dim=1) + scores = torch.einsum("qhd,khd->hqk", q.float(), k) * q.shape[-1] ** -0.5 + q_idx = torch.arange(q_len, device=q.device) + k_len - q_len + kv_idx = torch.arange(k_len, device=q.device) + rel_dist = q_idx[:, None] - kv_idx[None, :] + bias_idx = rel_dist.clamp(0, rel_bias.shape[-1] - 1) + bias = ( + rel_bias.float() + .permute(1, 0, 2) + .gather( + 2, + bias_idx[None].expand(q.shape[1], -1, -1), + ) + ) + bias.masked_fill_( + ~((rel_dist >= 0) & (rel_dist < rel_bias.shape[-1]))[None], + 0.0, + ) + scores += bias + if softcap > 0: + scores = torch.tanh(scores / softcap) * softcap + visible = torch.ones_like(rel_dist, dtype=torch.bool) + if causal: + visible &= rel_dist >= 0 + if window_size[0] is not None: + visible &= rel_dist <= window_size[0] + if window_size[1] is not None: + visible &= rel_dist >= -window_size[1] + scores.masked_fill_(~visible[None], -torch.inf) + outputs.append(torch.einsum("hqk,khd->qhd", torch.softmax(scores, dim=-1), v)) + return torch.cat(outputs) + + +@pytest.mark.parametrize( + ( + "dtype", + "num_q_heads", + "num_kv_heads", + "head_dim", + "head_dim_v", + "causal", + "window_size", + "softcap", + "has_sink", + "pack_gqa", + "preallocate_out", + ), + [ + pytest.param( + torch.bfloat16, + 4, + 4, + 32, + 32, + False, + (None, None), + 0.0, + False, + False, + False, + id="bf16-mha-hd32-global", + ), + pytest.param( + torch.float16, + 6, + 1, + 64, + 96, + True, + (None, None), + 0.0, + True, + True, + True, + id="fp16-mqa-hd64-hdv96-causal-sink", + ), + pytest.param( + torch.bfloat16, + 8, + 2, + 96, + 64, + False, + (47, 11), + 4.0, + False, + True, + False, + id="bf16-gqa-hd96-hdv64-local-softcap", + ), + pytest.param( + torch.bfloat16, + 8, + 2, + 128, + 128, + True, + (None, None), + 0.0, + True, + False, + False, + id="bf16-gqa-hd128-causal-sink-unpacked", + ), + pytest.param( + torch.bfloat16, + 6, + 1, + 192, + 128, + False, + (None, None), + 0.0, + False, + None, + False, + id="bf16-mqa-hd192-hdv128-global-auto-pack", + ), + pytest.param( + torch.float16, + 4, + 4, + 128, + 64, + False, + (63, 7), + 0.0, + False, + False, + True, + id="fp16-mha-hd128-hdv64-local-out", + ), + ], +) +def test_sm120_dense_feature_matrix_matches_reference( + dtype, + num_q_heads, + num_kv_heads, + head_dim, + head_dim_v, + causal, + window_size, + softcap, + has_sink, + pack_gqa, + preallocate_out, +): + """Cover dense MHA/GQA/MQA, dtype, shape, mask, sink, and output modes.""" + torch.manual_seed(20260801 + head_dim + head_dim_v) + batch_size, q_len, k_len = 2, 37, 141 + q = torch.randn( + batch_size, + q_len, + num_q_heads, + head_dim, + device="cuda", + dtype=dtype, + ) + k = torch.randn( + batch_size, + k_len, + num_kv_heads, + head_dim, + device="cuda", + dtype=dtype, + ) + v = torch.randn( + batch_size, + k_len, + num_kv_heads, + head_dim_v, + device="cuda", + dtype=dtype, + ) + sinks = ( + torch.randn(num_q_heads, device="cuda", dtype=torch.bfloat16) + if has_sink + else None + ) + out_buffer = torch.empty( + batch_size, + q_len, + num_q_heads, + head_dim_v, + device="cuda", + dtype=dtype, + ) + output, lse = flash_attn_varlen_func( + q, + k, + v, + causal=causal, + window_size=window_size, + softcap=softcap, + sinks=sinks, + pack_gqa=pack_gqa, + return_softmax_lse=True, + out=out_buffer if preallocate_out else None, + ) + references = [ + _attention_reference( + q[batch_idx], + k[batch_idx], + v[batch_idx], + batch_idx=batch_idx, + causal=causal, + window_size=window_size, + softcap=softcap, + sinks=sinks, + ) + for batch_idx in range(batch_size) + ] + output_reference = torch.stack([reference[0] for reference in references]) + lse_reference = torch.stack([reference[1] for reference in references]) + _assert_attention_close(output, output_reference, lse, lse_reference) + if preallocate_out: + assert output.data_ptr() == out_buffer.data_ptr() + + +@pytest.mark.parametrize( + ( + "dtype", + "num_q_heads", + "num_kv_heads", + "head_dim", + "head_dim_v", + "causal", + "window_size", + "pack_gqa", + "num_splits", + ), + [ + pytest.param( + torch.bfloat16, + 4, + 4, + 64, + 64, + False, + (None, None), + False, + 1, + id="bf16-mha-hd64-global", + ), + pytest.param( + torch.float16, + 8, + 2, + 80, + 48, + True, + (None, None), + True, + 2, + id="fp16-gqa-hd80-hdv48-causal-splitkv", + ), + pytest.param( + torch.bfloat16, + 6, + 1, + 128, + 96, + False, + (63, 9), + None, + 1, + id="bf16-mqa-hd128-hdv96-local-auto-pack", + ), + pytest.param( + torch.bfloat16, + 8, + 2, + 192, + 128, + True, + (127, 0), + False, + 2, + id="bf16-gqa-hd192-hdv128-causal-local-splitkv", + ), + ], +) +def test_sm120_varlen_feature_matrix_matches_reference( + dtype, + num_q_heads, + num_kv_heads, + head_dim, + head_dim_v, + causal, + window_size, + pack_gqa, + num_splits, +): + """Cover ragged batches across topology, asymmetric dimensions, and SplitKV.""" + torch.manual_seed(20260802 + head_dim + head_dim_v) + q_lengths = (19, 53) + k_lengths = (79, 151) + q_parts = [ + torch.randn( + length, + num_q_heads, + head_dim, + device="cuda", + dtype=dtype, + ) + for length in q_lengths + ] + k_parts = [ + torch.randn( + length, + num_kv_heads, + head_dim, + device="cuda", + dtype=dtype, + ) + for length in k_lengths + ] + v_parts = [ + torch.randn( + length, + num_kv_heads, + head_dim_v, + device="cuda", + dtype=dtype, + ) + for length in k_lengths + ] + cu_seqlens_q = torch.tensor( + [0, q_lengths[0], sum(q_lengths)], + device="cuda", + dtype=torch.int32, + ) + cu_seqlens_k = torch.tensor( + [0, k_lengths[0], sum(k_lengths)], + device="cuda", + dtype=torch.int32, + ) + output, lse = flash_attn_varlen_func( + torch.cat(q_parts), + torch.cat(k_parts), + torch.cat(v_parts), + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max(q_lengths), + max_seqlen_k=max(k_lengths), + causal=causal, + window_size=window_size, + num_splits=num_splits, + pack_gqa=pack_gqa, + return_softmax_lse=True, + ) + references = [ + _attention_reference( + q, + k, + v, + batch_idx=batch_idx, + causal=causal, + window_size=window_size, + ) + for batch_idx, (q, k, v) in enumerate(zip(q_parts, k_parts, v_parts)) + ] + output_reference = torch.cat([reference[0] for reference in references]) + lse_reference = torch.cat([reference[1] for reference in references], dim=1) + _assert_attention_close(output, output_reference, lse, lse_reference) + + +def test_sm120_dense_seqused_qk_matches_reference(): + """Dense storage with per-batch used lengths must ignore allocated padding.""" + torch.manual_seed(20260803) + batch_size, max_q, max_k = 3, 64, 160 + num_q_heads, num_kv_heads = 8, 2 + head_dim, head_dim_v = 64, 96 + q_lengths = (17, 41, 63) + k_lengths = (79, 129, 159) + q = torch.randn( + batch_size, + max_q, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k = torch.randn( + batch_size, + max_k, + num_kv_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + v = torch.randn( + batch_size, + max_k, + num_kv_heads, + head_dim_v, + device="cuda", + dtype=torch.bfloat16, + ) + seqused_q = torch.tensor(q_lengths, device="cuda", dtype=torch.int32) + seqused_k = torch.tensor(k_lengths, device="cuda", dtype=torch.int32) + output, lse = flash_attn_varlen_func( + q, + k, + v, + seqused_q=seqused_q, + seqused_k=seqused_k, + max_seqlen_q=max_q, + max_seqlen_k=max_k, + causal=True, + window_size=(95, 0), + pack_gqa=True, + return_softmax_lse=True, + ) + for batch_idx, (q_len, k_len) in enumerate(zip(q_lengths, k_lengths)): + reference, lse_reference = _attention_reference( + q[batch_idx, :q_len], + k[batch_idx, :k_len], + v[batch_idx, :k_len], + batch_idx=batch_idx, + causal=True, + window_size=(95, 0), + ) + _assert_attention_close( + output[batch_idx, :q_len], + reference, + lse[batch_idx, :, :q_len], + lse_reference, + ) + + +@pytest.mark.parametrize( + ("score_mod", "reference_kind", "has_aux"), + [ + pytest.param( + _coordinate_score_mod, + "coordinate", + False, + id="logical-coordinates", + ), + pytest.param(_aux_score_mod, "aux", True, id="aux-tensor"), + ], +) +def test_sm120_packed_gqa_score_mod_matches_reference( + score_mod, + reference_kind, + has_aux, +): + """Score modifiers must see logical packed-GQA indices and runtime aux data.""" + torch.manual_seed(20260804) + q_lengths = (17, 35) + k_lengths = (79, 143) + num_q_heads, num_kv_heads, head_dim = 8, 2, 64 + q_parts = [ + torch.randn( + length, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + for length in q_lengths + ] + k_parts = [ + torch.randn( + length, + num_kv_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + for length in k_lengths + ] + v_parts = [torch.randn_like(k) for k in k_parts] + cu_seqlens_q = torch.tensor( + [0, q_lengths[0], sum(q_lengths)], + device="cuda", + dtype=torch.int32, + ) + cu_seqlens_k = torch.tensor( + [0, k_lengths[0], sum(k_lengths)], + device="cuda", + dtype=torch.int32, + ) + aux_value = 0.125 + aux_tensors = ( + [torch.tensor([aux_value], device="cuda", dtype=torch.float32)] + if has_aux + else None + ) + output, lse = flash_attn_varlen_func( + torch.cat(q_parts), + torch.cat(k_parts), + torch.cat(v_parts), + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max(q_lengths), + max_seqlen_k=max(k_lengths), + causal=True, + pack_gqa=True, + score_mod=score_mod, + aux_tensors=aux_tensors, + return_softmax_lse=True, + ) + references = [ + _attention_reference( + q, + k, + v, + batch_idx=batch_idx, + causal=True, + score_mod=reference_kind, + aux_score=aux_value, + ) + for batch_idx, (q, k, v) in enumerate(zip(q_parts, k_parts, v_parts)) + ] + output_reference = torch.cat([reference[0] for reference in references]) + lse_reference = torch.cat([reference[1] for reference in references], dim=1) + _assert_attention_close(output, output_reference, lse, lse_reference) + + if has_aux: + updated_aux_value = -0.25 + aux_tensors[0].fill_(updated_aux_value) + updated_output, updated_lse = flash_attn_varlen_func( + torch.cat(q_parts), + torch.cat(k_parts), + torch.cat(v_parts), + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max(q_lengths), + max_seqlen_k=max(k_lengths), + causal=True, + pack_gqa=True, + score_mod=score_mod, + aux_tensors=aux_tensors, + return_softmax_lse=True, + ) + updated_references = [ + _attention_reference( + q, + k, + v, + batch_idx=batch_idx, + causal=True, + score_mod=reference_kind, + aux_score=updated_aux_value, + ) + for batch_idx, (q, k, v) in enumerate(zip(q_parts, k_parts, v_parts)) + ] + updated_output_reference = torch.cat( + [reference[0] for reference in updated_references] + ) + updated_lse_reference = torch.cat( + [reference[1] for reference in updated_references], + dim=1, + ) + _assert_attention_close( + updated_output, + updated_output_reference, + updated_lse, + updated_lse_reference, + ) + + +@pytest.mark.parametrize( + ( + "dtype", + "num_q_heads", + "num_kv_heads", + "head_dim", + "head_dim_v", + "causal", + "window_size", + "softcap", + "has_sink", + "pack_gqa", + "num_splits", + ), + [ + pytest.param( + torch.bfloat16, + 8, + 2, + 64, + 96, + True, + (None, None), + 0.0, + True, + True, + 1, + id="bf16-gqa-hd64-hdv96-causal-sink", + ), + pytest.param( + torch.float16, + 4, + 4, + 128, + 64, + False, + (255, 0), + 5.0, + False, + False, + 2, + id="fp16-mha-hd128-hdv64-local-softcap-splitkv", + ), + ], +) +def test_sm120_relative_bias_dense_matches_reference( + dtype, + num_q_heads, + num_kv_heads, + head_dim, + head_dim_v, + causal, + window_size, + softcap, + has_sink, + pack_gqa, + num_splits, +): + """Dense sheared bias composes with sink, softcap, SplitKV, LSE, and out.""" + torch.manual_seed(20260805 + head_dim + head_dim_v) + batch_size, q_len, k_len, rel_extent = 2, 33, 193, 256 + q = torch.randn( + batch_size, + q_len, + num_q_heads, + head_dim, + device="cuda", + dtype=dtype, + ) + k = torch.randn( + batch_size, + k_len, + num_kv_heads, + head_dim, + device="cuda", + dtype=dtype, + ) + v = torch.randn( + batch_size, + k_len, + num_kv_heads, + head_dim_v, + device="cuda", + dtype=dtype, + ) + rel_bias = ( + 0.1 + * torch.randn( + batch_size, + q_len, + num_q_heads, + rel_extent, + device="cuda", + ) + ).to(dtype) + sinks = ( + torch.randn(num_q_heads, device="cuda", dtype=torch.bfloat16) + if has_sink + else None + ) + out_buffer = torch.empty( + batch_size, + q_len, + num_q_heads, + head_dim_v, + device="cuda", + dtype=dtype, + ) + output, lse = flash_attn_varlen_func( + q, + k, + v, + causal=causal, + window_size=window_size, + softcap=softcap, + sinks=sinks, + num_splits=num_splits, + pack_gqa=pack_gqa, + rel_bias=rel_bias, + return_softmax_lse=True, + out=out_buffer, + ) + references = [ + _attention_reference( + q[batch_idx], + k[batch_idx], + v[batch_idx], + batch_idx=batch_idx, + causal=causal, + window_size=window_size, + softcap=softcap, + sinks=sinks, + rel_bias=rel_bias[batch_idx], + ) + for batch_idx in range(batch_size) + ] + output_reference = torch.stack([reference[0] for reference in references]) + lse_reference = torch.stack([reference[1] for reference in references]) + _assert_attention_close(output, output_reference, lse, lse_reference) + assert output.data_ptr() == out_buffer.data_ptr() + + +def test_sm120_relative_bias_cuda_graph_replays_and_eager_remains_reusable(): + """The shearing producer and attention consumer must capture as one graph.""" + torch.manual_seed(20260806) + q_len, k_len, rel_extent = 33, 193, 256 + num_q_heads, num_kv_heads, head_dim = 8, 2, 64 + q = torch.randn( + 1, + q_len, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k = torch.randn( + 1, + k_len, + num_kv_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + v = torch.randn_like(k) + rel_bias = ( + 0.1 + * torch.randn( + 1, + q_len, + num_q_heads, + rel_extent, + device="cuda", + ) + ).to(torch.bfloat16) + rel_bias_prep_cache = {} + + def run(): + return flash_attn_varlen_func( + q, + k, + v, + causal=True, + pack_gqa=True, + rel_bias=rel_bias, + rel_bias_prep_cache=rel_bias_prep_cache, + ) + + eager_before = run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = run() + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + eager_after = run() + torch.cuda.synchronize() + + reference, _ = _attention_reference( + q[0], + k[0], + v[0], + causal=True, + rel_bias=rel_bias[0], + ) + torch.testing.assert_close(eager_before, eager_after, atol=0.0, rtol=0.0) + torch.testing.assert_close(graph_output, eager_before, atol=0.0, rtol=0.0) + _assert_attention_close(graph_output[0], reference) + + +@pytest.mark.parametrize( + ( + "dtype", + "num_q_heads", + "num_kv_heads", + "head_dim", + "head_dim_v", + "q_len", + "page_size", + "causal", + "window_size", + "pack_gqa", + "num_splits", + ), + [ + pytest.param( + torch.float16, + 4, + 4, + 32, + 32, + 7, + 64, + True, + (None, None), + False, + 1, + id="fp16-mha-hd32-page64-causal", + ), + pytest.param( + torch.bfloat16, + 8, + 2, + 64, + 96, + 33, + 128, + False, + (127, 15), + True, + 2, + id="bf16-gqa-hd64-hdv96-page128-local-splitkv", + ), + pytest.param( + torch.bfloat16, + 6, + 1, + 128, + 64, + 5, + 64, + False, + (None, None), + None, + 2, + id="bf16-mqa-hd128-hdv64-page64-global-splitkv", + ), + pytest.param( + torch.bfloat16, + 8, + 2, + 96, + 128, + 17, + 32, + True, + (None, None), + False, + 1, + id="bf16-gqa-hd96-hdv128-page32-causal-unpacked", + ), + pytest.param( + torch.bfloat16, + 8, + 2, + 64, + 64, + 1, + 64, + True, + (None, None), + True, + 0, + id="bf16-gqa-hd64-page64-decode-auto-split", + ), + pytest.param( + torch.bfloat16, + 8, + 1, + 128, + 128, + 1, + 64, + True, + (None, None), + True, + 0, + id="bf16-mqa-hd128-page64-decode-auto-split", + ), + ], +) +def test_sm120_paged_kv_feature_matrix_matches_reference( + dtype, + num_q_heads, + num_kv_heads, + head_dim, + head_dim_v, + q_len, + page_size, + causal, + window_size, + pack_gqa, + num_splits, +): + """Cover paged MHA/GQA/MQA outside the HD256 decode specialization.""" + torch.manual_seed(20260807 + head_dim + head_dim_v + page_size) + batch_size, pages_per_sequence = 2, 6 + max_seqlen_k = pages_per_sequence * page_size + k_lengths = (max_seqlen_k - page_size - 11, max_seqlen_k - 1) + q = torch.randn( + batch_size, + q_len, + num_q_heads, + head_dim, + device="cuda", + dtype=dtype, + ) + k_cache = torch.randn( + batch_size * pages_per_sequence, + page_size, + num_kv_heads, + head_dim, + device="cuda", + dtype=dtype, + ) + v_cache = torch.randn( + batch_size * pages_per_sequence, + page_size, + num_kv_heads, + head_dim_v, + device="cuda", + dtype=dtype, + ) + page_table = ( + torch.randperm( + batch_size * pages_per_sequence, + device="cuda", + dtype=torch.int64, + ) + .to(torch.int32) + .view(batch_size, pages_per_sequence) + ) + seqused_k = torch.tensor(k_lengths, device="cuda", dtype=torch.int32) + output, lse = flash_attn_varlen_func( + q, + k_cache, + v_cache, + seqused_k=seqused_k, + max_seqlen_q=q_len, + max_seqlen_k=max_seqlen_k, + page_table=page_table, + causal=causal, + window_size=window_size, + num_splits=num_splits, + pack_gqa=pack_gqa, + return_softmax_lse=True, + ) + references = [] + for batch_idx, k_len in enumerate(k_lengths): + pages = page_table[batch_idx] + k = k_cache.index_select(0, pages).flatten(0, 1)[:k_len] + v = v_cache.index_select(0, pages).flatten(0, 1)[:k_len] + references.append( + _attention_reference( + q[batch_idx], + k, + v, + batch_idx=batch_idx, + causal=causal, + window_size=window_size, + ) + ) + output_reference = torch.stack([reference[0] for reference in references]) + lse_reference = torch.stack([reference[1] for reference in references]) + _assert_attention_close(output, output_reference, lse, lse_reference) + + if num_splits > 1: + unsplit_output = flash_attn_varlen_func( + q, + k_cache, + v_cache, + seqused_k=seqused_k, + max_seqlen_q=q_len, + max_seqlen_k=max_seqlen_k, + page_table=page_table, + causal=causal, + window_size=window_size, + num_splits=1, + pack_gqa=pack_gqa, + ) + torch.testing.assert_close(output, unsplit_output, atol=4e-3, rtol=0.0) + + +@pytest.mark.parametrize("head_dim", [64, 128]) +@pytest.mark.parametrize("num_kv_heads", [8, 2, 1]) +def test_sm120_low_hd_paged_decode_reuses_cached_host_plan( + monkeypatch, + head_dim, + num_kv_heads, +): + """Low-HD MHA/GQA/MQA decode must reuse its compiled TVM-FFI plan.""" + sm120_forward_host.clear_launch_plans() + cache_hits = [] + original_try_paged_decode = sm120_forward_host.try_paged_decode + + def record_cache_hit(**kwargs): + result = original_try_paged_decode(**kwargs) + cache_hits.append(result is not None) + return result + + monkeypatch.setattr( + sm120_forward_host, + "try_paged_decode", + record_cache_hit, + ) + torch.manual_seed(20260731 + head_dim + num_kv_heads) + batch_size, num_q_heads = 2, 8 + max_seqlen, page_size = 512, 64 + pages_per_request = max_seqlen // page_size + q = torch.randn( + batch_size, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k_cache = torch.randn( + batch_size * pages_per_request, + page_size, + num_kv_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + v_cache = torch.randn_like(k_cache) + page_table = torch.arange( + batch_size * pages_per_request, + device="cuda", + dtype=torch.int32, + ).view(batch_size, pages_per_request) + cache_seqlens = torch.full( + (batch_size,), + max_seqlen, + device="cuda", + dtype=torch.int32, + ) + cu_seqlens_q = torch.arange( + batch_size + 1, + device="cuda", + dtype=torch.int32, + ) + out = torch.empty_like(q) + + def run(): + return flash_attn_with_kvcache( + q=q, + k_cache=k_cache, + v_cache=v_cache, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + max_seqlen_k=max_seqlen, + causal=True, + num_splits=0, + pack_gqa=None, + out=out, + ) + + first = run().clone() + second = run().clone() + torch.cuda.synchronize() + + assert not any(cache_hits[:-1]) + assert cache_hits[-1] + torch.testing.assert_close(first, second, atol=0.0, rtol=0.0) + + +@pytest.mark.parametrize( + ("head_dim", "pack_gqa", "causal", "window_size", "softcap"), + [ + pytest.param(32, True, True, (None, None), 0.0, id="hd32-packed-causal"), + pytest.param( + 96, + False, + True, + (None, None), + 5.0, + id="hd96-unpacked-causal-softcap", + ), + pytest.param(128, True, False, (255, 0), 0.0, id="hd128-packed-local"), + ], +) +def test_sm120_relative_bias_varlen_matches_reference( + head_dim, + pack_gqa, + causal, + window_size, + softcap, +): + """Relative bias must follow logical Q/K positions for every SM120 tile.""" + torch.manual_seed(20260731 + head_dim) + q_lengths = (47, 73) + k_lengths = (193, 321) + num_q_heads, num_kv_heads, rel_extent = 8, 2, 256 + q_parts = [ + torch.randn( + length, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + for length in q_lengths + ] + k_parts = [ + torch.randn( + length, + num_kv_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + for length in k_lengths + ] + v_parts = [torch.randn_like(k) for k in k_parts] + bias_parts = [ + ( + 0.1 + * torch.randn( + length, + num_q_heads, + rel_extent, + device="cuda", + ) + ).to(torch.bfloat16) + for length in q_lengths + ] + cu_seqlens_q = torch.tensor( + [0, q_lengths[0], sum(q_lengths)], + device="cuda", + dtype=torch.int32, + ) + cu_seqlens_k = torch.tensor( + [0, k_lengths[0], sum(k_lengths)], + device="cuda", + dtype=torch.int32, + ) + + output = flash_attn_varlen_func( + torch.cat(q_parts), + torch.cat(k_parts), + torch.cat(v_parts), + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max(q_lengths), + max_seqlen_k=max(k_lengths), + causal=causal, + window_size=window_size, + softcap=softcap, + pack_gqa=pack_gqa, + rel_bias=torch.cat(bias_parts), + ) + reference = _relative_bias_reference( + q_parts, + k_parts, + v_parts, + bias_parts, + causal=causal, + window_size=window_size, + softcap=softcap, + ) + error = (output.float() - reference.float()).abs() + assert error.max().item() < 1e-2 + assert error.mean().item() < 5e-4 + + +def test_sm120_relative_bias_paged_splitkv_is_cache_order_independent(): + """Paged relative coordinates and SplitKV must not collide in the JIT cache.""" + torch.manual_seed(20260731) + q_lengths = (9, 19) + k_lengths = (383, 509) + num_q_heads, num_kv_heads, head_dim = 8, 2, 128 + rel_extent, page_size, pages_per_seq = 256, 128, 4 + q_parts = [ + torch.randn( + length, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + for length in q_lengths + ] + bias_parts = [ + ( + 0.1 + * torch.randn( + length, + num_q_heads, + rel_extent, + device="cuda", + ) + ).to(torch.bfloat16) + for length in q_lengths + ] + k_cache = torch.randn( + len(q_lengths) * pages_per_seq, + page_size, + num_kv_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + v_cache = torch.randn_like(k_cache) + page_table = torch.tensor( + [[3, 0, 7, 1], [6, 2, 5, 4]], + device="cuda", + dtype=torch.int32, + ) + k_parts = [ + k_cache.index_select(0, pages).flatten(0, 1)[:length] + for pages, length in zip(page_table, k_lengths) + ] + v_parts = [ + v_cache.index_select(0, pages).flatten(0, 1)[:length] + for pages, length in zip(page_table, k_lengths) + ] + q = torch.cat(q_parts) + rel_bias = torch.cat(bias_parts) + cu_seqlens_q = torch.tensor( + [0, q_lengths[0], sum(q_lengths)], + device="cuda", + dtype=torch.int32, + ) + seqused_k = torch.tensor(k_lengths, device="cuda", dtype=torch.int32) + + def run(num_splits): + return flash_attn_varlen_func( + q, + k_cache, + v_cache, + cu_seqlens_q=cu_seqlens_q, + seqused_k=seqused_k, + page_table=page_table, + max_seqlen_q=max(q_lengths), + max_seqlen_k=pages_per_seq * page_size, + causal=False, + window_size=(rel_extent - 1, 0), + num_splits=num_splits, + pack_gqa=True, + rel_bias=rel_bias, + ) + + split_output = run(2) + unsplit_output = run(1) + split_output_again = run(2) + reference = _relative_bias_reference( + q_parts, + k_parts, + v_parts, + bias_parts, + causal=False, + window_size=(rel_extent - 1, 0), + ) + + torch.testing.assert_close( + split_output, + split_output_again, + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + split_output, + unsplit_output, + atol=2e-3, + rtol=0.0, + ) + error = (split_output.float() - reference.float()).abs() + assert error.max().item() < 1e-2 + assert error.mean().item() < 5e-4 + + +@pytest.mark.parametrize("window_left", [None, 250]) +def test_sm120_varlen_mqa_hd256_learnable_sink(window_left): + """Cover Q6/KV1 head-dim-256 global and local prefill shapes.""" + torch.manual_seed(1234) + seq, num_q_heads, head_dim = 512, 6, 256 + q = torch.randn(seq, num_q_heads, head_dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn(seq, 1, head_dim, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + sinks = torch.randn(num_q_heads, device="cuda", dtype=torch.bfloat16) + cu_seqlens = torch.tensor([0, seq], dtype=torch.int32, device="cuda") + window_size = (None, None) if window_left is None else (window_left, 0) + out_ref, lse_ref = _reference(q, k, v, sinks, window_left) + + outputs = [] + for pack_gqa in (False, True, None): + out, lse = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=seq, + max_seqlen_k=seq, + softmax_scale=head_dim**-0.5, + causal=True, + window_size=window_size, + sinks=sinks, + pack_gqa=pack_gqa, + return_softmax_lse=True, + ) + output_error = (out.float() - out_ref).abs() + lse_error = (lse - lse_ref).abs() + assert output_error.max().item() < 1e-2 + assert output_error.mean().item() < 5e-4 + assert lse_error.max().item() < 5e-5 + outputs.append(out) + + torch.testing.assert_close(outputs[0], outputs[1], atol=0.0, rtol=0.0) + torch.testing.assert_close(outputs[1], outputs[2], atol=0.0, rtol=0.0) + + +def test_sm120_forward_only_varlen_cache_is_pack_order_independent(monkeypatch): + """MHA/GQA plans must be distinct and accept dynamic sequence offsets.""" + sm120_forward_host.clear_launch_plans() + cache_hits = [] + original_try_varlen = sm120_forward_host.try_varlen + + def record_cache_hit(**kwargs): + result = original_try_varlen(**kwargs) + cache_hits.append(result is not None) + return result + + monkeypatch.setattr(sm120_forward_host, "try_varlen", record_cache_hit) + torch.manual_seed(1234) + lengths = (128, 111, 97) + total, num_q_heads, head_dim = sum(lengths), 8, 64 + q = torch.randn( + total, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + offsets = [0] + for length in lengths: + offsets.append(offsets[-1] + length) + cu_seqlens = torch.tensor(offsets, dtype=torch.int32, device="cuda") + inputs = {} + for name, num_kv_heads in (("mha", 8), ("gqa", 2)): + k = torch.randn( + total, + num_kv_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + v = torch.randn_like(k) + reference = torch.cat( + [ + _attention_reference( + q[start:end], + k[start:end], + v[start:end], + causal=True, + )[0] + for start, end in zip(offsets[:-1], offsets[1:]) + ] + ) + inputs[name] = (k, v, reference) + + for name in ("mha", "gqa", "mha"): + k, v, reference = inputs[name] + output = torch.empty_like(q) + for _ in range(2): + result = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max(lengths), + max_seqlen_k=max(lengths), + softmax_scale=head_dim**-0.5, + causal=True, + pack_gqa=name == "gqa", + out=output, + ) + assert result.data_ptr() == output.data_ptr() + _assert_attention_close(result, reference) + + assert cache_hits == [False, True, False, True, True, True] + + # The compiled plan keys tensor metadata, not the contents of cu_seqlens. + # Repartition the same storage while retaining shape, total, and maximum. + alternate_lengths = (97, 111, 128) + alternate_offsets = [0] + for length in alternate_lengths: + alternate_offsets.append(alternate_offsets[-1] + length) + alternate_cu_seqlens = torch.tensor( + alternate_offsets, + dtype=torch.int32, + device="cuda", + ) + k, v, _ = inputs["mha"] + alternate_reference = torch.cat( + [ + _attention_reference( + q[start:end], + k[start:end], + v[start:end], + causal=True, + )[0] + for start, end in zip(alternate_offsets[:-1], alternate_offsets[1:]) + ] + ) + alternate_output = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=alternate_cu_seqlens, + cu_seqlens_k=alternate_cu_seqlens, + max_seqlen_q=max(alternate_lengths), + max_seqlen_k=max(alternate_lengths), + softmax_scale=head_dim**-0.5, + causal=True, + pack_gqa=False, + ) + assert cache_hits[-1] + _assert_attention_close(alternate_output, alternate_reference) + + +def test_sm120_varlen_padding_ctas_are_inert_across_tile_specializations(): + """Padding CTAs must not consume stale SMEM or write another batch's output.""" + num_sms = torch.cuda.get_device_properties(0).multi_processor_count + num_q_heads, head_dim = 6, 256 + + def make_inputs(lengths, seed): + torch.manual_seed(seed) + total = sum(lengths) + q = torch.randn( + total, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k = torch.randn(total, 1, head_dim, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + sinks = torch.randn(num_q_heads, device="cuda", dtype=torch.bfloat16) + cuts = [0] + for length in lengths: + cuts.append(cuts[-1] + length) + cu_seqlens = torch.tensor(cuts, device="cuda", dtype=torch.int32) + return q, k, v, sinks, cu_seqlens + + def run(inputs, lengths, pack_gqa): + q, k, v, sinks, cu_seqlens = inputs + return flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max(lengths), + max_seqlen_k=max(lengths), + softmax_scale=head_dim**-0.5, + causal=True, + sinks=sinks, + pack_gqa=pack_gqa, + return_softmax_lse=True, + ) + + # Exercise the three SM-normalized selector regions before the padding + # case. This is the order that exposed stale shared-memory state. + rows_per_sm = (58, 84, 100) + for seed, ratio in enumerate(rows_per_sm): + seq = math.ceil(ratio * num_sms / num_q_heads) + inputs = make_inputs([seq], seed) + for pack_gqa in (False, True, False): + run(inputs, [seq], pack_gqa) + torch.cuda.synchronize() + + # At 64 query rows/SM, two batches select M64 while the conservative + # varlen grid contains padding CTAs. + seq = math.ceil(32 * num_sms / num_q_heads) + lengths = [seq, seq] + inputs = make_inputs(lengths, 10) + q, k, v, sinks, _ = inputs + references = [ + _reference( + q[start : start + seq], + k[start : start + seq], + v[start : start + seq], + sinks, + None, + ) + for start in (0, seq) + ] + out_ref = torch.cat([reference[0] for reference in references], dim=0) + lse_ref = torch.cat([reference[1] for reference in references], dim=1) + + for pack_gqa in (False, True, None): + out, lse = run(inputs, lengths, pack_gqa) + output_error = (out.float() - out_ref).abs() + lse_error = (lse - lse_ref).abs() + assert output_error.max().item() < 1e-2 + assert output_error.mean().item() < 5e-4 + assert lse_error.max().item() < 5e-5 + + +@pytest.mark.parametrize("window_left", [None, 192]) +def test_sm120_paged_decode_ragged_splits_are_cache_order_independent( + monkeypatch, + window_left, +): + """Ragged SplitKV must remain correct across uniform/ragged cache reuse.""" + sm120_forward_host.clear_launch_plans() + cache_hits = [] + original_try_paged_decode = sm120_forward_host.try_paged_decode + + def record_cache_hit(**kwargs): + result = original_try_paged_decode(**kwargs) + cache_hits.append(result is not None) + return result + + monkeypatch.setattr( + sm120_forward_host, + "try_paged_decode", + record_cache_hit, + ) + torch.manual_seed(1234) + batch_size, num_q_heads, head_dim = 4, 6, 256 + max_seqlen, page_size = 1024, 64 + pages_per_request = max_seqlen // page_size + q = torch.randn( + batch_size, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k_cache = torch.randn( + batch_size * pages_per_request, + page_size, + 1, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + v_cache = torch.randn_like(k_cache) + page_table = torch.arange( + batch_size * pages_per_request, + device="cuda", + dtype=torch.int32, + ).view(batch_size, pages_per_request) + cu_seqlens_q = torch.arange( + batch_size + 1, + device="cuda", + dtype=torch.int32, + ) + sinks = torch.randn(num_q_heads, device="cuda", dtype=torch.bfloat16) + + def run(lengths): + return flash_attn_with_kvcache( + q=q, + k_cache=k_cache, + v_cache=v_cache, + page_table=page_table, + cache_seqlens=torch.tensor( + lengths, + device="cuda", + dtype=torch.int32, + ), + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + max_seqlen_k=max_seqlen, + causal=True, + window_size=((None, None) if window_left is None else (window_left, 0)), + num_splits=4, + pack_gqa=True, + sinks=sinks, + ) + + uniform_lengths = [max_seqlen] * batch_size + # Exercise both sides of every 64-token tile boundary. N-distributed QK + # has a different accumulator-column layout, so its mask must not use the + # ordinary QK path's R2P column mapping. + ragged_lengths = [1000, 513, 257, 63] + uniform_first = run(uniform_lengths) + ragged_first = run(ragged_lengths) + ragged_second = run(ragged_lengths) + uniform_second = run(uniform_lengths) + + # The public host fast path and the generic compile path may both probe an + # empty cache on the first call. Once registered, every reuse must hit. + assert not any(cache_hits[:-3]) + assert cache_hits[-3:] == [True, True, True] + torch.testing.assert_close(ragged_first, ragged_second, atol=0.0, rtol=0.0) + torch.testing.assert_close(uniform_first, uniform_second, atol=0.0, rtol=0.0) + + reference = torch.empty_like(ragged_first, dtype=torch.float32) + scale = head_dim**-0.5 + for batch_idx, length in enumerate(ragged_lengths): + pages = page_table[batch_idx] + start = 0 if window_left is None else max(0, length - 1 - window_left) + k = k_cache.index_select(0, pages).flatten(0, 1)[start:length, 0].float() + v = v_cache.index_select(0, pages).flatten(0, 1)[start:length, 0].float() + scores = q[batch_idx].float() @ k.T * scale + row_max = torch.maximum(scores.amax(dim=-1), sinks.float()) + weights = torch.exp(scores - row_max[:, None]) + denominator = weights.sum(dim=-1) + torch.exp(sinks.float() - row_max) + reference[batch_idx] = weights @ v / denominator[:, None] + + error = (ragged_first.float() - reference).abs() + assert error.max().item() < 1e-2 + assert error.mean().item() < 5e-4 + + +def test_sm120_paged_decode_transpose_is_cache_order_independent(): + """Gather and page-TMA transpose must not share a compiled specialization.""" + torch.manual_seed(20260729) + batch_size, num_q_heads, head_dim = 1, 6, 256 + max_seqlen, page_size = 2048, 64 + num_pages = max_seqlen // page_size + q = torch.randn( + batch_size, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k_cache = torch.randn( + num_pages, + page_size, + 1, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + v_cache = torch.randn_like(k_cache) + page_table = torch.randperm( + num_pages, + device="cuda", + dtype=torch.int64, + ).to( + torch.int32 + )[None] + cache_seqlens = torch.tensor( + [max_seqlen], + device="cuda", + dtype=torch.int32, + ) + cu_seqlens_q = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + sinks = torch.randn(num_q_heads, device="cuda", dtype=torch.bfloat16) + + def run(num_splits): + return flash_attn_with_kvcache( + q=q, + k_cache=k_cache, + v_cache=v_cache, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + max_seqlen_k=max_seqlen, + causal=True, + num_splits=num_splits, + pack_gqa=True, + sinks=sinks, + ) + + # S32 owns one KV tile per CTA and selects gather. S16 owns two and + # selects the full-transpose page-TMA class. Alternate both compile/cache + # entries in the order that previously exposed an illegal access. + normal_first = run(32) + transpose_after = run(16) + transpose_repeat = run(16) + normal_after = run(32) + torch.cuda.synchronize() + + torch.testing.assert_close(normal_first, normal_after, atol=0.0, rtol=0.0) + torch.testing.assert_close( + transpose_after, + transpose_repeat, + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + normal_first, + transpose_after, + atol=2e-3, + rtol=0.0, + ) + + pages = page_table[0] + k = k_cache.index_select(0, pages).flatten(0, 1)[:, 0].float() + v = v_cache.index_select(0, pages).flatten(0, 1)[:, 0].float() + scores = q[0].float() @ k.T * head_dim**-0.5 + row_max = torch.maximum(scores.amax(dim=-1), sinks.float()) + weights = torch.exp(scores - row_max[:, None]) + denominator = weights.sum(dim=-1) + torch.exp(sinks.float() - row_max) + reference = weights @ v / denominator[:, None] + error = (transpose_after[0].float() - reference).abs() + assert error.max().item() < 1e-2 + assert error.mean().item() < 5e-4 + + +@pytest.mark.parametrize( + ("num_q_heads", "pack_gqa", "expected_transpose", "expected_split_qk"), + [ + pytest.param(6, True, True, False, id="transpose"), + pytest.param(16, True, False, True, id="split-qk"), + pytest.param(1, None, False, True, id="auto-mha-split-qk"), + pytest.param(6, False, False, False, id="single-qk"), + ], +) +def test_sm120_paged_decode_graph_pdl_is_correct_and_eager_reusable( + monkeypatch, + num_q_heads, + pack_gqa, + expected_transpose, + expected_split_qk, +): + """Every SplitKV dataflow must safely launch its captured combine early.""" + sm120_forward_host.clear_launch_plans() + cache_hits = [] + original_try_paged_decode = sm120_forward_host.try_paged_decode + + def record_cache_hit(**kwargs): + result = original_try_paged_decode(**kwargs) + cache_hits.append(result is not None) + return result + + monkeypatch.setattr( + sm120_forward_host, + "try_paged_decode", + record_cache_hit, + ) + torch.manual_seed(20260730) + batch_size, head_dim = 1, 256 + max_seqlen, page_size = 2048, 64 + num_pages = max_seqlen // page_size + captured_plans = [] + original_resolve_plan = Sm120ForwardHost.resolve_plan + + def recording_resolve_plan(**kwargs): + plan = original_resolve_plan(**kwargs) + if kwargs["is_stream_capturing"]: + captured_plans.append(plan) + return plan + + monkeypatch.setattr( + Sm120ForwardHost, + "resolve_plan", + staticmethod(recording_resolve_plan), + ) + q = torch.randn( + batch_size, + num_q_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k_cache = torch.randn( + num_pages, + page_size, + 1, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + v_cache = torch.randn_like(k_cache) + page_table = torch.randperm( + num_pages, + device="cuda", + dtype=torch.int64, + ).to( + torch.int32 + )[None] + cache_seqlens = torch.tensor( + [max_seqlen], + device="cuda", + dtype=torch.int32, + ) + cu_seqlens_q = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + sinks = torch.randn(num_q_heads, device="cuda", dtype=torch.bfloat16) + + def run(): + return flash_attn_with_kvcache( + q=q, + k_cache=k_cache, + v_cache=v_cache, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + max_seqlen_k=max_seqlen, + causal=True, + num_splits=16, + pack_gqa=pack_gqa, + sinks=sinks, + ) + + eager_before = run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = run() + for _ in range(4): + graph.replay() + torch.cuda.synchronize() + eager_after = run() + torch.cuda.synchronize() + + # Eager warmup and graph capture may each probe through both the public + # bridge and the generic compile path. Captured execution must not reuse an + # eager launch plan, while the final eager call must. + assert not any(cache_hits[:-1]) + assert cache_hits[-1] + torch.testing.assert_close(eager_before, eager_after, atol=0.0, rtol=0.0) + torch.testing.assert_close( + graph_output, + eager_before, + atol=2e-3, + rtol=0.0, + ) + assert captured_plans + assert all(plan.num_splits > 1 for plan in captured_plans) + assert all(plan.launch_split_combine_early for plan in captured_plans) + assert all(plan.transpose_qk_pv is expected_transpose for plan in captured_plans) + assert all(plan.split_qk_n is expected_split_qk for plan in captured_plans) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__]))