diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 8e1ac39672..c04db3ec6f 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -81,6 +81,8 @@ def build_flash_attn_dualwave_swp_module( cross_seqlen=False, paged=False, kv_cache_layout="linear", + causal_lpt=True, + page_size=None, return_lse=False, has_bias=False, has_alibi=False, @@ -91,7 +93,8 @@ def build_flash_attn_dualwave_swp_module( Supports dense self-attention, varlen packed QKV, and paged-KV cache modes. Varlen uses cu_seqlens_q/kv with per-batch self-attention ranges. - Paged mode keeps Q/O dense and maps KV tiles through BlockTable pages. + Paged mode keeps Q/O dense and maps KV tiles through BlockTable pages; + ``page_size`` defaults to the kernel's BLOCK_N and may be any multiple of it. Varlen, paged, and split-K are mutually constrained by the caller. has_bias adds a dense elementwise attention bias and works on both the dense and varlen paths; see the HAS_BIAS block below for the per-mode shapes. @@ -146,6 +149,8 @@ def build_flash_attn_dualwave_swp_module( paged=paged, kv_cache_layout=kv_cache_layout, kv_vectorized=KV_VECTORIZED, + causal_lpt=causal_lpt, + page_size=page_size, return_lse=return_lse, xcd_swizzle=_xcd_swizzle, ) @@ -232,6 +237,8 @@ def flash_attn_dualwave_swp_gfx950_kernel( ctx.init_runtime_indices() ctx.init_lds(SharedStorage) ctx.init_thread_mapping() + if const_expr(traits.CAUSAL_LPT): + ctx.init_causal_lpt_order() ctx.init_sequence_lengths() ctx.init_descriptors() ctx.init_workspace() diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 13bccd3e76..ed9e543c1e 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -140,6 +140,25 @@ def _dense_routes_to_dualwave(batch: int, seq_len: int) -> bool: return seq_len >= _DENSE_DUALWAVE_MIN_SEQ +def _varlen_causal_lpt(causal_lpt, batch: int, max_seqlen_q: int) -> bool: + """Resolve causal_lpt=None for the varlen route; True/False stay as given. + + Varlen sizes grid_dim.y from the batch-wide max seqlen, so every sequence shorter + than that leaves inactive workgroups on the axis LPT reverses. While one sequence + spans at least twice as many q-blocks as the batch has sequences, reversing shortens + the causal tail as intended (measured 1.006-1.075x); past that the batch axis + dominates the dispatch and the reversal costs up to 4%. + + The factor of two is not slack. Sweeping B as a fraction of the q-block count at + max_seqlen_q 8k/16k/32k puts the crossover at 0.75-0.875 of it, consistently across + all three, so admitting the whole range would keep LPT on over a band that measures + 0.974-0.998x. Half is the widest round threshold whose entire range is a win. + """ + if causal_lpt is not None: + return bool(causal_lpt) + return 2 * batch <= -(-int(max_seqlen_q) // DUALWAVE_SWP_BLOCK_M) + + def _dense_light_cu(device: torch.device) -> int: try: return int(torch.cuda.get_device_properties(device.index).multi_processor_count) @@ -208,6 +227,7 @@ def _build_dense_dualwave( setprio: bool, debug_lazy_counts: bool, enable_stagger: bool, + causal_lpt: bool = True, return_lse: bool = False, has_bias: bool = False, has_alibi: bool = False, @@ -230,6 +250,7 @@ def _build_dense_dualwave( dualwave_swp_setprio=setprio, dualwave_swp_debug_lazy_counts=debug_lazy_counts, dualwave_swp_enable_stagger=enable_stagger, + causal_lpt=causal_lpt, return_lse=return_lse, has_bias=has_bias, has_alibi=has_alibi, @@ -305,6 +326,7 @@ def _build_varlen( setprio: bool, debug_lazy_counts: bool, enable_stagger: bool, + causal_lpt: bool = True, return_lse: bool = False, has_bias: bool = False, has_alibi: bool = False, @@ -327,6 +349,7 @@ def _build_varlen( dualwave_swp_setprio=setprio, dualwave_swp_debug_lazy_counts=debug_lazy_counts, dualwave_swp_enable_stagger=enable_stagger, + causal_lpt=causal_lpt, return_lse=return_lse, has_bias=has_bias, has_alibi=has_alibi, @@ -382,6 +405,7 @@ def _build_splitk( lazy_rescale: bool, setprio: bool, enable_stagger: bool, + causal_lpt: bool = True, return_lse: bool = False, has_bias: bool = False, has_alibi: bool = False, @@ -402,6 +426,7 @@ def _build_splitk( dualwave_swp_lazy_rescale=lazy_rescale, dualwave_swp_setprio=setprio, dualwave_swp_enable_stagger=enable_stagger, + causal_lpt=causal_lpt, return_lse=return_lse, has_bias=has_bias, has_alibi=has_alibi, @@ -425,6 +450,8 @@ def _build_paged( num_kv_splits: int = 1, varlen: bool = False, kv_cache_layout: str = "linear", + causal_lpt: bool = True, + page_size: int = 64, return_lse: bool = False, has_bias: bool = False, ): @@ -440,6 +467,9 @@ def _build_paged( ``kv_cache_layout`` selects the physical page layout: "linear" [NumBlocks,PageSize,Hkv,D] or "vectorized" (aiter 5D). + + ``page_size`` is the token count per page; pages larger than the 64-token KV tile + are walked tile by tile. """ from kernels.attention.flash_attn_gfx950 import build_flash_attn_dualwave_swp_module @@ -459,6 +489,8 @@ def _build_paged( dualwave_swp_lazy_rescale=lazy_rescale, dualwave_swp_setprio=setprio, dualwave_swp_enable_stagger=enable_stagger, + causal_lpt=causal_lpt, + page_size=page_size, return_lse=return_lse, has_bias=has_bias, ) @@ -466,11 +498,46 @@ def _build_paged( # ── paged-KV native path ──────────────────────────────────────────────────── -# gfx950 dualwave paged-KV currently supports exactly one configuration. -_PAGED_PAGE_SIZE = 64 +# KV tile length the dualwave pipeline consumes per step; a page holds one or more. +_PAGED_TILE = 64 +# Block-table entries (one per KV tile) staged in LDS per split. _PAGED_BT_LDS_SIZE = 2048 +def _paged_page_descriptor_fits(page_elems, elem_size): + """A page descriptor carries a signed-int32 element extent and a 32-bit byte span. + + Overflowing either wraps silently and reads the wrong rows. + """ + return page_elems < 2**31 and page_elems * elem_size < 2**32 + + +def _paged_kv_row_stride(k, v, page_size, num_kv_heads, head_dim): + """Element stride between KV rows of a [NumBlocks, PageSize, Hkv, D] page cache. + + None means "densely packed", and the kernel falls back to its Hkv*D default. vLLM + hands K and V in as views into one packed cache entry, so a row is wider than + Hkv*D; the kernel can take that stride at runtime as long as pages remain + page_size rows apart. The head stride is baked to D and only read when there is + more than one KV head, so a single-head cache is free to space its heads however + it likes. Layouts that do not fit are reported as dense and get copied by the + caller, as before. + """ + if k.stride() != v.stride(): + return None + block_stride, row_stride, head_stride, elem_stride = (int(s) for s in k.stride()) + dense_row = num_kv_heads * head_dim + if elem_stride != 1 or row_stride < dense_row or block_stride != page_size * row_stride: + return None + if num_kv_heads > 1 and head_stride != head_dim: + return None + if row_stride == dense_row: + return None + if not _paged_page_descriptor_fits(page_size * row_stride, k.element_size()): + return None + return row_stride + + def _flydsl_flash_attn_paged( q: torch.Tensor, k: torch.Tensor, @@ -494,13 +561,15 @@ def _flydsl_flash_attn_paged( dualwave_swp_lazy_rescale: bool, dualwave_swp_setprio: bool, dualwave_swp_enable_stagger: bool, + causal_lpt: Optional[bool], stream, ) -> torch.Tensor: """Native paged-KV attention on the gfx950 dualwave kernel. - Supported config ONLY (anything else raises): linear/vectorized cache layout - [NumBlocks, PageSize=64, NumKVHeads, HeadDim], vLLM lookup (block_table + - seqlen_k), causal, D=64/128, dtype bf16/f16. + Supported config ONLY (anything else raises): linear (page_size any multiple of the + 64-token KV tile) or vectorized (page_size 64) cache layout + [NumBlocks, PageSize, NumKVHeads, HeadDim], vLLM lookup (block_table + seqlen_k), + causal, D=64/128, dtype bf16/f16. - Dense 4D Q ``[B, Sq, H, D]``: split-K (num_kv_splits>1) supported (seq_len>=384). - Varlen packed Q ``[total_q, H, D]`` (cu_seqlens_q given): paged K/V looked up per kv-tile via block_table; split-K not supported (matches dense varlen). @@ -553,20 +622,35 @@ def _flydsl_flash_attn_paged( page_size = int(k.shape[1]) Hkv = int(k.shape[2]) k_head_dim = int(k.shape[3]) - if page_size != _PAGED_PAGE_SIZE: + if vectorized and page_size != _PAGED_TILE: + raise NotImplementedError( + f"flydsl_flash_attn_func: vectorized paged KV supports page_size={_PAGED_TILE} only, got {page_size}" + ) + if page_size < _PAGED_TILE or page_size % _PAGED_TILE != 0: raise NotImplementedError( - f"flydsl_flash_attn_func: native paged KV supports page_size={_PAGED_PAGE_SIZE} only, got {page_size}" + f"flydsl_flash_attn_func: native paged KV supports page_size as a multiple of " + f"{_PAGED_TILE}, got {page_size}" ) if D not in (64, 128): raise NotImplementedError(f"flydsl_flash_attn_func: native paged KV supports head_dim=64 or 128, got {D}") if k_head_dim != D: raise ValueError(f"flydsl_flash_attn_func: paged K head_dim ({k_head_dim}) must match q head_dim ({D})") - if num_kv_heads is None: num_kv_heads = Hkv if H % num_kv_heads != 0: raise ValueError(f"flydsl_flash_attn_func: num_heads ({H}) must be divisible by num_kv_heads ({num_kv_heads})") + # Dense is the narrowest pitch, so unlike a strided one it has no copy to fall back on. + # num_kv_heads, not Hkv: the kernel builds DEFAULT_STRIDE_KV_N from the resolved value. + dense_page_elems = page_size * num_kv_heads * D + if not _paged_page_descriptor_fits(dense_page_elems, k.element_size()): + raise NotImplementedError( + f"flydsl_flash_attn_func: paged page_size {page_size} with {num_kv_heads} KV heads and " + f"head_dim {D} needs a {dense_page_elems}-element " + f"({dense_page_elems * k.element_size()}-byte) page descriptor, which exceeds its " + f"limits of {2**31} elements and {2**32} bytes" + ) + # Split-K (paged, dense only): split the KV dimension across grid_z = B*num_kv_splits # workgroups + a combine pass. Fills the GPU for low-occupancy shapes (small B / few # heads), where single-split paged underutilizes the device. @@ -581,15 +665,18 @@ def _flydsl_flash_attn_paged( # paged always uses cross masking (per-batch seqlen_q/seqlen_kv come from cu_seqlens). _kv_lens = seqlen_k.reshape(-1).tolist() if max_seqlen_kv is None or (bias is not None and not varlen) else None skv = int(max_seqlen_kv) if max_seqlen_kv is not None else int(max(_kv_lens)) - max_kv_pages = (skv + page_size - 1) // page_size - max_pages_per_split = (max_kv_pages + int(num_kv_splits) - 1) // int(num_kv_splits) - if max_pages_per_split > _PAGED_BT_LDS_SIZE: - max_supported_kv = _PAGED_BT_LDS_SIZE * int(num_kv_splits) * page_size + # The staged block table is indexed per KV tile, not per page: a page_size + # above _PAGED_TILE spans several tiles. + max_kv_tiles = (skv + _PAGED_TILE - 1) // _PAGED_TILE + max_tiles_per_split = (max_kv_tiles + int(num_kv_splits) - 1) // int(num_kv_splits) + if max_tiles_per_split > _PAGED_BT_LDS_SIZE: + max_supported_kv = _PAGED_BT_LDS_SIZE * int(num_kv_splits) * _PAGED_TILE raise NotImplementedError( f"flydsl_flash_attn_func: paged KV length {skv} exceeds block-table LDS window " - f"({_PAGED_BT_LDS_SIZE} pages/split, max_kv_len={max_supported_kv} for " - f"num_kv_splits={num_kv_splits}, page_size={page_size})" + f"({_PAGED_BT_LDS_SIZE} tiles/split, max_kv_len={max_supported_kv} for " + f"num_kv_splits={num_kv_splits})" ) + stride_kv_n = None if vectorized else _paged_kv_row_stride(k, v, page_size, Hkv, D) if varlen: cross = bool(cross_seqlen) if cross_seqlen is not None else True else: @@ -628,9 +715,12 @@ def _flydsl_flash_attn_paged( launch_stream = torch.cuda.current_stream(q.device) if stream is None else stream # Short paged attention uses generic light; unsupported cases stay on dualwave. _arch = _gpu_arch(q.device) + # The generic light kernel assumes one KV tile per densely packed page. _paged_light_ok = ( (num_kv_splits <= 1) and bias is None # the light paged kernel has no bias path + and page_size == _PAGED_TILE + and stride_kv_n is None and D in (64, 128) and dtype_str in ("bf16", "f16") and (not _arch.startswith("gfx950") or Sq <= _VARLEN_LIGHT_MAX_SEQ) @@ -669,18 +759,24 @@ def _flydsl_flash_attn_paged( varlen=varlen, kv_cache_layout=kv_cache_layout, has_bias=bias is not None, + causal_lpt=_varlen_causal_lpt(causal_lpt, B, Sq) if varlen else causal_lpt is not False, + page_size=page_size, ) if out is None: out = torch.empty_like(q) # Keep tensors in natural shape; flattening can overflow int32 C-ABI dims. # The paged kernel rebuilds per-batch/page descriptors from base pointers. q_flat = q.contiguous() - k_flat = k.contiguous() - v_flat = v.contiguous() + # A strided cache is already in a layout the kernel addresses directly; + # copying it would both cost a cache-sized memcpy and drop the K/V interleave. + k_flat = k if stride_kv_n is not None else k.contiguous() + v_flat = v if stride_kv_n is not None else v.contiguous() o_flat = out.contiguous() kwargs = dict(block_table=block_table_i32, block_table_stride=block_table_stride, stream=launch_stream) if bias is not None: kwargs["bias"] = bias + if stride_kv_n is not None: + kwargs["stride_kv_n"] = stride_kv_n if varlen: kwargs["cu_seqlens_q"] = cu_seqlens_q kwargs["cu_seqlens_kv"] = cu_seqlens_kv @@ -792,6 +888,12 @@ def flydsl_flash_attn_func( # stay on one XCD instead of every XCD re-streaming that head's K/V. None # auto-selects on the shapes it helps; True/False force it. Dense non-fp8 only. dualwave_swp_xcd_swizzle: Optional[bool] = None, + # Reverse the causal q-block grid axis so the heaviest blocks issue first, + # shortening the makespan tail. Bit-identical, being a permutation of + # workgroup -> q-block. None auto-selects on varlen, where a batch holding more than + # half as many sequences as one sequence has q-blocks makes the reversal + # counterproductive; True/False force it. See flash_attn_utils._init_causal_lpt_order. + causal_lpt: Optional[bool] = None, # Debug: pass a pre-allocated float32[2] tensor to enable the lazy-rescale # branch counter (dualwave_swp_debug_lazy_counts=True). Only for dense mode. debug_counts: Optional[torch.Tensor] = None, @@ -808,7 +910,14 @@ def flydsl_flash_attn_func( v: Value tensor, same shape as k. Paged KV cache (future ABI): physical K/V cache tensors. Supported ``kv_cache_layout`` values: - - ``linear``: 4D paged K/V, ``[NumBlocks, PageSize, NumKVHeads, HeadDim]``. + - ``linear``: 4D paged K/V, ``[NumBlocks, PageSize, NumKVHeads, HeadDim]``, + page_size any multiple of the 64-token KV tile. K/V may be views into a + wider packed cache row as long as pages stay ``PageSize`` rows apart. + Consumed without a copy only when heads are packed NHD (``Hkv == 1`` + or ``head_stride == HeadDim``); a GQA cache laid out HND, mismatched + K/V strides, or a page pitch past the descriptor bound is repacked + via ``.contiguous()``. Not a general zero-copy promise for an + arbitrary vLLM cache. - ``linear3d``: page_size=1 special case, ``[NumBlocks, NumKVHeads, HeadDim]``. - ``vectorized``: aiter-style 5D K/V, where @@ -816,7 +925,8 @@ def flydsl_flash_attn_func( and ``V = [NumBlocks, NumKVHeads, PageSize / kVectorSize, HeadDim, kVectorSize]``. Here ``kVectorSize = 16 / element_size`` (bf16/fp16: 8, fp8: 16); - page_size and head_dim must be divisible by it. + page_size and head_dim must be divisible by it, and ``page_size`` is + 64 only (one KV tile per page); larger pages are rejected. causal: Bottom-right aligned causal mask when True. num_kv_heads: KV head count for GQA/MQA; defaults to q num_heads (MHA). cu_seqlens_q: Int32 ``[B+1]`` cumulative Q token counts (varlen). @@ -1012,6 +1122,7 @@ def flydsl_flash_attn_func( dualwave_swp_lazy_rescale=dualwave_swp_lazy_rescale, dualwave_swp_setprio=dualwave_swp_setprio, dualwave_swp_enable_stagger=dualwave_swp_enable_stagger, + causal_lpt=causal_lpt, stream=stream, ) @@ -1197,6 +1308,7 @@ def flydsl_flash_attn_func( lazy_rescale=dualwave_swp_lazy_rescale, setprio=dualwave_swp_setprio, enable_stagger=dualwave_swp_enable_stagger, + causal_lpt=causal_lpt is not False, return_lse=return_lse, has_bias=has_bias, has_alibi=has_alibi, @@ -1245,6 +1357,7 @@ def flydsl_flash_attn_func( setprio=dualwave_swp_setprio, debug_lazy_counts=debug_lazy, enable_stagger=dualwave_swp_enable_stagger, + causal_lpt=_varlen_causal_lpt(causal_lpt, B, Sq), return_lse=return_lse, has_bias=has_bias, has_alibi=has_alibi, @@ -1263,6 +1376,17 @@ def flydsl_flash_attn_func( ) # bias/ALiBi force dualwave: the generic dense kernel folds in neither. if debug_lazy or has_bias or has_alibi or has_sink or (can_dualwave and _dense_routes_to_dualwave(B, Sq)): + # Workgroups map to XCDs as linear_id % 8, and linear_id is + # bx + by*H + bz*H*nqb, so with H % 8 == 0 a head-fast grid pins + # head h to XCD h % 8. The ~256 resident workgroups span all H + # heads within one batch, leaving each XCD to juggle H/8 K/V + # streams against its L2 slice. The head-slow remap in + # _init_dualwave_thread_mapping puts the resident window inside a + # single head instead: 1 stream. Measured penalty for leaving it + # off tracks H/8 (-6% at 8 streams, -3% at 4, nil at <=2), not the + # hit rate and not traffic volume. Bijective, so output is + # unchanged. NB: that function's own comment states the opposite + # rationale ("scatter across all XCDs") and is wrong. num_q_blocks = -(-int(Sq) // DUALWAVE_SWP_BLOCK_M) if dualwave_swp_xcd_swizzle is None: xcd_swizzle = not causal and H % NUM_XCD_GFX950 == 0 and num_q_blocks >= MIN_Q_BLOCKS_XCD_SWIZZLE @@ -1281,6 +1405,7 @@ def flydsl_flash_attn_func( setprio=dualwave_swp_setprio, debug_lazy_counts=debug_lazy, enable_stagger=dualwave_swp_enable_stagger, + causal_lpt=causal_lpt is not False, return_lse=return_lse, has_bias=has_bias, has_alibi=has_alibi, diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 06ca141e4b..88f9d5647a 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -775,9 +775,15 @@ def _v_buf_base(traits, buf_id): def _kv_tile_addr(traits, tile_start, kv_gmem_elem_offset, kv_head_elem_offset, stride_kv_n_v): - """Return (src_base, soffset): dense uses tile_start*stride; paged folds page offset into the descriptor.""" + """Return (src_base, soffset): dense uses tile_start*stride; paged folds page offset into the descriptor. + + When a page holds more than one KV tile (PAGE_SIZE > BLOCK_N) the descriptor still + points at the page base, so the tile's row offset inside the page rides on soffset. + """ if const_expr(traits.PAGED): - return kv_head_elem_offset, 0 + if const_expr(traits.TILES_PER_PAGE == 1): + return kv_head_elem_offset, 0 + return kv_head_elem_offset, (tile_start % fx.Index(traits.PAGE_SIZE)) * stride_kv_n_v return kv_gmem_elem_offset, tile_start * stride_kv_n_v @@ -1008,6 +1014,9 @@ def _init_dualwave_thread_mapping(ctx): # so output is bit-identical; split-K's third grid axis would not survive it. # Non-causal only: under a causal mask q-block i does work proportional to i, so # making q_block the fast axis clusters unequal work and costs 7% (measured). + # Measured on causal paged prefill: opting causal into this remap costs 3-5% even + # when paired with CAUSAL_LPT below, which removes the load imbalance this comment + # blames. So the causal penalty is not only imbalance, and the non-causal gate stands. if const_expr( traits.XCD_SWIZZLE and not traits.SPLITK and not traits.CAUSAL and traits.NUM_HEADS_Q % NUM_XCD_GFX950 == 0 ): @@ -1057,6 +1066,27 @@ def _init_dualwave_thread_mapping(ctx): ctx.kv_head_idx = ctx.h_kv_idx +def _init_causal_lpt_order(ctx): + """Issue causal q-blocks longest-first by reversing the q-block grid axis. + + Causal work per q-block grows with the block index and workgroups dispatch in + flattened-id order, so the natural order issues the heaviest block last and the + makespan carries its tail. Must run after the thread mapping and before anything + that reads q_start (sequence lengths, tile bounds, q_row). + + The reversal is taken over ``grid_dim.y``, the extent the launcher actually + dispatched, which makes it a bijection on the dispatched q-block indices by + construction. Deriving it from ``seq_len`` instead would only agree while + ``seq_len`` is the same max the grid was sized from -- and under varlen it is a + per-batch maximum, so a mismatch there would silently drop or double-compute + q-blocks rather than merely reorder them. + + Shared by DualwaveKernelContext and DualwaveFp8KernelContext. + """ + ctx.q_block_idx = fx.Index(gpu.grid_dim.y) - fx.Index(1) - ctx.q_block_idx + ctx.q_start = ctx.q_block_idx * ctx.traits.BLOCK_M + + def _init_dualwave_q_row(ctx): """Set q_row / q_row_i32 / q_start_pos_i32 on a dualwave-style context.""" traits = ctx.traits @@ -1574,6 +1604,9 @@ class DualwaveSwpTraits: LGKMCNT_0_ONLY: int RETURN_LSE: bool = False XCD_SWIZZLE: bool = False + CAUSAL_LPT: bool = False + PAGE_SIZE: int = 0 + TILES_PER_PAGE: int = 1 @property def cache_tag(self): @@ -1598,6 +1631,8 @@ def cache_tag(self): self.KV_VECTORIZED, self.RETURN_LSE, self.XCD_SWIZZLE, + self.CAUSAL_LPT, + self.PAGE_SIZE, ) @@ -1621,6 +1656,8 @@ def _make_dualwave_swp_traits( kv_vectorized=None, return_lse=False, xcd_swizzle=False, + causal_lpt=True, + page_size=None, ): """Build gfx950 DUALWAVE_SWP compile-time layout traits.""" # Tile shape and wave geometry follow the gfx950 dual-wave 8-wave CTA. @@ -1633,6 +1670,20 @@ def _make_dualwave_swp_traits( block_size = num_waves * warp_size rows_per_wave = 32 + # A page holds PAGE_SIZE tokens while the KV pipeline consumes BLOCK_N at a time, so + # one page may cover several tiles. Only the page-table lookup and the in-page row + # offset depend on this; the tile loop itself stays BLOCK_N-granular. + if paged: + page_size = block_n if page_size is None else int(page_size) + if page_size < block_n or page_size % block_n != 0: + raise ValueError(f"paged page_size must be a positive multiple of {block_n}, got {page_size}") + tiles_per_page = page_size // block_n + if tiles_per_page > 1 and kv_cache_layout != "linear": + raise ValueError(f"page_size {page_size} > {block_n} requires the linear KV cache layout") + else: + page_size = 0 + tiles_per_page = 1 + # QK walks D in 16-wide MFMA K steps; PV consumes K_SUB_N in two 16-token steps. k_step_qk = 16 k_steps_qk = head_dim // k_step_qk @@ -1773,6 +1824,9 @@ def _make_dualwave_swp_traits( LGKMCNT_0_ONLY=0xC07F, RETURN_LSE=bool(return_lse), XCD_SWIZZLE=bool(xcd_swizzle), + CAUSAL_LPT=bool(causal_lpt) and bool(causal), + PAGE_SIZE=page_size, + TILES_PER_PAGE=tiles_per_page, ) @@ -3365,6 +3419,9 @@ def init_lds(self, shared_storage): def init_thread_mapping(self): _init_dualwave_thread_mapping(self) + def init_causal_lpt_order(self): + _init_causal_lpt_order(self) + def init_sequence_lengths(self, CuSeqQ=None, CuSeqKv=None): if CuSeqQ is None: CuSeqQ = self.CuSeqQ @@ -3458,7 +3515,7 @@ def init_descriptors( if const_expr(traits.PAGED): self.k_div = None self.v_div = None - page_elems = fx.Index(traits.BLOCK_N) * self.stride_kv_n_v + page_elems = fx.Index(traits.PAGE_SIZE) * self.stride_kv_n_v self.page_byte_stride = page_elems * fx.Index(traits.BF16_BYTES) self.page_nrec_bytes = fx.Int64(self.page_byte_stride) self.page_layout = fx.make_layout(fx.Int32(page_elems), fx.Int32(1)) @@ -3648,7 +3705,13 @@ def _load_block_table_to_lds(): dst = buffer_ops.get_element_ptr(lds_bt_base_ptr, byte_offset=byte_off, elem_type=T.i8) llvm.StoreOp(as_mlir_value(fx.Int32(0)), dst) if tile_idx < num_kv_tiles: - row_idx = batch_idx * block_table_stride_v + tile_idx + # The LDS copy stays tile-indexed; when a page spans several tiles + # its id is simply repeated, keeping load_page_id_lds unchanged. + if const_expr(traits.TILES_PER_PAGE == 1): + page_idx = tile_idx + else: + page_idx = tile_idx // fx.Index(traits.TILES_PER_PAGE) + row_idx = batch_idx * block_table_stride_v + page_idx v = fly.copy_atom_call_ssa([bt_v1i32], bt_atom, fx.slice(bt_div, (None, fx.Int32(row_idx)))) page_id_i32 = as_mlir_value(fx.Int32(Vec(v, (1,), fx.Int32)[0])) llvm.StoreOp(page_id_i32, dst) @@ -4441,17 +4504,7 @@ def init_runtime_indices(self): self.stride_o_n_v = fx.Index(traits.DEFAULT_STRIDE_O_N) def init_causal_lpt_order(self): - """Issue causal q-blocks longest-first by reversing the q-block grid axis. - - Causal work per q-block grows with the block index and workgroups dispatch in - flattened-id order, so the natural order issues the heaviest block last and the - makespan carries its tail. Must run after init_thread_mapping and before - init_sequence_lengths / init_tile_bounds / init_q_row read q_start. - """ - traits = self.traits - num_q_blocks = (self.seq_len_v + traits.BLOCK_M - 1) // traits.BLOCK_M - self.q_block_idx = num_q_blocks - fx.Index(1) - self.q_block_idx - self.q_start = self.q_block_idx * traits.BLOCK_M + _init_causal_lpt_order(self) def init_lds(self, shared_storage): lds = fx.SharedAllocator().allocate(shared_storage).peek() diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 818ed3f1e8..755f80e6b1 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -7,6 +7,7 @@ import argparse import csv import hashlib +import inspect import logging import math import random @@ -35,7 +36,17 @@ from flydsl.runtime.device import get_rocm_arch # noqa: E402 from kernels.attention import flash_attn_interface # noqa: E402 from kernels.attention.flash_attn_gfx950 import build_flash_attn_dualwave_swp_module # noqa: E402 -from kernels.attention.flash_attn_interface import flydsl_flash_attn_func # noqa: E402 +from kernels.attention.flash_attn_interface import ( # noqa: E402 + _PAGED_TILE, + _build_dense_dualwave, + _build_paged, + _build_splitk, + _build_varlen, + _paged_kv_row_stride, + _paged_page_descriptor_fits, + _varlen_causal_lpt, + flydsl_flash_attn_func, +) from kernels.attention.flash_attn_utils import ( # noqa: E402 BIAS_MAX_DESCRIPTOR_BYTES, BIAS_MAX_OFFSET_ELEMS, @@ -4600,6 +4611,27 @@ def run(flag): assert torch.equal(off, on) +@_requires_gfx950 +@pytest.mark.parametrize("H", [8, 16]) +def test_causal_lpt_is_bit_identical(H): + """LPT only reorders q-blocks, so the output must not move a bit. + + This is also the guard against a reversal that drops or double-computes + q-blocks, which would otherwise read as a speedup. + """ + B, S, D = 1, 8192, 128 + torch.manual_seed(H) + q = _rand_lse(B, S, H, D, dtype=torch.bfloat16) + k, v = torch.randn_like(q), torch.randn_like(q) + + def run(flag): + return flydsl_flash_attn_func(q, k, v, causal=True, causal_lpt=flag).clone() + + off, on = run(False), run(True) + torch.cuda.synchronize() + assert torch.equal(off, on) + + @_requires_gfx950 @pytest.mark.parametrize("xcd_swizzle", [None, True]) def test_xcd_swizzle_heads_not_multiple_of_xcd(xcd_swizzle): @@ -4626,6 +4658,499 @@ def test_xcd_swizzle_heads_not_multiple_of_xcd(xcd_swizzle): torch.testing.assert_close(out.float(), ref, atol=_ATOL_BF16, rtol=0) +@_requires_gfx950 +def test_causal_lpt_ignored_when_not_causal(): + """Non-causal has no work gradient across q-blocks, so the flag is inert.""" + B, S, H, D = 1, 4096, 8, 128 + torch.manual_seed(3) + q = _rand_lse(B, S, H, D, dtype=torch.bfloat16) + k, v = torch.randn_like(q), torch.randn_like(q) + + off = flydsl_flash_attn_func(q, k, v, causal=False, causal_lpt=False).clone() + on = flydsl_flash_attn_func(q, k, v, causal=False, causal_lpt=True).clone() + torch.cuda.synchronize() + assert torch.equal(off, on) + + +@_requires_gfx950 +@pytest.mark.parametrize("num_kv_splits", [2, 3]) +def test_causal_lpt_splitk_off_on_bit_identical(num_kv_splits): + """Split-K reaches the dualwave builder by its own route, so pin the flag there too. + + Split-K puts the splits on grid z and leaves q-blocks on grid y, so reversing y + stays a permutation; this checks that it is one in practice, not just in argument. + """ + B, S, H, D = 1, 512, 8, 128 + torch.manual_seed(num_kv_splits) + q = _rand_lse(B, S, H, D, dtype=torch.bfloat16) + k, v = torch.randn_like(q), torch.randn_like(q) + + def run(flag): + return flydsl_flash_attn_func(q, k, v, causal=True, num_kv_splits=num_kv_splits, causal_lpt=flag).clone() + + off, on = run(False), run(True) + torch.cuda.synchronize() + assert torch.equal(off, on) + + +@_requires_gfx950 +@pytest.mark.parametrize("seqs", [[1024, 1, 512], [4096, 1]]) +def test_causal_lpt_varlen_off_on_matches_reference(seqs): + """Long varlen takes _build_varlen; both flag settings must match the reference. + + Both batches mix a long prefill with a single-token decode row, the shape whose + dispatch order the auto-gate decides. Forcing the flag must stay correct either + way, which is what this pins; the gate itself only picks which way is faster. + """ + dtype = torch.bfloat16 + D, H, Hkv = 128, 4, 4 + cu = torch.tensor([0, *np.cumsum(seqs)], dtype=torch.int32, device="cuda") + max_seqlen_q = max(seqs) + total = int(cu[-1].item()) + torch.manual_seed(total) + q = _rand_lse(total, H, D, dtype=dtype) + k = _rand_lse(total, Hkv, D, dtype=dtype) + v = _rand_lse(total, Hkv, D, dtype=dtype) + + def run(flag): + return flydsl_flash_attn_func( + q, + k, + v, + causal=True, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=max_seqlen_q, + cross_seqlen=False, + causal_lpt=flag, + ).clone() + + off, on = run(False), run(True) + torch.cuda.synchronize() + assert torch.equal(off, on) + for s0, s1 in zip(cu[:-1].tolist(), cu[1:].tolist()): + qb, kb, vb = (t[s0:s1].unsqueeze(0) for t in (q, k, v)) + ref = flydsl_flash_attn_func(qb, kb, vb, causal=True) + torch.cuda.synchronize() + torch.testing.assert_close(off[s0:s1].unsqueeze(0), ref, atol=_ATOL_BF16, rtol=0) + + +def test_varlen_causal_lpt_auto_gate(): + """None auto-selects on the batch shape; an explicit setting is never overridden. + + grid_dim.y is sized from the batch-wide max seqlen, so LPT only helps while one + sequence still spans at least twice as many q-blocks as the batch has sequences. + BLOCK_M is 256, so max_seqlen_q=8192 gives 32 q-blocks, and the gate flips at B=16. + The measured crossover is 0.75-0.875 of the q-block count, so the boundary is pinned + below it rather than at it: B=24 and B=32 both regress with LPT forced on. + """ + assert _varlen_causal_lpt(None, 1, 8192) is True # pure prefill, 2 <= 32 + assert _varlen_causal_lpt(None, 16, 8192) is True # exactly at the boundary, 1.035x + assert _varlen_causal_lpt(None, 17, 8192) is False # first B past it + assert _varlen_causal_lpt(None, 32, 8192) is False # measured 0.985x forced on + assert _varlen_causal_lpt(None, 129, 8192) is False # long prefill + 128 decodes + assert _varlen_causal_lpt(None, 4, 1) is False # decode-only batch + # An explicit request wins in both directions, so an A/B stays trustworthy. + assert _varlen_causal_lpt(True, 1024, 8192) is True + assert _varlen_causal_lpt(False, 1, 8192) is False + + +@_requires_gfx950 +def test_varlen_causal_lpt_auto_matches_forced(): + """Whichever way the gate resolves, the output must be the reference output. + + The gate picks a dispatch order, not a result, so auto must agree bit-for-bit with + both forced settings on a batch either side of the threshold. + """ + for seqs in ([4096, 4096], [4096] + [1] * 64): + cu = torch.tensor([0, *np.cumsum(seqs)], dtype=torch.int32, device="cuda") + total, max_q = int(cu[-1].item()), max(seqs) + torch.manual_seed(total) + q = _rand_lse(total, 8, 128, dtype=torch.bfloat16) + k, v = torch.randn_like(q), torch.randn_like(q) + + def run(flag): + return flydsl_flash_attn_func( + q, + k, + v, + causal=True, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=max_q, + cross_seqlen=False, + causal_lpt=flag, + ).clone() + + auto, off, on = run(None), run(False), run(True) + torch.cuda.synchronize() + assert torch.equal(auto, off) and torch.equal(auto, on) + + +def test_causal_lpt_is_threaded_through_every_dualwave_route(): + """Every cached dualwave builder must take the flag, or its route silently defaults on. + + _build_splitk and _build_varlen reach build_flash_attn_dualwave_swp_module directly, + so a missing parameter here is not a type error -- it just restores the builder + default and makes the public override a no-op on that route. + """ + for builder in (_build_dense_dualwave, _build_varlen, _build_splitk, _build_paged): + assert "causal_lpt" in inspect.signature(builder).parameters, builder.__name__ + + +@_requires_gfx950 +def test_causal_lpt_reaches_the_builder_through_the_call_site(): + """The signature test above pins the parameter; this pins the argument. + + Neither is redundant, and on its own neither catches the original defect. Both + builders default causal_lpt=True, so dropping the kwarg at a call site is not a + type error, and because LPT is a permutation the outputs still match -- every + off/on test in this file passes with the flag ignored. What does not survive it + is the cache key: two calls differing only in causal_lpt must produce two + distinct builder entries, and collapse into one if the call site drops the flag. + """ + dtype = torch.bfloat16 + # max_seqlen_q above _VARLEN_LIGHT_MAX_SEQ so varlen takes the long builder, and + # a dense seqlen past the split-K eligibility floor so num_kv_splits is honored. + s_varlen, s_dense = 1024, 512 + cu = torch.tensor([0, s_varlen], dtype=torch.int32, device="cuda") + qv = _rand_lse(s_varlen, 8, 128, dtype=dtype) + qd = _rand_lse(1, s_dense, 8, 128, dtype=dtype) + routes = ( + ( + _build_varlen, + (qv, torch.randn_like(qv), torch.randn_like(qv)), + dict(cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=s_varlen, cross_seqlen=False), + ), + (_build_splitk, (qd, torch.randn_like(qd), torch.randn_like(qd)), dict(num_kv_splits=2)), + ) + + for builder, (q, k, v), kwargs in routes: + builder.cache_clear() + flydsl_flash_attn_func(q, k, v, causal=True, causal_lpt=False, **kwargs) + assert builder.cache_info().currsize == 1, builder.__name__ + flydsl_flash_attn_func(q, k, v, causal=True, causal_lpt=True, **kwargs) + torch.cuda.synchronize() + assert builder.cache_info().currsize == 2, f"{builder.__name__} ignored causal_lpt" + + +def _paged_cache_from_dense(k, v, page_size, packed=False): + """Scatter dense [B, S, Hkv, D] K/V into a shuffled page cache. + + ``packed=True`` mimics the vLLM cache, where K and V are views into one page + entry and a KV row is therefore ``2 * Hkv * D`` elements wide. + """ + B, S, Hkv, D = k.shape + assert S % page_size == 0 + num_pages = B * (S // page_size) + block_table = torch.randperm(num_pages, device=k.device).view(B, -1).to(torch.int32) + if packed: + cache = torch.empty(num_pages, page_size, 2, Hkv, D, dtype=k.dtype, device=k.device) + k_cache, v_cache = cache[:, :, 0], cache[:, :, 1] + else: + k_cache = torch.empty(num_pages, page_size, Hkv, D, dtype=k.dtype, device=k.device) + v_cache = torch.empty_like(k_cache) + pages = block_table.reshape(-1).long() + k_cache[pages] = k.reshape(num_pages, page_size, Hkv, D) + v_cache[pages] = v.reshape(num_pages, page_size, Hkv, D) + return k_cache, v_cache, block_table + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size", [64, 128, 256]) +def test_paged_page_size_matches_dense(page_size): + """A page wider than the 64-token KV tile is walked tile by tile.""" + B, S, H, Hkv, D = 2, 1024, 8, 2, 128 + torch.manual_seed(page_size) + q = _rand_lse(B, S, H, D, dtype=torch.bfloat16) + k = _rand_lse(B, S, Hkv, D, dtype=torch.bfloat16) + v = torch.randn_like(k) + k_cache, v_cache, block_table = _paged_cache_from_dense(k, v, page_size) + seqlen_k = torch.full((B,), S, dtype=torch.int32, device=q.device) + + ref = flydsl_flash_attn_func(q, k, v, causal=True) + out = flydsl_flash_attn_func( + q, + k_cache, + v_cache, + causal=True, + block_table=block_table, + seqlen_k=seqlen_k, + max_seqlen_kv=S, + ) + torch.cuda.synchronize() + torch.testing.assert_close(out, ref, atol=2e-2, rtol=1e-2) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size", [64, 128, 256]) +def test_paged_strided_kv_matches_contiguous(page_size): + """K/V views into a packed cache only change addressing, never the values read.""" + B, S, H, Hkv, D = 2, 1024, 8, 2, 128 + torch.manual_seed(page_size) + q = _rand_lse(B, S, H, D, dtype=torch.bfloat16) + k = _rand_lse(B, S, Hkv, D, dtype=torch.bfloat16) + v = torch.randn_like(k) + seqlen_k = torch.full((B,), S, dtype=torch.int32, device=q.device) + + def run(packed): + torch.manual_seed(0) # same block table for both layouts + k_cache, v_cache, block_table = _paged_cache_from_dense(k, v, page_size, packed=packed) + assert k_cache.is_contiguous() == (not packed) + return flydsl_flash_attn_func( + q, + k_cache, + v_cache, + causal=True, + block_table=block_table, + seqlen_k=seqlen_k, + max_seqlen_kv=S, + ).clone() + + dense, packed = run(False), run(True) + torch.cuda.synchronize() + assert torch.equal(dense, packed) + + +@_requires_gfx950 +@pytest.mark.parametrize("page_size", [64, 128, 256]) +def test_paged_vllm_cache_view(page_size): + """vLLM's cache is [NumBlocks, Hkv, PageSize, 2*D] split into two K/V views. + + Its head stride is PageSize*2*D rather than D, which the kernel only reads + with more than one KV head, so the single-head case must run in place. + """ + B, S, H, D = 2, 1024, 8, 128 + torch.manual_seed(page_size) + q = _rand_lse(B, S, H, D, dtype=torch.bfloat16) + k = _rand_lse(B, S, 1, D, dtype=torch.bfloat16) + v = torch.randn_like(k) + + num_pages = B * (S // page_size) + cache = torch.empty(num_pages, 1, page_size, 2 * D, dtype=k.dtype, device=k.device) + k_cache, v_cache = cache.transpose(1, 2).split(D, dim=-1) + block_table = torch.randperm(num_pages, device=k.device).view(B, -1).to(torch.int32) + pages = block_table.reshape(-1).long() + k_cache[pages] = k.reshape(num_pages, page_size, 1, D) + v_cache[pages] = v.reshape(num_pages, page_size, 1, D) + + ref = flydsl_flash_attn_func(q, k, v, causal=True) + out = flydsl_flash_attn_func( + q, + k_cache, + v_cache, + causal=True, + block_table=block_table, + seqlen_k=torch.full((B,), S, dtype=torch.int32, device=q.device), + max_seqlen_kv=S, + ) + torch.cuda.synchronize() + torch.testing.assert_close(out, ref, atol=2e-2, rtol=1e-2) + + +@_requires_gfx950 +def test_paged_rejected_pitch_still_computes_the_right_answer(monkeypatch): + """A rejected pitch must fall back to a correct copy, not merely be refused. + + _paged_kv_row_stride is unit-tested on stride tuples, and the pitches that trip its + bound span gigabytes, so the bound itself cannot be reached with a real cache here. + What is left untested by that is the consequence: returning None reports the cache as + dense, and the caller is then expected to copy it into a layout the descriptor can + address. This forces the rejection on a cache whose pitch is genuinely non-dense and + checks the output still matches, which fails if that path drops the strides instead of + copying. + """ + B, S, H, D, page_size = 2, 1024, 8, 128, 128 + torch.manual_seed(0) + q = _rand_lse(B, S, H, D, dtype=torch.bfloat16) + k = _rand_lse(B, S, 1, D, dtype=torch.bfloat16) + v = torch.randn_like(k) + + num_pages = B * (S // page_size) + cache = torch.empty(num_pages, 1, page_size, 2 * D, dtype=k.dtype, device=k.device) + k_cache, v_cache = cache.transpose(1, 2).split(D, dim=-1) + block_table = torch.randperm(num_pages, device=k.device).view(B, -1).to(torch.int32) + pages = block_table.reshape(-1).long() + k_cache[pages] = k.reshape(num_pages, page_size, 1, D) + v_cache[pages] = v.reshape(num_pages, page_size, 1, D) + # Guard the premise: this cache is only interesting because its pitch is accepted. + assert _paged_kv_row_stride(k_cache, v_cache, page_size, 1, D) == 2 * D + + paged_kw = dict( + causal=True, + block_table=block_table, + seqlen_k=torch.full((B,), S, dtype=torch.int32, device=q.device), + max_seqlen_kv=S, + ) + ref = flydsl_flash_attn_func(q, k, v, causal=True) + accepted = flydsl_flash_attn_func(q, k_cache, v_cache, **paged_kw) + + monkeypatch.setattr("kernels.attention.flash_attn_interface._paged_kv_row_stride", lambda *a, **kw: None) + rejected = flydsl_flash_attn_func(q, k_cache, v_cache, **paged_kw) + torch.cuda.synchronize() + + torch.testing.assert_close(rejected, ref, atol=2e-2, rtol=1e-2) + assert torch.equal(rejected, accepted), "the copy path disagrees with the in-place path" + + +class _FakeKV: + """Stride-only stand-in: the pitch gate reads nothing else off the tensor. + + The pitches under test span gigabytes, so a real allocation is out of the question. + """ + + def __init__(self, strides, elem_size=2): + self._strides = strides + self._elem_size = elem_size + + def stride(self): + return self._strides + + def element_size(self): + return self._elem_size + + +@pytest.mark.parametrize( + "row_stride, accepted", + [ + (256, True), # ordinary packed-cache pitch + (2**24 - 1, True), # widest page that still fits both descriptor fields + (2**24, False), # page_elems == 2**31, overflows the signed-int32 extent + (2**25, False), # comfortably past both bounds + ], +) +def test_paged_kv_row_stride_bounds_pitch(row_stride, accepted): + """A pitch whose page overruns the descriptor must be reported dense, not passed through. + + page_elems = page_size * row_stride is emitted as a signed-int32 extent and a + 32-bit byte span; at page_size 128 that caps row_stride at 2**24 - 1. Beyond it the + descriptor wraps and the kernel reads the wrong rows, so the gate returns None and + the caller copies instead. + """ + page_size, num_kv_heads, head_dim = 128, 1, 128 + strides = (page_size * row_stride, row_stride, head_dim, 1) + kv = _FakeKV(strides) + got = _paged_kv_row_stride(kv, kv, page_size, num_kv_heads, head_dim) + assert got == (row_stride if accepted else None) + + +def test_paged_kv_row_stride_dense_reports_none(): + """A densely packed cache has no runtime pitch to pass; the kernel uses its default.""" + page_size, num_kv_heads, head_dim = 128, 2, 128 + dense_row = num_kv_heads * head_dim + kv = _FakeKV((page_size * dense_row, dense_row, head_dim, 1)) + assert _paged_kv_row_stride(kv, kv, page_size, num_kv_heads, head_dim) is None + + +@pytest.mark.parametrize( + "page_size, num_kv_heads, head_dim, elem_size, fits", + [ + (128, 1, 128, 2, True), # ordinary vLLM page + (256, 8, 128, 2, True), # a 4-tile page with 8 KV heads is still tiny + (2**24 - 64, 1, 128, 2, True), # widest addressable page at Hkv=1, D=128 + (2**24, 1, 128, 2, False), # exactly 2**31 elements, wraps the int32 extent + (2**21, 8, 128, 2, False), # more KV heads reach the bound at a smaller page + (2**23, 1, 128, 4, False), # 2**30 elements fits, but 2**32 bytes does not + ], +) +def test_paged_page_descriptor_fits_is_exclusive_on_both_bounds(page_size, num_kv_heads, head_dim, elem_size, fits): + """The shared predicate's arithmetic, at the boundary from both directions. + + This pins the numbers only. That the dense pitch is actually put through it on the + launch path is test_paged_oversized_dense_page_is_rejected. + + A dense cache returns from the row-stride gate before its page-span check, so once + page_size became any multiple of the tile, page_size * Hkv * D was unbounded on that + path. It cannot fall back to a copy the way an overwide strided pitch does, since + dense is already the narrowest layout, so the entry point rejects it instead. The + fourth case is the boundary; the last shows the byte bound binding before the element + bound once an element is wider than two bytes. + """ + assert _paged_page_descriptor_fits(page_size * num_kv_heads * head_dim, elem_size) is fits + + +@_requires_gfx950 +def test_paged_oversized_dense_page_is_rejected(): + """An oversized dense page must be refused at the entry point rather than wrapped. + + A page wide enough to overrun the descriptor spans gigabytes, but nothing before the + check reads the pages themselves, so a stride-0 expand buys a real cache shape at no + allocation cost. That keeps the arguments the check is called with under test, which + a stubbed predicate would not. + + The second case is the same wrap reached through the head count rather than the page: + the kernel strides the dense case by the resolved num_kv_heads, so a cache whose own + page fits can still be oversized once the caller asks for more KV heads than it holds. + """ + D = 128 + q = _rand_lse(1, 256, 8, D, dtype=torch.bfloat16) + one = torch.empty(1, 1, 1, D, dtype=torch.bfloat16, device=q.device) + paged_kw = dict( + causal=True, + block_table=torch.zeros(1, 1, dtype=torch.int32, device=q.device), + seqlen_k=torch.full((1,), _PAGED_TILE, dtype=torch.int32, device=q.device), + max_seqlen_kv=_PAGED_TILE, + ) + + # 2**24 * 1 * 128 is exactly 2**31 elements, the first page that does not fit. + too_wide = one.expand(1, 2**24, 1, D) + with pytest.raises(NotImplementedError, match="page descriptor"): + flydsl_flash_attn_func(q, too_wide, too_wide, **paged_kw) + + # 2**21 * 1 * 128 fits; 2**21 * 8 * 128 is again exactly 2**31. + fits_alone = one.expand(1, 2**21, 1, D) + with pytest.raises(NotImplementedError, match="page descriptor"): + flydsl_flash_attn_func(q, fits_alone, fits_alone, num_kv_heads=8, **paged_kw) + + +@_requires_gfx950 +def test_paged_rejects_unsupported_page_size(): + B, S, H, Hkv, D = 1, 512, 8, 2, 128 + q = _rand_lse(B, S, H, D, dtype=torch.bfloat16) + k = _rand_lse(B, S, Hkv, D, dtype=torch.bfloat16) + v = torch.randn_like(k) + k_cache, v_cache, block_table = _paged_cache_from_dense(k, v, 32) + with pytest.raises(NotImplementedError): + flydsl_flash_attn_func( + q, + k_cache, + v_cache, + causal=True, + block_table=block_table, + seqlen_k=torch.full((B,), S, dtype=torch.int32, device=q.device), + max_seqlen_kv=S, + ) + + +@_requires_gfx950 +def test_paged_vectorized_rejects_multi_tile_page(): + """The vectorized layout keeps one tile per page in its address arithmetic. + + The linear layout walks a wide page tile by tile, but the vectorized 5D form still + folds the page into the descriptor, so a multi-tile page there reads the wrong rows + without failing. It has to be refused at the boundary instead. + """ + B, S, H, Hkv, D = 1, 512, 8, 2, 128 + page_size = 128 # two 64-token KV tiles + kvs = 8 # 16 bytes / bf16 + num_pages = B * (S // page_size) + q = _rand_lse(B, S, H, D, dtype=torch.bfloat16) + cache = torch.zeros(num_pages, Hkv, D // kvs, page_size, kvs, dtype=torch.bfloat16, device=q.device) + block_table = torch.arange(num_pages, dtype=torch.int32, device=q.device).view(B, -1) + + with pytest.raises(NotImplementedError, match="vectorized"): + flydsl_flash_attn_func( + q, + cache, + cache.clone(), + causal=True, + block_table=block_table, + seqlen_k=torch.full((B,), S, dtype=torch.int32, device=q.device), + max_seqlen_kv=S, + kv_cache_layout="vectorized", + ) + + if __name__ == "__main__": main()