diff --git a/python/cudnn/engines/manifest.py b/python/cudnn/engines/manifest.py index a225b66f7..71136aa28 100644 --- a/python/cudnn/engines/manifest.py +++ b/python/cudnn/engines/manifest.py @@ -169,7 +169,13 @@ def offered_ids(self) -> Dict[str, int]: # kda_cake declines in check_support unless the opt-in flag is set (its # forward numerics vs FLA are not reconciled yet); the slot itself is not # gated, since no LA family may withhold its only implementations. - slots={"kda_frost": EngineSlot(0), "kda_cutile": EngineSlot(1), "kda_summary_frost": EngineSlot(2), "kda_cake": EngineSlot(3)}, + slots={ + "kda_frost": EngineSlot(0), + "kda_cutile": EngineSlot(1), + "kda_summary_frost": EngineSlot(2), + "kda_cake": EngineSlot(3), + "kda_hopper": EngineSlot(4), + }, analyzer=("cudnn.linear_attention.graph_analyzer", "analyze"), ), EngineFamily( diff --git a/python/cudnn/linear_attention/__init__.py b/python/cudnn/linear_attention/__init__.py index afe1852ff..b9cdda553 100644 --- a/python/cudnn/linear_attention/__init__.py +++ b/python/cudnn/linear_attention/__init__.py @@ -73,13 +73,14 @@ def GdnEngines(ids): def KdaEngines(ids): - """The KDA family: frost, cuTile, and the CAKE C16 training route.""" + """The KDA family: frost, cuTile, the CAKE C16 training route, and the Hopper sm90 path.""" return _collect( ids, ("kda_frost", "cudnn.linear_attention.frost.kda_engine", "KdaFrostEngine"), ("kda_cutile", "cudnn.linear_attention.cutile.kda_engine", "KdaCuTileEngine"), ("kda_summary_frost", "cudnn.linear_attention.frost.kda_engine", "KdaSummaryFrostEngine"), ("kda_cake", "cudnn.linear_attention.cake.kda_engine", "KdaCakeEngine"), + ("kda_hopper", "cudnn.linear_attention.hopper.kda_engine", "KdaHopperEngine"), ) diff --git a/python/cudnn/linear_attention/hopper/__init__.py b/python/cudnn/linear_attention/hopper/__init__.py new file mode 100644 index 000000000..248ff7870 --- /dev/null +++ b/python/cudnn/linear_attention/hopper/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hopper (sm90) linear-attention engines.""" diff --git a/python/cudnn/linear_attention/hopper/kda_engine.py b/python/cudnn/linear_attention/hopper/kda_engine.py new file mode 100644 index 000000000..9c4dc7be5 --- /dev/null +++ b/python/cudnn/linear_attention/hopper/kda_engine.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hopper (sm90) KDA engine. + +The FROST KDA kernels are Blackwell-only by construction -- they are written on +``tcgen05`` MMA and Tensor Memory, neither of which exists on Hopper -- so +``frost/engine.py`` gates them to ``100 <= sm <= 103 or sm == 107`` and sm90 is +left with no linear-attention path at all (cuTile is the only other backend and +needs the ``cuda.tile`` runtime). + +This engine fills that hole with a Hopper-native CuTe DSL kernel: a +chunk-parallel PREP pass followed by a sequential SCAN over the [128, 128] +state, built on ``warpgroup`` (wgmma) against shared memory rather than TMEM. + +The kernel chunks at BT = 16 and arranges every exponent reaching ``exp2`` to be +``<= 0``, so the exponentials can only underflow to zero. That matters: the +production gate (``gate_lower_bound = -5``, mean log-decay ~ -2.5) makes a +64-token chunk span ~118 in the exponent, and fp32 overflows at 88. + +``initial_state`` IS supported -- the recurrence is seeded from it, so this +serves chunked-prefill continuation as well as whole-sequence prefill. A graph +that omits it is handed a zero seed. + +Scope is otherwise deliberately narrow, and everything outside it is DECLINED +rather than silently mis-served -- see :meth:`KdaHopperEngine.check_support`. +""" + +from typing import TYPE_CHECKING + +from cudnn import behavior_note +from cudnn.engines.base import BaseEngine, CompiledPlan, bind_ports +from cudnn.frost import buffers +from cudnn.graph_types import NodeType + +from ..graph_analyzer import analyze + +if TYPE_CHECKING: + from cudnn._pygraph import pygraph + +HOPPER_SM = 90 +HEAD_DIM = 128 + +# DLPack alignment hint per port. cu_seqlens is int32 and only 4-byte aligned; +# everything else is a 16-byte-aligned device allocation. +_ALIGN = {"cu_seqlens": 4} + + +class KdaHopperPlan(CompiledPlan): + """Bind the node's ports and hand them to the vendored sm90 kernel. + + The kernel is destination-passing: ``o`` and ``final_state`` are written in + place. It always produces a final state, so when the graph did not ask for + one we pass a scratch buffer and drop it. + """ + + takes_variant_pack = True + plan_name = "KdaHopperEngine" + + def __init__(self, graph): + from .kernel import kda_prefill_sm90 as kernel + + (node,) = graph.nodes + self.kernel = kernel + q, v, cu = (node.inputs[p] for p in ("q", "v", "cu_seqlens")) + self.total, self.h, self.k = (int(d) for d in q.dim) + self.v_dim = int(v.dim[2]) + self.n_seqs = int(cu.dim[0]) - 1 + self.want_state = "final_state" in node.outputs + self.ports = None + self._scratch_state = None + self._scratch_state_cute = None + self._zero_state = None + self._zero_state_cute = None + self._device = None + + def get_workspace_size(self) -> int: + # The kernel owns its own scratch (module-level cache keyed by shape and + # device), so nothing is carved out of the caller's workspace. + return 0 + + def execute(self, graph, variant_pack, ctx) -> None: + import torch + from cutlass.cute.runtime import from_dlpack + + if self.ports is None: + self.ports = bind_ports(graph, variant_pack) + (slots,) = self.ports.values() + self.names = list(slots.inputs) + list(slots.outputs) + self.indices = list(slots.inputs.values()) + list(slots.outputs.values()) + views = variant_pack.operands(self.indices) + # Operands arrive as cuDNN OperandBuffer views, which implement DLPack, + # so convert them STRAIGHT to CuTe tensors. Borrowing them as torch + # tensors first would cost two DLPack conversions per operand instead of + # one, and the kernel only ever wanted the CuTe form -- that double + # conversion was the single largest term in this engine's dispatch cost. + nb = {name: from_dlpack(view, assumed_align=_ALIGN.get(name, 16)) for name, view in zip(self.names, views)} + + if self._device is None: + # One-time: the scratch buffers below need a torch device, and the + # operand views do not carry one directly. + self._device = torch.from_dlpack(views[0]).device + + final_state = nb.get("final_state") + if final_state is None: + if self._scratch_state is None: + self._scratch_state = torch.empty( + self.n_seqs, + self.h, + self.v_dim, + self.k, + dtype=torch.float32, + device=self._device, + ) + self._scratch_state_cute = from_dlpack(self._scratch_state, assumed_align=16) + final_state = self._scratch_state_cute + + # The kernel always reads a seed. A graph without initial_state means + # "start from zero", so hand it a zero buffer rather than declining. + initial_state = nb.get("initial_state") + if initial_state is None: + if self._zero_state is None: + self._zero_state = torch.zeros( + self.n_seqs, + self.h, + self.v_dim, + self.k, + dtype=torch.float32, + device=self._device, + ) + self._zero_state_cute = from_dlpack(self._zero_state, assumed_align=16) + initial_state = self._zero_state_cute + + # Hand the stream down explicitly rather than pushing a torch stream + # context for the kernel to read back out of thread-local state. + stream_ptr = ctx.stream if ctx.stream else torch.cuda.current_stream().cuda_stream + self.kernel.run_cute( + nb["q"], + nb["k"], + nb["v"], + nb["g"], + nb["beta"], + nb["cu_seqlens"], + initial_state, + nb["O"], + final_state, + self.total, + self.h, + self.k, + self.n_seqs, + self._device, + stream_ptr, + ) + + +class KdaHopperEngine(BaseEngine): + """Hopper (sm90) CuTe DSL backend for single-node KDA forward graphs (THD).""" + + name = "kda_hopper" + behavior_notes = (behavior_note.RUNTIME_COMPILATION,) + + def check_support(self, graph: "pygraph") -> None: + import cudnn + + facts = graph._facts_for(analyze) + if facts is None or facts.op != "KDA": + raise NotImplementedError("KdaHopperEngine supports exactly one KDA node") + if facts.invalid: + raise NotImplementedError(f"KdaHopperEngine: {facts.invalid}") + + sm = buffers.current_sm() + if sm != HOPPER_SM: + raise NotImplementedError(f"KdaHopperEngine is the Hopper path and requires SM90 (found {sm})") + + try: + from .kernel import kda_prefill_sm90 # noqa: F401 -- availability probe + except ImportError as exc: + raise NotImplementedError(f"KdaHopperEngine requires the cutedsl extra: {exc}") from exc + + # --- scope. Each of these is a real limit of the vendored kernel, and + # declining is the point: a silently mis-served graph is worse than no + # Hopper path at all. + if facts.is_bwd: + raise NotImplementedError("KdaHopperEngine: forward only; there is no Hopper KDA backward kernel yet") + if facts.checkpoint_every_n_tokens: + raise NotImplementedError("KdaHopperEngine: state_checkpoints are not produced by the Hopper kernel") + if facts.safe_gate or facts.has_a_log or facts.has_dt_bias: + raise NotImplementedError("KdaHopperEngine: the kernel takes log-space g directly; safe_gate/a_log/dt_bias are unsupported") + if facts.use_beta_sigmoid: + raise NotImplementedError("KdaHopperEngine: beta must be post-sigmoid; use_beta_sigmoid_in_kernel is unsupported") + if facts.use_qk_l2norm: + raise NotImplementedError("KdaHopperEngine: q/k must be pre-normalized; use_qk_l2norm_in_kernel is unsupported") + if getattr(facts, "gate_domain", "log") != "log": + raise NotImplementedError("KdaHopperEngine: gate_domain='linear' is unsupported") + + # The kernel bakes in q * 1/sqrt(D); it takes no scale argument, so any + # other scale would be silently ignored and produce a wrong answer + # (caught by test_fwd_scale, which saw an rms ratio of 0.91). + if facts.scale is not None and abs(facts.scale - HEAD_DIM**-0.5) > 1e-9: + raise NotImplementedError(f"KdaHopperEngine: only the default scale 1/sqrt({HEAD_DIM}) is supported, got {facts.scale}") + if facts.cu_dtype not in (cudnn.data_type.INT32, None): + raise NotImplementedError(f"KdaHopperEngine: cu_seqlens must be int32, got {facts.cu_dtype}") + + if facts.d_qk != HEAD_DIM or facts.d_v != HEAD_DIM: + raise NotImplementedError(f"KdaHopperEngine: head dims must be {HEAD_DIM}, got K={facts.d_qk} V={facts.d_v}") + if not (facts.h_q == facts.h_k == facts.h_v): + raise NotImplementedError( + f"KdaHopperEngine: grouped heads are unsupported; q/k/v head counts must match " f"(got {facts.h_q}/{facts.h_k}/{facts.h_v})" + ) + if facts.io_dtype not in (cudnn.data_type.BFLOAT16, None): + raise NotImplementedError(f"KdaHopperEngine: q/k/v must be bf16, got {facts.io_dtype}") + if facts.g_dtype not in (cudnn.data_type.FLOAT, None): + raise NotImplementedError(f"KdaHopperEngine: 'g' must be fp32, got {facts.g_dtype}") + if facts.beta_dtype not in (cudnn.data_type.FLOAT, None): + raise NotImplementedError(f"KdaHopperEngine: 'beta' must be fp32, got {facts.beta_dtype}") + if facts.final_state_dtype not in (cudnn.data_type.FLOAT, None): + raise NotImplementedError(f"KdaHopperEngine: 'final_state' must be fp32, got {facts.final_state_dtype}") + # The kernel reads and writes the [128, 128] state as fp32 wgmma + # accumulators; a bf16 state pool would be reinterpreted, not converted. + if facts.state_dtype not in (cudnn.data_type.FLOAT, None): + raise NotImplementedError(f"KdaHopperEngine: 'initial_state' must be fp32, got {facts.state_dtype}") + if not facts.thd_layout: + raise NotImplementedError("KdaHopperEngine: q/k/v must be THD [total_T, heads, dim]") + + def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: + return KdaHopperPlan(graph) diff --git a/python/cudnn/linear_attention/hopper/kernel/__init__.py b/python/cudnn/linear_attention/hopper/kernel/__init__.py new file mode 100644 index 000000000..38d6838c4 --- /dev/null +++ b/python/cudnn/linear_attention/hopper/kernel/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Vendored Hopper (sm90) KDA prefill kernels.""" + +from .kda_prefill_sm90 import run # noqa: F401 diff --git a/python/cudnn/linear_attention/hopper/kernel/kda_prefill_sm90.py b/python/cudnn/linear_attention/hopper/kernel/kda_prefill_sm90.py new file mode 100644 index 000000000..244a3792b --- /dev/null +++ b/python/cudnn/linear_attention/hopper/kernel/kda_prefill_sm90.py @@ -0,0 +1,839 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Kimi Delta Attention (KDA) chunked linear-attention prefill for Hopper (sm90). + +PREP (grid = (max_chunks_per_seq, N, H), 128 thr) -- fully parallel over +16-token chunks. Per chunk it builds the pieces of the chunk's affine operator +in UT/WY form. BT = 16 plus "every exponent that reaches exp2 is <= 0" keeps +the production gate (g in (-5,0), mean -2.5) inside fp32 range: + + G = chunk-local inclusive cumsum of g (<= 0, decreasing) + E = exp(G) in (0,1] , a = E[L-1] + kw = k * E (<= |k|) + ub = beta * k / E (<= e^{|G_L|}) + kg = ub * a = beta*k*exp(G_L - G) (<= |k|) + qg = scale * q * E + T = I + strict_lower(kw @ ub^T) , Z = tril(qg @ ub^T, 0) + wt = T^{-1} kw , ut = T^{-1} v (forward substitution) + +SCAN (grid = (N, H), 256 thr = 2 warpgroups) -- sequential over chunks with the +state kept TRANSPOSED as an fp32 wgmma accumulator St = S^T [Dv=128, Dk=128], +which is exactly the V-major layout of initial_state / final_state: + + Rt = ut^T + St @ (-wt)^T (M=128, N=16, K=128) + Ot = St @ qg^T (M=128, N=16, K=128) + Ot += Rt @ Z^T (M=128, N=16, K=16) + St = St*a + Rt @ kg (M=128, N=128, K=16) + +Transposing the state makes every wgmma full-M (128 rows) instead of the 16-row +shape the untransposed recurrence would give. Both St and Rt are fed to WGMMA +as REGISTER A-operands (OperandSource.RMEM), so the [128,128] state never moves +through shared memory -- that removes ~96 KB of SMEM traffic and two CTA +barriers from every step of the serial chain. + +Provenance +---------- +Generated by a Kernel Factory campaign on the ``kda_chunked_prefill_sm90`` +operation (campaign ``m58kq4q63h0zn893p6sxywe9g4``, kernel +``b134720b4c5e787ac0d27c210c3faf9e5d6a8eacc9385236415175b97d067b7c``, +solution ``kda_sm90_segscan_leader_checkpoint``). + +Verified on H100 80GB HBM3 (SXM) at the production gate +(``gate_lower_bound = -5``) with a non-zero ``initial_state``, geomean over ten +shapes, two independent passes: + +=================== =========== ========== ========== +metric this kernel FlashKDA ratio +=================== =========== ========== ========== +per-call 398.1 us 442.4 us 1.11x +pipelined 260.8 us 401.0 us 1.53x +=================== =========== ========== ========== + +Accuracy against an fp64 oracle is 5.0e-03 to 1.0e-02, against FlashKDA's own +7.5e-03 to 1.0e-02 on the same inputs. + +The predecessor vendored here used ``BT = 64`` with a single mid-chunk anchor and +returned 100% NaN at ``gate_lower_bound = -5``: a 64-token chunk spans ~118 in +the exponent there and fp32 overflows at 88. Per-sub-block anchoring cannot fix +that, because KDA's decay is per-key-channel, so the cross-block correction +``exp(r_j[d] - r_i[d])`` lives inside the channel sum of ``Wn @ U^T`` and no +post-matmul scalar removes it. Every token in one matmul must share an anchor, +which forces ``BT <= 16``. This kernel goes further and arranges for every +exponent reaching ``exp2`` to be ``<= 0``, so the exponential can only underflow +to zero -- overflow is structurally impossible rather than merely budgeted. +""" + +import os + +os.environ.setdefault("CUTE_DSL_ARCH", "sm_90a") + +import math + +import torch + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.hopper_helpers as hh +from cutlass.cute.nvgpu import cpasync, warpgroup +from cutlass.cute.runtime import from_dlpack +import cuda.bindings.driver as cudadrv + +BF16 = cutlass.BFloat16 +F32 = cutlass.Float32 +I32 = cutlass.Int32 + +LOG2E = 1.4426950408889634 +DIM = 128 +NSTAGE = 6 +TXBYTES = 4 * (16 * 128 * 2) + 16 * 16 * 2 + 128 * 4 + + +def c2a_layout(c, a): + return cute.make_layout( + (a, c.shape[1], (c.shape[2], cute.size(c, mode=[0]) // cute.size(a))), + stride=( + c.stride[0], + c.stride[1], + (c.stride[2], cute.size(a, mode=[2]) * c.stride[0][2]), + ), + ) + + +@cute.jit +def acc_to_a(acc, a_shape): + operand = cute.make_rmem_tensor_like(c2a_layout(acc.layout, a_shape), BF16) + as_acc = cute.make_tensor(operand.iterator, acc.layout) + as_acc.store(acc.load().to(BF16)) + return operand + + +# --------------------------------------------------------------------------- +# PREP +# --------------------------------------------------------------------------- +@cute.kernel +def prep_kernel( + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mG: cute.Tensor, + mBeta: cute.Tensor, + mCu: cute.Tensor, + mMW: cute.Tensor, + mQG: cute.Tensor, + mKG: cute.Tensor, + mUT: cute.Tensor, + mZ: cute.Tensor, + mAv: cute.Tensor, + tiled_mma: cute.TiledMma, + scale: cutlass.Constexpr, + nch: cutlass.Constexpr, +): + tidx, _, _ = cute.arch.thread_idx() + lc, sq, hq = cute.arch.block_idx() + + s0 = mCu[sq] + s1 = mCu[sq + 1] + ncs = (s1 - s0 + 15) // 16 + + if lc < ncs: + sA_layout = cute.slice_( + hh.make_smem_layout_a(utils.LayoutEnum.ROW_MAJOR, (64, 16, 128), BF16, 1), + (None, None, 0), + ) + sB_layout = cute.slice_( + hh.make_smem_layout_b(utils.LayoutEnum.ROW_MAJOR, (64, 16, 128), BF16, 1), + (None, None, 0), + ) + smem = utils.SmemAllocator() + sA = smem.allocate_tensor(BF16, sA_layout.outer, byte_alignment=1024, swizzle=sA_layout.inner) + sB = smem.allocate_tensor(BF16, sB_layout.outer, byte_alignment=1024, swizzle=sB_layout.inner) + sM = smem.allocate_tensor(F32, cute.make_layout((64, 16)), byte_alignment=16) + + tok0 = s0 + lc * 16 + cb = hq * nch + s0 // 16 + sq + lc + d = tidx + + kw = [] + ub = [] + qg = [] + vv = [] + cs = F32(0.0) + ee = F32(1.0) + for i in cutlass.range_constexpr(16): + t = cutlass.min(tok0 + i, s1 - 1) + msk = F32(1.0) + if tok0 + i >= s1: + msk = F32(0.0) + kx = mK[t, hq, d].to(F32) * msk + qx = mQ[t, hq, d].to(F32) * msk + vx = mV[t, hq, d].to(F32) * msk + gx = mG[t, hq, d] * msk + bx = mBeta[t, hq] * msk + cs = cs + gx * F32(LOG2E) + ee = cute.math.exp2(cs, fastmath=True) + rr = cute.math.exp2(-cs, fastmath=True) + w = kx * ee + u = bx * kx * rr + qw = qx * ee * F32(scale) + kw.append(w) + ub.append(u) + qg.append(qw) + vv.append(vx) + sA[i, d] = w.to(BF16) + sA[16 + i, d] = qw.to(BF16) + sB[i, d] = u.to(BF16) + for i in cutlass.range_constexpr(32, 64): + sA[i, d] = BF16(0.0) + + aval = ee + cute.arch.barrier() + + thr_mma = tiled_mma.get_slice(0) + tCrA = tiled_mma.make_fragment_A(thr_mma.partition_A(sA)) + tCrB = tiled_mma.make_fragment_B(thr_mma.partition_B(sB)) + acc = cute.make_rmem_tensor(thr_mma.partition_C(cute.make_identity_tensor((64, 16))).shape, F32) + acc.fill(0.0) + tiled_mma.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.fence() + for kb in cutlass.range_constexpr(8): + cute.gemm(tiled_mma, acc, tCrA[None, None, kb], tCrB[None, None, kb], acc) + warpgroup.commit_group() + warpgroup.wait_group(0) + + copy_f32 = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), F32) + tc = cute.make_tiled_copy_C(copy_f32, tiled_mma) + thr_tc = tc.get_slice(tidx) + cute.autovec_copy(thr_tc.retile(acc), thr_tc.partition_D(sM)) + + cute.arch.barrier() + + wt = [kw[0]] + ut = [vv[0]] + for i in cutlass.range_constexpr(1, 16): + sw = kw[i] + su = vv[i] + for j in cutlass.range_constexpr(i): + m = sM[i, j] + sw = sw - m * wt[j] + su = su - m * ut[j] + wt.append(sw) + ut.append(su) + + for i in cutlass.range_constexpr(16): + mMW[cb * 16 + i, d] = (-wt[i]).to(BF16) + mQG[cb * 16 + i, d] = qg[i].to(BF16) + mKG[cb * 16 + i, d] = (ub[i] * aval).to(BF16) + mUT[cb * 16 + i, d] = ut[i].to(BF16) + mAv[cb, d] = aval + + zi = tidx // 16 + zj = tidx % 16 + for r in cutlass.range_constexpr(2): + ii = zi + r * 8 + zz = F32(0.0) + if zj <= ii: + zz = sM[16 + ii, zj] + mZ[cb * 16 + ii, zj] = zz.to(BF16) + + +# --------------------------------------------------------------------------- +# SEGMENT kernel: BUILD (segment affine operator) or EXPAND (outputs) +# +# chunk step (transposed state St [Dv, Dk]): +# Rt = ut^T + St @ (-wt)^T ; St = St*a + Rt @ kg +# so the chunk's affine map is St_new = St Mt_c + Ct_c with +# Mt_c = Diag(a) - wt^T kg , Ct_c = ut^T kg +# Composing over a segment obeys the SAME recurrence, with +# Mt: accumulator seeded with I and no ut term +# Ct: accumulator seeded with 0 and the ut term +# -- i.e. one kernel body serves the operator build and the expansion. +# --------------------------------------------------------------------------- +@cute.kernel +def seg_kernel( + mCu: cute.Tensor, + mIS: cute.Tensor, + mO: cute.Tensor, + mFS: cute.Tensor, + mSS: cute.Tensor, + mMT: cute.Tensor, + mCT: cute.Tensor, + atom_mw: cute.CopyAtom, + tMW: cute.Tensor, + atom_qg: cute.CopyAtom, + tQG: cute.Tensor, + atom_kg: cute.CopyAtom, + tKG: cute.Tensor, + atom_ut: cute.CopyAtom, + tUT: cute.Tensor, + atom_z: cute.CopyAtom, + tZ: cute.Tensor, + atom_av: cute.CopyAtom, + tAV: cute.Tensor, + mma_s: cute.TiledMma, + mma_n: cute.TiledMma, + nch: cutlass.Constexpr, + nheads: cutlass.Constexpr, + NSEG: cutlass.Constexpr, + MODE: cutlass.Constexpr, +): + BUILD = cutlass.const_expr(MODE < 2) + WHICH = cutlass.const_expr(1 if MODE != 0 else 0) + TXB = cutlass.const_expr(8704 + (4096 if MODE >= 1 else 0) + (4608 if MODE == 2 else 0)) + tidx, _, _ = cute.arch.thread_idx() + seg, hq, zz = cute.arch.block_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + if warp_idx == 0: + cpasync.prefetch_descriptor(atom_mw) + cpasync.prefetch_descriptor(atom_kg) + cpasync.prefetch_descriptor(atom_av) + if cutlass.const_expr(MODE >= 1): + cpasync.prefetch_descriptor(atom_ut) + if cutlass.const_expr(MODE == 2): + cpasync.prefetch_descriptor(atom_qg) + cpasync.prefetch_descriptor(atom_z) + + sq = zz + + s0 = mCu[sq] + s1 = mCu[sq + 1] + ncsq = (s1 - s0 + 15) // 16 + seglen = (ncsq + (NSEG - 1)) // NSEG + c0 = cutlass.min(seg * seglen, ncsq) + c1 = cutlass.min(c0 + seglen, ncsq) + ncs = c1 - c0 + cb = hq * nch + s0 // 16 + sq + c0 + opi = (sq * nheads + hq) * NSEG + seg + + sMW_l = hh.make_smem_layout_b(utils.LayoutEnum.ROW_MAJOR, (128, 16, 128), BF16, NSTAGE) + sKG_l = hh.make_smem_layout_b(utils.LayoutEnum.COL_MAJOR, (128, 128, 16), BF16, NSTAGE) + sZ_l = hh.make_smem_layout_b(utils.LayoutEnum.ROW_MAJOR, (128, 16, 16), BF16, NSTAGE) + sAV_l = cute.make_layout((1, 128, NSTAGE), stride=(128, 1, 128)) + + smem = utils.SmemAllocator() + mbar = smem.allocate_array(cutlass.Int64, NSTAGE * 2, byte_alignment=16) + sMW = smem.allocate_tensor(BF16, sMW_l.outer, byte_alignment=1024, swizzle=sMW_l.inner) + sQG = smem.allocate_tensor(BF16, sMW_l.outer, byte_alignment=1024, swizzle=sMW_l.inner) + sKG = smem.allocate_tensor(BF16, sKG_l.outer, byte_alignment=1024, swizzle=sKG_l.inner) + sZ = smem.allocate_tensor(BF16, sZ_l.outer, byte_alignment=1024, swizzle=sZ_l.inner) + sUT = smem.allocate_tensor(BF16, sKG_l.outer, byte_alignment=1024, swizzle=sKG_l.inner) + sAV = smem.allocate_tensor(F32, sAV_l, byte_alignment=1024) + + pipe = pipeline.PipelineTmaAsync.create( + barrier_storage=mbar, + num_stages=NSTAGE, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Warp, 8), + tx_count=TXB, + cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), + enable_multicast_signaling=True, + ) + + cta1 = cute.make_layout(1) + gMW = cute.local_tile(tMW, (16, 128), (None, 0)) + gQG = cute.local_tile(tQG, (16, 128), (None, 0)) + gKG = cute.local_tile(tKG, (128, 16), (0, None)) + gUT = cute.local_tile(tUT, (128, 16), (0, None)) + gZ = cute.local_tile(tZ, (16, 16), (None, 0)) + gAV = cute.local_tile(tAV, (1, 128), (None, 0)) + + aMWs, aMWg = cpasync.tma_partition(atom_mw, 0, cta1, cute.group_modes(sMW, 0, 2), cute.group_modes(gMW, 0, 2)) + aQGs, aQGg = cpasync.tma_partition(atom_qg, 0, cta1, cute.group_modes(sQG, 0, 2), cute.group_modes(gQG, 0, 2)) + aKGs, aKGg = cpasync.tma_partition(atom_kg, 0, cta1, cute.group_modes(sKG, 0, 2), cute.group_modes(gKG, 0, 2)) + aUTs, aUTg = cpasync.tma_partition(atom_ut, 0, cta1, cute.group_modes(sUT, 0, 2), cute.group_modes(gUT, 0, 2)) + aZs, aZg = cpasync.tma_partition(atom_z, 0, cta1, cute.group_modes(sZ, 0, 2), cute.group_modes(gZ, 0, 2)) + aAVs, aAVg = cpasync.tma_partition(atom_av, 0, cta1, cute.group_modes(sAV, 0, 2), cute.group_modes(gAV, 0, 2)) + + wgi = cute.arch.make_warp_uniform(tidx // 128) + thr_s = mma_s.get_slice(cute.make_layout(2, stride=128)(wgi)) + thr_n = mma_n.get_slice(cute.make_layout(2, stride=128)(wgi)) + + accS = cute.make_rmem_tensor(thr_s.partition_C(cute.make_identity_tensor((128, 128))).shape, F32) + accR = cute.make_rmem_tensor(thr_n.partition_C(cute.make_identity_tensor((128, 16))).shape, F32) + accO = cute.make_rmem_tensor(thr_n.partition_C(cute.make_identity_tensor((128, 16))).shape, F32) + + copy_f32 = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), F32) + tcS = cute.make_tiled_copy_C(copy_f32, mma_s) + thr_tcS = tcS.get_slice(tidx) + tcN = cute.make_tiled_copy_C(copy_f32, mma_n) + thr_tcN = tcN.get_slice(tidx) + + cS = thr_tcS.partition_D(cute.make_identity_tensor((128, 128))) + cO = thr_tcN.partition_D(cute.make_identity_tensor((128, 16))) + tAccS = thr_tcS.retile(accS) + tAccR = thr_tcN.retile(accR) + tAccO = thr_tcN.retile(accO) + + if cutlass.const_expr(BUILD): + accS.fill(0.0) + if cutlass.const_expr(WHICH == 0): + for e in cutlass.range_constexpr(cute.size(tAccS)): + if cS[e][0] == cS[e][1]: + tAccS[e] = F32(1.0) + else: + if cutlass.const_expr(NSEG == 1): + cute.autovec_copy(thr_tcS.partition_D(mIS[sq, hq, None, None]), tAccS) + else: + cute.autovec_copy(thr_tcS.partition_D(mSS[opi, None, None]), tAccS) + + prod = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, NSTAGE) + cons = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, NSTAGE) + + npre = cutlass.min(I32(NSTAGE), ncs) + if warp_idx == 0: + for pi in cutlass.range(npre, unroll=1): + pipe.producer_acquire(prod) + bp = pipe.producer_get_barrier(prod) + ti = cb + prod.count + si = prod.index + cute.copy(atom_mw, aMWg[(None, ti)], aMWs[(None, si)], tma_bar_ptr=bp) + cute.copy(atom_kg, aKGg[(None, ti)], aKGs[(None, si)], tma_bar_ptr=bp) + cute.copy(atom_av, aAVg[(None, ti)], aAVs[(None, si)], tma_bar_ptr=bp) + if cutlass.const_expr(MODE >= 1): + cute.copy(atom_ut, aUTg[(None, ti)], aUTs[(None, si)], tma_bar_ptr=bp) + if cutlass.const_expr(MODE == 2): + cute.copy(atom_qg, aQGg[(None, ti)], aQGs[(None, si)], tma_bar_ptr=bp) + cute.copy(atom_z, aZg[(None, ti)], aZs[(None, si)], tma_bar_ptr=bp) + prod.advance() + + mma_s.set(warpgroup.Field.ACCUMULATE, True) + mma_n.set(warpgroup.Field.ACCUMULATE, True) + a_shape_n = mma_n.tv_layout_A.shape[1] + + for ci in cutlass.range(ncs, unroll=1): + pipe.consumer_wait(cons) + st = cons.index + tok0 = s0 + (c0 + ci) * 16 + + accR.fill(0.0) + if cutlass.const_expr(WHICH == 1): + rUT = cute.make_rmem_tensor_like(tAccR, BF16) + cute.autovec_copy(thr_tcN.partition_D(sUT[None, None, st]), rUT) + tAccR.store(rUT.load().to(F32)) + + sAVb = cute.make_tensor(sAV[0, None, st].iterator, cute.make_layout((128, 128), stride=(0, 1))) + rAv = cute.make_rmem_tensor_like(tAccS, F32) + cute.autovec_copy(thr_tcS.partition_D(sAVb), rAv) + + aSt = acc_to_a(accS, a_shape_n) + tAccS.store(tAccS.load() * rAv.load()) + + tMWr = mma_n.make_fragment_B(thr_n.partition_B(sMW[None, None, st])) + + warpgroup.fence() + for kb in cutlass.range_constexpr(8): + cute.gemm(mma_n, accR, aSt[None, None, kb], tMWr[None, None, kb], accR) + warpgroup.commit_group() + if cutlass.const_expr(not BUILD): + accO.fill(0.0) + tQGr = mma_n.make_fragment_B(thr_n.partition_B(sQG[None, None, st])) + warpgroup.fence() + for kb in cutlass.range_constexpr(8): + cute.gemm(mma_n, accO, aSt[None, None, kb], tQGr[None, None, kb], accO) + warpgroup.commit_group() + warpgroup.wait_group(1) + else: + warpgroup.wait_group(0) + + aRt = acc_to_a(accR, a_shape_n) + tKGr = mma_s.make_fragment_B(thr_s.partition_B(sKG[None, None, st])) + + warpgroup.fence() + if cutlass.const_expr(not BUILD): + tZr = mma_n.make_fragment_B(thr_n.partition_B(sZ[None, None, st])) + cute.gemm(mma_n, accO, aRt[None, None, 0], tZr[None, None, 0], accO) + cute.gemm(mma_s, accS, aRt[None, None, 0], tKGr[None, None, 0], accS) + warpgroup.commit_group() + warpgroup.wait_group(0) + + pipe.consumer_release(cons) + cons.advance() + + if cutlass.const_expr(not BUILD): + rO = cute.make_rmem_tensor_like(tAccO, BF16) + rO.store(tAccO.load().to(BF16)) + for e in cutlass.range_constexpr(cute.size(tAccO)): + tt = tok0 + cO[e][1] + if tt < s1: + mO[tt, hq, cO[e][0]] = rO[e] + + nxt = ci + NSTAGE + if warp_idx == 0: + if nxt < ncs: + pipe.producer_acquire(prod) + bp = pipe.producer_get_barrier(prod) + ti = cb + nxt + si = prod.index + cute.copy(atom_mw, aMWg[(None, ti)], aMWs[(None, si)], tma_bar_ptr=bp) + cute.copy(atom_kg, aKGg[(None, ti)], aKGs[(None, si)], tma_bar_ptr=bp) + cute.copy(atom_av, aAVg[(None, ti)], aAVs[(None, si)], tma_bar_ptr=bp) + if cutlass.const_expr(MODE >= 1): + cute.copy(atom_ut, aUTg[(None, ti)], aUTs[(None, si)], tma_bar_ptr=bp) + if cutlass.const_expr(MODE == 2): + cute.copy(atom_qg, aQGg[(None, ti)], aQGs[(None, si)], tma_bar_ptr=bp) + cute.copy(atom_z, aZg[(None, ti)], aZs[(None, si)], tma_bar_ptr=bp) + prod.advance() + + if cutlass.const_expr(BUILD): + if cutlass.const_expr(WHICH == 0): + rM = cute.make_rmem_tensor_like(tAccS, BF16) + rM.store(tAccS.load().to(BF16)) + cute.autovec_copy(rM, thr_tcS.partition_D(mMT[opi, None, None])) + else: + cute.autovec_copy(tAccS, thr_tcS.partition_D(mCT[opi, None, None])) + else: + if c1 == ncsq: + cute.autovec_copy(tAccS, thr_tcS.partition_D(mFS[sq, hq, None, None])) + + +# --------------------------------------------------------------------------- +# COMBINE: short serial scan over the NSEG segment operators +# S_p = S_{p-1} Mt_p + Ct_p (128x128 x 128x128 gemm per step) +# --------------------------------------------------------------------------- +@cute.kernel +def comb_kernel( + mIS: cute.Tensor, + mSS: cute.Tensor, + mMT: cute.Tensor, + mCT: cute.Tensor, + mma_c: cute.TiledMma, + nheads: cutlass.Constexpr, + NSEG: cutlass.Constexpr, +): + tidx, _, _ = cute.arch.thread_idx() + sq, hq, _ = cute.arch.block_idx() + + sMt_l = cute.slice_( + hh.make_smem_layout_b(utils.LayoutEnum.COL_MAJOR, (128, 128, 128), BF16, 1), + (None, None, 0), + ) + smem = utils.SmemAllocator() + sMt = smem.allocate_tensor(BF16, sMt_l.outer, byte_alignment=1024, swizzle=sMt_l.inner) + + wgi = cute.arch.make_warp_uniform(tidx // 128) + thr_c = mma_c.get_slice(cute.make_layout(2, stride=128)(wgi)) + accS = cute.make_rmem_tensor(thr_c.partition_C(cute.make_identity_tensor((128, 128))).shape, F32) + copy_f32 = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), F32) + tcC = cute.make_tiled_copy_C(copy_f32, mma_c) + thr_tcC = tcC.get_slice(tidx) + tAccS = thr_tcC.retile(accS) + + cp_bf = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), BF16) + tv = cute.make_tiled_copy_tv(cp_bf, cute.make_layout((16, 16), stride=(1, 16)), cute.make_layout((8, 1))) + thr_tv = tv.get_slice(tidx) + sdst = [] + for j in cutlass.range_constexpr(8): + sdst.append(thr_tv.partition_D(cute.local_tile(sMt, (128, 16), (0, j)))) + + base = (sq * nheads + hq) * NSEG + mma_c.set(warpgroup.Field.ACCUMULATE, True) + a_shape_c = mma_c.tv_layout_A.shape[1] + + cute.autovec_copy(thr_tcC.partition_D(mIS[sq, hq, None, None]), tAccS) + + for p in cutlass.range(NSEG, unroll=1): + idx = base + p + cute.autovec_copy(tAccS, thr_tcC.partition_D(mSS[idx, None, None])) + if p + 1 < NSEG: + gMt = cute.make_tensor(mMT[idx, None, None].iterator, cute.make_layout((128, 128), stride=(1, 128))) + rs = [] + for j in cutlass.range_constexpr(8): + src = thr_tv.partition_S(cute.local_tile(gMt, (128, 16), (0, j))) + r = cute.make_rmem_tensor_like(src, BF16) + cute.autovec_copy(src, r) + rs.append(r) + aS = acc_to_a(accS, a_shape_c) + cute.autovec_copy(thr_tcC.partition_D(mCT[idx, None, None]), tAccS) + cute.arch.barrier() + for j in cutlass.range_constexpr(8): + cute.autovec_copy(rs[j], sdst[j]) + cute.arch.fence_view_async_shared() + cute.arch.barrier() + tMr = mma_c.make_fragment_B(thr_c.partition_B(sMt)) + warpgroup.fence() + for kb in cutlass.range_constexpr(8): + cute.gemm(mma_c, accS, aS[None, None, kb], tMr[None, None, kb], accS) + warpgroup.commit_group() + warpgroup.wait_group(0) + + +# --------------------------------------------------------------------------- +# host +# --------------------------------------------------------------------------- +@cute.jit +def kda_launch( + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mG: cute.Tensor, + mBeta: cute.Tensor, + mCu: cute.Tensor, + mIS: cute.Tensor, + mMW: cute.Tensor, + mQG: cute.Tensor, + mKG: cute.Tensor, + mUT: cute.Tensor, + mZ: cute.Tensor, + mAv: cute.Tensor, + mKGt: cute.Tensor, + mUTt: cute.Tensor, + mSS: cute.Tensor, + mMT: cute.Tensor, + mCT: cute.Tensor, + mO: cute.Tensor, + mFS: cute.Tensor, + nseq: cutlass.Constexpr, + nheads: cutlass.Constexpr, + max_chunks: cutlass.Constexpr, + nch: cutlass.Constexpr, + nseg: cutlass.Constexpr, + stream: cudadrv.CUstream, +): + scale = 1.0 / math.sqrt(DIM) + + mma_p = hh.make_trivial_tiled_mma( + BF16, + BF16, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.K, + F32, + (1, 1, 1), + tiler_mn=(64, 16), + ) + prep_kernel( + mQ, + mK, + mV, + mG, + mBeta, + mCu, + mMW, + mQG, + mKG, + mUT, + mZ, + mAv, + mma_p, + scale, + nch, + ).launch(grid=[max_chunks, nseq, nheads], block=[128, 1, 1], stream=stream) + + mma_s = hh.make_trivial_tiled_mma( + BF16, + BF16, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.MN, + F32, + (2, 1, 1), + tiler_mn=(64, 128), + a_source=warpgroup.OperandSource.RMEM, + ) + mma_n = hh.make_trivial_tiled_mma( + BF16, + BF16, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.K, + F32, + (2, 1, 1), + tiler_mn=(64, 16), + a_source=warpgroup.OperandSource.RMEM, + ) + + sMW_1 = cute.slice_(hh.make_smem_layout_b(utils.LayoutEnum.ROW_MAJOR, (128, 16, 128), BF16, NSTAGE), (None, None, 0)) + sKG_1 = cute.slice_(hh.make_smem_layout_b(utils.LayoutEnum.COL_MAJOR, (128, 128, 16), BF16, NSTAGE), (None, None, 0)) + sZ_1 = cute.slice_(hh.make_smem_layout_b(utils.LayoutEnum.ROW_MAJOR, (128, 16, 16), BF16, NSTAGE), (None, None, 0)) + sAV_1 = cute.make_layout((1, 128), stride=(128, 1)) + + op = cpasync.CopyBulkTensorTileG2SOp() + atom_mw, tMW = cpasync.make_tiled_tma_atom(op, mMW, sMW_1, (16, 128)) + atom_qg, tQG = cpasync.make_tiled_tma_atom(op, mQG, sMW_1, (16, 128)) + atom_kg, tKG = cpasync.make_tiled_tma_atom(op, mKGt, sKG_1, (128, 16)) + atom_ut, tUT = cpasync.make_tiled_tma_atom(op, mUTt, sKG_1, (128, 16)) + atom_z, tZ = cpasync.make_tiled_tma_atom(op, mZ, sZ_1, (16, 16)) + atom_av, tAV = cpasync.make_tiled_tma_atom(op, mAv, sAV_1, (1, 128)) + + if cutlass.const_expr(nseg > 1): + for md in cutlass.range_constexpr(2): + seg_kernel( + mCu, + mIS, + mO, + mFS, + mSS, + mMT, + mCT, + atom_mw, + tMW, + atom_qg, + tQG, + atom_kg, + tKG, + atom_ut, + tUT, + atom_z, + tZ, + atom_av, + tAV, + mma_s, + mma_n, + nch, + nheads, + nseg, + md, + ).launch(grid=[nseg, nheads, nseq], block=[256, 1, 1], stream=stream) + + mma_c = hh.make_trivial_tiled_mma( + BF16, + BF16, + cute.nvgpu.OperandMajorMode.K, + cute.nvgpu.OperandMajorMode.MN, + F32, + (2, 1, 1), + tiler_mn=(64, 128), + a_source=warpgroup.OperandSource.RMEM, + ) + comb_kernel( + mIS, + mSS, + mMT, + mCT, + mma_c, + nheads, + nseg, + ).launch(grid=[nseq, nheads, 1], block=[256, 1, 1], stream=stream) + + seg_kernel( + mCu, + mIS, + mO, + mFS, + mSS, + mMT, + mCT, + atom_mw, + tMW, + atom_qg, + tQG, + atom_kg, + tKG, + atom_ut, + tUT, + atom_z, + tZ, + atom_av, + tAV, + mma_s, + mma_n, + nch, + nheads, + nseg, + 2, + ).launch(grid=[nseg, nheads, nseq], block=[256, 1, 1], stream=stream) + + +_CACHE = {} +_WS = {} + + +def _pick_nseg(T, N, H): + L = T // max(N, 1) + ncs = (L + 15) // 16 + p = min(132 // (N * H), ncs // 4, 8) + if p < 3: + return 1 + return p + + +def _ws(nch, H, N, nseg, device): + """Scratch buffers for one shape, returned as ALREADY-CONVERTED CuTe tensors. + + These eleven buffers are keyed by shape and never change identity, so + re-running ``from_dlpack`` on them at every launch is pure overhead -- it was + 11 of the 20 conversions per call, and DLPack conversion dominated this + kernel's host time. Convert once, here, and hand the same CuTe tensors to + every launch; only the caller's own operands need per-call conversion. + + The torch buffers are kept alive in the cache entry: the CuTe tensors borrow + their memory and do not own it. + """ + key = (nch, H, N, nseg, str(device)) + w = _WS.get(key) + if w is None: + c = nch * H + kg = torch.empty((c * 16, 128), dtype=torch.bfloat16, device=device) + ut = torch.empty((c * 16, 128), dtype=torch.bfloat16, device=device) + nop = max(N * H * nseg, 1) + buffers = ( + torch.empty((c * 16, 128), dtype=torch.bfloat16, device=device), + torch.empty((c * 16, 128), dtype=torch.bfloat16, device=device), + kg, + ut, + torch.empty((c * 16, 16), dtype=torch.bfloat16, device=device), + torch.empty((c, 128), dtype=torch.float32, device=device), + kg.t(), + ut.t(), + torch.empty((nop, 128, 128), dtype=torch.float32, device=device), + torch.empty((nop, 128, 128), dtype=torch.bfloat16, device=device), + torch.empty((nop, 128, 128), dtype=torch.float32, device=device), + ) + w = (buffers, tuple(from_dlpack(b, assumed_align=16) for b in buffers)) + _WS[key] = w + return w[1] + + +def run(q, k, v, g, beta, cu_seqlens, initial_state, o, final_state): + """Launch from torch tensors (standalone / test entry point). + + Converts the nine operands to CuTe tensors and defers to :func:`run_cute`. + A caller that already holds CuTe tensors -- cuDNN's engine reaches them + straight off the variant pack -- should call that directly and skip the + torch round trip, which is two DLPack conversions per operand rather than + one. + """ + T, H, D = q.shape + N = initial_state.shape[0] + run_cute( + from_dlpack(q, assumed_align=16), + from_dlpack(k, assumed_align=16), + from_dlpack(v, assumed_align=16), + from_dlpack(g, assumed_align=16), + from_dlpack(beta, assumed_align=16), + from_dlpack(cu_seqlens, assumed_align=4), + from_dlpack(initial_state, assumed_align=16), + from_dlpack(o, assumed_align=16), + from_dlpack(final_state, assumed_align=16), + T, + H, + D, + N, + q.device, + torch.cuda.current_stream().cuda_stream, + ) + + +def run_cute(mQ, mK, mV, mG, mB, mCu, mIS, mO, mFS, T, H, D, N, device, stream_ptr): + """Launch from already-converted CuTe tensors on an explicit stream. + + ``stream_ptr`` is a raw CUDA stream handle, so the caller does not have to + push a torch stream context just so this function can read it back out of + thread-local state. + """ + nch = T // 16 + N + 1 + max_chunks = (T + 15) // 16 + nseg = _pick_nseg(T, N, H) + + # Already CuTe tensors, converted once per shape (see _ws). + ws = _ws(nch, H, N, nseg, device) + + key = (T, H, N, D) + fn = _CACHE.get(key) + stream = cudadrv.CUstream(stream_ptr) + + args = (mQ, mK, mV, mG, mB, mCu, mIS, *ws, mO, mFS) + if fn is None: + fn = cute.compile(kda_launch, *args, N, H, max_chunks, nch, nseg, stream) + _CACHE[key] = fn + fn(*args, stream) diff --git a/test/python/linear_attention/test_la.py b/test/python/linear_attention/test_la.py index 79f829aad..e4e19919a 100644 --- a/test/python/linear_attention/test_la.py +++ b/test/python/linear_attention/test_la.py @@ -145,10 +145,15 @@ def pinned_op(backend, variant): return functools.partial(op(variant), plan_name=backend.plan(variant)) -@pytest.fixture(params=("frost", "cutile")) +@pytest.fixture(params=("frost", "cutile", "hopper")) def backend(request): """One backend per run of each test; the tests pass its plan name to the - ops. The op graph caches are cleared around each test.""" + ops. The op graph caches are cleared around each test. + + ``hopper`` is the sm90 path and exists for KDA only, so every other variant + (and every architecture that is not Hopper) declines it and the test waives + itself through :func:`waive_unsupported` -- the same way ``cutile`` already + waives where it has no kernel.""" clear_caches() try: yield Backend(request.param) @@ -768,9 +773,12 @@ def test_bwd_split_initial_state(backend, variant): s0 = state0.detach().clone().requires_grad_(True) with waive_unsupported(backend, variant): o, _ = pinned_op(backend, variant)(*leaves, *op_tail(case), initial_state=s0, output_final_state=True, **kw) - if dO is None: - dO = torch.randn_like(o) - grads[tag] = torch.autograd.grad([o], leaves + [s0], [dO]) + if dO is None: + dO = torch.randn_like(o) + # Inside the waiver, matching test_bwd_split_d_final_state: a + # forward-only backend (hopper) serves the forward and declines the + # backward, which is a waive, not a failure. + grads[tag] = torch.autograd.grad([o], leaves + [s0], [dO]) for name, got, want in zip(list(tensors) + ["initial_state"], grads["split"], grads["uncut"]): assert_rms_close(f"d{name} split-vs-uncut", got, want.float(), BWD_TOL[torch.bfloat16])