diff --git a/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_delta_attn_utils.py b/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_delta_attn_utils.py index 81ebc93032c..8dff47880f0 100644 --- a/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_delta_attn_utils.py +++ b/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_delta_attn_utils.py @@ -185,16 +185,35 @@ def softplus(x): return tl.where(x < 20.0, tl.log(1.0 + tl.exp(x)), x) -def input_guard(fn: Callable) -> Callable: - """Ensure all tensor arguments are contiguous before kernel launch.""" - - @functools.wraps(fn) - def wrapper(*args: Any, **kwargs: Any) -> Any: - args = tuple(a.contiguous() if isinstance(a, torch.Tensor) else a for a in args) - kwargs = { - k: v.contiguous() if isinstance(v, torch.Tensor) else v - for k, v in kwargs.items() - } - return fn(*args, **kwargs) - - return wrapper +def input_guard(fn: Callable | None = None, *, skip: tuple[str, ...] = ()) -> Callable: + """Ensure tensor arguments are contiguous before kernel launch. + + ``skip`` names keyword arguments that keep their original storage. The + paged ``state_cache`` is one: each slot's ``[H, V, K]`` plane is dense, + but ``stride(0)`` may be padded, and packing the pool would copy it and + write a clone instead of the live cache. Those tensors must be indexed + with their own strides. + """ + skip_keys = frozenset(skip) + + def decorator(inner: Callable) -> Callable: + @functools.wraps(inner) + def wrapper(*args: Any, **kwargs: Any) -> Any: + args = tuple( + a.contiguous() if isinstance(a, torch.Tensor) else a for a in args + ) + kwargs = { + k: ( + v + if k in skip_keys or not isinstance(v, torch.Tensor) + else v.contiguous() + ) + for k, v in kwargs.items() + } + return inner(*args, **kwargs) + + return wrapper + + if fn is not None: + return decorator(fn) + return decorator diff --git a/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_fwd.py b/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_fwd.py index cec33f07d0a..5431138d52c 100644 --- a/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_fwd.py +++ b/aiter/ops/triton/_triton_kernels/chunk_delta_attn/chunk_fwd.py @@ -68,6 +68,10 @@ def chunk_delta_attn_fwd( use_qk_l2norm_in_kernel: bool = False, use_beta_sigmoid_in_kernel: bool = False, state_v_first: bool = False, + out: torch.Tensor | None = None, + state_cache: torch.Tensor | None = None, + state_indices: torch.Tensor | None = None, + has_initial_state: torch.Tensor | None = None, ) -> tuple: """ Forward pass for chunk_delta_attn. @@ -103,6 +107,9 @@ def chunk_delta_attn_fwd( state_v_first: Store the recurrent state V-first (``[V, K]``) instead of the default ``[K, V]``. Matches fla's option of the same name. + out: Optional contiguous output buffer, same shape as + ``v``. FlashKDA-only, like ``state_cache``. + state_cache: Optional paged fp32 V-first cache. FlashKDA-only. Returns: (o, final_state, g_cumsum, Aqk, Akk, w, u, qg, kg) @@ -151,6 +158,17 @@ def chunk_delta_attn_fwd( if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + if state_cache is not None and not use_flash_kda: + raise ValueError( + "paged state_cache is only implemented on the FlashKDA path; " + "gather into initial_state, or pass a call that flash_kda_supported accepts" + ) + if out is not None and not use_flash_kda: + raise ValueError( + "out is only implemented on the FlashKDA path; use the returned " + "tensor, or pass a call that flash_kda_supported accepts" + ) + if use_flash_kda: o, final_state = flash_kda_fwd( q=q, @@ -167,6 +185,10 @@ def chunk_delta_attn_fwd( state_v_first=state_v_first, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + out=out, + state_cache=state_cache, + state_indices=state_indices, + has_initial_state=has_initial_state, ) return o, final_state, None, None, None, None, None, None, None diff --git a/aiter/ops/triton/_triton_kernels/chunk_delta_attn/flash_kda.py b/aiter/ops/triton/_triton_kernels/chunk_delta_attn/flash_kda.py index 366a9f2a949..cedd5a26b83 100644 --- a/aiter/ops/triton/_triton_kernels/chunk_delta_attn/flash_kda.py +++ b/aiter/ops/triton/_triton_kernels/chunk_delta_attn/flash_kda.py @@ -44,6 +44,8 @@ intra kernel, so it answers that request by a different route and has no unpivoted mode to offer a caller who declined it. * Forward only, and no intermediates for a backward pass. + * Paged ``state_cache`` is FlashKDA-only, V-first, and is not packed: + each slot is dense ``[H, V, K]`` while ``stride(0)`` may be padded. """ import functools @@ -387,9 +389,26 @@ def _seg_occupancy_class(num_segs: int) -> int: # NUM_SEGS_CLASS is there for the same reason -- without it the segmented # and unsegmented schedules collide and whichever runs first sets the # config for both, a result the on-disk cache then keeps. - key=["H", "K", "W", "C", "HAS_V", "COMPUTE_OUTPUT", "NUM_SEGS_CLASS"], + key=[ + "H", + "K", + "W", + "C", + "HAS_V", + "COMPUTE_OUTPUT", + "NUM_SEGS_CLASS", + "PAGED_CACHE", + ], **autotune_cache_kwargs, ) +@triton.heuristics( + { + "PAGED_CACHE": lambda args: args["state_cache"] is not None, + "PAGED_H_IN": lambda args: ( + args["state_cache"] is not None and args["h_in"] is None + ), + } +) @triton.jit def _flash_kda_segment_kernel( ws_kd, @@ -425,6 +444,12 @@ def _flash_kda_segment_kernel( STORE_FINAL: tl.constexpr, STATE_V_FIRST: tl.constexpr, CM_OUT: tl.constexpr = "", + state_cache=None, + state_indices=None, + has_initial_state=None, + cache_stride=0, + PAGED_CACHE: tl.constexpr = False, + PAGED_H_IN: tl.constexpr = False, ): """Delta-rule recurrence over one segment of chunks. @@ -462,11 +487,30 @@ def _flash_kda_segment_kernel( b_h1 = tl.where(o_k1[:, None] == o_w[None, :], 1.0, 0.0) b_h2 = tl.where(o_k2[:, None] == o_w[None, :], 1.0, 0.0) elif HAS_H_IN: - s_off = (i_seg * H + i_h) * K * W + o_w[None, :] - b_h1 = tl.load(h_in + s_off + o_k1[:, None] * W, mask=m_w[None, :], other=0.0) - b_h2 = tl.load(h_in + s_off + o_k2[:, None] * W, mask=m_w[None, :], other=0.0) - b_h1 = b_h1.to(tl.float32) - b_h2 = b_h2.to(tl.float32) + if PAGED_H_IN: + i_n = tl.load(seg_seq + i_seg).to(tl.int64) + slot = tl.load(state_indices + i_n).to(tl.int64) + # A zero-start sequence's slot may hold a previous owner's state: + # mask the load so it is never read, not read and then discarded. + take = tl.load(has_initial_state + i_n).to(tl.int1) + m_h = m_w[None, :] & take + f_off = slot * cache_stride + i_h * V * K + o_w[None, :] * K + b_h1 = tl.load(state_cache + f_off + o_k1[:, None], mask=m_h, other=0.0).to( + tl.float32 + ) + b_h2 = tl.load(state_cache + f_off + o_k2[:, None], mask=m_h, other=0.0).to( + tl.float32 + ) + else: + s_off = (i_seg * H + i_h) * K * W + o_w[None, :] + b_h1 = tl.load( + h_in + s_off + o_k1[:, None] * W, mask=m_w[None, :], other=0.0 + ) + b_h2 = tl.load( + h_in + s_off + o_k2[:, None] * W, mask=m_w[None, :], other=0.0 + ) + b_h1 = b_h1.to(tl.float32) + b_h2 = b_h2.to(tl.float32) else: b_h1 = tl.zeros([64, BW], dtype=tl.float32) b_h2 = tl.zeros([64, BW], dtype=tl.float32) @@ -540,8 +584,21 @@ def _flash_kda_segment_kernel( if STORE_FINAL: # noqa: SIM102 if tl.load(seg_is_last + i_seg) == 1: i_n = tl.load(seg_seq + i_seg).to(tl.int64) - dt_s = final_state.dtype.element_ty - if STATE_V_FIRST: + dt_s = tl.float32 if PAGED_CACHE else final_state.dtype.element_ty + if PAGED_CACHE: + slot = tl.load(state_indices + i_n).to(tl.int64) + f_off = slot * cache_stride + i_h * V * K + o_w[None, :] * K + tl.store( + state_cache + f_off + o_k1[:, None], + b_h1.to(dt_s), + mask=m_w[None, :], + ) + tl.store( + state_cache + f_off + o_k2[:, None], + b_h2.to(dt_s), + mask=m_w[None, :], + ) + elif STATE_V_FIRST: f_off = (i_n * H + i_h) * V * K + o_w[None, :] * K tl.store( final_state + f_off + o_k1[:, None], @@ -567,7 +624,14 @@ def _flash_kda_segment_kernel( ) -@triton.heuristics({"HAS_H0": lambda args: args["h0"] is not None}) +@triton.heuristics( + { + "HAS_H0": lambda args: ( + args["h0"] is not None or args["state_cache"] is not None + ), + "PAGED_CACHE": lambda args: args["state_cache"] is not None, + } +) @triton.jit def _flash_kda_seg_scan_kernel( A_seg, @@ -580,6 +644,11 @@ def _flash_kda_seg_scan_kernel( V: tl.constexpr, BV: tl.constexpr, HAS_H0: tl.constexpr, + state_cache=None, + state_indices=None, + has_initial_state=None, + cache_stride=0, + PAGED_CACHE: tl.constexpr = False, ): """Serial scan across a sequence's segments: ``h <- A_seg h + b_seg``. @@ -599,13 +668,25 @@ def _flash_kda_seg_scan_kernel( m_v = o_v < V if HAS_H0: - base = (i_n * H + i_h) * K * V + o_v[None, :] - b_h1 = tl.load(h0 + base + o_k1[:, None] * V, mask=m_v[None, :], other=0.0).to( - tl.float32 - ) - b_h2 = tl.load(h0 + base + o_k2[:, None] * V, mask=m_v[None, :], other=0.0).to( - tl.float32 - ) + if PAGED_CACHE: + slot = tl.load(state_indices + i_n).to(tl.int64) + take = tl.load(has_initial_state + i_n).to(tl.int1) + m_h = m_v[None, :] & take + f_off = slot * cache_stride + i_h * V * K + o_v[None, :] * K + b_h1 = tl.load(state_cache + f_off + o_k1[:, None], mask=m_h, other=0.0).to( + tl.float32 + ) + b_h2 = tl.load(state_cache + f_off + o_k2[:, None], mask=m_h, other=0.0).to( + tl.float32 + ) + else: + base = (i_n * H + i_h) * K * V + o_v[None, :] + b_h1 = tl.load( + h0 + base + o_k1[:, None] * V, mask=m_v[None, :], other=0.0 + ).to(tl.float32) + b_h2 = tl.load( + h0 + base + o_k2[:, None] * V, mask=m_v[None, :], other=0.0 + ).to(tl.float32) else: b_h1 = tl.zeros([64, BV], dtype=tl.float32) b_h2 = tl.zeros([64, BV], dtype=tl.float32) @@ -849,7 +930,40 @@ def _build_segments( _seg_scan_fast = fast_launch(_flash_kda_seg_scan_kernel) -@input_guard +def _check_paged_state_cache( + state_cache: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + *, + N: int, + H: int, + V: int, + K: int, +) -> None: + """Paged layout: each slot is dense ``[H, V, K]``; ``stride(0)`` may pad.""" + if state_cache.dtype != torch.float32 or state_cache.dim() != 4: + raise ValueError("state_cache must be an fp32 [slots, H, V, K] tensor") + if tuple(state_cache.shape[1:]) != (H, V, K): + raise ValueError( + f"state_cache slot shape {tuple(state_cache.shape[1:])} != {(H, V, K)}" + ) + # The kernel addresses slot * stride(0), so a smaller outer stride (an + # expanded view is 0) would make neighbouring slots share storage. + if ( + state_cache.stride(-1) != 1 + or state_cache.stride(2) != K + or state_cache.stride(1) != V * K + or state_cache.stride(0) < H * V * K + ): + raise ValueError( + f"state_cache strides {tuple(state_cache.stride())} are not dense " + f"[H, V, K] slots; stride(0) may be padded but not below {H * V * K}" + ) + if state_indices.numel() != N or has_initial_state.numel() != N: + raise ValueError(f"state_indices and has_initial_state must have N={N} entries") + + +@input_guard(skip=("state_cache", "out")) def flash_kda_fwd( q: torch.Tensor, k: torch.Tensor, @@ -866,6 +980,10 @@ def flash_kda_fwd( cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.Tensor | None = None, chunks_per_seg: int | None = None, + out: torch.Tensor | None = None, + state_cache: torch.Tensor | None = None, + state_indices: torch.Tensor | None = None, + has_initial_state: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Two-kernel KDA forward, with an optional segmented (context parallel) scan. @@ -889,15 +1007,52 @@ def flash_kda_fwd( when omitted. chunks_per_seg: Segment length for the parallel scan. None picks a value from the head count; <= 0 disables it. + out: Optional output buffer, same shape as ``v``, and + contiguous: it is written in place rather than + packed by ``input_guard``. + state_cache: Optional paged fp32 V-first cache ``[slots, H, V, K]``. + Each slot's ``[H, V, K]`` plane must be dense; + ``stride(0)`` may be padded. This tensor is not + packed by ``input_guard``. + When set, the walk reads sequence ``n`` from row + ``state_indices[n]`` if ``has_initial_state[n]`` and + writes the final state back to the same row. + state_indices: Int32 ``[N]`` cache row per sequence, each in + ``[0, slots)``. Not range-checked: that would + need a host sync per call. + has_initial_state: Bool ``[N]``; false starts that sequence from zero. Returns: - ``(o, final_state)`` with ``o`` shaped like ``v``. + ``(o, final_state)`` with ``o`` shaped like ``v``. ``final_state`` is + ``None`` when ``state_cache`` is given, because the cache already holds + it. """ B, T, H, K = q.shape V = v.shape[-1] C = FLASH_KDA_CHUNK inv_block = min(FLASH_KDA_INV_BLOCK, C) dev = q.device + paged = state_cache is not None + if paged: + if initial_state is not None: + raise ValueError("state_cache replaces initial_state") + if state_indices is None or has_initial_state is None: + raise ValueError("state_cache needs state_indices and has_initial_state") + if not state_v_first: + raise ValueError("paged state_cache is V-first") + if out is None: + raise ValueError("paged state_cache needs an explicit out buffer") + if out is not None: + if out.shape != v.shape: + raise ValueError( + f"out shape {tuple(out.shape)} does not match v {tuple(v.shape)}" + ) + # K2 stores through packed [B, T, H, V] strides, and `out` is exempt + # from `input_guard` so the caller's buffer is what gets written. A + # strided view would need either its own strides in the kernel or a + # copy afterwards, and the point of taking `out` is to avoid the copy. + if not out.is_contiguous(): + raise ValueError("out must be contiguous") if cu_seqlens is not None: if B != 1: @@ -912,6 +1067,17 @@ def flash_kda_fwd( NT = triton.cdiv(T, C) total_tiles = B * NT + if paged: + _check_paged_state_cache( + state_cache, + state_indices, + has_initial_state, + N=N, + H=H, + V=V, + K=K, + ) + if cu_seqlens is not None: seqs = _seq_bounds(cu_seqlens) else: @@ -1024,13 +1190,16 @@ def flash_kda_fwd( h0 = h0.transpose(-1, -2) h0 = h0.contiguous() - o = torch.empty_like(v) + o = out if out is not None else torch.empty_like(v) final_state = None - if output_final_state: + store_final = output_final_state or paged + if output_final_state and not paged: shape = (N, H, V, K) if state_v_first else (N, H, K, V) state_dtype = h0.dtype if h0 is not None else torch.float32 final_state = torch.empty(shape, dtype=state_dtype, device=dev) + cache_stride = int(state_cache.stride(0)) if paged else 0 + common = { "ws_kd": ws_kd, "ws_qd": ws_qd, @@ -1053,15 +1222,29 @@ def flash_kda_fwd( "STATE_V_FIRST": state_v_first, "CM_OUT": CM_OUT_STORE, } + paged_args = { + "state_cache": state_cache, + "state_indices": None if not paged else state_indices.to(torch.int32), + "has_initial_state": has_initial_state, + "cache_stride": cache_stride, + } # Only pass A can go to Gluon, so an unsegmented shape reaches none of it. use_gluon_k2 = AITER_FDA_USE_GLUON_K2 and _gluon_k2_usable(C, K, V) and max_segs > 1 _log_route(use_gluon_k1, use_gluon_k2, C, K, V) def _launch_k2(*, W, **kw): + launch_args = { + **common, + "state_cache": None, + "state_indices": None, + "has_initial_state": None, + "cache_stride": 0, + **kw, + } return _segment_fast[ lambda meta, _w=W: (triton.cdiv(_w, meta["BW"]), num_segs * H) - ](W=W, **common, **kw) + ](W=W, **launch_args) if max_segs > 1: # Pass A: b_seg (affine part) and A_seg (linear part), both fully @@ -1129,10 +1312,11 @@ def _launch_k2(*, W, **kw): K=K, V=V, BV=BV_SCAN, + **paged_args, num_warps=SCAN_WARPS, ) else: - h_in = h0 + h_in = None if paged else h0 # Pass C: re-run each segment from its true incoming state, writing outputs. _launch_k2( @@ -1143,11 +1327,12 @@ def _launch_k2(*, W, **kw): final_state=final_state, W=V, INIT_IDENTITY=False, - HAS_H_IN=h_in is not None, + HAS_H_IN=h_in is not None or paged, HAS_V=True, COMPUTE_OUTPUT=True, STORE_H_OUT=False, - STORE_FINAL=output_final_state, + STORE_FINAL=store_final, + **(paged_args if paged else {}), ) return o, final_state diff --git a/aiter/ops/triton/kimi_delta_attn/chunk_delta_attn.py b/aiter/ops/triton/kimi_delta_attn/chunk_delta_attn.py index 245855a0964..91ec5e2dde4 100644 --- a/aiter/ops/triton/kimi_delta_attn/chunk_delta_attn.py +++ b/aiter/ops/triton/kimi_delta_attn/chunk_delta_attn.py @@ -61,6 +61,10 @@ def chunk_kimi_delta_attn( disable_recompute: bool = False, chunk_size: int | None = None, cu_seqlens: torch.LongTensor | None = None, + out: torch.Tensor | None = None, + state_cache: torch.Tensor | None = None, + state_indices: torch.Tensor | None = None, + has_initial_state: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: r""" Chunked Kimi Delta Attention forward pass using Triton (Forward only). @@ -140,12 +144,30 @@ def chunk_kimi_delta_attn( cu_seqlens (torch.LongTensor, optional): Cumulative sequence lengths of shape `[N+1]` for variable-length inputs, consistent with the FlashAttention API. Default: `None`. + out (torch.Tensor, optional): + Caller-owned output buffer, same shape as `v`, and contiguous. + Written in place rather than allocated and copied. Required with + `state_cache`. FlashKDA-only; the default pipeline rejects this + argument rather than ignoring it. + state_cache (torch.Tensor, optional): + Paged fp32 V-first cache `[slots, H, V, K]`. Replaces + `initial_state`: sequence `n` is row `state_indices[n]`. Each + slot's `[H, V, K]` plane must be dense; `stride(0)` may be padded. + FlashKDA-only; the default pipeline rejects this argument rather + than ignoring it. + state_indices (torch.Tensor, optional): + Int32 `[N]` cache row per sequence, each in `[0, slots)`. Required + with `state_cache`. Trusted as given, not range-checked. + has_initial_state (torch.Tensor, optional): + Bool `[N]`. False starts that sequence from zero. Required with + `state_cache`. Returns: tuple[torch.Tensor, torch.Tensor | None]: - o (torch.Tensor): Outputs of shape `[B, T, HV, V]`. - final_state (torch.Tensor | None): Final state if - `output_final_state=True` else `None`. + `output_final_state=True` else `None`. `None` when + `state_cache` is given, because the cache already holds it. Examples: >>> import torch @@ -247,6 +269,12 @@ def chunk_kimi_delta_attn( initial_state = initial_state.contiguous() if cu_seqlens is not None: cu_seqlens = cu_seqlens.contiguous() + if state_indices is not None: + state_indices = state_indices.contiguous() + if has_initial_state is not None: + has_initial_state = has_initial_state.contiguous() + # state_cache is not packed: stride(0) may be padded. Packing it would + # copy the pool and write a clone. o, final_state, *_ = chunk_delta_attn_fwd( q=q, @@ -270,5 +298,9 @@ def chunk_kimi_delta_attn( use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, state_v_first=state_v_first, + out=out, + state_cache=state_cache, + state_indices=state_indices, + has_initial_state=has_initial_state, ) return o.to(q.dtype), final_state diff --git a/op_tests/triton_tests/chunk_delta_attn/test_flash_kda.py b/op_tests/triton_tests/chunk_delta_attn/test_flash_kda.py index 42ac1e5d7d7..e84c75668c8 100644 --- a/op_tests/triton_tests/chunk_delta_attn/test_flash_kda.py +++ b/op_tests/triton_tests/chunk_delta_attn/test_flash_kda.py @@ -778,3 +778,235 @@ def counting(**kw): o_64, _ = chunk_kimi_delta_attn(chunk_size=64, **{**kwargs, "safe_gate": False}) assert len(calls) == 1 assert torch.equal(o_auto, o_64), "an ineligible call should have resolved to 64" + + +def _paged_pool(n, H, initial, pad_floats=0): + """V-first paged cache. ``pad_floats`` extra between slots.""" + V = K = K_DIM + inner = H * V * K + slot_stride = inner + pad_floats + storage = torch.zeros((n + 2) * slot_stride, device=device, dtype=torch.float32) + cache = torch.as_strided( + storage, + size=(n + 2, H, V, K), + stride=(slot_stride, V * K, K, 1), + ) + indices = torch.arange(n, 0, -1, device=device, dtype=torch.int32) + cache[indices] = initial + return cache, indices, storage + + +def _kimi_kwargs(q, k, v, g, beta, A_log, dt_bias, scale, **kw): + return { + "q": q, + "k": k, + "v": v, + "g": g, + "beta": beta, + "A_log": A_log, + "dt_bias": dt_bias, + "scale": scale, + "use_qk_l2norm_in_kernel": True, + "use_gate_in_kernel": True, + "use_beta_sigmoid_in_kernel": True, + "safe_gate": True, + "lower_bound": LOWER_BOUND, + "state_v_first": True, + **kw, + } + + +@pytest.mark.parametrize("T,chunks_per_seg", [(256, 0), (1024, 4)]) +def test_paged_cache_matches_dense(T, chunks_per_seg): + """In-kernel paged I/O matches gather into a dense V-first state.""" + H = 4 + q, k, v, g, beta, A_log, dt_bias, scale = make_inputs(1, T, H) + h0 = torch.randn(1, H, K_DIM, K_DIM, device=device, dtype=torch.float32) * 0.1 + common = { + "q": q, + "k": k, + "v": v, + "g": g, + "beta": beta, + "A_log": A_log, + "dt_bias": dt_bias, + "scale": scale, + "lower_bound": LOWER_BOUND, + "state_v_first": True, + "chunks_per_seg": chunks_per_seg, + } + o_dense, ht = flash_kda_fwd(**common, initial_state=h0, output_final_state=True) + cache, indices, _storage = _paged_pool(1, H, h0) + ptr_before = cache.data_ptr() + out = torch.empty_like(v) + paged_kw = { + **common, + "initial_state": None, + "output_final_state": False, + "out": out, + "state_cache": cache, + "state_indices": indices, + "has_initial_state": torch.ones(1, device=device, dtype=torch.bool), + } + # First paged launch autotunes a new PAGED_CACHE specialization; compare + # after that, not against a trial config. + flash_kda_fwd(**paged_kw) + cache[indices] = h0 + o_paged, ht_paged = flash_kda_fwd(**paged_kw) + assert ht_paged is None + assert o_paged.data_ptr() == out.data_ptr() + assert cache.data_ptr() == ptr_before, "paged cache must not be packed/copied" + assert torch.equal(o_paged, o_dense) + assert torch.equal(cache[indices], ht) + assert torch.count_nonzero(cache[[0, -1]]).item() == 0 + + +@pytest.mark.parametrize("seg", ["0", "4"]) +def test_paged_cache_has_initial_state_false(seg, monkeypatch): + """A zero-start slot is never read: its NaN contents must not leak. + + ``seg="4"`` cuts the 8-chunk sequence into two segments, so the load under + test is the segment kernel's incoming-state load as well as the scan's. + """ + monkeypatch.setenv("CHUNK_DELTA_ATTN_FLASH_KDA_SEG", seg) + H = 4 + args = make_inputs(1, 256, H) + v = args[2] + dirty = torch.full( + (1, H, K_DIM, K_DIM), float("nan"), device=device, dtype=torch.float32 + ) + o_zero, ht_zero = chunk_kimi_delta_attn( + **_kimi_kwargs( + *args, + initial_state=torch.zeros_like(dirty), + output_final_state=True, + ) + ) + cache, indices, _storage = _paged_pool(1, H, dirty) + out = torch.empty_like(v) + paged_kw = _kimi_kwargs( + *args, + out=out, + state_cache=cache, + state_indices=indices, + has_initial_state=torch.zeros(1, device=device, dtype=torch.bool), + ) + # Warm the paged specialization's autotune, then re-dirty the slot. + chunk_kimi_delta_attn(**paged_kw) + cache[indices] = dirty + o_paged, _ = chunk_kimi_delta_attn(**paged_kw) + assert torch.equal(o_paged, o_zero) + assert torch.equal(cache[indices], ht_zero) + + +def test_paged_cache_padded_stride_writes_live_pool(): + """stride(0) padding must not trigger a packed copy.""" + H = 4 + args = make_inputs(1, 256, H) + v = args[2] + h0 = torch.randn(1, H, K_DIM, K_DIM, device=device, dtype=torch.float32) * 0.1 + _, ht = chunk_kimi_delta_attn( + **_kimi_kwargs(*args, initial_state=h0, output_final_state=True) + ) + cache, indices, storage = _paged_pool(1, H, h0, pad_floats=128) + assert not cache.is_contiguous() + ptr = storage.data_ptr() + out = torch.empty_like(v) + chunk_kimi_delta_attn( + **_kimi_kwargs( + *args, + out=out, + state_cache=cache, + state_indices=indices, + has_initial_state=torch.ones(1, device=device, dtype=torch.bool), + ) + ) + assert storage.data_ptr() == ptr + assert torch.equal(cache[indices], ht) + # Padding between slots and the unused first/last rows stay zero. + inner = H * K_DIM * K_DIM + slot_stride = inner + 128 + for slot in (0, int(indices.item()) + 1): + row = storage[slot * slot_stride : (slot + 1) * slot_stride] + assert torch.count_nonzero(row).item() == 0 + + +@pytest.mark.parametrize("outer", ["expanded", "overlapping"]) +def test_paged_cache_rejects_overlapping_slots(outer): + """stride(0) below one [H, V, K] plane would alias neighbouring slots.""" + H = 4 + args = make_inputs(1, 128, H) + v = args[2] + plane = H * K_DIM * K_DIM + if outer == "expanded": + cache = torch.zeros(1, H, K_DIM, K_DIM, device=device).expand(3, -1, -1, -1) + else: + storage = torch.zeros(3 * plane, device=device) + cache = torch.as_strided( + storage, + size=(3, H, K_DIM, K_DIM), + stride=(plane // 2, K_DIM * K_DIM, K_DIM, 1), + ) + assert cache.stride(0) < plane + with pytest.raises(ValueError, match="not dense"): + chunk_kimi_delta_attn( + **_kimi_kwargs( + *args, + out=torch.empty_like(v), + state_cache=cache, + state_indices=torch.tensor([1], device=device, dtype=torch.int32), + has_initial_state=torch.ones(1, device=device, dtype=torch.bool), + ) + ) + + +def test_out_rejected_on_default_pipeline(): + """The default pipeline cannot honour `out`, so it must not drop it.""" + args = make_inputs(1, 128, 4) + v = args[2] + with ( + _force_default_pipeline(), + pytest.raises(ValueError, match="out is only implemented"), + ): + chunk_kimi_delta_attn(**_kimi_kwargs(*args, out=torch.empty_like(v))) + + +def test_out_must_be_contiguous(): + """A strided `out` is rejected, not silently packed into a clone.""" + args = make_inputs(1, 128, 4) + v = args[2] + B, T, H, V = v.shape + strided = torch.empty(B, T, H, 2 * V, device=device, dtype=v.dtype)[..., :V] + assert not strided.is_contiguous() + with pytest.raises(ValueError, match="out must be contiguous"): + chunk_kimi_delta_attn(**_kimi_kwargs(*args, out=strided)) + + +def test_out_is_written_in_place(): + """A contiguous `out` is the buffer the kernel writes, not a packed copy.""" + args = make_inputs(1, 128, 4) + ref, _ = chunk_kimi_delta_attn(**_kimi_kwargs(*args)) + out = torch.full_like(args[2], float("nan")) + o, _ = chunk_kimi_delta_attn(**_kimi_kwargs(*args, out=out)) + assert o.data_ptr() == out.data_ptr() + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + +def test_paged_cache_rejected_on_default_pipeline(): + args = make_inputs(1, 128, 4) + v = args[2] + h0 = torch.zeros(1, 4, K_DIM, K_DIM, device=device, dtype=torch.float32) + cache, indices, _ = _paged_pool(1, 4, h0) + with ( + _force_default_pipeline(), + pytest.raises(ValueError, match="paged state_cache is only implemented"), + ): + chunk_kimi_delta_attn( + **_kimi_kwargs( + *args, + out=torch.empty_like(v), + state_cache=cache, + state_indices=indices, + has_initial_state=torch.ones(1, device=device, dtype=torch.bool), + ) + )