From cd94f6e5a5ffe51412caaf90c6c2ec11a8f4bf43 Mon Sep 17 00:00:00 2001 From: Joshua Park Date: Mon, 17 Aug 2026 23:59:42 -0700 Subject: [PATCH 1/2] 1 --- python/cudnn/_pygraph.py | 7 +- python/cudnn/fla/gated_delta_rule.py | 4 +- python/cudnn/fla/kda.py | 4 +- .../linear_attention/cutile/gdn_engine.py | 7 +- .../linear_attention/cutile/kda_engine.py | 7 +- .../linear_attention/cutile/kernels/common.py | 29 -- .../linear_attention/frost/common/downcast.py | 155 ------- .../linear_attention/frost/common/thd.py | 19 +- .../linear_attention/frost/gdn2_engine.py | 22 +- .../linear_attention/frost/gdn_engine.py | 29 +- .../linear_attention/frost/kda_engine.py | 22 +- .../frost/kernel/gdn2_bprop_f16.py | 170 ++++---- .../frost/kernel/gdn2_prefill_f16.py | 126 ++++-- .../frost/kernel/gdn2_recompute_f16.py | 123 ++++-- .../frost/kernel/gdn_bprop_f16.py | 208 ++++------ .../frost/kernel/gdn_prefill_f16.py | 381 ++++++++++-------- .../frost/kernel/gdn_recompute_f16.py | 308 +++++++------- .../frost/kernel/kda_bprop_f16.py | 170 ++++---- .../frost/kernel/kda_prefill_f16.py | 124 ++++-- .../frost/kernel/kda_recompute_f16.py | 122 ++++-- .../cudnn/linear_attention/graph_analyzer.py | 13 + python/cudnn/linear_attention/ops/gdn.py | 24 +- python/cudnn/linear_attention/ops/gdn2.py | 24 +- python/cudnn/linear_attention/ops/kda.py | 24 +- .../frost/examples/01_gdn_prefill.py | 2 +- test/python/linear_attention/reference_gdn.py | 14 +- .../python/linear_attention/reference_gdn2.py | 16 +- test/python/linear_attention/reference_kda.py | 12 +- .../linear_attention/test_fla_compat.py | 34 ++ test/python/linear_attention/test_la.py | 253 ++++++++---- 30 files changed, 1252 insertions(+), 1201 deletions(-) delete mode 100644 python/cudnn/linear_attention/frost/common/downcast.py diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 1d5429359..fef68b781 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -2455,13 +2455,12 @@ def _moe_bwd_dweight_dims(node): def _linear_attention_final_state_dims(node): - # [N, HO, K, V]: N sequences (cu_seqlens carries N+1 boundaries), HO = - # max(q, v) heads — the recurrent state lives at the gate heads + # [N, HO, V, K] q, v = node.inputs["q"].dim, node.inputs["v"].dim cu = node.inputs.get("cu_seqlens") if cu is None or not cu.dim: return None - return [cu.dim[0] - 1, max(q[1], v[1]), q[2], v[2]] + return [cu.dim[0] - 1, max(q[1], v[1]), v[2], q[2]] def _linear_attention_state_checkpoints_dims(node): @@ -2470,7 +2469,7 @@ def _linear_attention_state_checkpoints_dims(node): cu = node.inputs["cu_seqlens"].dim if node.inputs.get("cu_seqlens") is not None else None if not n or not q or not v or not cu: return None - return [max(v[0] // n, 1), max(q[1], v[1]), q[2], v[2]] + return [max(v[0] // n + (cu[0] - 1), 1), max(q[1], v[1]), v[2], q[2]] def _linear_attention_o_dims(node): diff --git a/python/cudnn/fla/gated_delta_rule.py b/python/cudnn/fla/gated_delta_rule.py index d438d0d96..11ace7f7a 100644 --- a/python/cudnn/fla/gated_delta_rule.py +++ b/python/cudnn/fla/gated_delta_rule.py @@ -172,8 +172,8 @@ def fallback(reason): return fallback("variant") # state_v_first only changes the recurrent-state layout, so it is a no-op # for a stateless (training) call; decline only when a state is exchanged. - if state_v_first and (initial_state is not None or output_final_state): - return fallback("state_v_first") + if not state_v_first and (initial_state is not None or output_final_state): + return fallback("state_v_first=False") if not (q.is_cuda and torch.cuda.get_device_capability(q.device)[0] >= 10): return fallback("pre-Blackwell") try: diff --git a/python/cudnn/fla/kda.py b/python/cudnn/fla/kda.py index a8c90708f..398310ee7 100644 --- a/python/cudnn/fla/kda.py +++ b/python/cudnn/fla/kda.py @@ -183,8 +183,8 @@ def fallback(reason): if allow_neg_eigval or cp_context is not None or return_intermediate_states: return fallback("variant") - if state_v_first and (initial_state is not None or output_final_state): - return fallback("state_v_first") + if not state_v_first and (initial_state is not None or output_final_state): + return fallback("state_v_first=False") if not (q.is_cuda and torch.cuda.get_device_capability(q.device)[0] >= 10): return fallback("pre-Blackwell") try: diff --git a/python/cudnn/linear_attention/cutile/gdn_engine.py b/python/cudnn/linear_attention/cutile/gdn_engine.py index ce3013dbb..9980868ed 100644 --- a/python/cudnn/linear_attention/cutile/gdn_engine.py +++ b/python/cudnn/linear_attention/cutile/gdn_engine.py @@ -57,7 +57,7 @@ def __init__(self, graph): ("A", layout.add(total * HV * BT * isz), io, (total, HV, BT)), ("w", layout.add(total * HV * K * isz), io, (total, HV, K)), ("u", layout.add(total * HV * V * isz), io, (total, HV, V)), - ("state_checkpoints", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, K, V)), + ("state_checkpoints", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, V, K)), ("v_new", layout.add(total * HV * V * isz), io, (total, HV, V)), ] if l2norm: @@ -71,7 +71,7 @@ def __init__(self, graph): NK = common.cdiv(K, min(max(common.next_power_of_2(K), 16), 64)) regions += [ ("dv", layout.add(total * HV * V * isz), io, (total, HV, V)), - ("dstate", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, K, V)), + ("dstate", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, V, K)), ("dv2", layout.add(total * HV * V * isz), io, (total, HV, V)), ("dg_nk", layout.add(NK * total * HV * 4), f32, (NK, total, HV)), ("dw", layout.add(total * HV * K * isz), io, (total, HV, K)), @@ -132,7 +132,6 @@ def execute(self, graph, variant_pack, ctx) -> None: check_layouts_compact(self.plan_name, self.expect, self.names, views) nb = dict(zip(self.names, views)) stream = ctx.stream if ctx.stream is not None else 0 - self.common.ensure_cuda_context(stream) ws = Workspace.over(variant_pack, self.ws_bytes, self.plan_name) region = dict(zip(self.carve_names, ws.carve(self.carve))) self.common.build_chunk_table( @@ -167,6 +166,7 @@ def execute_fwd(self, nb, region, stream) -> None: cu_seqlens=nb["cu_seqlens"], chunk_indices=region["chunk_table"], bufs=region, + state_v_first=True, stream=stream, **gate, ) @@ -186,6 +186,7 @@ def execute_bwd(self, nb, region, stream) -> None: cu_seqlens=nb["cu_seqlens"], chunk_indices=region["chunk_table"], bufs=region, + state_v_first=True, stream=stream, ) diff --git a/python/cudnn/linear_attention/cutile/kda_engine.py b/python/cudnn/linear_attention/cutile/kda_engine.py index 6e844b58a..9b8f60d3e 100644 --- a/python/cudnn/linear_attention/cutile/kda_engine.py +++ b/python/cudnn/linear_attention/cutile/kda_engine.py @@ -64,7 +64,7 @@ def __init__(self, graph): ("u", layout.add(total * HV * V * isz), io, (total, HV, V)), ("qg", layout.add(total * HV * K * isz), io, (total, HV, K)), ("kg", layout.add(total * HV * K * isz), io, (total, HV, K)), - ("state_checkpoints", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, K, V)), + ("state_checkpoints", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, V, K)), ("v_new", layout.add(total * HV * V * isz), io, (total, HV, V)), ] if node.params.get("use_beta_sigmoid", False): @@ -82,7 +82,7 @@ def __init__(self, graph): regions += [ ("dAqk", layout.add(total * HV * BT * 4), f32, (total, HV, BT)), ("dv_dAv", layout.add(total * HV * V * isz), io, (total, HV, V)), - ("dstate", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, K, V)), + ("dstate", layout.add(NT_bound * HV * K * V * isz), io, (NT_bound, HV, V, K)), ("dv_dstate_u", layout.add(total * HV * V * isz), io, (total, HV, V)), ("dq", layout.add(total * HV * K * 4), f32, (total, HV, K)), ("dk", layout.add(total * HV * K * 4), f32, (total, HV, K)), @@ -158,7 +158,6 @@ def execute(self, graph, variant_pack, ctx) -> None: check_layouts_compact(self.plan_name, self.expect, self.names, views) nb = dict(zip(self.names, views)) stream = ctx.stream if ctx.stream is not None else 0 - self.common.ensure_cuda_context(stream) ws = Workspace.over(variant_pack, self.ws_bytes, self.plan_name) region = dict(zip(self.carve_names, ws.carve(self.carve))) self.common.build_chunk_table( @@ -197,6 +196,7 @@ def execute_fwd(self, nb, region, stream) -> None: cu_seqlens=nb["cu_seqlens"], chunk_indices=region["chunk_table"], bufs=region, + state_v_first=True, stream=stream, **gate, ) @@ -216,6 +216,7 @@ def execute_bwd(self, nb, region, stream) -> None: cu_seqlens=nb["cu_seqlens"], chunk_indices=region["chunk_table"], bufs=region, + state_v_first=True, stream=stream, ) diff --git a/python/cudnn/linear_attention/cutile/kernels/common.py b/python/cudnn/linear_attention/cutile/kernels/common.py index 3f24ebccf..f7cf25366 100644 --- a/python/cudnn/linear_attention/cutile/kernels/common.py +++ b/python/cudnn/linear_attention/cutile/kernels/common.py @@ -63,35 +63,6 @@ def opt(t, bufs, dtype_name: str = "float32"): return t -def ensure_cuda_context(stream=0) -> None: - """Bind a driver context to the calling thread when none is bound. - - ``ct.launch`` and the autotuner read the calling thread's context stack, - and an autograd backward runs on a worker thread where ``cudaSetDevice`` - has only moved the runtime's thread-local slot. Prefer the launch stream's - context, else retain the device's primary one. Best-effort: a context this - cannot establish fails at the launch, with the launch's own diagnostics.""" - try: - from cuda.bindings import driver as drv - - err, cur = drv.cuCtxGetCurrent() - if err == drv.CUresult.CUDA_SUCCESS and int(cur) != 0: - return - if stream: - err, sctx = drv.cuStreamGetCtx(stream) - if err == drv.CUresult.CUDA_SUCCESS: - drv.cuCtxSetCurrent(sctx) - return - device = current_device_id() - if device is None: - return - err, pctx = drv.cuDevicePrimaryCtxRetain(device) - if err == drv.CUresult.CUDA_SUCCESS: - drv.cuCtxSetCurrent(pctx) - except Exception: # noqa: BLE001 - pass - - # --- Launch tuning -------------------------------------------------------------------------------- diff --git a/python/cudnn/linear_attention/frost/common/downcast.py b/python/cudnn/linear_attention/frost/common/downcast.py deleted file mode 100644 index 0bb77d1c4..000000000 --- a/python/cudnn/linear_attention/frost/common/downcast.py +++ /dev/null @@ -1,155 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Initial-state staging for the FROST LA backward kernels: copy the caller's -``[N, HO, K, V]`` state (fp32 or io dtype, padded outer strides fine) into -the compact io-dtype buffer the per-(b,h) state descriptors read. - -Alignment is the caller's contract, as everywhere TMA is involved: 16-byte -aligned bases, and outer strides that keep every 8-element V chunk address -16-byte aligned (compact buffers trivially qualify).""" - -import functools - -import cuda.bindings.driver as cuda -import cutlass -import cutlass.cute as cute -from cutlass.cute.arch.nvvm_wrappers import inline_ptx -from cutlass.cute.runtime import from_dlpack - -from cudnn.frost.tile_dsl.pointwise import f16x2_to_f32, fp32_to_fp16 - - -@cute.kernel -def downcast_state_kernel( - mState0: cute.Tensor, - mOut: cute.Tensor, - n_k: cutlass.Int32, - threads_per_row: cutlass.Int32, - rows_per_cta: cutlass.Int32, -) -> None: - """Vectorized copy of the ``[N, HO, K, V]`` initial state into the - io-dtype buffer the backward's static state descriptor reads: grid - (K-tiles, HO, N), one 8-element V chunk per thread as 128-bit loads and - one 128-bit store, source read through its (dynamic) strides so padded - outer layouts stage zero-copy. fp32 sources convert through packed - ``cvt.rn.{f16,bf16}x2.f32``; same-dtype io sources copy words verbatim; - a 16-bit cross convert unpacks to fp32 pairs and repacks.""" - bid = cute.arch.block_idx() - tidx = cutlass.Int32(cute.arch.thread_idx()[0]) - k_idx = cutlass.Int32(bid[0]) * rows_per_cta + tidx // threads_per_row - v0 = (tidx % threads_per_row) * cutlass.Int32(8) - n_idx = cutlass.Int32(bid[2]) - h_idx = cutlass.Int32(bid[1]) - if k_idx < n_k: - src_elems = ( - cutlass.Int64(n_idx) * cutlass.Int64(mState0.stride[0]) - + cutlass.Int64(h_idx) * cutlass.Int64(mState0.stride[1]) - + cutlass.Int64(k_idx) * cutlass.Int64(mState0.stride[2]) - + cutlass.Int64(v0) - ) - dst_elems = ( - cutlass.Int64(n_idx) * cutlass.Int64(mOut.stride[0]) - + cutlass.Int64(h_idx) * cutlass.Int64(mOut.stride[1]) - + cutlass.Int64(k_idx) * cutlass.Int64(mOut.stride[2]) - + cutlass.Int64(v0) - ) - dst_addr = mOut.iterator.toint() + dst_elems * cutlass.Int64(2) - if cutlass.const_expr(mState0.element_type == cutlass.Float32): - src_addr = mState0.iterator.toint() + src_elems * cutlass.Int64(4) - f0, f1, f2, f3 = inline_ptx( - "ld.global.v4.f32 {$0, $1, $2, $3}, [$4];", - write_only_types=[cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Float32], - read_only_args=[src_addr], - ) - f4, f5, f6, f7 = inline_ptx( - "ld.global.v4.f32 {$0, $1, $2, $3}, [$4];", - write_only_types=[cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Float32], - read_only_args=[src_addr + cutlass.Int64(16)], - ) - w0 = fp32_to_fp16(f0, f1, dtype=mOut.element_type) - w1 = fp32_to_fp16(f2, f3, dtype=mOut.element_type) - w2 = fp32_to_fp16(f4, f5, dtype=mOut.element_type) - w3 = fp32_to_fp16(f6, f7, dtype=mOut.element_type) - else: - src_addr = mState0.iterator.toint() + src_elems * cutlass.Int64(2) - w0, w1, w2, w3 = inline_ptx( - "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", - write_only_types=[cutlass.Int32, cutlass.Int32, cutlass.Int32, cutlass.Int32], - read_only_args=[src_addr], - ) - if cutlass.const_expr(mState0.element_type != mOut.element_type): - lo0, hi0 = f16x2_to_f32(w0, dtype=mState0.element_type) - lo1, hi1 = f16x2_to_f32(w1, dtype=mState0.element_type) - lo2, hi2 = f16x2_to_f32(w2, dtype=mState0.element_type) - lo3, hi3 = f16x2_to_f32(w3, dtype=mState0.element_type) - w0 = fp32_to_fp16(lo0, hi0, dtype=mOut.element_type) - w1 = fp32_to_fp16(lo1, hi1, dtype=mOut.element_type) - w2 = fp32_to_fp16(lo2, hi2, dtype=mOut.element_type) - w3 = fp32_to_fp16(lo3, hi3, dtype=mOut.element_type) - inline_ptx( - "st.global.v4.b32 [$0], {$1, $2, $3, $4};", - read_only_args=[dst_addr, w0, w1, w2, w3], - ) - - -@cute.jit -def downcast_state_launch( - state0: cute.Tensor, - out: cute.Tensor, - n_k: cutlass.Int32, - threads_per_row: cutlass.Int32, - rows_per_cta: cutlass.Int32, - n_blocks: cutlass.Int32, - ho: cutlass.Int32, - n_seq: cutlass.Int32, - stream: cuda.CUstream, -): - downcast_state_kernel( - state0, - out, - n_k, - threads_per_row, - rows_per_cta, - ).launch(grid=(n_blocks, ho, n_seq), block=(128, 1, 1), stream=stream) - - -@functools.cache -def downcast_state_cache(key): - return {} - - -def downcast_state(initial_state, out, *, stream): - """Copy the initial state ``[N, HO, K, V]`` (fp32 or io dtype, padded - outer strides fine) into ``out`` (io dtype, same shape, compact) — the - buffer the backward's per-(b,h) state descriptors read.""" - if tuple(int(s_) for s_ in initial_state.shape) != tuple(int(s_) for s_ in out.shape): - raise ValueError(f"initial_state must match the io state buffer shape {tuple(out.shape)}; got {tuple(initial_state.shape)}") - n_seq, ho, k, v = (int(s_) for s_ in out.shape) - if v % 8 != 0: - raise ValueError(f"state V dim must be a multiple of 8 (8-element staging chunks); got {v}") - if v > 1024: - raise ValueError(f"state V dim must be <= 1024 (one 128-thread block stages a full row); got {v}") - threads_per_row = v // 8 - rows_per_cta = max(128 // threads_per_row, 1) - n_blocks = (k + rows_per_cta - 1) // rows_per_cta - key = (str(initial_state.dtype), str(out.dtype)) - cache = downcast_state_cache(key) - cu_stream = cuda.CUstream(int(stream)) - if "compiled" not in cache: - state0_c = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) - out_c = from_dlpack(out, assumed_align=16).mark_layout_dynamic(leading_dim=3) - cache["compiled"] = cute.compile( - downcast_state_launch, - state0_c, - out_c, - cutlass.Int32(k), - cutlass.Int32(threads_per_row), - cutlass.Int32(rows_per_cta), - cutlass.Int32(n_blocks), - cutlass.Int32(ho), - cutlass.Int32(n_seq), - cu_stream, - options="--enable-tvm-ffi", - ) - cache["compiled"](initial_state, out, k, threads_per_row, rows_per_cta, n_blocks, ho, n_seq, cu_stream) diff --git a/python/cudnn/linear_attention/frost/common/thd.py b/python/cudnn/linear_attention/frost/common/thd.py index e69ce580b..0ad38786e 100644 --- a/python/cudnn/linear_attention/frost/common/thd.py +++ b/python/cudnn/linear_attention/frost/common/thd.py @@ -12,9 +12,6 @@ * :func:`emit_checkpoint_seq_descs` — its per-chunk-checkpoint sibling; derives the per-sequence checkpoint offsets from the token ``cu_seqlens`` in place of a caller-computed prefix array. -* :func:`emit_copy_desc` — verbatim single-slot copy of a fully static - descriptor (dense ``[N, HO, K, V]`` state; batch and head are both load - coordinates). """ import cutlass @@ -83,7 +80,7 @@ def emit_checkpoint_seq_descs( """Per-BATCH descriptor array for the per-chunk checkpoint tensor with the head axis as a descriptor dimension (``(dv, dk, chunk, head)``). Derives the per-sequence checkpoint offsets from the TOKEN ``cu_seqlens`` on the fly - (``count_b = (seqlen_b - 1) // every_n``, running-prefix-summed) — an + (``count_b = (seqlen_b - 1) // every_n + 1``, running-prefix-summed) — an address fold no coordinate transform can express — and caps GLOBAL_DIM[``seq_ord``] to ``count_b``. The head index is a load coordinate. Runs on one electing thread; the calling warp elects and @@ -95,7 +92,7 @@ def emit_checkpoint_seq_descs( run = cutlass.Int32(0) for b in cutlass.range(0, n_batch, 1, unroll=1): s_tok = cutlass.Int32(cu[b + cutlass.Int32(1)]) - cutlass.Int32(cu[b]) - cnt = (s_tok - cutlass.Int32(1)) // every_n + cnt = (s_tok - cutlass.Int32(1)) // every_n + cutlass.Int32(1) cnt = cnt if s_tok > 0 else cutlass.Int32(0) checkpoint_base = run run = run + cnt @@ -114,15 +111,3 @@ def emit_checkpoint_seq_descs( new_value=cnt, ord=seq_ord, ) - - -@cute.jit -def emit_copy_desc(base_desc, desc_words) -> None: - """Verbatim single-slot copy of a fully static descriptor (e.g. the - dense ``[N, HO, K, V]`` initial state, whose batch and head are both - load coordinates). Runs on one electing thread; the calling warp - elects and fences.""" - desc_base = desc_words.iterator.raw_ptr() - src_words = Pointer(base_desc.get_ptr(), dtype=cutlass.Int64) - for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS): - (desc_base + i).store((src_words + i).load()) diff --git a/python/cudnn/linear_attention/frost/gdn2_engine.py b/python/cudnn/linear_attention/frost/gdn2_engine.py index 6ace8a082..0b7067fb1 100644 --- a/python/cudnn/linear_attention/frost/gdn2_engine.py +++ b/python/cudnn/linear_attention/frost/gdn2_engine.py @@ -130,8 +130,7 @@ def __init__(self, node, kernel_mod): q, g = node.inputs["q"], node.inputs["g"] self.b_t = kernel_mod.CFG.B_T - # cuts only for chunk-granular checkpoint cadences, never in batch-invariant mode - self.split = self.ckpt in (0, self.b_t) and not self.batch_invariant + self.split = self.ckpt % self.b_t == 0 and not self.batch_invariant total = q.dim[0] HO = g.dim[1] B = node.inputs["cu_seqlens"].dim[0] - 1 @@ -314,12 +313,10 @@ def __init__(self, node, bwd_mod, regen_mod): self.kcache = None self.regen_cache = None self.plan_name = "Gdn2FrostEngine (GDN2_BWD)" - from .common.downcast import downcast_state from .common.gate_bwd import GATE_BWD_BLOCKS, channel_gate_bwd from .common.head_reduce import head_group_reduce from .common.host import tensormap_workspace_bytes - self.downcast_state = downcast_state self.head_group_reduce = head_group_reduce self.channel_gate_bwd = channel_gate_bwd scale = node.params.get("scale") @@ -331,7 +328,6 @@ def __init__(self, node, bwd_mod, regen_mod): self.gate_lower_bound = float(glb) if glb is not None else bwd_mod.DEFAULT_GATE_LOWER_BOUND self.gate_bwd_blocks = GATE_BWD_BLOCKS self.has_state_checkpoints = "state_checkpoints" in node.inputs - self.has_state0 = "initial_state" in node.inputs self.has_dstate0 = "d_initial_state" in node.outputs q, g, v = node.inputs["q"], node.inputs["g"], node.inputs["v"] @@ -363,7 +359,6 @@ def __init__(self, node, bwd_mod, regen_mod): self.off_item_scratch = layout.add(self.work_item_rows * WORK_ITEM_FIELDS * 4) self.chunk_scratch_rows = chunk_scratch_rows(total, B, self.b_t) self.off_chunk_scratch = layout.add(self.chunk_scratch_rows * HO * 4) - self.off_state0_io = layout.add(B * HO * K * V * 2) if self.has_state0 else None if not self.has_state_checkpoints: self.state_checkpoints_rows = max(total // self.b_t + B, 1) self.off_state_checkpoints = layout.add(self.state_checkpoints_rows * HO * K * V * 2) @@ -398,10 +393,8 @@ def __init__(self, node, bwd_mod, regen_mod): if self.split: regions.append(("item_scratch", self.off_item_scratch, "int32", (self.work_item_rows, WORK_ITEM_FIELDS))) regions.append(("chunk_scratch", self.off_chunk_scratch, "float32", (self.chunk_scratch_rows, HO))) - if self.has_state0: - regions.append(("state0_io", self.off_state0_io, self.io_name, (B, HO, K, V))) if not self.has_state_checkpoints: - regions.append(("state_checkpoints", self.off_state_checkpoints, self.io_name, (self.state_checkpoints_rows, HO, K, V))) + regions.append(("state_checkpoints", self.off_state_checkpoints, self.io_name, (self.state_checkpoints_rows, HO, V, K))) regions.append(("regen_tensormaps", self.off_regen_tensormaps, "int64", (self.regen_tm_bytes // 8,))) if self.fold_dq: regions.append(("dq_ho", self.off_dq_ho, self.io_name, (total, HO, K))) @@ -489,10 +482,6 @@ def run(self, views, workspace, stream) -> None: region["sched_all"], stream, ) - state0_io = None - if state0 is not None: - state0_io = region["state0_io"] - self.downcast_state(state0, state0_io, stream=stream) if self.has_state_checkpoints: checkpoint_series = state_checkpoints else: @@ -539,7 +528,6 @@ def run(self, views, workspace, stream) -> None: db, dw, cu, - state0_io, dstate0 if self.has_dstate0 else None, dstate_in, work_items, @@ -587,10 +575,6 @@ def run(self, views, workspace, stream) -> None: stream=stream, ) - state0_io = None - if state0 is not None: - state0_io = region["state0_io"] - self.downcast_state(state0, state0_io, stream=stream) if self.has_state_checkpoints: checkpoint_series = state_checkpoints else: @@ -647,7 +631,7 @@ def run(self, views, workspace, stream) -> None: dw, cu, self.scale, - initial_state=state0_io, + use_initial_state=state0 is not None, d_initial_state=dstate0 if self.has_dstate0 else None, d_final_state=dstate_in, use_qk_l2norm_in_kernel=self.use_qk_l2norm, diff --git a/python/cudnn/linear_attention/frost/gdn_engine.py b/python/cudnn/linear_attention/frost/gdn_engine.py index c7747a0f2..197f93d25 100644 --- a/python/cudnn/linear_attention/frost/gdn_engine.py +++ b/python/cudnn/linear_attention/frost/gdn_engine.py @@ -14,7 +14,6 @@ from cudnn import behavior_note from cudnn.engines.base import BaseEngine, CompiledPlan -from cudnn.frost import buffers from cudnn.frost.device import build_device, current_device, multiprocessor_count from cudnn.frost.workspace import WorkspaceLayout, carve_plan from ..graph_analyzer import analyze @@ -134,8 +133,7 @@ def __init__(self, node, kernel_mod): self.ckpt = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) self.has_state_checkpoints = "state_checkpoints" in node.outputs self.batch_invariant = bool(node.params.get("batch_invariant", False)) - # cuts only for chunk-granular checkpoint cadences, never in batch-invariant mode - self.split = self.ckpt in (0, self.b_t) and not self.batch_invariant + self.split = self.ckpt % self.b_t == 0 and not self.batch_invariant layout = WorkspaceLayout() from .common.host import tensormap_workspace_bytes @@ -326,13 +324,11 @@ def __init__(self, node, bwd_mod, regen_mod): self.kcache = None self.regen_cache = None self.plan_name = "GdnFrostEngine (GDN_BWD)" - from .common.downcast import downcast_state from .common.gate_bwd import scalar_gate_bwd, scalar_gate_blocks from .common.head_reduce import head_group_reduce from .common.host import tensormap_workspace_bytes from .common.l2norm import l2norm_qk, l2norm_qk_bwd - self.downcast_state = downcast_state self.head_group_reduce = head_group_reduce self.scalar_gate_bwd = scalar_gate_bwd self.l2norm_qk = l2norm_qk @@ -352,7 +348,6 @@ def __init__(self, node, bwd_mod, regen_mod): HO = g.dim[1] B = node.inputs["cu_seqlens"].dim[0] - 1 self.has_state_checkpoints = "state_checkpoints" in node.inputs - self.has_state0 = "initial_state" in node.inputs self.io_name = "float16" if node.inputs["q"].get_data_type().name == "HALF" else "bfloat16" self.num_sm = multiprocessor_count(current_device()) @@ -364,7 +359,6 @@ def __init__(self, node, bwd_mod, regen_mod): self.off_sched = layout.add(16) self.tensormap_words = tensormap_workspace_bytes(bwd_mod, B) // 8 self.off_tensormaps = layout.add(self.tensormap_words * 8) - self.off_state0_io = layout.add(B * HO * K * V * 2) if self.has_state0 else None if not self.has_state_checkpoints: self.state_checkpoints_rows = max(total // self.b_t + B, 1) self.off_state_checkpoints = layout.add(self.state_checkpoints_rows * HO * K * V * 2) @@ -417,10 +411,8 @@ def __init__(self, node, bwd_mod, regen_mod): if self.split: regions.append(("item_scratch", self.off_item_scratch, "int32", (self.work_item_rows, WORK_ITEM_FIELDS))) regions.append(("chunk_scratch", self.off_chunk_scratch, "float32", (self.chunk_scratch_rows, HO))) - if self.has_state0: - regions.append(("state0_io", self.off_state0_io, self.io_name, (B, HO, K, V))) if not self.has_state_checkpoints: - regions.append(("state_checkpoints", self.off_state_checkpoints, self.io_name, (self.state_checkpoints_rows, HO, K, V))) + regions.append(("state_checkpoints", self.off_state_checkpoints, self.io_name, (self.state_checkpoints_rows, HO, V, K))) regions.append(("regen_tensormaps", self.off_regen_tensormaps, "int64", (self.regen_tensormap_words,))) if self.fold_dq: regions.append(("dq_ho", self.off_dq_ho, self.io_name, (total, HO, K))) @@ -512,13 +504,6 @@ def run(self, views, workspace, stream) -> None: region["sched_all"], stream, ) - state0_io = None - if state0 is not None: - if state0.dtype == self.io_name and buffers.is_contiguous(tuple(state0.shape), state0.stride()) and state0.data_ptr() % 16 == 0: - state0_io = state0 - else: - state0_io = region["state0_io"] - self.downcast_state(state0, state0_io, stream=stream) if self.has_state_checkpoints: checkpoint_series = state_checkpoints else: @@ -562,7 +547,6 @@ def run(self, views, workspace, stream) -> None: dg, db, cu, - state0_io, dstate0, dstate_in, work_items, @@ -609,13 +593,6 @@ def run(self, views, workspace, stream) -> None: stream=stream, ) - state0_io = None - if state0 is not None: - if state0.dtype == self.io_name and buffers.is_contiguous(tuple(state0.shape), state0.stride()) and state0.data_ptr() % 16 == 0: - state0_io = state0 - else: - state0_io = region["state0_io"] - self.downcast_state(state0, state0_io, stream=stream) if self.has_state_checkpoints: checkpoint_series = state_checkpoints else: @@ -667,7 +644,7 @@ def run(self, views, workspace, stream) -> None: db, cu, self.scale, - initial_state=state0_io, + use_initial_state=state0 is not None, d_initial_state=dstate0, d_final_state=dstate_in, safe_gate=self.safe_gate, diff --git a/python/cudnn/linear_attention/frost/kda_engine.py b/python/cudnn/linear_attention/frost/kda_engine.py index f2e95ab16..e661f0fda 100644 --- a/python/cudnn/linear_attention/frost/kda_engine.py +++ b/python/cudnn/linear_attention/frost/kda_engine.py @@ -128,8 +128,7 @@ def __init__(self, node, kernel_mod): q, g = node.inputs["q"], node.inputs["g"] self.b_t = kernel_mod.CFG.B_T - # cuts only for chunk-granular checkpoint cadences, never in batch-invariant mode - self.split = self.ckpt in (0, self.b_t) and not self.batch_invariant + self.split = self.ckpt % self.b_t == 0 and not self.batch_invariant total = q.dim[0] HO = g.dim[1] B = node.inputs["cu_seqlens"].dim[0] - 1 @@ -310,12 +309,10 @@ def __init__(self, node, bwd_mod, regen_mod): self.kcache = None self.regen_cache = None self.plan_name = "KdaFrostEngine (KDA_BWD)" - from .common.downcast import downcast_state from .common.gate_bwd import GATE_BWD_BLOCKS, channel_gate_bwd from .common.head_reduce import head_group_reduce from .common.host import tensormap_workspace_bytes - self.downcast_state = downcast_state self.head_group_reduce = head_group_reduce self.channel_gate_bwd = channel_gate_bwd scale = node.params.get("scale") @@ -327,7 +324,6 @@ def __init__(self, node, bwd_mod, regen_mod): self.gate_lower_bound = float(glb) if glb is not None else bwd_mod.DEFAULT_GATE_LOWER_BOUND self.gate_bwd_blocks = GATE_BWD_BLOCKS self.has_state_checkpoints = "state_checkpoints" in node.inputs - self.has_state0 = "initial_state" in node.inputs self.has_dstate0 = "d_initial_state" in node.outputs q, g, v = node.inputs["q"], node.inputs["g"], node.inputs["v"] @@ -359,7 +355,6 @@ def __init__(self, node, bwd_mod, regen_mod): self.off_item_scratch = layout.add(self.work_item_rows * WORK_ITEM_FIELDS * 4) self.chunk_scratch_rows = chunk_scratch_rows(total, B, self.b_t) self.off_chunk_scratch = layout.add(self.chunk_scratch_rows * HO * 4) - self.off_state0_io = layout.add(B * HO * K * V * 2) if self.has_state0 else None if not self.has_state_checkpoints: self.state_checkpoints_rows = max(total // self.b_t + B, 1) self.off_state_checkpoints = layout.add(self.state_checkpoints_rows * HO * K * V * 2) @@ -393,10 +388,8 @@ def __init__(self, node, bwd_mod, regen_mod): if self.split: regions.append(("item_scratch", self.off_item_scratch, "int32", (self.work_item_rows, WORK_ITEM_FIELDS))) regions.append(("chunk_scratch", self.off_chunk_scratch, "float32", (self.chunk_scratch_rows, HO))) - if self.has_state0: - regions.append(("state0_io", self.off_state0_io, self.io_name, (B, HO, K, V))) if not self.has_state_checkpoints: - regions.append(("state_checkpoints", self.off_state_checkpoints, self.io_name, (self.state_checkpoints_rows, HO, K, V))) + regions.append(("state_checkpoints", self.off_state_checkpoints, self.io_name, (self.state_checkpoints_rows, HO, V, K))) regions.append(("regen_tensormaps", self.off_regen_tensormaps, "int64", (self.regen_tm_bytes // 8,))) if self.fold_dq: regions.append(("dq_ho", self.off_dq_ho, self.io_name, (total, HO, K))) @@ -480,10 +473,6 @@ def run(self, views, workspace, stream) -> None: region["sched_all"], stream, ) - state0_io = None - if state0 is not None: - state0_io = region["state0_io"] - self.downcast_state(state0, state0_io, stream=stream) if self.has_state_checkpoints: checkpoint_series = state_checkpoints else: @@ -527,7 +516,6 @@ def run(self, views, workspace, stream) -> None: dg, db, cu, - state0_io, dstate0 if self.has_dstate0 else None, dstate_in, work_items, @@ -575,10 +563,6 @@ def run(self, views, workspace, stream) -> None: stream=stream, ) - state0_io = None - if state0 is not None: - state0_io = region["state0_io"] - self.downcast_state(state0, state0_io, stream=stream) if self.has_state_checkpoints: checkpoint_series = state_checkpoints else: @@ -632,7 +616,7 @@ def run(self, views, workspace, stream) -> None: db, cu, self.scale, - initial_state=state0_io, + use_initial_state=state0 is not None, d_initial_state=dstate0 if self.has_dstate0 else None, d_final_state=dstate_in, use_qk_l2norm_in_kernel=self.use_qk_l2norm, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py index 6f29deec3..be5552bd9 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py @@ -32,12 +32,12 @@ dGate = suffix-sum(dGate + dGate_last at row BT-1) (WG2 in-register reverse cumsum) dH <- diag-GEMM(eGl).dH + (scale.Q_decay)^T@dO - K_decay^T@dY -ABI: state_checkpoints `[total_checkpoints, HO, DK, DV]` (KV, v contiguous - the GDN checkpoint layout) io -dtype, the plain per-chunk series with NO initial-state slot (entry `c -- 1` = state entering chunk c >= 1; chunk 0 seeds from `initial_state`); beta `[T, HO, DK]` / w `[T, HO, DV]` io dtype +ABI: state_checkpoints `[total_checkpoints, HO, DV, DK]` (VK, k contiguous - the GDN checkpoint layout) io +dtype, the plain per-chunk series (entry `c` = state entering chunk c, so row 0 +is the initial state or zeros); beta `[T, HO, DK]` / w `[T, HO, DV]` io dtype (post-sigmoid); dq/dk/dv io at HO heads; dgate `[T, HO, DK]` fp32 (natural-log gate domain); dbeta/dw io dtype like beta/w; d_initial_state / d_final_state -fp32 `[N, HO, DK, DV]` (K-major, matching the prefill states). +fp32 `[N, HO, DV, DK]` (V-major, matching the prefill states). Warp assignments (16 warps = 512 threads): warps 0-3 : WG0 - Gate prefix scan + decay/restore operands (all chunks) @@ -64,7 +64,7 @@ from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr from cudnn.frost.device import current_device, multiprocessor_count -from ..common.thd import TENSOR_MAP_QWORDS, emit_copy_desc, emit_checkpoint_seq_descs, emit_seq_descs +from ..common.thd import TENSOR_MAP_QWORDS, emit_checkpoint_seq_descs, emit_seq_descs from .gdn2_bprop_config import CFG from cudnn.frost.tile_dsl.barrier import ( @@ -747,7 +747,6 @@ def super_mma_warp( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) sk_nt = cend - wstart @@ -928,6 +927,7 @@ def tcgen05_mma_warp( sSched, tmem_hold, sState_alt, + sState_direct, sK_decay_lead16, sK_inv_amaj, sK_restore_lead16, @@ -981,6 +981,26 @@ def tcgen05_mma_warp( idesc=idesc_state_k_at, kind=nvvm.Tcgen05MMAKind.F16, ) + idesc_state_k_kmaj = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.b_t, + m_dim=cfg.d_v, + ) + bmm_state_k_kmaj_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.d_k, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + atranspose=False, + cta_group=1, + idesc=idesc_state_k_kmaj, + kind=nvvm.Tcgen05MMAKind.F16, + ) bmm_dvinter_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, @@ -1251,13 +1271,13 @@ def tcgen05_mma_warp( d_dy_lead0 = sDy_lead16[0].desc() assert cfg.smem_state_stages == 1 d_state_alt0 = sState_alt[0].desc() + d_state_direct0 = sState_direct[0].desc() dstate0_index = PipelineState.start(phase=0) gbase = cutlass.Int32(0) sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) sk_nt = cend - wstart @@ -1290,8 +1310,8 @@ def tcgen05_mma_warp( if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_state_ready[state_index.idx].wait(state_index.phase) mma_ss( - bmm_state_k_desc, - d_state_alt0, + bmm_state_k_kmaj_desc, + d_state_direct0, d_kd_lead0 + decay_op_off, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_k_acc_offset), cutlass.Float32), accumulate=False, @@ -1589,7 +1609,6 @@ def tmaldg_warp( desc_beta_base, desc_w_base, desc_checkpoint_base, - desc_initial_state_base, bars, ) -> None: """TMA-LDG warp role (warp 14): persistent scheduler loop issuing every @@ -1691,7 +1710,6 @@ def tmaldg_warp( sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) next_tile, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) @@ -1708,7 +1726,6 @@ def tmaldg_warp( desc_beta_slot = (desc_beta_base + slot).tospace(cutlass.AddressSpace.generic) desc_w_slot = (desc_w_base + slot).tospace(cutlass.AddressSpace.generic) desc_checkpoint_slot = (desc_checkpoint_base + slot).tospace(cutlass.AddressSpace.generic) - desc_initial_state_slot = (desc_initial_state_base + cutlass.Int32(0)).tospace(cutlass.AddressSpace.generic) if elect_one: tma_tensormap_acquire(desc_q_slot) tma_tensormap_acquire(desc_k_slot) @@ -1718,8 +1735,6 @@ def tmaldg_warp( tma_tensormap_acquire(desc_beta_slot) tma_tensormap_acquire(desc_w_slot) tma_tensormap_acquire(desc_checkpoint_slot) - if cutlass.const_expr(cfg.use_initial_state): - tma_tensormap_acquire(desc_initial_state_slot) sk_nt = cend - wstart for rev_idx in cutlass.range(sk_nt, unroll=1): chunk_idx = cend - cutlass.Int32(1) - rev_idx @@ -1753,7 +1768,7 @@ def tmaldg_warp( beta_slice = tma_slice_runtime_desc(desc_beta_slot, cutlass.Int32(0), head_o, chunk_start) tma_load_tile(sBeta_tma[raw_index.idx], beta_slice, bars.mb_beta_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- entering state: checkpoint[c - 1], or initial_state for chunk 0 when given -- + # ---- entering state -- if chunk_idx >= FIRST_STATE_CHUNK: state_idx = state_index.idx bars.mb_state_cg0_done[state_idx].wait(state_index.phase) @@ -1761,16 +1776,8 @@ def tmaldg_warp( state_index = advance(state_index, cfg.smem_state_stages) if elect_one: bars.mb_state_ready[state_idx].arrive(n_bytes=cfg.tma_state_bytes) - if cutlass.const_expr(cfg.use_initial_state): - if chunk_idx == 0: - initial_state_slice = tma_slice_runtime_desc(desc_initial_state_slot, cutlass.Int32(0), cutlass.Int32(0), head_o, batch_idx) - tma_load_tile(sState_tma[state_idx], initial_state_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) - else: - state_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx - cutlass.Int32(1), head_o) - tma_load_tile(sState_tma[state_idx], state_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) - else: - state_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx - FIRST_STATE_CHUNK, head_o) - tma_load_tile(sState_tma[state_idx], state_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) + state_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx, head_o) + tma_load_tile(sState_tma[state_idx], state_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) # ---- dO load --------------------------------------------------------- bars.mb_do_done[raw_index.idx].wait(raw_index.phase) @@ -1862,7 +1869,6 @@ def compute0_warp_group( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) sk_nt = cend - wstart @@ -2183,8 +2189,17 @@ def compute0_warp_group( state_src = sState_raw.data_ptr() + state_index.idx * (cfg.d_k * cfg.d_v) for pl in cutlass.range_constexpr(2): for g8 in cutlass.range_constexpr(8): - state_frag = (state_src + pl * (cfg.d_k * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, g8 * 8, elem_bytes=2)).load( - count=8, alignment=16 + state_frag = cutlass.Vector.from_elements( + tuple( + ( + state_src + + (value_dim // 64) * (cfg.d_v * 64) + + (pl * 64 + g8 * 8 + e) * 64 + + swizzle_xor_128b(pl * 64 + g8 * 8 + e, value_dim % 64, elem_bytes=2) + ).load() + for e in range(8) + ), + cfg.io_dtype, ) nvvm.tcgen05_st( "32x32b", @@ -2253,7 +2268,6 @@ def compute1_warp_group( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) sk_nt = cend - wstart @@ -2266,12 +2280,14 @@ def compute1_warp_group( bars.mb_dstate_smem_cg2_done.wait(dstate_smem_done_index.phase) dstate_smem_done_index = advance(dstate_smem_done_index, 1) row_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 + dstate_src = (mDstate_in.iterator + mDstate_in.layout((batch_idx, head_idx, value_dim, 0))).raw_ptr() for sub in cutlass.range_constexpr(cfg.d_k // 16): seed_block = cutlass.Array(cutlass.Float32, 16, alignment=16) - for kk_i in cutlass.range_constexpr(16): - dval = mDstate_in[batch_idx, head_idx, sub * 16 + kk_i, value_dim].to(cutlass.Float32) - dval = dval if seed_true else cutlass.Float32(0.0) - seed_block[kk_i] = dval + for g in cutlass.range_constexpr(4): + seed_chunk = (dstate_src + sub * 16 + g * 4).load(count=4, alignment=16) + for t in cutlass.range_constexpr(4): + dval = seed_chunk[t].to(cutlass.Float32) + seed_block[g * 4 + t] = dval if seed_true else cutlass.Float32(0.0) nvvm.tcgen05_st( "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 16), cutlass.Float32), @@ -2551,20 +2567,24 @@ def compute1_warp_group( if sk_nt > 0: if wstart == 0: row_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 + dstate0_dst = (mDstate0.iterator + mDstate0.layout((batch_idx, head_idx, value_dim, 0))).raw_ptr() for sub in cutlass.range_constexpr(cfg.d_k // 32): dstate0_vec = nvvm.tcgen05_ld( "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 32), cutlass.Float32), num=32 ) - for kk_i in cutlass.range_constexpr(32): - mDstate0[batch_idx, head_idx, sub * 32 + kk_i, value_dim] = dstate0_vec[kk_i] + for g in cutlass.range_constexpr(8): + (dstate0_dst + sub * 32 + g * 4).store( + cutlass.Vector.from_elements(tuple(dstate0_vec[g * 4 + t] for t in range(4)), cutlass.Float32), + alignment=16, + ) else: for key_dim_base in cutlass.range_constexpr(0, cfg.d_k, 32): for kk_i in cutlass.range_constexpr(32): kd = key_dim_base + kk_i if cutlass.const_expr(cfg.use_dstate_in): - mDstate0[batch_idx, head_idx, kd, value_dim] = mDstate_in[batch_idx, head_idx, kd, value_dim] + mDstate0[batch_idx, head_idx, value_dim, kd] = mDstate_in[batch_idx, head_idx, value_dim, kd] else: - mDstate0[batch_idx, head_idx, kd, value_dim] = cutlass.Float32(0.0) + mDstate0[batch_idx, head_idx, value_dim, kd] = cutlass.Float32(0.0) bars.mb_dstate0_acc_stored.arrive() gbase += sk_nt tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @@ -2624,7 +2644,6 @@ def compute2_warp_group( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) sk_nt = cend - wstart @@ -2898,7 +2917,6 @@ def build_descs_body( base_dwo, base_dbo, base_checkpoint, - base_initial_state, desc_ws: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, @@ -2915,7 +2933,6 @@ def build_descs_body( dwo: cute.Tensor, dbo: cute.Tensor, state_checkpoints: cute.Tensor, - initial_state: cute.Tensor | None, n_batch: cutlass.Int32, q_rs: cutlass.Int32, k_rs: cutlass.Int32, @@ -2951,7 +2968,6 @@ def build_descs_body( sub11 = cute.make_tensor(desc_ws.iterator + 11 * arr_words, cute.make_layout((arr_words,), stride=(1,))) sub12 = cute.make_tensor(desc_ws.iterator + 12 * arr_words, cute.make_layout((arr_words,), stride=(1,))) sub13 = cute.make_tensor(desc_ws.iterator + 13 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub14 = cute.make_tensor(desc_ws.iterator + 14 * arr_words, cute.make_layout((cutlass.Int32(TENSOR_MAP_QWORDS),), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -3009,11 +3025,6 @@ def build_descs_body( if nvvm.elect_sync(): emit_checkpoint_seq_descs(base_checkpoint, sub13, cu_seqlens, state_checkpoints, n_batch, checkpoint_rs, checkpoint_every_n, 2) nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) - if cutlass.const_expr(initial_state is not None): - if widx == 14: - if nvvm.elect_sync(): - emit_copy_desc(base_initial_state, sub14) - nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) @cute.kernel @@ -3036,7 +3047,6 @@ def prologue_kernel( base_dwo: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_dbo: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], - base_initial_state: cutlass.GridConstant[cuda.tensor_map.TensorMap], desc_ws: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, @@ -3053,7 +3063,6 @@ def prologue_kernel( dwo: cute.Tensor, dbo: cute.Tensor, state_checkpoints: cute.Tensor, - initial_state: cute.Tensor | None, mStaging: cute.Tensor | None, mCount: cute.Tensor, mWorkItems: cute.Tensor, @@ -3122,7 +3131,6 @@ def prologue_kernel( base_dwo, base_dbo, base_checkpoint, - base_initial_state, desc_ws, cu_seqlens, q, @@ -3139,7 +3147,6 @@ def prologue_kernel( dwo, dbo, state_checkpoints, - initial_state, n_batch, q_rs, k_rs, @@ -3180,7 +3187,6 @@ def prologue( dwo: cute.Tensor, dbo: cute.Tensor, state_checkpoints: cute.Tensor, - initial_state: cute.Tensor | None, cu_seqlens: cute.Tensor, work_item_staging: cute.Tensor | None, work_count: cute.Tensor, @@ -3190,7 +3196,7 @@ def prologue( stream: cuda_driver.CUstream, ): """One-launch prologue: LPT-order the work items (when this kernel is - the table's first consumer) and build the 15 per-(batch, head) + the table's first consumer) and build the 14 per-(batch, head) TMA-descriptor arrays into ``tensormap_workspace``.""" h_q = q.shape[1] h_k = k.shape[1] @@ -3240,16 +3246,6 @@ def prologue( ), ) base_checkpoint = cuda.create_tensor_map_tiled_from_view(checkpoint_view, box_dims=(64, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) - base_initial_state = base_checkpoint - if cutlass.const_expr(initial_state is not None): - initial_state_view = cute.make_tensor( - initial_state.iterator, - cute.make_layout( - (d_v, d_k, ho, batch_size), - stride=(initial_state.stride[3], initial_state.stride[2], initial_state.stride[1], initial_state.stride[0]), - ), - ) - base_initial_state = cuda.create_tensor_map_tiled_from_view(initial_state_view, box_dims=(64, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) prologue_kernel( run_order, @@ -3270,7 +3266,6 @@ def prologue( base_dwo, base_dbo, base_checkpoint, - base_initial_state, tensormap_workspace, cu_seqlens, q, @@ -3287,7 +3282,6 @@ def prologue( dwo, dbo, state_checkpoints, - initial_state, work_item_staging, work_count, work_items, @@ -3315,7 +3309,6 @@ def prologue( def host( cfg: cutlass.Constexpr, state_checkpoints: cute.Tensor, - mState_init: cute.Tensor | None, a_log: cute.Tensor | None, dt_bias: cute.Tensor | None, dgate: cute.Tensor, @@ -3406,7 +3399,6 @@ def kernel( desc_dwo_base = desc_base_words + cutlass.Int32(11) * arr_words desc_db_base = desc_base_words + cutlass.Int32(12) * arr_words desc_checkpoint_base = desc_base_words + cutlass.Int32(13) * arr_words - desc_initial_state_base = desc_base_words + cutlass.Int32(14) * arr_words SMEM = cutlass.AddressSpace.smem bars = make_gdn2_bwd_bars(cfg) @@ -3455,6 +3447,14 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) + sState_direct = SmemTile( + base=sState_raw.data_ptr().toint(), + elems_per_stage=((cfg.state_cosize) // (cfg.smem_state_stages)) * bpe, + stages=cfg.smem_state_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) sK_decay_lead16 = SmemTile( base=sK_decay_raw.data_ptr().toint(), elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, @@ -3690,7 +3690,6 @@ def kernel( desc_beta_base, desc_w_base, desc_checkpoint_base, - desc_initial_state_base, bars, ) elif warp_idx == cfg.super_mma_warp_id: @@ -3721,6 +3720,7 @@ def kernel( sSched, tmem_hold, sState_alt, + sState_direct, sK_decay_lead16, sK_inv_amaj, sK_restore_lead16, @@ -4028,7 +4028,7 @@ def build_cfg( TENSORMAP_DESC_ARRAYS = 14 # per-batch runtime TMA descriptors: Q, K, V, Gate, dO, Beta, W, Checkpoint, dQ, dK, dV, dGate, dW_out, dBeta -TENSORMAP_STATIC_SLOTS = 1 # initial_state +TENSORMAP_STATIC_SLOTS = 0 # ---- Torch adapter / host-side compilation --------------------------------------- @@ -4074,7 +4074,7 @@ def chunk_gdn2_bwd_sm100( cu_seqlens, scale: float, *, - initial_state=None, + use_initial_state: bool = False, d_initial_state=None, d_final_state=None, use_qk_l2norm_in_kernel: bool = False, @@ -4107,10 +4107,10 @@ def chunk_gdn2_bwd_sm100( or logits when ``use_beta_sigmoid`` w: ``(total_tokens, HO, DV)`` io dtype post-sigmoid per-value write do: ``(total_tokens, HO, DV)`` io dtype - state_checkpoints: ``(total_checkpoints, HO, DK, DV)`` io dtype (KV, v contiguous - the GDN - checkpoint layout), the PLAIN per-chunk series with no initial-state - slot: sequence-local entry ``c - 1`` is the state ENTERING chunk c >= 1 - of sequence b; chunk 0 seeds from ``initial_state`` + state_checkpoints: ``(total_checkpoints, HO, DV, DK)`` io dtype (VK, k + contiguous - the GDN checkpoint layout), the PLAIN per-chunk series: + sequence-local entry ``c`` is the state ENTERING chunk c of sequence + b, so row 0 is the initial state (or zeros when the forward had none) dq/dk/dv: io dtype at ``HO = max(HQ, HV)`` heads, pre-allocated dgate: ``(total_tokens, HO, DK)`` fp32 (dL/d ln alpha; ``safe_gate`` leaves it in the transformed gate space), pre-allocated @@ -4119,10 +4119,10 @@ def chunk_gdn2_bwd_sm100( dw: ``(total_tokens, HO, DV)`` io dtype, pre-allocated cu_seqlens: ``(num_seqs + 1,)`` int32 scale: attention scale factor - initial_state: ``(num_seqs, HO, DK, DV)`` io dtype (KV) - the state - entering chunk 0 (engine-provided zeros when the graph has none) - d_initial_state: fp32 ``(num_seqs, HO, DK, DV)`` OUT (dL/dS0), or None - d_final_state: fp32 ``(num_seqs, HO, DK, DV)`` IN (dL/d final state) + use_initial_state: the forward ran with an initial state, so chunk 0 + has an entering state to load from ``state_checkpoints`` row 0 + d_initial_state: fp32 ``(num_seqs, HO, DV, DK)`` OUT (dL/dS0), or None + d_final_state: fp32 ``(num_seqs, HO, DV, DK)`` IN (dL/d final state) use_qk_l2norm_in_kernel: q/k arrive raw; the kernel normalizes for the recompute math and chains the L2-norm backward into dq/dk safe_gate: interpret ``gate`` through the safe-gate transform @@ -4145,16 +4145,13 @@ def chunk_gdn2_bwd_sm100( HO = max(HQ, HV) use_dstate_in = d_final_state is not None use_dstate0 = d_initial_state is not None - use_initial_state = initial_state is not None if work_items is None or work_count is None: raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") dyn_sched = sched_ctr is not None order_gen = work_item_scratch is None if order_in_prologue and sched_all is None: raise ValueError("order_in_prologue requires sched_all (the prologue zeroes both consumers' sched rings)") - for name, t in (("state_checkpoints", state_checkpoints), ("beta", beta), ("w", w), ("dbeta", dbeta), ("dw", dw)) + ( - (("initial_state", initial_state),) if use_initial_state else () - ): + for name, t in (("state_checkpoints", state_checkpoints), ("beta", beta), ("w", w), ("dbeta", dbeta), ("dw", dw)): if str(t.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: raise ValueError(f"{name} dtype must match the io dtype: got {t.dtype} with io {q.dtype}") for name, hh in (("HQ", HQ), ("HK", HK), ("HV", HV)): @@ -4224,9 +4221,6 @@ def chunk_gdn2_bwd_sm100( tensormap_ws_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() state_checkpoints_cute = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=len(state_checkpoints.shape) - 1) - initial_state_cute = ( - from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=len(initial_state.shape) - 1) if use_initial_state else None - ) a_log_cute = from_dlpack(a_log, assumed_align=4) if a_log is not None else None dt_bias_cute = from_dlpack(dt_bias, assumed_align=16) if dt_bias is not None else None dgate_cute = from_dlpack(dgate, assumed_align=16).mark_layout_dynamic(leading_dim=len(dgate.shape) - 1) @@ -4236,7 +4230,6 @@ def chunk_gdn2_bwd_sm100( host, cfg, state_checkpoints_cute, - initial_state_cute, a_log_cute, dt_bias_cute, dgate_cute, @@ -4270,7 +4263,6 @@ def chunk_gdn2_bwd_sm100( dwo_pl = from_dlpack(dw, assumed_align=16).mark_layout_dynamic(leading_dim=2) dbo_pl = from_dlpack(dbeta, assumed_align=16).mark_layout_dynamic(leading_dim=2) state_checkpoints_pl = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=3) - initial_state_pl = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) if use_initial_state else None cu_pl = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() ws_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() staging_pl = None @@ -4304,7 +4296,6 @@ def chunk_gdn2_bwd_sm100( dwo_pl, dbo_pl, state_checkpoints_pl, - initial_state_pl, cu_pl, staging_pl, work_count_pl, @@ -4329,7 +4320,6 @@ def chunk_gdn2_bwd_sm100( dw, dbeta, state_checkpoints, - initial_state, cu_seqlens, work_item_scratch if not order_gen else None, work_count, @@ -4340,7 +4330,6 @@ def chunk_gdn2_bwd_sm100( ) cache["compiled"]( state_checkpoints, - initial_state, a_log, dt_bias, dgate, @@ -4376,7 +4365,6 @@ def run_bwd( dbeta, dw, cu_seqlens, - initial_state, d_initial_state, d_final_state, work_items, @@ -4409,7 +4397,6 @@ def run_bwd( dw, dbeta, state_checkpoints, - initial_state, cu_seqlens, work_item_scratch, work_count, @@ -4420,7 +4407,6 @@ def run_bwd( ) cache["compiled"]( state_checkpoints, - initial_state, a_log, dt_bias, dgate, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py index 9b03aa0ef..98f9ff763 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py @@ -35,7 +35,7 @@ ABI: q `[T, HQ, DK]`, k `[T, HK, DK]`, v `[T, HV, DV]`, gate `[T, HO, DK]` fp32 (natural-log decay unless SAFE_GATE), beta `[T, HO, DK]` and w `[T, HO, DV]` in the io dtype, cu_seqlens int32, states/checkpoints -`[N, HO, DK, DV]` (KV, v contiguous). GQA/GVA head broadcast follows repeat_interleave: +`[N, HO, DV, DK]` (VK, k contiguous). GQA/GVA head broadcast follows repeat_interleave: source head = head_idx // (HO // H_x). State presence, L2NORM, SAFE_GATE, checkpoints, and the head ratios are compile-time specializations. @@ -948,6 +948,16 @@ def epilogue_warp( if elect_one: tma_tensormap_acquire(desc_o_slot) num_tile_chunks = wend - cstart + if cutlass.const_expr(cfg.enable_checkpoints): + if num_tile_chunks > 0 and wstart == 0: + checkpoint_stage = checkpoint_ready_index.idx + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) + checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) + checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0), head_o) + tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) + tma_store_commit() + tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() for local_chunk in cutlass.range(num_tile_chunks, unroll=1): chunk_idx = cstart + local_chunk global_chunk = global_chunk_base + local_chunk @@ -1038,7 +1048,7 @@ def epilogue_warp( if do_checkpoint: bars.mb_checkpoint_tmastg_ready[checkpoint_ready_index.idx].wait(checkpoint_ready_index.phase) checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) - checkpoint_entry = checkpoint_quot - cutlass.Int32(1) + checkpoint_entry = checkpoint_quot checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), checkpoint_entry, head_o) tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) tma_store_commit() @@ -1535,11 +1545,14 @@ def compute1_warp_group( seed_from_initial_state = cstart == 0 if cutlass.const_expr(mState_init is not None): if seed_from_initial_state: + seed_vw = 16 // (mState_init.element_type.width // 8) + seed_src = (mState_init.iterator + mState_init.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for col in cutlass.range_constexpr(32): - key_dim = key_block_start + col - state_block[col] = mState_init[batch_idx, head_o, key_dim, value_dim].to(cutlass.Float32) + for g in cutlass.range_constexpr(32 // seed_vw): + seed_chunk = (seed_src + key_block_start + g * seed_vw).load(count=seed_vw, alignment=16) + for t in cutlass.range_constexpr(seed_vw): + state_block[g * seed_vw + t] = seed_chunk[t].to(cutlass.Float32) nvvm.tcgen05_st( "32x32b", @@ -1588,8 +1601,48 @@ def compute1_warp_group( nvvm.tcgen05_wait("store") bars.mb_state_inp_ready.arrive() if cutlass.const_expr(cfg.enable_checkpoints): + if wstart == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + packs = tuple( + fp32_to_fp16(state_vecs[k_block][g * 8 + 2 * t], state_vecs[k_block][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) + for t in range(4) + ) + dk = k_block * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.tcgen05_wait("load") + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() bars.mb_state_acc_read_done.arrive() + if cutlass.const_expr(cfg.enable_checkpoints and mState_init is None): + if wstart == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + zero_packs = tuple(cutlass.Int32(0) for _ in range(4)) + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + dk = k_block * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(zero_packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + # ---- Y staging: Y = W*V - state*(Beta*K) ----------------------------- bars.mb_v_ready[raw_bar_index.idx].wait(raw_bar_index.phase) projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset @@ -1774,35 +1827,15 @@ def compute1_warp_group( checkpoint_swz_col0 = (checkpoint_vbase + frag_col_offset) % 64 checkpoint_swz_off = (checkpoint_vbase + 16 + frag_col_offset) // 64 * (cfg.d_k * 64) checkpoint_swz_col = (checkpoint_vbase + 16 + frag_col_offset) % 64 - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - checkpoint_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_block * 16, cutlass.Float32), num=2) - checkpoint_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + state_col_id + k_block * 16, cutlass.Float32), num=2) - checkpoint_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - checkpoint_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - for reg_idx in cutlass.range_constexpr(4): - checkpoint_pack0[reg_idx] = fp32_to_fp16(checkpoint_vec0[2 * reg_idx], checkpoint_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) - checkpoint_pack1[reg_idx] = fp32_to_fp16(checkpoint_vec1[2 * reg_idx], checkpoint_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) - checkpoint_row = k_block * 16 + frag_row_coord - nvvm.stmatrix( - sCheckpoint_ptr - + checkpoint_stage_base - + checkpoint_swz_off0 - + checkpoint_row * 64 - + swizzle_xor_128b(checkpoint_row, checkpoint_swz_col0, elem_bytes=2), - checkpoint_pack0.data_ptr().load(count=4, alignment=4), - nvvm.MMALayout.COL, - shape=nvvm.StoreShape.M8N8, - ) - nvvm.stmatrix( - sCheckpoint_ptr - + checkpoint_stage_base - + checkpoint_swz_off - + checkpoint_row * 64 - + swizzle_xor_128b(checkpoint_row, checkpoint_swz_col, elem_bytes=2), - checkpoint_pack1.data_ptr().load(count=4, alignment=4), - nvvm.MMALayout.COL, - shape=nvvm.StoreShape.M8N8, - ) + for k_base in cutlass.range_constexpr(0, cfg.d_k, 32): + checkpoint_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_base, cutlass.Float32), num=32) + for g in cutlass.range_constexpr(4): + packs = tuple(fp32_to_fp16(checkpoint_vec[g * 8 + 2 * t], checkpoint_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4)) + dk = k_base + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store(cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16) nvvm.tcgen05_wait("load") bars.mb_state_acc_read_done.arrive() nvvm.fence_proxy("async.shared", space="cta") @@ -2030,6 +2063,8 @@ def compute1_warp_group( if cutlass.const_expr(mState_out is not None): if seqlen_b > 0: if owns_final: + state_vw = 16 // (mState_out.element_type.width // 8) + state_dst = (mState_out.iterator + mState_out.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): loaded = nvvm.tcgen05_ld( "32x32b", @@ -2037,17 +2072,22 @@ def compute1_warp_group( num=32, ) - for col in cutlass.range_constexpr(32): - key_dim = key_block_start + col - mState_out[batch_idx, head_o, key_dim, value_dim] = loaded[col].to(mState_out.element_type) + for g in cutlass.range_constexpr(32 // state_vw): + (state_dst + key_block_start + g * state_vw).store( + cutlass.Vector.from_elements( + tuple(loaded[g * state_vw + t].to(mState_out.element_type) for t in range(state_vw)), + mState_out.element_type, + ), + alignment=16, + ) else: for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): for col in cutlass.range_constexpr(32): key_dim = key_block_start + col if cutlass.const_expr(mState_init is not None): - mState_out[batch_idx, head_o, key_dim, value_dim] = mState_init[batch_idx, head_o, key_dim, value_dim] + mState_out[batch_idx, head_o, value_dim, key_dim] = mState_init[batch_idx, head_o, value_dim, key_dim] else: - mState_out[batch_idx, head_o, key_dim, value_dim] = cutlass.Float32(0.0).to(mState_out.element_type) + mState_out[batch_idx, head_o, value_dim, key_dim] = cutlass.Float32(0.0).to(mState_out.element_type) global_chunk_base += num_tile_chunks tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @@ -2826,7 +2866,7 @@ def prologue( GLOBAL_ADDRESS (Int64) and caps the token GLOBAL_DIM to the sequence length, so the main kernel's coordinates are sequence-relative and tail chunks clip in hardware. The checkpoint descriptor is 3-D ``(dv, dk, entry)`` - over the packed ``[total_checkpoints, HO, DK, DV]`` series; its per-sequence entry + over the packed ``[total_checkpoints, HO, DV, DK]`` series; its per-sequence entry offsets ((seqlen-1)//N, prefix-summed) are derived on device and its entry extent is capped per sequence, so checkpoint store coordinates are sequence-local.""" @@ -3060,14 +3100,14 @@ def chunk_gdn2_sm100( w: ``(total_tokens, HO, DV)`` io dtype, channel-wise write gate output: ``(total_tokens, HO, DV)`` float16/bfloat16, pre-allocated cu_seqlens: ``(num_seqs + 1,)`` int32 - initial_state: ``(num_seqs, HO, DK, DV)`` float32/bfloat16, or None - output_state: ``(num_seqs, HO, DK, DV)`` float32/bfloat16, or None + initial_state: ``(num_seqs, HO, DV, DK)`` float32/bfloat16, or None + output_state: ``(num_seqs, HO, DV, DK)`` float32/bfloat16, or None scale: attention scale factor (must not be 0) checkpoint_every_n_tokens: emit a checkpoint entry every N tokens (0 = off). checkpoint[j] is the state after ``(j + 1) * N`` tokens, STRICTLY BEFORE the sequence end — the end-of-sequence state is only ``output_state``. - output_state_checkpoints: ``(total_checkpoints, HO, DK, DV)`` io-dtype (KV, v + output_state_checkpoints: ``(total_checkpoints, HO, DV, DK)`` io-dtype (VK, k contiguous — the GDN checkpoint layout); the per-sequence entry offsets are derived on device from ``cu_seqlens`` ((seqlen-1)//N, prefix-summed), so there is no cu_checkpoints array diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py index 66676ead1..8e5c5dd23 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py @@ -36,7 +36,7 @@ ABI: k `[T, HK, DK]`, v `[T, HV, DV]`, gate `[T, HO, DK]` fp32 (natural-log decay unless SAFE_GATE), beta `[T, HO, DK]` and w `[T, HO, DV]` in the io dtype, cu_seqlens int32, states/checkpoints -`[N, HO, DK, DV]` (KV, v contiguous). GQA/GVA head broadcast follows repeat_interleave: +`[N, HO, DV, DK]` (VK, k contiguous). GQA/GVA head broadcast follows repeat_interleave: source head = head_idx // (HO // H_x). State presence, L2NORM, SAFE_GATE, checkpoints, and the head ratios are compile-time specializations. @@ -848,6 +848,15 @@ def epilogue_warp( checkpoint_mod = (cstart + cutlass.Int32(1)) % checkpoint_chunks if elect_one: tma_tensormap_acquire(desc_checkpoint_slot) + if num_chunks_tile > 0 and wstart == 0: + checkpoint_stage = checkpoint_ready_index.idx + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) + checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) + checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0), head_o) + tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) + tma_store_commit() + tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() for local_chunk in cutlass.range(num_chunks_tile, unroll=1): chunk_idx = cstart + local_chunk if local_chunk > 0: @@ -858,7 +867,7 @@ def epilogue_warp( checkpoint_stage = checkpoint_ready_index.idx bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) - checkpoint_entry = checkpoint_quot - cutlass.Int32(1) + checkpoint_entry = checkpoint_quot checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), checkpoint_entry, head_o) tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) tma_store_commit() @@ -1259,11 +1268,14 @@ def compute1_warp_group( seed_from_initial_state = cstart == 0 if cutlass.const_expr(mState_init is not None): if seed_from_initial_state: + seed_vw = 16 // (mState_init.element_type.width // 8) + seed_src = (mState_init.iterator + mState_init.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for col in cutlass.range_constexpr(32): - key_dim = key_block_start + col - state_block[col] = mState_init[batch_idx, head_o, key_dim, value_dim].to(cutlass.Float32) + for g in cutlass.range_constexpr(32 // seed_vw): + seed_chunk = (seed_src + key_block_start + g * seed_vw).load(count=seed_vw, alignment=16) + for t in cutlass.range_constexpr(seed_vw): + state_block[g * seed_vw + t] = seed_chunk[t].to(cutlass.Float32) nvvm.tcgen05_st( "32x32b", @@ -1312,8 +1324,48 @@ def compute1_warp_group( nvvm.tcgen05_wait("store") bars.mb_state_inp_ready.arrive() if cutlass.const_expr(cfg.enable_checkpoints): + if wstart == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + packs = tuple( + fp32_to_fp16(state_vecs[k_block][g * 8 + 2 * t], state_vecs[k_block][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) + for t in range(4) + ) + dk = k_block * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.tcgen05_wait("load") + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() bars.mb_state_acc_read_done.arrive() + if cutlass.const_expr(cfg.enable_checkpoints and mState_init is None): + if wstart == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + zero_packs = tuple(cutlass.Int32(0) for _ in range(4)) + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + dk = k_block * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(zero_packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + # ---- Y staging: Y = W*V - state*(Beta*K) ----------------------------- bars.mb_v_ready[raw_bar_index.idx].wait(raw_bar_index.phase) projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset @@ -1493,35 +1545,15 @@ def compute1_warp_group( checkpoint_swz_col0 = (checkpoint_vbase + ldsm_col_offset) % 64 checkpoint_swz_off = (checkpoint_vbase + 16 + ldsm_col_offset) // 64 * (cfg.d_k * 64) checkpoint_swz_col = (checkpoint_vbase + 16 + ldsm_col_offset) % 64 - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - checkpoint_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_block * 16, cutlass.Float32), num=2) - checkpoint_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + state_col_id + k_block * 16, cutlass.Float32), num=2) - checkpoint_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - checkpoint_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - for reg_idx in cutlass.range_constexpr(4): - checkpoint_pack0[reg_idx] = fp32_to_fp16(checkpoint_vec0[2 * reg_idx], checkpoint_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) - checkpoint_pack1[reg_idx] = fp32_to_fp16(checkpoint_vec1[2 * reg_idx], checkpoint_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) - checkpoint_row = k_block * 16 + ldsm_row_coord - nvvm.stmatrix( - sCheckpoint_ptr - + checkpoint_stage_base - + checkpoint_swz_off0 - + checkpoint_row * 64 - + swizzle_xor_128b(checkpoint_row, checkpoint_swz_col0, elem_bytes=2), - checkpoint_pack0.data_ptr().load(count=4, alignment=4), - nvvm.MMALayout.COL, - shape=nvvm.StoreShape.M8N8, - ) - nvvm.stmatrix( - sCheckpoint_ptr - + checkpoint_stage_base - + checkpoint_swz_off - + checkpoint_row * 64 - + swizzle_xor_128b(checkpoint_row, checkpoint_swz_col, elem_bytes=2), - checkpoint_pack1.data_ptr().load(count=4, alignment=4), - nvvm.MMALayout.COL, - shape=nvvm.StoreShape.M8N8, - ) + for k_base in cutlass.range_constexpr(0, cfg.d_k, 32): + checkpoint_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_base, cutlass.Float32), num=32) + for g in cutlass.range_constexpr(4): + packs = tuple(fp32_to_fp16(checkpoint_vec[g * 8 + 2 * t], checkpoint_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4)) + dk = k_base + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store(cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16) nvvm.tcgen05_wait("load") bars.mb_state_acc_read_done.arrive() nvvm.fence_proxy("async.shared", space="cta") @@ -1656,6 +1688,8 @@ def compute1_warp_group( if cutlass.const_expr(mState_out is not None): if seqlen_b > 0: if owns_final: + state_vw = 16 // (mState_out.element_type.width // 8) + state_dst = (mState_out.iterator + mState_out.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): state_vec = nvvm.tcgen05_ld( "32x32b", @@ -1663,17 +1697,22 @@ def compute1_warp_group( num=32, ) - for col in cutlass.range_constexpr(32): - key_dim = key_block_start + col - mState_out[batch_idx, head_o, key_dim, value_dim] = state_vec[col].to(mState_out.element_type) + for g in cutlass.range_constexpr(32 // state_vw): + (state_dst + key_block_start + g * state_vw).store( + cutlass.Vector.from_elements( + tuple(state_vec[g * state_vw + t].to(mState_out.element_type) for t in range(state_vw)), + mState_out.element_type, + ), + alignment=16, + ) else: for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): for col in cutlass.range_constexpr(32): key_dim = key_block_start + col if cutlass.const_expr(mState_init is not None): - mState_out[batch_idx, head_o, key_dim, value_dim] = mState_init[batch_idx, head_o, key_dim, value_dim] + mState_out[batch_idx, head_o, value_dim, key_dim] = mState_init[batch_idx, head_o, value_dim, key_dim] else: - mState_out[batch_idx, head_o, key_dim, value_dim] = cutlass.Float32(0.0).to(mState_out.element_type) + mState_out[batch_idx, head_o, value_dim, key_dim] = cutlass.Float32(0.0).to(mState_out.element_type) tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) bars.mb_tmem_done[0].arrive() @@ -2555,13 +2594,13 @@ def chunk_gdn2_recompute_sm100( Post-sigmoid, or logits when ``use_beta_sigmoid`` w: ``(total_tokens, HO, DV)`` io dtype, channel-wise write gate cu_seqlens: ``(num_seqs + 1,)`` int32 - initial_state: ``(num_seqs, HO, DK, DV)`` float32/bfloat16, or None - output_state: ``(num_seqs, HO, DK, DV)`` float32/bfloat16, or None + initial_state: ``(num_seqs, HO, DV, DK)`` float32/bfloat16, or None + output_state: ``(num_seqs, HO, DV, DK)`` float32/bfloat16, or None checkpoint_every_n_tokens: emit a checkpoint entry every N tokens (0 = off). checkpoint[j] is the state after ``(j + 1) * N`` tokens, STRICTLY BEFORE the sequence end - the end-of-sequence state is only ``output_state``. - output_state_checkpoints: ``(total_checkpoints, HO, DK, DV)`` io-dtype (KV, V + output_state_checkpoints: ``(total_checkpoints, HO, DV, DK)`` io-dtype (VK, K contiguous - the GDN checkpoint layout); the per-sequence entry offsets are derived on device from ``cu_seqlens`` ((seqlen-1)//N, prefix-summed), so there is no cu_checkpoints array diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py index e620609fe..8f34c5461 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py @@ -116,7 +116,7 @@ import cutlass.experimental.cuda.tensor_map as tma from cutlass.cute.runtime import from_dlpack -from ..common.thd import emit_copy_desc, emit_checkpoint_seq_descs, emit_seq_descs, TENSOR_MAP_QWORDS +from ..common.thd import emit_checkpoint_seq_descs, emit_seq_descs, TENSOR_MAP_QWORDS from ..common.split_k import ORDER_CAPACITY, ORDER_ELEMS, ORDER_THREADS, decode_work_item, order_body from ..common.elementwise import softplus from ..common.host import get_dtype @@ -759,7 +759,6 @@ def gate_beta_warp( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) num_item_chunks = cend - wstart @@ -1110,6 +1109,26 @@ def mma_warp( idesc=idesc_k_state, kind=nvvm.Tcgen05MMAKind.F16, ) + idesc_k_state_kmaj = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.b_t, + m_dim=cfg.d_v, + ) + bmm_k_state_kmaj_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.d_k, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + atranspose=False, + cta_group=1, + idesc=idesc_k_state_kmaj, + kind=nvvm.Tcgen05MMAKind.F16, + ) idesc_du = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -1223,14 +1242,15 @@ def mma_warp( idesc=idesc_u, kind=nvvm.Tcgen05MMAKind.F16, ) - idesc_dqdk_inter = nvvm.Tcgen05InstrDesc.build( + idesc_dqdk_inter_at = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, b_dtype=cfg.io_dtype, n_dim=cfg.b_t, m_dim=cfg.d_k, + a_major=1, ) - bmm_dqdk_inter_desc = MmaDesc( + bmm_dqdk_inter_at_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.d_v, @@ -1238,9 +1258,9 @@ def mma_warp( bpe_b=bpe, tile_k_hw=16, btranspose=False, - atranspose=False, + atranspose=True, cta_group=1, - idesc=idesc_dqdk_inter, + idesc=idesc_dqdk_inter_at, kind=nvvm.Tcgen05MMAKind.F16, ) idesc_dka = nvvm.Tcgen05InstrDesc.build( @@ -1377,9 +1397,9 @@ def mma_warp( state_index = advance(state_index, cfg.smem_state_stages) bars.mb_kk_acc_done[0].wait(kk_acc_index.phase) kk_acc_index = advance(kk_acc_index, 1) - desc_state = d_state0 + desc_state = d_state_kmaj0 mma_ss( - bmm_k_state_desc, + bmm_k_state_kmaj_desc, desc_state, desc_k, k_state_acc_ptr, @@ -1419,10 +1439,10 @@ def mma_warp( if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_dq_acc_total_done[0].wait(dq_total_index.phase) dq_total_index = advance(dq_total_index, 1) - desc_state_kmaj_dq_inter = d_state_kmaj0 + desc_state_kmaj_dq_inter = d_state0 desc_do_kmaj_dq_inter = d_do_kmaj0 mma_ss( - bmm_dqdk_inter_desc, + bmm_dqdk_inter_at_desc, desc_state_kmaj_dq_inter, desc_do_kmaj_dq_inter, dq_acc_ptr, @@ -1518,11 +1538,11 @@ def mma_warp( bars.mb_dstate_smem_ready[0].wait(dstate_smem_index.phase) dstate_smem_index = advance(dstate_smem_index, 1) - desc_dstate_kmaj = d_dstate0 + desc_dstate = d_dstate0 desc_u_kmaj_dk_inter = d_v_kmaj0 mma_ss( - bmm_dqdk_inter_desc, - desc_dstate_kmaj, + bmm_dqdk_inter_at_desc, + desc_dstate, desc_u_kmaj_dk_inter, dvdk_acc_ptr, accumulate=False, @@ -1631,10 +1651,10 @@ def mma_warp( # ---- dK state-path = state(S) @ dY^T ----------------------------------------- if chunk_idx >= FIRST_STATE_CHUNK: - desc_state_kmaj_spath = d_state_kmaj0 + desc_state_kmaj_spath = d_state0 desc_dy_kmaj_spath = d_dv_kmaj0 mma_ss( - bmm_dqdk_inter_desc, + bmm_dqdk_inter_at_desc, desc_state_kmaj_spath, desc_dy_kmaj_spath, dk_state_path_acc_ptr, @@ -1703,7 +1723,6 @@ def tmaldg_warp( desc_v_base, desc_do_base, desc_checkpoint_base, - desc_initial_state_base, sSched, mSched, bars, @@ -1720,7 +1739,6 @@ def tmaldg_warp( sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 bpe = cfg.io_dtype.width // 8 granule_elems = 128 // bpe @@ -1798,16 +1816,12 @@ def tmaldg_warp( desc_v_slot = (desc_v_base + slot).tospace(cutlass.AddressSpace.generic) desc_do_slot = (desc_do_base + slot).tospace(cutlass.AddressSpace.generic) desc_checkpoint_slot = (desc_checkpoint_base + slot).tospace(cutlass.AddressSpace.generic) - if cutlass.const_expr(cfg.use_initial_state): - desc_initial_state_slot = (desc_initial_state_base + cutlass.Int32(0)).tospace(cutlass.AddressSpace.generic) if elect_one: tma_tensormap_acquire(desc_q_slot) tma_tensormap_acquire(desc_k_slot) tma_tensormap_acquire(desc_v_slot) tma_tensormap_acquire(desc_do_slot) tma_tensormap_acquire(desc_checkpoint_slot) - if cutlass.const_expr(cfg.use_initial_state): - tma_tensormap_acquire(desc_initial_state_slot) for rev_idx in cutlass.range(cend - wstart): chunk_idx = cend - 1 - rev_idx @@ -1851,23 +1865,15 @@ def tmaldg_warp( do_slice = tma_slice_runtime_desc(desc_do_slot, cutlass.Int32(0), head_o, tok_coord) tma_load_tile(sdO_tma[do_idx], do_slice, bars.mb_do_ready[do_idx].smem_ptr, acquire=False) - # ---- entering state: checkpoint c - 1, or initial_state for chunk 0 when given ---------- + # ---- entering state ---------- if chunk_idx >= FIRST_STATE_CHUNK: state_idx = state_index.idx bars.mb_state_mma_done[state_idx].wait(state_index.phase) state_index = advance(state_index, cfg.smem_state_stages) if elect_one: bars.mb_state_ready[state_idx].arrive(n_bytes=cfg.tma_state_bytes) - if cutlass.const_expr(cfg.use_initial_state): - if chunk_idx == 0: - initial_state_slice = tma_slice_runtime_desc(desc_initial_state_slot, cutlass.Int32(0), cutlass.Int32(0), head_o, batch_idx) - tma_load_tile(sCheckpoint_tma[state_idx], initial_state_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) - else: - checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx - 1, head_o) - tma_load_tile(sCheckpoint_tma[state_idx], checkpoint_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) - else: - checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx - FIRST_STATE_CHUNK, head_o) - tma_load_tile(sCheckpoint_tma[state_idx], checkpoint_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) + checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx, head_o) + tma_load_tile(sCheckpoint_tma[state_idx], checkpoint_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) @@ -1990,7 +1996,6 @@ def compute0_warp_group( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) num_item_chunks = cend - wstart @@ -2556,7 +2561,6 @@ def compute1_warp_group( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) num_item_chunks = cend - wstart @@ -2572,7 +2576,7 @@ def compute1_warp_group( for sub in cutlass.range_constexpr(num_state_subs): dstate_in_vals = [] for kk in cutlass.range_constexpr(ldtm_width): - v = gDstate_in[sub * ldtm_width + kk, cg1_tidx] + v = gDstate_in[cg1_tidx, sub * ldtm_width + kk] v = v if seed_from_dstate_in else cutlass.Float32(0.0) dstate_in_vals.append(v) nvvm.tcgen05_st( @@ -2589,26 +2593,21 @@ def compute1_warp_group( nvvm.tcgen05_wait("store") bars.mb_dstate_inp_ready[dstate_inp_idx].arrive() - for b in cutlass.range_constexpr(2): - for col_half in cutlass.range_constexpr(2): - dstate_smem_vec = nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + b * 16) << 16) + tmem_dstate_acc_col + col_half * 64, cutlass.Float32), - num=8, + for sub in cutlass.range_constexpr(num_state_subs): + dstate_smem_vec = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), + num=32, + ) + for g in cutlass.range_constexpr(ldtm_width // 8): + dstate_smem_pack = tuple( + fp32_to_fp16(dstate_smem_vec[g * 8 + 2 * t], dstate_smem_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dstate_dk = sub * ldtm_width + g * 8 + dstate_addr = (dstate_dk // 64) * (cfg.d_v * 64) + cg1_tidx * 64 + swizzle_xor_128b(cg1_tidx, dstate_dk % 64) + cutlass.inttoptr(sDstate_base_int + dstate_addr * 2, cutlass.AddressSpace.smem, cfg.io_dtype).store( + cutlass.Vector.from_elements(dstate_smem_pack, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 ) - dstate_smem_pack = [fp32_to_fp16(dstate_smem_vec[2 * j], dstate_smem_vec[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] - for c in cutlass.range_constexpr(4): - dstate_smem_row = col_half * 64 + frag_row + c * 16 - nvvm.stmatrix( - cutlass.inttoptr( - sDstate_base_int - + ((cg1_tidx // 64) * cfg.d_k * 64 + dstate_smem_row * 64 + swizzle_xor_128b(dstate_smem_row, frag_col + b * 16)) * 2, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [dstate_smem_pack[c * 4 + 0], dstate_smem_pack[c * 4 + 1], dstate_smem_pack[c * 4 + 2], dstate_smem_pack[c * 4 + 3]], - nvvm.MMALayout.COL, - ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dstate_smem_ready[0].arrive() @@ -3107,26 +3106,21 @@ def compute1_warp_group( if chunk_idx >= wstart + 1: bars.mb_state_dot_dstate_done[0].wait(cg1_state_dot_dstate_index.phase) cg1_state_dot_dstate_index = advance(cg1_state_dot_dstate_index, 1) - for b in cutlass.range_constexpr(2): - for col_half in cutlass.range_constexpr(2): - dstate_smem_vec = nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + b * 16) << 16) + tmem_dstate_acc_col + col_half * 64, cutlass.Float32), - num=8, + for sub in cutlass.range_constexpr(num_state_subs): + dstate_smem_vec = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), + num=32, + ) + for g in cutlass.range_constexpr(ldtm_width // 8): + dstate_smem_pack = tuple( + fp32_to_fp16(dstate_smem_vec[g * 8 + 2 * t], dstate_smem_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dstate_dk = sub * ldtm_width + g * 8 + dstate_addr = (dstate_dk // 64) * (cfg.d_v * 64) + cg1_tidx * 64 + swizzle_xor_128b(cg1_tidx, dstate_dk % 64) + cutlass.inttoptr(sDstate_base_int + dstate_addr * 2, cutlass.AddressSpace.smem, cfg.io_dtype).store( + cutlass.Vector.from_elements(dstate_smem_pack, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 ) - dstate_smem_pack = [fp32_to_fp16(dstate_smem_vec[2 * j], dstate_smem_vec[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] - for c in cutlass.range_constexpr(4): - dstate_smem_row = col_half * 64 + frag_row + c * 16 - nvvm.stmatrix( - cutlass.inttoptr( - sDstate_base_int - + ((cg1_tidx // 64) * cfg.d_k * 64 + dstate_smem_row * 64 + swizzle_xor_128b(dstate_smem_row, frag_col + b * 16)) * 2, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [dstate_smem_pack[c * 4 + 0], dstate_smem_pack[c * 4 + 1], dstate_smem_pack[c * 4 + 2], dstate_smem_pack[c * 4 + 3]], - nvvm.MMALayout.COL, - ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dstate_smem_ready[0].arrive() @@ -3146,7 +3140,7 @@ def compute1_warp_group( "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), num=32 ) for kk in cutlass.range_constexpr(32): - gDstate0[sub * ldtm_width + kk, cg1_tidx] = dstate0_vec[kk] + gDstate0[cg1_tidx, sub * ldtm_width + kk] = dstate0_vec[kk] if cutlass.const_expr(not cfg.use_dstate_in): bars.mb_dstate_scale_acc_done[dstate_idx].arrive() else: @@ -3158,11 +3152,11 @@ def compute1_warp_group( gDstate_in = mDstate_in[None, None, head_idx, batch_idx] for sub in cutlass.range_constexpr(num_state_subs): for kk in cutlass.range_constexpr(32): - gDstate0[sub * ldtm_width + kk, cg1_tidx] = gDstate_in[sub * ldtm_width + kk, cg1_tidx] + gDstate0[cg1_tidx, sub * ldtm_width + kk] = gDstate_in[cg1_tidx, sub * ldtm_width + kk] else: for sub in cutlass.range_constexpr(num_state_subs): for kk in cutlass.range_constexpr(32): - gDstate0[sub * ldtm_width + kk, cg1_tidx] = cutlass.Float32(0.0) + gDstate0[cg1_tidx, sub * ldtm_width + kk] = cutlass.Float32(0.0) tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @@ -3190,7 +3184,6 @@ def build_descs_body( base_dq, base_dk, base_dv, - base_initial_state, desc_ws: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, @@ -3201,7 +3194,6 @@ def build_descs_body( dq: cute.Tensor, dk: cute.Tensor, dv: cute.Tensor, - state0: Optional[cute.Tensor], n_batch: cutlass.Int32, q_rs: cutlass.Int32, k_rs: cutlass.Int32, @@ -3225,7 +3217,6 @@ def build_descs_body( sub5 = cute.make_tensor(desc_ws.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) sub6 = cute.make_tensor(desc_ws.iterator + 6 * arr_words, cute.make_layout((arr_words,), stride=(1,))) sub7 = cute.make_tensor(desc_ws.iterator + 7 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub8 = cute.make_tensor(desc_ws.iterator + 8 * arr_words, cute.make_layout((cutlass.Int32(TENSOR_MAP_QWORDS),), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -3259,11 +3250,6 @@ def build_descs_body( if nvvm.elect_sync(): emit_seq_descs(base_dv, sub7, cu_seqlens, dv, n_batch, dv_rs, 2) nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) - if cutlass.const_expr(state0 is not None): - if widx == 8: - if nvvm.elect_sync(): - emit_copy_desc(base_initial_state, sub8) - nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) @cute.kernel @@ -3279,7 +3265,6 @@ def prologue_kernel( base_dq: cutlass.GridConstant[tma.TensorMap], base_dk: cutlass.GridConstant[tma.TensorMap], base_dv: cutlass.GridConstant[tma.TensorMap], - base_initial_state: cutlass.GridConstant[tma.TensorMap], desc_ws: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, @@ -3290,7 +3275,6 @@ def prologue_kernel( dq: cute.Tensor, dk: cute.Tensor, dv: cute.Tensor, - state0: Optional[cute.Tensor], mStaging: Optional[cute.Tensor], mCount: cute.Tensor, mWorkItems: cute.Tensor, @@ -3347,7 +3331,6 @@ def prologue_kernel( base_dq, base_dk, base_dv, - base_initial_state, desc_ws, cu_seqlens, q, @@ -3358,7 +3341,6 @@ def prologue_kernel( dq, dk, dv, - state0, n_batch, q_rs, k_rs, @@ -3387,7 +3369,6 @@ def prologue( dv: cute.Tensor, state_checkpoints: cute.Tensor, cu_seqlens: cute.Tensor, - state0: Optional[cute.Tensor], work_item_staging: Optional[cute.Tensor], work_count: cute.Tensor, work_items: cute.Tensor, @@ -3398,8 +3379,7 @@ def prologue( """One-launch prologue: LPT-order the work items (with ``run_order``, when this kernel is the backward pair's first table consumer) and build the per-(b,h) TMA-descriptor arrays (Q, K, V, dO, checkpoint loads; dQ, dK, dV - stores; the io-dtype initial-state loads when ``state0`` is given) into - ``tensormap_workspace``.""" + stores) into ``tensormap_workspace``.""" h_q = q.shape[1] h_k = k.shape[1] h_v = v.shape[1] @@ -3447,17 +3427,6 @@ def prologue( ) base_desc_checkpoint = tma.create_tensor_map_tiled_from_view(checkpoint_view, box_dims=(64, d_k_state, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz128) - base_desc_state0 = base_desc_checkpoint - if cutlass.const_expr(state0 is not None): - initial_state_view = cute.make_tensor( - state0.iterator, - cute.make_layout( - (d_v_state, d_k_state, heads_out, batch_size), - stride=(state0.stride[3], state0.stride[2], state0.stride[1], state0.stride[0]), - ), - ) - base_desc_state0 = tma.create_tensor_map_tiled_from_view(initial_state_view, box_dims=(64, d_k_state, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz128) - prologue_kernel( run_order, order_gen, @@ -3470,7 +3439,6 @@ def prologue( base_desc_dq, base_desc_dk, base_desc_dv, - base_desc_state0, tensormap_workspace, cu_seqlens, q, @@ -3481,7 +3449,6 @@ def prologue( dq, dk, dv, - state0, work_item_staging, work_count, work_items, @@ -3759,7 +3726,6 @@ def kernel( desc_dq_base = desc_base_words + cutlass.Int32(5) * arr_words desc_dk_base = desc_base_words + cutlass.Int32(6) * arr_words desc_dv_base = desc_base_words + cutlass.Int32(7) * arr_words - desc_initial_state_base = desc_base_words + cutlass.Int32(8) * arr_words SMEM = cutlass.AddressSpace.smem bars = make_gdn_bars(cfg) @@ -3993,7 +3959,7 @@ def kernel( base=sDstate_raw.data_ptr().toint(), elems_per_stage=cfg.d_k * cfg.d_v * bpe, stages=1, - leading_byte_offset=LEAD, + leading_byte_offset=STATE_LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) @@ -4269,7 +4235,6 @@ def kernel( desc_v_base=desc_v_base, desc_do_base=desc_do_base, desc_checkpoint_base=desc_checkpoint_base, - desc_initial_state_base=desc_initial_state_base, sSched=sSched, bars=bars, ) @@ -4461,7 +4426,7 @@ def build_cfg( TENSORMAP_DESC_ARRAYS = 8 # per-batch runtime TMA descriptors: Q, K, V, dO, checkpoints, dQ, dK, dV -TENSORMAP_STATIC_SLOTS = 1 # initial_state +TENSORMAP_STATIC_SLOTS = 0 @functools.cache @@ -4586,7 +4551,7 @@ def chunk_gdn_bwd_sm100( cu_seqlens, scale: float, *, - initial_state=None, + use_initial_state: bool = False, d_initial_state=None, d_final_state=None, work_items=None, @@ -4607,10 +4572,10 @@ def chunk_gdn_bwd_sm100( Produces dQ/dK/dV/dGate/dBeta at ``HO = max(HQ, HV)`` heads (the caller reduces over the head group; dGate = dL/d(ln alpha)). With - ``initial_state`` (io dtype ``(num_seqs, HO, DK, DV)``, K-major — the - caller downcasts its fp32 state), chunk 0's forward state loads from it - through a dedicated per-(b,h) descriptor set; ``d_initial_state`` (fp32, - same shape) then also receives the initial-state gradient. The two go together. ``state_checkpoints`` is + ``use_initial_state``, chunk 0 reads its entering state from + ``state_checkpoints`` row 0 like every other chunk (the forward writes the + initial state there); ``d_initial_state`` (fp32) then also receives the + initial-state gradient. The two go together. ``state_checkpoints`` is always the PLAIN per-chunk checkpoint series. All tensors are contiguous, DLPack-compatible CUDA tensors on the same device. Compile-cache-and-replay. @@ -4626,7 +4591,7 @@ def chunk_gdn_bwd_sm100( beta: ``(total_tokens, HO)`` float32, update gate — post-sigmoid, or io-dtype logits when ``use_beta_sigmoid`` do: ``(total_tokens, HO, DV)`` float16/bfloat16, output gradient - state_checkpoints: ``(total_checkpoints, HO, DK, DV)`` io dtype, per-chunk + state_checkpoints: ``(total_checkpoints, HO, DV, DK)`` io dtype, per-chunk forward states from the prefill kernel's checkpoint output (``checkpoint_every_n_tokens=B_T``) dq/dk/dv: pre-allocated output gradients, shaped/typed like q/k/v at HO heads @@ -4635,8 +4600,8 @@ def chunk_gdn_bwd_sm100( dbeta: pre-allocated ``(total_tokens, HO)`` beta gradient; float32, or io dtype and wrt the raw logits under ``use_beta_sigmoid`` cu_seqlens: ``(num_seqs + 1,)`` int32 - initial_state: ``(num_seqs, HO, DK, DV)`` io dtype (matching ``state_checkpoints``), - or None + use_initial_state: the forward ran with an initial state, so chunk 0 + has an entering state to load from ``state_checkpoints`` row 0 scale: attention scale factor (must not be 0) safe_gate: interpret ``gate`` through the safe-gate transform a_log: ``(HO,)`` float32, safe-gate per-head log-amplitude (None = 0) @@ -4684,7 +4649,7 @@ def chunk_gdn_bwd_sm100( HK, HV, is_GQA, - initial_state is not None, + use_initial_state, d_final_state is not None, d_initial_state is not None, log_gate, @@ -4716,7 +4681,7 @@ def chunk_gdn_bwd_sm100( cache["compiled"] = compile( io_dtype, is_GQA, - use_initial_state=initial_state is not None, + use_initial_state=use_initial_state, use_dstate_in=d_final_state is not None, use_dstate0=d_initial_state is not None, log_gate=log_gate, @@ -4756,9 +4721,6 @@ def chunk_gdn_bwd_sm100( if "prologue" not in cache: checkpoints_pl = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=3) cu_pl = from_dlpack(cu_seqlens, assumed_align=4).mark_layout_dynamic() - state0_pl = None - if initial_state is not None: - state0_pl = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) staging_pl = None if run_order and not order_gen: staging_pl = from_dlpack(work_item_scratch, assumed_align=16) @@ -4785,7 +4747,6 @@ def chunk_gdn_bwd_sm100( from_dlpack(dv, assumed_align=16).mark_layout_dynamic(leading_dim=2), checkpoints_pl, cu_pl, - state0_pl, staging_pl, work_count_pl, work_items_pl, @@ -4804,7 +4765,6 @@ def chunk_gdn_bwd_sm100( dv, state_checkpoints, cu_seqlens, - initial_state, work_item_scratch if (run_order and not order_gen) else None, work_count, work_items, @@ -4854,7 +4814,6 @@ def run_bwd( dgate, dbeta, cu_seqlens, - initial_state, d_initial_state, d_final_state, work_items, @@ -4882,7 +4841,6 @@ def run_bwd( dv, state_checkpoints, cu_seqlens, - initial_state, work_item_scratch, work_count, work_items, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py index 2d17dee31..1e6aa5bcc 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py @@ -548,7 +548,6 @@ def tmastg_warp( while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) n_local = wend - cstart - n_padded = ((n_local + 1) // 2) * 2 head_o = head_idx slot = batch_idx * desc_qwords @@ -557,13 +556,25 @@ def tmastg_warp( tma_tensormap_acquire(desc_o_slot) if cutlass.const_expr(cfg.enable_checkpoints): desc_checkpoint_slot = (desc_checkpoint_base + slot).tospace(cutlass.AddressSpace.generic) - checkpoint_coord = wstart - 1 if wstart > 0 else cutlass.Int32(0) + checkpoint_coord = (wstart + ckpt_chunks - cutlass.Int32(1)) // ckpt_chunks checkpoint_mod = (cstart + cutlass.Int32(1)) % ckpt_chunks if elect_one: tma_tensormap_acquire(desc_checkpoint_slot) if n_local > 0: - for local_idx in cutlass.range(n_padded): + if cutlass.const_expr(cfg.enable_checkpoints): + if wstart == 0: + checkpoint_stage = checkpoint_store_cnt % cfg.smem_checkpoint_stages + checkpoint_phase = (checkpoint_store_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1) + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_phase) + checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), checkpoint_coord, head_o) + tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) + tma_store_commit() + tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() + checkpoint_coord += 1 + checkpoint_store_cnt = checkpoint_store_cnt + 1 + for local_idx in cutlass.range(n_local): chunk_idx = cstart + local_idx did_o = cutlass.Int32(0) @@ -651,14 +662,13 @@ def gate_beta_warp( while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) n_local = wend - cstart - n_padded = ((n_local + 1) // 2) * 2 if cutlass.const_expr(cfg.safe_gate): if n_local > 0: # per-head transform constants, fixed for the whole tile a_l2 = -cute.math.exp2(mA_log[head_idx].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) * cutlass.Float32(RCP_LN2) bias = mDt_bias[head_idx].to(cutlass.Float32) if n_local > 0: - for local_idx in cutlass.range(n_padded): + for local_idx in cutlass.range(n_local): # ---- Gate load: GMEM -> SMEM (OOB neutral) ----------------------- chunk_idx = cstart + local_idx n_cols = cfg.b_t // cfg.threads_per_warp @@ -888,19 +898,15 @@ def mma_warp( decay_u_inp_ptr = nvvm.make_tmem_ptr(tmem_inp_col + INP_SLOT_COLS, cutlass.Int8) k_state_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_cg1_acc_offset, cutlass.Float32) u_acc_ptr = k_state_acc_ptr - acc_cg0_0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f, cutlass.Float32) - acc_cg0_1 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + cfg.b_t, cutlass.Float32) sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) n_local = wend - cstart - n_padded = ((n_local + 1) // 2) * 2 - n_pairs = n_padded // 2 - # ---- fused KK^T/QK^T pair 0: both members issued ahead of the loop ------- - if n_pairs > 0: + # ---- fused KK^T/QK^T pair 0: each member issued ahead of the loop ------- + if n_local > 0: f0_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[f0_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) @@ -908,10 +914,12 @@ def mma_warp( bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) desc_kqf = sKQ[kqf_idx].desc() - mma_ss(bmm_qk_desc, desc_kqf, desc_kqf, acc_cg0_0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0_0, accumulate=True, k_count=KQ_HALF_K) + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + f0_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_qk_desc, desc_kqf, desc_kqf, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[f0_acc_idx].arrive(cta_group=1) + if n_local > 1: pend_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[pend_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) @@ -920,12 +928,13 @@ def mma_warp( kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) desc_kqf = sKQ[kqf_idx].desc() desc_kqf_b = desc_kqf + KQ_BOX - mma_ss(bmm_qk_desc, desc_kqf, desc_kqf_b, acc_cg0_1, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf_b + KQ_SEG, acc_cg0_1, accumulate=True, k_count=KQ_HALF_K) + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + pend_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_qk_desc, desc_kqf, desc_kqf_b, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf_b + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[pend_acc_idx].arrive(cta_group=1) - for local_idx in cutlass.range(n_padded): # noqa: B007 + for local_idx in cutlass.range(n_local): # noqa: B007 if cutlass.const_expr(cfg.use_initial_state): if local_idx == 0: if elect_one: @@ -956,7 +965,7 @@ def mma_warp( # ---- QK/KK lookahead (member 1) = [Q;K](S) @ K^T --------------------- if member == 1: - if (local_idx >> 1) + 1 < n_pairs: + if local_idx + 2 < n_local: pend_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[pend_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) @@ -965,8 +974,9 @@ def mma_warp( kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) desc_kqf = sKQ[kqf_idx].desc() desc_kqf_b = desc_kqf + KQ_BOX - mma_ss(bmm_qk_desc, desc_kqf, desc_kqf_b, acc_cg0_1, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf_b + KQ_SEG, acc_cg0_1, accumulate=True, k_count=KQ_HALF_K) + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + pend_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_qk_desc, desc_kqf, desc_kqf_b, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf_b + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[pend_acc_idx].arrive(cta_group=1) @@ -1005,7 +1015,7 @@ def mma_warp( # ---- KK/QK lookahead (member 0) = [K;Q](S) @ K^T --------------------- if member == 0: - if (local_idx >> 1) + 1 < n_pairs: + if local_idx + 2 < n_local: f0_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[f0_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) @@ -1013,8 +1023,9 @@ def mma_warp( bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) desc_kqf = sKQ[kqf_idx].desc() - mma_ss(bmm_qk_desc, desc_kqf, desc_kqf, acc_cg0_0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0_0, accumulate=True, k_count=KQ_HALF_K) + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + f0_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_qk_desc, desc_kqf, desc_kqf, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_qk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[f0_acc_idx].arrive(cta_group=1) @@ -1124,8 +1135,7 @@ def tmaldg_warp( tma_tensormap_acquire(desc_k_slot) tma_tensormap_acquire(desc_v_slot) - wend_padded = cstart + ((wend - cstart + 1) // 2) * 2 - if wend_padded > cstart: + if wend > cstart: kq_idx = kq_index.idx bars.mb_kq_done[kq_idx].wait(kq_index.phase) kq_index = advance(kq_index, cfg.smem_kq_stages) @@ -1137,7 +1147,7 @@ def tmaldg_warp( kq_tile = sKQ_lo_tma[kq_idx] tma_load_tile(kq_tile, k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) tma_load_tile(kq_tile.shifted(kq_box_elems), q_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) - for chunk_idx in cutlass.range(cstart + 1, wend_padded): + for chunk_idx in cutlass.range(cstart + 1, wend): tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) # ---- K + Q interleaved ------------------------------------------- @@ -1172,7 +1182,7 @@ def tmaldg_warp( v_index = advance(v_index, cfg.smem_v_stages) if elect_one: bars.mb_v_ready[v_idx].arrive(n_bytes=cfg.tma_v_bytes) - v_tok = (wend_padded - 1) * cutlass.Int32(cfg.b_t) + v_tok = (wend - cutlass.Int32(1)) * cutlass.Int32(cfg.b_t) v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, v_tok) tma_load_tile(sV_tma[v_idx], v_slice, bars.mb_v_ready[v_idx].smem_ptr, acquire=False) @@ -1249,16 +1259,23 @@ def compute0_warp_group( batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) n_local = wend - cstart n_pairs = (n_local + 1) // 2 - n_padded = n_pairs * 2 for pair_i in cutlass.range(n_pairs): + # An odd chunk count leaves the last pair with member 0 only; have_m1 + # is uniform across CG0, so the shared barriers below stay aligned. + have_m1 = pair_i * 2 + 1 < n_local + do_kk = have_m1 or pair_half == 0 + do_a = have_m1 or pair_half == 1 + # ---- Gate rows for this warp's KK / A member roles ------------------- gate0_idx = gate_index.idx bars.mb_gate_ready[gate0_idx].wait(gate_index.phase) gate_index = advance(gate_index, cfg.smem_gate_stages) - gate1_idx = gate_index.idx - bars.mb_gate_ready[gate1_idx].wait(gate_index.phase) - gate_index = advance(gate_index, cfg.smem_gate_stages) + gate1_idx = gate0_idx + if have_m1: + gate1_idx = gate_index.idx + bars.mb_gate_ready[gate1_idx].wait(gate_index.phase) + gate_index = advance(gate_index, cfg.smem_gate_stages) kk_gate_idx = gate1_idx if pair_half == 1 else gate0_idx a_gate_idx = gate0_idx if pair_half == 1 else gate1_idx @@ -1296,14 +1313,17 @@ def compute0_warp_group( decay_t_kk.append(cute.math.exp2(kk_row_cumsumlog - kk_cumsumlog_cols[col], fastmath=True) if is_lower else mask_zero) decay_t_a.append(cute.math.exp2(a_row_cumsumlog - a_cumsumlog_cols[col], fastmath=True) if is_lower else mask_zero) bars.mb_gate_done[gate0_idx].arrive() - bars.mb_gate_done[gate1_idx].arrive() + if have_m1: + bars.mb_gate_done[gate1_idx].arrive() beta0_idx = beta_index.idx bars.mb_beta_ready[beta0_idx].wait(beta_index.phase) beta_index = advance(beta_index, cfg.smem_beta_stages) - beta1_idx = beta_index.idx - bars.mb_beta_ready[beta1_idx].wait(beta_index.phase) - beta_index = advance(beta_index, cfg.smem_beta_stages) + beta1_idx = beta0_idx + if have_m1: + beta1_idx = beta_index.idx + bars.mb_beta_ready[beta1_idx].wait(beta_index.phase) + beta_index = advance(beta_index, cfg.smem_beta_stages) kk_beta_idx = beta1_idx if pair_half == 1 else beta0_idx kk_beta = [] for r in (row_u0_lo, row_u0_hi, row_u1_lo, row_u1_hi): @@ -1313,9 +1333,12 @@ def compute0_warp_group( acc0_idx = cg0_acc_ready.idx acc0_phase = cg0_acc_ready.phase cg0_acc_ready = advance(cg0_acc_ready, cfg.tmem_cg0_acc_stages) - acc1_idx = cg0_acc_ready.idx - acc1_phase = cg0_acc_ready.phase - cg0_acc_ready = advance(cg0_acc_ready, cfg.tmem_cg0_acc_stages) + acc1_idx = acc0_idx + acc1_phase = acc0_phase + if have_m1: + acc1_idx = cg0_acc_ready.idx + acc1_phase = cg0_acc_ready.phase + cg0_acc_ready = advance(cg0_acc_ready, cfg.tmem_cg0_acc_stages) kk_acc_idx = acc1_idx if pair_half == 1 else acc0_idx kk_acc_phase = acc1_phase if pair_half == 1 else acc0_phase a_acc_idx = acc0_idx if pair_half == 1 else acc1_idx @@ -1323,52 +1346,62 @@ def compute0_warp_group( tinv0_idx = tinv_index.idx tinv0_phase = tinv_index.phase tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) - tinv1_idx = tinv_index.idx - tinv1_phase = tinv_index.phase - tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) + tinv1_idx = tinv0_idx + tinv1_phase = tinv0_phase + if have_m1: + tinv1_idx = tinv_index.idx + tinv1_phase = tinv_index.phase + tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) kk_tinv_idx = tinv1_idx if pair_half == 1 else tinv0_idx kk_tinv_phase = tinv1_phase if pair_half == 1 else tinv0_phase - bars.mb_cg0_acc_ready[kk_acc_idx].wait(kk_acc_phase) tinv0_base = sTinv[tinv0_idx].base tinv1_base = sTinv[tinv1_idx].base kk_base = tinv1_base if pair_half == 1 else tinv0_base - kk_vec0 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) - kk_vec1 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + 16) << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) - bars.mb_t_inv_done[kk_tinv_idx].wait(kk_tinv_phase) - for u in cutlass.range_constexpr(2): - kk_vec = kk_vec1 if cutlass.const_expr(u == 1) else kk_vec0 - kk_pack = [] - for k in cutlass.range_constexpr(num_vals // 2): - b0 = kk_beta[u * 2 + 1] if cutlass.const_expr((k % 2) == 1) else kk_beta[u * 2] - p0, p1 = fmul2(kk_vec[2 * k], kk_vec[2 * k + 1], decay_t_kk[u * num_vals + 2 * k], decay_t_kk[u * num_vals + 2 * k + 1]) - v0, v1 = fmul2(p0, p1, b0, b0) - kk_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) - st_row = half_row_base + u * 16 + store_row_frag - for c in cutlass.range_constexpr(ACC_N_FRAGS): - nvvm.stmatrix( - cutlass.inttoptr( - kk_base + (st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [kk_pack[c * 4 + 0], kk_pack[c * 4 + 1], kk_pack[c * 4 + 2], kk_pack[c * 4 + 3]], - nvvm.MMALayout.ROW, - ) + if do_kk: + bars.mb_cg0_acc_ready[kk_acc_idx].wait(kk_acc_phase) + kk_vec0 = nvvm.tcgen05_ld( + "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 + ) + kk_vec1 = nvvm.tcgen05_ld( + "16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + 16) << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 + ) + bars.mb_t_inv_done[kk_tinv_idx].wait(kk_tinv_phase) + for u in cutlass.range_constexpr(2): + kk_vec = kk_vec1 if cutlass.const_expr(u == 1) else kk_vec0 + kk_pack = [] + for k in cutlass.range_constexpr(num_vals // 2): + b0 = kk_beta[u * 2 + 1] if cutlass.const_expr((k % 2) == 1) else kk_beta[u * 2] + p0, p1 = fmul2(kk_vec[2 * k], kk_vec[2 * k + 1], decay_t_kk[u * num_vals + 2 * k], decay_t_kk[u * num_vals + 2 * k + 1]) + v0, v1 = fmul2(p0, p1, b0, b0) + kk_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) + st_row = half_row_base + u * 16 + store_row_frag + for c in cutlass.range_constexpr(ACC_N_FRAGS): + nvvm.stmatrix( + cutlass.inttoptr( + kk_base + (st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS)) * bpe, + cutlass.AddressSpace.smem, + cutlass.BFloat16, + ), + [kk_pack[c * 4 + 0], kk_pack[c * 4 + 1], kk_pack[c * 4 + 2], kk_pack[c * 4 + 3]], + nvvm.MMALayout.ROW, + ) # ---- pair inverse: warps 0-1 own matrix 0, warps 2-3 matrix 1 -------- - inv_base = tinv1_base if warp_id >= 2 else tinv0_base + # With no member 1 there is one matrix, so warps 2-3 idle through the + # per-warp steps; the barriers below stay unconditional. + inv_base = tinv0_base + if have_m1: + inv_base = tinv1_base if warp_id >= 2 else tinv0_base + do_inv = have_m1 or warp_id < 2 # diagonal 8x8 Gauss-Jordan, all four warps nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - invert_diagonal_NxN(cfg, inv_base, (inverse_local_warp * cfg.threads_per_warp + lane_id) // 8, cg0_tidx, 8) + if do_inv: + invert_diagonal_NxN(cfg, inv_base, (inverse_local_warp * cfg.threads_per_warp + lane_id) // 8, cg0_tidx, 8) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, @@ -1376,20 +1409,24 @@ def compute0_warp_group( # 8x8 -> 16x16 (both matrices per warp) blockwise_diagonal_8x8_to_16x16(cfg, tinv0_base, warp_id * 16, lane_id) - blockwise_diagonal_8x8_to_16x16(cfg, tinv1_base, warp_id * 16, lane_id) + if have_m1: + blockwise_diagonal_8x8_to_16x16(cfg, tinv1_base, warp_id * 16, lane_id) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) # 16x16 -> 32x32, one tile per warp within the group - blockwise_diagonal_16x16_to_32x32(cfg, inv_base, inverse_local_warp * 32, lane_id) + if do_inv: + blockwise_diagonal_16x16_to_32x32(cfg, inv_base, inverse_local_warp * 32, lane_id) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) # 32x32 -> 64x64, two warps per matrix + # all four warps: this step carries its own CG0 barrier, and on the tail + # inv_base is matrix 0 for every warp so the duplicate band is identical blockwise_diagonal_32x32_to_64x64(cfg, inv_base, inverse_local_warp, lane_id) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, @@ -1432,82 +1469,87 @@ def compute0_warp_group( bars.mb_t_inv_ready[tinv0_idx].arrive() bars.mb_beta_done[beta0_idx].arrive() - # ---- Beta column-scaling + publish, stage 1 -------------------------- - beta_col = [] - for k in cutlass.range_constexpr(num_vals): - beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta1_idx]) - tinv_frags = [] - for c in cutlass.range_constexpr(ACC_N_FRAGS): - tinv_frags += list( - nvvm.ldmatrix( + if have_m1: + # ---- Beta column-scaling + publish, stage 1 -------------------------- + beta_col = [] + for k in cutlass.range_constexpr(num_vals): + beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta1_idx]) + tinv_frags = [] + for c in cutlass.range_constexpr(ACC_N_FRAGS): + tinv_frags += list( + nvvm.ldmatrix( + cutlass.inttoptr( + tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, + cutlass.AddressSpace.smem, + cutlass.BFloat16, + ), + 4, + nvvm.MMALayout.ROW, + ) + ) + tinv_pack = [] + for j in cutlass.range_constexpr(num_vals // 2): + lo, hi = f16x2_to_f32(tinv_frags[j], dtype=cfg.io_dtype) + s0, s1 = fmul2(lo, hi, beta_col[2 * j], beta_col[2 * j + 1]) + tinv_pack.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) + for c in cutlass.range_constexpr(ACC_N_FRAGS): + nvvm.stmatrix( cutlass.inttoptr( tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), - 4, + [tinv_pack[c * 4 + 0], tinv_pack[c * 4 + 1], tinv_pack[c * 4 + 2], tinv_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) - ) - tinv_pack = [] - for j in cutlass.range_constexpr(num_vals // 2): - lo, hi = f16x2_to_f32(tinv_frags[j], dtype=cfg.io_dtype) - s0, s1 = fmul2(lo, hi, beta_col[2 * j], beta_col[2 * j + 1]) - tinv_pack.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) - for c in cutlass.range_constexpr(ACC_N_FRAGS): - nvvm.stmatrix( - cutlass.inttoptr( - tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [tinv_pack[c * 4 + 0], tinv_pack[c * 4 + 1], tinv_pack[c * 4 + 2], tinv_pack[c * 4 + 3]], - nvvm.MMALayout.ROW, - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_t_inv_ready[tinv1_idx].arrive() - bars.mb_beta_done[beta1_idx].arrive() + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_t_inv_ready[tinv1_idx].arrive() + bars.mb_beta_done[beta1_idx].arrive() # ---- A_epi (opposite member, both halves in parallel) ---------------- a0_idx = a_index.idx a0_phase = a_index.phase a_index = advance(a_index, cfg.smem_a_stages) - a1_idx = a_index.idx - a1_phase = a_index.phase - a_index = advance(a_index, cfg.smem_a_stages) + a1_idx = a0_idx + a1_phase = a0_phase + if have_m1: + a1_idx = a_index.idx + a1_phase = a_index.phase + a_index = advance(a_index, cfg.smem_a_stages) my_a_idx = a0_idx if pair_half == 1 else a1_idx my_a_phase = a0_phase if pair_half == 1 else a1_phase - a_base = sA[my_a_idx].base - a_vec0 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + a_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) - a_vec1 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + 16) << 16) + tmem_cg0_acc_col + a_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) - nvvm.tcgen05_wait("load") - bars.mb_cg0_acc_done[a_acc_idx].arrive() - bars.mb_a_done[my_a_idx].wait(my_a_phase) - for u in cutlass.range_constexpr(2): - a_vec = a_vec1 if cutlass.const_expr(u == 1) else a_vec0 - a_pack = [] - for k in cutlass.range_constexpr(num_vals // 2): - p0, p1 = fmul2(a_vec[2 * k], a_vec[2 * k + 1], decay_t_a[u * num_vals + 2 * k], decay_t_a[u * num_vals + 2 * k + 1]) - v0, v1 = fmul2(p0, p1, scale, scale) - a_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) - st_row = half_row_base + u * 16 + store_row_frag - for c in cutlass.range_constexpr(ACC_N_FRAGS): - nvvm.stmatrix( - cutlass.inttoptr( - a_base + (st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [a_pack[c * 4 + 0], a_pack[c * 4 + 1], a_pack[c * 4 + 2], a_pack[c * 4 + 3]], - nvvm.MMALayout.ROW, - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_a_ready[my_a_idx].arrive() + if do_a: + a_base = sA[my_a_idx].base + a_vec0 = nvvm.tcgen05_ld( + "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + a_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 + ) + a_vec1 = nvvm.tcgen05_ld( + "16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + 16) << 16) + tmem_cg0_acc_col + a_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 + ) + nvvm.tcgen05_wait("load") + bars.mb_cg0_acc_done[a_acc_idx].arrive() + bars.mb_a_done[my_a_idx].wait(my_a_phase) + for u in cutlass.range_constexpr(2): + a_vec = a_vec1 if cutlass.const_expr(u == 1) else a_vec0 + a_pack = [] + for k in cutlass.range_constexpr(num_vals // 2): + p0, p1 = fmul2(a_vec[2 * k], a_vec[2 * k + 1], decay_t_a[u * num_vals + 2 * k], decay_t_a[u * num_vals + 2 * k + 1]) + v0, v1 = fmul2(p0, p1, scale, scale) + a_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) + st_row = half_row_base + u * 16 + store_row_frag + for c in cutlass.range_constexpr(ACC_N_FRAGS): + nvvm.stmatrix( + cutlass.inttoptr( + a_base + (st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS)) * bpe, + cutlass.AddressSpace.smem, + cutlass.BFloat16, + ), + [a_pack[c * 4 + 0], a_pack[c * 4 + 1], a_pack[c * 4 + 2], a_pack[c * 4 + 3]], + nvvm.MMALayout.ROW, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_a_ready[my_a_idx].arrive() tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) for _ in range(cfg.smem_t_inv_stages): @@ -1602,7 +1644,6 @@ def compute1_warp_group( while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) n_local = wend - cstart - n_padded = ((n_local + 1) // 2) * 2 if cutlass.const_expr(cfg.enable_checkpoints): ckpt_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) checkpoint_mod = cstart % ckpt_chunks @@ -1618,7 +1659,7 @@ def compute1_warp_group( for sub in cutlass.range_constexpr(num_state_subs): words = [] for k in cutlass.range_constexpr(32): - v = gState_init[sub * ldtm_width + k, cg1_tidx] + v = gState_init[cg1_tidx, sub * ldtm_width + k] if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): v = v.to(cfg.acc_dtype) words.append(v) @@ -1641,12 +1682,26 @@ def compute1_warp_group( thread_count=cfg.init_state_store_barrier_threads, ) - for local_idx in cutlass.range(n_padded): # noqa: B007 + for local_idx in cutlass.range(n_local): # noqa: B007 chunk_idx = cstart + local_idx if cutlass.const_expr(cfg.enable_checkpoints): do_checkpoint_now = checkpoint_mod == 0 checkpoint_mod = checkpoint_mod + cutlass.Int32(1) checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == ckpt_chunks else checkpoint_mod + if cutlass.const_expr(cfg.enable_checkpoints and not cfg.use_initial_state): + if chunk_idx == 0 and wstart == 0: + checkpoint_stage = checkpoint_cnt % cfg.smem_checkpoint_stages + checkpoint_phase_done = cutlass.Int32(1) ^ ((checkpoint_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1)) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_phase_done) + checkpoint_zero_ptr = cutlass.inttoptr( + sCheckpoint_base_int + checkpoint_stage * cfg.d_k * cfg.d_v * 2, cutlass.AddressSpace.smem, cutlass.Int32 + ) + for z in cutlass.range_constexpr(cfg.d_k * cfg.d_v // 2 // num_threads_cg1): + (checkpoint_zero_ptr + cg1_tidx + z * num_threads_cg1).store(cutlass.Int32(0)) + nvvm.fence_proxy("async.shared", space="cta") + if elect_one: + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + checkpoint_cnt = checkpoint_cnt + 1 valid_state = local_idx > 0 if cutlass.const_expr(cfg.use_initial_state): valid_state = cutlass.Boolean(True) @@ -1686,44 +1741,28 @@ def compute1_warp_group( if cutlass.const_expr(cfg.enable_checkpoints): # ---- state checkpoint ---------------------------------------- - do_checkpoint = do_checkpoint_now and chunk_idx > 0 and chunk_idx < wend + do_checkpoint = do_checkpoint_now and chunk_idx < wend do_checkpoint = do_checkpoint and chunk_idx >= wstart if do_checkpoint: - checkpoint_pack = [[cutlass.Int32(0) for _ in range(16)] for _ in range(4)] - for b in cutlass.range_constexpr(2): - for col_half in cutlass.range_constexpr(2): - checkpoint_vec = nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + b * 16) << 16) + tmem_state_col + col_half * 64, cutlass.Float32), - num=8, - ) - for j in cutlass.range_constexpr(16): - checkpoint_pack[b * 2 + col_half][j] = fp32_to_fp16( - checkpoint_vec[2 * j], checkpoint_vec[2 * j + 1], dtype=cfg.io_dtype - ) checkpoint_stage = checkpoint_cnt % cfg.smem_checkpoint_stages checkpoint_phase_done = cutlass.Int32(1) ^ ((checkpoint_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1)) bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_phase_done) - for b in cutlass.range_constexpr(2): - for col_half in cutlass.range_constexpr(2): - checkpoint_base = checkpoint_stage * cfg.d_k * cfg.d_v + (cg1_tidx // 64) * cfg.d_k * 64 - for c in cutlass.range_constexpr(4): - checkpoint_row = col_half * 64 + checkpoint_smem_row + c * 16 - nvvm.stmatrix( - cutlass.inttoptr( - sCheckpoint_base_int - + (checkpoint_base + checkpoint_row * 64 + swizzle_xor_128b(checkpoint_row, checkpoint_smem_col + b * 16)) * 2, - cutlass.AddressSpace.smem, - cfg.io_dtype, - ), - [ - checkpoint_pack[b * 2 + col_half][c * 4 + 0], - checkpoint_pack[b * 2 + col_half][c * 4 + 1], - checkpoint_pack[b * 2 + col_half][c * 4 + 2], - checkpoint_pack[b * 2 + col_half][c * 4 + 3], - ], - nvvm.MMALayout.COL, - ) + checkpoint_stage_base = checkpoint_stage * cfg.d_k * cfg.d_v + for sub in cutlass.range_constexpr(num_state_subs): + checkpoint_vec = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + num=32, + ) + for g in cutlass.range_constexpr(ldtm_width // 8): + packs = tuple( + fp32_to_fp16(checkpoint_vec[g * 8 + 2 * t], checkpoint_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dk = sub * ldtm_width + g * 8 + checkpoint_addr = checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + cg1_tidx * 64 + swizzle_xor_128b(cg1_tidx, dk % 64) + (sCheckpoint_raw.data_ptr() + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) nvvm.fence_proxy("async.shared", space="cta") if elect_one: bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() @@ -1918,7 +1957,7 @@ def compute1_warp_group( val = state_vec[k] if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): val = val.to(cfg.state_dtype) - gState_out[sub * ldtm_width + k, cg1_tidx] = val + gState_out[cg1_tidx, sub * ldtm_width + k] = val bars.mb_state_acc_scale_done[kv_last_idx].arrive() else: bars.mb_state_acc_scale_done[kv_last_idx].arrive() @@ -1931,11 +1970,11 @@ def compute1_warp_group( gState_in = mState_init[None, None, head_idx, batch_idx] for sub in cutlass.range_constexpr(num_state_subs): for k in cutlass.range_constexpr(32): - gState_out[sub * ldtm_width + k, cg1_tidx] = gState_in[sub * ldtm_width + k, cg1_tidx] + gState_out[cg1_tidx, sub * ldtm_width + k] = gState_in[cg1_tidx, sub * ldtm_width + k] else: for sub in cutlass.range_constexpr(num_state_subs): for k in cutlass.range_constexpr(32): - gState_out[sub * ldtm_width + k, cg1_tidx] = cutlass.Float32(0.0).to(cfg.state_dtype) + gState_out[cg1_tidx, sub * ldtm_width + k] = cutlass.Float32(0.0).to(cfg.state_dtype) tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py index 94a2d67ca..eb71a4b38 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py @@ -482,26 +482,25 @@ def tmastg_warp( while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) n_local = wend - cstart - n_padded = ((n_local + 1) // 2) * 2 head_o = head_idx slot = batch_idx * desc_qwords if cutlass.const_expr(cfg.enable_checkpoints): desc_checkpoint_slot = (desc_checkpoint_base + slot).tospace(cutlass.AddressSpace.generic) - checkpoint_coord = wstart - 1 if wstart > 0 else cutlass.Int32(0) - checkpoint_mod = (cstart + cutlass.Int32(1)) % ckpt_chunks + checkpoint_coord = (wstart + ckpt_chunks - cutlass.Int32(1)) // ckpt_chunks + checkpoint_mod = cstart % ckpt_chunks if elect_one: tma_tensormap_acquire(desc_checkpoint_slot) if n_local > 0: - for local_idx in cutlass.range(n_padded): + for local_idx in cutlass.range(n_local): chunk_idx = cstart + local_idx did_checkpoint = cutlass.Int32(0) if cutlass.const_expr(cfg.enable_checkpoints): checkpoint_stage = checkpoint_store_cnt % cfg.smem_checkpoint_stages checkpoint_phase = (checkpoint_store_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1) - if chunk_idx >= wstart - 1 and chunk_idx < wend - 1: + if chunk_idx >= wstart and chunk_idx < wend: if checkpoint_mod == 0: bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_phase) checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), checkpoint_coord, head_o) @@ -554,14 +553,13 @@ def gate_beta_warp( while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) n_local = wend - cstart - n_padded = ((n_local + 1) // 2) * 2 if cutlass.const_expr(cfg.safe_gate): if n_local > 0: # per-head transform constants, fixed for the whole tile a_l2 = -cute.math.exp2(mA_log[head_idx].to(cutlass.Float32) * cutlass.Float32(RCP_LN2), fastmath=True) * cutlass.Float32(RCP_LN2) bias = mDt_bias[head_idx].to(cutlass.Float32) if n_local > 0: - for local_idx in cutlass.range(n_padded): + for local_idx in cutlass.range(n_local): # ---- Gate load: GMEM -> SMEM (OOB neutral) ----------------------- chunk_idx = cstart + local_idx n_cols = cfg.b_t // cfg.threads_per_warp @@ -784,19 +782,15 @@ def mma_warp( decay_u_inp_ptr = nvvm.make_tmem_ptr(tmem_inp_col + INP_SLOT_COLS, cutlass.Int8) k_state_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_cg1_acc_offset, cutlass.Float32) u_acc_ptr = k_state_acc_ptr - acc_cg0_0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f, cutlass.Float32) - acc_cg0_1 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + cfg.b_t, cutlass.Float32) sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) n_local = wend - cstart - n_padded = ((n_local + 1) // 2) * 2 - n_pairs = n_padded // 2 - # ---- KK pair 0 = K(S) @ K^T (both members issued ahead of the loop) ------ - if n_pairs > 0: + # ---- KK pair 0 = K(S) @ K^T, each member issued ahead of the loop ------- + if n_local > 0: member0_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[member0_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) @@ -804,10 +798,12 @@ def mma_warp( bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) desc_kqf = sKQ[kqf_idx].desc() - mma_ss(bmm_kk_desc, desc_kqf, desc_kqf, acc_cg0_0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0_0, accumulate=True, k_count=KQ_HALF_K) + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + member0_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_kk_desc, desc_kqf, desc_kqf, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[member0_acc_idx].arrive(cta_group=1) + if n_local > 1: member1_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[member1_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) @@ -816,12 +812,13 @@ def mma_warp( kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) desc_kqf = sKQ[kqf_idx].desc() desc_kqf_member1 = desc_kqf + KQ_BOX - mma_ss(bmm_kk_desc, desc_kqf, desc_kqf_member1, acc_cg0_1, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf_member1 + KQ_SEG, acc_cg0_1, accumulate=True, k_count=KQ_HALF_K) + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + member1_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_kk_desc, desc_kqf, desc_kqf_member1, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf_member1 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[member1_acc_idx].arrive(cta_group=1) - for local_idx in cutlass.range(n_padded): # noqa: B007 + for local_idx in cutlass.range(n_local): # noqa: B007 if cutlass.const_expr(cfg.use_initial_state): if local_idx == 0: if elect_one: @@ -845,7 +842,7 @@ def mma_warp( # ---- KK pair lookahead (member 1) = K(S) @ K^T ----------------------- if member == 1: - if (local_idx >> 1) + 1 < n_pairs: + if local_idx + 2 < n_local: member1_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[member1_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) @@ -854,8 +851,9 @@ def mma_warp( kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) desc_kqf = sKQ[kqf_idx].desc() desc_kqf_member1 = desc_kqf + KQ_BOX - mma_ss(bmm_kk_desc, desc_kqf, desc_kqf_member1, acc_cg0_1, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf_member1 + KQ_SEG, acc_cg0_1, accumulate=True, k_count=KQ_HALF_K) + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + member1_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_kk_desc, desc_kqf, desc_kqf_member1, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf_member1 + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[member1_acc_idx].arrive(cta_group=1) @@ -884,7 +882,7 @@ def mma_warp( # ---- KK pair lookahead (member 0) = K(S) @ K^T ----------------------- if member == 0: - if (local_idx >> 1) + 1 < n_pairs: + if local_idx + 2 < n_local: member0_acc_idx = cg0_acc_index.idx bars.mb_cg0_acc_done[member0_acc_idx].wait(cg0_acc_index.phase) cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) @@ -892,8 +890,9 @@ def mma_warp( bars.mb_kq_ready[kqf_idx].wait(kq_fused_index.phase) kq_fused_index = advance(kq_fused_index, cfg.smem_kq_stages) desc_kqf = sKQ[kqf_idx].desc() - mma_ss(bmm_kk_desc, desc_kqf, desc_kqf, acc_cg0_0, accumulate=False, k_count=KQ_HALF_K) - mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0_0, accumulate=True, k_count=KQ_HALF_K) + acc_cg0 = nvvm.make_tmem_ptr(tmem_cg0_acc_col_f + member0_acc_idx * cfg.b_t, cutlass.Float32) + mma_ss(bmm_kk_desc, desc_kqf, desc_kqf, acc_cg0, accumulate=False, k_count=KQ_HALF_K) + mma_ss(bmm_kk_desc, desc_kqf + KQ_SEG, desc_kqf + KQ_SEG, acc_cg0, accumulate=True, k_count=KQ_HALF_K) if elect_one: bars.mb_cg0_acc_ready[member0_acc_idx].arrive(cta_group=1) @@ -985,8 +984,7 @@ def tmaldg_warp( tma_tensormap_acquire(desc_k_slot) tma_tensormap_acquire(desc_v_slot) - wend_padded = cstart + ((wend - cstart + 1) // 2) * 2 - if wend_padded > cstart: + if wend > cstart: kq_idx = kq_index.idx bars.mb_kq_done[kq_idx].wait(kq_index.phase) kq_index = advance(kq_index, cfg.smem_kq_stages) @@ -996,7 +994,7 @@ def tmaldg_warp( k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, tok_coord) kq_tile = sKQ_lo_tma[kq_idx] tma_load_tile(kq_tile, k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) - for chunk_idx in cutlass.range(cstart + 1, wend_padded): + for chunk_idx in cutlass.range(cstart + 1, wend): tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) # ---- K load ------------------------------------------------------ @@ -1028,7 +1026,7 @@ def tmaldg_warp( v_index = advance(v_index, cfg.smem_v_stages) if elect_one: bars.mb_v_ready[v_idx].arrive(n_bytes=cfg.tma_v_bytes) - v_tok = (wend_padded - 1) * cutlass.Int32(cfg.b_t) + v_tok = (wend - cutlass.Int32(1)) * cutlass.Int32(cfg.b_t) v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, v_tok) tma_load_tile(sV_tma[v_idx], v_slice, bars.mb_v_ready[v_idx].smem_ptr, acquire=False) @@ -1102,16 +1100,22 @@ def compute0_warp_group( batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) n_local = wend - cstart n_pairs = (n_local + 1) // 2 - n_padded = n_pairs * 2 for pair_i in cutlass.range(n_pairs): + # An odd chunk count leaves the last pair with member 0 only; have_m1 + # is uniform across CG0, so the shared barriers below stay aligned. + have_m1 = pair_i * 2 + 1 < n_local + do_kk = have_m1 or pair_half == 0 + # ---- Gate rows for this warp's KK member role ------------------------ gate0_idx = gate_index.idx bars.mb_gate_ready[gate0_idx].wait(gate_index.phase) gate_index = advance(gate_index, cfg.smem_gate_stages) - gate1_idx = gate_index.idx - bars.mb_gate_ready[gate1_idx].wait(gate_index.phase) - gate_index = advance(gate_index, cfg.smem_gate_stages) + gate1_idx = gate0_idx + if have_m1: + gate1_idx = gate_index.idx + bars.mb_gate_ready[gate1_idx].wait(gate_index.phase) + gate_index = advance(gate_index, cfg.smem_gate_stages) kk_gate_idx = gate1_idx if pair_half == 1 else gate0_idx row_u0_lo = half_row_base + lane_id // 4 @@ -1141,14 +1145,17 @@ def compute0_warp_group( col = (k // 4) * 2 + (k % 2) decay_t_kk.append(cute.math.exp2(row_cumsumlog - kk_col_cumsumlog[col], fastmath=True) if is_lower else mask_zero) bars.mb_gate_done[gate0_idx].arrive() - bars.mb_gate_done[gate1_idx].arrive() + if have_m1: + bars.mb_gate_done[gate1_idx].arrive() beta0_idx = beta_index.idx bars.mb_beta_ready[beta0_idx].wait(beta_index.phase) beta_index = advance(beta_index, cfg.smem_beta_stages) - beta1_idx = beta_index.idx - bars.mb_beta_ready[beta1_idx].wait(beta_index.phase) - beta_index = advance(beta_index, cfg.smem_beta_stages) + beta1_idx = beta0_idx + if have_m1: + beta1_idx = beta_index.idx + bars.mb_beta_ready[beta1_idx].wait(beta_index.phase) + beta_index = advance(beta_index, cfg.smem_beta_stages) kk_beta_idx = beta1_idx if pair_half == 1 else beta0_idx kk_beta = [] for r in (row_u0_lo, row_u0_hi, row_u1_lo, row_u1_hi): @@ -1158,65 +1165,76 @@ def compute0_warp_group( acc0_idx = cg0_acc_ready.idx acc0_phase = cg0_acc_ready.phase cg0_acc_ready = advance(cg0_acc_ready, cfg.tmem_cg0_acc_stages) - acc1_idx = cg0_acc_ready.idx - acc1_phase = cg0_acc_ready.phase - cg0_acc_ready = advance(cg0_acc_ready, cfg.tmem_cg0_acc_stages) + acc1_idx = acc0_idx + acc1_phase = acc0_phase + if have_m1: + acc1_idx = cg0_acc_ready.idx + acc1_phase = cg0_acc_ready.phase + cg0_acc_ready = advance(cg0_acc_ready, cfg.tmem_cg0_acc_stages) kk_acc_idx = acc1_idx if pair_half == 1 else acc0_idx kk_acc_phase = acc1_phase if pair_half == 1 else acc0_phase tinv0_idx = tinv_index.idx tinv0_phase = tinv_index.phase tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) - tinv1_idx = tinv_index.idx - tinv1_phase = tinv_index.phase - tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) + tinv1_idx = tinv0_idx + tinv1_phase = tinv0_phase + if have_m1: + tinv1_idx = tinv_index.idx + tinv1_phase = tinv_index.phase + tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) kk_tinv_idx = tinv1_idx if pair_half == 1 else tinv0_idx kk_tinv_phase = tinv1_phase if pair_half == 1 else tinv0_phase - bars.mb_cg0_acc_ready[kk_acc_idx].wait(kk_acc_phase) - tinv0_base = sTinv[tinv0_idx].base tinv1_base = sTinv[tinv1_idx].base kk_base = tinv1_base if pair_half == 1 else tinv0_base - - kk_vec0 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) - kk_vec1 = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + 16) << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 - ) - nvvm.tcgen05_wait("load") - bars.mb_cg0_acc_done[kk_acc_idx].arrive() - bars.mb_t_inv_done[kk_tinv_idx].wait(kk_tinv_phase) - for u in cutlass.range_constexpr(2): - kk_vec = kk_vec1 if cutlass.const_expr(u == 1) else kk_vec0 - kk_pack = [] - for k in cutlass.range_constexpr(num_vals // 2): - row_beta = kk_beta[u * 2 + 1] if cutlass.const_expr((k % 2) == 1) else kk_beta[u * 2] - p0, p1 = fmul2(kk_vec[2 * k], kk_vec[2 * k + 1], decay_t_kk[u * num_vals + 2 * k], decay_t_kk[u * num_vals + 2 * k + 1]) - v0, v1 = fmul2(p0, p1, row_beta, row_beta) - kk_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) - st_row = half_row_base + u * 16 + store_row_frag - for c in cutlass.range_constexpr(ACC_N_FRAGS): - nvvm.stmatrix( - cutlass.inttoptr( - kk_base + (st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [kk_pack[c * 4 + 0], kk_pack[c * 4 + 1], kk_pack[c * 4 + 2], kk_pack[c * 4 + 3]], - nvvm.MMALayout.ROW, - ) + if do_kk: + bars.mb_cg0_acc_ready[kk_acc_idx].wait(kk_acc_phase) + kk_vec0 = nvvm.tcgen05_ld( + "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 + ) + kk_vec1 = nvvm.tcgen05_ld( + "16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + 16) << 16) + tmem_cg0_acc_col + kk_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 + ) + nvvm.tcgen05_wait("load") + bars.mb_cg0_acc_done[kk_acc_idx].arrive() + bars.mb_t_inv_done[kk_tinv_idx].wait(kk_tinv_phase) + for u in cutlass.range_constexpr(2): + kk_vec = kk_vec1 if cutlass.const_expr(u == 1) else kk_vec0 + kk_pack = [] + for k in cutlass.range_constexpr(num_vals // 2): + row_beta = kk_beta[u * 2 + 1] if cutlass.const_expr((k % 2) == 1) else kk_beta[u * 2] + p0, p1 = fmul2(kk_vec[2 * k], kk_vec[2 * k + 1], decay_t_kk[u * num_vals + 2 * k], decay_t_kk[u * num_vals + 2 * k + 1]) + v0, v1 = fmul2(p0, p1, row_beta, row_beta) + kk_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) + st_row = half_row_base + u * 16 + store_row_frag + for c in cutlass.range_constexpr(ACC_N_FRAGS): + nvvm.stmatrix( + cutlass.inttoptr( + kk_base + (st_row * cfg.b_t + swizzle_xor_128b(st_row, store_col + c * FRAG_COLS)) * bpe, + cutlass.AddressSpace.smem, + cutlass.BFloat16, + ), + [kk_pack[c * 4 + 0], kk_pack[c * 4 + 1], kk_pack[c * 4 + 2], kk_pack[c * 4 + 3]], + nvvm.MMALayout.ROW, + ) # ---- pair inverse: warps 0-1 own matrix 0, warps 2-3 matrix 1 -------- - inv_base = tinv1_base if warp_id >= 2 else tinv0_base + # With no member 1 there is one matrix, so warps 2-3 idle through the + # per-warp steps; the barriers below stay unconditional. + inv_base = tinv0_base + if have_m1: + inv_base = tinv1_base if warp_id >= 2 else tinv0_base + do_inv = have_m1 or warp_id < 2 # diagonal 8x8 Gauss-Jordan, all four warps nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - invert_diagonal_NxN(cfg, inv_base, (inverse_local_warp * cfg.threads_per_warp + lane_id) // 8, cg0_tidx, 8) + if do_inv: + invert_diagonal_NxN(cfg, inv_base, (inverse_local_warp * cfg.threads_per_warp + lane_id) // 8, cg0_tidx, 8) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, @@ -1224,20 +1242,24 @@ def compute0_warp_group( # 8x8 -> 16x16 (both matrices per warp) blockwise_diagonal_8x8_to_16x16(cfg, tinv0_base, warp_id * 16, lane_id) - blockwise_diagonal_8x8_to_16x16(cfg, tinv1_base, warp_id * 16, lane_id) + if have_m1: + blockwise_diagonal_8x8_to_16x16(cfg, tinv1_base, warp_id * 16, lane_id) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) # 16x16 -> 32x32, one tile per warp within the group - blockwise_diagonal_16x16_to_32x32(cfg, inv_base, inverse_local_warp * 32, lane_id) + if do_inv: + blockwise_diagonal_16x16_to_32x32(cfg, inv_base, inverse_local_warp * 32, lane_id) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) # 32x32 -> 64x64, two warps per matrix + # all four warps: this step carries its own CG0 barrier, and on the tail + # inv_base is matrix 0 for every warp so the duplicate band is identical blockwise_diagonal_32x32_to_64x64(cfg, inv_base, inverse_local_warp, lane_id) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, @@ -1280,41 +1302,42 @@ def compute0_warp_group( bars.mb_t_inv_ready[tinv0_idx].arrive() bars.mb_beta_done[beta0_idx].arrive() - # ---- Beta column-scaling + publish, stage 1 -------------------------- - beta_col = [] - for k in cutlass.range_constexpr(num_vals): - beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta1_idx]) - tinv_frags = [] - for c in cutlass.range_constexpr(ACC_N_FRAGS): - tinv_frags += list( - nvvm.ldmatrix( + if have_m1: + # ---- Beta column-scaling + publish, stage 1 -------------------------- + beta_col = [] + for k in cutlass.range_constexpr(num_vals): + beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta1_idx]) + tinv_frags = [] + for c in cutlass.range_constexpr(ACC_N_FRAGS): + tinv_frags += list( + nvvm.ldmatrix( + cutlass.inttoptr( + tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, + cutlass.AddressSpace.smem, + cutlass.BFloat16, + ), + 4, + nvvm.MMALayout.ROW, + ) + ) + tinv_pack = [] + for j in cutlass.range_constexpr(num_vals // 2): + lo, hi = f16x2_to_f32(tinv_frags[j], dtype=cfg.io_dtype) + s0, s1 = fmul2(lo, hi, beta_col[2 * j], beta_col[2 * j + 1]) + tinv_pack.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) + for c in cutlass.range_constexpr(ACC_N_FRAGS): + nvvm.stmatrix( cutlass.inttoptr( tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), - 4, + [tinv_pack[c * 4 + 0], tinv_pack[c * 4 + 1], tinv_pack[c * 4 + 2], tinv_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) - ) - tinv_pack = [] - for j in cutlass.range_constexpr(num_vals // 2): - lo, hi = f16x2_to_f32(tinv_frags[j], dtype=cfg.io_dtype) - s0, s1 = fmul2(lo, hi, beta_col[2 * j], beta_col[2 * j + 1]) - tinv_pack.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) - for c in cutlass.range_constexpr(ACC_N_FRAGS): - nvvm.stmatrix( - cutlass.inttoptr( - tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [tinv_pack[c * 4 + 0], tinv_pack[c * 4 + 1], tinv_pack[c * 4 + 2], tinv_pack[c * 4 + 3]], - nvvm.MMALayout.ROW, - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_t_inv_ready[tinv1_idx].arrive() - bars.mb_beta_done[beta1_idx].arrive() + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_t_inv_ready[tinv1_idx].arrive() + bars.mb_beta_done[beta1_idx].arrive() tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) for _ in range(cfg.smem_t_inv_stages): @@ -1396,7 +1419,6 @@ def compute1_warp_group( while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) n_local = wend - cstart - n_padded = ((n_local + 1) // 2) * 2 if cutlass.const_expr(cfg.enable_checkpoints): ckpt_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) checkpoint_mod = cstart % ckpt_chunks @@ -1412,7 +1434,7 @@ def compute1_warp_group( for sub in cutlass.range_constexpr(num_state_subs): words = [] for k in cutlass.range_constexpr(32): - v = gState_init[sub * ldtm_width + k, cg1_tidx] + v = gState_init[cg1_tidx, sub * ldtm_width + k] if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): v = v.to(cfg.acc_dtype) words.append(v) @@ -1435,12 +1457,26 @@ def compute1_warp_group( thread_count=cfg.init_state_store_barrier_threads, ) - for local_idx in cutlass.range(n_padded): # noqa: B007 + for local_idx in cutlass.range(n_local): # noqa: B007 chunk_idx = cstart + local_idx if cutlass.const_expr(cfg.enable_checkpoints): do_checkpoint_now = checkpoint_mod == 0 checkpoint_mod = checkpoint_mod + cutlass.Int32(1) checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == ckpt_chunks else checkpoint_mod + if cutlass.const_expr(cfg.enable_checkpoints and not cfg.use_initial_state): + if chunk_idx == 0 and wstart == 0: + checkpoint_stage = checkpoint_cnt % cfg.smem_checkpoint_stages + checkpoint_phase_done = cutlass.Int32(1) ^ ((checkpoint_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1)) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_phase_done) + checkpoint_zero_ptr = cutlass.inttoptr( + sCheckpoint_base_int + checkpoint_stage * cfg.d_k * cfg.d_v * 2, cutlass.AddressSpace.smem, cutlass.Int32 + ) + for z in cutlass.range_constexpr(cfg.d_k * cfg.d_v // 2 // num_threads_cg1): + (checkpoint_zero_ptr + cg1_tidx + z * num_threads_cg1).store(cutlass.Int32(0)) + nvvm.fence_proxy("async.shared", space="cta") + if elect_one: + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + checkpoint_cnt = checkpoint_cnt + 1 valid_state = local_idx > 0 if cutlass.const_expr(cfg.use_initial_state): valid_state = cutlass.Boolean(True) @@ -1480,44 +1516,28 @@ def compute1_warp_group( if cutlass.const_expr(cfg.enable_checkpoints): # ---- state checkpoint ---------------------------------------- - do_checkpoint = do_checkpoint_now and chunk_idx > 0 and chunk_idx < wend + do_checkpoint = do_checkpoint_now and chunk_idx < wend do_checkpoint = do_checkpoint and chunk_idx >= wstart if do_checkpoint: - checkpoint_pack = [[cutlass.Int32(0) for _ in range(16)] for _ in range(4)] - for b in cutlass.range_constexpr(2): - for col_half in cutlass.range_constexpr(2): - checkpoint_vec = nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + b * 16) << 16) + tmem_state_col + col_half * 64, cutlass.Float32), - num=8, - ) - for j in cutlass.range_constexpr(16): - checkpoint_pack[b * 2 + col_half][j] = fp32_to_fp16( - checkpoint_vec[2 * j], checkpoint_vec[2 * j + 1], dtype=cfg.io_dtype - ) checkpoint_stage = checkpoint_cnt % cfg.smem_checkpoint_stages checkpoint_phase_done = cutlass.Int32(1) ^ ((checkpoint_cnt // cfg.smem_checkpoint_stages) & cutlass.Int32(1)) bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_phase_done) - for b in cutlass.range_constexpr(2): - for col_half in cutlass.range_constexpr(2): - checkpoint_base = checkpoint_stage * cfg.d_k * cfg.d_v + (cg1_tidx // 64) * cfg.d_k * 64 - for c in cutlass.range_constexpr(4): - checkpoint_row = col_half * 64 + checkpoint_frag_row + c * 16 - nvvm.stmatrix( - cutlass.inttoptr( - sCheckpoint_base_int - + (checkpoint_base + checkpoint_row * 64 + swizzle_xor_128b(checkpoint_row, checkpoint_frag_col + b * 16)) * 2, - cutlass.AddressSpace.smem, - cfg.io_dtype, - ), - [ - checkpoint_pack[b * 2 + col_half][c * 4 + 0], - checkpoint_pack[b * 2 + col_half][c * 4 + 1], - checkpoint_pack[b * 2 + col_half][c * 4 + 2], - checkpoint_pack[b * 2 + col_half][c * 4 + 3], - ], - nvvm.MMALayout.COL, - ) + checkpoint_stage_base = checkpoint_stage * cfg.d_k * cfg.d_v + for sub in cutlass.range_constexpr(num_state_subs): + checkpoint_vec = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + num=32, + ) + for g in cutlass.range_constexpr(ldtm_width // 8): + packs = tuple( + fp32_to_fp16(checkpoint_vec[g * 8 + 2 * t], checkpoint_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dk = sub * ldtm_width + g * 8 + checkpoint_addr = checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + cg1_tidx * 64 + swizzle_xor_128b(cg1_tidx, dk % 64) + (sCheckpoint_raw.data_ptr() + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) nvvm.fence_proxy("async.shared", space="cta") if elect_one: bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() @@ -1645,7 +1665,7 @@ def compute1_warp_group( val = state_vec[k] if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): val = val.to(cfg.state_dtype) - gState_out[sub * ldtm_width + k, cg1_tidx] = val + gState_out[cg1_tidx, sub * ldtm_width + k] = val bars.mb_state_acc_scale_done[kv_last_idx].arrive() else: bars.mb_state_acc_scale_done[kv_last_idx].arrive() @@ -1657,11 +1677,11 @@ def compute1_warp_group( if cutlass.const_expr(cfg.use_initial_state): gState_in = mState_init[None, None, head_idx, batch_idx] for r in cutlass.range(num_state_subs * ldtm_width): - gState_out[r, cg1_tidx] = gState_in[r, cg1_tidx] + gState_out[cg1_tidx, r] = gState_in[cg1_tidx, r] else: for sub in cutlass.range_constexpr(num_state_subs): for k in cutlass.range_constexpr(32): - gState_out[sub * ldtm_width + k, cg1_tidx] = cutlass.Float32(0.0).to(cfg.state_dtype) + gState_out[cg1_tidx, sub * ldtm_width + k] = cutlass.Float32(0.0).to(cfg.state_dtype) tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) diff --git a/python/cudnn/linear_attention/frost/kernel/kda_bprop_f16.py b/python/cudnn/linear_attention/frost/kernel/kda_bprop_f16.py index bdc9d6a3c..ea64a186c 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_bprop_f16.py @@ -32,16 +32,16 @@ dGate = suffix-sum(dGate + dGate_last at row BT-1) (WG2 in-register reverse cumsum) dH <- diag-GEMM(eGl).dH + (scale.Q_decay)^T@dO - K_decay^T@(Beta.dY) -ABI: state_checkpoints `[total_checkpoints, HO, DK, DV]` (KV, v contiguous) io -dtype, the plain per-chunk checkpoint series with NO initial-state slot (entry `c -- 1` = state entering chunk c >= 1; chunk 0 seeds from `initial_state`); dq/dk/dv io at HO heads; dgate `[T, HO, DK]` fp32 +ABI: state_checkpoints `[total_checkpoints, HO, DV, DK]` (VK, k contiguous) io +dtype, the plain per-chunk checkpoint series (entry `c` = state entering chunk +c, so row 0 is the initial state or zeros); dq/dk/dv io at HO heads; dgate `[T, HO, DK]` fp32 (natural-log gate domain; with SAFE_GATE the gradient stays wrt the transformed log-decay); dbeta `[T, HO]` fp32, io dtype with BETA_SIGMOID (post-sigmoid space, or wrt the raw logits under BETA_SIGMOID). Gate arrives natural-log fp32 unless SAFE_GATE (safe-gate transform from raw gate + a_log/dt_bias); beta arrives post-sigmoid fp32, or io-dtype logits with -BETA_SIGMOID; d_initial_state / d_final_state fp32 `[N, HO, DK, DV]` -(K-major). +BETA_SIGMOID; d_initial_state / d_final_state fp32 `[N, HO, DV, DK]` +(V-major). Warp assignments (16 warps = 512 threads): warps 0-3 : WG0 - Gate prefix scan + decay/restore operands (all chunks) + Beta scalar gather @@ -68,7 +68,7 @@ from ..common.host import get_dtype from cudnn.frost.buffers import data_ptr from cudnn.frost.device import current_device, multiprocessor_count -from ..common.thd import TENSOR_MAP_QWORDS, emit_copy_desc, emit_checkpoint_seq_descs, emit_seq_descs +from ..common.thd import TENSOR_MAP_QWORDS, emit_checkpoint_seq_descs, emit_seq_descs from .kda_bprop_config import CFG from cudnn.frost.tile_dsl.barrier import ( @@ -672,7 +672,6 @@ def super_mma_warp( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) num_compute_chunks = cend - wstart @@ -881,6 +880,7 @@ def tcgen05_mma_warp( sSched, tmem_base_holder, sState_alt, + sState_direct, sK_decay_lead16, sK_inv_lead16, sK_inv_amaj, @@ -935,6 +935,26 @@ def tcgen05_mma_warp( idesc=idesc_state_k_at, kind=nvvm.Tcgen05MMAKind.F16, ) + idesc_state_k_kmaj = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.b_t, + m_dim=cfg.d_v, + ) + bmm_state_k_kmaj_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.d_k, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + atranspose=False, + cta_group=1, + idesc=idesc_state_k_kmaj, + kind=nvvm.Tcgen05MMAKind.F16, + ) bmm_dvinter_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, @@ -1194,13 +1214,13 @@ def tcgen05_mma_warp( d_u_lead0 = sU_lead16[0].desc() assert cfg.smem_state_stages == 1 d_state_alt0 = sState_alt[0].desc() + d_state_direct0 = sState_direct[0].desc() dstate0_index = PipelineState.start(phase=0) chunk_serial_base = cutlass.Int32(0) sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) num_compute_chunks = cend - wstart @@ -1233,8 +1253,8 @@ def tcgen05_mma_warp( if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_state_ready[state_index.idx].wait(state_index.phase) mma_ss( - bmm_state_k_desc, - d_state_alt0, + bmm_state_k_kmaj_desc, + d_state_direct0, d_kd_lead0 + decay_op_off, nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_k_acc_offset), cutlass.Float32), accumulate=False, @@ -1520,7 +1540,6 @@ def tmaldg_warp( desc_gate_base, desc_do_base, desc_checkpoint_base, - desc_initial_state_base, bars, ) -> None: """TMA-LDG warp role (warp 14): persistent tile-scheduler loop issuing @@ -1600,7 +1619,6 @@ def tmaldg_warp( sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) next_tile, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) @@ -1615,7 +1633,6 @@ def tmaldg_warp( desc_gate_slot = (desc_gate_base + slot).tospace(cutlass.AddressSpace.generic) desc_do_slot = (desc_do_base + slot).tospace(cutlass.AddressSpace.generic) desc_checkpoint_slot = (desc_checkpoint_base + slot).tospace(cutlass.AddressSpace.generic) - desc_initial_state_slot = (desc_initial_state_base + cutlass.Int32(0)).tospace(cutlass.AddressSpace.generic) if elect_one: tma_tensormap_acquire(desc_q_slot) tma_tensormap_acquire(desc_k_slot) @@ -1623,8 +1640,6 @@ def tmaldg_warp( tma_tensormap_acquire(desc_gate_slot) tma_tensormap_acquire(desc_do_slot) tma_tensormap_acquire(desc_checkpoint_slot) - if cutlass.const_expr(cfg.use_initial_state): - tma_tensormap_acquire(desc_initial_state_slot) num_compute_chunks = cend - wstart for rev_idx in cutlass.range(num_compute_chunks, unroll=1): chunk_idx = cend - cutlass.Int32(1) - rev_idx @@ -1665,7 +1680,7 @@ def tmaldg_warp( v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, chunk_start) tma_load_tile(sV_tma[raw_index.idx], v_slice, bars.mb_v_ready[raw_index.idx].smem_ptr, acquire=False) - # ---- entering state: state_checkpoints[c - 1] (sequence-local), or initial_state for chunk 0 ---- + # ---- entering state ---- if chunk_idx >= FIRST_STATE_CHUNK: state_idx = state_index.idx bars.mb_state_cg0_done[state_idx].wait(state_index.phase) @@ -1673,16 +1688,8 @@ def tmaldg_warp( state_index = advance(state_index, cfg.smem_state_stages) if elect_one: bars.mb_state_ready[state_idx].arrive(n_bytes=cfg.tma_state_bytes) - if cutlass.const_expr(cfg.use_initial_state): - if chunk_idx == 0: - initial_state_slice = tma_slice_runtime_desc(desc_initial_state_slot, cutlass.Int32(0), cutlass.Int32(0), head_o, batch_idx) - tma_load_tile(sState_tma[state_idx], initial_state_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) - else: - state_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx - cutlass.Int32(1), head_o) - tma_load_tile(sState_tma[state_idx], state_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) - else: - state_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx - FIRST_STATE_CHUNK, head_o) - tma_load_tile(sState_tma[state_idx], state_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) + state_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx, head_o) + tma_load_tile(sState_tma[state_idx], state_slice, bars.mb_state_ready[state_idx].smem_ptr, acquire=False) raw_index = advance(raw_index, cfg.smem_raw_stages) tile_idx = next_tile @@ -1750,7 +1757,6 @@ def compute0_warp_group( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) num_compute_chunks = cend - wstart @@ -2072,8 +2078,17 @@ def compute0_warp_group( state_src = sState_raw.data_ptr() + state_index.idx * (cfg.d_k * cfg.d_v) for v_seg in cutlass.range_constexpr(2): for v_col8 in cutlass.range_constexpr(8): - state_frag = (state_src + v_seg * (cfg.d_k * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, v_col8 * 8, elem_bytes=2)).load( - count=8, alignment=16 + state_frag = cutlass.Vector.from_elements( + tuple( + ( + state_src + + (value_dim // 64) * (cfg.d_v * 64) + + (v_seg * 64 + v_col8 * 8 + e) * 64 + + swizzle_xor_128b(v_seg * 64 + v_col8 * 8 + e, value_dim % 64, elem_bytes=2) + ).load() + for e in range(8) + ), + cfg.io_dtype, ) nvvm.tcgen05_st( "32x32b", @@ -2143,7 +2158,6 @@ def compute1_warp_group( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) num_compute_chunks = cend - wstart @@ -2156,12 +2170,14 @@ def compute1_warp_group( bars.mb_dstate_smem_cg2_done.wait(dstate_smem_done_index.phase) dstate_smem_done_index = advance(dstate_smem_done_index, 1) row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + dstate_src = (mDstate_in.iterator + mDstate_in.layout((batch_idx, head_idx, value_dim, 0))).raw_ptr() for sub in cutlass.range_constexpr(cfg.d_k // 16): seed_block = cutlass.Array(cutlass.Float32, 16, alignment=16) - for kk_i in cutlass.range_constexpr(16): - dval = mDstate_in[batch_idx, head_idx, sub * 16 + kk_i, value_dim].to(cutlass.Float32) - dval = dval if seed_true else cutlass.Float32(0.0) - seed_block[kk_i] = dval + for g in cutlass.range_constexpr(4): + seed_chunk = (dstate_src + sub * 16 + g * 4).load(count=4, alignment=16) + for t in cutlass.range_constexpr(4): + dval = seed_chunk[t].to(cutlass.Float32) + seed_block[g * 4 + t] = dval if seed_true else cutlass.Float32(0.0) nvvm.tcgen05_st( "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 16), cutlass.Float32), @@ -2501,20 +2517,24 @@ def compute1_warp_group( if num_compute_chunks > 0: if wstart == 0: row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + dstate0_dst = (mDstate0.iterator + mDstate0.layout((batch_idx, head_idx, value_dim, 0))).raw_ptr() for sub in cutlass.range_constexpr(cfg.d_k // 32): dstate0_vec = nvvm.tcgen05_ld( "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 32), cutlass.Float32), num=32 ) - for kk_i in cutlass.range_constexpr(32): - mDstate0[batch_idx, head_idx, sub * 32 + kk_i, value_dim] = dstate0_vec[kk_i] + for g in cutlass.range_constexpr(8): + (dstate0_dst + sub * 32 + g * 4).store( + cutlass.Vector.from_elements(tuple(dstate0_vec[g * 4 + t] for t in range(4)), cutlass.Float32), + alignment=16, + ) else: for key_dim_base in cutlass.range_constexpr(0, cfg.d_k, 32): for kk_i in cutlass.range_constexpr(32): kd = key_dim_base + kk_i if cutlass.const_expr(cfg.use_dstate_in): - mDstate0[batch_idx, head_idx, kd, value_dim] = mDstate_in[batch_idx, head_idx, kd, value_dim] + mDstate0[batch_idx, head_idx, value_dim, kd] = mDstate_in[batch_idx, head_idx, value_dim, kd] else: - mDstate0[batch_idx, head_idx, kd, value_dim] = cutlass.Float32(0.0) + mDstate0[batch_idx, head_idx, value_dim, kd] = cutlass.Float32(0.0) bars.mb_dstate0_acc_stored.arrive() chunk_serial_base += num_compute_chunks tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @@ -2571,7 +2591,6 @@ def compute2_warp_group( sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 while tile_idx < total_tiles: batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) num_compute_chunks = cend - wstart @@ -2815,7 +2834,6 @@ def build_descs_body( base_dv, base_dgate, base_checkpoint, - base_initial_state, desc_ws: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, @@ -2828,7 +2846,6 @@ def build_descs_body( dv: cute.Tensor, dgate: cute.Tensor, state_checkpoints: cute.Tensor, - state0: cute.Tensor | None, n_batch: cutlass.Int32, q_row_stride: cutlass.Int32, k_row_stride: cutlass.Int32, @@ -2856,7 +2873,6 @@ def build_descs_body( sub7 = cute.make_tensor(desc_ws.iterator + 7 * arr_words, cute.make_layout((arr_words,), stride=(1,))) sub8 = cute.make_tensor(desc_ws.iterator + 8 * arr_words, cute.make_layout((arr_words,), stride=(1,))) sub9 = cute.make_tensor(desc_ws.iterator + 9 * arr_words, cute.make_layout((arr_words,), stride=(1,))) - sub10 = cute.make_tensor(desc_ws.iterator + 10 * arr_words, cute.make_layout((cutlass.Int32(TENSOR_MAP_QWORDS),), stride=(1,))) if widx == 0: if nvvm.elect_sync(): @@ -2898,11 +2914,6 @@ def build_descs_body( if nvvm.elect_sync(): emit_checkpoint_seq_descs(base_checkpoint, sub9, cu_seqlens, state_checkpoints, n_batch, checkpoint_row_stride, checkpoint_every_n, 2) nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) - if cutlass.const_expr(state0 is not None): - if widx == 10: - if nvvm.elect_sync(): - emit_copy_desc(base_initial_state, sub10) - nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) @cute.kernel @@ -2921,7 +2932,6 @@ def prologue_kernel( base_dv: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_dgate: cutlass.GridConstant[cuda.tensor_map.TensorMap], base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], - base_initial_state: cutlass.GridConstant[cuda.tensor_map.TensorMap], desc_ws: cute.Tensor, cu_seqlens: cute.Tensor, q: cute.Tensor, @@ -2934,7 +2944,6 @@ def prologue_kernel( dv: cute.Tensor, dgate: cute.Tensor, state_checkpoints: cute.Tensor, - state0: cute.Tensor | None, mStaging: cute.Tensor | None, mCount: cute.Tensor, mWorkItems: cute.Tensor | None, @@ -2995,7 +3004,6 @@ def prologue_kernel( base_dv, base_dgate, base_checkpoint, - base_initial_state, desc_ws, cu_seqlens, q, @@ -3008,7 +3016,6 @@ def prologue_kernel( dv, dgate, state_checkpoints, - state0, n_batch, q_row_stride, k_row_stride, @@ -3041,7 +3048,6 @@ def prologue( dv: cute.Tensor, dgate: cute.Tensor, state_checkpoints: cute.Tensor, - state0: cute.Tensor | None, cu_seqlens: cute.Tensor, work_item_staging: cute.Tensor | None, work_count: cute.Tensor, @@ -3094,16 +3100,6 @@ def prologue( ), ) base_checkpoint = cuda.create_tensor_map_tiled_from_view(checkpoint_view, box_dims=(64, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) - base_initial_state = base_checkpoint - if cutlass.const_expr(state0 is not None): - initial_state_view = cute.make_tensor( - state0.iterator, - cute.make_layout( - (d_v, d_k, ho, batch_size), - stride=(state0.stride[3], state0.stride[2], state0.stride[1], state0.stride[0]), - ), - ) - base_initial_state = cuda.create_tensor_map_tiled_from_view(initial_state_view, box_dims=(64, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) prologue_kernel( run_order, @@ -3120,7 +3116,6 @@ def prologue( base_dv, base_dgate, base_checkpoint, - base_initial_state, tensormap_workspace, cu_seqlens, q, @@ -3133,7 +3128,6 @@ def prologue( dv, dgate, state_checkpoints, - state0, work_item_staging, work_count, work_items, @@ -3160,7 +3154,6 @@ def host( dt_bias: cute.Tensor | None, beta: cute.Tensor, state_checkpoints: cute.Tensor, - mState_init: cute.Tensor | None, dgate: cute.Tensor, dbeta: cute.Tensor, cu_seqlens: cute.Tensor, @@ -3245,7 +3238,6 @@ def kernel( desc_dv_base = desc_base_words + cutlass.Int32(7) * arr_words desc_dgate_base = desc_base_words + cutlass.Int32(8) * arr_words desc_checkpoint_base = desc_base_words + cutlass.Int32(9) * arr_words - desc_initial_state_base = desc_base_words + cutlass.Int32(10) * arr_words SMEM = cutlass.AddressSpace.smem bars = make_kda_bwd_bars(cfg) @@ -3291,6 +3283,14 @@ def kernel( stride_byte_offset=STRIDE, layout=SWZ, ) + sState_direct = SmemTile( + base=sState_raw.data_ptr().toint(), + elems_per_stage=((cfg.state_cosize) // (cfg.smem_state_stages)) * bpe, + stages=cfg.smem_state_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) sK_decay_lead16 = SmemTile( base=sK_decay_raw.data_ptr().toint(), elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, @@ -3509,7 +3509,6 @@ def kernel( desc_gate_base, desc_do_base, desc_checkpoint_base, - desc_initial_state_base, bars, ) elif warp_idx == cfg.super_mma_warp_id: @@ -3542,6 +3541,7 @@ def kernel( sSched, tmem_base_holder, sState_alt, + sState_direct, sK_decay_lead16, sK_inv_lead16, sK_inv_amaj, @@ -3844,7 +3844,7 @@ def build_cfg( TENSORMAP_DESC_ARRAYS = 10 # per-batch runtime TMA descriptors: Q, K, V, Gate, dO, state_checkpoints, dQ, dK, dV, dGate -TENSORMAP_STATIC_SLOTS = 1 # initial_state +TENSORMAP_STATIC_SLOTS = 0 # ---- Torch adapter / host-side compilation --------------------------------------- @@ -3887,7 +3887,7 @@ def chunk_kda_bwd_sm100( cu_seqlens, scale: float, *, - initial_state=None, + use_initial_state: bool = False, d_initial_state=None, d_final_state=None, use_qk_l2norm_in_kernel: bool = False, @@ -3919,9 +3919,10 @@ def chunk_kda_bwd_sm100( beta: ``(total_tokens, HO)``. Post-sigmoid float32, or io-dtype logits when ``use_beta_sigmoid`` do: ``(total_tokens, HO, DV)`` io dtype - state_checkpoints: ``(total_checkpoints, HO, DK, DV)`` io dtype (KV, v contiguous), the PLAIN per-chunk checkpoint series with no initial-state - slot: sequence-local entry ``c - 1`` is the state ENTERING chunk c >= 1 - of sequence b; chunk 0 seeds from ``initial_state`` + state_checkpoints: ``(total_checkpoints, HO, DV, DK)`` io dtype (VK, k + contiguous), the PLAIN per-chunk checkpoint series: sequence-local + entry ``c`` is the state ENTERING chunk c of sequence b, so row 0 + is the initial state (or zeros when the forward had none) dq/dk/dv: io dtype at ``HO = max(HQ, HV)`` heads, pre-allocated dgate: ``(total_tokens, HO, DK)`` fp32 (dL/d ln alpha), pre-allocated. With ``safe_gate`` this stays the gradient wrt the TRANSFORMED @@ -3932,10 +3933,10 @@ def chunk_kda_bwd_sm100( folds the sigmoid derivative into its own dbeta write) cu_seqlens: ``(num_seqs + 1,)`` int32 scale: attention scale factor - initial_state: ``(num_seqs, HO, DK, DV)`` io dtype (KV) -- the state - entering chunk 0 (engine-provided zeros when the graph has none) - d_initial_state: fp32 ``(num_seqs, HO, DK, DV)`` OUT (dL/d initial state), or None - d_final_state: fp32 ``(num_seqs, HO, DK, DV)`` IN (dL/d final state) + use_initial_state: the forward ran with an initial state, so chunk 0 + has an entering state to load from ``state_checkpoints`` row 0 + d_initial_state: fp32 ``(num_seqs, HO, DV, DK)`` OUT (dL/d initial state), or None + d_final_state: fp32 ``(num_seqs, HO, DV, DK)`` IN (dL/d final state) use_qk_l2norm_in_kernel: q/k arrive raw; the kernel normalizes for the recompute math and chains the L2-norm backward into dq/dk safe_gate: interpret ``gate`` through the safe-gate transform @@ -3958,7 +3959,6 @@ def chunk_kda_bwd_sm100( HO = max(HQ, HV) use_dstate_in = d_final_state is not None use_dstate0 = d_initial_state is not None - use_initial_state = initial_state is not None if work_items is None or work_count is None: raise ValueError("work_items/work_count are required (the split-table stage builds them for every launch)") dyn_sched = sched_ctr is not None @@ -3966,9 +3966,8 @@ def chunk_kda_bwd_sm100( order_gen = order_in_prologue and work_item_scratch is None if run_order and sched_all is None: raise ValueError("order in the prologue requires sched_all (the prologue zeroes both consumers' sched rings)") - for name, t in (("state_checkpoints", state_checkpoints),) + ((("initial_state", initial_state),) if use_initial_state else ()): - if str(t.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: - raise ValueError(f"{name} dtype must match the io dtype: got {t.dtype} with io {q.dtype}") + if str(state_checkpoints.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: + raise ValueError(f"state_checkpoints dtype must match the io dtype: got {state_checkpoints.dtype} with io {q.dtype}") for name, hh in (("HQ", HQ), ("HK", HK), ("HV", HV)): if HO % hh != 0: raise ValueError(f"{name}={hh} must divide {HO}") @@ -4038,9 +4037,6 @@ def chunk_kda_bwd_sm100( dt_bias_cute = from_dlpack(dt_bias, assumed_align=16) if dt_bias is not None else None beta_cute = from_dlpack(beta, assumed_align=4).mark_layout_dynamic(leading_dim=len(beta.shape) - 1) state_checkpoints_cute = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=len(state_checkpoints.shape) - 1) - initial_state_cute = ( - from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=len(initial_state.shape) - 1) if use_initial_state else None - ) dgate_cute = from_dlpack(dgate, assumed_align=16).mark_layout_dynamic(leading_dim=len(dgate.shape) - 1) dbeta_cute = from_dlpack(dbeta, assumed_align=4).mark_layout_dynamic(leading_dim=len(dbeta.shape) - 1) cache["compiled"] = cute.compile( @@ -4050,7 +4046,6 @@ def chunk_kda_bwd_sm100( dt_bias_cute, beta_cute, state_checkpoints_cute, - initial_state_cute, dgate_cute, dbeta_cute, from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic(), @@ -4077,7 +4072,6 @@ def chunk_kda_bwd_sm100( dv_pl = from_dlpack(dv, assumed_align=16).mark_layout_dynamic(leading_dim=2) dgate_pl = from_dlpack(dgate, assumed_align=16).mark_layout_dynamic(leading_dim=2) state_checkpoints_pl = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=3) - initial_state_pl = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) if use_initial_state else None cu_pl = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() ws_pl = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() staging_pl = None @@ -4107,7 +4101,6 @@ def chunk_kda_bwd_sm100( dv_pl, dgate_pl, state_checkpoints_pl, - initial_state_pl, cu_pl, staging_pl, work_count_pl, @@ -4128,7 +4121,6 @@ def chunk_kda_bwd_sm100( dv, dgate, state_checkpoints, - initial_state, cu_seqlens, work_item_scratch if run_order else None, work_count, @@ -4142,7 +4134,6 @@ def chunk_kda_bwd_sm100( dt_bias, beta, state_checkpoints, - initial_state, dgate, dbeta, cu_seqlens, @@ -4173,7 +4164,6 @@ def run_bwd( dgate, dbeta, cu_seqlens, - initial_state, d_initial_state, d_final_state, work_items, @@ -4202,7 +4192,6 @@ def run_bwd( dv, dgate, state_checkpoints, - initial_state, cu_seqlens, work_item_scratch, work_count, @@ -4216,7 +4205,6 @@ def run_bwd( dt_bias, beta, state_checkpoints, - initial_state, dgate, dbeta, cu_seqlens, diff --git a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py index 433538856..3eb6c3724 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py @@ -39,7 +39,7 @@ ABI: q `[T, HQ, DK]`, k `[T, HK, DK]`, v `[T, HV, DV]`, gate `[T, HO, DK]` fp32 (natural-log decay unless SAFE_GATE, which applies the safe-gate transform from raw gate + a_log/dt_bias), beta `[T, HO]` fp32 -post-sigmoid, cu_seqlens int32, states/checkpoints `[N, HO, DK, DV]` (KV, v +post-sigmoid, cu_seqlens int32, states/checkpoints `[N, HO, DV, DK]` (VK, k contiguous). GQA/GVA head broadcast follows repeat_interleave: source head = head_idx // (HO // H_x). State presence, L2NORM, SAFE_GATE, checkpoints, and the @@ -933,6 +933,16 @@ def epilogue_warp( if elect_one: tma_tensormap_acquire(desc_o_slot) num_chunks_tile = wend - cstart + if cutlass.const_expr(cfg.enable_checkpoints): + if num_chunks_tile > 0 and wstart == 0: + checkpoint_stage = checkpoint_ready_index.idx + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) + checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) + checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0), head_o) + tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) + tma_store_commit() + tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): chunk_idx = cstart + local_chunk_idx cum_chunk = cum_chunk_base + local_chunk_idx @@ -1025,7 +1035,7 @@ def epilogue_warp( if do_checkpoint: bars.mb_checkpoint_tmastg_ready[checkpoint_ready_index.idx].wait(checkpoint_ready_index.phase) checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) - checkpoint_entry = checkpoint_quot - cutlass.Int32(1) + checkpoint_entry = checkpoint_quot checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), checkpoint_entry, head_o) tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) tma_store_commit() @@ -1527,11 +1537,14 @@ def compute1_warp_group( seed_from_initial_state = cstart == 0 if cutlass.const_expr(mState_init is not None): if seed_from_initial_state: + seed_vw = 16 // (mState_init.element_type.width // 8) + seed_src = (mState_init.iterator + mState_init.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for col in cutlass.range_constexpr(32): - key_dim = key_block_start + col - state_block[col] = mState_init[batch_idx, head_o, key_dim, value_dim].to(cutlass.Float32) + for g in cutlass.range_constexpr(32 // seed_vw): + seed_chunk = (seed_src + key_block_start + g * seed_vw).load(count=seed_vw, alignment=16) + for t in cutlass.range_constexpr(seed_vw): + state_block[g * seed_vw + t] = seed_chunk[t].to(cutlass.Float32) nvvm.tcgen05_st( "32x32b", @@ -1573,8 +1586,47 @@ def compute1_warp_group( nvvm.tcgen05_wait("store") bars.mb_state_inp_ready.arrive() if cutlass.const_expr(cfg.enable_checkpoints): + if wstart == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + for sub in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + packs = tuple( + fp32_to_fp16(state_vecs[sub][g * 8 + 2 * t], state_vecs[sub][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dk = sub * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.tcgen05_wait("load") + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() bars.mb_state_acc_read_done.arrive() + if cutlass.const_expr(cfg.enable_checkpoints and mState_init is None): + if wstart == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + zero_packs = tuple(cutlass.Int32(0) for _ in range(4)) + for sub in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + dk = sub * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(zero_packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + # ---- Y staging: Y = Beta * (V - state*K) ----------------------------- bars.mb_v_ready[raw_bar_index.idx].wait(raw_bar_index.phase) raw_v_frag0 = nvvm.ldmatrix( @@ -1728,35 +1780,16 @@ def compute1_warp_group( bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) - for slab in cutlass.range_constexpr(cfg.d_k // 16): - checkpoint_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + state_col_id + slab * 16, cutlass.Float32), num=2) - checkpoint_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + state_col_id + slab * 16, cutlass.Float32), num=2) - checkpoint_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - checkpoint_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - for reg_idx in cutlass.range_constexpr(4): - checkpoint_pack0[reg_idx] = fp32_to_fp16(checkpoint_vec0[2 * reg_idx], checkpoint_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) - checkpoint_pack1[reg_idx] = fp32_to_fp16(checkpoint_vec1[2 * reg_idx], checkpoint_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) - checkpoint_row = slab * 16 + ov_token_coord - nvvm.stmatrix( - sCheckpoint_ptr - + checkpoint_stage_base - + checkpoint_swz_off0 - + checkpoint_row * 64 - + swizzle_xor_128b(checkpoint_row, checkpoint_swz_col0, elem_bytes=2), - checkpoint_pack0.data_ptr().load(count=4, alignment=4), - nvvm.MMALayout.COL, - shape=nvvm.StoreShape.M8N8, - ) - nvvm.stmatrix( - sCheckpoint_ptr - + checkpoint_stage_base - + checkpoint_swz_off - + checkpoint_row * 64 - + swizzle_xor_128b(checkpoint_row, checkpoint_swz_col, elem_bytes=2), - checkpoint_pack1.data_ptr().load(count=4, alignment=4), - nvvm.MMALayout.COL, - shape=nvvm.StoreShape.M8N8, - ) + for sub in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + packs = tuple( + fp32_to_fp16(state_vecs[sub][g * 8 + 2 * t], state_vecs[sub][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dk = sub * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store(cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16) nvvm.tcgen05_wait("load") bars.mb_state_acc_read_done.arrive() nvvm.fence_proxy("async.shared", space="cta") @@ -1940,6 +1973,8 @@ def compute1_warp_group( if cutlass.const_expr(mState_out is not None): if seqlen_b > 0: if owns_final: + state_vw = 16 // (mState_out.element_type.width // 8) + state_dst = (mState_out.iterator + mState_out.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): loaded = nvvm.tcgen05_ld( "32x32b", @@ -1947,17 +1982,22 @@ def compute1_warp_group( num=32, ) - for col in cutlass.range_constexpr(32): - key_dim = key_block_start + col - mState_out[batch_idx, head_o, key_dim, value_dim] = loaded[col].to(mState_out.element_type) + for g in cutlass.range_constexpr(32 // state_vw): + (state_dst + key_block_start + g * state_vw).store( + cutlass.Vector.from_elements( + tuple(loaded[g * state_vw + t].to(mState_out.element_type) for t in range(state_vw)), + mState_out.element_type, + ), + alignment=16, + ) else: for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): for col in cutlass.range_constexpr(32): key_dim = key_block_start + col if cutlass.const_expr(mState_init is not None): - mState_out[batch_idx, head_o, key_dim, value_dim] = mState_init[batch_idx, head_o, key_dim, value_dim] + mState_out[batch_idx, head_o, value_dim, key_dim] = mState_init[batch_idx, head_o, value_dim, key_dim] else: - mState_out[batch_idx, head_o, key_dim, value_dim] = cutlass.Float32(0.0).to(mState_out.element_type) + mState_out[batch_idx, head_o, value_dim, key_dim] = cutlass.Float32(0.0).to(mState_out.element_type) cum_chunk_base += num_chunks_tile tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @@ -2892,15 +2932,15 @@ def chunk_kda_sm100( logits when ``use_beta_sigmoid_in_kernel`` output: ``(total_tokens, HO, DV)`` float16/bfloat16, pre-allocated cu_seqlens: ``(num_seqs + 1,)`` int32 - initial_state: ``(num_seqs, HO, DK, DV)`` float32/bfloat16, or None - output_state: ``(num_seqs, HO, DK, DV)`` float32/bfloat16, or None + initial_state: ``(num_seqs, HO, DV, DK)`` float32/bfloat16, or None + output_state: ``(num_seqs, HO, DV, DK)`` float32/bfloat16, or None scale: attention scale factor (must not be 0) checkpoint_every_n_tokens: emit a state checkpoint every N tokens (0 = off). state_checkpoints[j] is the state after ``(j + 1) * N`` tokens, STRICTLY BEFORE the sequence end — the end-of-sequence state is only ``output_state``. With ``N == B_T`` this is the per-chunk checkpoint series the backward pass consumes. - output_state_checkpoints: ``(total_checkpoints, HO, DK, DV)`` io-dtype (KV, V + output_state_checkpoints: ``(total_checkpoints, HO, DV, DK)`` io-dtype (VK, K contiguous); the per-sequence entry offsets are derived on device from ``cu_seqlens`` ((seqlen-1)//N, prefix-summed), so there is no cu_checkpoints array diff --git a/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py b/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py index 26dbe8405..404707b2d 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py @@ -39,7 +39,7 @@ ABI: k `[T, HK, DK]`, v `[T, HV, DV]`, gate `[T, HO, DK]` fp32 (natural-log decay unless SAFE_GATE, which applies the safe-gate transform from raw gate + a_log/dt_bias), beta `[T, HO]` fp32 -post-sigmoid, cu_seqlens int32, states/checkpoints `[N, HO, DK, DV]` (KV, v +post-sigmoid, cu_seqlens int32, states/checkpoints `[N, HO, DV, DK]` (VK, k contiguous). GQA/GVA head broadcast follows repeat_interleave: source head = head_idx // (HO // H_x). State presence, L2NORM, SAFE_GATE, checkpoints, and the @@ -814,6 +814,15 @@ def epilogue_warp( if elect_one: tma_tensormap_acquire(desc_checkpoint_slot) num_chunks_tile = wend - cstart + if num_chunks_tile > 0 and wstart == 0: + checkpoint_stage = checkpoint_ready_index.idx + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) + checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) + checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0), head_o) + tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) + tma_store_commit() + tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): chunk_idx = cstart + local_chunk_idx if local_chunk_idx > 0: @@ -824,7 +833,7 @@ def epilogue_warp( checkpoint_stage = checkpoint_ready_index.idx bars.mb_checkpoint_tmastg_ready[checkpoint_stage].wait(checkpoint_ready_index.phase) checkpoint_ready_index = advance(checkpoint_ready_index, cfg.smem_checkpoint_stages) - checkpoint_entry = checkpoint_quot - cutlass.Int32(1) + checkpoint_entry = checkpoint_quot checkpoint_slice = tma_slice_runtime_desc(desc_checkpoint_slot, cutlass.Int32(0), cutlass.Int32(0), checkpoint_entry, head_o) tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) tma_store_commit() @@ -1231,11 +1240,14 @@ def compute1_warp_group( seed_from_initial_state = cstart == 0 if cutlass.const_expr(mState_init is not None): if seed_from_initial_state: + seed_vw = 16 // (mState_init.element_type.width // 8) + seed_src = (mState_init.iterator + mState_init.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): state_block = cutlass.Array(cutlass.Float32, 32, alignment=16) - for col in cutlass.range_constexpr(32): - key_dim = key_block_start + col - state_block[col] = mState_init[batch_idx, head_o, key_dim, value_dim].to(cutlass.Float32) + for g in cutlass.range_constexpr(32 // seed_vw): + seed_chunk = (seed_src + key_block_start + g * seed_vw).load(count=seed_vw, alignment=16) + for t in cutlass.range_constexpr(seed_vw): + state_block[g * seed_vw + t] = seed_chunk[t].to(cutlass.Float32) nvvm.tcgen05_st( "32x32b", @@ -1278,8 +1290,47 @@ def compute1_warp_group( nvvm.tcgen05_wait("store") bars.mb_state_inp_ready.arrive() if cutlass.const_expr(cfg.enable_checkpoints): + if wstart == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + for sub in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + packs = tuple( + fp32_to_fp16(state_vecs[sub][g * 8 + 2 * t], state_vecs[sub][g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4) + ) + dk = sub * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.tcgen05_wait("load") + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() bars.mb_state_acc_read_done.arrive() + if cutlass.const_expr(cfg.enable_checkpoints and mState_init is None): + if wstart == 0: + checkpoint_stage = checkpoint_done_index.idx + bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) + checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) + checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) + zero_packs = tuple(cutlass.Int32(0) for _ in range(4)) + for sub in cutlass.range_constexpr(cfg.d_k // 16): + for g in cutlass.range_constexpr(2): + dk = sub * 16 + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store( + cutlass.Vector.from_elements(zero_packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16 + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + # ---- Y staging: Y = Beta * (V - State*K) ----------------------------- bars.mb_v_ready[raw_index.idx].wait(raw_index.phase) raw_v_frag0 = nvvm.ldmatrix( @@ -1426,35 +1477,15 @@ def compute1_warp_group( bars.mb_checkpoint_tmastg_done[checkpoint_stage].wait(checkpoint_done_index.phase) checkpoint_done_index = advance(checkpoint_done_index, cfg.smem_checkpoint_stages) checkpoint_stage_base = checkpoint_stage * (cfg.d_k * cfg.d_v) - for slab in cutlass.range_constexpr(cfg.d_k // 16): - checkpoint_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + state_col_id + slab * 16, cutlass.Float32), num=2) - checkpoint_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + state_col_id + slab * 16, cutlass.Float32), num=2) - checkpoint_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - checkpoint_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - for reg_idx in cutlass.range_constexpr(4): - checkpoint_pack0[reg_idx] = fp32_to_fp16(checkpoint_vec0[2 * reg_idx], checkpoint_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) - checkpoint_pack1[reg_idx] = fp32_to_fp16(checkpoint_vec1[2 * reg_idx], checkpoint_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) - checkpoint_row = slab * 16 + frag_row_coord - nvvm.stmatrix( - sCheckpoint_ptr - + checkpoint_stage_base - + checkpoint_swz_off0 - + checkpoint_row * 64 - + swizzle_xor_128b(checkpoint_row, checkpoint_swz_col0, elem_bytes=2), - checkpoint_pack0.data_ptr().load(count=4, alignment=4), - nvvm.MMALayout.COL, - shape=nvvm.StoreShape.M8N8, - ) - nvvm.stmatrix( - sCheckpoint_ptr - + checkpoint_stage_base - + checkpoint_swz_off - + checkpoint_row * 64 - + swizzle_xor_128b(checkpoint_row, checkpoint_swz_col, elem_bytes=2), - checkpoint_pack1.data_ptr().load(count=4, alignment=4), - nvvm.MMALayout.COL, - shape=nvvm.StoreShape.M8N8, - ) + for k_base in cutlass.range_constexpr(0, cfg.d_k, 32): + checkpoint_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_base, cutlass.Float32), num=32) + for g in cutlass.range_constexpr(4): + packs = tuple(fp32_to_fp16(checkpoint_vec[g * 8 + 2 * t], checkpoint_vec[g * 8 + 2 * t + 1], dtype=cfg.io_dtype) for t in range(4)) + dk = k_base + g * 8 + checkpoint_addr = ( + checkpoint_stage_base + (dk // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, dk % 64, elem_bytes=2) + ) + (sCheckpoint_ptr + checkpoint_addr).store(cutlass.Vector.from_elements(packs, cutlass.Int32).bitcast(cfg.io_dtype), alignment=16) nvvm.fence_proxy("async.shared", space="cta") bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() nvvm.tcgen05_wait("load") @@ -1569,6 +1600,8 @@ def compute1_warp_group( if cutlass.const_expr(mState_out is not None): if seqlen_b > 0: if owns_final: + state_vw = 16 // (mState_out.element_type.width // 8) + state_dst = (mState_out.iterator + mState_out.layout((batch_idx, head_o, value_dim, 0))).raw_ptr() for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): loaded = nvvm.tcgen05_ld( "32x32b", @@ -1576,17 +1609,22 @@ def compute1_warp_group( num=32, ) - for col in cutlass.range_constexpr(32): - key_dim = key_block_start + col - mState_out[batch_idx, head_o, key_dim, value_dim] = loaded[col].to(mState_out.element_type) + for g in cutlass.range_constexpr(32 // state_vw): + (state_dst + key_block_start + g * state_vw).store( + cutlass.Vector.from_elements( + tuple(loaded[g * state_vw + t].to(mState_out.element_type) for t in range(state_vw)), + mState_out.element_type, + ), + alignment=16, + ) else: for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): for col in cutlass.range_constexpr(32): key_dim = key_block_start + col if cutlass.const_expr(mState_init is not None): - mState_out[batch_idx, head_o, key_dim, value_dim] = mState_init[batch_idx, head_o, key_dim, value_dim] + mState_out[batch_idx, head_o, value_dim, key_dim] = mState_init[batch_idx, head_o, value_dim, key_dim] else: - mState_out[batch_idx, head_o, key_dim, value_dim] = cutlass.Float32(0.0).to(mState_out.element_type) + mState_out[batch_idx, head_o, value_dim, key_dim] = cutlass.Float32(0.0).to(mState_out.element_type) tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) bars.mb_tmem_done[0].arrive() @@ -2406,14 +2444,14 @@ def chunk_kda_recompute_sm100( beta: ``(total_tokens, HO)``. Post-sigmoid float32, or io-dtype logits when ``use_beta_sigmoid`` cu_seqlens: ``(num_seqs + 1,)`` int32 - initial_state: ``(num_seqs, HO, DK, DV)`` float32/bfloat16, or None - output_state: ``(num_seqs, HO, DK, DV)`` float32/bfloat16, or None + initial_state: ``(num_seqs, HO, DV, DK)`` float32/bfloat16, or None + output_state: ``(num_seqs, HO, DV, DK)`` float32/bfloat16, or None checkpoint_every_n_tokens: emit a state checkpoint every N tokens (0 = off). state_checkpoints[j] is the state after ``(j + 1) * N`` tokens, STRICTLY BEFORE the sequence end - the end-of-sequence state is only ``output_state``. With ``N == B_T`` this is the per-chunk checkpoint series the backward pass consumes. - output_state_checkpoints: ``(total_checkpoints, HO, DK, DV)`` io-dtype (KV, V + output_state_checkpoints: ``(total_checkpoints, HO, DV, DK)`` io-dtype (VK, K contiguous); the per-sequence entry offsets are derived on device from ``cu_seqlens`` ((seqlen-1)//N, prefix-summed), so there is no cu_checkpoints array diff --git a/python/cudnn/linear_attention/graph_analyzer.py b/python/cudnn/linear_attention/graph_analyzer.py index 2d6ee18f6..ef09ee556 100644 --- a/python/cudnn/linear_attention/graph_analyzer.py +++ b/python/cudnn/linear_attention/graph_analyzer.py @@ -65,6 +65,9 @@ class LaGraphFacts: h_o: int = 0 d_qk: int = 0 d_v: int = 0 + total_t: int = 0 # packed token count + n_seq: int = 0 + state_checkpoint_rows: int = 0 # rows declared on the checkpoint port (0 = absent or undeclared) gates_at_ho: bool = True # Gate/Beta(/W) carry HO = max(h_q, h_v) heads # dtypes (cudnn.data_type vocabulary; None = unset/inferred) @@ -175,6 +178,13 @@ def analyze(graph: "cudnn.pygraph") -> Optional[LaGraphFacts]: else: h_q = h_k = h_v = d_qk = d_v = 0 h_o = max(h_q, h_v) + total_t = int(q.dim[0]) if thd_layout else 0 + cu = ins["cu_seqlens"] + n_seq = int(cu.dim[0]) - 1 if cu.dim else 0 + ckpt_port = ins.get("state_checkpoints") + if ckpt_port is None: + ckpt_port = outs.get("state_checkpoints") + state_checkpoint_rows = int(ckpt_port.dim[0]) if ckpt_port is not None and ckpt_port.dim else 0 gates_at_ho = all(t is None or not t.dim or (len(t.dim) > 1 and int(t.dim[1]) == h_o) for t in (ins["g"], ins["beta"], ins.get("w"))) io_dtypes = {in_dt["q"], in_dt["k"], in_dt["v"]} - {None} state_dtypes = {in_dt.get("initial_state"), out_dt.get("final_state")} - {None} @@ -190,6 +200,9 @@ def analyze(graph: "cudnn.pygraph") -> Optional[LaGraphFacts]: h_o=h_o, d_qk=d_qk, d_v=d_v, + total_t=total_t, + n_seq=n_seq, + state_checkpoint_rows=state_checkpoint_rows, gates_at_ho=gates_at_ho, io_dtype=in_dt["q"], uniform_io=len(io_dtypes) <= 1, diff --git a/python/cudnn/linear_attention/ops/gdn.py b/python/cudnn/linear_attention/ops/gdn.py index b5b283f75..ffab63a8a 100644 --- a/python/cudnn/linear_attention/ops/gdn.py +++ b/python/cudnn/linear_attention/ops/gdn.py @@ -259,7 +259,7 @@ def _build_fprop_graph( cu_t = graph.tensor([N + 1], data_type=cu_dtype, name="cu_seqlens") state0_t = None if state_dtype is not None: - state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") + state0_t = graph.tensor([N, HO, V, K], data_type=state_dtype, name="initial_state") a_log_t = None dt_bias_t = None if safe_gate: @@ -443,12 +443,12 @@ def _gdn_fwd( variant_pack[t["dt_bias"]] = dt_bias final_state = torch.empty(0, dtype=torch.float32, device=device) if output_final_state: - final_state = torch.empty(N, HO, K, V, dtype=torch.float32, device=device) + final_state = torch.empty(N, HO, V, K, dtype=torch.float32, device=device) variant_pack[t["fs"]] = final_state state_checkpoints = torch.empty(0, dtype=q.dtype, device=device) if ckpt > 0: - total_checkpoints = max(total // ckpt, 1) - state_checkpoints = torch.empty(total_checkpoints, HO, K, V, dtype=q.dtype, device=device) + total_checkpoints = max(total // ckpt + N, 1) + state_checkpoints = torch.empty(total_checkpoints, HO, V, K, dtype=q.dtype, device=device) variant_pack[t["state_checkpoints"]] = state_checkpoints graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) return o, final_state, state_checkpoints @@ -486,10 +486,10 @@ def _gdn_fwd_fake( if initial_state is not None and initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") o = q.new_empty(total, HO, V) - final = q.new_empty((N, HO, K, V) if output_final_state else (0,), dtype=torch.float32) + final = q.new_empty((N, HO, V, K) if output_final_state else (0,), dtype=torch.float32) if checkpoint_every_n_tokens > 0: - total_checkpoints = max(total // int(checkpoint_every_n_tokens), 1) - state_checkpoints = q.new_empty(total_checkpoints, HO, K, V) + total_checkpoints = max(total // int(checkpoint_every_n_tokens) + N, 1) + state_checkpoints = q.new_empty(total_checkpoints, HO, V, K) else: state_checkpoints = q.new_empty(0) return o, final, state_checkpoints @@ -532,13 +532,13 @@ def _build_bprop_graph( dO_t = graph.tensor([total, HO, V], data_type=io_dtype, name="dO") state0_t = None if state_dtype is not None: - state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") + state0_t = graph.tensor([N, HO, V, K], data_type=state_dtype, name="initial_state") dfs_t = None if dstate_in_dtype is not None: - dfs_t = graph.tensor([N, HO, K, V], data_type=dstate_in_dtype, name="d_final_state") + dfs_t = graph.tensor([N, HO, V, K], data_type=dstate_in_dtype, name="d_final_state") ckpts_t = None if ckpt_rows is not None: - ckpts_t = graph.tensor([ckpt_rows, HO, K, V], data_type=io_dtype, name="state_checkpoints") + ckpts_t = graph.tensor([ckpt_rows, HO, V, K], data_type=io_dtype, name="state_checkpoints") a_log_t = None dt_bias_t = None if safe_gate: @@ -957,7 +957,7 @@ def gated_delta_net( HK = HV for canonical GQA: grouped K/V heads shared across query groups); v: ``[total_tokens, HV, V]`` g, beta: ``[total_tokens, HO]`` with ``HO = max(H, HV)``; cu_seqlens: ``[N+1]`` int32; O and the states live at HO heads - (initial_state / final_state: ``[N, HO, K, V]``) + (initial_state / final_state: ``[N, HO, V, K]``) A dense batch of N equal-length sequences is expressed as ``cu_seqlens = [0, T, 2T, ...]`` over the flattened tokens. @@ -989,7 +989,7 @@ def gated_delta_net( a_log: ``[HO]`` float32 safe-gate per-head log-amplitude. dt_bias: ``[HO]`` float32 safe-gate per-head bias. checkpoint_every_n_tokens: if ``> 0``, also return the per-chunk - recurrent state series ``state_checkpoints`` (``[total_checkpoints, HO, K, V]`` io dtype, + recurrent state series ``state_checkpoints`` (``[total_checkpoints, HO, V, K]`` io dtype, one entry per N tokens strictly before each sequence end; the FROST engine requires a positive multiple of the kernel chunk size, 64). The series is a non-differentiable dump. diff --git a/python/cudnn/linear_attention/ops/gdn2.py b/python/cudnn/linear_attention/ops/gdn2.py index bee0a35a5..6b36d3c10 100644 --- a/python/cudnn/linear_attention/ops/gdn2.py +++ b/python/cudnn/linear_attention/ops/gdn2.py @@ -264,7 +264,7 @@ def _build_fprop_graph( cu_t = graph.tensor([N + 1], data_type=cu_dtype, name="cu_seqlens") state0_t = None if state_dtype is not None: - state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") + state0_t = graph.tensor([N, HO, V, K], data_type=state_dtype, name="initial_state") a_log_t = None dt_bias_t = None if safe_gate: @@ -456,12 +456,12 @@ def _gdn2_fwd( variant_pack[t["dt_bias"]] = dt_bias final_state = torch.empty(0, dtype=torch.float32, device=device) if output_final_state: - final_state = torch.empty(N, HO, K, V, dtype=torch.float32, device=device) + final_state = torch.empty(N, HO, V, K, dtype=torch.float32, device=device) variant_pack[t["fs"]] = final_state state_checkpoints = torch.empty(0, dtype=q.dtype, device=device) if ckpt > 0: - total_checkpoints = max(total // ckpt, 1) - state_checkpoints = torch.empty(total_checkpoints, HO, K, V, dtype=q.dtype, device=device) + total_checkpoints = max(total // ckpt + N, 1) + state_checkpoints = torch.empty(total_checkpoints, HO, V, K, dtype=q.dtype, device=device) variant_pack[t["state_checkpoints"]] = state_checkpoints graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) return o, final_state, state_checkpoints @@ -501,10 +501,10 @@ def _gdn2_fwd_fake( if initial_state is not None and initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") o = q.new_empty(total, HO, V) - final = q.new_empty((N, HO, K, V) if output_final_state else (0,), dtype=torch.float32) + final = q.new_empty((N, HO, V, K) if output_final_state else (0,), dtype=torch.float32) if checkpoint_every_n_tokens > 0: - total_checkpoints = max(total // int(checkpoint_every_n_tokens), 1) - state_checkpoints = q.new_empty(total_checkpoints, HO, K, V) + total_checkpoints = max(total // int(checkpoint_every_n_tokens) + N, 1) + state_checkpoints = q.new_empty(total_checkpoints, HO, V, K) else: state_checkpoints = q.new_empty(0) return o, final, state_checkpoints @@ -549,13 +549,13 @@ def _build_bprop_graph( dO_t = graph.tensor([total, HO, V], data_type=io_dtype, name="dO") state0_t = None if state_dtype is not None: - state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") + state0_t = graph.tensor([N, HO, V, K], data_type=state_dtype, name="initial_state") dfs_t = None if dstate_in_dtype is not None: - dfs_t = graph.tensor([N, HO, K, V], data_type=dstate_in_dtype, name="d_final_state") + dfs_t = graph.tensor([N, HO, V, K], data_type=dstate_in_dtype, name="d_final_state") ckpts_t = None if ckpt_rows is not None: - ckpts_t = graph.tensor([ckpt_rows, HO, K, V], data_type=io_dtype, name="state_checkpoints") + ckpts_t = graph.tensor([ckpt_rows, HO, V, K], data_type=io_dtype, name="state_checkpoints") a_log_t = None dt_bias_t = None if safe_gate: @@ -995,7 +995,7 @@ def gated_delta_net_v2( HK = HV for canonical GQA: grouped k/v heads shared across query groups); v: ``[total_tokens, HV, V]``; g, beta: ``[total_tokens, HO, K]``; w: ``[total_tokens, HO, V]``; cu_seqlens: ``[N+1]`` int32; - initial_state / final_state: ``[N, HO, K, V]`` + initial_state / final_state: ``[N, HO, V, K]`` (``HO = max(H, HV)``: the gates, output, and state heads). A dense batch of N equal-length sequences is expressed as @@ -1028,7 +1028,7 @@ def gated_delta_net_v2( a_log: ``[HO]`` float32 safe-gate per-head log-amplitude. dt_bias: ``[HO, K]`` float32 safe-gate channel bias. checkpoint_every_n_tokens: if ``> 0``, also return the per-chunk - recurrent state series ``state_checkpoints`` (``[total_checkpoints, HO, K, V]`` io dtype, + recurrent state series ``state_checkpoints`` (``[total_checkpoints, HO, V, K]`` io dtype, one entry per N tokens strictly before each sequence end; the FROST engine requires a positive multiple of the kernel chunk size, 16). The series is a non-differentiable dump. diff --git a/python/cudnn/linear_attention/ops/kda.py b/python/cudnn/linear_attention/ops/kda.py index 2bdff166e..9be778f7b 100644 --- a/python/cudnn/linear_attention/ops/kda.py +++ b/python/cudnn/linear_attention/ops/kda.py @@ -270,7 +270,7 @@ def _build_fprop_graph( cu_t = graph.tensor([N + 1], data_type=cu_dtype, name="cu_seqlens") state0_t = None if state_dtype is not None: - state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") + state0_t = graph.tensor([N, HO, V, K], data_type=state_dtype, name="initial_state") a_log_t = None dt_bias_t = None if safe_gate: @@ -456,12 +456,12 @@ def _kda_fwd( variant_pack[t["dt_bias"]] = dt_bias final_state = torch.empty(0, dtype=torch.float32, device=device) if output_final_state: - final_state = torch.empty(N, HO, K, V, dtype=torch.float32, device=device) + final_state = torch.empty(N, HO, V, K, dtype=torch.float32, device=device) variant_pack[t["fs"]] = final_state state_checkpoints = torch.empty(0, dtype=q.dtype, device=device) if ckpt > 0: - total_checkpoints = max(total // ckpt, 1) - state_checkpoints = torch.empty(total_checkpoints, HO, K, V, dtype=q.dtype, device=device) + total_checkpoints = max(total // ckpt + N, 1) + state_checkpoints = torch.empty(total_checkpoints, HO, V, K, dtype=q.dtype, device=device) variant_pack[t["state_checkpoints"]] = state_checkpoints graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) return o, final_state, state_checkpoints @@ -500,10 +500,10 @@ def _kda_fwd_fake( if initial_state is not None and initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") o = q.new_empty(total, HO, V) - final = q.new_empty((N, HO, K, V) if output_final_state else (0,), dtype=torch.float32) + final = q.new_empty((N, HO, V, K) if output_final_state else (0,), dtype=torch.float32) if checkpoint_every_n_tokens > 0: - total_checkpoints = max(total // int(checkpoint_every_n_tokens), 1) - state_checkpoints = q.new_empty(total_checkpoints, HO, K, V) + total_checkpoints = max(total // int(checkpoint_every_n_tokens) + N, 1) + state_checkpoints = q.new_empty(total_checkpoints, HO, V, K) else: state_checkpoints = q.new_empty(0) return o, final, state_checkpoints @@ -547,13 +547,13 @@ def _build_bprop_graph( dO_t = graph.tensor([total, HO, V], data_type=io_dtype, name="dO") state0_t = None if state_dtype is not None: - state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") + state0_t = graph.tensor([N, HO, V, K], data_type=state_dtype, name="initial_state") dfs_t = None if dstate_in_dtype is not None: - dfs_t = graph.tensor([N, HO, K, V], data_type=dstate_in_dtype, name="d_final_state") + dfs_t = graph.tensor([N, HO, V, K], data_type=dstate_in_dtype, name="d_final_state") ckpts_t = None if ckpt_rows is not None: - ckpts_t = graph.tensor([ckpt_rows, HO, K, V], data_type=io_dtype, name="state_checkpoints") + ckpts_t = graph.tensor([ckpt_rows, HO, V, K], data_type=io_dtype, name="state_checkpoints") a_log_t = None dt_bias_t = None if safe_gate: @@ -983,7 +983,7 @@ def kimi_delta_attention( HK = HV for canonical GQA: grouped K/V heads shared across query groups); v: ``[total_tokens, HV, V]`` g: ``[total_tokens, HO, K]`` (per-key-channel log decay); beta: ``[total_tokens, HO]`` (scalar); cu_seqlens: ``[N+1]`` int32 - initial_state / final_state: ``[N, HO, K, V]`` + initial_state / final_state: ``[N, HO, V, K]`` (``HO = max(H, HV)``: the gates, output, and state heads) A dense batch of N equal-length sequences is expressed as @@ -1020,7 +1020,7 @@ def kimi_delta_attention( a_log: ``[HO]`` float32 safe-gate per-head log-amplitude. dt_bias: ``[HO, K]`` float32 safe-gate channel bias. checkpoint_every_n_tokens: if ``> 0``, also return the per-chunk - recurrent state series ``state_checkpoints`` (``[total_checkpoints, HO, K, V]`` io dtype, + recurrent state series ``state_checkpoints`` (``[total_checkpoints, HO, V, K]`` io dtype, one entry per N tokens strictly before each sequence end; the FROST engine requires a positive multiple of the kernel chunk size, 16). The series is a non-differentiable dump. diff --git a/test/python/linear_attention/frost/examples/01_gdn_prefill.py b/test/python/linear_attention/frost/examples/01_gdn_prefill.py index 547300e49..c3637df28 100644 --- a/test/python/linear_attention/frost/examples/01_gdn_prefill.py +++ b/test/python/linear_attention/frost/examples/01_gdn_prefill.py @@ -10,7 +10,7 @@ o_t = q_t S_t THD layout: token-packed ``[total, H, D]`` tensors plus ``cu_seqlens`` -sequence boundaries; the final state comes back K-major ``[N, H, K, V]``. +sequence boundaries; the final state comes back V-major ``[N, H, V, K]``. """ from __future__ import annotations diff --git a/test/python/linear_attention/reference_gdn.py b/test/python/linear_attention/reference_gdn.py index 3b0b9f6c2..b431945d7 100644 --- a/test/python/linear_attention/reference_gdn.py +++ b/test/python/linear_attention/reference_gdn.py @@ -71,7 +71,9 @@ def gdn_reference( q, k: ``[B, T, Hq/Hk, K]``; v: ``[B, T, Hv, V]``; g, beta: ``[B, T, Hg/Hb]``. Head counts must divide ``HO = max(Hq, Hv)``. scale: applied to q; defaults to ``1/sqrt(K)``. - initial_state: ``[B, HO, K, V]`` (or ``[N, HO, K, V]`` with cu_seqlens). + initial_state: ``[B, HO, V, K]`` (or ``[N, HO, V, K]`` with cu_seqlens), + V-major -- the kernel ABI. Transposed on entry; the fp64 + recurrence below runs in the natural ``S[k, v]``. cu_seqlens: packed varlen boundaries (requires B == 1). safe_gate: treat ``g`` as raw logits and use the log decay ``-exp(a_log) * softplus(g + dt_bias)`` (differentiable; a_log / @@ -80,7 +82,7 @@ def gdn_reference( Returns: ``(o, final_state)`` in fp64: o ``[B, T, HO, V]``, final_state - ``[B, HO, K, V]`` (``[N, HO, K, V]`` with cu_seqlens). + ``[B, HO, V, K]`` (``[N, HO, V, K]`` with cu_seqlens), V-major. """ K = q.shape[-1] if scale is None: @@ -125,9 +127,9 @@ def gdn_reference( if initial_state is None: state0 = torch.zeros(B, HV, K, V, dtype=torch.float64, device=q.device) else: - state0 = initial_state.double() + state0 = initial_state.double().transpose(-1, -2).contiguous() o, state = recurrent_dense(qf, kf, vf, alphaf, betaf, state0) - return o.permute(0, 2, 1, 3), state + return o.permute(0, 2, 1, 3), state.transpose(-1, -2).contiguous() assert q.shape[0] == 1, "cu_seqlens requires packed batch B == 1" bounds = cu_seqlens.tolist() @@ -137,7 +139,7 @@ def gdn_reference( if initial_state is None: state0 = torch.zeros(1, HV, K, V, dtype=torch.float64, device=q.device) else: - state0 = initial_state[n : n + 1].double() + state0 = initial_state[n : n + 1].double().transpose(-1, -2).contiguous() if e == s: states.append(state0) continue @@ -148,4 +150,4 @@ def gdn_reference( o = torch.cat(outs, dim=2).permute(0, 2, 1, 3) else: o = qf.new_zeros(1, 0, HV, V) - return o, torch.cat(states, dim=0) + return o, torch.cat(states, dim=0).transpose(-1, -2).contiguous() diff --git a/test/python/linear_attention/reference_gdn2.py b/test/python/linear_attention/reference_gdn2.py index 1bac23eab..258b6034a 100644 --- a/test/python/linear_attention/reference_gdn2.py +++ b/test/python/linear_attention/reference_gdn2.py @@ -27,7 +27,7 @@ All math runs in fp64 on the input device and is differentiable, so it doubles as the gradient oracle for the bprop tests. The recurrent state is -kept K-major ``[N, HO, K, V]`` here, matching the kernel ABI (KV, v +kept V-major ``[N, HO, V, K]`` here, matching the kernel ABI (VK, k contiguous). """ @@ -92,8 +92,8 @@ def gdn2_reference( per-key erase gate); w: ``[B, T, Hw, V]`` (per-value write gate). Head counts must divide ``HO = max(Hq, Hv)``. scale: applied to q; defaults to ``1/sqrt(K)``. - initial_state: ``[B, HO, K, V]`` (or ``[N, HO, K, V]`` with cu_seqlens), - K-major. + initial_state: ``[B, HO, V, K]`` (or ``[N, HO, V, K]`` with cu_seqlens), V-major, + V-major. cu_seqlens: packed varlen boundaries (requires B == 1). safe_gate: treat ``g`` as raw logits and use the log decay ``gate_lower_bound * sigmoid(exp(a_log) * (g + dt_bias))`` @@ -103,7 +103,7 @@ def gdn2_reference( Returns: ``(o, final_state)`` in fp64: o ``[B, T, HO, V]``, final_state - ``[B, HO, K, V]`` (``[N, HO, K, V]`` with cu_seqlens), K-major. + ``[B, HO, V, K]`` (``[N, HO, V, K]`` with cu_seqlens), V-major. """ K = q.shape[-1] if scale is None: @@ -153,9 +153,9 @@ def gdn2_reference( if initial_state is None: state0 = torch.zeros(B, HV, K, V, dtype=torch.float64, device=q.device) else: - state0 = initial_state.double() + state0 = initial_state.double().transpose(-1, -2).contiguous() o, state = recurrent_dense(qf, kf, vf, alphaf, betaf, wf, state0) - return o.permute(0, 2, 1, 3), state + return o.permute(0, 2, 1, 3), state.transpose(-1, -2).contiguous() assert q.shape[0] == 1, "cu_seqlens requires packed batch B == 1" bounds = cu_seqlens.tolist() @@ -165,7 +165,7 @@ def gdn2_reference( if initial_state is None: state0 = torch.zeros(1, HV, K, V, dtype=torch.float64, device=q.device) else: - state0 = initial_state[n : n + 1].double() + state0 = initial_state[n : n + 1].double().transpose(-1, -2).contiguous() if e == s: states.append(state0) continue @@ -176,4 +176,4 @@ def gdn2_reference( o = torch.cat(outs, dim=2).permute(0, 2, 1, 3) else: o = qf.new_zeros(1, 0, HV, V) - return o, torch.cat(states, dim=0) + return o, torch.cat(states, dim=0).transpose(-1, -2).contiguous() diff --git a/test/python/linear_attention/reference_kda.py b/test/python/linear_attention/reference_kda.py index 1f0894a4c..438fbf855 100644 --- a/test/python/linear_attention/reference_kda.py +++ b/test/python/linear_attention/reference_kda.py @@ -81,7 +81,7 @@ def kda_reference( (log-space per-channel decay); beta: ``[B, T, Hb]`` (scalar per token/head). Head counts must divide ``HO = max(Hq, Hv)``. scale: applied to q; defaults to ``1/sqrt(K)``. - initial_state: ``[B, HO, K, V]`` (or ``[N, HO, K, V]`` with cu_seqlens). + initial_state: ``[B, HO, V, K]`` (or ``[N, HO, V, K]`` with cu_seqlens), V-major. cu_seqlens: packed varlen boundaries (requires B == 1). safe_gate: treat ``g`` as raw logits and use the log decay ``gate_lower_bound * sigmoid(exp(a_log) * (g + dt_bias))`` @@ -91,7 +91,7 @@ def kda_reference( Returns: ``(o, final_state)`` in fp64: o ``[B, T, HO, V]``, final_state - ``[B, HO, K, V]`` (``[N, HO, K, V]`` with cu_seqlens). + ``[B, HO, V, K]`` (``[N, HO, V, K]`` with cu_seqlens), V-major. """ K = q.shape[-1] if scale is None: @@ -138,9 +138,9 @@ def kda_reference( if initial_state is None: state0 = torch.zeros(B, HV, K, V, dtype=torch.float64, device=q.device) else: - state0 = initial_state.double() + state0 = initial_state.double().transpose(-1, -2).contiguous() o, state = recurrent_dense(qf, kf, vf, alphaf, betaf, state0) - return o.permute(0, 2, 1, 3), state + return o.permute(0, 2, 1, 3), state.transpose(-1, -2).contiguous() assert q.shape[0] == 1, "cu_seqlens requires packed batch B == 1" bounds = cu_seqlens.tolist() @@ -150,7 +150,7 @@ def kda_reference( if initial_state is None: state0 = torch.zeros(1, HV, K, V, dtype=torch.float64, device=q.device) else: - state0 = initial_state[n : n + 1].double() + state0 = initial_state[n : n + 1].double().transpose(-1, -2).contiguous() if e == s: states.append(state0) continue @@ -161,4 +161,4 @@ def kda_reference( o = torch.cat(outs, dim=2).permute(0, 2, 1, 3) else: o = qf.new_zeros(1, 0, HV, V) - return o, torch.cat(states, dim=0) + return o, torch.cat(states, dim=0).transpose(-1, -2).contiguous() diff --git a/test/python/linear_attention/test_fla_compat.py b/test/python/linear_attention/test_fla_compat.py index b76b0aac4..c8f517ee0 100644 --- a/test/python/linear_attention/test_fla_compat.py +++ b/test/python/linear_attention/test_fla_compat.py @@ -293,6 +293,40 @@ def test_fallback_is_transparent(): torch.testing.assert_close(o_cud, o_fla, rtol=0, atol=0) +@pytest.mark.parametrize("state_v_first,expect_native", [(True, True), (False, False)]) +def test_state_v_first_routing(state_v_first, expect_native): + """cuDNN carries the recurrent state V-major, so it serves ``state_v_first=True`` + natively and declines the K-major request; a stateless call is layout-agnostic + and runs native either way.""" + m = _master(2, 256, 4, 4, 128, 128, seed=2) + lv = _leaves(m, torch.bfloat16) + o_cud, fs_cud = shim( + lv["q"], + lv["k"], + lv["v"], + lv["g"], + lv["beta"], + output_final_state=True, + state_v_first=state_v_first, + ) + got = last_path() + if expect_native: + assert got == "native", f"state_v_first={state_v_first}: expected native, got {got}" + assert fs_cud.shape == (m["q"].shape[0], m["v"].shape[2], m["v"].shape[3], m["q"].shape[3]) + else: + assert got.startswith("fallback"), f"state_v_first={state_v_first}: expected fallback, got {got}" + o_fla, fs_fla = chunk_gated_delta_rule( + lv["q"], + lv["k"], + lv["v"], + lv["g"], + lv["beta"], + output_final_state=True, + state_v_first=state_v_first, + ) + torch.testing.assert_close(o_cud, o_fla, rtol=0, atol=0) + + def test_accelerate_fla_patches_and_restores(): original = fla_gdr.chunk_gated_delta_rule try: diff --git a/test/python/linear_attention/test_la.py b/test/python/linear_attention/test_la.py index ebd086a38..7febb21d1 100644 --- a/test/python/linear_attention/test_la.py +++ b/test/python/linear_attention/test_la.py @@ -20,6 +20,7 @@ import contextlib import functools import math +import threading import pytest torch = pytest.importorskip("torch") @@ -53,6 +54,7 @@ RAGGED_SEQ_LENS = [ [256, 256], + [96, 32, 160, 1], [511, 501], [64, 128, 512], [31, 63, 93, 123, 150, 500], @@ -62,6 +64,14 @@ EDGE_LENS = [1, 15, 16, 17, 31, 63, 64, 65, 121, 251, 257] DETERMINISM_REPEATS = 8 +STRESS_REPEATS = 4 +STRESS_SHAPES = [ + dict(seq_lens=[96, 32, 160, 1]), + dict(seq_lens=[64, 64]), + dict(seq_lens=[255, 1]), + dict(seq_lens=[1]), + dict(B=8, T=192, H=64), +] SEED = 888 DTYPE_IDS = {torch.bfloat16: "bf16", torch.float16: "fp16"} @@ -257,18 +267,23 @@ def reference(case, *, scale=None, initial_state=None, l2norm=False, cu=None): return fn(*args, **kwargs) -def check(name, out, ref, tol): +def assert_rms_close(name, out, want, tol): out = out.float() assert torch.isfinite(out).all(), f"non-finite values in {name}" - r = rms_ratio(out.reshape(ref.shape), ref) + r = rms_ratio(out.reshape(want.shape), want) assert r < tol, f"{name} rms ratio {r:.4g} >= {tol}" -def check_fwd(case, o, fs, *, scale=None, initial_state=None, l2norm=False, tol_mult=1.0): - o_ref, fs_ref = reference(case, scale=scale, initial_state=initial_state, l2norm=l2norm) - check("o", o, o_ref, tol_mult * FWD_TOL[case.dtype]) +def assert_fwd_parity(backend, case, *, scale=None, use_initial_state=False, l2norm=False, seed=SEED + 1): + set_seed(seed) + state0 = None + if use_initial_state: + state0 = torch.randn(case.N, case.HO, case.V, case.K, device="cuda", dtype=torch.float32) * 0.05 + o, fs = run_fwd(backend, case, scale=scale, initial_state=state0, output_final_state=True, use_qk_l2norm_in_kernel=l2norm) + o_ref, fs_ref = reference(case, scale=scale, initial_state=state0, l2norm=l2norm) + assert_rms_close("o", o, o_ref, FWD_TOL[case.dtype]) if fs is not None and fs.numel(): - check("final_state", fs, fs_ref, tol_mult * STATE_TOL[case.dtype]) + assert_rms_close("final_state", fs, fs_ref, STATE_TOL[case.dtype]) # --------------------------------------------------------------------------- @@ -301,8 +316,7 @@ def test_fwd_basic(backend, variant, dtype, B, T, H, HV): if dtype == torch.float16 and (H, HV) not in HEAD_CONFIGS_SMALL: pytest.skip("fp16 runs the small head matrix") case = make_case(variant, dtype, B=B, T=T, H=H, HV=HV) - o, fs = run_fwd(backend, case, output_final_state=True) - check_fwd(case, o, fs) + assert_fwd_parity(backend, case) @pytest.mark.parametrize("alpha,beta,w", [(True, False, True), (False, True, True), (True, True, False)], ids=["no_beta", "no_alpha", "no_w"]) @@ -313,16 +327,14 @@ def test_fwd_gate_combinations(backend, variant, alpha, beta, w): if variant != "gdn2" and not w: pytest.skip("w is a GDN-2 gate") case = make_case(variant, torch.bfloat16, T=192, alpha=alpha, beta=beta, w=w) - o, fs = run_fwd(backend, case, output_final_state=True) - check_fwd(case, o, fs) + assert_fwd_parity(backend, case) @pytest.mark.parametrize("scale", [0.5, 1.0, None], ids=["half", "one", "auto"]) @pytest.mark.parametrize("variant", VARIANTS) def test_fwd_scale(backend, variant, scale): case = make_case(variant, torch.bfloat16, T=192) - o, fs = run_fwd(backend, case, scale=scale, output_final_state=True) - check_fwd(case, o, fs, scale=scale) + assert_fwd_parity(backend, case, scale=scale) @pytest.mark.parametrize("variant", VARIANTS) @@ -338,8 +350,7 @@ def test_fwd_default_scale_matches_explicit(backend, variant): def test_fwd_seqlen_edges(backend, variant, T): """Lengths straddling the kernels' chunk boundaries (16 and 64).""" case = make_case(variant, torch.bfloat16, T=T) - o, fs = run_fwd(backend, case, output_final_state=True) - check_fwd(case, o, fs) + assert_fwd_parity(backend, case) @pytest.mark.parametrize("H,HV", [(1, 1), (2, 4)]) @@ -347,8 +358,7 @@ def test_fwd_seqlen_edges(backend, variant, T): @pytest.mark.parametrize("variant", VARIANTS) def test_fwd_varlen_ragged(backend, variant, seq_lens, H, HV): case = make_case(variant, torch.bfloat16, seq_lens=seq_lens, H=H, HV=HV) - o, fs = run_fwd(backend, case, output_final_state=True) - check_fwd(case, o, fs) + assert_fwd_parity(backend, case) @pytest.mark.parametrize("variant", VARIANTS) @@ -359,8 +369,6 @@ def test_fwd_many_short_sequences(backend, variant): o, fs = run_fwd(backend, case, output_final_state=True) cu1 = torch.tensor([0, T], dtype=torch.int32, device="cuda") for n in (0, 1, 99, 199): - # clone: sliced views can start at non-16B-aligned offsets, which the - # kernels' buffer contract rejects sl = slice(T * n, T * (n + 1)) args = [ to_thd(case.q)[sl].clone(), @@ -373,8 +381,8 @@ def test_fwd_many_short_sequences(backend, variant): args.append(to_thd(case.gates["w"])[sl].clone()) with waive_unsupported(backend, variant): o_n, fs_n = pinned_op(backend, variant)(*args, cu1, output_final_state=True) - check(f"o[seq {n}]", o[sl], o_n.float(), FWD_TOL[torch.bfloat16]) - check(f"final_state[seq {n}]", fs[n], fs_n[0].float(), STATE_TOL[torch.bfloat16]) + assert_rms_close(f"o[seq {n}]", o[sl], o_n.float(), FWD_TOL[torch.bfloat16]) + assert_rms_close(f"final_state[seq {n}]", fs[n], fs_n[0].float(), STATE_TOL[torch.bfloat16]) @pytest.mark.parametrize("variant", VARIANTS) @@ -389,7 +397,7 @@ def test_fwd_zero_length_sequences(backend, variant): torch.testing.assert_close(fs[0], fs_base[0], atol=1e-3, rtol=1e-3) torch.testing.assert_close(fs[2], fs_base[1], atol=1e-3, rtol=1e-3) assert (fs[1] == 0).all() and (fs[3] == 0).all(), "zero-length sequence states must stay zero" - state0 = torch.randn(4, case.HO, case.K, case.V, device="cuda", dtype=torch.float32) * 0.05 + state0 = torch.randn(4, case.HO, case.V, case.K, device="cuda", dtype=torch.float32) * 0.05 o, fs_state0 = run_fwd(backend, case, cu=cu, initial_state=state0, output_final_state=True) torch.testing.assert_close(fs_state0[1], state0[1], atol=0.0, rtol=0.0) torch.testing.assert_close(fs_state0[3], state0[3], atol=0.0, rtol=0.0) @@ -399,19 +407,17 @@ def test_fwd_zero_length_sequences(backend, variant): @pytest.mark.parametrize("variant", VARIANTS) def test_fwd_initial_state(backend, variant, T): case = make_case(variant, torch.bfloat16, T=T) - state0 = torch.randn(case.N, case.HO, case.K, case.V, device="cuda", dtype=torch.float32) * 0.05 - o, fs = run_fwd(backend, case, initial_state=state0, output_final_state=True) - check_fwd(case, o, fs, initial_state=state0) + assert_fwd_parity(backend, case, use_initial_state=True) @pytest.mark.parametrize("T1,T2", [(128, 128), (64, 192), (192, 121)]) @pytest.mark.parametrize("variant", VARIANTS) def test_fwd_chunked_prefill(backend, variant, T1, T2): - """Two-phase prefill: part 1's final state feeds part 2; the concatenated + """Two-phase prefill: phase 1's final state feeds phase 2; the concatenated output matches a single-shot reference (state round-trips through fp32).""" case = make_case(variant, torch.bfloat16, B=2, T=T1 + T2) - def part(t0, t1, state0): + def run_phase(t0, t1, state0): sub = case.clone() sub.q, sub.k, sub.v = (x[:, t0:t1].contiguous() for x in (case.q, case.k, case.v)) sub.gates = {n: g[:, t0:t1].contiguous() for n, g in case.gates.items()} @@ -419,12 +425,12 @@ def part(t0, t1, state0): sub.cu = torch.arange(0, case.B + 1, dtype=torch.int32, device="cuda") * sub.T return run_fwd(backend, sub, initial_state=state0, output_final_state=True) - o1, fs1 = part(0, T1, None) - o2, fs2 = part(T1, T1 + T2, fs1) + o1, fs1 = run_phase(0, T1, None) + o2, fs2 = run_phase(T1, T1 + T2, fs1) o = torch.cat([o1.reshape(case.B, T1, case.HO, case.V), o2.reshape(case.B, T2, case.HO, case.V)], dim=1) o_ref, fs_ref = reference(case) - check("o", o, o_ref, 1.5 * FWD_TOL[case.dtype]) - check("final_state", fs2, fs_ref, 1.5 * STATE_TOL[case.dtype]) + assert_rms_close("o", o, o_ref, 1.5 * FWD_TOL[case.dtype]) + assert_rms_close("final_state", fs2, fs_ref, 1.5 * STATE_TOL[case.dtype]) @pytest.mark.parametrize("variant", VARIANTS) @@ -450,8 +456,7 @@ def test_fwd_packed_matches_per_sequence(backend, variant): def test_fwd_head_dims(backend, variant, K, V): """K/V head-dim variants; engines that only serve K = V = 128 decline.""" case = make_case(variant, torch.bfloat16, T=192, K=K, V=V) - o, fs = run_fwd(backend, case, output_final_state=True) - check_fwd(case, o, fs) + assert_fwd_parity(backend, case) @pytest.mark.parametrize("H,HK,HV", GQA_CONFIGS) @@ -460,8 +465,7 @@ def test_fwd_gqa(backend, variant, H, HK, HV): """Grouped heads: canonical GQA (native K at HK == HV), the expanded-k form, and shared-kv GVA (HK == HV > H); gates/O/states live at HO = max(H, HV).""" case = make_case(variant, torch.bfloat16, T=192, H=H, HK=HK, HV=HV) - o, fs = run_fwd(backend, case, output_final_state=True) - check_fwd(case, o, fs) + assert_fwd_parity(backend, case) @pytest.mark.parametrize("variant", VARIANTS) @@ -470,23 +474,20 @@ def test_fwd_multi_tile(backend, variant): to back, exercising the inter-tile state drain -> seed ordering the single-tile cases never reach.""" case = make_case(variant, torch.bfloat16, B=8, T=192, H=64) - o, fs = run_fwd(backend, case, output_final_state=True) - check_fwd(case, o, fs) + assert_fwd_parity(backend, case) @pytest.mark.parametrize("variant", VARIANTS) def test_fwd_qk_l2norm(backend, variant): """In-kernel Q/K L2 norm matches the reference on pre-normalized inputs.""" case = make_case(variant, torch.bfloat16, T=256) - o, fs = run_fwd(backend, case, output_final_state=True, use_qk_l2norm_in_kernel=True) - check_fwd(case, o, fs, l2norm=True) + assert_fwd_parity(backend, case, l2norm=True) @pytest.mark.parametrize("variant", VARIANTS) def test_fwd_strong_decay_varlen(backend, variant): case = make_case(variant, torch.bfloat16, seq_lens=[100, 2048, 0, 517], lo=0.1 if variant == "gdn" else 0.3) - o, fs = run_fwd(backend, case, output_final_state=True) - check_fwd(case, o, fs) + assert_fwd_parity(backend, case) @pytest.mark.parametrize("variant", VARIANTS) @@ -497,7 +498,7 @@ def test_fwd_output_contract(backend, variant): assert o.shape == (case.T, case.HO, case.V) and o.dtype == case.dtype assert fs.numel() == 0 o, fs = run_fwd(backend, case, output_final_state=True) - assert fs.shape == (case.N, case.HO, case.K, case.V) and fs.dtype == torch.float32 + assert fs.shape == (case.N, case.HO, case.V, case.K) and fs.dtype == torch.float32 # --------------------------------------------------------------------------- @@ -515,7 +516,7 @@ def assert_bwd_parity(backend, case, *, scale=None, use_initial_state=False, use set_seed(seed) state0_op = state0_ref = None if use_initial_state: - state0 = torch.randn(case.N, case.HO, case.K, case.V, device="cuda", dtype=torch.float32) * 0.05 + state0 = torch.randn(case.N, case.HO, case.V, case.K, device="cuda", dtype=torch.float32) * 0.05 state0_op = state0.detach().clone().requires_grad_(True) state0_ref = state0.detach().double().requires_grad_(True) @@ -560,7 +561,7 @@ def assert_bwd_parity(backend, case, *, scale=None, use_initial_state=False, use tol_n = gate_grad_tol else: tol_n = tol - check(f"d{name}", got, want, tol_n) + assert_rms_close(f"d{name}", got, want, tol_n) @pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL + [(16, 64)]) @@ -651,51 +652,57 @@ def test_bwd_with_checkpoints(backend, variant): @pytest.mark.parametrize("variant", VARIANTS) def test_checkpoints_match_prefix_final_states(backend, variant): - """state_checkpoints[j] is the state after (j+1)*ckpt tokens, strictly before the - end; rows are a shape-derived capacity bound, valid entries pack first.""" + """state_checkpoints[j] is the state AT token boundary j*ckpt, so row 0 is the + state entering the sequence and the end is excluded; rows are a shape-derived + capacity bound, valid entries pack first.""" ckpt = CHUNK[variant] T = 5 * ckpt case = make_case(variant, torch.bfloat16, T=T) o, fs, state_checkpoints = run_fwd(backend, case, output_final_state=True, checkpoint_every_n_tokens=ckpt) - valid = (T - 1) // ckpt - assert state_checkpoints.shape == (T // ckpt, case.HO, case.K, case.V) + valid = (T - 1) // ckpt + 1 + assert state_checkpoints.shape == (T // ckpt + 1, case.HO, case.V, case.K) assert state_checkpoints.dtype == case.dtype - for j in (0, valid - 1): - n = (j + 1) * ckpt + # row 0 is the incoming state, zero here since no initial_state was passed + assert not state_checkpoints[0].any(), "row 0 must be the (zero) incoming state" + for j in sorted({1, valid - 1}): + n = j * ckpt args = [to_thd(case.q)[:n], to_thd(case.k)[:n], to_thd(case.v)[:n], to_thd(case.gates["g"])[:n], to_thd(case.gates["beta"])[:n]] if variant == "gdn2": args.append(to_thd(case.gates["w"])[:n]) cu_n = torch.tensor([0, n], dtype=torch.int32, device="cuda") with waive_unsupported(backend, variant): o, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) - check(f"state_checkpoints[{j}]", state_checkpoints[j], fs_p[0], STATE_TOL[case.dtype]) + assert_rms_close(f"state_checkpoints[{j}]", state_checkpoints[j], fs_p[0], STATE_TOL[case.dtype]) @pytest.mark.parametrize("variant", VARIANTS) def test_checkpoints_varlen(backend, variant): - """Entries pack per sequence in order (one per ckpt tokens strictly before - each sequence end); each entry matches its sequence's truncated prefix.""" + """Entries pack per sequence in order (one per ckpt tokens from the sequence + start, end excluded); each entry matches its sequence's truncated prefix.""" ckpt = CHUNK[variant] seq_lens = [3 * ckpt + 5, ckpt - 1, 0, 2 * ckpt] case = make_case(variant, torch.bfloat16, seq_lens=seq_lens) o, fs, state_checkpoints = run_fwd(backend, case, output_final_state=True, checkpoint_every_n_tokens=ckpt) - counts = [max(sl - 1, 0) // ckpt for sl in seq_lens] + counts = [(sl - 1) // ckpt + 1 if sl > 0 else 0 for sl in seq_lens] # shape[0] is the shape-derived capacity bound; the packed prefix holds # sum(counts) real rows (per-sequence, in order), the tail is uninitialized - assert state_checkpoints.shape[0] == max(sum(seq_lens) // ckpt, 1) + assert state_checkpoints.shape[0] == max(sum(seq_lens) // ckpt + len(seq_lens), 1) bounds = case.cu.tolist() base = 0 for n, cnt in enumerate(counts): for j in sorted({0, cnt - 1} if cnt else set()): + if j == 0: + assert not state_checkpoints[base].any(), f"seq {n} row 0 must be the (zero) incoming state" + continue n0 = bounds[n] - ntok = (j + 1) * ckpt + ntok = j * ckpt args = [to_thd(t)[n0 : n0 + ntok].clone() for t in (case.q, case.k, case.v, case.gates["g"], case.gates["beta"])] if variant == "gdn2": args.append(to_thd(case.gates["w"])[n0 : n0 + ntok].clone()) cu_n = torch.tensor([0, ntok], dtype=torch.int32, device="cuda") with waive_unsupported(backend, variant): o_p, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) - check(f"state_checkpoints[seq {n}][{j}]", state_checkpoints[base + j], fs_p[0], STATE_TOL[case.dtype]) + assert_rms_close(f"state_checkpoints[seq {n}][{j}]", state_checkpoints[base + j], fs_p[0], STATE_TOL[case.dtype]) base += cnt @@ -717,8 +724,8 @@ def test_checkpoints_varlen_tight_capacity(backend, variant, recipe): reuse matches the recompute path bitwise.""" ckpt = CHUNK[variant] seq_lens = TIGHT_VARLEN_RECIPES[recipe](ckpt) - counts = [max(sl - 1, 0) // ckpt for sl in seq_lens] - assert sum(counts) == max(sum(seq_lens) // ckpt, 1), "recipe must fill the capacity bound exactly" + counts = [(sl - 1) // ckpt + 1 if sl > 0 else 0 for sl in seq_lens] + assert sum(counts) == max(sum(seq_lens) // ckpt + len(seq_lens), 1), "recipe must fill the capacity bound exactly" case = make_case(variant, torch.bfloat16, seq_lens=seq_lens) o, fs, state_checkpoints = run_fwd(backend, case, output_final_state=True, checkpoint_every_n_tokens=ckpt) assert state_checkpoints.shape[0] == sum(counts) @@ -727,8 +734,11 @@ def test_checkpoints_varlen_tight_capacity(backend, variant, recipe): base = 0 for n, cnt in enumerate(counts): for j in range(cnt): + if j == 0: + assert not state_checkpoints[base].any(), f"seq {n} row 0 must be the (zero) incoming state" + continue n0 = bounds[n] - ntok = (j + 1) * ckpt + ntok = j * ckpt ref_args = [t[:, n0 : n0 + ntok] for t in (case.q, case.k, case.v, case.gates["g"], case.gates["beta"])] args = [to_thd(t)[n0 : n0 + ntok].clone() for t in (case.q, case.k, case.v, case.gates["g"], case.gates["beta"])] if variant == "gdn2": @@ -736,11 +746,11 @@ def test_checkpoints_varlen_tight_capacity(backend, variant, recipe): args.append(to_thd(case.gates["w"])[n0 : n0 + ntok].clone()) with torch.no_grad(): _, fs_ref = ref_fn(*ref_args) - check(f"state_checkpoints[seq {n}][{j}] vs fp64 reference", state_checkpoints[base + j], fs_ref[0], STATE_TOL[case.dtype]) + assert_rms_close(f"state_checkpoints[seq {n}][{j}] vs fp64 reference", state_checkpoints[base + j], fs_ref[0], STATE_TOL[case.dtype]) cu_n = torch.tensor([0, ntok], dtype=torch.int32, device="cuda") with waive_unsupported(backend, variant): o_p, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) - check(f"state_checkpoints[seq {n}][{j}] vs solo prefix", state_checkpoints[base + j], fs_p[0], STATE_TOL[case.dtype]) + assert_rms_close(f"state_checkpoints[seq {n}][{j}] vs solo prefix", state_checkpoints[base + j], fs_p[0], STATE_TOL[case.dtype]) base += cnt grads_by_mode = [] for mode_ckpt in (0, ckpt): @@ -765,18 +775,19 @@ def test_checkpoints_coarse_cadence(backend, variant, ckpt_mult): T = 5 * ckpt case = make_case(variant, torch.bfloat16, T=T) o, fs, state_checkpoints = run_fwd(backend, case, output_final_state=True, checkpoint_every_n_tokens=ckpt) - valid = (T - 1) // ckpt - assert state_checkpoints.shape == (T // ckpt, case.HO, case.K, case.V) + valid = (T - 1) // ckpt + 1 + assert state_checkpoints.shape == (T // ckpt + 1, case.HO, case.V, case.K) assert state_checkpoints.dtype == case.dtype - for j in (0, valid - 1): - n = (j + 1) * ckpt + assert not state_checkpoints[0].any(), "row 0 must be the (zero) incoming state" + for j in sorted({1, valid - 1}): + n = j * ckpt args = [to_thd(case.q)[:n], to_thd(case.k)[:n], to_thd(case.v)[:n], to_thd(case.gates["g"])[:n], to_thd(case.gates["beta"])[:n]] if variant == "gdn2": args.append(to_thd(case.gates["w"])[:n]) cu_n = torch.tensor([0, n], dtype=torch.int32, device="cuda") with waive_unsupported(backend, variant): o, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) - check(f"state_checkpoints[{j}]", state_checkpoints[j], fs_p[0], STATE_TOL[case.dtype]) + assert_rms_close(f"state_checkpoints[{j}]", state_checkpoints[j], fs_p[0], STATE_TOL[case.dtype]) # --------------------------------------------------------------------------- @@ -822,7 +833,7 @@ def test_safe_gate_forward_parity(backend, variant): eff_case = case.clone(gates=dict(case.gates, g=g_eff, beta=eff_beta)) o_raw, fs_raw = run_fwd(backend, raw_case, **raw_kw) o_eff, fs_eff = run_fwd(backend, eff_case, **kw) - check("o", o_raw, o_eff.double(), 2e-2) + assert_rms_close("o", o_raw, o_eff.double(), 2e-2) assert rms_ratio(fs_raw, fs_eff) < 2e-2 @@ -877,7 +888,7 @@ def test_beta_sigmoid_in_kernel(backend): eff_case = case.clone(gates=dict(case.gates, beta=braw.float().sigmoid())) o_raw, fs_raw = run_fwd(backend, raw_case, output_final_state=True, use_beta_sigmoid_in_kernel=True) o_eff, fs_eff = run_fwd(backend, eff_case, output_final_state=True) - check("o", o_raw, o_eff.double(), 2e-2) + assert_rms_close("o", o_raw, o_eff.double(), 2e-2) assert rms_ratio(fs_raw, fs_eff) < 2e-2 @@ -982,7 +993,7 @@ def test_invalid_gate_dtype_raises(variant): @pytest.mark.parametrize("variant", VARIANTS) def test_invalid_initial_state_count_raises(variant): case = make_case(variant, torch.bfloat16, T=64) - state0 = torch.zeros(3, case.HO, case.K, case.V, device="cuda", dtype=torch.float32) + state0 = torch.zeros(3, case.HO, case.V, case.K, device="cuda", dtype=torch.float32) with pytest.raises(ValueError, match="initial"): op(variant)(*op_args(case), initial_state=state0) @@ -1021,7 +1032,7 @@ def assert_bitwise_runs(launch, repeats=DETERMINISM_REPEATS, label=""): @pytest.mark.parametrize("variant", VARIANTS) def test_determinism_fwd(backend, variant): case = make_case(variant, torch.bfloat16, seq_lens=[497, 16, 1, 480, 0, 253]) - state0 = torch.randn(case.N, case.HO, case.K, case.V, device="cuda", dtype=torch.float32) * 0.05 + state0 = torch.randn(case.N, case.HO, case.V, case.K, device="cuda", dtype=torch.float32) * 0.05 def launch(): o, fs = run_fwd(backend, case, initial_state=state0, output_final_state=True) @@ -1054,7 +1065,7 @@ def test_determinism_multi_tile_fwd(backend, variant): """Multi-tile grid (B*H >> SM count) with an initial state: bitwise stability across the inter-tile drain -> seed window.""" case = make_case(variant, torch.bfloat16, B=8, T=192, H=64) - state0 = torch.randn(case.N, case.HO, case.K, case.V, device="cuda", dtype=torch.float32) * 0.05 + state0 = torch.randn(case.N, case.HO, case.V, case.K, device="cuda", dtype=torch.float32) * 0.05 def launch(): o, fs = run_fwd(backend, case, initial_state=state0, output_final_state=True) @@ -1081,6 +1092,67 @@ def launch(): assert_bitwise_runs(launch, label=f"{variant} multi-tile bwd") +# --------------------------------------------------------------------------- +# Stress: repeated execution, and several live graphs, in one process +# --------------------------------------------------------------------------- + + +def _assert_stress_stable(backend, variant, shapes): + """Round robin over ``shapes`` for STRESS_REPEATS rounds; every shape must + reproduce its first round bitwise. Inputs are seeded, so any drift is the + kernel carrying state between executions.""" + ref = {} + for rnd in range(STRESS_REPEATS): + for si, shape in enumerate(shapes): + got = _stress_round(backend, variant, shape) + if si not in ref: + ref[si] = [t.detach().clone() for t in got] + continue + for name, want, have in zip(("o", "final_state", "dq", "dk"), ref[si], got): + assert torch.equal(want, have), f"shape {si} round {rnd}: {name} drifted between executions" + + +def _stress_round(backend, variant, shape): + """One fwd+bwd on freshly allocated operands; returns the outputs.""" + case = make_case(variant, torch.bfloat16, **shape) + leaves = [to_thd(case.q).detach().clone().requires_grad_(True), to_thd(case.k).detach().clone().requires_grad_(True)] + args = [leaves[0], leaves[1], to_thd(case.v), to_thd(case.gates["g"]), to_thd(case.gates["beta"])] + if variant == "gdn2": + args.append(to_thd(case.gates["w"])) + o, fs = pinned_op(backend, variant)(*args, case.cu, output_final_state=True) + # ones, not randn: the grad seed must not depend on global RNG ordering + grads = torch.autograd.grad([o, fs], leaves, [torch.ones_like(o), torch.ones_like(fs)]) + torch.cuda.synchronize() + return [o, fs, *grads] + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_replay_stress(backend, variant): + """Re-execute one cached plan many times in a single process. + + The ``backend`` fixture clears the op caches around every test, so no other + test runs a SECOND execution of a cached LA plan -- which is what every + training step does. Fresh operands each round also let the caching + allocator move them: an access that runs off the end of a packed tensor + only faults when that tensor happens to sit at the end of its mapping, so + a single execution on one layout proves nothing.""" + with waive_unsupported(backend, variant): + _assert_stress_stable(backend, variant, STRESS_SHAPES[:1]) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_multi_graph_stress(backend, variant): + """Several distinct LA graphs live at once, executed round robin. + + Each shape builds its own cached plan, so this is the multi-graph state a + training or serving loop runs in -- and the state the per-test cache + clearing hides. The dense entry also hands each CTA several work items, so + any per-work-item pipeline state that fails to return to its starting + parity shows up on the second item rather than the first.""" + with waive_unsupported(backend, variant): + _assert_stress_stable(backend, variant, STRESS_SHAPES) + + @pytest.mark.parametrize("backend", ["frost"], indirect=True) @pytest.mark.parametrize("variant", VARIANTS) def test_determinism_two_streams(backend, variant): @@ -1211,7 +1283,8 @@ def test_batch_invariance_with_coarse_checkpoints(backend, variant): T = 4 * ckpt case = make_case(variant, torch.bfloat16, T=T) o, fs, state_checkpoints = run_fwd(backend, case, output_final_state=True, batch_invariant=True, checkpoint_every_n_tokens=ckpt) - assert state_checkpoints.shape == (T // ckpt, case.HO, case.K, case.V) + assert state_checkpoints.shape == (T // ckpt + 1, case.HO, case.V, case.K) + assert not state_checkpoints[0].any(), "row 0 must be the (zero) incoming state" n = ckpt args = [to_thd(case.q)[:n], to_thd(case.k)[:n], to_thd(case.v)[:n], to_thd(case.gates["g"])[:n], to_thd(case.gates["beta"])[:n]] if variant == "gdn2": @@ -1219,7 +1292,41 @@ def test_batch_invariance_with_coarse_checkpoints(backend, variant): cu_n = torch.tensor([0, n], dtype=torch.int32, device="cuda") with waive_unsupported(backend, variant): o_p, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) - check("state_checkpoints[0]", state_checkpoints[0], fs_p[0], STATE_TOL[case.dtype]) + assert_rms_close("state_checkpoints[1]", state_checkpoints[1], fs_p[0], STATE_TOL[case.dtype]) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_execute_from_a_thread_with_no_cuda_context(backend, variant): + """A python engine reaches the driver directly, and the driver reads the + CALLING thread's context stack. A fresh thread is permanently cold, unlike + the autograd worker that the first backward in a process warms.""" + from cuda.bindings import driver as drv + + case = make_case(variant, torch.bfloat16, T=256) + leaves = [to_thd(case.q).detach().clone().requires_grad_(True), to_thd(case.k).detach().clone().requires_grad_(True)] + args = [leaves[0], leaves[1], to_thd(case.v), to_thd(case.gates["g"]), to_thd(case.gates["beta"])] + if variant == "gdn2": + args.append(to_thd(case.gates["w"])) + seen = {} + + def run_on_cold_thread(): + seen["before"] = int(drv.cuCtxGetCurrent()[1]) + try: + # batch_invariant skips the split-K table launch that would bind a context first + o, _ = pinned_op(backend, variant)(*args, case.cu, batch_invariant=True) + torch.autograd.grad([o], leaves, [torch.randn_like(o)]) + except BaseException as exc: # noqa: BLE001 + seen["exc"] = exc + seen["after"] = int(drv.cuCtxGetCurrent()[1]) + + worker = threading.Thread(target=run_on_cold_thread) + with waive_unsupported(backend, variant): + worker.start() + worker.join() + if "exc" in seen: + raise seen["exc"] + assert seen["before"] == 0, "the worker thread was already bound, so this no longer covers the cold path" + assert seen["after"] != 0, "execute left the calling thread with no CUDA context" # --------------------------------------------------------------------------- From 2aa2ad4bf201aec879d53168021b441eb55cbc7b Mon Sep 17 00:00:00 2001 From: Joshua Park Date: Tue, 18 Aug 2026 01:01:38 -0700 Subject: [PATCH 2/2] 2 --- test/python/linear_attention/reference_gdn2.py | 3 +-- test/python/linear_attention/test_fla_compat.py | 2 +- test/python/linear_attention/test_la.py | 14 +++++++------- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/test/python/linear_attention/reference_gdn2.py b/test/python/linear_attention/reference_gdn2.py index 258b6034a..2335c4bc8 100644 --- a/test/python/linear_attention/reference_gdn2.py +++ b/test/python/linear_attention/reference_gdn2.py @@ -92,8 +92,7 @@ def gdn2_reference( per-key erase gate); w: ``[B, T, Hw, V]`` (per-value write gate). Head counts must divide ``HO = max(Hq, Hv)``. scale: applied to q; defaults to ``1/sqrt(K)``. - initial_state: ``[B, HO, V, K]`` (or ``[N, HO, V, K]`` with cu_seqlens), V-major, - V-major. + initial_state: ``[B, HO, V, K]`` (or ``[N, HO, V, K]`` with cu_seqlens), V-major. cu_seqlens: packed varlen boundaries (requires B == 1). safe_gate: treat ``g`` as raw logits and use the log decay ``gate_lower_bound * sigmoid(exp(a_log) * (g + dt_bias))`` diff --git a/test/python/linear_attention/test_fla_compat.py b/test/python/linear_attention/test_fla_compat.py index c8f517ee0..f78c02b70 100644 --- a/test/python/linear_attention/test_fla_compat.py +++ b/test/python/linear_attention/test_fla_compat.py @@ -315,7 +315,7 @@ def test_state_v_first_routing(state_v_first, expect_native): assert fs_cud.shape == (m["q"].shape[0], m["v"].shape[2], m["v"].shape[3], m["q"].shape[3]) else: assert got.startswith("fallback"), f"state_v_first={state_v_first}: expected fallback, got {got}" - o_fla, fs_fla = chunk_gated_delta_rule( + o_fla, _ = chunk_gated_delta_rule( lv["q"], lv["k"], lv["v"], diff --git a/test/python/linear_attention/test_la.py b/test/python/linear_attention/test_la.py index 7febb21d1..02af57bde 100644 --- a/test/python/linear_attention/test_la.py +++ b/test/python/linear_attention/test_la.py @@ -658,7 +658,7 @@ def test_checkpoints_match_prefix_final_states(backend, variant): ckpt = CHUNK[variant] T = 5 * ckpt case = make_case(variant, torch.bfloat16, T=T) - o, fs, state_checkpoints = run_fwd(backend, case, output_final_state=True, checkpoint_every_n_tokens=ckpt) + _, _, state_checkpoints = run_fwd(backend, case, output_final_state=True, checkpoint_every_n_tokens=ckpt) valid = (T - 1) // ckpt + 1 assert state_checkpoints.shape == (T // ckpt + 1, case.HO, case.V, case.K) assert state_checkpoints.dtype == case.dtype @@ -671,7 +671,7 @@ def test_checkpoints_match_prefix_final_states(backend, variant): args.append(to_thd(case.gates["w"])[:n]) cu_n = torch.tensor([0, n], dtype=torch.int32, device="cuda") with waive_unsupported(backend, variant): - o, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) + _, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) assert_rms_close(f"state_checkpoints[{j}]", state_checkpoints[j], fs_p[0], STATE_TOL[case.dtype]) @@ -682,7 +682,7 @@ def test_checkpoints_varlen(backend, variant): ckpt = CHUNK[variant] seq_lens = [3 * ckpt + 5, ckpt - 1, 0, 2 * ckpt] case = make_case(variant, torch.bfloat16, seq_lens=seq_lens) - o, fs, state_checkpoints = run_fwd(backend, case, output_final_state=True, checkpoint_every_n_tokens=ckpt) + _, _, state_checkpoints = run_fwd(backend, case, output_final_state=True, checkpoint_every_n_tokens=ckpt) counts = [(sl - 1) // ckpt + 1 if sl > 0 else 0 for sl in seq_lens] # shape[0] is the shape-derived capacity bound; the packed prefix holds # sum(counts) real rows (per-sequence, in order), the tail is uninitialized @@ -701,7 +701,7 @@ def test_checkpoints_varlen(backend, variant): args.append(to_thd(case.gates["w"])[n0 : n0 + ntok].clone()) cu_n = torch.tensor([0, ntok], dtype=torch.int32, device="cuda") with waive_unsupported(backend, variant): - o_p, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) + _, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) assert_rms_close(f"state_checkpoints[seq {n}][{j}]", state_checkpoints[base + j], fs_p[0], STATE_TOL[case.dtype]) base += cnt @@ -749,7 +749,7 @@ def test_checkpoints_varlen_tight_capacity(backend, variant, recipe): assert_rms_close(f"state_checkpoints[seq {n}][{j}] vs fp64 reference", state_checkpoints[base + j], fs_ref[0], STATE_TOL[case.dtype]) cu_n = torch.tensor([0, ntok], dtype=torch.int32, device="cuda") with waive_unsupported(backend, variant): - o_p, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) + _, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) assert_rms_close(f"state_checkpoints[seq {n}][{j}] vs solo prefix", state_checkpoints[base + j], fs_p[0], STATE_TOL[case.dtype]) base += cnt grads_by_mode = [] @@ -786,7 +786,7 @@ def test_checkpoints_coarse_cadence(backend, variant, ckpt_mult): args.append(to_thd(case.gates["w"])[:n]) cu_n = torch.tensor([0, n], dtype=torch.int32, device="cuda") with waive_unsupported(backend, variant): - o, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) + _, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) assert_rms_close(f"state_checkpoints[{j}]", state_checkpoints[j], fs_p[0], STATE_TOL[case.dtype]) @@ -1291,7 +1291,7 @@ def test_batch_invariance_with_coarse_checkpoints(backend, variant): args.append(to_thd(case.gates["w"])[:n]) cu_n = torch.tensor([0, n], dtype=torch.int32, device="cuda") with waive_unsupported(backend, variant): - o_p, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) + _, fs_p = pinned_op(backend, variant)(*args, cu_n, output_final_state=True) assert_rms_close("state_checkpoints[1]", state_checkpoints[1], fs_p[0], STATE_TOL[case.dtype])