diff --git a/benchmark/linear_attention/Dockerfile b/benchmark/linear_attention/Dockerfile index 16f4889b8..04edf7e1b 100644 --- a/benchmark/linear_attention/Dockerfile +++ b/benchmark/linear_attention/Dockerfile @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -FROM nvcr.io/nvidia/pytorch:26.03-py3 +FROM nvcr.io/nvidia/pytorch:26.07-py3 # Set working directory WORKDIR /workspace @@ -26,5 +26,9 @@ RUN pip install nvidia-cutlass-dsl[cu13]==4.7.0 apache-tvm-ffi flash-linear-atte RUN git clone https://github.com/QwenLM/FlashQLA.git RUN pip install -v /workspace/FlashQLA +# Install FlashKDA from source. +RUN git clone https://github.com/MoonshotAI/FlashKDA.git +RUN pip install -v /workspace/FlashKDA + # Install the chart dependencies for plot_results.py RUN pip install pandas matplotlib seaborn diff --git a/benchmark/linear_attention/README.md b/benchmark/linear_attention/README.md index 676dc6177..7c2dac8c2 100644 --- a/benchmark/linear_attention/README.md +++ b/benchmark/linear_attention/README.md @@ -2,7 +2,7 @@ ## Introduction -This directory contains benchmarking tools for linear attention operations (Gated DeltaNet and its variants) across various backends. The benchmarks target training use cases with support for forward and backward passes, grouped-value attention (GVA), and the per-sequence recurrent state ports (initial state in, final state out). +This directory contains benchmarking tools for linear attention operations (GDN/KDA/GDN-2) across various backends. The benchmarks target training use cases with support for forward and backward passes. ## Contents @@ -36,7 +36,7 @@ python benchmark_single_linear_attention.py \ --la_backend cudnn --variant kda --data_type bfloat16 \ --skip_ref --profile_pass bwd -# cuDNN Frontend (GDN-2, forward only) +# cuDNN Frontend (GDN-2, forward pass) python benchmark_single_linear_attention.py \ --batch_size 1 --seqlen 8192 \ --num_q_heads 16 --num_kv_heads 16 --head_dim 128 \ @@ -64,8 +64,14 @@ python benchmark_single_linear_attention.py \ --la_backend flash_qla --variant gdn --data_type bfloat16 \ --skip_ref --fwd_bwd -# Recurrent state ports: seed with an initial state and request the final -# state (its gradient feeds the backward pass) +# FlashKDA comparison point (kda variant only, forward only, bf16) +python benchmark_single_linear_attention.py \ + --batch_size 1 --seqlen 8192 \ + --num_q_heads 32 --num_kv_heads 32 --head_dim 128 \ + --la_backend flash_kda --variant kda --data_type bfloat16 \ + --skip_ref + +# Input initial state and dump state for every chunk python benchmark_single_linear_attention.py \ --batch_size 1 --seqlen 8192 \ --num_q_heads 8 --num_kv_heads 64 --head_dim 128 \ @@ -75,33 +81,15 @@ python benchmark_single_linear_attention.py \ Run `python benchmark_single_linear_attention.py --help` for all options. -Dropping `--skip_ref` validates the forward output against FLA (the same way the SDPA benchmark validates against FlashAttention 4). +The `kda` and `gdn2` variants fuse q/k L2 normalization in-kernel on every backend; `gdn` runs unfused. ## Supported Backends | Backend | Description | |---------|-------------| | `cudnn` | cuDNN (native, via the cuDNN Frontend torch custom ops) | -| `fla` | FLA (flash-linear-attention, Triton) | +| `fla` | FLA (flash-linear-attention, Triton; `gdn`, `kda`, and `gdn2`) | | `flash_qla` | FlashQLA (TileLang fused GDN kernels, `gdn` variant only) | +| `flash_kda` | FlashKDA (`kda` forward variant only) | -The cuDNN backend routes through the pygraph engines: FROST (Cutlass DSL) on SM100-class devices, the cuTile engines elsewhere. Both passes run through autograd, exactly like a training step. - -## Supported Variants - -| Variant | Description | -|---------|-------------| -| `gdn` | Gated DeltaNet: scalar per-token decay and write strength | -| `kda` | Kimi Delta Attention: per-key-channel decay | -| `gdn2` | Gated DeltaNet v2: channel-wise decay/erase/write gates (forward only, cuDNN only) | - -The benchmark runs `kda` and `gdn2` with the in-kernel q/k L2 normalization off (`use_qk_l2norm_in_kernel=False`) on every backend, for an apples-to-apples comparison. - -Recent `fla` releases dispatch `chunk_gated_delta_rule` to FlashQLA whenever `flash_qla` is importable; the benchmark sets `FLA_DISABLE_BACKEND_DISPATCH=1` (unless already set) so the `fla` backend always measures FLA's own Triton kernels and the two backends stay distinct. - -## Notes - -- Head convention: `--num_q_heads` counts the query/key heads and `--num_kv_heads` counts the value heads; the gates, output, and recurrent state live at `max(num_q_heads, num_kv_heads)` heads. Both grouping directions are supported for `gdn`: grouped-value attention (`num_kv_heads > num_q_heads`, v-heads grouped over q-heads) and GQA (`num_q_heads > num_kv_heads`, q-heads grouped over v-heads, e.g. `--num_q_heads 64 --num_kv_heads 8`). The two counts must be equal or one a multiple of the other; `kda` and `gdn2` support the GVA direction only, and so does the `flash_qla` backend. -- The cuDNN ops use the THD (token-packed) layout internally; the benchmark expresses the dense batch as `cu_seqlens = [0, T, 2T, ...]`. -- `--initial_state` provides a per-sequence fp32 recurrent state (its gradient is produced in the backward pass); `--store_on` requests the per-sequence final state from the forward pass and feeds its gradient in the backward pass. Both are once-per-kernel I/O ports (one `[head_dim_qk, head_dim_vo]` tile per sequence per state head). -- Performance is measured with the torch profiler (device time of the matched kernels), with a 256 MB L2 flush before each timed iteration and the median reported. TFLOPS use the chunked-BMM FLOPs model documented in the script's `flops()`. +The cuDNN backend routes through the pygraph engines: FROST (Cutlass DSL) on SM100-class devices, the cuTile engines elsewhere. \ No newline at end of file diff --git a/benchmark/linear_attention/benchmark_single_linear_attention.py b/benchmark/linear_attention/benchmark_single_linear_attention.py index bd64dd231..a95f49373 100644 --- a/benchmark/linear_attention/benchmark_single_linear_attention.py +++ b/benchmark/linear_attention/benchmark_single_linear_attention.py @@ -36,9 +36,12 @@ "float16": 8192, } -# Chunk size of the chunked linear attention algorithms (both backends tile -# the sequence into 64-token chunks; the FLOPs model below depends on it). -_CHUNK_SIZE = 64 +# Chunk size of the chunked linear attention algorithms per variant (the +# FLOPs model below depends on it). +_CHUNK_SIZE = {"gdn": 64, "kda": 16, "gdn2": 16} + +# Safe-gate lower bound for kda. +_KDA_GATE_LOWER_BOUND = -5.0 def _peak_flops_per_clock_per_sm(dtype_str): @@ -177,13 +180,13 @@ def parse_args(): "--variant", default="gdn", type=str, - help="Linear attention variant to use. Can be 'gdn' (scalar decay), 'kda' (per-key-channel decay), or 'gdn2' (channel-wise decay/erase/write gates, forward only).", + help="Linear attention variant to use. Can be 'gdn' (scalar decay), 'kda' (per-key-channel decay), or 'gdn2' (channel-wise decay/erase/write gates).", choices=["gdn", "kda", "gdn2"], ) parser.add_argument( "--store_on", action="store_true", - help="Request the per-sequence final recurrent state from the forward pass (and feed its gradient in the backward pass)", + help="Dump the recurrent state every chunk plus the per-sequence final state from the forward pass (backends without a per-chunk state output are rejected; with backward, the final state's gradient feeds the backward pass)", ) parser.add_argument( "--initial_state", @@ -198,6 +201,7 @@ def parse_args(): choices=[ "fla", "flash_qla", + "flash_kda", "cudnn", ], ) @@ -250,13 +254,13 @@ def run_benchmark( head_dim_qk: Head dimension for Q/K (optional, for asymmetric) head_dim_vo: Head dimension for V/O (optional, for asymmetric) data_type: Data type ("bfloat16", "float16") - backend: Backend name ("cudnn", "fla", "flash_qla") + backend: Backend name ("cudnn", "fla", "flash_qla", "flash_kda") variant: Linear attention variant ("gdn", "kda", "gdn2") profile_pass: Which pass to profile ("fwd", "bwd", "both") num_iterations: Number of benchmark iterations num_warmup_iterations: Warmup iterations before measurement skip_ref: Skip reference validation - store_on: Request the final recurrent state from the forward pass + store_on: Dump per-chunk states plus the final recurrent state from the forward pass initial_state: Provide an initial recurrent state verbose: Print verbose output @@ -264,6 +268,7 @@ def run_benchmark( Dict with keys: - time_ms: Median time of the requested pass in milliseconds - tflops: TFLOPS for the requested pass + - bw_tb_per_sec: DRAM bandwidth (TB/s) for the requested pass - max_diff: Maximum difference vs reference - gpu_name: GPU name string - cudnn_version: cuDNN version (if available) @@ -335,7 +340,7 @@ def run_benchmark( raise RuntimeError(f"Benchmark failed with return code {result.returncode}.\n" f"stderr: {result.stderr}\n" f"stdout: {result.stdout}") # Parse CSV output - # Format: case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_time,bwd_time,fwd_tflops,bwd_tflops,max_diff,num_iters + # Format: case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_time,bwd_time,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw output_line = result.stdout.strip().split("\n")[-1] parts = output_line.split(",") @@ -361,13 +366,16 @@ def run_benchmark( if profile_pass == "fwd": time_ms = float(parts[8]) tflops = float(parts[10]) + bw_tb_per_sec = float(parts[14]) if len(parts) > 14 else 0.0 else: # "bwd" time_ms = float(parts[9]) tflops = float(parts[11]) + bw_tb_per_sec = float(parts[15]) if len(parts) > 15 else 0.0 return { "time_ms": time_ms, "tflops": tflops, + "bw_tb_per_sec": bw_tb_per_sec, "max_diff": float(parts[12]) if len(parts) > 12 else 0.0, "gpu_name": gpu_name, "cudnn_version": cudnn_version, @@ -431,17 +439,43 @@ def run_benchmark( # The gates, output, and recurrent state live at HO = max(q, v) heads: # GVA groups v-heads over q-heads, GQA (gdn only) the reverse. num_o_heads = max(num_q_heads, num_kv_heads) - if num_q_heads > num_kv_heads and args.variant != "gdn": - raise ValueError("GQA (num_q_heads > num_kv_heads) is only supported with the 'gdn' variant") - if args.variant == "gdn2" and run_bwd: - raise ValueError("gdn2 is forward only (the backward kernel is a stub); use --profile_pass fwd") - if args.variant == "gdn2" and args.la_backend == "fla": - raise ValueError("gdn2 is only supported with the 'cudnn' backend") + if num_q_heads > num_kv_heads and args.variant != "gdn" and args.la_backend != "cudnn": + raise ValueError(f"GQA (num_q_heads > num_kv_heads) with kda/gdn2 is only supported by the cudnn backend, not {args.la_backend}") + if args.variant == "gdn2" and args.la_backend == "fla" and num_q_heads != num_kv_heads: + raise ValueError("gdn2 with the 'fla' backend requires equal q/kv head counts") if args.la_backend == "flash_qla": if args.variant != "gdn": raise ValueError("flash_qla only supports the 'gdn' variant") if num_q_heads > num_kv_heads: raise ValueError("flash_qla does not support GQA (num_q_heads > num_kv_heads)") + if args.la_backend == "flash_kda": + if args.variant != "kda": + raise ValueError("flash_kda only supports the 'kda' variant") + if num_q_heads != num_kv_heads: + raise ValueError("flash_kda requires equal q/kv head counts") + if head_dim_qk != 128 or head_dim_vo != 128: + raise ValueError("flash_kda requires head_dim 128 for both qk and vo") + if args.data_type != "bfloat16": + raise ValueError("flash_kda only supports bfloat16") + if run_bwd: + raise ValueError("flash_kda is forward only; use --profile_pass fwd") + + if args.store_on: + if args.la_backend in ("flash_kda", "flash_qla"): + raise ValueError(f"--store_on dumps the state every chunk; {args.la_backend} only outputs the final state") + if args.la_backend == "fla" and args.variant == "gdn": + raise ValueError("--store_on dumps the state every chunk; fla's chunk_gated_delta_rule has no intermediate-state output") + if args.la_backend == "fla" and run_bwd: + raise ValueError("--store_on dumps the state every chunk; fla dumps intermediate states in inference mode only, use --profile_pass fwd") + + # FlashKDA always fuses q/k l2norm in-kernel, so kda/gdn2 fuse it on + # every backend for apples-to-apples comparisons; gdn runs unfused. + use_qk_l2norm = args.variant in ("kda", "gdn2") + + # Forward-only kda runs feed raw safe-gate/beta logits to every backend + # (the FlashKDA ABI); training runs keep post-activation gates (the + # cudnn/fla backwards then differentiate the same math). + kda_raw_gates = args.variant == "kda" and not run_bwd l2_flush_size_mb = 256 l2_flush_size = l2_flush_size_mb * 1024 * 1024 @@ -472,9 +506,22 @@ def run_benchmark( ## boundaries. A dense batch is cu_seqlens = [0, T, 2T, ...]. cu_seqlens = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * seqlen - def cudnn_linear_attention(query, key, value, gate, beta, write_gate, s0): + ## --store_on dumps the per-chunk state series alongside the final + ## state (the state_checkpoints output; one entry per kernel chunk). + ckpt_tokens = _CHUNK_SIZE[args.variant] if args.store_on else 0 + kda_safe_gate_kwargs = {} + if kda_raw_gates: + kda_safe_gate_kwargs = dict( + use_beta_sigmoid_in_kernel=True, + safe_gate=True, + gate_lower_bound=_KDA_GATE_LOWER_BOUND, + a_log=torch.zeros(num_o_heads, dtype=torch.float32, device=device), + dt_bias=torch.zeros(num_o_heads, head_dim_qk, dtype=torch.float32, device=device), + ) + + def cudnn_linear_attention(query, key, value, gate, beta, write_gate, state0): if args.variant == "gdn": - return gated_delta_net( + out = gated_delta_net( query, key, value, @@ -482,11 +529,13 @@ def cudnn_linear_attention(query, key, value, gate, beta, write_gate, s0): beta, cu_seqlens, scale=attn_scale, - initial_state=s0, + initial_state=state0, output_final_state=args.store_on, + use_qk_l2norm_in_kernel=use_qk_l2norm, + checkpoint_every_n_tokens=ckpt_tokens, ) elif args.variant == "kda": - return kimi_delta_attention( + out = kimi_delta_attention( query, key, value, @@ -494,12 +543,14 @@ def cudnn_linear_attention(query, key, value, gate, beta, write_gate, s0): beta, cu_seqlens, scale=attn_scale, - initial_state=s0, + initial_state=state0, output_final_state=args.store_on, - use_qk_l2norm_in_kernel=False, + use_qk_l2norm_in_kernel=use_qk_l2norm, + checkpoint_every_n_tokens=ckpt_tokens, + **kda_safe_gate_kwargs, ) else: # gdn2 - return gated_delta_net_v2( + out = gated_delta_net_v2( query, key, value, @@ -508,10 +559,12 @@ def cudnn_linear_attention(query, key, value, gate, beta, write_gate, s0): write_gate, cu_seqlens, scale=attn_scale, - initial_state=s0, + initial_state=state0, output_final_state=args.store_on, - use_qk_l2norm_in_kernel=False, + use_qk_l2norm_in_kernel=use_qk_l2norm, + checkpoint_every_n_tokens=ckpt_tokens, ) + return out[0], out[1] if args.la_backend == "flash_qla": attn_scale = head_dim_qk ** (-0.5) @@ -523,7 +576,7 @@ def cudnn_linear_attention(query, key, value, gate, beta, write_gate, s0): print(f"[INFO] FlashQLA Version: {getattr(flash_qla, '__version__', 'unknown')}") - def flash_qla_linear_attention(query, key, value, gate, beta, write_gate, s0): + def flash_qla_linear_attention(query, key, value, gate, beta, write_gate, state0): return fqla_chunk_gated_delta_rule( query, key, @@ -531,10 +584,39 @@ def flash_qla_linear_attention(query, key, value, gate, beta, write_gate, s0): gate, beta, scale=attn_scale, - initial_state=s0, + initial_state=state0, output_final_state=args.store_on, ) + if args.la_backend == "flash_kda": + attn_scale = head_dim_qk ** (-0.5) + + import flash_kda + + if args.verbose: + print(f"[INFO] FlashKDA Version: {getattr(flash_kda, '__version__', 'unknown')}") + + fkda_a_log = torch.zeros(num_o_heads, dtype=torch.float32, device=device) + fkda_dt_bias = torch.zeros(num_o_heads, head_dim_qk, dtype=torch.float32, device=device) + + def flash_kda_linear_attention(query, key, value, gate, beta, write_gate, state0): + out = torch.empty_like(value) + flash_kda.fwd( + query, + key, + value, + gate, + beta, + attn_scale, + out, + A_log=fkda_a_log, + dt_bias=fkda_dt_bias, + lower_bound=_KDA_GATE_LOWER_BOUND, + initial_state=state0, + final_state=None, + ) + return out, None + if args.la_backend == "fla" or (not args.skip_ref): attn_scale = head_dim_qk ** (-0.5) @@ -546,12 +628,28 @@ def flash_qla_linear_attention(query, key, value, gate, beta, write_gate, s0): try: from fla.ops.kda import chunk_kda except ImportError as e: - raise RuntimeError(f"The installed fla does not provide KDA (fla.ops.kda): {e}") + raise RuntimeError(f"The installed fla does not provide KDA (fla.ops.kda): {e}") from e else: # gdn2 - raise ValueError("gdn2 is only supported with the 'cudnn' backend (no fla implementation, so no reference either); use --skip_ref") + try: + from fla.ops.gdn2 import chunk_gdn2 + except ImportError as e: + raise RuntimeError(f"The installed fla does not provide GDN2 (fla.ops.gdn2): {e}") from e + + ## FLA takes dense (B, T, H, D) tensors; g is the log-space decay + ## (raw safe-gate logits in forward-only kda runs). --store_on adds + ## the per-chunk state dump (fla supports it in inference mode only). + fla_kda_kwargs = {} + if args.variant == "kda" and kda_raw_gates: + fla_kda_kwargs = dict( + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=_KDA_GATE_LOWER_BOUND, + use_beta_sigmoid_in_kernel=True, + A_log=torch.zeros(num_o_heads, dtype=torch.float32, device=device), + dt_bias=torch.zeros(num_o_heads * head_dim_qk, dtype=torch.float32, device=device), + ) - ## FLA takes dense (B, T, H, D) tensors; g is the log-space decay. - def fla_linear_attention(query, key, value, gate, beta, write_gate, s0): + def fla_linear_attention(query, key, value, gate, beta, write_gate, state0): if args.variant == "gdn": return chunk_gated_delta_rule( query, @@ -560,10 +658,27 @@ def fla_linear_attention(query, key, value, gate, beta, write_gate, s0): gate, beta, scale=attn_scale, - initial_state=s0, + initial_state=state0, output_final_state=args.store_on, + use_qk_l2norm_in_kernel=use_qk_l2norm, ) - else: # kda + elif args.variant == "kda": + if args.store_on: + with torch.inference_mode(): + o, fs, _state_checkpoints = chunk_kda( + query, + key, + value, + gate, + beta, + scale=attn_scale, + initial_state=state0, + output_final_state=True, + use_qk_l2norm_in_kernel=use_qk_l2norm, + return_intermediate_states=True, + **fla_kda_kwargs, + ) + return o, fs return chunk_kda( query, key, @@ -571,9 +686,39 @@ def fla_linear_attention(query, key, value, gate, beta, write_gate, s0): gate, beta, scale=attn_scale, - initial_state=s0, + initial_state=state0, output_final_state=args.store_on, - use_qk_l2norm_in_kernel=False, + use_qk_l2norm_in_kernel=use_qk_l2norm, + **fla_kda_kwargs, + ) + else: # gdn2 + if args.store_on: + with torch.inference_mode(): + o, fs, _state_checkpoints = chunk_gdn2( + query, + key, + value, + gate, + beta, + write_gate, + scale=attn_scale, + initial_state=state0, + output_final_state=True, + use_qk_l2norm_in_kernel=use_qk_l2norm, + return_intermediate_states=True, + ) + return o, fs + return chunk_gdn2( + query, + key, + value, + gate, + beta, + write_gate, + scale=attn_scale, + initial_state=state0, + output_final_state=args.store_on, + use_qk_l2norm_in_kernel=use_qk_l2norm, ) def get_linear_attention_function(backend): @@ -581,6 +726,8 @@ def get_linear_attention_function(backend): return fla_linear_attention elif backend == "flash_qla": return flash_qla_linear_attention + elif backend == "flash_kda": + return flash_kda_linear_attention elif backend == "cudnn": return cudnn_linear_attention else: @@ -595,20 +742,25 @@ def preprocess_qkv(query, key, value, backend): key.reshape(batch_size * seqlen, *key.shape[2:]), value.reshape(batch_size * seqlen, *value.shape[2:]), ) - elif backend in ("fla", "flash_qla"): + elif backend in ("fla", "flash_qla", "flash_kda"): return query, key, value else: raise ValueError(f"Invalid backend: {backend}") def preprocess_gates(gate, beta, write_gate, backend): if backend == "cudnn": + beta_t = beta.reshape(batch_size * seqlen, *beta.shape[2:]) + if kda_raw_gates: + beta_t = beta_t.to(target_dtype) return ( gate.reshape(batch_size * seqlen, *gate.shape[2:]), - beta.reshape(batch_size * seqlen, *beta.shape[2:]), + beta_t, write_gate.reshape(batch_size * seqlen, *write_gate.shape[2:]) if write_gate is not None else None, ) elif backend in ("fla", "flash_qla"): return gate, beta, write_gate + elif backend == "flash_kda": + return gate.to(target_dtype).contiguous(), beta.to(target_dtype).contiguous(), None else: raise ValueError(f"Invalid backend: {backend}") @@ -616,7 +768,7 @@ def preprocess_gates(gate, beta, write_gate, backend): def postprocess_o(output, backend): if backend == "cudnn": return output.reshape(batch_size, seqlen, num_o_heads, head_dim_vo) - elif backend in ("fla", "flash_qla"): + elif backend in ("fla", "flash_qla", "flash_kda"): return output else: raise ValueError(f"Invalid backend: {backend}") @@ -632,13 +784,13 @@ def flops( ): assert mode in ["fwd", "bwd", "fwd_bwd"] - # Chunked linear attention BMM FLOPs per 64-token chunk per (batch, + # Chunked linear attention BMM FLOPs per chunk per (batch, # state head), chunk size C, dims K (qk) and V (vo): # Forward: 5 BMM classes => # intra scores + WY prep (2 x C*C*K), WY apply (C*C*K + C*C*V), # intra output (C*C*V), inter output + state update (2 x C*K*V) # Backward: recompute + gradient chains, ~3x forward. - C = _CHUNK_SIZE + C = _CHUNK_SIZE[args.variant] num_chunks = ceil_div(seqlen, C) per_chunk = 2 * (3 * C * C * head_dim_qk + 2 * C * C * head_dim_vo + 2 * C * head_dim_qk * head_dim_vo) base = batch_size * num_kv_heads * num_chunks * per_chunk @@ -670,6 +822,63 @@ def tflops_per_sec( ) return f / time / 1e9 if not math.isnan(time) else 0.0 # Assume time is in msec + # Util functions for calculating DRAM bytes moved and bandwidth achieved + def dram_bytes( + batch_size, + seqlen, + head_dim_qk, + head_dim_vo, + num_q_heads, + num_kv_heads, + num_o_heads, + mode="fwd", + ): + assert mode in ["fwd", "bwd"] + io_bytes = 2 + f32_bytes = 4 + tokens = batch_size * seqlen + q_bytes = tokens * num_q_heads * head_dim_qk * io_bytes + k_bytes = tokens * num_q_heads * head_dim_qk * io_bytes + v_bytes = tokens * num_kv_heads * head_dim_vo * io_bytes + o_bytes = tokens * num_o_heads * head_dim_vo * io_bytes + if args.variant == "gdn": + gate_bytes = tokens * num_o_heads * 2 * f32_bytes + elif args.variant == "kda": + gate_bytes = tokens * num_o_heads * (head_dim_qk + 1) * f32_bytes + else: # gdn2 + gate_bytes = tokens * num_o_heads * (head_dim_qk * f32_bytes + (head_dim_qk + head_dim_vo) * io_bytes) + num_chunks = ceil_div(seqlen, _CHUNK_SIZE[args.variant]) + h_bytes = batch_size * num_o_heads * num_chunks * head_dim_qk * head_dim_vo * io_bytes + qkv_bytes = q_bytes + k_bytes + v_bytes + if mode == "fwd": + return qkv_bytes + gate_bytes + o_bytes + recompute_bytes = k_bytes + v_bytes + gate_bytes + h_bytes + return recompute_bytes + 2 * (qkv_bytes + gate_bytes) + o_bytes + h_bytes + + def tb_per_sec( + batch_size, + seqlen, + head_dim_qk, + head_dim_vo, + num_q_heads, + num_kv_heads, + num_o_heads, + time, + mode="fwd", + ): + assert mode in ["fwd", "bwd"] + b = dram_bytes( + batch_size, + seqlen, + head_dim_qk, + head_dim_vo, + num_q_heads, + num_kv_heads, + num_o_heads, + mode, + ) + return b / time / 1e9 if not math.isnan(time) else 0.0 # Assume time is in msec + ## Gate generators per variant. Decays are LOG-space (alpha = exp(g)), ## drawn from ranges the kernels' io-dtype arithmetic is conditioned for. def generate_gates(io_dtype): @@ -679,9 +888,15 @@ def generate_gates(io_dtype): beta = torch.rand(batch_size, seqlen, num_o_heads, device=device) write_gate = None elif args.variant == "kda": - # per-key-channel decay [B, T, HO, K] fp32 + post-sigmoid scalar beta + # per-key-channel decay [B, T, HO, K] fp32 + post-sigmoid scalar + # beta; forward-only runs feed raw logits with the same effective + # distributions (the in-kernel activations invert them) gate = torch.empty(batch_size, seqlen, num_o_heads, head_dim_qk, device=device).uniform_(0.5, 1.0).log() - beta = torch.rand(batch_size, seqlen, num_o_heads, device=device).sigmoid() + beta = torch.rand(batch_size, seqlen, num_o_heads, device=device) + if kda_raw_gates: + gate = torch.special.logit((gate / _KDA_GATE_LOWER_BOUND).clamp(1e-7, 1 - 1e-7)) + else: + beta = beta.sigmoid() write_gate = None else: # gdn2 # per-key decay/erase [B, T, HO, K] + per-value write gate [B, T, HO, V] @@ -732,15 +947,20 @@ def generate_gates(io_dtype): value.requires_grad_(True) gate.requires_grad_(True) beta.requires_grad_(True) + if write_gate is not None: + write_gate.requires_grad_(True) - # Per-sequence recurrent state ports (once-per-kernel I/O): the - # initial state seeds the recurrence; the final state is requested - # with --store_on and its gradient feeds the backward pass. - s0 = None + # Recurrent state ports: the initial state seeds the recurrence; + # --store_on dumps the per-chunk states plus the final state (whose + # gradient feeds the backward pass). + state0 = None if args.initial_state: - s0 = torch.randn(batch_size, num_o_heads, head_dim_qk, head_dim_vo, dtype=torch.float32, device=device) * 0.05 + state0 = torch.randn(batch_size, num_o_heads, head_dim_qk, head_dim_vo, dtype=torch.float32, device=device) * 0.05 + if args.la_backend == "flash_kda": + # FlashKDA state ports are V-major [B, H, V, K] + state0 = state0.transpose(-1, -2).contiguous() if run_bwd: - s0.requires_grad_(True) + state0.requires_grad_(True) if args.la_backend == "cudnn": dOutput = torch.randn(batch_size * seqlen, num_o_heads, head_dim_vo, dtype=target_dtype, device=device) else: @@ -755,7 +975,7 @@ def generate_gates(io_dtype): if run_fwd: with profile(activities=[ProfilerActivity.CUDA], record_shapes=True) as prof: with record_function("linear_attention.forward"): # Custom marker - output, final_state = la_function(query, key, value, gate, beta, write_gate, s0) + output, final_state = la_function(query, key, value, gate, beta, write_gate, state0) torch.cuda.synchronize() # Ensure all kernels finish # Filter profiler results by kernel name prefix @@ -766,6 +986,7 @@ def generate_gates(io_dtype): or item.key.startswith("kernel_cutlass") or item.key.startswith("triton_") or "chunk_" in item.key + or "_flash_kda_" in item.key or "l2norm" in item.key or "cutile" in item.key or "_kernel" in item.key @@ -777,7 +998,7 @@ def generate_gates(io_dtype): if i >= dry_run_iters: forward_times.append(fwd_time) else: - output, final_state = la_function(query, key, value, gate, beta, write_gate, s0) + output, final_state = la_function(query, key, value, gate, beta, write_gate, state0) torch.cuda.synchronize() if run_bwd: @@ -815,6 +1036,7 @@ def generate_gates(io_dtype): or item.key.startswith("kernel_cutlass") or item.key.startswith("triton_") or "chunk_" in item.key + or "_flash_kda_" in item.key or "l2norm" in item.key or "cutile" in item.key or "_kernel" in item.key @@ -835,11 +1057,23 @@ def generate_gates(io_dtype): if args.la_backend == "cudnn": gate_ref = gate.detach().reshape(batch_size, seqlen, *gate.shape[1:]) beta_ref = beta.detach().reshape(batch_size, seqlen, *beta.shape[1:]) + elif args.la_backend == "flash_kda": + # raw logits pass through; the fla reference applies the + # same in-kernel activations + gate_ref = gate.detach().float() + beta_ref = beta.detach().float() else: gate_ref = gate.detach() beta_ref = beta.detach() - s0_ref = s0.detach() if s0 is not None else None - output_ref, _ = fla_linear_attention(query_ref, key_ref, value_ref, gate_ref, beta_ref, None, s0_ref) + state0_ref = state0.detach() if state0 is not None else None + if state0_ref is not None and args.la_backend == "flash_kda": + state0_ref = state0_ref.transpose(-1, -2).contiguous() + wg_ref = None + if write_gate is not None: + wg_ref = write_gate.detach() + if args.la_backend == "cudnn": + wg_ref = wg_ref.reshape(batch_size, seqlen, *write_gate.shape[1:]) + output_ref, _ = fla_linear_attention(query_ref, key_ref, value_ref, gate_ref, beta_ref, wg_ref, state0_ref) torch.testing.assert_close(output.detach(), output_ref, rtol=1e-2, atol=1e-2) forward_diffs.append(torch.max(torch.abs(output.detach() - output_ref.detach())).item()) @@ -854,7 +1088,7 @@ def generate_gates(io_dtype): else: forward_diffs.append(0.0) - del query, key, value, gate, beta, write_gate, output, final_state, s0, dOutput, dFinal + del query, key, value, gate, beta, write_gate, output, final_state, state0, dOutput, dFinal _clock_sampler.stop() @@ -863,6 +1097,7 @@ def generate_gates(io_dtype): np.median(np.array(forward_times[5:])) if len(forward_times) > 5 else (np.median(np.array(forward_times)) if len(forward_times) > 0 else 0.0) ) fwd_tflops = 0.0 + fwd_bw = 0.0 if run_fwd and fwd_median_time > 0: fwd_tflops = tflops_per_sec( args.batch_size, @@ -873,11 +1108,23 @@ def generate_gates(io_dtype): fwd_median_time, "fwd", ) + fwd_bw = tb_per_sec( + args.batch_size, + args.seqlen, + head_dim_qk, + head_dim_vo, + args.num_q_heads, + args.num_kv_heads, + num_o_heads, + fwd_median_time, + "fwd", + ) bwd_median_time = ( np.median(np.array(backward_times[5:])) if len(backward_times) > 5 else (np.median(np.array(backward_times)) if len(backward_times) > 0 else 0.0) ) bwd_tflops = 0.0 + bwd_bw = 0.0 if run_bwd and bwd_median_time > 0: bwd_tflops = tflops_per_sec( args.batch_size, @@ -888,6 +1135,17 @@ def generate_gates(io_dtype): bwd_median_time, "bwd", ) + bwd_bw = tb_per_sec( + args.batch_size, + args.seqlen, + head_dim_qk, + head_dim_vo, + args.num_q_heads, + args.num_kv_heads, + num_o_heads, + bwd_median_time, + "bwd", + ) # Compute MMA SOL% using the per-arch FLOPs/clk/SM table and the actual # sampled boost clock observed during the benchmark window. @@ -904,16 +1162,21 @@ def generate_gates(io_dtype): fwd_sol_str = f", {fwd_tflops / _peak_mma_tflops * 100:.1f}% SOL" if _peak_mma_tflops and fwd_tflops > 0 else "" bwd_sol_str = f", {bwd_tflops / _peak_mma_tflops * 100:.1f}% SOL" if _peak_mma_tflops and bwd_tflops > 0 else "" + backend_tag = f"{args.la_backend}_state_on" if (args.store_on and args.la_backend == "cudnn") else args.la_backend if args.format_output: print( - f"{args.case_tag},{args.la_backend},{args.variant},{args.batch_size},{args.seqlen},{args.num_q_heads},{args.num_kv_heads},{head_dim_qk},{fwd_median_time:.3f},{bwd_median_time:.3f},{fwd_tflops:.0f},{bwd_tflops:.0f},{(np.max(np.array(forward_diffs[5:])) if len(forward_diffs) > 5 else (np.max(np.array(forward_diffs)) if len(forward_diffs) > 0 else 0.0)):.6f},{num_iters}" + f"{args.case_tag},{backend_tag},{args.variant},{args.batch_size},{args.seqlen},{args.num_q_heads},{args.num_kv_heads},{head_dim_qk},{fwd_median_time:.3f},{bwd_median_time:.3f},{fwd_tflops:.0f},{bwd_tflops:.0f},{(np.max(np.array(forward_diffs[5:])) if len(forward_diffs) > 5 else (np.max(np.array(forward_diffs)) if len(forward_diffs) > 0 else 0.0)):.6f},{num_iters},{fwd_bw:.2f},{bwd_bw:.2f}" ) else: if run_fwd and run_bwd: print( - f"{args.la_backend}/{args.variant}:: Median (fwd, bwd) Execution Times: {fwd_median_time:.3f} ms ({fwd_tflops:.0f} TFLOPS{fwd_sol_str}), {bwd_median_time:.3f} ms ({bwd_tflops:.0f} TFLOPS{bwd_sol_str})" + f"{args.la_backend}/{args.variant}:: Median (fwd, bwd) Execution Times: {fwd_median_time:.3f} ms ({fwd_tflops:.0f} TFLOPS{fwd_sol_str}, {fwd_bw:.2f} TB/s), {bwd_median_time:.3f} ms ({bwd_tflops:.0f} TFLOPS{bwd_sol_str}, {bwd_bw:.2f} TB/s)" ) elif run_fwd: - print(f"{args.la_backend}/{args.variant}:: Median (fwd) Execution Time: {fwd_median_time:.3f} ms ({fwd_tflops:.0f} TFLOPS{fwd_sol_str})") + print( + f"{args.la_backend}/{args.variant}:: Median (fwd) Execution Time: {fwd_median_time:.3f} ms ({fwd_tflops:.0f} TFLOPS{fwd_sol_str}, {fwd_bw:.2f} TB/s)" + ) elif run_bwd: - print(f"{args.la_backend}/{args.variant}:: Median (bwd) Execution Time: {bwd_median_time:.3f} ms ({bwd_tflops:.0f} TFLOPS{bwd_sol_str})") + print( + f"{args.la_backend}/{args.variant}:: Median (bwd) Execution Time: {bwd_median_time:.3f} ms ({bwd_tflops:.0f} TFLOPS{bwd_sol_str}, {bwd_bw:.2f} TB/s)" + ) diff --git a/benchmark/linear_attention/plot_results.py b/benchmark/linear_attention/plot_results.py index 654a93e90..e6254d489 100644 --- a/benchmark/linear_attention/plot_results.py +++ b/benchmark/linear_attention/plot_results.py @@ -9,12 +9,13 @@ size, with Forward and Backward TFLOPS panels side by side — same style as the SDPA training benchmark charts. - python plot_results.py results/gdn/b300/gdn_labench.csv \ - --output-dir results/gdn/b300 --gpu-name B300 --cudnn-version 9.24.0 + python plot_results.py results/gdn/gb300/gdn_20260813.csv \ + --output-dir results/gdn/gb300 --gpu-name GB300 --cudnn-version 9.24.0 """ import argparse from pathlib import Path +from typing import List, Optional import pandas as pd import matplotlib @@ -27,9 +28,14 @@ BACKEND_CONFIG = { "fla": {"name": "FLA (Triton)", "color": "#FF8C00", "order": 0}, "flash_qla": {"name": "FlashQLA (TileLang)", "color": "#6495ED", "order": 1}, - "cudnn": {"name": "cuDNN", "color": "#76b900", "order": 2}, + "flash_kda": {"name": "FlashKDA", "color": "#9370DB", "order": 2}, + "cudnn": {"name": "cuDNN (default)", "color": "#76b900", "order": 3}, + "cudnn_state_on": {"name": "cuDNN (state on)", "color": "#2f6e00", "order": 4}, } +# Backends dropped from every chart (rows may still exist in older CSVs). +UNAVAILABLE_BACKENDS = () + LABEL_FONT_SIZE = 10 LEGEND_FONT_SIZE = 8 TITLE_FONT_SIZE = 12 @@ -50,20 +56,35 @@ "bwd_tflops", "max_diff", "num_iters", + "fwd_bw", + "bwd_bw", ] +# One chart per metric: (fwd column, bwd column, y-axis label, bar label +# format, filename suffix). +METRIC_CONFIG = ( + ("fwd_tflops", "bwd_tflops", "TFLOPS", "%.0f", "_flops"), + ("fwd_bw", "bwd_bw", "DRAM Bandwidth (TB/s)", "%.2f", "_bw"), +) + -def get_backend_display_name(backend: str, cudnn_version: str = None) -> str: - base_name = BACKEND_CONFIG.get(backend, {}).get("name", backend) - if backend == "cudnn" and cudnn_version: - base_name = f"{base_name} {cudnn_version}" - return base_name +def get_backend_display_name(backend: str, cudnn_version: Optional[str] = None) -> str: + return BACKEND_CONFIG.get(backend, {}).get("name", backend) -def generate_charts(df: pd.DataFrame, output_dir: Path, gpu_name: str = "", cudnn_version: str = None, variant: str = "gdn", batch_sizes: list = None) -> list: +def generate_charts( + df: pd.DataFrame, + output_dir: Path, + gpu_name: str = "", + cudnn_version: Optional[str] = None, + variant: str = "gdn", + batch_sizes: Optional[List[int]] = None, + x_axis: str = "seqlen", +) -> list: output_dir.mkdir(parents=True, exist_ok=True) df = df[df["variant"] == variant].copy() + df = df[~df["backend"].isin(UNAVAILABLE_BACKENDS)].copy() if batch_sizes: df = df[df["batch_size"].isin(batch_sizes)].copy() if df.empty: @@ -76,66 +97,82 @@ def generate_charts(df: pd.DataFrame, output_dir: Path, gpu_name: str = "", cudn for _, row in df[["backend", "backend_display"]].drop_duplicates().iterrows(): palette[row["backend_display"]] = BACKEND_CONFIG.get(row["backend"], {}).get("color", "gray") + x_col, group_col = ("seqlen", "batch_size") if x_axis == "seqlen" else ("batch_size", "seqlen") + x_label = "Sequence Length" if x_axis == "seqlen" else "Batch Size" + saved_paths = [] - for batch_size in sorted(df["batch_size"].unique()): - sub = df[df["batch_size"] == batch_size].copy() - sub.sort_values(["seqlen", "backend_order"], inplace=True) - hue_order = list(sub.sort_values("backend_order")["backend_display"].drop_duplicates()) - - fwd_df = sub[sub["fwd_tflops"] > 0] - bwd_df = sub[sub["bwd_tflops"] > 0] - has_fwd = not fwd_df.empty - has_bwd = not bwd_df.empty - - if has_fwd and has_bwd: - fig, (ax_fwd, ax_bwd) = plt.subplots(1, 2, figsize=(14, 6), dpi=150) - elif has_fwd: - fig, ax_fwd = plt.subplots(1, 1, figsize=(10, 6), dpi=150) - ax_bwd = None - else: - fig, ax_bwd = plt.subplots(1, 1, figsize=(10, 6), dpi=150) - ax_fwd = None - - heads = sub["num_q_heads"].iloc[0] - head_dim = sub["head_dim"].iloc[0] - gpu_info = f" ({gpu_name})" if gpu_name else "" - fig.suptitle( - f"{variant.upper()} Linear Attention (BF16) — Batch = {batch_size}, Heads = {heads}, d = {head_dim}{gpu_info}", - fontsize=TITLE_FONT_SIZE, - ) - - for ax, pass_df, pass_name, y_col in ( - (ax_fwd, fwd_df, "Forward", "fwd_tflops"), - (ax_bwd, bwd_df, "Backward", "bwd_tflops"), - ): - if ax is None or pass_df.empty: + for group_val in sorted(df[group_col].unique()): + sub = df[df[group_col] == group_val].copy() + sub.sort_values([x_col, "backend_order"], inplace=True) + + for fwd_col, bwd_col, y_label, bar_fmt, file_suffix in METRIC_CONFIG: + if fwd_col not in sub.columns or bwd_col not in sub.columns: continue - sns.barplot( - data=pass_df, - x="seqlen", - y=y_col, - hue="backend_display", - hue_order=hue_order, - ax=ax, - palette=palette, - edgecolor="black", - linewidth=0.5, - errorbar=None, + fwd_df = sub[sub[fwd_col] > 0] + bwd_df = sub[sub[bwd_col] > 0] + has_fwd = not fwd_df.empty + has_bwd = not bwd_df.empty + if not has_fwd and not has_bwd: + continue + + if has_fwd and has_bwd: + fig, (ax_fwd, ax_bwd) = plt.subplots(1, 2, figsize=(14, 6), dpi=150) + elif has_fwd: + fig, ax_fwd = plt.subplots(1, 1, figsize=(10, 6), dpi=150) + ax_bwd = None + else: + fig, ax_bwd = plt.subplots(1, 1, figsize=(10, 6), dpi=150) + ax_fwd = None + + heads = sub["num_q_heads"].iloc[0] + head_dim = sub["head_dim"].iloc[0] + gpu_info = f" ({gpu_name})" if gpu_name else "" + group_label = f"Batch = {group_val}" if x_axis == "seqlen" else f"Sequence Length = {group_val}" + fig.suptitle( + f"{variant.upper()} Linear Attention (BF16) — {group_label}, Heads = {heads}, d = {head_dim}{gpu_info}", + fontsize=TITLE_FONT_SIZE, ) - ax.set_xlabel("Sequence Length", fontsize=LABEL_FONT_SIZE) - ax.set_ylabel("TFLOPS", fontsize=LABEL_FONT_SIZE) - ax.set_title(pass_name, fontsize=TITLE_FONT_SIZE) - ax.legend(title="Backend", fontsize=LEGEND_FONT_SIZE) - ax.tick_params(axis="x", rotation=45) - for container in ax.containers: - ax.bar_label(container, fmt="%.0f", fontsize=BAR_LABEL_FONT_SIZE) - - plt.tight_layout() - output_path = output_dir / f"{variant}_b{batch_size}.png" - plt.savefig(output_path, dpi=150, bbox_inches="tight") - plt.close() - saved_paths.append(output_path) - print(f"Chart saved to {output_path}") + + for ax, pass_df, pass_name, y_col in ( + (ax_fwd, fwd_df, "Forward", fwd_col), + (ax_bwd, bwd_df, "Backward", bwd_col), + ): + if ax is None or pass_df.empty: + continue + hue_order = list(pass_df.sort_values("backend_order")["backend_display"].drop_duplicates()) + sns.barplot( + data=pass_df, + x=x_col, + y=y_col, + hue="backend_display", + hue_order=hue_order, + ax=ax, + palette=palette, + edgecolor="black", + linewidth=0.5, + errorbar=None, + ) + ax.set_xlabel(x_label, fontsize=LABEL_FONT_SIZE) + ax.set_ylabel(y_label, fontsize=LABEL_FONT_SIZE) + ax.set_title(pass_name, fontsize=TITLE_FONT_SIZE) + ax.legend(title="Backend", fontsize=LEGEND_FONT_SIZE, loc="upper left") + ax.tick_params(axis="x", rotation=45) + for container in ax.containers: + ax.bar_label(container, fmt=bar_fmt, fontsize=BAR_LABEL_FONT_SIZE) + + plt.tight_layout() + gv = int(group_val) + if df[group_col].nunique() == 1: + # the sweep pinned the group dimension: fixed-batch (seqlen + # sweep) / fixed-seq (batch sweep) result-tree naming + stem = f"{variant}_fixed_batch" if x_axis == "seqlen" else f"{variant}_fixed_seq" + else: + stem = f"{variant}_b{gv}" if x_axis == "seqlen" else f"{variant}_t{gv}_bsweep" + output_path = output_dir / f"{stem}{file_suffix}.png" + plt.savefig(output_path, dpi=150, bbox_inches="tight") + plt.close() + saved_paths.append(output_path) + print(f"Chart saved to {output_path}") return saved_paths @@ -148,16 +185,20 @@ def main(): parser.add_argument("--cudnn-version", default=None, help="cuDNN backend version for the legend (e.g. 9.24.0)") parser.add_argument("--variant", default="gdn", help="Linear attention variant to plot") parser.add_argument("--batch-sizes", default=None, help="Comma-separated batch sizes to plot (default: all in the CSV)") + parser.add_argument( + "--x-axis", default="seqlen", choices=("seqlen", "batch"), help="Bar-group axis: seqlen (one chart per batch) or batch (one chart per seqlen)" + ) args = parser.parse_args() batch_sizes = [int(b) for b in args.batch_sizes.split(",")] if args.batch_sizes else None df = pd.read_csv(args.csv) - missing = [c for c in CSV_COLUMNS if c not in df.columns] + # fwd_bw/bwd_bw are newer columns; older CSVs simply skip the BW charts. + missing = [c for c in CSV_COLUMNS if c not in df.columns and c not in ("fwd_bw", "bwd_bw")] if missing: raise ValueError(f"CSV is missing expected columns: {missing}") output_dir = args.output_dir if args.output_dir is not None else args.csv.parent - generate_charts(df, output_dir, gpu_name=args.gpu_name, cudnn_version=args.cudnn_version, variant=args.variant, batch_sizes=batch_sizes) + generate_charts(df, output_dir, gpu_name=args.gpu_name, cudnn_version=args.cudnn_version, variant=args.variant, batch_sizes=batch_sizes, x_axis=args.x_axis) if __name__ == "__main__": diff --git a/benchmark/linear_attention/results/gdn/b300/gdn.png b/benchmark/linear_attention/results/gdn/b300/gdn.png deleted file mode 100644 index 43c09ec5b..000000000 Binary files a/benchmark/linear_attention/results/gdn/b300/gdn.png and /dev/null differ diff --git a/benchmark/linear_attention/results/gdn/b300/gdn_20260806.csv b/benchmark/linear_attention/results/gdn/b300/gdn_20260806.csv deleted file mode 100644 index 5e0625982..000000000 --- a/benchmark/linear_attention/results/gdn/b300/gdn_20260806.csv +++ /dev/null @@ -1,16 +0,0 @@ -case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters -fla_B1_T2048,fla,gdn,1,2048,64,64,128,0.217,0.649,89,89,0.000000,20 -flash_qla_B1_T2048,flash_qla,gdn,1,2048,64,64,128,0.102,0.352,190,165,0.000000,20 -cudnn_B1_T2048,cudnn,gdn,1,2048,64,64,128,0.098,0.279,196,208,0.000000,20 -fla_B1_T4096,fla,gdn,1,4096,64,64,128,0.416,1.260,93,92,0.000000,20 -flash_qla_B1_T4096,flash_qla,gdn,1,4096,64,64,128,0.186,0.677,208,171,0.000000,20 -cudnn_B1_T4096,cudnn,gdn,1,4096,64,64,128,0.129,0.428,301,271,0.000000,20 -fla_B1_T8192,fla,gdn,1,8192,64,64,128,0.796,2.490,97,93,0.000000,20 -flash_qla_B1_T8192,flash_qla,gdn,1,8192,64,64,128,0.347,1.342,223,173,0.000000,20 -cudnn_B1_T8192,cudnn,gdn,1,8192,64,64,128,0.187,0.721,414,321,0.000000,20 -fla_B1_T16384,fla,gdn,1,16384,64,64,128,1.567,4.957,99,94,0.000000,20 -flash_qla_B1_T16384,flash_qla,gdn,1,16384,64,64,128,0.674,2.648,229,175,0.000000,20 -cudnn_B1_T16384,cudnn,gdn,1,16384,64,64,128,0.305,1.309,506,354,0.000000,20 -fla_B1_T32768,fla,gdn,1,32768,64,64,128,3.119,9.976,99,93,0.000000,20 -flash_qla_B1_T32768,flash_qla,gdn,1,32768,64,64,128,1.329,5.294,233,175,0.000000,20 -cudnn_B1_T32768,cudnn,gdn,1,32768,64,64,128,0.529,2.494,585,372,0.000000,20 diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_20260814.csv b/benchmark/linear_attention/results/gdn/gb200/gdn_20260814.csv new file mode 100644 index 000000000..b95b666f2 --- /dev/null +++ b/benchmark/linear_attention/results/gdn/gb200/gdn_20260814.csv @@ -0,0 +1,41 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +gdn_h64,cudnn,gdn,4,2048,64,64,128,0.176,0.661,440,351,0.000000,20,3.08,2.66 +gdn_hon,cudnn_state_on,gdn,4,2048,64,64,128,0.227,0.522,340,444,0.000000,20,2.38,3.37 +gdn_h64,fla,gdn,4,2048,64,64,128,0.756,2.323,102,100,0.000000,20,0.72,0.76 +gdn_h64,flash_qla,gdn,4,2048,64,64,128,0.272,0.726,284,320,0.000000,20,1.99,2.42 +gdn_h64,cudnn,gdn,4,4096,64,64,128,0.305,1.271,507,365,0.000000,20,3.55,2.77 +gdn_hon,cudnn_state_on,gdn,4,4096,64,64,128,0.409,0.998,378,465,0.000000,20,2.65,3.52 +gdn_h64,fla,gdn,4,4096,64,64,128,1.478,4.596,105,101,0.000000,20,0.73,0.76 +gdn_h64,flash_qla,gdn,4,4096,64,64,128,0.513,1.409,301,329,0.000000,20,2.11,2.49 +gdn_h64,cudnn,gdn,4,8192,64,64,128,0.567,2.490,545,373,0.000000,20,3.82,2.82 +gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.789,1.953,392,475,0.000000,20,2.74,3.60 +gdn_h64,fla,gdn,4,8192,64,64,128,2.927,9.198,106,101,0.000000,20,0.74,0.76 +gdn_h64,flash_qla,gdn,4,8192,64,64,128,0.993,2.757,312,337,0.000000,20,2.18,2.55 +gdn_h64,cudnn,gdn,4,16384,64,64,128,1.062,4.930,582,376,0.000000,20,4.07,2.85 +gdn_hon,cudnn_state_on,gdn,4,16384,64,64,128,1.575,3.866,393,480,0.000000,20,2.75,3.64 +gdn_h64,fla,gdn,4,16384,64,64,128,5.842,18.483,106,100,0.000000,20,0.74,0.76 +gdn_h64,flash_qla,gdn,4,16384,64,64,128,1.961,5.502,315,337,0.000000,20,2.21,2.56 +gdn_h64,cudnn,gdn,4,32768,64,64,128,2.086,9.814,593,378,0.000000,20,4.15,2.87 +gdn_hon,cudnn_state_on,gdn,4,32768,64,64,128,3.256,7.689,380,483,0.000000,20,2.66,3.66 +gdn_h64,fla,gdn,4,32768,64,64,128,11.715,37.419,106,99,0.000000,20,0.74,0.75 +gdn_h64,flash_qla,gdn,4,32768,64,64,128,3.888,10.967,318,338,0.000000,20,2.23,2.56 +gdn_h64,cudnn,gdn,1,8192,64,64,128,0.178,0.667,434,347,0.000000,20,3.04,2.63 +gdn_hon,cudnn_state_on,gdn,1,8192,64,64,128,0.226,0.525,342,442,0.000000,20,2.39,3.35 +gdn_h64,fla,gdn,1,8192,64,64,128,0.821,2.482,94,93,0.000000,20,0.66,0.71 +gdn_h64,flash_qla,gdn,1,8192,64,64,128,0.342,1.313,226,177,0.000000,20,1.58,1.34 +gdn_h64,cudnn,gdn,2,8192,64,64,128,0.284,1.244,544,373,0.000000,20,3.81,2.82 +gdn_hon,cudnn_state_on,gdn,2,8192,64,64,128,0.387,0.973,400,477,0.000000,20,2.80,3.61 +gdn_h64,fla,gdn,2,8192,64,64,128,1.471,4.840,105,96,0.000000,20,0.74,0.73 +gdn_h64,flash_qla,gdn,2,8192,64,64,128,0.503,1.434,307,324,0.000000,20,2.15,2.45 +gdn_h64,cudnn,gdn,4,8192,64,64,128,0.565,2.489,547,373,0.000000,20,3.83,2.82 +gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.786,1.951,393,475,0.000000,20,2.75,3.60 +gdn_h64,fla,gdn,4,8192,64,64,128,2.927,9.200,106,101,0.000000,20,0.74,0.76 +gdn_h64,flash_qla,gdn,4,8192,64,64,128,0.992,2.769,312,335,0.000000,20,2.18,2.54 +gdn_h64,cudnn,gdn,8,8192,64,64,128,1.057,4.978,585,373,0.000000,20,4.10,2.82 +gdn_hon,cudnn_state_on,gdn,8,8192,64,64,128,1.555,3.876,398,479,0.000000,20,2.78,3.63 +gdn_h64,fla,gdn,8,8192,64,64,128,5.601,17.705,110,105,0.000000,20,0.77,0.79 +gdn_h64,flash_qla,gdn,8,8192,64,64,128,1.963,5.390,315,344,0.000000,20,2.20,2.61 +gdn_h64,cudnn,gdn,16,8192,64,64,128,1.891,9.164,654,405,0.000000,20,4.58,3.07 +gdn_hon,cudnn_state_on,gdn,16,8192,64,64,128,2.905,7.242,426,512,0.000000,20,2.98,3.88 +gdn_h64,fla,gdn,16,8192,64,64,128,11.050,35.154,112,106,0.000000,20,0.78,0.80 +gdn_h64,flash_qla,gdn,16,8192,64,64,128,3.622,9.670,342,384,0.000000,20,2.39,2.91 diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_bw.png b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_bw.png new file mode 100644 index 000000000..bd44211e5 Binary files /dev/null and b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_flops.png b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_flops.png new file mode 100644 index 000000000..5f4a442c7 Binary files /dev/null and b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_bw.png b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_bw.png new file mode 100644 index 000000000..4b834a1b5 Binary files /dev/null and b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_flops.png b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_flops.png new file mode 100644 index 000000000..4368b9288 Binary files /dev/null and b/benchmark/linear_attention/results/gdn/gb200/gdn_fixed_seq_flops.png differ diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_20260814.csv b/benchmark/linear_attention/results/gdn/gb300/gdn_20260814.csv new file mode 100644 index 000000000..a784e4c45 --- /dev/null +++ b/benchmark/linear_attention/results/gdn/gb300/gdn_20260814.csv @@ -0,0 +1,41 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +gdn_h64,cudnn,gdn,4,2048,64,64,128,0.165,0.660,469,352,0.000000,20,3.28,2.66 +gdn_hon,cudnn_state_on,gdn,4,2048,64,64,128,0.217,0.501,356,463,0.000000,20,2.49,3.51 +gdn_h64,fla,gdn,4,2048,64,64,128,0.740,2.284,104,102,0.000000,20,0.73,0.77 +gdn_h64,flash_qla,gdn,4,2048,64,64,128,0.273,0.729,283,318,0.000000,20,1.98,2.41 +gdn_h64,cudnn,gdn,4,4096,64,64,128,0.289,1.219,534,381,0.000000,20,3.74,2.88 +gdn_hon,cudnn_state_on,gdn,4,4096,64,64,128,0.396,0.954,390,486,0.000000,20,2.73,3.68 +gdn_h64,fla,gdn,4,4096,64,64,128,1.447,4.530,107,102,0.000000,20,0.75,0.78 +gdn_h64,flash_qla,gdn,4,4096,64,64,128,0.516,1.406,299,330,0.000000,20,2.10,2.50 +gdn_h64,cudnn,gdn,4,8192,64,64,128,0.546,2.389,566,388,0.000000,20,3.96,2.94 +gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.776,1.864,398,498,0.000000,20,2.79,3.77 +gdn_h64,fla,gdn,4,8192,64,64,128,2.865,9.067,108,102,0.000000,20,0.76,0.78 +gdn_h64,flash_qla,gdn,4,8192,64,64,128,1.000,2.766,309,335,0.000000,20,2.16,2.54 +gdn_h64,cudnn,gdn,4,16384,64,64,128,1.049,4.715,590,394,0.000000,20,4.13,2.98 +gdn_hon,cudnn_state_on,gdn,4,16384,64,64,128,1.557,3.675,397,505,0.000000,20,2.78,3.83 +gdn_h64,fla,gdn,4,16384,64,64,128,5.712,18.208,108,102,0.000000,20,0.76,0.77 +gdn_h64,flash_qla,gdn,4,16384,64,64,128,1.978,5.503,313,337,0.000000,20,2.19,2.55 +gdn_h64,cudnn,gdn,4,32768,64,64,128,2.073,9.381,597,396,0.000000,20,4.18,3.00 +gdn_hon,cudnn_state_on,gdn,4,32768,64,64,128,3.187,7.294,388,509,0.000000,20,2.72,3.85 +gdn_h64,fla,gdn,4,32768,64,64,128,11.451,36.798,108,101,0.000000,20,0.76,0.76 +gdn_h64,flash_qla,gdn,4,32768,64,64,128,3.918,10.903,316,340,0.000000,20,2.21,2.58 +gdn_h64,cudnn,gdn,1,8192,64,64,128,0.171,0.620,453,374,0.000000,20,3.17,2.83 +gdn_hon,cudnn_state_on,gdn,1,8192,64,64,128,0.221,0.480,350,483,0.000000,20,2.45,3.66 +gdn_h64,fla,gdn,1,8192,64,64,128,0.798,2.443,97,95,0.000000,20,0.68,0.72 +gdn_h64,flash_qla,gdn,1,8192,64,64,128,0.370,1.321,209,176,0.000000,20,1.46,1.33 +gdn_h64,cudnn,gdn,2,8192,64,64,128,0.278,1.146,556,405,0.000000,20,3.89,3.07 +gdn_hon,cudnn_state_on,gdn,2,8192,64,64,128,0.379,0.880,408,527,0.000000,20,2.85,4.00 +gdn_h64,fla,gdn,2,8192,64,64,128,1.447,4.768,107,97,0.000000,20,0.75,0.74 +gdn_h64,flash_qla,gdn,2,8192,64,64,128,0.505,1.418,306,327,0.000000,20,2.14,2.48 +gdn_h64,cudnn,gdn,4,8192,64,64,128,0.546,2.431,566,382,0.000000,20,3.96,2.89 +gdn_hon,cudnn_state_on,gdn,4,8192,64,64,128,0.778,1.865,398,497,0.000000,20,2.78,3.77 +gdn_h64,fla,gdn,4,8192,64,64,128,2.864,9.062,108,102,0.000000,20,0.76,0.78 +gdn_h64,flash_qla,gdn,4,8192,64,64,128,1.003,2.760,308,336,0.000000,20,2.16,2.55 +gdn_h64,cudnn,gdn,8,8192,64,64,128,1.042,4.935,593,376,0.000000,20,4.15,2.85 +gdn_hon,cudnn_state_on,gdn,8,8192,64,64,128,1.533,3.861,403,481,0.000000,20,2.82,3.64 +gdn_h64,fla,gdn,8,8192,64,64,128,5.505,17.444,112,106,0.000000,20,0.79,0.81 +gdn_h64,flash_qla,gdn,8,8192,64,64,128,1.977,5.358,313,346,0.000000,20,2.19,2.62 +gdn_h64,cudnn,gdn,16,8192,64,64,128,1.910,9.012,648,412,0.000000,20,4.53,3.12 +gdn_hon,cudnn_state_on,gdn,16,8192,64,64,128,2.866,7.139,432,520,0.000000,20,3.02,3.94 +gdn_h64,fla,gdn,16,8192,64,64,128,10.881,34.661,114,107,0.000000,20,0.80,0.81 +gdn_h64,flash_qla,gdn,16,8192,64,64,128,3.648,9.575,339,388,0.000000,20,2.37,2.94 diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_bw.png b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_bw.png new file mode 100644 index 000000000..1f427d9c4 Binary files /dev/null and b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_flops.png b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_flops.png new file mode 100644 index 000000000..eca28eefd Binary files /dev/null and b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_bw.png b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_bw.png new file mode 100644 index 000000000..113ca830d Binary files /dev/null and b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_flops.png b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_flops.png new file mode 100644 index 000000000..2afe662f9 Binary files /dev/null and b/benchmark/linear_attention/results/gdn/gb300/gdn_fixed_seq_flops.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_20260814.csv b/benchmark/linear_attention/results/gdn2/gb200/gdn2_20260814.csv new file mode 100644 index 000000000..f8bb549fe --- /dev/null +++ b/benchmark/linear_attention/results/gdn2/gb200/gdn2_20260814.csv @@ -0,0 +1,21 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +gdn2_h64,cudnn,gdn2,4,2048,64,64,128,0.257,2.050,176,66,0.000000,20,4.18,2.42 +gdn2_h64,fla,gdn2,4,2048,64,64,128,1.536,5.829,29,23,0.000000,20,0.70,0.85 +gdn2_h64,cudnn,gdn2,4,4096,64,64,128,0.474,4.075,190,66,0.000000,20,4.53,2.44 +gdn2_h64,fla,gdn2,4,4096,64,64,128,3.039,11.577,30,23,0.000000,20,0.71,0.86 +gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.908,8.136,199,67,0.000000,20,4.73,2.44 +gdn2_h64,fla,gdn2,4,8192,64,64,128,6.008,23.062,30,23,0.000000,20,0.71,0.86 +gdn2_h64,cudnn,gdn2,4,16384,64,64,128,1.770,16.278,204,66,0.000000,20,4.85,2.44 +gdn2_h64,fla,gdn2,4,16384,64,64,128,12.007,46.079,30,23,0.000000,20,0.72,0.86 +gdn2_h64,cudnn,gdn2,4,32768,64,64,128,3.508,32.541,206,67,0.000000,20,4.90,2.44 +gdn2_h64,fla,gdn2,4,32768,64,64,128,23.826,88.467,30,24,0.000000,20,0.72,0.90 +gdn2_h64,cudnn,gdn2,1,8192,64,64,128,0.476,3.672,95,37,0.000000,20,2.26,1.35 +gdn2_h64,fla,gdn2,1,8192,64,64,128,1.688,6.099,27,22,0.000000,20,0.64,0.81 +gdn2_h64,cudnn,gdn2,2,8192,64,64,128,0.461,4.031,196,67,0.000000,20,4.66,2.46 +gdn2_h64,fla,gdn2,2,8192,64,64,128,3.060,11.576,29,23,0.000000,20,0.70,0.86 +gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.906,8.130,199,67,0.000000,20,4.74,2.44 +gdn2_h64,fla,gdn2,4,8192,64,64,128,6.004,23.060,30,23,0.000000,20,0.72,0.86 +gdn2_h64,cudnn,gdn2,8,8192,64,64,128,1.718,16.113,210,67,0.000000,20,5.00,2.47 +gdn2_h64,fla,gdn2,8,8192,64,64,128,11.849,45.914,30,24,0.000000,20,0.72,0.87 +gdn2_h64,cudnn,gdn2,16,8192,64,64,128,3.001,28.796,240,75,0.000000,20,5.72,2.76 +gdn2_h64,fla,gdn2,16,8192,64,64,128,23.265,87.351,31,25,0.000000,20,0.74,0.91 diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_bw.png b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_bw.png new file mode 100644 index 000000000..07d543fd0 Binary files /dev/null and b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_flops.png b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_flops.png new file mode 100644 index 000000000..4b3fdcfb4 Binary files /dev/null and b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_bw.png b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_bw.png new file mode 100644 index 000000000..bc1065be0 Binary files /dev/null and b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_flops.png b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_flops.png new file mode 100644 index 000000000..f5e74978f Binary files /dev/null and b/benchmark/linear_attention/results/gdn2/gb200/gdn2_fixed_seq_flops.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_20260814.csv b/benchmark/linear_attention/results/gdn2/gb300/gdn2_20260814.csv new file mode 100644 index 000000000..83a2af408 --- /dev/null +++ b/benchmark/linear_attention/results/gdn2/gb300/gdn2_20260814.csv @@ -0,0 +1,21 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +gdn2_h64,cudnn,gdn2,4,2048,64,64,128,0.249,2.034,181,67,0.000000,20,4.32,2.44 +gdn2_h64,fla,gdn2,4,2048,64,64,128,1.479,5.669,30,24,0.000000,20,0.73,0.88 +gdn2_h64,cudnn,gdn2,4,4096,64,64,128,0.473,4.032,191,67,0.000000,20,4.54,2.46 +gdn2_h64,fla,gdn2,4,4096,64,64,128,2.916,11.276,31,24,0.000000,20,0.74,0.88 +gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.918,8.036,197,67,0.000000,20,4.68,2.47 +gdn2_h64,fla,gdn2,4,8192,64,64,128,5.772,22.466,31,24,0.000000,20,0.74,0.88 +gdn2_h64,cudnn,gdn2,4,16384,64,64,128,1.809,16.039,199,67,0.000000,20,4.75,2.48 +gdn2_h64,fla,gdn2,4,16384,64,64,128,11.469,44.851,31,24,0.000000,20,0.75,0.89 +gdn2_h64,cudnn,gdn2,4,32768,64,64,128,3.590,32.059,201,68,0.000000,20,4.79,2.48 +gdn2_h64,fla,gdn2,4,32768,64,64,128,22.873,86.102,32,25,0.000000,20,0.75,0.92 +gdn2_h64,cudnn,gdn2,1,8192,64,64,128,0.477,3.610,94,37,0.000000,20,2.25,1.38 +gdn2_h64,fla,gdn2,1,8192,64,64,128,1.622,5.937,28,23,0.000000,20,0.66,0.84 +gdn2_h64,cudnn,gdn2,2,8192,64,64,128,0.471,4.033,192,67,0.000000,20,4.56,2.46 +gdn2_h64,fla,gdn2,2,8192,64,64,128,2.931,11.275,31,24,0.000000,20,0.73,0.88 +gdn2_h64,cudnn,gdn2,4,8192,64,64,128,0.918,8.032,196,67,0.000000,20,4.68,2.47 +gdn2_h64,fla,gdn2,4,8192,64,64,128,5.764,22.456,31,24,0.000000,20,0.75,0.88 +gdn2_h64,cudnn,gdn2,8,8192,64,64,128,1.752,14.667,206,74,0.000000,20,4.90,2.71 +gdn2_h64,fla,gdn2,8,8192,64,64,128,11.361,44.740,32,24,0.000000,20,0.76,0.89 +gdn2_h64,cudnn,gdn2,16,8192,64,64,128,3.066,28.539,235,76,0.000000,20,5.60,2.78 +gdn2_h64,fla,gdn2,16,8192,64,64,128,22.279,84.983,32,25,0.000000,20,0.77,0.93 diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_bw.png b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_bw.png new file mode 100644 index 000000000..fe5da5fce Binary files /dev/null and b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_flops.png b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_flops.png new file mode 100644 index 000000000..d8ddb0f5c Binary files /dev/null and b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_bw.png b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_bw.png new file mode 100644 index 000000000..c89176339 Binary files /dev/null and b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_flops.png b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_flops.png new file mode 100644 index 000000000..73822db2e Binary files /dev/null and b/benchmark/linear_attention/results/gdn2/gb300/gdn2_fixed_seq_flops.png differ diff --git a/benchmark/linear_attention/results/kda/gb200/kda_20260814.csv b/benchmark/linear_attention/results/kda/gb200/kda_20260814.csv new file mode 100644 index 000000000..242c8cae6 --- /dev/null +++ b/benchmark/linear_attention/results/kda/gb200/kda_20260814.csv @@ -0,0 +1,41 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +kda_h64,cudnn,kda,4,2048,64,64,128,0.265,1.875,170,72,0.000000,20,3.05,2.22 +kda_hon,cudnn_state_on,kda,4,2048,64,64,128,0.354,1.566,127,86,0.000000,20,2.28,2.66 +kda_h64,fla,kda,4,2048,64,64,128,1.462,5.450,31,25,0.000000,20,0.55,0.76 +kda_h64,flash_kda,kda,4,2048,64,64,128,0.436,0.000,103,0,0.000000,20,1.85,0.00 +kda_h64,cudnn,kda,4,4096,64,64,128,0.506,3.727,178,73,0.000000,20,3.19,2.24 +kda_hon,cudnn_state_on,kda,4,4096,64,64,128,0.679,3.104,133,87,0.000000,20,2.38,2.68 +kda_h64,fla,kda,4,4096,64,64,128,2.892,10.862,31,25,0.000000,20,0.56,0.77 +kda_h64,flash_kda,kda,4,4096,64,64,128,0.851,0.000,106,0,0.000000,20,1.90,0.00 +kda_h64,cudnn,kda,4,8192,64,64,128,0.975,7.420,185,73,0.000000,20,3.31,2.25 +kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.337,6.181,135,88,0.000000,20,2.41,2.70 +kda_h64,fla,kda,4,8192,64,64,128,5.699,21.766,32,25,0.000000,20,0.57,0.77 +kda_h64,flash_kda,kda,4,8192,64,64,128,1.692,0.000,107,0,0.000000,20,1.91,0.00 +kda_h64,cudnn,kda,4,16384,64,64,128,1.923,14.847,188,73,0.000000,20,3.36,2.25 +kda_hon,cudnn_state_on,kda,4,16384,64,64,128,2.674,12.330,135,88,0.000000,20,2.42,2.70 +kda_h64,fla,kda,4,16384,64,64,128,11.386,43.630,32,25,0.000000,20,0.57,0.76 +kda_h64,flash_kda,kda,4,16384,64,64,128,3.368,0.000,107,0,0.000000,20,1.92,0.00 +kda_h64,cudnn,kda,4,32768,64,64,128,3.896,29.622,185,73,0.000000,20,3.32,2.25 +kda_hon,cudnn_state_on,kda,4,32768,64,64,128,5.256,24.661,137,88,0.000000,20,2.46,2.70 +kda_h64,fla,kda,4,32768,64,64,128,22.736,87.647,32,25,0.000000,20,0.57,0.76 +kda_h64,flash_kda,kda,4,32768,64,64,128,6.717,0.000,107,0,0.000000,20,1.92,0.00 +kda_h64,cudnn,kda,1,8192,64,64,128,0.485,2.861,93,47,0.000000,20,1.67,1.46 +kda_hon,cudnn_state_on,kda,1,8192,64,64,128,0.662,2.357,68,57,0.000000,20,1.22,1.77 +kda_h64,fla,kda,1,8192,64,64,128,1.609,5.757,28,24,0.000000,20,0.50,0.72 +kda_h64,flash_kda,kda,1,8192,64,64,128,0.902,0.000,50,0,0.000000,20,0.90,0.00 +kda_h64,cudnn,kda,2,8192,64,64,128,0.500,3.692,180,73,0.000000,20,3.23,2.26 +kda_hon,cudnn_state_on,kda,2,8192,64,64,128,0.668,3.127,135,87,0.000000,20,2.42,2.66 +kda_h64,fla,kda,2,8192,64,64,128,2.892,10.908,31,25,0.000000,20,0.56,0.76 +kda_h64,flash_kda,kda,2,8192,64,64,128,1.074,0.000,84,0,0.000000,20,1.50,0.00 +kda_h64,cudnn,kda,4,8192,64,64,128,0.984,7.425,183,73,0.000000,20,3.28,2.24 +kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.330,6.186,136,87,0.000000,20,2.43,2.69 +kda_h64,fla,kda,4,8192,64,64,128,5.710,21.769,32,25,0.000000,20,0.57,0.77 +kda_h64,flash_kda,kda,4,8192,64,64,128,1.692,0.000,107,0,0.000000,20,1.91,0.00 +kda_h64,cudnn,kda,8,8192,64,64,128,1.898,14.872,190,73,0.000000,20,3.40,2.24 +kda_hon,cudnn_state_on,kda,8,8192,64,64,128,2.706,12.279,133,88,0.000000,20,2.39,2.71 +kda_h64,fla,kda,8,8192,64,64,128,11.241,43.337,32,25,0.000000,20,0.57,0.77 +kda_h64,flash_kda,kda,8,8192,64,64,128,3.363,0.000,107,0,0.000000,20,1.92,0.00 +kda_h64,cudnn,kda,16,8192,64,64,128,3.292,26.193,219,83,0.000000,20,3.92,2.55 +kda_hon,cudnn_state_on,kda,16,8192,64,64,128,4.926,21.509,146,101,0.000000,20,2.62,3.10 +kda_h64,fla,kda,16,8192,64,64,128,22.024,85.810,33,25,0.000000,20,0.59,0.78 +kda_h64,flash_kda,kda,16,8192,64,64,128,6.369,0.000,113,0,0.000000,20,2.03,0.00 diff --git a/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_bw.png b/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_bw.png new file mode 100644 index 000000000..37db34ffe Binary files /dev/null and b/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_flops.png b/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_flops.png new file mode 100644 index 000000000..90bc35b48 Binary files /dev/null and b/benchmark/linear_attention/results/kda/gb200/kda_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_bw.png b/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_bw.png new file mode 100644 index 000000000..165b17bca Binary files /dev/null and b/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_flops.png b/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_flops.png new file mode 100644 index 000000000..b95798e4d Binary files /dev/null and b/benchmark/linear_attention/results/kda/gb200/kda_fixed_seq_flops.png differ diff --git a/benchmark/linear_attention/results/kda/gb300/kda_20260814.csv b/benchmark/linear_attention/results/kda/gb300/kda_20260814.csv new file mode 100644 index 000000000..5b9d21bf2 --- /dev/null +++ b/benchmark/linear_attention/results/kda/gb300/kda_20260814.csv @@ -0,0 +1,41 @@ +case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters,fwd_bw,bwd_bw +kda_h64,cudnn,kda,4,2048,64,64,128,0.261,1.867,172,72,0.000000,20,3.09,2.23 +kda_hon,cudnn_state_on,kda,4,2048,64,64,128,0.349,1.564,129,86,0.000000,20,2.31,2.66 +kda_h64,fla,kda,4,2048,64,64,128,1.421,5.352,32,25,0.000000,20,0.57,0.78 +kda_h64,flash_kda,kda,4,2048,64,64,128,0.433,0.000,104,0,0.000000,20,1.86,0.00 +kda_h64,cudnn,kda,4,4096,64,64,128,0.496,3.709,182,73,0.000000,20,3.26,2.25 +kda_hon,cudnn_state_on,kda,4,4096,64,64,128,0.673,3.097,134,87,0.000000,20,2.40,2.69 +kda_h64,fla,kda,4,4096,64,64,128,2.803,10.697,32,25,0.000000,20,0.58,0.78 +kda_h64,flash_kda,kda,4,4096,64,64,128,0.849,0.000,106,0,0.000000,20,1.90,0.00 +kda_h64,cudnn,kda,4,8192,64,64,128,0.965,7.391,187,73,0.000000,20,3.35,2.26 +kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.319,6.176,137,88,0.000000,20,2.45,2.70 +kda_h64,fla,kda,4,8192,64,64,128,5.533,21.361,33,25,0.000000,20,0.58,0.78 +kda_h64,flash_kda,kda,4,8192,64,64,128,1.687,0.000,107,0,0.000000,20,1.91,0.00 +kda_h64,cudnn,kda,4,16384,64,64,128,1.899,14.751,190,73,0.000000,20,3.40,2.26 +kda_hon,cudnn_state_on,kda,4,16384,64,64,128,2.607,12.312,138,88,0.000000,20,2.48,2.71 +kda_h64,fla,kda,4,16384,64,64,128,11.072,42.801,33,25,0.000000,20,0.58,0.78 +kda_h64,flash_kda,kda,4,16384,64,64,128,3.357,0.000,107,0,0.000000,20,1.92,0.00 +kda_h64,cudnn,kda,4,32768,64,64,128,3.844,29.469,188,73,0.000000,20,3.36,2.26 +kda_hon,cudnn_state_on,kda,4,32768,64,64,128,5.200,24.597,139,88,0.000000,20,2.48,2.71 +kda_h64,fla,kda,4,32768,64,64,128,22.044,85.871,33,25,0.000000,20,0.59,0.78 +kda_h64,flash_kda,kda,4,32768,64,64,128,6.700,0.000,108,0,0.000000,20,1.93,0.00 +kda_h64,cudnn,kda,1,8192,64,64,128,0.481,2.860,94,47,0.000000,20,1.68,1.46 +kda_hon,cudnn_state_on,kda,1,8192,64,64,128,0.655,2.524,69,54,0.000000,20,1.23,1.65 +kda_h64,fla,kda,1,8192,64,64,128,1.566,5.651,29,24,0.000000,20,0.52,0.74 +kda_h64,flash_kda,kda,1,8192,64,64,128,0.897,0.000,50,0,0.000000,20,0.90,0.00 +kda_h64,cudnn,kda,2,8192,64,64,128,0.499,3.676,181,74,0.000000,20,3.24,2.27 +kda_hon,cudnn_state_on,kda,2,8192,64,64,128,0.654,3.121,138,87,0.000000,20,2.47,2.67 +kda_h64,fla,kda,2,8192,64,64,128,2.814,10.715,32,25,0.000000,20,0.57,0.78 +kda_h64,flash_kda,kda,2,8192,64,64,128,1.070,0.000,84,0,0.000000,20,1.51,0.00 +kda_h64,cudnn,kda,4,8192,64,64,128,0.976,7.393,185,73,0.000000,20,3.31,2.25 +kda_hon,cudnn_state_on,kda,4,8192,64,64,128,1.319,6.175,137,88,0.000000,20,2.45,2.70 +kda_h64,fla,kda,4,8192,64,64,128,5.550,21.362,33,25,0.000000,20,0.58,0.78 +kda_h64,flash_kda,kda,4,8192,64,64,128,1.686,0.000,107,0,0.000000,20,1.92,0.00 +kda_h64,cudnn,kda,8,8192,64,64,128,1.855,13.464,194,80,0.000000,20,3.48,2.48 +kda_hon,cudnn_state_on,kda,8,8192,64,64,128,2.669,11.021,135,98,0.000000,20,2.42,3.02 +kda_h64,fla,kda,8,8192,64,64,128,10.939,42.524,33,25,0.000000,20,0.59,0.78 +kda_h64,flash_kda,kda,8,8192,64,64,128,3.351,0.000,108,0,0.000000,20,1.93,0.00 +kda_h64,cudnn,kda,16,8192,64,64,128,3.224,26.006,224,83,0.000000,20,4.01,2.56 +kda_hon,cudnn_state_on,kda,16,8192,64,64,128,4.767,21.475,151,101,0.000000,20,2.71,3.10 +kda_h64,fla,kda,16,8192,64,64,128,21.432,84.215,34,26,0.000000,20,0.60,0.79 +kda_h64,flash_kda,kda,16,8192,64,64,128,6.346,0.000,114,0,0.000000,20,2.04,0.00 diff --git a/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_bw.png b/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_bw.png new file mode 100644 index 000000000..0f39146ce Binary files /dev/null and b/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_bw.png differ diff --git a/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_flops.png b/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_flops.png new file mode 100644 index 000000000..331d8c900 Binary files /dev/null and b/benchmark/linear_attention/results/kda/gb300/kda_fixed_batch_flops.png differ diff --git a/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_bw.png b/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_bw.png new file mode 100644 index 000000000..bf4badc36 Binary files /dev/null and b/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_bw.png differ diff --git a/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_flops.png b/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_flops.png new file mode 100644 index 000000000..91a081c06 Binary files /dev/null and b/benchmark/linear_attention/results/kda/gb300/kda_fixed_seq_flops.png differ diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index a24c45d73..e327a475f 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -2447,11 +2447,11 @@ def _linear_attention_final_state_dims(node): return [cu.dim[0] - 1, max(q[1], v[1]), q[2], v[2]] -def _linear_attention_h_dims(node): - # [total_h, HO, K, V] at capacity: sum_b (sl_b - 1) // N <= total_T // N +def _linear_attention_state_checkpoints_dims(node): n = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) q, v = node.inputs["q"].dim, node.inputs["v"].dim - if not n or not q or not v: + 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]] @@ -2620,20 +2620,20 @@ def _training_phase(node): # norm stats exist only in TRAINING forward phase # ---- linear attention ---------------------------------------------------- "gdn": dict( node_type=NodeType.GDN, - inputs=("q", "k", "v", "g", "beta", "cu_seqlens", "initial_state"), - attrs=("scale", "output_final_state", "use_qk_l2norm", "checkpoint_every_n_tokens"), - outputs=("O", "final_state", "H"), + inputs=("q", "k", "v", "g", "beta", "cu_seqlens", "initial_state", "a_log", "dt_bias"), + attrs=("scale", "output_final_state", "use_qk_l2norm", "checkpoint_every_n_tokens", "safe_gate", "batch_invariant"), + outputs=("O", "final_state", "state_checkpoints"), maybe={ "final_state": lambda n: bool(n.params.get("output_final_state", False)), - "H": lambda n: bool(n.params.get("checkpoint_every_n_tokens") or 0), + "state_checkpoints": lambda n: bool(n.params.get("checkpoint_every_n_tokens") or 0), }, - infer={"O": _linear_attention_o_dims, "final_state": _linear_attention_final_state_dims, "H": _linear_attention_h_dims}, + infer={"O": _linear_attention_o_dims, "final_state": _linear_attention_final_state_dims, "state_checkpoints": _linear_attention_state_checkpoints_dims}, python_only=True, ), "gdn_bwd": dict( node_type=NodeType.GDN_BWD, - inputs=("q", "k", "v", "g", "beta", "cu_seqlens", "dO", "h", "initial_state", "d_final_state"), - attrs=("scale", "use_qk_l2norm"), + inputs=("q", "k", "v", "g", "beta", "cu_seqlens", "dO", "state_checkpoints", "initial_state", "d_final_state"), + attrs=("scale", "use_qk_l2norm", "batch_invariant"), outputs=("dQ", "dK", "dV", "dG", "dBeta", "d_initial_state"), maybe={"d_initial_state": lambda n: "initial_state" in n.inputs}, infer={"dQ": _like("q"), "dK": _like("k"), "dV": _like("v"), "dG": _like("g"), "dBeta": _like("beta"), "d_initial_state": _like("initial_state")}, @@ -2642,19 +2642,28 @@ def _training_phase(node): # norm stats exist only in TRAINING forward phase "kda": dict( node_type=NodeType.KDA, inputs=("q", "k", "v", "g", "beta", "cu_seqlens", "initial_state", "a_log", "dt_bias"), - attrs=("scale", "output_final_state", "use_qk_l2norm", "checkpoint_every_n_tokens", "use_beta_sigmoid", "safe_gate", "gate_lower_bound"), - outputs=("O", "final_state", "H"), + attrs=( + "scale", + "output_final_state", + "use_qk_l2norm", + "checkpoint_every_n_tokens", + "use_beta_sigmoid", + "safe_gate", + "gate_lower_bound", + "batch_invariant", + ), + outputs=("O", "final_state", "state_checkpoints"), maybe={ "final_state": lambda n: bool(n.params.get("output_final_state", False)), - "H": lambda n: bool(n.params.get("checkpoint_every_n_tokens") or 0), + "state_checkpoints": lambda n: bool(n.params.get("checkpoint_every_n_tokens") or 0), }, - infer={"O": _like("v"), "final_state": _linear_attention_final_state_dims, "H": _linear_attention_h_dims}, + infer={"O": _linear_attention_o_dims, "final_state": _linear_attention_final_state_dims, "state_checkpoints": _linear_attention_state_checkpoints_dims}, python_only=True, ), "kda_bwd": dict( node_type=NodeType.KDA_BWD, - inputs=("q", "k", "v", "g", "beta", "cu_seqlens", "dO", "h", "initial_state", "d_final_state"), - attrs=("scale", "use_qk_l2norm"), + inputs=("q", "k", "v", "g", "beta", "cu_seqlens", "dO", "state_checkpoints", "initial_state", "d_final_state"), + attrs=("scale", "use_qk_l2norm", "batch_invariant"), outputs=("dQ", "dK", "dV", "dG", "dBeta", "d_initial_state"), maybe={"d_initial_state": lambda n: "initial_state" in n.inputs}, infer={"dQ": _like("q"), "dK": _like("k"), "dV": _like("v"), "dG": _like("g"), "dBeta": _like("beta"), "d_initial_state": _like("initial_state")}, @@ -2662,20 +2671,20 @@ def _training_phase(node): # norm stats exist only in TRAINING forward phase ), "gdn2": dict( node_type=NodeType.GDN2, - inputs=("q", "k", "v", "g", "beta", "w", "cu_seqlens", "initial_state"), - attrs=("scale", "output_final_state", "use_qk_l2norm", "checkpoint_every_n_tokens", "use_beta_w_sigmoid"), - outputs=("O", "final_state", "H"), + inputs=("q", "k", "v", "g", "beta", "w", "cu_seqlens", "initial_state", "a_log", "dt_bias"), + attrs=("scale", "output_final_state", "use_qk_l2norm", "checkpoint_every_n_tokens", "safe_gate", "gate_lower_bound", "batch_invariant"), + outputs=("O", "final_state", "state_checkpoints"), maybe={ "final_state": lambda n: bool(n.params.get("output_final_state", False)), - "H": lambda n: bool(n.params.get("checkpoint_every_n_tokens") or 0), + "state_checkpoints": lambda n: bool(n.params.get("checkpoint_every_n_tokens") or 0), }, - infer={"O": _like("v"), "final_state": _linear_attention_final_state_dims, "H": _linear_attention_h_dims}, + infer={"O": _linear_attention_o_dims, "final_state": _linear_attention_final_state_dims, "state_checkpoints": _linear_attention_state_checkpoints_dims}, python_only=True, ), "gdn2_bwd": dict( node_type=NodeType.GDN2_BWD, - inputs=("q", "k", "v", "g", "beta", "w", "cu_seqlens", "dO", "h", "initial_state", "d_final_state"), - attrs=("scale",), + inputs=("q", "k", "v", "g", "beta", "w", "cu_seqlens", "dO", "state_checkpoints", "initial_state", "d_final_state"), + attrs=("scale", "use_qk_l2norm", "batch_invariant"), outputs=("dQ", "dK", "dV", "dG", "dBeta", "dW", "d_initial_state"), maybe={"d_initial_state": lambda n: "initial_state" in n.inputs}, infer={ diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index e3d5dd46a..85599938b 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -159,6 +159,12 @@ def all_contiguous(self): ok, offender = self.native.all_contiguous() return ok, (int(offender) if offender else -1) + def all_dense_layout(self): + """``(ok, slot)`` over every filled operand: the innermost size>1 dim + must be stride-1. Padded or permuted outer strides pass.""" + ok, offender = self.native.all_dense_layout() + return ok, (int(offender) if offender else -1) + @property def index_of_uid(self): if self._index_of is None: diff --git a/python/cudnn/engines/manifest.py b/python/cudnn/engines/manifest.py index fabdc4756..bb68f1db3 100644 --- a/python/cudnn/engines/manifest.py +++ b/python/cudnn/engines/manifest.py @@ -143,6 +143,7 @@ def offered_ids(self) -> Dict[str, int]: "cudnn.linear_attention", "GdnEngines", slots={"gdn_frost": EngineSlot(0), "gdn_cutile": EngineSlot(1)}, + analyzer=("cudnn.linear_attention.graph_analyzer", "analyze"), ), EngineFamily( KDA_ID_BASE, @@ -150,6 +151,7 @@ def offered_ids(self) -> Dict[str, int]: "cudnn.linear_attention", "KdaEngines", slots={"kda_frost": EngineSlot(0), "kda_cutile": EngineSlot(1)}, + analyzer=("cudnn.linear_attention.graph_analyzer", "analyze"), ), EngineFamily( GDN2_ID_BASE, @@ -157,6 +159,7 @@ def offered_ids(self) -> Dict[str, int]: "cudnn.linear_attention", "Gdn2Engines", slots={"gdn2_frost": EngineSlot(0)}, + analyzer=("cudnn.linear_attention.graph_analyzer", "analyze"), ), EngineFamily( FROST_GEMM_ID_BASE, diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index 014d12631..499561bb1 100644 --- a/python/cudnn/frost/buffers.py +++ b/python/cudnn/frost/buffers.py @@ -99,6 +99,14 @@ def dtype_name(buf) -> str: return str(buf.dtype).split(".")[-1] +def data_ptr(buf) -> int: + """Device address of a tensor-like (``data_ptr()`` or the CUDA array interface).""" + fn = getattr(buf, "data_ptr", None) + if fn is not None: + return fn() + return buf.__cuda_array_interface__["data"][0] + + class DeviceView: """Zero-copy DLPack view over a raw CUDA pointer. diff --git a/python/cudnn/frost/tile_dsl/mma.py b/python/cudnn/frost/tile_dsl/mma.py index b2bf2da36..5f3d37e33 100644 --- a/python/cudnn/frost/tile_dsl/mma.py +++ b/python/cudnn/frost/tile_dsl/mma.py @@ -479,6 +479,55 @@ def mma_step( acc[s_off + 3] = c3 +@cute.jit +def mma_step_k8( + acc, + a_frag, + b_frag, + *, + k_step: cutlass.Constexpr[int], + M: cutlass.Constexpr[int], + N: cutlass.Constexpr[int], + ab_dtype: cutlass.Constexpr[Type[cutlass.Numeric]] = cutlass.Float16, +): + if cutlass.const_expr(M % 16 != 0): + raise ValueError(f"mma_step_k8: M must be a multiple of 16, got M={M}") + if cutlass.const_expr(ab_dtype != cutlass.Float16 and ab_dtype != cutlass.BFloat16): + raise TypeError(f"mma_step_k8: ab_dtype must be Float16 or BFloat16, got {ab_dtype}") + M_BLOCKS = M // 16 + N_FRAGS = N // 8 + a_stride = len(a_frag) // M_BLOCKS + + ab_tag = "f16" if cutlass.const_expr(ab_dtype == cutlass.Float16) else "bf16" + mma_ptx = f"mma.sync.aligned.m16n8k8.row.col.f32.{ab_tag}.{ab_tag}.f32" " {$0,$1,$2,$3}, {$4,$5}, {$6}, {$7,$8,$9,$10};" + + for m_block in cutlass.range_constexpr(M_BLOCKS): + a_off = m_block * a_stride + k_step * 2 + a0 = a_frag[a_off + 0] + a1 = a_frag[a_off + 1] + acc_base = m_block * N_FRAGS * 4 + for n_frag in cutlass.range_constexpr(N_FRAGS): + b0 = b_frag[n_frag] + s_off = acc_base + n_frag * 4 + c0, c1, c2, c3 = inline_ptx( + mma_ptx, + write_only_types=[cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Float32], + read_only_args=[ + a0, + a1, + b0, + acc[s_off + 0], + acc[s_off + 1], + acc[s_off + 2], + acc[s_off + 3], + ], + ) + acc[s_off + 0] = c0 + acc[s_off + 1] = c1 + acc[s_off + 2] = c2 + acc[s_off + 3] = c3 + + @cute.jit def mma( acc, diff --git a/python/cudnn/frost/tile_dsl/swizzle.py b/python/cudnn/frost/tile_dsl/swizzle.py index e8b16adea..b660b95f5 100644 --- a/python/cudnn/frost/tile_dsl/swizzle.py +++ b/python/cudnn/frost/tile_dsl/swizzle.py @@ -8,11 +8,7 @@ @cute.jit def swizzle_xor_128b(row, col_elem, *, elem_bytes: cutlass.Constexpr[int] = 2): - chunk_elems = 16 // elem_bytes - chunk_idx = col_elem // chunk_elems - in_chunk = col_elem % chunk_elems - swz_chunk = chunk_idx ^ (row & 7) - return swz_chunk * chunk_elems + in_chunk + return col_elem ^ ((row & 7) * cutlass.const_expr(16 // elem_bytes)) @cute.jit diff --git a/python/cudnn/linear_attention/cutile/gdn_engine.py b/python/cudnn/linear_attention/cutile/gdn_engine.py index 6dae5642e..bf74b6da1 100644 --- a/python/cudnn/linear_attention/cutile/gdn_engine.py +++ b/python/cudnn/linear_attention/cutile/gdn_engine.py @@ -1,51 +1,50 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GDN (Gated DeltaNet) execution backend: GDN / GDN_BWD nodes on the -chunked cuTile kernels (``kernels/gdn_chunk_cutile``).""" +"""cuTile GDN engine: GDN / GDN_BWD nodes on the chunked cuTile kernels +(``kernels/gdn_chunk_cutile``).""" -from typing import TYPE_CHECKING, Any, Dict +from typing import TYPE_CHECKING from cudnn import behavior_note from cudnn.engines.base import BaseEngine, CompiledPlan, resolve_node_buffers from cudnn.graph_types import NodeType + from cudnn.frost import buffers from cudnn.frost.workspace import Workspace -from cudnn.linear_attention.engine_utils import _dtype_name +from ..graph_analyzer import check_layouts_compact, analyze, expect_table, to_buffer_dtype -if TYPE_CHECKING: - from cudnn.pygraph import pygraph +# entry base-alignment expectations; ports not listed assume 16 +CUTILE_ALIGN = {"cu_seqlens": 4, "a_log": 4, "beta": 4} -_REQUIRED_PORTS = { - NodeType.GDN: ("q", "k", "v", "g", "beta", "cu_seqlens"), - NodeType.GDN_BWD: ("q", "k", "v", "g", "beta", "cu_seqlens", "dO"), -} +if TYPE_CHECKING: + from cudnn._pygraph import pygraph -def _node_ws_layout(node): +def node_ws_layout(node): """Static carve plan for one node's pipeline intermediates: name -> (offset, dtype-name, shape). The chunk count is data-dependent (varlen), so chunk-indexed entries are sized and SHAPED at the bound ``cdiv(total, 64) + N`` — the device-built table's sentinel tail keeps bound-gridded launches inert past the real count. Terminal pipeline buffers (``o``/``final_state``; the backward's ``dq``/``dk`` finals, - ``wy_dv``, ``dg_cum``, ``db``, ``dh0``) are NOT carved — execute plants + ``wy_dv``, ``dg_cum``, ``db``, ``dstate0``) are NOT carved — execute plants the caller's output buffers under those names.""" - from .kernels.gdn_chunk_cutile import _BT, _cdiv, _next_power_of_2 + from .kernels.gdn_chunk_cutile import BT_CHUNK, cdiv, next_power_of_2 q, v, cu = (node.inputs[p] for p in ("q", "v", "cu_seqlens")) total, H, K = q.dim HV, V = v.dim[1], v.dim[2] N = cu.dim[0] - 1 - io = _dtype_name(q.get_data_type()) + io = to_buffer_dtype(q.get_data_type()) f32 = "float32" - NT_bound = _cdiv(total, _BT) + N + NT_bound = cdiv(total, BT_CHUNK) + N l2norm = bool(node.params.get("use_qk_l2norm", False)) size = 0 table = {} - def add(name, dtype, shape): + def carve(name, dtype, shape): nonlocal size nbytes = buffers.DTYPE_ITEMSIZE[dtype] for s in shape: @@ -53,161 +52,80 @@ def add(name, dtype, shape): table[name] = (size, dtype, tuple(int(s) for s in shape)) size += (nbytes + 127) & ~127 # 128B-aligned sequential carve - add("chunk_table", "int32", (NT_bound, 2)) - add("chunk_count", "int32", (1,)) - add("chunk_offsets", "int32", (N + 1,)) - add("dummy", "int32", (4,)) # inert stub backing for absent optional kernel args - add("g_cum", f32, (total, HV)) - add("A", io, (total, HV, _BT)) - add("w", io, (total, HV, K)) - add("u", io, (total, HV, V)) - add("h", io, (NT_bound, HV, K, V)) - add("v_new", io, (total, HV, V)) + carve("chunk_table", "int32", (NT_bound, 2)) + carve("chunk_count", "int32", (1,)) + carve("chunk_offsets", "int32", (N + 1,)) + carve("dummy", "int32", (4,)) # inert stub backing for absent optional kernel args + carve("g_cum", f32, (total, HV)) + carve("A", io, (total, HV, BT_CHUNK)) + carve("w", io, (total, HV, K)) + carve("u", io, (total, HV, V)) + carve("state_checkpoints", io, (NT_bound, HV, K, V)) + carve("v_new", io, (total, HV, V)) if l2norm: - add("q_norm", io, (total * H, K)) - add("q_rstd", f32, (total * H,)) - add("k_norm", io, (total * H, K)) - add("k_rstd", f32, (total * H,)) + carve("q_norm", io, (total * H, K)) + carve("q_rstd", f32, (total * H,)) + carve("k_norm", io, (total * H, K)) + carve("k_rstd", f32, (total * H,)) if node.node_type == NodeType.GDN_BWD: - add("dv", io, (total, HV, V)) - add("dh", io, (NT_bound, HV, K, V)) - add("dv2", io, (total, HV, V)) - NK = _cdiv(K, min(max(_next_power_of_2(K), 16), 64)) - add("dg_nk", f32, (NK, total, HV)) - add("dw", io, (total, HV, K)) + carve("dv", io, (total, HV, V)) + carve("dstate", io, (NT_bound, HV, K, V)) + carve("dv2", io, (total, HV, V)) + NK = cdiv(K, min(max(next_power_of_2(K), 16), 64)) + carve("dg_nk", f32, (NK, total, HV)) + carve("dw", io, (total, HV, K)) if HV != H or l2norm: - # dq/dk are finals only without l2norm on an MHA config; every + # dQ/dK are finals only without l2norm on an MHA config; every # other combination keeps them (or their head-reduced pair) as # pipeline intermediates - add("dq", io, (total, HV, K)) - add("dk", io, (total, HV, K)) + carve("dq", io, (total, HV, K)) + carve("dk", io, (total, HV, K)) if HV != H: - add("wy_dk_hred", io, (total, H, K)) + carve("wy_dk_hred", io, (total, H, K)) if l2norm: - add("dq_hred", io, (total, H, K)) - add("dk_hred", io, (total, H, K)) - add("dg", f32, (total, HV)) - add("wy_dk", io, (total, HV, K)) - add("wy_dg", f32, (total, HV)) + carve("dq_hred", io, (total, H, K)) + carve("dk_hred", io, (total, H, K)) + carve("dg", f32, (total, HV)) + carve("wy_dk", io, (total, HV, K)) + carve("wy_dg", f32, (total, HV)) return size, table -class _CuTilePlan(CompiledPlan): +class GdnCuTilePlan(CompiledPlan): """Carve plan over the caller's workspace: the layout is static per node; - the buffer arrives with every execute (the explicit-workspace convention).""" + the buffer arrives with every execute.""" - def __init__(self, engine, graph): - self._engine = engine - self._layouts = [(node, *_node_ws_layout(node)) for node in graph.nodes] + def __init__(self, graph): + self.layouts = [(node, *node_ws_layout(node)) for node in graph.nodes] + self.expects = {node: expect_table(node, CUTILE_ALIGN) for node in graph.nodes} # nodes execute sequentially, each re-carving the same buffer - self._ws_bytes = max(nbytes for _, nbytes, _ in self._layouts) + self.ws_bytes = max(nbytes for _, nbytes, _ in self.layouts) def get_workspace_size(self) -> int: - return self._ws_bytes + return self.ws_bytes def execute(self, graph, uid_to_data, ctx) -> None: - self._engine._execute(resolve_node_buffers(graph, uid_to_data), self, ctx) - - -class GdnCuTileEngine(BaseEngine): - """cuTile chunked-kernel backend for single-node GDN graphs (THD layout).""" - - name = "gdn_cutile" - behavior_notes = (behavior_note.RUNTIME_COMPILATION,) # JIT-compiled at build_plans() - - def check_support(self, graph: "pygraph") -> None: - if buffers.current_sm() is None: - raise NotImplementedError("GdnCuTileEngine requires a CUDA device") - try: - from cuda.bindings import runtime as _rt - - err, _cudart_version = _rt.cudaRuntimeGetVersion() - if int(err) != 0: - raise NotImplementedError(f"GdnCuTileEngine: cudaRuntimeGetVersion failed ({err})") - except ImportError as e: - raise NotImplementedError(f"GdnCuTileEngine requires cuda.bindings: {e}") - if _cudart_version < 13030: - raise NotImplementedError(f"GdnCuTileEngine requires CUDA 13.3+ (found {_cudart_version})") - try: - from .kernels.gdn_chunk_cutile import ( # noqa: F401 - chunk_gated_delta_rule, - ) - except ImportError as e: - raise NotImplementedError(f"GdnCuTileEngine requires the cuda.tile runtime: {e}") - - import cudnn - - supported_dtypes = (cudnn.data_type.HALF, cudnn.data_type.BFLOAT16, None) - if not graph.nodes: - raise NotImplementedError("GdnCuTileEngine: empty graph") - for node in graph.nodes: - required = _REQUIRED_PORTS.get(node.node_type) - if required is None: - raise NotImplementedError(f"GdnCuTileEngine only supports GDN/GDN_BWD nodes, got {node.node_type.name}") - for port in required: - if port not in node.inputs: - raise NotImplementedError(f"GdnCuTileEngine: {node.node_type.name} node '{node.name}' is missing input '{port}'") - if int(node.params.get("checkpoint_every_n_tokens", 0) or 0) > 0 or "H" in node.outputs: - raise NotImplementedError("GdnCuTileEngine: per-chunk H output is not supported") - q, k, v = (node.inputs[p] for p in ("q", "k", "v")) - for p in ("q", "k", "v"): - t = node.inputs[p] - if t.get_data_type() not in supported_dtypes: - raise NotImplementedError(f"GdnCuTileEngine: '{p}' must be fp16/bf16, got {t.get_data_type()}") - if t.dim and len(t.dim) != 3: - raise NotImplementedError(f"GdnCuTileEngine: '{p}' must be THD [total_T, heads, dim], got rank {len(t.dim)}") - if q.dim and k.dim and q.dim[1] != k.dim[1]: - raise NotImplementedError(f"GdnCuTileEngine: q and k head counts differ ({q.dim[1]} vs {k.dim[1]})") - if q.dim and v.dim and v.dim[1] % q.dim[1] != 0: - raise NotImplementedError( - f"GdnCuTileEngine: v heads ({v.dim[1]}) must be a multiple of q heads ({q.dim[1]}; GQA-style v broadcast is FROST-only)" - ) - if q.dim and q.dim[-1] > 256: - raise NotImplementedError(f"GdnCuTileEngine: head dim K must be <= 256, got {q.dim[-1]}") - if node.inputs["cu_seqlens"].get_data_type() not in (cudnn.data_type.INT32, None): - raise NotImplementedError("GdnCuTileEngine: cu_seqlens must be int32 (the device-side table builder reads it directly)") - io = q.get_data_type() - f32 = cudnn.data_type.FLOAT - if node.node_type == NodeType.GDN: - out_dtypes = {"O": io, "final_state": f32} - required_out = ("O",) - else: - beta_dt = node.inputs["beta"].get_data_type() - out_dtypes = {"dQ": io, "dK": io, "dV": io, "dG": f32, "dBeta": beta_dt, "d_initial_state": f32} - required_out = ("dQ", "dK", "dV", "dG", "dBeta") - if ("initial_state" in node.inputs) != ("d_initial_state" in node.outputs): - raise NotImplementedError("GdnCuTileEngine: d_initial_state output must be requested iff initial_state is given") - for port in required_out: - if port not in node.outputs: - raise NotImplementedError(f"GdnCuTileEngine: {node.node_type.name} node '{node.name}' is missing output '{port}'") - for port, want in out_dtypes.items(): - t = node.outputs.get(port) - if t is not None and t.get_data_type() not in (want, None): - raise NotImplementedError(f"GdnCuTileEngine: output '{port}' must be {want} (written in place), got {t.get_data_type()}") - - def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: - return _CuTilePlan(self, graph) - - def _execute(self, node_buffers, plan, ctx) -> None: - from .kernels.common import build_chunk_table - from .kernels.gdn_chunk_cutile import _BT - - stream = getattr(ctx, "stream", None) - stream = 0 if stream is None else stream - ws = Workspace(getattr(ctx, "workspace", None), plan._ws_bytes, "GdnCuTileEngine") - for node, _nbytes, table in plan._layouts: + from .kernels.common import build_chunk_table, ensure_cuda_context + from .kernels.gdn_chunk_cutile import BT_CHUNK + + node_buffers = resolve_node_buffers(graph, uid_to_data) + stream = ctx.stream if ctx.stream is not None else 0 + ensure_cuda_context(stream) + ws = Workspace(ctx.workspace, self.ws_bytes, "GdnCuTileEngine") + for node, _nbytes, table in self.layouts: nb = node_buffers[node] + check_layouts_compact("GdnCuTileEngine", self.expects[node], nb) cu_seqlens = nb.inputs["cu_seqlens"] N = node.inputs["cu_seqlens"].dim[0] - 1 bufs = {name: ws.view(off, dt, shape) for name, (off, dt, shape) in table.items()} bound = bufs["chunk_table"].shape[0] - build_chunk_table(bufs["chunk_table"], bufs["chunk_count"], bufs["chunk_offsets"], cu_seqlens, N, _BT, bound, stream=stream) + build_chunk_table(bufs["chunk_table"], bufs["chunk_count"], bufs["chunk_offsets"], cu_seqlens, N, BT_CHUNK, bound, stream=stream) if node.node_type == NodeType.GDN: - self._execute_fwd(node, nb, bufs, stream) + self.execute_fwd(node, nb, bufs, stream) else: - self._execute_bwd(node, nb, bufs, stream) + self.execute_bwd(node, nb, bufs, stream) - def _execute_fwd(self, node, nb, bufs, stream) -> None: + def execute_fwd(self, node, nb, bufs, stream) -> None: from .kernels.gdn_chunk_cutile import chunk_gated_delta_rule_fwd, l2norm_fwd want_state = "final_state" in node.outputs @@ -222,6 +140,9 @@ def _execute_fwd(self, node, nb, bufs, stream) -> None: bufs["o"] = nb.outputs["O"] if want_state: bufs["final_state"] = nb.outputs["final_state"] + gate_kwargs = {} + if node.params.get("safe_gate", False): + gate_kwargs = dict(use_gate_in_kernel=True, A_log=nb.inputs["a_log"], dt_bias=nb.inputs["dt_bias"]) chunk_gated_delta_rule_fwd( q=q, k=k, @@ -230,6 +151,7 @@ def _execute_fwd(self, node, nb, bufs, stream) -> None: beta=beta, scale=scale, initial_state=nb.inputs.get("initial_state"), + **gate_kwargs, output_final_state=want_state, cu_seqlens=nb.inputs["cu_seqlens"], chunk_indices=bufs["chunk_table"], @@ -237,17 +159,17 @@ def _execute_fwd(self, node, nb, bufs, stream) -> None: stream=stream, ) - def _execute_bwd(self, node, nb, bufs, stream) -> None: + def execute_bwd(self, node, nb, bufs, stream) -> None: + from .kernels.common import add_inplace, reshaped from .kernels.gdn_chunk_cutile import ( RCP_LN2, - _BT, + BT_CHUNK, chunk_gated_delta_rule_bwd, chunk_gated_delta_rule_fwd_intra, chunk_local_cumsum, l2norm_bwd, l2norm_fwd, ) - from .kernels.common import add_inplace, reshaped H, K = node.inputs["q"].dim[1], node.inputs["q"].dim[-1] HV = node.inputs["v"].dim[1] @@ -255,7 +177,7 @@ def _execute_bwd(self, node, nb, bufs, stream) -> None: g, beta, do = nb.inputs["g"], nb.inputs["beta"], nb.inputs["dO"] cu_seqlens = nb.inputs["cu_seqlens"] initial_state = nb.inputs.get("initial_state") - dht = nb.inputs.get("d_final_state") + dstate_in = nb.inputs.get("d_final_state") chunk_indices = bufs["chunk_table"] scale = node.params.get("scale") or K**-0.5 l2norm = bool(node.params.get("use_qk_l2norm", False)) @@ -263,8 +185,7 @@ def _execute_bwd(self, node, nb, bufs, stream) -> None: q, q_rstd = l2norm_fwd(q, out=bufs["q_norm"], rstd_out=bufs["q_rstd"], stream=stream) k, k_rstd = l2norm_fwd(k, out=bufs["k_norm"], rstd_out=bufs["k_rstd"], stream=stream) - # terminal pipeline buffers = the caller's output buffers (with - # l2norm, dq/dk stay carves and l2norm_bwd writes the caller's) + # terminal pipeline buffers = the caller's output buffers if not l2norm: bufs["dq" if HV == H else "dq_hred"] = nb.outputs["dQ"] bufs["dk" if HV == H else "dk_hred"] = nb.outputs["dK"] @@ -272,11 +193,10 @@ def _execute_bwd(self, node, nb, bufs, stream) -> None: bufs["dg_cum"] = nb.outputs["dG"] bufs["db"] = nb.outputs["dBeta"] if initial_state is not None: - bufs["dh0"] = nb.outputs["d_initial_state"] + bufs["dstate0"] = nb.outputs["d_initial_state"] # recompute the forward's cumulative gate and intra-chunk WY matrix - # (the backward entry expects them) - g_cum = chunk_local_cumsum(g, chunk_size=_BT, scale=RCP_LN2, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, out=bufs["g_cum"], stream=stream) + g_cum = chunk_local_cumsum(g, chunk_size=BT_CHUNK, scale=RCP_LN2, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, out=bufs["g_cum"], stream=stream) _, _, A = chunk_gated_delta_rule_fwd_intra( k=k, v=v, g=g_cum, beta=beta, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, bufs=bufs, compute_wu=False, stream=stream ) @@ -290,7 +210,7 @@ def _execute_bwd(self, node, nb, bufs, stream) -> None: scale=scale, initial_state=initial_state, do=do, - dht=dht, + dstate_in=dstate_in, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, bufs=bufs, @@ -300,8 +220,93 @@ def _execute_bwd(self, node, nb, bufs, stream) -> None: l2norm_bwd(q, q_rstd, dq, out=nb.outputs["dQ"], bufs=bufs, stream=stream) l2norm_bwd(k, k_rstd, dk, dy2=dk2, out=nb.outputs["dK"], bufs=bufs, stream=stream) else: - # dk/dk2 are the head-reduced finals for GVA, HV-head for MHA + # dK/dK2 are the head-reduced finals for GVA, HV-head for MHA n_dk = 1 - for s_ in dk.shape: - n_dk *= int(s_) + for s in dk.shape: + n_dk *= int(s) add_inplace(reshaped(dk, (n_dk,)), reshaped(dk2, (n_dk,)), n_dk, stream=stream) + + +class GdnCuTileEngine(BaseEngine): + """cuTile chunked-kernel backend for single-node GDN graphs (THD layout).""" + + name = "gdn_cutile" + behavior_notes = (behavior_note.RUNTIME_COMPILATION,) # JIT-compiled + autotuned on first execute per shape + + def check_support(self, graph: "pygraph") -> None: + import cudnn + + if buffers.current_sm() is None: + raise NotImplementedError("GdnCuTileEngine requires a CUDA device") + try: + from cuda.bindings import runtime + + err, cudart_version = runtime.cudaRuntimeGetVersion() + if int(err) != 0: + raise NotImplementedError(f"GdnCuTileEngine: cudaRuntimeGetVersion failed ({err})") + except ImportError as exc: + raise NotImplementedError(f"GdnCuTileEngine requires cuda.bindings: {exc}") from exc + if cudart_version < 13030: + raise NotImplementedError(f"GdnCuTileEngine requires CUDA 13.3+ (found {cudart_version})") + try: + from .kernels.gdn_chunk_cutile import chunk_gated_delta_rule # noqa: F401 — availability probe: ImportError = decline + except ImportError as exc: + raise NotImplementedError(f"GdnCuTileEngine requires the cuda.tile runtime: {exc}") from exc + + facts = graph._facts_for(analyze) + if facts is None or facts.op != "GDN": + raise NotImplementedError("GdnCuTileEngine supports exactly one GDN/GDN_BWD node") + if facts.invalid: + raise NotImplementedError(f"GdnCuTileEngine: {facts.invalid}") + if facts.checkpoint_every_n_tokens > 0 or facts.wants_state_checkpoints: + raise NotImplementedError("GdnCuTileEngine: per-chunk state_checkpoints output is not supported") + if facts.is_bwd and facts.safe_gate: + raise NotImplementedError("GdnCuTileEngine: safe_gate is forward-only") + f32 = cudnn.data_type.FLOAT + for port, got in ( + ("initial_state", facts.state_dtype), + ("final_state", facts.final_state_dtype), + ("d_final_state", facts.d_final_state_dtype), + ("d_initial_state", facts.d_initial_state_dtype), + ("a_log", facts.a_log_dtype), + ("dt_bias", facts.dt_bias_dtype), + ): + if got not in (f32, None): + raise NotImplementedError(f"GdnCuTileEngine: '{port}' must be fp32 (callers convert), got {got}") + if not facts.uniform_io: + raise NotImplementedError("GdnCuTileEngine: q/k/v dtypes must match") + if facts.io_dtype not in (cudnn.data_type.HALF, cudnn.data_type.BFLOAT16, None): + raise NotImplementedError(f"GdnCuTileEngine: q/k/v must be fp16/bf16, got {facts.io_dtype}") + if not facts.thd_layout: + raise NotImplementedError("GdnCuTileEngine: q/k/v must be THD [total_T, heads, dim]") + if facts.h_k != facts.h_q: + raise NotImplementedError(f"GdnCuTileEngine: q and k head counts differ ({facts.h_q} vs {facts.h_k})") + if facts.h_q and facts.h_v % facts.h_q != 0: + raise NotImplementedError( + f"GdnCuTileEngine: v heads ({facts.h_v}) must be a multiple of q heads ({facts.h_q}; GQA-style v broadcast is FROST-only)" + ) + if facts.d_qk > 256: + raise NotImplementedError(f"GdnCuTileEngine: head dim K must be <= 256, got {facts.d_qk}") + if facts.cu_dtype not in (cudnn.data_type.INT32, None): + raise NotImplementedError("GdnCuTileEngine: cu_seqlens must be int32 (the device-side table builder reads it directly)") + io = facts.io_dtype + f32 = cudnn.data_type.FLOAT + if not facts.is_bwd: + out_dtypes = {"O": (facts.o_dtype, io), "final_state": (facts.final_state_dtype, f32)} + else: + out_dtypes = { + "dQ": (facts.dq_dtype, io), + "dK": (facts.dk_dtype, io), + "dV": (facts.dv_dtype, io), + "dG": (facts.dg_dtype, f32), + "dBeta": (facts.dbeta_dtype, facts.beta_dtype), + "d_initial_state": (facts.d_initial_state_dtype, f32), + } + if facts.has_initial_state != facts.wants_d_initial_state: + raise NotImplementedError("GdnCuTileEngine: d_initial_state output must be requested iff initial_state is given") + for port, (got, want) in out_dtypes.items(): + if got is not None and got not in (want, None): + raise NotImplementedError(f"GdnCuTileEngine: output '{port}' must be {want} (written in place), got {got}") + + def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: + return GdnCuTilePlan(graph) diff --git a/python/cudnn/linear_attention/cutile/kda_engine.py b/python/cudnn/linear_attention/cutile/kda_engine.py index 94d96886d..ca5a769f3 100644 --- a/python/cudnn/linear_attention/cutile/kda_engine.py +++ b/python/cudnn/linear_attention/cutile/kda_engine.py @@ -1,269 +1,189 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""KDA (Kimi Delta Attention) execution backend: KDA / KDA_BWD nodes on the -chunked cuTile kernels (``kernels/kda_chunk_cutile``).""" +"""cuTile KDA engine: KDA / KDA_BWD nodes on the chunked cuTile kernels +(``kernels/kda_chunk_cutile``).""" -from typing import TYPE_CHECKING, Any, Dict +from typing import TYPE_CHECKING from cudnn import behavior_note from cudnn.engines.base import BaseEngine, CompiledPlan, resolve_node_buffers from cudnn.graph_types import NodeType + from cudnn.frost import buffers from cudnn.frost.workspace import Workspace -from cudnn.linear_attention.engine_utils import _dtype_name +from ..graph_analyzer import check_layouts_compact, analyze, expect_table, to_buffer_dtype -if TYPE_CHECKING: - from cudnn.pygraph import pygraph +# entry base-alignment expectations; ports not listed assume 16 +CUTILE_ALIGN = {"cu_seqlens": 4, "a_log": 4, "beta": 4} -_REQUIRED_PORTS = { - NodeType.KDA: ("q", "k", "v", "g", "beta", "cu_seqlens"), - NodeType.KDA_BWD: ("q", "k", "v", "g", "beta", "cu_seqlens", "dO"), -} +if TYPE_CHECKING: + from cudnn._pygraph import pygraph -def _node_ws_layout(node): +def node_ws_layout(node): """Static carve plan for one node's ``chunk_kda`` pipeline intermediates: - name -> (offset, dtype, shape). The chunk count is data-dependent - (varlen), so the 'h'/'dh' entries are SIZED with the upper bound - cdiv(total,64)+N and their shape carries None in the NT slot, substituted - at execute. A KDA_BWD node carries the union of the forward re-run's - intermediates and the backward temporaries, so both live disjointly in - one carve. Terminal pipeline buffers (the forward's ``o``/``fs``; the - backward's boundary casts / l2norm outputs, ``dv2`` and ``dh0``) are NOT - carved — execute plants the caller's output buffers under those names.""" - from .kernels.kda_chunk_cutile import _BT as BT, _cdiv, _next_power_of_2 + name -> (offset, dtype-name, shape). The chunk count is data-dependent + (varlen), so chunk-indexed entries are sized and SHAPED at the bound + ``cdiv(total, 64) + N`` — the device-built table's sentinel tail keeps + bound-gridded launches inert past the real count. A KDA_BWD node carries + the union of the forward re-run's intermediates and the backward + temporaries, so both live disjointly in one carve. Terminal pipeline + buffers (the forward's ``o``/``final_state``; the backward's boundary casts / + l2norm outputs, ``dv2`` and ``dstate0``) are NOT carved — execute plants the + caller's output buffers under those names.""" + from .kernels.kda_chunk_cutile import BT_CHUNK, cdiv, next_power_of_2 q, v, g, cu = (node.inputs[p] for p in ("q", "v", "g", "cu_seqlens")) total, H, K = q.dim HV, V = v.dim[1], v.dim[2] N = cu.dim[0] - 1 - io = _dtype_name(q.get_data_type()) + io = to_buffer_dtype(q.get_data_type()) f32 = "float32" BC = 32 if K >= 64 else 16 # fwd_intra sub-chunk (see chunk_kda_fwd_intra) - NK = _cdiv(K, min(64, _next_power_of_2(K))) # bwd_intra K-split (see chunk_kda_bwd_intra) - NT_bound = _cdiv(total, BT) + N + NT_bound = cdiv(total, BT_CHUNK) + N l2norm = bool(node.params.get("use_qk_l2norm", False)) size = 0 table = {} - def add(name, dtype, shape): + def carve(name, dtype, shape): nonlocal size nbytes = buffers.DTYPE_ITEMSIZE[dtype] - shape = tuple(NT_bound if s is None else int(s) for s in shape) for s in shape: - nbytes *= s - table[name] = (size, dtype, shape) + nbytes *= int(s) + table[name] = (size, dtype, tuple(int(s) for s in shape)) size += (nbytes + 127) & ~127 # 128B-aligned sequential carve - add("chunk_table", "int32", (NT_bound, 2)) - add("chunk_count", "int32", (1,)) - add("chunk_offsets", "int32", (N + 1,)) - add("dummy", "int32", (4,)) # inert stub backing for absent optional kernel args + carve("chunk_table", "int32", (NT_bound, 2)) + carve("chunk_count", "int32", (1,)) + carve("chunk_offsets", "int32", (N + 1,)) + carve("dummy", "int32", (4,)) # inert stub backing for absent optional kernel args # chunk_kda forward (also re-run inside KDA_BWD) - add("g_cum", f32, (total, HV, K)) - add("Aqk", io, (total, HV, BT)) - add("Akk", io, (total, HV, BT)) - add("Akkd", f32, (total, HV, BC)) - add("w", io, (total, HV, K)) - add("u", io, (total, HV, V)) - add("qg", io, (total, HV, K)) - add("kg", io, (total, HV, K)) - add("h", io, (None, HV, K, V)) - add("v_new", io, (total, HV, V)) + carve("g_cum", f32, (total, HV, K)) + carve("Aqk", io, (total, HV, BT_CHUNK)) + carve("Akk", io, (total, HV, BT_CHUNK)) + carve("Akkd", f32, (total, HV, BC)) + carve("w", io, (total, HV, K)) + carve("u", io, (total, HV, V)) + carve("qg", io, (total, HV, K)) + carve("kg", io, (total, HV, K)) + carve("state_checkpoints", io, (NT_bound, HV, K, V)) + carve("v_new", io, (total, HV, V)) + if node.params.get("use_beta_sigmoid", False): + carve("beta_sig", f32, (total, HV)) if node.node_type == NodeType.KDA_BWD: - add("o", io, (total, HV, V)) # discarded output of the forward re-run + carve("o", io, (total, HV, V)) # discarded output of the forward re-run if l2norm: - add("q_norm", io, (total * H, K)) - add("q_rstd", f32, (total * H,)) - add("k_norm", io, (total * H, K)) - add("k_rstd", f32, (total * H,)) + carve("q_norm", io, (total * H, K)) + carve("q_rstd", f32, (total * H,)) + carve("k_norm", io, (total * H, K)) + carve("k_rstd", f32, (total * H,)) if node.node_type == NodeType.KDA_BWD: - add("dAqk", f32, (total, HV, BT)) - add("dv_dAv", io, (total, HV, V)) - add("dh", io, (None, HV, K, V)) - add("dv_dhu", io, (total, HV, V)) - add("dq", f32, (total, HV, K)) - add("dk", f32, (total, HV, K)) - add("dg", f32, (total, HV, K)) - if _dtype_name(node.inputs["beta"].get_data_type()) != f32: - add("db", f32, (total, HV)) - add("dAkk", f32, (total, HV, BT)) - add("dq2", f32, (total, HV, K)) - add("dk2", f32, (total, HV, K)) + carve("dAqk", f32, (total, HV, BT_CHUNK)) + carve("dv_dAv", io, (total, HV, V)) + carve("dstate", io, (NT_bound, HV, K, V)) + carve("dv_dstate_u", io, (total, HV, V)) + carve("dq", f32, (total, HV, K)) + carve("dk", f32, (total, HV, K)) + carve("dg", f32, (total, HV, K)) + if to_buffer_dtype(node.inputs["beta"].get_data_type()) != f32: + carve("db", f32, (total, HV)) + carve("dAkk", f32, (total, HV, BT_CHUNK)) + carve("dq2", f32, (total, HV, K)) + carve("dk2", f32, (total, HV, K)) if HV != H: - add("dq_hred", f32, (total, H, K)) - add("dk_hred", f32, (total, H, K)) - add("db2", f32, (NK, total, HV)) - add("dg2", f32, (total, HV, K)) - if _dtype_name(g.get_data_type()) != f32: - add("dg_cum", f32, (total, HV, K)) + carve("dq_hred", f32, (total, H, K)) + carve("dk_hred", f32, (total, H, K)) + NK = cdiv(K, min(64, next_power_of_2(K))) # bwd_intra K-split (see chunk_kda_bwd_intra) + carve("db2", f32, (NK, total, HV)) + carve("dg2", f32, (total, HV, K)) + if to_buffer_dtype(g.get_data_type()) != f32: + carve("dg_cum", f32, (total, HV, K)) return size, table -class _KdaCuTilePlan(CompiledPlan): - """Carve plan over the caller's workspace (see GdnCuTileEngine's - ``_CuTilePlan``): static per-node layout; the buffer arrives with every - execute (the explicit-workspace convention).""" +class KdaCuTilePlan(CompiledPlan): + """Carve plan over the caller's workspace: the layout is static per node; + the buffer arrives with every execute.""" - def __init__(self, engine, graph): - self._engine = engine - self._layouts = [(node, *_node_ws_layout(node)) for node in graph.nodes] + def __init__(self, graph): + self.layouts = [(node, *node_ws_layout(node)) for node in graph.nodes] + self.expects = {node: expect_table(node, CUTILE_ALIGN) for node in graph.nodes} # nodes execute sequentially, each re-carving the same buffer - self._ws_bytes = max(nbytes for _, nbytes, _ in self._layouts) + self.ws_bytes = max(nbytes for _, nbytes, _ in self.layouts) def get_workspace_size(self) -> int: - return self._ws_bytes + return self.ws_bytes def execute(self, graph, uid_to_data, ctx) -> None: - self._engine._execute(resolve_node_buffers(graph, uid_to_data), self, ctx) - - -class KdaCuTileEngine(BaseEngine): - """cuTile chunked-kernel backend for single-node KDA graphs (THD layout).""" - - name = "kda_cutile" - behavior_notes = (behavior_note.RUNTIME_COMPILATION,) # JIT-compiled at build_plans() - - def check_support(self, graph: "pygraph") -> None: - if buffers.current_sm() is None: - raise NotImplementedError("KdaCuTileEngine requires a CUDA device") - try: - from cuda.bindings import runtime as _rt - - err, _cudart_version = _rt.cudaRuntimeGetVersion() - if int(err) != 0: - raise NotImplementedError(f"KdaCuTileEngine: cudaRuntimeGetVersion failed ({err})") - except ImportError as e: - raise NotImplementedError(f"KdaCuTileEngine requires cuda.bindings: {e}") - if _cudart_version < 13030: - raise NotImplementedError(f"KdaCuTileEngine requires CUDA 13.3+ (found {_cudart_version})") - try: - from .kernels.kda_chunk_cutile import ( # noqa: F401 - chunk_kda, - ) - except ImportError as e: - raise NotImplementedError(f"KdaCuTileEngine requires the cuda.tile runtime: {e}") - - import cudnn - - supported_dtypes = (cudnn.data_type.HALF, cudnn.data_type.BFLOAT16, None) - if not graph.nodes: - raise NotImplementedError("KdaCuTileEngine: empty graph") - for node in graph.nodes: - required = _REQUIRED_PORTS.get(node.node_type) - if required is None: - raise NotImplementedError(f"KdaCuTileEngine only supports KDA/KDA_BWD nodes, got {node.node_type.name}") - for port in required: - if port not in node.inputs: - raise NotImplementedError(f"KdaCuTileEngine: {node.node_type.name} node '{node.name}' is missing input '{port}'") - if int(node.params.get("checkpoint_every_n_tokens", 0) or 0) > 0 or "H" in node.outputs: - raise NotImplementedError("KdaCuTileEngine: per-chunk H output is not supported") - q, k, v = (node.inputs[p] for p in ("q", "k", "v")) - for p in ("q", "k", "v"): - t = node.inputs[p] - if t.get_data_type() not in supported_dtypes: - raise NotImplementedError(f"KdaCuTileEngine: '{p}' must be fp16/bf16, got {t.get_data_type()}") - if t.dim and len(t.dim) != 3: - raise NotImplementedError(f"KdaCuTileEngine: '{p}' must be THD [total_T, heads, dim], got rank {len(t.dim)}") - if q.dim and k.dim and q.dim[1] != k.dim[1]: - raise NotImplementedError(f"KdaCuTileEngine: q and k head counts differ ({q.dim[1]} vs {k.dim[1]})") - if q.dim and v.dim and v.dim[1] % q.dim[1] != 0: - raise NotImplementedError( - f"KdaCuTileEngine: v heads ({v.dim[1]}) must be a multiple of q heads ({q.dim[1]}; GQA-style v broadcast is FROST-only)" - ) - if q.dim and q.dim[-1] > 256: - raise NotImplementedError(f"KdaCuTileEngine: head dim K must be <= 256, got {q.dim[-1]}") - if node.inputs["cu_seqlens"].get_data_type() not in (cudnn.data_type.INT32, None): - raise NotImplementedError("KdaCuTileEngine: cu_seqlens must be int32 (the device-side table builder reads it directly)") - io = q.get_data_type() - f32 = cudnn.data_type.FLOAT - if node.node_type == NodeType.KDA: - out_dtypes = {"O": io, "final_state": f32} - required_out = ("O",) - else: - out_dtypes = { - "dQ": io, - "dK": io, - "dV": io, - "dG": node.inputs["g"].get_data_type(), - "dBeta": node.inputs["beta"].get_data_type(), - "d_initial_state": f32, - } - required_out = ("dQ", "dK", "dV", "dG", "dBeta") - if ("initial_state" in node.inputs) != ("d_initial_state" in node.outputs): - raise NotImplementedError("KdaCuTileEngine: d_initial_state output must be requested iff initial_state is given") - for port in required_out: - if port not in node.outputs: - raise NotImplementedError(f"KdaCuTileEngine: {node.node_type.name} node '{node.name}' is missing output '{port}'") - for port, want in out_dtypes.items(): - t = node.outputs.get(port) - if t is not None and t.get_data_type() not in (want, None): - raise NotImplementedError(f"KdaCuTileEngine: output '{port}' must be {want} (written in place), got {t.get_data_type()}") - - def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: - return _KdaCuTilePlan(self, graph) - - def _execute(self, node_buffers, plan, ctx) -> None: - from .kernels.common import build_chunk_table - from .kernels.kda_chunk_cutile import _BT - - stream = getattr(ctx, "stream", None) - stream = 0 if stream is None else stream - ws = Workspace(getattr(ctx, "workspace", None), plan._ws_bytes, "KdaCuTileEngine") - for node, _nbytes, table in plan._layouts: + from .kernels.common import build_chunk_table, ensure_cuda_context + from .kernels.kda_chunk_cutile import BT_CHUNK + + node_buffers = resolve_node_buffers(graph, uid_to_data) + stream = ctx.stream if ctx.stream is not None else 0 + ensure_cuda_context(stream) + ws = Workspace(ctx.workspace, self.ws_bytes, "KdaCuTileEngine") + for node, _nbytes, table in self.layouts: nb = node_buffers[node] + check_layouts_compact("KdaCuTileEngine", self.expects[node], nb) cu_seqlens = nb.inputs["cu_seqlens"] N = node.inputs["cu_seqlens"].dim[0] - 1 bufs = {name: ws.view(off, dt, shape) for name, (off, dt, shape) in table.items()} bound = bufs["chunk_table"].shape[0] - build_chunk_table(bufs["chunk_table"], bufs["chunk_count"], bufs["chunk_offsets"], cu_seqlens, N, _BT, bound, stream=stream) + build_chunk_table(bufs["chunk_table"], bufs["chunk_count"], bufs["chunk_offsets"], cu_seqlens, N, BT_CHUNK, bound, stream=stream) if node.node_type == NodeType.KDA: - self._execute_fwd(node, nb, bufs, stream) + self.execute_fwd(node, nb, bufs, stream) else: - self._execute_bwd(node, nb, bufs, stream) + self.execute_bwd(node, nb, bufs, stream) - @staticmethod - def _state_f32(s0): - # the kernel wants the recurrent state in fp32; callers convert - if s0 is not None and not str(s0.dtype).endswith("float32"): - raise ValueError("KdaCuTileEngine: state ports must be fp32 (callers convert)") - return s0 - - def _execute_fwd(self, node, nb, bufs, stream) -> None: + def execute_fwd(self, node, nb, bufs, stream) -> None: from .kernels.kda_chunk_cutile import chunk_kda want_state = "final_state" in node.outputs + q, k, v = nb.inputs["q"], nb.inputs["k"], nb.inputs["v"] + g, beta = nb.inputs["g"], nb.inputs["beta"] # terminal pipeline buffers = the caller's output buffers bufs["o"] = nb.outputs["O"] if want_state: - bufs["fs"] = nb.outputs["final_state"] + bufs["final_state"] = nb.outputs["final_state"] + raw_gate_kwargs = {} + if node.params.get("use_beta_sigmoid", False): + raw_gate_kwargs["use_beta_sigmoid_in_kernel"] = True + if node.params.get("safe_gate", False): + raw_gate_kwargs.update( + safe_gate=True, + use_gate_in_kernel=True, + lower_bound=float(node.params.get("gate_lower_bound") or -5.0), + A_log=nb.inputs["a_log"], + dt_bias=nb.inputs["dt_bias"], + ) chunk_kda( - nb.inputs["q"], - nb.inputs["k"], - nb.inputs["v"], - nb.inputs["g"], - nb.inputs["beta"], + q, + k, + v, + g, + beta, scale=node.params.get("scale"), - initial_state=self._state_f32(nb.inputs.get("initial_state")), + initial_state=nb.inputs.get("initial_state"), output_final_state=want_state, use_qk_l2norm_in_kernel=bool(node.params.get("use_qk_l2norm", False)), cu_seqlens=nb.inputs["cu_seqlens"], chunk_indices=bufs["chunk_table"], bufs=bufs, stream=stream, + **raw_gate_kwargs, ) - def _execute_bwd(self, node, nb, bufs, stream) -> None: + def execute_bwd(self, node, nb, bufs, stream) -> None: from .kernels.common import reshaped from .kernels.kda_chunk_cutile import chunk_kda_grad total, H, K = node.inputs["q"].dim cu_seqlens = nb.inputs["cu_seqlens"] initial_state = nb.inputs.get("initial_state") - dht = nb.inputs.get("d_final_state") + dstate_in = nb.inputs.get("d_final_state") do, q, k, v = nb.inputs["dO"], nb.inputs["q"], nb.inputs["k"], nb.inputs["v"] g, beta = nb.inputs["g"], nb.inputs["beta"] scale = node.params.get("scale") or K**-0.5 @@ -280,7 +200,7 @@ def _execute_bwd(self, node, nb, bufs, stream) -> None: bufs["dg_cast" if "dg_cum" in bufs else "dg_cum"] = nb.outputs["dG"] bufs["db_cast" if "db" in bufs else "db"] = nb.outputs["dBeta"] if initial_state is not None: - bufs["dh0"] = nb.outputs["d_initial_state"] + bufs["dstate0"] = nb.outputs["d_initial_state"] chunk_kda_grad( q, @@ -289,12 +209,101 @@ def _execute_bwd(self, node, nb, bufs, stream) -> None: g, beta, do, - dht=self._state_f32(dht), + dstate_in=dstate_in, scale=scale, - initial_state=self._state_f32(initial_state), + initial_state=initial_state, use_qk_l2norm_in_kernel=bool(node.params.get("use_qk_l2norm", False)), cu_seqlens=cu_seqlens, chunk_indices=bufs["chunk_table"], bufs=bufs, stream=stream, ) + + +class KdaCuTileEngine(BaseEngine): + """cuTile chunked-kernel backend for single-node KDA graphs (THD layout).""" + + name = "kda_cutile" + behavior_notes = (behavior_note.RUNTIME_COMPILATION,) # JIT-compiled + autotuned on first execute per shape + + def check_support(self, graph: "pygraph") -> None: + import cudnn + + if buffers.current_sm() is None: + raise NotImplementedError("KdaCuTileEngine requires a CUDA device") + try: + from cuda.bindings import runtime + + err, cudart_version = runtime.cudaRuntimeGetVersion() + if int(err) != 0: + raise NotImplementedError(f"KdaCuTileEngine: cudaRuntimeGetVersion failed ({err})") + except ImportError as exc: + raise NotImplementedError(f"KdaCuTileEngine requires cuda.bindings: {exc}") from exc + if cudart_version < 13030: + raise NotImplementedError(f"KdaCuTileEngine requires CUDA 13.3+ (found {cudart_version})") + try: + from .kernels.kda_chunk_cutile import chunk_kda # noqa: F401 — availability probe: ImportError = decline + except ImportError as exc: + raise NotImplementedError(f"KdaCuTileEngine requires the cuda.tile runtime: {exc}") from exc + + facts = graph._facts_for(analyze) + if facts is None or facts.op != "KDA": + raise NotImplementedError("KdaCuTileEngine supports exactly one KDA/KDA_BWD node") + if facts.invalid: + raise NotImplementedError(f"KdaCuTileEngine: {facts.invalid}") + if facts.checkpoint_every_n_tokens > 0 or facts.wants_state_checkpoints: + raise NotImplementedError("KdaCuTileEngine: per-chunk state_checkpoints output is not supported") + if facts.is_bwd and (facts.safe_gate or facts.use_beta_sigmoid): + raise NotImplementedError("KdaCuTileEngine: raw-logit gate modes (safe_gate / use_beta_sigmoid) are forward-only") + f32 = cudnn.data_type.FLOAT + for port, got in ( + ("initial_state", facts.state_dtype), + ("final_state", facts.final_state_dtype), + ("d_final_state", facts.d_final_state_dtype), + ("d_initial_state", facts.d_initial_state_dtype), + ("a_log", facts.a_log_dtype), + ("dt_bias", facts.dt_bias_dtype), + ): + if got not in (f32, None): + raise NotImplementedError(f"KdaCuTileEngine: '{port}' must be fp32 (callers convert), got {got}") + node = next(iter(graph.nodes), None) + glb = node.params.get("gate_lower_bound") if node is not None else None + if glb is not None and glb is not False and not (-5.0 <= float(glb) < 0): + raise NotImplementedError(f"KdaCuTileEngine: gate_lower_bound must be in [-5, 0) (chunk_kda log-gate floor), got {glb}") + if not facts.uniform_io: + raise NotImplementedError("KdaCuTileEngine: q/k/v dtypes must match") + if facts.io_dtype not in (cudnn.data_type.HALF, cudnn.data_type.BFLOAT16, None): + raise NotImplementedError(f"KdaCuTileEngine: q/k/v must be fp16/bf16, got {facts.io_dtype}") + if not facts.thd_layout: + raise NotImplementedError("KdaCuTileEngine: q/k/v must be THD [total_T, heads, dim]") + if facts.h_k != facts.h_q: + raise NotImplementedError(f"KdaCuTileEngine: q and k head counts differ ({facts.h_q} vs {facts.h_k})") + if facts.h_q and facts.h_v % facts.h_q != 0: + raise NotImplementedError( + f"KdaCuTileEngine: v heads ({facts.h_v}) must be a multiple of q heads ({facts.h_q}; GQA-style v broadcast is FROST-only)" + ) + if facts.d_qk > 256: + raise NotImplementedError(f"KdaCuTileEngine: head dim K must be <= 256, got {facts.d_qk}") + if facts.cu_dtype not in (cudnn.data_type.INT32, None): + raise NotImplementedError("KdaCuTileEngine: cu_seqlens must be int32 (the device-side table builder reads it directly)") + io = facts.io_dtype + f32 = cudnn.data_type.FLOAT + if not facts.is_bwd: + out_dtypes = {"O": (facts.o_dtype, io), "final_state": (facts.final_state_dtype, f32)} + else: + out_dtypes = { + "dQ": (facts.dq_dtype, io), + "dK": (facts.dk_dtype, io), + "dV": (facts.dv_dtype, io), + "dG": (facts.dg_dtype, facts.g_dtype), + "dBeta": (facts.dbeta_dtype, facts.beta_dtype), + "d_initial_state": (facts.d_initial_state_dtype, f32), + } + if facts.has_initial_state != facts.wants_d_initial_state: + raise NotImplementedError("KdaCuTileEngine: d_initial_state output must be requested iff initial_state is given") + for port, (got, want) in out_dtypes.items(): + if got is not None and got not in (want, None): + raise NotImplementedError(f"KdaCuTileEngine: output '{port}' must be {want} (written in place), got {got}") + + def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: + return KdaCuTilePlan(graph) diff --git a/python/cudnn/linear_attention/cutile/kernels/common.py b/python/cudnn/linear_attention/cutile/kernels/common.py index 0b1edb32f..d127ba4bb 100644 --- a/python/cudnn/linear_attention/cutile/kernels/common.py +++ b/python/cudnn/linear_attention/cutile/kernels/common.py @@ -33,10 +33,10 @@ ConstInt = ct.Constant[int] -_TILE = 2048 +TILE = 2048 -def _cdiv(a: int, b: int) -> int: +def cdiv(a: int, b: int) -> int: return (a + b - 1) // b @@ -51,20 +51,78 @@ def zero_fill(buf, *, stream) -> None: memset_zero_async(ptr, n, stream) -def reshaped(buf, shape): - """A ``shape``-d view of a contiguous device buffer: native ``reshape`` - when the object has one (torch & co.), else a fresh DLPack view over the - same pointer (the kernels derive their index rank from the array rank).""" - if hasattr(buf, "reshape"): - return buf.reshape(*shape) +def reshaped(buf, target_shape): + """A ``target_shape``-d DeviceView over the same pointer (contiguous by + the engine gate's contract; the kernels derive their index rank from the + array rank).""" from cudnn.frost.buffers import DeviceView, probe - ptr, _shape, _strides, dtype, dev = probe(buf) - return DeviceView(ptr, shape, dtype, dev) + ptr, shape, _strides, dtype, dev = probe(buf) + return DeviceView(ptr, shape, dtype, dev).reshape(tuple(target_shape)) + + +def dummy(dtype_name: str, bufs): + """Inert typed view over the workspace's 16-byte ``dummy`` carve, for + ABSENT optional kernel args (always paired with a flag==0, never + dereferenced). Dtype-bound so the compiled signature stays stable; the + library allocates nothing.""" + from cudnn.frost.buffers import DTYPE_ITEMSIZE, DeviceView + + d = bufs["dummy"] + return DeviceView(d.data_ptr(), (16 // DTYPE_ITEMSIZE[dtype_name],), dtype_name, d.__dlpack_device__()[1]) + + +def opt(t, bufs, dtype_name: str = "float32"): + """Resolve an optional tensor argument to a non-null cuTile launch arg: + the buffer if present (contiguous by the engine contract), else an inert + dummy (paired with a USE_*/HAS_* integer flag). cuTile never accepts None + in launch args, so this is the required dummy-tensor-plus-flag pattern.""" + if t is None: + return dummy(dtype_name, bufs) + return t + + +def dev_id(buf) -> int: + """Device ordinal of a DLPack/CAI buffer.""" + from cudnn.frost.buffers import probe + + return probe(buf)[4] + + +def ensure_cuda_context(stream=0) -> None: + """Make the calling thread's CUDA driver context current. + + cuTile launches and the autotuner's driver-API timing fail on threads + whose driver context stack is empty — e.g. autograd backward worker + threads, where cudaSetDevice alone binds nothing. Prefer the launch + stream's own context; else retain + set-current the current device's + primary context (retained only when no context is bound, so at most once + per thread). Best-effort: never fatal.""" + 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 + from cuda.bindings import runtime as rt + + err_d, dev = rt.cudaGetDevice() + if int(err_d) != 0: + return + err, pctx = drv.cuDevicePrimaryCtxRetain(dev) + if err == drv.CUresult.CUDA_SUCCESS: + drv.cuCtxSetCurrent(pctx) + except Exception: # noqa: BLE001 + pass @ct.kernel -def _add_inplace_kernel(dst, src, TILE: ConstInt): +def add_inplace_kernel(dst, src, TILE: ConstInt): pid = ct.bid(0) a = ct.load(dst, index=(pid,), shape=(TILE,), padding_mode=ct.PaddingMode.ZERO) b = ct.load(src, index=(pid,), shape=(TILE,), padding_mode=ct.PaddingMode.ZERO) @@ -73,11 +131,11 @@ def _add_inplace_kernel(dst, src, TILE: ConstInt): def add_inplace(dst, src, numel: int, *, stream) -> None: """``dst += src`` over ``numel`` flat elements (same dtype, contiguous).""" - ct.launch(stream, (_cdiv(numel, _TILE),), _add_inplace_kernel, (dst, src, _TILE)) + ct.launch(stream, (cdiv(numel, TILE),), add_inplace_kernel, (dst, src, TILE)) @ct.kernel -def _cast_copy_kernel(dst, src, TILE: ConstInt): +def cast_copy_kernel(dst, src, TILE: ConstInt): pid = ct.bid(0) t = ct.load(src, index=(pid,), shape=(TILE,), padding_mode=ct.PaddingMode.ZERO) ct.store(dst, index=(pid,), tile=ct.astype(t, dst.dtype)) @@ -86,11 +144,11 @@ def _cast_copy_kernel(dst, src, TILE: ConstInt): def cast_copy(dst, src, numel: int, *, stream) -> None: """``dst[:] = src`` over ``numel`` flat elements, converting to ``dst``'s dtype (a plain copy when the dtypes already match).""" - ct.launch(stream, (_cdiv(numel, _TILE),), _cast_copy_kernel, (dst, src, _TILE)) + ct.launch(stream, (cdiv(numel, TILE),), cast_copy_kernel, (dst, src, TILE)) @ct.kernel -def _sum_leading_kernel(dst, src, R: ConstInt, ACC: ConstInt, TILE: ConstInt): +def sum_leading_kernel(dst, src, R: ConstInt, ACC: ConstInt, TILE: ConstInt): pid = ct.bid(0) acc = ct.astype(ct.load(src, index=(0, pid), shape=(1, TILE), padding_mode=ct.PaddingMode.ZERO), ct.float32) for r in range(1, R): @@ -106,11 +164,11 @@ def sum_leading(dst, src, r: int, m: int, *, stream, accumulate: bool = False) - axis into ``dst`` (flat ``[m]``), accumulating in fp32. ``r`` is a compile-time constant (small fan-ins: split partials, head groups). ``accumulate`` adds the reduction onto ``dst`` instead of overwriting.""" - ct.launch(stream, (_cdiv(m, _TILE),), _sum_leading_kernel, (dst, src, r, int(accumulate), _TILE)) + ct.launch(stream, (cdiv(m, TILE),), sum_leading_kernel, (dst, src, r, int(accumulate), TILE)) @ct.kernel -def _build_chunk_table_kernel(cu_seqlens, table, count, offsets, N: ConstInt, CS: ConstInt, BOUND: ConstInt): +def build_chunk_table_kernel(cu_seqlens, table, count, offsets, N: ConstInt, CS: ConstInt, BOUND: ConstInt): run = 0 last = N - 1 ct.store(offsets, (0,), 0) @@ -147,11 +205,11 @@ def build_chunk_table(table, count, offsets, cu_seqlens, n_seqs: int, chunk_size chunk grids at ``bound`` unchanged. ``count`` (one int32) receives the real chunk count; ``offsets`` (int32 ``[n_seqs + 1]``) receives the per-sequence chunk prefix (``prepare_chunk_offsets`` semantics).""" - ct.launch(stream, (1,), _build_chunk_table_kernel, (cu_seqlens, reshaped(table, (2 * bound,)), count, offsets, n_seqs, chunk_size, bound)) + ct.launch(stream, (1,), build_chunk_table_kernel, (cu_seqlens, reshaped(table, (2 * bound,)), count, offsets, n_seqs, chunk_size, bound)) @ct.kernel -def _head_group_sum_kernel(dst, src, G: ConstInt, BT: ConstInt, BK: ConstInt): +def head_group_sum_kernel(dst, src, G: ConstInt, BT: ConstInt, BK: ConstInt): t = ct.bid(0) h = ct.bid(1) k = ct.bid(2) @@ -166,4 +224,4 @@ def head_group_sum(dst, src, t: int, h: int, g: int, k: int, *, stream) -> None: ``src`` is a 3-D ``[t, h*g, k]`` buffer, ``dst`` 3-D ``[t, h, k]``; fp32 accumulation, ``g`` consecutive heads per group (compile-time).""" BT, BK = 64, 128 # padded loads + clipped stores absorb ragged t/k - ct.launch(stream, (_cdiv(t, BT), h, _cdiv(k, BK)), _head_group_sum_kernel, (dst, src, g, BT, BK)) + ct.launch(stream, (cdiv(t, BT), h, cdiv(k, BK)), head_group_sum_kernel, (dst, src, g, BT, BK)) diff --git a/python/cudnn/linear_attention/cutile/kernels/gdn_chunk_cutile.py b/python/cudnn/linear_attention/cutile/kernels/gdn_chunk_cutile.py index 827006aa5..a6e2da2f0 100644 --- a/python/cudnn/linear_attention/cutile/kernels/gdn_chunk_cutile.py +++ b/python/cudnn/linear_attention/cutile/kernels/gdn_chunk_cutile.py @@ -20,9 +20,10 @@ from types import SimpleNamespace import cuda.tile as ct +from cuda.tile.tune import exhaustive_search -from .common import add_inplace, head_group_sum, reshaped, sum_leading, zero_fill -from cudnn.frost.buffers import dtype_name as _dtname +from .common import add_inplace, dev_id, dummy, ensure_cuda_context, head_group_sum, opt, reshaped, sum_leading, zero_fill +from cudnn.frost.buffers import dtype_name as dtname logger = logging.getLogger(__name__) @@ -31,36 +32,31 @@ RCP_LN2 = 1.4426950216 # 1/ln(2) # chunk size (BT tile) of these kernels; the engine's carve layout imports it -_BT = 64 +BT_CHUNK = 64 # Host-side utilities -def _cdiv(a: int, b: int) -> int: +def cdiv(a: int, b: int) -> int: return (a + b - 1) // b -def _next_power_of_2(n: int) -> int: +def next_power_of_2(n: int) -> int: return 1 << (n - 1).bit_length() # Launch-hint autotuning (occupancy x num_worker_warps) for hot kernels. -# These large unrolled tiles emit 255-reg kernels defaulting to occupancy=1, which -# starves the SM; a higher occupancy hint budgets fewer regs/thread -> more blocks/SM. -# Tuned once per (kernel, grid-shape) key and cached. -from cuda.tile.tune import exhaustive_search # noqa: E402 -_LAUNCH_HINT_CACHE: dict = {} +LAUNCH_HINT_CACHE: dict = {} -# occ=4/nww=4 = consistent winner; nww=8 and default kept as fallbacks. # NOTE: the upstream occupancy=2 config is EXCLUDED: on tileiras 13.2 (which # ignores num_worker_warps hints) a freshly compiled occ=2 # chunk_bwd_kernel_dqkwg deadlocks on its first launch (deterministic; # occ=1/4 are fine, and occ=2 is fine on the dhu/kkt kernels). Restore the # SimpleNamespace(occupancy=2, num_worker_warps=4) entry once the runtime is # on tileiras >= 13.3. -_LAUNCH_HINT_CONFIGS = [ +LAUNCH_HINT_CONFIGS = [ SimpleNamespace(occupancy=4, num_worker_warps=4), SimpleNamespace(occupancy=4, num_worker_warps=8), SimpleNamespace(occupancy=1, num_worker_warps=8), @@ -68,47 +64,23 @@ def _next_power_of_2(n: int) -> int: ] -def _ensure_cuda_context(): - """Make the calling thread's CUDA driver context current. - - exhaustive_search times configs through the driver API, which fails on threads - with an empty driver context stack (e.g. autograd backward worker threads). - Retain + set-current the device primary context. Best-effort: never fatal.""" - try: - from cuda.bindings import driver as _drv - - err, cur = _drv.cuCtxGetCurrent() - if err == _drv.CUresult.CUDA_SUCCESS and int(cur) != 0: - return - from cuda.bindings import runtime as _rt - - err_d, dev = _rt.cudaGetDevice() - if int(err_d) != 0: - return - err, ctx = _drv.cuDevicePrimaryCtxRetain(dev) - if err == _drv.CUresult.CUDA_SUCCESS: - _drv.cuCtxSetCurrent(ctx) - except Exception: # noqa: BLE001 - pass - - -def _device_attrs(): +def device_attrs(): """(sm_count, cc_major) of the current device via the runtime API (which auto-inits the primary context, so this works before any framework call).""" try: - from cuda.bindings import runtime as _rt + from cuda.bindings import runtime as rt - err, dev = _rt.cudaGetDevice() + err, dev = rt.cudaGetDevice() if int(err) != 0: return 0, 0 - err, sm = _rt.cudaDeviceGetAttribute(_rt.cudaDeviceAttr.cudaDevAttrMultiProcessorCount, dev) - err2, major = _rt.cudaDeviceGetAttribute(_rt.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, dev) + err, sm = rt.cudaDeviceGetAttribute(rt.cudaDeviceAttr.cudaDevAttrMultiProcessorCount, dev) + err2, major = rt.cudaDeviceGetAttribute(rt.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, dev) return (int(sm) if int(err) == 0 else 0), (int(major) if int(err2) == 0 else 0) except Exception: # noqa: BLE001 return 0, 0 -def _tuned_launch(kernel, stream, grid, args, cache_key, configs=None): +def tuned_launch(kernel, stream, grid, args, cache_key, configs=None): """Launch ``kernel`` with autotuned (occupancy, num_worker_warps) hints. Tunes once per ``cache_key`` via exhaustive_search, caches the specialized @@ -118,13 +90,13 @@ def _tuned_launch(kernel, stream, grid, args, cache_key, configs=None): """ kname = getattr(getattr(kernel, "_pyfunc", None), "__name__", repr(kernel)) key = (kname, cache_key) - if key not in _LAUNCH_HINT_CACHE: + if key not in LAUNCH_HINT_CACHE: tuned = None - _ensure_cuda_context() + ensure_cuda_context() try: with ct.compiler_timeout(20): result = exhaustive_search( - list(configs if configs is not None else _LAUNCH_HINT_CONFIGS), + list(configs if configs is not None else LAUNCH_HINT_CONFIGS), stream, lambda cfg: grid, kernel, @@ -136,8 +108,8 @@ def _tuned_launch(kernel, stream, grid, args, cache_key, configs=None): except Exception as e: # noqa: BLE001 logger.warning("launch-hint autotune failed for %s: %s; using default hints", kname, e) tuned = kernel - _LAUNCH_HINT_CACHE[key] = tuned - ct.launch(stream, grid, _LAUNCH_HINT_CACHE[key], args) + LAUNCH_HINT_CACHE[key] = tuned + ct.launch(stream, grid, LAUNCH_HINT_CACHE[key], args) def exp(x): @@ -153,21 +125,21 @@ def softplus(x): return ct.where(x < 20.0, ct.log(1.0 + ct.exp(x)), x) -def _tf32(a): +def tf32(a): """ct.mma/ct.matmul do not auto-cast fp32 operands to tf32; cast explicitly (allow-tf32 matmul semantics).""" return ct.astype(a, ct.tfloat32) if a.dtype == ct.float32 else a def safe_matmul(a, b): - return ct.matmul(_tf32(a), _tf32(b)) + return ct.matmul(tf32(a), tf32(b)) safe_dot = safe_matmul # alias: non-accumulating dot with fp32->tf32 guard def safe_mma(a, b, acc): - return ct.mma(_tf32(a), _tf32(b), acc) + return ct.mma(tf32(a), tf32(b), acc) def gather_flat(raw, off, *, mask=None, padding_value=0, check_bounds=False): @@ -178,7 +150,7 @@ def scatter_flat(raw, off, value, *, mask=None, check_bounds=False): raw.store_offset(off, value, mask=mask) -def _array_numel(arr): +def array_numel(arr): # Total element count; rank statically unrolled (<=5) since a Python for over # arr.shape is captured as a device loop and rejected. K<=256 bounds the rank. s = arr.shape @@ -196,7 +168,7 @@ def _array_numel(arr): # Bounds-checked flat access: 1-D bounds mask `0 <= off < numel` AND-ed with any -# caller mask. Pass `numel = _array_numel(arr)`. +# caller mask. Pass `numel = array_numel(arr)`. def gather_flat_cb(raw, off, numel, *, mask=None, padding_value=0): inb = (off >= 0) & (off < numel) m = inb if mask is None else (mask & inb) @@ -247,7 +219,7 @@ def l2norm_bwd_kernel1(y, rstd, dy, dy2, dx, eps, D, BD: ConstInt, HAS_DY2: Cons @ct.kernel def l2norm_fwd_kernel(x, y, rstd, eps, T, D: ConstInt, BD: ConstInt, BT: ConstInt): # D <= 512 path: BT rows per block, BD power-of-2 cols. Block-aligned -> - # ct.load with block index + ZERO padding (faithful to make_block_ptr). + # ct.load with block index + ZERO padding. i_t = ct.bid(0) b_x = ct.astype(ct.load(x, index=(i_t, 0), shape=(BT, BD), padding_mode=ct.PaddingMode.ZERO), ct.float32) b_rstd = ct.rsqrt(ct.sum(b_x * b_x, axis=1) + eps) @@ -413,7 +385,7 @@ def gdn_gate_bwd_kernel(g, A_log, dt_bias, dyg, dg, dA, db, T, H: ConstInt, BT: ct.scatter(db, i_t * H + i_h, b_db) -# chunk_gated_delta_rule_fwd_kkt_solve_kernel (fused kkt + solve, 4 sub-chunks) +# chunk_gated_delta_rule_fwd_kkt_solve_kernel (fused KK^T + solve, 4 sub-chunks) @ct.kernel def chunk_gated_delta_rule_fwd_kkt_solve_kernel( k, @@ -430,7 +402,7 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( BK: ConstInt, USE_G: ConstInt, ): - """Fused: beta * K @ K^T (lower triangular) + solve_tril (I+A)^{-1} in one pass.""" + """Fused: Beta * K @ K^T (lower triangular) + solve_tril (I+A)^{-1} in one pass.""" i_t = ct.bid(0) i_h = ct.bid(1) @@ -568,7 +540,7 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( b_A31 = ct.mma(b_k3, ct.transpose(b_k1), b_A31) b_A32 = ct.mma(b_k3, ct.transpose(b_k2), b_A32) - # -- Step 2: gate + beta scaling -- + # -- Step 2: Gate + Beta scaling -- m_d = ct.expand_dims(o_i, 1) > ct.expand_dims(o_i, 0) m_I = ct.expand_dims(o_i, 1) == ct.expand_dims(o_i, 0) @@ -821,13 +793,12 @@ def recompute_w_u_fwd_kernel( beta_seg = beta.slice(axis=0, start=bos, stop=eos) Z = ct.PaddingMode.ZERO - # Keep beta native bf16: scaling in bf16 avoids 2 extra ftof + # Keep Beta native bf16: scaling in bf16 avoids 2 extra ftof # converts/MMA-input that an f32 cast would force. b_b = ct.load(beta_seg, index=(i_t_loc, i_h), shape=(BT, 1), padding_mode=Z).reshape((BT,)) b_A = ct.load(A_seg, index=(i_t_loc, i_h, 0), shape=(BT, 1, BT), padding_mode=Z).reshape((BT, BT)) - # u = A @ (v * beta). latency=3 -> deeper pipeline of the per-tile TMA - # loads; small win over latency=2, no spill. + # U = A @ (V * Beta). latency=3 -> deeper pipeline of the per-tile TMA loads. for i_v in range(ct.cdiv(V, BV)): b_v = ct.load(v_seg, index=(i_t_loc, i_h, i_v), shape=(BT, 1, BV), padding_mode=Z, latency=3).reshape((BT, BV)) # bf16 * bf16 -> bf16. @@ -844,7 +815,7 @@ def recompute_w_u_fwd_kernel( ) b_g = exp2(b_g_val) - # w = A @ (k * beta * g) + # W = A @ (K * Beta * Gate) for i_k in range(ct.cdiv(K, BK)): b_k = ct.load(k_seg, index=(i_t_loc, i_kh, i_k), shape=(BT, 1, BK), padding_mode=Z, latency=3).reshape((BT, BK)) # bf16 * bf16 -> bf16; only the g-decay path promotes to f32. @@ -1122,7 +1093,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( o_bv = ct.arange(BV, dtype=ct.int32) m_v_lane = (i_v * BV + o_bv) < V - # --- Load initial state h0 -> b_h1..b_h4 --- + # --- Load initial state H0 -> b_h1..b_h4 --- if USE_INITIAL_STATE: if STATE_V_FIRST: row = (i_v * BV + o_bv)[:, None] @@ -1255,7 +1226,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( check_bounds=True, ) - # --- b_v = b_w @ b_h (accumulated over K blocks) --- + # --- V = W @ H (accumulated over K blocks) --- w_time = (t_base + i_t * BT + o_bt)[:, None] wmask_r = (i_t * BT + o_bt) < T bw1 = ct.gather( @@ -1409,7 +1380,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( b_v = ct.astype(b_v, k.dtype) - # --- b_h += b_k @ b_v (b_k is (64, BT) transposed load of k) --- + # --- H += K^T @ V (K loaded transposed as (64, BT)) --- k_time = (t_base + i_t * BT + o_bt)[None, :] kmask = ((0 + o64)[:, None] < K) & ((i_t * BT + o_bt)[None, :] < T) bk1 = ct.gather( @@ -1476,7 +1447,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( else: b_h4 = b_h4 + prod4 - # --- Store final state ht <- b_h* (ht (N, HV, *), per-sequence state) --- + # --- Store final state Ht <- b_h* (Ht (N, HV, *), per-sequence state) --- if STORE_FINAL_STATE: if STATE_V_FIRST: row = (i_v * BV + o_bv)[:, None] @@ -1513,7 +1484,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( w, g, gk, - dht, + dstate_in, dh0, do, dh, @@ -1570,34 +1541,34 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( dht_base = i_nh * K * V # Raw flat handles + element counts for 1-D flat gather/scatter (cuTile 1.4.0). - dhtf = dht.get_raw_memory() - dht_n = _array_numel(dht) + dhtf = dstate_in.get_raw_memory() + dht_n = array_numel(dstate_in) dhf = dh.get_raw_memory() - dh_n = _array_numel(dh) + dh_n = array_numel(dh) dh0f = dh0.get_raw_memory() - dh0_n = _array_numel(dh0) + dh0_n = array_numel(dh0) gf = g.get_raw_memory() - g_n = _array_numel(g) + g_n = array_numel(g) gkf = gk.get_raw_memory() - gk_n = _array_numel(gk) + gk_n = array_numel(gk) dof = do.get_raw_memory() - do_n = _array_numel(do) + do_n = array_numel(do) kf = k.get_raw_memory() - k_n = _array_numel(k) + k_n = array_numel(k) qf = q.get_raw_memory() - q_n = _array_numel(q) + q_n = array_numel(q) wf = w.get_raw_memory() - w_n = _array_numel(w) + w_n = array_numel(w) dvf = dv.get_raw_memory() - dv_n = _array_numel(dv) + dv_n = array_numel(dv) dv2f = dv2.get_raw_memory() - dv2_n = _array_numel(dv2) + dv2_n = array_numel(dv2) o64 = ct.arange(64, dtype=ct.int32) o_bt = ct.arange(BT, dtype=ct.int32) o_bv = ct.arange(BV, dtype=ct.int32) - # --- Load final-state gradient dht -> b_dh* --- + # --- Load final-state gradient dHt -> b_dh* --- if USE_FINAL_STATE_GRADIENT: if STATE_V_FIRST: row = (i_v * BV + o_bv)[:, None] @@ -1716,7 +1687,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( b_do = gather_flat_cb(dof, do_off, do_n, padding_value=0.0) b_do = ct.where(do_full, b_do, ct.zeros((BT, BV), dtype=b_do.dtype)) - # b_dv = sum_n b_k_n @ b_dh_n + # dV = sum_n K_n @ dH_n k_row = (i_t * BT + o_bt)[:, None] kmask_r = (i_t * BT + o_bt) < T bk1 = gather_flat_cb(kf, k_base + k_row * (H * K) + (0 + o64)[None, :], k_n, padding_value=0.0) @@ -1796,7 +1767,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( dv2_off = ct.where(do_full, dv2_off, dv2_oob) scatter_flat_cb(dv2f, dv2_off, ct.astype(b_dv, dv2.dtype), dv2_n) - # --- b_dh += trans?(q@do*scale - w@dv) per K block --- + # --- dH += trans?(Q @ dO * scale - W @ dV) per K block --- time = (i_t * BT + o_bt)[None, :] tmask = (i_t * BT + o_bt)[None, :] < T b_dv_cast = ct.astype(b_dv, ct.float32) @@ -1901,7 +1872,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( else: b_dh4 = b_dh4 + upd4 - # --- Store initial-state gradient dh0 <- b_dh* --- + # --- Store initial-state gradient dH0 <- b_dh* --- if USE_INITIAL_STATE: # Mask both the 64-wide K-block rows and BV-wide V cols; flat numel check # alone would let an over-range K row spill into the next head's dh0 slot. @@ -2028,16 +1999,15 @@ def chunk_fwd_kernel_o( b_A = ct.zeros((BT, BT), dtype=ct.float32) for i_k in range(ct.cdiv(K, BK)): - # latency=2 pipelining on this K-loop's TMA loads. b_q = ct.load(q_seg, index=(i_t, i_kh, i_k), shape=(BT, 1, BK), padding_mode=Z, latency=2).reshape((BT, BK)) b_k = ct.load(k_seg, index=(i_t, i_kh, i_k), shape=(BT, 1, BK), padding_mode=Z, latency=2).reshape((BT, BK)) if STATE_V_FIRST: - # b_h (BV, BK), o += q @ h^T + # b_h (BV, BK), O += Q @ H^T b_h = ct.load(h, index=(i_hc, i_v, i_k), shape=(1, BV, BK), padding_mode=Z, latency=2).reshape((BV, BK)) b_o = safe_mma(b_q, ct.transpose(b_h), b_o) else: - # b_h (BK, BV), o += q @ h + # b_h (BK, BV), O += Q @ H b_h = ct.load(h, index=(i_hc, i_k, i_v), shape=(1, BK, BV), padding_mode=Z, latency=2).reshape((BK, BV)) b_o = safe_mma(b_q, b_h, b_o) b_A = safe_mma(b_q, ct.transpose(b_k), b_A) @@ -2094,7 +2064,7 @@ def chunk_bwd_kernel_dqkwg( NV, # runtime cdiv(V,BV) -> rolled V-loop (avoids unroll reg-spill at K=256) ): # q/k/v/do/dv/dq/dk are (T,*,D) token-packed slabs, h/dh (NT_total*HV,*,*), - # dg (NK,T,HV); block-indexed ct.load -> TMA. latency=2. + # dg (NK,T,HV); block-indexed ct.load -> TMA. i_k = ct.bid(0) i_t = ct.bid(1) i_h = ct.bid(2) @@ -2331,30 +2301,6 @@ def chunk_bwd_kernel_dv_local( # an integer flag (USE_*/HAS_*). -def _dummy(dtype_name: str, bufs): - """Inert typed view over the workspace's 16-byte ``dummy`` carve, for - ABSENT optional kernel args (always paired with a flag==0, never - dereferenced). Dtype-bound so the compiled signature stays stable; the - library allocates nothing.""" - from cudnn.frost.buffers import DTYPE_ITEMSIZE, DeviceView - - d = bufs["dummy"] - return DeviceView(d.data_ptr(), (16 // DTYPE_ITEMSIZE[dtype_name],), dtype_name, d.__dlpack_device__()[1]) - - -def _i32(t): - if not str(t.dtype).endswith("int32"): - raise TypeError(f"index/boundary buffers must be int32 (callers convert), got {t.dtype}") - return t - - -def _opt(t, bufs, dtype_name: str = "float32"): - """The buffer if present (contiguous by the engine contract), else an inert dummy (paired with a flag).""" - if t is None: - return _dummy(dtype_name, bufs) - return t - - # explicit workspace: all device memory is the caller's — the engine plan (or a # standalone harness) pre-carves every pipeline intermediate as a named view in # ``bufs`` and passes outputs via ``out=``; nothing here allocates. @@ -2363,24 +2309,17 @@ def _opt(t, bufs, dtype_name: str = "float32"): # l2norm def l2norm_fwd(x, eps: float = 1e-6, out=None, rstd_out=None, stream=None): stream = 0 if stream is None else stream - if out is None or rstd_out is None: - raise ValueError("l2norm_fwd requires pre-allocated out= and rstd_out= buffers") x_shape_og = x.shape - x = x.reshape(-1, x.shape[-1]).contiguous() + x = reshaped(x, (-1, x.shape[-1])) y = reshaped(out, tuple(x.shape)) T, D = x.shape[0], x.shape[-1] MAX_FUSED_SIZE = 65536 // x.element_size() - BD = min(MAX_FUSED_SIZE, _next_power_of_2(D)) - if D > BD: - raise RuntimeError("This layer doesn't support feature dim >= 64KB.") + BD = min(MAX_FUSED_SIZE, next_power_of_2(D)) rstd = reshaped(rstd_out, (T,)) if D <= 512: BT = 32 - grid = (_cdiv(T, BT),) - # Memory-bound row reduction: default occupancy=1 starves the SM (the - # grid is huge). Autotune (occupancy, num_worker_warps) so the scheduler - # packs enough blocks/SM to hide DRAM latency. - _tuned_launch( + grid = (cdiv(T, BT),) + tuned_launch( l2norm_fwd_kernel, stream, grid, @@ -2404,26 +2343,18 @@ def l2norm_bwd( ): stream = 0 if stream is None else stream y_shape_og = y.shape - y = y.reshape(-1, dy.shape[-1]).contiguous() - dy = dy.reshape(-1, dy.shape[-1]).contiguous() - dy2_arg = dy2.reshape(-1, dy.shape[-1]).contiguous() if dy2 is not None else _dummy(_dtname(dy), bufs) - if out is None: - raise ValueError("l2norm_bwd requires a pre-allocated out= buffer") + y = y.reshape(-1, dy.shape[-1]) + dy = dy.reshape(-1, dy.shape[-1]) + dy2_arg = dy2.reshape(-1, dy.shape[-1]) if dy2 is not None else dummy(dtname(dy), bufs) dx = reshaped(out, tuple(y.shape)) T, D = y.shape[0], y.shape[-1] MAX_FUSED_SIZE = 65536 // y.element_size() - BD = min(MAX_FUSED_SIZE, _next_power_of_2(D)) - if D > BD: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - rstd_flat = rstd.reshape(-1).contiguous() + BD = min(MAX_FUSED_SIZE, next_power_of_2(D)) + rstd_flat = rstd.reshape(-1) if D <= 512: BT = 32 - grid = (_cdiv(T, BT),) - # Memory-bound row reduction (mirrors l2norm_fwd): default occupancy=1 - # starves the SM on the huge row grid. Autotune (occupancy, nww) so the - # scheduler packs enough blocks/SM to hide DRAM latency. This kernel was - # the largest non-dot backward kernel. - _tuned_launch( + grid = (cdiv(T, BT),) + tuned_launch( l2norm_bwd_kernel, stream, grid, @@ -2441,37 +2372,33 @@ def l2norm_bwd( # fused_beta_sigmoid -_BETA_SIGMOID_BLOCK_SIZE = 2048 +BETA_SIGMOID_BLOCK_SIZE = 2048 def fused_beta_sigmoid_fwd(x, scale: float = 1.0, out=None, stream=None): stream = 0 if stream is None else stream - if out is None: - raise ValueError("fused_beta_sigmoid_fwd requires a pre-allocated out= buffer (fp32, x-shaped)") y = reshaped(out, tuple(x.shape)) n = x.numel() - grid = (_cdiv(n, _BETA_SIGMOID_BLOCK_SIZE),) + grid = (cdiv(n, BETA_SIGMOID_BLOCK_SIZE),) ct.launch( stream, grid, fused_beta_sigmoid_fwd_kernel, - (x.reshape(-1), y.reshape(-1), float(scale), n, _BETA_SIGMOID_BLOCK_SIZE), + (reshaped(x, (-1,)), y.reshape(-1), float(scale), n, BETA_SIGMOID_BLOCK_SIZE), ) return y def fused_beta_sigmoid_bwd(x, dy, scale: float = 1.0, out=None, stream=None): stream = 0 if stream is None else stream - if out is None: - raise ValueError("fused_beta_sigmoid_bwd requires a pre-allocated out= buffer (x-shaped)") dx = reshaped(out, tuple(x.shape)) n = x.numel() - grid = (_cdiv(n, _BETA_SIGMOID_BLOCK_SIZE),) + grid = (cdiv(n, BETA_SIGMOID_BLOCK_SIZE),) ct.launch( stream, grid, fused_beta_sigmoid_bwd_kernel, - (x.reshape(-1), dy.reshape(-1), dx.reshape(-1), float(scale), n, _BETA_SIGMOID_BLOCK_SIZE), + (reshaped(x, (-1,)), reshaped(dy, (-1,)), dx.reshape(-1), float(scale), n, BETA_SIGMOID_BLOCK_SIZE), ) return dx @@ -2496,16 +2423,12 @@ def chunk_local_cumsum_scalar( stream = 0 if stream is None else stream T, H = g.shape BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") - if out is None: - raise ValueError("chunk_local_cumsum requires a pre-allocated out= buffer (g-shaped)") NT = len(chunk_indices) g_out = reshaped(out, (T, H)) scale_val = float(scale) if scale is not None else 0.0 has_scale = int(scale is not None) - cu_arg = _i32(cu_seqlens) - ci_arg = _i32(chunk_indices) + cu_arg = cu_seqlens + ci_arg = chunk_indices ct.launch( stream, (NT, H), @@ -2536,7 +2459,6 @@ def chunk_local_cumsum( stream=None, ): stream = 0 if stream is None else stream - assert len(g.shape) == 2, f"Unsupported input shape {g.shape}, expected (T, H)." return chunk_local_cumsum_scalar( g=g, chunk_size=chunk_size, @@ -2564,23 +2486,19 @@ def gdn_gate_chunk_cumsum( stream = 0 if stream is None else stream T, H = g.shape BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") - if out is None: - raise ValueError("gdn_gate_chunk_cumsum requires a pre-allocated out= buffer (fp32, g-shaped)") NT = len(chunk_indices) o = reshaped(out, (T, H)) - dt_arg = _opt(dt_bias, bufs, _dtname(A_log)).reshape(-1) + dt_arg = reshaped(opt(dt_bias, bufs, dtname(A_log)), (-1,)) scale_val = float(scale) if scale is not None else 0.0 - cu_arg = _i32(cu_seqlens) - ci_arg = _i32(chunk_indices) + cu_arg = cu_seqlens + ci_arg = chunk_indices ct.launch( stream, (NT, H), gdn_gate_chunk_cumsum_scalar_kernel, ( g, - A_log.reshape(-1), + reshaped(A_log, (-1,)), dt_arg, o, scale_val, @@ -2604,23 +2522,21 @@ def gdn_gate_bwd(g, A_log, dt_bias, dyg, dg_out=None, dA_out=None, dbias_out=Non H = g.shape[-1] T = g.numel() // H BT = 32 - NT = _cdiv(T, BT) - if dg_out is None or dA_out is None or (dt_bias is not None and dbias_out is None): - raise ValueError("gdn_gate_bwd requires pre-allocated dg_out=, dA_out= (and dbias_out= with dt_bias)") + NT = cdiv(T, BT) dg = reshaped(dg_out, tuple(g.shape)) dA_nt = reshaped(bufs["dA_gate"], (NT, H)) db_nt = reshaped(bufs["db_gate"], (NT, H)) if dt_bias is not None else None - dt_arg = _opt(dt_bias, bufs, _dtname(A_log)).reshape(-1) - db_arg = db_nt.reshape(-1) if db_nt is not None else _dummy("float32", bufs) + dt_arg = reshaped(opt(dt_bias, bufs, dtname(A_log)), (-1,)) + db_arg = db_nt.reshape(-1) if db_nt is not None else dummy("float32", bufs) ct.launch( stream, (NT, H), gdn_gate_bwd_kernel, ( - g.reshape(-1), - A_log.reshape(-1), + reshaped(g, (-1,)), + reshaped(A_log, (-1,)), dt_arg, - dyg.reshape(-1), + reshaped(dyg, (-1,)), dg.reshape(-1), dA_nt.reshape(-1), db_arg, @@ -2643,17 +2559,15 @@ def recompute_w_u_fwd(k, v, beta, A, g=None, cu_seqlens=None, chunk_indices=None BT = A.shape[-1] BK = 64 BV = 64 - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") NT = len(chunk_indices) w = reshaped(bufs["w"], (T, HV, K)) u = reshaped(bufs["u"], (T, HV, V)) - beta2 = beta.reshape(T, HV) + beta2 = reshaped(beta, (T, HV)) A3 = A.reshape(T, HV, BT) - g_arg = g.reshape(T, HV) if g is not None else _dummy("float32", bufs) - cu_arg = _i32(cu_seqlens) - ci_arg = _i32(chunk_indices) - _tuned_launch( + g_arg = g.reshape(T, HV) if g is not None else dummy("float32", bufs) + cu_arg = cu_seqlens + ci_arg = chunk_indices + tuned_launch( recompute_w_u_fwd_kernel, stream, (NT, HV), @@ -2676,7 +2590,7 @@ def recompute_w_u_fwd(k, v, beta, A, g=None, cu_seqlens=None, chunk_indices=None BV, int(g is not None), ), - cache_key=((NT, HV), K, V, BT, BK, BV, int(g is not None), str(k.dtype)), + cache_key=(H, HV, K, V, BT, BK, BV, int(g is not None), str(k.dtype), dev_id(k)), ) return w, u @@ -2685,21 +2599,19 @@ def prepare_wy_repr_bwd(k, v, beta, A, dw, du, g=None, cu_seqlens=None, chunk_in stream = 0 if stream is None else stream T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[1] BT = 64 - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") NT = len(chunk_indices) CONST_TILING = 64 - BK = min(max(_next_power_of_2(K), 16), CONST_TILING) - BV = min(max(_next_power_of_2(V), 16), CONST_TILING) + BK = min(max(next_power_of_2(K), 16), CONST_TILING) + BV = min(max(next_power_of_2(V), 16), CONST_TILING) dk = reshaped(bufs["wy_dk"], (T, HV, K)) dv = reshaped(bufs["wy_dv"], (T, HV, V)) dg = reshaped(bufs["wy_dg"], (T, HV)) if g is not None else None db = reshaped(bufs["db"], (T, HV)) - g_arg = _opt(g, bufs) - dg_arg = dg if dg is not None else _dummy("float32", bufs) - cu_arg = _i32(cu_seqlens) - ci_arg = _i32(chunk_indices) - _tuned_launch( + g_arg = opt(g, bufs) + dg_arg = dg if dg is not None else dummy("float32", bufs) + cu_arg = cu_seqlens + ci_arg = chunk_indices + tuned_launch( prepare_wy_repr_bwd_kernel, stream, (NT, HV), @@ -2725,10 +2637,10 @@ def prepare_wy_repr_bwd(k, v, beta, A, dw, du, g=None, cu_seqlens=None, chunk_in BK, BV, int(g is not None), - _cdiv(K, BK), - _cdiv(V, BV), + cdiv(K, BK), + cdiv(V, BV), ), - cache_key=((NT, HV), K, V, BT, BK, BV, int(g is not None), str(k.dtype)), + cache_key=(H, HV, K, V, BT, BK, BV, int(g is not None), str(k.dtype), dev_id(k)), ) if H != HV: dk_r = reshaped(bufs["wy_dk_hred"], (T, H, K)) @@ -2758,54 +2670,44 @@ def chunk_gated_delta_rule_fwd_h( stream = 0 if stream is None else stream T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[1] BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") N, NT = len(cu_seqlens) - 1, len(chunk_indices) chunk_offsets = bufs["chunk_offsets"] - assert K <= 256, "kernel does not support head dimension larger than 256." state_shape = (N, HV, V, K) if state_v_first else (N, HV, K, V) - h = reshaped(bufs["h"], (NT, HV) + state_shape[2:]) + h = reshaped(bufs["state_checkpoints"], (NT, HV) + state_shape[2:]) final_state = reshaped(bufs["final_state"], state_shape) if output_final_state else None if final_state is not None: zero_fill(final_state, stream=stream) v_new = reshaped(bufs["v_new"], (T, HV, V)) if save_new_value else None - vnew_arg = v_new if v_new is not None else _dummy(_dtname(u), bufs) - g_arg = _opt(g, bufs) - gk_arg = _opt(gk, bufs) - h0_arg = initial_state if initial_state is not None else _dummy("float32", bufs) - ht_arg = final_state if final_state is not None else _dummy("float32", bufs) - cu_arg = _i32(cu_seqlens) - co_arg = _i32(chunk_offsets) - - # BV = V-tile width. Shrinking BV widens the V grid (more parallel CTAs) without - # adding serial work; drop to 32 for low head counts where BV=64 grid-starves the - # GPU, else keep 64. Chosen up front (not swept) to avoid a tileiras compile stall. + vnew_arg = v_new if v_new is not None else dummy(dtname(u), bufs) + g_arg = opt(g, bufs) + gk_arg = opt(gk, bufs) + h0_arg = initial_state if initial_state is not None else dummy("float32", bufs) + ht_arg = final_state if final_state is not None else dummy("float32", bufs) + cu_arg = cu_seqlens + co_arg = chunk_offsets + + # BV = V-tile width; chosen up front (not swept) to avoid a tileiras compile stall. BV = 64 if V % 32 == 0: try: - sm_count, cc_major = _device_attrs() + sm_count, cc_major = device_attrs() except Exception: # noqa: BLE001 sm_count = 0 cc_major = 0 - grid_blocks = _cdiv(V, 64) * (N * HV) - # Split only when the doubled BV=32 grid still fits within one wave - # (grid_blocks*2 <= SMs); otherwise the halved MMA width is a net loss. + grid_blocks = cdiv(V, 64) * (N * HV) if sm_count and grid_blocks * 2 <= sm_count: - # Grid-saturation split (all arches): halve BV to double CTAs when the - # BV=64 grid under-occupies (V-lanes independent -> no extra serial work). + # Grid-saturation split (all arches). BV = 32 elif cc_major == 8 and K >= 192: # Register-spill split (sm_80 only): at K>=192 the 4 fp32 (64,BV) state # accumulators overflow the 255-reg budget at BV=64 -> heavy spill. - # BV=32 halves the state (~10x less spill), fwd_h 2853->2468us. sm_100 - # keeps BV=64 (BV=32 regresses its already-saturated grid). BV = 32 - grid = (_cdiv(V, BV), N * HV) + grid = (cdiv(V, BV), N * HV) # Multi-dim arrays passed as-is (per-dim indices + real strides). - _tuned_launch( + tuned_launch( chunk_gated_delta_rule_fwd_kernel_h_blockdim64, stream, grid, @@ -2835,7 +2737,8 @@ def chunk_gated_delta_rule_fwd_h( int(state_v_first), ), cache_key=( - grid, + H, + HV, K, V, BT, @@ -2847,6 +2750,7 @@ def chunk_gated_delta_rule_fwd_h( int(save_new_value), int(state_v_first), str(k.dtype), + dev_id(k), ), ) return h, v_new, final_state @@ -2861,7 +2765,7 @@ def chunk_gated_delta_rule_bwd_dhu( g=None, gk=None, h0=None, - dht=None, + dstate_in=None, scale=None, state_v_first=False, cu_seqlens=None, @@ -2873,28 +2777,25 @@ def chunk_gated_delta_rule_bwd_dhu( stream = 0 if stream is None else stream T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[1] BT = chunk_size - assert K <= 256, "kernel does not support head dimension being larger than 256." - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") N, NT = len(cu_seqlens) - 1, len(chunk_indices) chunk_offsets = bufs["chunk_offsets"] - dh = reshaped(bufs["dh"], (NT, HV, V, K) if state_v_first else (NT, HV, K, V)) - dh0 = reshaped(bufs["dh0"], tuple(h0.shape)) if h0 is not None else None + dh = reshaped(bufs["dstate"], (NT, HV, V, K) if state_v_first else (NT, HV, K, V)) + dh0 = reshaped(bufs["dstate0"], tuple(h0.shape)) if h0 is not None else None dv2 = reshaped(bufs["dv2"], (T, HV, V)) - # For K>128 the blockdim64 kernel's 4 (64xBV) fp32 dh accumulators overflow + # For K>128 the blockdim64 kernel's 4 (64xBV) fp32 dH accumulators overflow # tileiras allocation at BV=64; shrink BV to keep the live footprint in range. BV = 16 if K > 128 else 64 - grid = (_cdiv(V, BV), N * HV) - g_arg = _opt(g, bufs) - gk_arg = _opt(gk, bufs) - dht_arg = _opt(dht, bufs) - dh0_arg = dh0 if dh0 is not None else _dummy("float32", bufs) - cu_arg = _i32(cu_seqlens) - co_arg = _i32(chunk_offsets) - - _tuned_launch( + grid = (cdiv(V, BV), N * HV) + g_arg = opt(g, bufs) + gk_arg = opt(gk, bufs) + dht_arg = opt(dstate_in, bufs) + dh0_arg = dh0 if dh0 is not None else dummy("float32", bufs) + cu_arg = cu_seqlens + co_arg = chunk_offsets + + tuned_launch( chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64, stream, grid, @@ -2922,12 +2823,13 @@ def chunk_gated_delta_rule_bwd_dhu( int(g is not None), int(gk is not None), int(h0 is not None), - int(dht is not None), + int(dstate_in is not None), int(state_v_first), 1, ), cache_key=( - grid, + H, + HV, K, V, BT, @@ -2936,6 +2838,7 @@ def chunk_gated_delta_rule_bwd_dhu( int(gk is not None), int(state_v_first), str(q.dtype), + dev_id(q), ), ) return dh, dh0, dv2 @@ -2960,25 +2863,23 @@ def chunk_fwd_o( stream = 0 if stream is None else stream T, H, K, V, HV = *q.shape, v.shape[-1], v.shape[1] BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") NT = len(chunk_indices) if scale is None: scale = k.shape[-1] ** -0.5 o = reshaped(bufs["o"], (T, HV, V)) - BK = min(max(_next_power_of_2(K), 16), 64) + BK = min(max(next_power_of_2(K), 16), 64) BV = 64 - grid = (_cdiv(V, BV), NT, HV) - g_arg = g.reshape(T, HV) if g is not None else _dummy("float32", bufs) - gg_arg = _opt(g_gamma, bufs) - cu_arg = _i32(cu_seqlens) - ci_arg = _i32(chunk_indices) + grid = (cdiv(V, BV), NT, HV) + g_arg = g.reshape(T, HV) if g is not None else dummy("float32", bufs) + gg_arg = opt(g_gamma, bufs) + cu_arg = cu_seqlens + ci_arg = chunk_indices # h flattened to (NT*HV,*,*) slabs for block-indexed TMA. if state_v_first: h3 = h.reshape(NT * HV, V, K) else: h3 = h.reshape(NT * HV, K, V) - _tuned_launch( + tuned_launch( chunk_fwd_kernel_o, stream, grid, @@ -3005,7 +2906,8 @@ def chunk_fwd_o( int(state_v_first), ), cache_key=( - grid, + H, + HV, K, V, BT, @@ -3015,6 +2917,7 @@ def chunk_fwd_o( int(g_gamma is not None), int(state_v_first), str(q.dtype), + dev_id(q), ), ) return o @@ -3042,13 +2945,11 @@ def chunk_bwd_dqkwg( stream = 0 if stream is None else stream T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[1] BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") NT = len(chunk_indices) CONST_TILING = 64 - BK = min(max(_next_power_of_2(K), 16), CONST_TILING) - BV = min(max(_next_power_of_2(V), 16), CONST_TILING) - NK = _cdiv(K, BK) + BK = min(max(next_power_of_2(K), 16), CONST_TILING) + BV = min(max(next_power_of_2(V), 16), CONST_TILING) + NK = cdiv(K, BK) dq = reshaped(bufs["dq"], (T, HV, K)) dk = reshaped(bufs["dk"], (T, HV, K)) dg = reshaped(bufs["dg_nk"], (NK, T, HV)) if g is not None else None @@ -3061,15 +2962,15 @@ def chunk_bwd_dqkwg( else: h3 = h.reshape(NT * HV, K, V) dh3 = dh.reshape(NT * HV, K, V) - g_arg = g.reshape(T, HV) if g is not None else _dummy(_dtname(q), bufs) - gg_arg = _opt(g_gamma, bufs) - dw3 = dw if dw is not None else _dummy(_dtname(k), bufs) - dv3 = dv.reshape(T, HV, V) if dv is not None else _dummy(_dtname(k), bufs) - dg3 = dg if dg is not None else _dummy("float32", bufs) - cu_arg = _i32(cu_seqlens) - ci_arg = _i32(chunk_indices) + g_arg = g.reshape(T, HV) if g is not None else dummy(dtname(q), bufs) + gg_arg = opt(g_gamma, bufs) + dw3 = dw if dw is not None else dummy(dtname(k), bufs) + dv3 = dv.reshape(T, HV, V) if dv is not None else dummy(dtname(k), bufs) + dg3 = dg if dg is not None else dummy("float32", bufs) + cu_arg = cu_seqlens + ci_arg = chunk_indices use_dw = int(dw is not None and dv is not None) - _tuned_launch( + tuned_launch( chunk_bwd_kernel_dqkwg, stream, grid, @@ -3101,10 +3002,11 @@ def chunk_bwd_dqkwg( int(g_gamma is not None), use_dw, int(state_v_first), - _cdiv(V, BV), + cdiv(V, BV), ), cache_key=( - grid, + H, + HV, K, V, BT, @@ -3115,6 +3017,7 @@ def chunk_bwd_dqkwg( use_dw, int(state_v_first), str(q.dtype), + dev_id(q), ), ) if H != HV: @@ -3134,23 +3037,19 @@ def chunk_bwd_dv_local(q, k, do, g=None, g_gamma=None, A=None, scale=None, cu_se stream = 0 if stream is None else stream T, H, K, V, HV = *k.shape, do.shape[-1], do.shape[1] BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") CONST_TILING = 64 - BK = min(max(_next_power_of_2(K), 16), CONST_TILING) - BV = min(max(_next_power_of_2(V), 16), CONST_TILING) + BK = min(max(next_power_of_2(K), 16), CONST_TILING) + BV = min(max(next_power_of_2(V), 16), CONST_TILING) NT = len(chunk_indices) dv = reshaped(bufs["dv"], (T, HV, V)) grid = (NT, HV) - g_arg = _opt(g, bufs) - gg_arg = _opt(g_gamma, bufs) - A_arg = _opt(A, bufs) - cu_arg = _i32(cu_seqlens) - ci_arg = _i32(chunk_indices) + g_arg = opt(g, bufs) + gg_arg = opt(g_gamma, bufs) + A_arg = opt(A, bufs) + cu_arg = cu_seqlens + ci_arg = chunk_indices scale_val = float(scale) if scale is not None else 0.0 - # Autotune (occupancy, nww): an occ hint budgets - # fewer regs/thread -> more blocks/SM. - _tuned_launch( + tuned_launch( chunk_bwd_kernel_dv_local, stream, grid, @@ -3177,7 +3076,8 @@ def chunk_bwd_dv_local(q, k, do, g=None, g_gamma=None, A=None, scale=None, cu_se int(A is not None), ), cache_key=( - grid, + H, + HV, K, V, BT, @@ -3187,30 +3087,28 @@ def chunk_bwd_dv_local(q, k, do, g=None, g_gamma=None, A=None, scale=None, cu_se int(g_gamma is not None), int(A is not None), str(q.dtype), + dev_id(q), ), ) return dv -# fused intra (kkt + solve_tril + recompute_w_u) +# fused intra (KK^T + solve_tril + recompute_w_u) def chunk_gated_delta_rule_fwd_intra(k, v, g=None, beta=None, cu_seqlens=None, chunk_size=64, chunk_indices=None, bufs=None, compute_wu=True, stream=None): stream = 0 if stream is None else stream T, H, K, HV = *k.shape, beta.shape[1] BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") NT = len(chunk_indices) A = reshaped(bufs["A"], (T, HV, BT)) zero_fill(A, stream=stream) BK = 64 - g_arg = _opt(g, bufs) - cu_arg = _i32(cu_seqlens) - ci_arg = _i32(chunk_indices) - # Masked BC=16 4-sub-block kernel: masks the partial last chunk's tail - # (m_out / o_t < T). + g_arg = opt(g, bufs) + cu_arg = cu_seqlens + ci_arg = chunk_indices + # Masked BC=16 4-sub-block kernel: masks the partial last chunk's tail. BC = 16 - _tuned_launch( + tuned_launch( chunk_gated_delta_rule_fwd_kkt_solve_kernel, stream, (NT, HV), @@ -3229,7 +3127,7 @@ def chunk_gated_delta_rule_fwd_intra(k, v, g=None, beta=None, cu_seqlens=None, c BK, int(g is not None), ), - cache_key=(NT, HV, K, BT, BC, BK, int(g is not None), str(k.dtype)), + cache_key=(H, HV, K, BT, BC, BK, int(g is not None), str(k.dtype), dev_id(k)), ) if not compute_wu: return None, None, A @@ -3260,7 +3158,6 @@ def chunk_gated_delta_rule_fwd( stream=None, ): stream = 0 if stream is None else stream - assert cp_context is None, "context-parallel not supported in this port" g_input = g if use_gate_in_kernel else None if use_gate_in_kernel: g = gdn_gate_chunk_cumsum( @@ -3307,7 +3204,7 @@ def chunk_gated_delta_rule_bwd( scale, initial_state, do, - dht, + dstate_in, state_v_first=False, cu_seqlens=None, cp_context=None, @@ -3320,7 +3217,6 @@ def chunk_gated_delta_rule_bwd( stream=None, ): stream = 0 if stream is None else stream - assert cp_context is None, "context-parallel not supported in this port" w, u = recompute_w_u_fwd(k=k, v=v, beta=beta, A=A, g=g, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, bufs=bufs, stream=stream) h, v_new, _ = chunk_gated_delta_rule_fwd_h( k=k, @@ -3342,7 +3238,7 @@ def chunk_gated_delta_rule_bwd( w=w, g=g, h0=initial_state, - dht=dht, + dstate_in=dstate_in, do=do, dv=dv, scale=scale, @@ -3400,7 +3296,7 @@ def chunk_gated_delta_rule_grad( g, beta, do, - dht=None, + dstate_in=None, scale=None, initial_state=None, state_v_first=False, @@ -3420,7 +3316,7 @@ def chunk_gated_delta_rule_grad( ``q``/``k``/``v``/``do`` are ``[total_T, H, D]``, ``g``/``beta`` are ``[total_T, H]``; ``cu_seqlens`` and ``chunk_indices`` are required. - Recomputes the forward's prep (L2-normalized q/k + rstd, cumulative gate, + Recomputes the forward's prep (L2-normalized Q/K + rstd, cumulative gate, intra-chunk WY matrix) from the ORIGINAL inputs, then runs the backward kernels. @@ -3469,7 +3365,7 @@ def chunk_gated_delta_rule_grad( scale=scale, initial_state=initial_state, do=do, - dht=dht, + dstate_in=dstate_in, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, state_v_first=state_v_first, @@ -3521,7 +3417,6 @@ def chunk_gated_delta_rule( required. Returns ``(o, final_state)`` in THD layout, written into ``bufs['o']`` / ``bufs['final_state']``.""" stream = 0 if stream is None else stream - assert cp_context is None, "context-parallel not supported" if "transpose_state_layout" in kwargs: if state_v_first: diff --git a/python/cudnn/linear_attention/cutile/kernels/kda_chunk_cutile.py b/python/cudnn/linear_attention/cutile/kernels/kda_chunk_cutile.py index 8d297f05c..bd0189fc0 100644 --- a/python/cudnn/linear_attention/cutile/kernels/kda_chunk_cutile.py +++ b/python/cudnn/linear_attention/cutile/kernels/kda_chunk_cutile.py @@ -21,13 +21,11 @@ from types import SimpleNamespace import cuda.tile as ct - -from .common import add_inplace, head_group_sum, reshaped, sum_leading, zero_fill -from cudnn.frost.buffers import dtype_name as _dtname - -_F32_DEFAULT = "float32-default" from cuda.tile.tune import exhaustive_search +from .common import dev_id, dummy, head_group_sum, opt, reshaped, sum_leading, zero_fill +from cudnn.frost.buffers import dtype_name as dtname + logger = logging.getLogger(__name__) ConstInt = ct.Constant[int] @@ -35,22 +33,21 @@ RCP_LN2 = 1.4426950216 # 1/ln(2) # chunk size (BT tile) of these kernels; the engine's carve layout imports it -_BT = 64 +BT_CHUNK = 64 # --------------------------------------------------------------------------- -# Launch-hint autotune (mirrors the sibling GDN -# cuTile kernels). Only the @ct.kernel launch hints +# Launch-hint autotune. Only the @ct.kernel launch hints # (occupancy x num_worker_warps) are explored via kernel.replace_hints(...); # the grid / args / algorithm are kept UNCHANGED. The first config in every # grid equals the kernel default (occupancy=1, num_worker_warps=4) so any shape # that does not improve keeps the original behaviour (no regression). # --------------------------------------------------------------------------- -_DISABLE_TUNE = os.environ.get("DISABLE_TUNE", "") not in ("", "0") -_launch_hint_cache: dict = {} +DISABLE_TUNE = os.environ.get("DISABLE_TUNE", "") not in ("", "0") +launch_hint_cache: dict = {} -def _launch_hint_configs(occ_choices, nww_choices=(4, 8)): +def launch_hint_configs(occ_choices, nww_choices=(4, 8)): """occupancy x num_worker_warps grid; deduped, default (occ=1,nww=4) first.""" seen = set() cfgs = [] @@ -64,7 +61,7 @@ def _launch_hint_configs(occ_choices, nww_choices=(4, 8)): return cfgs -def _autotuned_launch(kernel, cache_key, grid, args, occ_choices=(1, 2, 3, 4), nww_choices=(4, 8), timeout=30, stream=None): +def autotuned_launch(kernel, cache_key, grid, args, occ_choices=(1, 2, 3, 4), nww_choices=(4, 8), timeout=30, stream=None): """Launch ``kernel`` with the best launch hints for ``cache_key``. Tune-once/cache/launch over launch hints only (grid, args and signature are @@ -73,14 +70,14 @@ def _autotuned_launch(kernel, cache_key, grid, args, occ_choices=(1, 2, 3, 4), n no-improvement shape keeps the base behaviour. """ stream = 0 if stream is None else stream - if _DISABLE_TUNE: + if DISABLE_TUNE: ct.launch(stream, grid, kernel, args) return - if cache_key not in _launch_hint_cache: + if cache_key not in launch_hint_cache: tuned = None try: - configs = _launch_hint_configs(occ_choices, nww_choices) + configs = launch_hint_configs(occ_choices, nww_choices) with ct.compiler_timeout(timeout): result = exhaustive_search( configs, @@ -94,16 +91,16 @@ def _autotuned_launch(kernel, cache_key, grid, args, occ_choices=(1, 2, 3, 4), n tuned = kernel.replace_hints(occupancy=best.occupancy, num_worker_warps=best.num_worker_warps) except Exception: tuned = None - _launch_hint_cache[cache_key] = tuned + launch_hint_cache[cache_key] = tuned - tuned = _launch_hint_cache[cache_key] + tuned = launch_hint_cache[cache_key] if tuned is None: ct.launch(stream, grid, kernel, args) else: ct.launch(stream, grid, tuned, args) -def _autotuned_launch_bv(kernel, cache_key, bv_choices, grid_fn, args_fn, timeout=40, stream=None): +def autotuned_launch_bv(kernel, cache_key, bv_choices, grid_fn, args_fn, timeout=40, stream=None): """Launch ``kernel`` autotuning over BV (the V-tile block width) only. BV is a kernel ConstInt that drives the grid V-fan-out and the V-tile shapes @@ -114,30 +111,21 @@ def _autotuned_launch_bv(kernel, cache_key, bv_choices, grid_fn, args_fn, timeou BV is cached and the BV-specialized kernel re-launched on every subsequent call. ``bv_choices[0]`` is the safe fallback on DISABLE_TUNE / tuning failure. - Why BV is the dominant fwd_h lever and why the sweep is BV-ONLY: the state - scan is register-bound (255 reg/thread -> ~1 block/SM) with a serial per-CTA - NT inter-chunk loop, so the only parallelism lever is the V-tile CTA count - (grid = (cdiv(V, BV), N*HV)). ncu on the dashboard shapes shows the optimum - is NOT monotone in BV: at (T=2048,V=128,N*HV=64) BV=64 (grid 128, warps_active - 14%) fwd_h=246us beats BV=32 (grid 256, 8%) 290us and BV=16 374us -- fewer, - fatter CTAs win once the grid already fills the SMs, because a smaller BV just - multiplies redundant k/w/g reloads; at (T=1024,V=64,N*HV=64) BV=32 (grid 128) - 65us beats BV=64 (grid 64, under-filled) 123us and BV=8 123us. Crucially the - occupancy/num_worker_warps hint is nearly inert here (per-BV best occ/nww is - within ~0.3% of default), so we do NOT cross occ x nww into this search: - doing so inflated it to 16 configs and the larger exhaustive_search do_bench - mis-ranked BV (it picked BV=32 for the 2048 shape even though a direct - BV=64/occ-tuned run is 15% faster). A clean 4-config BV-only sweep -- mirroring - the kda_chunk_fwd helper's ``_tuned_launch_fwdh_bv`` -- ranks BV correctly - (picks 64 for the 2048 shape, 32 for the 1024 shape). + Why the sweep is BV-ONLY: the state scan is register-bound (255 reg/thread + -> ~1 block/SM) with a serial per-CTA NT inter-chunk loop, so the only + parallelism lever is the V-tile CTA count (grid = (cdiv(V, BV), N*HV)). + The optimum is NOT monotone in BV (a smaller BV multiplies redundant K/W/G + reloads once the grid already fills the SMs), and the occupancy x + num_worker_warps hint is nearly inert per BV -- do NOT cross occ x nww into + this search (the larger sweep mis-ranks BV). """ stream = 0 if stream is None else stream - if _DISABLE_TUNE: + if DISABLE_TUNE: bv = bv_choices[0] ct.launch(stream, grid_fn(bv), kernel, args_fn(bv)) return - if cache_key not in _launch_hint_cache: + if cache_key not in launch_hint_cache: chosen = None try: configs = [SimpleNamespace(BV=bv) for bv in bv_choices] @@ -153,9 +141,9 @@ def _autotuned_launch_bv(kernel, cache_key, bv_choices, grid_fn, args_fn, timeou chosen = result.best.config.BV except Exception: chosen = None - _launch_hint_cache[cache_key] = chosen + launch_hint_cache[cache_key] = chosen - chosen = _launch_hint_cache[cache_key] + chosen = launch_hint_cache[cache_key] bv = bv_choices[0] if chosen is None else chosen ct.launch(stream, grid_fn(bv), kernel, args_fn(bv)) @@ -165,55 +153,20 @@ def _autotuned_launch_bv(kernel, cache_key, bv_choices, grid_fn, args_fn, timeou # =========================================================================== -def _cdiv(a: int, b: int) -> int: +def cdiv(a: int, b: int) -> int: return (a + b - 1) // b -def _next_power_of_2(n: int) -> int: +def next_power_of_2(n: int) -> int: return 1 << (n - 1).bit_length() -def _dev_id(buf) -> int: - """Device ordinal of a DLPack/CAI buffer.""" - from cudnn.frost.buffers import probe - - return probe(buf)[4] - - -def _dummy(dtype_name: str, bufs): - """Inert typed view over the workspace's 16-byte ``dummy`` carve, for - ABSENT optional kernel args (always paired with a flag==0, never - dereferenced). Dtype-bound so the compiled signature stays stable; the - library allocates nothing.""" - from cudnn.frost.buffers import DTYPE_ITEMSIZE, DeviceView - - d = bufs["dummy"] - return DeviceView(d.data_ptr(), (16 // DTYPE_ITEMSIZE[dtype_name],), dtype_name, d.__dlpack_device__()[1]) - - -def _i32(t): - if not str(t.dtype).endswith("int32"): - raise TypeError(f"index/boundary buffers must be int32 (callers convert), got {t.dtype}") - return t - - -def _opt(t, bufs, dtype_name: str = "float32"): - """Resolve an optional tensor argument to a non-null cuTile launch arg: - the buffer if present (contiguous by the engine contract), else an inert - dummy (paired with a USE_*/HAS_* integer flag). cuTile never accepts None - in launch args, so this is the required dummy-tensor-plus-flag pattern.""" - if t is None: - return _dummy(dtype_name, bufs) - return t - - -def _i32_flat(t): - """Required index buffer: int32 guard + flat 1-D view. Kernels index these +def i32_flat(t): + """Required index buffer as a flat 1-D view. Kernels index these via flat loads (e.g. chunk_indices[i_t*2]); chunk_indices is (NT, 2) and its row-major flat layout is [seg0, intra0, seg1, intra1, ...].""" from .common import reshaped - t = _i32(t) n = 1 for s_ in t.shape: n *= int(s_) @@ -225,7 +178,7 @@ def _i32_flat(t): # ``bufs`` and passes outputs via ``out=``; nothing here allocates. -def _cast(bufs, name, src, ref, stream=None): +def cast(bufs, name, src, ref, stream=None): """Dtype cast at the gradient boundary, through the ``bufs[name]`` carve.""" if str(src.dtype).split(".")[-1] == str(ref.dtype).split(".")[-1]: return src @@ -239,37 +192,6 @@ def _cast(bufs, name, src, ref, stream=None): return dst -def _i32(t): - if not str(t.dtype).endswith("int32"): - raise TypeError(f"index/boundary buffers must be int32 (callers convert), got {t.dtype}") - return t - - -def _opt(t, device_id, dtype_name: str = "float32"): - """Resolve an optional tensor argument to a non-null cuTile launch arg: - the buffer if present (contiguous by the engine contract), else an inert - dummy (paired with a USE_*/HAS_* integer flag). cuTile never accepts None - in launch args, so this is the required dummy-tensor-plus-flag pattern.""" - if t is None: - return _dummy(dtype_name, device_id) - return t - - -def _opt_i32(t, device_id): - if t is None: - return _dummy("int32", device_id) - # Kernels index these via flat 1-D loads (e.g. chunk_indices[i_t*2]); flatten - # so the cuTile load index rank (1) matches. chunk_indices is (NT, 2) and its - # row-major flat layout is [seg0, intra0, seg1, intra1, ...] as expected. - from .common import reshaped - - t = _i32(t) - n = 1 - for s_ in t.shape: - n *= int(s_) - return reshaped(t, (n,)) - - # =========================================================================== # Device helpers (plain `def` — NOT @ct.kernel) # =========================================================================== @@ -288,33 +210,33 @@ def softplus(x): return ct.where(x <= 20.0, ct.log(1.0 + ct.exp(x)), x) -def _tf32(a): +def tf32(a): """ct.mma/ct.matmul do not auto-cast fp32 operands to tf32; cast explicitly (allow-tf32 matmul semantics).""" return ct.astype(a, ct.tfloat32) if a.dtype == ct.float32 else a -def _mma_operands(a, b): +def mma_operands(a, b): # ct.mma/ct.matmul require matching operand dtypes. Forward callers always # pass matched dtypes (left unchanged here); some backward callers mix bf16 - # with fp32, so promote both to fp32 in that case. Then apply the tf32 guard. + # with fp32, so promote both to fp32 in that case. if a.dtype != b.dtype: a = ct.astype(a, ct.float32) b = ct.astype(b, ct.float32) - return _tf32(a), _tf32(b) + return tf32(a), tf32(b) def safe_matmul(a, b): # tf32 is only the multiply precision; ct.matmul returns a tf32 tile which # is a restricted dtype (no elementwise arithmetic). Cast the result back to # fp32 so downstream adds work (fp32 dot-accumulator semantics). - aa, bb = _mma_operands(a, b) + aa, bb = mma_operands(a, b) out = ct.matmul(aa, bb) return ct.astype(out, ct.float32) if out.dtype != ct.float32 else out def safe_mma(a, b, acc): - aa, bb = _mma_operands(a, b) + aa, bb = mma_operands(a, b) return ct.mma(aa, bb, acc) @@ -322,9 +244,7 @@ def bf16_mma(a, b, acc): # Like safe_mma but stages the MMA operands as bf16 (2 bytes) instead of # tf32 (4 bytes). cuTile stages tiny MMA operands into static SMEM; tf32 # staging doubles that footprint vs bf16, which on the SMEM-bound - # compute-pairs kernel pins it to 1 block/SM. The q/k inputs are bf16 - # upstream, so bf16 operands keep the multiply precision the inputs already - # carry; the fp32 accumulator is preserved. + # compute-pairs kernel pins it to 1 block/SM. return ct.mma(ct.astype(a, ct.bfloat16), ct.astype(b, ct.bfloat16), acc) @@ -334,10 +254,9 @@ def reg_matmul(a, b): # SMEM-staged HMMA path (~83KB static smem/block -> 1 block/SM, ~12% occ). # Keep these dots in registers instead, via a # broadcast-multiply-reduce over the contraction dim: out[i,j] = - # sum_k a[i,k]*b[k,j]. Inputs are rounded to tf32 first so the multiply - # precision matches the existing tf32-MMA path (preserves tolerance). - aa = _tf32(a) - bb = _tf32(b) + # sum_k a[i,k]*b[k,j]. Inputs are rounded to tf32 first. + aa = tf32(a) + bb = tf32(b) aa = ct.astype(aa, ct.float32) bb = ct.astype(bb, ct.float32) return ct.sum(aa[:, :, None] * bb[None, :, :], axis=1) @@ -356,7 +275,7 @@ def ct_min(a, b): # GATHER RATIONALE (l2norm): each row is loaded once, reduced once, written # once (single-pass streaming, no cooperative reuse) -> ct.gather/ct.scatter # avoids the smem staging tax of ct.load. Column mask handles BD-power-of-2 -# padding; partial-row handling via masks (replaces boundary_check). +# padding; partial-row handling via masks. @ct.kernel @@ -439,7 +358,7 @@ def fused_beta_sigmoid_bwd_kernel(x, dy, dx, scale, n_elements, BLOCK_SIZE: Cons # =========================================================================== -# vector chunk_local_cumsum (used for the vector gate g) +# vector chunk_local_cumsum (used for the vector gate G) # =========================================================================== # GATHER RATIONALE: head-interleaved [T, H, S] sub-view; loaded once, # cumsum'd along T, written back. s/o passed flattened (1-D) by the wrapper. @@ -494,7 +413,7 @@ def chunk_local_cumsum_vector_kernel( # =========================================================================== # kda_gate_chunk_cumsum_vector / kda_gate_bwd # =========================================================================== -# Vector gate: g is [B, T, H, S]; A_log is [H]; dt_bias is [H*S]. +# Vector gate: G is [B, T, H, S]; A_log is [H]; dt_bias is [H*S]. # GATHER RATIONALE: head-interleaved [B, T, H, S] sub-view streamed once. @@ -577,7 +496,7 @@ def kda_gate_bwd_kernel( HAS_BIAS: ConstInt, USE_LOWER_BOUND: ConstInt, ): - # g/dyg/dg layout [T, H, D] -> view (T, D) stride (H*D, 1) at base i_h*D. + # G/dYg/dG layout [T, H, D] -> view (T, D) stride (H*D, 1) at base i_h*D. i_t = ct.bid(0) i_h = ct.bid(1) @@ -623,11 +542,11 @@ def kda_gate_bwd_kernel( # =========================================================================== -# chunk_gla_fwd_kernel_o -- output projection o = qg @ h + A @ v +# chunk_gla_fwd_kernel_o -- output projection O = QG @ H + A @ V # =========================================================================== -# GATHER RATIONALE: q/g/v are reused across the i_k loop (K split), but with -# head-interleaved strides + a transposed STATE_V_FIRST h-view that TMA cannot -# express here, so the faithful non-TMA element-offset gather path is used. +# GATHER RATIONALE: Q/G/V are reused across the i_k loop (K split), but with +# head-interleaved strides + a transposed STATE_V_FIRST H-view that TMA cannot +# express here, so the non-TMA element-offset gather path is used. @ct.kernel @@ -655,8 +574,8 @@ def chunk_gla_fwd_kernel_o( i_hv = ct.bid(2) i_h = i_hv // (HV // H) - # grid dim-1 is the GLOBAL chunk index; h is laid out per global chunk - # (h-state kernel writes slot chunk_offsets[i_n] + local). Capture it + # grid dim-1 is the GLOBAL chunk index; H is laid out per global chunk + # (H-state kernel writes slot chunk_offsets[i_n] + local). Capture it # before i_t is reassigned to the per-sequence (local) chunk index. i_tg = i_t i_n = ct.load(chunk_indices, (i_t * 2,), shape=()).item() @@ -731,10 +650,9 @@ def chunk_gla_fwd_kernel_o( # =========================================================================== -# chunk_delta_h (state-h fwd + dhu bwd) -- SHARED with GDN +# chunk_delta_h (state-H fwd + dhu bwd) # =========================================================================== -# These two kernels are byte-for-byte identical (in math) to the GDN -# chunk_delta_h kernels. For KDA the gate carried is the per-key vector gate +# For KDA the gate carried is the per-key vector gate # `gk` (USE_GK=1, USE_G=0); the per-time scalar gate `g` path is inert here. @@ -766,12 +684,11 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( STATE_V_FIRST: ConstInt, ): # The KV state is carried as a SINGLE full-width tile ((BV, BK) when - # STATE_V_FIRST else (BK, BV), BK=next_pow2(K)) instead of a 4-way 64-wide - # sub-block unroll, so the kernel is general for ANY K (no K<=256 cap). Every - # K-axis load zero-pads cols [K:BK] and every K-axis store masks them off, so - # the tail is 0 on load, contributes 0 to every MMA, and stays 0 in the - # state; a padded gk tail loads as 0 so exp2(0)=1 leaves those zero cols - # unchanged. + # STATE_V_FIRST else (BK, BV), BK=next_pow2(K)), so the kernel is general + # for ANY K (no K<=256 cap). Every K-axis load zero-pads cols [K:BK] and + # every K-axis store masks them off, so the tail is 0 on load, contributes + # 0 to every MMA, and stays 0 in the state; a padded Gk tail loads as 0 so + # exp2(0)=1 leaves those zero cols unchanged. i_v = ct.bid(0) i_nh = ct.bid(1) i_n = i_nh // HV @@ -801,8 +718,8 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( # K-tile / V-tile validity masks. The kernel carries a single BK-wide K tile # (BK=next_pow2(K)) and BV-wide V tiles; when K or V is not a multiple of the - # tile width the extra lanes alias the neighbouring head's h slot. The matmul - # state rows are zeroed via K-masked k/gk loads, but the raw h/h0/ht + # tile width the extra lanes alias the neighbouring head's H slot. The matmul + # state rows are zeroed via K-masked K/Gk loads, but the raw H/H0/Ht # gather/scatter are also masked here -> no cross-head corruption when # K % BK != 0 (i.e. K < BK) or V % BV != 0. mkh = o_bk < K @@ -942,7 +859,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( w, g, gk, - dht, + dstate_in, dh0, do, dh, @@ -964,14 +881,12 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( USE_FINAL_STATE_GRADIENT: ConstInt, STATE_V_FIRST: ConstInt, ): - # - # dh state is a SINGLE full-width tile ((BV, BK) when STATE_V_FIRST else - # (BK, BV), BK=next_pow2(K)) instead of a 4-way 64-wide sub-block unroll, so - # the kernel is general for ANY K (no K<=256 cap). Every flat gather/scatter - # over the K axis uses oK=arange(BK) with a `(oK < K)` mask, so rows/cols - # [K:BK] are zero on load, contribute 0 to every MMA, and are masked out on - # store; a gk padding tail loads as 0 so exp2(0)=1 leaves those zero state - # cols unchanged. + # dH state is a SINGLE full-width tile ((BV, BK) when STATE_V_FIRST else + # (BK, BV), BK=next_pow2(K)), so the kernel is general for ANY K (no K<=256 + # cap). Every flat gather/scatter over the K axis uses oK=arange(BK) with a + # `(oK < K)` mask, so rows/cols [K:BK] are zero on load, contribute 0 to + # every MMA, and are masked out on store; a Gk padding tail loads as 0 so + # exp2(0)=1 leaves those zero state cols unchanged. i_v = ct.bid(0) i_nh = ct.bid(1) i_n = i_nh // HV @@ -1001,18 +916,18 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( o_bt = ct.arange(BT, dtype=ct.int32) o_bv = ct.arange(BV, dtype=ct.int32) - # V-boundary mask for the state-gradient (dh/dh0) scatters: a partial + # V-boundary mask for the state-gradient (dH/dH0) scatters: a partial # trailing V tile (V not a multiple of BV) must not write rows/cols >= V, # else the store spills into the neighbouring chunk/head's state-gradient # (boundary-checked store semantics). m_v = (i_v * BV + o_bv) < V - # --- Load final state gradient dht -> b_dh (single full-width tile; [K:BK] loads 0) --- + # --- Load final state gradient dHt -> b_dh (single full-width tile; [K:BK] loads 0) --- if USE_FINAL_STATE_GRADIENT: if STATE_V_FIRST: row = (i_v * BV + o_bv)[:, None] b_dh = b_dh + ct.gather( - dht, + dstate_in, dht_base + row * K + o_bk[None, :], mask=(o_bk < K)[None, :], check_bounds=True, @@ -1021,7 +936,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( else: col = (i_v * BV + o_bv)[None, :] b_dh = b_dh + ct.gather( - dht, + dstate_in, dht_base + o_bk[:, None] * V + col, mask=(o_bk < K)[:, None], check_bounds=True, @@ -1029,12 +944,12 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( ) # cuTile range() requires a positive step; iterate forward and reverse the - # index to preserve the original backward (NT-1 .. 0) chunk traversal. + # index to preserve the backward (NT-1 .. 0) chunk traversal. for _i_t in range(NT): i_t = NT - 1 - _i_t dh_chunk = dh_base + i_t * HV * K * V - # Store current b_dh to dh[i_t] (single full-width tile; [K:BK] masked out) + # Store current b_dh to dH[i_t] (single full-width tile; [K:BK] masked out) if STATE_V_FIRST: row = (i_v * BV + o_bv)[:, None] ct.scatter( @@ -1067,7 +982,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( bg_last_exp = ct.astype(0.0, ct.float32) b_g_exp = ct.zeros((BT,), dtype=ct.float32) - # do, dv, dv2 tiles + # dO, dV, dV2 tiles v_col = (i_v * BV + o_bv)[None, :] vmask_c = (i_v * BV + o_bv) < V v_full = (m_t[:, None]) & (vmask_c[None, :]) @@ -1086,7 +1001,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( ) bk = ct.where(m_t[:, None], bk, ct.zeros((BT, BK), dtype=bk.dtype)) if USE_GK: - # gk offset: base (bos*HV+i_h)*K ; then + last_idx*HV*K + o_k. + # Gk offset: base (bos*HV+i_h)*K ; then + last_idx*HV*K + o_k. # Padded tail [K:BK] loads as 0 -> exp2(0)=1 leaves those zero cols unchanged. gkl = (bos * HV + i_h) * K + last_idx * HV * K b_gk_last = ct.astype(ct.gather(gk, gkl + o_bk, mask=o_bk < K, check_bounds=True, padding_value=0.0), ct.float32) @@ -1097,7 +1012,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( decay = ct.where(m_t, exp2(bg_last - b_g), ct.zeros((BT,), dtype=ct.float32)) b_dv = b_dv * decay[:, None] - # b_dv += dv ; store to dv2 + # b_dv += dV ; store to dV2 dv_off = dv_base + (i_t * BT + o_bt)[:, None] * (HV * V) + v_col b_dv_load = ct.gather(dv, dv_off, check_bounds=True, padding_value=0.0) b_dv_load = ct.where(v_full, b_dv_load, ct.zeros((BT, BV), dtype=b_dv_load.dtype)) @@ -1157,7 +1072,7 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( # wy_fast (recompute_w_u_fwd) # =========================================================================== # 2D structured gather/scatter on reshaped flat views; head-interleaved -# strides folded into the (row, col) tuple indices. gk is the per-key gate. +# strides folded into the (row, col) tuple indices. Gk is the per-key gate. @ct.kernel @@ -1365,7 +1280,7 @@ def chunk_kda_fwd_kernel_intra_token_parallel( BH: ConstInt, BK: ConstInt, ): - # used (numerically identical). BK = next_power_of_2(K) is passed by host + # BK = next_power_of_2(K) is passed by host # (cuTile tile shapes must be compile-time constants). i_tg = ct.bid(0) i_hg = ct.bid(1) @@ -1406,7 +1321,7 @@ def chunk_kda_fwd_kernel_intra_token_parallel( # cuTile gather/scatter require the index tuple rank to match the array rank # (no raw pointer arithmetic). Arrays arrive pre-flattened from the - # host: q/k -> (B*T*H, K), g -> (B*T*HV, K), beta/Aqk/Akk -> 1-D. + # host: Q/K -> (B*T*H, K), G -> (B*T*HV, K), Beta/Aqk/Akk -> 1-D. o_hv = i_hg * BH + ct.arange(BH, dtype=ct.int32) o_h = o_hv // G o_k = ct.arange(BK, dtype=ct.int32) @@ -1416,12 +1331,12 @@ def chunk_kda_fwd_kernel_intra_token_parallel( col = ct.broadcast_to(o_k[None, :], (BH, BK)) - # q/k: row = (bos + token) * H + head; col = key + # Q/K: row = (bos + token) * H + head; col = key qk_row = ct.broadcast_to(((bos + i_t) * H + o_h)[:, None], (BH, BK)) b_q = ct.astype(ct.gather(q, (qk_row, col), mask=m_hk, check_bounds=False, padding_value=0.0), ct.float32) b_k = ct.astype(ct.gather(k, (qk_row, col), mask=m_hk, check_bounds=False, padding_value=0.0), ct.float32) - # g: row = (bos + token) * HV + head; beta: idx = (bos + token) * HV + head + # G: row = (bos + token) * HV + head; Beta: idx = (bos + token) * HV + head g_row = ct.broadcast_to(((bos + i_t) * HV + o_hv)[:, None], (BH, BK)) b_g = ct.astype(ct.gather(g, (g_row, col), mask=m_hk, check_bounds=False, padding_value=0.0), ct.float32) b_beta = ct.astype(ct.gather(beta, beta_base + i_t * HV + o_hv, mask=m_hv, check_bounds=False, padding_value=0.0), ct.float32) @@ -1430,8 +1345,7 @@ def chunk_kda_fwd_kernel_intra_token_parallel( # Counted loop over the (static) BC-wide sub-chunk window. A runtime-bounded # `for j in range(i_ts, j_hi)` lowers to per-iteration branches that tileiras # cannot unroll/predicate; a counted `for jj in range(BC)` with a runtime - # guard derives `j` from `jj` and stays fully unrolled (numerically identical: - # inactive iterations are masked out). See MEMORY counted-loop note. + # guard derives `j` from `jj` and stays fully unrolled. for jj in range(BC): j = i_ts + jj if j < i_t + 1 and j < T_eff and j < i_ts + BC: @@ -1465,7 +1379,7 @@ def chunk_kda_fwd_kernel_intra_token_parallel( # =========================================================================== # One BC sub-chunk per block: compute the diagonal Aqk/Akk blocks then do an # in-place forward-substitution triangular solve. GATHER RATIONALE: head- -# interleaved q/k/g sub-views + a transposed b_kgt; faithful gather path. +# interleaved Q/K/G sub-views + a transposed b_kgt. @ct.kernel @@ -1560,30 +1474,16 @@ def chunk_kda_fwd_kernel_inter_diag_compute_solve( m_aqk = m_c[:, None] & ((i_i * BC + o_bc) < BT)[None, :] ct.scatter(Aqk, aqk_off, ct.astype(b_Aqk, Aqk.dtype), mask=m_aqk, check_bounds=False) - # forward substitution on the diagonal Akk -> inverse, written to Akk(diag buf). + # diagonal Akk -> inverse via Neumann series by squaring, written to Akk(diag buf). # - # Akk here is strictly-lower-triangular and nilpotent (Akk^BC = 0), and the - # inverse we want is (I + Akk)^-1 = I - Akk + Akk^2 - ... (the serial sweep - # inits b_Ai = -Akk, so with N := -Akk this is (I - N)^-1 = I + N + N^2 + - # ... + N^(BC-1)). A per-row serial forward-substitution sweep - # (`for i in range(2, min(BC, ...))`) lowers to a chain of BC-1 - # serial SIMT rank-1 updates -- the dominant compute cost of this - # compute-bound kernel (ncu: 51% SM, HMMA only from the two upfront dots). + # Akk is strictly-lower-triangular and nilpotent (Akk^BC = 0); with + # N := -Akk, (I + Akk)^-1 = (I - N)^-1 = sum_{k=0}^{BC-1} N^k + # = prod_{j=0}^{log2(BC)-1} (I + N^(2^j)) -- log2(BC) block matmuls, no + # serial row dependency. Rows beyond T_eff have N = 0 (b_gq/b_gk were + # zeroed by m_c) and converge to the identity. # - # We replace the serial sweep with a Neumann-series-by-squaring block - # matmul: (I - N)^-1 = prod_{j=0}^{log2(BC)-1} (I + N^(2^j)). This is - # algebraically identical (telescopes to sum_{k=0}^{BC-1} N^k since - # N^BC = 0) and turns the BC-1 serial steps into log2(BC) block matmuls - # (the whole tile updates in parallel, no serial row dependency). Rows - # beyond T_eff already have N = 0 (b_gq/b_gk were zeroed by m_c), so those - # rows converge to the identity exactly as the guarded serial loop did. - # - # Precision: the squarings run at tf32 (SOLVE_TRIL_DOT_PRECISION on tf32- - # capable arch) via safe_matmul, giving real HMMA tensor-core ops at M = BC - # (>=16; fp32 operands would fall back to SIMT even at M=32). This matches - # the tolerance across the full fwd matrix (incl. raw non-L2-normed k in the - # dispatch suite, RMS ratio well under 2e-2) once the series is anchored on - # the correct sign (N = -Akk). + # Precision: the squarings run at tf32 via safe_matmul (fp32 operands + # would fall back to SIMT even at M=32). b_N = -b_Akk # N := -Akk, so (I + Akk)^-1 = (I - N)^-1 = sum_k N^k b_Ai = ct.astype(m_I, ct.float32) + b_N # (I + N) # Squaring stages: `range(2, BC)` traces as a compile-time-bounded loop with @@ -1607,7 +1507,7 @@ def chunk_kda_fwd_kernel_inter_diag_compute_solve( # --------------------------------------------------------------------------- # chunk_kda_fwd_kernel_inter_solve_fused -- off-diagonal Akk + merged solve # --------------------------------------------------------------------------- -def _load_bc_bk(arr, base, row0, col0, stride_row, BC: ConstInt, BK: ConstInt, T_eff, K: ConstInt): +def load_bc_bk(arr, base, row0, col0, stride_row, BC: ConstInt, BK: ConstInt, T_eff, K: ConstInt): # Boundary-checked (BC,BK) block load at (row0,col0) of the (T,K) view # (row stride stride_row), cast to fp32, on the flattened 1-D view. o_r = ct.arange(BC, dtype=ct.int32) @@ -1619,7 +1519,7 @@ def _load_bc_bk(arr, base, row0, col0, stride_row, BC: ConstInt, BK: ConstInt, T return ct.astype(ct.gather(arr, off, mask=mask, check_bounds=False, padding_value=0.0), ct.float32) -def _store_bc_bc(arr, base, row0, col0, stride_row, blk, BC: ConstInt, BT_or_BC: ConstInt, T_eff): +def store_bc_bc(arr, base, row0, col0, stride_row, blk, BC: ConstInt, BT_or_BC: ConstInt, T_eff): o_r = ct.arange(BC, dtype=ct.int32) o_c = ct.arange(BC, dtype=ct.int32) rows = row0 + o_r @@ -1629,37 +1529,6 @@ def _store_bc_bc(arr, base, row0, col0, stride_row, blk, BC: ConstInt, BT_or_BC: ct.scatter(arr, off, ct.astype(blk, arr.dtype), mask=mask, check_bounds=False) -def _cat_pow2(pieces, axis): - """Concatenate a power-of-two-length tuple of equal-shaped 16x16 leaf tiles - via a balanced binary tree (fully unrolled, no recursion). ct.cat takes at - most 2 operands and requires every intermediate shape to be a power of two, - so a balanced tree (16 -> 32 -> 64) is required (a left-fold would produce - illegal sizes like 48). Mirrors the proven helper in the sibling GDN - chunk kernels.""" - n = len(pieces) - if n == 1: - return pieces[0] - if n == 2: - return ct.cat((pieces[0], pieces[1]), axis) - # n == 4 - return ct.cat( - (ct.cat((pieces[0], pieces[1]), axis), ct.cat((pieces[2], pieces[3]), axis)), - axis, - ) - - -def _store_bt_bt(arr, base, row0, stride_row, blk, BT: ConstInt, T_eff): - # Single masked scatter of a full [BT, BT] tile to a head-interleaved flat - # view (row stride `stride_row` = HV*BT). Replaces the up-to-12 per-sub-block - # `_store_bc_bc` scatters, which serialise writes to the same Akk output. - o_r = ct.arange(BT, dtype=ct.int32) - o_c = ct.arange(BT, dtype=ct.int32) - rows = row0 + o_r - off = base + rows[:, None] * stride_row + o_c[None, :] - mask = ct.broadcast_to((rows < T_eff)[:, None], (BT, BT)) - ct.scatter(arr, off, ct.astype(blk, arr.dtype), mask=mask, check_bounds=False) - - @ct.kernel def chunk_kda_fwd_kernel_inter_solve_fused( q, @@ -1727,18 +1596,18 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Aqk32 = z b_Akk32 = z - # ---- off-diagonal blocks ---- + # ---- off-diagonal blocks ----------------------------------------------------- num_k = (K + BK - 1) // BK for i_k in range(num_k): kk = i_k * BK + o_k m_k = kk < K - b_k0 = _load_bc_bk(k, k_base, i_tc0, i_k * BK, H * K, BC, BK, T_eff, K) - b_g0 = _load_bc_bk(g, g_base, i_tc0, i_k * BK, HV * K, BC, BK, T_eff, K) + b_k0 = load_bc_bk(k, k_base, i_tc0, i_k * BK, H * K, BC, BK, T_eff, K) + b_g0 = load_bc_bk(g, g_base, i_tc0, i_k * BK, HV * K, BC, BK, T_eff, K) if i_tc1 < T_eff: - b_q1 = _load_bc_bk(q, q_base, i_tc1, i_k * BK, H * K, BC, BK, T_eff, K) - b_k1 = _load_bc_bk(k, k_base, i_tc1, i_k * BK, H * K, BC, BK, T_eff, K) - b_g1 = _load_bc_bk(g, g_base, i_tc1, i_k * BK, HV * K, BC, BK, T_eff, K) + b_q1 = load_bc_bk(q, q_base, i_tc1, i_k * BK, H * K, BC, BK, T_eff, K) + b_k1 = load_bc_bk(k, k_base, i_tc1, i_k * BK, H * K, BC, BK, T_eff, K) + b_g1 = load_bc_bk(g, g_base, i_tc1, i_k * BK, HV * K, BC, BK, T_eff, K) b_gn1 = ct.astype( ct.gather(g, g_base + i_tc1 * (HV * K) + kk, mask=m_k, check_bounds=False, padding_value=0.0), ct.float32, @@ -1749,9 +1618,9 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk10 = safe_mma(b_k1 * b_gqn, b_kgt, b_Akk10) if NC >= 3 and i_tc2 < T_eff: - b_q2 = _load_bc_bk(q, q_base, i_tc2, i_k * BK, H * K, BC, BK, T_eff, K) - b_k2 = _load_bc_bk(k, k_base, i_tc2, i_k * BK, H * K, BC, BK, T_eff, K) - b_g2 = _load_bc_bk(g, g_base, i_tc2, i_k * BK, HV * K, BC, BK, T_eff, K) + b_q2 = load_bc_bk(q, q_base, i_tc2, i_k * BK, H * K, BC, BK, T_eff, K) + b_k2 = load_bc_bk(k, k_base, i_tc2, i_k * BK, H * K, BC, BK, T_eff, K) + b_g2 = load_bc_bk(g, g_base, i_tc2, i_k * BK, HV * K, BC, BK, T_eff, K) b_gn2 = ct.astype( ct.gather(g, g_base + i_tc2 * (HV * K) + kk, mask=m_k, check_bounds=False, padding_value=0.0), ct.float32, @@ -1767,9 +1636,9 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk21 = safe_mma(b_kg2, b_kgt, b_Akk21) if NC >= 4 and i_tc3 < T_eff: - b_q3 = _load_bc_bk(q, q_base, i_tc3, i_k * BK, H * K, BC, BK, T_eff, K) - b_k3 = _load_bc_bk(k, k_base, i_tc3, i_k * BK, H * K, BC, BK, T_eff, K) - b_g3 = _load_bc_bk(g, g_base, i_tc3, i_k * BK, HV * K, BC, BK, T_eff, K) + b_q3 = load_bc_bk(q, q_base, i_tc3, i_k * BK, H * K, BC, BK, T_eff, K) + b_k3 = load_bc_bk(k, k_base, i_tc3, i_k * BK, H * K, BC, BK, T_eff, K) + b_g3 = load_bc_bk(g, g_base, i_tc3, i_k * BK, HV * K, BC, BK, T_eff, K) b_gn3 = ct.astype( ct.gather(g, g_base + i_tc3 * (HV * K) + kk, mask=m_k, check_bounds=False, padding_value=0.0), ct.float32, @@ -1787,17 +1656,17 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Aqk32 = safe_mma(b_qg3, b_kgt, b_Aqk32) b_Akk32 = safe_mma(b_kg3, b_kgt, b_Akk32) - # ---- save off-diagonal Aqk blocks, scale Akk by beta ---- + # ---- save off-diagonal Aqk blocks, scale Akk by Beta ------------------------- if i_tc1 < T_eff: - _store_bc_bc(Aqk, Aqk_base, i_tc1, 0, HV * BT, b_Aqk10 * scale, BC, BT, T_eff) + store_bc_bc(Aqk, Aqk_base, i_tc1, 0, HV * BT, b_Aqk10 * scale, BC, BT, T_eff) b_b1 = ct.astype( ct.gather(beta, beta_base + (i_tc1 + o_i) * HV, mask=m_tc1, check_bounds=False, padding_value=0.0), ct.float32, ) b_Akk10 = b_Akk10 * b_b1[:, None] if NC >= 3 and i_tc2 < T_eff: - _store_bc_bc(Aqk, Aqk_base, i_tc2, 0, HV * BT, b_Aqk20 * scale, BC, BT, T_eff) - _store_bc_bc(Aqk, Aqk_base, i_tc2, BC, HV * BT, b_Aqk21 * scale, BC, BT, T_eff) + store_bc_bc(Aqk, Aqk_base, i_tc2, 0, HV * BT, b_Aqk20 * scale, BC, BT, T_eff) + store_bc_bc(Aqk, Aqk_base, i_tc2, BC, HV * BT, b_Aqk21 * scale, BC, BT, T_eff) b_b2 = ct.astype( ct.gather(beta, beta_base + (i_tc2 + o_i) * HV, mask=m_tc2, check_bounds=False, padding_value=0.0), ct.float32, @@ -1805,9 +1674,9 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk20 = b_Akk20 * b_b2[:, None] b_Akk21 = b_Akk21 * b_b2[:, None] if NC >= 4 and i_tc3 < T_eff: - _store_bc_bc(Aqk, Aqk_base, i_tc3, 0, HV * BT, b_Aqk30 * scale, BC, BT, T_eff) - _store_bc_bc(Aqk, Aqk_base, i_tc3, BC, HV * BT, b_Aqk31 * scale, BC, BT, T_eff) - _store_bc_bc(Aqk, Aqk_base, i_tc3, 2 * BC, HV * BT, b_Aqk32 * scale, BC, BT, T_eff) + store_bc_bc(Aqk, Aqk_base, i_tc3, 0, HV * BT, b_Aqk30 * scale, BC, BT, T_eff) + store_bc_bc(Aqk, Aqk_base, i_tc3, BC, HV * BT, b_Aqk31 * scale, BC, BT, T_eff) + store_bc_bc(Aqk, Aqk_base, i_tc3, 2 * BC, HV * BT, b_Aqk32 * scale, BC, BT, T_eff) b_b3 = ct.astype( ct.gather(beta, beta_base + (i_tc3 + o_i) * HV, mask=m_tc3, check_bounds=False, padding_value=0.0), ct.float32, @@ -1816,13 +1685,13 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk31 = b_Akk31 * b_b3[:, None] b_Akk32 = b_Akk32 * b_b3[:, None] - # ---- load diagonal inverse blocks from Akkd (fp32) ---- - b_Ai00 = _load_bc_bk(Akkd, Akkd_base, i_tc0, 0, HV * BC, BC, BC, T_eff, BC) - b_Ai11 = _load_bc_bk(Akkd, Akkd_base, i_tc1, 0, HV * BC, BC, BC, T_eff, BC) - b_Ai22 = _load_bc_bk(Akkd, Akkd_base, i_tc2, 0, HV * BC, BC, BC, T_eff, BC) if NC >= 3 else z - b_Ai33 = _load_bc_bk(Akkd, Akkd_base, i_tc3, 0, HV * BC, BC, BC, T_eff, BC) if NC >= 4 else z + # ---- load diagonal inverse blocks from Akkd (fp32) --------------------------- + b_Ai00 = load_bc_bk(Akkd, Akkd_base, i_tc0, 0, HV * BC, BC, BC, T_eff, BC) + b_Ai11 = load_bc_bk(Akkd, Akkd_base, i_tc1, 0, HV * BC, BC, BC, T_eff, BC) + b_Ai22 = load_bc_bk(Akkd, Akkd_base, i_tc2, 0, HV * BC, BC, BC, T_eff, BC) if NC >= 3 else z + b_Ai33 = load_bc_bk(Akkd, Akkd_base, i_tc3, 0, HV * BC, BC, BC, T_eff, BC) if NC >= 4 else z - # ---- forward substitution on diagonals (only when gate not pre-solved) ---- + # ---- forward substitution on diagonals (only when gate not pre-solved) ------- if not USE_SAFE_GATE: m_A = o_i[:, None] > o_i[None, :] m_I = o_i[:, None] == o_i[None, :] @@ -1836,9 +1705,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # Counted loops over the static [BC] forward-substitution window with a # runtime guard (i < T_eff - i_tc0). A runtime upper bound (ct_min(..., # T_eff - i_tc0)) lowers to per-iteration branches tileiras can't - # unroll/predicate; the counted form stays fully unrolled and is - # numerically identical (inactive rows are skipped by the guard). See - # MEMORY counted-loop note. + # unroll/predicate; the counted form stays fully unrolled. for i in range(2, BC): if i < T_eff - i_tc0: b_a00 = -ct.astype( @@ -1885,7 +1752,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( if NC >= 4: b_Ai33 = b_Ai33 + ct.astype(m_I, ct.float32) - # ---- merged inverse using off-diagonals (tf32) ---- + # ---- merged inverse using off-diagonals (tf32) ------------------------------- b_Ai10 = -safe_matmul(safe_matmul(b_Ai11, b_Akk10), b_Ai00) b_Ai20 = z b_Ai21 = z @@ -1900,23 +1767,23 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Ai31 = -safe_matmul(b_Ai33, safe_matmul(b_Akk31, b_Ai11) + safe_matmul(b_Akk32, b_Ai21)) b_Ai30 = -safe_matmul(b_Ai33, safe_matmul(b_Akk30, b_Ai00) + safe_matmul(b_Akk31, b_Ai10) + safe_matmul(b_Akk32, b_Ai20)) - # ---- store full Akk_inv to Akk ---- - _store_bc_bc(Akk, Akk_base, i_tc0, 0, HV * BT, b_Ai00, BC, BT, T_eff) - _store_bc_bc(Akk, Akk_base, i_tc1, 0, HV * BT, b_Ai10, BC, BT, T_eff) - _store_bc_bc(Akk, Akk_base, i_tc1, BC, HV * BT, b_Ai11, BC, BT, T_eff) + # ---- store full Akk_inv to Akk ----------------------------------------------- + store_bc_bc(Akk, Akk_base, i_tc0, 0, HV * BT, b_Ai00, BC, BT, T_eff) + store_bc_bc(Akk, Akk_base, i_tc1, 0, HV * BT, b_Ai10, BC, BT, T_eff) + store_bc_bc(Akk, Akk_base, i_tc1, BC, HV * BT, b_Ai11, BC, BT, T_eff) if NC >= 3: - _store_bc_bc(Akk, Akk_base, i_tc2, 0, HV * BT, b_Ai20, BC, BT, T_eff) - _store_bc_bc(Akk, Akk_base, i_tc2, BC, HV * BT, b_Ai21, BC, BT, T_eff) - _store_bc_bc(Akk, Akk_base, i_tc2, 2 * BC, HV * BT, b_Ai22, BC, BT, T_eff) + store_bc_bc(Akk, Akk_base, i_tc2, 0, HV * BT, b_Ai20, BC, BT, T_eff) + store_bc_bc(Akk, Akk_base, i_tc2, BC, HV * BT, b_Ai21, BC, BT, T_eff) + store_bc_bc(Akk, Akk_base, i_tc2, 2 * BC, HV * BT, b_Ai22, BC, BT, T_eff) if NC >= 4: - _store_bc_bc(Akk, Akk_base, i_tc3, 0, HV * BT, b_Ai30, BC, BT, T_eff) - _store_bc_bc(Akk, Akk_base, i_tc3, BC, HV * BT, b_Ai31, BC, BT, T_eff) - _store_bc_bc(Akk, Akk_base, i_tc3, 2 * BC, HV * BT, b_Ai32, BC, BT, T_eff) - _store_bc_bc(Akk, Akk_base, i_tc3, 3 * BC, HV * BT, b_Ai33, BC, BT, T_eff) + store_bc_bc(Akk, Akk_base, i_tc3, 0, HV * BT, b_Ai30, BC, BT, T_eff) + store_bc_bc(Akk, Akk_base, i_tc3, BC, HV * BT, b_Ai31, BC, BT, T_eff) + store_bc_bc(Akk, Akk_base, i_tc3, 2 * BC, HV * BT, b_Ai32, BC, BT, T_eff) + store_bc_bc(Akk, Akk_base, i_tc3, 3 * BC, HV * BT, b_Ai33, BC, BT, T_eff) # =========================================================================== -# chunk_bwd (dAv) -- dAqk = do @ v.T ; dv = A @ do +# chunk_bwd (dAv) -- dAqk = dO @ V^T ; dV = A @ dO # =========================================================================== @@ -2032,7 +1899,7 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( i_hv = ct.bid(1) i_h = i_hv // (HV // H) - # grid dim-0 is the GLOBAL chunk index; h/dh are laid out per global + # grid dim-0 is the GLOBAL chunk index; H/dH are laid out per global # chunk (written with chunk_offsets+local). Capture it before i_t is # reassigned to the per-sequence (local) chunk index. i_tg = i_t @@ -2068,7 +1935,7 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( b_beta = ct.gather(beta, beta_base + o_t * HV, mask=m_t, check_bounds=False, padding_value=0.0) - # p_A transposed: view (BT, T) stride (1, HV*BT), block (BT, BT) at (0, off_t) + # A transposed: view (BT, T) stride (1, HV*BT), block (BT, BT) at (0, off_t) A_row = o_bt[:, None] A_time = (off_t + o_bt)[None, :] A_m = ct.broadcast_to((o_bt < BT)[:, None], (BT, BT)) & ((off_t + o_bt)[None, :] < T) @@ -2112,13 +1979,13 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( off_v = i_v * BV v_cols = (off_v + o_bv)[None, :] if STATE_V_FIRST: - # h/dh: view (V, K) stride (K, 1), block (BV, BK) at (off_v, off_k) + # H/dH: view (V, K) stride (K, 1), block (BV, BK) at (off_v, off_k) h_rows = (off_v + o_bv)[:, None] h_m = ct.broadcast_to((off_v + o_bv)[:, None] < V, (BV, BK)) & ct.broadcast_to(m_k[None, :], (BV, BK)) b_h = ct.gather(h, h_base + h_rows * K + o_k[None, :], mask=h_m, check_bounds=False, padding_value=0.0) b_dh = ct.gather(dh, dh_base + h_rows * K + o_k[None, :], mask=h_m, check_bounds=False, padding_value=0.0) else: - # h/dh transposed: view (V, K) stride (1, V), block (BV, BK) at (off_v, off_k) + # H/dH transposed: view (V, K) stride (1, V), block (BV, BK) at (off_v, off_k) h_row = (off_v + o_bv)[:, None] h_col = o_k[None, :] h_m = ct.broadcast_to((off_v + o_bv)[:, None] < V, (BV, BK)) & ct.broadcast_to(m_k[None, :], (BV, BK)) @@ -2152,9 +2019,7 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( b_dq = safe_mma(b_do, ct.astype(b_h, b_do.dtype), b_dq) b_dk = safe_mma(b_v_new, ct.astype(b_dh, b_v_new.dtype), b_dk) b_dw = safe_mma(ct.astype(b_dv, b_v_new.dtype), ct.astype(b_h, b_v_new.dtype), b_dw) - # TODO(cutile): an explicit debug-barrier here was - # dropped -- no cuTile equivalent; the b_v reuse below relies on - # program order within the block. + # the b_v reuse below relies on program order within the block. if i_k == 0: b_v = ct.gather( v, @@ -2236,7 +2101,7 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( # =========================================================================== -# chunk_bwd (intra) -- dq/dk/dg/db from dAqk/dAkk +# chunk_bwd (intra) -- dQ/dK/dG/dBeta from dAqk/dAkk # =========================================================================== @@ -2340,7 +2205,7 @@ def chunk_kda_bwd_kernel_intra( b_dk2 = b_dk2 * b_gqn o_i = o_bc - # current block q, k + # current block Q, K b_q = ct.gather(q, q_base + cur_rows * (H * K) + cur_cols, mask=m_cur, check_bounds=False, padding_value=0.0) b_k = ct.gather(k, k_base + cur_rows * (H * K) + cur_cols, mask=m_cur, check_bounds=False, padding_value=0.0) @@ -2405,7 +2270,7 @@ def chunk_kda_bwd_kernel_intra( b_dg2 = ct.astype(b_q, ct.float32) * b_dq2 - # dq2 = b_dq2 + dq + # dQ2 = b_dq2 + dQ dq_off = dq_base + cur_rows * (HV * K) + cur_cols b_dq_prev = ct.astype(ct.gather(dq, dq_off, mask=m_cur, check_bounds=False, padding_value=0.0), ct.float32) b_dq2_out = b_dq2 + b_dq_prev @@ -2488,8 +2353,7 @@ def chunk_kda_bwd_kernel_intra( # A maskless gather still emits a default `other` (padding_value=0) # operand; without a matching `mask` the tileiras LDGSTS lowering # asserts (llOthers non-empty, llMasks empty). Supply an explicit - # row-validity mask. Values where the mask is false are zeroed but - # never used (m_i = o_i <= j gates the contribution below). + # row-validity mask. m_row = (i_ti + o_bc) < T_eff b_dAqk = ct.gather(dAqk, dAqk_base + base_o + j * (HV * BT), mask=m_row, check_bounds=False, padding_value=0.0) b_dAkk = ct.gather(dAkk, dAkk_base + base_o + j * (HV * BT), mask=m_row, check_bounds=False, padding_value=0.0) @@ -2673,28 +2537,21 @@ def l2norm_fwd( stream=None, ): stream = 0 if stream is None else stream - if out is None or rstd_out is None: - raise ValueError("l2norm_fwd requires pre-allocated out= and rstd_out= buffers") x_shape_og = x.shape - x = x.reshape(-1, x.shape[-1]).contiguous() + x = reshaped(x, (-1, x.shape[-1])) y = reshaped(out, tuple(x.shape)) T, D = x.shape[0], x.shape[-1] MAX_FUSED_SIZE = 65536 // x.element_size() - BD = min(MAX_FUSED_SIZE, _next_power_of_2(D)) - if D > BD: - raise RuntimeError("This layer doesn't support feature dim >= 64KB.") + BD = min(MAX_FUSED_SIZE, next_power_of_2(D)) rstd = reshaped(rstd_out, (T,)) if D <= 512: BT = 32 - grid = (_cdiv(T, BT),) + grid = (cdiv(T, BT),) # l2norm_fwd is a MEMORY-bound row reduction (each row loaded once, - # reduced, written once). Fixed occ=1/nww=4 ran ~4x slower than - # tuned. Sweep occupancy x num_worker_warps (launch hints only, single - # kernel unchanged); the default (occ=1,nww=4) is tried first so a - # non-improving shape keeps the base behaviour. Memory-bound => native/ - # higher occupancy hides DRAM latency (do NOT force occ=1 here). - _l2_key = ("l2norm_fwd_kernel", int(D), int(BD), int(BT), str(x.dtype), _dev_id(x)) - _autotuned_launch(l2norm_fwd_kernel, _l2_key, grid, (x, y, rstd, float(eps), T, D, BD, BT), occ_choices=(1, 2, 4, 8), nww_choices=(4,), stream=stream) + # reduced, written once); higher occupancy hides DRAM latency (do NOT + # force occ=1 here). + _l2_key = ("l2norm_fwd_kernel", int(D), int(BD), int(BT), str(x.dtype), dev_id(x)) + autotuned_launch(l2norm_fwd_kernel, _l2_key, grid, (x, y, rstd, float(eps), T, D, BD, BT), occ_choices=(1, 2, 4, 8), nww_choices=(4,), stream=stream) else: ct.launch(stream, (T,), l2norm_fwd_kernel1, (x, y, rstd, float(eps), D, BD)) return y.view(x_shape_og), rstd.view(x_shape_og[:-1]) @@ -2703,20 +2560,16 @@ def l2norm_fwd( def l2norm_bwd(y, rstd, dy, eps: float = 1e-6, out=None, stream=None): stream = 0 if stream is None else stream y_shape_og = y.shape - y = y.reshape(-1, dy.shape[-1]).contiguous() - dy = dy.reshape(-1, dy.shape[-1]).contiguous() - if out is None: - raise ValueError("l2norm_bwd requires a pre-allocated out= buffer") + y = y.reshape(-1, dy.shape[-1]) + dy = dy.reshape(-1, dy.shape[-1]) dx = reshaped(out, tuple(y.shape)) T, D = y.shape[0], y.shape[-1] MAX_FUSED_SIZE = 65536 // y.element_size() - BD = min(MAX_FUSED_SIZE, _next_power_of_2(D)) - if D > BD: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - rstd_flat = rstd.reshape(-1).contiguous() + BD = min(MAX_FUSED_SIZE, next_power_of_2(D)) + rstd_flat = rstd.reshape(-1) if D <= 512: BT = 32 - grid = (_cdiv(T, BT),) + grid = (cdiv(T, BT),) ct.launch(stream, grid, l2norm_bwd_kernel, (y, rstd_flat, dy, dx, float(eps), T, D, BD, BT)) else: ct.launch(stream, (T,), l2norm_bwd_kernel1, (y, rstd_flat, dy, dx, float(eps), D, BD)) @@ -2726,37 +2579,33 @@ def l2norm_bwd(y, rstd, dy, eps: float = 1e-6, out=None, stream=None): # --------------------------------------------------------------------------- # fused_beta_sigmoid # --------------------------------------------------------------------------- -_BETA_SIGMOID_BLOCK_SIZE = 2048 +BETA_SIGMOID_BLOCK_SIZE = 2048 def fused_beta_sigmoid_fwd(x, scale: float = 1.0, out=None, stream=None): stream = 0 if stream is None else stream - if out is None: - raise ValueError("fused_beta_sigmoid_fwd requires a pre-allocated out= buffer (fp32, x-shaped)") y = reshaped(out, tuple(x.shape)) n = x.numel() - grid = (_cdiv(n, _BETA_SIGMOID_BLOCK_SIZE),) + grid = (cdiv(n, BETA_SIGMOID_BLOCK_SIZE),) ct.launch( stream, grid, fused_beta_sigmoid_fwd_kernel, - (x.reshape(-1), y.reshape(-1), float(scale), n, _BETA_SIGMOID_BLOCK_SIZE), + (reshaped(x, (-1,)), y.reshape(-1), float(scale), n, BETA_SIGMOID_BLOCK_SIZE), ) return y def fused_beta_sigmoid_bwd(x, dy, scale: float = 1.0, out=None, stream=None): stream = 0 if stream is None else stream - if out is None: - raise ValueError("fused_beta_sigmoid_bwd requires a pre-allocated out= buffer (x-shaped)") dx = reshaped(out, tuple(x.shape)) n = x.numel() - grid = (_cdiv(n, _BETA_SIGMOID_BLOCK_SIZE),) + grid = (cdiv(n, BETA_SIGMOID_BLOCK_SIZE),) ct.launch( stream, grid, fused_beta_sigmoid_bwd_kernel, - (x.reshape(-1), dy.reshape(-1), dx.reshape(-1), float(scale), n, _BETA_SIGMOID_BLOCK_SIZE), + (reshaped(x, (-1,)), dy.reshape(-1), dx.reshape(-1), float(scale), n, BETA_SIGMOID_BLOCK_SIZE), ) return dx @@ -2770,7 +2619,7 @@ def fused_beta_sigmoid(x, scale: float = 1.0, out=None, stream=None): # --------------------------------------------------------------------------- # chunk_local_cumsum (vector gate) / kda_gate_chunk_cumsum / kda_gate_bwd # --------------------------------------------------------------------------- -_BS_LIST_DEFAULT = 32 +BS_LIST_DEFAULT = 32 def chunk_local_cumsum_vector( @@ -2786,26 +2635,22 @@ def chunk_local_cumsum_vector( stream = 0 if stream is None else stream T, H, S = g.shape BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") - if out is None: - raise ValueError("chunk_local_cumsum requires a pre-allocated out= buffer (g-shaped)") NT = len(chunk_indices) assert chunk_size == 2 ** (chunk_size.bit_length() - 1), "chunk_size must be a power of 2" - BS = min(_BS_LIST_DEFAULT, _next_power_of_2(S)) + BS = min(BS_LIST_DEFAULT, next_power_of_2(S)) g_org = g g_out = reshaped(out, (T, H, S)) scale_val = float(scale) if scale is not None else 0.0 has_scale = int(scale is not None) - cu_arg = _i32_flat(cu_seqlens) - ci_arg = _i32_flat(chunk_indices) - grid = (_cdiv(S, BS), NT, H) + cu_arg = i32_flat(cu_seqlens) + ci_arg = i32_flat(chunk_indices) + grid = (cdiv(S, BS), NT, H) ct.launch( stream, grid, chunk_local_cumsum_vector_kernel, ( - g_org.reshape(-1), + reshaped(g_org, (-1,)), g_out.reshape(-1), scale_val, cu_arg, @@ -2832,7 +2677,6 @@ def chunk_local_cumsum( stream=None, ): stream = 0 if stream is None else stream - assert len(g.shape) == 3, f"KDA chunk_local_cumsum expects a (T, HV, K) vector gate, got shape {tuple(g.shape)}" return chunk_local_cumsum_vector( g=g, chunk_size=chunk_size, @@ -2861,27 +2705,23 @@ def kda_gate_chunk_cumsum( stream = 0 if stream is None else stream T, H, S = g.shape BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") - if out is None: - raise ValueError("kda_gate_chunk_cumsum requires a pre-allocated out= buffer (fp32, g-shaped)") NT = len(chunk_indices) assert chunk_size == 2 ** (chunk_size.bit_length() - 1), "chunk_size must be a power of 2" - BS = min(_BS_LIST_DEFAULT, _next_power_of_2(S)) + BS = min(BS_LIST_DEFAULT, next_power_of_2(S)) g_out = reshaped(out, (T, H, S)) - dt_arg = _opt(dt_bias, bufs, _dtname(A_log)).reshape(-1) + dt_arg = reshaped(opt(dt_bias, bufs, dtname(A_log)), (-1,)) scale_val = float(scale) if scale is not None else 0.0 lb_val = float(lower_bound) if lower_bound is not None else 0.0 - cu_arg = _i32_flat(cu_seqlens) - ci_arg = _i32_flat(chunk_indices) - grid = (_cdiv(S, BS), NT, H) + cu_arg = i32_flat(cu_seqlens) + ci_arg = i32_flat(chunk_indices) + grid = (cdiv(S, BS), NT, H) ct.launch( stream, grid, kda_gate_chunk_cumsum_vector_kernel, ( - g.reshape(-1), - A_log.reshape(-1), + reshaped(g, (-1,)), + reshaped(A_log, (-1,)), dt_arg, g_out.reshape(-1), scale_val, @@ -2909,15 +2749,13 @@ def kda_gate_bwd(g, A_log, dt_bias=None, dyg=None, lower_bound=None, dg_out=None H, K = g.shape[-2:] T = g.numel() // (H * K) BT = 32 - NT = _cdiv(T, BT) - BD = _next_power_of_2(K) - if dg_out is None or dA_out is None or (dt_bias is not None and dbias_out is None): - raise ValueError("kda_gate_bwd requires pre-allocated dg_out=, dA_out= (and dbias_out= with dt_bias)") + NT = cdiv(T, BT) + BD = next_power_of_2(K) dg = reshaped(dg_out, tuple(g.shape)) dA_nt = reshaped(bufs["dA_gate"], (NT, H)) db_nt = reshaped(bufs["db_gate"], (NT, H * K)) if dt_bias is not None else None - dt_arg = _opt(dt_bias, bufs, _dtname(A_log)).reshape(-1) - db_arg = db_nt.reshape(-1) if db_nt is not None else _dummy("float32", bufs) + dt_arg = reshaped(opt(dt_bias, bufs, dtname(A_log)), (-1,)) + db_arg = db_nt.reshape(-1) if db_nt is not None else dummy("float32", bufs) lb_val = float(lower_bound) if lower_bound is not None else 0.0 grid = (NT, H) ct.launch( @@ -2925,10 +2763,10 @@ def kda_gate_bwd(g, A_log, dt_bias=None, dyg=None, lower_bound=None, dg_out=None grid, kda_gate_bwd_kernel, ( - g.reshape(-1), - A_log.reshape(-1), + reshaped(g, (-1,)), + reshaped(A_log, (-1,)), dt_arg, - dyg.reshape(-1), + reshaped(dyg, (-1,)), dg.reshape(-1), dA_nt.reshape(-1), db_arg, @@ -2955,28 +2793,26 @@ def chunk_gla_fwd_o_gk(q, v, g, A, h, scale, state_v_first=False, cu_seqlens=Non stream = 0 if stream is None else stream T, H, K, HV, V = *q.shape, v.shape[1], v.shape[-1] BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") NT = len(chunk_indices) o = bufs["o"] zero_fill(o, stream=stream) - BK = min(max(_next_power_of_2(K), 16), 64) + BK = min(max(next_power_of_2(K), 16), 64) # BV=128 when V<=128 removes the V grid-split # (grid dim0 -> 1), doubling work/block but halving launched blocks. BV = 128 if V <= 128 else 64 - dev = _dev_id(q) - cu_arg = _i32_flat(cu_seqlens) - ci_arg = _i32_flat(chunk_indices) + dev = dev_id(q) + cu_arg = i32_flat(cu_seqlens) + ci_arg = i32_flat(chunk_indices) - grid = (_cdiv(V, BV), NT, HV) + grid = (cdiv(V, BV), NT, HV) # Kernel uses flat element-offset gather/scatter; pass 1-D views so the # cuTile index-tuple rank (1) matches the array rank. - _q_arg = q.reshape(-1) + _q_arg = reshaped(q, (-1,)) _v_arg = v.reshape(-1) _g_arg = g.reshape(-1) _h_arg = h.reshape(-1) - _o_arg = o.reshape(-1) + _o_arg = reshaped(o, (-1,)) _A_arg = A.reshape(-1) _o_args = ( _q_arg, @@ -2998,8 +2834,7 @@ def chunk_gla_fwd_o_gk(q, v, g, A, h, scale, state_v_first=False, cu_seqlens=Non int(state_v_first), ) # Launch-hint autotune (occupancy x num_worker_warps) on this - # output-projection kernel. Default config - # is tried first so non-improving shapes are unchanged. + # output-projection kernel. _o_key = ( "chunk_gla_fwd_kernel_o", int(H), @@ -3013,7 +2848,7 @@ def chunk_gla_fwd_o_gk(q, v, g, A, h, scale, state_v_first=False, cu_seqlens=Non str(q.dtype), str(dev), ) - _autotuned_launch(chunk_gla_fwd_kernel_o, _o_key, grid, _o_args, occ_choices=(1, 2, 4), nww_choices=(4,), stream=stream) + autotuned_launch(chunk_gla_fwd_kernel_o, _o_key, grid, _o_args, occ_choices=(1, 2, 4), nww_choices=(4,), stream=stream) return o @@ -3040,30 +2875,28 @@ def chunk_gated_delta_rule_fwd_h( stream = 0 if stream is None else stream T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[1] BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") N, NT = len(cu_seqlens) - 1, len(chunk_indices) chunk_offsets = bufs["chunk_offsets"] # Full-width state/K tile: BK = next_pow2(K). The blockdim64 kernel carries the # KV state as a single (BV, BK)/(BK, BV) tile (no K<=256 cap); K-axis loads # zero-pad [K:BK] and stores drop it, so any K is supported. - BK = _next_power_of_2(K) + BK = next_power_of_2(K) state_shape = (N, HV, V, K) if state_v_first else (N, HV, K, V) - h = reshaped(bufs["h"], (NT, HV) + state_shape[2:]) - final_state = reshaped(bufs["fs"], state_shape) if output_final_state else None + h = reshaped(bufs["state_checkpoints"], (NT, HV) + state_shape[2:]) + final_state = reshaped(bufs["final_state"], state_shape) if output_final_state else None if final_state is not None: zero_fill(final_state, stream=stream) v_new = reshaped(bufs["v_new"], (T, HV, V)) if save_new_value else None - dev = _dev_id(k) - vnew_arg = v_new if v_new is not None else _dummy(_dtname(u), bufs) - g_arg = _opt(g, bufs) - gk_arg = _opt(gk, bufs) - h0_arg = initial_state if initial_state is not None else _dummy("float32", bufs) - ht_arg = final_state if final_state is not None else _dummy("float32", bufs) - cu_arg = _i32_flat(cu_seqlens) - co_arg = _i32_flat(chunk_offsets) + dev = dev_id(k) + vnew_arg = v_new if v_new is not None else dummy(dtname(u), bufs) + g_arg = opt(g, bufs) + gk_arg = opt(gk, bufs) + h0_arg = initial_state if initial_state is not None else dummy("float32", bufs) + ht_arg = final_state if final_state is not None else dummy("float32", bufs) + cu_arg = i32_flat(cu_seqlens) + co_arg = i32_flat(chunk_offsets) # Kernel uses flat element-offset gather/scatter; pass 1-D views so the # cuTile index-tuple rank (1) matches the array rank. @@ -3072,22 +2905,17 @@ def chunk_gated_delta_rule_fwd_h( _w_arg = w.reshape(-1) _vnew_arg = vnew_arg.reshape(-1) _h_arg = h.reshape(-1) - _h01d = h0_arg.reshape(-1) - _ht1d = ht_arg.reshape(-1) + _h01d = reshaped(h0_arg, (-1,)) + _ht1d = reshaped(ht_arg, (-1,)) _g1d = g_arg.reshape(-1) _gk1d = gk_arg.reshape(-1) # BV is a V-tiling block width (drives the grid V-fan-out and the V-tile # shapes only; the K axis is always split into fixed 64-wide blocks, so BV - # never changes the numerics). ncu shows this latency-bound state scan is - # grid-under-filled at the previous hard-coded BV=32 (only cdiv(V,32)*N*HV - # CTAs, 1 block/SM) -- shrinking BV multiplies the CTA count and was the - # dominant fwd_h speedup (BV 32->8: 8192/K128 492->260us). - # Sweep BV over {16, 8, 32} (16 first => safe, beats baseline everywhere); - # the tuner picks the per-shape grid fill that best hides the inter-chunk - # latency. + # never changes the numerics). The tuner picks the per-shape grid fill that + # best hides the inter-chunk latency. def _grid_fn(bv): - return (_cdiv(V, bv), N * HV) + return (cdiv(V, bv), N * HV) def _args_fn(bv): return ( @@ -3137,19 +2965,13 @@ def _args_fn(bv): # capped at <= V so a single tile is never larger than V. The grid-fill # tradeoff is shape-dependent, not monotone: shrinking BV multiplies the # V-tile CTA count but each CTA is register-bound (255 reg -> ~1 block/SM) - # and redundantly reloads k/w/g, so past the point where the grid already + # and redundantly reloads K/W/G, so past the point where the grid already # fills the SMs a *larger* BV wins (fewer, fatter CTAs, higher warp - # occupancy). ncu on the dashboard shapes: at (T=2048,V=128,N*HV=64) BV=64 - # (grid 128, warps_active 14%) is 246us vs BV=32 (grid 256, 8%) 289us vs - # BV=16 374us -- so 64 must be a candidate; at (T=1024,V=64,N*HV=64) BV=32 - # (grid 128) 65us beats BV=64 (grid 64) 123us and BV=8 123us. The tuner - # benchmarks every candidate and picks the per-shape winner, so we offer the - # full {16, 8, 32, 64} sweep (mirrors the fwd-helper); 16 stays first as the - # safe DISABLE_TUNE/failure fallback. + # occupancy); 16 stays first as the safe DISABLE_TUNE/failure fallback. _bv_choices = tuple(bv for bv in (16, 8, 32, 64) if bv <= max(V, 8)) if not _bv_choices: _bv_choices = (min(32, V),) - _autotuned_launch_bv(chunk_gated_delta_rule_fwd_kernel_h_blockdim64, _h_key, _bv_choices, _grid_fn, _args_fn, stream=stream) + autotuned_launch_bv(chunk_gated_delta_rule_fwd_kernel_h_blockdim64, _h_key, _bv_choices, _grid_fn, _args_fn, stream=stream) return h, v_new, final_state @@ -3162,7 +2984,7 @@ def chunk_gated_delta_rule_bwd_dhu( g=None, gk=None, h0=None, - dht=None, + dstate_in=None, scale: float | None = None, state_v_first: bool = False, cu_seqlens=None, @@ -3175,21 +2997,19 @@ def chunk_gated_delta_rule_bwd_dhu( T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[1] BT = chunk_size # Full-width state/K tile: BK = next_pow2(K). The blockdim64 dhu kernel carries - # the dh state as a single (BV, BK)/(BK, BV) tile (no K<=256 cap); K-axis loads - # zero-pad/mask [K:BK] and stores drop it, so any K is supported. dh/dh0 + # the dH state as a single (BV, BK)/(BK, BV) tile (no K<=256 cap); K-axis loads + # zero-pad/mask [K:BK] and stores drop it, so any K is supported. dH/dH0 # carves stay K-wide. - BK = _next_power_of_2(K) - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") + BK = next_power_of_2(K) N, NT = len(cu_seqlens) - 1, len(chunk_indices) chunk_offsets = bufs["chunk_offsets"] - dh = reshaped(bufs["dh"], (NT, HV, V, K) if state_v_first else (NT, HV, K, V)) - dh0 = reshaped(bufs["dh0"], tuple(h0.shape)) if h0 is not None else None - dv2 = reshaped(bufs["dv_dhu"], (T, HV, V)) + dh = reshaped(bufs["dstate"], (NT, HV, V, K) if state_v_first else (NT, HV, K, V)) + dh0 = reshaped(bufs["dstate0"], tuple(h0.shape)) if h0 is not None else None + dv2 = reshaped(bufs["dv_dstate_u"], (T, HV, V)) BV = 64 - grid = (_cdiv(V, BV), N * HV) + grid = (cdiv(V, BV), N * HV) ct.launch( stream, @@ -3199,16 +3019,16 @@ def chunk_gated_delta_rule_bwd_dhu( q.reshape(-1), k.reshape(-1), w.reshape(-1), - _opt(g, bufs).reshape(-1), - _opt(gk, bufs).reshape(-1), - _opt(dht, bufs).reshape(-1), - (dh0 if dh0 is not None else _dummy("float32", bufs)).reshape(-1), - do.reshape(-1), + opt(g, bufs).reshape(-1), + opt(gk, bufs).reshape(-1), + reshaped(opt(dstate_in, bufs), (-1,)), + reshaped(dh0 if dh0 is not None else dummy("float32", bufs), (-1,)), + reshaped(do, (-1,)), dh.reshape(-1), dv.reshape(-1), dv2.reshape(-1), - _i32_flat(cu_seqlens), - _i32_flat(chunk_offsets), + i32_flat(cu_seqlens), + i32_flat(chunk_offsets), float(scale), H, HV, @@ -3220,7 +3040,7 @@ def chunk_gated_delta_rule_bwd_dhu( int(g is not None), int(gk is not None), int(h0 is not None), - int(dht is not None), + int(dstate_in is not None), int(state_v_first), ), ) @@ -3237,31 +3057,29 @@ def recompute_w_u_fwd(k, v, beta, A, gk, q=None, cu_seqlens=None, chunk_indices= BT = A.shape[-1] BK = 32 BV = 32 - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") NT = len(chunk_indices) - dev = _dev_id(k) + dev = dev_id(k) w = reshaped(bufs["w"], (T, HV, K)) u = reshaped(bufs["u"], (T, HV, V)) qg = reshaped(bufs["qg"], (T, HV, K)) if q is not None else None kg = reshaped(bufs["kg"], (T, HV, K)) - cu_arg = _i32_flat(cu_seqlens) - ci_arg = _i32_flat(chunk_indices) - q_arg = _opt(q, bufs, _dtname(k)) - qg_arg = qg if qg is not None else _dummy(_dtname(k), bufs) + cu_arg = i32_flat(cu_seqlens) + ci_arg = i32_flat(chunk_indices) + q_arg = opt(q, bufs, dtname(k)) + qg_arg = qg if qg is not None else dummy(dtname(k), bufs) # Kernel uses flat element-offset gather/scatter; pass 1-D views so the # cuTile index-tuple rank (1) matches the array rank. reshape(-1) on these - # contiguous tensors yields storage-aliasing views (outputs w/u/qg/kg too). + # contiguous tensors yields storage-aliasing views (outputs W/U/QG/KG too). _wu_grid = (NT, HV) _wu_args = ( - q_arg.reshape(-1), - k.reshape(-1), + reshaped(q_arg, (-1,)), + reshaped(k, (-1,)), qg_arg.reshape(-1), kg.reshape(-1), - v.reshape(-1), - beta.reshape(-1), + reshaped(v, (-1,)), + reshaped(beta, (-1,)), w.reshape(-1), u.reshape(-1), A.reshape(-1), @@ -3281,7 +3099,6 @@ def recompute_w_u_fwd(k, v, beta, A, gk, q=None, cu_seqlens=None, chunk_indices= ) # Launch-hint autotune (occupancy x num_worker_warps) on this wy_fast # recompute kernel. - # The default config is tried first so non-improving shapes are unchanged. _wu_key = ( "recompute_w_u_fwd_kda_kernel", int(H), @@ -3296,7 +3113,7 @@ def recompute_w_u_fwd(k, v, beta, A, gk, q=None, cu_seqlens=None, chunk_indices= str(k.dtype), str(dev), ) - _autotuned_launch(recompute_w_u_fwd_kda_kernel, _wu_key, _wu_grid, _wu_args, occ_choices=(1, 2, 4, 8), nww_choices=(4,), stream=stream) + autotuned_launch(recompute_w_u_fwd_kda_kernel, _wu_key, _wu_grid, _wu_args, occ_choices=(1, 2, 4, 8), nww_choices=(4,), stream=stream) return w, u, qg, kg @@ -3313,21 +3130,21 @@ def chunk_kda_fwd_intra_token_parallel(q, k, gk, beta, Aqk, Akk, scale, cu_seqle # BC-wide j-loop (gathers reused across heads share the same row indexing). # BH in {1,2,4,8}; HV must be divisible for the grid split. BH = 4 if (HV % 4 == 0) else (2 if (HV % 2 == 0) else 1) - BK = _next_power_of_2(K) - cu_arg = _i32_flat(cu_seqlens) - grid = (T, _cdiv(HV, BH)) + BK = next_power_of_2(K) + cu_arg = i32_flat(cu_seqlens) + grid = (T, cdiv(HV, BH)) # cuTile gather/scatter index-tuple rank must match the array rank, so pass - # pre-flattened views: q/k -> (T*H, K), gk -> (T*HV, K), beta/Aqk/Akk -> 1-D. + # pre-flattened views: Q/K -> (T*H, K), Gk -> (T*HV, K), Beta/Aqk/Akk -> 1-D. # Aqk/Akk are contiguous, so reshape(-1) is a view aliasing the original storage. ct.launch( stream, grid, chunk_kda_fwd_kernel_intra_token_parallel, ( - q.reshape(-1, K), - k.reshape(-1, K), + reshaped(q, (-1, K)), + reshaped(k, (-1, K)), gk.reshape(-1, K), - beta.reshape(-1), + reshaped(beta, (-1,)), Aqk.reshape(-1), Akk.reshape(-1), float(scale), @@ -3369,40 +3186,38 @@ def chunk_kda_fwd_intra( # BC=32 (NC=2) for BT=64,K>=64. With NC=2 there # is exactly ONE off-diagonal pair -> only 2 live [32,32] accumulators (vs 12 # [16,16] at NC=4), and its merged-inverse [32,32]@[32,32] ct.matmul lowers - # to HMMA (M=32) instead of SIMT scalar FADD/FMUL (M=16). See KDA fwd IR/SASS. + # to HMMA (M=32) instead of SIMT scalar FADD/FMUL (M=16). # K<64 falls back to BC=16/NC=4. BC = 32 if BT == 64 and K >= 64 else 16 - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") NT = len(chunk_indices) - NC = _cdiv(BT, BC) + NC = cdiv(BT, BC) # use_split_diag_compute_solve: pre-solve diagonals in # inter_diag_compute_solve so inter_solve_fused can SKIP forward-substitution. use_split_diag_compute_solve = (not safe_gate) and BT == 64 and K >= 64 use_solved_diagonal = safe_gate or use_split_diag_compute_solve - dev = _dev_id(k) + dev = dev_id(k) Aqk = reshaped(bufs["Aqk"], (T, HV, BT)) Akk = reshaped(bufs["Akk"], (T, HV, BT)) zero_fill(Akk, stream=stream) Akkd = reshaped(bufs["Akkd"], (T, HV, BC)) - cu_arg = _i32_flat(cu_seqlens) - ci_arg = _i32_flat(chunk_indices) + cu_arg = i32_flat(cu_seqlens) + ci_arg = i32_flat(chunk_indices) # Step 1: diagonal blocks into Akkd (fp32). When use_solved_diagonal is set # (safe_gate OR the split path) the diagonals are PRE-SOLVED # here so the inter-solve kernel can skip forward-substitution. if use_solved_diagonal: - BK = _next_power_of_2(K) + BK = next_power_of_2(K) # Kernel uses flat element-offset gather/scatter; pass 1-D views so the # cuTile index-tuple rank (1) matches the array rank. _diag_grid = (NT, NC, HV) _diag_args = ( - q.reshape(-1), - k.reshape(-1), + reshaped(q, (-1,)), + reshaped(k, (-1,)), gk.reshape(-1), - beta.reshape(-1), + reshaped(beta, (-1,)), Aqk.reshape(-1), Akkd.reshape(-1), float(scale), @@ -3416,8 +3231,7 @@ def chunk_kda_fwd_intra( BK, ) # Launch-hint autotune (occupancy x num_worker_warps) on this aux - # kernel; fixed default hints (occ=1, nww=4) ran ~2.9x slower. The - # default config is tried first so non-improving shapes are unchanged. + # kernel. _diag_key = ( "chunk_kda_fwd_kernel_inter_diag_compute_solve", int(H), @@ -3430,7 +3244,7 @@ def chunk_kda_fwd_intra( str(k.dtype), str(dev), ) - _autotuned_launch(chunk_kda_fwd_kernel_inter_diag_compute_solve, _diag_key, _diag_grid, _diag_args, stream=stream) + autotuned_launch(chunk_kda_fwd_kernel_inter_diag_compute_solve, _diag_key, _diag_grid, _diag_args, stream=stream) else: Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( q=q, k=k, gk=gk, beta=beta, Aqk=Aqk, Akk=Akkd, scale=scale, cu_seqlens=cu_seqlens, chunk_size=BT, sub_chunk_size=BC, stream=stream @@ -3441,16 +3255,15 @@ def chunk_kda_fwd_intra( # kernel. The fused kernel handles NC>=3/NC>=4 internally (block-triangular # forward-substitution over all sub-chunk pairs). With use_solved_diagonal # the forward-substitution is skipped. - BKf = _next_power_of_2(K) + BKf = next_power_of_2(K) # Launch-hint autotune (occupancy x num_worker_warps) on this fused - # solve kernel. The - # default config is tried first so non-improving shapes are unchanged. + # solve kernel. _isf_grid = (NT, HV) _isf_args = ( - q.reshape(-1), - k.reshape(-1), + reshaped(q, (-1,)), + reshaped(k, (-1,)), gk.reshape(-1), - beta.reshape(-1), + reshaped(beta, (-1,)), Aqk.reshape(-1), Akkd.reshape(-1), Akk.reshape(-1), @@ -3479,7 +3292,7 @@ def chunk_kda_fwd_intra( str(k.dtype), str(dev), ) - _autotuned_launch(chunk_kda_fwd_kernel_inter_solve_fused, _isf_key, _isf_grid, _isf_args, occ_choices=(1, 2, 4), nww_choices=(4,), stream=stream) + autotuned_launch(chunk_kda_fwd_kernel_inter_solve_fused, _isf_key, _isf_grid, _isf_args, occ_choices=(1, 2, 4), nww_choices=(4,), stream=stream) w, u, qg, kg = recompute_w_u_fwd( k=k, v=v, beta=beta, A=Akk, q=q if disable_recompute else None, gk=gk, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, bufs=bufs, stream=stream ) @@ -3493,25 +3306,22 @@ def chunk_kda_bwd_intra(q, k, g, beta, dAqk, dAkk, dq, dk, db, dg, cu_seqlens=No # Fast path: for BT >= 64 use larger # BC=32 / BK=64 sub-tiles and route through the SAFE_GATE matmul branch # instead of the BC-iteration scalar `for j` loops. The scalar path is the - # bwd bottleneck; the matmul branch is numerically - # equivalent (it is the same code reached when the user sets safe_gate=True). + # bwd bottleneck. use_fast_path = BT >= 64 BC = 32 if use_fast_path else min(16, BT) - BK = min(64 if use_fast_path else 32, _next_power_of_2(K)) + BK = min(64 if use_fast_path else 32, next_power_of_2(K)) safe_gate = safe_gate or use_fast_path - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") NT = len(chunk_indices) - NC = _cdiv(BT, BC) - NK = _cdiv(K, BK) + NC = cdiv(BT, BC) + NK = cdiv(K, BK) dq2 = reshaped(bufs["dq2"], (T, HV, K)) dk2 = reshaped(bufs["dk2"], (T, HV, K)) db2 = reshaped(bufs["db2"], (NK, T, HV)) dg2 = reshaped(bufs["dg2"], (T, HV, K)) - cu_arg = _i32_flat(cu_seqlens) - ci_arg = _i32_flat(chunk_indices) + cu_arg = i32_flat(cu_seqlens) + ci_arg = i32_flat(chunk_indices) grid = (NK * NC, NT, HV) # Kernel uses flat element-offset gather/scatter; pass 1-D views. ct.launch( @@ -3519,10 +3329,10 @@ def chunk_kda_bwd_intra(q, k, g, beta, dAqk, dAkk, dq, dk, db, dg, cu_seqlens=No grid, chunk_kda_bwd_kernel_intra, ( - q.reshape(-1), - k.reshape(-1), + reshaped(q, (-1,)), + reshaped(k, (-1,)), g.reshape(-1), - beta.reshape(-1), + reshaped(beta, (-1,)), dAqk.reshape(-1), dAkk.reshape(-1), dq.reshape(-1), @@ -3547,7 +3357,7 @@ def chunk_kda_bwd_intra(q, k, g, beta, dAqk, dAkk, dq, dk, db, dg, cu_seqlens=No ) dq = dq2 dk = dk2 - # db += sum_nk db2 (fp32 acc); the fan-in NK is a compile-time constant + # dBeta += sum_nk dBeta2 (fp32 acc); the fan-in NK is a compile-time constant sum_leading(reshaped(db, (T * HV,)), reshaped(db2, (NK, T * HV)), NK, T * HV, stream=stream, accumulate=True) dg = dg2 return dq, dk, db, dg @@ -3560,28 +3370,26 @@ def chunk_kda_bwd_dAv(q, k, v, do, A=None, scale=None, cu_seqlens=None, chunk_si stream = 0 if stream is None else stream T, H, K, HV, V = *k.shape, do.shape[1], do.shape[-1] BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") CONST_TILING = 64 - BK = min(max(_next_power_of_2(K), 16), CONST_TILING) - BV = min(max(_next_power_of_2(V), 16), CONST_TILING) + BK = min(max(next_power_of_2(K), 16), CONST_TILING) + BV = min(max(next_power_of_2(V), 16), CONST_TILING) NT = len(chunk_indices) dA = reshaped(bufs["dAqk"], (T, HV, BT)) dv = reshaped(bufs["dv_dAv"], (T, HV, V)) - cu_arg = _i32_flat(cu_seqlens) - ci_arg = _i32_flat(chunk_indices) + cu_arg = i32_flat(cu_seqlens) + ci_arg = i32_flat(chunk_indices) # Kernel uses flat element-offset gather/scatter; pass 1-D views. ct.launch( stream, (NT, HV), chunk_kda_bwd_kernel_dAv, ( - q.reshape(-1), - k.reshape(-1), + reshaped(q, (-1,)), + reshaped(k, (-1,)), v.reshape(-1), A.reshape(-1), - do.reshape(-1), + reshaped(do, (-1,)), dv.reshape(-1), dA.reshape(-1), cu_arg, @@ -3622,12 +3430,10 @@ def chunk_kda_bwd_wy_dqkg_fused( stream = 0 if stream is None else stream T, H, K, HV, V = *k.shape, v.shape[1], v.shape[-1] BT = chunk_size - if cu_seqlens is None or chunk_indices is None: - raise ValueError("cu_seqlens and chunk_indices are required (THD layout; callers build the (seq, intra) table)") NT = len(chunk_indices) CONST_TILING = 64 - BK = min(max(_next_power_of_2(K), 16), CONST_TILING) - BV = min(max(_next_power_of_2(V), 16), CONST_TILING) + BK = min(max(next_power_of_2(K), 16), CONST_TILING) + BV = min(max(next_power_of_2(V), 16), CONST_TILING) dq = reshaped(bufs["dq"], (T, HV, K)) dk = reshaped(bufs["dk"], (T, HV, K)) @@ -3643,15 +3449,15 @@ def chunk_kda_bwd_wy_dqkg_fused( grid, chunk_kda_bwd_kernel_wy_dqkg_fused, ( - q.reshape(-1), - k.reshape(-1), - v.reshape(-1), + reshaped(q, (-1,)), + reshaped(k, (-1,)), + reshaped(v, (-1,)), v_new.reshape(-1), g.reshape(-1), - beta.reshape(-1), + reshaped(beta, (-1,)), A.reshape(-1), h.reshape(-1), - do.reshape(-1), + reshaped(do, (-1,)), dh.reshape(-1), dq.reshape(-1), dk.reshape(-1), @@ -3660,8 +3466,8 @@ def chunk_kda_bwd_wy_dqkg_fused( dg.reshape(-1), db.reshape(-1), dA.reshape(-1), - _i32_flat(cu_seqlens), - _i32_flat(chunk_indices), + i32_flat(cu_seqlens), + i32_flat(chunk_indices), float(scale), H, HV, @@ -3703,12 +3509,12 @@ def chunk_kda_fwd( dt_bias=None, disable_recompute=False, return_intermediate_states=False, + compute_o=True, cp_context=None, bufs=None, stream=None, ): stream = 0 if stream is None else stream - assert cp_context is None, "context-parallel path is not supported in this self-contained port" g_org = None if use_gate_in_kernel: g_org = g @@ -3768,20 +3574,22 @@ def chunk_kda_fwd( stream=stream, ) - o = chunk_gla_fwd_o_gk( - q=q, - v=v_new, - g=g, - A=Aqk, - h=h, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - bufs=bufs, - stream=stream, - ) + o = None + if compute_o: + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + bufs=bufs, + stream=stream, + ) if disable_recompute is False: w, u, qg, kg, v_new = None, None, None, None, None if not return_intermediate_states: @@ -3801,7 +3609,7 @@ def chunk_kda_bwd( scale, initial_state, do, - dht, + dstate_in, g=None, g_org=None, state_v_first=False, @@ -3820,7 +3628,6 @@ def chunk_kda_bwd( **kwargs, ): stream = 0 if stream is None else stream - assert cp_context is None, "context-parallel path is not supported in this self-contained port" H, HV = q.shape[1], v.shape[1] G = HV // H @@ -3867,7 +3674,7 @@ def chunk_kda_bwd( w=w, gk=g, h0=initial_state, - dht=dht, + dstate_in=dstate_in, do=do, dv=dv, scale=scale, @@ -3919,7 +3726,7 @@ def chunk_kda_bwd( stream=stream, ) - # For GVA, reduce dq and dk from [T, HV, K] back to [T, H, K] + # For GVA, reduce dQ and dK from [T, HV, K] back to [T, H, K] if HV > H: T_, K_ = dq.shape[0], dq.shape[-1] dq_r = reshaped(bufs["dq_hred"], (T_, H, K_)) @@ -3959,7 +3766,7 @@ def chunk_kda_grad( g, beta, do, - dht=None, + dstate_in=None, scale=None, initial_state=None, use_qk_l2norm_in_kernel=False, @@ -3983,7 +3790,7 @@ def chunk_kda_grad( ``q``/``k``/``v``/``do`` are ``[total_T, H, D]``, ``g`` is ``[total_T, HV, K]``, ``beta`` is ``[total_T, HV]``; ``cu_seqlens`` and ``chunk_indices`` are required. Recomputes the forward's prep - (L2-normalized q/k + rstd, cumulative gate, WY factors, per-chunk + (L2-normalized Q/K + rstd, cumulative gate, WY factors, per-chunk states) from the inputs, then runs the backward kernels; intermediates live in the ``bufs`` carves when provided. @@ -3992,8 +3799,6 @@ def chunk_kda_grad( present). """ stream = 0 if stream is None else stream - if cu_seqlens is None: - raise ValueError("cu_seqlens is required (THD layout)") if scale is None: scale = q.shape[-1] ** -0.5 q_rstd, k_rstd = None, None @@ -4004,8 +3809,6 @@ def chunk_kda_grad( beta_raw = beta if use_beta_sigmoid_in_kernel: beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0, out=bufs["beta_sig"], stream=stream) - if cu_seqlens is not None and chunk_indices is None: - raise ValueError("varlen (cu_seqlens) requires chunk_indices — callers build the (seq, intra) table") _o, _fs, g_cumsum, Aqk, Akk, w, u, qg, kg, v_new, h, initial_state = chunk_kda_fwd( q=q_in, k=k_in, @@ -4025,6 +3828,8 @@ def chunk_kda_grad( dt_bias=dt_bias, chunk_size=chunk_size, state_v_first=state_v_first, + disable_recompute=True, + compute_o=False, bufs=bufs, stream=stream, ) @@ -4038,7 +3843,7 @@ def chunk_kda_grad( scale=scale, initial_state=initial_state, do=do, - dht=dht, + dstate_in=dstate_in, g=g_cumsum, g_org=g if use_gate_in_kernel else None, state_v_first=state_v_first, @@ -4050,6 +3855,7 @@ def chunk_kda_grad( use_gate_in_kernel=use_gate_in_kernel, A_log=A_log, dt_bias=dt_bias, + disable_recompute=True, bufs=bufs, w=w, u=u, @@ -4065,11 +3871,11 @@ def chunk_kda_grad( if use_beta_sigmoid_in_kernel: db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if allow_neg_eigval else 1.0, out=db, stream=stream) return ( - _cast(bufs, "dq_cast", dq, q), - _cast(bufs, "dk_cast", dk, k), - _cast(bufs, "dv_cast", dv, v), - _cast(bufs, "dg_cast", dg, g), - _cast(bufs, "db_cast", db, beta_raw), + cast(bufs, "dq_cast", dq, q), + cast(bufs, "dk_cast", dk, k), + cast(bufs, "dv_cast", dv, v), + cast(bufs, "dg_cast", dg, g), + cast(bufs, "db_cast", db, beta_raw), dh0, dA, dbias, @@ -4110,48 +3916,19 @@ def chunk_kda( THD layout, written into ``bufs['o']`` / ``bufs['fs']``.""" stream = 0 if stream is None else stream if "transpose_state_layout" in kwargs: - if state_v_first: - raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") state_v_first = kwargs.pop("transpose_state_layout") - assert cp_context is None, "context-parallel path is not supported in this self-contained port" - - if cu_seqlens is None: - raise ValueError("cu_seqlens is required (THD layout)") - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to equal the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if initial_state is not None: - assert str(initial_state.dtype).endswith("float32"), "initial_state must be in float32." - A_log, dt_bias = None, None if use_gate_in_kernel: - assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") chunk_size = kwargs.pop("chunk_size", 64) - if chunk_size not in (32, 64): - raise ValueError(f"`chunk_size` must be either 32 or 64 for KDA, got {chunk_size}.") if safe_gate and use_gate_in_kernel: - if lower_bound is None: - raise ValueError("`lower_bound` must be specified when `safe_gate=True` and `use_gate_in_kernel=True`.") if not (-5 <= lower_bound < 0): raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") - if allow_neg_eigval and not use_beta_sigmoid_in_kernel: - raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") - T, H, K, HV = *q.shape, v.shape[1] - if k.shape[1] != H: - raise ValueError(f"q and k must have the same number of heads, got q.shape[1]={H} and k.shape[1]={k.shape[1]}") - if HV % H != 0: - raise ValueError(f"For GVA, HV ({HV}) must be divisible by H ({H}), got HV % H = {HV % H}") - assert q.shape == k.shape, f"q and k must have the same shape, got q={q.shape} vs k={k.shape}" - assert tuple(g.shape) == (T, HV, K), f"g must have shape [T, HV, K]={[T, HV, K]}, got {list(g.shape)}" - assert tuple(beta.shape) == (T, HV), f"beta must have shape [T, HV]={[T, HV]}, got {list(beta.shape)}" if scale is None: scale = K**-0.5 @@ -4162,8 +3939,6 @@ def chunk_kda( k_in, _k_rstd = l2norm_fwd(k, out=bufs["k_norm"], rstd_out=bufs["k_rstd"], stream=stream) if use_beta_sigmoid_in_kernel: beta = fused_beta_sigmoid(beta, scale=2.0 if allow_neg_eigval else 1.0, out=bufs["beta_sig"], stream=stream) - if cu_seqlens is not None and chunk_indices is None: - raise ValueError("varlen (cu_seqlens) requires chunk_indices — callers build the (seq, intra) table") o, final_state, _gc, _Aqk, _Akk, _w, _u, _qg, _kg, _vn, h, _s0 = chunk_kda_fwd( q=q_in, k=k_in, diff --git a/python/cudnn/linear_attention/engine_utils.py b/python/cudnn/linear_attention/engine_utils.py deleted file mode 100644 index 73be1d59c..000000000 --- a/python/cudnn/linear_attention/engine_utils.py +++ /dev/null @@ -1,134 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Engine-side helpers shared by the linear-attention backends: check_support -dtype gates, the explicit-workspace carver, and the compiled-plan wrapper.""" - -from __future__ import annotations - -from cudnn.engines.base import CompiledPlan, NodeBuffers, bind_ports - -from cudnn.frost import buffers -from cudnn.frost.workspace import Workspace - - -def _dtype_name(dt) -> str: - import cudnn - - return {cudnn.data_type.HALF: "float16", cudnn.data_type.BFLOAT16: "bfloat16", cudnn.data_type.FLOAT: "float32"}[dt] - - -def _require_dtype(engine: str, node, port: str, want, *, out: bool = False) -> None: - import cudnn # noqa: F401 — `want` members come from cudnn.data_type - - t = (node.outputs if out else node.inputs).get(port) - if t is None: - return - got = t.get_data_type() - if got is None: - return # unset (e.g. inferred outputs): the kernel validates the buffer - wanted = want if isinstance(want, tuple) else (want,) - if got not in wanted: - names = "/".join(w.name for w in wanted) - raise NotImplementedError(f"{engine}: '{port}' must be {names} (the kernel-native dtype; no staging), got {got}") - - -def _require_state_pair(engine: str, node) -> None: - """The kernels require matching initial/final state dtypes.""" - s0 = node.inputs.get("initial_state") - fs = node.outputs.get("final_state") - if s0 is None or fs is None: - return - a, b = s0.get_data_type(), fs.get_data_type() - if a is not None and b is not None and a != b: - raise NotImplementedError(f"{engine}: initial_state and final_state dtypes must match (got {a} vs {b})") - - -class _FrostPlan(CompiledPlan): - """A compiled linear-attention kernel, driven from the normalized pack. - - The port-to-slot join is a property of the graph, so it happens once and is - kept; only the addresses change between executes. What the kernel receives - is built from the pack, never the caller's object — the geometry it is - checked against and the geometry it runs on are then the same reading. - """ - - takes_variant_pack = True - - def __init__(self, compiled): - self._compiled = compiled - self._ports = None - self._name = type(compiled).__name__ - - def get_workspace_size(self) -> int: - return self._compiled.workspace_bytes() - - def execute(self, graph, variant_pack, ctx) -> None: - ports = self._ports - if ports is None: - ports = self._ports = bind_ports(graph, variant_pack) - _check_contiguous(variant_pack, ports) - node_buffers = {} - for node, slots in ports.items(): - names = list(slots.inputs) + list(slots.outputs) - views = variant_pack.operands(list(slots.inputs.values()) + list(slots.outputs.values())) - split = len(slots.inputs) - node_buffers[node] = NodeBuffers(dict(zip(names[:split], views[:split])), dict(zip(names[split:], views[split:]))) - required = self._compiled.workspace_bytes() - workspace = Workspace.over(variant_pack, required, self._name) if required else None - self._compiled(node_buffers, workspace=workspace, stream=ctx.stream) - - -def _check_contiguous(variant_pack, ports) -> None: - """Contiguity gate over the whole pack, decided from the strides it holds. - - The dim and stride were taken from the caller's object once, at - normalization; probing each buffer again cost 8.6 us apiece — nine per GDN - forward — to learn what the pack already knows. The scan itself is in the - native pack, 0.24 us for eight operands, so only naming the offender costs - anything and that happens once, on the way to raising. - - One gate for every kernel rather than a call per compiled callable naming - its own ports: the rule was the same list every time, and a port added to a - node but forgotten there would have gone unchecked. - """ - ok, offender = variant_pack.all_contiguous() - if ok: - return - for node, slots in ports.items(): - for direction in (slots.inputs, slots.outputs): - for port, slot in direction.items(): - if slot == offender: - raise ValueError(f"cudnn.frost {node.name!r}: buffer for {port!r} must be contiguous (buffers pass straight to the kernel)") - raise ValueError(f"cudnn.frost: the buffer at variant-pack slot {offender} must be contiguous") - - -_pinned_engines = None # e.g. ("gdn_cutile",) -- set by a suite, None => the manifest decides - - -def pin_engines(names): - """Force planning onto the named engines for this process, or None to stop. - - A suite that means to validate ONE implementation says which, by name. It - is not a way to add an engine: every engine is in the manifest, and this - only narrows which of them a plan may land on. - """ - global _pinned_engines - previous = _pinned_engines - _pinned_engines = tuple(names) if names else None - return previous - - -def apply_pin(graph): - """Select the pinned engine's plan, if a pin is in force. - - Runs after create_execution_plans(), so it selects from the ranked list the - heuristics produced rather than replacing them. - """ - if not _pinned_engines: - return - names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] - index = next((i for i, n in enumerate(names) if any(n == p or n.startswith(p + "[") for p in _pinned_engines)), None) - if index is None: - raise AssertionError(f"pinned engines {list(_pinned_engines)} produced no plan; plans={names}") - graph.select_plan(index) diff --git a/python/cudnn/linear_attention/frost/__init__.py b/python/cudnn/linear_attention/frost/__init__.py index b923f883e..2ff16ba10 100644 --- a/python/cudnn/linear_attention/frost/__init__.py +++ b/python/cudnn/linear_attention/frost/__init__.py @@ -3,14 +3,11 @@ """cudnn.linear_attention.frost: the FROST linear-attention engines — Gated DeltaNet, Kimi Delta Attention, and Gated DeltaNet v2 on the SM100 -chunked kernels built on Cutlass primitives. ``GdnFrostEngine`` is the -default GDN engine on SM100/SM103; ``KdaFrostEngine`` and ``Gdn2FrostEngine`` -are forward-only (their backward kernels are stubs — KDA gradients run on -``KdaCuTileEngine``).""" - -# Lazy: importing one family's engine must not drag its neighbours in. The -# manifest's factories tolerate a missing optional dependency PER ENGINE, and -# eager imports here would have made one bad import cost all three. +chunked kernels built on Cutlass primitives. All three serve forward and +backward on SM100/SM103 and rank ahead of the cuTile fallbacks, except +GDN-2, which does not have a cuTile fallback.""" + +# Lazy: importing one family's engine must not drag its neighbours in. import importlib from typing import Any diff --git a/python/cudnn/linear_attention/frost/common/downcast.py b/python/cudnn/linear_attention/frost/common/downcast.py new file mode 100644 index 000000000..f6b19a9a4 --- /dev/null +++ b/python/cudnn/linear_attention/frost/common/downcast.py @@ -0,0 +1,100 @@ +# 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.""" + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + + +@cute.kernel +def downcast_state_f16_kernel( + mState0: cute.Tensor, + mOut: cute.Tensor, + n_k: cutlass.Int32, + threads_per_row: cutlass.Int32, + rows_per_cta: cutlass.Int32, +) -> None: + """Row-chunk 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, source read through its (dynamic) + strides so padded outer layouts stage zero-copy.""" + 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: + for i in cutlass.range_constexpr(8): + mOut[n_idx, h_idx, k_idx, v0 + i] = mState0[n_idx, h_idx, k_idx, v0 + i].to(mOut.element_type) + + +@cute.jit +def downcast_state_f16( + 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_f16_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. Stride-aware: + the source is read through its own layout, never reshaped or copied + host-side.""" + 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_f16, + 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/head_reduce.py b/python/cudnn/linear_attention/frost/common/head_reduce.py index 3e363cbaf..697d783c2 100644 --- a/python/cudnn/linear_attention/frost/common/head_reduce.py +++ b/python/cudnn/linear_attention/frost/common/head_reduce.py @@ -13,7 +13,7 @@ f16x2/bf16x2 pair, or one fp32 element), gathers it from all ``r`` group heads (coalesced, strided by ``inner_words``), accumulates in fp32, and stores one word back. Serves the f16/bf16 ``[total, HO, D]`` tensor grads -(dq/dk for GVA, dk/dv for GQA) and the fp32 ``[total, HO]`` gate/beta grads. +(dQ/dK for GVA, dK/dV for GQA) and the fp32 ``[total, HO]`` Gate/Beta grads. """ import cutlass @@ -22,16 +22,20 @@ from cutlass.cute.runtime import from_dlpack +from .host import get_dtype from cudnn.frost.tile_dsl.pointwise import f16x2_to_f32, fp32_to_fp16 BLOCK = 256 @cute.kernel -def _head_reduce_kernel( +def head_reduce_kernel( mIn: cute.Tensor, mOut: cute.Tensor, total_words: cutlass.Int64, + out_row_words: cutlass.Int64, + out_head_words: cutlass.Int64, + h_count: cutlass.Constexpr[int], r: cutlass.Constexpr[int], inner_words: cutlass.Constexpr[int], io_dtype: cutlass.Constexpr, @@ -43,13 +47,16 @@ def _head_reduce_kernel( seg = gw // cutlass.Int64(inner_words) w_off = gw - seg * cutlass.Int64(inner_words) base = seg * cutlass.Int64(r * inner_words) + w_off + t_idx = seg // cutlass.Int64(h_count) + h_idx = seg - t_idx * cutlass.Int64(h_count) + out_off = t_idx * out_row_words + h_idx * out_head_words + w_off if cutlass.const_expr(io_dtype == cutlass.Float32): in_p = cute.recast_ptr(mIn.iterator, dtype=cutlass.Float32) out_p = cute.recast_ptr(mOut.iterator, dtype=cutlass.Float32) acc = (in_p + base).load() for i in cutlass.range_constexpr(r - 1): acc = acc + (in_p + (base + (i + 1) * inner_words)).load() - (out_p + gw).store(acc) + (out_p + out_off).store(acc) else: in_p = cute.recast_ptr(mIn.iterator, dtype=cutlass.Int32) out_p = cute.recast_ptr(mOut.iterator, dtype=cutlass.Int32) @@ -58,39 +65,31 @@ def _head_reduce_kernel( lo, hi = f16x2_to_f32((in_p + (base + (i + 1) * inner_words)).load(), dtype=io_dtype) acc_lo = acc_lo + lo acc_hi = acc_hi + hi - (out_p + gw).store(fp32_to_fp16(acc_lo, acc_hi, dtype=io_dtype)) + (out_p + out_off).store(fp32_to_fp16(acc_lo, acc_hi, dtype=io_dtype)) @cute.jit -def _launch( +def launch( mIn: cute.Tensor, mOut: cute.Tensor, total_words: cutlass.Int64, + out_row_words: cutlass.Int64, + out_head_words: cutlass.Int64, grid_x: cutlass.Int32, + h_count: cutlass.Constexpr[int], r: cutlass.Constexpr[int], inner_words: cutlass.Constexpr[int], io_dtype: cutlass.Constexpr, stream: cuda.CUstream, ) -> None: - _head_reduce_kernel(mIn, mOut, total_words, r, inner_words, io_dtype).launch( + head_reduce_kernel(mIn, mOut, total_words, out_row_words, out_head_words, h_count, r, inner_words, io_dtype).launch( grid=(grid_x, 1, 1), block=(BLOCK, 1, 1), stream=stream, ) -_compiled_cache = {} - - -def _cutlass_io_dtype(dtype) -> type: - name = str(dtype).split(".")[-1] - if name == "bfloat16": - return cutlass.BFloat16 - if name == "float16": - return cutlass.Float16 - if name == "float32": - return cutlass.Float32 - raise ValueError(f"head_group_reduce: unsupported dtype {dtype} (float16/bfloat16/float32 only)") +compiled_cache = {} def head_group_reduce(src, dst, *, stream) -> None: @@ -98,7 +97,9 @@ def head_group_reduce(src, dst, *, stream) -> None: rank-2 ``(total, HO)`` into ``(total, H)`` — by summing each group of ``r = HO // H`` consecutive heads (fp32 accumulation). - Both tensors are contiguous, same-dtype (f16/bf16, or fp32), + ``src`` is contiguous (kernel-internal wide buffer); ``dst`` needs a + stride-1 innermost dim with free outer strides (f16/bf16 outer strides + must be even — word-pair stores). Same-dtype (f16/bf16, or fp32), DLPack-compatible CUDA tensors; the f16/bf16 inner extent ``D`` must be even. Compile-cache-and-replay per ``(dtype, HO, H, D)``.""" if len(src.shape) == 2: @@ -114,7 +115,7 @@ def head_group_reduce(src, dst, *, stream) -> None: raise ValueError(f"head_group_reduce: shape mismatch {tuple(src.shape)} -> {tuple(dst.shape)}") if H <= 0 or HO % H != 0 or HO <= H: raise ValueError(f"head_group_reduce: bad head group HO={HO} H={H}") - io_dtype = _cutlass_io_dtype(src.dtype) + io_dtype = get_dtype(src.dtype) if str(src.dtype).split(".")[-1] != str(dst.dtype).split(".")[-1]: raise ValueError(f"head_group_reduce: dtype mismatch {src.dtype} vs {dst.dtype}") is_fp32 = io_dtype == cutlass.Float32 @@ -126,24 +127,32 @@ def head_group_reduce(src, dst, *, stream) -> None: grid_x = -(-total_words // BLOCK) cu_stream = cuda.CUstream(int(stream)) - key = (str(src.dtype).split(".")[-1], HO, H, D) - if key not in _compiled_cache: + dst_strides = tuple(dst.stride()) + if not is_fp32 and any(st % 2 != 0 for st, sz in zip(dst_strides[:-1], dst.shape[:-1]) if sz != 1): + raise ValueError(f"head_group_reduce: f16/bf16 dst outer strides must be even (word-pair stores), got {dst_strides}") + if dst.shape[-1] != 1 and dst_strides[-1] != 1: + raise ValueError(f"head_group_reduce: dst innermost dim must be stride-1, got strides {dst_strides}") + out_row_words = dst_strides[0] if is_fp32 else dst_strides[0] // 2 + out_head_words = (dst_strides[1] if is_fp32 else dst_strides[1] // 2) if len(dst.shape) == 3 else 1 - def _tok(t): - c = from_dlpack(t, assumed_align=4) - c.mark_compact_shape_dynamic(mode=0, stride_order=tuple(range(len(t.shape))), divisibility=1) - return c - - _compiled_cache[key] = cute.compile( - _launch, - _tok(src), - _tok(dst), + key = (str(src.dtype).split(".")[-1], HO, H, D) + if key not in compiled_cache: + + src_c = from_dlpack(src, assumed_align=4) + src_c.mark_compact_shape_dynamic(mode=0, stride_order=tuple(range(len(src.shape))), divisibility=1) + compiled_cache[key] = cute.compile( + launch, + src_c, + from_dlpack(dst, assumed_align=4).mark_layout_dynamic(leading_dim=len(dst.shape) - 1), cutlass.Int64(total_words), + cutlass.Int64(out_row_words), + cutlass.Int64(out_head_words), cutlass.Int32(grid_x), + H, r, inner_words, io_dtype, cu_stream, options="--enable-tvm-ffi", ) - _compiled_cache[key](src, dst, total_words, grid_x, cu_stream) + compiled_cache[key](src, dst, total_words, out_row_words, out_head_words, grid_x, cu_stream) diff --git a/python/cudnn/linear_attention/frost/common/host.py b/python/cudnn/linear_attention/frost/common/host.py new file mode 100644 index 000000000..7235ebdec --- /dev/null +++ b/python/cudnn/linear_attention/frost/common/host.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Host-side helpers shared by the FROST LA kernel modules (engine-invoked).""" + +import cutlass + +from .thd import TENSOR_MAP_QWORDS + + +def get_dtype(dtype): + """dtype string -> cutlass DSL type (bf16/fp16 io, fp32/bf16 states).""" + name = str(dtype) + if "bfloat16" in name: + return cutlass.BFloat16 + if "float16" in name or "half" in name: + return cutlass.Float16 + if "float32" in name: + return cutlass.Float32 + raise ValueError(f"Unsupported dtype {dtype}, expected bfloat16, float16, or float32") + + +def tensormap_workspace_bytes(mod, B: int) -> int: + """Runtime TMA-descriptor block for a kernel module: per-batch arrays + + static slots + 128 alignment slack.""" + return TENSOR_MAP_QWORDS * 8 * (mod.TENSORMAP_DESC_ARRAYS * B + mod.TENSORMAP_STATIC_SLOTS) + 128 diff --git a/python/cudnn/linear_attention/frost/common/split_k.py b/python/cudnn/linear_attention/frost/common/split_k.py index 3d4146bed..878e20eef 100644 --- a/python/cudnn/linear_attention/frost/common/split_k.py +++ b/python/cudnn/linear_attention/frost/common/split_k.py @@ -16,15 +16,19 @@ Work-item table row (``WORK_ITEM_FIELDS`` x int32, chunk units):: - [batch_idx, head_idx, wstart, wend, cstart, cend] + [batch_idx, head_idx, wstart, wend, cstart, cend, batch_start, batch_end] + +``batch_start``/``batch_end`` are the token bounds ``cu_seqlens[b]`` / +``cu_seqlens[b+1]`` denormalized into the row: decode reads one 32-byte +vectorizable row instead of chasing a dependent ``cu_seqlens`` load pair. The item OWNS (writes outputs for) chunks ``[wstart, wend)``. The forward kernel COMPUTES chunks ``[cstart, wend)`` — ``[cstart, wstart)`` is the left warmup that rebuilds the incoming state from zero (accurate to ``2^log2_threshold`` because the gate decay over the window saturates). The backward kernel computes ``[wstart, cend)`` — ``[wend, cend)`` is the -right warmup for the reverse dH recurrence (the forward states come exactly -from the per-chunk H checkpoints). ``cstart == 0`` items seed the true +right warmup for the reverse dstate recurrence (the forward states come exactly +from the per-chunk state checkpoints). ``cstart == 0`` items seed the true initial state; ``cend == num_chunks`` items seed the true ``d_final_state`` — so the un-cut degenerate item ``(0, nc, 0, nc)`` reproduces the serial kernel exactly. @@ -56,6 +60,14 @@ consumes them in LPT order — the makespan tail is set by whatever starts last, so the big items must go first. This is what keeps ragged varlen batches balanced without cutting them. + +The order kernel also zeroes the main kernels' scheduler ticket rings +(dirty on exit), and with ``split=False`` it replaces the whole pipeline: +scan and walk never launch, and the order kernel synthesizes the uncut +whole-sequence item per (batch, head) from ``cu_seqlens`` alone, then +LPT-sorts those. That no-cuts table serves batch-invariant mode and +coarse checkpoint cadences (cuts may not cross a checkpoint period), so +ragged batches keep LPT scheduling at the cost of one single-CTA launch. """ import math @@ -68,7 +80,9 @@ from cutlass.cute.arch.nvvm_wrappers import inline_ptx from cutlass.cute.runtime import from_dlpack -WORK_ITEM_FIELDS = 6 +from cudnn.frost.buffers import data_ptr + +WORK_ITEM_FIELDS = 8 WARMUP_CAP_CHUNKS = 32 # hard warmup cap: a cut must saturate within one warp of chunks per side MAX_BLOCKS = 2048 # piece-count ceiling; host clamps ideal_chunks so the per-tile block count fits WARP_SIZE = 32 @@ -76,18 +90,18 @@ THREADS_PER_BLOCK = WARPS * WARP_SIZE SCAN_WARPS = 4 SCAN_THREADS = SCAN_WARPS * WARP_SIZE -SCAN_ROWS_PER_WARP = 4 # consecutive chunk rows per scan warp (amortizes the batch lookup + piece choice) -SCAN_TOKEN_STRIDE = 4 # sample every Nth token of a chunk: skipped tokens only RAISE the negative horizon sums (sound), and the safe-gate sigmoid is SFU-bound +SCAN_ROWS_PER_WARP = 4 # consecutive chunk rows per scan warp +SCAN_TOKEN_STRIDE = 4 # sample every Nth token of a chunk: skipped tokens only RAISE the negative horizon sums ORDER_THREADS = 1024 ORDER_ELEMS = 4 ORDER_CAPACITY = ( ORDER_THREADS * ORDER_ELEMS -) # sort capacity (32 KB SMEM); the kernel always launches — past this the device-side branch copies through unsorted (>25 items/SM, LPT is noise there) +) # sort capacity (32 KB SMEM); the kernel always launches — past this the device-side branch copies through unsorted OVERHEAD_TOKENS = 256 # per-item fixed cost for the piece model: state reseed + pipeline refill + typical warmup P_WINDOW = 16 # fill-regime piece-count search width P_BELOW = 8 # how far below the ideal-cap floor the fill-regime search may go -_DEFAULT_LOG2_THRESHOLD = -10.0 / math.log(2.0) # e^-10, in log2 units +DEFAULT_LOG2_THRESHOLD = -10.0 / math.log(2.0) # e^-10, in log2 units RCP_LN2 = 1.4426950408889634 # 1/ln(2): natural-log gates -> the scan's log2 domain @@ -119,40 +133,26 @@ def max_work_items(total_tokens: int, batch_size: int, n_heads_out: int, ideal_c @cute.jit -def decode_work_item(cfg, tile_idx, cu_seqlens, mWorkItems): - """Tile decode shared by every warp body of the main kernels. - - Legacy mode maps ``tile_idx`` to a full (batch, head) sequence; split-K - mode reads the work-item row. Returns +def decode_work_item(cfg, tile_idx, mWorkItems): + """Tile decode shared by every warp body of the main kernels: read the + work-item row (an uncut table row IS the whole sequence). Returns ``(batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend)`` with chunk-unit bounds.""" - if cutlass.const_expr(cfg.split_k): - batch_idx = mWorkItems[tile_idx, 0] - head_idx = mWorkItems[tile_idx, 1] - wstart = mWorkItems[tile_idx, 2] - wend = mWorkItems[tile_idx, 3] - cstart = mWorkItems[tile_idx, 4] - cend = mWorkItems[tile_idx, 5] - batch_start = cu_seqlens[batch_idx] - batch_end = cu_seqlens[batch_idx + 1] - seqlen_b = batch_end - batch_start - num_chunks_b = cute.ceil_div(seqlen_b, cfg.b_t) - else: - batch_idx = tile_idx // cfg.n_heads_out - head_idx = tile_idx % cfg.n_heads_out - batch_start = cu_seqlens[batch_idx] - batch_end = cu_seqlens[batch_idx + 1] - seqlen_b = batch_end - batch_start - num_chunks_b = cute.ceil_div(seqlen_b, cfg.b_t) - wstart = cutlass.Int32(0) - wend = num_chunks_b - cstart = cutlass.Int32(0) - cend = num_chunks_b + batch_idx = mWorkItems[tile_idx, 0] + head_idx = mWorkItems[tile_idx, 1] + wstart = mWorkItems[tile_idx, 2] + wend = mWorkItems[tile_idx, 3] + cstart = mWorkItems[tile_idx, 4] + cend = mWorkItems[tile_idx, 5] + batch_start = mWorkItems[tile_idx, 6] + batch_end = mWorkItems[tile_idx, 7] + seqlen_b = batch_end - batch_start + num_chunks_b = cute.ceil_div(seqlen_b, cfg.b_t) return batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend @cute.jit -def _emit_item(mWorkItems, mCount, batch_idx, head_idx, wstart, wend, cstart, cend): +def emit_item(mWorkItems, mCount, batch_idx, head_idx, wstart, wend, cstart, cend, batch_start, batch_end): count_addr = mCount.iterator.toint() wi = inline_ptx( "atom.global.add.s32 {$w0}, [{$r0}], 1;", @@ -165,10 +165,12 @@ def _emit_item(mWorkItems, mCount, batch_idx, head_idx, wstart, wend, cstart, ce mWorkItems[wi, 3] = wend mWorkItems[wi, 4] = cstart mWorkItems[wi, 5] = cend + mWorkItems[wi, 6] = batch_start + mWorkItems[wi, 7] = batch_end @cute.jit -def _clamped_log2(log_gate: cutlass.Constexpr[bool], gate_val: cutlass.Float32) -> cutlass.Float32: +def clamped_log2(log_gate: cutlass.Constexpr[bool], gate_val: cutlass.Float32) -> cutlass.Float32: """Log2-domain decay increment, clamped to <= 0 (a gate > 1 must not relax the horizon).""" if cutlass.const_expr(log_gate): @@ -179,7 +181,7 @@ def _clamped_log2(log_gate: cutlass.Constexpr[bool], gate_val: cutlass.Float32) @cute.jit -def _piece_choice(overhead_chunks: cutlass.Constexpr[int], num_chunks_b, n_tiles, num_sms, ideal_chunks): +def piece_choice(overhead_chunks: cutlass.Constexpr[int], num_chunks_b, n_tiles, num_sms, ideal_chunks): """Per-tile piece choice. Returns ``(span, num_blocks)``.""" # even spread: spans never exceed ideal_chunks (total work / SM count) p_hi = num_chunks_b if num_chunks_b < cutlass.Int32(MAX_BLOCKS) else cutlass.Int32(MAX_BLOCKS) @@ -189,8 +191,7 @@ def _piece_choice(overhead_chunks: cutlass.Constexpr[int], num_chunks_b, n_tiles p = p if p < p_hi else p_hi if n_tiles < cutlass.Int32(2) * num_sms: # fill regime: SMs to spare — search piece counts around the cap on - # the wave-quantized makespan estimate (the cap only binds in the - # many-tile regime, where its even spread kills varlen tails) + # the wave-quantized makespan estimate p_start = p - cutlass.Int32(P_BELOW) p_start = p_start if p_start > 0 else cutlass.Int32(1) best = cutlass.Int32(2147483647) @@ -203,8 +204,7 @@ def _piece_choice(overhead_chunks: cutlass.Constexpr[int], num_chunks_b, n_tiles hit = est < best best = est if hit else best p = cand if hit else p - # the estimate flatters marginal cuts (measured: <25% predicted gain - # loses to serial); cut only on a clear margin over uncut + # the estimate flatters marginal cuts; cut only on a clear margin over uncut est1 = ((n_tiles + num_sms - cutlass.Int32(1)) // num_sms) * (num_chunks_b + cutlass.Int32(overhead_chunks)) if cutlass.Int32(4) * best > cutlass.Int32(3) * est1: p = cutlass.Int32(1) @@ -217,23 +217,23 @@ def _piece_choice(overhead_chunks: cutlass.Constexpr[int], num_chunks_b, n_tiles @cute.jit -def _tile_spans(b_t: cutlass.Constexpr[int], overhead_chunks: cutlass.Constexpr[int], n_heads_out, n_tiles, num_sms, ideal_chunks, mCuSeqlens, tile): +def tile_spans(b_t: cutlass.Constexpr[int], overhead_chunks: cutlass.Constexpr[int], n_heads_out, n_tiles, num_sms, ideal_chunks, mCuSeqlens, tile): """Per-tile decode + piece choice. Returns ``(batch_idx, head_idx, batch_start, batch_end, num_chunks_b, cv_base, span, num_blocks)``; ``cv_base`` is the tile's row base in the GMEM chunk scratch.""" batch_idx = tile // n_heads_out head_idx = tile % n_heads_out - batch_start = mCuSeqlens[batch_idx] - batch_end = mCuSeqlens[batch_idx + 1] + batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) + batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) seqlen_b = batch_end - batch_start num_chunks_b = cute.ceil_div(seqlen_b, b_t) cv_base = batch_start // cutlass.Int32(b_t) + batch_idx - span, num_blocks = _piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) + span, num_blocks = piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) return batch_idx, head_idx, batch_start, batch_end, num_chunks_b, cv_base, span, num_blocks @cute.jit -def _near_boundary(c, span, num_blocks): +def near_boundary(c, span, num_blocks): """True iff chunk ``c`` lies in the walk's read window of a candidate boundary: suffix ``[j*span - W, j*span)`` or prefix ``[j*span, j*span + W)`` for some ``j`` in ``[1, num_blocks)``.""" @@ -246,13 +246,12 @@ def _near_boundary(c, span, num_blocks): @cute.kernel -def _scan_kernel( +def scan_kernel( b_t: cutlass.Constexpr[int], log_gate: cutlass.Constexpr[bool], safe_gate: cutlass.Constexpr[bool], gate_channels: cutlass.Constexpr[int], overhead_chunks: cutlass.Constexpr[int], - has_sched: cutlass.Constexpr[bool], n_heads_out: cutlass.Int32, n_tiles: cutlass.Int32, num_sms: cutlass.Int32, @@ -265,26 +264,19 @@ def _scan_kernel( mCuSeqlens: cute.Tensor, mChunkVals: cute.Tensor, mCount: cute.Tensor, - mSched: cute.Tensor | None, ): """Flat windowed chunk scan, per-channel gate (KDA / GDN-2): CTA ``(x, h)`` covers 16 chunk-scratch rows of head ``h``, one warp per chunk, lane ``l`` owning channels ``[l*cpl, (l+1)*cpl)``. Chunks outside every cut window — and whole tiles the piece choice leaves uncut — never touch the gate. CTA (0, 0) also zeroes the item count - for the walk and the main kernels' scheduler ticket ring (the kernels - leave the ring dirty on exit, so every launch needs a fresh build).""" + for the walk (the scheduler rings are the order kernel's job).""" tidx, _, _ = cute.arch.thread_idx() bidx = cute.arch.block_idx() tidx = cutlass.Int32(tidx) head_idx = cutlass.Int32(bidx[1]) if cutlass.Int32(bidx[0]) == 0 and head_idx == 0 and tidx == 0: mCount[0] = cutlass.Int32(0) - if cutlass.const_expr(has_sched): - i = cutlass.Int32(0) - while i < mSched.shape[0]: - mSched[i] = cutlass.Int32(0) - i = i + cutlass.Int32(1) lidx = tidx % cutlass.Int32(WARP_SIZE) widx = tidx // cutlass.Int32(WARP_SIZE) row0 = (cutlass.Int32(bidx[0]) * cutlass.Int32(SCAN_WARPS) + widx) * cutlass.Int32(SCAN_ROWS_PER_WARP) @@ -294,35 +286,35 @@ def _scan_kernel( hi = batch_size - cutlass.Int32(1) while lo < hi: mid = (lo + hi + cutlass.Int32(1)) // cutlass.Int32(2) - take = mCuSeqlens[mid] // cutlass.Int32(b_t) + mid <= row0 + take = cutlass.Int32(mCuSeqlens[mid]) // cutlass.Int32(b_t) + mid <= row0 lo = mid if take else lo hi = hi if take else mid - cutlass.Int32(1) batch_idx = lo - batch_start = mCuSeqlens[batch_idx] - batch_end = mCuSeqlens[batch_idx + 1] + batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) + batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) num_chunks_b = cute.ceil_div(batch_end - batch_start, b_t) cv_base = batch_start // cutlass.Int32(b_t) + batch_idx - # the piece choice is per batch: computed once here and again only when - # a row crosses into the next batch — NOT per row (on uncut tiles this - # is the scan's entire cost) - span, num_blocks = _piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) + # the piece choice is computed per batch, NOT per row + span, num_blocks = piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) for rr in cutlass.range_constexpr(SCAN_ROWS_PER_WARP): row = row0 + cutlass.Int32(rr) - while (batch_idx + cutlass.Int32(1) < batch_size) and (mCuSeqlens[batch_idx + 1] // cutlass.Int32(b_t) + batch_idx + cutlass.Int32(1) <= row): + while (batch_idx + cutlass.Int32(1) < batch_size) and ( + cutlass.Int32(mCuSeqlens[batch_idx + 1]) // cutlass.Int32(b_t) + batch_idx + cutlass.Int32(1) <= row + ): batch_idx = batch_idx + cutlass.Int32(1) - batch_start = mCuSeqlens[batch_idx] - batch_end = mCuSeqlens[batch_idx + 1] + batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) + batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) num_chunks_b = cute.ceil_div(batch_end - batch_start, b_t) cv_base = batch_start // cutlass.Int32(b_t) + batch_idx - span, num_blocks = _piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) + span, num_blocks = piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) c = row - cv_base if (c >= 0) and (c < num_chunks_b) and (num_blocks > 1): - if _near_boundary(c, span if span > 0 else cutlass.Int32(1), num_blocks): + if near_boundary(c, span if span > 0 else cutlass.Int32(1), num_blocks): # chunk value = max over channels of the per-channel # clamped-log2 sums (each lane owns a contiguous channel run) cpl = gate_channels // WARP_SIZE - row_elems = n_heads_out * cutlass.Int32(gate_channels) - lane_base = cutlass.Int64(head_idx * cutlass.Int32(gate_channels) + lidx * cutlass.Int32(cpl)) + row_elems = cutlass.Int32(mGate.stride[0]) + lane_base = cutlass.Int64(head_idx * cutlass.Int32(mGate.stride[1]) + lidx * cutlass.Int32(cpl)) gate_addr = mGate.iterator.toint() + lane_base * cutlass.Int64(4) gate_ptr = mGate.iterator + lane_base a_exp = cutlass.Float32(1.0) @@ -359,7 +351,7 @@ def _scan_kernel( contrib = contrib if inb else cutlass.Float32(0.0) else: gvq = gvq if inb else oob - contrib = _clamped_log2(log_gate, gvq) + contrib = clamped_log2(log_gate, gvq) ch_acc[q] = ch_acc[q] + contrib else: for q in cutlass.range_constexpr(cpl): @@ -371,7 +363,7 @@ def _scan_kernel( contrib = contrib if inb else cutlass.Float32(0.0) else: gv = gv if inb else oob - contrib = _clamped_log2(log_gate, gv) + contrib = clamped_log2(log_gate, gv) ch_acc[q] = ch_acc[q] + contrib m = ch_acc[0] for q in cutlass.range_constexpr(1, cpl): @@ -384,11 +376,10 @@ def _scan_kernel( @cute.kernel -def _scan_scalar_kernel( +def scan_scalar_kernel( b_t: cutlass.Constexpr[int], log_gate: cutlass.Constexpr[bool], overhead_chunks: cutlass.Constexpr[int], - has_sched: cutlass.Constexpr[bool], n_heads_out: cutlass.Int32, n_tiles: cutlass.Int32, num_sms: cutlass.Int32, @@ -398,23 +389,17 @@ def _scan_scalar_kernel( mCuSeqlens: cute.Tensor, mChunkVals: cute.Tensor, mCount: cute.Tensor, - mSched: cute.Tensor | None, ): """Scalar-gate scan (GDN): CTA ``(x, hg)`` covers 16 chunk-scratch rows for heads ``[hg*32, (hg+1)*32)``; lane ``l`` owns head ``hg*32 + l``, so gate reads and chunk-value writes are coalesced across lanes and every lane accumulates its own head — no reduction. CTA (0, 0) zeroes the - item count and the scheduler ring, as in the per-channel scan.""" + item count (the scheduler rings are the order kernel's job).""" tidx, _, _ = cute.arch.thread_idx() bidx = cute.arch.block_idx() tidx = cutlass.Int32(tidx) if cutlass.Int32(bidx[0]) == 0 and cutlass.Int32(bidx[1]) == 0 and tidx == 0: mCount[0] = cutlass.Int32(0) - if cutlass.const_expr(has_sched): - i = cutlass.Int32(0) - while i < mSched.shape[0]: - mSched[i] = cutlass.Int32(0) - i = i + cutlass.Int32(1) lidx = tidx % cutlass.Int32(WARP_SIZE) widx = tidx // cutlass.Int32(WARP_SIZE) h = cutlass.Int32(bidx[1]) * cutlass.Int32(WARP_SIZE) + lidx @@ -427,42 +412,44 @@ def _scan_scalar_kernel( hi = batch_size - cutlass.Int32(1) while lo < hi: mid = (lo + hi + cutlass.Int32(1)) // cutlass.Int32(2) - take = mCuSeqlens[mid] // cutlass.Int32(b_t) + mid <= row0 + take = cutlass.Int32(mCuSeqlens[mid]) // cutlass.Int32(b_t) + mid <= row0 lo = mid if take else lo hi = hi if take else mid - cutlass.Int32(1) batch_idx = lo - batch_start = mCuSeqlens[batch_idx] - batch_end = mCuSeqlens[batch_idx + 1] + batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) + batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) num_chunks_b = cute.ceil_div(batch_end - batch_start, b_t) cv_base = batch_start // cutlass.Int32(b_t) + batch_idx - span, num_blocks = _piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) + span, num_blocks = piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) for rr in cutlass.range_constexpr(SCAN_ROWS_PER_WARP): row = row0 + cutlass.Int32(rr) - while (batch_idx + cutlass.Int32(1) < batch_size) and (mCuSeqlens[batch_idx + 1] // cutlass.Int32(b_t) + batch_idx + cutlass.Int32(1) <= row): + while (batch_idx + cutlass.Int32(1) < batch_size) and ( + cutlass.Int32(mCuSeqlens[batch_idx + 1]) // cutlass.Int32(b_t) + batch_idx + cutlass.Int32(1) <= row + ): batch_idx = batch_idx + cutlass.Int32(1) - batch_start = mCuSeqlens[batch_idx] - batch_end = mCuSeqlens[batch_idx + 1] + batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) + batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) num_chunks_b = cute.ceil_div(batch_end - batch_start, b_t) cv_base = batch_start // cutlass.Int32(b_t) + batch_idx - span, num_blocks = _piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) + span, num_blocks = piece_choice(overhead_chunks, num_chunks_b, n_tiles, num_sms, ideal_chunks) c = row - cv_base if (c >= 0) and (c < num_chunks_b) and (num_blocks > 1): - if _near_boundary(c, span if span > 0 else cutlass.Int32(1), num_blocks): + if near_boundary(c, span if span > 0 else cutlass.Int32(1), num_blocks): oob = cutlass.Float32(0.0) if cutlass.const_expr(log_gate) else cutlass.Float32(1.0) acc = cutlass.Float32(0.0) for tt in cutlass.range(0, b_t, SCAN_TOKEN_STRIDE, unroll_full=True): pos = batch_start + c * cutlass.Int32(b_t) + cutlass.Int32(tt) inb = pos < batch_end pos_r = pos if inb else batch_start - gv = (mGate.iterator + cutlass.Int64(pos_r) * cutlass.Int64(n_heads_out) + h_r).load() + gv = (mGate.iterator + cutlass.Int64(pos_r) * cutlass.Int64(mGate.stride[0]) + h_r).load() gv = gv if inb else oob - acc = acc + _clamped_log2(log_gate, gv) + acc = acc + clamped_log2(log_gate, gv) if h_ok: mChunkVals[cv_base + c, h] = acc @cute.kernel -def _walk_kernel( +def walk_kernel( b_t: cutlass.Constexpr[int], overhead_chunks: cutlass.Constexpr[int], n_heads_out: cutlass.Int32, @@ -483,13 +470,13 @@ def _walk_kernel( bidx = cute.arch.block_idx()[0] tidx = cutlass.Int32(tidx) - batch_idx, head_idx, batch_start, batch_end, num_chunks_b, cv_base, span, num_blocks = _tile_spans( + batch_idx, head_idx, batch_start, batch_end, num_chunks_b, cv_base, span, num_blocks = tile_spans( b_t, overhead_chunks, n_heads_out, n_tiles, num_sms, ideal_chunks, mCuSeqlens, cutlass.Int32(bidx) ) if num_blocks <= 1: # single piece: no cuts, nothing scanned if tidx == 0: - _emit_item(mStaging, mCount, batch_idx, head_idx, cutlass.Int32(0), num_chunks_b, cutlass.Int32(0), num_chunks_b) + emit_item(mStaging, mCount, batch_idx, head_idx, cutlass.Int32(0), num_chunks_b, cutlass.Int32(0), num_chunks_b, batch_start, batch_end) else: # packed per-boundary probe results: warm_b | warm_f << 8, 0 = no cut @@ -545,31 +532,91 @@ def _walk_kernel( warm_f = r // cutlass.Int32(256) cend = wend + warm_f cend = cend if cend < num_chunks_b else num_chunks_b - _emit_item(mStaging, mCount, batch_idx, head_idx, prev_cut, wend, cur_cstart, cend) + emit_item(mStaging, mCount, batch_idx, head_idx, prev_cut, wend, cur_cstart, cend, batch_start, batch_end) cur_cstart = wend - warm_b prev_cut = wend jj = jj + cutlass.Int32(1) - _emit_item(mStaging, mCount, batch_idx, head_idx, prev_cut, num_chunks_b, cur_cstart, num_chunks_b) + emit_item(mStaging, mCount, batch_idx, head_idx, prev_cut, num_chunks_b, cur_cstart, num_chunks_b, batch_start, batch_end) + + +@cute.jit +def gen_item_bounds(b_t: cutlass.Constexpr[int], n_heads_out, mCuSeqlens, item): + """(batch, head, batch_start, batch_end, num_chunks) of the uncut + whole-sequence item ``item`` — a no-cuts table row is pure geometry.""" + batch_idx = item // n_heads_out + head_idx = item % n_heads_out + batch_start = cutlass.Int32(mCuSeqlens[batch_idx]) + batch_end = cutlass.Int32(mCuSeqlens[batch_idx + 1]) + num_chunks_b = cute.ceil_div(batch_end - batch_start, b_t) + return batch_idx, head_idx, batch_start, batch_end, num_chunks_b + + +@cute.jit +def write_item( + gen: cutlass.Constexpr[bool], + b_t: cutlass.Constexpr[int], + n_heads_out, + mCuSeqlens, + mStaging, + mWorkItems, + dst, + src, +): + """Final-table row ``dst`` from source item ``src``: the walk's staged + row, or (``gen``) the synthesized uncut item ``(0, nc, 0, nc)``.""" + if cutlass.const_expr(gen): + batch_idx, head_idx, batch_start, batch_end, num_chunks_b = gen_item_bounds(b_t, n_heads_out, mCuSeqlens, src) + mWorkItems[dst, 0] = batch_idx + mWorkItems[dst, 1] = head_idx + mWorkItems[dst, 2] = cutlass.Int32(0) + mWorkItems[dst, 3] = num_chunks_b + mWorkItems[dst, 4] = cutlass.Int32(0) + mWorkItems[dst, 5] = num_chunks_b + mWorkItems[dst, 6] = batch_start + mWorkItems[dst, 7] = batch_end + else: + for f in cutlass.range_constexpr(WORK_ITEM_FIELDS): + mWorkItems[dst, cutlass.Int32(f)] = mStaging[src, cutlass.Int32(f)] @cute.kernel -def _order_kernel( - mStaging: cute.Tensor, +def order_kernel( + gen: cutlass.Constexpr[bool], + has_sched: cutlass.Constexpr[bool], + b_t: cutlass.Constexpr[int], + n_heads_out: cutlass.Int32, + n_tiles: cutlass.Int32, + mCuSeqlens: cute.Tensor, + mStaging: cute.Tensor | None, mCount: cute.Tensor, mWorkItems: cute.Tensor, + mSched: cute.Tensor | None, ): - """LPT ordering (single CTA): bitonic-sort the staged items by span - ``cend - cstart``, longest first, and gather into the final table, so - the ticket scheduler starts the big items before the filler.""" + """LPT ordering (single CTA): bitonic-sort the items by span ``cend - + cstart``, longest first, into the final table, so the ticket scheduler + starts the big items before the filler. Sorts the walk's staged items, + or with ``gen`` synthesizes the uncut whole-sequence item per (batch, + head) from ``cu_seqlens`` directly — the no-cuts table. Thread 0 also + zeroes every ``sched_ctr`` cell (the main kernels' ticket rings, dirty + on exit): this kernel runs on every table build, split or not.""" tidx, _, _ = cute.arch.thread_idx() tidx = cutlass.Int32(tidx) - n = mCount[0] + if cutlass.const_expr(has_sched): + if tidx == 0: + si = cutlass.Int32(0) + while si < mSched.shape[0]: + mSched[si] = cutlass.Int32(0) + si = si + cutlass.Int32(1) + if cutlass.const_expr(gen): + n = n_tiles + if tidx == 0: + mCount[0] = n_tiles + else: + n = mCount[0] if n > cutlass.Int32(ORDER_CAPACITY): - # dozens of items per SM: LPT stops mattering, copy through i = tidx while i < n: - for f in cutlass.range_constexpr(WORK_ITEM_FIELDS): - mWorkItems[i, cutlass.Int32(f)] = mStaging[i, cutlass.Int32(f)] + write_item(gen, b_t, n_heads_out, mCuSeqlens, mStaging, mWorkItems, i, i) i = i + cutlass.Int32(ORDER_THREADS) else: sKey = cutlass.Array(cutlass.Int32, ORDER_CAPACITY, space=cutlass.AddressSpace.smem, alignment=16) @@ -587,7 +634,11 @@ def _order_kernel( for e in cutlass.range_constexpr(ORDER_ELEMS): i = tidx + cutlass.Int32(e * ORDER_THREADS) if i < n: - key = mStaging[i, 5] - mStaging[i, 4] + if cutlass.const_expr(gen): + batch_idx, head_idx, batch_start, batch_end, num_chunks_b = gen_item_bounds(b_t, n_heads_out, mCuSeqlens, i) + key = num_chunks_b + else: + key = mStaging[i, 5] - mStaging[i, 4] sKey[i] = key sIdx[i] = i kmin = kmin if kmin < key else key @@ -600,11 +651,10 @@ def _order_kernel( nvvm.atomicrmw("max", sSpread.data_ptr(1), kmax, space=nvvm.SharedSpace.shared_cta) nvvm.barrier_cta_sync() if sSpread[0] == sSpread[1]: - # every key equal (uniform batches): any order is LPT, copy through + # every key equal (uniform batches): copy through i2 = tidx while i2 < n: - for f in cutlass.range_constexpr(WORK_ITEM_FIELDS): - mWorkItems[i2, cutlass.Int32(f)] = mStaging[i2, cutlass.Int32(f)] + write_item(gen, b_t, n_heads_out, mCuSeqlens, mStaging, mWorkItems, i2, i2) i2 = i2 + cutlass.Int32(ORDER_THREADS) else: k = cutlass.Int32(2) @@ -634,12 +684,12 @@ def _order_kernel( i = tidx + cutlass.Int32(e * ORDER_THREADS) if i < n: src = sIdx[i] - for f in cutlass.range_constexpr(WORK_ITEM_FIELDS): - mWorkItems[i, cutlass.Int32(f)] = mStaging[src, cutlass.Int32(f)] + write_item(gen, b_t, n_heads_out, mCuSeqlens, mStaging, mWorkItems, i, src) @cute.jit -def _launch( +def launch( + split: cutlass.Constexpr[bool], b_t: cutlass.Constexpr[int], log_gate: cutlass.Constexpr[bool], safe_gate: cutlass.Constexpr[bool], @@ -653,12 +703,12 @@ def _launch( batch_size: cutlass.Int32, log2_thresh: cutlass.Float32, gate_scale_log2: cutlass.Float32, - mGate: cute.Tensor, + mGate: cute.Tensor | None, mALog: cute.Tensor | None, mDtBias: cute.Tensor | None, mCuSeqlens: cute.Tensor, - mChunkVals: cute.Tensor, - mStaging: cute.Tensor, + mChunkVals: cute.Tensor | None, + mStaging: cute.Tensor | None, mWorkItems: cute.Tensor, mCount: cute.Tensor, mSched: cute.Tensor | None, @@ -666,74 +716,78 @@ def _launch( n_walk_ctas: cutlass.Int32, stream: cuda.CUstream, ) -> None: - if cutlass.const_expr(gate_channels > 0): - _scan_kernel( - b_t, - log_gate, - safe_gate, - gate_channels, - overhead_chunks, - has_sched, - n_heads_out, - n_tiles, - num_sms, - ideal_chunks, - batch_size, - gate_scale_log2, - mGate, - mALog, - mDtBias, - mCuSeqlens, - mChunkVals, - mCount, - mSched, - ).launch( - grid=(n_scan_ctas, n_heads_out, 1), - block=(SCAN_THREADS, 1, 1), - stream=stream, - ) - else: - _scan_scalar_kernel( + if cutlass.const_expr(split): + if cutlass.const_expr(gate_channels > 0): + scan_kernel( + b_t, + log_gate, + safe_gate, + gate_channels, + overhead_chunks, + n_heads_out, + n_tiles, + num_sms, + ideal_chunks, + batch_size, + gate_scale_log2, + mGate, + mALog, + mDtBias, + mCuSeqlens, + mChunkVals, + mCount, + ).launch( + grid=(n_scan_ctas, n_heads_out, 1), + block=(SCAN_THREADS, 1, 1), + stream=stream, + ) + else: + scan_scalar_kernel( + b_t, + log_gate, + overhead_chunks, + n_heads_out, + n_tiles, + num_sms, + ideal_chunks, + batch_size, + mGate, + mCuSeqlens, + mChunkVals, + mCount, + ).launch( + grid=(n_scan_ctas, (n_heads_out + cutlass.Int32(WARP_SIZE - 1)) // cutlass.Int32(WARP_SIZE), 1), + block=(SCAN_THREADS, 1, 1), + stream=stream, + ) + walk_kernel( b_t, - log_gate, overhead_chunks, - has_sched, n_heads_out, n_tiles, num_sms, ideal_chunks, - batch_size, - mGate, + log2_thresh, mCuSeqlens, mChunkVals, + mStaging, mCount, - mSched, ).launch( - grid=(n_scan_ctas, (n_heads_out + cutlass.Int32(WARP_SIZE - 1)) // cutlass.Int32(WARP_SIZE), 1), - block=(SCAN_THREADS, 1, 1), + grid=(n_walk_ctas, 1, 1), + block=(THREADS_PER_BLOCK, 1, 1), stream=stream, ) - _walk_kernel( + order_kernel( + not split, + has_sched, b_t, - overhead_chunks, n_heads_out, n_tiles, - num_sms, - ideal_chunks, - log2_thresh, mCuSeqlens, - mChunkVals, - mStaging, - mCount, - ).launch( - grid=(n_walk_ctas, 1, 1), - block=(THREADS_PER_BLOCK, 1, 1), - stream=stream, - ) - _order_kernel( mStaging, mCount, mWorkItems, + mSched, ).launch( grid=(1, 1, 1), block=(ORDER_THREADS, 1, 1), @@ -741,7 +795,7 @@ def _launch( ) -_compiled_cache = {} +compiled_cache = {} def build_split_table( @@ -750,12 +804,12 @@ def build_split_table( work_items, work_count, *, - ideal_chunks, + ideal_chunks=None, n_tiles, num_sms, b_t, - chunk_scratch, - item_scratch, + chunk_scratch=None, + item_scratch=None, log2_threshold=None, log_gate=False, safe_gate=False, @@ -763,6 +817,7 @@ def build_split_table( dt_bias=None, gate_lower_bound=None, sched_ctr=None, + split=True, stream, ) -> None: """Fill ``work_items``/``work_count`` with the split-K partition of @@ -775,17 +830,26 @@ def build_split_table( ``gate_lower_bound * sigmoid(exp(a_log) * (g + dt_bias))`` per element, so cuts land on true decay values. + With ``split=False`` the scan and walk never launch: the order kernel + alone synthesizes the no-cuts table — the uncut whole-sequence item per + (batch, head), LPT-sorted by sequence length — and ``ideal_chunks`` / + ``chunk_scratch`` / ``item_scratch`` / the gate contents are unused. + Batch-invariant mode and coarse checkpoint cadences (cuts may not cross + a checkpoint period) take this path. + ``work_items`` and ``item_scratch`` are ``(max_items, - WORK_ITEM_FIELDS)`` int32 with ``max_items >= max_work_items(...)``; - ``work_count`` is ``(1,)`` int32 (zeroed here by the scan kernel, as - is every cell of ``sched_ctr`` when passed — the main kernels' int32 - ticket rings, ``(2,)`` per kernel launch that consumes this table, - since the kernels leave their ring dirty on exit); + WORK_ITEM_FIELDS)`` int32 with ``max_items >= max_work_items(...)`` + (``>= n_tiles`` rows and no ``item_scratch`` with ``split=False``); + ``work_count`` is ``(1,)`` int32. Every cell of ``sched_ctr`` when + passed — the main kernels' int32 ticket rings, ``(2,)`` per kernel + launch that consumes this table, dirty on exit — is zeroed by the + order kernel, which runs in both modes; the count is zeroed by the + scan (split) or written by the order kernel (non-split). ``chunk_scratch`` is ``(>= chunk_scratch_rows(total_tokens, B, b_t), HO)`` fp32 (contents managed here). Runs entirely on device — no host synchronization.""" if log2_threshold is None: - log2_threshold = _DEFAULT_LOG2_THRESHOLD + log2_threshold = DEFAULT_LOG2_THRESHOLD if len(gate.shape) not in (2, 3): raise ValueError(f"gate must be (total_tokens, HO) or (total_tokens, HO, DK), got {tuple(gate.shape)}") gate_channels = gate.shape[2] if len(gate.shape) == 3 else 0 @@ -796,35 +860,67 @@ def build_split_table( if not safe_gate: a_log = None dt_bias = None - if gate_channels and gate_channels % WARP_SIZE != 0: - raise ValueError(f"per-channel gate dim must be a multiple of {WARP_SIZE}, got {gate_channels}") - if gate_channels and gate_channels % 128 == 0 and gate.data_ptr() % 16 != 0: - raise ValueError("per-channel gate base must be 16-byte aligned (vectorized scan loads)") gate_scale_log2 = float(gate_lower_bound) * RCP_LN2 if safe_gate else 0.0 n_heads_out = gate.shape[1] batch_size = cu_seqlens.shape[0] - 1 - n_walk_ctas = batch_size * n_heads_out - need_rows = chunk_scratch_rows(gate.shape[0], batch_size, b_t) - n_scan_ctas = -(-need_rows // (SCAN_WARPS * SCAN_ROWS_PER_WARP)) - if len(chunk_scratch.shape) != 2 or chunk_scratch.shape[0] < need_rows or chunk_scratch.shape[1] != n_heads_out: - raise ValueError(f"chunk_scratch must be (>= {need_rows}, {n_heads_out}) fp32, got {tuple(chunk_scratch.shape)}") - if tuple(item_scratch.shape) != tuple(work_items.shape) or work_items.shape[1] != WORK_ITEM_FIELDS: - raise ValueError( - f"item_scratch must match work_items (max_items, {WORK_ITEM_FIELDS}) int32, got {tuple(item_scratch.shape)} vs {tuple(work_items.shape)}" - ) + if split: + if ideal_chunks is None or chunk_scratch is None or item_scratch is None: + raise ValueError("split=True requires ideal_chunks, chunk_scratch, and item_scratch") + if gate_channels and gate_channels % WARP_SIZE != 0: + raise ValueError(f"per-channel gate dim must be a multiple of {WARP_SIZE}, got {gate_channels}") + if gate_channels and gate_channels % 128 == 0 and data_ptr(gate) % 16 != 0: + raise ValueError("per-channel gate base must be 16-byte aligned (vectorized scan loads)") + n_walk_ctas = batch_size * n_heads_out + need_rows = chunk_scratch_rows(gate.shape[0], batch_size, b_t) + n_scan_ctas = -(-need_rows // (SCAN_WARPS * SCAN_ROWS_PER_WARP)) + if len(chunk_scratch.shape) != 2 or chunk_scratch.shape[0] < need_rows or chunk_scratch.shape[1] != n_heads_out: + raise ValueError(f"chunk_scratch must be (>= {need_rows}, {n_heads_out}) fp32, got {tuple(chunk_scratch.shape)}") + if tuple(item_scratch.shape) != tuple(work_items.shape) or work_items.shape[1] != WORK_ITEM_FIELDS: + raise ValueError( + f"item_scratch must match work_items (max_items, {WORK_ITEM_FIELDS}) int32, got {tuple(item_scratch.shape)} vs {tuple(work_items.shape)}" + ) + else: + # no-cuts table: the order kernel synthesizes it from cu_seqlens; the + # gate and every scan/walk operand are unused (normalized out of the + # compile key so all variants share one specialization per b_t/HO) + if work_items.shape[0] < n_tiles or work_items.shape[1] != WORK_ITEM_FIELDS: + raise ValueError(f"work_items must be (>= {n_tiles}, {WORK_ITEM_FIELDS}) int32, got {tuple(work_items.shape)}") + log_gate = False + safe_gate = False + gate_channels = 0 + a_log = None + dt_bias = None + gate_scale_log2 = 0.0 + gate = None + chunk_scratch = None + item_scratch = None + ideal_chunks = 0 + n_scan_ctas = 0 + n_walk_ctas = 0 overhead_chunks = max(1, OVERHEAD_TOKENS // b_t) cu_stream = cuda.CUstream(int(stream)) - key = (b_t, n_heads_out, bool(log_gate), bool(safe_gate), gate_channels, sched_ctr is not None) - if key not in _compiled_cache: - - def _dyn(t): - c = from_dlpack(t, assumed_align=4) - c.mark_compact_shape_dynamic(mode=0, stride_order=tuple(range(len(t.shape))), divisibility=1) - return c - - _compiled_cache[key] = cute.compile( - _launch, + key = (bool(split), b_t, n_heads_out, bool(log_gate), bool(safe_gate), gate_channels, sched_ctr is not None, str(cu_seqlens.dtype)) + if key not in compiled_cache: + + dt_bias_c = None + if safe_gate: + dt_bias_c = from_dlpack(dt_bias, assumed_align=4) + dt_bias_c.mark_compact_shape_dynamic(mode=0, stride_order=tuple(range(len(dt_bias.shape))), divisibility=1) + chunk_scratch_c = None + item_scratch_c = None + if split: + chunk_scratch_c = from_dlpack(chunk_scratch, assumed_align=4) + chunk_scratch_c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + item_scratch_c = from_dlpack(item_scratch, assumed_align=4) + item_scratch_c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + work_items_c = from_dlpack(work_items, assumed_align=4) + work_items_c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + work_count_c = from_dlpack(work_count, assumed_align=4) + work_count_c.mark_compact_shape_dynamic(mode=0, stride_order=(0,), divisibility=1) + compiled_cache[key] = cute.compile( + launch, + bool(split), b_t, bool(log_gate), bool(safe_gate), @@ -838,21 +934,21 @@ def _dyn(t): cutlass.Int32(batch_size), cutlass.Float32(log2_threshold), cutlass.Float32(gate_scale_log2), - _dyn(gate), + from_dlpack(gate, assumed_align=4).mark_layout_dynamic(leading_dim=len(gate.shape) - 1) if split else None, from_dlpack(a_log, assumed_align=4).mark_layout_dynamic() if safe_gate else None, - _dyn(dt_bias) if safe_gate else None, + dt_bias_c, from_dlpack(cu_seqlens, assumed_align=4).mark_layout_dynamic(), - _dyn(chunk_scratch), - _dyn(item_scratch), - _dyn(work_items), - _dyn(work_count), + chunk_scratch_c, + item_scratch_c, + work_items_c, + work_count_c, from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() if sched_ctr is not None else None, cutlass.Int32(n_scan_ctas), cutlass.Int32(n_walk_ctas), cu_stream, options="--enable-tvm-ffi", ) - _compiled_cache[key]( + compiled_cache[key]( n_heads_out, n_tiles, num_sms, diff --git a/python/cudnn/linear_attention/frost/common/thd.py b/python/cudnn/linear_attention/frost/common/thd.py index 70fb26991..c81f459be 100644 --- a/python/cudnn/linear_attention/frost/common/thd.py +++ b/python/cudnn/linear_attention/frost/common/thd.py @@ -4,214 +4,125 @@ """Shared THD / varlen (packed ``[T,H,D]`` + ``cu_seqlens``) device helpers. * :data:`TENSOR_MAP_QWORDS` — int64 words per 128-byte TMA descriptor. -* :func:`build_qkv_load_descs_kernel` — the separate device kernel that - builds a per-(batch x head) TMA-descriptor array in GMEM for varlen - loads/stores over a packed ``[T,H,D]`` tensor. -* :func:`build_h_descs_kernel` — its per-chunk-H sibling; derives the - per-sequence H offsets from the token ``cu_seqlens`` in place of a - caller-computed ``cu_h`` prefix array. -* :func:`build_state_descs_kernel` — per-(batch x head) descriptors over a - DENSE ``[N, HO, K, V]`` state tensor (one entry per slot). -* :func:`downcast_state_kernel` — elementwise fp32 -> io copy of the - initial state into the buffer the state descriptors read. +* :func:`emit_seq_descs` — device helper (one electing thread) that builds + a per-BATCH TMA-descriptor array in GMEM for varlen loads/stores over a + packed ``[T,H,D]`` tensor whose head axis is a load coordinate. Each op + calls it from a single combined ``build_all_descs_kernel`` launch (one + warp per array). +* :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 import cutlass.cute as cute import cutlass.experimental.primitives as nvvm -import cutlass.experimental.cuda.tensor_map as _tma from cutlass.base_dsl.typing import Pointer -import cuda.bindings.driver as _cuda_driver # noqa: F401 (cute.compile pulls cuda) TENSOR_MAP_QWORDS = 128 // 8 -@cute.kernel -def build_qkv_load_descs_kernel( - base_desc: cutlass.GridConstant[_tma.TensorMap], - desc_words: cute.Tensor, - cu_seqlens: cute.Tensor, - base_ptr: cute.Tensor, +@cute.jit +def emit_seq_descs( + base_desc, + desc_words, + cu_seqlens, + base_ptr, n_batch: cutlass.Int32, - n_heads: cutlass.Int32, - head_group: cutlass.Int32, - head_stride: cutlass.Int32, row_stride: cutlass.Int32, seq_ord: cutlass.Constexpr[int], ) -> None: - """Build a per-(batch x head) TMA-descriptor ARRAY for a VARLEN (THD) - Q/K/V/O tile over a packed ``[T,H,D]`` tensor + ``cu_seqlens``. It - re-points ``GLOBAL_ADDRESS`` per head so the flat output-head index - ``head_idx`` maps to its KV head (``head_idx // head_group``), - reproducing the GQA nested ``(h_r, h_v)`` head mode with stride-0 - replication. The cute-side ``tma_tensor[None, None, head_idx]`` head - indexing collapses to: - - * **identity (Q/O):** address head offset = ``head_idx * head_stride`` - (``head_group == 1``; ``head_idx // 1 == head_idx``). - * **grouped (K/V):** address head offset = - ``(head_idx // head_group) * head_stride`` — the stride-0 ``h_r`` - sub-mode means ``head_group`` consecutive Q heads share one KV head. - - Pass ``head_group == 1`` for identity and ``head_group == h_q // h_v`` - for K/V. The descriptor array is laid out ``[batch][head]`` (head-minor): - slot ``(b * n_heads + h)``. - - GLOBAL_ADDRESS folds BOTH the per-sequence token start - (``cu_seqlens[b] * row_stride``) AND the per-head offset - (``(head_idx // head_group) * head_stride``); GLOBAL_DIM[``seq_ord``] is - capped to the per-sequence token COUNT (``cu_seqlens[b+1] - cu_seqlens[b]``) - so a load box past the sequence end is OOB-clipped. Because the - GLOBAL_ADDRESS already points at the sequence's first token, the consumer - issues the TMA with a token coordinate of **0** (not the absolute packed - offset). - - Released via the GENERIC->TENSORMAP proxy fence; the consumer - (``tma_load_tile`` with ``gmem_slice.desc_ptr`` set) acquire-fences each - slot before the load. ``seq_ord`` is a compile-time ord - (``tensormap_replace`` ord must be a Python int).""" - if nvvm.elect_sync(): - desc_base = desc_words.iterator.raw_ptr() - src_words = Pointer(base_desc.get_ptr(), dtype=cutlass.Int64) - cu = cutlass.make_array_view(cu_seqlens) - base = base_ptr.iterator.raw_ptr() - for b in cutlass.range(0, n_batch, 1, unroll=1): - cu_b = cutlass.Int32(cu[b]) - s_b = cutlass.Int32(cu[b + cutlass.Int32(1)]) - cu_b - # Int64: the H descs cross 2^31 elements near cu[b] ~ 2k - # (row_stride = HO*DK*DV); inputs follow at larger B*T. - tok0 = cutlass.Int64(cu_b) * cutlass.Int64(row_stride) - for h in cutlass.range(0, n_heads, 1, unroll=1): - slot = b * n_heads + h - dptr = desc_base + slot * cutlass.Int32(TENSOR_MAP_QWORDS) - for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS): - (dptr + i).store((src_words + i).load()) - kv_h = h // head_group - head_off = cutlass.Int64(kv_h) * cutlass.Int64(head_stride) - addr = base + (tok0 + head_off) - nvvm.tensormap_replace( - nvvm.TensormapField.GLOBAL_ADDRESS, - dptr, - new_value=addr.toint(cutlass.Int64), - ) - nvvm.tensormap_replace( - nvvm.TensormapField.GLOBAL_DIM, - dptr, - new_value=s_b, - ord=seq_ord, - ) - nvvm.fence_proxy_release( - nvvm.MemScope.GPU, - from_proxy=nvvm.Proxy.GENERIC, - to_proxy=nvvm.Proxy.TENSORMAP, + """Per-BATCH TMA-descriptor array for a VARLEN (THD) tensor whose base + descriptor carries the head axis as a real dimension (``(d, head, + token)``); the head index is a load COORDINATE, so only the sequence + base and length are patched per slot. GLOBAL_ADDRESS folds + ``cu_seqlens[b] * row_stride`` (Int64); GLOBAL_DIM[``seq_ord``] is + capped to the per-sequence token count so tail loads zero-fill and tail + stores clip in hardware. GQA/GVA head grouping happens at the issue + site (``head_idx // group`` with a static group), not here. Runs on + one electing thread; the calling warp elects and release-fences + (GENERIC->TENSORMAP).""" + desc_base = desc_words.iterator.raw_ptr() + src_words = Pointer(base_desc.get_ptr(), dtype=cutlass.Int64) + cu = cutlass.make_array_view(cu_seqlens) + base = base_ptr.iterator.raw_ptr() + for b in cutlass.range(0, n_batch, 1, unroll=1): + cu_b = cutlass.Int32(cu[b]) + s_b = cutlass.Int32(cu[b + cutlass.Int32(1)]) - cu_b + dptr = desc_base + b * cutlass.Int32(TENSOR_MAP_QWORDS) + for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS): + (dptr + i).store((src_words + i).load()) + addr = base + cutlass.Int64(cu_b) * cutlass.Int64(row_stride) + nvvm.tensormap_replace( + nvvm.TensormapField.GLOBAL_ADDRESS, + dptr, + new_value=addr.toint(cutlass.Int64), ) - - -@cute.kernel -def downcast_state_kernel( - mS0: cute.Tensor, - mOut: cute.Tensor, - n: cutlass.Int32, -) -> None: - """Flat elementwise copy of the fp32 initial state into the io-dtype - buffer the backward's per-(b,h) state descriptors read (1-D views, one - element per thread).""" - idx = cutlass.Int32(cute.arch.block_idx()[0]) * cutlass.Int32(128) + cutlass.Int32(cute.arch.thread_idx()[0]) - if idx < n: - mOut[idx] = mS0[idx].to(mOut.element_type) - - -@cute.kernel -def build_state_descs_kernel( - base_desc: cutlass.GridConstant[_tma.TensorMap], - desc_words: cute.Tensor, - base_ptr: cute.Tensor, - n_batch: cutlass.Int32, - n_heads: cutlass.Int32, - ent_stride: cutlass.Int32, -) -> None: - """Per-(batch x head) TMA-descriptor array over a DENSE state tensor - ``[N, HO, K, V]``: slot ``(b * n_heads + h)`` gets GLOBAL_ADDRESS - pointing at its ``[K, V]`` tile (``slot * ent_stride`` elements in). - All dims come baked from the base descriptor (one entry per slot), so - only the address is patched.""" - if nvvm.elect_sync(): - desc_base = desc_words.iterator.raw_ptr() - src_words = Pointer(base_desc.get_ptr(), dtype=cutlass.Int64) - base = base_ptr.iterator.raw_ptr() - n_slots = n_batch * n_heads - for s in cutlass.range(0, n_slots, 1, unroll=1): - dptr = desc_base + s * cutlass.Int32(TENSOR_MAP_QWORDS) - for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS): - (dptr + i).store((src_words + i).load()) - addr = base + cutlass.Int64(s) * cutlass.Int64(ent_stride) - nvvm.tensormap_replace( - nvvm.TensormapField.GLOBAL_ADDRESS, - dptr, - new_value=addr.toint(cutlass.Int64), - ) - nvvm.fence_proxy_release( - nvvm.MemScope.GPU, - from_proxy=nvvm.Proxy.GENERIC, - to_proxy=nvvm.Proxy.TENSORMAP, + nvvm.tensormap_replace( + nvvm.TensormapField.GLOBAL_DIM, + dptr, + new_value=s_b, + ord=seq_ord, ) -@cute.kernel -def build_h_descs_kernel( - base_desc: cutlass.GridConstant[_tma.TensorMap], - desc_words: cute.Tensor, - cu_seqlens: cute.Tensor, - base_ptr: cute.Tensor, +@cute.jit +def emit_checkpoint_seq_descs( + base_desc, + desc_words, + cu_seqlens, + base_ptr, n_batch: cutlass.Int32, - n_heads: cutlass.Int32, - head_stride: cutlass.Int32, row_stride: cutlass.Int32, every_n: cutlass.Int32, seq_ord: cutlass.Constexpr[int], ) -> None: - """Per-(batch x head) TMA-descriptor array for the per-chunk H tensor. - - Unlike :func:`build_qkv_load_descs_kernel` this derives the per-sequence - H offsets from the TOKEN ``cu_seqlens`` on the fly instead of taking a - caller-computed ``cu_h`` prefix array: ``count_b = (seqlen_b - 1) // - every_n`` (0 for empty sequences), running-prefix-summed. - GLOBAL_DIM[``seq_ord``] is capped to ``count_b`` exactly as the - cu_h-differences cap did. H heads are identity-mapped (no GQA - grouping).""" - if nvvm.elect_sync(): - desc_base = desc_words.iterator.raw_ptr() - src_words = Pointer(base_desc.get_ptr(), dtype=cutlass.Int64) - cu = cutlass.make_array_view(cu_seqlens) - base = base_ptr.iterator.raw_ptr() - 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 = cnt if s_tok > 0 else cutlass.Int32(0) - h0 = run - run = run + cnt - ent0 = cutlass.Int64(h0) * cutlass.Int64(row_stride) - for h in cutlass.range(0, n_heads, 1, unroll=1): - slot = b * n_heads + h - dptr = desc_base + slot * cutlass.Int32(TENSOR_MAP_QWORDS) - for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS): - (dptr + i).store((src_words + i).load()) - head_off = cutlass.Int64(h) * cutlass.Int64(head_stride) - addr = base + (ent0 + head_off) - nvvm.tensormap_replace( - nvvm.TensormapField.GLOBAL_ADDRESS, - dptr, - new_value=addr.toint(cutlass.Int64), - ) - nvvm.tensormap_replace( - nvvm.TensormapField.GLOBAL_DIM, - dptr, - new_value=cnt, - ord=seq_ord, - ) - nvvm.fence_proxy_release( - nvvm.MemScope.GPU, - from_proxy=nvvm.Proxy.GENERIC, - to_proxy=nvvm.Proxy.TENSORMAP, + """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 + 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 + fences.""" + desc_base = desc_words.iterator.raw_ptr() + src_words = Pointer(base_desc.get_ptr(), dtype=cutlass.Int64) + cu = cutlass.make_array_view(cu_seqlens) + base = base_ptr.iterator.raw_ptr() + 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 = cnt if s_tok > 0 else cutlass.Int32(0) + checkpoint_base = run + run = run + cnt + dptr = desc_base + b * cutlass.Int32(TENSOR_MAP_QWORDS) + for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS): + (dptr + i).store((src_words + i).load()) + addr = base + cutlass.Int64(checkpoint_base) * cutlass.Int64(row_stride) + nvvm.tensormap_replace( + nvvm.TensormapField.GLOBAL_ADDRESS, + dptr, + new_value=addr.toint(cutlass.Int64), ) + nvvm.tensormap_replace( + nvvm.TensormapField.GLOBAL_DIM, + dptr, + 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 3bcd0fa8a..39487e44c 100644 --- a/python/cudnn/linear_attention/frost/gdn2_engine.py +++ b/python/cudnn/linear_attention/frost/gdn2_engine.py @@ -3,7 +3,8 @@ """FROST GDN-2 engine: GDN2 nodes on the chunked prefill kernel (``kernel/gdn2_prefill_f16.py``, Blackwell SM100/SM103, bf16/fp16, BT=16). -Forward only (GDN2_BWD declines); the only GDN-2 engine — no cuTile +Forward + backward (GDN2_BWD on ``kernel/gdn2_bprop_f16.py`` with a +checkpoint regen on ``kernel/gdn2_recompute_f16.py``); the only GDN-2 engine — no cuTile fallback.""" from __future__ import annotations @@ -14,61 +15,25 @@ from cudnn import behavior_note from cudnn.engines.base import BaseEngine, CompiledPlan -from cudnn.frost import buffers -from cudnn.frost.workspace import Workspace, WorkspaceLayout, carve_plan -from ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair - - -def _the_gdn2_node(graph): - nodes = list(graph.nodes) - if len(nodes) != 1: - return None - node = nodes[0] - if getattr(node.node_type, "name", None) not in ("GDN2", "GDN2_BWD"): - return None - return node - - -def _check_common(node) -> None: - """Shape/dtype gates for the prefill kernel.""" - import cudnn - - q, k, v = (node.inputs[p] for p in ("q", "k", "v")) - io_dtypes = {q.get_data_type(), k.get_data_type(), v.get_data_type()} - {None} - if len(io_dtypes) > 1: - raise NotImplementedError(f"Gdn2FrostEngine: q/k/v dtypes must match, got {io_dtypes}") - for p, t in (("q", q), ("k", k), ("v", v)): - if t.get_data_type() not in (cudnn.data_type.BFLOAT16, cudnn.data_type.HALF, None): - raise NotImplementedError(f"Gdn2FrostEngine: '{p}' must be bf16 or fp16, got {t.get_data_type()}") - if not t.dim or len(t.dim) != 3: - raise NotImplementedError(f"Gdn2FrostEngine: '{p}' must be THD [total_T, heads, dim]") - if q.dim[-1] != 128 or v.dim[-1] != 128: - raise NotImplementedError(f"Gdn2FrostEngine: head dims must be 128 (the recurrent state is 128x128), got K={q.dim[-1]} V={v.dim[-1]}") - if k.dim[1] != q.dim[1]: - raise NotImplementedError(f"Gdn2FrostEngine: q and k head counts differ ({q.dim[1]} vs {k.dim[1]})") - if v.dim[1] % q.dim[1] != 0: - raise NotImplementedError(f"Gdn2FrostEngine: v heads ({v.dim[1]}) must be a multiple of q heads ({q.dim[1]})") - - # kernel-native operand dtypes: buffers pass through without staging - fp32 = cudnn.data_type.FLOAT - io = q.get_data_type() - _require_dtype("Gdn2FrostEngine", node, "g", fp32) - if io is not None: - _require_dtype("Gdn2FrostEngine", node, "beta", io) - _require_dtype("Gdn2FrostEngine", node, "w", io) - _require_dtype("Gdn2FrostEngine", node, "cu_seqlens", cudnn.data_type.INT32) - _require_dtype("Gdn2FrostEngine", node, "initial_state", (fp32, cudnn.data_type.BFLOAT16)) - _require_dtype("Gdn2FrostEngine", node, "final_state", (fp32, cudnn.data_type.BFLOAT16), out=True) - _require_state_pair("Gdn2FrostEngine", node) +from cudnn.frost.buffers import current_device_id +from cudnn.frost.device import multiprocessor_count +from cudnn.frost.workspace import WorkspaceLayout, carve_plan +from ..graph_analyzer import FrostLaPlan, frost_la_gate, require, analyze def build_gdn2(graph): """The expensive step: import the kernel module (pulls in the Cutlass primitives; the cute.compile itself is cached inside the kernel per static config and runs on first execute, when the real buffers are known).""" - node = _the_gdn2_node(graph) - if node is None or node.node_type.name != "GDN2": - raise ValueError("build_gdn2: graph does not contain exactly one GDN2 node") + nodes = list(graph.nodes) + if len(nodes) != 1 or getattr(nodes[0].node_type, "name", None) not in ("GDN2", "GDN2_BWD"): + raise ValueError("build_gdn2: graph does not contain exactly one GDN2/GDN2_BWD node") + node = nodes[0] + if node.node_type.name == "GDN2_BWD": + from .kernel import gdn2_bprop_f16 as bwd_mod + from .kernel import gdn2_recompute_f16 as regen_mod + + return CompiledGdn2Bwd(node, bwd_mod, regen_mod) from .kernel import gdn2_prefill_f16 as kernel_mod return CompiledGdn2(node, kernel_mod) @@ -77,90 +42,122 @@ def build_gdn2(graph): class Gdn2FrostEngine(BaseEngine): """FROST chunked-kernel backend for single-node GDN-2 graphs (THD layout). - The only GDN-2 engine (SM100/SM103, forward only); declines ``GDN2_BWD`` - (the FROST backward kernel is a stub).""" + The only GDN-2 engine (SM100/SM103); GDN2_BWD runs on the FROST backward + kernel with a forward checkpoint recompute when the graph has no ``state_checkpoints`` input.""" name = "gdn2_frost" behavior_notes = (behavior_note.RUNTIME_COMPILATION,) # JIT-compiled at build_plans() def check_support(self, graph) -> None: - node = _the_gdn2_node(graph) - if node is None: - raise NotImplementedError("Gdn2FrostEngine supports exactly one GDN2 node") - if node.node_type.name == "GDN2_BWD": - raise NotImplementedError("Gdn2FrostEngine: the FROST GDN-2 backward kernel is a stub") - sm = buffers.current_sm() - if sm is None or not (100 <= sm <= 103): - raise NotImplementedError(f"Gdn2FrostEngine requires SM100-SM103 (found {sm})") - try: - import cutlass.experimental.primitives # noqa: F401 — availability probe: ImportError = decline - except ImportError as exc: - raise NotImplementedError(f"Gdn2FrostEngine requires the Cutlass DSL with cutlass.experimental.primitives: {exc}") from exc - for port in ("q", "k", "v", "g", "beta", "w", "cu_seqlens"): - if port not in node.inputs: - raise NotImplementedError(f"Gdn2FrostEngine: GDN2 node '{node.name}' is missing input '{port}'") - if int(node.params.get("checkpoint_every_n_tokens", 0) or 0) or "H" in node.outputs: - raise NotImplementedError("Gdn2FrostEngine: the per-chunk H output lands with the backward kernel (jopark/kda_gdn2_bprop)") - if node.node_type.name != "GDN2" and node.params.get("use_beta_w_sigmoid", False): - raise NotImplementedError("Gdn2FrostEngine: use_beta_w_sigmoid is a forward-node attribute") - _check_common(node) + import cudnn + + facts = graph._facts_for(analyze) + frost_la_gate("Gdn2FrostEngine", facts, "GDN2") + ckpt = facts.checkpoint_every_n_tokens + if ckpt and (facts.is_bwd or ckpt % 16 != 0): + raise NotImplementedError(f"Gdn2FrostEngine: checkpoint_every_n_tokens must be a positive multiple of 16 on the GDN-2 node (got {ckpt})") + if not facts.gates_at_ho: + raise NotImplementedError(f"Gdn2FrostEngine: g/beta/w must carry HO = max(q, v) heads ({facts.h_o})") + if facts.is_bwd and facts.safe_gate: + raise NotImplementedError("Gdn2FrostEngine: safe_gate is a forward-node attribute") + fp32 = cudnn.data_type.FLOAT + if facts.io_dtype is not None: + require("Gdn2FrostEngine", "beta", facts.beta_dtype, facts.io_dtype) + require("Gdn2FrostEngine", "w", facts.w_dtype, facts.io_dtype) + require("Gdn2FrostEngine", "a_log", facts.a_log_dtype, fp32) + require("Gdn2FrostEngine", "dt_bias", facts.dt_bias_dtype, fp32) + if facts.is_bwd: + for port, got in (("dO", facts.do_dtype), ("state_checkpoints", facts.state_checkpoints_dtype)): + if got not in (facts.io_dtype, None): + raise NotImplementedError(f"Gdn2FrostEngine: '{port}' must match the io dtype") + require("Gdn2FrostEngine", "initial_state", facts.state_dtype, fp32) + require("Gdn2FrostEngine", "d_final_state", facts.d_final_state_dtype, fp32) + require("Gdn2FrostEngine", "d_initial_state", facts.d_initial_state_dtype, fp32) + require("Gdn2FrostEngine", "dG", facts.dg_dtype, fp32) + if facts.io_dtype is not None: + require("Gdn2FrostEngine", "dBeta", facts.dbeta_dtype, facts.io_dtype) + require("Gdn2FrostEngine", "dW", facts.dw_dtype, facts.io_dtype) + else: + state_dtypes = (fp32, cudnn.data_type.BFLOAT16) + require("Gdn2FrostEngine", "initial_state", facts.state_dtype, state_dtypes) + require("Gdn2FrostEngine", "final_state", facts.final_state_dtype, state_dtypes) + if facts.io_dtype is not None: + require("Gdn2FrostEngine", "state_checkpoints", facts.state_checkpoints_out_dtype, facts.io_dtype) + if not facts.state_pair_match: + raise NotImplementedError("Gdn2FrostEngine: initial_state and final_state dtypes must match") def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: - return _FrostPlan(build_gdn2(graph)) + return FrostLaPlan(build_gdn2(graph)) class CompiledGdn2: """Compiled FROST GDN-2 plan: a callable over the resolved node buffers.""" def __init__(self, node, kernel_mod): - from .common.split_k import WORK_ITEM_FIELDS, chunk_scratch_rows, compute_ideal_chunks, max_work_items + from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items - self._node = node - self._kernel = kernel_mod + self.node = node + self.kernel = kernel_mod + self.build_split_table = build_split_table + self.plan_name = "Gdn2FrostEngine (GDN2)" scale = node.params.get("scale") - self._scale = float(scale) if scale is not None else 1.0 / math.sqrt(node.inputs["q"].dim[-1]) - self._use_qk_l2norm = bool(node.params.get("use_qk_l2norm", False)) - self._use_beta_w_sigmoid = bool(node.params.get("use_beta_w_sigmoid", False)) - self._has_fs = "final_state" in node.outputs + self.scale = float(scale) if scale is not None else 1.0 / math.sqrt(node.inputs["q"].dim[-1]) + self.use_qk_l2norm = bool(node.params.get("use_qk_l2norm", False)) + self.safe_gate = bool(node.params.get("safe_gate", False)) + glb = node.params.get("gate_lower_bound") + self.gate_lower_bound = float(glb) if glb is not None else kernel_mod.DEFAULT_GATE_LOWER_BOUND + self.has_final_state = "final_state" in node.outputs + self.has_state_checkpoints = "state_checkpoints" in node.outputs + self.ckpt = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) + self.batch_invariant = bool(node.params.get("batch_invariant", False)) q, g = node.inputs["q"], node.inputs["g"] - self._b_t = kernel_mod.CFG.B_T + 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 total = q.dim[0] HO = g.dim[1] B = node.inputs["cu_seqlens"].dim[0] - 1 layout = WorkspaceLayout() - self._off_sched = layout.add(8) # [ticket, done] for the dynamic scheduler - self._num_sm = kernel_mod._device_sm_count() - self._ideal = compute_ideal_chunks(total, HO, self._num_sm, self._b_t) - self._n_tiles = B * HO - self._work_item_rows = max_work_items(total, B, HO, self._ideal, self._b_t, self._num_sm) - self._n_heads_out = HO - self._off_work_items = layout.add(self._work_item_rows * WORK_ITEM_FIELDS * 4) - self._off_item_scratch = layout.add(self._work_item_rows * WORK_ITEM_FIELDS * 4) - self._off_work_count = layout.add(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._tensormap_bytes = kernel_mod.get_workspace_size(B, HO, HO) - self._off_tensormaps = layout.add(self._tensormap_bytes, align=128) - self._ws_bytes = layout.size - - self._carve = carve_plan( - "gdn2", - [ - (self._off_sched, "int32", (2,)), - (self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), - (self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), - (self._off_work_count, "int32", (1,)), - (self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, HO)), - (self._off_tensormaps, "int64", (self._tensormap_bytes // 8,)), - ], - ) + self.off_sched = layout.add(8) # [ticket, done] for the dynamic scheduler + self.num_sm = multiprocessor_count(current_device_id()) + self.n_tiles = B * HO + self.n_heads_out = HO + if self.split: + self.ideal = compute_ideal_chunks(total, HO, self.num_sm, self.b_t) + self.work_item_rows = max_work_items(total, B, HO, self.ideal, self.b_t, self.num_sm) + else: + self.ideal = None + self.work_item_rows = self.n_tiles + self.off_work_items = layout.add(self.work_item_rows * WORK_ITEM_FIELDS * 4) + self.off_work_count = layout.add(4) + if self.split: + 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) + from .common.host import tensormap_workspace_bytes + + self.tensormap_bytes = tensormap_workspace_bytes(kernel_mod, B) + self.off_tensormaps = layout.add(self.tensormap_bytes, align=128) + self.ws_bytes = layout.size + regions = [ + (self.off_sched, "int32", (2,)), + (self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), + (self.off_work_count, "int32", (1,)), + ] + if self.split: + regions += [ + (self.off_item_scratch, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), + (self.off_chunk_scratch, "float32", (self.chunk_scratch_rows, self.n_heads_out)), + ] + regions.append((self.off_tensormaps, "int64", (self.tensormap_bytes // 8,))) + self.carve = carve_plan("Gdn2FrostEngine (GDN2)", regions) def workspace_bytes(self) -> int: - return self._ws_bytes + return self.ws_bytes def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: - nb = node_buffers[self._node] + nb = node_buffers[self.node] q = nb.inputs["q"] k = nb.inputs["k"] v = nb.inputs["v"] @@ -168,33 +165,45 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: beta = nb.inputs["beta"] w = nb.inputs["w"] cu = nb.inputs["cu_seqlens"] - s0 = nb.inputs.get("initial_state") + state0 = nb.inputs.get("initial_state") o = nb.outputs["O"] - fs = nb.outputs["final_state"] if self._has_fs else None - + final_state = nb.outputs["final_state"] if self.has_final_state else None + state_checkpoints = nb.outputs["state_checkpoints"] if self.has_state_checkpoints else None + a_log = nb.inputs.get("a_log") + dt_bias = nb.inputs.get("dt_bias") stream = stream if stream is not None else 0 - ws = workspace - sched_ctr, work_items, item_scratch, work_count, chunk_scratch, tensormaps = ws.carve(self._carve) - from .common.split_k import build_split_table - - build_split_table( + if self.split: + sched_ctr, work_items, work_count, item_scratch, chunk_scratch, tensormaps = workspace.carve(self.carve) + else: + sched_ctr, work_items, work_count, tensormaps = workspace.carve(self.carve) + item_scratch = chunk_scratch = None + self.build_split_table( g, cu, work_items, work_count, - ideal_chunks=self._ideal, - n_tiles=self._n_tiles, - num_sms=self._num_sm, - b_t=self._b_t, + ideal_chunks=self.ideal, + n_tiles=self.n_tiles, + num_sms=self.num_sm, + b_t=self.b_t, chunk_scratch=chunk_scratch, item_scratch=item_scratch, log_gate=True, + safe_gate=self.safe_gate, + a_log=a_log, + dt_bias=dt_bias, + gate_lower_bound=self.gate_lower_bound if self.safe_gate else None, sched_ctr=sched_ctr, + split=self.split, stream=stream, ) - self._kernel.chunk_gdn2_sm100( + ckpt_kwargs = {} + if self.has_state_checkpoints: + # the kernel derives the per-sequence checkpoint entry offsets on device + ckpt_kwargs = dict(checkpoint_every_n_tokens=self.ckpt, output_state_checkpoints=state_checkpoints) + self.kernel.chunk_gdn2_sm100( q, k, v, @@ -203,15 +212,239 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: w, o, cu, - s0, - fs, - self._scale, - use_qk_l2norm_in_kernel=self._use_qk_l2norm, - use_beta_w_sigmoid_in_kernel=self._use_beta_w_sigmoid, + state0, + final_state, + self.scale, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, + safe_gate=self.safe_gate, + gate_lower_bound=self.gate_lower_bound, + a_log=a_log, + dt_bias=dt_bias, work_items=work_items, work_count=work_count, sched_ctr=sched_ctr, tensormap_workspace=tensormaps, + **ckpt_kwargs, + stream=stream, + ) + return None + + +class CompiledGdn2Bwd: + """Compiled FROST GDN-2 backward plan: the workspace holds the + regenerated per-chunk checkpoint series when the graph carries no ``state_checkpoints`` input, + plus GVA/GQA head scratch for dQ/dK/dV.""" + + def __init__(self, node, bwd_mod, regen_mod): + from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items + + self.node = node + self.bwd = bwd_mod + self.regen = regen_mod + self.build_split_table = build_split_table + self.plan_name = "Gdn2FrostEngine (GDN2_BWD)" + from .common.downcast import downcast_state + from .common.host import tensormap_workspace_bytes + + self.downcast_state = downcast_state + scale = node.params.get("scale") + self.scale = float(scale) if scale is not None else 1.0 / math.sqrt(node.inputs["q"].dim[-1]) + self.use_qk_l2norm = bool(node.params.get("use_qk_l2norm", False)) + 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"] + self.b_t = bwd_mod.CFG.B_T + total = q.dim[0] + HQ, HV = q.dim[1], v.dim[1] + HO = g.dim[1] + K, V = q.dim[-1], v.dim[-1] + B = node.inputs["cu_seqlens"].dim[0] - 1 + self.io_name = "float16" if node.inputs["q"].get_data_type().name == "HALF" else "bfloat16" + self.n_heads_out, self.total = HO, total + layout = WorkspaceLayout() + self.off_sched = layout.add(16) # one [ticket, done] ring each for the regen and bwd kernels + self.num_sm = multiprocessor_count(current_device_id()) + self.bwd_dyn_sched = B * HO <= self.num_sm + self.batch_invariant = bool(node.params.get("batch_invariant", False)) + # cuts never in batch-invariant mode: whole-sequence items keep each + # sequence's math independent of the batch composition + self.split = not self.batch_invariant + self.n_tiles = B * HO + if self.split: + self.ideal = compute_ideal_chunks(total, HO, self.num_sm, self.b_t) + self.work_item_rows = max_work_items(total, B, HO, self.ideal, self.b_t, self.num_sm) + else: + self.ideal = None + self.work_item_rows = self.n_tiles + self.off_work_items = layout.add(self.work_item_rows * WORK_ITEM_FIELDS * 4) + self.off_work_count = layout.add(4) + if self.split: + 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) + # chunk-0 entering state, io dtype (downcast initial_state or zeros) + 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) + self.regen_tm_bytes = tensormap_workspace_bytes(regen_mod, B) + self.off_regen_tensormaps = layout.add(self.regen_tm_bytes, align=128) + HK = node.inputs["k"].dim[1] + self.fold_dq = HQ < HO + self.fold_dk = HK < HO + self.fold_dv = HV < HO + if self.fold_dq: + self.off_dq_ho = layout.add(total * HO * K * 2) + if self.fold_dk: + self.off_dk_ho = layout.add(total * HO * K * 2) + if self.fold_dv: + self.off_dv_ho = layout.add(total * HO * V * 2) + self.bwd_tm_bytes = tensormap_workspace_bytes(bwd_mod, B) + self.off_bwd_tensormaps = layout.add(self.bwd_tm_bytes, align=128) + self.ws_bytes = layout.size + regions = [ + ("sched_regen", self.off_sched, "int32", (2,)), + ("sched_bwd", self.off_sched + 8, "int32", (2,)), + ("sched_all", self.off_sched, "int32", (4,)), + ("work_items", self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), + ("work_count", self.off_work_count, "int32", (1,)), + ("bwd_tensormaps", self.off_bwd_tensormaps, "int64", (self.bwd_tm_bytes // 8,)), + ] + 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(("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))) + if self.fold_dk: + regions.append(("dk_ho", self.off_dk_ho, self.io_name, (total, HO, K))) + if self.fold_dv: + regions.append(("dv_ho", self.off_dv_ho, self.io_name, (total, HO, V))) + self.carve_names = [name for name, _off, _dt, _shape in regions] + self.carve = carve_plan("Gdn2FrostEngine (GDN2_BWD)", [(off, dt, shape) for _name, off, dt, shape in regions]) + + def workspace_bytes(self) -> int: + return self.ws_bytes + + def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: + nb = node_buffers[self.node] + q = nb.inputs["q"] + k = nb.inputs["k"] + v = nb.inputs["v"] + g = nb.inputs["g"] + beta = nb.inputs["beta"] + w = nb.inputs["w"] + cu = nb.inputs["cu_seqlens"] + do = nb.inputs["dO"] + state_checkpoints = nb.inputs.get("state_checkpoints") + state0 = nb.inputs.get("initial_state") + dstate_in = nb.inputs.get("d_final_state") + dq = nb.outputs["dQ"] + dk = nb.outputs["dK"] + dv = nb.outputs["dV"] + dg = nb.outputs["dG"] + db = nb.outputs["dBeta"] + dw = nb.outputs["dW"] + dstate0 = nb.outputs.get("d_initial_state") + stream = stream if stream is not None else 0 + + HO, total = self.n_heads_out, self.total + K, V = q.shape[-1], v.shape[-1] + B = cu.shape[0] - 1 + region = dict(zip(self.carve_names, workspace.carve(self.carve))) + sched_regen = region["sched_regen"] + sched_bwd = region["sched_bwd"] + work_items = region["work_items"] + work_count = region["work_count"] + self.build_split_table( + g, + cu, + work_items, + work_count, + ideal_chunks=self.ideal, + n_tiles=self.n_tiles, + num_sms=self.num_sm, + b_t=self.b_t, + chunk_scratch=region.get("chunk_scratch"), + item_scratch=region.get("item_scratch"), + log_gate=True, + sched_ctr=region["sched_all"], + split=self.split, + 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: + checkpoint_series = region["state_checkpoints"] + self.regen.chunk_gdn2_recompute_sm100( + k, + v, + g, + beta, + w, + cu, + state0, + None, + checkpoint_every_n_tokens=self.b_t, + output_state_checkpoints=checkpoint_series, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, + work_items=work_items, + work_count=work_count, + sched_ctr=sched_regen, + tensormap_workspace=region["regen_tensormaps"], + stream=stream, + ) + + dq_out, dk_out, dv_out = dq, dk, dv + if self.fold_dq: + dq_out = region["dq_ho"] + if self.fold_dk: + dk_out = region["dk_ho"] + if self.fold_dv: + dv_out = region["dv_ho"] + + self.bwd.chunk_gdn2_bwd_sm100( + q, + k, + v, + g, + beta, + w, + do, + checkpoint_series, + dq_out, + dk_out, + dv_out, + dg, + db, + dw, + cu, + self.scale, + initial_state=state0_io, + d_initial_state=dstate0 if self.has_dstate0 else None, + d_final_state=dstate_in, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, + work_items=work_items, + work_count=work_count, + sched_ctr=sched_bwd if self.bwd_dyn_sched else None, + tensormap_workspace=region["bwd_tensormaps"], stream=stream, ) + if dq_out is not dq or dk_out is not dk or dv_out is not dv: + from .common.head_reduce import head_group_reduce + + for src_ho, dst in ((dq_out, dq), (dk_out, dk), (dv_out, dv)): + if src_ho is not dst: + head_group_reduce(src_ho, dst, stream=stream) return None diff --git a/python/cudnn/linear_attention/frost/gdn_engine.py b/python/cudnn/linear_attention/frost/gdn_engine.py index b5fd88d3b..d908011f2 100644 --- a/python/cudnn/linear_attention/frost/gdn_engine.py +++ b/python/cudnn/linear_attention/frost/gdn_engine.py @@ -1,9 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""FROST GDN engine: GDN / GDN_BWD nodes on the chunked prefill kernels -(``kernel/gdn_prefill_f16.py`` / ``kernel/gdn_bprop_f16.py``, Blackwell -SM100/SM103, bf16/fp16).""" +"""FROST GDN engine: GDN nodes on the chunked prefill kernel +(``kernel/gdn_prefill_f16.py``) and GDN_BWD nodes on the chunked backward +kernel (``kernel/gdn_bprop_f16.py``), Blackwell SM100/SM103, bf16/fp16. +The backward regenerates the per-chunk state checkpoints with the recompute kernel +(``kernel/gdn_recompute_f16.py``) when the graph does not provide one.""" from __future__ import annotations @@ -14,85 +16,25 @@ from cudnn.engines.base import BaseEngine, CompiledPlan from cudnn.frost import buffers -from cudnn.frost.workspace import Workspace, WorkspaceLayout, carve_plan -from ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair - - -def _the_gdn_node(graph): - nodes = list(graph.nodes) - if len(nodes) != 1: - return None - node = nodes[0] - if getattr(node.node_type, "name", None) not in ("GDN", "GDN_BWD"): - return None - return node - - -def _io_dtype_name(node) -> str: - """The io dtype's buffer-level name ('bfloat16' / 'float16').""" - import cudnn - - return "bfloat16" if node.inputs["q"].get_data_type() == cudnn.data_type.BFLOAT16 else "float16" - - -def _check_common(node) -> None: - """Shape/dtype gates shared by the fwd and bwd kernels.""" - import cudnn - - q, k, v = (node.inputs[p] for p in ("q", "k", "v")) - io_dtypes = {q.get_data_type(), k.get_data_type(), v.get_data_type()} - {None} - if len(io_dtypes) > 1: - raise NotImplementedError(f"GdnFrostEngine: q/k/v dtypes must match, got {io_dtypes}") - for p, t in (("q", q), ("k", k), ("v", v)): - if t.get_data_type() not in (cudnn.data_type.BFLOAT16, cudnn.data_type.HALF, None): - raise NotImplementedError(f"GdnFrostEngine: '{p}' must be bf16 or fp16, got {t.get_data_type()}") - if not t.dim or len(t.dim) != 3: - raise NotImplementedError(f"GdnFrostEngine: '{p}' must be THD [total_T, heads, dim]") - if q.dim[-1] != 128 or v.dim[-1] != 128: - raise NotImplementedError(f"GdnFrostEngine: head dims must be 128 (the recurrent state is 128x128), got K={q.dim[-1]} V={v.dim[-1]}") - hq, hk, hv = q.dim[1], k.dim[1], v.dim[1] - if hq != hk: - raise NotImplementedError(f"GdnFrostEngine: q and k head counts differ ({hq} vs {hk})") - if hv != hq and max(hq, hv) % min(hq, hv) != 0: - # GVA (v-heads grouped over q-heads) or GQA (q-heads grouped over - # v-heads, the kernel's opposite direction) - raise NotImplementedError(f"GdnFrostEngine: q heads ({hq}) and v heads ({hv}) must be equal or one a multiple of the other") - ho = hq if hq >= hv else hv - for p in ("g", "beta"): - t = node.inputs[p] - if t.dim and t.dim[1] != ho: - raise NotImplementedError(f"GdnFrostEngine: '{p}' must carry HO = max(q, v) heads ({ho}), got {t.dim[1]}") - - # kernel-native operand dtypes: buffers pass through without staging - fp32 = cudnn.data_type.FLOAT - io = (cudnn.data_type.BFLOAT16, cudnn.data_type.HALF) - state = (fp32, cudnn.data_type.BFLOAT16) - _require_dtype("GdnFrostEngine", node, "g", fp32) - _require_dtype("GdnFrostEngine", node, "beta", fp32) - _require_dtype("GdnFrostEngine", node, "cu_seqlens", cudnn.data_type.INT32) - _require_dtype("GdnFrostEngine", node, "initial_state", state) - _require_dtype("GdnFrostEngine", node, "final_state", state, out=True) - _require_state_pair("GdnFrostEngine", node) - if node.node_type.name == "GDN_BWD": - _require_dtype("GdnFrostEngine", node, "dO", io) - _require_dtype("GdnFrostEngine", node, "h", io) - _require_dtype("GdnFrostEngine", node, "d_final_state", fp32) - _require_dtype("GdnFrostEngine", node, "d_initial_state", fp32, out=True) - _require_dtype("GdnFrostEngine", node, "dG", fp32, out=True) - _require_dtype("GdnFrostEngine", node, "dBeta", fp32, out=True) +from cudnn.frost.buffers import current_device_id +from cudnn.frost.device import multiprocessor_count +from cudnn.frost.workspace import WorkspaceLayout, carve_plan +from ..graph_analyzer import FrostLaPlan, frost_la_gate, require, analyze def build_gdn(graph): """The expensive step: import the kernel module (pulls in the Cutlass primitives; the cute.compile itself is cached inside the kernel per static config and runs on first execute, when the real buffers are known).""" - node = _the_gdn_node(graph) - if node is None: + nodes = list(graph.nodes) + if len(nodes) != 1 or getattr(nodes[0].node_type, "name", None) not in ("GDN", "GDN_BWD"): raise ValueError("build_gdn: graph does not contain exactly one GDN/GDN_BWD node") + node = nodes[0] if node.node_type.name == "GDN_BWD": - from .kernel import gdn_bprop_f16 as kernel_mod + from .kernel import gdn_bprop_f16 as bwd_mod + from .kernel import gdn_recompute_f16 as regen_mod - return CompiledGdnBwd(node, kernel_mod) + return CompiledGdnBwd(node, bwd_mod, regen_mod) from .kernel import gdn_prefill_f16 as kernel_mod return CompiledGdn(node, kernel_mod) @@ -102,156 +44,147 @@ class GdnFrostEngine(BaseEngine): """FROST chunked-kernel backend for single-node GDN graphs (THD layout). Default GDN engine on SM100/SM103 (lowest GDN engine_id); declines - elsewhere so the router falls back to ``GdnCuTileEngine``.""" + elsewhere so ranking falls back to ``GdnCuTileEngine``.""" name = "gdn_frost" behavior_notes = (behavior_note.RUNTIME_COMPILATION,) # JIT-compiled at build_plans() def check_support(self, graph) -> None: - node = _the_gdn_node(graph) - if node is None: - raise NotImplementedError("GdnFrostEngine supports exactly one GDN/GDN_BWD node") - sm = buffers.current_sm() - if sm is None or not (100 <= sm <= 103): - raise NotImplementedError(f"GdnFrostEngine requires SM100-SM103 (found {sm})") - try: - import cutlass.experimental.primitives # noqa: F401 — availability probe: ImportError = decline - except ImportError as exc: - raise NotImplementedError(f"GdnFrostEngine requires the Cutlass DSL with cutlass.experimental.primitives: {exc}") from exc - if node.params.get("use_qk_l2norm", False): + import cudnn + + facts = graph._facts_for(analyze) + frost_la_gate("GdnFrostEngine", facts, "GDN") + if facts.use_qk_l2norm: raise NotImplementedError("GdnFrostEngine: use_qk_l2norm is not supported (the kernel takes q/k as given)") - ports = ("q", "k", "v", "g", "beta", "cu_seqlens") - if node.node_type.name == "GDN_BWD": - ports += ("dO",) - for port in ports: - if port not in node.inputs: - raise NotImplementedError(f"GdnFrostEngine: {node.node_type.name} node '{node.name}' is missing input '{port}'") - ckpt = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) - if ckpt and (node.node_type.name != "GDN" or ckpt % 64 != 0): + if facts.safe_gate: + raise NotImplementedError("GdnFrostEngine: safe_gate is not supported (no gate-activation path; the cuTile GDN engine serves it)") + ckpt = facts.checkpoint_every_n_tokens + if ckpt and (facts.is_bwd or ckpt % 64 != 0): raise NotImplementedError(f"GdnFrostEngine: checkpoint_every_n_tokens must be a positive multiple of 64 on the GDN node (got {ckpt})") - _check_common(node) + if not facts.gates_at_ho: + raise NotImplementedError(f"GdnFrostEngine: g/beta must carry HO = max(q, v) heads ({facts.h_o})") + fp32 = cudnn.data_type.FLOAT + io = (cudnn.data_type.BFLOAT16, cudnn.data_type.HALF) + state_dtypes = (fp32, cudnn.data_type.BFLOAT16) + require("GdnFrostEngine", "beta", facts.beta_dtype, fp32) + require("GdnFrostEngine", "initial_state", facts.state_dtype, state_dtypes) + require("GdnFrostEngine", "final_state", facts.final_state_dtype, state_dtypes) + if not facts.state_pair_match: + raise NotImplementedError("GdnFrostEngine: initial_state and final_state dtypes must match") + if facts.is_bwd: + require("GdnFrostEngine", "dO", facts.do_dtype, io) + if facts.io_dtype is not None: + require("GdnFrostEngine", "state_checkpoints", facts.state_checkpoints_dtype, facts.io_dtype) + require("GdnFrostEngine", "d_final_state", facts.d_final_state_dtype, fp32) + require("GdnFrostEngine", "d_initial_state", facts.d_initial_state_dtype, fp32) + require("GdnFrostEngine", "dG", facts.dg_dtype, fp32) + require("GdnFrostEngine", "dBeta", facts.dbeta_dtype, fp32) + elif facts.io_dtype is not None: + require("GdnFrostEngine", "state_checkpoints", facts.state_checkpoints_out_dtype, facts.io_dtype) def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: - return _FrostPlan(build_gdn(graph)) - - -def _splits_for(gate, cu_i32, work_items, item_scratch, chunk_scratch, work_count, ideal_chunks, n_tiles, num_sms, b_t, stream, log_gate=False, sched_ctr=None): - """Launch the split pipeline (stream-ordered; the scan zeroes the count - and the scheduler ticket ring).""" - from .common.split_k import build_split_table - - build_split_table( - gate, - cu_i32, - work_items, - work_count, - ideal_chunks=ideal_chunks, - n_tiles=n_tiles, - num_sms=num_sms, - b_t=b_t, - chunk_scratch=chunk_scratch, - item_scratch=item_scratch, - log_gate=log_gate, - sched_ctr=sched_ctr, - stream=stream, - ) - return work_items, work_count + return FrostLaPlan(build_gdn(graph)) class CompiledGdn: """Compiled FROST GDN plan: a callable over the resolved node buffers.""" def __init__(self, node, kernel_mod): - from .common.split_k import WORK_ITEM_FIELDS, chunk_scratch_rows, compute_ideal_chunks, max_work_items + from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items - self._node = node - self._kernel = kernel_mod + self.node = node + self.kernel = kernel_mod + self.build_split_table = build_split_table + self.plan_name = "GdnFrostEngine (GDN)" scale = node.params.get("scale") - self._scale = float(scale) if scale is not None else 1.0 / math.sqrt(node.inputs["q"].dim[-1]) + self.scale = float(scale) if scale is not None else 1.0 / math.sqrt(node.inputs["q"].dim[-1]) q, v, g = node.inputs["q"], node.inputs["v"], node.inputs["g"] - self._b_t = kernel_mod.CFG.B_T - total, K, V = q.dim[0], q.dim[2], v.dim[2] + self.b_t = kernel_mod.CFG.B_T + total = q.dim[0] HO = g.dim[1] B = node.inputs["cu_seqlens"].dim[0] - 1 - self._has_fs = "final_state" in node.outputs - self._ckpt = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) - self._has_h = "H" in node.outputs - # split work items assume chunk-granular state boundaries; other - # checkpoint cadences keep the serial per-(b,h) walk - self._split = self._ckpt in (0, self._b_t) + self.has_final_state = "final_state" in node.outputs + 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 layout = WorkspaceLayout() - self._off_tensormaps = layout.add(kernel_mod.get_workspace_size(B, q.dim[1], v.dim[1])) - self._tensormap_words = kernel_mod.get_workspace_size(B, q.dim[1], v.dim[1]) // 8 - self._off_sched = layout.add(8) # [ticket, done] for the dynamic scheduler - if self._split: - self._num_sm = kernel_mod._device_sm_count() - self._ideal = compute_ideal_chunks(total, HO, self._num_sm, self._b_t) - self._n_tiles = B * HO - self._n_heads_out = HO - self._work_item_rows = max_work_items(total, B, HO, self._ideal, self._b_t, self._num_sm) - self._off_work_items = layout.add(self._work_item_rows * WORK_ITEM_FIELDS * 4) - self._off_item_scratch = layout.add(self._work_item_rows * WORK_ITEM_FIELDS * 4) - self._off_work_count = layout.add(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._ws_bytes = layout.size - + from .common.host import tensormap_workspace_bytes + + self.tensormap_words = tensormap_workspace_bytes(kernel_mod, B) // 8 + self.off_tensormaps = layout.add(self.tensormap_words * 8) + self.off_sched = layout.add(8) # [ticket, done] for the dynamic scheduler + self.num_sm = multiprocessor_count(current_device_id()) + self.n_tiles = B * HO + self.n_heads_out = HO + if self.split: + self.ideal = compute_ideal_chunks(total, HO, self.num_sm, self.b_t) + self.work_item_rows = max_work_items(total, B, HO, self.ideal, self.b_t, self.num_sm) + else: + self.ideal = None + self.work_item_rows = self.n_tiles + self.off_work_items = layout.add(self.work_item_rows * WORK_ITEM_FIELDS * 4) + self.off_work_count = layout.add(4) + if self.split: + 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.ws_bytes = layout.size regions = [ - (self._off_tensormaps, "int64", (self._tensormap_words,)), - (self._off_sched, "int32", (2,)), + (self.off_tensormaps, "int64", (self.tensormap_words,)), + (self.off_sched, "int32", (2,)), + (self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), + (self.off_work_count, "int32", (1,)), ] - if self._split: + if self.split: regions += [ - (self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), - (self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), - (self._off_work_count, "int32", (1,)), - (self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, HO)), + (self.off_item_scratch, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), + (self.off_chunk_scratch, "float32", (self.chunk_scratch_rows, self.n_heads_out)), ] - self._carve = carve_plan("gdn", regions) + self.carve = carve_plan("GdnFrostEngine (GDN)", regions) def workspace_bytes(self) -> int: - return self._ws_bytes + return self.ws_bytes def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: - nb = node_buffers[self._node] + nb = node_buffers[self.node] q = nb.inputs["q"] k = nb.inputs["k"] v = nb.inputs["v"] g = nb.inputs["g"] beta = nb.inputs["beta"] cu = nb.inputs["cu_seqlens"] - s0 = nb.inputs.get("initial_state") + state0 = nb.inputs.get("initial_state") o = nb.outputs["O"] - fs = nb.outputs["final_state"] if self._has_fs else None - h = nb.outputs["H"] if self._has_h else None - - ws = workspace + final_state = nb.outputs["final_state"] if self.has_final_state else None + state_checkpoints = nb.outputs["state_checkpoints"] if self.has_state_checkpoints else None stream = stream if stream is not None else 0 - work_items = work_count = None - if self._split: - tensormaps, sched_ctr, work_items, item_scratch, work_count, chunk_scratch = ws.carve(self._carve) - _splits_for( - g, - cu, - work_items, - item_scratch, - chunk_scratch, - work_count, - self._ideal, - self._n_tiles, - self._num_sm, - self._b_t, - stream, - log_gate=True, - sched_ctr=sched_ctr, - ) + if self.split: + tensormaps, sched_ctr, work_items, work_count, item_scratch, chunk_scratch = workspace.carve(self.carve) else: - tensormaps, sched_ctr = ws.carve(self._carve) - buffers.memset_zero_async(sched_ctr.data_ptr(), sched_ctr.nbytes, stream) + tensormaps, sched_ctr, work_items, work_count = workspace.carve(self.carve) + item_scratch = chunk_scratch = None + self.build_split_table( + g, + cu, + work_items, + work_count, + ideal_chunks=self.ideal, + n_tiles=self.n_tiles, + num_sms=self.num_sm, + b_t=self.b_t, + chunk_scratch=chunk_scratch, + item_scratch=item_scratch, + log_gate=True, + sched_ctr=sched_ctr, + split=self.split, + stream=stream, + ) - self._kernel.chunk_gdn_sm100( + self.kernel.chunk_gdn_sm100( q, k, v, @@ -259,14 +192,14 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: beta, o, cu, - s0, - fs, - self._scale, + state0, + final_state, + self.scale, work_items=work_items, work_count=work_count, sched_ctr=sched_ctr, - checkpoint_every_n_tokens=self._ckpt, - output_h=h, + checkpoint_every_n_tokens=self.ckpt, + output_state_checkpoints=state_checkpoints, log_gate=True, workspace=tensormaps, stream=stream, @@ -278,93 +211,107 @@ class CompiledGdnBwd: """Compiled FROST GDN bprop plan: a callable over the resolved node buffers. Produces dQ/dK/dV/dG/dBeta; consumes the forward per-chunk states through - the node's ``h`` input, or regenerates them with a forward H-dump run when - the port is absent.""" + the node's ``state_checkpoints`` input, or regenerates them with the recompute + (checkpoint-only) kernel when the port is absent.""" - def __init__(self, node, kernel_mod): - from .common.split_k import WORK_ITEM_FIELDS, chunk_scratch_rows, compute_ideal_chunks, max_work_items + def __init__(self, node, bwd_mod, regen_mod): + from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items - self._node = node - self._kernel = kernel_mod - scale = node.params.get("scale") - self._scale = float(scale) if scale is not None else 1.0 / math.sqrt(node.inputs["q"].dim[-1]) + self.node = node + self.bwd = bwd_mod + self.regen = regen_mod + self.build_split_table = build_split_table + self.plan_name = "GdnFrostEngine (GDN_BWD)" + from .common.downcast import downcast_state + from .common.host import tensormap_workspace_bytes - from .kernel import gdn_prefill_f16 as fwd_mod + self.downcast_state = downcast_state + scale = node.params.get("scale") + self.scale = float(scale) if scale is not None else 1.0 / math.sqrt(node.inputs["q"].dim[-1]) - self._fwd = fwd_mod q, v, g = node.inputs["q"], node.inputs["v"], node.inputs["g"] - self._b_t = kernel_mod.CFG.B_T - total, K, V = q.dim[0], q.dim[2], v.dim[2] + self.b_t = bwd_mod.CFG.B_T + total = q.dim[0] + K, V = q.dim[-1], v.dim[-1] HQ, HV = q.dim[1], v.dim[1] HO = g.dim[1] B = node.inputs["cu_seqlens"].dim[0] - 1 - self._has_h = "h" in node.inputs - self._has_s0 = "initial_state" in node.inputs - self._has_dht = "d_final_state" in node.inputs - self._has_ds0 = "d_initial_state" in node.outputs - self._io_name = _io_dtype_name(node) - - # multi-wave grids pay ~2-3% for the ticket ring in the BACKWARD - # (its TMA-LDG publisher drives five load streams); static stride there - self._bwd_dyn_sched = B * HO <= kernel_mod._device_sm_count() + 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_id()) + self.bwd_dyn_sched = B * HO <= self.num_sm + self.batch_invariant = bool(node.params.get("batch_invariant", False)) + # cuts never in batch-invariant mode: whole-sequence items keep each + # sequence's math independent of the batch composition + self.split = not self.batch_invariant layout = WorkspaceLayout() - self._off_sched = layout.add(16) # one [ticket, done] ring each for the regen and bwd kernels - self._tensormap_words = kernel_mod.get_workspace_size(B, HQ, HV) // 8 - self._off_tensormaps = layout.add(self._tensormap_words * 8) - # regenerated H series (io dtype [n, HO, K, V]); the entry count is - # bounded statically so no device sync sizes it - self._h_rows = 0 if self._has_h else max(total // self._b_t + B, 1) - self._off_h = layout.add(self._h_rows * HO * K * V * 2) if self._h_rows else None - # io-downcast initial state (chunk 0's S reads it via its own - # descriptor set) - self._off_s0io = layout.add(B * HO * K * V * 2) if self._has_s0 else None - if not self._has_h: - self._fwd_tensormap_words = fwd_mod.get_workspace_size(B, HO, HO) // 8 - self._off_fwd_tensormaps = layout.add(self._fwd_tensormap_words * 8) - self._is_gva = HV > HQ - self._is_gqa = HQ > HV - if self._is_gva: - self._off_dq_ho = layout.add(total * HV * K * 2) - self._off_dk_ho = layout.add(total * HV * K * 2) - if self._is_gqa: - self._off_dv_ho = layout.add(total * HQ * V * 2) - self._num_sm = kernel_mod._device_sm_count() - self._ideal = compute_ideal_chunks(total, HO, self._num_sm, self._b_t) - self._n_tiles = B * HO - self._work_item_rows = max_work_items(total, B, HO, self._ideal, self._b_t, self._num_sm) - self._off_work_items = layout.add(self._work_item_rows * WORK_ITEM_FIELDS * 4) - self._off_item_scratch = layout.add(self._work_item_rows * WORK_ITEM_FIELDS * 4) - self._off_work_count = layout.add(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._shapes = (total, HQ, HV, HO, K, V, B) - self._ws_bytes = layout.size - - # The regions every backward carves. The rest (io state, regenerated H, - # head-group scratch) hang off build-time branches that are usually - # off, and stay on view() rather than turning this into index - # bookkeeping. The three scheduler views overlap on purpose: one ring - # each for the regen and bwd kernels, and both at once for the split - # pipeline that zeroes them. - self._carve = carve_plan( - "gdn_bwd", - [ - (self._off_sched, "int32", (2,)), - (self._off_sched + 8, "int32", (2,)), - (self._off_sched, "int32", (4,)), - (self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), - (self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), - (self._off_work_count, "int32", (1,)), - (self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, HO)), - (self._off_tensormaps, "int64", (self._tensormap_words,)), - ], - ) + self.off_sched = layout.add(16) # one [ticket, done] ring each for the regen and bwd kernels + self.tensormap_words = tensormap_workspace_bytes(bwd_mod, B) // 8 + self.off_tensormaps = layout.add(self.tensormap_words * 8) + # chunk-0 entering state, io dtype (downcast initial_state; absent = in-kernel zeros) + 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) + self.regen_tensormap_words = tensormap_workspace_bytes(regen_mod, B) // 8 + self.off_regen_tensormaps = layout.add(self.regen_tensormap_words * 8) + HK = node.inputs["k"].dim[1] + self.fold_dq = HQ < HO + self.fold_dk = HK < HO + self.fold_dv = HV < HO + if self.fold_dq: + self.off_dq_ho = layout.add(total * HO * K * 2) + if self.fold_dk: + self.off_dk_ho = layout.add(total * HO * K * 2) + if self.fold_dv: + self.off_dv_ho = layout.add(total * HO * V * 2) + self.n_tiles = B * HO + if self.split: + self.ideal = compute_ideal_chunks(total, HO, self.num_sm, self.b_t) + self.work_item_rows = max_work_items(total, B, HO, self.ideal, self.b_t, self.num_sm) + else: + self.ideal = None + self.work_item_rows = self.n_tiles + self.off_work_items = layout.add(self.work_item_rows * WORK_ITEM_FIELDS * 4) + self.off_work_count = layout.add(4) + if self.split: + 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.n_heads_out, self.total = HO, total + self.ws_bytes = layout.size + regions = [ + ("sched_regen", self.off_sched, "int32", (2,)), + ("sched_bwd", self.off_sched + 8, "int32", (2,)), + ("sched_all", self.off_sched, "int32", (4,)), + ("tensormaps", self.off_tensormaps, "int64", (self.tensormap_words,)), + ("work_items", self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), + ("work_count", self.off_work_count, "int32", (1,)), + ] + 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(("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))) + if self.fold_dk: + regions.append(("dk_ho", self.off_dk_ho, self.io_name, (total, HO, K))) + if self.fold_dv: + regions.append(("dv_ho", self.off_dv_ho, self.io_name, (total, HO, V))) + self.carve_names = [name for name, _off, _dt, _shape in regions] + self.carve = carve_plan("GdnFrostEngine (GDN_BWD)", [(off, dt, shape) for _name, off, dt, shape in regions]) def workspace_bytes(self) -> int: - return self._ws_bytes + return self.ws_bytes - def __call__(self, node_buffers, *, workspace=None, stream=None): - nb = node_buffers[self._node] + def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: + nb = node_buffers[self.node] q = nb.inputs["q"] k = nb.inputs["k"] v = nb.inputs["v"] @@ -372,103 +319,108 @@ def __call__(self, node_buffers, *, workspace=None, stream=None): beta = nb.inputs["beta"] cu = nb.inputs["cu_seqlens"] do = nb.inputs["dO"] - h_in = nb.inputs.get("h") if self._has_h else None - s0 = nb.inputs.get("initial_state") if self._has_s0 else None - dht = nb.inputs.get("d_final_state") if self._has_dht else None - ds0 = nb.outputs.get("d_initial_state") + state_checkpoints = nb.inputs.get("state_checkpoints") + state0 = nb.inputs.get("initial_state") + dstate_in = nb.inputs.get("d_final_state") dq = nb.outputs["dQ"] dk = nb.outputs["dK"] dv = nb.outputs["dV"] dg = nb.outputs["dG"] - dbeta = nb.outputs["dBeta"] - ws = workspace - total, HQ, HV, HO, K, V, B = self._shapes + db = nb.outputs["dBeta"] + dstate0 = nb.outputs.get("d_initial_state") + HO, total = self.n_heads_out, self.total + K, V = q.shape[-1], v.shape[-1] + B = cu.shape[0] - 1 stream = stream if stream is not None else 0 - sched_fwd, sched_bwd, sched_all, work_items, item_scratch, work_count, chunk_scratch, tensormaps = ws.carve(self._carve) - _splits_for( + region = dict(zip(self.carve_names, workspace.carve(self.carve))) + sched_regen = region["sched_regen"] + sched_bwd = region["sched_bwd"] + work_items = region["work_items"] + work_count = region["work_count"] + self.build_split_table( g, cu, work_items, - item_scratch, - chunk_scratch, work_count, - self._ideal, - self._n_tiles, - self._num_sm, - self._b_t, - stream, + ideal_chunks=self.ideal, + n_tiles=self.n_tiles, + num_sms=self.num_sm, + b_t=self.b_t, + chunk_scratch=region.get("chunk_scratch"), + item_scratch=region.get("item_scratch"), log_gate=True, - sched_ctr=sched_all, + sched_ctr=region["sched_all"], + split=self.split, + stream=stream, ) - s0_io = None - if s0 is not None: - s0_io = ws.view(self._off_s0io, self._io_name, (B, HO, K, V)) - self._fwd.downcast_state(s0, s0_io, stream=stream) - if h_in is not None: - h = h_in + 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: + # already the staged form: io dtype, compact, aligned — bind directly + 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: - h = ws.view(self._off_h, self._io_name, (self._h_rows, HO, K, V)) - self._fwd.chunk_gdn_sm100( - q, + checkpoint_series = region["state_checkpoints"] + self.regen.chunk_gdn_recompute_sm100( k, v, g, beta, - None, cu, - s0, + state0, None, - self._scale, - checkpoint_every_n_tokens=self._b_t, - output_h=h, + checkpoint_every_n_tokens=self.b_t, + output_state_checkpoints=checkpoint_series, work_items=work_items, work_count=work_count, - sched_ctr=sched_fwd, + sched_ctr=sched_regen, log_gate=True, - workspace=ws.view(self._off_fwd_tensormaps, "int64", (self._fwd_tensormap_words,)), + workspace=region["regen_tensormaps"], stream=stream, ) dq_out, dk_out, dv_out = dq, dk, dv - if self._is_gva: - dq_out = ws.view(self._off_dq_ho, self._io_name, (total, HV, K)) - dk_out = ws.view(self._off_dk_ho, self._io_name, (total, HV, K)) - if self._is_gqa: - dv_out = ws.view(self._off_dv_ho, self._io_name, (total, HQ, V)) - self._kernel.chunk_gdn_bwd_sm100( + if self.fold_dq: + dq_out = region["dq_ho"] + if self.fold_dk: + dk_out = region["dk_ho"] + if self.fold_dv: + dv_out = region["dv_ho"] + self.bwd.chunk_gdn_bwd_sm100( q, k, v, g, beta, do, - h, + checkpoint_series, dq_out, dk_out, dv_out, dg, - dbeta, + db, cu, - self._scale, - initial_state=s0_io, - d_initial_state=ds0, - d_final_state=dht, + self.scale, + initial_state=state0_io, + d_initial_state=dstate0, + d_final_state=dstate_in, work_items=work_items, work_count=work_count, - sched_ctr=sched_bwd if self._bwd_dyn_sched else None, + sched_ctr=sched_bwd if self.bwd_dyn_sched else None, log_gate=True, - workspace=tensormaps, + workspace=region["tensormaps"], stream=stream, ) - if self._is_gva: - from .common.head_reduce import head_group_reduce - - head_group_reduce(dq_out, dq, stream=stream) - head_group_reduce(dk_out, dk, stream=stream) - if self._is_gqa: + if dq_out is not dq or dk_out is not dk or dv_out is not dv: from .common.head_reduce import head_group_reduce - head_group_reduce(dv_out, dv, stream=stream) + for src_ho, dst in ((dq_out, dq), (dk_out, dk), (dv_out, dv)): + if src_ho is not dst: + head_group_reduce(src_ho, dst, stream=stream) return None diff --git a/python/cudnn/linear_attention/frost/kda_engine.py b/python/cudnn/linear_attention/frost/kda_engine.py index 85f58be11..471c9c5cd 100644 --- a/python/cudnn/linear_attention/frost/kda_engine.py +++ b/python/cudnn/linear_attention/frost/kda_engine.py @@ -2,8 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """FROST KDA engine: KDA nodes on the chunked prefill kernel -(``kernel/kda_prefill_f16.py``, Blackwell SM100/SM103, bf16/fp16, BT=16). -Forward only: KDA_BWD declines and falls back to ``KdaCuTileEngine``.""" +(``kernel/kda_prefill_f16.py``) and KDA_BWD nodes on the chunked backward +kernel (``kernel/kda_bprop_f16.py``), Blackwell SM100/SM103, bf16/fp16, +BT=16. The backward regenerates the per-chunk state checkpoints with the recompute +kernel (``kernel/kda_recompute_f16.py``) when the graph does not provide +one.""" from __future__ import annotations @@ -13,65 +16,25 @@ from cudnn import behavior_note from cudnn.engines.base import BaseEngine, CompiledPlan -from cudnn.frost import buffers -from cudnn.frost.workspace import Workspace, WorkspaceLayout, carve_plan -from ..engine_utils import _FrostPlan, _require_dtype, _require_state_pair - - -def _the_kda_node(graph): - nodes = list(graph.nodes) - if len(nodes) != 1: - return None - node = nodes[0] - if getattr(node.node_type, "name", None) not in ("KDA", "KDA_BWD"): - return None - return node - - -def _check_common(node) -> None: - """Shape/dtype gates for the prefill kernel.""" - import cudnn - - q, k, v = (node.inputs[p] for p in ("q", "k", "v")) - io_dtypes = {q.get_data_type(), k.get_data_type(), v.get_data_type()} - {None} - if len(io_dtypes) > 1: - raise NotImplementedError(f"KdaFrostEngine: q/k/v dtypes must match, got {io_dtypes}") - for p, t in (("q", q), ("k", k), ("v", v)): - if t.get_data_type() not in (cudnn.data_type.BFLOAT16, cudnn.data_type.HALF, None): - raise NotImplementedError(f"KdaFrostEngine: '{p}' must be bf16 or fp16, got {t.get_data_type()}") - if not t.dim or len(t.dim) != 3: - raise NotImplementedError(f"KdaFrostEngine: '{p}' must be THD [total_T, heads, dim]") - if q.dim[-1] != 128 or v.dim[-1] != 128: - raise NotImplementedError(f"KdaFrostEngine: head dims must be 128 (the recurrent state is 128x128), got K={q.dim[-1]} V={v.dim[-1]}") - if k.dim[1] != q.dim[1]: - raise NotImplementedError(f"KdaFrostEngine: q and k head counts differ ({q.dim[1]} vs {k.dim[1]})") - if v.dim[1] % q.dim[1] != 0: - raise NotImplementedError(f"KdaFrostEngine: v heads ({v.dim[1]}) must be a multiple of q heads ({q.dim[1]})") - - # kernel-native operand dtypes: buffers pass through without staging - fp32 = cudnn.data_type.FLOAT - _require_dtype("KdaFrostEngine", node, "g", fp32) - if node.node_type.name == "KDA" and node.params.get("use_beta_sigmoid", False): - # in-kernel sigmoid: beta arrives as io-dtype logits - if node.inputs["beta"].get_data_type() not in (q.get_data_type(), None): - raise NotImplementedError("KdaFrostEngine: use_beta_sigmoid takes io-dtype beta logits") - else: - _require_dtype("KdaFrostEngine", node, "beta", fp32) - _require_dtype("KdaFrostEngine", node, "a_log", fp32) - _require_dtype("KdaFrostEngine", node, "dt_bias", fp32) - _require_dtype("KdaFrostEngine", node, "cu_seqlens", cudnn.data_type.INT32) - _require_dtype("KdaFrostEngine", node, "initial_state", (fp32, cudnn.data_type.BFLOAT16)) - _require_dtype("KdaFrostEngine", node, "final_state", (fp32, cudnn.data_type.BFLOAT16), out=True) - _require_state_pair("KdaFrostEngine", node) +from cudnn.frost.buffers import current_device_id +from cudnn.frost.device import multiprocessor_count +from cudnn.frost.workspace import WorkspaceLayout, carve_plan +from ..graph_analyzer import FrostLaPlan, frost_la_gate, require, analyze def build_kda(graph): """The expensive step: import the kernel module (pulls in the Cutlass primitives; the cute.compile itself is cached inside the kernel per static config and runs on first execute, when the real buffers are known).""" - node = _the_kda_node(graph) - if node is None or node.node_type.name != "KDA": - raise ValueError("build_kda: graph does not contain exactly one KDA node") + nodes = list(graph.nodes) + if len(nodes) != 1 or getattr(nodes[0].node_type, "name", None) not in ("KDA", "KDA_BWD"): + raise ValueError("build_kda: graph does not contain exactly one KDA/KDA_BWD node") + node = nodes[0] + if node.node_type.name == "KDA_BWD": + from .kernel import kda_bprop_f16 as bwd_mod + from .kernel import kda_recompute_f16 as regen_mod + + return CompiledKdaBwd(node, bwd_mod, regen_mod) from .kernel import kda_prefill_f16 as kernel_mod return CompiledKda(node, kernel_mod) @@ -80,136 +43,169 @@ def build_kda(graph): class KdaFrostEngine(BaseEngine): """FROST chunked-kernel backend for single-node KDA graphs (THD layout). - Default KDA forward engine on SM100/SM103 (lowest KDA engine_id); - declines ``KDA_BWD`` so the router falls back to ``KdaCuTileEngine``.""" + Default KDA engine on SM100/SM103 (lowest KDA engine_id); serves KDA + forward and KDA_BWD (with a forward checkpoint recompute when ``state_checkpoints`` is absent).""" name = "kda_frost" behavior_notes = (behavior_note.RUNTIME_COMPILATION,) # JIT-compiled at build_plans() def check_support(self, graph) -> None: - node = _the_kda_node(graph) - if node is None: - raise NotImplementedError("KdaFrostEngine supports exactly one KDA node") - if node.node_type.name == "KDA_BWD": - raise NotImplementedError("KdaFrostEngine: the FROST KDA backward kernel is a stub; KdaCuTileEngine covers gradients") - sm = buffers.current_sm() - if sm is None or not (100 <= sm <= 103): - raise NotImplementedError(f"KdaFrostEngine requires SM100-SM103 (found {sm})") - try: - import cutlass.experimental.primitives # noqa: F401 — availability probe: ImportError = decline - except ImportError as exc: - raise NotImplementedError(f"KdaFrostEngine requires the Cutlass DSL with cutlass.experimental.primitives: {exc}") from exc - for port in ("q", "k", "v", "g", "beta", "cu_seqlens"): - if port not in node.inputs: - raise NotImplementedError(f"KdaFrostEngine: KDA node '{node.name}' is missing input '{port}'") - if int(node.params.get("checkpoint_every_n_tokens", 0) or 0) or "H" in node.outputs: - raise NotImplementedError("KdaFrostEngine: the per-chunk H output lands with the backward kernel (jopark/kda_gdn2_bprop)") - if node.node_type.name != "KDA" and (node.params.get("use_beta_sigmoid", False) or node.params.get("safe_gate", False)): + import cudnn + + facts = graph._facts_for(analyze) + frost_la_gate("KdaFrostEngine", facts, "KDA") + ckpt = facts.checkpoint_every_n_tokens + if ckpt and (facts.is_bwd or ckpt % 16 != 0): + raise NotImplementedError(f"KdaFrostEngine: checkpoint_every_n_tokens must be a positive multiple of 16 on the KDA node (got {ckpt})") + if not facts.gates_at_ho: + raise NotImplementedError(f"KdaFrostEngine: g/beta must carry HO = max(q, v) heads ({facts.h_o})") + if facts.is_bwd and (facts.use_beta_sigmoid or facts.safe_gate): raise NotImplementedError("KdaFrostEngine: use_beta_sigmoid/safe_gate are forward-node attributes") - if node.params.get("safe_gate", False): - for p in ("a_log", "dt_bias"): - if p not in node.inputs: - raise NotImplementedError(f"KdaFrostEngine: safe_gate requires input '{p}'") - elif "a_log" in node.inputs or "dt_bias" in node.inputs: - raise NotImplementedError("KdaFrostEngine: a_log/dt_bias require safe_gate=True") - _check_common(node) + fp32 = cudnn.data_type.FLOAT + if not facts.is_bwd and facts.use_beta_sigmoid: + # in-kernel sigmoid: Beta arrives as io-dtype logits + if facts.beta_dtype not in (facts.io_dtype, None): + raise NotImplementedError("KdaFrostEngine: use_beta_sigmoid takes io-dtype beta logits") + else: + require("KdaFrostEngine", "beta", facts.beta_dtype, fp32) + require("KdaFrostEngine", "a_log", facts.a_log_dtype, fp32) + require("KdaFrostEngine", "dt_bias", facts.dt_bias_dtype, fp32) + if facts.is_bwd: + for port, got in (("dO", facts.do_dtype), ("state_checkpoints", facts.state_checkpoints_dtype)): + if got not in (facts.io_dtype, None): + raise NotImplementedError(f"KdaFrostEngine: '{port}' must match the io dtype") + require("KdaFrostEngine", "initial_state", facts.state_dtype, fp32) + require("KdaFrostEngine", "d_final_state", facts.d_final_state_dtype, fp32) + require("KdaFrostEngine", "d_initial_state", facts.d_initial_state_dtype, fp32) + require("KdaFrostEngine", "dG", facts.dg_dtype, fp32) + require("KdaFrostEngine", "dBeta", facts.dbeta_dtype, fp32) + else: + state_dtypes = (fp32, cudnn.data_type.BFLOAT16) + require("KdaFrostEngine", "initial_state", facts.state_dtype, state_dtypes) + if facts.io_dtype is not None: + require("KdaFrostEngine", "state_checkpoints", facts.state_checkpoints_out_dtype, facts.io_dtype) + require("KdaFrostEngine", "final_state", facts.final_state_dtype, state_dtypes) + if not facts.state_pair_match: + raise NotImplementedError("KdaFrostEngine: initial_state and final_state dtypes must match") def build_plan(self, graph, plan, ctx=None) -> CompiledPlan: - return _FrostPlan(build_kda(graph)) + return FrostLaPlan(build_kda(graph)) class CompiledKda: """Compiled FROST KDA plan: a callable over the resolved node buffers.""" def __init__(self, node, kernel_mod): - from .common.split_k import WORK_ITEM_FIELDS, chunk_scratch_rows, compute_ideal_chunks, max_work_items + from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items - self._node = node - self._kernel = kernel_mod + self.node = node + self.kernel = kernel_mod + self.build_split_table = build_split_table + self.plan_name = "KdaFrostEngine (KDA)" scale = node.params.get("scale") - self._scale = float(scale) if scale is not None else 1.0 / math.sqrt(node.inputs["q"].dim[-1]) - self._use_qk_l2norm = bool(node.params.get("use_qk_l2norm", False)) - self._use_beta_sigmoid = bool(node.params.get("use_beta_sigmoid", False)) - self._safe_gate = bool(node.params.get("safe_gate", False)) + self.scale = float(scale) if scale is not None else 1.0 / math.sqrt(node.inputs["q"].dim[-1]) + self.use_qk_l2norm = bool(node.params.get("use_qk_l2norm", False)) + self.use_beta_sigmoid = bool(node.params.get("use_beta_sigmoid", False)) + self.safe_gate = bool(node.params.get("safe_gate", False)) glb = node.params.get("gate_lower_bound") - self._gate_lower_bound = float(glb) if glb is not None else kernel_mod.DEFAULT_GATE_LOWER_BOUND - self._has_fs = "final_state" in node.outputs + self.gate_lower_bound = float(glb) if glb is not None else kernel_mod.DEFAULT_GATE_LOWER_BOUND + self.has_final_state = "final_state" in node.outputs + self.has_state_checkpoints = "state_checkpoints" in node.outputs + self.ckpt = int(node.params.get("checkpoint_every_n_tokens", 0) or 0) + self.batch_invariant = bool(node.params.get("batch_invariant", False)) q, g = node.inputs["q"], node.inputs["g"] - self._b_t = kernel_mod.CFG.B_T + 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 total = q.dim[0] HO = g.dim[1] B = node.inputs["cu_seqlens"].dim[0] - 1 layout = WorkspaceLayout() - self._off_sched = layout.add(8) # [ticket, done] for the dynamic scheduler - self._num_sm = kernel_mod._device_sm_count() - self._ideal = compute_ideal_chunks(total, HO, self._num_sm, self._b_t) - self._n_tiles = B * HO - self._work_item_rows = max_work_items(total, B, HO, self._ideal, self._b_t, self._num_sm) - self._n_heads_out = HO - self._off_work_items = layout.add(self._work_item_rows * WORK_ITEM_FIELDS * 4) - self._off_item_scratch = layout.add(self._work_item_rows * WORK_ITEM_FIELDS * 4) - self._off_work_count = layout.add(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._tensormap_bytes = kernel_mod.get_workspace_size(B, HO, HO) - self._off_tensormaps = layout.add(self._tensormap_bytes, align=128) - self._ws_bytes = layout.size - - self._carve = carve_plan( - "kda", - [ - (self._off_sched, "int32", (2,)), - (self._off_work_items, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), - (self._off_item_scratch, "int32", (self._work_item_rows, WORK_ITEM_FIELDS)), - (self._off_work_count, "int32", (1,)), - (self._off_chunk_scratch, "float32", (self._chunk_scratch_rows, HO)), - (self._off_tensormaps, "int64", (self._tensormap_bytes // 8,)), - ], - ) + self.off_sched = layout.add(8) # [ticket, done] for the dynamic scheduler + self.num_sm = multiprocessor_count(current_device_id()) + self.n_tiles = B * HO + self.n_heads_out = HO + if self.split: + self.ideal = compute_ideal_chunks(total, HO, self.num_sm, self.b_t) + self.work_item_rows = max_work_items(total, B, HO, self.ideal, self.b_t, self.num_sm) + else: + self.ideal = None + self.work_item_rows = self.n_tiles + self.off_work_items = layout.add(self.work_item_rows * WORK_ITEM_FIELDS * 4) + self.off_work_count = layout.add(4) + if self.split: + 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) + from .common.host import tensormap_workspace_bytes + + self.tensormap_bytes = tensormap_workspace_bytes(kernel_mod, B) + self.off_tensormaps = layout.add(self.tensormap_bytes, align=128) + self.ws_bytes = layout.size + regions = [ + (self.off_sched, "int32", (2,)), + (self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), + (self.off_work_count, "int32", (1,)), + ] + if self.split: + regions += [ + (self.off_item_scratch, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), + (self.off_chunk_scratch, "float32", (self.chunk_scratch_rows, self.n_heads_out)), + ] + regions.append((self.off_tensormaps, "int64", (self.tensormap_bytes // 8,))) + self.carve = carve_plan("KdaFrostEngine (KDA)", regions) def workspace_bytes(self) -> int: - return self._ws_bytes + return self.ws_bytes def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: - nb = node_buffers[self._node] + nb = node_buffers[self.node] q = nb.inputs["q"] k = nb.inputs["k"] v = nb.inputs["v"] g = nb.inputs["g"] beta = nb.inputs["beta"] cu = nb.inputs["cu_seqlens"] - s0 = nb.inputs.get("initial_state") + state0 = nb.inputs.get("initial_state") o = nb.outputs["O"] - fs = nb.outputs["final_state"] if self._has_fs else None - + final_state = nb.outputs["final_state"] if self.has_final_state else None + state_checkpoints = nb.outputs["state_checkpoints"] if self.has_state_checkpoints else None + a_log = nb.inputs.get("a_log") + dt_bias = nb.inputs.get("dt_bias") stream = stream if stream is not None else 0 - ws = workspace - sched_ctr, work_items, item_scratch, work_count, chunk_scratch, tensormaps = ws.carve(self._carve) - from .common.split_k import build_split_table - - build_split_table( + if self.split: + sched_ctr, work_items, work_count, item_scratch, chunk_scratch, tensormaps = workspace.carve(self.carve) + else: + sched_ctr, work_items, work_count, tensormaps = workspace.carve(self.carve) + item_scratch = chunk_scratch = None + self.build_split_table( g, cu, work_items, work_count, - ideal_chunks=self._ideal, - n_tiles=self._n_tiles, - num_sms=self._num_sm, - b_t=self._b_t, + ideal_chunks=self.ideal, + n_tiles=self.n_tiles, + num_sms=self.num_sm, + b_t=self.b_t, chunk_scratch=chunk_scratch, item_scratch=item_scratch, log_gate=True, - safe_gate=self._safe_gate, - a_log=nb.inputs.get("a_log"), - dt_bias=nb.inputs.get("dt_bias"), - gate_lower_bound=self._gate_lower_bound if self._safe_gate else None, + safe_gate=self.safe_gate, + a_log=a_log, + dt_bias=dt_bias, + gate_lower_bound=self.gate_lower_bound if self.safe_gate else None, sched_ctr=sched_ctr, + split=self.split, stream=stream, ) - self._kernel.chunk_kda_sm100( + ckpt_kwargs = {} + if self.has_state_checkpoints: + # the kernel derives the per-sequence checkpoint entry offsets on device + ckpt_kwargs = dict(checkpoint_every_n_tokens=self.ckpt, output_state_checkpoints=state_checkpoints) + self.kernel.chunk_kda_sm100( q, k, v, @@ -217,19 +213,237 @@ def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: beta, o, cu, - s0, - fs, - self._scale, - use_qk_l2norm_in_kernel=self._use_qk_l2norm, - use_beta_sigmoid_in_kernel=self._use_beta_sigmoid, - safe_gate=self._safe_gate, - gate_lower_bound=self._gate_lower_bound, - a_log=nb.inputs.get("a_log"), - dt_bias=nb.inputs.get("dt_bias"), + state0, + final_state, + self.scale, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, + use_beta_sigmoid_in_kernel=self.use_beta_sigmoid, + safe_gate=self.safe_gate, + gate_lower_bound=self.gate_lower_bound, + a_log=a_log, + dt_bias=dt_bias, work_items=work_items, work_count=work_count, sched_ctr=sched_ctr, tensormap_workspace=tensormaps, + **ckpt_kwargs, + stream=stream, + ) + return None + + +class CompiledKdaBwd: + """Compiled FROST KDA backward plan. When the graph carries no ``state_checkpoints`` + input, one recompute pass over the forward inputs regenerates the per-chunk + checkpoint series into workspace (checkpoint stride ``b_t``). GVA/GQA + gradients land in HO-head scratch and are reduced back to the native head + counts.""" + + def __init__(self, node, bwd_mod, regen_mod): + from .common.split_k import WORK_ITEM_FIELDS, build_split_table, chunk_scratch_rows, compute_ideal_chunks, max_work_items + + self.node = node + self.bwd = bwd_mod + self.regen = regen_mod + self.build_split_table = build_split_table + self.plan_name = "KdaFrostEngine (KDA_BWD)" + from .common.downcast import downcast_state + from .common.host import tensormap_workspace_bytes + + self.downcast_state = downcast_state + scale = node.params.get("scale") + self.scale = float(scale) if scale is not None else 1.0 / math.sqrt(node.inputs["q"].dim[-1]) + self.use_qk_l2norm = bool(node.params.get("use_qk_l2norm", False)) + 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"] + self.b_t = bwd_mod.CFG.B_T + total = q.dim[0] + HQ, HV = q.dim[1], v.dim[1] + HO = g.dim[1] + K, V = q.dim[-1], v.dim[-1] + B = node.inputs["cu_seqlens"].dim[0] - 1 + self.io_name = "float16" if node.inputs["q"].get_data_type().name == "HALF" else "bfloat16" + self.n_heads_out, self.total = HO, total + layout = WorkspaceLayout() + self.off_sched = layout.add(16) # one [ticket, done] ring each for the regen and bwd kernels + self.num_sm = multiprocessor_count(current_device_id()) + self.bwd_dyn_sched = B * HO <= self.num_sm + self.batch_invariant = bool(node.params.get("batch_invariant", False)) + # cuts never in batch-invariant mode: whole-sequence items keep each + # sequence's math independent of the batch composition + self.split = not self.batch_invariant + self.n_tiles = B * HO + if self.split: + self.ideal = compute_ideal_chunks(total, HO, self.num_sm, self.b_t) + self.work_item_rows = max_work_items(total, B, HO, self.ideal, self.b_t, self.num_sm) + else: + self.ideal = None + self.work_item_rows = self.n_tiles + self.off_work_items = layout.add(self.work_item_rows * WORK_ITEM_FIELDS * 4) + self.off_work_count = layout.add(4) + if self.split: + 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) + # chunk-0 entering state, io dtype (downcast initial_state; absent = in-kernel zeros) + 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) + self.regen_tm_bytes = tensormap_workspace_bytes(regen_mod, B) + self.off_regen_tensormaps = layout.add(self.regen_tm_bytes, align=128) + HK = node.inputs["k"].dim[1] + self.fold_dq = HQ < HO + self.fold_dk = HK < HO + self.fold_dv = HV < HO + if self.fold_dq: + self.off_dq_ho = layout.add(total * HO * K * 2) + if self.fold_dk: + self.off_dk_ho = layout.add(total * HO * K * 2) + if self.fold_dv: + self.off_dv_ho = layout.add(total * HO * V * 2) + self.bwd_tm_bytes = tensormap_workspace_bytes(bwd_mod, B) + self.off_bwd_tensormaps = layout.add(self.bwd_tm_bytes, align=128) + self.ws_bytes = layout.size + regions = [ + ("sched_regen", self.off_sched, "int32", (2,)), + ("sched_bwd", self.off_sched + 8, "int32", (2,)), + ("sched_all", self.off_sched, "int32", (4,)), + ("work_items", self.off_work_items, "int32", (self.work_item_rows, WORK_ITEM_FIELDS)), + ("work_count", self.off_work_count, "int32", (1,)), + ("bwd_tensormaps", self.off_bwd_tensormaps, "int64", (self.bwd_tm_bytes // 8,)), + ] + 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(("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))) + if self.fold_dk: + regions.append(("dk_ho", self.off_dk_ho, self.io_name, (total, HO, K))) + if self.fold_dv: + regions.append(("dv_ho", self.off_dv_ho, self.io_name, (total, HO, V))) + self.carve_names = [name for name, _off, _dt, _shape in regions] + self.carve = carve_plan("KdaFrostEngine (KDA_BWD)", [(off, dt, shape) for _name, off, dt, shape in regions]) + + def workspace_bytes(self) -> int: + return self.ws_bytes + + def __call__(self, node_buffers, *, workspace=None, stream=None) -> Any: + nb = node_buffers[self.node] + q = nb.inputs["q"] + k = nb.inputs["k"] + v = nb.inputs["v"] + g = nb.inputs["g"] + beta = nb.inputs["beta"] + cu = nb.inputs["cu_seqlens"] + do = nb.inputs["dO"] + state_checkpoints = nb.inputs.get("state_checkpoints") + state0 = nb.inputs.get("initial_state") + dstate_in = nb.inputs.get("d_final_state") + dq = nb.outputs["dQ"] + dk = nb.outputs["dK"] + dv = nb.outputs["dV"] + dg = nb.outputs["dG"] + db = nb.outputs["dBeta"] + dstate0 = nb.outputs.get("d_initial_state") + stream = stream if stream is not None else 0 + + HO, total = self.n_heads_out, self.total + K, V = q.shape[-1], v.shape[-1] + B = cu.shape[0] - 1 + region = dict(zip(self.carve_names, workspace.carve(self.carve))) + sched_regen = region["sched_regen"] + sched_bwd = region["sched_bwd"] + work_items = region["work_items"] + work_count = region["work_count"] + self.build_split_table( + g, + cu, + work_items, + work_count, + ideal_chunks=self.ideal, + n_tiles=self.n_tiles, + num_sms=self.num_sm, + b_t=self.b_t, + chunk_scratch=region.get("chunk_scratch"), + item_scratch=region.get("item_scratch"), + log_gate=True, + sched_ctr=region["sched_all"], + split=self.split, 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: + checkpoint_series = region["state_checkpoints"] + self.regen.chunk_kda_recompute_sm100( + k, + v, + g, + beta, + cu, + state0, + None, + checkpoint_every_n_tokens=self.b_t, + output_state_checkpoints=checkpoint_series, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, + work_items=work_items, + work_count=work_count, + sched_ctr=sched_regen, + tensormap_workspace=region["regen_tensormaps"], + stream=stream, + ) + + dq_out, dk_out, dv_out = dq, dk, dv + if self.fold_dq: + dq_out = region["dq_ho"] + if self.fold_dk: + dk_out = region["dk_ho"] + if self.fold_dv: + dv_out = region["dv_ho"] + + self.bwd.chunk_kda_bwd_sm100( + q, + k, + v, + g, + beta, + do, + checkpoint_series, + dq_out, + dk_out, + dv_out, + dg, + db, + cu, + self.scale, + initial_state=state0_io, + d_initial_state=dstate0 if self.has_dstate0 else None, + d_final_state=dstate_in, + use_qk_l2norm_in_kernel=self.use_qk_l2norm, + work_items=work_items, + work_count=work_count, + sched_ctr=sched_bwd if self.bwd_dyn_sched else None, + tensormap_workspace=region["bwd_tensormaps"], + stream=stream, + ) + if dq_out is not dq or dk_out is not dk or dv_out is not dv: + from .common.head_reduce import head_group_reduce + + for src_ho, dst in ((dq_out, dq), (dk_out, dk), (dv_out, dv)): + if src_ho is not dst: + head_group_reduce(src_ho, dst, stream=stream) return None diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_config.py b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_config.py index 549f14119..3d2bccf39 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_config.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_config.py @@ -15,23 +15,64 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Gated DeltaNet v2 (GDN-2) Cutlass DSL backward kernel config — STUB. - -The FROST GDN-2 backward is not implemented yet (see ``gdn2_bprop_f16.py``); -these constants mirror the prefill config's tile shape for when the backward -kernel lands. +"""Gated DeltaNet v2 (GDN-2) Cutlass DSL backward kernel config (fixed +compile-time constants). The BT=16 backward (channel-wise erase gate beta + per-value write gate w) mirrors the prefill's 16-warp +(512-thread) specialization; the derived SMEM/TMEM sizes and offsets are +stamped by ``build_cfg`` in ``gdn2_bprop_f16.py``. Target arch: Blackwell SM100 (GB200) / SM103 (GB300). """ from dataclasses import dataclass +from typing import Tuple @dataclass(frozen=True) class Cfg: - B_T: int = 16 - D_K: int = 128 - D_V: int = 128 + # --- tile shape --- + B_T: int = 16 # chunk-inner token tile (BT=16 KDA schedule) + D_K: int = 128 # query/key head dim + D_V: int = 128 # value head dim + + # --- warp assignments (16 warps = 512 threads) --- + COMPUTE_GROUP_0_WARP_IDS: Tuple[int, ...] = (0, 1, 2, 3) # forward gate cumsum + decay-operand materialize + COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (4, 5, 6, 7) # value-side TMEM staging / restages / dH capture + COMPUTE_GROUP_2_WARP_IDS: Tuple[int, ...] = (8, 9, 10, 11) # dq/dk-bank drain, dG assembly + reverse cumsum + SUPER_MMA_WARP_ID: int = 12 # register-MMA KK/A/dA/dM + Neumann T_inv + TCGEN05_MMA_WARP_ID: int = 13 # tcgen05 GEMM schedule + TMA_WARP_ID: int = 14 # q/k/v/gate/do/state(checkpoint) TMA loads + EPILOGUE_WARP_ID: int = 15 # dq/dk/dv TMA stores only + + # --- register split --- + # WG1's drain runs at the register ceiling (spilled at 152) while WG0 + # sits near ~100 live regs, so WG0 funds WG1. Constraints: compute + # groups can't go below the 128-reg launch base (setmaxregister is + # INCREASE-only there), and warps 12-15 are one warpgroup so they must + # share a single setmaxregister value (56; super/epilogue peak ~R49). + # dht configs keep ~15 in-loop WG2 spills at 136 (152 doesn't fit). + NUM_REGS_COMPUTE_GROUP_0: int = 128 + NUM_REGS_COMPUTE_GROUP_1: int = 184 + NUM_REGS_COMPUTE_GROUP_2: int = 136 + NUM_REGS_OTHER: int = 64 # warpgroup-uniform; +8 donated by CG1 + + THREADS_PER_WARP: int = 32 + + BUFFER_ALIGN_BYTES: int = 1024 + + # --- SMEM / TMEM ring stage counts --- + SMEM_RAW_STAGES: int = 2 + SMEM_STATE_STAGES: int = 1 + SMEM_DECAY_STAGES: int = 2 + SMEM_INTERMEDIATE_STAGES: int = 2 + SMEM_STATE_SCALE_DIAG_STAGES: int = 2 + SMEM_DQ_STAGES: int = 1 + SMEM_DK_STAGES: int = 1 + SMEM_DGATE_STAGES: int = 1 + SMEM_DB_STAGES: int = 1 + SMEM_DV_STAGES: int = 2 + SMEM_DWO_STAGES: int = 2 + + CLUSTER_SHAPE_MNK: Tuple[int, int, int] = (1, 1, 1) CFG = Cfg() 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 012ba5e43..5ab67d122 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_bprop_f16.py @@ -1,37 +1,4034 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# This kernel is derived from cuDNN, NVIDIA Corporation. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Gated DeltaNet v2 (GDN-2) Cutlass DSL backward kernel — STUB. - -The FROST GDN-2 backward is not implemented yet. Because the GDN-2 chunk -size is small, the design recomputes the forward per-chunk states inside the -backward (no H store in the prefill kernel). Until the backward kernel -lands, this module raises ``NotImplementedError``. - -Target arch: Blackwell SM100 (GB200) / SM103 (GB300). +"""Chunked Gated DeltaNet v2 (GDN-2) BPROP kernel for Blackwell SM100/SM103 +(Cutlass DSL), BT=16 tiling with a per-key-channel decay. Framework-neutral +entry ``chunk_gdn2_bwd_sm100``. + +Algorithm overview (per chunk c, iterated c = NT-1 .. 0; within-chunk +log2-domain gate cumsum G[t,d], eG = 2^G, eGl = 2^G[BT-1]): + Inputs : Q/K[BT,DK], V/dO[BT,DV], S = checkpoint[c-1] (state ENTERING chunk c, KV), + Gate[BT,DK], Beta[BT,DK] (per-key erase), W[BT,DV] (per-value write) + State : dH[DV,DK] (state gradient, fp32 TMEM, accumulated backward) + + Operands (WG0, prefill recompute): K_decay = eG.(Beta.K) (erase key, Beta + folded), K_inv = K/eG, K_restore = (eGl/eG).K, Q_decay = eG.Q, diag(eGl). + + Forward recompute: T_inv = (I + strict_tril(K_decay@K_inv^T))^-1 (register Neumann, + Beta pre-folded); A = tril_incl(Q_decay@K_inv^T); Y = W.V - S^T K_decay; U = T_inv@Y. + + Backward math: + dU = K_restore@dH + A^T@dO (Q_decay carries scale, so A does too) + dY = T_inv^T@dU + dV = W.dY dW_out = V.dY (elementwise, WG1) + dA = tril_incl(dO@U^T) (unscaled) dM = dY@U^T dM_strict = +strict(dM) + dQ = eG.scale.(dO@S^T + dA@K_inv) + dK = Beta.eG.dK_decay + dK_inv/eG + (eGl/eG).dK_restore where (sign-flipped parts) + dK_decay part = dY@S^T + dM_strict@K_inv (= -dK_decay) + dK_inv part = dA^T@(scale.Q_decay) - dM_strict^T@K_decay (= dK_inv; one TMEM + acc, the minus rides the staged -dM_strict tile) + dK_restore part = U@dH^T (= +dK_restore) + dBeta[t,d] = k_n.eG.dK_decay = -k_n.eG.dK_decay part (per-channel, WG2) + dGate[t,d] = q_n.dQ_pre + Beta.dBeta + k_n.(dK_inv_part/eG + - (eGl/eG).dK_restore_part) + dGate_last[d] = eGl.sum_v(dH.S) + sum_t k_n.(eGl/eG).dK_restore_part + 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 +(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). + +Warp assignments (16 warps = 512 threads): + warps 0-3 : WG0 - Gate prefix scan + decay/restore operands (all chunks) + warps 4-7 : WG1 - value-side TMEM staging, restages, dstate capture, dV/dW_out + warps 8-11 : WG2 - dQ/dK part drain, dGate/dBeta assembly, reverse cumsum + warp 12 : super-MMA - register KK/A/dA/dM + Neumann inverse + warp 13 : tcgen05-MMA - the backward schedule + warp 14 : TMA load - Q/K/V/Gate/dO/state(checkpoint) loads + Beta/W tiles + warp 15 : epilogue - dQ/dK/dV/dGate/dBeta/dW_out TMA stores """ -from __future__ import annotations +from dataclasses import dataclass +from functools import lru_cache +from typing import NamedTuple, Optional, Type + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.experimental.cuda as cuda +import cutlass.experimental.primitives as nvvm +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + +from ..common.split_k import decode_work_item +from ..common.host import get_dtype +from cudnn.frost.buffers import current_device_id, data_ptr +from cudnn.frost.device import multiprocessor_count +from ..common.thd import TENSOR_MAP_QWORDS, emit_copy_desc, emit_checkpoint_seq_descs, emit_seq_descs +from .gdn2_bprop_config import CFG +from cudnn.frost.tile_dsl.barrier import ( + advance, + MBarrier, + PipelineState, + Producer, +) +from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, tma_slice_runtime_desc +from cudnn.frost.tile_dsl.mma import mma_ss, mma_step, mma_ts_step +from cudnn.frost.tile_dsl.swizzle import swizzle_lin_S, swizzle_xor_128b +from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire +from cudnn.frost.tile_dsl.pointwise import ( + opaque_f32_zero, + f16x2_to_f32, + fmul2, + ffma2, + movmatrix_16b, + mul_f16x2, + fp32_to_fp16, + sub_f16x2, +) + +LOG2_E: float = 1.4426950408889634 + +L2_NORM_EPS: float = 1.0e-12 + + +class Gdn2BwdBars(NamedTuple): + """Every inter-warp handoff as an ``MBarrier`` over its ring. Consumers + track ``(idx, phase)`` inline; the producer tag selects the arrive + lowering (``TMA_LOAD``/``MMA_COMMIT``/``THREAD``). + + Buffers read by both the MMA warp and a compute/warp group carry mixed + arrive counts (one MMA commit + N thread arrivers) so the producer only + reuses the slot once every reader is done.""" + + mb_q_ready: MBarrier + mb_q_done: MBarrier + mb_k_ready: MBarrier + mb_k_done: MBarrier + mb_gate_ready: MBarrier + mb_gate_done: MBarrier + mb_beta_ready: MBarrier + mb_beta_done: MBarrier + mb_do_ready: MBarrier + mb_do_done: MBarrier + mb_do_mma_done: MBarrier + mb_v_ready: MBarrier + mb_v_done: MBarrier + mb_w_ready: MBarrier + mb_w_done: MBarrier + mb_state_ready: MBarrier + mb_state_done: MBarrier + mb_state_cg0_done: MBarrier + mb_state_inp_ready: MBarrier + mb_state_inp_done: MBarrier + mb_state_inp_cg2_done: MBarrier + + mb_k_decay_inv_ready: MBarrier + mb_q_decay_k_restore_ready: MBarrier + mb_decay_done: MBarrier + + mb_t_inv_ready: MBarrier + mb_a_ready: MBarrier + mb_da_ready: MBarrier + mb_dm_ready: MBarrier + mb_a_done: MBarrier + mb_t_inv_done: MBarrier + mb_da_done: MBarrier + mb_dm_done: MBarrier + + mb_state_k_acc_ready: MBarrier + mb_y_inp_ready: MBarrier + mb_u_acc_ready: MBarrier + mb_u_smem_ready: MBarrier + mb_du_acc_ready: MBarrier + mb_du_inp_ready: MBarrier + mb_dy_acc_ready: MBarrier + mb_neg_dy_inp_ready: MBarrier + mb_dy_smem_ready: MBarrier + mb_dy_smem_done: MBarrier + mb_dstate_acc_ready: MBarrier + mb_dstate_inp_ready: MBarrier + mb_dstate_smem_ready: MBarrier + mb_dstate_smem_done: MBarrier + mb_dstate_smem_cg2_done: MBarrier + + mb_dq_acc_ready: MBarrier + mb_dk_decay_part_acc_ready: MBarrier + mb_dk_inv_part_acc_ready: MBarrier + mb_dk_restore_part_acc_ready: MBarrier + mb_dqk_acc_done: MBarrier + + mb_qk_raw_ready: MBarrier + mb_qk_raw_done: MBarrier + + mb_dq_tmastg_ready: MBarrier + mb_dq_tmastg_done: MBarrier + mb_dk_tmastg_ready: MBarrier + mb_dk_tmastg_done: MBarrier + mb_dv_tmastg_ready: MBarrier + mb_dv_tmastg_done: MBarrier + mb_dgate_tmastg_ready: MBarrier + mb_dgate_tmastg_done: MBarrier + mb_db_tmastg_ready: MBarrier + mb_db_tmastg_done: MBarrier + mb_dwo_tmastg_ready: MBarrier + mb_dwo_tmastg_done: MBarrier + mb_dstate0_acc_stored: MBarrier + mb_tmem_done: MBarrier + + mb_sched_ready: MBarrier + mb_sched_done: MBarrier + + +def make_gdn2_bwd_bars(cfg) -> Gdn2BwdBars: + """Bars factory. MUST be called from inside ``kernel`` (allocates the + mbarrier rings in SMEM ahead of the data buffers).""" + + def alloc(n): + return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=8) + + WARP = cfg.threads_per_warp + CG0 = len(cfg.compute_group_0_warp_ids) * WARP + CG2 = len(cfg.compute_group_2_warp_ids) * WARP + CG1 = len(cfg.compute_group_1_warp_ids) * WARP + MMA = 1 + + return Gdn2BwdBars( + mb_q_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_q_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0, producer=Producer.THREAD), + mb_k_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_k_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0, producer=Producer.THREAD), + mb_gate_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_gate_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0 + CG2, producer=Producer.THREAD), + mb_beta_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_beta_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0 + CG2, producer=Producer.THREAD), + mb_do_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_do_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=WARP, producer=Producer.THREAD), + mb_do_mma_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_v_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_v_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1, producer=Producer.THREAD), + mb_w_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_w_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1, producer=Producer.THREAD), + mb_state_ready=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_state_done=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_state_cg0_done=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=CG0, producer=Producer.THREAD), + mb_state_inp_ready=MBarrier(alloc(2), stages=2, init_count=CG0, producer=Producer.THREAD), + mb_state_inp_done=MBarrier(alloc(2), stages=2, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_state_inp_cg2_done=MBarrier(alloc(2), stages=2, init_count=CG2, producer=Producer.THREAD), + mb_k_decay_inv_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0, producer=Producer.THREAD), + mb_q_decay_k_restore_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0, producer=Producer.THREAD), + mb_decay_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_t_inv_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_a_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_da_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_dm_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_a_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_t_inv_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_da_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dm_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_u_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_du_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_du_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dy_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_neg_dy_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dy_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dy_smem_done=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dstate_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dstate_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dstate_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dstate_smem_done=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dstate_smem_cg2_done=MBarrier(alloc(1), stages=1, init_count=CG2, producer=Producer.THREAD), + mb_dq_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dk_decay_part_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dk_inv_part_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dk_restore_part_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dqk_acc_done=MBarrier(alloc(1), stages=1, init_count=CG2, producer=Producer.THREAD), + mb_qk_raw_ready=MBarrier(alloc(cfg.tmem_qk_raw_stages), stages=cfg.tmem_qk_raw_stages, init_count=CG0, producer=Producer.THREAD), + mb_qk_raw_done=MBarrier(alloc(cfg.tmem_qk_raw_stages), stages=cfg.tmem_qk_raw_stages, init_count=CG2, producer=Producer.THREAD), + mb_dq_tmastg_ready=MBarrier(alloc(cfg.smem_dq_stages), stages=cfg.smem_dq_stages, init_count=CG2, producer=Producer.THREAD), + mb_dq_tmastg_done=MBarrier(alloc(cfg.smem_dq_stages), stages=cfg.smem_dq_stages, init_count=WARP, producer=Producer.THREAD), + mb_dk_tmastg_ready=MBarrier(alloc(cfg.smem_dk_stages), stages=cfg.smem_dk_stages, init_count=CG2, producer=Producer.THREAD), + mb_dk_tmastg_done=MBarrier(alloc(cfg.smem_dk_stages), stages=cfg.smem_dk_stages, init_count=WARP, producer=Producer.THREAD), + mb_dv_tmastg_ready=MBarrier(alloc(cfg.smem_dv_stages), stages=cfg.smem_dv_stages, init_count=CG1, producer=Producer.THREAD), + mb_dv_tmastg_done=MBarrier(alloc(cfg.smem_dv_stages), stages=cfg.smem_dv_stages, init_count=WARP, producer=Producer.THREAD), + mb_dgate_tmastg_ready=MBarrier(alloc(cfg.smem_dgate_stages), stages=cfg.smem_dgate_stages, init_count=CG2, producer=Producer.THREAD), + mb_dgate_tmastg_done=MBarrier(alloc(cfg.smem_dgate_stages), stages=cfg.smem_dgate_stages, init_count=WARP, producer=Producer.THREAD), + mb_db_tmastg_ready=MBarrier(alloc(cfg.smem_db_stages), stages=cfg.smem_db_stages, init_count=CG2, producer=Producer.THREAD), + mb_db_tmastg_done=MBarrier(alloc(cfg.smem_db_stages), stages=cfg.smem_db_stages, init_count=WARP, producer=Producer.THREAD), + mb_dwo_tmastg_ready=MBarrier(alloc(cfg.smem_dwo_stages), stages=cfg.smem_dwo_stages, init_count=CG1, producer=Producer.THREAD), + mb_dwo_tmastg_done=MBarrier(alloc(cfg.smem_dwo_stages), stages=cfg.smem_dwo_stages, init_count=WARP, producer=Producer.THREAD), + mb_dstate0_acc_stored=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1 + CG2, producer=Producer.THREAD), + mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), + mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), + ) + + +# ---- Dynamic tile scheduler ------------------------------------------------------ + + +@cute.jit +def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): + """TMA-warp side: pull the next tile off the global ticket, publish it.""" + if cutlass.const_expr(cfg.dyn_sched): + bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) + if nvvm.elect_sync(): + fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sSched[sched_state.idx] = num_ctas + fetched + nvvm.bar_warp_sync(cute.arch.FULL_MASK) + next_tile = sSched[sched_state.idx] + if nvvm.elect_sync(): + bars.mb_sched_ready[sched_state.idx].arrive() + return next_tile, advance(sched_state, cfg.sched_stages) + return tile_idx + num_ctas, sched_state + + +@cute.jit +def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): + """Consumer side: read the TMA warp's published next tile.""" + if cutlass.const_expr(cfg.dyn_sched): + bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) + next_tile = sSched[sched_state.idx] + if nvvm.elect_sync(): + bars.mb_sched_done[sched_state.idx].arrive() + return next_tile, advance(sched_state, cfg.sched_stages) + return tile_idx + num_ctas, sched_state + + +# ---- Warp bodies ----------------------------------------------------------------- + + +@cute.jit +def epilogue_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_inv_raw, + sQ_decay_raw, + sDo_raw, + sU_raw, + sIntermediate_raw, + sDq_raw, + sDk_raw, + sDv_raw, + sDgate_raw, + sDb_raw, + sDwOut_raw, + desc_dq_base, + desc_dk_base, + desc_dv_base, + desc_dgate_base, + desc_db_base, + desc_dwo_base, + bars, +) -> None: + """Epilogue warp role (warp 15): register-MMA A/dA tiles and the + gradient TMA stores, in chunk order with a one-behind store ladder.""" + elect_one = nvvm.elect_sync() + + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + + # ---- ldmatrix/stmatrix lane decode ------------------------------------------- + rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) + rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) + lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) + lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) + stsm_row_coord = lane & 7 + stsm_col_coord = cutlass.Int32(0) + if (lane // 8) & 1: + stsm_row_coord = stsm_row_coord + cutlass.Int32(8) + if lane // 8 >= 2: + stsm_col_coord = cutlass.Int32(8) + stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + stsm_col_coord, bbits=1, mbase=3, sshift=3) + row_lo = lane // 4 + row_hi = row_lo + cutlass.Int32(8) + + # hoisted tril bitmask: bit i = row >= col for accum index i + tril_incl_mask = cutlass.Int32(0) + for accum_idx in cutlass.range_constexpr(8): + row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo + col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + if cutlass.const_expr(accum_idx % 2 == 1): + col_coord = col_coord + cutlass.Int32(1) + tril_incl_mask = tril_incl_mask | (cutlass.Int32(1 << accum_idx) if row_coord >= col_coord else cutlass.Int32(0)) + raw_index = PipelineState.start(phase=0) + u_index = PipelineState.start(phase=0) + gbase = cutlass.Int32(0) + + sDq_tma = SmemTile( + base=sDq_raw, + elems_per_stage=(cfg.b_t * cfg.d_k), + stages=cfg.smem_dq_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.b_t * 64, + ) + sDk_tma = SmemTile( + base=sDk_raw, + elems_per_stage=(cfg.b_t * cfg.d_k), + stages=cfg.smem_dk_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.b_t * 64, + ) + sDv_tma = SmemTile( + base=sDv_raw, + elems_per_stage=(cfg.b_t * cfg.d_v), + stages=cfg.smem_dv_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.b_t * 64, + ) + sDgate_tma = SmemTile( + base=sDgate_raw, + elems_per_stage=(cfg.b_t * cfg.d_k), + stages=cfg.smem_dgate_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 32), + tma_granu_elems=32, + tma_subtile_stride_elems=cfg.b_t * 32, + ) + sDb_tma = SmemTile( + base=sDb_raw, + elems_per_stage=(cfg.b_t * cfg.d_k), + stages=cfg.smem_db_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.b_t * 64, + ) + sDwOut_tma = SmemTile( + base=sDwOut_raw, + elems_per_stage=(cfg.b_t * cfg.d_v), + stages=cfg.smem_dwo_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.b_t * 64, + ) + dq_index = PipelineState.start(phase=0) + dk_index = PipelineState.start(phase=0) + dv_index = PipelineState.start(phase=0) + dgate_index = PipelineState.start(phase=0) + db_index = PipelineState.start(phase=0) + dwo_index = PipelineState.start(phase=0) + sched_state = PipelineState.start(phase=0) + tile_idx = cutlass.Int32(bidx) + FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 + 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) + head_o = head_idx + slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) + if elect_one: + desc_dq_slot = (desc_dq_base + slot).tospace(cutlass.AddressSpace.generic) + desc_dk_slot = (desc_dk_base + slot).tospace(cutlass.AddressSpace.generic) + desc_dv_slot = (desc_dv_base + slot).tospace(cutlass.AddressSpace.generic) + desc_dgate_slot = (desc_dgate_base + slot).tospace(cutlass.AddressSpace.generic) + desc_db_slot = (desc_db_base + slot).tospace(cutlass.AddressSpace.generic) + desc_dwo_slot = (desc_dwo_base + slot).tospace(cutlass.AddressSpace.generic) + tma_tensormap_acquire(desc_dq_slot) + tma_tensormap_acquire(desc_dk_slot) + tma_tensormap_acquire(desc_dv_slot) + tma_tensormap_acquire(desc_dgate_slot) + tma_tensormap_acquire(desc_db_slot) + tma_tensormap_acquire(desc_dwo_slot) + sk_nt = cend - wstart + pend_start = cutlass.Int32(0) + pend_writes = cutlass.Boolean(False) + for rev_idx in cutlass.range(sk_nt, unroll=1): + chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_start = chunk_idx * cfg.b_t + writes = chunk_idx < wend + gc = gbase + rev_idx + decay_stage = gc % cfg.smem_decay_stages + intermediate_stage = gc % cfg.smem_intermediate_stages + raw_stage = raw_index.idx + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sQ_decay_ptr = sQ_decay_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sDo_ptr = sDo_raw.data_ptr() + raw_stage * (cfg.d_v * cfg.b_t) + sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (cfg.intermediate_tiles * cfg.b_t * cfg.b_t) + + # ---- A = tril_incl(Q_decay @ K_inv^T) -------------------------------- + bars.mb_a_done[intermediate_stage].wait(((gc // cfg.smem_intermediate_stages) + 1) % 2) + bars.mb_q_decay_k_restore_ready[decay_stage].wait((gc // cfg.smem_decay_stages) % 2) + a_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + a_acc[accum_idx] = cutlass.Float32(0.0) + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + a_col = k_block * 16 + lhs_col_offset + a_seg = a_col // 64 + a_frag = nvvm.ldmatrix( + sQ_decay_ptr + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + b_col = k_block * 16 + rhs_col_offset + b_seg = b_col // 64 + b_frag = nvvm.ldmatrix( + sK_inv_ptr + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + mma_step( + a_acc, + (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), + (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + for accum_idx in cutlass.range_constexpr(8): + a_acc[accum_idx] = a_acc[accum_idx] if (tril_incl_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + nvvm.stmatrix( + sIntermediate_ptr + stsm_idx, + [ + fp32_to_fp16(a_acc[0], a_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[2], a_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[4], a_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[6], a_acc[7], dtype=cfg.io_dtype), + ], + nvvm.MMALayout.ROW, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_a_ready[intermediate_stage].arrive() + + # ---- dA = tril_incl(dO @ U^T) ---------------------------------------- + bars.mb_u_smem_ready.wait(u_index.phase) + u_index = advance(u_index, 1) + da_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + da_acc[accum_idx] = cutlass.Float32(0.0) + for k_block in cutlass.range_constexpr(cfg.d_v // 16): + a_col = k_block * 16 + lhs_col_offset + a_seg = a_col // 64 + a_frag = nvvm.ldmatrix( + sDo_ptr + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + b_col = k_block * 16 + rhs_col_offset + b_seg = b_col // 64 + b_frag = nvvm.ldmatrix( + sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + mma_step( + da_acc, + (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), + (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + # fence: the dO/U ldmatrix reads must complete before this release + # licenses the TMA reload (sDo) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_do_done[raw_stage].arrive() + for accum_idx in cutlass.range_constexpr(8): + da_acc[accum_idx] = da_acc[accum_idx] if (tril_incl_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + bars.mb_da_done[intermediate_stage].wait(((gc // cfg.smem_intermediate_stages) + 1) % 2) + nvvm.stmatrix( + sIntermediate_ptr + 2 * (cfg.b_t * cfg.b_t) + stsm_idx, + [ + fp32_to_fp16(da_acc[0], da_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(da_acc[2], da_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(da_acc[4], da_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(da_acc[6], da_acc[7], dtype=cfg.io_dtype), + ], + nvvm.MMALayout.ROW, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_da_ready[intermediate_stage].arrive() + raw_index = advance(raw_index, cfg.smem_raw_stages) + + # ---- dQ/dK/dGate/dBeta/dV/dW_out: previous chunk, one-behind store ladder ---- + if rev_idx > 0: + bars.mb_dq_tmastg_ready[dq_index.idx].wait(dq_index.phase) + if pend_writes: + desc_dq_slot = (desc_dq_base + slot).tospace(cutlass.AddressSpace.generic) + dq_slice = tma_slice_runtime_desc(desc_dq_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDq_tma[dq_index.idx], dq_slice, acquire=False) + tma_store_commit() + bars.mb_dk_tmastg_ready[dk_index.idx].wait(dk_index.phase) + if pend_writes: + desc_dk_slot = (desc_dk_base + slot).tospace(cutlass.AddressSpace.generic) + dk_slice = tma_slice_runtime_desc(desc_dk_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDk_tma[dk_index.idx], dk_slice, acquire=False) + tma_store_commit() + bars.mb_dgate_tmastg_ready[dgate_index.idx].wait(dgate_index.phase) + if pend_writes: + desc_dgate_slot = (desc_dgate_base + slot).tospace(cutlass.AddressSpace.generic) + dgate_slice = tma_slice_runtime_desc(desc_dgate_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDgate_tma[dgate_index.idx], dgate_slice, acquire=False) + tma_store_commit() + bars.mb_db_tmastg_ready[db_index.idx].wait(db_index.phase) + if pend_writes: + desc_db_slot = (desc_db_base + slot).tospace(cutlass.AddressSpace.generic) + db_slice = tma_slice_runtime_desc(desc_db_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDb_tma[db_index.idx], db_slice, acquire=False) + tma_store_commit() + bars.mb_dv_tmastg_ready[dv_index.idx].wait(dv_index.phase) + if pend_writes: + desc_dv_slot = (desc_dv_base + slot).tospace(cutlass.AddressSpace.generic) + dv_slice = tma_slice_runtime_desc(desc_dv_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDv_tma[dv_index.idx], dv_slice, acquire=False) + tma_store_commit() + bars.mb_dwo_tmastg_ready[dwo_index.idx].wait(dwo_index.phase) + if pend_writes: + desc_dwo_slot = (desc_dwo_base + slot).tospace(cutlass.AddressSpace.generic) + dwo_slice = tma_slice_runtime_desc(desc_dwo_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDwOut_tma[dwo_index.idx], dwo_slice, acquire=False) + tma_store_commit() + tma_store_wait(5) + bars.mb_dq_tmastg_done[dq_index.idx].arrive() + tma_store_wait(4) + bars.mb_dk_tmastg_done[dk_index.idx].arrive() + tma_store_wait(3) + bars.mb_dgate_tmastg_done[dgate_index.idx].arrive() + tma_store_wait(2) + bars.mb_db_tmastg_done[db_index.idx].arrive() + tma_store_wait(1) + bars.mb_dv_tmastg_done[dv_index.idx].arrive() + tma_store_wait(0) + bars.mb_dwo_tmastg_done[dwo_index.idx].arrive() + dq_index = advance(dq_index, cfg.smem_dq_stages) + dk_index = advance(dk_index, cfg.smem_dk_stages) + dgate_index = advance(dgate_index, cfg.smem_dgate_stages) + db_index = advance(db_index, cfg.smem_db_stages) + dv_index = advance(dv_index, cfg.smem_dv_stages) + dwo_index = advance(dwo_index, cfg.smem_dwo_stages) + pend_start = chunk_start + pend_writes = writes + + # ---- tile tail: drain the last chunk's dQ/dK/dGate/dBeta/dV/dW_out ------- + if sk_nt > 0: + bars.mb_dq_tmastg_ready[dq_index.idx].wait(dq_index.phase) + if pend_writes: + desc_dq_slot = (desc_dq_base + slot).tospace(cutlass.AddressSpace.generic) + dq_slice = tma_slice_runtime_desc(desc_dq_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDq_tma[dq_index.idx], dq_slice, acquire=False) + tma_store_commit() + bars.mb_dk_tmastg_ready[dk_index.idx].wait(dk_index.phase) + if pend_writes: + desc_dk_slot = (desc_dk_base + slot).tospace(cutlass.AddressSpace.generic) + dk_slice = tma_slice_runtime_desc(desc_dk_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDk_tma[dk_index.idx], dk_slice, acquire=False) + tma_store_commit() + bars.mb_dgate_tmastg_ready[dgate_index.idx].wait(dgate_index.phase) + if pend_writes: + desc_dgate_slot = (desc_dgate_base + slot).tospace(cutlass.AddressSpace.generic) + dgate_slice = tma_slice_runtime_desc(desc_dgate_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDgate_tma[dgate_index.idx], dgate_slice, acquire=False) + tma_store_commit() + bars.mb_db_tmastg_ready[db_index.idx].wait(db_index.phase) + if pend_writes: + desc_db_slot = (desc_db_base + slot).tospace(cutlass.AddressSpace.generic) + db_slice = tma_slice_runtime_desc(desc_db_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDb_tma[db_index.idx], db_slice, acquire=False) + tma_store_commit() + bars.mb_dv_tmastg_ready[dv_index.idx].wait(dv_index.phase) + if pend_writes: + desc_dv_slot = (desc_dv_base + slot).tospace(cutlass.AddressSpace.generic) + dv_slice = tma_slice_runtime_desc(desc_dv_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDv_tma[dv_index.idx], dv_slice, acquire=False) + tma_store_commit() + bars.mb_dwo_tmastg_ready[dwo_index.idx].wait(dwo_index.phase) + if pend_writes: + desc_dwo_slot = (desc_dwo_base + slot).tospace(cutlass.AddressSpace.generic) + dwo_slice = tma_slice_runtime_desc(desc_dwo_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDwOut_tma[dwo_index.idx], dwo_slice, acquire=False) + tma_store_commit() + tma_store_wait(5) + bars.mb_dq_tmastg_done[dq_index.idx].arrive() + tma_store_wait(4) + bars.mb_dk_tmastg_done[dk_index.idx].arrive() + tma_store_wait(3) + bars.mb_dgate_tmastg_done[dgate_index.idx].arrive() + tma_store_wait(2) + bars.mb_db_tmastg_done[db_index.idx].arrive() + tma_store_wait(1) + bars.mb_dv_tmastg_done[dv_index.idx].arrive() + tma_store_wait(0) + bars.mb_dwo_tmastg_done[dwo_index.idx].arrive() + dq_index = advance(dq_index, cfg.smem_dq_stages) + dk_index = advance(dk_index, cfg.smem_dk_stages) + dgate_index = advance(dgate_index, cfg.smem_dgate_stages) + db_index = advance(db_index, cfg.smem_db_stages) + dv_index = advance(dv_index, cfg.smem_dv_stages) + dwo_index = advance(dwo_index, cfg.smem_dwo_stages) + + gbase += sk_nt + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def super_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_decay_raw, + sK_inv_raw, + sU_raw, + sDy_raw, + sIntermediate_raw, + bars, +) -> None: + """Super-MMA warp role (warp 12): builds the Neumann T_inv and + strict-tril dM staging tiles, in chunk order.""" + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + sdy_index = PipelineState.start(phase=0) + + # ---- ldmatrix lane decode ---------------------------------------------------- + rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) + rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) + lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) + lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) + stsm_row_coord = lane & 7 + stsm_col_coord = cutlass.Int32(0) + if (lane // 8) & 1: + stsm_row_coord = stsm_row_coord + cutlass.Int32(8) + if lane // 8 >= 2: + stsm_col_coord = cutlass.Int32(8) + stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + stsm_col_coord, bbits=1, mbase=3, sshift=3) + row_lo = lane // 4 + row_hi = row_lo + cutlass.Int32(8) + + # hoisted tril bitmasks: bit i = row > col / row == col for accum index i + tril_strict_mask = cutlass.Int32(0) + eye_mask = cutlass.Int32(0) + for accum_idx in cutlass.range_constexpr(8): + row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo + col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + if cutlass.const_expr(accum_idx % 2 == 1): + col_coord = col_coord + cutlass.Int32(1) + tril_strict_mask = tril_strict_mask | (cutlass.Int32(1 << accum_idx) if row_coord > col_coord else cutlass.Int32(0)) + eye_mask = eye_mask | (cutlass.Int32(1 << accum_idx) if row_coord == col_coord else cutlass.Int32(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 + for rev_idx in cutlass.range(sk_nt, unroll=1): + gc = gbase + rev_idx + decay_stage = gc % cfg.smem_decay_stages + intermediate_stage = gc % cfg.smem_intermediate_stages + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (cfg.intermediate_tiles * cfg.b_t * cfg.b_t) + + bars.mb_t_inv_done[intermediate_stage].wait(((gc // cfg.smem_intermediate_stages) + 1) % 2) + + # ---- KK = K_decay @ K_inv^T ------------------------------------------ + bars.mb_k_decay_inv_ready[decay_stage].wait((gc // cfg.smem_decay_stages) % 2) + kk_lhs_row = lhs_row_coord + kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + kk_acc[accum_idx] = cutlass.Float32(0.0) + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + a_col = k_block * 16 + lhs_col_offset + a_seg = a_col // 64 + a_frag = nvvm.ldmatrix( + sK_decay_ptr + a_seg * (cfg.b_t * 64) + kk_lhs_row * 64 + swizzle_xor_128b(kk_lhs_row, a_col - a_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + b_col = k_block * 16 + rhs_col_offset + b_seg = b_col // 64 + b_frag = nvvm.ldmatrix( + sK_inv_ptr + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + mma_step( + kk_acc, + (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), + (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + + # ---- L = tril(KK, -1) ------------------------------------------------ + l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + lower = kk_acc[accum_idx] if (tril_strict_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + l_regs[accum_idx] = lower + l_a0 = fp32_to_fp16(l_regs[0], l_regs[1], dtype=cfg.io_dtype) + l_a1 = fp32_to_fp16(l_regs[2], l_regs[3], dtype=cfg.io_dtype) + l_a2 = fp32_to_fp16(l_regs[4], l_regs[5], dtype=cfg.io_dtype) + l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) + l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) + + tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + eye = cutlass.Float32(1.0) if (eye_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + tinv_acc[accum_idx] = eye - l_values[accum_idx] + + lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) + for _round in cutlass.range_constexpr(3): + sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + sq_acc[accum_idx] = cutlass.Float32(0.0) + mma_step( + sq_acc, + (lpow_a0, lpow_a1, lpow_a2, lpow_a3), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + lpow_a0 = fp32_to_fp16(sq_acc[0], sq_acc[1], dtype=cfg.io_dtype) + lpow_a1 = fp32_to_fp16(sq_acc[2], sq_acc[3], dtype=cfg.io_dtype) + lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) + lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) + upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + upd_acc[accum_idx] = cutlass.Float32(0.0) + tinv_p0 = fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype) + tinv_p1 = fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype) + tinv_p2 = fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype) + tinv_p3 = fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype) + mma_step( + upd_acc, + (tinv_p0, tinv_p1, tinv_p2, tinv_p3), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + tinv_lo0, tinv_hi0 = f16x2_to_f32(tinv_p0, dtype=cfg.io_dtype) + tinv_lo1, tinv_hi1 = f16x2_to_f32(tinv_p1, dtype=cfg.io_dtype) + tinv_lo2, tinv_hi2 = f16x2_to_f32(tinv_p2, dtype=cfg.io_dtype) + tinv_lo3, tinv_hi3 = f16x2_to_f32(tinv_p3, dtype=cfg.io_dtype) + tinv_acc[0] = tinv_lo0 + upd_acc[0] + tinv_acc[1] = tinv_hi0 + upd_acc[1] + tinv_acc[2] = tinv_lo1 + upd_acc[2] + tinv_acc[3] = tinv_hi1 + upd_acc[3] + tinv_acc[4] = tinv_lo2 + upd_acc[4] + tinv_acc[5] = tinv_hi2 + upd_acc[5] + tinv_acc[6] = tinv_lo3 + upd_acc[6] + tinv_acc[7] = tinv_hi3 + upd_acc[7] + + nvvm.stmatrix( + sIntermediate_ptr + 1 * (cfg.b_t * cfg.b_t) + stsm_idx, + [ + fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype), + ], + nvvm.MMALayout.ROW, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_t_inv_ready[intermediate_stage].arrive() + + # ---- dM = dY @ U^T --------------------------------------------------- + bars.mb_dm_done[intermediate_stage].wait(((gc // cfg.smem_intermediate_stages) + 1) % 2) + bars.mb_dy_smem_ready.wait(sdy_index.phase) + sdy_index = advance(sdy_index, 1) + dm_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + dm_acc[accum_idx] = cutlass.Float32(0.0) + for k_block in cutlass.range_constexpr(cfg.d_v // 16): + a_col = k_block * 16 + lhs_col_offset + a_seg = a_col // 64 + a_frag = nvvm.ldmatrix( + sDy_raw.data_ptr() + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + b_col = k_block * 16 + rhs_col_offset + b_seg = b_col // 64 + b_frag = nvvm.ldmatrix( + sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + mma_step( + dm_acc, + (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), + (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + dm_strict_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + dm_strict_regs[accum_idx] = dm_acc[accum_idx] if (tril_strict_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + w0 = fp32_to_fp16(dm_strict_regs[0], dm_strict_regs[1], dtype=cfg.io_dtype) + w1 = fp32_to_fp16(dm_strict_regs[2], dm_strict_regs[3], dtype=cfg.io_dtype) + w2 = fp32_to_fp16(dm_strict_regs[4], dm_strict_regs[5], dtype=cfg.io_dtype) + w3 = fp32_to_fp16(dm_strict_regs[6], dm_strict_regs[7], dtype=cfg.io_dtype) + nvvm.stmatrix(sIntermediate_ptr + 3 * (cfg.b_t * cfg.b_t) + stsm_idx, [w0, w1, w2, w3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) + nw0 = fp32_to_fp16(-dm_strict_regs[0], -dm_strict_regs[1], dtype=cfg.io_dtype) + nw1 = fp32_to_fp16(-dm_strict_regs[2], -dm_strict_regs[3], dtype=cfg.io_dtype) + nw2 = fp32_to_fp16(-dm_strict_regs[4], -dm_strict_regs[5], dtype=cfg.io_dtype) + nw3 = fp32_to_fp16(-dm_strict_regs[6], -dm_strict_regs[7], dtype=cfg.io_dtype) + nvvm.stmatrix(sIntermediate_ptr + 4 * (cfg.b_t * cfg.b_t) + stsm_idx, [nw0, nw1, nw2, nw3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dm_ready[intermediate_stage].arrive() + gbase += sk_nt + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def tcgen05_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + tmem_hold, + sState_alt, + sK_decay_lead16, + sK_inv_amaj, + sK_restore_lead16, + sDo_lead16, + sDo_amaj, + sQ_decay_trans, + sK_decay_trans, + sU_lead16, + sDy_lead16, + sDstate_alt, + sIntermediate, + sState_scale_diag, + bars, +) -> None: + """tcgen05-MMA warp role (warp 13): issues every tcgen05 GEMM and owns + the TMEM lifecycle.""" + elect_one = nvvm.elect_sync() + + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + nvvm.tcgen05_alloc(tmem_hold, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_hold.load() + bpe = cfg.io_dtype.width // 8 + + # ---- chunk-invariant GEMM descriptors ---------------------------------------- + idesc_mv_nt = 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, + ) + idesc_state_k_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_v, + a_major=1, + ) + bmm_state_k_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=True, + cta_group=1, + idesc=idesc_state_k_at, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_dvinter_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, + cta_group=1, + idesc=idesc_mv_nt, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_du_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_v, + a_major=1, + b_major=1, + ) + bmm_du_at_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + atranspose=True, + cta_group=1, + idesc=idesc_du_at, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dstate_q_at = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.d_k, + m_dim=cfg.d_v, + a_major=1, + b_major=1, + ) + bmm_dstate_q_at_desc = MmaDesc( + M=cfg.d_v, + N=cfg.d_k, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + atranspose=True, + cta_group=1, + idesc=idesc_dstate_q_at, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_qk_ts_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_mv_nt, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_mv_nt_t = 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, + b_major=1, + ) + bmm_qk_ts_t_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + cta_group=1, + idesc=idesc_mv_nt_t, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_diag = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=16, + m_dim=cfg.d_v, + ) + bmm_diag_desc = MmaDesc( + M=cfg.d_v, + N=16, + K=16, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_diag, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dstate_k = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.d_k, + m_dim=cfg.d_v, + b_major=1, + ) + bmm_dstate_k_desc = MmaDesc( + M=cfg.d_v, + N=cfg.d_k, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + cta_group=1, + idesc=idesc_dstate_k, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_state_ts = 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, + ) + bmm_state_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.d_v, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_state_ts, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_state_ts_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.d_v, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_state_ts, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dstate_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_dstate_at_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.d_v, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + atranspose=True, + cta_group=1, + idesc=idesc_dstate_at, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dgp = 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, + ) + bmm_dgrad_ts_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_dgp, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dgp_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_dgrad_at_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + atranspose=True, + cta_group=1, + idesc=idesc_dgp_at, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dgp_at_t = 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, + b_major=1, + ) + bmm_dgrad_at_t_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + atranspose=True, + cta_group=1, + idesc=idesc_dgp_at_t, + kind=nvvm.Tcgen05MMAKind.F16, + ) + + state_index = PipelineState.start(phase=0) + y_inp_index = PipelineState.start(phase=0) + dstate_inp_index = PipelineState.start(phase=0) + du_inp_index = PipelineState.start(phase=0) + neg_dy_index = PipelineState.start(phase=0) + u_smem_index = PipelineState.start(phase=0) + dstate_smem_index = PipelineState.start(phase=0) + parts_done_index = PipelineState.start(phase=1) + + do_seg = (cfg.b_t * cfg.d_v * (cfg.io_dtype.width // 8)) >> 4 + op_seg = (cfg.b_t * cfg.d_k * (cfg.io_dtype.width // 8)) >> 4 + intermediate_seg = (cfg.intermediate_tiles * cfg.b_t * cfg.b_t * (cfg.io_dtype.width // 8)) >> 4 + intermediate_slot = (cfg.b_t * cfg.b_t * (cfg.io_dtype.width // 8)) >> 4 + diag_seg = ((cfg.d_k // 16) * 256 * (cfg.io_dtype.width // 8)) >> 4 + d_do_amaj0 = sDo_amaj[0].desc() + d_qd_trans0 = sQ_decay_trans[0].desc() + d_kd_trans0 = sK_decay_trans[0].desc() + d_ki_amaj0 = sK_inv_amaj[0].desc() + d_int0 = sIntermediate[0].desc() + d_kd_lead0 = sK_decay_lead16[0].desc() + d_do_lead0 = sDo_lead16[0].desc() + d_kr_lead0 = sK_restore_lead16[0].desc() + d_diag0 = sState_scale_diag[0].desc() + d_dstate_alt0 = sDstate_alt[0].desc() + d_u_lead0 = sU_lead16[0].desc() + d_dy_lead0 = sDy_lead16[0].desc() + assert cfg.smem_state_stages == 1 + d_state_alt0 = sState_alt[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 + for rev_idx in cutlass.range(sk_nt, unroll=1): + gc = gbase + rev_idx + decay_stage = gc % cfg.smem_decay_stages + intermediate_stage = gc % cfg.smem_intermediate_stages + decay_phase = (gc // cfg.smem_decay_stages) % 2 + intermediate_phase = (gc // cfg.smem_intermediate_stages) % 2 + has_dstate = cutlass.Boolean(rev_idx > 0) + if cutlass.const_expr(cfg.use_dstate_in): + has_dstate = cutlass.Boolean(True) + raw_stage_idx = gc % cfg.smem_raw_stages + + # ---- stage-derived operand descriptors ------------------------------- + decay_op_off = decay_stage * op_seg + d_do_amaj = d_do_amaj0 + raw_stage_idx * do_seg + d_qd_trans = d_qd_trans0 + decay_op_off + d_kd_trans = d_kd_trans0 + decay_op_off + d_ki_amaj = d_ki_amaj0 + decay_op_off + d_int = d_int0 + intermediate_stage * intermediate_seg + d_int_tinv = d_int + intermediate_slot + d_int_da = d_int + 2 * intermediate_slot + d_int_dm = d_int + 3 * intermediate_slot + d_int_ndm = d_int + 4 * intermediate_slot + chunk_idx = cend - cutlass.Int32(1) - rev_idx + + # ---- state_k = state(S) @ K_decay^T ---------------------------------- + bars.mb_k_decay_inv_ready[decay_stage].wait(decay_phase) + 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, + d_kd_lead0 + decay_op_off, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_k_acc_offset), cutlass.Float32), + accumulate=False, + ) + if elect_one: + bars.mb_state_k_acc_ready.arrive(cta_group=1) + bars.mb_state_done[state_index.idx].arrive(cta_group=1) + state_index = advance(state_index, cfg.smem_state_stages) + + # ---- dQ inter = state(T) @ dO^T -------------------------------------- + bars.mb_dqk_acc_done.wait(parts_done_index.phase) + parts_done_index = advance(parts_done_index, 1) + bars.mb_state_inp_ready[gc % 2].wait((gc // 2) % 2) + bars.mb_do_ready[raw_stage_idx].wait((gc // cfg.smem_raw_stages) % 2) + if chunk_idx >= FIRST_STATE_CHUNK: + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_inp_offset + (gc % 2) * (cfg.d_v // 2)), cutlass.Int8) + b_desc = d_do_lead0 + raw_stage_idx * do_seg + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_state_ts_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_ts_desc.sps_B): + mma_ts_step( + bmm_state_ts_desc, + a_ptr.subview(sub * bmm_state_ts_desc.sps_B * bmm_state_ts_desc.tmem_advance_A), + b_desc + sub * (bmm_state_ts_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + + # ---- dU inter = dstate_inp(T) @ K_restore ---------------------------- + bars.mb_q_decay_k_restore_ready[decay_stage].wait(decay_phase) + if has_dstate: + bars.mb_dstate_inp_ready.wait(dstate_inp_index.phase) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_inp_offset), cutlass.Int8) + b_desc = d_kr_lead0 + decay_op_off + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_dvinter_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dvinter_desc.sps_B): + mma_ts_step( + bmm_dvinter_desc, + a_ptr.subview(sub * bmm_dvinter_desc.sps_B * bmm_dvinter_desc.tmem_advance_A), + b_desc + sub * (bmm_dvinter_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + + # ---- dstate decay = dstate_inp(T) @ diag(eGl) ------------------------ + if has_dstate: + desc_diag = d_diag0 + decay_stage * diag_seg + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_inp_offset) + k_block * 8, cutlass.Int8) + b_desc = desc_diag.advance_start_address(k_block * 256 * 2) + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset) + k_block * 16, cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_diag_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_diag_desc.sps_B): + mma_ts_step( + bmm_diag_desc, + a_ptr.subview(sub * bmm_diag_desc.sps_B * bmm_diag_desc.tmem_advance_A), + b_desc + sub * (bmm_diag_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + dstate_inp_index = advance(dstate_inp_index, 1) + + # ---- dU intra += dO^T(S) @ A ----------------------------------------- + bars.mb_a_ready[intermediate_stage].wait(intermediate_phase) + mma_ss( + bmm_du_at_desc, + d_do_amaj, + d_int, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_acc_offset), cutlass.Float32), + accumulate=has_dstate, + ) + if elect_one: + bars.mb_du_acc_ready.arrive(cta_group=1) + bars.mb_a_done[intermediate_stage].arrive(cta_group=1) + + # ---- dstate Q-term += dO^T(S) @ Q_decay ------------------------------ + mma_ss( + bmm_dstate_q_at_desc, + d_do_amaj, + d_qd_trans, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset), cutlass.Float32), + accumulate=has_dstate, + ) + + # ---- U = Y(T) @ T_inv ------------------------------------------------ + bars.mb_t_inv_ready[intermediate_stage].wait(intermediate_phase) + bars.mb_y_inp_ready.wait(y_inp_index.phase) + y_inp_index = advance(y_inp_index, 1) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_y_inp_offset), cutlass.Int8) + b_desc = d_int_tinv + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_u_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_qk_ts_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_qk_ts_desc.sps_B): + mma_ts_step( + bmm_qk_ts_desc, + a_ptr.subview(sub * bmm_qk_ts_desc.sps_B * bmm_qk_ts_desc.tmem_advance_A), + b_desc + sub * (bmm_qk_ts_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + if elect_one: + bars.mb_u_acc_ready.arrive(cta_group=1) + bars.mb_do_mma_done[raw_stage_idx].arrive(cta_group=1) + + # ---- dY = dU(T) @ T_inv ---------------------------------------------- + bars.mb_du_inp_ready.wait(du_inp_index.phase) + du_inp_index = advance(du_inp_index, 1) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_inp_offset), cutlass.Int8) + b_desc = d_int_tinv + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dy_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_qk_ts_t_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_qk_ts_t_desc.sps_B): + mma_ts_step( + bmm_qk_ts_t_desc, + a_ptr.subview(sub * bmm_qk_ts_t_desc.sps_B * bmm_qk_ts_t_desc.tmem_advance_A), + b_desc + sub * (bmm_qk_ts_t_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + if elect_one: + bars.mb_dy_acc_ready.arrive(cta_group=1) + bars.mb_t_inv_done[intermediate_stage].arrive(cta_group=1) + + # ---- dK_restore part = dstate(S) @ U^T ------------------------------- + bars.mb_u_smem_ready.wait(u_smem_index.phase) + u_smem_index = advance(u_smem_index, 1) + if has_dstate: + bars.mb_dstate_smem_ready.wait(dstate_smem_index.phase) + dstate_smem_index = advance(dstate_smem_index, 1) + mma_ss( + bmm_dstate_at_desc, + d_dstate_alt0, + d_u_lead0, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_restore_acc_offset), cutlass.Float32), + accumulate=False, + ) + if elect_one: + bars.mb_dk_restore_part_acc_ready.arrive(cta_group=1) + bars.mb_dstate_smem_done.arrive(cta_group=1) + + # ---- dstate K-term += -dY(T) @ K_decay ------------------------------- + bars.mb_neg_dy_inp_ready.wait(neg_dy_index.phase) + neg_dy_index = advance(neg_dy_index, 1) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_neg_dy_inp_offset), cutlass.Int8) + b_desc = d_kd_trans + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_dstate_k_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_k_desc.sps_B): + mma_ts_step( + bmm_dstate_k_desc, + a_ptr.subview(sub * bmm_dstate_k_desc.sps_B * bmm_dstate_k_desc.tmem_advance_A), + b_desc + sub * (bmm_dstate_k_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(True), + ) + if elect_one: + bars.mb_dstate_acc_ready.arrive(cta_group=1) + + # ---- dK_inv part = scale.Q_decay^T(S) @ dA --------------------------- + bars.mb_da_ready[intermediate_stage].wait(intermediate_phase) + mma_ss( + bmm_dgrad_at_t_desc, + d_qd_trans, + d_int_da, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), + accumulate=False, + ) + + # ---- dQ attn += K_inv^T(S) @ dA^T ------------------------------------ + if chunk_idx >= FIRST_STATE_CHUNK: + mma_ss( + bmm_dgrad_at_desc, + d_ki_amaj, + d_int_da, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32), + accumulate=True, + ) + if chunk_idx < FIRST_STATE_CHUNK: + mma_ss( + bmm_dgrad_at_desc, + d_ki_amaj, + d_int_da, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32), + accumulate=False, + ) + if elect_one: + bars.mb_dq_acc_ready.arrive(cta_group=1) + bars.mb_da_done[intermediate_stage].arrive(cta_group=1) + + # ---- dK_decay part = state(T) @ dY^T --------------------------------- + if chunk_idx >= FIRST_STATE_CHUNK: + bars.mb_dy_smem_ready.wait(gc % 2) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_inp_offset + (gc % 2) * (cfg.d_v // 2)), cutlass.Int8) + b_desc = d_dy_lead0 + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_state_ts_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_ts_desc.sps_B): + mma_ts_step( + bmm_state_ts_desc, + a_ptr.subview(sub * bmm_state_ts_desc.sps_B * bmm_state_ts_desc.tmem_advance_A), + b_desc + sub * (bmm_state_ts_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + if elect_one: + bars.mb_dy_smem_done.arrive(cta_group=1) + bars.mb_state_inp_done[gc % 2].arrive(cta_group=1) + + # ---- dK_inv part += K_decay^T(S) @ -dM_strict ------------------------ + bars.mb_dm_ready[intermediate_stage].wait(intermediate_phase) + mma_ss( + bmm_dgrad_at_t_desc, + d_kd_trans, + d_int_ndm, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), + accumulate=True, + ) + if elect_one: + bars.mb_dk_inv_part_acc_ready.arrive(cta_group=1) + + # ---- dK_decay part += K_inv^T(S) @ dM_strict^T ----------------------- + if chunk_idx >= FIRST_STATE_CHUNK: + mma_ss( + bmm_dgrad_at_desc, + d_ki_amaj, + d_int_dm, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), + accumulate=True, + ) + if chunk_idx < FIRST_STATE_CHUNK: + mma_ss( + bmm_dgrad_at_desc, + d_ki_amaj, + d_int_dm, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), + accumulate=False, + ) + if elect_one: + bars.mb_dk_decay_part_acc_ready.arrive(cta_group=1) + bars.mb_dm_done[intermediate_stage].arrive(cta_group=1) + bars.mb_decay_done[decay_stage].arrive(cta_group=1) + + # ---- tile end: WG1's dstate0 drain gates the next tile's dstate reuse ---- + bars.mb_dstate0_acc_stored.wait(dstate0_index.phase) + dstate0_index = advance(dstate0_index, 1) + gbase += sk_nt + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + bars.mb_tmem_done[0].wait(0) + nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) + nvvm.tcgen05_dealloc( + nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), + cutlass.Int32(512), + group=nvvm.CTAGroup.CTA_1, + ) + + +@cute.jit +def tmaldg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + mSched, + sSched, + q_tx_bytes, + k_tx_bytes, + gate_tx_bytes, + beta_tx_bytes, + do_tx_bytes, + v_tx_bytes, + w_tx_bytes, + sQ_raw, + sK_raw, + sV_raw, + sGate_raw, + sDo_raw, + sBeta_raw, + sW_raw, + sState_raw, + desc_q_base, + desc_k_base, + desc_v_base, + desc_gate_base, + desc_do_base, + 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 + G->S operand load.""" + elect_one = nvvm.elect_sync() + + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + + sQ_tma = SmemTile( + base=sQ_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sK_tma = SmemTile( + base=sK_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sV_tma = SmemTile( + base=sV_raw, + elems_per_stage=(cfg.d_v * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sGate_tma = SmemTile( + base=sGate_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 32), + tma_granu_elems=32, + tma_subtile_stride_elems=(cfg.b_t * 32), + ) + sDo_tma = SmemTile( + base=sDo_raw, + elems_per_stage=(cfg.d_v * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sBeta_tma = SmemTile( + base=sBeta_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sW_tma = SmemTile( + base=sW_raw, + elems_per_stage=(cfg.d_v * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sState_tma = SmemTile( + base=sState_raw, + elems_per_stage=(cfg.d_k * cfg.d_v), + stages=cfg.smem_state_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.d_k * 64, + ) + raw_index = PipelineState.start(phase=1) + state_index = PipelineState.start(phase=1) + 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) + head_o = head_idx + head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) + head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) + head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) + slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_q_slot = (desc_q_base + slot).tospace(cutlass.AddressSpace.generic) + desc_k_slot = (desc_k_base + slot).tospace(cutlass.AddressSpace.generic) + desc_v_slot = (desc_v_base + slot).tospace(cutlass.AddressSpace.generic) + desc_gate_slot = (desc_gate_base + slot).tospace(cutlass.AddressSpace.generic) + desc_do_slot = (desc_do_base + slot).tospace(cutlass.AddressSpace.generic) + 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) + tma_tensormap_acquire(desc_v_slot) + tma_tensormap_acquire(desc_gate_slot) + tma_tensormap_acquire(desc_do_slot) + 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 + chunk_start = chunk_idx * cfg.b_t + + # ---- Q load ---------------------------------------------------------- + bars.mb_q_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_q_ready[raw_index.idx].arrive(n_bytes=q_tx_bytes) + q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, chunk_start) + tma_load_tile(sQ_tma[raw_index.idx], q_slice, bars.mb_q_ready[raw_index.idx].smem_ptr, acquire=False) + + # ---- K load ---------------------------------------------------------- + bars.mb_k_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_k_ready[raw_index.idx].arrive(n_bytes=k_tx_bytes) + k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) + tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_index.idx].smem_ptr, acquire=False) + + # ---- Gate load ------------------------------------------------------- + bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_gate_ready[raw_index.idx].arrive(n_bytes=gate_tx_bytes) + gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_index.idx].smem_ptr, acquire=False) + + # ---- Beta load ------------------------------------------------------- + bars.mb_beta_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_beta_ready[raw_index.idx].arrive(n_bytes=beta_tx_bytes) + 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 -- + if chunk_idx >= FIRST_STATE_CHUNK: + state_idx = state_index.idx + bars.mb_state_cg0_done[state_idx].wait(state_index.phase) + bars.mb_state_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(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) + + # ---- dO load --------------------------------------------------------- + bars.mb_do_done[raw_index.idx].wait(raw_index.phase) + bars.mb_do_mma_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_do_ready[raw_index.idx].arrive(n_bytes=do_tx_bytes) + do_slice = tma_slice_runtime_desc(desc_do_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sDo_tma[raw_index.idx], do_slice, bars.mb_do_ready[raw_index.idx].smem_ptr, acquire=False) + + # ---- V load ---------------------------------------------------------- + bars.mb_v_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_v_ready[raw_index.idx].arrive(n_bytes=v_tx_bytes) + 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) + + # ---- W load ---------------------------------------------------------- + bars.mb_w_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_w_ready[raw_index.idx].arrive(n_bytes=w_tx_bytes) + w_slice = tma_slice_runtime_desc(desc_w_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sW_tma[raw_index.idx], w_slice, bars.mb_w_ready[raw_index.idx].smem_ptr, acquire=False) + raw_index = advance(raw_index, cfg.smem_raw_stages) + tile_idx = next_tile + + +@cute.jit +def compute0_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_hold, + warp_idx, + scale, + sK_inv_raw, + sGate_raw, + sK_raw, + sQ_raw, + sState_raw, + sV_raw, + sDo_raw, + sBeta_raw, + sW_raw, + sNorm_raw, + sK_decay_raw, + sK_restore_raw, + sQ_decay_raw, + sState_scale_diag_raw, + bars, +) -> None: + """WG0 warp role (warps 0-3): persistent tile-scheduler loop + gate prefix + scan and the decay/restore operand materialization into tcgen05 SMEM for + EVERY chunk (no ping-pong: the backward pipeline is drain-bound). Also + stashes the per-row Q/K inverse norms for WG2's dGate assembly and copies + H -> TMEM f16 at the chunk tail.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_hold.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + tmem_sp = warp_idx % (cfg.d_v // cfg.threads_per_warp) + value_dim = tmem_sp * cfg.threads_per_warp + lane + state_copy_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 + state_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 + for rev_idx in cutlass.range(sk_nt, unroll=1): + chunk_idx = cend - cutlass.Int32(1) - rev_idx + gc = gbase + rev_idx + chunk_start = chunk_idx * cfg.b_t + decay_stage = gc % cfg.smem_decay_stages + raw_stage = gc % cfg.smem_raw_stages + sQ_ptr = sQ_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sK_ptr = sK_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sBetaP_ptr = sBeta_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sWP_ptr = sW_raw.data_ptr() + raw_stage * (cfg.d_v * cfg.b_t) + sGate_ptr = sGate_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sQ_decay_ptr = sQ_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sK_restore_ptr = sK_restore_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sState_scale_diag_ptr = sState_scale_diag_raw.data_ptr() + decay_stage * ((cfg.d_k // 16) * 256) + + bars.mb_gate_ready[raw_stage].wait((gc // cfg.smem_raw_stages) % 2) + bars.mb_q_ready[raw_stage].wait((gc // cfg.smem_raw_stages) % 2) + bars.mb_k_ready[raw_stage].wait((gc // cfg.smem_raw_stages) % 2) + bars.mb_beta_ready[raw_stage].wait((gc // cfg.smem_raw_stages) % 2) + + row_group_start = cg0_warp * (cfg.b_t // len(cfg.compute_group_0_warp_ids)) + lane_row_group = lane // 8 + lane_in_row_group = lane - lane_row_group * 8 + decay_row = row_group_start + lane_row_group + + g_prefix_ptr = sGate_ptr + prefix_dim = cg0_warp * cfg.threads_per_warp + lane + # ---- gate prefix scan: cumulative log-gate per key channel ----------- + gate_raw = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + for row in cutlass.range_constexpr(cfg.b_t): + f32_segment = prefix_dim // 32 + f32_segment_dim = prefix_dim - f32_segment * 32 + prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) + gate_raw[row] = (sGate_ptr + prefix_idx).load() + g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + for row in cutlass.range_constexpr(cfg.b_t): + gate = gate_raw[row] + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) + if token_idx < seqlen_b: + gate = gate * cutlass.Float32(LOG2_E) + else: + gate = cutlass.Float32(0.0) + g_prefix_regs[row] = gate + + prefix_acc = cutlass.Float32(0.0) + for row_pair in cutlass.range_constexpr(cfg.b_t // 2): + row0 = row_pair * 2 + row1 = row0 + 1 + gate0 = g_prefix_regs[row0] + gate1 = g_prefix_regs[row1] + pair_vec = nvvm.add_packed_f32x2( + cutlass.Vector.from_elements((prefix_acc, gate0), cutlass.Float32), + cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32), + ftz=False, + rnd="rn", + ) + prefix0, row_pair_sum = cutlass.Float32(pair_vec[0]), cutlass.Float32(pair_vec[1]) + prefix1 = prefix_acc + row_pair_sum + g_prefix_regs[row0] = prefix0 + g_prefix_regs[row1] = prefix1 + prefix_acc = prefix1 + + for row in cutlass.range_constexpr(cfg.b_t): + g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) + + exp_g_last = g_prefix_regs[cfg.b_t - 1] + # ---- decay-slot guard: previous use fully consumed ------------------- + operand_done_phase = ((gc // cfg.smem_decay_stages) + 1) % 2 + bars.mb_decay_done[decay_stage].wait(operand_done_phase) + + for row in cutlass.range_constexpr(cfg.b_t): + f32_segment = prefix_dim // 32 + f32_segment_dim = prefix_dim - f32_segment * 32 + prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) + (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) + + # ---- state-scale diag: stage exp2(g_last) decay blocks --------------- + block = prefix_dim // cutlass.Int32(16) + coord = prefix_dim - block * cutlass.Int32(16) + linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + coord + diag_idx = swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) + sState_scale_diag_ptr[diag_idx] = exp_g_last.to(cfg.io_dtype) + + # ---- raw Q/K: SMEM -> TMEM ring (channel-major, for WG2) ------------- + qk_raw_stage = gc % cfg.tmem_qk_raw_stages + bars.mb_qk_raw_done[qk_raw_stage].wait(((gc // cfg.tmem_qk_raw_stages) + 1) % 2) + raw_seg = prefix_dim // 64 + raw_dim = prefix_dim - raw_seg * 64 + q_raw_words = cutlass.Array(cutlass.Int32, cfg.b_t // 2, alignment=16) + k_raw_words = cutlass.Array(cutlass.Int32, cfg.b_t // 2, alignment=16) + for t2 in cutlass.range_constexpr(cfg.b_t // 2): + t0 = 2 * t2 + ridx0 = raw_seg * (cfg.b_t * 64) + t0 * 64 + swizzle_xor_128b(t0, raw_dim, elem_bytes=2) + ridx1 = raw_seg * (cfg.b_t * 64) + (t0 + 1) * 64 + swizzle_xor_128b(t0 + 1, raw_dim, elem_bytes=2) + q0 = (sQ_ptr + ridx0).load().to(cutlass.Float32) + q1 = (sQ_ptr + ridx1).load().to(cutlass.Float32) + k0 = (sK_ptr + ridx0).load().to(cutlass.Float32) + k1 = (sK_ptr + ridx1).load().to(cutlass.Float32) + q_raw_words[t2] = fp32_to_fp16(q0, q1, dtype=cfg.io_dtype) + k_raw_words[t2] = fp32_to_fp16(k0, k1, dtype=cfg.io_dtype) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(state_copy_addr + (tmem_col + cfg.tmem_qraw_inp_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), + q_raw_words[0 : (cfg.b_t // 2)], + ) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(state_copy_addr + (tmem_col + cfg.tmem_kraw_inp_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), + k_raw_words[0 : (cfg.b_t // 2)], + ) + nvvm.tcgen05_wait("store") + bars.mb_qk_raw_ready[qk_raw_stage].arrive() + + nvvm.barrier_cta_sync(cfg.cg0_sync_barrier_id, thread_count=cfg.cg0_threads) + + k_inv_pack = cutlass.Array(cutlass.Int32, 2 * 4, alignment=16) + raw_q_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + raw_beta_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + # ---- optional Q/K L2-norm -------------------------------------------- + if cutlass.const_expr(cfg.l2norm): + qk0_lo = opaque_f32_zero() + qk0_hi = opaque_f32_zero() + qk1_lo = opaque_f32_zero() + qk1_hi = opaque_f32_zero() + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + raw_f16_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + raw_q_frag = (sQ_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_k_frag = (sK_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_beta_frag = (sBetaP_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_q_frag_f32 = raw_q_frag.to(cutlass.Float32) + raw_k_frag_f32 = raw_k_frag.to(cutlass.Float32) + raw_beta_frag_f32 = raw_beta_frag.to(cutlass.Float32) + for dim_offset in cutlass.range_constexpr(8): + q_val = raw_q_frag_f32[dim_offset] + k_val = raw_k_frag_f32[dim_offset] + raw_q_regs[reg_base + dim_offset] = q_val + raw_k_regs[reg_base + dim_offset] = k_val + raw_beta_regs[reg_base + dim_offset] = raw_beta_frag_f32[dim_offset] + if cutlass.const_expr(cfg.l2norm): + if cutlass.const_expr(dim_offset % 2 == 0): + qk0_lo, qk0_hi = ffma2(q_val, k_val, q_val, k_val, qk0_lo, qk0_hi) + else: + qk1_lo, qk1_hi = ffma2(q_val, k_val, q_val, k_val, qk1_lo, qk1_hi) + + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_q_done[raw_stage].arrive() + bars.mb_k_done[raw_stage].arrive() + bars.mb_beta_done[raw_stage].arrive() + + q_inv_norm = opaque_f32_zero() + cutlass.Float32(1.0) + k_inv_norm = opaque_f32_zero() + cutlass.Float32(1.0) + if cutlass.const_expr(cfg.l2norm): + q_sum_sq = qk0_lo + qk1_lo + k_sum_sq = qk0_hi + qk1_hi + q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 4, 31, kind=nvvm.Shfl.BFLY)) + q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 2, 31, kind=nvvm.Shfl.BFLY)) + q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 1, 31, kind=nvvm.Shfl.BFLY)) + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 4, 31, kind=nvvm.Shfl.BFLY)) + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 2, 31, kind=nvvm.Shfl.BFLY)) + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 1, 31, kind=nvvm.Shfl.BFLY)) + norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS) + q_inv_norm = cute.math.rsqrt(cute.math.max(q_sum_sq, norm_floor_sq), fastmath=True) + k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) + if lane_in_row_group == 0: + sNorm_raw[(gc % cfg.tmem_qk_raw_stages) * (2 * cfg.b_t) + decay_row] = q_inv_norm + sNorm_raw[(gc % cfg.tmem_qk_raw_stages) * (2 * cfg.b_t) + cfg.b_t + decay_row] = k_inv_norm + q_stage_norm = q_inv_norm * scale + + # ---- decay/restore operands: exp2(+-g) applied per key channel ------- + exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + for f32_group in cutlass.range_constexpr(2): + f32_dim_base = dim_base + f32_group * 4 + f32_segment = f32_dim_base // 32 + f32_segment_dim = f32_dim_base - f32_segment * 32 + g_prefix_idx = f32_segment * (cfg.b_t * 32) + decay_row * 32 + swizzle_xor_128b(decay_row, f32_segment_dim, elem_bytes=4) + exp_g_frag = (g_prefix_ptr + g_prefix_idx).load(count=4, alignment=16) + exp_g_last_idx = f32_segment * (cfg.b_t * 32) + (cfg.b_t - 1) * 32 + swizzle_xor_128b((cfg.b_t - 1), f32_segment_dim, elem_bytes=4) + exp_g_last_frag = (g_prefix_ptr + exp_g_last_idx).load(count=4, alignment=16) + f32_reg_base = reg_base + f32_group * 4 + exp_g_regs[f32_reg_base] = exp_g_frag[0] + exp_g_regs[f32_reg_base + 1] = exp_g_frag[1] + exp_g_regs[f32_reg_base + 2] = exp_g_frag[2] + exp_g_regs[f32_reg_base + 3] = exp_g_frag[3] + exp_g_last_regs[f32_reg_base] = exp_g_last_frag[0] + exp_g_last_regs[f32_reg_base + 1] = exp_g_last_frag[1] + exp_g_last_regs[f32_reg_base + 2] = exp_g_last_frag[2] + exp_g_last_regs[f32_reg_base + 3] = exp_g_last_frag[3] + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_gate_done[raw_stage].arrive() + + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + # ---- K decay + K inv operands: exp2(g) * Beta * K / exp2(-g) * K - + k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + for pair_idx in cutlass.range_constexpr(4): + dim0 = pair_idx * 2 + dim1 = dim0 + 1 + raw_reg_idx0 = reg_base + dim0 + raw_reg_idx1 = reg_base + dim1 + k_value0, k_value1 = fmul2(raw_k_regs[raw_reg_idx0], raw_k_regs[raw_reg_idx1], k_inv_norm, k_inv_norm) + k_beta0, k_beta1 = fmul2(k_value0, k_value1, raw_beta_regs[raw_reg_idx0], raw_beta_regs[raw_reg_idx1]) + k_pair = fp32_to_fp16(k_beta0, k_beta1, dtype=cfg.io_dtype) + exp_g_pair = fp32_to_fp16(exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1], dtype=cfg.io_dtype) + k_decay_pack[pair_idx] = mul_f16x2(k_pair, exp_g_pair, cfg.io_dtype) + exp_neg_g0 = cute.math.rcp(exp_g_regs[raw_reg_idx0], approx=True, ftz=True) + exp_neg_g1 = cute.math.rcp(exp_g_regs[raw_reg_idx1], approx=True, ftz=True) + exp_neg_pair = fp32_to_fp16(exp_neg_g0, exp_neg_g1, dtype=cfg.io_dtype) + k_norm_pair = fp32_to_fp16(k_value0, k_value1, dtype=cfg.io_dtype) + k_inv_pack[dim_half * 4 + pair_idx] = mul_f16x2(k_norm_pair, exp_neg_pair, cfg.io_dtype) + + k_inv_vec = cutlass.Vector.from_elements( + ( + k_inv_pack[dim_half * 4], + k_inv_pack[dim_half * 4 + 1], + k_inv_pack[dim_half * 4 + 2], + k_inv_pack[dim_half * 4 + 3], + ), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + k_decay_vec = cutlass.Vector.from_elements( + (k_decay_pack[0], k_decay_pack[1], k_decay_pack[2], k_decay_pack[3]), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + op_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + (sK_inv_ptr + op_idx).store(k_inv_vec, alignment=16) + (sK_decay_ptr + op_idx).store(k_decay_vec, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_k_decay_inv_ready[decay_stage].arrive() + + # ---- Q decay + K_restore operands ------------------------------------ + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + q_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + k_restore_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + for pair_idx in cutlass.range_constexpr(4): + dim0 = pair_idx * 2 + dim1 = dim0 + 1 + raw_reg_idx0 = reg_base + dim0 + raw_reg_idx1 = reg_base + dim1 + q_value0, q_value1 = fmul2(raw_q_regs[raw_reg_idx0], raw_q_regs[raw_reg_idx1], q_stage_norm, q_stage_norm) + q_pair = fp32_to_fp16(q_value0, q_value1, dtype=cfg.io_dtype) + exp_g_pair = fp32_to_fp16(exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1], dtype=cfg.io_dtype) + q_decay_pack[pair_idx] = mul_f16x2(q_pair, exp_g_pair, cfg.io_dtype) + exp_g_last_pair = fp32_to_fp16(exp_g_last_regs[raw_reg_idx0], exp_g_last_regs[raw_reg_idx1], dtype=cfg.io_dtype) + k_restore_pack[pair_idx] = mul_f16x2(k_inv_pack[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) + + q_decay_vec = cutlass.Vector.from_elements( + (q_decay_pack[0], q_decay_pack[1], q_decay_pack[2], q_decay_pack[3]), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + k_restore_vec = cutlass.Vector.from_elements( + (k_restore_pack[0], k_restore_pack[1], k_restore_pack[2], k_restore_pack[3]), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + op_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + (sQ_decay_ptr + op_idx).store(q_decay_vec, alignment=16) + (sK_restore_ptr + op_idx).store(k_restore_vec, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_q_decay_k_restore_ready[decay_stage].arrive() + + # ---- state copy: SMEM -> TMEM f16 ------------------------------------ + bars.mb_state_inp_done[gc % 2].wait(((gc // 2) + 1) % 2) + bars.mb_state_inp_cg2_done[gc % 2].wait(((gc // 2) + 1) % 2) + if chunk_idx >= FIRST_STATE_CHUNK: + bars.mb_state_ready[state_index.idx].wait(state_index.phase) + 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 + ) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr( + state_copy_addr + (tmem_col + cfg.tmem_state_inp_offset + (gc % 2) * (cfg.d_v // 2) + pl * 32 + g8 * 4), cutlass.Int8 + ), + state_frag.bitcast(cutlass.Int32), + ) + nvvm.tcgen05_wait("store") + bars.mb_state_cg0_done[state_index.idx].arrive() + state_index = advance(state_index, cfg.smem_state_stages) + bars.mb_state_inp_ready[gc % 2].arrive() + gbase += sk_nt + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def compute1_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_hold, + warp_idx, + mDstate0, + mDstate_in, + sV_raw, + sW_raw, + sDo_raw, + sU_raw, + sDy_raw, + sDv_raw, + sDwOut_raw, + sDstate_raw, + bars, +) -> None: + """WG1 warp role (warps 4-7): the value-side TMEM staging.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_hold.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + tmem_sp = warp_idx % (cfg.d_v // cfg.threads_per_warp) + ov_tok = (lane // 16) * 8 + (lane & 7) + ov_col = ((lane // 8) & 1) * 8 + value_dim = tmem_sp * cfg.threads_per_warp + lane + value_dim_base = tmem_sp * cfg.threads_per_warp + cg1_tidx = warp_idx % 4 * cfg.threads_per_warp + lane + + raw_index = PipelineState.start(phase=0) + state_k_index = PipelineState.start(phase=0) + u_acc_index = PipelineState.start(phase=0) + du_acc_index = PipelineState.start(phase=0) + dy_acc_index = PipelineState.start(phase=0) + sdy_done_index = PipelineState.start(phase=1) + dstate_ready_index = PipelineState.start(phase=0) + dstate_smem_done_index = PipelineState.start(phase=1) + dv_done_index = PipelineState.start(phase=1) + dwo_done_index = PipelineState.start(phase=1) + + 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 + + # ---- dht seeding: dH acc + dh_inp f16 + sdH ------------------------------ + if cutlass.const_expr(cfg.use_dstate_in): + if sk_nt > 0: + seed_true = cend == num_chunks_b + bars.mb_dstate_smem_done.wait(dstate_smem_done_index.phase) + 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 + 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 + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 16), cutlass.Float32), + seed_block[0:16], + ) + seed_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) + for pc in cutlass.range_constexpr(8): + seed_pack[pc] = fp32_to_fp16(seed_block[2 * pc], seed_block[2 * pc + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 8), cutlass.Int8), + seed_pack[0:8], + ) + nvvm.tcgen05_wait("store") + bars.mb_dstate_inp_ready.arrive() + + # ---- dht seed -> sdH: re-read dh_inp after the TMEM publish ------ + for sub in cutlass.range_constexpr(cfg.d_k // 16): + dstate_words = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 8), cutlass.Float32), num=8 + ) + for half in cutlass.range_constexpr(2): + d_base = sub * 16 + half * 8 + h_pack = cutlass.Vector.from_elements( + (dstate_words[half * 4], dstate_words[half * 4 + 1], dstate_words[half * 4 + 2], dstate_words[half * 4 + 3]), + cutlass.Float32, + ).bitcast(cfg.io_dtype) + h_addr = (d_base // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, d_base % 64, elem_bytes=2) + (sDstate_raw.data_ptr() + h_addr).store(h_pack, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dstate_smem_ready.arrive() + + for rev_idx in cutlass.range(sk_nt, unroll=1): + chunk_idx = cend - cutlass.Int32(1) - rev_idx + gc = gbase + rev_idx + has_dstate = cutlass.Boolean(rev_idx > 0) + if cutlass.const_expr(cfg.use_dstate_in): + has_dstate = cutlass.Boolean(True) + + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + row_addr_lo = tmem_row << 16 + row_addr_hi = (tmem_row + 16) << 16 + row_id0 = tmem_row + value_dim_base + row_id1 = row_id0 + 16 + + # ---- Y staging: Y = W*V - state_k -> TMEM f16 ------------------------ + bars.mb_v_ready[raw_index.idx].wait((gc // cfg.smem_raw_stages) % 2) + bars.mb_w_ready[raw_index.idx].wait((gc // cfg.smem_raw_stages) % 2) + projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset + input_col_id = tmem_col + cfg.tmem_y_inp_offset + raw_v_frag0 = nvvm.ldmatrix( + sV_ptr + + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) + + ov_tok * 64 + + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + raw_v_frag1 = nvvm.ldmatrix( + sV_ptr + + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) + + ov_tok * 64 + + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + raw_w_frag0 = nvvm.ldmatrix( + sW_ptr + + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) + + ov_tok * 64 + + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + raw_w_frag1 = nvvm.ldmatrix( + sW_ptr + + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) + + ov_tok * 64 + + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + + y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + if chunk_idx >= FIRST_STATE_CHUNK: + bars.mb_state_k_acc_ready.wait(state_k_index.phase) + state_k_index = advance(state_k_index, 1) + state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) + state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + state_k_pair = fp32_to_fp16(state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1], dtype=cfg.io_dtype) + wv_pair = mul_f16x2(raw_w_frag0[raw_matrix], raw_v_frag0[raw_matrix], cfg.io_dtype) + y_inp_pack0[reg_idx ^ 2] = sub_f16x2(wv_pair, state_k_pair, cfg.io_dtype) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + state_k_pair = fp32_to_fp16(state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1], dtype=cfg.io_dtype) + wv_pair = mul_f16x2(raw_w_frag1[raw_matrix], raw_v_frag1[raw_matrix], cfg.io_dtype) + y_inp_pack1[reg_idx ^ 2] = sub_f16x2(wv_pair, state_k_pair, cfg.io_dtype) + if chunk_idx < FIRST_STATE_CHUNK: + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + y_inp_pack0[reg_idx ^ 2] = mul_f16x2(raw_w_frag0[raw_matrix], raw_v_frag0[raw_matrix], cfg.io_dtype) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + y_inp_pack1[reg_idx ^ 2] = mul_f16x2(raw_w_frag1[raw_matrix], raw_v_frag1[raw_matrix], cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) + nvvm.tcgen05_wait("store") + bars.mb_y_inp_ready.arrive() + + # ---- dU restage: dU acc -> TMEM f16 A operand ------------------------ + bars.mb_du_acc_ready.wait(du_acc_index.phase) + du_acc_index = advance(du_acc_index, 1) + du_col_id = tmem_col + cfg.tmem_du_acc_offset + du_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + du_col_id, cutlass.Float32), num=2) + du_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + du_col_id, cutlass.Float32), num=2) + + du_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + du_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + frag_pair = reg_idx * 2 + du_pack0[reg_idx] = fp32_to_fp16(du_vec0[frag_pair], du_vec0[frag_pair + 1], dtype=cfg.io_dtype) + du_pack1[reg_idx] = fp32_to_fp16(du_vec1[frag_pair], du_vec1[frag_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + (tmem_col + cfg.tmem_du_inp_offset), cutlass.Int8), du_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + (tmem_col + cfg.tmem_du_inp_offset), cutlass.Int8), du_pack1[0:4]) + nvvm.tcgen05_wait("store") + bars.mb_du_inp_ready.arrive() + + # ---- U readback -> sU ------------------------------------------------ + bars.mb_u_acc_ready.wait(u_acc_index.phase) + u_acc_index = advance(u_acc_index, 1) + u_col_id = tmem_col + cfg.tmem_u_acc_offset + u_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + u_col_id, cutlass.Float32), num=2) + u_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + u_col_id, cutlass.Float32), num=2) + + u_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + u_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + u_pack0[reg_idx] = fp32_to_fp16(u_vec0[2 * reg_idx], u_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) + u_pack1[reg_idx] = fp32_to_fp16(u_vec1[2 * reg_idx], u_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) + nvvm.stmatrix( + sU_raw.data_ptr() + + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) + + ov_tok * 64 + + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + u_pack0.data_ptr().load(count=4, alignment=4), + nvvm.MMALayout.COL, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.stmatrix( + sU_raw.data_ptr() + + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) + + ov_tok * 64 + + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + u_pack1.data_ptr().load(count=4, alignment=4), + nvvm.MMALayout.COL, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_u_smem_ready.arrive() + + # ---- dY readback ------------------------------------------------------- + bars.mb_dy_acc_ready.wait(dy_acc_index.phase) + dy_acc_index = advance(dy_acc_index, 1) + dy_col_id = tmem_col + cfg.tmem_dy_acc_offset + dy_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + dy_col_id, cutlass.Float32), num=2) + dy_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + dy_col_id, cutlass.Float32), num=2) + + # ---- -dY -> TMEM: A operand of the dstate K-term ----------------------- + neg_dy_regs0 = cutlass.Array(cutlass.Float32, 8, alignment=16) + neg_dy_regs1 = cutlass.Array(cutlass.Float32, 8, alignment=16) + for e in cutlass.range_constexpr(8): + neg_dy_regs0[e] = -dy_vec0[e] + neg_dy_regs1[e] = -dy_vec1[e] + neg_dy_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + neg_dy_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + frag_pair = reg_idx * 2 + neg_dy_pack0[reg_idx] = fp32_to_fp16(neg_dy_regs0[frag_pair], neg_dy_regs0[frag_pair + 1], dtype=cfg.io_dtype) + neg_dy_pack1[reg_idx] = fp32_to_fp16(neg_dy_regs1[frag_pair], neg_dy_regs1[frag_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + (tmem_col + cfg.tmem_neg_dy_inp_offset), cutlass.Int8), neg_dy_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + (tmem_col + cfg.tmem_neg_dy_inp_offset), cutlass.Int8), neg_dy_pack1[0:4]) + nvvm.tcgen05_wait("store") + bars.mb_neg_dy_inp_ready.arrive() + + # ---- dY -> sdY: pack + store + publish (super dM + dV scalar operand) -- + addr_lo0 = (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) + ov_tok * 64 + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2) + addr_lo1 = ( + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) + + ov_tok * 64 + + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2) + ) + dy_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + dy_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + dy_pack0[reg_idx] = fp32_to_fp16(dy_vec0[2 * reg_idx], dy_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) + dy_pack1[reg_idx] = fp32_to_fp16(dy_vec1[2 * reg_idx], dy_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) + bars.mb_dy_smem_done.wait(sdy_done_index.phase) + sdy_done_index = advance(sdy_done_index, 1) + nvvm.stmatrix(sDy_raw.data_ptr() + addr_lo0, dy_pack0.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) + nvvm.stmatrix(sDy_raw.data_ptr() + addr_lo1, dy_pack1.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dy_smem_ready.arrive() + + # ---- scalar pass over own sdY: dV staging ---------------------------- + dv_stage = gc % cfg.smem_dv_stages + bars.mb_dv_tmastg_done[dv_stage].wait(dv_done_index.phase) + dv_done_index = advance(dv_done_index, cfg.smem_dv_stages) + dwo_stage = gc % cfg.smem_dwo_stages + bars.mb_dwo_tmastg_done[dwo_stage].wait(dwo_done_index.phase) + dwo_done_index = advance(dwo_done_index, cfg.smem_dwo_stages) + nvvm.barrier_cta_sync(cfg.cg1_sync_barrier_id, thread_count=cfg.cg1_threads) + sdv_stage_base = dv_stage * (cfg.b_t * cfg.d_v) + sdwo_stage_base = dwo_stage * (cfg.b_t * cfg.d_v) + c_seg = value_dim // 64 + c_dim = value_dim - c_seg * 64 + for t in cutlass.range_constexpr(cfg.b_t): + idx = c_seg * (cfg.b_t * 64) + t * 64 + swizzle_xor_128b(t, c_dim, elem_bytes=2) + dy_v = (sDy_raw.data_ptr() + idx).load().to(cutlass.Float32) + w_v = (sW_ptr + idx).load().to(cutlass.Float32) + v_v = (sV_ptr + idx).load().to(cutlass.Float32) + (sDv_raw.data_ptr() + sdv_stage_base + idx).store((w_v * dy_v).to(cfg.io_dtype)) + (sDwOut_raw.data_ptr() + sdwo_stage_base + idx).store((v_v * dy_v).to(cfg.io_dtype)) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dv_tmastg_ready[dv_stage].arrive() + bars.mb_dwo_tmastg_ready[dwo_stage].arrive() + bars.mb_v_done[raw_index.idx].arrive() + bars.mb_w_done[raw_index.idx].arrive() + + # ---- dH capture for the next ----------------------------------------- + bars.mb_dstate_acc_ready.wait(dstate_ready_index.phase) + dstate_ready_index = advance(dstate_ready_index, 1) + if rev_idx + cutlass.Int32(1) < sk_nt: + bars.mb_dstate_smem_done.wait(dstate_smem_done_index.phase) + 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 + for sub in cutlass.range_constexpr(cfg.d_k // 32): + dstate_vec = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 32), cutlass.Float32), num=32 + ) + dstate_pack = cutlass.Array(cutlass.Int32, 16, alignment=16) + for pc in cutlass.range_constexpr(16): + dstate_pack[pc] = fp32_to_fp16(dstate_vec[2 * pc], dstate_vec[2 * pc + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 16), cutlass.Int8), + dstate_pack[0:16], + ) + nvvm.tcgen05_wait("store") + bars.mb_dstate_inp_ready.arrive() + + # ---- dh_inp -> sdH: re-read after the TMEM publish --------------- + for sub in cutlass.range_constexpr(cfg.d_k // 32): + dstate_words = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 16), cutlass.Float32), num=16 + ) + for half in cutlass.range_constexpr(4): + d_base = sub * 32 + half * 8 + h_pack = cutlass.Vector.from_elements( + (dstate_words[half * 4], dstate_words[half * 4 + 1], dstate_words[half * 4 + 2], dstate_words[half * 4 + 3]), + cutlass.Float32, + ).bitcast(cfg.io_dtype) + h_addr = (d_base // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, d_base % 64, elem_bytes=2) + (sDstate_raw.data_ptr() + h_addr).store(h_pack, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dstate_smem_ready.arrive() + raw_index = advance(raw_index, cfg.smem_raw_stages) + + # ---- tile end: dS0 drain / zero-length pass-through ---------------------- + if cutlass.const_expr(mDstate0 is not None): + if sk_nt > 0: + if wstart == 0: + row_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 + 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] + 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] + else: + mDstate0[batch_idx, head_idx, kd, value_dim] = 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) + + bars.mb_tmem_done[0].arrive() + + +@cute.jit +def compute2_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_hold, + warp_idx, + sBeta_raw, + sGate_raw, + sNorm_raw, + sDq_raw, + sDk_raw, + sRed1_raw, + sDstate_raw, + sDgate_raw, + sDb_raw, + scale, + bars, +) -> None: + """WG2 warp role (warps 8-11): the gradient drain. Each thread owns one + key channel d for all 16 tokens: reads the dQ accumulator and the four dK + parts from TMEM, assembles dQ/dK with the per-channel gate factors (raw + Q/K arrive through WG0's TMEM ring), applies the in-kernel L2-norm + backward row projection, assembles the per-channel dGate including the + g_last terms, reverse-cumsums it in registers, stages dGate and dBeta + (db) for the epilogue's TMA stores, and stages dQ/dK for the epilogue's + TMA stores.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_2, nvvm.SetMaxRegisterAction.INCREASE) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_hold.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + wg1_sp = warp_idx % 4 + channel = wg1_sp * cfg.threads_per_warp + lane + row_addr = (tmem_row + wg1_sp * cfg.threads_per_warp) << 16 + cg2_tidx = channel + + raw_index = PipelineState.start(phase=0) + dq_acc_index = PipelineState.start(phase=0) + dk_decay_part_index = PipelineState.start(phase=0) + dk_inv_part_index = PipelineState.start(phase=0) + dk_restore_part_index = PipelineState.start(phase=0) + dgate_last_dstate_smem_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 + for rev_idx in cutlass.range(sk_nt, unroll=1): + chunk_idx = cend - cutlass.Int32(1) - rev_idx + gc = gbase + rev_idx + chunk_start = chunk_idx * cfg.b_t + raw_stage = gc % cfg.smem_raw_stages + decay_stage = gc % cfg.smem_decay_stages + has_dstate = cutlass.Boolean(rev_idx > 0) + if cutlass.const_expr(cfg.use_dstate_in): + has_dstate = cutlass.Boolean(True) + sBetaP_ptr = sBeta_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sGate_ptr = sGate_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + writes = chunk_idx < wend + + # ---- raw q/k/beta/gate landed: CG0 publishes the decay ring only after + # consuming them, so this wait is CG2's visibility guard --------------- + bars.mb_k_decay_inv_ready[decay_stage].wait((gc // cfg.smem_decay_stages) % 2) + + # ---- per-channel gate factors ---------------------------------------- + f32_seg = channel // 32 + f32_dim = channel - f32_seg * 32 + f16_seg = channel // 64 + f16_dim = channel - f16_seg * 64 + eg = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + for t in cutlass.range_constexpr(cfg.b_t): + eg[t] = (sGate_ptr + f32_seg * (cfg.b_t * 32) + t * 32 + swizzle_xor_128b(t, f32_dim, elem_bytes=4)).load() + egl = eg[cfg.b_t - 1] + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_gate_done[raw_stage].arrive() + + # ---- staged raw Q/K: TMEM ring cols for this chunk ------------------- + qk_raw_stage = gc % cfg.tmem_qk_raw_stages + qraw_col = tmem_col + cfg.tmem_qraw_inp_offset + qk_raw_stage * (cfg.b_t // 2) + kraw_col = tmem_col + cfg.tmem_kraw_inp_offset + qk_raw_stage * (cfg.b_t // 2) + norm_base = qk_raw_stage * (2 * cfg.b_t) + bars.mb_qk_raw_ready[qk_raw_stage].wait((gc // cfg.tmem_qk_raw_stages) % 2) + + # ---- dGate_last hdot: sum_v sdH[v, c] * S0[c, v] --------------------- + dgate_last_val = cutlass.Float32(0.0) + bars.mb_state_inp_ready[gc % 2].wait((gc // 2) % 2) + if has_dstate: + bars.mb_dstate_smem_ready.wait(dgate_last_dstate_smem_index.phase) + dgate_last_dstate_smem_index = advance(dgate_last_dstate_smem_index, 1) + for pl in cutlass.range_constexpr(2): + for row_half in cutlass.range_constexpr(2): + state_vec = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr( + row_addr + (tmem_col + cfg.tmem_state_inp_offset + (gc % 2) * (cfg.d_v // 2) + pl * 32 + row_half * 16), cutlass.Float32 + ), + num=16, + ) + hacc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for i in cutlass.range_constexpr(8): + hacc[i] = opaque_f32_zero() + for j in cutlass.range_constexpr(16): + v0 = pl * 64 + row_half * 32 + 2 * j + state_pair = cutlass.Vector.from_elements((state_vec[j],), cutlass.Float32).bitcast(cfg.io_dtype) + dstate_addr0 = (channel // 64) * (cfg.d_v * 64) + v0 * 64 + swizzle_xor_128b(v0, channel % 64, elem_bytes=2) + dstate_addr1 = (channel // 64) * (cfg.d_v * 64) + (v0 + 1) * 64 + swizzle_xor_128b(v0 + 1, channel % 64, elem_bytes=2) + hval0 = (sDstate_raw.data_ptr() + dstate_addr0).load().to(cutlass.Float32) + hval1 = (sDstate_raw.data_ptr() + dstate_addr1).load().to(cutlass.Float32) + hacc[(2 * j) % 8] = hacc[(2 * j) % 8] + hval0 * state_pair[0].to(cutlass.Float32) + hacc[(2 * j + 1) % 8] = hacc[(2 * j + 1) % 8] + hval1 * state_pair[1].to(cutlass.Float32) + part_a = (hacc[0] + hacc[4]) + (hacc[1] + hacc[5]) + part_b = (hacc[2] + hacc[6]) + (hacc[3] + hacc[7]) + dgate_last_val = dgate_last_val + (part_a + part_b) + bars.mb_dstate_smem_cg2_done.arrive() + bars.mb_state_inp_cg2_done[gc % 2].arrive() + + # ---- part-drain accumulators ------------------------------------------- + dq_n = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + dk_n = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + db_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + dgate_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + dgate_last_acc = cutlass.Array(cutlass.Float32, 4, alignment=16) + for i in cutlass.range_constexpr(4): + dgate_last_acc[i] = opaque_f32_zero() + for t in cutlass.range_constexpr(cfg.b_t): + dk_n[t] = cutlass.Float32(0.0) + + # ---- dK_restore part drain: (eGl/eG) scale + dGate_last k-dot ---------- + if has_dstate: + bars.mb_dk_restore_part_acc_ready.wait(dk_restore_part_index.phase) + dk_restore_part_index = advance(dk_restore_part_index, 1) + dk_restore_part_vec = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_restore_acc_offset), cutlass.Float32), num=cfg.b_t + ) + kr_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) + for t in cutlass.range_constexpr(cfg.b_t): + dk_hat = egl * cute.math.rcp(eg[t], approx=True, ftz=True) * dk_restore_part_vec[t] + dk_n[t] = dk_hat + k_pair = cutlass.Vector.from_elements((kr_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + k_v = k_pair[t % 2].to(cutlass.Float32) + if cutlass.const_expr(cfg.l2norm): + k_v = k_v * sNorm_raw[norm_base + cfg.b_t + t] + dgate_last_acc[t % 4] = dgate_last_acc[t % 4] + k_v * dk_hat + + # ---- dQ acc drain: eG.scale --------------------------------------------- + bars.mb_dq_acc_ready.wait(dq_acc_index.phase) + dq_acc_index = advance(dq_acc_index, 1) + dq_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dq_acc_offset), cutlass.Float32), num=cfg.b_t) + for t2 in cutlass.range_constexpr(cfg.b_t // 2): + t = 2 * t2 + es_lo, es_hi = fmul2(eg[t], eg[t + 1], scale, scale) + dq_n[t], dq_n[t + 1] = fmul2(es_lo, es_hi, dq_vec[t], dq_vec[t + 1]) + + # ---- dK_inv part drain: (dA - dM) term, 1/eG scale ---------------------- + bars.mb_dk_inv_part_acc_ready.wait(dk_inv_part_index.phase) + dk_inv_part_index = advance(dk_inv_part_index, 1) + dk_inv_part_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), num=cfg.b_t) + for t in cutlass.range_constexpr(cfg.b_t): + dk_n[t] = dk_n[t] + dk_inv_part_vec[t] * cute.math.rcp(eg[t], approx=True, ftz=True) + + # ---- dK_decay part drain: -eG scale, seeds dBeta and dGate -------------- + bars.mb_dk_decay_part_acc_ready.wait(dk_decay_part_index.phase) + dk_decay_part_index = advance(dk_decay_part_index, 1) + dk_decay_part_vec = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), num=cfg.b_t + ) + kd_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) + for t in cutlass.range_constexpr(cfg.b_t): + dk_decay = -eg[t] * dk_decay_part_vec[t] + k_pair = cutlass.Vector.from_elements((kd_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + k_v = k_pair[t % 2].to(cutlass.Float32) + if cutlass.const_expr(cfg.l2norm): + k_v = k_v * sNorm_raw[norm_base + cfg.b_t + t] + db_regs[t] = k_v * dk_decay + dgate_regs[t] = (sBetaP_ptr + f16_seg * (cfg.b_t * 64) + t * 64 + swizzle_xor_128b(t, f16_dim, elem_bytes=2)).load().to( + cutlass.Float32 + ) * dk_decay + dk_n[t] = dk_n[t] + dgate_regs[t] + + nvvm.tcgen05_wait("load") + bars.mb_dqk_acc_done.arrive() + + # ---- dGate finalize -------------------------------------------------- + qf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qraw_col, cutlass.Float32), num=cfg.b_t // 2) + kf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) + for t in cutlass.range_constexpr(cfg.b_t): + q_pair = cutlass.Vector.from_elements((qf_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + k_pair = cutlass.Vector.from_elements((kf_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + q_v = q_pair[t % 2].to(cutlass.Float32) + k_v = k_pair[t % 2].to(cutlass.Float32) + if cutlass.const_expr(cfg.l2norm): + q_v = q_v * sNorm_raw[norm_base + t] + k_v = k_v * sNorm_raw[norm_base + cfg.b_t + t] + dgate_regs[t] = ( + q_v * dq_n[t] + + (sBetaP_ptr + f16_seg * (cfg.b_t * 64) + t * 64 + swizzle_xor_128b(t, f16_dim, elem_bytes=2)).load().to(cutlass.Float32) * db_regs[t] + - k_v * (dk_n[t] - dgate_regs[t]) + ) + dgate_regs[cfg.b_t - 1] = dgate_regs[cfg.b_t - 1] + ((dgate_last_acc[0] + dgate_last_acc[1]) + (dgate_last_acc[2] + dgate_last_acc[3])) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_beta_done[raw_stage].arrive() + + # ---- L2-norm backward row projection --------------------------------- + if cutlass.const_expr(cfg.l2norm): + for grad, qk_col, inv_off in ((dq_n, qraw_col, 0), (dk_n, kraw_col, cfg.b_t)): + dots = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + for half in cutlass.range_constexpr(2): + p_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) + for tt in cutlass.range_constexpr(cfg.b_t // 2): + t = half * (cfg.b_t // 2) + tt + p_pair = cutlass.Vector.from_elements((p_words[tt // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + dots[t] = grad[t] * p_pair[tt % 2].to(cutlass.Float32) * sNorm_raw[norm_base + inv_off + t] + for off in cutlass.range_constexpr(5): + step = cutlass.const_expr(1 << off) + for t in cutlass.range_constexpr(cfg.b_t): + dots[t] = dots[t] + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, dots[t], step, 31, kind=nvvm.Shfl.BFLY)) + if lane == 0: + for t in cutlass.range_constexpr(cfg.b_t): + sRed1_raw[wg1_sp * cfg.b_t + t] = dots[t] + nvvm.barrier_cta_sync(cfg.cg2_sync_barrier_id, thread_count=cfg.cg2_threads) + for half in cutlass.range_constexpr(2): + a_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) + for tt in cutlass.range_constexpr(cfg.b_t // 2): + t = half * (cfg.b_t // 2) + tt + a_pair = cutlass.Vector.from_elements((a_words[tt // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + total_dot = sRed1_raw[t] + sRed1_raw[cfg.b_t + t] + sRed1_raw[2 * cfg.b_t + t] + sRed1_raw[3 * cfg.b_t + t] + norm_t = sNorm_raw[norm_base + inv_off + t] + grad[t] = (grad[t] - a_pair[tt % 2].to(cutlass.Float32) * norm_t * total_dot) * norm_t + nvvm.barrier_cta_sync(cfg.cg2_sync_barrier_id, thread_count=cfg.cg2_threads) + + nvvm.tcgen05_wait("load") + bars.mb_qk_raw_done[qk_raw_stage].arrive() + + # ---- stage dQ/dK for the epilogue TMA stores ------------------------- + dq_stage = gc % cfg.smem_dq_stages + dk_stage = gc % cfg.smem_dk_stages + bars.mb_dq_tmastg_done[dq_stage].wait(((gc // cfg.smem_dq_stages) + 1) % 2) + bars.mb_dk_tmastg_done[dk_stage].wait(((gc // cfg.smem_dk_stages) + 1) % 2) + dq_base = dq_stage * (cfg.b_t * cfg.d_k) + dk_base = dk_stage * (cfg.b_t * cfg.d_k) + for t in cutlass.range_constexpr(cfg.b_t): + out_idx = f16_seg * (cfg.b_t * 64) + t * 64 + swizzle_xor_128b(t, f16_dim, elem_bytes=2) + (sDq_raw.data_ptr() + dq_base + out_idx).store(dq_n[t].to(cfg.io_dtype)) + (sDk_raw.data_ptr() + dk_base + out_idx).store(dk_n[t].to(cfg.io_dtype)) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dq_tmastg_ready[dq_stage].arrive() + bars.mb_dk_tmastg_ready[dk_stage].arrive() + + # ---- dGate_last add -------------------------------------------------- + if has_dstate: + if chunk_idx >= FIRST_STATE_CHUNK: + dgate_regs[cfg.b_t - 1] = dgate_regs[cfg.b_t - 1] + egl * dgate_last_val + + # ---- dGate reverse cumsum -------------------------------------------- + suffix = cutlass.Float32(0.0) + for rt in cutlass.range_constexpr(cfg.b_t): + t = cfg.b_t - 1 - rt + suffix = suffix + dgate_regs[t] + dgate_regs[t] = suffix + + # ---- stage dGate + dBeta for the epilogue TMA stores ----------------- + dgate_stage = gc % cfg.smem_dgate_stages + bars.mb_dgate_tmastg_done[dgate_stage].wait(((gc // cfg.smem_dgate_stages) + 1) % 2) + db_stage = gc % cfg.smem_db_stages + bars.mb_db_tmastg_done[db_stage].wait(((gc // cfg.smem_db_stages) + 1) % 2) + for t in cutlass.range_constexpr(cfg.b_t): + dgate_idx = f32_seg * (cfg.b_t * 32) + t * 32 + swizzle_xor_128b(t, f32_dim, elem_bytes=4) + (sDgate_raw.data_ptr() + dgate_idx).store(dgate_regs[t]) + db_idx = f16_seg * (cfg.b_t * 64) + t * 64 + swizzle_xor_128b(t, f16_dim, elem_bytes=2) + (sDb_raw.data_ptr() + db_idx).store(db_regs[t].to(cfg.io_dtype)) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dgate_tmastg_ready[dgate_stage].arrive() + bars.mb_db_tmastg_ready[db_stage].arrive() + raw_index = advance(raw_index, cfg.smem_raw_stages) + gbase += sk_nt + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + bars.mb_tmem_done[0].arrive() + + +# --------------------------------------------------------------------------- +# Host-side assembly +# --------------------------------------------------------------------------- + + +@cute.kernel +def build_all_descs_kernel( + base_q: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_v: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_gate: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_do: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_beta: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_w: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_dq: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_dk: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_dv: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_dgate: cutlass.GridConstant[cuda.tensor_map.TensorMap], + 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, + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + do: cute.Tensor, + beta: cute.Tensor, + w: cute.Tensor, + dq: cute.Tensor, + dk: cute.Tensor, + dv: cute.Tensor, + dgate: cute.Tensor, + 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, + v_rs: cutlass.Int32, + g_rs: cutlass.Int32, + do_rs: cutlass.Int32, + beta_rs: cutlass.Int32, + w_rs: cutlass.Int32, + dq_rs: cutlass.Int32, + dk_rs: cutlass.Int32, + dv_rs: cutlass.Int32, + dgate_rs: cutlass.Int32, + dwo_rs: cutlass.Int32, + dbo_rs: cutlass.Int32, + checkpoint_rs: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, +) -> None: + """Single-launch builder for the per-batch TMA-descriptor arrays (one + warp per array).""" + tidx, _, _ = cute.arch.thread_idx() + widx = cutlass.Int32(tidx) // cutlass.Int32(32) + arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) + sub0 = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + 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((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((arr_words,), stride=(1,))) + 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(): + emit_seq_descs(base_q, sub0, cu_seqlens, q, n_batch, q_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 1: + if nvvm.elect_sync(): + emit_seq_descs(base_k, sub1, cu_seqlens, k, n_batch, k_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 2: + if nvvm.elect_sync(): + emit_seq_descs(base_v, sub2, cu_seqlens, v, n_batch, v_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 3: + if nvvm.elect_sync(): + emit_seq_descs(base_gate, sub3, cu_seqlens, gate, n_batch, g_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 4: + if nvvm.elect_sync(): + emit_seq_descs(base_do, sub4, cu_seqlens, do, n_batch, do_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 5: + if nvvm.elect_sync(): + emit_seq_descs(base_beta, sub5, cu_seqlens, beta, n_batch, beta_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 6: + if nvvm.elect_sync(): + emit_seq_descs(base_w, sub6, cu_seqlens, w, n_batch, w_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 7: + if nvvm.elect_sync(): + emit_seq_descs(base_dq, sub7, cu_seqlens, dq, n_batch, dq_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 8: + if nvvm.elect_sync(): + emit_seq_descs(base_dk, sub8, cu_seqlens, dk, n_batch, dk_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 9: + if nvvm.elect_sync(): + emit_seq_descs(base_dv, sub9, 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 widx == 10: + if nvvm.elect_sync(): + emit_seq_descs(base_dgate, sub10, cu_seqlens, dgate, n_batch, dgate_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 11: + if nvvm.elect_sync(): + emit_seq_descs(base_dwo, sub11, cu_seqlens, dwo, n_batch, dwo_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 12: + if nvvm.elect_sync(): + emit_seq_descs(base_dbo, sub12, cu_seqlens, dbo, n_batch, dbo_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 13: + 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.jit +def build_descs( + io_dtype: cutlass.Constexpr, + b_t: cutlass.Constexpr[int], + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + do: cute.Tensor, + beta: cute.Tensor, + w: cute.Tensor, + dq: cute.Tensor, + dk: cute.Tensor, + dv: cute.Tensor, + dgate: cute.Tensor, + dwo: cute.Tensor, + dbo: cute.Tensor, + state_checkpoints: cute.Tensor, + initial_state: cute.Tensor | None, + cu_seqlens: cute.Tensor, + tensormap_workspace: cute.Tensor, + stream: cuda_driver.CUstream, +): + """Build the 15 per-(batch, head) TMA-descriptor arrays into + ``tensormap_workspace``.""" + h_q = q.shape[1] + h_k = k.shape[1] + h_v = v.shape[1] + ho = gate.shape[1] + batch_size = cu_seqlens.shape[0] - 1 + d_k = q.shape[2] + d_v = v.shape[2] + bpe = io_dtype.width // 8 + granu = 128 // bpe + seqlen = q.shape[0] + + q_headed = cute.make_tensor(q.iterator, cute.make_layout((d_k, h_q, seqlen), stride=(1, q.stride[1], q.stride[0]))) + k_headed = cute.make_tensor(k.iterator, cute.make_layout((d_k, h_k, seqlen), stride=(1, k.stride[1], k.stride[0]))) + v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v.stride[1], v.stride[0]))) + gate_headed = cute.make_tensor(gate.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, gate.stride[1], gate.stride[0]))) + do_headed = cute.make_tensor(do.iterator, cute.make_layout((d_v, ho, seqlen), stride=(1, do.stride[1], do.stride[0]))) + beta_headed = cute.make_tensor(beta.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, beta.stride[1], beta.stride[0]))) + w_headed = cute.make_tensor(w.iterator, cute.make_layout((d_v, ho, seqlen), stride=(1, w.stride[1], w.stride[0]))) + dq_headed = cute.make_tensor(dq.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, dq.stride[1], dq.stride[0]))) + dk_headed = cute.make_tensor(dk.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, dk.stride[1], dk.stride[0]))) + dv_headed = cute.make_tensor(dv.iterator, cute.make_layout((d_v, ho, seqlen), stride=(1, dv.stride[1], dv.stride[0]))) + dgate_headed = cute.make_tensor(dgate.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, dgate.stride[1], dgate.stride[0]))) + dwo_headed = cute.make_tensor(dwo.iterator, cute.make_layout((d_v, ho, seqlen), stride=(1, dwo.stride[1], dwo.stride[0]))) + dbo_headed = cute.make_tensor(dbo.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, dbo.stride[1], dbo.stride[0]))) + + swz = cuda.TensorMapSwizzle.s128b + base_q = cuda.create_tensor_map_tiled_from_view(q_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_k = cuda.create_tensor_map_tiled_from_view(k_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_v = cuda.create_tensor_map_tiled_from_view(v_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_gate = cuda.create_tensor_map_tiled_from_view(gate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_do = cuda.create_tensor_map_tiled_from_view(do_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_beta = cuda.create_tensor_map_tiled_from_view(beta_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_w = cuda.create_tensor_map_tiled_from_view(w_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dq = cuda.create_tensor_map_tiled_from_view(dq_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dk = cuda.create_tensor_map_tiled_from_view(dk_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dv = cuda.create_tensor_map_tiled_from_view(dv_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dgate = cuda.create_tensor_map_tiled_from_view(dgate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dwo = cuda.create_tensor_map_tiled_from_view(dwo_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dbo = cuda.create_tensor_map_tiled_from_view(dbo_headed, box_dims=(granu, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + + checkpoint_view = cute.make_tensor( + state_checkpoints.iterator, + cute.make_layout( + (d_v, d_k, state_checkpoints.shape[0], ho), + stride=(state_checkpoints.stride[3], state_checkpoints.stride[2], state_checkpoints.stride[0], state_checkpoints.stride[1]), + ), + ) + 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) + + n_warps = 15 if initial_state is not None else 14 + build_all_descs_kernel( + base_q, + base_k, + base_v, + base_gate, + base_do, + base_beta, + base_w, + base_dq, + base_dk, + base_dv, + base_dgate, + base_dwo, + base_dbo, + base_checkpoint, + base_initial_state, + tensormap_workspace, + cu_seqlens, + q, + k, + v, + gate, + do, + beta, + w, + dq, + dk, + dv, + dgate, + dwo, + dbo, + state_checkpoints, + initial_state, + cutlass.Int32(batch_size), + cutlass.Int32(q.stride[0]), + cutlass.Int32(k.stride[0]), + cutlass.Int32(v.stride[0]), + cutlass.Int32(gate.stride[0]), + cutlass.Int32(do.stride[0]), + cutlass.Int32(beta.stride[0]), + cutlass.Int32(w.stride[0]), + cutlass.Int32(dq.stride[0]), + cutlass.Int32(dk.stride[0]), + cutlass.Int32(dv.stride[0]), + cutlass.Int32(dgate.stride[0]), + cutlass.Int32(dwo.stride[0]), + cutlass.Int32(dbo.stride[0]), + cutlass.Int32(state_checkpoints.stride[0]), + cutlass.Int32(b_t), + ).launch(grid=(1, 1, 1), block=(32 * n_warps, 1, 1), stream=stream) + + +@cute.jit +def host( + cfg: cutlass.Constexpr, + state_checkpoints: cute.Tensor, + mState_init: cute.Tensor | None, + dgate: cute.Tensor, + dbeta: cute.Tensor, + dw: cute.Tensor, + cu_seqlens: cute.Tensor, + d_initial_state: cute.Tensor | None, + d_final_state: cute.Tensor | None, + work_items: cute.Tensor | None, + work_count: cute.Tensor | None, + sched_ctr: cute.Tensor | None, + tensormap_workspace: cute.Tensor, + scale: cutlass.Float32, + stream, +) -> None: + num_sequences = cu_seqlens.shape[0] - 1 + + # ---- launch ------------------------------------------------------------------ + n_desc = num_sequences + grid_shape = (cfg.max_active_clusters, 1, 1) + kernel( + cfg, + tensormap_workspace, + n_desc, + cu_seqlens, + dgate, + dbeta, + dw, + d_initial_state, + d_final_state, + work_items, + work_count, + sched_ctr, + scale, + ).launch( + grid=grid_shape, + block=(cfg.threads_per_cta, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + +@cute.kernel +def kernel( + cfg: cutlass.Constexpr, + tensormap_workspace: cute.Tensor, + n_desc: cutlass.Int32, + cu_seqlens: cute.Tensor, + mDgate: cute.Tensor, + mDb: cute.Tensor, + mDw_out: cute.Tensor, + mDstate0: cute.Tensor | None, + mDstate_in: cute.Tensor | None, + mWorkItems: cute.Tensor, + mCount: cute.Tensor, + mSched: cute.Tensor | None, + scale: cutlass.Float32, +) -> None: + """BT=16 GDN-2 backward kernel (persistent, 16 warps).""" + tidx, _, _ = cute.arch.thread_idx() + bidx = cute.arch.block_idx()[0] + num_ctas = cute.arch.grid_dim()[0] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % cfg.threads_per_warp + + total_tiles = mCount[0] + assert cu_seqlens.element_type in (cutlass.Int32, cutlass.Int64) + assert mDgate.element_type == cutlass.Float32 + assert mDb.element_type == cfg.io_dtype and mDw_out.element_type == cfg.io_dtype + + desc_base_words = tensormap_workspace.iterator.raw_ptr() + arr_words = n_desc * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_q_base = desc_base_words + desc_k_base = desc_base_words + arr_words + desc_v_base = desc_base_words + cutlass.Int32(2) * arr_words + desc_gate_base = desc_base_words + cutlass.Int32(3) * arr_words + desc_do_base = desc_base_words + cutlass.Int32(4) * arr_words + desc_beta_base = desc_base_words + cutlass.Int32(5) * arr_words + desc_w_base = desc_base_words + cutlass.Int32(6) * arr_words + desc_dq_base = desc_base_words + cutlass.Int32(7) * arr_words + desc_dk_base = desc_base_words + cutlass.Int32(8) * arr_words + desc_dv_base = desc_base_words + cutlass.Int32(9) * arr_words + desc_dgate_base = desc_base_words + cutlass.Int32(10) * arr_words + 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) + tmem_hold = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) + sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + bpe = cfg.io_dtype.width // 8 + SWZ = 2 + LEAD = 16 + STRIDE = 8 * 128 + STATE_ALT_LEAD = cfg.d_v * 128 + + # sub-bank split: tcgen05-descriptor operands low, generic-client buffers high + sRed1_raw = cutlass.Array(cutlass.Float32, 4 * cfg.b_t, space=SMEM, alignment=64) + sNorm_raw = cutlass.Array(cutlass.Float32, cfg.tmem_qk_raw_stages * 2 * cfg.b_t, space=SMEM, alignment=64) + sState_raw = cutlass.Array(cfg.io_dtype, cfg.state_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDstate_raw = cutlass.Array(cfg.io_dtype, cfg.d_k * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_decay_raw = cutlass.Array(cfg.io_dtype, cfg.operand_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_inv_raw = cutlass.Array(cfg.io_dtype, cfg.operand_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_restore_raw = cutlass.Array(cfg.io_dtype, cfg.operand_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sQ_decay_raw = cutlass.Array(cfg.io_dtype, cfg.operand_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sState_scale_diag_raw = cutlass.Array(cfg.io_dtype, cfg.diag_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sIntermediate_raw = cutlass.Array(cfg.io_dtype, cfg.intermediate_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDo_raw = cutlass.Array(cfg.io_dtype, cfg.raw_v_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sBeta_raw = cutlass.Array(cfg.io_dtype, cfg.raw_qk_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + # sub-bank fill: high group starts at the 128KB midpoint + smem_bank_fill = cutlass.Array(cfg.io_dtype, 1024 // bpe, space=SMEM, alignment=cfg.buffer_align_bytes) + sQ_raw = cutlass.Array(cfg.io_dtype, cfg.raw_qk_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_raw = cutlass.Array(cfg.io_dtype, cfg.raw_qk_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sV_raw = cutlass.Array(cfg.io_dtype, cfg.raw_v_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sGate_raw = cutlass.Array(cutlass.Float32, cfg.raw_gate_cosize, space=SMEM, alignment=1024) + sW_raw = cutlass.Array(cfg.io_dtype, cfg.raw_v_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sU_raw = cutlass.Array(cfg.io_dtype, cfg.b_t * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) + sDy_raw = cutlass.Array(cfg.io_dtype, cfg.b_t * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) + sDq_raw = cutlass.Array(cfg.io_dtype, cfg.dq_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDk_raw = cutlass.Array(cfg.io_dtype, cfg.dk_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDv_raw = cutlass.Array(cfg.io_dtype, cfg.dv_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDgate_raw = cutlass.Array(cutlass.Float32, cfg.dgate_cosize, space=SMEM, alignment=1024) + sDwOut_raw = cutlass.Array(cfg.io_dtype, cfg.dwo_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDb_raw = cutlass.Array(cfg.io_dtype, cfg.db_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + + sState_alt = 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=STATE_ALT_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, + stages=cfg.smem_decay_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sK_restore_lead16 = SmemTile( + base=sK_restore_raw.data_ptr().toint(), + elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, + stages=cfg.smem_decay_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sDo_lead16 = SmemTile( + base=sDo_raw.data_ptr().toint(), + elems_per_stage=((cfg.raw_v_cosize) // (cfg.smem_raw_stages)) * bpe, + stages=cfg.smem_raw_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sDo_amaj = SmemTile( + base=sDo_raw.data_ptr().toint(), + elems_per_stage=((cfg.raw_v_cosize) // (cfg.smem_raw_stages)) * bpe, + stages=cfg.smem_raw_stages, + leading_byte_offset=cfg.b_t * 128, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sU_lead16 = SmemTile( + base=sU_raw.data_ptr().toint(), + elems_per_stage=((cfg.b_t * cfg.d_v) // (1)) * bpe, + stages=1, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sDy_lead16 = SmemTile( + base=sDy_raw.data_ptr().toint(), + elems_per_stage=((cfg.b_t * cfg.d_v) // (1)) * bpe, + stages=1, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sDstate_alt = SmemTile( + base=sDstate_raw.data_ptr().toint(), + elems_per_stage=((cfg.d_k * cfg.d_v) // (1)) * bpe, + stages=1, + leading_byte_offset=STATE_ALT_LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + + sQ_decay_trans = SmemTile( + base=sQ_decay_raw.data_ptr().toint(), + elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, + stages=cfg.smem_decay_stages, + leading_byte_offset=cfg.b_t * 128, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sK_inv_amaj = SmemTile( + base=sK_inv_raw.data_ptr().toint(), + elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, + stages=cfg.smem_decay_stages, + leading_byte_offset=cfg.b_t * 128, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sK_decay_trans = SmemTile( + base=sK_decay_raw.data_ptr().toint(), + elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, + stages=cfg.smem_decay_stages, + leading_byte_offset=cfg.b_t * 128, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sState_scale_diag = SmemTile( + base=sState_scale_diag_raw, + elems_per_stage=((cfg.d_k // 16) * 256), + stages=cfg.smem_decay_stages, + leading_byte_offset=16, + stride_byte_offset=(8 * 16 * 2), + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, + ) + sIntermediate = SmemTile( + base=sIntermediate_raw, + elems_per_stage=(cfg.intermediate_tiles * cfg.b_t * cfg.b_t), + stages=cfg.smem_intermediate_stages, + leading_byte_offset=16, + stride_byte_offset=(8 * cfg.b_t * 2), + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, + ) + q_tx_bytes = cutlass.const_expr(cfg.d_k * cfg.b_t * bpe) + k_tx_bytes = cutlass.const_expr(cfg.d_k * cfg.b_t * bpe) + gate_tx_bytes = cutlass.const_expr(cfg.d_k * cfg.b_t * 4) + beta_tx_bytes = cutlass.const_expr(cfg.d_k * cfg.b_t * bpe) + do_tx_bytes = cutlass.const_expr(cfg.d_v * cfg.b_t * bpe) + v_tx_bytes = cutlass.const_expr(cfg.d_v * cfg.b_t * bpe) + w_tx_bytes = cutlass.const_expr(cfg.d_v * cfg.b_t * bpe) + + elect_one = nvvm.elect_sync() + if warp_idx == cfg.tma_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_raw_stages): + bars.mb_q_ready[stage].init() + bars.mb_q_done[stage].init() + bars.mb_k_ready[stage].init() + bars.mb_k_done[stage].init() + bars.mb_gate_ready[stage].init() + bars.mb_gate_done[stage].init() + bars.mb_beta_ready[stage].init() + bars.mb_beta_done[stage].init() + bars.mb_do_ready[stage].init() + bars.mb_do_done[stage].init() + bars.mb_do_mma_done[stage].init() + bars.mb_v_ready[stage].init() + bars.mb_v_done[stage].init() + bars.mb_w_ready[stage].init() + bars.mb_w_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_state_stages): + bars.mb_state_ready[stage].init() + bars.mb_state_done[stage].init() + bars.mb_state_cg0_done[stage].init() + for stage in cutlass.range_constexpr(2): + bars.mb_state_inp_ready[stage].init() + bars.mb_state_inp_done[stage].init() + bars.mb_state_inp_cg2_done[stage].init() + elif warp_idx == cfg.tcgen05_mma_warp_id: + if elect_one: + bars.mb_state_k_acc_ready.init() + bars.mb_y_inp_ready.init() + bars.mb_u_acc_ready.init() + bars.mb_u_smem_ready.init() + bars.mb_du_acc_ready.init() + bars.mb_du_inp_ready.init() + bars.mb_dy_acc_ready.init() + bars.mb_neg_dy_inp_ready.init() + bars.mb_dy_smem_ready.init() + bars.mb_dy_smem_done.init() + bars.mb_dstate_acc_ready.init() + bars.mb_dstate_inp_ready.init() + bars.mb_dstate_smem_ready.init() + bars.mb_dstate_smem_done.init() + bars.mb_dstate_smem_cg2_done.init() + bars.mb_dq_acc_ready.init() + bars.mb_dk_decay_part_acc_ready.init() + bars.mb_dk_inv_part_acc_ready.init() + bars.mb_dk_restore_part_acc_ready.init() + bars.mb_dqk_acc_done.init() + bars.mb_dstate0_acc_stored.init() + bars.mb_tmem_done[0].init() + elif warp_idx == cfg.super_mma_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_decay_stages): + bars.mb_k_decay_inv_ready[stage].init() + bars.mb_q_decay_k_restore_ready[stage].init() + bars.mb_decay_done[stage].init() + for stage in cutlass.range_constexpr(cfg.tmem_qk_raw_stages): + bars.mb_qk_raw_ready[stage].init() + bars.mb_qk_raw_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_intermediate_stages): + bars.mb_t_inv_ready[stage].init() + bars.mb_a_ready[stage].init() + bars.mb_da_ready[stage].init() + bars.mb_dm_ready[stage].init() + bars.mb_a_done[stage].init() + bars.mb_t_inv_done[stage].init() + bars.mb_da_done[stage].init() + bars.mb_dm_done[stage].init() + elif warp_idx == cfg.epilogue_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_dq_stages): + bars.mb_dq_tmastg_ready[stage].init() + bars.mb_dq_tmastg_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_dk_stages): + bars.mb_dk_tmastg_ready[stage].init() + bars.mb_dk_tmastg_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_dgate_stages): + bars.mb_dgate_tmastg_ready[stage].init() + bars.mb_dgate_tmastg_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_db_stages): + bars.mb_db_tmastg_ready[stage].init() + bars.mb_db_tmastg_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_dv_stages): + bars.mb_dv_tmastg_ready[stage].init() + bars.mb_dv_tmastg_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_dwo_stages): + bars.mb_dwo_tmastg_ready[stage].init() + bars.mb_dwo_tmastg_done[stage].init() + for stage in cutlass.range_constexpr(cfg.sched_stages): + bars.mb_sched_ready[stage].init() + bars.mb_sched_done[stage].init() + diag_zero = cfg.io_dtype(0.0) + for diag_idx in cutlass.range(tidx, cfg.diag_cosize, cfg.threads_per_cta, unroll=1): + sState_scale_diag_raw[diag_idx] = diag_zero + nvvm.fence_mbarrier_init() + nvvm.barrier_cta_sync(0, thread_count=cfg.threads_per_cta) + if warp_idx == cfg.tma_warp_id: + tmaldg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + mSched, + sSched, + q_tx_bytes, + k_tx_bytes, + gate_tx_bytes, + beta_tx_bytes, + do_tx_bytes, + v_tx_bytes, + w_tx_bytes, + sQ_raw, + sK_raw, + sV_raw, + sGate_raw, + sDo_raw, + sBeta_raw, + sW_raw, + sState_raw, + desc_q_base, + desc_k_base, + desc_v_base, + desc_gate_base, + desc_do_base, + desc_beta_base, + desc_w_base, + desc_checkpoint_base, + desc_initial_state_base, + bars, + ) + elif warp_idx == cfg.super_mma_warp_id: + super_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_decay_raw, + sK_inv_raw, + sU_raw, + sDy_raw, + sIntermediate_raw, + bars, + ) + elif warp_idx == cfg.tcgen05_mma_warp_id: + tcgen05_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + tmem_hold, + sState_alt, + sK_decay_lead16, + sK_inv_amaj, + sK_restore_lead16, + sDo_lead16, + sDo_amaj, + sQ_decay_trans, + sK_decay_trans, + sU_lead16, + sDy_lead16, + sDstate_alt, + sIntermediate, + sState_scale_diag, + bars, + ) + elif warp_idx == cfg.epilogue_warp_id: + epilogue_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_inv_raw, + sQ_decay_raw, + sDo_raw, + sU_raw, + sIntermediate_raw, + sDq_raw, + sDk_raw, + sDv_raw, + sDgate_raw, + sDb_raw, + sDwOut_raw, + desc_dq_base, + desc_dk_base, + desc_dv_base, + desc_dgate_base, + desc_db_base, + desc_dwo_base, + bars, + ) + elif warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: + compute0_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_hold, + warp_idx, + scale, + sK_inv_raw, + sGate_raw, + sK_raw, + sQ_raw, + sState_raw, + sV_raw, + sDo_raw, + sBeta_raw, + sW_raw, + sNorm_raw, + sK_decay_raw, + sK_restore_raw, + sQ_decay_raw, + sState_scale_diag_raw, + bars, + ) + elif warp_idx >= cfg.compute_group_2_warp_ids[0] and warp_idx <= cfg.compute_group_2_warp_ids[-1]: + compute2_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_hold, + warp_idx, + sBeta_raw, + sGate_raw, + sNorm_raw, + sDq_raw, + sDk_raw, + sRed1_raw, + sDstate_raw, + sDgate_raw, + sDb_raw, + scale, + bars, + ) + elif warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]: + compute1_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_hold, + warp_idx, + mDstate0, + mDstate_in, + sV_raw, + sW_raw, + sDo_raw, + sU_raw, + sDy_raw, + sDv_raw, + sDwOut_raw, + sDstate_raw, + bars, + ) + + +@dataclass +class Gdn2BwdCfg: + """Kernel cfg (fixed BT=16 schedule constants; derived TMEM column offsets + and SMEM buffer cosizes are stamped by ``build_cfg``).""" + + io_dtype: Type[cutlass.Numeric] + use_dstate_in: bool + use_dstate0: bool + l2norm: bool + use_initial_state: bool + q_ratio: int + k_ratio: int + v_ratio: int + n_heads_out: int + max_active_clusters: int + dyn_sched: bool = False + sched_stages: int = 8 + + # ---- fixed constants stamped from CFG by build_cfg --------------------------- + compute_group_0_warp_ids: tuple = CFG.COMPUTE_GROUP_0_WARP_IDS + compute_group_2_warp_ids: tuple = CFG.COMPUTE_GROUP_2_WARP_IDS + compute_group_1_warp_ids: tuple = CFG.COMPUTE_GROUP_1_WARP_IDS + super_mma_warp_id: int = CFG.SUPER_MMA_WARP_ID + tcgen05_mma_warp_id: int = CFG.TCGEN05_MMA_WARP_ID + tma_warp_id: int = CFG.TMA_WARP_ID + epilogue_warp_id: int = CFG.EPILOGUE_WARP_ID + b_t: int = CFG.B_T + d_k: int = CFG.D_K + d_v: int = CFG.D_V + threads_per_warp: int = CFG.THREADS_PER_WARP + threads_per_cta: int = 0 + num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 + num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 + num_regs_compute_group_2: int = CFG.NUM_REGS_COMPUTE_GROUP_2 + num_regs_other: int = CFG.NUM_REGS_OTHER + + # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ------------------- + cg0_sync_barrier_id: int = 1 + cg0_threads: int = 0 + cg2_sync_barrier_id: int = 2 + cg2_threads: int = 0 + tmem_lifecycle_barrier_id: int = 3 + tmem_user_threads: int = 0 + cg1_sync_barrier_id: int = 4 + cg1_threads: int = 0 + + # ---- SMEM / TMEM stage counts + TMEM column offsets -------------------------- + smem_raw_stages: int = CFG.SMEM_RAW_STAGES + smem_state_stages: int = CFG.SMEM_STATE_STAGES + smem_decay_stages: int = CFG.SMEM_DECAY_STAGES + smem_intermediate_stages: int = CFG.SMEM_INTERMEDIATE_STAGES + smem_dq_stages: int = CFG.SMEM_DQ_STAGES + smem_dk_stages: int = CFG.SMEM_DK_STAGES + smem_dgate_stages: int = CFG.SMEM_DGATE_STAGES + smem_db_stages: int = CFG.SMEM_DB_STAGES + smem_dv_stages: int = CFG.SMEM_DV_STAGES + smem_dwo_stages: int = CFG.SMEM_DWO_STAGES + intermediate_tiles: int = 5 + tmem_dstate_acc_offset: int = 0 + tmem_dstate_inp_offset: int = 0 + tmem_state_k_acc_offset: int = 0 + tmem_u_acc_offset: int = 0 + tmem_du_acc_offset: int = 0 + tmem_dy_acc_offset: int = 0 + tmem_dq_acc_offset: int = 0 + tmem_dk_decay_acc_offset: int = 0 + tmem_dk_inv_acc_offset: int = 0 + tmem_dk_restore_acc_offset: int = 0 + tmem_qk_raw_stages: int = 4 + tmem_qraw_inp_offset: int = 0 + tmem_kraw_inp_offset: int = 0 + tmem_y_inp_offset: int = 0 + tmem_du_inp_offset: int = 0 + tmem_neg_dy_inp_offset: int = 0 + tmem_state_inp_offset: int = 0 + buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES + + # ---- buffer cosizes / TMA bytes stamped by build_cfg ------------------------- + raw_qk_cosize: int = 0 + raw_v_cosize: int = 0 + raw_gate_cosize: int = 0 + operand_cosize: int = 0 + diag_cosize: int = 0 + intermediate_cosize: int = 0 + state_cosize: int = 0 + dq_cosize: int = 0 + dk_cosize: int = 0 + dgate_cosize: int = 0 + db_cosize: int = 0 + dv_cosize: int = 0 + dwo_cosize: int = 0 + tma_state_bytes: int = 0 + + +def build_cfg( + io_dtype: Type[cutlass.Numeric], + *, + use_dstate_in: bool, + use_dstate0: bool, + l2norm: bool, + use_initial_state: bool, + q_ratio: int, + k_ratio: int, + v_ratio: int, + n_heads_out: int, + max_active_clusters: int, + dyn_sched: bool = False, +) -> Gdn2BwdCfg: + if io_dtype not in (cutlass.Float16, cutlass.BFloat16): + raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") + cfg = Gdn2BwdCfg( + io_dtype=io_dtype, + use_dstate_in=use_dstate_in, + use_dstate0=use_dstate0, + l2norm=l2norm, + use_initial_state=use_initial_state, + q_ratio=q_ratio, + k_ratio=k_ratio, + v_ratio=v_ratio, + n_heads_out=n_heads_out, + max_active_clusters=max_active_clusters, + dyn_sched=dyn_sched, + ) + cfg.threads_per_cta = 16 * cfg.threads_per_warp + cfg.cg0_threads = len(cfg.compute_group_0_warp_ids) * cfg.threads_per_warp + cfg.cg2_threads = len(cfg.compute_group_2_warp_ids) * cfg.threads_per_warp + cfg.cg1_threads = len(cfg.compute_group_1_warp_ids) * cfg.threads_per_warp + cfg.tmem_user_threads = ( + 1 + len(cfg.compute_group_2_warp_ids) + len(cfg.compute_group_1_warp_ids) + len(cfg.compute_group_0_warp_ids) + ) * cfg.threads_per_warp + + cfg.tmem_dstate_acc_offset = 0 + cfg.tmem_dstate_inp_offset = cfg.d_k + cfg.tmem_state_inp_offset = cfg.tmem_dstate_inp_offset + cfg.d_k // 2 + cfg.tmem_state_k_acc_offset = cfg.tmem_state_inp_offset + cfg.d_v + cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t + cfg.tmem_du_acc_offset = cfg.tmem_u_acc_offset + cfg.b_t + # dY overwrites the state_k slot: WG1's Y staging consumes state_k + # before the dY = dU @ T_inv MMA writes (du_inp chain), and the dY + # readback precedes state_k(c+1) = state @ K_decay^T via + # neg_dy_ready -> the -dY @ K_decay dstate MMA -> in-order MMA + cfg.tmem_dy_acc_offset = cfg.tmem_state_k_acc_offset + cfg.tmem_dq_acc_offset = cfg.tmem_du_acc_offset + cfg.b_t + cfg.tmem_dk_decay_acc_offset = cfg.tmem_dq_acc_offset + cfg.b_t + cfg.tmem_dk_inv_acc_offset = cfg.tmem_dk_decay_acc_offset + cfg.b_t + cfg.tmem_dk_restore_acc_offset = cfg.tmem_dk_inv_acc_offset + cfg.b_t + cfg.tmem_y_inp_offset = cfg.tmem_dk_restore_acc_offset + cfg.b_t + # -dY overwrites the y_inp slot: U = Y @ T_inv consumed Y before the dY + # block runs (u_acc_ready wait), and y_inp(c+1) is gated by + # state_k_acc_ready(c+1) whose commit covers the -dY @ K_decay MMA (c) + cfg.tmem_neg_dy_inp_offset = cfg.tmem_y_inp_offset + cfg.tmem_du_inp_offset = cfg.tmem_y_inp_offset + cfg.b_t // 2 + cfg.tmem_qraw_inp_offset = cfg.tmem_du_inp_offset + cfg.b_t // 2 + cfg.tmem_kraw_inp_offset = cfg.tmem_qraw_inp_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) + assert cfg.tmem_kraw_inp_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) <= 512 + + cfg.raw_qk_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t + cfg.raw_v_cosize = cfg.smem_raw_stages * cfg.d_v * cfg.b_t + cfg.raw_gate_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t + cfg.operand_cosize = cfg.smem_decay_stages * cfg.b_t * cfg.d_k + cfg.diag_cosize = cfg.smem_decay_stages * (cfg.d_k // 16) * 256 + cfg.intermediate_cosize = cfg.smem_intermediate_stages * cfg.intermediate_tiles * cfg.b_t * cfg.b_t + cfg.state_cosize = cfg.smem_state_stages * cfg.d_k * cfg.d_v + cfg.dq_cosize = cfg.smem_dq_stages * cfg.b_t * cfg.d_k + cfg.dk_cosize = cfg.smem_dk_stages * cfg.b_t * cfg.d_k + cfg.dgate_cosize = cfg.smem_dgate_stages * cfg.b_t * cfg.d_k + cfg.db_cosize = cfg.smem_db_stages * cfg.b_t * cfg.d_k + cfg.dv_cosize = cfg.smem_dv_stages * cfg.b_t * cfg.d_v + cfg.dwo_cosize = cfg.smem_dwo_stages * cfg.b_t * cfg.d_v + cfg.tma_state_bytes = cfg.d_k * cfg.d_v * (io_dtype.width // 8) + return 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 + + +# ---- Torch adapter / host-side compilation --------------------------------------- + + +@lru_cache(maxsize=None) +def get_compiled_cache( + io_dtype_str: str, + cu_dtype_str: str, + HQ: int, + HK: int, + HV: int, + use_dstate_in: bool, + use_dstate0: bool, + l2norm: bool, + use_initial_state: bool, + dyn_sched: bool, +): + return {} + + +def chunk_gdn2_bwd_sm100( + q, + k, + v, + gate, + beta, + w, + do, + state_checkpoints, + dq, + dk, + dv, + dgate, + dbeta, + dw, + cu_seqlens, + scale: float, + *, + initial_state=None, + d_initial_state=None, + d_final_state=None, + use_qk_l2norm_in_kernel: bool = False, + work_items=None, + work_count=None, + sched_ctr=None, + tensormap_workspace, + stream, +) -> None: + """Execute the Blackwell BT=16 chunked GDN-2 backward kernel. + + All tensors must be contiguous and on the same CUDA device. + + Args: + q: ``(total_tokens, HQ, DK)`` float16/bfloat16 + k: ``(total_tokens, HK, DK)`` float16/bfloat16 + v: ``(total_tokens, HV, DV)`` float16/bfloat16 + gate: ``(total_tokens, HO, DK)`` fp32 natural-log per-channel decay + beta: ``(total_tokens, HO, DK)`` io dtype post-sigmoid per-key erase + 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`` + 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 + dbeta: ``(total_tokens, HO, DK)`` io dtype, pre-allocated + 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_qk_l2norm_in_kernel: q/k arrive raw; the kernel normalizes for the + recompute math and chains the L2-norm backward into dq/dk + work_items/work_count: split-K table (``common/split_k.py``, REQUIRED; + an uncut table row is the whole (b, h) sequence); each item + computes chunks ``[wstart, cend)`` backward and writes + gradients only for ``[wstart, wend)`` + sched_ctr: ``(2,)`` int32 zeroed scratch enabling the dynamic + (work-stealing) tile scheduler + tensormap_workspace: ``tensormap_workspace_bytes(module, B)`` bytes, + 128-byte aligned, for the per-(batch, head) TMA-descriptor + arrays (tail chunks clip/zero-fill in hardware) + """ + HQ = q.shape[1] + HK = k.shape[1] + HV = v.shape[1] + 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 + 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 () + ): + 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)): + if HO % hh != 0: + raise ValueError(f"{name}={hh} must divide {HO}") + B = cu_seqlens.shape[0] - 1 + + cu_stream = cuda_driver.CUstream(int(stream)) + cache = get_compiled_cache( + str(q.dtype), + str(cu_seqlens.dtype), + HQ, + HK, + HV, + use_dstate_in, + use_dstate0, + use_qk_l2norm_in_kernel, + use_initial_state, + dyn_sched, + ) + + if "compiled" not in cache: + io_dtype = get_dtype(q.dtype) + cfg = build_cfg( + io_dtype, + use_dstate_in=use_dstate_in, + use_dstate0=use_dstate0, + l2norm=use_qk_l2norm_in_kernel, + use_initial_state=use_initial_state, + q_ratio=HO // HQ, + k_ratio=HO // HK, + v_ratio=HO // HV, + n_heads_out=HO, + max_active_clusters=multiprocessor_count(current_device_id()), + dyn_sched=dyn_sched, + ) + + dstate0_cute = None + if use_dstate0: + dstate0_cute = from_dlpack(d_initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) + dstate_in_cute = None + if use_dstate_in: + dstate_in_cute = from_dlpack(d_final_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) + wi_cute = from_dlpack(work_items, assumed_align=16) + wi_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + wc_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() + sc_cute = None + if dyn_sched: + sc_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + + 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 + ) + dgate_cute = from_dlpack(dgate, assumed_align=16).mark_layout_dynamic(leading_dim=len(dgate.shape) - 1) + dbeta_cute = from_dlpack(dbeta, assumed_align=16).mark_layout_dynamic(leading_dim=len(dbeta.shape) - 1) + dw_cute = from_dlpack(dw, assumed_align=16).mark_layout_dynamic(leading_dim=len(dw.shape) - 1) + cache["compiled"] = cute.compile( + host, + cfg, + state_checkpoints_cute, + initial_state_cute, + dgate_cute, + dbeta_cute, + dw_cute, + from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic(), + dstate0_cute, + dstate_in_cute, + wi_cute, + wc_cute, + sc_cute, + tensormap_ws_cute, + scale, + cu_stream, + options="--enable-tvm-ffi --opt-level 2", + ) -def get_workspace_size(B: int, HQ: int, HV: int) -> int: - return 0 + # ---- per-(batch, head) descriptor arrays: rebuild on input change ------------ + # desc build runs every execute by contract (cu contents are data; + # buffer pointers may change) - capture-safe, single tiny launch + if "build_descs" not in cache: + io_dtype = get_dtype(q.dtype) + q_bd = from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=2) + k_bd = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_bd = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_bd = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2) + do_bd = from_dlpack(do, assumed_align=16).mark_layout_dynamic(leading_dim=2) + beta_bd = from_dlpack(beta, assumed_align=16).mark_layout_dynamic(leading_dim=2) + w_bd = from_dlpack(w, assumed_align=16).mark_layout_dynamic(leading_dim=2) + dq_bd = from_dlpack(dq, assumed_align=16).mark_layout_dynamic(leading_dim=2) + dk_bd = from_dlpack(dk, assumed_align=16).mark_layout_dynamic(leading_dim=2) + dv_bd = from_dlpack(dv, assumed_align=16).mark_layout_dynamic(leading_dim=2) + dgate_bd = from_dlpack(dgate, assumed_align=16).mark_layout_dynamic(leading_dim=2) + dwo_bd = from_dlpack(dw, assumed_align=16).mark_layout_dynamic(leading_dim=2) + dbo_bd = from_dlpack(dbeta, assumed_align=16).mark_layout_dynamic(leading_dim=2) + state_checkpoints_bd = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=3) + initial_state_bd = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) if use_initial_state else None -def chunk_gdn2_bwd_sm100(*args, **kwargs) -> None: - """Not implemented — the FROST GDN-2 backward kernel is a stub.""" - raise NotImplementedError("FROST GDN-2 backward is not implemented yet (recompute-in-bprop kernel is a stub).") + cu_bd = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() + ws_bd = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + cache["build_descs"] = cute.compile( + build_descs, + io_dtype, + CFG.B_T, + q_bd, + k_bd, + v_bd, + gate_bd, + do_bd, + beta_bd, + w_bd, + dq_bd, + dk_bd, + dv_bd, + dgate_bd, + dwo_bd, + dbo_bd, + state_checkpoints_bd, + initial_state_bd, + cu_bd, + ws_bd, + cu_stream, + options="--enable-tvm-ffi", + ) + cache["build_descs"](q, k, v, gate, do, beta, w, dq, dk, dv, dgate, dw, dbeta, state_checkpoints, initial_state, cu_seqlens, tensormap_workspace, cu_stream) + cache["compiled"]( + state_checkpoints, + initial_state, + dgate, + dbeta, + dw, + cu_seqlens, + d_initial_state, + d_final_state, + work_items, + work_count, + sched_ctr, + tensormap_workspace, + scale, + cu_stream, + ) diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_config.py b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_config.py index 3749cb07c..495f6cf82 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_config.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_config.py @@ -38,10 +38,10 @@ class Cfg: # --- warp assignments (16 warps = 512 threads) --- COMPUTE_GROUP_0_WARP_IDS: Tuple[int, ...] = (0, 1, 2, 3, 4, 5, 6, 7) # decay/beta-operand materialize COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (8, 9, 10, 11) # value-side TMEM (w*v - erase) / epilogue - SUPER_MMA_WARP_ID: int = 12 # register-MMA kk/qk + Neumann inverse + SUPER_MMA_WARP_ID: int = 12 # register-MMA KK/A + Neumann T_inv TCGEN05_MMA_WARP_ID: int = 13 # tcgen05 state GEMMs TMA_WARP_ID: int = 14 # q/k/v/gate/beta/w TMA loads - EPILOGUE_WARP_ID: int = 15 # qk register-MMA + O store + EPILOGUE_WARP_ID: int = 15 # A register-MMA + O store # --- register split --- NUM_REGS_COMPUTE_GROUP_0: int = 160 @@ -57,9 +57,9 @@ class Cfg: SMEM_SCHED_STAGES: int = 8 SMEM_O_STAGES: int = 2 SMEM_DECAY_STAGES: int = 2 - SMEM_PAIRWISE_STAGES: int = 2 - SMEM_STATE_SCALE_DIAG_STAGES: int = 3 - QK_SCALE_READY_STAGES: int = 3 + SMEM_INTERMEDIATE_STAGES: int = 2 + SMEM_STATE_SCALE_DIAG_STAGES: int = 4 + QK_SCALE_READY_STAGES: int = 4 TMEM_Q_STATE_ACC_STAGES: int = 2 CLUSTER_SHAPE_MNK: Tuple[int, int, int] = (1, 1, 1) 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 4203786d8..2e97e2fc8 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py @@ -15,68 +15,68 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Chunked Gated DeltaNet v2 (GDN-2) prefill kernel for Blackwell SM100 +"""Chunked Gated DeltaNet v2 (GDN-2) prefill kernel for Blackwell SM100/SM103 (Cutlass DSL), BT=16 tiling with per-key-channel decay + per-key erase gate -(beta) + per-value write gate (w), using direct CUTLASS primitives. +(beta) + per-value write gate (W), using direct CUTLASS primitives. Framework-neutral entry ``chunk_gdn2_sm100``. Extends the KDA BT=16 schedule to the channel-wise gated delta rule: S_t = S_{t-1} * diag(exp(g_t)) - v_new_t = w_t * v_t - (beta_t * k_t)^T S_t - S_t += k_t (x) v_new_t - o_t = scale * q_t^T S_t + U_t = W_t * V_t - (beta_t * K_t)^T S_t + S_t += K_t (x) U_t + O_t = scale * Q_t^T S_t -vs KDA: the erase gate beta and write gate w are per-channel tensors. beta -is folded into the k_decay operand (feeds KK^T and state*k), the strict-lower -tile loses its per-row beta scale, and RHS becomes `w*v - state*k`; beta/w -arrive by TMA alongside q/k/v. +vs KDA: the erase gate beta and write gate W are per-channel tensors. beta +is folded into the k_decay operand (feeds KK^T and state*K), the strict-lower +tile loses its per-row beta scale, and Y becomes `W*V - state*K`; beta/W +arrive by TMA alongside Q/K/V. 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 int64, states/checkpoints -`[N, HO, DV, DK]` (VK). GQA/GVA head broadcast follows repeat_interleave: +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: source head = head_idx // (HO // H_x). State presence, L2NORM, SAFE_GATE, checkpoints, and the head ratios are compile-time specializations. Warp assignments (16 warps = 512 threads): - warps 0-7 : compute group 0 - gate prefix scan + decay/restore operands - warps 8-11 : compute group 1 - TMEM value side, o drain, state stores - warp 12 : super-MMA - register-MMA kk^T + Neumann inverse + warps 0-7 : compute group 0 - Gate prefix scan + decay/restore operands + warps 8-11 : compute group 1 - TMEM value side, O drain, state stores + warp 12 : super-MMA - register-MMA KK^T + Neumann inverse warp 13 : tcgen05-MMA - the six state GEMMs + the TMEM lifecycle warp 14 : TMA load - per-chunk input G->S loads - warp 15 : epilogue - register-MMA qk + the O TMA store + warp 15 : epilogue - register-MMA A + the O TMA store SMEM layout (~205 KB total): Buffer Bytes Stages - q / k / v raw 20480 5 <-- SW128 TMA ring (io dtype) - beta / w raw 2x 20480 5 <-- per-channel gates, same ring - gate raw 40960 5 <-- fp32 prefix-scan source + Q / K / V raw 20480 5 <-- SW128 TMA ring (io dtype) + Beta / W raw 2x 20480 5 <-- per-channel gates, same ring + Gate raw 40960 5 <-- fp32 prefix-scan source dt_bias (+a_log slot) 516 1 <-- SAFE_GATE only K_inv 8192 2 <-- token-major ldmatrix/tcgen05 B operand K decay / Q decay 2x 8192 2 <-- tcgen05 SW128 K-box-major A/B operands K restore 8192 2 <-- tcgen05 B operand for the state update state-scale diag 12288 3 <-- per-k-atom decay diagonal blocks - pairwise (A_inv / qk) 2048 2 <-- SW32 16x16 register-MMA tiles - o staging 8192 2 <-- W128 output drain + Intermediate (T_inv / A) 2048 2 <-- SW32 16x16 register-MMA tiles + O staging 8192 2 <-- W128 output drain TMEM layout (272 of 512 columns): Buffer Cols Purpose - state 0-127 S[DK,DV] fp32 recurrent state + state 0-127 state[DK,DV] fp32 recurrent state state inp 128-191 packed b16 A operand view of the state - q_state_acc 192-223 2-stage state*q -> o accumulator - state_k_acc 224-239 state*k fp32 accumulator - update_acc 240-255 update fp32 accumulator - rhs input 256-263 packed b16 A operand: w*v - state*(beta*k) - update input 264-271 packed b16 A operand: the update readback + q_state_acc 192-223 2-stage state*Q -> O accumulator + state_k_acc 224-239 state*K fp32 accumulator + u_acc 240-255 U fp32 accumulator + y_inp 256-263 packed b16 A operand: Y = W*V - state*(Beta*K) + u_inp 264-271 packed b16 A operand: the U readback GEMM schedule (tcgen05-MMA warp, in issue order per chunk): - state*k -> state_k_acc - state*q -> q_state_acc (the o acc) + state*K -> state_k_acc + state*Q -> q_state_acc (the O acc) state decay (diag blocks) - update = A_inv @ rhs -> update_acc - final_state += update @ k_restore - o += qk @ update -> q_state_acc + U = Y(T) @ T_inv -> u_acc + final_state += U @ K_restore + O += A @ U -> q_state_acc Requires a cutlass DSL build providing `cutlass.experimental.*`; not available in the pip nvidia-cutlass-dsl releases. @@ -94,28 +94,30 @@ from cutlass.cute.runtime import from_dlpack from ..common.split_k import decode_work_item -from ..common.thd import TENSOR_MAP_QWORDS, build_h_descs_kernel, build_qkv_load_descs_kernel +from ..common.host import get_dtype +from cudnn.frost.buffers import current_device_id, data_ptr +from cudnn.frost.device import multiprocessor_count +from ..common.thd import TENSOR_MAP_QWORDS, emit_checkpoint_seq_descs, emit_seq_descs from .gdn2_prefill_config import CFG from cudnn.frost.tile_dsl.barrier import ( advance, - arrive, MBarrier, PipelineState, Producer, - wait, ) from cudnn.frost.tile_dsl.handles import GmemTileTma, MmaDesc, SmemTile, tma_slice_runtime_desc -from cudnn.frost.tile_dsl.mma import mma_step, mma_ts +from cudnn.frost.tile_dsl.mma import mma_step, mma_ts_step from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_lin_S, swizzle_xor_128b from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire from cudnn.frost.tile_dsl.pointwise import ( + f16x2_to_f32, fadd2, fmul2, + ffma2, movmatrix_16b, mul_f16x2, opaque_f32_zero, fp32_to_fp16, - sigmoid_f16x2, sub_f16x2, ) @@ -132,39 +134,56 @@ class Gdn2Bars(NamedTuple): - """Every inter-warp handoff as an ``MBarrier`` over its ring (mirrors GDN's - ``GdnBars`` and KDA's ``KdaBars``).""" + """Every inter-warp handoff as an ``MBarrier`` over its ring.""" + + mb_q_ready: MBarrier + mb_q_done: MBarrier + mb_k_ready: MBarrier + mb_k_done: MBarrier + mb_v_ready: MBarrier + mb_v_done: MBarrier + mb_w_ready: MBarrier + mb_w_done: MBarrier + + mb_gate_ready: MBarrier + mb_gate_done: MBarrier + mb_beta_ready: MBarrier + mb_beta_done: MBarrier - mb_tma_done: MBarrier - mb_inputs_ready: MBarrier - mb_inputs_done: MBarrier mb_o_acc_ready: MBarrier mb_o_acc_done: MBarrier mb_state_k_acc_ready: MBarrier - mb_update_acc_ready: MBarrier + mb_u_acc_ready: MBarrier + mb_state_inp_ready: MBarrier - mb_state_scale_diag_done: MBarrier - mb_kk_qk_super_mma_done: MBarrier - mb_kk_qk_mma_done: MBarrier - mb_k_restore_done: MBarrier - mb_rhs_ready: MBarrier - mb_update_ready: MBarrier - mb_final_state_stored: MBarrier + mb_y_inp_ready: MBarrier + mb_u_inp_ready: MBarrier + + mb_t_inv_ready: MBarrier + mb_intermediate_done: MBarrier mb_a_ready: MBarrier - mb_qk_acc_ready: MBarrier - mb_a_done: MBarrier + mb_k_decay_inv_cg0_ready: MBarrier + mb_decay_tcgen05_done: MBarrier + mb_decay_super_done: MBarrier mb_qk_scale_ready: MBarrier - mb_k_decay_cg0_ready: MBarrier + mb_k_restore_acc_done: MBarrier + mb_state_scale_diag_done: MBarrier + + mb_tmem_done: MBarrier + mb_state_acc_read_done: MBarrier + mb_o_tmastg_ready: MBarrier mb_o_tmastg_done: MBarrier + + mb_checkpoint_tmastg_ready: MBarrier + mb_checkpoint_tmastg_done: MBarrier + mb_sched_ready: MBarrier mb_sched_done: MBarrier - mb_h_tmastg_ready: MBarrier - mb_h_tmastg_done: MBarrier def make_gdn2_bars(cfg) -> Gdn2Bars: - """Bars factory. MUST be called from inside ``_kernel`` (allocates the + """Bars factory. MUST be called from inside ``kernel`` (allocates the mbarrier rings in SMEM ahead of the data buffers).""" def alloc(n): @@ -175,71 +194,62 @@ def alloc(n): CG1_THREADS = len(cfg.compute_group_1_warp_ids) * WARP return Gdn2Bars( - mb_tma_done=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.TMA_LOAD), - mb_inputs_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=WARP, producer=Producer.THREAD), - mb_inputs_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_q_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_q_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_k_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_k_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_v_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_v_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_w_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_w_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_gate_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_gate_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_beta_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_beta_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), mb_o_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), mb_o_acc_done=MBarrier(alloc(cfg.tmem_q_state_acc_stages), stages=cfg.tmem_q_state_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), - mb_update_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), + mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), mb_state_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_state_scale_diag_done=MBarrier( - alloc(cfg.smem_state_scale_diag_stages), - stages=cfg.smem_state_scale_diag_stages, - init_count=1, - producer=Producer.MMA_COMMIT, - ), - mb_kk_qk_super_mma_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=2 * WARP, producer=Producer.THREAD), - mb_kk_qk_mma_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), - mb_k_restore_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), - mb_rhs_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_update_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_final_state_stored=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_a_ready=MBarrier(alloc(cfg.smem_pairwise_stages), stages=cfg.smem_pairwise_stages, init_count=WARP, producer=Producer.THREAD), - mb_qk_acc_ready=MBarrier(alloc(cfg.smem_pairwise_stages), stages=cfg.smem_pairwise_stages, init_count=WARP, producer=Producer.THREAD), - mb_a_done=MBarrier(alloc(cfg.smem_pairwise_stages), stages=cfg.smem_pairwise_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_t_inv_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_intermediate_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_a_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_k_decay_inv_cg0_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_decay_tcgen05_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_decay_super_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=2 * WARP, producer=Producer.THREAD), mb_qk_scale_ready=MBarrier( alloc(cfg.qk_scale_ready_stages), stages=cfg.qk_scale_ready_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD, ), - mb_k_decay_cg0_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_k_restore_acc_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_state_scale_diag_done=MBarrier( + alloc(cfg.smem_state_scale_diag_stages), + stages=cfg.smem_state_scale_diag_stages, + init_count=1, + producer=Producer.MMA_COMMIT, + ), + mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_state_acc_read_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_o_tmastg_ready=MBarrier(alloc(cfg.smem_o_stages), stages=cfg.smem_o_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_o_tmastg_done=MBarrier(alloc(cfg.smem_o_stages), stages=cfg.smem_o_stages, init_count=WARP, producer=Producer.THREAD), + mb_checkpoint_tmastg_ready=MBarrier( + alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=CG1_THREADS, producer=Producer.THREAD + ), + mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=WARP, producer=Producer.THREAD), mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), - # H staging handshake: CG1 fills sH (ready), the epilogue TMA-stores - # and frees it (done); single stage - mb_h_tmastg_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_h_tmastg_done=MBarrier(alloc(1), stages=1, init_count=WARP, producer=Producer.THREAD), ) -# --------------------------------------------------------------------------- -# Device-side helpers / warp bodies -# --------------------------------------------------------------------------- - - -@cute.jit -def _gate_log2(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: - """Map raw gate to the log2-domain decay increment used by KDA.""" - - if cutlass.const_expr(cfg.safe_gate): - half = cutlass.Float32(0.5) - sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half - return cfg.gate_scale_log2 * sigmoid - # Default ABI: gate arrives in natural-log space - return raw_gate * cutlass.Float32(LOG2_E) - - -# --------------------------------------------------------------------------- -# Dynamic tile scheduler: global-ticket ring -# --------------------------------------------------------------------------- +# ---- Dynamic tile scheduler ------------------------------------------------------ @cute.jit -def _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): """TMA-warp side: pull the next tile off the global ticket, publish it.""" if cutlass.const_expr(cfg.dyn_sched): bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) @@ -255,7 +265,7 @@ def _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ct @cute.jit -def _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): """Consumer side: read the TMA warp's published next tile.""" if cutlass.const_expr(cfg.dyn_sched): bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) @@ -267,26 +277,7 @@ def _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): @cute.jit -def _decay_key_dim(cfg, tok_coord, key_dim): - """Return the runtime key coordinate for tcgen05 SW128 decay operands.""" - - key_mask = cutlass.Int32(8) ^ ((tok_coord & cutlass.Int32(2)) * cutlass.Int32(16)) - return key_dim ^ key_mask - - -@cute.jit -def _diag_idx(cfg, key_dim): - """Return the SW32 index for one entry in the 8-block diagonal.""" - - block = key_dim // cutlass.Int32(16) - coord = key_dim - block * cutlass.Int32(16) - storage_col = coord ^ cutlass.Int32((cfg.b_t // 2)) - linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + storage_col - return swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) - - -@cute.jit -def _tmaldg_warp( +def tmaldg_warp( cfg, total_tiles, bidx, @@ -296,7 +287,6 @@ def _tmaldg_warp( mSched, sSched, lane, - tma_tx_bytes, sBeta_raw, sGate_raw, sK_raw, @@ -311,12 +301,9 @@ def _tmaldg_warp( desc_w_base, bars, ) -> None: - """TMA-LDG warp role (warp 14): persistent tile-scheduler loop + - per-chunk q/k/v/gate/beta/w G->S loads on one shared tx-count mbarrier. - Loads go through the per-(batch, head) descriptor array: head grouping - and the sequence base live in each descriptor and the token extent is - capped per sequence, so coordinates are sequence-relative and tail - chunks zero-fill in hardware.""" + """TMA-LDG warp role (warp 14): persistent scheduler loop issuing the + per-chunk Q/K/V/Beta/W/Gate G->S loads.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) sQ_tma = SmemTile( base=sQ_raw, @@ -384,22 +371,24 @@ def _tmaldg_warp( tma_granu_elems=32, tma_subtile_stride_elems=(cfg.b_t * 32), ) - tma_index = PipelineState.start(phase=0) raw_index = PipelineState.start(phase=1) + raw_bar_index = PipelineState.start(phase=0) sched_state = PipelineState.start(phase=1) 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, cu_seqlens, mWorkItems - ) - slot = (batch_idx * cutlass.Int32(cfg.n_heads_out) + head_idx) * cutlass.Int32(TENSOR_MAP_QWORDS) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) + head_o = head_idx + head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) + head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) + head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) + slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) desc_q_slot = (desc_q_base + slot).tospace(cutlass.AddressSpace.generic) desc_k_slot = (desc_k_base + slot).tospace(cutlass.AddressSpace.generic) desc_v_slot = (desc_v_base + slot).tospace(cutlass.AddressSpace.generic) desc_gate_slot = (desc_gate_base + slot).tospace(cutlass.AddressSpace.generic) desc_beta_slot = (desc_beta_base + slot).tospace(cutlass.AddressSpace.generic) desc_w_slot = (desc_w_base + slot).tospace(cutlass.AddressSpace.generic) - if nvvm.elect_sync(): + if elect_one: tma_tensormap_acquire(desc_q_slot) tma_tensormap_acquire(desc_k_slot) tma_tensormap_acquire(desc_v_slot) @@ -408,32 +397,55 @@ def _tmaldg_warp( tma_tensormap_acquire(desc_w_slot) for chunk_idx in cutlass.range(cstart, wend, 1, unroll=1): chunk_start = chunk_idx * cfg.b_t - bars.mb_inputs_done[raw_index.idx].wait(raw_index.phase) - # ---- q/k/v/gate/beta/w TMA loads ------------------------------------ - if nvvm.elect_sync(): - bars.mb_tma_done.arrive(n_bytes=tma_tx_bytes) - q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), chunk_start) - tma_load_tile(sQ_tma[raw_index.idx], q_slice, bars.mb_tma_done.smem_ptr, acquire=False) - k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), chunk_start) - tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_tma_done.smem_ptr, acquire=False) - v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), chunk_start) - tma_load_tile(sV_tma[raw_index.idx], v_slice, bars.mb_tma_done.smem_ptr, acquire=False) - beta_slice = tma_slice_runtime_desc(desc_beta_slot, cutlass.Int32(0), chunk_start) - tma_load_tile(sBeta_tma[raw_index.idx], beta_slice, bars.mb_tma_done.smem_ptr, acquire=False) - w_slice = tma_slice_runtime_desc(desc_w_slot, cutlass.Int32(0), chunk_start) - tma_load_tile(sW_tma[raw_index.idx], w_slice, bars.mb_tma_done.smem_ptr, acquire=False) - gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), chunk_start) - tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_tma_done.smem_ptr, acquire=False) - - bars.mb_tma_done.wait(tma_index.phase) - tma_index = advance(tma_index, 1) - bars.mb_inputs_ready[raw_index.idx].arrive() + + # ---- Q load ---------------------------------------------------------- + bars.mb_q_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_q_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_q_bytes) + q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, chunk_start) + tma_load_tile(sQ_tma[raw_index.idx], q_slice, bars.mb_q_ready[raw_bar_index.idx].smem_ptr, acquire=False) + + # ---- K load ---------------------------------------------------------- + bars.mb_k_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_k_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_k_bytes) + k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) + tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_bar_index.idx].smem_ptr, acquire=False) + + # ---- V load ---------------------------------------------------------- + bars.mb_v_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_v_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_v_bytes) + 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_bar_index.idx].smem_ptr, acquire=False) + + # ---- Beta load ------------------------------------------------------- + bars.mb_beta_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_beta_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_beta_bytes) + 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_bar_index.idx].smem_ptr, acquire=False) + + # ---- W load ---------------------------------------------------------- + bars.mb_w_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_w_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_w_bytes) + w_slice = tma_slice_runtime_desc(desc_w_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sW_tma[raw_index.idx], w_slice, bars.mb_w_ready[raw_bar_index.idx].smem_ptr, acquire=False) + + # ---- Gate load ------------------------------------------------------- + bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_gate_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_gate_bytes) + gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_bar_index.idx].smem_ptr, acquire=False) raw_index = advance(raw_index, cfg.smem_raw_stages) - tile_idx, sched_state = _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) + tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) @cute.jit -def _super_mma_warp( +def super_mma_warp( cfg, total_tiles, bidx, @@ -443,21 +455,20 @@ def _super_mma_warp( sSched, lane, sK_inv_raw, - sPairwise_raw, + sIntermediate_raw, sK_decay_raw, bars, ) -> None: - """Super-MMA warp role (warp 12): persistent tile-scheduler loop + - register-MMA kk^T, L = tril(kk, -1), and the Neumann-series A_inv, - staged to pairwise SMEM.""" + """Super-MMA warp role (warp 12): persistent scheduler loop computing the + Neumann-series T_inv by register MMA.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) - # ---- ldmatrix/stmatrix lane decode --------------------------------- + + # ---- ldmatrix/stmatrix lane decode ------------------------------------------- rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - decay_key_mask = cutlass.Int32(8) ^ ((lhs_row_coord & cutlass.Int32(2)) * cutlass.Int32(16)) - elems_per_128b = cutlass.Int32(64) + decay_key_mask = cutlass.Int32(8) stsm_row_coord = lane & 7 stsm_col_coord = cutlass.Int32(0) if (lane // 8) & 1: @@ -465,25 +476,23 @@ def _super_mma_warp( if lane // 8 >= 2: stsm_col_coord = cutlass.Int32(8) stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) - gbase = cutlass.Int32(0) + global_chunk_base = cutlass.Int32(0) 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, cu_seqlens, mWorkItems - ) - sk_nt = wend - cstart # processed chunks; ring bookkeeping runs on gbase + li - for li in cutlass.range(sk_nt, unroll=1): - gc = gbase + li - decay_stage = gc % cfg.smem_decay_stages - pairwise_stage = gc % cfg.smem_pairwise_stages + 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_tile_chunks = wend - cstart # processed chunks; ring bookkeeping runs on global_chunk_base + local_chunk + for local_chunk in cutlass.range(num_tile_chunks, unroll=1): + global_chunk = global_chunk_base + local_chunk + decay_stage = global_chunk % cfg.smem_decay_stages + intermediate_stage = global_chunk % cfg.smem_intermediate_stages sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) - sPairwise_ptr = sPairwise_raw.data_ptr() + pairwise_stage * (2 * cfg.b_t * cfg.b_t) + sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (2 * cfg.b_t * cfg.b_t) - bars.mb_a_done[pairwise_stage].wait(((gc // cfg.smem_pairwise_stages) + 1) % 2) - bars.mb_k_decay_cg0_ready[decay_stage].wait((gc // cfg.smem_decay_stages) % 2) - # ---- kk^T register MMA over the K blocks --------------------------- + bars.mb_k_decay_inv_cg0_ready[decay_stage].wait((global_chunk // cfg.smem_decay_stages) % 2) + + # ---- KK = K_decay @ K_inv^T ------------------------------------------ kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): kk_acc[accum_idx] = cutlass.Float32(0.0) @@ -492,7 +501,7 @@ def _super_mma_warp( # Load B operand k_inv_col = k_block * 16 + rhs_col_offset k_inv_segment = k_inv_col // 64 - rhs_vec = nvvm.ldmatrix( + rhs_frag = nvvm.ldmatrix( sK_inv_ptr + k_inv_segment * (cfg.b_t * 64) + rhs_row_coord * 64 @@ -502,35 +511,29 @@ def _super_mma_warp( ) # Load A operand storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask - storage_slice = storage_key // elems_per_128b - key_in_slice = storage_key - storage_slice * elems_per_128b - storage_phase = key_in_slice // cutlass.Int32(16) - byte_in_slice = ( - lhs_row_coord * cutlass.Int32(128) - + storage_phase * cutlass.Int32(32) - + (key_in_slice - storage_phase * cutlass.Int32(16)) * cutlass.Int32(2) - ) - kk_lhs_vec = nvvm.ldmatrix( + storage_slice = storage_key // 64 + kk_lhs_frag = nvvm.ldmatrix( sK_decay_ptr - + storage_slice * cutlass.Int32(cfg.b_t) * elems_per_128b - + ((byte_in_slice ^ ((lhs_row_coord & cutlass.Int32(7)) << 4)) // cutlass.Int32(2)), + + storage_slice * (cfg.b_t * 64) + + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( kk_acc, - (kk_lhs_vec[0], kk_lhs_vec[1], kk_lhs_vec[2], kk_lhs_vec[3]), - (rhs_vec[0], rhs_vec[1], rhs_vec[2], rhs_vec[3]), + (kk_lhs_frag[0], kk_lhs_frag[1], kk_lhs_frag[2], kk_lhs_frag[3]), + (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # ---- L = tril(kk, -1) fragment -------------------------------------- + + # ---- L = tril(KK, -1) fragment --------------------------------------- row_lo = lane // 4 row_hi = row_lo + cutlass.Int32(8) - l_frag = cutlass.Array(cutlass.Float32, 8, alignment=16) + l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_lo if cutlass.const_expr(accum_idx % 4 >= 2): @@ -538,15 +541,15 @@ def _super_mma_warp( col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) - l_frag[accum_idx] = kk_acc[accum_idx] if row_coord > col_coord else cutlass.Float32(0.0) - l_a0 = fp32_to_fp16(l_frag[0], l_frag[1], dtype=cfg.io_dtype) - l_a1 = fp32_to_fp16(l_frag[2], l_frag[3], dtype=cfg.io_dtype) - l_a2 = fp32_to_fp16(l_frag[4], l_frag[5], dtype=cfg.io_dtype) - l_a3 = fp32_to_fp16(l_frag[6], l_frag[7], dtype=cfg.io_dtype) + l_regs[accum_idx] = kk_acc[accum_idx] if row_coord > col_coord else cutlass.Float32(0.0) + l_a0 = fp32_to_fp16(l_regs[0], l_regs[1], dtype=cfg.io_dtype) + l_a1 = fp32_to_fp16(l_regs[2], l_regs[3], dtype=cfg.io_dtype) + l_a2 = fp32_to_fp16(l_regs[4], l_regs[5], dtype=cfg.io_dtype) + l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) - # ---- A_inv = I - L, then three Neumann doubling rounds ------------- - ainv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + # ---- T_inv = I - L, then three Neumann doubling rounds --------------- + tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_lo if cutlass.const_expr(accum_idx % 4 >= 2): @@ -555,18 +558,19 @@ def _super_mma_warp( if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) eye = cutlass.Float32(1.0) if row_coord == col_coord else cutlass.Float32(0.0) - ainv_acc[accum_idx] = eye - l_values[accum_idx] + tinv_acc[accum_idx] = eye - l_values[accum_idx] lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) for _round in cutlass.range_constexpr(3): - # Lpow <- Lpow @ Lpow (packed A-layout fragments, B via movmatrix) + # ---- Lpow = Lpow @ Lpow ------------------------------------------ sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): sq_acc[accum_idx] = cutlass.Float32(0.0) mma_step( sq_acc, (lpow_a0, lpow_a1, lpow_a2, lpow_a3), - (movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3)), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), k_step=0, M=16, N=16, @@ -576,47 +580,58 @@ def _super_mma_warp( lpow_a1 = fp32_to_fp16(sq_acc[2], sq_acc[3], dtype=cfg.io_dtype) lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) - # A_inv <- A_inv + A_inv @ Lpow, keeping A_inv in registers + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) + # ---- T_inv += T_inv @ Lpow --------------------------------------- upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): upd_acc[accum_idx] = cutlass.Float32(0.0) + tinv_p0 = fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype) + tinv_p1 = fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype) + tinv_p2 = fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype) + tinv_p3 = fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype) mma_step( upd_acc, - ( - fp32_to_fp16(ainv_acc[0], ainv_acc[1], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[2], ainv_acc[3], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[4], ainv_acc[5], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[6], ainv_acc[7], dtype=cfg.io_dtype), - ), - (movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3)), + (tinv_p0, tinv_p1, tinv_p2, tinv_p3), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - for accum_idx in cutlass.range_constexpr(8): - ainv_acc[accum_idx] = ainv_acc[accum_idx].to(cfg.io_dtype).to(cutlass.Float32) + upd_acc[accum_idx] - + tinv_lo0, tinv_hi0 = f16x2_to_f32(tinv_p0, dtype=cfg.io_dtype) + tinv_lo1, tinv_hi1 = f16x2_to_f32(tinv_p1, dtype=cfg.io_dtype) + tinv_lo2, tinv_hi2 = f16x2_to_f32(tinv_p2, dtype=cfg.io_dtype) + tinv_lo3, tinv_hi3 = f16x2_to_f32(tinv_p3, dtype=cfg.io_dtype) + tinv_acc[0] = tinv_lo0 + upd_acc[0] + tinv_acc[1] = tinv_hi0 + upd_acc[1] + tinv_acc[2] = tinv_lo1 + upd_acc[2] + tinv_acc[3] = tinv_hi1 + upd_acc[3] + tinv_acc[4] = tinv_lo2 + upd_acc[4] + tinv_acc[5] = tinv_hi2 + upd_acc[5] + tinv_acc[6] = tinv_lo3 + upd_acc[6] + tinv_acc[7] = tinv_hi3 + upd_acc[7] + + bars.mb_intermediate_done[intermediate_stage].wait(((global_chunk // cfg.smem_intermediate_stages) + 1) % 2) nvvm.stmatrix( - sPairwise_ptr + (cfg.b_t * cfg.b_t) + stsm_idx, + sIntermediate_ptr + (cfg.b_t * cfg.b_t) + stsm_idx, [ - fp32_to_fp16(ainv_acc[0], ainv_acc[1], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[2], ainv_acc[3], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[4], ainv_acc[5], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[6], ainv_acc[7], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype), ], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8, ) nvvm.fence_proxy("async.shared", space="cta") - bars.mb_a_ready[pairwise_stage].arrive() - bars.mb_kk_qk_super_mma_done[decay_stage].arrive() - gbase += sk_nt - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + bars.mb_t_inv_ready[intermediate_stage].arrive() + bars.mb_decay_super_done[decay_stage].arrive() + global_chunk_base += num_tile_chunks + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @cute.jit -def _tcgen05_mma_warp( +def tcgen05_mma_warp( cfg, total_tiles, bidx, @@ -624,24 +639,35 @@ def _tcgen05_mma_warp( cu_seqlens, mWorkItems, sSched, - tmem_hold, - sPairwise, + sTmem_base, + sIntermediate, sK_decay, sK_restore, sQ_decay, sState_scale_diag, bars, ) -> None: - """tcgen05-MMA warp role (warp 13): persistent tile-scheduler loop, issues - all six state GEMMs in dependency order and owns the TMEM lifecycle.""" + """tcgen05-MMA warp role (warp 13): persistent scheduler loop issuing + every tcgen05 GEMM.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) - tmem_base = tmem_hold.load() + nvvm.tcgen05_alloc(sTmem_base, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = sTmem_base.load() + state_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) + state_dsts = tuple(nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset + k * 16, cutlass.Float32) for k in range(cfg.d_k // 16)) + state_k_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_k_acc_offset, cutlass.Float32) + u_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_acc_offset, cutlass.Float32) + y_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_inp_offset, cutlass.Int8) + u_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_inp_offset, cutlass.Int8) + state_dst_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset, cutlass.Float32) state_inp_index = PipelineState.start(phase=0) - rhs_index = PipelineState.start(phase=0) - update_index = PipelineState.start(phase=0) - final_state_index = PipelineState.start(phase=0) + state_read_index = PipelineState.start(phase=0) + y_inp_index = PipelineState.start(phase=0) + u_inp_index = PipelineState.start(phase=0) qk_scale_index = PipelineState.start(phase=0) - # ---- chunk-invariant GEMM descriptors ------------------------------ + + # ---- chunk-invariant GEMM descriptors ---------------------------------------- bpe = cfg.io_dtype.width // 8 idesc_acc = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -691,7 +717,7 @@ def _tcgen05_mma_warp( idesc=idesc_diag, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_pairwise_desc = MmaDesc( + bmm_qk_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -715,128 +741,127 @@ def _tcgen05_mma_warp( idesc=idesc_final_state, kind=nvvm.Tcgen05MMAKind.F16, ) - gbase = cutlass.Int32(0) + STATE_A_SEG = bmm_state_desc.sps_B * bmm_state_desc.tmem_advance_A + STATE_B_SEG = bmm_state_desc.smem_subtile_B >> 4 + global_chunk_base = cutlass.Int32(0) 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, cu_seqlens, mWorkItems - ) - sk_nt = wend - cstart - for li in cutlass.range(sk_nt, unroll=1): - gc = gbase + li - q_state_acc_stage = gc % cfg.tmem_q_state_acc_stages - decay_stage = gc % cfg.smem_decay_stages + 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_tile_chunks = wend - cstart + for local_chunk in cutlass.range(num_tile_chunks, unroll=1): + global_chunk = global_chunk_base + local_chunk + have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_chunk > 0 + q_state_acc_stage = global_chunk % cfg.tmem_q_state_acc_stages + decay_stage = global_chunk % cfg.smem_decay_stages state_scale_diag_stage = qk_scale_index.idx - pairwise_stage = gc % cfg.smem_pairwise_stages + intermediate_stage = global_chunk % cfg.smem_intermediate_stages sK_decay_stage = sK_decay[decay_stage] sQ_decay_stage = sQ_decay[decay_stage] sK_restore_stage = sK_restore[decay_stage] sState_scale_diag_stage = sState_scale_diag[state_scale_diag_stage] - sPairwise_stage = sPairwise[pairwise_stage] - - # ---- state*k -> state_k_acc ---------------------------------------- - bars.mb_k_decay_cg0_ready[decay_stage].wait((gc // cfg.smem_decay_stages) % 2) - bars.mb_state_inp_ready.wait(state_inp_index.phase) - state_inp_index = advance(state_inp_index, 1) - desc_k_decay = sK_decay_stage.desc() - - state_a_tmem_base = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) - mma_ts( - bmm_state_desc, - state_a_tmem_base, - desc_k_decay, - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_k_acc_offset, cutlass.Float32), - accumulate=False, - ) + sIntermediate_stage = sIntermediate[intermediate_stage] + + # ---- state_k = state(T) @ K_decay^T ---------------------------------- + bars.mb_k_decay_inv_cg0_ready[decay_stage].wait((global_chunk // cfg.smem_decay_stages) % 2) + if have_state: + bars.mb_state_inp_ready.wait(state_inp_index.phase) + state_inp_index = advance(state_inp_index, 1) + desc_k_decay = sK_decay_stage.desc() + + for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + mma_ts_step( + bmm_state_desc, + state_inp_ptr.subview(s * STATE_A_SEG), + desc_k_decay + s * STATE_B_SEG, + state_k_acc_ptr, + k, + cutlass.Boolean(s + k > 0), + ) - if nvvm.elect_sync(): - bars.mb_state_k_acc_ready.arrive(cta_group=1) + if elect_one: + bars.mb_state_k_acc_ready.arrive(cta_group=1) - # ---- state*q -> q_state_acc (stays live until qk@update fuses into o) + # ---- q_state = state(T) @ Q_decay^T ---------------------------------- bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) - bars.mb_o_acc_done[q_state_acc_stage].wait(((gc // cfg.tmem_q_state_acc_stages + cutlass.Int32(1)) % cutlass.Int32(2))) - desc_q_decay = sQ_decay_stage.desc() - - state_a_tmem_base = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) - mma_ts( - bmm_state_desc, - state_a_tmem_base, - desc_q_decay, - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_q_state_acc_offset + q_state_acc_stage * cfg.b_t, cutlass.Float32), - accumulate=False, - ) - - if nvvm.elect_sync(): - bars.mb_kk_qk_mma_done[decay_stage].arrive(cta_group=1) + bars.mb_o_acc_done[q_state_acc_stage].wait(((global_chunk // cfg.tmem_q_state_acc_stages + cutlass.Int32(1)) % cutlass.Int32(2))) + q_state_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_q_state_acc_offset + q_state_acc_stage * cfg.b_t, cutlass.Float32) + if have_state: + desc_q_decay = sQ_decay_stage.desc() + for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + mma_ts_step( + bmm_state_desc, + state_inp_ptr.subview(s * STATE_A_SEG), + desc_q_decay + s * STATE_B_SEG, + q_state_acc_ptr, + k, + cutlass.Boolean(s + k > 0), + ) - # ---- state decay (per-k-atom diag blocks) ---------------------------- - desc_diag = sState_scale_diag_stage.desc() + if elect_one: + bars.mb_decay_tcgen05_done[decay_stage].arrive(cta_group=1) - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - mma_ts( - bmm_diag_desc, - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset + k_block * 8, cutlass.Int8), - desc_diag.advance_start_address(k_block * 256 * 2), - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_offset + k_block * 16, cutlass.Float32), - accumulate=False, - ) + if cutlass.const_expr(cfg.enable_checkpoints): + if have_state: + bars.mb_state_acc_read_done.wait(state_read_index.phase) + state_read_index = advance(state_read_index, 1) + + # ---- state decay = state(T) @ diag(exp2(g_last)) (per-k-atom blocks) ---- + if have_state: + desc_diag = sState_scale_diag_stage.desc() + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + mma_ts_step( + bmm_diag_desc, + state_inp_ptr.subview(k_block * bmm_diag_desc.tmem_advance_A), + desc_diag.advance_start_address(k_block * 256 * 2), + state_dsts[k_block], + 0, + cutlass.Boolean(False), + ) - if nvvm.elect_sync(): + if elect_one: bars.mb_state_scale_diag_done[state_scale_diag_stage].arrive(cta_group=1) - # ---- update = A_inv @ rhs -> update_acc ------------------------------ - bars.mb_a_ready[pairwise_stage].wait((gc // cfg.smem_pairwise_stages) % 2) - bars.mb_rhs_ready.wait(rhs_index.phase) - rhs_index = advance(rhs_index, 1) - lhs_tmem = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_rhs_inp_offset, cutlass.Int8) - desc_pairwise = sPairwise_stage.shifted((cfg.b_t * cfg.b_t)).desc() - mma_ts( - bmm_pairwise_desc, - lhs_tmem, - desc_pairwise, - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_update_acc_offset, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_update_acc_ready.arrive(cta_group=1) - - # ---- final_state += update @ k_restore ------------------------------- - bars.mb_update_ready.wait(update_index.phase) - update_index = advance(update_index, 1) - update_tmem = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_update_inp_offset, cutlass.Int8) + # ---- u_acc = Y(T) @ T_inv -------------------------------------------- + bars.mb_t_inv_ready[intermediate_stage].wait((global_chunk // cfg.smem_intermediate_stages) % 2) + bars.mb_y_inp_ready.wait(y_inp_index.phase) + y_inp_index = advance(y_inp_index, 1) + desc_qk = sIntermediate_stage.shifted((cfg.b_t * cfg.b_t)).desc() + mma_ts_step(bmm_qk_desc, y_inp_ptr, desc_qk, u_acc_ptr, 0, cutlass.Boolean(False)) + if elect_one: + bars.mb_u_acc_ready.arrive(cta_group=1) + + # ---- final_state += U(T) @ K_restore --------------------------------- + bars.mb_u_inp_ready.wait(u_inp_index.phase) + u_inp_index = advance(u_inp_index, 1) desc_k_restore = sK_restore_stage.desc() - mma_ts( - bmm_final_state_desc, - update_tmem, - desc_k_restore, - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_offset, cutlass.Float32), - accumulate=True, - ) - if nvvm.elect_sync(): - bars.mb_k_restore_done[decay_stage].arrive(cta_group=1) - - # ---- o += qk @ update -> q_state_acc ---------------------------------- - bars.mb_qk_acc_ready[pairwise_stage].wait((gc // cfg.smem_pairwise_stages) % 2) - lhs_tmem = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_update_inp_offset, cutlass.Int8) - desc_pairwise = sPairwise_stage.desc() - mma_ts( - bmm_pairwise_desc, - lhs_tmem, - desc_pairwise, + mma_ts_step(bmm_final_state_desc, u_inp_ptr, desc_k_restore, state_dst_ptr, 0, have_state) + if elect_one: + bars.mb_k_restore_acc_done[decay_stage].arrive(cta_group=1) + + # ---- O += U(T) @ A --------------------------------------------------- + bars.mb_a_ready[intermediate_stage].wait((global_chunk // cfg.smem_intermediate_stages) % 2) + desc_qk = sIntermediate_stage.desc() + mma_ts_step( + bmm_qk_desc, + u_inp_ptr, + desc_qk, nvvm.make_tmem_ptr(tmem_base + cfg.tmem_q_state_acc_offset + q_state_acc_stage * cfg.b_t, cutlass.Float32), - accumulate=True, + 0, + have_state, ) - if nvvm.elect_sync(): + if elect_one: bars.mb_o_acc_ready.arrive(cta_group=1) - bars.mb_a_done[pairwise_stage].arrive(cta_group=1) + bars.mb_intermediate_done[intermediate_stage].arrive(cta_group=1) qk_scale_index = advance(qk_scale_index, cfg.smem_state_scale_diag_stages) - bars.mb_final_state_stored.wait(final_state_index.phase) - final_state_index = advance(final_state_index, 1) - gbase += sk_nt - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + global_chunk_base += num_tile_chunks + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + bars.mb_tmem_done[0].wait(0) + nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), cutlass.Int32(512), @@ -845,7 +870,7 @@ def _tcgen05_mma_warp( @cute.jit -def _epilogue_warp( +def epilogue_warp( cfg, total_tiles, bidx, @@ -857,22 +882,23 @@ def _epilogue_warp( mO, sK_inv_raw, sO_raw, - sPairwise_raw, + sIntermediate_raw, sQ_decay_raw, - sH_raw, + sCheckpoint_raw, desc_o_base, - desc_h_base, + desc_checkpoint_base, checkpoint_every_n_tokens, bars, ) -> None: - """Epilogue warp role (warp 15): register-MMA qk (causal), A_inv qk - staging, and the O TMA store drain.""" + """Epilogue warp role (warp 15): register-MMA A (causal) and the O TMA + store drain.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(cfg.enable_checkpoints): - sH_tma = SmemTile( - base=sH_raw, + sCheckpoint_tma = SmemTile( + base=sCheckpoint_raw, elems_per_stage=(cfg.d_k * cfg.d_v), - stages=1, + stages=cfg.smem_checkpoint_stages, leading_byte_offset=0, stride_byte_offset=0, layout=0, @@ -880,7 +906,7 @@ def _epilogue_warp( tma_granu_elems=64, tma_subtile_stride_elems=cfg.d_k * 64, ) - h_ready_index = PipelineState.start(phase=0) + checkpoint_ready_index = PipelineState.start(phase=0) sO_tma = SmemTile( base=sO_raw, elems_per_stage=(cfg.b_t * cfg.d_v), @@ -893,13 +919,13 @@ def _epilogue_warp( tma_subtile_stride_elems=cfg.b_t * 64, ) qk_scale_index = PipelineState.start(phase=0) - # ---- ldmatrix/stmatrix lane decode --------------------------------- + + # ---- ldmatrix/stmatrix lane decode ------------------------------------------- rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - decay_key_mask = cutlass.Int32(8) ^ ((lhs_row_coord & cutlass.Int32(2)) * cutlass.Int32(16)) - elems_per_128b = cutlass.Int32(64) + decay_key_mask = cutlass.Int32(8) stsm_row_coord = lane & 7 stsm_col_coord = cutlass.Int32(0) if (lane // 8) & 1: @@ -907,44 +933,45 @@ def _epilogue_warp( if lane // 8 >= 2: stsm_col_coord = cutlass.Int32(8) stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) - gbase = cutlass.Int32(0) + global_chunk_base = cutlass.Int32(0) 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, cu_seqlens, mWorkItems - ) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) head_o = head_idx - o_slot = (batch_idx * cutlass.Int32(cfg.n_heads_out) + head_idx) * cutlass.Int32(TENSOR_MAP_QWORDS) + o_slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) desc_o_slot = (desc_o_base + o_slot).tospace(cutlass.AddressSpace.generic) if cutlass.const_expr(cfg.enable_checkpoints): - desc_h_slot = (desc_h_base + o_slot).tospace(cutlass.AddressSpace.generic) - if nvvm.elect_sync(): - tma_tensormap_acquire(desc_h_slot) - if nvvm.elect_sync(): + desc_checkpoint_slot = (desc_checkpoint_base + o_slot).tospace(cutlass.AddressSpace.generic) + checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + checkpoint_quot = (cstart + cutlass.Int32(1)) // checkpoint_chunks + checkpoint_mod = (cstart + cutlass.Int32(1)) % checkpoint_chunks + if elect_one: + tma_tensormap_acquire(desc_checkpoint_slot) + if elect_one: tma_tensormap_acquire(desc_o_slot) - sk_nt = wend - cstart - for li in cutlass.range(sk_nt, unroll=1): - chunk_idx = cstart + li - gc = gbase + li - decay_stage = gc % cfg.smem_decay_stages - pairwise_stage = gc % cfg.smem_pairwise_stages + num_tile_chunks = wend - cstart + for local_chunk in cutlass.range(num_tile_chunks, unroll=1): + chunk_idx = cstart + local_chunk + global_chunk = global_chunk_base + local_chunk + decay_stage = global_chunk % cfg.smem_decay_stages + intermediate_stage = global_chunk % cfg.smem_intermediate_stages sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) sQ_decay_ptr = sQ_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) - sPairwise_ptr = sPairwise_raw.data_ptr() + pairwise_stage * (2 * cfg.b_t * cfg.b_t) + sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (2 * cfg.b_t * cfg.b_t) - bars.mb_a_done[pairwise_stage].wait(((gc // cfg.smem_pairwise_stages) + 1) % 2) bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) - # ---- qk register MMA (inclusive-causal) ---------------------------- - qk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + + # ---- A = Q_decay @ K_inv^T ------------------------------------------ + a_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): - qk_acc[accum_idx] = cutlass.Float32(0.0) + a_acc[accum_idx] = cutlass.Float32(0.0) for k_block in cutlass.range_constexpr((cfg.d_k // 16)): # Load B operand k_inv_col = k_block * 16 + rhs_col_offset k_inv_segment = k_inv_col // 64 - rhs_vec = nvvm.ldmatrix( + rhs_frag = nvvm.ldmatrix( sK_inv_ptr + k_inv_segment * (cfg.b_t * 64) + rhs_row_coord * 64 @@ -954,26 +981,19 @@ def _epilogue_warp( ) # Load A operand storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask - storage_slice = storage_key // elems_per_128b - key_in_slice = storage_key - storage_slice * elems_per_128b - storage_phase = key_in_slice // cutlass.Int32(16) - byte_in_slice = ( - lhs_row_coord * cutlass.Int32(128) - + storage_phase * cutlass.Int32(32) - + (key_in_slice - storage_phase * cutlass.Int32(16)) * cutlass.Int32(2) - ) - qk_lhs_vec = nvvm.ldmatrix( + storage_slice = storage_key // 64 + a_lhs_frag = nvvm.ldmatrix( sQ_decay_ptr - + storage_slice * cutlass.Int32(cfg.b_t) * elems_per_128b - + ((byte_in_slice ^ ((lhs_row_coord & cutlass.Int32(7)) << 4)) // cutlass.Int32(2)), + + storage_slice * (cfg.b_t * 64) + + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( - qk_acc, - (qk_lhs_vec[0], qk_lhs_vec[1], qk_lhs_vec[2], qk_lhs_vec[3]), - (rhs_vec[0], rhs_vec[1], rhs_vec[2], rhs_vec[3]), + a_acc, + (a_lhs_frag[0], a_lhs_frag[1], a_lhs_frag[2], a_lhs_frag[3]), + (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), k_step=0, M=16, N=16, @@ -987,78 +1007,104 @@ def _epilogue_warp( col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) - qk_acc[accum_idx] = qk_acc[accum_idx] if row_coord >= col_coord else cutlass.Float32(0.0) + a_acc[accum_idx] = a_acc[accum_idx] if row_coord >= col_coord else cutlass.Float32(0.0) + bars.mb_intermediate_done[intermediate_stage].wait(((global_chunk // cfg.smem_intermediate_stages) + 1) % 2) nvvm.stmatrix( - sPairwise_ptr + stsm_idx, + sIntermediate_ptr + stsm_idx, [ - fp32_to_fp16(qk_acc[0], qk_acc[1], dtype=cfg.io_dtype), - fp32_to_fp16(qk_acc[2], qk_acc[3], dtype=cfg.io_dtype), - fp32_to_fp16(qk_acc[4], qk_acc[5], dtype=cfg.io_dtype), - fp32_to_fp16(qk_acc[6], qk_acc[7], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[0], a_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[2], a_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[4], a_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[6], a_acc[7], dtype=cfg.io_dtype), ], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8, ) nvvm.fence_proxy("async.shared", space="cta") - bars.mb_qk_acc_ready[pairwise_stage].arrive() - bars.mb_kk_qk_super_mma_done[decay_stage].arrive() + bars.mb_a_ready[intermediate_stage].arrive() + bars.mb_decay_super_done[decay_stage].arrive() qk_scale_index = advance(qk_scale_index, cfg.qk_scale_ready_stages) - # ---- O drain: staged output tile -> GMEM TMA store ----------------- - if li > 0: + # ---- checkpoint + O drain: checkpoint stores first (CG1 stages checkpoint before O) -- + if local_chunk > 0: output_chunk = chunk_idx - cutlass.Int32(1) output_chunk_start = output_chunk * cfg.b_t - o_stage = (gc - cutlass.Int32(1)) % cfg.smem_o_stages - bars.mb_o_tmastg_ready[o_stage].wait(((gc - cutlass.Int32(1)) // cfg.smem_o_stages) % 2) - o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), output_chunk_start) - if cutlass.const_expr(cfg.split_k): - # warmup chunks stage O to SMEM but never store it - if output_chunk >= wstart: - tma_store_tile(sO_tma[o_stage], o_slice, acquire=False) + o_stage = (global_chunk - cutlass.Int32(1)) % cfg.smem_o_stages + did_checkpoint = cutlass.Int32(0) + checkpoint_stage = cutlass.Int32(0) + if cutlass.const_expr(cfg.enable_checkpoints): + # ---- checkpoint store ---------------------------------------- + do_checkpoint = checkpoint_mod == 0 + do_checkpoint = do_checkpoint and chunk_idx >= wstart + checkpoint_stage = checkpoint_ready_index.idx + 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_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() - else: + did_checkpoint = cutlass.Int32(1) + checkpoint_mod = checkpoint_mod + cutlass.Int32(1) + if checkpoint_mod == checkpoint_chunks: + checkpoint_mod = cutlass.Int32(0) + checkpoint_quot = checkpoint_quot + cutlass.Int32(1) + bars.mb_o_tmastg_ready[o_stage].wait(((global_chunk - cutlass.Int32(1)) // cfg.smem_o_stages) % 2) + o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), head_o, output_chunk_start) + did_o = cutlass.Int32(0) + if output_chunk >= wstart: tma_store_tile(sO_tma[o_stage], o_slice, acquire=False) tma_store_commit() - tma_store_wait(0) - bars.mb_o_tmastg_done[o_stage].arrive() - # ---- H store: CG1 staged the state entering chunk_idx ---------- - if cutlass.const_expr(cfg.enable_checkpoints): - if li > 0: - tokens_done = chunk_idx * cutlass.Int32(cfg.b_t) - do_h = tokens_done % checkpoint_every_n_tokens == 0 - if cutlass.const_expr(cfg.split_k): - do_h = do_h and chunk_idx >= wstart - if do_h: - bars.mb_h_tmastg_ready.wait(h_ready_index.phase) - h_ready_index = advance(h_ready_index, 1) - # sequence-local entry: the per-(b,h) descriptor folds the - # sequence base into GLOBAL_ADDRESS and caps the extent - h_entry = tokens_done // checkpoint_every_n_tokens - cutlass.Int32(1) - h_slice = tma_slice_runtime_desc(desc_h_slot, cutlass.Int32(0), cutlass.Int32(0), h_entry) - tma_store_tile(sH_tma[0], h_slice, acquire=False) - tma_store_commit() + did_o = cutlass.Int32(1) + if cutlass.const_expr(cfg.enable_checkpoints): + if did_checkpoint == 1 and did_o == 1: + tma_store_wait(1) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() + tma_store_wait(0) + bars.mb_o_tmastg_done[o_stage].arrive() + if did_checkpoint == 1 and did_o == 0: tma_store_wait(0) - bars.mb_h_tmastg_done.arrive() - # ---- last computed chunk drain (always owned: it is wend - 1) ------ - if sk_nt > 0: + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() + bars.mb_o_tmastg_done[o_stage].arrive() + if did_checkpoint == 0: + if did_o == 1: + tma_store_wait(0) + bars.mb_o_tmastg_done[o_stage].arrive() + else: + tma_store_wait(0) + bars.mb_o_tmastg_done[o_stage].arrive() + + # ---- last computed chunk drain (always owned: it is wend - 1) ------------ + if num_tile_chunks > 0: output_chunk = wend - cutlass.Int32(1) - og = gbase + sk_nt - cutlass.Int32(1) + last_global_chunk = global_chunk_base + num_tile_chunks - cutlass.Int32(1) output_chunk_start = output_chunk * cfg.b_t - o_stage = og % cfg.smem_o_stages - bars.mb_o_tmastg_ready[o_stage].wait((og // cfg.smem_o_stages) % 2) - # a partial last chunk is clipped by the descriptor's token extent - o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), output_chunk_start) + o_stage = last_global_chunk % cfg.smem_o_stages + bars.mb_o_tmastg_ready[o_stage].wait((last_global_chunk // cfg.smem_o_stages) % 2) + o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), head_o, output_chunk_start) tma_store_tile(sO_tma[o_stage], o_slice, acquire=False) tma_store_commit() tma_store_wait(0) bars.mb_o_tmastg_done[o_stage].arrive() - gbase += sk_nt - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + global_chunk_base += num_tile_chunks + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @cute.jit -def _compute0_warp_group( +def gate_scale(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: + """Map raw gate to the log2-domain decay increment.""" + + if cutlass.const_expr(cfg.safe_gate): + half = cutlass.Float32(0.5) + sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half + return cfg.gate_scale_log2 * sigmoid + # Default ABI: gate arrives in natural-log space + return raw_gate * cutlass.Float32(LOG2_E) + + +@cute.jit +def compute0_warp_group( cfg, total_tiles, bidx, @@ -1084,9 +1130,8 @@ def _compute0_warp_group( sState_scale_diag_raw, bars, ) -> None: - """CG0 warp role (warps 0-7, two ping-pong groups): persistent - tile-scheduler loop + gate prefix scan and the decay/restore operand - materialization into tcgen05 SMEM.""" + """CG0 warp role (warps 0-7, two ping-pong groups): Gate prefix scan and + the decay/restore operand materialization.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] cg0_group_id = cg0_warp // cfg.cg0_warps_per_group @@ -1094,28 +1139,33 @@ def _compute0_warp_group( prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane cg0_a_log_exp = cutlass.Float32(1.0) cg0_dt_bias_value = cutlass.Float32(0.0) - gbase = cutlass.Int32(0) + global_chunk_base = cutlass.Int32(0) sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) + opaque_one = opaque_f32_zero() + cutlass.Float32(1.0) 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, cu_seqlens, mWorkItems - ) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) head_o = head_idx - sk_nt = wend - cstart + num_tile_chunks = wend - cstart if cutlass.const_expr(cfg.safe_gate): - # per-head safe-gate constants straight from GMEM (the head - # changes per tile, so there is no SMEM staging to handshake) - if sk_nt > 0: + if num_tile_chunks > 0: cg0_a_log_exp = cute.math.exp2(mA_log[head_o].to(cutlass.Float32) * LOG2_E, fastmath=True) cg0_dt_bias_value = mDt_bias[head_o, prefix_dim].to(cutlass.Float32) - for li in cutlass.range(cg0_group_id, sk_nt, cfg.cg0_group_count, unroll=1): - chunk_idx = cstart + li - gc = gbase + li + # tile entry: both ping-pong groups inherit each other's delivery proofs (parity-swap guard) + nvvm.barrier_cta_sync(cfg.cg0_tile_entry_barrier_id, thread_count=cfg.cg0_group_count * cfg.cg0_threads_per_group) + cg0_first_global_chunk = global_chunk_base + cutlass.Int32(cg0_group_id) + diag_ring_stage = cg0_first_global_chunk % cutlass.Int32(cfg.smem_state_scale_diag_stages) + diag_ring_phase = (cg0_first_global_chunk // cutlass.Int32(cfg.smem_state_scale_diag_stages)) % cutlass.Int32(2) + raw_ring_stage = cg0_first_global_chunk % cutlass.Int32(cfg.smem_raw_stages) + raw_bar_stage = cg0_first_global_chunk % cutlass.Int32(cfg.smem_raw_bar_stages) + raw_bar_phase = (cg0_first_global_chunk // cutlass.Int32(cfg.smem_raw_bar_stages)) % cutlass.Int32(2) + for local_chunk in cutlass.range(cg0_group_id, num_tile_chunks, cfg.cg0_group_count, unroll=1): + chunk_idx = cstart + local_chunk + global_chunk = global_chunk_base + local_chunk chunk_start = chunk_idx * cfg.b_t - decay_stage = gc % cfg.smem_decay_stages - raw_stage = gc % cfg.smem_raw_stages - state_scale_diag_stage = gc % cfg.smem_state_scale_diag_stages + decay_stage = global_chunk % cfg.smem_decay_stages + raw_stage = raw_ring_stage + state_scale_diag_stage = diag_ring_stage qk_scale_ready_stage = state_scale_diag_stage sQ_ptr = sQ_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) sK_ptr = sK_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) @@ -1129,60 +1179,20 @@ def _compute0_warp_group( sK_restore_ptr = sK_restore_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) sState_scale_diag_ptr = sState_scale_diag_raw.data_ptr() + state_scale_diag_stage * ((cfg.d_k // 16) * 256) - bars.mb_inputs_ready[raw_stage].wait((gc // cfg.smem_raw_stages) % 2) - - # ---- tail chunk: zero-fill raw staging past seqlen ----------------- - if chunk_start + cutlass.Int32(cfg.b_t) > seqlen_b: - if cg0_local_warp == 0: - f16_zero = mQ.element_type(0.0) - f16_zero_vec = cutlass.Vector.from_elements( - ( - f16_zero, - f16_zero, - f16_zero, - f16_zero, - f16_zero, - f16_zero, - f16_zero, - f16_zero, - ), - mQ.element_type, - ) - f32_zero = cutlass.Float32(0.0) - f32_zero_vec = cutlass.Vector.from_elements( - (f32_zero, f32_zero, f32_zero, f32_zero), - cutlass.Float32, - ) - for row in cutlass.range_constexpr(cfg.b_t): - token_idx = chunk_start + cutlass.Int32(row) - if token_idx >= seqlen_b: - if lane < (cfg.d_k // 8): - f16_dim_base = lane * 8 - f16_segment = f16_dim_base // 64 - f16_segment_dim = f16_dim_base - f16_segment * 64 - f16_idx = f16_segment * (cfg.b_t * 64) + row * 64 + swizzle_xor_128b(row, f16_segment_dim, elem_bytes=2) - (sQ_ptr + f16_idx).store(f16_zero_vec, alignment=16) - (sK_ptr + f16_idx).store(f16_zero_vec, alignment=16) - (sV_ptr + f16_idx).store(f16_zero_vec, alignment=16) - (sBeta_ptr + f16_idx).store(f16_zero_vec, alignment=16) - (sW_ptr + f16_idx).store(f16_zero_vec, alignment=16) - if lane < (cfg.d_k // 4): - f32_dim_base = lane * 4 - f32_segment = f32_dim_base // 32 - f32_segment_dim = f32_dim_base - f32_segment * 32 - f32_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) - (sGate_ptr + f32_idx).store(f32_zero_vec, alignment=16) - nvvm.barrier_cta_sync(cfg.nbar_cg0_group0_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) + bars.mb_gate_ready[raw_bar_stage].wait(raw_bar_phase) row_group_start = cg0_local_warp * (cfg.b_t // cfg.cg0_warps_per_group) lane_row_group = lane // 8 lane_in_row_group = lane - lane_row_group * 8 decay_row = row_group_start + lane_row_group - - g_prefix_ptr = sGate_ptr + decay_key_mask = cutlass.Int32(8) prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane - # ---- gate prefix scan: cumulative log-gate per key channel -------- + + # ---- Gate prefix scan ----------------------------------------------- + f32_segment = prefix_dim // 32 + prefix_seg_base = f32_segment * (cfg.b_t * 32) + prefix_col = prefix_dim - f32_segment * 32 g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) if cutlass.const_expr(cfg.safe_gate): valid_rows = seqlen_b - chunk_idx * cutlass.Int32(cfg.b_t) @@ -1190,21 +1200,17 @@ def _compute0_warp_group( for row_pair in cutlass.range_constexpr(cfg.b_t // 2): row0 = row_pair * 2 row1 = row0 + 1 - f32_segment = prefix_dim // 32 - f32_segment_dim = prefix_dim - f32_segment * 32 - prefix_idx0 = f32_segment * (cfg.b_t * 32) + row0 * 32 + swizzle_xor_128b(row0, f32_segment_dim, elem_bytes=4) - f32_segment = prefix_dim // 32 - f32_segment_dim = prefix_dim - f32_segment * 32 - prefix_idx1 = f32_segment * (cfg.b_t * 32) + row1 * 32 + swizzle_xor_128b(row1, f32_segment_dim, elem_bytes=4) + prefix_idx0 = prefix_seg_base + swizzle_xor_128b(row0, row0 * 32 + prefix_col, elem_bytes=4) + prefix_idx1 = prefix_seg_base + swizzle_xor_128b(row1, row1 * 32 + prefix_col, elem_bytes=4) gate0 = (sGate_ptr + prefix_idx0).load() gate1 = (sGate_ptr + prefix_idx1).load() gate0 = cg0_a_log_exp * (gate0 + cg0_dt_bias_value) gate1 = cg0_a_log_exp * (gate1 + cg0_dt_bias_value) - gate0 = _gate_log2( + gate0 = gate_scale( cfg, gate0, ) - gate1 = _gate_log2( + gate1 = gate_scale( cfg, gate1, ) @@ -1214,13 +1220,11 @@ def _compute0_warp_group( g_prefix_regs[row1] = gate_pair[1] else: for row in cutlass.range_constexpr(cfg.b_t): - f32_segment = prefix_dim // 32 - f32_segment_dim = prefix_dim - f32_segment * 32 - prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) + prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) gate = (sGate_ptr + prefix_idx).load() token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) if token_idx < seqlen_b: - gate = _gate_log2( + gate = gate_scale( cfg, gate, ) @@ -1240,62 +1244,70 @@ def _compute0_warp_group( g_prefix_regs[row1] = prefix1 prefix_acc = prefix1 + # ---- exp2(g): stage prefixes + final-token decay --------------------- for row in cutlass.range_constexpr(cfg.b_t): g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) - # ---- exp2(g): stage prefixes + final-token decay ------------------ exp_g_last = g_prefix_regs[cfg.b_t - 1] for row in cutlass.range_constexpr(cfg.b_t): - f32_segment = prefix_dim // 32 - f32_segment_dim = prefix_dim - f32_segment * 32 - prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) + prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) - # ---- state-scale diag: stage exp2(g_last) decay blocks ------------- - bars.mb_state_scale_diag_done[state_scale_diag_stage].wait((gc // cfg.smem_state_scale_diag_stages + 1) % 2) - diag_idx = _diag_idx(cfg, prefix_dim) + # ---- state-scale diag: stage exp2(g_last) decay blocks --------------- + bars.mb_state_scale_diag_done[state_scale_diag_stage].wait(diag_ring_phase ^ cutlass.Int32(1)) + block = prefix_dim // cutlass.Int32(16) + coord = prefix_dim - block * cutlass.Int32(16) + storage_col = coord ^ cutlass.Int32((cfg.b_t // 2)) + linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + storage_col + diag_idx = swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) sState_scale_diag_ptr[diag_idx] = exp_g_last.to(cfg.io_dtype) - nvvm.barrier_cta_sync(cfg.nbar_cg0_group0_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) + nvvm.barrier_cta_sync(cfg.cg0_group_sync_barrier_base_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) - k_inv_words = cutlass.Array(cutlass.Int32, 2 * 4, alignment=16) + bars.mb_q_ready[raw_bar_stage].wait(raw_bar_phase) + bars.mb_k_ready[raw_bar_stage].wait(raw_bar_phase) + bars.mb_beta_ready[raw_bar_stage].wait(raw_bar_phase) + k_inv_pack = cutlass.Array(cutlass.Int32, 2 * 4, alignment=16) raw_q_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) raw_beta_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) - # ---- optional q/k L2-norm + K_inv staging -------------------------- - q_sum_sq = cutlass.Float32(0.0) - k_sum_sq = cutlass.Float32(0.0) + + # ---- optional Q/K L2-norm ------------------------------------------- + if cutlass.const_expr(cfg.l2norm): + qk0_lo = opaque_f32_zero() + qk0_hi = opaque_f32_zero() + qk1_lo = opaque_f32_zero() + qk1_hi = opaque_f32_zero() for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 f16_segment = dim_base // 64 f16_segment_dim = dim_base - f16_segment * 64 raw_f16_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) - raw_q_vec = (sQ_ptr + raw_f16_idx).load(count=8, alignment=16) - raw_k_vec = (sK_ptr + raw_f16_idx).load(count=8, alignment=16) - raw_beta_vec = (sBeta_ptr + raw_f16_idx).load(count=8, alignment=16) - raw_q_vec_f32 = raw_q_vec.to(cutlass.Float32) - raw_k_vec_f32 = raw_k_vec.to(cutlass.Float32) - raw_beta_vec_f32 = raw_beta_vec.to(cutlass.Float32) + raw_q_frag = (sQ_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_k_frag = (sK_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_beta_frag = (sBeta_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_q_vec_f32 = raw_q_frag.to(cutlass.Float32) + raw_k_vec_f32 = raw_k_frag.to(cutlass.Float32) + raw_beta_vec_f32 = raw_beta_frag.to(cutlass.Float32) for dim_offset in cutlass.range_constexpr(8): q_val = raw_q_vec_f32[dim_offset] k_val = raw_k_vec_f32[dim_offset] raw_q_regs[reg_base + dim_offset] = q_val raw_k_regs[reg_base + dim_offset] = k_val beta_val = raw_beta_vec_f32[dim_offset] - if cutlass.const_expr(cfg.beta_w_sigmoid): - # Roundtrip through the io dtype to match host-side beta.sigmoid() - half = cutlass.Float32(0.5) - beta_val = (cute.math.tanh(beta_val * half, approx=True) * half + half).to(cfg.io_dtype).to(cutlass.Float32) raw_beta_regs[reg_base + dim_offset] = beta_val - q_sum_sq = q_sum_sq + q_val * q_val - k_sum_sq = k_sum_sq + k_val * k_val + if cutlass.const_expr(cfg.l2norm): + if cutlass.const_expr(dim_offset % 2 == 0): + qk0_lo, qk0_hi = ffma2(q_val, k_val, q_val, k_val, qk0_lo, qk0_hi) + else: + qk1_lo, qk1_hi = ffma2(q_val, k_val, q_val, k_val, qk1_lo, qk1_hi) - # opaque 1.0: keeps the no-l2norm packed-mul operands out of libNVVM's - # constant folder (the documented inline_ptx "n"-constraint ICE) - q_inv_norm = opaque_f32_zero() + cutlass.Float32(1.0) - k_inv_norm = opaque_f32_zero() + cutlass.Float32(1.0) + q_inv_norm = opaque_one + k_inv_norm = opaque_one if cutlass.const_expr(cfg.l2norm): + q_sum_sq = qk0_lo + qk1_lo + k_sum_sq = qk0_hi + qk1_hi q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 4, 31, kind=nvvm.Shfl.BFLY)) q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 2, 31, kind=nvvm.Shfl.BFLY)) q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 1, 31, kind=nvvm.Shfl.BFLY)) @@ -1306,7 +1318,7 @@ def _compute0_warp_group( q_inv_norm = cute.math.rsqrt(cute.math.max(q_sum_sq, norm_floor_sq), fastmath=True) k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) - # ---- decay/restore operands: exp2(+-g) applied per key channel ----- + # ---- decay/restore operands: exp2(+-g) applied per key channel ------- exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) for dim_half in cutlass.range_constexpr(2): @@ -1318,28 +1330,28 @@ def _compute0_warp_group( f32_segment = f32_dim_base // 32 f32_segment_dim = f32_dim_base - f32_segment * 32 g_prefix_idx = f32_segment * (cfg.b_t * 32) + decay_row * 32 + swizzle_xor_128b(decay_row, f32_segment_dim, elem_bytes=4) - exp_g_vec = (g_prefix_ptr + g_prefix_idx).load(count=4, alignment=16) + exp_g_frag = (sGate_ptr + g_prefix_idx).load(count=4, alignment=16) f32_segment = f32_dim_base // 32 f32_segment_dim = f32_dim_base - f32_segment * 32 exp_g_last_idx = f32_segment * (cfg.b_t * 32) + (cfg.b_t - 1) * 32 + swizzle_xor_128b((cfg.b_t - 1), f32_segment_dim, elem_bytes=4) - exp_g_last_vec = (g_prefix_ptr + exp_g_last_idx).load(count=4, alignment=16) + exp_g_last_frag = (sGate_ptr + exp_g_last_idx).load(count=4, alignment=16) half_reg_base = f32_group * 4 f32_reg_base = reg_base + half_reg_base - exp_g_regs[f32_reg_base] = exp_g_vec[0] - exp_g_regs[f32_reg_base + 1] = exp_g_vec[1] - exp_g_regs[f32_reg_base + 2] = exp_g_vec[2] - exp_g_regs[f32_reg_base + 3] = exp_g_vec[3] - exp_neg_g_regs[half_reg_base] = cute.math.rcp(exp_g_vec[0], approx=True, ftz=True) - exp_neg_g_regs[half_reg_base + 1] = cute.math.rcp(exp_g_vec[1], approx=True, ftz=True) - exp_neg_g_regs[half_reg_base + 2] = cute.math.rcp(exp_g_vec[2], approx=True, ftz=True) - exp_neg_g_regs[half_reg_base + 3] = cute.math.rcp(exp_g_vec[3], approx=True, ftz=True) - exp_g_last_regs[f32_reg_base] = exp_g_last_vec[0] - exp_g_last_regs[f32_reg_base + 1] = exp_g_last_vec[1] - exp_g_last_regs[f32_reg_base + 2] = exp_g_last_vec[2] - exp_g_last_regs[f32_reg_base + 3] = exp_g_last_vec[3] - - # ---- k decay operand: exp2(g) * k ------------------------------ - k_decay_words = cutlass.Array(cutlass.Int32, 4, alignment=16) + exp_g_regs[f32_reg_base] = exp_g_frag[0] + exp_g_regs[f32_reg_base + 1] = exp_g_frag[1] + exp_g_regs[f32_reg_base + 2] = exp_g_frag[2] + exp_g_regs[f32_reg_base + 3] = exp_g_frag[3] + exp_neg_g_regs[half_reg_base] = cute.math.rcp(exp_g_frag[0], approx=True, ftz=True) + exp_neg_g_regs[half_reg_base + 1] = cute.math.rcp(exp_g_frag[1], approx=True, ftz=True) + exp_neg_g_regs[half_reg_base + 2] = cute.math.rcp(exp_g_frag[2], approx=True, ftz=True) + exp_neg_g_regs[half_reg_base + 3] = cute.math.rcp(exp_g_frag[3], approx=True, ftz=True) + exp_g_last_regs[f32_reg_base] = exp_g_last_frag[0] + exp_g_last_regs[f32_reg_base + 1] = exp_g_last_frag[1] + exp_g_last_regs[f32_reg_base + 2] = exp_g_last_frag[2] + exp_g_last_regs[f32_reg_base + 3] = exp_g_last_frag[3] + + # ---- K decay + K_inv operands: K * exp2(+g) and K * exp2(-g) ----- + k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 dim1 = dim0 + 1 @@ -1349,59 +1361,55 @@ def _compute0_warp_group( k_beta0, k_beta1 = fmul2(k_value0, k_value1, raw_beta_regs[raw_reg_idx0], raw_beta_regs[raw_reg_idx1]) k_pair = fp32_to_fp16(k_beta0, k_beta1, dtype=cfg.io_dtype) exp_g_pair = fp32_to_fp16(exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1], dtype=cfg.io_dtype) - k_decay_words[pair_idx] = mul_f16x2(k_pair, exp_g_pair, cfg.io_dtype) - k_inv_words[dim_half * 4 + pair_idx] = fp32_to_fp16(k_value0 * exp_neg_g_regs[dim0], k_value1 * exp_neg_g_regs[dim1], dtype=cfg.io_dtype) + k_decay_pack[pair_idx] = mul_f16x2(k_pair, exp_g_pair, cfg.io_dtype) + exp_neg_pair = fp32_to_fp16(exp_neg_g_regs[dim0], exp_neg_g_regs[dim1], dtype=cfg.io_dtype) + k_norm_pair = fp32_to_fp16(k_value0, k_value1, dtype=cfg.io_dtype) + k_inv_pack[dim_half * 4 + pair_idx] = mul_f16x2(k_norm_pair, exp_neg_pair, cfg.io_dtype) k_inv_vec = cutlass.Vector.from_elements( ( - k_inv_words[dim_half * 4], - k_inv_words[dim_half * 4 + 1], - k_inv_words[dim_half * 4 + 2], - k_inv_words[dim_half * 4 + 3], + k_inv_pack[dim_half * 4], + k_inv_pack[dim_half * 4 + 1], + k_inv_pack[dim_half * 4 + 2], + k_inv_pack[dim_half * 4 + 3], ), cutlass.Int32, ).bitcast(cfg.io_dtype) k_decay_vec = cutlass.Vector.from_elements( ( - k_decay_words[0], - k_decay_words[1], - k_decay_words[2], - k_decay_words[3], + k_decay_pack[0], + k_decay_pack[1], + k_decay_pack[2], + k_decay_pack[3], ), cutlass.Int32, ).bitcast(cfg.io_dtype) if cutlass.const_expr(dim_half == 0): - operand_done_phase = ((gc // cfg.smem_decay_stages) + 1) % 2 - bars.mb_kk_qk_super_mma_done[decay_stage].wait(operand_done_phase) - bars.mb_kk_qk_mma_done[decay_stage].wait(operand_done_phase) + operand_done_phase = ((global_chunk // cfg.smem_decay_stages) + 1) % 2 + bars.mb_decay_super_done[decay_stage].wait(operand_done_phase) + bars.mb_decay_tcgen05_done[decay_stage].wait(operand_done_phase) f16_segment = dim_base // 64 f16_segment_dim = dim_base - f16_segment * 64 k_inv_swizzled_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) (sK_inv_ptr + k_inv_swizzled_idx).store(k_inv_vec, alignment=16) - decay_storage_dim_base = _decay_key_dim( - cfg, - decay_row, - dim_base, + storage_key = dim_base ^ decay_key_mask + storage_slice = storage_key // 64 + decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( + decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 ) - decay_linear_idx_base = decay_row * cfg.d_k + decay_storage_dim_base - sw128_elems_per_128b = 128 // 2 - sw128_row = decay_linear_idx_base // cfg.d_k - sw128_col = decay_linear_idx_base - sw128_row * cfg.d_k - sw128_slice = sw128_col // sw128_elems_per_128b - sw128_col_in_slice = sw128_col - sw128_slice * sw128_elems_per_128b - sw128_slice_linear = sw128_row * sw128_elems_per_128b + sw128_col_in_slice - sw128_byte = sw128_slice_linear * 2 - sw128_mask = (sw128_byte >> 7 & 7) << 4 - decay_swizzled_idx_base = sw128_slice * cfg.b_t * sw128_elems_per_128b + (sw128_byte ^ sw128_mask) // 2 - (sK_decay_ptr + decay_swizzled_idx_base).store(k_decay_vec, alignment=16) + (sK_decay_ptr + decay_swizzled_idx).store(k_decay_vec, alignment=16) nvvm.fence_proxy("async.shared", space="cta") - bars.mb_k_decay_cg0_ready[decay_stage].arrive() + bars.mb_k_decay_inv_cg0_ready[decay_stage].arrive() + bars.mb_q_done[raw_stage].arrive() + bars.mb_k_done[raw_stage].arrive() + bars.mb_gate_done[raw_stage].arrive() + bars.mb_beta_done[raw_stage].arrive() - # ---- q decay operand ----------------------------------------------- + # ---- Q_decay operand: Q * q_inv_norm -------------------------------- for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 - q_decay_words = cutlass.Array(cutlass.Int32, 4, alignment=16) + q_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 dim1 = dim0 + 1 @@ -1410,68 +1418,68 @@ def _compute0_warp_group( q_value0, q_value1 = fmul2(raw_q_regs[raw_reg_idx0], raw_q_regs[raw_reg_idx1], q_inv_norm, q_inv_norm) q_pair = fp32_to_fp16(q_value0, q_value1, dtype=cfg.io_dtype) exp_g_pair = fp32_to_fp16(exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1], dtype=cfg.io_dtype) - q_decay_words[pair_idx] = mul_f16x2(q_pair, exp_g_pair, cfg.io_dtype) + q_decay_pack[pair_idx] = mul_f16x2(q_pair, exp_g_pair, cfg.io_dtype) q_decay_vec = cutlass.Vector.from_elements( ( - q_decay_words[0], - q_decay_words[1], - q_decay_words[2], - q_decay_words[3], + q_decay_pack[0], + q_decay_pack[1], + q_decay_pack[2], + q_decay_pack[3], ), cutlass.Int32, ).bitcast(cfg.io_dtype) - decay_storage_dim_base = _decay_key_dim( - cfg, - decay_row, - dim_base, + storage_key = dim_base ^ decay_key_mask + storage_slice = storage_key // 64 + decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( + decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 ) - decay_linear_idx_base = decay_row * cfg.d_k + decay_storage_dim_base - sw128_elems_per_128b = 128 // 2 - sw128_row = decay_linear_idx_base // cfg.d_k - sw128_col = decay_linear_idx_base - sw128_row * cfg.d_k - sw128_slice = sw128_col // sw128_elems_per_128b - sw128_col_in_slice = sw128_col - sw128_slice * sw128_elems_per_128b - sw128_slice_linear = sw128_row * sw128_elems_per_128b + sw128_col_in_slice - sw128_byte = sw128_slice_linear * 2 - sw128_mask = (sw128_byte >> 7 & 7) << 4 - decay_swizzled_idx_base = sw128_slice * cfg.b_t * sw128_elems_per_128b + (sw128_byte ^ sw128_mask) // 2 - (sQ_decay_ptr + decay_swizzled_idx_base).store(q_decay_vec, alignment=16) - - bars.mb_k_restore_done[decay_stage].wait(((gc // cfg.smem_decay_stages + 1) % 2)) - - # ---- k_restore operand ---------------------------------------------- + (sQ_decay_ptr + decay_swizzled_idx).store(q_decay_vec, alignment=16) + + # ---- K_restore operand: K_inv * exp_g_last -------------------------- + bars.mb_k_restore_acc_done[decay_stage].wait(((global_chunk // cfg.smem_decay_stages + 1) % 2)) for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 - k_restore_words = cutlass.Array(cutlass.Int32, 4, alignment=16) + k_restore_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 dim1 = dim0 + 1 exp_g_last_pair = fp32_to_fp16(exp_g_last_regs[reg_base + dim0], exp_g_last_regs[reg_base + dim1], dtype=cfg.io_dtype) - k_restore_words[pair_idx] = mul_f16x2(k_inv_words[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) + k_restore_pack[pair_idx] = mul_f16x2(k_inv_pack[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) storage_row = decay_row ^ (cfg.b_t // 2) f16_segment = dim_base // 64 f16_segment_dim = dim_base - f16_segment * 64 k_restore_idx = f16_segment * (cfg.b_t * 64) + storage_row * 64 + swizzle_xor_128b(storage_row, f16_segment_dim, elem_bytes=2) k_restore_vec = cutlass.Vector.from_elements( ( - k_restore_words[0], - k_restore_words[1], - k_restore_words[2], - k_restore_words[3], + k_restore_pack[0], + k_restore_pack[1], + k_restore_pack[2], + k_restore_pack[3], ), cutlass.Int32, ).bitcast(cfg.io_dtype) (sK_restore_ptr + k_restore_idx).store(k_restore_vec, alignment=16) nvvm.fence_proxy("async.shared", space="cta") bars.mb_qk_scale_ready[qk_scale_ready_stage].arrive() - gbase += sk_nt - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + diag_ring_stage = diag_ring_stage + cutlass.Int32(cfg.cg0_group_count) + wrapped = diag_ring_stage >= cutlass.Int32(cfg.smem_state_scale_diag_stages) + diag_ring_stage = diag_ring_stage - cutlass.Int32(cfg.smem_state_scale_diag_stages) if wrapped else diag_ring_stage + diag_ring_phase = diag_ring_phase ^ (cutlass.Int32(1) if wrapped else cutlass.Int32(0)) + raw_ring_stage = raw_ring_stage + cutlass.Int32(cfg.cg0_group_count) + raw_ring_wrapped = raw_ring_stage >= cutlass.Int32(cfg.smem_raw_stages) + raw_ring_stage = raw_ring_stage - cutlass.Int32(cfg.smem_raw_stages) if raw_ring_wrapped else raw_ring_stage + raw_bar_stage = raw_bar_stage + cutlass.Int32(cfg.cg0_group_count) + raw_bar_wrapped = raw_bar_stage >= cutlass.Int32(cfg.smem_raw_bar_stages) + raw_bar_stage = raw_bar_stage - cutlass.Int32(cfg.smem_raw_bar_stages) if raw_bar_wrapped else raw_bar_stage + raw_bar_phase = raw_bar_phase ^ (cutlass.Int32(1) if raw_bar_wrapped else cutlass.Int32(0)) + global_chunk_base += num_tile_chunks + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @cute.jit -def _compute1_warp_group( +def compute1_warp_group( cfg, total_tiles, bidx, @@ -1480,549 +1488,574 @@ def _compute1_warp_group( mWorkItems, sSched, lane, - tmem_hold, + sTmem_base, warp_idx, - mS_out, - mS_init, + mState_out, + mState_init, mO, sO_raw, sV_raw, sW_raw, - sH_raw, + sCheckpoint_raw, checkpoint_every_n_tokens, scale, bars, ) -> None: - """CG1 warp role (warps 8-11): persistent tile-scheduler loop + value-side - TMEM staging (state input, rhs, update), output drain to SMEM, and - checkpoint/final-state stores (split-K: only owned entries).""" + """CG1 warp role (warps 8-11): persistent scheduler loop running the + value-side TMEM staging, O drain, and checkpoint/final-state stores.""" nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) sO_ptr = sO_raw.data_ptr() - sH_ptr = sH_raw.data_ptr() if cutlass.const_expr(cfg.enable_checkpoints) else sO_raw.data_ptr() - h_done_index = PipelineState.start(phase=1) # sH starts free - tmem_base = tmem_hold.load() + sCheckpoint_ptr = sCheckpoint_raw.data_ptr() if cutlass.const_expr(cfg.enable_checkpoints) else sO_raw.data_ptr() + checkpoint_done_index = PipelineState.start(phase=1) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = sTmem_base.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 - tmem_sp = warp_idx % (cfg.d_v // cfg.threads_per_warp) - # ldmatrix.x4/stmatrix.x4 COL lane decode shared by the v loads and o stores - ov_tok = (lane // 16) * 8 + (lane & 7) - ov_col = ((lane // 8) & 1) * 8 - row_id = tmem_row + tmem_sp * cfg.threads_per_warp - value_dim = tmem_sp * cfg.threads_per_warp + lane + tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) + frag_row_coord = (lane // 16) * 8 + (lane & 7) + frag_col_offset = ((lane // 8) & 1) * 8 + row_id = tmem_row + tmem_subpartition * cfg.threads_per_warp + value_dim = tmem_subpartition * cfg.threads_per_warp + lane state_k_acc_index = PipelineState.start(phase=0) - update_acc_index = PipelineState.start(phase=0) + u_acc_index = PipelineState.start(phase=0) o_acc_index = PipelineState.start(phase=0) - kr_index = PipelineState.start(phase=0) # CG1's per-chunk mb_k_restore_done wait slot + k_restore_index = PipelineState.start(phase=0) # CG1's per-chunk mb_k_restore_acc_done wait slot raw_index = PipelineState.start(phase=0) # raw-ring slot for the sV/sW reads + inputs_done arrives - gbase = cutlass.Int32(0) + raw_bar_index = PipelineState.start(phase=0) # even-depth ready-ring slot (decoupled from the data ring) + global_chunk_base = cutlass.Int32(0) 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, cu_seqlens, mWorkItems - ) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) head_o = head_idx - sk_nt = wend - cstart - - if sk_nt > 0: - # ---- first chunk: seed state TMEM from mS_init (else zeros); - # split-K warmup items (cstart > 0) rebuild their state from zero - seed_from_s0 = cstart == 0 - 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_value = cutlass.Float32(0.0) - if cutlass.const_expr(mS_init is not None): - state_value = mS_init[batch_idx, head_o, key_dim, value_dim].to(cutlass.Float32) - if cutlass.const_expr(cfg.split_k): - state_value = state_value if seed_from_s0 else cutlass.Float32(0.0) - state_block[col] = state_value + num_tile_chunks = wend - cstart - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_offset + key_block_start), cutlass.Float32), - state_block[0:32], - ) + if num_tile_chunks > 0: + # ---- first chunk: seed state TMEM from mState_init ---------- + seed_from_initial_state = cstart == 0 + if cutlass.const_expr(mState_init is not None): + if seed_from_initial_state: + 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) - nvvm.tcgen05_wait("store") + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + state_block[0:32], + ) + else: + 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): + state_block[col] = cutlass.Float32(0.0) + + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + state_block[0:32], + ) + if cutlass.const_expr(mState_init is not None): + nvvm.tcgen05_wait("store") sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) - row_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 - state_col_id = tmem_col + cfg.tmem_state_offset - # ---- state -> packed b16 A operand (TMEM roundtrip) ---------------- - state_blocks = [] - for sub in cutlass.range_constexpr(cfg.d_k // 16): - state_blocks.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) - nvvm.tcgen05_wait("load") + row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + state_col_id = tmem_col + cfg.tmem_state_acc_offset + # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- packed_col_id = tmem_col + cfg.tmem_state_inp_offset - for sub in cutlass.range_constexpr(cfg.d_k // 16): - packed_state = cutlass.Array(cutlass.Int32, 8, alignment=16) - for packed_col in cutlass.range_constexpr(8): - source_pair = packed_col ^ 4 - packed_state[packed_col] = fp32_to_fp16(state_blocks[sub][2 * source_pair], state_blocks[sub][2 * source_pair + 1], dtype=cfg.io_dtype) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), - packed_state[0:8], - ) + if cutlass.const_expr(mState_init is not None): + state_vecs = [] + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_block * 16, cutlass.Float32), num=16)) + + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) + for packed_col in cutlass.range_constexpr(8): + source_pair = packed_col ^ 4 + state_pack[packed_col] = fp32_to_fp16( + state_vecs[k_block][2 * source_pair], state_vecs[k_block][2 * source_pair + 1], dtype=cfg.io_dtype + ) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + k_block * 8, cutlass.Int8), + state_pack[0:8], + ) - nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready.arrive() + nvvm.tcgen05_wait("store") + bars.mb_state_inp_ready.arrive() + if cutlass.const_expr(cfg.enable_checkpoints): + bars.mb_state_acc_read_done.arrive() - # ---- rhs staging: rhs input = w*v - state*(beta*k) ----------------- - bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + # ---- 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 - input_col_id = tmem_col + cfg.tmem_rhs_inp_offset - value_dim_base = tmem_sp * cfg.threads_per_warp + input_col_id = tmem_col + cfg.tmem_y_inp_offset + value_dim_base = tmem_subpartition * cfg.threads_per_warp - # ---- read back state*k acc + raw v fragments ----------------------- + # ---- raw V fragments, then W, then the state*K acc readback ---------- row_id0 = tmem_row + value_dim_base - state_k0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - row_id1 = row_id0 + 16 - state_k1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) - - raw_v_regs0 = nvvm.ldmatrix( + raw_v_frag0 = nvvm.ldmatrix( sV_ptr - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_v_regs1 = nvvm.ldmatrix( + raw_v_frag1 = nvvm.ldmatrix( sV_ptr - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_w_regs0 = nvvm.ldmatrix( + bars.mb_w_ready[raw_bar_index.idx].wait(raw_bar_index.phase) + raw_w_frag0 = nvvm.ldmatrix( sW_ptr - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_w_regs1 = nvvm.ldmatrix( + raw_w_frag1 = nvvm.ldmatrix( sW_ptr - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - nvvm.tcgen05_wait("load") + if cutlass.const_expr(mState_init is not None): + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) + state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) - packed_rhs0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k0[frag_pair], state_k0[frag_pair + 1] - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - w_pair = raw_w_regs0[raw_matrix] - if cutlass.const_expr(cfg.beta_w_sigmoid): - w_gate0, w_gate1 = sigmoid_f16x2(w_pair, cfg.io_dtype) - w_pair = fp32_to_fp16(w_gate0, w_gate1, dtype=cfg.io_dtype) + w_pair = raw_w_frag0[raw_matrix] wv_pair = mul_f16x2( w_pair, - raw_v_regs0[raw_matrix], - cfg.io_dtype, - ) - packed_rhs0[reg_idx] = sub_f16x2( - wv_pair, - state_k_pair, + raw_v_frag0[raw_matrix], cfg.io_dtype, ) + if cutlass.const_expr(mState_init is not None): + state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + y_inp_pack0[reg_idx] = sub_f16x2( + wv_pair, + state_k_pair, + cfg.io_dtype, + ) + else: + y_inp_pack0[reg_idx] = wv_pair - packed_rhs1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k1[frag_pair], state_k1[frag_pair + 1] - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - w_pair = raw_w_regs1[raw_matrix] - if cutlass.const_expr(cfg.beta_w_sigmoid): - w_gate0, w_gate1 = sigmoid_f16x2(w_pair, cfg.io_dtype) - w_pair = fp32_to_fp16(w_gate0, w_gate1, dtype=cfg.io_dtype) + w_pair = raw_w_frag1[raw_matrix] wv_pair = mul_f16x2( w_pair, - raw_v_regs1[raw_matrix], - cfg.io_dtype, - ) - packed_rhs1[reg_idx] = sub_f16x2( - wv_pair, - state_k_pair, + raw_v_frag1[raw_matrix], cfg.io_dtype, ) + if cutlass.const_expr(mState_init is not None): + state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + y_inp_pack1[reg_idx] = sub_f16x2( + wv_pair, + state_k_pair, + cfg.io_dtype, + ) + else: + y_inp_pack1[reg_idx] = wv_pair - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), packed_rhs0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), packed_rhs1[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) nvvm.tcgen05_wait("store") - state_k_acc_index = advance(state_k_acc_index, 1) - bars.mb_inputs_done[raw_index.idx].arrive() - bars.mb_rhs_ready.arrive() - - # ---- update readback -> packed b16 A operand ----------------------- - bars.mb_update_acc_ready.wait(update_acc_index.phase) - update = nvvm.tcgen05_ld( + if cutlass.const_expr(mState_init is not None): + state_k_acc_index = advance(state_k_acc_index, 1) + bars.mb_v_done[raw_index.idx].arrive() + bars.mb_w_done[raw_index.idx].arrive() + bars.mb_y_inp_ready.arrive() + + # ---- U repack: acc TMEM -> packed b16 TMEM -------------------------- + bars.mb_u_acc_ready.wait(u_acc_index.phase) + u_acc_vals = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((tmem_row + tmem_sp * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_update_acc_offset), cutlass.Float32), + nvvm.make_tmem_ptr((tmem_row + tmem_subpartition * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), num=cfg.b_t, ) - nvvm.tcgen05_wait("load") - packed_update = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): source_pair = packed_col ^ 4 token0 = source_pair * 2 token1 = token0 + 1 - packed_update[packed_col] = fp32_to_fp16(update[token0], update[token1], dtype=cfg.io_dtype) + u_inp_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_update_inp_offset), cutlass.Int8), - packed_update[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_u_inp_offset), cutlass.Int8), + u_inp_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") - update_acc_index = advance(update_acc_index, 1) - bars.mb_update_ready.arrive() + u_acc_index = advance(u_acc_index, 1) + bars.mb_u_inp_ready.arrive() - bars.mb_k_restore_done[kr_index.idx].wait(kr_index.phase) - kr_index = advance(kr_index, cfg.smem_decay_stages) + bars.mb_k_restore_acc_done[k_restore_index.idx].wait(k_restore_index.phase) + k_restore_index = advance(k_restore_index, cfg.smem_decay_stages) raw_index = advance(raw_index, cfg.smem_raw_stages) + raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) - # the first chunk is peeled above so this steady-state loop always - # drains the prior chunk's output - for li in cutlass.range(1, sk_nt, 1, unroll=1): - chunk_idx = cstart + li - gc = gbase + li + if cutlass.const_expr(cfg.enable_checkpoints): + cg1_checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + cg1_checkpoint_mod = (cstart + cutlass.Int32(1)) % cg1_checkpoint_chunks + for local_chunk in cutlass.range(1, num_tile_chunks, 1, unroll=1): + chunk_idx = cstart + local_chunk + global_chunk = global_chunk_base + local_chunk sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) prev_output_chunk = chunk_idx - cutlass.Int32(1) - prev_og = gc - cutlass.Int32(1) - prev_o_stage = prev_og % cfg.smem_o_stages - prev_q_state_acc_stage = prev_og % cfg.tmem_q_state_acc_stages + prev_global_chunk = global_chunk - cutlass.Int32(1) + prev_o_stage = prev_global_chunk % cfg.smem_o_stages + prev_q_state_acc_stage = prev_global_chunk % cfg.tmem_q_state_acc_stages prev_o_stage_base = prev_o_stage * (cfg.b_t * cfg.d_v) - # H entry gate: the state read by this restage entered chunk_idx, - # i.e. the state after chunk_idx * b_t tokens (strictly before the - # sequence end -- the end state is only final_state); split-K - # warmup reconstructions below wstart belong to the previous item - do_h = False + do_checkpoint = False if cutlass.const_expr(cfg.enable_checkpoints): - do_h = (chunk_idx * cutlass.Int32(cfg.b_t)) % checkpoint_every_n_tokens == 0 - if cutlass.const_expr(cfg.split_k): - do_h = do_h and chunk_idx >= wstart - if do_h: - # sH is free once the epilogue's previous TMA store retired - bars.mb_h_tmastg_done.wait(h_done_index.phase) - h_done_index = advance(h_done_index, 1) - row_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 - state_col_id = tmem_col + cfg.tmem_state_offset - # ---- state -> packed b16 A operand (TMEM roundtrip) ---------------- - state_blocks = [] - for sub in cutlass.range_constexpr(cfg.d_k // 16): - state_blocks.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) - bars.mb_o_tmastg_done[prev_o_stage].wait(((prev_og // cfg.smem_o_stages) + 1) % 2) - nvvm.tcgen05_wait("load") + do_checkpoint = cg1_checkpoint_mod == 0 + cg1_checkpoint_mod = cg1_checkpoint_mod + cutlass.Int32(1) + cg1_checkpoint_mod = cutlass.Int32(0) if cg1_checkpoint_mod == cg1_checkpoint_chunks else cg1_checkpoint_mod + do_checkpoint = do_checkpoint and chunk_idx >= wstart + row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + state_col_id = tmem_col + cfg.tmem_state_acc_offset + + # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + state_vecs = [] + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_block * 16, cutlass.Float32), num=16)) packed_col_id = tmem_col + cfg.tmem_state_inp_offset - for sub in cutlass.range_constexpr(cfg.d_k // 16): - packed_state = cutlass.Array(cutlass.Int32, 8, alignment=16) + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) for packed_col in cutlass.range_constexpr(8): source_pair = packed_col ^ 4 - packed_state[packed_col] = fp32_to_fp16(state_blocks[sub][2 * source_pair], state_blocks[sub][2 * source_pair + 1], dtype=cfg.io_dtype) + state_pack[packed_col] = fp32_to_fp16(state_vecs[k_block][2 * source_pair], state_vecs[k_block][2 * source_pair + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), - packed_state[0:8], + nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + k_block * 8, cutlass.Int8), + state_pack[0:8], ) - if cutlass.const_expr(cfg.enable_checkpoints): - # stage this sub's state to sH TRANSPOSED (KV: k rows, v - # contiguous in 64-v slabs, swizzled — the GDN H layout); - # each thread scatters its 16 k values down one v column - if do_h: - h_col = value_dim % cutlass.Int32(64) - h_seg = (value_dim // cutlass.Int32(64)) * (cfg.d_k * cutlass.Int32(64)) - k_base = sub * 16 - for j in cutlass.range_constexpr(16): - hv = state_blocks[sub][j].to(cfg.io_dtype) - (sH_ptr + h_seg + cutlass.Int32((k_base + j) * 64) + swizzle_xor_128b(cutlass.Int32(k_base + j), h_col, elem_bytes=2)).store(hv) - nvvm.tcgen05_wait("store") bars.mb_state_inp_ready.arrive() + + # ---- checkpoint store ----------------------------------------------- if cutlass.const_expr(cfg.enable_checkpoints): - if do_h: + if do_checkpoint: + 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) + row16_addr = ((tmem_row + tmem_subpartition * cfg.threads_per_warp) + 16) << 16 + checkpoint_vbase = tmem_subpartition * cfg.threads_per_warp + checkpoint_swz_off0 = (checkpoint_vbase + frag_col_offset) // 64 * (cfg.d_k * 64) + 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, + ) + nvvm.tcgen05_wait("load") + bars.mb_state_acc_read_done.arrive() nvvm.fence_proxy("async.shared", space="cta") - bars.mb_h_tmastg_ready.arrive() + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + else: + bars.mb_state_acc_read_done.arrive() bars.mb_o_acc_ready.wait(o_acc_index.phase) o_acc_index = advance(o_acc_index, 1) projection_col_id = tmem_col + cfg.tmem_q_state_acc_offset + prev_q_state_acc_stage * cfg.b_t - value_dim_base = tmem_sp * cfg.threads_per_warp + value_dim_base = tmem_subpartition * cfg.threads_per_warp row_id0 = tmem_row + value_dim_base row_id1 = row_id0 + 16 - loaded0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - loaded1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + loaded_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) + loaded_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) - # ---- output drain: q_state_acc -> scaled b16 -> SMEM stmatrix ------- - stsm_regs0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - stsm_regs1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + # ---- output drain: O acc TMEM -> scaled b16 SMEM -------------------- + stsm_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + stsm_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - scaled0_0, scaled0_1 = fmul2(loaded0[2 * reg_idx], loaded0[2 * reg_idx + 1], scale, scale) - scaled1_0, scaled1_1 = fmul2(loaded1[2 * reg_idx], loaded1[2 * reg_idx + 1], scale, scale) - stsm_regs0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) - stsm_regs1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + scaled0_0, scaled0_1 = fmul2(loaded_vec0[2 * reg_idx], loaded_vec0[2 * reg_idx + 1], scale, scale) + scaled1_0, scaled1_1 = fmul2(loaded_vec1[2 * reg_idx], loaded_vec1[2 * reg_idx + 1], scale, scale) + stsm_pack0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) + stsm_pack1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + bars.mb_o_tmastg_done[prev_o_stage].wait(((prev_global_chunk // cfg.smem_o_stages) + 1) % 2) nvvm.stmatrix( sO_ptr + prev_o_stage_base - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), - stsm_regs0.data_ptr().load(count=4, alignment=4), + + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), + stsm_pack0.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( sO_ptr + prev_o_stage_base - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), - stsm_regs1.data_ptr().load(count=4, alignment=4), + + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), + stsm_pack1.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) - # release only after the stmatrix pair: the STSM->F2FP->FMUL2->LDTM - # register chain pins the TMEM reads complete without a wait("load") bars.mb_o_acc_done[prev_q_state_acc_stage].arrive() nvvm.fence_proxy("async.shared", space="cta") bars.mb_o_tmastg_ready[prev_o_stage].arrive() - bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - # ---- rhs staging: rhs input = w*v - state*(beta*k) ----------------- + # ---- 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 - input_col_id = tmem_col + cfg.tmem_rhs_inp_offset - value_dim_base = tmem_sp * cfg.threads_per_warp + input_col_id = tmem_col + cfg.tmem_y_inp_offset + value_dim_base = tmem_subpartition * cfg.threads_per_warp - # ---- read back state*k acc + raw v fragments ----------------------- + # ---- raw V fragments, then W, then the state*K acc readback ---------- row_id0 = tmem_row + value_dim_base - state_k0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - row_id1 = row_id0 + 16 - state_k1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) - - raw_v_regs0 = nvvm.ldmatrix( + raw_v_frag0 = nvvm.ldmatrix( sV_ptr - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_v_regs1 = nvvm.ldmatrix( + raw_v_frag1 = nvvm.ldmatrix( sV_ptr - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_w_regs0 = nvvm.ldmatrix( + bars.mb_w_ready[raw_bar_index.idx].wait(raw_bar_index.phase) + raw_w_frag0 = nvvm.ldmatrix( sW_ptr - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - raw_w_regs1 = nvvm.ldmatrix( + raw_w_frag1 = nvvm.ldmatrix( sW_ptr - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), 4, nvvm.MMALayout.COL, ) - nvvm.tcgen05_wait("load") - packed_rhs0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) + state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + + y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k0[frag_pair], state_k0[frag_pair + 1] + state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - w_pair = raw_w_regs0[raw_matrix] - if cutlass.const_expr(cfg.beta_w_sigmoid): - w_gate0, w_gate1 = sigmoid_f16x2(w_pair, cfg.io_dtype) - w_pair = fp32_to_fp16(w_gate0, w_gate1, dtype=cfg.io_dtype) + w_pair = raw_w_frag0[raw_matrix] wv_pair = mul_f16x2( w_pair, - raw_v_regs0[raw_matrix], + raw_v_frag0[raw_matrix], cfg.io_dtype, ) - packed_rhs0[reg_idx] = sub_f16x2( + y_inp_pack0[reg_idx] = sub_f16x2( wv_pair, state_k_pair, cfg.io_dtype, ) - packed_rhs1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k1[frag_pair], state_k1[frag_pair + 1] + state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - w_pair = raw_w_regs1[raw_matrix] - if cutlass.const_expr(cfg.beta_w_sigmoid): - w_gate0, w_gate1 = sigmoid_f16x2(w_pair, cfg.io_dtype) - w_pair = fp32_to_fp16(w_gate0, w_gate1, dtype=cfg.io_dtype) + w_pair = raw_w_frag1[raw_matrix] wv_pair = mul_f16x2( w_pair, - raw_v_regs1[raw_matrix], + raw_v_frag1[raw_matrix], cfg.io_dtype, ) - packed_rhs1[reg_idx] = sub_f16x2( + y_inp_pack1[reg_idx] = sub_f16x2( wv_pair, state_k_pair, cfg.io_dtype, ) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), packed_rhs0[0:4]) - - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), packed_rhs1[0:4]) - + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) nvvm.tcgen05_wait("store") state_k_acc_index = advance(state_k_acc_index, 1) - bars.mb_inputs_done[raw_index.idx].arrive() - bars.mb_rhs_ready.arrive() + bars.mb_v_done[raw_index.idx].arrive() + bars.mb_w_done[raw_index.idx].arrive() + bars.mb_y_inp_ready.arrive() - # ---- update readback -> packed b16 A operand ----------------------- - bars.mb_update_acc_ready.wait(update_acc_index.phase) - update = nvvm.tcgen05_ld( + # ---- U repack: acc TMEM -> packed b16 TMEM -------------------------- + bars.mb_u_acc_ready.wait(u_acc_index.phase) + u_acc_vals = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((tmem_row + tmem_sp * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_update_acc_offset), cutlass.Float32), + nvvm.make_tmem_ptr((tmem_row + tmem_subpartition * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), num=cfg.b_t, ) - nvvm.tcgen05_wait("load") - packed_update = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): source_pair = packed_col ^ 4 token0 = source_pair * 2 token1 = token0 + 1 - packed_update[packed_col] = fp32_to_fp16(update[token0], update[token1], dtype=cfg.io_dtype) + u_inp_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_update_inp_offset), cutlass.Int8), - packed_update[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_u_inp_offset), cutlass.Int8), + u_inp_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") - update_acc_index = advance(update_acc_index, 1) - bars.mb_update_ready.arrive() + u_acc_index = advance(u_acc_index, 1) + bars.mb_u_inp_ready.arrive() - bars.mb_k_restore_done[kr_index.idx].wait(kr_index.phase) - kr_index = advance(kr_index, cfg.smem_decay_stages) + bars.mb_k_restore_acc_done[k_restore_index.idx].wait(k_restore_index.phase) + k_restore_index = advance(k_restore_index, cfg.smem_decay_stages) raw_index = advance(raw_index, cfg.smem_raw_stages) + raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) - if sk_nt > 0: - og = gbase + sk_nt - cutlass.Int32(1) + if num_tile_chunks > 0: + last_global_chunk = global_chunk_base + num_tile_chunks - cutlass.Int32(1) output_chunk = wend - cutlass.Int32(1) - final_o_stage = og % cfg.smem_o_stages - final_q_state_acc_stage = og % cfg.tmem_q_state_acc_stages + final_o_stage = last_global_chunk % cfg.smem_o_stages + final_q_state_acc_stage = last_global_chunk % cfg.tmem_q_state_acc_stages final_o_stage_base = final_o_stage * (cfg.b_t * cfg.d_v) - bars.mb_o_tmastg_done[final_o_stage].wait(((og // cfg.smem_o_stages) + 1) % 2) bars.mb_o_acc_ready.wait(o_acc_index.phase) o_acc_index = advance(o_acc_index, 1) projection_col_id = tmem_col + cfg.tmem_q_state_acc_offset + final_q_state_acc_stage * cfg.b_t - value_dim_base = tmem_sp * cfg.threads_per_warp + value_dim_base = tmem_subpartition * cfg.threads_per_warp row_id0 = tmem_row + value_dim_base row_id1 = row_id0 + 16 - loaded0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - loaded1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + loaded_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) + loaded_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) - # ---- output drain: q_state_acc -> scaled b16 -> SMEM stmatrix ------- - stsm_regs0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - stsm_regs1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + # ---- output drain: O acc TMEM -> scaled b16 SMEM -------------------- + stsm_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + stsm_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - scaled0_0, scaled0_1 = fmul2(loaded0[2 * reg_idx], loaded0[2 * reg_idx + 1], scale, scale) - scaled1_0, scaled1_1 = fmul2(loaded1[2 * reg_idx], loaded1[2 * reg_idx + 1], scale, scale) - stsm_regs0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) - stsm_regs1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + scaled0_0, scaled0_1 = fmul2(loaded_vec0[2 * reg_idx], loaded_vec0[2 * reg_idx + 1], scale, scale) + scaled1_0, scaled1_1 = fmul2(loaded_vec1[2 * reg_idx], loaded_vec1[2 * reg_idx + 1], scale, scale) + stsm_pack0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) + stsm_pack1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + bars.mb_o_tmastg_done[final_o_stage].wait(((last_global_chunk // cfg.smem_o_stages) + 1) % 2) nvvm.stmatrix( sO_ptr + final_o_stage_base - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), - stsm_regs0.data_ptr().load(count=4, alignment=4), + + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2), + stsm_pack0.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( sO_ptr + final_o_stage_base - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), - stsm_regs1.data_ptr().load(count=4, alignment=4), + + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2), + stsm_pack1.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) - # release only after the stmatrix pair: the STSM->F2FP->FMUL2->LDTM - # register chain pins the TMEM reads complete without a wait("load") bars.mb_o_acc_done[final_q_state_acc_stage].arrive() nvvm.fence_proxy("async.shared", space="cta") bars.mb_o_tmastg_ready[final_o_stage].arrive() - # split-K: only the item owning the sequence's last chunk holds the - # true end-of-sequence state (legacy tiles always do: wend == nc) owns_final = wend == num_chunks_b - # ---- final-state drain: final_state acc -> GMEM -------------------- - if cutlass.const_expr(mS_out is not None): + + # ---- final-state drain: state acc TMEM -> GMEM --------------------------- + if cutlass.const_expr(mState_out is not None): if seqlen_b > 0: if owns_final: for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): loaded = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_offset + key_block_start), cutlass.Float32), + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), num=32, ) - nvvm.tcgen05_wait("load") for col in cutlass.range_constexpr(32): key_dim = key_block_start + col - mS_out[batch_idx, head_o, key_dim, value_dim] = loaded[col].to(mS_out.element_type) + mState_out[batch_idx, head_o, key_dim, value_dim] = loaded[col].to(mState_out.element_type) else: - # zero-length sequence: the state passes through untouched - # (S0 when seeded, zeros otherwise); pure GMEM, no TMEM 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(mS_init is not None): - mS_out[batch_idx, head_o, key_dim, value_dim] = mS_init[batch_idx, head_o, key_dim, value_dim] + 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] else: - mS_out[batch_idx, head_o, key_dim, value_dim] = cutlass.Float32(0.0).to(mS_out.element_type) - bars.mb_final_state_stored.arrive() - gbase += sk_nt - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + mState_out[batch_idx, head_o, key_dim, value_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) + + bars.mb_tmem_done[0].arrive() @cute.jit -def _host( +def host( cfg: cutlass.Constexpr, q: cute.Tensor, k: cute.Tensor, @@ -2045,15 +2078,11 @@ def _host( stream, ) -> None: num_sequences = cu_seqlens.shape[0] - 1 - ho = raw_gate.shape[1] - # ---- persistent launch: the grid only needs to cover the tiles ------ - total_tiles = num_sequences * ho - # CUDA-graph-stable launch: fixed SM-count grid grid_shape = (cfg.max_active_clusters, 1, 1) - _kernel( + kernel( cfg, tensormap_workspace, - cutlass.Int32(num_sequences * ho), + cutlass.Int32(num_sequences), q, k, v, @@ -2069,7 +2098,6 @@ def _host( work_items, work_count, sched_ctr, - total_tiles, scale, checkpoint_every_n_tokens, ).launch( @@ -2081,7 +2109,7 @@ def _host( @cute.kernel -def _kernel( +def kernel( cfg: cutlass.Constexpr, tensormap_workspace: cute.Tensor, n_desc: cutlass.Int32, @@ -2094,13 +2122,12 @@ def _kernel( mBeta: cute.Tensor, mW: cute.Tensor, cu_seqlens: cute.Tensor, - mS_init: cute.Tensor | None, + mState_init: cute.Tensor | None, mO: cute.Tensor, - mS_out: cute.Tensor | None, - mWorkItems: cute.Tensor | None, - mCount: cute.Tensor | None, + mState_out: cute.Tensor | None, + mWorkItems: cute.Tensor, + mCount: cute.Tensor, mSched: cute.Tensor | None, - total_tiles: cutlass.Int32, scale: cutlass.Float32, checkpoint_every_n_tokens: cutlass.Int32, ) -> None: @@ -2118,8 +2145,7 @@ def _kernel( warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) lane = tidx % cfg.threads_per_warp - if cutlass.const_expr(cfg.split_k): - total_tiles = mCount[0] + total_tiles = mCount[0] if cutlass.const_expr(cfg.dyn_sched): assert mSched is not None and mSched.element_type == cutlass.Int32 assert mQ.element_type == cfg.io_dtype and mK.element_type == cfg.io_dtype and mV.element_type == cfg.io_dtype @@ -2127,16 +2153,16 @@ def _kernel( assert mBeta.element_type == cfg.io_dtype and mW.element_type == cfg.io_dtype, "channel-wise beta/w must match the io dtype" assert cu_seqlens.element_type in (cutlass.Int32, cutlass.Int64) if cutlass.const_expr(cfg.use_initial_state): - assert mS_init is not None and mS_init.element_type in (cutlass.BFloat16, cutlass.Float32) + assert mState_init is not None and mState_init.element_type in (cutlass.BFloat16, cutlass.Float32) else: - assert mS_init is None, "mS_init must be None if use_initial_state is False" + assert mState_init is None, "mState_init must be None if use_initial_state is False" if cutlass.const_expr(cfg.store_final_state): - assert mS_out is not None and mS_out.element_type in (cutlass.BFloat16, cutlass.Float32) + assert mState_out is not None and mState_out.element_type in (cutlass.BFloat16, cutlass.Float32) else: - assert mS_out is None, "mS_out must be None if store_final_state is False" - if cutlass.const_expr(mS_init is not None and mS_out is not None): - assert mS_init.element_type == mS_out.element_type - # per-(batch, head) TMA-descriptor arrays: [q, k, v, gate, beta, w, o] + assert mState_out is None, "mState_out must be None if store_final_state is False" + if cutlass.const_expr(mState_init is not None and mState_out is not None): + assert mState_init.element_type == mState_out.element_type + # per-BATCH TMA-descriptor arrays (heads are load coordinates): [Q, K, V, Gate, Beta, W, O] desc_base_words = tensormap_workspace.iterator.raw_ptr() arr_words = n_desc * cutlass.Int32(TENSOR_MAP_QWORDS) desc_q_base = desc_base_words @@ -2146,19 +2172,17 @@ def _kernel( desc_beta_base = desc_base_words + cutlass.Int32(4) * arr_words desc_w_base = desc_base_words + cutlass.Int32(5) * arr_words desc_o_base = desc_base_words + cutlass.Int32(6) * arr_words - desc_h_base = desc_base_words + cutlass.Int32(7) * arr_words + desc_checkpoint_base = desc_base_words + cutlass.Int32(7) * arr_words # Buffers are declaration-ordered and intentionally non-aliased. SMEM = cutlass.AddressSpace.smem bars = make_gdn2_bars(cfg) - # The hand-written K-box-major SW128 mapping is normalized to phase 0, - # so both tcgen05 and ldmatrix can share 1KB-aligned operand buffers. - tmem_hold = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) + sTmem_base = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) sK_decay_raw = cutlass.Array(cfg.io_dtype, cfg.k_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sQ_decay_raw = cutlass.Array(cfg.io_dtype, cfg.q_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sK_restore_raw = cutlass.Array(cfg.io_dtype, cfg.k_restore_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) - sPairwise_raw = cutlass.Array(cfg.io_dtype, cfg.pairwise_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sIntermediate_raw = cutlass.Array(cfg.io_dtype, cfg.intermediate_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sQ_raw = cutlass.Array(mQ.element_type, cfg.q_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sK_raw = cutlass.Array(mK.element_type, cfg.k_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sV_raw = cutlass.Array(mV.element_type, cfg.v_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) @@ -2169,18 +2193,15 @@ def _kernel( mO.element_type, cfg.o_cosize, space=SMEM, - # The scalar CG1 store computes W128 offsets relative to this buffer. - # Align to the full s128b period so absolute SMEM address bits do not - # add a hidden phase to the TMA store-side swizzle. alignment=cfg.buffer_align_bytes, ) sBeta_raw = cutlass.Array(mBeta.element_type, cfg.beta_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sW_raw = cutlass.Array(mW.element_type, cfg.w_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) - sH_raw = ( - cutlass.Array(cfg.io_dtype, cfg.d_k * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) if cutlass.const_expr(cfg.enable_checkpoints) else sO_raw + sCheckpoint_raw = ( + cutlass.Array(cfg.io_dtype, cfg.smem_checkpoint_stages * cfg.d_k * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) + if cutlass.const_expr(cfg.enable_checkpoints) + else sO_raw ) - # K-box-major SW128 staging: 16B leading offset, 1KB stride; the decay - # stores apply a row-group key xor so tcgen05 B reads logical [DK, BT]. sK_decay = SmemTile( base=sK_decay_raw, elems_per_stage=(cfg.d_k * cfg.b_t), @@ -2213,58 +2234,61 @@ def _kernel( stride_byte_offset=(8 * 16 * 2), layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, ) - sPairwise = SmemTile( - base=sPairwise_raw, + sIntermediate = SmemTile( + base=sIntermediate_raw, elems_per_stage=(2 * cfg.b_t * cfg.b_t), - stages=cfg.smem_pairwise_stages, + stages=cfg.smem_intermediate_stages, leading_byte_offset=16, stride_byte_offset=(8 * cfg.b_t * 2), layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, ) - tma_tx_bytes = cutlass.const_expr( - cfg.d_k * cfg.b_t * mQ.element_type.width // 8 - + cfg.d_k * cfg.b_t * mK.element_type.width // 8 - + cfg.d_v * cfg.b_t * mV.element_type.width // 8 - + cfg.d_k * cfg.b_t * mGate.element_type.width // 8 - + cfg.d_k * cfg.b_t * mBeta.element_type.width // 8 - + cfg.d_v * cfg.b_t * mW.element_type.width // 8 - ) + elect_one = nvvm.elect_sync() if warp_idx == cfg.tma_warp_id: - if nvvm.elect_sync(): - bars.mb_tma_done.init() + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_raw_bar_stages): + bars.mb_q_ready[stage].init() + bars.mb_k_ready[stage].init() + bars.mb_gate_ready[stage].init() + bars.mb_beta_ready[stage].init() + bars.mb_v_ready[stage].init() + bars.mb_w_ready[stage].init() for stage in cutlass.range_constexpr(cfg.smem_raw_stages): - bars.mb_inputs_ready[stage].init() - bars.mb_inputs_done[stage].init() + bars.mb_q_done[stage].init() + bars.mb_k_done[stage].init() + bars.mb_gate_done[stage].init() + bars.mb_beta_done[stage].init() + bars.mb_v_done[stage].init() + bars.mb_w_done[stage].init() elif warp_idx == cfg.tcgen05_mma_warp_id: - if nvvm.elect_sync(): + if elect_one: bars.mb_o_acc_ready.init() for stage in cutlass.range_constexpr(cfg.tmem_q_state_acc_stages): bars.mb_o_acc_done[stage].init() bars.mb_state_k_acc_ready.init() - bars.mb_update_acc_ready.init() + bars.mb_u_acc_ready.init() bars.mb_state_inp_ready.init() for stage in cutlass.range_constexpr(cfg.smem_state_scale_diag_stages): bars.mb_state_scale_diag_done[stage].init() for stage in cutlass.range_constexpr(cfg.smem_decay_stages): - bars.mb_kk_qk_super_mma_done[stage].init() - bars.mb_kk_qk_mma_done[stage].init() - bars.mb_k_restore_done[stage].init() - bars.mb_rhs_ready.init() - bars.mb_update_ready.init() - bars.mb_final_state_stored.init() + bars.mb_decay_tcgen05_done[stage].init() + bars.mb_decay_super_done[stage].init() + bars.mb_k_restore_acc_done[stage].init() + bars.mb_y_inp_ready.init() + bars.mb_u_inp_ready.init() + bars.mb_tmem_done[0].init() elif warp_idx == cfg.super_mma_warp_id: - if nvvm.elect_sync(): - for stage in cutlass.range_constexpr(cfg.smem_pairwise_stages): + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_intermediate_stages): + bars.mb_t_inv_ready[stage].init() bars.mb_a_ready[stage].init() - bars.mb_qk_acc_ready[stage].init() - bars.mb_a_done[stage].init() + bars.mb_intermediate_done[stage].init() for stage in cutlass.range_constexpr(cfg.qk_scale_ready_stages): bars.mb_qk_scale_ready[stage].init() for stage in cutlass.range_constexpr(cfg.smem_decay_stages): - bars.mb_k_decay_cg0_ready[stage].init() + bars.mb_k_decay_inv_cg0_ready[stage].init() elif warp_idx == cfg.epilogue_warp_id: - if nvvm.elect_sync(): + if elect_one: for stage in cutlass.range_constexpr(cfg.smem_o_stages): bars.mb_o_tmastg_ready[stage].init() bars.mb_o_tmastg_done[stage].init() @@ -2272,39 +2296,17 @@ def _kernel( bars.mb_sched_ready[stage].init() bars.mb_sched_done[stage].init() if cutlass.const_expr(cfg.enable_checkpoints): - bars.mb_h_tmastg_ready.init() - bars.mb_h_tmastg_done.init() + for stage in cutlass.range_constexpr(cfg.smem_checkpoint_stages): + bars.mb_checkpoint_tmastg_ready[stage].init() + bars.mb_checkpoint_tmastg_done[stage].init() + bars.mb_state_acc_read_done.init() diag_zero = cfg.io_dtype(0.0) for diag_idx in cutlass.range(tidx, cfg.state_scale_diag_cosize, cfg.threads_per_cta, unroll=1): sState_scale_diag_raw[diag_idx] = diag_zero nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync(0, thread_count=cfg.threads_per_cta) - if (warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]) or warp_idx == cfg.tcgen05_mma_warp_id: - if warp_idx == cfg.tcgen05_mma_warp_id: - nvvm.tcgen05_alloc(tmem_hold, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) - nvvm.barrier_cta_sync(cfg.nbar_tmem_lifecycle_id, thread_count=cfg.tmem_user_threads) - if warp_idx == cfg.tcgen05_mma_warp_id: - nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) - nvvm.barrier_cta_sync(cfg.nbar_tmem_lifecycle_id, thread_count=cfg.tmem_user_threads) - - # Actual SMEM/TMEM buffers for the BT=16 schedule: - # q/k/v : 16 x 128 each - # gate_log2 : 16 x 128 - # beta : 16 - # q/k inverse norm : 16 each, staged once per decay stage - # exp_g_last : 128, CG0-local and staged once per decay stage - # state-scale diag : 8 x 16 x 16 input dtype, zeroed once in the prologue - # q_decay/k_decay : tcgen05 SW128 operands shared with super-MMA - # k_restore : tcgen05 SW128 N-major final-state operand - # k_inv : 16 x 128 token-major for super-MMA RHS - # A inverse/QK : 16 x 16 each, plus transposed tcgen05 operands - # state : external/kernel ABI is VK `[DV, DK]`; reference - # math can view it as KV `[DK, DV]` by transposing. - # The TS A-staging path keeps VK in TMEM so state*k - # is `[DV, DK] @ [DK, BT] -> [DV, BT]` with M=128. - if warp_idx == cfg.tma_warp_id: - _tmaldg_warp( + tmaldg_warp( cfg, total_tiles, bidx, @@ -2314,7 +2316,6 @@ def _kernel( mSched, sSched, lane, - tma_tx_bytes, sBeta_raw, sGate_raw, sK_raw, @@ -2330,7 +2331,7 @@ def _kernel( bars, ) elif warp_idx == cfg.super_mma_warp_id: - _super_mma_warp( + super_mma_warp( cfg, total_tiles, bidx, @@ -2340,12 +2341,12 @@ def _kernel( sSched, lane, sK_inv_raw, - sPairwise_raw, + sIntermediate_raw, sK_decay_raw, bars, ) elif warp_idx == cfg.tcgen05_mma_warp_id: - _tcgen05_mma_warp( + tcgen05_mma_warp( cfg, total_tiles, bidx, @@ -2353,8 +2354,8 @@ def _kernel( cu_seqlens, mWorkItems, sSched, - tmem_hold, - sPairwise, + sTmem_base, + sIntermediate, sK_decay, sK_restore, sQ_decay, @@ -2362,7 +2363,7 @@ def _kernel( bars, ) elif warp_idx == cfg.epilogue_warp_id: - _epilogue_warp( + epilogue_warp( cfg, total_tiles, bidx, @@ -2374,16 +2375,16 @@ def _kernel( mO, sK_inv_raw, sO_raw, - sPairwise_raw, + sIntermediate_raw, sQ_decay_raw, - sH_raw, + sCheckpoint_raw, desc_o_base, - desc_h_base, + desc_checkpoint_base, checkpoint_every_n_tokens, bars, ) elif warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: - _compute0_warp_group( + compute0_warp_group( cfg, total_tiles, bidx, @@ -2410,7 +2411,7 @@ def _kernel( bars, ) elif warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]: - _compute1_warp_group( + compute1_warp_group( cfg, total_tiles, bidx, @@ -2419,15 +2420,15 @@ def _kernel( mWorkItems, sSched, lane, - tmem_hold, + sTmem_base, warp_idx, - mS_out, - mS_init, + mState_out, + mState_init, mO, sO_raw, sV_raw, sW_raw, - sH_raw, + sCheckpoint_raw, checkpoint_every_n_tokens, scale, bars, @@ -2439,8 +2440,7 @@ class Gdn2Cfg: """Kernel cfg (fixed BT=16 schedule constants; derived TMEM column offsets and SMEM buffer cosizes are stamped by ``build_cfg``; per-stage sizes are inlined at the use sites). Passed ``cfg``-first (a ``cutlass.Constexpr``) - into ``_host`` / ``_kernel`` and every warp body, mirroring GDN's - ``GdnCfg``.""" + into ``host`` / ``kernel`` and every warp body.""" io_dtype: Type[cutlass.Numeric] state_dtype: Type[cutlass.Numeric] @@ -2450,15 +2450,11 @@ class Gdn2Cfg: l2norm: bool safe_gate: bool gate_scale_log2: float - beta_w_sigmoid: bool q_ratio: int k_ratio: int v_ratio: int n_heads_out: int max_active_clusters: int - # split-K: tiles come from a work-item table (see common/split_k.py); - # each item computes chunks [cstart, wend) and writes only [wstart, wend) - split_k: bool = False dyn_sched: bool = False sched_stages: int = CFG.SMEM_SCHED_STAGES @@ -2477,32 +2473,35 @@ class Gdn2Cfg: cg0_group_count: int = 2 cg0_warps_per_group: int = 4 cg0_threads_per_group: int = 0 - nbar_cg0_group0_id: int = 1 # CG0 group g syncs on nbar id 1 + g + cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on nbar id 1 + g + cg0_tile_entry_barrier_id: int = 5 # CG0-wide (both groups) work-item entry sync tmem_user_threads: int = 0 - nbar_tmem_lifecycle_id: int = 3 + tmem_lifecycle_barrier_id: int = 3 num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 num_regs_other: int = CFG.NUM_REGS_OTHER - # --- SMEM / TMEM ring stage counts --- + # ---- SMEM / TMEM ring stage counts ------------------------------------------- smem_raw_stages: int = CFG.SMEM_RAW_STAGES + smem_raw_bar_stages: int = 0 # ready-ring mbar depth: raw rounded up to even (CG0 ping-pong parity) + smem_checkpoint_stages: int = 1 smem_o_stages: int = CFG.SMEM_O_STAGES smem_decay_stages: int = CFG.SMEM_DECAY_STAGES - smem_pairwise_stages: int = CFG.SMEM_PAIRWISE_STAGES + smem_intermediate_stages: int = CFG.SMEM_INTERMEDIATE_STAGES smem_state_scale_diag_stages: int = CFG.SMEM_STATE_SCALE_DIAG_STAGES qk_scale_ready_stages: int = CFG.QK_SCALE_READY_STAGES tmem_q_state_acc_stages: int = CFG.TMEM_Q_STATE_ACC_STAGES - # --- TMEM column offsets (state doubles as the final_state acc) --- - tmem_state_offset: int = 0 + # ---- TMEM column offsets (state doubles as the final_state acc) -------------- + tmem_state_acc_offset: int = 0 tmem_state_inp_offset: int = 0 tmem_q_state_acc_offset: int = 0 tmem_state_k_acc_offset: int = 0 - tmem_update_acc_offset: int = 0 - tmem_rhs_inp_offset: int = 0 - tmem_update_inp_offset: int = 0 + tmem_u_acc_offset: int = 0 + tmem_y_inp_offset: int = 0 + tmem_u_inp_offset: int = 0 - # --- SMEM buffer cosizes --- + # ---- SMEM buffer cosizes ----------------------------------------------------- q_cosize: int = 0 k_cosize: int = 0 v_cosize: int = 0 @@ -2515,7 +2514,15 @@ class Gdn2Cfg: k_restore_cosize: int = 0 state_scale_diag_cosize: int = 0 o_cosize: int = 0 - pairwise_cosize: int = 0 + + # TMA transaction bytes per stage + tma_q_bytes: int = 0 + tma_k_bytes: int = 0 + tma_gate_bytes: int = 0 + tma_beta_bytes: int = 0 + tma_v_bytes: int = 0 + tma_w_bytes: int = 0 + intermediate_cosize: int = 0 def build_cfg( @@ -2528,13 +2535,11 @@ def build_cfg( l2norm: bool, safe_gate: bool, gate_scale_log2: float, - beta_w_sigmoid: bool, q_ratio: int, k_ratio: int, v_ratio: int, n_heads_out: int, max_active_clusters: int, - split_k: bool = False, dyn_sched: bool = False, ) -> Gdn2Cfg: """Build the per-compile ``Gdn2Cfg`` (io_dtype in {Float16, BFloat16}); @@ -2550,32 +2555,30 @@ def build_cfg( l2norm=l2norm, safe_gate=safe_gate, gate_scale_log2=gate_scale_log2, - beta_w_sigmoid=beta_w_sigmoid, q_ratio=q_ratio, k_ratio=k_ratio, v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=max_active_clusters, - split_k=split_k, dyn_sched=dyn_sched, ) if enable_checkpoints: - # the 32 KB H staging buffer must fit next to the raw ring: trim the - # q/k/v/gate/beta/w TMA lookahead for H compiles - cfg.smem_raw_stages = 4 + cfg.smem_raw_stages = 3 + cfg.smem_checkpoint_stages = 2 + cfg.smem_raw_bar_stages = cfg.smem_raw_stages + (cfg.smem_raw_stages % 2) cfg.threads_per_cta = 16 * cfg.threads_per_warp cfg.cg0_threads_per_group = cfg.cg0_warps_per_group * cfg.threads_per_warp cfg.tmem_user_threads = (1 + len(cfg.compute_group_1_warp_ids)) * cfg.threads_per_warp if cfg.smem_state_scale_diag_stages != cfg.qk_scale_ready_stages: raise ValueError("diag and qk-scale ready rings must share their rolling stage") - cfg.tmem_state_inp_offset = cfg.tmem_state_offset + cfg.d_k + cfg.tmem_state_inp_offset = cfg.tmem_state_acc_offset + cfg.d_k cfg.tmem_q_state_acc_offset = cfg.tmem_state_inp_offset + (cfg.d_k // 2) cfg.tmem_state_k_acc_offset = cfg.tmem_q_state_acc_offset + cfg.tmem_q_state_acc_stages * cfg.b_t - cfg.tmem_update_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t - cfg.tmem_rhs_inp_offset = cfg.tmem_update_acc_offset + cfg.b_t - cfg.tmem_update_inp_offset = cfg.tmem_rhs_inp_offset + (cfg.b_t // 2) - assert (cfg.tmem_update_inp_offset + (cfg.b_t // 2)) <= 512 + cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t + cfg.tmem_y_inp_offset = cfg.tmem_u_acc_offset + cfg.b_t + cfg.tmem_u_inp_offset = cfg.tmem_y_inp_offset + (cfg.b_t // 2) + assert (cfg.tmem_u_inp_offset + (cfg.b_t // 2)) <= 512 cfg.q_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t cfg.k_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t @@ -2589,19 +2592,104 @@ def build_cfg( cfg.k_restore_cosize = cfg.smem_decay_stages * cfg.d_k * cfg.b_t cfg.state_scale_diag_cosize = cfg.smem_state_scale_diag_stages * (cfg.d_k // 16) * 256 cfg.o_cosize = cfg.smem_o_stages * cfg.b_t * cfg.d_v - cfg.pairwise_cosize = cfg.smem_pairwise_stages * 2 * cfg.b_t * cfg.b_t + cfg.intermediate_cosize = cfg.smem_intermediate_stages * 2 * cfg.b_t * cfg.b_t + cfg.tma_q_bytes = cfg.d_k * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_k_bytes = cfg.d_k * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_gate_bytes = cfg.d_k * cfg.b_t * 4 + cfg.tma_beta_bytes = cfg.d_k * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_v_bytes = cfg.d_v * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_w_bytes = cfg.d_v * cfg.b_t * (cfg.io_dtype.width // 8) return cfg -def get_workspace_size(B: int, HQ: int, HV: int) -> int: - """Bytes for the per-(batch, head) TMA-descriptor arrays (q, k, v, gate, - beta, w, o, h) + 128 alignment slack.""" - HO = HQ if HQ >= HV else HV - return TENSOR_MAP_QWORDS * 8 * (8 * B * HO) + 128 +TENSORMAP_DESC_ARRAYS = 8 # per-batch runtime TMA descriptors: Q, K, V, Gate, Beta, W, O, Checkpoint +TENSORMAP_STATIC_SLOTS = 0 + + +@cute.kernel +def build_all_descs_kernel( + base_q: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_v: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_gate: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_beta: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_w: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_o: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], + desc_ws: cute.Tensor, + cu_seqlens: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + beta: cute.Tensor, + w: cute.Tensor, + o: cute.Tensor, + state_checkpoints: cute.Tensor | None, + n_batch: cutlass.Int32, + q_row_stride: cutlass.Int32, + k_row_stride: cutlass.Int32, + v_row_stride: cutlass.Int32, + gate_row_stride: cutlass.Int32, + beta_row_stride: cutlass.Int32, + w_row_stride: cutlass.Int32, + o_row_stride: cutlass.Int32, + checkpoint_row_stride: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, +) -> None: + """Single-launch builder for the per-BATCH descriptor arrays (one warp + per array; warp ``i`` emits array ``i`` and release-fences its slots). + Heads are load coordinates, so only the sequence base and token extent + are patched per slot.""" + tidx, _, _ = cute.arch.thread_idx() + widx = cutlass.Int32(tidx) // cutlass.Int32(32) + arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) + sub0 = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + 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,))) + + if widx == 0: + if nvvm.elect_sync(): + emit_seq_descs(base_q, sub0, cu_seqlens, q, n_batch, q_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 1: + if nvvm.elect_sync(): + emit_seq_descs(base_k, sub1, cu_seqlens, k, n_batch, k_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 2: + if nvvm.elect_sync(): + emit_seq_descs(base_v, sub2, cu_seqlens, v, n_batch, v_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 3: + if nvvm.elect_sync(): + emit_seq_descs(base_gate, sub3, cu_seqlens, gate, n_batch, gate_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 4: + if nvvm.elect_sync(): + emit_seq_descs(base_beta, sub4, cu_seqlens, beta, n_batch, beta_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 5: + if nvvm.elect_sync(): + emit_seq_descs(base_w, sub5, cu_seqlens, w, n_batch, w_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 6: + if nvvm.elect_sync(): + emit_seq_descs(base_o, sub6, cu_seqlens, o, n_batch, o_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if cutlass.const_expr(state_checkpoints is not None): + if widx == 7: + if nvvm.elect_sync(): + emit_checkpoint_seq_descs(base_checkpoint, sub7, 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) @cute.jit -def _build_descs( +def build_descs( io_dtype: cutlass.Constexpr, b_t: cutlass.Constexpr[int], q: cute.Tensor, @@ -2611,148 +2699,101 @@ def _build_descs( beta: cute.Tensor, w: cute.Tensor, o: cute.Tensor, - h: cute.Tensor | None, + state_checkpoints: cute.Tensor | None, cu_seqlens: cute.Tensor, tensormap_workspace: cute.Tensor, - h_every_n: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, stream: cuda_driver.CUstream, ): """Build the 8 per-(batch, head) TMA-descriptor arrays (q, k, v, gate, - beta, w, o, h) into ``tensormap_workspace``. + beta, w, o, state_checkpoints) into ``tensormap_workspace``. Launched on every execute: the descriptors fold cu_seqlens contents into GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. Each descriptor folds the sequence base + head offset into 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 H descriptor is 3-D ``(dv, dk, entry)`` - over the packed ``[total_h, HO, DK, DV]`` series; its per-sequence entry + 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 offsets ((seqlen-1)//N, prefix-summed) are derived on device and its - entry extent is capped per sequence, so H store coordinates are + entry extent is capped per sequence, so checkpoint store coordinates are sequence-local.""" h_q = q.shape[1] h_k = k.shape[1] h_v = v.shape[1] - ho = gate.shape[1] + n_heads_out = gate.shape[1] batch_size = cu_seqlens.shape[0] - 1 d_k = q.shape[2] d_v = v.shape[2] bpe = io_dtype.width // 8 - granu = 128 // bpe + tma_granu_elems = 128 // bpe seqlen = q.shape[0] - def _head0(t, dim, heads): - # 2-D (dim, token) head-0 view: box (granu, b_t) matches the main - # kernel's SMEM staging byte-for-byte - return cute.make_tensor(t.iterator, cute.make_layout((dim, seqlen), stride=(1, heads * dim))) + q_headed = cute.make_tensor(q.iterator, cute.make_layout((d_k, h_q, seqlen), stride=(1, q.stride[1], q.stride[0]))) + k_headed = cute.make_tensor(k.iterator, cute.make_layout((d_k, h_k, seqlen), stride=(1, k.stride[1], k.stride[0]))) + v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v.stride[1], v.stride[0]))) + gate_headed = cute.make_tensor(gate.iterator, cute.make_layout((d_k, n_heads_out, seqlen), stride=(1, gate.stride[1], gate.stride[0]))) + beta_headed = cute.make_tensor(beta.iterator, cute.make_layout((d_k, n_heads_out, seqlen), stride=(1, beta.stride[1], beta.stride[0]))) + w_headed = cute.make_tensor(w.iterator, cute.make_layout((d_v, n_heads_out, seqlen), stride=(1, w.stride[1], w.stride[0]))) + o_headed = cute.make_tensor(o.iterator, cute.make_layout((d_v, n_heads_out, seqlen), stride=(1, o.stride[1], o.stride[0]))) swz = cuda.TensorMapSwizzle.s128b - base_q = cuda.create_tensor_map_tiled_from_view(_head0(q, d_k, h_q), box_dims=(granu, b_t), stride_order=(0, 1), swizzle=swz) - base_k = cuda.create_tensor_map_tiled_from_view(_head0(k, d_k, h_k), box_dims=(granu, b_t), stride_order=(0, 1), swizzle=swz) - base_v = cuda.create_tensor_map_tiled_from_view(_head0(v, d_v, h_v), box_dims=(granu, b_t), stride_order=(0, 1), swizzle=swz) - base_gate = cuda.create_tensor_map_tiled_from_view(_head0(gate, d_k, ho), box_dims=(32, b_t), stride_order=(0, 1), swizzle=swz) - base_beta = cuda.create_tensor_map_tiled_from_view(_head0(beta, d_k, ho), box_dims=(granu, b_t), stride_order=(0, 1), swizzle=swz) - base_w = cuda.create_tensor_map_tiled_from_view(_head0(w, d_v, ho), box_dims=(granu, b_t), stride_order=(0, 1), swizzle=swz) - base_o = cuda.create_tensor_map_tiled_from_view(_head0(o, d_v, ho), box_dims=(granu, b_t), stride_order=(0, 1), swizzle=swz) - - arr_words = (batch_size * ho) * TENSOR_MAP_QWORDS - ws_iter = tensormap_workspace.iterator - - def _sub(i): - return cute.make_tensor(ws_iter + i * arr_words, cute.make_layout((arr_words,), stride=(1,))) - - build_qkv_load_descs_kernel( - base_q, _sub(0), cu_seqlens, q, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(ho // h_q), cutlass.Int32(d_k), cutlass.Int32(h_q * d_k), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_k, _sub(1), cu_seqlens, k, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(ho // h_k), cutlass.Int32(d_k), cutlass.Int32(h_k * d_k), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_v, _sub(2), cu_seqlens, v, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(ho // h_v), cutlass.Int32(d_v), cutlass.Int32(h_v * d_v), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_gate, _sub(3), cu_seqlens, gate, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(1), cutlass.Int32(d_k), cutlass.Int32(ho * d_k), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_beta, _sub(4), cu_seqlens, beta, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(1), cutlass.Int32(d_k), cutlass.Int32(ho * d_k), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_w, _sub(5), cu_seqlens, w, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(1), cutlass.Int32(d_v), cutlass.Int32(ho * d_v), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_o, _sub(6), cu_seqlens, o, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(1), cutlass.Int32(d_v), cutlass.Int32(ho * d_v), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - if cutlass.const_expr(h is not None): - h_view = cute.make_tensor( - h.iterator, + base_q = cuda.create_tensor_map_tiled_from_view(q_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_k = cuda.create_tensor_map_tiled_from_view(k_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_v = cuda.create_tensor_map_tiled_from_view(v_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_gate = cuda.create_tensor_map_tiled_from_view(gate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_beta = cuda.create_tensor_map_tiled_from_view(beta_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_w = cuda.create_tensor_map_tiled_from_view(w_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_o = cuda.create_tensor_map_tiled_from_view(o_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + + base_checkpoint = base_o + if cutlass.const_expr(state_checkpoints is not None): + checkpoint_view = cute.make_tensor( + state_checkpoints.iterator, cute.make_layout( - (h.shape[3], h.shape[2], h.shape[0]), - stride=(h.stride[3], h.stride[2], h.stride[0]), + (state_checkpoints.shape[3], state_checkpoints.shape[2], state_checkpoints.shape[0], n_heads_out), + stride=(state_checkpoints.stride[3], state_checkpoints.stride[2], state_checkpoints.stride[0], state_checkpoints.stride[1]), ), ) - base_h = cuda.create_tensor_map_tiled_from_view(h_view, box_dims=(granu, d_k, 1), stride_order=(0, 1, 2), swizzle=swz) - build_h_descs_kernel( - base_h, - _sub(7), - cu_seqlens, - h, - cutlass.Int32(batch_size), - cutlass.Int32(ho), - cutlass.Int32(h.stride[1]), - cutlass.Int32(h.stride[0]), - h_every_n, - 2, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - - -# --------------------------------------------------------------------------- -# Torch adapter / host-side compilation -# --------------------------------------------------------------------------- - - -def _device_sm_count() -> int: - """Multiprocessor count of the current device (runtime API: auto-inits - the primary context, so this works before any other CUDA call).""" - from cuda.bindings import runtime as _rt - - err, dev = _rt.cudaGetDevice() - if int(err) != 0: - raise RuntimeError(f"cudaGetDevice failed: {err}") - err, count = _rt.cudaDeviceGetAttribute(_rt.cudaDeviceAttr.cudaDevAttrMultiProcessorCount, dev) - if int(err) != 0: - raise RuntimeError(f"cudaDeviceGetAttribute failed: {err}") - return count - - -def _data_ptr(t) -> int: - """Device address of a tensor-like (``data_ptr()`` or the CUDA array - interface).""" - fn = getattr(t, "data_ptr", None) - if fn is not None: - return fn() - return t.__cuda_array_interface__["data"][0] - - -def _cutlass_io_dtype(dtype): - name = str(dtype) - if "bfloat16" in name: - return cutlass.BFloat16 - if "float16" in name or "half" in name: - return cutlass.Float16 - raise ValueError(f"Unsupported dtype {dtype}, expected bfloat16 or float16") - - -def _cutlass_state_dtype(dtype): - name = str(dtype) - if "bfloat16" in name: - return cutlass.BFloat16 - if "float32" in name: - return cutlass.Float32 - raise ValueError(f"Unsupported state dtype {dtype}, expected float32 or bfloat16") + base_checkpoint = cuda.create_tensor_map_tiled_from_view(checkpoint_view, box_dims=(tma_granu_elems, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) + n_warps = 8 if state_checkpoints is not None else 7 + build_all_descs_kernel( + base_q, + base_k, + base_v, + base_gate, + base_beta, + base_w, + base_o, + base_checkpoint, + tensormap_workspace, + cu_seqlens, + q, + k, + v, + gate, + beta, + w, + o, + state_checkpoints, + cutlass.Int32(batch_size), + cutlass.Int32(q.stride[0]), + cutlass.Int32(k.stride[0]), + cutlass.Int32(v.stride[0]), + cutlass.Int32(gate.stride[0]), + cutlass.Int32(beta.stride[0]), + cutlass.Int32(w.stride[0]), + cutlass.Int32(o.stride[0]), + cutlass.Int32(state_checkpoints.stride[0] if state_checkpoints is not None else 0), + checkpoint_every_n, + ).launch(grid=(1, 1, 1), block=(32 * n_warps, 1, 1), stream=stream) + + +# ---- Torch adapter / host-side compilation --------------------------------------- @lru_cache(maxsize=None) -def _get_compiled_cache( +def get_compiled_cache( io_dtype_str: str, state_dtype_str: str, cu_dtype_str: str, @@ -2765,8 +2806,6 @@ def _get_compiled_cache( l2norm: bool, safe_gate: bool, gate_lower_bound: float, - beta_w_sigmoid: bool, - split_k: bool, dyn_sched: bool, ): """Return a mutable dict that lazily stores the compiled kernel.""" @@ -2782,12 +2821,10 @@ def compile( l2norm: bool, safe_gate: bool, gate_scale_log2: float, - beta_w_sigmoid: bool, q_ratio: int, k_ratio: int, v_ratio: int, n_heads_out: int, - split_k: bool = False, dyn_sched: bool = False, *, num_sm: int, @@ -2800,9 +2837,9 @@ def compile( beta_cute, w_cute, cu_seqlens_cute, - s_in_cute, + state_in_cute, o_cute, - s_out_cute, + state_out_cute, work_items_cute=None, work_count_cute=None, sched_ctr_cute=None, @@ -2821,18 +2858,16 @@ def compile( l2norm=l2norm, safe_gate=safe_gate, gate_scale_log2=gate_scale_log2, - beta_w_sigmoid=beta_w_sigmoid, q_ratio=q_ratio, k_ratio=k_ratio, v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=num_sm, - split_k=split_k, dyn_sched=dyn_sched, ) return cute.compile( - _host, + host, cfg, q_cute, k_cute, @@ -2843,9 +2878,9 @@ def compile( beta_cute, w_cute, cu_seqlens_cute, - s_in_cute, + state_in_cute, o_cute, - s_out_cute, + state_out_cute, work_items_cute, work_count_cute, sched_ctr_cute, @@ -2870,13 +2905,12 @@ def chunk_gdn2_sm100( output_state, scale: float, checkpoint_every_n_tokens: int = 0, - output_checkpoints=None, + output_state_checkpoints=None, use_qk_l2norm_in_kernel: bool = False, safe_gate: bool = False, gate_lower_bound: float = DEFAULT_GATE_LOWER_BOUND, a_log=None, dt_bias=None, - use_beta_w_sigmoid_in_kernel: bool = False, work_items=None, work_count=None, sched_ctr=None, @@ -2886,7 +2920,9 @@ def chunk_gdn2_sm100( ) -> None: """Execute the Blackwell BT=16 chunked GDN-2 prefill kernel. - All tensors must be contiguous and on the same CUDA device. + All tensors must be on the same CUDA device with a stride-1 innermost + dim; outer strides are free (padded / permuted views are read through + the TMA descriptors and dynamic layouts). Args: q: ``(total_tokens, HQ, DK)`` float16/bfloat16 @@ -2897,32 +2933,28 @@ def chunk_gdn2_sm100( ``lower_bound * sigmoid(exp(a_log) * (gate + dt_bias))``. beta: ``(total_tokens, HO, DK)`` io dtype, channel-wise erase gate w: ``(total_tokens, HO, DV)`` io dtype, channel-wise write gate - (both raw logits when ``use_beta_w_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 scale: attention scale factor (must not be 0) - checkpoint_every_n_tokens: emit an H entry every N tokens (0 = off). - H[j] is the state after ``(j + 1) * N`` tokens, STRICTLY BEFORE + 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_checkpoints: ``(total_h, HO, DK, DV)`` io-dtype (KV, v - contiguous — the GDN H layout); the per-sequence entry offsets + output_state_checkpoints: ``(total_checkpoints, HO, DK, DV)`` io-dtype (KV, v + 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 use_qk_l2norm_in_kernel: L2-normalize q/k rows inside the kernel safe_gate: interpret ``gate`` through the safe-gate transform a_log: ``(HO,)`` float32, safe-gate per-head log-amplitude (None = 0) dt_bias: ``(HO, DK)`` float32, safe-gate channel bias (None = 0) - use_beta_w_sigmoid_in_kernel: ``beta``/``w`` hold logits; sigmoid in-kernel - work_items: ``(max_items, 6)`` int32 split-K work-item table from - ``common/split_k.py``, or None for the one-tile-per-(b,h) - layout. With a table, each item computes chunks - ``[cstart, wend)`` and writes O/checkpoints only for - ``[wstart, wend)``. - work_count: ``(1,)`` int32 device-side item count (required with - work_items) + work_items: ``(max_items, 8)`` int32 work-item table from + ``common/split_k.py`` (REQUIRED; an uncut table row is the whole + (b, h) sequence). Each item computes chunks ``[cstart, wend)`` + and writes O/checkpoints only for ``[wstart, wend)``. + work_count: ``(1,)`` int32 device-side item count (REQUIRED) """ HQ = q.shape[1] HK = k.shape[1] @@ -2932,21 +2964,15 @@ def chunk_gdn2_sm100( store_final_state = output_state is not None enable_checkpoints = checkpoint_every_n_tokens > 0 if enable_checkpoints: - if output_checkpoints is None: - raise ValueError("checkpoint_every_n_tokens > 0 requires output_checkpoints") - if str(output_checkpoints.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: + if output_state_checkpoints is None: + raise ValueError("checkpoint_every_n_tokens > 0 requires output_state_checkpoints") + if str(output_state_checkpoints.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: raise ValueError( - f"output_checkpoints dtype must match the io dtype (fp32 state belongs to output_state): got {output_checkpoints.dtype} with io {q.dtype}" + f"output_state_checkpoints dtype must match the io dtype (fp32 state belongs to output_state): got {output_state_checkpoints.dtype} with io {q.dtype}" ) - split_k = work_items 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 - if split_k: - if work_count is None: - raise ValueError("work_count is required with work_items") - if enable_checkpoints and checkpoint_every_n_tokens != CFG.B_T: - raise ValueError(f"split-K checkpoints require checkpoint_every_n_tokens == {CFG.B_T}, got {checkpoint_every_n_tokens}") - elif work_count is not None: - raise ValueError("work_count must be None without work_items") if initial_state is not None: state_dtype_src = initial_state.dtype @@ -2968,11 +2994,9 @@ def chunk_gdn2_sm100( if not safe_gate: a_log = None dt_bias = None - if _data_ptr(tensormap_workspace) % 128 != 0: - raise ValueError("tensormap_workspace must be 128-byte aligned") cu_stream = cuda_driver.CUstream(int(stream)) - cache = _get_compiled_cache( + cache = get_compiled_cache( str(q.dtype), str(state_dtype_src), str(cu_seqlens.dtype), @@ -2985,48 +3009,34 @@ def chunk_gdn2_sm100( use_qk_l2norm_in_kernel, safe_gate, gate_lower_bound, - use_beta_w_sigmoid_in_kernel, - split_k, dyn_sched, ) if "compiled" not in cache: - io_dtype = _cutlass_io_dtype(q.dtype) - state_dtype = _cutlass_state_dtype(state_dtype_src) - q_cute = from_dlpack(q, assumed_align=16) - q_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - k_cute = from_dlpack(k, assumed_align=16) - k_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - v_cute = from_dlpack(v, assumed_align=16) - v_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - gate_cute = from_dlpack(gate, assumed_align=16) - gate_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + io_dtype = get_dtype(q.dtype) + state_dtype = get_dtype(state_dtype_src) + q_cute = from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=2) + k_cute = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_cute = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_cute = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2) 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 - beta_cute = from_dlpack(beta, assumed_align=4) - beta_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - w_cute = from_dlpack(w, assumed_align=16) - w_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - o_cute = from_dlpack(output, assumed_align=16) - o_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + beta_cute = from_dlpack(beta, assumed_align=4).mark_layout_dynamic(leading_dim=2) + w_cute = from_dlpack(w, assumed_align=16).mark_layout_dynamic(leading_dim=2) + o_cute = from_dlpack(output, assumed_align=16).mark_layout_dynamic(leading_dim=2) cu_seqlens_cute = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() - s_in_cute = None + state_in_cute = None if use_initial_state: - s_in_cute = from_dlpack(initial_state, assumed_align=16) - s_in_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=CFG.D_K) + state_in_cute = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) - s_out_cute = None + state_out_cute = None if store_final_state: - s_out_cute = from_dlpack(output_state, assumed_align=16) - s_out_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=CFG.D_K) + state_out_cute = from_dlpack(output_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) - work_items_cute = None - work_count_cute = None - if split_k: - work_items_cute = from_dlpack(work_items, assumed_align=4) - work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) - work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() + work_items_cute = from_dlpack(work_items, assumed_align=16) + work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() sched_ctr_cute = None if dyn_sched: @@ -3043,14 +3053,12 @@ def chunk_gdn2_sm100( use_qk_l2norm_in_kernel, safe_gate, gate_scale_log2, - use_beta_w_sigmoid_in_kernel, q_ratio, k_ratio, v_ratio, HO, - split_k, dyn_sched, - num_sm=_device_sm_count(), + num_sm=multiprocessor_count(current_device_id()), q_cute=q_cute, k_cute=k_cute, v_cute=v_cute, @@ -3060,9 +3068,9 @@ def chunk_gdn2_sm100( beta_cute=beta_cute, w_cute=w_cute, cu_seqlens_cute=cu_seqlens_cute, - s_in_cute=s_in_cute, + state_in_cute=state_in_cute, o_cute=o_cute, - s_out_cute=s_out_cute, + state_out_cute=state_out_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, sched_ctr_cute=sched_ctr_cute, @@ -3073,43 +3081,29 @@ def chunk_gdn2_sm100( ) compiled = cache["compiled"] - - # The descriptors encode cu_seqlens' CONTENTS, which no key built from the - # buffers can track. The skip this replaces asked torch's _version counter, - # so it was sound for a torch caller and silently stale for every other - # producer. Rebuilding unconditionally measures free: 131 vs 135 us of host - # time, and 157 either way once the launches are waited on. - h_for_descs = output_checkpoints if enable_checkpoints else None - if cache.get("build_descs_has_h") != (h_for_descs is not None): + state_checkpoints_for_descs = output_state_checkpoints if enable_checkpoints else None + # desc build runs every execute by contract (cu contents are data; + # buffer pointers may change) — capture-safe, single tiny launch + if cache.get("build_descs_has_state_checkpoints") != (state_checkpoints_for_descs is not None): cache.pop("build_descs", None) - cache["build_descs_has_h"] = h_for_descs is not None + cache["build_descs_has_state_checkpoints"] = state_checkpoints_for_descs is not None if "build_descs" not in cache: - io_dtype = _cutlass_io_dtype(q.dtype) - - def _bd3(t): - c = from_dlpack(t, assumed_align=16) - c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - return c - - def _bd4(t): - c = from_dlpack(t, assumed_align=16) - c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) - return c + io_dtype = get_dtype(q.dtype) cu_bd = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() ws_bd = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() cache["build_descs"] = cute.compile( - _build_descs, + build_descs, io_dtype, CFG.B_T, - _bd3(q), - _bd3(k), - _bd3(v), - _bd3(gate), - _bd3(beta), - _bd3(w), - _bd3(output), - None if h_for_descs is None else _bd4(h_for_descs), + from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(beta, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(w, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(output, assumed_align=16).mark_layout_dynamic(leading_dim=2), + None if state_checkpoints_for_descs is None else from_dlpack(state_checkpoints_for_descs, assumed_align=16).mark_layout_dynamic(leading_dim=3), cu_bd, ws_bd, cutlass.Int32(checkpoint_every_n_tokens), @@ -3124,13 +3118,12 @@ def _bd4(t): beta, w, output, - h_for_descs, + state_checkpoints_for_descs, cu_seqlens, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, ) - compiled( q, k, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_config.py b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_config.py new file mode 100644 index 000000000..2b723e05b --- /dev/null +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_config.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# This kernel is derived from cuDNN, NVIDIA Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Gated DeltaNet v2 (GDN-2) Cutlass DSL recompute (state/checkpoint-only) kernel +config (fixed compile-time constants). A fork of the BT=16 KDA schedule +extended with the per-key erase gate (beta) and per-value write gate (w); the +derived SMEM/TMEM sizes and offsets are stamped by ``build_cfg`` in +``gdn2_recompute_f16.py``. + +Target arch: Blackwell SM100 (GB200) / SM103 (GB300). +""" + +from dataclasses import dataclass +from typing import Tuple + + +@dataclass(frozen=True) +class Cfg: + # --- tile shape --- + B_T: int = 16 # chunk-inner token tile (BT=16 schedule) + D_K: int = 128 # query/key head dim + D_V: int = 128 # value head dim + + # --- warp assignments (16 warps = 512 threads) --- + COMPUTE_GROUP_0_WARP_IDS: Tuple[int, ...] = (0, 1, 2, 3, 4, 5, 6, 7) # decay/beta-operand materialize + COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (8, 9, 10, 11) # value-side TMEM (w*v - erase) + SUPER_MMA_WARP_ID: int = 12 # register-MMA KK + Neumann T_inv + TCGEN05_MMA_WARP_ID: int = 13 # tcgen05 state GEMMs + TMA_WARP_ID: int = 14 # k/v/gate/beta/w TMA loads + EPILOGUE_WARP_ID: int = 15 # checkpoint TMA store + + # --- register split --- + NUM_REGS_COMPUTE_GROUP_0: int = 160 + NUM_REGS_COMPUTE_GROUP_1: int = 136 + NUM_REGS_OTHER: int = 56 + + THREADS_PER_WARP: int = 32 + + BUFFER_ALIGN_BYTES: int = 1024 + + # --- SMEM / TMEM ring stage counts --- + SMEM_RAW_STAGES: int = 5 + SMEM_SCHED_STAGES: int = 8 + SMEM_DECAY_STAGES: int = 2 + SMEM_INTERMEDIATE_STAGES: int = 2 + SMEM_STATE_SCALE_DIAG_STAGES: int = 4 + QK_SCALE_READY_STAGES: int = 4 + + CLUSTER_SHAPE_MNK: Tuple[int, int, int] = (1, 1, 1) + + +CFG = Cfg() diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py new file mode 100644 index 000000000..399602fe8 --- /dev/null +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_recompute_f16.py @@ -0,0 +1,2629 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# This kernel is derived from cuDNN, NVIDIA Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Chunked Gated DeltaNet v2 (GDN-2) recompute (state/checkpoint-only) kernel for +Blackwell SM100/SM103 (Cutlass DSL), BT=16 tiling with per-key-channel decay + +per-key erase gate (beta) + per-value write gate (w), using direct CUTLASS +primitives. Framework-neutral entry ``chunk_gdn2_recompute_sm100``. + +A copy of the prefill (``gdn2_prefill_f16.py``) with the Q/O path removed: +it reproduces the per-chunk checkpoint series and the final state +byte-for-byte while skipping the attention output. + + S_t = S_{t-1} * diag(exp(g_t)) + U_t = W_t * V_t - (Beta_t * K_t)^T S_t + S_t += K_t (x) U_t + +vs KDA: the erase gate Beta and write gate W are per-channel tensors. Beta +is folded into the K_decay operand (feeds KK^T and state*K), the strict-lower +tile loses its per-row Beta scale, and Y becomes `W*V - state*K`; Beta/W +arrive by TMA alongside K/V. + +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: +source head = head_idx // (HO // H_x). State presence, L2NORM, SAFE_GATE, +checkpoints, and the head ratios are compile-time specializations. + +Warp assignments (16 warps = 512 threads): + warps 0-7 : compute group 0 - Gate prefix scan + decay/restore operands + warps 8-11 : compute group 1 - TMEM value side, state stores + warp 12 : super-MMA - register-MMA KK^T + Neumann inverse + warp 13 : tcgen05-MMA - the four state GEMMs + the TMEM lifecycle + warp 14 : TMA load - per-chunk input G->S loads + warp 15 : epilogue - the checkpoint TMA store + +SMEM layout: + Buffer Bytes Stages + K / V raw 20480 5 <-- SW128 TMA ring (io dtype) + Beta / W raw 2x 20480 5 <-- per-channel gates, same ring + Gate raw 40960 5 <-- fp32 prefix-scan source + dt_bias (+a_log slot) 516 1 <-- SAFE_GATE only + K_inv 8192 2 <-- token-major ldmatrix/tcgen05 B operand + K decay 8192 2 <-- tcgen05 SW128 K-box-major A/B operands + K restore 8192 2 <-- tcgen05 B operand for the state update + state-scale diag 12288 3 <-- per-k-atom decay diagonal blocks + intermediate (T_inv) 1024 2 <-- SW32 16x16 register-MMA tiles + +TMEM layout (240 of 512 columns): + Buffer Cols Purpose + state 0-127 state[DK,DV] fp32 recurrent state + state inp 128-191 packed b16 A operand view of the state + state_k_acc 192-207 state*K fp32 accumulator + u_acc 208-223 U fp32 accumulator + Y 224-231 packed b16 A operand: W*V - state*(Beta*K) + U input 232-239 packed b16 A operand: the U readback + +GEMM schedule (tcgen05-MMA warp, in issue order per chunk): + state*K -> state_k_acc + state decay (diag blocks) + U = T_inv @ Y -> u_acc + final_state += U @ K_restore + +Requires a cutlass DSL build providing `cutlass.experimental.*`; not available +in the pip nvidia-cutlass-dsl releases. +""" + +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable, NamedTuple, Optional, Type + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.experimental.cuda as cuda +import cutlass.experimental.primitives as nvvm +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + +from ..common.split_k import decode_work_item +from ..common.host import get_dtype +from cudnn.frost.buffers import current_device_id, data_ptr +from cudnn.frost.device import multiprocessor_count +from ..common.thd import TENSOR_MAP_QWORDS, emit_checkpoint_seq_descs, emit_seq_descs +from .gdn2_recompute_config import CFG +from cudnn.frost.tile_dsl.barrier import ( + advance, + MBarrier, + PipelineState, + Producer, +) +from cudnn.frost.tile_dsl.handles import GmemTileTma, MmaDesc, SmemTile, tma_slice_runtime_desc +from cudnn.frost.tile_dsl.mma import mma_step, mma_ts_step +from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_lin_S, swizzle_xor_128b +from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire +from cudnn.frost.tile_dsl.pointwise import ( + f16x2_to_f32, + fadd2, + fmul2, + ffma2, + movmatrix_16b, + mul_f16x2, + opaque_f32_zero, + fp32_to_fp16, + sub_f16x2, +) + +LOG2_E: float = 1.4426950408889634 + + +DEFAULT_GATE_LOWER_BOUND: float = -5.0 + + +# Host-side API defaults. + + +L2_NORM_EPS: float = 1.0e-12 + + +class Gdn2Bars(NamedTuple): + """Every inter-warp handoff as an ``MBarrier`` over its ring.""" + + mb_k_ready: MBarrier + mb_k_done: MBarrier + mb_v_ready: MBarrier + mb_v_done: MBarrier + mb_w_ready: MBarrier + mb_w_done: MBarrier + + mb_gate_ready: MBarrier + mb_gate_done: MBarrier + mb_beta_ready: MBarrier + mb_beta_done: MBarrier + + mb_state_k_acc_ready: MBarrier + mb_u_acc_ready: MBarrier + + mb_state_inp_ready: MBarrier + mb_y_inp_ready: MBarrier + mb_u_inp_ready: MBarrier + + mb_k_decay_inv_cg0_ready: MBarrier + mb_decay_tcgen05_done: MBarrier + mb_decay_super_done: MBarrier + mb_k_restore_acc_done: MBarrier + mb_qk_scale_ready: MBarrier + mb_state_scale_diag_done: MBarrier + mb_t_inv_ready: MBarrier + mb_t_inv_done: MBarrier + + mb_state_acc_read_done: MBarrier + mb_tmem_done: MBarrier + + mb_checkpoint_tmastg_ready: MBarrier + mb_checkpoint_tmastg_done: MBarrier + + mb_sched_ready: MBarrier + mb_sched_done: MBarrier + + +def make_gdn2_bars(cfg) -> Gdn2Bars: + """Bars factory. MUST be called from inside ``kernel`` (allocates the + mbarrier rings in SMEM ahead of the data buffers).""" + + def alloc(n): + return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=8) + + WARP = cfg.threads_per_warp + CG0_GROUP_THREADS = cfg.cg0_warps_per_group * WARP + CG1_THREADS = len(cfg.compute_group_1_warp_ids) * WARP + + return Gdn2Bars( + mb_k_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_k_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_v_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_v_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_w_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_w_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_gate_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_gate_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_beta_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_beta_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), + mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), + mb_state_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_k_decay_inv_cg0_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_decay_tcgen05_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_decay_super_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=WARP, producer=Producer.THREAD), + mb_k_restore_acc_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_qk_scale_ready=MBarrier( + alloc(cfg.qk_scale_ready_stages), + stages=cfg.qk_scale_ready_stages, + init_count=CG0_GROUP_THREADS, + producer=Producer.THREAD, + ), + mb_state_scale_diag_done=MBarrier( + alloc(cfg.smem_state_scale_diag_stages), + stages=cfg.smem_state_scale_diag_stages, + init_count=1, + producer=Producer.MMA_COMMIT, + ), + mb_t_inv_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_t_inv_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_state_acc_read_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_checkpoint_tmastg_ready=MBarrier( + alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=CG1_THREADS, producer=Producer.THREAD + ), + mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=WARP, producer=Producer.THREAD), + mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), + mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), + ) + + +# ---- Dynamic tile scheduler ------------------------------------------------------ + + +@cute.jit +def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): + """TMA-warp side: pull the next tile off the global ticket, publish it.""" + if cutlass.const_expr(cfg.dyn_sched): + bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) + if nvvm.elect_sync(): + fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sSched[sched_state.idx] = num_ctas + fetched + nvvm.bar_warp_sync(cute.arch.FULL_MASK) + next_tile = sSched[sched_state.idx] + if nvvm.elect_sync(): + bars.mb_sched_ready[sched_state.idx].arrive() + return next_tile, advance(sched_state, cfg.sched_stages) + return tile_idx + num_ctas, sched_state + + +@cute.jit +def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): + """Consumer side: read the TMA warp's published next tile.""" + if cutlass.const_expr(cfg.dyn_sched): + bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) + next_tile = sSched[sched_state.idx] + if nvvm.elect_sync(): + bars.mb_sched_done[sched_state.idx].arrive() + return next_tile, advance(sched_state, cfg.sched_stages) + return tile_idx + num_ctas, sched_state + + +@cute.jit +def diag_idx(cfg, key_dim): + """Return the SW32 index for one entry in the 8-block diagonal.""" + + block = key_dim // cutlass.Int32(16) + coord = key_dim - block * cutlass.Int32(16) + storage_col = coord ^ cutlass.Int32((cfg.b_t // 2)) + linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + storage_col + return swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) + + +@cute.jit +def tmaldg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + mSched, + sSched, + lane, + sBeta_raw, + sGate_raw, + sK_raw, + sV_raw, + sW_raw, + desc_k_base, + desc_v_base, + desc_gate_base, + desc_beta_base, + desc_w_base, + bars, +) -> None: + """TMA-LDG warp role (warp 14): persistent scheduler loop issuing the + per-chunk K/V/Beta/W/Gate G->S loads.""" + elect_one = nvvm.elect_sync() + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + sK_tma = SmemTile( + base=sK_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sV_tma = SmemTile( + base=sV_raw, + elems_per_stage=(cfg.d_v * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sBeta_tma = SmemTile( + base=sBeta_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sW_tma = SmemTile( + base=sW_raw, + elems_per_stage=(cfg.d_v * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sGate_tma = SmemTile( + base=sGate_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 32), + tma_granu_elems=32, + tma_subtile_stride_elems=(cfg.b_t * 32), + ) + raw_index = PipelineState.start(phase=1) + raw_bar_index = PipelineState.start(phase=0) + sched_state = PipelineState.start(phase=1) + 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) + head_o = head_idx + head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) + head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) + slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_k_slot = (desc_k_base + slot).tospace(cutlass.AddressSpace.generic) + desc_v_slot = (desc_v_base + slot).tospace(cutlass.AddressSpace.generic) + desc_gate_slot = (desc_gate_base + slot).tospace(cutlass.AddressSpace.generic) + desc_beta_slot = (desc_beta_base + slot).tospace(cutlass.AddressSpace.generic) + desc_w_slot = (desc_w_base + slot).tospace(cutlass.AddressSpace.generic) + if elect_one: + tma_tensormap_acquire(desc_k_slot) + tma_tensormap_acquire(desc_v_slot) + tma_tensormap_acquire(desc_gate_slot) + tma_tensormap_acquire(desc_beta_slot) + tma_tensormap_acquire(desc_w_slot) + for chunk_idx in cutlass.range(cstart, wend, 1, unroll=1): + chunk_start = chunk_idx * cfg.b_t + + # ---- K load ---------------------------------------------------------- + bars.mb_k_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_k_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_k_bytes) + k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) + tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_bar_index.idx].smem_ptr, acquire=False) + + # ---- V load ---------------------------------------------------------- + bars.mb_v_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_v_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_v_bytes) + 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_bar_index.idx].smem_ptr, acquire=False) + + # ---- Beta load ------------------------------------------------------- + bars.mb_beta_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_beta_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_beta_bytes) + 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_bar_index.idx].smem_ptr, acquire=False) + + # ---- W load ---------------------------------------------------------- + bars.mb_w_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_w_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_w_bytes) + w_slice = tma_slice_runtime_desc(desc_w_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sW_tma[raw_index.idx], w_slice, bars.mb_w_ready[raw_bar_index.idx].smem_ptr, acquire=False) + + # ---- Gate load ------------------------------------------------------- + bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_gate_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_gate_bytes) + gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_bar_index.idx].smem_ptr, acquire=False) + raw_index = advance(raw_index, cfg.smem_raw_stages) + raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) + tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def super_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_inv_raw, + sIntermediate_raw, + sK_decay_raw, + bars, +) -> None: + """Super-MMA warp role (warp 12): persistent scheduler loop computing the + register-MMA T_inv.""" + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + + # ---- ldmatrix/stmatrix lane decode ------------------------------------------- + rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) + rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) + lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) + lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) + decay_key_mask = cutlass.Int32(8) + stsm_row_coord = lane & 7 + stsm_col_coord = cutlass.Int32(0) + if (lane // 8) & 1: + stsm_row_coord = stsm_row_coord + cutlass.Int32(8) + if lane // 8 >= 2: + stsm_col_coord = cutlass.Int32(8) + stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) + global_chunk_base = cutlass.Int32(0) + 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) + num_chunks_tile = wend - cstart # processed chunks; ring bookkeeping runs on global_chunk_base + local_chunk + for local_chunk in cutlass.range(num_chunks_tile, unroll=1): + global_chunk = global_chunk_base + local_chunk + decay_stage = global_chunk % cfg.smem_decay_stages + intermediate_stage = global_chunk % cfg.smem_intermediate_stages + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (cfg.b_t * cfg.b_t) + + bars.mb_k_decay_inv_cg0_ready[decay_stage].wait((global_chunk // cfg.smem_decay_stages) % 2) + + # ---- KK = K_decay @ K_inv^T ------------------------------------------ + kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + kk_acc[accum_idx] = cutlass.Float32(0.0) + + for k_block in cutlass.range_constexpr((cfg.d_k // 16)): + # Load B operand + k_inv_col = k_block * 16 + rhs_col_offset + k_inv_segment = k_inv_col // 64 + rhs_frag = nvvm.ldmatrix( + sK_inv_ptr + + k_inv_segment * (cfg.b_t * 64) + + rhs_row_coord * 64 + + swizzle_xor_128b(rhs_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + # Load A operand + storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask + storage_slice = storage_key // 64 + kk_lhs_frag = nvvm.ldmatrix( + sK_decay_ptr + + storage_slice * (cfg.b_t * 64) + + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + + mma_step( + kk_acc, + (kk_lhs_frag[0], kk_lhs_frag[1], kk_lhs_frag[2], kk_lhs_frag[3]), + (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + + # ---- L = tril(KK, -1) fragment --------------------------------------- + row_lo = lane // 4 + row_hi = row_lo + cutlass.Int32(8) + l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + row_coord = row_lo + if cutlass.const_expr(accum_idx % 4 >= 2): + row_coord = row_hi + col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + if cutlass.const_expr(accum_idx % 2 == 1): + col_coord = col_coord + cutlass.Int32(1) + l_regs[accum_idx] = kk_acc[accum_idx] if row_coord > col_coord else cutlass.Float32(0.0) + l_a0 = fp32_to_fp16(l_regs[0], l_regs[1], dtype=cfg.io_dtype) + l_a1 = fp32_to_fp16(l_regs[2], l_regs[3], dtype=cfg.io_dtype) + l_a2 = fp32_to_fp16(l_regs[4], l_regs[5], dtype=cfg.io_dtype) + l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) + l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) + + # ---- T_inv = I - L, then three Neumann doubling rounds --------------- + tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + row_coord = row_lo + if cutlass.const_expr(accum_idx % 4 >= 2): + row_coord = row_hi + col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + if cutlass.const_expr(accum_idx % 2 == 1): + col_coord = col_coord + cutlass.Int32(1) + eye = cutlass.Float32(1.0) if row_coord == col_coord else cutlass.Float32(0.0) + tinv_acc[accum_idx] = eye - l_values[accum_idx] + + lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) + for _round in cutlass.range_constexpr(3): + # ---- Lpow = Lpow @ Lpow ------------------------------------------ + sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + sq_acc[accum_idx] = cutlass.Float32(0.0) + mma_step( + sq_acc, + (lpow_a0, lpow_a1, lpow_a2, lpow_a3), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + lpow_a0 = fp32_to_fp16(sq_acc[0], sq_acc[1], dtype=cfg.io_dtype) + lpow_a1 = fp32_to_fp16(sq_acc[2], sq_acc[3], dtype=cfg.io_dtype) + lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) + lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) + # ---- T_inv += T_inv @ Lpow --------------------------------------- + upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + upd_acc[accum_idx] = cutlass.Float32(0.0) + tinv_p0 = fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype) + tinv_p1 = fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype) + tinv_p2 = fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype) + tinv_p3 = fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype) + mma_step( + upd_acc, + (tinv_p0, tinv_p1, tinv_p2, tinv_p3), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + tinv_lo0, tinv_hi0 = f16x2_to_f32(tinv_p0, dtype=cfg.io_dtype) + tinv_lo1, tinv_hi1 = f16x2_to_f32(tinv_p1, dtype=cfg.io_dtype) + tinv_lo2, tinv_hi2 = f16x2_to_f32(tinv_p2, dtype=cfg.io_dtype) + tinv_lo3, tinv_hi3 = f16x2_to_f32(tinv_p3, dtype=cfg.io_dtype) + tinv_acc[0] = tinv_lo0 + upd_acc[0] + tinv_acc[1] = tinv_hi0 + upd_acc[1] + tinv_acc[2] = tinv_lo1 + upd_acc[2] + tinv_acc[3] = tinv_hi1 + upd_acc[3] + tinv_acc[4] = tinv_lo2 + upd_acc[4] + tinv_acc[5] = tinv_hi2 + upd_acc[5] + tinv_acc[6] = tinv_lo3 + upd_acc[6] + tinv_acc[7] = tinv_hi3 + upd_acc[7] + + bars.mb_t_inv_done[intermediate_stage].wait(((global_chunk // cfg.smem_intermediate_stages) + 1) % 2) + nvvm.stmatrix( + sIntermediate_ptr + stsm_idx, + [ + fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype), + ], + nvvm.MMALayout.ROW, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_t_inv_ready[intermediate_stage].arrive() + bars.mb_decay_super_done[decay_stage].arrive() + global_chunk_base += num_chunks_tile + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def tcgen05_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + tmem_base_holder, + sIntermediate, + sK_decay, + sK_restore, + sState_scale_diag, + bars, +) -> None: + """tcgen05-MMA warp role (warp 13): persistent scheduler loop issuing every + tcgen05 GEMM.""" + elect_one = nvvm.elect_sync() + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + nvvm.tcgen05_alloc(tmem_base_holder, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_holder.load() + state_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) + state_dsts = tuple(nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset + k * 16, cutlass.Float32) for k in range(cfg.d_k // 16)) + state_k_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_k_acc_offset, cutlass.Float32) + u_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_acc_offset, cutlass.Float32) + y_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_inp_offset, cutlass.Int8) + u_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_inp_offset, cutlass.Int8) + state_dst_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset, cutlass.Float32) + state_inp_index = PipelineState.start(phase=0) + state_read_index = PipelineState.start(phase=0) + y_inp_index = PipelineState.start(phase=0) + u_inp_index = PipelineState.start(phase=0) + qk_scale_index = PipelineState.start(phase=0) + + # ---- chunk-invariant GEMM descriptors ---------------------------------------- + bpe = cfg.io_dtype.width // 8 + idesc_acc = 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, + b_major=0, + ) + idesc_diag = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=16, + m_dim=cfg.d_v, + b_major=0, + ) + idesc_final_state = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.d_k, + m_dim=cfg.d_v, + b_major=1, + ) + bmm_state_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, + cta_group=1, + idesc=idesc_acc, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_diag_desc = MmaDesc( + M=cfg.d_v, + N=16, + K=16, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_diag, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_t_inv_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_acc, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_final_state_desc = MmaDesc( + M=cfg.d_v, + N=cfg.d_k, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + cta_group=1, + idesc=idesc_final_state, + kind=nvvm.Tcgen05MMAKind.F16, + ) + STATE_A_SEG = bmm_state_desc.sps_B * bmm_state_desc.tmem_advance_A + STATE_B_SEG = bmm_state_desc.smem_subtile_B >> 4 + global_chunk_base = cutlass.Int32(0) + 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) + num_chunks_tile = wend - cstart + for local_chunk in cutlass.range(num_chunks_tile, unroll=1): + global_chunk = global_chunk_base + local_chunk + have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_chunk > 0 + decay_stage = global_chunk % cfg.smem_decay_stages + state_scale_diag_stage = qk_scale_index.idx + intermediate_stage = global_chunk % cfg.smem_intermediate_stages + sK_decay_stage = sK_decay[decay_stage] + sK_restore_stage = sK_restore[decay_stage] + sState_scale_diag_stage = sState_scale_diag[state_scale_diag_stage] + sIntermediate_stage = sIntermediate[intermediate_stage] + + # ---- state_k_acc = state(T) @ K_decay^T ------------------------------ + bars.mb_k_decay_inv_cg0_ready[decay_stage].wait((global_chunk // cfg.smem_decay_stages) % 2) + if have_state: + bars.mb_state_inp_ready.wait(state_inp_index.phase) + state_inp_index = advance(state_inp_index, 1) + desc_k_decay = sK_decay_stage.desc() + + for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + mma_ts_step( + bmm_state_desc, + state_inp_ptr.subview(s * STATE_A_SEG), + desc_k_decay + s * STATE_B_SEG, + state_k_acc_ptr, + k, + cutlass.Boolean(s + k > 0), + ) + + if elect_one: + bars.mb_state_k_acc_ready.arrive(cta_group=1) + + if elect_one: + bars.mb_decay_tcgen05_done[decay_stage].arrive(cta_group=1) + + if cutlass.const_expr(cfg.enable_checkpoints): + if have_state: + bars.mb_state_acc_read_done.wait(state_read_index.phase) + state_read_index = advance(state_read_index, 1) + + # ---- state decay = state(T) @ diag(exp2(g_last)) (per-k-atom blocks) -- + bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) + if have_state: + desc_diag = sState_scale_diag_stage.desc() + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + mma_ts_step( + bmm_diag_desc, + state_inp_ptr.subview(k_block * bmm_diag_desc.tmem_advance_A), + desc_diag.advance_start_address(k_block * 256 * 2), + state_dsts[k_block], + 0, + cutlass.Boolean(False), + ) + + if elect_one: + bars.mb_state_scale_diag_done[state_scale_diag_stage].arrive(cta_group=1) + + # ---- u_acc = Y(T) @ T_inv^T ------------------------------------------ + bars.mb_t_inv_ready[intermediate_stage].wait((global_chunk // cfg.smem_intermediate_stages) % 2) + bars.mb_y_inp_ready.wait(y_inp_index.phase) + y_inp_index = advance(y_inp_index, 1) + desc_t_inv = sIntermediate_stage.desc() + mma_ts_step(bmm_t_inv_desc, y_inp_ptr, desc_t_inv, u_acc_ptr, 0, cutlass.Boolean(False)) + if elect_one: + bars.mb_u_acc_ready.arrive(cta_group=1) + bars.mb_t_inv_done[intermediate_stage].arrive(cta_group=1) + + # ---- final_state += U(T) @ K_restore --------------------------------- + bars.mb_u_inp_ready.wait(u_inp_index.phase) + u_inp_index = advance(u_inp_index, 1) + desc_k_restore = sK_restore_stage.desc() + + mma_ts_step(bmm_final_state_desc, u_inp_ptr, desc_k_restore, state_dst_ptr, 0, have_state) + if elect_one: + bars.mb_k_restore_acc_done[decay_stage].arrive(cta_group=1) + + qk_scale_index = advance(qk_scale_index, cfg.smem_state_scale_diag_stages) + + global_chunk_base += num_chunks_tile + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + bars.mb_tmem_done[0].wait(0) + nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) + nvvm.tcgen05_dealloc( + nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), + cutlass.Int32(512), + group=nvvm.CTAGroup.CTA_1, + ) + + +@cute.jit +def epilogue_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + sCheckpoint_raw, + desc_checkpoint_base, + checkpoint_every_n_tokens, + bars, +) -> None: + """Epilogue warp role (warp 15): the checkpoint TMA store drain.""" + elect_one = nvvm.elect_sync() + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + if cutlass.const_expr(cfg.enable_checkpoints): + sCheckpoint_tma = SmemTile( + base=sCheckpoint_raw, + elems_per_stage=(cfg.d_k * cfg.d_v), + stages=cfg.smem_checkpoint_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.d_k * 64, + ) + checkpoint_ready_index = PipelineState.start(phase=0) + 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) + head_o = head_idx + num_chunks_tile = wend - cstart + if cutlass.const_expr(cfg.enable_checkpoints): + checkpoint_slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_checkpoint_slot = (desc_checkpoint_base + checkpoint_slot).tospace(cutlass.AddressSpace.generic) + checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + checkpoint_quot = (cstart + cutlass.Int32(1)) // checkpoint_chunks + checkpoint_mod = (cstart + cutlass.Int32(1)) % checkpoint_chunks + if elect_one: + tma_tensormap_acquire(desc_checkpoint_slot) + for local_chunk in cutlass.range(num_chunks_tile, unroll=1): + chunk_idx = cstart + local_chunk + if local_chunk > 0: + # ---- checkpoint store ---------------------------------------- + do_checkpoint = checkpoint_mod == 0 + do_checkpoint = do_checkpoint and chunk_idx >= wstart + if do_checkpoint: + 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_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() + tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() + checkpoint_mod = checkpoint_mod + cutlass.Int32(1) + if checkpoint_mod == checkpoint_chunks: + checkpoint_mod = cutlass.Int32(0) + checkpoint_quot = checkpoint_quot + cutlass.Int32(1) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def gate_scale(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: + """Map raw gate to the log2-domain decay increment.""" + + if cutlass.const_expr(cfg.safe_gate): + half = cutlass.Float32(0.5) + sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half + return cfg.gate_scale_log2 * sigmoid + # Default ABI: Gate arrives in natural-log space + return raw_gate * cutlass.Float32(LOG2_E) + + +@cute.jit +def compute0_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + warp_idx, + mA_log, + mDt_bias, + sK_inv_raw, + sBeta_raw, + sGate_raw, + sK_raw, + sV_raw, + sW_raw, + sK_decay_raw, + sK_restore_raw, + sState_scale_diag_raw, + bars, +) -> None: + """CG0 warp-group role (warps 0-7): persistent scheduler loop computing the + Gate prefix scan and the decay/restore operands.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] + cg0_group_id = cg0_warp // cfg.cg0_warps_per_group + cg0_local_warp = cg0_warp % cfg.cg0_warps_per_group + prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + cg0_a_log_exp = cutlass.Float32(1.0) + cg0_dt_bias_value = cutlass.Float32(0.0) + global_chunk_base = cutlass.Int32(0) + sched_state = PipelineState.start(phase=0) + tile_idx = cutlass.Int32(bidx) + opaque_one = opaque_f32_zero() + cutlass.Float32(1.0) + 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) + head_o = head_idx + num_chunks_tile = wend - cstart + if cutlass.const_expr(cfg.safe_gate): + if num_chunks_tile > 0: + cg0_a_log_exp = cute.math.exp2(mA_log[head_o].to(cutlass.Float32) * LOG2_E, fastmath=True) + cg0_dt_bias_value = mDt_bias[head_o, prefix_dim].to(cutlass.Float32) + # non-pow2 ring cursors: one divmod per tile each, +group-stride with wrap per chunk + # tile entry: both ping-pong groups inherit each other's delivery proofs (parity-swap guard) + nvvm.barrier_cta_sync(cfg.cg0_tile_entry_barrier_id, thread_count=cfg.cg0_group_count * cfg.cg0_threads_per_group) + group_chunk_base = global_chunk_base + cutlass.Int32(cg0_group_id) + diag_stage_cursor = group_chunk_base % cutlass.Int32(cfg.smem_state_scale_diag_stages) + diag_phase = (group_chunk_base // cutlass.Int32(cfg.smem_state_scale_diag_stages)) % cutlass.Int32(2) + raw_stage_cursor = group_chunk_base % cutlass.Int32(cfg.smem_raw_stages) + raw_bar_cursor = group_chunk_base % cutlass.Int32(cfg.smem_raw_bar_stages) + raw_bar_phase = (group_chunk_base // cutlass.Int32(cfg.smem_raw_bar_stages)) % cutlass.Int32(2) + for local_chunk in cutlass.range(cg0_group_id, num_chunks_tile, cfg.cg0_group_count, unroll=1): + chunk_idx = cstart + local_chunk + global_chunk = global_chunk_base + local_chunk + chunk_start = chunk_idx * cfg.b_t + decay_stage = global_chunk % cfg.smem_decay_stages + raw_stage = raw_stage_cursor + state_scale_diag_stage = diag_stage_cursor + qk_scale_ready_stage = state_scale_diag_stage + sK_ptr = sK_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sV_ptr = sV_raw.data_ptr() + raw_stage * (cfg.d_v * cfg.b_t) + sGate_ptr = sGate_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sBeta_ptr = sBeta_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sW_ptr = sW_raw.data_ptr() + raw_stage * (cfg.d_v * cfg.b_t) + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sK_restore_ptr = sK_restore_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sState_scale_diag_ptr = sState_scale_diag_raw.data_ptr() + state_scale_diag_stage * ((cfg.d_k // 16) * 256) + + bars.mb_gate_ready[raw_bar_cursor].wait(raw_bar_phase) + + row_group_start = cg0_local_warp * (cfg.b_t // cfg.cg0_warps_per_group) + lane_row_group = lane // 8 + lane_in_row_group = lane - lane_row_group * 8 + decay_row = row_group_start + lane_row_group + decay_key_mask = cutlass.Int32(8) + + prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + + # ---- Gate prefix scan ----------------------------------------------- + f32_segment = prefix_dim // 32 + prefix_seg_base = f32_segment * (cfg.b_t * 32) + prefix_col = prefix_dim - f32_segment * 32 + g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + if cutlass.const_expr(cfg.safe_gate): + valid_rows = seqlen_b - chunk_idx * cutlass.Int32(cfg.b_t) + valid_mask = cutlass.vector.create_mask([cfg.b_t], [valid_rows]) + for row_pair in cutlass.range_constexpr(cfg.b_t // 2): + row0 = row_pair * 2 + row1 = row0 + 1 + prefix_idx0 = prefix_seg_base + swizzle_xor_128b(row0, row0 * 32 + prefix_col, elem_bytes=4) + prefix_idx1 = prefix_seg_base + swizzle_xor_128b(row1, row1 * 32 + prefix_col, elem_bytes=4) + gate0 = (sGate_ptr + prefix_idx0).load() + gate1 = (sGate_ptr + prefix_idx1).load() + gate0 = cg0_a_log_exp * (gate0 + cg0_dt_bias_value) + gate1 = cg0_a_log_exp * (gate1 + cg0_dt_bias_value) + gate0 = gate_scale( + cfg, + gate0, + ) + gate1 = gate_scale( + cfg, + gate1, + ) + gate_pair = cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32) + gate_pair = cutlass.vector.where(valid_mask[row0 : row1 + 1], gate_pair, 0.0) + g_prefix_regs[row0] = gate_pair[0] + g_prefix_regs[row1] = gate_pair[1] + else: + for row in cutlass.range_constexpr(cfg.b_t): + prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) + gate = (sGate_ptr + prefix_idx).load() + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) + if token_idx < seqlen_b: + gate = gate_scale( + cfg, + gate, + ) + else: + gate = cutlass.Float32(0.0) + g_prefix_regs[row] = gate + + prefix_acc = cutlass.Float32(0.0) + for row_pair in cutlass.range_constexpr(cfg.b_t // 2): + row0 = row_pair * 2 + row1 = row0 + 1 + gate0 = g_prefix_regs[row0] + gate1 = g_prefix_regs[row1] + prefix0, row_pair_sum = fadd2(prefix_acc, gate0, gate0, gate1) + prefix1 = prefix_acc + row_pair_sum + g_prefix_regs[row0] = prefix0 + g_prefix_regs[row1] = prefix1 + prefix_acc = prefix1 + + # ---- exp2(g): stage prefixes + final-token decay --------------------- + for row in cutlass.range_constexpr(cfg.b_t): + g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) + + exp_g_last = g_prefix_regs[cfg.b_t - 1] + for row in cutlass.range_constexpr(cfg.b_t): + prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) + (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) + + # ---- state-scale diag: stage exp2(g_last) decay blocks --------------- + bars.mb_state_scale_diag_done[state_scale_diag_stage].wait(diag_phase ^ cutlass.Int32(1)) + sState_scale_diag_ptr[diag_idx(cfg, prefix_dim)] = exp_g_last.to(cfg.io_dtype) + + nvvm.barrier_cta_sync(cfg.cg0_group_sync_barrier_base_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) + + bars.mb_k_ready[raw_bar_cursor].wait(raw_bar_phase) + bars.mb_beta_ready[raw_bar_cursor].wait(raw_bar_phase) + k_inv_pack = cutlass.Array(cutlass.Int32, 2 * 4, alignment=16) + raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + raw_beta_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + + # ---- optional K L2-norm + K_inv staging ------------------------------ + if cutlass.const_expr(cfg.l2norm): + kk0_lo = opaque_f32_zero() + kk0_hi = opaque_f32_zero() + kk1_lo = opaque_f32_zero() + kk1_hi = opaque_f32_zero() + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + raw_f16_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + raw_k_frag = (sK_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_beta_frag = (sBeta_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_k_frag_f32 = raw_k_frag.to(cutlass.Float32) + raw_beta_frag_f32 = raw_beta_frag.to(cutlass.Float32) + for dim_offset in cutlass.range_constexpr(8): + k_val = raw_k_frag_f32[dim_offset] + raw_k_regs[reg_base + dim_offset] = k_val + beta_val = raw_beta_frag_f32[dim_offset] + raw_beta_regs[reg_base + dim_offset] = beta_val + if cutlass.const_expr(cfg.l2norm): + if cutlass.const_expr(dim_offset % 2 == 0): + kk0_lo, kk0_hi = ffma2(k_val, k_val, k_val, k_val, kk0_lo, kk0_hi) + else: + kk1_lo, kk1_hi = ffma2(k_val, k_val, k_val, k_val, kk1_lo, kk1_hi) + + k_inv_norm = opaque_one + if cutlass.const_expr(cfg.l2norm): + k_sum_sq = kk0_hi + kk1_hi + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 4, 31, kind=nvvm.Shfl.BFLY)) + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 2, 31, kind=nvvm.Shfl.BFLY)) + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 1, 31, kind=nvvm.Shfl.BFLY)) + norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS) + k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) + + # ---- decay/restore operands: exp2(+-g) applied per key channel ------- + exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + exp_neg_g_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) + for f32_group in cutlass.range_constexpr(2): + f32_dim_base = dim_base + f32_group * 4 + f32_segment = f32_dim_base // 32 + f32_segment_dim = f32_dim_base - f32_segment * 32 + g_prefix_idx = f32_segment * (cfg.b_t * 32) + decay_row * 32 + swizzle_xor_128b(decay_row, f32_segment_dim, elem_bytes=4) + exp_g_frag = (sGate_ptr + g_prefix_idx).load(count=4, alignment=16) + f32_segment = f32_dim_base // 32 + f32_segment_dim = f32_dim_base - f32_segment * 32 + exp_g_last_idx = f32_segment * (cfg.b_t * 32) + (cfg.b_t - 1) * 32 + swizzle_xor_128b((cfg.b_t - 1), f32_segment_dim, elem_bytes=4) + exp_g_last_frag = (sGate_ptr + exp_g_last_idx).load(count=4, alignment=16) + half_reg_base = f32_group * 4 + f32_reg_base = reg_base + half_reg_base + exp_g_regs[f32_reg_base] = exp_g_frag[0] + exp_g_regs[f32_reg_base + 1] = exp_g_frag[1] + exp_g_regs[f32_reg_base + 2] = exp_g_frag[2] + exp_g_regs[f32_reg_base + 3] = exp_g_frag[3] + exp_neg_g_regs[half_reg_base] = cute.math.rcp(exp_g_frag[0], approx=True, ftz=True) + exp_neg_g_regs[half_reg_base + 1] = cute.math.rcp(exp_g_frag[1], approx=True, ftz=True) + exp_neg_g_regs[half_reg_base + 2] = cute.math.rcp(exp_g_frag[2], approx=True, ftz=True) + exp_neg_g_regs[half_reg_base + 3] = cute.math.rcp(exp_g_frag[3], approx=True, ftz=True) + exp_g_last_regs[f32_reg_base] = exp_g_last_frag[0] + exp_g_last_regs[f32_reg_base + 1] = exp_g_last_frag[1] + exp_g_last_regs[f32_reg_base + 2] = exp_g_last_frag[2] + exp_g_last_regs[f32_reg_base + 3] = exp_g_last_frag[3] + + # ---- K_decay + K_inv operands: K * exp2(+g) and K * exp2(-g) ----- + k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + for pair_idx in cutlass.range_constexpr(4): + dim0 = pair_idx * 2 + dim1 = dim0 + 1 + raw_reg_idx0 = reg_base + dim0 + raw_reg_idx1 = reg_base + dim1 + k_value0, k_value1 = fmul2(raw_k_regs[raw_reg_idx0], raw_k_regs[raw_reg_idx1], k_inv_norm, k_inv_norm) + k_beta0, k_beta1 = fmul2(k_value0, k_value1, raw_beta_regs[raw_reg_idx0], raw_beta_regs[raw_reg_idx1]) + k_pair = fp32_to_fp16(k_beta0, k_beta1, dtype=cfg.io_dtype) + exp_g_pair = fp32_to_fp16(exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1], dtype=cfg.io_dtype) + k_decay_pack[pair_idx] = mul_f16x2(k_pair, exp_g_pair, cfg.io_dtype) + exp_neg_pair = fp32_to_fp16(exp_neg_g_regs[dim0], exp_neg_g_regs[dim1], dtype=cfg.io_dtype) + k_norm_pair = fp32_to_fp16(k_value0, k_value1, dtype=cfg.io_dtype) + k_inv_pack[dim_half * 4 + pair_idx] = mul_f16x2(k_norm_pair, exp_neg_pair, cfg.io_dtype) + + k_inv_vec = cutlass.Vector.from_elements( + ( + k_inv_pack[dim_half * 4], + k_inv_pack[dim_half * 4 + 1], + k_inv_pack[dim_half * 4 + 2], + k_inv_pack[dim_half * 4 + 3], + ), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + k_decay_vec = cutlass.Vector.from_elements( + ( + k_decay_pack[0], + k_decay_pack[1], + k_decay_pack[2], + k_decay_pack[3], + ), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + if cutlass.const_expr(dim_half == 0): + operand_done_phase = ((global_chunk // cfg.smem_decay_stages) + 1) % 2 + bars.mb_decay_super_done[decay_stage].wait(operand_done_phase) + bars.mb_decay_tcgen05_done[decay_stage].wait(operand_done_phase) + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + k_inv_swizzled_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + (sK_inv_ptr + k_inv_swizzled_idx).store(k_inv_vec, alignment=16) + storage_key = dim_base ^ decay_key_mask + storage_slice = storage_key // 64 + decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( + decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 + ) + (sK_decay_ptr + decay_swizzled_idx).store(k_decay_vec, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_k_decay_inv_cg0_ready[decay_stage].arrive() + bars.mb_k_done[raw_stage].arrive() + bars.mb_gate_done[raw_stage].arrive() + bars.mb_beta_done[raw_stage].arrive() + + # ---- K_restore operand: K_inv * exp_g_last -------------------------- + bars.mb_k_restore_acc_done[decay_stage].wait(((global_chunk // cfg.smem_decay_stages + 1) % 2)) + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + k_restore_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + for pair_idx in cutlass.range_constexpr(4): + dim0 = pair_idx * 2 + dim1 = dim0 + 1 + exp_g_last_pair = fp32_to_fp16(exp_g_last_regs[reg_base + dim0], exp_g_last_regs[reg_base + dim1], dtype=cfg.io_dtype) + k_restore_pack[pair_idx] = mul_f16x2(k_inv_pack[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) + storage_row = decay_row ^ (cfg.b_t // 2) + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + k_restore_idx = f16_segment * (cfg.b_t * 64) + storage_row * 64 + swizzle_xor_128b(storage_row, f16_segment_dim, elem_bytes=2) + k_restore_vec = cutlass.Vector.from_elements( + ( + k_restore_pack[0], + k_restore_pack[1], + k_restore_pack[2], + k_restore_pack[3], + ), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + (sK_restore_ptr + k_restore_idx).store(k_restore_vec, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_qk_scale_ready[qk_scale_ready_stage].arrive() + diag_stage_cursor = diag_stage_cursor + cutlass.Int32(cfg.cg0_group_count) + diag_wrapped = diag_stage_cursor >= cutlass.Int32(cfg.smem_state_scale_diag_stages) + diag_stage_cursor = diag_stage_cursor - cutlass.Int32(cfg.smem_state_scale_diag_stages) if diag_wrapped else diag_stage_cursor + diag_phase = diag_phase ^ (cutlass.Int32(1) if diag_wrapped else cutlass.Int32(0)) + raw_stage_cursor = raw_stage_cursor + cutlass.Int32(cfg.cg0_group_count) + raw_wrapped = raw_stage_cursor >= cutlass.Int32(cfg.smem_raw_stages) + raw_stage_cursor = raw_stage_cursor - cutlass.Int32(cfg.smem_raw_stages) if raw_wrapped else raw_stage_cursor + raw_bar_cursor = raw_bar_cursor + cutlass.Int32(cfg.cg0_group_count) + raw_bar_wrapped = raw_bar_cursor >= cutlass.Int32(cfg.smem_raw_bar_stages) + raw_bar_cursor = raw_bar_cursor - cutlass.Int32(cfg.smem_raw_bar_stages) if raw_bar_wrapped else raw_bar_cursor + raw_bar_phase = raw_bar_phase ^ (cutlass.Int32(1) if raw_bar_wrapped else cutlass.Int32(0)) + global_chunk_base += num_chunks_tile + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def compute1_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_base_holder, + warp_idx, + mState_out, + mState_init, + sV_raw, + sW_raw, + sCheckpoint_raw, + checkpoint_every_n_tokens, + bars, +) -> None: + """CG1 warp-group role (warps 8-11): persistent scheduler loop staging the + value-side TMEM operands and storing checkpoints/final state.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) + sCheckpoint_ptr = sCheckpoint_raw.data_ptr() + checkpoint_done_index = PipelineState.start(phase=1) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_holder.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) + # ldmatrix.x4/stmatrix.x4 COL lane decode for the V/W loads + ldsm_row_coord = (lane // 16) * 8 + (lane & 7) + ldsm_col_offset = ((lane // 8) & 1) * 8 + row_id = tmem_row + tmem_subpartition * cfg.threads_per_warp + value_dim = tmem_subpartition * cfg.threads_per_warp + lane + state_k_acc_index = PipelineState.start(phase=0) + u_acc_index = PipelineState.start(phase=0) + k_restore_index = PipelineState.start(phase=0) # CG1's per-chunk mb_k_restore_acc_done wait slot + raw_index = PipelineState.start(phase=0) # raw-ring slot for the sV/sW reads + inputs_done arrives + raw_bar_index = PipelineState.start(phase=0) # even-depth ready-ring slot (decoupled from the data ring) + 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) + head_o = head_idx + num_chunks_tile = wend - cstart + + if num_chunks_tile > 0: + # ---- first chunk: seed state TMEM from mState_init ---------- + seed_from_initial_state = cstart == 0 + if cutlass.const_expr(mState_init is not None): + if seed_from_initial_state: + 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) + + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + state_block[0:32], + ) + else: + 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): + state_block[col] = cutlass.Float32(0.0) + + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + state_block[0:32], + ) + if cutlass.const_expr(mState_init is not None): + nvvm.tcgen05_wait("store") + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + + row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + state_col_id = tmem_col + cfg.tmem_state_acc_offset + + # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + packed_col_id = tmem_col + cfg.tmem_state_inp_offset + if cutlass.const_expr(mState_init is not None): + state_vecs = [] + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_block * 16, cutlass.Float32), num=16)) + + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) + for packed_col in cutlass.range_constexpr(8): + source_pair = packed_col ^ 4 + state_pack[packed_col] = fp32_to_fp16( + state_vecs[k_block][2 * source_pair], state_vecs[k_block][2 * source_pair + 1], dtype=cfg.io_dtype + ) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + k_block * 8, cutlass.Int8), + state_pack[0:8], + ) + + nvvm.tcgen05_wait("store") + bars.mb_state_inp_ready.arrive() + if cutlass.const_expr(cfg.enable_checkpoints): + bars.mb_state_acc_read_done.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 + input_col_id = tmem_col + cfg.tmem_y_inp_offset + value_dim_base = tmem_subpartition * cfg.threads_per_warp + + # ---- raw V fragments, then W, then the state*K acc readback ---------- + row_id0 = tmem_row + value_dim_base + row_id1 = row_id0 + 16 + raw_v_frag0 = nvvm.ldmatrix( + sV_ptr + + (value_dim_base + ldsm_col_offset) // 64 * (cfg.b_t * 64) + + ldsm_row_coord * 64 + + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + ldsm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + raw_v_frag1 = nvvm.ldmatrix( + sV_ptr + + (value_dim_base + 16 + ldsm_col_offset) // 64 * (cfg.b_t * 64) + + ldsm_row_coord * 64 + + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + 16 + ldsm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + bars.mb_w_ready[raw_bar_index.idx].wait(raw_bar_index.phase) + raw_w_frag0 = nvvm.ldmatrix( + sW_ptr + + (value_dim_base + ldsm_col_offset) // 64 * (cfg.b_t * 64) + + ldsm_row_coord * 64 + + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + ldsm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + raw_w_frag1 = nvvm.ldmatrix( + sW_ptr + + (value_dim_base + 16 + ldsm_col_offset) // 64 * (cfg.b_t * 64) + + ldsm_row_coord * 64 + + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + 16 + ldsm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + + if cutlass.const_expr(mState_init is not None): + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) + state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + + y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + w_pair = raw_w_frag0[raw_matrix] + wv_pair = mul_f16x2( + w_pair, + raw_v_frag0[raw_matrix], + cfg.io_dtype, + ) + if cutlass.const_expr(mState_init is not None): + state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + y_inp_pack0[reg_idx] = sub_f16x2( + wv_pair, + state_k_pair, + cfg.io_dtype, + ) + else: + y_inp_pack0[reg_idx] = wv_pair + + y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + w_pair = raw_w_frag1[raw_matrix] + wv_pair = mul_f16x2( + w_pair, + raw_v_frag1[raw_matrix], + cfg.io_dtype, + ) + if cutlass.const_expr(mState_init is not None): + state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + y_inp_pack1[reg_idx] = sub_f16x2( + wv_pair, + state_k_pair, + cfg.io_dtype, + ) + else: + y_inp_pack1[reg_idx] = wv_pair + + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) + + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) + + nvvm.tcgen05_wait("store") + if cutlass.const_expr(mState_init is not None): + state_k_acc_index = advance(state_k_acc_index, 1) + bars.mb_v_done[raw_index.idx].arrive() + bars.mb_w_done[raw_index.idx].arrive() + bars.mb_y_inp_ready.arrive() + + # ---- U repack: acc TMEM -> packed b16 TMEM -------------------------- + bars.mb_u_acc_ready.wait(u_acc_index.phase) + u_vals = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr((tmem_row + tmem_subpartition * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), + num=cfg.b_t, + ) + + u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): + source_pair = packed_col ^ 4 + token0 = source_pair * 2 + token1 = token0 + 1 + u_inp_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) + + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_u_inp_offset), cutlass.Int8), + u_inp_pack[0 : (cfg.b_t // 2)], + ) + nvvm.tcgen05_wait("store") + u_acc_index = advance(u_acc_index, 1) + bars.mb_u_inp_ready.arrive() + + bars.mb_k_restore_acc_done[k_restore_index.idx].wait(k_restore_index.phase) + k_restore_index = advance(k_restore_index, cfg.smem_decay_stages) + raw_index = advance(raw_index, cfg.smem_raw_stages) + raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) + + if cutlass.const_expr(cfg.enable_checkpoints): + cg1_checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + cg1_checkpoint_mod = (cstart + cutlass.Int32(1)) % cg1_checkpoint_chunks + for local_chunk in cutlass.range(1, num_chunks_tile, 1, unroll=1): + chunk_idx = cstart + local_chunk + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sW_ptr = sW_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + + do_checkpoint = False + if cutlass.const_expr(cfg.enable_checkpoints): + do_checkpoint = cg1_checkpoint_mod == 0 + cg1_checkpoint_mod = cg1_checkpoint_mod + cutlass.Int32(1) + cg1_checkpoint_mod = cutlass.Int32(0) if cg1_checkpoint_mod == cg1_checkpoint_chunks else cg1_checkpoint_mod + do_checkpoint = do_checkpoint and chunk_idx >= wstart + row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + state_col_id = tmem_col + cfg.tmem_state_acc_offset + + # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + state_vecs = [] + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + k_block * 16, cutlass.Float32), num=16)) + + packed_col_id = tmem_col + cfg.tmem_state_inp_offset + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) + for packed_col in cutlass.range_constexpr(8): + source_pair = packed_col ^ 4 + state_pack[packed_col] = fp32_to_fp16(state_vecs[k_block][2 * source_pair], state_vecs[k_block][2 * source_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + k_block * 8, cutlass.Int8), + state_pack[0:8], + ) + nvvm.tcgen05_wait("store") + bars.mb_state_inp_ready.arrive() + + # ---- checkpoint: post-publish f32 fragment read; the decay GEMM's ------ + if cutlass.const_expr(cfg.enable_checkpoints): + if do_checkpoint: + 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) + row16_addr = ((tmem_row + tmem_subpartition * cfg.threads_per_warp) + 16) << 16 + checkpoint_vbase = tmem_subpartition * cfg.threads_per_warp + checkpoint_swz_off0 = (checkpoint_vbase + ldsm_col_offset) // 64 * (cfg.d_k * 64) + 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, + ) + nvvm.tcgen05_wait("load") + bars.mb_state_acc_read_done.arrive() + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + else: + bars.mb_state_acc_read_done.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 + input_col_id = tmem_col + cfg.tmem_y_inp_offset + value_dim_base = tmem_subpartition * cfg.threads_per_warp + + # ---- raw V fragments, then W, then the state*K acc readback ---------- + row_id0 = tmem_row + value_dim_base + row_id1 = row_id0 + 16 + raw_v_frag0 = nvvm.ldmatrix( + sV_ptr + + (value_dim_base + ldsm_col_offset) // 64 * (cfg.b_t * 64) + + ldsm_row_coord * 64 + + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + ldsm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + raw_v_frag1 = nvvm.ldmatrix( + sV_ptr + + (value_dim_base + 16 + ldsm_col_offset) // 64 * (cfg.b_t * 64) + + ldsm_row_coord * 64 + + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + 16 + ldsm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + bars.mb_w_ready[raw_bar_index.idx].wait(raw_bar_index.phase) + raw_w_frag0 = nvvm.ldmatrix( + sW_ptr + + (value_dim_base + ldsm_col_offset) // 64 * (cfg.b_t * 64) + + ldsm_row_coord * 64 + + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + ldsm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + raw_w_frag1 = nvvm.ldmatrix( + sW_ptr + + (value_dim_base + 16 + ldsm_col_offset) // 64 * (cfg.b_t * 64) + + ldsm_row_coord * 64 + + swizzle_xor_128b(ldsm_row_coord, (value_dim_base + 16 + ldsm_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) + state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + + y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + w_pair = raw_w_frag0[raw_matrix] + wv_pair = mul_f16x2( + w_pair, + raw_v_frag0[raw_matrix], + cfg.io_dtype, + ) + y_inp_pack0[reg_idx] = sub_f16x2( + wv_pair, + state_k_pair, + cfg.io_dtype, + ) + + y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + w_pair = raw_w_frag1[raw_matrix] + wv_pair = mul_f16x2( + w_pair, + raw_v_frag1[raw_matrix], + cfg.io_dtype, + ) + y_inp_pack1[reg_idx] = sub_f16x2( + wv_pair, + state_k_pair, + cfg.io_dtype, + ) + + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) + + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) + + nvvm.tcgen05_wait("store") + state_k_acc_index = advance(state_k_acc_index, 1) + bars.mb_v_done[raw_index.idx].arrive() + bars.mb_w_done[raw_index.idx].arrive() + bars.mb_y_inp_ready.arrive() + + # ---- U repack: acc TMEM -> packed b16 TMEM -------------------------- + bars.mb_u_acc_ready.wait(u_acc_index.phase) + u_vals = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr((tmem_row + tmem_subpartition * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_u_acc_offset), cutlass.Float32), + num=cfg.b_t, + ) + + u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): + source_pair = packed_col ^ 4 + token0 = source_pair * 2 + token1 = token0 + 1 + u_inp_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) + + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_u_inp_offset), cutlass.Int8), + u_inp_pack[0 : (cfg.b_t // 2)], + ) + nvvm.tcgen05_wait("store") + u_acc_index = advance(u_acc_index, 1) + bars.mb_u_inp_ready.arrive() + + bars.mb_k_restore_acc_done[k_restore_index.idx].wait(k_restore_index.phase) + k_restore_index = advance(k_restore_index, cfg.smem_decay_stages) + raw_index = advance(raw_index, cfg.smem_raw_stages) + raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) + + owns_final = wend == num_chunks_b + + # ---- final-state drain: state acc TMEM -> GMEM --------------------------- + if cutlass.const_expr(mState_out is not None): + if seqlen_b > 0: + if owns_final: + for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): + state_vec = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + 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) + 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] + else: + mState_out[batch_idx, head_o, key_dim, value_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() + + +@cute.jit +def host( + cfg: cutlass.Constexpr, + k: cute.Tensor, + v: cute.Tensor, + raw_gate: cute.Tensor, + a_log: cute.Tensor | None, + dt_bias: cute.Tensor | None, + beta: cute.Tensor, + w: cute.Tensor, + cu_seqlens: cute.Tensor, + initial_state: cute.Tensor | None, + final_state: cute.Tensor | None, + work_items: cute.Tensor | None, + work_count: cute.Tensor | None, + sched_ctr: cute.Tensor | None, + tensormap_workspace: cute.Tensor, + checkpoint_every_n_tokens: cutlass.Int32, + stream, +) -> None: + num_sequences = cu_seqlens.shape[0] - 1 + grid_shape = (cfg.max_active_clusters, 1, 1) + kernel( + cfg, + tensormap_workspace, + cutlass.Int32(num_sequences), + k, + v, + raw_gate, + a_log, + dt_bias, + beta, + w, + cu_seqlens, + initial_state, + final_state, + work_items, + work_count, + sched_ctr, + checkpoint_every_n_tokens, + ).launch( + grid=grid_shape, + block=(cfg.threads_per_cta, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + +@cute.kernel +def kernel( + cfg: cutlass.Constexpr, + tensormap_workspace: cute.Tensor, + n_desc: cutlass.Int32, + mK: cute.Tensor, + mV: cute.Tensor, + mGate: cute.Tensor, + mA_log: cute.Tensor | None, + mDt_bias: cute.Tensor | None, + mBeta: cute.Tensor, + mW: cute.Tensor, + cu_seqlens: cute.Tensor, + mState_init: cute.Tensor | None, + mState_out: cute.Tensor | None, + mWorkItems: cute.Tensor, + mCount: cute.Tensor, + mSched: cute.Tensor | None, + checkpoint_every_n_tokens: cutlass.Int32, +) -> None: + """BT=16 GDN-2 recompute device kernel (persistent); grid + `(min(tiles, SM count), 1, 1)`.""" + + tidx, _, _ = cute.arch.thread_idx() + bidx = cute.arch.block_idx()[0] + num_ctas = cute.arch.grid_dim()[0] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % cfg.threads_per_warp + + total_tiles = mCount[0] + if cutlass.const_expr(cfg.dyn_sched): + assert mSched is not None and mSched.element_type == cutlass.Int32 + assert mK.element_type == cfg.io_dtype and mV.element_type == cfg.io_dtype + assert mGate.element_type == cutlass.Float32 + assert mBeta.element_type == cfg.io_dtype and mW.element_type == cfg.io_dtype, "channel-wise beta/w must match the io dtype" + assert cu_seqlens.element_type in (cutlass.Int32, cutlass.Int64) + if cutlass.const_expr(cfg.use_initial_state): + assert mState_init is not None and mState_init.element_type in (cutlass.BFloat16, cutlass.Float32) + else: + assert mState_init is None, "mState_init must be None if use_initial_state is False" + if cutlass.const_expr(cfg.store_final_state): + assert mState_out is not None and mState_out.element_type in (cutlass.BFloat16, cutlass.Float32) + else: + assert mState_out is None, "mState_out must be None if store_final_state is False" + if cutlass.const_expr(mState_init is not None and mState_out is not None): + assert mState_init.element_type == mState_out.element_type + # per-BATCH TMA-descriptor arrays (heads are load coordinates): [K, V, Gate, Beta, W, Checkpoint] + desc_base_words = tensormap_workspace.iterator.raw_ptr() + arr_words = n_desc * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_k_base = desc_base_words + desc_v_base = desc_base_words + arr_words + desc_gate_base = desc_base_words + cutlass.Int32(2) * arr_words + desc_beta_base = desc_base_words + cutlass.Int32(3) * arr_words + desc_w_base = desc_base_words + cutlass.Int32(4) * arr_words + desc_checkpoint_base = desc_base_words + cutlass.Int32(5) * arr_words + + # Buffers are declaration-ordered and intentionally non-aliased. + SMEM = cutlass.AddressSpace.smem + bars = make_gdn2_bars(cfg) + tmem_base_holder = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) + sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + sK_decay_raw = cutlass.Array(cfg.io_dtype, cfg.k_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_restore_raw = cutlass.Array(cfg.io_dtype, cfg.k_restore_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sIntermediate_raw = cutlass.Array(cfg.io_dtype, cfg.intermediate_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_raw = cutlass.Array(mK.element_type, cfg.k_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sV_raw = cutlass.Array(mV.element_type, cfg.v_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sGate_raw = cutlass.Array(cutlass.Float32, cfg.gate_cosize, space=SMEM, alignment=1024) + sState_scale_diag_raw = cutlass.Array(cfg.io_dtype, cfg.state_scale_diag_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_inv_raw = cutlass.Array(cfg.io_dtype, cfg.k_inv_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sBeta_raw = cutlass.Array(mBeta.element_type, cfg.beta_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sW_raw = cutlass.Array(mW.element_type, cfg.w_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sCheckpoint_raw = ( + cutlass.Array(cfg.io_dtype, cfg.smem_checkpoint_stages * cfg.d_k * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) + if cutlass.const_expr(cfg.enable_checkpoints) + else sV_raw + ) + sK_decay = SmemTile( + base=sK_decay_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_decay_stages, + leading_byte_offset=16, + stride_byte_offset=1024, + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) + sK_restore = SmemTile( + base=sK_restore_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_decay_stages, + leading_byte_offset=(cfg.b_t * (cfg.d_v // 2) * 2), + stride_byte_offset=(8 * (cfg.d_v // 2) * 2), + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) + sState_scale_diag = SmemTile( + base=sState_scale_diag_raw, + elems_per_stage=((cfg.d_k // 16) * 256), + stages=cfg.smem_state_scale_diag_stages, + leading_byte_offset=16, + stride_byte_offset=(8 * 16 * 2), + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, + ) + sIntermediate = SmemTile( + base=sIntermediate_raw, + elems_per_stage=(cfg.b_t * cfg.b_t), + stages=cfg.smem_intermediate_stages, + leading_byte_offset=16, + stride_byte_offset=(8 * cfg.b_t * 2), + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, + ) + + elect_one = nvvm.elect_sync() + if warp_idx == cfg.tma_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_raw_bar_stages): + bars.mb_k_ready[stage].init() + bars.mb_gate_ready[stage].init() + bars.mb_beta_ready[stage].init() + bars.mb_v_ready[stage].init() + bars.mb_w_ready[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_raw_stages): + bars.mb_k_done[stage].init() + bars.mb_gate_done[stage].init() + bars.mb_beta_done[stage].init() + bars.mb_v_done[stage].init() + bars.mb_w_done[stage].init() + elif warp_idx == cfg.tcgen05_mma_warp_id: + if elect_one: + bars.mb_state_k_acc_ready.init() + bars.mb_u_acc_ready.init() + bars.mb_state_inp_ready.init() + for stage in cutlass.range_constexpr(cfg.smem_state_scale_diag_stages): + bars.mb_state_scale_diag_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_decay_stages): + bars.mb_decay_tcgen05_done[stage].init() + bars.mb_decay_super_done[stage].init() + bars.mb_k_restore_acc_done[stage].init() + bars.mb_y_inp_ready.init() + bars.mb_u_inp_ready.init() + bars.mb_tmem_done[0].init() + elif warp_idx == cfg.super_mma_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_intermediate_stages): + bars.mb_t_inv_ready[stage].init() + bars.mb_t_inv_done[stage].init() + for stage in cutlass.range_constexpr(cfg.qk_scale_ready_stages): + bars.mb_qk_scale_ready[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_decay_stages): + bars.mb_k_decay_inv_cg0_ready[stage].init() + elif warp_idx == cfg.epilogue_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.sched_stages): + bars.mb_sched_ready[stage].init() + bars.mb_sched_done[stage].init() + if cutlass.const_expr(cfg.enable_checkpoints): + for stage in cutlass.range_constexpr(cfg.smem_checkpoint_stages): + bars.mb_checkpoint_tmastg_ready[stage].init() + bars.mb_checkpoint_tmastg_done[stage].init() + bars.mb_state_acc_read_done.init() + diag_zero = cfg.io_dtype(0.0) + for diag_idx in cutlass.range(tidx, cfg.state_scale_diag_cosize, cfg.threads_per_cta, unroll=1): + sState_scale_diag_raw[diag_idx] = diag_zero + nvvm.fence_mbarrier_init() + nvvm.barrier_cta_sync(0, thread_count=cfg.threads_per_cta) + if warp_idx == cfg.tma_warp_id: + tmaldg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + mSched, + sSched, + lane, + sBeta_raw, + sGate_raw, + sK_raw, + sV_raw, + sW_raw, + desc_k_base, + desc_v_base, + desc_gate_base, + desc_beta_base, + desc_w_base, + bars, + ) + elif warp_idx == cfg.super_mma_warp_id: + super_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_inv_raw, + sIntermediate_raw, + sK_decay_raw, + bars, + ) + elif warp_idx == cfg.tcgen05_mma_warp_id: + tcgen05_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + tmem_base_holder, + sIntermediate, + sK_decay, + sK_restore, + sState_scale_diag, + bars, + ) + elif warp_idx == cfg.epilogue_warp_id: + epilogue_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + sCheckpoint_raw, + desc_checkpoint_base, + checkpoint_every_n_tokens, + bars, + ) + elif warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: + compute0_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + warp_idx, + mA_log, + mDt_bias, + sK_inv_raw, + sBeta_raw, + sGate_raw, + sK_raw, + sV_raw, + sW_raw, + sK_decay_raw, + sK_restore_raw, + sState_scale_diag_raw, + bars, + ) + elif warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]: + compute1_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_base_holder, + warp_idx, + mState_out, + mState_init, + sV_raw, + sW_raw, + sCheckpoint_raw, + checkpoint_every_n_tokens, + bars, + ) + + +@dataclass +class Gdn2RecomputeCfg: + """Kernel cfg (fixed BT=16 schedule constants; derived TMEM column offsets + and SMEM buffer cosizes are stamped by ``build_cfg``; per-stage sizes are + inlined at the use sites). Passed ``cfg``-first (a ``cutlass.Constexpr``) + into ``host`` / ``kernel`` and every warp body.""" + + io_dtype: Type[cutlass.Numeric] + state_dtype: Type[cutlass.Numeric] + use_initial_state: bool + store_final_state: bool + enable_checkpoints: bool + l2norm: bool + safe_gate: bool + gate_scale_log2: float + k_ratio: int + v_ratio: int + n_heads_out: int + max_active_clusters: int + dyn_sched: bool = False + sched_stages: int = CFG.SMEM_SCHED_STAGES + + compute_group_0_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_0_WARP_IDS + compute_group_1_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS + super_mma_warp_id: int = CFG.SUPER_MMA_WARP_ID + tcgen05_mma_warp_id: int = CFG.TCGEN05_MMA_WARP_ID + tma_warp_id: int = CFG.TMA_WARP_ID + epilogue_warp_id: int = CFG.EPILOGUE_WARP_ID + b_t: int = CFG.B_T + d_k: int = CFG.D_K + d_v: int = CFG.D_V + threads_per_warp: int = CFG.THREADS_PER_WARP + buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES + threads_per_cta: int = 0 + cg0_group_count: int = 2 + cg0_warps_per_group: int = 4 + cg0_threads_per_group: int = 0 + cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on nbar id 1 + g + cg0_tile_entry_barrier_id: int = 5 # CG0-wide (both groups) work-item entry sync + tmem_user_threads: int = 0 + tmem_lifecycle_barrier_id: int = 3 + num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 + num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 + num_regs_other: int = CFG.NUM_REGS_OTHER + + # ---- SMEM / TMEM ring stage counts ------------------------------------------- + smem_raw_stages: int = CFG.SMEM_RAW_STAGES + smem_raw_bar_stages: int = 0 # ready-ring mbar depth: raw rounded up to even (CG0 ping-pong parity) + smem_checkpoint_stages: int = 1 + smem_decay_stages: int = CFG.SMEM_DECAY_STAGES + smem_intermediate_stages: int = CFG.SMEM_INTERMEDIATE_STAGES + smem_state_scale_diag_stages: int = CFG.SMEM_STATE_SCALE_DIAG_STAGES + qk_scale_ready_stages: int = CFG.QK_SCALE_READY_STAGES + + # ---- TMEM column offsets (state doubles as the final_state acc) -------------- + tmem_state_acc_offset: int = 0 + tmem_state_inp_offset: int = 0 + tmem_state_k_acc_offset: int = 0 + tmem_u_acc_offset: int = 0 + tmem_y_inp_offset: int = 0 + tmem_u_inp_offset: int = 0 + + # ---- SMEM buffer cosizes ----------------------------------------------------- + k_cosize: int = 0 + v_cosize: int = 0 + gate_cosize: int = 0 + beta_cosize: int = 0 + w_cosize: int = 0 + k_inv_cosize: int = 0 + k_decay_cosize: int = 0 + k_restore_cosize: int = 0 + state_scale_diag_cosize: int = 0 + + # TMA transaction bytes per stage + tma_k_bytes: int = 0 + tma_gate_bytes: int = 0 + tma_beta_bytes: int = 0 + tma_v_bytes: int = 0 + tma_w_bytes: int = 0 + intermediate_cosize: int = 0 + + +def build_cfg( + io_dtype: Type[cutlass.Numeric], + state_dtype: Type[cutlass.Numeric], + *, + use_initial_state: bool, + store_final_state: bool, + enable_checkpoints: bool, + l2norm: bool, + safe_gate: bool, + gate_scale_log2: float, + k_ratio: int, + v_ratio: int, + n_heads_out: int, + max_active_clusters: int, + dyn_sched: bool = False, +) -> Gdn2RecomputeCfg: + """Build the per-compile ``Gdn2RecomputeCfg`` (io_dtype in {Float16, BFloat16}); + fills the derived TMEM column offsets and SMEM buffer cosizes.""" + if io_dtype not in (cutlass.Float16, cutlass.BFloat16): + raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") + cfg = Gdn2RecomputeCfg( + io_dtype=io_dtype, + state_dtype=state_dtype, + use_initial_state=use_initial_state, + store_final_state=store_final_state, + enable_checkpoints=enable_checkpoints, + l2norm=l2norm, + safe_gate=safe_gate, + gate_scale_log2=gate_scale_log2, + k_ratio=k_ratio, + v_ratio=v_ratio, + n_heads_out=n_heads_out, + max_active_clusters=max_active_clusters, + dyn_sched=dyn_sched, + ) + if enable_checkpoints: + cfg.smem_raw_stages = 4 + cfg.smem_checkpoint_stages = 2 + cfg.smem_raw_bar_stages = cfg.smem_raw_stages + (cfg.smem_raw_stages % 2) + cfg.threads_per_cta = 16 * cfg.threads_per_warp + cfg.cg0_threads_per_group = cfg.cg0_warps_per_group * cfg.threads_per_warp + cfg.tmem_user_threads = (1 + len(cfg.compute_group_1_warp_ids)) * cfg.threads_per_warp + if cfg.smem_state_scale_diag_stages != cfg.qk_scale_ready_stages: + raise ValueError("diag and qk-scale ready rings must share their rolling stage") + + cfg.tmem_state_inp_offset = cfg.tmem_state_acc_offset + cfg.d_k + cfg.tmem_state_k_acc_offset = cfg.tmem_state_inp_offset + (cfg.d_k // 2) + cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t + cfg.tmem_y_inp_offset = cfg.tmem_u_acc_offset + cfg.b_t + cfg.tmem_u_inp_offset = cfg.tmem_y_inp_offset + (cfg.b_t // 2) + assert (cfg.tmem_u_inp_offset + (cfg.b_t // 2)) <= 512 + + cfg.k_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t + cfg.v_cosize = cfg.smem_raw_stages * cfg.d_v * cfg.b_t + cfg.gate_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t + cfg.beta_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t + cfg.w_cosize = cfg.smem_raw_stages * cfg.d_v * cfg.b_t + cfg.k_inv_cosize = cfg.smem_decay_stages * cfg.b_t * cfg.d_k + cfg.k_decay_cosize = cfg.smem_decay_stages * cfg.d_k * cfg.b_t + cfg.k_restore_cosize = cfg.smem_decay_stages * cfg.d_k * cfg.b_t + cfg.state_scale_diag_cosize = cfg.smem_state_scale_diag_stages * (cfg.d_k // 16) * 256 + cfg.intermediate_cosize = cfg.smem_intermediate_stages * cfg.b_t * cfg.b_t + cfg.tma_k_bytes = cfg.d_k * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_gate_bytes = cfg.d_k * cfg.b_t * 4 + cfg.tma_beta_bytes = cfg.d_k * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_v_bytes = cfg.d_v * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_w_bytes = cfg.d_v * cfg.b_t * (cfg.io_dtype.width // 8) + return cfg + + +TENSORMAP_DESC_ARRAYS = 6 # per-batch runtime TMA descriptors: K, V, Gate, Beta, W, Checkpoint +TENSORMAP_STATIC_SLOTS = 0 + + +@cute.kernel +def build_all_descs_kernel( + base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_v: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_gate: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_beta: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_w: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], + tensormap_workspace: cute.Tensor, + cu_seqlens: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + beta: cute.Tensor, + w: cute.Tensor, + state_checkpoints: cute.Tensor | None, + n_batch: cutlass.Int32, + k_row_stride: cutlass.Int32, + v_row_stride: cutlass.Int32, + gate_row_stride: cutlass.Int32, + beta_row_stride: cutlass.Int32, + w_row_stride: cutlass.Int32, + checkpoint_row_stride: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, +) -> None: + """Single-launch builder for the per-batch TMA-descriptor arrays.""" + tidx, _, _ = cute.arch.thread_idx() + widx = cutlass.Int32(tidx) // cutlass.Int32(32) + arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) + sub0 = cute.make_tensor(tensormap_workspace.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(tensormap_workspace.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(tensormap_workspace.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(tensormap_workspace.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(tensormap_workspace.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub5 = cute.make_tensor(tensormap_workspace.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + + if widx == 0: + if nvvm.elect_sync(): + emit_seq_descs(base_k, sub0, cu_seqlens, k, n_batch, k_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 1: + if nvvm.elect_sync(): + emit_seq_descs(base_v, sub1, cu_seqlens, v, n_batch, v_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 2: + if nvvm.elect_sync(): + emit_seq_descs(base_gate, sub2, cu_seqlens, gate, n_batch, gate_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 3: + if nvvm.elect_sync(): + emit_seq_descs(base_beta, sub3, cu_seqlens, beta, n_batch, beta_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 4: + if nvvm.elect_sync(): + emit_seq_descs(base_w, sub4, cu_seqlens, w, n_batch, w_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if cutlass.const_expr(state_checkpoints is not None): + if widx == 5: + if nvvm.elect_sync(): + emit_checkpoint_seq_descs(base_checkpoint, sub5, 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) + + +@cute.jit +def build_descs( + io_dtype: cutlass.Constexpr, + b_t: cutlass.Constexpr[int], + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + beta: cute.Tensor, + w: cute.Tensor, + state_checkpoints: cute.Tensor | None, + cu_seqlens: cute.Tensor, + tensormap_workspace: cute.Tensor, + checkpoint_every_n: cutlass.Int32, + stream: cuda_driver.CUstream, +): + """Build the 6 per-batch TMA-descriptor arrays (k, v, gate, beta, w, state_checkpoints) + into ``tensormap_workspace``.""" + h_k = k.shape[1] + h_v = v.shape[1] + ho = gate.shape[1] + batch_size = cu_seqlens.shape[0] - 1 + d_k = k.shape[2] + d_v = v.shape[2] + bpe = io_dtype.width // 8 + tma_box_elems = 128 // bpe + seqlen = k.shape[0] + + k_headed = cute.make_tensor(k.iterator, cute.make_layout((d_k, h_k, seqlen), stride=(1, k.stride[1], k.stride[0]))) + v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v.stride[1], v.stride[0]))) + gate_headed = cute.make_tensor(gate.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, gate.stride[1], gate.stride[0]))) + beta_headed = cute.make_tensor(beta.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, beta.stride[1], beta.stride[0]))) + w_headed = cute.make_tensor(w.iterator, cute.make_layout((d_v, ho, seqlen), stride=(1, w.stride[1], w.stride[0]))) + + swz = cuda.TensorMapSwizzle.s128b + base_k = cuda.create_tensor_map_tiled_from_view(k_headed, box_dims=(tma_box_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_v = cuda.create_tensor_map_tiled_from_view(v_headed, box_dims=(tma_box_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_gate = cuda.create_tensor_map_tiled_from_view(gate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_beta = cuda.create_tensor_map_tiled_from_view(beta_headed, box_dims=(tma_box_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_w = cuda.create_tensor_map_tiled_from_view(w_headed, box_dims=(tma_box_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + + base_checkpoint = base_k + if cutlass.const_expr(state_checkpoints is not None): + checkpoint_view = cute.make_tensor( + state_checkpoints.iterator, + cute.make_layout( + (state_checkpoints.shape[3], state_checkpoints.shape[2], state_checkpoints.shape[0], ho), + stride=(state_checkpoints.stride[3], state_checkpoints.stride[2], state_checkpoints.stride[0], state_checkpoints.stride[1]), + ), + ) + base_checkpoint = cuda.create_tensor_map_tiled_from_view(checkpoint_view, box_dims=(tma_box_elems, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) + n_warps = 6 if state_checkpoints is not None else 5 + build_all_descs_kernel( + base_k, + base_v, + base_gate, + base_beta, + base_w, + base_checkpoint, + tensormap_workspace, + cu_seqlens, + k, + v, + gate, + beta, + w, + state_checkpoints, + cutlass.Int32(batch_size), + cutlass.Int32(k.stride[0]), + cutlass.Int32(v.stride[0]), + cutlass.Int32(gate.stride[0]), + cutlass.Int32(beta.stride[0]), + cutlass.Int32(w.stride[0]), + cutlass.Int32(state_checkpoints.stride[0] if state_checkpoints is not None else 0), + checkpoint_every_n, + ).launch(grid=(1, 1, 1), block=(32 * n_warps, 1, 1), stream=stream) + + +# ---- Torch adapter / host-side compilation --------------------------------------- + + +@lru_cache(maxsize=None) +def get_compiled_cache( + io_dtype_str: str, + state_dtype_str: str, + cu_dtype_str: str, + HO: int, + HK: int, + HV: int, + use_initial_state: bool, + store_final_state: bool, + enable_checkpoints: bool, + l2norm: bool, + safe_gate: bool, + gate_lower_bound: float, + dyn_sched: bool, +): + """Return a mutable dict that lazily stores the compiled kernel.""" + return {} + + +def compile( + io_dtype, + state_dtype, + use_initial_state: bool, + store_final_state: bool, + enable_checkpoints: bool, + l2norm: bool, + safe_gate: bool, + gate_scale_log2: float, + k_ratio: int, + v_ratio: int, + n_heads_out: int, + dyn_sched: bool = False, + *, + num_sm: int, + k_cute, + v_cute, + gate_cute, + a_log_cute, + dt_bias_cute, + beta_cute, + w_cute, + cu_seqlens_cute, + state_in_cute, + state_out_cute, + work_items_cute=None, + work_count_cute=None, + sched_ctr_cute=None, + tensormap_ws_cute, + checkpoint_every_n_tokens, + stream, +): + """JIT-compile the chunked GDN-2 recompute kernel for one static config.""" + cfg = build_cfg( + io_dtype, + state_dtype, + use_initial_state=use_initial_state, + store_final_state=store_final_state, + enable_checkpoints=enable_checkpoints, + l2norm=l2norm, + safe_gate=safe_gate, + gate_scale_log2=gate_scale_log2, + k_ratio=k_ratio, + v_ratio=v_ratio, + n_heads_out=n_heads_out, + max_active_clusters=num_sm, + dyn_sched=dyn_sched, + ) + + return cute.compile( + host, + cfg, + k_cute, + v_cute, + gate_cute, + a_log_cute, + dt_bias_cute, + beta_cute, + w_cute, + cu_seqlens_cute, + state_in_cute, + state_out_cute, + work_items_cute, + work_count_cute, + sched_ctr_cute, + tensormap_ws_cute, + checkpoint_every_n_tokens, + stream, + options="--enable-tvm-ffi --opt-level 2", + ) + + +def chunk_gdn2_recompute_sm100( + k, + v, + gate, + beta, + w, + cu_seqlens, + initial_state, + output_state, + checkpoint_every_n_tokens: int = 0, + output_state_checkpoints=None, + use_qk_l2norm_in_kernel: bool = False, + safe_gate: bool = False, + gate_lower_bound: float = DEFAULT_GATE_LOWER_BOUND, + a_log=None, + dt_bias=None, + work_items=None, + work_count=None, + sched_ctr=None, + *, + tensormap_workspace, + stream, +) -> None: + """Execute the Blackwell BT=16 chunked GDN-2 recompute (state/checkpoint-only) kernel. + + All tensors must be on the same CUDA device with a stride-1 innermost + dim; outer strides are free (padded / permuted views are read through + the TMA descriptors and dynamic layouts). + + Args: + k: ``(total_tokens, HK, DK)`` float16/bfloat16 + v: ``(total_tokens, HV, DV)`` float16/bfloat16 + gate: ``(total_tokens, HO, DK)`` float32. Natural-log decay unless + ``safe_gate``, which applies the safe-gate transform + ``lower_bound * sigmoid(exp(a_log) * (gate + dt_bias))``. + beta: ``(total_tokens, HO, DK)`` io dtype, channel-wise erase gate + 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 + 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 + 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 + use_qk_l2norm_in_kernel: L2-normalize K rows inside the kernel + safe_gate: interpret ``gate`` through the safe-gate transform + a_log: ``(HO,)`` float32, safe-gate per-head log-amplitude (None = 0) + dt_bias: ``(HO, DK)`` float32, safe-gate channel bias (None = 0) + work_items: ``(max_items, 8)`` int32 work-item table from + ``common/split_k.py`` (REQUIRED; an uncut table row is the whole + (b, h) sequence). Each item computes chunks ``[cstart, wend)`` + and writes checkpoints only for ``[wstart, wend)``. + work_count: ``(1,)`` int32 device-side item count (REQUIRED) + """ + HK = k.shape[1] + HV = v.shape[1] + HO = gate.shape[1] + use_initial_state = initial_state is not None + store_final_state = output_state is not None + enable_checkpoints = checkpoint_every_n_tokens > 0 + if enable_checkpoints: + if output_state_checkpoints is None: + raise ValueError("checkpoint_every_n_tokens > 0 requires output_state_checkpoints") + if str(output_state_checkpoints.dtype).split(".")[-1] != str(k.dtype).split(".")[-1]: + raise ValueError( + f"output_state_checkpoints dtype must match the io dtype (fp32 state belongs to output_state): got {output_state_checkpoints.dtype} with io {k.dtype}" + ) + 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 + + if initial_state is not None: + state_dtype_src = initial_state.dtype + elif output_state is not None: + state_dtype_src = output_state.dtype + else: + state_dtype_src = "float32" + + for name, h in (("HK", HK), ("HV", HV)): + if HO % h != 0: + raise ValueError(f"{name}={h} must divide sab heads {HO}") + k_ratio = HO // HK + v_ratio = HO // HV + gate_scale_log2 = gate_lower_bound * LOG2_E + + if safe_gate and (a_log is None or dt_bias is None): + raise ValueError("safe_gate requires a_log and dt_bias") + if not safe_gate: + a_log = None + dt_bias = None + cu_stream = cuda_driver.CUstream(int(stream)) + + cache = get_compiled_cache( + str(k.dtype), + str(state_dtype_src), + str(cu_seqlens.dtype), + HO, + HK, + HV, + use_initial_state, + store_final_state, + enable_checkpoints, + use_qk_l2norm_in_kernel, + safe_gate, + gate_lower_bound, + dyn_sched, + ) + + if "compiled" not in cache: + io_dtype = get_dtype(k.dtype) + state_dtype = get_dtype(state_dtype_src) + k_cute = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_cute = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_cute = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2) + 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 + beta_cute = from_dlpack(beta, assumed_align=4).mark_layout_dynamic(leading_dim=2) + w_cute = from_dlpack(w, assumed_align=16).mark_layout_dynamic(leading_dim=2) + cu_seqlens_cute = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() + + state_in_cute = None + if use_initial_state: + state_in_cute = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) + + state_out_cute = None + if store_final_state: + state_out_cute = from_dlpack(output_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) + + work_items_cute = from_dlpack(work_items, assumed_align=16) + work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() + + sched_ctr_cute = None + if dyn_sched: + sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + + tensormap_ws_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + + cache["compiled"] = compile( + io_dtype, + state_dtype, + use_initial_state, + store_final_state, + enable_checkpoints, + use_qk_l2norm_in_kernel, + safe_gate, + gate_scale_log2, + k_ratio, + v_ratio, + HO, + dyn_sched, + num_sm=multiprocessor_count(current_device_id()), + k_cute=k_cute, + v_cute=v_cute, + gate_cute=gate_cute, + a_log_cute=a_log_cute, + dt_bias_cute=dt_bias_cute, + beta_cute=beta_cute, + w_cute=w_cute, + cu_seqlens_cute=cu_seqlens_cute, + state_in_cute=state_in_cute, + state_out_cute=state_out_cute, + work_items_cute=work_items_cute, + work_count_cute=work_count_cute, + sched_ctr_cute=sched_ctr_cute, + tensormap_ws_cute=tensormap_ws_cute, + checkpoint_every_n_tokens=checkpoint_every_n_tokens, + stream=cu_stream, + ) + + compiled = cache["compiled"] + state_checkpoints_for_descs = output_state_checkpoints if enable_checkpoints else None + # desc build runs every execute by contract (cu contents are data; + # buffer pointers may change) - capture-safe, single tiny launch + if cache.get("build_descs_has_state_checkpoints") != (state_checkpoints_for_descs is not None): + cache.pop("build_descs", None) + cache["build_descs_has_state_checkpoints"] = state_checkpoints_for_descs is not None + if "build_descs" not in cache: + io_dtype = get_dtype(k.dtype) + + cu_bd = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() + ws_bd = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + cache["build_descs"] = cute.compile( + build_descs, + io_dtype, + CFG.B_T, + from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(beta, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(w, assumed_align=16).mark_layout_dynamic(leading_dim=2), + None if state_checkpoints_for_descs is None else from_dlpack(state_checkpoints_for_descs, assumed_align=16).mark_layout_dynamic(leading_dim=3), + cu_bd, + ws_bd, + cutlass.Int32(checkpoint_every_n_tokens), + cu_stream, + options="--enable-tvm-ffi", + ) + cache["build_descs"]( + k, + v, + gate, + beta, + w, + state_checkpoints_for_descs, + cu_seqlens, + tensormap_workspace, + checkpoint_every_n_tokens, + cu_stream, + ) + compiled( + k, + v, + gate, + a_log, + dt_bias, + beta, + w, + cu_seqlens, + initial_state if use_initial_state else None, + output_state if store_final_state else None, + work_items, + work_count, + sched_ctr, + tensormap_workspace, + checkpoint_every_n_tokens, + cu_stream, + ) diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_config.py b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_config.py index 0fcf93fda..9536c5a85 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_config.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_config.py @@ -34,7 +34,6 @@ class Cfg: D_V: int = 128 # value head dim # --- TMA descriptor pool --- - BYTES_PER_TENSORMAP: int = 128 # --- warp assignments (12 warps total) --- COMPUTE_GROUP_0_WARP_IDS: Tuple[int, ...] = (0, 1, 2, 3) # T-pairwise / kk_epi / qk_epi / inverse / parts @@ -46,8 +45,8 @@ class Cfg: # --- register split --- NUM_REGS_COMPUTE_GROUP_0: int = 224 - NUM_REGS_COMPUTE_GROUP_1: int = 256 - NUM_REGS_OTHER: int = 24 + NUM_REGS_COMPUTE_GROUP_1: int = 248 + NUM_REGS_OTHER: int = 32 THREADS_PER_WARP: int = 32 @@ -58,8 +57,8 @@ class Cfg: SMEM_Q_STAGES: int = 1 SMEM_K_STAGES: int = 2 SMEM_V_STAGES: int = 1 - SMEM_AINV_STAGES: int = 1 - SMEM_QK_STAGES: int = 1 + SMEM_T_INV_STAGES: int = 1 + SMEM_A_STAGES: int = 1 # --- TMEM stage counts --- TMEM_DH_ACC_STAGES: int = 1 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 0b36946fb..27fbda843 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py @@ -19,88 +19,88 @@ Chunked Gated Delta Net (GDN) BPROP kernel for Blackwell SM100 (Cutlass primitives). Algorithm overview (per chunk c, iterated c = NT-1 .. 0): - Inputs : Q[BT,DK], K[BT,DK], V[BT,DV], dO[BT,DV], S_c[DK,DV] (= H[c-1], - the forward state ENTERING chunk c), gate[BT], beta[BT] - State : dH[DK,DV] (state gradient, held in TMEM, accumulated backward) + Inputs : Q[BT,DK], K[BT,DK], V[BT,DV], dO[BT,DV], S_c[DK,DV] (= checkpoint entry c-1, + the forward state ENTERING chunk c), Gate[BT], Beta[BT] + State : dstate[DK,DV] (state gradient, held in TMEM, accumulated backward) MMA order. A = the staged attention matrix - (CG0's qk epilogue, sQk); dO' = dO * gCumprod * scale (CG1 restage): - kk : W_kk[BT,BT] = K @ K^T -> shared acc (for T) - qk : W_qk[BT,BT] = Q @ K^T -> shared acc - dV inter : dV[DV,BT] = dH^T(TMEM) @ K^T -> the dV/dK slot - (acc=False; CG1 then scales it by gDecayScale IN PLACE) - ks : KS[BT,DV] = S^T(SMEM) @ K^T -> shared acc + (CG0's A epilogue, sA); dO' = dO * cumprod_vals * scale (CG1 restage): + KK : W_kk[BT,BT] = K @ K^T -> shared acc (for T) + QK : W_qk[BT,BT] = Q @ K^T -> shared acc (for the A tile) + dV inter : dV[DV,BT] = dstate^T(TMEM) @ K^T -> the dV/dK slot + (acc=False; CG1 then scales it by decay_scale_vals IN PLACE) + KS : KS[BT,DV] = S^T(SMEM) @ K^T -> shared acc dU intra : dV/dK slot += dO^T(SMEM) @ A(SMEM) (waits CG1's in-place scale -> the accumulate is the reference's dv2 = dv_intra + exp2(g_last-g)*(k@dh)) - dH q-term : dH += dO'^T(TMEM) @ Q (acc=False on the - first backward chunk = dH init; the reference's + dstate Q-term : dstate += dO'^T(TMEM) @ Q (acc=False on the + first backward chunk = dstate init; the reference's dh = dh*exp2(g_last) + (q*scale*exp2(g))^T @ do) dY : dY[DV,BT] = dU^T(TMEM f16) @ T(SMEM) -> shared acc - (dY == dV == d(delta); T read TRANSPOSED vs prefill) - dA : dA_eff[BT,BT]= dO(SMEM) @ delta^T(SMEM) -> shared acc - (sV holds delta; CG0 masks it -> sDa for dQ/dK) + (dY == dV == d(Y); T read TRANSPOSED vs prefill) + dA : dA_eff[BT,BT]= dO(SMEM) @ U^T(SMEM) -> shared acc + (sV holds U; CG0 masks it -> sDa for dQ/dK) dM core : dM[BT,BT] = dY(SMEM) @ U^T(SMEM) -- the WY inverse backward T^T dT T^T collapsed via - T^T dU = dY and T Y = U (beta folds through - sAinv's column scale into both factors); CG0 + T^T dU = dY and T Y = U (Beta folds through + sTinv's column scale into both factors); CG0 applies -strict(2^{g_i-g_j}) -> sDm (the reference's dA22; both betas cancel against the K rows in the dM-terms) - dH ds-term : dH += dY'^T(TMEM f16) @ K (dY' = -gCumprod - * dY: the -(w^T @ d(delta)) term; ready commits HERE) + dstate update-term : dstate += dY'^T(TMEM f16) @ K (dY' = -cumprod_vals + * dY: the -(W^T @ dY) term; ready commits HERE) dK dM-terms : dV/dK slot += K^T @ dM^T + K^T @ dM (the reference's dk += dA22^T@K_beta / dk_beta += dA22@K folded, the beta row scale cancelling with K_beta = beta*K) - dK s-path : SY[BT,DV] = S^T(SMEM) @ dY^T(SMEM) -> the shared f16 + dK state-path : SY[BT,DV] = S^T(SMEM) @ dY^T(SMEM) -> the shared f16 input columns (dO'/dU/dY' dead by then); CG1 reads it as - -gCumprod[t] * (dY @ S^T)[t,:] and adds it to the banked + -cumprod_vals[t] * (dY @ S^T)[t,:] and adds it to the banked inter+attn dK terms Per chunk CG1: restages dO' and dU and dY' -> the f16 input columns - (dY' overwrites dU), computes delta = V - gCumprod*(K @ S) in registers and stages - Y (= delta) and gks to their dedicated f16 TMEM slots (the u-GEMM and + (dY' overwrites dU), computes Y = V - cumprod_vals*(K @ state) in registers and stages + Y and g_k_state to their dedicated f16 TMEM slots (the U GEMM and the dV-pass read them), stages Q^T over the Y slot after the dV pass, loads dY and stages it plain to sdV (the dV output). SMEM layout (~226 KB of the 227 KB SM100 cap): Buffer Size (B) Stages - q 16384 1 - k 32768 2 <-- double-buffered (prefetch next chunk) - v 16384 1 <-- overwritten in place by u + Q 16384 1 + K 32768 2 <-- double-buffered (prefetch next chunk) + V 16384 1 <-- overwritten in place by U dO 16384 1 - S (forward state H[c-1]) 32768 1 <-- bf16 [DK,DV], TMA-loaded - A_inverse / T 8192 1 <-- inverse OUTPUT (upper tri = kernel-start zeros) - KK (pristine M_kk) 8192 1 <-- kk_epi's only store; inverse input + dG/dBeta - QK staging / sDa 8192 1 <-- ALIAS: A then the masked dA + state (checkpoint entry c-1) 32768 1 <-- io-dtype [DK,DV], TMA-loaded + T_inv 8192 1 <-- inverse OUTPUT (upper tri = kernel-start zeros) + KK (pristine M_kk) 8192 1 <-- KK epi's only store; inverse input + dGate/dBeta + A staging / sDa 8192 1 <-- ALIAS: A then the masked dA dM staging (sDm) 8192 1 <-- Step 8 -> dK dM-terms - dH_entry (sdH) 32768 1 <-- f16 restage, dK-inter's A + dstate_entry (sDstate) 32768 1 <-- f16 restage, dK-inter's A dQ store staging 16384 1 dK store staging 16384 1 dV store staging 16384 1 - cumsumlog / cumprod / beta 3 x 512 2 <-- in-place dG/dBeta staging + cumsumlog / cumprod / Beta 3 x 512 2 <-- in-place dGate/dBeta staging -TMEM layout (512 cols; EXACTLY the prefill map, S->dH, O->dV): - cols 0-128 : dH accumulator (fp32) <-- prefill: state (S) +TMEM layout (512 cols; EXACTLY the prefill map, state->dstate, O->dV): + cols 0-128 : dstate accumulator (fp32) <-- prefill: state acc cols 128-192 : dV/dK accumulator (fp32) <-- one slot, five sequential per-chunk productions (dV inter -> dU intra -> dK inter -> dK attn -> dK dM-terms), each with its own mbar pair - cols 192-256 : dH input (f16 packed) <-- prefill: state_inp - cols 256-384 : shared accumulators x2 <-- kk / qk / ks / u / dY / dA + cols 192-256 : dstate input (f16 packed) <-- prefill: state_inp + cols 256-384 : shared accumulators x2 <-- KK / A / k_state / U / dY / dA / dM core cols 384-448 : shared inputs x2 (f16 packed) <-- dO' / dU / dY'; the dK - s-path acc overwrites them after their last GEMM reads - (CG1's s-path readout precedes its next-chunk restages) - cols 448-512 : Y (448) + gks (480) f16 slots until the dV pass, then + state-path acc overwrites them after their last GEMM reads + (CG1's state-path readout precedes its next-chunk restages) + cols 448-512 : Y (448) + g_k_state (480) f16 slots until the dV pass, then Q^T (448) until CG1's dQ dot reads it Warp assignments (12 warps = 384 threads): - warps 0-3 : compute group 0 - T-pairwise, kk_epi, qk_epi, inverse - warps 4-7 : compute group 1 - dV epilogue, dH scale + f16 restage - (later: V-K*S -> SMEM, dY staging) + warps 0-3 : compute group 0 - T-pairwise, KK epi, A epi, inverse + warps 4-7 : compute group 1 - dV epilogue, dstate scale + f16 restage + (later: Y = V - K*state -> TMEM, dY staging) warp 8 : MMA warp - issues the GEMMs - warp 9 : TMA load warp - loads q, k (double-buf), v, dO, S(=H) - warp 10 : gate warp - loads gate/beta (double-buffered, - backward order) + stores dG/dBeta + warp 9 : TMA load warp - loads Q, K (double-buf), V, dO, state(=checkpoint) + warp 10 : gate warp - loads Gate/Beta (double-buffered, + backward order) + stores dGate/dBeta warp 11 : epilogue warp - store dQ, dK, dV to global memory """ @@ -113,13 +113,14 @@ import cutlass import cutlass.cute as cute import cutlass.experimental.primitives as nvvm -import cutlass.experimental.cuda.tensor_map as _tma -from cutlass.cute.arch.nvvm_wrappers import inline_ptx +import cutlass.experimental.cuda.tensor_map as tma from cutlass.cute.runtime import from_dlpack -from cutlass.cutlass_dsl import min as _cutlass_min -from ..common.thd import build_h_descs_kernel, build_qkv_load_descs_kernel, build_state_descs_kernel, TENSOR_MAP_QWORDS +from ..common.thd import emit_copy_desc, emit_checkpoint_seq_descs, emit_seq_descs, TENSOR_MAP_QWORDS from ..common.split_k import decode_work_item +from ..common.host import get_dtype +from cudnn.frost.buffers import current_device_id, data_ptr +from cudnn.frost.device import multiprocessor_count RCP_LN2 = 1.4426950408889634 # 1/ln(2): natural-log gates -> the kernel's log2 domain from cudnn.frost.tile_dsl.barrier import ( @@ -127,10 +128,9 @@ Producer, PipelineState, advance, - arrive, ) from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, tma_slice_runtime_desc -from cudnn.frost.tile_dsl.mma import mma_ss, mma_ts, mma_step +from cudnn.frost.tile_dsl.mma import mma_ss, mma_step_k8, mma_ts_step, mma_step from cudnn.frost.tile_dsl.pointwise import fp32_to_fp16, f16x2_to_f32, fmul2, fadd2, ffma2, opaque_f32_zero, sub_f16x2 from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_xor_128b from cudnn.frost.tile_dsl.tma import ( @@ -165,47 +165,59 @@ class GdnBwdBars(NamedTuple): mb_k_cg0_done: MBarrier mb_v_ready: MBarrier mb_v_mma_done: MBarrier - mb_v_cg1_done: MBarrier mb_do_ready: MBarrier mb_do_mma_done: MBarrier - mb_do_cg1_done: MBarrier - mb_s_ready: MBarrier - mb_s_done: MBarrier + mb_state_ready: MBarrier + mb_state_mma_done: MBarrier mb_gate_ready: MBarrier mb_gate_done: MBarrier mb_beta_ready: MBarrier mb_beta_done: MBarrier - mb_dh_acc_ready: MBarrier - mb_dh_acc_done: MBarrier - mb_du_scale_ready: MBarrier - mb_du_scale_done: MBarrier - mb_du_total_ready: MBarrier - mb_dk_scale_ready: MBarrier - mb_dk_scale_done: MBarrier - mb_dk_attn_ready: MBarrier - mb_dk_attn_done: MBarrier - mb_dk_total_ready: MBarrier - mb_dk_total_done: MBarrier + mb_dstate_acc_ready: MBarrier + mb_dstate_scale_acc_done: MBarrier + mb_du_scale_acc_ready: MBarrier + mb_du_scale_acc_done: MBarrier + mb_du_total_acc_ready: MBarrier + mb_dk_scale_acc_ready: MBarrier + mb_dk_scale_acc_done: MBarrier + mb_dk_attn_acc_ready: MBarrier + mb_dk_attn_acc_done: MBarrier + mb_dk_total_acc_ready: MBarrier + mb_dk_total_acc_done: MBarrier + mb_dq_acc_scale_ready: MBarrier + mb_dq_acc_scale_done: MBarrier + mb_dq_acc_total_ready: MBarrier + mb_dq_acc_total_done: MBarrier mb_kk_acc_ready: MBarrier mb_kk_acc_done: MBarrier mb_a_acc_ready: MBarrier - mb_ks_acc_ready: MBarrier + mb_k_state_acc_ready: MBarrier mb_u_acc_ready: MBarrier mb_dy_acc_ready: MBarrier + mb_da_acc_ready: MBarrier + mb_dm_acc_ready: MBarrier + mb_dm_acc_done: MBarrier + mb_dk_state_path_acc_ready: MBarrier - mb_ainv_ready: MBarrier - mb_ainv_done: MBarrier - mb_qk_ready: MBarrier - mb_qk_done: MBarrier - mb_dh_inp_ready: MBarrier - mb_dh_inp_done: MBarrier - mb_dop_inp_ready: MBarrier - mb_dop_inp_done: MBarrier + mb_dstate_inp_ready: MBarrier + mb_dstate_inp_done: MBarrier + mb_do_prime_inp_ready: MBarrier mb_du_inp_ready: MBarrier mb_dyp_inp_ready: MBarrier - mb_dyp_inp_done: MBarrier + mb_y_ready: MBarrier + + mb_t_inv_ready: MBarrier + mb_a_ready: MBarrier + mb_a_done: MBarrier + mb_u_ready: MBarrier + mb_dstate_smem_ready: MBarrier + mb_state_dot_dstate_done: MBarrier + mb_da_ready: MBarrier + + mb_dbeta_cg1_ready: MBarrier + mb_dgate_cg1_ready: MBarrier mb_dq_tmastg_ready: MBarrier mb_dq_tmastg_done: MBarrier @@ -213,31 +225,15 @@ class GdnBwdBars(NamedTuple): mb_dk_tmastg_done: MBarrier mb_dv_tmastg_ready: MBarrier mb_dv_tmastg_done: MBarrier - - mb_y_ready: MBarrier mb_sdv_done: MBarrier - mb_u_ready: MBarrier - mb_dhs_ready: MBarrier - mb_dhs_done: MBarrier - mb_da_ready: MBarrier - mb_dq_acc_scale_ready: MBarrier - mb_dq_acc_scale_done: MBarrier - mb_dq_acc_total_ready: MBarrier - mb_dq_acc_total_done: MBarrier - mb_da_acc_ready: MBarrier - mb_dm_ready: MBarrier - mb_dm_done: MBarrier - mb_dbeta_cg1_ready: MBarrier - mb_dgate_cg1_ready: MBarrier - mb_hdh_done: MBarrier - mb_dk_spath_ready: MBarrier + mb_tmem_done: MBarrier mb_sched_ready: MBarrier mb_sched_done: MBarrier def make_gdn_bars(cfg) -> GdnBwdBars: - """GdnBwdBars factory. MUST be called from inside ``_kernel`` (allocates SMEM; + """GdnBwdBars factory. MUST be called from inside ``kernel`` (allocates SMEM; the mbar rings sit ahead of the gate scalar arrays and data buffers).""" ONE_LANE = 1 MMA_ARRIVERS = len([cfg.mma_warp_id]) @@ -259,84 +255,82 @@ def alloc(n): mb_k_cg0_done=MBarrier(alloc(cfg.smem_k_stages), stages=cfg.smem_k_stages, init_count=CG0_THREADS, producer=Producer.THREAD), mb_v_ready=MBarrier(alloc(cfg.smem_v_stages), stages=cfg.smem_v_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), mb_v_mma_done=MBarrier(alloc(cfg.smem_v_stages), stages=cfg.smem_v_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_v_cg1_done=MBarrier(alloc(cfg.smem_v_stages), stages=cfg.smem_v_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_do_ready=MBarrier(alloc(cfg.smem_do_stages), stages=cfg.smem_do_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), mb_do_mma_done=MBarrier(alloc(cfg.smem_do_stages), stages=cfg.smem_do_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_do_cg1_done=MBarrier(alloc(cfg.smem_do_stages), stages=cfg.smem_do_stages, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_s_ready=MBarrier(alloc(cfg.smem_s_stages), stages=cfg.smem_s_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), - mb_s_done=MBarrier(alloc(cfg.smem_s_stages), stages=cfg.smem_s_stages, init_count=ONE_LANE + CG0_THREADS, producer=Producer.MMA_COMMIT), + mb_state_ready=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), + mb_state_mma_done=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=1, producer=Producer.MMA_COMMIT), mb_gate_ready=MBarrier(alloc(cfg.smem_gate_stages), stages=cfg.smem_gate_stages, init_count=GATE_WARP, producer=Producer.THREAD), mb_gate_done=MBarrier(alloc(cfg.smem_gate_stages), stages=cfg.smem_gate_stages, init_count=CG0_PLUS_CG1, producer=Producer.THREAD), mb_beta_ready=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=GATE_WARP, producer=Producer.THREAD), mb_beta_done=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=CG0_PLUS_CG1, producer=Producer.THREAD), - mb_dh_acc_ready=MBarrier(alloc(cfg.tmem_dh_acc_stages), stages=cfg.tmem_dh_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dh_acc_done=MBarrier(alloc(cfg.tmem_dh_acc_stages), stages=cfg.tmem_dh_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_dstate_acc_ready=MBarrier( + alloc(cfg.tmem_dstate_acc_stages), stages=cfg.tmem_dstate_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT + ), + mb_dstate_scale_acc_done=MBarrier( + alloc(cfg.tmem_dstate_acc_stages), stages=cfg.tmem_dstate_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD + ), # five sequential per-chunk productions share the dV/dK TMEM slot - mb_du_scale_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_du_scale_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_du_total_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dk_scale_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dk_scale_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_dk_attn_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dk_attn_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_dk_total_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dk_total_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_du_scale_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_du_scale_acc_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_du_total_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_dk_scale_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_dk_scale_acc_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_dk_attn_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_dk_attn_acc_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_dk_total_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_dk_total_acc_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_dq_acc_scale_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_dq_acc_scale_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_dq_acc_total_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_dq_acc_total_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), # shared accumulators at STATIC columns: group A holds - # kk -> ks -> dY -> dM core, group B holds A -> U -> dA. + # KK -> k_state -> dY -> dM core, group B holds A -> U -> dA. mb_kk_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_kk_acc_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), mb_a_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_ks_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_k_state_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_dy_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_ainv_ready=MBarrier(alloc(cfg.smem_ainv_stages), stages=cfg.smem_ainv_stages, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_ainv_done=MBarrier(alloc(cfg.smem_ainv_stages), stages=cfg.smem_ainv_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_qk_ready=MBarrier(alloc(cfg.smem_qk_stages), stages=cfg.smem_qk_stages, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_qk_done=MBarrier(alloc(cfg.smem_qk_stages), stages=cfg.smem_qk_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dh_inp_ready=MBarrier(alloc(cfg.tmem_dh_inp_stages), stages=cfg.tmem_dh_inp_stages, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_dh_inp_done=MBarrier(alloc(cfg.tmem_dh_inp_stages), stages=cfg.tmem_dh_inp_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_da_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_dm_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_dm_acc_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_dk_state_path_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_dstate_inp_ready=MBarrier(alloc(cfg.tmem_dstate_inp_stages), stages=cfg.tmem_dstate_inp_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_dstate_inp_done=MBarrier( + alloc(cfg.tmem_dstate_inp_stages), stages=cfg.tmem_dstate_inp_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT + ), # f16 input restages at STATIC columns: dO' alone; dU and dY' overlap - mb_dop_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_dop_inp_done=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_do_prime_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_du_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_dyp_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_dyp_inp_done=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_y_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_t_inv_ready=MBarrier(alloc(cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_a_ready=MBarrier(alloc(cfg.smem_a_stages), stages=cfg.smem_a_stages, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_a_done=MBarrier(alloc(cfg.smem_a_stages), stages=cfg.smem_a_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_u_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_dstate_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_state_dot_dstate_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_da_ready=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_dbeta_cg1_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_dgate_cg1_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_dq_tmastg_ready=MBarrier(alloc(cfg.smem_dq_stages), stages=cfg.smem_dq_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_dq_tmastg_done=MBarrier(alloc(cfg.smem_dq_stages), stages=cfg.smem_dq_stages, init_count=EPI_WARP, producer=Producer.THREAD), mb_dk_tmastg_ready=MBarrier(alloc(cfg.smem_dk_stages), stages=cfg.smem_dk_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_dk_tmastg_done=MBarrier(alloc(cfg.smem_dk_stages), stages=cfg.smem_dk_stages, init_count=EPI_WARP, producer=Producer.THREAD), mb_dv_tmastg_ready=MBarrier(alloc(cfg.smem_dv_stages), stages=cfg.smem_dv_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_dv_tmastg_done=MBarrier(alloc(cfg.smem_dv_stages), stages=cfg.smem_dv_stages, init_count=EPI_WARP, producer=Producer.THREAD), - mb_y_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_sdv_done=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_u_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_dhs_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_dhs_done=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_da_ready=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_dq_acc_scale_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dq_acc_scale_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_dq_acc_total_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dq_acc_total_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_da_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dm_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_dm_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_dbeta_cg1_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_dgate_cg1_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_hdh_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_dk_spath_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG0_THREADS + CG1_THREADS, producer=Producer.THREAD), mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=ONE_LANE, producer=Producer.THREAD), mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=11, producer=Producer.THREAD), ) -# --------------------------------------------------------------------------- -# Dynamic tile scheduler: global-ticket work-stealing ring -# --------------------------------------------------------------------------- +# ---- Dynamic tile scheduler ------------------------------------------------------ @cute.jit -def _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): """TMA-LDG-warp side: pull the next tile off the global ticket, publish it.""" if cutlass.const_expr(cfg.dyn_sched): bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) @@ -352,7 +346,7 @@ def _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ct @cute.jit -def _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): """Consumer side: read the TMA-LDG warp's published next tile.""" if cutlass.const_expr(cfg.dyn_sched): bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) @@ -363,20 +357,9 @@ def _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): return tile_idx + num_ctas, sched_state -# --------------------------------------------------------------------------- -# Device-side helpers / warp bodies -# --------------------------------------------------------------------------- - - @cute.jit -def _invert_diagonal_NxN(cfg, in_base, out_base, d, tidx, N: int = 8): - """Stage 1: Gauss-Jordan inversion of one diagonal NxN block, reading the - raw block from ``in_base`` and writing the inverse to ``out_base`` (same - swizzled offsets). - - The tile swizzle re-homes whole rows only, so a diagonal block's row stays - a contiguous N-element run at ``swz(row_lin_base)``. - """ +def invert_diagonal_NxN(cfg, in_base, out_base, d, tidx, N: int = 8): + """Gauss-Jordan inversion of one diagonal NxN block, ``in_base`` -> ``out_base`` (f16 SMEM).""" tidx_in_group = tidx % N BT = cfg.b_t @@ -416,23 +399,9 @@ def _invert_diagonal_NxN(cfg, in_base, out_base, d, tidx, N: int = 8): @cute.jit -def _mma_m16n8k8(a0, a1, b0, c_regs, dtype: cutlass.Constexpr): - """One m16n8k8 reg-reg mma (``mma_step`` only emits the k16 form), - accumulating into the ``c_regs`` buffer in place.""" - tag = "f16" if cutlass.const_expr(dtype == cutlass.Float16) else "bf16" - c_regs[0], c_regs[1], c_regs[2], c_regs[3] = inline_ptx( - f"mma.sync.aligned.m16n8k8.row.col.f32.{tag}.{tag}.f32" " {$0,$1,$2,$3}, {$4,$5}, {$6}, {$7,$8,$9,$10};", - write_only_types=[cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Float32], - read_only_args=[a0, a1, b0, c_regs[0], c_regs[1], c_regs[2], c_regs[3]], - ) - - -@cute.jit -def _warp_reduce_scatter_frag16(vals, lane_id): +def warp_reduce_scatter_frag_16_elems(vals, lane_id): """Reduce-scatter 16 fragment token-partials (tcol = (lane%4)*2 + - (j//2)*8 + (j%2)) over the 8 lane-groups: step k exchanges with - lane^(4< 16x16 (C <- -D^{-1} C A^{-1}). - Raw C blocks read from ``raw_base`` (sKK); corrected blocks and all - writes on ``base_int`` (sAinv). - - Keep the per-lane ldmatrix offset (``lds1``/``lds4``) a SEPARATE sum term - from the warp-uniform (row, col) origin in all the diagonal helpers — - folding it into the row term costs ~2% (per-lane address datapath). - """ +def blockwise_diagonal_8x8_to_16x16(cfg, base_int, raw_base, d0, lane_id): + """Off-diagonal correction 8x8 -> 16x16 (C <- -D^{-1} C A^{-1}); raw C from ``raw_base``, writes on ``base_int``.""" bpe = cfg.io_dtype.width // 8 - lds1 = (lane_id % 8) * 64 + ldsm_x1_off = (lane_id % 8) * 64 d = nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + 8 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + cutlass.inttoptr( + base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + 8 + ldsm_x1_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), 1, nvvm.MMALayout.ROW, ) c = nvvm.ldmatrix( - cutlass.inttoptr(raw_base + swizzle_lin_128b((d0 + 8) * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + cutlass.inttoptr(raw_base + swizzle_lin_128b((d0 + 8) * 64 + d0 + ldsm_x1_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), 1, nvvm.MMALayout.COL, ) + + # ---- T = -(D^{-1} @ C) ------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(4): c_regs[i] = cutlass.Float32(0.0) - _mma_m16n8k8(d, d, c, c_regs, cfg.io_dtype) + mma_step_k8(c_regs, [d, d], [c], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(4): c_regs[i] = -c_regs[i] - a_f16 = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(2)] + a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(2)] + + # ---- C = T @ A^{-1} ---------------------------------------------------------- ai = nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + ldsm_x1_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), 1, nvvm.MMALayout.COL, ) o_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(4): o_regs[i] = cutlass.Float32(0.0) - _mma_m16n8k8(a_f16[0], a_f16[1], ai, o_regs, cfg.io_dtype) - o_f16 = fp32_to_fp16(o_regs[0], o_regs[1], dtype=cfg.io_dtype) + mma_step_k8(o_regs, a_pack, [ai], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) + o_pack = fp32_to_fp16(o_regs[0], o_regs[1], dtype=cfg.io_dtype) + + # ---- store corrected C ------------------------------------------------------- nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - o_f16, + cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + ldsm_x1_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + o_pack, nvvm.MMALayout.ROW, ) @cute.jit -def _blockwise_diagonal_16x16_to_32x32(cfg, base_int, raw_base, d0, lane_id): - """Stage 3: off-diagonal correction 16x16 -> 32x32 (raw C from - ``raw_base``).""" +def blockwise_diagonal_16x16_to_32x32(cfg, base_int, raw_base, d0, lane_id): + """Off-diagonal correction 16x16 -> 32x32 (raw C from ``raw_base``).""" bpe = cfg.io_dtype.width // 8 - lds4 = (lane_id % 16) * 64 + (lane_id // 16) * 8 + ldsm_x4_off = (lane_id % 16) * 64 + (lane_id // 16) * 8 d = list( nvvm.ldmatrix( cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + 16 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + 16 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 ), 4, nvvm.MMALayout.ROW, @@ -513,21 +482,27 @@ def _blockwise_diagonal_16x16_to_32x32(cfg, base_int, raw_base, d0, lane_id): ) c = list( nvvm.ldmatrix( - cutlass.inttoptr(raw_base + swizzle_lin_128b((d0 + 16) * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + cutlass.inttoptr( + raw_base + swizzle_lin_128b((d0 + 16) * 64 + d0 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), 4, nvvm.MMALayout.COL, ) ) + + # ---- T = -(D^{-1} @ C) ------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(8): c_regs[i] = cutlass.Float32(0.0) mma_step(c_regs, d, c, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(8): c_regs[i] = -c_regs[i] - a_f16 = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + + # ---- C = T @ A^{-1} ---------------------------------------------------------- ai = list( nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), 4, nvvm.MMALayout.COL, ) @@ -535,28 +510,29 @@ def _blockwise_diagonal_16x16_to_32x32(cfg, base_int, raw_base, d0, lane_id): o_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(8): o_regs[i] = cutlass.Float32(0.0) - mma_step(o_regs, a_f16, ai, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) - ow = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + mma_step(o_regs, a_pack, ai, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) + o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + + # ---- store corrected C ------------------------------------------------------- nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - ow, + cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + o_pack, nvvm.MMALayout.ROW, ) @cute.jit -def _blockwise_diagonal_32x32_to_64x64(cfg, base_int, raw_base, warp_id, lane_id): - """Stage 4: off-diagonal correction 32x32 -> 64x64 (2 warps, one 16-row - M-band each; raw C from ``raw_base``).""" +def blockwise_diagonal_32x32_to_64x64(cfg, base_int, raw_base, warp_id, lane_id): + """Off-diagonal correction 32x32 -> 64x64 (2 warps, one 16-row M-band each; raw C from ``raw_base``).""" band = warp_id % 2 bpe = cfg.io_dtype.width // 8 - lds4 = (lane_id % 16) * 64 + (lane_id // 16) * 8 - a_regs = [] + ldsm_x4_off = (lane_id % 16) * 64 + (lane_id // 16) * 8 + a_frags = [] for vs in cutlass.range_constexpr(2): - a_regs += list( + a_frags += list( nvvm.ldmatrix( cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + band * 16) * 64 + 32 + vs * 16 + lds4, row_stride_log2=6) * bpe, + base_int + swizzle_lin_128b((32 + band * 16) * 64 + 32 + vs * 16 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), @@ -564,12 +540,12 @@ def _blockwise_diagonal_32x32_to_64x64(cfg, base_int, raw_base, warp_id, lane_id nvvm.MMALayout.ROW, ) ) - b_regs = [] + b_frags = [] for vs in cutlass.range_constexpr(4): - b_regs += list( + b_frags += list( nvvm.ldmatrix( cutlass.inttoptr( - raw_base + swizzle_lin_128b((32 + (vs // 2) * 16) * 64 + (vs % 2) * 16 + lds4, row_stride_log2=6) * bpe, + raw_base + swizzle_lin_128b((32 + (vs // 2) * 16) * 64 + (vs % 2) * 16 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), @@ -577,20 +553,24 @@ def _blockwise_diagonal_32x32_to_64x64(cfg, base_int, raw_base, warp_id, lane_id nvvm.MMALayout.COL, ) ) + + # ---- T = -(D^{-1} @ C) ------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 16, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(16): c_regs[i] = cutlass.Float32(0.0) for ks in cutlass.range_constexpr(2): - mma_step(c_regs, a_regs, b_regs[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + mma_step(c_regs, a_frags, b_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(16): c_regs[i] = -c_regs[i] - a_f16 = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] - ai_regs = [] + a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] + + # ---- C = T @ A^{-1} ---------------------------------------------------------- + ai_frags = [] for vs in cutlass.range_constexpr(4): - ai_regs += list( + ai_frags += list( nvvm.ldmatrix( cutlass.inttoptr( - base_int + swizzle_lin_128b(((vs // 2) * 16) * 64 + (vs % 2) * 16 + lds4, row_stride_log2=6) * bpe, + base_int + swizzle_lin_128b(((vs // 2) * 16) * 64 + (vs % 2) * 16 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), @@ -602,26 +582,32 @@ def _blockwise_diagonal_32x32_to_64x64(cfg, base_int, raw_base, warp_id, lane_id for i in cutlass.range_constexpr(16): o_regs[i] = cutlass.Float32(0.0) for ks in cutlass.range_constexpr(2): - mma_step(o_regs, a_f16, ai_regs[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) - ow = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] + mma_step(o_regs, a_pack, ai_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] + + # ---- store corrected C ------------------------------------------------------- nvvm.barrier_cta_sync_aligned( cfg.inverse_inner_barrier_id, thread_count=cfg.inverse_inner_barrier_threads, ) nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((32 + band * 16) * 64 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - ow[0:4], + cutlass.inttoptr( + base_int + swizzle_lin_128b((32 + band * 16) * 64 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), + o_pack[0:4], nvvm.MMALayout.ROW, ) nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((32 + band * 16) * 64 + 16 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - ow[4:8], + cutlass.inttoptr( + base_int + swizzle_lin_128b((32 + band * 16) * 64 + 16 + ldsm_x4_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), + o_pack[4:8], nvvm.MMALayout.ROW, ) @cute.jit -def _tmastg_warp( +def tmastg_warp( cfg, total_tiles, bidx, @@ -639,17 +625,18 @@ def _tmastg_warp( ): """TMA-STG warp role (warp 11): persistent tile-scheduler loop + per-chunk dQ/dK/dV TMA bulk-stores from the SMEM staging buffers to global memory.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) dq_index = PipelineState.start(phase=0) dk_index = PipelineState.start(phase=0) dv_index = PipelineState.start(phase=0) + sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) - S_MIN = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 + FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 bpe = cfg.io_dtype.width // 8 - granu = 128 // bpe + granule_elems = 128 // bpe sdQ_tma = SmemTile( base=sdQ_raw, elems_per_stage=(cfg.dq_cosize // cfg.smem_dq_stages), @@ -658,7 +645,7 @@ def _tmastg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, + tma_granu_elems=granule_elems, tma_subtile_stride_elems=4096, ) sdK_tma = SmemTile( @@ -669,7 +656,7 @@ def _tmastg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, + tma_granu_elems=granule_elems, tma_subtile_stride_elems=4096, ) sdV_tma = SmemTile( @@ -680,22 +667,21 @@ def _tmastg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, + tma_granu_elems=granule_elems, tma_subtile_stride_elems=4096, ) heads_out = cutlass.Int32(cfg.n_heads_out) desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) 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, cu_seqlens, mWorkItems - ) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - slot = (batch_idx * heads_out + head_idx) * desc_qwords + head_o = head_idx + slot = batch_idx * desc_qwords desc_dq_slot = (desc_dq_base + slot).tospace(cutlass.AddressSpace.generic) desc_dk_slot = (desc_dk_base + slot).tospace(cutlass.AddressSpace.generic) desc_dv_slot = (desc_dv_base + slot).tospace(cutlass.AddressSpace.generic) - if nvvm.elect_sync(): + if elect_one: tma_tensormap_acquire(desc_dq_slot) tma_tensormap_acquire(desc_dk_slot) tma_tensormap_acquire(desc_dv_slot) @@ -707,37 +693,24 @@ def _tmastg_warp( dv_idx = dv_index.idx bars.mb_dv_tmastg_ready[dv_idx].wait(dv_index.phase) dv_index = advance(dv_index, cfg.smem_dv_stages) - dv_slice = tma_slice_runtime_desc(desc_dv_slot, cutlass.Int32(0), tok_coord) - if cutlass.const_expr(cfg.split_k): - # right-warmup chunks stage grads to SMEM but never store them - if chunk_idx < wend: - tma_store_tile(sdV_tma[dv_idx], dv_slice, acquire=False) - tma_store_commit() - else: + dv_slice = tma_slice_runtime_desc(desc_dv_slot, cutlass.Int32(0), head_o, tok_coord) + if chunk_idx < wend: tma_store_tile(sdV_tma[dv_idx], dv_slice, acquire=False) tma_store_commit() dq_idx = dq_index.idx bars.mb_dq_tmastg_ready[dq_idx].wait(dq_index.phase) dq_index = advance(dq_index, cfg.smem_dq_stages) - dq_slice = tma_slice_runtime_desc(desc_dq_slot, cutlass.Int32(0), tok_coord) - if cutlass.const_expr(cfg.split_k): - if chunk_idx < wend: - tma_store_tile(sdQ_tma[dq_idx], dq_slice, acquire=False) - tma_store_commit() - else: + dq_slice = tma_slice_runtime_desc(desc_dq_slot, cutlass.Int32(0), head_o, tok_coord) + if chunk_idx < wend: tma_store_tile(sdQ_tma[dq_idx], dq_slice, acquire=False) tma_store_commit() dk_idx = dk_index.idx bars.mb_dk_tmastg_ready[dk_idx].wait(dk_index.phase) dk_index = advance(dk_index, cfg.smem_dk_stages) - dk_slice = tma_slice_runtime_desc(desc_dk_slot, cutlass.Int32(0), tok_coord) - if cutlass.const_expr(cfg.split_k): - if chunk_idx < wend: - tma_store_tile(sdK_tma[dk_idx], dk_slice, acquire=False) - tma_store_commit() - else: + dk_slice = tma_slice_runtime_desc(desc_dk_slot, cutlass.Int32(0), head_o, tok_coord) + if chunk_idx < wend: tma_store_tile(sdK_tma[dk_idx], dk_slice, acquire=False) tma_store_commit() @@ -747,11 +720,11 @@ def _tmastg_warp( bars.mb_dq_tmastg_done[dq_idx].arrive() tma_store_wait(0) bars.mb_dk_tmastg_done[dk_idx].arrive() - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @cute.jit -def _gate_beta_warp( +def gate_beta_warp( cfg, total_tiles, bidx, @@ -761,7 +734,7 @@ def _gate_beta_warp( tidx, mGate, mBeta, - mDg, + mDgate, mDbeta, sCumsumlog, sCumprod, @@ -769,8 +742,8 @@ def _gate_beta_warp( sSched, bars, ): - """Gate/beta LOAD + STORE warp role (warp 10): per-chunk gate/beta G->S - loads (BACKWARD order, one-chunk prefetch) and the dG/dBeta stores.""" + """Gate/beta LOAD + STORE warp role (warp 10): per-chunk Gate/Beta G->S + loads and the dGate/dBeta stores.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) gate_index = PipelineState.start(phase=1) beta_index = PipelineState.start(phase=1) @@ -780,21 +753,17 @@ def _gate_beta_warp( n_cols = cfg.b_t // cfg.threads_per_warp sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) - S_MIN = 0 if cfg.use_initial_state else 1 + 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, cu_seqlens, mWorkItems - ) - sk_nt = cend - wstart - if cutlass.const_expr(cfg.split_k): - # dG/dBeta ownership: mask stores past the item's write range - write_end = batch_start + wend * cfg.b_t - write_end = write_end if write_end < batch_end else batch_end - else: - write_end = batch_end - # ---- prefetch: the FIRST backward chunk's gate/beta ------------ - if sk_nt > 0: + 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 + # dGate/dBeta ownership: mask stores past the item's write range + write_end = batch_start + wend * cfg.b_t + write_end = write_end if write_end < batch_end else batch_end + + # ---- prefetch: the FIRST backward chunk's Gate/Beta ---------------------- + if num_item_chunks > 0: chunk_offset = batch_start + (cend - 1) * cfg.b_t gGate = cute.domain_offset((chunk_offset,), mGate[None, head_idx]) gBeta = cute.domain_offset((chunk_offset,), mBeta[None, head_idx]) @@ -805,42 +774,42 @@ def _gate_beta_warp( pos = lidx + col * cfg.threads_per_warp pos_valid[col] = cute.elem_less(chunk_offset + pos, batch_end) - # --- Gate load (OOB neutral: 1.0 -> log2 = 0.0) --- - tGrGate = [cutlass.Float32(0.0)] * n_cols + # ---- Gate load: GMEM -> SMEM (OOB neutral: 1.0 -> log2 = 0.0) -------- + gate_vals = [cutlass.Float32(0.0)] * n_cols for col in cutlass.range_constexpr(n_cols): pos = lidx + col * cfg.threads_per_warp oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) - tGrGate[col] = gGate[pos] if pos_valid[col] else oob_neutral + gate_vals[col] = gGate[pos] if pos_valid[col] else oob_neutral if cutlass.const_expr(cfg.log_gate): for col in cutlass.range_constexpr(n_cols): - tGrGate[col] = tGrGate[col] * cutlass.Float32(RCP_LN2) + gate_vals[col] = gate_vals[col] * cutlass.Float32(RCP_LN2) else: for col in cutlass.range_constexpr(n_cols): - tGrGate[col] = cute.math.log2(tGrGate[col] + 1e-10, fastmath=True) + gate_vals[col] = cute.math.log2(gate_vals[col] + 1e-10, fastmath=True) for offset in [1, 2, 4, 8, 16]: for col in cutlass.range_constexpr(n_cols): - n = nvvm.shfl_sync(0xFFFFFFFF, tGrGate[col], offset, 0, kind=nvvm.Shfl.UP) + n = nvvm.shfl_sync(0xFFFFFFFF, gate_vals[col], offset, 0, kind=nvvm.Shfl.UP) if lidx >= offset: - tGrGate[col] = tGrGate[col] + n + gate_vals[col] = gate_vals[col] + n for col in cutlass.range_constexpr(1, n_cols): last_v = nvvm.shfl_sync( 0xFFFFFFFF, - tGrGate[col - 1], + gate_vals[col - 1], cfg.threads_per_warp - 1, cfg.threads_per_warp - 1, kind=nvvm.Shfl.IDX, ) - tGrGate[col] += last_v + gate_vals[col] += last_v for col in cutlass.range_constexpr(n_cols): pos = lidx + col * cfg.threads_per_warp - sCumsumlog[pos, 0, gate_idx] = tGrGate[col] - sCumprod[pos, 0, gate_idx] = cute.math.exp2(tGrGate[col], fastmath=True) + sCumsumlog[pos, 0, gate_idx] = gate_vals[col] + sCumprod[pos, 0, gate_idx] = cute.math.exp2(gate_vals[col], fastmath=True) bars.mb_gate_ready[gate_idx].arrive() - # --- Beta load (per-element async G->S cp.async) --- + # ---- Beta load: GMEM -> SMEM (per-element cp.async) ------------------ beta_idx = beta_index.idx beta_index = advance(beta_index, cfg.smem_beta_stages) for col in cutlass.range_constexpr(n_cols): @@ -851,9 +820,9 @@ def _gate_beta_warp( nvvm.cp_async_shared_global(dst, src, 4, nvvm.LoadCacheModifier.CA, cp_size=cp_size) nvvm.cp_async_mbarrier_arrive(bars.mb_beta_ready[beta_idx].smem_ptr, noinc=True) - for rev_idx in cutlass.range(sk_nt): - # ---- prefetch the NEXT chunk's gate/beta ------------------- - if rev_idx + 1 < sk_nt: + for rev_idx in cutlass.range(num_item_chunks): + # ---- prefetch the NEXT chunk's Gate/Beta ----------------------------- + if rev_idx + 1 < num_item_chunks: chunk_offset = batch_start + (cend - 2 - rev_idx) * cfg.b_t gGate = cute.domain_offset((chunk_offset,), mGate[None, head_idx]) gBeta = cute.domain_offset((chunk_offset,), mBeta[None, head_idx]) @@ -864,37 +833,37 @@ def _gate_beta_warp( pos = lidx + col * cfg.threads_per_warp pos_valid[col] = cute.elem_less(chunk_offset + pos, batch_end) - tGrGate = [cutlass.Float32(0.0)] * n_cols + gate_vals = [cutlass.Float32(0.0)] * n_cols for col in cutlass.range_constexpr(n_cols): pos = lidx + col * cfg.threads_per_warp oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) - tGrGate[col] = gGate[pos] if pos_valid[col] else oob_neutral + gate_vals[col] = gGate[pos] if pos_valid[col] else oob_neutral if cutlass.const_expr(cfg.log_gate): for col in cutlass.range_constexpr(n_cols): - tGrGate[col] = tGrGate[col] * cutlass.Float32(RCP_LN2) + gate_vals[col] = gate_vals[col] * cutlass.Float32(RCP_LN2) else: for col in cutlass.range_constexpr(n_cols): - tGrGate[col] = cute.math.log2(tGrGate[col] + 1e-10, fastmath=True) + gate_vals[col] = cute.math.log2(gate_vals[col] + 1e-10, fastmath=True) for offset in [1, 2, 4, 8, 16]: for col in cutlass.range_constexpr(n_cols): - n = nvvm.shfl_sync(0xFFFFFFFF, tGrGate[col], offset, 0, kind=nvvm.Shfl.UP) + n = nvvm.shfl_sync(0xFFFFFFFF, gate_vals[col], offset, 0, kind=nvvm.Shfl.UP) if lidx >= offset: - tGrGate[col] = tGrGate[col] + n + gate_vals[col] = gate_vals[col] + n for col in cutlass.range_constexpr(1, n_cols): last_v = nvvm.shfl_sync( 0xFFFFFFFF, - tGrGate[col - 1], + gate_vals[col - 1], cfg.threads_per_warp - 1, cfg.threads_per_warp - 1, kind=nvvm.Shfl.IDX, ) - tGrGate[col] += last_v + gate_vals[col] += last_v for col in cutlass.range_constexpr(n_cols): pos = lidx + col * cfg.threads_per_warp - sCumsumlog[pos, 0, gate_idx] = tGrGate[col] - sCumprod[pos, 0, gate_idx] = cute.math.exp2(tGrGate[col], fastmath=True) + sCumsumlog[pos, 0, gate_idx] = gate_vals[col] + sCumprod[pos, 0, gate_idx] = cute.math.exp2(gate_vals[col], fastmath=True) bars.mb_gate_ready[gate_idx].arrive() @@ -908,30 +877,30 @@ def _gate_beta_warp( nvvm.cp_async_shared_global(dst, src, 4, nvvm.LoadCacheModifier.CA, cp_size=cp_size) nvvm.cp_async_mbarrier_arrive(bars.mb_beta_ready[beta_idx].smem_ptr, noinc=True) - # ---- store-ready wait + in-place store back ---------------- + # ---- store-ready wait + in-place store back -------------------------- st_offset = batch_start + (cend - 1 - rev_idx) * cfg.b_t - gGate_st = cute.domain_offset((st_offset,), mDg[None, head_idx]) + gGate_st = cute.domain_offset((st_offset,), mDgate[None, head_idx]) gBeta_st = cute.domain_offset((st_offset,), mDbeta[None, head_idx]) g_st_idx = gate_store_index.idx bars.mb_gate_done[g_st_idx].wait(gate_store_index.phase) gate_store_index = advance(gate_store_index, cfg.smem_gate_stages) - tGrDg = [cutlass.Float32(0.0)] * n_cols + dgate_vals = [cutlass.Float32(0.0)] * n_cols for col in cutlass.range_constexpr(n_cols): pos = lidx + col * cfg.threads_per_warp - tGrDg[col] = sCumsumlog[pos, 0, g_st_idx] + dgate_vals[col] = sCumsumlog[pos, 0, g_st_idx] for offset in [1, 2, 4, 8, 16]: for col in cutlass.range_constexpr(n_cols): - n = nvvm.shfl_sync(0xFFFFFFFF, tGrDg[col], offset, 31, kind=nvvm.Shfl.DOWN) + n = nvvm.shfl_sync(0xFFFFFFFF, dgate_vals[col], offset, 31, kind=nvvm.Shfl.DOWN) if lidx < cfg.threads_per_warp - offset: - tGrDg[col] = tGrDg[col] + n + dgate_vals[col] = dgate_vals[col] + n for col in cutlass.range_constexpr(n_cols - 1): - cc = cutlass.const_expr(n_cols - 2 - col) - later_total = nvvm.shfl_sync(0xFFFFFFFF, tGrDg[cc + 1], 0, 0, kind=nvvm.Shfl.IDX) - tGrDg[cc] = tGrDg[cc] + later_total + rev_col = cutlass.const_expr(n_cols - 2 - col) + later_total = nvvm.shfl_sync(0xFFFFFFFF, dgate_vals[rev_col + 1], 0, 0, kind=nvvm.Shfl.IDX) + dgate_vals[rev_col] = dgate_vals[rev_col] + later_total for col in cutlass.range_constexpr(n_cols): pos = lidx + col * cfg.threads_per_warp if cute.elem_less(st_offset + pos, write_end): - gGate_st[pos] = tGrDg[col] + gGate_st[pos] = dgate_vals[col] b_st_idx = beta_store_index.idx bars.mb_beta_done[b_st_idx].wait(beta_store_index.phase) beta_store_index = advance(beta_store_index, cfg.smem_beta_stages) @@ -939,18 +908,18 @@ def _gate_beta_warp( pos = lidx + col * cfg.threads_per_warp if cute.elem_less(st_offset + pos, write_end): gBeta_st[pos] = sBeta[pos, 0, b_st_idx] - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @cute.jit -def _mma_warp( +def mma_warp( cfg, total_tiles, bidx, num_ctas, cu_seqlens, mWorkItems, - tmem_hold, + tmem_base_slot, sQ, sQ_trans, sK, @@ -959,62 +928,60 @@ def _mma_warp( sV_kmaj, sdO, sdO_kmaj, - sS, - sS_kmaj, - sAinv, - sAinv_trans, - sQk, - sQk_trans, + sState, + sState_kmaj, + sTinv, + sTinv_trans, + sA, + sA_trans, sDa, sDa_trans, - sdH, + sDstate, sDm, sDm_trans, - sdV, sdV_kmaj, sSched, bars, ): - """MMA (UMMA issuer) warp role (warp 8): persistent scheduler loop + - per-chunk (BACKWARD order) issue of the full GEMM pipeline (see the - module docstring for the per-chunk MMA order); owns the TMEM lifecycle - (alloc up front, dealloc once CG0 and CG1 both signal mb_tmem_done).""" + """MMA issuer role (warp 8): persistent scheduler loop issuing every + tcgen05 GEMM.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) - kk_read = PipelineState.start(phase=0) - dk_total_free = PipelineState.start(phase=1) - du_scale_read = PipelineState.start(phase=0) - dk_scale_read = PipelineState.start(phase=0) - dk_attn_read = PipelineState.start(phase=0) - dh_acc_index = PipelineState.start(phase=0 if cfg.use_dht else 1) + kk_acc_index = PipelineState.start(phase=0) + dk_total_index = PipelineState.start(phase=1) + du_scale_index = PipelineState.start(phase=0) + dk_scale_index = PipelineState.start(phase=0) + dk_attn_index = PipelineState.start(phase=0) + dstate_acc_index = PipelineState.start(phase=0 if cfg.use_dstate_in else 1) k_index = PipelineState.start(phase=0) q_index = PipelineState.start(phase=0) - s_index = PipelineState.start(phase=0) + state_index = PipelineState.start(phase=0) v_index = PipelineState.start(phase=0) - ainv_index = PipelineState.start(phase=0) + tinv_index = PipelineState.start(phase=0) do_index = PipelineState.start(phase=0) - y_rdy_index = PipelineState.start(phase=0) + y_ready_index = PipelineState.start(phase=0) u_index = PipelineState.start(phase=0) - da_rdy_index = PipelineState.start(phase=0) - dv_rdy_index = PipelineState.start(phase=0) + da_ready_index = PipelineState.start(phase=0) + dv_ready_index = PipelineState.start(phase=0) dm_index = PipelineState.start(phase=0) - dhs_index = PipelineState.start(phase=0) + dstate_smem_index = PipelineState.start(phase=0) dq_scale_index = PipelineState.start(phase=0) dq_total_index = PipelineState.start(phase=1) - qk_index = PipelineState.start(phase=0) - dop_inp_rdy = PipelineState.start(phase=0) - du_inp_rdy = PipelineState.start(phase=0) - dyp_inp_rdy = PipelineState.start(phase=0) - dh_inp_index = PipelineState.start(phase=0) + a_index = PipelineState.start(phase=0) + do_prime_inp_ready = PipelineState.start(phase=0) + du_inp_ready = PipelineState.start(phase=0) + dyp_inp_ready = PipelineState.start(phase=0) + dstate_inp_index = PipelineState.start(phase=0) - nvvm.tcgen05_alloc(tmem_hold, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) + nvvm.tcgen05_alloc(tmem_base_slot, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) nvvm.barrier_cta_sync_aligned( cfg.tmem_alloc_barrier_id, thread_count=cfg.tmem_alloc_barrier_threads, ) - tmem_base = tmem_hold.load() + tmem_base = tmem_base_slot.load() - # ---- chunk-invariant GEMM descriptors ------------------------------ + # ---- chunk-invariant GEMM descriptors ---------------------------------------- bpe = cfg.io_dtype.width // 8 idesc_qk = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -1038,20 +1005,20 @@ def _mma_warp( tmem_shared_acc_col = tmem_base + cfg.tmem_shared_acc_offset tmem_shared_inp_col = tmem_base + cfg.tmem_shared_inp_offset SHARED_INP_STAGE_COLS = cfg.b_t // 2 - tmem_dop_col = tmem_shared_inp_col + tmem_do_prime_col = tmem_shared_inp_col tmem_du_col = tmem_shared_inp_col + SHARED_INP_STAGE_COLS tmem_dyp_col = tmem_du_col ACC_STAGE_COLS = cfg.b_t tmem_acc_a = tmem_shared_acc_col tmem_acc_b = tmem_shared_acc_col + ACC_STAGE_COLS tmem_kk_col = tmem_acc_a - tmem_ks_col = tmem_acc_a + tmem_k_state_col = tmem_acc_a tmem_dy_col = tmem_acc_a tmem_dm_core_col = tmem_acc_a tmem_a_col = tmem_acc_b tmem_u_col = tmem_acc_b tmem_da_col = tmem_acc_b - tmem_dk_spath_col = tmem_shared_inp_col + tmem_dk_state_path_col = tmem_shared_inp_col idesc_dv = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -1073,11 +1040,11 @@ def _mma_warp( idesc=idesc_dv, kind=nvvm.Tcgen05MMAKind.F16, ) - tmem_dh_inp_col = tmem_base + cfg.tmem_dh_inp_offset - tmem_dvdk_col = tmem_base + cfg.tmem_dvdk_offset - DH_INP_STAGE_COLS = cfg.d_k // 2 + tmem_dstate_inp_col = tmem_base + cfg.tmem_dstate_inp_offset + tmem_dvdk_acc_col = tmem_base + cfg.tmem_dvdk_acc_offset + DSTATE_INP_STAGE_COLS = cfg.d_k // 2 - idesc_ks = nvvm.Tcgen05InstrDesc.build( + idesc_k_state = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, b_dtype=cfg.io_dtype, @@ -1085,7 +1052,7 @@ def _mma_warp( m_dim=cfg.d_v, a_major=1, ) - bmm_ks_desc = MmaDesc( + bmm_k_state_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -1095,7 +1062,7 @@ def _mma_warp( btranspose=False, atranspose=True, cta_group=1, - idesc=idesc_ks, + idesc=idesc_k_state, kind=nvvm.Tcgen05MMAKind.F16, ) @@ -1121,7 +1088,7 @@ def _mma_warp( idesc=idesc_du, kind=nvvm.Tcgen05MMAKind.F16, ) - idesc_dh = nvvm.Tcgen05InstrDesc.build( + idesc_dstate_upd = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=cfg.io_dtype, b_dtype=cfg.io_dtype, @@ -1129,7 +1096,7 @@ def _mma_warp( m_dim=cfg.d_k, b_major=1, ) - bmm_dh_desc = MmaDesc( + bmm_dstate_upd_desc = MmaDesc( M=cfg.d_k, N=cfg.d_v, K=cfg.b_t, @@ -1139,15 +1106,15 @@ def _mma_warp( btranspose=True, atranspose=False, cta_group=1, - idesc=idesc_dh, + idesc=idesc_dstate_upd, kind=nvvm.Tcgen05MMAKind.F16, ) SHARED_INP_STAGE_COLS = cfg.b_t // 2 - tmem_dop_col = tmem_shared_inp_col + tmem_do_prime_col = tmem_shared_inp_col tmem_du_col = tmem_shared_inp_col + SHARED_INP_STAGE_COLS tmem_dyp_col = tmem_du_col tmem_shared_inp_col = tmem_base + cfg.tmem_shared_inp_offset - tmem_dh_col = tmem_base + cfg.tmem_dh_offset + tmem_dstate_acc_col = tmem_base + cfg.tmem_dstate_acc_offset tmem_y_col = tmem_base + cfg.tmem_y_offset idesc_dy = nvvm.Tcgen05InstrDesc.build( @@ -1211,14 +1178,14 @@ def _mma_warp( idesc=idesc_u, kind=nvvm.Tcgen05MMAKind.F16, ) - idesc_dstate = nvvm.Tcgen05InstrDesc.build( + idesc_dqdk_inter = 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, ) - bmm_dstate_desc = MmaDesc( + bmm_dqdk_inter_desc = MmaDesc( M=cfg.d_k, N=cfg.b_t, K=cfg.d_v, @@ -1228,7 +1195,7 @@ def _mma_warp( btranspose=False, atranspose=False, cta_group=1, - idesc=idesc_dstate, + idesc=idesc_dqdk_inter, kind=nvvm.Tcgen05MMAKind.F16, ) idesc_dka = nvvm.Tcgen05InstrDesc.build( @@ -1275,673 +1242,394 @@ def _mma_warp( kind=nvvm.Tcgen05MMAKind.F16, ) + do_prime_inp_ptr = nvvm.make_tmem_ptr(tmem_do_prime_col, cutlass.Int8) + y_inp_ptr = nvvm.make_tmem_ptr(tmem_y_col, cutlass.Int8) + du_inp_ptr = nvvm.make_tmem_ptr(tmem_du_col, cutlass.Int8) + dyp_inp_ptr = nvvm.make_tmem_ptr(tmem_dyp_col, cutlass.Int8) + kk_acc_ptr = nvvm.make_tmem_ptr(tmem_kk_col, cutlass.Float32) + a_acc_ptr = nvvm.make_tmem_ptr(tmem_a_col, cutlass.Float32) + k_state_acc_ptr = nvvm.make_tmem_ptr(tmem_k_state_col, cutlass.Float32) + u_acc_ptr = nvvm.make_tmem_ptr(tmem_u_col, cutlass.Float32) + dy_acc_ptr = nvvm.make_tmem_ptr(tmem_dy_col, cutlass.Float32) + da_acc_ptr = nvvm.make_tmem_ptr(tmem_da_col, cutlass.Float32) + dm_core_acc_ptr = nvvm.make_tmem_ptr(tmem_dm_core_col, cutlass.Float32) + dk_state_path_acc_ptr = nvvm.make_tmem_ptr(tmem_dk_state_path_col, cutlass.Float32) + dstate_acc_ptr = nvvm.make_tmem_ptr(tmem_dstate_acc_col, cutlass.Float32) + dq_acc_ptr = nvvm.make_tmem_ptr(tmem_dstate_inp_col, cutlass.Float32) + dvdk_acc_ptr = nvvm.make_tmem_ptr(tmem_dvdk_acc_col, cutlass.Float32) + + # ---- warp-top descriptors (1-stage tiles are loop-constant; K advances) ---- + d_q0 = sQ[0].desc() + d_k0 = sK[0].desc() + d_k_trans0 = sK_trans[0].desc() + d_q_trans0 = sQ_trans[0].desc() + d_do0 = sdO[0].desc() + d_do_kmaj0 = sdO_kmaj[0].desc() + d_state0 = sState[0].desc() + d_state_kmaj0 = sState_kmaj[0].desc() + d_tinv0 = sTinv[0].desc() + d_tinv_trans0 = sTinv_trans[0].desc() + d_a_trans0 = sA_trans[0].desc() + d_v_kmaj0 = sV_kmaj[0].desc() + d_dv_kmaj0 = sdV_kmaj[0].desc() + d_da0 = sDa[0].desc() + d_da_trans0 = sDa_trans[0].desc() + d_dm0 = sDm[0].desc() + d_dm_trans0 = sDm_trans[0].desc() + d_dstate0 = sDstate[0].desc() + K_STAGE_BYTES = (cfg.k_cosize // cfg.smem_k_stages) * (cfg.io_dtype.width // 8) + sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) - S_MIN = 0 if cfg.use_initial_state else 1 - SFIRST_MIN = 1 if cfg.use_initial_state else 2 + FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 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, cu_seqlens, mWorkItems - ) - sk_nt = cend - wstart - - # ---- first backward chunk (c = NT-1): no dH yet; with a dht seed - # every chunk takes the steady path ---- - if cutlass.const_expr(not cfg.use_dht): - if sk_nt > 0: - k_idx = k_index.idx - bars.mb_k_ready[k_idx].wait(k_index.phase) - k_index = advance(k_index, cfg.smem_k_stages) - - desc_k = sK[k_idx].desc() - mma_ss( - bmm_qk_desc, - desc_k, - desc_k, - nvvm.make_tmem_ptr(tmem_kk_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_kk_acc_ready[0].arrive(cta_group=1) - - q_idx = q_index.idx - bars.mb_q_ready[q_idx].wait(q_index.phase) - q_index = advance(q_index, cfg.smem_q_stages) - - desc_q = sQ[q_idx].desc() - mma_ss( - bmm_qk_desc, - desc_q, - desc_k, - nvvm.make_tmem_ptr(tmem_a_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_a_acc_ready[0].arrive(cta_group=1) - - # ---- ks (K @ S) first, then dQ inter (S @ dO) ---------------------- - s_idx = s_index.idx - if cend >= SFIRST_MIN: - bars.mb_s_ready[s_idx].wait(s_index.phase) - s_index = advance(s_index, cfg.smem_s_stages) - bars.mb_kk_acc_done[0].wait(kk_read.phase) - kk_read = advance(kk_read, 1) - desc_s = sS[s_idx].desc() - mma_ss( - bmm_ks_desc, - desc_s, - desc_k, - nvvm.make_tmem_ptr(tmem_ks_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_ks_acc_ready[0].arrive(cta_group=1) - do_idx = do_index.idx - bars.mb_do_ready[do_idx].wait(do_index.phase) - do_index = advance(do_index, cfg.smem_do_stages) - if cend >= SFIRST_MIN: - bars.mb_dq_acc_total_done[0].wait(dq_total_index.phase) - dq_total_index = advance(dq_total_index, 1) - desc_s_k = sS_kmaj[s_idx].desc() - desc_do_k2 = sdO_kmaj[do_idx].desc() - mma_ss( - bmm_dstate_desc, - desc_s_k, - desc_do_k2, - nvvm.make_tmem_ptr(tmem_dh_inp_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_dq_acc_scale_ready[0].arrive(cta_group=1) - - # ---- dU intra: dO^T @ A -> the dV/dK slot ------------------ - du_qk_idx = qk_index.idx - bars.mb_qk_ready[du_qk_idx].wait(qk_index.phase) - qk_index = advance(qk_index, cfg.smem_qk_stages) - bars.mb_dk_total_done[0].wait(dk_total_free.phase) - dk_total_free = advance(dk_total_free, 1) - - desc_do_mn = sdO[do_idx].desc() - desc_qk_t = sQk_trans[du_qk_idx].desc() - mma_ss( - bmm_du_desc, - desc_do_mn, - desc_qk_t, - nvvm.make_tmem_ptr(tmem_dvdk_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_du_total_ready[0].arrive(cta_group=1) - - # ---- dH init: dO'^T @ Q ------------------------------------- - bars.mb_dop_inp_ready[0].wait(dop_inp_rdy.phase) - dop_inp_rdy = advance(dop_inp_rdy, 1) - dh_idx = dh_acc_index.idx - bars.mb_dh_acc_done[dh_idx].wait(dh_acc_index.phase) - dh_acc_index = advance(dh_acc_index, cfg.tmem_dh_acc_stages) - - do_a_ptr = nvvm.make_tmem_ptr( - tmem_dop_col, - cutlass.Int8, - ) - desc_q_t = sQ_trans[q_idx].desc() - mma_ts( - bmm_dh_desc, - do_a_ptr, - desc_q_t, - nvvm.make_tmem_ptr(tmem_dh_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_dop_inp_done[0].arrive(cta_group=1) - - # ---- Y operand acquire: y_ready = CG1's delta -------------- - bars.mb_y_ready[0].wait(y_rdy_index.phase) - y_rdy_index = advance(y_rdy_index, 1) - v_idx = v_index.idx - v_index = advance(v_index, cfg.smem_v_stages) - - # ---- U recompute: U^T = Y^T @ T^T -> shared acc ------------ - ainv_idx = ainv_index.idx - bars.mb_ainv_ready[ainv_idx].wait(ainv_index.phase) - ainv_index = advance(ainv_index, cfg.smem_ainv_stages) - - y_a_ptr = nvvm.make_tmem_ptr(tmem_y_col, cutlass.Int8) - desc_t_plain = sAinv[ainv_idx].desc() - mma_ts( - bmm_u_desc, - y_a_ptr, - desc_t_plain, - nvvm.make_tmem_ptr(tmem_u_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_u_acc_ready[0].arrive(cta_group=1) - - # ---- dY: dU^T(TMEM f16) @ T --------------------------------- - bars.mb_du_inp_ready[0].wait(du_inp_rdy.phase) - du_inp_rdy = advance(du_inp_rdy, 1) - - du_a_ptr = nvvm.make_tmem_ptr( - tmem_du_col, - cutlass.Int8, - ) - desc_ainv_t = sAinv_trans[ainv_idx].desc() - mma_ts( - bmm_dy_desc, - du_a_ptr, - desc_ainv_t, - nvvm.make_tmem_ptr(tmem_dy_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_dy_acc_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_ainv_done[ainv_idx].arrive(cta_group=1) - - # ---- dA_eff: dO @ U^T -> shared acc ------------------------- - bars.mb_u_ready[0].wait(u_index.phase) - u_index = advance(u_index, 1) - - desc_do_k = sdO_kmaj[do_idx].desc() - desc_u = sV_kmaj[v_idx].desc() - mma_ss( - bmm_da_desc, - desc_do_k, - desc_u, - nvvm.make_tmem_ptr(tmem_da_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_da_acc_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_do_mma_done[do_idx].arrive(cta_group=1) - - # ---- dM core: dY @ U^T -------------------------------------- - dv_stg_x_idx = dv_rdy_index.idx - bars.mb_dv_tmastg_ready[dv_stg_x_idx].wait(dv_rdy_index.phase) - dv_rdy_index = advance(dv_rdy_index, cfg.smem_dv_stages) - - desc_dy_dm = sdV_kmaj[0].desc() - desc_u_dm = sV_kmaj[v_idx].desc() - mma_ss( - bmm_da_desc, - desc_dy_dm, - desc_u_dm, - nvvm.make_tmem_ptr(tmem_dm_core_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_dm_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_v_mma_done[v_idx].arrive(cta_group=1) - - # ---- dH ds-term: dH += dY'^T @ K ---------------------------- - bars.mb_dyp_inp_ready[0].wait(dyp_inp_rdy.phase) - dyp_inp_rdy = advance(dyp_inp_rdy, 1) - - dyp_a_ptr = nvvm.make_tmem_ptr( - tmem_dyp_col, - cutlass.Int8, - ) - desc_k_t = sK_trans[k_idx].desc() - mma_ts( - bmm_dh_desc, - dyp_a_ptr, - desc_k_t, - nvvm.make_tmem_ptr(tmem_dh_col, cutlass.Float32), - accumulate=True, - ) - if nvvm.elect_sync(): - bars.mb_dh_acc_ready[dh_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_dyp_inp_done[0].arrive(cta_group=1) - - # ---- dK attn: Q^T @ dA(sDa) -> the dV/dK slot -------------- - bars.mb_da_ready[0].wait(da_rdy_index.phase) - da_rdy_index = advance(da_rdy_index, 1) - - desc_q_mn = sQ_trans[q_idx].desc() - desc_da_t = sDa_trans[0].desc() - mma_ss( - bmm_dka_desc, - desc_q_mn, - desc_da_t, - nvvm.make_tmem_ptr(tmem_dvdk_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_dk_attn_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_q_mma_done[q_idx].arrive(cta_group=1) - - # ---- dQ attn: K^T @ dA onto the rescaled dQ acc ------------------- - if cend >= SFIRST_MIN: - bars.mb_dq_acc_scale_done[0].wait(dq_scale_index.phase) - dq_scale_index = advance(dq_scale_index, 1) - if cend < SFIRST_MIN: - bars.mb_dq_acc_total_done[0].wait(dq_total_index.phase) - dq_total_index = advance(dq_total_index, 1) - desc_k_mn = sK_trans[k_idx].desc() - desc_da_p = sDa[0].desc() - if cend >= SFIRST_MIN: - mma_ss( - bmm_dqa_desc, - desc_k_mn, - desc_da_p, - nvvm.make_tmem_ptr(tmem_dh_inp_col, cutlass.Float32), - accumulate=True, - ) - if cend < SFIRST_MIN: - mma_ss( - bmm_dqa_desc, - desc_k_mn, - desc_da_p, - nvvm.make_tmem_ptr(tmem_dh_inp_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_dq_acc_total_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_qk_done[du_qk_idx].arrive(cta_group=1) - - # ---- dK s-path: S @ dY^T -> shared acc, before the dM-terms -------- - if cend >= SFIRST_MIN: - desc_s_k5 = sS_kmaj[s_idx].desc() - desc_dy_k5 = sdV_kmaj[0].desc() - mma_ss( - bmm_dstate_desc, - desc_s_k5, - desc_dy_k5, - nvvm.make_tmem_ptr(tmem_dk_spath_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_dk_spath_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_s_done[s_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_sdv_done[0].arrive(cta_group=1) - - # ---- dK dM-terms: K^T @ dM^T + K^T @ dM onto the slot ------ - bars.mb_dm_done[0].wait(dm_index.phase) - dm_index = advance(dm_index, 1) - bars.mb_dk_attn_done[0].wait(dk_attn_read.phase) - dk_attn_read = advance(dk_attn_read, 1) - desc_k_mn2 = sK_trans[k_idx].desc() - desc_dm_t = sDm_trans[0].desc() - mma_ss( - bmm_dka_desc, - desc_k_mn2, - desc_dm_t, - nvvm.make_tmem_ptr(tmem_dvdk_col, cutlass.Float32), - accumulate=True, - ) - desc_dm_p = sDm[0].desc() - mma_ss( - bmm_dqa_desc, - desc_k_mn2, - desc_dm_p, - nvvm.make_tmem_ptr(tmem_dvdk_col, cutlass.Float32), - accumulate=True, - ) - if nvvm.elect_sync(): - bars.mb_dk_total_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_k_mma_done[k_idx].arrive(cta_group=1) + 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 - # ---- chunks NT-2 .. 0 (backward): full body ---------------------- - for rev_idx in cutlass.range(0 if cfg.use_dht else 1, sk_nt): + # ---- chunks NT-2 .. 0 (backward): full body ------------------------------ + for rev_idx in cutlass.range(num_item_chunks): chunk_idx = cend - 1 - rev_idx - # ---- kk: K @ K^T -> shared acc ------------------------------ + have_dstate = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_dstate_in) else rev_idx > 0 + + # ---- KK = K(S) @ K^T ------------------------------------------------- k_idx = k_index.idx bars.mb_k_ready[k_idx].wait(k_index.phase) k_index = advance(k_index, cfg.smem_k_stages) - desc_k = sK[k_idx].desc() + desc_k = d_k0.advance_start_address(k_idx * K_STAGE_BYTES) mma_ss( bmm_qk_desc, desc_k, desc_k, - nvvm.make_tmem_ptr(tmem_kk_col, cutlass.Float32), + kk_acc_ptr, accumulate=False, ) - if nvvm.elect_sync(): + if elect_one: bars.mb_kk_acc_ready[0].arrive(cta_group=1) - # ---- qk: Q @ K^T -> shared acc ------------------------------ + # ---- QK = Q(S) @ K^T ------------------------------------------------- q_idx = q_index.idx bars.mb_q_ready[q_idx].wait(q_index.phase) q_index = advance(q_index, cfg.smem_q_stages) - desc_q = sQ[q_idx].desc() + desc_q = d_q0 mma_ss( bmm_qk_desc, desc_q, desc_k, - nvvm.make_tmem_ptr(tmem_a_col, cutlass.Float32), + a_acc_ptr, accumulate=False, ) - if nvvm.elect_sync(): + if elect_one: bars.mb_a_acc_ready[0].arrive(cta_group=1) - # ---- ks (K @ S): hoisted above dV inter so CG1's delta chain - # is not serialized behind the prev chunk's dK readout ---------- - s_idx = s_index.idx - if chunk_idx >= S_MIN: - bars.mb_s_ready[s_idx].wait(s_index.phase) - s_index = advance(s_index, cfg.smem_s_stages) - bars.mb_kk_acc_done[0].wait(kk_read.phase) - kk_read = advance(kk_read, 1) - desc_s = sS[s_idx].desc() + # ---- k_state = state^T(S) @ K^T ----------------------------------------------- + state_idx = state_index.idx + if chunk_idx >= FIRST_STATE_CHUNK: + bars.mb_state_ready[state_idx].wait(state_index.phase) + 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 mma_ss( - bmm_ks_desc, - desc_s, + bmm_k_state_desc, + desc_state, desc_k, - nvvm.make_tmem_ptr(tmem_ks_col, cutlass.Float32), + k_state_acc_ptr, accumulate=False, ) - if nvvm.elect_sync(): - bars.mb_ks_acc_ready[0].arrive(cta_group=1) - - # ---- dV inter: dH^T @ K -> the dV/dK slot ------------------- - dh_inp_idx = dh_inp_index.idx - bars.mb_dh_inp_ready[dh_inp_idx].wait(dh_inp_index.phase) - dh_inp_index = advance(dh_inp_index, cfg.tmem_dh_inp_stages) - bars.mb_dk_total_done[0].wait(dk_total_free.phase) - dk_total_free = advance(dk_total_free, 1) - - dh_a_ptr = nvvm.make_tmem_ptr(tmem_dh_inp_col + dh_inp_idx * DH_INP_STAGE_COLS, cutlass.Int8) - mma_ts( - bmm_dv_desc, - dh_a_ptr, - desc_k, - nvvm.make_tmem_ptr(tmem_dvdk_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_du_scale_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_dh_inp_done[dh_inp_idx].arrive(cta_group=1) + if elect_one: + bars.mb_k_state_acc_ready[0].arrive(cta_group=1) + + # ---- dV inter = dstate^T(T) @ K ------------------------------------------ + dstate_inp_idx = dstate_inp_index.idx + if have_dstate: + bars.mb_dstate_inp_ready[dstate_inp_idx].wait(dstate_inp_index.phase) + dstate_inp_index = advance(dstate_inp_index, cfg.tmem_dstate_inp_stages) + bars.mb_dk_total_acc_done[0].wait(dk_total_index.phase) + dk_total_index = advance(dk_total_index, 1) + + if have_dstate: + dstate_a_ptr = nvvm.make_tmem_ptr(tmem_dstate_inp_col + dstate_inp_idx * DSTATE_INP_STAGE_COLS, cutlass.Int8) + for sub in cutlass.range_constexpr(bmm_dv_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dv_desc.sps_B): + mma_ts_step( + bmm_dv_desc, + dstate_a_ptr.subview(sub * bmm_dv_desc.sps_B * bmm_dv_desc.tmem_advance_A), + desc_k + sub * (bmm_dv_desc.smem_subtile_B >> 4), + dvdk_acc_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + if elect_one: + bars.mb_du_scale_acc_ready[0].arrive(cta_group=1) + bars.mb_dstate_inp_done[dstate_inp_idx].arrive(cta_group=1) + # ---- dQ inter = state(S) @ dO^T ------------------------------------------ do_idx = do_index.idx bars.mb_do_ready[do_idx].wait(do_index.phase) do_index = advance(do_index, cfg.smem_do_stages) - if chunk_idx >= S_MIN: + 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_s_k = sS_kmaj[s_idx].desc() - desc_do_k2 = sdO_kmaj[do_idx].desc() + desc_state_kmaj_dq_inter = d_state_kmaj0 + desc_do_kmaj_dq_inter = d_do_kmaj0 mma_ss( - bmm_dstate_desc, - desc_s_k, - desc_do_k2, - nvvm.make_tmem_ptr(tmem_dh_inp_col, cutlass.Float32), + bmm_dqdk_inter_desc, + desc_state_kmaj_dq_inter, + desc_do_kmaj_dq_inter, + dq_acc_ptr, accumulate=False, ) - if nvvm.elect_sync(): + if elect_one: bars.mb_dq_acc_scale_ready[0].arrive(cta_group=1) - # ---- dU intra: dO^T @ A accumulated onto the scaled dV inter ---- - du_qk_idx = qk_index.idx - bars.mb_qk_ready[du_qk_idx].wait(qk_index.phase) - qk_index = advance(qk_index, cfg.smem_qk_stages) - bars.mb_du_scale_done[0].wait(du_scale_read.phase) - du_scale_read = advance(du_scale_read, 1) + # ---- dU intra += dO^T(S) @ A ----------------------------------------- + du_a_idx = a_index.idx + bars.mb_a_ready[du_a_idx].wait(a_index.phase) + a_index = advance(a_index, cfg.smem_a_stages) + if have_dstate: + bars.mb_du_scale_acc_done[0].wait(du_scale_index.phase) + du_scale_index = advance(du_scale_index, 1) - desc_do_mn = sdO[do_idx].desc() - desc_qk_t = sQk_trans[du_qk_idx].desc() + desc_do_mnmaj = d_do0 + desc_a_t = d_a_trans0 mma_ss( bmm_du_desc, - desc_do_mn, - desc_qk_t, - nvvm.make_tmem_ptr(tmem_dvdk_col, cutlass.Float32), - accumulate=True, - ) - if nvvm.elect_sync(): - bars.mb_du_total_ready[0].arrive(cta_group=1) - - # ---- dH update: dO'^T @ Q ----------------------------------- - bars.mb_dop_inp_ready[0].wait(dop_inp_rdy.phase) - dop_inp_rdy = advance(dop_inp_rdy, 1) - dh_idx = dh_acc_index.idx - bars.mb_dh_acc_done[dh_idx].wait(dh_acc_index.phase) - dh_acc_index = advance(dh_acc_index, cfg.tmem_dh_acc_stages) - - do_a_ptr = nvvm.make_tmem_ptr( - tmem_dop_col, - cutlass.Int8, - ) - desc_q_t = sQ_trans[q_idx].desc() - mma_ts( - bmm_dh_desc, - do_a_ptr, - desc_q_t, - nvvm.make_tmem_ptr(tmem_dh_col, cutlass.Float32), - accumulate=True, + desc_do_mnmaj, + desc_a_t, + dvdk_acc_ptr, + accumulate=have_dstate, ) - if nvvm.elect_sync(): - bars.mb_dop_inp_done[0].arrive(cta_group=1) + if elect_one: + bars.mb_du_total_acc_ready[0].arrive(cta_group=1) + + # ---- dstate update += dO'^T(T) @ Q --------------------------------------- + bars.mb_do_prime_inp_ready[0].wait(do_prime_inp_ready.phase) + do_prime_inp_ready = advance(do_prime_inp_ready, 1) + dstate_idx = dstate_acc_index.idx + bars.mb_dstate_scale_acc_done[dstate_idx].wait(dstate_acc_index.phase) + dstate_acc_index = advance(dstate_acc_index, cfg.tmem_dstate_acc_stages) + + desc_q_t = d_q_trans0 + for sub in cutlass.range_constexpr(bmm_dstate_upd_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_upd_desc.sps_B): + mma_ts_step( + bmm_dstate_upd_desc, + do_prime_inp_ptr.subview(sub * bmm_dstate_upd_desc.sps_B * bmm_dstate_upd_desc.tmem_advance_A), + desc_q_t + sub * (bmm_dstate_upd_desc.smem_subtile_B >> 4), + dstate_acc_ptr, + k, + cutlass.Boolean(True) if cutlass.const_expr(sub + k > 0) else have_dstate, + ) - # ---- Y operand acquire: y_ready = CG1's delta -------------- - bars.mb_y_ready[0].wait(y_rdy_index.phase) - y_rdy_index = advance(y_rdy_index, 1) + # ---- U^T recompute = Y^T(T) @ T^T ------------------------------------ + tinv_idx = tinv_index.idx + bars.mb_t_inv_ready[tinv_idx].wait(tinv_index.phase) + tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) + bars.mb_y_ready[0].wait(y_ready_index.phase) + y_ready_index = advance(y_ready_index, 1) v_idx = v_index.idx v_index = advance(v_index, cfg.smem_v_stages) - # ---- U recompute: U^T = Y^T @ T^T -> shared acc ------------ - ainv_idx = ainv_index.idx - bars.mb_ainv_ready[ainv_idx].wait(ainv_index.phase) - ainv_index = advance(ainv_index, cfg.smem_ainv_stages) - - y_a_ptr = nvvm.make_tmem_ptr(tmem_y_col, cutlass.Int8) - desc_t_plain = sAinv[ainv_idx].desc() - mma_ts( - bmm_u_desc, - y_a_ptr, - desc_t_plain, - nvvm.make_tmem_ptr(tmem_u_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): + desc_tinv = d_tinv0 + for sub in cutlass.range_constexpr(bmm_u_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_u_desc.sps_B): + mma_ts_step( + bmm_u_desc, + y_inp_ptr.subview(sub * bmm_u_desc.sps_B * bmm_u_desc.tmem_advance_A), + desc_tinv + sub * (bmm_u_desc.smem_subtile_B >> 4), + u_acc_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + if elect_one: bars.mb_u_acc_ready[0].arrive(cta_group=1) - # ---- dY: dU^T(TMEM f16) @ T --------------------------------- - bars.mb_du_inp_ready[0].wait(du_inp_rdy.phase) - du_inp_rdy = advance(du_inp_rdy, 1) - - du_a_ptr = nvvm.make_tmem_ptr( - tmem_du_col, - cutlass.Int8, - ) - desc_ainv_t = sAinv_trans[ainv_idx].desc() - mma_ts( - bmm_dy_desc, - du_a_ptr, - desc_ainv_t, - nvvm.make_tmem_ptr(tmem_dy_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): + # ---- dY = dU^T(T) @ T ------------------------------------------------ + bars.mb_du_inp_ready[0].wait(du_inp_ready.phase) + du_inp_ready = advance(du_inp_ready, 1) + + desc_tinv_t = d_tinv_trans0 + for sub in cutlass.range_constexpr(bmm_dy_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dy_desc.sps_B): + mma_ts_step( + bmm_dy_desc, + du_inp_ptr.subview(sub * bmm_dy_desc.sps_B * bmm_dy_desc.tmem_advance_A), + desc_tinv_t + sub * (bmm_dy_desc.smem_subtile_B >> 4), + dy_acc_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + if elect_one: bars.mb_dy_acc_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_ainv_done[ainv_idx].arrive(cta_group=1) - # ---- dK inter: dH_entry(SMEM) @ U^T -> the dV/dK slot ------- + # ---- dK_inter = dstate_entry(S) @ U^T ------------------------------------ bars.mb_u_ready[0].wait(u_index.phase) u_index = advance(u_index, 1) - bars.mb_dhs_ready[0].wait(dhs_index.phase) - dhs_index = advance(dhs_index, 1) + if have_dstate: + bars.mb_dstate_smem_ready[0].wait(dstate_smem_index.phase) + dstate_smem_index = advance(dstate_smem_index, 1) - desc_dh_k = sdH[0].desc() - desc_u_t = sV_kmaj[v_idx].desc() - mma_ss( - bmm_dstate_desc, - desc_dh_k, - desc_u_t, - nvvm.make_tmem_ptr(tmem_dvdk_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_dk_scale_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_dhs_done[0].arrive(cta_group=1) + desc_dstate_kmaj = d_dstate0 + desc_u_kmaj_dk_inter = d_v_kmaj0 + mma_ss( + bmm_dqdk_inter_desc, + desc_dstate_kmaj, + desc_u_kmaj_dk_inter, + dvdk_acc_ptr, + accumulate=False, + ) + if elect_one: + bars.mb_dk_scale_acc_ready[0].arrive(cta_group=1) - # ---- dA_eff: dO @ U^T -> shared acc (U gated at dK-inter) ---- - desc_do_k = sdO_kmaj[do_idx].desc() - desc_u = sV_kmaj[v_idx].desc() + # ---- dA_eff = dO(S) @ U^T -------------------------------------------- + desc_do_kmaj_da = d_do_kmaj0 + desc_u_kmaj_da = d_v_kmaj0 mma_ss( bmm_da_desc, - desc_do_k, - desc_u, - nvvm.make_tmem_ptr(tmem_da_col, cutlass.Float32), + desc_do_kmaj_da, + desc_u_kmaj_da, + da_acc_ptr, accumulate=False, ) - if nvvm.elect_sync(): + if elect_one: bars.mb_da_acc_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): bars.mb_do_mma_done[do_idx].arrive(cta_group=1) - # ---- dM core: dY @ U^T -------------------------------------- - dv_stg_x_idx = dv_rdy_index.idx - bars.mb_dv_tmastg_ready[dv_stg_x_idx].wait(dv_rdy_index.phase) - dv_rdy_index = advance(dv_rdy_index, cfg.smem_dv_stages) + # ---- dM core = dY(S) @ U^T ------------------------------------------- + dv_ready_idx = dv_ready_index.idx + bars.mb_dv_tmastg_ready[dv_ready_idx].wait(dv_ready_index.phase) + dv_ready_index = advance(dv_ready_index, cfg.smem_dv_stages) - desc_dy_dm = sdV_kmaj[0].desc() - desc_u_dm = sV_kmaj[v_idx].desc() + desc_dy_kmaj_dm = d_dv_kmaj0 + desc_u_kmaj_dm = d_v_kmaj0 mma_ss( bmm_da_desc, - desc_dy_dm, - desc_u_dm, - nvvm.make_tmem_ptr(tmem_dm_core_col, cutlass.Float32), + desc_dy_kmaj_dm, + desc_u_kmaj_dm, + dm_core_acc_ptr, accumulate=False, ) - if nvvm.elect_sync(): - bars.mb_dm_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): + if elect_one: + bars.mb_dm_acc_ready[0].arrive(cta_group=1) bars.mb_v_mma_done[v_idx].arrive(cta_group=1) - # ---- dH ds-term: dH += dY'^T @ K ---------------------------- - bars.mb_dyp_inp_ready[0].wait(dyp_inp_rdy.phase) - dyp_inp_rdy = advance(dyp_inp_rdy, 1) - - dyp_a_ptr = nvvm.make_tmem_ptr( - tmem_dyp_col, - cutlass.Int8, - ) - desc_k_t = sK_trans[k_idx].desc() - mma_ts( - bmm_dh_desc, - dyp_a_ptr, - desc_k_t, - nvvm.make_tmem_ptr(tmem_dh_col, cutlass.Float32), - accumulate=True, - ) - if nvvm.elect_sync(): - bars.mb_dh_acc_ready[dh_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_dyp_inp_done[0].arrive(cta_group=1) - - # ---- dK attn: Q^T @ dA(sDa) onto the rescaled dK inter ------ - bars.mb_da_ready[0].wait(da_rdy_index.phase) - da_rdy_index = advance(da_rdy_index, 1) - bars.mb_dk_scale_done[0].wait(dk_scale_read.phase) - dk_scale_read = advance(dk_scale_read, 1) - - desc_q_mn = sQ_trans[q_idx].desc() - desc_da_t = sDa_trans[0].desc() + # ---- dstate update-term += dY'^T(T) @ K -------------------------------------- + bars.mb_dyp_inp_ready[0].wait(dyp_inp_ready.phase) + dyp_inp_ready = advance(dyp_inp_ready, 1) + + desc_k_t = d_k_trans0.advance_start_address(k_idx * K_STAGE_BYTES) + for sub in cutlass.range_constexpr(bmm_dstate_upd_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_upd_desc.sps_B): + mma_ts_step( + bmm_dstate_upd_desc, + dyp_inp_ptr.subview(sub * bmm_dstate_upd_desc.sps_B * bmm_dstate_upd_desc.tmem_advance_A), + desc_k_t + sub * (bmm_dstate_upd_desc.smem_subtile_B >> 4), + dstate_acc_ptr, + k, + cutlass.Boolean(True), + ) + if elect_one: + bars.mb_dstate_acc_ready[dstate_idx].arrive(cta_group=1) + + # ---- dK attn += Q^T(S) @ dA ------------------------------------------ + bars.mb_da_ready[0].wait(da_ready_index.phase) + da_ready_index = advance(da_ready_index, 1) + if have_dstate: + bars.mb_dk_scale_acc_done[0].wait(dk_scale_index.phase) + dk_scale_index = advance(dk_scale_index, 1) + + desc_q_mnmaj_dk_attn = d_q_trans0 + desc_da_t = d_da_trans0 mma_ss( bmm_dka_desc, - desc_q_mn, + desc_q_mnmaj_dk_attn, desc_da_t, - nvvm.make_tmem_ptr(tmem_dvdk_col, cutlass.Float32), - accumulate=True, + dvdk_acc_ptr, + accumulate=have_dstate, ) - if nvvm.elect_sync(): - bars.mb_dk_attn_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): + if elect_one: + bars.mb_dk_attn_acc_ready[0].arrive(cta_group=1) bars.mb_q_mma_done[q_idx].arrive(cta_group=1) - # ---- dQ attn: K^T @ dA onto the rescaled dQ acc ------------------- - if chunk_idx >= S_MIN: + # ---- dQ attn += K^T(S) @ dA ------------------------------------------ + if chunk_idx >= FIRST_STATE_CHUNK: bars.mb_dq_acc_scale_done[0].wait(dq_scale_index.phase) dq_scale_index = advance(dq_scale_index, 1) - if chunk_idx < S_MIN: + 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_k_mn = sK_trans[k_idx].desc() - desc_da_p = sDa[0].desc() - if chunk_idx >= S_MIN: + desc_k_mnmaj_dq_attn = d_k_trans0.advance_start_address(k_idx * K_STAGE_BYTES) + desc_da = d_da0 + if chunk_idx >= FIRST_STATE_CHUNK: mma_ss( bmm_dqa_desc, - desc_k_mn, - desc_da_p, - nvvm.make_tmem_ptr(tmem_dh_inp_col, cutlass.Float32), + desc_k_mnmaj_dq_attn, + desc_da, + dq_acc_ptr, accumulate=True, ) - if chunk_idx < S_MIN: + if chunk_idx < FIRST_STATE_CHUNK: mma_ss( bmm_dqa_desc, - desc_k_mn, - desc_da_p, - nvvm.make_tmem_ptr(tmem_dh_inp_col, cutlass.Float32), + desc_k_mnmaj_dq_attn, + desc_da, + dq_acc_ptr, accumulate=False, ) - if nvvm.elect_sync(): + if elect_one: bars.mb_dq_acc_total_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_qk_done[du_qk_idx].arrive(cta_group=1) + bars.mb_a_done[du_a_idx].arrive(cta_group=1) - # ---- dK s-path: S @ dY^T -> shared acc, before the dM-terms -------- - if chunk_idx >= S_MIN: - desc_s_k5 = sS_kmaj[s_idx].desc() - desc_dy_k5 = sdV_kmaj[0].desc() + # ---- dK state-path = state(S) @ dY^T ----------------------------------------- + if chunk_idx >= FIRST_STATE_CHUNK: + desc_state_kmaj_spath = d_state_kmaj0 + desc_dy_kmaj_spath = d_dv_kmaj0 mma_ss( - bmm_dstate_desc, - desc_s_k5, - desc_dy_k5, - nvvm.make_tmem_ptr(tmem_dk_spath_col, cutlass.Float32), + bmm_dqdk_inter_desc, + desc_state_kmaj_spath, + desc_dy_kmaj_spath, + dk_state_path_acc_ptr, accumulate=False, ) - if nvvm.elect_sync(): - bars.mb_dk_spath_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_s_done[s_idx].arrive(cta_group=1) - if nvvm.elect_sync(): + if elect_one: + bars.mb_dk_state_path_acc_ready[0].arrive(cta_group=1) + bars.mb_state_mma_done[state_idx].arrive(cta_group=1) + if elect_one: bars.mb_sdv_done[0].arrive(cta_group=1) - # ---- dK dM-terms: K^T @ dM^T + K^T @ dM onto the slot ------- - bars.mb_dm_done[0].wait(dm_index.phase) + # ---- dK dM-terms += K^T(S) @ dM^T + K^T(S) @ dM ---------------------- + bars.mb_dm_acc_done[0].wait(dm_index.phase) dm_index = advance(dm_index, 1) - bars.mb_dk_attn_done[0].wait(dk_attn_read.phase) - dk_attn_read = advance(dk_attn_read, 1) - desc_k_mn2 = sK_trans[k_idx].desc() - desc_dm_t = sDm_trans[0].desc() + bars.mb_dk_attn_acc_done[0].wait(dk_attn_index.phase) + dk_attn_index = advance(dk_attn_index, 1) + desc_k_mnmaj_dm_terms = d_k_trans0.advance_start_address(k_idx * K_STAGE_BYTES) + desc_dm_t = d_dm_trans0 mma_ss( bmm_dka_desc, - desc_k_mn2, + desc_k_mnmaj_dm_terms, desc_dm_t, - nvvm.make_tmem_ptr(tmem_dvdk_col, cutlass.Float32), + dvdk_acc_ptr, accumulate=True, ) - desc_dm_p = sDm[0].desc() + desc_dm = d_dm0 mma_ss( bmm_dqa_desc, - desc_k_mn2, - desc_dm_p, - nvvm.make_tmem_ptr(tmem_dvdk_col, cutlass.Float32), + desc_k_mnmaj_dm_terms, + desc_dm, + dvdk_acc_ptr, accumulate=True, ) - if nvvm.elect_sync(): - bars.mb_dk_total_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): + if elect_one: + bars.mb_dk_total_acc_ready[0].arrive(cta_group=1) bars.mb_k_mma_done[k_idx].arrive(cta_group=1) - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) - bars.mb_dk_total_done[0].wait(dk_total_free.phase) + bars.mb_dk_total_acc_done[0].wait(dk_total_index.phase) bars.mb_tmem_done[0].wait(0) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) @@ -1953,7 +1641,7 @@ def _mma_warp( @cute.jit -def _tmaldg_warp( +def tmaldg_warp( cfg, total_tiles, bidx, @@ -1964,36 +1652,33 @@ def _tmaldg_warp( sK_raw, sV_raw, sdO_raw, - sS_raw, + sState_raw, desc_q_base, desc_k_base, desc_v_base, desc_do_base, - desc_s_base, - desc_s0_base, + desc_checkpoint_base, + desc_initial_state_base, sSched, mSched, bars, ): - """TMA-LDG warp role (warp 9): persistent tile loop + per-chunk (BACKWARD - order) Q/K/V/dO loads, plus S = H[c-1] for chunks c >= S_MIN (with an - initial state, chunk 0's S comes from the io-dtype S0 descriptors - instead). The per-(b,h) runtime descriptors fold the sequence start and - head, so Q/K/V/dO load at token ``chunk_idx*BT`` and S at the - sequence-local index ``c - 1``.""" + """TMA-LDG warp role (warp 9): persistent tile loop issuing every + Q/K/V/dO/state TMA load.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) q_index = PipelineState.start(phase=1) k_index = PipelineState.start(phase=1) v_index = PipelineState.start(phase=1) do_index = PipelineState.start(phase=1) - s_index = PipelineState.start(phase=1) + state_index = PipelineState.start(phase=1) sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) - S_MIN = 0 if cfg.use_initial_state else 1 + 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 - granu = 128 // bpe + granule_elems = 128 // bpe bt = cfg.b_t q_stage_elems = cfg.q_cosize // cfg.smem_q_stages k_stage_elems = cfg.k_cosize // cfg.smem_k_stages @@ -2005,8 +1690,8 @@ def _tmaldg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, - tma_subtile_stride_elems=bt * granu, + tma_granu_elems=granule_elems, + tma_subtile_stride_elems=bt * granule_elems, ) sK_tma = SmemTile( base=sK_raw, @@ -2016,8 +1701,8 @@ def _tmaldg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, - tma_subtile_stride_elems=bt * granu, + tma_granu_elems=granule_elems, + tma_subtile_stride_elems=bt * granule_elems, ) sV_tma = SmemTile( base=sV_raw, @@ -2027,7 +1712,7 @@ def _tmaldg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, + tma_granu_elems=granule_elems, tma_subtile_stride_elems=4096, ) sdO_tma = SmemTile( @@ -2038,13 +1723,13 @@ def _tmaldg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, + tma_granu_elems=granule_elems, tma_subtile_stride_elems=4096, ) - sS_tma = SmemTile( - base=sS_raw, - elems_per_stage=(cfg.s_cosize // cfg.smem_s_stages), - stages=cfg.smem_s_stages, + sCheckpoint_tma = SmemTile( + base=sState_raw, + elems_per_stage=(cfg.state_cosize // cfg.smem_state_stages), + stages=cfg.smem_state_stages, leading_byte_offset=0, stride_byte_offset=0, layout=0, @@ -2056,101 +1741,90 @@ def _tmaldg_warp( desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) 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, cu_seqlens, mWorkItems - ) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - slot = (batch_idx * heads_out + head_idx) * desc_qwords + head_o = head_idx + head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) + head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) + head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) + slot = batch_idx * desc_qwords desc_q_slot = (desc_q_base + slot).tospace(cutlass.AddressSpace.generic) desc_k_slot = (desc_k_base + slot).tospace(cutlass.AddressSpace.generic) desc_v_slot = (desc_v_base + slot).tospace(cutlass.AddressSpace.generic) desc_do_slot = (desc_do_base + slot).tospace(cutlass.AddressSpace.generic) - desc_s_slot = (desc_s_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_s0_slot = (desc_s0_base + slot).tospace(cutlass.AddressSpace.generic) - if nvvm.elect_sync(): + 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_s_slot) + tma_tensormap_acquire(desc_checkpoint_slot) if cutlass.const_expr(cfg.use_initial_state): - tma_tensormap_acquire(desc_s0_slot) + tma_tensormap_acquire(desc_initial_state_slot) for rev_idx in cutlass.range(cend - wstart): chunk_idx = cend - 1 - rev_idx tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) - # ---------------------------------------------------------- - # K (B operand of GEMMs 1/2/3, double-buffered) - # ---------------------------------------------------------- + # ---- K load ---------------------------------------------------------- k_idx = k_index.idx bars.mb_k_mma_done[k_idx].wait(k_index.phase) bars.mb_k_cg0_done[k_idx].wait(k_index.phase) k_index = advance(k_index, cfg.smem_k_stages) - if nvvm.elect_sync(): + if elect_one: bars.mb_k_ready[k_idx].arrive(n_bytes=cfg.tma_k_bytes) - k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), tok_coord) + k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, tok_coord) tma_load_tile(sK_tma[k_idx], k_slice, bars.mb_k_ready[k_idx].smem_ptr, acquire=False) - # ---------------------------------------------------------- - # Q (A operand of the qk GEMM, single-buffered) - # ---------------------------------------------------------- + # ---- Q load ---------------------------------------------------------- q_idx = q_index.idx bars.mb_q_mma_done[q_idx].wait(q_index.phase) bars.mb_q_cg1_done[q_idx].wait(q_index.phase) q_index = advance(q_index, cfg.smem_q_stages) - if nvvm.elect_sync(): + if elect_one: bars.mb_q_ready[q_idx].arrive(n_bytes=cfg.tma_q_bytes) - q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), tok_coord) + q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, tok_coord) tma_load_tile(sQ_tma[q_idx], q_slice, bars.mb_q_ready[q_idx].smem_ptr, acquire=False) - # ---------------------------------------------------------- - # V (transposed [DV, T] descriptor) - # ---------------------------------------------------------- + # ---- V load ---------------------------------------------------------- v_idx = v_index.idx bars.mb_v_mma_done[v_idx].wait(v_index.phase) - bars.mb_v_cg1_done[v_idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) - if nvvm.elect_sync(): + if elect_one: bars.mb_v_ready[v_idx].arrive(n_bytes=cfg.tma_v_bytes) - v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), tok_coord) + v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), head_v, tok_coord) tma_load_tile(sV_tma[v_idx], v_slice, bars.mb_v_ready[v_idx].smem_ptr, acquire=False) - # ---------------------------------------------------------- - # dO (transposed [DV, T] descriptor, like V) - # ---------------------------------------------------------- + # ---- dO load --------------------------------------------------------- do_idx = do_index.idx bars.mb_do_mma_done[do_idx].wait(do_index.phase) - bars.mb_do_cg1_done[do_idx].wait(do_index.phase) do_index = advance(do_index, cfg.smem_do_stages) - if nvvm.elect_sync(): + if elect_one: bars.mb_do_ready[do_idx].arrive(n_bytes=cfg.tma_do_bytes) - do_slice = tma_slice_runtime_desc(desc_do_slot, cutlass.Int32(0), tok_coord) + 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) - # ---------------------------------------------------------- - # S = H[c-1] (forward state entering chunk c; S0 for c = 0 - # when an initial state is given, none otherwise) - # ---------------------------------------------------------- - if chunk_idx >= S_MIN: - s_idx = s_index.idx - bars.mb_s_done[s_idx].wait(s_index.phase) - s_index = advance(s_index, cfg.smem_s_stages) - if nvvm.elect_sync(): - bars.mb_s_ready[s_idx].arrive(n_bytes=cfg.tma_s_bytes) + # ---- entering state: checkpoint c - 1, or initial_state for chunk 0 when given ---------- + 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: - s0_slice = tma_slice_runtime_desc(desc_s0_slot, cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0)) - tma_load_tile(sS_tma[s_idx], s0_slice, bars.mb_s_ready[s_idx].smem_ptr, acquire=False) + 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: - s_slice = tma_slice_runtime_desc(desc_s_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx - 1) - tma_load_tile(sS_tma[s_idx], s_slice, bars.mb_s_ready[s_idx].smem_ptr, acquire=False) + 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: - s_slice = tma_slice_runtime_desc(desc_s_slot, cutlass.Int32(0), cutlass.Int32(0), chunk_idx - S_MIN) - tma_load_tile(sS_tma[s_idx], s_slice, bars.mb_s_ready[s_idx].smem_ptr, acquire=False) + 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) - tile_idx, sched_state = _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) for _ in range(cfg.smem_q_stages): bars.mb_q_mma_done[q_index.idx].wait(q_index.phase) @@ -2162,16 +1836,14 @@ def _tmaldg_warp( k_index = advance(k_index, cfg.smem_k_stages) for _ in range(cfg.smem_v_stages): bars.mb_v_mma_done[v_index.idx].wait(v_index.phase) - bars.mb_v_cg1_done[v_index.idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) for _ in range(cfg.smem_do_stages): bars.mb_do_mma_done[do_index.idx].wait(do_index.phase) - bars.mb_do_cg1_done[do_index.idx].wait(do_index.phase) do_index = advance(do_index, cfg.smem_do_stages) @cute.jit -def _compute0_warp( +def compute0_warp_group( cfg, total_tiles, bidx, @@ -2179,41 +1851,41 @@ def _compute0_warp( cu_seqlens, mWorkItems, tidx, - tmem_hold, + tmem_base_slot, scale, sCumsumlog, sCumprod, sBeta, - sAinv, + sTinv, sKK, - sQk, + sA, sDa, sDm, sK, sdQ, - sdH, - ss_flat, - sdh_flat, + sDstate, + sstate_flat, + sdstate_flat, sSched, bars, ): - """Compute warp-group 0 role (warps 0-3): persistent scheduler loop + - per-chunk T-pairwise, kk_epi, hierarchical blockwise inverse, qk_epi.""" + """Compute warp-group 0 role (warps 0-3): persistent scheduler loop + building each chunk's blockwise-inverse T and attention matrices.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) gate_index = PipelineState.start(phase=0) beta_index = PipelineState.start(phase=0) - qk_index = PipelineState.start(phase=1) - ainv_index = PipelineState.start(phase=1) + a_index = PipelineState.start(phase=1) + tinv_index = PipelineState.start(phase=1) da_acc_index = PipelineState.start(phase=0) - dm_rdy_index = PipelineState.start(phase=0) + dm_ready_index = PipelineState.start(phase=0) cg0_dbeta_index = PipelineState.start(phase=0) nvvm.barrier_cta_sync_aligned( cfg.tmem_alloc_barrier_id, thread_count=cfg.tmem_alloc_barrier_threads, ) - tmem_base = tmem_hold.load() + tmem_base = tmem_base_slot.load() num_threads_cg0 = cfg.threads_per_warp * len(cfg.compute_group_0_warp_ids) cg0_tidx = tidx % num_threads_cg0 @@ -2229,76 +1901,80 @@ def _compute0_warp( tmem_shared_acc_col = tmem_base + cfg.tmem_shared_acc_offset tmem_shared_inp_col = tmem_base + cfg.tmem_shared_inp_offset SHARED_INP_STAGE_COLS = cfg.b_t // 2 - tmem_dop_col = tmem_shared_inp_col + tmem_do_prime_col = tmem_shared_inp_col tmem_du_col = tmem_shared_inp_col + SHARED_INP_STAGE_COLS tmem_dyp_col = tmem_du_col ACC_STAGE_COLS = cfg.b_t tmem_acc_a = tmem_shared_acc_col tmem_acc_b = tmem_shared_acc_col + ACC_STAGE_COLS tmem_kk_col = tmem_acc_a - tmem_ks_col = tmem_acc_a + tmem_k_state_col = tmem_acc_a tmem_dy_col = tmem_acc_a tmem_dm_core_col = tmem_acc_a tmem_a_col = tmem_acc_b tmem_u_col = tmem_acc_b tmem_da_col = tmem_acc_b - tmem_dk_spath_col = tmem_shared_inp_col + tmem_dk_state_path_col = tmem_shared_inp_col acc_zero = cfg.acc_dtype(0.0) mask_zero = opaque_f32_zero() - ov_tok = cg0_tidx % 8 + (cg0_tidx // 16 % 2) * 8 - ov_col = (cg0_tidx // 8 % 2) * 8 + (cg0_tidx // 32 % 2) * 32 - ov_slab = (cg0_tidx // 64) * 4096 + frag_row = cg0_tidx % 8 + (cg0_tidx // 16 % 2) * 8 + frag_col = (cg0_tidx // 8 % 2) * 8 + (cg0_tidx // 32 % 2) * 32 + frag_slab_off = (cg0_tidx // 64) * 4096 sK_base_p = cute.make_ptr(cfg.io_dtype, sK[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) k_stage_elems_cg0 = cfg.k_cosize // cfg.smem_k_stages sdQ_base = cute.make_ptr(cfg.io_dtype, sdQ[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - sdH_parts_p = cute.make_ptr(cfg.io_dtype, sdH[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - # dG fold scratch in sKK (NOT sAinv: its upper triangle is kernel-start - # zeros, never rewritten); kk_epi fully rewrites sKK every chunk and its - # readers (inverse, M-terms dot) precede the fold in CG0 program order + sdH_parts_p = cute.make_ptr(cfg.io_dtype, sDstate[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + # dGate fold scratch in sKK skk_red = cute.make_ptr(cutlass.Float32, sKK[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) cg0_k_index = PipelineState.start(phase=0) cg0_dgate_index = PipelineState.start(phase=0) - tmem_dvdk_col = tmem_base + cfg.tmem_dvdk_offset - tmem_dh_inp_col = tmem_base + cfg.tmem_dh_inp_offset - cg0_kk_rdy = PipelineState.start(phase=0) - cg0_a_rdy = PipelineState.start(phase=0) - cg0_dk_scale_rdy = PipelineState.start(phase=0) - cg0_dq_scale_rdy = PipelineState.start(phase=0) - cg0_dhs_rdy = PipelineState.start(phase=0) - cg0_dk_attn_rdy = PipelineState.start(phase=0) - DHT0 = 1 if cfg.use_dht else 0 - - ainv_zero_ptr = cute.make_ptr(cutlass.Int32, sAinv[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - for z in cutlass.range_constexpr(cfg.ainv_cosize * bpe // 4 // num_threads_cg0): - (ainv_zero_ptr + cg0_tidx + z * num_threads_cg0).store(cutlass.Int32(0)) + tmem_dvdk_acc_col = tmem_base + cfg.tmem_dvdk_acc_offset + tmem_dstate_inp_col = tmem_base + cfg.tmem_dstate_inp_offset + cg0_kk_ready = PipelineState.start(phase=0) + cg0_a_ready = PipelineState.start(phase=0) + cg0_dk_scale_ready = PipelineState.start(phase=0) + cg0_dq_scale_ready = PipelineState.start(phase=0) + cg0_dstate_smem_ready = PipelineState.start(phase=0) + cg0_dk_attn_ready = PipelineState.start(phase=0) + DSTATE_IN0 = 1 if cfg.use_dstate_in else 0 + + tinv_zero_ptr = cute.make_ptr(cutlass.Int32, sTinv[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + for z in cutlass.range_constexpr(cfg.t_inv_cosize * bpe // 4 // num_threads_cg0): + (tinv_zero_ptr + cg0_tidx + z * num_threads_cg0).store(cutlass.Int32(0)) sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) - S_MIN = 0 if cfg.use_initial_state else 1 + 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, cu_seqlens, mWorkItems - ) - sk_nt = cend - wstart + 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 - for chunk_idx in cutlass.range(sk_nt): - # ---- Step 1: T-pairwise ------------------------------------ + for chunk_idx in cutlass.range(num_item_chunks): + # ---- T-pairwise ------------------------------------------------------ gate_idx = gate_index.idx bars.mb_gate_ready[gate_idx].wait(gate_index.phase) gate_index = advance(gate_index, cfg.smem_gate_stages) - gT = [] - gT_strict = [] + row_cs = [] + for r in cutlass.range_constexpr(2): + row_cs.append(sCumsumlog[warp_id * 16 + lane_id // 4 + r * 8, 0, gate_idx]) + col_cs = [] + for g in cutlass.range_constexpr(8): + for b in cutlass.range_constexpr(2): + col_cs.append(sCumsumlog[(lane_id % 4) * 2 + g * 8 + b, 0, gate_idx]) + decay_t = [] + decay_t_strict = [] for k in cutlass.range_constexpr(num_vals): crow = warp_id * 16 + lane_id // 4 + ((k // 2) % 2) * 8 ccol = (lane_id % 4) * 2 + ((k // 4) * 8 + k % 2) - gT.append(cute.math.exp2(sCumsumlog[crow, 0, gate_idx] - sCumsumlog[ccol, 0, gate_idx], fastmath=True) if crow >= ccol else mask_zero) - gT_strict.append(mask_zero if crow == ccol else gT[k]) - gDecayScale = [] + decay_t.append(cute.math.exp2(row_cs[(k // 2) % 2] - col_cs[(k // 4) * 2 + (k % 2)], fastmath=True) if crow >= ccol else mask_zero) + decay_t_strict.append(mask_zero if crow == ccol else decay_t[k]) last_cs = sCumsumlog[cfg.b_t - 1, 0, gate_idx] - for k in cutlass.range_constexpr(num_vals): - gDecayScale.append(cute.math.exp2(last_cs - sCumsumlog[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx], fastmath=True)) + decay_scale_fp32 = [] + for i in cutlass.range_constexpr(16): + decay_scale_fp32.append(cute.math.exp2(last_cs - col_cs[i], fastmath=True)) + decay_scale_vals = [decay_scale_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] cumprod_total = sCumprod[sCumprod.shape[0] - 1, 0, gate_idx] beta_idx = beta_index.idx @@ -2310,21 +1986,20 @@ def _compute0_warp( crow = warp_id * 16 + lane_id // 4 + ((k // 2) % 2) * 8 gBeta.append(sBeta[crow, 0, beta_idx]) - # ---- Step 2: kk_epi: M_kk[i,j] = W_kk[i,j] * T[i,j] * beta[i] ---- - ainv_idx = ainv_index.idx - ainv_phase = ainv_index.phase - ainv_index = advance(ainv_index, cfg.smem_ainv_stages) - bars.mb_kk_acc_ready[0].wait(cg0_kk_rdy.phase) - cg0_kk_rdy = advance(cg0_kk_rdy, 1) + # ---- KK epi: M_kk[i,j] = W_kk[i,j] * T[i,j] * Beta[i] --------------- + tinv_idx = tinv_index.idx + tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) + bars.mb_kk_acc_ready[0].wait(cg0_kk_ready.phase) + cg0_kk_ready = advance(cg0_kk_ready, 1) - ainv_base = sAinv[ainv_idx].base - kk_base = sKK[ainv_idx].base + tinv_base = sTinv[tinv_idx].base + kk_base = sKK[tinv_idx].base kk_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_kk_col, cutlass.Float32), num=8) - kk_f16 = [] + kk_pack = [] for k in cutlass.range_constexpr(num_vals // 2): - p0, p1 = fmul2(kk_vec[2 * k], kk_vec[2 * k + 1], gT[2 * k], gT[2 * k + 1]) + p0, p1 = fmul2(kk_vec[2 * k], kk_vec[2 * k + 1], decay_t[2 * k], decay_t[2 * k + 1]) v0, v1 = fmul2(p0, p1, gBeta[2 * k], gBeta[2 * k + 1]) - kk_f16.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) + kk_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( cutlass.inttoptr( @@ -2332,82 +2007,82 @@ def _compute0_warp( cutlass.AddressSpace.smem, cutlass.BFloat16, ), - [kk_f16[c * 4 + 0], kk_f16[c * 4 + 1], kk_f16[c * 4 + 2], kk_f16[c * 4 + 3]], + [kk_pack[c * 4 + 0], kk_pack[c * 4 + 1], kk_pack[c * 4 + 2], kk_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) - if chunk_idx < cend - S_MIN: + if chunk_idx < cend - FIRST_STATE_CHUNK: bars.mb_kk_acc_done[0].arrive() - # ---- Step 3: qk_epi: W_qkv[i,j] = W_qk[i,j] * T[i,j] * scale ---- - qk_idx = qk_index.idx - bars.mb_qk_done[qk_idx].wait(qk_index.phase) - qk_index = advance(qk_index, cfg.smem_qk_stages) - bars.mb_a_acc_ready[0].wait(cg0_a_rdy.phase) - cg0_a_rdy = advance(cg0_a_rdy, 1) + # ---- A epi: A[i,j] = W_qk[i,j] * T[i,j] * scale --------------------- + a_idx = a_index.idx + a_phase = a_index.phase + a_index = advance(a_index, cfg.smem_a_stages) + bars.mb_a_acc_ready[0].wait(cg0_a_ready.phase) + cg0_a_ready = advance(cg0_a_ready, 1) - qk_base = sQk[qk_idx].base - qk_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_a_col, cutlass.Float32), num=8) - qk_f16 = [] + a_base = sA[a_idx].base + a_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_a_col, cutlass.Float32), num=8) + a_pack = [] for k in cutlass.range_constexpr(num_vals // 2): - p0, p1 = fmul2(qk_vec[2 * k], qk_vec[2 * k + 1], gT[2 * k], gT[2 * k + 1]) + p0, p1 = fmul2(a_vec[2 * k], a_vec[2 * k + 1], decay_t[2 * k], decay_t[2 * k + 1]) v0, v1 = fmul2(p0, p1, scale, scale) - qk_f16.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) + a_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) + bars.mb_a_done[a_idx].wait(a_phase) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( cutlass.inttoptr( - qk_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, + a_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), - [qk_f16[c * 4 + 0], qk_f16[c * 4 + 1], qk_f16[c * 4 + 2], qk_f16[c * 4 + 3]], + [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_qk_ready[qk_idx].arrive() + bars.mb_a_ready[a_idx].arrive() - # ---- Step 4: blockwise inverse: A_inv = ------------------- - bars.mb_ainv_done[ainv_idx].wait(ainv_phase) + # ---- blockwise inverse: T_inv = ------------------------------------- nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) if warp_id < 2: - _invert_diagonal_NxN(cfg, kk_base, ainv_base, cg0_tidx // 8, cg0_tidx, 8) + invert_diagonal_NxN(cfg, kk_base, tinv_base, cg0_tidx // 8, cg0_tidx, 8) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - _blockwise_diagonal_8x8_to_16x16(cfg, ainv_base, kk_base, warp_id * 16, lane_id) + blockwise_diagonal_8x8_to_16x16(cfg, tinv_base, kk_base, warp_id * 16, lane_id) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) if warp_id < 2: - _blockwise_diagonal_16x16_to_32x32(cfg, ainv_base, kk_base, warp_id * 32, lane_id) + blockwise_diagonal_16x16_to_32x32(cfg, tinv_base, kk_base, warp_id * 32, lane_id) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) if warp_id < 2: - _blockwise_diagonal_32x32_to_64x64(cfg, ainv_base, kk_base, warp_id, lane_id) + blockwise_diagonal_32x32_to_64x64(cfg, tinv_base, kk_base, warp_id, lane_id) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - # ---- beta column-scaling in place: A_inv[i,j] *= beta[j] ---- + # ---- Beta column-scaling in place: T_inv[i,j] *= Beta[j] ------------ beta_col = [] for k in cutlass.range_constexpr(num_vals): beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta_idx]) - ainv_f16 = [] + tinv_frags = [] for c in cutlass.range_constexpr(ACC_N_FRAGS): - ainv_f16 += list( + tinv_frags += list( nvvm.ldmatrix( cutlass.inttoptr( - ainv_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, + tinv_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), @@ -2415,125 +2090,127 @@ def _compute0_warp( nvvm.MMALayout.ROW, ) ) - ainv_scaled = [] + tinv_pack = [] for j in cutlass.range_constexpr(num_vals // 2): - lo, hi = f16x2_to_f32(ainv_f16[j], dtype=cfg.io_dtype) + 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]) - ainv_scaled.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) + 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( - ainv_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, + tinv_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), - [ainv_scaled[c * 4 + 0], ainv_scaled[c * 4 + 1], ainv_scaled[c * 4 + 2], ainv_scaled[c * 4 + 3]], + [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_ainv_ready[ainv_idx].arrive() - - # ---- dQ inter rescale: gCumprod * scale in place (CG0-owned; - # the scale_ready wait also covers the hdh sS read via ks) ---- - if chunk_idx < cend - S_MIN: - gCumprod = [] - for k in cutlass.range_constexpr(num_vals): - gCumprod.append(sCumprod[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) - bars.mb_dq_acc_scale_ready[0].wait(cg0_dq_scale_rdy.phase) - cg0_dq_scale_rdy = advance(cg0_dq_scale_rdy, 1) - # all loads issue before the first store: a TMEM store between - # loads pins ptxas to one load latency per sub + bars.mb_t_inv_ready[tinv_idx].arrive() + + # ---- dQ inter rescale ------------------------------------------------ + if chunk_idx < cend - FIRST_STATE_CHUNK: + cumprod_fp32 = [] + for g in cutlass.range_constexpr(8): + for b in cutlass.range_constexpr(2): + cumprod_fp32.append(sCumprod[(lane_id % 4) * 2 + g * 8 + b, 0, gate_idx]) + cumprod_vals = [cumprod_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] + bars.mb_dq_acc_scale_ready[0].wait(cg0_dq_scale_ready.phase) + cg0_dq_scale_ready = advance(cg0_dq_scale_ready, 1) dqi_ptrs = [] dqi_vecs = [] for sub in cutlass.range_constexpr(2): - dqi_ptrs.append(nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dh_inp_col, cutlass.Float32)) + dqi_ptrs.append(nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dstate_inp_col, cutlass.Float32)) dqi_vecs.append(nvvm.tcgen05_ld("16x256b", dqi_ptrs[sub], num=8)) for sub in cutlass.range_constexpr(2): dqi_scaled = [] for j in cutlass.range_constexpr(16): - p0, p1 = fmul2(dqi_vecs[sub][2 * j], dqi_vecs[sub][2 * j + 1], gCumprod[2 * j], gCumprod[2 * j + 1]) + p0, p1 = fmul2(dqi_vecs[sub][2 * j], dqi_vecs[sub][2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) s0, s1 = fmul2(p0, p1, scale, scale) dqi_scaled += [s0, s1] nvvm.tcgen05_st("16x256b", dqi_ptrs[sub], cutlass.Vector.from_elements(tuple(dqi_scaled), cutlass.Float32)) nvvm.tcgen05_wait("store") bars.mb_dq_acc_scale_done[0].arrive() - # ---- dg_last H ⊙ dH term (octet-vectorized LDS.128); the sS - # read is released via the mb_s_done arrive below. The sdH - # staging is gated by a dhs_ready mirror-wait (non-consuming) ---- - if chunk_idx + DHT0 >= 1: - bars.mb_dhs_ready[0].wait(cg0_dhs_rdy.phase) - cg0_dhs_rdy = advance(cg0_dhs_rdy, 1) - sdh_base = sdh_flat.iterator.toint() - ss_base = ss_flat.iterator.toint() - hdh_lo = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] - hdh_hi = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] + # ---- dGate_last state dot dstate term ---------------------------------------- + if chunk_idx + DSTATE_IN0 >= 1: + bars.mb_dstate_smem_ready[0].wait(cg0_dstate_smem_ready.phase) + cg0_dstate_smem_ready = advance(cg0_dstate_smem_ready, 1) + sdstate_base = sdstate_flat.iterator.toint() + sstate_base = sstate_flat.iterator.toint() + state_dot_dstate_lo = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] + state_dot_dstate_hi = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] for oct_ in cutlass.range_constexpr(cfg.d_v // 8): - hs0 = (oct_ // 8) * (cfg.d_k * 64) + cg0_tidx * 64 + swizzle_xor_128b(cg0_tidx, (oct_ % 8) * 8) - dw = cute.make_tensor( - cute.make_ptr(cutlass.Int32, sdh_base + hs0 * bpe, mem_space=cute.AddressSpace.smem, assumed_align=16), + state_dot_off = (oct_ // 8) * (cfg.d_k * 64) + cg0_tidx * 64 + swizzle_xor_128b(cg0_tidx, (oct_ % 8) * 8) + dstate_frag = cute.make_tensor( + cute.make_ptr(cutlass.Int32, sdstate_base + state_dot_off * bpe, mem_space=cute.AddressSpace.smem, assumed_align=16), cute.make_layout(4), ).load() - sw = cute.make_tensor( - cute.make_ptr(cutlass.Int32, ss_base + hs0 * bpe, mem_space=cute.AddressSpace.smem, assumed_align=16), + state_frag = cute.make_tensor( + cute.make_ptr(cutlass.Int32, sstate_base + state_dot_off * bpe, mem_space=cute.AddressSpace.smem, assumed_align=16), cute.make_layout(4), ).load() for w in cutlass.range_constexpr(4): - d_lo, d_hi = f16x2_to_f32(dw[w], dtype=cfg.io_dtype) - s_lo, s_hi = f16x2_to_f32(sw[w], dtype=cfg.io_dtype) - hdh_lo[w], hdh_hi[w] = ffma2(d_lo, d_hi, s_lo, s_hi, hdh_lo[w], hdh_hi[w]) - hdh = ((hdh_lo[0] + hdh_lo[1]) + (hdh_lo[2] + hdh_lo[3])) + ((hdh_hi[0] + hdh_hi[1]) + (hdh_hi[2] + hdh_hi[3])) + dstate_lo, dstate_hi = f16x2_to_f32(dstate_frag[w], dtype=cfg.io_dtype) + state_lo, state_hi = f16x2_to_f32(state_frag[w], dtype=cfg.io_dtype) + state_dot_dstate_lo[w], state_dot_dstate_hi[w] = ffma2( + dstate_lo, dstate_hi, state_lo, state_hi, state_dot_dstate_lo[w], state_dot_dstate_hi[w] + ) + state_dot_dstate = ((state_dot_dstate_lo[0] + state_dot_dstate_lo[1]) + (state_dot_dstate_lo[2] + state_dot_dstate_lo[3])) + ( + (state_dot_dstate_hi[0] + state_dot_dstate_hi[1]) + (state_dot_dstate_hi[2] + state_dot_dstate_hi[3]) + ) for off in [1, 2, 4, 8, 16]: - hdh += nvvm.shfl_sync(0xFFFFFFFF, hdh, off, 31, kind=nvvm.Shfl.BFLY) - bars.mb_hdh_done[0].arrive() - if chunk_idx < cend - S_MIN: - nvvm.fence_proxy("async.shared", space="cta") - nvvm.mbarrier_arrive(bars.mb_s_done[0].smem_ptr) + state_dot_dstate += nvvm.shfl_sync(0xFFFFFFFF, state_dot_dstate, off, 31, kind=nvvm.Shfl.BFLY) + bars.mb_state_dot_dstate_done[0].arrive() - # ---- dK inter rescale: gDecayScale in place + the skd dot ---- + # ---- dK inter rescale ------------------------------------------------ cg0_k_idx = cg0_k_index.idx cg0_k_index = advance(cg0_k_index, cfg.smem_k_stages) - skd = cutlass.Float32(0.0) - if chunk_idx + DHT0 >= 1: - bars.mb_dk_scale_ready[0].wait(cg0_dk_scale_rdy.phase) - cg0_dk_scale_rdy = advance(cg0_dk_scale_rdy, 1) - skd_lo = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] - skd_hi = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] + k_dot_dk_inter = cutlass.Float32(0.0) + if chunk_idx + DSTATE_IN0 >= 1: + bars.mb_dk_scale_acc_ready[0].wait(cg0_dk_scale_ready.phase) + cg0_dk_scale_ready = advance(cg0_dk_scale_ready, 1) + k_dot_dk_inter_lo = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] + k_dot_dk_inter_hi = [opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero(), opaque_f32_zero()] dki_ptrs = [] dki_vecs = [] for sub in cutlass.range_constexpr(2): - dki_ptrs.append(nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_col, cutlass.Float32)) + dki_ptrs.append(nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32)) dki_vecs.append(nvvm.tcgen05_ld("16x256b", dki_ptrs[sub], num=8)) for sub in cutlass.range_constexpr(2): dki_scaled = [] for j in cutlass.range_constexpr(16): - s0, s1 = fmul2(dki_vecs[sub][2 * j], dki_vecs[sub][2 * j + 1], gDecayScale[2 * j], gDecayScale[2 * j + 1]) + s0, s1 = fmul2(dki_vecs[sub][2 * j], dki_vecs[sub][2 * j + 1], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1]) dki_scaled += [s0, s1] nvvm.tcgen05_st("16x256b", dki_ptrs[sub], cutlass.Vector.from_elements(tuple(dki_scaled), cutlass.Float32)) for m0 in cutlass.range_constexpr(4): - frag_addr = ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - k_f16 = nvvm.ldmatrix((sK_base_p + cg0_k_idx * k_stage_elems_cg0 + frag_addr).raw_ptr(), 4, nvvm.MMALayout.COL) + frag_addr = frag_slab_off + (frag_row + m0 * 16) * 64 + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) + k_frag = nvvm.ldmatrix((sK_base_p + cg0_k_idx * k_stage_elems_cg0 + frag_addr).raw_ptr(), 4, nvvm.MMALayout.COL) for i in cutlass.range_constexpr(4): - k_lo, k_hi = f16x2_to_f32(k_f16[i], dtype=cfg.io_dtype) - skd_lo[i], skd_hi[i] = ffma2(dki_scaled[8 * m0 + 2 * i], dki_scaled[8 * m0 + 2 * i + 1], k_lo, k_hi, skd_lo[i], skd_hi[i]) + k_lo, k_hi = f16x2_to_f32(k_frag[i], dtype=cfg.io_dtype) + k_dot_dk_inter_lo[i], k_dot_dk_inter_hi[i] = ffma2( + dki_scaled[8 * m0 + 2 * i], dki_scaled[8 * m0 + 2 * i + 1], k_lo, k_hi, k_dot_dk_inter_lo[i], k_dot_dk_inter_hi[i] + ) nvvm.tcgen05_wait("store") - bars.mb_dk_scale_done[0].arrive() - skd = ((skd_lo[0] + skd_lo[1]) + (skd_lo[2] + skd_lo[3])) + ((skd_hi[0] + skd_hi[1]) + (skd_hi[2] + skd_hi[3])) + bars.mb_dk_scale_acc_done[0].arrive() + k_dot_dk_inter = ((k_dot_dk_inter_lo[0] + k_dot_dk_inter_lo[1]) + (k_dot_dk_inter_lo[2] + k_dot_dk_inter_lo[3])) + ( + (k_dot_dk_inter_hi[0] + k_dot_dk_inter_hi[1]) + (k_dot_dk_inter_hi[2] + k_dot_dk_inter_hi[3]) + ) for off in [1, 2, 4, 8, 16]: - skd += nvvm.shfl_sync(0xFFFFFFFF, skd, off, 31, kind=nvvm.Shfl.BFLY) + k_dot_dk_inter += nvvm.shfl_sync(0xFFFFFFFF, k_dot_dk_inter, off, 31, kind=nvvm.Shfl.BFLY) - # ---- Step 6: dA epilogue ----------------------------------- + # ---- dA epilogue ----------------------------------------------------- bars.mb_da_acc_ready[0].wait(da_acc_index.phase) da_acc_index = advance(da_acc_index, 1) da_base = sDa[0].base da_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_da_col, cutlass.Float32), num=8) - da_f16 = [] + da_pack = [] for k in cutlass.range_constexpr(num_vals // 2): - p0, p1 = fmul2(da_vec[2 * k], da_vec[2 * k + 1], gT[2 * k], gT[2 * k + 1]) + p0, p1 = fmul2(da_vec[2 * k], da_vec[2 * k + 1], decay_t[2 * k], decay_t[2 * k + 1]) v0, v1 = fmul2(p0, p1, scale, scale) - da_f16.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) + da_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( cutlass.inttoptr( @@ -2541,24 +2218,24 @@ def _compute0_warp( cutlass.AddressSpace.smem, cutlass.BFloat16, ), - [da_f16[c * 4 + 0], da_f16[c * 4 + 1], da_f16[c * 4 + 2], da_f16[c * 4 + 3]], + [da_pack[c * 4 + 0], da_pack[c * 4 + 1], da_pack[c * 4 + 2], da_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_da_ready[0].arrive() - # ---- Step 8: dM epilogue ----------------------------------- - bars.mb_dm_ready[0].wait(dm_rdy_index.phase) - dm_rdy_index = advance(dm_rdy_index, 1) + # ---- dM epilogue ----------------------------------------------------- + bars.mb_dm_acc_ready[0].wait(dm_ready_index.phase) + dm_ready_index = advance(dm_ready_index, 1) dm_base = sDm[0].base dm_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dm_core_col, cutlass.Float32), num=8) - dm_f16 = [] + dm_pack = [] for k in cutlass.range_constexpr(num_vals // 2): - p0, p1 = fmul2(dm_vec[2 * k], dm_vec[2 * k + 1], gT_strict[2 * k], gT_strict[2 * k + 1]) + p0, p1 = fmul2(dm_vec[2 * k], dm_vec[2 * k + 1], decay_t_strict[2 * k], decay_t_strict[2 * k + 1]) v0 = cutlass.Float32(0.0) - p0 v1 = cutlass.Float32(0.0) - p1 - dm_f16.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) + dm_pack.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) for c in cutlass.range_constexpr(ACC_N_FRAGS): nvvm.stmatrix( cutlass.inttoptr( @@ -2566,27 +2243,27 @@ def _compute0_warp( cutlass.AddressSpace.smem, cutlass.BFloat16, ), - [dm_f16[c * 4 + 0], dm_f16[c * 4 + 1], dm_f16[c * 4 + 2], dm_f16[c * 4 + 3]], + [dm_pack[c * 4 + 0], dm_pack[c * 4 + 1], dm_pack[c * 4 + 2], dm_pack[c * 4 + 3]], nvvm.MMALayout.ROW, ) nvvm.fence_proxy("async.shared", space="cta") - bars.mb_dm_done[0].arrive() + bars.mb_dm_acc_done[0].arrive() - # ---- dK attn read: 16x256b fragment view for the part_k dot ---- - bars.mb_dk_attn_ready[0].wait(cg0_dk_attn_rdy.phase) - cg0_dk_attn_rdy = advance(cg0_dk_attn_rdy, 1) - dks_frag = [] + # ---- dK attn read ---------------------------------------------------- + bars.mb_dk_attn_acc_ready[0].wait(cg0_dk_attn_ready.phase) + cg0_dk_attn_ready = advance(cg0_dk_attn_ready, 1) + dks_regs = [] for sub in cutlass.range_constexpr(2): dks_vec = nvvm.tcgen05_ld( "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_col, cutlass.Float32), + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), num=8, ) - dks_frag.append([dks_vec[k] for k in range(32)]) + dks_regs.append([dks_vec[k] for k in range(32)]) nvvm.tcgen05_wait("load") - bars.mb_dk_attn_done[0].arrive() + bars.mb_dk_attn_acc_done[0].arrive() - # ---- dBeta/dGate M-terms: E = strict ⊙ dm_core ⊙ M_kk(sKK). ---- + # ---- dBeta/dGate M-terms: E = strict ⊙ dM_core ⊙ M_kk(sKK). ---------- kk_frag = [] for c in cutlass.range_constexpr(ACC_N_FRAGS): kk_frag += list( @@ -2625,40 +2302,42 @@ def _compute0_warp( for off in [1, 2]: for rp in cutlass.range_constexpr(2): row_part[rp] += nvvm.shfl_sync(0xFFFFFFFF, row_part[rp], off, 31, kind=nvvm.Shfl.BFLY) - # CG1's v-terms assign the dBeta base into sBeta; fold the k-term on top bars.mb_dbeta_cg1_ready[0].wait(cg0_dbeta_index.phase) cg0_dbeta_index = advance(cg0_dbeta_index, 1) if lane_id % 4 == 0: for rp in cutlass.range_constexpr(2): crow_r = warp_id * 16 + lane_id // 4 + rp * 8 sBeta[crow_r, 0, beta_idx] = sBeta[crow_r, 0, beta_idx] - row_part[rp] * binv_row[rp] - # ---- part reductions: fragment dot (K frags from sK, dks from - # the 16x256b r4 view), col_part folded in, reduce-scatter ---- - if chunk_idx + DHT0 >= S_MIN: + + # ---- part reductions ------------------------------------------------- + if chunk_idx + DSTATE_IN0 >= FIRST_STATE_CHUNK: part_k = [acc_zero] * 16 for sub in cutlass.range_constexpr(2): for m0 in cutlass.range_constexpr(4): - frag_addr = ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - k_f16 = nvvm.ldmatrix((sK_base_p + cg0_k_idx * k_stage_elems_cg0 + frag_addr).raw_ptr(), 4, nvvm.MMALayout.COL) + frag_addr = frag_slab_off + (frag_row + m0 * 16) * 64 + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) + k_frag = nvvm.ldmatrix((sK_base_p + cg0_k_idx * k_stage_elems_cg0 + frag_addr).raw_ptr(), 4, nvvm.MMALayout.COL) for i in cutlass.range_constexpr(4): - k_lo, k_hi = f16x2_to_f32(k_f16[i], dtype=cfg.io_dtype) - kk0 = cutlass.const_expr(8 * m0 + 2 * i) - j0 = cutlass.const_expr((kk0 // 4) * 2 + (kk0 % 2)) + k_lo, k_hi = f16x2_to_f32(k_frag[i], dtype=cfg.io_dtype) + frag_e0 = cutlass.const_expr(8 * m0 + 2 * i) + part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) if cutlass.const_expr(sub == 0 and i % 2 == 0): - part_k[j0], part_k[j0 + 1] = fmul2(dks_frag[sub][kk0], dks_frag[sub][kk0 + 1], k_lo, k_hi) + part_k[part_e0], part_k[part_e0 + 1] = fmul2(dks_regs[sub][frag_e0], dks_regs[sub][frag_e0 + 1], k_lo, k_hi) else: - part_k[j0], part_k[j0 + 1] = ffma2(dks_frag[sub][kk0], dks_frag[sub][kk0 + 1], k_lo, k_hi, part_k[j0], part_k[j0 + 1]) + part_k[part_e0], part_k[part_e0 + 1] = ffma2( + dks_regs[sub][frag_e0], dks_regs[sub][frag_e0 + 1], k_lo, k_hi, part_k[part_e0], part_k[part_e0 + 1] + ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_k_cg0_done[cg0_k_idx].arrive() - am = [col_part[j] - part_k[j] for j in range(16)] - am_lo, am_hi = _warp_reduce_scatter_frag16(am, lane_id) - dg_last_w = skd + ((cumprod_total * hdh if chunk_idx + DHT0 >= 1 else acc_zero) if chunk_idx < cend - S_MIN else acc_zero) + dgate_part = [col_part[j] - part_k[j] for j in range(16)] + dgate_part_lo, dgate_part_hi = warp_reduce_scatter_frag_16_elems(dgate_part, lane_id) + dgate_last_w = k_dot_dk_inter + ( + (cumprod_total * state_dot_dstate if chunk_idx + DSTATE_IN0 >= 1 else acc_zero) if chunk_idx < cend - FIRST_STATE_CHUNK else acc_zero + ) tok0 = (lane_id // 4) * 8 + (lane_id % 4) * 2 - (skk_red + warp_id * 64 + tok0).store(am_lo) - (skk_red + warp_id * 64 + tok0 + 1).store(am_hi + dg_last_w if lane_id == 31 else am_hi) + (skk_red + warp_id * 64 + tok0).store(dgate_part_lo) + (skk_red + warp_id * 64 + tok0 + 1).store(dgate_part_hi + dgate_last_w if lane_id == 31 else dgate_part_hi) - # CG1 adds part_q to dGate first; fold CG0's terms on top bars.mb_dgate_cg1_ready[0].wait(cg0_dgate_index.phase) cg0_dgate_index = advance(cg0_dgate_index, 1) if lane_id % 4 == 0: @@ -2667,35 +2346,36 @@ def _compute0_warp( sCumsumlog[crow_r, 0, gate_idx] = sCumsumlog[crow_r, 0, gate_idx] - row_part[rp] nvvm.barrier_cta_sync_aligned(cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads) if cg0_tidx < 64: - dg_sum = ( + dgate_sum = ( (skk_red + cg0_tidx).load() + (skk_red + 64 + cg0_tidx).load() + (skk_red + 128 + cg0_tidx).load() + (skk_red + 192 + cg0_tidx).load() ) - sCumsumlog[cg0_tidx, 0, gate_idx] = sCumsumlog[cg0_tidx, 0, gate_idx] + dg_sum + sCumsumlog[cg0_tidx, 0, gate_idx] = sCumsumlog[cg0_tidx, 0, gate_idx] + dgate_sum - if chunk_idx + DHT0 < S_MIN: + if chunk_idx + DSTATE_IN0 < FIRST_STATE_CHUNK: part_k = [acc_zero] * 16 for sub in cutlass.range_constexpr(2): for m0 in cutlass.range_constexpr(4): - frag_addr = ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - k_f16 = nvvm.ldmatrix((sK_base_p + cg0_k_idx * k_stage_elems_cg0 + frag_addr).raw_ptr(), 4, nvvm.MMALayout.COL) + frag_addr = frag_slab_off + (frag_row + m0 * 16) * 64 + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) + k_frag = nvvm.ldmatrix((sK_base_p + cg0_k_idx * k_stage_elems_cg0 + frag_addr).raw_ptr(), 4, nvvm.MMALayout.COL) for i in cutlass.range_constexpr(4): - k_lo, k_hi = f16x2_to_f32(k_f16[i], dtype=cfg.io_dtype) - kk0 = cutlass.const_expr(8 * m0 + 2 * i) - j0 = cutlass.const_expr((kk0 // 4) * 2 + (kk0 % 2)) + k_lo, k_hi = f16x2_to_f32(k_frag[i], dtype=cfg.io_dtype) + frag_e0 = cutlass.const_expr(8 * m0 + 2 * i) + part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) if cutlass.const_expr(sub == 0 and i % 2 == 0): - part_k[j0], part_k[j0 + 1] = fmul2(dks_frag[sub][kk0], dks_frag[sub][kk0 + 1], k_lo, k_hi) + part_k[part_e0], part_k[part_e0 + 1] = fmul2(dks_regs[sub][frag_e0], dks_regs[sub][frag_e0 + 1], k_lo, k_hi) else: - part_k[j0], part_k[j0 + 1] = ffma2(dks_frag[sub][kk0], dks_frag[sub][kk0 + 1], k_lo, k_hi, part_k[j0], part_k[j0 + 1]) + part_k[part_e0], part_k[part_e0 + 1] = ffma2( + dks_regs[sub][frag_e0], dks_regs[sub][frag_e0 + 1], k_lo, k_hi, part_k[part_e0], part_k[part_e0 + 1] + ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_k_cg0_done[cg0_k_idx].arrive() - am = [col_part[j] - part_k[j] for j in range(16)] - am_lo, am_hi = _warp_reduce_scatter_frag16(am, lane_id) + dgate_part = [col_part[j] - part_k[j] for j in range(16)] + dgate_part_lo, dgate_part_hi = warp_reduce_scatter_frag_16_elems(dgate_part, lane_id) tok0 = (lane_id // 4) * 8 + (lane_id % 4) * 2 - (skk_red + warp_id * 64 + tok0).store(am_lo) - (skk_red + warp_id * 64 + tok0 + 1).store(am_hi) + (skk_red + warp_id * 64 + tok0).store(dgate_part_lo) + (skk_red + warp_id * 64 + tok0 + 1).store(dgate_part_hi) - # CG1 adds part_q to dGate first; fold CG0's terms on top bars.mb_dgate_cg1_ready[0].wait(cg0_dgate_index.phase) cg0_dgate_index = advance(cg0_dgate_index, 1) if lane_id % 4 == 0: @@ -2704,38 +2384,34 @@ def _compute0_warp( sCumsumlog[crow_r, 0, gate_idx] = sCumsumlog[crow_r, 0, gate_idx] - row_part[rp] nvvm.barrier_cta_sync_aligned(cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads) if cg0_tidx < 64: - dg_sum = ( + dgate_sum = ( (skk_red + cg0_tidx).load() + (skk_red + 64 + cg0_tidx).load() + (skk_red + 128 + cg0_tidx).load() + (skk_red + 192 + cg0_tidx).load() ) - sCumsumlog[cg0_tidx, 0, gate_idx] = sCumsumlog[cg0_tidx, 0, gate_idx] + dg_sum + sCumsumlog[cg0_tidx, 0, gate_idx] = sCumsumlog[cg0_tidx, 0, gate_idx] + dgate_sum bars.mb_gate_done[gate_idx].arrive() bars.mb_beta_done[beta_idx].arrive() - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) - for _ in range(cfg.smem_qk_stages): - bars.mb_qk_done[qk_index.idx].wait(qk_index.phase) - qk_index = advance(qk_index, cfg.smem_qk_stages) - for _ in range(cfg.smem_ainv_stages): - bars.mb_ainv_done[ainv_index.idx].wait(ainv_index.phase) - ainv_index = advance(ainv_index, cfg.smem_ainv_stages) - - # CG0 done with TMEM: release the MMA warp's dealloc + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + for _ in range(cfg.smem_a_stages): + bars.mb_a_done[a_index.idx].wait(a_index.phase) + a_index = advance(a_index, cfg.smem_a_stages) + bars.mb_tmem_done[0].arrive() @cute.jit -def _compute1_warp( +def compute1_warp_group( cfg, total_tiles, bidx, num_ctas, cu_seqlens, mWorkItems, - mDs0_out, - mDht, + mDstate0_out, + mDstate_in, tidx, warp_idx, - tmem_hold, + tmem_base_slot, scale, sQ, sK, @@ -2747,38 +2423,32 @@ def _compute1_warp( sdQ, sdK, sdV, - sdH, + sDstate, sDa, sDm, sSched, bars, ): - """Compute warp-group 1 role (warps 4-7): persistent scheduler loop and - per-chunk (BACKWARD order): dV-acc / dQ-acc - in-place decay scales, dO'/dU/dY' TMEM restages, Y and gks staging, - u readout over sV, the dV-pass dG/dBeta v-terms, dV/dK staging with the - dK s-path fold, and the next-chunk dH prep at the chunk tail.""" + """Compute warp-group 1 role (warps 4-7): persistent scheduler loop + running each chunk's gradient epilogues and stagings.""" v_index = PipelineState.start(phase=0) do_index = PipelineState.start(phase=0) gate_index = PipelineState.start(phase=0) - cg1_ks_rdy = PipelineState.start(phase=0) - cg1_u_rdy = PipelineState.start(phase=0) - cg1_dy_rdy = PipelineState.start(phase=0) - cg1_du_scale_rdy = PipelineState.start(phase=0) - cg1_du_total_rdy = PipelineState.start(phase=0) - cg1_dk_total_rdy = PipelineState.start(phase=0) - dh_acc_index = PipelineState.start(phase=0) - dop_inp_free = PipelineState.start(phase=1) - dyp_inp_free = PipelineState.start(phase=1) - dhs_index = PipelineState.start(phase=1) - cg1_hdh_index = PipelineState.start(phase=0) + cg1_k_state_ready = PipelineState.start(phase=0) + cg1_u_ready = PipelineState.start(phase=0) + cg1_dy_ready = PipelineState.start(phase=0) + cg1_du_scale_ready = PipelineState.start(phase=0) + cg1_du_total_ready = PipelineState.start(phase=0) + cg1_dk_total_ready = PipelineState.start(phase=0) + dstate_acc_index = PipelineState.start(phase=0) + cg1_state_dot_dstate_index = PipelineState.start(phase=0) dq_index = PipelineState.start(phase=1) cg1_beta_index = PipelineState.start(phase=0) sdv_done_index = PipelineState.start(phase=1) - cg1_dk_spath_rdy = PipelineState.start(phase=0) - dq_total_rdy_index = PipelineState.start(phase=0) - dh_inp_index = PipelineState.start(phase=1) + cg1_dk_state_path_ready = PipelineState.start(phase=0) + dq_total_ready_index = PipelineState.start(phase=0) + dstate_inp_index = PipelineState.start(phase=1) dk_index = PipelineState.start(phase=1) dv_index = PipelineState.start(phase=1) @@ -2787,7 +2457,7 @@ def _compute1_warp( cfg.tmem_alloc_barrier_id, thread_count=cfg.tmem_alloc_barrier_threads, ) - tmem_base = tmem_hold.load() + tmem_base = tmem_base_slot.load() num_threads_cg1 = cfg.threads_per_warp * len(cfg.compute_group_1_warp_ids) cg1_tidx = tidx % num_threads_cg1 @@ -2796,31 +2466,31 @@ def _compute1_warp( ldtm_width = 32 sttm_width = ldtm_width // 2 num_state_subs = cutlass.const_expr(cfg.d_v // ldtm_width) - tmem_dh_col = tmem_base + cfg.tmem_dh_offset - tmem_dh_inp_col = tmem_base + cfg.tmem_dh_inp_offset - tmem_dvdk_col = tmem_base + cfg.tmem_dvdk_offset + tmem_dstate_acc_col = tmem_base + cfg.tmem_dstate_acc_offset + tmem_dstate_inp_col = tmem_base + cfg.tmem_dstate_inp_offset + tmem_dvdk_acc_col = tmem_base + cfg.tmem_dvdk_acc_offset tmem_shared_acc_col = tmem_base + cfg.tmem_shared_acc_offset tmem_shared_inp_col = tmem_base + cfg.tmem_shared_inp_offset SHARED_INP_STAGE_COLS = cfg.b_t // 2 - tmem_dop_col = tmem_shared_inp_col + tmem_do_prime_col = tmem_shared_inp_col tmem_du_col = tmem_shared_inp_col + SHARED_INP_STAGE_COLS tmem_dyp_col = tmem_du_col tmem_y_col = tmem_base + cfg.tmem_y_offset - tmem_gks_col = tmem_y_col + SHARED_INP_STAGE_COLS + tmem_g_k_state_col = tmem_y_col + SHARED_INP_STAGE_COLS ACC_STAGE_COLS = cfg.b_t tmem_acc_a = tmem_shared_acc_col tmem_acc_b = tmem_shared_acc_col + ACC_STAGE_COLS tmem_kk_col = tmem_acc_a - tmem_ks_col = tmem_acc_a + tmem_k_state_col = tmem_acc_a tmem_dy_col = tmem_acc_a tmem_dm_core_col = tmem_acc_a tmem_a_col = tmem_acc_b tmem_u_col = tmem_acc_b tmem_da_col = tmem_acc_b - tmem_dk_spath_col = tmem_shared_inp_col - ov_tok = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 - ov_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 - ov_slab = (cg1_tidx // 64) * 4096 + tmem_dk_state_path_col = tmem_shared_inp_col + frag_row = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 + frag_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 + frag_slab_off = (cg1_tidx // 64) * 4096 dv_stage_elems = cfg.dv_cosize // cfg.smem_dv_stages sdV_base = cute.make_ptr(cfg.io_dtype, sdV[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) v_stage_elems = cfg.v_cosize // cfg.smem_v_stages @@ -2832,601 +2502,110 @@ def _compute1_warp( dq_stage_elems = cfg.dq_cosize // cfg.smem_dq_stages sdQ_base = cute.make_ptr(cfg.io_dtype, sdQ[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) sdK_base = cute.make_ptr(cfg.io_dtype, sdK[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - sdH_base_int = sdH[0].base - # dBeta/dGate reduction scratch: f32 view of the owned dQ stage (rebased per chunk) + sDstate_base_int = sDstate[0].base sred_base = cute.make_ptr(cutlass.Float32, sdQ[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - # part_q scratch in sdH: consumed by dK-inter (covered by total_ready = - # the dQ-attn commit) and re-staged by CG1 itself later in the iteration - sdh_red = cute.make_ptr(cutlass.Float32, sdH[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) - dh_done_idx = cutlass.Int32(0) - cg1w = cg1_tidx // cfg.threads_per_warp + sdstate_red = cute.make_ptr(cutlass.Float32, sDstate[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + dstate_done_idx = cutlass.Int32(0) + cg1_warp_id = cg1_tidx // cfg.threads_per_warp sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) - S_MIN = 0 if cfg.use_initial_state else 1 + 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, cu_seqlens, mWorkItems - ) - sk_nt = cend - wstart - # ---- d_final_state prologue: seed the dH acc (f32) + the f16 - # entry staging (dh_inp + sdH) ---- - if cutlass.const_expr(cfg.use_dht): - if sk_nt > 0: - # split-K: only the item owning the sequence tail receives - # the true d_final_state; warmup items rebuild dH from zero - gDht = mDht[None, None, head_idx, batch_idx] - seed_from_dht = cend == num_chunks_b - dh_inp_idx = dh_inp_index.idx - bars.mb_dh_inp_done[dh_inp_idx].wait(dh_inp_index.phase) - dh_inp_index = advance(dh_inp_index, cfg.tmem_dh_inp_stages) + 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 + + # ---- d_final_state prologue ---------------------------------------------- + if cutlass.const_expr(cfg.use_dstate_in): + if num_item_chunks > 0: + gDstate_in = mDstate_in[None, None, head_idx, batch_idx] + seed_from_dstate_in = cend == num_chunks_b + dstate_inp_idx = dstate_inp_index.idx + bars.mb_dstate_inp_done[dstate_inp_idx].wait(dstate_inp_index.phase) + dstate_inp_index = advance(dstate_inp_index, cfg.tmem_dstate_inp_stages) for sub in cutlass.range_constexpr(num_state_subs): - dht_vals = [] + dstate_in_vals = [] for kk in cutlass.range_constexpr(ldtm_width): - v = gDht[sub * ldtm_width + kk, cg1_tidx] - if cutlass.const_expr(cfg.split_k): - v = v if seed_from_dht else cutlass.Float32(0.0) - dht_vals.append(v) + v = gDstate_in[sub * ldtm_width + kk, cg1_tidx] + v = v if seed_from_dstate_in else cutlass.Float32(0.0) + dstate_in_vals.append(v) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dh_col + sub * ldtm_width, cutlass.Float32), - cutlass.Vector.from_elements(tuple(dht_vals), cutlass.Float32), + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), + cutlass.Vector.from_elements(tuple(dstate_in_vals), cutlass.Float32), ) - dht_f16 = [fp32_to_fp16(dht_vals[2 * j], dht_vals[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] + dstate_in_pack = [fp32_to_fp16(dstate_in_vals[2 * j], dstate_in_vals[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dh_inp_col + sub * sttm_width, cutlass.Int32), - cutlass.Vector.from_elements(tuple(dht_f16), cutlass.Int32), + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_inp_col + sub * sttm_width, cutlass.Int32), + cutlass.Vector.from_elements(tuple(dstate_in_pack), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_dh_inp_ready[dh_inp_idx].arrive() + bars.mb_dstate_inp_ready[dstate_inp_idx].arrive() - # sdH staging read back from the just-seeded acc - bars.mb_dhs_done[0].wait(dhs_index.phase) - dhs_index = advance(dhs_index, 1) - dhs_vecs = [] for b in cutlass.range_constexpr(2): - for hh in cutlass.range_constexpr(2): - dhs_vecs.append( - nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + b * 16) << 16) + tmem_dh_col + hh * 64, cutlass.Float32), - num=8, - ) + 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 b in cutlass.range_constexpr(2): - for hh in cutlass.range_constexpr(2): - dhs_vec = dhs_vecs[b * 2 + hh] - dhs_f16 = [fp32_to_fp16(dhs_vec[2 * j], dhs_vec[2 * j + 1], dtype=cfg.io_dtype) for j in range(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): - dhs_row = hh * 64 + ov_tok + c * 16 + dstate_smem_row = col_half * 64 + frag_row + c * 16 nvvm.stmatrix( cutlass.inttoptr( - sdH_base_int + ((cg1_tidx // 64) * cfg.d_k * 64 + dhs_row * 64 + swizzle_xor_128b(dhs_row, ov_col + b * 16)) * 2, + 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, ), - [dhs_f16[c * 4 + 0], dhs_f16[c * 4 + 1], dhs_f16[c * 4 + 2], dhs_f16[c * 4 + 3]], + [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_dhs_ready[0].arrive() + bars.mb_dstate_smem_ready[0].arrive() - if cutlass.const_expr(not cfg.use_dht): - if sk_nt > 0: - # ---- first backward chunk (c = NT-1): no dH yet ------------------- - gate_idx = gate_index.idx - bars.mb_gate_ready[gate_idx].wait(gate_index.phase) - gate_index = advance(gate_index, cfg.smem_gate_stages) - beta_idx = cg1_beta_index.idx - bars.mb_beta_ready[beta_idx].wait(cg1_beta_index.phase) - cg1_beta_index = advance(cg1_beta_index, cfg.smem_beta_stages) - num_vals = 32 - gCumprod = [] - for k in cutlass.range_constexpr(num_vals): - gCumprod.append(sCumprod[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) - - # ---- dO' restage: dO * gCumprod * scale -> shared_inp TMEM ---- - bars.mb_dop_inp_done[0].wait(dop_inp_free.phase) - dop_inp_free = advance(dop_inp_free, 1) - do_idx = do_index.idx - bars.mb_do_ready[do_idx].wait(do_index.phase) - do_index = advance(do_index, cfg.smem_do_stages) - do_regs = [[cutlass.Float32(0.0), cutlass.Float32(0.0)] for _ in range(32)] - for c in cutlass.range_constexpr(8): - m0 = cutlass.const_expr(c % 4) - sub = cutlass.const_expr(c // 4) - do_f16 = nvvm.ldmatrix( - ( - sdO_base + do_idx * do_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - ).raw_ptr(), - 4, - nvvm.MMALayout.COL, + # ---- chunks NT-1 .. 0 (backward) ------------------------------------------ + for rev_idx in cutlass.range(num_item_chunks): + chunk_idx = cend - 1 - rev_idx + have_dstate = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_dstate_in) else rev_idx > 0 + gate_idx = gate_index.idx + bars.mb_gate_ready[gate_idx].wait(gate_index.phase) + gate_index = advance(gate_index, cfg.smem_gate_stages) + + beta_idx = cg1_beta_index.idx + bars.mb_beta_ready[beta_idx].wait(cg1_beta_index.phase) + cg1_beta_index = advance(cg1_beta_index, cfg.smem_beta_stages) + + # ---- dstate rescale: dstate *= this chunk's cumprod -------------------------- + if have_dstate: + cumprod_top = sCumprod[sCumprod.shape[0] - 1, 0, gate_idx] + for sub in cutlass.range_constexpr(num_state_subs): + dstate_rescale_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 i in cutlass.range_constexpr(4): - lo, hi = f16x2_to_f32(do_f16[i], dtype=cfg.io_dtype) - p0, p1 = fmul2(lo, hi, gCumprod[8 * m0 + 2 * i], gCumprod[8 * m0 + 2 * i + 1]) - do_regs[8 * m0 + 2 * i][sub], do_regs[8 * m0 + 2 * i + 1][sub] = fmul2(p0, p1, scale, scale) - for sub in cutlass.range_constexpr(2): - do_pack = [fp32_to_fp16(do_regs[2 * j][sub], do_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + dstate_rescaled = [] + for j in cutlass.range_constexpr(16): + h0, h1 = fmul2(dstate_rescale_vec[2 * j], dstate_rescale_vec[2 * j + 1], cumprod_top, cumprod_top) + dstate_rescaled += [h0, h1] nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dop_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(do_pack), cutlass.Int32), + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_acc_col + sub * ldtm_width, cutlass.Float32), + cutlass.Vector.from_elements(tuple(dstate_rescaled), cutlass.Float32), ) nvvm.tcgen05_wait("store") - bars.mb_dop_inp_ready[0].arrive() - - # ---- v - k*state: delta = V - cumprod*(K @ S), in place ---- - v_idx = v_index.idx - bars.mb_v_ready[v_idx].wait(v_index.phase) - v_index = advance(v_index, cfg.smem_v_stages) - if cend >= SFIRST_MIN: - # V stays PACKED in the io dtype; the scaled KS is packed - # (needed for its own TMEM store anyway) and subtracted - # with packed 16-bit ops. - v_words = [[cutlass.Int32(0), cutlass.Int32(0)] for _ in range(16)] - for c in cutlass.range_constexpr(8): - m0 = cutlass.const_expr(c % 4) - sub = cutlass.const_expr(c // 4) - v_f16 = nvvm.ldmatrix( - ( - sV_base + v_idx * v_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - ).raw_ptr(), - 4, - nvvm.MMALayout.COL, - ) - for i in cutlass.range_constexpr(4): - v_words[4 * m0 + i][sub] = v_f16[i] - bars.mb_ks_acc_ready[0].wait(cg1_ks_rdy.phase) - cg1_ks_rdy = advance(cg1_ks_rdy, 1) - ks_vecs = [] - for sub in cutlass.range_constexpr(2): - ks_vecs.append( - nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_ks_col, cutlass.Float32), - num=8, - ) - ) - for sub in cutlass.range_constexpr(2): - gks_pack = [] - y_pack = [] - for j in cutlass.range_constexpr(16): - g0, g1 = fmul2(ks_vecs[sub][2 * j], ks_vecs[sub][2 * j + 1], gCumprod[2 * j], gCumprod[2 * j + 1]) - gks_word = fp32_to_fp16(g0, g1, dtype=cfg.io_dtype) - gks_pack.append(gks_word) - y_pack.append(sub_f16x2(v_words[j][sub], gks_word, cfg.io_dtype)) - nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(y_pack), cutlass.Int32), - ) - nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_gks_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(gks_pack), cutlass.Int32), - ) - if cend < SFIRST_MIN: - for sub in cutlass.range_constexpr(2): - v_pack = [] - for m0 in cutlass.range_constexpr(4): - v_f16 = nvvm.ldmatrix( - ( - sV_base + v_idx * v_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - ).raw_ptr(), - 4, - nvvm.MMALayout.COL, - ) - for i in cutlass.range_constexpr(4): - v_pack.append(v_f16[i]) - nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(v_pack), cutlass.Int32), - ) - nvvm.tcgen05_wait("store") - bars.mb_y_ready[0].arrive() - - # ---- dU restage: dv acc ------------------------------------ - bars.mb_dyp_inp_done[0].wait(dyp_inp_free.phase) - dyp_inp_free = advance(dyp_inp_free, 1) - bars.mb_du_total_ready[0].wait(cg1_du_total_rdy.phase) - cg1_du_total_rdy = advance(cg1_du_total_rdy, 1) - du_vecs = [] - for sub in cutlass.range_constexpr(2): - du_vecs.append(nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_col, cutlass.Float32), num=8)) - for sub in cutlass.range_constexpr(2): - du_pack = [fp32_to_fp16(du_vecs[sub][2 * j], du_vecs[sub][2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] - nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_du_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(du_pack), cutlass.Int32), - ) - nvvm.tcgen05_wait("store") - bars.mb_du_inp_ready[0].arrive() - bars.mb_do_cg1_done[do_idx].arrive() - - # ---- U readout --------------------------------------------- - bars.mb_u_acc_ready[0].wait(cg1_u_rdy.phase) - cg1_u_rdy = advance(cg1_u_rdy, 1) - u_regs = [] - for sub in cutlass.range_constexpr(2): - u_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_u_col, cutlass.Float32), num=8) - u_regs.append([u_vec[k] for k in range(32)]) - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - u_f16 = [fp32_to_fp16(u_regs[sub][8 * m0 + 2 * j], u_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - nvvm.stmatrix( - ( - sV_base + v_idx * v_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - ).raw_ptr(), - u_f16, - nvvm.MMALayout.COL, - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_u_ready[0].arrive() - bars.mb_v_cg1_done[v_idx].arrive() - - # ---- dY ---------------------------------------------------- - bars.mb_dy_acc_ready[0].wait(cg1_dy_rdy.phase) - cg1_dy_rdy = advance(cg1_dy_rdy, 1) - dy_regs = [] - for sub in cutlass.range_constexpr(2): - dy_vec = nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dy_col, cutlass.Float32), - num=8, - ) - dy_regs.append([dy_vec[k] for k in range(32)]) - - # ---- dY' = -gCumprod * dY -> f16 shared_inp ---------------- - neg_one = cutlass.Float32(-1.0) - gCumprodNeg = [gCumprod[k] * neg_one for k in range(32)] - for sub in cutlass.range_constexpr(2): - dyp = [] - for j in cutlass.range_constexpr(16): - n0, n1 = fmul2(dy_regs[sub][2 * j], dy_regs[sub][2 * j + 1], gCumprodNeg[2 * j], gCumprodNeg[2 * j + 1]) - dyp += [n0, n1] - dyp_pack = [fp32_to_fp16(dyp[2 * j], dyp[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] - nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dyp_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(dyp_pack), cutlass.Int32), - ) - nvvm.tcgen05_wait("store") - bars.mb_dyp_inp_ready[0].arrive() - - # ---- dV staging: dV = dY -> the sdV slot ------------------- - dv_stg_idx = dv_index.idx - bars.mb_dv_tmastg_done[dv_stg_idx].wait(dv_index.phase) - dv_index = advance(dv_index, cfg.smem_dv_stages) - bars.mb_sdv_done[0].wait(sdv_done_index.phase) - sdv_done_index = advance(sdv_done_index, 1) - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - dv_f16 = [fp32_to_fp16(dy_regs[sub][8 * m0 + 2 * j], dy_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - nvvm.stmatrix( - ( - sdV_base - + dv_stg_idx * dv_stage_elems - + ov_slab - + (ov_tok + m0 * 16) * 64 - + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - ).raw_ptr(), - dv_f16, - nvvm.MMALayout.COL, - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_dv_tmastg_ready[dv_stg_idx].arrive() - dk_stg_idx = dk_index.idx - bars.mb_dk_tmastg_done[dk_stg_idx].wait(dk_index.phase) - dk_index = advance(dk_index, cfg.smem_dk_stages) - dq_stg_idx = dq_index.idx - bars.mb_dq_tmastg_done[dq_stg_idx].wait(dq_index.phase) - dq_index = advance(dq_index, cfg.smem_dq_stages) - sred = sred_base + dq_stg_idx * (dq_stage_elems // 2) - - # ---- dBeta/dGate v-terms: dbeta_t += rowsum(dV ⊙ Y)_t / beta_t ---- - part_y = [cutlass.Float32(0.0)] * 16 - for sub in cutlass.range_constexpr(2): - y_pk = nvvm.tcgen05_ld("16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), num=8) - for j in cutlass.range_constexpr(16): - lo, hi = f16x2_to_f32(y_pk[j], dtype=cfg.io_dtype) - kk0 = cutlass.const_expr(2 * j) - j0 = cutlass.const_expr((kk0 // 4) * 2 + (kk0 % 2)) - if cutlass.const_expr(sub == 0 and j % 2 == 0): - part_y[j0], part_y[j0 + 1] = fmul2(dy_regs[sub][kk0], dy_regs[sub][kk0 + 1], lo, hi) - else: - part_y[j0], part_y[j0 + 1] = ffma2(dy_regs[sub][kk0], dy_regs[sub][kk0 + 1], lo, hi, part_y[j0], part_y[j0 + 1]) - py_lo, py_hi = _warp_reduce_scatter_frag16(part_y, lane_id) - vt_tok0 = (lane_id // 4) * 8 + (lane_id % 4) * 2 - if cend >= SFIRST_MIN: - part_g = [cutlass.Float32(0.0)] * 16 - for sub in cutlass.range_constexpr(2): - g_pk = nvvm.tcgen05_ld("16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_gks_col, cutlass.Int32), num=8) - for j in cutlass.range_constexpr(16): - lo, hi = f16x2_to_f32(g_pk[j], dtype=cfg.io_dtype) - kk0 = cutlass.const_expr(2 * j) - j0 = cutlass.const_expr((kk0 // 4) * 2 + (kk0 % 2)) - if cutlass.const_expr(sub == 0 and j % 2 == 0): - part_g[j0], part_g[j0 + 1] = fmul2(dy_regs[sub][kk0], dy_regs[sub][kk0 + 1], lo, hi) - else: - part_g[j0], part_g[j0 + 1] = ffma2(dy_regs[sub][kk0], dy_regs[sub][kk0 + 1], lo, hi, part_g[j0], part_g[j0 + 1]) - pg_lo, pg_hi = _warp_reduce_scatter_frag16(part_g, lane_id) - (sred + cg1w * 64 + vt_tok0).store(py_lo) - (sred + cg1w * 64 + vt_tok0 + 1).store(py_hi) - (sred + 256 + cg1w * 64 + vt_tok0).store(pg_lo) - (sred + 256 + cg1w * 64 + vt_tok0 + 1).store(pg_hi) - nvvm.barrier_cta_sync_aligned(cfg.cg1_barrier_id, thread_count=cfg.cg1_barrier_threads) - if cg1_tidx < 64: - binv_t = cute.math.rcp(sBeta[cg1_tidx, 0, beta_idx] + cutlass.Float32(1e-10), approx=True, ftz=True) - ysum = (sred + cg1_tidx).load() + (sred + 64 + cg1_tidx).load() + (sred + 128 + cg1_tidx).load() + (sred + 192 + cg1_tidx).load() - gsum = (sred + 256 + cg1_tidx).load() + (sred + 320 + cg1_tidx).load() + (sred + 384 + cg1_tidx).load() + (sred + 448 + cg1_tidx).load() - sBeta[cg1_tidx, 0, beta_idx] = ysum * binv_t - sCumsumlog[cg1_tidx, 0, gate_idx] = cutlass.Float32(0.0) - gsum - if cend < SFIRST_MIN: - (sred + cg1w * 64 + vt_tok0).store(py_lo) - (sred + cg1w * 64 + vt_tok0 + 1).store(py_hi) - nvvm.barrier_cta_sync_aligned(cfg.cg1_barrier_id, thread_count=cfg.cg1_barrier_threads) - if cg1_tidx < 64: - binv_t = cute.math.rcp(sBeta[cg1_tidx, 0, beta_idx] + cutlass.Float32(1e-10), approx=True, ftz=True) - ysum = (sred + cg1_tidx).load() + (sred + 64 + cg1_tidx).load() + (sred + 128 + cg1_tidx).load() + (sred + 192 + cg1_tidx).load() - sBeta[cg1_tidx, 0, beta_idx] = ysum * binv_t - sCumsumlog[cg1_tidx, 0, gate_idx] = cutlass.Float32(0.0) - bars.mb_beta_done[beta_idx].arrive() - bars.mb_dbeta_cg1_ready[0].arrive() - # sred reads must land before the dq stmatrix below reuses the stage - nvvm.barrier_cta_sync_aligned(cfg.cg1_barrier_id, thread_count=cfg.cg1_barrier_threads) - # ---- Q fragments held in registers over the dq stage (sQ - # free once read; the dQ dot consumes them -- no TMEM trip) ---- - q_frag = [] - for sub in cutlass.range_constexpr(2): - q_words = [] - for m0 in cutlass.range_constexpr(4): - q_f16 = nvvm.ldmatrix( - (sQ_base + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16)).raw_ptr(), - 4, - nvvm.MMALayout.COL, - ) - for i in cutlass.range_constexpr(4): - q_words.append(q_f16[i]) - q_frag.append(q_words) - # the store+wait order the sQ LDSM reads before the release arrive - # (a bare register consume can be hoisted past by ptxas); the dot - # below reads the registers, so the TMEM read-back is still gone - nvvm.tcgen05_st( - "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(q_words), cutlass.Int32), - ) - nvvm.tcgen05_wait("store") - bars.mb_q_cg1_done[0].arrive() - - # ---- dq final read -> sdQ (output staging; the fragments are - # held for the dQ dot below) --------- - bars.mb_dq_acc_total_ready[0].wait(dq_total_rdy_index.phase) - dq_total_rdy_index = advance(dq_total_rdy_index, 1) - dq_frag = [] - for sub in cutlass.range_constexpr(2): - dqv = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dh_inp_col, cutlass.Float32), num=8) - dq_frag.append([dqv[k] for k in range(32)]) - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - frag_addr = ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - dq_f16 = [fp32_to_fp16(dq_frag[sub][8 * m0 + 2 * j], dq_frag[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - nvvm.stmatrix((sdQ_base + dq_stg_idx * dq_stage_elems + frag_addr).raw_ptr(), dq_f16, nvvm.MMALayout.COL) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_dq_acc_total_done[0].arrive() - bars.mb_dq_tmastg_ready[dq_stg_idx].arrive() - - # ---- dQ dot (part_q): fragment dot of the held dq acc with the - # staged Q^T, added to dGate FIRST; CG0 adds after parts_ready ---- - part_q = [cutlass.Float32(0.0)] * 16 - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - for i in cutlass.range_constexpr(4): - q_lo, q_hi = f16x2_to_f32(q_frag[sub][4 * m0 + i], dtype=cfg.io_dtype) - kk0 = cutlass.const_expr(8 * m0 + 2 * i) - j0 = cutlass.const_expr((kk0 // 4) * 2 + (kk0 % 2)) - if cutlass.const_expr(sub == 0 and i % 2 == 0): - part_q[j0], part_q[j0 + 1] = fmul2(dq_frag[sub][kk0], dq_frag[sub][kk0 + 1], q_lo, q_hi) - else: - part_q[j0], part_q[j0 + 1] = ffma2(dq_frag[sub][kk0], dq_frag[sub][kk0 + 1], q_lo, q_hi, part_q[j0], part_q[j0 + 1]) - amq_lo, amq_hi = _warp_reduce_scatter_frag16(part_q, lane_id) - tok0 = (lane_id // 4) * 8 + (lane_id % 4) * 2 - (sdh_red + (cg1_tidx // 32) * 64 + tok0).store(amq_lo) - (sdh_red + (cg1_tidx // 32) * 64 + tok0 + 1).store(amq_hi) - nvvm.barrier_cta_sync_aligned(cfg.cg1_barrier_id, thread_count=cfg.cg1_barrier_threads) - if cg1_tidx < 64: - pq_sum = ( - (sdh_red + cg1_tidx).load() + (sdh_red + 64 + cg1_tidx).load() + (sdh_red + 128 + cg1_tidx).load() + (sdh_red + 192 + cg1_tidx).load() - ) - sCumsumlog[cg1_tidx, 0, gate_idx] = sCumsumlog[cg1_tidx, 0, gate_idx] + pq_sum - bars.mb_gate_done[gate_idx].arrive() - bars.mb_dgate_cg1_ready[0].arrive() - - # ---- NEXT-CHUNK dH prep (>= 2: single-chunk tiles have no - # consumer and the dh_acc_ready consume would starve the drain) ---- - if sk_nt >= 2: - dh_idx = dh_acc_index.idx - bars.mb_dh_acc_ready[dh_idx].wait(dh_acc_index.phase) - dh_acc_index = advance(dh_acc_index, cfg.tmem_dh_acc_stages) - dh_done_idx = dh_idx - dh_inp_idx = dh_inp_index.idx - bars.mb_dh_inp_done[dh_inp_idx].wait(dh_inp_index.phase) - dh_inp_index = advance(dh_inp_index, cfg.tmem_dh_inp_stages) - dh_regs = [[cutlass.Float32(0.0) for _ in range(num_state_subs)] for _ in range(32)] - # all loads issue before the first store: a TMEM store - # between loads pins ptxas to one load latency per sub - dh_vecs = [] - for sub in cutlass.range_constexpr(num_state_subs): - dh_vecs.append( - nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dh_col + sub * ldtm_width, cutlass.Float32), num=32) - ) - for sub in cutlass.range_constexpr(num_state_subs): - for k in cutlass.range_constexpr(32): - dh_regs[k][sub] = dh_vecs[sub][k] - - dh_f16 = [fp32_to_fp16(dh_regs[2 * j][sub], dh_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dh_inp_col + sub * sttm_width, cutlass.Int32), - cutlass.Vector.from_elements(tuple(dh_f16), cutlass.Int32), - ) - nvvm.tcgen05_wait("store") - bars.mb_dh_inp_ready[dh_inp_idx].arrive() - - # ---- dK s-path fold: read while the dM-terms GEMMs run ------------ - if cend >= SFIRST_MIN: - bars.mb_dk_spath_ready[0].wait(cg1_dk_spath_rdy.phase) - cg1_dk_spath_rdy = advance(cg1_dk_spath_rdy, 1) - dk_spath_vecs = [] - for sub in cutlass.range_constexpr(2): - dk_spath_vecs.append( - nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dk_spath_col, cutlass.Float32), - num=8, - ) - ) - nvvm.tcgen05_wait("load") - dk_spath_regs = [] - for sub in cutlass.range_constexpr(2): - dk_spath_row = [] - for j in cutlass.range_constexpr(16): - n0, n1 = fmul2(dk_spath_vecs[sub][2 * j], dk_spath_vecs[sub][2 * j + 1], gCumprodNeg[2 * j], gCumprodNeg[2 * j + 1]) - dk_spath_row += [n0, n1] - dk_spath_regs.append(dk_spath_row) - - bars.mb_dk_total_ready[0].wait(cg1_dk_total_rdy.phase) - cg1_dk_total_rdy = advance(cg1_dk_total_rdy, 1) - dmr_vecs = [] - for sub in cutlass.range_constexpr(2): - dmr_vecs.append( - nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_col, cutlass.Float32), num=8) - ) - nvvm.tcgen05_wait("load") - bars.mb_dk_total_done[0].arrive() - for sub in cutlass.range_constexpr(2): - dk_sum = [dmr_vecs[sub][k] + dk_spath_regs[sub][k] for k in range(32)] - for m0 in cutlass.range_constexpr(4): - dk_f16 = [fp32_to_fp16(dk_sum[8 * m0 + 2 * j], dk_sum[8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - nvvm.stmatrix( - ( - sdK_base - + dk_stg_idx * dk_stage_elems - + ov_slab - + (ov_tok + m0 * 16) * 64 - + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - ).raw_ptr(), - dk_f16, - nvvm.MMALayout.COL, - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_dk_tmastg_ready[dk_stg_idx].arrive() - if cend < SFIRST_MIN: - bars.mb_dk_total_ready[0].wait(cg1_dk_total_rdy.phase) - cg1_dk_total_rdy = advance(cg1_dk_total_rdy, 1) - dmr_vecs = [] - for sub in cutlass.range_constexpr(2): - dmr_vecs.append( - nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_col, cutlass.Float32), num=8) - ) - nvvm.tcgen05_wait("load") - bars.mb_dk_total_done[0].arrive() - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - dk_f16 = [fp32_to_fp16(dmr_vecs[sub][8 * m0 + 2 * j], dmr_vecs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - nvvm.stmatrix( - ( - sdK_base - + dk_stg_idx * dk_stage_elems - + ov_slab - + (ov_tok + m0 * 16) * 64 - + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - ).raw_ptr(), - dk_f16, - nvvm.MMALayout.COL, - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_dk_tmastg_ready[dk_stg_idx].arrive() - - # ---- dH prep, sdH restage half (the dK readout above fires - # dk_total_done before this sweep) -------------------------- - if sk_nt >= 2: - # the sdH overwrite below waits only CG0's hdh read - bars.mb_hdh_done[0].wait(cg1_hdh_index.phase) - cg1_hdh_index = advance(cg1_hdh_index, 1) - bars.mb_dhs_done[0].wait(dhs_index.phase) - dhs_index = advance(dhs_index, 1) - for b in cutlass.range_constexpr(2): - for hh in cutlass.range_constexpr(2): - dhs_vec = nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + b * 16) << 16) + tmem_dh_col + hh * 64, cutlass.Float32), - num=8, - ) - dhs_f16 = [fp32_to_fp16(dhs_vec[2 * j], dhs_vec[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] - for c in cutlass.range_constexpr(4): - dhs_row = hh * 64 + ov_tok + c * 16 - nvvm.stmatrix( - cutlass.inttoptr( - sdH_base_int + ((cg1_tidx // 64) * cfg.d_k * 64 + dhs_row * 64 + swizzle_xor_128b(dhs_row, ov_col + b * 16)) * 2, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [dhs_f16[c * 4 + 0], dhs_f16[c * 4 + 1], dhs_f16[c * 4 + 2], dhs_f16[c * 4 + 3]], - nvvm.MMALayout.COL, - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_dhs_ready[0].arrive() - - if sk_nt < 2: - cg1_hdh_index = advance(cg1_hdh_index, 1) - - # ---- chunks NT-2 .. 0 (backward): full body ---------------------- - for rev_idx in cutlass.range(0 if cfg.use_dht else 1, sk_nt): - chunk_idx = cend - 1 - rev_idx - gate_idx = gate_index.idx - bars.mb_gate_ready[gate_idx].wait(gate_index.phase) - gate_index = advance(gate_index, cfg.smem_gate_stages) - - beta_idx = cg1_beta_index.idx - bars.mb_beta_ready[beta_idx].wait(cg1_beta_index.phase) - cg1_beta_index = advance(cg1_beta_index, cfg.smem_beta_stages) - - # ---- dH rescale: dH *= this chunk's cumprod ---- - cumprod_top = sCumprod[sCumprod.shape[0] - 1, 0, gate_idx] - dhr_vecs = [] - for sub in cutlass.range_constexpr(num_state_subs): - dhr_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dh_col + sub * ldtm_width, cutlass.Float32), num=32)) - for sub in cutlass.range_constexpr(num_state_subs): - dhr_scaled = [] - for j in cutlass.range_constexpr(16): - h0, h1 = fmul2(dhr_vecs[sub][2 * j], dhr_vecs[sub][2 * j + 1], cumprod_top, cumprod_top) - dhr_scaled += [h0, h1] - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dh_col + sub * ldtm_width, cutlass.Float32), - cutlass.Vector.from_elements(tuple(dhr_scaled), cutlass.Float32), - ) - nvvm.tcgen05_wait("store") - bars.mb_dh_acc_done[dh_done_idx].arrive() + bars.mb_dstate_scale_acc_done[dstate_done_idx].arrive() num_vals = 32 - gDecayScale = [] - last_cumsumlog = sCumsumlog[cfg.b_t - 1, 0, gate_idx] - for k in cutlass.range_constexpr(num_vals): - gDecayScale.append(cute.math.exp2(last_cumsumlog - sCumsumlog[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx], fastmath=True)) - gCumprod = [] - for k in cutlass.range_constexpr(num_vals): - gCumprod.append(sCumprod[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) + cumprod_fp32 = [] + for g in cutlass.range_constexpr(8): + for b in cutlass.range_constexpr(2): + cumprod_fp32.append(sCumprod[(lane_id % 4) * 2 + g * 8 + b, 0, gate_idx]) + cumprod_vals = [cumprod_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] - # ---- dO' restage: dO * gCumprod * scale -> shared_inp TMEM ---- - bars.mb_dop_inp_done[0].wait(dop_inp_free.phase) - dop_inp_free = advance(dop_inp_free, 1) + # ---- dO' restage: dO * cumprod_vals * scale -> shared_inp TMEM ----------- do_idx = do_index.idx bars.mb_do_ready[do_idx].wait(do_index.phase) do_index = advance(do_index, cfg.smem_do_stages) @@ -3434,102 +2613,122 @@ def _compute1_warp( for c in cutlass.range_constexpr(8): m0 = cutlass.const_expr(c % 4) sub = cutlass.const_expr(c // 4) - do_f16 = nvvm.ldmatrix( - (sdO_base + do_idx * do_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16)).raw_ptr(), + do_frag = nvvm.ldmatrix( + ( + sdO_base + + do_idx * do_stage_elems + + frag_slab_off + + (frag_row + m0 * 16) * 64 + + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) + ).raw_ptr(), 4, nvvm.MMALayout.COL, ) for i in cutlass.range_constexpr(4): - lo, hi = f16x2_to_f32(do_f16[i], dtype=cfg.io_dtype) - p0, p1 = fmul2(lo, hi, gCumprod[8 * m0 + 2 * i], gCumprod[8 * m0 + 2 * i + 1]) + lo, hi = f16x2_to_f32(do_frag[i], dtype=cfg.io_dtype) + p0, p1 = fmul2(lo, hi, cumprod_vals[8 * m0 + 2 * i], cumprod_vals[8 * m0 + 2 * i + 1]) do_regs[8 * m0 + 2 * i][sub], do_regs[8 * m0 + 2 * i + 1][sub] = fmul2(p0, p1, scale, scale) for sub in cutlass.range_constexpr(2): do_pack = [fp32_to_fp16(do_regs[2 * j][sub], do_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dop_col, cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_do_prime_col, cutlass.Int32), cutlass.Vector.from_elements(tuple(do_pack), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_dop_inp_ready[0].arrive() + bars.mb_do_prime_inp_ready[0].arrive() - # ---- dV inter: in-place decay scale ------------------------- - bars.mb_du_scale_ready[0].wait(cg1_du_scale_rdy.phase) - cg1_du_scale_rdy = advance(cg1_du_scale_rdy, 1) - dv_ptrs = [] - dv_vecs = [] - for sub in cutlass.range_constexpr(2): - dv_ptrs.append(nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_col, cutlass.Float32)) - dv_vecs.append(nvvm.tcgen05_ld("16x256b", dv_ptrs[sub], num=8)) - for sub in cutlass.range_constexpr(2): - dv_scaled = [] - for j in cutlass.range_constexpr(16): - s0, s1 = fmul2(dv_vecs[sub][2 * j], dv_vecs[sub][2 * j + 1], gDecayScale[2 * j], gDecayScale[2 * j + 1]) - dv_scaled += [s0, s1] - nvvm.tcgen05_st("16x256b", dv_ptrs[sub], cutlass.Vector.from_elements(tuple(dv_scaled), cutlass.Float32)) - nvvm.tcgen05_wait("store") - bars.mb_du_scale_done[0].arrive() + # ---- dV inter: in-place decay scale ---------------------------------- + if have_dstate: + last_cumsumlog = sCumsumlog[cfg.b_t - 1, 0, gate_idx] + col_cs_fp32 = [] + for g in cutlass.range_constexpr(8): + for b in cutlass.range_constexpr(2): + col_cs_fp32.append(sCumsumlog[(lane_id % 4) * 2 + g * 8 + b, 0, gate_idx]) + decay_scale_fp32 = [] + for i in cutlass.range_constexpr(16): + decay_scale_fp32.append(cute.math.exp2(last_cumsumlog - col_cs_fp32[i], fastmath=True)) + decay_scale_vals = [decay_scale_fp32[(k // 4) * 2 + (k % 2)] for k in range(num_vals)] + bars.mb_du_scale_acc_ready[0].wait(cg1_du_scale_ready.phase) + cg1_du_scale_ready = advance(cg1_du_scale_ready, 1) + dv_ptrs = [] + dv_vecs = [] + for sub in cutlass.range_constexpr(2): + dv_ptrs.append(nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32)) + dv_vecs.append(nvvm.tcgen05_ld("16x256b", dv_ptrs[sub], num=8)) + for sub in cutlass.range_constexpr(2): + dv_scaled = [] + for j in cutlass.range_constexpr(16): + s0, s1 = fmul2(dv_vecs[sub][2 * j], dv_vecs[sub][2 * j + 1], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1]) + dv_scaled += [s0, s1] + nvvm.tcgen05_st("16x256b", dv_ptrs[sub], cutlass.Vector.from_elements(tuple(dv_scaled), cutlass.Float32)) + nvvm.tcgen05_wait("store") + bars.mb_du_scale_acc_done[0].arrive() - # ---- v - k*state: delta = V - cumprod*(K @ S), in place ---- + # ---- Y staging: Y = V - cumprod*(K @ state) -> f16 TMEM slots ----------- v_idx = v_index.idx bars.mb_v_ready[v_idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) - if chunk_idx >= S_MIN: - v_regs = [[cutlass.Float32(0.0), cutlass.Float32(0.0)] for _ in range(32)] - gks_regs = [[cutlass.Float32(0.0), cutlass.Float32(0.0)] for _ in range(32)] + if chunk_idx >= FIRST_STATE_CHUNK: + v_frags = [[cutlass.Int32(0), cutlass.Int32(0)] for _ in range(16)] for c in cutlass.range_constexpr(8): m0 = cutlass.const_expr(c % 4) sub = cutlass.const_expr(c // 4) - v_f16 = nvvm.ldmatrix( - (sV_base + v_idx * v_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16)).raw_ptr(), + v_frag = nvvm.ldmatrix( + ( + sV_base + + v_idx * v_stage_elems + + frag_slab_off + + (frag_row + m0 * 16) * 64 + + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) + ).raw_ptr(), 4, nvvm.MMALayout.COL, ) for i in cutlass.range_constexpr(4): - lo, hi = f16x2_to_f32(v_f16[i], dtype=cfg.io_dtype) - v_regs[8 * m0 + 2 * i][sub] = lo - v_regs[8 * m0 + 2 * i + 1][sub] = hi - bars.mb_ks_acc_ready[0].wait(cg1_ks_rdy.phase) - cg1_ks_rdy = advance(cg1_ks_rdy, 1) + v_frags[4 * m0 + i][sub] = v_frag[i] + bars.mb_k_state_acc_ready[0].wait(cg1_k_state_ready.phase) + cg1_k_state_ready = advance(cg1_k_state_ready, 1) for sub in cutlass.range_constexpr(2): - ks_vec = nvvm.tcgen05_ld( + k_state_vec = nvvm.tcgen05_ld( "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_ks_col, cutlass.Float32), + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_k_state_col, cutlass.Float32), num=8, ) + g_k_state_pack = [] + y_pack = [] for j in cutlass.range_constexpr(16): - g0, g1 = fmul2(ks_vec[2 * j], ks_vec[2 * j + 1], gCumprod[2 * j], gCumprod[2 * j + 1]) - gks_regs[2 * j][sub] = g0 - gks_regs[2 * j + 1][sub] = g1 - v_regs[2 * j][sub] = v_regs[2 * j][sub] - g0 - v_regs[2 * j + 1][sub] = v_regs[2 * j + 1][sub] - g1 - for sub in cutlass.range_constexpr(2): - y_pack = [fp32_to_fp16(v_regs[2 * j][sub], v_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + g0, g1 = fmul2(k_state_vec[2 * j], k_state_vec[2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) + g_k_state_word = fp32_to_fp16(g0, g1, dtype=cfg.io_dtype) + g_k_state_pack.append(g_k_state_word) + y_pack.append(sub_f16x2(v_frags[j][sub], g_k_state_word, cfg.io_dtype)) nvvm.tcgen05_st( "16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), cutlass.Vector.from_elements(tuple(y_pack), cutlass.Int32), ) - gks_pack = [fp32_to_fp16(gks_regs[2 * j][sub], gks_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_gks_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(gks_pack), cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_g_k_state_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(g_k_state_pack), cutlass.Int32), ) - if chunk_idx < S_MIN: + if chunk_idx < FIRST_STATE_CHUNK: for sub in cutlass.range_constexpr(2): v_pack = [] for m0 in cutlass.range_constexpr(4): - v_f16 = nvvm.ldmatrix( + v_frag = nvvm.ldmatrix( ( - sV_base + v_idx * v_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) + sV_base + + v_idx * v_stage_elems + + frag_slab_off + + (frag_row + m0 * 16) * 64 + + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) ).raw_ptr(), 4, nvvm.MMALayout.COL, ) for i in cutlass.range_constexpr(4): - v_lo, v_hi = f16x2_to_f32(v_f16[i], dtype=cfg.io_dtype) - v_pack.append(fp32_to_fp16(v_lo, v_hi, dtype=cfg.io_dtype)) + v_pack.append(v_frag[i]) nvvm.tcgen05_st( "16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), @@ -3538,13 +2737,11 @@ def _compute1_warp( nvvm.tcgen05_wait("store") bars.mb_y_ready[0].arrive() - # ---- dU restage: dv acc ------------------------------------ - bars.mb_dyp_inp_done[0].wait(dyp_inp_free.phase) - dyp_inp_free = advance(dyp_inp_free, 1) - bars.mb_du_total_ready[0].wait(cg1_du_total_rdy.phase) - cg1_du_total_rdy = advance(cg1_du_total_rdy, 1) + # ---- dU restage: dV acc ---------------------------------------------- + bars.mb_du_total_acc_ready[0].wait(cg1_du_total_ready.phase) + cg1_du_total_ready = advance(cg1_du_total_ready, 1) for sub in cutlass.range_constexpr(2): - du_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_col, cutlass.Float32), num=8) + du_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), num=8) du_pack = [fp32_to_fp16(du_vec[2 * j], du_vec[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "16x128b", @@ -3553,30 +2750,34 @@ def _compute1_warp( ) nvvm.tcgen05_wait("store") bars.mb_du_inp_ready[0].arrive() - bars.mb_do_cg1_done[do_idx].arrive() - # ---- U readout --------------------------------------------- - bars.mb_u_acc_ready[0].wait(cg1_u_rdy.phase) - cg1_u_rdy = advance(cg1_u_rdy, 1) + # ---- U readout ------------------------------------------------------- + bars.mb_u_acc_ready[0].wait(cg1_u_ready.phase) + cg1_u_ready = advance(cg1_u_ready, 1) u_regs = [] for sub in cutlass.range_constexpr(2): u_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_u_col, cutlass.Float32), num=8) u_regs.append([u_vec[k] for k in range(32)]) for sub in cutlass.range_constexpr(2): for m0 in cutlass.range_constexpr(4): - u_f16 = [fp32_to_fp16(u_regs[sub][8 * m0 + 2 * j], u_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + u_pack = [fp32_to_fp16(u_regs[sub][8 * m0 + 2 * j], u_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] nvvm.stmatrix( - (sV_base + v_idx * v_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16)).raw_ptr(), - u_f16, + ( + sV_base + + v_idx * v_stage_elems + + frag_slab_off + + (frag_row + m0 * 16) * 64 + + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) + ).raw_ptr(), + u_pack, nvvm.MMALayout.COL, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_u_ready[0].arrive() - bars.mb_v_cg1_done[v_idx].arrive() - # ---- dY ---------------------------------------------------- - bars.mb_dy_acc_ready[0].wait(cg1_dy_rdy.phase) - cg1_dy_rdy = advance(cg1_dy_rdy, 1) + # ---- dY -------------------------------------------------------------- + bars.mb_dy_acc_ready[0].wait(cg1_dy_ready.phase) + cg1_dy_ready = advance(cg1_dy_ready, 1) dy_regs = [] for sub in cutlass.range_constexpr(2): dy_vec = nvvm.tcgen05_ld( @@ -3586,13 +2787,13 @@ def _compute1_warp( ) dy_regs.append([dy_vec[k] for k in range(32)]) - # ---- dY' = -gCumprod * dY -> f16 shared_inp ---------------- + # ---- dY' = -cumprod_vals * dY -> f16 shared_inp -------------------------- neg_one = cutlass.Float32(-1.0) - gCumprodNeg = [gCumprod[k] * neg_one for k in range(32)] + cumprod_neg_vals = [cumprod_vals[k] * neg_one for k in range(32)] for sub in cutlass.range_constexpr(2): dyp = [] for j in cutlass.range_constexpr(16): - n0, n1 = fmul2(dy_regs[sub][2 * j], dy_regs[sub][2 * j + 1], gCumprodNeg[2 * j], gCumprodNeg[2 * j + 1]) + n0, n1 = fmul2(dy_regs[sub][2 * j], dy_regs[sub][2 * j + 1], cumprod_neg_vals[2 * j], cumprod_neg_vals[2 * j + 1]) dyp += [n0, n1] dyp_pack = [fp32_to_fp16(dyp[2 * j], dyp[2 * j + 1], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( @@ -3603,7 +2804,7 @@ def _compute1_warp( nvvm.tcgen05_wait("store") bars.mb_dyp_inp_ready[0].arrive() - # ---- dV staging: dV = dY -> the sdV slot ------------------- + # ---- dV staging: dV = dY -> the sdV slot ----------------------------- dv_stg_idx = dv_index.idx bars.mb_dv_tmastg_done[dv_stg_idx].wait(dv_index.phase) dv_index = advance(dv_index, cfg.smem_dv_stages) @@ -3611,55 +2812,62 @@ def _compute1_warp( sdv_done_index = advance(sdv_done_index, 1) for sub in cutlass.range_constexpr(2): for m0 in cutlass.range_constexpr(4): - dv_f16 = [fp32_to_fp16(dy_regs[sub][8 * m0 + 2 * j], dy_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + dv_pack = [fp32_to_fp16(dy_regs[sub][8 * m0 + 2 * j], dy_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] nvvm.stmatrix( ( - sdV_base + dv_stg_idx * dv_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) + sdV_base + + dv_stg_idx * dv_stage_elems + + frag_slab_off + + (frag_row + m0 * 16) * 64 + + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) ).raw_ptr(), - dv_f16, + dv_pack, nvvm.MMALayout.COL, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dv_tmastg_ready[dv_stg_idx].arrive() - dk_stg_idx = dk_index.idx - bars.mb_dk_tmastg_done[dk_stg_idx].wait(dk_index.phase) - dk_index = advance(dk_index, cfg.smem_dk_stages) dq_stg_idx = dq_index.idx bars.mb_dq_tmastg_done[dq_stg_idx].wait(dq_index.phase) dq_index = advance(dq_index, cfg.smem_dq_stages) sred = sred_base + dq_stg_idx * (dq_stage_elems // 2) - # ---- dBeta/dGate v-terms: dbeta_t += rowsum(dV ⊙ Y)_t / beta_t ---- + # ---- dBeta/dGate V-terms: dBeta_t += rowsum(dV ⊙ Y)_t / Beta_t ------- part_y = [cutlass.Float32(0.0)] * 16 for sub in cutlass.range_constexpr(2): - y_pk = nvvm.tcgen05_ld("16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), num=8) + y_vec = nvvm.tcgen05_ld("16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), num=8) for j in cutlass.range_constexpr(16): - lo, hi = f16x2_to_f32(y_pk[j], dtype=cfg.io_dtype) - kk0 = cutlass.const_expr(2 * j) - j0 = cutlass.const_expr((kk0 // 4) * 2 + (kk0 % 2)) + lo, hi = f16x2_to_f32(y_vec[j], dtype=cfg.io_dtype) + frag_e0 = cutlass.const_expr(2 * j) + part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) if cutlass.const_expr(sub == 0 and j % 2 == 0): - part_y[j0], part_y[j0 + 1] = fmul2(dy_regs[sub][kk0], dy_regs[sub][kk0 + 1], lo, hi) + part_y[part_e0], part_y[part_e0 + 1] = fmul2(dy_regs[sub][frag_e0], dy_regs[sub][frag_e0 + 1], lo, hi) else: - part_y[j0], part_y[j0 + 1] = ffma2(dy_regs[sub][kk0], dy_regs[sub][kk0 + 1], lo, hi, part_y[j0], part_y[j0 + 1]) - py_lo, py_hi = _warp_reduce_scatter_frag16(part_y, lane_id) + part_y[part_e0], part_y[part_e0 + 1] = ffma2( + dy_regs[sub][frag_e0], dy_regs[sub][frag_e0 + 1], lo, hi, part_y[part_e0], part_y[part_e0 + 1] + ) + py_lo, py_hi = warp_reduce_scatter_frag_16_elems(part_y, lane_id) vt_tok0 = (lane_id // 4) * 8 + (lane_id % 4) * 2 - if chunk_idx >= S_MIN: + if chunk_idx >= FIRST_STATE_CHUNK: part_g = [cutlass.Float32(0.0)] * 16 for sub in cutlass.range_constexpr(2): - g_pk = nvvm.tcgen05_ld("16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_gks_col, cutlass.Int32), num=8) + g_k_state_vec = nvvm.tcgen05_ld( + "16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_g_k_state_col, cutlass.Int32), num=8 + ) for j in cutlass.range_constexpr(16): - lo, hi = f16x2_to_f32(g_pk[j], dtype=cfg.io_dtype) - kk0 = cutlass.const_expr(2 * j) - j0 = cutlass.const_expr((kk0 // 4) * 2 + (kk0 % 2)) + lo, hi = f16x2_to_f32(g_k_state_vec[j], dtype=cfg.io_dtype) + frag_e0 = cutlass.const_expr(2 * j) + part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) if cutlass.const_expr(sub == 0 and j % 2 == 0): - part_g[j0], part_g[j0 + 1] = fmul2(dy_regs[sub][kk0], dy_regs[sub][kk0 + 1], lo, hi) + part_g[part_e0], part_g[part_e0 + 1] = fmul2(dy_regs[sub][frag_e0], dy_regs[sub][frag_e0 + 1], lo, hi) else: - part_g[j0], part_g[j0 + 1] = ffma2(dy_regs[sub][kk0], dy_regs[sub][kk0 + 1], lo, hi, part_g[j0], part_g[j0 + 1]) - pg_lo, pg_hi = _warp_reduce_scatter_frag16(part_g, lane_id) - (sred + cg1w * 64 + vt_tok0).store(py_lo) - (sred + cg1w * 64 + vt_tok0 + 1).store(py_hi) - (sred + 256 + cg1w * 64 + vt_tok0).store(pg_lo) - (sred + 256 + cg1w * 64 + vt_tok0 + 1).store(pg_hi) + part_g[part_e0], part_g[part_e0 + 1] = ffma2( + dy_regs[sub][frag_e0], dy_regs[sub][frag_e0 + 1], lo, hi, part_g[part_e0], part_g[part_e0 + 1] + ) + pg_lo, pg_hi = warp_reduce_scatter_frag_16_elems(part_g, lane_id) + (sred + cg1_warp_id * 64 + vt_tok0).store(py_lo) + (sred + cg1_warp_id * 64 + vt_tok0 + 1).store(py_hi) + (sred + 256 + cg1_warp_id * 64 + vt_tok0).store(pg_lo) + (sred + 256 + cg1_warp_id * 64 + vt_tok0 + 1).store(pg_hi) nvvm.barrier_cta_sync_aligned(cfg.cg1_barrier_id, thread_count=cfg.cg1_barrier_threads) if cg1_tidx < 64: binv_t = cute.math.rcp(sBeta[cg1_tidx, 0, beta_idx] + cutlass.Float32(1e-10), approx=True, ftz=True) @@ -3667,9 +2875,9 @@ def _compute1_warp( gsum = (sred + 256 + cg1_tidx).load() + (sred + 320 + cg1_tidx).load() + (sred + 384 + cg1_tidx).load() + (sred + 448 + cg1_tidx).load() sBeta[cg1_tidx, 0, beta_idx] = ysum * binv_t sCumsumlog[cg1_tidx, 0, gate_idx] = cutlass.Float32(0.0) - gsum - if chunk_idx < S_MIN: - (sred + cg1w * 64 + vt_tok0).store(py_lo) - (sred + cg1w * 64 + vt_tok0 + 1).store(py_hi) + if chunk_idx < FIRST_STATE_CHUNK: + (sred + cg1_warp_id * 64 + vt_tok0).store(py_lo) + (sred + cg1_warp_id * 64 + vt_tok0 + 1).store(py_hi) nvvm.barrier_cta_sync_aligned(cfg.cg1_barrier_id, thread_count=cfg.cg1_barrier_threads) if cg1_tidx < 64: binv_t = cute.math.rcp(sBeta[cg1_tidx, 0, beta_idx] + cutlass.Float32(1e-10), approx=True, ftz=True) @@ -3678,26 +2886,21 @@ def _compute1_warp( sCumsumlog[cg1_tidx, 0, gate_idx] = cutlass.Float32(0.0) bars.mb_beta_done[beta_idx].arrive() bars.mb_dbeta_cg1_ready[0].arrive() - # sred reads must land before the dq stmatrix below reuses the stage nvvm.barrier_cta_sync_aligned(cfg.cg1_barrier_id, thread_count=cfg.cg1_barrier_threads) - # ---- Q fragments held in registers over the dq stage (sQ - # free once read; the dQ dot consumes them -- no TMEM trip) ---- + + # ---- Q fragments held in registers ----------------------------------- q_frag = [] for sub in cutlass.range_constexpr(2): q_words = [] for m0 in cutlass.range_constexpr(4): q_f16 = nvvm.ldmatrix( - (sQ_base + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16)).raw_ptr(), + (sQ_base + frag_slab_off + (frag_row + m0 * 16) * 64 + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16)).raw_ptr(), 4, nvvm.MMALayout.COL, ) for i in cutlass.range_constexpr(4): - q_lo, q_hi = f16x2_to_f32(q_f16[i], dtype=cfg.io_dtype) - q_words.append(fp32_to_fp16(q_lo, q_hi, dtype=cfg.io_dtype)) + q_words.append(q_f16[i]) q_frag.append(q_words) - # the store+wait order the sQ LDSM reads before the release arrive - # (a bare register consume can be hoisted past by ptxas); the dot - # below reads the registers, so the TMEM read-back is still gone nvvm.tcgen05_st( "16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_col, cutlass.Int32), @@ -3706,231 +2909,223 @@ def _compute1_warp( nvvm.tcgen05_wait("store") bars.mb_q_cg1_done[0].arrive() - # ---- dq final read -> sdQ (output staging; the fragments are - # held for the dQ dot below) --------- - bars.mb_dq_acc_total_ready[0].wait(dq_total_rdy_index.phase) - dq_total_rdy_index = advance(dq_total_rdy_index, 1) - dq_frag = [] + # ---- dQ final read -> sdQ -------------------------------------------- + bars.mb_dq_acc_total_ready[0].wait(dq_total_ready_index.phase) + dq_total_ready_index = advance(dq_total_ready_index, 1) + dq_regs = [] for sub in cutlass.range_constexpr(2): - dqv = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dh_inp_col, cutlass.Float32), num=8) - dq_frag.append([dqv[k] for k in range(32)]) + dq_vec = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dstate_inp_col, cutlass.Float32), num=8) + dq_regs.append([dq_vec[k] for k in range(32)]) for m0 in cutlass.range_constexpr(4): - frag_addr = ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - dq_f16 = [fp32_to_fp16(dqv[8 * m0 + 2 * j], dqv[8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - nvvm.stmatrix((sdQ_base + dq_stg_idx * dq_stage_elems + frag_addr).raw_ptr(), dq_f16, nvvm.MMALayout.COL) + frag_addr = frag_slab_off + (frag_row + m0 * 16) * 64 + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) + dq_pack = [fp32_to_fp16(dq_vec[8 * m0 + 2 * j], dq_vec[8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + nvvm.stmatrix((sdQ_base + dq_stg_idx * dq_stage_elems + frag_addr).raw_ptr(), dq_pack, nvvm.MMALayout.COL) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dq_acc_total_done[0].arrive() bars.mb_dq_tmastg_ready[dq_stg_idx].arrive() - # ---- dQ dot (part_q): fragment dot of the held dq acc with the - # staged Q^T, added to dGate FIRST; CG0 adds after parts_ready ---- + # ---- dQ dot (part_q) ------------------------------------------------- part_q = [cutlass.Float32(0.0)] * 16 for sub in cutlass.range_constexpr(2): for m0 in cutlass.range_constexpr(4): for i in cutlass.range_constexpr(4): q_lo, q_hi = f16x2_to_f32(q_frag[sub][4 * m0 + i], dtype=cfg.io_dtype) - kk0 = cutlass.const_expr(8 * m0 + 2 * i) - j0 = cutlass.const_expr((kk0 // 4) * 2 + (kk0 % 2)) + frag_e0 = cutlass.const_expr(8 * m0 + 2 * i) + part_e0 = cutlass.const_expr((frag_e0 // 4) * 2 + (frag_e0 % 2)) if cutlass.const_expr(sub == 0 and i % 2 == 0): - part_q[j0], part_q[j0 + 1] = fmul2(dq_frag[sub][kk0], dq_frag[sub][kk0 + 1], q_lo, q_hi) + part_q[part_e0], part_q[part_e0 + 1] = fmul2(dq_regs[sub][frag_e0], dq_regs[sub][frag_e0 + 1], q_lo, q_hi) else: - part_q[j0], part_q[j0 + 1] = ffma2(dq_frag[sub][kk0], dq_frag[sub][kk0 + 1], q_lo, q_hi, part_q[j0], part_q[j0 + 1]) - amq_lo, amq_hi = _warp_reduce_scatter_frag16(part_q, lane_id) + part_q[part_e0], part_q[part_e0 + 1] = ffma2( + dq_regs[sub][frag_e0], dq_regs[sub][frag_e0 + 1], q_lo, q_hi, part_q[part_e0], part_q[part_e0 + 1] + ) + part_q_lo, part_q_hi = warp_reduce_scatter_frag_16_elems(part_q, lane_id) tok0 = (lane_id // 4) * 8 + (lane_id % 4) * 2 - (sdh_red + (cg1_tidx // 32) * 64 + tok0).store(amq_lo) - (sdh_red + (cg1_tidx // 32) * 64 + tok0 + 1).store(amq_hi) + (sdstate_red + (cg1_tidx // 32) * 64 + tok0).store(part_q_lo) + (sdstate_red + (cg1_tidx // 32) * 64 + tok0 + 1).store(part_q_hi) nvvm.barrier_cta_sync_aligned(cfg.cg1_barrier_id, thread_count=cfg.cg1_barrier_threads) if cg1_tidx < 64: - pq_sum = (sdh_red + cg1_tidx).load() + (sdh_red + 64 + cg1_tidx).load() + (sdh_red + 128 + cg1_tidx).load() + (sdh_red + 192 + cg1_tidx).load() + pq_sum = ( + (sdstate_red + cg1_tidx).load() + + (sdstate_red + 64 + cg1_tidx).load() + + (sdstate_red + 128 + cg1_tidx).load() + + (sdstate_red + 192 + cg1_tidx).load() + ) sCumsumlog[cg1_tidx, 0, gate_idx] = sCumsumlog[cg1_tidx, 0, gate_idx] + pq_sum bars.mb_gate_done[gate_idx].arrive() bars.mb_dgate_cg1_ready[0].arrive() - # ---- NEXT-CHUNK dH prep ------------------------------------ + # ---- NEXT-CHUNK dstate prep ---------------------------------------------- if chunk_idx >= wstart + 1: - dh_idx = dh_acc_index.idx - bars.mb_dh_acc_ready[dh_idx].wait(dh_acc_index.phase) - dh_acc_index = advance(dh_acc_index, cfg.tmem_dh_acc_stages) - dh_done_idx = dh_idx - dh_inp_idx = dh_inp_index.idx - bars.mb_dh_inp_done[dh_inp_idx].wait(dh_inp_index.phase) - dh_inp_index = advance(dh_inp_index, cfg.tmem_dh_inp_stages) - dh_regs = [[cutlass.Float32(0.0) for _ in range(num_state_subs)] for _ in range(32)] + dstate_idx = dstate_acc_index.idx + bars.mb_dstate_acc_ready[dstate_idx].wait(dstate_acc_index.phase) + dstate_acc_index = advance(dstate_acc_index, cfg.tmem_dstate_acc_stages) + dstate_done_idx = dstate_idx + dstate_inp_idx = dstate_inp_index.idx + bars.mb_dstate_inp_done[dstate_inp_idx].wait(dstate_inp_index.phase) + dstate_inp_index = advance(dstate_inp_index, cfg.tmem_dstate_inp_stages) + dstate_regs = [[cutlass.Float32(0.0) for _ in range(num_state_subs)] for _ in range(32)] for sub in cutlass.range_constexpr(num_state_subs): - dh_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dh_col + sub * ldtm_width, cutlass.Float32), num=32) + dstate_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 k in cutlass.range_constexpr(32): - dh_regs[k][sub] = dh_vec[k] + dstate_regs[k][sub] = dstate_vec[k] - dh_f16 = [fp32_to_fp16(dh_regs[2 * j][sub], dh_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + dstate_pack = [fp32_to_fp16(dstate_regs[2 * j][sub], dstate_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dh_inp_col + sub * sttm_width, cutlass.Int32), - cutlass.Vector.from_elements(tuple(dh_f16), cutlass.Int32), + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dstate_inp_col + sub * sttm_width, cutlass.Int32), + cutlass.Vector.from_elements(tuple(dstate_pack), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_dh_inp_ready[dh_inp_idx].arrive() + bars.mb_dstate_inp_ready[dstate_inp_idx].arrive() - # ---- dK s-path fold: read while the dM-terms GEMMs run ------------ - if chunk_idx >= S_MIN: - bars.mb_dk_spath_ready[0].wait(cg1_dk_spath_rdy.phase) - cg1_dk_spath_rdy = advance(cg1_dk_spath_rdy, 1) - dk_spath_vecs = [] + # ---- dK fold --------------------------------------------------------- + dk_stg_idx = dk_index.idx + bars.mb_dk_tmastg_done[dk_stg_idx].wait(dk_index.phase) + dk_index = advance(dk_index, cfg.smem_dk_stages) + if chunk_idx >= FIRST_STATE_CHUNK: + bars.mb_dk_state_path_acc_ready[0].wait(cg1_dk_state_path_ready.phase) + cg1_dk_state_path_ready = advance(cg1_dk_state_path_ready, 1) + dk_state_path_vecs = [] for sub in cutlass.range_constexpr(2): - dk_spath_vecs.append( + dk_state_path_vecs.append( nvvm.tcgen05_ld( "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dk_spath_col, cutlass.Float32), + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dk_state_path_col, cutlass.Float32), num=8, ) ) nvvm.tcgen05_wait("load") - dk_spath_regs = [] + dk_state_path_regs = [] for sub in cutlass.range_constexpr(2): - dk_spath_row = [] + dk_state_path_row = [] for j in cutlass.range_constexpr(16): - n0, n1 = fmul2(dk_spath_vecs[sub][2 * j], dk_spath_vecs[sub][2 * j + 1], gCumprodNeg[2 * j], gCumprodNeg[2 * j + 1]) - dk_spath_row += [n0, n1] - dk_spath_regs.append(dk_spath_row) + n0, n1 = fmul2(dk_state_path_vecs[sub][2 * j], dk_state_path_vecs[sub][2 * j + 1], cumprod_neg_vals[2 * j], cumprod_neg_vals[2 * j + 1]) + dk_state_path_row += [n0, n1] + dk_state_path_regs.append(dk_state_path_row) - bars.mb_dk_total_ready[0].wait(cg1_dk_total_rdy.phase) - cg1_dk_total_rdy = advance(cg1_dk_total_rdy, 1) + bars.mb_dk_total_acc_ready[0].wait(cg1_dk_total_ready.phase) + cg1_dk_total_ready = advance(cg1_dk_total_ready, 1) dmr_vecs = [] for sub in cutlass.range_constexpr(2): - dmr_vecs.append(nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_col, cutlass.Float32), num=8)) + dmr_vecs.append( + nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), num=8) + ) nvvm.tcgen05_wait("load") - bars.mb_dk_total_done[0].arrive() + bars.mb_dk_total_acc_done[0].arrive() for sub in cutlass.range_constexpr(2): - dk_sum = [dmr_vecs[sub][k] + dk_spath_regs[sub][k] for k in range(32)] + dk_sum = [dmr_vecs[sub][k] + dk_state_path_regs[sub][k] for k in range(32)] for m0 in cutlass.range_constexpr(4): - dk_f16 = [fp32_to_fp16(dk_sum[8 * m0 + 2 * j], dk_sum[8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + dk_pack = [fp32_to_fp16(dk_sum[8 * m0 + 2 * j], dk_sum[8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] nvvm.stmatrix( ( sdK_base + dk_stg_idx * dk_stage_elems - + ov_slab - + (ov_tok + m0 * 16) * 64 - + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) + + frag_slab_off + + (frag_row + m0 * 16) * 64 + + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) ).raw_ptr(), - dk_f16, + dk_pack, nvvm.MMALayout.COL, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dk_tmastg_ready[dk_stg_idx].arrive() - if chunk_idx < S_MIN: - bars.mb_dk_total_ready[0].wait(cg1_dk_total_rdy.phase) - cg1_dk_total_rdy = advance(cg1_dk_total_rdy, 1) + if chunk_idx < FIRST_STATE_CHUNK: + bars.mb_dk_total_acc_ready[0].wait(cg1_dk_total_ready.phase) + cg1_dk_total_ready = advance(cg1_dk_total_ready, 1) dmr_vecs = [] for sub in cutlass.range_constexpr(2): - dmr_vecs.append(nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_col, cutlass.Float32), num=8)) + dmr_vecs.append( + nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_dvdk_acc_col, cutlass.Float32), num=8) + ) nvvm.tcgen05_wait("load") - bars.mb_dk_total_done[0].arrive() + bars.mb_dk_total_acc_done[0].arrive() for sub in cutlass.range_constexpr(2): for m0 in cutlass.range_constexpr(4): - dk_f16 = [fp32_to_fp16(dmr_vecs[sub][8 * m0 + 2 * j], dmr_vecs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + dk_pack = [fp32_to_fp16(dmr_vecs[sub][8 * m0 + 2 * j], dmr_vecs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] nvvm.stmatrix( ( sdK_base + dk_stg_idx * dk_stage_elems - + ov_slab - + (ov_tok + m0 * 16) * 64 - + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) + + frag_slab_off + + (frag_row + m0 * 16) * 64 + + swizzle_xor_128b(frag_row + m0 * 16, frag_col + sub * 16) ).raw_ptr(), - dk_f16, + dk_pack, nvvm.MMALayout.COL, ) nvvm.fence_proxy("async.shared", space="cta") bars.mb_dk_tmastg_ready[dk_stg_idx].arrive() - # ---- dH prep, sdH restage half (the dK readout above fires - # dk_total_done before this sweep) -------------------------- + # ---- dstate prep --------------------------------------------------------- if chunk_idx >= wstart + 1: - # the sdH overwrite below waits only CG0's hdh read - bars.mb_hdh_done[0].wait(cg1_hdh_index.phase) - cg1_hdh_index = advance(cg1_hdh_index, 1) - bars.mb_dhs_done[0].wait(dhs_index.phase) - dhs_index = advance(dhs_index, 1) - dhs_vecs = [] + 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 hh in cutlass.range_constexpr(2): - dhs_vecs.append( - nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + b * 16) << 16) + tmem_dh_col + hh * 64, cutlass.Float32), - num=8, - ) + 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 b in cutlass.range_constexpr(2): - for hh in cutlass.range_constexpr(2): - dhs_vec = dhs_vecs[b * 2 + hh] - dhs_f16 = [fp32_to_fp16(dhs_vec[2 * j], dhs_vec[2 * j + 1], dtype=cfg.io_dtype) for j in range(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): - dhs_row = hh * 64 + ov_tok + c * 16 + dstate_smem_row = col_half * 64 + frag_row + c * 16 nvvm.stmatrix( cutlass.inttoptr( - sdH_base_int + ((cg1_tidx // 64) * cfg.d_k * 64 + dhs_row * 64 + swizzle_xor_128b(dhs_row, ov_col + b * 16)) * 2, + 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, ), - [dhs_f16[c * 4 + 0], dhs_f16[c * 4 + 1], dhs_f16[c * 4 + 2], dhs_f16[c * 4 + 3]], + [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_dhs_ready[0].arrive() + bars.mb_dstate_smem_ready[0].arrive() if chunk_idx < wstart + 1: - cg1_hdh_index = advance(cg1_hdh_index, 1) - - # ---- dH drain: with an initial state this is dL/dS0 ---- - if sk_nt > 0: - dh_idx = dh_acc_index.idx - bars.mb_dh_acc_ready[dh_idx].wait(dh_acc_index.phase) - dh_acc_index = advance(dh_acc_index, cfg.tmem_dh_acc_stages) - if cutlass.const_expr(cfg.use_initial_state): - # split-K: only the item owning chunk 0 drains dL/dS0 - if cutlass.const_expr(cfg.split_k): - if wstart == 0: - gDs0 = mDs0_out[None, None, head_idx, batch_idx] - for sub in cutlass.range_constexpr(num_state_subs): - ds0_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dh_col + sub * ldtm_width, cutlass.Float32), num=32 - ) - for kk in cutlass.range_constexpr(32): - gDs0[sub * ldtm_width + kk, cg1_tidx] = ds0_vec[kk] - else: - gDs0 = mDs0_out[None, None, head_idx, batch_idx] + cg1_state_dot_dstate_index = advance(cg1_state_dot_dstate_index, 1) + + # ---- dstate drain: with an initial state this is d_initial_state ---------------------- + if num_item_chunks > 0: + dstate_idx = dstate_acc_index.idx + bars.mb_dstate_acc_ready[dstate_idx].wait(dstate_acc_index.phase) + dstate_acc_index = advance(dstate_acc_index, cfg.tmem_dstate_acc_stages) + if cutlass.const_expr(cfg.use_dstate0): + if wstart == 0: + gDstate0 = mDstate0_out[None, None, head_idx, batch_idx] for sub in cutlass.range_constexpr(num_state_subs): - ds0_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_dh_col + sub * ldtm_width, cutlass.Float32), num=32) + dstate0_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 kk in cutlass.range_constexpr(32): - gDs0[sub * ldtm_width + kk, cg1_tidx] = ds0_vec[kk] - if cutlass.const_expr(not cfg.use_dht): - bars.mb_dh_acc_done[dh_idx].arrive() + gDstate0[sub * ldtm_width + kk, cg1_tidx] = dstate0_vec[kk] + if cutlass.const_expr(not cfg.use_dstate_in): + bars.mb_dstate_scale_acc_done[dstate_idx].arrive() else: - # zero-length sequence: the gradient passes straight through - # (dS0 = dHt when given, zeros otherwise); pure GMEM, no TMEM - if cutlass.const_expr(cfg.use_initial_state): - write_passthrough = True - if cutlass.const_expr(cfg.split_k): - write_passthrough = wstart == 0 + if cutlass.const_expr(cfg.use_dstate0): + write_passthrough = wstart == 0 if write_passthrough: - gDs0 = mDs0_out[None, None, head_idx, batch_idx] - if cutlass.const_expr(cfg.use_dht): - gDht = mDht[None, None, head_idx, batch_idx] + gDstate0 = mDstate0_out[None, None, head_idx, batch_idx] + if cutlass.const_expr(cfg.use_dstate_in): + 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): - gDs0[sub * ldtm_width + kk, cg1_tidx] = gDht[sub * ldtm_width + kk, cg1_tidx] + gDstate0[sub * ldtm_width + kk, cg1_tidx] = gDstate_in[sub * ldtm_width + kk, cg1_tidx] else: for sub in cutlass.range_constexpr(num_state_subs): for kk in cutlass.range_constexpr(32): - gDs0[sub * ldtm_width + kk, cg1_tidx] = cutlass.Float32(0.0) + gDstate0[sub * ldtm_width + kk, cg1_tidx] = cutlass.Float32(0.0) - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) - # CG1 done with TMEM: release the MMA warp's dealloc bars.mb_tmem_done[0].arrive() - for _ in range(cfg.tmem_dh_inp_stages): - bars.mb_dh_inp_done[dh_inp_index.idx].wait(dh_inp_index.phase) - dh_inp_index = advance(dh_inp_index, cfg.tmem_dh_inp_stages) + for _ in range(cfg.tmem_dstate_inp_stages): + bars.mb_dstate_inp_done[dstate_inp_index.idx].wait(dstate_inp_index.phase) + dstate_inp_index = advance(dstate_inp_index, cfg.tmem_dstate_inp_stages) for _ in range(cfg.smem_dk_stages): bars.mb_dk_tmastg_done[dk_index.idx].wait(dk_index.phase) dk_index = advance(dk_index, cfg.smem_dk_stages) @@ -3939,8 +3134,94 @@ def _compute1_warp( dv_index = advance(dv_index, cfg.smem_dv_stages) +@cute.kernel +def build_all_descs_kernel( + base_q: cutlass.GridConstant[tma.TensorMap], + base_k: cutlass.GridConstant[tma.TensorMap], + base_v: cutlass.GridConstant[tma.TensorMap], + base_do: cutlass.GridConstant[tma.TensorMap], + base_checkpoint: cutlass.GridConstant[tma.TensorMap], + 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, + k: cute.Tensor, + v: cute.Tensor, + do_: cute.Tensor, + state_checkpoints: cute.Tensor, + 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, + v_rs: cutlass.Int32, + do_rs: cutlass.Int32, + checkpoint_rs: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, + dq_rs: cutlass.Int32, + dk_rs: cutlass.Int32, + dv_rs: cutlass.Int32, +) -> None: + """Single-launch builder for the per-BATCH descriptor arrays (one warp per array).""" + tidx, _, _ = cute.arch.thread_idx() + widx = cutlass.Int32(tidx) // cutlass.Int32(32) + arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) + sub0 = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + 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(): + emit_seq_descs(base_q, sub0, cu_seqlens, q, n_batch, q_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 1: + if nvvm.elect_sync(): + emit_seq_descs(base_k, sub1, cu_seqlens, k, n_batch, k_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 2: + if nvvm.elect_sync(): + emit_seq_descs(base_v, sub2, cu_seqlens, v, n_batch, v_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 3: + if nvvm.elect_sync(): + emit_seq_descs(base_do, sub3, cu_seqlens, do_, n_batch, do_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 4: + if nvvm.elect_sync(): + emit_checkpoint_seq_descs(base_checkpoint, sub4, 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 widx == 5: + if nvvm.elect_sync(): + emit_seq_descs(base_dq, sub5, cu_seqlens, dq, n_batch, dq_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 6: + if nvvm.elect_sync(): + emit_seq_descs(base_dk, sub6, cu_seqlens, dk, n_batch, dk_rs, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 7: + 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.jit -def _build_descs( +def build_descs( io_dtype: cutlass.Constexpr, b_t: cutlass.Constexpr[int], q: cute.Tensor, @@ -3950,40 +3231,25 @@ def _build_descs( dq: cute.Tensor, dk: cute.Tensor, dv: cute.Tensor, - h: cute.Tensor, + state_checkpoints: cute.Tensor, cu_seqlens: cute.Tensor, - s0: Optional[cute.Tensor], + state0: Optional[cute.Tensor], tensormap_workspace: cute.Tensor, stream: cuda.CUstream, ): - """Build the per-(b,h) TMA-descriptor arrays (Q, K, V, dO, H loads; - dQ, dK, dV stores; the io-dtype S0 loads when ``s0`` is given) into - ``tensormap_workspace``. - - Launched on every execute: the descriptors fold cu_seqlens contents into - GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. - - The H descriptor is 3-D ``(dv, dk, h)`` over the packed - ``[total_h, HO, DK, DV]`` H tensor; ``build_h_descs_kernel`` derives the - per-sequence H offset from the token ``cu_seqlens`` ((seqlen-1)//B_T, - prefix-summed), folds it and the head into GLOBAL_ADDRESS, and caps - GLOBAL_DIM[2] to the per-sequence H count, so the load coordinate is the - sequence-local H index. The S0 descriptors share the H descriptor - format (one ``[DK, DV]`` entry per slot over the dense - ``[N, HO, DK, DV]`` buffer), so the load path is interchangeable.""" + """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``.""" h_q = q.shape[1] h_k = k.shape[1] h_v = v.shape[1] batch_size = cu_seqlens.shape[0] - 1 heads_out = h_q if h_q >= h_v else h_v - q_group = heads_out // h_q - k_group = heads_out // h_k - v_group = heads_out // h_v d_v = v.shape[2] - d_k_h = h.shape[2] - d_v_h = h.shape[3] + d_k_state = state_checkpoints.shape[2] + d_v_state = state_checkpoints.shape[3] bpe = io_dtype.width // 8 - granu = 128 // bpe + granule_elems = 128 // bpe bt = b_t q_row_stride, q_head_stride = q.stride[0], q.stride[1] @@ -3994,176 +3260,96 @@ def _build_descs( dk_row_stride, dk_head_stride = dk.stride[0], dk.stride[1] dv_row_stride, dv_head_stride = dv.stride[0], dv.stride[1] - q_head0 = q[None, 0, None] - k_head0 = k[None, 0, None] - - def _trans_head0(t, d): - view = t[None, 0, None] - return cute.make_tensor( - view.iterator, - cute.make_layout((d, view.shape[0]), stride=(view.stride[1], view.stride[0])), - ) - - v_head0 = _trans_head0(v, d_v) - do_head0 = _trans_head0(do_, d_v) - dq_head0 = _trans_head0(dq, dq.shape[2]) - dk_head0 = _trans_head0(dk, dk.shape[2]) - dv_head0 = _trans_head0(dv, d_v) - swz128 = _tma.TensorMapSwizzle.s128b - base_desc_q = _tma.create_tensor_map_tiled_from_view(q_head0, box_dims=(bt, granu), stride_order=(1, 0), swizzle=swz128) - base_desc_k = _tma.create_tensor_map_tiled_from_view(k_head0, box_dims=(bt, granu), stride_order=(1, 0), swizzle=swz128) - base_desc_v = _tma.create_tensor_map_tiled_from_view(v_head0, box_dims=(granu, bt), stride_order=(0, 1), swizzle=swz128) - base_desc_do = _tma.create_tensor_map_tiled_from_view(do_head0, box_dims=(granu, bt), stride_order=(0, 1), swizzle=swz128) - base_desc_dq = _tma.create_tensor_map_tiled_from_view(dq_head0, box_dims=(granu, bt), stride_order=(0, 1), swizzle=swz128) - base_desc_dk = _tma.create_tensor_map_tiled_from_view(dk_head0, box_dims=(granu, bt), stride_order=(0, 1), swizzle=swz128) - base_desc_dv = _tma.create_tensor_map_tiled_from_view(dv_head0, box_dims=(granu, bt), stride_order=(0, 1), swizzle=swz128) - h_view = cute.make_tensor( - h.iterator, + seqlen = q.shape[0] + d_k = q.shape[2] + q_headed = cute.make_tensor(q.iterator, cute.make_layout((seqlen, h_q, d_k), stride=(q_row_stride, q_head_stride, 1))) + k_headed = cute.make_tensor(k.iterator, cute.make_layout((seqlen, h_k, d_k), stride=(k_row_stride, k_head_stride, 1))) + + v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v.stride[1], v.stride[0]))) + do_headed = cute.make_tensor(do_.iterator, cute.make_layout((d_v, heads_out, seqlen), stride=(1, do_.stride[1], do_.stride[0]))) + dq_headed = cute.make_tensor(dq.iterator, cute.make_layout((dq.shape[2], heads_out, seqlen), stride=(1, dq.stride[1], dq.stride[0]))) + dk_headed = cute.make_tensor(dk.iterator, cute.make_layout((dk.shape[2], heads_out, seqlen), stride=(1, dk.stride[1], dk.stride[0]))) + dv_headed = cute.make_tensor(dv.iterator, cute.make_layout((d_v, heads_out, seqlen), stride=(1, dv.stride[1], dv.stride[0]))) + swz128 = tma.TensorMapSwizzle.s128b + base_desc_q = tma.create_tensor_map_tiled_from_view(q_headed, box_dims=(bt, 1, granule_elems), stride_order=(2, 1, 0), swizzle=swz128) + base_desc_k = tma.create_tensor_map_tiled_from_view(k_headed, box_dims=(bt, 1, granule_elems), stride_order=(2, 1, 0), swizzle=swz128) + base_desc_v = tma.create_tensor_map_tiled_from_view(v_headed, box_dims=(granule_elems, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_do = tma.create_tensor_map_tiled_from_view(do_headed, box_dims=(granule_elems, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_dq = tma.create_tensor_map_tiled_from_view(dq_headed, box_dims=(granule_elems, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_dk = tma.create_tensor_map_tiled_from_view(dk_headed, box_dims=(granule_elems, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + base_desc_dv = tma.create_tensor_map_tiled_from_view(dv_headed, box_dims=(granule_elems, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + checkpoint_view = cute.make_tensor( + state_checkpoints.iterator, cute.make_layout( - (d_v_h, d_k_h, h.shape[0]), - stride=(h.stride[3], h.stride[2], h.stride[0]), + (d_v_state, d_k_state, state_checkpoints.shape[0], heads_out), + stride=(state_checkpoints.stride[3], state_checkpoints.stride[2], state_checkpoints.stride[0], state_checkpoints.stride[1]), ), ) - base_desc_h = _tma.create_tensor_map_tiled_from_view(h_view, box_dims=(64, d_k_h, 1), stride_order=(0, 1, 2), swizzle=swz128) + 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) - arr_words = (batch_size * heads_out) * TENSOR_MAP_QWORDS - ws_iter = tensormap_workspace.iterator - - def sub_array(i): - return cute.make_tensor(ws_iter + i * arr_words, cute.make_layout((arr_words,), stride=(1,))) + 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) - build_qkv_load_descs_kernel( + n_warps = 9 if state0 is not None else 8 + build_all_descs_kernel( base_desc_q, - sub_array(0), - cu_seqlens, - q, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(q_group), - cutlass.Int32(q_head_stride), - cutlass.Int32(q_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( base_desc_k, - sub_array(1), - cu_seqlens, - k, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(k_group), - cutlass.Int32(k_head_stride), - cutlass.Int32(k_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( base_desc_v, - sub_array(2), - cu_seqlens, - v, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(v_group), - cutlass.Int32(v_head_stride), - cutlass.Int32(v_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( base_desc_do, - sub_array(3), + base_desc_checkpoint, + base_desc_dq, + base_desc_dk, + base_desc_dv, + base_desc_state0, + tensormap_workspace, cu_seqlens, + q, + k, + v, do_, + state_checkpoints, + dq, + dk, + dv, + state0, cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(1), - cutlass.Int32(do_head_stride), + cutlass.Int32(q_row_stride), + cutlass.Int32(k_row_stride), + cutlass.Int32(v_row_stride), cutlass.Int32(do_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_h_descs_kernel( - base_desc_h, - sub_array(4), - cu_seqlens, - h, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(h.stride[1]), - cutlass.Int32(h.stride[0]), + cutlass.Int32(state_checkpoints.stride[0]), cutlass.Int32(b_t), - 2, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_desc_dq, - sub_array(5), - cu_seqlens, - dq, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(1), - cutlass.Int32(dq_head_stride), cutlass.Int32(dq_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_desc_dk, - sub_array(6), - cu_seqlens, - dk, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(1), - cutlass.Int32(dk_head_stride), cutlass.Int32(dk_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_desc_dv, - sub_array(7), - cu_seqlens, - dv, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(1), - cutlass.Int32(dv_head_stride), cutlass.Int32(dv_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - if cutlass.const_expr(s0 is not None): - s0_view = cute.make_tensor( - s0.iterator, - cute.make_layout( - (d_v_h, d_k_h, 1), - stride=(s0.stride[3], s0.stride[2], s0.stride[1]), - ), - ) - base_desc_s0 = _tma.create_tensor_map_tiled_from_view(s0_view, box_dims=(64, d_k_h, 1), stride_order=(0, 1, 2), swizzle=swz128) - build_state_descs_kernel( - base_desc_s0, - sub_array(8), - s0, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(s0.stride[1]), - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) + ).launch(grid=(1, 1, 1), block=(32 * n_warps, 1, 1), stream=stream) @cute.jit -def _host( +def host( cfg: cutlass.Constexpr, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, gate: cute.Tensor, beta: cute.Tensor, - dg: cute.Tensor, + dgate: cute.Tensor, dbeta: cute.Tensor, do_: cute.Tensor, dq: cute.Tensor, dk: cute.Tensor, dv: cute.Tensor, cu_seqlens: cute.Tensor, - ds0: Optional[cute.Tensor], - dht: Optional[cute.Tensor], + dstate0: Optional[cute.Tensor], + dstate_in: Optional[cute.Tensor], work_items: Optional[cute.Tensor], work_count: Optional[cute.Tensor], sched_ctr: Optional[cute.Tensor], @@ -4171,125 +3357,21 @@ def _host( tensormap_workspace: cute.Tensor, stream: cuda.CUstream, ): - h_q = q.shape[1] - h_v = v.shape[1] + h_q = cfg.h_q + h_k = cfg.h_k + h_v = cfg.h_v batch_size = cu_seqlens.shape[0] - 1 heads_out = h_q if h_q >= h_v else h_v - # ---- GQA reshapes: fold the head group into a ---------------------- - if cutlass.const_expr(cfg.is_GQA): - h_r = h_q // h_v - h_qv = h_v - q = cute.make_tensor( - q.iterator, - cute.make_layout( - (q.shape[0], q.shape[2], (h_r, h_v)), - stride=(q.stride[0], q.stride[2], (q.stride[1], h_r * q.stride[1])), - ), - ) - k = cute.make_tensor( - k.iterator, - cute.make_layout( - (k.shape[0], k.shape[2], (h_r, h_v)), - stride=(k.stride[0], k.stride[2], (0, k.stride[1])), - ), - ) - v = cute.make_tensor( - v.iterator, - cute.make_layout( - (v.shape[2], v.shape[0], (h_r, h_v)), - stride=(v.stride[2], v.stride[0], (0, v.stride[1])), - ), - ) - else: - h_r = h_v // h_q - h_qv = h_q - q = cute.make_tensor( - q.iterator, - cute.make_layout( - (q.shape[0], q.shape[2], (h_r, h_q)), - stride=(q.stride[0], q.stride[2], (0, q.stride[1])), - ), - ) - k = cute.make_tensor( - k.iterator, - cute.make_layout( - (k.shape[0], k.shape[2], (h_r, h_q)), - stride=(k.stride[0], k.stride[2], (0, k.stride[1])), - ), - ) - v = cute.make_tensor( - v.iterator, - cute.make_layout( - (v.shape[2], v.shape[0], (h_r, h_q)), - stride=(v.stride[2], v.stride[0], (v.stride[1], h_r * v.stride[1])), - ), - ) - - gate = cute.make_tensor( - gate.iterator, - cute.make_layout( - (gate.shape[0], (h_r, h_qv)), - stride=(gate.stride[0], (gate.stride[1], h_r * gate.stride[1])), - ), - ) - beta = cute.make_tensor( - beta.iterator, - cute.make_layout( - (beta.shape[0], (h_r, h_qv)), - stride=(beta.stride[0], (beta.stride[1], h_r * beta.stride[1])), - ), - ) - dg = cute.make_tensor( - dg.iterator, - cute.make_layout( - (dg.shape[0], (h_r, h_qv)), - stride=(dg.stride[0], (dg.stride[1], h_r * dg.stride[1])), - ), - ) - dbeta = cute.make_tensor( - dbeta.iterator, - cute.make_layout( - (dbeta.shape[0], (h_r, h_qv)), - stride=(dbeta.stride[0], (dbeta.stride[1], h_r * dbeta.stride[1])), - ), - ) - if cutlass.const_expr(ds0 is not None): - ds0 = cute.make_tensor( - ds0.iterator, - cute.make_layout( - (ds0.shape[2], ds0.shape[3], (h_r, h_qv), ds0.shape[0]), - stride=( - ds0.stride[2], - ds0.stride[3], - (ds0.stride[1], h_r * ds0.stride[1]), - ds0.stride[0], - ), - ), - ) - if cutlass.const_expr(dht is not None): - dht = cute.make_tensor( - dht.iterator, - cute.make_layout( - (dht.shape[2], dht.shape[3], (h_r, h_qv), dht.shape[0]), - stride=( - dht.stride[2], - dht.stride[3], - (dht.stride[1], h_r * dht.stride[1]), - dht.stride[0], - ), - ), - ) - - # ---- SMEM sizing: per-buffer element cosizes ----------------------- + # ---- SMEM sizing: per-buffer element cosizes --------------------------------- bpe = cfg.io_dtype.width // 8 q_tile_elems = cfg.b_t * cfg.d_k k_tile_elems = cfg.b_t * cfg.d_k v_tile_elems = cfg.d_v * cfg.b_t do_tile_elems = cfg.d_v * cfg.b_t - s_tile_elems = cfg.d_k * cfg.d_v - ainv_tile_elems = cfg.b_t * cfg.b_t - qk_tile_elems = cfg.b_t * cfg.b_t + state_tile_elems = cfg.d_k * cfg.d_v + tinv_tile_elems = cfg.b_t * cfg.b_t + a_tile_elems = cfg.b_t * cfg.b_t dq_tile_elems = cfg.b_t * cfg.d_k dk_tile_elems = cfg.b_t * cfg.d_k dv_tile_elems = cfg.d_v * cfg.b_t @@ -4297,9 +3379,9 @@ def _host( cfg.k_cosize = k_tile_elems * cfg.smem_k_stages cfg.v_cosize = v_tile_elems * cfg.smem_v_stages cfg.do_cosize = do_tile_elems * cfg.smem_do_stages - cfg.s_cosize = s_tile_elems * cfg.smem_s_stages - cfg.ainv_cosize = ainv_tile_elems * cfg.smem_ainv_stages - cfg.qk_cosize = qk_tile_elems * cfg.smem_qk_stages + cfg.state_cosize = state_tile_elems * cfg.smem_state_stages + cfg.t_inv_cosize = tinv_tile_elems * cfg.smem_t_inv_stages + cfg.a_cosize = a_tile_elems * cfg.smem_a_stages cfg.dq_cosize = dq_tile_elems * cfg.smem_dq_stages cfg.dk_cosize = dk_tile_elems * cfg.smem_dk_stages cfg.dv_cosize = dv_tile_elems * cfg.smem_dv_stages @@ -4311,27 +3393,27 @@ def _host( cfg.tma_k_bytes = k_tile_elems * bpe cfg.tma_v_bytes = v_tile_elems * bpe cfg.tma_do_bytes = do_tile_elems * bpe - cfg.tma_s_bytes = s_tile_elems * bpe + cfg.tma_state_bytes = state_tile_elems * bpe cfg.n_heads_out = heads_out - num_descs = batch_size * heads_out + cfg.q_ratio = heads_out // h_q + cfg.k_ratio = heads_out // h_k + cfg.v_ratio = heads_out // h_v + num_descs = batch_size - # ---- launch -------------------------------------------------------- - total_tiles = batch_size * heads_out - # CUDA-graph-stable launch: fixed SM-count grid + # ---- launch ------------------------------------------------------------------ grid_shape = (cfg.max_active_clusters, 1, 1) - _kernel( + kernel( cfg, gate, beta, - dg, + dgate, dbeta, cu_seqlens, scale, cumsumlog_smem_layout_staged, beta_smem_layout_staged, - total_tiles, q, k, v, @@ -4339,8 +3421,8 @@ def _host( dq, dk, dv, - ds0, - dht, + dstate0, + dstate_in, work_items, work_count, sched_ctr, @@ -4356,17 +3438,16 @@ def _host( @cute.kernel -def _kernel( +def kernel( cfg: cutlass.Constexpr, mGate: cute.Tensor, mBeta: cute.Tensor, - mDg: cute.Tensor, + mDgate: cute.Tensor, mDbeta: cute.Tensor, cu_seqlens: cute.Tensor, scale: cutlass.Float32, cumsumlog_smem_layout_staged: cute.Layout, beta_smem_layout_staged: cute.Layout, - total_tiles: cutlass.Int32, mQ, mK, mV, @@ -4374,31 +3455,127 @@ def _kernel( mdQ, mdK, mdV, - mDs0, - mDht, - mWorkItems: Optional[cute.Tensor], - mCount: Optional[cute.Tensor], + mDstate0, + mDstate_in, + mWorkItems: cute.Tensor, + mCount: cute.Tensor, mSched: Optional[cute.Tensor], tensormap_workspace: cute.Tensor, n_desc: cutlass.Int32, ): - """ - Main GDN bprop chunked kernel. - - Warp specialization is the outermost control flow: each warp role owns - its own persistent tile-scheduler loop, iterating over (batch, head) - tiles and then over chunks within each tile in BACKWARD order. - """ + """Main GDN bprop chunked kernel (warp-specialized persistent body).""" tidx, _, _ = cute.arch.thread_idx() warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) bidx = cute.arch.block_idx()[0] num_ctas = cute.arch.grid_dim()[0] - if cutlass.const_expr(cfg.split_k): - total_tiles = mCount[0] + total_tiles = mCount[0] if cutlass.const_expr(cfg.dyn_sched): assert mSched is not None, "mSched must be provided if dyn_sched is True" + if cutlass.const_expr(cfg.is_GQA): + h_r = cfg.h_q // cfg.h_v + h_qv = cfg.h_v + mQ = cute.make_tensor( + mQ.iterator, + cute.make_layout( + (mQ.shape[0], mQ.shape[2], (h_r, h_qv)), + stride=(mQ.stride[0], mQ.stride[2], (mQ.stride[1], h_r * mQ.stride[1])), + ), + ) + mK = cute.make_tensor( + mK.iterator, + cute.make_layout( + (mK.shape[0], mK.shape[2], (h_r, h_qv)), + stride=(mK.stride[0], mK.stride[2], (0, mK.stride[1])), + ), + ) + mV = cute.make_tensor( + mV.iterator, + cute.make_layout( + (mV.shape[2], mV.shape[0], (h_r, h_qv)), + stride=(mV.stride[2], mV.stride[0], (0, mV.stride[1])), + ), + ) + else: + h_r = cfg.h_v // cfg.h_q + h_qv = cfg.h_q + mQ = cute.make_tensor( + mQ.iterator, + cute.make_layout( + (mQ.shape[0], mQ.shape[2], (h_r, h_qv)), + stride=(mQ.stride[0], mQ.stride[2], (0, mQ.stride[1])), + ), + ) + mK = cute.make_tensor( + mK.iterator, + cute.make_layout( + (mK.shape[0], mK.shape[2], (h_r, h_qv)), + stride=(mK.stride[0], mK.stride[2], (0, mK.stride[1])), + ), + ) + mV = cute.make_tensor( + mV.iterator, + cute.make_layout( + (mV.shape[2], mV.shape[0], (h_r, h_qv)), + stride=(mV.stride[2], mV.stride[0], (mV.stride[1], h_r * mV.stride[1])), + ), + ) + mGate = cute.make_tensor( + mGate.iterator, + cute.make_layout( + (mGate.shape[0], (h_r, h_qv)), + stride=(mGate.stride[0], (mGate.stride[1], h_r * mGate.stride[1])), + ), + ) + mBeta = cute.make_tensor( + mBeta.iterator, + cute.make_layout( + (mBeta.shape[0], (h_r, h_qv)), + stride=(mBeta.stride[0], (mBeta.stride[1], h_r * mBeta.stride[1])), + ), + ) + mDgate = cute.make_tensor( + mDgate.iterator, + cute.make_layout( + (mDgate.shape[0], (h_r, h_qv)), + stride=(mDgate.stride[0], (mDgate.stride[1], h_r * mDgate.stride[1])), + ), + ) + mDbeta = cute.make_tensor( + mDbeta.iterator, + cute.make_layout( + (mDbeta.shape[0], (h_r, h_qv)), + stride=(mDbeta.stride[0], (mDbeta.stride[1], h_r * mDbeta.stride[1])), + ), + ) + if cutlass.const_expr(mDstate0 is not None): + mDstate0 = cute.make_tensor( + mDstate0.iterator, + cute.make_layout( + (mDstate0.shape[2], mDstate0.shape[3], (h_r, h_qv), mDstate0.shape[0]), + stride=( + mDstate0.stride[2], + mDstate0.stride[3], + (mDstate0.stride[1], h_r * mDstate0.stride[1]), + mDstate0.stride[0], + ), + ), + ) + if cutlass.const_expr(mDstate_in is not None): + mDstate_in = cute.make_tensor( + mDstate_in.iterator, + cute.make_layout( + (mDstate_in.shape[2], mDstate_in.shape[3], (h_r, h_qv), mDstate_in.shape[0]), + stride=( + mDstate_in.stride[2], + mDstate_in.stride[3], + (mDstate_in.stride[1], h_r * mDstate_in.stride[1]), + mDstate_in.stride[0], + ), + ), + ) + desc_base_words = tensormap_workspace.iterator.raw_ptr() desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) arr_words = n_desc * desc_qwords @@ -4406,16 +3583,16 @@ def _kernel( desc_k_base = desc_base_words + arr_words desc_v_base = desc_base_words + cutlass.Int32(2) * arr_words desc_do_base = desc_base_words + cutlass.Int32(3) * arr_words - desc_s_base = desc_base_words + cutlass.Int32(4) * arr_words + desc_checkpoint_base = desc_base_words + cutlass.Int32(4) * arr_words 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_s0_base = desc_base_words + cutlass.Int32(8) * arr_words + desc_initial_state_base = desc_base_words + cutlass.Int32(8) * arr_words SMEM = cutlass.AddressSpace.smem bars = make_gdn_bars(cfg) sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=cutlass.AddressSpace.smem, alignment=16) - tmem_hold = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=16) + tmem_base_slot = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=16) cumsumlog_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) cumprod_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) beta_raw = cutlass.Array(cutlass.Float32, cute.cosize(beta_smem_layout_staged), space=SMEM, alignment=128) @@ -4426,7 +3603,7 @@ def _kernel( STRIDE = 8 * 128 KT_LEAD = (cfg.d_v // 2) * 128 V_LEAD = (cfg.d_v // 2) * 128 - S_LEAD = cfg.d_k * 128 + STATE_LEAD = cfg.d_k * 128 sQ_raw = cutlass.Array( cfg.io_dtype, cfg.q_cosize, @@ -4493,97 +3670,97 @@ def _kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sS_raw = cutlass.Array( + sState_raw = cutlass.Array( cfg.io_dtype, - cfg.s_cosize, + cfg.state_cosize, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sS = SmemTile( - base=sS_raw.data_ptr().toint(), - elems_per_stage=(cfg.s_cosize // cfg.smem_s_stages) * bpe, - stages=cfg.smem_s_stages, - leading_byte_offset=S_LEAD, + sState = 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=STATE_LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sS_kmaj = SmemTile( - base=sS_raw.data_ptr().toint(), - elems_per_stage=(cfg.s_cosize // cfg.smem_s_stages) * bpe, - stages=cfg.smem_s_stages, + sState_kmaj = 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, ) - sAinv_raw = cutlass.Array( + sTinv_raw = cutlass.Array( cfg.io_dtype, - cfg.ainv_cosize, + cfg.t_inv_cosize, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sAinv = SmemTile( - base=sAinv_raw.data_ptr().toint(), - elems_per_stage=(cfg.ainv_cosize // cfg.smem_ainv_stages) * bpe, - stages=cfg.smem_ainv_stages, + sTinv = SmemTile( + base=sTinv_raw.data_ptr().toint(), + elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages) * bpe, + stages=cfg.smem_t_inv_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sAinv_trans = SmemTile( - base=sAinv_raw.data_ptr().toint(), - elems_per_stage=(cfg.ainv_cosize // cfg.smem_ainv_stages) * bpe, - stages=cfg.smem_ainv_stages, + sTinv_trans = SmemTile( + base=sTinv_raw.data_ptr().toint(), + elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages) * bpe, + stages=cfg.smem_t_inv_stages, leading_byte_offset=(cfg.b_t // 2) * 128, stride_byte_offset=STRIDE, layout=SWZ, ) sKK_raw = cutlass.Array( cfg.io_dtype, - cfg.ainv_cosize, + cfg.t_inv_cosize, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) sKK = SmemTile( base=sKK_raw.data_ptr().toint(), - elems_per_stage=(cfg.ainv_cosize // cfg.smem_ainv_stages) * bpe, - stages=cfg.smem_ainv_stages, + elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages) * bpe, + stages=cfg.smem_t_inv_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sQk_raw = cutlass.Array( + sA_raw = cutlass.Array( cfg.io_dtype, - cfg.qk_cosize, + cfg.a_cosize, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sQk = SmemTile( - base=sQk_raw.data_ptr().toint(), - elems_per_stage=(cfg.qk_cosize // cfg.smem_qk_stages) * bpe, - stages=cfg.smem_qk_stages, + sA = SmemTile( + base=sA_raw.data_ptr().toint(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, + stages=cfg.smem_a_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sQk_trans = SmemTile( - base=sQk_raw.data_ptr().toint(), - elems_per_stage=(cfg.qk_cosize // cfg.smem_qk_stages) * bpe, - stages=cfg.smem_qk_stages, + sA_trans = SmemTile( + base=sA_raw.data_ptr().toint(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, + stages=cfg.smem_a_stages, leading_byte_offset=(cfg.b_t // 2) * 128, stride_byte_offset=STRIDE, layout=SWZ, ) sDa = SmemTile( - base=sQk_raw.data_ptr().toint(), - elems_per_stage=(cfg.qk_cosize // cfg.smem_qk_stages) * bpe, + base=sA_raw.data_ptr().toint(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, stages=1, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) sDa_trans = SmemTile( - base=sQk_raw.data_ptr().toint(), - elems_per_stage=(cfg.qk_cosize // cfg.smem_qk_stages) * bpe, + base=sA_raw.data_ptr().toint(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, stages=1, leading_byte_offset=(cfg.b_t // 2) * 128, stride_byte_offset=STRIDE, @@ -4591,13 +3768,13 @@ def _kernel( ) sDm_raw = cutlass.Array( cfg.io_dtype, - cfg.qk_cosize // cfg.smem_qk_stages, + cfg.a_cosize // cfg.smem_a_stages, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) sDm = SmemTile( base=sDm_raw.data_ptr().toint(), - elems_per_stage=(cfg.qk_cosize // cfg.smem_qk_stages) * bpe, + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, stages=1, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, @@ -4605,14 +3782,13 @@ def _kernel( ) sDm_trans = SmemTile( base=sDm_raw.data_ptr().toint(), - elems_per_stage=(cfg.qk_cosize // cfg.smem_qk_stages) * bpe, + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, stages=1, leading_byte_offset=(cfg.b_t // 2) * 128, stride_byte_offset=STRIDE, layout=SWZ, ) - # sub-bank split: V + sdH + dQ + dK + dV (the LDSM/LDS/STS/STSM-heavy - # set) allocated last -> all past the 128KB sub-bank boundary + # sub-bank split: V + sDstate + dQ + dK + dV allocated last sV_raw = cutlass.Array( cfg.io_dtype, cfg.v_cosize, @@ -4635,14 +3811,14 @@ def _kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sdH_raw = cutlass.Array( + sDstate_raw = cutlass.Array( cfg.io_dtype, cfg.d_k * cfg.d_v, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sdH = SmemTile( - base=sdH_raw.data_ptr().toint(), + sDstate = SmemTile( + base=sDstate_raw.data_ptr().toint(), elems_per_stage=cfg.d_k * cfg.d_v * bpe, stages=1, leading_byte_offset=LEAD, @@ -4691,21 +3867,14 @@ def _kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sdV_kmaj = SmemTile( - base=sdV_raw.data_ptr().toint(), - elems_per_stage=(cfg.dv_cosize // cfg.smem_dv_stages) * bpe, - stages=cfg.smem_dv_stages, - leading_byte_offset=LEAD, - stride_byte_offset=STRIDE, - layout=SWZ, - ) - sdh_flat = cute.make_tensor( - cute.make_ptr(cfg.io_dtype, sdH_raw.data_ptr().toint(), mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes), + sdV_kmaj = sdV + sdstate_flat = cute.make_tensor( + cute.make_ptr(cfg.io_dtype, sDstate_raw.data_ptr().toint(), mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes), cute.make_layout(cfg.d_k * cfg.d_v), ) - ss_flat = cute.make_tensor( - cute.make_ptr(cfg.io_dtype, sS_raw.data_ptr().toint(), mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes), - cute.make_layout(cfg.s_cosize), + sstate_flat = cute.make_tensor( + cute.make_ptr(cfg.io_dtype, sState_raw.data_ptr().toint(), mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes), + cute.make_layout(cfg.state_cosize), ) sCumsumlog = cute.make_tensor( cute.make_ptr(cutlass.Float32, cumsumlog_raw.data_ptr().toint(), mem_space=cute.AddressSpace.smem, assumed_align=128), @@ -4720,9 +3889,7 @@ def _kernel( beta_smem_layout_staged, ) - # ------------------------------------------------------------------ - # mbarrier init (all threads) - # ------------------------------------------------------------------ + # ---- mbarrier init (all threads) --------------------------------------------- for s_ in range(cfg.sched_stages): bars.mb_sched_ready[s_].init() bars.mb_sched_done[s_].init() @@ -4737,59 +3904,54 @@ def _kernel( for s in range(cfg.smem_v_stages): bars.mb_v_ready[s].init() bars.mb_v_mma_done[s].init() - bars.mb_v_cg1_done[s].init() for s in range(cfg.smem_do_stages): bars.mb_do_ready[s].init() bars.mb_do_mma_done[s].init() - bars.mb_do_cg1_done[s].init() - for s in range(cfg.smem_s_stages): - bars.mb_s_ready[s].init() - bars.mb_s_done[s].init() + for s in range(cfg.smem_state_stages): + bars.mb_state_ready[s].init() + bars.mb_state_mma_done[s].init() for s in range(cfg.smem_gate_stages): bars.mb_gate_ready[s].init() bars.mb_gate_done[s].init() for s in range(cfg.smem_beta_stages): bars.mb_beta_ready[s].init() bars.mb_beta_done[s].init() - for s in range(cfg.tmem_dh_acc_stages): - bars.mb_dh_acc_ready[s].init() - bars.mb_dh_acc_done[s].init() + for s in range(cfg.tmem_dstate_acc_stages): + bars.mb_dstate_acc_ready[s].init() + bars.mb_dstate_scale_acc_done[s].init() for b in ( - bars.mb_du_scale_ready, - bars.mb_du_scale_done, - bars.mb_du_total_ready, - bars.mb_dk_scale_ready, - bars.mb_dk_scale_done, - bars.mb_dk_attn_ready, - bars.mb_dk_attn_done, - bars.mb_dk_total_ready, - bars.mb_dk_total_done, + bars.mb_du_scale_acc_ready, + bars.mb_du_scale_acc_done, + bars.mb_du_total_acc_ready, + bars.mb_dk_scale_acc_ready, + bars.mb_dk_scale_acc_done, + bars.mb_dk_attn_acc_ready, + bars.mb_dk_attn_acc_done, + bars.mb_dk_total_acc_ready, + bars.mb_dk_total_acc_done, ): b[0].init() for b in ( bars.mb_kk_acc_ready, bars.mb_kk_acc_done, bars.mb_a_acc_ready, - bars.mb_ks_acc_ready, + bars.mb_k_state_acc_ready, bars.mb_u_acc_ready, bars.mb_dy_acc_ready, ): b[0].init() - for s in range(cfg.smem_ainv_stages): - bars.mb_ainv_ready[s].init() - bars.mb_ainv_done[s].init() - for s in range(cfg.smem_qk_stages): - bars.mb_qk_ready[s].init() - bars.mb_qk_done[s].init() - for s in range(cfg.tmem_dh_inp_stages): - bars.mb_dh_inp_ready[s].init() - bars.mb_dh_inp_done[s].init() + for s in range(cfg.smem_t_inv_stages): + bars.mb_t_inv_ready[s].init() + for s in range(cfg.smem_a_stages): + bars.mb_a_ready[s].init() + bars.mb_a_done[s].init() + for s in range(cfg.tmem_dstate_inp_stages): + bars.mb_dstate_inp_ready[s].init() + bars.mb_dstate_inp_done[s].init() for b in ( - bars.mb_dop_inp_ready, - bars.mb_dop_inp_done, + bars.mb_do_prime_inp_ready, bars.mb_du_inp_ready, bars.mb_dyp_inp_ready, - bars.mb_dyp_inp_done, ): b[0].init() for s in range(cfg.smem_dq_stages): @@ -4804,31 +3966,28 @@ def _kernel( bars.mb_y_ready[0].init() bars.mb_sdv_done[0].init() bars.mb_u_ready[0].init() - bars.mb_dhs_ready[0].init() - bars.mb_dhs_done[0].init() + bars.mb_dstate_smem_ready[0].init() bars.mb_da_ready[0].init() bars.mb_dq_acc_scale_ready[0].init() bars.mb_dq_acc_scale_done[0].init() bars.mb_dq_acc_total_ready[0].init() bars.mb_dq_acc_total_done[0].init() bars.mb_da_acc_ready[0].init() - bars.mb_dm_ready[0].init() - bars.mb_dm_done[0].init() + bars.mb_dm_acc_ready[0].init() + bars.mb_dm_acc_done[0].init() bars.mb_dbeta_cg1_ready[0].init() bars.mb_dgate_cg1_ready[0].init() - bars.mb_hdh_done[0].init() - bars.mb_dk_spath_ready[0].init() + bars.mb_state_dot_dstate_done[0].init() + bars.mb_dk_state_path_acc_ready[0].init() bars.mb_tmem_done[0].init() nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync() - # ------------------------------------------------------------------ - # 2. Warp specialization - each warp role owns its own scheduler loop - # ------------------------------------------------------------------ + # ---- warp specialization ----------------------------------------------------- if warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: - _compute0_warp( + compute0_warp_group( cfg, total_tiles, bidx, @@ -4836,38 +3995,38 @@ def _kernel( cu_seqlens, mWorkItems, tidx, - tmem_hold=tmem_hold, + tmem_base_slot=tmem_base_slot, scale=scale, sCumsumlog=sCumsumlog, sCumprod=sCumprod, sBeta=sBeta, - sAinv=sAinv, + sTinv=sTinv, sKK=sKK, - sQk=sQk, + sA=sA, sDa=sDa, sDm=sDm, sK=sK, sdQ=sdQ, - sdH=sdH, - ss_flat=ss_flat, - sdh_flat=sdh_flat, + sDstate=sDstate, + sstate_flat=sstate_flat, + sdstate_flat=sdstate_flat, sSched=sSched, bars=bars, ) if warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]: - _compute1_warp( + compute1_warp_group( cfg, total_tiles, bidx, num_ctas, cu_seqlens, mWorkItems, - mDs0, - mDht, + mDstate0, + mDstate_in, tidx, warp_idx=warp_idx, - tmem_hold=tmem_hold, + tmem_base_slot=tmem_base_slot, scale=scale, sQ=sQ, sK=sK, @@ -4879,7 +4038,7 @@ def _kernel( sdQ=sdQ, sdK=sdK, sdV=sdV, - sdH=sdH, + sDstate=sDstate, sDa=sDa, sDm=sDm, sSched=sSched, @@ -4887,14 +4046,14 @@ def _kernel( ) elif warp_idx == cfg.mma_warp_id: - _mma_warp( + mma_warp( cfg, total_tiles, bidx, num_ctas, cu_seqlens, mWorkItems, - tmem_hold=tmem_hold, + tmem_base_slot=tmem_base_slot, sQ=sQ, sQ_trans=sQ_trans, sK=sK, @@ -4903,25 +4062,24 @@ def _kernel( sV_kmaj=sV_kmaj, sdO=sdO, sdO_kmaj=sdO_kmaj, - sS=sS, - sS_kmaj=sS_kmaj, - sAinv=sAinv, - sAinv_trans=sAinv_trans, - sQk=sQk, - sQk_trans=sQk_trans, + sState=sState, + sState_kmaj=sState_kmaj, + sTinv=sTinv, + sTinv_trans=sTinv_trans, + sA=sA, + sA_trans=sA_trans, sDa=sDa, sDa_trans=sDa_trans, - sdH=sdH, + sDstate=sDstate, sDm=sDm, sDm_trans=sDm_trans, - sdV=sdV, sdV_kmaj=sdV_kmaj, sSched=sSched, bars=bars, ) elif warp_idx == cfg.tma_qkv_warp_id: - _tmaldg_warp( + tmaldg_warp( cfg, total_tiles, bidx, @@ -4933,19 +4091,19 @@ def _kernel( sK_raw=sK_raw, sV_raw=sV_raw, sdO_raw=sdO_raw, - sS_raw=sS_raw, + sState_raw=sState_raw, desc_q_base=desc_q_base, desc_k_base=desc_k_base, desc_v_base=desc_v_base, desc_do_base=desc_do_base, - desc_s_base=desc_s_base, - desc_s0_base=desc_s0_base, + desc_checkpoint_base=desc_checkpoint_base, + desc_initial_state_base=desc_initial_state_base, sSched=sSched, bars=bars, ) if warp_idx == cfg.load_gate_beta_warp_id: - _gate_beta_warp( + gate_beta_warp( cfg, total_tiles, bidx, @@ -4955,7 +4113,7 @@ def _kernel( tidx, mGate=mGate, mBeta=mBeta, - mDg=mDg, + mDgate=mDgate, mDbeta=mDbeta, sCumsumlog=sCumsumlog, sCumprod=sCumprod, @@ -4964,7 +4122,7 @@ def _kernel( bars=bars, ) if warp_idx == cfg.epilogue_warp_id: - _tmastg_warp( + tmastg_warp( cfg, total_tiles, bidx, @@ -4988,23 +4146,19 @@ class GdnBwdCfg: The per-compile parameters (dtypes, GQA) are the ``cute.compile`` cache keys; the rest is derived from the module-global ``CFG`` constants. - ``_host`` stamps the shape-derived fields at trace time. + ``host`` stamps the shape-derived fields at trace time. """ use_initial_state: bool - use_dht: bool + use_dstate_in: bool + use_dstate0: bool io_dtype: Type[cutlass.Numeric] acc_dtype: Type[cutlass.Numeric] max_active_clusters: int is_GQA: bool - # split-K: tiles come from a work-item table (see common/split_k.py); - # each item computes chunks [wstart, cend) backward, writes [wstart, wend) - split_k: bool = False - # gate input domain: natural-log decay instead of raw - # linear alpha; the gate warp then skips its log2 and rescales by 1/ln2 log_gate: bool = False - # --- fixed constants stamped from CFG by build_cfg --- + # ---- fixed constants stamped from CFG by build_cfg --------------------------- b_t: int = CFG.B_T d_k: int = CFG.D_K d_v: int = CFG.D_V @@ -5023,7 +4177,7 @@ class GdnBwdCfg: dyn_sched: bool = False sched_stages: int = CFG.SMEM_SCHED_STAGES - # --- named barrier slots (ids 1-6; 0 is the CTA-wide sync) --- + # ---- named barrier slots (ids 1-6; 0 is the CTA-wide sync) ------------------- tmem_alloc_barrier_id: int = 1 tmem_alloc_barrier_threads: int = 0 inverse_barrier_id: int = 2 @@ -5035,40 +4189,40 @@ class GdnBwdCfg: cg1_barrier_id: int = 5 cg1_barrier_threads: int = 0 - # --- SMEM / TMEM stage counts + TMEM column offsets --- + # ---- SMEM / TMEM stage counts + TMEM column offsets -------------------------- smem_q_stages: int = CFG.SMEM_Q_STAGES smem_k_stages: int = CFG.SMEM_K_STAGES smem_v_stages: int = CFG.SMEM_V_STAGES smem_do_stages: int = 1 - smem_s_stages: int = 1 - smem_ainv_stages: int = CFG.SMEM_AINV_STAGES - smem_qk_stages: int = CFG.SMEM_QK_STAGES + smem_state_stages: int = 1 + smem_t_inv_stages: int = CFG.SMEM_T_INV_STAGES + smem_a_stages: int = CFG.SMEM_A_STAGES smem_dq_stages: int = 1 smem_dk_stages: int = 1 smem_dv_stages: int = 1 smem_gate_stages: int = 2 smem_beta_stages: int = 2 - tmem_dh_acc_stages: int = CFG.TMEM_DH_ACC_STAGES + tmem_dstate_acc_stages: int = CFG.TMEM_DH_ACC_STAGES tmem_dvdk_acc_stages: int = CFG.TMEM_DVDK_ACC_STAGES - tmem_dh_inp_stages: int = CFG.TMEM_DH_INP_STAGES + tmem_dstate_inp_stages: int = CFG.TMEM_DH_INP_STAGES tmem_shared_inp_stages: int = CFG.TMEM_SHARED_INP_STAGES tmem_shared_acc_stages: int = CFG.TMEM_SHARED_ACC_STAGES - tmem_dh_offset: int = 0 - tmem_dvdk_offset: int = 0 - tmem_dh_inp_offset: int = 0 + tmem_dstate_acc_offset: int = 0 + tmem_dvdk_acc_offset: int = 0 + tmem_dstate_inp_offset: int = 0 tmem_shared_acc_offset: int = 0 tmem_shared_inp_offset: int = 0 tmem_y_offset: int = 0 buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES - # --- stamped by _host at trace time (shape-derived) --- + # ---- stamped by host at trace time (shape-derived) -------------------------- q_cosize: int = 0 k_cosize: int = 0 v_cosize: int = 0 do_cosize: int = 0 - s_cosize: int = 0 - ainv_cosize: int = 0 - qk_cosize: int = 0 + state_cosize: int = 0 + t_inv_cosize: int = 0 + a_cosize: int = 0 dq_cosize: int = 0 dk_cosize: int = 0 dv_cosize: int = 0 @@ -5076,8 +4230,11 @@ class GdnBwdCfg: tma_k_bytes: int = 0 tma_v_bytes: int = 0 tma_do_bytes: int = 0 - tma_s_bytes: int = 0 + tma_state_bytes: int = 0 n_heads_out: int = 0 + q_ratio: int = 1 + k_ratio: int = 1 + v_ratio: int = 1 def build_cfg( @@ -5086,8 +4243,8 @@ def build_cfg( max_active_clusters: int, is_GQA: bool, use_initial_state: bool = False, - use_dht: bool = False, - split_k: bool = False, + use_dstate_in: bool = False, + use_dstate0: bool = False, log_gate: bool = False, dyn_sched: bool = False, ) -> GdnBwdCfg: @@ -5097,12 +4254,12 @@ def build_cfg( raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") cfg = GdnBwdCfg( use_initial_state=use_initial_state, - use_dht=use_dht, + use_dstate_in=use_dstate_in, + use_dstate0=use_dstate0, io_dtype=io_dtype, acc_dtype=cutlass.Float32, max_active_clusters=max_active_clusters, is_GQA=is_GQA, - split_k=split_k, log_gate=log_gate, dyn_sched=dyn_sched, ) @@ -5114,67 +4271,30 @@ def build_cfg( cfg.inverse_inner_barrier_threads = cfg.threads_per_warp * 2 cfg.init_state_store_barrier_threads = cfg.threads_per_warp * n_cg1 cfg.cg1_barrier_threads = cfg.threads_per_warp * n_cg1 - cfg.tmem_dh_offset = 0 - cfg.tmem_dvdk_offset = cfg.tmem_dh_offset + cfg.tmem_dh_acc_stages * 128 - cfg.tmem_dh_inp_offset = cfg.tmem_dvdk_offset + cfg.tmem_dvdk_acc_stages * 64 - cfg.tmem_shared_acc_offset = cfg.tmem_dh_inp_offset + cfg.tmem_dh_inp_stages * 64 + cfg.tmem_dstate_acc_offset = 0 + cfg.tmem_dvdk_acc_offset = cfg.tmem_dstate_acc_offset + cfg.tmem_dstate_acc_stages * 128 + cfg.tmem_dstate_inp_offset = cfg.tmem_dvdk_acc_offset + cfg.tmem_dvdk_acc_stages * 64 + cfg.tmem_shared_acc_offset = cfg.tmem_dstate_inp_offset + cfg.tmem_dstate_inp_stages * 64 cfg.tmem_shared_inp_offset = cfg.tmem_shared_acc_offset + cfg.tmem_shared_acc_stages * 64 cfg.tmem_y_offset = cfg.tmem_shared_inp_offset + cfg.tmem_shared_inp_stages * (cfg.b_t // 2) return cfg -def get_workspace_size(B: int, HQ: int, HV: int): - HO = HQ if HQ >= HV else HV - return CFG.BYTES_PER_TENSORMAP * (9 * B * HO) + 128 - - -def _check_cuda(err): - if err != cuda.CUresult.CUDA_SUCCESS: - raise RuntimeError(f"CUDA driver call failed: {err}") - - -def _data_ptr(t) -> int: - """Device address of a tensor-like (``data_ptr()`` or the CUDA array - interface).""" - fn = getattr(t, "data_ptr", None) - if fn is not None: - return fn() - return t.__cuda_array_interface__["data"][0] - - -def _device_sm_count() -> int: - """Multiprocessor count of the current device (runtime API: auto-inits - the primary context, so this works before any other CUDA call).""" - from cuda.bindings import runtime as _rt - - err, dev = _rt.cudaGetDevice() - if int(err) != 0: - raise RuntimeError(f"cudaGetDevice failed: {err}") - err, count = _rt.cudaDeviceGetAttribute(_rt.cudaDeviceAttr.cudaDevAttrMultiProcessorCount, dev) - if int(err) != 0: - raise RuntimeError(f"cudaDeviceGetAttribute failed: {err}") - return count - - -def _cutlass_io_dtype(dtype): - name = str(dtype) - if "bfloat16" in name: - return cutlass.BFloat16 - if "float16" in name or "half" in name: - return cutlass.Float16 - raise ValueError(f"Unsupported dtype {dtype}, expected bfloat16 or float16") +TENSORMAP_DESC_ARRAYS = 8 # per-batch runtime TMA descriptors: Q, K, V, dO, checkpoints, dQ, dK, dV +TENSORMAP_STATIC_SLOTS = 1 # initial_state @functools.cache -def _get_compiled_cache( +def get_compiled_cache( io_dtype_str: str, + cu_dtype_str: str, HQ: int, HK: int, HV: int, is_GQA: bool, use_initial_state: bool = False, - use_dht: bool = False, - split_k: bool = False, + use_dstate_in: bool = False, + use_dstate0: bool = False, log_gate: bool = False, dyn_sched: bool = False, ): @@ -5186,26 +4306,29 @@ def compile( io_dtype, is_GQA: bool, use_initial_state: bool = False, - use_dht: bool = False, - split_k: bool = False, + use_dstate_in: bool = False, + use_dstate0: bool = False, log_gate: bool = False, dyn_sched: bool = False, *, num_sm: int, + h_q: int, + h_k: int, + h_v: int, q_cute, k_cute, v_cute, gate_cute, beta_cute, - dg_cute, + dgate_cute, dbeta_cute, do_cute, dq_cute, dk_cute, dv_cute, cu_seqlens_cute, - ds0_cute=None, - dht_cute=None, + dstate0_cute=None, + dstate_in_cute=None, work_items_cute=None, work_count_cute=None, sched_ctr_cute=None, @@ -5219,29 +4342,32 @@ def compile( max_active_clusters=num_sm, is_GQA=is_GQA, use_initial_state=use_initial_state, - use_dht=use_dht, - split_k=split_k, + use_dstate_in=use_dstate_in, + use_dstate0=use_dstate0, log_gate=log_gate, dyn_sched=dyn_sched, ) + cfg.h_q = h_q + cfg.h_k = h_k + cfg.h_v = h_v return cute.compile( - _host, + host, cfg, q_cute, k_cute, v_cute, gate_cute, beta_cute, - dg_cute, + dgate_cute, dbeta_cute, do_cute, dq_cute, dk_cute, dv_cute, cu_seqlens_cute, - ds0_cute, - dht_cute, + dstate0_cute, + dstate_in_cute, work_items_cute, work_count_cute, sched_ctr_cute, @@ -5259,11 +4385,11 @@ def chunk_gdn_bwd_sm100( gate, beta, do, - h, + state_checkpoints, dq, dk, dv, - dg, + dgate, dbeta, cu_seqlens, scale: float, @@ -5280,13 +4406,13 @@ def chunk_gdn_bwd_sm100( ) -> None: """Execute the Blackwell chunked GDN bprop kernel (THD / varlen entry). - Produces dQ/dK/dV/dG/dBeta at ``HO = max(HQ, HV)`` heads (the caller - reduces over the head group; dgate = dL/d(ln alpha)). With + 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 dL/dS0. The two go together. ``h`` is - always the PLAIN per-chunk series. All tensors are contiguous, + same shape) 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. @@ -5298,112 +4424,98 @@ def chunk_gdn_bwd_sm100( alpha, or the natural-log decay when ``log_gate`` beta: ``(total_tokens, HO)`` float32, update gate do: ``(total_tokens, HO, DV)`` float16/bfloat16, output gradient - h: ``(total_h, HO, DK, DV)`` bfloat16, per-chunk forward states from - the prefill kernel's H output (``checkpoint_every_n_tokens=B_T``) + state_checkpoints: ``(total_checkpoints, HO, DK, DV)`` 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 - dg/dbeta: pre-allocated ``(total_tokens, HO)`` float32 gate/beta + dgate/dbeta: pre-allocated ``(total_tokens, HO)`` float32 gate/beta gradients cu_seqlens: ``(num_seqs + 1,)`` int32 - initial_state: ``(num_seqs, HO, DK, DV)`` io dtype (matching ``h``), + initial_state: ``(num_seqs, HO, DK, DV)`` io dtype (matching ``state_checkpoints``), or None scale: attention scale factor (must not be 0) - work_items: ``(max_items, 6)`` int32 split-K work-item table from - ``common/split_k.py``, or None for the one-tile-per-(b,h) - layout. With a table, each item computes chunks - ``[wstart, cend)`` backward and writes gradients only for - ``[wstart, wend)``. - work_count: ``(1,)`` int32 device-side item count (required with - work_items) - workspace: ``(>= get_workspace_size(B, HQ, HV) // 8,)`` int64, + work_items: ``(max_items, 8)`` int32 work-item table from + ``common/split_k.py`` (REQUIRED; an uncut table row is the whole + (b, h) sequence). Each item computes chunks ``[wstart, cend)`` + backward and writes gradients only for ``[wstart, wend)``. + work_count: ``(1,)`` int32 device-side item count (REQUIRED) + workspace: ``(>= tensormap_workspace_bytes(module, B) // 8,)`` int64, 128-byte aligned; holds the per-(b,h) TMA descriptors stream: CUDA stream handle (``cudaStream_t`` as an int) """ HQ = q.shape[1] + HK = k.shape[1] HV = v.shape[1] + HO = max(HQ, HV) DK = q.shape[2] B = cu_seqlens.shape[0] - 1 is_GQA = HQ >= HV - split_k = work_items is not None - io_dtype = _cutlass_io_dtype(q.dtype) - if str(h.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: - raise ValueError(f"h dtype must match the io dtype (the prefill H output): got {h.dtype} with io {q.dtype}") - if (initial_state is None) != (d_initial_state is None): - raise ValueError("initial_state and d_initial_state go together (chunk 0 reads S0; the drain produces dS0)") - if initial_state is not None and str(initial_state.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: - raise ValueError(f"initial_state must be io dtype (the caller downcasts): got {initial_state.dtype} with io {q.dtype}") - if split_k: - if work_count is None: - raise ValueError("work_count is required with work_items") - elif work_count is not None: - raise ValueError("work_count must be None without work_items") - - ws_words = get_workspace_size(B, HQ, HV) // 8 - if workspace.shape[0] < ws_words: - raise ValueError(f"workspace too small: need {ws_words} int64 words, " f"got {workspace.shape[0]}") - if _data_ptr(workspace) % 128 != 0: - raise ValueError("workspace must be 128-byte aligned") + 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)") + io_dtype = get_dtype(q.dtype) + for name, hh in (("HQ", HQ), ("HK", HK), ("HV", HV)): + if HO % hh != 0: + raise ValueError(f"{name}={hh} must divide {HO}") + cu_stream = cuda.CUstream(int(stream)) dyn_sched = sched_ctr is not None - cache = _get_compiled_cache(str(q.dtype), HQ, k.shape[1], HV, is_GQA, d_initial_state is not None, d_final_state is not None, split_k, log_gate, dyn_sched) + cache = get_compiled_cache( + str(q.dtype), + str(cu_seqlens.dtype), + HQ, + HK, + HV, + is_GQA, + initial_state is not None, + d_final_state is not None, + d_initial_state is not None, + log_gate, + dyn_sched, + ) if "compiled" not in cache: - - def _tok3(t): - c = from_dlpack(t, assumed_align=16) - c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - return c - - def _tok2(t): - c = from_dlpack(t, assumed_align=16) - c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) - return c - cu_seqlens_cute = from_dlpack(cu_seqlens, assumed_align=4).mark_layout_dynamic() workspace_cute = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() - ds0_cute = None + dstate0_cute = None if d_initial_state is not None: - # prefill s_out marking (the drain reuses the fs-store indexing) - ds0_cute = from_dlpack(d_initial_state, assumed_align=16) - ds0_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=CFG.D_K) - dht_cute = None + dstate0_cute = from_dlpack(d_initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) + dstate_in_cute = None if d_final_state is not None: - dht_cute = from_dlpack(d_final_state, assumed_align=16) - dht_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=CFG.D_K) - work_items_cute = None - work_count_cute = None - if split_k: - work_items_cute = from_dlpack(work_items, assumed_align=4) - work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) - work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() + dstate_in_cute = from_dlpack(d_final_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) + work_items_cute = from_dlpack(work_items, assumed_align=16) + work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() sched_ctr_cute = None if dyn_sched: sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() cache["compiled"] = compile( io_dtype, is_GQA, - use_initial_state=d_initial_state is not None, - use_dht=d_final_state is not None, - split_k=split_k, + use_initial_state=initial_state is not None, + use_dstate_in=d_final_state is not None, + use_dstate0=d_initial_state is not None, log_gate=log_gate, dyn_sched=dyn_sched, - num_sm=_device_sm_count(), - q_cute=_tok3(q), - k_cute=_tok3(k), - v_cute=_tok3(v), - gate_cute=_tok2(gate), - beta_cute=_tok2(beta), - dg_cute=_tok2(dg), - dbeta_cute=_tok2(dbeta), - do_cute=_tok3(do), - dq_cute=_tok3(dq), - dk_cute=_tok3(dk), - dv_cute=_tok3(dv), + num_sm=multiprocessor_count(current_device_id()), + h_q=HQ, + h_k=HK, + h_v=HV, + q_cute=from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=2), + k_cute=from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2), + v_cute=from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2), + gate_cute=from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=1), + beta_cute=from_dlpack(beta, assumed_align=16).mark_layout_dynamic(leading_dim=1), + dgate_cute=from_dlpack(dgate, assumed_align=16).mark_layout_dynamic(leading_dim=1), + dbeta_cute=from_dlpack(dbeta, assumed_align=16).mark_layout_dynamic(leading_dim=1), + do_cute=from_dlpack(do, assumed_align=16).mark_layout_dynamic(leading_dim=2), + dq_cute=from_dlpack(dq, assumed_align=16).mark_layout_dynamic(leading_dim=2), + dk_cute=from_dlpack(dk, assumed_align=16).mark_layout_dynamic(leading_dim=2), + dv_cute=from_dlpack(dv, assumed_align=16).mark_layout_dynamic(leading_dim=2), cu_seqlens_cute=cu_seqlens_cute, - ds0_cute=ds0_cute, - dht_cute=dht_cute, + dstate0_cute=dstate0_cute, + dstate_in_cute=dstate_in_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, sched_ctr_cute=sched_ctr_cute, @@ -5414,54 +4526,43 @@ def _tok2(t): compiled = cache["compiled"] - # The descriptors encode cu_seqlens' CONTENTS, which no key built from the - # buffers can track. The skip this replaces asked torch's _version counter, - # so it was sound for a torch caller and silently stale for every other - # producer. Rebuilding unconditionally measures free: 131 vs 135 us of host - # time, and 157 either way once the launches are waited on. + # desc build runs every execute by contract (cu contents are data; + # buffer pointers may change) — capture-safe, single tiny launch if "build_descs" not in cache: - def _tok3_bc(t): - c = from_dlpack(t, assumed_align=16) - c.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - return c - - h_bc = from_dlpack(h, assumed_align=16) - h_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + checkpoints_bc = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=3) cu_bc = from_dlpack(cu_seqlens, assumed_align=4).mark_layout_dynamic() - s0_bc = None + state0_bc = None if initial_state is not None: - s0_bc = from_dlpack(initial_state, assumed_align=16) - s0_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + state0_bc = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) ws_bc = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() cache["build_descs"] = cute.compile( - _build_descs, + build_descs, io_dtype, CFG.B_T, - _tok3_bc(q), - _tok3_bc(k), - _tok3_bc(v), - _tok3_bc(do), - _tok3_bc(dq), - _tok3_bc(dk), - _tok3_bc(dv), - h_bc, + from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(do, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(dq, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(dk, assumed_align=16).mark_layout_dynamic(leading_dim=2), + from_dlpack(dv, assumed_align=16).mark_layout_dynamic(leading_dim=2), + checkpoints_bc, cu_bc, - s0_bc, + state0_bc, ws_bc, cu_stream, options="--enable-tvm-ffi", ) - cache["build_descs"](q, k, v, do, dq, dk, dv, h, cu_seqlens, initial_state, workspace, cu_stream) - + cache["build_descs"](q, k, v, do, dq, dk, dv, state_checkpoints, cu_seqlens, initial_state, workspace, cu_stream) compiled( q, k, v, gate, beta, - dg, + dgate, dbeta, do, dq, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_config.py b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_config.py index bce18e37c..0860db59f 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_config.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_config.py @@ -33,17 +33,14 @@ class Cfg: D_V: int = 128 # value head dim (M of GEMMs 3-6, N of GEMM 7) # --- TMA descriptor pool --- - BYTES_PER_TENSORMAP: int = 128 # --- warp assignments (12 warps total) --- COMPUTE_GROUP_0_WARP_IDS: Tuple[int, ...] = (0, 1, 2, 3) # T-pairwise / kk_epi / qk_epi / inverse COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (4, 5, 6, 7) # kv_decay_v / v-k*state / epi ops - MMA_WARP_ID: int = 8 # CG0 issuer: KK/QK per pair + LOAD_GATE_BETA_WARP_ID: int = 8 # gate/beta chunk loads + TMEM lifecycle TMA_QKV_WARP_ID: int = 9 - MMA_CG1_WARP_ID: int = 10 # CG1 issuer: KS/QS/NV/QKV/KV per chunk - # The lightly loaded O/H epilogue warp also loads gate/beta. + MMA_WARP_ID: int = 10 # sole tcgen05 issuer: fused KK/QK pairs + KS/QS/U/QKV/KV per chunk EPILOGUE_WARP_ID: int = 11 - LOAD_GATE_BETA_WARP_ID: int = 11 # --- register split --- NUM_REGS_COMPUTE_GROUP_0: int = 224 @@ -56,12 +53,11 @@ class Cfg: # --- SMEM stage counts --- SMEM_SCHED_STAGES: int = 2 - SMEM_Q_STAGES: int = 3 - SMEM_K_STAGES: int = 4 - SMEM_V_STAGES: int = 3 - SMEM_AINV_STAGES: int = 2 - SMEM_QK_STAGES: int = 2 - SMEM_O_STAGES: int = 2 + SMEM_KQ_STAGES: int = 4 + SMEM_V_STAGES: int = 2 + SMEM_T_INV_STAGES: int = 3 + SMEM_A_STAGES: int = 3 + SMEM_O_STAGES: int = 1 SMEM_GATE_STAGES: int = 3 SMEM_BETA_STAGES: int = 3 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 a19bc6864..cec6bd144 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_prefill_f16.py @@ -17,10 +17,10 @@ """ Chunked Gated Delta Net (GDN) prefill kernel for Blackwell SM100 (Cutlass primitives) -with optional per-chunk state (H) output. +with optional per-chunk state-checkpoint output. Algorithm overview (per chunk c, tokens [cC, (c+1)C)): - Inputs : Q[BT,DK], K[BT,DK], V[BT,DV], gate[BT] (scalar gate), beta[BT] (scalar LR) + Inputs : Q[BT,DK], K[BT,DK], V[BT,DV], Gate[BT] (scalar gate), Beta[BT] (scalar LR) State : S_prev[DK,DV] (recurrent state, held in TMEM) Preprocessing (compute warp group 0): @@ -29,56 +29,56 @@ T_pairwise[i,j] = cumprod[i] / cumprod[j] (i>=j) inter-token transfer weights (stored in registers; 128 regs/thread) - GEMM 1 - kk : W_kk[BT,BT] = K @ K^T (lower-triangular intra scores) - GEMM 2 - qk : W_qk[BT,BT] = Q @ K^T (output attention scores) - GEMM 3 - k*state : KS[BT,DV] = K @ S_prev (key applied to state) - GEMM 4 - q*state : QS[BT,DV] = Q @ S_prev (inter-chunk output, before T scaling) - GEMM 5 - new v : NV[BT,DV] = A_inv @ V (corrected value vectors) - where A_inv = (I + M_kk)^{-1}, M_kk[i,j] = T[i,j]*beta[i]*W_kk[i,j] (lower-tri, hierarchical blockwise inverse) - GEMM 6 - qkv : O_intra[BT,DV] = W_qkv @ NV (intra-chunk output) - where W_qkv = T*beta*W_qk (scaled qk scores) - GEMM 7 - kv update : dS[DK,DV] = K^T @ delta (state update, BT contraction) - where delta[BT,DV] = V - KS (delta rule residuals, after decay) + GEMM 1 - KK : W_kk[BT,BT] = K @ K^T (lower-triangular intra scores) + GEMM 2 - QK : W_qk[BT,BT] = Q @ K^T (output attention scores) + GEMM 3 - K*state : KS[BT,DV] = K @ S_prev (key applied to state) + GEMM 4 - Q*state : QS[BT,DV] = Q @ S_prev (inter-chunk output, before T scaling) + GEMM 5 - U : U[BT,DV] = T_inv @ Y (corrected value vectors) + where T_inv = (I + M_kk)^{-1}, M_kk[i,j] = T[i,j]*Beta[i]*W_kk[i,j] (lower-tri, hierarchical blockwise inverse) + GEMM 6 - QKV : O_intra[BT,DV] = W_qkv @ U (intra-chunk output) + where W_qkv = T*Beta*W_qk (the A tile) + GEMM 7 - KV update : S_upd[DK,DV] = K^T @ (decay .* U) (state update, BT contraction) + where Y[BT,DV] = V - KS (delta rule residuals, after decay) Epilogue: O[BT,DV] = O_intra + T_col * QS (combine intra + inter) - S_next = cumprod[BT-1] * S_prev + dS (update state in TMEM) + S_next = cumprod[BT-1] * S_prev + S_upd (update state in TMEM) Chunks run in PAIRS (CG0 warp halves invert chunk 0 / chunk 1 in parallel); odd counts pad with a neutral zero-filled chunk. SMEM layout (227 KB = full; stage counts live in gdn_prefill_config.py; -enable_h compiles trim K/V stages to fit the H buffer): +enable_checkpoints compiles trim K/V stages to fit the checkpoint buffer): Buffer Size (B) Stages - q 16384 3 - k 16384 4 - v 16384 3 - A_inverse / new_v 8192 2 - QK output 8192 2 + Q 16384 3 + K 16384 4 + V 16384 3 + T_inv 8192 3 + A tile output 8192 3 O store 16384 2 - H staging DK*DV*2 1 <-- enable_h only - cumsumlog / cumprod / beta 256 3 + checkpoint staging DK*DV*2 1 <-- enable_checkpoints only + cumsumlog / cumprod / Beta 256 3 sched ticket ring 4 2 <-- dyn_sched publish ring TMEM layout (512 columns): Buffer Cols - state (S) 128 <-- DKxDV fp32 = 128x128x4B - q*state / O acc 64 <-- BTxDV fp32 accumulator + state 128 <-- DKxDV fp32 = 128x128x4B + Q*state / O acc 64 <-- BTxDV fp32 accumulator state inp 64 <-- fp16 state staging (GEMMs 3/4 A operand) cg0 shared acc 128 <-- 2-stage ring: KK0/KK1 then QK0/QK1 - cg1 shared acc 64 <-- 1-stage ring: KS then NV - vks+nv / decay_v inp 64 <-- slot 0 = VKS then NV, slot 1 = decay_v + cg1 shared acc 64 <-- 1-stage ring: KS then U + Y + U input / decayed-U inp 64 <-- slot 0 = Y then U input, slot 1 = decayed U (b16) Warp assignments (12 warps = 384 threads): - warps 0-3 : compute group 0 - T-pairwise x2, kk_epi x2, pair inverse, - qk_epi x2 - warps 4-7 : compute group 1 - state restage/rescale, v-k*state, - state*q_epi, new_v_epi, qkv_epilogue - warp 8 : CG0 MMA issuer - KK0/KK1/QK0/QK1 per pair; TMEM lifecycle - warp 9 : TMA load warp - loads q, k, v - warp 10 : CG1 MMA issuer - KS/QS/NV/QKV/KV per chunk - warp 11 : epilogue warp - gate/beta loads (4-chunk lookahead) + - O then H TMA stores + warps 0-3 : compute group 0 - T-pairwise x2, KK_epi x2, pair inverse, + A_epi x2 + warps 4-7 : compute group 1 - state restage/rescale, Y = V - K*state, + state*Q_epi, U_epi, QKV_epilogue + warp 8 : Gate/Beta loads + warp 9 : TMA load warp - loads Q, K, V + warp 10 : MMA warp - fused KK/QK pairs + K*state/Q*state/U/QKV/KV per chunk; + TMEM lifecycle + warp 11 : epilogue warp - O then checkpoint TMA stores """ import functools @@ -90,13 +90,15 @@ import cutlass import cutlass.cute as cute import cutlass.experimental.primitives as nvvm -import cutlass.experimental.cuda.tensor_map as _tma -from cutlass.cute.arch.nvvm_wrappers import inline_ptx +import cutlass.experimental.cuda.tensor_map as tma from cutlass.cute.runtime import from_dlpack -from cutlass.cutlass_dsl import min as _cutlass_min +from cutlass.cutlass_dsl import min -from ..common.thd import build_h_descs_kernel, build_qkv_load_descs_kernel, downcast_state_kernel, TENSOR_MAP_QWORDS +from ..common.thd import emit_checkpoint_seq_descs, emit_seq_descs, TENSOR_MAP_QWORDS from ..common.split_k import decode_work_item +from ..common.host import get_dtype +from cudnn.frost.buffers import current_device_id, data_ptr +from cudnn.frost.device import multiprocessor_count RCP_LN2 = 1.4426950408889634 # 1/ln(2): natural-log gates -> the kernel's log2 domain from cudnn.frost.tile_dsl.barrier import ( @@ -104,10 +106,9 @@ Producer, PipelineState, advance, - arrive, ) from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, tma_slice_runtime_desc -from cudnn.frost.tile_dsl.mma import mma_ss, mma_ts, mma_step +from cudnn.frost.tile_dsl.mma import mma_ss, mma_step_k8, mma_ts_step, mma_step from cudnn.frost.tile_dsl.pointwise import fadd2, fp32_to_fp16, f16x2_to_f32, fmul2, opaque_f32_zero, sub_f16x2 from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_xor_128b from cudnn.frost.tile_dsl.tma import ( @@ -129,10 +130,8 @@ class GdnBars(NamedTuple): slot by arriving ``_done``. """ - mb_q_ready: MBarrier - mb_q_done: MBarrier - mb_k_ready: MBarrier - mb_k_done: MBarrier + mb_kq_ready: MBarrier + mb_kq_done: MBarrier mb_v_ready: MBarrier mb_v_done: MBarrier @@ -141,43 +140,44 @@ class GdnBars(NamedTuple): mb_beta_ready: MBarrier mb_beta_done: MBarrier - mb_kv_acc_ready: MBarrier - mb_kv_acc_scale_done: MBarrier + mb_state_acc_ready: MBarrier + mb_state_acc_scale_done: MBarrier mb_o_acc_ready: MBarrier - mb_o_acc_done: MBarrier - mb_o_state_scale_acc_ready: MBarrier + mb_o_final_acc_ready: MBarrier mb_o_state_scale_acc_done: MBarrier mb_cg0_acc_ready: MBarrier mb_cg0_acc_done: MBarrier - mb_ks_ready: MBarrier - mb_nv_ready: MBarrier - mb_ainv_ready: MBarrier - mb_ainv_done: MBarrier - mb_qk_ready: MBarrier - mb_qk_done: MBarrier mb_state_inp_ready: MBarrier - mb_vks_inp_ready: MBarrier - mb_nv_inp_ready: MBarrier - mb_decay_v_inp_ready: MBarrier + mb_y_inp_ready: MBarrier + mb_u_inp_ready: MBarrier + mb_decay_u_inp_ready: MBarrier + + mb_t_inv_ready: MBarrier + mb_t_inv_done: MBarrier + mb_a_ready: MBarrier + mb_a_done: MBarrier + + mb_k_state_acc_ready: MBarrier + mb_u_acc_ready: MBarrier mb_o_tmastg_ready: MBarrier mb_o_tmastg_done: MBarrier - mb_h_tmastg_ready: MBarrier - mb_h_tmastg_done: MBarrier + mb_checkpoint_tmastg_ready: MBarrier + mb_checkpoint_tmastg_done: MBarrier mb_tmem_done: MBarrier + mb_sched_ready: MBarrier mb_sched_done: MBarrier def make_gdn_bars(cfg) -> GdnBars: - """GdnBars factory. MUST be called from inside ``_kernel`` (allocates SMEM; - the mbar rings sit ahead of the gate scalar arrays and data buffers).""" + """GdnBars factory. MUST be called from inside ``kernel`` (allocates SMEM).""" ONE_LANE = 1 MMA_ARRIVERS = len([cfg.mma_warp_id]) - BOTH_ISSUERS = len([cfg.mma_warp_id, cfg.mma_cg1_warp_id]) + KQ_RELEASE_SITES = 1 GATE_WARP = cfg.threads_per_warp * len([cfg.load_gate_beta_warp_id]) EPI_WARP = cfg.threads_per_warp * len([cfg.epilogue_warp_id]) CG0_THREADS = cfg.threads_per_warp * len(cfg.compute_group_0_warp_ids) @@ -188,60 +188,50 @@ def alloc(n): return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=16) return GdnBars( - mb_q_ready=MBarrier(alloc(cfg.smem_q_stages), stages=cfg.smem_q_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), - mb_q_done=MBarrier(alloc(cfg.smem_q_stages), stages=cfg.smem_q_stages, init_count=BOTH_ISSUERS, producer=Producer.MMA_COMMIT), - mb_k_ready=MBarrier(alloc(cfg.smem_k_stages), stages=cfg.smem_k_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), - mb_k_done=MBarrier(alloc(cfg.smem_k_stages), stages=cfg.smem_k_stages, init_count=BOTH_ISSUERS, producer=Producer.MMA_COMMIT), + mb_kq_ready=MBarrier(alloc(cfg.smem_kq_stages), stages=cfg.smem_kq_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), + mb_kq_done=MBarrier(alloc(cfg.smem_kq_stages), stages=cfg.smem_kq_stages, init_count=KQ_RELEASE_SITES, producer=Producer.MMA_COMMIT), mb_v_ready=MBarrier(alloc(cfg.smem_v_stages), stages=cfg.smem_v_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), mb_v_done=MBarrier(alloc(cfg.smem_v_stages), stages=cfg.smem_v_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_gate_ready=MBarrier(alloc(cfg.smem_gate_stages), stages=cfg.smem_gate_stages, init_count=GATE_WARP, producer=Producer.THREAD), mb_gate_done=MBarrier(alloc(cfg.smem_gate_stages), stages=cfg.smem_gate_stages, init_count=CG0_PLUS_CG1, producer=Producer.THREAD), mb_beta_ready=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=GATE_WARP, producer=Producer.THREAD), mb_beta_done=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_kv_acc_ready=MBarrier(alloc(cfg.tmem_kv_acc_stages), stages=cfg.tmem_kv_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_kv_acc_scale_done=MBarrier(alloc(cfg.tmem_kv_acc_stages), stages=cfg.tmem_kv_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_state_acc_ready=MBarrier(alloc(cfg.tmem_state_acc_stages), stages=cfg.tmem_state_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_state_acc_scale_done=MBarrier(alloc(cfg.tmem_state_acc_stages), stages=cfg.tmem_state_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_o_acc_ready=MBarrier(alloc(cfg.tmem_q_state_acc_stages), stages=cfg.tmem_q_state_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_o_acc_done=MBarrier(alloc(cfg.tmem_q_state_acc_stages), stages=cfg.tmem_q_state_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_o_state_scale_acc_ready=MBarrier( + mb_o_final_acc_ready=MBarrier( alloc(cfg.tmem_q_state_acc_stages), stages=cfg.tmem_q_state_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT ), mb_o_state_scale_acc_done=MBarrier( alloc(cfg.tmem_q_state_acc_stages), stages=cfg.tmem_q_state_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD ), mb_cg0_acc_ready=MBarrier(alloc(cfg.tmem_cg0_acc_stages), stages=cfg.tmem_cg0_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_cg0_acc_done=MBarrier(alloc(cfg.tmem_cg0_acc_stages), stages=cfg.tmem_cg0_acc_stages, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_ks_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_nv_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_ainv_ready=MBarrier(alloc(cfg.smem_ainv_stages), stages=cfg.smem_ainv_stages, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_ainv_done=MBarrier(alloc(cfg.smem_ainv_stages), stages=cfg.smem_ainv_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), - mb_qk_ready=MBarrier(alloc(cfg.smem_qk_stages), stages=cfg.smem_qk_stages, init_count=CG0_THREADS, producer=Producer.THREAD), - mb_qk_done=MBarrier(alloc(cfg.smem_qk_stages), stages=cfg.smem_qk_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_cg0_acc_done=MBarrier(alloc(cfg.tmem_cg0_acc_stages), stages=cfg.tmem_cg0_acc_stages, init_count=CG0_THREADS // 2, producer=Producer.THREAD), mb_state_inp_ready=MBarrier(alloc(cfg.tmem_state_inp_stages), stages=cfg.tmem_state_inp_stages, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_vks_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_nv_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_decay_v_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_decay_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_t_inv_ready=MBarrier(alloc(cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_t_inv_done=MBarrier(alloc(cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_a_ready=MBarrier(alloc(cfg.smem_a_stages), stages=cfg.smem_a_stages, init_count=CG0_THREADS // 2, producer=Producer.THREAD), + mb_a_done=MBarrier(alloc(cfg.smem_a_stages), stages=cfg.smem_a_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_k_state_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), mb_o_tmastg_ready=MBarrier(alloc(cfg.smem_o_stages), stages=cfg.smem_o_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_o_tmastg_done=MBarrier(alloc(cfg.smem_o_stages), stages=cfg.smem_o_stages, init_count=EPI_WARP, producer=Producer.THREAD), - mb_h_tmastg_ready=MBarrier(alloc(cfg.smem_h_stages), stages=cfg.smem_h_stages, init_count=len(cfg.compute_group_1_warp_ids), producer=Producer.THREAD), - mb_h_tmastg_done=MBarrier(alloc(cfg.smem_h_stages), stages=cfg.smem_h_stages, init_count=EPI_WARP, producer=Producer.THREAD), + mb_checkpoint_tmastg_ready=MBarrier( + alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=len(cfg.compute_group_1_warp_ids), producer=Producer.THREAD + ), + mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=EPI_WARP, producer=Producer.THREAD), mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=11, producer=Producer.THREAD), ) -# --------------------------------------------------------------------------- -# Device-side helpers / warp bodies -# --------------------------------------------------------------------------- - - @cute.jit -def _invert_diagonal_NxN(cfg, base_int, d, tidx, N: int = 8): - """Stage 1: Gauss-Jordan inversion of one diagonal NxN block in-place (f16 SMEM). - - The tile swizzle re-homes whole rows only, so a diagonal block's row stays - a contiguous N-element run at ``swz(row_lin_base)``. - """ +def invert_diagonal_NxN(cfg, base_int, d, tidx, N: int = 8): + """Gauss-Jordan inversion of one diagonal NxN block in-place (f16 SMEM).""" tidx_in_group = tidx % N BT = cfg.b_t @@ -272,70 +262,65 @@ def _invert_diagonal_NxN(cfg, base_int, d, tidx, N: int = 8): @cute.jit -def _mma_m16n8k8(a0, a1, b0, c_regs, dtype: cutlass.Constexpr): - """One m16n8k8 reg-reg mma (``mma_step`` only emits the k16 form), - accumulating into the ``c_regs`` buffer in place.""" - tag = "f16" if cutlass.const_expr(dtype == cutlass.Float16) else "bf16" - c_regs[0], c_regs[1], c_regs[2], c_regs[3] = inline_ptx( - f"mma.sync.aligned.m16n8k8.row.col.f32.{tag}.{tag}.f32" " {$0,$1,$2,$3}, {$4,$5}, {$6}, {$7,$8,$9,$10};", - write_only_types=[cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Float32], - read_only_args=[a0, a1, b0, c_regs[0], c_regs[1], c_regs[2], c_regs[3]], - ) - - -@cute.jit -def _blockwise_diagonal_8x8_to_16x16(cfg, base_int, d0, lane_id): - """Stage 2: off-diagonal correction 8x8 -> 16x16 (C <- -D^{-1} C A^{-1}). - - Keep the per-lane ldmatrix offset (``lds1``/``lds4``) a SEPARATE sum term - from the warp-uniform (row, col) origin in all the diagonal helpers — - folding it into the row term costs ~2% (per-lane address datapath). - """ +def blockwise_diagonal_8x8_to_16x16(cfg, base_int, d0, lane_id): + """Off-diagonal correction 8x8 -> 16x16 (C <- -D^{-1} C A^{-1}).""" bpe = cfg.io_dtype.width // 8 - lds1 = (lane_id % 8) * 64 + ldsm_x1_lane_off = (lane_id % 8) * 64 d = nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + 8 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + cutlass.inttoptr( + base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + 8 + ldsm_x1_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), 1, nvvm.MMALayout.ROW, ) c = nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + cutlass.inttoptr( + base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + ldsm_x1_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), 1, nvvm.MMALayout.COL, ) + + # ---- T = -(D^{-1} @ C) ------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(4): c_regs[i] = cutlass.Float32(0.0) - _mma_m16n8k8(d, d, c, c_regs, cfg.io_dtype) + mma_step_k8(c_regs, [d, d], [c], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(4): c_regs[i] = -c_regs[i] - a_f16 = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(2)] - ai = nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(2)] + + # ---- C = T @ A^{-1} ---------------------------------------------------------- + ai_frag = nvvm.ldmatrix( + cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + ldsm_x1_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), 1, nvvm.MMALayout.COL, ) o_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(4): o_regs[i] = cutlass.Float32(0.0) - _mma_m16n8k8(a_f16[0], a_f16[1], ai, o_regs, cfg.io_dtype) - o_f16 = fp32_to_fp16(o_regs[0], o_regs[1], dtype=cfg.io_dtype) + mma_step_k8(o_regs, a_pack, [ai_frag], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) + o_pack = fp32_to_fp16(o_regs[0], o_regs[1], dtype=cfg.io_dtype) + + # ---- store corrected C ------------------------------------------------------- nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - o_f16, + cutlass.inttoptr( + base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + ldsm_x1_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), + o_pack, nvvm.MMALayout.ROW, ) @cute.jit -def _blockwise_diagonal_16x16_to_32x32(cfg, base_int, d0, lane_id): - """Stage 3: off-diagonal correction 16x16 -> 32x32.""" +def blockwise_diagonal_16x16_to_32x32(cfg, base_int, d0, lane_id): + """Off-diagonal correction 16x16 -> 32x32.""" bpe = cfg.io_dtype.width // 8 - lds4 = (lane_id % 16) * 64 + (lane_id // 16) * 8 + ldsm_x4_lane_off = (lane_id % 16) * 64 + (lane_id // 16) * 8 d = list( nvvm.ldmatrix( cutlass.inttoptr( - base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + 16 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + 16 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 ), 4, nvvm.MMALayout.ROW, @@ -343,21 +328,29 @@ def _blockwise_diagonal_16x16_to_32x32(cfg, base_int, d0, lane_id): ) c = list( nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + cutlass.inttoptr( + base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), 4, nvvm.MMALayout.COL, ) ) + + # ---- T = -(D^{-1} @ C) ------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(8): c_regs[i] = cutlass.Float32(0.0) mma_step(c_regs, d, c, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(8): c_regs[i] = -c_regs[i] - a_f16 = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - ai = list( + a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + + # ---- C = T @ A^{-1} ---------------------------------------------------------- + ai_frag = list( nvvm.ldmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + cutlass.inttoptr( + base_int + swizzle_lin_128b(d0 * 64 + d0 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), 4, nvvm.MMALayout.COL, ) @@ -365,28 +358,31 @@ def _blockwise_diagonal_16x16_to_32x32(cfg, base_int, d0, lane_id): o_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(8): o_regs[i] = cutlass.Float32(0.0) - mma_step(o_regs, a_f16, ai, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) - ow = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + mma_step(o_regs, a_pack, ai_frag, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) + o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + + # ---- store corrected C ------------------------------------------------------- nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - ow, + cutlass.inttoptr( + base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), + o_pack, nvvm.MMALayout.ROW, ) @cute.jit -def _blockwise_diagonal_32x32_to_64x64(cfg, base_int, warp_id, lane_id): - """Stage 4: off-diagonal correction 32x32 -> 64x64 (2 warps, one 16-row - M-band each).""" +def blockwise_diagonal_32x32_to_64x64(cfg, base_int, warp_id, lane_id): + """Off-diagonal correction 32x32 -> 64x64 (2 warps, one 16-row M-band each).""" band = warp_id % 2 bpe = cfg.io_dtype.width // 8 - lds4 = (lane_id % 16) * 64 + (lane_id // 16) * 8 - a_regs = [] + ldsm_x4_lane_off = (lane_id % 16) * 64 + (lane_id // 16) * 8 + a_frags = [] for vs in cutlass.range_constexpr(2): - a_regs += list( + a_frags += list( nvvm.ldmatrix( cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + band * 16) * 64 + 32 + vs * 16 + lds4, row_stride_log2=6) * bpe, + base_int + swizzle_lin_128b((32 + band * 16) * 64 + 32 + vs * 16 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), @@ -394,12 +390,12 @@ def _blockwise_diagonal_32x32_to_64x64(cfg, base_int, warp_id, lane_id): nvvm.MMALayout.ROW, ) ) - b_regs = [] + b_frags = [] for vs in cutlass.range_constexpr(4): - b_regs += list( + b_frags += list( nvvm.ldmatrix( cutlass.inttoptr( - base_int + swizzle_lin_128b((32 + (vs // 2) * 16) * 64 + (vs % 2) * 16 + lds4, row_stride_log2=6) * bpe, + base_int + swizzle_lin_128b((32 + (vs // 2) * 16) * 64 + (vs % 2) * 16 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), @@ -407,20 +403,24 @@ def _blockwise_diagonal_32x32_to_64x64(cfg, base_int, warp_id, lane_id): nvvm.MMALayout.COL, ) ) + + # ---- T = -(D^{-1} @ C) ------------------------------------------------------- c_regs = cutlass.Array(cutlass.Float32, 16, alignment=16, space=cutlass.AddressSpace.rmem) for i in cutlass.range_constexpr(16): c_regs[i] = cutlass.Float32(0.0) for ks in cutlass.range_constexpr(2): - mma_step(c_regs, a_regs, b_regs[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + mma_step(c_regs, a_frags, b_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) for i in cutlass.range_constexpr(16): c_regs[i] = -c_regs[i] - a_f16 = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] - ai_regs = [] + a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] + + # ---- C = T @ A^{-1} ---------------------------------------------------------- + ai_frags = [] for vs in cutlass.range_constexpr(4): - ai_regs += list( + ai_frags += list( nvvm.ldmatrix( cutlass.inttoptr( - base_int + swizzle_lin_128b(((vs // 2) * 16) * 64 + (vs % 2) * 16 + lds4, row_stride_log2=6) * bpe, + base_int + swizzle_lin_128b(((vs // 2) * 16) * 64 + (vs % 2) * 16 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), @@ -432,115 +432,35 @@ def _blockwise_diagonal_32x32_to_64x64(cfg, base_int, warp_id, lane_id): for i in cutlass.range_constexpr(16): o_regs[i] = cutlass.Float32(0.0) for ks in cutlass.range_constexpr(2): - mma_step(o_regs, a_f16, ai_regs[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) - ow = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] + mma_step(o_regs, a_pack, ai_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] + + # ---- store corrected C ------------------------------------------------------- nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((32 + band * 16) * 64 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - ow[0:4], + cutlass.inttoptr( + base_int + swizzle_lin_128b((32 + band * 16) * 64 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), + o_pack[0:4], nvvm.MMALayout.ROW, ) nvvm.stmatrix( - cutlass.inttoptr(base_int + swizzle_lin_128b((32 + band * 16) * 64 + 16 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), - ow[4:8], + cutlass.inttoptr( + base_int + swizzle_lin_128b((32 + band * 16) * 64 + 16 + ldsm_x4_lane_off, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), + o_pack[4:8], nvvm.MMALayout.ROW, ) -@cute.jit -def _load_gate_beta_chunk( - cfg, - lidx, - mGate, - mBeta, - sCumsumlog, - sCumprod, - sBeta, - gate_index, - beta_index, - head_idx, - batch_start, - batch_end, - chunk_idx, - bars, -): - """Load gate[BT]/beta[BT] for one chunk (epilogue warp). - - The OOB predicate is RUNTIME (covers the last valid chunk AND padded - pair chunks), keeping one body in SASS. OOB gate positions read a - clamped in-bounds address and select the neutral value (gate=1 -> - log 0); OOB beta lanes zero-fill via cp_size=0.""" - n_cols = cfg.b_t // cfg.threads_per_warp - chunk_offset = batch_start + chunk_idx * cfg.b_t - gGateSeq = mGate[None, head_idx] - gBeta = cute.domain_offset((chunk_offset,), mBeta[None, head_idx]) - - gate_idx = gate_index.idx - gate_phase = gate_index.phase - gate_index = advance(gate_index, cfg.smem_gate_stages) - - pos_valid = [None] * n_cols - tGrGate = [cutlass.Float32(0.0)] * n_cols - oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) - for col in cutlass.range_constexpr(n_cols): - tok = chunk_offset + lidx + col * cfg.threads_per_warp - pos_valid[col] = cute.elem_less(tok, batch_end) - # batch_end >= 1 whenever chunks exist, so the clamp stays in bounds - tok_clamped = _cutlass_min(tok, batch_end - 1) - tGrGate[col] = gGateSeq[tok_clamped] if pos_valid[col] else oob_neutral - - if cutlass.const_expr(cfg.log_gate): - for col in cutlass.range_constexpr(n_cols): - tGrGate[col] = tGrGate[col] * cutlass.Float32(RCP_LN2) - else: - for col in cutlass.range_constexpr(n_cols): - tGrGate[col] = cute.math.log2(tGrGate[col] + 1e-10, fastmath=True) - for offset in [1, 2, 4, 8, 16]: - for col in cutlass.range_constexpr(n_cols): - n = nvvm.shfl_sync(0xFFFFFFFF, tGrGate[col], offset, 0, kind=nvvm.Shfl.UP) - if lidx >= offset: - tGrGate[col] = tGrGate[col] + n - for col in cutlass.range_constexpr(1, n_cols): - last_v = nvvm.shfl_sync( - 0xFFFFFFFF, - tGrGate[col - 1], - cfg.threads_per_warp - 1, - cfg.threads_per_warp - 1, - kind=nvvm.Shfl.IDX, - ) - tGrGate[col] += last_v - - bars.mb_gate_done[gate_idx].wait(gate_phase) - for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp - sCumsumlog[pos, 0, gate_idx] = tGrGate[col] - sCumprod[pos, 0, gate_idx] = cute.math.exp2(tGrGate[col], fastmath=True) - - bars.mb_gate_ready[gate_idx].arrive() - - # --- Beta load (per-element async G->S cp.async) --- - beta_idx = beta_index.idx - bars.mb_beta_done[beta_idx].wait(beta_index.phase) - beta_index = advance(beta_index, cfg.smem_beta_stages) - for col in cutlass.range_constexpr(n_cols): - pos = lidx + col * cfg.threads_per_warp - src = gBeta.iterator + gBeta.layout((pos,)) - dst = sBeta.iterator + sBeta.layout((pos, 0, beta_idx)) - cp_size = cutlass.Int32(4) * cutlass.Int32(pos_valid[col]) - nvvm.cp_async_shared_global(dst, src, 4, nvvm.LoadCacheModifier.CA, cp_size=cp_size) - nvvm.cp_async_mbarrier_arrive(bars.mb_beta_ready[beta_idx].smem_ptr, noinc=True) - return gate_index, beta_index - - -# --------------------------------------------------------------------------- -# Dynamic tile scheduler: global-ticket work-stealing ring +# ---- Dynamic tile scheduler ------------------------------------------------------ @cute.jit -def _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): """TMA-LDG-warp side: pull the next tile off the global ticket, publish it.""" if cutlass.const_expr(cfg.dyn_sched): bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) @@ -556,7 +476,7 @@ def _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ct @cute.jit -def _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): """Consumer side: read the TMA-LDG warp's published next tile.""" if cutlass.const_expr(cfg.dyn_sched): bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) @@ -568,7 +488,7 @@ def _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): @cute.jit -def _tmastg_warp( +def tmastg_warp( cfg, total_tiles, bidx, @@ -577,31 +497,24 @@ def _tmastg_warp( mWorkItems, checkpoint_every_n_tokens, tidx, - mGate, - mBeta, - sCumsumlog, - sCumprod, - sBeta, sO_raw, - sH_raw, + sCheckpoint_raw, desc_o_base, - desc_h_base, + desc_checkpoint_base, sSched, bars, ): - """Epilogue warp role (warp 11): persistent tile-scheduler loop; loads - gate/beta with a four-chunk lookahead and issues the per-chunk O and - H-state TMA bulk-stores from SMEM staging to global memory.""" + """Epilogue warp role (warp 11): persistent scheduler loop issuing the + per-chunk O and state-checkpoint TMA stores.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) o_index = PipelineState.start(phase=0) - gate_index = PipelineState.start(phase=1) - beta_index = PipelineState.start(phase=1) lidx = tidx % cfg.threads_per_warp sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) bpe = cfg.io_dtype.width // 8 - granu = 128 // bpe + elems_per_128b = 128 // bpe sO_tma = SmemTile( base=sO_raw, elems_per_stage=(cfg.o_cosize // cfg.smem_o_stages), @@ -610,362 +523,270 @@ def _tmastg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, + tma_granu_elems=elems_per_128b, tma_subtile_stride_elems=4096, ) - if cutlass.const_expr(cfg.enable_h): - h_granu = 64 - sH_tma = SmemTile( - base=sH_raw, - elems_per_stage=(cfg.h_cosize // cfg.smem_h_stages), - stages=cfg.smem_h_stages, + if cutlass.const_expr(cfg.enable_checkpoints): + checkpoint_elems_per_128b = 64 + sCheckpoint_tma = SmemTile( + base=sCheckpoint_raw, + elems_per_stage=(cfg.checkpoint_cosize // cfg.smem_checkpoint_stages), + stages=cfg.smem_checkpoint_stages, leading_byte_offset=0, stride_byte_offset=0, layout=0, - tma_loads_per_tile=cfg.d_v // h_granu, - tma_granu_elems=h_granu, - tma_subtile_stride_elems=cfg.d_k * h_granu, + tma_loads_per_tile=cfg.d_v // checkpoint_elems_per_128b, + tma_granu_elems=checkpoint_elems_per_128b, + tma_subtile_stride_elems=cfg.d_k * checkpoint_elems_per_128b, ) - hr_cnt = cutlass.Int32(0) + checkpoint_store_cnt = cutlass.Int32(0) + ckpt_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) heads_out = cutlass.Int32(cfg.n_heads_out) desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) 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, cu_seqlens, mWorkItems - ) + 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 - slot = (batch_idx * heads_out + head_idx) * desc_qwords - if cutlass.const_expr(cfg.enable_o): - desc_o_slot = (desc_o_base + slot).tospace(cutlass.AddressSpace.generic) - if nvvm.elect_sync(): - tma_tensormap_acquire(desc_o_slot) - if cutlass.const_expr(cfg.enable_h): - desc_h_slot = (desc_h_base + slot).tospace(cutlass.AddressSpace.generic) - # sequence-local H index of this item's first owned entry - h_coord = wstart - 1 if wstart > 0 else cutlass.Int32(0) - if nvvm.elect_sync(): - tma_tensormap_acquire(desc_h_slot) + head_o = head_idx + slot = batch_idx * desc_qwords + desc_o_slot = (desc_o_base + slot).tospace(cutlass.AddressSpace.generic) + if elect_one: + 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_mod = (cstart + cutlass.Int32(1)) % ckpt_chunks + if elect_one: + tma_tensormap_acquire(desc_checkpoint_slot) if n_local > 0: - # ---- gate/beta lookahead: chunk count is padded even, so the - # first two stages always exist and chunks 2/3 come together ---- - for pf in range(2): - gate_index, beta_index = _load_gate_beta_chunk( - cfg, lidx, mGate, mBeta, sCumsumlog, sCumprod, sBeta, gate_index, beta_index, head_idx, batch_start, batch_end, cstart + pf, bars - ) - if n_padded > 2: - for pf in range(2, 4): - gate_index, beta_index = _load_gate_beta_chunk( - cfg, - lidx, - mGate, - mBeta, - sCumsumlog, - sCumprod, - sBeta, - gate_index, - beta_index, - head_idx, - batch_start, - batch_end, - cstart + pf, - bars, - ) - for local_idx in cutlass.range(n_padded): - if local_idx + 4 < n_padded: - gate_index, beta_index = _load_gate_beta_chunk( - cfg, - lidx, - mGate, - mBeta, - sCumsumlog, - sCumprod, - sBeta, - gate_index, - beta_index, - head_idx, - batch_start, - batch_end, - cstart + local_idx + 4, - bars, - ) chunk_idx = cstart + local_idx did_o = cutlass.Int32(0) - if cutlass.const_expr(cfg.enable_o): - o_idx = o_index.idx - bars.mb_o_tmastg_ready[o_idx].wait(o_index.phase) - o_index = advance(o_index, cfg.smem_o_stages) - - # padded / warmup chunks stage O but never store it - if chunk_idx >= wstart and chunk_idx < wend: - tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) - o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), tok_coord) - tma_store_tile(sO_tma[o_idx], o_slice, acquire=False) - tma_store_commit() - did_o = cutlass.Int32(1) - - did_h = cutlass.Int32(0) - if cutlass.const_expr(cfg.enable_h): + o_idx = o_index.idx + bars.mb_o_tmastg_ready[o_idx].wait(o_index.phase) + o_index = advance(o_index, cfg.smem_o_stages) + + if chunk_idx >= wstart and chunk_idx < wend: + tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) + o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), head_o, tok_coord) + tma_store_tile(sO_tma[o_idx], o_slice, acquire=False) + tma_store_commit() + did_o = cutlass.Int32(1) + + 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 (cfg.b_t * (chunk_idx + 1)) % checkpoint_every_n_tokens == 0: - bars.mb_h_tmastg_ready[hr_cnt % cfg.smem_h_stages].wait((hr_cnt // cfg.smem_h_stages) & cutlass.Int32(1)) - h_slice = tma_slice_runtime_desc(desc_h_slot, cutlass.Int32(0), cutlass.Int32(0), h_coord) - tma_store_tile(sH_tma[hr_cnt % cfg.smem_h_stages], h_slice, acquire=False) + 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) + tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) tma_store_commit() - h_coord += 1 - did_h = cutlass.Int32(1) - - if cutlass.const_expr(cfg.enable_o): - if cutlass.const_expr(cfg.enable_h): - if did_o == 1 and did_h == 1: - tma_store_wait(1) - bars.mb_o_tmastg_done[o_idx].arrive() - tma_store_wait(0) - bars.mb_h_tmastg_done[hr_cnt % cfg.smem_h_stages].arrive() - hr_cnt = hr_cnt + 1 - if did_o == 1 and did_h == 0: - tma_store_wait(0) - bars.mb_o_tmastg_done[o_idx].arrive() - if did_o == 0: - if did_h == 1: - tma_store_wait(0) - bars.mb_h_tmastg_done[hr_cnt % cfg.smem_h_stages].arrive() - hr_cnt = hr_cnt + 1 - bars.mb_o_tmastg_done[o_idx].arrive() - else: - if did_o == 1: + checkpoint_coord += 1 + did_checkpoint = cutlass.Int32(1) + 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): + if did_o == 1 and did_checkpoint == 1: + tma_store_wait(1) + bars.mb_o_tmastg_done[o_idx].arrive() + tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() + checkpoint_store_cnt = checkpoint_store_cnt + 1 + if did_o == 1 and did_checkpoint == 0: + tma_store_wait(0) + bars.mb_o_tmastg_done[o_idx].arrive() + if did_o == 0: + if did_checkpoint == 1: tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() + checkpoint_store_cnt = checkpoint_store_cnt + 1 bars.mb_o_tmastg_done[o_idx].arrive() else: - if cutlass.const_expr(cfg.enable_h): - if did_h == 1: - tma_store_wait(0) - bars.mb_h_tmastg_done[hr_cnt % cfg.smem_h_stages].arrive() - hr_cnt = hr_cnt + 1 + if did_o == 1: + tma_store_wait(0) + bars.mb_o_tmastg_done[o_idx].arrive() - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) - - for _ in range(cfg.smem_gate_stages): - bars.mb_gate_done[gate_index.idx].wait(gate_index.phase) - gate_index = advance(gate_index, cfg.smem_gate_stages) - for _ in range(cfg.smem_beta_stages): - bars.mb_beta_done[beta_index.idx].wait(beta_index.phase) - beta_index = advance(beta_index, cfg.smem_beta_stages) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @cute.jit -def _mma0_warp( +def gate_beta_warp( cfg, total_tiles, bidx, num_ctas, cu_seqlens, mWorkItems, - tmem_hold, - sQ, - sK, + tidx, + mGate, + mBeta, + sCumsumlog, + sCumprod, + sBeta, sSched, bars, ): - """CG0 MMA issuer role (warp 8): persistent scheduler loop + per-pair - KK0/KK1/QK0/QK1 issue into CG0's private accumulator ring; owns the TMEM - lifecycle (alloc up front, dealloc once CG1 signals mb_tmem_done).""" + """Gate/Beta producer (warp 8): persistent scheduler loop + the + cumsum/cumprod/Beta chunk loads.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) - cg0_acc_index = PipelineState.start(phase=1) - k_index = PipelineState.start(phase=0) - q_index = PipelineState.start(phase=0) - - nvvm.tcgen05_alloc(tmem_hold, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) - nvvm.barrier_cta_sync_aligned( - cfg.tmem_alloc_barrier_id, - thread_count=cfg.tmem_alloc_barrier_threads, - ) - tmem_base = tmem_hold.load() - - bpe = cfg.io_dtype.width // 8 - idesc_qk = 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.b_t, - ) - bmm_qk_desc = MmaDesc( - M=cfg.b_t, - N=cfg.b_t, - K=cfg.d_k, - bpe_a=bpe, - bpe_b=bpe, - tile_k_hw=16, - btranspose=False, - cta_group=1, - idesc=idesc_qk, - kind=nvvm.Tcgen05MMAKind.F16, - ) - tmem_cg0_acc_col = tmem_base + cfg.tmem_cg0_acc_offset - ACC_STAGE_COLS = cfg.b_t + gate_index = PipelineState.start(phase=1) + beta_index = PipelineState.start(phase=1) + lidx = tidx % cfg.threads_per_warp 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, cu_seqlens, mWorkItems - ) - n_pairs = (wend - cstart + 1) // 2 - for _pair in cutlass.range(n_pairs): # noqa: B007 - # ---- GEMM 1 (chunk 0 of pair): kk0 ------------------------- - kk0_acc_idx = cg0_acc_index.idx - bars.mb_cg0_acc_done[kk0_acc_idx].wait(cg0_acc_index.phase) - cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - k0_idx = k_index.idx - bars.mb_k_ready[k0_idx].wait(k_index.phase) - k_index = advance(k_index, cfg.smem_k_stages) - - desc_k0 = sK[k0_idx].desc() - mma_ss( - bmm_qk_desc, - desc_k0, - desc_k0, - nvvm.make_tmem_ptr(tmem_cg0_acc_col + kk0_acc_idx * ACC_STAGE_COLS, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_cg0_acc_ready[kk0_acc_idx].arrive(cta_group=1) - - # ---- GEMM 1 (chunk 1 of pair): kk1 ------------------------- - kk1_acc_idx = cg0_acc_index.idx - bars.mb_cg0_acc_done[kk1_acc_idx].wait(cg0_acc_index.phase) - cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - k1_idx = k_index.idx - bars.mb_k_ready[k1_idx].wait(k_index.phase) - k_index = advance(k_index, cfg.smem_k_stages) - - desc_k1 = sK[k1_idx].desc() - mma_ss( - bmm_qk_desc, - desc_k1, - desc_k1, - nvvm.make_tmem_ptr(tmem_cg0_acc_col + kk1_acc_idx * ACC_STAGE_COLS, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_cg0_acc_ready[kk1_acc_idx].arrive(cta_group=1) - - # ---- GEMM 2 (chunk 0 of pair): qk0 ------------------------- - q0_idx = q_index.idx - bars.mb_q_ready[q0_idx].wait(q_index.phase) - q_index = advance(q_index, cfg.smem_q_stages) - qk0_acc_idx = cg0_acc_index.idx - bars.mb_cg0_acc_done[qk0_acc_idx].wait(cg0_acc_index.phase) - cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - - desc_q0 = sQ[q0_idx].desc() - mma_ss( - bmm_qk_desc, - desc_q0, - desc_k0, - nvvm.make_tmem_ptr(tmem_cg0_acc_col + qk0_acc_idx * ACC_STAGE_COLS, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_cg0_acc_ready[qk0_acc_idx].arrive(cta_group=1) - - # ---- GEMM 2 (chunk 1 of pair): qk1 ------------------------- - q1_idx = q_index.idx - bars.mb_q_ready[q1_idx].wait(q_index.phase) - q_index = advance(q_index, cfg.smem_q_stages) - qk1_acc_idx = cg0_acc_index.idx - bars.mb_cg0_acc_done[qk1_acc_idx].wait(cg0_acc_index.phase) - cg0_acc_index = advance(cg0_acc_index, cfg.tmem_cg0_acc_stages) - - desc_q1 = sQ[q1_idx].desc() - mma_ss( - bmm_qk_desc, - desc_q1, - desc_k1, - nvvm.make_tmem_ptr(tmem_cg0_acc_col + qk1_acc_idx * ACC_STAGE_COLS, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_cg0_acc_ready[qk1_acc_idx].arrive(cta_group=1) + 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 n_local > 0: + for local_idx in cutlass.range(n_padded): + # ---- Gate load: GMEM -> SMEM (OOB neutral) ----------------------- + chunk_idx = cstart + local_idx + n_cols = cfg.b_t // cfg.threads_per_warp + chunk_offset = batch_start + chunk_idx * cfg.b_t + gGateSeq = mGate[None, head_idx] + gBeta = cute.domain_offset((chunk_offset,), mBeta[None, head_idx]) - # this issuer's Q/K releases (the CG1 issuer commits its own) - if nvvm.elect_sync(): - bars.mb_q_done[q0_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_q_done[q1_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_k_done[k0_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_k_done[k1_idx].arrive(cta_group=1) + gate_idx = gate_index.idx + gate_phase = gate_index.phase + gate_index = advance(gate_index, cfg.smem_gate_stages) - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + pos_valid = [None] * n_cols + gate_vals = [cutlass.Float32(0.0)] * n_cols + oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) + for col in cutlass.range_constexpr(n_cols): + tok = chunk_offset + lidx + col * cfg.threads_per_warp + pos_valid[col] = cute.elem_less(tok, batch_end) + tok_clamped = min(tok, batch_end - 1) + gate_vals[col] = gGateSeq[tok_clamped] if pos_valid[col] else oob_neutral + + if cutlass.const_expr(cfg.log_gate): + for col in cutlass.range_constexpr(n_cols): + gate_vals[col] = gate_vals[col] * cutlass.Float32(RCP_LN2) + else: + for col in cutlass.range_constexpr(n_cols): + gate_vals[col] = cute.math.log2(gate_vals[col] + 1e-10, fastmath=True) + for offset in [1, 2, 4, 8, 16]: + for col in cutlass.range_constexpr(n_cols): + n = nvvm.shfl_sync(0xFFFFFFFF, gate_vals[col], offset, 0, kind=nvvm.Shfl.UP) + if lidx >= offset: + gate_vals[col] = gate_vals[col] + n + for col in cutlass.range_constexpr(1, n_cols): + last_v = nvvm.shfl_sync( + 0xFFFFFFFF, + gate_vals[col - 1], + cfg.threads_per_warp - 1, + cfg.threads_per_warp - 1, + kind=nvvm.Shfl.IDX, + ) + gate_vals[col] += last_v + + bars.mb_gate_done[gate_idx].wait(gate_phase) + for col in cutlass.range_constexpr(n_cols): + pos = lidx + col * cfg.threads_per_warp + sCumsumlog[pos, 0, gate_idx] = gate_vals[col] + sCumprod[pos, 0, gate_idx] = cute.math.exp2(gate_vals[col], fastmath=True) + bars.mb_gate_ready[gate_idx].arrive() + + # ---- Beta load: GMEM -> SMEM (per-element cp.async) -------------------------- + beta_idx = beta_index.idx + bars.mb_beta_done[beta_idx].wait(beta_index.phase) + beta_index = advance(beta_index, cfg.smem_beta_stages) + for col in cutlass.range_constexpr(n_cols): + pos = lidx + col * cfg.threads_per_warp + src = gBeta.iterator + gBeta.layout((pos,)) + dst = sBeta.iterator + sBeta.layout((pos, 0, beta_idx)) + cp_size = cutlass.Int32(4) * cutlass.Int32(pos_valid[col]) + nvvm.cp_async_shared_global(dst, src, 4, nvvm.LoadCacheModifier.CA, cp_size=cp_size) + nvvm.cp_async_mbarrier_arrive(bars.mb_beta_ready[beta_idx].smem_ptr, noinc=True) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) - bars.mb_tmem_done[0].wait(0) - nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) - nvvm.tcgen05_dealloc( - nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), - cutlass.Int32(512), - group=nvvm.CTAGroup.CTA_1, - ) + for _ in range(cfg.smem_gate_stages): + bars.mb_gate_done[gate_index.idx].wait(gate_index.phase) + gate_index = advance(gate_index, cfg.smem_gate_stages) + for _ in range(cfg.smem_beta_stages): + bars.mb_beta_done[beta_index.idx].wait(beta_index.phase) + beta_index = advance(beta_index, cfg.smem_beta_stages) @cute.jit -def _mma1_warp( +def mma_warp( cfg, total_tiles, bidx, num_ctas, cu_seqlens, mWorkItems, - tmem_hold, - sQ, - sK, - sK_trans, - sAinv, - sQk, + tmem_base_slot, + sKQ, + sKQ_trans, + sTinv, + sA, sSched, bars, ): - """CG1 MMA issuer role (warp 10): persistent scheduler loop + per-chunk - issue of the five state/output GEMMs (KS/QS/NV/QKV/KV) in dependency - order.""" + """MMA issuer role (warp 10): persistent scheduler loop issuing every + tcgen05 GEMM.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) o_acc_index = PipelineState.start(phase=1) o_state_scale_index = PipelineState.start(phase=0) kv_acc_index = PipelineState.start(phase=1) - k_index = PipelineState.start(phase=0) - q_index = PipelineState.start(phase=0) - ainv_index = PipelineState.start(phase=0) - qk_index = PipelineState.start(phase=0) + kq_index = PipelineState.start(phase=0) + cg0_acc_index = PipelineState.start(phase=1) + kq_fused_index = PipelineState.start(phase=0) + tinv_index = PipelineState.start(phase=0) + a_index = PipelineState.start(phase=0) state_inp_index = PipelineState.start(phase=0) - vks_inp_rdy = PipelineState.start(phase=0) - nv_inp_rdy = PipelineState.start(phase=0) - decay_v_inp_rdy = PipelineState.start(phase=0) + y_inp_ready = PipelineState.start(phase=0) + u_inp_ready = PipelineState.start(phase=0) + decay_u_inp_ready = PipelineState.start(phase=0) + nvvm.tcgen05_alloc(tmem_base_slot, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) nvvm.barrier_cta_sync_aligned( cfg.tmem_alloc_barrier_id, thread_count=cfg.tmem_alloc_barrier_threads, ) - tmem_base = tmem_hold.load() - # ---- chunk-invariant GEMM descriptors ------------------------------ + # ---- chunk-invariant GEMM descriptors ---------------------------------------- + idesc_qk = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.b_t, + m_dim=2 * cfg.b_t, + ) + bmm_qk_desc = MmaDesc( + M=2 * cfg.b_t, + N=cfg.b_t, + K=cfg.d_k, + bpe_a=cfg.io_dtype.width // 8, + bpe_b=cfg.io_dtype.width // 8, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_qk, + kind=nvvm.Tcgen05MMAKind.F16, + ) bpe = cfg.io_dtype.width // 8 - idesc_qs = nvvm.Tcgen05InstrDesc.build( + idesc_q_state = 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_qs_desc = MmaDesc( + bmm_q_state_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.d_k, @@ -975,7 +796,7 @@ def _mma1_warp( btranspose=False, atranspose=False, cta_group=1, - idesc=idesc_qs, + idesc=idesc_q_state, kind=nvvm.Tcgen05MMAKind.F16, ) idesc_qkv_ts = nvvm.Tcgen05InstrDesc.build( @@ -1019,241 +840,201 @@ def _mma1_warp( idesc=idesc_kv, kind=nvvm.Tcgen05MMAKind.F16, ) + KQ_SEG = (2 * cfg.b_t * 64 * bpe) >> 4 + KQ_BOX = (cfg.b_t * 64 * bpe) >> 4 + KQ_HALF_K = (cfg.d_k // 16) // 2 + KQ_A_HALF = KQ_HALF_K * bmm_q_state_desc.tmem_advance_A - tmem_state_col = tmem_base + cfg.tmem_state_offset - tmem_q_state_col = tmem_base + cfg.tmem_q_state_offset - tmem_state_inp_col = tmem_base + cfg.tmem_state_inp_offset - tmem_inp_col = tmem_base + cfg.tmem_inp_offset ACC_STAGE_COLS = cfg.b_t KV_ACC_STAGE_COLS = cfg.d_v STATE_INP_STAGE_COLS = cfg.d_k // 2 INP_SLOT_COLS = cfg.b_t // 2 - tmem_ks_col = tmem_base + cfg.tmem_cg1_acc_offset - tmem_nv_col = tmem_ks_col - tmem_vks_col = tmem_inp_col - # NV overwrites the VKS slot once GEMM 5 has consumed it - tmem_nv_inp_col = tmem_inp_col - tmem_decay_v_col = tmem_inp_col + INP_SLOT_COLS + + tmem_base = tmem_base_slot.load() + tmem_cg0_acc_col_f = tmem_base + cfg.tmem_cg0_acc_offset + tmem_state_col = tmem_base + cfg.tmem_state_acc_offset + tmem_q_state_col = tmem_base + cfg.tmem_q_state_acc_offset + tmem_state_inp_col = tmem_base + cfg.tmem_state_inp_offset + tmem_inp_col = tmem_base + cfg.tmem_y_decay_u_inp_offset + y_inp_ptr = nvvm.make_tmem_ptr(tmem_inp_col, cutlass.Int8) + u_inp_ptr = nvvm.make_tmem_ptr(tmem_inp_col, cutlass.Int8) + 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, cu_seqlens, mWorkItems - ) + 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 - first_loop = 0 - if cutlass.const_expr(not cfg.use_initial_state): - # ---- peeled first chunk: S_prev = 0, GEMMs 3/4 skipped ----- - if n_local > 0: - k_idx = k_index.idx - bars.mb_k_ready[k_idx].wait(k_index.phase) - k_index = advance(k_index, cfg.smem_k_stages) - q_idx = q_index.idx - bars.mb_q_ready[q_idx].wait(q_index.phase) - q_index = advance(q_index, cfg.smem_q_stages) - if nvvm.elect_sync(): - bars.mb_q_done[q_idx].arrive(cta_group=1) - - # ---- GEMM 5: new_v --------------------------------- - bars.mb_vks_inp_ready[0].wait(vks_inp_rdy.phase) - vks_inp_rdy = advance(vks_inp_rdy, 1) - ainv_idx = ainv_index.idx - bars.mb_ainv_ready[ainv_idx].wait(ainv_index.phase) - ainv_index = advance(ainv_index, cfg.smem_ainv_stages) - - desc_ainv = sAinv[ainv_idx].desc() - vks_a_ptr = nvvm.make_tmem_ptr(tmem_vks_col, cutlass.Int8) - mma_ts( - bmm_qkv_ts_desc, - vks_a_ptr, - desc_ainv, - nvvm.make_tmem_ptr(tmem_nv_col, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_nv_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_ainv_done[ainv_idx].arrive(cta_group=1) - - # ---- GEMM 6: qkv ----------------------------------- - bars.mb_nv_inp_ready[0].wait(nv_inp_rdy.phase) - nv_inp_rdy = advance(nv_inp_rdy, 1) - qk_idx = qk_index.idx - bars.mb_qk_ready[qk_idx].wait(qk_index.phase) - qk_index = advance(qk_index, cfg.smem_qk_stages) - qs_idx = o_acc_index.idx - bars.mb_o_acc_done[qs_idx].wait(o_acc_index.phase) - o_acc_index = advance(o_acc_index, cfg.tmem_q_state_acc_stages) - - qkv_a_ptr = nvvm.make_tmem_ptr(tmem_nv_inp_col, cutlass.Int8) - desc_nv = sQk[qk_idx].desc() - mma_ts( - bmm_qkv_ts_desc, - qkv_a_ptr, - desc_nv, - nvvm.make_tmem_ptr(tmem_q_state_col + qs_idx * ACC_STAGE_COLS, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_qk_done[qk_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_o_state_scale_acc_ready[qs_idx].arrive(cta_group=1) - - # ---- GEMM 7: kv_update --------------------------------------- - bars.mb_decay_v_inp_ready[0].wait(decay_v_inp_rdy.phase) - decay_v_inp_rdy = advance(decay_v_inp_rdy, 1) - kv_acc_idx = kv_acc_index.idx - bars.mb_kv_acc_scale_done[kv_acc_idx].wait(kv_acc_index.phase) - kv_acc_index = advance(kv_acc_index, cfg.tmem_kv_acc_stages) - - delta_a_ptr = nvvm.make_tmem_ptr(tmem_decay_v_col, cutlass.Int8) - desc_kt = sK_trans[k_idx].desc() - mma_ts( - bmm_kv_desc, - delta_a_ptr, - desc_kt, - nvvm.make_tmem_ptr(tmem_state_col + kv_acc_idx * KV_ACC_STAGE_COLS, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_kv_acc_ready[kv_acc_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_k_done[k_idx].arrive(cta_group=1) - - first_loop = 1 - - for local_idx in cutlass.range(first_loop, n_padded): # noqa: B007 + # ---- fused KK^T/QK^T pair 0: both members issued ahead of the loop ------- + if n_pairs > 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) + kqf_idx = kq_fused_index.idx + 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) + if elect_one: + bars.mb_cg0_acc_ready[f0_acc_idx].arrive(cta_group=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) + kqf_idx = kq_fused_index.idx + 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() + 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) + 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 if cutlass.const_expr(cfg.use_initial_state): if local_idx == 0: - kv_acc_index = advance(kv_acc_index, cfg.tmem_kv_acc_stages) - - k_idx = k_index.idx - q_idx = q_index.idx - s_idx = state_inp_index.idx + if elect_one: + bars.mb_state_acc_ready[kv_acc_index.idx].arrive(cta_group=1) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) + have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_idx > 0 + + kq_idx = kq_index.idx + member = local_idx & 1 + state_inp_idx = state_inp_index.idx q_state_acc_idx = o_acc_index.idx - ainv_idx = ainv_index.idx - qk_idx = qk_index.idx - qs2_idx = o_state_scale_index.idx + tinv_idx = tinv_index.idx + a_idx = a_index.idx + o_scale_idx = o_state_scale_index.idx kv_acc_idx = kv_acc_index.idx - desc_k = sK[k_idx].desc() - desc_q = sQ[q_idx].desc() - desc_ainv = sAinv[ainv_idx].desc() - desc_nv = sQk[qk_idx].desc() - desc_kt = sK_trans[k_idx].desc() - state_a_ptr = nvvm.make_tmem_ptr(tmem_state_inp_col + s_idx * STATE_INP_STAGE_COLS, cutlass.Int8) - vks_a_ptr = nvvm.make_tmem_ptr(tmem_vks_col, cutlass.Int8) - qkv_a_ptr = nvvm.make_tmem_ptr(tmem_nv_inp_col, cutlass.Int8) - delta_a_ptr = nvvm.make_tmem_ptr(tmem_decay_v_col, cutlass.Int8) - ks_acc_ptr = nvvm.make_tmem_ptr(tmem_ks_col, cutlass.Float32) - qs_acc_ptr = nvvm.make_tmem_ptr(tmem_q_state_col + q_state_acc_idx * ACC_STAGE_COLS, cutlass.Float32) - nv_acc_ptr = nvvm.make_tmem_ptr(tmem_nv_col, cutlass.Float32) - qkv_acc_ptr = nvvm.make_tmem_ptr(tmem_q_state_col + qs2_idx * ACC_STAGE_COLS, cutlass.Float32) - kv_acc_ptr = nvvm.make_tmem_ptr(tmem_state_col + kv_acc_idx * KV_ACC_STAGE_COLS, cutlass.Float32) - - bars.mb_k_ready[k_idx].wait(k_index.phase) - k_index = advance(k_index, cfg.smem_k_stages) - bars.mb_q_ready[q_idx].wait(q_index.phase) - q_index = advance(q_index, cfg.smem_q_stages) - - # ---- GEMM 3: k*state --------------------------------------- - bars.mb_state_inp_ready[s_idx].wait(state_inp_index.phase) - state_inp_index = advance(state_inp_index, cfg.tmem_state_inp_stages) - - mma_ts( - bmm_qs_desc, - state_a_ptr, - desc_k, - ks_acc_ptr, - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_ks_ready[0].arrive(cta_group=1) - - # ---- GEMM 4: q*state --------------------------------------- - bars.mb_o_acc_done[q_state_acc_idx].wait(o_acc_index.phase) + kq_member_off = member * KQ_BOX + desc_k = sKQ[kq_idx].desc() + kq_member_off + desc_q = sKQ[kq_idx].desc() + (KQ_BOX - kq_member_off) + desc_tinv = sTinv[tinv_idx].desc() + desc_a = sA[a_idx].desc() + desc_kt = sKQ_trans[kq_idx].desc() + kq_member_off + state_a_ptr = nvvm.make_tmem_ptr(tmem_state_inp_col + state_inp_idx * STATE_INP_STAGE_COLS, cutlass.Int8) + q_state_acc_ptr = nvvm.make_tmem_ptr(tmem_q_state_col + q_state_acc_idx * ACC_STAGE_COLS, cutlass.Float32) + qkv_acc_ptr = nvvm.make_tmem_ptr(tmem_q_state_col + o_scale_idx * ACC_STAGE_COLS, cutlass.Float32) + state_acc_ptr = nvvm.make_tmem_ptr(tmem_state_col + kv_acc_idx * KV_ACC_STAGE_COLS, cutlass.Float32) + + kq_index = advance(kq_index, cfg.smem_kq_stages) + + # ---- QK/KK lookahead (member 1) = [Q;K](S) @ K^T --------------------- + if member == 1: + if (local_idx >> 1) + 1 < n_pairs: + 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) + kqf_idx = kq_fused_index.idx + 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() + 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) + if elect_one: + bars.mb_cg0_acc_ready[pend_acc_idx].arrive(cta_group=1) + + # ---- K*state^T (GEMM 3) = state^T(T) @ K^T ------------------------------------ + if have_state: + bars.mb_state_inp_ready[state_inp_idx].wait(state_inp_index.phase) + state_inp_index = advance(state_inp_index, cfg.tmem_state_inp_stages) + + for k in cutlass.range_constexpr(KQ_HALF_K): + mma_ts_step(bmm_q_state_desc, state_a_ptr, desc_k, k_state_acc_ptr, k, cutlass.Boolean(k > 0)) + for k in cutlass.range_constexpr(KQ_HALF_K): + mma_ts_step(bmm_q_state_desc, state_a_ptr.subview(KQ_A_HALF), desc_k + KQ_SEG, k_state_acc_ptr, k, cutlass.Boolean(True)) + if elect_one: + bars.mb_k_state_acc_ready[0].arrive(cta_group=1) + + # ---- Q*state^T (GEMM 4) = state^T(T) @ Q^T ------------------------------------ o_acc_index = advance(o_acc_index, cfg.tmem_q_state_acc_stages) - mma_ts( - bmm_qs_desc, - state_a_ptr, - desc_q, - qs_acc_ptr, - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_o_acc_ready[q_state_acc_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_q_done[q_idx].arrive(cta_group=1) - - # ---- GEMM 5: new_v ----------------------------------------- - # vks_inp_ready also proves CG1 read KS out of the shared column - bars.mb_vks_inp_ready[0].wait(vks_inp_rdy.phase) - vks_inp_rdy = advance(vks_inp_rdy, 1) - bars.mb_ainv_ready[ainv_idx].wait(ainv_index.phase) - ainv_index = advance(ainv_index, cfg.smem_ainv_stages) - mma_ts( - bmm_qkv_ts_desc, - vks_a_ptr, - desc_ainv, - nv_acc_ptr, - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_nv_ready[0].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_ainv_done[ainv_idx].arrive(cta_group=1) - - # ---- GEMM 6: qkv ------------------------------------------- - bars.mb_nv_inp_ready[0].wait(nv_inp_rdy.phase) - nv_inp_rdy = advance(nv_inp_rdy, 1) - bars.mb_qk_ready[qk_idx].wait(qk_index.phase) - qk_index = advance(qk_index, cfg.smem_qk_stages) - bars.mb_o_state_scale_acc_done[qs2_idx].wait(o_state_scale_index.phase) - o_state_scale_index = advance(o_state_scale_index, cfg.tmem_q_state_acc_stages) - mma_ts( - bmm_qkv_ts_desc, - qkv_a_ptr, - desc_nv, - qkv_acc_ptr, - accumulate=True, - ) - if nvvm.elect_sync(): - bars.mb_qk_done[qk_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_o_state_scale_acc_ready[qs2_idx].arrive(cta_group=1) - - # ---- GEMM 7: kv_update ------------------------------------------- - bars.mb_decay_v_inp_ready[0].wait(decay_v_inp_rdy.phase) - decay_v_inp_rdy = advance(decay_v_inp_rdy, 1) - bars.mb_kv_acc_scale_done[kv_acc_idx].wait(kv_acc_index.phase) - kv_acc_index = advance(kv_acc_index, cfg.tmem_kv_acc_stages) - mma_ts( - bmm_kv_desc, - delta_a_ptr, - desc_kt, - kv_acc_ptr, - accumulate=True, - ) - if nvvm.elect_sync(): - bars.mb_kv_acc_ready[kv_acc_idx].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_k_done[k_idx].arrive(cta_group=1) + if have_state: + for k in cutlass.range_constexpr(KQ_HALF_K): + mma_ts_step(bmm_q_state_desc, state_a_ptr, desc_q, q_state_acc_ptr, k, cutlass.Boolean(k > 0)) + for k in cutlass.range_constexpr(KQ_HALF_K): + mma_ts_step(bmm_q_state_desc, state_a_ptr.subview(KQ_A_HALF), desc_q + KQ_SEG, q_state_acc_ptr, k, cutlass.Boolean(True)) + if elect_one: + bars.mb_o_acc_ready[q_state_acc_idx].arrive(cta_group=1) + + # ---- U^T (GEMM 5) = Y^T(T) @ T_inv^T --------------------------------------- + bars.mb_t_inv_ready[tinv_idx].wait(tinv_index.phase) + tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) + bars.mb_y_inp_ready[0].wait(y_inp_ready.phase) + y_inp_ready = advance(y_inp_ready, 1) + for k in cutlass.range_constexpr(cfg.b_t // 16): + mma_ts_step(bmm_qkv_ts_desc, y_inp_ptr, desc_tinv, u_acc_ptr, k, cutlass.Boolean(k > 0)) + if elect_one: + bars.mb_u_acc_ready[0].arrive(cta_group=1) + bars.mb_t_inv_done[tinv_idx].arrive(cta_group=1) + + # ---- KK/QK lookahead (member 0) = [K;Q](S) @ K^T --------------------- + if member == 0: + if (local_idx >> 1) + 1 < n_pairs: + 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) + kqf_idx = kq_fused_index.idx + 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) + if elect_one: + bars.mb_cg0_acc_ready[f0_acc_idx].arrive(cta_group=1) + + # ---- O^T (GEMM 6) += U input^T(T) @ A^T ------------------------------ + bars.mb_a_ready[a_idx].wait(a_index.phase) + a_index = advance(a_index, cfg.smem_a_stages) + if have_state: + bars.mb_o_state_scale_acc_done[o_scale_idx].wait(o_state_scale_index.phase) + o_state_scale_index = advance(o_state_scale_index, cfg.tmem_q_state_acc_stages) + bars.mb_u_inp_ready[0].wait(u_inp_ready.phase) + u_inp_ready = advance(u_inp_ready, 1) + for k in cutlass.range_constexpr(cfg.b_t // 16): + mma_ts_step(bmm_qkv_ts_desc, u_inp_ptr, desc_a, qkv_acc_ptr, k, cutlass.Boolean(True) if cutlass.const_expr(k > 0) else have_state) + if elect_one: + bars.mb_a_done[a_idx].arrive(cta_group=1) + bars.mb_o_final_acc_ready[o_scale_idx].arrive(cta_group=1) + + # ---- state^T (GEMM 7) += decayed U^T(T) @ K ------------------------------ + bars.mb_decay_u_inp_ready[0].wait(decay_u_inp_ready.phase) + decay_u_inp_ready = advance(decay_u_inp_ready, 1) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) + for k in cutlass.range_constexpr(cfg.b_t // 16): + mma_ts_step(bmm_kv_desc, decay_u_inp_ptr, desc_kt, state_acc_ptr, k, cutlass.Boolean(True) if cutlass.const_expr(k > 0) else have_state) + if elect_one: + bars.mb_state_acc_ready[kv_acc_idx].arrive(cta_group=1) + bars.mb_kq_done[kq_idx].arrive(cta_group=1) + + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + bars.mb_tmem_done[0].wait(0) + nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) + nvvm.tcgen05_dealloc( + nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), + cutlass.Int32(512), + group=nvvm.CTAGroup.CTA_1, + ) @cute.jit -def _tmaldg_warp( +def tmaldg_warp( cfg, total_tiles, bidx, num_ctas, cu_seqlens, mWorkItems, - sQ_raw, - sK_raw, + sKQ_raw, sV_raw, desc_q_base, desc_k_base, @@ -1263,44 +1044,29 @@ def _tmaldg_warp( bars, ): """TMA-LDG warp role (warp 9): persistent scheduler loop + per-chunk - Q/K/V G->S TMA loads. - - Each load goes through a per-(b,h) runtime descriptor (built once on - host) whose GLOBAL_ADDRESS already folds the sequence start and head, so - the only load coordinate is token = ``chunk_idx*BT``.""" + Q/K/V G->S TMA loads.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) - q_index = PipelineState.start(phase=1) - k_index = PipelineState.start(phase=1) + kq_index = PipelineState.start(phase=1) v_index = PipelineState.start(phase=1) sched_state = PipelineState.start(phase=1) tile_idx = cutlass.Int32(bidx) bpe = cfg.io_dtype.width // 8 - granu = 128 // bpe + elems_per_128b = 128 // bpe bt = cfg.b_t - q_stage_elems = cfg.q_cosize // cfg.smem_q_stages - k_stage_elems = cfg.k_cosize // cfg.smem_k_stages - sQ_tma = SmemTile( - base=sQ_raw, - elems_per_stage=q_stage_elems, - stages=cfg.smem_q_stages, + kq_stage_elems = cfg.kq_cosize // cfg.smem_kq_stages + kq_box_elems = kq_stage_elems // 4 + sKQ_lo_tma = SmemTile( + base=sKQ_raw, + elems_per_stage=kq_stage_elems, + stages=cfg.smem_kq_stages, leading_byte_offset=0, stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, - tma_subtile_stride_elems=bt * granu, - ) - sK_tma = SmemTile( - base=sK_raw, - elems_per_stage=k_stage_elems, - stages=cfg.smem_k_stages, - leading_byte_offset=0, - stride_byte_offset=0, - layout=0, - tma_loads_per_tile=2, - tma_granu_elems=granu, - tma_subtile_stride_elems=bt * granu, + tma_granu_elems=elems_per_128b, + tma_subtile_stride_elems=2 * bt * elems_per_128b, ) sV_tma = SmemTile( base=sV_raw, @@ -1310,79 +1076,91 @@ def _tmaldg_warp( stride_byte_offset=0, layout=0, tma_loads_per_tile=2, - tma_granu_elems=granu, + tma_granu_elems=elems_per_128b, tma_subtile_stride_elems=4096, ) heads_out = cutlass.Int32(cfg.n_heads_out) desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) 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, cu_seqlens, mWorkItems - ) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) - slot = (batch_idx * heads_out + head_idx) * desc_qwords + head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) + head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) + head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) + slot = batch_idx * desc_qwords desc_q_slot = (desc_q_base + slot).tospace(cutlass.AddressSpace.generic) desc_k_slot = (desc_k_base + slot).tospace(cutlass.AddressSpace.generic) desc_v_slot = (desc_v_base + slot).tospace(cutlass.AddressSpace.generic) - if nvvm.elect_sync(): + if elect_one: tma_tensormap_acquire(desc_q_slot) tma_tensormap_acquire(desc_k_slot) tma_tensormap_acquire(desc_v_slot) - # padded loads zero-fill (descriptor token extent is capped) wend_padded = cstart + ((wend - cstart + 1) // 2) * 2 - for chunk_idx in cutlass.range(cstart, wend_padded): - tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) - - # ---------------------------------------------------------- - # K (B operand of GEMM-kk / GEMM-qk, double-buffered) - # ---------------------------------------------------------- - k_idx = k_index.idx - bars.mb_k_done[k_idx].wait(k_index.phase) - k_index = advance(k_index, cfg.smem_k_stages) - if nvvm.elect_sync(): - bars.mb_k_ready[k_idx].arrive(n_bytes=cfg.tma_k_bytes) - k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), tok_coord) - tma_load_tile(sK_tma[k_idx], k_slice, bars.mb_k_ready[k_idx].smem_ptr, acquire=False) - - # ---------------------------------------------------------- - # Q (A operand of GEMM-qk, single-buffered) - # ---------------------------------------------------------- - q_idx = q_index.idx - bars.mb_q_done[q_idx].wait(q_index.phase) - q_index = advance(q_index, cfg.smem_q_stages) - if nvvm.elect_sync(): - bars.mb_q_ready[q_idx].arrive(n_bytes=cfg.tma_q_bytes) - q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), tok_coord) - tma_load_tile(sQ_tma[q_idx], q_slice, bars.mb_q_ready[q_idx].smem_ptr, acquire=False) + if wend_padded > 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) + if elect_one: + bars.mb_kq_ready[kq_idx].arrive(n_bytes=cfg.tma_kq_bytes) + tok_coord = cstart * cutlass.Int32(cfg.b_t) + k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, tok_coord) + q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, 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) + 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): + tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) + + # ---- K + Q interleaved ------------------------------------------- + kq_idx = kq_index.idx + bars.mb_kq_done[kq_idx].wait(kq_index.phase) + kq_index = advance(kq_index, cfg.smem_kq_stages) + if elect_one: + bars.mb_kq_ready[kq_idx].arrive(n_bytes=cfg.tma_kq_bytes) + member = (chunk_idx - cstart) & 1 + k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, tok_coord) + q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, tok_coord) + kq_tile = sKQ_lo_tma[kq_idx] + if member == 0: + 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) + else: + tma_load_tile(kq_tile, q_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) + tma_load_tile(kq_tile.shifted(kq_box_elems), k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) + + # ---- V load ------------------------------------------------------ + v_idx = v_index.idx + bars.mb_v_done[v_idx].wait(v_index.phase) + 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 = (chunk_idx - 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) - # ---------------------------------------------------------- - # V (A operand of GEMM-new_v; transposed [DV, T] descriptor) - # ---------------------------------------------------------- v_idx = v_index.idx bars.mb_v_done[v_idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) - if nvvm.elect_sync(): + if elect_one: bars.mb_v_ready[v_idx].arrive(n_bytes=cfg.tma_v_bytes) - v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), tok_coord) + v_tok = (wend_padded - 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) - tile_idx, sched_state = _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) - for _ in range(cfg.smem_q_stages): - bars.mb_q_done[q_index.idx].wait(q_index.phase) - q_index = advance(q_index, cfg.smem_q_stages) - for _ in range(cfg.smem_k_stages): - bars.mb_k_done[k_index.idx].wait(k_index.phase) - k_index = advance(k_index, cfg.smem_k_stages) + for _ in range(cfg.smem_kq_stages): + bars.mb_kq_done[kq_index.idx].wait(kq_index.phase) + kq_index = advance(kq_index, cfg.smem_kq_stages) for _ in range(cfg.smem_v_stages): bars.mb_v_done[v_index.idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) @cute.jit -def _compute0_warp( +def compute0_warp_group( cfg, total_tiles, bidx, @@ -1390,45 +1168,46 @@ def _compute0_warp( cu_seqlens, mWorkItems, tidx, - tmem_hold, + tmem_base_slot, scale, sCumsumlog, sBeta, - sAinv, - sQk, - sH_raw, + sTinv, + sA, + sCheckpoint_raw, checkpoint_every_n_tokens, sSched, bars, ): - """Compute warp-group 0 role (warps 0-3): persistent scheduler loop + - per-PAIR T-pairwise x2, kk_epi x2, pair inverse (warps 0-1 invert chunk - 0's matrix while warps 2-3 invert chunk 1's), qk_epi x2, and (enable_h) - the H checkpoint readouts state TMEM -> sH woven around the QK epilogues.""" + """Compute warp-group 0 role (warps 0-3): persistent scheduler loop + computing each chunk pair's T_inv and A epilogues.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) gate_index = PipelineState.start(phase=0) beta_index = PipelineState.start(phase=0) - cg0_acc_rdy = PipelineState.start(phase=0) - ainv_index = PipelineState.start(phase=1) - qk_index = PipelineState.start(phase=1) + cg0_acc_ready = PipelineState.start(phase=0) + tinv_index = PipelineState.start(phase=1) + a_index = PipelineState.start(phase=1) nvvm.barrier_cta_sync_aligned( cfg.tmem_alloc_barrier_id, thread_count=cfg.tmem_alloc_barrier_threads, ) - tmem_base = tmem_hold.load() + tmem_base = tmem_base_slot.load() num_threads_cg0 = cfg.threads_per_warp * len(cfg.compute_group_0_warp_ids) cg0_tidx = tidx % num_threads_cg0 warp_id = cg0_tidx // cfg.threads_per_warp lane_id = cg0_tidx % cfg.threads_per_warp inverse_local_warp = warp_id % 2 + pair_half = warp_id // 2 + half_row_base = inverse_local_warp * 32 bpe = cfg.io_dtype.width // 8 num_vals = 32 FRAG_COLS = 16 ACC_N_FRAGS = cfg.b_t // FRAG_COLS store_row = warp_id * 16 + lane_id % 16 + store_row_frag = lane_id % 16 store_col = (lane_id // 16) * 8 tmem_warp_row = warp_id * cfg.threads_per_warp tmem_cg0_acc_col = tmem_base + cfg.tmem_cg0_acc_offset @@ -1440,40 +1219,55 @@ def _compute0_warp( 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, cu_seqlens, mWorkItems - ) + 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): - chunk0 = cstart + pair_i * 2 - chunk1 = chunk0 + 1 - # ---- Step 1: T-pairwise, both chunks in one traversal -------- + # ---- 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) - - row_cs0_lo = sCumsumlog[crow_lo, 0, gate0_idx] - row_cs0_hi = sCumsumlog[crow_hi, 0, gate0_idx] - row_cs1_lo = sCumsumlog[crow_lo, 0, gate1_idx] - row_cs1_hi = sCumsumlog[crow_hi, 0, gate1_idx] - - gT0 = [] - gT1 = [] - for k in cutlass.range_constexpr(num_vals): - hi_row = cutlass.const_expr(((k // 2) % 2) == 1) - crow = crow_hi if cutlass.const_expr(hi_row) else crow_lo - ccol = (lane_id % 4) * 2 + ((k // 4) * 8 + k % 2) - is_lower = crow >= ccol - cs0 = row_cs0_hi if cutlass.const_expr(hi_row) else row_cs0_lo - cs1 = row_cs1_hi if cutlass.const_expr(hi_row) else row_cs1_lo - gT0.append(cute.math.exp2(cs0 - sCumsumlog[ccol, 0, gate0_idx], fastmath=True) if is_lower else mask_zero) - gT1.append(cute.math.exp2(cs1 - sCumsumlog[ccol, 0, gate1_idx], fastmath=True) if is_lower else mask_zero) + kk_gate_idx = gate1_idx if pair_half == 1 else gate0_idx + a_gate_idx = gate0_idx if pair_half == 1 else gate1_idx + + row_u0_lo = half_row_base + lane_id // 4 + row_u0_hi = row_u0_lo + 8 + row_u1_lo = row_u0_lo + 16 + row_u1_hi = row_u0_lo + 24 + + kk_cumsumlog_rows = [] + a_cumsumlog_rows = [] + for r in (row_u0_lo, row_u0_hi, row_u1_lo, row_u1_hi): + kk_cumsumlog_rows.append(sCumsumlog[r, 0, kk_gate_idx]) + a_cumsumlog_rows.append(sCumsumlog[r, 0, a_gate_idx]) + kk_cumsumlog_cols = [] + a_cumsumlog_cols = [] + for g in cutlass.range_constexpr(8): + for b in cutlass.range_constexpr(2): + ccol = (lane_id % 4) * 2 + g * 8 + b + kk_cumsumlog_cols.append(sCumsumlog[ccol, 0, kk_gate_idx]) + a_cumsumlog_cols.append(sCumsumlog[ccol, 0, a_gate_idx]) + + decay_t_kk = [] + decay_t_a = [] + for u in cutlass.range_constexpr(2): + for k in cutlass.range_constexpr(num_vals): + hi_row = ((k // 2) % 2) == 1 + crow_u0 = row_u0_hi if cutlass.const_expr(hi_row) else row_u0_lo + crow_u1 = row_u1_hi if cutlass.const_expr(hi_row) else row_u1_lo + crow = crow_u1 if cutlass.const_expr(u == 1) else crow_u0 + ccol = (lane_id % 4) * 2 + ((k // 4) * 8 + k % 2) + is_lower = crow >= ccol + kk_row_cumsumlog = kk_cumsumlog_rows[u * 2 + (1 if hi_row else 0)] + a_row_cumsumlog = a_cumsumlog_rows[u * 2 + (1 if hi_row else 0)] + col = (k // 4) * 2 + (k % 2) + 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() @@ -1483,119 +1277,108 @@ def _compute0_warp( beta1_idx = beta_index.idx bars.mb_beta_ready[beta1_idx].wait(beta_index.phase) beta_index = advance(beta_index, cfg.smem_beta_stages) - # beta row scaling only needs the two per-thread row scalars - beta0_lo = sBeta[crow_lo, 0, beta0_idx] - beta0_hi = sBeta[crow_hi, 0, beta0_idx] - beta1_lo = sBeta[crow_lo, 0, beta1_idx] - beta1_hi = sBeta[crow_hi, 0, beta1_idx] - - # ---- Step 2: kk_epi0 + kk_epi1 ------------------------------ - ainv0_idx = ainv_index.idx - bars.mb_ainv_done[ainv0_idx].wait(ainv_index.phase) - ainv_index = advance(ainv_index, cfg.smem_ainv_stages) - kk0_acc_idx = cg0_acc_rdy.idx - bars.mb_cg0_acc_ready[kk0_acc_idx].wait(cg0_acc_rdy.phase) - cg0_acc_rdy = advance(cg0_acc_rdy, cfg.tmem_cg0_acc_stages) - - ainv0_base = sAinv[ainv0_idx].base - kk_vec = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + kk0_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 + 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): + kk_beta.append(sBeta[r, 0, kk_beta_idx]) + + # ---- KK_epi (each warp pair stages its own member) ------------------- + 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) + 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 + + 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) + 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 ) - nvvm.tcgen05_wait("load") - bars.mb_cg0_acc_done[kk0_acc_idx].arrive() - kk_f16 = [] - for k in cutlass.range_constexpr(num_vals // 2): - b0 = beta0_hi if cutlass.const_expr((k % 2) == 1) else beta0_lo - p0, p1 = fmul2(kk_vec[2 * k], kk_vec[2 * k + 1], gT0[2 * k], gT0[2 * k + 1]) - v0, v1 = fmul2(p0, p1, b0, b0) - kk_f16.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) - for c in cutlass.range_constexpr(ACC_N_FRAGS): - nvvm.stmatrix( - cutlass.inttoptr( - ainv0_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [kk_f16[c * 4 + 0], kk_f16[c * 4 + 1], kk_f16[c * 4 + 2], kk_f16[c * 4 + 3]], - nvvm.MMALayout.ROW, - ) - - ainv1_idx = ainv_index.idx - bars.mb_ainv_done[ainv1_idx].wait(ainv_index.phase) - ainv_index = advance(ainv_index, cfg.smem_ainv_stages) - kk1_acc_idx = cg0_acc_rdy.idx - bars.mb_cg0_acc_ready[kk1_acc_idx].wait(cg0_acc_rdy.phase) - cg0_acc_rdy = advance(cg0_acc_rdy, cfg.tmem_cg0_acc_stages) - - ainv1_base = sAinv[ainv1_idx].base - kk_vec = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + kk1_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[kk1_acc_idx].arrive() - kk_f16 = [] - for k in cutlass.range_constexpr(num_vals // 2): - b1 = beta1_hi if cutlass.const_expr((k % 2) == 1) else beta1_lo - p0, p1 = fmul2(kk_vec[2 * k], kk_vec[2 * k + 1], gT1[2 * k], gT1[2 * k + 1]) - v0, v1 = fmul2(p0, p1, b1, b1) - kk_f16.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) - for c in cutlass.range_constexpr(ACC_N_FRAGS): - nvvm.stmatrix( - cutlass.inttoptr( - ainv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [kk_f16[c * 4 + 0], kk_f16[c * 4 + 1], kk_f16[c * 4 + 2], kk_f16[c * 4 + 3]], - nvvm.MMALayout.ROW, - ) + 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 = ainv1_base if warp_id >= 2 else ainv0_base + # ---- pair inverse: warps 0-1 own matrix 0, warps 2-3 matrix 1 -------- + inv_base = tinv1_base if warp_id >= 2 else tinv0_base - # Stage 1: diagonal 8x8 Gauss-Jordan, all four warps + # 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) + 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, ) - # Stage 2: 8x8 -> 16x16 (both matrices per warp) - _blockwise_diagonal_8x8_to_16x16(cfg, ainv0_base, warp_id * 16, lane_id) - _blockwise_diagonal_8x8_to_16x16(cfg, ainv1_base, warp_id * 16, lane_id) + # 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) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - # Stage 3: 16x16 -> 32x32, one tile per warp within the group - _blockwise_diagonal_16x16_to_32x32(cfg, inv_base, inverse_local_warp * 32, lane_id) + # 16x16 -> 32x32, one tile per warp within the group + 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, ) - # Stage 4: 32x32 -> 64x64, two warps per matrix - _blockwise_diagonal_32x32_to_64x64(cfg, inv_base, inverse_local_warp, lane_id) + # 32x32 -> 64x64, two warps per matrix + blockwise_diagonal_32x32_to_64x64(cfg, inv_base, inverse_local_warp, lane_id) nvvm.barrier_cta_sync_aligned( cfg.inverse_barrier_id, thread_count=cfg.inverse_barrier_threads, ) - # ---- beta column-scaling + publish, stage 0 then stage 1 ---- + # ---- Beta column-scaling + publish, stage 0 -------------------------- beta_col = [] for k in cutlass.range_constexpr(num_vals): beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta0_idx]) - ainv_f16 = [] + tinv_frags = [] for c in cutlass.range_constexpr(ACC_N_FRAGS): - ainv_f16 += list( + tinv_frags += list( nvvm.ldmatrix( cutlass.inttoptr( - ainv0_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, + tinv0_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), @@ -1603,34 +1386,35 @@ def _compute0_warp( nvvm.MMALayout.ROW, ) ) - ainv_scaled = [] + tinv_pack = [] for j in cutlass.range_constexpr(num_vals // 2): - lo, hi = f16x2_to_f32(ainv_f16[j], dtype=cfg.io_dtype) + 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]) - ainv_scaled.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) + 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( - ainv0_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, + tinv0_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), - [ainv_scaled[c * 4 + 0], ainv_scaled[c * 4 + 1], ainv_scaled[c * 4 + 2], ainv_scaled[c * 4 + 3]], + [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_ainv_ready[ainv0_idx].arrive() + 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]) - ainv_f16 = [] + tinv_frags = [] for c in cutlass.range_constexpr(ACC_N_FRAGS): - ainv_f16 += list( + tinv_frags += list( nvvm.ldmatrix( cutlass.inttoptr( - ainv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, + tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), @@ -1638,100 +1422,77 @@ def _compute0_warp( nvvm.MMALayout.ROW, ) ) - ainv_scaled = [] + tinv_pack = [] for j in cutlass.range_constexpr(num_vals // 2): - lo, hi = f16x2_to_f32(ainv_f16[j], dtype=cfg.io_dtype) + 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]) - ainv_scaled.append(fp32_to_fp16(s0, s1, dtype=cfg.io_dtype)) + 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( - ainv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, + tinv1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16, ), - [ainv_scaled[c * 4 + 0], ainv_scaled[c * 4 + 1], ainv_scaled[c * 4 + 2], ainv_scaled[c * 4 + 3]], + [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_ainv_ready[ainv1_idx].arrive() + bars.mb_t_inv_ready[tinv1_idx].arrive() bars.mb_beta_done[beta1_idx].arrive() - # ---- Step 3: qk_epi0 ---------------------------------------- - qk0_idx = qk_index.idx - bars.mb_qk_done[qk0_idx].wait(qk_index.phase) - qk_index = advance(qk_index, cfg.smem_qk_stages) - qk0_acc_idx = cg0_acc_rdy.idx - bars.mb_cg0_acc_ready[qk0_acc_idx].wait(cg0_acc_rdy.phase) - cg0_acc_rdy = advance(cg0_acc_rdy, cfg.tmem_cg0_acc_stages) - - qk0_base = sQk[qk0_idx].base - qk_vec = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + qk0_acc_idx * ACC_STAGE_COLS, cutlass.Float32), num=8 + # ---- 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) + 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 ) - nvvm.tcgen05_wait("load") - bars.mb_cg0_acc_done[qk0_acc_idx].arrive() - qk_f16 = [] - for k in cutlass.range_constexpr(num_vals // 2): - p0, p1 = fmul2(qk_vec[2 * k], qk_vec[2 * k + 1], gT0[2 * k], gT0[2 * k + 1]) - v0, v1 = fmul2(p0, p1, scale, scale) - qk_f16.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) - for c in cutlass.range_constexpr(ACC_N_FRAGS): - nvvm.stmatrix( - cutlass.inttoptr( - qk0_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [qk_f16[c * 4 + 0], qk_f16[c * 4 + 1], qk_f16[c * 4 + 2], qk_f16[c * 4 + 3]], - nvvm.MMALayout.ROW, - ) - nvvm.fence_proxy("async.shared", space="cta") - bars.mb_qk_ready[qk0_idx].arrive() - - # ---- Step 4: qk_epi1 ---------------------------------------- - qk1_idx = qk_index.idx - bars.mb_qk_done[qk1_idx].wait(qk_index.phase) - qk_index = advance(qk_index, cfg.smem_qk_stages) - qk1_acc_idx = cg0_acc_rdy.idx - bars.mb_cg0_acc_ready[qk1_acc_idx].wait(cg0_acc_rdy.phase) - cg0_acc_rdy = advance(cg0_acc_rdy, cfg.tmem_cg0_acc_stages) - - qk1_base = sQk[qk1_idx].base - qk_vec = nvvm.tcgen05_ld( - "16x256b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_cg0_acc_col + qk1_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[qk1_acc_idx].arrive() - qk_f16 = [] - for k in cutlass.range_constexpr(num_vals // 2): - p0, p1 = fmul2(qk_vec[2 * k], qk_vec[2 * k + 1], gT1[2 * k], gT1[2 * k + 1]) - v0, v1 = fmul2(p0, p1, scale, scale) - qk_f16.append(fp32_to_fp16(v0, v1, dtype=cfg.io_dtype)) - for c in cutlass.range_constexpr(ACC_N_FRAGS): - nvvm.stmatrix( - cutlass.inttoptr( - qk1_base + (store_row * cfg.b_t + swizzle_xor_128b(store_row, store_col + c * FRAG_COLS)) * bpe, - cutlass.AddressSpace.smem, - cutlass.BFloat16, - ), - [qk_f16[c * 4 + 0], qk_f16[c * 4 + 1], qk_f16[c * 4 + 2], qk_f16[c * 4 + 3]], - nvvm.MMALayout.ROW, - ) + 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_qk_ready[qk1_idx].arrive() + 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_ainv_stages): - bars.mb_ainv_done[ainv_index.idx].wait(ainv_index.phase) - ainv_index = advance(ainv_index, cfg.smem_ainv_stages) - for _ in range(cfg.smem_qk_stages): - bars.mb_qk_done[qk_index.idx].wait(qk_index.phase) - qk_index = advance(qk_index, cfg.smem_qk_stages) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + for _ in range(cfg.smem_t_inv_stages): + bars.mb_t_inv_done[tinv_index.idx].wait(tinv_index.phase) + tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) + for _ in range(cfg.smem_a_stages): + bars.mb_a_done[a_index.idx].wait(a_index.phase) + a_index = advance(a_index, cfg.smem_a_stages) @cute.jit -def _compute1_warp( +def compute1_warp_group( cfg, total_tiles, bidx, @@ -1740,35 +1501,34 @@ def _compute1_warp( mWorkItems, tidx, warp_idx, - tmem_hold, + tmem_base_slot, scale, sV, sCumsumlog, sCumprod, sBeta, sO, - sH_raw, - mS_init, - mS_out, + sCheckpoint_raw, + mState_init, + mState_out, checkpoint_every_n_tokens, sSched, bars, ): - """Compute warp-group 1 role (warps 4-7): persistent scheduler loop, - initial-state seed, then one uniform per-chunk body: fused state - restage+rescale, V-K*S, state*q_epi, new_v_epi + decay_v publish, qkv - epilogue; final-state store per item.""" + """Compute warp-group 1 role (warps 4-7): persistent scheduler loop + running the per-chunk state-update and output epilogues.""" + elect_one = nvvm.elect_sync() v_index = PipelineState.start(phase=0) gate_index = PipelineState.start(phase=0) kv_acc_index = PipelineState.start(phase=0) - o_acc_rdy_index = PipelineState.start(phase=0) - o_scale_rdy_index = PipelineState.start(phase=0) - ks_rdy_index = PipelineState.start(phase=0) - nv_rdy_index = PipelineState.start(phase=0) - kv_acc_seed_index = PipelineState.start(phase=1) + o_acc_ready_index = PipelineState.start(phase=0) + o_final_acc_ready_index = PipelineState.start(phase=0) + k_state_ready_index = PipelineState.start(phase=0) + u_acc_ready_index = PipelineState.start(phase=0) + state_acc_seed_index = PipelineState.start(phase=1) o_index = PipelineState.start(phase=1) - si_cnt = cutlass.Int32(0) + state_inp_cnt = cutlass.Int32(0) kv_done_idx = cutlass.Int32(0) nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) @@ -1776,7 +1536,7 @@ def _compute1_warp( cfg.tmem_alloc_barrier_id, thread_count=cfg.tmem_alloc_barrier_threads, ) - tmem_base = tmem_hold.load() + tmem_base = tmem_base_slot.load() num_threads_cg1 = cfg.threads_per_warp * len(cfg.compute_group_1_warp_ids) cg1_tidx = tidx % num_threads_cg1 @@ -1785,93 +1545,100 @@ def _compute1_warp( ldtm_width = 32 sttm_width = ldtm_width // 2 num_state_subs = cutlass.const_expr(cfg.d_v // ldtm_width) - tmem_state_col = tmem_base + cfg.tmem_state_offset + tmem_state_col = tmem_base + cfg.tmem_state_acc_offset tmem_state_inp_col = tmem_base + cfg.tmem_state_inp_offset - tmem_q_state_col = tmem_base + cfg.tmem_q_state_offset - tmem_inp_col = tmem_base + cfg.tmem_inp_offset + tmem_q_state_col = tmem_base + cfg.tmem_q_state_acc_offset + tmem_inp_col = tmem_base + cfg.tmem_y_decay_u_inp_offset ACC_STAGE_COLS = cfg.b_t INP_SLOT_COLS = cfg.b_t // 2 - tmem_ks_col = tmem_base + cfg.tmem_cg1_acc_offset - tmem_nv_col = tmem_ks_col - tmem_vks_col = tmem_inp_col - tmem_nv_inp_col = tmem_inp_col + tmem_k_state_col = tmem_base + cfg.tmem_cg1_acc_offset + tmem_u_acc_col = tmem_k_state_col + tmem_y_inp_col = tmem_inp_col + tmem_u_inp_col = tmem_inp_col tmem_decay_v_col = tmem_inp_col + INP_SLOT_COLS - ov_tok = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 - ov_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 - ov_slab = (cg1_tidx // 64) * 4096 + v_o_smem_tok = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 + v_o_smem_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 + v_o_smem_subtile_off = (cg1_tidx // 64) * 4096 v_stage_elems = cfg.v_cosize // cfg.smem_v_stages o_stage_elems = cfg.o_cosize // cfg.smem_o_stages sV_base = cute.make_ptr(cfg.io_dtype, sV[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) sO_base = cute.make_ptr(cfg.io_dtype, sO[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) num_vals = 32 - if cutlass.const_expr(cfg.enable_h): - sH_base_int = sH_raw.data_ptr().toint() - h_cnt = cutlass.Int32(0) - h_ov_tok = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 - h_ov_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 + if cutlass.const_expr(cfg.enable_checkpoints): + sCheckpoint_base_int = sCheckpoint_raw.data_ptr().toint() + checkpoint_cnt = cutlass.Int32(0) + checkpoint_smem_row = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 + checkpoint_smem_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 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, cu_seqlens, mWorkItems - ) - sk_nt = wend - cstart - n_padded = ((sk_nt + 1) // 2) * 2 - if sk_nt > 0: + 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 + if n_local > 0: if cutlass.const_expr(cfg.use_initial_state): - # ---- initial-state seed: S_init GMEM -> state TMEM - # (split-K warmup items seed zeros through the same path) ---- - gS_init = mS_init[None, None, head_idx, batch_idx] - kv_init_idx = kv_acc_seed_index.idx - bars.mb_kv_acc_scale_done[kv_init_idx].wait(kv_acc_seed_index.phase) - kv_acc_seed_index = advance(kv_acc_seed_index, cfg.tmem_kv_acc_stages) - seed_from_s0 = cstart == 0 - for sub in cutlass.range_constexpr(num_state_subs): - words = [] - for k in cutlass.range_constexpr(32): - v = gS_init[sub * ldtm_width + k, cg1_tidx] - if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): - v = v.to(cfg.acc_dtype) - if cutlass.const_expr(cfg.split_k): - v = v if seed_from_s0 else cutlass.Float32(0.0) - words.append(v) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), - cutlass.Vector.from_elements(tuple(words), cutlass.Float32), - ) + # ---- initial-state seed: initial_state GMEM -> state TMEM --------------- + gState_init = mState_init[None, None, head_idx, batch_idx] + kv_init_idx = state_acc_seed_index.idx + bars.mb_state_acc_scale_done[kv_init_idx].wait(state_acc_seed_index.phase) + state_acc_seed_index = advance(state_acc_seed_index, cfg.tmem_state_acc_stages) + seed_from_initial_state = cstart == 0 + if seed_from_initial_state: + 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] + if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): + v = v.to(cfg.acc_dtype) + words.append(v) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + cutlass.Vector.from_elements(tuple(words), cutlass.Float32), + ) + else: + for sub in cutlass.range_constexpr(num_state_subs): + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + cutlass.Vector.from_elements(tuple(cutlass.Float32(0.0) for _ in range(32)), cutlass.Float32), + ) nvvm.tcgen05_wait("store") nvvm.barrier_cta_sync_aligned( cfg.init_state_store_barrier_id, thread_count=cfg.init_state_store_barrier_threads, ) - if cg1_tidx == 0: - arrive(bars.mb_kv_acc_ready[kv_init_idx].smem_ptr) for local_idx in cutlass.range(n_padded): # 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 valid_state = local_idx > 0 if cutlass.const_expr(cfg.use_initial_state): valid_state = cutlass.Boolean(True) - kv_acc_seed_index = advance(kv_acc_seed_index, cfg.tmem_kv_acc_stages) + state_acc_seed_index = advance(state_acc_seed_index, cfg.tmem_state_acc_stages) gate_idx = gate_index.idx bars.mb_gate_ready[gate_idx].wait(gate_index.phase) gate_index = advance(gate_index, cfg.smem_gate_stages) cumprod_total = sCumprod[sCumprod.shape[0] - 1, 0, gate_idx] - # ---- fused state restage + rescale (one read serves both) ---- + # ---- state restage + rescale ------------------------------------- if valid_state: kv_idx = kv_acc_index.idx - bars.mb_kv_acc_ready[kv_idx].wait(kv_acc_index.phase) - kv_acc_index = advance(kv_acc_index, cfg.tmem_kv_acc_stages) + bars.mb_state_acc_ready[kv_idx].wait(kv_acc_index.phase) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) kv_done_idx = kv_idx state_regs = [[cutlass.Float32(0.0) for _ in range(num_state_subs)] for _ in range(32)] - siu_idx = si_cnt % cfg.tmem_state_inp_stages - # loads before stores: a TMEM st blocks ld hoisting + state_inp_stage_idx = state_inp_cnt % cfg.tmem_state_inp_stages state_vecs = [] for sub in cutlass.range_constexpr(num_state_subs): state_vecs.append( @@ -1880,57 +1647,60 @@ def _compute1_warp( for sub in cutlass.range_constexpr(num_state_subs): for k in cutlass.range_constexpr(32): state_regs[k][sub] = state_vecs[sub][k] - state_f16 = [fp32_to_fp16(state_regs[2 * j][sub], state_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + state_pack = [fp32_to_fp16(state_regs[2 * j][sub], state_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_inp_col + sub * sttm_width, cutlass.Int32), - cutlass.Vector.from_elements(tuple(state_f16), cutlass.Int32), + cutlass.Vector.from_elements(tuple(state_pack), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready[siu_idx].arrive() - si_cnt = si_cnt + 1 - - if cutlass.const_expr(cfg.enable_h): - # ---- H checkpoint: state TMEM -> sH before the decayed - # write-back overwrites it ---- - do_h = (cfg.b_t * chunk_idx - 1) % checkpoint_every_n_tokens == checkpoint_every_n_tokens - 1 and chunk_idx < wend - if cutlass.const_expr(cfg.split_k): - do_h = do_h and chunk_idx >= wstart - if do_h: - h_regs = [[cutlass.Int32(0) for _ in range(16)] for _ in range(4)] + bars.mb_state_inp_ready[state_inp_stage_idx].arrive() + state_inp_cnt = state_inp_cnt + 1 + + 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 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 hh in cutlass.range_constexpr(2): - h_vec = nvvm.tcgen05_ld( + 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 + hh * 64, cutlass.Float32), + 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): - h_regs[b * 2 + hh][j] = fp32_to_fp16(h_vec[2 * j], h_vec[2 * j + 1], dtype=cfg.io_dtype) - bars.mb_h_tmastg_done[h_cnt % cfg.smem_h_stages].wait(cutlass.Int32(1) ^ ((h_cnt // cfg.smem_h_stages) & cutlass.Int32(1))) + 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 hh in cutlass.range_constexpr(2): - h_base = (h_cnt % cfg.smem_h_stages) * cfg.d_k * cfg.d_v + (cg1_tidx // 64) * cfg.d_k * 64 + 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): - h_row = hh * 64 + h_ov_tok + c * 16 + checkpoint_row = col_half * 64 + checkpoint_smem_row + c * 16 nvvm.stmatrix( cutlass.inttoptr( - sH_base_int + (h_base + h_row * 64 + swizzle_xor_128b(h_row, h_ov_col + b * 16)) * 2, + 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, ), [ - h_regs[b * 2 + hh][c * 4 + 0], - h_regs[b * 2 + hh][c * 4 + 1], - h_regs[b * 2 + hh][c * 4 + 2], - h_regs[b * 2 + hh][c * 4 + 3], + 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, ) nvvm.fence_proxy("async.shared", space="cta") - if nvvm.elect_sync(): - bars.mb_h_tmastg_ready[h_cnt % cfg.smem_h_stages].arrive() - h_cnt = h_cnt + 1 + if elect_one: + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + checkpoint_cnt = checkpoint_cnt + 1 for sub in cutlass.range_constexpr(num_state_subs): state_scaled = [] @@ -1943,184 +1713,176 @@ def _compute1_warp( cutlass.Vector.from_elements(tuple(state_scaled), cutlass.Float32), ) nvvm.tcgen05_wait("store") - bars.mb_kv_acc_scale_done[kv_idx].arrive() + bars.mb_state_acc_scale_done[kv_idx].arrive() - # ---- deferred per-row gate register builds --------------- - gCumprod = [] + # ---- per-row Gate register builds -------------------------------- + cumprod_vals = [] for k in cutlass.range_constexpr(num_vals): - gCumprod.append(sCumprod[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) + cumprod_vals.append(sCumprod[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) last_cumsumlog = sCumsumlog[cfg.b_t - 1, 0, gate_idx] - # gathers before math: interleaving pins one LDS latency per pair - gCs = [] + cumsumlog_vals = [] for k in cutlass.range_constexpr(num_vals): - gCs.append(sCumsumlog[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) - gDecayScale = [] + cumsumlog_vals.append(sCumsumlog[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) + decay_scale_vals = [] for k in cutlass.range_constexpr(0, num_vals, 2): - d0, d1 = fadd2(last_cumsumlog, last_cumsumlog, -gCs[k], -gCs[k + 1]) - gDecayScale.append(cute.math.exp2(d0, fastmath=True)) - gDecayScale.append(cute.math.exp2(d1, fastmath=True)) + d0, d1 = fadd2(last_cumsumlog, last_cumsumlog, -cumsumlog_vals[k], -cumsumlog_vals[k + 1]) + decay_scale_vals.append(cute.math.exp2(d0, fastmath=True)) + decay_scale_vals.append(cute.math.exp2(d1, fastmath=True)) bars.mb_gate_done[gate_idx].arrive() - # ---- v - k*state (packed 16-bit; V ring cursor survives - # item boundaries, so no fixed SMEM stage assumption) ------- + # ---- Y = V - K*state (packed 16-bit) ----------------------------- v_idx = v_index.idx bars.mb_v_ready[v_idx].wait(v_index.phase) v_index = advance(v_index, cfg.smem_v_stages) - v_words = [[cutlass.Int32(0), cutlass.Int32(0)] for _ in range(16)] + v_frags = [[cutlass.Int32(0), cutlass.Int32(0)] for _ in range(16)] for c in cutlass.range_constexpr(8): m0 = cutlass.const_expr(c % 4) sub = cutlass.const_expr(c // 4) - v_f16 = nvvm.ldmatrix( - (sV_base + v_idx * v_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16)).raw_ptr(), + v_frag = nvvm.ldmatrix( + ( + sV_base + + v_idx * v_stage_elems + + v_o_smem_subtile_off + + (v_o_smem_tok + m0 * 16) * 64 + + swizzle_xor_128b(v_o_smem_tok + m0 * 16, v_o_smem_col + sub * 16) + ).raw_ptr(), 4, nvvm.MMALayout.COL, ) for i in cutlass.range_constexpr(4): - v_words[4 * m0 + i][sub] = v_f16[i] + v_frags[4 * m0 + i][sub] = v_frag[i] if valid_state: - bars.mb_ks_ready[0].wait(ks_rdy_index.phase) - ks_rdy_index = advance(ks_rdy_index, 1) + bars.mb_k_state_acc_ready[0].wait(k_state_ready_index.phase) + k_state_ready_index = advance(k_state_ready_index, 1) for sub in cutlass.range_constexpr(2): - ks_vec = nvvm.tcgen05_ld( + k_state_vec = nvvm.tcgen05_ld( "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_ks_col, cutlass.Float32), + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_k_state_col, cutlass.Float32), num=8, ) for j in cutlass.range_constexpr(16): - s0, s1 = fmul2(ks_vec[2 * j], ks_vec[2 * j + 1], gCumprod[2 * j], gCumprod[2 * j + 1]) - ks_word = fp32_to_fp16(s0, s1, dtype=cfg.io_dtype) - v_words[j][sub] = sub_f16x2(v_words[j][sub], ks_word, cfg.io_dtype) + s0, s1 = fmul2(k_state_vec[2 * j], k_state_vec[2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) + k_state_word = fp32_to_fp16(s0, s1, dtype=cfg.io_dtype) + v_frags[j][sub] = sub_f16x2(v_frags[j][sub], k_state_word, cfg.io_dtype) for sub in cutlass.range_constexpr(2): nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_vks_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(v_words[j][sub] for j in range(16)), cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_inp_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(v_frags[j][sub] for j in range(16)), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_vks_inp_ready[0].arrive() + bars.mb_y_inp_ready[0].arrive() - # ---- state*q_epi: QS *= cumprod * scale, in place ------- + # ---- state*Q_epi: Q*state *= cumprod * scale ---------------- if valid_state: - qs_idx = o_acc_rdy_index.idx - bars.mb_o_acc_ready[qs_idx].wait(o_acc_rdy_index.phase) - o_acc_rdy_index = advance(o_acc_rdy_index, cfg.tmem_q_state_acc_stages) + q_state_idx = o_acc_ready_index.idx + bars.mb_o_acc_ready[q_state_idx].wait(o_acc_ready_index.phase) + o_acc_ready_index = advance(o_acc_ready_index, cfg.tmem_q_state_acc_stages) - qs_ptrs = [] - qs_vecs = [] + q_state_ptrs = [] + q_state_vecs = [] for sub in cutlass.range_constexpr(2): - qs_ptrs.append(nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_q_state_col + qs_idx * ACC_STAGE_COLS, cutlass.Float32)) - qs_vecs.append(nvvm.tcgen05_ld("16x256b", qs_ptrs[sub], num=8)) + q_state_ptrs.append( + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_q_state_col + q_state_idx * ACC_STAGE_COLS, cutlass.Float32) + ) + q_state_vecs.append(nvvm.tcgen05_ld("16x256b", q_state_ptrs[sub], num=8)) for sub in cutlass.range_constexpr(2): - qs_scaled = [] + q_state_scaled = [] for j in cutlass.range_constexpr(16): - p0, p1 = fmul2(qs_vecs[sub][2 * j], qs_vecs[sub][2 * j + 1], gCumprod[2 * j], gCumprod[2 * j + 1]) + p0, p1 = fmul2(q_state_vecs[sub][2 * j], q_state_vecs[sub][2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) s0, s1 = fmul2(p0, p1, scale, scale) - qs_scaled += [s0, s1] - nvvm.tcgen05_st("16x256b", qs_ptrs[sub], cutlass.Vector.from_elements(tuple(qs_scaled), cutlass.Float32)) + q_state_scaled += [s0, s1] + nvvm.tcgen05_st("16x256b", q_state_ptrs[sub], cutlass.Vector.from_elements(tuple(q_state_scaled), cutlass.Float32)) nvvm.tcgen05_wait("store") - bars.mb_o_state_scale_acc_done[qs_idx].arrive() + bars.mb_o_state_scale_acc_done[q_state_idx].arrive() - # ---- new_v_epi + decay_v publish ------------------------- - bars.mb_nv_ready[0].wait(nv_rdy_index.phase) - nv_rdy_index = advance(nv_rdy_index, 1) + # ---- U_epi + decayed-U publish ----------------------------------- + bars.mb_u_acc_ready[0].wait(u_acc_ready_index.phase) + u_acc_ready_index = advance(u_acc_ready_index, 1) bars.mb_v_done[v_idx].arrive() - nv_regs = [[cutlass.Float32(0.0), cutlass.Float32(0.0)] for _ in range(32)] - # NV reuses the VKS slot; both publishes share one fence - nv_vecs = [] + u_regs = [[cutlass.Float32(0.0), cutlass.Float32(0.0)] for _ in range(32)] + u_vecs = [] for sub in cutlass.range_constexpr(2): - nv_vecs.append( + u_vecs.append( nvvm.tcgen05_ld( "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_nv_col, cutlass.Float32), + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_u_acc_col, cutlass.Float32), num=8, ) ) for sub in cutlass.range_constexpr(2): for k in cutlass.range_constexpr(32): - nv_regs[k][sub] = nv_vecs[sub][k] + u_regs[k][sub] = u_vecs[sub][k] - nv_f16 = [fp32_to_fp16(nv_regs[2 * j][sub], nv_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + u_pack = [fp32_to_fp16(u_regs[2 * j][sub], u_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "16x128b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_nv_inp_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(nv_f16), cutlass.Int32), + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_u_inp_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(u_pack), cutlass.Int32), ) + nvvm.tcgen05_wait("store") + bars.mb_u_inp_ready[0].arrive() + + for sub in cutlass.range_constexpr(2): for j in cutlass.range_constexpr(16): - nv_regs[2 * j][sub], nv_regs[2 * j + 1][sub] = fmul2( - nv_regs[2 * j][sub], nv_regs[2 * j + 1][sub], gDecayScale[2 * j], gDecayScale[2 * j + 1] + u_regs[2 * j][sub], u_regs[2 * j + 1][sub] = fmul2( + u_regs[2 * j][sub], u_regs[2 * j + 1][sub], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1] ) - decay_f16 = [fp32_to_fp16(nv_regs[2 * j][sub], nv_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + decay_pack = [fp32_to_fp16(u_regs[2 * j][sub], u_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] nvvm.tcgen05_st( "16x128b", nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_decay_v_col, cutlass.Int32), - cutlass.Vector.from_elements(tuple(decay_f16), cutlass.Int32), + cutlass.Vector.from_elements(tuple(decay_pack), cutlass.Int32), ) nvvm.tcgen05_wait("store") - bars.mb_nv_inp_ready[0].arrive() - bars.mb_decay_v_inp_ready[0].arrive() - - # ---- qkv_epilogue: O acc -> sO --------------------------- - if cutlass.const_expr(cfg.enable_o): - o_idx = o_index.idx - bars.mb_o_tmastg_done[o_idx].wait(o_index.phase) - o_index = advance(o_index, cfg.smem_o_stages) - qs2_idx = o_scale_rdy_index.idx - bars.mb_o_state_scale_acc_ready[qs2_idx].wait(o_scale_rdy_index.phase) - o_scale_rdy_index = advance(o_scale_rdy_index, cfg.tmem_q_state_acc_stages) - - if cutlass.const_expr(cfg.enable_o): - o_regs = [] - for sub in cutlass.range_constexpr(2): - o_vec = nvvm.tcgen05_ld( - "16x256b", - nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_q_state_col + qs2_idx * ACC_STAGE_COLS, cutlass.Float32), - num=8, - ) - o_regs.append([o_vec[k] for k in range(32)]) - # REQUIRED: an arrive does not order in-flight tcgen05 loads + bars.mb_decay_u_inp_ready[0].arrive() + + # ---- QKV_epilogue: O acc TMEM -> sO SMEM ------------------------- + o_scale_idx = o_final_acc_ready_index.idx + bars.mb_o_final_acc_ready[o_scale_idx].wait(o_final_acc_ready_index.phase) + o_final_acc_ready_index = advance(o_final_acc_ready_index, cfg.tmem_q_state_acc_stages) + + o_regs = [] + for sub in cutlass.range_constexpr(2): + o_vec = nvvm.tcgen05_ld( + "16x256b", + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_q_state_col + o_scale_idx * ACC_STAGE_COLS, cutlass.Float32), + num=8, + ) + o_regs.append([o_vec[k] for k in range(32)]) nvvm.tcgen05_wait("load") - bars.mb_o_acc_done[qs2_idx].arrive() - if cutlass.const_expr(cfg.enable_o): - for sub in cutlass.range_constexpr(2): - for m0 in cutlass.range_constexpr(4): - o_f16 = [fp32_to_fp16(o_regs[sub][8 * m0 + 2 * j], o_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] - nvvm.stmatrix( - ( - sO_base + o_idx * o_stage_elems + ov_slab + (ov_tok + m0 * 16) * 64 + swizzle_xor_128b(ov_tok + m0 * 16, ov_col + sub * 16) - ).raw_ptr(), - o_f16, - nvvm.MMALayout.COL, - ) - nvvm.fence_proxy("async.shared", space="cta") - - bars.mb_o_tmastg_ready[o_idx].arrive() - - # ---- final state S: state TMEM -> GMEM ------------------------- - if sk_nt > 0: - # required even unstored: next item's decay_v must not race GEMM 7 + o_idx = o_index.idx + bars.mb_o_tmastg_done[o_idx].wait(o_index.phase) + o_index = advance(o_index, cfg.smem_o_stages) + for sub in cutlass.range_constexpr(2): + for m0 in cutlass.range_constexpr(4): + o_pack = [fp32_to_fp16(o_regs[sub][8 * m0 + 2 * j], o_regs[sub][8 * m0 + 2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + nvvm.stmatrix( + ( + sO_base + + o_idx * o_stage_elems + + v_o_smem_subtile_off + + (v_o_smem_tok + m0 * 16) * 64 + + swizzle_xor_128b(v_o_smem_tok + m0 * 16, v_o_smem_col + sub * 16) + ).raw_ptr(), + o_pack, + nvvm.MMALayout.COL, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_o_tmastg_ready[o_idx].arrive() + + # ---- final state: state TMEM -> GMEM ----------------------------------- + if n_local > 0: kv_last_idx = kv_acc_index.idx - bars.mb_kv_acc_ready[kv_last_idx].wait(kv_acc_index.phase) - kv_acc_index = advance(kv_acc_index, cfg.tmem_kv_acc_stages) + bars.mb_state_acc_ready[kv_last_idx].wait(kv_acc_index.phase) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) if cutlass.const_expr(cfg.store_final_state): - # split-K: only the last chunk's owner holds the final state - if cutlass.const_expr(cfg.split_k): - if wend == num_chunks_b: - gS_out = mS_out[None, None, head_idx, batch_idx] - for sub in cutlass.range_constexpr(num_state_subs): - state_vec = nvvm.tcgen05_ld( - "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), num=32 - ) - for k in cutlass.range_constexpr(32): - val = state_vec[k] - if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): - val = val.to(cfg.state_dtype) - gS_out[sub * ldtm_width + k, cg1_tidx] = val - else: - gS_out = mS_out[None, None, head_idx, batch_idx] + if wend == num_chunks_b: + gState_out = mState_out[None, None, head_idx, batch_idx] for sub in cutlass.range_constexpr(num_state_subs): state_vec = nvvm.tcgen05_ld( "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), num=32 @@ -2129,45 +1891,99 @@ def _compute1_warp( val = state_vec[k] if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): val = val.to(cfg.state_dtype) - gS_out[sub * ldtm_width + k, cg1_tidx] = val - bars.mb_kv_acc_scale_done[kv_last_idx].arrive() + gState_out[sub * ldtm_width + k, cg1_tidx] = val + bars.mb_state_acc_scale_done[kv_last_idx].arrive() else: - bars.mb_kv_acc_scale_done[kv_last_idx].arrive() + bars.mb_state_acc_scale_done[kv_last_idx].arrive() else: - # zero-length sequence: state passes through, pure GMEM if cutlass.const_expr(cfg.store_final_state): - write_passthrough = True - if cutlass.const_expr(cfg.split_k): - write_passthrough = wend == num_chunks_b + write_passthrough = wend == num_chunks_b if write_passthrough: - gS_out = mS_out[None, None, head_idx, batch_idx] + gState_out = mState_out[None, None, head_idx, batch_idx] if cutlass.const_expr(cfg.use_initial_state): - gS_in = mS_init[None, None, head_idx, batch_idx] + 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): - gS_out[sub * ldtm_width + k, cg1_tidx] = gS_in[sub * ldtm_width + k, cg1_tidx] + gState_out[sub * ldtm_width + k, cg1_tidx] = gState_in[sub * ldtm_width + k, cg1_tidx] else: for sub in cutlass.range_constexpr(num_state_subs): for k in cutlass.range_constexpr(32): - gS_out[sub * ldtm_width + k, cg1_tidx] = cutlass.Float32(0.0).to(cfg.state_dtype) + gState_out[sub * ldtm_width + k, cg1_tidx] = cutlass.Float32(0.0).to(cfg.state_dtype) - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) - # CG1 done with TMEM: release the MMA warp's dealloc bars.mb_tmem_done[0].arrive() - if cutlass.const_expr(cfg.enable_o): - for _ in range(cfg.smem_o_stages): - bars.mb_o_tmastg_done[o_index.idx].wait(o_index.phase) - o_index = advance(o_index, cfg.smem_o_stages) - if cutlass.const_expr(cfg.enable_h): - for _ in range(cfg.smem_h_stages): - bars.mb_h_tmastg_done[h_cnt % cfg.smem_h_stages].wait(cutlass.Int32(1) ^ ((h_cnt // cfg.smem_h_stages) & cutlass.Int32(1))) - h_cnt = h_cnt + 1 + for _ in range(cfg.smem_o_stages): + bars.mb_o_tmastg_done[o_index.idx].wait(o_index.phase) + o_index = advance(o_index, cfg.smem_o_stages) + if cutlass.const_expr(cfg.enable_checkpoints): + for _ in range(cfg.smem_checkpoint_stages): + 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_cnt = checkpoint_cnt + 1 + + +@cute.kernel +def build_all_descs_kernel( + base_q: cutlass.GridConstant[tma.TensorMap], + base_k: cutlass.GridConstant[tma.TensorMap], + base_v: cutlass.GridConstant[tma.TensorMap], + base_o: cutlass.GridConstant[tma.TensorMap], + base_checkpoint: cutlass.GridConstant[tma.TensorMap], + desc_ws: cute.Tensor, + cu_seqlens: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + o: Optional[cute.Tensor], + state_checkpoints_out: Optional[cute.Tensor], + n_batch: cutlass.Int32, + q_row_stride: cutlass.Int32, + k_row_stride: cutlass.Int32, + v_row_stride: cutlass.Int32, + o_row_stride: cutlass.Int32, + checkpoint_row_stride: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, +) -> None: + """Single-launch builder for the per-batch descriptor arrays (one warp + per array).""" + tidx, _, _ = cute.arch.thread_idx() + widx = cutlass.Int32(tidx) // cutlass.Int32(32) + arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) + sub0 = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + + if widx == 0: + if nvvm.elect_sync(): + emit_seq_descs(base_q, sub0, cu_seqlens, q, n_batch, q_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 1: + if nvvm.elect_sync(): + emit_seq_descs(base_k, sub1, cu_seqlens, k, n_batch, k_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 2: + if nvvm.elect_sync(): + emit_seq_descs(base_v, sub2, cu_seqlens, v, n_batch, v_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if cutlass.const_expr(o is not None): + if widx == 3: + if nvvm.elect_sync(): + emit_seq_descs(base_o, sub3, cu_seqlens, o, n_batch, o_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if cutlass.const_expr(state_checkpoints_out is not None): + if widx == 4: + if nvvm.elect_sync(): + emit_checkpoint_seq_descs(base_checkpoint, sub4, cu_seqlens, state_checkpoints_out, 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) @cute.jit -def _build_descs( +def build_descs( io_dtype: cutlass.Constexpr, b_t: cutlass.Constexpr[int], q: cute.Tensor, @@ -2175,191 +1991,84 @@ def _build_descs( v: cute.Tensor, o: Optional[cute.Tensor], cu_seqlens: cute.Tensor, - h_out: Optional[cute.Tensor], - h_every_n: cutlass.Int32, + state_checkpoints_out: Optional[cute.Tensor], + checkpoint_every_n: cutlass.Int32, tensormap_workspace: cute.Tensor, stream: cuda.CUstream, ): - """Build the 5 per-(b,h) TMA-descriptor arrays (Q, K, V, O, S) into - ``tensormap_workspace``. Compiled and launched separately from the main - kernel, and launched on every execute: the descriptors fold ``cu_seqlens`` - contents into GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read - without a D2H sync that CUDA-graph capture forbids. - - The H descriptor is 3-D ``(dv, dk, h)`` over the packed - ``[total_h, HO, DK, DV]`` H tensor; ``build_h_descs_kernel`` derives the - per-sequence H offset from the token ``cu_seqlens`` ((seqlen-1)//h_every_n, - prefix-summed), folds it and the head into GLOBAL_ADDRESS, and caps - GLOBAL_DIM[2] to the per-sequence H count, so the store coordinate is the - sequence-local H index.""" + """Build the 5 per-(b,h) TMA-descriptor arrays (Q, K, V, O, checkpoints) into + ``tensormap_workspace``.""" h_q = q.shape[1] h_k = k.shape[1] h_v = v.shape[1] batch_size = cu_seqlens.shape[0] - 1 heads_out = h_q if h_q >= h_v else h_v - q_group = heads_out // h_q - k_group = heads_out // h_k - v_group = heads_out // h_v d_v = v.shape[2] bpe = io_dtype.width // 8 - granu = 128 // bpe + elems_per_128b = 128 // bpe bt = b_t q_row_stride, q_head_stride = q.stride[0], q.stride[1] k_row_stride, k_head_stride = k.stride[0], k.stride[1] v_row_stride, v_head_stride = v.stride[0], v.stride[1] - q_head0 = q[None, 0, None] - k_head0 = k[None, 0, None] - v_view = v[None, 0, None] - v_head0 = cute.make_tensor( - v_view.iterator, - cute.make_layout((d_v, v_view.shape[0]), stride=(v_view.stride[1], v_view.stride[0])), - ) - swz128 = _tma.TensorMapSwizzle.s128b - base_desc_q = _tma.create_tensor_map_tiled_from_view(q_head0, box_dims=(bt, granu), stride_order=(1, 0), swizzle=swz128) - base_desc_k = _tma.create_tensor_map_tiled_from_view(k_head0, box_dims=(bt, granu), stride_order=(1, 0), swizzle=swz128) - base_desc_v = _tma.create_tensor_map_tiled_from_view(v_head0, box_dims=(granu, bt), stride_order=(0, 1), swizzle=swz128) - - arr_words = (batch_size * heads_out) * TENSOR_MAP_QWORDS - ws_iter = tensormap_workspace.iterator - - def sub_array(k): - return cute.make_tensor(ws_iter + k * arr_words, cute.make_layout((arr_words,), stride=(1,))) + seqlen = q.shape[0] + d_k = q.shape[2] + q_headed = cute.make_tensor(q.iterator, cute.make_layout((seqlen, h_q, d_k), stride=(q_row_stride, q_head_stride, 1))) + k_headed = cute.make_tensor(k.iterator, cute.make_layout((seqlen, h_k, d_k), stride=(k_row_stride, k_head_stride, 1))) + v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v_head_stride, v_row_stride))) + swz128 = tma.TensorMapSwizzle.s128b + base_desc_q = tma.create_tensor_map_tiled_from_view(q_headed, box_dims=(bt, 1, elems_per_128b), stride_order=(2, 1, 0), swizzle=swz128) + base_desc_k = tma.create_tensor_map_tiled_from_view(k_headed, box_dims=(bt, 1, elems_per_128b), stride_order=(2, 1, 0), swizzle=swz128) + base_desc_v = tma.create_tensor_map_tiled_from_view(v_headed, box_dims=(elems_per_128b, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + + base_desc_o = base_desc_v + if cutlass.const_expr(o is not None): + o_headed = cute.make_tensor(o.iterator, cute.make_layout((d_v, heads_out, seqlen), stride=(1, o.stride[1], o.stride[0]))) + base_desc_o = tma.create_tensor_map_tiled_from_view(o_headed, box_dims=(elems_per_128b, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + + base_desc_checkpoint = base_desc_v + if cutlass.const_expr(state_checkpoints_out is not None): + d_k_state = state_checkpoints_out.shape[2] + d_v_state = state_checkpoints_out.shape[3] + checkpoint_elems_per_128b = 128 // (state_checkpoints_out.element_type.width // 8) + checkpoint_view = cute.make_tensor( + state_checkpoints_out.iterator, + cute.make_layout( + (d_v_state, d_k_state, state_checkpoints_out.shape[0], heads_out), + stride=(state_checkpoints_out.stride[3], state_checkpoints_out.stride[2], state_checkpoints_out.stride[0], state_checkpoints_out.stride[1]), + ), + ) + base_desc_checkpoint = tma.create_tensor_map_tiled_from_view( + checkpoint_view, box_dims=(checkpoint_elems_per_128b, d_k_state, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz128 + ) - build_qkv_load_descs_kernel( + n_warps = 5 if state_checkpoints_out is not None else (4 if o is not None else 3) + build_all_descs_kernel( base_desc_q, - sub_array(0), - cu_seqlens, - q, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(q_group), - cutlass.Int32(q_head_stride), - cutlass.Int32(q_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( base_desc_k, - sub_array(1), - cu_seqlens, - k, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(k_group), - cutlass.Int32(k_head_stride), - cutlass.Int32(k_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( base_desc_v, - sub_array(2), + base_desc_o, + base_desc_checkpoint, + tensormap_workspace, cu_seqlens, + q, + k, v, + o, + state_checkpoints_out, cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(v_group), - cutlass.Int32(v_head_stride), + cutlass.Int32(q_row_stride), + cutlass.Int32(k_row_stride), cutlass.Int32(v_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - if cutlass.const_expr(o is not None): - o_row_stride, o_head_stride = o.stride[0], o.stride[1] - o_view = o[None, 0, None] - o_head0 = cute.make_tensor( - o_view.iterator, - cute.make_layout((d_v, o_view.shape[0]), stride=(o_view.stride[1], o_view.stride[0])), - ) - base_desc_o = _tma.create_tensor_map_tiled_from_view(o_head0, box_dims=(granu, bt), stride_order=(0, 1), swizzle=swz128) - build_qkv_load_descs_kernel( - base_desc_o, - sub_array(3), - cu_seqlens, - o, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(1), - cutlass.Int32(o_head_stride), - cutlass.Int32(o_row_stride), - 1, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - - if cutlass.const_expr(h_out is not None): - d_k_s = h_out.shape[2] - d_v_s = h_out.shape[3] - h_granu = 128 // (h_out.element_type.width // 8) - h_view = cute.make_tensor( - h_out.iterator, - cute.make_layout( - (d_v_s, d_k_s, h_out.shape[0]), - stride=(h_out.stride[3], h_out.stride[2], h_out.stride[0]), - ), - ) - base_desc_h = _tma.create_tensor_map_tiled_from_view(h_view, box_dims=(h_granu, d_k_s, 1), stride_order=(0, 1, 2), swizzle=swz128) - build_h_descs_kernel( - base_desc_h, - sub_array(4), - cu_seqlens, - h_out, - cutlass.Int32(batch_size), - cutlass.Int32(heads_out), - cutlass.Int32(h_out.stride[1]), - cutlass.Int32(h_out.stride[0]), - h_every_n, - 2, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) + cutlass.Int32(o.stride[0] if o is not None else 0), + cutlass.Int32(state_checkpoints_out.stride[0] if state_checkpoints_out is not None else 0), + checkpoint_every_n, + ).launch(grid=(1, 1, 1), block=(32 * n_warps, 1, 1), stream=stream) @cute.jit -def _downcast_state( - s0: cute.Tensor, - out: cute.Tensor, - n: cutlass.Int32, - n_blocks: cutlass.Int32, - stream: cuda.CUstream, -): - downcast_state_kernel( - s0, - out, - n, - ).launch(grid=(n_blocks, 1, 1), block=(128, 1, 1), stream=stream) - - -@functools.cache -def _downcast_state_cache(key): - return {} - - -def downcast_state(initial_state, out, *, stream): - """Copy the fp32 initial state ``[N, HO, K, V]`` into ``out`` (io dtype, - same shape) — the buffer the backward's per-(b,h) state descriptors - read.""" - n = 1 - for s_ in out.shape: - n *= int(s_) - key = (str(initial_state.dtype), str(out.dtype)) - cache = _downcast_state_cache(key) - cu_stream = cuda.CUstream(int(stream)) - n_blocks = (n + 127) // 128 - s0_flat = initial_state.reshape(n) - out_flat = out.reshape(n) - if "compiled" not in cache: - s0_c = from_dlpack(s0_flat, assumed_align=16).mark_layout_dynamic() - out_c = from_dlpack(out_flat, assumed_align=16).mark_layout_dynamic() - cache["compiled"] = cute.compile( - _downcast_state, - s0_c, - out_c, - cutlass.Int32(n), - cutlass.Int32(n_blocks), - cu_stream, - options="--enable-tvm-ffi", - ) - cache["compiled"](s0_flat, out_flat, n, n_blocks, cu_stream) - - -@cute.jit -def _host( +def host( cfg: cutlass.Constexpr, q: cute.Tensor, k: cute.Tensor, @@ -2368,8 +2077,8 @@ def _host( beta: cute.Tensor, o: Optional[cute.Tensor], cu_seqlens: cute.Tensor, - s_in: Optional[cute.Tensor], - s_out: Optional[cute.Tensor], + state_in: Optional[cute.Tensor], + state_out: Optional[cute.Tensor], work_items: Optional[cute.Tensor], work_count: Optional[cute.Tensor], sched_ctr: Optional[cute.Tensor], @@ -2379,11 +2088,12 @@ def _host( stream: cuda.CUstream, ): h_q = q.shape[1] + h_k = k.shape[1] h_v = v.shape[1] batch_size = cu_seqlens.shape[0] - 1 heads_out = h_q if h_q >= h_v else h_v - # ---- GQA reshapes: fold the head group into a ---------------------- + # ---- GQA reshapes: fold the head group into a -------------------------------- if cutlass.const_expr(cfg.is_GQA): h_r = h_q // h_v h_qv = h_v @@ -2455,71 +2165,71 @@ def _host( stride=(o.stride[2], o.stride[0], (o.stride[1], h_r * o.stride[1])), ), ) - if cutlass.const_expr(s_in is not None): - s_in = cute.make_tensor( - s_in.iterator, + if cutlass.const_expr(state_in is not None): + state_in = cute.make_tensor( + state_in.iterator, cute.make_layout( - (s_in.shape[2], s_in.shape[3], (h_r, h_qv), s_in.shape[0]), + (state_in.shape[2], state_in.shape[3], (h_r, h_qv), state_in.shape[0]), stride=( - s_in.stride[2], - s_in.stride[3], - (s_in.stride[1], h_r * s_in.stride[1]), - s_in.stride[0], + state_in.stride[2], + state_in.stride[3], + (state_in.stride[1], h_r * state_in.stride[1]), + state_in.stride[0], ), ), ) - if cutlass.const_expr(s_out is not None): - s_out = cute.make_tensor( - s_out.iterator, + if cutlass.const_expr(state_out is not None): + state_out = cute.make_tensor( + state_out.iterator, cute.make_layout( - (s_out.shape[2], s_out.shape[3], (h_r, h_qv), s_out.shape[0]), + (state_out.shape[2], state_out.shape[3], (h_r, h_qv), state_out.shape[0]), stride=( - s_out.stride[2], - s_out.stride[3], - (s_out.stride[1], h_r * s_out.stride[1]), - s_out.stride[0], + state_out.stride[2], + state_out.stride[3], + (state_out.stride[1], h_r * state_out.stride[1]), + state_out.stride[0], ), ), ) - # ---- SMEM sizing: per-buffer element cosizes ----------------------- + + # ---- SMEM sizing: per-buffer element cosizes --------------------------------- bpe = cfg.io_dtype.width // 8 - q_tile_elems = cfg.b_t * cfg.d_k - k_tile_elems = cfg.b_t * cfg.d_k + kq_tile_elems = 2 * cfg.b_t * cfg.d_k v_tile_elems = cfg.d_v * cfg.b_t - ainv_tile_elems = cfg.b_t * cfg.b_t - qk_tile_elems = cfg.b_t * cfg.b_t + tinv_tile_elems = cfg.b_t * cfg.b_t + a_tile_elems = cfg.b_t * cfg.b_t o_tile_elems = cfg.d_v * cfg.b_t - cfg.q_cosize = q_tile_elems * cfg.smem_q_stages - cfg.k_cosize = k_tile_elems * cfg.smem_k_stages + cfg.kq_cosize = kq_tile_elems * cfg.smem_kq_stages cfg.v_cosize = v_tile_elems * cfg.smem_v_stages - cfg.ainv_cosize = ainv_tile_elems * cfg.smem_ainv_stages - cfg.qk_cosize = qk_tile_elems * cfg.smem_qk_stages + cfg.t_inv_cosize = tinv_tile_elems * cfg.smem_t_inv_stages + cfg.a_cosize = a_tile_elems * cfg.smem_a_stages cfg.o_cosize = o_tile_elems * cfg.smem_o_stages - cfg.h_cosize = cfg.d_k * cfg.d_v * cfg.smem_h_stages + cfg.checkpoint_cosize = cfg.d_k * cfg.d_v * cfg.smem_checkpoint_stages cumsumlog_smem_layout_staged = cute.make_layout((cfg.b_t, 1, cfg.smem_gate_stages)) beta_smem_layout_staged = cute.make_layout((cfg.b_t, 1, cfg.smem_beta_stages)) - cfg.tma_q_bytes = q_tile_elems * bpe - cfg.tma_k_bytes = k_tile_elems * bpe + cfg.tma_kq_bytes = kq_tile_elems * bpe cfg.tma_v_bytes = v_tile_elems * bpe cfg.tma_o_bytes = o_tile_elems * bpe cfg.n_heads_out = heads_out - num_descs = batch_size * heads_out + cfg.q_ratio = heads_out // h_q + cfg.k_ratio = heads_out // h_k + cfg.v_ratio = heads_out // h_v + num_descs = batch_size - # ---- launch -------------------------------------------------------- - total_tiles = batch_size * heads_out + # ---- launch ------------------------------------------------------------------ # CUDA-graph-stable launch: fixed SM-count grid; shapes ride on buffer contents grid_shape = (cfg.max_active_clusters, 1, 1) - _kernel( + kernel( cfg, gate, beta, cu_seqlens, - s_in, - s_out, + state_in, + state_out, work_items, work_count, sched_ctr, @@ -2527,7 +2237,6 @@ def _host( scale, cumsumlog_smem_layout_staged, beta_smem_layout_staged, - total_tiles, q, k, v, @@ -2544,21 +2253,20 @@ def _host( @cute.kernel -def _kernel( +def kernel( cfg: cutlass.Constexpr, mGate: cute.Tensor, mBeta: cute.Tensor, cu_seqlens: cute.Tensor, - mS_init: Optional[cute.Tensor], - mS_out: Optional[cute.Tensor], - mWorkItems: Optional[cute.Tensor], - mCount: Optional[cute.Tensor], + mState_init: Optional[cute.Tensor], + mState_out: Optional[cute.Tensor], + mWorkItems: cute.Tensor, + mCount: cute.Tensor, mSched: Optional[cute.Tensor], checkpoint_every_n_tokens: cutlass.Int32, scale: cutlass.Float32, cumsumlog_smem_layout_staged: cute.Layout, beta_smem_layout_staged: cute.Layout, - total_tiles: cutlass.Int32, mQ, mK, mV, @@ -2566,31 +2274,24 @@ def _kernel( tensormap_workspace: cute.Tensor, n_desc: cutlass.Int32, ): - """ - Main GDN chunked kernel. - - Warp specialization is the outermost control flow: each warp role owns - its own persistent tile-scheduler loop, iterating over (batch, head) - tiles and then over chunks within each tile. - """ + """Main GDN chunked kernel: warp-specialized persistent tile loop.""" tidx, _, _ = cute.arch.thread_idx() warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) bidx = cute.arch.block_idx()[0] num_ctas = cute.arch.grid_dim()[0] - if cutlass.const_expr(cfg.split_k): - total_tiles = mCount[0] + total_tiles = mCount[0] if cutlass.const_expr(cfg.dyn_sched): assert mSched is not None, "mSched must be provided if dyn_sched is True" if cutlass.const_expr(cfg.use_initial_state): - assert mS_init is not None, "mS_init must be provided if use_initial_state is True" + assert mState_init is not None, "mState_init must be provided if use_initial_state is True" else: - assert mS_init is None, "mS_init must be None if use_initial_state is False" + assert mState_init is None, "mState_init must be None if use_initial_state is False" if cutlass.const_expr(cfg.store_final_state): - assert mS_out is not None, "mS_out must be provided if store_final_state is True" + assert mState_out is not None, "mState_out must be provided if store_final_state is True" else: - assert mS_out is None, "mS_out must be None if store_final_state is False" + assert mState_out is None, "mState_out must be None if store_final_state is False" desc_base_words = tensormap_workspace.iterator.raw_ptr() desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) @@ -2599,15 +2300,9 @@ def _kernel( desc_k_base = desc_base_words + arr_words desc_v_base = desc_base_words + cutlass.Int32(2) * arr_words desc_o_base = desc_base_words + cutlass.Int32(3) * arr_words - desc_h_base = desc_base_words + cutlass.Int32(4) * arr_words + desc_checkpoint_base = desc_base_words + cutlass.Int32(4) * arr_words SMEM = cutlass.AddressSpace.smem - bars = make_gdn_bars(cfg) - tmem_hold = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=16) - sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) - cumsumlog_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) - cumprod_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) - beta_raw = cutlass.Array(cutlass.Float32, cute.cosize(beta_smem_layout_staged), space=SMEM, alignment=128) bpe = cfg.io_dtype.width // 8 SWZ = 2 @@ -2615,71 +2310,85 @@ def _kernel( STRIDE = 8 * 128 KT_LEAD = (cfg.d_v // 2) * 128 V_LEAD = (cfg.d_v // 2) * 128 - sQ_raw = cutlass.Array( + sO_raw = cutlass.Array( cfg.io_dtype, - cfg.q_cosize, + cfg.o_cosize, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sQ = SmemTile( - base=sQ_raw.data_ptr().toint(), - elems_per_stage=(cfg.q_cosize // cfg.smem_q_stages) * bpe, - stages=cfg.smem_q_stages, + sO = SmemTile( + base=sO_raw.data_ptr().toint(), + elems_per_stage=(cfg.o_cosize // cfg.smem_o_stages) * bpe, + stages=cfg.smem_o_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sK_raw = cutlass.Array( + if cutlass.const_expr(cfg.enable_checkpoints): + sCheckpoint_raw = cutlass.Array( + cfg.io_dtype, + cfg.checkpoint_cosize, + space=cutlass.AddressSpace.smem, + alignment=cfg.buffer_align_bytes, + ) + else: + sCheckpoint_raw = None + sKQ_raw = cutlass.Array( cfg.io_dtype, - cfg.k_cosize, + cfg.kq_cosize, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sK = SmemTile( - base=sK_raw.data_ptr().toint(), - elems_per_stage=(cfg.k_cosize // cfg.smem_k_stages) * bpe, - stages=cfg.smem_k_stages, + sKQ = SmemTile( + base=sKQ_raw.data_ptr().toint(), + elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages) * bpe, + stages=cfg.smem_kq_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sK_trans = SmemTile( - base=sK_raw.data_ptr().toint(), - elems_per_stage=(cfg.k_cosize // cfg.smem_k_stages) * bpe, - stages=cfg.smem_k_stages, - leading_byte_offset=KT_LEAD, + sKQ_trans = SmemTile( + base=sKQ_raw.data_ptr().toint(), + elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages) * bpe, + stages=cfg.smem_kq_stages, + leading_byte_offset=2 * KT_LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sAinv_raw = cutlass.Array( + bars = make_gdn_bars(cfg) + tmem_base_slot = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=16) + sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + cumsumlog_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) + cumprod_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) + beta_raw = cutlass.Array(cutlass.Float32, cute.cosize(beta_smem_layout_staged), space=SMEM, alignment=128) + sTinv_raw = cutlass.Array( cfg.io_dtype, - cfg.ainv_cosize, + cfg.t_inv_cosize, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sAinv = SmemTile( - base=sAinv_raw.data_ptr().toint(), - elems_per_stage=(cfg.ainv_cosize // cfg.smem_ainv_stages) * bpe, - stages=cfg.smem_ainv_stages, + sTinv = SmemTile( + base=sTinv_raw.data_ptr().toint(), + elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages) * bpe, + stages=cfg.smem_t_inv_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - sQk_raw = cutlass.Array( + sA_raw = cutlass.Array( cfg.io_dtype, - cfg.qk_cosize, + cfg.a_cosize, space=cutlass.AddressSpace.smem, alignment=cfg.buffer_align_bytes, ) - sQk = SmemTile( - base=sQk_raw.data_ptr().toint(), - elems_per_stage=(cfg.qk_cosize // cfg.smem_qk_stages) * bpe, - stages=cfg.smem_qk_stages, + sA = SmemTile( + base=sA_raw.data_ptr().toint(), + elems_per_stage=(cfg.a_cosize // cfg.smem_a_stages) * bpe, + stages=cfg.smem_a_stages, leading_byte_offset=LEAD, stride_byte_offset=STRIDE, layout=SWZ, ) - # descriptor operands total 136 KB: V/O/H land past the sub-bank midpoint sV_raw = cutlass.Array( cfg.io_dtype, cfg.v_cosize, @@ -2694,29 +2403,6 @@ def _kernel( stride_byte_offset=STRIDE, layout=SWZ, ) - sO_raw = cutlass.Array( - cfg.io_dtype, - cfg.o_cosize, - space=cutlass.AddressSpace.smem, - alignment=cfg.buffer_align_bytes, - ) - sO = SmemTile( - base=sO_raw.data_ptr().toint(), - elems_per_stage=(cfg.o_cosize // cfg.smem_o_stages) * bpe, - stages=cfg.smem_o_stages, - leading_byte_offset=LEAD, - stride_byte_offset=STRIDE, - layout=SWZ, - ) - if cutlass.const_expr(cfg.enable_h): - sH_raw = cutlass.Array( - cfg.io_dtype, - cfg.h_cosize, - space=cutlass.AddressSpace.smem, - alignment=cfg.buffer_align_bytes, - ) - else: - sH_raw = None sCumsumlog = cute.make_tensor( cute.make_ptr(cutlass.Float32, cumsumlog_raw.data_ptr().toint(), mem_space=cute.AddressSpace.smem, assumed_align=128), cumsumlog_smem_layout_staged, @@ -2730,15 +2416,10 @@ def _kernel( beta_smem_layout_staged, ) - # ------------------------------------------------------------------ - # mbarrier init (all threads) - # ------------------------------------------------------------------ - for s in range(cfg.smem_q_stages): - bars.mb_q_ready[s].init() - bars.mb_q_done[s].init() - for s in range(cfg.smem_k_stages): - bars.mb_k_ready[s].init() - bars.mb_k_done[s].init() + # ---- mbarrier init (all threads) --------------------------------------------- + for s in range(cfg.smem_kq_stages): + bars.mb_kq_ready[s].init() + bars.mb_kq_done[s].init() for s in range(cfg.smem_v_stages): bars.mb_v_ready[s].init() bars.mb_v_done[s].init() @@ -2748,35 +2429,34 @@ def _kernel( for s in range(cfg.smem_beta_stages): bars.mb_beta_ready[s].init() bars.mb_beta_done[s].init() - for s in range(cfg.tmem_kv_acc_stages): - bars.mb_kv_acc_ready[s].init() - bars.mb_kv_acc_scale_done[s].init() + for s in range(cfg.tmem_state_acc_stages): + bars.mb_state_acc_ready[s].init() + bars.mb_state_acc_scale_done[s].init() for s in range(cfg.tmem_q_state_acc_stages): bars.mb_o_acc_ready[s].init() - bars.mb_o_acc_done[s].init() - bars.mb_o_state_scale_acc_ready[s].init() + bars.mb_o_final_acc_ready[s].init() bars.mb_o_state_scale_acc_done[s].init() for s in range(cfg.tmem_cg0_acc_stages): bars.mb_cg0_acc_ready[s].init() bars.mb_cg0_acc_done[s].init() - bars.mb_ks_ready[0].init() - bars.mb_nv_ready[0].init() - for s in range(cfg.smem_ainv_stages): - bars.mb_ainv_ready[s].init() - bars.mb_ainv_done[s].init() - for s in range(cfg.smem_qk_stages): - bars.mb_qk_ready[s].init() - bars.mb_qk_done[s].init() + bars.mb_k_state_acc_ready[0].init() + bars.mb_u_acc_ready[0].init() + for s in range(cfg.smem_t_inv_stages): + bars.mb_t_inv_ready[s].init() + bars.mb_t_inv_done[s].init() + for s in range(cfg.smem_a_stages): + bars.mb_a_ready[s].init() + bars.mb_a_done[s].init() for s in range(cfg.tmem_state_inp_stages): bars.mb_state_inp_ready[s].init() - for b in (bars.mb_vks_inp_ready, bars.mb_nv_inp_ready, bars.mb_decay_v_inp_ready): + for b in (bars.mb_y_inp_ready, bars.mb_u_inp_ready, bars.mb_decay_u_inp_ready): b[0].init() for s in range(cfg.smem_o_stages): bars.mb_o_tmastg_ready[s].init() bars.mb_o_tmastg_done[s].init() - for s in range(cfg.smem_h_stages): - bars.mb_h_tmastg_ready[s].init() - bars.mb_h_tmastg_done[s].init() + for s in range(cfg.smem_checkpoint_stages): + bars.mb_checkpoint_tmastg_ready[s].init() + bars.mb_checkpoint_tmastg_done[s].init() for s_ in range(cfg.sched_stages): bars.mb_sched_ready[s_].init() bars.mb_sched_done[s_].init() @@ -2785,12 +2465,10 @@ def _kernel( nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync() - # ------------------------------------------------------------------ - # 2. Warp specialization - each warp role owns its own scheduler loop - # ------------------------------------------------------------------ + # ---- warp specialization ----------------------------------------------------- if warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: - _compute0_warp( + compute0_warp_group( cfg, total_tiles, bidx, @@ -2798,20 +2476,20 @@ def _kernel( cu_seqlens, mWorkItems, tidx, - tmem_hold=tmem_hold, + tmem_base_slot=tmem_base_slot, scale=scale, sCumsumlog=sCumsumlog, sBeta=sBeta, - sAinv=sAinv, - sQk=sQk, - sH_raw=sH_raw, + sTinv=sTinv, + sA=sA, + sCheckpoint_raw=sCheckpoint_raw, checkpoint_every_n_tokens=checkpoint_every_n_tokens, sSched=sSched, bars=bars, ) if warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]: - _compute1_warp( + compute1_warp_group( cfg, total_tiles, bidx, @@ -2820,64 +2498,65 @@ def _kernel( mWorkItems, tidx, warp_idx=warp_idx, - tmem_hold=tmem_hold, + tmem_base_slot=tmem_base_slot, scale=scale, sV=sV, sCumsumlog=sCumsumlog, sCumprod=sCumprod, sBeta=sBeta, sO=sO, - sH_raw=sH_raw, - mS_init=mS_init, - mS_out=mS_out, + sCheckpoint_raw=sCheckpoint_raw, + mState_init=mState_init, + mState_out=mState_out, checkpoint_every_n_tokens=checkpoint_every_n_tokens, sSched=sSched, bars=bars, ) - elif warp_idx == cfg.mma_warp_id: - _mma0_warp( + elif warp_idx == cfg.load_gate_beta_warp_id: + gate_beta_warp( cfg, total_tiles, bidx, num_ctas, cu_seqlens, mWorkItems, - tmem_hold=tmem_hold, - sQ=sQ, - sK=sK, + tidx=tidx, + mGate=mGate, + mBeta=mBeta, + sCumsumlog=sCumsumlog, + sCumprod=sCumprod, + sBeta=sBeta, sSched=sSched, bars=bars, ) - elif warp_idx == cfg.mma_cg1_warp_id: - _mma1_warp( + elif warp_idx == cfg.mma_warp_id: + mma_warp( cfg, total_tiles, bidx, num_ctas, cu_seqlens, mWorkItems, - tmem_hold=tmem_hold, - sQ=sQ, - sK=sK, - sK_trans=sK_trans, - sAinv=sAinv, - sQk=sQk, + tmem_base_slot=tmem_base_slot, + sKQ=sKQ, + sKQ_trans=sKQ_trans, + sTinv=sTinv, + sA=sA, sSched=sSched, bars=bars, ) elif warp_idx == cfg.tma_qkv_warp_id: - _tmaldg_warp( + tmaldg_warp( cfg, total_tiles, bidx, num_ctas, cu_seqlens, mWorkItems, - sQ_raw=sQ_raw, - sK_raw=sK_raw, + sKQ_raw=sKQ_raw, sV_raw=sV_raw, desc_q_base=desc_q_base, desc_k_base=desc_k_base, @@ -2888,7 +2567,7 @@ def _kernel( ) if warp_idx == cfg.epilogue_warp_id: - _tmastg_warp( + tmastg_warp( cfg, total_tiles, bidx, @@ -2897,15 +2576,10 @@ def _kernel( mWorkItems, checkpoint_every_n_tokens=checkpoint_every_n_tokens, tidx=tidx, - mGate=mGate, - mBeta=mBeta, - sCumsumlog=sCumsumlog, - sCumprod=sCumprod, - sBeta=sBeta, sO_raw=sO_raw, - sH_raw=sH_raw, + sCheckpoint_raw=sCheckpoint_raw, desc_o_base=desc_o_base, - desc_h_base=desc_h_base, + desc_checkpoint_base=desc_checkpoint_base, sSched=sSched, bars=bars, ) @@ -2917,9 +2591,9 @@ class GdnCfg: The per-compile parameters (dtypes, GQA, state flags) are the ``cute.compile`` cache keys; the rest is derived from the module-global - ``CFG`` constants. ``_host`` stamps the shape-derived fields at trace - time. Passed ``cfg``-first (a ``cutlass.Constexpr``) into ``_host`` / - ``_kernel`` and every warp body. + ``CFG`` constants. ``host`` stamps the shape-derived fields at trace + time. Passed ``cfg``-first (a ``cutlass.Constexpr``) into ``host`` / + ``kernel`` and every warp body. """ io_dtype: Type[cutlass.Numeric] @@ -2929,23 +2603,20 @@ class GdnCfg: is_GQA: bool use_initial_state: bool store_final_state: bool - enable_h: bool - enable_o: bool - split_k: bool = False + enable_checkpoints: bool log_gate: bool = False dyn_sched: bool = False sched_stages: int = CFG.SMEM_SCHED_STAGES - # --- fixed constants stamped from CFG by build_cfg --- + # ---- fixed constants stamped from CFG by build_cfg --------------------------- b_t: int = CFG.B_T d_k: int = CFG.D_K d_v: int = CFG.D_V compute_group_0_warp_ids: Tuple[int, ...] = CFG.COMPUTE_GROUP_0_WARP_IDS compute_group_1_warp_ids: Tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS - mma_warp_id: int = CFG.MMA_WARP_ID - tma_qkv_warp_id: int = CFG.TMA_QKV_WARP_ID - mma_cg1_warp_id: int = CFG.MMA_CG1_WARP_ID load_gate_beta_warp_id: int = CFG.LOAD_GATE_BETA_WARP_ID + tma_qkv_warp_id: int = CFG.TMA_QKV_WARP_ID + mma_warp_id: int = CFG.MMA_WARP_ID epilogue_warp_id: int = CFG.EPILOGUE_WARP_ID num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 @@ -2954,7 +2625,7 @@ class GdnCfg: threads_per_cta: int = 0 cluster_shape_mnk: Tuple[int, int, int] = CFG.CLUSTER_SHAPE_MNK - # --- named barrier slots (ids 1-4; 0 is the CTA-wide sync) --- + # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ------------------- tmem_alloc_barrier_id: int = 1 tmem_alloc_barrier_threads: int = 0 inverse_barrier_id: int = 2 @@ -2962,42 +2633,42 @@ class GdnCfg: init_state_store_barrier_id: int = 4 init_state_store_barrier_threads: int = 0 - # --- SMEM / TMEM stage counts + TMEM column offsets --- - smem_q_stages: int = CFG.SMEM_Q_STAGES - smem_k_stages: int = CFG.SMEM_K_STAGES + # ---- SMEM / TMEM stage counts + TMEM column offsets -------------------------- + smem_kq_stages: int = CFG.SMEM_KQ_STAGES smem_v_stages: int = CFG.SMEM_V_STAGES - smem_ainv_stages: int = CFG.SMEM_AINV_STAGES - smem_qk_stages: int = CFG.SMEM_QK_STAGES + smem_t_inv_stages: int = CFG.SMEM_T_INV_STAGES + smem_a_stages: int = CFG.SMEM_A_STAGES smem_o_stages: int = CFG.SMEM_O_STAGES - smem_h_stages: int = 1 + smem_checkpoint_stages: int = 1 smem_gate_stages: int = CFG.SMEM_GATE_STAGES smem_beta_stages: int = CFG.SMEM_BETA_STAGES - tmem_kv_acc_stages: int = CFG.TMEM_KV_ACC_STAGES + tmem_state_acc_stages: int = CFG.TMEM_KV_ACC_STAGES tmem_q_state_acc_stages: int = CFG.TMEM_Q_STATE_ACC_STAGES tmem_state_inp_stages: int = CFG.TMEM_STATE_INP_STAGES tmem_cg0_acc_stages: int = CFG.TMEM_CG0_ACC_STAGES tmem_cg1_acc_stages: int = CFG.TMEM_CG1_ACC_STAGES - tmem_state_offset: int = 0 - tmem_q_state_offset: int = 0 + tmem_state_acc_offset: int = 0 + tmem_q_state_acc_offset: int = 0 tmem_state_inp_offset: int = 0 tmem_cg0_acc_offset: int = 0 tmem_cg1_acc_offset: int = 0 - tmem_inp_offset: int = 0 + tmem_y_decay_u_inp_offset: int = 0 buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES - # --- stamped by _host at trace time (shape-derived) --- - q_cosize: int = 0 - k_cosize: int = 0 + # ---- stamped by host at trace time (shape-derived) -------------------------- + kq_cosize: int = 0 v_cosize: int = 0 - ainv_cosize: int = 0 - qk_cosize: int = 0 + t_inv_cosize: int = 0 + a_cosize: int = 0 o_cosize: int = 0 - h_cosize: int = 0 - tma_q_bytes: int = 0 - tma_k_bytes: int = 0 + checkpoint_cosize: int = 0 + tma_kq_bytes: int = 0 tma_v_bytes: int = 0 tma_o_bytes: int = 0 n_heads_out: int = 0 + q_ratio: int = 1 + k_ratio: int = 1 + v_ratio: int = 1 def build_cfg( @@ -3008,15 +2679,12 @@ def build_cfg( is_GQA: bool, use_initial_state: bool, store_final_state: bool = True, - enable_h: bool = False, - enable_o: bool = True, - split_k: bool = False, + enable_checkpoints: bool = False, log_gate: bool = False, dyn_sched: bool = False, ) -> GdnCfg: """Build the per-compile ``GdnCfg`` (io_dtype ∈ {Float16, BFloat16}; - acc is always Float32). Fills the derived thread / barrier counts and - the TMEM column offsets.""" + acc is always Float32).""" if io_dtype not in (cutlass.Float16, cutlass.BFloat16): raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") cfg = GdnCfg( @@ -3027,104 +2695,50 @@ def build_cfg( is_GQA=is_GQA, use_initial_state=use_initial_state, store_final_state=store_final_state, - enable_h=enable_h, - enable_o=enable_o, - split_k=split_k, + enable_checkpoints=enable_checkpoints, log_gate=log_gate, dyn_sched=dyn_sched, ) - cfg.smem_h_stages = 1 - if enable_h: - # trim K/V lookahead so the 32 KB H buffer fits - cfg.smem_k_stages = 3 - cfg.smem_v_stages = 2 + cfg.smem_checkpoint_stages = 1 + if enable_checkpoints: + cfg.smem_kq_stages = 3 if not use_initial_state: - # fund the lightweight warps from CG1 (peeled chunk carries more cursors) cfg.num_regs_compute_group_1 = 232 cfg.num_regs_other = 48 n_cg0 = len(cfg.compute_group_0_warp_ids) n_cg1 = len(cfg.compute_group_1_warp_ids) cfg.threads_per_cta = cfg.threads_per_warp * (4 + n_cg0 + n_cg1) - # both MMA issuer warps join the TMEM alloc barrier - cfg.tmem_alloc_barrier_threads = cfg.threads_per_warp * (2 + n_cg0 + n_cg1) + cfg.tmem_alloc_barrier_threads = cfg.threads_per_warp * (1 + n_cg0 + n_cg1) cfg.inverse_barrier_threads = cfg.threads_per_warp * n_cg0 cfg.init_state_store_barrier_threads = cfg.threads_per_warp * n_cg1 - cfg.tmem_state_offset = 0 - cfg.tmem_q_state_offset = cfg.tmem_state_offset + cfg.tmem_kv_acc_stages * 128 - cfg.tmem_state_inp_offset = cfg.tmem_q_state_offset + cfg.tmem_q_state_acc_stages * 64 + cfg.tmem_state_acc_offset = 0 + cfg.tmem_q_state_acc_offset = cfg.tmem_state_acc_offset + cfg.tmem_state_acc_stages * 128 + cfg.tmem_state_inp_offset = cfg.tmem_q_state_acc_offset + cfg.tmem_q_state_acc_stages * 64 cfg.tmem_cg0_acc_offset = cfg.tmem_state_inp_offset + cfg.tmem_state_inp_stages * 64 cfg.tmem_cg1_acc_offset = cfg.tmem_cg0_acc_offset + cfg.tmem_cg0_acc_stages * 64 - cfg.tmem_inp_offset = cfg.tmem_cg1_acc_offset + cfg.tmem_cg1_acc_stages * 64 + cfg.tmem_y_decay_u_inp_offset = cfg.tmem_cg1_acc_offset + cfg.tmem_cg1_acc_stages * 64 return cfg -def get_workspace_size(B: int, HQ: int, HV: int): - HO = HQ if HQ >= HV else HV - return CFG.BYTES_PER_TENSORMAP * (5 * B * HO) + 128 +TENSORMAP_DESC_ARRAYS = 5 # per-batch runtime TMA descriptors: Q, K, V, O, checkpoints +TENSORMAP_STATIC_SLOTS = 0 # --------------------------------------------------------------------------- -def _check_cuda(err): - if err != cuda.CUresult.CUDA_SUCCESS: - raise RuntimeError(f"CUDA driver call failed: {err}") - - -def _data_ptr(t) -> int: - """Device address of a tensor-like (``data_ptr()`` or the CUDA array - interface).""" - fn = getattr(t, "data_ptr", None) - if fn is not None: - return fn() - return t.__cuda_array_interface__["data"][0] - - -def _device_sm_count() -> int: - """Multiprocessor count of the current device (runtime API: auto-inits - the primary context, so this works before any other CUDA call).""" - from cuda.bindings import runtime as _rt - - err, dev = _rt.cudaGetDevice() - if int(err) != 0: - raise RuntimeError(f"cudaGetDevice failed: {err}") - err, count = _rt.cudaDeviceGetAttribute(_rt.cudaDeviceAttr.cudaDevAttrMultiProcessorCount, dev) - if int(err) != 0: - raise RuntimeError(f"cudaDeviceGetAttribute failed: {err}") - return count - - -def _cutlass_io_dtype(dtype): - name = str(dtype) - if "bfloat16" in name: - return cutlass.BFloat16 - if "float16" in name or "half" in name: - return cutlass.Float16 - raise ValueError(f"Unsupported dtype {dtype}, expected bfloat16 or float16") - - -def _cutlass_state_dtype(dtype): - name = str(dtype) - if "bfloat16" in name: - return cutlass.BFloat16 - if "float32" in name: - return cutlass.Float32 - raise ValueError(f"Unsupported state dtype {dtype}, expected float32 or bfloat16") - - @functools.cache -def _get_compiled_cache( +def get_compiled_cache( io_dtype_str: str, state_dtype_str: str, + cu_dtype_str: str, HQ: int, HK: int, HV: int, is_GQA: bool, use_initial_state: bool, store_final_state: bool, - enable_h: bool, - enable_o: bool, - split_k: bool, + enable_checkpoints: bool, log_gate: bool, dyn_sched: bool, ): @@ -3138,9 +2752,7 @@ def compile( is_GQA: bool, use_initial_state: bool, store_final_state: bool, - enable_h: bool, - enable_o: bool, - split_k: bool = False, + enable_checkpoints: bool, log_gate: bool = False, dyn_sched: bool = False, *, @@ -3152,8 +2764,8 @@ def compile( beta_cute, o_cute, cu_seqlens_cute, - s_in_cute, - s_out_cute, + state_in_cute, + state_out_cute, work_items_cute=None, work_count_cute=None, sched_ctr_cute=None, @@ -3170,15 +2782,13 @@ def compile( is_GQA=is_GQA, use_initial_state=use_initial_state, store_final_state=store_final_state, - enable_h=enable_h, - enable_o=enable_o, - split_k=split_k, + enable_checkpoints=enable_checkpoints, log_gate=log_gate, dyn_sched=dyn_sched, ) return cute.compile( - _host, + host, cfg, q_cute, k_cute, @@ -3187,8 +2797,8 @@ def compile( beta_cute, o_cute, cu_seqlens_cute, - s_in_cute, - s_out_cute, + state_in_cute, + state_out_cute, work_items_cute, work_count_cute, sched_ctr_cute, @@ -3212,7 +2822,7 @@ def chunk_gdn_sm100( output_state, scale: float, checkpoint_every_n_tokens: int = 0, - output_h=None, + output_state_checkpoints=None, work_items=None, work_count=None, sched_ctr=None, @@ -3234,44 +2844,30 @@ def chunk_gdn_sm100( gate: ``(total_tokens, HO)`` float32, forget gate — raw linear alpha, or the natural-log decay when ``log_gate`` beta: ``(total_tokens, HO)`` float32, update gate - output: ``(total_tokens, HO, DK)`` float16/bfloat16, pre-allocated, - or None to skip the O output path entirely (state/H-only run; - requires output_h or output_state) + output: ``(total_tokens, HO, DK)`` float16/bfloat16, pre-allocated cu_seqlens: ``(num_seqs + 1,)`` int32 initial_state: ``(num_seqs, HO, DK, DK)`` float32/bfloat16, or None output_state: ``(num_seqs, HO, DK, DK)`` float32/bfloat16, or None scale: attention scale factor (must not be 0) - checkpoint_every_n_tokens: emit an H entry every N tokens (0 = off) - output_h: ``(total_h, HO, DK, DK)`` bfloat16, or None. H[j] is the + checkpoint_every_n_tokens: emit a checkpoint entry every N tokens (0 = off) + output_state_checkpoints: ``(total_checkpoints, HO, DK, DK)`` io dtype, or None. Entry j is the state after ``(j + 1) * N`` tokens, STRICTLY BEFORE the sequence - end -- the end-of-sequence state is only ``output_state`` (S, - fp32-capable). With ``N == B_T`` this is the per-chunk state + end -- the end-of-sequence state is only ``output_state`` + (fp32-capable). With ``N == B_T`` this is the per-chunk checkpoint series the backward pass consumes. - work_items: ``(max_items, 6)`` int32 split-K work-item table from - ``common/split_k.py``, or None for the one-tile-per-(b,h) - layout. With a table, each item computes chunks - ``[cstart, wend)`` and writes O/H only for ``[wstart, wend)``. - work_count: ``(1,)`` int32 device-side item count (required with - work_items) + work_items: ``(max_items, 8)`` int32 work-item table from + ``common/split_k.py`` (REQUIRED; an uncut table row is the whole + (b, h) sequence). Each item computes chunks ``[cstart, wend)`` + and writes O/checkpoints only for ``[wstart, wend)``. + work_count: ``(1,)`` int32 device-side item count (REQUIRED) log_gate: ``gate`` holds natural-log decay values; the gate warp skips its log2 (rescales by 1/ln2) instead of exponentiating upstream - workspace: ``(>= get_workspace_size(B, HQ, HV) // 8,)`` int64, + workspace: ``(>= tensormap_workspace_bytes(module, B) // 8,)`` int64, 128-byte aligned; holds the per-(b,h) TMA descriptors (contents managed here — reuse the same buffer across calls) stream: CUDA stream handle (``cudaStream_t`` as an int) """ - if checkpoint_every_n_tokens < 0: - raise ValueError("checkpoint_every_n_tokens must be non-negative") - if checkpoint_every_n_tokens > 0: - if checkpoint_every_n_tokens % CFG.B_T != 0: - raise ValueError(f"checkpoint_every_n_tokens must be a multiple of the chunk " f"size ({CFG.B_T}), got {checkpoint_every_n_tokens}") - if output_h is None: - raise ValueError("output_h must be provided when checkpoint_every_n_tokens > 0") - if str(output_h.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: - raise ValueError(f"output_h dtype must match the io dtype (fp32 state belongs to " f"output_state): got {output_h.dtype} with io {q.dtype}") - elif output_h is not None: - raise ValueError("output_h must be None when checkpoint_every_n_tokens == 0") HQ = q.shape[1] HV = v.shape[1] DK = q.shape[2] @@ -3279,20 +2875,11 @@ def chunk_gdn_sm100( is_GQA = HQ >= HV use_initial_state = initial_state is not None store_final_state = output_state is not None - enable_h = checkpoint_every_n_tokens > 0 - enable_o = output is not None - split_k = work_items is not None + enable_checkpoints = checkpoint_every_n_tokens > 0 + 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 - if not enable_o and not (enable_h or store_final_state): - raise ValueError("output=None requires output_h or output_state") - if split_k: - if work_count is None: - raise ValueError("work_count is required with work_items") - if enable_h and checkpoint_every_n_tokens != CFG.B_T: - raise ValueError(f"split-K H checkpoints require checkpoint_every_n_tokens == {CFG.B_T}, got {checkpoint_every_n_tokens}") - elif work_count is not None: - raise ValueError("work_count must be None without work_items") - io_dtype = _cutlass_io_dtype(q.dtype) + io_dtype = get_dtype(q.dtype) if initial_state is not None: state_dtype_src = initial_state.dtype @@ -3300,27 +2887,21 @@ def chunk_gdn_sm100( state_dtype_src = output_state.dtype else: state_dtype_src = None - state_dtype = _cutlass_state_dtype(state_dtype_src) if state_dtype_src is not None else cutlass.Float32 + state_dtype = get_dtype(state_dtype_src) if state_dtype_src is not None else cutlass.Float32 - ws_words = get_workspace_size(B, HQ, HV) // 8 - if workspace.shape[0] < ws_words: - raise ValueError(f"workspace too small: need {ws_words} int64 words, " f"got {workspace.shape[0]}") - if _data_ptr(workspace) % 128 != 0: - raise ValueError("workspace must be 128-byte aligned") cu_stream = cuda.CUstream(int(stream)) - cache = _get_compiled_cache( + cache = get_compiled_cache( str(q.dtype), str(state_dtype_src), + str(cu_seqlens.dtype), HQ, k.shape[1], HV, is_GQA, use_initial_state, store_final_state, - enable_h, - enable_o, - split_k, + enable_checkpoints, log_gate, dyn_sched, ) @@ -3336,30 +2917,25 @@ def chunk_gdn_sm100( gate_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) beta_cute = from_dlpack(beta, assumed_align=16) beta_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) - o_cute = None - if enable_o: - o_cute = from_dlpack(output, assumed_align=16) - o_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - cu_seqlens_cute = from_dlpack(cu_seqlens, assumed_align=4).mark_layout_dynamic() + o_cute = from_dlpack(output, assumed_align=16) + o_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + cu_seqlens_cute = from_dlpack(cu_seqlens, assumed_align=8 if str(cu_seqlens.dtype).endswith("int64") else 4).mark_layout_dynamic() - s_in_cute = None + state_in_cute = None if use_initial_state: - s_in_cute = from_dlpack(initial_state, assumed_align=16) - s_in_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=DK) + state_in_cute = from_dlpack(initial_state, assumed_align=16) + state_in_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=DK) - s_out_cute = None + state_out_cute = None if store_final_state: - s_out_cute = from_dlpack(output_state, assumed_align=16) - s_out_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=DK) + state_out_cute = from_dlpack(output_state, assumed_align=16) + state_out_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=DK) workspace_cute = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() - work_items_cute = None - work_count_cute = None - if split_k: - work_items_cute = from_dlpack(work_items, assumed_align=4) - work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) - work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() + work_items_cute = from_dlpack(work_items, assumed_align=16) + work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() sched_ctr_cute = None if dyn_sched: @@ -3371,12 +2947,10 @@ def chunk_gdn_sm100( is_GQA, use_initial_state, store_final_state, - enable_h, - enable_o, - split_k, + enable_checkpoints, log_gate, dyn_sched, - num_sm=_device_sm_count(), + num_sm=multiprocessor_count(current_device_id()), q_cute=q_cute, k_cute=k_cute, v_cute=v_cute, @@ -3384,8 +2958,8 @@ def chunk_gdn_sm100( beta_cute=beta_cute, o_cute=o_cute, cu_seqlens_cute=cu_seqlens_cute, - s_in_cute=s_in_cute, - s_out_cute=s_out_cute, + state_in_cute=state_in_cute, + state_out_cute=state_out_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, sched_ctr_cute=sched_ctr_cute, @@ -3397,11 +2971,8 @@ def chunk_gdn_sm100( compiled = cache["compiled"] - # The descriptors encode cu_seqlens' CONTENTS, which no key built from the - # buffers can track. The skip this replaces asked torch's _version counter, - # so it was sound for torch callers and silently stale for every other - # producer. Rebuilding unconditionally measured free: 131 vs 135 us of host - # time, 157 either way once the launches are waited on. + # desc build runs every execute by contract (cu contents are data; + # buffer pointers may change) — capture-safe, single tiny launch if "build_descs" not in cache: q_bc = from_dlpack(q, assumed_align=16) q_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) @@ -3409,18 +2980,17 @@ def chunk_gdn_sm100( k_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) v_bc = from_dlpack(v, assumed_align=16) v_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - o_bc = None - if enable_o: - o_bc = from_dlpack(output, assumed_align=16) - o_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - cu_bc = from_dlpack(cu_seqlens, assumed_align=4).mark_layout_dynamic() - s_bc = None - if enable_h: - s_bc = from_dlpack(output_h, assumed_align=16) - s_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + o_bc = from_dlpack(output, assumed_align=16) + o_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + cu_bc = from_dlpack(cu_seqlens, assumed_align=8 if str(cu_seqlens.dtype).endswith("int64") else 4).mark_layout_dynamic() + checkpoints_bc = None + cu_ckpt_bc = None + if enable_checkpoints: + checkpoints_bc = from_dlpack(output_state_checkpoints, assumed_align=16) + checkpoints_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) ws_bc = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() cache["build_descs"] = cute.compile( - _build_descs, + build_descs, io_dtype, CFG.B_T, q_bc, @@ -3428,8 +2998,8 @@ def chunk_gdn_sm100( v_bc, o_bc, cu_bc, - s_bc, - cutlass.Int32(checkpoint_every_n_tokens if enable_h else 1), + checkpoints_bc, + cutlass.Int32(checkpoint_every_n_tokens if enable_checkpoints else 1), ws_bc, cu_stream, options="--enable-tvm-ffi", @@ -3440,12 +3010,11 @@ def chunk_gdn_sm100( v, output, cu_seqlens, - output_h, - checkpoint_every_n_tokens if enable_h else 1, + output_state_checkpoints, + checkpoint_every_n_tokens if enable_checkpoints else 1, workspace, cu_stream, ) - compiled( q, k, diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_recompute_config.py b/python/cudnn/linear_attention/frost/kernel/gdn_recompute_config.py new file mode 100644 index 000000000..30675f5a2 --- /dev/null +++ b/python/cudnn/linear_attention/frost/kernel/gdn_recompute_config.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# This kernel is derived from cuDNN, NVIDIA Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Gated DeltaNet (GDN) Cutlass-primitives recompute (state/H-only) kernel config (fixed +compile-time constants; the per-compile attributes live on ``GdnCfg`` in the kernel file). + +Target arch: Blackwell SM100 (GB200) / SM103 (GB300). +""" + +from dataclasses import dataclass +from typing import Tuple + + +@dataclass(frozen=True) +class Cfg: + # --- tile shape --- + B_T: int = 64 # chunk size / token tile (the mma N or K of every GEMM) + D_K: int = 128 # key head dim (contraction of GEMMs 1/3, M of GEMM 7) + D_V: int = 128 # value head dim (M of GEMMs 3/5, N of GEMM 7) + + # --- TMA descriptor pool --- + + # --- warp assignments (12 warps total) --- + COMPUTE_GROUP_0_WARP_IDS: Tuple[int, ...] = (0, 1, 2, 3) # T-pairwise / kk_epi / inverse + COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (4, 5, 6, 7) # kv_decay_v / v-k*state / epi ops + LOAD_GATE_BETA_WARP_ID: int = 8 # gate/beta chunk loads + TMEM lifecycle + TMA_KV_WARP_ID: int = 9 + MMA_WARP_ID: int = 10 # sole tcgen05 issuer: KK pairs + KS/U/KV per chunk + EPILOGUE_WARP_ID: int = 11 + + # --- register split --- + NUM_REGS_COMPUTE_GROUP_0: int = 224 + NUM_REGS_COMPUTE_GROUP_1: int = 256 + NUM_REGS_OTHER: int = 24 + + THREADS_PER_WARP: int = 32 + + CLUSTER_SHAPE_MNK: Tuple[int, int, int] = (1, 1, 1) + + # --- SMEM stage counts --- + SMEM_SCHED_STAGES: int = 2 + SMEM_KQ_STAGES: int = 4 + SMEM_V_STAGES: int = 2 + SMEM_T_INV_STAGES: int = 3 + SMEM_GATE_STAGES: int = 3 + SMEM_BETA_STAGES: int = 3 + + # --- TMEM stage counts --- + TMEM_KV_ACC_STAGES: int = 1 + TMEM_STATE_INP_STAGES: int = 1 + TMEM_CG0_ACC_STAGES: int = 2 + TMEM_CG1_ACC_STAGES: int = 1 + + BUFFER_ALIGN_BYTES: int = 1024 + + +CFG = Cfg() diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py new file mode 100644 index 000000000..4ff922c18 --- /dev/null +++ b/python/cudnn/linear_attention/frost/kernel/gdn_recompute_f16.py @@ -0,0 +1,2604 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# This kernel is derived from cuDNN, NVIDIA Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Chunked Gated Delta Net (GDN) recompute (state/checkpoint-only) kernel for Blackwell SM100 +(Cutlass primitives): the prefill state/checkpoint pipeline with the Q/O path removed. + +Algorithm overview (per chunk c, tokens [cC, (c+1)C)): + Inputs : K[BT,DK], V[BT,DV], Gate[BT] (scalar gate), Beta[BT] (scalar LR) + State : S_prev[DK,DV] (recurrent state, held in TMEM) + + Preprocessing (compute warp group 0): + cumsumlog[t] = sum_{l=0}^{t} log(Gate_l) cumulative log of gates + cumprod[t] = exp(cumsumlog[t]) cumulative product of gates + T_pairwise[i,j] = cumprod[i] / cumprod[j] (i>=j) inter-token transfer weights + (stored in registers; 128 regs/thread) + + GEMM 1 - KK : W_kk[BT,BT] = K @ K^T (lower-triangular intra scores) + GEMM 3 - K*state : KS[BT,DV] = K @ S_prev (key applied to state) + GEMM 5 - U : U[BT,DV] = T_inv @ Y (corrected value vectors) + where T_inv = (I + M_kk)^{-1}, M_kk[i,j] = T[i,j]*Beta[i]*W_kk[i,j] (lower-tri, hierarchical blockwise inverse) + GEMM 7 - KV update : S_upd[DK,DV] = K^T @ (decay .* U) (state update, BT contraction) + where Y[BT,DV] = V - KS (delta rule residuals, after decay) + + Epilogue: + S_next = cumprod[BT-1] * S_prev + S_upd (update state in TMEM) + +Chunks run in PAIRS (CG0 warp halves invert chunk 0 / chunk 1 in parallel); +odd counts pad with a neutral zero-filled chunk. + +SMEM layout (stage counts live in gdn_recompute_config.py; +enable_checkpoints compiles trim K stages to fit the checkpoint buffer): + Buffer Size (B) Stages + K (two-box stage) 32768 4 + V 16384 3 + T_inv 8192 2 + checkpoint staging DK*DV*2 1 <-- enable_checkpoints only + cumsumlog / cumprod / Beta 256 3 + sched ticket ring 4 2 <-- dyn_sched publish ring + +TMEM layout (512 columns): + Buffer Cols + state 128 <-- DKxDV fp32 = 128x128x4B + state inp 64 <-- fp16 state staging (GEMM 3 A operand) + cg0 shared acc 128 <-- 2-stage ring: KK0/KK1 + cg1 shared acc 64 <-- 1-stage ring: KS then U + Y / decayed-U inp 64 <-- slot 0 = Y (V - K*state), slot 1 = decayed U (b16) + +Warp assignments (12 warps = 384 threads): + warps 0-3 : compute group 0 - T-pairwise x2, KK_epi x2, pair inverse + warps 4-7 : compute group 1 - state restage/rescale, Y = V - K*state, + U epilogue + warp 8 : Gate/Beta loads + warp 9 : TMA load warp - loads K, V + warp 10 : MMA warp - fused KK pairs + K*state/U/KV per chunk; + TMEM lifecycle + warp 11 : epilogue warp - checkpoint TMA stores +""" + +import functools +from dataclasses import dataclass +from typing import NamedTuple, Optional, Type, Tuple + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.experimental.primitives as nvvm +import cutlass.experimental.cuda.tensor_map as tma +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import min + +from ..common.thd import emit_checkpoint_seq_descs, emit_seq_descs, TENSOR_MAP_QWORDS +from ..common.split_k import decode_work_item +from ..common.host import get_dtype +from cudnn.frost.buffers import current_device_id, data_ptr +from cudnn.frost.device import multiprocessor_count + +RCP_LN2 = 1.4426950408889634 # 1/ln(2): natural-log gates -> the kernel's log2 domain +from cudnn.frost.tile_dsl.barrier import ( + MBarrier, + Producer, + PipelineState, + advance, +) +from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, tma_slice_runtime_desc +from cudnn.frost.tile_dsl.mma import mma_ss, mma_step_k8, mma_ts_step, mma_step +from cudnn.frost.tile_dsl.pointwise import fadd2, fp32_to_fp16, f16x2_to_f32, fmul2, opaque_f32_zero, sub_f16x2 +from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_xor_128b +from cudnn.frost.tile_dsl.tma import ( + tma_load_tile, + tma_store_tile, + tma_store_commit, + tma_store_wait, + tma_tensormap_acquire, +) +from .gdn_recompute_config import CFG + + +class GdnBars(NamedTuple): + """GDN pipeline mbarrier inventory. + + Every pipeline is a ``_ready``/``_done`` MBarrier pair over one ring: a + slot is acquired for filling by waiting ``_done`` and committed by + arriving ``_ready``; the reading side waits ``_ready`` and releases the + slot by arriving ``_done``. + """ + + mb_kq_ready: MBarrier + mb_kq_done: MBarrier + mb_v_ready: MBarrier + mb_v_done: MBarrier + + mb_gate_ready: MBarrier + mb_gate_done: MBarrier + mb_beta_ready: MBarrier + mb_beta_done: MBarrier + + mb_state_acc_ready: MBarrier + mb_state_acc_scale_done: MBarrier + mb_cg0_acc_ready: MBarrier + mb_cg0_acc_done: MBarrier + mb_k_state_acc_ready: MBarrier + mb_u_acc_ready: MBarrier + + mb_t_inv_ready: MBarrier + mb_t_inv_done: MBarrier + mb_state_inp_ready: MBarrier + mb_y_inp_ready: MBarrier + mb_decay_u_inp_ready: MBarrier + + mb_checkpoint_tmastg_ready: MBarrier + mb_checkpoint_tmastg_done: MBarrier + + mb_tmem_done: MBarrier + mb_sched_ready: MBarrier + mb_sched_done: MBarrier + + +def make_gdn_bars(cfg) -> GdnBars: + """GdnBars factory. MUST be called from inside ``kernel`` (allocates SMEM).""" + ONE_LANE = 1 + MMA_ARRIVERS = len([cfg.mma_warp_id]) + KQ_RELEASE_SITES = 1 + GATE_WARP = cfg.threads_per_warp * len([cfg.load_gate_beta_warp_id]) + EPI_WARP = cfg.threads_per_warp * len([cfg.epilogue_warp_id]) + CG0_THREADS = cfg.threads_per_warp * len(cfg.compute_group_0_warp_ids) + CG1_THREADS = cfg.threads_per_warp * len(cfg.compute_group_1_warp_ids) + CG0_PLUS_CG1 = CG0_THREADS + CG1_THREADS + + def alloc(n): + return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=16) + + return GdnBars( + mb_kq_ready=MBarrier(alloc(cfg.smem_kq_stages), stages=cfg.smem_kq_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), + mb_kq_done=MBarrier(alloc(cfg.smem_kq_stages), stages=cfg.smem_kq_stages, init_count=KQ_RELEASE_SITES, producer=Producer.MMA_COMMIT), + mb_v_ready=MBarrier(alloc(cfg.smem_v_stages), stages=cfg.smem_v_stages, init_count=ONE_LANE, producer=Producer.TMA_LOAD), + mb_v_done=MBarrier(alloc(cfg.smem_v_stages), stages=cfg.smem_v_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_gate_ready=MBarrier(alloc(cfg.smem_gate_stages), stages=cfg.smem_gate_stages, init_count=GATE_WARP, producer=Producer.THREAD), + mb_gate_done=MBarrier(alloc(cfg.smem_gate_stages), stages=cfg.smem_gate_stages, init_count=CG0_PLUS_CG1, producer=Producer.THREAD), + mb_beta_ready=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=GATE_WARP, producer=Producer.THREAD), + mb_beta_done=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_state_acc_ready=MBarrier(alloc(cfg.tmem_state_acc_stages), stages=cfg.tmem_state_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_state_acc_scale_done=MBarrier(alloc(cfg.tmem_state_acc_stages), stages=cfg.tmem_state_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_cg0_acc_ready=MBarrier(alloc(cfg.tmem_cg0_acc_stages), stages=cfg.tmem_cg0_acc_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_cg0_acc_done=MBarrier(alloc(cfg.tmem_cg0_acc_stages), stages=cfg.tmem_cg0_acc_stages, init_count=CG0_THREADS // 2, producer=Producer.THREAD), + mb_k_state_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_t_inv_ready=MBarrier(alloc(cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, init_count=CG0_THREADS, producer=Producer.THREAD), + mb_t_inv_done=MBarrier(alloc(cfg.smem_t_inv_stages), stages=cfg.smem_t_inv_stages, init_count=MMA_ARRIVERS, producer=Producer.MMA_COMMIT), + mb_state_inp_ready=MBarrier(alloc(cfg.tmem_state_inp_stages), stages=cfg.tmem_state_inp_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_decay_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_checkpoint_tmastg_ready=MBarrier( + alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=len(cfg.compute_group_1_warp_ids), producer=Producer.THREAD + ), + mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=EPI_WARP, producer=Producer.THREAD), + mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), + mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=11, producer=Producer.THREAD), + ) + + +@cute.jit +def invert_diagonal_NxN(cfg, base_int, d, tidx, N: int = 8): + """Gauss-Jordan inversion of one diagonal NxN block in-place (f16 SMEM).""" + tidx_in_group = tidx % N + BT = cfg.b_t + + row_lin_base = (d * N + tidx_in_group) * BT + d * N + row_phys = swizzle_lin_128b(row_lin_base, row_stride_log2=6) + row_ptr = ( + cute.make_ptr( + cfg.io_dtype, + base_int, + mem_space=cute.AddressSpace.smem, + assumed_align=cfg.buffer_align_bytes, + ) + + row_phys + ) + + row = [(row_ptr + j).load().to(cutlass.Float32) for j in range(N)] + for i in cutlass.range_constexpr(N): + row[i] = cutlass.Float32(1.0) if tidx_in_group == i else row[i] + for src_row in cutlass.range_constexpr(N - 1): + row_scale = -row[src_row] + for i in cutlass.range_constexpr(src_row): + shfl_val = nvvm.shfl_sync(0xFFFFFFFF, row[i], src_row, 0b1100000011111, kind=nvvm.Shfl.IDX) + row[i] = row[i] + row_scale * shfl_val if tidx_in_group > src_row else row[i] + row[src_row] = row_scale if tidx_in_group > src_row else row[src_row] + + for j in cutlass.range_constexpr(N): + (row_ptr + j).store(row[j].to(cfg.io_dtype)) + + +@cute.jit +def blockwise_diagonal_8x8_to_16x16(cfg, base_int, d0, lane_id): + """Off-diagonal correction 8x8 -> 16x16 (C <- -D^{-1} C A^{-1}).""" + bpe = cfg.io_dtype.width // 8 + lds1 = (lane_id % 8) * 64 + d = nvvm.ldmatrix( + cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + 8 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + 1, + nvvm.MMALayout.ROW, + ) + c = nvvm.ldmatrix( + cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + 1, + nvvm.MMALayout.COL, + ) + + # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + c_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(4): + c_regs[i] = cutlass.Float32(0.0) + mma_step_k8(c_regs, [d, d], [c], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) + for i in cutlass.range_constexpr(4): + c_regs[i] = -c_regs[i] + a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(2)] + + # ---- C = T @ A^{-1} ---------------------------------------------------------- + ai = nvvm.ldmatrix( + cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + 1, + nvvm.MMALayout.COL, + ) + o_regs = cutlass.Array(cutlass.Float32, 4, alignment=16, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(4): + o_regs[i] = cutlass.Float32(0.0) + mma_step_k8(o_regs, a_pack, [ai], k_step=0, M=16, N=8, ab_dtype=cfg.io_dtype) + o_pack = fp32_to_fp16(o_regs[0], o_regs[1], dtype=cfg.io_dtype) + + # ---- store corrected C ------------------------------------------------------- + nvvm.stmatrix( + cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 8) * 64 + d0 + lds1, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + o_pack, + nvvm.MMALayout.ROW, + ) + + +@cute.jit +def blockwise_diagonal_16x16_to_32x32(cfg, base_int, d0, lane_id): + """Off-diagonal correction 16x16 -> 32x32.""" + bpe = cfg.io_dtype.width // 8 + lds4 = (lane_id % 16) * 64 + (lane_id // 16) * 8 + d = list( + nvvm.ldmatrix( + cutlass.inttoptr( + base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + 16 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16 + ), + 4, + nvvm.MMALayout.ROW, + ) + ) + c = list( + nvvm.ldmatrix( + cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + 4, + nvvm.MMALayout.COL, + ) + ) + + # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + c_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(8): + c_regs[i] = cutlass.Float32(0.0) + mma_step(c_regs, d, c, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) + for i in cutlass.range_constexpr(8): + c_regs[i] = -c_regs[i] + a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + + # ---- C = T @ A^{-1} ---------------------------------------------------------- + ai = list( + nvvm.ldmatrix( + cutlass.inttoptr(base_int + swizzle_lin_128b(d0 * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + 4, + nvvm.MMALayout.COL, + ) + ) + o_regs = cutlass.Array(cutlass.Float32, 8, alignment=16, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(8): + o_regs[i] = cutlass.Float32(0.0) + mma_step(o_regs, a_pack, ai, k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype) + o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(4)] + + # ---- store corrected C ------------------------------------------------------- + nvvm.stmatrix( + cutlass.inttoptr(base_int + swizzle_lin_128b((d0 + 16) * 64 + d0 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + o_pack, + nvvm.MMALayout.ROW, + ) + + +@cute.jit +def blockwise_diagonal_32x32_to_64x64(cfg, base_int, warp_id, lane_id): + """Off-diagonal correction 32x32 -> 64x64 (2 warps, one 16-row M-band each).""" + band = warp_id % 2 + bpe = cfg.io_dtype.width // 8 + lds4 = (lane_id % 16) * 64 + (lane_id // 16) * 8 + a_frags = [] + for vs in cutlass.range_constexpr(2): + a_frags += list( + nvvm.ldmatrix( + cutlass.inttoptr( + base_int + swizzle_lin_128b((32 + band * 16) * 64 + 32 + vs * 16 + lds4, row_stride_log2=6) * bpe, + cutlass.AddressSpace.smem, + cutlass.BFloat16, + ), + 4, + nvvm.MMALayout.ROW, + ) + ) + b_frags = [] + for vs in cutlass.range_constexpr(4): + b_frags += list( + nvvm.ldmatrix( + cutlass.inttoptr( + base_int + swizzle_lin_128b((32 + (vs // 2) * 16) * 64 + (vs % 2) * 16 + lds4, row_stride_log2=6) * bpe, + cutlass.AddressSpace.smem, + cutlass.BFloat16, + ), + 4, + nvvm.MMALayout.COL, + ) + ) + + # ---- T = -(D^{-1} @ C) ------------------------------------------------------- + c_regs = cutlass.Array(cutlass.Float32, 16, alignment=16, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(16): + c_regs[i] = cutlass.Float32(0.0) + for ks in cutlass.range_constexpr(2): + mma_step(c_regs, a_frags, b_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + for i in cutlass.range_constexpr(16): + c_regs[i] = -c_regs[i] + a_pack = [fp32_to_fp16(c_regs[2 * j], c_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] + + # ---- C = T @ A^{-1} ---------------------------------------------------------- + ai_frags = [] + for vs in cutlass.range_constexpr(4): + ai_frags += list( + nvvm.ldmatrix( + cutlass.inttoptr( + base_int + swizzle_lin_128b(((vs // 2) * 16) * 64 + (vs % 2) * 16 + lds4, row_stride_log2=6) * bpe, + cutlass.AddressSpace.smem, + cutlass.BFloat16, + ), + 4, + nvvm.MMALayout.COL, + ) + ) + o_regs = cutlass.Array(cutlass.Float32, 16, alignment=16, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(16): + o_regs[i] = cutlass.Float32(0.0) + for ks in cutlass.range_constexpr(2): + mma_step(o_regs, a_pack, ai_frags[ks * 8 : ks * 8 + 8], k_step=ks, M=16, N=32, ab_dtype=cfg.io_dtype) + o_pack = [fp32_to_fp16(o_regs[2 * j], o_regs[2 * j + 1], dtype=cfg.io_dtype) for j in range(8)] + + # ---- store corrected C ------------------------------------------------------- + nvvm.barrier_cta_sync_aligned( + cfg.inverse_barrier_id, + thread_count=cfg.inverse_barrier_threads, + ) + nvvm.stmatrix( + cutlass.inttoptr(base_int + swizzle_lin_128b((32 + band * 16) * 64 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + o_pack[0:4], + nvvm.MMALayout.ROW, + ) + nvvm.stmatrix( + cutlass.inttoptr(base_int + swizzle_lin_128b((32 + band * 16) * 64 + 16 + lds4, row_stride_log2=6) * bpe, cutlass.AddressSpace.smem, cutlass.BFloat16), + o_pack[4:8], + nvvm.MMALayout.ROW, + ) + + +# ---- Dynamic tile scheduler ------------------------------------------------------ + + +@cute.jit +def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): + """TMA-LDG-warp side: pull the next tile off the global ticket, publish it.""" + if cutlass.const_expr(cfg.dyn_sched): + bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) + if nvvm.elect_sync(): + fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sSched[sched_state.idx] = num_ctas + fetched + nvvm.bar_warp_sync(cute.arch.FULL_MASK) + next_tile = sSched[sched_state.idx] + if nvvm.elect_sync(): + bars.mb_sched_ready[sched_state.idx].arrive() + return next_tile, advance(sched_state, cfg.sched_stages) + return tile_idx + num_ctas, sched_state + + +@cute.jit +def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): + """Consumer side: read the TMA-LDG warp's published next tile.""" + if cutlass.const_expr(cfg.dyn_sched): + bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) + next_tile = sSched[sched_state.idx] + if nvvm.elect_sync(): + bars.mb_sched_done[sched_state.idx].arrive() + return next_tile, advance(sched_state, cfg.sched_stages) + return tile_idx + num_ctas, sched_state + + +@cute.jit +def tmastg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + checkpoint_every_n_tokens, + tidx, + sCheckpoint_raw, + desc_checkpoint_base, + sSched, + bars, +): + """Epilogue warp role (warp 11): persistent scheduler loop issuing the + per-chunk checkpoint TMA stores.""" + elect_one = nvvm.elect_sync() + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + lidx = tidx % cfg.threads_per_warp + sched_state = PipelineState.start(phase=0) + tile_idx = cutlass.Int32(bidx) + + if cutlass.const_expr(cfg.enable_checkpoints): + checkpoint_granu = 64 + sCheckpoint_tma = SmemTile( + base=sCheckpoint_raw, + elems_per_stage=(cfg.checkpoint_cosize // cfg.smem_checkpoint_stages), + stages=cfg.smem_checkpoint_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=cfg.d_v // checkpoint_granu, + tma_granu_elems=checkpoint_granu, + tma_subtile_stride_elems=cfg.d_k * checkpoint_granu, + ) + checkpoint_store_cnt = cutlass.Int32(0) + ckpt_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + heads_out = cutlass.Int32(cfg.n_heads_out) + desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) + + 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 + if elect_one: + tma_tensormap_acquire(desc_checkpoint_slot) + + if n_local > 0: + for local_idx in cutlass.range(n_padded): + 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 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) + tma_store_tile(sCheckpoint_tma[checkpoint_stage], checkpoint_slice, acquire=False) + tma_store_commit() + checkpoint_coord += 1 + did_checkpoint = cutlass.Int32(1) + checkpoint_mod = checkpoint_mod + cutlass.Int32(1) + checkpoint_mod = cutlass.Int32(0) if checkpoint_mod == ckpt_chunks else checkpoint_mod + + if did_checkpoint == 1: + tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() + checkpoint_store_cnt = checkpoint_store_cnt + 1 + + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def gate_beta_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + tidx, + mGate, + mBeta, + sCumsumlog, + sCumprod, + sBeta, + sSched, + bars, +): + """Gate/Beta producer (warp 8): persistent scheduler loop + the + cumsum/cumprod/Beta chunk loads.""" + + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + gate_index = PipelineState.start(phase=1) + beta_index = PipelineState.start(phase=1) + lidx = tidx % cfg.threads_per_warp + + 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 + if n_local > 0: + for local_idx in cutlass.range(n_padded): + # ---- Gate load: GMEM -> SMEM (OOB neutral) ----------------------- + chunk_idx = cstart + local_idx + n_cols = cfg.b_t // cfg.threads_per_warp + chunk_offset = batch_start + chunk_idx * cfg.b_t + gGateSeq = mGate[None, head_idx] + gBeta = cute.domain_offset((chunk_offset,), mBeta[None, head_idx]) + + gate_idx = gate_index.idx + gate_phase = gate_index.phase + gate_index = advance(gate_index, cfg.smem_gate_stages) + + pos_valid = [None] * n_cols + gate_vals = [cutlass.Float32(0.0)] * n_cols + oob_neutral = cutlass.Float32(0.0) if cutlass.const_expr(cfg.log_gate) else cutlass.Float32(1.0) + for col in cutlass.range_constexpr(n_cols): + tok = chunk_offset + lidx + col * cfg.threads_per_warp + pos_valid[col] = cute.elem_less(tok, batch_end) + tok_clamped = min(tok, batch_end - 1) + gate_vals[col] = gGateSeq[tok_clamped] if pos_valid[col] else oob_neutral + + if cutlass.const_expr(cfg.log_gate): + for col in cutlass.range_constexpr(n_cols): + gate_vals[col] = gate_vals[col] * cutlass.Float32(RCP_LN2) + else: + for col in cutlass.range_constexpr(n_cols): + gate_vals[col] = cute.math.log2(gate_vals[col] + 1e-10, fastmath=True) + for offset in [1, 2, 4, 8, 16]: + for col in cutlass.range_constexpr(n_cols): + n = nvvm.shfl_sync(0xFFFFFFFF, gate_vals[col], offset, 0, kind=nvvm.Shfl.UP) + if lidx >= offset: + gate_vals[col] = gate_vals[col] + n + for col in cutlass.range_constexpr(1, n_cols): + last_v = nvvm.shfl_sync( + 0xFFFFFFFF, + gate_vals[col - 1], + cfg.threads_per_warp - 1, + cfg.threads_per_warp - 1, + kind=nvvm.Shfl.IDX, + ) + gate_vals[col] += last_v + + bars.mb_gate_done[gate_idx].wait(gate_phase) + for col in cutlass.range_constexpr(n_cols): + pos = lidx + col * cfg.threads_per_warp + sCumsumlog[pos, 0, gate_idx] = gate_vals[col] + sCumprod[pos, 0, gate_idx] = cute.math.exp2(gate_vals[col], fastmath=True) + + bars.mb_gate_ready[gate_idx].arrive() + + # ---- Beta load: GMEM -> SMEM (per-element cp.async) -------------------------- + beta_idx = beta_index.idx + bars.mb_beta_done[beta_idx].wait(beta_index.phase) + beta_index = advance(beta_index, cfg.smem_beta_stages) + for col in cutlass.range_constexpr(n_cols): + pos = lidx + col * cfg.threads_per_warp + src = gBeta.iterator + gBeta.layout((pos,)) + dst = sBeta.iterator + sBeta.layout((pos, 0, beta_idx)) + cp_size = cutlass.Int32(4) * cutlass.Int32(pos_valid[col]) + nvvm.cp_async_shared_global(dst, src, 4, nvvm.LoadCacheModifier.CA, cp_size=cp_size) + nvvm.cp_async_mbarrier_arrive(bars.mb_beta_ready[beta_idx].smem_ptr, noinc=True) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + for _ in range(cfg.smem_gate_stages): + bars.mb_gate_done[gate_index.idx].wait(gate_index.phase) + gate_index = advance(gate_index, cfg.smem_gate_stages) + for _ in range(cfg.smem_beta_stages): + bars.mb_beta_done[beta_index.idx].wait(beta_index.phase) + beta_index = advance(beta_index, cfg.smem_beta_stages) + + +@cute.jit +def mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + tmem_hold, + sKQ, + sKQ_trans, + sTinv, + sSched, + bars, +): + """MMA issuer role (warp 10): persistent scheduler loop issuing every + tcgen05 GEMM.""" + elect_one = nvvm.elect_sync() + + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + kv_acc_index = PipelineState.start(phase=1) + kq_index = PipelineState.start(phase=0) + cg0_acc_index = PipelineState.start(phase=1) + kq_fused_index = PipelineState.start(phase=0) + tinv_index = PipelineState.start(phase=0) + state_inp_index = PipelineState.start(phase=0) + y_inp_ready = PipelineState.start(phase=0) + decay_u_inp_ready = PipelineState.start(phase=0) + + nvvm.tcgen05_alloc(tmem_hold, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) + nvvm.barrier_cta_sync_aligned( + cfg.tmem_alloc_barrier_id, + thread_count=cfg.tmem_alloc_barrier_threads, + ) + + # ---- chunk-invariant GEMM descriptors ---------------------------------------- + idesc_kk = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.b_t, + m_dim=2 * cfg.b_t, + ) + bmm_kk_desc = MmaDesc( + M=2 * cfg.b_t, + N=cfg.b_t, + K=cfg.d_k, + bpe_a=cfg.io_dtype.width // 8, + bpe_b=cfg.io_dtype.width // 8, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_kk, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bpe = cfg.io_dtype.width // 8 + idesc_k_state = 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_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, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_u_ts = 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_u_ts_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + atranspose=False, + cta_group=1, + idesc=idesc_u_ts, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_kv = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.d_v, + m_dim=cfg.d_k, + b_major=1, + ) + bmm_kv_desc = MmaDesc( + M=cfg.d_k, + N=cfg.d_v, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + atranspose=False, + cta_group=1, + idesc=idesc_kv, + kind=nvvm.Tcgen05MMAKind.F16, + ) + KQ_SEG = (2 * cfg.b_t * 64 * bpe) >> 4 + KQ_BOX = (cfg.b_t * 64 * bpe) >> 4 + KQ_HALF_K = (cfg.d_k // 16) // 2 + KQ_A_HALF = KQ_HALF_K * bmm_k_state_desc.tmem_advance_A + + KV_ACC_STAGE_COLS = cfg.d_v + STATE_INP_STAGE_COLS = cfg.d_k // 2 + INP_SLOT_COLS = cfg.b_t // 2 + + tmem_base = tmem_hold.load() + tmem_cg0_acc_col_f = tmem_base + cfg.tmem_cg0_acc_offset + tmem_state_col = tmem_base + cfg.tmem_state_acc_offset + tmem_state_inp_col = tmem_base + cfg.tmem_state_inp_offset + tmem_inp_col = tmem_base + cfg.tmem_y_decay_u_inp_offset + y_inp_ptr = nvvm.make_tmem_ptr(tmem_inp_col, cutlass.Int8) + 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: + 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) + kqf_idx = kq_fused_index.idx + 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) + if elect_one: + bars.mb_cg0_acc_ready[member0_acc_idx].arrive(cta_group=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) + kqf_idx = kq_fused_index.idx + 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() + 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) + 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 + if cutlass.const_expr(cfg.use_initial_state): + if local_idx == 0: + if elect_one: + bars.mb_state_acc_ready[kv_acc_index.idx].arrive(cta_group=1) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) + have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_idx > 0 + + kq_idx = kq_index.idx + member = local_idx & 1 + state_inp_idx = state_inp_index.idx + tinv_idx = tinv_index.idx + kv_acc_idx = kv_acc_index.idx + kq_member_off = member * KQ_BOX + desc_k = sKQ[kq_idx].desc() + kq_member_off + desc_tinv = sTinv[tinv_idx].desc() + desc_kt = sKQ_trans[kq_idx].desc() + kq_member_off + state_a_ptr = nvvm.make_tmem_ptr(tmem_state_inp_col + state_inp_idx * STATE_INP_STAGE_COLS, cutlass.Int8) + state_acc_ptr = nvvm.make_tmem_ptr(tmem_state_col + kv_acc_idx * KV_ACC_STAGE_COLS, cutlass.Float32) + + kq_index = advance(kq_index, cfg.smem_kq_stages) + + # ---- KK pair lookahead (member 1) = K(S) @ K^T ----------------------- + if member == 1: + if (local_idx >> 1) + 1 < n_pairs: + 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) + kqf_idx = kq_fused_index.idx + 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() + 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) + if elect_one: + bars.mb_cg0_acc_ready[member1_acc_idx].arrive(cta_group=1) + + # ---- K*state = state(T) @ K^T (GEMM 3) ---------------------------------------- + if have_state: + bars.mb_state_inp_ready[state_inp_idx].wait(state_inp_index.phase) + state_inp_index = advance(state_inp_index, cfg.tmem_state_inp_stages) + + for k in cutlass.range_constexpr(KQ_HALF_K): + mma_ts_step(bmm_k_state_desc, state_a_ptr, desc_k, k_state_acc_ptr, k, cutlass.Boolean(k > 0)) + for k in cutlass.range_constexpr(KQ_HALF_K): + mma_ts_step(bmm_k_state_desc, state_a_ptr.subview(KQ_A_HALF), desc_k + KQ_SEG, k_state_acc_ptr, k, cutlass.Boolean(True)) + if elect_one: + bars.mb_k_state_acc_ready[0].arrive(cta_group=1) + + # ---- U = Y(T) @ T_inv (GEMM 5) --------------------------------------------- + bars.mb_t_inv_ready[tinv_idx].wait(tinv_index.phase) + tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) + bars.mb_y_inp_ready[0].wait(y_inp_ready.phase) + y_inp_ready = advance(y_inp_ready, 1) + for k in cutlass.range_constexpr(cfg.b_t // 16): + mma_ts_step(bmm_u_ts_desc, y_inp_ptr, desc_tinv, u_acc_ptr, k, cutlass.Boolean(k > 0)) + if elect_one: + bars.mb_u_acc_ready[0].arrive(cta_group=1) + bars.mb_t_inv_done[tinv_idx].arrive(cta_group=1) + + # ---- KK pair lookahead (member 0) = K(S) @ K^T ----------------------- + if member == 0: + if (local_idx >> 1) + 1 < n_pairs: + 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) + kqf_idx = kq_fused_index.idx + 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) + if elect_one: + bars.mb_cg0_acc_ready[member0_acc_idx].arrive(cta_group=1) + + # ---- state += decayed U(T) @ K (GEMM 7) ---------------------------------- + bars.mb_decay_u_inp_ready[0].wait(decay_u_inp_ready.phase) + decay_u_inp_ready = advance(decay_u_inp_ready, 1) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) + for k in cutlass.range_constexpr(cfg.b_t // 16): + mma_ts_step(bmm_kv_desc, decay_u_inp_ptr, desc_kt, state_acc_ptr, k, cutlass.Boolean(True) if cutlass.const_expr(k > 0) else have_state) + if elect_one: + bars.mb_state_acc_ready[kv_acc_idx].arrive(cta_group=1) + bars.mb_kq_done[kq_idx].arrive(cta_group=1) + + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + bars.mb_tmem_done[0].wait(0) + nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) + nvvm.tcgen05_dealloc( + nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), + cutlass.Int32(512), + group=nvvm.CTAGroup.CTA_1, + ) + + +@cute.jit +def tmaldg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sKQ_raw, + sV_raw, + desc_k_base, + desc_v_base, + mSched, + sSched, + bars, +): + """TMA-LDG warp role (warp 9): persistent scheduler loop + per-chunk + K/V G->S TMA loads.""" + elect_one = nvvm.elect_sync() + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + kq_index = PipelineState.start(phase=1) + v_index = PipelineState.start(phase=1) + sched_state = PipelineState.start(phase=1) + tile_idx = cutlass.Int32(bidx) + + bpe = cfg.io_dtype.width // 8 + granu = 128 // bpe + bt = cfg.b_t + kq_stage_elems = cfg.kq_cosize // cfg.smem_kq_stages + kq_box_elems = kq_stage_elems // 4 + sKQ_lo_tma = SmemTile( + base=sKQ_raw, + elems_per_stage=kq_stage_elems, + stages=cfg.smem_kq_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=2, + tma_granu_elems=granu, + tma_subtile_stride_elems=2 * bt * granu, + ) + sV_tma = SmemTile( + base=sV_raw, + elems_per_stage=cfg.v_cosize // cfg.smem_v_stages, + stages=cfg.smem_v_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=2, + tma_granu_elems=granu, + tma_subtile_stride_elems=4096, + ) + heads_out = cutlass.Int32(cfg.n_heads_out) + desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) + + 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) + + head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) + head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) + slot = batch_idx * desc_qwords + desc_k_slot = (desc_k_base + slot).tospace(cutlass.AddressSpace.generic) + desc_v_slot = (desc_v_base + slot).tospace(cutlass.AddressSpace.generic) + if elect_one: + tma_tensormap_acquire(desc_k_slot) + tma_tensormap_acquire(desc_v_slot) + + wend_padded = cstart + ((wend - cstart + 1) // 2) * 2 + if wend_padded > 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) + if elect_one: + bars.mb_kq_ready[kq_idx].arrive(n_bytes=cfg.tma_kq_bytes) + tok_coord = cstart * cutlass.Int32(cfg.b_t) + 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): + tok_coord = chunk_idx * cutlass.Int32(cfg.b_t) + + # ---- K load ------------------------------------------------------ + kq_idx = kq_index.idx + bars.mb_kq_done[kq_idx].wait(kq_index.phase) + kq_index = advance(kq_index, cfg.smem_kq_stages) + if elect_one: + bars.mb_kq_ready[kq_idx].arrive(n_bytes=cfg.tma_kq_bytes) + member = (chunk_idx - cstart) & 1 + k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, tok_coord) + kq_tile = sKQ_lo_tma[kq_idx] + if member == 0: + tma_load_tile(kq_tile, k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) + else: + tma_load_tile(kq_tile.shifted(kq_box_elems), k_slice, bars.mb_kq_ready[kq_idx].smem_ptr, acquire=False) + + # ---- V load ------------------------------------------------------ + v_idx = v_index.idx + bars.mb_v_done[v_idx].wait(v_index.phase) + 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 = (chunk_idx - 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) + + v_idx = v_index.idx + bars.mb_v_done[v_idx].wait(v_index.phase) + 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_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) + + tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + + for _ in range(cfg.smem_kq_stages): + bars.mb_kq_done[kq_index.idx].wait(kq_index.phase) + kq_index = advance(kq_index, cfg.smem_kq_stages) + for _ in range(cfg.smem_v_stages): + bars.mb_v_done[v_index.idx].wait(v_index.phase) + v_index = advance(v_index, cfg.smem_v_stages) + + +@cute.jit +def compute0_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + tidx, + tmem_hold, + sCumsumlog, + sBeta, + sTinv, + sCheckpoint_raw, + checkpoint_every_n_tokens, + sSched, + bars, +): + """Compute warp-group 0 role (warps 0-3): persistent scheduler loop + building the per-chunk beta-scaled T_inv operands.""" + + nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + gate_index = PipelineState.start(phase=0) + beta_index = PipelineState.start(phase=0) + cg0_acc_ready = PipelineState.start(phase=0) + tinv_index = PipelineState.start(phase=1) + + nvvm.barrier_cta_sync_aligned( + cfg.tmem_alloc_barrier_id, + thread_count=cfg.tmem_alloc_barrier_threads, + ) + tmem_base = tmem_hold.load() + + num_threads_cg0 = cfg.threads_per_warp * len(cfg.compute_group_0_warp_ids) + cg0_tidx = tidx % num_threads_cg0 + warp_id = cg0_tidx // cfg.threads_per_warp + lane_id = cg0_tidx % cfg.threads_per_warp + inverse_local_warp = warp_id % 2 + pair_half = warp_id // 2 + half_row_base = inverse_local_warp * 32 + bpe = cfg.io_dtype.width // 8 + num_vals = 32 + FRAG_COLS = 16 + ACC_N_FRAGS = cfg.b_t // FRAG_COLS + store_row = warp_id * 16 + lane_id % 16 + store_row_frag = lane_id % 16 + store_col = (lane_id // 16) * 8 + tmem_warp_row = warp_id * cfg.threads_per_warp + tmem_cg0_acc_col = tmem_base + cfg.tmem_cg0_acc_offset + ACC_STAGE_COLS = cfg.b_t + mask_zero = opaque_f32_zero() + crow_lo = warp_id * 16 + lane_id // 4 + crow_hi = crow_lo + 8 + + 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_pairs = (n_local + 1) // 2 + n_padded = n_pairs * 2 + + for pair_i in cutlass.range(n_pairs): + # ---- 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) + kk_gate_idx = gate1_idx if pair_half == 1 else gate0_idx + + row_u0_lo = half_row_base + lane_id // 4 + row_u0_hi = row_u0_lo + 8 + row_u1_lo = row_u0_lo + 16 + row_u1_hi = row_u0_lo + 24 + + kk_row_cumsumlog = [] + for r in (row_u0_lo, row_u0_hi, row_u1_lo, row_u1_hi): + kk_row_cumsumlog.append(sCumsumlog[r, 0, kk_gate_idx]) + kk_col_cumsumlog = [] + for g in cutlass.range_constexpr(8): + for b in cutlass.range_constexpr(2): + ccol = (lane_id % 4) * 2 + g * 8 + b + kk_col_cumsumlog.append(sCumsumlog[ccol, 0, kk_gate_idx]) + + decay_t_kk = [] + for u in cutlass.range_constexpr(2): + for k in cutlass.range_constexpr(num_vals): + hi_row = ((k // 2) % 2) == 1 + crow_u0 = row_u0_hi if cutlass.const_expr(hi_row) else row_u0_lo + crow_u1 = row_u1_hi if cutlass.const_expr(hi_row) else row_u1_lo + crow = crow_u1 if cutlass.const_expr(u == 1) else crow_u0 + ccol = (lane_id % 4) * 2 + ((k // 4) * 8 + k % 2) + is_lower = crow >= ccol + row_cumsumlog = kk_row_cumsumlog[u * 2 + (1 if hi_row else 0)] + 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() + + 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) + 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): + kk_beta.append(sBeta[r, 0, kk_beta_idx]) + + # ---- KK_epi (each warp pair stages its own member) ------------------- + 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) + 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) + 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, + ) + + # ---- pair inverse: warps 0-1 own matrix 0, warps 2-3 matrix 1 -------- + inv_base = tinv1_base if warp_id >= 2 else tinv0_base + + # 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) + nvvm.barrier_cta_sync_aligned( + cfg.inverse_barrier_id, + thread_count=cfg.inverse_barrier_threads, + ) + + # 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) + 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) + nvvm.barrier_cta_sync_aligned( + cfg.inverse_barrier_id, + thread_count=cfg.inverse_barrier_threads, + ) + + # 32x32 -> 64x64, two warps per matrix + blockwise_diagonal_32x32_to_64x64(cfg, inv_base, inverse_local_warp, lane_id) + nvvm.barrier_cta_sync_aligned( + cfg.inverse_barrier_id, + thread_count=cfg.inverse_barrier_threads, + ) + + # ---- Beta column-scaling + publish, stage 0 -------------------------- + beta_col = [] + for k in cutlass.range_constexpr(num_vals): + beta_col.append(sBeta[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, beta0_idx]) + tinv_frags = [] + for c in cutlass.range_constexpr(ACC_N_FRAGS): + tinv_frags += list( + nvvm.ldmatrix( + cutlass.inttoptr( + tinv0_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( + tinv0_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[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( + 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, + ), + [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() + + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + for _ in range(cfg.smem_t_inv_stages): + bars.mb_t_inv_done[tinv_index.idx].wait(tinv_index.phase) + tinv_index = advance(tinv_index, cfg.smem_t_inv_stages) + + +@cute.jit +def compute1_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + tidx, + warp_idx, + tmem_hold, + sV, + sCumsumlog, + sCumprod, + sBeta, + sCheckpoint_raw, + mState_init, + mState_out, + checkpoint_every_n_tokens, + sSched, + bars, +): + """Compute warp-group 1 role (warps 4-7): persistent scheduler loop + owning the recurrent state from seed to final store.""" + elect_one = nvvm.elect_sync() + + v_index = PipelineState.start(phase=0) + gate_index = PipelineState.start(phase=0) + kv_acc_index = PipelineState.start(phase=0) + k_state_ready_index = PipelineState.start(phase=0) + u_acc_ready_index = PipelineState.start(phase=0) + state_acc_seed_index = PipelineState.start(phase=1) + state_inp_cnt = cutlass.Int32(0) + kv_done_idx = cutlass.Int32(0) + + nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) + nvvm.barrier_cta_sync_aligned( + cfg.tmem_alloc_barrier_id, + thread_count=cfg.tmem_alloc_barrier_threads, + ) + tmem_base = tmem_hold.load() + + num_threads_cg1 = cfg.threads_per_warp * len(cfg.compute_group_1_warp_ids) + cg1_tidx = tidx % num_threads_cg1 + lane_id = cg1_tidx % cfg.threads_per_warp + tmem_warp_row = (cg1_tidx // cfg.threads_per_warp) * cfg.threads_per_warp + ldtm_width = 32 + sttm_width = ldtm_width // 2 + num_state_subs = cutlass.const_expr(cfg.d_v // ldtm_width) + tmem_state_col = tmem_base + cfg.tmem_state_acc_offset + tmem_state_inp_col = tmem_base + cfg.tmem_state_inp_offset + tmem_inp_col = tmem_base + cfg.tmem_y_decay_u_inp_offset + INP_SLOT_COLS = cfg.b_t // 2 + tmem_k_state_col = tmem_base + cfg.tmem_cg1_acc_offset + tmem_u_acc_col = tmem_k_state_col + tmem_y_inp_col = tmem_inp_col + tmem_decay_v_col = tmem_inp_col + INP_SLOT_COLS + v_frag_tok = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 + v_frag_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 + v_frag_slab = (cg1_tidx // 64) * 4096 + v_stage_elems = cfg.v_cosize // cfg.smem_v_stages + sV_base = cute.make_ptr(cfg.io_dtype, sV[0].base, mem_space=cute.AddressSpace.smem, assumed_align=cfg.buffer_align_bytes) + num_vals = 32 + if cutlass.const_expr(cfg.enable_checkpoints): + sCheckpoint_base_int = sCheckpoint_raw.data_ptr().toint() + checkpoint_cnt = cutlass.Int32(0) + checkpoint_frag_row = cg1_tidx % 8 + (cg1_tidx // 16 % 2) * 8 + checkpoint_frag_col = (cg1_tidx // 8 % 2) * 8 + (cg1_tidx // 32 % 2) * 32 + + 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 + if cutlass.const_expr(cfg.enable_checkpoints): + ckpt_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + checkpoint_mod = cstart % ckpt_chunks + if n_local > 0: + if cutlass.const_expr(cfg.use_initial_state): + # ---- initial-state seed: initial_state GMEM -> state TMEM --------------- + gState_init = mState_init[None, None, head_idx, batch_idx] + kv_init_idx = state_acc_seed_index.idx + bars.mb_state_acc_scale_done[kv_init_idx].wait(state_acc_seed_index.phase) + state_acc_seed_index = advance(state_acc_seed_index, cfg.tmem_state_acc_stages) + seed_from_initial_state = cstart == 0 + if seed_from_initial_state: + 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] + if cutlass.const_expr(cfg.state_dtype != cfg.acc_dtype): + v = v.to(cfg.acc_dtype) + words.append(v) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + cutlass.Vector.from_elements(tuple(words), cutlass.Float32), + ) + else: + for sub in cutlass.range_constexpr(num_state_subs): + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + cutlass.Vector.from_elements(tuple(cutlass.Float32(0.0) for _ in range(32)), cutlass.Float32), + ) + nvvm.tcgen05_wait("store") + + nvvm.barrier_cta_sync_aligned( + cfg.init_state_store_barrier_id, + thread_count=cfg.init_state_store_barrier_threads, + ) + + for local_idx in cutlass.range(n_padded): # 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 + valid_state = local_idx > 0 + if cutlass.const_expr(cfg.use_initial_state): + valid_state = cutlass.Boolean(True) + state_acc_seed_index = advance(state_acc_seed_index, cfg.tmem_state_acc_stages) + + gate_idx = gate_index.idx + bars.mb_gate_ready[gate_idx].wait(gate_index.phase) + gate_index = advance(gate_index, cfg.smem_gate_stages) + cumprod_total = sCumprod[sCumprod.shape[0] - 1, 0, gate_idx] + + # ---- state restage + rescale ------------------------------------- + if valid_state: + kv_idx = kv_acc_index.idx + bars.mb_state_acc_ready[kv_idx].wait(kv_acc_index.phase) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) + kv_done_idx = kv_idx + + state_regs = [[cutlass.Float32(0.0) for _ in range(num_state_subs)] for _ in range(32)] + state_inp_stage_idx = state_inp_cnt % cfg.tmem_state_inp_stages + state_vecs = [] + for sub in cutlass.range_constexpr(num_state_subs): + state_vecs.append( + nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), num=32) + ) + for sub in cutlass.range_constexpr(num_state_subs): + for k in cutlass.range_constexpr(32): + state_regs[k][sub] = state_vecs[sub][k] + state_pack = [fp32_to_fp16(state_regs[2 * j][sub], state_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_inp_col + sub * sttm_width, cutlass.Int32), + cutlass.Vector.from_elements(tuple(state_pack), cutlass.Int32), + ) + nvvm.tcgen05_wait("store") + bars.mb_state_inp_ready[state_inp_stage_idx].arrive() + state_inp_cnt = state_inp_cnt + 1 + + 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 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, + ) + nvvm.fence_proxy("async.shared", space="cta") + if elect_one: + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + checkpoint_cnt = checkpoint_cnt + 1 + + for sub in cutlass.range_constexpr(num_state_subs): + state_scaled = [] + for j in cutlass.range_constexpr(16): + s0, s1 = fmul2(state_regs[2 * j][sub], state_regs[2 * j + 1][sub], cumprod_total, cumprod_total) + state_scaled += [s0, s1] + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), + cutlass.Vector.from_elements(tuple(state_scaled), cutlass.Float32), + ) + nvvm.tcgen05_wait("store") + bars.mb_state_acc_scale_done[kv_idx].arrive() + + # ---- per-row Gate register builds -------------------------------- + cumprod_vals = [] + for k in cutlass.range_constexpr(num_vals): + cumprod_vals.append(sCumprod[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) + last_cumsumlog = sCumsumlog[cfg.b_t - 1, 0, gate_idx] + cumsumlog_vals = [] + for k in cutlass.range_constexpr(num_vals): + cumsumlog_vals.append(sCumsumlog[(lane_id % 4) * 2 + ((k // 4) * 8 + k % 2), 0, gate_idx]) + decay_scale_vals = [] + for k in cutlass.range_constexpr(0, num_vals, 2): + d0, d1 = fadd2(last_cumsumlog, last_cumsumlog, -cumsumlog_vals[k], -cumsumlog_vals[k + 1]) + decay_scale_vals.append(cute.math.exp2(d0, fastmath=True)) + decay_scale_vals.append(cute.math.exp2(d1, fastmath=True)) + bars.mb_gate_done[gate_idx].arrive() + + # ---- Y = V - K*state (packed 16-bit) ----------------------------- + v_idx = v_index.idx + bars.mb_v_ready[v_idx].wait(v_index.phase) + v_index = advance(v_index, cfg.smem_v_stages) + + v_frags = [[cutlass.Int32(0), cutlass.Int32(0)] for _ in range(16)] + for c in cutlass.range_constexpr(8): + tok_block = cutlass.const_expr(c % 4) + sub = cutlass.const_expr(c // 4) + v_frag = nvvm.ldmatrix( + ( + sV_base + + v_idx * v_stage_elems + + v_frag_slab + + (v_frag_tok + tok_block * 16) * 64 + + swizzle_xor_128b(v_frag_tok + tok_block * 16, v_frag_col + sub * 16) + ).raw_ptr(), + 4, + nvvm.MMALayout.COL, + ) + for i in cutlass.range_constexpr(4): + v_frags[4 * tok_block + i][sub] = v_frag[i] + if valid_state: + bars.mb_k_state_acc_ready[0].wait(k_state_ready_index.phase) + k_state_ready_index = advance(k_state_ready_index, 1) + + for sub in cutlass.range_constexpr(2): + k_state_vec = nvvm.tcgen05_ld( + "16x256b", + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_k_state_col, cutlass.Float32), + num=8, + ) + for j in cutlass.range_constexpr(16): + s0, s1 = fmul2(k_state_vec[2 * j], k_state_vec[2 * j + 1], cumprod_vals[2 * j], cumprod_vals[2 * j + 1]) + k_state_pack = fp32_to_fp16(s0, s1, dtype=cfg.io_dtype) + v_frags[j][sub] = sub_f16x2(v_frags[j][sub], k_state_pack, cfg.io_dtype) + for sub in cutlass.range_constexpr(2): + nvvm.tcgen05_st( + "16x128b", + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_y_inp_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(v_frags[j][sub] for j in range(16)), cutlass.Int32), + ) + nvvm.tcgen05_wait("store") + bars.mb_y_inp_ready[0].arrive() + + # ---- U epilogue + decayed-U publish ------------------------------ + bars.mb_u_acc_ready[0].wait(u_acc_ready_index.phase) + u_acc_ready_index = advance(u_acc_ready_index, 1) + bars.mb_v_done[v_idx].arrive() + + u_acc_regs = [[cutlass.Float32(0.0), cutlass.Float32(0.0)] for _ in range(32)] + u_acc_vecs = [] + for sub in cutlass.range_constexpr(2): + u_acc_vecs.append( + nvvm.tcgen05_ld( + "16x256b", + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_u_acc_col, cutlass.Float32), + num=8, + ) + ) + for sub in cutlass.range_constexpr(2): + for k in cutlass.range_constexpr(32): + u_acc_regs[k][sub] = u_acc_vecs[sub][k] + + for sub in cutlass.range_constexpr(2): + for j in cutlass.range_constexpr(16): + u_acc_regs[2 * j][sub], u_acc_regs[2 * j + 1][sub] = fmul2( + u_acc_regs[2 * j][sub], u_acc_regs[2 * j + 1][sub], decay_scale_vals[2 * j], decay_scale_vals[2 * j + 1] + ) + decay_pack = [fp32_to_fp16(u_acc_regs[2 * j][sub], u_acc_regs[2 * j + 1][sub], dtype=cfg.io_dtype) for j in range(16)] + nvvm.tcgen05_st( + "16x128b", + nvvm.make_tmem_ptr(((tmem_warp_row + sub * 16) << 16) + tmem_decay_v_col, cutlass.Int32), + cutlass.Vector.from_elements(tuple(decay_pack), cutlass.Int32), + ) + nvvm.tcgen05_wait("store") + bars.mb_decay_u_inp_ready[0].arrive() + + # ---- final state: state TMEM -> GMEM ----------------------------------- + if n_local > 0: + kv_last_idx = kv_acc_index.idx + bars.mb_state_acc_ready[kv_last_idx].wait(kv_acc_index.phase) + kv_acc_index = advance(kv_acc_index, cfg.tmem_state_acc_stages) + if cutlass.const_expr(cfg.store_final_state): + if wend == num_chunks_b: + gState_out = mState_out[None, None, head_idx, batch_idx] + for sub in cutlass.range_constexpr(num_state_subs): + state_vec = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr((tmem_warp_row << 16) + tmem_state_col + sub * ldtm_width, cutlass.Float32), num=32 + ) + for k in cutlass.range_constexpr(32): + 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 + bars.mb_state_acc_scale_done[kv_last_idx].arrive() + else: + bars.mb_state_acc_scale_done[kv_last_idx].arrive() + else: + if cutlass.const_expr(cfg.store_final_state): + write_passthrough = wend == num_chunks_b + if write_passthrough: + gState_out = mState_out[None, None, head_idx, batch_idx] + 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] + 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) + + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + bars.mb_tmem_done[0].arrive() + + if cutlass.const_expr(cfg.enable_checkpoints): + for _ in range(cfg.smem_checkpoint_stages): + 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_cnt = checkpoint_cnt + 1 + + +@cute.kernel +def build_all_descs_kernel( + base_k: cutlass.GridConstant[tma.TensorMap], + base_v: cutlass.GridConstant[tma.TensorMap], + base_checkpoint: cutlass.GridConstant[tma.TensorMap], + desc_ws: cute.Tensor, + cu_seqlens: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + state_checkpoints_out: Optional[cute.Tensor], + n_batch: cutlass.Int32, + k_row_stride: cutlass.Int32, + v_row_stride: cutlass.Int32, + checkpoint_row_stride: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, +) -> None: + """Single-launch builder for the per-BATCH descriptor arrays (one warp + per array).""" + tidx, _, _ = cute.arch.thread_idx() + widx = cutlass.Int32(tidx) // cutlass.Int32(32) + arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_k_arr = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) + desc_v_arr = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_checkpoint_arr = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + + if widx == 0: + if nvvm.elect_sync(): + emit_seq_descs(base_k, desc_k_arr, cu_seqlens, k, n_batch, k_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 1: + if nvvm.elect_sync(): + emit_seq_descs(base_v, desc_v_arr, cu_seqlens, v, n_batch, v_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if cutlass.const_expr(state_checkpoints_out is not None): + if widx == 2: + if nvvm.elect_sync(): + emit_checkpoint_seq_descs( + base_checkpoint, desc_checkpoint_arr, cu_seqlens, state_checkpoints_out, 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) + + +@cute.jit +def build_descs( + io_dtype: cutlass.Constexpr, + b_t: cutlass.Constexpr[int], + k: cute.Tensor, + v: cute.Tensor, + cu_seqlens: cute.Tensor, + state_checkpoints_out: Optional[cute.Tensor], + checkpoint_every_n: cutlass.Int32, + tensormap_workspace: cute.Tensor, + stream: cuda.CUstream, +): + """Build the 3 per-batch TMA-descriptor arrays (K, V, checkpoints) into + ``tensormap_workspace``.""" + h_k = k.shape[1] + h_v = v.shape[1] + batch_size = cu_seqlens.shape[0] - 1 + d_v = v.shape[2] + bpe = io_dtype.width // 8 + granu = 128 // bpe + bt = b_t + + k_row_stride, k_head_stride = k.stride[0], k.stride[1] + v_row_stride, v_head_stride = v.stride[0], v.stride[1] + + seqlen = k.shape[0] + d_k = k.shape[2] + k_headed = cute.make_tensor(k.iterator, cute.make_layout((seqlen, h_k, d_k), stride=(k_row_stride, k_head_stride, 1))) + v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v_head_stride, v_row_stride))) + swz128 = tma.TensorMapSwizzle.s128b + base_desc_k = tma.create_tensor_map_tiled_from_view(k_headed, box_dims=(bt, 1, granu), stride_order=(2, 1, 0), swizzle=swz128) + base_desc_v = tma.create_tensor_map_tiled_from_view(v_headed, box_dims=(granu, 1, bt), stride_order=(0, 1, 2), swizzle=swz128) + + base_desc_checkpoint = base_desc_v + if cutlass.const_expr(state_checkpoints_out is not None): + d_k_state = state_checkpoints_out.shape[2] + d_v_state = state_checkpoints_out.shape[3] + checkpoint_granu = 128 // (state_checkpoints_out.element_type.width // 8) + checkpoint_view = cute.make_tensor( + state_checkpoints_out.iterator, + cute.make_layout( + (d_v_state, d_k_state, state_checkpoints_out.shape[0], state_checkpoints_out.shape[1]), + stride=(state_checkpoints_out.stride[3], state_checkpoints_out.stride[2], state_checkpoints_out.stride[0], state_checkpoints_out.stride[1]), + ), + ) + base_desc_checkpoint = tma.create_tensor_map_tiled_from_view( + checkpoint_view, box_dims=(checkpoint_granu, d_k_state, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz128 + ) + + n_warps = 3 if state_checkpoints_out is not None else 2 + build_all_descs_kernel( + base_desc_k, + base_desc_v, + base_desc_checkpoint, + tensormap_workspace, + cu_seqlens, + k, + v, + state_checkpoints_out, + cutlass.Int32(batch_size), + cutlass.Int32(k_row_stride), + cutlass.Int32(v_row_stride), + cutlass.Int32(state_checkpoints_out.stride[0] if state_checkpoints_out is not None else 0), + checkpoint_every_n, + ).launch(grid=(1, 1, 1), block=(32 * n_warps, 1, 1), stream=stream) + + +@cute.jit +def host( + cfg: cutlass.Constexpr, + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + beta: cute.Tensor, + cu_seqlens: cute.Tensor, + state_in: Optional[cute.Tensor], + state_out: Optional[cute.Tensor], + work_items: Optional[cute.Tensor], + work_count: Optional[cute.Tensor], + sched_ctr: Optional[cute.Tensor], + checkpoint_every_n_tokens: cutlass.Int32, + tensormap_workspace: cute.Tensor, + stream: cuda.CUstream, +): + h_k = k.shape[1] + h_v = v.shape[1] + batch_size = cu_seqlens.shape[0] - 1 + heads_out = gate.shape[1] + + # ---- GQA reshapes: fold the head group into a -------------------------------- + if cutlass.const_expr(cfg.is_GQA): + h_ratio = heads_out // h_v + h_native = h_v + k = cute.make_tensor( + k.iterator, + cute.make_layout( + (k.shape[0], k.shape[2], (h_ratio, h_v)), + stride=(k.stride[0], k.stride[2], (0, k.stride[1])), + ), + ) + v = cute.make_tensor( + v.iterator, + cute.make_layout( + (v.shape[2], v.shape[0], (h_ratio, h_v)), + stride=(v.stride[2], v.stride[0], (0, v.stride[1])), + ), + ) + else: + h_ratio = h_v // h_k + h_native = h_k + k = cute.make_tensor( + k.iterator, + cute.make_layout( + (k.shape[0], k.shape[2], (h_ratio, h_k)), + stride=(k.stride[0], k.stride[2], (0, k.stride[1])), + ), + ) + v = cute.make_tensor( + v.iterator, + cute.make_layout( + (v.shape[2], v.shape[0], (h_ratio, h_k)), + stride=(v.stride[2], v.stride[0], (v.stride[1], h_ratio * v.stride[1])), + ), + ) + + gate = cute.make_tensor( + gate.iterator, + cute.make_layout( + (gate.shape[0], (h_ratio, h_native)), + stride=(gate.stride[0], (gate.stride[1], h_ratio * gate.stride[1])), + ), + ) + beta = cute.make_tensor( + beta.iterator, + cute.make_layout( + (beta.shape[0], (h_ratio, h_native)), + stride=(beta.stride[0], (beta.stride[1], h_ratio * beta.stride[1])), + ), + ) + if cutlass.const_expr(state_in is not None): + state_in = cute.make_tensor( + state_in.iterator, + cute.make_layout( + (state_in.shape[2], state_in.shape[3], (h_ratio, h_native), state_in.shape[0]), + stride=( + state_in.stride[2], + state_in.stride[3], + (state_in.stride[1], h_ratio * state_in.stride[1]), + state_in.stride[0], + ), + ), + ) + if cutlass.const_expr(state_out is not None): + state_out = cute.make_tensor( + state_out.iterator, + cute.make_layout( + (state_out.shape[2], state_out.shape[3], (h_ratio, h_native), state_out.shape[0]), + stride=( + state_out.stride[2], + state_out.stride[3], + (state_out.stride[1], h_ratio * state_out.stride[1]), + state_out.stride[0], + ), + ), + ) + + # ---- SMEM sizing: per-buffer element cosizes --------------------------------- + bpe = cfg.io_dtype.width // 8 + kq_tile_elems = 2 * cfg.b_t * cfg.d_k + v_tile_elems = cfg.d_v * cfg.b_t + tinv_tile_elems = cfg.b_t * cfg.b_t + cfg.kq_cosize = kq_tile_elems * cfg.smem_kq_stages + cfg.v_cosize = v_tile_elems * cfg.smem_v_stages + cfg.t_inv_cosize = tinv_tile_elems * cfg.smem_t_inv_stages + cfg.checkpoint_cosize = cfg.d_k * cfg.d_v * cfg.smem_checkpoint_stages + + cumsumlog_smem_layout_staged = cute.make_layout((cfg.b_t, 1, cfg.smem_gate_stages)) + beta_smem_layout_staged = cute.make_layout((cfg.b_t, 1, cfg.smem_beta_stages)) + + cfg.tma_kq_bytes = (kq_tile_elems // 2) * bpe + cfg.tma_v_bytes = v_tile_elems * bpe + + cfg.n_heads_out = heads_out + cfg.k_ratio = heads_out // h_k + cfg.v_ratio = heads_out // h_v + num_descs = batch_size + + # ---- launch ------------------------------------------------------------------ + grid_shape = (cfg.max_active_clusters, 1, 1) + + kernel( + cfg, + gate, + beta, + cu_seqlens, + state_in, + state_out, + work_items, + work_count, + sched_ctr, + checkpoint_every_n_tokens, + cumsumlog_smem_layout_staged, + beta_smem_layout_staged, + k, + v, + tensormap_workspace, + cutlass.Int32(num_descs), + ).launch( + grid=grid_shape, + block=(cfg.threads_per_cta, 1, 1), + cluster=cfg.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=1, + ) + + +@cute.kernel +def kernel( + cfg: cutlass.Constexpr, + mGate: cute.Tensor, + mBeta: cute.Tensor, + cu_seqlens: cute.Tensor, + mState_init: Optional[cute.Tensor], + mState_out: Optional[cute.Tensor], + mWorkItems: cute.Tensor, + mCount: cute.Tensor, + mSched: Optional[cute.Tensor], + checkpoint_every_n_tokens: cutlass.Int32, + cumsumlog_smem_layout_staged: cute.Layout, + beta_smem_layout_staged: cute.Layout, + mK, + mV, + tensormap_workspace: cute.Tensor, + n_desc: cutlass.Int32, +): + """Main GDN chunked kernel: warp-specialized dispatch over (batch, head) + tiles.""" + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + bidx = cute.arch.block_idx()[0] + num_ctas = cute.arch.grid_dim()[0] + + total_tiles = mCount[0] + if cutlass.const_expr(cfg.dyn_sched): + assert mSched is not None, "mSched must be provided if dyn_sched is True" + + if cutlass.const_expr(cfg.use_initial_state): + assert mState_init is not None, "mState_init must be provided if use_initial_state is True" + else: + assert mState_init is None, "mState_init must be None if use_initial_state is False" + if cutlass.const_expr(cfg.store_final_state): + assert mState_out is not None, "mState_out must be provided if store_final_state is True" + else: + assert mState_out is None, "mState_out must be None if store_final_state is False" + + desc_base_words = tensormap_workspace.iterator.raw_ptr() + desc_qwords = cutlass.Int32(TENSOR_MAP_QWORDS) + arr_words = n_desc * desc_qwords + desc_k_base = desc_base_words + desc_v_base = desc_base_words + arr_words + desc_checkpoint_base = desc_base_words + cutlass.Int32(2) * arr_words + + SMEM = cutlass.AddressSpace.smem + + bpe = cfg.io_dtype.width // 8 + SWZ = 2 + LEAD = 16 + STRIDE = 8 * 128 + KT_LEAD = (cfg.d_v // 2) * 128 + V_LEAD = (cfg.d_v // 2) * 128 + if cutlass.const_expr(cfg.enable_checkpoints): + sCheckpoint_raw = cutlass.Array( + cfg.io_dtype, + cfg.checkpoint_cosize, + space=cutlass.AddressSpace.smem, + alignment=cfg.buffer_align_bytes, + ) + else: + sCheckpoint_raw = None + sKQ_raw = cutlass.Array( + cfg.io_dtype, + cfg.kq_cosize, + space=cutlass.AddressSpace.smem, + alignment=cfg.buffer_align_bytes, + ) + sKQ = SmemTile( + base=sKQ_raw.data_ptr().toint(), + elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages) * bpe, + stages=cfg.smem_kq_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sKQ_trans = SmemTile( + base=sKQ_raw.data_ptr().toint(), + elems_per_stage=(cfg.kq_cosize // cfg.smem_kq_stages) * bpe, + stages=cfg.smem_kq_stages, + leading_byte_offset=2 * KT_LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + bars = make_gdn_bars(cfg) + tmem_hold = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=16) + sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + cumsumlog_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) + cumprod_raw = cutlass.Array(cutlass.Float32, cute.cosize(cumsumlog_smem_layout_staged), space=SMEM, alignment=128) + beta_raw = cutlass.Array(cutlass.Float32, cute.cosize(beta_smem_layout_staged), space=SMEM, alignment=128) + sTinv_raw = cutlass.Array( + cfg.io_dtype, + cfg.t_inv_cosize, + space=cutlass.AddressSpace.smem, + alignment=cfg.buffer_align_bytes, + ) + sTinv = SmemTile( + base=sTinv_raw.data_ptr().toint(), + elems_per_stage=(cfg.t_inv_cosize // cfg.smem_t_inv_stages) * bpe, + stages=cfg.smem_t_inv_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sV_raw = cutlass.Array( + cfg.io_dtype, + cfg.v_cosize, + space=cutlass.AddressSpace.smem, + alignment=cfg.buffer_align_bytes, + ) + sV = SmemTile( + base=sV_raw.data_ptr().toint(), + elems_per_stage=(cfg.v_cosize // cfg.smem_v_stages) * bpe, + stages=cfg.smem_v_stages, + leading_byte_offset=V_LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sCumsumlog = cute.make_tensor( + cute.make_ptr(cutlass.Float32, cumsumlog_raw.data_ptr().toint(), mem_space=cute.AddressSpace.smem, assumed_align=128), + cumsumlog_smem_layout_staged, + ) + sCumprod = cute.make_tensor( + cute.make_ptr(cutlass.Float32, cumprod_raw.data_ptr().toint(), mem_space=cute.AddressSpace.smem, assumed_align=128), + cumsumlog_smem_layout_staged, + ) + sBeta = cute.make_tensor( + cute.make_ptr(cutlass.Float32, beta_raw.data_ptr().toint(), mem_space=cute.AddressSpace.smem, assumed_align=128), + beta_smem_layout_staged, + ) + + # ---- mbarrier init (all threads) --------------------------------------------- + for s in range(cfg.smem_kq_stages): + bars.mb_kq_ready[s].init() + bars.mb_kq_done[s].init() + for s in range(cfg.smem_v_stages): + bars.mb_v_ready[s].init() + bars.mb_v_done[s].init() + for s in range(cfg.smem_gate_stages): + bars.mb_gate_ready[s].init() + bars.mb_gate_done[s].init() + for s in range(cfg.smem_beta_stages): + bars.mb_beta_ready[s].init() + bars.mb_beta_done[s].init() + for s in range(cfg.tmem_state_acc_stages): + bars.mb_state_acc_ready[s].init() + bars.mb_state_acc_scale_done[s].init() + for s in range(cfg.tmem_cg0_acc_stages): + bars.mb_cg0_acc_ready[s].init() + bars.mb_cg0_acc_done[s].init() + bars.mb_k_state_acc_ready[0].init() + bars.mb_u_acc_ready[0].init() + for s in range(cfg.smem_t_inv_stages): + bars.mb_t_inv_ready[s].init() + bars.mb_t_inv_done[s].init() + for s in range(cfg.tmem_state_inp_stages): + bars.mb_state_inp_ready[s].init() + for b in (bars.mb_y_inp_ready, bars.mb_decay_u_inp_ready): + b[0].init() + for s in range(cfg.smem_checkpoint_stages): + bars.mb_checkpoint_tmastg_ready[s].init() + bars.mb_checkpoint_tmastg_done[s].init() + for s_ in range(cfg.sched_stages): + bars.mb_sched_ready[s_].init() + bars.mb_sched_done[s_].init() + bars.mb_tmem_done[0].init() + + nvvm.fence_mbarrier_init() + nvvm.barrier_cta_sync() + + # ---- warp specialization ----------------------------------------------------- + + if warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: + compute0_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + tidx, + tmem_hold=tmem_hold, + sCumsumlog=sCumsumlog, + sBeta=sBeta, + sTinv=sTinv, + sCheckpoint_raw=sCheckpoint_raw, + checkpoint_every_n_tokens=checkpoint_every_n_tokens, + sSched=sSched, + bars=bars, + ) + + if warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]: + compute1_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + tidx, + warp_idx=warp_idx, + tmem_hold=tmem_hold, + sV=sV, + sCumsumlog=sCumsumlog, + sCumprod=sCumprod, + sBeta=sBeta, + sCheckpoint_raw=sCheckpoint_raw, + mState_init=mState_init, + mState_out=mState_out, + checkpoint_every_n_tokens=checkpoint_every_n_tokens, + sSched=sSched, + bars=bars, + ) + + elif warp_idx == cfg.load_gate_beta_warp_id: + gate_beta_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + tidx=tidx, + mGate=mGate, + mBeta=mBeta, + sCumsumlog=sCumsumlog, + sCumprod=sCumprod, + sBeta=sBeta, + sSched=sSched, + bars=bars, + ) + + elif warp_idx == cfg.mma_warp_id: + mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + tmem_hold=tmem_hold, + sKQ=sKQ, + sKQ_trans=sKQ_trans, + sTinv=sTinv, + sSched=sSched, + bars=bars, + ) + + elif warp_idx == cfg.tma_kv_warp_id: + tmaldg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sKQ_raw=sKQ_raw, + sV_raw=sV_raw, + desc_k_base=desc_k_base, + desc_v_base=desc_v_base, + mSched=mSched, + sSched=sSched, + bars=bars, + ) + + if warp_idx == cfg.epilogue_warp_id: + tmastg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + checkpoint_every_n_tokens=checkpoint_every_n_tokens, + tidx=tidx, + sCheckpoint_raw=sCheckpoint_raw, + desc_checkpoint_base=desc_checkpoint_base, + sSched=sSched, + bars=bars, + ) + + +@dataclass +class GdnRecomputeCfg: + """Per-compile GDN kernel knob, built by ``build_cfg``. + + The per-compile parameters (dtypes, GQA, state flags) are the + ``cute.compile`` cache keys; the rest is derived from the module-global + ``CFG`` constants. ``host`` stamps the shape-derived fields at trace + time. Passed ``cfg``-first (a ``cutlass.Constexpr``) into ``host`` / + ``kernel`` and every warp body. + """ + + io_dtype: Type[cutlass.Numeric] + acc_dtype: Type[cutlass.Numeric] + state_dtype: Type[cutlass.Numeric] + max_active_clusters: int + is_GQA: bool + use_initial_state: bool + store_final_state: bool + enable_checkpoints: bool + log_gate: bool = False + dyn_sched: bool = False + sched_stages: int = CFG.SMEM_SCHED_STAGES + + # ---- fixed constants stamped from CFG by build_cfg --------------------------- + b_t: int = CFG.B_T + d_k: int = CFG.D_K + d_v: int = CFG.D_V + compute_group_0_warp_ids: Tuple[int, ...] = CFG.COMPUTE_GROUP_0_WARP_IDS + compute_group_1_warp_ids: Tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS + load_gate_beta_warp_id: int = CFG.LOAD_GATE_BETA_WARP_ID + tma_kv_warp_id: int = CFG.TMA_KV_WARP_ID + mma_warp_id: int = CFG.MMA_WARP_ID + epilogue_warp_id: int = CFG.EPILOGUE_WARP_ID + num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 + num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 + num_regs_other: int = CFG.NUM_REGS_OTHER + threads_per_warp: int = CFG.THREADS_PER_WARP + threads_per_cta: int = 0 + cluster_shape_mnk: Tuple[int, int, int] = CFG.CLUSTER_SHAPE_MNK + + # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ------------------- + tmem_alloc_barrier_id: int = 1 + tmem_alloc_barrier_threads: int = 0 + inverse_barrier_id: int = 2 + inverse_barrier_threads: int = 0 + init_state_store_barrier_id: int = 4 + init_state_store_barrier_threads: int = 0 + + # ---- SMEM / TMEM stage counts + TMEM column offsets -------------------------- + smem_kq_stages: int = CFG.SMEM_KQ_STAGES + smem_v_stages: int = CFG.SMEM_V_STAGES + smem_t_inv_stages: int = CFG.SMEM_T_INV_STAGES + smem_checkpoint_stages: int = 1 + smem_gate_stages: int = CFG.SMEM_GATE_STAGES + smem_beta_stages: int = CFG.SMEM_BETA_STAGES + tmem_state_acc_stages: int = CFG.TMEM_KV_ACC_STAGES + tmem_state_inp_stages: int = CFG.TMEM_STATE_INP_STAGES + tmem_cg0_acc_stages: int = CFG.TMEM_CG0_ACC_STAGES + tmem_cg1_acc_stages: int = CFG.TMEM_CG1_ACC_STAGES + tmem_state_acc_offset: int = 0 + tmem_state_inp_offset: int = 0 + tmem_cg0_acc_offset: int = 0 + tmem_cg1_acc_offset: int = 0 + tmem_y_decay_u_inp_offset: int = 0 + buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES + + # ---- stamped by host at trace time (shape-derived) -------------------------- + kq_cosize: int = 0 + v_cosize: int = 0 + t_inv_cosize: int = 0 + checkpoint_cosize: int = 0 + tma_kq_bytes: int = 0 + tma_v_bytes: int = 0 + n_heads_out: int = 0 + k_ratio: int = 1 + v_ratio: int = 1 + + +def build_cfg( + io_dtype: Type[cutlass.Numeric], + state_dtype: Type[cutlass.Numeric], + *, + max_active_clusters: int, + is_GQA: bool, + use_initial_state: bool, + store_final_state: bool = True, + enable_checkpoints: bool = False, + log_gate: bool = False, + dyn_sched: bool = False, +) -> GdnRecomputeCfg: + """Build the per-compile ``GdnRecomputeCfg`` (io_dtype ∈ {Float16, BFloat16}; + acc is always Float32).""" + if io_dtype not in (cutlass.Float16, cutlass.BFloat16): + raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") + cfg = GdnRecomputeCfg( + io_dtype=io_dtype, + acc_dtype=cutlass.Float32, + state_dtype=state_dtype, + max_active_clusters=max_active_clusters, + is_GQA=is_GQA, + use_initial_state=use_initial_state, + store_final_state=store_final_state, + enable_checkpoints=enable_checkpoints, + log_gate=log_gate, + dyn_sched=dyn_sched, + ) + cfg.smem_checkpoint_stages = 1 + if enable_checkpoints: + cfg.smem_kq_stages = 3 + if not use_initial_state: + cfg.num_regs_compute_group_1 = 232 + cfg.num_regs_other = 48 + n_cg0 = len(cfg.compute_group_0_warp_ids) + n_cg1 = len(cfg.compute_group_1_warp_ids) + cfg.threads_per_cta = cfg.threads_per_warp * (4 + n_cg0 + n_cg1) + cfg.tmem_alloc_barrier_threads = cfg.threads_per_warp * (1 + n_cg0 + n_cg1) + cfg.inverse_barrier_threads = cfg.threads_per_warp * n_cg0 + cfg.init_state_store_barrier_threads = cfg.threads_per_warp * n_cg1 + cfg.tmem_state_acc_offset = 0 + cfg.tmem_state_inp_offset = cfg.tmem_state_acc_offset + cfg.tmem_state_acc_stages * 128 + cfg.tmem_cg0_acc_offset = cfg.tmem_state_inp_offset + cfg.tmem_state_inp_stages * 64 + cfg.tmem_cg1_acc_offset = cfg.tmem_cg0_acc_offset + cfg.tmem_cg0_acc_stages * 64 + cfg.tmem_y_decay_u_inp_offset = cfg.tmem_cg1_acc_offset + cfg.tmem_cg1_acc_stages * 64 + return cfg + + +TENSORMAP_DESC_ARRAYS = 3 # per-batch runtime TMA descriptors: K, V, checkpoints +TENSORMAP_STATIC_SLOTS = 0 + + +# --------------------------------------------------------------------------- + + +@functools.cache +def get_compiled_cache( + io_dtype_str: str, + state_dtype_str: str, + cu_dtype_str: str, + HK: int, + HV: int, + HO: int, + is_GQA: bool, + use_initial_state: bool, + store_final_state: bool, + enable_checkpoints: bool, + log_gate: bool, + dyn_sched: bool, +): + """Return a mutable dict that lazily stores the compiled kernel.""" + return {} + + +def compile( + io_dtype, + state_dtype, + is_GQA: bool, + use_initial_state: bool, + store_final_state: bool, + enable_checkpoints: bool, + log_gate: bool = False, + dyn_sched: bool = False, + *, + num_sm: int, + k_cute, + v_cute, + gate_cute, + beta_cute, + cu_seqlens_cute, + state_in_cute, + state_out_cute, + work_items_cute=None, + work_count_cute=None, + sched_ctr_cute=None, + checkpoint_every_n_tokens, + workspace_cute, + stream, +): + """JIT-compile the chunked GDN recompute kernel for one static config.""" + cfg = build_cfg( + io_dtype, + state_dtype, + max_active_clusters=num_sm, + is_GQA=is_GQA, + use_initial_state=use_initial_state, + store_final_state=store_final_state, + enable_checkpoints=enable_checkpoints, + log_gate=log_gate, + dyn_sched=dyn_sched, + ) + + return cute.compile( + host, + cfg, + k_cute, + v_cute, + gate_cute, + beta_cute, + cu_seqlens_cute, + state_in_cute, + state_out_cute, + work_items_cute, + work_count_cute, + sched_ctr_cute, + checkpoint_every_n_tokens, + workspace_cute, + stream, + options="--enable-tvm-ffi --opt-level 3", + ) + + +def chunk_gdn_recompute_sm100( + k, + v, + gate, + beta, + cu_seqlens, + initial_state, + output_state, + checkpoint_every_n_tokens: int = 0, + output_state_checkpoints=None, + work_items=None, + work_count=None, + sched_ctr=None, + log_gate: bool = False, + *, + workspace, + stream, +) -> None: + """Execute the Blackwell chunked GDN recompute kernel (state/checkpoint-only, + THD / varlen entry). + + All tensors are contiguous, DLPack-compatible CUDA tensors on the same + device. Compile-cache-and-replay: the kernel is compiled once per static + config (dtypes, head counts, state flags) and replayed afterwards. + + Args: + k: ``(total_tokens, HK, DK)`` float16/bfloat16 + v: ``(total_tokens, HV, DK)`` float16/bfloat16 + gate: ``(total_tokens, HO)`` float32, forget gate — raw linear + alpha, or the natural-log decay when ``log_gate`` + beta: ``(total_tokens, HO)`` float32, update gate + cu_seqlens: ``(num_seqs + 1,)`` int32 + initial_state: ``(num_seqs, HO, DK, DK)`` float32/bfloat16, or None + output_state: ``(num_seqs, HO, DK, DK)`` float32/bfloat16, or None + checkpoint_every_n_tokens: emit a checkpoint entry every N tokens (0 = off) + output_state_checkpoints: ``(total_checkpoints, HO, DK, DK)`` io dtype, or None. Entry j is the + state after ``(j + 1) * N`` tokens, STRICTLY BEFORE the sequence + end -- the end-of-sequence state is only ``output_state`` + (fp32-capable). With ``N == B_T`` this is the per-chunk checkpoint + series the backward pass consumes. + work_items: ``(max_items, 8)`` int32 work-item table from + ``common/split_k.py`` (REQUIRED; an uncut table row is the whole + (b, h) sequence). Each item computes chunks ``[cstart, wend)`` + and writes checkpoints only for ``[wstart, wend)``. + work_count: ``(1,)`` int32 device-side item count (REQUIRED) + log_gate: ``gate`` holds natural-log decay values; the gate warp + skips its log2 (rescales by 1/ln2) instead of exponentiating + upstream + workspace: ``(>= tensormap_workspace_bytes(module, B) // 8,)`` int64, + 128-byte aligned; holds the per-(b,h) TMA descriptors (contents + managed here — reuse the same buffer across calls) + stream: CUDA stream handle (``cudaStream_t`` as an int) + """ + HK = k.shape[1] + HV = v.shape[1] + HO = gate.shape[1] + DK = k.shape[2] + B = cu_seqlens.shape[0] - 1 + is_GQA = HK >= HV + use_initial_state = initial_state is not None + store_final_state = output_state is not None + enable_checkpoints = checkpoint_every_n_tokens > 0 + 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 + if not (enable_checkpoints or store_final_state): + raise ValueError("output_state_checkpoints or output_state is required") + io_dtype = get_dtype(k.dtype) + + if initial_state is not None: + state_dtype_src = initial_state.dtype + elif output_state is not None: + state_dtype_src = output_state.dtype + else: + state_dtype_src = None + state_dtype = get_dtype(state_dtype_src) if state_dtype_src is not None else cutlass.Float32 + + cu_stream = cuda.CUstream(int(stream)) + + cache = get_compiled_cache( + str(k.dtype), + str(state_dtype_src), + str(cu_seqlens.dtype), + HK, + HV, + HO, + is_GQA, + use_initial_state, + store_final_state, + enable_checkpoints, + log_gate, + dyn_sched, + ) + + if "compiled" not in cache: + k_cute = from_dlpack(k, assumed_align=16) + k_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + v_cute = from_dlpack(v, assumed_align=16) + v_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + gate_cute = from_dlpack(gate, assumed_align=16) + gate_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + beta_cute = from_dlpack(beta, assumed_align=16) + beta_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + cu_seqlens_cute = from_dlpack(cu_seqlens, assumed_align=8 if str(cu_seqlens.dtype).endswith("int64") else 4).mark_layout_dynamic() + + state_in_cute = None + if use_initial_state: + state_in_cute = from_dlpack(initial_state, assumed_align=16) + state_in_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=DK) + + state_out_cute = None + if store_final_state: + state_out_cute = from_dlpack(output_state, assumed_align=16) + state_out_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=DK) + + workspace_cute = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() + + work_items_cute = from_dlpack(work_items, assumed_align=16) + work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() + + sched_ctr_cute = None + if dyn_sched: + sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + + cache["compiled"] = compile( + io_dtype, + state_dtype, + is_GQA, + use_initial_state, + store_final_state, + enable_checkpoints, + log_gate, + dyn_sched, + num_sm=multiprocessor_count(current_device_id()), + k_cute=k_cute, + v_cute=v_cute, + gate_cute=gate_cute, + beta_cute=beta_cute, + cu_seqlens_cute=cu_seqlens_cute, + state_in_cute=state_in_cute, + state_out_cute=state_out_cute, + work_items_cute=work_items_cute, + work_count_cute=work_count_cute, + sched_ctr_cute=sched_ctr_cute, + checkpoint_every_n_tokens=checkpoint_every_n_tokens, + workspace_cute=workspace_cute, + stream=cu_stream, + ) + + compiled = cache["compiled"] + + # desc build runs every execute by contract (cu contents are data; + # buffer pointers may change) — capture-safe, single tiny launch + if "build_descs" not in cache: + k_bc = from_dlpack(k, assumed_align=16) + k_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + v_bc = from_dlpack(v, assumed_align=16) + v_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + cu_bc = from_dlpack(cu_seqlens, assumed_align=8 if str(cu_seqlens.dtype).endswith("int64") else 4).mark_layout_dynamic() + checkpoints_bc = None + cu_ckpt_bc = None + if enable_checkpoints: + checkpoints_bc = from_dlpack(output_state_checkpoints, assumed_align=16) + checkpoints_bc.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + ws_bc = from_dlpack(workspace, assumed_align=128).mark_layout_dynamic() + cache["build_descs"] = cute.compile( + build_descs, + io_dtype, + CFG.B_T, + k_bc, + v_bc, + cu_bc, + checkpoints_bc, + cutlass.Int32(checkpoint_every_n_tokens if enable_checkpoints else 1), + ws_bc, + cu_stream, + options="--enable-tvm-ffi", + ) + cache["build_descs"]( + k, + v, + cu_seqlens, + output_state_checkpoints, + checkpoint_every_n_tokens if enable_checkpoints else 1, + workspace, + cu_stream, + ) + compiled( + k, + v, + gate, + beta, + cu_seqlens, + initial_state, + output_state, + work_items, + work_count, + sched_ctr, + checkpoint_every_n_tokens, + workspace, + cu_stream, + ) diff --git a/python/cudnn/linear_attention/frost/kernel/kda_bprop_config.py b/python/cudnn/linear_attention/frost/kernel/kda_bprop_config.py index 7814626c5..a1cc79914 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_bprop_config.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_bprop_config.py @@ -15,23 +15,56 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Kimi Delta Attention (KDA) Cutlass DSL backward kernel config — STUB. - -The FROST KDA backward is not implemented yet (see ``kda_bprop_f16.py``); -these constants mirror the prefill config's tile shape for when the backward -kernel lands. +"""Kimi Delta Attention (KDA) Cutlass DSL backward kernel config (fixed +compile-time constants). The BT=16 backward mirrors the prefill's 16-warp +(512-thread) specialization; the derived SMEM/TMEM sizes and offsets are +stamped by ``build_cfg`` in ``kda_bprop_f16.py``. Target arch: Blackwell SM100 (GB200) / SM103 (GB300). """ from dataclasses import dataclass +from typing import Tuple @dataclass(frozen=True) class Cfg: - B_T: int = 16 - D_K: int = 128 - D_V: int = 128 + # --- tile shape --- + B_T: int = 16 # chunk-inner token tile (BT=16 KDA schedule) + D_K: int = 128 # query/key head dim + D_V: int = 128 # value head dim + + # --- warp assignments (16 warps = 512 threads) --- + COMPUTE_GROUP_0_WARP_IDS: Tuple[int, ...] = (0, 1, 2, 3) # forward gate cumsum + decay-operand materialize + COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (4, 5, 6, 7) # value-side TMEM staging / restages / dH capture + COMPUTE_GROUP_2_WARP_IDS: Tuple[int, ...] = (8, 9, 10, 11) # dq/dk-bank drain, dG assembly + reverse cumsum + SUPER_MMA_WARP_ID: int = 12 # register-MMA KK/A/dA/dM + Neumann T_inv + TCGEN05_MMA_WARP_ID: int = 13 # tcgen05 GEMM schedule + TMA_WARP_ID: int = 14 # q/k/v/gate/do/S(H) TMA loads + EPILOGUE_WARP_ID: int = 15 # dq/dk/dv TMA stores only + + # --- register split --- + NUM_REGS_COMPUTE_GROUP_0: int = 144 + NUM_REGS_COMPUTE_GROUP_1: int = 168 + NUM_REGS_COMPUTE_GROUP_2: int = 144 + NUM_REGS_OTHER: int = 56 + + THREADS_PER_WARP: int = 32 + + BUFFER_ALIGN_BYTES: int = 1024 + + # --- SMEM / TMEM ring stage counts --- + SMEM_RAW_STAGES: int = 2 + SMEM_S_STAGES: int = 1 + SMEM_DECAY_STAGES: int = 2 + SMEM_INTERMEDIATE_STAGES: int = 2 + SMEM_STATE_SCALE_DIAG_STAGES: int = 2 + SMEM_DQ_STAGES: int = 1 + SMEM_DK_STAGES: int = 1 + SMEM_DGATE_STAGES: int = 1 + SMEM_DV_STAGES: int = 2 + + CLUSTER_SHAPE_MNK: Tuple[int, int, int] = (1, 1, 1) CFG = Cfg() 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 fd4c8a513..e9576c212 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_bprop_f16.py @@ -1,41 +1,3854 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# This kernel is derived from cuDNN, NVIDIA Corporation. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Kimi Delta Attention (KDA) Cutlass DSL backward kernel — STUB. - -The FROST KDA backward is not implemented yet. Because the KDA chunk size is -small, the design recomputes the forward per-chunk states inside the backward -(no H store in the prefill kernel). Until the backward kernel lands, this -module raises ``NotImplementedError``; the cuTile KDA engine -(``cudnn.engines.KdaCuTileEngine``) provides KDA gradients. - -Target arch: Blackwell SM100 (GB200) / SM103 (GB300). +"""Chunked Kimi Delta Attention (KDA) BPROP kernel for Blackwell SM100/SM103 +(Cutlass DSL), BT=16 tiling with a per-key-channel decay. Framework-neutral +entry ``chunk_kda_bwd_sm100``. + +Algorithm overview (per chunk c, iterated c = NT-1 .. 0; within-chunk +log2-domain gate cumsum G[t,d], eG = 2^G, eGl = 2^G[BT-1]): + Inputs : Q/K[BT,DK], V/dO[BT,DV], S = state_checkpoints[c-1] (state ENTERING chunk c, KV), + Gate[BT,DK], Beta[BT] + State : dH[DV,DK] (state gradient, fp32 TMEM, accumulated backward) + + Operands (WG0, prefill recompute): K_decay = eG.K, K_inv = K/eG, + K_restore = (eGl/eG).K, Q_decay = eG.Q, diag(eGl). + + Forward recompute: A_kk = K_decay@K_inv^T; T_inv = (I + Beta.tril(A_kk,-1))^-1 + (register Neumann); A = tril_incl(Q_decay@K_inv^T); Y = Beta.(V - S^T K_decay); + U = T_inv@Y. + + Backward math: + dU = K_restore@dH + A^T@dO (Q_decay carries scale, so A does too) + dY = T_inv^T@dU dV = Beta.dY + dA = tril_incl(dO@U^T) (unscaled) dM = dY@U^T dM_strict = Beta_row.strict(dM) + dQ = eG.scale.(dO@S^T + dA@K_inv) + dK = eG.dK_decay + dK_inv/eG + (eGl/eG).dK_restore where (sign-flipped parts) + dK_decay part = (Beta.dY)@S^T + dM_strict@K_inv (= -dK_decay) + dK_inv part = dA^T@(scale.Q_decay) - dM_strict^T@K_decay (= dK_inv; one TMEM + acc, the minus rides the staged -dM_strict tile) + dK_restore part = U@dH^T (= +dK_restore) + dBeta = sum_v dY.Y / Beta - sum_{j= 1; chunk 0 seeds from `initial_state`); dq/dk/dv io at HO heads; dgate `[T, HO, DK]` fp32 +(natural-log gate domain); dbeta `[T, HO]` fp32; d_initial_state / +d_final_state fp32 `[N, HO, DK, DV]` (K-major). + +Warp assignments (16 warps = 512 threads): + warps 0-3 : WG0 - Gate prefix scan + decay/restore operands (all chunks) + Beta scalar gather + warps 4-7 : WG1 - value-side TMEM staging, dstate capture, dBeta + warps 8-11 : WG2 - dQ/dK part drain, dGate assembly, reverse cumsum, dGate + warp 12 : super-MMA - register KK/A/dA/dM + Neumann inverse + warp 13 : tcgen05-MMA - the 15-GEMM backward schedule + warp 14 : TMA load - Q/K/V/Gate/dO/entering-state loads + warp 15 : epilogue - dQ/dK/dV/dGate TMA stores """ -from __future__ import annotations +from dataclasses import dataclass +from functools import lru_cache +from typing import NamedTuple, Optional, Type + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.experimental.cuda as cuda +import cutlass.experimental.primitives as nvvm +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + +from ..common.split_k import decode_work_item +from ..common.host import get_dtype +from cudnn.frost.buffers import current_device_id, data_ptr +from cudnn.frost.device import multiprocessor_count +from ..common.thd import TENSOR_MAP_QWORDS, emit_copy_desc, emit_checkpoint_seq_descs, emit_seq_descs +from .kda_bprop_config import CFG +from cudnn.frost.tile_dsl.barrier import ( + advance, + MBarrier, + PipelineState, + Producer, +) +from cudnn.frost.tile_dsl.handles import MmaDesc, SmemTile, tma_slice_runtime_desc +from cudnn.frost.tile_dsl.mma import mma_ss, mma_step, mma_ts_step +from cudnn.frost.tile_dsl.swizzle import swizzle_lin_S, swizzle_xor_128b +from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire +from cudnn.frost.tile_dsl.pointwise import ( + opaque_f32_zero, + f16x2_to_f32, + fmul2, + ffma2, + movmatrix_16b, + mul_f16x2, + fp32_to_fp16, + sub_f16x2, +) + +LOG2_E: float = 1.4426950408889634 + +L2_NORM_EPS: float = 1.0e-12 + + +class KdaBwdBars(NamedTuple): + """Every inter-warp handoff as an ``MBarrier`` over its ring. Consumers + track ``(idx, phase)`` inline; the producer tag selects the arrive + lowering (``TMA_LOAD``/``MMA_COMMIT``/``THREAD``). + + Buffers read by both the MMA warp and a compute/warp group carry mixed + arrive counts (one MMA commit + N thread arrivers) so the producer only + reuses the slot once every reader is done.""" + + mb_q_ready: MBarrier + mb_q_done: MBarrier + mb_k_ready: MBarrier + mb_k_done: MBarrier + mb_gate_ready: MBarrier + mb_gate_done: MBarrier + mb_do_ready: MBarrier + mb_do_done: MBarrier + mb_v_ready: MBarrier + mb_v_done: MBarrier + mb_state_ready: MBarrier + mb_state_done: MBarrier + mb_state_cg0_done: MBarrier + + mb_beta_ready: MBarrier + mb_beta_done: MBarrier + + mb_state_k_acc_ready: MBarrier + mb_du_acc_ready: MBarrier + mb_u_acc_ready: MBarrier + mb_dy_acc_ready: MBarrier + mb_dq_acc_ready: MBarrier + mb_dk_decay_part_acc_ready: MBarrier + mb_dk_inv_part_acc_ready: MBarrier + mb_dk_restore_part_acc_ready: MBarrier + mb_dqk_acc_done: MBarrier + + mb_qk_raw_ready: MBarrier + mb_qk_raw_done: MBarrier + mb_state_inp_ready: MBarrier + mb_state_inp_done: MBarrier + mb_state_inp_cg2_done: MBarrier + mb_y_inp_ready: MBarrier + mb_du_inp_ready: MBarrier + mb_neg_beta_dy_inp_ready: MBarrier + + mb_k_decay_inv_ready: MBarrier + mb_q_decay_k_restore_ready: MBarrier + mb_decay_done: MBarrier + mb_t_inv_ready: MBarrier + mb_t_inv_done: MBarrier + mb_a_ready: MBarrier + mb_a_done: MBarrier + mb_da_ready: MBarrier + mb_da_done: MBarrier + mb_dm_ready: MBarrier + mb_dm_done: MBarrier + mb_u_smem_ready: MBarrier + mb_dy_smem_ready: MBarrier + mb_dbeta_m_ready: MBarrier + + mb_dstate_acc_ready: MBarrier + mb_dstate_inp_ready: MBarrier + mb_dstate_smem_ready: MBarrier + mb_dstate_smem_done: MBarrier + mb_dstate_smem_cg2_done: MBarrier + + mb_dq_tmastg_ready: MBarrier + mb_dq_tmastg_done: MBarrier + mb_dk_tmastg_ready: MBarrier + mb_dk_tmastg_done: MBarrier + mb_dv_tmastg_ready: MBarrier + mb_dv_tmastg_done: MBarrier + mb_dgate_tmastg_ready: MBarrier + mb_dgate_tmastg_done: MBarrier + + mb_dstate0_acc_stored: MBarrier + mb_tmem_done: MBarrier + + mb_sched_ready: MBarrier + mb_sched_done: MBarrier + + +def make_kda_bwd_bars(cfg) -> KdaBwdBars: + """Bars factory. MUST be called from inside ``kernel`` (allocates the + mbarrier rings in SMEM ahead of the data buffers).""" + + def alloc(n): + return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=8) + + WARP = cfg.threads_per_warp + CG0 = len(cfg.compute_group_0_warp_ids) * WARP + CG2 = len(cfg.compute_group_2_warp_ids) * WARP + CG1 = len(cfg.compute_group_1_warp_ids) * WARP + MMA = 1 + + return KdaBwdBars( + mb_q_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_q_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0, producer=Producer.THREAD), + mb_k_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_k_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0, producer=Producer.THREAD), + mb_gate_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_gate_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0 + CG2, producer=Producer.THREAD), + mb_do_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_do_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=WARP, producer=Producer.THREAD), + mb_v_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_v_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1, producer=Producer.THREAD), + mb_state_ready=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_state_done=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_state_cg0_done=MBarrier(alloc(cfg.smem_state_stages), stages=cfg.smem_state_stages, init_count=CG0, producer=Producer.THREAD), + mb_beta_ready=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=WARP, producer=Producer.THREAD), + mb_beta_done=MBarrier(alloc(cfg.smem_beta_stages), stages=cfg.smem_beta_stages, init_count=WARP + CG1, producer=Producer.THREAD), + mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_du_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dy_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dq_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dk_decay_part_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dk_inv_part_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dk_restore_part_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dqk_acc_done=MBarrier(alloc(1), stages=1, init_count=CG2, producer=Producer.THREAD), + mb_qk_raw_ready=MBarrier(alloc(cfg.tmem_qk_raw_stages), stages=cfg.tmem_qk_raw_stages, init_count=CG0, producer=Producer.THREAD), + mb_qk_raw_done=MBarrier(alloc(cfg.tmem_qk_raw_stages), stages=cfg.tmem_qk_raw_stages, init_count=CG2, producer=Producer.THREAD), + mb_state_inp_ready=MBarrier(alloc(2), stages=2, init_count=CG0, producer=Producer.THREAD), + mb_state_inp_done=MBarrier(alloc(2), stages=2, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_state_inp_cg2_done=MBarrier(alloc(2), stages=2, init_count=CG2, producer=Producer.THREAD), + mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_du_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_neg_beta_dy_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_k_decay_inv_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0, producer=Producer.THREAD), + mb_q_decay_k_restore_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0, producer=Producer.THREAD), + mb_decay_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_t_inv_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_t_inv_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_a_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_a_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_da_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_da_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dm_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_dm_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_u_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dy_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dbeta_m_ready=MBarrier(alloc(1), stages=1, init_count=WARP, producer=Producer.THREAD), + mb_dstate_acc_ready=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dstate_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dstate_smem_ready=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_dstate_smem_done=MBarrier(alloc(1), stages=1, init_count=MMA, producer=Producer.MMA_COMMIT), + mb_dstate_smem_cg2_done=MBarrier(alloc(1), stages=1, init_count=CG2, producer=Producer.THREAD), + mb_dq_tmastg_ready=MBarrier(alloc(cfg.smem_dq_stages), stages=cfg.smem_dq_stages, init_count=CG2, producer=Producer.THREAD), + mb_dq_tmastg_done=MBarrier(alloc(cfg.smem_dq_stages), stages=cfg.smem_dq_stages, init_count=WARP, producer=Producer.THREAD), + mb_dk_tmastg_ready=MBarrier(alloc(cfg.smem_dk_stages), stages=cfg.smem_dk_stages, init_count=CG2, producer=Producer.THREAD), + mb_dk_tmastg_done=MBarrier(alloc(cfg.smem_dk_stages), stages=cfg.smem_dk_stages, init_count=WARP, producer=Producer.THREAD), + mb_dv_tmastg_ready=MBarrier(alloc(cfg.smem_dv_stages), stages=cfg.smem_dv_stages, init_count=CG1, producer=Producer.THREAD), + mb_dv_tmastg_done=MBarrier(alloc(cfg.smem_dv_stages), stages=cfg.smem_dv_stages, init_count=WARP, producer=Producer.THREAD), + mb_dgate_tmastg_ready=MBarrier(alloc(cfg.smem_dgate_stages), stages=cfg.smem_dgate_stages, init_count=CG2, producer=Producer.THREAD), + mb_dgate_tmastg_done=MBarrier(alloc(cfg.smem_dgate_stages), stages=cfg.smem_dgate_stages, init_count=WARP, producer=Producer.THREAD), + mb_dstate0_acc_stored=MBarrier(alloc(1), stages=1, init_count=CG1, producer=Producer.THREAD), + mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1 + CG2, producer=Producer.THREAD), + mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), + mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), + ) + + +# ---- Dynamic tile scheduler ------------------------------------------------------ + + +@cute.jit +def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): + """TMA-warp side: pull the next tile off the global ticket, publish it.""" + if cutlass.const_expr(cfg.dyn_sched): + bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) + if nvvm.elect_sync(): + fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sSched[sched_state.idx] = num_ctas + fetched + nvvm.bar_warp_sync(cute.arch.FULL_MASK) + next_tile = sSched[sched_state.idx] + if nvvm.elect_sync(): + bars.mb_sched_ready[sched_state.idx].arrive() + return next_tile, advance(sched_state, cfg.sched_stages) + return tile_idx + num_ctas, sched_state + + +@cute.jit +def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): + """Consumer side: read the TMA warp's published next tile.""" + if cutlass.const_expr(cfg.dyn_sched): + bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) + next_tile = sSched[sched_state.idx] + if nvvm.elect_sync(): + bars.mb_sched_done[sched_state.idx].arrive() + return next_tile, advance(sched_state, cfg.sched_stages) + return tile_idx + num_ctas, sched_state + + +# ---- Warp bodies ----------------------------------------------------------------- + + +@cute.jit +def epilogue_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_inv_raw, + sQ_decay_raw, + sDo_raw, + sU_raw, + sIntermediate_raw, + sDq_raw, + sDk_raw, + sDv_raw, + sDgate_raw, + desc_dq_base, + desc_dk_base, + desc_dv_base, + desc_dgate_base, + bars, +) -> None: + """Epilogue warp role (warp 15): the register-MMA A/dA tiles and the + dQ/dK/dV/dGate TMA stores, in chunk order with a one-behind store + ladder.""" + elect_one = nvvm.elect_sync() + + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + + # ---- ldmatrix/stmatrix lane decode ------------------------------------------- + rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) + rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) + lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) + lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) + stsm_row_coord = lane & 7 + stsm_col_coord = cutlass.Int32(0) + if (lane // 8) & 1: + stsm_row_coord = stsm_row_coord + cutlass.Int32(8) + if lane // 8 >= 2: + stsm_col_coord = cutlass.Int32(8) + stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + stsm_col_coord, bbits=1, mbase=3, sshift=3) + row_lo = lane // 4 + row_hi = row_lo + cutlass.Int32(8) + + tril_incl_mask = cutlass.Int32(0) + for accum_idx in cutlass.range_constexpr(8): + row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo + col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + if cutlass.const_expr(accum_idx % 2 == 1): + col_coord = col_coord + cutlass.Int32(1) + tril_incl_mask = tril_incl_mask | (cutlass.Int32(1 << accum_idx) if row_coord >= col_coord else cutlass.Int32(0)) + raw_index = PipelineState.start(phase=0) + u_index = PipelineState.start(phase=0) + chunk_serial_base = cutlass.Int32(0) + + sDq_tma = SmemTile( + base=sDq_raw, + elems_per_stage=(cfg.b_t * cfg.d_k), + stages=cfg.smem_dq_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.b_t * 64, + ) + sDk_tma = SmemTile( + base=sDk_raw, + elems_per_stage=(cfg.b_t * cfg.d_k), + stages=cfg.smem_dk_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.b_t * 64, + ) + sDv_tma = SmemTile( + base=sDv_raw, + elems_per_stage=(cfg.b_t * cfg.d_v), + stages=cfg.smem_dv_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.b_t * 64, + ) + sDgate_tma = SmemTile( + base=sDgate_raw, + elems_per_stage=(cfg.b_t * cfg.d_k), + stages=cfg.smem_dgate_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 32), + tma_granu_elems=32, + tma_subtile_stride_elems=cfg.b_t * 32, + ) + dq_index = PipelineState.start(phase=0) + dk_index = PipelineState.start(phase=0) + dv_index = PipelineState.start(phase=0) + dgate_index = PipelineState.start(phase=0) + sched_state = PipelineState.start(phase=0) + tile_idx = cutlass.Int32(bidx) + FIRST_STATE_CHUNK = 0 if cfg.use_initial_state else 1 + 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) + head_o = head_idx + slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) + if elect_one: + desc_dq_slot = (desc_dq_base + slot).tospace(cutlass.AddressSpace.generic) + desc_dk_slot = (desc_dk_base + slot).tospace(cutlass.AddressSpace.generic) + desc_dv_slot = (desc_dv_base + slot).tospace(cutlass.AddressSpace.generic) + desc_dgate_slot = (desc_dgate_base + slot).tospace(cutlass.AddressSpace.generic) + tma_tensormap_acquire(desc_dq_slot) + tma_tensormap_acquire(desc_dk_slot) + tma_tensormap_acquire(desc_dv_slot) + tma_tensormap_acquire(desc_dgate_slot) + num_compute_chunks = cend - wstart + pend_start = cutlass.Int32(0) + pend_writes = cutlass.Boolean(False) + for rev_idx in cutlass.range(num_compute_chunks, unroll=1): + chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_start = chunk_idx * cfg.b_t + writes = chunk_idx < wend + chunk_serial = chunk_serial_base + rev_idx + decay_stage = chunk_serial % cfg.smem_decay_stages + intermediate_stage = chunk_serial % cfg.smem_intermediate_stages + raw_stage = raw_index.idx + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sQ_decay_ptr = sQ_decay_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sDo_ptr = sDo_raw.data_ptr() + raw_stage * (cfg.d_v * cfg.b_t) + sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (cfg.intermediate_tiles * cfg.b_t * cfg.b_t) + + # ---- A = tril_incl(Q_decay @ K_inv^T) -------------------------------- + bars.mb_a_done[intermediate_stage].wait(((chunk_serial // cfg.smem_intermediate_stages) + 1) % 2) + bars.mb_q_decay_k_restore_ready[decay_stage].wait((chunk_serial // cfg.smem_decay_stages) % 2) + a_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + a_acc[accum_idx] = cutlass.Float32(0.0) + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + a_col = k_block * 16 + lhs_col_offset + a_seg = a_col // 64 + a_frag = nvvm.ldmatrix( + sQ_decay_ptr + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + b_col = k_block * 16 + rhs_col_offset + b_seg = b_col // 64 + b_frag = nvvm.ldmatrix( + sK_inv_ptr + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + mma_step( + a_acc, + (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), + (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + for accum_idx in cutlass.range_constexpr(8): + a_acc[accum_idx] = a_acc[accum_idx] if (tril_incl_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + nvvm.stmatrix( + sIntermediate_ptr + stsm_idx, + [ + fp32_to_fp16(a_acc[0], a_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[2], a_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[4], a_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[6], a_acc[7], dtype=cfg.io_dtype), + ], + nvvm.MMALayout.ROW, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_a_ready[intermediate_stage].arrive() + + # ---- dA = tril_incl(dO @ U^T) ---------------------------------------- + bars.mb_u_smem_ready.wait(u_index.phase) + u_index = advance(u_index, 1) + da_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + da_acc[accum_idx] = cutlass.Float32(0.0) + for k_block in cutlass.range_constexpr(cfg.d_v // 16): + a_col = k_block * 16 + lhs_col_offset + a_seg = a_col // 64 + a_frag = nvvm.ldmatrix( + sDo_ptr + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + b_col = k_block * 16 + rhs_col_offset + b_seg = b_col // 64 + b_frag = nvvm.ldmatrix( + sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + mma_step( + da_acc, + (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), + (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + # fence: the dO/U ldmatrix reads must complete before this release + # licenses the TMA reload (sDo) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_do_done[raw_stage].arrive() + for accum_idx in cutlass.range_constexpr(8): + da_acc[accum_idx] = da_acc[accum_idx] if (tril_incl_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + bars.mb_da_done[intermediate_stage].wait(((chunk_serial // cfg.smem_intermediate_stages) + 1) % 2) + nvvm.stmatrix( + sIntermediate_ptr + 2 * (cfg.b_t * cfg.b_t) + stsm_idx, + [ + fp32_to_fp16(da_acc[0], da_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(da_acc[2], da_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(da_acc[4], da_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(da_acc[6], da_acc[7], dtype=cfg.io_dtype), + ], + nvvm.MMALayout.ROW, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_da_ready[intermediate_stage].arrive() + raw_index = advance(raw_index, cfg.smem_raw_stages) + + # ---- dQ/dK/dGate/dV: previous chunk, one-behind store ladder --------- + if rev_idx > 0: + bars.mb_dq_tmastg_ready[dq_index.idx].wait(dq_index.phase) + if pend_writes: + desc_dq_slot = (desc_dq_base + slot).tospace(cutlass.AddressSpace.generic) + dq_slice = tma_slice_runtime_desc(desc_dq_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDq_tma[dq_index.idx], dq_slice, acquire=False) + tma_store_commit() + bars.mb_dk_tmastg_ready[dk_index.idx].wait(dk_index.phase) + if pend_writes: + desc_dk_slot = (desc_dk_base + slot).tospace(cutlass.AddressSpace.generic) + dk_slice = tma_slice_runtime_desc(desc_dk_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDk_tma[dk_index.idx], dk_slice, acquire=False) + tma_store_commit() + bars.mb_dgate_tmastg_ready[dgate_index.idx].wait(dgate_index.phase) + if pend_writes: + desc_dgate_slot = (desc_dgate_base + slot).tospace(cutlass.AddressSpace.generic) + dgate_slice = tma_slice_runtime_desc(desc_dgate_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDgate_tma[dgate_index.idx], dgate_slice, acquire=False) + tma_store_commit() + bars.mb_dv_tmastg_ready[dv_index.idx].wait(dv_index.phase) + if pend_writes: + desc_dv_slot = (desc_dv_base + slot).tospace(cutlass.AddressSpace.generic) + dv_slice = tma_slice_runtime_desc(desc_dv_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDv_tma[dv_index.idx], dv_slice, acquire=False) + tma_store_commit() + tma_store_wait(3) + bars.mb_dq_tmastg_done[dq_index.idx].arrive() + tma_store_wait(2) + bars.mb_dk_tmastg_done[dk_index.idx].arrive() + tma_store_wait(1) + bars.mb_dgate_tmastg_done[dgate_index.idx].arrive() + tma_store_wait(0) + bars.mb_dv_tmastg_done[dv_index.idx].arrive() + dq_index = advance(dq_index, cfg.smem_dq_stages) + dk_index = advance(dk_index, cfg.smem_dk_stages) + dgate_index = advance(dgate_index, cfg.smem_dgate_stages) + dv_index = advance(dv_index, cfg.smem_dv_stages) + pend_start = chunk_start + pend_writes = writes + + # ---- tile tail: drain the last chunk's dQ/dK/dGate/dV -------------------- + if num_compute_chunks > 0: + bars.mb_dq_tmastg_ready[dq_index.idx].wait(dq_index.phase) + if pend_writes: + desc_dq_slot = (desc_dq_base + slot).tospace(cutlass.AddressSpace.generic) + dq_slice = tma_slice_runtime_desc(desc_dq_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDq_tma[dq_index.idx], dq_slice, acquire=False) + tma_store_commit() + bars.mb_dk_tmastg_ready[dk_index.idx].wait(dk_index.phase) + if pend_writes: + desc_dk_slot = (desc_dk_base + slot).tospace(cutlass.AddressSpace.generic) + dk_slice = tma_slice_runtime_desc(desc_dk_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDk_tma[dk_index.idx], dk_slice, acquire=False) + tma_store_commit() + bars.mb_dgate_tmastg_ready[dgate_index.idx].wait(dgate_index.phase) + if pend_writes: + desc_dgate_slot = (desc_dgate_base + slot).tospace(cutlass.AddressSpace.generic) + dgate_slice = tma_slice_runtime_desc(desc_dgate_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDgate_tma[dgate_index.idx], dgate_slice, acquire=False) + tma_store_commit() + bars.mb_dv_tmastg_ready[dv_index.idx].wait(dv_index.phase) + if pend_writes: + desc_dv_slot = (desc_dv_base + slot).tospace(cutlass.AddressSpace.generic) + dv_slice = tma_slice_runtime_desc(desc_dv_slot, cutlass.Int32(0), head_o, pend_start) + tma_store_tile(sDv_tma[dv_index.idx], dv_slice, acquire=False) + tma_store_commit() + tma_store_wait(3) + bars.mb_dq_tmastg_done[dq_index.idx].arrive() + tma_store_wait(2) + bars.mb_dk_tmastg_done[dk_index.idx].arrive() + tma_store_wait(1) + bars.mb_dgate_tmastg_done[dgate_index.idx].arrive() + tma_store_wait(0) + bars.mb_dv_tmastg_done[dv_index.idx].arrive() + dq_index = advance(dq_index, cfg.smem_dq_stages) + dk_index = advance(dk_index, cfg.smem_dk_stages) + dgate_index = advance(dgate_index, cfg.smem_dgate_stages) + dv_index = advance(dv_index, cfg.smem_dv_stages) + + chunk_serial_base += num_compute_chunks + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def super_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_decay_raw, + sK_inv_raw, + sU_raw, + sDy_raw, + sIntermediate_raw, + sBeta_raw, + sBetaM_raw, + bars, +) -> None: + """Super-MMA warp role (warp 12): the Neumann T_inv and dM register MMAs + plus the dBeta M-term row sums, in chunk order.""" + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + sdy_index = PipelineState.start(phase=0) + + # ---- ldmatrix lane decode ---------------------------------------------------- + rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) + rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) + lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) + lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) + stsm_row_coord = lane & 7 + stsm_col_coord = cutlass.Int32(0) + if (lane // 8) & 1: + stsm_row_coord = stsm_row_coord + cutlass.Int32(8) + if lane // 8 >= 2: + stsm_col_coord = cutlass.Int32(8) + stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + stsm_col_coord, bbits=1, mbase=3, sshift=3) + row_lo = lane // 4 + row_hi = row_lo + cutlass.Int32(8) + + # hoisted tril bitmasks: bit i = row > col / row == col for accum index i + tril_strict_mask = cutlass.Int32(0) + eye_mask = cutlass.Int32(0) + for accum_idx in cutlass.range_constexpr(8): + row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo + col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + if cutlass.const_expr(accum_idx % 2 == 1): + col_coord = col_coord + cutlass.Int32(1) + tril_strict_mask = tril_strict_mask | (cutlass.Int32(1 << accum_idx) if row_coord > col_coord else cutlass.Int32(0)) + eye_mask = eye_mask | (cutlass.Int32(1 << accum_idx) if row_coord == col_coord else cutlass.Int32(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 + for rev_idx in cutlass.range(num_compute_chunks, unroll=1): + chunk_serial = chunk_serial_base + rev_idx + decay_stage = chunk_serial % cfg.smem_decay_stages + intermediate_stage = chunk_serial % cfg.smem_intermediate_stages + sBeta_ptr = sBeta_raw.data_ptr() + (chunk_serial % cfg.smem_beta_stages) * cfg.b_t + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (cfg.intermediate_tiles * cfg.b_t * cfg.b_t) + + bars.mb_t_inv_done[intermediate_stage].wait(((chunk_serial // cfg.smem_intermediate_stages) + 1) % 2) + + # ---- KK = K_decay @ K_inv^T ------------------------------------------ + bars.mb_k_decay_inv_ready[decay_stage].wait((chunk_serial // cfg.smem_decay_stages) % 2) + kk_lhs_row = lhs_row_coord + kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + kk_acc[accum_idx] = cutlass.Float32(0.0) + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + a_col = k_block * 16 + lhs_col_offset + a_seg = a_col // 64 + a_frag = nvvm.ldmatrix( + sK_decay_ptr + a_seg * (cfg.b_t * 64) + kk_lhs_row * 64 + swizzle_xor_128b(kk_lhs_row, a_col - a_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + b_col = k_block * 16 + rhs_col_offset + b_seg = b_col // 64 + b_frag = nvvm.ldmatrix( + sK_inv_ptr + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + mma_step( + kk_acc, + (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), + (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + + # ---- L = Beta * tril(KK, -1) ----------------------------------------- + bars.mb_beta_ready[chunk_serial % cfg.smem_beta_stages].wait((chunk_serial // cfg.smem_beta_stages) % 2) + beta_lo = (sBeta_ptr + row_lo).load().to(cutlass.Float32) + beta_hi = (sBeta_ptr + row_hi).load().to(cutlass.Float32) + bars.mb_beta_done[chunk_serial % cfg.smem_beta_stages].arrive() + l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + beta_scale = beta_lo if accum_idx % 4 < 2 else beta_hi + lower = kk_acc[accum_idx] if (tril_strict_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + l_regs[accum_idx] = lower * beta_scale + l_a0 = fp32_to_fp16(l_regs[0], l_regs[1], dtype=cfg.io_dtype) + l_a1 = fp32_to_fp16(l_regs[2], l_regs[3], dtype=cfg.io_dtype) + l_a2 = fp32_to_fp16(l_regs[4], l_regs[5], dtype=cfg.io_dtype) + l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) + l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) + + tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + eye = cutlass.Float32(1.0) if (eye_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + tinv_acc[accum_idx] = eye - l_values[accum_idx] + + lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) + for _round in cutlass.range_constexpr(3): + sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + sq_acc[accum_idx] = cutlass.Float32(0.0) + mma_step( + sq_acc, + (lpow_a0, lpow_a1, lpow_a2, lpow_a3), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + lpow_a0 = fp32_to_fp16(sq_acc[0], sq_acc[1], dtype=cfg.io_dtype) + lpow_a1 = fp32_to_fp16(sq_acc[2], sq_acc[3], dtype=cfg.io_dtype) + lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) + lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) + upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + upd_acc[accum_idx] = cutlass.Float32(0.0) + tinv_p0 = fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype) + tinv_p1 = fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype) + tinv_p2 = fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype) + tinv_p3 = fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype) + mma_step( + upd_acc, + (tinv_p0, tinv_p1, tinv_p2, tinv_p3), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + tinv_lo0, tinv_hi0 = f16x2_to_f32(tinv_p0, dtype=cfg.io_dtype) + tinv_lo1, tinv_hi1 = f16x2_to_f32(tinv_p1, dtype=cfg.io_dtype) + tinv_lo2, tinv_hi2 = f16x2_to_f32(tinv_p2, dtype=cfg.io_dtype) + tinv_lo3, tinv_hi3 = f16x2_to_f32(tinv_p3, dtype=cfg.io_dtype) + tinv_acc[0] = tinv_lo0 + upd_acc[0] + tinv_acc[1] = tinv_hi0 + upd_acc[1] + tinv_acc[2] = tinv_lo1 + upd_acc[2] + tinv_acc[3] = tinv_hi1 + upd_acc[3] + tinv_acc[4] = tinv_lo2 + upd_acc[4] + tinv_acc[5] = tinv_hi2 + upd_acc[5] + tinv_acc[6] = tinv_lo3 + upd_acc[6] + tinv_acc[7] = tinv_hi3 + upd_acc[7] + + nvvm.stmatrix( + sIntermediate_ptr + 1 * (cfg.b_t * cfg.b_t) + stsm_idx, + [ + fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype), + ], + nvvm.MMALayout.ROW, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_t_inv_ready[intermediate_stage].arrive() + + # ---- dM = dY @ U^T --------------------------------------------------- + bars.mb_dm_done[intermediate_stage].wait(((chunk_serial // cfg.smem_intermediate_stages) + 1) % 2) + bars.mb_dy_smem_ready.wait(sdy_index.phase) + sdy_index = advance(sdy_index, 1) + dm_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + dm_acc[accum_idx] = cutlass.Float32(0.0) + for k_block in cutlass.range_constexpr(cfg.d_v // 16): + a_col = k_block * 16 + lhs_col_offset + a_seg = a_col // 64 + a_frag = nvvm.ldmatrix( + sDy_raw.data_ptr() + a_seg * (cfg.b_t * 64) + lhs_row_coord * 64 + swizzle_xor_128b(lhs_row_coord, a_col - a_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + b_col = k_block * 16 + rhs_col_offset + b_seg = b_col // 64 + b_frag = nvvm.ldmatrix( + sU_raw.data_ptr() + b_seg * (cfg.b_t * 64) + rhs_row_coord * 64 + swizzle_xor_128b(rhs_row_coord, b_col - b_seg * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + mma_step( + dm_acc, + (a_frag[0], a_frag[1], a_frag[2], a_frag[3]), + (b_frag[0], b_frag[1], b_frag[2], b_frag[3]), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + # ---- dM_strict = Beta_row . strict(dM) ------------------------------- + dm_strict_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + beta_scale = beta_lo if accum_idx % 4 < 2 else beta_hi + val = dm_acc[accum_idx] * beta_scale if (tril_strict_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + dm_strict_regs[accum_idx] = val + w0 = fp32_to_fp16(dm_strict_regs[0], dm_strict_regs[1], dtype=cfg.io_dtype) + w1 = fp32_to_fp16(dm_strict_regs[2], dm_strict_regs[3], dtype=cfg.io_dtype) + w2 = fp32_to_fp16(dm_strict_regs[4], dm_strict_regs[5], dtype=cfg.io_dtype) + w3 = fp32_to_fp16(dm_strict_regs[6], dm_strict_regs[7], dtype=cfg.io_dtype) + nvvm.stmatrix(sIntermediate_ptr + 3 * (cfg.b_t * cfg.b_t) + stsm_idx, [w0, w1, w2, w3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) + nw0 = fp32_to_fp16(-dm_strict_regs[0], -dm_strict_regs[1], dtype=cfg.io_dtype) + nw1 = fp32_to_fp16(-dm_strict_regs[2], -dm_strict_regs[3], dtype=cfg.io_dtype) + nw2 = fp32_to_fp16(-dm_strict_regs[4], -dm_strict_regs[5], dtype=cfg.io_dtype) + nw3 = fp32_to_fp16(-dm_strict_regs[6], -dm_strict_regs[7], dtype=cfg.io_dtype) + nvvm.stmatrix(sIntermediate_ptr + 4 * (cfg.b_t * cfg.b_t) + stsm_idx, [nw0, nw1, nw2, nw3], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dm_ready[intermediate_stage].arrive() + + # ---- M-term: bsum = sum strict(dM . KK) ------------------------------ + bsum_lo = cutlass.Float32(0.0) + bsum_hi = cutlass.Float32(0.0) + for accum_idx in cutlass.range_constexpr(8): + e = dm_acc[accum_idx] * kk_acc[accum_idx] if (tril_strict_mask >> accum_idx) & 1 else cutlass.Float32(0.0) + if cutlass.const_expr(accum_idx % 4 < 2): + bsum_lo = bsum_lo + e + else: + bsum_hi = bsum_hi + e + bsum_lo = bsum_lo + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, bsum_lo, 1, 31, kind=nvvm.Shfl.BFLY)) + bsum_lo = bsum_lo + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, bsum_lo, 2, 31, kind=nvvm.Shfl.BFLY)) + bsum_hi = bsum_hi + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, bsum_hi, 1, 31, kind=nvvm.Shfl.BFLY)) + bsum_hi = bsum_hi + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, bsum_hi, 2, 31, kind=nvvm.Shfl.BFLY)) + if lane % 4 == 0: + sBetaM_raw[row_lo] = -bsum_lo + sBetaM_raw[row_hi] = -bsum_hi + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dbeta_m_ready.arrive() + chunk_serial_base += num_compute_chunks + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def tcgen05_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + tmem_base_holder, + sState_alt, + sK_decay_lead16, + sK_inv_lead16, + sK_inv_amaj, + sK_restore_lead16, + sDo_lead16, + sDo_amaj, + sQ_decay_trans, + sK_decay_trans, + sU_lead16, + sDv_lead16, + sDstate_alt, + sIntermediate, + sState_scale_diag, + bars, +) -> None: + """tcgen05-MMA warp role (warp 13): issues every tcgen05 GEMM and owns + the TMEM lifecycle.""" + elect_one = nvvm.elect_sync() + + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + nvvm.tcgen05_alloc(tmem_base_holder, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_holder.load() + bpe = cfg.io_dtype.width // 8 + + # ---- chunk-invariant GEMM descriptors ---------------------------------------- + idesc_mv_nt = 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, + ) + idesc_state_k_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_v, + a_major=1, + ) + bmm_state_k_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=True, + cta_group=1, + idesc=idesc_state_k_at, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_dvinter_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, + cta_group=1, + idesc=idesc_mv_nt, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_du_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_v, + a_major=1, + b_major=1, + ) + bmm_du_at_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + atranspose=True, + cta_group=1, + idesc=idesc_du_at, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dstate_q_at = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.d_k, + m_dim=cfg.d_v, + a_major=1, + b_major=1, + ) + bmm_dstate_q_at_desc = MmaDesc( + M=cfg.d_v, + N=cfg.d_k, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + atranspose=True, + cta_group=1, + idesc=idesc_dstate_q_at, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_qk_ts_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_mv_nt, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_mv_nt_t = 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, + b_major=1, + ) + bmm_qk_ts_t_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + cta_group=1, + idesc=idesc_mv_nt_t, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_diag = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=16, + m_dim=cfg.d_v, + ) + bmm_diag_desc = MmaDesc( + M=cfg.d_v, + N=16, + K=16, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_diag, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dstate_k = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.d_k, + m_dim=cfg.d_v, + b_major=1, + ) + bmm_dstate_k_desc = MmaDesc( + M=cfg.d_v, + N=cfg.d_k, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + cta_group=1, + idesc=idesc_dstate_k, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dstate = 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, + ) + bmm_dstate_ts_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.d_v, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_dstate, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dstate_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_dstate_at_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.d_v, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + atranspose=True, + cta_group=1, + idesc=idesc_dstate_at, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dgp = 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, + ) + bmm_dgrad_ts_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_dgp, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dgp_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_dgrad_at_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + atranspose=True, + cta_group=1, + idesc=idesc_dgp_at, + kind=nvvm.Tcgen05MMAKind.F16, + ) + idesc_dgp_at_t = 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, + b_major=1, + ) + bmm_dgrad_at_t_desc = MmaDesc( + M=cfg.d_k, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + atranspose=True, + cta_group=1, + idesc=idesc_dgp_at_t, + kind=nvvm.Tcgen05MMAKind.F16, + ) + + state_index = PipelineState.start(phase=0) + y_inp_index = PipelineState.start(phase=0) + dstate_inp_index = PipelineState.start(phase=0) + du_inp_index = PipelineState.start(phase=0) + neg_beta_dy_index = PipelineState.start(phase=0) + u_smem_index = PipelineState.start(phase=0) + dstate_smem_index = PipelineState.start(phase=0) + parts_done_index = PipelineState.start(phase=1) + + do_seg = (cfg.b_t * cfg.d_v * (cfg.io_dtype.width // 8)) >> 4 + op_seg = (cfg.b_t * cfg.d_k * (cfg.io_dtype.width // 8)) >> 4 + intermediate_seg = (cfg.intermediate_tiles * cfg.b_t * cfg.b_t * (cfg.io_dtype.width // 8)) >> 4 + intermediate_slot = (cfg.b_t * cfg.b_t * (cfg.io_dtype.width // 8)) >> 4 + diag_seg = ((cfg.d_k // 16) * 256 * (cfg.io_dtype.width // 8)) >> 4 + dv_seg = (cfg.b_t * cfg.d_v * (cfg.io_dtype.width // 8)) >> 4 + d_do_amaj0 = sDo_amaj[0].desc() + d_qd_trans0 = sQ_decay_trans[0].desc() + d_kd_trans0 = sK_decay_trans[0].desc() + d_ki_amaj0 = sK_inv_amaj[0].desc() + d_int0 = sIntermediate[0].desc() + d_kd_lead0 = sK_decay_lead16[0].desc() + d_do_lead0 = sDo_lead16[0].desc() + d_kr_lead0 = sK_restore_lead16[0].desc() + d_diag0 = sState_scale_diag[0].desc() + d_dv_lead0 = sDv_lead16[0].desc() + d_dstate_alt0 = sDstate_alt[0].desc() + d_u_lead0 = sU_lead16[0].desc() + assert cfg.smem_state_stages == 1 + d_state_alt0 = sState_alt[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 + for rev_idx in cutlass.range(num_compute_chunks, unroll=1): + chunk_serial = chunk_serial_base + rev_idx + decay_stage = chunk_serial % cfg.smem_decay_stages + intermediate_stage = chunk_serial % cfg.smem_intermediate_stages + decay_phase = (chunk_serial // cfg.smem_decay_stages) % 2 + intermediate_phase = (chunk_serial // cfg.smem_intermediate_stages) % 2 + has_dstate = cutlass.Boolean(rev_idx > 0) + if cutlass.const_expr(cfg.use_dstate_in): + has_dstate = cutlass.Boolean(True) + raw_stage_idx = chunk_serial % cfg.smem_raw_stages + + # ---- stage-derived operand descriptors ------------------------------- + decay_op_off = decay_stage * op_seg + d_do_amaj = d_do_amaj0 + raw_stage_idx * do_seg + d_qd_trans = d_qd_trans0 + decay_op_off + d_kd_trans = d_kd_trans0 + decay_op_off + d_ki_amaj = d_ki_amaj0 + decay_op_off + d_int = d_int0 + intermediate_stage * intermediate_seg + d_int_tinv = d_int + intermediate_slot + d_int_da = d_int + 2 * intermediate_slot + d_int_dm = d_int + 3 * intermediate_slot + d_int_ndm = d_int + 4 * intermediate_slot + chunk_idx = cend - cutlass.Int32(1) - rev_idx + + # ---- state_k = state(S) @ K_decay^T -------------------------------------- + bars.mb_k_decay_inv_ready[decay_stage].wait(decay_phase) + 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, + d_kd_lead0 + decay_op_off, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_k_acc_offset), cutlass.Float32), + accumulate=False, + ) + if elect_one: + bars.mb_state_k_acc_ready.arrive(cta_group=1) + bars.mb_state_done[state_index.idx].arrive(cta_group=1) + state_index = advance(state_index, cfg.smem_state_stages) + + # ---- dQ inter = state(T) @ dO^T ------------------------------------------ + bars.mb_dqk_acc_done.wait(parts_done_index.phase) + parts_done_index = advance(parts_done_index, 1) + bars.mb_state_inp_ready[chunk_serial % 2].wait((chunk_serial // 2) % 2) + bars.mb_do_ready[raw_stage_idx].wait((chunk_serial // cfg.smem_raw_stages) % 2) + if chunk_idx >= FIRST_STATE_CHUNK: + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_inp_offset + (chunk_serial % 2) * (cfg.d_v // 2)), cutlass.Int8) + b_desc = d_do_lead0 + raw_stage_idx * do_seg + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_dstate_ts_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_ts_desc.sps_B): + mma_ts_step( + bmm_dstate_ts_desc, + a_ptr.subview(sub * bmm_dstate_ts_desc.sps_B * bmm_dstate_ts_desc.tmem_advance_A), + b_desc + sub * (bmm_dstate_ts_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + + # ---- dU inter = dstate(T) @ K_restore ------------------------------------ + bars.mb_q_decay_k_restore_ready[decay_stage].wait(decay_phase) + if has_dstate: + bars.mb_dstate_inp_ready.wait(dstate_inp_index.phase) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_inp_offset), cutlass.Int8) + b_desc = d_kr_lead0 + decay_op_off + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_dvinter_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dvinter_desc.sps_B): + mma_ts_step( + bmm_dvinter_desc, + a_ptr.subview(sub * bmm_dvinter_desc.sps_B * bmm_dvinter_desc.tmem_advance_A), + b_desc + sub * (bmm_dvinter_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + + # ---- dstate decay = dstate(T) @ diag(eGl) ------------------------------------ + if has_dstate: + desc_diag = d_diag0 + decay_stage * diag_seg + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_inp_offset) + k_block * 8, cutlass.Int8) + b_desc = desc_diag.advance_start_address(k_block * 256 * 2) + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset) + k_block * 16, cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_diag_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_diag_desc.sps_B): + mma_ts_step( + bmm_diag_desc, + a_ptr.subview(sub * bmm_diag_desc.sps_B * bmm_diag_desc.tmem_advance_A), + b_desc + sub * (bmm_diag_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + dstate_inp_index = advance(dstate_inp_index, 1) + + # ---- dU intra += dO^T(S) @ A ----------------------------------------- + bars.mb_a_ready[intermediate_stage].wait(intermediate_phase) + mma_ss( + bmm_du_at_desc, + d_do_amaj, + d_int, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_acc_offset), cutlass.Float32), + accumulate=has_dstate, + ) + if elect_one: + bars.mb_du_acc_ready.arrive(cta_group=1) + bars.mb_a_done[intermediate_stage].arrive(cta_group=1) + + # ---- dstate q-term += dO^T(S) @ Q_decay ---------------------------------- + mma_ss( + bmm_dstate_q_at_desc, + d_do_amaj, + d_qd_trans, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset), cutlass.Float32), + accumulate=has_dstate, + ) + + # ---- U = Y(T) @ T_inv ------------------------------------------------ + bars.mb_t_inv_ready[intermediate_stage].wait(intermediate_phase) + bars.mb_y_inp_ready.wait(y_inp_index.phase) + y_inp_index = advance(y_inp_index, 1) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_y_inp_offset), cutlass.Int8) + b_desc = d_int_tinv + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_u_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_qk_ts_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_qk_ts_desc.sps_B): + mma_ts_step( + bmm_qk_ts_desc, + a_ptr.subview(sub * bmm_qk_ts_desc.sps_B * bmm_qk_ts_desc.tmem_advance_A), + b_desc + sub * (bmm_qk_ts_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + if elect_one: + bars.mb_u_acc_ready.arrive(cta_group=1) + + # ---- dY = dU(T) @ T_inv ---------------------------------------------- + bars.mb_du_inp_ready.wait(du_inp_index.phase) + du_inp_index = advance(du_inp_index, 1) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_du_inp_offset), cutlass.Int8) + dy_b_desc = d_int_tinv + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dy_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_qk_ts_t_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_qk_ts_t_desc.sps_B): + mma_ts_step( + bmm_qk_ts_t_desc, + a_ptr.subview(sub * bmm_qk_ts_t_desc.sps_B * bmm_qk_ts_t_desc.tmem_advance_A), + dy_b_desc + sub * (bmm_qk_ts_t_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + if elect_one: + bars.mb_dy_acc_ready.arrive(cta_group=1) + bars.mb_t_inv_done[intermediate_stage].arrive(cta_group=1) + + # ---- dK_restore part = dstate(S) @ U^T ----------------------------------- + bars.mb_u_smem_ready.wait(u_smem_index.phase) + u_smem_index = advance(u_smem_index, 1) + if has_dstate: + bars.mb_dstate_smem_ready.wait(dstate_smem_index.phase) + dstate_smem_index = advance(dstate_smem_index, 1) + mma_ss( + bmm_dstate_at_desc, + d_dstate_alt0, + d_u_lead0, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_restore_acc_offset), cutlass.Float32), + accumulate=False, + ) + if elect_one: + bars.mb_dk_restore_part_acc_ready.arrive(cta_group=1) + bars.mb_dstate_smem_done.arrive(cta_group=1) + + # ---- dstate dY-term += -Beta.dY(T) @ K_decay ----------------------------- + bars.mb_neg_beta_dy_inp_ready.wait(neg_beta_dy_index.phase) + neg_beta_dy_index = advance(neg_beta_dy_index, 1) + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_neg_beta_dy_inp_offset), cutlass.Int8) + b_desc = d_kd_trans + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dstate_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_dstate_k_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_k_desc.sps_B): + mma_ts_step( + bmm_dstate_k_desc, + a_ptr.subview(sub * bmm_dstate_k_desc.sps_B * bmm_dstate_k_desc.tmem_advance_A), + b_desc + sub * (bmm_dstate_k_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(True), + ) + if elect_one: + bars.mb_dstate_acc_ready.arrive(cta_group=1) + + # ---- dK_inv part = scale.Q_decay^T(S) @ dA --------------------------- + bars.mb_da_ready[intermediate_stage].wait(intermediate_phase) + mma_ss( + bmm_dgrad_at_t_desc, + d_qd_trans, + d_int_da, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), + accumulate=False, + ) + + # ---- dQ attn += K_inv^T(S) @ dA^T ------------------------------------ + if chunk_idx >= FIRST_STATE_CHUNK: + mma_ss( + bmm_dgrad_at_desc, + d_ki_amaj, + d_int_da, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32), + accumulate=True, + ) + if chunk_idx < FIRST_STATE_CHUNK: + mma_ss( + bmm_dgrad_at_desc, + d_ki_amaj, + d_int_da, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dq_acc_offset), cutlass.Float32), + accumulate=False, + ) + if elect_one: + bars.mb_dq_acc_ready.arrive(cta_group=1) + bars.mb_da_done[intermediate_stage].arrive(cta_group=1) + + # ---- dK_decay part = state(T) @ (Beta.dY)^T ------------------------------ + bars.mb_dv_tmastg_ready[chunk_serial % cfg.smem_dv_stages].wait((chunk_serial // cfg.smem_dv_stages) % 2) + if chunk_idx >= FIRST_STATE_CHUNK: + a_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_state_inp_offset + (chunk_serial % 2) * (cfg.d_v // 2)), cutlass.Int8) + b_desc = d_dv_lead0 + (chunk_serial % cfg.smem_dv_stages) * dv_seg + c_ptr = nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32) + for sub in cutlass.range_constexpr(bmm_dstate_ts_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_dstate_ts_desc.sps_B): + mma_ts_step( + bmm_dstate_ts_desc, + a_ptr.subview(sub * bmm_dstate_ts_desc.sps_B * bmm_dstate_ts_desc.tmem_advance_A), + b_desc + sub * (bmm_dstate_ts_desc.smem_subtile_B >> 4), + c_ptr, + k, + cutlass.Boolean(sub + k > 0), + ) + if elect_one: + bars.mb_state_inp_done[chunk_serial % 2].arrive(cta_group=1) + + # ---- dK_inv part += K_decay^T(S) @ -dM_strict ------------------------ + bars.mb_dm_ready[intermediate_stage].wait(intermediate_phase) + mma_ss( + bmm_dgrad_at_t_desc, + d_kd_trans, + d_int_ndm, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), + accumulate=True, + ) + if elect_one: + bars.mb_dk_inv_part_acc_ready.arrive(cta_group=1) + + # ---- dK_decay part += K_inv^T(S) @ dM_strict^T ----------------------- + if chunk_idx >= FIRST_STATE_CHUNK: + mma_ss( + bmm_dgrad_at_desc, + d_ki_amaj, + d_int_dm, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), + accumulate=True, + ) + if chunk_idx < FIRST_STATE_CHUNK: + mma_ss( + bmm_dgrad_at_desc, + d_ki_amaj, + d_int_dm, + nvvm.make_tmem_ptr((tmem_base + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), + accumulate=False, + ) + if elect_one: + bars.mb_dk_decay_part_acc_ready.arrive(cta_group=1) + bars.mb_dm_done[intermediate_stage].arrive(cta_group=1) + bars.mb_decay_done[decay_stage].arrive(cta_group=1) + + # ---- tile end: WG1's dstate0 drain gates the next tile's dstate reuse ------------ + bars.mb_dstate0_acc_stored.wait(dstate0_index.phase) + dstate0_index = advance(dstate0_index, 1) + chunk_serial_base += num_compute_chunks + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + bars.mb_tmem_done[0].wait(0) + nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) + nvvm.tcgen05_dealloc( + nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), + cutlass.Int32(512), + group=nvvm.CTAGroup.CTA_1, + ) + + +@cute.jit +def tmaldg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + mSched, + sSched, + lane, + sQ_raw, + sK_raw, + sV_raw, + sGate_raw, + sDo_raw, + sState_raw, + desc_q_base, + desc_k_base, + desc_v_base, + 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 + every G->S TMA load.""" + elect_one = nvvm.elect_sync() + + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + + sQ_tma = SmemTile( + base=sQ_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sK_tma = SmemTile( + base=sK_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sV_tma = SmemTile( + base=sV_raw, + elems_per_stage=(cfg.d_v * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sGate_tma = SmemTile( + base=sGate_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 32), + tma_granu_elems=32, + tma_subtile_stride_elems=(cfg.b_t * 32), + ) + sDo_tma = SmemTile( + base=sDo_raw, + elems_per_stage=(cfg.d_v * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sState_tma = SmemTile( + base=sState_raw, + elems_per_stage=(cfg.d_k * cfg.d_v), + stages=cfg.smem_state_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.d_k * 64, + ) + raw_index = PipelineState.start(phase=1) + state_index = PipelineState.start(phase=1) + 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) + head_o = head_idx + head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) + head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) + head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) + slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_q_slot = (desc_q_base + slot).tospace(cutlass.AddressSpace.generic) + desc_k_slot = (desc_k_base + slot).tospace(cutlass.AddressSpace.generic) + desc_v_slot = (desc_v_base + slot).tospace(cutlass.AddressSpace.generic) + 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) + tma_tensormap_acquire(desc_v_slot) + 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 + chunk_start = chunk_idx * cfg.b_t + + # ---- Q load ---------------------------------------------------------- + bars.mb_q_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_q_ready[raw_index.idx].arrive(n_bytes=cfg.tma_q_bytes) + q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, chunk_start) + tma_load_tile(sQ_tma[raw_index.idx], q_slice, bars.mb_q_ready[raw_index.idx].smem_ptr, acquire=False) + + # ---- K load ---------------------------------------------------------- + bars.mb_k_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_k_ready[raw_index.idx].arrive(n_bytes=cfg.tma_k_bytes) + k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) + tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_index.idx].smem_ptr, acquire=False) + + # ---- Gate load ------------------------------------------------------- + bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_gate_ready[raw_index.idx].arrive(n_bytes=cfg.tma_gate_bytes) + gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_index.idx].smem_ptr, acquire=False) + + # ---- dO load --------------------------------------------------------- + bars.mb_do_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_do_ready[raw_index.idx].arrive(n_bytes=cfg.tma_do_bytes) + do_slice = tma_slice_runtime_desc(desc_do_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sDo_tma[raw_index.idx], do_slice, bars.mb_do_ready[raw_index.idx].smem_ptr, acquire=False) + + # ---- V load ---------------------------------------------------------- + bars.mb_v_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_v_ready[raw_index.idx].arrive(n_bytes=cfg.tma_v_bytes) + 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 ---- + if chunk_idx >= FIRST_STATE_CHUNK: + state_idx = state_index.idx + bars.mb_state_cg0_done[state_idx].wait(state_index.phase) + bars.mb_state_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(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) + raw_index = advance(raw_index, cfg.smem_raw_stages) + tile_idx = next_tile + + +@cute.jit +def compute0_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_base_holder, + warp_idx, + scale, + mBeta, + sBeta_raw, + sK_inv_raw, + sGate_raw, + sK_raw, + sQ_raw, + sState_raw, + sNorm_raw, + sK_decay_raw, + sK_restore_raw, + sQ_decay_raw, + sState_scale_diag_raw, + bars, +) -> None: + """WG0 warp role (warps 0-3): persistent tile-scheduler loop + gate prefix + scan and the decay/restore operand materialization into tcgen05 SMEM for + EVERY chunk (no ping-pong: the backward pipeline is drain-bound). Also + stashes the per-row Q/K inverse norms for WG2's dGate assembly and copies + H -> TMEM f16 at the chunk tail.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_holder.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) + value_dim = tmem_subpartition * cfg.threads_per_warp + lane + row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + state_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 + for rev_idx in cutlass.range(num_compute_chunks, unroll=1): + chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_serial = chunk_serial_base + rev_idx + chunk_start = chunk_idx * cfg.b_t + decay_stage = chunk_serial % cfg.smem_decay_stages + raw_stage = chunk_serial % cfg.smem_raw_stages + sQ_ptr = sQ_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sK_ptr = sK_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sGate_ptr = sGate_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sQ_decay_ptr = sQ_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sK_restore_ptr = sK_restore_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sState_scale_diag_ptr = sState_scale_diag_raw.data_ptr() + decay_stage * ((cfg.d_k // 16) * 256) + + # ---- beta scalars: gathered in the inputs-wait shadow ---------------- + if cg0_warp == 0: + beta_stage = chunk_serial % cfg.smem_beta_stages + bars.mb_beta_done[beta_stage].wait(((chunk_serial // cfg.smem_beta_stages) + 1) % 2) + if lane < cfg.b_t: + token_idx = chunk_start + lane + beta_value = cutlass.Float32(0.0) + if token_idx < seqlen_b: + beta_value = mBeta[batch_start + token_idx, head_idx].to(cutlass.Float32) + sBeta_raw[beta_stage * cfg.b_t + lane] = beta_value + bars.mb_beta_ready[beta_stage].arrive() + bars.mb_gate_ready[raw_stage].wait((chunk_serial // cfg.smem_raw_stages) % 2) + bars.mb_q_ready[raw_stage].wait((chunk_serial // cfg.smem_raw_stages) % 2) + bars.mb_k_ready[raw_stage].wait((chunk_serial // cfg.smem_raw_stages) % 2) -def get_workspace_size(B: int, HQ: int, HV: int) -> int: - return 0 + row_group_start = cg0_warp * (cfg.b_t // len(cfg.compute_group_0_warp_ids)) + lane_row_group = lane // 8 + lane_in_row_group = lane - lane_row_group * 8 + decay_row = row_group_start + lane_row_group + g_prefix_ptr = sGate_ptr + prefix_dim = cg0_warp * cfg.threads_per_warp + lane + # ---- gate prefix scan: cumulative log-gate per key channel ----------- + gate_raw = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + for row in cutlass.range_constexpr(cfg.b_t): + f32_segment = prefix_dim // 32 + f32_segment_dim = prefix_dim - f32_segment * 32 + prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) + gate_raw[row] = (sGate_ptr + prefix_idx).load() + g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + for row in cutlass.range_constexpr(cfg.b_t): + gate = gate_raw[row] + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) + if token_idx < seqlen_b: + gate = gate * cutlass.Float32(LOG2_E) + else: + gate = cutlass.Float32(0.0) + g_prefix_regs[row] = gate -def chunk_kda_bwd_sm100(*args, **kwargs) -> None: - """Not implemented — the FROST KDA backward kernel is a stub.""" - raise NotImplementedError( - "FROST KDA backward is not implemented yet (recompute-in-bprop kernel is a stub); " - "use the cuTile KDA engine (cudnn.engines.KdaCuTileEngine) for KDA gradients." + prefix_acc = cutlass.Float32(0.0) + for row_pair in cutlass.range_constexpr(cfg.b_t // 2): + row0 = row_pair * 2 + row1 = row0 + 1 + gate0 = g_prefix_regs[row0] + gate1 = g_prefix_regs[row1] + pair_vec = nvvm.add_packed_f32x2( + cutlass.Vector.from_elements((prefix_acc, gate0), cutlass.Float32), + cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32), + ftz=False, + rnd="rn", + ) + prefix0, row_pair_sum = cutlass.Float32(pair_vec[0]), cutlass.Float32(pair_vec[1]) + prefix1 = prefix_acc + row_pair_sum + g_prefix_regs[row0] = prefix0 + g_prefix_regs[row1] = prefix1 + prefix_acc = prefix1 + + for row in cutlass.range_constexpr(cfg.b_t): + g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) + + exp_g_last = g_prefix_regs[cfg.b_t - 1] + # ---- decay-slot guard: previous use fully consumed ------------------- + operand_done_phase = ((chunk_serial // cfg.smem_decay_stages) + 1) % 2 + bars.mb_decay_done[decay_stage].wait(operand_done_phase) + + for row in cutlass.range_constexpr(cfg.b_t): + f32_segment = prefix_dim // 32 + f32_segment_dim = prefix_dim - f32_segment * 32 + prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) + (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) + + # ---- state-scale diag: stage exp2(g_last) decay blocks --------------- + block = prefix_dim // cutlass.Int32(16) + coord = prefix_dim - block * cutlass.Int32(16) + linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + coord + diag_idx = swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) + sState_scale_diag_ptr[diag_idx] = exp_g_last.to(cfg.io_dtype) + + # ---- raw Q/K: SMEM -> TMEM ring (channel-major, for WG2) ------------- + qk_raw_stage = chunk_serial % cfg.tmem_qk_raw_stages + bars.mb_qk_raw_done[qk_raw_stage].wait(((chunk_serial // cfg.tmem_qk_raw_stages) + 1) % 2) + raw_seg = prefix_dim // 64 + raw_dim = prefix_dim - raw_seg * 64 + q_raw_words = cutlass.Array(cutlass.Int32, cfg.b_t // 2, alignment=16) + k_raw_words = cutlass.Array(cutlass.Int32, cfg.b_t // 2, alignment=16) + for t2 in cutlass.range_constexpr(cfg.b_t // 2): + t0 = 2 * t2 + ridx0 = raw_seg * (cfg.b_t * 64) + t0 * 64 + swizzle_xor_128b(t0, raw_dim, elem_bytes=2) + ridx1 = raw_seg * (cfg.b_t * 64) + (t0 + 1) * 64 + swizzle_xor_128b(t0 + 1, raw_dim, elem_bytes=2) + q0 = (sQ_ptr + ridx0).load().to(cutlass.Float32) + q1 = (sQ_ptr + ridx1).load().to(cutlass.Float32) + k0 = (sK_ptr + ridx0).load().to(cutlass.Float32) + k1 = (sK_ptr + ridx1).load().to(cutlass.Float32) + q_raw_words[t2] = fp32_to_fp16(q0, q1, dtype=cfg.io_dtype) + k_raw_words[t2] = fp32_to_fp16(k0, k1, dtype=cfg.io_dtype) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_qraw_inp_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), + q_raw_words[0 : (cfg.b_t // 2)], + ) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_kraw_inp_offset + qk_raw_stage * (cfg.b_t // 2)), cutlass.Int8), + k_raw_words[0 : (cfg.b_t // 2)], + ) + nvvm.tcgen05_wait("store") + bars.mb_qk_raw_ready[qk_raw_stage].arrive() + + nvvm.barrier_cta_sync(cfg.cg0_sync_barrier_id, thread_count=cfg.cg0_threads) + + k_inv_pack = cutlass.Array(cutlass.Int32, 2 * 4, alignment=16) + raw_q_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + # ---- optional Q/K L2-norm -------------------------------------------- + if cutlass.const_expr(cfg.l2norm): + qk0_lo = opaque_f32_zero() + qk0_hi = opaque_f32_zero() + qk1_lo = opaque_f32_zero() + qk1_hi = opaque_f32_zero() + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + raw_f16_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + raw_q_frag = (sQ_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_k_frag = (sK_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_q_frag_f32 = raw_q_frag.to(cutlass.Float32) + raw_k_frag_f32 = raw_k_frag.to(cutlass.Float32) + for dim_offset in cutlass.range_constexpr(8): + q_val = raw_q_frag_f32[dim_offset] + k_val = raw_k_frag_f32[dim_offset] + raw_q_regs[reg_base + dim_offset] = q_val + raw_k_regs[reg_base + dim_offset] = k_val + if cutlass.const_expr(cfg.l2norm): + if cutlass.const_expr(dim_offset % 2 == 0): + qk0_lo, qk0_hi = ffma2(q_val, k_val, q_val, k_val, qk0_lo, qk0_hi) + else: + qk1_lo, qk1_hi = ffma2(q_val, k_val, q_val, k_val, qk1_lo, qk1_hi) + + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_q_done[raw_stage].arrive() + bars.mb_k_done[raw_stage].arrive() + + q_inv_norm = opaque_f32_zero() + cutlass.Float32(1.0) + k_inv_norm = opaque_f32_zero() + cutlass.Float32(1.0) + if cutlass.const_expr(cfg.l2norm): + q_sum_sq = qk0_lo + qk1_lo + k_sum_sq = qk0_hi + qk1_hi + q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 4, 31, kind=nvvm.Shfl.BFLY)) + q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 2, 31, kind=nvvm.Shfl.BFLY)) + q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 1, 31, kind=nvvm.Shfl.BFLY)) + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 4, 31, kind=nvvm.Shfl.BFLY)) + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 2, 31, kind=nvvm.Shfl.BFLY)) + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 1, 31, kind=nvvm.Shfl.BFLY)) + norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS) + q_inv_norm = cute.math.rsqrt(cute.math.max(q_sum_sq, norm_floor_sq), fastmath=True) + k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) + if lane_in_row_group == 0: + sNorm_raw[(chunk_serial % cfg.tmem_qk_raw_stages) * (2 * cfg.b_t) + decay_row] = q_inv_norm + sNorm_raw[(chunk_serial % cfg.tmem_qk_raw_stages) * (2 * cfg.b_t) + cfg.b_t + decay_row] = k_inv_norm + q_stage_norm = q_inv_norm * scale + + # ---- decay/restore operands: exp2(+-g) applied per key channel ------- + exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + for f32_group in cutlass.range_constexpr(2): + f32_dim_base = dim_base + f32_group * 4 + f32_segment = f32_dim_base // 32 + f32_segment_dim = f32_dim_base - f32_segment * 32 + g_prefix_idx = f32_segment * (cfg.b_t * 32) + decay_row * 32 + swizzle_xor_128b(decay_row, f32_segment_dim, elem_bytes=4) + exp_g_frag = (g_prefix_ptr + g_prefix_idx).load(count=4, alignment=16) + exp_g_last_idx = f32_segment * (cfg.b_t * 32) + (cfg.b_t - 1) * 32 + swizzle_xor_128b((cfg.b_t - 1), f32_segment_dim, elem_bytes=4) + exp_g_last_frag = (g_prefix_ptr + exp_g_last_idx).load(count=4, alignment=16) + f32_reg_base = reg_base + f32_group * 4 + exp_g_regs[f32_reg_base] = exp_g_frag[0] + exp_g_regs[f32_reg_base + 1] = exp_g_frag[1] + exp_g_regs[f32_reg_base + 2] = exp_g_frag[2] + exp_g_regs[f32_reg_base + 3] = exp_g_frag[3] + exp_g_last_regs[f32_reg_base] = exp_g_last_frag[0] + exp_g_last_regs[f32_reg_base + 1] = exp_g_last_frag[1] + exp_g_last_regs[f32_reg_base + 2] = exp_g_last_frag[2] + exp_g_last_regs[f32_reg_base + 3] = exp_g_last_frag[3] + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_gate_done[raw_stage].arrive() + + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + # ---- K decay + K inv operands: exp2(+g) * K / exp2(-g) * K ------- + k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + for pair_idx in cutlass.range_constexpr(4): + dim0 = pair_idx * 2 + dim1 = dim0 + 1 + raw_reg_idx0 = reg_base + dim0 + raw_reg_idx1 = reg_base + dim1 + k_value0, k_value1 = fmul2(raw_k_regs[raw_reg_idx0], raw_k_regs[raw_reg_idx1], k_inv_norm, k_inv_norm) + k_pair = fp32_to_fp16(k_value0, k_value1, dtype=cfg.io_dtype) + exp_g_pair = fp32_to_fp16(exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1], dtype=cfg.io_dtype) + k_decay_pack[pair_idx] = mul_f16x2(k_pair, exp_g_pair, cfg.io_dtype) + exp_neg_g0 = cute.math.rcp(exp_g_regs[raw_reg_idx0], approx=True, ftz=True) + exp_neg_g1 = cute.math.rcp(exp_g_regs[raw_reg_idx1], approx=True, ftz=True) + exp_neg_pair = fp32_to_fp16(exp_neg_g0, exp_neg_g1, dtype=cfg.io_dtype) + k_inv_pack[dim_half * 4 + pair_idx] = mul_f16x2(k_pair, exp_neg_pair, cfg.io_dtype) + + k_inv_vec = cutlass.Vector.from_elements( + ( + k_inv_pack[dim_half * 4], + k_inv_pack[dim_half * 4 + 1], + k_inv_pack[dim_half * 4 + 2], + k_inv_pack[dim_half * 4 + 3], + ), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + k_decay_vec = cutlass.Vector.from_elements( + (k_decay_pack[0], k_decay_pack[1], k_decay_pack[2], k_decay_pack[3]), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + op_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + (sK_inv_ptr + op_idx).store(k_inv_vec, alignment=16) + (sK_decay_ptr + op_idx).store(k_decay_vec, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_k_decay_inv_ready[decay_stage].arrive() + + # ---- Q decay + K_restore operands ------------------------------------ + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + q_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + k_restore_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + for pair_idx in cutlass.range_constexpr(4): + dim0 = pair_idx * 2 + dim1 = dim0 + 1 + raw_reg_idx0 = reg_base + dim0 + raw_reg_idx1 = reg_base + dim1 + q_value0, q_value1 = fmul2(raw_q_regs[raw_reg_idx0], raw_q_regs[raw_reg_idx1], q_stage_norm, q_stage_norm) + q_pair = fp32_to_fp16(q_value0, q_value1, dtype=cfg.io_dtype) + exp_g_pair = fp32_to_fp16(exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1], dtype=cfg.io_dtype) + q_decay_pack[pair_idx] = mul_f16x2(q_pair, exp_g_pair, cfg.io_dtype) + exp_g_last_pair = fp32_to_fp16(exp_g_last_regs[raw_reg_idx0], exp_g_last_regs[raw_reg_idx1], dtype=cfg.io_dtype) + k_restore_pack[pair_idx] = mul_f16x2(k_inv_pack[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) + + q_decay_vec = cutlass.Vector.from_elements( + (q_decay_pack[0], q_decay_pack[1], q_decay_pack[2], q_decay_pack[3]), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + k_restore_vec = cutlass.Vector.from_elements( + (k_restore_pack[0], k_restore_pack[1], k_restore_pack[2], k_restore_pack[3]), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + op_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + (sQ_decay_ptr + op_idx).store(q_decay_vec, alignment=16) + (sK_restore_ptr + op_idx).store(k_restore_vec, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_q_decay_k_restore_ready[decay_stage].arrive() + + # ---- state copy: SMEM -> TMEM f16 ---------------------------------------- + bars.mb_state_inp_done[chunk_serial % 2].wait(((chunk_serial // 2) + 1) % 2) + bars.mb_state_inp_cg2_done[chunk_serial % 2].wait(((chunk_serial // 2) + 1) % 2) + if chunk_idx >= FIRST_STATE_CHUNK: + bars.mb_state_ready[state_index.idx].wait(state_index.phase) + 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 + ) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr( + row_addr + (tmem_col + cfg.tmem_state_inp_offset + (chunk_serial % 2) * (cfg.d_v // 2) + v_seg * 32 + v_col8 * 4), cutlass.Int8 + ), + state_frag.bitcast(cutlass.Int32), + ) + nvvm.tcgen05_wait("store") + bars.mb_state_cg0_done[state_index.idx].arrive() + state_index = advance(state_index, cfg.smem_state_stages) + bars.mb_state_inp_ready[chunk_serial % 2].arrive() + chunk_serial_base += num_compute_chunks + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def compute1_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_base_holder, + warp_idx, + mDbeta, + mDstate0, + mDstate_in, + sV_raw, + sDo_raw, + sBeta_raw, + sU_raw, + sDy_raw, + sDv_raw, + sDstate_raw, + sRed_raw, + sBetaM_raw, + bars, +) -> None: + """WG1 warp role (warps 4-7): the value-side TMEM staging.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_holder.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) + token_row_coord = (lane // 16) * 8 + (lane & 7) + value_col_offset = ((lane // 8) & 1) * 8 + value_dim = tmem_subpartition * cfg.threads_per_warp + lane + value_dim_base = tmem_subpartition * cfg.threads_per_warp + cg1_tidx = warp_idx % 4 * cfg.threads_per_warp + lane + + raw_index = PipelineState.start(phase=0) + state_k_index = PipelineState.start(phase=0) + u_acc_index = PipelineState.start(phase=0) + du_acc_index = PipelineState.start(phase=0) + dy_acc_index = PipelineState.start(phase=0) + dstate_ready_index = PipelineState.start(phase=0) + dstate_smem_done_index = PipelineState.start(phase=1) + dbeta_m_index = PipelineState.start(phase=0) + dv_done_index = PipelineState.start(phase=1) + + 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 + + # ---- dht seeding: dh acc + dh_inp f16 + sdH ------------------------------ + if cutlass.const_expr(cfg.use_dstate_in): + if num_compute_chunks > 0: + seed_true = cend == num_chunks_b + bars.mb_dstate_smem_done.wait(dstate_smem_done_index.phase) + 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 + 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 + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 16), cutlass.Float32), + seed_block[0:16], + ) + seed_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) + for pc in cutlass.range_constexpr(8): + seed_pack[pc] = fp32_to_fp16(seed_block[2 * pc], seed_block[2 * pc + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 8), cutlass.Int8), + seed_pack[0:8], + ) + nvvm.tcgen05_wait("store") + bars.mb_dstate_inp_ready.arrive() + + # ---- dht seed -> sdH: re-read dh_inp after the TMEM publish ------ + for sub in cutlass.range_constexpr(cfg.d_k // 16): + dstate_words = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 8), cutlass.Float32), num=8 + ) + for half in cutlass.range_constexpr(2): + d_base = sub * 16 + half * 8 + h_vec = cutlass.Vector.from_elements( + (dstate_words[half * 4], dstate_words[half * 4 + 1], dstate_words[half * 4 + 2], dstate_words[half * 4 + 3]), + cutlass.Float32, + ).bitcast(cfg.io_dtype) + h_addr = (d_base // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, d_base % 64, elem_bytes=2) + (sDstate_raw.data_ptr() + h_addr).store(h_vec, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dstate_smem_ready.arrive() + + for rev_idx in cutlass.range(num_compute_chunks, unroll=1): + chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_serial = chunk_serial_base + rev_idx + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sDo_ptr = sDo_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sBeta_ptr = sBeta_raw.data_ptr() + (chunk_serial % cfg.smem_beta_stages) * cfg.b_t + row_addr_lo = tmem_row << 16 + row_addr_hi = (tmem_row + 16) << 16 + row_id0 = tmem_row + value_dim_base + row_id1 = row_id0 + 16 + has_dstate = cutlass.Boolean(rev_idx > 0) + if cutlass.const_expr(cfg.use_dstate_in): + has_dstate = cutlass.Boolean(True) + + # ---- Y staging: Y = Beta * (V - state_k) ----------------------------- + bars.mb_v_ready[raw_index.idx].wait((chunk_serial // cfg.smem_raw_stages) % 2) + projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset + input_col_id = tmem_col + cfg.tmem_y_inp_offset + raw_v_frag0 = nvvm.ldmatrix( + sV_ptr + + (value_dim_base + value_col_offset) // 64 * (cfg.b_t * 64) + + token_row_coord * 64 + + swizzle_xor_128b(token_row_coord, (value_dim_base + value_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + raw_v_frag1 = nvvm.ldmatrix( + sV_ptr + + (value_dim_base + 16 + value_col_offset) // 64 * (cfg.b_t * 64) + + token_row_coord * 64 + + swizzle_xor_128b(token_row_coord, (value_dim_base + 16 + value_col_offset) % 64, elem_bytes=2), + 4, + nvvm.MMALayout.COL, + ) + # fence: the V ldmatrix reads must complete before this release + # licenses the TMA reload (sV) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_v_done[raw_index.idx].arrive() + + bars.mb_beta_ready[chunk_serial % cfg.smem_beta_stages].wait((chunk_serial // cfg.smem_beta_stages) % 2) + beta_pairs = cutlass.Array(cutlass.Int32, 2, space=cutlass.AddressSpace.rmem) + for half in cutlass.range_constexpr(2): + token0 = ((half * 4 + (lane & 3)) ^ 4) * 2 + beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) + beta1 = (sBeta_ptr + token0 + 1).load().to(cutlass.Float32) + beta_pairs[half] = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) + diff_w0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + diff_w1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + if chunk_idx >= FIRST_STATE_CHUNK: + bars.mb_state_k_acc_ready.wait(state_k_index.phase) + state_k_index = advance(state_k_index, 1) + state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) + state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + state_k_pair = fp32_to_fp16(state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1], dtype=cfg.io_dtype) + diff_pair = sub_f16x2(raw_v_frag0[raw_matrix], state_k_pair, cfg.io_dtype) + diff_w0[reg_idx ^ 2] = diff_pair + y_inp_pack0[reg_idx ^ 2] = mul_f16x2(beta_pairs[reg_idx // 2], diff_pair, cfg.io_dtype) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + state_k_pair = fp32_to_fp16(state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1], dtype=cfg.io_dtype) + diff_pair = sub_f16x2(raw_v_frag1[raw_matrix], state_k_pair, cfg.io_dtype) + diff_w1[reg_idx ^ 2] = diff_pair + y_inp_pack1[reg_idx ^ 2] = mul_f16x2(beta_pairs[reg_idx // 2], diff_pair, cfg.io_dtype) + if chunk_idx < FIRST_STATE_CHUNK: + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + diff_w0[reg_idx ^ 2] = raw_v_frag0[raw_matrix] + y_inp_pack0[reg_idx ^ 2] = mul_f16x2(beta_pairs[reg_idx // 2], raw_v_frag0[raw_matrix], cfg.io_dtype) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + diff_w1[reg_idx ^ 2] = raw_v_frag1[raw_matrix] + y_inp_pack1[reg_idx ^ 2] = mul_f16x2(beta_pairs[reg_idx // 2], raw_v_frag1[raw_matrix], cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + input_col_id, cutlass.Int8), y_inp_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + input_col_id, cutlass.Int8), y_inp_pack1[0:4]) + nvvm.tcgen05_wait("store") + bars.mb_y_inp_ready.arrive() + + # ---- dU restage: dU acc -> TMEM f16 A operand ------------------------ + bars.mb_du_acc_ready.wait(du_acc_index.phase) + du_acc_index = advance(du_acc_index, 1) + du_col_id = tmem_col + cfg.tmem_du_acc_offset + du_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + du_col_id, cutlass.Float32), num=2) + du_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + du_col_id, cutlass.Float32), num=2) + + du_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + du_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + frag_pair = reg_idx * 2 + du_pack0[reg_idx] = fp32_to_fp16(du_vec0[frag_pair], du_vec0[frag_pair + 1], dtype=cfg.io_dtype) + du_pack1[reg_idx] = fp32_to_fp16(du_vec1[frag_pair], du_vec1[frag_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + (tmem_col + cfg.tmem_du_inp_offset), cutlass.Int8), du_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + (tmem_col + cfg.tmem_du_inp_offset), cutlass.Int8), du_pack1[0:4]) + nvvm.tcgen05_wait("store") + bars.mb_du_inp_ready.arrive() + + # ---- U readback -> sU ------------------------------------------------ + bars.mb_u_acc_ready.wait(u_acc_index.phase) + u_acc_index = advance(u_acc_index, 1) + u_col_id = tmem_col + cfg.tmem_u_acc_offset + u_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + u_col_id, cutlass.Float32), num=2) + u_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + u_col_id, cutlass.Float32), num=2) + + u_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + u_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + u_pack0[reg_idx] = fp32_to_fp16(u_vec0[2 * reg_idx], u_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) + u_pack1[reg_idx] = fp32_to_fp16(u_vec1[2 * reg_idx], u_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) + nvvm.stmatrix( + sU_raw.data_ptr() + + (value_dim_base + value_col_offset) // 64 * (cfg.b_t * 64) + + token_row_coord * 64 + + swizzle_xor_128b(token_row_coord, (value_dim_base + value_col_offset) % 64, elem_bytes=2), + u_pack0.data_ptr().load(count=4, alignment=4), + nvvm.MMALayout.COL, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.stmatrix( + sU_raw.data_ptr() + + (value_dim_base + 16 + value_col_offset) // 64 * (cfg.b_t * 64) + + token_row_coord * 64 + + swizzle_xor_128b(token_row_coord, (value_dim_base + 16 + value_col_offset) % 64, elem_bytes=2), + u_pack1.data_ptr().load(count=4, alignment=4), + nvvm.MMALayout.COL, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_u_smem_ready.arrive() + + # ---- dY readback ------------------------------------------------------- + bars.mb_dy_acc_ready.wait(dy_acc_index.phase) + dy_acc_index = advance(dy_acc_index, 1) + dy_col_id = tmem_col + cfg.tmem_dy_acc_offset + dy_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + dy_col_id, cutlass.Float32), num=2) + dy_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + dy_col_id, cutlass.Float32), num=2) + + # ---- dY -> sdY: pack + store + publish (super warp's dM operand) ------- + addr_lo0 = ( + (value_dim_base + value_col_offset) // 64 * (cfg.b_t * 64) + + token_row_coord * 64 + + swizzle_xor_128b(token_row_coord, (value_dim_base + value_col_offset) % 64, elem_bytes=2) + ) + addr_lo1 = ( + (value_dim_base + 16 + value_col_offset) // 64 * (cfg.b_t * 64) + + token_row_coord * 64 + + swizzle_xor_128b(token_row_coord, (value_dim_base + 16 + value_col_offset) % 64, elem_bytes=2) + ) + dy_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + dy_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + dy_pack0[reg_idx] = fp32_to_fp16(dy_vec0[2 * reg_idx], dy_vec0[2 * reg_idx + 1], dtype=cfg.io_dtype) + dy_pack1[reg_idx] = fp32_to_fp16(dy_vec1[2 * reg_idx], dy_vec1[2 * reg_idx + 1], dtype=cfg.io_dtype) + nvvm.stmatrix(sDy_raw.data_ptr() + addr_lo0, dy_pack0.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) + nvvm.stmatrix(sDy_raw.data_ptr() + addr_lo1, dy_pack1.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dy_smem_ready.arrive() + + # ---- beta scalars -> beta.dY -> sdV (epilogue TMA + dK_decay MMA operand) --- + beta_c0 = (sBeta_ptr + (lane % 4) * 2).load().to(cutlass.Float32) + beta_c1 = (sBeta_ptr + (lane % 4) * 2 + 1).load().to(cutlass.Float32) + beta_c8 = (sBeta_ptr + (lane % 4) * 2 + 8).load().to(cutlass.Float32) + beta_c9 = (sBeta_ptr + (lane % 4) * 2 + 9).load().to(cutlass.Float32) + bars.mb_beta_done[chunk_serial % cfg.smem_beta_stages].arrive() + beta_dy_regs0 = cutlass.Array(cutlass.Float32, 8, alignment=16) + beta_dy_regs1 = cutlass.Array(cutlass.Float32, 8, alignment=16) + for e2 in cutlass.range_constexpr(4): + e = 2 * e2 + b_lo = beta_c8 if cutlass.const_expr(e >= 4) else beta_c0 + b_hi = beta_c9 if cutlass.const_expr(e >= 4) else beta_c1 + beta_dy_regs0[e], beta_dy_regs0[e + 1] = fmul2(dy_vec0[e], dy_vec0[e + 1], b_lo, b_hi) + beta_dy_regs1[e], beta_dy_regs1[e + 1] = fmul2(dy_vec1[e], dy_vec1[e + 1], b_lo, b_hi) + + # ---- -beta.dY -> TMEM: A operand of the dH ds-term --------------------- + neg_beta_dy_regs0 = cutlass.Array(cutlass.Float32, 8, alignment=16) + neg_beta_dy_regs1 = cutlass.Array(cutlass.Float32, 8, alignment=16) + for e in cutlass.range_constexpr(8): + neg_beta_dy_regs0[e] = -beta_dy_regs0[e] + neg_beta_dy_regs1[e] = -beta_dy_regs1[e] + neg_beta_dy_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + neg_beta_dy_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + frag_pair = reg_idx * 2 + neg_beta_dy_pack0[reg_idx] = fp32_to_fp16(neg_beta_dy_regs0[frag_pair], neg_beta_dy_regs0[frag_pair + 1], dtype=cfg.io_dtype) + neg_beta_dy_pack1[reg_idx] = fp32_to_fp16(neg_beta_dy_regs1[frag_pair], neg_beta_dy_regs1[frag_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_lo + (tmem_col + cfg.tmem_neg_beta_dy_inp_offset), cutlass.Int8), neg_beta_dy_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(row_addr_hi + (tmem_col + cfg.tmem_neg_beta_dy_inp_offset), cutlass.Int8), neg_beta_dy_pack1[0:4]) + nvvm.tcgen05_wait("store") + bars.mb_neg_beta_dy_inp_ready.arrive() + + # ---- dBeta v-term parts: dY.(V - state_k), diff_w from Y staging ------- + tok4 = cutlass.Array(cutlass.Float32, 4, alignment=16) + for s in cutlass.range_constexpr(4): + tok4[s] = cutlass.Float32(0.0) + for j in cutlass.range_constexpr(4): + d0_lo, d0_hi = f16x2_to_f32(diff_w0[j], dtype=cfg.io_dtype) + d1_lo, d1_hi = f16x2_to_f32(diff_w1[j], dtype=cfg.io_dtype) + s_lo = cutlass.const_expr(2 * (j // 2)) + s_hi = cutlass.const_expr(2 * (j // 2) + 1) + t_lo, t_hi = ffma2(dy_vec0[2 * j], dy_vec0[2 * j + 1], d0_lo, d0_hi, tok4[s_lo], tok4[s_hi]) + t_lo, t_hi = ffma2(dy_vec1[2 * j], dy_vec1[2 * j + 1], d1_lo, d1_hi, t_lo, t_hi) + tok4[s_lo] = t_lo + tok4[s_hi] = t_hi + + # ---- beta.dY -> sdV (epilogue TMA + dK_decay MMA operand) ------------ + beta_dy_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + beta_dy_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + beta_dy_pack0[reg_idx] = fp32_to_fp16(beta_dy_regs0[2 * reg_idx], beta_dy_regs0[2 * reg_idx + 1], dtype=cfg.io_dtype) + beta_dy_pack1[reg_idx] = fp32_to_fp16(beta_dy_regs1[2 * reg_idx], beta_dy_regs1[2 * reg_idx + 1], dtype=cfg.io_dtype) + dv_stage = chunk_serial % cfg.smem_dv_stages + sdv_stage_base = dv_stage * (cfg.b_t * cfg.d_v) + bars.mb_dv_tmastg_done[dv_stage].wait(dv_done_index.phase) + dv_done_index = advance(dv_done_index, cfg.smem_dv_stages) + nvvm.stmatrix( + sDv_raw.data_ptr() + sdv_stage_base + addr_lo0, + beta_dy_pack0.data_ptr().load(count=4, alignment=4), + nvvm.MMALayout.COL, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.stmatrix( + sDv_raw.data_ptr() + sdv_stage_base + addr_lo1, + beta_dy_pack1.data_ptr().load(count=4, alignment=4), + nvvm.MMALayout.COL, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dv_tmastg_ready[dv_stage].arrive() + + # ---- dBeta = sum_v dY.(V - state_k) + M-term ------------------------- + nvvm.barrier_cta_sync(cfg.cg1_sync_barrier_id, thread_count=cfg.cg1_threads) + bars.mb_dbeta_m_ready.wait(dbeta_m_index.phase) + dbeta_m_index = advance(dbeta_m_index, 1) + for off in cutlass.range_constexpr(3): + step = cutlass.const_expr(4 << off) + for s in cutlass.range_constexpr(4): + tok4[s] = tok4[s] + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, tok4[s], step, 31, kind=nvvm.Shfl.BFLY)) + if lane < 4: + for s in cutlass.range_constexpr(4): + sRed_raw[(warp_idx % 4) * cfg.b_t + (lane % 4) * 2 + (s % 2) + 8 * (s // 2)] = tok4[s] + nvvm.barrier_cta_sync(cfg.cg1_sync_barrier_id, thread_count=cfg.cg1_threads) + if cg1_tidx < cfg.b_t: + acc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + acc = acc + sRed_raw[w * cfg.b_t + cg1_tidx] + db_val = acc + sBetaM_raw[cg1_tidx] + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cg1_tidx + if token_idx < seqlen_b and chunk_idx < wend: + mDbeta[batch_start + token_idx, head_idx] = db_val + nvvm.barrier_cta_sync(cfg.cg1_sync_barrier_id, thread_count=cfg.cg1_threads) + + # ---- dH capture for the next ----------------------------------------- + bars.mb_dstate_acc_ready.wait(dstate_ready_index.phase) + dstate_ready_index = advance(dstate_ready_index, 1) + if rev_idx + cutlass.Int32(1) < num_compute_chunks: + bars.mb_dstate_smem_done.wait(dstate_smem_done_index.phase) + 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 + for sub in cutlass.range_constexpr(cfg.d_k // 32): + dstate_vec = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_acc_offset + sub * 32), cutlass.Float32), num=32 + ) + dstate_pack = cutlass.Array(cutlass.Int32, 16, alignment=16) + for pc in cutlass.range_constexpr(16): + dstate_pack[pc] = fp32_to_fp16(dstate_vec[2 * pc], dstate_vec[2 * pc + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 16), cutlass.Int8), + dstate_pack[0:16], + ) + nvvm.tcgen05_wait("store") + bars.mb_dstate_inp_ready.arrive() + + # ---- dh_inp -> sdH: re-read after the TMEM publish --------------- + for sub in cutlass.range_constexpr(cfg.d_k // 32): + dstate_words = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dstate_inp_offset + sub * 16), cutlass.Float32), num=16 + ) + for half in cutlass.range_constexpr(4): + d_base = sub * 32 + half * 8 + h_vec = cutlass.Vector.from_elements( + (dstate_words[half * 4], dstate_words[half * 4 + 1], dstate_words[half * 4 + 2], dstate_words[half * 4 + 3]), + cutlass.Float32, + ).bitcast(cfg.io_dtype) + h_addr = (d_base // 64) * (cfg.d_v * 64) + value_dim * 64 + swizzle_xor_128b(value_dim, d_base % 64, elem_bytes=2) + (sDstate_raw.data_ptr() + h_addr).store(h_vec, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dstate_smem_ready.arrive() + raw_index = advance(raw_index, cfg.smem_raw_stages) + + # ---- tile end: dS0 drain / zero-length pass-through ---------------------- + if cutlass.const_expr(mDstate0 is not None): + if num_compute_chunks > 0: + if wstart == 0: + row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + 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] + 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] + else: + mDstate0[batch_idx, head_idx, kd, value_dim] = 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) + + bars.mb_tmem_done[0].arrive() + + +@cute.jit +def compute2_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_base_holder, + warp_idx, + sGate_raw, + sNorm_raw, + sDq_raw, + sDk_raw, + sRed1_raw, + sDstate_raw, + sDgate_raw, + scale, + bars, +) -> None: + """WG2 warp role (warps 8-11): the gradient drain. Each thread owns one + key channel d for all 16 tokens: reads the dQ accumulator and the four dK + parts from TMEM, assembles dQ/dK with the per-channel gate factors (raw + Q/K arrive through WG0's TMEM ring), applies the in-kernel L2-norm + backward row projection, assembles the per-channel dGate including the + g_last terms, reverse-cumsums it in registers, stages dGate for the + epilogue's TMA store, and stages dQ/dK for the epilogue's TMA stores.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_2, nvvm.SetMaxRegisterAction.INCREASE) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_holder.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + tmem_subpartition = warp_idx % 4 + channel = tmem_subpartition * cfg.threads_per_warp + lane + row_addr = (tmem_row + tmem_subpartition * cfg.threads_per_warp) << 16 + cg2_tidx = channel + + raw_index = PipelineState.start(phase=0) + dq_acc_index = PipelineState.start(phase=0) + dk_decay_part_index = PipelineState.start(phase=0) + dk_inv_part_index = PipelineState.start(phase=0) + dk_restore_part_index = PipelineState.start(phase=0) + dgate_last_dstate_smem_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 + for rev_idx in cutlass.range(num_compute_chunks, unroll=1): + chunk_idx = cend - cutlass.Int32(1) - rev_idx + chunk_serial = chunk_serial_base + rev_idx + chunk_start = chunk_idx * cfg.b_t + raw_stage = chunk_serial % cfg.smem_raw_stages + decay_stage = chunk_serial % cfg.smem_decay_stages + has_dstate = cutlass.Boolean(rev_idx > 0) + if cutlass.const_expr(cfg.use_dstate_in): + has_dstate = cutlass.Boolean(True) + sGate_ptr = sGate_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + writes = chunk_idx < wend + + # ---- gate landed: CG0 publishes the decay ring only after consuming + # the gate TMA, so this wait is CG2's visibility guard for sGate ------ + bars.mb_k_decay_inv_ready[decay_stage].wait((chunk_serial // cfg.smem_decay_stages) % 2) + + # ---- per-channel gate factors ---------------------------------------- + f32_seg = channel // 32 + f32_dim = channel - f32_seg * 32 + f16_seg = channel // 64 + f16_dim = channel - f16_seg * 64 + eg = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + for t in cutlass.range_constexpr(cfg.b_t): + eg[t] = (sGate_ptr + f32_seg * (cfg.b_t * 32) + t * 32 + swizzle_xor_128b(t, f32_dim, elem_bytes=4)).load() + egl = eg[cfg.b_t - 1] + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_gate_done[raw_stage].arrive() + + # ---- staged raw Q/K: TMEM ring cols for this chunk ------------------- + qk_raw_stage = chunk_serial % cfg.tmem_qk_raw_stages + qraw_col = tmem_col + cfg.tmem_qraw_inp_offset + qk_raw_stage * (cfg.b_t // 2) + kraw_col = tmem_col + cfg.tmem_kraw_inp_offset + qk_raw_stage * (cfg.b_t // 2) + norm_base = qk_raw_stage * (2 * cfg.b_t) + bars.mb_qk_raw_ready[qk_raw_stage].wait((chunk_serial // cfg.tmem_qk_raw_stages) % 2) + + # ---- dGate_last hdot: sum_v sdH[v, c] * S0[c, v] --------------------- + dgate_last_val = cutlass.Float32(0.0) + bars.mb_state_inp_ready[chunk_serial % 2].wait((chunk_serial // 2) % 2) + if has_dstate: + bars.mb_dstate_smem_ready.wait(dgate_last_dstate_smem_index.phase) + dgate_last_dstate_smem_index = advance(dgate_last_dstate_smem_index, 1) + for pl in cutlass.range_constexpr(2): + for row_half in cutlass.range_constexpr(2): + state_vec = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr( + row_addr + (tmem_col + cfg.tmem_state_inp_offset + (chunk_serial % 2) * (cfg.d_v // 2) + pl * 32 + row_half * 16), + cutlass.Float32, + ), + num=16, + ) + hacc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for i in cutlass.range_constexpr(8): + hacc[i] = opaque_f32_zero() + for j in cutlass.range_constexpr(16): + v0 = pl * 64 + row_half * 32 + 2 * j + state_pair = cutlass.Vector.from_elements((state_vec[j],), cutlass.Float32).bitcast(cfg.io_dtype) + dstate_addr0 = (channel // 64) * (cfg.d_v * 64) + v0 * 64 + swizzle_xor_128b(v0, channel % 64, elem_bytes=2) + dstate_addr1 = (channel // 64) * (cfg.d_v * 64) + (v0 + 1) * 64 + swizzle_xor_128b(v0 + 1, channel % 64, elem_bytes=2) + hval0 = (sDstate_raw.data_ptr() + dstate_addr0).load().to(cutlass.Float32) + hval1 = (sDstate_raw.data_ptr() + dstate_addr1).load().to(cutlass.Float32) + hacc[(2 * j) % 8] = hacc[(2 * j) % 8] + hval0 * state_pair[0].to(cutlass.Float32) + hacc[(2 * j + 1) % 8] = hacc[(2 * j + 1) % 8] + hval1 * state_pair[1].to(cutlass.Float32) + part_a = (hacc[0] + hacc[4]) + (hacc[1] + hacc[5]) + part_b = (hacc[2] + hacc[6]) + (hacc[3] + hacc[7]) + dgate_last_val = dgate_last_val + (part_a + part_b) + bars.mb_dstate_smem_cg2_done.arrive() + bars.mb_state_inp_cg2_done[chunk_serial % 2].arrive() + + # ---- part-drain accumulators ------------------------------------------- + dq_n = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + dk_n = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + dgate_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + dgate_last_acc = cutlass.Array(cutlass.Float32, 4, alignment=16) + for i in cutlass.range_constexpr(4): + dgate_last_acc[i] = opaque_f32_zero() + for t in cutlass.range_constexpr(cfg.b_t): + dk_n[t] = cutlass.Float32(0.0) + + # ---- dK_restore part drain: (eGl/eG) scale + dGate_last K-dot ---------- + if has_dstate: + bars.mb_dk_restore_part_acc_ready.wait(dk_restore_part_index.phase) + dk_restore_part_index = advance(dk_restore_part_index, 1) + dk_restore_part_vec = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_restore_acc_offset), cutlass.Float32), num=cfg.b_t + ) + kr_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) + for t in cutlass.range_constexpr(cfg.b_t): + dk_hat = egl * cute.math.rcp(eg[t], approx=True, ftz=True) * dk_restore_part_vec[t] + dk_n[t] = dk_hat + k_pair = cutlass.Vector.from_elements((kr_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + k_v = k_pair[t % 2].to(cutlass.Float32) + if cutlass.const_expr(cfg.l2norm): + k_v = k_v * sNorm_raw[norm_base + cfg.b_t + t] + dgate_last_acc[t % 4] = dgate_last_acc[t % 4] + k_v * dk_hat + + # ---- dQ acc drain: eG.scale --------------------------------------------- + bars.mb_dq_acc_ready.wait(dq_acc_index.phase) + dq_acc_index = advance(dq_acc_index, 1) + dq_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dq_acc_offset), cutlass.Float32), num=cfg.b_t) + for t2 in cutlass.range_constexpr(cfg.b_t // 2): + t = 2 * t2 + es_lo, es_hi = fmul2(eg[t], eg[t + 1], scale, scale) + dq_n[t], dq_n[t + 1] = fmul2(es_lo, es_hi, dq_vec[t], dq_vec[t + 1]) + + # ---- dK_inv part drain: (dA - dM) term, 1/eG scale ---------------------- + bars.mb_dk_inv_part_acc_ready.wait(dk_inv_part_index.phase) + dk_inv_part_index = advance(dk_inv_part_index, 1) + dk_inv_part_vec = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_inv_acc_offset), cutlass.Float32), num=cfg.b_t) + for t in cutlass.range_constexpr(cfg.b_t): + dk_n[t] = dk_n[t] + dk_inv_part_vec[t] * cute.math.rcp(eg[t], approx=True, ftz=True) + + # ---- dK_decay part drain: -eG scale, seeds dGate ------------------------ + bars.mb_dk_decay_part_acc_ready.wait(dk_decay_part_index.phase) + dk_decay_part_index = advance(dk_decay_part_index, 1) + dk_decay_part_vec = nvvm.tcgen05_ld( + "32x32b", nvvm.make_tmem_ptr(row_addr + (tmem_col + cfg.tmem_dk_decay_acc_offset), cutlass.Float32), num=cfg.b_t + ) + for t in cutlass.range_constexpr(cfg.b_t): + dgate_regs[t] = -eg[t] * dk_decay_part_vec[t] + dk_n[t] = dk_n[t] + dgate_regs[t] + + nvvm.tcgen05_wait("load") + bars.mb_dqk_acc_done.arrive() + + # ---- dGate finalize -------------------------------------------------- + qf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qraw_col, cutlass.Float32), num=cfg.b_t // 2) + kf_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + kraw_col, cutlass.Float32), num=cfg.b_t // 2) + for t in cutlass.range_constexpr(cfg.b_t): + q_pair = cutlass.Vector.from_elements((qf_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + k_pair = cutlass.Vector.from_elements((kf_words[t // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + q_v = q_pair[t % 2].to(cutlass.Float32) + k_v = k_pair[t % 2].to(cutlass.Float32) + if cutlass.const_expr(cfg.l2norm): + q_v = q_v * sNorm_raw[norm_base + t] + k_v = k_v * sNorm_raw[norm_base + cfg.b_t + t] + dgate_regs[t] = q_v * dq_n[t] + k_v * (cutlass.Float32(2.0) * dgate_regs[t] - dk_n[t]) + dgate_regs[cfg.b_t - 1] = dgate_regs[cfg.b_t - 1] + ((dgate_last_acc[0] + dgate_last_acc[1]) + (dgate_last_acc[2] + dgate_last_acc[3])) + + # ---- L2-norm backward row projection --------------------------------- + if cutlass.const_expr(cfg.l2norm): + for grad, qk_col, inv_off in ((dq_n, qraw_col, 0), (dk_n, kraw_col, cfg.b_t)): + dots = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + for half in cutlass.range_constexpr(2): + p_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) + for tt in cutlass.range_constexpr(cfg.b_t // 2): + t = half * (cfg.b_t // 2) + tt + p_pair = cutlass.Vector.from_elements((p_words[tt // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + dots[t] = grad[t] * p_pair[tt % 2].to(cutlass.Float32) * sNorm_raw[norm_base + inv_off + t] + for off in cutlass.range_constexpr(5): + step = cutlass.const_expr(1 << off) + for t in cutlass.range_constexpr(cfg.b_t): + dots[t] = dots[t] + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, dots[t], step, 31, kind=nvvm.Shfl.BFLY)) + if lane == 0: + for t in cutlass.range_constexpr(cfg.b_t): + sRed1_raw[tmem_subpartition * cfg.b_t + t] = dots[t] + nvvm.barrier_cta_sync(cfg.cg2_sync_barrier_id, thread_count=cfg.cg2_threads) + for half in cutlass.range_constexpr(2): + a_words = nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + qk_col + half * (cfg.b_t // 4), cutlass.Float32), num=cfg.b_t // 4) + for tt in cutlass.range_constexpr(cfg.b_t // 2): + t = half * (cfg.b_t // 2) + tt + a_pair = cutlass.Vector.from_elements((a_words[tt // 2],), cutlass.Float32).bitcast(cfg.io_dtype) + total_dot = sRed1_raw[t] + sRed1_raw[cfg.b_t + t] + sRed1_raw[2 * cfg.b_t + t] + sRed1_raw[3 * cfg.b_t + t] + norm_t = sNorm_raw[norm_base + inv_off + t] + grad[t] = (grad[t] - a_pair[tt % 2].to(cutlass.Float32) * norm_t * total_dot) * norm_t + nvvm.barrier_cta_sync(cfg.cg2_sync_barrier_id, thread_count=cfg.cg2_threads) + + bars.mb_qk_raw_done[qk_raw_stage].arrive() + + # ---- stage dQ/dK for the epilogue TMA stores ------------------------- + dq_stage = chunk_serial % cfg.smem_dq_stages + dk_stage = chunk_serial % cfg.smem_dk_stages + bars.mb_dq_tmastg_done[dq_stage].wait(((chunk_serial // cfg.smem_dq_stages) + 1) % 2) + bars.mb_dk_tmastg_done[dk_stage].wait(((chunk_serial // cfg.smem_dk_stages) + 1) % 2) + dq_base = dq_stage * (cfg.b_t * cfg.d_k) + dk_base = dk_stage * (cfg.b_t * cfg.d_k) + for t in cutlass.range_constexpr(cfg.b_t): + out_idx = f16_seg * (cfg.b_t * 64) + t * 64 + swizzle_xor_128b(t, f16_dim, elem_bytes=2) + (sDq_raw.data_ptr() + dq_base + out_idx).store(dq_n[t].to(cfg.io_dtype)) + (sDk_raw.data_ptr() + dk_base + out_idx).store(dk_n[t].to(cfg.io_dtype)) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dq_tmastg_ready[dq_stage].arrive() + bars.mb_dk_tmastg_ready[dk_stage].arrive() + + # ---- dGate_last add -------------------------------------------------- + if has_dstate: + if chunk_idx >= FIRST_STATE_CHUNK: + dgate_regs[cfg.b_t - 1] = dgate_regs[cfg.b_t - 1] + egl * dgate_last_val + + # ---- dGate reverse cumsum -------------------------------------------- + suffix = cutlass.Float32(0.0) + for rt in cutlass.range_constexpr(cfg.b_t): + t = cfg.b_t - 1 - rt + suffix = suffix + dgate_regs[t] + dgate_regs[t] = suffix + + # ---- stage dGate for the epilogue TMA store -------------------------- + dgate_stage = chunk_serial % cfg.smem_dgate_stages + bars.mb_dgate_tmastg_done[dgate_stage].wait(((chunk_serial // cfg.smem_dgate_stages) + 1) % 2) + for t in cutlass.range_constexpr(cfg.b_t): + dgate_idx = f32_seg * (cfg.b_t * 32) + t * 32 + swizzle_xor_128b(t, f32_dim, elem_bytes=4) + (sDgate_raw.data_ptr() + dgate_idx).store(dgate_regs[t]) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_dgate_tmastg_ready[dgate_stage].arrive() + raw_index = advance(raw_index, cfg.smem_raw_stages) + chunk_serial_base += num_compute_chunks + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + bars.mb_tmem_done[0].arrive() + + +# --------------------------------------------------------------------------- +# Host-side assembly +# --------------------------------------------------------------------------- + + +@cute.kernel +def build_all_descs_kernel( + base_q: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_v: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_gate: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_do: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_dq: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_dk: cutlass.GridConstant[cuda.tensor_map.TensorMap], + 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, + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + do: cute.Tensor, + dq: cute.Tensor, + dk: cute.Tensor, + 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, + v_row_stride: cutlass.Int32, + gate_row_stride: cutlass.Int32, + do_row_stride: cutlass.Int32, + dq_row_stride: cutlass.Int32, + dk_row_stride: cutlass.Int32, + dv_row_stride: cutlass.Int32, + dgate_row_stride: cutlass.Int32, + checkpoint_row_stride: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, +) -> None: + """Single-launch builder for the per-batch descriptor arrays (one warp + per array).""" + tidx, _, _ = cute.arch.thread_idx() + widx = cutlass.Int32(tidx) // cutlass.Int32(32) + arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) + sub0 = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) + sub1 = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub2 = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub3 = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + sub4 = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + 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((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(): + emit_seq_descs(base_q, sub0, cu_seqlens, q, n_batch, q_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 1: + if nvvm.elect_sync(): + emit_seq_descs(base_k, sub1, cu_seqlens, k, n_batch, k_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 2: + if nvvm.elect_sync(): + emit_seq_descs(base_v, sub2, cu_seqlens, v, n_batch, v_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 3: + if nvvm.elect_sync(): + emit_seq_descs(base_gate, sub3, cu_seqlens, gate, n_batch, gate_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 4: + if nvvm.elect_sync(): + emit_seq_descs(base_do, sub4, cu_seqlens, do, n_batch, do_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 5: + if nvvm.elect_sync(): + emit_seq_descs(base_dq, sub5, cu_seqlens, dq, n_batch, dq_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 6: + if nvvm.elect_sync(): + emit_seq_descs(base_dk, sub6, cu_seqlens, dk, n_batch, dk_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 7: + if nvvm.elect_sync(): + emit_seq_descs(base_dv, sub7, cu_seqlens, dv, n_batch, dv_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 8: + if nvvm.elect_sync(): + emit_seq_descs(base_dgate, sub8, cu_seqlens, dgate, n_batch, dgate_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 9: + 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.jit +def build_descs( + io_dtype: cutlass.Constexpr, + b_t: cutlass.Constexpr[int], + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + do: cute.Tensor, + dq: cute.Tensor, + dk: cute.Tensor, + dv: cute.Tensor, + dgate: cute.Tensor, + state_checkpoints: cute.Tensor, + state0: cute.Tensor | None, + cu_seqlens: cute.Tensor, + tensormap_workspace: cute.Tensor, + stream: cuda_driver.CUstream, +): + """Build the 11 per-(batch, head) capped TMA-descriptor arrays into + ``tensormap_workspace`` (sequence-relative coordinates; tail loads + zero-fill and tail stores clip in hardware).""" + h_q = q.shape[1] + h_k = k.shape[1] + h_v = v.shape[1] + ho = gate.shape[1] + batch_size = cu_seqlens.shape[0] - 1 + d_k = q.shape[2] + d_v = v.shape[2] + bpe = io_dtype.width // 8 + tma_granu_elems = 128 // bpe + seqlen = q.shape[0] + + q_headed = cute.make_tensor(q.iterator, cute.make_layout((d_k, h_q, seqlen), stride=(1, q.stride[1], q.stride[0]))) + k_headed = cute.make_tensor(k.iterator, cute.make_layout((d_k, h_k, seqlen), stride=(1, k.stride[1], k.stride[0]))) + v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v.stride[1], v.stride[0]))) + gate_headed = cute.make_tensor(gate.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, gate.stride[1], gate.stride[0]))) + do_headed = cute.make_tensor(do.iterator, cute.make_layout((d_v, ho, seqlen), stride=(1, do.stride[1], do.stride[0]))) + dq_headed = cute.make_tensor(dq.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, dq.stride[1], dq.stride[0]))) + dk_headed = cute.make_tensor(dk.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, dk.stride[1], dk.stride[0]))) + dv_headed = cute.make_tensor(dv.iterator, cute.make_layout((d_v, ho, seqlen), stride=(1, dv.stride[1], dv.stride[0]))) + dgate_headed = cute.make_tensor(dgate.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, dgate.stride[1], dgate.stride[0]))) + + swz = cuda.TensorMapSwizzle.s128b + base_q = cuda.create_tensor_map_tiled_from_view(q_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_k = cuda.create_tensor_map_tiled_from_view(k_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_v = cuda.create_tensor_map_tiled_from_view(v_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_gate = cuda.create_tensor_map_tiled_from_view(gate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_do = cuda.create_tensor_map_tiled_from_view(do_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dq = cuda.create_tensor_map_tiled_from_view(dq_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dk = cuda.create_tensor_map_tiled_from_view(dk_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dv = cuda.create_tensor_map_tiled_from_view(dv_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_dgate = cuda.create_tensor_map_tiled_from_view(dgate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + + checkpoint_view = cute.make_tensor( + state_checkpoints.iterator, + cute.make_layout( + (d_v, d_k, state_checkpoints.shape[0], ho), + stride=(state_checkpoints.stride[3], state_checkpoints.stride[2], state_checkpoints.stride[0], state_checkpoints.stride[1]), + ), + ) + 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) + + n_warps = 11 if state0 is not None else 10 + build_all_descs_kernel( + base_q, + base_k, + base_v, + base_gate, + base_do, + base_dq, + base_dk, + base_dv, + base_dgate, + base_checkpoint, + base_initial_state, + tensormap_workspace, + cu_seqlens, + q, + k, + v, + gate, + do, + dq, + dk, + dv, + dgate, + state_checkpoints, + state0, + cutlass.Int32(batch_size), + cutlass.Int32(q.stride[0]), + cutlass.Int32(k.stride[0]), + cutlass.Int32(v.stride[0]), + cutlass.Int32(gate.stride[0]), + cutlass.Int32(do.stride[0]), + cutlass.Int32(dq.stride[0]), + cutlass.Int32(dk.stride[0]), + cutlass.Int32(dv.stride[0]), + cutlass.Int32(dgate.stride[0]), + cutlass.Int32(state_checkpoints.stride[0]), + cutlass.Int32(b_t), + ).launch(grid=(1, 1, 1), block=(32 * n_warps, 1, 1), stream=stream) + + +@cute.jit +def host( + cfg: cutlass.Constexpr, + beta: cute.Tensor, + state_checkpoints: cute.Tensor, + mState_init: cute.Tensor | None, + dgate: cute.Tensor, + dbeta: cute.Tensor, + cu_seqlens: cute.Tensor, + d_initial_state: cute.Tensor | None, + d_final_state: cute.Tensor | None, + work_items: cute.Tensor | None, + work_count: cute.Tensor | None, + sched_ctr: cute.Tensor | None, + tensormap_workspace: cute.Tensor, + scale: cutlass.Float32, + stream, +) -> None: + num_sequences = cu_seqlens.shape[0] - 1 + + # ---- launch ------------------------------------------------------------------ + n_desc = num_sequences + grid_shape = (cfg.max_active_clusters, 1, 1) + kernel( + cfg, + tensormap_workspace, + n_desc, + beta, + cu_seqlens, + dgate, + dbeta, + d_initial_state, + d_final_state, + work_items, + work_count, + sched_ctr, + scale, + ).launch( + grid=grid_shape, + block=(cfg.threads_per_cta, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + +@cute.kernel +def kernel( + cfg: cutlass.Constexpr, + tensormap_workspace: cute.Tensor, + n_desc: cutlass.Int32, + mBeta: cute.Tensor, + cu_seqlens: cute.Tensor, + mDgate: cute.Tensor, + mDbeta: cute.Tensor, + mDstate0: cute.Tensor | None, + mDstate_in: cute.Tensor | None, + mWorkItems: cute.Tensor, + mCount: cute.Tensor, + mSched: cute.Tensor | None, + scale: cutlass.Float32, +) -> None: + """BT=16 KDA backward kernel (persistent, 16 warps).""" + tidx, _, _ = cute.arch.thread_idx() + bidx = cute.arch.block_idx()[0] + num_ctas = cute.arch.grid_dim()[0] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % cfg.threads_per_warp + + total_tiles = mCount[0] + assert mBeta.element_type == cutlass.Float32 + assert cu_seqlens.element_type in (cutlass.Int32, cutlass.Int64) + assert mDgate.element_type == cutlass.Float32 and mDbeta.element_type == cutlass.Float32 + + desc_base_words = tensormap_workspace.iterator.raw_ptr() + arr_words = n_desc * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_q_base = desc_base_words + desc_k_base = desc_base_words + arr_words + desc_v_base = desc_base_words + cutlass.Int32(2) * arr_words + desc_gate_base = desc_base_words + cutlass.Int32(3) * arr_words + desc_do_base = desc_base_words + cutlass.Int32(4) * arr_words + 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_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) + tmem_base_holder = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) + sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + bpe = cfg.io_dtype.width // 8 + SWZ = 2 + LEAD = 16 + STRIDE = 8 * 128 + STATE_ALT_LEAD = cfg.d_v * 128 + + # sub-bank split: tcgen05-descriptor operands low, generic-client buffers high + sBeta_raw = cutlass.Array(cutlass.Float32, cfg.smem_beta_stages * cfg.b_t, space=SMEM, alignment=64) + sNorm_raw = cutlass.Array(cutlass.Float32, cfg.tmem_qk_raw_stages * 2 * cfg.b_t, space=SMEM, alignment=64) + sRed_raw = cutlass.Array(cutlass.Float32, 4 * cfg.b_t, space=SMEM, alignment=64) + sRed1_raw = cutlass.Array(cutlass.Float32, 4 * cfg.b_t, space=SMEM, alignment=64) + sBetaM_raw = cutlass.Array(cutlass.Float32, cfg.b_t, space=SMEM, alignment=64) + sState_raw = cutlass.Array(cfg.io_dtype, cfg.state_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDstate_raw = cutlass.Array(cfg.io_dtype, cfg.d_k * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_decay_raw = cutlass.Array(cfg.io_dtype, cfg.operand_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_inv_raw = cutlass.Array(cfg.io_dtype, cfg.operand_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_restore_raw = cutlass.Array(cfg.io_dtype, cfg.operand_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sQ_decay_raw = cutlass.Array(cfg.io_dtype, cfg.operand_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sState_scale_diag_raw = cutlass.Array(cfg.io_dtype, cfg.diag_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sIntermediate_raw = cutlass.Array(cfg.io_dtype, cfg.intermediate_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDo_raw = cutlass.Array(cfg.io_dtype, cfg.raw_v_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDv_raw = cutlass.Array(cfg.io_dtype, cfg.dv_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sQ_raw = cutlass.Array(cfg.io_dtype, cfg.raw_qk_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_raw = cutlass.Array(cfg.io_dtype, cfg.raw_qk_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sV_raw = cutlass.Array(cfg.io_dtype, cfg.raw_v_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sGate_raw = cutlass.Array(cutlass.Float32, cfg.raw_gate_cosize, space=SMEM, alignment=1024) + sDy_raw = cutlass.Array(cfg.io_dtype, cfg.b_t * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) + sU_raw = cutlass.Array(cfg.io_dtype, cfg.b_t * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) + sDq_raw = cutlass.Array(cfg.io_dtype, cfg.dq_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDk_raw = cutlass.Array(cfg.io_dtype, cfg.dk_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sDgate_raw = cutlass.Array(cutlass.Float32, cfg.dgate_cosize, space=SMEM, alignment=1024) + + sState_alt = 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=STATE_ALT_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, + stages=cfg.smem_decay_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sK_inv_lead16 = SmemTile( + base=sK_inv_raw.data_ptr().toint(), + elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, + stages=cfg.smem_decay_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sK_restore_lead16 = SmemTile( + base=sK_restore_raw.data_ptr().toint(), + elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, + stages=cfg.smem_decay_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sDo_lead16 = SmemTile( + base=sDo_raw.data_ptr().toint(), + elems_per_stage=((cfg.raw_v_cosize) // (cfg.smem_raw_stages)) * bpe, + stages=cfg.smem_raw_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sDo_amaj = SmemTile( + base=sDo_raw.data_ptr().toint(), + elems_per_stage=((cfg.raw_v_cosize) // (cfg.smem_raw_stages)) * bpe, + stages=cfg.smem_raw_stages, + leading_byte_offset=cfg.b_t * 128, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sU_lead16 = SmemTile( + base=sU_raw.data_ptr().toint(), + elems_per_stage=((cfg.b_t * cfg.d_v) // (1)) * bpe, + stages=1, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sDv_lead16 = SmemTile( + base=sDv_raw.data_ptr().toint(), + elems_per_stage=((cfg.dv_cosize) // (cfg.smem_dv_stages)) * bpe, + stages=cfg.smem_dv_stages, + leading_byte_offset=LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sDstate_alt = SmemTile( + base=sDstate_raw.data_ptr().toint(), + elems_per_stage=((cfg.d_k * cfg.d_v) // (1)) * bpe, + stages=1, + leading_byte_offset=STATE_ALT_LEAD, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + + sQ_decay_trans = SmemTile( + base=sQ_decay_raw.data_ptr().toint(), + elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, + stages=cfg.smem_decay_stages, + leading_byte_offset=cfg.b_t * 128, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sK_inv_amaj = SmemTile( + base=sK_inv_raw.data_ptr().toint(), + elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, + stages=cfg.smem_decay_stages, + leading_byte_offset=cfg.b_t * 128, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sK_decay_trans = SmemTile( + base=sK_decay_raw.data_ptr().toint(), + elems_per_stage=((cfg.operand_cosize) // (cfg.smem_decay_stages)) * bpe, + stages=cfg.smem_decay_stages, + leading_byte_offset=cfg.b_t * 128, + stride_byte_offset=STRIDE, + layout=SWZ, + ) + sState_scale_diag = SmemTile( + base=sState_scale_diag_raw, + elems_per_stage=((cfg.d_k // 16) * 256), + stages=cfg.smem_decay_stages, + leading_byte_offset=16, + stride_byte_offset=(8 * 16 * 2), + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, + ) + sIntermediate = SmemTile( + base=sIntermediate_raw, + elems_per_stage=(cfg.intermediate_tiles * cfg.b_t * cfg.b_t), + stages=cfg.smem_intermediate_stages, + leading_byte_offset=16, + stride_byte_offset=(8 * cfg.b_t * 2), + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, + ) + + elect_one = nvvm.elect_sync() + if warp_idx == cfg.tma_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_raw_stages): + bars.mb_q_ready[stage].init() + bars.mb_q_done[stage].init() + bars.mb_k_ready[stage].init() + bars.mb_k_done[stage].init() + bars.mb_gate_ready[stage].init() + bars.mb_gate_done[stage].init() + bars.mb_do_ready[stage].init() + bars.mb_do_done[stage].init() + bars.mb_v_ready[stage].init() + bars.mb_v_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_beta_stages): + bars.mb_beta_ready[stage].init() + bars.mb_beta_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_state_stages): + bars.mb_state_ready[stage].init() + bars.mb_state_done[stage].init() + bars.mb_state_cg0_done[stage].init() + for stage in cutlass.range_constexpr(2): + bars.mb_state_inp_ready[stage].init() + bars.mb_state_inp_done[stage].init() + bars.mb_state_inp_cg2_done[stage].init() + elif warp_idx == cfg.tcgen05_mma_warp_id: + if elect_one: + bars.mb_state_k_acc_ready.init() + bars.mb_y_inp_ready.init() + bars.mb_u_acc_ready.init() + bars.mb_u_smem_ready.init() + bars.mb_du_acc_ready.init() + bars.mb_du_inp_ready.init() + bars.mb_dy_acc_ready.init() + bars.mb_neg_beta_dy_inp_ready.init() + bars.mb_dy_smem_ready.init() + bars.mb_dstate_acc_ready.init() + bars.mb_dstate_inp_ready.init() + bars.mb_dstate_smem_ready.init() + bars.mb_dstate_smem_done.init() + bars.mb_dstate_smem_cg2_done.init() + bars.mb_dq_acc_ready.init() + bars.mb_dk_decay_part_acc_ready.init() + bars.mb_dk_inv_part_acc_ready.init() + bars.mb_dk_restore_part_acc_ready.init() + bars.mb_dqk_acc_done.init() + bars.mb_dbeta_m_ready.init() + bars.mb_dstate0_acc_stored.init() + bars.mb_tmem_done[0].init() + elif warp_idx == cfg.super_mma_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_decay_stages): + bars.mb_k_decay_inv_ready[stage].init() + bars.mb_q_decay_k_restore_ready[stage].init() + bars.mb_decay_done[stage].init() + for stage in cutlass.range_constexpr(cfg.tmem_qk_raw_stages): + bars.mb_qk_raw_ready[stage].init() + bars.mb_qk_raw_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_intermediate_stages): + bars.mb_t_inv_ready[stage].init() + bars.mb_a_ready[stage].init() + bars.mb_da_ready[stage].init() + bars.mb_dm_ready[stage].init() + bars.mb_a_done[stage].init() + bars.mb_t_inv_done[stage].init() + bars.mb_da_done[stage].init() + bars.mb_dm_done[stage].init() + elif warp_idx == cfg.epilogue_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_dq_stages): + bars.mb_dq_tmastg_ready[stage].init() + bars.mb_dq_tmastg_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_dk_stages): + bars.mb_dk_tmastg_ready[stage].init() + bars.mb_dk_tmastg_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_dgate_stages): + bars.mb_dgate_tmastg_ready[stage].init() + bars.mb_dgate_tmastg_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_dv_stages): + bars.mb_dv_tmastg_ready[stage].init() + bars.mb_dv_tmastg_done[stage].init() + for stage in cutlass.range_constexpr(cfg.sched_stages): + bars.mb_sched_ready[stage].init() + bars.mb_sched_done[stage].init() + diag_zero = cfg.io_dtype(0.0) + for diag_idx in cutlass.range(tidx, cfg.diag_cosize, cfg.threads_per_cta, unroll=1): + sState_scale_diag_raw[diag_idx] = diag_zero + nvvm.fence_mbarrier_init() + nvvm.barrier_cta_sync(0, thread_count=cfg.threads_per_cta) + if warp_idx == cfg.tma_warp_id: + tmaldg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + mSched, + sSched, + lane, + sQ_raw, + sK_raw, + sV_raw, + sGate_raw, + sDo_raw, + sState_raw, + desc_q_base, + desc_k_base, + desc_v_base, + desc_gate_base, + desc_do_base, + desc_checkpoint_base, + desc_initial_state_base, + bars, + ) + elif warp_idx == cfg.super_mma_warp_id: + super_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_decay_raw, + sK_inv_raw, + sU_raw, + sDy_raw, + sIntermediate_raw, + sBeta_raw, + sBetaM_raw, + bars, + ) + elif warp_idx == cfg.tcgen05_mma_warp_id: + tcgen05_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + tmem_base_holder, + sState_alt, + sK_decay_lead16, + sK_inv_lead16, + sK_inv_amaj, + sK_restore_lead16, + sDo_lead16, + sDo_amaj, + sQ_decay_trans, + sK_decay_trans, + sU_lead16, + sDv_lead16, + sDstate_alt, + sIntermediate, + sState_scale_diag, + bars, + ) + elif warp_idx == cfg.epilogue_warp_id: + epilogue_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_inv_raw, + sQ_decay_raw, + sDo_raw, + sU_raw, + sIntermediate_raw, + sDq_raw, + sDk_raw, + sDv_raw, + sDgate_raw, + desc_dq_base, + desc_dk_base, + desc_dv_base, + desc_dgate_base, + bars, + ) + elif warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: + compute0_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_base_holder, + warp_idx, + scale, + mBeta, + sBeta_raw, + sK_inv_raw, + sGate_raw, + sK_raw, + sQ_raw, + sState_raw, + sNorm_raw, + sK_decay_raw, + sK_restore_raw, + sQ_decay_raw, + sState_scale_diag_raw, + bars, + ) + elif warp_idx >= cfg.compute_group_2_warp_ids[0] and warp_idx <= cfg.compute_group_2_warp_ids[-1]: + compute2_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_base_holder, + warp_idx, + sGate_raw, + sNorm_raw, + sDq_raw, + sDk_raw, + sRed1_raw, + sDstate_raw, + sDgate_raw, + scale, + bars, + ) + elif warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]: + compute1_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + tmem_base_holder, + warp_idx, + mDbeta, + mDstate0, + mDstate_in, + sV_raw, + sDo_raw, + sBeta_raw, + sU_raw, + sDy_raw, + sDv_raw, + sDstate_raw, + sRed_raw, + sBetaM_raw, + bars, + ) + + +@dataclass +class KdaBwdCfg: + """Kernel cfg (fixed BT=16 schedule constants; derived TMEM column offsets + and SMEM buffer cosizes are stamped by ``build_cfg``).""" + + io_dtype: Type[cutlass.Numeric] + use_dstate_in: bool + use_dstate0: bool + l2norm: bool + use_initial_state: bool + q_ratio: int + k_ratio: int + v_ratio: int + n_heads_out: int + max_active_clusters: int + dyn_sched: bool = False + sched_stages: int = 8 + + # ---- fixed constants stamped from CFG by build_cfg --------------------------- + compute_group_0_warp_ids: tuple = CFG.COMPUTE_GROUP_0_WARP_IDS + compute_group_2_warp_ids: tuple = CFG.COMPUTE_GROUP_2_WARP_IDS + compute_group_1_warp_ids: tuple = CFG.COMPUTE_GROUP_1_WARP_IDS + super_mma_warp_id: int = CFG.SUPER_MMA_WARP_ID + tcgen05_mma_warp_id: int = CFG.TCGEN05_MMA_WARP_ID + tma_warp_id: int = CFG.TMA_WARP_ID + epilogue_warp_id: int = CFG.EPILOGUE_WARP_ID + b_t: int = CFG.B_T + d_k: int = CFG.D_K + d_v: int = CFG.D_V + threads_per_warp: int = CFG.THREADS_PER_WARP + threads_per_cta: int = 0 + num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 + num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 + num_regs_compute_group_2: int = CFG.NUM_REGS_COMPUTE_GROUP_2 + num_regs_other: int = CFG.NUM_REGS_OTHER + + # ---- named barrier slots (ids 1-4; 0 is the CTA-wide sync) ------------------- + cg0_sync_barrier_id: int = 1 + cg0_threads: int = 0 + cg2_sync_barrier_id: int = 2 + cg2_threads: int = 0 + tmem_lifecycle_barrier_id: int = 3 + tmem_user_threads: int = 0 + cg1_sync_barrier_id: int = 4 + cg1_threads: int = 0 + + # ---- SMEM / TMEM stage counts + TMEM column offsets -------------------------- + smem_raw_stages: int = CFG.SMEM_RAW_STAGES + smem_state_stages: int = CFG.SMEM_S_STAGES + smem_decay_stages: int = CFG.SMEM_DECAY_STAGES + smem_intermediate_stages: int = CFG.SMEM_INTERMEDIATE_STAGES + smem_dq_stages: int = CFG.SMEM_DQ_STAGES + smem_dk_stages: int = CFG.SMEM_DK_STAGES + smem_dgate_stages: int = CFG.SMEM_DGATE_STAGES + smem_dv_stages: int = CFG.SMEM_DV_STAGES + smem_beta_stages: int = 4 + intermediate_tiles: int = 5 + tmem_dstate_acc_offset: int = 0 + tmem_dstate_inp_offset: int = 0 + tmem_state_k_acc_offset: int = 0 + tmem_u_acc_offset: int = 0 + tmem_du_acc_offset: int = 0 + tmem_dy_acc_offset: int = 0 + tmem_dq_acc_offset: int = 0 + tmem_dk_decay_acc_offset: int = 0 + tmem_dk_inv_acc_offset: int = 0 + tmem_dk_restore_acc_offset: int = 0 + tmem_qk_raw_stages: int = 4 + tmem_qraw_inp_offset: int = 0 + tmem_kraw_inp_offset: int = 0 + tmem_y_inp_offset: int = 0 + tmem_du_inp_offset: int = 0 + tmem_neg_beta_dy_inp_offset: int = 0 + tmem_state_inp_offset: int = 0 + buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES + + # ---- buffer cosizes / TMA bytes stamped by build_cfg ------------------------- + raw_qk_cosize: int = 0 + raw_v_cosize: int = 0 + raw_gate_cosize: int = 0 + operand_cosize: int = 0 + diag_cosize: int = 0 + intermediate_cosize: int = 0 + state_cosize: int = 0 + dq_cosize: int = 0 + dk_cosize: int = 0 + dgate_cosize: int = 0 + dv_cosize: int = 0 + + # TMA transaction bytes per stage + tma_q_bytes: int = 0 + tma_k_bytes: int = 0 + tma_gate_bytes: int = 0 + tma_do_bytes: int = 0 + tma_v_bytes: int = 0 + tma_state_bytes: int = 0 + + +def build_cfg( + io_dtype: Type[cutlass.Numeric], + *, + use_dstate_in: bool, + use_dstate0: bool, + l2norm: bool, + use_initial_state: bool, + q_ratio: int, + k_ratio: int, + v_ratio: int, + n_heads_out: int, + max_active_clusters: int, + dyn_sched: bool = False, +) -> KdaBwdCfg: + if io_dtype not in (cutlass.Float16, cutlass.BFloat16): + raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") + cfg = KdaBwdCfg( + io_dtype=io_dtype, + use_dstate_in=use_dstate_in, + use_dstate0=use_dstate0, + l2norm=l2norm, + use_initial_state=use_initial_state, + q_ratio=q_ratio, + k_ratio=k_ratio, + v_ratio=v_ratio, + n_heads_out=n_heads_out, + max_active_clusters=max_active_clusters, + dyn_sched=dyn_sched, ) + cfg.threads_per_cta = 16 * cfg.threads_per_warp + cfg.cg0_threads = len(cfg.compute_group_0_warp_ids) * cfg.threads_per_warp + cfg.cg2_threads = len(cfg.compute_group_2_warp_ids) * cfg.threads_per_warp + cfg.cg1_threads = len(cfg.compute_group_1_warp_ids) * cfg.threads_per_warp + cfg.tmem_user_threads = ( + 1 + len(cfg.compute_group_2_warp_ids) + len(cfg.compute_group_1_warp_ids) + len(cfg.compute_group_0_warp_ids) + ) * cfg.threads_per_warp + + cfg.tmem_dstate_acc_offset = 0 + cfg.tmem_dstate_inp_offset = cfg.d_k + cfg.tmem_state_inp_offset = cfg.tmem_dstate_inp_offset + cfg.d_k // 2 + cfg.tmem_state_k_acc_offset = cfg.tmem_state_inp_offset + cfg.d_v + cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t + cfg.tmem_du_acc_offset = cfg.tmem_u_acc_offset + cfg.b_t + cfg.tmem_dy_acc_offset = cfg.tmem_state_k_acc_offset + cfg.tmem_dq_acc_offset = cfg.tmem_du_acc_offset + cfg.b_t + cfg.tmem_dk_decay_acc_offset = cfg.tmem_dq_acc_offset + cfg.b_t + cfg.tmem_dk_inv_acc_offset = cfg.tmem_dk_decay_acc_offset + cfg.b_t + cfg.tmem_dk_restore_acc_offset = cfg.tmem_dk_inv_acc_offset + cfg.b_t + cfg.tmem_y_inp_offset = cfg.tmem_dk_restore_acc_offset + cfg.b_t + cfg.tmem_neg_beta_dy_inp_offset = cfg.tmem_y_inp_offset + cfg.tmem_du_inp_offset = cfg.tmem_y_inp_offset + cfg.b_t // 2 + cfg.tmem_qraw_inp_offset = cfg.tmem_du_inp_offset + cfg.b_t // 2 + cfg.tmem_kraw_inp_offset = cfg.tmem_qraw_inp_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) + assert cfg.tmem_kraw_inp_offset + cfg.tmem_qk_raw_stages * (cfg.b_t // 2) <= 512 + + cfg.raw_qk_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t + cfg.raw_v_cosize = cfg.smem_raw_stages * cfg.d_v * cfg.b_t + cfg.raw_gate_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t + cfg.operand_cosize = cfg.smem_decay_stages * cfg.b_t * cfg.d_k + cfg.diag_cosize = cfg.smem_decay_stages * (cfg.d_k // 16) * 256 + cfg.intermediate_cosize = cfg.smem_intermediate_stages * cfg.intermediate_tiles * cfg.b_t * cfg.b_t + cfg.state_cosize = cfg.smem_state_stages * cfg.d_k * cfg.d_v + cfg.dq_cosize = cfg.smem_dq_stages * cfg.b_t * cfg.d_k + cfg.dk_cosize = cfg.smem_dk_stages * cfg.b_t * cfg.d_k + cfg.dgate_cosize = cfg.smem_dgate_stages * cfg.b_t * cfg.d_k + cfg.dv_cosize = cfg.smem_dv_stages * cfg.b_t * cfg.d_v + cfg.tma_state_bytes = cfg.d_k * cfg.d_v * (io_dtype.width // 8) + cfg.tma_q_bytes = cfg.d_k * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_k_bytes = cfg.d_k * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_gate_bytes = cfg.d_k * cfg.b_t * 4 + cfg.tma_do_bytes = cfg.d_v * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_v_bytes = cfg.d_v * cfg.b_t * (cfg.io_dtype.width // 8) + return 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 + + +# ---- Torch adapter / host-side compilation --------------------------------------- + + +@lru_cache(maxsize=None) +def get_compiled_cache( + io_dtype_str: str, + cu_dtype_str: str, + HQ: int, + HK: int, + HV: int, + use_dstate_in: bool, + use_dstate0: bool, + l2norm: bool, + use_initial_state: bool, + dyn_sched: bool, +): + return {} + + +def chunk_kda_bwd_sm100( + q, + k, + v, + gate, + beta, + do, + state_checkpoints, + dq, + dk, + dv, + dgate, + dbeta, + cu_seqlens, + scale: float, + *, + initial_state=None, + d_initial_state=None, + d_final_state=None, + use_qk_l2norm_in_kernel: bool = False, + work_items=None, + work_count=None, + sched_ctr=None, + tensormap_workspace, + stream, +) -> None: + """Execute the Blackwell BT=16 chunked KDA backward kernel. + + All tensors must be contiguous and on the same CUDA device. + + Args: + q: ``(total_tokens, HQ, DK)`` float16/bfloat16 + k: ``(total_tokens, HK, DK)`` float16/bfloat16 + v: ``(total_tokens, HV, DV)`` float16/bfloat16 + gate: ``(total_tokens, HO, DK)`` fp32 natural-log per-channel decay + beta: ``(total_tokens, HO)`` fp32 post-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`` + 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 + dbeta: ``(total_tokens, HO)`` fp32, 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/d initial state), or None + d_final_state: fp32 ``(num_seqs, HO, DK, DV)`` 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 + work_items/work_count: split-K table (``common/split_k.py``, REQUIRED; + an uncut table row is the whole (b, h) sequence); each item + computes chunks ``[wstart, cend)`` backward and writes + gradients only for ``[wstart, wend)`` + sched_ctr: ``(2,)`` int32 zeroed scratch enabling the dynamic + (work-stealing) tile scheduler + tensormap_workspace: ``tensormap_workspace_bytes(module, B)`` bytes, + 128-byte aligned, for the per-(batch, head) TMA-descriptor + arrays (tail chunks clip/zero-fill in hardware) + """ + HQ = q.shape[1] + HK = k.shape[1] + HV = v.shape[1] + 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 + 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}") + for name, hh in (("HQ", HQ), ("HK", HK), ("HV", HV)): + if HO % hh != 0: + raise ValueError(f"{name}={hh} must divide {HO}") + B = cu_seqlens.shape[0] - 1 + + cu_stream = cuda_driver.CUstream(int(stream)) + cache = get_compiled_cache( + str(q.dtype), + str(cu_seqlens.dtype), + HQ, + HK, + HV, + use_dstate_in, + use_dstate0, + use_qk_l2norm_in_kernel, + use_initial_state, + dyn_sched, + ) + + if "compiled" not in cache: + io_dtype = get_dtype(q.dtype) + cfg = build_cfg( + io_dtype, + use_dstate_in=use_dstate_in, + use_dstate0=use_dstate0, + l2norm=use_qk_l2norm_in_kernel, + use_initial_state=use_initial_state, + q_ratio=HO // HQ, + k_ratio=HO // HK, + v_ratio=HO // HV, + n_heads_out=HO, + max_active_clusters=multiprocessor_count(current_device_id()), + dyn_sched=dyn_sched, + ) + + dstate0_cute = None + if use_dstate0: + dstate0_cute = from_dlpack(d_initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) + dstate_in_cute = None + if use_dstate_in: + dstate_in_cute = from_dlpack(d_final_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) + wi_cute = from_dlpack(work_items, assumed_align=16) + wi_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + wc_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() + sc_cute = None + if dyn_sched: + sc_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + + tensormap_ws_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + + 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( + host, + cfg, + beta_cute, + state_checkpoints_cute, + initial_state_cute, + dgate_cute, + dbeta_cute, + from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic(), + dstate0_cute, + dstate_in_cute, + wi_cute, + wc_cute, + sc_cute, + tensormap_ws_cute, + scale, + cu_stream, + options="--enable-tvm-ffi --opt-level 2", + ) + + # ---- per-(batch, head) descriptor arrays: rebuild on input change ------------ + # desc build runs every execute by contract (cu contents are data; + # buffer pointers may change) -- capture-safe, single tiny launch + if "build_descs" not in cache: + io_dtype = get_dtype(q.dtype) + + q_bd = from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=2) + k_bd = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_bd = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_bd = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2) + do_bd = from_dlpack(do, assumed_align=16).mark_layout_dynamic(leading_dim=2) + dq_bd = from_dlpack(dq, assumed_align=16).mark_layout_dynamic(leading_dim=2) + dk_bd = from_dlpack(dk, assumed_align=16).mark_layout_dynamic(leading_dim=2) + dv_bd = from_dlpack(dv, assumed_align=16).mark_layout_dynamic(leading_dim=2) + dgate_bd = from_dlpack(dgate, assumed_align=16).mark_layout_dynamic(leading_dim=2) + state_checkpoints_bd = from_dlpack(state_checkpoints, assumed_align=16).mark_layout_dynamic(leading_dim=3) + initial_state_bd = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) if use_initial_state else None + + cu_bd = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() + ws_bd = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + cache["build_descs"] = cute.compile( + build_descs, + io_dtype, + CFG.B_T, + q_bd, + k_bd, + v_bd, + gate_bd, + do_bd, + dq_bd, + dk_bd, + dv_bd, + dgate_bd, + state_checkpoints_bd, + initial_state_bd, + cu_bd, + ws_bd, + cu_stream, + options="--enable-tvm-ffi", + ) + cache["build_descs"](q, k, v, gate, do, dq, dk, dv, dgate, state_checkpoints, initial_state, cu_seqlens, tensormap_workspace, cu_stream) + cache["compiled"]( + beta, + state_checkpoints, + initial_state, + dgate, + dbeta, + cu_seqlens, + d_initial_state, + d_final_state, + work_items, + work_count, + sched_ctr, + tensormap_workspace, + scale, + cu_stream, + ) + + +# ---- Engine-side helpers: checkpoint-series bounds + entry-0 state seeding ---------------- diff --git a/python/cudnn/linear_attention/frost/kernel/kda_prefill_config.py b/python/cudnn/linear_attention/frost/kernel/kda_prefill_config.py index c7900f207..b77fece84 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_prefill_config.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_prefill_config.py @@ -37,10 +37,10 @@ class Cfg: # --- warp assignments (16 warps = 512 threads) --- COMPUTE_GROUP_0_WARP_IDS: Tuple[int, ...] = (0, 1, 2, 3, 4, 5, 6, 7) # decay-operand materialize (2-group ping-pong) COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (8, 9, 10, 11) # value-side TMEM / epilogue staging - SUPER_MMA_WARP_ID: int = 12 # register-MMA kk/qk + Neumann inverse + SUPER_MMA_WARP_ID: int = 12 # register-MMA KK/A + Neumann T_inv TCGEN05_MMA_WARP_ID: int = 13 # tcgen05 state GEMMs TMA_WARP_ID: int = 14 # q/k/v/gate TMA loads - EPILOGUE_WARP_ID: int = 15 # qk register-MMA + O store + EPILOGUE_WARP_ID: int = 15 # A register-MMA + O store # --- register split --- NUM_REGS_COMPUTE_GROUP_0: int = 160 @@ -56,9 +56,9 @@ class Cfg: SMEM_SCHED_STAGES: int = 8 SMEM_O_STAGES: int = 2 SMEM_DECAY_STAGES: int = 2 - SMEM_PAIRWISE_STAGES: int = 2 - SMEM_STATE_SCALE_DIAG_STAGES: int = 3 - QK_SCALE_READY_STAGES: int = 3 + SMEM_INTERMEDIATE_STAGES: int = 2 + SMEM_STATE_SCALE_DIAG_STAGES: int = 4 + QK_SCALE_READY_STAGES: int = 4 TMEM_Q_STATE_ACC_STAGES: int = 2 CLUSTER_SHAPE_MNK: Tuple[int, int, int] = (1, 1, 1) 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 71263b08b..b3819a9fd 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py @@ -15,7 +15,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Chunked Kimi Delta Attention (KDA) prefill kernel for Blackwell SM100 +"""Chunked Kimi Delta Attention (KDA) prefill kernel for Blackwell SM100/SM103 (Cutlass DSL), BT=16 tiling with a per-key-channel decay. Framework-neutral entry ``chunk_kda_sm100``. @@ -29,59 +29,59 @@ Pipeline (direct CUTLASS primitives, chunk_idx-size 16 KDA schedule): - load q/k/v/gate/beta - optional in-kernel L2-norm of q/k (L2NORM specialization) - exp2(g), exp2(-g), stage final-token exp2(g) as exp2(g_last) - super-MMA: kk/qk/Neumann inverse + apply beta - tcgen05-MMA: state*k/state*q/new-v/kv-update/qkv - store o, periodic state checkpoints, final state + load Q/K/V/Gate/Beta + optional in-kernel L2-norm of Q/K (L2NORM specialization) + exp2(G), exp2(-G), stage final-token exp2(G) as exp2(G_last) + super-MMA: KK/A/Neumann inverse (T_inv) + apply Beta + tcgen05-MMA: state*K/state*Q/U solve/state update/O + store O, periodic state checkpoints, final state 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, DV, DK]` (VK). +post-sigmoid, cu_seqlens int32, states/checkpoints `[N, HO, DK, DV]` (KV, v +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. Warp assignments (16 warps = 512 threads): - warps 0-7 : compute group 0 - gate prefix scan + decay/restore operands - warps 8-11 : compute group 1 - TMEM value side, o drain, state stores - warp 12 : super-MMA - register-MMA kk^T + Neumann inverse + warps 0-7 : compute group 0 - Gate prefix scan + decay/restore operands + warps 8-11 : compute group 1 - TMEM value side, O drain, state stores + warp 12 : super-MMA - register-MMA KK^T + Neumann inverse warp 13 : tcgen05-MMA - the six state GEMMs + the TMEM lifecycle warp 14 : TMA load - per-chunk input G->S loads - warp 15 : epilogue - register-MMA qk + the O TMA store + warp 15 : epilogue - register-MMA A + the O TMA store SMEM layout (~221 KB total): Buffer Bytes Stages - q / k / v raw 32768 8 <-- SW128 TMA ring (io dtype) + Q / K / V raw 32768 8 <-- SW128 TMA ring (io dtype) gate raw 65536 8 <-- fp32 prefix-scan source beta 512 8 <-- fp32 per-token scalars - dt_bias (+a_log slot) 516 1 <-- SAFE_GATE only K_inv 8192 2 <-- token-major ldmatrix/tcgen05 B operand K decay / Q decay 2x 8192 2 <-- tcgen05 SW128 K-box-major A/B operands K restore 8192 2 <-- tcgen05 B operand for the state update - state-scale diag 12288 3 <-- per-k-atom decay diagonal blocks - pairwise (A_inv / qk) 2048 2 <-- SW32 16x16 register-MMA tiles - o staging 8192 2 <-- W128 output drain + state-scale diag 12288 3 <-- per-K-atom decay diagonal blocks + intermediate (A / T_inv) 2048 2 <-- SW32 16x16 register-MMA tiles + O staging 8192 2 <-- W128 output drain TMEM layout (272 of 512 columns): Buffer Cols Purpose - state 0-127 S[DK,DV] fp32 recurrent state + state 0-127 state[DK,DV] fp32 recurrent state state inp 128-191 packed b16 A operand view of the state - q_state_acc 192-223 2-stage state*q -> o accumulator - state_k_acc 224-239 state*k fp32 accumulator - update_acc 240-255 update fp32 accumulator - rhs input 256-263 packed b16 A operand: beta * (v - state*k) - update input 264-271 packed b16 A operand: the update readback + q_state_acc 192-223 2-stage state*Q -> O accumulator + state_k_acc 224-239 state*K fp32 accumulator + u_acc 240-255 U fp32 accumulator + y_inp 256-263 packed b16 Y staging: Beta * (V - state*K) + u_inp 264-271 packed b16 U input (b16 U repack) GEMM schedule (tcgen05-MMA warp, in issue order per chunk): - state*k -> state_k_acc - state*q -> q_state_acc (the o acc) + state*K -> state_k_acc + state*Q -> q_state_acc (the O acc) state decay (diag blocks) - update = A_inv @ rhs -> update_acc - final_state += update @ k_restore - o += qk @ update -> q_state_acc + U = Y(T) @ T_inv -> u_acc + final_state += U @ K_restore + O += A @ U -> q_state_acc Requires a cutlass DSL build providing `cutlass.experimental.*`; not available in the pip nvidia-cutlass-dsl releases. @@ -99,24 +99,27 @@ from cutlass.cute.runtime import from_dlpack from ..common.split_k import decode_work_item -from ..common.thd import TENSOR_MAP_QWORDS, build_h_descs_kernel, build_qkv_load_descs_kernel +from ..common.host import get_dtype +from cudnn.frost.buffers import current_device_id, data_ptr +from cudnn.frost.device import multiprocessor_count +from ..common.thd import TENSOR_MAP_QWORDS, emit_checkpoint_seq_descs, emit_seq_descs from .kda_prefill_config import CFG from cudnn.frost.tile_dsl.barrier import ( advance, - arrive, MBarrier, PipelineState, Producer, - wait, ) from cudnn.frost.tile_dsl.handles import GmemTileTma, MmaDesc, SmemTile, tma_slice_runtime_desc -from cudnn.frost.tile_dsl.mma import mma_step, mma_ts +from cudnn.frost.tile_dsl.mma import mma_step, mma_ts_step from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_lin_S, swizzle_xor_128b from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire from cudnn.frost.tile_dsl.pointwise import ( opaque_f32_zero, + f16x2_to_f32, fadd2, fmul2, + ffma2, movmatrix_16b, mul_f16x2, fp32_to_fp16, @@ -136,40 +139,58 @@ class KdaBars(NamedTuple): - """Every inter-warp handoff as an ``MBarrier`` over its ring (mirrors GDN's - ``GdnBars``). Consumers track ``(idx, phase)`` inline; the producer tag selects + """Every inter-warp handoff as an ``MBarrier`` over its ring. Consumers + track ``(idx, phase)`` inline; the producer tag selects the arrive lowering (``TMA_LOAD``/``MMA_COMMIT``/``THREAD``).""" - mb_tma_done: MBarrier - mb_inputs_ready: MBarrier - mb_inputs_done: MBarrier + mb_q_ready: MBarrier + mb_q_done: MBarrier + mb_k_ready: MBarrier + mb_k_done: MBarrier + mb_v_ready: MBarrier + mb_v_done: MBarrier + mb_gate_ready: MBarrier + mb_gate_done: MBarrier + + mb_beta_ready: MBarrier + mb_beta_done: MBarrier + mb_o_acc_ready: MBarrier mb_o_acc_done: MBarrier mb_state_k_acc_ready: MBarrier - mb_update_acc_ready: MBarrier + mb_u_acc_ready: MBarrier + mb_state_inp_ready: MBarrier - mb_state_scale_diag_done: MBarrier - mb_kk_qk_super_mma_done: MBarrier - mb_kk_qk_mma_done: MBarrier - mb_k_restore_done: MBarrier - mb_rhs_ready: MBarrier - mb_update_ready: MBarrier - mb_final_state_stored: MBarrier + mb_y_inp_ready: MBarrier + mb_u_inp_ready: MBarrier + + mb_t_inv_ready: MBarrier + mb_t_inv_done: MBarrier mb_a_ready: MBarrier - mb_qk_acc_ready: MBarrier mb_a_done: MBarrier mb_qk_scale_ready: MBarrier - mb_k_decay_cg0_ready: MBarrier + mb_state_scale_diag_done: MBarrier + mb_k_decay_inv_cg0_ready: MBarrier + mb_decay_tcgen05_done: MBarrier + mb_decay_super_done: MBarrier + mb_k_restore_done: MBarrier + + mb_state_acc_read_done: MBarrier + mb_state_acc_done: MBarrier + mb_tmem_done: MBarrier + mb_o_tmastg_ready: MBarrier mb_o_tmastg_done: MBarrier + + mb_checkpoint_tmastg_ready: MBarrier + mb_checkpoint_tmastg_done: MBarrier + mb_sched_ready: MBarrier mb_sched_done: MBarrier - mb_h_tmastg_ready: MBarrier - mb_h_tmastg_done: MBarrier def make_kda_bars(cfg) -> KdaBars: - """Bars factory. MUST be called from inside ``_kernel`` (allocates the + """Bars factory. MUST be called from inside ``kernel`` (allocates the mbarrier rings in SMEM ahead of the data buffers).""" def alloc(n): @@ -180,71 +201,62 @@ def alloc(n): CG1_THREADS = len(cfg.compute_group_1_warp_ids) * WARP return KdaBars( - mb_tma_done=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.TMA_LOAD), - mb_inputs_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=WARP, producer=Producer.THREAD), - mb_inputs_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_q_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_q_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_k_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_k_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_v_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_v_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_gate_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_gate_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_beta_ready=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=WARP, producer=Producer.THREAD), + mb_beta_done=MBarrier(alloc(cfg.smem_raw_bar_stages), stages=cfg.smem_raw_bar_stages, init_count=WARP + CG1_THREADS, producer=Producer.THREAD), mb_o_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), mb_o_acc_done=MBarrier(alloc(cfg.tmem_q_state_acc_stages), stages=cfg.tmem_q_state_acc_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), - mb_update_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), + mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), mb_state_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_t_inv_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_t_inv_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_a_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_a_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_qk_scale_ready=MBarrier( + alloc(cfg.qk_scale_ready_stages), + stages=cfg.qk_scale_ready_stages, + init_count=CG0_GROUP_THREADS, + producer=Producer.THREAD, + ), mb_state_scale_diag_done=MBarrier( alloc(cfg.smem_state_scale_diag_stages), stages=cfg.smem_state_scale_diag_stages, init_count=1, producer=Producer.MMA_COMMIT, ), - mb_kk_qk_super_mma_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=2 * WARP, producer=Producer.THREAD), - mb_kk_qk_mma_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_k_decay_inv_cg0_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_decay_tcgen05_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_decay_super_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=2 * WARP, producer=Producer.THREAD), mb_k_restore_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), - mb_rhs_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_update_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_final_state_stored=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_a_ready=MBarrier(alloc(cfg.smem_pairwise_stages), stages=cfg.smem_pairwise_stages, init_count=WARP, producer=Producer.THREAD), - mb_qk_acc_ready=MBarrier(alloc(cfg.smem_pairwise_stages), stages=cfg.smem_pairwise_stages, init_count=WARP, producer=Producer.THREAD), - mb_a_done=MBarrier(alloc(cfg.smem_pairwise_stages), stages=cfg.smem_pairwise_stages, init_count=1, producer=Producer.MMA_COMMIT), - mb_qk_scale_ready=MBarrier( - alloc(cfg.qk_scale_ready_stages), - stages=cfg.qk_scale_ready_stages, - init_count=CG0_GROUP_THREADS, - producer=Producer.THREAD, - ), - mb_k_decay_cg0_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_state_acc_read_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_state_acc_done=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), + mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), mb_o_tmastg_ready=MBarrier(alloc(cfg.smem_o_stages), stages=cfg.smem_o_stages, init_count=CG1_THREADS, producer=Producer.THREAD), mb_o_tmastg_done=MBarrier(alloc(cfg.smem_o_stages), stages=cfg.smem_o_stages, init_count=WARP, producer=Producer.THREAD), + mb_checkpoint_tmastg_ready=MBarrier( + alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=CG1_THREADS, producer=Producer.THREAD + ), + mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=WARP, producer=Producer.THREAD), mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), - # H staging handshake: CG1 fills sH (ready), the epilogue TMA-stores - # and frees it (done); single stage - mb_h_tmastg_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), - mb_h_tmastg_done=MBarrier(alloc(1), stages=1, init_count=WARP, producer=Producer.THREAD), ) -# --------------------------------------------------------------------------- -# Device-side helpers / warp bodies -# --------------------------------------------------------------------------- - - -@cute.jit -def _gate_log2(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: - """Map raw gate to the log2-domain decay increment used by KDA.""" - - if cutlass.const_expr(cfg.safe_gate): - half = cutlass.Float32(0.5) - sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half - return cfg.gate_scale_log2 * sigmoid - # Default ABI: gate arrives in natural-log space - return raw_gate * cutlass.Float32(LOG2_E) - - -# --------------------------------------------------------------------------- -# Dynamic tile scheduler: global-ticket ring -# --------------------------------------------------------------------------- +# ---- Dynamic tile scheduler ------------------------------------------------------ @cute.jit -def _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): +def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): """TMA-warp side: pull the next tile off the global ticket, publish it.""" if cutlass.const_expr(cfg.dyn_sched): bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) @@ -260,7 +272,7 @@ def _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ct @cute.jit -def _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): +def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): """Consumer side: read the TMA warp's published next tile.""" if cutlass.const_expr(cfg.dyn_sched): bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) @@ -272,7 +284,7 @@ def _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): @cute.jit -def _tmaldg_warp( +def tmaldg_warp( cfg, total_tiles, bidx, @@ -281,26 +293,20 @@ def _tmaldg_warp( mWorkItems, mSched, sSched, - tma_tx_bytes, lane, - mBeta, sQ_raw, sK_raw, sV_raw, sGate_raw, - sBeta_raw, desc_q_base, desc_k_base, desc_v_base, desc_gate_base, bars, ) -> None: - """TMA-LDG warp role (warp 14): persistent tile-scheduler loop + per-chunk - q/k/v/gate G->S loads on one shared tx-count mbarrier plus the per-token - beta scalar stage. Loads go through the per-(batch, head) descriptor - array: head grouping and the sequence base live in each descriptor and - the token extent is capped per sequence, so coordinates are - sequence-relative and tail chunks zero-fill in hardware.""" + """TMA-LDG warp role (warp 14): persistent scheduler loop issuing the + per-chunk Q/K/V/Gate G->S loads.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) sQ_tma = SmemTile( base=sQ_raw, @@ -346,60 +352,64 @@ def _tmaldg_warp( tma_granu_elems=32, tma_subtile_stride_elems=(cfg.b_t * 32), ) - tma_index = PipelineState.start(phase=0) raw_index = PipelineState.start(phase=1) + raw_bar_index = PipelineState.start(phase=0) sched_state = PipelineState.start(phase=1) 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, cu_seqlens, mWorkItems - ) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) head_o = head_idx - slot = (batch_idx * cutlass.Int32(cfg.n_heads_out) + head_idx) * cutlass.Int32(TENSOR_MAP_QWORDS) + head_q = head_idx if cfg.q_ratio == 1 else head_idx // cutlass.Int32(cfg.q_ratio) + head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) + head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) + slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) desc_q_slot = (desc_q_base + slot).tospace(cutlass.AddressSpace.generic) desc_k_slot = (desc_k_base + slot).tospace(cutlass.AddressSpace.generic) desc_v_slot = (desc_v_base + slot).tospace(cutlass.AddressSpace.generic) desc_gate_slot = (desc_gate_base + slot).tospace(cutlass.AddressSpace.generic) - if nvvm.elect_sync(): + 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_gate_slot) for chunk_idx in cutlass.range(cstart, wend, 1, unroll=1): chunk_start = chunk_idx * cfg.b_t - bars.mb_inputs_done[raw_index.idx].wait(raw_index.phase) - # ---- q/k/v/gate TMA loads + per-token beta scalar stage ------------ - if nvvm.elect_sync(): - bars.mb_tma_done.arrive(n_bytes=tma_tx_bytes) - q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), chunk_start) - tma_load_tile(sQ_tma[raw_index.idx], q_slice, bars.mb_tma_done.smem_ptr, acquire=False) - k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), chunk_start) - tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_tma_done.smem_ptr, acquire=False) - v_slice = tma_slice_runtime_desc(desc_v_slot, cutlass.Int32(0), chunk_start) - tma_load_tile(sV_tma[raw_index.idx], v_slice, bars.mb_tma_done.smem_ptr, acquire=False) - gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), chunk_start) - tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_tma_done.smem_ptr, acquire=False) - - if lane < cfg.b_t: - token_idx = chunk_start + lane - beta_value = cutlass.Float32(0.0) - if token_idx < seqlen_b: - beta_value = mBeta[batch_start + token_idx, head_o].to(cutlass.Float32) - if cutlass.const_expr(cfg.beta_sigmoid): - # Roundtrip through the io dtype to bit-match host-side mBeta.sigmoid() - half = cutlass.Float32(0.5) - beta_value = (cute.math.tanh(beta_value * half, approx=True) * half + half).to(mBeta.element_type).to(cutlass.Float32) - sBeta_raw[raw_index.idx * cfg.b_t + lane] = beta_value - - bars.mb_tma_done.wait(tma_index.phase) - tma_index = advance(tma_index, 1) - bars.mb_inputs_ready[raw_index.idx].arrive() + + # ---- Q load ---------------------------------------------------------- + bars.mb_q_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_q_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_q_bytes) + q_slice = tma_slice_runtime_desc(desc_q_slot, cutlass.Int32(0), head_q, chunk_start) + tma_load_tile(sQ_tma[raw_index.idx], q_slice, bars.mb_q_ready[raw_bar_index.idx].smem_ptr, acquire=False) + + # ---- K load ---------------------------------------------------------- + bars.mb_k_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_k_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_k_bytes) + k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) + tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_bar_index.idx].smem_ptr, acquire=False) + + # ---- Gate load ------------------------------------------------------- + bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_gate_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_gate_bytes) + gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_bar_index.idx].smem_ptr, acquire=False) + + # ---- V load ---------------------------------------------------------- + bars.mb_v_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_v_ready[raw_bar_index.idx].arrive(n_bytes=cfg.tma_v_bytes) + 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_bar_index.idx].smem_ptr, acquire=False) + raw_index = advance(raw_index, cfg.smem_raw_stages) - tile_idx, sched_state = _sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) + tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) @cute.jit -def _super_mma_warp( +def super_mma_warp( cfg, total_tiles, bidx, @@ -409,23 +419,24 @@ def _super_mma_warp( sSched, lane, sK_inv_raw, - sPairwise_raw, + sIntermediate_raw, sBeta_raw, sK_decay_raw, bars, ) -> None: - """Super-MMA warp role (warp 12): persistent tile-scheduler loop + - register-MMA kk^T, L = beta*tril(kk), and the Neumann-series A_inv, - staged to pairwise SMEM.""" + """Super-MMA warp role (warp 12): persistent scheduler loop computing the + Neumann-series T_inv via register MMA.""" nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) raw_index = PipelineState.start(phase=0) - # ---- ldmatrix/stmatrix lane decode --------------------------------- + t_inv_free = PipelineState.start(phase=1) + k_decay_ready = PipelineState.start(phase=0) + + # ---- ldmatrix/stmatrix lane decode ------------------------------------------- rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - decay_key_mask = cutlass.Int32(8) ^ ((lhs_row_coord & cutlass.Int32(2)) * cutlass.Int32(16)) - elems_per_128b = cutlass.Int32(64) + decay_key_mask = cutlass.Int32(8) stsm_row_coord = lane & 7 stsm_col_coord = cutlass.Int32(0) if (lane // 8) & 1: @@ -433,26 +444,25 @@ def _super_mma_warp( if lane // 8 >= 2: stsm_col_coord = cutlass.Int32(8) stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) - gbase = cutlass.Int32(0) + cum_chunk_base = cutlass.Int32(0) 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, cu_seqlens, mWorkItems - ) - sk_nt = wend - cstart # processed chunks; ring bookkeeping runs on gbase + li - for li in cutlass.range(sk_nt, unroll=1): - gc = gbase + li - decay_stage = gc % cfg.smem_decay_stages - pairwise_stage = gc % cfg.smem_pairwise_stages + 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_chunks_tile = wend - cstart # processed chunks; ring bookkeeping runs on cum_chunk_base + local_chunk_idx + for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): + cum_chunk = cum_chunk_base + local_chunk_idx + decay_stage = k_decay_ready.idx + intermediate_stage = t_inv_free.idx sBeta_ptr = sBeta_raw.data_ptr() + raw_index.idx * cfg.b_t sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) - sPairwise_ptr = sPairwise_raw.data_ptr() + pairwise_stage * (2 * cfg.b_t * cfg.b_t) + sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (2 * cfg.b_t * cfg.b_t) - bars.mb_a_done[pairwise_stage].wait(((gc // cfg.smem_pairwise_stages) + 1) % 2) - bars.mb_k_decay_cg0_ready[decay_stage].wait((gc // cfg.smem_decay_stages) % 2) - # ---- kk^T register MMA over the K blocks --------------------------- + bars.mb_k_decay_inv_cg0_ready[decay_stage].wait(k_decay_ready.phase) + k_decay_ready = advance(k_decay_ready, cfg.smem_decay_stages) + + # ---- KK = K_decay @ K_inv^T ------------------------------------------ kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): kk_acc[accum_idx] = cutlass.Float32(0.0) @@ -461,7 +471,7 @@ def _super_mma_warp( # Load B operand k_inv_col = k_block * 16 + rhs_col_offset k_inv_segment = k_inv_col // 64 - rhs_vec = nvvm.ldmatrix( + rhs_frag = nvvm.ldmatrix( sK_inv_ptr + k_inv_segment * (cfg.b_t * 64) + rhs_row_coord * 64 @@ -471,56 +481,50 @@ def _super_mma_warp( ) # Load A operand storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask - storage_slice = storage_key // elems_per_128b - key_in_slice = storage_key - storage_slice * elems_per_128b - storage_phase = key_in_slice // cutlass.Int32(16) - byte_in_slice = ( - lhs_row_coord * cutlass.Int32(128) - + storage_phase * cutlass.Int32(32) - + (key_in_slice - storage_phase * cutlass.Int32(16)) * cutlass.Int32(2) - ) - kk_lhs_vec = nvvm.ldmatrix( + storage_slice = storage_key // 64 + kk_lhs_frag = nvvm.ldmatrix( sK_decay_ptr - + storage_slice * cutlass.Int32(cfg.b_t) * elems_per_128b - + ((byte_in_slice ^ ((lhs_row_coord & cutlass.Int32(7)) << 4)) // cutlass.Int32(2)), + + storage_slice * (cfg.b_t * 64) + + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( kk_acc, - (kk_lhs_vec[0], kk_lhs_vec[1], kk_lhs_vec[2], kk_lhs_vec[3]), - (rhs_vec[0], rhs_vec[1], rhs_vec[2], rhs_vec[3]), + (kk_lhs_frag[0], kk_lhs_frag[1], kk_lhs_frag[2], kk_lhs_frag[3]), + (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - # ---- L = beta * tril(kk, -1) fragment ------------------------------ + + # ---- L = Beta * tril(KK, -1) fragment -------------------------------- + bars.mb_beta_ready[raw_index.idx].wait(raw_index.phase) row_lo = lane // 4 row_hi = row_lo + cutlass.Int32(8) beta_lo = (sBeta_ptr + row_lo).load().to(cutlass.Float32) beta_hi = (sBeta_ptr + row_hi).load().to(cutlass.Float32) - l_frag = cutlass.Array(cutlass.Float32, 8, alignment=16) + l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): - row_coord = row_lo - beta_scale = beta_lo - if cutlass.const_expr(accum_idx % 4 >= 2): - row_coord = row_hi - beta_scale = beta_hi + row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) - lower = kk_acc[accum_idx] if row_coord > col_coord else cutlass.Float32(0.0) - l_frag[accum_idx] = lower * beta_scale - l_a0 = fp32_to_fp16(l_frag[0], l_frag[1], dtype=cfg.io_dtype) - l_a1 = fp32_to_fp16(l_frag[2], l_frag[3], dtype=cfg.io_dtype) - l_a2 = fp32_to_fp16(l_frag[4], l_frag[5], dtype=cfg.io_dtype) - l_a3 = fp32_to_fp16(l_frag[6], l_frag[7], dtype=cfg.io_dtype) + l_regs[accum_idx] = kk_acc[accum_idx] if row_coord > col_coord else cutlass.Float32(0.0) + for pair in cutlass.range_constexpr(4): + beta_scale = beta_hi if cutlass.const_expr(pair % 2 == 1) else beta_lo + l_regs[2 * pair], l_regs[2 * pair + 1] = fmul2(l_regs[2 * pair], l_regs[2 * pair + 1], beta_scale, beta_scale) + bars.mb_beta_done[raw_index.idx].arrive() + l_a0 = fp32_to_fp16(l_regs[0], l_regs[1], dtype=cfg.io_dtype) + l_a1 = fp32_to_fp16(l_regs[2], l_regs[3], dtype=cfg.io_dtype) + l_a2 = fp32_to_fp16(l_regs[4], l_regs[5], dtype=cfg.io_dtype) + l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) - # ---- A_inv = I - L, then three Neumann doubling rounds ------------- - ainv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + # ---- T_inv = I - L, then three Neumann doubling rounds --------------- + tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): row_coord = row_lo if cutlass.const_expr(accum_idx % 4 >= 2): @@ -529,18 +533,19 @@ def _super_mma_warp( if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) eye = cutlass.Float32(1.0) if row_coord == col_coord else cutlass.Float32(0.0) - ainv_acc[accum_idx] = eye - l_values[accum_idx] + tinv_acc[accum_idx] = eye - l_values[accum_idx] lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) for _round in cutlass.range_constexpr(3): - # Lpow <- Lpow @ Lpow (packed A-layout fragments, B via movmatrix) + # ---- Lpow = Lpow @ Lpow ------------------------------------------ sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): sq_acc[accum_idx] = cutlass.Float32(0.0) mma_step( sq_acc, (lpow_a0, lpow_a1, lpow_a2, lpow_a3), - (movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3)), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), k_step=0, M=16, N=16, @@ -550,48 +555,60 @@ def _super_mma_warp( lpow_a1 = fp32_to_fp16(sq_acc[2], sq_acc[3], dtype=cfg.io_dtype) lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) - # A_inv <- A_inv + A_inv @ Lpow, keeping A_inv in registers + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) + # ---- T_inv += T_inv @ Lpow --------------------------------------- upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): upd_acc[accum_idx] = cutlass.Float32(0.0) + tinv_p0 = fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype) + tinv_p1 = fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype) + tinv_p2 = fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype) + tinv_p3 = fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype) mma_step( upd_acc, - ( - fp32_to_fp16(ainv_acc[0], ainv_acc[1], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[2], ainv_acc[3], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[4], ainv_acc[5], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[6], ainv_acc[7], dtype=cfg.io_dtype), - ), - (movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3)), + (tinv_p0, tinv_p1, tinv_p2, tinv_p3), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), k_step=0, M=16, N=16, ab_dtype=cfg.io_dtype, ) - for accum_idx in cutlass.range_constexpr(8): - ainv_acc[accum_idx] = ainv_acc[accum_idx].to(cfg.io_dtype).to(cutlass.Float32) + upd_acc[accum_idx] - + tinv_lo0, tinv_hi0 = f16x2_to_f32(tinv_p0, dtype=cfg.io_dtype) + tinv_lo1, tinv_hi1 = f16x2_to_f32(tinv_p1, dtype=cfg.io_dtype) + tinv_lo2, tinv_hi2 = f16x2_to_f32(tinv_p2, dtype=cfg.io_dtype) + tinv_lo3, tinv_hi3 = f16x2_to_f32(tinv_p3, dtype=cfg.io_dtype) + tinv_acc[0] = tinv_lo0 + upd_acc[0] + tinv_acc[1] = tinv_hi0 + upd_acc[1] + tinv_acc[2] = tinv_lo1 + upd_acc[2] + tinv_acc[3] = tinv_hi1 + upd_acc[3] + tinv_acc[4] = tinv_lo2 + upd_acc[4] + tinv_acc[5] = tinv_hi2 + upd_acc[5] + tinv_acc[6] = tinv_lo3 + upd_acc[6] + tinv_acc[7] = tinv_hi3 + upd_acc[7] + + bars.mb_t_inv_done[intermediate_stage].wait(t_inv_free.phase) + t_inv_free = advance(t_inv_free, cfg.smem_intermediate_stages) nvvm.stmatrix( - sPairwise_ptr + (cfg.b_t * cfg.b_t) + stsm_idx, + sIntermediate_ptr + (cfg.b_t * cfg.b_t) + stsm_idx, [ - fp32_to_fp16(ainv_acc[0], ainv_acc[1], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[2], ainv_acc[3], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[4], ainv_acc[5], dtype=cfg.io_dtype), - fp32_to_fp16(ainv_acc[6], ainv_acc[7], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype), ], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8, ) nvvm.fence_proxy("async.shared", space="cta") - bars.mb_a_ready[pairwise_stage].arrive() - bars.mb_kk_qk_super_mma_done[decay_stage].arrive() - raw_index = advance(raw_index, cfg.smem_raw_stages) - gbase += sk_nt - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + bars.mb_t_inv_ready[intermediate_stage].arrive() + bars.mb_decay_super_done[decay_stage].arrive() + raw_index = advance(raw_index, cfg.smem_raw_bar_stages) + cum_chunk_base += num_chunks_tile + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @cute.jit -def _tcgen05_mma_warp( +def tcgen05_mma_warp( cfg, total_tiles, bidx, @@ -599,24 +616,38 @@ def _tcgen05_mma_warp( cu_seqlens, mWorkItems, sSched, - tmem_hold, - sPairwise, + tmem_base_slot, + sIntermediate, sK_decay, sK_restore, sQ_decay, sState_scale_diag, bars, ) -> None: - """tcgen05-MMA warp role (warp 13): persistent tile-scheduler loop, issues - all six state GEMMs in dependency order and owns the TMEM lifecycle.""" + """tcgen05-MMA warp role (warp 13): persistent scheduler loop issuing + every tcgen05 GEMM and owning the TMEM lifecycle.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) - tmem_base = tmem_hold.load() + nvvm.tcgen05_alloc(tmem_base_slot, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_slot.load() + state_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) + state_dsts = tuple(nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset + k * 16, cutlass.Float32) for k in range(cfg.d_k // 16)) + state_k_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_k_acc_offset, cutlass.Float32) + u_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_acc_offset, cutlass.Float32) + y_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_inp_offset, cutlass.Int8) + u_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_inp_offset, cutlass.Int8) + state_dst_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset, cutlass.Float32) state_inp_index = PipelineState.start(phase=0) - rhs_index = PipelineState.start(phase=0) - update_index = PipelineState.start(phase=0) - final_state_index = PipelineState.start(phase=0) + state_read_index = PipelineState.start(phase=0) + y_inp_index = PipelineState.start(phase=0) + u_inp_index = PipelineState.start(phase=0) qk_scale_index = PipelineState.start(phase=0) - # ---- chunk-invariant GEMM descriptors ------------------------------ + k_decay_ready = PipelineState.start(phase=0) + intermediate_ready = PipelineState.start(phase=0) + o_acc_free = PipelineState.start(phase=1) + + # ---- chunk-invariant GEMM descriptors ---------------------------------------- bpe = cfg.io_dtype.width // 8 idesc_acc = nvvm.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, @@ -666,7 +697,7 @@ def _tcgen05_mma_warp( idesc=idesc_diag, kind=nvvm.Tcgen05MMAKind.F16, ) - bmm_pairwise_desc = MmaDesc( + bmm_intermediate_desc = MmaDesc( M=cfg.d_v, N=cfg.b_t, K=cfg.b_t, @@ -690,128 +721,131 @@ def _tcgen05_mma_warp( idesc=idesc_final_state, kind=nvvm.Tcgen05MMAKind.F16, ) - gbase = cutlass.Int32(0) + STATE_A_SEG = bmm_state_desc.sps_B * bmm_state_desc.tmem_advance_A + STATE_B_SEG = bmm_state_desc.smem_subtile_B >> 4 + cum_chunk_base = cutlass.Int32(0) 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, cu_seqlens, mWorkItems - ) - sk_nt = wend - cstart - for li in cutlass.range(sk_nt, unroll=1): - gc = gbase + li - q_state_acc_stage = gc % cfg.tmem_q_state_acc_stages - decay_stage = gc % cfg.smem_decay_stages + 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_chunks_tile = wend - cstart + for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): + cum_chunk = cum_chunk_base + local_chunk_idx + have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_chunk_idx > 0 + q_state_acc_stage = o_acc_free.idx + decay_stage = k_decay_ready.idx state_scale_diag_stage = qk_scale_index.idx - pairwise_stage = gc % cfg.smem_pairwise_stages + intermediate_stage = intermediate_ready.idx sK_decay_stage = sK_decay[decay_stage] sQ_decay_stage = sQ_decay[decay_stage] sK_restore_stage = sK_restore[decay_stage] sState_scale_diag_stage = sState_scale_diag[state_scale_diag_stage] - sPairwise_stage = sPairwise[pairwise_stage] - - # ---- state*k -> state_k_acc ---------------------------------------- - bars.mb_k_decay_cg0_ready[decay_stage].wait((gc // cfg.smem_decay_stages) % 2) - bars.mb_state_inp_ready.wait(state_inp_index.phase) - state_inp_index = advance(state_inp_index, 1) - desc_k_decay = sK_decay_stage.desc() - - state_a_tmem_base = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) - mma_ts( - bmm_state_desc, - state_a_tmem_base, - desc_k_decay, - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_k_acc_offset, cutlass.Float32), - accumulate=False, - ) + sIntermediate_stage = sIntermediate[intermediate_stage] + + # ---- state_k = S(T) @ K_decay^T -------------------------------------- + bars.mb_k_decay_inv_cg0_ready[decay_stage].wait(k_decay_ready.phase) + k_decay_ready = advance(k_decay_ready, cfg.smem_decay_stages) + if have_state: + bars.mb_state_inp_ready.wait(state_inp_index.phase) + state_inp_index = advance(state_inp_index, 1) + desc_k_decay = sK_decay_stage.desc() + + for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + mma_ts_step( + bmm_state_desc, + state_inp_ptr.subview(s * STATE_A_SEG), + desc_k_decay + s * STATE_B_SEG, + state_k_acc_ptr, + k, + cutlass.Boolean(s + k > 0), + ) - if nvvm.elect_sync(): - bars.mb_state_k_acc_ready.arrive(cta_group=1) + if elect_one: + bars.mb_state_k_acc_ready.arrive(cta_group=1) - # ---- state*q -> q_state_acc (stays live until qk@update fuses into o) + # ---- q_state = state(T) @ Q_decay^T --------------------------------- bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) - bars.mb_o_acc_done[q_state_acc_stage].wait(((gc // cfg.tmem_q_state_acc_stages + cutlass.Int32(1)) % cutlass.Int32(2))) - desc_q_decay = sQ_decay_stage.desc() - - state_a_tmem_base = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) - mma_ts( - bmm_state_desc, - state_a_tmem_base, - desc_q_decay, - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_q_state_acc_offset + q_state_acc_stage * cfg.b_t, cutlass.Float32), - accumulate=False, - ) + bars.mb_o_acc_done[q_state_acc_stage].wait(o_acc_free.phase) + o_acc_free = advance(o_acc_free, cfg.tmem_q_state_acc_stages) + q_state_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_q_state_acc_offset + q_state_acc_stage * cfg.b_t, cutlass.Float32) + if have_state: + desc_q_decay = sQ_decay_stage.desc() + for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + mma_ts_step( + bmm_state_desc, + state_inp_ptr.subview(s * STATE_A_SEG), + desc_q_decay + s * STATE_B_SEG, + q_state_acc_ptr, + k, + cutlass.Boolean(s + k > 0), + ) - if nvvm.elect_sync(): - bars.mb_kk_qk_mma_done[decay_stage].arrive(cta_group=1) - - # ---- state decay (per-k-atom diag blocks) ---------------------------- - desc_diag = sState_scale_diag_stage.desc() - - for k_block in cutlass.range_constexpr(cfg.d_k // 16): - mma_ts( - bmm_diag_desc, - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset + k_block * 8, cutlass.Int8), - desc_diag.advance_start_address(k_block * 256 * 2), - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_offset + k_block * 16, cutlass.Float32), - accumulate=False, - ) + if elect_one: + bars.mb_decay_tcgen05_done[decay_stage].arrive(cta_group=1) - if nvvm.elect_sync(): - bars.mb_state_scale_diag_done[state_scale_diag_stage].arrive(cta_group=1) + # ---- S decay = S(T) @ diag(exp2(G_last)) --------- + if cutlass.const_expr(cfg.enable_checkpoints): + if have_state: + bars.mb_state_acc_read_done.wait(state_read_index.phase) + state_read_index = advance(state_read_index, 1) + if have_state: + desc_diag = sState_scale_diag_stage.desc() + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + mma_ts_step( + bmm_diag_desc, + state_inp_ptr.subview(k_block * bmm_diag_desc.tmem_advance_A), + desc_diag.advance_start_address(k_block * 256 * 2), + state_dsts[k_block], + 0, + cutlass.Boolean(False), + ) - # ---- update = A_inv @ rhs -> update_acc ------------------------------ - bars.mb_a_ready[pairwise_stage].wait((gc // cfg.smem_pairwise_stages) % 2) - bars.mb_rhs_ready.wait(rhs_index.phase) - rhs_index = advance(rhs_index, 1) - lhs_tmem = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_rhs_inp_offset, cutlass.Int8) - desc_pairwise = sPairwise_stage.shifted((cfg.b_t * cfg.b_t)).desc() - mma_ts( - bmm_pairwise_desc, - lhs_tmem, - desc_pairwise, - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_update_acc_offset, cutlass.Float32), - accumulate=False, - ) - if nvvm.elect_sync(): - bars.mb_update_acc_ready.arrive(cta_group=1) + if elect_one: + bars.mb_state_scale_diag_done[state_scale_diag_stage].arrive(cta_group=1) - # ---- final_state += update @ k_restore ------------------------------- - bars.mb_update_ready.wait(update_index.phase) - update_index = advance(update_index, 1) - update_tmem = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_update_inp_offset, cutlass.Int8) + # ---- U = Y(T) @ T_inv ------------------------------------------------ + bars.mb_t_inv_ready[intermediate_stage].wait(intermediate_ready.phase) + bars.mb_y_inp_ready.wait(y_inp_index.phase) + y_inp_index = advance(y_inp_index, 1) + d_int = sIntermediate_stage.shifted((cfg.b_t * cfg.b_t)).desc() + mma_ts_step(bmm_intermediate_desc, y_inp_ptr, d_int, u_acc_ptr, 0, cutlass.Boolean(False)) + if elect_one: + bars.mb_t_inv_done[intermediate_stage].arrive(cta_group=1) + bars.mb_u_acc_ready.arrive(cta_group=1) + + # ---- final_state += U(T) @ K_restore --------------------------------- + bars.mb_u_inp_ready.wait(u_inp_index.phase) + u_inp_index = advance(u_inp_index, 1) desc_k_restore = sK_restore_stage.desc() - mma_ts( - bmm_final_state_desc, - update_tmem, - desc_k_restore, - nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_offset, cutlass.Float32), - accumulate=True, - ) - if nvvm.elect_sync(): + mma_ts_step(bmm_final_state_desc, u_inp_ptr, desc_k_restore, state_dst_ptr, 0, have_state) + if elect_one: bars.mb_k_restore_done[decay_stage].arrive(cta_group=1) - - # ---- o += qk @ update -> q_state_acc ---------------------------------- - bars.mb_qk_acc_ready[pairwise_stage].wait((gc // cfg.smem_pairwise_stages) % 2) - lhs_tmem = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_update_inp_offset, cutlass.Int8) - desc_pairwise = sPairwise_stage.desc() - mma_ts( - bmm_pairwise_desc, - lhs_tmem, - desc_pairwise, + bars.mb_state_acc_done.arrive(cta_group=1) + + # ---- O += U(T) @ A --------------------------------------------------- + bars.mb_a_ready[intermediate_stage].wait(intermediate_ready.phase) + intermediate_ready = advance(intermediate_ready, cfg.smem_intermediate_stages) + d_int = sIntermediate_stage.desc() + mma_ts_step( + bmm_intermediate_desc, + u_inp_ptr, + d_int, nvvm.make_tmem_ptr(tmem_base + cfg.tmem_q_state_acc_offset + q_state_acc_stage * cfg.b_t, cutlass.Float32), - accumulate=True, + 0, + have_state, ) - if nvvm.elect_sync(): + if elect_one: bars.mb_o_acc_ready.arrive(cta_group=1) - bars.mb_a_done[pairwise_stage].arrive(cta_group=1) + bars.mb_a_done[intermediate_stage].arrive(cta_group=1) qk_scale_index = advance(qk_scale_index, cfg.smem_state_scale_diag_stages) - bars.mb_final_state_stored.wait(final_state_index.phase) - final_state_index = advance(final_state_index, 1) - gbase += sk_nt - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + cum_chunk_base += num_chunks_tile + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + bars.mb_tmem_done[0].wait(0) + nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) nvvm.tcgen05_dealloc( nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), cutlass.Int32(512), @@ -820,7 +854,7 @@ def _tcgen05_mma_warp( @cute.jit -def _epilogue_warp( +def epilogue_warp( cfg, total_tiles, bidx, @@ -832,27 +866,23 @@ def _epilogue_warp( mO, sK_inv_raw, sO_raw, - sPairwise_raw, + sIntermediate_raw, sQ_decay_raw, - sH_raw, + sCheckpoint_raw, desc_o_base, - desc_h_base, + desc_checkpoint_base, checkpoint_every_n_tokens, bars, ) -> None: - """Epilogue warp role (warp 15): persistent tile-scheduler loop + - register-MMA qk (causal), A_inv qk staging, the O TMA store drain - (split-K warmup chunks drain the SMEM stage but never store), and the - per-chunk H TMA store when checkpoints are enabled. O and H stores go - through the per-(batch, head) descriptor arrays: the token / entry - extents are capped per sequence, so partial tails are clipped by the - hardware and the H coordinate is sequence-local.""" + """Epilogue warp role (warp 15): persistent scheduler loop computing the + causal A tile via register MMA and draining the O/checkpoint TMA stores.""" + elect_one = nvvm.elect_sync() nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(cfg.enable_checkpoints): - sH_tma = SmemTile( - base=sH_raw, + sCheckpoint_tma = SmemTile( + base=sCheckpoint_raw, elems_per_stage=(cfg.d_k * cfg.d_v), - stages=1, + stages=cfg.smem_checkpoint_stages, leading_byte_offset=0, stride_byte_offset=0, layout=0, @@ -860,7 +890,7 @@ def _epilogue_warp( tma_granu_elems=64, tma_subtile_stride_elems=cfg.d_k * 64, ) - h_ready_index = PipelineState.start(phase=0) + checkpoint_ready_index = PipelineState.start(phase=0) sO_tma = SmemTile( base=sO_raw, elems_per_stage=(cfg.b_t * cfg.d_v), @@ -873,13 +903,14 @@ def _epilogue_warp( tma_subtile_stride_elems=cfg.b_t * 64, ) qk_scale_index = PipelineState.start(phase=0) - # ---- ldmatrix/stmatrix lane decode --------------------------------- + a_free = PipelineState.start(phase=1) + + # ---- ldmatrix/stmatrix lane decode ------------------------------------------- rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) - decay_key_mask = cutlass.Int32(8) ^ ((lhs_row_coord & cutlass.Int32(2)) * cutlass.Int32(16)) - elems_per_128b = cutlass.Int32(64) + decay_key_mask = cutlass.Int32(8) stsm_row_coord = lane & 7 stsm_col_coord = cutlass.Int32(0) if (lane // 8) & 1: @@ -887,44 +918,46 @@ def _epilogue_warp( if lane // 8 >= 2: stsm_col_coord = cutlass.Int32(8) stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) - gbase = cutlass.Int32(0) + cum_chunk_base = cutlass.Int32(0) 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, cu_seqlens, mWorkItems - ) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) head_o = head_idx - o_slot = (batch_idx * cutlass.Int32(cfg.n_heads_out) + head_idx) * cutlass.Int32(TENSOR_MAP_QWORDS) + o_slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) desc_o_slot = (desc_o_base + o_slot).tospace(cutlass.AddressSpace.generic) if cutlass.const_expr(cfg.enable_checkpoints): - desc_h_slot = (desc_h_base + o_slot).tospace(cutlass.AddressSpace.generic) - if nvvm.elect_sync(): - tma_tensormap_acquire(desc_h_slot) - if nvvm.elect_sync(): + desc_checkpoint_slot = (desc_checkpoint_base + o_slot).tospace(cutlass.AddressSpace.generic) + checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + checkpoint_quot = (cstart + cutlass.Int32(1)) // checkpoint_chunks + checkpoint_mod = (cstart + cutlass.Int32(1)) % checkpoint_chunks + if elect_one: + tma_tensormap_acquire(desc_checkpoint_slot) + if elect_one: tma_tensormap_acquire(desc_o_slot) - sk_nt = wend - cstart - for li in cutlass.range(sk_nt, unroll=1): - chunk_idx = cstart + li - gc = gbase + li - decay_stage = gc % cfg.smem_decay_stages - pairwise_stage = gc % cfg.smem_pairwise_stages + num_chunks_tile = wend - cstart + 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 + decay_stage = cum_chunk % cfg.smem_decay_stages + intermediate_stage = a_free.idx + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) sQ_decay_ptr = sQ_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) - sPairwise_ptr = sPairwise_raw.data_ptr() + pairwise_stage * (2 * cfg.b_t * cfg.b_t) + sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (2 * cfg.b_t * cfg.b_t) - bars.mb_a_done[pairwise_stage].wait(((gc // cfg.smem_pairwise_stages) + 1) % 2) bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) - # ---- qk register MMA (inclusive-causal) ---------------------------- - qk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + + # ---- A = Q_decay @ K_inv^T ------------------------------------------ + a_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) for accum_idx in cutlass.range_constexpr(8): - qk_acc[accum_idx] = cutlass.Float32(0.0) + a_acc[accum_idx] = cutlass.Float32(0.0) for k_block in cutlass.range_constexpr((cfg.d_k // 16)): # Load B operand k_inv_col = k_block * 16 + rhs_col_offset k_inv_segment = k_inv_col // 64 - rhs_vec = nvvm.ldmatrix( + rhs_frag = nvvm.ldmatrix( sK_inv_ptr + k_inv_segment * (cfg.b_t * 64) + rhs_row_coord * 64 @@ -934,26 +967,19 @@ def _epilogue_warp( ) # Load A operand storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask - storage_slice = storage_key // elems_per_128b - key_in_slice = storage_key - storage_slice * elems_per_128b - storage_phase = key_in_slice // cutlass.Int32(16) - byte_in_slice = ( - lhs_row_coord * cutlass.Int32(128) - + storage_phase * cutlass.Int32(32) - + (key_in_slice - storage_phase * cutlass.Int32(16)) * cutlass.Int32(2) - ) - qk_lhs_vec = nvvm.ldmatrix( + storage_slice = storage_key // 64 + a_lhs_frag = nvvm.ldmatrix( sQ_decay_ptr - + storage_slice * cutlass.Int32(cfg.b_t) * elems_per_128b - + ((byte_in_slice ^ ((lhs_row_coord & cutlass.Int32(7)) << 4)) // cutlass.Int32(2)), + + storage_slice * (cfg.b_t * 64) + + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), 4, nvvm.MMALayout.ROW, ) mma_step( - qk_acc, - (qk_lhs_vec[0], qk_lhs_vec[1], qk_lhs_vec[2], qk_lhs_vec[3]), - (rhs_vec[0], rhs_vec[1], rhs_vec[2], rhs_vec[3]), + a_acc, + (a_lhs_frag[0], a_lhs_frag[1], a_lhs_frag[2], a_lhs_frag[3]), + (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), k_step=0, M=16, N=16, @@ -967,78 +993,104 @@ def _epilogue_warp( col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) if cutlass.const_expr(accum_idx % 2 == 1): col_coord = col_coord + cutlass.Int32(1) - qk_acc[accum_idx] = qk_acc[accum_idx] if row_coord >= col_coord else cutlass.Float32(0.0) + a_acc[accum_idx] = a_acc[accum_idx] if row_coord >= col_coord else cutlass.Float32(0.0) + bars.mb_a_done[intermediate_stage].wait(a_free.phase) + a_free = advance(a_free, cfg.smem_intermediate_stages) nvvm.stmatrix( - sPairwise_ptr + stsm_idx, + sIntermediate_ptr + stsm_idx, [ - fp32_to_fp16(qk_acc[0], qk_acc[1], dtype=cfg.io_dtype), - fp32_to_fp16(qk_acc[2], qk_acc[3], dtype=cfg.io_dtype), - fp32_to_fp16(qk_acc[4], qk_acc[5], dtype=cfg.io_dtype), - fp32_to_fp16(qk_acc[6], qk_acc[7], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[0], a_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[2], a_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[4], a_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(a_acc[6], a_acc[7], dtype=cfg.io_dtype), ], nvvm.MMALayout.ROW, shape=nvvm.StoreShape.M8N8, ) nvvm.fence_proxy("async.shared", space="cta") - bars.mb_qk_acc_ready[pairwise_stage].arrive() - bars.mb_kk_qk_super_mma_done[decay_stage].arrive() + bars.mb_a_ready[intermediate_stage].arrive() + bars.mb_decay_super_done[decay_stage].arrive() qk_scale_index = advance(qk_scale_index, cfg.qk_scale_ready_stages) - # ---- O drain: staged output tile -> GMEM TMA store ----------------- - if li > 0: + # ---- checkpoint + O drain: checkpoint stores first (CG1 stages the checkpoint before O) ------------- + if local_chunk_idx > 0: output_chunk = chunk_idx - cutlass.Int32(1) output_chunk_start = output_chunk * cfg.b_t - o_stage = (gc - cutlass.Int32(1)) % cfg.smem_o_stages - bars.mb_o_tmastg_ready[o_stage].wait(((gc - cutlass.Int32(1)) // cfg.smem_o_stages) % 2) - o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), output_chunk_start) - if cutlass.const_expr(cfg.split_k): - # warmup chunks stage O to SMEM but never store it - if output_chunk >= wstart: - tma_store_tile(sO_tma[o_stage], o_slice, acquire=False) + o_stage = (cum_chunk - cutlass.Int32(1)) % cfg.smem_o_stages + did_checkpoint = cutlass.Int32(0) + checkpoint_stage = cutlass.Int32(0) + if cutlass.const_expr(cfg.enable_checkpoints): + # ---- checkpoint store ---------------------------------------- + do_checkpoint = checkpoint_mod == 0 + do_checkpoint = do_checkpoint and chunk_idx >= wstart + checkpoint_stage = checkpoint_ready_index.idx + 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_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() - else: + did_checkpoint = cutlass.Int32(1) + checkpoint_mod = checkpoint_mod + cutlass.Int32(1) + if checkpoint_mod == checkpoint_chunks: + checkpoint_mod = cutlass.Int32(0) + checkpoint_quot = checkpoint_quot + cutlass.Int32(1) + bars.mb_o_tmastg_ready[o_stage].wait(((cum_chunk - cutlass.Int32(1)) // cfg.smem_o_stages) % 2) + o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), head_o, output_chunk_start) + did_o = cutlass.Int32(0) + if output_chunk >= wstart: tma_store_tile(sO_tma[o_stage], o_slice, acquire=False) tma_store_commit() - tma_store_wait(0) - bars.mb_o_tmastg_done[o_stage].arrive() - # ---- H store: CG1 staged the state entering chunk_idx ---------- - if cutlass.const_expr(cfg.enable_checkpoints): - if li > 0: - tokens_done = chunk_idx * cutlass.Int32(cfg.b_t) - do_h = tokens_done % checkpoint_every_n_tokens == 0 - if cutlass.const_expr(cfg.split_k): - do_h = do_h and chunk_idx >= wstart - if do_h: - bars.mb_h_tmastg_ready.wait(h_ready_index.phase) - h_ready_index = advance(h_ready_index, 1) - # sequence-local entry: the per-(b,h) descriptor folds the - # sequence base into GLOBAL_ADDRESS and caps the extent - h_entry = tokens_done // checkpoint_every_n_tokens - cutlass.Int32(1) - h_slice = tma_slice_runtime_desc(desc_h_slot, cutlass.Int32(0), cutlass.Int32(0), h_entry) - tma_store_tile(sH_tma[0], h_slice, acquire=False) - tma_store_commit() + did_o = cutlass.Int32(1) + if cutlass.const_expr(cfg.enable_checkpoints): + if did_checkpoint == 1 and did_o == 1: + tma_store_wait(1) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() tma_store_wait(0) - bars.mb_h_tmastg_done.arrive() - # ---- last computed chunk drain (always owned: it is wend - 1) ------ - if sk_nt > 0: + bars.mb_o_tmastg_done[o_stage].arrive() + if did_checkpoint == 1 and did_o == 0: + tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() + bars.mb_o_tmastg_done[o_stage].arrive() + if did_checkpoint == 0: + if did_o == 1: + tma_store_wait(0) + bars.mb_o_tmastg_done[o_stage].arrive() + else: + tma_store_wait(0) + bars.mb_o_tmastg_done[o_stage].arrive() + + # ---- last computed chunk drain (always owned: it is wend - 1) ------------ + if num_chunks_tile > 0: output_chunk = wend - cutlass.Int32(1) - og = gbase + sk_nt - cutlass.Int32(1) + last_cum_chunk = cum_chunk_base + num_chunks_tile - cutlass.Int32(1) output_chunk_start = output_chunk * cfg.b_t - o_stage = og % cfg.smem_o_stages - bars.mb_o_tmastg_ready[o_stage].wait((og // cfg.smem_o_stages) % 2) - # a partial last chunk is clipped by the descriptor's token extent - o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), output_chunk_start) + o_stage = last_cum_chunk % cfg.smem_o_stages + bars.mb_o_tmastg_ready[o_stage].wait((last_cum_chunk // cfg.smem_o_stages) % 2) + o_slice = tma_slice_runtime_desc(desc_o_slot, cutlass.Int32(0), head_o, output_chunk_start) tma_store_tile(sO_tma[o_stage], o_slice, acquire=False) tma_store_commit() tma_store_wait(0) bars.mb_o_tmastg_done[o_stage].arrive() - gbase += sk_nt - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + cum_chunk_base += num_chunks_tile + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def gate_scale(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: + """Map raw gate to the log2-domain decay increment used by KDA.""" + + if cutlass.const_expr(cfg.safe_gate): + half = cutlass.Float32(0.5) + sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half + return cfg.gate_scale_log2 * sigmoid + return raw_gate * cutlass.Float32(LOG2_E) @cute.jit -def _compute0_warp_group( +def compute0_warp_group( cfg, total_tiles, bidx, @@ -1053,18 +1105,18 @@ def _compute0_warp_group( mDt_bias, sK_inv_raw, sGate_raw, + mBeta, + sBeta_raw, sK_raw, sQ_raw, - sV_raw, sK_decay_raw, sK_restore_raw, sQ_decay_raw, sState_scale_diag_raw, bars, ) -> None: - """CG0 warp role (warps 0-7, two ping-pong groups): persistent - tile-scheduler loop + gate prefix scan and the decay/restore operand - materialization into tcgen05 SMEM.""" + """CG0 warp role (warps 0-7, two ping-pong groups): persistent scheduler + loop for the Gate prefix scan and decay/restore operand materialization.""" nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] cg0_group_id = cg0_warp // cfg.cg0_warps_per_group @@ -1072,32 +1124,34 @@ def _compute0_warp_group( prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane cg0_a_log_exp = cutlass.Float32(1.0) cg0_dt_bias_value = cutlass.Float32(0.0) - gbase = cutlass.Int32(0) + cum_chunk_base = cutlass.Int32(0) sched_state = PipelineState.start(phase=0) tile_idx = cutlass.Int32(bidx) + opaque_one = opaque_f32_zero() + cutlass.Float32(1.0) 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, cu_seqlens, mWorkItems - ) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) head_o = head_idx - sk_nt = wend - cstart + num_chunks_tile = wend - cstart if cutlass.const_expr(cfg.safe_gate): - # per-head safe-gate constants straight from GMEM (the head - # changes per tile, so there is no SMEM staging to handshake) - if sk_nt > 0: + if num_chunks_tile > 0: cg0_a_log_exp = cute.math.exp2(mA_log[head_o].to(cutlass.Float32) * LOG2_E, fastmath=True) cg0_dt_bias_value = mDt_bias[head_o, prefix_dim].to(cutlass.Float32) - for li in cutlass.range(cg0_group_id, sk_nt, cfg.cg0_group_count, unroll=1): - chunk_idx = cstart + li - gc = gbase + li + # tile entry: both ping-pong groups inherit each other's delivery proofs (parity-swap guard) + nvvm.barrier_cta_sync(cfg.cg0_tile_entry_barrier_id, thread_count=cfg.cg0_group_count * cfg.cg0_threads_per_group) + group_cum_chunk_start = cum_chunk_base + cutlass.Int32(cg0_group_id) + diag_ring_idx = group_cum_chunk_start % cutlass.Int32(cfg.smem_state_scale_diag_stages) + diag_ring_phase = (group_cum_chunk_start // cutlass.Int32(cfg.smem_state_scale_diag_stages)) % cutlass.Int32(2) + for local_chunk_idx in cutlass.range(cg0_group_id, num_chunks_tile, cfg.cg0_group_count, unroll=1): + chunk_idx = cstart + local_chunk_idx + cum_chunk = cum_chunk_base + local_chunk_idx chunk_start = chunk_idx * cfg.b_t - decay_stage = gc % cfg.smem_decay_stages - raw_stage = gc % cfg.smem_raw_stages - state_scale_diag_stage = gc % cfg.smem_state_scale_diag_stages + decay_stage = cum_chunk % cfg.smem_decay_stages + raw_stage = cum_chunk % cfg.smem_raw_stages + raw_bar_stage = cum_chunk % cfg.smem_raw_bar_stages + state_scale_diag_stage = diag_ring_idx qk_scale_ready_stage = state_scale_diag_stage sQ_ptr = sQ_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) sK_ptr = sK_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) - sV_ptr = sV_raw.data_ptr() + raw_stage * (cfg.d_v * cfg.b_t) sGate_ptr = sGate_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) @@ -1105,65 +1159,36 @@ def _compute0_warp_group( sK_restore_ptr = sK_restore_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) sState_scale_diag_ptr = sState_scale_diag_raw.data_ptr() + state_scale_diag_stage * ((cfg.d_k // 16) * 256) - bars.mb_inputs_ready[raw_stage].wait((gc // cfg.smem_raw_stages) % 2) - - # ---- tail chunk: zero-fill raw staging past seqlen ----------------- - if chunk_start + cutlass.Int32(cfg.b_t) > seqlen_b: - if cg0_local_warp == 0: - f16_zero = mQ.element_type(0.0) - f16_zero_vec = cutlass.Vector.from_elements( - ( - f16_zero, - f16_zero, - f16_zero, - f16_zero, - f16_zero, - f16_zero, - f16_zero, - f16_zero, - ), - mQ.element_type, - ) - f32_zero = cutlass.Float32(0.0) - f32_zero_vec = cutlass.Vector.from_elements( - (f32_zero, f32_zero, f32_zero, f32_zero), - cutlass.Float32, - ) - for row in cutlass.range_constexpr(cfg.b_t): - token_idx = chunk_start + cutlass.Int32(row) - if token_idx >= seqlen_b: - if lane < (cfg.d_k // 8): - f16_dim_base = lane * 8 - f16_segment = f16_dim_base // 64 - f16_segment_dim = f16_dim_base - f16_segment * 64 - f16_idx = f16_segment * (cfg.b_t * 64) + row * 64 + swizzle_xor_128b(row, f16_segment_dim, elem_bytes=2) - (sQ_ptr + f16_idx).store(f16_zero_vec, alignment=16) - (sK_ptr + f16_idx).store(f16_zero_vec, alignment=16) - (sV_ptr + f16_idx).store(f16_zero_vec, alignment=16) - if lane < (cfg.d_k // 4): - f32_dim_base = lane * 4 - f32_segment = f32_dim_base // 32 - f32_segment_dim = f32_dim_base - f32_segment * 32 - f32_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) - (sGate_ptr + f32_idx).store(f32_zero_vec, alignment=16) - nvvm.barrier_cta_sync(cfg.nbar_cg0_group0_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) + # ---- Beta scalars --------------------------------------------------- + if cg0_local_warp == 0: + bars.mb_beta_done[raw_bar_stage].wait(((cum_chunk // cfg.smem_raw_bar_stages) + 1) % 2) + if lane < cfg.b_t: + token_idx = chunk_idx * cfg.b_t + lane + beta_value = cutlass.Float32(0.0) + if token_idx < seqlen_b: + beta_value = mBeta[batch_start + token_idx, head_o].to(cutlass.Float32) + if cutlass.const_expr(cfg.beta_sigmoid): + half = cutlass.Float32(0.5) + beta_value = (cute.math.tanh(beta_value * half, approx=True) * half + half).to(mBeta.element_type).to(cutlass.Float32) + sBeta_raw[raw_bar_stage * cfg.b_t + lane] = beta_value + bars.mb_beta_ready[raw_bar_stage].arrive() + bars.mb_gate_ready[raw_bar_stage].wait((cum_chunk // cfg.smem_raw_bar_stages) % 2) row_group_start = cg0_local_warp * (cfg.b_t // cfg.cg0_warps_per_group) lane_row_group = lane // 8 lane_in_row_group = lane - lane_row_group * 8 decay_row = row_group_start + lane_row_group - - g_prefix_ptr = sGate_ptr + decay_key_mask = cutlass.Int32(8) prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane - # ---- gate prefix scan: cumulative log-gate per key channel -------- - # gathers first, math second: a load consumed in the same iteration - # serializes on LDS latency + + # ---- Gate prefix scan ----------------------------------------------- + f32_segment = prefix_dim // 32 + prefix_seg_base = f32_segment * (cfg.b_t * 32) + prefix_col = prefix_dim - f32_segment * 32 gate_raw = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) for row in cutlass.range_constexpr(cfg.b_t): - f32_segment = prefix_dim // 32 - f32_segment_dim = prefix_dim - f32_segment * 32 - prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) + prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) gate_raw[row] = (sGate_ptr + prefix_idx).load() g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) if cutlass.const_expr(cfg.safe_gate): @@ -1174,11 +1199,11 @@ def _compute0_warp_group( row1 = row0 + 1 gate0 = cg0_a_log_exp * (gate_raw[row0] + cg0_dt_bias_value) gate1 = cg0_a_log_exp * (gate_raw[row1] + cg0_dt_bias_value) - gate0 = _gate_log2( + gate0 = gate_scale( cfg, gate0, ) - gate1 = _gate_log2( + gate1 = gate_scale( cfg, gate1, ) @@ -1191,7 +1216,7 @@ def _compute0_warp_group( gate = gate_raw[row] token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) if token_idx < seqlen_b: - gate = _gate_log2( + gate = gate_scale( cfg, gate, ) @@ -1211,19 +1236,17 @@ def _compute0_warp_group( g_prefix_regs[row1] = prefix1 prefix_acc = prefix1 + # ---- exp2(G): stage prefixes + final-token decay --------------------- for row in cutlass.range_constexpr(cfg.b_t): g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) - # ---- exp2(g): stage prefixes + final-token decay ------------------ exp_g_last = g_prefix_regs[cfg.b_t - 1] for row in cutlass.range_constexpr(cfg.b_t): - f32_segment = prefix_dim // 32 - f32_segment_dim = prefix_dim - f32_segment * 32 - prefix_idx = f32_segment * (cfg.b_t * 32) + row * 32 + swizzle_xor_128b(row, f32_segment_dim, elem_bytes=4) + prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) - # ---- state-scale diag: stage exp2(g_last) decay blocks ------------- - bars.mb_state_scale_diag_done[state_scale_diag_stage].wait((gc // cfg.smem_state_scale_diag_stages + 1) % 2) + # ---- state-scale diag: stage exp2(G_last) decay blocks --------------- + bars.mb_state_scale_diag_done[state_scale_diag_stage].wait(diag_ring_phase ^ cutlass.Int32(1)) block = prefix_dim // cutlass.Int32(16) coord = prefix_dim - block * cutlass.Int32(16) storage_col = coord ^ cutlass.Int32((cfg.b_t // 2)) @@ -1231,37 +1254,46 @@ def _compute0_warp_group( diag_idx = swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) sState_scale_diag_ptr[diag_idx] = exp_g_last.to(cfg.io_dtype) - nvvm.barrier_cta_sync(cfg.nbar_cg0_group0_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) + nvvm.barrier_cta_sync(cfg.cg0_group_sync_barrier_base_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) - k_inv_words = cutlass.Array(cutlass.Int32, 2 * 4, alignment=16) + bars.mb_q_ready[raw_bar_stage].wait((cum_chunk // cfg.smem_raw_bar_stages) % 2) + bars.mb_k_ready[raw_bar_stage].wait((cum_chunk // cfg.smem_raw_bar_stages) % 2) + k_inv_pack = cutlass.Array(cutlass.Int32, 2 * 4, alignment=16) raw_q_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) - # ---- optional q/k L2-norm + K_inv staging -------------------------- - q_sum_sq = cutlass.Float32(0.0) - k_sum_sq = cutlass.Float32(0.0) + + # ---- optional Q/K L2-norm ------------------------------------------- + if cutlass.const_expr(cfg.l2norm): + q_sq_even = opaque_f32_zero() + k_sq_even = opaque_f32_zero() + q_sq_odd = opaque_f32_zero() + k_sq_odd = opaque_f32_zero() for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 f16_segment = dim_base // 64 f16_segment_dim = dim_base - f16_segment * 64 raw_f16_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) - raw_q_vec = (sQ_ptr + raw_f16_idx).load(count=8, alignment=16) - raw_k_vec = (sK_ptr + raw_f16_idx).load(count=8, alignment=16) - raw_q_vec_f32 = raw_q_vec.to(cutlass.Float32) - raw_k_vec_f32 = raw_k_vec.to(cutlass.Float32) + raw_q_frag = (sQ_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_k_frag = (sK_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_q_vec_f32 = raw_q_frag.to(cutlass.Float32) + raw_k_vec_f32 = raw_k_frag.to(cutlass.Float32) for dim_offset in cutlass.range_constexpr(8): q_val = raw_q_vec_f32[dim_offset] k_val = raw_k_vec_f32[dim_offset] raw_q_regs[reg_base + dim_offset] = q_val raw_k_regs[reg_base + dim_offset] = k_val - q_sum_sq = q_sum_sq + q_val * q_val - k_sum_sq = k_sum_sq + k_val * k_val + if cutlass.const_expr(cfg.l2norm): + if cutlass.const_expr(dim_offset % 2 == 0): + q_sq_even, k_sq_even = ffma2(q_val, k_val, q_val, k_val, q_sq_even, k_sq_even) + else: + q_sq_odd, k_sq_odd = ffma2(q_val, k_val, q_val, k_val, q_sq_odd, k_sq_odd) - # opaque 1.0: keeps the no-l2norm packed-mul operands out of libNVVM's - # constant folder (the documented inline_ptx "n"-constraint ICE) - q_inv_norm = opaque_f32_zero() + cutlass.Float32(1.0) - k_inv_norm = opaque_f32_zero() + cutlass.Float32(1.0) + q_inv_norm = opaque_one + k_inv_norm = opaque_one if cutlass.const_expr(cfg.l2norm): + q_sum_sq = q_sq_even + q_sq_odd + k_sum_sq = k_sq_even + k_sq_odd q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 4, 31, kind=nvvm.Shfl.BFLY)) q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 2, 31, kind=nvvm.Shfl.BFLY)) q_sum_sq = q_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, q_sum_sq, 1, 31, kind=nvvm.Shfl.BFLY)) @@ -1272,11 +1304,9 @@ def _compute0_warp_group( q_inv_norm = cute.math.rsqrt(cute.math.max(q_sum_sq, norm_floor_sq), fastmath=True) k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) - # ---- decay/restore operands: exp2(+-g) applied per key channel ----- + # ---- decay/restore operands: exp2(+-G) ------------------------------ exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) - # both halves' prefix gathers up front: the sK_inv/sK_decay stores - # below would otherwise pin the second half's loads behind them for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 @@ -1285,24 +1315,20 @@ def _compute0_warp_group( f32_segment = f32_dim_base // 32 f32_segment_dim = f32_dim_base - f32_segment * 32 g_prefix_idx = f32_segment * (cfg.b_t * 32) + decay_row * 32 + swizzle_xor_128b(decay_row, f32_segment_dim, elem_bytes=4) - exp_g_vec = (g_prefix_ptr + g_prefix_idx).load(count=4, alignment=16) + exp_g_frag = (sGate_ptr + g_prefix_idx).load(count=4, alignment=16) exp_g_last_idx = f32_segment * (cfg.b_t * 32) + (cfg.b_t - 1) * 32 + swizzle_xor_128b((cfg.b_t - 1), f32_segment_dim, elem_bytes=4) - exp_g_last_vec = (g_prefix_ptr + exp_g_last_idx).load(count=4, alignment=16) + exp_g_last_frag = (sGate_ptr + exp_g_last_idx).load(count=4, alignment=16) f32_reg_base = reg_base + f32_group * 4 - exp_g_regs[f32_reg_base] = exp_g_vec[0] - exp_g_regs[f32_reg_base + 1] = exp_g_vec[1] - exp_g_regs[f32_reg_base + 2] = exp_g_vec[2] - exp_g_regs[f32_reg_base + 3] = exp_g_vec[3] - exp_g_last_regs[f32_reg_base] = exp_g_last_vec[0] - exp_g_last_regs[f32_reg_base + 1] = exp_g_last_vec[1] - exp_g_last_regs[f32_reg_base + 2] = exp_g_last_vec[2] - exp_g_last_regs[f32_reg_base + 3] = exp_g_last_vec[3] + for j in cutlass.range_constexpr(4): + exp_g_regs[f32_reg_base + j] = exp_g_frag[j] + exp_g_last_regs[f32_reg_base + j] = exp_g_last_frag[j] for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 - # ---- k decay operand: exp2(g) * k ------------------------------ - k_decay_words = cutlass.Array(cutlass.Int32, 4, alignment=16) + + # ---- K decay + K_inv operands: K * exp2(+G) and K * exp2(-G) ----- + k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 dim1 = dim0 + 1 @@ -1311,58 +1337,55 @@ def _compute0_warp_group( k_value0, k_value1 = fmul2(raw_k_regs[raw_reg_idx0], raw_k_regs[raw_reg_idx1], k_inv_norm, k_inv_norm) k_pair = fp32_to_fp16(k_value0, k_value1, dtype=cfg.io_dtype) exp_g_pair = fp32_to_fp16(exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1], dtype=cfg.io_dtype) - k_decay_words[pair_idx] = mul_f16x2(k_pair, exp_g_pair, cfg.io_dtype) + k_decay_pack[pair_idx] = mul_f16x2(k_pair, exp_g_pair, cfg.io_dtype) exp_neg_g0 = cute.math.rcp(exp_g_regs[raw_reg_idx0], approx=True, ftz=True) exp_neg_g1 = cute.math.rcp(exp_g_regs[raw_reg_idx1], approx=True, ftz=True) - k_inv_words[dim_half * 4 + pair_idx] = fp32_to_fp16(k_value0 * exp_neg_g0, k_value1 * exp_neg_g1, dtype=cfg.io_dtype) + exp_neg_pair = fp32_to_fp16(exp_neg_g0, exp_neg_g1, dtype=cfg.io_dtype) + k_inv_pack[dim_half * 4 + pair_idx] = mul_f16x2(k_pair, exp_neg_pair, cfg.io_dtype) k_inv_vec = cutlass.Vector.from_elements( ( - k_inv_words[dim_half * 4], - k_inv_words[dim_half * 4 + 1], - k_inv_words[dim_half * 4 + 2], - k_inv_words[dim_half * 4 + 3], + k_inv_pack[dim_half * 4], + k_inv_pack[dim_half * 4 + 1], + k_inv_pack[dim_half * 4 + 2], + k_inv_pack[dim_half * 4 + 3], ), cutlass.Int32, ).bitcast(cfg.io_dtype) k_decay_vec = cutlass.Vector.from_elements( ( - k_decay_words[0], - k_decay_words[1], - k_decay_words[2], - k_decay_words[3], + k_decay_pack[0], + k_decay_pack[1], + k_decay_pack[2], + k_decay_pack[3], ), cutlass.Int32, ).bitcast(cfg.io_dtype) if cutlass.const_expr(dim_half == 0): - operand_done_phase = ((gc // cfg.smem_decay_stages) + 1) % 2 - bars.mb_kk_qk_super_mma_done[decay_stage].wait(operand_done_phase) - bars.mb_kk_qk_mma_done[decay_stage].wait(operand_done_phase) + operand_done_phase = ((cum_chunk // cfg.smem_decay_stages) + 1) % 2 + bars.mb_decay_super_done[decay_stage].wait(operand_done_phase) + bars.mb_decay_tcgen05_done[decay_stage].wait(operand_done_phase) f16_segment = dim_base // 64 f16_segment_dim = dim_base - f16_segment * 64 k_inv_swizzled_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) (sK_inv_ptr + k_inv_swizzled_idx).store(k_inv_vec, alignment=16) - key_mask = cutlass.Int32(8) ^ (decay_row & cutlass.Int32(2)) * cutlass.Int32(16) - decay_storage_dim_base = dim_base ^ key_mask - decay_linear_idx_base = decay_row * cfg.d_k + decay_storage_dim_base - sw128_elems_per_128b = 128 // 2 - sw128_row = decay_linear_idx_base // cfg.d_k - sw128_col = decay_linear_idx_base - sw128_row * cfg.d_k - sw128_slice = sw128_col // sw128_elems_per_128b - sw128_col_in_slice = sw128_col - sw128_slice * sw128_elems_per_128b - sw128_slice_linear = sw128_row * sw128_elems_per_128b + sw128_col_in_slice - sw128_byte = sw128_slice_linear * 2 - sw128_mask = (sw128_byte >> 7 & 7) << 4 - decay_swizzled_idx_base = sw128_slice * cfg.b_t * sw128_elems_per_128b + (sw128_byte ^ sw128_mask) // 2 - (sK_decay_ptr + decay_swizzled_idx_base).store(k_decay_vec, alignment=16) + storage_key = dim_base ^ decay_key_mask + storage_slice = storage_key // 64 + decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( + decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 + ) + (sK_decay_ptr + decay_swizzled_idx).store(k_decay_vec, alignment=16) nvvm.fence_proxy("async.shared", space="cta") - bars.mb_k_decay_cg0_ready[decay_stage].arrive() + bars.mb_k_decay_inv_cg0_ready[decay_stage].arrive() + bars.mb_q_done[raw_stage].arrive() + bars.mb_k_done[raw_stage].arrive() + bars.mb_gate_done[raw_stage].arrive() - # ---- q decay operand ----------------------------------------------- + # ---- Q_decay operand: Q * q_inv_norm -------------------------------- for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 - q_decay_words = cutlass.Array(cutlass.Int32, 4, alignment=16) + q_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 dim1 = dim0 + 1 @@ -1371,65 +1394,61 @@ def _compute0_warp_group( q_value0, q_value1 = fmul2(raw_q_regs[raw_reg_idx0], raw_q_regs[raw_reg_idx1], q_inv_norm, q_inv_norm) q_pair = fp32_to_fp16(q_value0, q_value1, dtype=cfg.io_dtype) exp_g_pair = fp32_to_fp16(exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1], dtype=cfg.io_dtype) - q_decay_words[pair_idx] = mul_f16x2(q_pair, exp_g_pair, cfg.io_dtype) + q_decay_pack[pair_idx] = mul_f16x2(q_pair, exp_g_pair, cfg.io_dtype) q_decay_vec = cutlass.Vector.from_elements( ( - q_decay_words[0], - q_decay_words[1], - q_decay_words[2], - q_decay_words[3], + q_decay_pack[0], + q_decay_pack[1], + q_decay_pack[2], + q_decay_pack[3], ), cutlass.Int32, ).bitcast(cfg.io_dtype) - key_mask = cutlass.Int32(8) ^ (decay_row & cutlass.Int32(2)) * cutlass.Int32(16) - decay_storage_dim_base = dim_base ^ key_mask - decay_linear_idx_base = decay_row * cfg.d_k + decay_storage_dim_base - sw128_elems_per_128b = 128 // 2 - sw128_row = decay_linear_idx_base // cfg.d_k - sw128_col = decay_linear_idx_base - sw128_row * cfg.d_k - sw128_slice = sw128_col // sw128_elems_per_128b - sw128_col_in_slice = sw128_col - sw128_slice * sw128_elems_per_128b - sw128_slice_linear = sw128_row * sw128_elems_per_128b + sw128_col_in_slice - sw128_byte = sw128_slice_linear * 2 - sw128_mask = (sw128_byte >> 7 & 7) << 4 - decay_swizzled_idx_base = sw128_slice * cfg.b_t * sw128_elems_per_128b + (sw128_byte ^ sw128_mask) // 2 - (sQ_decay_ptr + decay_swizzled_idx_base).store(q_decay_vec, alignment=16) - - bars.mb_k_restore_done[decay_stage].wait(((gc // cfg.smem_decay_stages + 1) % 2)) - - # ---- k_restore operand ---------------------------------------------- + storage_key = dim_base ^ decay_key_mask + storage_slice = storage_key // 64 + decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( + decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 + ) + (sQ_decay_ptr + decay_swizzled_idx).store(q_decay_vec, alignment=16) + + # ---- K_restore operand: K_inv * exp_g_last -------------------------- + bars.mb_k_restore_done[decay_stage].wait(((cum_chunk // cfg.smem_decay_stages + 1) % 2)) for dim_half in cutlass.range_constexpr(2): dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 reg_base = dim_half * 8 - k_restore_words = cutlass.Array(cutlass.Int32, 4, alignment=16) + k_restore_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) for pair_idx in cutlass.range_constexpr(4): dim0 = pair_idx * 2 dim1 = dim0 + 1 exp_g_last_pair = fp32_to_fp16(exp_g_last_regs[reg_base + dim0], exp_g_last_regs[reg_base + dim1], dtype=cfg.io_dtype) - k_restore_words[pair_idx] = mul_f16x2(k_inv_words[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) + k_restore_pack[pair_idx] = mul_f16x2(k_inv_pack[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) storage_row = decay_row ^ (cfg.b_t // 2) f16_segment = dim_base // 64 f16_segment_dim = dim_base - f16_segment * 64 k_restore_idx = f16_segment * (cfg.b_t * 64) + storage_row * 64 + swizzle_xor_128b(storage_row, f16_segment_dim, elem_bytes=2) k_restore_vec = cutlass.Vector.from_elements( ( - k_restore_words[0], - k_restore_words[1], - k_restore_words[2], - k_restore_words[3], + k_restore_pack[0], + k_restore_pack[1], + k_restore_pack[2], + k_restore_pack[3], ), cutlass.Int32, ).bitcast(cfg.io_dtype) (sK_restore_ptr + k_restore_idx).store(k_restore_vec, alignment=16) nvvm.fence_proxy("async.shared", space="cta") bars.mb_qk_scale_ready[qk_scale_ready_stage].arrive() - gbase += sk_nt - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + diag_ring_idx = diag_ring_idx + cutlass.Int32(cfg.cg0_group_count) + wrapped = diag_ring_idx >= cutlass.Int32(cfg.smem_state_scale_diag_stages) + diag_ring_idx = diag_ring_idx - cutlass.Int32(cfg.smem_state_scale_diag_stages) if wrapped else diag_ring_idx + diag_ring_phase = diag_ring_phase ^ (cutlass.Int32(1) if wrapped else cutlass.Int32(0)) + cum_chunk_base += num_chunks_tile + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) @cute.jit -def _compute1_warp_group( +def compute1_warp_group( cfg, total_tiles, bidx, @@ -1438,528 +1457,518 @@ def _compute1_warp_group( mWorkItems, sSched, lane, - tmem_hold, + tmem_base_slot, warp_idx, - mS_out, - mS_init, + mState_out, + mState_init, mO, sO_raw, sBeta_raw, sV_raw, - sH_raw, + sCheckpoint_raw, checkpoint_every_n_tokens, scale, bars, ) -> None: - """CG1 warp role (warps 8-11): persistent tile-scheduler loop + value-side - TMEM staging (state input, rhs, update), output drain to SMEM, and - checkpoint/final-state stores (split-K: only owned entries).""" + """CG1 warp role (warps 8-11): persistent scheduler loop for the + value-side TMEM staging, output drain, and state stores.""" nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) sO_ptr = sO_raw.data_ptr() - sH_ptr = sH_raw.data_ptr() if cutlass.const_expr(cfg.enable_checkpoints) else sO_raw.data_ptr() - h_done_index = PipelineState.start(phase=1) # sH starts free - tmem_base = tmem_hold.load() + sCheckpoint_ptr = sCheckpoint_raw.data_ptr() if cutlass.const_expr(cfg.enable_checkpoints) else sO_raw.data_ptr() + checkpoint_done_index = PipelineState.start(phase=1) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = tmem_base_slot.load() tmem_col = tmem_base & 0xFFFF tmem_row = tmem_base >> 16 - tmem_sp = warp_idx % (cfg.d_v // cfg.threads_per_warp) - # ldmatrix.x4/stmatrix.x4 COL lane decode shared by the v loads and o stores - ov_tok = (lane // 16) * 8 + (lane & 7) - ov_col = ((lane // 8) & 1) * 8 - row_id = tmem_row + tmem_sp * cfg.threads_per_warp - value_dim = tmem_sp * cfg.threads_per_warp + lane + tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) + ov_token_coord = (lane // 16) * 8 + (lane & 7) + ov_col_coord = ((lane // 8) & 1) * 8 + row_id = tmem_row + tmem_subpartition * cfg.threads_per_warp + value_dim = tmem_subpartition * cfg.threads_per_warp + lane + value_dim_base = tmem_subpartition * cfg.threads_per_warp + row_addr = row_id << 16 + row16_addr = (row_id + 16) << 16 + st_row_addr = tmem_row << 16 + st_row16_addr = (tmem_row + 16) << 16 + state_col_id = tmem_col + cfg.tmem_state_acc_offset + packed_col_id = tmem_col + cfg.tmem_state_inp_offset + state_k_col_id = tmem_col + cfg.tmem_state_k_acc_offset + y_inp_col_id = tmem_col + cfg.tmem_y_inp_offset + u_acc_addr = row_addr + tmem_col + cfg.tmem_u_acc_offset + u_inp_addr = st_row_addr + tmem_col + cfg.tmem_u_inp_offset + q_state_col_base = tmem_col + cfg.tmem_q_state_acc_offset + ov_swz_off0 = ( + (value_dim_base + ov_col_coord) // 64 * (cfg.b_t * 64) + + ov_token_coord * 64 + + swizzle_xor_128b(ov_token_coord, (value_dim_base + ov_col_coord) % 64, elem_bytes=2) + ) + ov_swz_off = ( + (value_dim_base + 16 + ov_col_coord) // 64 * (cfg.b_t * 64) + + ov_token_coord * 64 + + swizzle_xor_128b(ov_token_coord, (value_dim_base + 16 + ov_col_coord) % 64, elem_bytes=2) + ) + checkpoint_swz_off0 = (value_dim_base + ov_col_coord) // 64 * (cfg.d_k * 64) + checkpoint_swz_col0 = (value_dim_base + ov_col_coord) % 64 + checkpoint_swz_off = (value_dim_base + 16 + ov_col_coord) // 64 * (cfg.d_k * 64) + checkpoint_swz_col = (value_dim_base + 16 + ov_col_coord) % 64 state_k_acc_index = PipelineState.start(phase=0) - update_acc_index = PipelineState.start(phase=0) + u_acc_index = PipelineState.start(phase=0) o_acc_index = PipelineState.start(phase=0) - kr_index = PipelineState.start(phase=0) # CG1's per-chunk mb_k_restore_done wait slot - raw_index = PipelineState.start(phase=0) # raw-ring slot for the sV/sBeta reads + inputs_done arrives - gbase = cutlass.Int32(0) + state_upd_index = PipelineState.start(phase=0) + raw_index = PipelineState.start(phase=0) + raw_bar_index = PipelineState.start(phase=0) # even-depth ready/beta-ring slot (decoupled from the data ring) + cum_chunk_base = cutlass.Int32(0) 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, cu_seqlens, mWorkItems - ) + batch_idx, head_idx, batch_start, batch_end, seqlen_b, num_chunks_b, wstart, wend, cstart, cend = decode_work_item(cfg, tile_idx, mWorkItems) head_o = head_idx - sk_nt = wend - cstart - - if sk_nt > 0: - # ---- first chunk: seed state TMEM from mS_init (else zeros); - # split-K warmup items (cstart > 0) rebuild their state from zero - seed_from_s0 = cstart == 0 - 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_value = cutlass.Float32(0.0) - if cutlass.const_expr(mS_init is not None): - state_value = mS_init[batch_idx, head_o, key_dim, value_dim].to(cutlass.Float32) - if cutlass.const_expr(cfg.split_k): - state_value = state_value if seed_from_s0 else cutlass.Float32(0.0) - state_block[col] = state_value - - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_offset + key_block_start), cutlass.Float32), - state_block[0:32], - ) + num_chunks_tile = wend - cstart - nvvm.tcgen05_wait("store") - sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) - sBeta_ptr = sBeta_raw.data_ptr() + raw_index.idx * cfg.b_t - - row_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 - state_col_id = tmem_col + cfg.tmem_state_offset - # ---- state -> packed b16 A operand (TMEM roundtrip) ---------------- - state_blocks = [] - for sub in cutlass.range_constexpr(cfg.d_k // 16): - state_blocks.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) - nvvm.tcgen05_wait("load") + if num_chunks_tile > 0: + # ---- first chunk: seed state TMEM from mState_init ---------- + seed_from_initial_state = cstart == 0 + if cutlass.const_expr(mState_init is not None): + if seed_from_initial_state: + 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) - packed_col_id = tmem_col + cfg.tmem_state_inp_offset - for sub in cutlass.range_constexpr(cfg.d_k // 16): - packed_state = cutlass.Array(cutlass.Int32, 8, alignment=16) - for packed_col in cutlass.range_constexpr(8): - source_pair = packed_col ^ 4 - packed_state[packed_col] = fp32_to_fp16(state_blocks[sub][2 * source_pair], state_blocks[sub][2 * source_pair + 1], dtype=cfg.io_dtype) - nvvm.tcgen05_st( - "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), - packed_state[0:8], - ) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + state_block[0:32], + ) + else: + 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): + state_block[col] = cutlass.Float32(0.0) - nvvm.tcgen05_wait("store") - bars.mb_state_inp_ready.arrive() + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + state_block[0:32], + ) + if cutlass.const_expr(mState_init is not None): + nvvm.tcgen05_wait("store") + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sBeta_ptr = sBeta_raw.data_ptr() + raw_bar_index.idx * cfg.b_t + + # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + if cutlass.const_expr(mState_init is not None): + state_vecs = [] + for sub in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) + + for sub in cutlass.range_constexpr(cfg.d_k // 16): + packed_state = cutlass.Array(cutlass.Int32, 8, alignment=16) + for packed_col in cutlass.range_constexpr(8): + source_pair = packed_col ^ 4 + packed_state[packed_col] = fp32_to_fp16(state_vecs[sub][2 * source_pair], state_vecs[sub][2 * source_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), + packed_state[0:8], + ) + nvvm.tcgen05_wait("store") + bars.mb_state_inp_ready.arrive() + if cutlass.const_expr(cfg.enable_checkpoints): + bars.mb_state_acc_read_done.arrive() - # ---- rhs staging: rhs input = beta * (v - state*k) ----------------- - bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset - input_col_id = tmem_col + cfg.tmem_rhs_inp_offset - value_dim_base = tmem_sp * cfg.threads_per_warp - - # ---- read back state*k acc + raw v fragments ----------------------- - row_id0 = tmem_row + value_dim_base - state_k0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - - row_id1 = row_id0 + 16 - state_k1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) - - raw_v_regs0 = nvvm.ldmatrix( - sV_ptr - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + # ---- 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( + sV_ptr + ov_swz_off0, 4, nvvm.MMALayout.COL, ) - raw_v_regs1 = nvvm.ldmatrix( - sV_ptr - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + raw_v_frag1 = nvvm.ldmatrix( + sV_ptr + ov_swz_off, 4, nvvm.MMALayout.COL, ) - nvvm.tcgen05_wait("load") + bars.mb_beta_ready[raw_bar_index.idx].wait(raw_bar_index.phase) + if cutlass.const_expr(mState_init is not None): + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + + state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + state_k_col_id, cutlass.Float32), num=2) + state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + state_k_col_id, cutlass.Float32), num=2) - packed_rhs0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + beta_pack = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - packed_col = (reg_idx // 2) * 4 + (lane & 3) - source_pair = packed_col ^ 4 - token0 = source_pair * 2 - token1 = token0 + 1 + token0 = (((reg_idx // 2) * 4 + (lane & 3)) ^ 4) * 2 beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) - beta1 = (sBeta_ptr + token1).load().to(cutlass.Float32) + beta1 = (sBeta_ptr + token0 + 1).load().to(cutlass.Float32) + beta_pack[reg_idx] = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) + y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k0[frag_pair], state_k0[frag_pair + 1] - beta_pair = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - diff_pair = sub_f16x2( - raw_v_regs0[raw_matrix], - state_k_pair, - cfg.io_dtype, - ) - packed_rhs0[reg_idx] = mul_f16x2( - beta_pair, + if cutlass.const_expr(mState_init is not None): + state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + diff_pair = sub_f16x2( + raw_v_frag0[raw_matrix], + state_k_pair, + cfg.io_dtype, + ) + else: + diff_pair = raw_v_frag0[raw_matrix] + y_inp_pack0[reg_idx] = mul_f16x2( + beta_pack[reg_idx], diff_pair, cfg.io_dtype, ) - packed_rhs1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - packed_col = (reg_idx // 2) * 4 + (lane & 3) - source_pair = packed_col ^ 4 - token0 = source_pair * 2 - token1 = token0 + 1 - beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) - beta1 = (sBeta_ptr + token1).load().to(cutlass.Float32) raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k1[frag_pair], state_k1[frag_pair + 1] - beta_pair = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) - state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) - diff_pair = sub_f16x2( - raw_v_regs1[raw_matrix], - state_k_pair, - cfg.io_dtype, - ) - packed_rhs1[reg_idx] = mul_f16x2( - beta_pair, + if cutlass.const_expr(mState_init is not None): + state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + diff_pair = sub_f16x2( + raw_v_frag1[raw_matrix], + state_k_pair, + cfg.io_dtype, + ) + else: + diff_pair = raw_v_frag1[raw_matrix] + y_inp_pack1[reg_idx] = mul_f16x2( + beta_pack[reg_idx], diff_pair, cfg.io_dtype, ) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), packed_rhs0[0:4]) - - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), packed_rhs1[0:4]) - + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row_addr + y_inp_col_id, cutlass.Int8), y_inp_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row16_addr + y_inp_col_id, cutlass.Int8), y_inp_pack1[0:4]) nvvm.tcgen05_wait("store") - state_k_acc_index = advance(state_k_acc_index, 1) - bars.mb_inputs_done[raw_index.idx].arrive() - bars.mb_rhs_ready.arrive() - - # ---- update readback -> packed b16 A operand ----------------------- - bars.mb_update_acc_ready.wait(update_acc_index.phase) - update = nvvm.tcgen05_ld( + if cutlass.const_expr(mState_init is not None): + state_k_acc_index = advance(state_k_acc_index, 1) + bars.mb_v_done[raw_index.idx].arrive() + bars.mb_beta_done[raw_bar_index.idx].arrive() + bars.mb_y_inp_ready.arrive() + + # ---- U repack: u_acc TMEM -> packed b16 U input TMEM ---------------- + bars.mb_u_acc_ready.wait(u_acc_index.phase) + u_vals = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((tmem_row + tmem_sp * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_update_acc_offset), cutlass.Float32), + nvvm.make_tmem_ptr(u_acc_addr, cutlass.Float32), num=cfg.b_t, ) - nvvm.tcgen05_wait("load") - packed_update = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): source_pair = packed_col ^ 4 token0 = source_pair * 2 token1 = token0 + 1 - packed_update[packed_col] = fp32_to_fp16(update[token0], update[token1], dtype=cfg.io_dtype) + u_inp_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_update_inp_offset), cutlass.Int8), - packed_update[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr(u_inp_addr, cutlass.Int8), + u_inp_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") - update_acc_index = advance(update_acc_index, 1) - bars.mb_update_ready.arrive() + u_acc_index = advance(u_acc_index, 1) + bars.mb_u_inp_ready.arrive() + if cutlass.const_expr(cfg.enable_checkpoints): + bars.mb_state_acc_done.wait(state_upd_index.phase) + state_upd_index = advance(state_upd_index, 1) - bars.mb_k_restore_done[kr_index.idx].wait(kr_index.phase) - kr_index = advance(kr_index, cfg.smem_decay_stages) raw_index = advance(raw_index, cfg.smem_raw_stages) + raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) - # the first chunk is peeled above so this steady-state loop always - # drains the prior chunk's output - for li in cutlass.range(1, sk_nt, 1, unroll=1): - chunk_idx = cstart + li - gc = gbase + li + if cutlass.const_expr(cfg.enable_checkpoints): + cg1_checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + cg1_checkpoint_mod = (cstart + cutlass.Int32(1)) % cg1_checkpoint_chunks + for local_chunk_idx in cutlass.range(1, num_chunks_tile, 1, unroll=1): + chunk_idx = cstart + local_chunk_idx + cum_chunk = cum_chunk_base + local_chunk_idx sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) - sBeta_ptr = sBeta_raw.data_ptr() + raw_index.idx * cfg.b_t + sBeta_ptr = sBeta_raw.data_ptr() + raw_bar_index.idx * cfg.b_t prev_output_chunk = chunk_idx - cutlass.Int32(1) - prev_og = gc - cutlass.Int32(1) - prev_o_stage = prev_og % cfg.smem_o_stages - prev_q_state_acc_stage = prev_og % cfg.tmem_q_state_acc_stages + prev_cum_chunk = cum_chunk - cutlass.Int32(1) + prev_o_stage = prev_cum_chunk % cfg.smem_o_stages + prev_q_state_acc_stage = prev_cum_chunk % cfg.tmem_q_state_acc_stages prev_o_stage_base = prev_o_stage * (cfg.b_t * cfg.d_v) - # H entry gate: the state read by this restage entered chunk_idx, - # i.e. the state after chunk_idx * b_t tokens (strictly before the - # sequence end -- the end state is only final_state); split-K - # warmup reconstructions below wstart belong to the previous item - do_h = False + do_checkpoint = False if cutlass.const_expr(cfg.enable_checkpoints): - do_h = (chunk_idx * cutlass.Int32(cfg.b_t)) % checkpoint_every_n_tokens == 0 - if cutlass.const_expr(cfg.split_k): - do_h = do_h and chunk_idx >= wstart - if do_h: - # sH is free once the epilogue's previous TMA store retired - bars.mb_h_tmastg_done.wait(h_done_index.phase) - h_done_index = advance(h_done_index, 1) - row_addr = (tmem_row + tmem_sp * cfg.threads_per_warp) << 16 - state_col_id = tmem_col + cfg.tmem_state_offset - # ---- state -> packed b16 A operand (TMEM roundtrip) ---------------- - state_blocks = [] + do_checkpoint = cg1_checkpoint_mod == 0 + cg1_checkpoint_mod = cg1_checkpoint_mod + cutlass.Int32(1) + cg1_checkpoint_mod = cutlass.Int32(0) if cg1_checkpoint_mod == cg1_checkpoint_chunks else cg1_checkpoint_mod + do_checkpoint = do_checkpoint and chunk_idx >= wstart + + # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + if cutlass.const_expr(not cfg.enable_checkpoints): + bars.mb_state_acc_done.wait(state_upd_index.phase) + state_upd_index = advance(state_upd_index, 1) + state_vecs = [] for sub in cutlass.range_constexpr(cfg.d_k // 16): - state_blocks.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) - bars.mb_o_tmastg_done[prev_o_stage].wait(((prev_og // cfg.smem_o_stages) + 1) % 2) - nvvm.tcgen05_wait("load") + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) - packed_col_id = tmem_col + cfg.tmem_state_inp_offset for sub in cutlass.range_constexpr(cfg.d_k // 16): packed_state = cutlass.Array(cutlass.Int32, 8, alignment=16) for packed_col in cutlass.range_constexpr(8): source_pair = packed_col ^ 4 - packed_state[packed_col] = fp32_to_fp16(state_blocks[sub][2 * source_pair], state_blocks[sub][2 * source_pair + 1], dtype=cfg.io_dtype) + packed_state[packed_col] = fp32_to_fp16(state_vecs[sub][2 * source_pair], state_vecs[sub][2 * source_pair + 1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), packed_state[0:8], ) - if cutlass.const_expr(cfg.enable_checkpoints): - # stage this sub's state to sH TRANSPOSED (KV: k rows, v - # contiguous in 64-v slabs, swizzled — the GDN H layout); - # each thread scatters its 16 k values down one v column - if do_h: - h_col = value_dim % cutlass.Int32(64) - h_seg = (value_dim // cutlass.Int32(64)) * (cfg.d_k * cutlass.Int32(64)) - k_base = sub * 16 - for j in cutlass.range_constexpr(16): - hv = state_blocks[sub][j].to(cfg.io_dtype) - (sH_ptr + h_seg + cutlass.Int32((k_base + j) * 64) + swizzle_xor_128b(cutlass.Int32(k_base + j), h_col, elem_bytes=2)).store(hv) - nvvm.tcgen05_wait("store") bars.mb_state_inp_ready.arrive() + + # ---- checkpoint store ----------------------------------------------- if cutlass.const_expr(cfg.enable_checkpoints): - if do_h: + if do_checkpoint: + 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 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, + ) + nvvm.tcgen05_wait("load") + bars.mb_state_acc_read_done.arrive() nvvm.fence_proxy("async.shared", space="cta") - bars.mb_h_tmastg_ready.arrive() + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + else: + bars.mb_state_acc_read_done.arrive() + bars.mb_o_acc_ready.wait(o_acc_index.phase) o_acc_index = advance(o_acc_index, 1) - projection_col_id = tmem_col + cfg.tmem_q_state_acc_offset + prev_q_state_acc_stage * cfg.b_t - value_dim_base = tmem_sp * cfg.threads_per_warp + projection_col_id = q_state_col_base + prev_q_state_acc_stage * cfg.b_t + loaded_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + projection_col_id, cutlass.Float32), num=2) + loaded_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + projection_col_id, cutlass.Float32), num=2) - row_id0 = tmem_row + value_dim_base - row_id1 = row_id0 + 16 - loaded0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - loaded1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) - - # ---- output drain: q_state_acc -> scaled b16 -> SMEM stmatrix ------- - stsm_regs0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - stsm_regs1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + # ---- output drain: O acc TMEM -> scaled b16 SMEM -------------------- + stsm_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + stsm_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - scaled0_0, scaled0_1 = fmul2(loaded0[2 * reg_idx], loaded0[2 * reg_idx + 1], scale, scale) - scaled1_0, scaled1_1 = fmul2(loaded1[2 * reg_idx], loaded1[2 * reg_idx + 1], scale, scale) - stsm_regs0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) - stsm_regs1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + scaled0_0, scaled0_1 = fmul2(loaded_vec0[2 * reg_idx], loaded_vec0[2 * reg_idx + 1], scale, scale) + scaled1_0, scaled1_1 = fmul2(loaded_vec1[2 * reg_idx], loaded_vec1[2 * reg_idx + 1], scale, scale) + stsm_pack0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) + stsm_pack1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + bars.mb_o_tmastg_done[prev_o_stage].wait(((prev_cum_chunk // cfg.smem_o_stages) + 1) % 2) nvvm.stmatrix( - sO_ptr - + prev_o_stage_base - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), - stsm_regs0.data_ptr().load(count=4, alignment=4), + sO_ptr + prev_o_stage_base + ov_swz_off0, + stsm_pack0.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( - sO_ptr - + prev_o_stage_base - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), - stsm_regs1.data_ptr().load(count=4, alignment=4), + sO_ptr + prev_o_stage_base + ov_swz_off, + stsm_pack1.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) - # release only after the stmatrix pair: the STSM->F2FP->FMUL2->LDTM - # register chain pins the TMEM reads complete without a wait("load") bars.mb_o_acc_done[prev_q_state_acc_stage].arrive() nvvm.fence_proxy("async.shared", space="cta") bars.mb_o_tmastg_ready[prev_o_stage].arrive() - bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) - # ---- rhs staging: rhs input = beta * (v - state*k) ----------------- - projection_col_id = tmem_col + cfg.tmem_state_k_acc_offset - input_col_id = tmem_col + cfg.tmem_rhs_inp_offset - value_dim_base = tmem_sp * cfg.threads_per_warp - - # ---- read back state*k acc + raw v fragments ----------------------- - row_id0 = tmem_row + value_dim_base - state_k0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - - row_id1 = row_id0 + 16 - state_k1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) - - raw_v_regs0 = nvvm.ldmatrix( - sV_ptr - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), + # ---- 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( + sV_ptr + ov_swz_off0, 4, nvvm.MMALayout.COL, ) - raw_v_regs1 = nvvm.ldmatrix( - sV_ptr - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), + raw_v_frag1 = nvvm.ldmatrix( + sV_ptr + ov_swz_off, 4, nvvm.MMALayout.COL, ) - nvvm.tcgen05_wait("load") + bars.mb_beta_ready[raw_bar_index.idx].wait(raw_bar_index.phase) + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + + state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + state_k_col_id, cutlass.Float32), num=2) + state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + state_k_col_id, cutlass.Float32), num=2) - packed_rhs0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + beta_pack = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - packed_col = (reg_idx // 2) * 4 + (lane & 3) - source_pair = packed_col ^ 4 - token0 = source_pair * 2 - token1 = token0 + 1 + token0 = (((reg_idx // 2) * 4 + (lane & 3)) ^ 4) * 2 beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) - beta1 = (sBeta_ptr + token1).load().to(cutlass.Float32) + beta1 = (sBeta_ptr + token0 + 1).load().to(cutlass.Float32) + beta_pack[reg_idx] = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) + y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k0[frag_pair], state_k0[frag_pair + 1] - beta_pair = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) + state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) diff_pair = sub_f16x2( - raw_v_regs0[raw_matrix], + raw_v_frag0[raw_matrix], state_k_pair, cfg.io_dtype, ) - packed_rhs0[reg_idx] = mul_f16x2( - beta_pair, + y_inp_pack0[reg_idx] = mul_f16x2( + beta_pack[reg_idx], diff_pair, cfg.io_dtype, ) - packed_rhs1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - packed_col = (reg_idx // 2) * 4 + (lane & 3) - source_pair = packed_col ^ 4 - token0 = source_pair * 2 - token1 = token0 + 1 - beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) - beta1 = (sBeta_ptr + token1).load().to(cutlass.Float32) raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) frag_pair = (reg_idx ^ 2) * 2 - state_k_val0, state_k_val1 = state_k1[frag_pair], state_k1[frag_pair + 1] - beta_pair = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) + state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) diff_pair = sub_f16x2( - raw_v_regs1[raw_matrix], + raw_v_frag1[raw_matrix], state_k_pair, cfg.io_dtype, ) - packed_rhs1[reg_idx] = mul_f16x2( - beta_pair, + y_inp_pack1[reg_idx] = mul_f16x2( + beta_pack[reg_idx], diff_pair, cfg.io_dtype, ) - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row << 16) + input_col_id, cutlass.Int8), packed_rhs0[0:4]) - - nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr((tmem_row + 16 << 16) + input_col_id, cutlass.Int8), packed_rhs1[0:4]) - + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row_addr + y_inp_col_id, cutlass.Int8), y_inp_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row16_addr + y_inp_col_id, cutlass.Int8), y_inp_pack1[0:4]) nvvm.tcgen05_wait("store") state_k_acc_index = advance(state_k_acc_index, 1) - bars.mb_inputs_done[raw_index.idx].arrive() - bars.mb_rhs_ready.arrive() + bars.mb_v_done[raw_index.idx].arrive() + bars.mb_beta_done[raw_bar_index.idx].arrive() + bars.mb_y_inp_ready.arrive() - # ---- update readback -> packed b16 A operand ----------------------- - bars.mb_update_acc_ready.wait(update_acc_index.phase) - update = nvvm.tcgen05_ld( + # ---- U repack: u_acc TMEM -> packed b16 U input TMEM ---------------- + bars.mb_u_acc_ready.wait(u_acc_index.phase) + u_vals = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((tmem_row + tmem_sp * cfg.threads_per_warp << 16) + (tmem_col + cfg.tmem_update_acc_offset), cutlass.Float32), + nvvm.make_tmem_ptr(u_acc_addr, cutlass.Float32), num=cfg.b_t, ) - nvvm.tcgen05_wait("load") - packed_update = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): source_pair = packed_col ^ 4 token0 = source_pair * 2 token1 = token0 + 1 - packed_update[packed_col] = fp32_to_fp16(update[token0], update[token1], dtype=cfg.io_dtype) + u_inp_pack[packed_col] = fp32_to_fp16(u_vals[token0], u_vals[token1], dtype=cfg.io_dtype) nvvm.tcgen05_st( "32x32b", - nvvm.make_tmem_ptr((tmem_row << 16) + (tmem_col + cfg.tmem_update_inp_offset), cutlass.Int8), - packed_update[0 : (cfg.b_t // 2)], + nvvm.make_tmem_ptr(u_inp_addr, cutlass.Int8), + u_inp_pack[0 : (cfg.b_t // 2)], ) nvvm.tcgen05_wait("store") - update_acc_index = advance(update_acc_index, 1) - bars.mb_update_ready.arrive() + u_acc_index = advance(u_acc_index, 1) + bars.mb_u_inp_ready.arrive() - bars.mb_k_restore_done[kr_index.idx].wait(kr_index.phase) - kr_index = advance(kr_index, cfg.smem_decay_stages) + if cutlass.const_expr(cfg.enable_checkpoints): + bars.mb_state_acc_done.wait(state_upd_index.phase) + state_upd_index = advance(state_upd_index, 1) raw_index = advance(raw_index, cfg.smem_raw_stages) - - if sk_nt > 0: - og = gbase + sk_nt - cutlass.Int32(1) - final_o_stage = og % cfg.smem_o_stages - final_q_state_acc_stage = og % cfg.tmem_q_state_acc_stages + raw_bar_index = advance(raw_bar_index, cfg.smem_raw_bar_stages) + + if num_chunks_tile > 0: + if cutlass.const_expr(not cfg.enable_checkpoints): + bars.mb_state_acc_done.wait(state_upd_index.phase) + state_upd_index = advance(state_upd_index, 1) + last_cum_chunk = cum_chunk_base + num_chunks_tile - cutlass.Int32(1) + final_o_stage = last_cum_chunk % cfg.smem_o_stages + final_q_state_acc_stage = last_cum_chunk % cfg.tmem_q_state_acc_stages final_o_stage_base = final_o_stage * (cfg.b_t * cfg.d_v) - bars.mb_o_tmastg_done[final_o_stage].wait(((og // cfg.smem_o_stages) + 1) % 2) bars.mb_o_acc_ready.wait(o_acc_index.phase) o_acc_index = advance(o_acc_index, 1) - projection_col_id = tmem_col + cfg.tmem_q_state_acc_offset + final_q_state_acc_stage * cfg.b_t - value_dim_base = tmem_sp * cfg.threads_per_warp + projection_col_id = q_state_col_base + final_q_state_acc_stage * cfg.b_t - row_id0 = tmem_row + value_dim_base - row_id1 = row_id0 + 16 - loaded0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id0 << 16) + projection_col_id, cutlass.Float32), num=2) - loaded1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr((row_id1 << 16) + projection_col_id, cutlass.Float32), num=2) + loaded_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + projection_col_id, cutlass.Float32), num=2) + loaded_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + projection_col_id, cutlass.Float32), num=2) - # ---- output drain: q_state_acc -> scaled b16 -> SMEM stmatrix ------- - stsm_regs0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) - stsm_regs1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + # ---- output drain: O acc TMEM -> scaled b16 SMEM -------------------- + stsm_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + stsm_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(4): - scaled0_0, scaled0_1 = fmul2(loaded0[2 * reg_idx], loaded0[2 * reg_idx + 1], scale, scale) - scaled1_0, scaled1_1 = fmul2(loaded1[2 * reg_idx], loaded1[2 * reg_idx + 1], scale, scale) - stsm_regs0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) - stsm_regs1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + scaled0_0, scaled0_1 = fmul2(loaded_vec0[2 * reg_idx], loaded_vec0[2 * reg_idx + 1], scale, scale) + scaled1_0, scaled1_1 = fmul2(loaded_vec1[2 * reg_idx], loaded_vec1[2 * reg_idx + 1], scale, scale) + stsm_pack0[reg_idx] = fp32_to_fp16(scaled0_0, scaled0_1, dtype=mO.element_type) + stsm_pack1[reg_idx] = fp32_to_fp16(scaled1_0, scaled1_1, dtype=mO.element_type) + bars.mb_o_tmastg_done[final_o_stage].wait(((last_cum_chunk // cfg.smem_o_stages) + 1) % 2) nvvm.stmatrix( - sO_ptr - + final_o_stage_base - + (value_dim_base + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + ov_col) % 64, elem_bytes=2), - stsm_regs0.data_ptr().load(count=4, alignment=4), + sO_ptr + final_o_stage_base + ov_swz_off0, + stsm_pack0.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) nvvm.stmatrix( - sO_ptr - + final_o_stage_base - + (value_dim_base + 16 + ov_col) // 64 * (cfg.b_t * 64) - + ov_tok * 64 - + swizzle_xor_128b(ov_tok, (value_dim_base + 16 + ov_col) % 64, elem_bytes=2), - stsm_regs1.data_ptr().load(count=4, alignment=4), + sO_ptr + final_o_stage_base + ov_swz_off, + stsm_pack1.data_ptr().load(count=4, alignment=4), nvvm.MMALayout.COL, shape=nvvm.StoreShape.M8N8, ) - # release only after the stmatrix pair: the STSM->F2FP->FMUL2->LDTM - # register chain pins the TMEM reads complete without a wait("load") bars.mb_o_acc_done[final_q_state_acc_stage].arrive() nvvm.fence_proxy("async.shared", space="cta") bars.mb_o_tmastg_ready[final_o_stage].arrive() - # split-K: only the item owning the sequence's last chunk holds the - # true end-of-sequence state (legacy tiles always do: wend == nc) owns_final = wend == num_chunks_b - # ---- final-state drain: final_state acc -> GMEM -------------------- - if cutlass.const_expr(mS_out is not None): + + # ---- final-state drain: state acc TMEM -> GMEM --------------------------- + if cutlass.const_expr(mState_out is not None): if seqlen_b > 0: if owns_final: for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): loaded = nvvm.tcgen05_ld( "32x32b", - nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_offset + key_block_start), cutlass.Float32), + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), num=32, ) - nvvm.tcgen05_wait("load") for col in cutlass.range_constexpr(32): key_dim = key_block_start + col - mS_out[batch_idx, head_o, key_dim, value_dim] = loaded[col].to(mS_out.element_type) + mState_out[batch_idx, head_o, key_dim, value_dim] = loaded[col].to(mState_out.element_type) else: - # zero-length sequence: the state passes through untouched - # (S0 when seeded, zeros otherwise); pure GMEM, no TMEM 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(mS_init is not None): - mS_out[batch_idx, head_o, key_dim, value_dim] = mS_init[batch_idx, head_o, key_dim, value_dim] + 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] else: - mS_out[batch_idx, head_o, key_dim, value_dim] = cutlass.Float32(0.0).to(mS_out.element_type) - bars.mb_final_state_stored.arrive() - gbase += sk_nt - tile_idx, sched_state = _sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + mState_out[batch_idx, head_o, key_dim, value_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) + + bars.mb_tmem_done[0].arrive() @cute.jit -def _host( +def host( cfg: cutlass.Constexpr, q: cute.Tensor, k: cute.Tensor, @@ -1981,15 +1990,11 @@ def _host( stream, ) -> None: num_sequences = cu_seqlens.shape[0] - 1 - ho = raw_gate.shape[1] - # ---- persistent launch: the grid only needs to cover the tiles ------ - total_tiles = num_sequences * ho - # CUDA-graph-stable launch: fixed SM-count grid grid_shape = (cfg.max_active_clusters, 1, 1) - _kernel( + kernel( cfg, tensormap_workspace, - cutlass.Int32(num_sequences * ho), + cutlass.Int32(num_sequences), q, k, v, @@ -2004,7 +2009,6 @@ def _host( work_items, work_count, sched_ctr, - total_tiles, scale, checkpoint_every_n_tokens, ).launch( @@ -2016,7 +2020,7 @@ def _host( @cute.kernel -def _kernel( +def kernel( cfg: cutlass.Constexpr, tensormap_workspace: cute.Tensor, n_desc: cutlass.Int32, @@ -2028,23 +2032,16 @@ def _kernel( mDt_bias: cute.Tensor | None, mBeta: cute.Tensor, cu_seqlens: cute.Tensor, - mS_init: cute.Tensor | None, + mState_init: cute.Tensor | None, mO: cute.Tensor, - mS_out: cute.Tensor | None, - mWorkItems: cute.Tensor | None, - mCount: cute.Tensor | None, + mState_out: cute.Tensor | None, + mWorkItems: cute.Tensor, + mCount: cute.Tensor, mSched: cute.Tensor | None, - total_tiles: cutlass.Int32, scale: cutlass.Float32, checkpoint_every_n_tokens: cutlass.Int32, ) -> None: - """BT=16 KDA forward kernel (persistent). - - Grid: `(min(tiles, SM count), 1, 1)`. Every warp role runs a - tile-scheduler loop over the tiles — one packed sequence/head each, or - one split-K work item — iterating its chunks in order. Source heads - follow repeat_interleave: head_x = head_idx // X_RATIO. - """ + """BT=16 KDA forward kernel (persistent, tile-scheduled).""" tidx, _, _ = cute.arch.thread_idx() bidx = cute.arch.block_idx()[0] @@ -2052,8 +2049,7 @@ def _kernel( warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) lane = tidx % cfg.threads_per_warp - if cutlass.const_expr(cfg.split_k): - total_tiles = mCount[0] + total_tiles = mCount[0] if cutlass.const_expr(cfg.dyn_sched): assert mSched is not None and mSched.element_type == cutlass.Int32 assert mQ.element_type == cfg.io_dtype and mK.element_type == cfg.io_dtype and mV.element_type == cfg.io_dtype @@ -2062,16 +2058,15 @@ def _kernel( assert mBeta.element_type == beta_expected assert cu_seqlens.element_type in (cutlass.Int32, cutlass.Int64) if cutlass.const_expr(cfg.use_initial_state): - assert mS_init is not None and mS_init.element_type in (cutlass.BFloat16, cutlass.Float32) + assert mState_init is not None and mState_init.element_type in (cutlass.BFloat16, cutlass.Float32) else: - assert mS_init is None, "mS_init must be None if use_initial_state is False" + assert mState_init is None, "mState_init must be None if use_initial_state is False" if cutlass.const_expr(cfg.store_final_state): - assert mS_out is not None and mS_out.element_type in (cutlass.BFloat16, cutlass.Float32) + assert mState_out is not None and mState_out.element_type in (cutlass.BFloat16, cutlass.Float32) else: - assert mS_out is None, "mS_out must be None if store_final_state is False" - if cutlass.const_expr(mS_init is not None and mS_out is not None): - assert mS_init.element_type == mS_out.element_type - # per-(batch, head) TMA-descriptor arrays: [q, k, v, gate, o] + assert mState_out is None, "mState_out must be None if store_final_state is False" + if cutlass.const_expr(mState_init is not None and mState_out is not None): + assert mState_init.element_type == mState_out.element_type desc_base_words = tensormap_workspace.iterator.raw_ptr() arr_words = n_desc * cutlass.Int32(TENSOR_MAP_QWORDS) desc_q_base = desc_base_words @@ -2079,19 +2074,17 @@ def _kernel( desc_v_base = desc_base_words + cutlass.Int32(2) * arr_words desc_gate_base = desc_base_words + cutlass.Int32(3) * arr_words desc_o_base = desc_base_words + cutlass.Int32(4) * arr_words - desc_h_base = desc_base_words + cutlass.Int32(5) * arr_words + desc_checkpoint_base = desc_base_words + cutlass.Int32(5) * arr_words # Buffers are declaration-ordered and intentionally non-aliased. SMEM = cutlass.AddressSpace.smem bars = make_kda_bars(cfg) - # The hand-written K-box-major SW128 mapping is normalized to phase 0, - # so both tcgen05 and ldmatrix can share 1KB-aligned operand buffers. - tmem_hold = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) + tmem_base_slot = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) sK_decay_raw = cutlass.Array(cfg.io_dtype, cfg.k_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sQ_decay_raw = cutlass.Array(cfg.io_dtype, cfg.q_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sK_restore_raw = cutlass.Array(cfg.io_dtype, cfg.k_restore_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) - sPairwise_raw = cutlass.Array(cfg.io_dtype, cfg.pairwise_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sIntermediate_raw = cutlass.Array(cfg.io_dtype, cfg.intermediate_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sQ_raw = cutlass.Array(mQ.element_type, cfg.q_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sK_raw = cutlass.Array(mK.element_type, cfg.k_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) sV_raw = cutlass.Array(mV.element_type, cfg.v_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) @@ -2102,17 +2095,14 @@ def _kernel( mO.element_type, cfg.o_cosize, space=SMEM, - # The scalar CG1 store computes W128 offsets relative to this buffer. - # Align to the full s128b period so absolute SMEM address bits do not - # add a hidden phase to the TMA store-side swizzle. alignment=cfg.buffer_align_bytes, ) sBeta_raw = cutlass.Array(cutlass.Float32, cfg.beta_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) - sH_raw = ( - cutlass.Array(cfg.io_dtype, cfg.d_k * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) if cutlass.const_expr(cfg.enable_checkpoints) else sO_raw + sCheckpoint_raw = ( + cutlass.Array(cfg.io_dtype, cfg.smem_checkpoint_stages * cfg.d_k * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) + if cutlass.const_expr(cfg.enable_checkpoints) + else sO_raw ) - # K-box-major SW128 staging: 16B leading offset, 1KB stride; the decay - # stores apply a row-group key xor so tcgen05 B reads logical [DK, BT]. sK_decay = SmemTile( base=sK_decay_raw, elems_per_stage=(cfg.d_k * cfg.b_t), @@ -2145,56 +2135,61 @@ def _kernel( stride_byte_offset=(8 * 16 * 2), layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, ) - sPairwise = SmemTile( - base=sPairwise_raw, + sIntermediate = SmemTile( + base=sIntermediate_raw, elems_per_stage=(2 * cfg.b_t * cfg.b_t), - stages=cfg.smem_pairwise_stages, + stages=cfg.smem_intermediate_stages, leading_byte_offset=16, stride_byte_offset=(8 * cfg.b_t * 2), layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, ) - tma_tx_bytes = cutlass.const_expr( - cfg.d_k * cfg.b_t * mQ.element_type.width // 8 - + cfg.d_k * cfg.b_t * mK.element_type.width // 8 - + cfg.d_v * cfg.b_t * mV.element_type.width // 8 - + cfg.d_k * cfg.b_t * mGate.element_type.width // 8 - ) + elect_one = nvvm.elect_sync() if warp_idx == cfg.tma_warp_id: - if nvvm.elect_sync(): - bars.mb_tma_done.init() + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_raw_bar_stages): + bars.mb_q_ready[stage].init() + bars.mb_k_ready[stage].init() + bars.mb_v_ready[stage].init() + bars.mb_gate_ready[stage].init() + bars.mb_beta_ready[stage].init() + bars.mb_beta_done[stage].init() for stage in cutlass.range_constexpr(cfg.smem_raw_stages): - bars.mb_inputs_ready[stage].init() - bars.mb_inputs_done[stage].init() + bars.mb_q_done[stage].init() + bars.mb_k_done[stage].init() + bars.mb_v_done[stage].init() + bars.mb_gate_done[stage].init() elif warp_idx == cfg.tcgen05_mma_warp_id: - if nvvm.elect_sync(): + if elect_one: bars.mb_o_acc_ready.init() for stage in cutlass.range_constexpr(cfg.tmem_q_state_acc_stages): bars.mb_o_acc_done[stage].init() bars.mb_state_k_acc_ready.init() - bars.mb_update_acc_ready.init() + bars.mb_u_acc_ready.init() + bars.mb_state_acc_done.init() bars.mb_state_inp_ready.init() for stage in cutlass.range_constexpr(cfg.smem_state_scale_diag_stages): bars.mb_state_scale_diag_done[stage].init() for stage in cutlass.range_constexpr(cfg.smem_decay_stages): - bars.mb_kk_qk_super_mma_done[stage].init() - bars.mb_kk_qk_mma_done[stage].init() + bars.mb_decay_tcgen05_done[stage].init() + bars.mb_decay_super_done[stage].init() bars.mb_k_restore_done[stage].init() - bars.mb_rhs_ready.init() - bars.mb_update_ready.init() - bars.mb_final_state_stored.init() + bars.mb_y_inp_ready.init() + bars.mb_u_inp_ready.init() + bars.mb_tmem_done[0].init() elif warp_idx == cfg.super_mma_warp_id: - if nvvm.elect_sync(): - for stage in cutlass.range_constexpr(cfg.smem_pairwise_stages): + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_intermediate_stages): + bars.mb_t_inv_ready[stage].init() bars.mb_a_ready[stage].init() - bars.mb_qk_acc_ready[stage].init() + bars.mb_t_inv_done[stage].init() bars.mb_a_done[stage].init() for stage in cutlass.range_constexpr(cfg.qk_scale_ready_stages): bars.mb_qk_scale_ready[stage].init() for stage in cutlass.range_constexpr(cfg.smem_decay_stages): - bars.mb_k_decay_cg0_ready[stage].init() + bars.mb_k_decay_inv_cg0_ready[stage].init() elif warp_idx == cfg.epilogue_warp_id: - if nvvm.elect_sync(): + if elect_one: for stage in cutlass.range_constexpr(cfg.smem_o_stages): bars.mb_o_tmastg_ready[stage].init() bars.mb_o_tmastg_done[stage].init() @@ -2202,39 +2197,17 @@ def _kernel( bars.mb_sched_ready[stage].init() bars.mb_sched_done[stage].init() if cutlass.const_expr(cfg.enable_checkpoints): - bars.mb_h_tmastg_ready.init() - bars.mb_h_tmastg_done.init() + for stage in cutlass.range_constexpr(cfg.smem_checkpoint_stages): + bars.mb_checkpoint_tmastg_ready[stage].init() + bars.mb_checkpoint_tmastg_done[stage].init() + bars.mb_state_acc_read_done.init() diag_zero = cfg.io_dtype(0.0) for diag_idx in cutlass.range(tidx, cfg.state_scale_diag_cosize, cfg.threads_per_cta, unroll=1): sState_scale_diag_raw[diag_idx] = diag_zero nvvm.fence_mbarrier_init() nvvm.barrier_cta_sync(0, thread_count=cfg.threads_per_cta) - if (warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]) or warp_idx == cfg.tcgen05_mma_warp_id: - if warp_idx == cfg.tcgen05_mma_warp_id: - nvvm.tcgen05_alloc(tmem_hold, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) - nvvm.barrier_cta_sync(cfg.nbar_tmem_lifecycle_id, thread_count=cfg.tmem_user_threads) - if warp_idx == cfg.tcgen05_mma_warp_id: - nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) - nvvm.barrier_cta_sync(cfg.nbar_tmem_lifecycle_id, thread_count=cfg.tmem_user_threads) - - # Actual SMEM/TMEM buffers for the BT=16 schedule: - # q/k/v : 16 x 128 each - # gate_log2 : 16 x 128 - # beta : 16 - # q/k inverse norm : 16 each, staged once per decay stage - # exp_g_last : 128, CG0-local and staged once per decay stage - # state-scale diag : 8 x 16 x 16 input dtype, zeroed once in the prologue - # q_decay/k_decay : tcgen05 SW128 operands shared with super-MMA - # k_restore : tcgen05 SW128 N-major final-state operand - # k_inv : 16 x 128 token-major for super-MMA RHS - # A inverse/QK : 16 x 16 each, plus transposed tcgen05 operands - # state : external/kernel ABI is VK `[DV, DK]`; reference - # math can view it as KV `[DK, DV]` by transposing. - # The TS A-staging path keeps VK in TMEM so state*k - # is `[DV, DK] @ [DK, BT] -> [DV, BT]` with M=128. - if warp_idx == cfg.tma_warp_id: - _tmaldg_warp( + tmaldg_warp( cfg, total_tiles, bidx, @@ -2243,14 +2216,11 @@ def _kernel( mWorkItems, mSched, sSched, - tma_tx_bytes, lane, - mBeta, sQ_raw, sK_raw, sV_raw, sGate_raw, - sBeta_raw, desc_q_base, desc_k_base, desc_v_base, @@ -2258,7 +2228,7 @@ def _kernel( bars, ) elif warp_idx == cfg.super_mma_warp_id: - _super_mma_warp( + super_mma_warp( cfg, total_tiles, bidx, @@ -2268,13 +2238,13 @@ def _kernel( sSched, lane, sK_inv_raw, - sPairwise_raw, + sIntermediate_raw, sBeta_raw, sK_decay_raw, bars, ) elif warp_idx == cfg.tcgen05_mma_warp_id: - _tcgen05_mma_warp( + tcgen05_mma_warp( cfg, total_tiles, bidx, @@ -2282,8 +2252,8 @@ def _kernel( cu_seqlens, mWorkItems, sSched, - tmem_hold, - sPairwise, + tmem_base_slot, + sIntermediate, sK_decay, sK_restore, sQ_decay, @@ -2291,7 +2261,7 @@ def _kernel( bars, ) elif warp_idx == cfg.epilogue_warp_id: - _epilogue_warp( + epilogue_warp( cfg, total_tiles, bidx, @@ -2303,16 +2273,16 @@ def _kernel( mO, sK_inv_raw, sO_raw, - sPairwise_raw, + sIntermediate_raw, sQ_decay_raw, - sH_raw, + sCheckpoint_raw, desc_o_base, - desc_h_base, + desc_checkpoint_base, checkpoint_every_n_tokens, bars, ) elif warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: - _compute0_warp_group( + compute0_warp_group( cfg, total_tiles, bidx, @@ -2327,9 +2297,10 @@ def _kernel( mDt_bias, sK_inv_raw, sGate_raw, + mBeta, + sBeta_raw, sK_raw, sQ_raw, - sV_raw, sK_decay_raw, sK_restore_raw, sQ_decay_raw, @@ -2337,7 +2308,7 @@ def _kernel( bars, ) elif warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]: - _compute1_warp_group( + compute1_warp_group( cfg, total_tiles, bidx, @@ -2346,15 +2317,15 @@ def _kernel( mWorkItems, sSched, lane, - tmem_hold, + tmem_base_slot, warp_idx, - mS_out, - mS_init, + mState_out, + mState_init, mO, sO_raw, sBeta_raw, sV_raw, - sH_raw, + sCheckpoint_raw, checkpoint_every_n_tokens, scale, bars, @@ -2366,8 +2337,7 @@ class KdaCfg: """Kernel cfg (fixed BT=16 schedule constants; derived TMEM column offsets and SMEM buffer cosizes are stamped by ``build_cfg``; per-stage sizes are inlined at the use sites). Passed ``cfg``-first (a ``cutlass.Constexpr``) - into ``_host`` / ``_kernel`` and every warp body, mirroring GDN's - ``GdnCfg``.""" + into ``host`` / ``kernel`` and every warp body.""" io_dtype: Type[cutlass.Numeric] state_dtype: Type[cutlass.Numeric] @@ -2383,9 +2353,6 @@ class KdaCfg: v_ratio: int n_heads_out: int max_active_clusters: int - # split-K: tiles come from a work-item table (see common/split_k.py); - # each item computes chunks [cstart, wend) and writes only [wstart, wend) - split_k: bool = False dyn_sched: bool = False sched_stages: int = CFG.SMEM_SCHED_STAGES @@ -2404,32 +2371,35 @@ class KdaCfg: cg0_group_count: int = 2 cg0_warps_per_group: int = 4 cg0_threads_per_group: int = 0 - nbar_cg0_group0_id: int = 1 # CG0 group g syncs on nbar id 1 + g + cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on nbar id 1 + g + cg0_tile_entry_barrier_id: int = 5 # CG0-wide (both groups) work-item entry sync tmem_user_threads: int = 0 - nbar_tmem_lifecycle_id: int = 3 + tmem_lifecycle_barrier_id: int = 3 num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 num_regs_other: int = CFG.NUM_REGS_OTHER - # --- SMEM / TMEM ring stage counts --- + # ---- SMEM / TMEM ring stage counts ------------------------------------------- smem_raw_stages: int = CFG.SMEM_RAW_STAGES + smem_raw_bar_stages: int = 0 # ready/beta-ring mbar depth: raw rounded up to even (CG0 ping-pong parity) smem_o_stages: int = CFG.SMEM_O_STAGES + smem_checkpoint_stages: int = 1 smem_decay_stages: int = CFG.SMEM_DECAY_STAGES - smem_pairwise_stages: int = CFG.SMEM_PAIRWISE_STAGES + smem_intermediate_stages: int = CFG.SMEM_INTERMEDIATE_STAGES smem_state_scale_diag_stages: int = CFG.SMEM_STATE_SCALE_DIAG_STAGES qk_scale_ready_stages: int = CFG.QK_SCALE_READY_STAGES tmem_q_state_acc_stages: int = CFG.TMEM_Q_STATE_ACC_STAGES - # --- TMEM column offsets (state doubles as the final_state acc) --- - tmem_state_offset: int = 0 + # ---- TMEM column offsets (state doubles as the final_state acc) -------------- + tmem_state_acc_offset: int = 0 tmem_state_inp_offset: int = 0 tmem_q_state_acc_offset: int = 0 tmem_state_k_acc_offset: int = 0 - tmem_update_acc_offset: int = 0 - tmem_rhs_inp_offset: int = 0 - tmem_update_inp_offset: int = 0 + tmem_u_acc_offset: int = 0 + tmem_y_inp_offset: int = 0 + tmem_u_inp_offset: int = 0 - # --- SMEM buffer cosizes --- + # ---- SMEM buffer cosizes ----------------------------------------------------- q_cosize: int = 0 k_cosize: int = 0 v_cosize: int = 0 @@ -2441,7 +2411,13 @@ class KdaCfg: k_restore_cosize: int = 0 state_scale_diag_cosize: int = 0 o_cosize: int = 0 - pairwise_cosize: int = 0 + + # TMA transaction bytes per stage + tma_q_bytes: int = 0 + tma_k_bytes: int = 0 + tma_v_bytes: int = 0 + tma_gate_bytes: int = 0 + intermediate_cosize: int = 0 def build_cfg( @@ -2460,7 +2436,6 @@ def build_cfg( v_ratio: int, n_heads_out: int, max_active_clusters: int, - split_k: bool = False, dyn_sched: bool = False, ) -> KdaCfg: """Build the per-compile ``KdaCfg`` (io_dtype in {Float16, BFloat16}); @@ -2482,51 +2457,117 @@ def build_cfg( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=max_active_clusters, - split_k=split_k, dyn_sched=dyn_sched, ) if enable_checkpoints: - # the 32 KB H staging buffer must fit next to the raw ring: trim the - # q/k/v/gate TMA lookahead for H compiles - cfg.smem_raw_stages = 6 + cfg.smem_raw_stages = 5 + cfg.smem_checkpoint_stages = 2 + cfg.smem_raw_bar_stages = cfg.smem_raw_stages + (cfg.smem_raw_stages % 2) cfg.threads_per_cta = 16 * cfg.threads_per_warp cfg.cg0_threads_per_group = cfg.cg0_warps_per_group * cfg.threads_per_warp cfg.tmem_user_threads = (1 + len(cfg.compute_group_1_warp_ids)) * cfg.threads_per_warp if cfg.smem_state_scale_diag_stages != cfg.qk_scale_ready_stages: raise ValueError("diag and qk-scale ready rings must share their rolling stage") - cfg.tmem_state_inp_offset = cfg.tmem_state_offset + cfg.d_k + cfg.tmem_state_inp_offset = cfg.tmem_state_acc_offset + cfg.d_k cfg.tmem_q_state_acc_offset = cfg.tmem_state_inp_offset + (cfg.d_k // 2) cfg.tmem_state_k_acc_offset = cfg.tmem_q_state_acc_offset + cfg.tmem_q_state_acc_stages * cfg.b_t - cfg.tmem_update_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t - cfg.tmem_rhs_inp_offset = cfg.tmem_update_acc_offset + cfg.b_t - cfg.tmem_update_inp_offset = cfg.tmem_rhs_inp_offset + (cfg.b_t // 2) - assert (cfg.tmem_update_inp_offset + (cfg.b_t // 2)) <= 512 + cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t + cfg.tmem_y_inp_offset = cfg.tmem_u_acc_offset + cfg.b_t + cfg.tmem_u_inp_offset = cfg.tmem_y_inp_offset + (cfg.b_t // 2) + assert (cfg.tmem_u_inp_offset + (cfg.b_t // 2)) <= 512 cfg.q_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t cfg.k_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t cfg.v_cosize = cfg.smem_raw_stages * cfg.d_v * cfg.b_t cfg.gate_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t - cfg.beta_cosize = cfg.smem_raw_stages * cfg.b_t + cfg.beta_cosize = cfg.smem_raw_bar_stages * cfg.b_t cfg.k_inv_cosize = cfg.smem_decay_stages * cfg.b_t * cfg.d_k cfg.k_decay_cosize = cfg.smem_decay_stages * cfg.d_k * cfg.b_t cfg.q_decay_cosize = cfg.smem_decay_stages * cfg.d_k * cfg.b_t cfg.k_restore_cosize = cfg.smem_decay_stages * cfg.d_k * cfg.b_t cfg.state_scale_diag_cosize = cfg.smem_state_scale_diag_stages * (cfg.d_k // 16) * 256 cfg.o_cosize = cfg.smem_o_stages * cfg.b_t * cfg.d_v - cfg.pairwise_cosize = cfg.smem_pairwise_stages * 2 * cfg.b_t * cfg.b_t + cfg.intermediate_cosize = cfg.smem_intermediate_stages * 2 * cfg.b_t * cfg.b_t + cfg.tma_q_bytes = cfg.d_k * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_k_bytes = cfg.d_k * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_v_bytes = cfg.d_v * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_gate_bytes = cfg.d_k * cfg.b_t * 4 return cfg -def get_workspace_size(B: int, HQ: int, HV: int) -> int: - """Bytes for the per-(batch, head) TMA-descriptor arrays (q, k, v, gate, - o, h) + 128 alignment slack.""" - HO = HQ if HQ >= HV else HV - return TENSOR_MAP_QWORDS * 8 * (6 * B * HO) + 128 +TENSORMAP_DESC_ARRAYS = 6 # per-batch runtime TMA descriptors: Q, K, V, Gate, O, state_checkpoints +TENSORMAP_STATIC_SLOTS = 0 + + +@cute.kernel +def build_all_descs_kernel( + base_q: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_v: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_gate: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_o: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], + desc_ws: cute.Tensor, + cu_seqlens: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + o: cute.Tensor, + state_checkpoints: cute.Tensor | None, + n_batch: cutlass.Int32, + q_token_stride: cutlass.Int32, + k_token_stride: cutlass.Int32, + v_token_stride: cutlass.Int32, + gate_token_stride: cutlass.Int32, + o_token_stride: cutlass.Int32, + checkpoint_entry_stride: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, +) -> None: + """Single-launch builder for the per-BATCH descriptor arrays (one warp + per array).""" + tidx, _, _ = cute.arch.thread_idx() + widx = cutlass.Int32(tidx) // cutlass.Int32(32) + arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_q_arr = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) + desc_k_arr = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_v_arr = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_gate_arr = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_o_arr = cute.make_tensor(desc_ws.iterator + 4 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_checkpoint_arr = cute.make_tensor(desc_ws.iterator + 5 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + + if widx == 0: + if nvvm.elect_sync(): + emit_seq_descs(base_q, desc_q_arr, cu_seqlens, q, n_batch, q_token_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 1: + if nvvm.elect_sync(): + emit_seq_descs(base_k, desc_k_arr, cu_seqlens, k, n_batch, k_token_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 2: + if nvvm.elect_sync(): + emit_seq_descs(base_v, desc_v_arr, cu_seqlens, v, n_batch, v_token_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 3: + if nvvm.elect_sync(): + emit_seq_descs(base_gate, desc_gate_arr, cu_seqlens, gate, n_batch, gate_token_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 4: + if nvvm.elect_sync(): + emit_seq_descs(base_o, desc_o_arr, cu_seqlens, o, n_batch, o_token_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if cutlass.const_expr(state_checkpoints is not None): + if widx == 5: + if nvvm.elect_sync(): + emit_checkpoint_seq_descs( + base_checkpoint, desc_checkpoint_arr, cu_seqlens, state_checkpoints, n_batch, checkpoint_entry_stride, checkpoint_every_n, 2 + ) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) @cute.jit -def _build_descs( +def build_descs( io_dtype: cutlass.Constexpr, b_t: cutlass.Constexpr[int], q: cute.Tensor, @@ -2534,139 +2575,80 @@ def _build_descs( v: cute.Tensor, gate: cute.Tensor, o: cute.Tensor, - h: cute.Tensor | None, + state_checkpoints: cute.Tensor | None, cu_seqlens: cute.Tensor, tensormap_workspace: cute.Tensor, - h_every_n: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, stream: cuda_driver.CUstream, ): """Build the 6 per-(batch, head) TMA-descriptor arrays (q, k, v, gate, - o, h) into ``tensormap_workspace``. - - Launched on every execute: the descriptors fold cu_seqlens contents into - GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. Each descriptor - folds the sequence base + head offset into 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 - H descriptor is 3-D ``(dv, dk, entry)`` over the packed ``[total_h, HO, - DK, DV]`` series; its per-sequence entry offsets ((seqlen-1)//N, - prefix-summed) are derived on device and its entry extent is capped per - sequence, so H store coordinates are sequence-local.""" + o, state_checkpoints) into ``tensormap_workspace``.""" h_q = q.shape[1] h_k = k.shape[1] h_v = v.shape[1] - ho = gate.shape[1] + n_heads_out = gate.shape[1] batch_size = cu_seqlens.shape[0] - 1 d_k = q.shape[2] d_v = v.shape[2] bpe = io_dtype.width // 8 - granu = 128 // bpe + tma_granu_elems = 128 // bpe seqlen = q.shape[0] - def _head0(t, dim, heads): - # 2-D (dim, token) head-0 view: box (granu, b_t) matches the main - # kernel's SMEM staging byte-for-byte - return cute.make_tensor(t.iterator, cute.make_layout((dim, seqlen), stride=(1, heads * dim))) + q_headed = cute.make_tensor(q.iterator, cute.make_layout((d_k, h_q, seqlen), stride=(1, q.stride[1], q.stride[0]))) + k_headed = cute.make_tensor(k.iterator, cute.make_layout((d_k, h_k, seqlen), stride=(1, k.stride[1], k.stride[0]))) + v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v.stride[1], v.stride[0]))) + gate_headed = cute.make_tensor(gate.iterator, cute.make_layout((d_k, n_heads_out, seqlen), stride=(1, gate.stride[1], gate.stride[0]))) + o_headed = cute.make_tensor(o.iterator, cute.make_layout((d_v, n_heads_out, seqlen), stride=(1, o.stride[1], o.stride[0]))) swz = cuda.TensorMapSwizzle.s128b - base_q = cuda.create_tensor_map_tiled_from_view(_head0(q, d_k, h_q), box_dims=(granu, b_t), stride_order=(0, 1), swizzle=swz) - base_k = cuda.create_tensor_map_tiled_from_view(_head0(k, d_k, h_k), box_dims=(granu, b_t), stride_order=(0, 1), swizzle=swz) - base_v = cuda.create_tensor_map_tiled_from_view(_head0(v, d_v, h_v), box_dims=(granu, b_t), stride_order=(0, 1), swizzle=swz) - base_gate = cuda.create_tensor_map_tiled_from_view(_head0(gate, d_k, ho), box_dims=(32, b_t), stride_order=(0, 1), swizzle=swz) - base_o = cuda.create_tensor_map_tiled_from_view(_head0(o, d_v, ho), box_dims=(granu, b_t), stride_order=(0, 1), swizzle=swz) - - arr_words = (batch_size * ho) * TENSOR_MAP_QWORDS - ws_iter = tensormap_workspace.iterator - - def _sub(i): - return cute.make_tensor(ws_iter + i * arr_words, cute.make_layout((arr_words,), stride=(1,))) - - build_qkv_load_descs_kernel( - base_q, _sub(0), cu_seqlens, q, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(ho // h_q), cutlass.Int32(d_k), cutlass.Int32(h_q * d_k), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_k, _sub(1), cu_seqlens, k, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(ho // h_k), cutlass.Int32(d_k), cutlass.Int32(h_k * d_k), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_v, _sub(2), cu_seqlens, v, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(ho // h_v), cutlass.Int32(d_v), cutlass.Int32(h_v * d_v), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_gate, _sub(3), cu_seqlens, gate, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(1), cutlass.Int32(d_k), cutlass.Int32(ho * d_k), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - build_qkv_load_descs_kernel( - base_o, _sub(4), cu_seqlens, o, cutlass.Int32(batch_size), cutlass.Int32(ho), cutlass.Int32(1), cutlass.Int32(d_v), cutlass.Int32(ho * d_v), 1 - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - if cutlass.const_expr(h is not None): - h_view = cute.make_tensor( - h.iterator, + base_q = cuda.create_tensor_map_tiled_from_view(q_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_k = cuda.create_tensor_map_tiled_from_view(k_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_v = cuda.create_tensor_map_tiled_from_view(v_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_gate = cuda.create_tensor_map_tiled_from_view(gate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_o = cuda.create_tensor_map_tiled_from_view(o_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + + base_checkpoint = base_o + if cutlass.const_expr(state_checkpoints is not None): + checkpoint_view = cute.make_tensor( + state_checkpoints.iterator, cute.make_layout( - (h.shape[3], h.shape[2], h.shape[0]), - stride=(h.stride[3], h.stride[2], h.stride[0]), + (state_checkpoints.shape[3], state_checkpoints.shape[2], state_checkpoints.shape[0], n_heads_out), + stride=(state_checkpoints.stride[3], state_checkpoints.stride[2], state_checkpoints.stride[0], state_checkpoints.stride[1]), ), ) - base_h = cuda.create_tensor_map_tiled_from_view(h_view, box_dims=(granu, d_k, 1), stride_order=(0, 1, 2), swizzle=swz) - build_h_descs_kernel( - base_h, - _sub(5), - cu_seqlens, - h, - cutlass.Int32(batch_size), - cutlass.Int32(ho), - cutlass.Int32(h.stride[1]), - cutlass.Int32(h.stride[0]), - h_every_n, - 2, - ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) - - -# --------------------------------------------------------------------------- -# Torch adapter / host-side compilation -# --------------------------------------------------------------------------- - - -def _device_sm_count() -> int: - """Multiprocessor count of the current device (runtime API: auto-inits - the primary context, so this works before any other CUDA call).""" - from cuda.bindings import runtime as _rt - - err, dev = _rt.cudaGetDevice() - if int(err) != 0: - raise RuntimeError(f"cudaGetDevice failed: {err}") - err, count = _rt.cudaDeviceGetAttribute(_rt.cudaDeviceAttr.cudaDevAttrMultiProcessorCount, dev) - if int(err) != 0: - raise RuntimeError(f"cudaDeviceGetAttribute failed: {err}") - return count - - -def _data_ptr(t) -> int: - """Device address of a tensor-like (``data_ptr()`` or the CUDA array - interface).""" - fn = getattr(t, "data_ptr", None) - if fn is not None: - return fn() - return t.__cuda_array_interface__["data"][0] - - -def _cutlass_io_dtype(dtype): - name = str(dtype) - if "bfloat16" in name: - return cutlass.BFloat16 - if "float16" in name or "half" in name: - return cutlass.Float16 - raise ValueError(f"Unsupported dtype {dtype}, expected bfloat16 or float16") - - -def _cutlass_state_dtype(dtype): - name = str(dtype) - if "bfloat16" in name: - return cutlass.BFloat16 - if "float32" in name: - return cutlass.Float32 - raise ValueError(f"Unsupported state dtype {dtype}, expected float32 or bfloat16") + base_checkpoint = cuda.create_tensor_map_tiled_from_view(checkpoint_view, box_dims=(tma_granu_elems, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) + n_warps = 6 if state_checkpoints is not None else 5 + build_all_descs_kernel( + base_q, + base_k, + base_v, + base_gate, + base_o, + base_checkpoint, + tensormap_workspace, + cu_seqlens, + q, + k, + v, + gate, + o, + state_checkpoints, + cutlass.Int32(batch_size), + cutlass.Int32(q.stride[0]), + cutlass.Int32(k.stride[0]), + cutlass.Int32(v.stride[0]), + cutlass.Int32(gate.stride[0]), + cutlass.Int32(o.stride[0]), + cutlass.Int32(state_checkpoints.stride[0] if state_checkpoints is not None else 0), + checkpoint_every_n, + ).launch(grid=(1, 1, 1), block=(32 * n_warps, 1, 1), stream=stream) + + +# ---- Torch adapter / host-side compilation --------------------------------------- @lru_cache(maxsize=None) -def _get_compiled_cache( +def get_compiled_cache( io_dtype_str: str, state_dtype_str: str, cu_dtype_str: str, @@ -2680,7 +2662,6 @@ def _get_compiled_cache( safe_gate: bool, gate_lower_bound: float, beta_sigmoid: bool, - split_k: bool, dyn_sched: bool, ): """Return a mutable dict that lazily stores the compiled kernel.""" @@ -2701,7 +2682,6 @@ def compile( k_ratio: int, v_ratio: int, n_heads_out: int, - split_k: bool = False, dyn_sched: bool = False, *, num_sm: int, @@ -2713,9 +2693,9 @@ def compile( dt_bias_cute, beta_cute, cu_seqlens_cute, - s_in_cute, + state_in_cute, o_cute, - s_out_cute, + state_out_cute, work_items_cute=None, work_count_cute=None, sched_ctr_cute=None, @@ -2740,12 +2720,11 @@ def compile( v_ratio=v_ratio, n_heads_out=n_heads_out, max_active_clusters=num_sm, - split_k=split_k, dyn_sched=dyn_sched, ) return cute.compile( - _host, + host, cfg, q_cute, k_cute, @@ -2755,9 +2734,9 @@ def compile( dt_bias_cute, beta_cute, cu_seqlens_cute, - s_in_cute, + state_in_cute, o_cute, - s_out_cute, + state_out_cute, work_items_cute, work_count_cute, sched_ctr_cute, @@ -2781,7 +2760,7 @@ def chunk_kda_sm100( output_state, scale: float, checkpoint_every_n_tokens: int = 0, - output_checkpoints=None, + output_state_checkpoints=None, use_qk_l2norm_in_kernel: bool = False, safe_gate: bool = False, gate_lower_bound: float = DEFAULT_GATE_LOWER_BOUND, @@ -2797,7 +2776,9 @@ def chunk_kda_sm100( ) -> None: """Execute the Blackwell BT=16 chunked KDA prefill kernel. - All tensors must be contiguous and on the same CUDA device. + All tensors must be on the same CUDA device with a stride-1 innermost + dim; outer strides are free (padded / permuted views are read through + the TMA descriptors and dynamic layouts). Args: q: ``(total_tokens, HQ, DK)`` float16/bfloat16 @@ -2813,13 +2794,13 @@ def chunk_kda_sm100( initial_state: ``(num_seqs, HO, DK, DV)`` float32/bfloat16, or None output_state: ``(num_seqs, HO, DK, DV)`` float32/bfloat16, or None scale: attention scale factor (must not be 0) - checkpoint_every_n_tokens: emit an H entry every N tokens (0 = off). - H[j] is the state after ``(j + 1) * N`` tokens, STRICTLY BEFORE + 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 state + ``output_state``. With ``N == B_T`` this is the per-chunk checkpoint series the backward pass consumes. - output_checkpoints: ``(total_h, HO, DK, DV)`` io-dtype (KV, v - contiguous — the GDN H layout); the per-sequence entry offsets + output_state_checkpoints: ``(total_checkpoints, HO, DK, DV)`` io-dtype (KV, V + 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 use_qk_l2norm_in_kernel: L2-normalize q/k rows inside the kernel @@ -2827,13 +2808,11 @@ def chunk_kda_sm100( a_log: ``(HO,)`` float32, safe-gate per-head log-amplitude (None = 0) dt_bias: ``(HO, DK)`` float32, safe-gate channel bias (None = 0) use_beta_sigmoid_in_kernel: ``beta`` holds logits; sigmoid in-kernel - work_items: ``(max_items, 6)`` int32 split-K work-item table from - ``common/split_k.py``, or None for the one-tile-per-(b,h) - layout. With a table, each item computes chunks - ``[cstart, wend)`` and writes O/checkpoints only for - ``[wstart, wend)``. - work_count: ``(1,)`` int32 device-side item count (required with - work_items) + work_items: ``(max_items, 8)`` int32 work-item table from + ``common/split_k.py`` (REQUIRED; an uncut table row is the whole + (b, h) sequence). Each item computes chunks ``[cstart, wend)`` + and writes O/checkpoints only for ``[wstart, wend)``. + work_count: ``(1,)`` int32 device-side item count (REQUIRED) sched_ctr: ``(2,)`` int32 device scratch ``[ticket, done]`` enabling the dynamic (work-stealing) tile scheduler; must be zeroed before every launch (``build_split_table`` does this when it is passed as @@ -2847,21 +2826,15 @@ def chunk_kda_sm100( store_final_state = output_state is not None enable_checkpoints = checkpoint_every_n_tokens > 0 if enable_checkpoints: - if output_checkpoints is None: - raise ValueError("checkpoint_every_n_tokens > 0 requires output_checkpoints") - if str(output_checkpoints.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: + if output_state_checkpoints is None: + raise ValueError("checkpoint_every_n_tokens > 0 requires output_state_checkpoints") + if str(output_state_checkpoints.dtype).split(".")[-1] != str(q.dtype).split(".")[-1]: raise ValueError( - f"output_checkpoints dtype must match the io dtype (fp32 state belongs to output_state): got {output_checkpoints.dtype} with io {q.dtype}" + f"output_state_checkpoints dtype must match the io dtype (fp32 state belongs to output_state): got {output_state_checkpoints.dtype} with io {q.dtype}" ) - split_k = work_items 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 - if split_k: - if work_count is None: - raise ValueError("work_count is required with work_items") - if enable_checkpoints and checkpoint_every_n_tokens != CFG.B_T: - raise ValueError(f"split-K checkpoints require checkpoint_every_n_tokens == {CFG.B_T}, got {checkpoint_every_n_tokens}") - elif work_count is not None: - raise ValueError("work_count must be None without work_items") if initial_state is not None: state_dtype_src = initial_state.dtype @@ -2883,11 +2856,9 @@ def chunk_kda_sm100( if not safe_gate: a_log = None dt_bias = None - if _data_ptr(tensormap_workspace) % 128 != 0: - raise ValueError("tensormap_workspace must be 128-byte aligned") cu_stream = cuda_driver.CUstream(int(stream)) - cache = _get_compiled_cache( + cache = get_compiled_cache( str(q.dtype), str(state_dtype_src), str(cu_seqlens.dtype), @@ -2901,45 +2872,33 @@ def chunk_kda_sm100( safe_gate, gate_lower_bound, use_beta_sigmoid_in_kernel, - split_k, dyn_sched, ) if "compiled" not in cache: - io_dtype = _cutlass_io_dtype(q.dtype) - state_dtype = _cutlass_state_dtype(state_dtype_src) - q_cute = from_dlpack(q, assumed_align=16) - q_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - k_cute = from_dlpack(k, assumed_align=16) - k_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - v_cute = from_dlpack(v, assumed_align=16) - v_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - gate_cute = from_dlpack(gate, assumed_align=16) - gate_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + io_dtype = get_dtype(q.dtype) + state_dtype = get_dtype(state_dtype_src) + q_cute = from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=2) + k_cute = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_cute = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_cute = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2) 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 - beta_cute = from_dlpack(beta, assumed_align=4) - beta_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) - o_cute = from_dlpack(output, assumed_align=16) - o_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + beta_cute = from_dlpack(beta, assumed_align=4).mark_layout_dynamic(leading_dim=1) + o_cute = from_dlpack(output, assumed_align=16).mark_layout_dynamic(leading_dim=2) cu_seqlens_cute = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() - s_in_cute = None + state_in_cute = None if use_initial_state: - s_in_cute = from_dlpack(initial_state, assumed_align=16) - s_in_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=CFG.D_K) + state_in_cute = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) - s_out_cute = None + state_out_cute = None if store_final_state: - s_out_cute = from_dlpack(output_state, assumed_align=16) - s_out_cute.mark_layout_dynamic().mark_compact_shape_dynamic(mode=3, stride_order=(0, 1, 2, 3), divisibility=CFG.D_K) + state_out_cute = from_dlpack(output_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) - work_items_cute = None - work_count_cute = None - if split_k: - work_items_cute = from_dlpack(work_items, assumed_align=4) - work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) - work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() + work_items_cute = from_dlpack(work_items, assumed_align=16) + work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() sched_ctr_cute = None if dyn_sched: @@ -2961,9 +2920,8 @@ def chunk_kda_sm100( k_ratio, v_ratio, HO, - split_k, dyn_sched, - num_sm=_device_sm_count(), + num_sm=multiprocessor_count(current_device_id()), q_cute=q_cute, k_cute=k_cute, v_cute=v_cute, @@ -2972,9 +2930,9 @@ def chunk_kda_sm100( dt_bias_cute=dt_bias_cute, beta_cute=beta_cute, cu_seqlens_cute=cu_seqlens_cute, - s_in_cute=s_in_cute, + state_in_cute=state_in_cute, o_cute=o_cute, - s_out_cute=s_out_cute, + state_out_cute=state_out_cute, work_items_cute=work_items_cute, work_count_cute=work_count_cute, sched_ctr_cute=sched_ctr_cute, @@ -2985,36 +2943,26 @@ def chunk_kda_sm100( ) compiled = cache["compiled"] - - # The descriptors encode cu_seqlens' CONTENTS, which no key built from the - # buffers can track. The skip this replaces asked torch's _version counter, - # so it was sound for a torch caller and silently stale for every other - # producer. Rebuilding unconditionally measures free: 131 vs 135 us of host - # time, and 157 either way once the launches are waited on. - h_for_descs = output_checkpoints if enable_checkpoints else None - if cache.get("build_descs_has_h") != (h_for_descs is not None): + state_checkpoints_for_descs = output_state_checkpoints if enable_checkpoints else None + # desc build runs every execute by contract (cu contents are data; + # buffer pointers may change) — capture-safe, single tiny launch + if cache.get("build_descs_has_state_checkpoints") != (state_checkpoints_for_descs is not None): cache.pop("build_descs", None) - cache["build_descs_has_h"] = h_for_descs is not None + cache["build_descs_has_state_checkpoints"] = state_checkpoints_for_descs is not None if "build_descs" not in cache: - io_dtype = _cutlass_io_dtype(q.dtype) - q_bd = from_dlpack(q, assumed_align=16) - q_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - k_bd = from_dlpack(k, assumed_align=16) - k_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - v_bd = from_dlpack(v, assumed_align=16) - v_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - gate_bd = from_dlpack(gate, assumed_align=16) - gate_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) - o_bd = from_dlpack(output, assumed_align=16) - o_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2), divisibility=1) + io_dtype = get_dtype(q.dtype) + q_bd = from_dlpack(q, assumed_align=16).mark_layout_dynamic(leading_dim=2) + k_bd = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_bd = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_bd = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2) + o_bd = from_dlpack(output, assumed_align=16).mark_layout_dynamic(leading_dim=2) cu_bd = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() ws_bd = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() - h_bd = None - if h_for_descs is not None: - h_bd = from_dlpack(h_for_descs, assumed_align=16) - h_bd.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1, 2, 3), divisibility=1) + state_checkpoints_bd = None + if state_checkpoints_for_descs is not None: + state_checkpoints_bd = from_dlpack(state_checkpoints_for_descs, assumed_align=16).mark_layout_dynamic(leading_dim=3) cache["build_descs"] = cute.compile( - _build_descs, + build_descs, io_dtype, CFG.B_T, q_bd, @@ -3022,7 +2970,7 @@ def chunk_kda_sm100( v_bd, gate_bd, o_bd, - h_bd, + state_checkpoints_bd, cu_bd, ws_bd, cutlass.Int32(checkpoint_every_n_tokens), @@ -3035,13 +2983,12 @@ def chunk_kda_sm100( v, gate, output, - h_for_descs, + state_checkpoints_for_descs, cu_seqlens, tensormap_workspace, checkpoint_every_n_tokens, cu_stream, ) - compiled( q, k, diff --git a/python/cudnn/linear_attention/frost/kernel/kda_recompute_config.py b/python/cudnn/linear_attention/frost/kernel/kda_recompute_config.py new file mode 100644 index 000000000..0e39486e8 --- /dev/null +++ b/python/cudnn/linear_attention/frost/kernel/kda_recompute_config.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# This kernel is derived from cuDNN, NVIDIA Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Kimi Delta Attention (KDA) Cutlass DSL recompute (state/H-only) kernel +config (fixed compile-time constants). The BT=16 KDA schedule uses a 16-warp +(512-thread) specialization with a per-key-channel decay; the derived +SMEM/TMEM sizes and offsets are stamped by ``build_cfg`` in +``kda_recompute_f16.py``. + +Target arch: Blackwell SM100 (GB200) / SM103 (GB300). +""" + +from dataclasses import dataclass +from typing import Tuple + + +@dataclass(frozen=True) +class Cfg: + # --- tile shape --- + B_T: int = 16 # chunk-inner token tile (BT=16 KDA schedule) + D_K: int = 128 # query/key head dim + D_V: int = 128 # value head dim + + # --- warp assignments (16 warps = 512 threads) --- + COMPUTE_GROUP_0_WARP_IDS: Tuple[int, ...] = (0, 1, 2, 3, 4, 5, 6, 7) # decay-operand materialize (2-group ping-pong) + COMPUTE_GROUP_1_WARP_IDS: Tuple[int, ...] = (8, 9, 10, 11) # value-side TMEM / state staging + SUPER_MMA_WARP_ID: int = 12 # register-MMA KK + Neumann T_inv + TCGEN05_MMA_WARP_ID: int = 13 # tcgen05 state GEMMs + TMA_WARP_ID: int = 14 # k/v/gate TMA loads + EPILOGUE_WARP_ID: int = 15 # H TMA store + + # --- register split --- + NUM_REGS_COMPUTE_GROUP_0: int = 160 + NUM_REGS_COMPUTE_GROUP_1: int = 136 + NUM_REGS_OTHER: int = 56 + + THREADS_PER_WARP: int = 32 + + BUFFER_ALIGN_BYTES: int = 1024 + + # --- SMEM / TMEM ring stage counts --- + SMEM_RAW_STAGES: int = 8 + SMEM_SCHED_STAGES: int = 8 + SMEM_DECAY_STAGES: int = 2 + SMEM_INTERMEDIATE_STAGES: int = 2 + SMEM_STATE_SCALE_DIAG_STAGES: int = 4 + QK_SCALE_READY_STAGES: int = 4 + + CLUSTER_SHAPE_MNK: Tuple[int, int, int] = (1, 1, 1) + + +CFG = Cfg() diff --git a/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py b/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py new file mode 100644 index 000000000..c78fbf232 --- /dev/null +++ b/python/cudnn/linear_attention/frost/kernel/kda_recompute_f16.py @@ -0,0 +1,2514 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# This kernel is derived from cuDNN, NVIDIA Corporation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Chunked Kimi Delta Attention (KDA) recompute (state/checkpoints-only) kernel for +Blackwell SM100/SM103 (Cutlass DSL), BT=16 tiling with a per-key-channel decay. +Framework-neutral entry ``chunk_kda_recompute_sm100``. + +Persistent kernel: the grid is the SM count and every warp role +runs a tile-scheduler loop (``decode_work_item``); a tile is one (batch, +head) sequence, or one split-K work item computing chunks ``[cstart, wend)`` +and writing checkpoints only for the owned ``[wstart, wend)`` (see +``common/split_k.py``; warmup chunks rebuild the incoming state from +zero). All ring stage/phase bookkeeping runs on cumulative per-CTA chunk +counters so pipelines flow seamlessly across tiles. + +Pipeline (direct CUTLASS primitives, chunk_idx-size 16 KDA schedule): + + load K/V/Gate/Beta + optional in-kernel L2-norm of K (L2NORM specialization) + exp2(g), exp2(-g), stage final-token exp2(g) as exp2(g_last) + super-MMA: KK/Neumann inverse + apply Beta + tcgen05-MMA: State*K / U / state update + store periodic state checkpoints, final state + +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 +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. + +Warp assignments (16 warps = 512 threads): + warps 0-7 : compute group 0 - Gate prefix scan + decay/restore operands + warps 8-11 : compute group 1 - TMEM value side, state stores + warp 12 : super-MMA - register-MMA KK^T + Neumann inverse + warp 13 : tcgen05-MMA - the four state GEMMs + the TMEM lifecycle + warp 14 : TMA load - per-chunk input G->S loads + warp 15 : epilogue - the checkpoint TMA store + +SMEM layout: + Buffer Bytes Stages + K / V raw 32768 8 <-- SW128 TMA ring (io dtype) + Gate raw 65536 8 <-- fp32 prefix-scan source + Beta 512 8 <-- fp32 per-token scalars + K_inv 8192 2 <-- token-major ldmatrix/tcgen05 B operand + K decay 8192 2 <-- tcgen05 SW128 K-box-major A/B operands + K restore 8192 2 <-- tcgen05 B operand for the state update + state-scale diag 12288 3 <-- per-k-atom decay diagonal blocks + T_inv 2048 2 <-- SW32 16x16 register-MMA tiles + +TMEM layout (240 of 512 columns): + Buffer Cols Purpose + state 0-127 state[DK,DV] fp32 recurrent state + state inp 128-191 packed b16 A operand view of the state + state_k_acc 192-207 State*K fp32 accumulator + u_acc 208-223 U fp32 accumulator + Y 224-231 packed b16 A operand: Beta * (V - State*K) + U input 232-239 packed b16 A operand: the b16 U repack + +GEMM schedule (tcgen05-MMA warp, in issue order per chunk): + State*K -> state_k_acc + State decay (diag blocks) + U = Y(T) @ T_inv -> u_acc + final_state += U @ K_restore + +Requires a cutlass DSL build providing `cutlass.experimental.*`; not +available in the pip nvidia-cutlass-dsl releases. +""" + +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable, NamedTuple, Optional, Type + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.experimental.cuda as cuda +import cutlass.experimental.primitives as nvvm +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + +from ..common.split_k import decode_work_item +from ..common.host import get_dtype +from cudnn.frost.buffers import current_device_id, data_ptr +from cudnn.frost.device import multiprocessor_count +from ..common.thd import TENSOR_MAP_QWORDS, emit_checkpoint_seq_descs, emit_seq_descs +from .kda_recompute_config import CFG +from cudnn.frost.tile_dsl.barrier import ( + advance, + MBarrier, + PipelineState, + Producer, +) +from cudnn.frost.tile_dsl.handles import GmemTileTma, MmaDesc, SmemTile, tma_slice_runtime_desc +from cudnn.frost.tile_dsl.mma import mma_step, mma_ts_step +from cudnn.frost.tile_dsl.swizzle import swizzle_lin_128b, swizzle_lin_S, swizzle_xor_128b +from cudnn.frost.tile_dsl.tma import tma_load_tile, tma_store_commit, tma_store_tile, tma_store_wait, tma_tensormap_acquire +from cudnn.frost.tile_dsl.pointwise import ( + opaque_f32_zero, + f16x2_to_f32, + fadd2, + fmul2, + ffma2, + movmatrix_16b, + mul_f16x2, + fp32_to_fp16, + sub_f16x2, +) + +LOG2_E: float = 1.4426950408889634 + + +DEFAULT_GATE_LOWER_BOUND: float = -5.0 + + +# Host-side API defaults. + + +L2_NORM_EPS: float = 1.0e-12 + + +class KdaBars(NamedTuple): + """Every inter-warp handoff as an ``MBarrier`` over its ring. Consumers track ``(idx, phase)`` inline; the producer tag selects + the arrive lowering (``TMA_LOAD``/``MMA_COMMIT``/``THREAD``).""" + + mb_k_ready: MBarrier + mb_k_done: MBarrier + mb_v_ready: MBarrier + mb_v_done: MBarrier + mb_gate_ready: MBarrier + mb_gate_done: MBarrier + + mb_beta_ready: MBarrier + mb_beta_done: MBarrier + + mb_state_k_acc_ready: MBarrier + mb_u_acc_ready: MBarrier + + mb_state_inp_ready: MBarrier + mb_y_inp_ready: MBarrier + mb_u_inp_ready: MBarrier + + mb_t_inv_ready: MBarrier + mb_t_inv_done: MBarrier + mb_qk_scale_ready: MBarrier + mb_state_scale_diag_done: MBarrier + mb_k_decay_inv_cg0_ready: MBarrier + mb_decay_tcgen05_done: MBarrier + mb_decay_super_done: MBarrier + mb_k_restore_done: MBarrier + + mb_state_acc_done: MBarrier + mb_state_acc_read_done: MBarrier + mb_tmem_done: MBarrier + + mb_checkpoint_tmastg_ready: MBarrier + mb_checkpoint_tmastg_done: MBarrier + + mb_sched_ready: MBarrier + mb_sched_done: MBarrier + + +def make_kda_bars(cfg) -> KdaBars: + """Bars factory. MUST be called from inside ``kernel`` (allocates the + mbarrier rings in SMEM ahead of the data buffers).""" + + def alloc(n): + return cutlass.Array(cutlass.Int64, n, space=cutlass.AddressSpace.smem, alignment=8) + + WARP = cfg.threads_per_warp + CG0_GROUP_THREADS = cfg.cg0_warps_per_group * WARP + CG1_THREADS = len(cfg.compute_group_1_warp_ids) * WARP + + return KdaBars( + mb_k_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_k_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_v_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_v_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_gate_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=1, producer=Producer.TMA_LOAD), + mb_gate_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_beta_ready=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=WARP, producer=Producer.THREAD), + mb_beta_done=MBarrier(alloc(cfg.smem_raw_stages), stages=cfg.smem_raw_stages, init_count=WARP + CG1_THREADS, producer=Producer.THREAD), + mb_state_k_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), + mb_u_acc_ready=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), + mb_state_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_y_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_u_inp_ready=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_t_inv_ready=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=WARP, producer=Producer.THREAD), + mb_t_inv_done=MBarrier(alloc(cfg.smem_intermediate_stages), stages=cfg.smem_intermediate_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_qk_scale_ready=MBarrier( + alloc(cfg.qk_scale_ready_stages), + stages=cfg.qk_scale_ready_stages, + init_count=CG0_GROUP_THREADS, + producer=Producer.THREAD, + ), + mb_state_scale_diag_done=MBarrier( + alloc(cfg.smem_state_scale_diag_stages), + stages=cfg.smem_state_scale_diag_stages, + init_count=1, + producer=Producer.MMA_COMMIT, + ), + mb_k_decay_inv_cg0_ready=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=CG0_GROUP_THREADS, producer=Producer.THREAD), + mb_decay_tcgen05_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_decay_super_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=WARP, producer=Producer.THREAD), + mb_k_restore_done=MBarrier(alloc(cfg.smem_decay_stages), stages=cfg.smem_decay_stages, init_count=1, producer=Producer.MMA_COMMIT), + mb_state_acc_done=MBarrier(alloc(1), stages=1, init_count=1, producer=Producer.MMA_COMMIT), + mb_state_acc_read_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_tmem_done=MBarrier(alloc(1), stages=1, init_count=CG1_THREADS, producer=Producer.THREAD), + mb_checkpoint_tmastg_ready=MBarrier( + alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=CG1_THREADS, producer=Producer.THREAD + ), + mb_checkpoint_tmastg_done=MBarrier(alloc(cfg.smem_checkpoint_stages), stages=cfg.smem_checkpoint_stages, init_count=WARP, producer=Producer.THREAD), + mb_sched_ready=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=1, producer=Producer.THREAD), + mb_sched_done=MBarrier(alloc(cfg.sched_stages), stages=cfg.sched_stages, init_count=15, producer=Producer.THREAD), + ) + + +# ---- Dynamic tile scheduler ------------------------------------------------------ + + +@cute.jit +def sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas): + """TMA-warp side: pull the next tile off the global ticket, publish it.""" + if cutlass.const_expr(cfg.dyn_sched): + bars.mb_sched_done[sched_state.idx].wait(sched_state.phase) + if nvvm.elect_sync(): + fetched = cutlass.Int32(nvvm.atomicrmw("add", mSched.iterator, cutlass.Int32(1), mem_order="relaxed", syncscope="gpu")) + sSched[sched_state.idx] = num_ctas + fetched + nvvm.bar_warp_sync(cute.arch.FULL_MASK) + next_tile = sSched[sched_state.idx] + if nvvm.elect_sync(): + bars.mb_sched_ready[sched_state.idx].arrive() + return next_tile, advance(sched_state, cfg.sched_stages) + return tile_idx + num_ctas, sched_state + + +@cute.jit +def sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas): + """Consumer side: read the TMA warp's published next tile.""" + if cutlass.const_expr(cfg.dyn_sched): + bars.mb_sched_ready[sched_state.idx].wait(sched_state.phase) + next_tile = sSched[sched_state.idx] + if nvvm.elect_sync(): + bars.mb_sched_done[sched_state.idx].arrive() + return next_tile, advance(sched_state, cfg.sched_stages) + return tile_idx + num_ctas, sched_state + + +@cute.jit +def tmaldg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + mSched, + sSched, + lane, + sK_raw, + sV_raw, + sGate_raw, + desc_k_base, + desc_v_base, + desc_gate_base, + bars, +) -> None: + """TMA-LDG warp role (warp 14): persistent scheduler loop issuing the + per-chunk K/V/Gate G->S loads.""" + elect_one = nvvm.elect_sync() + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + sK_tma = SmemTile( + base=sK_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sV_tma = SmemTile( + base=sV_raw, + elems_per_stage=(cfg.d_v * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=(cfg.b_t * 64), + ) + sGate_tma = SmemTile( + base=sGate_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_raw_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_k // 32), + tma_granu_elems=32, + tma_subtile_stride_elems=(cfg.b_t * 32), + ) + raw_index = PipelineState.start(phase=1) + sched_state = PipelineState.start(phase=1) + 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) + head_o = head_idx + head_k = head_idx if cfg.k_ratio == 1 else head_idx // cutlass.Int32(cfg.k_ratio) + head_v = head_idx if cfg.v_ratio == 1 else head_idx // cutlass.Int32(cfg.v_ratio) + slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_k_slot = (desc_k_base + slot).tospace(cutlass.AddressSpace.generic) + desc_v_slot = (desc_v_base + slot).tospace(cutlass.AddressSpace.generic) + desc_gate_slot = (desc_gate_base + slot).tospace(cutlass.AddressSpace.generic) + if elect_one: + tma_tensormap_acquire(desc_k_slot) + tma_tensormap_acquire(desc_v_slot) + tma_tensormap_acquire(desc_gate_slot) + for chunk_idx in cutlass.range(cstart, wend, 1, unroll=1): + chunk_start = chunk_idx * cfg.b_t + + # ---- K load ---------------------------------------------------------- + bars.mb_k_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_k_ready[raw_index.idx].arrive(n_bytes=cfg.tma_k_bytes) + k_slice = tma_slice_runtime_desc(desc_k_slot, cutlass.Int32(0), head_k, chunk_start) + tma_load_tile(sK_tma[raw_index.idx], k_slice, bars.mb_k_ready[raw_index.idx].smem_ptr, acquire=False) + + # ---- Gate load ------------------------------------------------------- + bars.mb_gate_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_gate_ready[raw_index.idx].arrive(n_bytes=cfg.tma_gate_bytes) + gate_slice = tma_slice_runtime_desc(desc_gate_slot, cutlass.Int32(0), head_o, chunk_start) + tma_load_tile(sGate_tma[raw_index.idx], gate_slice, bars.mb_gate_ready[raw_index.idx].smem_ptr, acquire=False) + + # ---- V load ---------------------------------------------------------- + bars.mb_v_done[raw_index.idx].wait(raw_index.phase) + if elect_one: + bars.mb_v_ready[raw_index.idx].arrive(n_bytes=cfg.tma_v_bytes) + 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) + + raw_index = advance(raw_index, cfg.smem_raw_stages) + tile_idx, sched_state = sched_publish_next(cfg, bars, sSched, mSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def super_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_inv_raw, + sIntermediate_raw, + sBeta_raw, + sK_decay_raw, + bars, +) -> None: + """Super-MMA warp role (warp 12): persistent scheduler loop computing the + register-MMA Neumann-series T_inv.""" + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + raw_index = PipelineState.start(phase=0) + t_inv_free = PipelineState.start(phase=1) + k_decay_ready = PipelineState.start(phase=0) + + # ---- ldmatrix/stmatrix lane decode ------------------------------------------- + rhs_row_coord = lane % 8 + (cutlass.Int32(8) if (lane // 16) else cutlass.Int32(0)) + rhs_col_offset = cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0) + lhs_row_coord = lane % 8 + (cutlass.Int32(8) if ((lane // 8) % 2) else cutlass.Int32(0)) + lhs_col_offset = cutlass.Int32(8) if ((lane // 8) // 2) else cutlass.Int32(0) + decay_key_mask = cutlass.Int32(8) + stsm_row_coord = lane & 7 + stsm_col_coord = cutlass.Int32(0) + if (lane // 8) & 1: + stsm_row_coord = stsm_row_coord + cutlass.Int32(8) + if lane // 8 >= 2: + stsm_col_coord = cutlass.Int32(8) + stsm_idx = swizzle_lin_S(stsm_row_coord * cfg.b_t + (stsm_col_coord ^ (cfg.b_t // 2)), bbits=1, mbase=3, sshift=3) + cum_chunk_base = cutlass.Int32(0) + 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) + num_chunks_tile = wend - cstart # processed chunks; ring bookkeeping runs on cum_chunk_base + local_chunk_idx + for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): + cum_chunk = cum_chunk_base + local_chunk_idx + decay_stage = k_decay_ready.idx + intermediate_stage = t_inv_free.idx + sBeta_ptr = sBeta_raw.data_ptr() + raw_index.idx * cfg.b_t + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sIntermediate_ptr = sIntermediate_raw.data_ptr() + intermediate_stage * (2 * cfg.b_t * cfg.b_t) + + bars.mb_k_decay_inv_cg0_ready[decay_stage].wait(k_decay_ready.phase) + k_decay_ready = advance(k_decay_ready, cfg.smem_decay_stages) + + # ---- KK = K_decay @ K_inv^T ------------------------------------------ + kk_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + kk_acc[accum_idx] = cutlass.Float32(0.0) + + for k_block in cutlass.range_constexpr((cfg.d_k // 16)): + # Load B operand + k_inv_col = k_block * 16 + rhs_col_offset + k_inv_segment = k_inv_col // 64 + rhs_frag = nvvm.ldmatrix( + sK_inv_ptr + + k_inv_segment * (cfg.b_t * 64) + + rhs_row_coord * 64 + + swizzle_xor_128b(rhs_row_coord, k_inv_col - k_inv_segment * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + # Load A operand + storage_key = (k_block * 16 + lhs_col_offset) ^ decay_key_mask + storage_slice = storage_key // 64 + kk_lhs_frag = nvvm.ldmatrix( + sK_decay_ptr + + storage_slice * (cfg.b_t * 64) + + swizzle_xor_128b(lhs_row_coord, lhs_row_coord * 64 + storage_key - storage_slice * 64, elem_bytes=2), + 4, + nvvm.MMALayout.ROW, + ) + + mma_step( + kk_acc, + (kk_lhs_frag[0], kk_lhs_frag[1], kk_lhs_frag[2], kk_lhs_frag[3]), + (rhs_frag[0], rhs_frag[1], rhs_frag[2], rhs_frag[3]), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + + # ---- L = Beta * tril(KK, -1) fragment -------------------------------- + bars.mb_beta_ready[raw_index.idx].wait(raw_index.phase) + row_lo = lane // 4 + row_hi = row_lo + cutlass.Int32(8) + beta_lo = (sBeta_ptr + row_lo).load().to(cutlass.Float32) + beta_hi = (sBeta_ptr + row_hi).load().to(cutlass.Float32) + l_regs = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + row_coord = row_hi if cutlass.const_expr(accum_idx % 4 >= 2) else row_lo + col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + if cutlass.const_expr(accum_idx % 2 == 1): + col_coord = col_coord + cutlass.Int32(1) + l_regs[accum_idx] = kk_acc[accum_idx] if row_coord > col_coord else cutlass.Float32(0.0) + for pair in cutlass.range_constexpr(4): + beta_scale = beta_hi if cutlass.const_expr(pair % 2 == 1) else beta_lo + l_regs[2 * pair], l_regs[2 * pair + 1] = fmul2(l_regs[2 * pair], l_regs[2 * pair + 1], beta_scale, beta_scale) + bars.mb_beta_done[raw_index.idx].arrive() + l_a0 = fp32_to_fp16(l_regs[0], l_regs[1], dtype=cfg.io_dtype) + l_a1 = fp32_to_fp16(l_regs[2], l_regs[3], dtype=cfg.io_dtype) + l_a2 = fp32_to_fp16(l_regs[4], l_regs[5], dtype=cfg.io_dtype) + l_a3 = fp32_to_fp16(l_regs[6], l_regs[7], dtype=cfg.io_dtype) + l_values = cutlass.Vector.from_elements((l_a0, l_a1, l_a2, l_a3), cutlass.Int32).bitcast(cfg.io_dtype).to(cutlass.Float32) + + # ---- T_inv = I - L, then three Neumann doubling rounds --------------- + tinv_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + row_coord = row_lo + if cutlass.const_expr(accum_idx % 4 >= 2): + row_coord = row_hi + col_coord = (accum_idx // 4) * 8 + 2 * (lane % 4) + if cutlass.const_expr(accum_idx % 2 == 1): + col_coord = col_coord + cutlass.Int32(1) + eye = cutlass.Float32(1.0) if row_coord == col_coord else cutlass.Float32(0.0) + tinv_acc[accum_idx] = eye - l_values[accum_idx] + + lpow_a0, lpow_a1, lpow_a2, lpow_a3 = l_a0, l_a1, l_a2, l_a3 + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(l_a0), movmatrix_16b(l_a1), movmatrix_16b(l_a2), movmatrix_16b(l_a3) + for _round in cutlass.range_constexpr(3): + # ---- Lpow = Lpow @ Lpow ------------------------------------------ + sq_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + sq_acc[accum_idx] = cutlass.Float32(0.0) + mma_step( + sq_acc, + (lpow_a0, lpow_a1, lpow_a2, lpow_a3), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + lpow_a0 = fp32_to_fp16(sq_acc[0], sq_acc[1], dtype=cfg.io_dtype) + lpow_a1 = fp32_to_fp16(sq_acc[2], sq_acc[3], dtype=cfg.io_dtype) + lpow_a2 = fp32_to_fp16(sq_acc[4], sq_acc[5], dtype=cfg.io_dtype) + lpow_a3 = fp32_to_fp16(sq_acc[6], sq_acc[7], dtype=cfg.io_dtype) + mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3 = movmatrix_16b(lpow_a0), movmatrix_16b(lpow_a1), movmatrix_16b(lpow_a2), movmatrix_16b(lpow_a3) + # ---- T_inv += T_inv @ Lpow --------------------------------------- + upd_acc = cutlass.Array(cutlass.Float32, 8, alignment=16) + for accum_idx in cutlass.range_constexpr(8): + upd_acc[accum_idx] = cutlass.Float32(0.0) + tinv_p0 = fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype) + tinv_p1 = fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype) + tinv_p2 = fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype) + tinv_p3 = fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype) + mma_step( + upd_acc, + (tinv_p0, tinv_p1, tinv_p2, tinv_p3), + (mov_lpow0, mov_lpow1, mov_lpow2, mov_lpow3), + k_step=0, + M=16, + N=16, + ab_dtype=cfg.io_dtype, + ) + tinv_lo0, tinv_hi0 = f16x2_to_f32(tinv_p0, dtype=cfg.io_dtype) + tinv_lo1, tinv_hi1 = f16x2_to_f32(tinv_p1, dtype=cfg.io_dtype) + tinv_lo2, tinv_hi2 = f16x2_to_f32(tinv_p2, dtype=cfg.io_dtype) + tinv_lo3, tinv_hi3 = f16x2_to_f32(tinv_p3, dtype=cfg.io_dtype) + tinv_acc[0] = tinv_lo0 + upd_acc[0] + tinv_acc[1] = tinv_hi0 + upd_acc[1] + tinv_acc[2] = tinv_lo1 + upd_acc[2] + tinv_acc[3] = tinv_hi1 + upd_acc[3] + tinv_acc[4] = tinv_lo2 + upd_acc[4] + tinv_acc[5] = tinv_hi2 + upd_acc[5] + tinv_acc[6] = tinv_lo3 + upd_acc[6] + tinv_acc[7] = tinv_hi3 + upd_acc[7] + + bars.mb_t_inv_done[intermediate_stage].wait(t_inv_free.phase) + t_inv_free = advance(t_inv_free, cfg.smem_intermediate_stages) + nvvm.stmatrix( + sIntermediate_ptr + (cfg.b_t * cfg.b_t) + stsm_idx, + [ + fp32_to_fp16(tinv_acc[0], tinv_acc[1], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[2], tinv_acc[3], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[4], tinv_acc[5], dtype=cfg.io_dtype), + fp32_to_fp16(tinv_acc[6], tinv_acc[7], dtype=cfg.io_dtype), + ], + nvvm.MMALayout.ROW, + shape=nvvm.StoreShape.M8N8, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_t_inv_ready[intermediate_stage].arrive() + bars.mb_decay_super_done[decay_stage].arrive() + raw_index = advance(raw_index, cfg.smem_raw_stages) + cum_chunk_base += num_chunks_tile + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def tcgen05_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + sTmem_base, + sIntermediate, + sK_decay, + sK_restore, + sState_scale_diag, + bars, +) -> None: + """tcgen05-MMA warp role (warp 13): persistent scheduler loop issuing + every state GEMM and owning the TMEM lifecycle.""" + elect_one = nvvm.elect_sync() + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + nvvm.tcgen05_alloc(sTmem_base, cutlass.Int32(512), group=nvvm.CTAGroup.CTA_1) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = sTmem_base.load() + state_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_inp_offset, cutlass.Int8) + state_dsts = tuple(nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset + k * 16, cutlass.Float32) for k in range(cfg.d_k // 16)) + state_k_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_k_acc_offset, cutlass.Float32) + u_acc_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_acc_offset, cutlass.Float32) + y_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_y_inp_offset, cutlass.Int8) + u_inp_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_u_inp_offset, cutlass.Int8) + state_dst_ptr = nvvm.make_tmem_ptr(tmem_base + cfg.tmem_state_acc_offset, cutlass.Float32) + state_inp_index = PipelineState.start(phase=0) + state_read_index = PipelineState.start(phase=0) + y_inp_index = PipelineState.start(phase=0) + u_inp_index = PipelineState.start(phase=0) + qk_scale_index = PipelineState.start(phase=0) + k_decay_ready = PipelineState.start(phase=0) + t_inv_ready = PipelineState.start(phase=0) + + # ---- chunk-invariant GEMM descriptors ---------------------------------------- + bpe = cfg.io_dtype.width // 8 + idesc_acc = 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, + b_major=0, + ) + idesc_diag = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=16, + m_dim=cfg.d_v, + b_major=0, + ) + idesc_final_state = nvvm.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cfg.io_dtype, + b_dtype=cfg.io_dtype, + n_dim=cfg.d_k, + m_dim=cfg.d_v, + b_major=1, + ) + bmm_state_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, + cta_group=1, + idesc=idesc_acc, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_diag_desc = MmaDesc( + M=cfg.d_v, + N=16, + K=16, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_diag, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_qk_desc = MmaDesc( + M=cfg.d_v, + N=cfg.b_t, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=False, + cta_group=1, + idesc=idesc_acc, + kind=nvvm.Tcgen05MMAKind.F16, + ) + bmm_final_state_desc = MmaDesc( + M=cfg.d_v, + N=cfg.d_k, + K=cfg.b_t, + bpe_a=bpe, + bpe_b=bpe, + tile_k_hw=16, + btranspose=True, + cta_group=1, + idesc=idesc_final_state, + kind=nvvm.Tcgen05MMAKind.F16, + ) + STATE_A_SEG = bmm_state_desc.sps_B * bmm_state_desc.tmem_advance_A + STATE_B_SEG = bmm_state_desc.smem_subtile_B >> 4 + cum_chunk_base = cutlass.Int32(0) + 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) + num_chunks_tile = wend - cstart + for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): + cum_chunk = cum_chunk_base + local_chunk_idx + have_state = cutlass.Boolean(True) if cutlass.const_expr(cfg.use_initial_state) else local_chunk_idx > 0 + decay_stage = k_decay_ready.idx + state_scale_diag_stage = qk_scale_index.idx + intermediate_stage = t_inv_ready.idx + sK_decay_stage = sK_decay[decay_stage] + sK_restore_stage = sK_restore[decay_stage] + sState_scale_diag_stage = sState_scale_diag[state_scale_diag_stage] + sIntermediate_stage = sIntermediate[intermediate_stage] + + # ---- State*K = State(T) @ K_decay^T ---------------------------------- + bars.mb_k_decay_inv_cg0_ready[decay_stage].wait(k_decay_ready.phase) + k_decay_ready = advance(k_decay_ready, cfg.smem_decay_stages) + if have_state: + bars.mb_state_inp_ready.wait(state_inp_index.phase) + state_inp_index = advance(state_inp_index, 1) + desc_k_decay = sK_decay_stage.desc() + + for s in cutlass.range_constexpr(bmm_state_desc.num_subtiles_B): + for k in cutlass.range_constexpr(bmm_state_desc.sps_B): + mma_ts_step( + bmm_state_desc, + state_inp_ptr.subview(s * STATE_A_SEG), + desc_k_decay + s * STATE_B_SEG, + state_k_acc_ptr, + k, + cutlass.Boolean(s + k > 0), + ) + + if elect_one: + bars.mb_state_k_acc_ready.arrive(cta_group=1) + + if elect_one: + bars.mb_decay_tcgen05_done[decay_stage].arrive(cta_group=1) + + bars.mb_qk_scale_ready[qk_scale_index.idx].wait(qk_scale_index.phase) + if cutlass.const_expr(cfg.enable_checkpoints): + if have_state: + bars.mb_state_acc_read_done.wait(state_read_index.phase) + state_read_index = advance(state_read_index, 1) + + # ---- State decay = State(T) @ exp2(g_last) diag (per-k-atom blocks) ---- + if have_state: + desc_diag = sState_scale_diag_stage.desc() + for k_block in cutlass.range_constexpr(cfg.d_k // 16): + mma_ts_step( + bmm_diag_desc, + state_inp_ptr.subview(k_block * bmm_diag_desc.tmem_advance_A), + desc_diag.advance_start_address(k_block * 256 * 2), + state_dsts[k_block], + 0, + cutlass.Boolean(False), + ) + + if elect_one: + bars.mb_state_scale_diag_done[state_scale_diag_stage].arrive(cta_group=1) + + # ---- U = Y(T) @ T_inv ------------------------------------------------ + bars.mb_t_inv_ready[intermediate_stage].wait(t_inv_ready.phase) + bars.mb_y_inp_ready.wait(y_inp_index.phase) + y_inp_index = advance(y_inp_index, 1) + desc_t_inv = sIntermediate_stage.shifted((cfg.b_t * cfg.b_t)).desc() + mma_ts_step(bmm_qk_desc, y_inp_ptr, desc_t_inv, u_acc_ptr, 0, cutlass.Boolean(False)) + if elect_one: + bars.mb_t_inv_done[intermediate_stage].arrive(cta_group=1) + bars.mb_u_acc_ready.arrive(cta_group=1) + + # ---- final_state += U(T) @ K_restore --------------------------------- + bars.mb_u_inp_ready.wait(u_inp_index.phase) + u_inp_index = advance(u_inp_index, 1) + desc_k_restore = sK_restore_stage.desc() + + mma_ts_step(bmm_final_state_desc, u_inp_ptr, desc_k_restore, state_dst_ptr, 0, have_state) + if elect_one: + bars.mb_k_restore_done[decay_stage].arrive(cta_group=1) + bars.mb_state_acc_done.arrive(cta_group=1) + + t_inv_ready = advance(t_inv_ready, cfg.smem_intermediate_stages) + qk_scale_index = advance(qk_scale_index, cfg.smem_state_scale_diag_stages) + + cum_chunk_base += num_chunks_tile + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + bars.mb_tmem_done[0].wait(0) + nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) + nvvm.tcgen05_dealloc( + nvvm.make_tmem_ptr(tmem_base, cutlass.Int8), + cutlass.Int32(512), + group=nvvm.CTAGroup.CTA_1, + ) + + +@cute.jit +def epilogue_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + sCheckpoint_raw, + desc_checkpoint_base, + checkpoint_every_n_tokens, + bars, +) -> None: + """Epilogue warp role (warp 15): persistent scheduler loop issuing the + per-chunk checkpoint TMA stores.""" + elect_one = nvvm.elect_sync() + nvvm.setmaxregister(cfg.num_regs_other, nvvm.SetMaxRegisterAction.DECREASE) + if cutlass.const_expr(cfg.enable_checkpoints): + sCheckpoint_tma = SmemTile( + base=sCheckpoint_raw, + elems_per_stage=(cfg.d_k * cfg.d_v), + stages=cfg.smem_checkpoint_stages, + leading_byte_offset=0, + stride_byte_offset=0, + layout=0, + tma_loads_per_tile=(cfg.d_v // 64), + tma_granu_elems=64, + tma_subtile_stride_elems=cfg.d_k * 64, + ) + checkpoint_ready_index = PipelineState.start(phase=0) + 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) + if cutlass.const_expr(cfg.enable_checkpoints): + head_o = head_idx + checkpoint_slot = batch_idx * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_checkpoint_slot = (desc_checkpoint_base + checkpoint_slot).tospace(cutlass.AddressSpace.generic) + checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + checkpoint_quot = (cstart + cutlass.Int32(1)) // checkpoint_chunks + checkpoint_mod = (cstart + cutlass.Int32(1)) % checkpoint_chunks + if elect_one: + tma_tensormap_acquire(desc_checkpoint_slot) + num_chunks_tile = wend - cstart + for local_chunk_idx in cutlass.range(num_chunks_tile, unroll=1): + chunk_idx = cstart + local_chunk_idx + if local_chunk_idx > 0: + # ---- checkpoint store ---------------------------------------- + do_checkpoint = checkpoint_mod == 0 + do_checkpoint = do_checkpoint and chunk_idx >= wstart + if do_checkpoint: + 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_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() + tma_store_wait(0) + bars.mb_checkpoint_tmastg_done[checkpoint_stage].arrive() + checkpoint_mod = checkpoint_mod + cutlass.Int32(1) + if checkpoint_mod == checkpoint_chunks: + checkpoint_mod = cutlass.Int32(0) + checkpoint_quot = checkpoint_quot + cutlass.Int32(1) + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def gate_scale(cfg, raw_gate: cutlass.Float32) -> cutlass.Float32: + """Map raw gate to the log2-domain decay increment used by KDA.""" + + if cutlass.const_expr(cfg.safe_gate): + half = cutlass.Float32(0.5) + sigmoid = cute.math.tanh(raw_gate * half, approx=True) * half + half + return cfg.gate_scale_log2 * sigmoid + # Default ABI: Gate arrives in natural-log space + return raw_gate * cutlass.Float32(LOG2_E) + + +@cute.jit +def compute0_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + warp_idx, + mA_log, + mDt_bias, + sK_inv_raw, + sGate_raw, + mBeta, + sBeta_raw, + sK_raw, + sK_decay_raw, + sK_restore_raw, + sState_scale_diag_raw, + bars, +) -> None: + """CG0 warp-group role (warps 0-7): persistent scheduler loop running the + Gate prefix scan and staging the decay/restore operands.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_0, nvvm.SetMaxRegisterAction.INCREASE) + cg0_warp = warp_idx - cfg.compute_group_0_warp_ids[0] + cg0_group_id = cg0_warp // cfg.cg0_warps_per_group + cg0_local_warp = cg0_warp % cfg.cg0_warps_per_group + prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + cg0_a_log_exp = cutlass.Float32(1.0) + cg0_dt_bias_value = cutlass.Float32(0.0) + cum_chunk_base = cutlass.Int32(0) + sched_state = PipelineState.start(phase=0) + tile_idx = cutlass.Int32(bidx) + opaque_one = opaque_f32_zero() + cutlass.Float32(1.0) + 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) + head_o = head_idx + num_chunks_tile = wend - cstart + if cutlass.const_expr(cfg.safe_gate): + if num_chunks_tile > 0: + cg0_a_log_exp = cute.math.exp2(mA_log[head_o].to(cutlass.Float32) * LOG2_E, fastmath=True) + cg0_dt_bias_value = mDt_bias[head_o, prefix_dim].to(cutlass.Float32) + # tile entry: both ping-pong groups inherit each other's delivery proofs (parity-swap guard) + nvvm.barrier_cta_sync(cfg.cg0_tile_entry_barrier_id, thread_count=cfg.cg0_group_count * cfg.cg0_threads_per_group) + group_cum_chunk_start = cum_chunk_base + cutlass.Int32(cg0_group_id) + diag_ring_stage = group_cum_chunk_start % cutlass.Int32(cfg.smem_state_scale_diag_stages) + diag_ring_phase = (group_cum_chunk_start // cutlass.Int32(cfg.smem_state_scale_diag_stages)) % cutlass.Int32(2) + for local_chunk_idx in cutlass.range(cg0_group_id, num_chunks_tile, cfg.cg0_group_count, unroll=1): + chunk_idx = cstart + local_chunk_idx + cum_chunk = cum_chunk_base + local_chunk_idx + chunk_start = chunk_idx * cfg.b_t + decay_stage = cum_chunk % cfg.smem_decay_stages + raw_stage = cum_chunk % cfg.smem_raw_stages + state_scale_diag_stage = diag_ring_stage + qk_scale_ready_stage = state_scale_diag_stage + sK_ptr = sK_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sGate_ptr = sGate_raw.data_ptr() + raw_stage * (cfg.d_k * cfg.b_t) + sK_inv_ptr = sK_inv_raw.data_ptr() + decay_stage * (cfg.b_t * cfg.d_k) + sK_decay_ptr = sK_decay_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sK_restore_ptr = sK_restore_raw.data_ptr() + decay_stage * (cfg.d_k * cfg.b_t) + sState_scale_diag_ptr = sState_scale_diag_raw.data_ptr() + state_scale_diag_stage * ((cfg.d_k // 16) * 256) + + # ---- Beta scalars --------------------------------------------------- + if cg0_local_warp == 0: + bars.mb_beta_done[raw_stage].wait(((cum_chunk // cfg.smem_raw_stages) + 1) % 2) + if lane < cfg.b_t: + token_idx = chunk_idx * cfg.b_t + lane + beta_value = cutlass.Float32(0.0) + if token_idx < seqlen_b: + beta_value = mBeta[batch_start + token_idx, head_o].to(cutlass.Float32) + if cutlass.const_expr(cfg.beta_sigmoid): + half = cutlass.Float32(0.5) + beta_value = (cute.math.tanh(beta_value * half, approx=True) * half + half).to(mBeta.element_type).to(cutlass.Float32) + sBeta_raw[raw_stage * cfg.b_t + lane] = beta_value + bars.mb_beta_ready[raw_stage].arrive() + bars.mb_gate_ready[raw_stage].wait((cum_chunk // cfg.smem_raw_stages) % 2) + + row_group_start = cg0_local_warp * (cfg.b_t // cfg.cg0_warps_per_group) + lane_row_group = lane // 8 + lane_in_row_group = lane - lane_row_group * 8 + decay_row = row_group_start + lane_row_group + decay_key_mask = cutlass.Int32(8) + + prefix_dim = cg0_local_warp * cfg.threads_per_warp + lane + + # ---- Gate prefix scan ----------------------------------------------- + f32_segment = prefix_dim // 32 + prefix_seg_base = f32_segment * (cfg.b_t * 32) + prefix_col = prefix_dim - f32_segment * 32 + gate_raw = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + for row in cutlass.range_constexpr(cfg.b_t): + prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) + gate_raw[row] = (sGate_ptr + prefix_idx).load() + g_prefix_regs = cutlass.Array(cutlass.Float32, cfg.b_t, alignment=16) + if cutlass.const_expr(cfg.safe_gate): + valid_rows = seqlen_b - chunk_idx * cutlass.Int32(cfg.b_t) + valid_mask = cutlass.vector.create_mask([cfg.b_t], [valid_rows]) + for row_pair in cutlass.range_constexpr(cfg.b_t // 2): + row0 = row_pair * 2 + row1 = row0 + 1 + gate0 = cg0_a_log_exp * (gate_raw[row0] + cg0_dt_bias_value) + gate1 = cg0_a_log_exp * (gate_raw[row1] + cg0_dt_bias_value) + gate0 = gate_scale( + cfg, + gate0, + ) + gate1 = gate_scale( + cfg, + gate1, + ) + gate_pair = cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32) + gate_pair = cutlass.vector.where(valid_mask[row0 : row1 + 1], gate_pair, 0.0) + g_prefix_regs[row0] = gate_pair[0] + g_prefix_regs[row1] = gate_pair[1] + else: + for row in cutlass.range_constexpr(cfg.b_t): + gate = gate_raw[row] + token_idx = chunk_idx * cutlass.Int32(cfg.b_t) + cutlass.Int32(row) + if token_idx < seqlen_b: + gate = gate_scale( + cfg, + gate, + ) + else: + gate = cutlass.Float32(0.0) + g_prefix_regs[row] = gate + + prefix_acc = cutlass.Float32(0.0) + for row_pair in cutlass.range_constexpr(cfg.b_t // 2): + row0 = row_pair * 2 + row1 = row0 + 1 + gate0 = g_prefix_regs[row0] + gate1 = g_prefix_regs[row1] + prefix0, row_pair_sum = fadd2(prefix_acc, gate0, gate0, gate1) + prefix1 = prefix_acc + row_pair_sum + g_prefix_regs[row0] = prefix0 + g_prefix_regs[row1] = prefix1 + prefix_acc = prefix1 + + # ---- exp2(g): stage prefixes + final-token decay --------------------- + for row in cutlass.range_constexpr(cfg.b_t): + g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True) + + exp_g_last = g_prefix_regs[cfg.b_t - 1] + for row in cutlass.range_constexpr(cfg.b_t): + prefix_idx = prefix_seg_base + swizzle_xor_128b(row, row * 32 + prefix_col, elem_bytes=4) + (sGate_ptr + prefix_idx).store(g_prefix_regs[row]) + + # ---- state-scale diag: stage exp2(g_last) decay blocks --------------- + bars.mb_state_scale_diag_done[state_scale_diag_stage].wait(diag_ring_phase ^ cutlass.Int32(1)) + block = prefix_dim // cutlass.Int32(16) + coord = prefix_dim - block * cutlass.Int32(16) + storage_col = coord ^ cutlass.Int32((cfg.b_t // 2)) + linear_idx = block * cutlass.Int32(256) + coord * cutlass.Int32(16) + storage_col + diag_idx = swizzle_lin_S(linear_idx, bbits=1, mbase=3, sshift=3) + sState_scale_diag_ptr[diag_idx] = exp_g_last.to(cfg.io_dtype) + + nvvm.barrier_cta_sync(cfg.cg0_group_sync_barrier_base_id + cg0_group_id, thread_count=cfg.cg0_threads_per_group) + + bars.mb_k_ready[raw_stage].wait((cum_chunk // cfg.smem_raw_stages) % 2) + k_inv_pack = cutlass.Array(cutlass.Int32, 2 * 4, alignment=16) + raw_k_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + + # ---- optional K L2-norm + K_inv staging ------------------------------ + if cutlass.const_expr(cfg.l2norm): + kk0_lo = opaque_f32_zero() + kk0_hi = opaque_f32_zero() + kk1_lo = opaque_f32_zero() + kk1_hi = opaque_f32_zero() + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + raw_f16_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + raw_k_frag = (sK_ptr + raw_f16_idx).load(count=8, alignment=16) + raw_k_vec_f32 = raw_k_frag.to(cutlass.Float32) + for dim_offset in cutlass.range_constexpr(8): + k_val = raw_k_vec_f32[dim_offset] + raw_k_regs[reg_base + dim_offset] = k_val + if cutlass.const_expr(cfg.l2norm): + if cutlass.const_expr(dim_offset % 2 == 0): + kk0_lo, kk0_hi = ffma2(k_val, k_val, k_val, k_val, kk0_lo, kk0_hi) + else: + kk1_lo, kk1_hi = ffma2(k_val, k_val, k_val, k_val, kk1_lo, kk1_hi) + + k_inv_norm = opaque_one + if cutlass.const_expr(cfg.l2norm): + k_sum_sq = kk0_hi + kk1_hi + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 4, 31, kind=nvvm.Shfl.BFLY)) + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 2, 31, kind=nvvm.Shfl.BFLY)) + k_sum_sq = k_sum_sq + cutlass.Float32(nvvm.shfl_sync(0xFFFFFFFF, k_sum_sq, 1, 31, kind=nvvm.Shfl.BFLY)) + norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS) + k_inv_norm = cute.math.rsqrt(cute.math.max(k_sum_sq, norm_floor_sq), fastmath=True) + + # ---- decay/restore operands: exp2(+-g) applied per key channel ------- + exp_g_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + exp_g_last_regs = cutlass.Array(cutlass.Float32, 2 * 8, alignment=16) + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + for f32_group in cutlass.range_constexpr(2): + f32_dim_base = dim_base + f32_group * 4 + f32_segment = f32_dim_base // 32 + f32_segment_dim = f32_dim_base - f32_segment * 32 + g_prefix_idx = f32_segment * (cfg.b_t * 32) + decay_row * 32 + swizzle_xor_128b(decay_row, f32_segment_dim, elem_bytes=4) + exp_g_frag = (sGate_ptr + g_prefix_idx).load(count=4, alignment=16) + exp_g_last_idx = f32_segment * (cfg.b_t * 32) + (cfg.b_t - 1) * 32 + swizzle_xor_128b((cfg.b_t - 1), f32_segment_dim, elem_bytes=4) + exp_g_last_frag = (sGate_ptr + exp_g_last_idx).load(count=4, alignment=16) + f32_reg_base = reg_base + f32_group * 4 + for j in cutlass.range_constexpr(4): + exp_g_regs[f32_reg_base + j] = exp_g_frag[j] + exp_g_last_regs[f32_reg_base + j] = exp_g_last_frag[j] + + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + + # ---- K decay + K_inv operands: K * exp2(+g) and K * exp2(-g) ----- + k_decay_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + for pair_idx in cutlass.range_constexpr(4): + dim0 = pair_idx * 2 + dim1 = dim0 + 1 + raw_reg_idx0 = reg_base + dim0 + raw_reg_idx1 = reg_base + dim1 + k_value0, k_value1 = fmul2(raw_k_regs[raw_reg_idx0], raw_k_regs[raw_reg_idx1], k_inv_norm, k_inv_norm) + k_pair = fp32_to_fp16(k_value0, k_value1, dtype=cfg.io_dtype) + exp_g_pair = fp32_to_fp16(exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1], dtype=cfg.io_dtype) + k_decay_pack[pair_idx] = mul_f16x2(k_pair, exp_g_pair, cfg.io_dtype) + exp_neg_g0 = cute.math.rcp(exp_g_regs[raw_reg_idx0], approx=True, ftz=True) + exp_neg_g1 = cute.math.rcp(exp_g_regs[raw_reg_idx1], approx=True, ftz=True) + exp_neg_pair = fp32_to_fp16(exp_neg_g0, exp_neg_g1, dtype=cfg.io_dtype) + k_inv_pack[dim_half * 4 + pair_idx] = mul_f16x2(k_pair, exp_neg_pair, cfg.io_dtype) + + k_inv_vec = cutlass.Vector.from_elements( + ( + k_inv_pack[dim_half * 4], + k_inv_pack[dim_half * 4 + 1], + k_inv_pack[dim_half * 4 + 2], + k_inv_pack[dim_half * 4 + 3], + ), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + k_decay_vec = cutlass.Vector.from_elements( + ( + k_decay_pack[0], + k_decay_pack[1], + k_decay_pack[2], + k_decay_pack[3], + ), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + if cutlass.const_expr(dim_half == 0): + operand_done_phase = ((cum_chunk // cfg.smem_decay_stages) + 1) % 2 + bars.mb_decay_super_done[decay_stage].wait(operand_done_phase) + bars.mb_decay_tcgen05_done[decay_stage].wait(operand_done_phase) + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + k_inv_swizzled_idx = f16_segment * (cfg.b_t * 64) + decay_row * 64 + swizzle_xor_128b(decay_row, f16_segment_dim, elem_bytes=2) + (sK_inv_ptr + k_inv_swizzled_idx).store(k_inv_vec, alignment=16) + storage_key = dim_base ^ decay_key_mask + storage_slice = storage_key // 64 + decay_swizzled_idx = storage_slice * (cfg.b_t * 64) + swizzle_xor_128b( + decay_row, decay_row * 64 + storage_key - storage_slice * 64, elem_bytes=2 + ) + (sK_decay_ptr + decay_swizzled_idx).store(k_decay_vec, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_k_decay_inv_cg0_ready[decay_stage].arrive() + + # ---- K_restore operand: K_inv * exp_g_last -------------------------- + bars.mb_k_restore_done[decay_stage].wait(((cum_chunk // cfg.smem_decay_stages + 1) % 2)) + for dim_half in cutlass.range_constexpr(2): + dim_base = dim_half * (cfg.d_k // 2) + lane_in_row_group * 8 + reg_base = dim_half * 8 + k_restore_pack = cutlass.Array(cutlass.Int32, 4, alignment=16) + for pair_idx in cutlass.range_constexpr(4): + dim0 = pair_idx * 2 + dim1 = dim0 + 1 + exp_g_last_pair = fp32_to_fp16(exp_g_last_regs[reg_base + dim0], exp_g_last_regs[reg_base + dim1], dtype=cfg.io_dtype) + k_restore_pack[pair_idx] = mul_f16x2(k_inv_pack[dim_half * 4 + pair_idx], exp_g_last_pair, cfg.io_dtype) + storage_row = decay_row ^ (cfg.b_t // 2) + f16_segment = dim_base // 64 + f16_segment_dim = dim_base - f16_segment * 64 + k_restore_idx = f16_segment * (cfg.b_t * 64) + storage_row * 64 + swizzle_xor_128b(storage_row, f16_segment_dim, elem_bytes=2) + k_restore_vec = cutlass.Vector.from_elements( + ( + k_restore_pack[0], + k_restore_pack[1], + k_restore_pack[2], + k_restore_pack[3], + ), + cutlass.Int32, + ).bitcast(cfg.io_dtype) + (sK_restore_ptr + k_restore_idx).store(k_restore_vec, alignment=16) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_qk_scale_ready[qk_scale_ready_stage].arrive() + bars.mb_k_done[raw_stage].arrive() + bars.mb_gate_done[raw_stage].arrive() + diag_ring_stage = diag_ring_stage + cutlass.Int32(cfg.cg0_group_count) + wrapped = diag_ring_stage >= cutlass.Int32(cfg.smem_state_scale_diag_stages) + diag_ring_stage = diag_ring_stage - cutlass.Int32(cfg.smem_state_scale_diag_stages) if wrapped else diag_ring_stage + diag_ring_phase = diag_ring_phase ^ (cutlass.Int32(1) if wrapped else cutlass.Int32(0)) + cum_chunk_base += num_chunks_tile + tile_idx, sched_state = sched_next_tile(cfg, bars, sSched, sched_state, tile_idx, num_ctas) + + +@cute.jit +def compute1_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sTmem_base, + warp_idx, + mState_out, + mState_init, + sBeta_raw, + sV_raw, + sCheckpoint_raw, + checkpoint_every_n_tokens, + bars, +) -> None: + """CG1 warp-group role (warps 8-11): persistent scheduler loop staging the + value-side TMEM operands and storing the checkpoint/final states.""" + nvvm.setmaxregister(cfg.num_regs_compute_group_1, nvvm.SetMaxRegisterAction.INCREASE) + sCheckpoint_ptr = sCheckpoint_raw.data_ptr() if cutlass.const_expr(cfg.enable_checkpoints) else sV_raw.data_ptr() + checkpoint_done_index = PipelineState.start(phase=1) + nvvm.barrier_cta_sync(cfg.tmem_lifecycle_barrier_id, thread_count=cfg.tmem_user_threads) + tmem_base = sTmem_base.load() + tmem_col = tmem_base & 0xFFFF + tmem_row = tmem_base >> 16 + tmem_subpartition = warp_idx % (cfg.d_v // cfg.threads_per_warp) + # ldmatrix.x4 COL lane decode for the V loads + frag_row_coord = (lane // 16) * 8 + (lane & 7) + frag_col_offset = ((lane // 8) & 1) * 8 + row_id = tmem_row + tmem_subpartition * cfg.threads_per_warp + value_dim = tmem_subpartition * cfg.threads_per_warp + lane + value_dim_base = tmem_subpartition * cfg.threads_per_warp + row_addr = row_id << 16 + row16_addr = (row_id + 16) << 16 + st_row_addr = tmem_row << 16 + st_row16_addr = (tmem_row + 16) << 16 + state_col_id = tmem_col + cfg.tmem_state_acc_offset + packed_col_id = tmem_col + cfg.tmem_state_inp_offset + statek_col_id = tmem_col + cfg.tmem_state_k_acc_offset + y_inp_col_id = tmem_col + cfg.tmem_y_inp_offset + u_acc_addr = row_addr + tmem_col + cfg.tmem_u_acc_offset + u_inp_addr = st_row_addr + tmem_col + cfg.tmem_u_inp_offset + v_swz_off0 = ( + (value_dim_base + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + frag_col_offset) % 64, elem_bytes=2) + ) + v_swz_off = ( + (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.b_t * 64) + + frag_row_coord * 64 + + swizzle_xor_128b(frag_row_coord, (value_dim_base + 16 + frag_col_offset) % 64, elem_bytes=2) + ) + checkpoint_swz_off0 = (value_dim_base + frag_col_offset) // 64 * (cfg.d_k * 64) + checkpoint_swz_col0 = (value_dim_base + frag_col_offset) % 64 + checkpoint_swz_off = (value_dim_base + 16 + frag_col_offset) // 64 * (cfg.d_k * 64) + checkpoint_swz_col = (value_dim_base + 16 + frag_col_offset) % 64 + state_k_acc_index = PipelineState.start(phase=0) + u_acc_index = PipelineState.start(phase=0) + state_upd_index = PipelineState.start(phase=0) + raw_index = PipelineState.start(phase=0) + 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) + head_o = head_idx + num_chunks_tile = wend - cstart + + if num_chunks_tile > 0: + # ---- first chunk: seed state TMEM from mState_init ---------- + seed_from_initial_state = cstart == 0 + if cutlass.const_expr(mState_init is not None): + if seed_from_initial_state: + 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) + + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + state_block[0:32], + ) + else: + 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): + state_block[col] = cutlass.Float32(0.0) + + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + state_block[0:32], + ) + if cutlass.const_expr(mState_init is not None): + nvvm.tcgen05_wait("store") + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sBeta_ptr = sBeta_raw.data_ptr() + raw_index.idx * cfg.b_t + + # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + if cutlass.const_expr(mState_init is not None): + state_vecs = [] + for sub in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) + + for sub in cutlass.range_constexpr(cfg.d_k // 16): + state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) + for packed_col in cutlass.range_constexpr(8): + source_pair = packed_col ^ 4 + state_pack[packed_col] = fp32_to_fp16(state_vecs[sub][2 * source_pair], state_vecs[sub][2 * source_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), + state_pack[0:8], + ) + + nvvm.tcgen05_wait("store") + bars.mb_state_inp_ready.arrive() + if cutlass.const_expr(cfg.enable_checkpoints): + bars.mb_state_acc_read_done.arrive() + + # ---- Y staging: Y = Beta * (V - State*K) ----------------------------- + bars.mb_v_ready[raw_index.idx].wait(raw_index.phase) + raw_v_frag0 = nvvm.ldmatrix( + sV_ptr + v_swz_off0, + 4, + nvvm.MMALayout.COL, + ) + raw_v_frag1 = nvvm.ldmatrix( + sV_ptr + v_swz_off, + 4, + nvvm.MMALayout.COL, + ) + bars.mb_beta_ready[raw_index.idx].wait(raw_index.phase) + if cutlass.const_expr(mState_init is not None): + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + + state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + statek_col_id, cutlass.Float32), num=2) + state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + statek_col_id, cutlass.Float32), num=2) + + beta_pack = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + token0 = (((reg_idx // 2) * 4 + (lane & 3)) ^ 4) * 2 + beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) + beta1 = (sBeta_ptr + token0 + 1).load().to(cutlass.Float32) + beta_pack[reg_idx] = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) + y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + if cutlass.const_expr(mState_init is not None): + state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + diff_pair = sub_f16x2( + raw_v_frag0[raw_matrix], + state_k_pair, + cfg.io_dtype, + ) + else: + diff_pair = raw_v_frag0[raw_matrix] + y_inp_pack0[reg_idx] = mul_f16x2( + beta_pack[reg_idx], + diff_pair, + cfg.io_dtype, + ) + + y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + if cutlass.const_expr(mState_init is not None): + state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + diff_pair = sub_f16x2( + raw_v_frag1[raw_matrix], + state_k_pair, + cfg.io_dtype, + ) + else: + diff_pair = raw_v_frag1[raw_matrix] + y_inp_pack1[reg_idx] = mul_f16x2( + beta_pack[reg_idx], + diff_pair, + cfg.io_dtype, + ) + + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row_addr + y_inp_col_id, cutlass.Int8), y_inp_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row16_addr + y_inp_col_id, cutlass.Int8), y_inp_pack1[0:4]) + nvvm.tcgen05_wait("store") + if cutlass.const_expr(mState_init is not None): + state_k_acc_index = advance(state_k_acc_index, 1) + bars.mb_v_done[raw_index.idx].arrive() + bars.mb_beta_done[raw_index.idx].arrive() + bars.mb_y_inp_ready.arrive() + + # ---- U repack: u_acc TMEM -> packed b16 (U input) TMEM -------------- + bars.mb_u_acc_ready.wait(u_acc_index.phase) + u_acc_vals = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr(u_acc_addr, cutlass.Float32), + num=cfg.b_t, + ) + + u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): + source_pair = packed_col ^ 4 + token0 = source_pair * 2 + token1 = token0 + 1 + u_inp_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) + + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(u_inp_addr, cutlass.Int8), + u_inp_pack[0 : (cfg.b_t // 2)], + ) + nvvm.tcgen05_wait("store") + u_acc_index = advance(u_acc_index, 1) + bars.mb_u_inp_ready.arrive() + if cutlass.const_expr(cfg.enable_checkpoints): + bars.mb_state_acc_done.wait(state_upd_index.phase) + state_upd_index = advance(state_upd_index, 1) + + raw_index = advance(raw_index, cfg.smem_raw_stages) + + if cutlass.const_expr(cfg.enable_checkpoints): + cg1_checkpoint_chunks = checkpoint_every_n_tokens // cutlass.Int32(cfg.b_t) + cg1_checkpoint_mod = (cstart + cutlass.Int32(1)) % cg1_checkpoint_chunks + for local_chunk_idx in cutlass.range(1, num_chunks_tile, 1, unroll=1): + chunk_idx = cstart + local_chunk_idx + sV_ptr = sV_raw.data_ptr() + raw_index.idx * (cfg.d_v * cfg.b_t) + sBeta_ptr = sBeta_raw.data_ptr() + raw_index.idx * cfg.b_t + + do_checkpoint = False + if cutlass.const_expr(cfg.enable_checkpoints): + do_checkpoint = cg1_checkpoint_mod == 0 + cg1_checkpoint_mod = cg1_checkpoint_mod + cutlass.Int32(1) + cg1_checkpoint_mod = cutlass.Int32(0) if cg1_checkpoint_mod == cg1_checkpoint_chunks else cg1_checkpoint_mod + do_checkpoint = do_checkpoint and chunk_idx >= wstart + + # ---- state repack: acc TMEM -> packed b16 TMEM ---------------------- + if cutlass.const_expr(not cfg.enable_checkpoints): + bars.mb_state_acc_done.wait(state_upd_index.phase) + state_upd_index = advance(state_upd_index, 1) + state_vecs = [] + for sub in cutlass.range_constexpr(cfg.d_k // 16): + state_vecs.append(nvvm.tcgen05_ld("32x32b", nvvm.make_tmem_ptr(row_addr + state_col_id + sub * 16, cutlass.Float32), num=16)) + + for sub in cutlass.range_constexpr(cfg.d_k // 16): + state_pack = cutlass.Array(cutlass.Int32, 8, alignment=16) + for packed_col in cutlass.range_constexpr(8): + source_pair = packed_col ^ 4 + state_pack[packed_col] = fp32_to_fp16(state_vecs[sub][2 * source_pair], state_vecs[sub][2 * source_pair + 1], dtype=cfg.io_dtype) + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr((tmem_row << 16) + packed_col_id + sub * 8, cutlass.Int8), + state_pack[0:8], + ) + nvvm.tcgen05_wait("store") + bars.mb_state_inp_ready.arrive() + + # ---- checkpoint store ----------------------------------------------- + if cutlass.const_expr(cfg.enable_checkpoints): + if do_checkpoint: + 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 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, + ) + nvvm.fence_proxy("async.shared", space="cta") + bars.mb_checkpoint_tmastg_ready[checkpoint_stage].arrive() + nvvm.tcgen05_wait("load") + bars.mb_state_acc_read_done.arrive() + else: + bars.mb_state_acc_read_done.arrive() + + # ---- Y staging: Y = Beta * (V - State*K) ----------------------------- + bars.mb_v_ready[raw_index.idx].wait(raw_index.phase) + raw_v_frag0 = nvvm.ldmatrix( + sV_ptr + v_swz_off0, + 4, + nvvm.MMALayout.COL, + ) + raw_v_frag1 = nvvm.ldmatrix( + sV_ptr + v_swz_off, + 4, + nvvm.MMALayout.COL, + ) + bars.mb_beta_ready[raw_index.idx].wait(raw_index.phase) + + # ---- read back State*K acc ------------------------------------------- + bars.mb_state_k_acc_ready.wait(state_k_acc_index.phase) + state_k_vec0 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row_addr + statek_col_id, cutlass.Float32), num=2) + state_k_vec1 = nvvm.tcgen05_ld("16x256b", nvvm.make_tmem_ptr(row16_addr + statek_col_id, cutlass.Float32), num=2) + + beta_pack = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + token0 = (((reg_idx // 2) * 4 + (lane & 3)) ^ 4) * 2 + beta0 = (sBeta_ptr + token0).load().to(cutlass.Float32) + beta1 = (sBeta_ptr + token0 + 1).load().to(cutlass.Float32) + beta_pack[reg_idx] = fp32_to_fp16(beta0, beta1, dtype=cfg.io_dtype) + y_inp_pack0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + state_k_val0, state_k_val1 = state_k_vec0[frag_pair], state_k_vec0[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + diff_pair = sub_f16x2( + raw_v_frag0[raw_matrix], + state_k_pair, + cfg.io_dtype, + ) + y_inp_pack0[reg_idx] = mul_f16x2( + beta_pack[reg_idx], + diff_pair, + cfg.io_dtype, + ) + + y_inp_pack1 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1) + frag_pair = (reg_idx ^ 2) * 2 + state_k_val0, state_k_val1 = state_k_vec1[frag_pair], state_k_vec1[frag_pair + 1] + state_k_pair = fp32_to_fp16(state_k_val0, state_k_val1, dtype=cfg.io_dtype) + diff_pair = sub_f16x2( + raw_v_frag1[raw_matrix], + state_k_pair, + cfg.io_dtype, + ) + y_inp_pack1[reg_idx] = mul_f16x2( + beta_pack[reg_idx], + diff_pair, + cfg.io_dtype, + ) + + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row_addr + y_inp_col_id, cutlass.Int8), y_inp_pack0[0:4]) + nvvm.tcgen05_st("16x128b", nvvm.make_tmem_ptr(st_row16_addr + y_inp_col_id, cutlass.Int8), y_inp_pack1[0:4]) + nvvm.tcgen05_wait("store") + state_k_acc_index = advance(state_k_acc_index, 1) + bars.mb_v_done[raw_index.idx].arrive() + bars.mb_beta_done[raw_index.idx].arrive() + bars.mb_y_inp_ready.arrive() + + # ---- U repack: u_acc TMEM -> packed b16 (U input) TMEM -------------- + bars.mb_u_acc_ready.wait(u_acc_index.phase) + u_acc_vals = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr(u_acc_addr, cutlass.Float32), + num=cfg.b_t, + ) + + u_inp_pack = cutlass.Array(cutlass.Int32, (cfg.b_t // 2), alignment=16) + for packed_col in cutlass.range_constexpr((cfg.b_t // 2)): + source_pair = packed_col ^ 4 + token0 = source_pair * 2 + token1 = token0 + 1 + u_inp_pack[packed_col] = fp32_to_fp16(u_acc_vals[token0], u_acc_vals[token1], dtype=cfg.io_dtype) + + nvvm.tcgen05_st( + "32x32b", + nvvm.make_tmem_ptr(u_inp_addr, cutlass.Int8), + u_inp_pack[0 : (cfg.b_t // 2)], + ) + nvvm.tcgen05_wait("store") + u_acc_index = advance(u_acc_index, 1) + bars.mb_u_inp_ready.arrive() + + if cutlass.const_expr(cfg.enable_checkpoints): + bars.mb_state_acc_done.wait(state_upd_index.phase) + state_upd_index = advance(state_upd_index, 1) + raw_index = advance(raw_index, cfg.smem_raw_stages) + + if num_chunks_tile > 0: + if cutlass.const_expr(not cfg.enable_checkpoints): + bars.mb_state_acc_done.wait(state_upd_index.phase) + state_upd_index = advance(state_upd_index, 1) + + owns_final = wend == num_chunks_b + + # ---- final-state drain: state acc TMEM -> GMEM --------------------------- + if cutlass.const_expr(mState_out is not None): + if seqlen_b > 0: + if owns_final: + for key_block_start in cutlass.range_constexpr(0, cfg.d_k, 32): + loaded = nvvm.tcgen05_ld( + "32x32b", + nvvm.make_tmem_ptr((row_id << 16) + (tmem_col + cfg.tmem_state_acc_offset + key_block_start), cutlass.Float32), + 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) + 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] + else: + mState_out[batch_idx, head_o, key_dim, value_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() + + +@cute.jit +def host( + cfg: cutlass.Constexpr, + k: cute.Tensor, + v: cute.Tensor, + raw_gate: cute.Tensor, + a_log: cute.Tensor | None, + dt_bias: cute.Tensor | None, + beta: cute.Tensor, + cu_seqlens: cute.Tensor, + initial_state: cute.Tensor | None, + final_state: cute.Tensor | None, + work_items: cute.Tensor | None, + work_count: cute.Tensor | None, + sched_ctr: cute.Tensor | None, + tensormap_workspace: cute.Tensor, + checkpoint_every_n_tokens: cutlass.Int32, + stream, +) -> None: + num_sequences = cu_seqlens.shape[0] - 1 + grid_shape = (cfg.max_active_clusters, 1, 1) + kernel( + cfg, + tensormap_workspace, + cutlass.Int32(num_sequences), + k, + v, + raw_gate, + a_log, + dt_bias, + beta, + cu_seqlens, + initial_state, + final_state, + work_items, + work_count, + sched_ctr, + checkpoint_every_n_tokens, + ).launch( + grid=grid_shape, + block=(cfg.threads_per_cta, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + +@cute.kernel +def kernel( + cfg: cutlass.Constexpr, + tensormap_workspace: cute.Tensor, + n_desc: cutlass.Int32, + mK: cute.Tensor, + mV: cute.Tensor, + mGate: cute.Tensor, + mA_log: cute.Tensor | None, + mDt_bias: cute.Tensor | None, + mBeta: cute.Tensor, + cu_seqlens: cute.Tensor, + mState_init: cute.Tensor | None, + mState_out: cute.Tensor | None, + mWorkItems: cute.Tensor, + mCount: cute.Tensor, + mSched: cute.Tensor | None, + checkpoint_every_n_tokens: cutlass.Int32, +) -> None: + """BT=16 KDA recompute (state/checkpoints-only) persistent kernel body: every warp + role runs a tile-scheduler loop over the tiles.""" + + tidx, _, _ = cute.arch.thread_idx() + bidx = cute.arch.block_idx()[0] + num_ctas = cute.arch.grid_dim()[0] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % cfg.threads_per_warp + + total_tiles = mCount[0] + if cutlass.const_expr(cfg.dyn_sched): + assert mSched is not None and mSched.element_type == cutlass.Int32 + assert mK.element_type == cfg.io_dtype and mV.element_type == cfg.io_dtype + assert mGate.element_type == cutlass.Float32 + beta_expected = cfg.io_dtype if cutlass.const_expr(cfg.beta_sigmoid) else cutlass.Float32 + assert mBeta.element_type == beta_expected + assert cu_seqlens.element_type in (cutlass.Int32, cutlass.Int64) + if cutlass.const_expr(cfg.use_initial_state): + assert mState_init is not None and mState_init.element_type in (cutlass.BFloat16, cutlass.Float32) + else: + assert mState_init is None, "mState_init must be None if use_initial_state is False" + if cutlass.const_expr(cfg.store_final_state): + assert mState_out is not None and mState_out.element_type in (cutlass.BFloat16, cutlass.Float32) + else: + assert mState_out is None, "mState_out must be None if store_final_state is False" + if cutlass.const_expr(mState_init is not None and mState_out is not None): + assert mState_init.element_type == mState_out.element_type + desc_base_words = tensormap_workspace.iterator.raw_ptr() + arr_words = n_desc * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_k_base = desc_base_words + desc_v_base = desc_base_words + arr_words + desc_gate_base = desc_base_words + cutlass.Int32(2) * arr_words + desc_checkpoint_base = desc_base_words + cutlass.Int32(3) * arr_words + + # Buffers are declaration-ordered and intentionally non-aliased. + SMEM = cutlass.AddressSpace.smem + bars = make_kda_bars(cfg) + sTmem_base = cutlass.Array(cutlass.Int32, 1, space=SMEM, alignment=4) + sSched = cutlass.Array(cutlass.Int32, cfg.sched_stages, space=SMEM, alignment=16) + sK_decay_raw = cutlass.Array(cfg.io_dtype, cfg.k_decay_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_restore_raw = cutlass.Array(cfg.io_dtype, cfg.k_restore_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sIntermediate_raw = cutlass.Array(cfg.io_dtype, cfg.intermediate_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_raw = cutlass.Array(mK.element_type, cfg.k_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sV_raw = cutlass.Array(mV.element_type, cfg.v_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sGate_raw = cutlass.Array(cutlass.Float32, cfg.gate_cosize, space=SMEM, alignment=1024) + sState_scale_diag_raw = cutlass.Array(cfg.io_dtype, cfg.state_scale_diag_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sK_inv_raw = cutlass.Array(cfg.io_dtype, cfg.k_inv_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sBeta_raw = cutlass.Array(cutlass.Float32, cfg.beta_cosize, space=SMEM, alignment=cfg.buffer_align_bytes) + sCheckpoint_raw = ( + cutlass.Array(cfg.io_dtype, cfg.smem_checkpoint_stages * cfg.d_k * cfg.d_v, space=SMEM, alignment=cfg.buffer_align_bytes) + if cutlass.const_expr(cfg.enable_checkpoints) + else sV_raw + ) + sK_decay = SmemTile( + base=sK_decay_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_decay_stages, + leading_byte_offset=16, + stride_byte_offset=1024, + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) + sK_restore = SmemTile( + base=sK_restore_raw, + elems_per_stage=(cfg.d_k * cfg.b_t), + stages=cfg.smem_decay_stages, + leading_byte_offset=(cfg.b_t * (cfg.d_v // 2) * 2), + stride_byte_offset=(8 * (cfg.d_v // 2) * 2), + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) + sState_scale_diag = SmemTile( + base=sState_scale_diag_raw, + elems_per_stage=((cfg.d_k // 16) * 256), + stages=cfg.smem_state_scale_diag_stages, + leading_byte_offset=16, + stride_byte_offset=(8 * 16 * 2), + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, + ) + sIntermediate = SmemTile( + base=sIntermediate_raw, + elems_per_stage=(2 * cfg.b_t * cfg.b_t), + stages=cfg.smem_intermediate_stages, + leading_byte_offset=16, + stride_byte_offset=(8 * cfg.b_t * 2), + layout=nvvm.Tcgen05SmemSwizzle.SWIZZLE_32B, + ) + + elect_one = nvvm.elect_sync() + if warp_idx == cfg.tma_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_raw_stages): + bars.mb_k_ready[stage].init() + bars.mb_v_ready[stage].init() + bars.mb_gate_ready[stage].init() + bars.mb_beta_ready[stage].init() + bars.mb_beta_done[stage].init() + bars.mb_k_done[stage].init() + bars.mb_v_done[stage].init() + bars.mb_gate_done[stage].init() + elif warp_idx == cfg.tcgen05_mma_warp_id: + if elect_one: + bars.mb_state_k_acc_ready.init() + bars.mb_u_acc_ready.init() + bars.mb_state_acc_done.init() + bars.mb_state_inp_ready.init() + for stage in cutlass.range_constexpr(cfg.smem_state_scale_diag_stages): + bars.mb_state_scale_diag_done[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_decay_stages): + bars.mb_decay_tcgen05_done[stage].init() + bars.mb_decay_super_done[stage].init() + bars.mb_k_restore_done[stage].init() + bars.mb_y_inp_ready.init() + bars.mb_u_inp_ready.init() + bars.mb_tmem_done[0].init() + elif warp_idx == cfg.super_mma_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.smem_intermediate_stages): + bars.mb_t_inv_ready[stage].init() + bars.mb_t_inv_done[stage].init() + for stage in cutlass.range_constexpr(cfg.qk_scale_ready_stages): + bars.mb_qk_scale_ready[stage].init() + for stage in cutlass.range_constexpr(cfg.smem_decay_stages): + bars.mb_k_decay_inv_cg0_ready[stage].init() + elif warp_idx == cfg.epilogue_warp_id: + if elect_one: + for stage in cutlass.range_constexpr(cfg.sched_stages): + bars.mb_sched_ready[stage].init() + bars.mb_sched_done[stage].init() + if cutlass.const_expr(cfg.enable_checkpoints): + for stage in cutlass.range_constexpr(cfg.smem_checkpoint_stages): + bars.mb_checkpoint_tmastg_ready[stage].init() + bars.mb_checkpoint_tmastg_done[stage].init() + bars.mb_state_acc_read_done.init() + diag_zero = cfg.io_dtype(0.0) + for diag_idx in cutlass.range(tidx, cfg.state_scale_diag_cosize, cfg.threads_per_cta, unroll=1): + sState_scale_diag_raw[diag_idx] = diag_zero + nvvm.fence_mbarrier_init() + nvvm.barrier_cta_sync(0, thread_count=cfg.threads_per_cta) + if warp_idx == cfg.tma_warp_id: + tmaldg_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + mSched, + sSched, + lane, + sK_raw, + sV_raw, + sGate_raw, + desc_k_base, + desc_v_base, + desc_gate_base, + bars, + ) + elif warp_idx == cfg.super_mma_warp_id: + super_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sK_inv_raw, + sIntermediate_raw, + sBeta_raw, + sK_decay_raw, + bars, + ) + elif warp_idx == cfg.tcgen05_mma_warp_id: + tcgen05_mma_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + sTmem_base, + sIntermediate, + sK_decay, + sK_restore, + sState_scale_diag, + bars, + ) + elif warp_idx == cfg.epilogue_warp_id: + epilogue_warp( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + sCheckpoint_raw, + desc_checkpoint_base, + checkpoint_every_n_tokens, + bars, + ) + elif warp_idx >= cfg.compute_group_0_warp_ids[0] and warp_idx <= cfg.compute_group_0_warp_ids[-1]: + compute0_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + warp_idx, + mA_log, + mDt_bias, + sK_inv_raw, + sGate_raw, + mBeta, + sBeta_raw, + sK_raw, + sK_decay_raw, + sK_restore_raw, + sState_scale_diag_raw, + bars, + ) + elif warp_idx >= cfg.compute_group_1_warp_ids[0] and warp_idx <= cfg.compute_group_1_warp_ids[-1]: + compute1_warp_group( + cfg, + total_tiles, + bidx, + num_ctas, + cu_seqlens, + mWorkItems, + sSched, + lane, + sTmem_base, + warp_idx, + mState_out, + mState_init, + sBeta_raw, + sV_raw, + sCheckpoint_raw, + checkpoint_every_n_tokens, + bars, + ) + + +@dataclass +class KdaRecomputeCfg: + """Kernel cfg (fixed BT=16 schedule constants; derived TMEM column offsets + and SMEM buffer cosizes are stamped by ``build_cfg``; per-stage sizes are + inlined at the use sites). Passed ``cfg``-first (a ``cutlass.Constexpr``) + into ``host`` / ``kernel`` and every warp body.""" + + io_dtype: Type[cutlass.Numeric] + state_dtype: Type[cutlass.Numeric] + use_initial_state: bool + store_final_state: bool + enable_checkpoints: bool + l2norm: bool + safe_gate: bool + gate_scale_log2: float + beta_sigmoid: bool + k_ratio: int + v_ratio: int + n_heads_out: int + max_active_clusters: int + dyn_sched: bool = False + sched_stages: int = CFG.SMEM_SCHED_STAGES + + compute_group_0_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_0_WARP_IDS + compute_group_1_warp_ids: tuple[int, ...] = CFG.COMPUTE_GROUP_1_WARP_IDS + super_mma_warp_id: int = CFG.SUPER_MMA_WARP_ID + tcgen05_mma_warp_id: int = CFG.TCGEN05_MMA_WARP_ID + tma_warp_id: int = CFG.TMA_WARP_ID + epilogue_warp_id: int = CFG.EPILOGUE_WARP_ID + b_t: int = CFG.B_T + d_k: int = CFG.D_K + d_v: int = CFG.D_V + threads_per_warp: int = CFG.THREADS_PER_WARP + buffer_align_bytes: int = CFG.BUFFER_ALIGN_BYTES + threads_per_cta: int = 0 + cg0_group_count: int = 2 + cg0_warps_per_group: int = 4 + cg0_threads_per_group: int = 0 + cg0_group_sync_barrier_base_id: int = 1 # CG0 group g syncs on nbar id 1 + g + cg0_tile_entry_barrier_id: int = 5 # CG0-wide (both groups) work-item entry sync + tmem_user_threads: int = 0 + tmem_lifecycle_barrier_id: int = 3 + num_regs_compute_group_0: int = CFG.NUM_REGS_COMPUTE_GROUP_0 + num_regs_compute_group_1: int = CFG.NUM_REGS_COMPUTE_GROUP_1 + num_regs_other: int = CFG.NUM_REGS_OTHER + + # ---- SMEM / TMEM ring stage counts ------------------------------------------- + smem_raw_stages: int = CFG.SMEM_RAW_STAGES + smem_checkpoint_stages: int = 1 + smem_decay_stages: int = CFG.SMEM_DECAY_STAGES + smem_intermediate_stages: int = CFG.SMEM_INTERMEDIATE_STAGES + smem_state_scale_diag_stages: int = CFG.SMEM_STATE_SCALE_DIAG_STAGES + qk_scale_ready_stages: int = CFG.QK_SCALE_READY_STAGES + + # ---- TMEM column offsets (state doubles as the final_state acc) -------------- + tmem_state_acc_offset: int = 0 + tmem_state_inp_offset: int = 0 + tmem_state_k_acc_offset: int = 0 + tmem_u_acc_offset: int = 0 + tmem_y_inp_offset: int = 0 + tmem_u_inp_offset: int = 0 + + # ---- SMEM buffer cosizes ----------------------------------------------------- + k_cosize: int = 0 + v_cosize: int = 0 + gate_cosize: int = 0 + beta_cosize: int = 0 + k_inv_cosize: int = 0 + k_decay_cosize: int = 0 + k_restore_cosize: int = 0 + state_scale_diag_cosize: int = 0 + + # TMA transaction bytes per stage + tma_k_bytes: int = 0 + tma_v_bytes: int = 0 + tma_gate_bytes: int = 0 + intermediate_cosize: int = 0 + + +def build_cfg( + io_dtype: Type[cutlass.Numeric], + state_dtype: Type[cutlass.Numeric], + *, + use_initial_state: bool, + store_final_state: bool, + enable_checkpoints: bool, + l2norm: bool, + safe_gate: bool, + gate_scale_log2: float, + beta_sigmoid: bool, + k_ratio: int, + v_ratio: int, + n_heads_out: int, + max_active_clusters: int, + dyn_sched: bool = False, +) -> KdaRecomputeCfg: + """Build the per-compile ``KdaRecomputeCfg`` (io_dtype in {Float16, BFloat16}); + fills the derived TMEM column offsets and SMEM buffer cosizes.""" + if io_dtype not in (cutlass.Float16, cutlass.BFloat16): + raise ValueError(f"io_dtype={io_dtype} not supported; only Float16 and BFloat16 are supported") + cfg = KdaRecomputeCfg( + io_dtype=io_dtype, + state_dtype=state_dtype, + use_initial_state=use_initial_state, + store_final_state=store_final_state, + enable_checkpoints=enable_checkpoints, + l2norm=l2norm, + safe_gate=safe_gate, + gate_scale_log2=gate_scale_log2, + beta_sigmoid=beta_sigmoid, + k_ratio=k_ratio, + v_ratio=v_ratio, + n_heads_out=n_heads_out, + max_active_clusters=max_active_clusters, + dyn_sched=dyn_sched, + ) + if enable_checkpoints: + cfg.smem_raw_stages = 6 + cfg.smem_checkpoint_stages = 2 + if cfg.smem_raw_stages % 2 != 0: + raise ValueError("smem_raw_stages must be even: the CG0 ping-pong groups alias parity waits on odd rings") + cfg.threads_per_cta = 16 * cfg.threads_per_warp + cfg.cg0_threads_per_group = cfg.cg0_warps_per_group * cfg.threads_per_warp + cfg.tmem_user_threads = (1 + len(cfg.compute_group_1_warp_ids)) * cfg.threads_per_warp + if cfg.smem_state_scale_diag_stages != cfg.qk_scale_ready_stages: + raise ValueError("diag and qk-scale ready rings must share their rolling stage") + + cfg.tmem_state_inp_offset = cfg.tmem_state_acc_offset + cfg.d_k + cfg.tmem_state_k_acc_offset = cfg.tmem_state_inp_offset + (cfg.d_k // 2) + cfg.tmem_u_acc_offset = cfg.tmem_state_k_acc_offset + cfg.b_t + cfg.tmem_y_inp_offset = cfg.tmem_u_acc_offset + cfg.b_t + cfg.tmem_u_inp_offset = cfg.tmem_y_inp_offset + (cfg.b_t // 2) + assert (cfg.tmem_u_inp_offset + (cfg.b_t // 2)) <= 512 + + cfg.k_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t + cfg.v_cosize = cfg.smem_raw_stages * cfg.d_v * cfg.b_t + cfg.gate_cosize = cfg.smem_raw_stages * cfg.d_k * cfg.b_t + cfg.beta_cosize = cfg.smem_raw_stages * cfg.b_t + cfg.k_inv_cosize = cfg.smem_decay_stages * cfg.b_t * cfg.d_k + cfg.k_decay_cosize = cfg.smem_decay_stages * cfg.d_k * cfg.b_t + cfg.k_restore_cosize = cfg.smem_decay_stages * cfg.d_k * cfg.b_t + cfg.state_scale_diag_cosize = cfg.smem_state_scale_diag_stages * (cfg.d_k // 16) * 256 + cfg.intermediate_cosize = cfg.smem_intermediate_stages * 2 * cfg.b_t * cfg.b_t + cfg.tma_k_bytes = cfg.d_k * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_v_bytes = cfg.d_v * cfg.b_t * (cfg.io_dtype.width // 8) + cfg.tma_gate_bytes = cfg.d_k * cfg.b_t * 4 + return cfg + + +TENSORMAP_DESC_ARRAYS = 4 # per-batch runtime TMA descriptors: K, V, Gate, state_checkpoints +TENSORMAP_STATIC_SLOTS = 0 + + +@cute.kernel +def build_all_descs_kernel( + base_k: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_v: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_gate: cutlass.GridConstant[cuda.tensor_map.TensorMap], + base_checkpoint: cutlass.GridConstant[cuda.tensor_map.TensorMap], + desc_ws: cute.Tensor, + cu_seqlens: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + state_checkpoints: cute.Tensor | None, + n_batch: cutlass.Int32, + k_row_stride: cutlass.Int32, + v_row_stride: cutlass.Int32, + gate_row_stride: cutlass.Int32, + checkpoint_row_stride: cutlass.Int32, + checkpoint_every_n: cutlass.Int32, +) -> None: + """Single-launch builder kernel: one warp emits each per-batch TMA + descriptor array.""" + tidx, _, _ = cute.arch.thread_idx() + widx = cutlass.Int32(tidx) // cutlass.Int32(32) + arr_words = n_batch * cutlass.Int32(TENSOR_MAP_QWORDS) + desc_words_k = cute.make_tensor(desc_ws.iterator, cute.make_layout((arr_words,), stride=(1,))) + desc_words_v = cute.make_tensor(desc_ws.iterator + arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_words_gate = cute.make_tensor(desc_ws.iterator + 2 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + desc_words_checkpoint = cute.make_tensor(desc_ws.iterator + 3 * arr_words, cute.make_layout((arr_words,), stride=(1,))) + + if widx == 0: + if nvvm.elect_sync(): + emit_seq_descs(base_k, desc_words_k, cu_seqlens, k, n_batch, k_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 1: + if nvvm.elect_sync(): + emit_seq_descs(base_v, desc_words_v, cu_seqlens, v, n_batch, v_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if widx == 2: + if nvvm.elect_sync(): + emit_seq_descs(base_gate, desc_words_gate, cu_seqlens, gate, n_batch, gate_row_stride, 2) + nvvm.fence_proxy_release(nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, to_proxy=nvvm.Proxy.TENSORMAP) + if cutlass.const_expr(state_checkpoints is not None): + if widx == 3: + if nvvm.elect_sync(): + emit_checkpoint_seq_descs( + base_checkpoint, desc_words_checkpoint, 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) + + +@cute.jit +def build_descs( + io_dtype: cutlass.Constexpr, + b_t: cutlass.Constexpr[int], + k: cute.Tensor, + v: cute.Tensor, + gate: cute.Tensor, + state_checkpoints: cute.Tensor | None, + cu_seqlens: cute.Tensor, + tensormap_workspace: cute.Tensor, + checkpoint_every_n: cutlass.Int32, + stream: cuda_driver.CUstream, +): + """Build the per-batch K/V/Gate/checkpoint TMA-descriptor arrays into + ``tensormap_workspace``.""" + h_k = k.shape[1] + h_v = v.shape[1] + ho = gate.shape[1] + batch_size = cu_seqlens.shape[0] - 1 + d_k = k.shape[2] + d_v = v.shape[2] + bpe = io_dtype.width // 8 + tma_granu_elems = 128 // bpe + seqlen = k.shape[0] + + k_headed = cute.make_tensor(k.iterator, cute.make_layout((d_k, h_k, seqlen), stride=(1, k.stride[1], k.stride[0]))) + v_headed = cute.make_tensor(v.iterator, cute.make_layout((d_v, h_v, seqlen), stride=(1, v.stride[1], v.stride[0]))) + gate_headed = cute.make_tensor(gate.iterator, cute.make_layout((d_k, ho, seqlen), stride=(1, gate.stride[1], gate.stride[0]))) + + swz = cuda.TensorMapSwizzle.s128b + base_k = cuda.create_tensor_map_tiled_from_view(k_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_v = cuda.create_tensor_map_tiled_from_view(v_headed, box_dims=(tma_granu_elems, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + base_gate = cuda.create_tensor_map_tiled_from_view(gate_headed, box_dims=(32, 1, b_t), stride_order=(0, 1, 2), swizzle=swz) + + base_checkpoint = base_gate + if cutlass.const_expr(state_checkpoints is not None): + checkpoint_view = cute.make_tensor( + state_checkpoints.iterator, + cute.make_layout( + (state_checkpoints.shape[3], state_checkpoints.shape[2], state_checkpoints.shape[0], ho), + stride=(state_checkpoints.stride[3], state_checkpoints.stride[2], state_checkpoints.stride[0], state_checkpoints.stride[1]), + ), + ) + base_checkpoint = cuda.create_tensor_map_tiled_from_view(checkpoint_view, box_dims=(tma_granu_elems, d_k, 1, 1), stride_order=(0, 1, 2, 3), swizzle=swz) + n_warps = 4 if state_checkpoints is not None else 3 + build_all_descs_kernel( + base_k, + base_v, + base_gate, + base_checkpoint, + tensormap_workspace, + cu_seqlens, + k, + v, + gate, + state_checkpoints, + cutlass.Int32(batch_size), + cutlass.Int32(k.stride[0]), + cutlass.Int32(v.stride[0]), + cutlass.Int32(gate.stride[0]), + cutlass.Int32(state_checkpoints.stride[0] if state_checkpoints is not None else 0), + checkpoint_every_n, + ).launch(grid=(1, 1, 1), block=(32 * n_warps, 1, 1), stream=stream) + + +# ---- Torch adapter / host-side compilation --------------------------------------- + + +@lru_cache(maxsize=None) +def get_compiled_cache( + io_dtype_str: str, + state_dtype_str: str, + cu_dtype_str: str, + HO: int, + HK: int, + HV: int, + use_initial_state: bool, + store_final_state: bool, + enable_checkpoints: bool, + l2norm: bool, + safe_gate: bool, + gate_lower_bound: float, + beta_sigmoid: bool, + dyn_sched: bool, +): + """Return a mutable dict that lazily stores the compiled kernel.""" + return {} + + +def compile( + io_dtype, + state_dtype, + use_initial_state: bool, + store_final_state: bool, + enable_checkpoints: bool, + l2norm: bool, + safe_gate: bool, + gate_scale_log2: float, + beta_sigmoid: bool, + k_ratio: int, + v_ratio: int, + n_heads_out: int, + dyn_sched: bool = False, + *, + num_sm: int, + k_cute, + v_cute, + gate_cute, + a_log_cute, + dt_bias_cute, + beta_cute, + cu_seqlens_cute, + state_in_cute, + state_out_cute, + work_items_cute=None, + work_count_cute=None, + sched_ctr_cute=None, + tensormap_ws_cute, + checkpoint_every_n_tokens, + stream, +): + """JIT-compile the chunked KDA recompute kernel for one static config.""" + cfg = build_cfg( + io_dtype, + state_dtype, + use_initial_state=use_initial_state, + store_final_state=store_final_state, + enable_checkpoints=enable_checkpoints, + l2norm=l2norm, + safe_gate=safe_gate, + gate_scale_log2=gate_scale_log2, + beta_sigmoid=beta_sigmoid, + k_ratio=k_ratio, + v_ratio=v_ratio, + n_heads_out=n_heads_out, + max_active_clusters=num_sm, + dyn_sched=dyn_sched, + ) + + return cute.compile( + host, + cfg, + k_cute, + v_cute, + gate_cute, + a_log_cute, + dt_bias_cute, + beta_cute, + cu_seqlens_cute, + state_in_cute, + state_out_cute, + work_items_cute, + work_count_cute, + sched_ctr_cute, + tensormap_ws_cute, + checkpoint_every_n_tokens, + stream, + options="--enable-tvm-ffi --opt-level 2", + ) + + +def chunk_kda_recompute_sm100( + k, + v, + gate, + beta, + cu_seqlens, + initial_state, + output_state, + checkpoint_every_n_tokens: int = 0, + output_state_checkpoints=None, + use_qk_l2norm_in_kernel: bool = False, + safe_gate: bool = False, + gate_lower_bound: float = DEFAULT_GATE_LOWER_BOUND, + a_log=None, + dt_bias=None, + use_beta_sigmoid_in_kernel: bool = False, + work_items=None, + work_count=None, + sched_ctr=None, + *, + tensormap_workspace, + stream, +) -> None: + """Execute the Blackwell BT=16 chunked KDA recompute (state/checkpoints-only) + kernel. + + All tensors must be on the same CUDA device with a stride-1 innermost + dim; outer strides are free (padded / permuted views are read through + the TMA descriptors and dynamic layouts). + + Args: + k: ``(total_tokens, HK, DK)`` float16/bfloat16 + v: ``(total_tokens, HV, DV)`` float16/bfloat16 + gate: ``(total_tokens, HO, DK)`` float32. Natural-log decay unless + ``safe_gate``, which applies the safe-gate transform + ``lower_bound * sigmoid(exp(a_log) * (gate + dt_bias))``. + beta: ``(total_tokens, HO)``. Post-sigmoid float32, or io-dtype + logits when ``use_beta_sigmoid_in_kernel`` + 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 + 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 + 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 + use_qk_l2norm_in_kernel: L2-normalize k rows inside the kernel + safe_gate: interpret ``gate`` through the safe-gate transform + a_log: ``(HO,)`` float32, safe-gate per-head log-amplitude (None = 0) + dt_bias: ``(HO, DK)`` float32, safe-gate channel bias (None = 0) + use_beta_sigmoid_in_kernel: ``beta`` holds logits; sigmoid in-kernel + work_items: ``(max_items, 8)`` int32 work-item table from + ``common/split_k.py`` (REQUIRED; an uncut table row is the whole + (b, h) sequence). Each item computes chunks ``[cstart, wend)`` + and writes checkpoints only for ``[wstart, wend)``. + work_count: ``(1,)`` int32 device-side item count (REQUIRED) + sched_ctr: ``(2,)`` int32 device scratch ``[ticket, done]`` enabling + the dynamic (work-stealing) tile scheduler; must be zeroed before + every launch (``build_split_table`` does this when it is passed as + ``sched_ctr``). None keeps the static CTA stride. + """ + HK = k.shape[1] + HV = v.shape[1] + HO = gate.shape[1] + use_initial_state = initial_state is not None + store_final_state = output_state is not None + enable_checkpoints = checkpoint_every_n_tokens > 0 + if enable_checkpoints: + if output_state_checkpoints is None: + raise ValueError("checkpoint_every_n_tokens > 0 requires output_state_checkpoints") + if str(output_state_checkpoints.dtype).split(".")[-1] != str(k.dtype).split(".")[-1]: + raise ValueError( + f"output_state_checkpoints dtype must match the io dtype (fp32 state belongs to output_state): got {output_state_checkpoints.dtype} with io {k.dtype}" + ) + 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 + + if initial_state is not None: + state_dtype_src = initial_state.dtype + elif output_state is not None: + state_dtype_src = output_state.dtype + else: + state_dtype_src = "float32" + + for name, h in (("HK", HK), ("HV", HV)): + if HO % h != 0: + raise ValueError(f"{name}={h} must divide sab heads {HO}") + k_ratio = HO // HK + v_ratio = HO // HV + gate_scale_log2 = gate_lower_bound * LOG2_E + + if safe_gate and (a_log is None or dt_bias is None): + raise ValueError("safe_gate requires a_log and dt_bias") + if not safe_gate: + a_log = None + dt_bias = None + cu_stream = cuda_driver.CUstream(int(stream)) + + cache = get_compiled_cache( + str(k.dtype), + str(state_dtype_src), + str(cu_seqlens.dtype), + HO, + HK, + HV, + use_initial_state, + store_final_state, + enable_checkpoints, + use_qk_l2norm_in_kernel, + safe_gate, + gate_lower_bound, + use_beta_sigmoid_in_kernel, + dyn_sched, + ) + + if "compiled" not in cache: + io_dtype = get_dtype(k.dtype) + state_dtype = get_dtype(state_dtype_src) + k_cute = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_cute = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_cute = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2) + 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 + beta_cute = from_dlpack(beta, assumed_align=4).mark_layout_dynamic(leading_dim=1) + cu_seqlens_cute = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() + + state_in_cute = None + if use_initial_state: + state_in_cute = from_dlpack(initial_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) + + state_out_cute = None + if store_final_state: + state_out_cute = from_dlpack(output_state, assumed_align=16).mark_layout_dynamic(leading_dim=3) + + work_items_cute = from_dlpack(work_items, assumed_align=16) + work_items_cute.mark_compact_shape_dynamic(mode=0, stride_order=(0, 1), divisibility=1) + work_count_cute = from_dlpack(work_count, assumed_align=4).mark_layout_dynamic() + + sched_ctr_cute = None + if dyn_sched: + sched_ctr_cute = from_dlpack(sched_ctr, assumed_align=4).mark_layout_dynamic() + + tensormap_ws_cute = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + + cache["compiled"] = compile( + io_dtype, + state_dtype, + use_initial_state, + store_final_state, + enable_checkpoints, + use_qk_l2norm_in_kernel, + safe_gate, + gate_scale_log2, + use_beta_sigmoid_in_kernel, + k_ratio, + v_ratio, + HO, + dyn_sched, + num_sm=multiprocessor_count(current_device_id()), + k_cute=k_cute, + v_cute=v_cute, + gate_cute=gate_cute, + a_log_cute=a_log_cute, + dt_bias_cute=dt_bias_cute, + beta_cute=beta_cute, + cu_seqlens_cute=cu_seqlens_cute, + state_in_cute=state_in_cute, + state_out_cute=state_out_cute, + work_items_cute=work_items_cute, + work_count_cute=work_count_cute, + sched_ctr_cute=sched_ctr_cute, + tensormap_ws_cute=tensormap_ws_cute, + checkpoint_every_n_tokens=checkpoint_every_n_tokens, + stream=cu_stream, + ) + + compiled = cache["compiled"] + state_checkpoints_for_descs = output_state_checkpoints if enable_checkpoints else None + # desc build runs every execute by contract (cu contents are data; + # buffer pointers may change) - capture-safe, single tiny launch + if cache.get("build_descs_has_state_checkpoints") != (state_checkpoints_for_descs is not None): + cache.pop("build_descs", None) + cache["build_descs_has_state_checkpoints"] = state_checkpoints_for_descs is not None + if "build_descs" not in cache: + io_dtype = get_dtype(k.dtype) + k_bd = from_dlpack(k, assumed_align=16).mark_layout_dynamic(leading_dim=2) + v_bd = from_dlpack(v, assumed_align=16).mark_layout_dynamic(leading_dim=2) + gate_bd = from_dlpack(gate, assumed_align=16).mark_layout_dynamic(leading_dim=2) + cu_bd = from_dlpack(cu_seqlens, assumed_align=8).mark_layout_dynamic() + ws_bd = from_dlpack(tensormap_workspace, assumed_align=128).mark_layout_dynamic() + state_checkpoints_bd = None + if state_checkpoints_for_descs is not None: + state_checkpoints_bd = from_dlpack(state_checkpoints_for_descs, assumed_align=16).mark_layout_dynamic(leading_dim=3) + cache["build_descs"] = cute.compile( + build_descs, + io_dtype, + CFG.B_T, + k_bd, + v_bd, + gate_bd, + state_checkpoints_bd, + cu_bd, + ws_bd, + cutlass.Int32(checkpoint_every_n_tokens), + cu_stream, + options="--enable-tvm-ffi", + ) + cache["build_descs"]( + k, + v, + gate, + state_checkpoints_for_descs, + cu_seqlens, + tensormap_workspace, + checkpoint_every_n_tokens, + cu_stream, + ) + compiled( + k, + v, + gate, + a_log, + dt_bias, + beta, + cu_seqlens, + initial_state if use_initial_state else None, + output_state if store_final_state else None, + work_items, + work_count, + sched_ctr, + tensormap_workspace, + checkpoint_every_n_tokens, + cu_stream, + ) diff --git a/python/cudnn/linear_attention/graph_analyzer.py b/python/cudnn/linear_attention/graph_analyzer.py new file mode 100644 index 000000000..e3b7fa6a3 --- /dev/null +++ b/python/cudnn/linear_attention/graph_analyzer.py @@ -0,0 +1,356 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Engine-agnostic linear-attention graph analysis: ``graph.nodes`` -> :class:`LaGraphFacts`. + +One analyzer serves the three LA families (gdn / kda / gdn2 — single dedicated +nodes sharing the THD port vocabulary). :func:`analyze` is the callable each +family names in ``engines/manifest.py``; PLANNING runs it once per frozen +graph and attaches the record, so the family's engines read that same record +back instead of each parsing the node. + +Also hosts the engine-side helpers shared by the LA engines: the +check_support gates over the facts record, the execute-time buffer-layout +gate, and the compiled-plan wrapper. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional + +import cudnn +from cudnn.engines.base import CompiledPlan, NodeBuffers, bind_ports +from cudnn.frost import buffers +from cudnn.frost.workspace import Workspace + +BUFFER_NAME_FROM_CUDNN = { + cudnn.data_type.HALF: "float16", + cudnn.data_type.BFLOAT16: "bfloat16", + cudnn.data_type.FLOAT: "float32", + cudnn.data_type.INT32: "int32", + cudnn.data_type.INT64: "int64", +} + + +def to_buffer_dtype(dt) -> str: + """cudnn.data_type -> the buffer-level dtype name (``buffers.DTYPES`` vocabulary).""" + return BUFFER_NAME_FROM_CUDNN[dt] + + +# node type -> (op family, is_bwd) +LA_NODE_OPS = { + cudnn.NodeType.GDN: ("GDN", False), + cudnn.NodeType.GDN_BWD: ("GDN", True), + cudnn.NodeType.KDA: ("KDA", False), + cudnn.NodeType.KDA_BWD: ("KDA", True), + cudnn.NodeType.GDN2: ("GDN2", False), + cudnn.NodeType.GDN2_BWD: ("GDN2", True), +} + + +@dataclass(frozen=True) +class LaGraphFacts: + """What a single-LA graph asks for. Pure description — no support + judgment: head-dim limits, dtype sets, and feature coverage are per-engine + knowledge, matched in each engine's ``check_support``. ``invalid`` is the + one exception: a graph-consistency error (malformed regardless of which + kernel would run — a missing required port, ``d_initial_state`` without + ``initial_state``, safe-gate inputs without the attribute); when set, + every engine is ineligible.""" + + invalid: Optional[str] = None + + op: str = "" # "GDN" | "KDA" | "GDN2" + is_bwd: bool = False + + # geometry (THD; zeros when the port ranks are not declared) + thd_layout: bool = True # Q/K/V are rank-3 [total_T, heads, dim] + h_q: int = 0 + h_k: int = 0 + h_v: int = 0 + h_o: int = 0 + d_qk: int = 0 + d_v: int = 0 + gates_at_ho: bool = True # Gate/Beta(/W) carry HO = max(h_q, h_v) heads + + # dtypes (cudnn.data_type vocabulary; None = unset/inferred) + io_dtype: Any = None # Q's dtype + uniform_io: bool = True # Q/K/V dtypes agree + g_dtype: Any = None + beta_dtype: Any = None + w_dtype: Any = None + cu_dtype: Any = None + a_log_dtype: Any = None + dt_bias_dtype: Any = None + do_dtype: Any = None + state_checkpoints_dtype: Any = None + state_checkpoints_out_dtype: Any = None # fwd checkpoint OUTPUT port dtype + d_final_state_dtype: Any = None + state_dtype: Any = None + final_state_dtype: Any = None + state_pair_match: bool = True + o_dtype: Any = None + dq_dtype: Any = None + dk_dtype: Any = None + dv_dtype: Any = None + dg_dtype: Any = None + dbeta_dtype: Any = None + dw_dtype: Any = None + d_initial_state_dtype: Any = None + + # ports present / requested + has_initial_state: bool = False + wants_d_initial_state: bool = False + wants_state_checkpoints: bool = False # fwd checkpoint-series output + + # attributes + scale: Optional[float] = None + use_qk_l2norm: bool = False + safe_gate: bool = False + use_beta_sigmoid: bool = False + checkpoint_every_n_tokens: int = 0 + batch_invariant: bool = False + + +def analyze(graph: "cudnn.pygraph") -> Optional[LaGraphFacts]: + """Facts for a single-LA graph, or None if the graph is anything else. + + Pure: attaching and caching is the graph's job (create_execution_plans + -> _attach_facts). This is the callable the LA families name in their + manifest ``analyzer`` entries.""" + nodes = list(graph.nodes) + if len(nodes) != 1: + return None + node = nodes[0] + kind = LA_NODE_OPS.get(node.node_type) + if kind is None: + return None + op, is_bwd = kind + ins, outs, params = node.inputs, node.outputs, node.params + + required_in = ["q", "k", "v", "g", "beta", "cu_seqlens"] + if op == "GDN2": + required_in.append("w") + if is_bwd: + required_in.append("dO") + required_out = (["dQ", "dK", "dV", "dG", "dBeta"] + (["dW"] if op == "GDN2" else [])) if is_bwd else ["O"] + + safe_gate = bool(params.get("safe_gate", False)) + ckpt = int(params.get("checkpoint_every_n_tokens", 0) or 0) + invalid = None + missing_in = [p for p in required_in if p not in ins] + missing_out = [p for p in required_out if p not in outs] + if missing_in: + invalid = f"{node.node_type.name} node '{node.name}' is missing input(s) {missing_in}" + elif missing_out: + invalid = f"{node.node_type.name} node '{node.name}' is missing output(s) {missing_out}" + elif "d_initial_state" in outs and "initial_state" not in ins: + invalid = "d_initial_state requires initial_state" + elif safe_gate and ("a_log" not in ins or "dt_bias" not in ins): + invalid = "safe_gate requires a_log and dt_bias inputs" + elif not safe_gate and ("a_log" in ins or "dt_bias" in ins): + invalid = "a_log/dt_bias require safe_gate=True" + elif params.get("gate_lower_bound") is not None and not safe_gate: + invalid = "gate_lower_bound requires safe_gate=True" + elif ckpt < 0: + invalid = "checkpoint_every_n_tokens must be non-negative" + elif not is_bwd and ckpt > 0 and "state_checkpoints" not in outs: + invalid = "checkpoint_every_n_tokens > 0 requires the state_checkpoints output" + elif not is_bwd and ckpt == 0 and "state_checkpoints" in outs: + invalid = "state_checkpoints output requires checkpoint_every_n_tokens > 0" + if invalid is not None: + return LaGraphFacts(invalid=invalid, op=op, is_bwd=is_bwd) + + in_dt = {name: t.get_data_type() for name, t in ins.items()} + out_dt = {name: t.get_data_type() for name, t in outs.items()} + q, k, v = ins["q"], ins["k"], ins["v"] + + thd_layout = all(t.dim and len(t.dim) == 3 for t in (q, k, v)) + if thd_layout: + _, h_q, d_qk = (int(d) for d in q.dim) + h_k, h_v, d_v = int(k.dim[1]), int(v.dim[1]), int(v.dim[2]) + else: + h_q = h_k = h_v = d_qk = d_v = 0 + h_o = max(h_q, h_v) + 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} + scale = params.get("scale") + + return LaGraphFacts( + op=op, + is_bwd=is_bwd, + thd_layout=thd_layout, + h_q=h_q, + h_k=h_k, + h_v=h_v, + h_o=h_o, + d_qk=d_qk, + d_v=d_v, + gates_at_ho=gates_at_ho, + io_dtype=in_dt["q"], + uniform_io=len(io_dtypes) <= 1, + g_dtype=in_dt["g"], + beta_dtype=in_dt["beta"], + w_dtype=in_dt.get("w"), + cu_dtype=in_dt["cu_seqlens"], + a_log_dtype=in_dt.get("a_log"), + dt_bias_dtype=in_dt.get("dt_bias"), + do_dtype=in_dt.get("dO"), + state_checkpoints_dtype=in_dt.get("state_checkpoints"), + state_checkpoints_out_dtype=out_dt.get("state_checkpoints"), + d_final_state_dtype=in_dt.get("d_final_state"), + state_dtype=in_dt.get("initial_state"), + final_state_dtype=out_dt.get("final_state"), + state_pair_match=len(state_dtypes) <= 1, + o_dtype=out_dt.get("O"), + dq_dtype=out_dt.get("dQ"), + dk_dtype=out_dt.get("dK"), + dv_dtype=out_dt.get("dV"), + dg_dtype=out_dt.get("dG"), + dbeta_dtype=out_dt.get("dBeta"), + dw_dtype=out_dt.get("dW"), + d_initial_state_dtype=out_dt.get("d_initial_state"), + has_initial_state="initial_state" in ins, + wants_d_initial_state="d_initial_state" in outs, + wants_state_checkpoints="state_checkpoints" in outs, + scale=float(scale) if scale is not None else None, + use_qk_l2norm=bool(params.get("use_qk_l2norm", False)), + safe_gate=safe_gate, + use_beta_sigmoid=bool(params.get("use_beta_sigmoid", False)), + checkpoint_every_n_tokens=ckpt, + batch_invariant=bool(params.get("batch_invariant", False)), + ) + + +# --------------------------------------------------------------------------- +# Engine-side helpers shared by the LA engines +# --------------------------------------------------------------------------- + + +def require(engine: str, port: str, got, want) -> None: + """check_support dtype gate over a facts field: unset passes (the kernel + validates the buffer), anything else must be the kernel-native dtype.""" + if got is None: + return + wanted = want if isinstance(want, tuple) else (want,) + if got not in wanted: + names = "/".join(w.name for w in wanted) + raise NotImplementedError(f"{engine}: '{port}' must be {names} (the kernel-native dtype; no staging), got {got}") + + +def frost_la_gate(engine: str, facts, op: str) -> None: + """The FROST LA engines' shared check_support core: the analyzer record, + the device/DSL environment, and the gates common to all three kernels.""" + if facts is None or facts.op != op: + raise NotImplementedError(f"{engine} supports exactly one {op}/{op}_BWD node") + if facts.invalid: + raise NotImplementedError(f"{engine}: {facts.invalid}") + sm = buffers.current_sm() + if sm is None or not (100 <= sm <= 103): + raise NotImplementedError(f"{engine} requires SM100-SM103 (found {sm})") + installed, version = buffers.cutedsl_state() + if not installed: + raise NotImplementedError(f"{engine} requires the cutedsl extra (nvidia-cutlass-dsl), which is not installed") + if buffers.cutedsl_too_old(version): + want = ".".join(str(v) for v in buffers.CUTEDSL_MIN_VERSION) + raise NotImplementedError(f"{engine} requires nvidia-cutlass-dsl >= {want}; found {version[1]}") + if not facts.uniform_io: + raise NotImplementedError(f"{engine}: q/k/v dtypes must match") + require(engine, "q/k/v", facts.io_dtype, (cudnn.data_type.BFLOAT16, cudnn.data_type.HALF)) + if not facts.thd_layout: + raise NotImplementedError(f"{engine}: q/k/v must be THD [total_T, heads, dim]") + if facts.d_qk != 128 or facts.d_v != 128: + raise NotImplementedError(f"{engine}: head dims must be 128 (the recurrent state is 128x128), got K={facts.d_qk} V={facts.d_v}") + if facts.h_k not in (facts.h_q, facts.h_v): + raise NotImplementedError(f"{engine}: k heads ({facts.h_k}) must match q's ({facts.h_q}) or v's ({facts.h_v}; canonical GQA shares grouped k/v heads)") + if facts.h_v != facts.h_q and max(facts.h_q, facts.h_v) % min(facts.h_q, facts.h_v) != 0: + raise NotImplementedError(f"{engine}: q heads ({facts.h_q}) and v heads ({facts.h_v}) must be equal or one a multiple of the other") + require(engine, "g", facts.g_dtype, cudnn.data_type.FLOAT) + require(engine, "cu_seqlens", facts.cu_dtype, (cudnn.data_type.INT32, cudnn.data_type.INT64)) + + +class FrostLaPlan(CompiledPlan): + """A compiled LA executor, driven from the normalized variant pack: the + port-to-slot join is a property of the graph, so it happens once and is + kept; between executes only the buffer addresses move.""" + + takes_variant_pack = True + + def __init__(self, compiled): + self.compiled = compiled + self.ports = None + + def get_workspace_size(self) -> int: + return self.compiled.workspace_bytes() + + def execute(self, graph, variant_pack, ctx) -> None: + ports = self.ports + if ports is None: + ports = self.ports = bind_ports(graph, variant_pack) + ok, offender = variant_pack.all_dense_layout() + if not ok: + raise ValueError(dense_layout_message(self.compiled.plan_name, ports, offender)) + node_buffers = {} + for node, slots in ports.items(): + names = list(slots.inputs) + list(slots.outputs) + views = variant_pack.operands(list(slots.inputs.values()) + list(slots.outputs.values())) + split = len(slots.inputs) + node_buffers[node] = NodeBuffers(dict(zip(names[:split], views[:split])), dict(zip(names[split:], views[split:]))) + workspace = Workspace.over(variant_pack, self.compiled.workspace_bytes(), type(self.compiled).__name__) + self.compiled(node_buffers, workspace=workspace, stream=ctx.stream) + + +def expect_table(node, align) -> dict: + """Build-time ``{port: (dims, dtype_name, align_bytes)}`` for + :func:`check_layouts`: bound buffers must match the node's frozen + geometry exactly (one graph per shape), and base pointers must satisfy + the kernel entry's ``assumed_align`` claim. ``align`` maps port name -> + bytes (family table read from the entry's from_dlpack calls; absent + ports default to 16).""" + table = {} + for ports in (node.inputs, node.outputs): + for name, t in ports.items(): + if t is None: + continue + dims = tuple(int(d) for d in t.dim) if t.dim else None + table[name] = (dims, BUFFER_NAME_FROM_CUDNN.get(t.get_data_type()), align.get(name, 16)) + return table + + +def check_layouts_compact(plan_name: str, expect, nb) -> None: + """Execute-time gate for the cuTile backend: every bound buffer must be + CONTIGUOUS (the kernels stage rank-merged views and whole-buffer zero + fills), and must match the node's build-time dims/dtype and base + alignment per ``expect`` (see :func:`expect_table`).""" + for ports in (nb.inputs, nb.outputs): + for name, b in ports.items(): + if b is None: + continue + ptr, shape, strides, dtype, _dev = buffers.probe(b) + exp = expect.get(name) if expect else None + if exp is not None: + dims, dtype_name, align = exp + if dims is not None and tuple(shape) != dims: + raise ValueError(f"{plan_name}: buffer for {name!r} must match the graph's build-time dims {dims}; got {tuple(shape)}") + if dtype_name is not None and dtype != dtype_name: + raise ValueError(f"{plan_name}: buffer for {name!r} must be {dtype_name} (the node's declared dtype); got {dtype}") + if align and ptr % align != 0: + raise ValueError(f"{plan_name}: buffer for {name!r} base pointer must be {align}-byte aligned; got 0x{ptr:x}") + if not buffers.is_contiguous(shape, strides): + raise ValueError( + f"{plan_name}: buffer for {name!r} must be contiguous (the cuTile backend stages rank-merged views); got shape {shape} strides {strides}" + ) + + +def dense_layout_message(plan_name, ports, offender) -> str: + """Name the port behind ``all_dense_layout``'s failing slot. Buffers pass + straight to the stride-plumbed kernels, so the one execute-time rule is a + stride-1 innermost dim; this walk only runs on the way to raising.""" + for slots in ports.values(): + for direction in (slots.inputs, slots.outputs): + for port, slot in direction.items(): + if slot == offender: + return f"{plan_name}: buffer for {port!r} must have a stride-1 innermost dim (buffers pass straight to the kernel)" + return f"{plan_name}: the buffer at variant-pack slot {offender} must have a stride-1 innermost dim" diff --git a/python/cudnn/linear_attention/ops/gdn.py b/python/cudnn/linear_attention/ops/gdn.py index a27186e4e..33a3f7666 100644 --- a/python/cudnn/linear_attention/ops/gdn.py +++ b/python/cudnn/linear_attention/ops/gdn.py @@ -10,130 +10,254 @@ S_t = alpha_t (I - beta_t k_t^T k_t) S_{t-1} + beta_t k_t^T v_t, o_t = q_t S_t, -where ``alpha_t`` is a scalar per-token decay in ``(0, 1]`` and ``beta_t`` -is a scalar per-token write strength. +where ``S_t`` is the recurrent state, ``alpha_t`` a scalar per-token decay +in ``(0, 1]``, and ``beta_t`` a scalar per-token write strength. Layout follows the graph-API GDN node: THD — token-packed ``[total_tokens, heads, dim]`` tensors plus ``cu_seqlens`` sequence boundaries. -The op is a thin adapter over the graph API (the SDPA-op pattern): forward +The op is a thin adapter over the graph API: forward and backward execute cached single-node ``GDN`` / ``GDN_BWD`` pygraphs. -Engine selection happens at graph planning time over the registered python +Engine selection happens at graph planning time over the manifest's python engines: ``GdnFrostEngine`` (default on SM100/SM103) with ``GdnCuTileEngine`` as the fallback everywhere else. Registered through ``torch.library.custom_op`` so it composes with autograd, ``torch.compile``, and DDP. -The backward graph recomputes the forward's cheap intermediates (cumulative -gate, intra-chunk WY factor) — the ``GDN_BWD`` node contract keeps them off -the autograd wire. +Graph caching ensures cuDNN graphs are built once per unique configuration +and reused across calls. """ -from __future__ import annotations - import math from typing import Dict, Optional, Tuple import torch - import cudnn -from cudnn.linear_attention import engine_utils -_OP_NAMESPACE = "cudnn" -_OP_NAME = "gated_delta_net" +# --------------------------------------------------------------------------- +# Module-level state +# --------------------------------------------------------------------------- + -_TORCH_TO_CUDNN_DTYPE = { +_TORCH_DTYPE_TO_CUDNN = { torch.float16: cudnn.data_type.HALF, torch.bfloat16: cudnn.data_type.BFLOAT16, torch.float32: cudnn.data_type.FLOAT, + torch.int32: cudnn.data_type.INT32, + torch.int64: cudnn.data_type.INT64, } # one graph per static configuration (shapes, dtypes, scale, flags, device) -_fwd_graph_cache: Dict[tuple, tuple] = {} -_bwd_graph_cache: Dict[tuple, tuple] = {} -_ws_cache: Dict[int, "torch.Tensor"] = {} +_fprop_cache: Dict[tuple, tuple] = {} +_bprop_cache: Dict[tuple, tuple] = {} +_cudnn_handles: Dict[int, int] = {} + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def select_plan(graph, plan_name): + """Pin one execution plan by name on a freshly built graph (create the + plans, select by name, check support); ``None`` keeps default routing.""" + if plan_name is None: + return + graph.create_execution_plans() + names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] + matches = [i for i, n in enumerate(names) if n == plan_name or n.startswith(plan_name + "[")] + if not matches: + raise cudnn.cudnnGraphNotSupportedError(f"no {plan_name} plan for this graph (offered: {names})") + graph.select_plan(matches[0]) + graph.check_support() def _graph_workspace(graph, device): - """Caller-side workspace for a compiled graph (the explicit-workspace - convention: query the plan's size, allocate, pass to execute).""" + """Caller-side workspace for a compiled graph (grow-only, held on the + graph object itself — same lifetime by construction, no id()-keyed + side table).""" if not graph._is_built: - # mirror execute()'s auto-build: plan via the router first (a bare - # build() would lower GDN to the backend, which has no lowering) + # plan first: a bare build() would lower GDN to the backend, + # which has no lowering if not graph._planning_done: graph.create_execution_plans() - engine_utils.apply_pin(graph) if graph.selected_engine is None: graph.build() else: graph.build_plans() size = graph.get_workspace_size() - ws = _ws_cache.get(id(graph)) + ws = getattr(graph, "_la_ops_workspace", None) if ws is None or ws.numel() < size or ws.device != device: ws = torch.empty(max(size, 1), dtype=torch.uint8, device=device) - _ws_cache[id(graph)] = ws + graph._la_ops_workspace = ws return ws -_handle_cache: Dict[int, int] = {} - - -def _graph_handle(device): - """Per-device cuDNN handle carrying the caller's current stream - (classic ``set_stream`` semantics).""" +def _get_handle(device): + """Per-device cuDNN handle carrying the caller's current stream.""" idx = device.index if device.index is not None else torch.cuda.current_device() - handle = _handle_cache.get(idx) + handle = _cudnn_handles.get(idx) if handle is None: with torch.cuda.device(idx): handle = cudnn.create_handle() - _handle_cache[idx] = handle + _cudnn_handles[idx] = handle cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) return handle -def _cudnn_dtype(dtype: Optional[torch.dtype]): - return _TORCH_TO_CUDNN_DTYPE[dtype] if dtype is not None else None - - -# --------------------------------------------------------------------------- -# Forward -# --------------------------------------------------------------------------- +def _torch_dtype_to_cudnn(dtype: torch.dtype): + """Map a PyTorch dtype to a cuDNN data_type enum.""" + return _TORCH_DTYPE_TO_CUDNN[dtype] def _check_dtype(name, t, want) -> None: if t.dtype != want: - raise TypeError(f"{_OP_NAME}: {name} must be {want} (kernel-native; callers convert), got {t.dtype}") + raise TypeError(f"gated_delta_net: {name} must be {want} (kernel-native; callers convert), got {t.dtype}") + + +def _make_fprop_cache_key( + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + k_shape, + v_shape, + cu_dtype, + scale, + output_final_state, + use_qk_l2norm, + batch_invariant, + has_initial_state, + ckpt, + device, + plan_name, +): + return ( + "fprop", + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + k_shape, + v_shape, + cu_dtype, + float(scale), + bool(output_final_state), + bool(use_qk_l2norm), + bool(batch_invariant), + bool(has_initial_state), + ckpt, + device, + plan_name, + ) + + +def _make_bprop_cache_key( + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + do_dtype, + k_shape, + v_shape, + cu_dtype, + has_initial_state, + has_d_final_state, + ckpt_rows, + scale, + use_qk_l2norm, + batch_invariant, + device, + plan_name, +): + return ( + "bprop", + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + do_dtype, + k_shape, + v_shape, + cu_dtype, + bool(has_initial_state), + bool(has_d_final_state), + ckpt_rows, + float(scale), + bool(use_qk_l2norm), + bool(batch_invariant), + device, + plan_name, + ) -def _build_fwd_graph(total, N, H, HV, K, V, io_dtype, g_dtype, beta_dtype, state_dtype, scale, output_final_state, use_qk_l2norm): +# --------------------------------------------------------------------------- +# Forward graph builder +# --------------------------------------------------------------------------- + + +def _build_fprop_graph( + total, N, H, HK, HV, K, V, io_dtype, g_dtype, beta_dtype, state_dtype, cu_dtype, scale, output_final_state, use_qk_l2norm, batch_invariant, ckpt +): graph = cudnn.pygraph() HO = max(H, HV) q_t = graph.tensor([total, H, K], data_type=io_dtype, name="q") - k_t = graph.tensor([total, H, K], data_type=io_dtype, name="k") + k_t = graph.tensor([total, HK, K], data_type=io_dtype, name="k") v_t = graph.tensor([total, HV, V], data_type=io_dtype, name="v") g_t = graph.tensor([total, HO], data_type=g_dtype, name="g") beta_t = graph.tensor([total, HO], data_type=beta_dtype, name="beta") - cu_t = graph.tensor([N + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") - s0_t = None + cu_t = graph.tensor([N + 1], data_type=cu_dtype, name="cu_seqlens") + state0_t = None if state_dtype is not None: - s0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") - O_t, fs_t, _h_t = graph.gdn( + state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") + O_t, fs_t, state_checkpoints_t = graph.gdn( q=q_t, k=k_t, v=v_t, g=g_t, beta=beta_t, cu_seqlens=cu_t, - initial_state=s0_t, + initial_state=state0_t, scale=scale, output_final_state=output_final_state, use_qk_l2norm=use_qk_l2norm, + batch_invariant=batch_invariant, + checkpoint_every_n_tokens=ckpt, name="gdn", ) - return graph, dict(q=q_t, k=k_t, v=v_t, g=g_t, beta=beta_t, cu=cu_t, s0=s0_t, O=O_t, fs=fs_t) + return graph, dict(q=q_t, k=k_t, v=v_t, g=g_t, beta=beta_t, cu=cu_t, state0=state0_t, O=O_t, fs=fs_t, state_checkpoints=state_checkpoints_t) -@torch.library.custom_op(f"{_OP_NAMESPACE}::{_OP_NAME}_fwd", mutates_args=()) +# --------------------------------------------------------------------------- +# Forward custom op +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("cudnn::gated_delta_net_fwd", mutates_args=()) def _gdn_fwd( q: torch.Tensor, k: torch.Tensor, @@ -145,117 +269,177 @@ def _gdn_fwd( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor]: - """GDN forward via a cached single-node GDN pygraph (THD layout). - - Returns ``(o, final_state)``; ``final_state`` is a zero-size tensor when - ``output_final_state`` is ``False``. + batch_invariant: bool = False, + checkpoint_every_n_tokens: int = 0, + plan_name: Optional[str] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """GDN forward (internal): a cached single-node GDN pygraph, THD layout. + + Returns ``(o, final_state, state_checkpoints)``; ``final_state`` / ``state_checkpoints`` are zero-size + tensors when ``output_final_state`` is ``False`` / + ``checkpoint_every_n_tokens`` is ``0``. """ total, H, K = q.shape - if k.shape[1] != H: - raise ValueError(f"k must carry the same head count as q ({H}), got {k.shape[1]}") + HK = k.shape[1] HV, V = v.shape[1], v.shape[2] + if HK not in (H, HV): + raise ValueError(f"k head count ({HK}) must match q's ({H}) or v's ({HV}); canonical GQA shares grouped k/v heads") N = cu_seqlens.shape[0] - 1 device = q.device - cu = cu_seqlens.to(torch.int32).contiguous() + if cu_seqlens.dtype not in (torch.int32, torch.int64): + raise ValueError(f"gated_delta_net: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") + cu = cu_seqlens _check_dtype("g", g, torch.float32) _check_dtype("beta", beta, torch.float32) if initial_state is not None: _check_dtype("initial_state", initial_state, torch.float32) if initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") - g32 = g.contiguous() - beta32 = beta.contiguous() - s0 = initial_state.contiguous() if initial_state is not None else None - - key = ( + for _name, _t in (("k", k), ("v", v), ("g", g), ("beta", beta), ("cu_seqlens", cu_seqlens), ("initial_state", initial_state)): + if _t is not None and _t.device != device: + raise ValueError(f"gated_delta_net: {_name} must be on q's device ({device}); got {_t.device}") + g32 = g + beta32 = beta + state0 = initial_state if initial_state is not None else None + ckpt = int(checkpoint_every_n_tokens) + + cache_key = _make_fprop_cache_key( total, N, H, + HK, HV, K, V, q.dtype, - float(scale), - bool(output_final_state), - bool(use_qk_l2norm_in_kernel), - bool(s0 is not None), + k.dtype, + v.dtype, + tuple(k.shape), + tuple(v.shape), + cu_seqlens.dtype, + scale, + output_final_state, + use_qk_l2norm_in_kernel, + batch_invariant, + state0 is not None, + ckpt, device, + plan_name, ) - if key not in _fwd_graph_cache: - _fwd_graph_cache[key] = _build_fwd_graph( + if cache_key not in _fprop_cache: + _fprop_cache[cache_key] = _build_fprop_graph( total, N, H, + HK, HV, K, V, - _cudnn_dtype(q.dtype), + _torch_dtype_to_cudnn(q.dtype), cudnn.data_type.FLOAT, cudnn.data_type.FLOAT, - cudnn.data_type.FLOAT if s0 is not None else None, + cudnn.data_type.FLOAT if state0 is not None else None, + _torch_dtype_to_cudnn(cu_seqlens.dtype), float(scale), bool(output_final_state), bool(use_qk_l2norm_in_kernel), + bool(batch_invariant), + ckpt, ) - graph, t = _fwd_graph_cache[key] + select_plan(_fprop_cache[cache_key][0], plan_name) + + graph, t = _fprop_cache[cache_key] HO = max(H, HV) o = torch.empty(total, HO, V, dtype=q.dtype, device=device) variant_pack = { - t["q"]: q.contiguous(), - t["k"]: k.contiguous(), - t["v"]: v.contiguous(), + t["q"]: q, + t["k"]: k, + t["v"]: v, t["g"]: g32, t["beta"]: beta32, t["cu"]: cu, t["O"]: o, } - if s0 is not None: - variant_pack[t["s0"]] = s0 + if state0 is not None: + variant_pack[t["state0"]] = state0 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) variant_pack[t["fs"]] = final_state - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_graph_handle(device)) - return o, 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) + 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 @_gdn_fwd.register_fake -def _gdn_fwd_fake(q, k, v, g, beta, cu_seqlens, scale, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False): +def _gdn_fwd_fake( + q, + k, + v, + g, + beta, + cu_seqlens, + scale, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + batch_invariant=False, + checkpoint_every_n_tokens=0, + plan_name: Optional[str] = None, +): total, H, K = q.shape - if k.shape[1] != H: - raise ValueError(f"k must carry the same head count as q ({H}), got {k.shape[1]}") + HK = k.shape[1] HV, V = v.shape[1], v.shape[2] + if HK not in (H, HV): + raise ValueError(f"k head count ({HK}) must match q's ({H}) or v's ({HV}); canonical GQA shares grouped k/v heads") HO = max(H, HV) N = cu_seqlens.shape[0] - 1 + if cu_seqlens.dtype not in (torch.int32, torch.int64): + raise ValueError(f"gated_delta_net: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") + 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) - return o, final + 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) + else: + state_checkpoints = q.new_empty(0) + return o, final, state_checkpoints # --------------------------------------------------------------------------- -# Backward +# Backward graph builder # --------------------------------------------------------------------------- -def _build_bwd_graph(total, N, H, HV, K, V, io_dtype, g_dtype, beta_dtype, state_dtype, dht_dtype, scale, use_qk_l2norm): +def _build_bprop_graph( + total, N, H, HK, HV, K, V, io_dtype, g_dtype, beta_dtype, state_dtype, dstate_in_dtype, cu_dtype, ckpt_rows, scale, use_qk_l2norm, batch_invariant +): graph = cudnn.pygraph() HO = max(H, HV) q_t = graph.tensor([total, H, K], data_type=io_dtype, name="q") - k_t = graph.tensor([total, H, K], data_type=io_dtype, name="k") + k_t = graph.tensor([total, HK, K], data_type=io_dtype, name="k") v_t = graph.tensor([total, HV, V], data_type=io_dtype, name="v") g_t = graph.tensor([total, HO], data_type=g_dtype, name="g") beta_t = graph.tensor([total, HO], data_type=beta_dtype, name="beta") - cu_t = graph.tensor([N + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") + cu_t = graph.tensor([N + 1], data_type=cu_dtype, name="cu_seqlens") dO_t = graph.tensor([total, HO, V], data_type=io_dtype, name="dO") - s0_t = None + state0_t = None if state_dtype is not None: - s0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") + state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") dfs_t = None - if dht_dtype is not None: - dfs_t = graph.tensor([N, HO, K, V], data_type=dht_dtype, name="d_final_state") - dQ_t, dK_t, dV_t, dG_t, dBeta_t, dS0_t = graph.gdn_bwd( + if dstate_in_dtype is not None: + dfs_t = graph.tensor([N, HO, K, V], 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") + dQ_t, dK_t, dV_t, dG_t, dBeta_t, dstate0_t = graph.gdn_bwd( q=q_t, k=k_t, v=v_t, @@ -263,10 +447,12 @@ def _build_bwd_graph(total, N, H, HV, K, V, io_dtype, g_dtype, beta_dtype, state beta=beta_t, cu_seqlens=cu_t, dO=dO_t, - initial_state=s0_t, + state_checkpoints=ckpts_t, + initial_state=state0_t, d_final_state=dfs_t, scale=scale, use_qk_l2norm=use_qk_l2norm, + batch_invariant=batch_invariant, name="gdn_bwd", ) return graph, dict( @@ -277,18 +463,24 @@ def _build_bwd_graph(total, N, H, HV, K, V, io_dtype, g_dtype, beta_dtype, state beta=beta_t, cu=cu_t, dO=dO_t, - s0=s0_t, + state0=state0_t, dfs=dfs_t, dQ=dQ_t, dK=dK_t, dV=dV_t, dG=dG_t, dBeta=dBeta_t, - dS0=dS0_t, + dstate0=dstate0_t, + ckpts=ckpts_t, ) -@torch.library.custom_op(f"{_OP_NAMESPACE}::{_OP_NAME}_bwd", mutates_args=()) +# --------------------------------------------------------------------------- +# Backward custom op +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("cudnn::gated_delta_net_bwd", mutates_args=()) def _gdn_bwd( dO: torch.Tensor, q: torch.Tensor, @@ -300,107 +492,170 @@ def _gdn_bwd( scale: float, initial_state: Optional[torch.Tensor] = None, d_final_state: Optional[torch.Tensor] = None, + state_checkpoints: Optional[torch.Tensor] = None, use_qk_l2norm_in_kernel: bool = False, + batch_invariant: bool = False, + plan_name: Optional[str] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """GDN backward via a cached single-node GDN_BWD pygraph (THD layout). + """GDN backward (internal): a cached single-node GDN_BWD pygraph, THD layout. - Returns ``(dq, dk, dv, dg, dbeta, d_initial_state)``; ``d_initial_state`` - is a zero-size tensor when ``initial_state`` is ``None``. + ``state_checkpoints`` is the forward's per-chunk state series (io dtype, + chunk cadence); when given, the engine consumes it instead of running + the checkpoint recompute pass. Returns ``(dq, dk, dv, dg, dbeta, + d_initial_state)``; ``d_initial_state`` is a zero-size tensor when + ``initial_state`` is ``None``. """ total, H, K = q.shape - if k.shape[1] != H: - raise ValueError(f"k must carry the same head count as q ({H}), got {k.shape[1]}") + # autograd materializes reduction grads as broadcast (stride-0) + # views; densify ONLY those (dense callers pass through untouched) + if 0 in dO.stride(): + dO = dO.contiguous() + if d_final_state is not None and 0 in d_final_state.stride(): + d_final_state = d_final_state.contiguous() + HK = k.shape[1] HV, V = v.shape[1], v.shape[2] + if HK not in (H, HV): + raise ValueError(f"k head count ({HK}) must match q's ({H}) or v's ({HV}); canonical GQA shares grouped k/v heads") N = cu_seqlens.shape[0] - 1 device = q.device - cu = cu_seqlens.to(torch.int32).contiguous() + if cu_seqlens.dtype not in (torch.int32, torch.int64): + raise ValueError(f"gated_delta_net: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") + cu = cu_seqlens _check_dtype("g", g, torch.float32) _check_dtype("beta", beta, torch.float32) if initial_state is not None: _check_dtype("initial_state", initial_state, torch.float32) if initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") - g32 = g.contiguous() - beta32 = beta.contiguous() - s0 = initial_state.contiguous() if initial_state is not None else None + if state_checkpoints is not None: + _check_dtype("state_checkpoints", state_checkpoints, q.dtype) + for _name, _t in ( + ("k", k), + ("v", v), + ("g", g), + ("beta", beta), + ("cu_seqlens", cu_seqlens), + ("dO", dO), + ("d_final_state", d_final_state), + ("state_checkpoints", state_checkpoints), + ): + if _t is not None and _t.device != device: + raise ValueError(f"gated_delta_net: {_name} must be on q's device ({device}); got {_t.device}") + g32 = g + beta32 = beta + state0 = initial_state if initial_state is not None else None if d_final_state is not None: _check_dtype("d_final_state", d_final_state, torch.float32) - dht = d_final_state.contiguous() if d_final_state is not None else None + dstate_in = d_final_state if d_final_state is not None else None - key = ( + cache_key = _make_bprop_cache_key( total, N, H, + HK, HV, K, V, q.dtype, - bool(s0 is not None), - bool(dht is not None), - float(scale), - bool(use_qk_l2norm_in_kernel), + k.dtype, + v.dtype, + dO.dtype, + tuple(k.shape), + tuple(v.shape), + cu_seqlens.dtype, + state0 is not None, + dstate_in is not None, + state_checkpoints.shape[0] if state_checkpoints is not None else None, + scale, + use_qk_l2norm_in_kernel, + batch_invariant, device, + plan_name, ) - if key not in _bwd_graph_cache: - _bwd_graph_cache[key] = _build_bwd_graph( + if cache_key not in _bprop_cache: + _bprop_cache[cache_key] = _build_bprop_graph( total, N, H, + HK, HV, K, V, - _cudnn_dtype(q.dtype), + _torch_dtype_to_cudnn(q.dtype), cudnn.data_type.FLOAT, cudnn.data_type.FLOAT, - cudnn.data_type.FLOAT if s0 is not None else None, - cudnn.data_type.FLOAT if dht is not None else None, + cudnn.data_type.FLOAT if state0 is not None else None, + cudnn.data_type.FLOAT if dstate_in is not None else None, + _torch_dtype_to_cudnn(cu_seqlens.dtype), + state_checkpoints.shape[0] if state_checkpoints is not None else None, float(scale), bool(use_qk_l2norm_in_kernel), + bool(batch_invariant), ) - graph, t = _bwd_graph_cache[key] + select_plan(_bprop_cache[cache_key][0], plan_name) + + graph, t = _bprop_cache[cache_key] HO = max(H, HV) dq = torch.empty(total, H, K, dtype=q.dtype, device=device) - dk = torch.empty(total, H, K, dtype=k.dtype, device=device) - dv = torch.empty(total, HV, V, dtype=v.dtype, device=device) + dk = torch.empty(total, HK, K, dtype=q.dtype, device=device) + dv = torch.empty(total, HV, V, dtype=q.dtype, device=device) dg32 = torch.empty(total, HO, dtype=torch.float32, device=device) dbeta32 = torch.empty(total, HO, dtype=torch.float32, device=device) variant_pack = { - t["q"]: q.contiguous(), - t["k"]: k.contiguous(), - t["v"]: v.contiguous(), + t["q"]: q, + t["k"]: k, + t["v"]: v, t["g"]: g32, t["beta"]: beta32, t["cu"]: cu, - t["dO"]: dO.contiguous(), + t["dO"]: dO, t["dQ"]: dq, t["dK"]: dk, t["dV"]: dv, t["dG"]: dg32, t["dBeta"]: dbeta32, } - dh032 = None - if s0 is not None: - variant_pack[t["s0"]] = s0 - dh032 = torch.empty_like(s0) - variant_pack[t["dS0"]] = dh032 - if dht is not None: - variant_pack[t["dfs"]] = dht - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_graph_handle(device)) - dh0 = dh032 if dh032 is not None else torch.empty(0, dtype=torch.float32, device=device) - return dq, dk, dv, dg32, dbeta32, dh0 + dstate0 = None + if state0 is not None: + variant_pack[t["state0"]] = state0 + dstate0 = torch.empty_like(state0) + variant_pack[t["dstate0"]] = dstate0 + if dstate_in is not None: + variant_pack[t["dfs"]] = dstate_in + if state_checkpoints is not None: + variant_pack[t["ckpts"]] = state_checkpoints + graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) + if dstate0 is None: + dstate0 = torch.empty(0, dtype=torch.float32, device=device) + return dq, dk, dv, dg32, dbeta32, dstate0 @_gdn_bwd.register_fake -def _gdn_bwd_fake(dO, q, k, v, g, beta, cu_seqlens, scale, initial_state=None, d_final_state=None, use_qk_l2norm_in_kernel=False): - dh0 = torch.empty_like(initial_state) if initial_state is not None else q.new_empty(0, dtype=torch.float32) +def _gdn_bwd_fake( + dO, + q, + k, + v, + g, + beta, + cu_seqlens, + scale, + initial_state=None, + d_final_state=None, + state_checkpoints=None, + use_qk_l2norm_in_kernel=False, + batch_invariant=False, + plan_name=None, +): + dstate0 = torch.empty_like(initial_state) if initial_state is not None else q.new_empty(0, dtype=torch.float32) return ( torch.empty_like(q), torch.empty_like(k), torch.empty_like(v), torch.empty_like(g), torch.empty_like(beta), - dh0, + dstate0, ) @@ -410,21 +665,37 @@ def _gdn_bwd_fake(dO, q, k, v, g, beta, cu_seqlens, scale, initial_state=None, d def _gdn_setup_context(ctx, inputs, output): - q, k, v, g, beta, cu_seqlens, scale, initial_state, output_final_state, use_qk_l2norm_in_kernel = inputs + q, k, v, g, beta, cu_seqlens, scale, initial_state, output_final_state, use_qk_l2norm_in_kernel, batch_invariant, checkpoint_every_n_tokens, plan_name = ( + inputs + ) # save_for_backward cannot hold None; keep initial_state as an attribute. - ctx.save_for_backward(q, k, v, g, beta, cu_seqlens) + saved = [q, k, v, g, beta, cu_seqlens] + ctx.ckpt_reuse = checkpoint_every_n_tokens == 64 and output[2].numel() > 0 + if ctx.ckpt_reuse: + saved.append(output[2]) + ctx.save_for_backward(*saved) ctx.initial_state = initial_state ctx.scale = scale ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - - -def _gdn_backward(ctx, dO, dFinal): - q, k, v, g, beta, cu_seqlens = ctx.saved_tensors + ctx.batch_invariant = batch_invariant + ctx.plan_name = plan_name + ctx.set_materialize_grads(False) + ctx.mark_non_differentiable(output[2]) + + +def _gdn_backward(ctx, dO, dFinal, _dstate_checkpoints): + if ctx.ckpt_reuse: + q, k, v, g, beta, cu_seqlens, state_checkpoints = ctx.saved_tensors + else: + q, k, v, g, beta, cu_seqlens = ctx.saved_tensors + state_checkpoints = None initial_state = ctx.initial_state - dht = dFinal if (dFinal is not None and dFinal.numel() > 0) else None - dq, dk, dv, dg, dbeta, dh0 = torch.ops.cudnn.gated_delta_net_bwd( - dO.contiguous(), + if dO is None: + dO = torch.zeros(q.shape[0], max(q.shape[1], v.shape[1]), v.shape[2], dtype=q.dtype, device=q.device) + dstate_in = dFinal if (dFinal is not None and dFinal.numel() > 0) else None + dq, dk, dv, dg, dbeta, dstate0 = torch.ops.cudnn.gated_delta_net_bwd( + dO, q, k, v, @@ -433,11 +704,14 @@ def _gdn_backward(ctx, dO, dFinal): cu_seqlens, ctx.scale, initial_state=initial_state, - d_final_state=dht, + d_final_state=dstate_in, + state_checkpoints=state_checkpoints, use_qk_l2norm_in_kernel=ctx.use_qk_l2norm_in_kernel, + batch_invariant=ctx.batch_invariant, + plan_name=ctx.plan_name, ) # q, k, v, g, beta, cu_seqlens, scale, initial_state, output_final_state, - # use_qk_l2norm_in_kernel + # use_qk_l2norm_in_kernel, batch_invariant, checkpoint_every_n_tokens, plan_name return ( dq, dk, @@ -446,14 +720,17 @@ def _gdn_backward(ctx, dO, dFinal): dbeta, None, None, - dh0 if initial_state is not None else None, + dstate0 if initial_state is not None else None, + None, + None, + None, None, None, ) torch.library.register_autograd( - f"{_OP_NAMESPACE}::{_OP_NAME}_fwd", + "cudnn::gated_delta_net_fwd", _gdn_backward, setup_context=_gdn_setup_context, ) @@ -475,14 +752,18 @@ def gated_delta_net( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, + batch_invariant: bool = False, + checkpoint_every_n_tokens: int = 0, + plan_name: Optional[str] = None, ): """Gated DeltaNet (GDN) linear attention. THD layout (matches the graph-API GDN node): - q, k: ``[total_tokens, H, K]``; v: ``[total_tokens, HV, V]`` + q: ``[total_tokens, H, K]``; k: ``[total_tokens, HK, K]`` (HK = H, or + 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 + cu_seqlens: ``[N+1]`` int32; O and the states live at HO heads (initial_state / final_state: ``[N, HO, K, V]``) A dense batch of N equal-length sequences is expressed as @@ -500,18 +781,30 @@ def gated_delta_net( initial_state: optional recurrent state (otherwise zero). output_final_state: if ``True``, also return the per-sequence state after the last token. - use_qk_l2norm_in_kernel: if ``True``, L2-normalize the q/k rows inside + use_qk_l2norm_in_kernel: if ``True``, L2-normalize the Q/K rows inside the kernel. Engines that cannot honor it decline the graph. - + batch_invariant: if ``True``, each sequence's results are bitwise + independent of the batch composition (whole-sequence scheduling; + disables split-K load balancing). + checkpoint_every_n_tokens: if ``> 0``, also return the per-chunk + recurrent state series ``state_checkpoints`` (``[total_checkpoints, HO, K, V]`` 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. + + plan_name: optionally pin one execution plan by name (the plan + API's ``get_plan_name_at_index`` names, e.g. ``gdn_frost``); a + graph offering no such plan raises ``cudnnGraphNotSupportedError``. Returns: - ``(o, final_state)`` with ``o`` shaped like ``v``. ``final_state`` is - empty unless ``output_final_state=True``. + ``(o, final_state)`` with ``o`` shaped like ``v``, or + ``(o, final_state, state_checkpoints)`` when ``checkpoint_every_n_tokens > 0``. + ``final_state`` is empty unless ``output_final_state=True``. """ if q.dim() != 3: raise ValueError("expected THD [total_tokens, heads, dim] tensors") if scale is None: scale = 1.0 / math.sqrt(q.shape[-1]) - return torch.ops.cudnn.gated_delta_net_fwd( + o, final_state, state_checkpoints = torch.ops.cudnn.gated_delta_net_fwd( q, k, v, @@ -522,4 +815,10 @@ def gated_delta_net( initial_state=initial_state, output_final_state=bool(output_final_state), use_qk_l2norm_in_kernel=bool(use_qk_l2norm_in_kernel), + batch_invariant=bool(batch_invariant), + checkpoint_every_n_tokens=int(checkpoint_every_n_tokens), + plan_name=plan_name, ) + if checkpoint_every_n_tokens > 0: + return o, final_state, state_checkpoints + return o, final_state diff --git a/python/cudnn/linear_attention/ops/gdn2.py b/python/cudnn/linear_attention/ops/gdn2.py index 57ca61dd3..33da40d16 100644 --- a/python/cudnn/linear_attention/ops/gdn2.py +++ b/python/cudnn/linear_attention/ops/gdn2.py @@ -9,106 +9,259 @@ S_t = (I - k_t (beta_t . k_t)^T) Diag(exp(g_t)) S_{t-1} + k_t (w_t . v_t)^T, o_t = scale * q_t S_t, -with per-key-channel log decay ``g_t in R^K``, per-key erase gate +where ``S_t`` is the recurrent state, with per-key-channel log decay +``g_t in R^K``, per-key erase gate ``beta_t in R^K`` (applied inside the erase read of the decayed state), and a per-value write gate ``w_t in R^V`` (applied to the written value). Layout follows the graph-API GDN-2 node: THD — token-packed -``[total_tokens, heads, dim]`` q/k/v, ``g``/``beta`` ``[total_tokens, HO, K]``, +``[total_tokens, heads, dim]`` Q/K/V, ``g``/``beta`` ``[total_tokens, HO, K]``, ``w`` ``[total_tokens, HO, V]``, plus ``cu_seqlens`` boundaries. -The op is a thin adapter over the graph API (the SDPA-op pattern): forward -executes a cached single-node ``GDN2`` pygraph. Engine selection happens at -graph planning time over the registered python engines — ``Gdn2FrostEngine`` -is the only GDN-2 engine (SM100/SM103). The op is **forward only** (the -FROST GDN-2 backward kernel is a stub), so no autograd is registered; -differentiating through it raises. Registered through -``torch.library.custom_op`` so it composes with ``torch.compile``. -""" +The op is a thin adapter over the graph API: forward +and backward execute cached single-node ``GDN2`` / ``GDN2_BWD`` pygraphs. +Engine selection happens at graph planning time over the registered python +engines — ``Gdn2FrostEngine`` is the only GDN-2 engine (SM100/SM103). +Registered through ``torch.library.custom_op`` so it composes with autograd, +``torch.compile``, and DDP. -from __future__ import annotations +Graph caching ensures cuDNN graphs are built once per unique configuration +and reused across calls. +""" import math from typing import Dict, Optional, Tuple import torch - import cudnn -from cudnn.linear_attention import engine_utils -_OP_NAMESPACE = "cudnn" -_OP_NAME = "gated_delta_net_v2" +# --------------------------------------------------------------------------- +# Module-level state +# --------------------------------------------------------------------------- + -_TORCH_TO_CUDNN_DTYPE = { +_TORCH_DTYPE_TO_CUDNN = { torch.float16: cudnn.data_type.HALF, torch.bfloat16: cudnn.data_type.BFLOAT16, torch.float32: cudnn.data_type.FLOAT, + torch.int32: cudnn.data_type.INT32, + torch.int64: cudnn.data_type.INT64, } # one graph per static configuration (shapes, dtypes, scale, flags, device) -_fwd_graph_cache: Dict[tuple, tuple] = {} -_ws_cache: Dict[int, "torch.Tensor"] = {} +_fprop_cache: Dict[tuple, tuple] = {} +_bprop_cache: Dict[tuple, tuple] = {} +_cudnn_handles: Dict[int, int] = {} + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def select_plan(graph, plan_name): + """Pin one execution plan by name on a freshly built graph (create the + plans, select by name, check support); ``None`` keeps default routing.""" + if plan_name is None: + return + graph.create_execution_plans() + names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] + matches = [i for i, n in enumerate(names) if n == plan_name or n.startswith(plan_name + "[")] + if not matches: + raise cudnn.cudnnGraphNotSupportedError(f"no {plan_name} plan for this graph (offered: {names})") + graph.select_plan(matches[0]) + graph.check_support() def _graph_workspace(graph, device): - """Caller-side workspace for a compiled graph (the explicit-workspace - convention: query the plan's size, allocate, pass to execute).""" + """Caller-side workspace for a compiled graph.""" if not graph._is_built: - # mirror execute()'s auto-build: plan via the router first (a bare - # build() would lower GDN2 to the backend, which has no lowering) + # mirror execute()'s auto-build: plan first (a bare build() would + # lower GDN2 to the backend, which has no lowering) if not graph._planning_done: graph.create_execution_plans() - engine_utils.apply_pin(graph) if graph.selected_engine is None: graph.build() else: graph.build_plans() size = graph.get_workspace_size() - ws = _ws_cache.get(id(graph)) + ws = getattr(graph, "_la_ops_workspace", None) if ws is None or ws.numel() < size or ws.device != device: ws = torch.empty(max(size, 1), dtype=torch.uint8, device=device) - _ws_cache[id(graph)] = ws + graph._la_ops_workspace = ws return ws -_handle_cache: Dict[int, int] = {} - - -def _graph_handle(device): - """Per-device cuDNN handle carrying the caller's current stream - (classic ``set_stream`` semantics).""" +def _get_handle(device): + """Per-device cuDNN handle carrying the caller's current stream.""" idx = device.index if device.index is not None else torch.cuda.current_device() - handle = _handle_cache.get(idx) + handle = _cudnn_handles.get(idx) if handle is None: with torch.cuda.device(idx): handle = cudnn.create_handle() - _handle_cache[idx] = handle + _cudnn_handles[idx] = handle cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) return handle -def _cudnn_dtype(dtype: Optional[torch.dtype]): - return _TORCH_TO_CUDNN_DTYPE[dtype] if dtype is not None else None +def _torch_dtype_to_cudnn(dtype: torch.dtype): + """Map a PyTorch dtype to a cuDNN data_type enum.""" + return _TORCH_DTYPE_TO_CUDNN[dtype] def _check_dtype(name, t, want) -> None: if t.dtype != want: - raise TypeError(f"{_OP_NAME}: {name} must be {want} (kernel-native; callers convert), got {t.dtype}") + raise TypeError(f"gated_delta_net_v2: {name} must be {want} (kernel-native; callers convert), got {t.dtype}") + + +def _make_fprop_cache_key( + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + k_shape, + v_shape, + cu_dtype, + scale, + output_final_state, + use_qk_l2norm, + batch_invariant, + safe_gate, + gate_lower_bound, + has_initial_state, + ckpt, + device, + plan_name, +): + return ( + "fprop", + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + k_shape, + v_shape, + cu_dtype, + float(scale), + bool(output_final_state), + bool(use_qk_l2norm), + bool(batch_invariant), + bool(safe_gate), + float(gate_lower_bound) if gate_lower_bound is not None else None, + bool(has_initial_state), + ckpt, + device, + plan_name, + ) + + +def _make_bprop_cache_key( + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + do_dtype, + k_shape, + v_shape, + cu_dtype, + state_dtype, + dstate_in_dtype, + ckpt_rows, + scale, + use_qk_l2norm, + batch_invariant, + device, + plan_name, +): + return ( + "bprop", + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + do_dtype, + k_shape, + v_shape, + cu_dtype, + state_dtype, + dstate_in_dtype, + ckpt_rows, + float(scale), + bool(use_qk_l2norm), + bool(batch_invariant), + device, + plan_name, + ) + + +# --------------------------------------------------------------------------- +# Forward graph builder +# --------------------------------------------------------------------------- -def _build_fwd_graph(total, N, H, HV, K, V, io_dtype, g_dtype, gate_dtype, state_dtype, scale, output_final_state, use_qk_l2norm): +def _build_fprop_graph( + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + g_dtype, + gate_dtype, + state_dtype, + cu_dtype, + scale, + output_final_state, + use_qk_l2norm, + batch_invariant, + safe_gate, + gate_lower_bound, + ckpt, +): graph = cudnn.pygraph() + HO = max(H, HV) q_t = graph.tensor([total, H, K], data_type=io_dtype, name="q") - k_t = graph.tensor([total, H, K], data_type=io_dtype, name="k") + k_t = graph.tensor([total, HK, K], data_type=io_dtype, name="k") v_t = graph.tensor([total, HV, V], data_type=io_dtype, name="v") - g_t = graph.tensor([total, HV, K], data_type=g_dtype, name="g") - beta_t = graph.tensor([total, HV, K], data_type=gate_dtype, name="beta") - w_t = graph.tensor([total, HV, V], data_type=gate_dtype, name="w") - cu_t = graph.tensor([N + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") - s0_t = None + g_t = graph.tensor([total, HO, K], data_type=g_dtype, name="g") + beta_t = graph.tensor([total, HO, K], data_type=gate_dtype, name="beta") + w_t = graph.tensor([total, HO, V], data_type=gate_dtype, name="w") + cu_t = graph.tensor([N + 1], data_type=cu_dtype, name="cu_seqlens") + state0_t = None if state_dtype is not None: - s0_t = graph.tensor([N, HV, K, V], data_type=state_dtype, name="initial_state") - O_t, fs_t, _h_t = graph.gdn2( + state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") + a_log_t = None + dt_bias_t = None + if safe_gate: + a_log_t = graph.tensor([HO], data_type=cudnn.data_type.FLOAT, name="a_log") + dt_bias_t = graph.tensor([HO, K], data_type=cudnn.data_type.FLOAT, name="dt_bias") + O_t, fs_t, state_checkpoints_t = graph.gdn2( q=q_t, k=k_t, v=v_t, @@ -116,16 +269,41 @@ def _build_fwd_graph(total, N, H, HV, K, V, io_dtype, g_dtype, gate_dtype, state beta=beta_t, w=w_t, cu_seqlens=cu_t, - initial_state=s0_t, + initial_state=state0_t, + a_log=a_log_t, + dt_bias=dt_bias_t, scale=scale, output_final_state=output_final_state, use_qk_l2norm=use_qk_l2norm, + batch_invariant=batch_invariant, + safe_gate=safe_gate, + gate_lower_bound=gate_lower_bound, + checkpoint_every_n_tokens=ckpt, name="gdn2", ) - return graph, dict(q=q_t, k=k_t, v=v_t, g=g_t, beta=beta_t, w=w_t, cu=cu_t, s0=s0_t, O=O_t, fs=fs_t) + return graph, dict( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + w=w_t, + cu=cu_t, + state0=state0_t, + a_log=a_log_t, + dt_bias=dt_bias_t, + O=O_t, + fs=fs_t, + state_checkpoints=state_checkpoints_t, + ) -@torch.library.custom_op(f"{_OP_NAMESPACE}::{_OP_NAME}_fwd", mutates_args=()) +# --------------------------------------------------------------------------- +# Forward custom op +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("cudnn::gated_delta_net_v2_fwd", mutates_args=()) def _gdn2_fwd( q: torch.Tensor, k: torch.Tensor, @@ -138,92 +316,545 @@ def _gdn2_fwd( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor]: - """GDN-2 forward via a cached single-node GDN2 pygraph (THD layout). - - Returns ``(o, final_state)``; ``final_state`` is a zero-size tensor when - ``output_final_state`` is ``False``. + batch_invariant: bool = False, + safe_gate: bool = False, + gate_lower_bound: Optional[float] = None, + a_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + checkpoint_every_n_tokens: int = 0, + plan_name: Optional[str] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """GDN-2 forward (internal): a cached single-node GDN2 pygraph, THD layout. + + Returns ``(o, final_state, state_checkpoints)``; ``final_state`` / ``state_checkpoints`` are zero-size + tensors when ``output_final_state`` is ``False`` / + ``checkpoint_every_n_tokens`` is ``0``. """ total, H, K = q.shape - if k.shape[1] != H: - raise ValueError(f"k must carry the same head count as q ({H}), got {k.shape[1]}") + HK = k.shape[1] HV, V = v.shape[1], v.shape[2] + if HK not in (H, HV): + raise ValueError(f"k head count ({HK}) must match q's ({H}) or v's ({HV}); canonical GQA shares grouped k/v heads") + HO = max(H, HV) N = cu_seqlens.shape[0] - 1 device = q.device - cu = cu_seqlens.to(torch.int32).contiguous() + if cu_seqlens.dtype not in (torch.int32, torch.int64): + raise ValueError(f"gated_delta_net_v2: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") + cu = cu_seqlens _check_dtype("g", g, torch.float32) _check_dtype("beta", beta, q.dtype) _check_dtype("w", w, q.dtype) + if safe_gate: + if a_log is None or dt_bias is None: + raise ValueError("gated_delta_net_v2: safe_gate requires a_log and dt_bias") + _check_dtype("a_log", a_log, torch.float32) + _check_dtype("dt_bias", dt_bias, torch.float32) + elif a_log is not None or dt_bias is not None: + raise ValueError("gated_delta_net_v2: a_log/dt_bias require safe_gate=True") if initial_state is not None: _check_dtype("initial_state", initial_state, torch.float32) if initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") - g32 = g.contiguous() - beta_io = beta.contiguous() - w_io = w.contiguous() - s0 = initial_state.contiguous() if initial_state is not None else None - - key = ( + g32 = g + beta_io = beta + w_io = w + for _name, _t in ( + ("k", k), + ("v", v), + ("g", g), + ("beta", beta), + ("w", w), + ("cu_seqlens", cu_seqlens), + ("initial_state", initial_state), + ("a_log", a_log), + ("dt_bias", dt_bias), + ): + if _t is not None and _t.device != device: + raise ValueError(f"gated_delta_net_v2: {_name} must be on q's device ({device}); got {_t.device}") + state0 = initial_state if initial_state is not None else None + ckpt = int(checkpoint_every_n_tokens) + + cache_key = _make_fprop_cache_key( total, N, H, + HK, HV, K, V, q.dtype, - bool(s0 is not None), - float(scale), - bool(output_final_state), - bool(use_qk_l2norm_in_kernel), + k.dtype, + v.dtype, + tuple(k.shape), + tuple(v.shape), + cu_seqlens.dtype, + scale, + output_final_state, + use_qk_l2norm_in_kernel, + batch_invariant, + safe_gate, + gate_lower_bound, + state0 is not None, + ckpt, device, + plan_name, ) - if key not in _fwd_graph_cache: - _fwd_graph_cache[key] = _build_fwd_graph( + if cache_key not in _fprop_cache: + _fprop_cache[cache_key] = _build_fprop_graph( total, N, H, + HK, HV, K, V, - _cudnn_dtype(q.dtype), + _torch_dtype_to_cudnn(q.dtype), cudnn.data_type.FLOAT, - _cudnn_dtype(q.dtype), - cudnn.data_type.FLOAT if s0 is not None else None, + _torch_dtype_to_cudnn(q.dtype), + cudnn.data_type.FLOAT if state0 is not None else None, + _torch_dtype_to_cudnn(cu_seqlens.dtype), float(scale), bool(output_final_state), bool(use_qk_l2norm_in_kernel), + bool(batch_invariant), + bool(safe_gate), + float(gate_lower_bound) if gate_lower_bound is not None else None, + ckpt, ) - graph, t = _fwd_graph_cache[key] + select_plan(_fprop_cache[cache_key][0], plan_name) + + graph, t = _fprop_cache[cache_key] - o = torch.empty(total, HV, V, dtype=q.dtype, device=device) + o = torch.empty(total, HO, V, dtype=q.dtype, device=device) variant_pack = { - t["q"]: q.contiguous(), - t["k"]: k.contiguous(), - t["v"]: v.contiguous(), + t["q"]: q, + t["k"]: k, + t["v"]: v, t["g"]: g32, t["beta"]: beta_io, t["w"]: w_io, t["cu"]: cu, t["O"]: o, } - if s0 is not None: - variant_pack[t["s0"]] = s0 + if state0 is not None: + variant_pack[t["state0"]] = state0 + if safe_gate: + variant_pack[t["a_log"]] = a_log + 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, HV, K, V, dtype=torch.float32, device=device) + final_state = torch.empty(N, HO, K, V, dtype=torch.float32, device=device) variant_pack[t["fs"]] = final_state - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_graph_handle(device)) - return o, 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) + 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 @_gdn2_fwd.register_fake -def _gdn2_fwd_fake(q, k, v, g, beta, w, cu_seqlens, scale, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False): - total, _H, K = q.shape +def _gdn2_fwd_fake( + q, + k, + v, + g, + beta, + w, + cu_seqlens, + scale, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + batch_invariant=False, + safe_gate=False, + gate_lower_bound=None, + a_log=None, + dt_bias=None, + checkpoint_every_n_tokens=0, + plan_name: Optional[str] = None, +): + total, H, K = q.shape + HK = k.shape[1] + HV, V = v.shape[1], v.shape[2] + if HK not in (H, HV): + raise ValueError(f"k head count ({HK}) must match q's ({H}) or v's ({HV}); canonical GQA shares grouped k/v heads") + HO = max(H, HV) + N = cu_seqlens.shape[0] - 1 + if cu_seqlens.dtype not in (torch.int32, torch.int64): + raise ValueError(f"gated_delta_net_v2: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") + 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) + 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) + else: + state_checkpoints = q.new_empty(0) + return o, final, state_checkpoints + + +# --------------------------------------------------------------------------- +# Backward graph builder +# --------------------------------------------------------------------------- + + +def _build_bprop_graph( + total, N, H, HK, HV, K, V, io_dtype, g_dtype, gate_dtype, state_dtype, dstate_in_dtype, cu_dtype, ckpt_rows, scale, use_qk_l2norm, batch_invariant +): + graph = cudnn.pygraph() + HO = max(H, HV) + q_t = graph.tensor([total, H, K], data_type=io_dtype, name="q") + k_t = graph.tensor([total, HK, K], data_type=io_dtype, name="k") + v_t = graph.tensor([total, HV, V], data_type=io_dtype, name="v") + g_t = graph.tensor([total, HO, K], data_type=g_dtype, name="g") + beta_t = graph.tensor([total, HO, K], data_type=gate_dtype, name="beta") + w_t = graph.tensor([total, HO, V], data_type=gate_dtype, name="w") + cu_t = graph.tensor([N + 1], data_type=cu_dtype, name="cu_seqlens") + 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") + 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") + 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") + dQ_t, dK_t, dV_t, dG_t, dBeta_t, dW_t, dstate0_t = graph.gdn2_bwd( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + w=w_t, + cu_seqlens=cu_t, + dO=dO_t, + state_checkpoints=ckpts_t, + initial_state=state0_t, + d_final_state=dfs_t, + scale=scale, + use_qk_l2norm=use_qk_l2norm, + batch_invariant=batch_invariant, + name="gdn2_bwd", + ) + return graph, dict( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + w=w_t, + cu=cu_t, + dO=dO_t, + state0=state0_t, + dfs=dfs_t, + dQ=dQ_t, + dK=dK_t, + dV=dV_t, + dG=dG_t, + dBeta=dBeta_t, + dW=dW_t, + dstate0=dstate0_t, + ckpts=ckpts_t, + ) + + +# --------------------------------------------------------------------------- +# Backward custom op +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("cudnn::gated_delta_net_v2_bwd", mutates_args=()) +def _gdn2_bwd( + dO: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + w: torch.Tensor, + cu_seqlens: torch.Tensor, + scale: float, + initial_state: Optional[torch.Tensor] = None, + d_final_state: Optional[torch.Tensor] = None, + state_checkpoints: Optional[torch.Tensor] = None, + use_qk_l2norm_in_kernel: bool = False, + batch_invariant: bool = False, + plan_name: Optional[str] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """GDN-2 backward (internal): a cached single-node GDN2_BWD pygraph, THD layout. + + ``state_checkpoints`` is the forward's per-chunk state series (io dtype, + chunk cadence); when given, the engine consumes it instead of running + the checkpoint recompute pass. Returns ``(dq, dk, dv, dg, dbeta, dw, + d_initial_state)``; ``d_initial_state`` is a zero-size tensor when + ``initial_state`` is ``None``. + """ + total, H, K = q.shape + # autograd materializes reduction grads as broadcast (stride-0) + # views; densify ONLY those (dense callers pass through untouched) + if 0 in dO.stride(): + dO = dO.contiguous() + if d_final_state is not None and 0 in d_final_state.stride(): + d_final_state = d_final_state.contiguous() + HK = k.shape[1] HV, V = v.shape[1], v.shape[2] + if HK not in (H, HV): + raise ValueError(f"k head count ({HK}) must match q's ({H}) or v's ({HV}); canonical GQA shares grouped k/v heads") + HO = max(H, HV) N = cu_seqlens.shape[0] - 1 - o = q.new_empty(total, HV, V) - final = q.new_empty((N, HV, K, V) if output_final_state else (0,), dtype=torch.float32) - return o, final + device = q.device + if cu_seqlens.dtype not in (torch.int32, torch.int64): + raise ValueError(f"gated_delta_net_v2: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") + cu = cu_seqlens + _check_dtype("g", g, torch.float32) + _check_dtype("beta", beta, q.dtype) + _check_dtype("w", w, q.dtype) + if initial_state is not None: + _check_dtype("initial_state", initial_state, torch.float32) + if initial_state.shape[0] != N: + raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") + if d_final_state is not None: + _check_dtype("d_final_state", d_final_state, torch.float32) + if state_checkpoints is not None: + _check_dtype("state_checkpoints", state_checkpoints, q.dtype) + for _name, _t in ( + ("k", k), + ("v", v), + ("g", g), + ("beta", beta), + ("w", w), + ("cu_seqlens", cu_seqlens), + ("dO", dO), + ("d_final_state", d_final_state), + ("state_checkpoints", state_checkpoints), + ): + if _t is not None and _t.device != device: + raise ValueError(f"gated_delta_net_v2: {_name} must be on q's device ({device}); got {_t.device}") + state0 = initial_state if initial_state is not None else None + dstate_in = d_final_state if d_final_state is not None else None + + cache_key = _make_bprop_cache_key( + total, + N, + H, + HK, + HV, + K, + V, + q.dtype, + k.dtype, + v.dtype, + dO.dtype, + tuple(k.shape), + tuple(v.shape), + cu_seqlens.dtype, + state0.dtype if state0 is not None else None, + dstate_in.dtype if dstate_in is not None else None, + state_checkpoints.shape[0] if state_checkpoints is not None else None, + scale, + use_qk_l2norm_in_kernel, + batch_invariant, + device, + plan_name, + ) + if cache_key not in _bprop_cache: + _bprop_cache[cache_key] = _build_bprop_graph( + total, + N, + H, + HK, + HV, + K, + V, + _torch_dtype_to_cudnn(q.dtype), + cudnn.data_type.FLOAT, + _torch_dtype_to_cudnn(q.dtype), + _torch_dtype_to_cudnn(state0.dtype) if state0 is not None else None, + _torch_dtype_to_cudnn(dstate_in.dtype) if dstate_in is not None else None, + _torch_dtype_to_cudnn(cu_seqlens.dtype), + state_checkpoints.shape[0] if state_checkpoints is not None else None, + float(scale), + bool(use_qk_l2norm_in_kernel), + bool(batch_invariant), + ) + select_plan(_bprop_cache[cache_key][0], plan_name) + + graph, t = _bprop_cache[cache_key] + + dq = torch.empty(total, H, K, dtype=q.dtype, device=device) + dk = torch.empty(total, HK, K, dtype=q.dtype, device=device) + dv = torch.empty(total, HV, V, dtype=q.dtype, device=device) + dg = torch.empty(total, HO, K, dtype=torch.float32, device=device) + dbeta = torch.empty(total, HO, K, dtype=beta.dtype, device=device) + dw = torch.empty(total, HO, V, dtype=w.dtype, device=device) + variant_pack = { + t["q"]: q, + t["k"]: k, + t["v"]: v, + t["g"]: g, + t["beta"]: beta, + t["w"]: w, + t["cu"]: cu, + t["dO"]: dO, + t["dQ"]: dq, + t["dK"]: dk, + t["dV"]: dv, + t["dG"]: dg, + t["dBeta"]: dbeta, + t["dW"]: dw, + } + dstate0 = None + if state0 is not None: + variant_pack[t["state0"]] = state0 + dstate0 = torch.empty_like(state0) + variant_pack[t["dstate0"]] = dstate0 + if dstate_in is not None: + variant_pack[t["dfs"]] = dstate_in + if state_checkpoints is not None: + variant_pack[t["ckpts"]] = state_checkpoints + graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) + if dstate0 is None: + dstate0 = torch.empty(0, dtype=torch.float32, device=device) + return dq, dk, dv, dg, dbeta, dw, dstate0 + + +@_gdn2_bwd.register_fake +def _gdn2_bwd_fake( + dO, + q, + k, + v, + g, + beta, + w, + cu_seqlens, + scale, + initial_state=None, + d_final_state=None, + state_checkpoints=None, + use_qk_l2norm_in_kernel=False, + batch_invariant=False, + plan_name=None, +): + dstate0 = torch.empty_like(initial_state) if initial_state is not None else q.new_empty(0, dtype=torch.float32) + return ( + torch.empty_like(q), + torch.empty_like(k), + torch.empty_like(v), + g.new_empty(g.shape, dtype=torch.float32), + torch.empty_like(beta), + torch.empty_like(w), + dstate0, + ) + + +# --------------------------------------------------------------------------- +# Autograd registration +# --------------------------------------------------------------------------- + + +def _gdn2_setup_context(ctx, inputs, output): + ( + q, + k, + v, + g, + beta, + w, + cu_seqlens, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + batch_invariant, + safe_gate, + gate_lower_bound, + a_log, + dt_bias, + checkpoint_every_n_tokens, + plan_name, + ) = inputs + # save_for_backward cannot hold None; keep initial_state as an attribute. + saved = [q, k, v, g, beta, w, cu_seqlens] + ctx.ckpt_reuse = checkpoint_every_n_tokens == 16 and output[2].numel() > 0 + if ctx.ckpt_reuse: + saved.append(output[2]) + ctx.save_for_backward(*saved) + ctx.initial_state = initial_state + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.batch_invariant = batch_invariant + ctx.plan_name = plan_name + ctx.safe_gate = bool(safe_gate) + ctx.set_materialize_grads(False) + ctx.mark_non_differentiable(output[2]) + + +def _gdn2_backward(ctx, dO, dFinal, _dstate_checkpoints): + if ctx.safe_gate: + raise NotImplementedError("gated_delta_net_v2: safe_gate is forward-only (GDN2_BWD takes post-activation gates)") + if ctx.ckpt_reuse: + q, k, v, g, beta, w, cu_seqlens, state_checkpoints = ctx.saved_tensors + else: + q, k, v, g, beta, w, cu_seqlens = ctx.saved_tensors + state_checkpoints = None + initial_state = ctx.initial_state + + if dO is None: + dO = torch.zeros(q.shape[0], max(q.shape[1], v.shape[1]), v.shape[2], dtype=q.dtype, device=q.device) + dstate_in = dFinal if (dFinal is not None and dFinal.numel() > 0) else None + dq, dk, dv, dg, dbeta, dw, dstate0 = torch.ops.cudnn.gated_delta_net_v2_bwd( + dO, + q, + k, + v, + g, + beta, + w, + cu_seqlens, + ctx.scale, + initial_state=initial_state, + d_final_state=dstate_in, + state_checkpoints=state_checkpoints, + use_qk_l2norm_in_kernel=ctx.use_qk_l2norm_in_kernel, + batch_invariant=ctx.batch_invariant, + plan_name=ctx.plan_name, + ) + # q, k, v, g, beta, w, cu_seqlens, scale, initial_state, + # output_final_state, use_qk_l2norm_in_kernel, batch_invariant, + # safe_gate, gate_lower_bound, a_log, dt_bias, + # checkpoint_every_n_tokens, plan_name + return ( + dq, + dk, + dv, + dg, + dbeta, + dw, + None, + None, + dstate0 if initial_state is not None else None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +torch.library.register_autograd( + "cudnn::gated_delta_net_v2_fwd", + _gdn2_backward, + setup_context=_gdn2_setup_context, +) # --------------------------------------------------------------------------- @@ -243,41 +874,69 @@ def gated_delta_net_v2( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Gated DeltaNet v2 (GDN-2) linear attention — forward only. + batch_invariant: bool = False, + safe_gate: bool = False, + gate_lower_bound: Optional[float] = None, + a_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + checkpoint_every_n_tokens: int = 0, + plan_name: Optional[str] = None, +): + """Gated DeltaNet v2 (GDN-2) linear attention. THD layout (matches the graph-API GDN-2 node): - q, k: ``[total_tokens, H, K]``; v: ``[total_tokens, HV, V]``; - g, beta: ``[total_tokens, HV, K]``; w: ``[total_tokens, HV, V]``; + q: ``[total_tokens, H, K]``; k: ``[total_tokens, HK, K]`` (HK = H, or + 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, HV, K, V]``. + initial_state / final_state: ``[N, HO, K, V]`` + (``HO = max(H, HV)``: the gates, output, and state heads). A dense batch of N equal-length sequences is expressed as ``cu_seqlens = [0, T, 2T, ...]`` over the flattened tokens. Args: - g: per-key-channel log-space decay (``alpha = exp(g)``). - beta: per-key erase gate. - w: per-value write gate. + g: per-key-channel log-space decay (``alpha = exp(g)``), or raw + pre-activation logits when ``safe_gate=True``. + beta: per-key erase gate (io dtype, post-activation). + w: per-value write gate (io dtype, post-activation). cu_seqlens: ``[N+1]`` int32 sequence boundaries over the packed tokens. scale: attention scale applied to ``q``. Defaults to ``1 / sqrt(K)``. initial_state: optional recurrent state (otherwise zero). output_final_state: if ``True``, also return the per-sequence state after the last token. - use_qk_l2norm_in_kernel: if ``True``, L2-normalize the q/k rows inside - the kernel; if ``False``, pass q/k as given (the caller owns their + use_qk_l2norm_in_kernel: if ``True``, L2-normalize the Q/K rows inside + the kernel; if ``False``, pass Q/K as given (the caller owns their conditioning). - + batch_invariant: if ``True``, each sequence's results are bitwise + independent of the batch composition (whole-sequence scheduling; + disables split-K load balancing). + safe_gate: interpret ``g`` through the safe-gate transform + ``gate_lower_bound * sigmoid(exp(a_log) * (g + dt_bias))``. + Requires ``a_log`` and ``dt_bias``. Forward-only. + gate_lower_bound: safe-gate lower bound in log space (default -5.0). + 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, + 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. + + plan_name: optionally pin one execution plan by name (the plan + API's ``get_plan_name_at_index`` names, e.g. ``gdn2_frost``); a + graph offering no such plan raises ``cudnnGraphNotSupportedError``. Returns: - ``(o, final_state)`` with ``o`` shaped like ``v``. ``final_state`` is - empty unless ``output_final_state=True``. + ``(o, final_state)`` with ``o`` shaped like ``v``, or + ``(o, final_state, state_checkpoints)`` when ``checkpoint_every_n_tokens > 0``. + ``final_state`` is empty unless ``output_final_state=True``. """ if q.dim() != 3: raise ValueError("expected THD [total_tokens, heads, dim] tensors") if scale is None: scale = 1.0 / math.sqrt(q.shape[-1]) - return torch.ops.cudnn.gated_delta_net_v2_fwd( + o, final_state, state_checkpoints = torch.ops.cudnn.gated_delta_net_v2_fwd( q, k, v, @@ -289,4 +948,14 @@ def gated_delta_net_v2( initial_state=initial_state, output_final_state=bool(output_final_state), use_qk_l2norm_in_kernel=bool(use_qk_l2norm_in_kernel), + batch_invariant=bool(batch_invariant), + safe_gate=bool(safe_gate), + gate_lower_bound=float(gate_lower_bound) if gate_lower_bound is not None else None, + a_log=a_log, + dt_bias=dt_bias, + checkpoint_every_n_tokens=int(checkpoint_every_n_tokens), + plan_name=plan_name, ) + if checkpoint_every_n_tokens > 0: + return o, final_state, state_checkpoints + return o, final_state diff --git a/python/cudnn/linear_attention/ops/kda.py b/python/cudnn/linear_attention/ops/kda.py index 7801da3f3..7362071f2 100644 --- a/python/cudnn/linear_attention/ops/kda.py +++ b/python/cudnn/linear_attention/ops/kda.py @@ -20,124 +20,298 @@ per-key-channel log decay ``[total_tokens, heads, dim]``; ``beta`` is scalar ``[total_tokens, heads]``. -The op is a thin adapter over the graph API (the SDPA-op pattern): forward +The op is a thin adapter over the graph API: forward and backward execute cached single-node ``KDA`` / ``KDA_BWD`` pygraphs. -Engine selection happens at graph planning time over the registered python -engines: ``KdaFrostEngine`` (forward, default on SM100/SM103) with -``KdaCuTileEngine`` as the fallback — and the backward engine everywhere -(the FROST KDA backward kernel is a stub). Registered through +Engine selection happens at graph planning time over the manifest's python +engines: ``KdaFrostEngine`` (default on SM100/SM103, forward and backward) +with ``KdaCuTileEngine`` as the fallback. Registered through ``torch.library.custom_op`` so it composes with autograd, ``torch.compile``, and DDP. -The backward graph recomputes the forward's cheap intermediates (cumulative -gate, intra-chunk WY factors) — the ``KDA_BWD`` node contract keeps them off -the autograd wire. +Graph caching ensures cuDNN graphs are built once per unique configuration +and reused across calls. """ -from __future__ import annotations - import math from typing import Dict, Optional, Tuple import torch - import cudnn -from cudnn.linear_attention import engine_utils -_OP_NAMESPACE = "cudnn" -_OP_NAME = "kimi_delta_attention" +# --------------------------------------------------------------------------- +# Module-level state +# --------------------------------------------------------------------------- + -_TORCH_TO_CUDNN_DTYPE = { +_TORCH_DTYPE_TO_CUDNN = { torch.float16: cudnn.data_type.HALF, torch.bfloat16: cudnn.data_type.BFLOAT16, torch.float32: cudnn.data_type.FLOAT, + torch.int32: cudnn.data_type.INT32, + torch.int64: cudnn.data_type.INT64, } # one graph per static configuration (shapes, dtypes, scale, flags, device) -_fwd_graph_cache: Dict[tuple, tuple] = {} -_bwd_graph_cache: Dict[tuple, tuple] = {} -_ws_cache: Dict[int, "torch.Tensor"] = {} +_fprop_cache: Dict[tuple, tuple] = {} +_bprop_cache: Dict[tuple, tuple] = {} +_cudnn_handles: Dict[int, int] = {} + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def select_plan(graph, plan_name): + """Pin one execution plan by name on a freshly built graph (create the + plans, select by name, check support); ``None`` keeps default routing.""" + if plan_name is None: + return + graph.create_execution_plans() + names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] + matches = [i for i, n in enumerate(names) if n == plan_name or n.startswith(plan_name + "[")] + if not matches: + raise cudnn.cudnnGraphNotSupportedError(f"no {plan_name} plan for this graph (offered: {names})") + graph.select_plan(matches[0]) + graph.check_support() def _graph_workspace(graph, device): - """Caller-side workspace for a compiled graph (the explicit-workspace - convention: query the plan's size, allocate, pass to execute).""" + """Caller-side workspace for a compiled graph.""" if not graph._is_built: - # mirror execute()'s auto-build: plan via the router first (a bare - # build() would lower KDA to the backend, which has no lowering) + # plan first (a bare build() would lower KDA to the backend, + # which has no lowering) if not graph._planning_done: graph.create_execution_plans() - engine_utils.apply_pin(graph) if graph.selected_engine is None: graph.build() else: graph.build_plans() size = graph.get_workspace_size() - ws = _ws_cache.get(id(graph)) + ws = getattr(graph, "_la_ops_workspace", None) if ws is None or ws.numel() < size or ws.device != device: ws = torch.empty(max(size, 1), dtype=torch.uint8, device=device) - _ws_cache[id(graph)] = ws + graph._la_ops_workspace = ws return ws -_handle_cache: Dict[int, int] = {} - - -def _graph_handle(device): - """Per-device cuDNN handle carrying the caller's current stream - (classic ``set_stream`` semantics).""" +def _get_handle(device): + """Per-device cuDNN handle carrying the caller's current stream.""" idx = device.index if device.index is not None else torch.cuda.current_device() - handle = _handle_cache.get(idx) + handle = _cudnn_handles.get(idx) if handle is None: with torch.cuda.device(idx): handle = cudnn.create_handle() - _handle_cache[idx] = handle + _cudnn_handles[idx] = handle cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) return handle -def _cudnn_dtype(dtype: Optional[torch.dtype]): - return _TORCH_TO_CUDNN_DTYPE[dtype] if dtype is not None else None - - -# --------------------------------------------------------------------------- -# Forward -# --------------------------------------------------------------------------- +def _torch_dtype_to_cudnn(dtype: torch.dtype): + """Map a PyTorch dtype to a cuDNN data_type enum.""" + return _TORCH_DTYPE_TO_CUDNN[dtype] def _check_dtype(name, t, want) -> None: if t.dtype != want: - raise TypeError(f"{_OP_NAME}: {name} must be {want} (kernel-native; callers convert), got {t.dtype}") + raise TypeError(f"kimi_delta_attention: {name} must be {want} (kernel-native; callers convert), got {t.dtype}") + + +def _make_fprop_cache_key( + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + k_shape, + v_shape, + cu_dtype, + scale, + output_final_state, + use_qk_l2norm, + batch_invariant, + use_beta_sigmoid, + safe_gate, + gate_lower_bound, + has_initial_state, + ckpt, + device, + plan_name, +): + return ( + "fprop", + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + k_shape, + v_shape, + cu_dtype, + float(scale), + bool(output_final_state), + bool(use_qk_l2norm), + bool(batch_invariant), + bool(use_beta_sigmoid), + bool(safe_gate), + float(gate_lower_bound) if gate_lower_bound is not None else None, + bool(has_initial_state), + ckpt, + device, + plan_name, + ) + + +def _make_bprop_cache_key( + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + do_dtype, + k_shape, + v_shape, + cu_dtype, + g_dtype, + beta_dtype, + state_dtype, + dstate_in_dtype, + ckpt_rows, + scale, + use_qk_l2norm, + batch_invariant, + device, + plan_name, +): + return ( + "bprop", + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + k_dtype, + v_dtype, + do_dtype, + k_shape, + v_shape, + cu_dtype, + g_dtype, + beta_dtype, + state_dtype, + dstate_in_dtype, + ckpt_rows, + float(scale), + bool(use_qk_l2norm), + bool(batch_invariant), + device, + plan_name, + ) -def _build_fwd_graph(total, N, H, HV, K, V, io_dtype, g_dtype, beta_dtype, state_dtype, scale, output_final_state, use_qk_l2norm): +# --------------------------------------------------------------------------- +# Forward graph builder +# --------------------------------------------------------------------------- + + +def _build_fprop_graph( + total, + N, + H, + HK, + HV, + K, + V, + io_dtype, + g_dtype, + beta_dtype, + state_dtype, + cu_dtype, + scale, + output_final_state, + use_qk_l2norm, + batch_invariant, + use_beta_sigmoid, + safe_gate, + gate_lower_bound, + ckpt, +): graph = cudnn.pygraph() + HO = max(H, HV) q_t = graph.tensor([total, H, K], data_type=io_dtype, name="q") - k_t = graph.tensor([total, H, K], data_type=io_dtype, name="k") + k_t = graph.tensor([total, HK, K], data_type=io_dtype, name="k") v_t = graph.tensor([total, HV, V], data_type=io_dtype, name="v") - g_t = graph.tensor([total, HV, K], data_type=g_dtype, name="g") - beta_t = graph.tensor([total, HV], data_type=beta_dtype, name="beta") - cu_t = graph.tensor([N + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") - s0_t = None + g_t = graph.tensor([total, HO, K], data_type=g_dtype, name="g") + beta_t = graph.tensor([total, HO], data_type=beta_dtype, name="beta") + cu_t = graph.tensor([N + 1], data_type=cu_dtype, name="cu_seqlens") + state0_t = None if state_dtype is not None: - s0_t = graph.tensor([N, HV, K, V], data_type=state_dtype, name="initial_state") - O_t, fs_t, _h_t = graph.kda( + state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") + a_log_t = None + dt_bias_t = None + if safe_gate: + a_log_t = graph.tensor([HO], data_type=cudnn.data_type.FLOAT, name="a_log") + dt_bias_t = graph.tensor([HO, K], data_type=cudnn.data_type.FLOAT, name="dt_bias") + O_t, fs_t, state_checkpoints_t = graph.kda( q=q_t, k=k_t, v=v_t, g=g_t, beta=beta_t, cu_seqlens=cu_t, - initial_state=s0_t, + initial_state=state0_t, + a_log=a_log_t, + dt_bias=dt_bias_t, scale=scale, output_final_state=output_final_state, use_qk_l2norm=use_qk_l2norm, + batch_invariant=batch_invariant, + use_beta_sigmoid=use_beta_sigmoid, + safe_gate=safe_gate, + gate_lower_bound=gate_lower_bound, + checkpoint_every_n_tokens=ckpt, name="kda", ) - return graph, dict(q=q_t, k=k_t, v=v_t, g=g_t, beta=beta_t, cu=cu_t, s0=s0_t, O=O_t, fs=fs_t) + return graph, dict( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + cu=cu_t, + state0=state0_t, + a_log=a_log_t, + dt_bias=dt_bias_t, + O=O_t, + fs=fs_t, + state_checkpoints=state_checkpoints_t, + ) + + +# --------------------------------------------------------------------------- +# Forward custom op +# --------------------------------------------------------------------------- -@torch.library.custom_op(f"{_OP_NAMESPACE}::{_OP_NAME}_fwd", mutates_args=()) +@torch.library.custom_op("cudnn::kimi_delta_attention_fwd", mutates_args=()) def _kda_fwd( q: torch.Tensor, k: torch.Tensor, @@ -149,112 +323,213 @@ def _kda_fwd( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor]: - """KDA forward via a cached single-node KDA pygraph (THD layout). - - Returns ``(o, final_state)``; ``final_state`` is a zero-size tensor when - ``output_final_state`` is ``False``. + batch_invariant: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + safe_gate: bool = False, + gate_lower_bound: Optional[float] = None, + a_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + checkpoint_every_n_tokens: int = 0, + plan_name: Optional[str] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """KDA forward (internal): a cached single-node KDA pygraph, THD layout. + + Returns ``(o, final_state, state_checkpoints)``; ``final_state`` / ``state_checkpoints`` are zero-size + tensors when ``output_final_state`` is ``False`` / + ``checkpoint_every_n_tokens`` is ``0``. """ total, H, K = q.shape - if k.shape[1] != H: - raise ValueError(f"k must carry the same head count as q ({H}), got {k.shape[1]}") + HK = k.shape[1] HV, V = v.shape[1], v.shape[2] + if HK not in (H, HV): + raise ValueError(f"k head count ({HK}) must match q's ({H}) or v's ({HV}); canonical GQA shares grouped k/v heads") + HO = max(H, HV) N = cu_seqlens.shape[0] - 1 device = q.device - cu = cu_seqlens.to(torch.int32).contiguous() + if cu_seqlens.dtype not in (torch.int32, torch.int64): + raise ValueError(f"kimi_delta_attention: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") + cu = cu_seqlens _check_dtype("g", g, torch.float32) - _check_dtype("beta", beta, torch.float32) + if use_beta_sigmoid_in_kernel: + _check_dtype("beta", beta, q.dtype) + else: + _check_dtype("beta", beta, torch.float32) + if safe_gate: + if a_log is None or dt_bias is None: + raise ValueError("kimi_delta_attention: safe_gate requires a_log and dt_bias") + _check_dtype("a_log", a_log, torch.float32) + _check_dtype("dt_bias", dt_bias, torch.float32) + elif a_log is not None or dt_bias is not None: + raise ValueError("kimi_delta_attention: a_log/dt_bias require safe_gate=True") if initial_state is not None: _check_dtype("initial_state", initial_state, torch.float32) if initial_state.shape[0] != N: raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") - g32 = g.contiguous() - beta32 = beta.contiguous() - s0 = initial_state.contiguous() if initial_state is not None else None - - key = ( + for _name, _t in ( + ("k", k), + ("v", v), + ("g", g), + ("beta", beta), + ("cu_seqlens", cu_seqlens), + ("initial_state", initial_state), + ("a_log", a_log), + ("dt_bias", dt_bias), + ): + if _t is not None and _t.device != device: + raise ValueError(f"kimi_delta_attention: {_name} must be on q's device ({device}); got {_t.device}") + state0 = initial_state if initial_state is not None else None + ckpt = int(checkpoint_every_n_tokens) + + cache_key = _make_fprop_cache_key( total, N, H, + HK, HV, K, V, q.dtype, - bool(s0 is not None), - float(scale), - bool(output_final_state), - bool(use_qk_l2norm_in_kernel), + k.dtype, + v.dtype, + tuple(k.shape), + tuple(v.shape), + cu_seqlens.dtype, + scale, + output_final_state, + use_qk_l2norm_in_kernel, + batch_invariant, + use_beta_sigmoid_in_kernel, + safe_gate, + gate_lower_bound, + state0 is not None, + ckpt, device, + plan_name, ) - if key not in _fwd_graph_cache: - _fwd_graph_cache[key] = _build_fwd_graph( + if cache_key not in _fprop_cache: + _fprop_cache[cache_key] = _build_fprop_graph( total, N, H, + HK, HV, K, V, - _cudnn_dtype(q.dtype), - cudnn.data_type.FLOAT, + _torch_dtype_to_cudnn(q.dtype), cudnn.data_type.FLOAT, - cudnn.data_type.FLOAT if s0 is not None else None, + _torch_dtype_to_cudnn(beta.dtype), + cudnn.data_type.FLOAT if state0 is not None else None, + _torch_dtype_to_cudnn(cu_seqlens.dtype), float(scale), bool(output_final_state), bool(use_qk_l2norm_in_kernel), + bool(batch_invariant), + bool(use_beta_sigmoid_in_kernel), + bool(safe_gate), + float(gate_lower_bound) if gate_lower_bound is not None else None, + ckpt, ) - graph, t = _fwd_graph_cache[key] + select_plan(_fprop_cache[cache_key][0], plan_name) + + graph, t = _fprop_cache[cache_key] - o = torch.empty(total, HV, V, dtype=q.dtype, device=device) + o = torch.empty(total, HO, V, dtype=q.dtype, device=device) variant_pack = { - t["q"]: q.contiguous(), - t["k"]: k.contiguous(), - t["v"]: v.contiguous(), - t["g"]: g32, - t["beta"]: beta32, + t["q"]: q, + t["k"]: k, + t["v"]: v, + t["g"]: g, + t["beta"]: beta, t["cu"]: cu, t["O"]: o, } - if s0 is not None: - variant_pack[t["s0"]] = s0 + if state0 is not None: + variant_pack[t["state0"]] = state0 + if safe_gate: + variant_pack[t["a_log"]] = a_log + 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, HV, K, V, dtype=torch.float32, device=device) + final_state = torch.empty(N, HO, K, V, dtype=torch.float32, device=device) variant_pack[t["fs"]] = final_state - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_graph_handle(device)) - return o, 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) + 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 @_kda_fwd.register_fake -def _kda_fwd_fake(q, k, v, g, beta, cu_seqlens, scale, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False): - total, _H, K = q.shape +def _kda_fwd_fake( + q, + k, + v, + g, + beta, + cu_seqlens, + scale, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + batch_invariant=False, + use_beta_sigmoid_in_kernel=False, + safe_gate=False, + gate_lower_bound=None, + a_log=None, + dt_bias=None, + checkpoint_every_n_tokens=0, + plan_name: Optional[str] = None, +): + total, H, K = q.shape + HK = k.shape[1] HV, V = v.shape[1], v.shape[2] + if HK not in (H, HV): + raise ValueError(f"k head count ({HK}) must match q's ({H}) or v's ({HV}); canonical GQA shares grouped k/v heads") + HO = max(H, HV) N = cu_seqlens.shape[0] - 1 - o = q.new_empty(total, HV, V) - final = q.new_empty((N, HV, K, V) if output_final_state else (0,), dtype=torch.float32) - return o, final + if cu_seqlens.dtype not in (torch.int32, torch.int64): + raise ValueError(f"kimi_delta_attention: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") + 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) + 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) + else: + state_checkpoints = q.new_empty(0) + return o, final, state_checkpoints # --------------------------------------------------------------------------- -# Backward +# Backward graph builder # --------------------------------------------------------------------------- -def _build_bwd_graph(total, N, H, HV, K, V, io_dtype, g_dtype, beta_dtype, state_dtype, dht_dtype, scale, use_qk_l2norm): +def _build_bprop_graph( + total, N, H, HK, HV, K, V, io_dtype, g_dtype, beta_dtype, state_dtype, dstate_in_dtype, cu_dtype, ckpt_rows, scale, use_qk_l2norm, batch_invariant +): graph = cudnn.pygraph() + HO = max(H, HV) q_t = graph.tensor([total, H, K], data_type=io_dtype, name="q") - k_t = graph.tensor([total, H, K], data_type=io_dtype, name="k") + k_t = graph.tensor([total, HK, K], data_type=io_dtype, name="k") v_t = graph.tensor([total, HV, V], data_type=io_dtype, name="v") - g_t = graph.tensor([total, HV, K], data_type=g_dtype, name="g") - beta_t = graph.tensor([total, HV], data_type=beta_dtype, name="beta") - cu_t = graph.tensor([N + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") - dO_t = graph.tensor([total, HV, V], data_type=io_dtype, name="dO") - s0_t = None + g_t = graph.tensor([total, HO, K], data_type=g_dtype, name="g") + beta_t = graph.tensor([total, HO], data_type=beta_dtype, name="beta") + cu_t = graph.tensor([N + 1], data_type=cu_dtype, name="cu_seqlens") + dO_t = graph.tensor([total, HO, V], data_type=io_dtype, name="dO") + state0_t = None if state_dtype is not None: - s0_t = graph.tensor([N, HV, K, V], data_type=state_dtype, name="initial_state") + state0_t = graph.tensor([N, HO, K, V], data_type=state_dtype, name="initial_state") dfs_t = None - if dht_dtype is not None: - dfs_t = graph.tensor([N, HV, K, V], data_type=dht_dtype, name="d_final_state") - dQ_t, dK_t, dV_t, dG_t, dBeta_t, dS0_t = graph.kda_bwd( + if dstate_in_dtype is not None: + dfs_t = graph.tensor([N, HO, K, V], 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") + dQ_t, dK_t, dV_t, dG_t, dBeta_t, dstate0_t = graph.kda_bwd( q=q_t, k=k_t, v=v_t, @@ -262,10 +537,12 @@ def _build_bwd_graph(total, N, H, HV, K, V, io_dtype, g_dtype, beta_dtype, state beta=beta_t, cu_seqlens=cu_t, dO=dO_t, - initial_state=s0_t, + state_checkpoints=ckpts_t, + initial_state=state0_t, d_final_state=dfs_t, scale=scale, use_qk_l2norm=use_qk_l2norm, + batch_invariant=batch_invariant, name="kda_bwd", ) return graph, dict( @@ -276,18 +553,24 @@ def _build_bwd_graph(total, N, H, HV, K, V, io_dtype, g_dtype, beta_dtype, state beta=beta_t, cu=cu_t, dO=dO_t, - s0=s0_t, + state0=state0_t, dfs=dfs_t, dQ=dQ_t, dK=dK_t, dV=dV_t, dG=dG_t, dBeta=dBeta_t, - dS0=dS0_t, + dstate0=dstate0_t, + ckpts=ckpts_t, ) -@torch.library.custom_op(f"{_OP_NAMESPACE}::{_OP_NAME}_bwd", mutates_args=()) +# --------------------------------------------------------------------------- +# Backward custom op +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("cudnn::kimi_delta_attention_bwd", mutates_args=()) def _kda_bwd( dO: torch.Tensor, q: torch.Tensor, @@ -299,20 +582,36 @@ def _kda_bwd( scale: float, initial_state: Optional[torch.Tensor] = None, d_final_state: Optional[torch.Tensor] = None, + state_checkpoints: Optional[torch.Tensor] = None, use_qk_l2norm_in_kernel: bool = False, + batch_invariant: bool = False, + plan_name: Optional[str] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """KDA backward via a cached single-node KDA_BWD pygraph (THD layout). + """KDA backward (internal): a cached single-node KDA_BWD pygraph, THD layout. - Returns ``(dq, dk, dv, dg, dbeta, d_initial_state)``; ``d_initial_state`` - is a zero-size tensor when ``initial_state`` is ``None``. + ``state_checkpoints`` is the forward's per-chunk state series (io dtype, + chunk cadence); when given, the engine consumes it instead of running + the checkpoint recompute pass. Returns ``(dq, dk, dv, dg, dbeta, + d_initial_state)``; ``d_initial_state`` is a zero-size tensor when + ``initial_state`` is ``None``. """ total, H, K = q.shape - if k.shape[1] != H: - raise ValueError(f"k must carry the same head count as q ({H}), got {k.shape[1]}") + # autograd materializes reduction grads as broadcast (stride-0) + # views; densify ONLY those (dense callers pass through untouched) + if 0 in dO.stride(): + dO = dO.contiguous() + if d_final_state is not None and 0 in d_final_state.stride(): + d_final_state = d_final_state.contiguous() + HK = k.shape[1] HV, V = v.shape[1], v.shape[2] + if HK not in (H, HV): + raise ValueError(f"k head count ({HK}) must match q's ({H}) or v's ({HV}); canonical GQA shares grouped k/v heads") + HO = max(H, HV) N = cu_seqlens.shape[0] - 1 device = q.device - cu = cu_seqlens.to(torch.int32).contiguous() + if cu_seqlens.dtype not in (torch.int32, torch.int64): + raise ValueError(f"kimi_delta_attention: cu_seqlens must be int32 or int64; got {cu_seqlens.dtype}") + cu = cu_seqlens _check_dtype("g", g, torch.float32) _check_dtype("beta", beta, torch.float32) if initial_state is not None: @@ -321,84 +620,132 @@ def _kda_bwd( raise ValueError(f"initial_state must carry one state per sequence: got {initial_state.shape[0]} for {N} sequences") if d_final_state is not None: _check_dtype("d_final_state", d_final_state, torch.float32) - s0 = initial_state.contiguous() if initial_state is not None else None - dht = d_final_state.contiguous() if d_final_state is not None else None - - key = ( + if state_checkpoints is not None: + _check_dtype("state_checkpoints", state_checkpoints, q.dtype) + for _name, _t in ( + ("k", k), + ("v", v), + ("g", g), + ("beta", beta), + ("cu_seqlens", cu_seqlens), + ("dO", dO), + ("d_final_state", d_final_state), + ("state_checkpoints", state_checkpoints), + ): + if _t is not None and _t.device != device: + raise ValueError(f"kimi_delta_attention: {_name} must be on q's device ({device}); got {_t.device}") + state0 = initial_state if initial_state is not None else None + dstate_in = d_final_state if d_final_state is not None else None + + cache_key = _make_bprop_cache_key( total, N, H, + HK, HV, K, V, q.dtype, + k.dtype, + v.dtype, + dO.dtype, + tuple(k.shape), + tuple(v.shape), + cu_seqlens.dtype, g.dtype, beta.dtype, - s0.dtype if s0 is not None else None, - dht.dtype if dht is not None else None, - float(scale), - bool(use_qk_l2norm_in_kernel), + state0.dtype if state0 is not None else None, + dstate_in.dtype if dstate_in is not None else None, + state_checkpoints.shape[0] if state_checkpoints is not None else None, + scale, + use_qk_l2norm_in_kernel, + batch_invariant, device, + plan_name, ) - if key not in _bwd_graph_cache: - _bwd_graph_cache[key] = _build_bwd_graph( + if cache_key not in _bprop_cache: + _bprop_cache[cache_key] = _build_bprop_graph( total, N, H, + HK, HV, K, V, - _cudnn_dtype(q.dtype), - _cudnn_dtype(g.dtype), - _cudnn_dtype(beta.dtype), - _cudnn_dtype(s0.dtype) if s0 is not None else None, - _cudnn_dtype(dht.dtype) if dht is not None else None, + _torch_dtype_to_cudnn(q.dtype), + _torch_dtype_to_cudnn(g.dtype), + _torch_dtype_to_cudnn(beta.dtype), + _torch_dtype_to_cudnn(state0.dtype) if state0 is not None else None, + _torch_dtype_to_cudnn(dstate_in.dtype) if dstate_in is not None else None, + _torch_dtype_to_cudnn(cu_seqlens.dtype), + state_checkpoints.shape[0] if state_checkpoints is not None else None, float(scale), bool(use_qk_l2norm_in_kernel), + bool(batch_invariant), ) - graph, t = _bwd_graph_cache[key] + select_plan(_bprop_cache[cache_key][0], plan_name) + + graph, t = _bprop_cache[cache_key] dq = torch.empty(total, H, K, dtype=q.dtype, device=device) - dk = torch.empty(total, H, K, dtype=k.dtype, device=device) - dv = torch.empty(total, HV, V, dtype=v.dtype, device=device) - dg = torch.empty(total, HV, K, dtype=g.dtype, device=device) - dbeta = torch.empty(total, HV, dtype=beta.dtype, device=device) + dk = torch.empty(total, HK, K, dtype=q.dtype, device=device) + dv = torch.empty(total, HV, V, dtype=q.dtype, device=device) + dg = torch.empty(total, HO, K, dtype=g.dtype, device=device) + dbeta = torch.empty(total, HO, dtype=beta.dtype, device=device) variant_pack = { - t["q"]: q.contiguous(), - t["k"]: k.contiguous(), - t["v"]: v.contiguous(), - t["g"]: g.contiguous(), - t["beta"]: beta.contiguous(), + t["q"]: q, + t["k"]: k, + t["v"]: v, + t["g"]: g, + t["beta"]: beta, t["cu"]: cu, - t["dO"]: dO.contiguous(), + t["dO"]: dO, t["dQ"]: dq, t["dK"]: dk, t["dV"]: dv, t["dG"]: dg, t["dBeta"]: dbeta, } - dh032 = None - if s0 is not None: - variant_pack[t["s0"]] = s0 - dh032 = torch.empty_like(s0) - variant_pack[t["dS0"]] = dh032 - if dht is not None: - variant_pack[t["dfs"]] = dht - graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_graph_handle(device)) - dh0 = dh032 if dh032 is not None else torch.empty(0, dtype=torch.float32, device=device) - return dq, dk, dv, dg, dbeta, dh0 + dstate0 = None + if state0 is not None: + variant_pack[t["state0"]] = state0 + dstate0 = torch.empty_like(state0) + variant_pack[t["dstate0"]] = dstate0 + if dstate_in is not None: + variant_pack[t["dfs"]] = dstate_in + if state_checkpoints is not None: + variant_pack[t["ckpts"]] = state_checkpoints + graph.execute(variant_pack, workspace=_graph_workspace(graph, device), handle=_get_handle(device)) + if dstate0 is None: + dstate0 = torch.empty(0, dtype=torch.float32, device=device) + return dq, dk, dv, dg, dbeta, dstate0 @_kda_bwd.register_fake -def _kda_bwd_fake(dO, q, k, v, g, beta, cu_seqlens, scale, initial_state=None, d_final_state=None, use_qk_l2norm_in_kernel=False): - dh0 = torch.empty_like(initial_state) if initial_state is not None else q.new_empty(0, dtype=torch.float32) +def _kda_bwd_fake( + dO, + q, + k, + v, + g, + beta, + cu_seqlens, + scale, + initial_state=None, + d_final_state=None, + state_checkpoints=None, + use_qk_l2norm_in_kernel=False, + batch_invariant=False, + plan_name=None, +): + dstate0 = torch.empty_like(initial_state) if initial_state is not None else q.new_empty(0, dtype=torch.float32) return ( torch.empty_like(q), torch.empty_like(k), torch.empty_like(v), torch.empty_like(g), torch.empty_like(beta), - dh0, + dstate0, ) @@ -408,21 +755,58 @@ def _kda_bwd_fake(dO, q, k, v, g, beta, cu_seqlens, scale, initial_state=None, d def _kda_setup_context(ctx, inputs, output): - q, k, v, g, beta, cu_seqlens, scale, initial_state, output_final_state, use_qk_l2norm_in_kernel = inputs + ( + q, + k, + v, + g, + beta, + cu_seqlens, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + batch_invariant, + use_beta_sigmoid_in_kernel, + safe_gate, + gate_lower_bound, + a_log, + dt_bias, + checkpoint_every_n_tokens, + plan_name, + ) = inputs # save_for_backward cannot hold None; keep initial_state as an attribute. - ctx.save_for_backward(q, k, v, g, beta, cu_seqlens) + saved = [q, k, v, g, beta, cu_seqlens] + ctx.ckpt_reuse = checkpoint_every_n_tokens == 16 and output[2].numel() > 0 + if ctx.ckpt_reuse: + saved.append(output[2]) + ctx.save_for_backward(*saved) ctx.initial_state = initial_state ctx.scale = scale ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - - -def _kda_backward(ctx, dO, dFinal): - q, k, v, g, beta, cu_seqlens = ctx.saved_tensors + ctx.batch_invariant = batch_invariant + ctx.plan_name = plan_name + ctx.use_beta_sigmoid_in_kernel = bool(use_beta_sigmoid_in_kernel) + ctx.safe_gate = bool(safe_gate) + ctx.set_materialize_grads(False) + ctx.mark_non_differentiable(output[2]) + + +def _kda_backward(ctx, dO, dFinal, _dstate_checkpoints): + if ctx.use_beta_sigmoid_in_kernel or ctx.safe_gate: + raise NotImplementedError("kimi_delta_attention: safe_gate/use_beta_sigmoid_in_kernel are forward-only (KDA_BWD takes post-activation gates)") + if ctx.ckpt_reuse: + q, k, v, g, beta, cu_seqlens, state_checkpoints = ctx.saved_tensors + else: + q, k, v, g, beta, cu_seqlens = ctx.saved_tensors + state_checkpoints = None initial_state = ctx.initial_state - dht = dFinal if (dFinal is not None and dFinal.numel() > 0) else None - dq, dk, dv, dg, dbeta, dh0 = torch.ops.cudnn.kimi_delta_attention_bwd( - dO.contiguous(), + if dO is None: + dO = torch.zeros(q.shape[0], max(q.shape[1], v.shape[1]), v.shape[2], dtype=q.dtype, device=q.device) + dstate_in = dFinal if (dFinal is not None and dFinal.numel() > 0) else None + dq, dk, dv, dg, dbeta, dstate0 = torch.ops.cudnn.kimi_delta_attention_bwd( + dO, q, k, v, @@ -431,11 +815,16 @@ def _kda_backward(ctx, dO, dFinal): cu_seqlens, ctx.scale, initial_state=initial_state, - d_final_state=dht, + d_final_state=dstate_in, + state_checkpoints=state_checkpoints, use_qk_l2norm_in_kernel=ctx.use_qk_l2norm_in_kernel, + batch_invariant=ctx.batch_invariant, + plan_name=ctx.plan_name, ) # q, k, v, g, beta, cu_seqlens, scale, initial_state, output_final_state, - # use_qk_l2norm_in_kernel + # use_qk_l2norm_in_kernel, batch_invariant, use_beta_sigmoid_in_kernel, + # safe_gate, gate_lower_bound, a_log, dt_bias, checkpoint_every_n_tokens, + # plan_name return ( dq, dk, @@ -444,14 +833,22 @@ def _kda_backward(ctx, dO, dFinal): dbeta, None, None, - dh0 if initial_state is not None else None, + dstate0 if initial_state is not None else None, + None, + None, + None, + None, + None, + None, + None, + None, None, None, ) torch.library.register_autograd( - f"{_OP_NAMESPACE}::{_OP_NAME}_fwd", + "cudnn::kimi_delta_attention_fwd", _kda_backward, setup_context=_kda_setup_context, ) @@ -473,15 +870,25 @@ def kimi_delta_attention( initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, + batch_invariant: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + safe_gate: bool = False, + gate_lower_bound: Optional[float] = None, + a_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + checkpoint_every_n_tokens: int = 0, + plan_name: Optional[str] = None, ): """Kimi Delta Attention (KDA) linear attention. THD layout (matches the graph-API KDA node): - q, k: ``[total_tokens, H, K]``; v: ``[total_tokens, HV, V]`` - g: ``[total_tokens, HV, K]`` (per-key-channel log decay); - beta: ``[total_tokens, HV]`` (scalar); cu_seqlens: ``[N+1]`` int32 - initial_state / final_state: ``[N, HV, K, V]`` + q: ``[total_tokens, H, K]``; k: ``[total_tokens, HK, K]`` (HK = H, or + 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]`` + (``HO = max(H, HV)``: the gates, output, and state heads) A dense batch of N equal-length sequences is expressed as ``cu_seqlens = [0, T, 2T, ...]`` over the flattened tokens. @@ -491,26 +898,48 @@ def kimi_delta_attention( ``d_initial_state`` are returned in float32. Args: - g: per-key-channel log-space decay (``alpha = exp(g) in (0, 1]^K``). - beta: per-token scalar write strength. + g: per-key-channel log-space decay (``alpha = exp(g) in (0, 1]^K``), + or raw pre-activation logits when ``safe_gate=True``. + beta: per-token scalar write strength (float32 post-sigmoid), or + io-dtype logits when ``use_beta_sigmoid_in_kernel=True``. cu_seqlens: ``[N+1]`` int32 sequence boundaries over the packed tokens. scale: attention scale applied to ``q``. Defaults to ``1 / sqrt(K)``. initial_state: optional recurrent state (otherwise zero). output_final_state: if ``True``, also return the per-sequence state after the last token. - use_qk_l2norm_in_kernel: if ``True``, L2-normalize the q/k rows inside - the kernel (the KDA model's feature map); if ``False``, pass q/k + use_qk_l2norm_in_kernel: if ``True``, L2-normalize the Q/K rows inside + the kernel (the KDA model's feature map); if ``False``, pass Q/K as given (the caller owns their conditioning). - + batch_invariant: if ``True``, each sequence's results are bitwise + independent of the batch composition (whole-sequence scheduling; + disables split-K load balancing). + use_beta_sigmoid_in_kernel: apply ``sigmoid(beta)`` inside the kernel. + Forward-only. + safe_gate: interpret ``g`` through the safe-gate transform + ``gate_lower_bound * sigmoid(exp(a_log) * (g + dt_bias))``. + Requires ``a_log`` and ``dt_bias``. Forward-only. + gate_lower_bound: safe-gate lower bound in log space (default -5.0). + 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, + 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. + + plan_name: optionally pin one execution plan by name (the plan + API's ``get_plan_name_at_index`` names, e.g. ``kda_frost``); a + graph offering no such plan raises ``cudnnGraphNotSupportedError``. Returns: - ``(o, final_state)`` with ``o`` shaped like ``v``. ``final_state`` is - empty unless ``output_final_state=True``. + ``(o, final_state)`` with ``o`` shaped like ``v``, or + ``(o, final_state, state_checkpoints)`` when ``checkpoint_every_n_tokens > 0``. + ``final_state`` is empty unless ``output_final_state=True``. """ if q.dim() != 3: raise ValueError("expected THD [total_tokens, heads, dim] tensors") if scale is None: scale = 1.0 / math.sqrt(q.shape[-1]) - return torch.ops.cudnn.kimi_delta_attention_fwd( + o, final_state, state_checkpoints = torch.ops.cudnn.kimi_delta_attention_fwd( q, k, v, @@ -521,4 +950,15 @@ def kimi_delta_attention( initial_state=initial_state, output_final_state=bool(output_final_state), use_qk_l2norm_in_kernel=bool(use_qk_l2norm_in_kernel), + batch_invariant=bool(batch_invariant), + use_beta_sigmoid_in_kernel=bool(use_beta_sigmoid_in_kernel), + safe_gate=bool(safe_gate), + gate_lower_bound=float(gate_lower_bound) if gate_lower_bound is not None else None, + a_log=a_log, + dt_bias=dt_bias, + checkpoint_every_n_tokens=int(checkpoint_every_n_tokens), + plan_name=plan_name, ) + if checkpoint_every_n_tokens > 0: + return o, final_state, state_checkpoints + return o, final_state diff --git a/python/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp index 8b9b0b351..3eeef0662 100644 --- a/python/pygraph/variant_pack.cpp +++ b/python/pygraph/variant_pack.cpp @@ -556,6 +556,23 @@ class VariantPackNative { return true; } + bool + all_dense_layout(std::string &offender) const { + for (size_t i = 0; i < operands_.size(); i++) { + const Operand &operand = operands_[i]; + if (!operand.filled || operand.stride.empty()) continue; + for (int d = operand.ndim - 1; d >= 0; d--) { + if (operand.shape[d] == 1) continue; + if (operand.stride[d] != 1) { + offender = std::to_string(i); + return false; + } + break; + } + } + return true; + } + bool operand_contiguous(size_t index) const { const Operand &operand = operands_.at(index); @@ -840,9 +857,15 @@ its parts. .def("operands", &VariantPackNative::operands) .def_property_readonly("address", &VariantPackNative::pointer_array) .def("__len__", &VariantPackNative::size) - .def("all_contiguous", [](const VariantPackNative &self) { + .def("all_contiguous", + [](const VariantPackNative &self) { + std::string offender; + bool ok = self.all_contiguous(offender); + return py::make_tuple(ok, offender); + }) + .def("all_dense_layout", [](const VariantPackNative &self) { std::string offender; - bool ok = self.all_contiguous(offender); + bool ok = self.all_dense_layout(offender); return py::make_tuple(ok, offender); }); } diff --git a/test/python/linear_attention/common.py b/test/python/linear_attention/common.py deleted file mode 100644 index a00cd27e8..000000000 --- a/test/python/linear_attention/common.py +++ /dev/null @@ -1,208 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Shared helpers for the GDN (Gated DeltaNet) test suite. - -Availability guards, run dispatcher, and comparison tolerances shared by the -fprop and bprop test files. -""" - -from __future__ import annotations - -import os -from typing import Optional - -import pytest -import torch - -_HAS_CUDA = torch.cuda.is_available() - -try: - import cuda.tile # noqa: F401 — the engines' own availability gate - - from cudnn.linear_attention.ops import gated_delta_net, kimi_delta_attention - - _HAS_CUTILE = True -except ImportError: - _HAS_CUTILE = False - -GDN_MARKS = [ - pytest.mark.L0, - pytest.mark.skipif(not _HAS_CUDA, reason="needs CUDA"), - pytest.mark.skipif(not _HAS_CUTILE, reason="needs the cuda.tile runtime"), -] - -KDA_MARKS = GDN_MARKS # same gate: both ops are lazy exports of the same package - -FWD_TOL = {torch.bfloat16: 2e-2, torch.float16: 1e-2} -STATE_TOL = {torch.bfloat16: 2e-2, torch.float16: 1e-2} -BWD_TOL = {torch.bfloat16: 4e-2, torch.float16: 2e-2} - -# (H, HV) pairs: H = q/k heads, HV = v/g/beta heads (GVA when HV > H). -HEAD_CONFIGS = [(1, 1), (3, 3), (1, 2), (2, 4), (16, 32), (16, 64)] -HEAD_CONFIGS_SMALL = [(1, 1), (2, 4)] - - -def assert_engine_declines(graph, engine_name: str) -> None: - """``engine_name`` must not serve ``graph``. - - Asserted against the ranked plan list rather than through a failing - ``build()``: a decline only advances the walk, so a sibling engine - (cuTile vs FROST) may still serve the graph — and does wherever both are - installed, which is why the build-fails form passes on a box missing one - of them and fails in CI.""" - try: - graph.create_execution_plans() - except Exception: - return # nothing claimed it at all, which is a stronger decline - names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] - assert not any(n.startswith(engine_name) for n in names), f"{engine_name} claimed a graph it must decline; plans={names}" - - -def run_gdn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - *, - scale: Optional[float] = None, - initial_state: Optional[torch.Tensor] = None, - output_final_state: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, -): - """Run GDN through the public custom op (the op is THD-only: dense - ``[B, T, ...]`` inputs are flattened with ``cu_seqlens = [0, T, 2T, ...]``; - packed ``[1, total, ...]`` inputs are squeezed to the op's rank-3 THD - layout). Returns ``(o, final_state)`` with ``final_state`` normalized to - ``None`` when not requested.""" - g = g.float() # the op requires kernel-native fp32 gates; callers convert - beta = beta.float() - if cu_seqlens is None: - B, T = q.shape[0], q.shape[1] - cu = torch.arange(0, B + 1, dtype=torch.int32, device=q.device) * T - o, fs = gated_delta_net( - q.reshape(B * T, *q.shape[2:]), - k.reshape(B * T, *k.shape[2:]), - v.reshape(B * T, *v.shape[2:]), - g.reshape(B * T, *g.shape[2:]), - beta.reshape(B * T, *beta.shape[2:]), - cu, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - ) - o = o.reshape(B, T, *o.shape[1:]) - else: - o, fs = gated_delta_net( - q.squeeze(0), - k.squeeze(0), - v.squeeze(0), - g.squeeze(0), - beta.squeeze(0), - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - ) - o = o.unsqueeze(0) - if fs is None or fs.numel() == 0: - fs = None - return o, fs - - -def run_kda( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - *, - scale: Optional[float] = None, - initial_state: Optional[torch.Tensor] = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: Optional[torch.Tensor] = None, -): - """Run KDA through the public custom op (the op is THD-only: dense - ``[B, T, ...]`` inputs are flattened with ``cu_seqlens = [0, T, 2T, ...]``). - ``g`` is the per-key-channel log decay ([..., HV, K]); ``beta`` is scalar - ([..., HV]). Returns ``(o, final_state)`` with ``final_state`` normalized - to ``None`` when not requested.""" - g = g.float() # the op requires kernel-native fp32 gates; callers convert - beta = beta.float() - if cu_seqlens is None: - B, T = q.shape[0], q.shape[1] - cu = torch.arange(0, B + 1, dtype=torch.int32, device=q.device) * T - o, fs = kimi_delta_attention( - q.reshape(B * T, *q.shape[2:]), - k.reshape(B * T, *k.shape[2:]), - v.reshape(B * T, *v.shape[2:]), - g.reshape(B * T, *g.shape[2:]), - beta.reshape(B * T, *beta.shape[2:]), - cu, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - ) - o = o.reshape(B, T, *o.shape[1:]) - else: - o, fs = kimi_delta_attention( - q.squeeze(0), - k.squeeze(0), - v.squeeze(0), - g.squeeze(0), - beta.squeeze(0), - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - cu_seqlens=cu_seqlens, - ) - o = o.unsqueeze(0) - if fs is None or fs.numel() == 0: - fs = None - return o, fs - - -DETERMINISM_REPEATS = int(os.environ.get("DETERMINISM_REPEATS", "8")) - - -def bitwise_bits(t: torch.Tensor) -> torch.Tensor: - return t.contiguous().view(torch.uint8) - - -def assert_bitwise_runs(launch, repeats=DETERMINISM_REPEATS, label=""): - """``launch()`` returns a tuple of freshly-written output tensors. Launch - ``repeats`` times back to back (single sync at the end) and require every - run to match run 0 bit for bit (barrier/fence races are timing-dependent; - any mismatching bit is a failure, there is no tolerance).""" - runs = [launch() for _ in range(repeats)] - torch.cuda.synchronize() - for out in runs[0]: - assert torch.isfinite(out.float()).all(), f"{label}: non-finite output in run 0" - for r, outs in enumerate(runs[1:], start=1): - for i, (a, b) in enumerate(zip(runs[0], outs)): - assert torch.equal(bitwise_bits(a), bitwise_bits(b)), f"{label}: output {i} differs between run 0 and run {r}" - - -def assert_concurrent_stream_runs(launch_a, launch_b, s1, s2, repeats=DETERMINISM_REPEATS): - """Two concurrent kernel instances on separate streams must not perturb - each other: every repeat must match its own single-stream baseline.""" - with torch.cuda.stream(s1): - base_a = launch_a() - torch.cuda.synchronize() - with torch.cuda.stream(s2): - base_b = launch_b() - torch.cuda.synchronize() - for r in range(repeats): - with torch.cuda.stream(s1): - out_a = launch_a() - with torch.cuda.stream(s2): - out_b = launch_b() - torch.cuda.synchronize() - for label, base, outs in (("A", base_a, out_a), ("B", base_b, out_b)): - for i, (x, y) in enumerate(zip(base, outs)): - assert torch.equal(bitwise_bits(x), bitwise_bits(y)), f"stream {label} output {i} differs on concurrent run {r}" diff --git a/test/python/linear_attention/conftest.py b/test/python/linear_attention/conftest.py index 00cccd912..147770ea1 100644 --- a/test/python/linear_attention/conftest.py +++ b/test/python/linear_attention/conftest.py @@ -2,16 +2,32 @@ # SPDX-License-Identifier: Apache-2.0 """ -Fixtures for the GDN test suite: seeded input factories for q/k/v and the -gate tensors (g = log decay, beta = write strength). +Fixtures for the linear-attention test suite: seeded input factories for +Q/K/V and the gate tensors (G = log decay, Beta = write strength). + +Also overlays the source ``python/cudnn`` dir onto the installed ``cudnn`` +package's ``__path__``: the built wheel may lack the engine subtrees the +suite pins (``cudnn.linear_attention.frost``). Unnecessary once the engines +ship in the built frontend package. """ from __future__ import annotations +from pathlib import Path + import pytest import torch import torch.nn.functional as F +try: + import cudnn + + _SRC_CUDNN = Path(__file__).resolve().parents[3] / "python" / "cudnn" + if _SRC_CUDNN.is_dir() and str(_SRC_CUDNN) not in cudnn.__path__: + cudnn.__path__.append(str(_SRC_CUDNN)) +except ImportError: + pass # test_la.py skips via importorskip + def multidist_randu(num_dists, dim, *, device, mean_std=0.05, lower=-0.25, upper=0.25): """Rows drawn from per-row uniform distributions with normally-distributed @@ -21,9 +37,9 @@ def multidist_randu(num_dists, dim, *, device, mean_std=0.05, lower=-0.25, upper def gen_qkv(B, T, H, HV, K, V, dtype, device="cuda"): - """Dense [B, T, heads, dim] q/k/v. k is l2-normalized along the feature - dim so the delta-rule update (I - beta k k^T) stays contractive for - beta in (0, 1].""" + """Dense [B, T, heads, dim] Q/K/V. K is l2-normalized along the feature + dim so the delta-rule update (I - Beta K K^T) stays contractive for + Beta in (0, 1].""" q = multidist_randu(B * T * H, K, device=device).reshape(B, T, H, K) k = multidist_randu(B * T * H, K, device=device).reshape(B, T, H, K) k = F.normalize(k, p=2.0, dim=-1) @@ -36,8 +52,8 @@ def gen_qkv(B, T, H, HV, K, V, dtype, device="cuda"): def gen_gates(B, T, HV, dtype, device="cuda", alpha=True, beta=True): - """Gate tensors [B, T, HV]. alpha off -> g = 0 (decay factor exactly 1); - beta off -> beta = 1 (plain delta rule).""" + """Gate tensors [B, T, HV]. alpha off -> G = 0 (decay factor exactly 1); + beta off -> Beta = 1 (plain delta rule).""" if alpha: a = torch.empty(B, T, HV, device=device, dtype=torch.float32).uniform_(0.1, 1.0) g = a.log() @@ -51,11 +67,11 @@ def gen_gates(B, T, HV, dtype, device="cuda", alpha=True, beta=True): def gen_kda_gates(B, T, HV, K, dtype, device="cuda", alpha=True, beta=True, lo=0.9): - """KDA gate tensors: g [B, T, HV, K] fp32 per-key-channel log decay, beta + """KDA gate tensors: G [B, T, HV, K] fp32 per-key-channel log decay, Beta [B, T, HV] scalar write strength. ``g`` stays fp32 (the per-channel decay kernel keeps it fp32); the per-channel cumulative product over a 64-token - chunk stays well within bf16 range for ``alpha >= lo``. alpha off -> g = 0 - (no decay); beta off -> beta = 1 (plain delta rule).""" + chunk stays well within bf16 range for ``alpha >= lo``. alpha off -> G = 0 + (no decay); beta off -> Beta = 1 (plain delta rule).""" if alpha: a = torch.empty(B, T, HV, K, device=device, dtype=torch.float32).uniform_(lo, 1.0) g = a.log() @@ -69,10 +85,10 @@ def gen_kda_gates(B, T, HV, K, dtype, device="cuda", alpha=True, beta=True, lo=0 def gen_gdn2_gates(B, T, HO, K, V, dtype, device="cuda", alpha=True, beta=True, w=True, lo=0.5): - """GDN-2 gate tensors: g [B, T, HO, K] fp32 per-key-channel log decay, - beta [B, T, HO, K] per-key erase gate, w [B, T, HO, V] per-value write - gate. beta/w are io-dtype (rounded before both kernel and reference see - them). alpha off -> g = 0; beta/w off -> ones.""" + """GDN-2 gate tensors: G [B, T, HO, K] fp32 per-key-channel log decay, + Beta [B, T, HO, K] per-key erase gate, W [B, T, HO, V] per-value write + gate. Beta/W are io-dtype (rounded before both kernel and reference see + them). alpha off -> G = 0; beta/w off -> ones.""" if alpha: a = torch.empty(B, T, HO, K, device=device, dtype=torch.float32).uniform_(lo, 1.0) g = a.log() diff --git a/test/python/linear_attention/cutile/__init__.py b/test/python/linear_attention/cutile/__init__.py deleted file mode 100644 index 52a7a9daf..000000000 --- a/test/python/linear_attention/cutile/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 diff --git a/test/python/linear_attention/cutile/conftest.py b/test/python/linear_attention/cutile/conftest.py deleted file mode 100644 index 1caa0845b..000000000 --- a/test/python/linear_attention/cutile/conftest.py +++ /dev/null @@ -1,32 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""The cutile suite validates the cuTile engines and their kernels, so op -calls made here pin them — under default ranking the FROST engines would serve -the FROST-eligible shapes instead. - -The pin is enforced where it is applied: engine_utils.apply_pin() raises if the -pinned engine produced no plan, so a pin that stops working fails the first op -call rather than silently testing another engine. This suite once ran for -months against whichever engine the router picked, because the seam it pinned -through was dead and nothing checked.""" - -from __future__ import annotations - -import pytest - - -@pytest.fixture(autouse=True, scope="package") -def _pin_cutile_engines(): - from cudnn.linear_attention import engine_utils - from cudnn.linear_attention.ops import gdn, kda - - saved = engine_utils.pin_engines(("gdn_cutile", "kda_cutile")) - for m in (gdn, kda): - m._fwd_graph_cache.clear() - m._bwd_graph_cache.clear() - yield - engine_utils.pin_engines(saved) - for m in (gdn, kda): - m._fwd_graph_cache.clear() - m._bwd_graph_cache.clear() diff --git a/test/python/linear_attention/cutile/test_gdn_bprop.py b/test/python/linear_attention/cutile/test_gdn_bprop.py deleted file mode 100644 index 76e455af3..000000000 --- a/test/python/linear_attention/cutile/test_gdn_bprop.py +++ /dev/null @@ -1,272 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Backward (bprop) tests for the GDN cuTile kernel: gradients of q, k, v, g, -beta (and the initial state) are compared against autograd through the fp64 -recurrent reference in ``reference_gdn`` using a shared random upstream -gradient. - -Covers: fp16/bf16, MHA and GVA head configs, full-chunk and partial-chunk -sequence lengths, ragged varlen batches (cu_seqlens), zero-length sequences, -explicit/auto scale, alpha (decay gate) and beta (write strength) on/off, -chunked prefill with state carry-over, initial-state (dh0) and final-state -(dht) gradient paths, and head-dim variants including K != V and K = 256. -""" - -from __future__ import annotations - -import math -import random - -import pytest -import torch - -from ..common import ( - BWD_TOL, - FWD_TOL, - GDN_MARKS, - HEAD_CONFIGS, - HEAD_CONFIGS_SMALL, - run_gdn, -) -from ..conftest import gen_gates, gen_qkv -from ..reference_gdn import gdn_reference, rms_ratio - -pytestmark = GDN_MARKS - -_SEED = 42 - - -def _seed_all(seed=_SEED): - random.seed(seed) - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - - -def _run_bprop_case( - dtype, - H, - HV, - B, - T, - K=128, - V=128, - scale=None, - alpha=True, - beta=True, - initial_state=False, - state_grad=False, - cu_seqlens=None, - total_T=None, -): - """Run kernel fwd+bwd and reference fwd+bwd with a shared upstream - gradient; assert forward parity and per-input gradient parity.""" - _seed_all() - Teff = total_T or T - q0, k0, v0 = gen_qkv(B, Teff, H, HV, K, V, dtype) - g0, b0 = gen_gates(B, Teff, HV, dtype, alpha=alpha, beta=beta) - - S0_data = None - if initial_state: - N = B if cu_seqlens is None else cu_seqlens.numel() - 1 - S0_data = torch.randn(N, HV, K, V, device="cuda", dtype=torch.float32) * 0.05 - - w = torch.randn(B, Teff, HV, V, device="cuda", dtype=torch.float32) - if state_grad: - N = B if cu_seqlens is None else cu_seqlens.numel() - 1 - wf = torch.randn(N, HV, K, V, device="cuda", dtype=torch.float32) - - # --- kernel --- - leaves = { - "q": q0.clone().requires_grad_(), - "k": k0.clone().requires_grad_(), - "v": v0.clone().requires_grad_(), - "g": g0.clone().requires_grad_(), - "beta": b0.clone().requires_grad_(), - } - S0 = S0_data.clone().requires_grad_() if initial_state else None - o, fs = run_gdn( - leaves["q"], - leaves["k"], - leaves["v"], - leaves["g"], - leaves["beta"], - scale=scale, - initial_state=S0, - output_final_state=state_grad, - cu_seqlens=cu_seqlens, - ) - loss = (o.float() * w).sum() - if state_grad: - loss = loss + (fs.float() * wf).sum() - loss.backward() - torch.cuda.synchronize() - - # --- fp64 reference --- - ref_leaves = {n: t.detach().double().requires_grad_() for n, t in leaves.items()} - S0_ref = S0_data.detach().double().requires_grad_() if initial_state else None - o_ref, fs_ref = gdn_reference( - ref_leaves["q"], - ref_leaves["k"], - ref_leaves["v"], - ref_leaves["g"], - ref_leaves["beta"], - scale=scale, - initial_state=S0_ref, - cu_seqlens=cu_seqlens, - ) - loss_ref = (o_ref * w.double()).sum() - if state_grad: - loss_ref = loss_ref + (fs_ref * wf.double()).sum() - loss_ref.backward() - - r_o = rms_ratio(o, o_ref) - assert r_o < FWD_TOL[dtype], f"forward o rms ratio {r_o:.4g} >= {FWD_TOL[dtype]}" - - pairs = [(f"d{n}", leaves[n].grad, ref_leaves[n].grad) for n in leaves] - if initial_state: - pairs.append(("dh0", S0.grad, S0_ref.grad)) - for name, got, ref in pairs: - assert got is not None, f"no gradient for {name}" - assert torch.isfinite(got).all(), f"non-finite gradient for {name}" - r = rms_ratio(got, ref) - assert r < BWD_TOL[dtype], f"{name} rms ratio {r:.4g} >= {BWD_TOL[dtype]}" - - -@pytest.mark.parametrize("alpha,beta", [(True, True), (True, False), (False, True)]) -@pytest.mark.parametrize("scale", [1.0, "auto"]) -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS) -@pytest.mark.parametrize("B,T", [(1, 64), (1, 128), (1, 256), (2, 256)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_basic(dtype, B, T, H, HV, scale, alpha, beta): - scale = 1.0 / math.sqrt(128) if scale == "auto" else scale - _run_bprop_case(dtype, H, HV, B, T, scale=scale, alpha=alpha, beta=beta) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (2, 4), (16, 64)]) -@pytest.mark.parametrize("T", [31, 61, 91, 121, 251]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_nonfull(dtype, T, H, HV): - """Backward through sequence lengths that are not chunk multiples.""" - _run_bprop_case(dtype, H, HV, 1, T) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (2, 4), (16, 64)]) -@pytest.mark.parametrize( - "seq_lens", - [[256, 256], [511, 501], [64, 128, 512], [31, 63, 93, 123, 150, 500], [255, 257]], -) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_varlen_ragged(dtype, seq_lens, H, HV): - """Backward through ragged varlen batches (cu_seqlens path).""" - bounds = [0] - for sl in seq_lens: - bounds.append(bounds[-1] + sl) - cu = torch.tensor(bounds, dtype=torch.int32, device="cuda") - _run_bprop_case(dtype, H, HV, 1, None, cu_seqlens=cu, total_T=bounds[-1]) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (16, 64)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_zero_length_sequence(dtype, H, HV, T=64): - """A zero-length sequence in the varlen batch must not perturb the - gradients of the others.""" - _seed_all() - q, k, v = gen_qkv(1, T, H, HV, 128, 128, dtype) - g, b = gen_gates(1, T, HV, dtype) - w = torch.randn(1, T, HV, 128, device="cuda", dtype=torch.float32) - - def run(cu): - leaves = { - "q": q.clone().requires_grad_(), - "k": k.clone().requires_grad_(), - "v": v.clone().requires_grad_(), - "g": g.clone().requires_grad_(), - "beta": b.clone().requires_grad_(), - } - o, _ = run_gdn(leaves["q"], leaves["k"], leaves["v"], leaves["g"], leaves["beta"], cu_seqlens=cu) - (o.float() * w).sum().backward() - torch.cuda.synchronize() - return {n: t.grad for n, t in leaves.items()} - - ref = run(torch.tensor([0, T], dtype=torch.int32, device="cuda")) - got = run(torch.tensor([0, T, T], dtype=torch.int32, device="cuda")) - for name in ref: - torch.testing.assert_close(got[name], ref[name], atol=1e-3, rtol=1e-3, msg=f"d{name} perturbed by the zero-length sequence") - - -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL) -@pytest.mark.parametrize("T1,T2", [(128, 128), (64, 192), (192, 121)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_chunked_prefill(dtype, T1, T2, H, HV, B=2, K=128, V=128): - """Backward through a two-phase prefill: the state carried from part 1 - into part 2 chains the gradient (dht -> dh0) across the boundary; the - leaf gradients must match a single-shot fp64 reference.""" - _seed_all() - T = T1 + T2 - q0, k0, v0 = gen_qkv(B, T, H, HV, K, V, dtype) - g0, b0 = gen_gates(B, T, HV, dtype) - w = torch.randn(B, T, HV, V, device="cuda", dtype=torch.float32) - - leaves = { - "q": q0.clone().requires_grad_(), - "k": k0.clone().requires_grad_(), - "v": v0.clone().requires_grad_(), - "g": g0.clone().requires_grad_(), - "beta": b0.clone().requires_grad_(), - } - - def part(t0, t1, S0, output_final_state): - return run_gdn( - leaves["q"][:, t0:t1], - leaves["k"][:, t0:t1], - leaves["v"][:, t0:t1], - leaves["g"][:, t0:t1], - leaves["beta"][:, t0:t1], - initial_state=S0, - output_final_state=output_final_state, - ) - - o1, fs1 = part(0, T1, None, True) - o2, _ = part(T1, T, fs1, False) - o = torch.cat([o1, o2], dim=1) - (o.float() * w).sum().backward() - torch.cuda.synchronize() - - ref_leaves = {n: t.detach().double().requires_grad_() for n, t in leaves.items()} - o_ref, _ = gdn_reference(ref_leaves["q"], ref_leaves["k"], ref_leaves["v"], ref_leaves["g"], ref_leaves["beta"]) - (o_ref * w.double()).sum().backward() - - # The carried state round-trips through the op's output dtype, so allow - # slightly more than the single-shot backward tolerance. - tol = 1.5 * BWD_TOL[dtype] - for name in leaves: - r = rms_ratio(leaves[name].grad, ref_leaves[name].grad) - assert r < tol, f"chunked-prefill d{name} rms ratio {r:.4g} >= {tol}" - - -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL) -@pytest.mark.parametrize("T", [128, 251]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_initial_state(dtype, T, H, HV): - """Gradient flow through a random (non-zero) initial recurrent state - (dh0 path, no final-state gradient).""" - _run_bprop_case(dtype, H, HV, 1, T, initial_state=True) - - -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL) -@pytest.mark.parametrize("T", [128, 192, 251]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_state_grads(dtype, T, H, HV): - """Initial-state gradient (dh0) and final-state gradient (dht) paths: - the loss includes the final state and the initial state requires grad.""" - _run_bprop_case(dtype, H, HV, 1, T, initial_state=True, state_grad=True) - - -@pytest.mark.parametrize("K,V", [(64, 64), (64, 128), (128, 128), (256, 128)]) -@pytest.mark.parametrize("T", [128, 251]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_bprop_head_dims(dtype, T, K, V, H=2, HV=2): - """Backward through head-dim variants: K != V and the K = 256 bound.""" - _run_bprop_case(dtype, H, HV, 1, T, K=K, V=V) diff --git a/test/python/linear_attention/cutile/test_gdn_fprop.py b/test/python/linear_attention/cutile/test_gdn_fprop.py deleted file mode 100644 index fad0ef7c5..000000000 --- a/test/python/linear_attention/cutile/test_gdn_fprop.py +++ /dev/null @@ -1,236 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Forward (fprop) tests for the GDN cuTile kernel, validated against the fp64 -recurrent reference in ``reference_gdn``. - -Covers: fp16/bf16, MHA and GVA head configs, full-chunk and partial-chunk -sequence lengths, ragged varlen batches (cu_seqlens), zero-length sequences, -explicit/auto scale, alpha (decay gate) and beta (write strength) on/off, -chunked prefill with state carry-over, initial state, and head-dim variants -including K != V and K = 256. -""" - -from __future__ import annotations - -import math -import random - -import pytest -import torch - -from ..common import ( - FWD_TOL, - GDN_MARKS, - HEAD_CONFIGS, - HEAD_CONFIGS_SMALL, - STATE_TOL, - run_gdn, -) -from ..conftest import gen_gates, gen_qkv -from ..reference_gdn import gdn_reference, rms_ratio - -pytestmark = GDN_MARKS - -_SEED = 42 - - -def _seed_all(seed=_SEED): - random.seed(seed) - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - - -def _check_fwd(o, o_ref, dtype, what="o"): - assert torch.isfinite(o).all(), f"non-finite values in {what}" - r = rms_ratio(o, o_ref) - assert r < FWD_TOL[dtype], f"{what} rms ratio {r:.4g} >= {FWD_TOL[dtype]}" - - -def _check_state(fs, fs_ref, dtype): - assert torch.isfinite(fs).all(), "non-finite values in final_state" - r = rms_ratio(fs, fs_ref) - assert r < STATE_TOL[dtype], f"final_state rms ratio {r:.4g} >= {STATE_TOL[dtype]}" - - -def _run_fprop_case( - dtype, - H, - HV, - B, - T, - K=128, - V=128, - scale=None, - alpha=True, - beta=True, - initial_state=False, - cu_seqlens=None, - total_T=None, -): - """Build inputs, run the kernel, and compare o + final_state against the - fp64 recurrent reference. ``cu_seqlens`` implies a packed batch with - B == 1 and T == total_T.""" - _seed_all() - q, k, v = gen_qkv(B, total_T or T, H, HV, K, V, dtype) - g, b = gen_gates(B, total_T or T, HV, dtype, alpha=alpha, beta=beta) - - S0 = None - if initial_state: - N = B if cu_seqlens is None else cu_seqlens.numel() - 1 - S0 = torch.randn(N, HV, K, V, device="cuda", dtype=torch.float32) * 0.05 - - o, fs = run_gdn(q, k, v, g, b, scale=scale, initial_state=S0, output_final_state=True, cu_seqlens=cu_seqlens) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = gdn_reference(q, k, v, g, b, scale=scale, initial_state=S0, cu_seqlens=cu_seqlens) - - _check_fwd(o, o_ref, dtype) - assert fs is not None - _check_state(fs, fs_ref, dtype) - - -@pytest.mark.parametrize("beta", [False, True]) -@pytest.mark.parametrize("alpha", [False, True]) -@pytest.mark.parametrize("scale", [1.0, "auto"]) -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS) -@pytest.mark.parametrize("B,T", [(1, 64), (1, 128), (1, 256), (2, 256)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_basic(dtype, B, T, H, HV, scale, alpha, beta): - if not alpha and not beta: - pytest.skip("output amplitude grows unbounded along the token dimension") - scale = 1.0 / math.sqrt(128) if scale == "auto" else scale - _run_fprop_case(dtype, H, HV, B, T, scale=scale, alpha=alpha, beta=beta) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (2, 4), (16, 64)]) -@pytest.mark.parametrize("T", [31, 61, 91, 121, 251]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_nonfull(dtype, T, H, HV): - """Sequence lengths that are not a multiple of the 64-token chunk.""" - _run_fprop_case(dtype, H, HV, 1, T) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (2, 4), (16, 64)]) -@pytest.mark.parametrize( - "seq_lens", - [[256, 256], [511, 501], [64, 128, 512], [31, 63, 93, 123, 150, 500]], -) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_varlen_ragged(dtype, seq_lens, H, HV): - """Ragged multi-sequence batches through the packed cu_seqlens path.""" - bounds = [0] - for sl in seq_lens: - bounds.append(bounds[-1] + sl) - cu = torch.tensor(bounds, dtype=torch.int32, device="cuda") - _run_fprop_case(dtype, H, HV, 1, None, cu_seqlens=cu, total_T=bounds[-1]) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (16, 64)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_zero_length_sequence(dtype, H, HV, T=64): - """A zero-length sequence in the varlen batch must not perturb the others.""" - _seed_all() - q, k, v = gen_qkv(1, T, H, HV, 128, 128, dtype) - g, b = gen_gates(1, T, HV, dtype) - cu = torch.tensor([0, T], dtype=torch.int32, device="cuda") - cu_with_empty = torch.tensor([0, T, T], dtype=torch.int32, device="cuda") - - o_ref, fs_ref = run_gdn(q, k, v, g, b, output_final_state=True, cu_seqlens=cu) - o, fs = run_gdn(q, k, v, g, b, output_final_state=True, cu_seqlens=cu_with_empty) - torch.cuda.synchronize() - - torch.testing.assert_close(o, o_ref, atol=1e-3, rtol=1e-3) - torch.testing.assert_close(fs[0], fs_ref[0], atol=1e-3, rtol=1e-3) - assert (fs[1] == 0).all(), "state of a zero-length sequence must stay zero" - - -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL) -@pytest.mark.parametrize("T1,T2", [(128, 128), (64, 192), (192, 121)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_chunked_prefill(dtype, T1, T2, H, HV, B=2, K=128, V=128): - """Two-phase prefill: the final state of part 1 feeds part 2 as its - initial state; the concatenated output must match a single-shot run.""" - _seed_all() - T = T1 + T2 - q, k, v = gen_qkv(B, T, H, HV, K, V, dtype) - g, b = gen_gates(B, T, HV, dtype) - - def part(t0, t1, S0): - return run_gdn( - q[:, t0:t1].contiguous(), - k[:, t0:t1].contiguous(), - v[:, t0:t1].contiguous(), - g[:, t0:t1].contiguous(), - b[:, t0:t1].contiguous(), - initial_state=S0, - output_final_state=True, - ) - - o1, fs1 = part(0, T1, None) - o2, fs2 = part(T1, T, fs1) - o = torch.cat([o1, o2], dim=1) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = gdn_reference(q, k, v, g, b) - - # The state round-trips through the op's output dtype between the two - # calls, so allow slightly more than the single-shot forward tolerance. - tol = 1.5 * FWD_TOL[dtype] - r_o = rms_ratio(o, o_ref) - r_s = rms_ratio(fs2, fs_ref) - assert r_o < tol, f"chunked-prefill o rms ratio {r_o:.4g} >= {tol}" - assert r_s < tol, f"chunked-prefill final_state rms ratio {r_s:.4g} >= {tol}" - - -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL) -@pytest.mark.parametrize("T", [128, 251]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_initial_state(dtype, T, H, HV): - """Random (non-zero) initial recurrent state.""" - _run_fprop_case(dtype, H, HV, 1, T, initial_state=True) - - -@pytest.mark.parametrize("K,V", [(64, 64), (64, 128), (128, 128), (256, 128)]) -@pytest.mark.parametrize("T", [128, 251]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_fprop_head_dims(dtype, T, K, V, H=2, HV=2): - """Head-dim variants: K != V and the K = 256 upper bound.""" - _run_fprop_case(dtype, H, HV, 1, T, K=K, V=V) - - -# --------------------------------------------------------------------------- -# Argument validation -# --------------------------------------------------------------------------- - - -def _dummy_inputs(H=2, HV=2, T=64, K=128, V=128): - q, k, v = gen_qkv(1, T, H, HV, K, V, torch.bfloat16) - g, b = gen_gates(1, T, HV, torch.bfloat16) - return q[0], k[0], v[0], g[0], b[0] - - -def _cu(*bounds): - return torch.tensor(bounds, dtype=torch.int32, device="cuda") - - -def test_invalid_qk_head_mismatch(): - q, k, v, g, b = _dummy_inputs() - with pytest.raises(Exception, match="head|No valid engine"): - run_gdn(q.unsqueeze(0), k[:, :1].unsqueeze(0), v.unsqueeze(0), g.unsqueeze(0), b.unsqueeze(0), cu_seqlens=_cu(0, 64)) - - -def test_invalid_gva_head_ratio(): - q, k, v, g, b = _dummy_inputs(H=2, HV=2) - with pytest.raises(Exception, match="divisible|multiple|head|No valid engine"): - run_gdn(q.unsqueeze(0), k.unsqueeze(0), v.repeat(1, 3, 1)[:, :3].unsqueeze(0), g.unsqueeze(0), b.unsqueeze(0), cu_seqlens=_cu(0, 64)) - - -def test_invalid_initial_state_count(): - q, k, v, g, b = _dummy_inputs(T=128) - S0 = torch.zeros(1, 2, 128, 128, device="cuda", dtype=torch.float32) - with pytest.raises(Exception, match="initial"): - run_gdn(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), g.unsqueeze(0), b.unsqueeze(0), cu_seqlens=_cu(0, 64, 128), initial_state=S0) diff --git a/test/python/linear_attention/cutile/test_kda_bprop.py b/test/python/linear_attention/cutile/test_kda_bprop.py deleted file mode 100644 index 658c032c7..000000000 --- a/test/python/linear_attention/cutile/test_kda_bprop.py +++ /dev/null @@ -1,273 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Backward (bprop) tests for the KDA cuTile kernel: gradients of q, k, v, g, -beta (and the initial state) are compared against autograd through the fp64 -recurrent reference in ``reference_kda`` using a shared random upstream -gradient. - -Covers: fp16/bf16, MHA and GVA head configs, full-chunk and partial-chunk -sequence lengths, ragged varlen batches (cu_seqlens), zero-length sequences, -explicit/auto scale, alpha (per-channel decay gate) and beta (write strength) -on/off, chunked prefill with state carry-over, initial-state (dh0) and -final-state (dht) gradient paths, and head-dim variants including K != V and -K = 256. -""" - -from __future__ import annotations - -import math -import random - -import pytest -import torch - -from ..common import ( - BWD_TOL, - FWD_TOL, - HEAD_CONFIGS, - HEAD_CONFIGS_SMALL, - KDA_MARKS, - run_kda, -) -from ..conftest import gen_kda_gates, gen_qkv -from ..reference_kda import kda_reference, rms_ratio - -pytestmark = KDA_MARKS - -_SEED = 42 - - -def _seed_all(seed=_SEED): - random.seed(seed) - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - - -def _run_bprop_case( - dtype, - H, - HV, - B, - T, - K=128, - V=128, - scale=None, - alpha=True, - beta=True, - initial_state=False, - state_grad=False, - cu_seqlens=None, - total_T=None, -): - """Run kernel fwd+bwd and reference fwd+bwd with a shared upstream - gradient; assert forward parity and per-input gradient parity.""" - _seed_all() - Teff = total_T or T - q0, k0, v0 = gen_qkv(B, Teff, H, HV, K, V, dtype) - g0, b0 = gen_kda_gates(B, Teff, HV, K, dtype, alpha=alpha, beta=beta) - - S0_data = None - if initial_state: - N = B if cu_seqlens is None else cu_seqlens.numel() - 1 - S0_data = torch.randn(N, HV, K, V, device="cuda", dtype=torch.float32) * 0.05 - - w = torch.randn(B, Teff, HV, V, device="cuda", dtype=torch.float32) - if state_grad: - N = B if cu_seqlens is None else cu_seqlens.numel() - 1 - wf = torch.randn(N, HV, K, V, device="cuda", dtype=torch.float32) - - # --- kernel --- - leaves = { - "q": q0.clone().requires_grad_(), - "k": k0.clone().requires_grad_(), - "v": v0.clone().requires_grad_(), - "g": g0.clone().requires_grad_(), - "beta": b0.clone().requires_grad_(), - } - S0 = S0_data.clone().requires_grad_() if initial_state else None - o, fs = run_kda( - leaves["q"], - leaves["k"], - leaves["v"], - leaves["g"], - leaves["beta"], - scale=scale, - initial_state=S0, - output_final_state=state_grad, - cu_seqlens=cu_seqlens, - ) - loss = (o.float() * w).sum() - if state_grad: - loss = loss + (fs.float() * wf).sum() - loss.backward() - torch.cuda.synchronize() - - # --- fp64 reference --- - ref_leaves = {n: t.detach().double().requires_grad_() for n, t in leaves.items()} - S0_ref = S0_data.detach().double().requires_grad_() if initial_state else None - o_ref, fs_ref = kda_reference( - ref_leaves["q"], - ref_leaves["k"], - ref_leaves["v"], - ref_leaves["g"], - ref_leaves["beta"], - scale=scale, - initial_state=S0_ref, - cu_seqlens=cu_seqlens, - ) - loss_ref = (o_ref * w.double()).sum() - if state_grad: - loss_ref = loss_ref + (fs_ref * wf.double()).sum() - loss_ref.backward() - - r_o = rms_ratio(o, o_ref) - assert r_o < FWD_TOL[dtype], f"forward o rms ratio {r_o:.4g} >= {FWD_TOL[dtype]}" - - pairs = [(f"d{n}", leaves[n].grad, ref_leaves[n].grad) for n in leaves] - if initial_state: - pairs.append(("dh0", S0.grad, S0_ref.grad)) - for name, got, ref in pairs: - assert got is not None, f"no gradient for {name}" - assert torch.isfinite(got).all(), f"non-finite gradient for {name}" - r = rms_ratio(got, ref) - assert r < BWD_TOL[dtype], f"{name} rms ratio {r:.4g} >= {BWD_TOL[dtype]}" - - -@pytest.mark.parametrize("alpha,beta", [(True, True), (True, False), (False, True)]) -@pytest.mark.parametrize("scale", [1.0, "auto"]) -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS) -@pytest.mark.parametrize("B,T", [(1, 64), (1, 128), (1, 256), (2, 256)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_basic(dtype, B, T, H, HV, scale, alpha, beta): - scale = 1.0 / math.sqrt(128) if scale == "auto" else scale - _run_bprop_case(dtype, H, HV, B, T, scale=scale, alpha=alpha, beta=beta) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (2, 4), (16, 64)]) -@pytest.mark.parametrize("T", [31, 61, 91, 121, 251]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_nonfull(dtype, T, H, HV): - """Backward through sequence lengths that are not chunk multiples.""" - _run_bprop_case(dtype, H, HV, 1, T) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (2, 4), (16, 64)]) -@pytest.mark.parametrize( - "seq_lens", - [[256, 256], [511, 501], [64, 128, 512], [31, 63, 93, 123, 150, 500], [255, 257]], -) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_varlen_ragged(dtype, seq_lens, H, HV): - """Backward through ragged varlen batches (cu_seqlens path).""" - bounds = [0] - for sl in seq_lens: - bounds.append(bounds[-1] + sl) - cu = torch.tensor(bounds, dtype=torch.int32, device="cuda") - _run_bprop_case(dtype, H, HV, 1, None, cu_seqlens=cu, total_T=bounds[-1]) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (16, 64)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_zero_length_sequence(dtype, H, HV, T=64): - """A zero-length sequence in the varlen batch must not perturb the - gradients of the others.""" - _seed_all() - q, k, v = gen_qkv(1, T, H, HV, 128, 128, dtype) - g, b = gen_kda_gates(1, T, HV, 128, dtype) - w = torch.randn(1, T, HV, 128, device="cuda", dtype=torch.float32) - - def run(cu): - leaves = { - "q": q.clone().requires_grad_(), - "k": k.clone().requires_grad_(), - "v": v.clone().requires_grad_(), - "g": g.clone().requires_grad_(), - "beta": b.clone().requires_grad_(), - } - o, _ = run_kda(leaves["q"], leaves["k"], leaves["v"], leaves["g"], leaves["beta"], cu_seqlens=cu) - (o.float() * w).sum().backward() - torch.cuda.synchronize() - return {n: t.grad for n, t in leaves.items()} - - ref = run(torch.tensor([0, T], dtype=torch.int32, device="cuda")) - got = run(torch.tensor([0, T, T], dtype=torch.int32, device="cuda")) - for name in ref: - torch.testing.assert_close(got[name], ref[name], atol=1e-3, rtol=1e-3, msg=f"d{name} perturbed by the zero-length sequence") - - -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL) -@pytest.mark.parametrize("T1,T2", [(128, 128), (64, 192), (192, 121)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_chunked_prefill(dtype, T1, T2, H, HV, B=2, K=128, V=128): - """Backward through a two-phase prefill: the state carried from part 1 - into part 2 chains the gradient (dht -> dh0) across the boundary; the - leaf gradients must match a single-shot fp64 reference.""" - _seed_all() - T = T1 + T2 - q0, k0, v0 = gen_qkv(B, T, H, HV, K, V, dtype) - g0, b0 = gen_kda_gates(B, T, HV, K, dtype) - w = torch.randn(B, T, HV, V, device="cuda", dtype=torch.float32) - - leaves = { - "q": q0.clone().requires_grad_(), - "k": k0.clone().requires_grad_(), - "v": v0.clone().requires_grad_(), - "g": g0.clone().requires_grad_(), - "beta": b0.clone().requires_grad_(), - } - - def part(t0, t1, S0, output_final_state): - return run_kda( - leaves["q"][:, t0:t1], - leaves["k"][:, t0:t1], - leaves["v"][:, t0:t1], - leaves["g"][:, t0:t1], - leaves["beta"][:, t0:t1], - initial_state=S0, - output_final_state=output_final_state, - ) - - o1, fs1 = part(0, T1, None, True) - o2, _ = part(T1, T, fs1, False) - o = torch.cat([o1, o2], dim=1) - (o.float() * w).sum().backward() - torch.cuda.synchronize() - - ref_leaves = {n: t.detach().double().requires_grad_() for n, t in leaves.items()} - o_ref, _ = kda_reference(ref_leaves["q"], ref_leaves["k"], ref_leaves["v"], ref_leaves["g"], ref_leaves["beta"]) - (o_ref * w.double()).sum().backward() - - # The carried state round-trips through the op's output dtype, so allow - # slightly more than the single-shot backward tolerance. - tol = 1.5 * BWD_TOL[dtype] - for name in leaves: - r = rms_ratio(leaves[name].grad, ref_leaves[name].grad) - assert r < tol, f"chunked-prefill d{name} rms ratio {r:.4g} >= {tol}" - - -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL) -@pytest.mark.parametrize("T", [128, 251]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_initial_state(dtype, T, H, HV): - """Gradient flow through a random (non-zero) initial recurrent state - (dh0 path, no final-state gradient).""" - _run_bprop_case(dtype, H, HV, 1, T, initial_state=True) - - -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL) -@pytest.mark.parametrize("T", [128, 192, 251]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_bprop_state_grads(dtype, T, H, HV): - """Initial-state gradient (dh0) and final-state gradient (dht) paths: - the loss includes the final state and the initial state requires grad.""" - _run_bprop_case(dtype, H, HV, 1, T, initial_state=True, state_grad=True) - - -@pytest.mark.parametrize("K,V", [(64, 64), (64, 128), (128, 128), (256, 128)]) -@pytest.mark.parametrize("T", [128, 251]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_bprop_head_dims(dtype, T, K, V, H=2, HV=2): - """Backward through head-dim variants: K != V and the K = 256 bound.""" - _run_bprop_case(dtype, H, HV, 1, T, K=K, V=V) diff --git a/test/python/linear_attention/cutile/test_kda_fprop.py b/test/python/linear_attention/cutile/test_kda_fprop.py deleted file mode 100644 index e7ce2ed3d..000000000 --- a/test/python/linear_attention/cutile/test_kda_fprop.py +++ /dev/null @@ -1,236 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Forward (fprop) tests for the KDA cuTile kernel, validated against the fp64 -recurrent reference in ``reference_kda``. - -Covers: fp16/bf16, MHA and GVA head configs, full-chunk and partial-chunk -sequence lengths, ragged varlen batches (cu_seqlens), zero-length sequences, -explicit/auto scale, alpha (per-channel decay gate) and beta (write strength) -on/off, chunked prefill with state carry-over, initial state, and head-dim -variants including K != V and K = 256. -""" - -from __future__ import annotations - -import math -import random - -import pytest -import torch - -from ..common import ( - FWD_TOL, - HEAD_CONFIGS, - HEAD_CONFIGS_SMALL, - KDA_MARKS, - STATE_TOL, - run_kda, -) -from ..conftest import gen_kda_gates, gen_qkv -from ..reference_kda import kda_reference, rms_ratio - -pytestmark = KDA_MARKS - -_SEED = 42 - - -def _seed_all(seed=_SEED): - random.seed(seed) - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - - -def _check_fwd(o, o_ref, dtype, what="o"): - assert torch.isfinite(o).all(), f"non-finite values in {what}" - r = rms_ratio(o, o_ref) - assert r < FWD_TOL[dtype], f"{what} rms ratio {r:.4g} >= {FWD_TOL[dtype]}" - - -def _check_state(fs, fs_ref, dtype): - assert torch.isfinite(fs).all(), "non-finite values in final_state" - r = rms_ratio(fs, fs_ref) - assert r < STATE_TOL[dtype], f"final_state rms ratio {r:.4g} >= {STATE_TOL[dtype]}" - - -def _run_fprop_case( - dtype, - H, - HV, - B, - T, - K=128, - V=128, - scale=None, - alpha=True, - beta=True, - initial_state=False, - cu_seqlens=None, - total_T=None, -): - """Build inputs, run the kernel, and compare o + final_state against the - fp64 recurrent reference. ``cu_seqlens`` implies a packed batch with - B == 1 and T == total_T.""" - _seed_all() - q, k, v = gen_qkv(B, total_T or T, H, HV, K, V, dtype) - g, b = gen_kda_gates(B, total_T or T, HV, K, dtype, alpha=alpha, beta=beta) - - S0 = None - if initial_state: - N = B if cu_seqlens is None else cu_seqlens.numel() - 1 - S0 = torch.randn(N, HV, K, V, device="cuda", dtype=torch.float32) * 0.05 - - o, fs = run_kda(q, k, v, g, b, scale=scale, initial_state=S0, output_final_state=True, cu_seqlens=cu_seqlens) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = kda_reference(q, k, v, g, b, scale=scale, initial_state=S0, cu_seqlens=cu_seqlens) - - _check_fwd(o, o_ref, dtype) - assert fs is not None - _check_state(fs, fs_ref, dtype) - - -@pytest.mark.parametrize("beta", [False, True]) -@pytest.mark.parametrize("alpha", [False, True]) -@pytest.mark.parametrize("scale", [1.0, "auto"]) -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS) -@pytest.mark.parametrize("B,T", [(1, 64), (1, 128), (1, 256), (2, 256)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_basic(dtype, B, T, H, HV, scale, alpha, beta): - if not alpha and not beta: - pytest.skip("output amplitude grows unbounded along the token dimension") - scale = 1.0 / math.sqrt(128) if scale == "auto" else scale - _run_fprop_case(dtype, H, HV, B, T, scale=scale, alpha=alpha, beta=beta) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (2, 4), (16, 64)]) -@pytest.mark.parametrize("T", [31, 61, 91, 121, 251]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_nonfull(dtype, T, H, HV): - """Sequence lengths that are not a multiple of the 64-token chunk.""" - _run_fprop_case(dtype, H, HV, 1, T) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (2, 4), (16, 64)]) -@pytest.mark.parametrize( - "seq_lens", - [[256, 256], [511, 501], [64, 128, 512], [31, 63, 93, 123, 150, 500]], -) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_varlen_ragged(dtype, seq_lens, H, HV): - """Ragged multi-sequence batches through the packed cu_seqlens path.""" - bounds = [0] - for sl in seq_lens: - bounds.append(bounds[-1] + sl) - cu = torch.tensor(bounds, dtype=torch.int32, device="cuda") - _run_fprop_case(dtype, H, HV, 1, None, cu_seqlens=cu, total_T=bounds[-1]) - - -@pytest.mark.parametrize("H,HV", [(1, 1), (16, 64)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_zero_length_sequence(dtype, H, HV, T=64): - """A zero-length sequence in the varlen batch must not perturb the others.""" - _seed_all() - q, k, v = gen_qkv(1, T, H, HV, 128, 128, dtype) - g, b = gen_kda_gates(1, T, HV, 128, dtype) - cu = torch.tensor([0, T], dtype=torch.int32, device="cuda") - cu_with_empty = torch.tensor([0, T, T], dtype=torch.int32, device="cuda") - - o_ref, fs_ref = run_kda(q, k, v, g, b, output_final_state=True, cu_seqlens=cu) - o, fs = run_kda(q, k, v, g, b, output_final_state=True, cu_seqlens=cu_with_empty) - torch.cuda.synchronize() - - torch.testing.assert_close(o, o_ref, atol=1e-3, rtol=1e-3) - torch.testing.assert_close(fs[0], fs_ref[0], atol=1e-3, rtol=1e-3) - assert (fs[1] == 0).all(), "state of a zero-length sequence must stay zero" - - -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL) -@pytest.mark.parametrize("T1,T2", [(128, 128), (64, 192), (192, 121)]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_chunked_prefill(dtype, T1, T2, H, HV, B=2, K=128, V=128): - """Two-phase prefill: the final state of part 1 feeds part 2 as its - initial state; the concatenated output must match a single-shot run.""" - _seed_all() - T = T1 + T2 - q, k, v = gen_qkv(B, T, H, HV, K, V, dtype) - g, b = gen_kda_gates(B, T, HV, K, dtype) - - def part(t0, t1, S0): - return run_kda( - q[:, t0:t1].contiguous(), - k[:, t0:t1].contiguous(), - v[:, t0:t1].contiguous(), - g[:, t0:t1].contiguous(), - b[:, t0:t1].contiguous(), - initial_state=S0, - output_final_state=True, - ) - - o1, fs1 = part(0, T1, None) - o2, fs2 = part(T1, T, fs1) - o = torch.cat([o1, o2], dim=1) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = kda_reference(q, k, v, g, b) - - # The state round-trips through the op's output dtype between the two - # calls, so allow slightly more than the single-shot forward tolerance. - tol = 1.5 * FWD_TOL[dtype] - r_o = rms_ratio(o, o_ref) - r_s = rms_ratio(fs2, fs_ref) - assert r_o < tol, f"chunked-prefill o rms ratio {r_o:.4g} >= {tol}" - assert r_s < tol, f"chunked-prefill final_state rms ratio {r_s:.4g} >= {tol}" - - -@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL) -@pytest.mark.parametrize("T", [128, 251]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fprop_initial_state(dtype, T, H, HV): - """Random (non-zero) initial recurrent state.""" - _run_fprop_case(dtype, H, HV, 1, T, initial_state=True) - - -@pytest.mark.parametrize("K,V", [(64, 64), (64, 128), (128, 128), (256, 128)]) -@pytest.mark.parametrize("T", [128, 251]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_fprop_head_dims(dtype, T, K, V, H=2, HV=2): - """Head-dim variants: K != V and the K = 256 upper bound.""" - _run_fprop_case(dtype, H, HV, 1, T, K=K, V=V) - - -# --------------------------------------------------------------------------- -# Argument validation -# --------------------------------------------------------------------------- - - -def _dummy_inputs(H=2, HV=2, T=64, K=128, V=128): - q, k, v = gen_qkv(1, T, H, HV, K, V, torch.bfloat16) - g, b = gen_kda_gates(1, T, HV, K, torch.bfloat16) - return q[0], k[0], v[0], g[0], b[0] - - -def _cu(*bounds): - return torch.tensor(bounds, dtype=torch.int32, device="cuda") - - -def test_invalid_qk_head_mismatch(): - q, k, v, g, b = _dummy_inputs() - with pytest.raises(Exception, match="head|No valid engine"): - run_kda(q.unsqueeze(0), k[:, :1].unsqueeze(0), v.unsqueeze(0), g.unsqueeze(0), b.unsqueeze(0), cu_seqlens=_cu(0, 64)) - - -def test_invalid_gva_head_ratio(): - q, k, v, g, b = _dummy_inputs(H=2, HV=2) - with pytest.raises(Exception, match="divisible|multiple|head|No valid engine"): - run_kda(q.unsqueeze(0), k.unsqueeze(0), v.repeat(1, 3, 1)[:, :3].unsqueeze(0), g.unsqueeze(0), b.unsqueeze(0), cu_seqlens=_cu(0, 64)) - - -def test_invalid_initial_state_count(): - q, k, v, g, b = _dummy_inputs(T=128) - S0 = torch.zeros(1, 2, 128, 128, device="cuda", dtype=torch.float32) - with pytest.raises(Exception, match="initial"): - run_kda(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), g.unsqueeze(0), b.unsqueeze(0), cu_seqlens=_cu(0, 64, 128), initial_state=S0) diff --git a/test/python/linear_attention/frost/__init__.py b/test/python/linear_attention/frost/__init__.py deleted file mode 100644 index 52a7a9daf..000000000 --- a/test/python/linear_attention/frost/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 diff --git a/test/python/linear_attention/frost/conftest.py b/test/python/linear_attention/frost/conftest.py deleted file mode 100644 index 3a76f4b73..000000000 --- a/test/python/linear_attention/frost/conftest.py +++ /dev/null @@ -1,23 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Make ``cudnn.linear_attention.frost`` importable: append the source ``python/cudnn`` dir -to the installed ``cudnn`` package's ``__path__`` (the wheel lacks the ``FROST`` -subtree). Unnecessary once the engine ships in the built frontend package.""" - -from __future__ import annotations - -import sys -from pathlib import Path - -import cudnn - -_SRC_CUDNN = Path(__file__).resolve().parents[4] / "python" / "cudnn" -if _SRC_CUDNN.is_dir() and str(_SRC_CUDNN) not in cudnn.__path__: - cudnn.__path__.append(str(_SRC_CUDNN)) - -# The shared helpers live in the linear_attention test package root -# (test/python/linear_attention); pytest only prepends this file's own directory. -_TEST_PY = Path(__file__).resolve().parents[2] -if str(_TEST_PY) not in sys.path: - sys.path.insert(0, str(_TEST_PY)) diff --git a/test/python/linear_attention/frost/examples/01_gdn_prefill.py b/test/python/linear_attention/frost/examples/01_gdn_prefill.py new file mode 100644 index 000000000..547300e49 --- /dev/null +++ b/test/python/linear_attention/frost/examples/01_gdn_prefill.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Example 01: GDN (Gated DeltaNet) prefill (pure cuDNN frontend API). + +Per-token recurrence with scalar decay ``alpha_t = exp(g_t)`` and write +strength ``beta_t``:: + + S_t = alpha_t (I - beta_t k_t^T k_t) S_{t-1} + beta_t k_t^T v_t + 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]``. +""" + +from __future__ import annotations + +import math + +import cudnn +import torch + + +def _build_plans(g) -> None: + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + g.select_plan(names.index("gdn_frost")) # pin the FROST entry + g.check_support() + g.build_plans() + + +def _rms_ratio(out, ref): + out, ref = out.detach().double(), ref.detach().double() + return ((out - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-12)).item() + + +def _reference(q, k, v, g, beta, cu, scale): + """fp64 per-token recurrence over the packed batch. Returns (o, final_state).""" + total, H, D = q.shape + V = v.shape[2] + q, k, v, g, beta = (x.double() for x in (q, k, v, g, beta)) + o = torch.zeros(total, H, V, dtype=torch.float64, device=q.device) + fs = torch.zeros(cu.numel() - 1, H, D, V, dtype=torch.float64, device=q.device) + for n in range(cu.numel() - 1): + S = torch.zeros(H, D, V, dtype=torch.float64, device=q.device) + for t in range(int(cu[n]), int(cu[n + 1])): + a, b = g[t].exp(), beta[t] # [H] scalar decay / write strength + residual = v[t] - a[:, None] * torch.einsum("hd,hdv->hv", k[t], S) + S = a[:, None, None] * S + b[:, None, None] * torch.einsum("hd,hv->hdv", k[t], residual) + o[t] = torch.einsum("hd,hdv->hv", q[t] * scale, S) + fs[n] = S + return o, fs + + +def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: + torch.manual_seed(0) + device = "cuda" + total, num_seqs = sum(seq_lens), len(seq_lens) + scale = 1.0 / math.sqrt(D) + + q = torch.randn(total, H, D, device=device).bfloat16() + k = torch.nn.functional.normalize(torch.randn(total, H, D, device=device), dim=-1).bfloat16() + v = torch.randn(total, H, D, device=device).bfloat16() + gate = torch.empty(total, H, device=device).uniform_(0.1, 1.0).log().contiguous() + beta = torch.rand(total, H, device=device).contiguous() + cu = torch.tensor([0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32, device=device) + + g = cudnn.pygraph() + q_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="q") + k_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="k") + v_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="v") + g_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="g") + beta_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="beta") + cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") + O_t, fs_t, _h_t = g.gdn( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + cu_seqlens=cu_t, + scale=scale, + output_final_state=True, + name="gdn", + ) + O_t.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) + fs_t.set_output(True).set_data_type(cudnn.data_type.FLOAT) + _build_plans(g) + + o = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + fs = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device=device) + pack = {q_t: q, k_t: k, v_t: v, g_t: gate, beta_t: beta, cu_t: cu, O_t: o, fs_t: fs} + g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) + torch.cuda.synchronize() + + o_ref, fs_ref = _reference(q, k, v, gate, beta, cu, scale) + r_o, r_s = _rms_ratio(o, o_ref), _rms_ratio(fs, fs_ref) + assert r_o < 2e-2, f"o rms ratio {r_o:.4g}" + assert r_s < 2e-2, f"final_state rms ratio {r_s:.4g}" + print(f"[01] PASS gdn prefill seq_lens={list(seq_lens)} H={H} D={D} (o rms {r_o:.2e}, fs rms {r_s:.2e})") + + +if __name__ == "__main__": + main() diff --git a/test/python/linear_attention/frost/examples/02_gdn_backward.py b/test/python/linear_attention/frost/examples/02_gdn_backward.py new file mode 100644 index 000000000..6cc05443b --- /dev/null +++ b/test/python/linear_attention/frost/examples/02_gdn_backward.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Example 02: GDN (Gated DeltaNet) backward (pure cuDNN frontend API). + +The GDN_BWD node takes the forward inputs plus ``dO`` and returns +``(dQ, dK, dV, dG, dBeta, dS0)``. Without the optional per-chunk ``h`` input +the engine recomputes the forward state pass internally. Gradients are +checked against fp64 autograd through the per-token recurrence. +""" + +from __future__ import annotations + +import math + +import cudnn +import torch + + +def _build_plans(g) -> None: + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + g.select_plan(names.index("gdn_frost")) # pin the FROST entry + g.check_support() + g.build_plans() + + +def _rms_ratio(out, ref): + out, ref = out.detach().double(), ref.detach().double() + return ((out - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-12)).item() + + +def _reference_o(q, k, v, g, beta, cu, scale): + """Differentiable fp64 per-token recurrence; returns o.""" + total, H, _D = q.shape + V = v.shape[2] + outs = [] + for n in range(cu.numel() - 1): + S = torch.zeros(H, q.shape[2], V, dtype=torch.float64, device=q.device) + for t in range(int(cu[n]), int(cu[n + 1])): + a, b = g[t].exp(), beta[t] + residual = v[t] - a[:, None] * torch.einsum("hd,hdv->hv", k[t], S) + S = a[:, None, None] * S + b[:, None, None] * torch.einsum("hd,hv->hdv", k[t], residual) + outs.append(torch.einsum("hd,hdv->hv", q[t] * scale, S)) + return torch.stack(outs, dim=0) + + +def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: + torch.manual_seed(0) + device = "cuda" + total, num_seqs = sum(seq_lens), len(seq_lens) + scale = 1.0 / math.sqrt(D) + + q = torch.randn(total, H, D, device=device).bfloat16() + k = torch.nn.functional.normalize(torch.randn(total, H, D, device=device), dim=-1).bfloat16() + v = torch.randn(total, H, D, device=device).bfloat16() + gate = torch.empty(total, H, device=device).uniform_(0.1, 1.0).log().contiguous() + beta = torch.rand(total, H, device=device).contiguous() + do = torch.randn(total, H, D, device=device).bfloat16() + cu = torch.tensor([0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32, device=device) + + g = cudnn.pygraph() + q_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="q") + k_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="k") + v_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="v") + g_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="g") + beta_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="beta") + cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") + do_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="dO") + dQ_t, dK_t, dV_t, dG_t, dBeta_t, _dS0_t = g.gdn_bwd( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + cu_seqlens=cu_t, + dO=do_t, + scale=scale, + name="gdn_bwd", + ) + for t_, dt in ( + (dQ_t, cudnn.data_type.BFLOAT16), + (dK_t, cudnn.data_type.BFLOAT16), + (dV_t, cudnn.data_type.BFLOAT16), + (dG_t, cudnn.data_type.FLOAT), + (dBeta_t, cudnn.data_type.FLOAT), + ): + t_.set_output(True).set_data_type(dt) + _build_plans(g) + + dq = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + dk = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + dv = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + dg = torch.empty(total, H, dtype=torch.float32, device=device) + db = torch.empty(total, H, dtype=torch.float32, device=device) + pack = {q_t: q, k_t: k, v_t: v, g_t: gate, beta_t: beta, cu_t: cu, do_t: do, dQ_t: dq, dK_t: dk, dV_t: dv, dG_t: dg, dBeta_t: db} + g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) + torch.cuda.synchronize() + + leaves = [x.double().requires_grad_(True) for x in (q, k, v, gate, beta)] + o_ref = _reference_o(*leaves, cu, scale) + grads = torch.autograd.grad((o_ref * do.double()).sum(), leaves) + for name, out, ref, tol in ( + ("dQ", dq, grads[0], 5e-2), + ("dK", dk, grads[1], 5e-2), + ("dV", dv, grads[2], 5e-2), + ("dG", dg, grads[3], 5e-2), + ("dBeta", db, grads[4], 5e-2), + ): + r = _rms_ratio(out, ref) + assert r < tol, f"{name} rms ratio {r:.4g}" + print(f"[02] PASS gdn backward (recompute) seq_lens={list(seq_lens)} H={H} D={D}") + + +if __name__ == "__main__": + main() diff --git a/test/python/linear_attention/frost/examples/03_kda_prefill.py b/test/python/linear_attention/frost/examples/03_kda_prefill.py new file mode 100644 index 000000000..393870680 --- /dev/null +++ b/test/python/linear_attention/frost/examples/03_kda_prefill.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Example 03: KDA (Kimi Delta Attention) prefill (pure cuDNN frontend API). + +KDA replaces GDN's scalar decay with a per-key-channel decay +``a_t = exp(g_t) in R^K`` (the decayed state feeds the delta-rule residual):: + + S' = Diag(a_t) S_{t-1} + S_t = S' + beta_t k_t^T (v_t - k_t S') + o_t = q_t S_t + +``use_qk_l2norm=False`` passes q/k through as given, so this example feeds +pre-normalized rows. +""" + +from __future__ import annotations + +import math + +import cudnn +import torch + + +def _build_plans(g) -> None: + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + g.select_plan(names.index("kda_frost")) # pin the FROST entry + g.check_support() + g.build_plans() + + +def _rms_ratio(out, ref): + out, ref = out.detach().double(), ref.detach().double() + return ((out - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-12)).item() + + +def _randu(rows, dim, device): + """Per-row uniform [-0.25, 0.25) with normally-distributed means: mildly + heterogeneous data that keeps the recurrence stable.""" + means = torch.randn(rows, 1, device=device) * 0.05 + return means + torch.rand(rows, dim, device=device) * 0.5 - 0.25 + + +def _reference(q, k, v, g, beta, cu, scale): + """fp64 per-token recurrence over the packed batch. Returns (o, final_state).""" + total, H, D = q.shape + V = v.shape[2] + q, k, v, g, beta = (x.double() for x in (q, k, v, g, beta)) + o = torch.zeros(total, H, V, dtype=torch.float64, device=q.device) + fs = torch.zeros(cu.numel() - 1, H, D, V, dtype=torch.float64, device=q.device) + for n in range(cu.numel() - 1): + S = torch.zeros(H, D, V, dtype=torch.float64, device=q.device) + for t in range(int(cu[n]), int(cu[n + 1])): + S = g[t].exp()[..., None] * S # per-key-channel decay first + residual = v[t] - torch.einsum("hd,hdv->hv", k[t], S) + S = S + beta[t][:, None, None] * torch.einsum("hd,hv->hdv", k[t], residual) + o[t] = torch.einsum("hd,hdv->hv", q[t] * scale, S) + fs[n] = S + return o, fs + + +def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: + torch.manual_seed(0) + device = "cuda" + total, num_seqs = sum(seq_lens), len(seq_lens) + scale = 1.0 / math.sqrt(D) + + q = torch.nn.functional.normalize(_randu(total * H, D, device), dim=-1).reshape(total, H, D).bfloat16() + k = torch.nn.functional.normalize(_randu(total * H, D, device), dim=-1).reshape(total, H, D).bfloat16() + v = _randu(total * H, D, device).reshape(total, H, D).bfloat16() + gate = torch.empty(total, H, D, device=device).uniform_(0.5, 1.0).log().contiguous() + beta = torch.rand(total, H, device=device).sigmoid().contiguous() + cu = torch.tensor([0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32, device=device) + + g = cudnn.pygraph() + q_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="q") + k_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="k") + v_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="v") + g_t = g.tensor([total, H, D], data_type=cudnn.data_type.FLOAT, name="g") + beta_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="beta") + cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") + O_t, fs_t, _h_t = g.kda( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + cu_seqlens=cu_t, + scale=scale, + output_final_state=True, + use_qk_l2norm=False, + name="kda", + ) + O_t.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) + fs_t.set_output(True).set_data_type(cudnn.data_type.FLOAT) + _build_plans(g) + + o = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + fs = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device=device) + pack = {q_t: q, k_t: k, v_t: v, g_t: gate, beta_t: beta, cu_t: cu, O_t: o, fs_t: fs} + g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) + torch.cuda.synchronize() + + o_ref, fs_ref = _reference(q, k, v, gate, beta, cu, scale) + r_o = _rms_ratio(o, o_ref) + assert r_o < 5e-2, f"o rms ratio {r_o:.4g}" + r_s = _rms_ratio(fs, fs_ref) + assert r_s < 5e-2, f"final_state rms ratio {r_s:.4g}" + print(f"[03] PASS kda prefill seq_lens={list(seq_lens)} H={H} D={D} (fs rms {r_s:.2e})") + + +if __name__ == "__main__": + main() diff --git a/test/python/linear_attention/frost/examples/04_kda_backward.py b/test/python/linear_attention/frost/examples/04_kda_backward.py new file mode 100644 index 000000000..4a8b24360 --- /dev/null +++ b/test/python/linear_attention/frost/examples/04_kda_backward.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Example 04: KDA (Kimi Delta Attention) backward (pure cuDNN frontend API). + +The KDA_BWD node takes the forward inputs plus ``dO`` and returns +``(dQ, dK, dV, dG, dBeta)`` (per-key-channel ``dG``). Without the optional +per-chunk ``h`` input the engine recomputes the forward state pass +internally. Gradients are checked against fp64 autograd through the +per-token recurrence. +""" + +from __future__ import annotations + +import math + +import cudnn +import torch + + +def _build_plans(g) -> None: + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + g.select_plan(names.index("kda_frost")) # pin the FROST entry + g.check_support() + g.build_plans() + + +def _rms_ratio(out, ref): + out, ref = out.detach().double(), ref.detach().double() + return ((out - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-12)).item() + + +def _reference_o(q, k, v, g, beta, cu, scale): + """Differentiable fp64 per-token recurrence; returns o.""" + total, H, D = q.shape + V = v.shape[2] + outs = [] + for n in range(cu.numel() - 1): + S = torch.zeros(H, D, V, dtype=torch.float64, device=q.device) + for t in range(int(cu[n]), int(cu[n + 1])): + S = g[t].exp()[..., None] * S + residual = v[t] - torch.einsum("hd,hdv->hv", k[t], S) + S = S + beta[t][:, None, None] * torch.einsum("hd,hv->hdv", k[t], residual) + outs.append(torch.einsum("hd,hdv->hv", q[t] * scale, S)) + return torch.stack(outs, dim=0) + + +def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: + torch.manual_seed(0) + device = "cuda" + total, num_seqs = sum(seq_lens), len(seq_lens) + scale = 1.0 / math.sqrt(D) + + q = torch.nn.functional.normalize(torch.randn(total, H, D, device=device), dim=-1).bfloat16() + k = torch.nn.functional.normalize(torch.randn(total, H, D, device=device), dim=-1).bfloat16() + v = torch.randn(total, H, D, device=device).bfloat16() + gate = torch.empty(total, H, D, device=device).uniform_(0.5, 1.0).log().contiguous() + beta = torch.rand(total, H, device=device).sigmoid().contiguous() + do = torch.randn(total, H, D, device=device).bfloat16() + cu = torch.tensor([0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32, device=device) + + g = cudnn.pygraph() + q_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="q") + k_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="k") + v_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="v") + g_t = g.tensor([total, H, D], data_type=cudnn.data_type.FLOAT, name="g") + beta_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="beta") + cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") + do_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="dO") + outs = g.kda_bwd( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + cu_seqlens=cu_t, + dO=do_t, + scale=scale, + use_qk_l2norm=False, + name="kda_bwd", + ) + dtypes = (cudnn.data_type.BFLOAT16,) * 3 + (cudnn.data_type.FLOAT,) * 2 + grads_t = [out.set_output(True).set_data_type(dt) for out, dt in zip(outs, dtypes)] + _build_plans(g) + + dq = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + dk = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + dv = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + dg = torch.empty(total, H, D, dtype=torch.float32, device=device) + db = torch.empty(total, H, dtype=torch.float32, device=device) + pack = {q_t: q, k_t: k, v_t: v, g_t: gate, beta_t: beta, cu_t: cu, do_t: do} + pack.update(dict(zip(grads_t, (dq, dk, dv, dg, db)))) + g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) + torch.cuda.synchronize() + + leaves = [x.double().requires_grad_(True) for x in (q, k, v, gate, beta)] + o_ref = _reference_o(*leaves, cu, scale) + grads = torch.autograd.grad((o_ref * do.double()).sum(), leaves) + for name, out, ref in (("dQ", dq, grads[0]), ("dK", dk, grads[1]), ("dV", dv, grads[2]), ("dG", dg, grads[3]), ("dBeta", db, grads[4])): + r = _rms_ratio(out, ref) + assert r < 5e-2, f"{name} rms ratio {r:.4g}" + print(f"[04] PASS kda backward (recompute) seq_lens={list(seq_lens)} H={H} D={D}") + + +if __name__ == "__main__": + main() diff --git a/test/python/linear_attention/frost/examples/05_gdn2_prefill.py b/test/python/linear_attention/frost/examples/05_gdn2_prefill.py new file mode 100644 index 000000000..19da0445b --- /dev/null +++ b/test/python/linear_attention/frost/examples/05_gdn2_prefill.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Example 05: GDN-2 (Gated DeltaNet v2) prefill (pure cuDNN frontend API). + +GDN-2 gates every channel: per-key decay ``a_t = exp(g_t) in R^K``, per-key +erase gate ``beta_t in R^K``, and per-value write gate ``w_t in R^V``:: + + S' = Diag(a_t) S_{t-1} + S_t = S' + k_t^T (w_t . v_t - (beta_t . k_t) S') + o_t = q_t S_t + +``beta``/``w`` are io-dtype post-sigmoid tensors; ``use_qk_l2norm=False`` +passes q/k through as given, so this example feeds pre-normalized rows. +""" + +from __future__ import annotations + +import math + +import cudnn +import torch + + +def _build_plans(g) -> None: + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + g.select_plan(names.index("gdn2_frost")) # pin the FROST entry + g.check_support() + g.build_plans() + + +def _rms_ratio(out, ref): + out, ref = out.detach().double(), ref.detach().double() + return ((out - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-12)).item() + + +def _randu(rows, dim, device): + """Per-row uniform [-0.25, 0.25) with normally-distributed means: mildly + heterogeneous data that keeps the recurrence stable.""" + means = torch.randn(rows, 1, device=device) * 0.05 + return means + torch.rand(rows, dim, device=device) * 0.5 - 0.25 + + +def _reference(q, k, v, g, beta, w, cu, scale): + """fp64 per-token recurrence over the packed batch. Returns (o, final_state).""" + total, H, D = q.shape + V = v.shape[2] + q, k, v, g, beta, w = (x.double() for x in (q, k, v, g, beta, w)) + o = torch.zeros(total, H, V, dtype=torch.float64, device=q.device) + fs = torch.zeros(cu.numel() - 1, H, D, V, dtype=torch.float64, device=q.device) + for n in range(cu.numel() - 1): + S = torch.zeros(H, D, V, dtype=torch.float64, device=q.device) + for t in range(int(cu[n]), int(cu[n + 1])): + S = g[t].exp()[..., None] * S # per-key-channel decay first + erase = torch.einsum("hd,hdv->hv", beta[t] * k[t], S) + v_new = w[t] * v[t] - erase + S = S + torch.einsum("hd,hv->hdv", k[t], v_new) + o[t] = torch.einsum("hd,hdv->hv", q[t] * scale, S) + fs[n] = S + return o, fs + + +def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: + torch.manual_seed(0) + device = "cuda" + total, num_seqs = sum(seq_lens), len(seq_lens) + scale = 1.0 / math.sqrt(D) + + q = torch.nn.functional.normalize(_randu(total * H, D, device), dim=-1).reshape(total, H, D).bfloat16() + k = torch.nn.functional.normalize(_randu(total * H, D, device), dim=-1).reshape(total, H, D).bfloat16() + v = _randu(total * H, D, device).reshape(total, H, D).bfloat16() + gate = torch.empty(total, H, D, device=device).uniform_(0.5, 1.0).log().contiguous() + beta = (torch.rand(total, H, D, device=device).sigmoid() * 2.0).bfloat16().contiguous() + w = torch.rand(total, H, D, device=device).sigmoid().bfloat16().contiguous() + cu = torch.tensor([0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32, device=device) + + g = cudnn.pygraph() + q_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="q") + k_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="k") + v_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="v") + g_t = g.tensor([total, H, D], data_type=cudnn.data_type.FLOAT, name="g") + beta_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="beta") + w_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="w") + cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") + O_t, fs_t, _h_t = g.gdn2( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + w=w_t, + cu_seqlens=cu_t, + scale=scale, + output_final_state=True, + use_qk_l2norm=False, + name="gdn2", + ) + O_t.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) + fs_t.set_output(True).set_data_type(cudnn.data_type.FLOAT) + _build_plans(g) + + o = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + fs = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device=device) + pack = {q_t: q, k_t: k, v_t: v, g_t: gate, beta_t: beta, w_t: w, cu_t: cu, O_t: o, fs_t: fs} + g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) + torch.cuda.synchronize() + + o_ref, fs_ref = _reference(q, k, v, gate, beta, w, cu, scale) + r_o = _rms_ratio(o, o_ref) + assert r_o < 5e-2, f"o rms ratio {r_o:.4g}" + r_s = _rms_ratio(fs, fs_ref) + assert r_s < 5e-2, f"final_state rms ratio {r_s:.4g}" + print(f"[05] PASS gdn2 prefill seq_lens={list(seq_lens)} H={H} D={D} (fs rms {r_s:.2e})") + + +if __name__ == "__main__": + main() diff --git a/test/python/linear_attention/frost/examples/06_gdn2_backward.py b/test/python/linear_attention/frost/examples/06_gdn2_backward.py new file mode 100644 index 000000000..8bfb6a7d4 --- /dev/null +++ b/test/python/linear_attention/frost/examples/06_gdn2_backward.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Example 06: GDN-2 (Gated DeltaNet v2) backward (pure cuDNN frontend API). + +The GDN2_BWD node takes the forward inputs plus ``dO`` and returns +``(dQ, dK, dV, dG, dBeta, dW)`` (per-key-channel ``dG``/``dBeta``, per-value +``dW``; ``dBeta``/``dW`` in io dtype). Without the optional per-chunk ``h`` +input the engine recomputes the forward state pass internally. Gradients are +checked against fp64 autograd through the per-token recurrence. +""" + +from __future__ import annotations + +import math + +import cudnn +import torch + + +def _build_plans(g) -> None: + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + g.select_plan(names.index("gdn2_frost")) # pin the FROST entry + g.check_support() + g.build_plans() + + +def _rms_ratio(out, ref): + out, ref = out.detach().double(), ref.detach().double() + return ((out - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-12)).item() + + +def _reference_o(q, k, v, g, beta, w, cu, scale): + """Differentiable fp64 per-token recurrence; returns o.""" + total, H, D = q.shape + V = v.shape[2] + outs = [] + for n in range(cu.numel() - 1): + S = torch.zeros(H, D, V, dtype=torch.float64, device=q.device) + for t in range(int(cu[n]), int(cu[n + 1])): + S = g[t].exp()[..., None] * S + erase = torch.einsum("hd,hdv->hv", beta[t] * k[t], S) + v_new = w[t] * v[t] - erase + S = S + torch.einsum("hd,hv->hdv", k[t], v_new) + outs.append(torch.einsum("hd,hdv->hv", q[t] * scale, S)) + return torch.stack(outs, dim=0) + + +def main(seq_lens=(192, 320), H: int = 2, D: int = 128) -> None: + torch.manual_seed(0) + device = "cuda" + total, num_seqs = sum(seq_lens), len(seq_lens) + scale = 1.0 / math.sqrt(D) + + q = torch.nn.functional.normalize(torch.randn(total, H, D, device=device), dim=-1).bfloat16() + k = torch.nn.functional.normalize(torch.randn(total, H, D, device=device), dim=-1).bfloat16() + v = torch.randn(total, H, D, device=device).bfloat16() + gate = torch.empty(total, H, D, device=device).uniform_(0.5, 1.0).log().contiguous() + beta = (torch.rand(total, H, D, device=device).sigmoid() * 2.0).bfloat16().contiguous() + w = torch.rand(total, H, D, device=device).sigmoid().bfloat16().contiguous() + do = torch.randn(total, H, D, device=device).bfloat16() + cu = torch.tensor([0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32, device=device) + + g = cudnn.pygraph() + q_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="q") + k_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="k") + v_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="v") + g_t = g.tensor([total, H, D], data_type=cudnn.data_type.FLOAT, name="g") + beta_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="beta") + w_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="w") + cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") + do_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="dO") + outs = g.gdn2_bwd( + q=q_t, + k=k_t, + v=v_t, + g=g_t, + beta=beta_t, + w=w_t, + cu_seqlens=cu_t, + dO=do_t, + scale=scale, + name="gdn2_bwd", + ) + dtypes = (cudnn.data_type.BFLOAT16,) * 3 + (cudnn.data_type.FLOAT, cudnn.data_type.BFLOAT16, cudnn.data_type.BFLOAT16) + grads_t = [out.set_output(True).set_data_type(dt) for out, dt in zip(outs, dtypes)] + _build_plans(g) + + dq = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + dk = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + dv = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + dg = torch.empty(total, H, D, dtype=torch.float32, device=device) + db = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + dw = torch.empty(total, H, D, dtype=torch.bfloat16, device=device) + pack = {q_t: q, k_t: k, v_t: v, g_t: gate, beta_t: beta, w_t: w, cu_t: cu, do_t: do} + pack.update(dict(zip(grads_t, (dq, dk, dv, dg, db, dw)))) + g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) + torch.cuda.synchronize() + + leaves = [x.double().requires_grad_(True) for x in (q, k, v, gate, beta, w)] + o_ref = _reference_o(*leaves, cu, scale) + grads = torch.autograd.grad((o_ref * do.double()).sum(), leaves) + for name, out, ref in ( + ("dQ", dq, grads[0]), + ("dK", dk, grads[1]), + ("dV", dv, grads[2]), + ("dG", dg, grads[3]), + ("dBeta", db, grads[4]), + ("dW", dw, grads[5]), + ): + r = _rms_ratio(out, ref) + assert r < 5e-2, f"{name} rms ratio {r:.4g}" + print(f"[06] PASS gdn2 backward (recompute) seq_lens={list(seq_lens)} H={H} D={D}") + + +if __name__ == "__main__": + main() diff --git a/test/python/linear_attention/frost/test_gdn2_bprop_kernel.py b/test/python/linear_attention/frost/test_gdn2_bprop_kernel.py deleted file mode 100644 index 1c23dbfdb..000000000 --- a/test/python/linear_attention/frost/test_gdn2_bprop_kernel.py +++ /dev/null @@ -1,44 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""FROST GDN-2 backward: currently a STUB on this branch (the small-chunk -design recomputes the forward states in the backward once the kernel lands). -The contract: ``Gdn2FrostEngine`` declines ``GDN2_BWD`` -graphs so the router can fall back.""" - -from __future__ import annotations - -import pytest - -import cudnn - -from linear_attention.common import assert_engine_declines - -pytestmark = pytest.mark.L0 - - -def test_gdn2_bwd_frost_engine_declines(): - - total, H, D = 256, 2, 128 - g = cudnn.pygraph() - q_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="q") - k_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="k") - v_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="v") - g_t = g.tensor([total, H, D], data_type=cudnn.data_type.FLOAT, name="g") - beta_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="beta") - w_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="w") - cu_t = g.tensor([2], data_type=cudnn.data_type.INT32, name="cu_seqlens") - dO_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="dO") - g.gdn2_bwd( - q=q_t, - k=k_t, - v=v_t, - g=g_t, - beta=beta_t, - w=w_t, - cu_seqlens=cu_t, - dO=dO_t, - scale=0.125, - name="gdn2_bwd", - ) - assert_engine_declines(g, "gdn2_frost") diff --git a/test/python/linear_attention/frost/test_gdn2_prefill_kernel.py b/test/python/linear_attention/frost/test_gdn2_prefill_kernel.py deleted file mode 100644 index 0fdefc555..000000000 --- a/test/python/linear_attention/frost/test_gdn2_prefill_kernel.py +++ /dev/null @@ -1,676 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""FROST GDN-2 prefill tests (``pygraph`` + ``Gdn2FrostEngine``) against the -fp64 recurrent reference.""" - -from __future__ import annotations - -import math -import random -from itertools import accumulate - -import pytest -import torch -import torch.nn.functional as F - -import cudnn # noqa: F401 (conftest extends cudnn.__path__ with the source tree) -from cudnn.linear_attention.frost import Gdn2FrostEngine - -from linear_attention.common import assert_bitwise_runs, assert_concurrent_stream_runs, assert_engine_declines -from linear_attention.conftest import multidist_randu -from linear_attention.reference_gdn2 import gdn2_reference, rms_ratio - -pytestmark = pytest.mark.L0 - -SEED = 42 - - -def _sm100_dsl_available() -> bool: - if not torch.cuda.is_available(): - return False - major, _minor = torch.cuda.get_device_capability() - if major != 10: - return False - try: - import cutlass.experimental.primitives # noqa: F401 -- often a sys.modules alias, invisible to find_spec - except ImportError: - return False - return True - - -requires_runtime = pytest.mark.skipif(not _sm100_dsl_available(), reason="needs an SM100-class GPU and the Cutlass DSL") - - -def _seed(seed=SEED): - random.seed(seed) - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - - -def _cu(seq_lens, device="cuda"): - return torch.tensor([0] + list(accumulate(seq_lens)), dtype=torch.int32, device=device) - - -def _gen_thd(seq_lens, H, HV, head_size, dtype, HK=None): - HK = H if HK is None else HK - total = sum(seq_lens) - q = multidist_randu(total * H, head_size, device="cuda").reshape(total, H, head_size) - k = F.normalize(multidist_randu(total * HK, head_size, device="cuda").reshape(total, HK, head_size), p=2.0, dim=-1) - v = multidist_randu(total * HV, head_size, device="cuda").reshape(total, HV, head_size) - return q.to(dtype).contiguous(), k.to(dtype).contiguous(), v.to(dtype).contiguous() - - -_DT = {torch.bfloat16: cudnn.data_type.BFLOAT16, torch.float16: cudnn.data_type.HALF, torch.float32: cudnn.data_type.FLOAT} - - -def _build_gdn2_engine_graph( - total, - H, - D, - num_seqs, - scale, - *, - io_dt=None, - HV=None, - use_qk_l2norm=True, - with_s0=False, - s0_dt=None, - fs_dt=None, - h_dt=None, - checkpoint_every_n_tokens=0, - use_beta_w_sigmoid=False, - bwd=False, -): - - io_dt = io_dt or cudnn.data_type.BFLOAT16 - HV = HV or H - HO = max(H, HV) - g = cudnn.pygraph() - q_t = g.tensor([total, H, D], data_type=io_dt, name="q") - k_t = g.tensor([total, H, D], data_type=io_dt, name="k") - v_t = g.tensor([total, HV, D], data_type=io_dt, name="v") - g_t = g.tensor([total, HO, D], data_type=cudnn.data_type.FLOAT, name="g") - beta_t = g.tensor([total, HO, D], data_type=io_dt, name="beta") - w_t = g.tensor([total, HO, D], data_type=io_dt, name="w") - cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") - t = dict(q=q_t, k=k_t, v=v_t, g=g_t, beta=beta_t, w=w_t, cu=cu_t) - if with_s0: - t["s0"] = g.tensor([num_seqs, HO, D, D], data_type=s0_dt or cudnn.data_type.FLOAT, name="initial_state") - if bwd: - t["dO"] = g.tensor([total, HO, D], data_type=io_dt, name="dO") - g.gdn2_bwd( - q=q_t, - k=k_t, - v=v_t, - g=g_t, - beta=beta_t, - w=w_t, - cu_seqlens=cu_t, - dO=t["dO"], - scale=scale, - name="gdn2_bwd", - ) - return g, t - O_t, fs_t, h_t = g.gdn2( - q=q_t, - k=k_t, - v=v_t, - g=g_t, - beta=beta_t, - w=w_t, - cu_seqlens=cu_t, - initial_state=t.get("s0"), - scale=scale, - output_final_state=True, - use_qk_l2norm=use_qk_l2norm, - checkpoint_every_n_tokens=checkpoint_every_n_tokens, - use_beta_w_sigmoid=use_beta_w_sigmoid, - name="gdn2", - ) - O_t.set_output(True).set_data_type(io_dt) - fs_t.set_output(True).set_data_type(fs_dt or cudnn.data_type.FLOAT) - t["O"], t["fs"] = O_t, fs_t - if h_t is not None: - h_t.set_output(True).set_data_type(h_dt or io_dt) - t["H"] = h_t - return g, t - - -def _run_gdn2(q, k, v, gate, beta, w, scale, cu, initial_state=None, output_state=None, out_h=None, every_n=0, use_qk_l2norm=True): - """Torch adapter over the graph. ``gate`` is the per-key- - channel natural-log decay (fp32); ``beta``/``w`` are the io-dtype - channel-wise gates; state ports are K-major ``[N, HO, K, V]``. GQA ``k`` - (and ``HV < HQ`` ``v``) are pre-broadcast: the node serves HK == HQ, HV a - multiple of HQ.""" - device = q.device - H, HV = q.shape[1], v.shape[1] - if k.shape[1] != H: - k = k.repeat_interleave(H // k.shape[1], dim=1) - if HV < H: - v = v.repeat_interleave(H // HV, dim=1) - HV = H - total, D = q.shape[0], q.shape[2] - HO = max(H, HV) - num_seqs = cu.shape[0] - 1 - if output_state is None: - output_state = torch.empty(num_seqs, HO, D, D, dtype=torch.float32, device=device) - g, t = _build_gdn2_engine_graph( - total, - H, - D, - num_seqs, - scale, - io_dt=_DT[q.dtype], - HV=HV, - use_qk_l2norm=use_qk_l2norm, - with_s0=initial_state is not None, - s0_dt=None if initial_state is None else _DT[initial_state.dtype], - fs_dt=_DT[output_state.dtype], - h_dt=None if out_h is None else _DT[out_h.dtype], - checkpoint_every_n_tokens=every_n, - ) - g.build() - output = torch.empty(total, HO, D, dtype=q.dtype, device=device) - pack = { - t["q"]: q.contiguous(), - t["k"]: k.contiguous(), - t["v"]: v.contiguous(), - t["g"]: gate.float().contiguous(), - t["beta"]: beta.contiguous(), - t["w"]: w.contiguous(), - t["cu"]: cu.to(torch.int32).contiguous(), - t["O"]: output, - t["fs"]: output_state, - } - if initial_state is not None: - pack[t["s0"]] = initial_state.contiguous() - if out_h is not None: - pack[t["H"]] = out_h - g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) - return output, output_state - - -def _run_and_check(dtype, H, HV, seq_lens, with_s0=False, HK=None, gate_lo=None): - """Run the engine on random inputs and compare against the fp64 recurrent - reference (the kernel L2-normalizes q/k internally). fp16 io needs a - higher gate floor: k * exp2(-cumsum(g)) grazes the fp16 range at 0.5.""" - _seed() - head_size = 128 - total = sum(seq_lens) - HO = max(H, HV) - num_seqs = len(seq_lens) - q, k, v = _gen_thd(seq_lens, H, HV, head_size, dtype, HK=HK) - lo = gate_lo if gate_lo is not None else (0.6 if dtype == torch.float16 else 0.5) - gate = torch.empty(total, HO, head_size, device="cuda").uniform_(lo, 1.0).log() - beta = (torch.rand(total, HO, head_size, device="cuda").sigmoid() * 2.0).to(dtype) - w = torch.rand(total, HO, head_size, device="cuda").sigmoid().to(dtype) - s0 = (torch.randn(num_seqs, HO, head_size, head_size, dtype=torch.float32, device="cuda") * 0.05).contiguous() if with_s0 else None - output_state = torch.full((num_seqs, HO, head_size, head_size), float("nan"), dtype=torch.float32, device="cuda") - - scale = 1.0 / math.sqrt(head_size) - output, _ = _run_gdn2(q, k, v, gate, beta, w, scale, _cu(seq_lens), initial_state=s0, output_state=output_state) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = gdn2_reference( - F.normalize(q.float(), dim=-1).unsqueeze(0), - F.normalize(k.float(), dim=-1).unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - w.unsqueeze(0), - scale=scale, - initial_state=s0, - cu_seqlens=_cu(seq_lens), - ) - tol = 1.2e-1 if dtype == torch.float16 else 1e-1 - torch.testing.assert_close(output.float(), o_ref.squeeze(0).float(), atol=tol, rtol=tol) - assert rms_ratio(output_state, fs_ref) < 5e-2 # engine state ports are K-major, like the reference - - -@requires_runtime -@pytest.mark.parametrize("H,HV", [(1, 1), (2, 2), (2, 4)]) -@pytest.mark.parametrize("seq_lens", [[256], [256, 256], [64, 128, 512]]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_gdn2_kernel_basic(dtype, H, HV, seq_lens): - _run_and_check(dtype, H, HV, seq_lens) - - -@requires_runtime -@pytest.mark.parametrize( - "seq_lens", - [[1], [15], [16], [17], [63, 65], [240, 255, 257], [7] * 24 + [1] * 8, [2048], [33] * 200], - ids=lambda s: f"{len(s)}seqs_{sum(s)}tok", -) -def test_gdn2_kernel_seqlen_edges(seq_lens): - """Chunk-boundary lengths (BT=16 +/- 1), single-token, many-short-seq - packs, a long sequence (many mbarrier ring wraps), and more tiles than - SMs (persistent CTAs walking several tiles).""" - _run_and_check(torch.bfloat16, 2, 4, seq_lens) - - -@requires_runtime -@pytest.mark.parametrize("seq_lens", [[240, 255, 257], [7] * 24 + [1] * 8]) -def test_gdn2_kernel_seqlen_edges_fp16(seq_lens): - _run_and_check(torch.float16, 2, 4, seq_lens) - - -@requires_runtime -@pytest.mark.parametrize("seq_lens", [[256], [64, 129, 512]]) -def test_gdn2_kernel_initial_state(seq_lens): - _run_and_check(torch.bfloat16, 2, 2, seq_lens, with_s0=True) - - -@requires_runtime -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_gdn2_beta_w_sigmoid_in_kernel(dtype): - """use_beta_w_sigmoid: io-dtype beta/w logits + in-kernel sigmoid must - match the host-side sigmoid path (approx-tanh sigmoid and the io-dtype - roundtrip differ by at most ~1 ulp per gate).""" - _seed() - seq_lens, H, D = [256, 129], 2, 128 - total, num_seqs = sum(seq_lens), len(seq_lens) - q, k, v = _gen_thd(seq_lens, H, H, D, dtype) - gate = torch.empty(total, H, D, device="cuda").uniform_(0.5, 1.0).log().float().contiguous() - beta_logits = torch.randn(total, H, D, device="cuda") - w_logits = torch.randn(total, H, D, device="cuda") - scale = 1.0 / math.sqrt(D) - cu = _cu(seq_lens) - - outs = [] - for in_kernel in (True, False): - if in_kernel: - beta, w = beta_logits.to(dtype).contiguous(), w_logits.to(dtype).contiguous() - else: - beta = beta_logits.to(dtype).float().sigmoid().to(dtype).contiguous() - w = w_logits.to(dtype).float().sigmoid().to(dtype).contiguous() - g, t = _build_gdn2_engine_graph(total, H, D, num_seqs, scale, io_dt=_DT[dtype], use_beta_w_sigmoid=in_kernel) - g.build() - o = torch.empty(total, H, D, dtype=dtype, device="cuda") - fs = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device="cuda") - pack = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta, t["w"]: w, t["cu"]: cu, t["O"]: o, t["fs"]: fs} - g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda")) - torch.cuda.synchronize() - outs.append((o, fs)) - torch.testing.assert_close(outs[0][0].float(), outs[1][0].float(), atol=2.5e-2, rtol=2.5e-2) - assert rms_ratio(outs[0][1], outs[1][1]) < 2e-2 - - -@requires_runtime -@pytest.mark.parametrize("with_initial_state", [False, True]) -def test_gdn2_zero_length_sequence_state_passthrough(with_initial_state): - """A zero-length sequence's final-state slot gets the passthrough value: - its initial state when seeded, zeros otherwise.""" - _seed() - seq_len, H, D, sentinel = 256, 2, 128, 123.0 - q, k, v = _gen_thd([seq_len], H, H, D, torch.bfloat16) - gate = torch.empty(seq_len, H, D, device="cuda").uniform_(0.5, 1.0).log().float().contiguous() - beta = (torch.rand(seq_len, H, D, device="cuda").sigmoid() * 2.0).to(torch.bfloat16) - w = torch.rand(seq_len, H, D, device="cuda").sigmoid().to(torch.bfloat16) - s0 = torch.randn(2, H, D, D, dtype=torch.float32, device="cuda") if with_initial_state else None - fs = torch.full((2, H, D, D), sentinel, dtype=torch.float32, device="cuda") - _run_gdn2(q, k, v, gate, beta, w, 1.0 / math.sqrt(D), _cu([seq_len, 0]), initial_state=s0, output_state=fs) - torch.cuda.synchronize() - expected = s0[1] if with_initial_state else torch.zeros_like(fs[1]) - torch.testing.assert_close(fs[1], expected, atol=0, rtol=0) - - -@requires_runtime -@pytest.mark.parametrize("HQ,HK,HV", [(1, 1, 1), (4, 1, 1), (3, 3, 3), (6, 2, 2), (1, 1, 2), (2, 2, 4), (16, 16, 32), (16, 16, 64)]) -def test_gdn2_kernel_head_configs(HQ, HK, HV): - """The FlashInfer head-config matrix: GQA (HK < HQ), GVA (HV > HQ), odd - head counts, and production-sized 16/32/64-head grids (the adapter - pre-broadcasts k and any HV < HQ v to the node's HK == HQ contract).""" - _run_and_check(torch.bfloat16, HQ, HV, [64, 128, 512], HK=HK) - - -@requires_runtime -@pytest.mark.parametrize("HQ,HK,HV", [(2, 2, 4), (16, 16, 64)]) -def test_gdn2_kernel_head_configs_fp16(HQ, HK, HV): - _run_and_check(torch.float16, HQ, HV, [64, 128, 512], HK=HK) - - -@requires_runtime -def test_gdn2_chunked_prefill(): - """Splitting a sequence at a chunk boundary and reseeding from the fp32 - final state must reproduce the single-shot run (state roundtrip; the two - runs get different split tables, so the match is approximate).""" - _seed() - H, D, T1, T2 = 2, 128, 256, 192 - dtype = torch.bfloat16 - q, k, v = _gen_thd([T1 + T2], H, H, D, dtype) - gate = torch.empty(T1 + T2, H, D, device="cuda").uniform_(0.5, 1.0).log() - beta = (torch.rand(T1 + T2, H, D, device="cuda").sigmoid() * 2.0).to(dtype) - w = torch.rand(T1 + T2, H, D, device="cuda").sigmoid().to(dtype) - scale = 1.0 / math.sqrt(D) - - fs_full = torch.full((1, H, D, D), float("nan"), dtype=torch.float32, device="cuda") - o_full, _ = _run_gdn2(q, k, v, gate, beta, w, scale, _cu([T1 + T2]), output_state=fs_full) - - fs1 = torch.full((1, H, D, D), float("nan"), dtype=torch.float32, device="cuda") - o1, _ = _run_gdn2(q[:T1], k[:T1], v[:T1], gate[:T1], beta[:T1], w[:T1], scale, _cu([T1]), output_state=fs1) - fs2 = torch.full((1, H, D, D), float("nan"), dtype=torch.float32, device="cuda") - o2, _ = _run_gdn2( - q[T1:].contiguous(), - k[T1:].contiguous(), - v[T1:].contiguous(), - gate[T1:].contiguous(), - beta[T1:].contiguous(), - w[T1:].contiguous(), - scale, - _cu([T2]), - initial_state=fs1, - output_state=fs2, - ) - torch.cuda.synchronize() - - torch.testing.assert_close(o1.float(), o_full[:T1].float(), atol=2e-3, rtol=2e-3) - torch.testing.assert_close(o2.float(), o_full[T1:].float(), atol=2e-3, rtol=2e-3) - torch.testing.assert_close(fs2, fs_full, atol=2e-3, rtol=2e-3) - - -@requires_runtime -def test_gdn2_kernel_state_dtype_bf16(): - """bf16 initial/final state buffers (the io-downcast S0 path).""" - _seed() - seq_lens, H, D = [64, 128, 512], 2, 128 - total, num_seqs = sum(seq_lens), len(seq_lens) - dtype = torch.bfloat16 - q, k, v = _gen_thd(seq_lens, H, H, D, dtype) - gate = torch.empty(total, H, D, device="cuda").uniform_(0.5, 1.0).log() - beta = (torch.rand(total, H, D, device="cuda").sigmoid() * 2.0).to(dtype) - w = torch.rand(total, H, D, device="cuda").sigmoid().to(dtype) - scale = 1.0 / math.sqrt(D) - s0 = (torch.randn(num_seqs, H, D, D, dtype=torch.float32, device="cuda") * 0.05).to(dtype).contiguous() - fs = torch.full((num_seqs, H, D, D), float("nan"), dtype=dtype, device="cuda") - o, _ = _run_gdn2(q, k, v, gate, beta, w, scale, _cu(seq_lens), initial_state=s0, output_state=fs) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = gdn2_reference( - F.normalize(q.float(), dim=-1).unsqueeze(0), - F.normalize(k.float(), dim=-1).unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - w.unsqueeze(0), - scale=scale, - initial_state=s0.float(), - cu_seqlens=_cu(seq_lens), - ) - torch.testing.assert_close(o.float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - assert rms_ratio(fs.float(), fs_ref) < 5e-2 - - -# --------------------------------------------------------------------------- -# Gdn2FrostEngine: graph-level coverage through the router -# --------------------------------------------------------------------------- - - -def _gdn2_engine_inputs(seq_lens, H, D): - from linear_attention.conftest import gen_gdn2_gates, gen_qkv - - total = sum(seq_lens) - q, k, v = gen_qkv(1, total, H, H, D, D, torch.bfloat16) - gate, beta, w = gen_gdn2_gates(1, total, H, D, D, torch.bfloat16) - cu = torch.tensor([0] + list(accumulate(seq_lens)), dtype=torch.int32, device="cuda") - return (x.squeeze(0).contiguous() for x in (q, k, v, gate, beta, w)), cu - - -@requires_runtime -@pytest.mark.parametrize("seq_lens", [[256], [512, 512], [64, 128, 512]]) -def test_gdn2_frost_engine_matches_reference(seq_lens, H=2, D=128): - - _seed() - (q, k, v, gate, beta, w), cu = _gdn2_engine_inputs(seq_lens, H, D) - total, num_seqs = sum(seq_lens), len(seq_lens) - scale = 1.0 / math.sqrt(D) - - g, t = _build_gdn2_engine_graph(total, H, D, num_seqs, scale) - g.build() - assert isinstance(g.selected_engine, Gdn2FrostEngine) - assert g.get_workspace_size() > 0 # split-K work-item table + scheduler counters - - o_buf = torch.empty(total, H, D, dtype=torch.bfloat16, device="cuda") - fs_buf = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device="cuda") - pack = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta, t["w"]: w, t["cu"]: cu, t["O"]: o_buf, t["fs"]: fs_buf} - g.execute(pack, torch.empty(g.get_workspace_size(), dtype=torch.uint8, device="cuda")) - torch.cuda.synchronize() - - # the kernel L2-normalizes q/k internally - with torch.no_grad(): - o_ref, fs_ref = gdn2_reference( - F.normalize(q.float(), dim=-1).unsqueeze(0), - F.normalize(k.float(), dim=-1).unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - w.unsqueeze(0), - scale=scale, - cu_seqlens=cu, - ) - torch.testing.assert_close(o_buf.float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - r_s = rms_ratio(fs_buf, fs_ref) # engine state ports are K-major - assert r_s < 5e-2, f"final_state rms ratio {r_s:.4g}" - - -@requires_runtime -def test_gdn2_frost_engine_no_l2norm_matches_reference(seq_lens=(256,), H=2, D=128): - """use_qk_l2norm=False passes q/k through as given, so the test feeds - pre-normalized rows (the kernel's io-dtype arithmetic needs them).""" - - _seed() - (q, k, v, gate, beta, w), cu = _gdn2_engine_inputs(list(seq_lens), H, D) - q = F.normalize(q.float(), dim=-1).to(q.dtype) - k = F.normalize(k.float(), dim=-1).to(k.dtype) - total, num_seqs = sum(seq_lens), len(seq_lens) - scale = 1.0 / math.sqrt(D) - - g, t = _build_gdn2_engine_graph(total, H, D, num_seqs, scale, use_qk_l2norm=False) - g.build() - assert isinstance(g.selected_engine, Gdn2FrostEngine) - - o_buf = torch.empty(total, H, D, dtype=torch.bfloat16, device="cuda") - fs_buf = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device="cuda") - pack = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta, t["w"]: w, t["cu"]: cu, t["O"]: o_buf, t["fs"]: fs_buf} - g.execute(pack, torch.empty(g.get_workspace_size(), dtype=torch.uint8, device="cuda")) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = gdn2_reference( - q.float().unsqueeze(0), - k.float().unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - w.unsqueeze(0), - scale=scale, - cu_seqlens=cu, - ) - torch.testing.assert_close(o_buf.float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - r_s = rms_ratio(fs_buf, fs_ref) # engine state ports are K-major - assert r_s < 5e-2, f"final_state rms ratio {r_s:.4g}" - - -@requires_runtime -def test_gdn2_frost_engine_initial_state(seq_lens=(128, 256), H=2, D=128): - - _seed() - (q, k, v, gate, beta, w), cu = _gdn2_engine_inputs(seq_lens, H, D) - total, num_seqs = sum(seq_lens), len(seq_lens) - scale = 1.0 / math.sqrt(D) - s0 = torch.randn(num_seqs, H, D, D, dtype=torch.float32, device="cuda") * 0.05 - - g, t = _build_gdn2_engine_graph(total, H, D, num_seqs, scale, with_s0=True) - g.build() - assert isinstance(g.selected_engine, Gdn2FrostEngine) - - o_buf = torch.empty(total, H, D, dtype=torch.bfloat16, device="cuda") - fs_buf = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device="cuda") - pack = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta, t["w"]: w, t["cu"]: cu, t["s0"]: s0, t["O"]: o_buf, t["fs"]: fs_buf} - g.execute(pack, torch.empty(g.get_workspace_size(), dtype=torch.uint8, device="cuda")) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = gdn2_reference( - F.normalize(q.float(), dim=-1).unsqueeze(0), - F.normalize(k.float(), dim=-1).unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - w.unsqueeze(0), - scale=scale, - initial_state=s0, - cu_seqlens=cu, - ) - torch.testing.assert_close(o_buf.float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - r_s = rms_ratio(fs_buf, fs_ref) # engine state ports are K-major - assert r_s < 5e-2, f"final_state rms ratio {r_s:.4g}" - - -def test_gdn2_frost_engine_declines_bwd(): - """GDN2_BWD declines (stub backward kernel on this branch).""" - g, _t = _build_gdn2_engine_graph(256, 2, 128, 1, 0.125, bwd=True) - assert_engine_declines(g, "gdn2_frost") - - -@requires_runtime -def test_gdn2_frost_engine_declines_wrong_head_dim(): - g, _t = _build_gdn2_engine_graph(256, 2, 64, 1, 0.125) - assert_engine_declines(g, "gdn2_frost") - - -# --------------------------------------------------------------------------- -# Determinism stress: bitwise repeat runs + two-stream co-residency -# --------------------------------------------------------------------------- - -DET_VARLEN_MIX = [497, 16, 1, 480, 0, 253] # zero-length + single-token + odd tails - - -def _det_launch(seq_lens, H, HV, with_s0=False, stream=None): - q, k, v = _gen_thd(seq_lens, H, HV, 128, torch.bfloat16) - total, HO = sum(seq_lens), max(H, HV) - gate = torch.empty(total, HO, 128, device="cuda").uniform_(0.5, 1.0).log().float().contiguous() - beta = (torch.rand(total, HO, 128, device="cuda").sigmoid() * 2.0).to(torch.bfloat16).contiguous() - w = torch.rand(total, HO, 128, device="cuda").sigmoid().to(torch.bfloat16).contiguous() - s0 = (torch.randn(len(seq_lens), HO, 128, 128, dtype=torch.float32, device="cuda") * 0.05).contiguous() if with_s0 else None - cu = _cu(seq_lens) - scale = 1.0 / math.sqrt(128) - - def launch(): - if stream is not None: - with torch.cuda.stream(stream): - return _run_gdn2(q, k, v, gate, beta, w, scale, cu, initial_state=s0) - return _run_gdn2(q, k, v, gate, beta, w, scale, cu, initial_state=s0) - - return launch - - -@requires_runtime -@pytest.mark.parametrize("seq_lens", [DET_VARLEN_MIX, [4096]], ids=["varlen_mix", "long"]) -def test_gdn2_prefill_determinism(seq_lens): - _seed() - assert_bitwise_runs(_det_launch(seq_lens, 2, 4), label="gdn2") - - -@requires_runtime -def test_gdn2_prefill_determinism_initial_state(): - _seed() - assert_bitwise_runs(_det_launch([256, 640, 0, 33], 2, 2, with_s0=True), label="gdn2+s0") - - -@requires_runtime -def test_gdn2_concurrent_streams_determinism(): - _seed() - s1, s2 = torch.cuda.Stream(), torch.cuda.Stream() - assert_concurrent_stream_runs(_det_launch([1024, 31, 512], 2, 4, stream=s1), _det_launch([256, 999, 1, 128], 2, 4, stream=s2), s1, s2) - - -# --------------------------------------------------------------------------- -# split-K: strong decay on a long-sequence pack, so the table actually cuts -# (the per-channel gate scan itself is covered in test_kda_prefill_kernel.py) -# --------------------------------------------------------------------------- - - -@requires_runtime -@pytest.mark.parametrize("with_s0", [False, True]) -@pytest.mark.parametrize("heads", [(2, 2), (1, 4)]) # MHA, GVA -def test_gdn2_prefill_split_strong_decay(heads, with_s0): - """Strong decay (gate floor 0.3) saturates the scan's warmup threshold on - the 2048-token sequence, so the split table cuts it into several work - items; checked against the fp64 recurrent reference.""" - HQ, HV = heads - _run_and_check(torch.bfloat16, HQ, HV, [100, 2048, 0, 517], with_s0=with_s0, gate_lo=0.3) - - -# --------------------------------------------------------------------------- -# CUDA graph capture/replay across dynamic shapes (fixed SM-count grid) -# --------------------------------------------------------------------------- - - -@requires_runtime -def test_gdn2_prefill_cuda_graph_replay(): - """Capture the ENGINE execute once, replay across CHANGED effective - shapes: capacity buffers, cu_seqlens with zero-length tails; everything - the engine launches (sched memset, split table, desc rebuilds, kernel) - must be capture-safe, and every replay must match an eager engine launch - on the same data bit for bit.""" - T_cap, B_cap, H, D = 768, 4, 2, 128 - dev = "cuda" - _seed() - scale = 1.0 / math.sqrt(D) - q = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - k = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - v = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - gate = torch.zeros(T_cap, H, D, dtype=torch.float32, device=dev) - beta = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - w = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - cu = torch.zeros(B_cap + 1, dtype=torch.int32, device=dev) - o_graph = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - fs_graph = torch.zeros(B_cap, H, D, D, dtype=torch.float32, device=dev) - o_eager = torch.zeros_like(o_graph) - fs_eager = torch.zeros_like(fs_graph) - - g, t = _build_gdn2_engine_graph(T_cap, H, D, B_cap, scale) - g.build() - ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=dev) - base = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta, t["w"]: w, t["cu"]: cu} - pack_graph = {**base, t["O"]: o_graph, t["fs"]: fs_graph} - pack_eager = {**base, t["O"]: o_eager, t["fs"]: fs_eager} - - def fill(seq_lens): - total = sum(seq_lens) - q[:total] = torch.randn(total, H, D, device=dev).bfloat16() * 0.5 - k[:total] = F.normalize(torch.randn(total, H, D, device=dev), dim=-1).bfloat16() - v[:total] = torch.randn(total, H, D, device=dev).bfloat16() * 0.5 - gate[:total] = torch.empty(total, H, D, device=dev).uniform_(0.5, 1.0).log() - beta[:total] = (torch.rand(total, H, D, device=dev).sigmoid() * 2.0).bfloat16() - w[:total] = torch.rand(total, H, D, device=dev).sigmoid().bfloat16() - bounds = [0] + list(accumulate(seq_lens)) - bounds += [bounds[-1]] * (B_cap + 1 - len(bounds)) - cu.copy_(torch.tensor(bounds, dtype=torch.int32)) - return total - - fill([256, 512, 0, 0]) - stream = torch.cuda.Stream() - handle = cudnn.create_handle() - cudnn.set_stream(handle, stream.cuda_stream) - with torch.cuda.stream(stream): - g.execute(pack_graph, ws, handle=handle) # warmup: compile + caches - torch.cuda.synchronize() - - cg = torch.cuda.CUDAGraph() - with torch.cuda.graph(cg, stream=stream): - g.execute(pack_graph, ws, handle=handle) - - eager_handle = cudnn.create_handle() - cudnn.set_stream(eager_handle, torch.cuda.current_stream().cuda_stream) - for seq_lens in ([256, 512, 0, 0], [100, 200, 56, 0], [768], [16, 0, 16, 736 - 32]): - total = fill(seq_lens) - cg.replay() - torch.cuda.synchronize() - g.execute(pack_eager, ws, handle=eager_handle) - torch.cuda.synchronize() - assert torch.equal(o_graph[:total], o_eager[:total]), f"graph o diverges from eager at {seq_lens}" - assert torch.equal(fs_graph, fs_eager), f"graph final_state diverges from eager at {seq_lens}" diff --git a/test/python/linear_attention/frost/test_gdn_bprop_kernel.py b/test/python/linear_attention/frost/test_gdn_bprop_kernel.py deleted file mode 100644 index 0594a4e10..000000000 --- a/test/python/linear_attention/frost/test_gdn_bprop_kernel.py +++ /dev/null @@ -1,964 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""FROST GDN backward tests (pygraph -GDN_BWD node on GdnFrostEngine) against fp64 autograd oracles.""" - -from __future__ import annotations - -import math -import random -from itertools import accumulate - -import pytest -import torch -import torch.nn.functional as F - -import cudnn # noqa: F401 (conftest extends cudnn.__path__ with the source tree) -from cudnn.linear_attention.frost import GdnFrostEngine - -from linear_attention.common import assert_bitwise_runs -from linear_attention.conftest import multidist_randu -from linear_attention.reference_gdn import gdn_reference, rms_ratio - -pytestmark = pytest.mark.L0 - -SEED = 42 - -BWD_TOL = 3e-2 - -HEAD_CONFIGS = [ - (1, 1, 1), - (4, 1, 1), - (3, 3, 3), - (6, 2, 2), - (1, 1, 2), - (2, 2, 4), - (16, 16, 32), - (16, 16, 64), -] - - -def _sm100_dsl_available() -> bool: - if not torch.cuda.is_available(): - return False - major, _minor = torch.cuda.get_device_capability() - if major != 10: - return False - try: - import cutlass.experimental.primitives # noqa: F401 — often a sys.modules alias, invisible to find_spec - except ImportError: - return False - return True - - -requires_runtime = pytest.mark.skipif( - not _sm100_dsl_available(), - reason="needs an SM100-class GPU and the Cutlass DSL", -) - - -def _seed(seed=SEED): - random.seed(seed) - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - - -def _cu(seq_lens, device="cuda"): - return torch.tensor([0] + list(accumulate(seq_lens)), dtype=torch.int32, device=device) - - -def _cu_h(seq_lens, device="cuda"): - counts = [max(sl - 1, 0) // 64 for sl in seq_lens] - return torch.tensor([0] + list(accumulate(counts)), dtype=torch.int32, device=device), sum(counts) - - -def _gen_case(seq_lens, HQ=2, HK=None, HV=None, head_size=128, dtype=torch.bfloat16, alpha_on=True, beta_on=True): - """THD inputs at native head counts; do/alpha/beta at HO = max(HQ, HV).""" - HK = HQ if HK is None else HK - HV = HQ if HV is None else HV - HO = max(HQ, HV) - _seed() - total = sum(seq_lens) - q = multidist_randu(total * HQ, head_size, device="cuda").reshape(total, HQ, head_size) - k = multidist_randu(total * HK, head_size, device="cuda").reshape(total, HK, head_size) - k = F.normalize(k, p=2.0, dim=-1) - v = multidist_randu(total * HV, head_size, device="cuda").reshape(total, HV, head_size) - do = multidist_randu(total * HO, head_size, device="cuda").reshape(total, HO, head_size) - alpha = torch.empty(total, HO, device="cuda").uniform_(0.1, 1.0) if alpha_on else torch.ones(total, HO, device="cuda") - beta = torch.rand(total, HO, device="cuda") if beta_on else torch.ones(total, HO, device="cuda") - return ( - q.to(dtype).contiguous(), - k.to(dtype).contiguous(), - v.to(dtype).contiguous(), - do.to(dtype).contiguous(), - alpha.contiguous(), - beta.contiguous(), - ) - - -def _build_bwd_graph(total, HQ, HV, D, num_seqs, scale, io_dt, *, h_shape=None, s0=False, dht=False): - - HO = max(HQ, HV) - g = cudnn.pygraph() - t = dict( - q=g.tensor([total, HQ, D], data_type=io_dt, name="q"), - k=g.tensor([total, HQ, D], data_type=io_dt, name="k"), - v=g.tensor([total, HV, D], data_type=io_dt, name="v"), - g=g.tensor([total, HO], data_type=cudnn.data_type.FLOAT, name="g"), - beta=g.tensor([total, HO], data_type=cudnn.data_type.FLOAT, name="beta"), - cu=g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens"), - dO=g.tensor([total, HO, D], data_type=io_dt, name="dO"), - ) - kwargs = {} - if h_shape is not None: - t["h"] = kwargs["h"] = g.tensor(list(h_shape), data_type=io_dt, name="h") - if s0: - t["s0"] = kwargs["initial_state"] = g.tensor([num_seqs, HO, D, D], data_type=cudnn.data_type.FLOAT, name="initial_state") - if dht: - t["dht"] = kwargs["d_final_state"] = g.tensor([num_seqs, HO, D, D], data_type=cudnn.data_type.FLOAT, name="d_final_state") - dQ_t, dK_t, dV_t, dG_t, dBeta_t, dS0_t = g.gdn_bwd( - q=t["q"], - k=t["k"], - v=t["v"], - g=t["g"], - beta=t["beta"], - cu_seqlens=t["cu"], - dO=t["dO"], - scale=float(scale), - name="gdn_bwd", - **kwargs, - ) - outs = [(dQ_t, io_dt), (dK_t, io_dt), (dV_t, io_dt), (dG_t, cudnn.data_type.FLOAT), (dBeta_t, cudnn.data_type.FLOAT)] - if s0: - outs.append((dS0_t, cudnn.data_type.FLOAT)) - for t_, dt in outs: - t_.set_output(True).set_data_type(dt) - t["grads"] = [o for o, _ in outs] - g.build() - return g, t - - -def _fwd_h(q, k, v, alpha, beta, scale, cu, seq_lens, s0=None, every_n=64): - """Forward through the engine with the per-chunk H output; returns h.""" - - device = q.device - total, HQ, HV, D = q.shape[0], q.shape[1], v.shape[1], q.shape[2] - HO = max(HQ, HV) - num_seqs = cu.shape[0] - 1 - io_dt = cudnn.data_type.BFLOAT16 if q.dtype == torch.bfloat16 else cudnn.data_type.HALF - g = cudnn.pygraph() - q_t = g.tensor([total, HQ, D], data_type=io_dt, name="q") - k_t = g.tensor([total, HQ, D], data_type=io_dt, name="k") - v_t = g.tensor([total, HV, D], data_type=io_dt, name="v") - g_t = g.tensor([total, HO], data_type=cudnn.data_type.FLOAT, name="g") - b_t = g.tensor([total, HO], data_type=cudnn.data_type.FLOAT, name="beta") - cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") - s0_t = g.tensor([num_seqs, HO, D, D], data_type=cudnn.data_type.FLOAT, name="initial_state") if s0 is not None else None - O_t, _fs_t, h_t = g.gdn( - q=q_t, - k=k_t, - v=v_t, - g=g_t, - beta=b_t, - cu_seqlens=cu_t, - initial_state=s0_t, - scale=float(scale), - checkpoint_every_n_tokens=every_n, - name="gdn", - ) - O_t.set_output(True).set_data_type(io_dt) - h_t.set_output(True).set_data_type(io_dt) - g.build() - total_h = sum(max(sl - 1, 0) // every_n for sl in seq_lens) - o = torch.empty(total, HO, D, dtype=q.dtype, device=device) - h = torch.full((max(total_h, 1), HO, D, D), float("nan"), dtype=q.dtype, device=device) - pack = {q_t: q, k_t: k, v_t: v, g_t: alpha.float().log().contiguous(), b_t: beta.float().contiguous(), cu_t: cu, O_t: o, h_t: h} - if s0 is not None: - pack[s0_t] = s0.contiguous() - g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) - torch.cuda.synchronize() - return h - - -def _run_bwd(q, k, v, alpha, beta, do, h, scale, cu, s0=None, dht=None): - """Engine-driven backward; ``h=None`` exercises the recompute path. - Returns (dq, dk, dv, dg, db) and appends ds0 when ``s0`` is given. - ``k`` with fewer heads than ``q`` is pre-broadcast to the node contract - and ``dk`` is group-reduced back to the caller's granularity.""" - device = q.device - total, HQ, HV, D = q.shape[0], q.shape[1], v.shape[1], q.shape[2] - HK = k.shape[1] - if HK != HQ: - k = k.repeat_interleave(HQ // HK, dim=1).contiguous() - HO = max(HQ, HV) - num_seqs = cu.shape[0] - 1 - io_dt = cudnn.data_type.BFLOAT16 if q.dtype == torch.bfloat16 else cudnn.data_type.HALF - g, t = _build_bwd_graph( - total, - HQ, - HV, - D, - num_seqs, - scale, - io_dt, - h_shape=None if h is None else h.shape, - s0=s0 is not None, - dht=dht is not None, - ) - dq = torch.full((total, HQ, D), float("nan"), dtype=q.dtype, device=device) - dk = torch.full((total, HQ, D), float("nan"), dtype=q.dtype, device=device) - dv = torch.full((total, HV, D), float("nan"), dtype=q.dtype, device=device) - dg = torch.full((total, HO), float("nan"), dtype=torch.float32, device=device) - db = torch.full((total, HO), float("nan"), dtype=torch.float32, device=device) - bufs = [dq, dk, dv, dg, db] - if s0 is not None: - bufs.append(torch.full((num_seqs, HO, D, D), float("nan"), dtype=torch.float32, device=device)) - pack = { - t["q"]: q, - t["k"]: k, - t["v"]: v, - t["g"]: alpha.float().log().contiguous(), - t["beta"]: beta.float().contiguous(), - t["cu"]: cu, - t["dO"]: do, - } - for ot, buf in zip(t["grads"], bufs): - pack[ot] = buf - if h is not None: - pack[t["h"]] = h - if s0 is not None: - pack[t["s0"]] = s0.contiguous() - if dht is not None: - pack[t["dht"]] = dht.contiguous() - g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) - torch.cuda.synchronize() - if HK != HQ: - bufs[1] = _reduce_ref(bufs[1], HK) - return tuple(bufs) - - -def _reference_grads(q, k, v, alpha, beta, do, scale, seq_lens): - """fp64 autograd oracle: grads of (o * dO).sum() at HO-head granularity. - - q/k/v leaves are pre-expanded to HO with the kernel's repeat_interleave - head mapping; the gate leaf is NATURAL-LOG (the kernel's dgate - convention) and alpha/beta are already HO-shaped.""" - HO = max(q.shape[1], v.shape[1]) - - def leaf(x): - r = HO // x.shape[1] - x = x.double().repeat_interleave(r, dim=1) if r > 1 else x.double() - return x.requires_grad_(True) - - qq, kk, vv = leaf(q), leaf(k), leaf(v) - gate = alpha.double().log().requires_grad_(True) - bb = beta.double().requires_grad_(True) - o, _fs = gdn_reference( - qq.unsqueeze(0), - kk.unsqueeze(0), - vv.unsqueeze(0), - gate.unsqueeze(0), - bb.unsqueeze(0), - scale=scale, - initial_state=None, - cu_seqlens=_cu(seq_lens, q.device), - ) - (o.squeeze(0) * do.double()).sum().backward() - return qq.grad, kk.grad, vv.grad, gate.grad, bb.grad - - -def _reduce_ref(g_ho, native_h): - ho = g_ho.shape[1] - if ho == native_h: - return g_ho - return g_ho.view(g_ho.shape[0], native_h, ho // native_h, *g_ho.shape[2:]).sum(2) - - -def _run_bprop_case(seq_lens, HQ=2, HK=None, HV=None, head_size=128, dtype=torch.bfloat16, scale=None, alpha_on=True, beta_on=True): - """Prefill H -> backward kernel -> compare all five gradients against the - fp64 autograd oracle at HO granularity.""" - scale = 1.0 / math.sqrt(head_size) if scale is None else scale - q, k, v, do, alpha, beta = _gen_case(seq_lens, HQ, HK, HV, head_size, dtype, alpha_on, beta_on) - cu = _cu(seq_lens) - dq, dk, dv, dg, db = _run_bwd(q, k, v, alpha, beta, do, None, scale, cu) - torch.cuda.synchronize() - dq_ref, dk_ref, dv_ref, dg_ref, db_ref = _reference_grads(q, k, v, alpha, beta, do, scale, seq_lens) - for name, got, ref, tol in ( - ("dq", dq, _reduce_ref(dq_ref, q.shape[1]), BWD_TOL), - ("dk", dk, _reduce_ref(dk_ref, k.shape[1]), BWD_TOL), - ("dv", dv, _reduce_ref(dv_ref, v.shape[1]), BWD_TOL), - ("dg", dg, dg_ref, BWD_TOL), - ("db", db, db_ref, BWD_TOL), - ): - assert torch.isfinite(got.float()).all(), f"non-finite values in {name}" - r = rms_ratio(got.float(), ref.float()) - assert r < tol, f"{name} rms ratio {r:.4g} >= {tol}" - - -# --------------------------------------------------------------------------- -# Comprehensive correctness (all five gradients per case) -# --------------------------------------------------------------------------- - - -@requires_runtime -@pytest.mark.parametrize("num_q_heads, num_k_heads, num_v_heads", HEAD_CONFIGS) -@pytest.mark.parametrize("seq_lens", [[256], [256, 256], [64, 128, 512]]) -@pytest.mark.parametrize( - "dtype", - [ - "float16", - "bfloat16", - ], -) -def test_bprop_kernel_basic(dtype, num_q_heads, num_k_heads, num_v_heads, seq_lens): - _run_bprop_case(seq_lens, num_q_heads, num_k_heads, num_v_heads, dtype=getattr(torch, dtype)) - - -@requires_runtime -@pytest.mark.parametrize("beta", [False, True]) -@pytest.mark.parametrize("alpha", [False, True]) -@pytest.mark.parametrize("scale", [1.0, "auto"]) -def test_bprop_kernel_gates_and_scale(scale, alpha, beta): - if not alpha and not beta: - pytest.skip("large diff due to output value amplitude explosion along token dimension") - scale = 1.0 / math.sqrt(128) if scale == "auto" else scale - _run_bprop_case([64, 128, 512], 3, 3, 3, scale=scale, alpha_on=alpha, beta_on=beta) - - -@requires_runtime -@pytest.mark.parametrize("num_q_heads, num_k_heads, num_v_heads", [(3, 3, 3), (4, 1, 1), (2, 2, 4)]) -@pytest.mark.parametrize("seq_lens", [[31], [251], [511, 501], [31, 63, 93, 123, 150, 500]]) -@pytest.mark.parametrize("dtype", ["bfloat16", "float16"]) -def test_bprop_kernel_nonfull(dtype, num_q_heads, num_k_heads, num_v_heads, seq_lens): - _run_bprop_case(seq_lens, num_q_heads, num_k_heads, num_v_heads, dtype=getattr(torch, dtype)) - - -@requires_runtime -@pytest.mark.parametrize("num_q_heads, num_k_heads, num_v_heads", [(1, 1, 1), (16, 16, 64)]) -@pytest.mark.parametrize("seq_len", [256, 255]) -def test_bprop_kernel_zero_length_sequence(num_q_heads, num_k_heads, num_v_heads, seq_len): - """A trailing zero-length sequence neither perturbs the gradients nor hangs.""" - head_size = 128 - scale = 1.0 / math.sqrt(head_size) - q, k, v, do, alpha, beta = _gen_case([seq_len], num_q_heads, num_k_heads, num_v_heads, head_size) - - def run(seq_lens): - cu = _cu(seq_lens) - out = _run_bwd(q, k, v, alpha, beta, do, None, scale, cu) - torch.cuda.synchronize() - return out - - ref = run([seq_len]) - got = run([seq_len, 0]) - for name, g, r in zip(("dq", "dk", "dv", "dg", "db"), got, ref): - torch.testing.assert_close(g, r, atol=2e-2, rtol=2e-2, msg=f"{name} perturbed by the zero-length sequence") - - -# --------------------------------------------------------------------------- -# FROST engine: GDN_BWD through the graph -# --------------------------------------------------------------------------- - - -def test_frost_gdn_bwd_engine_no_h(seq_lens=(128, 256), H=2, D=128): - """GDN_BWD without the h input: the engine reruns the forward with H - dumping and must match the explicit-h path bit for bit.""" - - _seed() - seq_lens = list(seq_lens) - q, k, v, do, alpha, beta = _gen_case(seq_lens, HQ=H, head_size=D) - device = q.device - scale = 1.0 / math.sqrt(D) - num_seqs = len(seq_lens) - total = q.shape[0] - cu = _cu(seq_lens, device) - cu_h_plain, total_h = _cu_h(seq_lens) - - h = _fwd_h(q, k, v, alpha, beta, scale, cu, seq_lens) - - results = {} - for with_h in (True, False): - g = cudnn.pygraph() - q_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="q") - k_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="k") - v_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="v") - g_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="g") - beta_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="beta") - cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") - do_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="dO") - kwargs = {} - if with_h: - kwargs["h"] = g.tensor([max(total_h, 1), H, D, D], data_type=cudnn.data_type.BFLOAT16, name="h") - dQ_t, dK_t, dV_t, dG_t, dBeta_t, _ = g.gdn_bwd( - q=q_t, - k=k_t, - v=v_t, - g=g_t, - beta=beta_t, - cu_seqlens=cu_t, - dO=do_t, - scale=scale, - name="gdn_bwd", - **kwargs, - ) - for t_, dt in ( - (dQ_t, cudnn.data_type.BFLOAT16), - (dK_t, cudnn.data_type.BFLOAT16), - (dV_t, cudnn.data_type.BFLOAT16), - (dG_t, cudnn.data_type.FLOAT), - (dBeta_t, cudnn.data_type.FLOAT), - ): - t_.set_output(True).set_data_type(dt) - g.build() - assert isinstance(g.selected_engine, GdnFrostEngine), f"frost engine must accept with_h={with_h}" - - dq = torch.empty(total, H, D, dtype=q.dtype, device=device) - dk = torch.empty(total, H, D, dtype=q.dtype, device=device) - dv = torch.empty(total, H, D, dtype=q.dtype, device=device) - dg = torch.empty(total, H, dtype=torch.float32, device=device) - dbeta = torch.empty(total, H, dtype=torch.float32, device=device) - pack = { - q_t: q, - k_t: k, - v_t: v, - g_t: alpha.log().contiguous(), - beta_t: beta.contiguous(), - cu_t: cu, - do_t: do, - dQ_t: dq, - dK_t: dk, - dV_t: dv, - dG_t: dg, - dBeta_t: dbeta, - } - if with_h: - pack[kwargs["h"]] = h - g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) - torch.cuda.synchronize() - results[with_h] = (dq, dk, dv, dg, dbeta) - - # both paths run the identical kernels on the identical log-decay inputs - for a, b, name in zip(results[True], results[False], ("dq", "dk", "dv", "dg", "dbeta")): - assert torch.equal(a, b), f"no-h recompute path diverges from explicit h on {name}" - - -def _reference_grads_s0(q, k, v, alpha, beta, do, scale, seq_lens, s0): - """fp64 autograd oracle including the initial-state leaf.""" - HO = max(q.shape[1], v.shape[1]) - - def leaf(x): - r = HO // x.shape[1] - x = x.double().repeat_interleave(r, dim=1) if r > 1 else x.double() - return x.requires_grad_(True) - - qq, kk, vv = leaf(q), leaf(k), leaf(v) - gate = alpha.double().log().requires_grad_(True) - bb = beta.double().requires_grad_(True) - ss0 = s0.double().requires_grad_(True) - o, _fs = gdn_reference( - qq.unsqueeze(0), - kk.unsqueeze(0), - vv.unsqueeze(0), - gate.unsqueeze(0), - bb.unsqueeze(0), - scale=scale, - initial_state=ss0, - cu_seqlens=_cu(seq_lens, q.device), - ) - (o.squeeze(0) * do.double()).sum().backward() - return qq.grad, kk.grad, vv.grad, gate.grad, bb.grad, ss0.grad - - -@requires_runtime -@pytest.mark.parametrize("heads", [(2, 2, 2), (2, 2, 4)]) -@pytest.mark.parametrize("seq_lens", [[256], [128, 512]]) -def test_bprop_d_initial_state(seq_lens, heads): - """dL/dS0 vs the fp64 autograd oracle with a nonzero initial state. - - The backward takes the plain per-chunk h plus the io-downcast S0 - (``initial_state``, read through its own descriptor set for chunk 0).""" - _seed() - q, k, v, do, alpha, beta = _gen_case(seq_lens, *heads) - HO = max(q.shape[1], v.shape[1]) - D = q.shape[2] - device = q.device - scale = 1.0 / (D**0.5) - num_seqs = len(seq_lens) - s0 = (torch.randn(num_seqs, HO, D, D, device=device, dtype=torch.float32) * 0.05).contiguous() - - cu = _cu(seq_lens, device) - dq, dk, dv, dg, db, ds0 = _run_bwd(q, k, v, alpha, beta, do, None, scale, cu, s0=s0) - - rq, rk, rv, rg, rb, rs0 = _reference_grads_s0(q, k, v, alpha, beta, do, scale, seq_lens, s0) - - def rms(a, b): - return (a.double() - b).pow(2).mean().sqrt().item() - - assert not ds0.isnan().any(), "d_initial_state has unwritten slots" - rq, rk, rv = _reduce_ref(rq, q.shape[1]), _reduce_ref(rk, k.shape[1]), _reduce_ref(rv, v.shape[1]) - assert rms(dq, rq) < 6e-2 and rms(dk, rk) < 6e-2 and rms(dv, rv) < 6e-2 - assert rms(dg, rg) < 6e-2 and rms(db, rb) < 6e-2 - assert rms(ds0, rs0) < 6e-2, f"ds0 rms {rms(ds0, rs0)}" - - -def _reference_grads_dht(q, k, v, alpha, beta, do, scale, seq_lens, dht): - """fp64 autograd oracle with a final-state-gradient loss term (dht).""" - HO = max(q.shape[1], v.shape[1]) - - def leaf(x): - r = HO // x.shape[1] - x = x.double().repeat_interleave(r, dim=1) if r > 1 else x.double() - return x.requires_grad_(True) - - qq, kk, vv = leaf(q), leaf(k), leaf(v) - gate = alpha.double().log().requires_grad_(True) - bb = beta.double().requires_grad_(True) - o, fs = gdn_reference( - qq.unsqueeze(0), - kk.unsqueeze(0), - vv.unsqueeze(0), - gate.unsqueeze(0), - bb.unsqueeze(0), - scale=scale, - initial_state=None, - cu_seqlens=_cu(seq_lens, q.device), - ) - ((o.squeeze(0) * do.double()).sum() + (fs * dht.double()).sum()).backward() - return qq.grad, kk.grad, vv.grad, gate.grad, bb.grad - - -@requires_runtime -@pytest.mark.parametrize("heads", [(2, 2, 2), (2, 2, 4)]) -@pytest.mark.parametrize("seq_lens", [[64], [256], [128, 512]]) -def test_bprop_d_final_state(seq_lens, heads): - """d_final_state seeds the backward dH at the first processed chunk; - every gradient picks up the propagated term.""" - _seed() - q, k, v, do, alpha, beta = _gen_case(seq_lens, *heads) - HO = max(q.shape[1], v.shape[1]) - D = q.shape[2] - device = q.device - scale = 1.0 / (D**0.5) - num_seqs = len(seq_lens) - dht = (torch.randn(num_seqs, HO, D, D, device=device, dtype=torch.float32) * 0.05).contiguous() - - cu = _cu(seq_lens, device) - dq, dk, dv, dg, db = _run_bwd(q, k, v, alpha, beta, do, None, scale, cu, dht=dht) - - rq, rk, rv, rg, rb = _reference_grads_dht(q, k, v, alpha, beta, do, scale, seq_lens, dht) - - def rms(a, b): - return (a.double() - b).pow(2).mean().sqrt().item() - - rq, rk, rv = _reduce_ref(rq, q.shape[1]), _reduce_ref(rk, k.shape[1]), _reduce_ref(rv, v.shape[1]) - for name, got, ref in (("dq", dq, rq), ("dk", dk, rk), ("dv", dv, rv), ("dg", dg, rg), ("db", db, rb)): - assert torch.isfinite(got.float()).all(), f"non-finite values in {name}" - assert rms(got, ref) < 6e-2, f"{name} rms {rms(got, ref):.4g}" - - -@requires_runtime -def test_frost_gdn_bwd_engine_initial_state(seq_lens=(128, 256), H=2, D=128): - """GDN_BWD through the pygraph engine path with initial_state: the engine - extends h with the per-sequence S0 entries; the node's state ports are - K-major [N, H, K, V] (the engine converts to the kernel orientation).""" - - _seed() - seq_lens = list(seq_lens) - q, k, v, do, alpha, beta = _gen_case(seq_lens, HQ=H, head_size=D) - device = q.device - scale = 1.0 / math.sqrt(D) - num_seqs = len(seq_lens) - total = q.shape[0] - s0 = (torch.randn(num_seqs, H, D, D, device=device, dtype=torch.float32) * 0.05).contiguous() - - cu = _cu(seq_lens, device) - _cu_h_plain, total_h = _cu_h(seq_lens) - h = _fwd_h(q, k, v, alpha, beta, scale, cu, seq_lens, s0=s0) - - g = cudnn.pygraph() - q_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="q") - k_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="k") - v_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="v") - g_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="g") - beta_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="beta") - cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") - do_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="dO") - h_t = g.tensor([max(total_h, 1), H, D, D], data_type=cudnn.data_type.BFLOAT16, name="h") - s0_t = g.tensor([num_seqs, H, D, D], data_type=cudnn.data_type.FLOAT, name="initial_state") - dQ_t, dK_t, dV_t, dG_t, dBeta_t, dS0_t = g.gdn_bwd( - q=q_t, - k=k_t, - v=v_t, - g=g_t, - beta=beta_t, - cu_seqlens=cu_t, - dO=do_t, - h=h_t, - initial_state=s0_t, - scale=scale, - name="gdn_bwd", - ) - for t_, dt in ( - (dQ_t, cudnn.data_type.BFLOAT16), - (dK_t, cudnn.data_type.BFLOAT16), - (dV_t, cudnn.data_type.BFLOAT16), - (dG_t, cudnn.data_type.FLOAT), - (dBeta_t, cudnn.data_type.FLOAT), - (dS0_t, cudnn.data_type.FLOAT), - ): - t_.set_output(True).set_data_type(dt) - - g.build() - assert isinstance(g.selected_engine, GdnFrostEngine), "frost engine must accept initial_state" - - dq = torch.empty(total, H, D, dtype=q.dtype, device=device) - dk = torch.empty(total, H, D, dtype=q.dtype, device=device) - dv = torch.empty(total, H, D, dtype=q.dtype, device=device) - dg = torch.empty(total, H, dtype=torch.float32, device=device) - dbeta = torch.empty(total, H, dtype=torch.float32, device=device) - ds0 = torch.full((num_seqs, H, D, D), float("nan"), dtype=torch.float32, device=device) - wsb = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device) - g.execute( - { - q_t: q, - k_t: k, - v_t: v, - g_t: alpha.log().contiguous(), - beta_t: beta.contiguous(), - cu_t: cu, - do_t: do, - h_t: h, - s0_t: s0, - dQ_t: dq, - dK_t: dk, - dV_t: dv, - dG_t: dg, - dBeta_t: dbeta, - dS0_t: ds0, - }, - wsb, - ) - torch.cuda.synchronize() - - rq, rk, rv, rg, rb, rs0 = _reference_grads_s0(q, k, v, alpha, beta, do, scale, seq_lens, s0) - - def rms(a, b): - return (a.double() - b).pow(2).mean().sqrt().item() - - assert not ds0.isnan().any(), "d_initial_state has unwritten slots" - assert rms(dq, rq) < BWD_TOL and rms(dk, rk) < BWD_TOL and rms(dv, rv) < BWD_TOL - assert rms(dg, rg) < BWD_TOL and rms(dbeta, rb) < BWD_TOL - assert rms(ds0, rs0) < BWD_TOL, f"ds0 rms {rms(ds0, rs0)}" - - -# --------------------------------------------------------------------------- -# GVA: head-group reduction back to native q/k heads -# --------------------------------------------------------------------------- - - -def _reference_grads_native(q, k, v, alpha, beta, do, scale, seq_lens): - """fp64 autograd oracle at NATIVE head counts: gdn_reference expands the - head groups internally, so autograd sums each group's gradient.""" - qq = q.double().requires_grad_(True) - kk = k.double().requires_grad_(True) - vv = v.double().requires_grad_(True) - gate = alpha.double().log().requires_grad_(True) - bb = beta.double().requires_grad_(True) - o, _fs = gdn_reference( - qq.unsqueeze(0), - kk.unsqueeze(0), - vv.unsqueeze(0), - gate.unsqueeze(0), - bb.unsqueeze(0), - scale=scale, - initial_state=None, - cu_seqlens=_cu(seq_lens, q.device), - ) - (o.squeeze(0) * do.double()).sum().backward() - return qq.grad, kk.grad, vv.grad, gate.grad, bb.grad - - -@requires_runtime -@pytest.mark.parametrize("num_q_heads, num_v_heads", [(2, 4), (16, 64), (3, 6)]) -@pytest.mark.parametrize("seq_lens", [[256], [128, 512], [77, 178], [1]]) -def test_bprop_kernel_gva_native_heads(num_q_heads, num_v_heads, seq_lens): - """GVA through the node surface: the engine returns native-head grads.""" - scale = 1.0 / math.sqrt(128) - q, k, v, do, alpha, beta = _gen_case(seq_lens, num_q_heads, num_q_heads, num_v_heads) - cu = _cu(seq_lens) - dq, dk, dv, dg, db = _run_bwd(q, k, v, alpha, beta, do, None, scale, cu) - refs = _reference_grads_native(q, k, v, alpha, beta, do, scale, seq_lens) - for name, got, ref in (("dq", dq, refs[0]), ("dk", dk, refs[1]), ("dv", dv, refs[2]), ("dg", dg, refs[3]), ("db", db, refs[4])): - assert torch.isfinite(got.float()).all(), f"non-finite values in {name}" - r = rms_ratio(got.float(), ref.float()) - assert r < BWD_TOL, f"{name} rms ratio {r:.4g} >= {BWD_TOL}" - - -@requires_runtime -@pytest.mark.parametrize("num_q_heads, num_v_heads", [(2, 4), (16, 64)]) -def test_frost_gdn_bwd_engine_gva(num_q_heads, num_v_heads, seq_lens=(128, 256), D=128): - """GDN_BWD through the FROST engine with grouped value heads: dQ/dK come - back at the node's native q/k head counts via the head-group reduction.""" - - _seed() - seq_lens = list(seq_lens) - HQ, HV = num_q_heads, num_v_heads - q, k, v, do, alpha, beta = _gen_case(seq_lens, HQ=HQ, HV=HV, head_size=D) - device = q.device - scale = 1.0 / math.sqrt(D) - num_seqs = len(seq_lens) - total = q.shape[0] - cu = _cu(seq_lens, device) - cu_h, total_h = _cu_h(seq_lens) - h = _fwd_h(q, k, v, alpha, beta, scale, cu, seq_lens) - - g = cudnn.pygraph() - q_t = g.tensor([total, HQ, D], data_type=cudnn.data_type.BFLOAT16, name="q") - k_t = g.tensor([total, HQ, D], data_type=cudnn.data_type.BFLOAT16, name="k") - v_t = g.tensor([total, HV, D], data_type=cudnn.data_type.BFLOAT16, name="v") - g_t = g.tensor([total, HV], data_type=cudnn.data_type.FLOAT, name="g") - beta_t = g.tensor([total, HV], data_type=cudnn.data_type.FLOAT, name="beta") - cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") - do_t = g.tensor([total, HV, D], data_type=cudnn.data_type.BFLOAT16, name="dO") - h_t = g.tensor([max(total_h, 1), HV, D, D], data_type=cudnn.data_type.BFLOAT16, name="h") - dQ_t, dK_t, dV_t, dG_t, dBeta_t, _ = g.gdn_bwd( - q=q_t, - k=k_t, - v=v_t, - g=g_t, - beta=beta_t, - cu_seqlens=cu_t, - dO=do_t, - h=h_t, - scale=scale, - name="gdn_bwd", - ) - for t_, dt in ( - (dQ_t, cudnn.data_type.BFLOAT16), - (dK_t, cudnn.data_type.BFLOAT16), - (dV_t, cudnn.data_type.BFLOAT16), - (dG_t, cudnn.data_type.FLOAT), - (dBeta_t, cudnn.data_type.FLOAT), - ): - t_.set_output(True).set_data_type(dt) - g.build() - assert isinstance(g.selected_engine, GdnFrostEngine), "frost engine must accept GVA" - - dq = torch.full((total, HQ, D), float("nan"), dtype=q.dtype, device=device) - dk = torch.full((total, HQ, D), float("nan"), dtype=q.dtype, device=device) - dv = torch.full((total, HV, D), float("nan"), dtype=q.dtype, device=device) - dg = torch.full((total, HV), float("nan"), dtype=torch.float32, device=device) - dbeta = torch.full((total, HV), float("nan"), dtype=torch.float32, device=device) - g.execute( - { - q_t: q, - k_t: k, - v_t: v, - g_t: alpha.log().contiguous(), - beta_t: beta.contiguous(), - cu_t: cu, - do_t: do, - h_t: h, - dQ_t: dq, - dK_t: dk, - dV_t: dv, - dG_t: dg, - dBeta_t: dbeta, - }, - torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device), - ) - torch.cuda.synchronize() - - refs = _reference_grads_native(q, k, v, alpha, beta, do, scale, seq_lens) - for name, got, ref in (("dq", dq, refs[0]), ("dk", dk, refs[1]), ("dv", dv, refs[2]), ("dg", dg, refs[3]), ("db", dbeta, refs[4])): - assert torch.isfinite(got.float()).all(), f"non-finite values in {name}" - r = rms_ratio(got.float(), ref.float()) - assert r < BWD_TOL, f"{name} rms ratio {r:.4g} >= {BWD_TOL}" - - -@requires_runtime -@pytest.mark.parametrize("num_q_heads, num_k_heads, num_v_heads", [(4, 1, 1), (6, 2, 2)]) -@pytest.mark.parametrize("seq_lens", [[256], [128, 512]]) -def test_bprop_kernel_gqa_native_heads(num_q_heads, num_k_heads, num_v_heads, seq_lens): - """GQA through the node surface (q/k at HQ heads, v at native HV): the - engine returns native-head grads; dk/dg/db reduce to the HK/HV oracle.""" - HQ, HK, HV = num_q_heads, num_k_heads, num_v_heads - HO = HQ - D = 128 - scale = 1.0 / math.sqrt(D) - _seed() - total = sum(seq_lens) - q = multidist_randu(total * HQ, D, device="cuda").reshape(total, HQ, D) - k = F.normalize(multidist_randu(total * HK, D, device="cuda").reshape(total, HK, D), p=2.0, dim=-1) - v = multidist_randu(total * HV, D, device="cuda").reshape(total, HV, D) - do = multidist_randu(total * HO, D, device="cuda").reshape(total, HO, D) - q, k, v, do = (t.bfloat16().contiguous() for t in (q, k, v, do)) - alpha_hv = torch.empty(total, HV, device="cuda").uniform_(0.1, 1.0) - beta_hv = torch.rand(total, HV, device="cuda") - alpha = alpha_hv.repeat_interleave(HO // HV, dim=1).contiguous() - beta = beta_hv.repeat_interleave(HO // HV, dim=1).contiguous() - cu = _cu(seq_lens) - k_hq = k.repeat_interleave(HQ // HK, dim=1).contiguous() - dq, dk_hq, dv, dg_ho, db_ho = _run_bwd(q, k_hq, v, alpha, beta, do, None, scale, cu) - refs = _reference_grads_native(q, k, v, alpha_hv, beta_hv, do, scale, seq_lens) - for name, got, ref in ( - ("dq", dq, refs[0]), - ("dk", _reduce_ref(dk_hq, HK), refs[1]), - ("dv", dv, refs[2]), - ("dg", _reduce_ref(dg_ho, HV), refs[3]), - ("db", _reduce_ref(db_ho, HV), refs[4]), - ): - assert torch.isfinite(got.float()).all(), f"non-finite values in {name}" - r = rms_ratio(got.float(), ref.float()) - assert r < BWD_TOL, f"{name} rms ratio {r:.4g} >= {BWD_TOL}" - - -# --------------------------------------------------------------------------- -# split-K: strong decay on a long-sequence pack, so the backward's table -# actually cuts (see the partition-table unit tests in -# test_gdn_prefill_kernel.py) -# --------------------------------------------------------------------------- - - -@requires_runtime -@pytest.mark.parametrize("with_s0", [False, True]) -@pytest.mark.parametrize("with_dht", [False, True]) -def test_bprop_split_strong_decay(with_s0, with_dht): - """Strong decay saturates the scan's warmup threshold on the 2048-token - sequence, so the split table cuts it into several work items; all - gradients against the fp64 autograd oracle.""" - _seed() - seq_lens = [100, 2048, 0, 517] - H = 2 - total = sum(seq_lens) - q = torch.randn(total, H, 128, dtype=torch.bfloat16, device="cuda") * 0.5 - k = F.normalize(torch.randn(total, H, 128, device="cuda"), dim=-1).bfloat16() - v = torch.randn(total, H, 128, dtype=torch.bfloat16, device="cuda") * 0.5 - do = torch.randn(total, H, 128, dtype=torch.bfloat16, device="cuda") * 0.5 - alpha = (torch.rand(total, H, device="cuda") * 0.9 + 0.05).float() - beta = torch.rand(total, H, dtype=torch.float32, device="cuda") - cu = _cu(seq_lens) - B = len(seq_lens) - scale = 1.0 / math.sqrt(128) - s0 = (torch.randn(B, H, 128, 128, dtype=torch.float32, device="cuda") * 0.05) if with_s0 else None - dht = (torch.randn(B, H, 128, 128, dtype=torch.float32, device="cuda") * 0.05) if with_dht else None - - got = _run_bwd(q, k, v, alpha, beta, do, None, scale, cu, s0=s0, dht=dht) - - qq = q.double().requires_grad_(True) - kk = k.double().requires_grad_(True) - vv = v.double().requires_grad_(True) - gl = alpha.double().log().requires_grad_(True) - bb = beta.double().requires_grad_(True) - ss0 = s0.double().requires_grad_(True) if with_s0 else None - o, fs = gdn_reference( - qq.unsqueeze(0), - kk.unsqueeze(0), - vv.unsqueeze(0), - gl.unsqueeze(0), - bb.unsqueeze(0), - scale=scale, - initial_state=ss0, - cu_seqlens=cu, - ) - loss = (o.squeeze(0) * do.double()).sum() - if with_dht: - loss = loss + (fs * dht.double()).sum() - loss.backward() - ref = (qq.grad, kk.grad, vv.grad, gl.grad, bb.grad) + ((ss0.grad,) if with_s0 else ()) - - names = ("dq", "dk", "dv", "dg", "dbeta") + (("ds0",) if with_s0 else ()) - for name, a_, b_ in zip(names, got, ref): - assert not torch.isnan(a_.float()).any(), f"unwritten {name}" - assert rms_ratio(a_.float(), b_.float()) < 6e-2, name - if with_s0: - # a zero-length sequence passes the gradient through: ds0 = dht or 0 - zi = seq_lens.index(0) - expected = dht[zi] if with_dht else torch.zeros_like(got[5][zi]) - torch.testing.assert_close(got[5][zi], expected, atol=0, rtol=0) - - -# --------------------------------------------------------------------------- -# Determinism stress: bitwise repeat runs (fixed forward, repeated backward) -# --------------------------------------------------------------------------- - - -@requires_runtime -def test_gdn_bprop_determinism(): - seq_lens = [497, 16, 480, 256] - q, k, v, do, alpha, beta = _gen_case(seq_lens, HQ=2, HK=2, HV=2) - cu = _cu(seq_lens) - scale = 1.0 / math.sqrt(128) - assert_bitwise_runs(lambda: _run_bwd(q, k, v, alpha, beta, do, None, scale, cu), label="gdn_bwd") - - -# --------------------------------------------------------------------------- -# CUDA graph capture/replay across dynamic shapes (fixed SM-count grid) -# --------------------------------------------------------------------------- - - -@requires_runtime -def test_gdn_bprop_cuda_graph_replay(): - """Capture the ENGINE backward once (h=None: the forward-state recompute - happens INSIDE the capture), replay across CHANGED effective shapes; - every replay must match an eager engine backward bit for bit.""" - T_cap, B_cap, H, D = 768, 4, 2, 128 - dev = "cuda" - _seed() - scale = 1.0 / math.sqrt(D) - q = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - k = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - v = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - do = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - log_g = torch.zeros(T_cap, H, dtype=torch.float32, device=dev) - beta = torch.zeros(T_cap, H, dtype=torch.float32, device=dev) - cu = torch.zeros(B_cap + 1, dtype=torch.int32, device=dev) - - g, t = _build_bwd_graph(T_cap, H, H, D, B_cap, scale, cudnn.data_type.BFLOAT16) - ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=dev) - base = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: log_g, t["beta"]: beta, t["cu"]: cu, t["dO"]: do} - - def bufs(): - return [torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) for _ in range(3)] + [ - torch.zeros(T_cap, H, dtype=torch.float32, device=dev) for _ in range(2) - ] - - grads_graph, grads_eager = bufs(), bufs() - pack_graph = {**base, **dict(zip(t["grads"], grads_graph))} - pack_eager = {**base, **dict(zip(t["grads"], grads_eager))} - - def fill(seq_lens): - total = sum(seq_lens) - q[:total] = torch.randn(total, H, D, device=dev).bfloat16() * 0.5 - k[:total] = F.normalize(torch.randn(total, H, D, device=dev), dim=-1).bfloat16() - v[:total] = torch.randn(total, H, D, device=dev).bfloat16() * 0.5 - do[:total] = torch.randn(total, H, D, device=dev).bfloat16() * 0.5 - log_g[:total] = torch.empty(total, H, device=dev).uniform_(0.1, 1.0).log() - beta[:total] = torch.rand(total, H, device=dev) - bounds = [0] + list(accumulate(seq_lens)) - bounds += [bounds[-1]] * (B_cap + 1 - len(bounds)) - cu.copy_(torch.tensor(bounds, dtype=torch.int32)) - return total - - fill([256, 448, 0, 0]) - stream = torch.cuda.Stream() - handle = cudnn.create_handle() - cudnn.set_stream(handle, stream.cuda_stream) - with torch.cuda.stream(stream): - g.execute(pack_graph, ws, handle=handle) # warmup: compile + caches - torch.cuda.synchronize() - - cg = torch.cuda.CUDAGraph() - with torch.cuda.graph(cg, stream=stream): - g.execute(pack_graph, ws, handle=handle) - - eager_handle = cudnn.create_handle() - cudnn.set_stream(eager_handle, torch.cuda.current_stream().cuda_stream) - for seq_lens in ([256, 448, 0, 0], [100, 200, 56, 0], [768], [64, 0, 64, 512]): - total = fill(seq_lens) - cg.replay() - torch.cuda.synchronize() - g.execute(pack_eager, ws, handle=eager_handle) - torch.cuda.synchronize() - for name, a_, b_ in zip(("dq", "dk", "dv", "dg", "db"), grads_graph, grads_eager): - assert torch.equal(a_[:total], b_[:total]), f"graph {name} diverges from eager at {seq_lens}" diff --git a/test/python/linear_attention/frost/test_gdn_prefill_kernel.py b/test/python/linear_attention/frost/test_gdn_prefill_kernel.py deleted file mode 100644 index 9d9f24752..000000000 --- a/test/python/linear_attention/frost/test_gdn_prefill_kernel.py +++ /dev/null @@ -1,644 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""FROST GDN prefill tests (pygraph GDN node on GdnFrostEngine) against the -fp64 recurrent reference.""" - -from __future__ import annotations - -import math -import random -from itertools import accumulate - -import pytest -import torch -import torch.nn.functional as F - -import cudnn # noqa: F401 (conftest extends cudnn.__path__ with the source tree) - -from linear_attention.common import FWD_TOL, STATE_TOL, assert_bitwise_runs, assert_concurrent_stream_runs, assert_engine_declines -from linear_attention.conftest import multidist_randu -from linear_attention.reference_gdn import gdn_reference, rms_ratio - -pytestmark = pytest.mark.L0 - -SEED = 42 - - -def _sm100_dsl_available() -> bool: - if not torch.cuda.is_available(): - return False - major, _minor = torch.cuda.get_device_capability() - if major != 10: - return False - try: - import cutlass.experimental.primitives # noqa: F401 — often a sys.modules alias, invisible to find_spec - except ImportError: - return False - return True - - -requires_runtime = pytest.mark.skipif( - not _sm100_dsl_available(), - reason="needs an SM100-class GPU and the Cutlass DSL", -) - - -def _seed(seed=SEED): - random.seed(seed) - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - - -def _cu(seq_lens, device="cuda"): - return torch.tensor([0] + list(accumulate(seq_lens)), dtype=torch.int32, device=device) - - -def _gen_thd(seq_lens, num_q_heads, num_k_heads, num_v_heads, head_size, dtype): - total = sum(seq_lens) - q = multidist_randu(total * num_q_heads, head_size, device="cuda").reshape(total, num_q_heads, head_size) - k = multidist_randu(total * num_k_heads, head_size, device="cuda").reshape(total, num_k_heads, head_size) - k = F.normalize(k, p=2.0, dim=-1) - v = multidist_randu(total * num_v_heads, head_size, device="cuda").reshape(total, num_v_heads, head_size) - return q.to(dtype).contiguous(), k.to(dtype).contiguous(), v.to(dtype).contiguous() - - -def _build_gdn_graph( - total, HQ, HV, head_size, num_seqs, scale, io_dt, *, output_final_state, s0_shape=None, s0_dt=None, fs_dt=None, checkpoint_every_n_tokens=0 -): - - HO = max(HQ, HV) - g = cudnn.pygraph() - t = dict( - q=g.tensor([total, HQ, head_size], data_type=io_dt, name="q"), - k=g.tensor([total, HQ, head_size], data_type=io_dt, name="k"), - v=g.tensor([total, HV, head_size], data_type=io_dt, name="v"), - g=g.tensor([total, HO], data_type=cudnn.data_type.FLOAT, name="g"), - beta=g.tensor([total, HO], data_type=cudnn.data_type.FLOAT, name="beta"), - cu=g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens"), - ) - if s0_shape is not None: - t["s0"] = g.tensor(list(s0_shape), data_type=s0_dt, name="initial_state") - O_t, fs_t, h_t = g.gdn( - q=t["q"], - k=t["k"], - v=t["v"], - g=t["g"], - beta=t["beta"], - cu_seqlens=t["cu"], - initial_state=t.get("s0"), - scale=float(scale), - output_final_state=output_final_state, - checkpoint_every_n_tokens=checkpoint_every_n_tokens, - name="gdn", - ) - if h_t is not None: - h_t.set_output(True).set_data_type(io_dt) - t["H"] = h_t - O_t.set_output(True).set_data_type(io_dt) - t["O"] = O_t - if output_final_state: - fs_t.set_output(True).set_data_type(fs_dt) - t["fs"] = fs_t - g.build() - return g, t - - -def chunk_gated_delta_rule( - q, - k, - v, - alpha, - beta, - scale, - initial_state, - output_final_state, - cu_seqlens, - output=None, - output_state=None, -): - """Torch adapter over the graph. ``alpha`` / ``beta`` are the - raw linear gates (``None`` -> ones); the node takes natural-log decay.""" - device = q.device - total, HQ = q.shape[0], q.shape[1] - HV = v.shape[1] - HO = max(HQ, HV) - head_size = q.shape[2] - num_seqs = cu_seqlens.shape[0] - 1 - io_dt = cudnn.data_type.BFLOAT16 if q.dtype == torch.bfloat16 else cudnn.data_type.HALF - - gate = (alpha if alpha is not None else torch.ones(total, HO, device=device)).float() - log_g = gate.log().contiguous() - beta_f = (beta if beta is not None else torch.ones(total, HO, device=device)).float().contiguous() - cu = cu_seqlens.to(torch.int32).contiguous() - if output is None: - output = torch.empty(total, HO, head_size, dtype=q.dtype, device=device) - if output_final_state and output_state is None: - output_state = torch.empty(num_seqs, HO, head_size, head_size, dtype=torch.float32, device=device) - if not output_final_state: - output_state = None - - s0_dt = None - if initial_state is not None: - s0_dt = cudnn.data_type.BFLOAT16 if initial_state.dtype == torch.bfloat16 else cudnn.data_type.FLOAT - fs_dt = None - if output_final_state: - fs_dt = cudnn.data_type.BFLOAT16 if output_state.dtype == torch.bfloat16 else cudnn.data_type.FLOAT - g, t = _build_gdn_graph( - total, - HQ, - HV, - head_size, - num_seqs, - scale, - io_dt, - output_final_state=output_final_state, - s0_shape=None if initial_state is None else initial_state.shape, - s0_dt=s0_dt, - fs_dt=fs_dt, - ) - pack = {t["q"]: q.contiguous(), t["k"]: k.contiguous(), t["v"]: v.contiguous(), t["g"]: log_g, t["beta"]: beta_f, t["cu"]: cu, t["O"]: output} - if initial_state is not None: - pack[t["s0"]] = initial_state.contiguous() - if output_final_state: - pack[t["fs"]] = output_state - ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device) - g.execute(pack, ws) - return output, output_state - - -def _reference(q, k, v, alpha, beta, scale, seq_lens, initial_state=None): - """fp64 oracle on THD inputs; gates are raw linear (``None`` -> ones).""" - total = q.shape[0] - HO = max(q.shape[1], v.shape[1]) - device = q.device - gate = alpha if alpha is not None else torch.ones(total, HO, device=device) - beta_f = beta if beta is not None else torch.ones(total, HO, device=device) - with torch.no_grad(): - o, fs = gdn_reference( - q.unsqueeze(0), - k.unsqueeze(0), - v.unsqueeze(0), - gate.log().unsqueeze(0), - beta_f.unsqueeze(0), - scale=scale, - initial_state=initial_state, - cu_seqlens=_cu(seq_lens, device), - ) - return o.squeeze(0), fs - - -HEAD_CONFIGS = [ - (1, 1, 1), - (4, 1, 1), - (3, 3, 3), - (6, 2, 2), - (1, 1, 2), - (2, 2, 4), - (16, 16, 32), - (16, 16, 64), -] - - -def _run_prefill_case( - dtype, - num_q_heads, - num_k_heads, - num_v_heads, - seq_lens, - scale, - alpha, - beta, - head_size=128, -): - _seed() - total = sum(seq_lens) - num_seqs = len(seq_lens) - HO = max(num_q_heads, num_v_heads) - - dtype = getattr(torch, dtype) - q, k, v = _gen_thd(seq_lens, num_q_heads, num_k_heads, num_v_heads, head_size, dtype) - alpha_t = torch.empty(total, HO, device="cuda").uniform_(0.1, 1.0) if alpha else None - beta_t = torch.rand(total, HO, device="cuda") if beta else None - - our_o = torch.full((total, HO, head_size), float("nan"), dtype=dtype, device="cuda") - our_state = torch.full((num_seqs, HO, head_size, head_size), float("nan"), dtype=torch.float32, device="cuda") - chunk_gated_delta_rule(q, k, v, alpha_t, beta_t, scale, None, True, _cu(seq_lens), output=our_o, output_state=our_state) - torch.cuda.synchronize() - - ref_o, ref_state = _reference(q, k, v, alpha_t, beta_t, scale, seq_lens) - assert rms_ratio(our_o, ref_o) < FWD_TOL[dtype] - # The kernel state is K-major [N,HO,K,V], matching the reference. - assert rms_ratio(our_state, ref_state) < STATE_TOL[dtype] - - -@requires_runtime -@pytest.mark.parametrize("num_q_heads, num_k_heads, num_v_heads", HEAD_CONFIGS) -@pytest.mark.parametrize("seq_lens", [[256], [256, 256], [64, 128, 512]]) -@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) -def test_prefill_kernel_basic(dtype, num_q_heads, num_k_heads, num_v_heads, seq_lens): - _run_prefill_case( - dtype, - num_q_heads, - num_k_heads, - num_v_heads, - seq_lens, - scale=1.0 / math.sqrt(128), - alpha=True, - beta=True, - ) - - -@requires_runtime -@pytest.mark.parametrize("beta", [False, True]) -@pytest.mark.parametrize("alpha", [False, True]) -@pytest.mark.parametrize("scale", [1.0, "auto"]) -def test_prefill_kernel_gates_and_scale(scale, alpha, beta): - if not alpha and not beta: - pytest.skip("large diff due to output value amplitude explosion along token dimension") - scale = 1.0 / math.sqrt(128) if scale == "auto" else scale - _run_prefill_case("bfloat16", 3, 3, 3, [64, 128, 512], scale=scale, alpha=alpha, beta=beta) - - -@requires_runtime -@pytest.mark.parametrize("num_q_heads, num_k_heads, num_v_heads", [(3, 3, 3), (4, 1, 1), (2, 2, 4)]) -@pytest.mark.parametrize("seq_lens", [[31], [251], [511, 501], [31, 63, 93, 123, 150, 500]]) -@pytest.mark.parametrize("dtype", ["bfloat16"]) -def test_prefill_kernel_nonfull(dtype, num_q_heads, num_k_heads, num_v_heads, seq_lens): - _run_prefill_case( - dtype, - num_q_heads, - num_k_heads, - num_v_heads, - seq_lens, - scale=1.0 / math.sqrt(128), - alpha=True, - beta=True, - ) - - -@requires_runtime -@pytest.mark.parametrize("num_q_heads, num_k_heads, num_v_heads", [(1, 1, 1), (16, 16, 64)]) -@pytest.mark.parametrize("seq_len", [256, 255]) -def test_prefill_kernel_zero_length_sequence(num_q_heads, num_k_heads, num_v_heads, seq_len): - """A trailing zero-length sequence neither changes the output nor hangs.""" - _seed() - head_size = 128 - HO = max(num_q_heads, num_v_heads) - q, k, v = _gen_thd([seq_len], num_q_heads, num_k_heads, num_v_heads, head_size, torch.bfloat16) - alpha = torch.rand(seq_len, HO, device="cuda") - beta = torch.rand(seq_len, HO, device="cuda") - - ref_o, _ = chunk_gated_delta_rule(q, k, v, alpha, beta, 0.1, None, False, _cu([seq_len])) - our_o, _ = chunk_gated_delta_rule(q, k, v, alpha, beta, 0.1, None, False, _cu([seq_len, 0])) - torch.cuda.synchronize() - torch.testing.assert_close(our_o, ref_o, atol=2e-2, rtol=2e-2) - - -@requires_runtime -@pytest.mark.parametrize("with_initial_state", [False, True]) -def test_prefill_zero_length_sequence_state_passthrough(with_initial_state): - """A zero-length sequence's final-state slot gets the passthrough value: - its initial state when seeded, zeros otherwise.""" - _seed() - seq_len, head_size, num_heads, sentinel = 256, 128, 1, 123.0 - q, k, v = _gen_thd([seq_len], num_heads, num_heads, num_heads, head_size, torch.bfloat16) - alpha = torch.rand(seq_len, num_heads, device="cuda") - beta = torch.rand(seq_len, num_heads, device="cuda") - - s0 = torch.randn(2, num_heads, head_size, head_size, dtype=torch.float32, device="cuda") if with_initial_state else None - our_state = torch.full((2, num_heads, head_size, head_size), sentinel, dtype=torch.float32, device="cuda") - chunk_gated_delta_rule(q, k, v, alpha, beta, 0.1, s0, True, _cu([seq_len, 0]), output_state=our_state) - torch.cuda.synchronize() - expected = s0[1] if with_initial_state else torch.zeros_like(our_state[1]) - torch.testing.assert_close(our_state[1], expected, atol=0, rtol=0) - - -@requires_runtime -@pytest.mark.parametrize("num_q_heads, num_k_heads, num_v_heads", [(6, 2, 2), (2, 2, 4)]) -@pytest.mark.parametrize( - "seq_lens1, seq_lens2", - [([61], [128]), ([256, 256], [511, 501]), ([64, 128, 512], [123, 150, 500])], -) -def test_chunked_prefill(num_q_heads, num_k_heads, num_v_heads, seq_lens1, seq_lens2): - """Two-phase prefill carrying the state matches a single-shot reference.""" - _seed() - head_size = 128 - dtype = torch.bfloat16 - num_seqs = len(seq_lens1) - assert num_seqs == len(seq_lens2) - HO = max(num_q_heads, num_v_heads) - q1, k1, v1 = _gen_thd(seq_lens1, num_q_heads, num_k_heads, num_v_heads, head_size, dtype) - q2, k2, v2 = _gen_thd(seq_lens2, num_q_heads, num_k_heads, num_v_heads, head_size, dtype) - alpha1 = torch.empty(sum(seq_lens1), HO, device="cuda").uniform_(0.1, 1.0) - alpha2 = torch.empty(sum(seq_lens2), HO, device="cuda").uniform_(0.1, 1.0) - beta1 = torch.rand(sum(seq_lens1), HO, device="cuda") - beta2 = torch.rand(sum(seq_lens2), HO, device="cuda") - - scale = 1.0 / math.sqrt(head_size) - o1, state1 = chunk_gated_delta_rule(q1, k1, v1, alpha1, beta1, scale, None, True, _cu(seq_lens1)) - o2, state2 = chunk_gated_delta_rule(q2, k2, v2, alpha2, beta2, scale, state1, True, _cu(seq_lens2)) - torch.cuda.synchronize() - - def concat_varlen(t1, cua, t2, cub): - out = [] - for i in range(cua.size(0) - 1): - out.append(t1[cua[i] : cua[i + 1]]) - out.append(t2[cub[i] : cub[i + 1]]) - return torch.concat(out) - - cu1c, cu2c = _cu(seq_lens1).cpu(), _cu(seq_lens2).cpu() - our_o = concat_varlen(o1, cu1c, o2, cu2c) - q = concat_varlen(q1, cu1c, q2, cu2c) - k = concat_varlen(k1, cu1c, k2, cu2c) - v = concat_varlen(v1, cu1c, v2, cu2c) - alpha = concat_varlen(alpha1, cu1c, alpha2, cu2c) - beta = concat_varlen(beta1, cu1c, beta2, cu2c) - seq_lens = [a + b for a, b in zip(seq_lens1, seq_lens2)] - - ref_o, ref_state = _reference(q, k, v, alpha, beta, scale, seq_lens) - assert rms_ratio(our_o, ref_o) < FWD_TOL[dtype] - assert rms_ratio(state2, ref_state) < STATE_TOL[dtype] - - -@requires_runtime -@pytest.mark.parametrize("seq_lens", [[64], [256, 256], [64, 128, 512]]) -def test_prefill_kernel_state_dtype_bf16(seq_lens): - """bf16 recurrent state (initial + final) against the fp64 reference.""" - _seed() - head_size = 128 - num_heads = 3 - num_seqs = len(seq_lens) - total = sum(seq_lens) - q, k, v = _gen_thd(seq_lens, num_heads, num_heads, num_heads, head_size, torch.bfloat16) - alpha = torch.empty(total, num_heads, device="cuda").uniform_(0.1, 1.0) - beta = torch.rand(total, num_heads, device="cuda") - initial_state_ref = (torch.randn(num_seqs, num_heads, head_size, head_size, dtype=torch.float32, device="cuda") * 0.01).to(torch.bfloat16) - initial_state = initial_state_ref.contiguous() - - scale = 1.0 / math.sqrt(head_size) - our_state = torch.zeros(num_seqs, num_heads, head_size, head_size, dtype=torch.bfloat16, device="cuda") - our_o, _ = chunk_gated_delta_rule(q, k, v, alpha, beta, scale, initial_state, True, _cu(seq_lens), output_state=our_state) - torch.cuda.synchronize() - - ref_o, ref_state = _reference(q, k, v, alpha, beta, scale, seq_lens, initial_state=initial_state_ref.float()) - assert rms_ratio(our_o, ref_o) < 5e-2 - assert rms_ratio(our_state.float(), ref_state) < 5e-2 - - -# --------------------------------------------------------------------------- -# Per-chunk H output (fwd node surface; the GDN_BWD node's ``h`` source) -# --------------------------------------------------------------------------- - - -def _reference_h(q, k, v, alpha, beta, scale, seq_lens, every_n): - """fp64-chained per-chunk states: entry j of a sequence is the state after - (j + 1) * every_n tokens, strictly before the sequence end (bf16, K-major).""" - HO = max(q.shape[1], v.shape[1]) - - def expand(x): - r = HO // x.shape[1] - return (x.double().repeat_interleave(r, dim=1) if r > 1 else x.double()).unsqueeze(0) - - qq, kk, vv = expand(q), expand(k), expand(v) - gate = alpha.double().log().unsqueeze(0) - bb = beta.double().unsqueeze(0) - cu_piece = torch.tensor([0, every_n], dtype=torch.int32, device=q.device) - hs, off = [], 0 - for sl in seq_lens: - state = torch.zeros(1, HO, q.shape[2], v.shape[2], dtype=torch.float64, device=q.device) - for j in range(max(sl - 1, 0) // every_n): - a, b = off + j * every_n, off + (j + 1) * every_n - _o, state = gdn_reference( - qq[:, a:b], - kk[:, a:b], - vv[:, a:b], - gate[:, a:b], - bb[:, a:b], - scale=scale, - initial_state=state, - cu_seqlens=cu_piece, - ) - hs.append(state.squeeze(0)) - off += sl - return torch.stack(hs).to(q.dtype) if hs else torch.zeros(0, HO, q.shape[2], v.shape[2], dtype=q.dtype, device=q.device) - - -def _run_fwd_h(q, k, v, alpha, beta, scale, cu_seqlens, seq_lens, every_n): - """Engine-driven forward with the per-chunk H output; returns (o, fs, h).""" - device = q.device - total, HQ, HV, D = q.shape[0], q.shape[1], v.shape[1], q.shape[2] - HO = max(HQ, HV) - num_seqs = cu_seqlens.shape[0] - 1 - io_dt = cudnn.data_type.BFLOAT16 if q.dtype == torch.bfloat16 else cudnn.data_type.HALF - g, t = _build_gdn_graph( - total, - HQ, - HV, - D, - num_seqs, - scale, - io_dt, - output_final_state=True, - fs_dt=cudnn.data_type.FLOAT, - checkpoint_every_n_tokens=every_n, - ) - total_h = sum(max(sl - 1, 0) // every_n for sl in seq_lens) - o = torch.empty(total, HO, D, dtype=q.dtype, device=device) - fs = torch.empty(num_seqs, HO, D, D, dtype=torch.float32, device=device) - h = torch.full((max(total_h, 1), HO, D, D), float("nan"), dtype=q.dtype, device=device) - gate = (alpha if alpha is not None else torch.ones(total, HO, device=device)).float().log().contiguous() - beta_f = (beta if beta is not None else torch.ones(total, HO, device=device)).float().contiguous() - pack = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta_f, t["cu"]: cu_seqlens, t["O"]: o, t["fs"]: fs, t["H"]: h} - g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) - torch.cuda.synchronize() - return o, fs, h[:total_h] - - -@requires_runtime -@pytest.mark.parametrize("every_n", [64, 128]) -@pytest.mark.parametrize("seq_lens", [[256], [64, 128, 512], [255, 0, 511]]) -def test_fwd_h_matches_reference(seq_lens, every_n): - _seed() - q, k, v = _gen_thd(seq_lens, 2, 2, 2, 128, torch.bfloat16) - total = sum(seq_lens) - alpha = torch.empty(total, 2, device="cuda").uniform_(0.1, 1.0) - beta = torch.rand(total, 2, device="cuda") - scale = 1.0 / math.sqrt(128) - _o, _fs, h = _run_fwd_h(q, k, v, alpha, beta, scale, _cu(seq_lens), seq_lens, every_n) - assert not h.float().isnan().any(), "H has unwritten entries" - h_ref = _reference_h(q, k, v, alpha, beta, scale, seq_lens, every_n) - assert h.shape == h_ref.shape - if h.numel(): - r = rms_ratio(h.float(), h_ref.float()) - assert r < 5e-2, f"H rms ratio {r:.4g}" - - -@requires_runtime -def test_fwd_h_cutile_declines(): - - g = cudnn.pygraph() - q_t = g.tensor([256, 2, 128], data_type=cudnn.data_type.BFLOAT16, name="q") - k_t = g.tensor([256, 2, 128], data_type=cudnn.data_type.BFLOAT16, name="k") - v_t = g.tensor([256, 2, 128], data_type=cudnn.data_type.BFLOAT16, name="v") - g_t = g.tensor([256, 2], data_type=cudnn.data_type.FLOAT, name="g") - b_t = g.tensor([256, 2], data_type=cudnn.data_type.FLOAT, name="beta") - cu_t = g.tensor([2], data_type=cudnn.data_type.INT32, name="cu_seqlens") - O_t, _fs, h_t = g.gdn(q=q_t, k=k_t, v=v_t, g=g_t, beta=b_t, cu_seqlens=cu_t, scale=1.0, checkpoint_every_n_tokens=64, name="gdn") - O_t.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) - h_t.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) - assert_engine_declines(g, "gdn_cutile") # cuTile has no per-chunk H output - - -# --------------------------------------------------------------------------- -# split-K: strong decay on a long-sequence pack, so the table actually cuts -# --------------------------------------------------------------------------- - - -def _gen_split_case(seq_lens, HQ, HV): - total = sum(seq_lens) - HO = max(HQ, HV) - q = torch.randn(total, HQ, 128, dtype=torch.bfloat16, device="cuda") * 0.5 - k = F.normalize(torch.randn(total, HQ, 128, device="cuda"), dim=-1).bfloat16() - v = torch.randn(total, HV, 128, dtype=torch.bfloat16, device="cuda") * 0.5 - gate = (torch.rand(total, HO, device="cuda") * 0.9 + 0.05).float() - beta = torch.rand(total, HO, dtype=torch.float32, device="cuda") - return q, k, v, gate, beta, _cu(seq_lens) - - -@requires_runtime -@pytest.mark.parametrize("with_s0", [False, True]) -@pytest.mark.parametrize("heads", [(2, 2), (1, 4), (4, 1)]) # MHA, GVA, GQA -def test_prefill_split_strong_decay(heads, with_s0): - """Strong decay saturates the scan's warmup threshold on the 2048-token - sequence, so the split table cuts it into several work items; checked - against the fp64 recurrent reference.""" - _seed() - HQ, HV = heads - seq_lens = [100, 2048, 0, 517] - q, k, v, gate, beta, cu = _gen_split_case(seq_lens, HQ, HV) - HO = max(HQ, HV) - B = len(seq_lens) - s0 = (torch.randn(B, HO, 128, 128, dtype=torch.float32, device="cuda") * 0.05) if with_s0 else None - scale = 1.0 / math.sqrt(128) - - o = torch.full((sum(seq_lens), HO, 128), float("nan"), dtype=q.dtype, device="cuda") - fs = torch.full((B, HO, 128, 128), float("nan"), dtype=torch.float32, device="cuda") - chunk_gated_delta_rule(q, k, v, gate, beta, scale, s0, True, cu, output=o, output_state=fs) - torch.cuda.synchronize() - - o_ref, fs_ref = _reference(q, k, v, gate, beta, scale, seq_lens, initial_state=s0) - # zero-length sequences leave their final-state slot untouched - nz = torch.tensor([sl > 0 for sl in seq_lens], device="cuda") - assert not torch.isnan(o).any() and not torch.isnan(fs[nz]).any() - assert rms_ratio(o.float(), o_ref.float()) < FWD_TOL[torch.bfloat16] - assert rms_ratio(fs[nz], fs_ref[nz]) < STATE_TOL[torch.bfloat16] - - -# --------------------------------------------------------------------------- -# Determinism stress: bitwise repeat runs + two-stream co-residency -# --------------------------------------------------------------------------- - -DET_VARLEN_MIX = [497, 16, 1, 480, 0, 253] # zero-length + single-token + odd tails - - -def _det_launch(seq_lens, heads=(2, 2, 4), stream=None): - HQ, HK, HV = heads - HO = max(HQ, HV) - total = sum(seq_lens) - q, k, v = _gen_thd(seq_lens, HQ, HK, HV, 128, torch.bfloat16) - alpha = torch.empty(total, HO, device="cuda").uniform_(0.9, 1.0) - beta = torch.rand(total, HO, device="cuda").sigmoid() - cu = _cu(seq_lens) - scale = 1.0 / math.sqrt(128) - - def launch(): - if stream is not None: - with torch.cuda.stream(stream): - return chunk_gated_delta_rule(q, k, v, alpha, beta, scale, None, True, cu) - return chunk_gated_delta_rule(q, k, v, alpha, beta, scale, None, True, cu) - - return launch - - -@requires_runtime -@pytest.mark.parametrize("seq_lens", [DET_VARLEN_MIX, [4096]], ids=["varlen_mix", "long"]) -def test_gdn_prefill_determinism(seq_lens): - _seed() - assert_bitwise_runs(_det_launch(seq_lens), label="gdn") - - -@requires_runtime -def test_gdn_concurrent_streams_determinism(): - _seed() - s1, s2 = torch.cuda.Stream(), torch.cuda.Stream() - assert_concurrent_stream_runs(_det_launch([1024, 31, 512], stream=s1), _det_launch([256, 999, 1, 128], stream=s2), s1, s2) - - -# --------------------------------------------------------------------------- -# CUDA graph capture/replay across dynamic shapes (fixed SM-count grid) -# --------------------------------------------------------------------------- - - -@requires_runtime -def test_gdn_prefill_cuda_graph_replay(): - """Capture the ENGINE execute once, replay across CHANGED effective - shapes: capacity buffers, cu_seqlens with zero-length tails; everything - the engine launches (sched memset, split table, desc rebuilds, kernel) - must be capture-safe, and every replay must match an eager engine launch - on the same data bit for bit.""" - T_cap, B_cap, H, D = 768, 4, 2, 128 - dev = "cuda" - _seed() - scale = 1.0 / math.sqrt(D) - q = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - k = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - v = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - log_g = torch.zeros(T_cap, H, dtype=torch.float32, device=dev) - beta = torch.zeros(T_cap, H, dtype=torch.float32, device=dev) - cu = torch.zeros(B_cap + 1, dtype=torch.int32, device=dev) - o_graph = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - fs_graph = torch.zeros(B_cap, H, D, D, dtype=torch.float32, device=dev) - o_eager = torch.zeros_like(o_graph) - fs_eager = torch.zeros_like(fs_graph) - - g, t = _build_gdn_graph(T_cap, H, H, D, B_cap, scale, cudnn.data_type.BFLOAT16, output_final_state=True, fs_dt=cudnn.data_type.FLOAT) - ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=dev) - base = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: log_g, t["beta"]: beta, t["cu"]: cu} - pack_graph = {**base, t["O"]: o_graph, t["fs"]: fs_graph} - pack_eager = {**base, t["O"]: o_eager, t["fs"]: fs_eager} - - def fill(seq_lens): - total = sum(seq_lens) - q[:total] = torch.randn(total, H, D, device=dev).bfloat16() * 0.5 - k[:total] = F.normalize(torch.randn(total, H, D, device=dev), dim=-1).bfloat16() - v[:total] = torch.randn(total, H, D, device=dev).bfloat16() * 0.5 - log_g[:total] = torch.empty(total, H, device=dev).uniform_(0.1, 1.0).log() - beta[:total] = torch.rand(total, H, device=dev) - bounds = [0] + list(accumulate(seq_lens)) - bounds += [bounds[-1]] * (B_cap + 1 - len(bounds)) - cu.copy_(torch.tensor(bounds, dtype=torch.int32)) - return total - - fill([256, 448, 0, 0]) - stream = torch.cuda.Stream() - handle = cudnn.create_handle() - cudnn.set_stream(handle, stream.cuda_stream) - with torch.cuda.stream(stream): - g.execute(pack_graph, ws, handle=handle) # warmup: compile + caches - torch.cuda.synchronize() - - cg = torch.cuda.CUDAGraph() - with torch.cuda.graph(cg, stream=stream): - g.execute(pack_graph, ws, handle=handle) - - eager_handle = cudnn.create_handle() - cudnn.set_stream(eager_handle, torch.cuda.current_stream().cuda_stream) - for seq_lens in ([256, 448, 0, 0], [100, 200, 56, 0], [768], [64, 0, 64, 512]): - total = fill(seq_lens) - cg.replay() - torch.cuda.synchronize() - g.execute(pack_eager, ws, handle=eager_handle) - torch.cuda.synchronize() - assert torch.equal(o_graph[:total], o_eager[:total]), f"graph o diverges from eager at {seq_lens}" - assert torch.equal(fs_graph, fs_eager), f"graph final_state diverges from eager at {seq_lens}" diff --git a/test/python/linear_attention/frost/test_kda_bprop_kernel.py b/test/python/linear_attention/frost/test_kda_bprop_kernel.py deleted file mode 100644 index 3f14fd0f8..000000000 --- a/test/python/linear_attention/frost/test_kda_bprop_kernel.py +++ /dev/null @@ -1,42 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""FROST KDA backward: currently a STUB on this branch (the small-chunk -design recomputes the forward states in the backward once the kernel lands). -The contract: ``KdaFrostEngine`` declines ``KDA_BWD`` -graphs so the router can fall back.""" - -from __future__ import annotations - -import pytest - -import cudnn # noqa: F401 (conftest extends cudnn.__path__ with the source tree) - -from linear_attention.common import assert_engine_declines - -pytestmark = pytest.mark.L0 - - -def test_kda_bwd_frost_engine_declines(): - - total, H, D = 256, 2, 128 - g = cudnn.pygraph() - q_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="q") - k_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="k") - v_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="v") - g_t = g.tensor([total, H, D], data_type=cudnn.data_type.FLOAT, name="g") - beta_t = g.tensor([total, H], data_type=cudnn.data_type.FLOAT, name="beta") - cu_t = g.tensor([2], data_type=cudnn.data_type.INT32, name="cu_seqlens") - dO_t = g.tensor([total, H, D], data_type=cudnn.data_type.BFLOAT16, name="dO") - g.kda_bwd( - q=q_t, - k=k_t, - v=v_t, - g=g_t, - beta=beta_t, - cu_seqlens=cu_t, - dO=dO_t, - scale=0.125, - name="kda_bwd", - ) - assert_engine_declines(g, "kda_frost") # stub backward kernel on this branch diff --git a/test/python/linear_attention/frost/test_kda_prefill_kernel.py b/test/python/linear_attention/frost/test_kda_prefill_kernel.py deleted file mode 100644 index 0cf8b28aa..000000000 --- a/test/python/linear_attention/frost/test_kda_prefill_kernel.py +++ /dev/null @@ -1,757 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""FROST KDA prefill tests (``pygraph`` + ``KdaFrostEngine``) against the -fp64 recurrent reference.""" - -from __future__ import annotations - -import math -import random -from itertools import accumulate - -import pytest -import torch -import torch.nn.functional as F - -import cudnn # noqa: F401 (conftest extends cudnn.__path__ with the source tree) -from cudnn.linear_attention.frost import KdaFrostEngine - -from linear_attention.common import assert_bitwise_runs, assert_concurrent_stream_runs, assert_engine_declines -from linear_attention.conftest import multidist_randu -from linear_attention.reference_kda import kda_reference, rms_ratio - -pytestmark = pytest.mark.L0 - -SEED = 42 - - -def _sm100_dsl_available() -> bool: - if not torch.cuda.is_available(): - return False - major, _minor = torch.cuda.get_device_capability() - if major != 10: - return False - try: - import cutlass.experimental.primitives # noqa: F401 -- often a sys.modules alias, invisible to find_spec - except ImportError: - return False - return True - - -requires_runtime = pytest.mark.skipif(not _sm100_dsl_available(), reason="needs an SM100-class GPU and the Cutlass DSL") - - -def _seed(seed=SEED): - random.seed(seed) - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - - -def _cu(seq_lens, device="cuda"): - return torch.tensor([0] + list(accumulate(seq_lens)), dtype=torch.int32, device=device) - - -def _gen_thd(seq_lens, H, HV, head_size, dtype, HK=None): - HK = H if HK is None else HK - total = sum(seq_lens) - q = multidist_randu(total * H, head_size, device="cuda").reshape(total, H, head_size) - k = multidist_randu(total * HK, head_size, device="cuda").reshape(total, HK, head_size) - k = F.normalize(k, p=2.0, dim=-1) - v = multidist_randu(total * HV, head_size, device="cuda").reshape(total, HV, head_size) - return q.to(dtype).contiguous(), k.to(dtype).contiguous(), v.to(dtype).contiguous() - - -_DT = {torch.bfloat16: cudnn.data_type.BFLOAT16, torch.float16: cudnn.data_type.HALF, torch.float32: cudnn.data_type.FLOAT} - - -def _build_kda_engine_graph( - total, - H, - D, - num_seqs, - scale, - *, - io_dt=None, - HV=None, - use_qk_l2norm=True, - with_s0=False, - s0_dt=None, - fs_dt=None, - h_dt=None, - checkpoint_every_n_tokens=0, - use_beta_sigmoid=False, - safe_gate=False, - gate_lower_bound=None, - bwd=False, -): - - io_dt = io_dt or cudnn.data_type.BFLOAT16 - HV = HV or H - HO = max(H, HV) - g = cudnn.pygraph() - q_t = g.tensor([total, H, D], data_type=io_dt, name="q") - k_t = g.tensor([total, H, D], data_type=io_dt, name="k") - v_t = g.tensor([total, HV, D], data_type=io_dt, name="v") - g_t = g.tensor([total, HO, D], data_type=cudnn.data_type.FLOAT, name="g") - beta_t = g.tensor([total, HO], data_type=io_dt if use_beta_sigmoid else cudnn.data_type.FLOAT, name="beta") - cu_t = g.tensor([num_seqs + 1], data_type=cudnn.data_type.INT32, name="cu_seqlens") - t = dict(q=q_t, k=k_t, v=v_t, g=g_t, beta=beta_t, cu=cu_t) - if safe_gate: - t["a_log"] = g.tensor([HO], data_type=cudnn.data_type.FLOAT, name="a_log") - t["dt_bias"] = g.tensor([HO, D], data_type=cudnn.data_type.FLOAT, name="dt_bias") - if with_s0: - t["s0"] = g.tensor([num_seqs, HO, D, D], data_type=s0_dt or cudnn.data_type.FLOAT, name="initial_state") - if bwd: - t["dO"] = g.tensor([total, HO, D], data_type=io_dt, name="dO") - g.kda_bwd( - q=q_t, - k=k_t, - v=v_t, - g=g_t, - beta=beta_t, - cu_seqlens=cu_t, - dO=t["dO"], - scale=scale, - use_qk_l2norm=use_qk_l2norm, - name="kda_bwd", - ) - return g, t - O_t, fs_t, h_t = g.kda( - q=q_t, - k=k_t, - v=v_t, - g=g_t, - beta=beta_t, - cu_seqlens=cu_t, - initial_state=t.get("s0"), - scale=scale, - output_final_state=True, - use_qk_l2norm=use_qk_l2norm, - checkpoint_every_n_tokens=checkpoint_every_n_tokens, - use_beta_sigmoid=use_beta_sigmoid, - safe_gate=safe_gate, - gate_lower_bound=gate_lower_bound, - a_log=t.get("a_log"), - dt_bias=t.get("dt_bias"), - name="kda", - ) - O_t.set_output(True).set_data_type(io_dt) - fs_t.set_output(True).set_data_type(fs_dt or cudnn.data_type.FLOAT) - t["O"], t["fs"] = O_t, fs_t - if h_t is not None: - h_t.set_output(True).set_data_type(h_dt or io_dt) - t["H"] = h_t - return g, t - - -def _run_kda(q, k, v, gate, beta, scale, cu, initial_state=None, output_state=None, out_h=None, every_n=0, use_qk_l2norm=True): - """Torch adapter over the graph. ``gate`` is the per-key- - channel natural-log decay (fp32), ``beta`` the post-sigmoid scalar (fp32); - state ports are K-major ``[N, HO, K, V]``. GQA ``k`` (and ``HV < HQ`` - ``v``) are pre-broadcast: the node serves HK == HQ, HV a multiple of HQ.""" - device = q.device - H, HV = q.shape[1], v.shape[1] - if k.shape[1] != H: - k = k.repeat_interleave(H // k.shape[1], dim=1) - if HV < H: - v = v.repeat_interleave(H // HV, dim=1) - HV = H - total, D = q.shape[0], q.shape[2] - HO = max(H, HV) - num_seqs = cu.shape[0] - 1 - if output_state is None: - output_state = torch.empty(num_seqs, HO, D, D, dtype=torch.float32, device=device) - g, t = _build_kda_engine_graph( - total, - H, - D, - num_seqs, - scale, - io_dt=_DT[q.dtype], - HV=HV, - use_qk_l2norm=use_qk_l2norm, - with_s0=initial_state is not None, - s0_dt=None if initial_state is None else _DT[initial_state.dtype], - fs_dt=_DT[output_state.dtype], - h_dt=None if out_h is None else _DT[out_h.dtype], - checkpoint_every_n_tokens=every_n, - ) - g.build() - output = torch.empty(total, HO, D, dtype=q.dtype, device=device) - pack = { - t["q"]: q.contiguous(), - t["k"]: k.contiguous(), - t["v"]: v.contiguous(), - t["g"]: gate.float().contiguous(), - t["beta"]: beta.float().contiguous(), - t["cu"]: cu.to(torch.int32).contiguous(), - t["O"]: output, - t["fs"]: output_state, - } - if initial_state is not None: - pack[t["s0"]] = initial_state.contiguous() - if out_h is not None: - pack[t["H"]] = out_h - g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=device)) - return output, output_state - - -def _run_and_check(dtype, H, HV, seq_lens, with_s0=False, HK=None, gate_lo=None): - """Run the engine on random inputs and compare against the fp64 recurrent - reference (the kernel L2-normalizes q/k internally). fp16 io needs a - higher gate floor: k * exp2(-cumsum(g)) grazes the fp16 range at 0.5.""" - _seed() - head_size = 128 - total = sum(seq_lens) - HO = max(H, HV) - num_seqs = len(seq_lens) - q, k, v = _gen_thd(seq_lens, H, HV, head_size, dtype, HK=HK) - # BT=16 io-dtype Neumann inverse: stronger decay + post-sigmoid beta. - lo = gate_lo if gate_lo is not None else (0.6 if dtype == torch.float16 else 0.5) - gate = torch.empty(total, HO, head_size, device="cuda").uniform_(lo, 1.0).log() - beta = torch.rand(total, HO, device="cuda").sigmoid() - scale = 1.0 / math.sqrt(head_size) - s0 = (torch.randn(num_seqs, HO, head_size, head_size, dtype=torch.float32, device="cuda") * 0.05).contiguous() if with_s0 else None - output_state = torch.full((num_seqs, HO, head_size, head_size), float("nan"), dtype=torch.float32, device="cuda") - - o, fs = _run_kda(q, k, v, gate, beta, scale, _cu(seq_lens), initial_state=s0, output_state=output_state) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = kda_reference( - F.normalize(q.float(), dim=-1).unsqueeze(0), - F.normalize(k.float(), dim=-1).unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - scale=scale, - initial_state=s0, - cu_seqlens=_cu(seq_lens), - ) - tol = 1.2e-1 if dtype == torch.float16 else 1e-1 - torch.testing.assert_close(o.float(), o_ref.squeeze(0).float(), atol=tol, rtol=tol) - assert rms_ratio(fs, fs_ref) < 5e-2 # kernel state is K-major, like the reference - - -@requires_runtime -@pytest.mark.parametrize("H,HV", [(1, 1), (2, 2), (2, 4)]) -@pytest.mark.parametrize("seq_lens", [[256], [256, 256], [64, 128, 512]]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_kda_kernel_basic(dtype, H, HV, seq_lens): - _run_and_check(dtype, H, HV, seq_lens) - - -@requires_runtime -@pytest.mark.parametrize( - "seq_lens", - [[1], [15], [16], [17], [63, 65], [240, 255, 257], [7] * 24 + [1] * 8, [2048], [33] * 200], - ids=lambda s: f"{len(s)}seqs_{sum(s)}tok", -) -def test_kda_kernel_seqlen_edges(seq_lens): - """Chunk-boundary lengths (BT=16 +/- 1), single-token, many-short-seq - packs, a long sequence (many mbarrier ring wraps), and more tiles than - SMs (persistent CTAs walking several tiles).""" - _run_and_check(torch.bfloat16, 2, 4, seq_lens) - - -@requires_runtime -@pytest.mark.parametrize("seq_lens", [[240, 255, 257], [7] * 24 + [1] * 8]) -def test_kda_kernel_seqlen_edges_fp16(seq_lens): - _run_and_check(torch.float16, 2, 4, seq_lens) - - -@requires_runtime -@pytest.mark.parametrize("seq_lens", [[256], [64, 129, 512]]) -def test_kda_kernel_initial_state(seq_lens): - _run_and_check(torch.bfloat16, 2, 2, seq_lens, with_s0=True) - - -@requires_runtime -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_kda_beta_sigmoid_in_kernel(dtype): - """use_beta_sigmoid: io-dtype beta logits + in-kernel sigmoid must match - the host-side sigmoid path (the kernel roundtrips through the io dtype; - the approx-tanh sigmoid differs by at most ~1 io-dtype ulp).""" - _seed() - seq_lens, H, D = [256, 129], 2, 128 - total, num_seqs = sum(seq_lens), len(seq_lens) - q, k, v = _gen_thd(seq_lens, H, H, D, dtype) - gate = torch.empty(total, H, D, device="cuda").uniform_(0.5, 1.0).log() - logits = torch.randn(total, H, device="cuda") - scale = 1.0 / math.sqrt(D) - cu = _cu(seq_lens) - - outs = [] - for in_kernel in (True, False): - beta = logits.to(dtype).contiguous() if in_kernel else logits.to(dtype).float().sigmoid().to(dtype).float().contiguous() - g, t = _build_kda_engine_graph(total, H, D, num_seqs, scale, io_dt=_DT[dtype], use_beta_sigmoid=in_kernel) - g.build() - o = torch.empty(total, H, D, dtype=dtype, device="cuda") - fs = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device="cuda") - pack = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate.float().contiguous(), t["beta"]: beta, t["cu"]: cu, t["O"]: o, t["fs"]: fs} - g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda")) - torch.cuda.synchronize() - outs.append((o, fs)) - torch.testing.assert_close(outs[0][0].float(), outs[1][0].float(), atol=2.5e-2, rtol=2.5e-2) - assert rms_ratio(outs[0][1], outs[1][1]) < 2e-2 - - -@requires_runtime -def test_kda_safe_gate(dtype=torch.bfloat16): - """safe_gate: the in-kernel transform lower_bound * sigmoid(exp(a_log) * - (g + dt_bias)) must match feeding the host-computed transform as a plain - natural-log gate.""" - _seed() - seq_lens, H, D = [256, 129], 2, 128 - total, num_seqs = sum(seq_lens), len(seq_lens) - q, k, v = _gen_thd(seq_lens, H, H, D, dtype) - raw_gate = torch.randn(total, H, D, device="cuda").contiguous() - a_log = (torch.randn(H, device="cuda") * 0.3).contiguous() - dt_bias = (torch.randn(H, D, device="cuda") * 0.3).contiguous() - beta = torch.rand(total, H, device="cuda").sigmoid().contiguous() - scale = 1.0 / math.sqrt(D) - cu = _cu(seq_lens) - lower_bound = -5.0 # the kernel default, pinned through the attr - - outs = [] - for safe_gate in (True, False): - if safe_gate: - gate = raw_gate - else: - gate = (lower_bound * torch.sigmoid(a_log.exp().view(1, H, 1) * (raw_gate + dt_bias.view(1, H, D)))).contiguous() - g, t = _build_kda_engine_graph(total, H, D, num_seqs, scale, io_dt=_DT[dtype], safe_gate=safe_gate, gate_lower_bound=lower_bound if safe_gate else None) - g.build() - o = torch.empty(total, H, D, dtype=dtype, device="cuda") - fs = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device="cuda") - pack = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta, t["cu"]: cu, t["O"]: o, t["fs"]: fs} - if safe_gate: - pack[t["a_log"]] = a_log - pack[t["dt_bias"]] = dt_bias - g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda")) - torch.cuda.synchronize() - outs.append((o, fs)) - torch.testing.assert_close(outs[0][0].float(), outs[1][0].float(), atol=2.5e-2, rtol=2.5e-2) - assert rms_ratio(outs[0][1], outs[1][1]) < 2e-2 - - -@requires_runtime -@pytest.mark.parametrize("with_initial_state", [False, True]) -def test_kda_zero_length_sequence_state_passthrough(with_initial_state): - """A zero-length sequence's final-state slot gets the passthrough value: - its initial state when seeded, zeros otherwise.""" - _seed() - seq_len, H, D, sentinel = 256, 2, 128, 123.0 - q, k, v = _gen_thd([seq_len], H, H, D, torch.bfloat16) - gate = torch.empty(seq_len, H, D, device="cuda").uniform_(0.5, 1.0).log() - beta = torch.rand(seq_len, H, device="cuda").sigmoid() - s0 = torch.randn(2, H, D, D, dtype=torch.float32, device="cuda") if with_initial_state else None - fs = torch.full((2, H, D, D), sentinel, dtype=torch.float32, device="cuda") - _run_kda(q, k, v, gate, beta, 1.0 / math.sqrt(D), _cu([seq_len, 0]), initial_state=s0, output_state=fs) - torch.cuda.synchronize() - expected = s0[1] if with_initial_state else torch.zeros_like(fs[1]) - torch.testing.assert_close(fs[1], expected, atol=0, rtol=0) - - -@requires_runtime -@pytest.mark.parametrize("HQ,HK,HV", [(1, 1, 1), (4, 1, 1), (3, 3, 3), (6, 2, 2), (1, 1, 2), (2, 2, 4), (16, 16, 32), (16, 16, 64)]) -def test_kda_kernel_head_configs(HQ, HK, HV): - """The FlashInfer head-config matrix: GQA (HK < HQ), GVA (HV > HQ), odd - head counts, and production-sized 16/32/64-head grids (the adapter - pre-broadcasts k and any HV < HQ v to the node's HK == HQ contract).""" - _run_and_check(torch.bfloat16, HQ, HV, [64, 128, 512], HK=HK) - - -@requires_runtime -@pytest.mark.parametrize("HQ,HK,HV", [(2, 2, 4), (16, 16, 64)]) -def test_kda_kernel_head_configs_fp16(HQ, HK, HV): - _run_and_check(torch.float16, HQ, HV, [64, 128, 512], HK=HK) - - -@requires_runtime -def test_kda_chunked_prefill(): - """Splitting a sequence at a chunk boundary and reseeding from the fp32 - final state must reproduce the single-shot run (state roundtrip; the two - runs get different split tables, so the match is approximate).""" - _seed() - H, D, T1, T2 = 2, 128, 256, 192 - dtype = torch.bfloat16 - q, k, v = _gen_thd([T1 + T2], H, H, D, dtype) - gate = torch.empty(T1 + T2, H, D, device="cuda").uniform_(0.5, 1.0).log() - beta = torch.rand(T1 + T2, H, device="cuda").sigmoid() - scale = 1.0 / math.sqrt(D) - - fs_full = torch.full((1, H, D, D), float("nan"), dtype=torch.float32, device="cuda") - o_full, _ = _run_kda(q, k, v, gate, beta, scale, _cu([T1 + T2]), output_state=fs_full) - - fs1 = torch.full((1, H, D, D), float("nan"), dtype=torch.float32, device="cuda") - o1, _ = _run_kda(q[:T1], k[:T1], v[:T1], gate[:T1], beta[:T1], scale, _cu([T1]), output_state=fs1) - fs2 = torch.full((1, H, D, D), float("nan"), dtype=torch.float32, device="cuda") - o2, _ = _run_kda( - q[T1:].contiguous(), - k[T1:].contiguous(), - v[T1:].contiguous(), - gate[T1:].contiguous(), - beta[T1:].contiguous(), - scale, - _cu([T2]), - initial_state=fs1, - output_state=fs2, - ) - torch.cuda.synchronize() - - torch.testing.assert_close(o1.float(), o_full[:T1].float(), atol=2e-3, rtol=2e-3) - torch.testing.assert_close(o2.float(), o_full[T1:].float(), atol=2e-3, rtol=2e-3) - torch.testing.assert_close(fs2, fs_full, atol=2e-3, rtol=2e-3) - - -@requires_runtime -def test_kda_kernel_state_dtype_bf16(): - """bf16 initial/final state buffers (the io-downcast S0 path).""" - _seed() - seq_lens, H, D = [64, 128, 512], 2, 128 - total, num_seqs = sum(seq_lens), len(seq_lens) - q, k, v = _gen_thd(seq_lens, H, H, D, torch.bfloat16) - gate = torch.empty(total, H, D, device="cuda").uniform_(0.5, 1.0).log() - beta = torch.rand(total, H, device="cuda").sigmoid() - scale = 1.0 / math.sqrt(D) - s0 = (torch.randn(num_seqs, H, D, D, dtype=torch.float32, device="cuda") * 0.05).to(torch.bfloat16).contiguous() - fs = torch.full((num_seqs, H, D, D), float("nan"), dtype=torch.bfloat16, device="cuda") - o, _ = _run_kda(q, k, v, gate, beta, scale, _cu(seq_lens), initial_state=s0, output_state=fs) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = kda_reference( - F.normalize(q.float(), dim=-1).unsqueeze(0), - F.normalize(k.float(), dim=-1).unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - scale=scale, - initial_state=s0.float(), - cu_seqlens=_cu(seq_lens), - ) - torch.testing.assert_close(o.float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - assert rms_ratio(fs.float(), fs_ref) < 5e-2 - - -# --------------------------------------------------------------------------- -# KdaFrostEngine: graph-level coverage through the router -# --------------------------------------------------------------------------- - - -def _kda_engine_inputs(seq_lens, H, D): - from linear_attention.conftest import gen_kda_gates, gen_qkv - - total = sum(seq_lens) - q, k, v = gen_qkv(1, total, H, H, D, D, torch.bfloat16) - # BT=16 io-dtype Neumann inverse: stronger decay + post-sigmoid beta. - gate, beta = gen_kda_gates(1, total, H, D, torch.bfloat16, lo=0.5) - beta = beta.float().sigmoid() - cu = torch.tensor([0] + list(accumulate(seq_lens)), dtype=torch.int32, device="cuda") - return (x.squeeze(0).contiguous() for x in (q, k, v, gate, beta)), cu - - -@requires_runtime -@pytest.mark.parametrize("seq_lens", [[256], [512, 512], [64, 128, 512]]) -def test_kda_frost_engine_matches_reference(seq_lens, H=2, D=128): - - _seed() - (q, k, v, gate, beta), cu = _kda_engine_inputs(seq_lens, H, D) - total, num_seqs = sum(seq_lens), len(seq_lens) - scale = 1.0 / math.sqrt(D) - - g, t = _build_kda_engine_graph(total, H, D, num_seqs, scale) - g.build() - assert isinstance(g.selected_engine, KdaFrostEngine) - assert g.get_workspace_size() > 0 # split-K work-item table + scheduler counters - - o_buf = torch.empty(total, H, D, dtype=torch.bfloat16, device="cuda") - fs_buf = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device="cuda") - pack = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta, t["cu"]: cu, t["O"]: o_buf, t["fs"]: fs_buf} - g.execute(pack, torch.empty(g.get_workspace_size(), dtype=torch.uint8, device="cuda")) - torch.cuda.synchronize() - - # the kernel L2-normalizes q/k internally (use_qk_l2norm) - with torch.no_grad(): - o_ref, fs_ref = kda_reference( - F.normalize(q.float(), dim=-1).unsqueeze(0), - F.normalize(k.float(), dim=-1).unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - scale=scale, - cu_seqlens=cu, - ) - torch.testing.assert_close(o_buf.float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - r_s = rms_ratio(fs_buf, fs_ref) # engine state ports are K-major - assert r_s < 5e-2, f"final_state rms ratio {r_s:.4g}" - - -@requires_runtime -def test_kda_frost_engine_initial_state(seq_lens=(128, 256), H=2, D=128): - - _seed() - (q, k, v, gate, beta), cu = _kda_engine_inputs(seq_lens, H, D) - total, num_seqs = sum(seq_lens), len(seq_lens) - scale = 1.0 / math.sqrt(D) - s0 = torch.randn(num_seqs, H, D, D, dtype=torch.float32, device="cuda") * 0.05 - - g, t = _build_kda_engine_graph(total, H, D, num_seqs, scale, with_s0=True) - g.build() - assert isinstance(g.selected_engine, KdaFrostEngine) - - o_buf = torch.empty(total, H, D, dtype=torch.bfloat16, device="cuda") - fs_buf = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device="cuda") - pack = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta, t["cu"]: cu, t["s0"]: s0, t["O"]: o_buf, t["fs"]: fs_buf} - g.execute(pack, torch.empty(g.get_workspace_size(), dtype=torch.uint8, device="cuda")) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = kda_reference( - F.normalize(q.float(), dim=-1).unsqueeze(0), - F.normalize(k.float(), dim=-1).unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - scale=scale, - initial_state=s0, - cu_seqlens=cu, - ) - torch.testing.assert_close(o_buf.float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - r_s = rms_ratio(fs_buf, fs_ref) # engine state ports are K-major - assert r_s < 5e-2, f"final_state rms ratio {r_s:.4g}" - - -@requires_runtime -def test_kda_frost_engine_no_l2norm_matches_reference(seq_lens=(256,), H=2, D=128): - """use_qk_l2norm=False passes q/k through as given, so the test feeds - pre-normalized rows (the kernel's io-dtype arithmetic needs them).""" - - _seed() - (q, k, v, gate, beta), cu = _kda_engine_inputs(list(seq_lens), H, D) - q = F.normalize(q.float(), dim=-1).to(q.dtype) - k = F.normalize(k.float(), dim=-1).to(k.dtype) - total, num_seqs = sum(seq_lens), len(seq_lens) - scale = 1.0 / math.sqrt(D) - - g, t = _build_kda_engine_graph(total, H, D, num_seqs, scale, use_qk_l2norm=False) - g.build() - assert isinstance(g.selected_engine, KdaFrostEngine) - - o_buf = torch.empty(total, H, D, dtype=torch.bfloat16, device="cuda") - fs_buf = torch.empty(num_seqs, H, D, D, dtype=torch.float32, device="cuda") - pack = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta, t["cu"]: cu, t["O"]: o_buf, t["fs"]: fs_buf} - g.execute(pack, torch.empty(g.get_workspace_size(), dtype=torch.uint8, device="cuda")) - torch.cuda.synchronize() - - with torch.no_grad(): - o_ref, fs_ref = kda_reference( - q.float().unsqueeze(0), - k.float().unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - scale=scale, - cu_seqlens=cu, - ) - torch.testing.assert_close(o_buf.float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - r_s = rms_ratio(fs_buf, fs_ref) # engine state ports are K-major - assert r_s < 5e-2, f"final_state rms ratio {r_s:.4g}" - - -def test_kda_frost_engine_declines_bwd(): - """KDA_BWD declines (stub backward kernel on this branch).""" - g, _t = _build_kda_engine_graph(256, 2, 128, 1, 0.125, bwd=True) - assert_engine_declines(g, "kda_frost") - - -@requires_runtime -def test_kda_frost_engine_declines_wrong_head_dim(): - g, _t = _build_kda_engine_graph(256, 2, 64, 1, 0.125) - assert_engine_declines(g, "kda_frost") - - -# --------------------------------------------------------------------------- -# Determinism stress: bitwise repeat runs + two-stream co-residency -# --------------------------------------------------------------------------- - -DET_VARLEN_MIX = [497, 16, 1, 480, 0, 253] # zero-length + single-token + odd tails - - -def _det_launch(seq_lens, H, HV, with_s0=False, stream=None): - q, k, v = _gen_thd(seq_lens, H, HV, 128, torch.bfloat16) - total, HO = sum(seq_lens), max(H, HV) - gate = torch.empty(total, HO, 128, device="cuda").uniform_(0.5, 1.0).log().float().contiguous() - beta = torch.rand(total, HO, device="cuda").sigmoid().float().contiguous() - s0 = (torch.randn(len(seq_lens), HO, 128, 128, dtype=torch.float32, device="cuda") * 0.05).contiguous() if with_s0 else None - cu = _cu(seq_lens) - scale = 1.0 / math.sqrt(128) - - def launch(): - if stream is not None: - with torch.cuda.stream(stream): - return _run_kda(q, k, v, gate, beta, scale, cu, initial_state=s0) - return _run_kda(q, k, v, gate, beta, scale, cu, initial_state=s0) - - return launch - - -@requires_runtime -@pytest.mark.parametrize("seq_lens", [DET_VARLEN_MIX, [4096]], ids=["varlen_mix", "long"]) -def test_kda_prefill_determinism(seq_lens): - _seed() - assert_bitwise_runs(_det_launch(seq_lens, 2, 4), label="kda") - - -@requires_runtime -def test_kda_prefill_determinism_initial_state(): - _seed() - assert_bitwise_runs(_det_launch([256, 640, 0, 33], 2, 2, with_s0=True), label="kda+s0") - - -@requires_runtime -def test_kda_concurrent_streams_determinism(): - _seed() - s1, s2 = torch.cuda.Stream(), torch.cuda.Stream() - assert_concurrent_stream_runs(_det_launch([1024, 31, 512], 2, 4, stream=s1), _det_launch([256, 999, 1, 128], 2, 4, stream=s2), s1, s2) - - -# --------------------------------------------------------------------------- -# split-K: strong decay on a long-sequence pack, so the table actually cuts -# --------------------------------------------------------------------------- - - -@requires_runtime -@pytest.mark.parametrize("with_s0", [False, True]) -@pytest.mark.parametrize("heads", [(2, 2), (1, 4)]) # MHA, GVA -def test_kda_prefill_split_strong_decay(heads, with_s0): - """Strong decay (gate floor 0.3) saturates the scan's warmup threshold on - the 2048-token sequence, so the split table cuts it into several work - items; checked against the fp64 recurrent reference.""" - HQ, HV = heads - _run_and_check(torch.bfloat16, HQ, HV, [100, 2048, 0, 517], with_s0=with_s0, gate_lo=0.3) - - -@requires_runtime -def test_kda_prefill_split_safe_gate(dtype=torch.bfloat16): - """safe_gate + split-K: the partition scan applies the gate transform - itself, so cuts land on true decay values; checked against the fp64 - reference fed the host-computed transform.""" - _seed() - HO = 2 - seq_lens = [100, 2048, 0, 517] - total = sum(seq_lens) - B = len(seq_lens) - q, k, v = _gen_thd(seq_lens, HO, HO, 128, dtype) - raw_gate = (torch.randn(total, HO, 128, device="cuda") + 1.0).contiguous() - a_log = (torch.randn(HO, device="cuda") * 0.3).contiguous() - dt_bias = (torch.randn(HO, 128, device="cuda") * 0.3).contiguous() - beta = torch.rand(total, HO, device="cuda").sigmoid().float() - cu = _cu(seq_lens) - scale = 1.0 / math.sqrt(128) - lower_bound = -5.0 - - g, t = _build_kda_engine_graph(total, HO, 128, B, scale, safe_gate=True, gate_lower_bound=lower_bound) - g.build() - o = torch.full((total, HO, 128), float("nan"), dtype=dtype, device="cuda") - fs = torch.full((B, HO, 128, 128), float("nan"), dtype=torch.float32, device="cuda") - pack = { - t["q"]: q, - t["k"]: k, - t["v"]: v, - t["g"]: raw_gate, - t["beta"]: beta, - t["cu"]: cu, - t["O"]: o, - t["fs"]: fs, - t["a_log"]: a_log, - t["dt_bias"]: dt_bias, - } - g.execute(pack, torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda")) - torch.cuda.synchronize() - assert not torch.isnan(o).any() and not torch.isnan(fs).any() - - gate = lower_bound * torch.sigmoid(a_log.exp().view(1, HO, 1) * (raw_gate + dt_bias.view(1, HO, 128))) - with torch.no_grad(): - o_ref, fs_ref = kda_reference( - F.normalize(q.float(), dim=-1).unsqueeze(0), - F.normalize(k.float(), dim=-1).unsqueeze(0), - v.unsqueeze(0), - gate.unsqueeze(0), - beta.unsqueeze(0), - scale=scale, - cu_seqlens=cu, - ) - torch.testing.assert_close(o.float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - assert rms_ratio(fs, fs_ref) < 5e-2 - - -# --------------------------------------------------------------------------- -# CUDA graph capture/replay across dynamic shapes (fixed SM-count grid) -# --------------------------------------------------------------------------- - - -@requires_runtime -def test_kda_prefill_cuda_graph_replay(): - """Capture the ENGINE execute once, replay across CHANGED effective - shapes: capacity buffers, cu_seqlens with zero-length tails; everything - the engine launches (sched memset, split table, desc rebuilds, kernel) - must be capture-safe, and every replay must match an eager engine launch - on the same data bit for bit.""" - T_cap, B_cap, H, D = 768, 4, 2, 128 - dev = "cuda" - _seed() - scale = 1.0 / math.sqrt(D) - q = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - k = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - v = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - gate = torch.zeros(T_cap, H, D, dtype=torch.float32, device=dev) - beta = torch.zeros(T_cap, H, dtype=torch.float32, device=dev) - cu = torch.zeros(B_cap + 1, dtype=torch.int32, device=dev) - o_graph = torch.zeros(T_cap, H, D, dtype=torch.bfloat16, device=dev) - fs_graph = torch.zeros(B_cap, H, D, D, dtype=torch.float32, device=dev) - o_eager = torch.zeros_like(o_graph) - fs_eager = torch.zeros_like(fs_graph) - - g, t = _build_kda_engine_graph(T_cap, H, D, B_cap, scale) - g.build() - ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device=dev) - base = {t["q"]: q, t["k"]: k, t["v"]: v, t["g"]: gate, t["beta"]: beta, t["cu"]: cu} - pack_graph = {**base, t["O"]: o_graph, t["fs"]: fs_graph} - pack_eager = {**base, t["O"]: o_eager, t["fs"]: fs_eager} - - def fill(seq_lens): - total = sum(seq_lens) - q[:total] = torch.randn(total, H, D, device=dev).bfloat16() * 0.5 - k[:total] = F.normalize(torch.randn(total, H, D, device=dev), dim=-1).bfloat16() - v[:total] = torch.randn(total, H, D, device=dev).bfloat16() * 0.5 - gate[:total] = torch.empty(total, H, D, device=dev).uniform_(0.5, 1.0).log() - beta[:total] = torch.rand(total, H, device=dev) - bounds = [0] + list(accumulate(seq_lens)) - bounds += [bounds[-1]] * (B_cap + 1 - len(bounds)) - cu.copy_(torch.tensor(bounds, dtype=torch.int32)) - return total - - fill([256, 512, 0, 0]) - stream = torch.cuda.Stream() - handle = cudnn.create_handle() - cudnn.set_stream(handle, stream.cuda_stream) - with torch.cuda.stream(stream): - g.execute(pack_graph, ws, handle=handle) # warmup: compile + caches - torch.cuda.synchronize() - - cg = torch.cuda.CUDAGraph() - with torch.cuda.graph(cg, stream=stream): - g.execute(pack_graph, ws, handle=handle) - - eager_handle = cudnn.create_handle() - cudnn.set_stream(eager_handle, torch.cuda.current_stream().cuda_stream) - for seq_lens in ([256, 512, 0, 0], [100, 200, 56, 0], [768], [16, 0, 16, 736 - 32]): - total = fill(seq_lens) - cg.replay() - torch.cuda.synchronize() - g.execute(pack_eager, ws, handle=eager_handle) - torch.cuda.synchronize() - assert torch.equal(o_graph[:total], o_eager[:total]), f"graph o diverges from eager at {seq_lens}" - assert torch.equal(fs_graph, fs_eager), f"graph final_state diverges from eager at {seq_lens}" diff --git a/test/python/linear_attention/ops/__init__.py b/test/python/linear_attention/ops/__init__.py deleted file mode 100644 index 52a7a9daf..000000000 --- a/test/python/linear_attention/ops/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 diff --git a/test/python/linear_attention/ops/test_gdn2_op.py b/test/python/linear_attention/ops/test_gdn2_op.py deleted file mode 100644 index 225bca7fa..000000000 --- a/test/python/linear_attention/ops/test_gdn2_op.py +++ /dev/null @@ -1,136 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""GDN-2 custom-op tests (``cudnn.linear_attention.ops.gated_delta_net_v2``): -THD layout and torch.compile against the fp64 recurrent reference. GDN-2 runs -on the FROST SM100 engine only (forward), so these tests require an -SM100-class GPU with the Cutlass DSL runtime.""" - -from __future__ import annotations - -import importlib.util -import math - -import pytest -import torch -import torch.nn.functional as F - -from cudnn.linear_attention.ops import gated_delta_net_v2 - -from ..conftest import gen_gdn2_gates, gen_qkv -from ..reference_gdn2 import gdn2_reference, rms_ratio - -pytestmark = pytest.mark.L0 - - -def _frost_gdn2_available() -> bool: - if not torch.cuda.is_available(): - return False - major, _minor = torch.cuda.get_device_capability() - if major != 10 or importlib.util.find_spec("cutlass") is None: - return False - try: - import cutlass.experimental.primitives # noqa: F401 — the engine's own availability gate - except Exception: # noqa: BLE001 - return False - return True - - -requires_runtime = pytest.mark.skipif(not _frost_gdn2_available(), reason="needs an SM100-class GPU and the Cutlass DSL GDN-2 prefill kernel runtime") - - -def _inputs(B, T, H, K=128, V=128, dtype=torch.bfloat16, seed=0): - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - q, k, v = gen_qkv(B, T, H, H, K, V, dtype) - g, beta, w = gen_gdn2_gates(B, T, H, K, V, dtype) - return q, k, v, g, beta, w - - -def _thd(x): - return x.reshape(-1, *x.shape[2:]) - - -def _cu(B, T): - return torch.arange(0, B + 1, dtype=torch.int32, device="cuda") * T - - -@requires_runtime -class TestGdn2Op: - def test_forward_parity(self, B=1, T=256, H=2): - q, k, v, g, beta, w = _inputs(B, T, H) - scale = 1.0 / math.sqrt(128) - o, _ = gated_delta_net_v2(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _thd(w), _cu(B, T), scale=scale, use_qk_l2norm_in_kernel=True) - with torch.no_grad(): - o_ref, _ = gdn2_reference(F.normalize(q.float(), dim=-1), F.normalize(k.float(), dim=-1), v, g, beta, w, scale=scale) - torch.testing.assert_close(o.view_as(o_ref).float(), o_ref.float(), atol=1e-1, rtol=1e-1) - - def test_forward_parity_no_l2norm(self, B=1, T=256, H=2): - q, k, v, g, beta, w = _inputs(B, T, H) - q = F.normalize(q.float(), dim=-1).to(q.dtype) - k = F.normalize(k.float(), dim=-1).to(k.dtype) - scale = 1.0 / math.sqrt(128) - o, _ = gated_delta_net_v2(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _thd(w), _cu(B, T), scale=scale, use_qk_l2norm_in_kernel=False) - with torch.no_grad(): - o_ref, _ = gdn2_reference(q.float(), k.float(), v, g, beta, w, scale=scale) - torch.testing.assert_close(o.view_as(o_ref).float(), o_ref.float(), atol=1e-1, rtol=1e-1) - - def test_forward_parity_thd(self): - seq_lens = [64, 192] - total = sum(seq_lens) - q, k, v, g, beta, w = (x.squeeze(0) for x in _inputs(1, total, 2)) - cu = torch.tensor([0, 64, 256], dtype=torch.int32, device="cuda") - o, _ = gated_delta_net_v2(q, k, v, g, beta, w, cu, use_qk_l2norm_in_kernel=True) - with torch.no_grad(): - o_ref, _ = gdn2_reference( - F.normalize(q.float(), dim=-1).unsqueeze(0), - F.normalize(k.float(), dim=-1).unsqueeze(0), - v.unsqueeze(0), - g.unsqueeze(0), - beta.unsqueeze(0), - w.unsqueeze(0), - cu_seqlens=cu, - ) - torch.testing.assert_close(o.float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - - def test_initial_state(self, B=1, T=256, H=2): - q, k, v, g, beta, w = _inputs(B, T, H) - S0 = torch.randn(B, H, 128, 128, dtype=torch.float32, device="cuda") * 0.05 - scale = 1.0 / math.sqrt(128) - o, fs = gated_delta_net_v2( - _thd(q), - _thd(k), - _thd(v), - _thd(g), - _thd(beta), - _thd(w), - _cu(B, T), - scale=scale, - initial_state=S0, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - ) - with torch.no_grad(): - o_ref, fs_ref = gdn2_reference(F.normalize(q.float(), dim=-1), F.normalize(k.float(), dim=-1), v, g, beta, w, scale=scale, initial_state=S0) - torch.testing.assert_close(o.view_as(o_ref).float(), o_ref.float(), atol=1e-1, rtol=1e-1) - assert rms_ratio(fs, fs_ref) < 5e-2 - - def test_default_scale(self): - q, k, v, g, beta, w = (_thd(x) for x in _inputs(1, 128, 1)) - cu = _cu(1, 128) - o_default, _ = gated_delta_net_v2(q, k, v, g, beta, w, cu, use_qk_l2norm_in_kernel=True) - o_explicit, _ = gated_delta_net_v2(q, k, v, g, beta, w, cu, scale=1.0 / math.sqrt(128), use_qk_l2norm_in_kernel=True) - torch.testing.assert_close(o_default, o_explicit) - - def test_no_final_state_returns_empty(self): - q, k, v, g, beta, w = (_thd(x) for x in _inputs(1, 128, 1)) - _o, final = gated_delta_net_v2(q, k, v, g, beta, w, _cu(1, 128), output_final_state=False, use_qk_l2norm_in_kernel=True) - assert final.numel() == 0 - - def test_torch_compile_forward(self): - q, k, v, g, beta, w = (_thd(x) for x in _inputs(1, 128, 1)) - cu = _cu(1, 128) - compiled = torch.compile(gated_delta_net_v2, fullgraph=True) - o_eager, _ = gated_delta_net_v2(q, k, v, g, beta, w, cu, use_qk_l2norm_in_kernel=True) - o_comp, _ = compiled(q, k, v, g, beta, w, cu, use_qk_l2norm_in_kernel=True) - torch.testing.assert_close(o_eager, o_comp) diff --git a/test/python/linear_attention/ops/test_gdn_op.py b/test/python/linear_attention/ops/test_gdn_op.py deleted file mode 100644 index 7063385b4..000000000 --- a/test/python/linear_attention/ops/test_gdn_op.py +++ /dev/null @@ -1,205 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""GDN custom-op tests (``cudnn.linear_attention.gated_delta_net``): THD -layout, autograd, and torch.compile against the fp64 recurrent reference.""" - -from __future__ import annotations - -import importlib.util -import math - -import pytest -import torch - -from cudnn.linear_attention.ops import gated_delta_net - -from ..common import FWD_TOL, GDN_MARKS, STATE_TOL -from ..conftest import gen_gates, gen_qkv -from ..reference_gdn import gdn_reference, rms_ratio - -pytestmark = GDN_MARKS - - -def _frost_gdn_available() -> bool: - if not torch.cuda.is_available(): - return False - major, _minor = torch.cuda.get_device_capability() - if major != 10 or importlib.util.find_spec("cutlass") is None: - return False - try: - import cutlass.experimental.primitives # noqa: F401 — the engine's own availability gate - except Exception: # noqa: BLE001 - return False - return True - - -requires_frost = pytest.mark.skipif(not _frost_gdn_available(), reason="GDN with HV < HQ (GQA-style v broadcast) is served by the FROST engine only") - - -def _inputs(B, T, H, HV, K, V, dtype=torch.bfloat16, seed=0): - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - q, k, v = gen_qkv(B, T, H, HV, K, V, dtype) - g, beta = gen_gates(B, T, HV, torch.float32) # the op requires kernel-native fp32 gates - return q, k, v, g, beta - - -def _thd(x): - return x.reshape(-1, *x.shape[2:]) - - -def _cu(B, T): - return torch.arange(0, B + 1, dtype=torch.int32, device="cuda") * T - - -class TestGdnOp: - @pytest.mark.parametrize( - "B,T,H,K,V", - [ - (1, 128, 1, 64, 64), - (2, 256, 4, 64, 64), - (1, 128, 2, 64, 128), # K != V - ], - ) - def test_forward_parity(self, B, T, H, K, V): - q, k, v, g, beta = _inputs(B, T, H, H, K, V) - scale = 1.0 / math.sqrt(K) - o, _ = gated_delta_net(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), scale=scale) - with torch.no_grad(): - o_ref, _ = gdn_reference(q, k, v, g, beta, scale=scale) - assert rms_ratio(o.view_as(o_ref), o_ref) < FWD_TOL[torch.bfloat16] - - def test_forward_parity_gva(self, B=1, T=192, H=2, HV=4, K=64, V=64): - q, k, v, g, beta = _inputs(B, T, H, HV, K, V) - o, fs = gated_delta_net(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), output_final_state=True) - with torch.no_grad(): - o_ref, fs_ref = gdn_reference(q, k, v, g, beta) - assert rms_ratio(o.view_as(o_ref), o_ref) < FWD_TOL[torch.bfloat16] - assert rms_ratio(fs, fs_ref) < STATE_TOL[torch.bfloat16] - - @pytest.mark.parametrize("B,T,H,K,V", [(1, 128, 1, 64, 64), (2, 128, 2, 64, 64)]) - def test_backward_runs(self, B, T, H, K, V): - q, k, v, g, beta = (_thd(x).detach().requires_grad_(True) for x in _inputs(B, T, H, H, K, V)) - o, _ = gated_delta_net(q, k, v, g, beta, _cu(B, T)) - o.sum().backward() - for name, t in [("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta)]: - assert t.grad is not None, f"no grad for {name}" - assert torch.isfinite(t.grad).all(), f"non-finite grad for {name}" - - def test_default_scale(self): - q, k, v, g, beta = (_thd(x) for x in _inputs(1, 128, 1, 1, 64, 64)) - cu = _cu(1, 128) - o_default, _ = gated_delta_net(q, k, v, g, beta, cu) - o_explicit, _ = gated_delta_net(q, k, v, g, beta, cu, scale=1.0 / math.sqrt(64)) - torch.testing.assert_close(o_default, o_explicit) - - def test_no_final_state_returns_empty(self): - q, k, v, g, beta = (_thd(x) for x in _inputs(1, 128, 1, 1, 64, 64)) - _o, final = gated_delta_net(q, k, v, g, beta, _cu(1, 128), output_final_state=False) - assert final.numel() == 0 - - def test_initial_state(self): - B, T, H, K, V = 1, 128, 2, 64, 64 - q, k, v, g, beta = _inputs(B, T, H, H, K, V) - S0 = torch.randn(B, H, K, V, dtype=torch.float32, device="cuda") * 0.05 - o, fs = gated_delta_net(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), initial_state=S0, output_final_state=True) - with torch.no_grad(): - o_ref, fs_ref = gdn_reference(q, k, v, g, beta, initial_state=S0) - assert rms_ratio(o.view_as(o_ref), o_ref) < FWD_TOL[torch.bfloat16] - assert rms_ratio(fs, fs_ref) < STATE_TOL[torch.bfloat16] - - def test_packed_matches_per_sequence(self): - B, T, H, K, V = 2, 128, 2, 64, 64 - q, k, v, g, beta = _inputs(B, T, H, H, K, V) - o, fs = gated_delta_net(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), output_final_state=True) - for b in range(B): - o_b, fs_b = gated_delta_net(q[b], k[b], v[b], g[b], beta[b], _cu(1, T), output_final_state=True) - torch.testing.assert_close(o[b * T : (b + 1) * T], o_b) - torch.testing.assert_close(fs[b], fs_b[0]) - - def test_thd_ragged_parity_and_backward(self): - seq_lens = [64, 192] - H, K, V = 2, 64, 64 - total = sum(seq_lens) - q, k, v, g, beta = (x.squeeze(0).detach().requires_grad_(True) for x in _inputs(1, total, H, H, K, V)) - cu = torch.tensor([0, 64, 256], dtype=torch.int32, device="cuda") - - o, _ = gated_delta_net(q, k, v, g, beta, cu_seqlens=cu) - with torch.no_grad(): - o_ref, _ = gdn_reference(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), g.unsqueeze(0), beta.unsqueeze(0), cu_seqlens=cu) - assert rms_ratio(o, o_ref.squeeze(0)) < FWD_TOL[torch.bfloat16] - - o.sum().backward() - for name, t in [("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta)]: - assert t.grad is not None and torch.isfinite(t.grad).all(), f"bad grad for {name}" - - @requires_frost - def test_forward_parity_gqa(self, B=1, T=192, H=4, HV=1, K=128, V=128): - """GQA (q heads group over v heads): gates/o/state at HO = H; the - reference expands v over the head group (each output head runs its - own q/k against the shared v).""" - torch.random.manual_seed(0) - torch.cuda.manual_seed(0) - q, k, v = gen_qkv(B, T, H, HV, K, V, torch.bfloat16) - g, beta = gen_gates(B, T, H, torch.float32) # the op requires kernel-native fp32 gates - o, fs = gated_delta_net(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), output_final_state=True) - v_exp = v.repeat_interleave(H // HV, dim=2) - with torch.no_grad(): - o_ref, fs_ref = gdn_reference(q, k, v_exp, g, beta) - assert rms_ratio(o.view_as(o_ref), o_ref) < FWD_TOL[torch.bfloat16] - assert rms_ratio(fs, fs_ref) < STATE_TOL[torch.bfloat16] - - @requires_frost - def test_backward_parity_gqa(self, B=1, T=128, H=4, HV=2, K=128, V=128): - """GQA gradients vs the op's own (validated) MHA path on the expanded - v: dQ/dK/dG/dBeta match directly, dV matches the head-group sum.""" - torch.random.manual_seed(0) - torch.cuda.manual_seed(0) - q, k, v = gen_qkv(B, T, H, HV, K, V, torch.bfloat16) - g, beta = gen_gates(B, T, H, torch.float32) # the op requires kernel-native fp32 gates - r = H // HV - cu = _cu(B, T) - - gqa = {n: _thd(x).detach().requires_grad_(True) for n, x in (("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta))} - o, _ = gated_delta_net(gqa["q"], gqa["k"], gqa["v"], gqa["g"], gqa["beta"], cu) - w = torch.randn_like(o.float()) - (o.float() * w).sum().backward() - - mha = {n: _thd(x).detach().requires_grad_(True) for n, x in (("q", q), ("k", k), ("v", v.repeat_interleave(r, dim=2)), ("g", g), ("beta", beta))} - o_ref, _ = gated_delta_net(mha["q"], mha["k"], mha["v"], mha["g"], mha["beta"], cu) - (o_ref.float() * w).sum().backward() - - torch.testing.assert_close(o, o_ref, atol=1e-3, rtol=1e-3) - for name in ("q", "k", "g", "beta"): - torch.testing.assert_close(gqa[name].grad, mha[name].grad, atol=1e-3, rtol=1e-3, msg=f"d{name} mismatch") - dv_ref = mha["v"].grad.view(B * T, HV, r, V).sum(2) - torch.testing.assert_close(gqa["v"].grad.float(), dv_ref.float(), atol=1e-2, rtol=1e-2, msg="dV mismatch") - - def test_forward_parity_qk_l2norm(self, B=1, T=192, H=2, K=64, V=64): - import torch.nn.functional as F - - q, k, v, g, beta = _inputs(B, T, H, H, K, V) - scale = 1.0 / math.sqrt(K) - o, fs = gated_delta_net(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), scale=scale, output_final_state=True, use_qk_l2norm_in_kernel=True) - with torch.no_grad(): - o_ref, fs_ref = gdn_reference(F.normalize(q.float(), dim=-1), F.normalize(k.float(), dim=-1), v, g, beta, scale=scale) - assert rms_ratio(o.view_as(o_ref), o_ref) < FWD_TOL[torch.bfloat16] - assert rms_ratio(fs, fs_ref) < STATE_TOL[torch.bfloat16] - - @pytest.mark.parametrize("H,HV", [(2, 2), (2, 4)]) - def test_backward_runs_qk_l2norm(self, H, HV, B=1, T=128, K=64, V=64): - q, k, v, g, beta = (_thd(x).detach().requires_grad_(True) for x in _inputs(B, T, H, HV, K, V)) - o, _ = gated_delta_net(q, k, v, g, beta, _cu(B, T), use_qk_l2norm_in_kernel=True) - o.sum().backward() - for name, t in [("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta)]: - assert t.grad is not None, f"no grad for {name}" - assert torch.isfinite(t.grad).all(), f"non-finite grad for {name}" - - def test_torch_compile_forward(self): - q, k, v, g, beta = (_thd(x) for x in _inputs(1, 128, 1, 1, 64, 64)) - cu = _cu(1, 128) - compiled = torch.compile(gated_delta_net, fullgraph=True) - o_eager, _ = gated_delta_net(q, k, v, g, beta, cu) - o_comp, _ = compiled(q, k, v, g, beta, cu) - torch.testing.assert_close(o_eager, o_comp) diff --git a/test/python/linear_attention/ops/test_kda_op.py b/test/python/linear_attention/ops/test_kda_op.py deleted file mode 100644 index 20369fc4a..000000000 --- a/test/python/linear_attention/ops/test_kda_op.py +++ /dev/null @@ -1,234 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""KDA custom-op tests (``cudnn.linear_attention.ops.kimi_delta_attention``): -THD layout, autograd, and torch.compile against the fp64 recurrent -reference.""" - -from __future__ import annotations - -import importlib.util -import math - -import pytest -import torch - -from cudnn.linear_attention.ops import kimi_delta_attention - -from ..common import FWD_TOL, KDA_MARKS, STATE_TOL -from ..conftest import gen_kda_gates, gen_qkv -from ..reference_kda import kda_reference, rms_ratio - -pytestmark = KDA_MARKS - - -def _frost_kda_available() -> bool: - if not torch.cuda.is_available(): - return False - major, _minor = torch.cuda.get_device_capability() - if major != 10 or importlib.util.find_spec("cutlass") is None: - return False - try: - import cutlass.experimental.primitives # noqa: F401 — the engine's own availability gate - except Exception: # noqa: BLE001 - return False - return True - - -def _inputs(B, T, H, HV, K, V, dtype=torch.bfloat16, seed=0): - torch.random.manual_seed(seed) - torch.cuda.manual_seed(seed) - q, k, v = gen_qkv(B, T, H, HV, K, V, dtype) - g, beta = gen_kda_gates(B, T, HV, K, torch.float32) # the op requires kernel-native fp32 gates - return q, k, v, g, beta - - -def _thd(x): - return x.reshape(-1, *x.shape[2:]) - - -def _cu(B, T): - return torch.arange(0, B + 1, dtype=torch.int32, device="cuda") * T - - -class TestKdaOp: - @pytest.mark.parametrize( - "B,T,H,K,V", - [ - (1, 128, 1, 64, 64), - (2, 256, 4, 64, 64), - (1, 128, 2, 64, 128), # K != V - ], - ) - def test_forward_parity(self, B, T, H, K, V): - q, k, v, g, beta = _inputs(B, T, H, H, K, V) - scale = 1.0 / math.sqrt(K) - o, _ = kimi_delta_attention(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), scale=scale) - with torch.no_grad(): - o_ref, _ = kda_reference(q, k, v, g, beta, scale=scale) - assert rms_ratio(o.view_as(o_ref), o_ref) < FWD_TOL[torch.bfloat16] - - def test_forward_parity_gva(self, B=1, T=192, H=2, HV=4, K=64, V=64): - q, k, v, g, beta = _inputs(B, T, H, HV, K, V) - o, fs = kimi_delta_attention(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), output_final_state=True) - with torch.no_grad(): - o_ref, fs_ref = kda_reference(q, k, v, g, beta) - assert rms_ratio(o.view_as(o_ref), o_ref) < FWD_TOL[torch.bfloat16] - assert rms_ratio(fs, fs_ref) < STATE_TOL[torch.bfloat16] - - @pytest.mark.parametrize("B,T,H,K,V", [(1, 128, 1, 64, 64), (2, 128, 2, 64, 64)]) - def test_backward_runs(self, B, T, H, K, V): - q, k, v, g, beta = (_thd(x).detach().requires_grad_(True) for x in _inputs(B, T, H, H, K, V)) - o, _ = kimi_delta_attention(q, k, v, g, beta, _cu(B, T)) - o.sum().backward() - for name, t in [("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta)]: - assert t.grad is not None, f"no grad for {name}" - assert torch.isfinite(t.grad).all(), f"non-finite grad for {name}" - - def test_default_scale(self): - q, k, v, g, beta = (_thd(x) for x in _inputs(1, 128, 1, 1, 64, 64)) - cu = _cu(1, 128) - o_default, _ = kimi_delta_attention(q, k, v, g, beta, cu) - o_explicit, _ = kimi_delta_attention(q, k, v, g, beta, cu, scale=1.0 / math.sqrt(64)) - torch.testing.assert_close(o_default, o_explicit) - - def test_no_final_state_returns_empty(self): - q, k, v, g, beta = (_thd(x) for x in _inputs(1, 128, 1, 1, 64, 64)) - _o, final = kimi_delta_attention(q, k, v, g, beta, _cu(1, 128), output_final_state=False) - assert final.numel() == 0 - - def test_initial_state(self): - B, T, H, K, V = 1, 128, 2, 64, 64 - q, k, v, g, beta = _inputs(B, T, H, H, K, V) - S0 = torch.randn(B, H, K, V, dtype=torch.float32, device="cuda") * 0.05 - o, fs = kimi_delta_attention(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), initial_state=S0, output_final_state=True) - with torch.no_grad(): - o_ref, fs_ref = kda_reference(q, k, v, g, beta, initial_state=S0) - assert rms_ratio(o.view_as(o_ref), o_ref) < FWD_TOL[torch.bfloat16] - assert rms_ratio(fs, fs_ref) < STATE_TOL[torch.bfloat16] - - def test_packed_matches_per_sequence(self): - B, T, H, K, V = 2, 128, 2, 64, 64 - q, k, v, g, beta = _inputs(B, T, H, H, K, V) - o, fs = kimi_delta_attention(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), output_final_state=True) - for b in range(B): - o_b, fs_b = kimi_delta_attention(q[b], k[b], v[b], g[b], beta[b], _cu(1, T), output_final_state=True) - torch.testing.assert_close(o[b * T : (b + 1) * T], o_b) - torch.testing.assert_close(fs[b], fs_b[0]) - - def test_thd_ragged_parity_and_backward(self): - seq_lens = [64, 192] - H, K, V = 2, 64, 64 - total = sum(seq_lens) - q, k, v, g, beta = (x.squeeze(0).detach().requires_grad_(True) for x in _inputs(1, total, H, H, K, V)) - cu = torch.tensor([0, 64, 256], dtype=torch.int32, device="cuda") - - o, _ = kimi_delta_attention(q, k, v, g, beta, cu_seqlens=cu) - with torch.no_grad(): - o_ref, _ = kda_reference(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), g.unsqueeze(0), beta.unsqueeze(0), cu_seqlens=cu) - assert rms_ratio(o, o_ref.squeeze(0)) < FWD_TOL[torch.bfloat16] - - o.sum().backward() - for name, t in [("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta)]: - assert t.grad is not None and torch.isfinite(t.grad).all(), f"bad grad for {name}" - - def test_torch_compile_forward(self): - q, k, v, g, beta = (_thd(x) for x in _inputs(1, 128, 1, 1, 64, 64)) - cu = _cu(1, 128) - compiled = torch.compile(kimi_delta_attention, fullgraph=True) - o_eager, _ = kimi_delta_attention(q, k, v, g, beta, cu) - o_comp, _ = compiled(q, k, v, g, beta, cu) - torch.testing.assert_close(o_eager, o_comp) - - def test_forward_parity_qk_l2norm(self, B=1, T=256, H=2, K=128, V=128): - """D=128 + use_qk_l2norm: routes to KdaFrostEngine on SM100/SM103 (the - cuTile engine serves it elsewhere — both honor the in-kernel norm).""" - import torch.nn.functional as F - - q, k, v, g, beta = _inputs(B, T, H, H, K, V) - # the FROST BT=16 kernel's Neumann inverse wants stronger decay and a - # post-sigmoid beta (see its kernel suite) - g = torch.empty(B, T, H, K, device="cuda").uniform_(0.5, 1.0).log() - beta = beta.float().sigmoid().to(beta.dtype) - scale = 1.0 / math.sqrt(K) - o, fs = kimi_delta_attention( - _thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), scale=scale, output_final_state=True, use_qk_l2norm_in_kernel=True - ) - with torch.no_grad(): - o_ref, fs_ref = kda_reference(F.normalize(q.float(), dim=-1), F.normalize(k.float(), dim=-1), v, g, beta, scale=scale) - torch.testing.assert_close(o.view_as(o_ref).float(), o_ref.float(), atol=1e-1, rtol=1e-1) - assert rms_ratio(fs, fs_ref) < 5e-2 - - def test_forward_parity_no_l2norm_frost(self, B=1, T=256, H=2, K=128, V=128): - """D=128 without the in-kernel norm routes to KdaFrostEngine on - SM100/SM103, so q/k are pre-normalized here.""" - import torch.nn.functional as F - - q, k, v, g, beta = _inputs(B, T, H, H, K, V) - g = torch.empty(B, T, H, K, device="cuda").uniform_(0.5, 1.0).log() - beta = beta.float().sigmoid().to(beta.dtype) - q = F.normalize(q.float(), dim=-1).to(q.dtype) - k = F.normalize(k.float(), dim=-1).to(k.dtype) - scale = 1.0 / math.sqrt(K) - o, fs = kimi_delta_attention(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), scale=scale, output_final_state=True) - with torch.no_grad(): - o_ref, fs_ref = kda_reference(q.float(), k.float(), v, g, beta, scale=scale) - torch.testing.assert_close(o.view_as(o_ref).float(), o_ref.float(), atol=1e-1, rtol=1e-1) - assert rms_ratio(fs, fs_ref) < 5e-2 - - def test_backward_runs_qk_l2norm(self, B=1, T=128, H=2, K=128, V=128): - """D=128 + use_qk_l2norm: forward may run on KdaFrostEngine, backward - always falls back to the cuTile engine (FROST KDA_BWD is a stub).""" - q, k, v, g, beta = (_thd(x).detach().requires_grad_(True) for x in _inputs(B, T, H, H, K, V)) - o, _ = kimi_delta_attention(q, k, v, g, beta, _cu(B, T), use_qk_l2norm_in_kernel=True) - o.sum().backward() - for name, t in [("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta)]: - assert t.grad is not None, f"no grad for {name}" - assert torch.isfinite(t.grad).all(), f"non-finite grad for {name}" - - def test_thd_ragged_parity_frost(self): - """D=128 ragged varlen: the FROST path on SM100/SM103 (cuTile elsewhere).""" - import torch.nn.functional as F - - seq_lens = [64, 192] - total, H, K, V = sum(seq_lens), 2, 128, 128 - q, k, v, g, beta = _inputs(1, total, H, H, K, V) - g = torch.empty(1, total, H, K, device="cuda").uniform_(0.5, 1.0).log() - beta = beta.float().sigmoid().to(beta.dtype) - cu = torch.tensor([0, 64, 256], dtype=torch.int32, device="cuda") - o, fs = kimi_delta_attention(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), cu, output_final_state=True, use_qk_l2norm_in_kernel=True) - with torch.no_grad(): - o_ref, fs_ref = kda_reference(F.normalize(q.float(), dim=-1), F.normalize(k.float(), dim=-1), v, g, beta, cu_seqlens=cu) - torch.testing.assert_close(o.view_as(o_ref.squeeze(0)).float(), o_ref.squeeze(0).float(), atol=1e-1, rtol=1e-1) - assert rms_ratio(fs, fs_ref) < 5e-2 - - def test_initial_state_frost(self, B=1, T=256, H=2, K=128, V=128): - """D=128 + initial state on the FROST path.""" - import torch.nn.functional as F - - q, k, v, g, beta = _inputs(B, T, H, H, K, V) - g = torch.empty(B, T, H, K, device="cuda").uniform_(0.5, 1.0).log() - beta = beta.float().sigmoid().to(beta.dtype) - S0 = torch.randn(B, H, K, V, dtype=torch.float32, device="cuda") * 0.05 - o, fs = kimi_delta_attention( - _thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(B, T), initial_state=S0, output_final_state=True, use_qk_l2norm_in_kernel=True - ) - with torch.no_grad(): - o_ref, fs_ref = kda_reference(F.normalize(q.float(), dim=-1), F.normalize(k.float(), dim=-1), v, g, beta, initial_state=S0) - torch.testing.assert_close(o.view_as(o_ref).float(), o_ref.float(), atol=1e-1, rtol=1e-1) - assert rms_ratio(fs, fs_ref) < 5e-2 - - def test_frost_engine_selected(self): - """On SM100/SM103 with the DSL runtime, D=128 + l2norm graphs must - actually lower to KdaFrostEngine (guards the router ranking).""" - if not _frost_kda_available(): - pytest.skip("needs an SM100-class GPU and the Cutlass DSL KDA prefill kernel runtime") - from cudnn.linear_attention.frost.kda_engine import KdaFrostEngine - from cudnn.linear_attention.ops import kda as kda_ops - - kda_ops._fwd_graph_cache.clear() - q, k, v, g, beta = _inputs(1, 128, 2, 2, 128, 128) - g = torch.empty(1, 128, 2, 128, device="cuda").uniform_(0.5, 1.0).log() - beta = beta.float().sigmoid().to(beta.dtype) - kimi_delta_attention(_thd(q), _thd(k), _thd(v), _thd(g), _thd(beta), _cu(1, 128), use_qk_l2norm_in_kernel=True) - assert any(isinstance(graph.selected_engine, KdaFrostEngine) for graph, _t in kda_ops._fwd_graph_cache.values()) diff --git a/test/python/linear_attention/reference_gdn.py b/test/python/linear_attention/reference_gdn.py index a828f58d6..4df5f1120 100644 --- a/test/python/linear_attention/reference_gdn.py +++ b/test/python/linear_attention/reference_gdn.py @@ -9,8 +9,8 @@ S_t = alpha_t (I - beta_t k_t^T k_t) S_{t-1} + beta_t k_t^T v_t o_t = q_t S_t -with scalar per-token decay ``alpha_t = exp(g_t)`` and write strength -``beta_t``. Supports grouped heads (every input's head count must divide +where ``S_t`` is the recurrent state, with scalar per-token decay +``alpha_t = exp(g_t)`` and write strength ``beta_t``. Supports grouped heads (every input's head count must divide ``HO = max(Hq, Hv)``; heads are replicated onto the HO output heads), an optional initial state, and varlen packed batches via ``cu_seqlens``. @@ -33,21 +33,21 @@ def rms_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: return ((out - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-12)).item() -def _recurrent_dense(q, k, v, alpha, beta, S0): - """Dense recurrence in [B, HV, T, *] layout, fp64. Returns (o, S_T).""" +def recurrent_dense(q, k, v, alpha, beta, state0): + """Dense recurrence in [B, HV, T, *] layout, fp64. Returns (o, final state).""" T = q.shape[2] - S = S0 + state = state0 outs = [] for t in range(T): kt = k[:, :, t, :] vt = v[:, :, t, :] at = alpha[:, :, t] bt = beta[:, :, t] - kt_S = (kt.unsqueeze(-2) @ S).squeeze(-2) - residual = vt - at[..., None] * kt_S - S = at[..., None, None] * S + bt[..., None, None] * (kt.unsqueeze(-1) @ residual.unsqueeze(-2)) - outs.append((q[:, :, t, :].unsqueeze(-2) @ S).squeeze(-2)) - return torch.stack(outs, dim=2), S + kt_state = (kt.unsqueeze(-2) @ state).squeeze(-2) + residual = vt - at[..., None] * kt_state + state = at[..., None, None] * state + bt[..., None, None] * (kt.unsqueeze(-1) @ residual.unsqueeze(-2)) + outs.append((q[:, :, t, :].unsqueeze(-2) @ state).squeeze(-2)) + return torch.stack(outs, dim=2), state def gdn_reference( @@ -80,15 +80,22 @@ def gdn_reference( HO = max(q.shape[2], v.shape[2]) - def expand(x): - r = HO // x.shape[2] - return x.repeat_interleave(r, dim=2) if r > 1 else x - - qf = expand(q.double() * scale) - kf = expand(k.double()) - vf = expand(v.double()) - alphaf = expand(g.double().exp()) - betaf = expand(beta.double()) + qf = q.double() * scale + kf = k.double() + vf = v.double() + alphaf = g.double().exp() + betaf = beta.double() + # expand tensors for grouped heads (view, no copy), as in the sdpa references + if q.shape[2] != HO: + qf = qf.unsqueeze(3).expand(-1, -1, -1, HO // q.shape[2], -1).reshape(q.shape[0], q.shape[1], HO, -1) + if k.shape[2] != HO: + kf = kf.unsqueeze(3).expand(-1, -1, -1, HO // k.shape[2], -1).reshape(k.shape[0], k.shape[1], HO, -1) + if v.shape[2] != HO: + vf = vf.unsqueeze(3).expand(-1, -1, -1, HO // v.shape[2], -1).reshape(v.shape[0], v.shape[1], HO, -1) + if g.shape[2] != HO: + alphaf = alphaf.unsqueeze(3).expand(-1, -1, -1, HO // g.shape[2]).reshape(g.shape[0], g.shape[1], HO) + if beta.shape[2] != HO: + betaf = betaf.unsqueeze(3).expand(-1, -1, -1, HO // beta.shape[2]).reshape(beta.shape[0], beta.shape[1], HO) HV = HO # [B, T, HV, *] -> [B, HV, T, *] @@ -103,11 +110,11 @@ def expand(x): if cu_seqlens is None: B = q.shape[0] if initial_state is None: - S0 = torch.zeros(B, HV, K, V, dtype=torch.float64, device=q.device) + state0 = torch.zeros(B, HV, K, V, dtype=torch.float64, device=q.device) else: - S0 = initial_state.double() - o, S = _recurrent_dense(qf, kf, vf, alphaf, betaf, S0) - return o.permute(0, 2, 1, 3), S + state0 = initial_state.double() + o, state = recurrent_dense(qf, kf, vf, alphaf, betaf, state0) + return o.permute(0, 2, 1, 3), state assert q.shape[0] == 1, "cu_seqlens requires packed batch B == 1" bounds = cu_seqlens.tolist() @@ -115,15 +122,15 @@ def expand(x): for n in range(len(bounds) - 1): s, e = bounds[n], bounds[n + 1] if initial_state is None: - S0 = torch.zeros(1, HV, K, V, dtype=torch.float64, device=q.device) + state0 = torch.zeros(1, HV, K, V, dtype=torch.float64, device=q.device) else: - S0 = initial_state[n : n + 1].double() + state0 = initial_state[n : n + 1].double() if e == s: - states.append(S0) + states.append(state0) continue - o_n, S_n = _recurrent_dense(qf[:, :, s:e], kf[:, :, s:e], vf[:, :, s:e], alphaf[:, :, s:e], betaf[:, :, s:e], S0) + o_n, state_n = recurrent_dense(qf[:, :, s:e], kf[:, :, s:e], vf[:, :, s:e], alphaf[:, :, s:e], betaf[:, :, s:e], state0) outs.append(o_n) - states.append(S_n) + states.append(state_n) if outs: o = torch.cat(outs, dim=2).permute(0, 2, 1, 3) else: diff --git a/test/python/linear_attention/reference_gdn2.py b/test/python/linear_attention/reference_gdn2.py index b8d9ecf26..f91e599a4 100644 --- a/test/python/linear_attention/reference_gdn2.py +++ b/test/python/linear_attention/reference_gdn2.py @@ -6,8 +6,8 @@ GDN-2 generalizes GDN's scalar gates to three channel-wise gates: a per-key decay ``alpha_t = exp(g_t) in (0, 1]^K``, a per-key erase gate ``beta_t in -R^K``, and a NEW per-value write gate ``w_t in R^V`` (``k``/``q`` already -feature-mapped, ``q`` pre-scaled): +R^K``, and a NEW per-value write gate ``w_t in R^V`` (``S_t`` is the +recurrent state; ``k``/``q`` already feature-mapped, ``q`` pre-scaled): S_t = (I - k_t (beta_t . k_t)^T) Diag(alpha_t) S_{t-1} + k_t (w_t . v_t)^T o_t = q_t S_t @@ -27,8 +27,8 @@ 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; the kernel keeps it V-major -``[N, HO, V, K]`` (transpose at the boundary). +kept K-major ``[N, HO, K, V]`` here, matching the kernel ABI (KV, v +contiguous). """ from __future__ import annotations @@ -46,12 +46,12 @@ def rms_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: return ((out - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-12)).item() -def _recurrent_dense(q, k, v, alpha, beta, w, S0): - """Dense recurrence in [B, HV, T, *] layout, fp64. Returns (o, S_T). +def recurrent_dense(q, k, v, alpha, beta, w, state0): + """Dense recurrence in [B, HV, T, *] layout, fp64. Returns (o, final state). alpha, beta: per-key-channel [B, HV, T, K]; w: per-value-channel [B, HV, T, V].""" T = q.shape[2] - S = S0 + state = state0 outs = [] for t in range(T): kt = k[:, :, t, :] # [B, HV, K] @@ -59,12 +59,12 @@ def _recurrent_dense(q, k, v, alpha, beta, w, S0): at = alpha[:, :, t, :] # [B, HV, K] (= exp(g_t)) bt = beta[:, :, t, :] # [B, HV, K] wt = w[:, :, t, :] # [B, HV, V] - S = at[..., None] * S # per-K-channel decay first - erase = ((bt * kt).unsqueeze(-2) @ S).squeeze(-2) # (beta . k)^T S_dec: [B, HV, V] + state = at[..., None] * state # per-K-channel decay first + erase = ((bt * kt).unsqueeze(-2) @ state).squeeze(-2) # (beta . k)^T on the decayed state: [B, HV, V] v_new = wt * vt - erase - S = S + kt.unsqueeze(-1) @ v_new.unsqueeze(-2) # k (x) v_new - outs.append((q[:, :, t, :].unsqueeze(-2) @ S).squeeze(-2)) - return torch.stack(outs, dim=2), S + state = state + kt.unsqueeze(-1) @ v_new.unsqueeze(-2) # k (x) v_new + outs.append((q[:, :, t, :].unsqueeze(-2) @ state).squeeze(-2)) + return torch.stack(outs, dim=2), state def gdn2_reference( @@ -101,16 +101,25 @@ def gdn2_reference( HO = max(q.shape[2], v.shape[2]) - def expand(x): - r = HO // x.shape[2] - return x.repeat_interleave(r, dim=2) if r > 1 else x - - qf = expand(q.double() * scale) - kf = expand(k.double()) - vf = expand(v.double()) - alphaf = expand(g.double().exp()) # [B, T, HO, K] - betaf = expand(beta.double()) # [B, T, HO, K] - wf = expand(w.double()) # [B, T, HO, V] + qf = q.double() * scale + kf = k.double() + vf = v.double() + alphaf = g.double().exp() # [B, T, HO, K] + betaf = beta.double() # [B, T, HO, K] + wf = w.double() # [B, T, HO, V] + # expand tensors for grouped heads (view, no copy), as in the sdpa references + if q.shape[2] != HO: + qf = qf.unsqueeze(3).expand(-1, -1, -1, HO // q.shape[2], -1).reshape(q.shape[0], q.shape[1], HO, -1) + if k.shape[2] != HO: + kf = kf.unsqueeze(3).expand(-1, -1, -1, HO // k.shape[2], -1).reshape(k.shape[0], k.shape[1], HO, -1) + if v.shape[2] != HO: + vf = vf.unsqueeze(3).expand(-1, -1, -1, HO // v.shape[2], -1).reshape(v.shape[0], v.shape[1], HO, -1) + if g.shape[2] != HO: + alphaf = alphaf.unsqueeze(3).expand(-1, -1, -1, HO // g.shape[2], -1).reshape(g.shape[0], g.shape[1], HO, -1) + if beta.shape[2] != HO: + betaf = betaf.unsqueeze(3).expand(-1, -1, -1, HO // beta.shape[2], -1).reshape(beta.shape[0], beta.shape[1], HO, -1) + if w.shape[2] != HO: + wf = wf.unsqueeze(3).expand(-1, -1, -1, HO // w.shape[2], -1).reshape(w.shape[0], w.shape[1], HO, -1) # [B, T, HO, *] -> [B, HO, T, *] qf = qf.permute(0, 2, 1, 3) @@ -126,11 +135,11 @@ def expand(x): if cu_seqlens is None: B = q.shape[0] if initial_state is None: - S0 = torch.zeros(B, HV, K, V, dtype=torch.float64, device=q.device) + state0 = torch.zeros(B, HV, K, V, dtype=torch.float64, device=q.device) else: - S0 = initial_state.double() - o, S = _recurrent_dense(qf, kf, vf, alphaf, betaf, wf, S0) - return o.permute(0, 2, 1, 3), S + state0 = initial_state.double() + o, state = recurrent_dense(qf, kf, vf, alphaf, betaf, wf, state0) + return o.permute(0, 2, 1, 3), state assert q.shape[0] == 1, "cu_seqlens requires packed batch B == 1" bounds = cu_seqlens.tolist() @@ -138,15 +147,15 @@ def expand(x): for n in range(len(bounds) - 1): s, e = bounds[n], bounds[n + 1] if initial_state is None: - S0 = torch.zeros(1, HV, K, V, dtype=torch.float64, device=q.device) + state0 = torch.zeros(1, HV, K, V, dtype=torch.float64, device=q.device) else: - S0 = initial_state[n : n + 1].double() + state0 = initial_state[n : n + 1].double() if e == s: - states.append(S0) + states.append(state0) continue - o_n, S_n = _recurrent_dense(qf[:, :, s:e], kf[:, :, s:e], vf[:, :, s:e], alphaf[:, :, s:e], betaf[:, :, s:e], wf[:, :, s:e], S0) + o_n, state_n = recurrent_dense(qf[:, :, s:e], kf[:, :, s:e], vf[:, :, s:e], alphaf[:, :, s:e], betaf[:, :, s:e], wf[:, :, s:e], state0) outs.append(o_n) - states.append(S_n) + states.append(state_n) if outs: o = torch.cat(outs, dim=2).permute(0, 2, 1, 3) else: diff --git a/test/python/linear_attention/reference_kda.py b/test/python/linear_attention/reference_kda.py index 5e8f5ffd3..28ed0e9db 100644 --- a/test/python/linear_attention/reference_kda.py +++ b/test/python/linear_attention/reference_kda.py @@ -12,9 +12,9 @@ S_t = (I - beta_t k_t^T k_t) Diag(alpha_t) S_{t-1} + beta_t k_t^T v_t o_t = q_t S_t -with scalar per-token write strength ``beta_t``. The decay is applied FIRST -(``S_dec = Diag(alpha_t) S_{t-1}``) and the delta-rule correction reads the -already-decayed state, so unlike GDN the order matters. Supports grouped +where ``S_t`` is the recurrent state and ``beta_t`` the scalar per-token +write strength. The decay is applied first and the delta-rule correction +reads the already-decayed state, so unlike GDN the order matters. Supports grouped heads (every input's head count must divide ``HO = max(Hq, Hv)``; heads are replicated onto the HO output heads), an optional initial state, and varlen packed batches via ``cu_seqlens``. @@ -38,24 +38,24 @@ def rms_ratio(out: torch.Tensor, ref: torch.Tensor) -> float: return ((out - ref).pow(2).mean().sqrt() / ref.pow(2).mean().sqrt().clamp_min(1e-12)).item() -def _recurrent_dense(q, k, v, alpha, beta, S0): - """Dense recurrence in [B, HV, T, *] layout, fp64. Returns (o, S_T). +def recurrent_dense(q, k, v, alpha, beta, state0): + """Dense recurrence in [B, HV, T, *] layout, fp64. Returns (o, final state). ``alpha`` is per-key-channel: [B, HV, T, K] (GDN's is scalar [B, HV, T]).""" T = q.shape[2] - S = S0 + state = state0 outs = [] for t in range(T): kt = k[:, :, t, :] vt = v[:, :, t, :] at = alpha[:, :, t, :] # [B, HV, K] per-channel decay bt = beta[:, :, t] - S = at[..., None] * S # decay first: Diag(alpha) S, per-K-row scaling - kt_S = (kt.unsqueeze(-2) @ S).squeeze(-2) - residual = vt - kt_S # reads the already-decayed state (no scalar alpha here) - S = S + bt[..., None, None] * (kt.unsqueeze(-1) @ residual.unsqueeze(-2)) - outs.append((q[:, :, t, :].unsqueeze(-2) @ S).squeeze(-2)) - return torch.stack(outs, dim=2), S + state = at[..., None] * state # decay first: Diag(alpha) state, per-K-row scaling + kt_state = (kt.unsqueeze(-2) @ state).squeeze(-2) + residual = vt - kt_state # reads the already-decayed state (no scalar alpha here) + state = state + bt[..., None, None] * (kt.unsqueeze(-1) @ residual.unsqueeze(-2)) + outs.append((q[:, :, t, :].unsqueeze(-2) @ state).squeeze(-2)) + return torch.stack(outs, dim=2), state def kda_reference( @@ -89,15 +89,23 @@ def kda_reference( HO = max(q.shape[2], v.shape[2]) - def expand(x): - r = HO // x.shape[2] - return x.repeat_interleave(r, dim=2) if r > 1 else x - - qf = expand(q.double() * scale) - kf = expand(k.double()) - vf = expand(v.double()) - alphaf = expand(g.double().exp()) # [B, T, HO, K] - betaf = expand(beta.double()) + qf = q.double() * scale + kf = k.double() + vf = v.double() + alphaf = g.double().exp() # [B, T, HO, K] + betaf = beta.double() + # expand tensors for grouped heads (view, no copy), as in the sdpa references + if q.shape[2] != HO: + qf = qf.unsqueeze(3).expand(-1, -1, -1, HO // q.shape[2], -1).reshape(q.shape[0], q.shape[1], HO, -1) + if k.shape[2] != HO: + kf = kf.unsqueeze(3).expand(-1, -1, -1, HO // k.shape[2], -1).reshape(k.shape[0], k.shape[1], HO, -1) + if v.shape[2] != HO: + vf = vf.unsqueeze(3).expand(-1, -1, -1, HO // v.shape[2], -1).reshape(v.shape[0], v.shape[1], HO, -1) + if g.shape[2] != HO: + alphaf = alphaf.unsqueeze(3).expand(-1, -1, -1, HO // g.shape[2], -1).reshape(g.shape[0], g.shape[1], HO, -1) + if beta.shape[2] != HO: + # direct-reference callers only: the op and the kernels take beta at HO heads + betaf = betaf.unsqueeze(3).expand(-1, -1, -1, HO // beta.shape[2]).reshape(beta.shape[0], beta.shape[1], HO) # [B, T, HO, *] -> [B, HO, T, *] qf = qf.permute(0, 2, 1, 3) @@ -112,11 +120,11 @@ def expand(x): if cu_seqlens is None: B = q.shape[0] if initial_state is None: - S0 = torch.zeros(B, HV, K, V, dtype=torch.float64, device=q.device) + state0 = torch.zeros(B, HV, K, V, dtype=torch.float64, device=q.device) else: - S0 = initial_state.double() - o, S = _recurrent_dense(qf, kf, vf, alphaf, betaf, S0) - return o.permute(0, 2, 1, 3), S + state0 = initial_state.double() + o, state = recurrent_dense(qf, kf, vf, alphaf, betaf, state0) + return o.permute(0, 2, 1, 3), state assert q.shape[0] == 1, "cu_seqlens requires packed batch B == 1" bounds = cu_seqlens.tolist() @@ -124,15 +132,15 @@ def expand(x): for n in range(len(bounds) - 1): s, e = bounds[n], bounds[n + 1] if initial_state is None: - S0 = torch.zeros(1, HV, K, V, dtype=torch.float64, device=q.device) + state0 = torch.zeros(1, HV, K, V, dtype=torch.float64, device=q.device) else: - S0 = initial_state[n : n + 1].double() + state0 = initial_state[n : n + 1].double() if e == s: - states.append(S0) + states.append(state0) continue - o_n, S_n = _recurrent_dense(qf[:, :, s:e], kf[:, :, s:e], vf[:, :, s:e], alphaf[:, :, s:e], betaf[:, :, s:e], S0) + o_n, state_n = recurrent_dense(qf[:, :, s:e], kf[:, :, s:e], vf[:, :, s:e], alphaf[:, :, s:e], betaf[:, :, s:e], state0) outs.append(o_n) - states.append(S_n) + states.append(state_n) if outs: o = torch.cat(outs, dim=2).permute(0, 2, 1, 3) else: diff --git a/test/python/linear_attention/test_la.py b/test/python/linear_attention/test_la.py new file mode 100644 index 000000000..271f8a92a --- /dev/null +++ b/test/python/linear_attention/test_la.py @@ -0,0 +1,1173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Backend-parametrized test suite for the linear-attention ops. + +One suite for the public API (``gated_delta_net`` / ``kimi_delta_attention`` / +``gated_delta_net_v2``), parametrized over the python engine backends that can +serve it. Each test pins one backend's plan onto the ops' graphs (plan-API +``select_plan`` by name) and validates against the fp64 recurrent references. A pinned engine that declines a +configuration waives the test (``cudnnGraphNotSupportedError`` -> skip), so +the support surface is owned by the engines' ``check_support``, not by a +suite-side matrix; a backend that is not installed skips the same way. + +Determinism soaks and CUDA-graph replay run on the backends that promise +those contracts (currently FROST only). +""" + +from __future__ import annotations + +import contextlib +import functools +import math +import pytest + +torch = pytest.importorskip("torch") +cudnn = pytest.importorskip("cudnn") +la_ops = pytest.importorskip("cudnn.linear_attention.ops") + +import torch.nn.functional as F # noqa: E402 + +from .conftest import gen_qkv # noqa: E402 +from .reference_gdn import gdn_reference, rms_ratio # noqa: E402 +from .reference_gdn2 import gdn2_reference # noqa: E402 +from .reference_kda import kda_reference # noqa: E402 + +pytestmark = [ + pytest.mark.L0, + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"), +] + +VARIANTS = ("gdn", "kda", "gdn2") +CHUNK = {"gdn": 64, "kda": 16, "gdn2": 16} + +FWD_TOL = {torch.bfloat16: 2e-2, torch.float16: 1e-2} +STATE_TOL = {torch.bfloat16: 2e-2, torch.float16: 1e-2} +BWD_TOL = {torch.bfloat16: 4e-2, torch.float16: 3e-2} +STATE_GRAD_TOL = 6e-2 + +# (H, HV) pairs: H = Q/K heads, HV = V heads; gates/O/states live at HO = max. +HEAD_CONFIGS = [(1, 1), (3, 3), (1, 2), (2, 4), (16, 32), (16, 64)] +HEAD_CONFIGS_SMALL = [(1, 1), (2, 4)] +GQA_CONFIGS = [(4, 4, 1), (6, 6, 2), (4, 1, 1), (6, 2, 2), (1, 2, 2), (2, 4, 4)] + +RAGGED_SEQ_LENS = [ + [256, 256], + [511, 501], + [64, 128, 512], + [31, 63, 93, 123, 150, 500], + [7] * 24 + [1] * 8, + [2048], +] +EDGE_LENS = [1, 15, 16, 17, 31, 63, 64, 65, 121, 251, 257] + +DETERMINISM_REPEATS = 8 +SEED = 888 + +DTYPE_IDS = {torch.bfloat16: "bf16", torch.float16: "fp16"} + + +# --------------------------------------------------------------------------- +# Backend pinning +# --------------------------------------------------------------------------- + + +def op_modules(): + from cudnn.linear_attention.ops import gdn, gdn2, kda + + return {"gdn": gdn, "kda": kda, "gdn2": gdn2} + + +def family_engines(backend_name): + """The backend's engine instances per family, ids assigned by the manifest.""" + from cudnn.engines import manifest + + suffix = "_" + backend_name + families = {} + for family in manifest.MANIFEST: + if family.name in VARIANTS: + engines = manifest.instantiate(family, family.offered_ids()) + families[family.name] = [e for e in engines if e.name.endswith(suffix)] + return families + + +def clear_op_caches(): + for mod in op_modules().values(): + mod._fprop_cache.clear() + mod._bprop_cache.clear() + + +class Case: + """One test configuration: inputs, gates, cu_seqlens and the geometry.""" + + __slots__ = ("variant", "dtype", "q", "k", "v", "gates", "cu", "B", "T", "N", "H", "HK", "HV", "HO", "K", "V", "varlen") + + def __init__(self, **fields): + for name in self.__slots__: + setattr(self, name, fields.pop(name)) + assert not fields, f"unknown Case fields: {sorted(fields)}" + + def clone(self, **overrides): + fields = {name: getattr(self, name) for name in self.__slots__} + fields.update(overrides) + return Case(**fields) + + +class Backend: + """The pinned backend: its name, its engine instances per family, and the + plan name a graph must offer (``_``, e.g. ``gdn_frost``).""" + + __slots__ = ("name", "engines") + + def __init__(self, name, engines): + self.name = name + self.engines = engines + + def plan(self, variant): + return f"{variant}_{self.name}" + + +def pinned_op(backend, variant): + """The variant's op with the backend's plan pinned (ops-level + ``plan_name``, the examples' ``select_plan`` paradigm).""" + return functools.partial(op(variant), plan_name=backend.plan(variant)) + + +@pytest.fixture(params=("frost", "cutile")) +def backend(request): + """One backend per run of each test; the tests pass its plan name to the + ops. The op graph caches are cleared around each test so the pin + assertions only ever see this test's graphs.""" + name = request.param + families = family_engines(name) + missing = [v for v in VARIANTS if v not in families] + if missing: + pytest.fail(f"the engine manifest offers no {missing} families — a stale installed cudnn package is likely shadowing the source tree") + clear_op_caches() + try: + yield Backend(name, families) + finally: + clear_op_caches() + + +@contextlib.contextmanager +def waive_unsupported(backend, variant): + """A backend with no engine for the family, or an engine decline (no plan + offered / check_support raise), waives the test; the engines own the + support surface.""" + if not backend.engines[variant]: + pytest.skip(f"the {backend.name} backend has no {variant} engine") + try: + yield + except cudnn.cudnnGraphNotSupportedError as exc: + pytest.skip(f"{backend.name} {variant} declined: {exc}") + + +# --------------------------------------------------------------------------- +# Case generation and dispatch +# --------------------------------------------------------------------------- + + +def set_seed(seed=SEED): + torch.random.manual_seed(seed) + torch.cuda.manual_seed(seed) + + +def gate_lo(dtype): + return 0.6 if dtype == torch.float16 else 0.5 + + +def gen_gates(variant, B, T, HO, K, V, dtype, *, alpha=True, beta=True, w=True, lo=None, device="cuda"): + if lo is None: + lo = gate_lo(dtype) + gshape = (B, T, HO) if variant == "gdn" else (B, T, HO, K) + if alpha: + g = torch.empty(gshape, device=device, dtype=torch.float32).uniform_(lo, 1.0).log() + else: + g = torch.zeros(gshape, device=device, dtype=torch.float32) + if variant == "gdn2": + b = (torch.rand(B, T, HO, K, device=device).sigmoid() * 2.0).to(dtype) if beta else torch.ones(B, T, HO, K, device=device, dtype=dtype) + wt = torch.rand(B, T, HO, V, device=device).sigmoid().to(dtype) if w else torch.ones(B, T, HO, V, device=device, dtype=dtype) + return {"g": g, "beta": b, "w": wt} + b = torch.rand(B, T, HO, device=device) if beta else torch.ones(B, T, HO, device=device) + return {"g": g, "beta": b} + + +def make_case(variant, dtype, *, B=1, T=None, seq_lens=None, H=2, HK=None, HV=None, K=128, V=128, alpha=True, beta=True, w=True, lo=None, seed=SEED): + """Dense ``(B, T)`` or packed varlen (``seq_lens``, B == 1) inputs plus the + matching ``cu_seqlens``. ``HK`` defaults to ``H``; ``HK == HV < H`` is + canonical (native grouped K) GQA.""" + set_seed(seed) + HV = H if HV is None else HV + HK = H if HK is None else HK + HO = max(H, HV) + if seq_lens is not None: + total = sum(seq_lens) + bounds = [0] + for sl in seq_lens: + bounds.append(bounds[-1] + sl) + cu = torch.tensor(bounds, dtype=torch.int32, device="cuda") + B, T, N, varlen = 1, total, len(seq_lens), True + else: + cu = torch.arange(0, B + 1, dtype=torch.int32, device="cuda") * T + N, varlen = B, False + q, k, v = gen_qkv(B, T, H, HV, K, V, dtype) + if HK != H: + from .conftest import multidist_randu + + k = F.normalize(multidist_randu(B * T * HK, K, device="cuda").reshape(B, T, HK, K), p=2.0, dim=-1).to(dtype).contiguous() + gates = gen_gates(variant, B, T, HO, K, V, dtype, alpha=alpha, beta=beta, w=w, lo=lo) + return Case(variant=variant, dtype=dtype, q=q, k=k, v=v, gates=gates, cu=cu, B=B, T=T, N=N, H=H, HK=HK, HV=HV, HO=HO, K=K, V=V, varlen=varlen) + + +def to_thd(x): + return x.reshape(-1, *x.shape[2:]) + + +def op(variant): + return {"gdn": la_ops.gated_delta_net, "kda": la_ops.kimi_delta_attention, "gdn2": la_ops.gated_delta_net_v2}[variant] + + +def op_args(case, cu=None): + args = [to_thd(case.q), to_thd(case.k), to_thd(case.v), to_thd(case.gates["g"]), to_thd(case.gates["beta"])] + if case.variant == "gdn2": + args.append(to_thd(case.gates["w"])) + args.append(case.cu if cu is None else cu) + return args + + +def run_fwd(backend, case, *, cu=None, **kw): + with waive_unsupported(backend, case.variant): + return pinned_op(backend, case.variant)(*op_args(case, cu=cu), **kw) + + +def reference(case, *, scale=None, initial_state=None, l2norm=False, cu=None): + fn = {"gdn": gdn_reference, "kda": kda_reference, "gdn2": gdn2_reference}[case.variant] + q, k = case.q, case.k + if l2norm: + q = F.normalize(q.float(), dim=-1) + k = F.normalize(k.float(), dim=-1) + args = [q, k, case.v, case.gates["g"], case.gates["beta"]] + if case.variant == "gdn2": + args.append(case.gates["w"]) + kwargs = dict(scale=scale, initial_state=initial_state) + if case.varlen or cu is not None: + kwargs["cu_seqlens"] = case.cu if cu is None else cu + with torch.no_grad(): + return fn(*args, **kwargs) + + +def check(name, out, ref, tol): + out = out.float() + assert torch.isfinite(out).all(), f"non-finite values in {name}" + r = rms_ratio(out.reshape(ref.shape), ref) + 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]) + if fs is not None and fs.numel(): + check("final_state", fs, fs_ref, tol_mult * STATE_TOL[case.dtype]) + + +# --------------------------------------------------------------------------- +# Backend pin seam +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_backend_pin_selects_engine(backend, variant): + """The pinned engine actually serves the graph — a dead pin would silently + validate whatever the default routing picks.""" + case = make_case(variant, torch.bfloat16, T=4 * CHUNK[variant]) + with waive_unsupported(backend, variant): + pinned_op(backend, variant)(*op_args(case)) + mod = op_modules()[variant] + names = {g.selected_engine.name for g, entry in mod._fprop_cache.values() if g.selected_engine is not None} + assert names == {f"{variant}_{backend.name}"}, f"expected only {variant}_{backend.name} to serve, got {names}" + + +# --------------------------------------------------------------------------- +# Forward parity +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("H,HV", HEAD_CONFIGS) +@pytest.mark.parametrize("B,T", [(1, 64), (1, 128), (2, 256)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=DTYPE_IDS.get) +@pytest.mark.parametrize("variant", VARIANTS) +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) + + +@pytest.mark.parametrize("alpha,beta,w", [(True, False, True), (False, True, True), (True, True, False)], ids=["no_beta", "no_alpha", "no_w"]) +@pytest.mark.parametrize("variant", VARIANTS) +def test_fwd_gate_combinations(backend, variant, alpha, beta, w): + if variant != "gdn" and not alpha: + pytest.skip("the delta-rule inverse needs decay (matches FI's use_g=False skip)") + 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) + + +@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) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_fwd_default_scale_matches_explicit(backend, variant): + case = make_case(variant, torch.bfloat16, T=128) + o_default, _ = run_fwd(backend, case) + o_explicit, _ = run_fwd(backend, case, scale=1.0 / math.sqrt(case.K)) + torch.testing.assert_close(o_default, o_explicit) + + +@pytest.mark.parametrize("T", EDGE_LENS) +@pytest.mark.parametrize("variant", VARIANTS) +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) + + +@pytest.mark.parametrize("H,HV", [(1, 1), (2, 4)]) +@pytest.mark.parametrize("seq_lens", RAGGED_SEQ_LENS, ids=lambda sl: f"{len(sl)}seqs_{sum(sl)}tok") +@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) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_fwd_many_short_sequences(backend, variant): + """A 200-sequence packed batch matches the same sequences run one by one.""" + T = 33 + case = make_case(variant, torch.bfloat16, seq_lens=[T] * 200) + 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(), + to_thd(case.k)[sl].clone(), + to_thd(case.v)[sl].clone(), + to_thd(case.gates["g"])[sl].clone(), + to_thd(case.gates["beta"])[sl].clone(), + ] + if case.variant == "gdn2": + 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]) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_fwd_zero_length_sequences(backend, variant): + """Empty sequences must not perturb their neighbors; their state rows stay + zero (or pass the initial state through when one is given).""" + case = make_case(variant, torch.bfloat16, seq_lens=[64, 128]) + o_base, fs_base = run_fwd(backend, case, output_final_state=True) + cu = torch.tensor([0, 64, 64, 192, 192], dtype=torch.int32, device="cuda") + o, fs = run_fwd(backend, case, cu=cu, output_final_state=True) + torch.testing.assert_close(o, o_base, atol=1e-3, rtol=1e-3) + 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 + 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) + + +@pytest.mark.parametrize("T", [128, 251]) +@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) + + +@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 + 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): + 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()} + sub.T = t1 - t0 + 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) + 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]) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_fwd_packed_matches_per_sequence(backend, variant): + B, T = 3, 128 + case = make_case(variant, torch.bfloat16, B=B, T=T) + o, fs = run_fwd(backend, case, output_final_state=True) + cu1 = torch.tensor([0, T], dtype=torch.int32, device="cuda") + for b in range(B): + args = [case.q[b], case.k[b], case.v[b], case.gates["g"][b], case.gates["beta"][b]] + if variant == "gdn2": + args.append(case.gates["w"][b]) + with waive_unsupported(backend, variant): + o_b, fs_b = pinned_op(backend, variant)(*args, cu1, output_final_state=True) + torch.testing.assert_close(o[b * T : (b + 1) * T], o_b) + torch.testing.assert_close(fs[b], fs_b[0]) + + +@pytest.mark.parametrize( + "variant,K,V", + [("gdn", 64, 64), ("gdn", 64, 128), ("gdn", 128, 128), ("gdn", 256, 128), ("kda", 64, 64), ("kda", 64, 128), ("kda", 128, 128), ("gdn2", 128, 128)], +) +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) + + +@pytest.mark.parametrize("H,HK,HV", GQA_CONFIGS) +@pytest.mark.parametrize("variant", VARIANTS) +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) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_fwd_multi_tile(backend, variant): + """B*H well above the SM count: each CTA walks several (b, h) tiles back + 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) + + +@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) + + +@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) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_fwd_output_contract(backend, variant): + """O is io-dtype at HO heads; final_state is fp32 and empty unless requested.""" + case = make_case(variant, torch.bfloat16, T=128, H=2, HV=4) + o, fs = run_fwd(backend, case) + 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 + + +# --------------------------------------------------------------------------- +# Backward parity (oracle: fp64 autograd through the references) +# --------------------------------------------------------------------------- + + +def assert_bwd_parity(backend, case, *, scale=None, use_initial_state=False, use_dfs=False, l2norm=False, gate_grad_tol=None, seed=SEED + 1): + variant, tol = case.variant, BWD_TOL[case.dtype] + tensors = {"q": case.q, "k": case.k, "v": case.v, "g": case.gates["g"], "beta": case.gates["beta"]} + if variant == "gdn2": + tensors["w"] = case.gates["w"] + op_leaves = {n: to_thd(t).detach().clone().requires_grad_(True) for n, t in tensors.items()} + ref_leaves = {n: t.detach().double().requires_grad_(True) for n, t in tensors.items()} + 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_op = state0.detach().clone().requires_grad_(True) + state0_ref = state0.detach().double().requires_grad_(True) + + with waive_unsupported(backend, variant): + args = [op_leaves["q"], op_leaves["k"], op_leaves["v"], op_leaves["g"], op_leaves["beta"]] + if variant == "gdn2": + args.append(op_leaves["w"]) + args.append(case.cu) + o, fs = pinned_op(backend, variant)(*args, scale=scale, initial_state=state0_op, output_final_state=True, use_qk_l2norm_in_kernel=l2norm) + dO = torch.randn_like(o) + outputs, grad_outputs = [o], [dO] + dFS = None + if use_dfs: + dFS = torch.randn_like(fs) * 0.1 + outputs.append(fs) + grad_outputs.append(dFS) + grad_inputs = list(op_leaves.values()) + ([state0_op] if use_initial_state else []) + grads = torch.autograd.grad(outputs, grad_inputs, grad_outputs) + + qd, kd = ref_leaves["q"], ref_leaves["k"] + if l2norm: + qd, kd = F.normalize(qd, dim=-1), F.normalize(kd, dim=-1) + ref_fn = {"gdn": gdn_reference, "kda": kda_reference, "gdn2": gdn2_reference}[variant] + ref_args = [qd, kd, ref_leaves["v"], ref_leaves["g"], ref_leaves["beta"]] + if variant == "gdn2": + ref_args.append(ref_leaves["w"]) + ref_kwargs = dict(scale=scale, initial_state=state0_ref) + if case.varlen: + ref_kwargs["cu_seqlens"] = case.cu + o_ref, fs_ref = ref_fn(*ref_args, **ref_kwargs) + ref_outputs, ref_gos = [o_ref], [dO.double().reshape(o_ref.shape)] + if use_dfs: + ref_outputs.append(fs_ref) + ref_gos.append(dFS.double().reshape(fs_ref.shape)) + ref_grads = torch.autograd.grad(ref_outputs, list(ref_leaves.values()) + ([state0_ref] if use_initial_state else []), ref_gos) + + names = list(op_leaves) + (["initial_state"] if use_initial_state else []) + for name, got, want in zip(names, grads, ref_grads): + if name == "initial_state": + tol_n = STATE_GRAD_TOL + elif name in ("g", "beta", "w") and gate_grad_tol is not None: + tol_n = gate_grad_tol + else: + tol_n = tol + check(f"d{name}", got, want, tol_n) + + +@pytest.mark.parametrize("H,HV", HEAD_CONFIGS_SMALL + [(16, 64)]) +@pytest.mark.parametrize("T", [64, 128, 251]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=DTYPE_IDS.get) +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_parity(backend, variant, dtype, T, H, HV): + if dtype == torch.float16 and (T != 128 or (H, HV) != (1, 1)): + pytest.skip("fp16 runs one representative backward config") + if (H, HV) == (16, 64) and T != 128: + pytest.skip("the large GVA config runs one length") + assert_bwd_parity(backend, make_case(variant, dtype, T=T, H=H, HV=HV)) + + +@pytest.mark.parametrize("H,HK,HV", GQA_CONFIGS) +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_gqa(backend, variant, H, HK, HV): + assert_bwd_parity(backend, make_case(variant, torch.bfloat16, T=128, H=H, HK=HK, HV=HV)) + + +@pytest.mark.parametrize("seq_lens", [[64, 192], [31, 63, 93, 123]], ids=["two", "ragged"]) +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_varlen(backend, variant, seq_lens): + assert_bwd_parity(backend, make_case(variant, torch.bfloat16, seq_lens=seq_lens)) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_zero_length_sequence(backend, variant): + assert_bwd_parity(backend, make_case(variant, torch.bfloat16, seq_lens=[64, 0, 128])) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_initial_state(backend, variant): + assert_bwd_parity(backend, make_case(variant, torch.bfloat16, T=128), use_initial_state=True) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_d_final_state(backend, variant): + assert_bwd_parity(backend, make_case(variant, torch.bfloat16, T=128), use_initial_state=True, use_dfs=True) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_d_final_state_partial_chunk(backend, variant): + assert_bwd_parity(backend, make_case(variant, torch.bfloat16, T=251), use_dfs=True) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_scale(backend, variant): + """A non-default scale must reach the backward path (the engines carry an + independent 1/sqrt(K) default that would mask a dropped plumb).""" + assert_bwd_parity(backend, make_case(variant, torch.bfloat16, T=128), scale=1.0) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_qk_l2norm(backend, variant): + """dQ/dK must include the in-kernel normalization's own backward.""" + assert_bwd_parity(backend, make_case(variant, torch.bfloat16, T=128), l2norm=True) + + +def test_bwd_no_decay_gate_grad_floor(backend): + """GDN with alpha off: dGate/dBeta are cancelling-reduction noise floors; + the data grads stay at full tolerance.""" + assert_bwd_parity(backend, make_case("gdn", torch.bfloat16, T=192, alpha=False), gate_grad_tol=0.3) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_with_checkpoints(backend, variant): + """The checkpoint dump is non-differentiable and must not block backward.""" + ckpt = CHUNK[variant] + case = make_case(variant, torch.bfloat16, T=4 * ckpt) + q_t = to_thd(case.q).detach().clone().requires_grad_(True) + g_t = to_thd(case.gates["g"]).detach().clone().requires_grad_(True) + args = [q_t, to_thd(case.k), to_thd(case.v), g_t, to_thd(case.gates["beta"])] + if variant == "gdn2": + args.append(to_thd(case.gates["w"])) + with waive_unsupported(backend, variant): + o, fs, state_checkpoints = pinned_op(backend, variant)(*args, case.cu, output_final_state=True, checkpoint_every_n_tokens=ckpt) + assert not state_checkpoints.requires_grad + (o.sum() + fs.sum()).backward() + for name, t in (("q", q_t), ("g", g_t)): + assert t.grad is not None and torch.isfinite(t.grad).all(), f"bad grad for {name}" + + +# --------------------------------------------------------------------------- +# Checkpoints (per-chunk state series) +# --------------------------------------------------------------------------- + + +@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.""" + 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) + assert state_checkpoints.dtype == case.dtype + for j in (0, valid - 1): + n = (j + 1) * 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]) + + +@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.""" + 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] + # 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) + bounds = case.cu.tolist() + base = 0 + for n, cnt in enumerate(counts): + for j in sorted({0, cnt - 1} if cnt else set()): + n0 = bounds[n] + ntok = (j + 1) * 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]) + base += cnt + + +TIGHT_VARLEN_RECIPES = { + "pair+quarter": lambda c: [c + c // 4] * 2, + "triple+1": lambda c: [c + 1] * 3, + "single-2c+1": lambda c: [2 * c + 1], +} + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("recipe", sorted(TIGHT_VARLEN_RECIPES)) +@pytest.mark.parametrize("variant", VARIANTS) +def test_checkpoints_varlen_tight_capacity(backend, variant, recipe): + """Non-multiple varlen lengths where the packed entries fill the + host-computable capacity exactly (sum over seqs of ceil(T/ckpt) - 1 == + total // ckpt): every entry lands in bounds and matches the fp64 + recurrence and its sequence's solo prefix run; the chunk-cadence bwd + 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" + 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) + ref_fn = {"gdn": gdn_reference, "kda": kda_reference, "gdn2": gdn2_reference}[variant] + bounds = case.cu.tolist() + base = 0 + for n, cnt in enumerate(counts): + for j in range(cnt): + n0 = bounds[n] + ntok = (j + 1) * 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": + ref_args.append(case.gates["w"][:, n0 : n0 + ntok]) + 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]) + 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]) + base += cnt + grads_by_mode = [] + for mode_ckpt in (0, ckpt): + 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"])) + with waive_unsupported(backend, variant): + out = pinned_op(backend, variant)(*args, case.cu, checkpoint_every_n_tokens=mode_ckpt) + set_seed(SEED + 5) + dO = torch.randn_like(out[0]) + grads_by_mode.append(torch.autograd.grad([out[0]], leaves, [dO])) + for gr, gc in zip(grads_by_mode[0], grads_by_mode[1]): + assert torch.equal(bits(gr), bits(gc)), "checkpoint-reuse grads differ from the recompute path" + + +@pytest.mark.parametrize("ckpt_mult", [2, 3]) +@pytest.mark.parametrize("variant", VARIANTS) +def test_checkpoints_coarse_cadence(backend, variant, ckpt_mult): + """Coarser cadences (multiples of the base chunk) keep the prefix contract.""" + ckpt = CHUNK[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) + assert state_checkpoints.dtype == case.dtype + for j in (0, valid - 1): + n = (j + 1) * 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]) + + +# --------------------------------------------------------------------------- +# Raw-logit gate modes (safe gate, in-kernel Beta sigmoid) +# --------------------------------------------------------------------------- + + +def safe_gate_case(variant, T=256, H=2, K=128, V=128, seed=SEED + 7): + case = make_case(variant, torch.bfloat16, T=T, H=H, K=K, V=V, seed=seed) + set_seed(seed + 1) + graw = torch.randn(1, T, case.HO, K, device="cuda", dtype=torch.float32) + a_log = torch.zeros(case.HO, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(case.HO, K, dtype=torch.float32, device="cuda") + return case, graw, a_log, dt_bias + + +@pytest.mark.parametrize("variant", ["kda", "gdn2"]) +def test_safe_gate_forward_parity(backend, variant): + """Raw logits with a_log = 0 / dt_bias = 0 match the post-activation path + fed the host-side transform ``lb * sigmoid(g)``.""" + lb = -5.0 + case, graw, a_log, dt_bias = safe_gate_case(variant) + kw = dict(output_final_state=True, use_qk_l2norm_in_kernel=True) + raw_kw = dict(kw, safe_gate=True, gate_lower_bound=lb, a_log=a_log, dt_bias=dt_bias) + raw_gates = dict(case.gates, g=graw) + if variant == "kda": + braw = torch.randn(1, case.T, case.HO, device="cuda").to(case.dtype) + raw_gates["beta"] = braw + raw_kw["use_beta_sigmoid_in_kernel"] = True + eff_beta = braw.float().sigmoid() + else: + eff_beta = case.gates["beta"] + raw_case = case.clone(gates=raw_gates) + eff_case = case.clone(gates=dict(case.gates, g=lb * torch.sigmoid(graw), 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_ratio(fs_raw, fs_eff) < 2e-2 + + +@pytest.mark.parametrize("variant", ["kda", "gdn2"]) +def test_safe_gate_backward_raises(backend, variant): + """Raw-logit gate modes are forward-only by contract.""" + lb = -5.0 + case, graw, a_log, dt_bias = safe_gate_case(variant, T=128) + raw_gates = dict(case.gates, g=graw) + kw = dict(safe_gate=True, gate_lower_bound=lb, a_log=a_log, dt_bias=dt_bias, use_qk_l2norm_in_kernel=True) + if variant == "kda": + raw_gates["beta"] = torch.randn(1, case.T, case.HO, device="cuda").to(case.dtype) + kw["use_beta_sigmoid_in_kernel"] = True + raw_case = case.clone(gates=raw_gates) + g_leaf = to_thd(raw_gates["g"]).detach().clone().requires_grad_(True) + args = [to_thd(raw_case.q).detach().clone().requires_grad_(True), to_thd(raw_case.k), to_thd(raw_case.v), g_leaf, to_thd(raw_gates["beta"])] + if variant == "gdn2": + args.append(to_thd(raw_gates["w"])) + with waive_unsupported(backend, variant): + o, _ = pinned_op(backend, variant)(*args, case.cu, **kw) + with pytest.raises(NotImplementedError, match="forward-only"): + o.sum().backward() + + +def test_beta_sigmoid_in_kernel(backend): + """KDA: io-dtype Beta logits with the in-kernel sigmoid match the + post-activation fp32 path.""" + case = make_case("kda", torch.bfloat16, T=256) + set_seed(SEED + 11) + braw = torch.randn(1, case.T, case.HO, device="cuda").to(case.dtype) + raw_case = case.clone(gates=dict(case.gates, beta=braw)) + 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_ratio(fs_raw, fs_eff) < 2e-2 + + +# --------------------------------------------------------------------------- +# torch.compile +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_torch_compile_forward(backend, variant): + case = make_case(variant, torch.bfloat16, T=128) + with waive_unsupported(backend, variant): + o_eager, _ = pinned_op(backend, variant)(*op_args(case)) + compiled = torch.compile(pinned_op(backend, variant), fullgraph=True) + o_comp, _ = compiled(*op_args(case)) + torch.testing.assert_close(o_eager, o_comp) + + +# --------------------------------------------------------------------------- +# Argument validation (op-level contract; raises before engine selection) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_invalid_rank_raises(variant): + case = make_case(variant, torch.bfloat16, T=64) + with pytest.raises(ValueError, match="THD"): + op(variant)(case.q, to_thd(case.k), to_thd(case.v), *[to_thd(case.gates[n]) for n in case.gates], case.cu) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_invalid_qk_head_mismatch_raises(variant): + case = make_case(variant, torch.bfloat16, T=64, H=2) + args = [to_thd(case.q), to_thd(case.k)[:, :1].contiguous(), to_thd(case.v), to_thd(case.gates["g"]), to_thd(case.gates["beta"])] + if variant == "gdn2": + args.append(to_thd(case.gates["w"])) + with pytest.raises(ValueError, match="head count"): + op(variant)(*args, case.cu) + + +@pytest.mark.parametrize("variant", VARIANTS) +def test_invalid_gate_dtype_raises(variant): + case = make_case(variant, torch.bfloat16, T=64) + args = [to_thd(case.q), to_thd(case.k), to_thd(case.v), to_thd(case.gates["g"]).to(torch.bfloat16), to_thd(case.gates["beta"])] + if variant == "gdn2": + args.append(to_thd(case.gates["w"])) + with pytest.raises(TypeError, match="must be"): + op(variant)(*args, case.cu) + + +@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) + with pytest.raises(ValueError, match="initial"): + op(variant)(*op_args(case), initial_state=state0) + + +@pytest.mark.parametrize("variant", ["kda", "gdn2"]) +def test_invalid_safe_gate_args_raise(variant): + case = make_case(variant, torch.bfloat16, T=64) + with pytest.raises(ValueError, match="safe_gate"): + op(variant)(*op_args(case), safe_gate=True) + with pytest.raises(ValueError, match="safe_gate"): + op(variant)(*op_args(case), a_log=torch.zeros(case.HO, device="cuda")) + + +# --------------------------------------------------------------------------- +# Determinism (contract held by the FROST backend) +# --------------------------------------------------------------------------- + + +def bits(t): + return t.contiguous().view(torch.uint8) + + +def assert_bitwise_runs(launch, repeats=DETERMINISM_REPEATS, label=""): + """Back-to-back launches (single sync) must match run 0 bit for bit — + barrier/fence races are timing-dependent, so there is no tolerance.""" + runs = [launch() for _ in range(repeats)] + torch.cuda.synchronize() + for out in runs[0]: + assert torch.isfinite(out.float()).all(), f"{label}: non-finite output in run 0" + for r, outs in enumerate(runs[1:], start=1): + for i, (a, b) in enumerate(zip(runs[0], outs)): + assert torch.equal(bits(a), bits(b)), f"{label}: output {i} differs between run 0 and run {r}" + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@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 + + def launch(): + o, fs = run_fwd(backend, case, initial_state=state0, output_final_state=True) + return o, fs + + assert_bitwise_runs(launch, label=f"{variant} fwd") + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_determinism_bwd(backend, variant): + case = make_case(variant, torch.bfloat16, seq_lens=[497, 16, 1, 480, 0, 253]) + 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"])) + with waive_unsupported(backend, variant): + o, fs = pinned_op(backend, variant)(*args, case.cu) + dO = torch.randn_like(o) + + def launch(): + return torch.autograd.grad([o], leaves, [dO], retain_graph=True) + + assert_bitwise_runs(launch, label=f"{variant} bwd") + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +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 + + def launch(): + o, fs = run_fwd(backend, case, initial_state=state0, output_final_state=True) + return o, fs + + assert_bitwise_runs(launch, label=f"{variant} multi-tile fwd") + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_determinism_multi_tile_bwd(backend, variant): + case = make_case(variant, torch.bfloat16, B=8, T=192, H=64) + 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"])) + with waive_unsupported(backend, variant): + o, fs = pinned_op(backend, variant)(*args, case.cu) + dO = torch.randn_like(o) + + def launch(): + return torch.autograd.grad([o], leaves, [dO], retain_graph=True) + + assert_bitwise_runs(launch, label=f"{variant} multi-tile bwd") + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_determinism_two_streams(backend, variant): + """Two concurrent instances on separate streams must not perturb each + other: every repeat matches its own single-stream baseline.""" + case_a = make_case(variant, torch.bfloat16, seq_lens=[497, 16, 1, 480, 0, 253], seed=SEED) + case_b = make_case(variant, torch.bfloat16, B=2, T=512, seed=SEED + 1) + launch_a = lambda: run_fwd(backend, case_a, output_final_state=True) # noqa: E731 + launch_b = lambda: run_fwd(backend, case_b, output_final_state=True) # noqa: E731 + s1, s2 = torch.cuda.Stream(), torch.cuda.Stream() + # order the side streams behind the input generation (default stream) + torch.cuda.synchronize() + + with torch.cuda.stream(s1): + base_a = launch_a() + torch.cuda.synchronize() + with torch.cuda.stream(s2): + base_b = launch_b() + torch.cuda.synchronize() + for r in range(DETERMINISM_REPEATS): + with torch.cuda.stream(s1): + out_a = launch_a() + with torch.cuda.stream(s2): + out_b = launch_b() + torch.cuda.synchronize() + for label, base, outs in (("A", base_a, out_a), ("B", base_b, out_b)): + for i, (x, y) in enumerate(zip(base, outs)): + assert torch.equal(bits(x), bits(y)), f"stream {label} output {i} differs on concurrent run {r}" + + +# --------------------------------------------------------------------------- +# Batch invariance (whole-sequence work items; packed == solo, bitwise) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_batch_invariance_fwd(backend, variant): + """batch_invariant=True: each packed sequence matches its solo B = 1 run bitwise.""" + case = make_case(variant, torch.bfloat16, seq_lens=[497, 16, 1, 480, 0, 253]) + o, fs = run_fwd(backend, case, output_final_state=True, batch_invariant=True) + bounds = case.cu.tolist() + for n in range(case.N): + s, e = bounds[n], bounds[n + 1] + if s == e: + continue + args = [to_thd(t)[s:e].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"])[s:e].clone()) + cu1 = torch.tensor([0, e - s], dtype=torch.int32, device="cuda") + with waive_unsupported(backend, variant): + o_solo, fs_solo = pinned_op(backend, variant)(*args, cu1, output_final_state=True, batch_invariant=True) + assert torch.equal(bits(o[s:e]), bits(o_solo)), f"seq {n}: packed o differs from solo" + assert torch.equal(bits(fs[n]), bits(fs_solo[0])), f"seq {n}: packed final state differs from solo" + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_batch_invariance_bwd(backend, variant): + """batch_invariant=True: a sequence's grads match its solo B = 1 run bitwise.""" + case = make_case(variant, torch.bfloat16, seq_lens=[497, 16, 1, 480, 0, 253]) + 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"])) + s, e = 0, 497 + with waive_unsupported(backend, variant): + o, fs = pinned_op(backend, variant)(*args, case.cu, batch_invariant=True) + dO = torch.randn_like(o) + grads_packed = torch.autograd.grad([o], leaves, [dO], retain_graph=True) + solo_args = [to_thd(t)[s:e].clone() for t in (case.q, case.k, case.v, case.gates["g"], case.gates["beta"])] + if variant == "gdn2": + solo_args.append(to_thd(case.gates["w"])[s:e].clone()) + solo_leaves = [solo_args[0].requires_grad_(True), solo_args[1].requires_grad_(True)] + cu1 = torch.tensor([0, e - s], dtype=torch.int32, device="cuda") + o_solo, fs_solo = pinned_op(backend, variant)(*solo_args, cu1, batch_invariant=True) + grads_solo = torch.autograd.grad([o_solo], solo_leaves, [dO[s:e].clone()]) + for gp, gs in zip(grads_packed, grads_solo): + assert torch.equal(bits(gp[s:e]), bits(gs)), "packed grad slice differs from solo grad" + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_bwd_checkpoint_reuse(backend, variant): + """Training with chunk-cadence checkpoints: the bwd consumes the fwd's + series instead of recomputing it, and the grads match bitwise.""" + case = make_case(variant, torch.bfloat16, seq_lens=[497, 16, 1, 480, 0, 253]) + grads_by_mode = [] + for ckpt in (0, CHUNK[variant]): + 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"])) + with waive_unsupported(backend, variant): + out = pinned_op(backend, variant)(*args, case.cu, checkpoint_every_n_tokens=ckpt) + o = out[0] + set_seed(SEED + 5) + dO = torch.randn_like(o) + grads_by_mode.append(torch.autograd.grad([o], leaves, [dO])) + for gr, gc in zip(grads_by_mode[0], grads_by_mode[1]): + assert torch.equal(bits(gr), bits(gc)), "checkpoint-reuse grads differ from the recompute path" + + +@pytest.mark.parametrize("backend", ["cutile"], indirect=True) +@pytest.mark.parametrize("variant", ["gdn", "kda"]) +def test_batch_invariance_cutile(backend, variant): + """cuTile is batch-invariant by construction; the flag must hold there too.""" + case = make_case(variant, torch.bfloat16, seq_lens=[497, 16, 1, 480, 0, 253]) + o, fs = run_fwd(backend, case, output_final_state=True, batch_invariant=True) + bounds = case.cu.tolist() + for n in range(case.N): + s, e = bounds[n], bounds[n + 1] + if s == e: + continue + args = [to_thd(t)[s:e].clone() for t in (case.q, case.k, case.v, case.gates["g"], case.gates["beta"])] + cu1 = torch.tensor([0, e - s], dtype=torch.int32, device="cuda") + with waive_unsupported(backend, variant): + o_solo, fs_solo = pinned_op(backend, variant)(*args, cu1, output_final_state=True, batch_invariant=True) + assert torch.equal(bits(o[s:e]), bits(o_solo)), f"seq {n}: packed o differs from solo" + assert torch.equal(bits(fs[n]), bits(fs_solo[0])), f"seq {n}: packed final state differs from solo" + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_batch_invariance_with_coarse_checkpoints(backend, variant): + """batch_invariant=True composes with a coarser checkpoint cadence.""" + ckpt = CHUNK[variant] * 2 + 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) + 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": + 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) + check("state_checkpoints[0]", state_checkpoints[0], fs_p[0], STATE_TOL[case.dtype]) + + +# --------------------------------------------------------------------------- +# CUDA-graph replay (contract held by the FROST backend) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("backend", ["frost"], indirect=True) +@pytest.mark.parametrize("variant", VARIANTS) +def test_cuda_graph_replay_fwd(backend, variant): + case = make_case(variant, torch.bfloat16, B=2, T=256) + + def launch(): + return pinned_op(backend, variant)(*op_args(case), output_final_state=True) + + with waive_unsupported(backend, variant): + eager = launch() + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + launch() + torch.cuda.current_stream().wait_stream(warmup) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = launch() + graph.replay() + torch.cuda.synchronize() + for i, (a, b) in enumerate(zip(eager, captured)): + assert torch.equal(bits(a), bits(b)), f"replayed output {i} differs from eager"